2026年9月7日 月曜日
AI時短ラボ
研究· 約15

SGLangが「allgather+scatter」を1回のall-to-allに統合──DeepSeek-V4-Proのデコード時間が320μsから169μsに

推論エンジンSGLangはv0.5.18で、pure-DP dp-attention構成におけるTP LMHeadの通信処理を、allgatherとscatterの2段階から、単一のall-to-all通信へと統合した。公式リリースノートによれば、DeepSeek-V4-Pro・B200のデコード処理でLMHead部分の所要時間が320マイクロ秒から169マイクロ秒に半減し、TPOT(トークンあたりの出力時間)も36.97msから35.67msへ改善している。

SGLangが「allgather+scatter」を1回のall-to-allに統合──DeepSeek-V4-Proのデコード時間が320μsから169μsに
執筆・編集:
目次

推論エンジンSGLangのv0.5.18(2026年8月22日公開)には、大規模モデルを複数GPUに分散させる際の通信パターンを1つに統合する最適化が含まれている。「TP LMHead with All-to-All」と呼ばれる変更だ。実装したのはPR #32313で、2026年8月18日にマージされている。

3行まとめ

  1. SGLang v0.5.18は、pure-DP dp-attention構成のTP LMHead(テンソル並列のLMヘッド)が行っていたallgather+scatterの2段階通信を、単一のall-to-all通信に置き換えるPR #32313(2026-08-18マージ)を含む。
  2. DeepSeek-V4-Pro・B200・dp=8・バッチサイズ36の実測プロファイルでは、LMHead部分の通信コストが151マイクロ秒(allgather)から46マイクロ秒(all-to-all)に縮小し、LMHead全体では320マイクロ秒→169マイクロ秒、TPOTは36.97ms→35.67msに改善したとPR本文・リリースノートは報告している。
  3. この最適化はpure-DP dp-attention構成専用で、dp-lm-head・context parallel・dp-tp混在並列とは同時に有効化できない、とPR本文に明記されている。

何を1つの通信にまとめたのか

公式リリースノートの説明を見る。

"TP LMHead with All-to-All: The TP LMHead's allgather + scatter becomes a single all-to-all for pure-DP dp-attention. On DeepSeek-V4-Pro B200 decode, LMHead time drops 320us to 169us and TPOT improves 36.97ms to 35.67ms"

(TP LMHeadをAll-to-Allで最適化:TP LMHead(テンソル並列化されたLMヘッド)のallgather+scatterが、pure-DP dp-attention構成において、単一のall-to-all通信に統合された。DeepSeek-V4-Pro・B200のデコード処理では、LMHeadの所要時間が320マイクロ秒から169マイクロ秒に短縮し、TPOTは36.97msから35.67msに改善する)

LMHead(Language Model Head)は、モデルの内部表現を、語彙全体に対する確率分布に変換する最終段の層だ。TP(テンソル並列化)構成では、この層の計算・出力が複数のGPUに分散しているため、各GPUが持つ部分的な結果を集約する通信処理が必要になる。従来は「allgather(全ノードの結果を全ノードに集める)」と「scatter(データを各ノードに配り直す)」という2段階の通信を行っていたが、pure-DP dp-attention(データ並列のみで構成されたattention処理)という特定の構成では、この2段階を1回のall-to-all通信にまとめられる、という最適化だ。

PR本文が明かす内訳:320μsのうち通信は何割か

PR #32313のMotivation欄によれば、そもそもの出発点は「pure DPかつdp-attention構成で--enable-dp-lm-headを使うと、各rankのバッチサイズが小さいときGEMM(行列積)効率が悪化する」という問題だった。DeepSeek-V4-Proでdp=8・--enable-dp-lm-headを使う場合、各dp rankのバッチサイズが36のとき、LMHeadのGEMM自体だけで320マイクロ秒かかっていたという。この状況では「TP LMHead」(テンソル並列でLMHeadを分割する方式)の方が有利になるが、従来のTP LMHead実装は語彙サイズ(vocab_size)に対するallgather通信を伴い、DeepSeek-V4では各rankが[bs×8, vocab_size//8×7]=約31MBのデータを送受信し、これに約150マイクロ秒かかっていた、とPRは説明している。

PRが提案するのは、このallgather+scatterをAll2All(all-to-all)通信に置き換えることで、各rankが送受信するデータ量を[bs×7, vocab_size//8]の送信=約3.88MB、[bs, vocab_size//8×7]の受信=約3.88MBまで削減する方式だ。PR本文に貼られたプロファイル結果(DeepSeek-V4-Pro・B200・PD分離構成のデコードノード、dp=8・バッチサイズ36)を表にすると次の通り。

構成 GEMM時間 通信時間 LMHead合計時間
dp-lm-head(データ並列、比較対象) 320μs 320μs
tp-lm-head-allgather(従来のTP方式) 71μs 151μs 278μs
tp-lm-head-all2all(今回の変更) 71μs 46μs 169μs

GEMM自体の時間は71マイクロ秒で共通しており、削減されているのは通信部分(151マイクロ秒→46マイクロ秒)だとPRのプロファイル表からわかる。PR本文には、DeepSeek-V4でのGSM8Kテストで精度に問題がないことを確認した旨も記載されている。PRの変更規模は7ファイル・追加226行・削除11行で、著者はSYChen123氏。

数字が示す改善幅

LMHead単体の処理時間は320マイクロ秒から169マイクロ秒へと、ほぼ半分になっている。これがTPOT(Time Per Output Token、1トークンを出力するのにかかる時間)全体では36.97msから35.67msへの改善にとどまっている点は、LMHead部分がデコード処理全体に占める割合が、時間にして1.3ms程度(全体の3〜4%程度)であることを示唆している。局所的には大きな改善でも、全体のレイテンシに対するインパクトは相対的に小さい、という構造だ。

「pure-DP dp-attention」という前提条件

この最適化は、すべての並列化構成に適用されるわけではなく、「pure-DP dp-attention」という特定の構成を対象にしている。dp-attention(data-parallel attention)は、attention層の計算をデータ並列で行う設計であり、SGLangはこれを含む複数の並列化戦略(テンソル並列・データ並列・エキスパート並列など)を組み合わせられる。PR本文は「dp-tp混在並列でこの最適化を有意な改善とともにサポートするのは困難」と明記しており、実装上もdp-lm-head・context parallel(cp)・dp-tp混在並列を同時に有効化することはサポートしない、という制約が明示されている。

B200という新しいGPU世代での計測

計測に使われているB200は、NVIDIAのBlackwell世代のGPUだ。同じv0.5.18のリリースノートには、後述するFlashInfer MNNVL allreduceの最適化でもBlackwell世代GPUでの計測値が使われており、SGLangの開発チームが最新世代のGPU上でのボトルネック解消に力を入れていることがうかがえる。

320マイクロ秒を、積み重ねの規模で捉え直す

筆者はGPU間の通信最適化を自分で実装した経験はないが、「allgather+scatterという2段階の通信を、1回のall-to-allにまとめる」という発想自体は、分散処理全般に共通する最適化パターンだと理解している。通信のラウンド数を減らすことは、単純に計算量を減らすこと以上に、ネットワークのレイテンシ(往復時間)を減らす効果がある。320マイクロ秒という数字自体は人間の感覚では一瞬だが、1秒間に何千回とトークンを生成する推論サーバーにとっては、積み重なると無視できない差になる、という規模感で捉えるべき数字だと考えている。

PRの実コード差分を読んで分かったこと

https://github.com/sgl-project/sglang/pull/32313.diffを取得し、実際に変更された7ファイルのコード差分を読んだ。

  • 実際のCLIフラグ名を確認できたserver_args.py(PRのdiff・v0.5.18タグのraw.githubusercontent.com両方で確認)にはenable_tp_lm_head_all_to_allというOptional[bool]型の引数が新規追加されており、CLIでは--enable-tp-lm-head-all-to-allargparse.BooleanOptionalActionのため--no-enable-tp-lm-head-all-to-allで明示的にオフにもできる)。既定値はNoneで、明示指定がない場合は後述の条件に基づき自動的に有効・無効が決まる。
  • 自動有効化の条件は6つoverrides.pyに追加された_tp_lm_head_all_to_all_default関数によれば、「PD分離構成のdecodeノードであること(disaggregation_mode == "decode")」「enable_dp_attentionが有効」「dp_size > 1」「tp_size == dp_size」「attn_cp_size == 1」「enable_dp_lm_headが無効」の6条件がすべて揃った場合にのみ、明示指定なしでこの機能が自動的にオンになる。prefill専用ノードや、prefillとdecodeが同居するノードでは既定で無効のままだとdocstringに明記されている。
  • 実行時のフォールバック条件も具体的に書かれているlogits_processor.pyに追加された_can_use_tp_lm_head_all_to_all()は、LMHeadのTP設定・バッチのレイアウト(CUDA graph実行時のパディング済み等分割、またはeager実行時の全rank同一トークン数)を毎回チェックしており、条件を満たさないバッチでは自動的に既存のallgather経路にフォールバックする設計になっている。
  • 実装の中身_get_logits内でself.use_attn_tp_group(dp-lm-head用の経路)と_can_use_tp_lm_head_all_to_all()のどちらにも該当しない場合のみ、従来の_logits_gatherer(allgather)が使われる。all-to-all自体はget_tp_group().all_to_all_single()で1回呼び出され、送信元rank順に並んだ出力を、_reassemble_tp_lm_head_all_to_all_output()という関数がviewpermutereshapeでrank×行×語彙シャードの並びから行優先の全語彙logitsに組み直している。
  • CUDA graphキャプチャ前のウォームアップが追加されているbootstrap.pyには_prewarm_tp_lm_head_all_to_all()という関数が新設され、GPUメモリの空き容量計算(KVキャッシュのサイジングに使う)より前に、PyNCCLのP2P通信リソースを実際に1回動かして確保しておく処理が入っている。ウォームアップに使うサイズはコード内コメントによれば「DeepSeek-V4-Proでdpあたりのバッチサイズ120・vocab_size 129,280と仮定すると、各peerが送るチャンクは120×129,280/8≒1.849MB」という計算に基づく、1peerあたり4MB(4 << 20)の固定値。
  • 本題と無関係な修正が1件含まれている:同じPRのdeep_gemm_wrapper/compile_utils.pyには、事前コンパイル時のGPUメモリ空き容量取得をgpu_id=0固定からtorch.cuda.current_device()に修正し、しきい値をmax_m > 4096からmax_m > 2048に変更する差分も含まれていた。TP LMHead all-to-allの説明文には出てこない変更で、マルチGPU環境での事前コンパイル時のメモリ見積もりに関する別件の修正とみられる。

PR本文は「dp-tp混在並列でこの最適化をサポートするのは困難」と述べているのみで、今後対応が広がる予定があるのかは、今回確認したPR本文・リリースノートのいずれにも記載がなかった。またDeepSeek-V4-Pro・B200以外のモデル・GPU世代の組み合わせでの実測値も、公式リリースノート・PR本文・diffの範囲では見つからなかった。PR本文中の画像(プロファイルのスクリーンショット4点)はGitHub上の添付画像であり、この記事では本文のテキスト部分とコード差分のみを根拠にしている。SGLangを自分の環境・B200上で動かして実測値を再現したわけでもない。

関連記事

シェア: ポスト はてブ

出典・参照資料

AIニュースの解説を動画でも

YouTubeでは注目ニュースの背景を解説し、Xでは新着記事をお知らせしています。

コメント

まだコメントはありません。最初のコメントを書いてみませんか?

AIについて聞きたいことはありますか?

質問箱で無料で受け付けています。回答は公開され、他の方の参考にもなります。

質問箱を見る →

新しい記事をメールで受け取る

AIの新しい発表を、出典付きで整理して届けます。

関連記事