跳到正文
vLLM Blog· Mistral and the vLLM Team·· 16 小时前精选AI 评分72

vLLM 在发布当天支持 Mistral Large 4

Day 0 support for Mistral Large 4 on vLLM

AI 导读

vLLM 宣布在 Mistral Large 4 发布当天提供支持,该模型是 1 万亿参数、每 token 激活 49B 的 MoE 原生多模态模型,采用 MLA 加 DSA 稀疏注意力并内置 DSpark 投机解码草稿器。

推荐理由

原文给出 Mistral Large 4 在 vLLM 上的完整支持细节与实测数字,读者可据此判断长上下文与多模态部署的取舍。

正文 · AI 翻译

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

支持它不仅仅需要一个模型文件。ML4 运行在 vLLM 的模型运行器 V2(MRV2)上,并支持可中断的 CUDA 图。其构建模块位于新的 vllm/models/mistral/ 包中,这项工作还带来了对解码上下文并行(DCP)、预填充上下文并行(PCP)、DSpark、基于 NIXL 的分离式服务和 HiSparse 主机内存 KV 卸载的通用改进。

亮点

  • 架构:72 层 MoE,128 个路由专家(top-4)加 1 个共享专家,隐藏维度 6144。注意力为 MLA(KV LoRA rank 512,RoPE 维度 64),配有 DSA 稀疏索引器,每个查询选择 top 2048 个键。总参数 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:在搭载 ROCm 的 AMD Instinct MI350 系列 GPU 上实现 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 倍),并持有约 1.5 倍的 KV 缓存。

*核心评测中 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 以在 8 个 rank 间分片潜在 KV 缓存。它要求使用默认的 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 rank 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 层为 dense;第 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 层为 dense。*

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

稀疏 MLA + DSA indexer。 ML4 复用了 vLLM 的稀疏 MLA 栈(带 fp8_ds_mla latent cache 的 FlashMLA sparse,以及带 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_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 构建的层规格,而不是共享代码中的模型特定分支。

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

可断开的 CUDA 图和融合。 ML4 在无 torch.compile 的情况下运行。它使用可断开的 CUDA 图,该架构会自动启用(当文本模型位于 多模态包装器内时也是如此)。相比 eager 模式,它们带来了约 6 倍的加速。 手写融合替代了编译器原本会执行的操作,同时支持进一步优化:

  • 第一轮:SSMax、MLA RoPE + SSMax、indexer Q 路径和 indexer K 路径。这些将解码 TPOT 降低了 8–10%。
  • 解码轮次:FP32 router GEMM、为 indexer query 使用单次启动的 Triton LayerNorm、在一个 FlashInfer kernel 中完成 all-reduce + residual + RMSNorm,以及将 indexer K 直接写入 FP8 cache。在 TP8 FP8 上,c1 TPOT 从 10.22 降至 9.00 ms(-12%),c8 从 16.01 降至 15.08 ms。

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

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

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

并发 128k / 256k / 1M-token 请求,其 KV cache 在 8× GB300 上以 ML4 NVFP4 可容纳,TP8 与 TP8 DCP8 对比:floor(KV pool / input length) - 3.60M 和 28.76M tokens,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 中,选择器就地读取。这替代了每个 indexer 层的一次 NCCL all-gather 和一次 6-pass radix select。
  • Top-k 压缩:对于 32+ token 的批次,每个 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 output 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%:c1 TPOT 在 8k 输入时为 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 耗尽并攀升至 15 s。

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

  • 在 PCP 下,MLA KV-cache 写入首先会在 PCP rank 之间对其输入执行 all-gather,而该 gather 取决于批次的 prefill/decode 划分。在分段图中捕获时,该划分在捕获时被冻结。现在只有 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.* 个张量和 mtp 块,位于 params.json 中)。DSpark 是一种并行 drafter:一个小型非因果 transformer 在一次前向传播中提出一个块的所有 draft 位置。它以从目标层 1、18、36、54 和 71 提取的隐藏状态为条件。ML4 的 drafter 增加了:

  • RNN Markov 头:一个门控循环状态(秩 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 个真实提示,贪婪(温度 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)

该表使用贪婪解码。使用采样时,接受率保持接近:在温度 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 预填充节点可以移交给 TP 解码节点,两侧均启用 DSpark。NIXL 从一个预填充 rank 拉取复制的草稿窗口,并以 DCP 切片拉取目标 KV。任何 P/D 布局下的接受率相同:在 8k / 32k / 60k 输入下约为 2.98 / 2.86 / 2.78。

DCP 下的滑动窗口 MLA 草稿层和 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 从第 0 天起就在 AMD Instinct MI350 系列 GPU 上使用 vLLM 的 ROCm 构建运行。稀疏 MLA 运行在 AITER 的 MLA 解码内核上,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 副本,每个副本剩余的 HBM 约为 FP8 专家在同样四块 GPU 上所能留下的 6 倍,可用于 KV 缓存和激活。

特性。 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(在两阶段 FlyDSL 内核上运行 MXFP4 专家)、块 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 在 AITER 的 Triton 内核上运行稀疏 MLA,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 层没有可供 AITER 内核读取的 RoPE 部分,因此它们运行在 vLLM 的 ragged Triton 稀疏注意力内核上。

AITER 上的 DSA 索引器。一个融合 Triton 内核将每个索引器键量化为 FP8(使用单一缩放因子)并写入索引器缓存,每个索引器层每个 token 占 132 字节。在一个缓存块内,键以 16×16 瓦片存储,这正是 AITER 分页读取器所期望的布局。写入方和每个读取方(插入、预填充 gather 和分页 logits)共享同一条布局规则,因此它们不会出现不一致。预填充打分运行在 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 存储,使用单一逐张量缩放因子,每层每 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 时,检查点加载器保持不变:它仍然加载带块缩放因子的块 FP8 专家权重。这些参数起始于 meta 设备,逐层加载器一次物化一层。某一层的专家加载完成后,每个专家立即被反量化为 bf16 并重新量化为 MXFP4(E2M1 值,每 32 个值一个 E8M0 缩放因子),同时释放该层的 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 内核,通过跨层 KV 缓存共享扩展 MLA。我们将 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 在启动时接纳 8.77M 个 KV token,稳态下最多可保持 11.25M 个,而仅使用 GPU 时为 1.97M(4.5 倍到 5.7 倍)。GSM8K-200 保持不变(两者均为 0.975),大海捞针测试在约 98k(10/10)和约 248k(5/5)token 下通过。

HiSparse KV capacity
HiSparse KV 容量

单个 GB300 节点上一个 TP4 服务器的 KV token 容量(NVFP4,max_model_len 262k,前缀缓存关闭):仅 GPU 时为 1.97M token;使用 HiSparse 并由四个 TP rank 共享一个 440 GiB 主机内存池时,启动时接纳 8.77M(4.5 倍),稳态下最多可达 11.25M(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 Blog · vllm.ai