vLLM Kimi-K3 DSpark 推測解碼實作

概述

vLLM 已成功為 2.8T 參數的前沿模型 Kimi K3 訓練並部署了 DSpark 推測器。透過利用 Speculators 訓練函式庫與 GB300 NVL72 硬體,此實作將數學推理工作負載的單流互動性從約 110 提升至 435 tok/s/user,並在併發負載下提供高達 3.5 倍的輸出吞吐量。

DSpark 演算法

DSpark 是 DFlash 區塊級推測解碼演算法的擴展,旨在解決「後綴衰減 (suffix decay)」問題,即在平行預測的區塊中,單一錯誤會導致剩餘的 token 均失效。DSpark 維持了 DFlash 的平行骨幹架構,但引入了兩個特定組件以提升 token 間的連貫性:

  • Markov logit-bias head: 此組件會依序採樣 token,並使用低秩轉換矩陣根據先前選擇的 token 來調整 logits,在不需要額外 transformer pass 的情況下恢復局部依賴性。
  • Confidence head: 此組件會估計 token 被接受的機率,允許硬體感知調度器在低負載時驗證較長的前綴,並在高系統負載時修剪不太可能的後綴。

與 DFlash 相比,DSpark 在 Qwen3 目標模型上報告了 16–18% 更長的接受序列,且比 EAGLE-3 的序列長度高出 27–31%。

推論性能與能力

Kimi K3 DSpark 推測器使用一個五層、五十億參數的草稿模型 (draft model),在每次解碼步驟中提出 8 個 token。

基準測試與吞吐量

在九個評估領域中,該模型在每次驗證輪次中實現了 4.11 個 token 的宏觀平均接受長度。性能在結構化任務中表現最高:

  • Mathematical reasoning: 6.42 tokens
  • HumanEval: 4.96 tokens
  • Translation: 4.65 tokens

長文本性能

在 LongBench-v2 資料集上,該推測器在 378K-token 的提示詞 (prompt) 下,每次解碼迭代中達到高達 5.31 個輸出 token。前 10% 的請求維持了至少 3.76 tokens 每迭代,顯示出推測解碼在極端上下文長度下依然有效。

併發與延遲

隨著併發數從 1 增加到 16,總體輸出吞吐量從每秒 177 個 token 增加到 683 個 token。儘管併發請求增加了 16 倍,中位數首字延遲 (TTFT) 保持穩定,僅增加了 100 毫秒 (從 379 增加到 479 ms)。

硬體與訓練基礎設施

GB300 NVL72 配置

訓練是在 GB300 機架上進行的,使用 Ubuntu 24.04.4 LTS 與 NVIDIA 的 64K-page kernel 6.14。環境使用了 NVIDIA 的 610.57.04 開源核心 GPU 驅動程式 (R610) 與 CUDA 13.4.0 Developer Preview,這是第一個包含 Rubin 支援 (sm_107) 的工具包。

Mooncake Hidden-State 傳輸

由於草稿模型通常需要目標模型的隱藏狀態 (hidden states) 來對齊預測,vLLM 實作了 MooncakeHiddenStatesConnector。此系統實現了分離式訓練與隱藏狀態提取,這對於像 Kimi K3 (2.8T 參數) 這樣即使在 4-bit 量化下也超過單節點配置 VRAM 限制的模型而言是必要的。

  • Mechanism: 一個主 Mooncake proxy 處理程序會管理 vLLM 與訓練實例之間的通訊。訓練程序透過 vLLM 前端請求隱藏狀態,接收一個 Mooncake store key,接著由 Mooncake master 透過 RDMA 或 TCP 進行傳輸。
  • Topology: 為 Kimi K3 找到的最佳配置涉及三節點一組:兩個節點專用於 vLLM 推論,一個節點專用於訓練。

部署

Kimi K3 DSpark 推測器可透過 Hugging Face (RedHatAI/Kimi-K3-speculator.dspark) 取得,並可使用 vLLM 搭配以下推測配置進行部署:

{
  "model": "RedHatAI/Kimi-K3-speculator.dspark",
  "num_speculative_tokens": 8,
  "num_speculative_tokens": 8,
  "method": "dspark",
  "draft_sample_method": "probabilistic",
  "rejection_sample_method": "block"
}

Sources