对于运行多个自定义 AI 模型（尤其是近期流行的混合专家模型，即 MoE 模型系列）的组织和个人来说，当单个模型的流量不足以填满专用计算端点时，往往会面临为空闲 GPU 容量买单的挑战。为了解决这个问题，我们与 vLLM 社区合作，针对 GPT-OSS 或 Qwen 等主流开源 MoE 模型开发了一种高效的多低秩自适应（Multi-LoRA）服务解决方案。Multi-LoRA 是一种流行的模型微调方法。它不是重训练整个模型权重，而是保持原始权重冻结，并将小的、可训练的适配器注入到模型的层中。使用 Multi-LoRA，在推理时，多个自定义模型共享同一个 GPU，每次请求只需切换适配器即可。例如，五个客户各占用一个专用 GPU 的 10%，通过 Multi-LoRA 可以将它们全部服务于单个 GPU，从而将五个利用率不足的 GPU 转变为一个高效共享的 GPU。

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

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

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

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

每个专家都是一个小型前馈网络，分两个阶段处理令牌的隐藏状态。首先，gate_up 投影将紧凑的隐藏状态（例如 4096 维）扩展为更大的中间空间（例如 11008 维）。这种扩展是必要的，因为紧凑空间中的特征纠缠在一起——更大的空间让网络有余地将它们拉开、转换并有选择地控制哪些特征重要。其次，down 投影将结果压缩回原始维度。这有助于保持输出与模型其余部分的兼容性，并充当瓶颈，迫使网络仅保留最有用的特征。总之，这种“先扩展再压缩”的模式使每个专家能够应用丰富的转换，同时保持一致的输出大小。vLLM 使用 fused_moe 内核将这些投影执行为组通用矩阵乘法（Group GEMM）操作——即为分配给给定令牌的每个专家执行一次 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 适配器为投影添加了两个操作。收缩操作计算 z=xA，将输入从 h_in 维度减少到 r 维度。扩展操作采用该 r 维结果，并通过将 z 与 B 相乘，将其投影回 h_out 维度。这在图 1 的右侧进行了说明。

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

每个专家有两个权重投影：gate_up 和 down。当应用 LoRA 适配器时，它会为每个投影增加两个低秩操作，即收缩（shrink）和扩展（expand）。这意味着每个专家总共需要四个 LoRA 内核操作：gate_up 的收缩和扩展，以及 down 的收缩和扩展。在 Multi-LoRA 服务设置中，由于多个 LoRA 适配器同时服务于不同的用户或任务，系统必须高效地管理每个请求中、每个适配器、每个专家的这四个操作。这成为 MoE 模型的一个关键性能瓶颈。这四个操作涉及矩阵，其中一个维度（LoRA 秩 r）比另一个维度（例如隐藏状态和中间表示维度）小 100-300 倍。标准的 GEMM 内核专为近似正方形的矩阵设计，在细长矩阵上表现不佳，这就是本文后续描述的内核优化之所以必要的原因。除了必须针对细长矩阵进行优化外，为 MoE 模型添加 Multi-LoRA 支持还带来了两个技术挑战。首先，vLLM 缺乏在 MoE 层上执行 LoRA 的内核，因为现有的密集型 Multi-LoRA 内核无法处理专家路由。其次，MoE LoRA 结合了两种稀疏性：专家路由（分配给不同专家的令牌）和适配器选择（使用不同 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) 对这四个内核操作的计算和内存吞吐量进行了性能分析：gate_up_shrink、gate_up_expand、down_shrink 和 down_expand。这些发现促使我们针对这四个内核开发了执行优化、内核级优化和调整后的配置。

执行优化

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

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

内核优化

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）的求和来解决这个问题，这些线程组并行计算部分和，然后合并结果。这些部分结果需要一个原子加法来产生最终和。由于我们执行的是不含额外逻辑的纯原子加法，我们通过将原子加法操作的参数 sem="relaxed" 设置为来利用 Triton 编译器的优化自由度。

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

我们还从收缩和扩展 LoRA 内核中删除了不必要的遮罩和点积操作。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 调整文档。我们同时调整了四个 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 社区的合作，我们实现并开源了 MoE 模型（包括 GPT-OSS、Qwen3 MoE、DeepSeek 和 Llama 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 的优化有助于跨模型提供额外的延迟改进，例如 GPT-OSS 20B 在 vLLM 0.15.0 的基础上 OTPS 快了 19%，TTFT 好 8%。若要开始在 Amazon 上托管自定义模型，请参阅 Amazon SageMaker AI 托管和 Amazon Bedrock 文档。

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

致谢

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

同时发布于 AWS 博客。