语义路由系统面临着扩展性挑战。当每个分类请求都需要独立运行多个微调模型时，计算成本会随着模型数量的增加而线性增长。本文将探讨 vLLM 语义路由器的 Rust 分类层最近的一次重构，如何通过架构模块化、低秩适应（LoRA）和并发优化来解决这个问题。

背景：从 BERT 到模块化系统

之前的实现主要依赖 BERT 和 ModernBERT 进行意图分类和越狱检测。虽然 ModernBERT 在英文文本分类任务上表现出色，但它有以下局限性：

  • 语言覆盖范围：与在更多样化数据集上训练的模型相比，原始 ModernBERT 的多语言支持有限。（注：mmBERT 是 ModernBERT 的一个大规模多语言变体，支持超过 1800 种语言，它在本次重构开始后发布，代表了解决多语言挑战的另一种方法）
  • 上下文长度：虽然 ModernBERT 使用 RoPE 将上下文扩展到 8,192 个 token（来源），但像 Qwen3-Embedding 这样的模型支持高达 32,768 个 token，这对于处理超长文档非常有利。
  • 模型耦合：分类逻辑与特定的模型架构紧密耦合，使得添加新模型变得困难。

这些限制促使我们进行了一次更广泛的重构，旨在让系统能够在保持性能的同时支持多种模型类型。模块化架构意味着像 mmBERT 这样的新模型可以与 Qwen3-Embedding 和 EmbeddingGemma 一起集成，从而让路由器能够为每个任务选择最合适的模型。

架构重组

这次重构在 candle-binding crate 中引入了一个分层架构。这种结构实现了关注点分离：核心功能与具体模型保持独立，而添加新的模型架构则无需修改现有代码。`DualPathUnifiedClassifier` 实现了一种路由逻辑，可以根据任务需求在传统的微调模型和 LoRA 适配模型之间进行选择。

长上下文嵌入模型

两个新的嵌入模型解决了上下文长度的限制：

Qwen3-Embedding

Qwen3-Embedding 支持高达 32,768 个 token 的上下文长度（Hugging Face 模型卡）。其实现使用了 RoPE（旋转位置嵌入），通过在长距离上提高频率分辨率来实现这种扩展的上下文处理能力。

Qwen3-Embedding 在超过 100 种语言的文本上进行了训练（Hugging Face 模型卡），这使其非常适合多语言路由场景，而之前仅使用 ModernBERT 的方法在这种场景下会遇到困难。

EmbeddingGemma-300M

谷歌的 EmbeddingGemma-300M 采取了不同的方法，专注于在保持高质量的同时减小模型尺寸。该模型支持 2,048 个 token 的上下文长度，并实现了 Matryoshka 表示学习，这意味着嵌入向量可以被截断到 768、512、256 或 128 维而无需重新训练（Hugging Face 模型卡）。

该架构使用了多查询注意力（MQA），包含 3 个查询头和 1 个键值头，从而降低了内存带宽需求。一个显著的特点是在 Transformer 模块之后应用了一个密集的瓶颈层（768 → 3072 → 768），这根据 Matryoshka 训练方法提高了嵌入质量。

用于多任务分类的低秩适应（LoRA）

LoRA 解决了一个先前系统中根本性的低效问题。当一个分类系统需要确定意图、检测个人可识别信息（PII）并检查安全问题时，最直接的方法是运行三个独立的微调模型：

每个模型都需要让输入数据通过其整个网络，包括计算成本高昂的基础 Transformer 层。这导致了 O(n) 的复杂度，其中 n 是分类任务的数量。

LoRA 通过共享基础模型的计算来改变这一状况：

基础模型只运行一次，生成中间表示。然后，每个 LoRA 适配器应用特定于任务的低秩权重更新来定制输出。由于 LoRA 适配器通常只修改不到 1% 的模型参数，因此这最后一步比运行完整的模型要快得多。

在 `parallel_engine.rs` 中的实现使用了 Rayon 进行数据并行，可以并发处理多个 LoRA 适配器。对于一个需要进行三次分类的请求，这将工作负载从三次完整的前向传播变为一次完整的前向传播加上三次轻量级的适配器应用。

通过 `OnceLock` 实现并发

之前的实现使用 `lazy_static` 来管理全局分类器状态，这在并发负载下会引入锁竞争。重构后，我们用 Rust 标准库中的 `OnceLock` 替换了它。

`OnceLock` 在初始化后提供无锁读取。首次初始化之后，所有后续访问都只是简单的指针读取，没有同步开销。`oncelock_concurrent_test.rs` 中的测试通过 10 个并发线程执行总共 30 次分类来验证了这一点，确认了吞吐量会随着线程数的增加而线性扩展。

当路由器处理多个传入请求时，这一点至关重要。使用 `lazy_static`，并发请求会排队等待一个互斥锁。而使用 `OnceLock`，它们可以无竞争地并行执行。

用于 GPU 加速的 Flash Attention

Flash Attention 2 支持作为 CUDA 构建的一个可选功能提供，但它要求使用 Ampere 代或更新的 GPU（计算能力 ≥ 8.0）。Flash Attention 通过将计算分块处理，使其能容纳在高速的片上 SRAM 内存中，从而避免了从较慢的 GPU DRAM 中重复读取数据，优化了注意力机制。

ModernBERT 和 Qwen3 都从 Flash Attention 的集成中受益：

  • ModernBERT：自注意力计算速度提升高达 3 倍，同时显著减少了内存使用（来源）。该模型还使用交替的注意力模式（每三层使用全局注意力，其余使用局部滑动窗口注意力）来平衡效率和上下文保留（来源）。

  • Qwen3：集成 FlashAttention-2 使注意力操作的速度提升高达 4 倍。对于 14B 变体，这意味着推理速度达到 70-110 tokens/秒，而没有它时为 30-35 tokens/秒——这种性能提升随着上下文变长而更加显著（来源）。

Rust 的实现通过 Cargo features 使 Flash Attention 成为可选功能，这使得它可以在没有兼容 GPU 的系统上部署，同时在硬件支持时能够实现显著的性能提升。

为云原生生态系统进行跨语言集成

选择 Rust 作为核心分类引擎，并结合 Go FFI（外部函数接口）绑定，解决了云原生环境中的一个实际部署挑战。

为什么选择 Rust 进行机器学习推理？

Rust 为分类层提供了几个优势：

  • 性能：接近 C 语言的性能和零成本抽象，对低延迟推理至关重要。
  • 内存安全：编译时保证防止了诸如缓冲区溢出和悬垂指针之类的常见错误。
  • 并发性：所有权系统防止了数据竞争，使得通过 Rayon 进行安全的并行处理成为可能。
  • 无垃圾回收：可预测的延迟，没有会影响请求处理的垃圾回收暂停。

Candle 框架利用了 Rust 的这些优点，同时为机器学习模型开发提供了一个熟悉的 API。

为什么 Go FFI 绑定很重要？

虽然 Rust 在计算密集型的机器学习推理方面表现出色，但 Go 在云原生基础设施生态系统中占据主导地位。FFI 层将这两个世界连接起来。这种集成使得在以 Go 为主要语言的环境中进行部署成为可能：

  • Envoy 代理集成：语义路由器作为一个用 Go 编写的 Envoy 外部处理过滤器 运行。FFI 允许 Go 过滤器利用高性能的 Rust 分类功能，而无需重写整个 Envoy 集成层。
  • Kubernetes Operators：云原生 Operator 通常使用 controller-runtime 以 Go 语言编写。FFI 使这些 Operator 能够直接嵌入分类逻辑，而不是通过网络调用分离的服务。
  • 服务网格：像 Istio、Linkerd 和 Consul 这样的项目都是基于 Go 的。FFI 允许路由决策使用基于机器学习的分类，同时保持与现有网格控制平面的兼容性。
  • API 网关：许多 API 网关（如 Kong、Tyk）都有 Go 组件。FFI 使得在网关层实现语义路由成为可能，而无需引入额外的微服务。

部署灵活性

这种双语言架构提供了多种部署选项：

  • 嵌入式模式：Go 服务通过 CGO 直接链接到 Rust 库，最大限度地减少了延迟和部署复杂性。
  • 进程隔离：分类层可以作为一个独立的进程运行，通过 gRPC 或 Unix 套接字进行通信，以提供额外的故障隔离。
  • 混合工作负载：服务可以结合 Go 在网络和编排方面的优势以及 Rust 在机器学习推理方面的性能。

语义路由器广泛利用了这种模式。主要的路由逻辑、配置管理和缓存实现都在 Go 中，而计算密集型的分类则在 Rust 中运行。这种分离使得每个组件都可以使用最合适的语言，同时通过 FFI 层保持清晰的接口。

性能特点

这种架构的好处因工作负载而异：

  • 单任务与多任务分类：由于没有基础模型共享，LoRA 带来的好处微乎其微。传统的微调模型可能更快。当对同一输入进行多次分类时，LoRA 显示出明显优势。由于基础模型只运行一次，每个任务只执行 LoRA 适配器，与运行多个完整模型相比，开销大幅降低。实际的加速比取决于基础模型计算与适配器计算的比例。
  • 长上下文输入：Qwen3-Embedding 能够对长达 32K token 的文档进行路由决策而无需截断，超越了 ModernBERT 对超长文档的 8K 限制。在兼容的 GPU 上启用 Flash Attention 2 后，随着上下文长度的增加，性能优势变得更加显著。
  • 多语言路由：模型现在可以处理 ModernBERT 训练数据有限的语言的路由决策。
  • 高并发性：`OnceLock` 消除了锁竞争，使得分类操作的吞吐量能够随着 CPU 核心数的增加而扩展。
  • GPU 加速：启用 Flash Attention 2 后，注意力操作的运行速度提高了 3-4 倍，并且在序列长度较长时加速效果更为明显。这使得 GPU 部署在高吞吐量场景中尤其具有优势。

未来方向

模块化架构为未来的扩展提供了可能：

  • 通过实现 `CoreModel` trait 来添加更多的嵌入模型。
  • 当 Candle 支持 Flash Attention 3 时，集成该功能。
  • 支持量化（4 位、8 位）以减少内存占用。
  • 为特定领域的路由定制 LoRA 适配器。
  • 为其他语言（如 Python、Java、C++）提供 FFI 绑定，以扩展集成可能性。

该系统现在有了一个坚实的基础，可以在不进行架构变更的情况下融入新的研究进展。FFI 层提供了一个稳定的接口，允许 Rust 实现独立演进，同时保持与现有基于 Go 的部署的兼容性。

资源