跳到正文
vLLM 官方博客· Inferact and the vLLM Team·· 4 小时前精选AI 评分65

vLLM 实现 MiniMax-M3.1-Flash 的 Day-0 支持

vLLM Delivers Day-0 Support for MiniMax-M3.1-Flash

AI 导读

vLLM 在发布当天为 MiniMax-M3.1-Flash 提供支持,覆盖 NVIDIA Blackwell、Rubin 和 AMD MI350 系列 GPU。

推荐理由

原文给出多平台内核优化细节和 AgentX 实测数据,读者可以据此评估该模型在不同 GPU 上的推理部署收益。

正文 · AI 翻译

MiniMax-M3.1-Flash 是 MiniMax-M3 家族中的最新模型,在推理、编码、智能体和视觉任务上均有改进。它还针对推理效率进行了构建,采用原生 NVFP4 KV 缓存、NVFP4 量化的 MoE 权重,以及用于推测解码的 DSpark 草稿模型。

vLLM 在 Day 0 即支持 MiniMax-M3.1-Flash,可在 NVIDIA Blackwell GPU 上运行,在配备专用 Rubin 内核的 NVIDIA Rubin GPU 上运行,也可在 AMD MI350 系列 GPU 上运行,在智能体工作负载上表现出强劲性能。

TL;DR

  • 带 NVFP4 KV 缓存的 MiniMax 稀疏注意力(MSA),同时用于索引器和注意力路径,外加 NVFP4 MoE 权重和 DSpark 草稿模型。
  • Blackwell MSA 内核使用新的 CUDA 13.4 QMUL4 指令在内核内部将 NVFP4 反量化为 FP8:索引器 QK 解码最高提速 1.7 倍,稀疏注意力解码提速 1.3 倍。
  • 带 MXFP8 共享专家融合和提前 TMEM 释放的 NVFP4 MegaMoE:预填充阶段内核延迟降低 20%,预填充吞吐量提升 7%。
  • 索引器解码上下文并行(DCP):在 4× GB300 上、512K 上下文时,每个稀疏层的索引器最高提速 3.3 倍。
  • QKV GEMM + norm + RoPE + KV-insert 融合:在 512+ tokens 时预填充提速 1.1–1.6 倍,TP4 下 P50 TTFT 约降低 19%。
  • FlashInfer PrimiTS 后端用于 DSpark 草稿模型注意力:比 FlashAttention 4 快 1.3–2 倍,总前向时间降低 5%。
  • Rubin(SM107)MSA 内核也已集成到 vLLM 中,性能达到业界领先水平。
  • AMD MI350 系列:NVFP4 原生 MSA 内核、索引器解码上下文并行、双批次重叠以及 mono 内核,带来全面的 2–3 倍性能提升。
  • 智能体服务:在 NVIDIA GB300 上,延迟敏感端最高每用户 418.9 tok/s;在 NVIDIA Rubin 上,最高每用户 724.9 tok/s 以及每 GPU 总计 191.5K tok/s;在 AMD MI355X 上,在 SemiAnalysis AgentX 基准测试中,按 TCO 计算最高每美元 1.333 亿总 tokens。
  • Day-0 RL 支持:在 NVFP4 检查点上使用 veRL、FSDP2 和 LoRA 进行端到端 RL。

快速开始

两家厂商均提供预构建镜像:

  • NVIDIA:vllm/vllm-openai:minimax-m31-flash
  • AMD:vllm/vllm-openai-rocm:minimax-m31-flash

请注意,该模型只有在代码经审查并合并到上游后,才会出现在 nightly PyPI wheel 中,这通常需要约一周时间。完整配方见 vLLM Recipes,检查点见 MiniMaxAI/MiniMax-M3.1-Flash。

NVIDIA 上的 TP4

我们推荐使用 CUDA 13.4 工具包以获得最佳性能。

vllm serve MiniMaxAI/MiniMax-M3.1-Flash \
  -tp 4 \
  --reasoning-parser minimax_m3 \
  --block-size 128 \
  --kv-cache-dtype nvfp4 \
  --attention-config '{"indexer_kv_dtype": "nvfp4"}' \
  --speculative-config '{"method":"dspark","num_speculative_tokens":7,"attention_backend":"FLASHINFER_TS","kv_cache_dtype":"fp8","draft_sample_method":"probabilistic"}' \
  --disable-hybrid-kv-cache-manager

AMD 上的 TP4

vllm serve MiniMaxAI/MiniMax-M3.1-Flash \
  -tp 4 \
  --reasoning-parser minimax_m3 \
  --block-size 128 \
  --kv-cache-dtype nvfp4 \
  --attention-config '{"indexer_kv_dtype": "nvfp4", "indexer_decode_cp": true}' \
  --quantization-config.moe mxfp8 \
  --speculative-config '{"method":"dspark","num_speculative_tokens":3,"attention_backend":"TRITON_ATTN","kv_cache_dtype":"fp8","draft_sample_method":"probabilistic"}' \
  --disable-hybrid-kv-cache-manager

MSA 入门

MiniMax-M3 模型家族用块稀疏注意力的一种变体(arXiv 2606.13392)替代全注意力,以提高推理效率。KV 缓存被拆分为 128 个 token 的块。一个轻量级索引器为每个 GQA 组使用一个索引查询,针对单个共享的索引键头,对上下文中的每个键进行评分,然后将这些分数最大池化为每个块一个分数。主注意力随后精确运行,但仅作用于排名前 16 的块,并且始终保留最近的块。这样无论上下文多长,每个查询最多只处理 2,048 个键。

收益最大的是解码阶段,因为它受内存带宽限制。每一步仍会扫描整个上下文的索引键,但每层每个 token 仅 128 字节,是稠密注意力所读取的 1,024 字节 K/V 的八分之一(FP8,4 个 KV 头 × 128 维)。随后仅为所选的 2,048 个 token 读取完整的 K/V。在长上下文解码时,每步的 HBM 流量降至稠密注意力所读取的约 12.5%。

图 1:MiniMax 稀疏注意力。每个 GQA 组的索引器挑选其 top KV 块,该组内的每个查询头都恰好只关注这些块。逐步查看。

MiniMax-M3.1 架构中的新特性

MiniMax-M3.1 相比上一代大幅提升推理效率:它针对 NVFP4 KV 缓存进行了原生的量化感知训练,其路由 MoE 权重也以 NVFP4 存储。NVFP4 专家将检查点从 451.5 GB(MiniMax-M3-MXFP8)缩小到 250.1 GB,减小了 44.6%。NVFP4 KV 缓存每 token 需要 34,560 字节,而不是 FP8 下的 61,440 字节,每 token 节省 43.75%。综合来看,在四块 NVIDIA B300 GPU 上采用 DEP4 部署和 256K token 序列时,M3.1 可容纳 76 个并发序列,而使用 FP8 KV 缓存的 M3 仅能容纳 28 个,在相同硬件上约为 2.7 倍,这直接提升了服务吞吐量。

图 2:MiniMax-M3.1 经过量化感知训练,为索引器和稀疏注意力均使用 NVFP4 KV 缓存。更改 KV 缓存 dtype、并行方式、GPU 或上下文长度,以查看可容纳的序列数。

性能优化

vLLM 在多个硬件平台上为 MiniMax-M3.1 提供了强劲性能。在 NVIDIA GPU 上,我们提供了针对 NVFP4 KV 缓存优化的内核。我们的 MSA 代码库(vllm-project/MSA)包含针对不同工作负载和操作调优的内核。这些包括 MSA 索引器 prefill 和 decode、稀疏注意力 prefill 和 decode,以及 MSA 专用的 top-k 内核。

图 3:Blackwell 上解码时(M ≤ 32)的一个 MiniMax-M3.1 稀疏解码器层,逐个内核展示。带编号的徽章标记了本文中的优化项;点击其中一个可高亮其内核。

NVIDIA Blackwell 优化的 MSA 内核

KV 缓存以 NVFP4 存储,而查询保持 FP8,因此我们无法使用原生 NVFP4 张量核心指令进行矩阵乘法。相反,该内核将 NVFP4 KV 缓存反量化为 FP8,并运行 FP8 MMA。为保持反量化成本较低,我们使用 CUDA 13.4 中引入的新 Blackwell 专用 PTX 指令 mul.e4m3x4.e2m1x4.e4m3x4(SASS QMUL4)。它在单条指令中按比例因子缩放四个 NVFP4 值并生成四个 FP8 值。这比之前的 FP16 往返路径所需的 CUDA 核心指令少得多。结果是在索引器 QK 解码内核上最高获得 1.7 倍加速,在稀疏注意力解码内核上获得 1.3 倍加速。

NVFP4 MegaMoE

MegaMoE 是现代专家并行(EP)部署的关键构建块。我们将 DeepGEMM 的 MegaMoE 扩展到 NVFP4 精度,并融合了 MXFP8 共享专家,已在我们 fork 中开源(vllm-project/DeepGEMM)。为 MegaMoE 添加 NVFP4 支持不仅仅是更换 dtype。与之前的 FP8×FP4 路径相比,NVFP4 将每元素字节数减半,并将张量核心吞吐量翻倍。内核的其他部分也必须随之改变,否则流水线就会失去平衡。

其中一项改动是在 FC1 尾声提前释放 TMEM 阶段。此前,TMEM 阶段只有在尾声完成后才会释放。这种等待是不必要的,因为尾声完全在寄存器中运行。一旦累加器加载到寄存器中,在任何算术运算发生之前,TMEM 就可以被释放。这在 prefill 中把总内核延迟降低了 20%,而更快的 NVFP4 计算在 prefill 中最重要。它还将整体 prefill 吞吐量提高了 7%。

NVFP4 indexer DCP(解码上下文并行)

在低延迟服务中,每个 GPU 对完整上下文计算自己的索引头的分数。在四个 GPU 上,这意味着每个 indexer 层都要读取完整索引缓存四次。Indexer DCP 改为按上下文拆分工作。每个 GPU 对四分之一页计算全部四个头的分数。然后它通过 NVLink 将每个头的本地 top-16 发送给拥有该头的 GPU,由后者将候选合并为最终的 top-16。选择结果保持精确,在 4× NVIDIA GB300 GPU 上、512K 上下文时,indexer 每个稀疏层的运行速度最高提升 3.3×。

QKV GEMM + norm + RoPE + KV 插入融合

MiniMax-M3.1 的 QKV 和 indexer 投影过去会把完整的 QKV 张量写入 HBM。随后第二个内核再把它读回来,用于 norm、RoPE 和 KV 缓存插入。我们将整条链路融合进单个 GEMM 内核,其尾声逐个头处理 FP32 累加器。它对 query 和 index query 应用 norm 和 RoPE。它还将 K、V 和 index key 量化为 NVFP4,并直接写入各自的缓存。QKV 张量完全不接触 HBM。在 GB300 上,512 个或更多 token 的 prefill 批次运行速度提升 1.1–1.6×,在 TP4 下 P50 首 token 时间(TTFT)下降约 19%。

用于 DSpark drafter 的 FlashInfer PrimiTS

MiniMax-M3.1 使用 DSpark 作为其投机解码 drafter,它运行带 4096 token 滑动窗口的非因果 GQA。此前,vLLM 在 FlashAttention 4 上运行该 drafter,而后者并未针对解码形状进行优化。我们将其切换到 FlashInfer 的 PrimiTS 后端。PrimiTS 比 FlashAttention 4 快 1.3–2×,并将整体前向时间缩短 5%。

NVIDIA GB300 上的智能体工作负载基准测试

vLLM 在智能体服务负载上提供了强劲性能。我们用 SemiAnalysis AgentX 基准测试(详见我们之前的文章)在 NVIDIA GB300 上对我们的 day-0 镜像进行了基准测试。在延迟敏感的前沿一端,MiniMax-M3.1-Flash 每用户保持 418.9 tokens 每秒。沿着前沿进一步,它在每 GPU 服务 38.6K 总 tokens 每秒的同时,让每个用户保持在 200 tokens 每秒以上。在吞吐量一端,它达到每 GPU 153.8K tokens 每秒,约为最具交互性点的 24×,同时每个用户仍能看到约 127 tokens 每秒。

图 4:NVIDIA GB300 上 MiniMax-M3.1-Flash 的 AgentX 帕累托前沿,使用 7 个投机 token 和 3.16–3.26 个 token 的合成接受长度(包括第一个 token)。悬停某个点可查看其指标。

NVIDIA Rubin 优化的 MSA

我们还在 NVIDIA Rubin GPU 上优化了此模型。为缓解 NVFP4 到 FP8 的转换开销,我们将 NVFP4 直接转换为 FP16,并以 FP16 为 indexer 运行 score MMA。这减少了约三分之一的去量化指令,而且分数变得精确,而不再为了运行 FP8 MMA 而携带 E4M3 的重新舍入。

在查询数 ≤4 时,单头索引器解码内核更进一步:借助 SM107 的打包 fma.rn.f32x2.f16x2,它对每个 MMA 结果应用一次块缩放,而不是对每个 K 值应用一次。这还通过每个 SM 运行三个 CTA 实现了更多的延迟隐藏。Rubin 更大的共享内存(328 KB,而 Blackwell 为 228 KB)以及额外的 Tensor Memory 列数(576,而 Blackwell 为 512)让这成为可能。

最后,由于 Rubin 的 SM 数量多于 Blackwell,工作规划器根据 SM 数量来确定瓦片大小,因此工作均匀分布在 Rubin 的全部 212 个 SM 上。

预填充稀疏注意力也得益于局部性域优化。我们通过 CUDA 虚拟内存 API 分配部分缓冲区,使 2 MB 页面在两个域之间交替,合并后的内核将每个瓦片交给一个 CTA,其 SM 位于持有该瓦片数据的域中。在 Rubin 上,预填充与解码注意力内核使用密集 FP8 MMA,每条指令 K=64,这将 MMA 指令数减半,并以 Blackwell 两倍的 FP8 速率运行。

NVIDIA Rubin 上的智能体工作负载基准测试

在 NVIDIA Rubin 上,MiniMax-M3.1-Flash 在同一基准测试中将前沿进一步推进。在延迟敏感端,它保持每用户每秒 724.9 个 token。沿前沿继续,它在每 GPU 每秒总计服务 157.3K 个 token 的同时,让每个用户保持在每秒 200 个 token 以上。在吞吐量端,它达到每 GPU 每秒 191.5K 个 token,同时每个用户仍能看到约每秒 134 个 token,在相同每用户速率下约为 GB300 的 1.5 倍。

图 5:MiniMax-M3.1-Flash 在 NVIDIA Rubin 上的 AgentX 帕累托前沿。悬停在某个点上可查看其配置和指标。

AMD GPU 支持

在 AMD GPU 上,我们带来了 NVFP4 KV 缓存支持,因此 AMD 部署也可以使用低精度格式来削减 KV 缓存内存。对于解码,我们添加了索引器 DCP 以加速长上下文解码,以及一个加速低延迟推理的 mono 内核。对于预填充,序列并行和双批次重叠削减了引擎内部的通信,新的 KV 缓存布局让每个 KV 块通过一次传输即可移动到解码节点。MSA 和 mono 内核构建在 FlyDSL 中,这是 AMD 用于 GPU 内核的 Python DSL,并得到了 FlyDSL 团队的帮助。

NVFP4 MSA 内核

MiniMax-M3.1 针对主注意力和索引器均使用 NVFP4 KV 缓存(QAT)进行训练,但 CDNA4 没有原生指令来反量化 NVFP4。为了保持 MSA 内核的高速,我们将反量化转变为查表。E4M3 缩放因子只有 256 种可能,E2M1 幅值有 8 种,因此我们将每个缩放值预先计算好(已舍入到 FP8),存入 LDS 中一个 2 KB 的表。一次 v_perm_b32 字节置换加上一次对符号位的 XOR,即可将四个 FP4 值转换为 FP8。该内核还对接共享前缀页面的最多 64 个请求只反量化一次。综合起来,这些使稀疏解码注意力比将选定块转换为 FP8 并运行 FP8 注意力内核快最多 2.5 倍。

NVFP4 索引器 DCP

我们在 AMD 上使用与 NVIDIA 相同的索引器解码上下文并行方案,而在 CDNA4 上它还带来了额外的收益。没有 DCP 时,TP4 下的每块 GPU 都要在完整上下文上为自己的索引头打分,因此索引缓存要被读取四次。这也会让矩阵核心利用率不足。一个 DSpark-3 解码步骤每个请求只有 4 个查询行,仅填满 16×16×128 FP8 MFMA 的 16 列中的 4 列。使用 DCP 后,每块 GPU 在四分之一上下文上为全部四个头打分,于是 4 个头 × 4 行填满全部 16 列。单个 kernel 交换索引查询、对其分块打分,并通过 Infinity Fabric 将每个头的分数直接写入所属 GPU 的内存。设备端握手让这在 HIP graphs 内也能正常工作。随后所有者像以前一样挑选前 16 个,而由于分数与单 GPU 版本逐位相同,选择是精确的。在 4× MI355X 上,针对 45 个 30K–210K token 的请求,索引器在每个稀疏层上最高快 2.8 倍。

SP + DBO 预填充与一次读取 KV 传输

在 TP4 下,MI355X 上的长提示预填充受通信限制,既包括节点内 GPU 之间,也包括预填充节点与解码节点之间。在引擎内部,双批次重叠(DBO)将张量并行集合通信隐藏到另一个微批次的计算之后。序列并行(SP)则直接削减流量。每块 GPU 保留四分之一的 token,注意力输出以 MXFP8 进行全对全交换,使每块 GPU 在自己的 token 上运行完整的 o_proj,而 MoE 输出则进行 reduce-scatter。这将每块 GPU 每 token 每层的流量从约 36 KB 降至 20 KB,并且 norm、量化和路由器仅需在四分之一的 token 上运行。

更快的预填充引擎只有在 KV 缓存也能快速到达解码节点时才有用。vLLM 的逐层 KV 缓存布局将每个 KV 块拆分到 125 个小页中,每个页都需要自己的 NIXL RDMA 读取。我们改用块外层布局,将每个块存储在一个连续的 1.82 MB 行中,因此 NIXL 只需一次读取即可移动它。综合起来,引擎和 KV 传输的改动使 4× MI355X 上 16K+ token 的预填充批次快 1.8–2.1 倍。

低延迟单 kernel

在低并发下,AMD GPU 上的解码受限于启动开销:kernel 启动开销比数学计算耗时更长。单解码通过将每个稀疏 MoE 解码层作为单个 FlyDSL kernel 启动来解决这一问题。那一次启动涵盖了从 QKV 投影和稀疏注意力到张量并行全归约以及 MXFP4 MoE 的所有内容,替代了约 1,170 次 kernel 启动。在 TP4 下的 MI350X 上,这使解码步骤最高快 2.3 倍,并使低并发下 8K 输入/1K 输出吞吐量翻倍。在智能体编程轨迹上,它将交互性提升了 50–65%。

AMD MI355X 上的智能体工作负载基准测试

我们的 AMD 实现还提供了强劲的每美元性能。在 InferenceX AgentX 基准测试中,MiniMax-M3.1-Flash 在前沿的延迟敏感端保持每用户每秒 331.0 个 token。在前沿更远处,它让每个用户保持在每秒 170 个 token 以上,同时实现每美元 TCO 总计 3460 万个 token。在成本高效端,它达到每美元 1.333 亿个 token,接近最具交互性点的 10 倍,而每个用户仍能获得约每秒 79 个 token。

图6:MiniMax-M3.1-Flash 在 AMD MI355X 上的 AgentX 帕累托前沿,以每美元总拥有成本(TCO)对应的总 token 数计。c1 和 c4 使用 3.2 token 的合成接受长度与 7 个投机 token;其他点使用 2.45 token 与 3 个投机 token。悬停在某个点上可查看其指标。

Day-0 RL 支持

我们从 MiniMax-M3.1 的 NVFP4 检查点开始对其进行 RL,并且只训练秩为 32 的 LoRA 适配器。veRL 训练器以引擎服务它时相同的方式量化基础模型:路由专家上使用 NVFP4,其他投影上使用 MXFP8,并在每一步之后只将适配器同步到引擎。在引擎侧,专家 LoRA 运行在 FlashInfer 的 TRTLLM NVFP4 MoE 内核中,因此基础专家仍然运行在正常的服务内核上。

在训练之前,我们从三个方面检查了训练器和引擎是否一致。引擎接收到的权重与训练器发送的完全相同,两者将每个 token 路由到相同的专家(路由重放),并且两者为相同 token 分配了接近的 logprob。

剩余的差距用 token 级重要性采样(TIS)来纠正,重要性比率上限为 2。因为每个 batch 只进行一次优化器步进,训练器使用 rollout logprob 作为旧 logprob,并跳过会重新计算它们的额外前向传播。

在 DAPO-Math-17K 上,使用 16 块 GB300 GPU、200 步 GRPO 将 AIME-2024 准确率从 36.3% 提升到 67.5%。在同一轮运行中,训练奖励从 −0.06 上升到 0.67,train/rollout KL 保持较低,从 0.040 下降到 0.024。

图7:在 DAPO-Math-17K 上使用 veRL 对 MiniMax-M3.1-Flash 进行 Day-0 RL(16× GB300,200 步 GRPO):AIME-2024 准确率、训练奖励和 train/rollout KL。使用选项卡在它们之间切换。

致谢

我们衷心感谢 MiniMax 团队开源 MiniMax-M3.1-Flash。我们也感谢 Inferact、NVIDIA 和 AMD(Hongxia Yang 负责 mono 内核,Felix Marty、Andy Luo、Felix Li、Hattie Wu、Gilbert Lei 和 Peng Sun)在 vLLM 中为该模型提供 day-0 支持,以及 FlyDSL 团队在 AMD 内核方面的帮助和 vLLM 社区。

来源:vLLM 官方博客 · vllm.ai