vLLM 首日支持 Mistral Large 4,覆盖 NVIDIA、AMD 与 TPU
Day 0 support for Mistral Large 4 on vLLM
vLLM 宣布在发布当天支持 Mistral Large 4,这是一个 1 万亿参数、每 token 激活 49B 的 MoE 原生多模态模型,提供 FP8 与 NVFP4 两种 checkpoint。
原文给出 Mistral Large 4 在 vLLM 上的首日支持细节与各硬件配置的实测吞吐,可据此判断部署选型。
我们很高兴地宣布,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/ 包中,这项工作还带来了对解码上下文并行(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 激活 49 亿(含嵌入和输出层为 52 亿)。
- 检查点: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-4FP8,单节点 8 GPU,或两个 GB200/GB300 节点各 4 GPU。权重需要约 1.1 TB 的 GPU 内存。
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 节点上,添加 --nnodes 2 --node-rank <0|1> --master-addr <node0> --distributed-executor-backend mp,并在节点 1 上添加 --headless。
单节点 8-GPU 配置已在 B200 上验证:TP8 加 DSpark 在 GSM8K 上的得分与分离式运行在噪声范围内相当,而在 8+8 B200 上采用 NVFP4 1P1D、P DP8/EP → D TP8,在 32k 输入时达到 164 输出 tok/s/GPU(100 tok/s/用户),60k 时为 83(150 tok/s/用户)。相同布局下 B200 的每 GPU 吞吐约为 GB300 的 0.8 倍。
NVFP4。NVFP4 检查点每 GPU 使用约 76 GiB 权重,而 FP8 约为 128 GiB,因此可为 KV 缓存留出更多内存。
vllm serve mistralai/Mistral-Large-4-1T-A52B-NVFP4 \
--tensor-parallel-size 8 \
--speculative-config '{"method": "dspark", "num_speculative_tokens": 3}'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 消耗对比。误差条表示通过 bootstrapping 估计的 95% 置信区间。在所有领域,包括多模态推理(GDP.pdf、MathVision)、长上下文(Graphwalks 512k)、编码智能体(DeepSWE、SWE-bench Pro、Terminal-Bench 4)、业务自动化(AutomationBench)和网络安全(CyberGym),NVFP4 checkpoint 均与 FP8 行为一致。
长上下文(DCP)。 添加 --decode-context-parallel-size 8 以将潜在 KV cache 分片到 8 个 rank 上。它需要默认的 FP8 KV cache。
多模态。 传入 --limit-mm-per-prompt '{"image": N}' 以允许每个 prompt 最多 N 张图片。
有关已测试的部署配置(单节点、多节点、分离式 P/D 以及特定硬件设置),请参阅 Mistral Large 4 recipe。
实现细节
模型架构
| 组件 | 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 层的潜在 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 层的潜在 KV,且仅第 0 层为 dense。*
在 vLLM 中支持每个部分所需的工作:
稀疏 MLA + DSA indexer。 ML4 复用 vLLM 的稀疏 MLA 栈(带 fp8_ds_mla 潜在 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 潜在表示,5,808 字节用于 indexer。
跨层 KV 共享。 九个 NoPE 层读取第 35 层写入的潜在 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 图与融合。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 head,因此在 TP 下缓存会在每个 rank 上复制。DCP 则按 token 对其进行分片。在 TP8 上,DCP8 将 KV 容量从 2.38M 提升至 18.88M tokens(FP8),从 3.60M 提升至 28.76M tokens(NVFP4)。

并发的 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 稀疏(已合并为 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 槽位,并用
topk_length限定 kernel 边界,而不是扫描 2048 个充满-1的槽位。在 5k-token 预填充 profile 中,attention 从 165 ms 降至 38 ms。 - 一个融合 kernel 中和空行(即某个 rank 不拥有某查询的任何 top-k 槽位)。
DCP 专属开销降至每步 3.2 ms。结合解码融合,在 TP8 DCP8 FP8 上(60k 输入 / 1k 输出):
| 之前 | 之后 | |
|---|---|---|
| c1 TTFT p50 | 6.89 s | 4.73 s |
| c1 TPOT p50 | 17.82 ms | 12.50 ms |
| c64 output tok/s | 139.1 | 204.0 (+47%) |

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 对 8.7 ms,在 60k 下为 11.4 对 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 依赖于批次的预填充/解码拆分。在分段图中捕获时,该拆分在捕获时被冻结。现在只有 gather 作为 eager 断点运行,写入固定大小的缓冲区,而 cache 写入本身保留在图中。
- 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.* 个张量以及 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 个真实提示,贪婪(温度 0),最多 512 个输出 token:
| FP8 基础 TPOT | FP8 DSpark TPOT (AL) | NVFP4 基础 TPOT | NVFP4 DSpark TPOT (AL) | |
|---|---|---|---|---|
| c1 | 8.84 ms | 3.50 ms (2.90) | 8.21 ms | 2.95 ms (2.93) |
| c8 | 15.28 ms | 7.08 ms (2.85) | 12.85 ms | 5.80 ms (2.86) |
| c32 | 29.92 ms | 11.57 ms (2.85) | 25.38 ms | 9.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,使用时 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 降至
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。
滑动窗口 MLA 草稿层和 DCP 下的 DSpark 正在 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(8 GPU)的最佳输出 tok/s:
| P → D | 8k,≥40 tok/s/用户 | 8k,≥100 | 32k,≥40 | 32k,≥100 | 60k,≥40 / ≥100 | c1 TTFT 8k / 32k / 60k |
|---|---|---|---|---|---|---|
| P DP4/EP → D DP4/EP | 718.2 | 176.7 | 202.6 | 139.2 | 108.8 | 0.49 / 1.90 / 3.65 s |
| P DP4/EP → D TP4 | 459.2 | 194.2 | 205.2 | 150.0 | 107.2 | 0.46 / 1.88 / 3.66 s |
| P PCP4×DCP4 → D TP4 | 350.6 | 192.6 | 91.5 | 91.5 | 49.8 | 0.35 / 1.37 / 2.57 s |
| P TP4 → D TP4 | 259.2 | 185.3 | 62.8 | 62.8 | 32.3 | 0.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 checkpoint(Mistral-Large-4-1T-A52B-Quark-MXFP4)。Quark 仅对路由专家进行重新量化,采用静态权重和动态激活缩放。注意力、共享专家、稠密前馈层、路由器、输出头、视觉编码器和 DSpark drafter 保持原始精度。该 checkpoint 从约 1,061 GB 缩小至 578 GB,因此可在四块 MI350 系列 GPU(每块 288 GB)上提供服务。服务命令、评估结果和复现配方见模型卡。
- 在加载时重新量化 FP8 checkpoint(
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 cache 和激活。
特性。 ROCm 上支持 TP8 和 TP4、专家并行、通过 Pixtral 编码器输入图像,以及使用 checkpoint 内置 drafter 的 DSpark 投机解码。在 MI350X 上以 TP4 和 FP8 KV cache 运行时,DSpark 在真实提示上每步接受 3.1 个 token,并在一个请求时将每 token 解码延迟降低 2.7 倍,8 个请求时降低 1.9 倍,32 个请求时降低 1.8 倍(输出吞吐量分别为 2.1 倍、1.6 倍和 1.4 倍)。在 ROCm 上,潜在 KV cache 默认是 FP8(E4M3),其 KV 容量是 bf16 的 1.86 倍;--kv-cache-dtype bfloat16 可选择退出。
AITER。VLLM_ROCM_USE_AITER=1 在 AMD 的内核库 AITER 上运行 vLLM 的热路径:融合 MoE(MXFP4 专家运行在两阶段 FlyDSL 内核上)、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 将大型 prefill 步骤(最高 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 倍。更大的批次和 prefill 运行标准路径。它运行在默认 FP8 KV cache 上,需要 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 mistralROCm 实现细节
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 分页读取器所期望的布局。写入方和所有读取方(插入、预填充收集和分页 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 都在 kernel 内运行,通过四块 GPU 之间的对等内存进行,因此一层只需一次启动,无需主机往返。这些 kernel 针对 1、2、4 行构建,这正是 vLLM 完整的 CUDA 图捕获尺寸,并在这些图内重放。其他任何步骤都运行 vLLM 的层:预填充、混合批次、更宽的批次。该 kernel 改编自 ROCm ATOM 的 GLM-5.2 融合解码,并针对 ML4 的层类型进行了扩展:
- 无查询 LoRA;
- 查询和索引查询上的 SSMax;
- softmax 路由;
- NoPE 层;
- 读取第 35 层潜在缓存的 KV 共享层;
- 索引共享层。
故障停止。 等待对等 GPU 的持久 kernel 绝不能挂起服务器。每次 kernel 内等待在 10 秒后放弃,放弃会在每个 rank 上设置一个粘性故障字,从而立即结束其他 rank 的等待。每次前向之后,设备端断言会在计算 logits 之前检查该字。如果该字已设置,该步骤的剩余工作将永不运行,worker 中止,引擎停止,因此失败步骤产生的任何 token 都不会到达客户端。该检查每次启动约耗时 1 µs。
该路径同样是默认拒绝。如果部署使用了 kernel 不提供的任何内容,启动会失败,并给出一个列出所有原因的单一错误,且所有 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_e4m3TPU 实现依赖 vllm-torchtpu 的稀疏 MLA 和 DSA 索引器 Pallas kernel,并通过跨层 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}}'对于 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 时通过。

单个 GB300 节点上一台 TP4 服务器的 KV token 容量(NVFP4,max_model_len 262k,前缀缓存关闭):仅 GPU 上就有 1.97M tokens;借助 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