AI・機械学習
上級

メモリ効率的ロングコンテキスト推論(メモリコウリツテキロングコンテキストスイロン)

LLMの長系列推論時に発生する大量のKVキャッシュメモリを削減する技術群の総称。PagedAttention、KVキャッシュ圧縮、Multi-Query Attention(MQA)、Grouped-Query Attention(GQA)などの手法により、128Kトークン超の推論をコンシューマGPUでも実行可能にする。

0 回閲覧
0 いいね

メモリ効率的ロングコンテキスト推論とは

メモリ効率的ロングコンテキスト推論とは、LLMが長いシーケンスを処理する際に発生するKVキャッシュ(Key-Value Cache)のメモリ消費を削減する技術群である。標準的なTransformerでは、シーケンス長に比例してKVキャッシュが増大し、128Kトークンの推論では数十GBのGPUメモリが必要になる。この問題を解決する技術は、ロングコンテキストLLMの実用化に不可欠である。

KVキャッシュのメモリ問題

KVキャッシュのメモリ消費量は以下の式で計算される:

KVキャッシュサイズ = 2 × 層数 × ヘッド数 × ヘッド次元 × シーケンス長 × バイト数

具体例(Llama 3.1 70B、FP16):

  • 2 × 80層 × 8 KVヘッド(GQA) × 128次元 × 128,000トークン × 2バイト = 約32.8GB
  • バッチサイズ4では131.2GBとなり、A100 80GBでも1枚では処理不可能

この問題に対する主要な解決策:

  • Attention構造の改良: MQA/GQA でKVヘッド数を削減
  • KVキャッシュ管理: PagedAttention でメモリの断片化を解消
  • KVキャッシュ圧縮: 量子化・蒸留・プルーニングでキャッシュサイズを削減
  • アーキテクチャレベル: SWA・Ring Attention・StreamingLLM で根本的にメモリ消費パターンを変更

主要技術の比較

技術メモリ削減率精度影響実装先適用レベル
GQA(Grouped-Query Attention)KVヘッド1/4〜1/8最小(<1%)Llama 3.1, Mistralモデルアーキテクチャ
MQA(Multi-Query Attention)KVヘッド1/n小(1〜2%)PaLM 2, Falconモデルアーキテクチャ
PagedAttention断片化解消(20〜40%削減)なし(完全等価)vLLM推論エンジン
KV量子化(FP16→INT4)4倍削減小(0.5〜2%)KIVI, QServe推論最適化
H2O(Heavy-Hitter Oracle)10〜20倍削減中(タスク依存)研究段階KVプルーニング
StreamingLLM固定メモリ中(長距離精度低下)研究→実用化中ストリーミング推論
Ring AttentionGPU数に比例分散なし(完全等価)研究→実用化中分散推論

PagedAttention(vLLM)

UC BerkeleyのKwon et al.が2023年に提案し、vLLMで実用化されたKVキャッシュ管理技術:

  • OSの仮想メモリ管理と同様の「ページング」概念をKVキャッシュに適用
  • KVキャッシュを固定サイズのブロック(ページ)に分割し、非連続なメモリ領域に配置
  • 従来のKVキャッシュは連続メモリ確保が必要で、最大シーケンス長分を事前確保→大量の無駄
  • PagedAttentionでは実際に使用した分だけ動的にページを割り当て、メモリ使用率を最適化
  • 結果: 同一GPUで2〜4倍のスループット向上。vLLMの中核技術として業界標準に

GQA(Grouped-Query Attention)

MHA(Multi-Head Attention)とMQA(Multi-Query Attention)の中間に位置する手法:

  • MHA: Q/K/Vすべてが同数のヘッド(例: 64ヘッド)→KVキャッシュ最大
  • MQA: K/Vが1ヘッドに共有→KVキャッシュ最小だが精度低下
  • GQA: K/Vを8ヘッドにグループ化(Qは64ヘッド)→KVキャッシュ1/8で精度はMHAに近い
  • Llama 3.1(70B/405B)、Mistral Large 2、Qwen2.5-72BなどがGQA採用
  • 128Kコンテキストでのメモリ節約効果: 70Bモデルで約262GB→32.8GB(FP16)

KVキャッシュ量子化

KVキャッシュの値をFP16からINT8/INT4に量子化してメモリを削減:

  • KIVI(2024): Key/Valueのチャネルごとにスケールを計算し、Per-channel INT4量子化を適用。Perplexity増加0.5%未満で4倍のメモリ削減
  • QServe(2024): W4A8KV4方式(重みINT4、活性化INT8、KVキャッシュINT4)の統合量子化フレームワーク
  • KV-Compress(2024): 重要度の低いKVペアを選択的に破棄しつつ、残りを量子化するハイブリッド手法
  • 実用的な精度: INT8量子化ではperplexity増加0.1%未満、INT4でも0.5〜1%の範囲

StreamingLLM

MIT Han Labが提案した、無限長テキストのストリーミング処理技術:

  • 先頭のAttention Sink Token(通常4トークン)と直近のウィンドウ(例: 4,096トークン)のみを保持
  • 中間部のKVキャッシュを破棄することで、メモリ消費を固定化
  • Attention Sinkの発見: 先頭トークンはコンテンツに関係なく高い注意重みを受ける。これを除去するとAttentionの数値安定性が崩壊
  • 制約: 破棄された中間部の情報は失われるため、要約や検索には不向き。チャットボットの長時間対話など、直近の文脈が重要なタスクに適する

よくある質問(FAQ)

Q1: コンシューマGPU(RTX 4090 24GB)でLlama 3.1 70Bの128K推論は可能か? A: 標準ではモデル重みだけで約140GB(FP16)必要なため不可能。4bit量子化(約35GB)+ GQAのKVキャッシュ圧縮(約4GB@INT4)+ PagedAttentionでも合計約39GBとなり、2枚構成(48GB)が最低ライン。8Bモデルなら1枚で128K推論が可能。

Q2: PagedAttentionとFlash Attentionは競合するのか? A: 競合しない。Flash Attentionは計算カーネルの最適化(HBMアクセス削減)、PagedAttentionはKVキャッシュのメモリ管理最適化であり、両者は直交する技術。vLLMは両方を同時に使用している。

Q3: GQAとMQAのどちらが主流か? A: 2026年現在、GQAが圧倒的に主流。MQAは精度低下が大きすぎ、GQAは精度とメモリ効率のバランスが優れている。Llama 3.1、Mistral Large 2、Qwen2.5、Gemma 2など主要モデルのほぼ全てがGQAを採用。

まとめ

  • KVキャッシュのメモリ消費はロングコンテキスト推論の最大のボトルネック
  • GQAによるKVヘッド数削減(1/4〜1/8)が事実上の標準アーキテクチャ
  • PagedAttention(vLLM)がKVキャッシュメモリ管理の業界標準
  • INT4 KV量子化で精度低下0.5%未満・メモリ4倍削減が実用化
  • StreamingLLMは無限長ストリーミング処理を固定メモリで実現するが、中間情報は失われる