> [!abstract] 概要(arXiv abstract の日本語訳) > テンソル並列を用いた大規模言語モデル(LLM)の分散推論は、NVLink のような高速 GPU インターコネクトを介していても 20% の通信オーバーヘッドを生じうる。これらのオーバーヘッドを緩和するため、計算をより小さなタスクに分解して通信とオーバーラップさせる複数の技術が提案されてきた。しかし、これらの技術のいずれも vLLM・SGLang・TensorRT-LLM のようなシステムにおけるテンソル並列サービングでデフォルトでは有効化されていない。これは、1 イテレーションあたりに処理されるトークン数が低レイテンシサービングを支えるために小さく保たれるのが通例であり、そのような小規模なワークロードを通信オーバーラップのために分解するとかえって性能が悪化するためである。さらに、通信自体が計算に使えたはずの多くの streaming multiprocessor(SM)を消費し、オーバーヘッドを増大させる。本稿では、トークン長がわずか 1024 という小ささでもテンソル並列モデル推論の効率的な計算-通信オーバーラップを実現する初のシステム TokenWeave を提示する。TokenWeave は、これまで見過ごされてきた演算である RMSNorm を重要な要素として特定し、新規の融合 AllReduce–RMSNorm カーネルを実装することで通信とともに最適化する。さらに、このカーネルは最新 GPU(Hopper・Blackwell 等)で利用可能な NVSHARP/Multimem 機能を活用し、8×H100 DGX システム上でわずか 2–8 の streaming multiprocessor(SM)のみを用いて通信と RMSNorm を効率的に同時実行する。評価の結果、複数のモデル・ワークロードにわたって最大 1.28 倍のレイテンシ高速化(baseline÷ours)、最大 1.19 倍の高いスループット(ours÷baseline)を達成した。複数の設定において、TokenWeave はすべての通信を除去した等価モデルよりも高い性能を発揮する。ソースコードは https://github.com/microsoft/tokenweave で公開されている。 ## 論文情報 - タイトル: TokenWeave: Efficient Compute-Communication Overlap for Distributed LLM Inference - 著者: Raja Gond・Nipun Kwatra・Ramachandran Ramjee(全員 [[Microsoft Research]] India) - 媒体: 9th MLSys Conference(MLSys 2026, Bellevue, WA, USA)採択。arXiv 初出は 2025-05-16、v5(掲載版)は 2026-05-01。 - arXiv ID: 2505.11329 (cs.DC) - コード: https://github.com/microsoft/tokenweave ## 概要 テンソル並列(TP)推論では、AllReduce 通信が NVLink 接続の 8×H100 DGX でも 9–23% のレイテンシオーバーヘッドを生む。TokenWeave は、(1) RMSNorm を AllReduce と融合した専用 CUDA カーネル、(2) wave 数を意識した 2 分割(smart-splitting)による計算-通信オーバーラップ、(3) NVSHARP/Multimem を利用した極小 SM 数での通信実行、の 3 点を組み合わせ、従来手法が対象にできなかった小トークン数(1K 前後)の推論イテレーションでも通信オーバーラップを実用化した。vLLM V1 に実装し、Llama-3.3-70B・Qwen2.5-72B・Mixtral-8x22B で評価している。 ## 問題設定 - **入力**: テンソル並列度 N の GPU 群上で、あるイテレーションに含まれるトークン数 T(prefill・decode・chunked-prefill 混在の hybrid バッチを含む)。 - **前提**: 各 Transformer 層は Attention Block・MLP Block それぞれの後に AllReduce と RMSNorm(+ residual 加算、以下 RMSNorm と略記)を実行する(Figure 7a)。低レイテンシサービングでは T は小さく保たれる(例: vLLM 0.8.5 のデフォルトチャンクサイズ 2048)。 - **課題**: 既存の compute-communication overlap 手法(fine-grained tile 分解・coarse-grained token 分解いずれも)は、通信を隠すために計算を細分化するが、GPU の wave quantization 効果により分割オーバーヘッドがボトルネックになる。この分割オーバーヘッドは問題サイズが小さいほど深刻化するため、既存手法は 8K+ トークンの大バッチでしか有効に機能しない(Table 1・Figure 6・Figure 9 が裏付ける)。 ## 提案手法 ### アーキテクチャ **Figure 7: TokenWeave の全体像(vanilla TP との対比)** ![[_attachments/arxiv-2505.11329/fig07-architecture-overview.png]] (Figure 7. (a) Vanilla Tensor Parallelism: すべての計算・通信操作が逐次実行される。(b) TokenWeave: 入力バッチを 2 分割(Split 0 / Split 1)し、AllReduce を RMSNorm と融合した Fused-AR RMSNorm カーネルとして扱う。片方の split の計算が、もう片方の split の Fused-AR RMSNorm(通信+正規化)とオーバーラップして実行される。独立した Compute Stream と Communication Stream が weave するように交互に実行を進める構成であることが図名 TokenWeave の由来。Source: Adapted from Figure 7.) - **Coarse-Grained Token-Splitting**: 入力バッチ(T トークン)を prefix-split(Ta トークン)と suffix-split(Tb トークン)の 2 分割にする。Attention 以外の全操作はトークン単位で独立に処理できるが、Attention は suffix 側のトークンが prefix 側に依存するため、chunked attention 実装(Sarathi-Serve 由来)を用いて prefix-split を先に処理する順序制約を課す。 - **Overlapped Execution**: Compute stream と Communication stream の 2 本の CUDA stream を用い、`torch.cuda.stream wait` による軽量同期でデータ依存を扱う。Split 0 の AllReduce+RMSNorm 実行中に Split 1 の Attention/FFN 計算を進め、以降も互い違い(weave)に処理する。 ### アルゴリズム/手法の詳細 **なぜ 2 分割か**: パイプラインを作り依存関係を回避するには最低 2 分割が必要であり、3 分割以上は分割オーバーヘッドを増やすだけで追加のオーバーラップ機会を提供しないため採用しない(§4)。 **Wave-Aware Smart-Splitting(§4.1.1)**: H100(132 SM)上で 300 CTA(Cooperative Thread Array)の GEMM を実行すると 2 full wave + 1 partial wave(計 3 wave 分の実行時間)になる。これを均等に 150 CTA ずつ 2 分割すると、各分割が「1 full wave(132 SM)+ 1 partial wave(18 SM)」の 2 wave を要し、合計 4 wave に増加してしまう(Figure 8 の équal split)。Smart-splitting は、両分割の合計 wave 数が元の未分割カーネルの wave 数を超えないよう split offset を調整する。上記の例では、片方を 132 CTA(ちょうど 1 full wave)、もう片方を 168 CTA(1 full wave + 1 partial wave)に分けることで合計 3 wave に抑える(Figure 8 の smart split)。CTA-SM 割り当ては GEMM 形状・カーネル実装・タイルサイズ等の決定論的関数だが、TokenWeave は高性能なクローズドソース cuBLAS カーネルを用いるため解析的に split offset を求めることができず、代わりに複数バッチサイズ・シーケンス長に対するオフラインプロファイリングスイープ(Algorithm 1)で最適な split offset を事前計算し、実行時に参照する。 **RMSNorm Reordering(§4.2)**: 通常の実装は AllReduce の後で各 GPU が独立に RMSNorm を計算するため、全 GPU が同一のトークン埋め込みに対して冗長に計算している。AllReduce を ReduceScatter → AllGather に分解し、ReduceScatter 完了直後(各 GPU が担当する 1/N シャードのみ保持する時点)で RMSNorm を実行すれば、計算量を GPU 数 N 分の 1 に削減できる。ただし Table 1 が示すように、この単純な reorder + fusion(Simple Fusion)は AllReduce を RS/AG に分割するオーバーヘッド(Figure 4)がRMSNorm 削減分の利得を相殺してしまい、512–8K トークン域ではむしろ性能が悪化する(相対性能 0.92–0.98)。 **Fused AllReduce–RMSNorm 実装(§4.3, Figure 18)**: NVSHARP/Multimem を活用し、ReduceScatter・RMSNorm・AllGather を単一 CUDA カーネルに融合する。各 GPU は multimem ld reduce add でスイッチ内リダクション済みの担当シャードを取得し、そのまま SM レジスタ上で分散(variance)を計算(通常必要な 1 回目の HBM 読み出しを省略)、正規化後の値を multimem st で直接 AllGather 用に書き出す(通常必要な最終 HBM 書き込みを省略)。Residual 加算も同カーネルに融合する。これにより Table 1 のとおり、逐次 AllReduce+RMSNorm 比で 1.34–1.39 倍の一貫した高速化を得る。この融合カーネルは 2–8 SM で近最適な性能に達し(Figure 10)、残る大多数の SM を計算に回せる。 ### 実装上の工夫 - vLLM V1(0.8.5)上に実装。PyTorch 2.6.0 + CUDA 12.4、Triton 3.2.0 + カスタム CUDA 拡張。PyTorch の SymmetricMemory API で NVLS(NVSHARP/Multimem)を利用し、NCCL 呼び出しを介さず Triton/CUDA カーネル内から直接メモリ操作する。 - Attention backend は FlashAttention-3。 - Selective Enabling(§A.3, Figure 3): トークン数が閾値未満の場合は splitting/overlap を無効化し、Fused AllReduce–RMSNorm カーネルのみを使う(分割オーバーヘッドが利得を上回るため)。閾値はオフラインプロファイリングにより Llama・Qwen(dense)で 1K、Mixtral(MoE)で 4K に設定。 ## 新規性 - 既存の fine-grained タイル分解手法(Wang et al. 2022・Flux・TileLink)は GEMM カーネル内部でのオーバーラップに限定され、Attention のような非 GEMM 演算はオーバーラップできない。加えて AllReduce を RS+AG に分割する構造的オーバーヘッド(Figure 4 で最大 50%超)と、タイル粒度の通信が大転送より非効率(Figure 6)という 2 つの問題を抱える。TokenWeave は融合カーネルで RS+AG 分割オーバーヘッドを解消し、Attention・FFN いずれもオーバーラップ対象にできる。 - NanoFlow(coarse-grained・カーネル粒度の nanobatch 分割)は大バッチ前提でオーバーヘッドが許容されるが、TokenWeave の smart-splitting は wave quantization を明示的に最小化することで、1K トークンという小規模バッチでもオーバーヘッドを実質無視できるレベルに抑える点が異なる(Figure 9)。 - RMSNorm を「これまで見過ごされてきた重要な演算」として初めて通信オーバーラップの対象に組み込んだ点(TileLink・Flux・NanoFlow はいずれも RMSNorm を通信と同時に最適化しない)。 ## 実験設定 - **ハードウェア**: 8×H100 NVIDIA DGX(NVSHARP 対応、128 CPU コア、800GB ホストメモリ)。追加で 4×H100(TP-4)、8×B200 DGX(192GB HBM、PyTorch 2.10.0、CUDA 13.0、vLLM 0.14.1、FlashInfer 0.5.3、Appendix C)でも評価。GPU クロックは TDP 周波数に固定して測定のばらつきを排除(B200 は共有サーバのため root 権限がなくクロック固定なし)。 - **モデル**: Llama-3.3-70B・Qwen2.5-72B(dense)、Mixtral-8x22B(MoE)、Qwen3-235B-A22B(Appendix B.2)。いずれも instruction-tuned 版。 - **比較対象**: vLLM-Default(標準 AllReduce)・vLLM-Multimem(NVSHARP/Multimem 最適化 AllReduce、主 baseline)・vLLM-nocomm(通信を完全除去した非正当だが性能上限の参照値)・TileLink(タイル中心の SoTA compute-communication overlap、Flux を上回るとされる)・NanoFlow(nanobatch 粒度のスケジューリング型手法。A100 専用実装を著者らが H100 向けに移植)。 - **ワークロード**: ShareGPT・arXiv(要約データセット由来)の実トレース、および固定 (入力長, 出力長) の合成トレース。chunked-prefills(デフォルトチャンクサイズ 2K、Mixtral は 4K)を使用。 ## 実験結果 ### 定量評価 - **エンドツーエンドスループット(Figure 11)**: ShareGPT・arXiv・固定長トレースにわたり、8×H100 上で dense モデルは約 1.19 倍、Mixtral は 1.10–1.11 倍のスループット改善。arXiv トレースでは 1.15 倍。 - **チャンクサイズ変化への頑健性(Figure 12)**: チャンクサイズ 1024–8192 の範囲で 1.14–1.26 倍のスループット改善を一貫して達成。 - **単一イテレーションレイテンシ(Figure 13)**: dense モデルは 1K トークンから 1.16–1.28 倍のレイテンシ改善。Mixtral は小トークン域(1K・2K)で smart-splitting 込みだと純オーバーヘッドが生じるため、この場合のみ融合カーネルのみを使い overlap を無効化(§A.3)する。 - **TileLink との比較(Figure 14, 単層比較)**: TileLink は 1K トークンでネット悪化・4K以上でようやく改善し 1.2 倍で頭打ちになるのに対し、TokenWeave は 1K トークンで既に 1.20 倍、最大 1.35 倍まで到達。 - **NanoFlow との比較(Figure 15)**: NanoFlow は 1.04–1.09 倍(NanoFlow 論文が報告する通信改善 1.07 倍と整合)にとどまるのに対し、TokenWeave は幅広いシナリオで約 1.19 倍。 - **融合カーネル単体のアブレーション(Table 1, Figure 16)**: TokenWeave-fuseonly(融合カーネルのみ、分割・オーバーラップなし)は 1.04–1.09 倍。フル TokenWeave はこれに加えてオーバーラップ由来の追加ゲインを得る。融合カーネル自体は逐次 AR+RMSNorm 比で全トークン長域(64〜32K)にわたり 1.34–1.39 倍。 - **smart-splitting のアブレーション(Figure 17)**: smart-splitting なし(equal-split)ではシーケンス長依存のジッター(wave quantization の有無による)が生じるが、smart-splitting はこのジッターをほぼ完全に除去し、一貫した高いゲインを提供する。 - **8×B200 での検証**: 融合カーネルは B200 でも少数 SM(Figure 25)で近最適、decode(Figure 26)・prefill(Figure 27)双方でレイテンシ改善を確認しており、H100 固有の結果でないことを示す。 ### アブレーション - Table 1: Simple Fusion(RS+RMSNorm+AG の単純な融合)は 512–8K トークンで逐次実装比 0.92–0.98 倍(悪化)。TokenWeave の Fused カーネルは同じトークン範囲で 1.34–1.39 倍。両者の差は HBM トラフィックの扱いに起因(§4.3)。 - Figure 17: equal-split と smart-split の比較により、wave quantization の抑制効果を定量的に確認。 ### 定性評価 - Figure 13・Figure 21・Figure 22 で、TokenWeave が「通信を完全に除去した」非正当な参照実装 vLLM-nocomm に近い、あるいはそれを上回る性能に到達する設定が複数示されている。これは通信オーバーラップに加え RMSNorm の重複計算除去自体が正味の性能改善に寄与するため。 ## 考察 - 4K トークン以上では、TokenWeave は通信オーバーヘッドの回収に加え、融合カーネルによる RMSNorm 最適化の分だけ「通信ゼロ」の理論上限を上回る(§1, Figure 2)。これは分割・オーバーラップと融合カーネルという 2 つの独立した最適化軸が相加的に効くためである。 - Mixtral(MoE)での効果が dense モデルより小さいのは、そもそもの通信オーバーヘッド比率が低い(Figure 1)ことに加え、MLP 計算がエキスパートに分散されメモリバウンドになりやすく、粗粒度の 2 分割がさらなるオーバーヘッドを生みやすいため(§5.2.2)。 - disaggregated serving(prefill/decode 分離)設定では、decode-only の小バッチには融合カーネル単体が有効、prefill-only の大バッチにはフル TokenWeave が有効という住み分けが成立する(§1 の contribution 4)。 ## 強み / 弱点・課題 - **強み**: (1) 既存手法が対象外としていた 1K トークン級の小規模イテレーションで通信オーバーラップを実用化した初のシステムである点。(2) RMSNorm という見落とされがちな演算を通信最適化の対象に組み込んだ点。(3) vLLM V1 という実運用サービングスタックに統合し、TileLink・NanoFlow のような研究プロトタイプ比較でも end-to-end 評価が可能な点。(4) H100・B200 双方で検証し将来世代への一般化可能性を示した点。 - **弱点・限界**: (1) 融合 AllReduce–RMSNorm カーネルは NVSHARP/Multimem(Hopper・Blackwell 世代以降)に強く依存し、これらの機能を持たない GPU(A100 以前や他ベンダー)には適用できない。(2) smart-splitting のオフセットは解析的に求まらずオフラインプロファイリングに依存するため、新しいモデル形状・GPU 構成ごとに事前のプロファイリングコストが発生する。(3) 現在の実装は BFloat16 のみサポート(Figure 18 注記)。(4) Mixtral のような MoE モデルでは小トークン域でオーバーラップがかえって不利になるケースがあり、モデルアーキテクチャ依存の閾値チューニング(§A.3)が必要になる。