LLM推論のVRAM計算 — 重みより先にKVキャッシュが溢れます

はじめに — 「このモデル、うちのGPUに載りますか」

この質問に「動かしてみないと分かりません」と答えると、インフラ予算の会議で言うことがなくなります。幸い、推論メモリは学習と違ってほぼすべて算数で予測できます。オプティマイザ状態もグラディエントもないからです。

ところが、たいていの計算は途中で止まります。「7Bモデルはfp16で14GBだから24GBのカードに余裕で入る」までは合っています。その次にバッチを32に上げてコンテキストを4千トークンに伸ばした瞬間、メモリ不足で落ちます。抜けている項がKVキャッシュで、この項は重みと違ってユーザートラフィックの形に応じて育ちます。

この記事では推論VRAMを四つの塊(重み、KVキャッシュ、活性化、フレームワークオーバーヘッド)に分け、それぞれの計算方法を整理します。数字はすべて自分で検算できるように式で示すので、ご自身のモデルのconfigファイルを開いたまま読んでいただければと思います。

重みのメモリ — 掛け算一回で終わる部分

重みは正直です。パラメータ数にパラメータあたりのバイト数を掛ければ終わりです。

BYTES_PER_PARAM = {
    "fp32": 4.0,
    "fp16": 2.0,   # bf16も同じ
    "int8": 1.0,
    "fp8": 1.0,
    "int4": 0.5,   # 4bit、実際にはスケール/ゼロポイントが少し加わる
}

GIB = 1024 ** 3

def weight_gib(params_billions: float, dtype: str = "fp16") -> float:
    return params_billions * 1e9 * BYTES_PER_PARAM[dtype] / GIB

for n in (7, 13, 70):
    print(n, [round(weight_gib(n, d), 1) for d in ("fp16", "int8", "int4")])
# 7  [13.0, 6.5, 3.3]
# 13 [24.2, 12.1, 6.1]
# 70 [130.4, 65.2, 32.6]

ここで単位を一度整理しておきます。70億×2バイトは140億バイトで、これを10の9乗で割ると14GB、2の30乗で割ると13.0GiBです。カードのスペックシートにある「80GB」はたいてい10進表記で、nvidia-smiが見せる数字は2進表記なので、この7パーセントの差のせいで計算がぎりぎり合わないことが起きます。以下の計算はすべてGiB基準に統一します。

4bitの項目の数字は理想的な値です。実際の4bitチェックポイントはグループごとにスケールとゼロポイントをfp16で持つため、グループサイズが128ならパラメータあたりおよそ0.5ビットが追加で乗ります。パラメータあたり0.5バイトではなく0.53バイト程度で見積もるほうが安全です。さらに埋め込み層と出力層は量子化せずに残す実装が多く、語彙サイズが大きいモデルほどこの差が広がります。

KVキャッシュ — 実際に溢れる場所

自己回帰デコーディングは、トークンを一つ作るたびに以前のトークンすべてに対するアテンションを計算します。すでに計算したキーとバリューを捨てると毎ステップで全シーケンスを計算し直すことになるので、すべて保持します。それがKVキャッシュです。

トークン一つが占めるバイト数は次のとおりです。

def kv_bytes_per_token(layers: int, kv_heads: int, head_dim: int,
                       bytes_per_elem: float = 2.0) -> float:
    # 2はKとVの二つを保存するため
    return 2 * layers * kv_heads * head_dim * bytes_per_elem

# 7BクラスのMHA構成: 32レイヤー、32 KVヘッド、ヘッド次元128
print(kv_bytes_per_token(32, 32, 128) / 1024, "KiB")   # 512.0 KiB
# 70BクラスのGQA構成: 80レイヤー、8 KVヘッド、ヘッド次元128
print(kv_bytes_per_token(80, 8, 128) / 1024, "KiB")    # 320.0 KiB

トークン一つで512KiB。この数字がぴんと来なければ、こう見ていただければと思います。7BクラスのMHAモデルでコンテキスト4096トークンのリクエスト一つは、KVキャッシュに2GiBを使います。同時リクエストが32件なら64GiBです。重み13GiBの五倍です。

この表は縦ではなく横に読むべきものです。重みは固定値で、最後の列はトラフィックに比例します。コンテキストを4Kから32Kに伸ばすと最後の列だけが8倍になります。長いコンテキストを謳うサービスが実際には同時処理数を大きく下げて運用している理由がこれです。

注意点が一つあります。KVキャッシュはアクティブなシーケンス数ではなく生きているトークンの総量に比例します。リクエスト100件がそれぞれ500トークンなら、リクエスト5件がそれぞれ10,000トークンの場合と同じメモリを使います。キャパシティ計画は秒間リクエスト数ではなく同時トークン数で立てる必要があります。

GQAとMQA — KVヘッドだけを減らす構造変更

上の表で7BのMHA行とGQA行の差がちょうど4倍なのは、KVヘッドが32個から8個に減ったからです。クエリヘッドは依然として32個で、4個ずつまとめて一つのKVヘッドを共有します。

# クエリヘッドはそのまま、KVヘッドだけを減らす
n_heads, head_dim, layers = 32, 128, 32

for kv_heads, name in [(32, "MHA"), (8, "GQA-8"), (1, "MQA")]:
    per_token = kv_bytes_per_token(layers, kv_heads, head_dim)
    ratio = n_heads // kv_heads
    print(f"{name:6} kv_heads={kv_heads:2}  {per_token/1024:6.1f} KiB/token  ({ratio}배 절감)")
# MHA    kv_heads=32   512.0 KiB/token  (1배 절감)
# GQA-8  kv_heads= 8   128.0 KiB/token  (4배 절감)
# MQA    kv_heads= 1    16.0 KiB/token  (32배 절감)

削減比率はクエリヘッド数をKVヘッド数で割った値と正確に一致します。これは近似ではなく定義上そうなります。

品質面では、GQAはおおむね安い取引と評価されています。原論文は少数のKVヘッドでもMHAに近い品質を保つと報告しており、その後の多くの公開モデルが8個前後のKVヘッドを採用しました。一方MQAはKVヘッドが一つしかないため、品質低下が観測されたという報告のほうが多くなっています。ただしこれは事前学習段階のアーキテクチャ選択なので、すでにデプロイされたモデルに対して変えられる値ではありません。皆さんが選べるのは「GQAを使うモデルを選ぶ」までです。

モデルのconfig.jsonでnum_key_value_headsがnum_attention_headsより小さければGQAです。計算時にこの二つを取り違えると結果が4倍ずれます。

活性化、オーバーヘッド、そして全体の計算機

残りの二つの塊は状況によって大きさが変わるため、正確な公式を示すのが難しい部分です。代わりに、どの変数に比例するのかが分かれば十分です。

活性化メモリはデコーディング段階ではほとんど問題になりません。ステップごとにバッチあたり一トークンずつしか流れないからです。問題はプリフィルです。プロンプト全体が一度に通るので、中間テンソルがバッチ×プロンプト長に比例します。

def prefill_activation_gib(batch, prompt_len, hidden, intermediate,
                           live_buffers=4, bytes_per_elem=2.0):
    """プリフィルのピーク活性化の粗い下限。live_buffersはフレームワークごとに違うので測定が必要。"""
    per_token = (hidden + intermediate) * bytes_per_elem
    return batch * prompt_len * per_token * live_buffers / GIB

# 7Bクラス: hidden 4096, intermediate 11008
print(round(prefill_activation_gib(8, 4096, 4096, 11008), 2), "GiB")  # 3.69 GiB

live_buffers=4は私が任意に置いた定数です。実際に同時に生きているバッファ数はカーネル融合の程度とメモリプランナの実装によって変わるので、この値は必ず測定で補正する必要があります。確かなのはバッチとプロンプト長の積に比例するという事実だけで、だからこそチャンク単位のプリフィル(長いプロンプトを複数の断片に分けて入れる方式)がこの項を定数に押さえてくれます。

フレームワークオーバーヘッドはCUDAコンテキスト、cuBLASワークスペース、通信バッファ、アロケータの余裕分を合わせたものです。プロセスあたり1GiB前後を見ておけばだいたい合います。そこにアロケータの断片化を見込んで全体に10〜15パーセントの余裕を掛けるのが実務的です。

これで全部を合わせた計算機を作ります。

def total_vram_gib(params_b, weight_dtype, layers, kv_heads, head_dim,
                   seq_len, batch, kv_bytes=2.0, overhead=1.15, fixed_gib=1.0):
    w = weight_gib(params_b, weight_dtype)
    kv = kv_bytes_per_token(layers, kv_heads, head_dim, kv_bytes) * seq_len * batch / GIB
    return (w + kv) * overhead + fixed_gib

# 70B、4bit重み、GQA-8、8Kコンテキスト、同時16
print(round(total_vram_gib(70, "int4", 80, 8, 128, 8192, 16), 1), "GiB")
# 84.6 GiB  → 80GiBカード一枚には入らない

重みは32.6GiBなのにKVキャッシュが40GiBです。重みを4bitに押し込んだところでKVキャッシュのほうが大きいのです。この点を見落とすと「量子化したのになぜまだ載らないのか」という状態に閉じ込められます。

実務でより頻繁に使うのは逆方向の計算です。カードが決まっているときに同時に何件受けられるかを問うほうです。

def max_concurrent(vram_gib, params_b, weight_dtype, layers, kv_heads, head_dim,
                   seq_len, util=0.9, kv_bytes=2.0, fixed_gib=1.0):
    usable = vram_gib * util - fixed_gib - weight_gib(params_b, weight_dtype)
    if usable <= 0:
        return 0
    per_seq_gib = kv_bytes_per_token(layers, kv_heads, head_dim, kv_bytes) * seq_len / GIB
    return int(usable / per_seq_gib)

print(max_concurrent(80, 70, "int4", 80, 8, 128, seq_len=8192))   # 15
print(max_concurrent(80, 70, "int4", 80, 8, 128, seq_len=2048))   # 61
print(max_concurrent(80, 70, "int4", 80, 8, 128, seq_len=8192, kv_bytes=1.0))  # 30

三行がそれぞれ一つの運用判断です。コンテキスト上限を8Kから2Kに下げればスループットが4倍になり、KVキャッシュをfp8で保存すれば2倍になります。どちらの判断も品質に影響しますが、少なくともいくら得ていくら賭けているのかが目に見えます。

PagedAttention — 計算が合っていても使えないメモリ

ここまでの計算は「必要な分だけ正確に使う」を前提にしていました。初期のサービング実装はそうできていませんでした。リクエストが来ると、そのリクエストが到達しうる最大長ぶんのKVキャッシュを連続した領域にあらかじめ確保していたからです。最大長が4096なのに実際の出力が200トークンなら、残りはそのまま捨てられます。

vLLMの論文はこの無駄を二種類に分けて説明しています。予約して使わない内部断片化と、ブロックの間に残って誰も使えない外部断片化です。論文が測定したワークロードでは実際に有効に使われたKVメモリの比率はかなり低く、残りはこの二つの断片化で消えていました。具体的な数値はワークロード依存なので、そのまま引用するよりも「出力長の分散が大きいほど無駄が増える」という構造だけ覚えていただければ十分です。

PagedAttentionはこの問題をオペレーティングシステムのやり方で解きます。KVキャッシュを固定サイズのブロック(通常16トークン)に切り分け、論理的に連続したシーケンスを物理的に散らばったブロックにマッピングします。結果として無駄はシーケンスあたり最後のブロックの空き、つまり最大15トークンに制限されます。

副次効果がもう一つあり、実はこちらのほうが大きいのです。ブロック単位で管理すると複数のシーケンスが同じブロックを共有できます。同一のシステムプロンプトを使うリクエスト100件は、その接頭辞に当たるブロックを一組だけ持って参照カウントで共有します。接頭辞が長くリクエスト数が多いサービスでは、この接頭辞共有が断片化の解消よりも大きな削減を生みます。

# vLLMで実際に調整することになる値
vllm serve <모델경로> \
  --max-model-len 8192 \            # コンテキスト上限。KVキャッシュの上限を直接決める
  --gpu-memory-utilization 0.90 \   # 全VRAM中の使用比率。残りはオーバーヘッド用
  --kv-cache-dtype fp8 \            # KVキャッシュだけ8bitに。容量は半分
  --max-num-seqs 64 \               # 同時シーケンス上限
  --enable-prefix-caching           # 共通接頭辞ブロックの再利用

# 起動ログでこの行を確認してください。上の計算と合っているはずです。
#   "GPU KV cache size: 129,024 tokens"
#   "Maximum concurrency for 8192 tokens per request: 15.75x"

起動ログに出るKVキャッシュのトークン数が先ほどのmax_concurrentの計算と大きく食い違うなら、どちらかが間違っています。たいていはKVヘッド数をクエリヘッド数と取り違えた場合です。

量子化はタダではありません

重みを4bitに減らせばメモリは4分の1になります。これは事実です。問題は、この文がたいていここで終わることです。

第一に、品質の損失は平均指標には現れにくいのです。WikiTextのパープレキシティが0.1上がったと要約される変化が、実際には長い出力の後半、コード生成、多言語、フォーマット遵守といった裾の領域ではるかに大きく現れることがよくあります。パープレキシティは平均的なトークン予測の難しさであってタスクの成功率ではありません。量子化するかどうかを決めるときは、必ずご自身のタスクの評価セットで測る必要があります。その評価セットの作り方は勘でやらないLLM評価の回で別途扱いました。

第二に、4bitが常に速いわけではありません。重みのみの量子化は保存だけが4bitで演算はfp16なので、行列積のたびに逆量子化が入ります。バッチが小さくメモリ帯域がボトルネックのときは読むバイト数が4分の1なので確実に速くなります。逆にバッチが大きく演算がボトルネックになると逆量子化のオーバーヘッドだけが残り、fp16より遅くなることがあります。「量子化すれば速くなる」はバッチ1のベンチマークの話です。

第三に、KVキャッシュの量子化は重みの量子化とは別の判断であり、おおむねより安全です。fp8のKVキャッシュは損失が小さいという報告が多い一方、4bitのKVは長いコンテキストで劣化が観測されるという報告があります。先ほどの計算で見たとおり長いコンテキストのサービスでメモリを支配するのはKV側なので、順序を逆にしてKVキャッシュから先にfp8へ落とし、重みはfp16か8bitで置く組み合わせのほうが良い場合が少なくありません。

まとめると選択の順序はこうなります。KVキャッシュのdtypeをまず見て、次にコンテキスト上限を見て、重みの量子化は最後に見ます。ほとんどのチームは正反対の順序で近づいて時間を使ってしまいます。

おわりに — 重みは定数、KVキャッシュは変数

覚えるべきことは一行です。重みのメモリはモデルを選んだ瞬間に確定する定数で、KVキャッシュは皆さんが受けると決めたトラフィックの形が決める変数です。容量不足はほぼ常に変数の側から来ます。

ですからデプロイ前に計算すべき値は二つで足ります。トークンあたりのKVバイト数(2×レイヤー数×KVヘッド数×ヘッド次元×要素バイト数)と、残りのVRAMをその値で割った総トークン数です。この二つの数字があれば、コンテキスト上限と同時実行上限の交換を会議室でその場で計算できます。量子化の話はその後で構いません。