对于运行多个自定义 AI 模型（尤其是近期的混合专家模型，MoE）的企业和个人来说，如果单个模型的流量不足以充分利用专用的计算终端节点，他们往往需要为闲置的 GPU 容量买单。为了解决这个问题，我们与 vLLM 社区合作，针对 GPT-OSS 或 Qwen 等流行的开源 MoE 模型开发了一种高效的 Multi-LoRA（多低秩自适应）服务方案。Multi-LoRA 是一种常用的模型微调方法。它不重训练整个模型权重，而是保持原始权重冻结，并在模型的层中注入小型的、可训练的适配器（adapters）。通过 Multi-LoRA，在推理时，多个自定义模型可以共享同一个 GPU，每次请求仅需切换适配器即可。例如，原本 5 个客户各自使用 10% 的专用 GPU，现在可以通过 Multi-LoRA 在单块 GPU 上提供服务，将 5 个利用率不足的 GPU 转变为 1 个高效共享的 GPU。

在这篇文章中，我们将解释如何在 vLLM 中为混合专家模型（MoE）实现 Multi-LoRA 推理，描述我们执行的内核级优化，并展示您如何从中受益。全文将以 GPT-OSS 20B 为主要示例。

您可以从 0.15.0 及更高版本开始，在本地 vLLM 部署中使用这些改进功能。Multi-LoRA 服务现已适用于包括 GPT-OSS、Qwen3-MoE、DeepSeek 和 Llama MoE 在内的 MoE 模型系列。我们的优化还改善了密集模型（如 Llama3.3 70B 或 Qwen3 32B）的 Multi-LoRA 托管性能。针对 Amazon 的特定优化在 vLLM 0.15.0 的基础上带来了进一步的延迟改善，例如，GPT-OSS 20B 的输出 Token 每秒（OTPS，即模型生成输出的速度）提升了 19%，首 Token 延迟（TTFT，即模型开始生成输出前的等待时间）降低了 8%。要受益于这些优化，请将您的 LoRA 自定义模型托管在 Amazon SageMaker AI 或 Amazon Bedrock 上。

在 vLLM 中实现 MoE 模型的 Multi-LoRA 推理

在深入探讨我们为 vLLM 中的 MoE 模型实现 Multi-LoRA 推理的过程之前，我们先提供一些关于 MoE 模型和 LoRA 微调的背景信息，这对理解我们优化方案背后的逻辑至关重要。MoE 模型包含多个被称为“专家”的专用神经网络。路由算法将每个输入 Token 指向最相关的专家，然后将其输出进行汇总。这种稀疏架构能以更少的计算资源处理更大的模型，因为每个 Token 仅激活模型总参数的一小部分（详见下图 1 的可视化）。

每个专家都是一个小型的全连接网络，分两个阶段处理 Token 的隐藏状态。首先，gate_up 投影将紧凑的隐藏状态（例如 4096 维度）扩展为更大的中间空间（例如 11008 维度）。这种扩展是必要的，因为紧凑空间中的特征纠缠紧密，较大的空间为网络提供了拉伸、转换并有选择性地控制重要特征的余地。其次，down 投影将结果压缩回原始维度。这有助于保持输出与模型其余部分的兼容性，并作为瓶颈，强制网络仅保留最有用的特征。总之，这种“先扩展再压缩”的模式使每个专家既能执行丰富的转换，又能保持输出大小一致。vLLM 使用 fused_moe 内核将这些投影作为分组通用矩阵乘法（Group GEMM）操作来执行——为分配给特定 Token 的每个专家执行一个 GEMM。Multi-LoRA 微调保持基础模型权重 W（例如 gate_up 投影的 W_gate_up）冻结，并训练两个小矩阵 A 和 B，它们共同构成一个适配器。对于形状为 h_in × h_out 的基础权重 W，LoRA 训练形状为 h_in × r 的 A 和形状为 r × h_out 的 B，其中 r 是 LoRA 秩（通常为 16-64）。微调后的输出变为 y = xW + xAB。每个 LoRA 适配器为投影增加了两个操作。收缩（shrink）操作计算 z=xA，将输入从 h_in 维度降至 r 维度。扩展（expand）操作将该 r 维结果与 B 相乘，投影回 h_out 维度。这在图 1 的右侧进行了说明。

图 1：MoE-LoRA 模型的工作原理示意图，示例隐藏状态维度为 4096，中间表示维度为 11008，LoRA 秩 r = 32。

每个专家都有两个权重投影：gate_up 和 down。应用 LoRA 适配器时，它会为每个投影增加两个低秩操作，即收缩和扩展。这意味着每个专家总共需要 4 个 LoRA 内核操作：gate_up 的收缩和扩展，以及 down 的收缩和扩展。在多用户或多任务同时使用多个 LoRA 适配器的 Multi-LoRA 服务设置中，系统必须高效地管理每个专家、每个适配器、每次请求的这 4 个操作。这是 MoE 模型的核心性能瓶颈。这 4 个操作涉及矩阵运算，其中一个维度（LoRA 秩 r）比另一个维度（如隐藏状态和中间表示维度）小 100-300 倍。标准的 GEMM 内核是为近似方形的矩阵设计的，在细长矩阵上的表现较差，这就是为什么本文稍后描述的内核优化是必要的。除了针对细长矩阵进行优化外，为 MoE 模型添加 Multi-LoRA 支持还存在两个技术挑战。首先，vLLM 缺乏在 MoE 层上执行 LoRA 的内核，因为现有的密集型 Multi-LoRA 内核无法处理专家路由。其次，MoE LoRA 结合了两种稀疏性来源：专家路由（Token 被分配给不同的专家）和适配器选择（请求使用不同的 LoRA 适配器）。这种复合稀疏性需要专门的内核设计。为了应对这些挑战，我们创建了一个 fused_moe_lora 内核，将 LoRA 操作集成到 fused_moe 内核中。该新内核为 gate_up 和 down 投影执行 LoRA 收缩和扩展 GEMM。fused_moe_lora 内核遵循与 fused_moe 内核相同的逻辑，并为相应的已激活 LoRA 适配器在网格中增加了一个维度。

改进 vLLM 中的 Multi-LoRA 推理性能

在完成初步实现后，我们使用 NVIDIA Nsight Systems (Nsys) 来识别瓶颈，并发现 fused_moe_lora 内核是延迟最高的部分。然后，我们使用 NVIDIA Nsight Compute (NCU) 对 4 个内核操作的计算和内存吞吐量进行了性能分析：gate_up_shrink、gate_up_expand、down_shrink 和 down_expand。这些发现促使我们开发了执行优化、内核级优化以及针对这 4 个内核的调优配置。

执行优化

在我们的初步实现中，Multi-LoRA 的 TTFT 比基础模型（即 GPT-OSS 20B 的公开发布版本）高出 10 倍（表现更差）。我们的性能分析显示，Triton 编译器将依赖于输入长度的变量视为编译时常量，导致 fused_moe_lora 内核在每次新的上下文长度下都被重新编译，而不是被重用。这在图 2 中很明显：在每次 fused_moe_lora 内核执行之前的 cuModuleLoadData 调用表明 GPU 正在加载新编译的内核二进制文件，而不是重用缓存的二进制文件，内核启动时间之间的巨大间隙显示了 GPU 在重新编译期间处于闲置状态。这种开销导致了相比基础模型 10 倍的 TTFT 回归。我们通过为这些变量添加 do_not_specialize 编译器提示解决了这个问题，指示 Triton 将内核编译一次并将其重用于所有上下文长度。

图 2：在执行优化之前 fused_moe_lora 内核的性能分析结果。

内核（Kernel）优化

Split-K 是一种工作分解策略，有助于改善细长矩阵的负载均衡。LoRA 收缩计算 xA，其中 x 的维度为 1×h_in，A 的维度为 h_in×r。r 个输出元素中的每一个都需要对 h_in 次乘法进行求和。标准 GEMM 内核将不同的线程组（共享片上内存的 GPU 线程批次）分配给不同的输出元素，但每个线程组都是按顺序计算其 h_in 求和的。由于 r 在几十个量级而 h_in 在几千个量级，导致并行化的输出元素很少，而每个元素都需要进行冗长的序列求和。Split-K 通过将 GEMM 内部维度 K（本例中 K=h_in）上的求和拆分到多个线程组来解决这个问题，这些线程组并行计算部分和，然后结合它们的结果。这些部分结果需要原子加法（atomic add）来生成最终和。由于我们执行的是纯原子加法且没有额外逻辑，我们通过将原子加法操作的参数 sem="relaxed" 进行设置，利用了 Triton 编译器的优化自由度。

GPU 调度程序将多个线程组分配给同一个输出元素，并同时运行针对不同输出元素的线程组。对于 lora_shrink，每个输出元素都需要读取 A 的一列，该列跨越 h_in 行。由于 h_in 在数千量级，每一列触及的缓存行分散在较大的内存区域中。相邻的列共享相同的行并在缓存中重叠，因此处理相邻列的线程组可以从重用彼此加载的数据中受益。协同线程阵列（CTA）洗牌（swizzling）重新排序了调度，使得处理相邻列的线程组同时运行，从而增加了 L2 缓存重用率。我们将 CTA 洗牌应用于 lora_shrink 操作。

我们还从 LoRA 收缩和扩展内核中删除了不必要的掩码（masking）和点积运算。Triton 内核以固定大小的块加载数据，但矩阵维度可能无法被这些块大小均匀整除。例如，如果 BLOCK_SIZE_K 为 64，但矩阵维度 K 为 100，第二个块将尝试读取 28 个无效的内存位置。掩码有助于通过检查每个索引在加载前是否在边界内来防止这些非法内存访问。然而，这些条件检查在每次加载操作上执行，即使元素是有效的，也会增加开销。我们引入了一个 EVEN_K 参数，用于检查 K 是否能被 BLOCK_SIZE_K 均匀整除。当为真时，加载有效，可以完全跳过掩码检查，从而有助于减少掩码开销和不必要的点积计算。

最后，我们将 LoRA 权重的加法与基础模型权重的加法融合到了 LoRA 扩展内核中。这种优化有助于减少内核启动开销。这些内核优化使我们能够达到 GPT-OSS 20B 144 OTPS 和 135 ms TTFT 的性能。

针对 Amazon SageMaker AI 和 Amazon Bedrock 的内核配置调优

Triton 内核需要调优诸如块大小（BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K）等参数，这些参数控制矩阵计算如何在线程组间进行分配。高级参数包括 GROUP_SIZE_M（控制线程组排序以优化缓存局部性）和 SPLIT_K（在内部矩阵维度上并行化求和）。

我们发现，使用针对标准融合 MoE 优化的默认配置的 MoE LoRA 内核，在 Multi-LoRA 服务中表现不佳。这些默认配置没有考虑对应 LoRA 索引的额外网格维度以及来自多个适配器的复合稀疏性。为了解决这一瓶颈，我们添加了支持，允许用户通过提供文件夹路径来加载自定义调优配置。有关更多信息，请参阅 vLLM LoRA 调优文档。我们同时调优了 4 个 fused_moe_lora 操作（gate_up_shrink, gate_up_expand, down_shrink, down_expand），因为它们共享相同的 BLOCK_SIZE_M 参数。Amazon SageMaker AI 和 Bedrock 的客户现在可以使用这些自动加载的调优配置，GPT-OSS 20B 的性能达到了 171 OTPS 和 124 ms TTFT。

结果与结论

通过与 vLLM 社区的合作，我们实现了包括 GPT-OSS、Qwen3 MoE、DeepSeek 和 Llama MoE 在内的 MoE 模型的 Multi-LoRA 服务并进行了开源。随后，我们应用了优化，例如在 vLLM 0.15.0 中与 vLLM 0.11.1rc3 相比，GPT-OSS 20B 的 OTPS 提升了 454%，TTFT 降低了 87%。一些优化（特别是内核调优和 CTA 洗牌）也改善了密集模型的性能，例如 Qwen3 32B 的 OTPS 提升了 99%。要在本地部署中利用这些工作，请使用 vLLM 0.15.0 或更高版本。Amazon Bedrock 和 Amazon SageMaker AI 中提供的 Amazon 特有优化有助于在各模型中实现额外的延迟改进，例如与 vLLM 0.15.0 相比，GPT-OSS 20B 的 OTPS 快 19%，TTFT 好 8%。要开始在 Amazon 上托管自定义模型，请参阅 Amazon SageMaker AI 托管和 Amazon Bedrock 文档。

图 3：GPT-OSS 20B Multi-LoRA 推理的每秒输出 Token 数 (OTPS) 和首 Token 延迟 (TTFT)：1/ vLLM 0.11.1rc3 中的初始实现；2/ vLLM 0.15.0；3/ vLLM 0.15.0 配合 AWS 自定义内核调优。实验使用了 1600 个输入 Token 和 600 个输出 Token，LoRA 秩为 32，并并行加载 8 个适配器。

致谢

我们要感谢来自 vLLM 社区的贡献者和合作者：Jie Li、Chen Wu、Varun Sundar Rabindranath、Simon Mo 和 Robert Shaw，以及我们的团队成员：Xin Yang、Sadaf Fardeen、Ashish Khetan 和 George Karypis。

也发布在 AWS 博客上。