SGLang-TorchTPUはTPU v7で8160億パラメータMoEをどう高速化したか

導入

大規模モデルの推論性能は、行列積の速度だけでは決まらない。特にバッチが小さいデコードでは、演算器が処理を終えても、重みがHBMから届くのを待ったり、KVキャッシュを読み出したり、デバイス間通信が終わるのを待ったりする。さらに、層ごとに別のカーネルを起動すると、その境界で重みのストリーミングと計算の重なりが途切れてしまう。

PyTorchの技術ブログは、この問題に対するSGLang-TorchTPUの実装を説明している。SGLangがスケジューリング、KVキャッシュ、モデルサービングを担当し、TorchTPUをPyTorchバックエンドとして使う。ホットパスにはPallasカーネルを配置し、4チップのTPU v7ホスト1台でMetaのMuse Spark Openを動かす。4チップは8デバイスとして扱われ、テンソル並列で8分割される。

Muse Spark Openは、総パラメータ数8160億のMixture-of-Expertsモデルで、1トークンあたり約420億パラメータが動作する。62個のデコーダ層はすべてMoEであり、各層は256個のエキスパートから8個を選ぶ。注意機構は、直近2048トークンを見るスライディングウィンドウと、最大約100万トークンまで扱うフルコンテキストを組み合わせている。そのため、計算だけでなく重みの搬送とKVキャッシュ容量も重要になる。

主な最適化

• ステップ全体を一つのMegakernelに統合。 従来は各層で複数のカーネルと集合通信を実行していた。新しいPallasプログラムは62層と最終正規化を、デバイスごとの一つの呼び出しにまとめる。埋め込み、ドラフトモデル、LMヘッド、サンプリングなどは外側に残るが、中心となるデコードまたは検証処理は連続した実行領域になる。
• データ転送をパイプライン化。 密な重み、エキスパート重み、KV、部分和、デバイス間転送を、利用される前からDMAで進める。転送を単に速くするのではなく、MXUやベクトルユニットの計算と重ねて、待ち時間をクリティカルパスから隠すことが狙いだ。
• デコード経路を細かく改善。 集合通信の折りたたみ、KVチャンクを先読みするmax-free注意ループ、通常経路でのレジスタスピルを抑えるエキスパートループを採用する。DFlash2推測デコードでは、小さなモデルが複数トークンを予測し、大きなモデルが一度に検証する。
• KVキャッシュを多層化。 Megakernelではシーケンスをレーンに割り当てるKVレイアウトを使う。さらにTPUのL2 HiCacheによってホストDRAMを追加のキャッシュ層として利用し、セッションを考慮した追い出しで、長い会話や共有プレフィックスに関係するデータを残しやすくする。
• Prefillも融合。 融合MoEカーネルとシーケンス並列化により、長い入力を処理するprefill経路を改善する。これは逐次的なデコードとは異なるボトルネックを持つため、別の最適化が必要になる。

測定結果と意味

以前のMuse SparkチェックポイントでDFlash2を使った検証を測定したところ、Megakernelの単一リクエストでは出力トークン当たり1.95ミリ秒、従来経路では3.78ミリ秒だった。約1.94倍の改善である。2リクエストを同時に処理した場合の改善幅は1.85倍だった。並列度が高くなると差が縮まるのは、バッチ化によって演算器の利用率が上がり、カーネル境界のコストが相対的に小さくなるためだ。

検証ステップの時間は、そのステップで必要な重みをHBMから読み出すだけの時間の約1.5〜1.7倍に収まった。これは、層ごとの起動や集合通信のオーバーヘッドに支配される状態から、メモリ帯域が支配する状態に近づいたことを示す。

チームは、エージェント型コーディングセッションを基にした100万トークンのAgentXトレースセットでも評価した。4チップのTPU v7で、計算にはFP8、エキスパート重みにはNVFP4を使用している。結果は、高い対話性を重視する低スループットの条件から、1リクエスト当たりの速度を抑えつつ全体処理量を高める条件まで、複数の運用点を示した。ただし、コストや比較対象の条件は特定の構成に依存するため、すべてのモデルやクラウド環境にそのまま一般化できるものではない。

この事例の重要性は、推論最適化の単位を個別オペレーターから処理ステップ全体へ移した点にある。MoEと長文脈モデルでは、重み、KV、通信、スケジューラの境界が同時に性能を左右する。これらを一つのパイプライン可能なプログラムに組み込む方法は有望だが、TPUのレイアウト、VMEM容量、ページ型KVキャッシュ、動的なリクエストを同時に扱うため、実装は複雑になる。今後の汎用化には、コンパイラ、ランタイム、モデル実装、サービングフレームワークの連携が欠かせない。

出典：PyTorch Blog