跳到正文
北京时间
原文
vLLM 官方博客· Mistral and the vLLM Team·· 17 小时前精选AI 评分71

vLLM 发布 Mistral Large 4 Day 0 支持

Day 0 support for Mistral Large 4 on vLLM

AI 导读

vLLM 团队宣布对 Mistral 新旗舰模型 Mistral Large 4 提供 Day 0 支持,该模型为 1T 参数、每 token 激活 49B 的原生多模态 MoE 模型,采用 MLA 稀疏注意力加 DSA top-k 索引器、跨层 KV 共享、NoPE、SSMax 和 LeRoPE 等新架构,并内置 DSpark 推测解码草稿器与 Pixtral 视觉编码器。

推荐理由

vLLM 团队详解支持 Mistral Large 4 所需的架构适配与量化、长上下文、推测解码等实测数据,可作部署参考。

正文 · AI 翻译

我们很高兴地宣布,vLLM 已实现 Mistral Large 4 的 day-0 支持。Mistral Large 4 是 Mistral 推出的新一代前沿混合专家(MoE)原生多模态模型。它把 DeepSeek 风格的稀疏注意力栈(MLA 加 DSA top-k 索引器)与多项新架构特性结合在一起:跨层 KV 共享、NoPE 层、SSMax 和 LeRoPE。模型自带 DSpark 投机解码草稿器和 Pixtral 视觉编码器。

支持它远不止加一个模型文件那么简单。ML4 运行在 vLLM 的模型运行器 V2(MRV2)上,并支持可断开的 CUDA 图。它的各个构建模块放在新的 vllm/models/mistral/ 包中,这项工作还顺带改进了 decode 上下文并行(DCP)、prefill 上下文并行(PCP)、DSpark、基于 NIXL 的分离式服务以及 HiSparse 主机内存 KV 卸载。

亮点

  • 架构:72 层 MoE,128 个路由专家(top-4)加 1 个共享专家,隐藏维度 6144。注意力采用 MLA(KV LoRA 秩 512,RoPE 维度 64),并配有 DSA 稀疏索引器,每个查询选取 top 2048 个 key。总参数 1 万亿,每 token 激活 490 亿(含嵌入层和输出层为 520 亿)。
  • 检查点:FP8(块缩放权重)和 NVFP4(专家用 NVFP4,注意力/共享专家用 FP8)。两者默认都使用 FP8 潜在 KV 缓存(fp8_ds_mla)。
  • 多模态:Pixtral 视觉编码器处理图像输入,支持每个提示词包含多张图片。
  • 投机解码:内置 DSpark 草稿器(3 层滑动窗口 MLA,深度 3),每步接受约 2.9 个 token,在 TP8 上输出吞吐约提升 2 倍。
  • 长上下文:DCP 在 TP8 上带来约 7.9 倍的 KV 容量。PCP+DCP 预填充把首 token 时间(TTFT)最多缩短 1.6 倍。DSpark 在两种模式下都能工作。
  • 分离式服务:NIXL 预填充/解码(P/D)分离,支持 DP/EP、TP 和 PCP×DCP 预填充节点,两侧都启用 DSpark。
  • HiSparse:主机内存稀疏 KV 卸载,现已支持跨层 KV 共享。
  • AMD:在 AMD Instinct MI350 系列 GPU 上通过 ROCm 实现 day-0 支持,包括 DSpark、图像输入和 MXFP4 专家的专家并行。

用法

Mistral Large 4 的支持包含在 vllm/vllm-openai:mistral-large-4 Docker 镜像中,该镜像还包含 mistral-common 1.13。下面的命令在镜像内运行 vllm serve:

docker pull vllm/vllm-openai:mistral-large-4

8 块 GPU 上的 FP8:

vllm serve mistralai/Mistral-Large-4-1T-A52B \
  --tensor-parallel-size 8 \
  --speculative-config '{"method": "dspark", "num_speculative_tokens": 3}' \
  --tool-call-parser mistral --enable-auto-tool-choice \
  --reasoning-parser mistral

4 块 GPU 上的 NVFP4(Blackwell;在 8 块 GPU 上使用 --tensor-parallel-size 8):

vllm serve mistralai/Mistral-Large-4-1T-A52B-NVFP4 \
  --tensor-parallel-size 4 \
  --speculative-config '{"method": "dspark", "num_speculative_tokens": 3}' \
  --tool-call-parser mistral --enable-auto-tool-choice \
  --reasoning-parser mistral

NVFP4 不以质量换速度。在多模态、长上下文、智能体编程、业务自动化和网络安全评测中,它在每项任务上与 FP8 的差距都在 2.1 分以内,平均差异为 −0.3 分。token 消耗也相当:平均补全 token 数与 FP8 相差约 4% 以内。它还更快、更省:在 TP8 下,不使用投机解码时 NVFP4 每 token 解码速度比 FP8 快 1.08–1.19 倍(启用 DSpark 时为 1.18–1.22 倍),KV 缓存容量约为 1.5 倍。

*FP8 与 NVFP4 在核心评测上的得分和 token 消耗对比。误差条为通过自助法估计的 95% 置信区间。在所有领域,包括多模态推理(GDP.pdf、MathVision)、长上下文(Graphwalks 512k)、编程智能体(DeepSWE、SWE-bench Pro、Terminal-Bench 4)、业务自动化(AutomationBench)和网络安全(CyberGym),NVFP4 检查点的表现都与 FP8 一致。

长上下文(DCP)。添加 --decode-context-parallel-size 8 可将潜在 KV 缓存分片到 8 个 rank 上。这要求使用默认的 FP8 KV 缓存。

多模态。传入 --limit-mm-per-prompt '{"image": N}' 可允许每个提示词最多 N 张图片。

多节点、分离式 P/D 以及其他已测试的配置,参见 Mistral Large 4 配方。

实现细节

模型架构

组件Mistral Large 4
层数 / 隐藏维度 / 注意力头数72 / 6144 / 64
注意力MLA,无 query LoRA,KV LoRA 秩 512,RoPE 维度 64
稀疏注意力DSA indexer,top-k 2048,32 个 index 头;44 层各自拥有一个 indexer,28 层复用最近的前一个 indexer 层的 top-k
位置编码YaRN + LeRoPE(按频率学习缩放);每第 4 层为 NoPE
query 缩放每个注意力层和 indexer 都使用 SSMax(随长度变化的 query 缩放)
KV 共享9 层(39、43、…、71)读取第 35 层的 latent cache(每 4 层 1 层)
MoE第 0 层为稠密层;第 1–71 层:128 个专家,top-4,1 个共享专家
输出logit soft cap 30
视觉Pixtral 编码器(30 层,隐藏维度 2048,patch 14)

*72 个解码器层:每第 4 层为 NoPE,28 层复用最近的前一个 indexer 层的 top-k,NoPE 层 39–71 读取第 35 层的 latent KV,且只有第 0 层是稠密层。*

在 vLLM 中支持各部分所需的工作:

稀疏 MLA + DSA indexer。 ML4 复用了 vLLM 的稀疏 MLA 栈(带 fp8_ds_mla latent cache 的 FlashMLA 稀疏实现,外加带 FP8 indexer cache 的稀疏 indexer)。ML4 特有的差异叠加在其上:共享 indexer 层复用最近的前一个 indexer 层的 top-k;indexer 拥有自己的 SSMax;NoPE 层上使用 NoPE indexer。每个 token 的 KV cache 存储 47,136 字节:41,328 字节用于 MLA latent,5,808 字节用于 indexer。

跨层 KV 共享。 九个 NoPE 层读取第 35 层写入的 latent KV,而不是写入自己的 KV。我们在 kv_sharing_target_layer_name 到 MLAAttention 以及 AttentionLayerBase 中加入了支持。消费层跳过 cache 写入,V2 runner 将它们解析为目标层的 KV cache 规格。消费层不占用 KV 内存:72 层中有 63 层拥有自己的 cache。 该特性与模型无关,正在 vllm#64095 中向上游合入。

NoPE、SSMax 与 LeRoPE。 NoPE 层以 RoPE 维度为零的 MLA 形式表达,因此使用相同的 kernel 和 cache 布局。LeRoPE 被折叠进 YaRN 频率。SSMax 被融合进 MLA RoPE kernel(见下文融合部分)。

Mistral 模型包。 移植没有在 deepseek_v2.py 中增加开关,而是基于 vllm/models/mistral/ 下可复用的 Mistral 组件构建:注意力、indexer、SSMax、MLP、MoE、融合算子以及逐层规格(common/layer_spec.py)。ML4 模型本身是 vllm/models/mistral_large4/nvidia/model.py,DSpark drafter 紧挨着它 (dspark.py)。逐层特性(RoPE/NoPE、indexer 拥有者/共享者、KV 拥有者/消费者)来自基于 params.json 构建的层规格,而不是共享代码中的模型特定分支。

模型 runner V2。 ML4 仅在 MRV2 上运行。V2 注意力工具会发现 KV 共享层并查找其目标规格。DSpark 接入 V2 的 DFlash/DSpark speculator。

可断开的 CUDA graph 与融合。 ML4 运行时无需 torch.compile。它使用可断开的 CUDA graph,该特性会针对该架构自动启用(当文本模型位于多模态封装内时同样如此)。相比 eager 模式,它们带来了约 6 倍加速。 手写融合替代了编译器原本会做的工作,同时支持进一步优化:

  • 第一轮:SSMax、MLA RoPE + SSMax、indexer Q 路径和 indexer K 路径。这些将 decode TPOT 降低了 8–10%。
  • 解码环节:FP32 router GEMM;indexer query 的 LayerNorm 用一次 Triton 启动完成;all-reduce、残差和 RMSNorm 合并进一个 FlashInfer kernel;indexer K 直接写入 FP8 cache。在 TP8 FP8 下,c1 TPOT 从 10.22 ms 降到 9.00 ms(-12%),c8 从 16.01 ms 降到 15.08 ms。

长上下文:解码与预填充的上下文并行

DCP。 MLA 只保留一个 latent KV head,因此在 TP 下 cache 会在每个 rank 上复制一份。DCP 改为按 token 分片。在 TP8 上,DCP8 把 KV 容量从 2.38M 提升到 18.88M token(FP8),从 3.60M 提升到 28.76M token(NVFP4)。

Long-context capacity on 8x GB300: TP8 vs TP8 DCP8
8× GB300 上的长上下文容量:TP8 与 TP8 DCP8 对比

并发 128k / 256k / 1M token 请求,KV cache 在 8× GB300 上以 ML4 NVFP4 运行,TP8 与 TP8 DCP8 对比:floor(KV pool / 输入长度) - 3.60M 与 28.76M token,gpu-memory-utilization 0.85。

开箱即用时,搭配稀疏 MLA 的 DCP 在每个解码步都要付一笔不小的固定开销。在 60k 上下文下做 profile 显示,每步约有 7 ms 花在 DCP 专属工作上:

  • 对 LSE 和 indexer top-k 候选做 NCCL all-gather(共 44 个 indexer 层);
  • reduce-scatter;
  • 稳定的 top-k 合并;
  • query gather;
  • 小型 mask kernel 和启动间隙。

我们把这些替换为:

  • a2a combine,用于 FlashMLA sparse(已合入 vllm#61285):用对称内存的 all-to-all LSE combine 取代 all-gather + reduce-scatter。ML4 默认启用。
  • 直接 top-k 交换:每个 rank 把打包好的 (score, id) 候选推送到对端的对称内存 inbox,由 selector 就地读取。这取代了每个 indexer 层的一次 NCCL all-gather 和一次 6-pass radix select。
  • Top-k 压缩:对 32 个 token 以上的 batch,每个 rank 压缩自己有效的 top-k 槽位,并用 topk_length 限定 kernel 范围,而不是扫描 2048 个塞满 -1 的槽位。在 5k token 的预填充 profile 中,attention 从 165 ms 降到 38 ms。
  • 一个融合 kernel 负责处理空行(即某个 rank 不拥有某条 query 的任何 top-k 槽位)。

DCP 专属开销降到每步 3.2 ms。再加上解码融合,在 TP8 DCP8 FP8 下(60k 输入 / 1k 输出):

之前之后
c1 TTFT p506.89 s4.73 s
c1 TPOT p5017.82 ms12.50 ms
c64 输出 tok/s139.1204.0 (+47%)
DCP8 before and after the collectives port
DCP8 在集合通信移植前后的对比

TP8 DCP8 下 TTFT 和 TPOT p50 随并发变化(60k 输入 / 1k 输出,无投机解码),对比 DCP 集合通信移植与解码融合之前(a2db3f0b9d)和之后(bbe5926511)。

这些 DCP 优化随 Mistral Large 4 支持在 day 0 一同上线。

何时使用 DCP。DCP 用每步延迟换 KV 容量。在同样的 8x GB300 NVFP4 部署和工作负载下(无投机解码),DCP8 每 token 的解码速度比 TP8 慢约 16–28%:8k 输入时 c1 TPOT 为 11.1 vs 8.7 ms,60k 时为 11.4 vs 9.1 ms。在最多 128 个 8k–60k token 的并发请求下,TP8 的 3.6M token 池从来不是瓶颈,因此 TP8 在那里吞吐更高。一旦 KV 成为限制,DCP8 才划算:超长上下文,或大量并发的长生命周期会话。在我们的 agentic P/D 运行中(约 100k token 会话),一个 TP4 DCP4 解码节点在 32 个会话下把 TTFT 保持在 0.7 s,而普通 TP4 解码节点 KV 耗尽,TTFT 攀升到 15 s。

PCP + DCP 预填充。对长 prompt,PCP 把一个 prompt 的 token 拆分到各 rank。PCP+DCP 现在在 ML4 上支持 CUDA graphs,包括 DCP > PCP 和 DSpark。两项修复让这成为可能:

  • 在 PCP 下,MLA KV-cache 写入会先跨 PCP rank 对输入做 all-gather,而这次 gather 取决于该批次的 prefill/decode 划分。在分段图(piecewise graph)中捕获时,这个划分在捕获时就被固定了。现在只有 gather 作为 eager 断点运行,写入固定大小的缓冲区,而 cache 写入本身仍留在图里。
  • DCP 组的 rank 顺序与集合通信返回数据的顺序保持一致。

在 TP1×PCP4×DCP4 上启用图后,GSM8K-200 从 0.315 提升到 0.975。配合 DSpark 可达 0.980。作为 P/D 的 prefill 节点,PCP4×DCP4 给出我们测到的最低 TTFT(见 P/D 一节)。

DSpark 投机解码

ML4 在 checkpoint 中内置了一个 DSpark drafter(mtp_dspark.* 个张量,以及 params.json 中的一个 mtp 块)。DSpark 是一种并行 drafter:一个小型非因果 transformer 在一次前向中提出一个块内所有 draft 位置。它以从目标模型第 1、18、36、54、71 层引出的隐藏状态为条件。ML4 的 drafter 还增加了:

  • RNN Markov 头:一个门控循环状态(rank 512)在各 draft 位置间运行,因此每个位置的 logits 都依赖于它之前 draft 出的 token;
  • 位置专属的 mask 输入:每个 draft 槽位都有自己的学习输入嵌入,而不是共用一个 mask token;
  • 滑动窗口 MLA draft 层:3 个 MLA 层,窗口 2048 token,BF16 KV cache,位于各自的 KV cache 组。Triton MLA 后端为它们新增了滑动窗口解码。

*一个 DSpark 步骤:目标模型的引出状态喂给 draft 的上下文 K/V,[anchor, mask_1, mask_2] 块一次通过 3 个滑动窗口 MLA 层,RNN 头依次输出 3 个 draft token,再由目标模型验证。*

用 --speculative-config '{"method": "dspark", "num_speculative_tokens": 3}' 启用。

结果。TP8,200 条真实提示词,贪心解码(temperature 0),最多输出 512 token:

FP8 基线 TPOTFP8 DSpark TPOT(AL)NVFP4 基线 TPOTNVFP4 DSpark TPOT(AL)
c18.84 ms3.50 ms (2.90)8.21 ms2.95 ms (2.93)
c815.28 ms7.08 ms (2.85)12.85 ms5.80 ms (2.86)
c3229.92 ms11.57 ms (2.85)25.38 ms9.80 ms (2.86)

该表使用贪心解码。采样时接受率也保持接近:在 temperature 1.0 / top-p 0.95 下,并发 8 时 AL 约为 2.8。温度更高时,在投机配置中设置 "draft_sample_method": "probabilistic":T=1.5 时它把 AL 从 2.04 提升到 2.36–2.53(吞吐量 +12–22%)。得益于标准的拒绝采样,两种模式都是无损的。

在相同并发下,DSpark 在 FP8 上把输出吞吐量提升 1.9–2.1 倍,在 NVFP4 上提升 1.9–2.2 倍。它不改变模型生成的内容:每个 draft token 都会对照目标模型验证,我们也做了实测确认。当基线模型重新给 DSpark 的输出打分时,这些输出与基线模型自身贪心选择的一致程度,和基线模型自己的输出一样。在 GSM8K(1,319 题)上,FP8 不用 DSpark 得 0.909,用 DSpark 得 0.918。

*每 GPU 输出 tok/s 对每用户 tok/s(1000 / TPOT p50),虚线为不用 DSpark,实线为用 DSpark(深度 3)。Short:上表中的 200 条提示词 TP8 运行。8k / 32k / 60k:ML4 NVFP4 共置于 8× GB300,TP8 和 TP8 DCP8,`vllm bench serve` 使用真实文本提示词和 1k 输出 token;60k 时只有 DSpark 运行。*

DSpark + DCP。在此之前,滑动窗口 MLA 草稿模型无法在 DCP 下运行。现在,草稿模型的小窗口缓存(3 层 × 约 2k token,每请求约 7 MB)会在每个 DCP rank 上复制一份,而目标缓存仍按 DCP 分片。这把已有支持(vllm#56723)扩展到了带窗口的 MLA 草稿模型。在 TP8 DCP8、60k 上下文下,c1 TPOT 从 11.82 ms 降到 5.22 ms(2.26 倍,AL 2.85),GSM8K-200 保持在 0.965。 代价是 KV 容量:在 DCP8 下,每个 64 token 的草稿块持有一个块 ID,其目标页覆盖 512 个 token。在 64k 上下文下,容量降到纯 DCP8 的约 0.69 倍。在 max_num_seqs 128 下,平均上下文低于约 99k 时这一限制不会触发。

P/D 中的 DSpark。DCP prefill 节点可以向 TP decode 节点交接,两侧都启用 DSpark。NIXL 从一个 prefill rank 拉取复制的草稿窗口,并以 DCP 分片形式拉取目标 KV。无论采用哪种 P/D 布局,接受率都相同:8k / 32k / 60k 输入下约为 2.98 / 2.86 / 2.78。

滑动窗口 MLA 草稿层以及 DCP 下的 DSpark 正在 vllm#63795 和 vllm#63796 中合入上游。

基于 NIXL 的分离式服务

我们在 4+4 块 GB300 GPU(一个 NVL72 机架)上用 NIXL 对 1P1D 做了基准测试。配置如下:

  • 两侧均启用 NVFP4 与 DSpark,全面使用 CUDA graphs,不启用前缀缓存;
  • vllm bench serve 使用固定长度提示词:每个输入长度 320 条唯一的单轮聊天提示词,每条恰好 8,192 / 32,768 / 60,000 个 token。每条提示词都是一段自然文本(WikiText-2 和 GSM8K 题目)后接一个问题,提示词之间没有共享前缀,以 token ID 形式发送;
  • 每请求 1,024 个输出 token(贪心解码,忽略 EOS);
  • 并发扫描和固定 QPS 扫描。

所有配置经代理在 GSM8K-200 上得分 0.965–0.980,NIXL 传输零失败,请求零失败。

*每 GPU 输出 tok/s 对比每用户 tok/s,1P1D,4+4 GB300。实线为并发扫描;空心/虚线标记为固定 QPS 扫描。*

在给定最低每用户速度下,每 GPU(8 块 GPU)的最佳输出 tok/s:

P → D8k,≥40 tok/s/用户8k,≥10032k,≥4032k,≥10060k,≥40 / ≥100c1 TTFT 8k / 32k / 60k
P DP4/EP → D DP4/EP718.2176.7202.6139.2108.80.49 / 1.90 / 3.65 s
P DP4/EP → D TP4459.2194.2205.2150.0107.20.46 / 1.88 / 3.66 s
P PCP4×DCP4 → D TP4350.6192.691.591.549.80.35 / 1.37 / 2.57 s
P TP4 → D TP4259.2185.362.862.832.30.52 / 2.07 / 4.02 s

AMD

Mistral Large 4 从发布首日起就通过 vLLM 的 ROCm 构建在 AMD Instinct MI350 系列 GPU 上运行。稀疏 MLA 运行在 AITER 的 MLA decode 内核上,DSA 索引器运行在 AITER 的 gfx950 索引器上。

MXFP4 专家。在 ROCm 上,路由专家以 OCP MXFP4 形式运行在 AITER 的融合 MoE 上。有两种方式可以获得它们:

  • AMD Quark MXFP4 检查点(Mistral-Large-4-1T-A52B-Quark-MXFP4)。Quark 只重新量化路由专家,权重使用静态缩放、激活使用动态缩放。注意力、共享专家、稠密前馈层、路由器、输出头、视觉编码器和 DSpark 草稿模型都保持原始精度。检查点从约 1,061 GB 缩小到 578 GB,因此可以在四块 MI350 系列 GPU(每块 288 GB)上服务。服务命令、评测结果和复现步骤见模型卡。
  • 加载时重新量化 FP8 检查点(VLLM_FP8_MOE_REQUANT=mxfp4)。专家在加载过程中逐层转换,因此峰值内存也足够低,TP4 同样可用。TP8 下每个 GPU 的专家权重从约 125 GiB 降到 69 GiB,在 1–64 并发请求下,输出吞吐量比直接服务 FP8 专家提升 22–71%。VLLM_FP8_MOE_REQUANT_ACT 设置 MXFP4 专家计算时使用的激活格式:a8w4(FP8 激活,默认)或 a4w4(MXFP4 激活)。

无论哪种方式,一个 8-GPU 节点都能跑两个 TP4 副本,每个副本留给 KV 缓存和激活的 HBM 约为 FP8 专家在同样四块 GPU 上所能留下的 6 倍。

功能。 ROCm 上支持 TP8 和 TP4、专家并行、通过 Pixtral 编码器输入图像,以及使用检查点内置草稿模型的 DSpark 投机解码。在 MI350X 上以 TP4 配合 FP8 KV 缓存运行时,DSpark 在真实提示词上每步接受 3.1 个 token,单请求下每 token 解码延迟降低 2.7 倍,8 请求下降低 1.9 倍,32 请求下降低 1.8 倍(输出吞吐量分别为 2.1 倍、1.6 倍和 1.4 倍)。ROCm 上潜在 KV 缓存默认是 FP8(E4M3),KV 容量是 bf16 的 1.86 倍;--kv-cache-dtype bfloat16 可将其关闭。

AITER。VLLM_ROCM_USE_AITER=1 在 AMD 的内核库 AITER 上运行 vLLM 的热点路径:融合 MoE(MXFP4 专家跑在两阶段 FlyDSL 内核上)、块 FP8 和 bf16 GEMM、MLA 解码内核、RMSNorm 以及张量并行自定义 all-reduce。ROCm 镜像自带针对该模型 MoE 和 GEMM 形状的 AITER 调优表,AITER 在启动时会将其与自己的调优表合并;在 TP4 下,1–8 并发请求时输出吞吐量提升 9–16%。还有两项设置能进一步加分:VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA=1 让稀疏 MLA 跑在 AITER 的 Triton 内核上,AITER_CUSTOM_AR_MAX_SIZE=134217728 让大型预填充步骤(最大 128 MiB)的 all-reduce 留在 AITER 的自定义 all-reduce 上。两者合计在 TP4 下再提升 2–5% 的输出吞吐量,并将首 token 时间降低 5–7%。下面两条命令都会启用这三项设置。

专用解码。对于延迟敏感的服务,--additional-config '{"decode_backend": "latency_specialized"}' 会把 1、2 或 4 请求解码步骤中的每个 MoE 层作为一次常驻 FlyDSL 内核启动来运行。该启动融合了注意力、索引器、路由器、MXFP4 专家以及两次张量并行 all-reduce。在 MI350X 上以 TP4 运行时,与同样 GPU 上 vLLM 的标准解码相比,单请求下每 token 解码延迟降低 1.5 倍,四请求下降低 1.1 倍。它可以与 DSpark 组合:单请求的验证块跑在融合内核上,每 token 延迟比不做投机的标准解码低 2.9 倍。更大的批次和预填充走标准路径。它运行在默认 FP8 KV 缓存上,要求 TP4 和 VLLM_FP8_MOE_REQUANT_ACT=a4w4,否则会拒绝启动并给出原因。该内核改编自 ROCm ATOM 的 GLM-5.2 融合解码。

在 MI350 系列上服务。两条命令都会启用 AITER,并在加载时将 FP8 专家重新量化为 MXFP4:

# TP8, one replica per node
VLLM_ROCM_USE_AITER=1 VLLM_FP8_MOE_REQUANT=mxfp4 VLLM_FP8_MOE_REQUANT_ACT=a8w4 \
VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA=1 AITER_CUSTOM_AR_MAX_SIZE=134217728 \
vllm serve mistralai/Mistral-Large-4-1T-A52B --tensor-parallel-size 8 \
  --speculative-config '{"method": "dspark", "num_speculative_tokens": 3}' \
  --tool-call-parser mistral --enable-auto-tool-choice --reasoning-parser mistral
 
# TP4 with the specialized decode, for few concurrent users; two replicas per node (second on GPUs 4-7)
VLLM_ROCM_USE_AITER=1 VLLM_FP8_MOE_REQUANT=mxfp4 VLLM_FP8_MOE_REQUANT_ACT=a4w4 \
VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA=1 AITER_CUSTOM_AR_MAX_SIZE=134217728 \
vllm serve mistralai/Mistral-Large-4-1T-A52B --tensor-parallel-size 4 \
  --additional-config '{"decode_backend": "latency_specialized"}' --max-num-seqs 64 \
  --speculative-config '{"method": "dspark", "num_speculative_tokens": 3}' \
  --tool-call-parser mistral --enable-auto-tool-choice --reasoning-parser mistral

ROCm 实现细节

AITER 上的稀疏 MLA。在 ROCm 上,ML4 的稀疏 MLA 运行在 ROCM_AITER_MLA_SPARSE 后端。索引器选出的 top-2048 位置会被转换成全局缓存槽位,每个查询 token 在槽位流中各自占一条记录。随后 AITER 的 MLA 解码内核以常驻模式读取这条流,工作切分元数据提前构建好,以便在 CU 之间均衡负载。AITER 的 MLA 内核要求头数为 16 的倍数:TP4 下每块 GPU 正好分到 ML4 的 64 个头中的 16 个,而 TP8 下的 8 个头会被补齐到 16。VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA=1 则改用 AITER 的 Triton 稀疏 MLA 内核,它不需要补齐头数,也不需要单独一份 FP8 的查询副本。bf16 缓存上的 NoPE 层没有 RoPE 部分可供 AITER 内核读取,因此它们运行在 vLLM 的 ragged Triton 稀疏注意力内核上。

AITER 上的 DSA 索引器。一个融合 Triton 内核把每个索引器键量化为 FP8(使用单一 scale)并写入索引器缓存,每个索引器层每 token 占 132 字节。在一个缓存块内,键按 16×16 的瓦片存储,这正是 AITER 分页读取器期望的布局。写入方和所有读取方(插入、prefill gather 和分页 logits)共用同一条布局规则,因此不会出现不一致。Prefill 打分运行在 AITER 的 fp8_mqa_logits 上,解码打分运行在 AITER 的分页 MQA-logits 内核上,两者都读取 FP8 键。top-2048 在形状受支持时由 AITER 的索引器 top-k 给出,否则由 vLLM 的 top_k_per_row 给出。共享索引器层复用最近的前一个索引器层的 top-k,与其他平台上的做法一致。

FP8 潜在缓存。ROCm 没有 fp8_ds_mla 布局,因此当检查点要求 FP8 KV 缓存时,ROCm 改用普通的 E4M3 缓存。全部 576 个潜在维度(512 个 NoPE 和 64 个 RoPE)都以 FP8 存储,使用一个逐张量 scale,每层每 token 占 576 字节,而 fp8_ds_mla 为 656 字节,bf16 为 1,152 字节。查询也被量化为 FP8,因此 AITER 的内核执行的是 FP8 × FP8。算上 63 个拥有缓存的层加上索引器,一个 token 占 42,096 字节,而 bf16 下为 78,384 字节:这就是 AMD 部分所说的 1.86× KV 容量。

加载时逐层将专家量化为 MXFP4。使用 VLLM_FP8_MOE_REQUANT=mxfp4 时,检查点加载器保持不变:它仍然加载带块 scale 的 block-FP8 专家权重。这些参数起初位于 meta 设备上,逐层加载器一次只物化一层。某一层的专家一加载完成,每个专家就会被反量化为 bf16,再重新量化为 MXFP4(E2M1 数值,每 32 个数值一个 E8M0 scale),随后释放该层的 FP8 副本。峰值内存是 MXFP4 模型加上一层的 FP8 专家,这正是 1T 模型能装进四块 GPU 的原因。权重误差与直接把 bf16 权重量化为 MXFP4 大致相当。之后专家运行在 AITER 的融合 MoE 上:a8w4 使用带 FP8 激活的两阶段 FlyDSL 内核,a4w4 则使用 MXFP4 × MXFP4。

专用解码:每层一个常驻内核。使用 decode_backend=latency_specialized 时,1、2 或 4 个请求的纯解码步骤(或单个请求的 DSpark 验证块)会把 71 个 MoE 层中的每一层作为一次 FlyDSL 内核启动来运行。该启动使用 256 个 CTA,全部同时常驻在 GPU 的 256 个 CU 上。其内部依次运行:

  • 输入归一化,以及带 SSMax 和 LeRoPE 的查询投影与潜在 KV 投影;
  • 整个 DSA 索引器:投影、键缓存写入、打分以及精确的 top-2048;
  • 在潜在缓存上做稀疏 MLA,以 bf16 或 FP8 原地读取;
  • 输出投影、softmax top-4 路由,以及共享专家和路由 MXFP4 专家。

嵌入、稠密层 0 和最终归一化仍留在 vLLM 的模块上。

两次张量并行 all-reduce 都在核内完成,走的是四块 GPU 之间的对等内存,所以一层只需一次启动,不用回主机往返。这些核按 1、2、4 行构建,正好覆盖 vLLM 的 CUDA 图捕获尺寸,并在这些图内重放。其他步骤仍走 vLLM 的层:预填充、混合批次、更宽的批次。该核改编自 ROCm ATOM 的 GLM-5.2 融合解码,并针对 ML4 的层类型做了扩展:

  • 查询不使用 LoRA;
  • 查询和索引查询上使用 SSMax;
  • softmax 路由;
  • NoPE 层;
  • 读取第 35 层潜在缓存的 KV 共享层;
  • 索引共享层。

故障即停。 等待对等 GPU 的常驻核绝不能把服务器挂死。核内每次等待超过 10 秒就放弃,放弃时会在每个 rank 上置一个粘性故障字,其他 rank 的等待随即结束。每次前向之后,设备端断言会在计算 logits 之前检查该字。若该字已置位,本步剩余工作不再执行,worker 中止、引擎停止,失败步骤产生的任何 token 都不会到达客户端。这项检查每次启动约耗时 1 µs。

这条路径同样默认拒绝。如果部署用到了核不支持的任何东西,启动就会失败,并给出一个错误列出全部原因;所有 TP rank 在任何集合通信开始前都会对每项检查达成一致。不支持的特性包括其他 TP 尺寸、专家并行或上下文并行、LoRA、KV 连接器、其他投机方法,以及未以 a4w4 形式存储的专家。

TPU

Mistral Large 4 从第 0 天起就运行在 TPU v7x 上,使用 vllm-torchtpu,并支持 Pixtral 视觉编码器。我们验证过的配置把主机内张量并行、跨主机数据并行与专家并行结合起来,共使用 16 块 TPU v7x 芯片。

在 TPU v7x 上服务: 以下命令在一台主机上启动单个数据并行 worker:

vllm serve mistralai/Mistral-Large-4-1T-A52B \
  --tensor-parallel-size 8 --data-parallel-size 4 --data-parallel-size-local 1 \
  --data-parallel-start-rank "$TPU_WORKER_ID" \
  --data-parallel-address "$HEAD_IP" --data-parallel-rpc-port 29696 \
  --enable-expert-parallel --enable-ep-weight-filter \
  --dtype bfloat16 --kv-cache-dtype fp8_e4m3

TPU 实现依赖 vllm-torchtpu 的稀疏 MLA 和 DSA 索引器 Pallas 核,并把 MLA 扩展为跨层 KV 缓存共享。MLA 潜在 KV 缓存以 FP8 存储,稠密线性投影支持 FP8 权重和激活量化。路由专家使用仅权重量化 FP4(W4A16)。

HiSparse:主机内存稀疏 KV 卸载

稀疏注意力每个查询只读取 2048 个键,所以长请求的大部分 KV 不必留在 GPU 内存里。HiSparse 把完整的潜在 KV 放在锁页主机内存中,保留一份常驻 GPU 工作集加每请求热缓冲区,未命中时从主机拷贝选中的行。用 HiSparseConnector 启用:

--kv-transfer-config '{"kv_connector": "HiSparseConnector", "kv_role": "kv_both",
                       "kv_connector_extra_config": {"host_pool_gib": 440}}'

已知问题。 当张量并行组位于单节点时,还要设置 VLLM_FLASHINFER_ALLREDUCE_BACKEND=trtllm。一旦 HiSparse 把 KV 溢出到主机,默认的单节点 FlashInfer all-reduce(MNNVL)可能以 CUDA error: unspecified launch failure 让引擎崩溃(vllm#64327)。修复在 flashinfer#7782 中。

针对 ML4,我们让 HiSparse 支持了跨层 KV 共享。九个消费层绑定到第 35 层的常驻视图和主机副本,只保留自己的热缓冲区和 LRU(每层有自己的索引器)。没有数据重复,也不需要环境变量。与把共享层实体化相比,同样的主机池能多容纳 14.3% 的 token。我们还修复了:

  • 从多模态模型的文本配置中读取 index_topk;
  • CUDA 图捕获期间出现双重主机镜像。

容量。在单台 GB300 节点上(TP4、NVFP4),四个 TP rank 共享同一个去重后的主机内存池,因此 host_pool_gib 是总量。在 440 GiB 的池下,HiSparse 启动时可接纳 877 万 KV token,稳态下最高可容纳 1125 万,而仅靠 GPU 只有 197 万(4.5 倍到 5.7 倍)。GSM8K-200 结果不变(两者均为 0.975),大海捞针测试在约 9.8 万(10/10)和约 24.8 万(5/5)token 下均通过。

HiSparse KV capacity
HiSparse KV 容量

单台 GB300 节点上一个 TP4 服务器的 KV token 容量(NVFP4,max_model_len 262k,关闭前缀缓存):仅 GPU 为 197 万 token;启用 HiSparse 并由四个 TP rank 共享 440 GiB 主机内存池后,启动时可接纳 877 万(4.5 倍),稳态下最高 1125 万(5.7 倍)。复用第 35 层 KV 的层在主机上只存一份。

下一步

  • 上游合入:MLA 跨层 KV 共享、DSpark 滑动窗口 MLA 草稿,以及 DCP 下的 DSpark 和 HiSparse KV 共享支持。
  • 一个在 P/D 中运行 DSpark 的 DCP 解码节点,以及独立的草稿块池,使 DCP 下的 DSpark 保持完整 KV 容量。
  • HiSparse 支持 NIXL 导入的 GPU 落地,以及面向多节点 TP 的节点共享主机内存池。
  • 更多面向低并发延迟的解码融合。

致谢

vLLM 对 Mistral Large 4 的支持由 Nicolò Lucchesi 和 Mistral Inference 团队牵头合作完成。我们还要感谢审阅、设计并合入大量底层工作的 vLLM 维护者:Lucas Wilkinson、Flora Feng、Matt Bonanni 以及 Red Hat 团队的其他成员;整个 Inferact 团队;Cyrus Leung 以及整个 vLLM 社区。

在 AMD 方面,我们还要感谢 Hongxia Yang、Felix Marty、Andy Luo、Indira Vats 在 AMD Instinct GPU 上的 Day 0 支持与优化,以及 Bowen Bao、Chuan Li、Peng Sun 的总体支持。

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