🤖 machine learning

Accelerating Sparse Transformer Inference on GPU

本論文は、効率的なマルチヘッドアテンションマッピングのための解析的モデリングと、演算フュージョンを動的に最適化するための2段階探索戦略を活用してスパースTransformer推論を高速化するGPUフレームワークSTOFを提示し、MHA計算およびエンドツーエンド推論においてそれぞれ最大1.6倍および1.4倍の高速化を実現する。

原著者： Wenhao Dai, Haodong Deng, Mengfei Rong, Xinyu Yang, Hongyu Liu, Fangxin Liu, Hailong Yang, Qianwen Cao, Qingxiao Sun

公開日 2026-05-20

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

CC BY 4.0

原著者： Wenhao Dai, Haodong Deng, Mengfei Rong, Xinyu Yang, Hongyu Liu, Fangxin Liu, Hailong Yang, Qianwen Cao, Qingxiao Sun

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

巨大な図書館（大規模言語モデル）に質問への答えを見つけるために、膨大な数の本を読もうとしていると想像してください。この図書館はトランスフォーマーと呼ばれる部屋に整理されており、各部屋には、あなたの質問に関連する特定の文を見つけるために数千ページをスキャンしなければならない司書（マルチヘッドアテンション機構）がいます。

問題は、多くの質問において、ページの大部分は無関係だということです。司書は、無関係なページや不要なページをめくるのに時間を浪費してしまいます。ここでスパース性が役立ちます。これは、無関係なページに「読む必要あり」のシールを貼るようなものです。

しかし、現在の司書たち（既存のソフトウェア）は、これらのシールの使い方が下手です。彼らは「読む必要あり」のページを無視して通り過ぎたり、シールが奇妙でランダムなパターンで配置されていると混乱したりします。さらに、図書館には要約やフォーマットなど、通常は別々に行われる他のタスクもあり、タスク間の移動時間が追加されます。

ここで登場するのが、研究者たちが提案した新しいシステムSTOFです。STOF は、これらの「スパースな」図書館に特化して設計された、超効率的で賢い図書館管理システムだと考えてください。その仕組みを簡単な部分に分けて説明します。

1. 賢い司書（統一された MHA カーネル）

研究者たちは、異なる「読む必要あり」のパターンには異なる戦略が必要だと気づきました。

問題点: 一部のシールのパターンは、整然とした行（スライディングウィンドウのようなもの）ですが、他のパターンはランダムに散らばっています（宝くじのチケットのようなもの）。古いシステムは、これらすべてに「万能」な手法を使おうとしていましたが、それは遅いものでした。
STOF の解決策: STOF は、作業に最適なツールを選ぶ賢い司書のように機能します。
- シールが整然とした小さなクラスターにある場合、司書は**「行単位（Row-wise）」**のアプローチを採用します。つまり、本の行全体を一度に掴み、素早くスキャンします。
- シールが散らばっている場合や図書館が巨大な場合は、**「ブロック単位（Block-wise）」**のアプローチを使用します。本を小さく管理しやすいチャンクに分割し、有効なシールがついている特定のチャンクだけを開きます。
結果: 無視するのではなく「読む必要あり」のページを完全にスキップすることで、司書ははるかに速く作業できます。

2. 組立ライン（オペレータ融合）

通常の図書館では、司書は読み終えた後、要約するために別の机へ歩き、答えをフォーマットするためにさらに別の机へ歩きます。この移動（メモリとプロセッサ間のデータ移動）は遅いです。

問題点: 現在のシステムは、しばしば単純なタスクのみを結合します。複雑な計算などの重労働を別の手順に任せており、これにより渋滞が発生します。
STOF の解決策: STOF は、カスタム組立ラインを構築します。それはプロセス全体を見て、「これらの手順を結合できるか？」と問います。
- 単に 2 つの単純なタスクを接着するのではなく、複雑な数学的タスクとフォーマットタスクを結合する完璧な方法を考え出します。
- 「検索エンジン」を使用して、これらのタスクを結合するさまざまな方法（異なる組立ラインのレイアウトを試すようなもの）を試み、読み進めている図書館のサイズに最も適した、最も速く動く組み合わせを見つけます。

3. オートパイロット（階層的検索）

すべての本のサイズと質問の種類に対して、完璧な組立ラインを手動で設計することはできません。組み合わせが多すぎます。

STOF の解決策: STOF には、その場で学習するオートパイロットが備わっています。
- フェーズ 1（地図）: 図書館の構造を見て、「読む必要あり」のシールの位置の概略図を描きます。
- フェーズ 2（最適化）: 2 段階の検索を実行します。まず、組立ラインの境界を広げて、どこまで進めることができるかを確認します。次に、以前の試みがどの程度うまくいったかに基づいて、作業者（パラメータ）の速度を微調整します。
- 何が機能したかを記憶（キャッシュ）し、同じ遅いアイデアを再テストする時間を無駄にしません。

結果：どれほど速くなったか？

研究者たちは、人気のある AI モデル（BERT、GPT、LLaMA など）を使用して、高性能なグラフィックカード（GPU）上で STOF をテストしました。

速度: 既存の最良の方法と比較して、STOF はコアとなる読書タスク（MHA）を最大1.6 倍高速化しました。
全体速度: 質問への回答というプロセス全体（エンドツーエンド）を見た場合、最大1.4 倍高速化しました。
巨大な図書館: 図書館が大きいほど（テキストシーケンスが長いほど）、STOF の効果は顕著でした。なぜなら、無意味な作業をこれほど多くスキップできたからです。

まとめ

STOFとは、AI が読む必要のないページを読む時間を浪費するのを止め、机間を行き来するのを防ぐシステムだと考えてください。それは、不要なものをスキップし、有用な手順を 1 つの滑らかで高速な動作に結合するための、賢く適応的な戦略を使用します。これにより、特に長く複雑なテキストを扱う場合、AI モデルの動作が大幅に高速化されます。

技術概要：GPU 上での疎性トランスフォーマー推論の高速化

問題定義

大規模言語モデル（LLM）はトランスフォーマーアーキテクチャに大きく依存しており、マルチヘッドアテンション（MHA）が中核的な計算コンポーネントである。計算量を削減するために、因果的、スライディングウィンドウ、ランダムなどのスパース性を誘発するマスク層が広く用いられているが、これらは GPU 上で顕著なパフォーマンス上の課題をもたらす：

断片化された計算：スパースなマスクは断片化されたデータアクセスパターンを生み出し、メモリ帯域幅のボトルネックを悪化させ、GPU リソースの効率的な利用を妨げる。
柔軟性の欠如した最適化：既存の高性能 MHA 実装（例：FlashAttention）は、通常、因果的やスライディングウィンドウのような特定の連続的なマスクパターンのみをサポートする。離散的または非構造化のスパース性（例：Bigbird、ランダムアテンション）には対応できず、計算後にスコア行列をリセットする非効率的な手法に依存せざるを得ない。
静的な演算子融合：深層学習フレームワークにおける現在の演算子融合スキームは、主に静的またはルール駆動型である。多様なアプリケーションシナリオ、可変のシーケンス長、異なるモデルハイパーパラメータに適応できない。さらに、融合演算子の最適な設定は、個別にチューニングされた演算子の単純な合計とは大きく異なることが多く、逐次的なチューニングでは最適化が不十分となる。

手法：STOF フレームワーク

著者は、柔軟なマスクサポートと適応的演算子融合を統合することで、GPU 上でのスパース性トランスフォーマー推論を高速化するフレームワーク「STOF」を提案する。STOF は 2 つの主要モジュールから構成される：

1. 統合 MHA モジュール

このモジュールは、多様なマスクパターンにわたって効率的にスパースアテンションを計算するという課題に対処する。

2 段階ストレージ形式：任意のマスクパターン（構造化および非構造化の両方）を処理するために、STOF はブロック圧縮疎行列（BSR）とビットマップを組み合わせた新しいストレージ形式を導入する。
- 外部タイル（OT）：グローバルにスキップされるブロックを表す。「フル」OT（すべての要素が有効）と「パート」OT（有効/無効が混在）はポインタ配列を通じて追跡される。
- 内部タイル（IT）：ブロック内の要素分布を表す。「パート」OT 内の有効要素は、64 ビットビットマップ（8x8 ブロックあたり 1 つの uint64）でエンコードされ、完全な疎行列を具体化することなく正確なマスクを可能にする。
デュアルカーネル戦略：STOF は分析モデルに基づき、2 つの異なるカーネルタイプを実装する。
- 行方向カーネル：短いシーケンスを持つ高度にスパースなシナリオ向けに最適化されている。クエリ（ $Q$ ）を行単位にスライスし、ワープレベルのシャッフルを利用し、データ局所性を最大化するためにワープ間同期を排除する。
- ブロック方向カーネル：より広範なシナリオ向けの汎用カーネルである。 $Q$ 、 $K$ 、 $V$ をサブブロックに分割し、共有メモリ（SMEM）と非同期データコピーを活用して、メモリ転送と計算をオーバーラップさせる。
カーネル選択：分析モデルが、マスクのスパース性とシーケンス長に基づいて行方向カーネルかブロック方向カーネルのどちらを使用するかを決定し、最も効率的な実行パスを確保する。

2. 演算子融合モジュール

このモジュールは、MHA レヤーを超えた適応的融合の必要性に対処する。

融合スキームの変換：STOF は、ハッシュエンコーディングを使用して融合スキームを数値表現に変換する。計算グラフ内の隣接する演算子は、融合境界を表すバイナリ配列にマッピングされ、その後、特定のコンパイルテンプレート（Triton および TileLang を通じて実装）にデコードされる。
階層型検索エンジン：静的ルールに代わり、STOF は最適な融合境界とカーネルパラメータを決定するために 2 段階の検索メカニズムを採用する。
1. 初期化：ニューラルハッシュと事前定義されたルールを使用して、頻出するサブグラフパターンを特定し、初期融合スキームを確立する。
2. 2 段階チューニング：
  - ステージ 1（融合拡張）：報酬ベースのアルゴリズム（拡張、奪取、競合）を使用して融合境界を反復的に拡張し、過去の結果のキャッシュに対してパフォーマンス向上を検証する。
  - ステージ 2（パラメータサンプリング）：境界が固定された後、特定の融合スキームに対する最適な設定を見つけるために、カーネルパラメータ（ブロックサイズ、パイプラインステージなど）をサンプリングする。

主要な貢献

包括的な分析：著者は、多様なマスクパターンと推論設定の影響を分析し、MHA 計算とダウンストリーム演算子融合の両方における最適化の機会を明らかにした。
統合 MHA モジュール：固有のストレージ形式（BSR + ビットマップ）と最適化を実装し、分析選択モデルによって管理される、行方向およびブロック方向カーネルを実装する新しいモジュール。
適応的演算子融合：数値デコーディングを介して融合スキームをコンパイルテンプレートに変換し、2 段階検索エンジンを利用して動的に最適な設定を決定するモジュール。
STOF フレームワーク：GPU 上で柔軟なマスクパターンと最適な演算子融合を可能にする完全な推論システム。

実験結果

著者は、PyTorch Native、FlashAttention-2（FA2）、FlexAttention、ByteTransformer、MCFuser、Bolt などの最先端のベースラインに対して、NVIDIA RTX 4090（Ada）および A100（Ampere）GPU 上で STOF を評価した。

MHA パフォーマンス：STOF は、MHA 計算において、既存の最良の MHA 実装（例：FlexAttention）に対して最大1.6 倍の高速化を達成した。因果的、スライディングウィンドウ、Longformer、Bigbird などのさまざまなマスクパターンと、最大 4,096 のシーケンス長にわたって一貫した優位性を示した。
エンドツーエンド推論：BERT、GPT-2、LLaMA、T5、ViT などのモデルにおいて、エンドツーエンド推論で最大1.4 倍の高速化を達成した。
チューニング効率：STOF のチューニングプロセスは、競合他社（MCFuser および Bolt）よりも大幅に高速であり、報酬ベースのサンプリングとキャッシングメカニズムにより、大規模シナリオでは6.7 倍から 6.9 倍高速であった。
スケーラビリティ：超長シーケンス（最大 16k）を用いた実験において、STOF は顕著なパフォーマンス向上（PyTorch Compile に対して最大 16.8 倍）を示し、メモリ効率も向上させた。ベースラインが失敗したシーケンス長においても、Out-of-Memory（OOM）エラーを回避した。

意義と主張

本論文は、STOF がスパース性トランスフォーマーシナリオに対して、柔軟なマスクパターンと多様な演算子融合スキームを同時に可能にする最初のシステムであると主張している。その意義は、手動チューニングされたスパース固有カーネルと、自動的な密演算子融合の間のギャップを埋める点にある。カスタム MHA カーネルと生成コンパイルテンプレートを統合する完全なスタックを提供することで、STOF は、静的またはルール駆動型のアプローチでは以前は達成不可能だった、より広範な最適化の機会を確立する。著者は、彼らのアプローチが新興の LLM アーキテクチャ（例：Mixture-of-Experts）や新しい GPU 世代にも拡張可能であり、大規模モデル推論における持続的なパフォーマンス向上への道筋を提供すると強調している。