LMSYS、Kimi K3モデルのSGLangとMilesへのDay-0サポートを追加
本文の状態
日本語全文を表示中
詳細モードで約41分の本文を読めます。
SGLang と Miles が Kimi K3 の Day-0 サポートを開始し、2.8T パラメータのハイブリッドアーキテクチャに対応した推論と RL 学習環境を提供する。
AI深層分析を開く2026年8月1日 01:34
AI深層分析
キーポイント
Kimi K3 の新アーキテクチャ対応
SGLang と Miles は、69 層の KDA リニアアテンションと 24 層の MLA を組み合わせた Kimi K3 のハイブリッド構造を完全にサポートする。
メモリ管理と推論性能の革新
再帰状態の書き込みやプレフィックスキャッシュに対応した新しいメモリ管理設計により、バッチサイズ 1 で約 113 トークン/秒の推論速度を達成する。
DSpark による推測的デコーディング
K3 用に訓練したドラフトモデルを用いた DSpark 推測的デコーディングにより、バッチサイズ 1 で約 423 トークン/秒の速度を記録する。
RL 学習における Miles の活用
Miles はネイティブ MXFP4 チェックポイント上で LoRA RL を実行し、12 時間のトレーニングで AIME-2024 のスコアを 43.3% から 76.7% に向上させた。
Kimi K3 のアーキテクチャ的特徴
K3 は 2.8T パラメータ、ハイブリッド注意機構(69 層の KDA リニア注意と 24 層の MLA)、および LatentMoE を採用している。
重要な引用
K3 is the first open-source model in the 3-trillion-parameter class, and its hybrid architecture departs from convention in almost every place a serving stack has assumptions.
A kernel ladder reaching ~113 tok/s at batch 1, before speculation.
DSpark speculative decoding... ReplaySSM handles the KDA state, replaying raw inputs instead of snapshotting per step, a roughly 32x cut in draft-window memory.
K3's hybrid KDA and MLA architecture creates two memory management challenges.
編集コメントを表示
編集コメント
Kimi K3 のような超大規模かつ複雑なアーキテクチャを持つモデルが、Day-0 で主要な推論・学習フレームワークでサポートされるのは極めて異例である。これはオープンソースコミュニティとハードウェアベンダーの緊密な連携が、最新モデルの実用化を加速させる好例と言える。
Source Article
元記事を日本語で読む
本文に関係しない購読案内、埋め込み通知、サイト内プロモーションは除いています。
SGLang と Miles で Kimi K3 の Day-0 サポートを開始しました
このたび、SGLang および Miles において、Kimi K3 の Day-0 サポートを提供できることを発表できることを嬉しく思います。K3 は 3 兆パラメータクラス初のオープンソースモデルであり、そのハイブリッドアーキテクチャは、推論スタックが通常持つ前提条件をほぼすべて覆すものとなっています。
Moonshot AI と NVIDIA のチームとの協力のもと、発売日には SGLang を推論用、Miles を RL 学習用に使い分け、K3 を完全にサポートします。本記事では、その実現に至った背景と技術的な取り組みについて解説します。
主なポイント
- 新しいハイブリッドアーキテクチャ:2.8T パラメータ、69 KDA(Linear Attention)層を 24 MLA、LatentMoE、Attention Residuals と交互に配置。推論スタックの多くの前提条件がここで崩れます。
- 2 種類の状態管理:自己上書きする再帰的状態を含むメモリ管理。プレフィックスキャッシュ、オーバーラップスケジューリング、ページングをその上に再構築し、最終的なサイズ見積もり不要な統合プール設計を採用しました。
- カーネルの最適化により、バッチサイズ 1 で推測なしでも約 113 tok/s のスループットを実現。
- DSpark によるスペキュレーションデコーディング:K3 用に訓練したドラフトモデルを使用。バッチサイズ 1 で約 423 tok/s を達成。ReplaySSM が KDA 状態を処理し、ステップごとのスナップショットではなく生入力を再生することで、ドラフトウィンドウのメモリ使用量を約 32 分の 1 に削減。
- フェーズ別の並列化:チャンク化パイプライン並列によるプレフィルとテンソル並列によるデコードを組み合わせ、PD 分離構成下で GPU あたり 2,808 tok/s を実現。
Miles を用いた LoRA RL を、ネイティブ MXFP4 チェックポイント上で実行し、トレーニングとロールアウトを同一 GPU 上に配置することで、12 時間の学習で AIME-2024 のスコアを 43.3% から 76.7% に引き上げました。
起動コマンドや各ワークロードの設定ガイドは、Kimi K3 クックブック で確認できます。
Kimi K3 の導入
Kimi K3 は、3 兆パラメータクラスのモデルとして初めてオープンソース化されたものです。2.8T パラメータと 100 万トークンのコンテキストウィンドウを備え、ネイティブな視覚理解機能を搭載しています。
先行する K2.5 は、SGLang がすでに良好にサポートしていたアーキテクチャの親戚のようなモデルだったため、既存のスタックをそのまま適用できました。しかし K3 はそうではありません。従来の常識から複数の独立した点で大きく逸脱しているのです。
| Kimi K2.5 | Kimi K3 | |
|---|---|---|
| スケール | ~1T パラメータ | 2.8T パラメータ |
| アテンション | all-MLA、61 レイヤー | ハイブリッド:69 の KDA リニアアテンション+24 の MLA、計 93 レイヤー |
| アテンションの残差 | なし | すべてのアテンション出力を保存し、ブロックごとに集約 |
| MoE | 384 エキスパート、トップ 8 | LatentMoE:896 エキスパート、トップ 16、3584 次元の潜在空間で動作 |
| エキスパート活性化関数 | SwiGLU | SiTU |
| ビジョンタワー | MoonViT | MoonViT3d(K3 専用の新スタック) |
各行は、それぞれがサービング上の課題を表しています。1M コンテキストを現実的なコストで実現可能にしているのはハイブリッド・アテンションスタックですが、これによりサーバーは同時に 2 種類の状態を保持することになります。つまり、MLA のトークンごとの KV(Key-Value)に加え、各リクエストごとに固定サイズの KDA 状態が並存するのです。以下のメモリアクセス管理セクションでは、まさにこの点に焦点を当てます。
アテンションの残差(Residuals)は、スタック全体を通じてアテンション出力のバンクを流すことで処理されますが、これは SGLang の標準的なレイヤー接続が前提とする仮定を崩し、DP アテンションなどの特定の箇所で K3 専用のパスが必要となりました。LatentMoE は、896 個のエクスパートのうち 16 個を、SiTU 活性化関数(SwiGLU の代わりに使用)を用いた下位投影された潜在空間内でルーティングします。そのため、既存の MoE カーネルをそのまま適用することはできません。さらにビジョンパス(タワー、プロジェクター、プロセッサ、そして Kimi の XTML メディアフォーマット)はゼロから構築されました。以下の各セクションは、この基盤構築の上に成り立っています。
ハイブリッド KDA メモリ管理
K3 のハイブリッド KDA と MLA アーキテクチャは、2 つのメモリアクセス管理上の課題を生み出しています。1 つ目は、可変的な KDA 再帰状態に対してプリフィックスキャッシングを安全かつ効率的に行うことです。2 つ目は、KDA 状態と MLA の KV を動的に容量を共有することです。
安全で効率的な状態の再利用
アテンションの KV は書き込み専用(append-only)です。一度計算されると変更されないため、スケジューラはラディックスツリー上でキャッシュされたプリフィックスを複数のリクエスト間で安全に共有できます。一方、KDA 状態は異なります。各レイヤーは固定サイズの再帰バッファを保持しており、トークンごとにインプレースで上書きされるため、プリフィックスキャッシングでは不変の KV ではなく、可変の状態を管理する必要があります。
SGLang では、ラジックスツリーとリクエストの作業スロットの間で、可変な KDA 状態を安全にキャッシュして再利用するために、3 つの明示的な状態移動を行います。
コピー・オン・ライト (Copy-on-write) は、フォワードパスが状態を変更する前に共有チェックポイントをプライベートなスロットへ復元します。スナップショット (Snapshot) は進化した状態をキャプチャし、ドネート (donate) はそのスナップショットをツリーへ転送します。
*3 つの状態移動。 コピー・オン・ライト、スナップショット、ドネート。それぞれがシリアルなフォワードストリーム上でどの位置に配置されているかを示します。
この設計は、競合状態を構造的に排除しています。復元とスナップショットのコピーは、状態を生成する操作と消費する操作の間にフォワードストリーム上にキューイングされるため、ストリームの順序性がインプレイース更新との競合を防ぎます。スナップショットは、ピンポンバッファの 2 つのスロット間で交互に実行されるため、ツリーへ付与されている状態を上書きすることはありません。ドネートではスロットインデックスのみが転送され、状態そのもののコピーは行われません。
2 つ目のスロットは境界点で遅延割り当てされ、直後に解放されます。これにより、リクエストごとに永続的なスロットを確保する必要がありません。この仕組みには、デバイス全体の同期やホットパスにおけるロックも不要です。
キャッシュのヒット率はチェックポイントの配置によって決定されます。再帰的な状態は逆方向に実行できず、より古いチェックポイントから順方向に再生する必要があります。私たちは、プリフィル時にラジックスツリーの整列されたノードおよびチャンク境界で、デコード時には固定トークン間隔でチェックポイントを作成し、パスごとの上限と LRU 方式に基づいてスパースなセットを維持します。
Marconi に着想を得て、私たちは分岐点を優先しています。なぜなら、これらの点のプレフィックスはすべての子ブランチで共有されるからです。リクエストがエッジ途中で分岐した場合、その直上の最も近いチェックポイントから再生し、整列されたフォーク地点に新しいチェックポイントを配置します。これにより、後のブランチはそこから直接復元できるようになります。
*ラジックスツリー上のチェックポイント。 スパースなチェックポイントのオーバーレイと分岐点。*
競合のない状態移動と分岐を考慮したチェックポイント配置により、高効率な KDA プレフィックスキャッシングが実現され、これは重なりスケジューリング、推測的デコーディング、ページ化 KV と組み合わせても問題なく動作します。
統一メモリ:両方の状態タイプに共通のプール
各状態タイプはそれぞれ独立したプールで管理されており、このプールの設計こそが最終的な解決策です。2 つの割り当て単位には 3 桁の違いがあります。1 つはリクエストごとの大規模な KDA 状態ブロック(TP=8 の場合約 54 MB で、69 レベルの KDA をすべてカバー)で、もう 1 つはトークンごとの小規模な MLA KV ブロック(約 27 KB で、24 レベルの MLA をすべてカバー)です。そのため現在ではこれらは起動時にサイズが設定された 2 つの別々のプールに存在しています。このサイズ設定はトラフィックへの賭けであり、その予測が外れると、あるプールでメモリ不足が発生する一方で、別のプールにはまだ空き容量が残るという事態になります。
統一メモリはこの 2 つのプールを 1 つに統合します。KDA 状態は一方から、MLA KV ブロックは他方から埋められ、その間に残った未使用バイトが単一の自由領域となります。
統一メモリ。上段:現在では 2 つの状態タイプは起動時にサイズ設定された別々のプールを持っており、片方が空しくなっている間に他方が満杯になることがあります。中段:統一メモリを使用すると両方の状態タイプが同じプールから割り当てられ、反対側から成長して中央に単一の自由領域を残します。下段:中央で状態ブロックを解放すると隙間が生じますが、端からのブロックをその隙間に移動させることで、自由領域は一つのまとまりとして維持されます。
解放処理も同様にシンプルです。リクエストが完了したり、中止されたり、圧力下で取り消されたりした場合、そのブロックは解放されます。もしこれが中央に隙間を作ってしまう場合、端からのブロックをその隙間に移動させることで、自由領域は一つのまとまりとして維持されます。
このレイアウトにより、柔軟なページサイズの実現とメモリの断片化防止が可能になります。54 MB の KDA 状態ブロックと 27 KB の MLA KV ブロックは、共通のページサイズを強制されることなく同じバイト領域から割り当てられます。前述の操作により両端が詰まった状態で維持され、状態移動にかかるコストは無視できるほど小さいため、空き領域は常にどちらの種類にも利用可能な連続した領域として保たれます。したがって、容量は起動時のフラグではなくワークロードに応じて動的に決定されます。短いリクエストが多い場合は状態ブロックでプールが埋まり、長いコンテキストが多い場合は KV で埋まりますが、いずれの場合も再構成の必要はありません。
統一メモリ機能は、--enable-unified-memory オプションを指定した場合のみ有効になります(オプトイン方式)。実装の詳細については、後日別の記事で解説する予定です。
DSpark を活用した予測デコーディング
K3 には、DSpark ブロックによる予測デコーディング機能が標準搭載されています。これは K3 のために訓練したドラフトモデル に基づいています。この統合において特に注目すべきは、検証の予算を効果的な箇所に集中させることと、予測処理においても K3 の再帰型 KDA 状態を維持できる点です。
検証する価値がある箇所のみを検証する
DSpark は、1 ステップごとにドラフトトークンのブロックを提案し、ターゲットモデルはそれを 1 回の順伝播で検証します。バッチサイズが 1 の場合、追加の検証位置にかかるコストは実質的にゼロです。この段階ではレイテンシがボトルネックとなっており、わずかなトークンが増えたところで影響はありません。しかし、バッチが埋まってくると状況は一変します。
検証トークンは、そのステップで他のすべてのリクエストと競合するようになります。しかも、多くの場合、それらのトークンは検証に失敗してしまいます。チャットワークロードでは平均的な受容長さが約 2.7 であるため、典型的なステップで 8 つの検証位置のうち 5 つが拒否されることになります。すべてのトークンを検証しても、サーバー側で捨てられるトークンに対しては完全なコストを支払うことになり、非効率です。
この問題を解決するための要素は、すでにシステム内に存在していました。ドラフトモデルには訓練済みの信頼度ヘッド(confidence head)が含まれており、各位置のトークンが検証を通過する確率を予測します。一方、サーバー側では、負荷レベルごとに追加の検証トークンが実際にどの程度のコストがかかるかを記録したプロファイルを用意しています。
これら 2 つの要素を結合し、ステップごとのプランナーが最適化を行います。各リクエストは、期待される価値が限界コストを上回る間だけ検証トークンを保持し、それ以外のウィンドウ部分はターゲットモデルによる順伝播の実行前に切り捨てられます。残されたトークンは従来通り厳密に検証されるため、出力の損失はありません。このトレードオフにより、わずかに受容長さが短くなる代わりに、1 ステップあたりのコストを下げることができます。
負荷下では「トリミング」が有利に働きます。チャットパネル(許容トークン数約 2.7、左)と Few-shot 数学パネル(許容トークン数約 5.0、右)におけるデコードスループットを比較すると、「全検証(verify-all)」対「トリミング」の結果が示されています。バッチサイズ 8 までは両者の差はほとんどありませんが、それを超えると開きが生じ、バッチサイズ 256 ではそれぞれ +68% と +24% の向上が見られました。この間、許容トークン数は 2.7 から 2.2、5.0 から 4.3 に低下しつつも、スループットは上昇し続けています。
測定されたコスト曲線こそが最も興味深い点でした。検証トークンの限界コストは滑らかではなく、階段状になっています。あるトークンが現在のカーネル波にほぼ無負荷で収まる「平坦な棚」と、新しいカーネルを開始する際に急激にコストが跳ね上がる「急勾配」が交互に現れます。プランナーはこのコスト曲面を読み解きます。リクエストの次のトークンが安価な棚に乗っていれば保持し、急勾配に乗りそうであればそこで切り捨てます。その結果、許容トーク数曲線には小さな揺らぎが見られますが、スループットは滑らかで単調に推移します。これはプランナーがノイズではなく、実際のハードウェアの特性(棚)を反映していることを示しています。
バッチサイズ 8 を下回る領域では負荷軽減の圧力がほとんどないため、トリミングの効果は限定的か、わずかにマイナスになることもあります。この領域での早期終了機能の実装は、今後のプランナーにおける既知の課題となっています。
ReplaySSM: KDA 状態のための生入力リプレイ
スペキュレーティブ・ディコーディングでは γ+1 個のドラフトトークンを一度に検証しますが、受け入れられるのはプレフィックスのみとなるため、状態を巻き戻せる必要があります。KDA レイヤーの状態はトークンごとに上書きされるため、従来手法では各ドラフトステップ完了後に K×V の全状態をスナップショットとして保存します。K=V=128 の場合、レイヤーとヘッドごとにリクエストあたり 64 KB を必要とし、これを γ+1 ステップ分確保する必要があります。K3 では 69 層の KDA レイヤーが存在するため、実行中のリクエストごとにこの一時領域を確保する必要があり、恒久的な状態プールと競合して並行処理数の上限を制限していました。
ReplaySSM では、各ドラフトステップの生入力 Sᵢ = (vᵢ, kᵢ, gkᵢ, βᵢ) を約 1 KB のサイズで保持します。これは検証カーネルが通過する際に書き込まれます。サンプラーが受け入れられるトークン長を確定すると、全レイヤーとヘッドを対象とした 1 つのフォールド・カーネルがコミット済みチェックポイントから受け入れられたプレフィックスのみを読み込み、その場で状態を進化させます。このフォールド処理は検証再帰の完全な複製であり、ゲート値を再計算するのではなく、検証カーネルが保存した値を利用します。そのため、再構築された状態は、再帰ベースラインがコミットしていたものとビット単位で完全に一致します。
- メモリ: ドラフトウィンドウのサイズが 512 KB から 16 KB に削減され、約 32 倍の改善となります。このメモリを状態プールに返却することで、並行処理数の上限が数倍向上します。
- スピード: 検証カーネルはステップごとにスナップショットを書き込まなくなるためメモリアクセスが削減され、フォールド処理は全レイヤーとヘッドに対して 1 つの融合された起動で実行されるため、バッチサイズが小さくてもステップあたりの処理時間が改善します。
KDA ReplaySSM は、チェックポイント h₀を読み取る際に書き込みを行わず、各レイヤーとヘッドごとに個別のバッファにドラフトステップごとの生データを保存します。その後、fold カーネルは受け入れられたプレフィックスのみを再生し、そのスロットをその場で上書きします。
カーネル
2.8T 規模のハイブリッドモデルにおける単一シーケンスのデコードステップは、計算量にボトルネックがあるわけではありません。むしろ、カーネル起動回数とレイテンシが課題となります。トークンあたりには 93 層のアテンション(KDA が 69 層、MLA が 24 層)と 92 層の潜在変数 MoE が存在し、初期段階では数百もの微小なカーネルが起動していました。この最適化キャンペーンはプロファイリング主導で進められました。「一つ機能を融合させ、固定されたプロトコル上で A/B テストを行い、GSM8K の結果を閾値として採用するか判断する」というサイクルを繰り返したのです。
起動とコピーの排除(P1–P4、+19.9 tok/s)。カーネル自体を高速化するのではなく、ステップそのものを小さくします。MoE の前段処理を一つの GEMM に集約し、KDA のペアになった細い射影も統合します(P1)。プロファイルに基づく探索により、レイヤーごとのアップキャストやコピー、不要な起動が排除されます(P2, P3)。ルーティングは [M, 896] のロジットに対してレジスタ内で完結するラジックス選択としてワンパスで実行されます(P4)。
NVIDIA 計算カーネル (P5–P8、+10.3 tok/s)
デフォルト設定がバッチサイズ 1 の形状に不適切だったため、4 つのステージで NVIDIA と共同開発または提供されたカーネルを切り替えています。具体的には、融合 KDA デコードカーネル、ルーティングバイパス背後にある trtllm-gen W4A8 SiTU MoE の cubin、TMA アテンション残差集約、そして小規模な M において cuBLAS を置き換える CuTe-DSL TGV bf16 GEMM です。これらすべては、他のコンポーネントと同様に A/B テストと精度ゲートを通じて導入されました。
通信融合 (P9, P12, P13、+27.6 tok/s)
最も大きな改善をもたらすのが、CustomAllReduceV2 の対称メモリ平面における MNNVL ファブリック用の融合された all-reduce ファミリーです。小規模なメッセージにはワンショットマルチキャストストアを、大規模なものにはスイッチ内 NVLS による集約を使用します。残差加算と RMSNorm は、この集合演算(P9)の内部で実行されます。MoE の最終処理は次に集合演算のステージングパスへ移行し(P12)、up-projection は複製されたものから列並列 GEMM へと変更され、マルチキャスト all-gather で完了します(P13)。
オーバーラップとプロローグ融合 (P10, P11, P14, P15、+10.4 tok/s)
残りの部分は、ボトルネックとなるクリティカルチェーンの削減に注力しています。ストライド入力 MXFP8 量子化(P10)、独立したサイドストリーム上で上流カーネルの末尾へ融合された残差書き戻し(P11)、KDA の GEMV チェーンと qkvg GEMM のオーバーラップ(P14)、そして MLA デコードのプロローグを 1 つのカーネルに融合し、その背後にあるアテンションカーネルを PDL が準備する(P15)といった最適化です。
上記の 4 つの改善項目に対応する時系列の進捗は以下の通りです。
「一般化可能な教訓」として、All-reduce は同期ポイントであるため、そこで節約された 1 マイクロ秒はステップ時間の 1 マイクロ秒に直接換算されます。一方、別のストリームのオーバーラップの隙間に配置されたカーネルが節約する時間は、およそ 1/10 の効しかありません。このキャンペーンにおいて最も効果の高い習慣は、カーネルを実装する前にそのカーネルがクリティカルパスに含まれるかをトレースで確認することでした。
¹ リベースにより P5 ベースラインは 64.2 から 63.8 に変更されました。ステップ数はリベース後のベースラインに対して測定されます。
² P13 はマージウィンドウ後の再較正された標準ベースラインであり、単一の PR の寄与を示すものではありません。
³ P1〜P4 に続くセグメントには、後に置き換えられた中間最適化も含まれています(例:All-reduce の結合→P9、Tiny/1-CTA GEMV→P8、Radix Router v1→P4、Marlin top-k-sum→P6/P12、早期アテンション加算→P7)。これらの一時的な効果は曲線には現れていますが、特定の名称を持つポイントはありません。
K3 の並列化
K3 のハイブリッドアーキテクチャでは、従来の並列化手法の選択肢が通用しません。Tensor Parallelism(テンソル並列)では MLA の KV キャッシュをシャードできません(KV ヘッドは 1 つで分割対象がないため)、すべてのランクが完全なコピーを保持する必要があります。また、すべての GEMM を 8 分割して処理し、レイヤーごとに集約通信のコストが発生します。
一方、Pure DP Attention(純粋なデータ並列アテンション)では、アテンション重みをすべてのランクに複製します。これには KDA に約 61 GB、MLA に 11 GB のメモリが必要で、KV キャッシュと KDA 状態の保存にもメモリを消費します。
これらのコストにより、Prefill(事前処理)と Decode(デコード)は異なる課題に直面するため、K3 はフェーズごとに回答を分割して対応しています。具体的には、Prefill では Chunked Pipeline Parallelism(チャンク化パイプライン並列)、Decode には Context Parallelism(コンテキスト並列)を採用します。
Prefill: チャンク化パイプライン並列
TP(Tensor Parallelism)の事前計算フェーズでは、各レイヤーの最後に AllReduce 演算が発生し、これが計算と並行して実行できないバリアとなります。一方、パイプライン並列化はモデルをレイヤー単位で分割します。K3 の場合、93 レイヤーが 8 つのステージに分割され、プロンプトもチャンクに切り分けられてこれらを通ってストリーミングされます。
チャンク化されたパイプライン並列事前計算。各ステージは異なるチャンクを同時に処理します。ステージ間の引き継ぎは、次のチャンクの計算中に行われるため、K3 ではその 91% が隠蔽されます。一方、TP(下段の帯)では、すべてのランクが待機する AllReduce が各レイヤーの最後に発生します。
このアプローチには三つの利点があります。残された通信はステージ間の引き継ぎのみとなり、これは次のチャンクの計算に隠れることができます。各ランクは完全なレイヤーを実行するため、GEMM(行列乗算)演算が 8 倍幅広くなり、効率が向上します。また、各ステージが保持する KV キャッシュと活性化値は約 12 レイヤー分に限定されるため、非常に長いプロンプトの事前計算も快適に行えます。ただし、パイプラインはある程度深くする必要があります。浅い PP4×TP2 の構成では引き継ぎをカバーできず、かつ TP2 の AllReduce コストを負うことになるため、TEP8 と比べてベンチマーク性能は向上しません。
2 台の GB300(計 8 GPU)における 8K トークンの事前計算を測定した結果です。トポロジーのみを変数としています。
Deep PP は両方の軸で勝利します。左のグラフは、c1〜c4 の交差点を過ぎた後の性能を示しており、PP8×TP1 は TEP8 の上限を約 1.7 倍上回り、TTFT(Time to First Token)も低くなります。右のグラフは、各ランクでの FLOPs が等しい場合の 1,000 トークンあたりのプレフィルコストを示しています。PP4×TP2 と TEP8 は計算資源と通信コストをトレードオフして互角ですが、PP8 は両方の指標で最も安価です。
PP8 の欠点は単一リクエスト時のみ発生します。また、アイドル状態のプレフィルワーカーは設定ミスであるため問題外です。K3 の分離型サービングでは、プレフィルノードが PP8 で動作します。各ノードのプレフィル容量は TEP8 ノードの 1.45〜1.72 倍あり、1 つのプレフィルノードで複数のデコードノードを賄うことができます。一方、デコードノードでは TP(Tensor Parallelism)または DCP(Decode Context Parallelism)が採用され、これについては次節で詳しく解説します。
デコード:コンテキスト並列化
デコード領域では、複製された KV キャッシュがボトルネックとなります。TP の場合、リクエストを 1 つ追加するか、あるいはコンテキスト長を 1,000 トークン増やすと、すべてのランクで同じ量のバイト数が必要になります。一方、Decode Context Parallelism(DCP)は、MLA の KV をヘッド単位ではなくトークンの位置に基づいてシャードします。具体的には、ランク r が位置 p を所有するのは「p mod N == r」の場合です。これにより、各ランクはすべてのリクエストのコンテキストを 1/N に分割してインターリーブした形で保持し、アテンションカーネルの上層ではこのシャード化が隠蔽されます。
*位置ベースのシャード化。 4 つのランクにまたがる同じ 16 トークン位置において、TP は 64 の物理コピーを保持するのに対し、DCP は 16 のみで済みます。解放されたバイト数が論理的な KV キャパシティとなり、K3 では DCP8 を使用すると約 7.9 倍の容量増を実現します。
位置シャード化はソフトマックスを破綻させます。各ランクがキーの 1/N しか見ず、部分ソフトマックスを加算できないためです。
解決策は FlashAttention の手法と同じです。各ランクは部分アテンション出力とヘッドごとの log-sum-exp を返します。その後、層ごとに一度だけ all-to-all コミュニケーションを行い、各ランクが全コンテキストにわたる 1/N のヘッドを受け取ります。log-sum-exp による局所的なマージを行うことで正確な結果が得られ、その出力は TP(Tensor Parallelism)の下で出力投影が期待するヘッドレイアウトそのものになります。
通信コストは層あたり 1 回の集合演算のみです。
***デコードステップ(層ごと)。各ランクは全ヘッドのクエリをローカルで投影し、自分が担当する位置のみに対してアテンションを実行します。その後、部分出力とその log-sum-exp を単一のパケットにまとめて all-to-all で転送します。ローカルでのマージにより標準的な TP ヘッドレイアウトが得られ、アテンション以降の処理には一切変更がありません。
その他の部分はすべてそのままです。DCP グループは TP グループ内部に構築されるため、TP8 に DCP8 を組み合わせても GPU 数は 8 台のままです。MoE(Mixture of Experts)も従来通り並列化を実行します。
KDA はこのルールを説明する例外です。KDA の状態はトークンごとではなくリクエストごとに固定サイズの行列として保持されるため、シャード化する位置軸が存在しません。そのため KDA 層はヘッド単位で TP シャード化されたまま維持されます。
この機能は単一のフラグ --dcp-size N で有効化できます。
この技術がもたらす恩恵は、エージェントによるトラフィックにおいて顕著に現れます。ここでは数十万トークン単位のセッションがキャッシュに蓄積されます。
2×4 GB300環境で実際のコーディングエージェントのセッションを再生し、ホストメモリ KV 階層を両アームに配置した上で、DCP を唯一の違いとして比較しました。
DCP はアクティブセットの壁を取り払います。 ホスト階層は再プリフィル(re-prefill)トラフィックを救済しますが、並行セッション数が 16 に達すると、アクティブな作業セットが TP8 のデバイス KV を超え、スループットが急落します。一方、DCP8 は論理 KV の容量を 150 万トークンから 1220 万トークンに引き上げ、48 セッションという負荷下でも 541 トークン/秒の性能を維持します。
DCP はスタック全体と調和して動作します。DSpark の検証ステップはデコードステップとして機能するため、同じ複製された Q と one-all-to-all パスを利用します。PD(Prefill-Decode)分離構成では、プリフィル側は DCP を意識する必要がなく、各デコードランクは転送境界で自身が担当する位置情報を引き出すだけで済みます。これにより、PP や TP によるプリフィル処理が DCP デコードにスムーズに連携できます。
残るボトルネックは KDA です。KDA のリクエストごとの状態情報は位置シャード化できないため、DCP が MLA(Multi-Head Latent Attention)の壁を突破した時点で、実行中のリクエスト数の上限が新たな制約条件となります。上記のメモリセクションで説明した統一メモリ設計が、この課題に対応します。
すべての戦略を組み合わせた場合
これらの要素がどのように組み合わさるかは、最終的には計測によって決まります。PD 分離をループに組み込み、プリフィルトポロジ、デコードトポロジ、およびプリフィル対デコードの比率を変化させることで、サービス性能のフロンティアを明らかにできます。
*Serving の最前線*
スループット性能の観点では、FP4 アーム上で 1 つの PP8 プリフェッチワーカーが 1 つの TP8 デコードノードを駆動する構成で、GPU あたり 2,808 トークン/秒を達成しています。一方、DCP 構成を採用し、2 つの PP8 プリフェッチワーカーが 2 つの DCP8 デコードノードを制御する構成では、わずかに下回る 2,633 トークン/秒となりました。
次に注目すべきは、プリフェッチとデコードの比率を調整する「knob」です。1 つのプリフェッチワーカーが 2 つ、3 つ、そして 4 つの独立したデコードインスタンスを制御する構成では、総スループットは低下しますが、ユーザーあたりの速度は劇的に向上し、1 ユーザーあたり 116 トークン/秒を超える性能を発揮します。
RL: ネイティブ MXFP4 ベース上での LoRA 学習
K3 における Day-0 の強化学習(RL)は、Miles 環境上で LoRA 学習をネイティブに実行する構成です。これは、Miles の Megatron バックエンド上で動作する BF16 トレーナーと、同じ 64 GB H100 を共有するネイティブパッキング対応の MXFP4 SGLang ランタイムエンジンによって実現されています。
このバックエンドは、KDA(Key-Value Attention)、NoPE-MLA、アテンション残差バンク、そして潜在空間 MoE(Mixture of Experts)を包括的にサポートしており、TP/SP/PP/CP/EP のすべての並列化戦略に対応しています。
LoRA 推論と重み同期
エンジンが読み込むチェックポイントは、そのままの状態で提供され、一切書き換えられません。各ステップでは BF16 形式の LoRA アダプターのみが転送され、エンジンはこれを量子化されたベース GEMM の上に別個の BF16 B(Ax) 項として適用します。これにより、ポリシー更新は 4 ビットベース重みに対してフル精度で推論に反映されます。
密な投影層は SGLang の Triton LoRA バックエンドを経由し、896 個のルーティング専門家には Marlin パス上の統合 MoE-LoRA カーネルが用いられます。また、共有専門家のデルタ値は、統合 MoE フロント GEMM に折り込まれます。
アダプターは GPU メモリプールに格納されており、各同期処理でその場所ですべて書き換えられます。そのため、ロールアウト側には BF16 形式のコピーが常駐せず、重みの完全な再同期やループ内の量子化ステップも不要となります。
パラレル処理
パイプライン並列処理。 K3 のアテンション残差スナップショットバンクはステージ境界をまたぐ必要がありますが、Megatron のポイントツーポイント通信では 1 つの隠れ状態テンソルしか伝送されません。そのため、ステージ境界では [prefix_sum, bank] をこのテンソルにパックし、エントリ時にアンパックします。Megatron 自体は変更されていません。
コンテキスト並列性について。アテンションと残差演算がトークン単位で実行されるため、銀行(バッチ)のシャードは自動的に解放されます。MLA においても K3 に特化したコードは不要です。Megatron の MLA で CP(コンテキスト並列)対応が必要な作業はロータリーテーブルのスライスだけですが、K3 はロータリー埋め込みを持たないため、その射影演算は標準の TE アテンションコアに置かれ、CP を継承します。
問題となるのは KDA です。再帰状態と畳み込みのハロ部を処理する際、fla の CP コンテキストを経由する必要があります。KDA は連続したランクローカルのチャンクを要求しますが、Megatron が保存しているのは環状アテンションが期待するジグザグ順序です。このため、再配置(relayout)処理は KDA 周辺のみで実行されます。
エキスパート並列性について。アダプターは潜空間側の因子を、ルーティングされるすべての 896 エキスパート間で共有し、他のパラメータはエキスパートごとに保持します。これはエンジンの融合 MoE-LoRA の契約と一致しています。この共有因子は EP(エキスパート並列) across で複製されますが、エキスパート並列としてタグ付けされているため、DDP による縮小はエキスパート DP の範囲のみで行われます。リリースで独自に追加される勾配縮小は、EP 和の計算のみです。
メモリ
ネイティブ MXFP4 の展開では GPU あたり最大約 225 GiB に達し、BF16 トレーナーは初期化時点で約 155 GiB です(277 GiB のカード上)。つまり、両者は同時にメモリ上に存在することはなく、一方が実行されている間にもう一方は待機状態になります。共存が可能なのは、すべてのハンドオフで残滓を残さない場合に限られます。
ベースモデルの重みは不変です。エンジンは実行中ずっとベース重みを GPU 内に保持し、トレーナーが動作している間は KV キャッシュと CUDA グラフのみを解放します。ベースは一度も解放されないため、復元されることもありません。LoRA パスではベースのバイト自体が転送されないため、ベース同期は完全にスキップされます。
プロセスグループは稼働状態を維持します。トレーナーの NCCL コミュニケータバッファは torch-allocator のメモリではないため、オフローディングしても解放されません。そのため、明確な解決策として、スリープ時にグループを破棄し、ウェイクアップ時に再構築することがあります。これは、常駐する数 GiB を、サイクルごとの再構築と EP(Expert Parallelism)のウォームアップに置き換えるトレードオフです。エンジンがこれらのバイトを再び必要とする場合のみ、このコストを支払う価値があります。しかし今回は不要です。なぜなら、ベースウェイトは更新パスを通過しないからです。したがって、プロセスグループは稼働状態を維持します。
アダプターの転送はチャンクごとに制限されます。アダプターは約 2,800 のテンソルで構成されています。各チャンクは単一の平坦化された CUDA IPC バケットとして送信され、合計 278 個あります。これは、受信側が確認し、すべてのプロデューサーランクがエンジングループのバリアを通過した時点で収集されます。これにより、転送全体にわたってメモリが統合されるのを防ぎ、チャンクごとの一時的な IPC メモリを制限できます。ピーク時には GPU あたり約 48 GiB の節約になります。
ホストへのコピーは、実際にそれを読み取る場所でのみ存在します。エンジンは、アダプターが GPU プールにインストールされると、その CPU コピーを解放します。これにより、スケジューラの RSS(Resident Set Size)は 76〜88 GiB から約 17 GiB に減少します。CPU コピーが存在しないアダプターを再インストールすることはできないため、プールからの退去はエラーとして拒否され、古くなったスロットが黙って提供されることはありません。トレーナー側では、DDP バッファはライフタイムによって分割されます。アダプターのパラメータバッファは、トレーナーがスリープしている間にウェイト更新で読み込まれるため CPU ベック領域に留まります。一方、勾配バッファは再構築可能であり、スリープ時に破棄されるバックアップなしの領域へ配置されます。
検証
強化学習(RL)のサポートスタックにおいて、トレーニングの正しさが最も重要です。私たちはレシピ作業を行う前に、これらのチェックを構築しました。
Train/rollout KL は、ロールアウトごとにログされる指標で、サンプリングされたトークン全体にわたる KL(rollout ‖ train) をシュルマン K3 推定値として計算します。これはエンジンが返した対数尤度と、トレーナーが再計算した値から導き出されます。これは目的関数の一部ではなく診断ツールであり、KL ペナルティは課されません。その下限は約 2e-3 で、これはベースモデルを MXFP4 でサービス提供することで設定されています。ステップ数に伴うこの値の増加こそが、発散の兆候です。
A canary lockstep probe(カニエール・ロックステッププローブ) は、最初のロールアウトから 1 つの軌道を選択し固定します。その後、各ステップで生アダプターの下にあるエンジンと、同じポリシーバージョンにおけるトレーナーの両方がスコアを付けます。2 つの曲線の連動こそが、学習されたアダプターが推論に到達したことを証明するものです。転送チェックサムは、バイトが正しく到着したことを確認するものであり、中身が読まれているかどうかまでは保証しません。
Tensor-level dump comparison(テンレベルダンプ比較) では、同じトークンを 2 つのビルド、あるいは 2 つの並列レイアウトで実行し、すべての順伝播活性化値とすべてのパラメータ勾配をダンプします。その後、測定されたノイズフロアに対して相対 L2 ノルムとコサイン類似度で比較を行います。ビット完全性(bit-exactness)が基準となるわけではありません。異なる GPU セット上で同じ演算を行っても、すでに ulp レベルで差異が生じるためです。この手法こそが、パイプライン境界とコンテキスト並列レイアウトの検証に用いられています。CP=2 と CP=1 のビルドで呼び出しが同一であれば 2 つはビット単位で一致し、フォワードカーネルを交換した場合でもコサイン類似度の中央値は 0.996 に達します。
重み同期のアサーション
トレーニング側と推論エンジン間で、テンソルごとの SHA256 マニフェストを照合します。エクスポートされたテンソルが一つも更新されていないアダプタのアップデートを検知してエラーとするバリデーター、バージョン 1 ですべての B ファクターがゼロであることを確認するチェック、そしてアダプタ勾配とオプティマイザステップに関する検証を実装しています。
学習結果
報告された実験は、DAPO(Direct Preference Optimization)を用いた数学問題の学習です。16 ノード × 4 GB300 の構成で、BF16 精度のトレーニングを TP8 / PP8 / EP8 で並列化し、ネイティブ MXFP4 ロールアウトエンジンとコロケートして実行しました。応答長は 4096 トークン、ロールアウトごとのサンプリング数は 64、各ロールアウトでオプティマイザステップを 1 回行います。LoRA のランクは 32、αは 64、学習率は 1e-5 です。KL 項なしの GRPO を採用し、壁時計時間で 12 時間まで実行しました。
AIME-2024 のグリディ評価では、60 ステップで正答率が 43.3% から 76.7% に上昇。30 問中 13 問をクリアし、カットオフ時点でもまだ改善傾向にあります。一方、トレーニングとロールアウトの KL 散逸は、MXFP4 の量子化フロアである約 2e-3 で全期間を通じて一定に保たれています。
*12 時間の DAPO 実行結果。 左:ロールアウトごとにサンプリングした 64 件の応答に対する平均報酬(応答チャネルで評価)。各ロールアウトは異なるプロンプトバッチを抽出するため、トレンドがシグナルとなり、ステップごとの値はノイズとして扱われます。中央:AIME-2024 のグリディパス率。トレーニングと同じ 4096 トークン制限の下で、10 ロールアウトごとに評価されます。この制限を超えて正解に至った場合はミスとみなされます。右:サンプリングされたトークンにおける KL(ロールアウト ‖ トレーニング) のシュルマン k3 推定値。これは損失関数に追加されることはなく、報告のみ行われます。量子化フロアで平坦であることが目標であり、単調な増加が収束の兆候となります。
謝辞
本プロジェクトは、RadixArk の SGLang & Miles チームと Moonshot AI チームが協力し、NVIDIA、AMD、Approaching AI、Baseten、Modal とともに実現したものです。
AMD: Wun-guo Huang, Xinyi Song, Hai Xiao, Soga Lin, Duyi Wang, Thomas Wang
Approaching AI: Huanming Shen, Xiaohao Zhang, Nan Li, Mingxing Zhang
テスト環境として AMD インスタンスを提供いただいた DigitalOcean にも感謝いたします。
また、Kimi K3 を SGLang で提供いただいている Google Cloud、DigitalOcean、Nebius、fal、RunPod、DeepInfra、GMI Cloud の皆様にも厚く御礼申し上げます。
原文を表示
We are excited to announce Day-0 support for Kimi K3
in SGLang and Miles. K3 is the first open-source model in the 3-trillion-parameter class,
and its hybrid architecture departs from convention in almost every place a serving stack
has assumptions. In collaboration with the Moonshot AI and NVIDIA teams, the two cover K3
in full on launch day: SGLang for inference, Miles for RL training. This post covers what
it took.
Highlights
- A new hybrid architecture. 2.8T parameters, 69 KDA linear-attention layers
interleaved with 24 MLA, LatentMoE, and Attention Residuals; most of a serving stack's
assumptions break somewhere.
- Memory management for two kinds of state, including a recurrent state that
overwrites itself in place, with prefix caching, overlap scheduling and paging rebuilt
on top, and a unified pool design that removes the last sizing guess.
- A kernel ladder reaching ~113 tok/s at batch 1, before speculation.
- DSpark speculative decoding, with a draft model we trained for K3: batch-1 decode
reaches ~423 tok/s. ReplaySSM handles the KDA state, replaying raw inputs instead of
snapshotting per step, a roughly 32x cut in draft-window memory.
- Parallelism split by phase: chunked pipeline-parallel prefill and tensor-parallel
decode, composing under PD disaggregation to 2,808 tok/s per GPU.
- LoRA RL with Miles on the native MXFP4 checkpoint, colocated trainer and rollout
on the same GPUs: AIME-2024 43.3% to 76.7% over a 12-hour run.
Launch commands and per-workload configuration guidance live in the
Bringing up Kimi K3
Kimi K3 is the first open-source model in
the 3-trillion-parameter class: 2.8T parameters, a 1M-token context window, and native visual
understanding. Its predecessor K2.5 was architecturally a close
relative of models SGLang already served well, so most of the stack simply applied. K3 is not like
that — it departs from convention in several independent places at once:
Kimi K2.5Kimi K3
Scale~1T parameters2.8T parameters
Attentionall-MLA, 61 layershybrid: 69 KDA linear attention + 24 MLA, 93 layers
Attention Residualsnoneevery attention output banked, aggregated per block
MoE384 experts, top-8LatentMoE: 896 experts, top-16, run in a 3584-dim latent space
Expert activationSwiGLUSiTU
Vision towerMoonViTMoonViT3d, a new K3-specific stack
Each of these rows is a serving problem. The hybrid attention stack is what makes the 1M context
affordable, but it means the server holds two kinds of state at once: one fixed-size KDA state per
request next to MLA's per-token KV — the memory-management sections below are about exactly this.
Attention Residuals thread a bank of attention outputs through the whole stack, which breaks the
assumptions SGLang's standard layer plumbing makes and forced K3-specific paths in places like DP
attention. LatentMoE routes 16 of 896 experts inside a down-projected latent space with the SiTU
activation instead of SwiGLU, so no existing MoE kernel applied out of the box. And the vision path —
tower, projector, processor, and Kimi's XTML media format — was brought up from scratch. The sections
below all build on this bring-up.
Hybrid KDA Memory Management
K3's hybrid KDA and MLA architecture creates two memory management challenges. The first is making prefix caching safe and efficient for mutable KDA recurrent state. The second is dynamically sharing capacity between KDA state and MLA KV.
Safe and efficient state reuse
Attention KV is append-only. Once computed, it never changes, so the scheduler can safely share cached prefixes across requests in the radix tree. KDA state is different. Each layer maintains a fixed-size recurrent buffer that is overwritten in place at every token, so prefix caching must manage mutable state rather than immutable KV.
In SGLang, we use three explicit state moves between the radix tree and a request's working slot to cache and reuse mutable KDA state safely. Copy-on-write restores a shared checkpoint into a private slot before the forward mutates it. Snapshot captures the advanced state, and donate transfers that snapshot to the tree.
*The three state moves. Copy-on-write, snapshot, and donate, and where each
one sits relative to the serial forward stream.*
The design is race-free by construction. Restore and snapshot copies are enqueued on the forward stream between the operations that produce and consume the state, so stream ordering prevents races with in-place updates. Snapshots alternate between the two slots of the ping-pong extra buffer, so the next snapshot cannot overwrite state being attached to the tree. Donate transfers only a slot index and copies no state. The second slot is allocated lazily at a boundary and released immediately afterward, avoiding a permanent slot per request. This requires neither device-wide synchronization nor a lock on the hot path.
Checkpoint placement determines the cache hit rate. A recurrent state cannot run backwards and must be replayed forward from an earlier checkpoint. We take checkpoints only at aligned radix tree nodes, at chunk boundaries during prefill and at a fixed token interval during decode, and keep a sparse set under a per-path cap and LRU. Inspired by Marconi, we prioritize branching points because their prefixes are shared by all child branches. When a request diverges mid-edge, it replays from the nearest checkpoint above and plants one at the aligned fork, allowing later branches to restore there directly.
*Checkpoints on the radix tree. The sparse checkpoint overlay and the
branching point.*
Race-free state movement and branch-aware checkpoint placement provide high-reuse KDA prefix caching that composes cleanly with the overlap scheduler, speculative decoding, and paged KV.
Unified memory: one pool for both kinds of state
Everything above manages each kind of state inside its own pool, and the pools themselves
are the one guess left in the design. The two allocation units are three orders of magnitude
apart: one large KDA state block per request (about 54 MB under TP=8, covering all 69 KDA
layers) and one small MLA KV block per token (about 27 KB, covering all 24 MLA layers), so
today they live in two separate pools, sized at startup. That sizing is a bet on the
traffic, and when the bet is off the server runs out of memory in one pool while the other
still has room to spare.
Unified memory replaces the two pools with one: KDA states fill in from one end, MLA KV blocks from the
other, and the unused bytes between them form a single free region.
*Unified memory. Top: today the two kinds of state have separate pools, sized at startup, so
one can sit idle while the other is full. Middle: with unified memory both kinds allocate from the same
pool, growing in from opposite ends with a single free region between them. Bottom: freeing a state
block in the middle leaves a gap, and a block from the end is moved into it so the free region stays in
one piece.*
Freeing is just as simple. When a request finishes, is aborted, or is retracted under pressure, its
blocks are released. If that leaves a gap in the middle, a block from the end is moved into it so the
free region stays in one piece.
This layout gives us flexible page sizes and no memory fragmentation: a 54 MB KDA state block and a 27 KB MLA KV
block draw from the same bytes with no common page size forced on them, and the moves above keep each
end packed at a negligible cost for state movement, so the free space is always one contiguous region
usable by either kind. Capacity therefore follows the workload rather than a
startup flag: many short requests fill the pool with state blocks, a few long contexts fill it with KV,
and neither case needs anything reconfigured.
Unified memory ships opt-in behind --enable-unified-memory. A follow-up post will go through the
implementation in detail.
Speculative decoding with DSpark
K3 ships with DSpark block speculative decoding, driven by
a draft model we trained for K3.
Two parts of the integration deserve their own story: spending the verification budget
only where it pays, and making K3's recurrent KDA state survive speculation at all.
Verify only what is worth verifying
DSpark proposes a block of draft tokens per step, and the target verifies the whole block
in one forward. At batch size 1 the extra verify positions are essentially free: the step
is latency-bound, and a few more tokens ride along. As the batch fills, that stops being
true. Verify tokens now compete with every other request for the same step time, and most
of them lose their bet anyway: on a chat workload the accept length is around 2.7, so five
of the eight verified positions on a typical step get rejected. Verifying everything pays
full price for tokens the server then throws away.
The pieces to fix this were already in the system. The draft carries a trained confidence
head that predicts, per position, how likely each token is to survive verification. On the
other side, a one-time profile of the server records what an extra verify token actually
costs at each load level. A per-step planner joins the two: each request keeps verify
tokens only while their expected value covers the marginal cost, and the rest of the
window is trimmed before the target forward launches. What remains is verified exactly as
before, so outputs stay lossless. The trade is a slightly shorter accepted run for a
cheaper step.
*Trimming wins under load. Decode throughput, verify-all vs trim, on a chat
panel (accept ~2.7, left) and a few-shot math panel (accept ~5.0, right). Break-even
through bs 8, then the gap opens with batch size: +68% and +24% at bs 256, with accept
length easing from 2.7 to 2.2 and from 5.0 to 4.3 while throughput climbs.*
The measured cost curve turned out to be the interesting part. The marginal cost of a
verify token is not smooth. It is a staircase: flat shelves where another token slots into
the current kernel waves for almost nothing, and steep risers where it starts a new one.
The planner reads this surface. When a request's next tokens sit on a cheap shelf it keeps
them; when they would start a riser it cuts there. In the data this shows up as small
wiggles in the accept-length curve while throughput stays smooth and monotone: the planner
is surfing real hardware shelves, not noise.
Below batch size 8 there is little pressure to relieve, so trimming is break-even to
mildly negative there; a small-batch early-exit in the planner is the known follow-up.
ReplaySSM: raw-input replay for the KDA state
Speculative decoding verifies γ+1 draft tokens at once and may accept only a prefix, so the state has to be rewindable. A KDA layer's state overwrites itself every token, so the baseline snapshots the whole K×V state after every draft step, 64 KB per request, layer and head at K=V=128, times γ+1 steps. Across K3's 69 KDA layers that scratch is reserved per running request, so it competes with the persistent state pool and caps concurrency.
ReplaySSM keeps each draft step's raw inputs Sᵢ = (vᵢ, kᵢ, gkᵢ, βᵢ) instead, about 1 KB, written by the verify kernel on its way through. Once the sampler fixes the accepted length, one fold kernel covering every layer and head replays just the accepted prefix from the committed checkpoint and advances it in place. The fold is a verbatim clone of the verify recurrence and consumes the gate values the verify kernel stored rather than recomputing them, so the rebuilt state is bit-identical to what the recurrent baseline would have committed.
- Memory: the draft window goes from 512 KB to 16 KB, roughly 32×. Handing that memory back to the state pool lifts the concurrency ceiling several-fold.
- Speed: the verify kernel no longer writes a snapshot per step, which cuts its memory access, and the replay is one fused launch for all layers and heads, so per-step time improves even at small batch sizes.
*KDA ReplaySSM. Verify reads the checkpoint h₀ without writing it and stores
each draft step's raw inputs into a per-slot buffer, one per layer and head. The fold
kernel then replays only the accepted prefix and overwrites the slot in place.*
Kernels
A single-sequence decode step on a 2.8T hybrid model is not compute-bound —
it is a launch-count and latency problem: 93 attention layers (69 KDA + 24
MLA) plus 92 latent-MoE layers per token, and at bring-up the step fired
hundreds of tiny kernels. The campaign was profile-driven: fuse one thing,
A/B it on the fixed protocol, gate on GSM8K, repeat.
Launch & copy elimination (P1–P4, +19.9 tok/s). Make the step smaller,
not the kernels faster: the MoE front collapses into one GEMM and KDA's
paired skinny projections merge (P1); a profile-guided sweep removes
per-layer upcasts, copies and spare launches (P2, P3); routing becomes a
one-pass register-resident radix select over the [M, 896] logits (P4).
NVIDIA compute kernels (P5–P8, +10.3 tok/s). Four stages swap in kernels
co-developed with or provided by NVIDIA where the default was wrong for the
bs=1 shape: the fused KDA decode kernel, the trtllm-gen W4A8 SiTU MoE cubins
behind our routing bypass, the TMA attention-residual aggregation, and the
CuTe-DSL TGV bf16 GEMMs that replace cuBLAS at small M — each through the
same A/B and accuracy gate as everything else.
Communication fusion (P9, P12, P13, +27.6 tok/s). The largest era: a
fused all-reduce family for the MNNVL fabric on CustomAllReduceV2's
symmetric-memory plane — one-shot multicast stores for small messages, NVLS
in-switch reduction for large ones, with the residual add and RMSNorm riding
inside the collective (P9). The MoE finalize then moves into the collective's
staging pass (P12), and the up-projection flips from a replicated to a
column-parallel GEMM completed by a multicast all-gather (P13).
Overlap & prologue fusion (P10, P11, P14, P15, +10.4 tok/s). The rest
trims the remaining critical chain: a strided-input MXFP8 quant (P10),
residual write-back fused into upstream kernel tails with independent
branches on a side stream (P11), KDA's GEMV chain overlapped with the qkvg
GEMM (P14), and the MLA decode prologue fused into one kernel with PDL
arming the attention kernel behind it (P15).
The chronological ladder behind the four bars:
The lesson that generalizes. All-reduce is a synchronization point, so a
microsecond saved there converts one-for-one into step time; a kernel sitting
in another stream's overlap slack converts at roughly one-tenth. Checking
critical-path membership in the trace before writing a kernel was the single
highest-leverage habit of this campaign.
¹ A rebase moved the P5 baseline 64.2 → 63.8; the step is measured against
the post-rebase base.
² P13 is a re-calibrated canonical baseline after a merge window, not a
single-PR attribution.
³ Segments leading into P1–P4 also carry interim optimizations later
superseded (concat all-reduce → P9; tiny/1-CTA GEMV → P8; radix router v1 →
P4; Marlin top-k-sum → P6/P12; early attn-res add → P7); their transient
gains are in the curve but have no named point.
Parallelizing K3
For K3's hybrid architecture, the conventional parallelism choices do not hold up. Tensor
parallelism cannot shard MLA's KV cache (one KV head, nothing to split by), so every rank
holds a full copy; it also slices every GEMM eight ways and pays a collective per layer.
Pure DP attention instead replicates the attention weights on every rank, around 61 GB for
KDA plus 11 GB for MLA, memory the KV cache and KDA states need. Prefill and decode fail
differently under these costs, so K3 splits the answer by phase: chunked pipeline
parallelism for prefill, context parallelism for decode.
Prefill: chunked pipeline parallelism
In TP prefill, every layer ends in an AllReduce, a barrier that cannot overlap with
compute. Pipeline parallelism cuts the model by layer instead: K3's 93 layers become 8
stages, and the prompt is cut into chunks that stream through them:
*Chunked pipeline-parallel prefill. The stages work on different chunks at the
same time. The hand-off between stages runs while the stage computes its next chunk, so
91% of it is hidden on K3. Under TP (bottom strip), every layer ends in an AllReduce
that all ranks wait for.*
This wins three ways. The only communication left, the hand-off to the next stage, hides
behind the next chunk's compute. Each rank runs whole layers, so the GEMMs are eight times
wider and more efficient. And each stage holds KV and activations for only its ~12 layers,
which makes very long prompts comfortable to prefill. The pipeline does have to be deep:
a shallow PP4×TP2 cannot cover its hand-offs and still pays TP2's AllReduce, and it
benchmarks no better than TEP8.
Measured on 8K prefill across 2×4 GB300, with topology as the only variable:
*Deep PP wins on both axes. Left: past the c1–c4 crossover, PP8×TP1 climbs to
about 1.7× TEP8's ceiling, with lower TTFT. Right: cost per 1k prefill tokens at equal
per-rank FLOPs; PP4×TP2 and TEP8 trade compute against communication and tie, PP8 is
cheapest on both.*
PP8 loses only at a single request, and a prefill worker that idle is misconfigured
anyway. In K3's disaggregated serving, the prefill nodes run PP8. Each of them has 1.45 to
1.72 times the prefill capacity of a TEP8 node, so one prefill node can keep several
decode nodes fed. The decode nodes run TP or DCP, which the next part covers.
Decode: context parallelism
Decode is where the replicated KV cache binds: under TP, admitting one more request or one
more thousand tokens of context costs the same bytes on every rank. Decode Context
Parallelism (DCP) shards MLA's KV by token position instead of by head. Rank r owns
position p when p mod N equals r, so every rank holds an interleaved 1/N of every
request's context, and the sharding is invisible above the attention kernel:
*Sharding by position. The same 16 token positions on 4 ranks: TP stores 64
physical copies, DCP stores 16. The freed bytes become logical KV capacity, about 7.9×
on K3 with DCP8.*
Position-sharding breaks the softmax, since each rank sees only 1/N of the keys and
partial softmaxes do not add. The fix is FlashAttention's own: each rank returns its
partial attention output together with a log-sum-exp per head, and one all-to-all per
layer exchanges them so that each rank ends up with 1/N of the heads over the full
context. A local merge by log-sum-exp is then exact, and the result is already the head
layout the output projection expects under TP. One collective per layer is the entire
communication cost:
*The decode step, per layer. Each rank projects the full-head query locally,
attends over only the positions it owns, and ships partial outputs and their
log-sum-exps in a single packed all-to-all. A local merge leaves the standard TP head
layout; nothing downstream of attention changes.*
Everything else is left alone. DCP groups are built inside the TP group, so TP8 with DCP8
is still 8 GPUs, and the MoE runs whatever parallelism it already had. KDA is the
exception that explains the rule: its state is one fixed-size matrix per request rather
than per token, so there is no position axis to shard, and the KDA layers stay TP-sharded
by head. The whole feature is one flag, --dcp-size N.
What this buys shows up on agentic traffic, where hundred-thousand-token sessions pile up
in the cache. Replaying real coding-agent sessions on 2×4 GB300, with a host-memory KV
tier on both arms and DCP as the only difference:
*DCP removes the active-set wall. The host tier rescues re-prefill traffic,
but at 16 concurrent sessions the active working set outgrows TP8's device KV and
throughput collapses. DCP8 lifts logical KV from 1.5M to 12.2M tokens and carries the
same workload to 541 tok/s at 48 sessions.*
DCP composes with the rest of the stack. A DSpark verify step is a decode step, so it
rides the same replicated-Q, one-all-to-all path; under PD disaggregation the prefill
side stays DCP-unaware and each decode rank simply pulls the positions it owns at the
transfer boundary, which is what lets PP or TP prefill feed DCP decode. The remaining
ceiling is KDA: its per-request state cannot be position-sharded, so once DCP lifts the
MLA wall the running-request cap becomes the binding limit; the unified-memory design in the memory section above takes that on.
Combining all strategies together
How the pieces compose is ultimately a measurement. Putting PD disaggregation in the
loop, with the prefill topology, decode topology and prefill:decode ratio varying, gives
the serving frontier:
*The serving frontier. At the throughput end one PP8 prefill worker feeding
one TP8 decode node delivers 2,808 tok/s per GPU on the fp4 arm, with the DCP
composition, two PP8 prefill workers feeding two DCP8 decode nodes, just behind at
2,633. Moving right is the prefill:decode knob at work: one prefill worker feeding two,
three, then four independent decode instances trades aggregate throughput for per-user
speed, out past 116 tok/s per user.*
RL: LoRA Training on the Native MXFP4 Base
Day-0 RL for K3 is colocated LoRA training with Miles: a BF16 trainer on Miles's Megatron backend and native packed MXFP4 SGLang rollout engines sharing the same 64 GB300s. The backend covers KDA, NoPE-MLA, the attention-residual bank and the latent MoE, with TP/SP/PP/CP/EP.
LoRA serving and weight sync
The engines serve the checkpoint as shipped and never rewrite it. Each step transfers only BF16 LoRA adapters, and the engine applies the delta as a separate BF16 B(Ax) term on top of the quantized base GEMM, so the policy update reaches inference at full precision over 4-bit base weights. Dense projections go through SGLang's Triton LoRA backend, the 896 routed experts through a fused MoE-LoRA kernel on the Marlin path, and the shared-expert delta is folded into the fused MoE front GEMM. Adapters live in a GPU memory pool that each sync replaces in place, so there is no resident BF16 copy on the rollout side, no full-weight resync, and no requantization step in the loop.
Parallelism
Pipeline parallelism. K3's attention-residual snapshot bank has to cross stage boundaries, but Megatron's point-to-point carries one hidden-states tensor, so the stage boundary packs [prefix_sum, bank] into it and unpacks on entry. Megatron is untouched.
Context parallelism. The bank shards for free, since every attention-residual operation is per-token. MLA needs no K3-specific code either: rotary-table slicing is the only CP-aware work in Megatron's MLA and K3 has no rotary embedding, so its projections sit on the stock TE attention core and inherit CP. KDA is the part that needs work, going through fla's CP context for the recurrent state and the convolution halo. It wants a contiguous rank-local chunk where Megatron stores the zigzag order ring attention expects, so the relayout runs only around KDA.
Expert parallelism. The adapter shares its latent-side factor across all 896 routed experts and keeps the other per-expert, matching the engine's fused MoE-LoRA contract. That shared factor is replicated across EP but tagged expert-parallel, so DDP reduces it only over expert-DP, and the EP sum is the one gradient reduction the release adds itself.
Memory
A native-MXFP4 rollout peaks near 225 GiB/GPU and the BF16 trainer near 155 at initialization, on a 277 GiB card, so the two are never both resident: each sleeps while the other runs, and colocation depends on every handoff leaving nothing behind.
The base never moves. The engines keep their base weights resident on the GPU for the whole run, and only the KV cache and CUDA graphs are released while the trainer works. Since the base is never released it is never restored either, and the LoRA path ships no base bytes at all: base sync is skipped entirely.
The process groups stay up. The trainer's NCCL communicator buffers are not torch-allocator memory, so offloading cannot release them, and the obvious answer is to destroy the groups on sleep and rebuild them on wake. That trades a resident few GiB for a per-cycle rebuild and EP warmup, which is only worth paying when the engines need those bytes back — and they do not here, because no base weights cross the update path. The groups stay up.
Adapter transfer is bounded per chunk. The adapter is about 2,800 tensors. Each chunk ships as one flattened CUDA IPC bucket, 278 in total, and is collected as soon as the receiver acknowledges it and every producer rank has crossed an engine-group barrier. That bounds transient IPC memory per chunk instead of letting it integrate across the transfer, worth about 48 GiB/GPU at the peak.
Host copies exist only where something reads them. The engine releases each adapter's CPU copy once it is installed into the GPU pool, taking scheduler RSS from 76–88 GiB to about 17; an adapter whose CPU copy is gone cannot be reinstalled, so pool eviction is refused with an error rather than silently serving a stale slot. On the trainer side the DDP buffers split by lifetime: adapter parameter buffers stay in the CPU-backed region because the weight update reads them while the trainer sleeps, while gradient buffers are rebuildable and go to a no-backup region that sleep discards.
Validation
Training correctness is the most important part of an RL support stack. We built these checks before the recipe work.
Train/rollout KL, logged every rollout, is a Schulman k3 estimate of KL(rollout ‖ train) over the sampled tokens, computed from the log-probabilities the engine returned and the trainer's recomputed ones. It is a diagnostic, not part of the objective, which carries no KL penalty. Its floor is about 2e-3, set by serving the base in MXFP4; growth with step count is the divergence signature.
A canary lockstep probe. One trajectory is pinned from the first rollout and scored every step by both the engine, under the live adapter, and the trainer, at the same policy version. Co-movement of the two curves is what proves the trained adapter reaches inference: a transfer checksum proves the bytes arrived, not that anything reads them.
Tensor-level dump comparison. The same tokens run through two builds or two parallel layouts, dumping every forward activation and every parameter gradient, compared by relative L2 and cosine against a measured noise floor. Bit-exactness is not the bar, since the same arithmetic on two different GPU sets already differs at the ulp level. This is how the pipeline boundary and the context-parallel layout were validated: where CP=2 and CP=1 build the identical call the two are bit-identical, and a forward-kernel swap lands at cosine median 0.996.
Weight-sync assertions. Per-tensor SHA256 manifests between trainer and engine, a validator that errors if an adapter update changes no exported tensor, a check that every B factor is zero at version 1, and adapter gradient and optimizer-step checks.
Training result
The reported run is DAPO math on 16 nodes × 4 GB300: BF16 trainer at TP8 / PP8 / EP8 colocated with native MXFP4 rollout engines, 4096-token responses, 64 samples per rollout, one optimizer step per rollout, rank-32 / α-64 LoRA at lr 1e-5, GRPO with no KL term, run to a 12 h wall-clock limit. AIME-2024 greedy eval climbs 43.3% to 76.7% over 60 steps, 13 of 30 problems to 23, and is still rising at the cutoff, while train/rollout KL holds flat at the ~2e-3 MXFP4 floor for the whole run.
*The 12-hour DAPO run. Left: mean reward over the 64 sampled responses per
rollout, graded on the response channel; each rollout draws a different prompt batch, so
the trend is the signal and the per-step value is noise. Middle: AIME-2024 greedy pass
rate, evaluated every 10 rollouts under the same 4096-token limit as training, so a
correct solution that runs past the limit scores as a miss. Right: a Schulman k3 estimate
of KL(rollout ‖ train) at the sampled tokens, reported only and never added to the loss;
flat at the quantization floor is the target, and monotonic growth is the divergence
signature.*
Acknowledgments
This work was a collaboration between the SGLang & Miles team at RadixArk and the Moonshot AI team, together with NVIDIA, AMD, Approaching AI, Baseten, and Modal.
AMD: Wun-guo Huang, Xinyi Song, Hai Xiao, Soga Lin, Duyi Wang, Thomas Wang
Approaching AI: Huanming Shen, Xiaohao Zhang, Nan Li, Mingxing Zhang
Thanks to DigitalOcean for providing AMD instances for our testing.
Thanks to Google Cloud, DigitalOcean, Nebius, fal, RunPod, DeepInfra and GMI Cloud for serving Kimi K3 on SGLang.
同じ出来事を4媒体で確認
同じ出来事を扱う別媒体の記事です。見出しと公開時刻を比較できます。
関連記事
今日のまとめ
AIデイリーブリーフで今日の重要ニュースをまとめ読み