我们很高兴地宣布,在 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 秩 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-4Docker 镜像发布,该镜像还包含mistral-common1.13。以下命令在该镜像vllm serve内运行:

dockerpullvllm/vllm-openai:mistral-large-4

8 GPU 上的 FP8:

vllmservemistralai/Mistral-Large-4-1T-A52B\
--tensor-parallel-size8\
--speculative-config'{"method": "dspark", "num_speculative_tokens": 3}'\
--tool-call-parsermistral--enable-auto-tool-choice\
--reasoning-parsermistral

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

vllmservemistralai/Mistral-Large-4-1T-A52B-NVFP4\
--tensor-parallel-size4\
--speculative-config'{"method": "dspark", "num_speculative_tokens": 3}'\
--tool-call-parsermistral--enable-auto-tool-choice\
--reasoning-parsermistral

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

*FP8 与 NVFP4 在核心评测上的得分和 token 消耗对比。误差条表示通过自助法(bootstrapping)估计的 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 rank 512,RoPE 维度 64
稀疏注意力DSA 索引器,top-k 2048,32 个索引头;44 层拥有索引器,28 层复用最近的更早索引器层的 top-k
位置编码YaRN + LeRoPE(按频率学习缩放);每第 4 层为 NoPE
query 缩放每个注意力层和索引器均使用 SSMax(长度相关的 query 缩放)
KV 共享9 层(39、43、…、71)读取第 35 层的潜在缓存(每 4 层 1 次)
MoE第 0 层为稠密层;第 1–71 层:128 个专家,top-4,1 个共享专家
输出logit 软上限 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 栈(带 latent cache 的 FlashMLA 稀疏版本,fp8_ds_mla外加带 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,因此它们使用相同的内核和 cache 布局。LeRoPE 被折叠进 YaRN 频率中。SSMax 被融合进 MLA RoPE 内核(见下文融合部分)。

一个 Mistral 模型包。移植没有向deepseek_v2.py添加标志,而是基于vllm/models/mistral/下可复用的 Mistral 组件构建:attention、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 attention 工具会发现 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 缓存。在 TP8 FP8 上,c1 TPOT 从 10.22 降至 9.00 ms(-12%),c8 从 16.01 降至 15.08 ms。

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

DCP。MLA 只保留一个 latent KV 头,因此在 TP 下缓存会在每个 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 缓存在 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) 候选推送到对等节点的对称内存收件箱中,选择器就地读取它们。这替代了每个 indexer 层的一次 NCCL all-gather 和一次 6 轮基数选择。
  • Top-k 压缩:对于 32+ token 的批次,每个 rank 压缩其有效的 top-k 槽位,并用以下方式约束 kernel 边界topk_length,而不是扫描 2048 个充满-1的槽位。在 5k-token 的预填充 profile 中,注意力从 165 ms 降至 38 ms。
  • 融合内核消除了空行(即不拥有查询任何 top-k 槽位的 rank)。

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 优化在 day 0 随 Mistral Large 4 支持一同发布。

何时使用 DCP。DCP 以每步延迟换取 KV 容量。在相同的 8x GB300 NVFP4 部署和工作负载下(无投机解码),DCP8 的每 token 解码速度比 TP8 慢约 16–28%: c1 TPOT 在 8k 输入时为 11.1 对 8.7 ms,在 60k 时为 11.4 对 9.1 ms。在最多 128 个并发的 8k–60k token 请求下,TP8 的 3.6M token 池从不是瓶颈,因此 TP8 在那里提供更高的 吞吐量。一旦 KV 成为限制,DCP8 就能带来收益:超长上下文,或大量并发的 长生命周期会话。在我们的智能体 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 依赖于批次的预填充/解码拆分。在分段图中捕获时,该拆分在捕获时被冻结。现在只有 gather 作为 eager 断点运行,写入固定大小的缓冲区, 而缓存写入本身保留在图中。
  • DCP 组的 rank 排序与集合通信返回数据的顺序一致。

在 TP1×PCP4×DCP4 上启用 graphs 后,GSM8K-200 从 0.315 提升到 0.975。配合 DSpark 则达到 0.980。 作为 P/D 预填充节点,PCP4×DCP4 给出了我们测得的最低 TTFT(见 P/D 部分)。

DSpark 投机解码

ML4 在检查点内附带一个 DSpark drafter(mtp_dspark.*张量和mtp块params.json)。DSpark 是一个并行草稿器:一个小型非因果 transformer 在一次前向传播中提出一个块的所有草稿位置。它以从目标层 1、18、36、54 和 71 提取的隐藏状态为条件。ML4 的草稿器增加了:

  • 一个 RNN 马尔可夫头:一个门控循环状态(秩 512)在各草稿位置上运行,因此每个位置的 logits 依赖于在它之前草拟的 token;
  • 位置特定的掩码输入:每个草稿槽获得自己学习的输入嵌入,而不是一个共享的掩码 token;
  • 滑动窗口 MLA 草稿层:3 个 MLA 层,具有 2048 token 窗口和 BF16 KV 缓存,位于它们自己的 KV 缓存组中。Triton MLA 后端为它们增加了滑动窗口解码。

*一个 DSpark 步骤:目标提取馈入草稿的上下文 K/V,[anchor, mask_1, mask_2] 块在一次传递中通过 3 个滑动窗口 MLA 层,RNN 头按顺序发出 3 个草稿 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 倍。它不改变模型生成的内容:每个草稿 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 草稿。在 60k 上下文下的 TP8 DCP8 上,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_seqs128 以下、平均上下文低于约 99k 时不会生效。

P/D 中的 DSpark。一个 DCP 预填充节点可以将任务移交给一个 TP 解码节点,两侧均启用 DSpark。NIXL 从一个预填充 rank 拉取复制的草稿窗口,并以 DCP 切片拉取目标 KV。任何 P/D 布局下的接受率都相同:在 8k / 32k / 60k 输入下约为 2.98 / 2.86 / 2.78。

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

使用 NIXL 的分离式服务

我们在 4+4 块 GB300 GPU(一个 NVL72 机架)上使用 NIXL 对 1P1D 进行了基准测试。设置:

  • 两侧均启用 DSpark 的 NVFP4,处处使用 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,4+4 块 GB300 上的 1P1D。实线为并发扫描;空心/点线标记为固定 QPS 扫描。*

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

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 从发布首日起即可在 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在 AITER(AMD 的内核库)上运行 vLLM 的热路径:融合 MoE(在两阶段 FlyDSL 内核上运行 MXFP4 专家)、block-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,每节点一个副本
VLLM_ROCM_USE_AITER=1VLLM_FP8_MOE_REQUANT=mxfp4VLLM_FP8_MOE_REQUANT_ACT=a8w4\
VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA=1AITER_CUSTOM_AR_MAX_SIZE=134217728\
vllmservemistralai/Mistral-Large-4-1T-A52B--tensor-parallel-size8\
--speculative-config'{"method": "dspark", "num_speculative_tokens": 3}'\
--tool-call-parsermistral--enable-auto-tool-choice--reasoning-parsermistral

# TP4 搭配专用解码,适用于少量并发用户;每节点两个副本(第二个位于 GPU 4-7)
VLLM_ROCM_USE_AITER=1VLLM_FP8_MOE_REQUANT=mxfp4VLLM_FP8_MOE_REQUANT_ACT=a4w4\
VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA=1AITER_CUSTOM_AR_MAX_SIZE=134217728\
vllmservemistralai/Mistral-Large-4-1T-A52B--tensor-parallel-size4\
--additional-config'{"decode_backend": "latency_specialized"}'--max-num-seqs64\
--speculative-config'{"method": "dspark", "num_speculative_tokens": 3}'\
--tool-call-parsermistral--enable-auto-tool-choice--reasoning-parsermistral

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 存储,使用一个 per-tensor 缩放因子,每 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时,检查点加载器保持不变:它仍然加载带块缩放因子的 block-FP8 专家权重。其参数从 meta 设备开始,逐层加载器一次物化一层。某一层的专家加载完成后,每个专家立即被反量化为 bf16 并重新量化为 MXFP4(E2M1 值,每 32 个值一个 E8M0 缩放因子),同时释放该层的 FP8 副本。峰值内存为 MXFP4 模型加上一层 FP8 专家,这正是 1T 模型能加载到四块 GPU 上的原因。权重误差与直接将 bf16 权重量化为 MXFP4 大致相同。随后专家在 AITER 的融合 MoE 上运行:两阶段 FlyDSL 内核,对a8w4使用 FP8 激活,对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 都在 kernel 内、通过四块 GPU 之间的 peer memory 运行,因此一层只需一次 launch,没有 host 往返。这些 kernel 针对 1、2、4 行构建,正是 vLLM 的完整 CUDA graph 捕获尺寸,并在这些 graph 内重放。其他任何 step 都走 vLLM 的层:prefill、混合 batch、更宽的 batch。该 kernel 改编自 ROCm ATOM 的 GLM-5.2 融合 decode,并针对 ML4 的层类型做了扩展:

  • 无 query LoRA;
  • query 和 index query 上的 SSMax;
  • softmax 路由;
  • NoPE 层;
  • 读取第 35 层 latent cache 的 KV 共享层;
  • index 共享层。

Fail-stop。等待 peer GPU 的持久化 kernel 绝不能把服务器挂死。每次 kernel 内等待在 10 秒后放弃,放弃会在每个 rank 上设置一个粘性故障字,从而立即结束其他 rank 的等待。每次 forward 之后,在计算 logits 之前,一个 device 侧 assert 会检查该字。如果该字被设置,该 step 的剩余工作不再运行,worker 中止、引擎停止,因此失败 step 的任何 token 都不会到达客户端。该检查每次 launch 约耗时 1 µs。

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

TPU

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

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

vllmservemistralai/Mistral-Large-4-1T-A52B\
--tensor-parallel-size8--data-parallel-size4--data-parallel-size-local1\
--data-parallel-start-rank"$TPU_WORKER_ID"\
--data-parallel-address"$HEAD_IP"--data-parallel-rpc-port29696\
--enable-expert-parallel--enable-ep-weight-filter\
--dtypebfloat16--kv-cache-dtypefp8_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。默认的单节点 FlashInfer all-reduce(MNNVL)可能在 HiSparse 将 KV 溢出到主机后导致引擎崩溃(CUDA error: unspecified launch failure一旦 HiSparse 将 KV 溢出到主机(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 和 HiSparse KV 共享支持下的 DSpark。
  • 在 P/D 中带 DSpark 的 DCP 解码节点,以及独立的草稿块池,使 DCP 下的 DSpark 保持完整的 KV 容量。
  • HiSparse 支持 NIXL 导入的 GPU 落地,以及面向多节点 TP 的节点共享主机池。
  • 更多解码融合以降低低并发延迟。

致谢

vLLM 对 Mistral Large 4 的支持是由 Nicolò Lucchesi 和 Mistral 推理团队主导的合作项目。我们还要感谢审阅、设计并落地了大量底层工作的 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 的总体支持。