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

vLLM 首日支持 Mistral Large 4,覆盖 NVIDIA、AMD 与 TPU

Day 0 support for Mistral Large 4 on vLLM

AI 导读

vLLM 宣布在发布当天支持 Mistral Large 4,这是一个 1 万亿参数、每 token 激活 49B 的 MoE 原生多模态模型,提供 FP8 与 NVFP4 两种 checkpoint。

推荐理由

原文给出 Mistral Large 4 在 vLLM 上的首日支持细节与各硬件配置的实测吞吐,可据此判断部署选型。

正文 · 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/ 包中,这项工作还带来了对解码上下文并行(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-4

FP8,单节点 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)。

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 稀疏(已合并为 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 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 对 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 基础 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,使用时 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 → 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 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 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 分页读取器所期望的布局。写入方和所有读取方(插入、预填充收集和分页 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_e4m3

TPU 实现依赖 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 时通过。

HiSparse KV capacity
HiSparse KV 容量

单个 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