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

Kimi K3.1 Pro 是 Moonshot AI 最强大的开放权重模型：一个拥有 2.8T 参数（104B 活跃参数）的多模态 MoE，具有 1M 令牌上下文，在编码、代理和知识基准测试中与最强的闭源模型相匹敌。

主要变化是注意力机制：K3.1 Pro 为所有 24 个全注意力层添加了稀疏注意力 (DSA) 以提高效率。vLLM 在 NVIDIA GPU、AMD GPU、Google TPUs 和阿里云 T-Head 神武 PPU 上从第一天就支持它。

简而言之

• 混合模型上的稀疏注意力： K3.1 Pro 保留了 K3 的 KDA + MLA 混合结构，并为所有 24 个 MLA 层添加了 DSA 索引器。在长上下文中，解码步骤读取每层大约四分之一字节，而密集 MLA 读取。
• 一种 KV 缓存用于三种状态： vLLM 将每个块的 MLA 潜在值、索引器键和 KDA 状态打包在一起。打包的缓存与 FlashInfer 和 FlashMLA 稀疏内核一起运行，使用 FP8 或 NVFP4 KV 和可选的 MXFP4 索引键，以及 P/D 解聚合将其不变地传输。
• 令牌分片的预填充注意力： 在 8× GB300 上，与头分片预填充相比，第一个令牌的时间 (TTFT) 快 1.69 倍，令牌数为 340K。
• 带有稀疏索引器的解码上下文并行性 (DCP)： 精确的 top-k 合并使用 ~10 倍更少的流量 使注意力模块在 DCP8 中 快 1.3–1.5 倍。DCP 预填充注意力每层运行速度提高 4.1–4.3 倍，DCP8 持有 TP8 的 7.2 倍 KV 令牌。
• 融合的低延迟解码： 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。
• 零不匹配的 RL 从第一天开始： 首次从第一天开始支持零不匹配 RL：vLLM 和 Megatron 对基础模型和 LoRA 适配器都产生位相同的对数概率，因此 PPO KL 在 100 步完整模型运行的每一步都精确为 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 注意力。索引器为每个令牌保留一个 FP8 键（包括其缩放在内的 132 字节），并使用 32 个轻量级头为每个查询评分每个先前的令牌。然后它保留前 2,048 个令牌。MLA 注意力然后精确运行，但仅限于这 2,048 个令牌。DSA 在 DeepSeek-V3.2 中引入；我们的 DeepSeek-V3.2 文章 更详细地介绍了索引器。

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

图 1：Kimi K3.1 Pro 中的稀疏注意力 (DSA)。索引器为每个过去的令牌评分，保留前 2,048 个，并且每个查询头仅关注这些令牌的潜在值。逐步查看。

KV 缓存占用

索引键是每个 token 的第二个缓存，仅次于 MLA 潜在值。在张量并行（TP）下，两者都在每个 GPU 上复制；KDA 状态是每个请求的，与 K3 一样。

潜在值和索引键可以独立量化。仅 NVFP4 潜在值就足以使 K3.1 Pro 低于 K3 的 FP8。选择 MXFP4 索引键可以多容纳 11% 的 token，并在长上下文中使索引器速度提升 10–20%。两者结合，八 GB300 GPU 持有的 token 数是 FP8 的 1.7 倍，而 DCP8 将其乘以七：约可同时处理 36 个完整的 1M-token 上下文，而 TP8 的 FP8 只能处理 3 个。

图 2：在 8× GB300 上测量的 KV 缓存容量。更改潜在值和索引键格式、并行性或上下文长度，以查看可以容纳多少个完整上下文。

用于三种状态类型的 KV 缓存

每个 DSA 层存储两个宽度不同的 per-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 个头填充到它们的最小瓦片大小。
• 索引器根本没有被按头分割：每个 GPU 评分每个 token，因此它的二次成本被重复了八次。

对于较大的预填充步骤，vLLM 按照token对注意力进行分片。所有对所有的操作将每个GPU的查询从“所有token，12个头”移动到“1/8的token，所有96个头”。每个GPU运行它自己的token的索引器和稀疏MLA，然后第二个所有对所有的操作将输出返回到头部布局。由于KV缓存已经在TP下复制，所以其他地方没有变化。MoE、序列并行性和缓存布局保持不变。为了平衡因果注意力，每个请求被切成16块，GPU r 拿到 r 和 15 − r 这两块，这与预填充上下文并行性相同的平衡分割。两个所有对所有的操作都直接写入对等方的对称内存。

图4：四个GPU上的token分片预填充，前后对比。按头分割，每个GPU重复索引器并填充它的头瓦片；在一次所有对所有的操作后，每个GPU只用每个头对token的一个平衡切片进行评分。逐步查看。

在8× GB300上使用TP8、真实权重和一次一个请求的结果：

在长缓存的prefix上的短追加，在代理转向中很常见，也能受益。对于在150K-token缓存的prefix上的2K-token追加，索引器和稀疏MLA内核每个GPU大约需要15 ms而不是56 ms（8× GB300，FlashInfer稀疏MLA）。代价是交换缓冲区，每个GPU大约需要0.8 GiB的KV缓存。

长上下文：使用稀疏索引器进行解码上下文并行性

TP在每个GPU上复制MLA潜空间，所以增加GPU不会增加KV容量。解码上下文并行性 (DCP) 按token对缓存进行分片：GPU r of W 存储每个 W-th token。对于密集MLA，每个GPU关注它的分片，部分输出通过它们的log-sum-exp (LSE)组合。

DSA增加了一步：每个GPU只能对其自己的1/W的token进行排序，所以全局前2048必须在每一层在注意力开始之前跨GPU合并。

一个精确的前k合并，流量减少10倍

通用合并将每个GPU的完整本地前2048，作为FP32分数和token id对，发送到其他每个GPU。在8× GB300上的DCP8，有64个类似AgentX的请求，每个请求有5个查询token，注意力模块需要258 µs的671 µs。

vLLM发送的流量要少得多。每个GPU只从BF16分数中挑选其最佳1.5 × 2048 / W候选者，并将每个候选者与其分数打包成32位，大约减少10倍的字节。然后每个GPU从收集的候选者中通过基数选择相同的全局前2048。合并仍然是精确的：如果一个GPU持有的前2048比它发送的多，选择会检测到这一点，并且那个GPU会重新发送该行的完整本地前k。

图5：在8× GB300上使用5个查询token每个请求的DCP8的一个DSA注意力模块。精确的单轮合并比通用合并快1.27–1.51倍。悬停在一个条形上以查看其值。

从批处理32开始，以及在128K上下文中，K3.1 Pro的DCP8下的稀疏注意力现在比K3的密集MLA快。

部分注意力输出然后被交换并通过对称内存上的两个CuTe DSL内核用LSE组合：比现有的CUDA实现快3.7倍（在8× GB300上的320 BF16形状），输出位相同。这有助于DCP下的每个MLA模型。

没有合并的DCP预填充

一个普通的 DCP 预填充会对每个提示 token 支付合并和 LSE 组合，并在每个 GPU 上重复注意力工作。当 DCP 组是 TP 组时，vLLM 反而会从所有 GPU 收集步骤的上下文，包括 MLA 潜在值和索引键，到一个对称内存工作区（每个 GPU 大约 0.7 GB，每 1M tokens，使用 FP8）。然后它在该工作区上运行上一节的 token 分片预填充。不再有合并或组合，输出与没有 DCP 的 token 分片预填充逐位相同。

端到端，单个请求在 DCP8 下的预填充速度与在 TP8 下一样快：在 27K、108K 和 340K tokens 下分别为 1.10、4.40 和 14.61 秒，与 1.10、4.39 和 14.50 秒（8× GB300，FP8 KV）相比。

何时使用 DCP

DCP 用每步延迟换取 KV 容量：DCP8 将每个 GPU 的 MLA 和索引键缓存减少八倍，而 KDA 状态已经按头进行分割。当上下文适合时，普通的 TP 在低并发下解码更快；当 KV 是限制因素时，DCP 会带来收益，例如非常长的上下文或许多长时间存在的 agentic 会话。在 8× GB300 上，DCP8 持有 21.8M tokens 的 FP8 KV，而 TP8 为 3.04M（使用 NVFP4 潜在值为 32.1M 对 4.49M）。在并发 1 时，TPOT 在 32K 上下文（9.45 对 11.58 ms）下使用 TP8 为 9.35 ms，使用 DCP8 为 11.48 ms；融合解码路径将 DCP8 提升到 5.96 和 6.04 ms。

8 个 GPU 在 TP8 下也服务完整的 1M-token 请求：在 8× GB300 上使用 FP8 KV，缓存持有 3.04M tokens，2.9 个并发 1M-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 稀疏；在 DCP 下批量 1–8，一个为形状构建的 CuTe DSL 内核（每层 11.9 µs 对比批量 1 的 21.4 µs）；
• 值投影和 sigmoid 门在一个内核中，然后是 K3 的输出投影及其融合的全归约。

在 DCP 下，查询发布到每个 GPU，每个 GPU 的 top-k 候选者精确交换和合并，部分输出通过 LSE 结合，所有这些通过内核内的 NVLink 邮箱完成，路径上没有 NCCL 集体操作。

图 6：在 8× GB300 上使用 TP8 和 1K-token 输入的每个输出 token 的时间。融合解码路径在批量 1 时快 1.63 倍，在批量 64 时仍然快 1.18 倍。悬停条形图可查看其值。

使用 DSpark 推测解码（7 个草稿标记），融合路径将并发 1 下的每个标记时间从 2.86 毫秒降低到 2.14 毫秒。GSM8K 准确率保持不变。在 DCP8 下，融合路径将并发 1 下的 TPOT 从 11.48 毫秒降低到 5.96 毫秒（32K 上下文）和从 11.58 毫秒降低到 6.04 毫秒（128K）。

使用 DSpark 进行推测解码

我们为 K3.1 Pro 训练了一个 DSpark 推测器，并发布了一个预览检查点：Inferact/Kimi-K3.1-Pro-DSpark-preview。使用真实采样（温度 1.0，top-p 0.95），预览版本在 GSM8K 上使用 7 个草稿标记已经达到了 4.597 的平均接受长度。在 8× GB300 上的 TP8 和批大小为 1 的情况下，这是 3.0 倍的解码加速。该推测器还保持了其在 K3.1 Pro 的完整上下文窗口（高达 1M 标记）中的准确率，这是在长上下文 BEAM 基准测试中测量的。

这个检查点只是一个预览。我们目前正在训练一个更准确、更高效的推测器，并很快会发布它。

自主服务：AgentX

vLLM 在自主服务工作负载上表现出色。我们在 NVIDIA GB300 和 NVIDIA Rubin 上使用 SemiAnalysis AgentX 基准测试来评估我们的初始图像。

NVIDIA GB300

在延迟敏感的前沿，Kimi K3.1 Pro 每用户保持 375.9 个标记每秒。沿着前沿进一步，它在每个 GPU 上服务 19.2K 总标记每秒的同时，保持每个用户在 100 个标记每秒以上。在吞吐量端，它达到每个 GPU 37.8K 标记每秒，大约是最交互点的 26 倍，而每个用户仍然看到大约 54 个标记每秒。

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

NVIDIA Rubin

在 NVIDIA Rubin 上，Kimi K3.1 Pro 在延迟敏感的前沿每用户保持 815.6 个标记每秒。沿着前沿进一步，它在每个 GPU 上服务 46.0K 总标记每秒的同时，保持每个用户在 100 个标记每秒以上。在吞吐量端，它达到每个 GPU 68.7K 标记每秒，大约是最交互点的 51 倍，而每个用户仍然看到大约 75 个标记每秒。

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

零不匹配的初始 RL

首次，vLLM 的初始支持将 RL 带来了零不匹配：使用 verl 的 LoRA RL，vLLM 作为部署引擎，Megatron 作为训练器，其中训练器的日志概率与部署的比特完全一致。vLLM 不是通过重要性采样或路由重放来缩小差距，而是在内核级别消除了它：训练器的正向传递以比特为单位重现了部署，这将 RL 训练稳定性提升到了一个新的水平：

• 在整个模型上的 LoRA RL。 排名 32 的 LoRA 覆盖了所有注意力投影和所有 896 个专家，并且基础模型和适配器在 vLLM 和 Megatron 之间都是比特完全一致的。
• 批大小不变的内核，原生支持。 使用 VLLM_BATCH_INVARIANT=1 一个标记的日志概率不依赖于其批大小，并且 vLLM 的越来越多的高性能内核在设计上是批大小不变的，例如 DeepSelect top-k。

图 9：相同的标记，相同的内核，相同的日志概率。切换到默认部署进行比较。

完整模型超过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 个用于注意力的标记之前扫描上下文。MXFP4 将每个索引键从 FP8 中的 132 字节减少到 68 字节，包括缩放，将索引缓存存储减少了 48.5%。 评分使用原生 FP4 矩阵指令，并使用 FP32 累加和跨头的加权缩减。在推测验证期间，键在查询行之间重复使用以减少重复读取。

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

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

在预填充期间，GPU 收集索引键并将查询行分配给自己。每个 GPU 对其分配的行与完整上下文进行评分，并直接选择全局 top-k，避免了每个排名候选的合并。主要注意力可以分别收集 KV，当这减少通信并适合工作空间预算时，然后处理所有查询行用于本地头。

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

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

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

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

解码：融合低延迟解码

对于小的推测验证批次，DSA 索引和稀疏注意力单独运行，然后是一个持久的单核，它融合了注意力输出投影、AttnRes 和完整的潜在 MoE 计算，包括张量并行全归约，到一个单独的启动中。这减少了内核启动开销，并将 MoE 权重加载与前面的归约和残差处理重叠。专用路径支持每步最多八个验证标记，包括一个包含七个草稿标记的单个请求。

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

通信：高效的跨 GPU 通信

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

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

TPU

vLLM通过vllm-torchtpu在TPU上支持Kimi K3.1 Pro，并在TPU v7x的16芯片切片上验证。布局为TP32，MoE使用EP32，DCP将MLA潜力和索引键按标记分片以支持长上下文。DSA的工作被分配到TPU的两种计算单元：TensorCore评分FP8索引键，然后SparseCore选择前2,048个标记并将它们的潜力收集到一个密集缓冲区，供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的帮助。