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

背景：从 BERT 到模块化系统

之前的实现主要依赖 BERT 和 ModernBERT 进行意图识别和越狱（jailbreak）分类。虽然 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

Google 的 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 个 token，而没有它时为 30-35 个 token —— 随着上下文长度的增加，这种性能提升变得更加明显（来源）。

Rust 实现通过 Cargo features 使 Flash Attention 成为可选，允许在没有兼容 GPU 的系统上部署，同时在硬件支持时实现实质性的性能提升。

云原生生态系统的跨语言集成

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

为什么在 ML 推理中使用 Rust

Rust 为分类层提供了多项优势：

  • 性能：具有零成本抽象的接近 C 语言的性能，这对于低延迟推理至关重要
  • 内存安全：编译时保证可防止常见的错误，如缓冲区溢出和释放后使用错误
  • 并发性：所有权系统防止数据竞争，从而实现使用 Rayon 的安全并行处理
  • 无垃圾回收：可预测的延迟，没有会影响请求处理的 GC 停顿

Candle 框架利用了 Rust 的这些优势，同时为 ML 模型开发提供了熟悉的 API。

为什么 Go FFI 绑定很重要

虽然 Rust 擅长计算密集型的 ML 推理，但 Go 在云原生基础设施生态系统中占据主导地位。FFI 层桥接了这两个世界。这种集成使得可以在以 Go 为主要语言的环境中部署：

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

部署灵活性

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

  • 嵌入模式：Go 服务通过 CGO 直接链接到 Rust 库，最大限度地减少延迟和部署复杂性
  • 进程隔离：分类层可以作为独立进程运行，通过 gRPC 或 Unix 套接字通信，以实现额外的故障隔离
  • 混合工作负载：服务可以将 Go 的网络和编排优势与 Rust 的 ML 推理性能相结合

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

性能特征

这种架构的优势因工作负载而异：

  • 单任务 vs 多任务分类：由于没有基础模型共享，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-bit, 8-bit）以减少内存占用
  • 用于领域特定路由的自定义 LoRA 适配器
  • 为更多语言（Python、Java、C++）提供 FFI 绑定，以扩大集成可能性

该系统现在已经有了集成新研究成果的基础，而无需更改架构。FFI 层提供了一个稳定的接口，允许 Rust 实现在保持与现有基于 Go 部署的兼容性的同时独立演进。

资源