🤖 machine learning

Tile-Level Activation Overlap for Efficient LLM Inference

本論文は、中間テンソルの実体化オーバーヘッドを排除するために、タイルレベルでSwiGLU活性化関数と行列乗算を融合させた2つの特化したCUTLASSベースのSM90カーネルを紹介するものであり、これによりNVIDIA H100 GPUにおいて最大2.47倍の高速化と79.5%のピーク利用率を達成し、効率性と数値精度の両面でPyTorchおよびcuBLASを凌駕している。

原著者： Abhinav Jangda, Tyler Sorensen, Sebastian Burckhardt, Jianlan YE, Chaoyin Li, Atul Gupta

公開日 2026-07-07

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Abhinav Jangda, Tyler Sorensen, Sebastian Burckhardt, Jianlan YE, Chaoyin Li, Atul Gupta

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、複雑なロボット（チャットボットを動かしているような大規模言語モデル）を製造する高速工場を運営していると想像してください。各ロボットを組み立てるために、あなたの工場には特定の組立ライン、すなわちMLP（ロボットの脳の核心部分）があります。

長い間、この組立ラインには大きな非効率性が存在していました。以下に、この問題と、この論文で提示されている解決策を分かりやすく説明します。

問題点：「仲介役」によるボトルネック

現代のロボット工場では、ロボットをより賢くするためにSwiGLUと呼ばれる特定のステップが使用されます。SwiGLUを、2つの別々の計算を必要とする品質チェックだと考えてください。

計算A： ロボットの現在の状態を測定する。
計算B： ロボットの潜在能力を測定する。
接着（Glue）： これら2つの測定結果を組み合わせて、最終的な結果を得る。

旧来の方法（ボトルネック）：
標準的な工場設定（PyTorchなどで使用されているもの）では、計算Aを行った後、作業員は結果を廊下にある巨大なホワイトボード（高帯域幅メモリ）に書き留めなければなりません。その後、彼らはホワイトボードまで歩いて行き、それを読み取り、計算Bを行い、その結果を2つ目のホワイトボードに書き込み、最後に「接着」ステップを行うために両方のホワイトボードを読み取るために戻ってこなければなりません。

この論文によれば、小型のロボット（小規模なAIモデル）の場合、この「ホワイトボードへ歩いて行く」作業が、総時間の**30%から37%**を占めていることが分かりました。これは、シェフが料理をする時間よりも、スパイスを一つ取りにパントリーまで歩いている時間の方が長いようなものです。

解決策：2つの新しい「スーパーキッチン」

著者たちは、ホワイトボードを完全に排除した、カスタム設計された2つの新しいキッチン（カーネル）を構築しました。結果を書き留めてから戻ってくるのではなく、作業員は材料を手に持ち（レジスタ内）、すべてを一連の連続した動作で行います。

彼らは、工場の規模に合わせて2つの異なる戦略を作成しました。

1. 「ピンポン」キッチン（Kernel-1）

仕組み： リレーレースを想像してください。一人の作業員が次の材料を取りに行っている間に、別の作業員はすでに現在のバッチの混合作業を行っています。
トリック： 彼らは「ピンポン」方式のスケジュールを使用しています。マシンが2つ目の材料のセットを取り込んでいる間に、作業員はその同じ時間を使って、最初のセットに対して「接着」の計算を行います。
最適： 巨大なロボット（大規模モデル）を作る工場、または多くのロボットを同時に稼働させている（大規模バッチ）場合に適しています。これは、大きな機械をフル稼働させるのに十分な数の作業員がいる、大規模な組立ラインのようなものです。

2. 「インターリーブ（交互配置）」キッチン（Kernel-2）

仕組み： 箱詰め作業を想像してください。赤いアイテムをすべて詰めてから次に青いアイテムを詰めるのではなく、赤、青、赤と、完璧に交互に詰めていきます。
トリック： 計算のための「レシピ」である2つの重み行列を、開始前に混ぜ合わせます。これにより、作業員は「赤」と「青」の材料を同時に掴み、それらを即座に一緒に処理することができます。
最適： 小型のロボットを作っている工場、または一度に動かしているロボットが少ない（小規模バッチ）場合に適しています。この方法は非常に効率的であり、誰もが立ち往生して待っていることがないよう、工場のフロアを常に作業員で満たし続けます。

結果：スピードと精度

著者たちは、これらの新しいキッチンを、NVIDIA H100チップ（利用可能な最も強力なAIプロセッサ）上で、様々なサイズのロボット（0.5Bの極小モデルから72Bの巨大モデルまで）に対してテストしました。

劇的なスピードアップ： 小型ロボットの場合、新しいキッチンは標準的な方法よりも2.47倍高速でした。これは、10分のタスクを4分に短縮することに相当します。
ボトルネックの転換： 旧システムは「メモリ・バウンド（メモリ制限）」（ホワイトボードへ歩いて行くことに時間を費やしすぎている状態）でした。新しいシステムは「コンピュート・バウンド（計算制限）」（実際に計算を行うことに全力を注いでいる状態）になりました。彼らはチップの最大理論速度の**79.5%**に到達しました。
コンパイラでは不可能： 著者たちは、標準的な「オートパイロット」ソフトウェア（PyTorchのtorch.compile）を使用して、自動的に修正を試みました。しかし、それは失敗しました。オートパイロットは、彼らのカスタムキッチンよりも3倍から7倍遅かったのです。これは、この特定の問題においては、人間の専門家による手書きの解決策が必要であることを証明しています。コンピュータはまだ自分自身でこれを見つけ出すことはできません。
より高い精度： 驚くべきことに、新しいカスタムキッチンは標準的な方法よりも精度も高かったです。標準的な方法では、結果の4.5%から11%に小さな計算誤差がありましたが、新しいキッチンではエラーはゼロでした。

まとめ

この論文は、工場のフロアの動き方、具体的には中間的なメモを読み書きするために作業員が頻繁にホワイトボードへ歩いて行くことを止めることで、AIモデルを大幅に高速化できることを示しています。特に、エッジデバイス（スマートフォンやノートパソコンなど）で使用される小型モデルにおいて効果的です。彼らは、標準的なソフトウェアツールではこの効率性を再現できないことを証明し、これらの結果を達成するには専門化された手書きのコードが必要であることを明らかにしました。

技術要約：効率的なLLM推論のためのタイルレベル・アクティベーション・オーバーラップ

問題提起

Qwen-2.5、LLaMA、Mistral、Gemmaを含む現代の大規模言語モデル（LLM）は、主にSwiGLU活性化関数を利用している。このゲート付きMLP構造は、 $Y = \text{SiLU}(A \times W_1) \odot (A \times W_2)$ を計算し、2つの独立した行列積（GeMM）と、それに続く要素ごとのゲート付き活性化を必要とする。

標準的な実装では、このプロセスにおいて、GeMMと活性化のステージ間で2つの完全な中間テンソル（$Gate $および$ Up$）が高帯域幅メモリ（HBM）上に実体化される。NVIDIA H100 GPUでのプロファイリングによれば、この中間テンソルの実体化とそれに伴う活性化計算は、MLP実行時間の**9〜37%**を消費している。このオーバーヘッドは、より小さなモデル（例：Qwen-2.5 0.5B）で最も深刻であり、そこではSwiGLUがMLP時間の30%以上を占めている。量子化（FP8、INT4）によってテンソルコアの算術密度が増加するにつれ、これらのメモリ帯域に制約される操作の相対的なコストが増大し、重大なボトルネックを生み出している。

既存のコンパイラ・インフラストラクチャ、特にPyTorchのtorch.compileは、この問題に対処できていない。これは、異なる重み行列を持つ2つの別々のGeMM間での融合（fusion）が、グラフレベルの融合パスにおける根本的な制限により不可能であるためである。実験の結果、torch.compileは、このパターンに対してEager PyTorchの性能の34〜94%しか達成できず、明示的な融合ヒントを与えても改善はわずかであることが示されている。

手法

著者らは、SwiGLUをタイルレベルでGeMMに融合させ、中間的なHBMへの実体化を排除する、2つの補完的なCUTLASSベースのSM90カーネルを提示している。両方のカーネルは、H100（SM90）で利用可能なTensor Memory Accelerator（TMA）とワープ特化型スケジューリングを活用している。

1. カーネル1：Pingpongスケジュール・オーバーラップによる同期SwiGLU

メカニズム: このカーネルは、データをロードするプロデューサー（TMA経由）と、MMAを実行するコンシューマーにワープを分割するCUTLASS Pingpongワープ特化型スケジュールを利用する。
戦略: $[M, N]$ スレッドブロックグリッドを作成する。コンシューマーフェーズにおいて、Gateタイルが計算された後、Upウェイトタイルが到着するまでの間に、コンシューマーワープがアイドル状態になる短いウィンドウが存在する。カーネルはこのウィンドウを利用して、レジスタに留まっているGateアキュムレータに対してSwish活性化（ $\text{SiLU}(Gate)$ ）を計算する。
オーバーラップ: Swish計算は、プロデューサーによるUpウェイトタイルのTMAロードと時間的にオーバーラップする。最終的な要素ごとの乗算（ $\text{SiLU}(Gate) \odot Up$ ）およびストアはエピローグで行われる。
最適化: この設計は、 $[M, N]$ グリッドがGPUのStreaming Multiprocessor（SM）を十分に飽和させる大きなバッチサイズ（ $M \ge 2048$ ）向けに最適化されている。

2. カーネル2：カスタムEpilogue Visitor TreeによるインターリーブSwiGLU

メカニズム: このカーネルは、PairMulStoreと呼ばれるカスタムのEpilogue Visitor Tree (EVT) ノードを導入している。
戦略: 重み行列 $W_1$ と $W_2$ を、列を交互に配置した単一の融合行列 $W_{fused} \in \mathbb{R}^{K \times 2N}$ （偶数列がGate、奇数列がUp）へと結合する。カーネルは $[M, 2N, K]$ 上で標準的なGeMMを実行し、 $[M, 2N]$ スレッドブロックグリッドを作成する。
融合: PairMulStore ノードは、各スレッドブロック内の $[M, 2N]$ アキュムレータ上で動作する。これは隣接する列ペア $(D[m, 2n], D[m, 2n+1])$ を読み取り、 $\text{SiLU}(D[m, 2n]) \odot D[m, 2n+1]$ を適用し、結果をストアする。
オーバーラップ: この設計は、SwiGLU計算をタイルストアとインターリーブさせる。現在のタイルの算術演算は、前のタイルの非同期TMAストアとオーバーラップし、計算とストアのオーバーラップを実現する。
最適化: スレッドブロック数の2倍の増加は、小さなバッチサイズ（ $M \le 2048$ ）におけるSM占有率を大幅に向上させるため、ほとんどの設定において推奨されるデフォルトとなる。

主な貢献

タイルレベルの融合: SwiGLUをタイルレベルでGeMMに融合させ、推論ステップあたりの4つの中間HBM操作（GateとUpの読み書き）を排除した最初の実装である。
デュアルカーネル戦略: 大規模バッチ用のカーネル1（逐次的なGate→Swish→Up）と、小規模バッチ用のカーネル2（インターリーブされた重みとカスタムEVT）を用いた補完的なアプローチにより、占有率とスケジューリングの制約にそれぞれ対処している。
数値的優位性: 融合カーネルは、FP32リファレンスと比較してミスマッチゼロ（相対誤差1%を超える要素が0%）を達成している。これに対し、ベースラインのcuBLAS実装は、タイル間での積極的なBF16部分累積により、4.5〜11%のミスマッチを示す。
コンパイラの限界の検証: torch.compileがこの融合を再現できず、手作りのカーネルよりも3〜7倍遅いままであることを示す経験的証拠。

結果

NVIDIA H100上のQwen-2.5モデル（0.5Bから72B）で評価：

スピードアップ: カーネルは、PyTorch Eager（cuBLAS + 個別のSwiGLU）に対して最大2.47倍のスピードアップを達成した。20個のベンチマーク構成のうち、カーネル2が13個で勝利し、カーネル1が7個（主に高バッチサイズの大型モデル）で勝利した。
レジームシフト: ルーフライン分析により、融合がワークロードをメモリ帯域制約から計算量制約へとシフトさせることが確認された。小規模モデルでは、算術強度が最大247%増加し、動作点がルーフラインのリッジを横切る。
利用率: ピークBF16利用率はH100上で**79.5%**に達し、小規模モデルにおけるベースラインの利用率をほぼ倍増させている。
効率性: 融合効率（実際のスピードアップ vs 理論的なメモリ節約）は、小規模モデルにおいて100%に近い。カーネル2については、計算とストアのオーバーラップの恩恵により、効率が100%を超える構成もある。
スケーラビリティ: スピードアップはモデルのサイズとともに減少する（0.5Bで2.47倍 vs 高バッチサイズでの72Bで約1.0倍）。これは、GeMM計算が総時間に占める割合が支配的になり、中間的なメモリトラフィックの相対的な割合が減少するためである。

意義

本論文は、現代のLLM推論、特にエッジ展開やモバイルサイズ（0.5B–14B）のモデルにおいて、ボトルネックが生の計算能力から中間テンソルの実体化によるメモリトラフィックへと移行していることを主張している。著者らは、手作りのハードウェア認識型タイルレベル設計が、現在のコンパイラ・インフラストラクチャの限界を克服するために必要であることを証明した。中間HBMアクセスを排除し、データ移動と活性化計算をオーバーラップさせることで、これらのカーネルは推論レイテンシを大幅に削減し、計算利用率を向上させる。これは、レイテンシとスループットが制約となるプロダクション・サービング・システムにとって極めて重要な最適化である。