はじめに：アテンションがボトルネックになった時代
AIエージェントや長文コンテキストのワークロードが一般化するにつれ、モデルが処理するコンテキスト長は増え続けています。問題は、この過程でアテンション演算が推論時間に占める割合が急激に大きくなる点です。以前はアテンションは「ただ存在するレイヤー」でしたが、今やアテンションをどう設計するかがモデルの推論性能を左右します。
本記事はNVIDIA Developer Blogの分析を基に、アテンション性能を決定する4つの主要変数—Group Size(G)、Head Dimension(Hsz)、Sequence Length、Tensor Parallelism—を整理し、すぐに適用可能な4つのガイドラインに集約します。モデルを直接学習させる方だけでなく、LLMサービングの性能チューニングを行うバックエンドエンジニアやMLOpsエンジニアにとっても実務的な示唆が得られる内容です。
💡 前提: すべての分析はGEMM Shapeの算術とFP8ベースのPrefill/Decodeカーネル測定値に基づいています。感覚ではなく数値で語ります。
PrefillとDecodeは完全に異なる問題である
アテンション最適化の第一歩は、PrefillとDecodeを分離して捉えることです。これらを一括りに考える方は多いですが、両フェーズはボトルネックが正反対です。
| 区分 | Prefill | Decode |
|---|---|---|
| Query長 | 全入力 (ISLトークン) | 1トークン |
| KVコンテキスト | プロンプト (ISLトークン) | 全KVキャッシュ (KVSLトークン) |
| Attention GEMM-M | ISL × G (大) | G (小) |
| 主要ボトルネック | 演算 (matmul + softmax) | HBM帯域 (メモリ) |
Prefillはプロンプト全体を並列処理するためCompute-bound(演算集約型)、Decodeはトークンを1つずつ生成するためMemory-bound(メモリ集約型)です。ここでSpeculative Decodingを用いるとDecodeのGEMM-Mが大きくなり、Compute-bound側へ移行できます。
FlashAttentionが行う処理
FlashAttentionはアテンション行列全体をメモリに実体化せず、Q/K/VタイルをHBMからSRAMへストリーミングしながら3ステップを1パスで融合します。
- BMM1: QueryとKeyのスコアリング
- Online Softmax: running max/sumによる正規化
- BMM2: Valueへの重み付け
BMMはTensor Coreで、softmaxの指数演算はSFU(Special-Function Unit)で実行されます。この構造が後述する算術強度分析の基盤となります。
Group Size(G)：Decodeの主要レバー
Group Sizeは1つのKV headを共有するQuery headの数です。MHAはG=1、GQAはG=4/8/16、MQAはG=QHです。
# 算術強度(Arithmetic Intensity)の公式 (FP8 KVキャッシュ想定)
# Prefill: AI = 2 * G * ISL / (G + 1) → Gが大きくなると2*ISLに収束
# Decode: AI = 2 * G * KVSL / (G + KVSL) ≈ 2 * G (KVSL >> Gのとき)
def arithmetic_intensity_prefill(G, ISL):
return 2 * G * ISL / (G + 1)
def arithmetic_intensity_decode(G, KVSL):
return 2 * G * KVSL / (G + KVSL)
# 例: Gを8から16に上げてもPrefillは6%未満の改善
print(arithmetic_intensity_prefill(8, 32768)) # ~15,990
print(arithmetic_intensity_prefill(16, 32768)) # ~15,998
# 一方DecodeはGにほぼ線形比例
print(arithmetic_intensity_decode(8, 32768)) # ~15.99
print(arithmetic_intensity_decode(16, 32768)) # ~31.99
重要ポイント: PrefillはGに鈍感、DecodeはGに敏感です。Gを2倍にするとDecodeの算術強度も2倍になり、メモリトラフィックが減ってGPU利用率が向上します。実際、G=1(MHA)からG=8(GQA)に変えるとDecodeランタイムが約2倍高速化されます。
Head Dimension(Hsz)：128または256が正解
Hszは算術強度自体を変えません。FLOPsとBytesが同時に2倍になるからです。それでもランタイムが増えるのは、カーネルが3種類の処理を異なる比率で実行するためです。
- Matmul: Hszに比例するが、128アラインメント単位で処理されるためHsz=64は128タイルのコストをそのまま支払います。
- Memory(KV state): HBMが128バイト単位で動くため、Hszが128の倍数のときに効率的です。
- Softmax: Hszと無関係！Attention-score行列(Query × Key)にのみ依存します。
# Softmax演算量はHszと無関係
def softmax_ops_prefill(PB, QH, ISL):
return PB * QH * ISL ** 2
def softmax_ops_decode(DB, QH, KVSL):
return DB * QH * KVSL
# Hszを64→128に上げてもsoftmaxコストは不変
print(softmax_ops_prefill(1, 32, 8192)) # 同一
PrefillではHszが大きくなるほどmatmulの比重が増し、softmaxが相対的に隠れます。DecodeではHszが大きくなるとKVバイトが増えてランタイムが増加します。したがって128と256がスイートスポットです。Hsz=64は128タイルのコストを支払い、Hsz≥512はTMEM容量の限界に達します。
🔗 コンテキストエンジニアリングの観点からモデル設計を見たい方はバックグラウンドコーディングエージェントのためのコンテンツエンジニアリングガイドも併せてご覧ください。
Sequence Length：Prefillは二乗、Decodeは線形
シーケンス長はPrefillとDecodeに全く異なる形で影響します。
- Prefill: ISL²の処理 (全トークンが全トークンを参照)。KVトラフィックはISLに比例。→ 算術強度がISLに線形増加 → Compute-boundを維持
- Decode: 各ステップでKVSL全体を読み込み。算術強度は2×G付近で固定 → Memory-boundを維持
つまり、ISLを2倍にするとPrefillは約4倍遅くなり、KVSLを2倍にするとDecodeは2倍遅くなります。短いシーケンスではsetup/post-processingのオーバーヘッドによりこの比率より緩やかですが、長くなると理論値に収束します。
Tensor Parallelism：KHが戦略を決定する
TPはアテンションheadをGPUに分散させます。各GPUはQH/TP個のQuery headとKH/TP個のKV headを持ちます。問題はTP > KHになる瞬間です。1つのグループが複数ランクに跨がり、KV headの複製が発生してメモリと帯域のオーバーヘッドだけが増えます。
# TPの制約条件
def validate_tp_strategy(num_kv_heads, tp_size):
if tp_size > num_kv_heads:
print(f"⚠️ TP={tp_size} > KH={num_kv_heads}: KV head複製が発生")
print(" → ADP(Attention Data Parallelism)またはKVP(KV Parallelism)を検討")
return False
print(f"✅ TP={tp_size} ≤ KH={num_kv_heads}: 正常シャーディング")
return True
validate_tp_strategy(2, 4) # ⚠️ 複製発生
validate_tp_strategy(8, 4) # ✅ OK
Nemotron 3のようにKV headが2つしかないモデルはTPをすぐに使い切ります。その場合はWide EP(ADP + EP)やHelix Parallelism(KVP + EP)に切り替える必要があります。TensorRT-LLMはこの両戦略をサポートしています。
⚠️ 注意点と限界
- Speculative Decodingの落とし穴: ドラフトトークン数(D)が十分大きくなるとDecodeがCompute-boundへ移行し、既存の最適化前提が崩れる可能性があります。
- Prefix Caching環境: エージェント/マルチターンアプリで一般的ですが、短いISL + 長いprefix cacheの組み合わせではPrefillがDecodeのように振る舞います。この場合、GとHszのチューニング戦略を再設計する必要があります。
- 測定環境への依存性: 本記事の数値はFP8 KVキャッシュと特定のGPUアーキテクチャ基準です。BF16や他のアーキテクチャでは最適点が移動する可能性があります。
- モデル再学習コスト: GやHszを変更するにはモデルの再学習が必要です。サービング段階で変更できるのはTP戦略とKVキャッシュ圧縮程度です。
🔗 データレイヤーの一貫性とコンテキストアーキテクチャ設計に興味がある方はAIエージェントの沈黙の毒データレイヤー一貫性、今こそコンテキストアーキテクトになるときをご参照ください。
まとめ：4つのガイドライン
本記事のすべての分析は、以下の4つの実践チェックリストに集約されます。
| # | ガイドライン | 核心的な理由 |
|---|---|---|
| 1 | GはDecode基準で高く | PrefillはGに鈍感、Decode算術強度 ≈ 2×G |
| 2 | Hsz = 128または256 | GPUタイル/128B転送アラインメント + TMEM予算 |
| 3 | 有効KV状態を削減 | PrefillはISL²、DecodeはKVSL線形増加 |
| 4 | KHが並列化戦略を決定 | TP ≤ KHを維持、不足時はWide EP/Helix |
日本開発エコシステムにおける適用文脈
国内のLLMサービング環境は大きく2つに分かれます。大手企業/クラウド事業者はH100/H200クラスタを保有しTensorRT-LLMやvLLMベースの最適化を直接実施できますが、スタートアップや中小SIerはA100やL40Sといった限られたGPUでマルチテナントサービングを行うケースが多いです。後者の場合、Gを高めたGQAモデル(例: Llama 3系)を選択するだけで、同一ハードウェア上でのDecodeスループットを大きく引き上げられます。また、国内の金融/医療ドメインのように長文コンテキスト + オンプレミスの組み合わせが必要な場面では、KVキャッシュ圧縮やsliding-window attentionの適用が選択肢ではなく必須要件になりつつあります。
次のステップ学習方向
- Sparse Attention: 本記事はdense attention基準です。Sparse attentionが登場するとボトルネック構造が再び変わります。
- vLLM/TensorRT-LLM実習: 実際のサービングフレームワークでGQA/MQA設定がどう反映されるか、コードレベルで確認してみてください。
- Rooflineモデルを自分で描く: 自身がサービングするモデルの算術強度を計算し、ridge pointに対してどこに位置するかを把握すれば、どの最適化が有効かが見えてきます。
アテンションはもはや「実装ディテール」ではなくアーキテクチャ設計の意思決定です。モデルを作る人もサービングする人も、この4軸を理解していれば同じGPUでより多くのトークンをより速く生成できます。🚀
根拠資料: NVIDIA Developer Blog - Co-Designing AI Model Attention for Fast, Interactive Long-Context Inference
併せて読みたい記事