vLLM、Mooncakeで隠れ状態をノード間ストリーミングし、2.8TのKimi K3向けDSparkドラフトモデルを訓練

vLLMチームは、2.8兆パラメータのKimi K3向けDSpark speculative modelを訓練して公開し、大規模モデルを単一マシンに収容しにくいという問題を、独立してスケール可能な推論プレーンと訓練プレーンに分割した。Kimi K3の重みを4-bitに量子化しても、隠れ状態の抽出には少なくとも2つのGB300ノードが必要となる。そのためチームは、2つの4 GPUノードでターゲットモデルを実行し、別の4 GPUノードでドラフトモデルを訓練したうえで、Mooncakeを使用し、RDMAまたはTCP経由で中間状態とターゲットlogitsをストリーミングした。

公開されたドラフトモデルは5層、約50億パラメータで、1ラウンドにつき8 tokenを提案する。DSparkはEAGLE-3のようにドラフトをtoken単位で生成するのではなく、まずブロック全体を並列に予測し、その後、低ランクMarkov headによって隣接token間の条件関係を補完する。さらに別のconfidence headが、各位置のtokenがターゲットモデルに受理される確率を推定することで、スケジューラーは負荷に応じて検証プレフィックスを短縮または延長できる。9種類の評価における平均受理長は1ラウンド当たり4.11 tokenで、数学的推論では6.42、HumanEvalでは4.96だった。

チームが公開した数学ワークロードでは、単一ストリームの速度が約110 token/sから435 token/sへ向上した。並列数を1から16に増やすと、総スループットは177 token/sから683 token/sへ上昇する一方、time to first token（TTFT）の中央値は379ミリ秒から479ミリ秒への増加にとどまった。378K-tokenのLongBench-v2ケースでは、1ラウンド当たり5.31個の出力tokenを達成しており、長いコンテキストが必ずしもドラフトの受理率を急落させるわけではないことを示している。

エンジニアリング上の価値は、単一のcheckpointにとどまらない。Apache 2.0ライセンスのSpeculatorsには、DSpark、ノード間hidden-state connector、訓練コマンド、vLLMから直接読み込めるHugging Face形式が追加されている。ただし、性能値は同じチームがGB300 NVL72と入念に調整されたソフトウェアスタックを用いて得たものだ。プロンプト分布、サンプリング設定、旧世代GPU、ネットワークインターコネクトが異なれば、ドラフトモデルによる追加計算が性能上の利点を相殺する可能性がある。今後は、第三者が費用対効果を再現できるか、またMooncakeによる転送と5層のドラフトモデルがマルチテナント環境でのサービスでも引き続き採算に合うかが焦点となる。