SGLangとMilesがKimi K3にDay-0サポートを提供

SGLangとMilesがKimi K3にDay-0サポートを提供

はじめに

SGLangチームとMilesは、Kimi K3へのDay-0サポートを正式に発表した。K3は初のオープンソース3兆パラメータ級モデルであり、そのハイブリッドアーキテクチャは従来のServingスタックの前提を複数の点で覆すものである。Moonshot AIおよびNVIDIAとの連携により、SGLangが推論を、MilesがRL学習を担当し、リリース日の全シナリオを完全にカバーする。

bs=1 optimization-category waterfall

主要ハイライト

  • 新ハイブリッドアーキテクチャ:2.8Tパラメータ、69層のKDA線形アテンションと24層のMLAを交互配置し、LatentMoEとAttention Residualsをサポート。
  • デュアル状態メモリ管理:プレフィックスキャッシュの再構築、オーバーラップスケジューリングおよびページングメカニズムを採用し、統一メモリプールによりサイジングの推測を排除。
  • カーネル段階的最適化:バッチサイズ1で約113 tok/sを達成。
  • DSpark投機的デコード:専用ドラフトモデルを導入し、バッチサイズ1のデコードスループットを約423 tok/sに向上。
  • フェーズ並列化:チャンク化パイプライン並列プリフィルとテンソル並列デコードを組み合わせ、GPU当たり2808 tok/sを実現。
  • MilesがネイティブMXFP4チェックポイントのLoRA RL学習をサポート。
bs=1 decode throughput ladder

アーキテクチャ適応の課題

K3は前世代のK2.5と大きく異なる:パラメータ規模が2.8Tに拡大し、アテンションにKDAとMLAのハイブリッドを採用、Attention ResidualとLatentMoEを導入、活性化関数をSiTUに変更、ビジョンタワーをMoonViT3dにアップグレードしている。これらの変更により、SGLangは複数のコアパスを再構築することが求められた。

The three KDA state moves, copy-on-write, snapshot and donate, and where each sits relative to the serial forward stream.

ハイブリッドKDAメモリ管理

KDA状態はインプレースで上書き可能であり、コピーオンライト・スナップショット・ドネートの3種類の明示的状態移動により安全な再利用を実現する。チェックポイントはRadixツリーの分岐点にのみ配置され、高いヒット率と競合なしを保証する。

Sparse KDA state checkpoints overlaid on the radix tree, and the branching point where a new request diverges from a cached prefix.

統一メモリプール

単一のメモリプールがKDA状態ブロックとMLA KVブロックの両方を収容し、両端から向かい合う形で増加する。解放後は自動的にフラグメントを整理し、ワークロードに応じた動的割り当てを実現する。

Unified memory: two static pools versus one pool holding both kinds of state, and the freeing sequence

DSpark投機的デコード

専用ドラフトモデルと信頼度スケジューリングトリムを組み合わせ、検証に値する位置のみを検証する。ReplaySSMによりKDA状態のメモリオーバーヘッドを約32分の1に削減する。

Decode throughput vs batch size, verify-all vs confidence-scheduled trim, on a chat panel and a few-shot math panel.KDA ReplaySSM: the verify kernel reads the committed checkpoint and stores each draft step

並列化とパイプライン最適化

プリフィルにはチャンク化パイプライン並列を、デコードにはテンソル並列を採用し、両者をPD分離(PD disaggregation)で組み合わせることで、スループットを大幅に向上させ通信オーバーヘッドを低減する。

Chunked pipeline-parallel prefill: a long prompt cut into chunks streaming through pipeline stages, with P2P hand-off arrows running under the next chunkLeft: prefill throughput per GPU vs concurrency, TEP8 vs PP8xTP1. Right: cost decomposition per 1k tokens, compute plus exposed communication, for TP8 / PP4xTP2 / TEP8 / PP8xTP1.TP replicates the same token positions on every rank; DCP stripes them round-robin, so the same GPUs hold N times the logical context.

まとめ

完全な起動コマンドと設定ガイドはKimi K3 cookbookを参照されたい。SGLangとMilesのDay-0サポートは、超大規模ハイブリッドアーキテクチャモデルのServingと学習における新たな基準を打ち立てるものである。