vLLM 发布 Kimi K3.1 Pro Day-0 支持,引入稀疏注意力优化

作者:Inferact and the vLLM Team原标题:vLLM Delivers Day-0 Support for Kimi K3.1 Pro
85AI 研判分

vLLM 宣布对 Moonshot AI 最新开源模型 Kimi K3.1 Pro 提供 Day-0 支持。该模型采用 2.8T 参数混合专家架构并引入稀疏注意力(DSA),vLLM 通过 Token 分片预填充和解码上下文并行等技术,在长上下文场景下显著提升推理效率与吞吐量。

报道正文

本文目录26 个章节

vLLM 为 Kimi K3.1 Pro 提供 Day-0 支持

Kimi K3.1 Pro 是 Moonshot AI 迄今能力最强的开放权重模型:这是一个拥有 2.8T 参数(激活参数 104B)的多模态 MoE 模型,具备 1M token 的上下文窗口,在编程、智能体(agentic)和知识基准测试中表现可与最强的闭源模型相媲美。

主要变化在于注意力机制:K3.1 Pro 在所有 24 层全注意力层中引入了稀疏注意力(DSA)以提升效率。vLLM 在 NVIDIA GPU、AMD GPU、Google TPU 以及阿里巴巴平头哥 PPU 上均提供了 Day-0 支持。

TL;DR

  • 混合模型上的稀疏注意力: K3.1 Pro 保留了 K3 的 KDA + MLA 混合架构,并在所有 24 层 MLA 层中添加了 DSA 索引器。在长上下文场景下,解码步骤每层读取的字节数约为密集 MLA 的四分之一。
  • Token 分片预填充注意力: 在 8× GB300 上处理 340K token 时,首字生成时间(TTFT)比头部(head)分片预填充快高达 1.69 倍。
  • 带稀疏索引器的解码上下文并行(DCP): 通过精确的 top-k 合并实现流量减少约 10 倍,使注意力模块在 DCP8 配置下速度提升 1.3–1.5 倍。DCP 预填充注意力每层运行速度快 4.1–4.3 倍,且 DCP8 可容纳的 KV token 数量是 TP8 的 7.2 倍。
  • 融合低延迟解码: 源自 Kimi K3 的专用 CuTe DSL 解码路径现已覆盖 DSA 层,使得在 8× GB300 上解码速度提升 1.2–1.6 倍,在 DCP8 配置下提升 1.9 倍。
  • 智能体服务性能: 在 NVIDIA GB300 上,对延迟敏感的场景下单用户吞吐量高达 375.9 tok/s,高吞吐量场景下单 GPU 总吞吐量达 37.8K tok/s;在 SemiAnalysis AgentX 基准测试中,NVIDIA Rubin 平台上的单用户吞吐量高达 815.6 tok/s,单 GPU 总吞吐量达 68.7K tok/s。
  • Day-0 零失配强化学习(RL): 首次实现 Day-0 支持伴随零失配 RL:vLLM 和 Megatron 针对基础模型和 LoRA 适配器生成的 log-probs 完全逐位一致(bit-identical),因此在 100 步的全模型训练过程中,每一步的 PPO KL 散度严格为 0。

快速开始

在 NVIDIA GPU 上:

文本
# See the Kimi K3.1 recipe for the Docker image and multi-node commands.
vllm serve moonshotai/Kimi-K3.1-Pro \
  --tensor-parallel-size 8 \
  --trust-remote-code \
  --load-format fastsafetensors \
  --enable-auto-tool-choice \
  --tool-call-parser kimi_k3 \
  --reasoning-parser kimi_k3 \
  --speculative-config '{"model":"Inferact/Kimi-K3.1-Pro-DSpark-preview","num_speculative_tokens":5,"method":"dspark","draft_sample_method":"probabilistic"}'

在 AMD GPU 上:

文本
VLLM_ROCM_USE_AITER=1 SAFETENSORS_FAST_GPU=1 \
vllm serve moonshotai/Kimi-K3.1-Pro \
      --served-model-name kimi-k31-pro \
      --tensor-parallel-size 8 \
      --trust-remote-code \
      --load-format fastsafetensors \
      --max-num-seqs 16 \
      --gpu-memory-utilization 0.85 \
      --speculative-config '{"model":"Inferact/Kimi-K3.1-Pro-DSpark-preview","num_speculative_tokens":5,"method":"dspark","draft_sample_method":"probabilistic"}'

Kimi K3.1 Pro 的新特性

Kimi K3Kimi K3.1 Pro
层数93:69 KDA + 24 全注意力93:69 KDA + 24 全注意力
全注意力层稠密 MLAMLA + DSA 索引器(top-2048),NoPE
索引器无32 头 × 128 维,FP8 键
MoE896 专家,16 激活,MXFP4不变
上下文窗口1M1M
总参数量2.8T2.8T

K3.1 Pro 中的稀疏注意力

每个稀疏注意力层在其 MLA 注意力旁边都有一个小型索引器。该索引器为每个 token 保留一个 FP8 键(含缩放因子共 132 字节),并使用 32 个轻量级头对每个查询的所有先前 token 进行评分。然后它保留得分最高的前 2,048 个 token。随后,MLA 注意力仅在这 2,048 个 token 上精确运行。DSA 最初在 DeepSeek-V3.2 中引入;我们的 DeepSeek-V3.2 博客文章更详细地介绍了索引器。

这一改进在长上下文解码阶段收益最大,因为该阶段受限于内存带宽。稠密 MLA 每层每个上下文 token 需读取 576 字节的 FP8 潜在向量。DSA 仍需扫描整个上下文,但只读取 132 字节的索引键。之后,它仅为这 2,048 个 token 读取完整的潜在向量。在长上下文情况下,这大约相当于稠密 MLA 每层读取字节数的 23%。

图 1:Kimi K3.1 Pro 中的稀疏注意力(DSA)。索引器对所有历史 token 进行评分,保留前 2,048 个,并且每个查询头仅对这些 token 的潜在向量进行精确注意力计算。请逐步查看。

KV Cache 占用空间

索引键是位于 MLA 潜在向量旁边的第二个 per-token 缓存。在张量并行(TP)下,两者在每个 GPU 上均被复制;KDA 状态按请求存储,与 K3 相同。

每 GPU 每 token 的 KV Cache(24 个全注意力层)字节数
Kimi K3,FP8 MLA 潜在向量13,824
Kimi K3.1 Pro,FP8 潜在向量 + FP8 索引键16,992
Kimi K3.1 Pro,NVFP4 潜在向量 + MXFP4 索引键10,080

Latent 和 index keys 可以独立进行量化。仅将 latent 量化为 NVFP4,就能使 K3.1 Pro 的性能超越使用 FP8 的 K3。若进一步选择将 index keys 量化为 MXFP4,则可额外容纳 11% 的 token,并在长上下文场景下使 indexer 速度提升 10–20%。结合这两种优化,八块 GB300 GPU 可容纳的 token 数量是使用 FP8 时的 1.7 倍;而采用 DCP8(分布式上下文并行度为 8)后,这一容量再乘以七:可同时支持约 36 个完整的 1M-token 上下文,相比之下,使用 TP8(张量并行度为 8)配合 FP8 仅能支持 3 个。

图 2:在 8× GB300 上实测的 KV cache 容量。更改 latent 和 index-key 格式、并行策略或上下文长度,即可查看可容纳的完整上下文数量。

面向三种状态类型的 KV Cache

每个 DSA 层存储两种不同宽度的 per-token cache:MLA latent(FP8 下为 576 字节,NVFP4 下为 352 字节)和 index keys(FP8 下为 132 字节,MXFP4 下为 68 字节)。两者的宽度互不为倍数。此外,69 个 KDA 层各自为每个请求添加一个循环状态(recurrent state),其大小远大于单层的一个 token page。如果缓存结构要求所有层的 page 大小相同,则必须将较小的 page 填充至最大尺寸,造成空间浪费。

vLLM 改为按 block 而非按 layer 对 cache 进行打包。cache 中的一个 block 是一条连续的内存行,包含该 block 在所有 24 个 DSA 层中的 MLA pages 和 index-key pages,以及共享该 block 的 KDA states。同一行内的 pages 可以具有不同大小,唯一的约束是该行必须容纳对应的 KDA states。block 大小为 2 的幂次,因此运行在不同 TP 规模下的 prefill 和 decode 实例的 blocks 可以整除对齐,便于进行 KV transfer。

Attention kernels 仍然看到它们所期望的小型固定大小 pages。vLLM 在每一步之前将 block tables 和 cache slots 转换为 kernel-page 地址,使得 FlashInfer 和 FlashMLA 的 sparse MLA kernels 以及 DeepGEMM 的 indexer kernels 无需修改即可读取打包后的 cache。

这种布局使整个技术栈的其他部分得以正常工作:

  • 任意 sparse MLA kernel 与 KV 格式兼容。 在 Blackwell 架构上,默认使用 FlashInfer 的 sparse MLA kernel 配合 FP8 latent。FlashInfer 新推出的 PrimsTS kernels 也支持 NVFP4 latent。
  • 分离式服务(Disaggregated serving)。 每个 block 是一段连续的内存区域,因此 P/D 之间的传输只需一次大型 RDMA 读取即可完成,而非数十次针对各层的小规模读取。
  • 共享基础设施。 相同的 packed planner 同时服务于 GLM-5.3-Flash 和 DeepSeek-V4,它们的 cache 都面临混合 page 大小的问题(参见 DeepSeek V4 博客文章)。

图 3:在 TP8(张量并行度为 8)下,打包缓存的布局方式。注意力组与三个 KDA 组共享同一个块行池;选定一个 token、一个 DSA 层和一个 KDA 层,即可看到它们分别落在哪个块和页中。

Prefill:按 Token 而非 Head 对注意力进行分片

在 TP8 配置下,每个 GPU 拥有 K3.1 Pro 96 个注意力头中的 12 个。这与稀疏 MLA prefill 阶段并不契合,原因与我们在 GLM-5.3 博客文章中所述一致:

  • MLA 的潜在 KV 由所有头共享,因此每个 GPU 需要读取整个缓存来仅服务这 12 个头,且稀疏内核会将这 12 个头填充至其最小 tile 大小。
  • Indexer 完全未按头拆分:每个 GPU 都对所有 token 进行评分,导致其二次方计算成本重复了八次。

对于大型 prefill 步骤,vLLM 改为按 token 对注意力进行分片。通过一次 all-to-all 通信,将每个 GPU 的查询从“所有 token,12 个头”转换为“1/8 的 token,全部 96 个头”。每个 GPU 针对其负责的 token 运行 indexer 和稀疏 MLA,随后第二次 all-to-all 通信将输出结果还原为按头布局。由于 KV 缓存在 TP 模式下已是复制状态,其他部分保持不变。MoE、序列并行以及缓存布局均维持原样。为了平衡因果注意力,每个请求被切分为 16 份,GPU r 负责第 r 份和第 15 − r 份,这与 prefill 上下文并行中的均衡划分策略相同。两次 all-to-all 通信均直接写入对等节点的对称内存。

图 4:在四张 GPU 上,token 分片 prefill 的前后对比。若按头分片,每个 GPU 都会重复执行 indexer 并填充其头 tile;经过 all-to-all 通信后,每个 GPU 仅使用所有头对均衡划分的 token 子集进行评分。可逐步查看该过程。

在 8× GB300 上使用 TP8、真实权重且每次处理一个请求的结果如下:

Prompt tokensTTFT,按头分片TTFT,按 token 分片加速比
27K0.82 s0.71 s1.16×
108K3.72 s2.89 s1.29×
340K15.89 s9.42 s1.69×

在长缓存前缀上进行短追加操作(常见于 agentic 轮次)也能受益。对于在 150K token 缓存前缀上追加 2K token 的情况,indexer 和稀疏 MLA 内核在每个 GPU 上的耗时约为 15 ms,而此前为 56 ms(8× GB300,FlashInfer 稀疏 MLA)。代价是交换缓冲区,每个 GPU 约占用 0.8 GiB 的 KV 缓存空间。

长上下文:结合稀疏 Indexer 的解码上下文并行

TP 会在每个 GPU 上复制 MLA 潜在变量,因此增加 GPU 数量并不会增加 KV 容量。解码上下文并行(DCP)改为按 token 对缓存进行分片:W 个 GPU 中的第 r 个存储每第 W 个 token。对于稠密 MLA,每个 GPU 仅对其分片执行注意力计算,并通过 log-sum-exp (LSE) 合并部分输出。

DSA 增加了一个步骤:每个 GPU 只能对自己拥有的 1/W 的 token 进行排序,因此在注意力计算开始前,必须在每一层跨 GPU 合并全局 top-2048。

流量减少 10 倍的精确 top-k 合并

通用合并方法会将每个 GPU 完整的本地 top-2048(以 FP32 分数和 token id 对的形式)发送给其他所有 GPU。在 8× GB300 上的 DCP8 配置下,处理 64 个类似 AgentX 的请求(每个请求包含 5 个 query token),该操作耗时 258 µs,占整个 671 µs 注意力模块时间的很大一部分。

vLLM 发送的数据量要少得多。每个 GPU 直接从 BF16 分数中挑选其最佳的 1.5 × 2048 / W 个候选项,并将每个候选项与其分数打包为 32 位,数据量减少约 10 倍。随后,每个 GPU 从收集到的候选项中通过基数选择算法选出相同的全局 top-2048。这种合并保持精确性:如果某个 GPU 持有的 top-2048 数量超过其发送的数量,选择过程会检测到这一情况,该 GPU 将针对该行重新发送其完整的本地 top-k。

图 5:在 8× GB300 上的 DCP8 配置下,每个请求包含 5 个 query token 的一个 DSA 注意力模块。精确的单轮合并比通用合并快 1.27–1.51 倍。悬停条形可查看具体数值。

从 batch size 32 开始,且在 128K 上下文长度下,K3.1 Pro 在 DCP8 下的稀疏注意力现在比 K3 的稠密 MLA 更快。

随后,部分注意力输出通过对称内存上的两个 CuTe DSL kernel 进行交换并通过 LSE 合并:相比现有的 CUDA 实现(在 8× GB300 上测试了 320 种 BF16 形状),速度提升高达 3.7 倍,且输出结果逐位一致。这有助于所有在 DCP 下运行的 MLA 模型。

无需合并的 DCP Prefill

普通的 DCP prefill 会为每个 prompt token 支付合并和 LSE 组合的成本,并在每个 GPU 上重复注意力计算工作。当 DCP 组即为 TP 组时,vLLM 转而将所有 GPU 上的当前步上下文(包括 MLA 潜在变量和索引键)收集到对称内存工作区中(使用 FP8 时,每 GPU 每 1M token 约占 0.7 GB)。然后在该工作区上运行上一节提到的按 token 分片的 prefill。这样既不需要合并也不需要组合,且输出结果与不使用 DCP 的按 token 分片 prefill 逐位一致。

单个 DSA 层,8× GB300 上的 DCP8,AgentX 混合 16K token 步长DCP PrefillGathered-context Prefill
FP8, FlashInfer trtllm-gen24.8 ms5.8 ms (4.3×)
FP8, FlashInfer PrimsTS24.8 ms5.7 ms (4.3×)
NVFP4, FlashInfer PrimsTS25.7 ms6.2 ms (4.1×)

端到端来看,在 DCP8 下单个请求的 Prefill 速度与 TP8 相当:在 27K、108K 和 340K token 时分别为 1.10、4.40 和 14.61 秒,而 TP8 下为 1.10、4.39 和 14.50 秒(8× GB300,FP8 KV)。

何时使用 DCP

DCP 以牺牲每步延迟为代价换取 KV 容量:DCP8 将每个 GPU 上的 MLA 和索引键缓存减少至八分之一,且 KDA 状态已按头部分片。当上下文适合时,在低并发下普通 TP 解码更快;而当 KV 成为瓶颈时(如极长上下文或大量长期运行的 Agent 会话),DCP 则更具优势。在 8× GB300 上,DCP8 可容纳 2180 万 token 的 FP8 KV,而 TP8 仅为 304 万 token(若使用 NVFP4 latent,则分别为 3210 万与 449 万 token)。在并发度为 1 时,TP8 的 TPOT 为 9.35 ms,DCP8 在 32K 上下文时为 11.48 ms(在 128K 上下文时分别为 9.45 vs 11.58 ms);融合解码路径可将 DCP8 的 TPOT 降至 5.96 和 6.04 ms。

8 张 GPU 以 TP8 模式也能服务完整的 1M token 请求:在 8× GB300 上使用 FP8 KV 时,缓存可容纳 304 万 token,即支持约 2.9 个并发的 1M token 请求。

低延迟解码:融合的 DSA 解码路径

对于 Kimi K3,vLLM 提供了一个可选的解码路径,该路径用针对小批量调优的 CuTe DSL 内核替换了几乎每一个算子。它将通信融合进计算中,重构了 MoE 调度并将权重预取到 L2(参见 Hyper Optimizing Kimi K3 in vLLM)。对于 K3.1 Pro,此路径现已覆盖 DSA 层。一个 DSA 层的运行流程如下:

  • 一个打包的输入投影,用于查询、KV 潜在变量、输出门控以及索引器的键和头权重;
  • 一个融合的索引器前端:在一个内核中完成查询归一化、索引器的查询投影、其 FP8 或 MXFP4 量化(并折叠入头权重)以及新 token 的索引键;
  • MLA 查询(上投影和吸收式键投影,将新 token 的潜在变量写入 KV 缓存),与索引器重叠执行而非在其之前运行;
  • 一个融合的索引器:由两个内核组成,分别实现键平衡评分器和精确 top-k 选择,输出注意力机制读取的物理 KV 行;
  • 稀疏注意力:FlashInfer 稀疏 MLA(根据批次大小选择 trtllm-gen 或 PrimsTS)或 FlashMLA 稀疏;在 DCP 下且批次大小为 1–8 时,使用针对该形状构建的 CuTe DSL 内核(在批次为 1 时每层耗时 11.9 µs,对比 21.4 µs);
  • 在一个内核中完成值投影和 sigmoid 门控,随后是 K3 的输出投影及其融合的 all-reduce。

在 DCP 下,查询被发布到每个 GPU,每个 GPU 的 top-k 候选项被精确交换并合并,部分输出通过 LSE 组合,所有操作均通过内核内的 NVLink 邮箱完成,路径上没有 NCCL 集合通信。

图 6:在 8× GB300 上使用 TP8 和 1K token 输入时的每输出 token 时间。融合解码路径在批次为 1 时快 1.63 倍,在批次为 64 时仍快 1.18 倍。悬停条形可查看具体数值。

结合 DSpark 推测解码(7 个草稿 token),融合路径将并发度为 1 时的每 token 时间从 2.86 ms 降低至 2.14 ms。GSM8K 准确率保持不变。在 DCP8 下,融合路径将并发度为 1 时的 TPOT 在 32K 上下文长度下从 11.48 ms 降至 5.96 ms,在 128K 上下文长度下从 11.58 ms 降至 6.04 ms。

使用 DSpark 进行推测解码

我们为 K3.1 Pro 训练了一个 DSpark 推测器,并发布了预览检查点:Inferact/Kimi-K3.1-Pro-DSpark-preview。在真实采样条件下(温度 1.0,top-p 0.95),该预览版在 GSM8K 上使用 7 个草稿 token 已达到平均接受长度 4.597。在 8× GB300 上使用 TP8 且批次大小为 1 时,这带来了 3.0 倍的解码加速。该推测器在 K3.1 Pro 高达 1M token 的完整上下文窗口内也能保持其准确率,这在长上下文 BEAM 基准测试中得到验证。

此检查点仅为预览版。我们目前正在训练一个更准确、更高效的推测器,并将很快发布。

智能体服务:AgentX

vLLM 在智能体(Agentic)服务负载上展现出强劲的性能。我们使用 SemiAnalysis AgentX 基准测试,对 NVIDIA GB300 和 NVIDIA Rubin 平台上的 Day-0 镜像进行了评测。

NVIDIA GB300

在对延迟敏感的前沿区域,Kimi K3.1 Pro 实现了每用户每秒 375.9 个 token 的吞吐量。沿着前沿曲线进一步延伸,它在保持每用户每秒超过 100 个 token 的同时,实现了每 GPU 每秒 19.2K 总 token 的服务能力。在吞吐量极限端,它达到了每 GPU 每秒 37.8K 个 token,约为最高交互性点的 26 倍,而每个用户仍能获得约每秒 54 个 token 的体验。

图 7:NVIDIA GB300 上 Kimi K3.1 Pro 的 AgentX 帕累托前沿,采用 4 个推测 token 和 3.36 个 token 的合成接受长度。将鼠标悬停在点上可查看其配置和指标。

NVIDIA Rubin

在 NVIDIA Rubin 平台上,Kimi K3.1 Pro 在对延迟敏感的前沿区域实现了每用户每秒 815.6 个 token 的吞吐量。沿着前沿曲线进一步延伸,它在保持每用户每秒超过 100 个 token 的同时,实现了每 GPU 每秒 46.0K 总 token 的服务能力。在吞吐量极限端,它达到了每 GPU 每秒 68.7K 个 token,约为最高交互性点的 51 倍,而每个用户仍能获得约每秒 75 个 token 的体验。

图 8:NVIDIA Rubin 上 Kimi K3.1 Pro 的 AgentX 帕累托前沿。并发度不超过 64 的点使用 7 个推测 token 和 3.84 个 token 的合成接受长度;更高并发度的点使用 3 个推测 token 和 3.0 的合成接受长度。将鼠标悬停在点上可查看其配置和指标。

Day-0 支持零失配强化学习

首次,vLLM 的 Day-0 支持推出了零失配(zero-mismatch)强化学习方案:结合 verl 实现 LoRA RL,其中 vLLM 作为 rollout 引擎,Megatron 作为训练器,且训练器的 log-probs 与 rollout 引擎的结果逐位一致。vLLM 不再通过重要性采样或路由重放来限制差距,而是在内核层面彻底消除这一差距:训练器的前向传播逐位复现了 rollout 结果,从而将 RL 训练的稳定性提升到了新的高度:

  • 全模型 LoRA RL。 Rank-32 的 LoRA 覆盖了所有注意力投影层以及全部 896 个专家模块,且基础模型和适配器在 vLLM 与 Megatron 之间均保持逐位一致。
  • 原生支持批次不变性内核。 设置 VLLM_BATCH_INVARIANT=1 后,token 的 log-prob 不再依赖于其所在的批次。vLLM 越来越多的高性能内核在设计上即具备批次不变性,例如 DeepSelect top-k。

图 9:相同的 token、相同的内核、相同的 log-probs。切换至默认 rollout 以进行对比。

在完整模型上运行超过 100 步后,rollout/trainer KL 在每一步都精确为 0.0,而 AIME mean@8 从 0.451 上升至 0.664(图 10)。

图 10:完整模型上的 100 步 LoRA RL 运行。

针对 batch-invariant 路径的性能优化正在进行中。源代码位于 Inferact/zero-mismatch-rl。

AMD

在 AMD Instinct MI355X 上,vLLM 将 MXFP4 索引与稀疏注意力及解码上下文并行相结合。这些优化降低了长上下文的内存需求,并提升了 prefill 和 decode 的效率。

量化:原生 MXFP4 索引器

索引器在扫描上下文后,会选取 2,048 个 token 进行注意力计算。MXFP4 将每个索引键从 FP8 的 132 字节减少到 68 字节(包含缩放因子),从而将索引缓存存储量减少了 48.5%。评分过程使用原生 FP4 矩阵指令,配合 FP32 累加以及跨头的加权归约。在推测验证期间,查询行之间复用键值以减少重复读取。

更小的索引缓存为缓存上下文留出了更多 GPU 内存,有利于高并发下的长上下文服务。在 AgentX 测试中,MXFP4-indexer 配置在并发数为 48 时,总工作负载吞吐量提高了 5.3%。

Prefill:上下文并行稀疏注意力

在 prefill 阶段,GPU 收集索引键并在彼此之间分配查询行。每个 GPU 针对完整上下文对其分配的行进行评分,并直接选择全局 top-k,避免了合并各 rank 候选项的过程。当这能减少通信并符合工作空间预算时,主注意力可以单独收集 KV,然后处理本地头的所有查询行。

结合 DCP 的 prefill 优化在开发基准测试中提升了吞吐量和延迟表现。

输入 / 输出 token 数总吞吐量 (tokens/s)输出吞吐量 (tokens/s)中位 TTFT (s)中位 TPOT (ms)
8K / 1K5,357.8 → 6,015.4 (+12%)595.3 → 668.42.462 → 1.87524.571 → 22.063
128K / 1006,950 → 17,678 (+154%)5.30 → 13.48268.7 → 105.5331.5 → 129.6

设置:MI355X,TP8/DCP8,并发数 16,64 条提示,禁用前缀缓存。

FlyDSL 稀疏 MLA 直接从 FP8 KV Cache 中读取选定的 token。每个 workgroup 处理一个 query row,覆盖其本地的所有 attention head,从而在 prefill 步骤中的大量 query rows 之间提供并行性。对于较小的 batch,则保留 Triton 路径,该路径在每个 row 内部暴露出额外的并行度。

实现方式基准测试时长 (s)TTFT 中位数 (s)输出吞吐量 (tokens/s)
Triton432.896.214.79
FlyDSL404.489.815.83
提升幅度降低 6.6%降低 6.6%提高 7.0%

配置:MI355X, TP8/DCP8, 并发数 16, 64 个 prompts, ISL/OSL 128K/100

Decode:融合低延迟解码

对于小型推测验证(speculative verification)batch,DSA 索引和稀疏注意力分别运行,随后通过一个持久化 monokernel 将注意力输出投影、AttnRes 以及完整的 latent MoE 计算(包括张量并行 all-reduce)融合到单次启动中。这减少了 kernel 启动开销,并使 MoE 权重加载与之前的 reduction 及残差处理重叠。该专用路径支持每步最多八个验证 token,包括单个请求携带七个 draft token 的情况。

指标常规路径专用解码路径提升幅度
TPOT 均值6.02 ms4.40 ms27.0%

通信:高效的跨 GPU 通信

在 decode 阶段,每个 GPU 交换候选分数(candidate scores),同时保持 token ID 本地存储。通过共享阈值和一致的平局处理机制来确定选择哪些本地候选项。对于 2,048 个候选项,仅交换分数即可将每个 query 每个 GPU 的候选负载从 16 KiB 减半至 8 KiB。

AITER 通过专用的 GPU kernel 加速 DCP all-gather,该 kernel 通过 HIP IPC 读取对端缓冲区并直接写入所需的输出布局。这将数据收集与布局转换相结合,避免了单独的重排序拷贝。在八块 MI355X GPU 上的独立微基准测试中,针对所测试的 query 和候选 tensor 形状,AITER 路径相比 RCCL 将 all-gather 延迟降低了 53–61%。

TPU

vLLM 通过 vllm-torchtpu 在 TPU 上实现了对 Kimi K3.1 Pro 的 Day-0 支持,并在 16 芯片的 TPU v7x 切片上完成了验证。其部署布局为 TP32,MoE 部分采用 EP32,并启用 DCP(分布式上下文并行),针对长上下文场景按 token 对 MLA 潜在向量和索引键进行分片。DSA 的计算任务被拆分到 TPU 的两类计算单元中:TensorCore 负责对 FP8 格式的索引键进行打分,随后 SparseCore 选取前 2,048 个 token,并将其潜在向量收集到一个密集缓冲区中,供 TensorCore 的 MXU 执行注意力计算。与 GPU 上的情况类似,DSA 和 KDA 状态共享同一个基于块索引的 KV 池,从而确保混合缓存中的前缀缓存功能正常工作。

Alibaba T-HEAD

Alibaba T-HEAD 团队通过 vllm-sail 插件,增加了对在 Alibaba T-HEAD Zhenwu PPU(包括 M890 和 810E)上运行 Kimi K3.1 Pro 的支持。已验证的配置如下:

  • 810E:双节点,TP32 / EP32,使用 W4A16 MoE 后端
  • M890:四节点,TP8 / DP4 / EP32,使用 DeepEP

这标志着 Zhenwu PPU 首次为主流开源大模型提供 Day-0 支持,后续将持续提供支持。

致谢

我们衷心感谢 Moonshot 团队开源了 Kimi K3.1 Pro。同时,我们也感谢 Inferact、AMD 和 Alibaba T-Head 在 vLLM 中为该模型提供了 Day-0 支持,以及 NVIDIA 提供的帮助。