SGLangが2週間でGLM-5.2を最適化：8×B300上で500 TPS超を達成

SGLangチームがGLM-5.2-NVFP4のサービス最適化レビューを公開した。2週間でday-0サポート版をより安定した本番向け推論パスへ進化させ、8×B300・バッチサイズ=1の低レイテンシ環境で500 TPS超を達成した。

図1：GLM-5.2 NVFP4（8×B300上）における Day-0 と v0.5.15.post1 のパフォーマンス比較。

主な結論

• 8×B300・bs=1環境でスループット500 TPS超を達成。
• GLM-5.2 MTPに対して同期オーバーヘッドゼロのspeculative decodingを実現。
• SGLangにIndexShare MTPを内蔵し、Spec V2と接続。
• TopK-V2は80K ISL環境でTopK-V1より2.33×高速。
• Indexer Prologue Fusionにより関連カーネル数を12個から4個に削減。
• BF16層向けにCuTe DSL GEMM最適化を導入。

背景：GLM-5.2のアーキテクチャ変更

GLM-5.2は従来のGLMチェックポイントの主要構造を継承している。DeepSeek-V3スタイルのMoEに、sparse-attentionインデクサーを持つDSAを重ねた構成だ。初期バージョンとの比較では、2つの重要な変更点が加わっている。DSA向けのIndexShareと、IndexShareとKVShareを組み合わせたMTPだ。

SGLangはday 0からGrace Blackwell / Blackwellハードウェア上でGLM-5.2-NVFP4チェックポイントをサポートし、sparse attentionとMoEの処理にはtrtllm-genカーネルを使用している。今回の最適化の目標は、当初の動作可能なスタックをより高速・安定した本番デプロイ向けのサービスパスへアップグレードすることだった。

ランタイム最適化：Spec V2とゼロオーバーヘッドスケジューリング

Spec V2は、SGLangのspeculative decoding向けオーバーラップランタイムだ。基本的な考え方は、GPUがforward streamで現在のモデルforwardを実行している間に、CPUとplan streamが次ステップに必要なKVアロケーションとメタデータを並行して準備することで、CPUのブックキーピングをGPU実行時間内に隠蔽するというものだ。

SGLangは最近Spec V2をデフォルトで有効化した。理論上、オーバーラップスケジューラーによりCPUはGPUが現在のイテレーションで処理中に次ステップの準備を完了でき、2回のrun_batchの間にほぼバブルが生じない。しかし実際の実装では、チームはいくつかの同期ポイントを解消する必要があった。

• DSA draft-extendパスがCUDA graphをサポートするようにする。
• seq_lens_cpuをDSAに対してオプションにし、D2H syncを排除する。
• 残存するH2D syncを削除する。
• _apply_cuda_graph_metadata内の小型eagerメタデータ操作を融合する。

これらのGPUバブルを排除した結果、エンドツーエンドのTPSが約11%向上した。

図2・図3：Spec V2最適化前後のdecode traceの比較。最適化後はrun_batchのイテレーション間に明確なバブルが見られなくなった。

SGLangにおけるIndexShare MTPの実装

GLM-5.2は比較的強力なMTPヘッドを備えており、accept lengthが5以上に達することが多い。これはエージェンティックコーディングなどの低レイテンシワークロードにとって非常に重要であり、受容長が長いほどターゲットモデルの検証頻度を大幅に下げ、インタラクティブなスループットを向上できる。

GLM-5.2のMTP動作を正しくサポートするため、SGLangはspeculative decodingランタイムに2つの改修を加えた。まず、IndexShareはSGLangが複数のdraftステップ間でDSAインデクサーのtop-kを再利用することを要求する。draft step 0で計算されたtop-kを保持し、後続のdraft stepに引き継ぐことで、各ステップがインデクサーを再度実行する必要がなくなる。長いコンテキストでは、これによりdraft-stepのコストを最大約1.9×削減でき、出力品質にも影響しない。

次に、top-kのseedは正確な位置から取得しなければならない。SGLangでは、それは前回のrun_batchイテレーションのdraft-extendから取得される。Spec V2のステップは非同期実行されるため、チームはこのseedをオーバーラップスケジューラーのリレーバッファを通して受け渡し、イテレーション間で失われないよう保証する必要があった。

カーネル最適化その1：TopK-V2

DSAインデクサーは各クエリを過去のKV位置に対するスコアに変換し、sparse attentionに必要な候補位置を選択する。以前のDeepSeek-V4最適化で提案したLightning-TopKの考え方を継承し、SGLangは既存のDSA TopK-V1をTopK-V2へアップグレードした。TopK-V2はTopKをソーティング問題ではなく選択問題として扱う。

図4：TopK-V2は8つのCTAからなるクラスターで長いスコア行を処理し、各CTAがローカルの10ビットヒストグラムを構築し、クラスター全体のreductionで境界binを特定する。

図5：TopK-V2はヒストグラム構築時にFP32スコアをFP16に丸め、数値順序を保持するunsignedキーに変換する。上位10ビットで1024個のbinを選択する。

TopK-V2は長さに応じて異なる処理パスを採用している。短い行と中程度の長さの行にはregister-residentまたはsingle-CTAのストリーミングパスを使用し、長い行は8つのCTAからなるクラスターがローカルの10ビットradixヒストグラムを構築し、クラスター内でreductionを行って第2048番目のスコアを含むthreshold binを特定する。

threshold binより高い値は直接出力され、境界領域の候補は精密なFP32 radix selectionへ進む。その後、論理位置は物理インデクサーKVキャッシュスロットに変換される。このカーネルはランタイムでkを最大2048までサポートし、精密なtie-breakにより厳密にkエントリを返す。

さらに、planningカーネルはバッチのシーケンス長分布に基づいてcluster cutoffを選択し、persistent cluster poolのワークリストを構築する。生成されたplanはforward毎に作成され、全DSA層で共有される。TopK-V2はさらにselectionとpage-table変換を単一カーネルに融合し、レイテンシをさらに低減している。

図6：TopK-V1とTopK-V2のカーネルレイテンシ比較。テスト環境はバッチサイズ1・draft token 6個のターゲットモデル検証シナリオ。

ベンチマーク結果では、80K ISL環境においてTopK-V2の平均カーネルレイテンシは40.7 μsから17.5 μsへ短縮され、2.33×の高速化を達成した。コンテキスト長が増加するほど優位性はさらに拡大し、1M ISL環境ではレイテンシが372.1 μsから36.6 μsへ短縮され、10.17×の高速化に達した。これはTopK-V2が長いコンテキストのワークロードに対してスケーラビリティが大幅に優れていることを示している。

カーネル最適化その2：Indexer Prologue Fusion

図7：DSA Indexer Prologueカーネル融合前後の依存チェーンの変化。

DSAインデクサーprologueは2種類のデータを準備する。インデクサーKVキャッシュに格納するキー表現と、sparse attention候補の計算に使用するクエリ表現だ。元の実装は一連の小さなカーネルとprojectionで構成されていた。

融合前、キーブランチはwk、LayerNorm、RoPE、Hadamard transform、FP8量子化、キャッシュストアを順次実行し、クエリブランチはwq_b、RoPE、Hadamard transform、FP8量子化、head-gate scalingを実行していた。またweights_projはper-head gateを生成する独立したprojectionだった。

PR #27705は2つの方向からこの依存チェーンを短縮した。まず、wkとweights_projを単一のBF16 projectionに融合した。wk_weights_projだ。その出力はキーactivationとraw head-gate weightsに分割され、インデクサーパスから小型GEMMを1つ削除できるとともに、head-gate weightsを融合後のクエリカーネルが直接再利用できるようになった。

次に、チームはelementwise tailを融合した。

• キーパス：LayerNorm + RoPE + FP8量子化 + pagedインデクサーKVキャッシュストア。
• クエリパス：RoPE + FP8量子化 + head-gate scaling。

融合後のスケジュール構造はより短く明確になった。キーブランチはキャッシュストアを含む単一カーネルとして実行でき、クエリブランチは別の融合カーネルとして実行でき、両者は並行してオーバーラップ実行できる。最終的に、関連カーネル数は12個から4個へ削減された。

融合パスではHadamard transformも削除された。QとKに同一のorthonormal transformを適用することで量子化前の内積が保持されるため、その主な影響は量子化表現に集中しており、融合後のパスでは変換前のactivationを直接量子化する方式へ変更された。

カーネル数の削減はdecodeスループットの向上に直結しており、特に小さいバッチサイズで顕著だ。バッチサイズ1ではdecodeスループットが約8%向上し、バッチサイズ128では改善幅は小さいが安定して約5%の向上が見られた。

カーネル最適化その3：BF16 GEMMの改善

図8：CuTe DSL BF16 GEMMのcuBLAS GEMMに対する各バッチサイズでの高速化効果。

GLM-5.2のすべての行列積がNVFP4で実行されるわけではない。精度を保護するため、チェックポイントの量化戦略はattention projectionとshared-expert MLPにはBF16を維持し、routed expertのみを量子化している。

PR #30117はこれらのBF16層専用のオプションバックエンドとして、FlashinferのTGV GEMMをベースにしたCuTe DSL BF16 GEMMバックエンドを導入した。このカーネルは処理を異なるwarpに分担させる設計で、一部のwarpがメモリからデータをロードし、1つのwarpが行列積を担当し、残りのwarpが結果を書き戻す。ロード・演算・ストアを並行してオーバーラップできるため、従来のパスと比較して待機時間を削減でき、複数のバッチサイズでcuBLASを上回る性能を発揮した。

総合的な効果と長いコンテキストでのパフォーマンス

Spec V2、IndexShare MTP、TopK-V2、Indexer Prologue Fusion、GEMM最適化を組み合わせた結果、SGLangはGLM-5.2-NVFP4においてスループットとレイテンシのトレードオフを大幅に改善した。特にエージェンティックワークロードに多い低バッチ・長コンテキストのシナリオにおいて、これらの最適化はGPUバブル、CPU/GPU同期、小カーネルの起動オーバーヘッド、長シーケンスでのTopKコストをまとめて削減した。

図9：SGLang上のGLM NVFP4のパフォーマンスPareto curve。各種設定におけるスループットとインタラクティブ性のトレードオフを示す。

図10：4枚のGB300 GPU上で実施したinput sequence length ablation。入力長の変化がパフォーマンスに与える影響を示す。

まとめ

今回の最適化の本質は単なるカーネルの単点高速化にとどまらず、ランタイムオーバーラップ、MTPセマンティクス、DSAインデクサー、長コンテキストTopK、BF16 GEMM、スケジューリング依存チェーンにわたるシステムレベルの再設計にある。最終的に、GLM-5.2-NVFP4はSGLang上でday-0の動作可能版から本番サービス向けのより効率的な推論パスへと進化し、8×B300上で500 TPS超の低レイテンシスループットを達成した。