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

Kimi K3.1 Pro 是 Moonshot AI 能力最强的开放权重模型：一个 2.8T 参数（104B 激活）的多模态 MoE，拥有 1M token 上下文，在编程、智能体和知识基准上可与最强的闭源模型匹敌。

主要变化在注意力：K3.1 Pro 为全部 24 个全注意力层加入了稀疏注意力（DSA）以提升效率。vLLM 在第 0 天即支持它，覆盖 NVIDIA GPU、AMD GPU、Google TPU 和阿里巴巴平头哥 PPU。

TL;DR

• 混合模型上的稀疏注意力：K3.1 Pro 保留 K3 的 KDA + MLA 混合结构，并为全部 24 个 MLA 层加入 DSA 索引器。在长上下文下，一个解码步骤每层读取的字节数约为稠密 MLA 的四分之一。
• 一份 KV 缓存承载三种状态：vLLM 将每个块的 MLA latent、索引器键和 KDA 状态打包在一起。打包后的缓存使用 FlashInfer 和 FlashMLA 稀疏内核运行，支持 FP8 或 NVFP4 KV 以及可选的 MXFP4 索引键，P/D 分离可原样传输该缓存。
• Token 分片的预填充注意力：在 8× GB300 上、340K token 时，首 token 时间（TTFT）比头分片预填充快至多 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；在 NVIDIA Rubin 上，SemiAnalysis AgentX 基准中每用户最高 815.6 tok/s，每 GPU 总计 68.7K tok/s。
• 第 0 天零失配 RL：首次实现第 0 天支持即附带零失配 RL：vLLM 与 Megatron 对基础模型和 LoRA 适配器均产生逐位相同的 log-prob，因此在 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 的新变化

K3.1 Pro 中的稀疏注意力

每个稀疏注意力层在其 MLA 注意力旁都有一个小型索引器。索引器为每个 token 保留一个 FP8 键（含缩放共 132 字节），并用 32 个轻量头为每个查询对之前所有 token 打分。然后保留 top 2,048 个 token。MLA 注意力随后精确运行，但只在这 2,048 个 token 上运行。DSA 在 DeepSeek-V3.2 中引入；我们的DeepSeek-V3.2 文章更详细地介绍了索引器。

收益在长上下文解码中最大，而长上下文解码受内存带宽限制。稠密 MLA 每层每个上下文 token 读取 576 字节的 FP8 latent。DSA 仍会扫描整个上下文，但只读取 132 字节的索引键。然后它只为 2,048 个 token 读取完整 latent。在长上下文下，这约为稠密 MLA 每层读取字节数的 23%。

图 1：Kimi K3.1 Pro 中的稀疏注意力（DSA）。索引器为每个过去 token 打分，保留 top 2,048，每个查询头只精确关注这些 token 的 latent。逐步查看。

KV 缓存占用

索引键是 MLA 潜在表示旁边的第二个逐 token 缓存。在张量并行（TP）下，两者都会在每个 GPU 上复制；KDA 状态按请求保存，与 K3 相同。

潜在表示和索引键可以独立量化。仅使用 NVFP4 潜在表示，K3.1 Pro 就已经低于使用 FP8 的 K3。可选的 MXFP4 索引键在此基础上还能多容纳 11% 的 token，并使索引器在长上下文下快 10–20%。两者同时使用时，八块 GB300 GPU 可容纳的 token 数量是 FP8 的 1.7 倍，而 DCP8 再将其乘以七：一次约 36 个完整的 1M token 上下文，而 TP8 配 FP8 只有 3 个。

图 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 潜在表示。
• 分离式服务。每个块是一段连续内存，因此 P/D 传输只需一次大型 RDMA 读取即可移动它，而不是数十次小的逐层读取。
• 共享基础设施。同一个打包规划器也服务于 GLM-5.3-Flash 和 DeepSeek-V4，它们的缓存同样存在混合页大小的问题（见 DeepSeek V4 文章）。

图 3：TP8 下打包缓存的布局。注意力组和三个 KDA 组共享一个块行池；选择一个 token、一个 DSA 层和一个 KDA 层，即可查看各自落在哪个块和页中。

预填充：按 token 而非按头分片注意力

在 TP8 下，每块 GPU 拥有 K3.1 Pro 的 96 个注意力头中的 12 个。这对稀疏 MLA 预填充来说并不合适，原因与我们在 GLM-5.3 文章中所述相同：

• MLA 的潜在 KV 由所有头共享，因此每块 GPU 都要读取完整缓存却只服务 12 个头，而稀疏内核会把这 12 个头填充到其最小 tile。
• 索引器完全不按头拆分：每块 GPU 都为每个 token 打分，因此其二次方成本被重复了八次。

对于较大的 prefill 步骤，vLLM 改为按 token 对注意力进行分片。一次 all-to-all 将每个 GPU 的 query 从“全部 token，12 个头”变为“1/8 的 token，全部 96 个头”。每个 GPU 为自己的 token 运行 indexer 和稀疏 MLA，第二次 all-to-all 将输出恢复到按头布局。KV cache 在 TP 下已经复制，因此其他方面没有变化。MoE、序列并行和 cache 布局保持不变。为了平衡因果注意力，每个请求被切成 16 片，GPU r 取第 r 片和第 15 − r 片，与 prefill 上下文并行采用相同的均衡切分。两次 all-to-all 都直接写入对端的对称内存。

图 4：四块 GPU 上按 token 分片的 prefill，前后对比。按头切分时，每个 GPU 都重复运行 indexer 并对其头分块做填充；经过一次 all-to-all 后，每个 GPU 只用全部头对均衡切分后的一段 token 计算分数。逐步查看。

在 8× GB300、TP8、真实权重、一次一个请求下的结果：

在长缓存前缀上追加短内容（agentic 轮次中很常见）也能受益。在 150K token 的缓存前缀上追加 2K token 时，indexer 和稀疏 MLA 内核每 GPU 约耗时 15 ms，而非 56 ms（8× GB300，FlashInfer 稀疏 MLA）。代价是交换缓冲区，每 GPU 约 0.8 GiB 的 KV cache。

长上下文：带稀疏 indexer 的 decode 上下文并行

TP 会在每个 GPU 上复制 MLA latent，因此增加 GPU 并不会增加 KV 容量。Decode 上下文并行（DCP）改为按 token 对 cache 分片：W 个 GPU 中的 GPU r 存储每第 W 个 token。对于稠密 MLA，每个 GPU 在其分片上做注意力，部分输出通过 log-sum-exp（LSE）合并。

DSA 多了一步：每个 GPU 只能对自己那 1/W 的 token 排序，因此全局 top-2048 必须在每一层、注意力开始之前跨 GPU 合并。

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

通用合并会把每个 GPU 完整的本地 top-2048（FP32 分数与 token id 对）发送给其他所有 GPU。在 8× GB300 上以 DCP8 运行 64 个类 AgentX 请求、每个 5 个 query token 时，它占用了 671 µs 注意力模块中的 258 µs。

vLLM 发送的数据少得多。每个 GPU 直接从 BF16 分数中只挑选最好的 1.5 × 2048 / W 个候选，并将每个候选与其分数打包进 32 位，字节数约减少 10×。然后每个 GPU 从收集到的候选中基数选择出相同的全局 top-2048。合并保持精确：如果某个 GPU 持有的 top-2048 多于它发送的数量，选择过程会检测到，该 GPU 会为该行重新发送其完整的本地 top-k。

图 5：在 8× GB300 上以 DCP8 运行的一个 DSA 注意力模块，每个请求 5 个 query token。精确的单轮合并比通用合并快 1.27–1.51×。将鼠标悬停在柱上可查看其数值。

从 batch 32 起，且在 128K 上下文下，K3.1 Pro 在 DCP8 下的稀疏注意力现已快于 K3 的稠密 MLA。

随后，部分注意力输出通过对称内存上的两个 CuTe DSL 内核进行交换并用 LSE 合并：比现有 CUDA 实现快最多 3.7×（8× GB300 上的 320 个 BF16 形状），且输出逐位一致。这对 DCP 下的每个 MLA 模型都有帮助。

无需合并的 DCP prefill

普通的 DCP prefill 会为每个 prompt token 付出 merge 和 LSE combine 的代价，并在每块 GPU 上重复注意力计算。当 DCP 组就是 TP 组时，vLLM 改为把该 step 的上下文（包括 MLA latent 和 index keys）从所有 GPU 收集到一个对称内存工作区（FP8 下每 GPU 每 1M token 约 0.7 GB），然后在其上运行上一节所述的 token-sharded prefill。这样不再有 merge 或 combine，且输出与不带 DCP 的 token-sharded prefill 逐位一致。

端到端来看，单个请求在 DCP8 下的 prefill 速度与 TP8 相当：27K、108K 和 340K token 时分别为 1.10、4.40 和 14.61 s，对比 1.10、4.39 和 14.50 s（8× GB300，FP8 KV）。

何时使用 DCP

DCP 用每步延迟换取 KV 容量：DCP8 将每 GPU 的 MLA 和 index-key 缓存削减为八分之一，而 KDA 状态本就按 head 拆分。当上下文能装下时，普通 TP 在低并发下解码更快；当 KV 成为瓶颈时（如超长上下文或大量长生命周期的 agentic 会话），DCP 才划算。在 8× GB300 上，DCP8 可容纳 21.8M token 的 FP8 KV，而 TP8 为 3.04M（使用 NVFP4 latent 时为 32.1M 对 4.49M）。并发为 1 时，32K 上下文下 TP8 的 TPOT 为 9.35 ms，DCP8 为 11.48 ms（128K 时为 9.45 对 11.58 ms）；融合解码路径将 DCP8 降至 5.96 和 6.04 ms。

TP8 下的八块 GPU 也能服务完整的 1M token 请求：在 8× GB300 配 FP8 KV 时，缓存可容纳 3.04M token，即 2.9 个并发的 1M token 请求。

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

针对 Kimi K3，vLLM 提供了一条可选的解码路径，几乎用为小批量调优的 CuTe DSL kernel 替换了所有算子。它将通信融合进计算，重做了 MoE 调度，并把权重预取到 L2（见 Hyper Optimizing Kimi K3 in vLLM）。对于 K3.1 Pro，该路径现已覆盖 DSA 层。一个 DSA 层的执行流程为：

• 为 query、KV latent、output gate 以及 indexer 的 key 和 head 权重做一次打包的输入投影；
• 融合的 indexer 前端：query norm、indexer 的 query 投影、其 FP8 或 MXFP4 量化（并入 head 权重）以及新 token 的 index key，在一个 kernel 中完成；
• MLA query（up-projection 和 absorbed key projection，将新 token 的 latent 写入 KV 缓存），与 indexer 重叠执行而非在其之前运行；
• 融合的 indexer：一个 key-balanced scorer 和一个精确 top-k 选择，分两个 kernel 完成，输出注意力读取的物理 KV 行；
• 稀疏注意力：FlashInfer sparse MLA（trtllm-gen 或 PrimsTS，按 batch size 选择）或 FlashMLA sparse；在 DCP 下 batch 1–8 时，使用为该形状构建的 CuTe DSL kernel（batch 1 时每层 11.9 µs 对 21.4 µs）；
• value projection 和 sigmoid gate 在一个 kernel 中完成，随后是 K3 的 output projection 及其融合 all-reduce。

在 DCP 下，query 被发布到每块 GPU，各 GPU 的 top-k 候选被精确交换和合并，部分输出通过 LSE 合并，全部经由 kernel 内的 NVLink mailbox 完成，路径上没有 NCCL 集合通信。

图 6：8× GB300、TP8、1K token 输入下每个输出 token 的耗时。融合解码路径在 batch 1 时快 1.63×，在 batch 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。在真实采样（temperature 1.0，top-p 0.95）下，该预览版在 GSM8K 上使用 7 个草稿 token 时已达到 4.597 的平均接受长度。在 batch size 为 1、TP8 于 8× GB300 上时，解码加速达 3.0×。该投机器在 K3.1 Pro 最长 1M token 的完整上下文窗口内也保持准确率，该结果在长上下文 BEAM 基准上测得。

该检查点仅为预览版。我们正在训练一个更准确、更高效的投机器，并将很快发布。

Agentic 服务：AgentX

vLLM 在 agentic 服务负载上表现出色。我们在 NVIDIA GB300 和 NVIDIA Rubin 上使用 SemiAnalysis AgentX 基准测试了我们的 day-0 镜像。

NVIDIA GB300

在延迟敏感的前沿端，Kimi K3.1 Pro 保持每用户每秒 375.9 个 token。沿前沿进一步推进，它在每 GPU 每秒服务 19.2K 总 token 的同时，仍让每个用户保持每秒 100 个 token 以上。在吞吐端，它达到每 GPU 每秒 37.8K 个 token，约为最具交互性点的 26×，而每个用户仍能看到约每秒 54 个 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。沿前沿进一步推进，它在每 GPU 每秒服务 46.0K 总 token 的同时，仍让每个用户保持每秒 100 个 token 以上。在吞吐端，它达到每 GPU 每秒 68.7K 个 token，约为最具交互性点的 51×，而每个用户仍能看到约每秒 75 个 token。

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

day 0 零失配 RL

vLLM 的 day-0 支持首次实现了零失配的 RL：使用 verl 的 LoRA RL，vLLM 作为 rollout 引擎，Megatron 作为训练器，训练器的 log-prob 与 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-prob。切换到默认 rollout 以进行比较。

在完整模型上运行超过 100 步，每一步的 rollout/trainer KL 都恰好为 0.0，而 AIME mean@8 从 0.451 上升到 0.664（图 10）。

图 10：在完整模型上进行的 100 步 LoRA RL 运行。

批处理不变路径的性能优化仍在进行中。源代码位于 Inferact/zero-mismatch-rl。

AMD

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

量化：原生 MXFP4 索引器

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

更小的索引缓存为缓存上下文留出更多 GPU 内存，有利于高并发下的长上下文服务。在 AgentX 扫描中，MXFP4 索引器配置在并发 48 时实现了高出 5.3% 的总工作负载吞吐量。

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

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

组合的 DCP 预填充优化在开发基准测试中提升了吞吐量和延迟。

设置：MI355X，TP8/DCP8，并发 16，64 个提示，前缀缓存已禁用。

FlyDSL 稀疏 MLA 直接从 FP8 KV 缓存中读取选定的 token。每个工作组在其本地注意力头上处理一个查询行，从而在预填充步骤中的众多查询行之间提供并行性。较小的批次保留 Triton 路径，该路径在每一行内提供额外的并行性。

设置：MI355X，TP8/DCP8，并发 16，64 个提示，ISL/OSL 128K/100

解码：融合低延迟解码

对于小型推测验证批次，DSA 索引和稀疏注意力分别运行，随后是一个持久化单内核，将注意力输出投影、AttnRes 和完整的潜在 MoE 计算（包括张量并行 all-reduce）融合为单次启动。这减少了内核启动开销，并将 MoE 权重加载与之前的归约和残差处理重叠。专用路径每步最多支持八个验证 token，包括一个带有七个草稿 token 的单个请求。

设置：TP8/DCP1，FP8 索引器，8K/1K，c1，草稿 token 7，合成 AL 3.84

通信：高效的跨 GPU 通信

在解码过程中，每个 GPU 交换候选分数，同时将 token ID 保留在本地。共享阈值和一致的平局处理方式决定哪些本地候选被选中。对于 2,048 个候选，仅交换分数就能将每个 GPU 每次查询的候选负载从 16 KiB 减半至 8 KiB。

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

TPU

vLLM 通过 vllm-torchtpu 在 TPU 上于第 0 天支持 Kimi K3.1 Pro，并在 TPU v7x 的 16 芯片切片上完成验证。布局为 TP32，MoE 采用 EP32，DCP 按 token 对 MLA latent 和索引键进行分片以支持长上下文。DSA 的工作被拆分到 TPU 的两类计算单元上：TensorCore 对 FP8 索引键进行评分，随后 SparseCore 选出前 2,048 个 token 并将其 latent 收集到密集缓冲区中，供 TensorCore 的 MXU 执行注意力计算。与 GPU 上一样，DSA 和 KDA 状态共享一个按块索引的 KV 池，从而使前缀缓存在混合缓存中继续生效。

阿里巴巴 T-HEAD

阿里巴巴 T-HEAD 团队通过 vllm-sail 插件新增了对在阿里巴巴 T-HEAD 真武 PPU（M890 和 810E）上运行 Kimi K3.1 Pro 的支持。已验证的配置如下：

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

这标志着真武 PPU 首次为主流开源大模型提供第 0 天支持，持续支持即将推出。

致谢

我们衷心感谢 Moonshot 团队开源 Kimi K3.1 Pro。我们也感谢 Inferact、AMD 和阿里巴巴 T-HEAD 在 vLLM 中为该模型提供第 0 天支持，并感谢 NVIDIA 的帮助。