跳到正文
vLLM Blog· Helen Zhao, Fynn Schmitt-Ulms, Yuchen Fama, Antonio J. Dominguez, and Kevin Li·· 20 天前精选AI 评分62

vLLM 如何用 GB300 NVL72 训练 Kimi K3 的 DSpark 投机解码模型

How we trained the fastest DSpark for Kimi-K3 using GB300 NVL72

AI 导读

vLLM 团队基于 Speculators 训练库为 2.8T 参数的 Kimi K3 训练出 DSpark 投机解码草稿模型,在数学推理上把单流交互速度从约 110 提升到约 435 tok/s/user,并发负载下同等交互性时输出吞吐最高提升约 3.5 倍。

推荐理由

vLLM 团队公开了在 GB300 上训练 Kimi K3 投机解码草稿模型的完整方案,含跨节点隐藏状态传输与部署命令,可迁移到其他大模型。

正文 · AI 翻译

今年 6 月,DeepSeek 发布了 DSpark,这是对 DFlash 块级投机解码算法的扩展。新算法承诺更强的 token 间连贯性,从而带来更好的接受长度。但对开源社区而言,真正的问题始终如一:你能否训练它、打包它并部署它,而无需一名博士生守着检查点?

得益于 vLLM 项目的 Speculators 训练库,答案是响亮的“可以”!借助 Speculators,你可以轻松地以标准的、兼容 Hugging Face 的格式训练、打包和部署 DSpark 草稿模型,vLLM 可直接加载。该实现已在 Qwen3.6-35B-A3B、Gemma-4-31B-it、GLM-5.2 等模型上得到验证。本博客文章将探讨我们如何扩展训练库以支持 Kimi K3——一个 2.8T 参数的前沿模型。我们新的 DSpark 投机器在数学推理任务上将单流交互性从约 110 提升至约 435 tok/s/user,同时在并发负载下、在相同交互性水平下提供高达 约 3.5 倍的输出吞吐量提升。

Kimi K3 output throughput versus interactivity with and without the DSpark speculator
Kimi K3 在使用与不使用 DSpark 投机器时的输出吞吐量与交互性对比

图 1. Kimi K3 DSpark 投机器在数学推理工作负载上同时提升了单流交互性和总输出吞吐量。

什么是 DSpark?

大语言模型每次前向传播生成一个 token。投机解码通过让一个轻量级草稿模型提出多个 token,再由完整的目标模型一并验证,从而加速这一过程。

EAGLE-3 是一个强大的基线,但它仍然以自回归方式起草:提出七个 token 需要七个顺序起草步骤。DFlash 则相反,通过一次非因果主干前向传播预测整个块,报告称其加速比 EAGLE-3 高达 2.5 倍。

代价是并行位置之间无法相互条件化。对于提示词“Thank you!”,模型可能在第一个位置独立地同时偏好“Of”和“No”两种回复,在下一个位置同时偏好“course”和“problem”,从而产生诸如“Of problem”这样的不匹配。由于验证在第一个被拒绝的 token 处停止,一个错误也会使剩余后缀失效——DSpark 论文将这一问题称为后缀衰减。

DSpark 保留了 DFlash 的并行主干,同时增加了两个轻量级组件:

  • 马尔可夫 logit 偏置头按顺序采样 token,并使用先前选定的 token 调整每个位置的 logits。其低秩转移矩阵无需另一次 transformer 前向传播即可恢复重要的局部依赖关系。
  • 置信度头估计每个 token 被接受的概率。硬件感知调度器利用这些估计,在轻负载下验证更长的前缀,并在系统繁忙时裁剪不太可能的后缀。

因此,DSpark 通过继承单次主干前向传播,保留了并行起草的主要优势,同时恢复了自回归生成的部分连贯性。在 Qwen3 目标模型上,它报告称其接受序列比 DFlash 长 16–18%,比 EAGLE-3 长 27–31%。在 DeepSeek-V4 生产服务中,它在相同吞吐量下将每用户生成速度较之前的 MTP-1 基线提升了 60–85%。

DSpark architecture combining parallel drafting, sequential correction, confidence scoring, and hardware-aware verification
DSpark 架构,融合了并行起草、顺序校正、置信度评分与硬件感知验证

图 2. DSpark 在目标模型验证之前,将并行块与顺序校正及硬件感知前缀调度相结合。

推理期间的性能

DSpark 的半自回归设计只有在额外的顺序工作仍远低于另一次草稿模型前向传播的成本时,才能提升推理性能。发布的 Kimi K3 DSpark 推测器 使用一个五层、五十亿参数的草稿模型,并在每个解码步骤中提出八个 token。

在九个评估领域中,它实现了每轮验证平均 4.11 个 token 的宏平均接受长度。在结构化任务上表现最强:数学推理为 6.42 个 token,HumanEval 为 4.96,翻译为 4.65。在低请求数下,加速效果尤为显著。

该模型尤其适用于长上下文提示。在 LongBench-v2 这类具有挑战性的领域特定数据集上,我们的 Kimi K3 DSpark 在 378K token 的提示上,每个解码迭代最多可输出 5.31 个 token。即使在更广泛的工作负载中,排名前 10% 的请求每次迭代也至少达到 3.76 个 token,这表明在真正的长上下文长度下,深度推测运行仍然可行。

随着同时到达的请求增多,它也能有效扩展。将并发数从 1 增加到 16,可将总输出吞吐量从每秒 177 个 token 提升至 683 个 token。

至关重要的是,在负载下响应启动仍然很快。尽管同时服务的并发请求增加了 16 倍,首 token 中位时间仅增加 100 毫秒,从 379 毫秒增至 479 毫秒。

部署很简单。请根据你的硬件和用例,遵循官方 vLLM 配方:

docker run --gpus all \
  --privileged --ipc=host -p 8000:8000 \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  -e GLOO_SOCKET_IFNAME=$IFACE_NAME \
  -e NCCL_SOCKET_IFNAME=$IFACE_NAME \
  -e VLLM_ALLREDUCE_USE_FLASHINFER=1 \
  -e VLLM_ENGINE_READY_TIMEOUT_S=3600 \
  -e VLLM_USE_V2_MODEL_RUNNER=1 \
  -e VLLM_USE_RUST_FRONTEND=1 \
  vllm/vllm-openai:latest moonshotai/Kimi-K3 \
  --trust-remote-code \
  --gpu-memory-utilization 0.95 \
  --tensor-parallel-size 16 \
  --nnodes 2 \
  --node-rank 0 \
  --master-addr $HEAD_IP \
  --load-format fastsafetensors \
  --no-enable-flashinfer-autotune \
  --max-model-len 1048576 \
  --kv-cache-dtype fp8 \
  --attention-config '{"use_prefill_query_quantization":true,"mla_prefill_backend":"TOKENSPEED_MLA"}' \
  --enable-prefix-caching \
  --attention-backend TOKENSPEED_MLA \
  --prefix-match-unit 128 \
  --reasoning-parser kimi_k3 \
  --language-model-only \
  --speculative-config '{"model":"RedHatAI/Kimi-K3-speculator.dspark", "num_speculative_tokens":8, "method":"dspark", "draft_sample_method":"probabilistic", "rejection_sample_method":"block"}'

硬件设置

该模型在 Verda 慷慨提供的 GB300 机架上训练。Verda 是一家欧洲 AI 云服务商,端到端拥有从数据中心到托管服务的整个技术栈,并运行在 100% 可再生能源上。Verda 是欧洲首批部署 GB300 NVL72 机架的提供商之一。其内部 AI 实验室也在同一批机架上进行推理研究。

GB300 机架保留 NVIDIA 参考设计,并以裸金属方式运行,不使用虚拟化,在 NVIDIA 的 64K 页内核 6.14 上使用 Ubuntu 24.04.4 LTS。

Verda 将优化工作集中在最重要的系统级配置选择上,为研究人员提供前沿级环境。该系统使用 NVIDIA 的 610.57.04 开源内核 GPU 驱动(R610)。CUDA 13.4.0 开发者预览版与 CUDA 13.1 和 12.9 一同安装,而 NCCL 2.31.2 可在系统范围内使用。

选择 CUDA 13.4.0 是因为其新的 Blackwell 编程特性,并且它是首个包含 Rubin 支持的工具包(sm_107)。这有助于确保今天在 GB300 上编写和分析的代码,能够为下一代硬件构建,而不会出现意外的兼容性问题。基于硬件计数器的 GPU 性能分析也无需 root 权限即可使用,让每位研究人员都能对训练和推理工作负载进行分析。

我们一直与 Verda 合作,帮助开源训练和推理基础设施扩展到最新的 GPU 架构,目前重点聚焦于涵盖 GB300 到 VR200 的机架级系统。我们很高兴看到这一合作的成果正在成形。

推测解码草稿模型尽管体积小,却如此强大的部分原因在于,它们通常以目标模型的隐藏状态作为输入来辅助自身的预测。这极大地改善了它们的工作上下文,使草稿模型能够将其预测与目标模型紧密对齐。

需要注意的是,训练这些草稿模型需要一个包含隐藏状态输入和目标模型对数概率输出的数据集。幸运的是,vLLM 拥有一个隐藏状态提取系统,可以按需为数据集样本获取目标模型的内部隐藏状态。该系统使用一个虚拟草稿模型,复用 vLLM 的草稿模型管道来接收目标隐藏状态,并将其插入虚拟注意力层的 KV 缓存中。随后,一个实现 KVConnector 接口的类可以从 vLLM 中检索并传输这些隐藏状态。目前 vLLM 附带了一个 ExampleHiddenStatesConnector 正是做这件事的,它会将隐藏状态异步写入磁盘。

该系统在单节点 vLLM 与训练配置下运行良好。例如,节点上一半的 GPU 可以专用于训练,另一半则在 vLLM 中服务目标模型并按需提取隐藏状态。该系统已在 Speculators 和 vLLM 中得到数月的良好支持。然而,对于像 Kimi K3 这样拥有 2.8T 参数的模型,即使模型权重采用 4 位量化,最先进的加速器也开始遇到显存限制。我们需要一个能够超越单节点训练、实现训练与隐藏状态提取解耦的系统。

Mooncake hidden-state transfer between Speculators training and vLLM inference processes
Speculators 训练与 vLLM 推理进程之间的 Mooncake 隐藏状态传输

图 3. Mooncake 连接器将控制路径与隐藏状态数据路径分离,使用 RDMA 或 TCP 在 vLLM 与 Speculators 数据加载器之间传输。

基于这些需求,我们构建了 MooncakeHiddenStatesConnector,它使用 Mooncake 传输引擎作为后端,在进程之间和跨节点流式传输隐藏状态。新系统使用一个主 Mooncake 代理进程来管理与 vLLM 以及将自身注册为客户端的训练实例之间的通信。设置完成后,训练进程可以向 vLLM 前端发送请求,并收到一个 Mooncake 存储键作为响应。该键随后被提供给 Mooncake 主节点,由它代理从 vLLM 引擎到 Speculators 数据加载器的传输。所有这些都自动发生,由 Mooncake 服务器管理,并且根据配置,使用高速 RDMA 传输或常规 TCP 在同一节点或不同节点的进程之间发送数据。

训练 Kimi K3 DSpark

Twelve-node Kimi K3 DSpark training topology with disaggregated vLLM inference and Mooncake transfers
十二节点 Kimi K3 DSpark 训练拓扑,采用解耦的 vLLM 推理和 Mooncake 传输

图 4. 每个四节点组将一个四 GPU 节点专用于训练,两个四 GPU 节点用于 vLLM 推理,由 Mooncake 传输隐藏状态。

有了新的 MooncakeHiddenStatesConnector,我们现在拥有一个能够扩展到大型模型、多节点训练的系统。即使将 Kimi K3 量化为 4 位,该模型仍需要至少两个 GB300 节点(每个节点四块 GPU)来服务。我们尝试了不同的训练和 vLLM 配置,发现三节点组合——两个用于推理,一个用于训练——提供了最佳吞吐量。使用 Mooncake 连接器的 Speculators 的另一个优势是,可以轻松地独立扩展任一组件以获得最佳性能。

结论

Speculators 是一个开源库,用于构建、训练、评估和共享可直接与 vLLM 等推理引擎集成的投机解码模型。无论你是在训练新的草稿模型、探索 DFlash、DSpark 或 DFlash2 等算法,还是提升生产环境的推理性能,我们都欢迎你的想法和贡献。加入 vLLM Community Slack,在 #speculators 和 #feat-spec-decode 中找到我们,提出问题、分享结果、讨论新算法,并与社区协作。欢迎对代码、文档、示例、模型支持和评估工具做出贡献。

来源:vLLM Blog · vllm.ai