大数跨境

vLLM 公开 Kimi K3.1 Pro 首日适配细节

vLLM 公开 Kimi K3.1 Pro 首日适配细节 AINLP
2026-11-07
2
导读:不止是能跑起来

Kimi K3.1 Pro 开源后,vLLM 团队公开了首日适配的工程细节:如何处理稀疏注意力与混合缓存,如何降低长上下文推理延迟,以及怎样让推理引擎和训练器给出逐比特一致的结果,以下是全文翻译。


Kimi K3.1 Pro 是月之暗面目前能力最强的开放权重模型:这款多模态 MoE 模型拥有 2.8 万亿总参数、1040 亿激活参数和 100 万 token 上下文,在编程、Agent 与知识类基准上可与最强的闭源模型比肩。

主要变化在注意力机制:为了提高效率,K3.1 Pro 在全部 24 个全注意力层中加入了稀疏注意力(DSA)。vLLM 在发布首日便支持其运行于 NVIDIA GPU、AMD GPU、Google TPU 和阿里平头哥真武 PPU。

要点速览

  • 混合模型中的稀疏注意力:K3.1 Pro 保留了 K3 的 KDA + MLA 混合架构,并为全部 24 个 MLA 层加入 DSA 索引器。在长上下文下,每层单次解码需要读取的字节数约为稠密 MLA 的四分之一。
  • 用一套 KV 缓存存放三类状态:vLLM 将每个块的 MLA 潜变量、索引键和 KDA 状态打包存放。这种缓存布局支持 FlashInfer 和 FlashMLA 的稀疏内核,支持 FP8 或 NVFP4 KV,以及可选的 MXFP4 索引键;预填充与解码分离部署时,可以原样传输。
  • 按 token 切分的预填充注意力:在 8 张 GB300 上处理 340K token 输入时,与按注意力头切分的预填充相比,首 token 延迟(TTFT)的加速比最高达到 1.69 倍。
  • 带稀疏索引器的解码上下文并行(DCP):精确 top-k 合并将通信量降至约十分之一,使 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 倍。
  • Agent 推理服务:在 SemiAnalysis AgentX 基准中,NVIDIA GB300 在低延迟端的单用户速度最高达到 375.9 token/s,在高吞吐端的每 GPU 总吞吐达到 37.8K token/s;NVIDIA Rubin 的对应结果最高分别为 815.6 token/s 和 68.7K token/s。
  • 首日即支持零偏差强化学习:这是首次在发布首日提供零偏差 RL 支持。对于基础模型和 LoRA 适配器,vLLM 与 Megatron 给出的对数概率(log-probs)都逐比特一致,因此,在覆盖整个模型的 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 K3
Kimi K3.1 Pro
层数
93 层:69 个 KDA 层 + 24 个全注意力层
93 层:69 个 KDA 层 + 24 个全注意力层
全注意力层
稠密 MLA
MLA + DSA 索引器(top-2048),NoPE
索引器
无
32 个头 × 128 维,FP8 索引键
MoE
896 个专家,激活 16 个,MXFP4
不变
上下文窗口
1M
1M
总参数量
2.8T
2.8T

K3.1 Pro 中的稀疏注意力

每个稀疏注意力层都在 MLA 注意力旁配有一个小型索引器。索引器为每个 token 保存一个 FP8 索引键,连同缩放因子共 132 字节;对于每个查询,它用 32 个轻量级头为所有先前 token 打分,然后保留得分最高的 2048 个。随后,MLA 只在这 2048 个 token 上执行精确注意力计算。DSA 最初由 DeepSeek-V3.2 引入,我们关于 DeepSeek-V3.2 的文章更详细地介绍了该索引器。

收益最大的是受显存带宽限制的长上下文解码。稠密 MLA 在每一层都要为上下文中的每个 token 读取 576 字节的 FP8 潜变量。DSA 仍会扫描整个上下文,但扫描的只是每个 132 字节的索引键,之后仅为选出的 2048 个 token 读取完整潜变量。在长上下文下,每层读取的字节数约为稠密 MLA 的 23%。

原文图 1:Kimi K3.1 Pro 的稀疏注意力(DSA)。索引器为所有先前 token 打分,保留最高分的 2048 个,每个查询头只对这些 token 的潜变量计算精确注意力。可在原文中逐步查看这一过程。

原文图 1:Kimi K3.1 Pro 的稀疏注意力(DSA)。索引器为所有先前 token 打分,保留最高分的 2048 个,每个查询头只对这些 token 的潜变量计算精确注意力。可在原文中逐步查看这一过程。

KV 缓存占用

除了 MLA 潜变量,索引键还需要第二份逐 token 缓存。在张量并行(TP)下,两者都会在每张 GPU 上各存一份;KDA 状态则与 K3 一样,按请求保存。

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

潜变量和索引键可以独立量化。仅将潜变量改为 NVFP4,就能让 K3.1 Pro 的占用低于采用 FP8 的 K3。在此基础上,可选的 MXFP4 索引键还能多容纳 11% 的 token,并在长上下文下使索引器加速 10%—20%。两者同时启用时,8 张 GB300 能容纳的 token 数是 FP8 配置的 1.7 倍;DCP8 又将这一容量扩大约 7 倍,可以同时容纳约 36 个完整的 100 万 token 上下文,而 TP8 搭配 FP8 时约为 3 个。

原文图 2:8 张 GB300 上实测的 KV 缓存容量。可在原文中切换潜变量与索引键格式、并行方式或上下文长度,查看能容纳多少个完整上下文。

原文图 2:8 张 GB300 上实测的 KV 缓存容量。可在原文中切换潜变量与索引键格式、并行方式或上下文长度,查看能容纳多少个完整上下文。

用一套 KV 缓存存放三类状态

每个 DSA 层都要保存两种宽度不同的逐 token 缓存:MLA 潜变量采用 FP8 时为 576 字节,采用 NVFP4 时为 352 字节;索引键采用 FP8 时为 132 字节,采用 MXFP4 时为 68 字节。这两种宽度互不构成整数倍关系。另外,69 个 KDA 层中的每一层都要为每个请求保存一份循环状态,其大小远超单层的一页 token。如果缓存为所有层分配相同大小的页,就不得不把较小的页填充到最大页的大小。

vLLM 改为按块打包缓存,而不是按层组织。一个缓存块对应一整行连续内存,其中放着该块在全部 24 个 DSA 层上的 MLA 页和索引键页,以及共享该块的 KDA 状态。同一行中的页可以大小不同,唯一的约束是这一行能容纳 KDA 状态。块大小取 2 的幂,使采用不同 TP 规模的预填充与解码实例的块能够整除对应,配对完成 KV 传输。

注意力内核看到的仍是它们所需的小型定长页。vLLM 会在每一步之前,将块表和缓存槽位转换为内核页地址,因此 FlashInfer、FlashMLA 的稀疏 MLA 内核以及 DeepGEMM 的索引器内核无需改动,就能读取打包后的缓存。

这一布局也让其余系统组件能够正常协作:

  • 支持各种稀疏 MLA 内核和 KV 格式。在 Blackwell 上,默认使用 FlashInfer 稀疏 MLA 和 FP8 潜变量;FlashInfer 新的 PrimsTS 内核也支持 NVFP4 潜变量。
  • 支持分离式推理服务。每个块都是一段连续内存,因此预填充与解码之间的传输只需一次大块 RDMA 读取,无需按层执行数十次小块读取。
  • 共享基础设施。同一个打包规划器也服务于 GLM-5.3-Flash 和 DeepSeek-V4,这些模型的缓存同样存在页大小混杂的问题,详见 DeepSeek V4 相关文章。
原文图 3:TP8 下的打包缓存布局。注意力组和三个 KDA 组共享同一块行内存池;可在原文中选择 token、DSA 层和 KDA 层,查看它们分别落入哪个块和页。

原文图 3:TP8 下的打包缓存布局。注意力组和三个 KDA 组共享同一块行内存池;可在原文中选择 token、DSA 层和 KDA 层,查看它们分别落入哪个块和页。

预填充按 token 切分注意力

TP8 下,每张 GPU 负责 K3.1 Pro 的 96 个注意力头中的 12 个。这并不适合稀疏 MLA 预填充,原因与我们在 GLM-5.3 文章中介绍的相同:

  • MLA 的潜在 KV 由所有注意力头共享,因此每张 GPU 都要读取完整缓存,却只服务 12 个头;稀疏内核还会把这 12 个头填充到自身最小计算块的大小。
  • 索引器根本不按头切分:每张 GPU 都会为所有 token 打分,因此其二次复杂度的计算被重复了 8 遍。

对于较大的预填充步,vLLM 改为按 token 切分注意力。第一次 all-to-all 通信,将每张 GPU 的查询从“全部 token、12 个头”转换为“八分之一的 token、全部 96 个头”。每张 GPU 为自己负责的 token 运行索引器和稀疏 MLA,再通过第二次 all-to-all 将输出恢复为按头划分的布局。TP 下 KV 缓存原本就已复制,因此其他部分不需要改变:MoE、序列并行和缓存布局都保持不变。为了平衡因果注意力的负载,每个请求被切成 16 份,GPU r 负责第 r 份和第 15 − r 份,与预填充上下文并行采用的均衡切分方式相同。两次 all-to-all 都直接写入其他 GPU 的对称内存。

原文图 4:4 张 GPU 上按 token 切分预填充的前后对比。按头切分时,每张 GPU 都会重复运行索引器,并为头维度的计算块做填充;经过一次 all-to-all 后,每张 GPU 只对均衡分配的一部分 token 计算,但使用全部注意力头。可在原文中逐步查看。

原文图 4:4 张 GPU 上按 token 切分预填充的前后对比。按头切分时,每张 GPU 都会重复运行索引器,并为头维度的计算块做填充;经过一次 all-to-all 后,每张 GPU 只对均衡分配的一部分 token 计算,但使用全部注意力头。可在原文中逐步查看。

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

输入 token 数
按头切分的 TTFT
按 token 切分的 TTFT
加速比
27K
0.82 s
0.71 s
1.16×
108K
3.72 s
2.89 s
1.29×
340K
15.89 s
9.42 s
1.69×

Agent 多轮交互中常见的另一种情况,也能受益:在较长的已缓存前缀后追加一小段输入。对于 150K token 的缓存前缀,追加 2K token 时,索引器和稀疏 MLA 内核的耗时从每 GPU 约 56 ms 降至 15 ms,配置为 8 张 GB300、FlashInfer 稀疏 MLA。代价是需要交换缓冲区,会占用每张 GPU 约 0.8 GiB 原本可用于 KV 缓存的空间。

长上下文中的稀疏索引与解码上下文并行

TP 会在每张 GPU 上复制 MLA 潜变量,因此增加 GPU 并不会增加 KV 容量。解码上下文并行(DCP)则按 token 切分缓存:W 张 GPU 中的第 r 张保存每隔 W 个 token 分配给自己的那些 token。对于稠密 MLA,每张 GPU 在自己的分片上计算注意力,再通过 log-sum-exp(LSE)合并部分输出。

DSA 多了一步:每张 GPU 只能对自己持有的那 1/W token 排序,因此每一层在计算注意力之前,都必须跨 GPU 合并出全局 top-2048。

通信量降至十分之一的精确 top-k 合并

通用合并方法会把每张 GPU 的完整本地 top-2048,以 FP32 分数与 token ID 对的形式,发给其他所有 GPU。在 8 张 GB300 上采用 DCP8,处理 64 个类似 AgentX 的请求、每个请求含 5 个查询 token 时,这一步占用了整个 671 µs 注意力模块中的 258 µs。

vLLM 传输的数据少得多。每张 GPU 直接从 BF16 分数中,只选出排名最高的 1.5 × 2048 / W 个候选,把每个候选与其分数打包为 32 位,传输字节数约降至十分之一。随后,每张 GPU 都从收集到的候选中,使用基数选择算法得到同一个全局 top-2048。合并仍然是精确的:如果某张 GPU 实际持有的全局 top-2048 条目比它发出的还多,选择过程会检测到这一情况,该 GPU 就会为这一行重新发送完整的本地 top-k。

原文图 5:8 张 GB300、DCP8 下的单个 DSA 注意力模块,每个请求含 5 个查询 token。精确的单轮合并比通用合并快 1.27—1.51 倍,可在原文中悬停柱形查看数值。

原文图 5:8 张 GB300、DCP8 下的单个 DSA 注意力模块,每个请求含 5 个查询 token。精确的单轮合并比通用合并快 1.27—1.51 倍,可在原文中悬停柱形查看数值。

在 128K 上下文、批量大小从 32 起的条件下,DCP8 下 K3.1 Pro 的稀疏注意力现在已经快于 K3 的稠密 MLA。

随后,系统使用两个基于对称内存的 CuTe DSL 内核交换部分注意力输出,并通过 LSE 合并。在 8 张 GB300 上测试 320 种 BF16 形状时,与现有 CUDA 实现相比,加速最高达到 3.7 倍,输出逐比特一致。所有在 DCP 下运行的 MLA 模型都能受益。

无需合并的 DCP 预填充

直接实现 DCP 预填充,会为每个输入 token 支付一次 top-k 合并和 LSE 合并的开销,并在每张 GPU 上重复注意力计算。当 DCP 组与 TP 组相同时,vLLM 改为从所有 GPU 收集这一步的上下文,包括 MLA 潜变量和索引键,放入对称内存工作区;FP8 配置下,每 100 万 token 需要每 GPU 约 0.7 GB 工作区。随后,在这份上下文上运行上一节介绍的按 token 切分预填充,不再需要合并。其输出与不启用 DCP 的按 token 切分预填充逐比特一致。

单个 DSA 层,8 张 GB300、DCP8,AgentX 混合 16K token 步
DCP 预填充
收集完整上下文后的预填充
FP8,FlashInfer trtllm-gen
24.8 ms
5.8 ms(4.3×)
FP8,FlashInfer PrimsTS
24.8 ms
5.7 ms(4.3×)
NVFP4,FlashInfer PrimsTS
25.7 ms
6.2 ms(4.1×)

端到端来看,DCP8 下单个请求的预填充速度与 TP8 相当:在 27K、108K 和 340K token 输入下,DCP8 分别为 1.10、4.40 和 14.61 s,TP8 分别为 1.10、4.39 和 14.50 s。配置为 8 张 GB300、FP8 KV。

何时使用 DCP

DCP 用逐步执行的延迟换取 KV 容量:DCP8 把每张 GPU 的 MLA 和索引键缓存缩小到八分之一,而 KDA 状态本来就按头切分。当上下文放得下时,单纯 TP 在低并发下解码更快;当 KV 容量成为限制时,例如超长上下文或大量长期存活的 Agent 会话,DCP 才更划算。在 8 张 GB300 上,DCP8 可以容纳 21.8M token 的 FP8 KV,而 TP8 为 3.04M;采用 NVFP4 潜变量时,分别为 32.1M 和 4.49M。并发为 1、上下文为 32K 时,TP8 的每输出 token 延迟(TPOT)为 9.35 ms,DCP8 为 11.48 ms;128K 时分别为 9.45 ms 和 11.58 ms。融合解码路径则将 DCP8 的这两个结果降至 5.96 ms 和 6.04 ms。

8 张 GPU 采用 TP8,也能处理完整的 100 万 token 请求:在 8 张 GB300、FP8 KV 下,缓存可容纳 3.04M token,相当于同时容纳 2.9 个 100 万 token 请求。

低延迟解码的 DSA 融合路径

对于 Kimi K3,vLLM 提供了一条可选的解码路径,用针对小批量调优的 CuTe DSL 内核替代几乎所有算子。它把通信融合进计算,重新安排 MoE 调度,并将权重预取到 L2,详见《在 vLLM 中深度优化 Kimi K3》。对于 K3.1 Pro,这条路径现已覆盖 DSA 层。一个 DSA 层按以下方式运行:

  • 一次打包的输入投影,同时处理查询、KV 潜变量、输出门控,以及索引器的键和头权重。
  • 融合的索引器前端:在一个内核中完成查询归一化、索引器查询投影、融入头权重的 FP8 或 MXFP4 量化,以及新 token 的索引键计算。
  • MLA 查询计算,包括升维投影和已吸收的键投影,并将新 token 的潜变量写入 KV 缓存;它与索引器重叠执行,而不是先于索引器执行。
  • 融合的索引器:用两个内核完成按键均衡的打分和精确 top-k 选择,直接输出注意力所需读取的物理 KV 行地址。
  • 稀疏注意力:使用 FlashInfer 稀疏 MLA,根据批量大小选择 trtllm-gen 或 PrimsTS,也可使用 FlashMLA sparse;在 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 倍。可在原文中悬停柱形查看数值。

原文图 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、32K 上下文时的 TPOT 从 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,使用 7 个草稿 token 时,预览版在 GSM8K 上的平均接受长度已达到 4.597。在 8 张 GB300、TP8、批量为 1 时,这相当于 3.0 倍的解码加速。长上下文 BEAM 基准的测量还表明,该草稿模型在 K3.1 Pro 最高 100 万 token 的完整上下文窗口内,都能保持准确性。

这还只是一个预览检查点。我们正在训练更准确、更高效的草稿模型,很快就会发布。

Agent 推理服务与 AgentX

vLLM 在 Agent 推理服务负载中表现出很强的性能。我们使用 SemiAnalysis AgentX 基准,在 NVIDIA GB300 和 NVIDIA Rubin 上测试首日发布的镜像。

NVIDIA GB300

在帕累托前沿偏重低延迟的一端,Kimi K3.1 Pro 的单用户速度保持在 375.9 token/s。沿前沿向更高吞吐移动时,它能在每 GPU 总吞吐达到 19.2K token/s 的同时,让每位用户的速度保持在 100 token/s 以上。在高吞吐端,每 GPU 达到 37.8K token/s,约为交互速度最高的那个点的 26 倍;此时,每位用户仍能获得约 54 token/s。

原文图 7:Kimi K3.1 Pro 在 NVIDIA GB300 上的 AgentX 帕累托前沿。使用 4 个投机 token,合成接受长度为 3.36 token。可在原文中悬停数据点查看配置与指标。

原文图 7:Kimi K3.1 Pro 在 NVIDIA GB300 上的 AgentX 帕累托前沿。使用 4 个投机 token,合成接受长度为 3.36 token。可在原文中悬停数据点查看配置与指标。

NVIDIA Rubin

在 NVIDIA Rubin 上,Kimi K3.1 Pro 在帕累托前沿偏重低延迟的一端,单用户速度保持在 815.6 token/s。沿前沿向更高吞吐移动时,它能在每 GPU 总吞吐达到 46.0K token/s 的同时,让每位用户的速度保持在 100 token/s 以上。在高吞吐端,每 GPU 达到 68.7K token/s,约为交互速度最高的那个点的 51 倍;此时,每位用户仍能获得约 75 token/s。

原文图 8:Kimi K3.1 Pro 在 NVIDIA Rubin 上的 AgentX 帕累托前沿。并发不超过 64 的数据点使用 7 个投机 token,合成接受长度为 3.84 token;更高并发使用 3 个投机 token,合成接受长度为 3.0 token。可在原文中悬停数据点查看配置与指标。

原文图 8:Kimi K3.1 Pro 在 NVIDIA Rubin 上的 AgentX 帕累托前沿。并发不超过 64 的数据点使用 7 个投机 token,合成接受长度为 3.84 token;更高并发使用 3 个投机 token,合成接受长度为 3.0 token。可在原文中悬停数据点查看配置与指标。

首日支持零偏差强化学习

vLLM 首次在模型发布首日提供了零偏差 RL 支持:使用 verl 进行 LoRA 强化学习,由 vLLM 充当 rollout 引擎、Megatron 充当训练器,训练器给出的对数概率与 rollout 逐比特一致。vLLM 不再依靠重要性采样或路由重放来约束两者之间的差距,而是在内核层面消除差距,让训练器的前向计算逐比特复现 rollout,将 RL 训练稳定性推向新的水平:

  • 覆盖整个模型的 LoRA RL。秩为 32 的 LoRA 覆盖每个注意力投影和全部 896 个专家;无论基础模型还是适配器,在 vLLM 与 Megatron 之间都逐比特一致。
  • 原生支持批量不变内核。设置 VLLM_BATCH_INVARIANT=1 后,某个 token 的对数概率不再取决于它所在的批次。越来越多 vLLM 高性能内核在设计上就具有批量不变性,例如 DeepSelect top-k。
原文图 9:相同的 token、相同的内核、相同的对数概率。可在原文中切换到默认 rollout 进行比较。

原文图 9:相同的 token、相同的内核、相同的对数概率。可在原文中切换到默认 rollout 进行比较。

在覆盖整个模型的 100 步运行中,rollout 与训练器之间的 KL 每一步都精确为 0.0,而 AIME mean@8 从 0.451 上升至 0.664,见图 10。

原文图 10:在整个模型上进行 100 步 LoRA RL 的结果。

原文图 10:在整个模型上进行 100 步 LoRA RL 的结果。

批量不变路径的性能优化仍在继续。源代码位于 Inferact/zero-mismatch-rl。

AMD

在 AMD Instinct MI355X 上,vLLM 将 MXFP4 索引、稀疏注意力和解码上下文并行结合起来,降低长上下文的显存需求,并提高预填充与解码效率。

量化与原生 MXFP4 索引器

索引器先扫描上下文,再选出 2048 个 token 计算注意力。计入缩放因子后,MXFP4 把每个索引键从 FP8 的 132 字节降至 68 字节,索引缓存占用减少 48.5%。打分使用原生 FP4 矩阵指令,以 FP32 累加,并跨头做加权归约。在投机验证阶段,不同查询行会复用键,以减少重复读取。

更小的索引缓存让 GPU 有更多显存用于缓存上下文,有利于高并发长上下文推理服务。在 AgentX 的参数扫描中,采用 MXFP4 索引器的配置在并发为 48 时,总负载吞吐提高了 5.3%。

预填充与上下文并行稀疏注意力

预填充期间,GPU 收集索引键,并在彼此之间划分查询行。每张 GPU 为自己负责的行针对完整上下文打分,直接选出全局 top-k,无需合并各 rank 的候选。如果能够减少通信且工作区预算允许,主注意力还可以单独收集 KV,随后为本地注意力头处理全部查询行。

综合这些 DCP 预填充优化后,开发阶段基准中的吞吐与延迟都有改善。

输入/输出 token 数
总吞吐(token/s)
输出吞吐(token/s)
TTFT 中位数(s)
TPOT 中位数(ms)
8K / 1K
5,357.8 → 6,015.4(+12%)
595.3 → 668.4
2.462 → 1.875
24.571 → 22.063
128K / 100
6,950 → 17,678(+154%)
5.30 → 13.48
268.7 → 105.5
331.5 → 129.6

配置:MI355X,TP8/DCP8,并发 16,64 条输入,关闭前缀缓存。

FlyDSL 稀疏 MLA 直接从 FP8 KV 缓存读取选中的 token。每个工作组针对本地注意力头处理一行查询,从而利用预填充步骤中的大量查询行提供并行性。较小的批次仍保留 Triton 路径,以便在每行内部提供更多并行性。

实现
基准运行时长(s)
TTFT 中位数(s)
输出吞吐(token/s)
Triton
432.8
96.2
14.79
FlyDSL
404.4
89.8
15.83
改善
降低 6.6%
降低 6.6%
提高 7.0%

配置:MI355X,TP8/DCP8,并发 16,64 条输入,输入/输出序列长度为 128K/100。

解码与融合式低延迟路径

对于较小的投机验证批次,DSA 索引和稀疏注意力分别运行;之后由一个持久化单体内核,在一次启动中融合注意力输出投影、AttnRes 和完整的潜在 MoE 计算,包括张量并行 all-reduce。这减少了内核启动开销,并让 MoE 权重加载与前面的归约、残差处理重叠执行。该专用路径每步最多支持 8 个验证 token,包括单个请求搭配 7 个草稿 token 的情况。

指标
常规路径
专用解码路径
改善
平均 TPOT
6.02 ms
4.40 ms
27.0%

配置:TP8/DCP1,FP8 索引器,输入/输出为 8K/1K,并发 1,7 个草稿 token,合成接受长度 3.84。

高效的跨 GPU 通信

解码期间,各 GPU 交换候选分数,但把 token ID 留在本地。共享阈值与一致的同分处理方式,决定各 GPU 选中哪些本地候选。对于 2048 个候选,只交换分数,可以把每查询每 GPU 的候选载荷从 16 KiB 减半至 8 KiB。

AITER 通过专用 GPU 内核加速 DCP all-gather:借助 HIP IPC 读取其他 GPU 的缓冲区,并直接写入所需输出布局,将数据收集与布局转换合并,省去单独的重排拷贝。在 8 张 MI355X GPU 上的独立微基准中,对于所测试的查询与候选张量形状,AITER 路径相较 RCCL 将 all-gather 延迟降低了 53%—61%。

TPU

vLLM 通过 vllm-torchtpu,在发布首日便支持 Kimi K3.1 Pro 在 TPU 上运行,并已在由 16 颗 TPU v7x 芯片组成的切片上验证。布局采用 TP32,MoE 使用 EP32;对于长上下文,再通过 DCP 按 token 切分 MLA 潜变量和索引键。DSA 的工作分配给 TPU 的两类计算单元:TensorCore 为 FP8 索引键打分,SparseCore 随后选出最高分的 2048 个 token,并把它们的潜变量收集到稠密缓冲区,再由 TensorCore 的 MXU 执行注意力计算。与 GPU 上一样,DSA 与 KDA 状态共享按块索引的 KV 池,使前缀缓存能够在混合缓存上正常工作。

阿里平头哥

阿里平头哥团队通过 vllm-sail 插件,增加了在真武 PPU 上运行 Kimi K3.1 Pro 的支持,涵盖 M890 与 810E。已验证的配置如下:

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

这是真武 PPU 首次为主流开源大模型提供首日支持,后续还将持续跟进。

致谢

我们衷心感谢 Moonshot 团队开源 Kimi K3.1 Pro,也感谢 Inferact、AMD 和阿里平头哥为 vLLM 提供模型首日支持,以及 NVIDIA 的帮助。

参考链接

  • 原文:https://vllm.ai/blog/2026-11-04-kimi-k31-pro
  • DeepSeek-V3.2 适配详解:https://vllm.ai/blog/2025-09-29-deepseek-v3-2
  • DeepSeek V4 适配详解:https://vllm.ai/blog/2026-04-24-deepseek-v4
  • GLM-5.3 上下文并行详解:https://vllm.ai/blog/2026-10-23-glm53-part2-pcp-dcp-hisparse
  • 解码上下文并行:https://vllm.ai/blog/2026-08-07-decode-context-parallelism
  • 在 vLLM 中深度优化 Kimi K3:https://vllm.ai/blog/2026-10-16-kimi-k3-megakernel
  • Inferact/Kimi-K3.1-Pro-DSpark-preview:https://huggingface.co/Inferact/Kimi-K3.1-Pro-DSpark-preview
  • verl:https://github.com/volcengine/verl
  • Inferact/zero-mismatch-rl:https://github.com/Inferact/zero-mismatch-rl
  • vllm-sail:https://github.com/vllm-project/vllm-sail

【声明】内容源于网络
0
0
AINLP
一个有趣有AI的自然语言处理公众号:关注AI、NLP、大模型LLM、机器学习、推荐系统、计算广告等相关技术。公众号可直接对话双语聊天机器人,尝试对对联、作诗机、藏头诗生成器、自动写作等,查询相似词,测试NLP相关工具包。
内容 6107
粉丝 0
AINLP 一个有趣有AI的自然语言处理公众号:关注AI、NLP、大模型LLM、机器学习、推荐系统、计算广告等相关技术。公众号可直接对话双语聊天机器人,尝试对对联、作诗机、藏头诗生成器、自动写作等,查询相似词,测试NLP相关工具包。
总阅读39.6k
粉丝0
内容6.1k