8月1日、NVIDIA Developerが「Co-Designing AI Model Attention for Fast, Interactive Long-Context Inference」と題した記事を公開した。LLMの長コンテキスト推論においてAttention機構をGPUアーキテクチャと共同設計することで、スループットとインタラクティブ性を両立する手法について、数式と実測データの両面から詳しく分析した内容だ。

エージェント型ワークロードの普及とともに、LLMのコンテキスト長は急速に伸びている。それに伴い、推論コスト全体に占めるAttentionの割合が増大しており、「Attentionをどう実装するか」以上に「Attentionをどう設計するか」が推論性能を左右するようになってきた。

本記事はNVIDIAによるAIモデル共同設計（Co-Design）シリーズの一環で、グループサイズ（G）・ヘッド次元（Hsz）・シーケンス長の3つのパラメータが、PrefillとDecodeそれぞれのフェーズにどう影響するかを分析している。最終的に、モデル開発者向けの4つの実践的ガイドラインに集約される。

PrefillとDecodeは、まったく別の問題だ

まずおさえておくべき前提として、Attention推論には性質の異なる2フェーズが存在する。

Prefillはプロンプト全体を並列処理するため大規模な行列積となり、演算律速になる。Decodeは1トークンずつ生成しながらKVキャッシュ全体をHBM（High Bandwidth Memory）から読み出すため、メモリ律速になる。この非対称性が、以降のすべての分析の出発点だ。

なお、プレフィックスキャッシュを使うマルチターンのアプリでは、新しいターンのISLが短くても長い過去のKVキャッシュを参照するため、PrefillがDecodeのように振る舞うケースもある。

グループサイズGは「Decode専用の最適化レバー」

最も実践的なインパクトがあるのが、グループサイズ G（1つのKVヘッドを共有するクエリヘッドの数）だ。Attentionのバリアントとの対応は以下のとおりだ。

• **MHA（Multi-Head Attention）**：G=1。クエリヘッドとKVヘッドが1対1
• **GQA（Grouped-Query Attention）**：G=4/8/16など。複数のクエリヘッドがKVヘッドを共有
• **MQA（Multi-Query Attention）**：G=クエリヘッド数（QH）。全クエリヘッドが単一のKVヘッドを共有

GQA・MQAはKVキャッシュのサイズと帯域消費を削減するために広く採用されており、LLaMA 3・Mistral・Gemmaなど主要なオープンモデルの多くがGQAを採用している。

Decodeへの効果：Gを2倍にするとDecodeが約2倍速くなる

Decode時の演算強度はおよそ以下で表される：

Arithmetic Intensity ≈ 2 × G（KVSLがGより十分大きい場合）

この式は、Decode時に各KVヘッドのロードをG個のクエリヘッドで再利用できることから導かれる。KVSLがGに比べて十分大きいという前提のもとで、KVキャッシュのロードコストがボトルネックとなる場合の近似式だ。Gを1から8に増やすと演算強度は8倍になる。これはKVヘッド数が減ることでトークンあたりのKVロード量が減り、かつ各ロードをより多くのクエリヘッドに分散できるためだ。実測でもGを2倍にするたびにDecodeのランタイムが約1/2になっている（KVSL=128Kの場合）。

元記事ではNVIDIA Nemotron 3がGQAを採用しKVヘッド数を絞った設計を採用していることが紹介されており、Decodeの効率化が主な設計動機として挙げられている。

Prefillへの効果：ほぼ無影響

一方、Prefill時の演算強度はおよそ以下で表される：

Arithmetic Intensity ≈ 2 × ISL × G/(G+1)

この式はGが大きくなるにつれ2 × ISLに収束する性質を持つ。ISL=32Kの条件でGを8から16に倍増させても演算強度の改善は6%未満。実測でも、GをMHA（G=1）からMQA（G=64）まで変えてもPrefillランタイムの変化は1%以内だ。

ガイドライン1：GはDecodeの効率を基準に、できるだけ高く設定せよ。Prefillのランタイムはほぼ変わらない。

ヘッド次元Hszは128か256が最適解

ヘッド次元（Hsz）とは、各Attentionヘッドが扱うベクトルの次元数を指す。全体の隠れ次元をヘッド数で割った値に相当し、モデルの表現力と計算コストに影響する基本的なハイパーパラメータだ。

Hszは演算強度自体には影響しない。FLOPsとバイト数がHszと同じ割合で増えるため、両者の比（＝演算強度）は変わらないからだ。

ただし実際のランタイムはHszとともに増加する。その理由は、Attentionカーネルが3種類の処理で構成されており、それぞれHszへの依存度が異なるためだ。

• 行列積（BMM）：Hszに比例して増加するが、GPUタイルサイズの整数倍でのみ効率的に動作する
• KVキャッシュのメモリアクセス：同様にHsz倍増でメモリ転送量が増加
• Softmax：アテンションスコア行列（クエリ数×キー数）に作用するため、Hszに無依存

特に重要なのは次の2点だ：

• Hsz=64は内部的に128幅タイルを使うため、実質128の料金を払っている
• Hsz≥512はTensor Memory（TMEM）容量の限界に近づく

ガイドライン2：HszはGPUタイル幅に揃え、128か256を選べ。

シーケンス長のコストは非対称——Prefillは二乗、Decodeは線形

シーケンス長の増加がコストに与える影響は、フェーズによって大きく異なる。

• Prefill：全トークンが全トークンに注目するため演算量はISL²に比例。ISLが2倍になるとランタイムはおよそ4倍
• Decode：1ステップごとにKVキャッシュ全体を読むため、KVSLが2倍になるとランタイムはおよそ2倍

Decodeが常にメモリ律速であるのに対し、Prefillは演算律速のまま長いコンテキストでも保たれる。この非対称性から、実用上の指針が導かれる。

KVキャッシュの実効サイズを削減するアプローチとしては、以下のような手法が挙げられる。

• KVキャッシュ圧縮：量子化や低ランク近似でKVキャッシュのメモリフットプリントを削減する
• スパースアテンション：全トークンではなく重要なトークンのみに注目する
• スライディングウィンドウアテンション：直近の一定範囲のトークンのみをKVとして保持する（Mistralなどが採用）
• ハイブリッド設計：グローバルなKV状態を持つレイヤーを一部に限定し、残りはローカルな注意機構で代替する（Nemotron 3が採用するアプローチ）

ガイドライン3：KVキャッシュの実効サイズを減らせ。PrefillはISL²、DecodeはKVSLで線形にコストが伸びるという非対称性を踏まえ、長コンテキストではDecodeのメモリ帯域圧迫がボトルネックになりやすい。KVキャッシュ圧縮・スパースアテンション・スライディングウィンドウアテンション・ハイブリッド設計などを組み合わせ、KVの実効サイズを抑えることが重要だ。

テンソル並列化の上限はKVヘッド数で決まる

テンソル並列化（TP）はAttentionヘッドをGPU間で分割する手法だ。各GPUはQH/TP個のクエリヘッドとKH/TP個のKVヘッドを担当する。

TP > KHになると、1グループのクエリヘッドが複数のGPUにまたがり、各GPUがKVヘッドのコピーを持つ必要が生じる。これはKVの重複保持を招き、メモリと帯域幅を無駄遣いするだけだ。

KVヘッドが少ないモデル（MQAのKH=1、GQAでKHを大幅に絞った設計など）はTPがすぐ上限に達する。この場合、以下の並列化戦略の組み合わせが有効だ。

• Attention Data Parallelism（ADP）：Attentionレイヤーをデータ並列で処理する手法。KVヘッドが少なくてもスケールアウトできる
• KV Parallelism（KVP）：KVキャッシュ自体を複数GPUに分散して保持・処理する手法
• Expert Parallelism（EP）：MoEモデルにおいてFFN（エキスパート層）を並列分散する手法

TensorRT-LLMはこれらをWide EP（ADP+EP）やHelix Parallelism（KVP+EP）として実装しており、KVヘッドが少ないモデルでも大規模GPU環境でのスループットを確保できる。

ガイドライン4：TP≤KHを守れ。KVヘッドが少ないモデルはADPかKVPとEPの組み合わせに切り替えよ。

4つのガイドラインのまとめ

詳細はCo-Designing AI Model Attention for Fast, Interactive Long-Context Inferenceを参照していただきたい。