GPUの演算器は速いのに、LLMはメモリの読み書きと待ち時間で遅くなる。FlashAttention、GQA・MQA・MLA、スライディング窓、PagedAttention、連続バッチング、MoE、投機的デコーディングが、それぞれ「何を減らしているのか」を、図を動かしながら整理します。『GPUとLLM』第5章（最終章）。
『GPUとLLM』の最終章です。第3章では、GPUの演算器はとても速いのに、LLMの推論、とくに1トークンずつ生成する decode では、HBM からデータを読む時間と待ち時間のほうが壁になることを見ました。第4章では、それでもモデルを載せるための量子化と並列化を扱いました。
この章では、大規模モデルの設計と運用がその壁をどう削っているかを見ていきます。図の4本のバーが削る対象です。HBMとの往復、KVキャッシュの量、1トークンあたりの計算、そして逐次の待ち。薄い部分が削った量のイメージで、右にそれぞれを削る手法を並べました。以降のページでは、タイトルの下の色付きの札で、いまどれを削っているかを示します。
まずは HBM との往復です。第2章で見た Attention は、Q と K の積でスコアの表 S を作り、softmax で確率の表 P にしてから V を掛けます。S と P はどちらも T×T（T は文脈の長さ）の大きさです。素直に実装すると、図の下の段の赤い正方形のように、この表をいったん HBM に書き出し、softmax のためにまた読み、P を書いて、V を掛けるためにまた読みます。
T＝8,192 なら、表のマスは約6,710万個。2バイトで持つと1ヘッド・1層だけで約134MB になり、これをヘッドと層の数だけ往復させます。掛け算そのものより読み書きに時間がかかる、第3章で見た「メモリ律速」の典型です。上の段の共有メモリは速いのですが、SM ごとに数百KBしかないので、表全体は入りません。
FlashAttention（Dao ら, 2022）は、表全体を作るのをやめます。Q・K・V を行ごとの小さなブロックに分け、図のように Q のブロックと、K・V のブロックを1つずつ共有メモリに読み込みます。その組み合わせで計算できるのは、表 S のうちの小さなタイル1枚ぶんだけです。
タイルのスコアはその場で softmax の途中計算に使い、すぐ V のブロックと掛けて出力 O に足し込みます。用が済んだタイルは HBM に書かずに捨てます。下の段の S が白く抜けたのは、T×T の表がもうどこにも置かれないことを表しています。赤い枠は「いまどのタイルを計算しているか」の目印です。HBM に書き戻すのは、最後にできあがる O（T×d）だけです。
softmax は、行全体の最大値と和が分からないと計算できません。そこで行ごとに「ここまでの最大値 m」と「ここまでの指数の和 ℓ」だけを持ち歩き、新しいタイルが来るたびに右の式で更新します。最大値が大きくなったら、それまでの和と出力に exp(m − m新) を掛けて縮め直すのがポイントです。最後に O を ℓ で割れば、表全体で計算したのと同じ答えになり、近似ではありません。
図では赤いタイルが右へ流れ、K・V のブロックも下へ進んでいます。持ち歩くのは O と m、ℓ だけなので、余分なメモリは O(T²) から O(T) になります。論文は、HBM の読み書きが Θ(T·d ＋ T²) から Θ(T²·d²÷M)（M は共有メモリの大きさ）に減ることを示しています。
次は KVキャッシュです。第2章で見たとおり、その大きさは上の掛け算で決まります。2 は K と V の2つぶん、L は層数、H_kv は K・V のヘッド数、d はヘッドの次元、T は文脈の長さ、最後が1つの数のバイト数です。ここからの手法は、それぞれこのどれかの項を小さくします。
ふつうの MHA（Multi-Head Attention）では、Q のヘッド1本ごとに専用の K・V ヘッドがあり、中段の模式図では8本ずつが1対1でつながっています。下のバーは、Qヘッド128本・ヘッド次元128・60層（DeepSeek-V2 と同じ大きさ）で、1層・1トークンあたりに保存する数です。MHA なら 2×128×128＝32,768個。BF16 なら1トークン約3.9MB、12.8万トークンの文脈で約515GB になります。
GQA（Grouped-Query Attention, Ainslie ら 2023）は、Q のヘッドをいくつかの組に分け、組ごとに1組の K・V を共有させます。図では、薄いオレンジの枠で囲んだ 4本ずつの Q が1組の K・V を使うように、下の段の箱が寄り集まりました。式の中で小さくなるのは、オレンジに光った H_kv の項だけです。
バーは K・V ヘッドを8本にした場合で、2×8×128＝2,048個、MHA の16分の1です。12.8万トークンでも約32GB に収まります。Q のヘッド数はそのままなので、問いかけの多様さは残ります。公開されている設定では、Llama 2 70B が Q 64本に対して K・V 8本、Mistral 7B が Q 32本に対して K・V 8本という GQA の構成です。
共有をとことん進めると、すべての Q ヘッドが1組の K・V を使う形になります。これが MQA（Multi-Query Attention, Shazeer 2019）です。図では2つの組がさらに1つにまとまり、K・V の箱が1つだけ残りました。GQA は、組の数を変えることで MHA（組の数＝ヘッド数）と MQA（組の数＝1）のあいだをつなぐ方式とも言えます。
バーはほとんど見えないほど短くなり、1層・1トークンあたり256個、MHA の128分の1です。ただし GQA の論文は、MQA では品質が落ちることがあると指摘しています。すべての Q が同じ K・V を見るので、表現の幅が狭くなるからです。そのため、品質と節約のバランスをとった GQA を採用する公開モデルが多くあります。
DeepSeek-V2（2024）で報告された MLA（Multi-head Latent Attention）は、「ヘッド数 × ヘッド次元」の部分を別の方法で小さくします。K・V をヘッドごとに保存する代わりに、それらを作るもとになる低次元の潜在ベクトル c だけを保存します。図の紫の帯がそれで、DeepSeek-V2 では c の512次元に、位置情報（RoPE）用の64次元を足した576個です。各ヘッドの K・V は、c に行列を掛ければ復元できます。
576個は MHA の約57分の1で、論文は「グループ数2.25の GQA と同じ量」と書いています。MQA より大きいものの、MHA より良い性能だったと報告されている点が特徴です。60層・BF16 なら1トークン約69KB、12.8万トークンで約9GB です。
スライディングウィンドウ注意は、式の T の項に手をつけます。各層が見るのを直近の W トークンに限れば、それより古い K・V は捨てられるので、キャッシュは T ではなく min(T, W) に比例します。図の上の段では、オレンジの窓の中だけがキャッシュに残ります。
下のグラフでは、青い線（全部を見る）が文脈とともに伸び続けるのに対し、オレンジの線（窓を使う）は W を超えると増えなくなります。Mistral 7B は W＝4,096 で、3.2万トークンの列ではキャッシュが8分の1になると報告しています。層を重ねると情報は窓の外へも少しずつ伝わるので、遠い過去が完全に消えるわけではありません。Gemma 2 のように、窓の層と全体を見る層を1層おきに並べる設計もあります。
KVキャッシュには、大きさとは別の問題もあります。置き方です。リクエストが最終的に何トークン生成するかは、終わるまで分かりません。そこで従来の推論サーバーは、図のようにリクエストごとに最大長ぶんの連続した領域を前もって予約していました。濃い色のブロックが実際に使っている部分、薄い色が予約したまま空いている部分です。
こうなると、新しいリクエスト C は入る場所がなく待たされます。vLLM の論文（Kwon ら, SOSP 2023）は、既存のシステムで KVキャッシュ用のメモリのうち実際のトークンが入っていたのは20.4〜38.2%だけだったと測定しています。下のバーの赤い部分が、予約の空きと断片化で失われていた量です。第3章で見たとおり decode は一度にまとめる数が多いほど効率が上がるので、この無駄はそのまま速さの損になります。
PagedAttention は、OS の仮想メモリと同じ考え方を使います。KVキャッシュを固定サイズのブロック（vLLM の既定では16トークンぶん）に区切り、必要になったときに空いているブロックを1つずつ割り当てます。ブロックは連続している必要がなく、どこにあるかはブロック表で覚えておきます。
図では予約の帯が消え、A と B のブロックが飛び飛びの場所に移り、待っていた C も空いた場所に入りました。無駄になるのは各リクエストの最後の1ブロックの中だけ（半端な A3 がそれです）で、vLLM のブログでは4%未満と説明されています。より多くのリクエストを同時に載せられるので、論文では同じ待ち時間で2〜4倍のスループットが報告されています。
第3章で見たとおり、decode はメモリ律速なので、多くのリクエストをまとめて（バッチで）処理するほど、重みを1回読むあいだにできる仕事が増えます。問題は、リクエストごとに生成する長さがばらばらなことです。
図は、同時に処理できる4つの「席」を時間の流れに沿って描いたものです。リクエスト単位でバッチを組む素朴な方式では、短い返事が先に終わっても、一番長い依頼4が終わるまで次のバッチを始められません。赤い部分は、席が空いたまま GPU が働いていない時間です。下の段では、次の依頼たちが、席が空いているのに待たされています。この図では、席が埋まっているのは全体の約59%だけです。
Orca（Yu ら, OSDI 2022）は、スケジュールの単位をリクエスト全体から1ステップ（1トークンの生成）に変えました。毎ステップ、終わったリクエストを外し、待っているリクエストを空いた席に入れます。図では赤い空き時間が消え、依頼5〜7がすぐ席に入りました。論文は GPT-3 175B で、同じ待ち時間のまま最大36.9倍のスループットを報告しています。
下の図は、さらに一歩進んだ運用です。prefill は計算律速、decode はメモリ律速と性質が違うので、同じ GPU で混ぜると互いの足を引っぱります。そこで2つを別の GPU 群に分け、KVキャッシュを受け渡す方法が提案されています（DistServe など。DeepSeek-V3 の技術報告も、2つの段階を分けて配置しています）。
3つ目は、1トークンあたりの計算です。MoE（Mixture of Experts, 混合エキスパート）では、層の中の FFN を、専門家と呼ぶ小さな FFN の集まりに置き換えます。トークンが来ると、ルーターが専門家それぞれに点数をつけ、上位 k 人だけにそのトークンを計算させます。出力は、選ばれた専門家の出力を点数で重みづけして足したもの（Σ）です。
図は Mixtral 8x7B と同じ「8人から2人」の模式図で、点数のバーが長い2人だけが紫に光っています。誰が選ばれるかは、トークンごと・層ごとに変わります。もっと大きな実例には、DeepSeek-V3（256人から8人＋全員が通る共通の1人）や Muse Spark Open（256人から8人、共通の専門家なし）があります。
MoE には大きさを表す数が2つあります。全体の重みの数（総パラメータ）と、1トークンの計算で実際に使う重みの数（活性パラメータ）です。Mixtral 8x7B では、それぞれ約47B と約13B と報告されています。図では専門家が横一列に並んで上のバーになり、選ばれた2人が下の段へ降りてきました。
重み1個につき掛け算と足し算を1回ずつするので、1トークンの計算 ≈ 2 × 活性パラメータ FLOPs です。一方、誰が呼ばれるかは前もって分からないので、メモリには全員分を置く必要があり、ここで第4章の低精度化が効きます。DeepSeek-V3 は FP8 での学習を大規模に検証したと報告し、gpt-oss-120b は MoE の重みを4ビット（MXFP4）にして80GB の GPU 1枚に収めています。
最後は逐次の待ちです。decode では、1トークン進むたびに大きなモデルの重みを HBM から読み直します（第3章）。次のトークンは前のトークンが決まらないと作れないので、この読み込みを1つずつ順番に待つしかありません。
投機的デコーディング（Leviathan ら, ICML 2023 など）は、小さくて速い下書きモデルに、先の k トークンを書かせます。図では k＝4 で、緑の4つが下書きです。下書きモデルは軽いので、4回呼んでもすぐ終わります。そのあいだ大きなモデルはまだ何もしていません。下の段のように、ふつうなら大きなモデルの重みを5回読むところを、うまくいけば1回の読み込みで済ませるのが狙いです。
大きなモデルは、下書きの位置を1回の順伝播でまとめて計算します。これは prefill と同じ形で、重みを1回読むあいだに複数トークンぶん計算できるので、第3章でいう算術強度が上がります。図では3つが受理され、4つ目は却下されて、大きなモデル自身の答え「です」が入りました。受理の規則は、出力の分布が大きなモデル単独のときと変わらないように作られています。
下書き1つが確率 α で受理されるとすると、1回の検証で進むトークン数の期待値は (1 − α^(k+1)) ÷ (1 − α) です。全部受理されれば k＋1 個進みます。下のバーは k＝4 の場合で、α が高いほど上限に近づきます。論文は T5-XXL で、出力を変えずに2〜3倍速くなったと報告しています。
第1章では、GPU が数千のコアと Tensor Core、そして共有メモリから HBM へと続くメモリの階層でできていることを見ました。第2章では、LLM の計算のほとんどが行列積であること、第3章では、その行列積でも decode のように同じデータを使い回せない場面では、演算器ではなくメモリの帯域が速さを決めることを見ました。
第4章は、1枚に載らないモデルを量子化と並列化で載せる工夫、そしてこの第5章は、HBM との往復・KVキャッシュ・1トークンの計算・逐次の待ちを削る工夫でした。図の5つの小さな絵を左から右へたどると、フロンティアモデルの工夫の多くが「演算器を待たせないように、データの置き方と動かし方を設計する」ことだと分かります。新しい手法に出会ったら、まず「これは何を減らしているのか」を問うと、位置づけが見えやすくなります。
