第2章 — KVキャッシュという課題
LLM Primer VI: AIシステムのスケーリング を章ごとに紹介していくウォークスルー、第2回。重みが手に入れる機会を得る前に、あらゆるサービングシステムのVRAMを食い尽くすデータ構造に名前をつける章です。
なぜこの章があるのか
第1章の自己回帰ループは、記憶することで二次的な仕事を避けています。デコードの各ステップでモデルは過去の全トークンに注意を向ける必要があり、それを毎回ゼロからやり直すことは補完長にわたってプリフィル全体を反復すること — N倍の乗数 — を意味します。KVキャッシュはその算術を節約するワークスペースです。各層でトークンごとのKとV射影を保存し、以降のトークンは自分のKとVだけを計算して過去の分を読み戻す。キャッシュのおかげでデコードはおおよそO(N²)ではなくO(N)に留まる。そしてキャッシュは、私が見てきたあらゆる本番システムで、単一で最もVRAMを食う消費者でもあります。第2章はそのサイズを支配する算式、それを縮めるアーキテクチャの変種、そしてナイーブな割当を壊す断片化問題を歩きます。
2.1 キャッシュサイズを支配する算式
1つの系列のKVメモリは 2 × L × H_kv × D × S × バイト。2つのテンソル(KとV)、L層、H_kv個のキー/バリューヘッド、ヘッド次元D、系列長S、精度あたりのバイト数。艦隊レベルの数字にはバッチサイズを掛ける。全ての軸は乗算的です。Llama-3-70Bを L=80、H_kv=8(GQA)、D=128、S=8,192、BF16で回すと、1系列につきおよそ40 GBのキャッシュが載る — バッチ以前で。バッチ32ではキャッシュだけで約1.3 TBに近づき、重み自体の数倍、H100のVRAMの数倍。Llama-2-70B、最後の主要なMHAモデルは H_kv=8ではなく64で、キャッシュは同じ構成でも8倍になります。この算式は物珍しさではなく、1台のチップが同時に何ユーザーを保持できるかを決める数字です。箱の上の他のメモリ予算 — アクティベーション、フレームワークオーバーヘッド、一時バッファ — は全てその周りに収める必要があります。
2.2 MHA、GQA、MQAはキャッシュサイズ設計の選択である
H_kvの項こそ、現代のアーキテクチャが工夫して小さくしてきた項です。マルチヘッドアテンション(MHA)はクエリヘッドごとに1つのKとVヘッドを持ちます。キャッシュは太く、品質は最高、各ヘッドが何に注意を向けるかを専門化できる。マルチクエリアテンション(MQA)は共有された1組のKとVヘッドに潰します。キャッシュはH倍縮みますが、ヘッドは専門化を失い、長いコンテキストで測定可能な品質低下が出ます。2023年に登場し、いまはLlama-3、Mistral、Mixtral、Qwen、DeepSeekの主流となったグループ化クエリアテンション(GQA)は、クエリヘッドをGグループに分割してKとVを共有します。典型的な G = H/8 は8倍のキャッシュ削減を、評価スイートが確実には検知できない品質コストで実現します。GQAは無料ではなく — K/V射影のパラメータ数も減らし、その予算はモデルの他所に再配分されます — 経験的に、その再配分はほとんど害を及ぼしません。MLA(マルチラテントアテンション、DeepSeek-V2)は低ランクの潜在キャッシュでさらに押し進めます。研究方向としては有望ですが、2026年の本番配備はGQAが優勢です。
2.3 ナイーブな割当は予算の大半を無駄にする
KVメモリを割り当てる素直な方法 — 系列ごとに最大可能長のサイズの連続スラブを予約する — は、実トラフィックとの接触で失敗します。ほとんどのリクエストは最大より遥かに手前で完了し、予約された末尾は死んだメモリになる。Llama-3-70Bで32K上限のとき、1,000トークンの補完は31スラブ分のKVキャッシュを無駄にします。バッチ32、上限8K、平均長800だと、任意の瞬間に予約されたKVキャッシュのおよそ90パーセントが未使用です。さらに悪いことに、レイアウトは柔軟性がない。新規到着は他の系列の未使用な末尾を使えません。末尾はすでにコミットされているからです。同時実行数は平均トラフィックが支えられる量ではなく、最悪ケースの予算が許す量に潰れる。問題は内部断片化 — 1960年代にオペレーティングシステムがページングで解決したのと同じ失敗モードです。第8章でPagedAttentionがその解を持ち込む様子を見ます。
第2章が敷いたもの
キャッシュの算式とその軸は、以降の章が1台のGPUに何人のユーザーが載るかを推論するたびに再浮上します。第3章は、サービング用GPUを選ぶ際に何故FLOPsよりVRAM容量が重要かを説明するのに使います。第7章は、連続バッチングがモデルではなくKVキャッシュに束縛される理由を説明するのに使います。第8章はPagedAttention、H2Oエビクション、プレフィックスキャッシュを正当化するのに使います — この3つの技法の目的はまさに、キャッシュを予約されたバイトのスラブではなくページ化された仮想メモリのように振る舞わせることです。算式は、以降の本書のメモリ議論の文法です。
次回 — 第3章: 生成AI向けのデータセンターGPU。 重みとKVキャッシュの両方を保持しなければならないシリコンを、スペックシートではなく機構で読み解く。
