ML 研究向け GPU · KYC 不要の暗号通貨決済
IteraGPU
推論メモリ · コンテキストと並行性

KV キャッシュにどれだけメモリを確保すべきか?

均一な密なキャッシュの場合、各レイヤー、KV ヘッド、保持される位置、同時シーケンスごとに 2 つのテンソル K と V を数えます。重みの形式ではなく、実際のキャッシュの形式を使用してください。この計算は、メモリ領域の理論上の容量を示すものであり、GPU の総ピーク、レイテンシ、回答の品質を予測するものではありません。その後、実際の負荷で割り当て戦略を確認してください。

01 /

メモリに残るものを説明する

自己回帰生成中、すでに処理されたトークンのキーと値を、以降のステップのために保持できます。このキャッシュはアテンション層に属します。その容量は、パラメータ数だけでなく、モデルと保持される履歴に依存します。会話のコンテキストには、指示、以前のメッセージ、アプリケーションが追加したドキュメントも含まれます。

最初の判断は運用上のものです。いくつのシーケンスを、どのくらいの長さまでアクティブに保つ必要があるか?20件のリクエストのうち2件が同時に実行されるキューは、必ずしも20個の常駐キャッシュを意味しません。実際のリクエストの受け入れ状況と、生成によって生じる可能性のある追加シーケンスを確認します。

モデルのリビジョン、アテンション層、KV ヘッド、キーと値の次元、その dtype、キャッシュ戦略を含むシートを用意します。トークナイザーと会話テンプレート適用後のトークンを数えます。文字数の制限では、この割り当ては表せません。

技術的な出典: Hugging Face — レイヤーごとのキャッシュの動作と形式

02 /

KV ヘッドを使う、特に GQA の場合

クエリヘッドの数 Q と KV ヘッドの数は異なる場合があります。従来のマルチヘッドアテンションでは一致します。MQA では単一の KV ヘッドが共有され、GQA では複数の Q ヘッドが 1 つの KV ヘッドにまとめられます。コンフィギュレーションに該当フィールドがある場合は num_key_value_heads を読み、アーキテクチャにとっての意味を確認してください。

例えば、Qヘッドが40、KVヘッドが8の場合、KVグループあたりQヘッドは5つになります。保存キャッシュの計算式には40ではなく8を使います。この比率はアテンションのメモリ全体を表すものではありません。演算や変換によって一時的な領域が生じることがあります。

学習済みモデルのメモリ需要を減らすために、単にこの数値を変更しないでください。アテンションのスキームはそのアーキテクチャの一部です。ヘッド数が異なる2つのモデルが、容量計算によって等価な変種になるわけではありません。それらの品質は個別に評価する必要があります。

技術的な出典: Hugging Face — LlamaConfigのフィールドとMHA・MQA・GQAの区別 · PyTorch — Q/K/Vの次元とGQAアテンションの制約

03 /

計算式とその単位を定める

均一な場合、Lは層数、HkvはKVヘッド数、Dはその次元、Tはシーケンスごとに保持される位置数、Bはシーケンス数、qは値あたりのバイト数です。係数2はKとVを数えています。この近似は、キーと値が同じ次元と同じ形式であり、圧縮もプレフィックス共有もないことを前提としています。

最終的な結果のみを変換してください。1 GiB は 1 073 741 824 バイト、1 十進 GB は 1 000 000 000 バイトです。丸めや単位の変更で差が隠れないよう、バイト数のまま記録に残しておきましょう。

パディングなしで長さが異なる場合は、B × Tを実際に保存される位置の合計に置き換えます。層が異質な場合は、層ごとに合計します。パディング付きの密なストレージ、ブロック単位の割り当て、静的な予約では、割り当て済みのスロットを数える必要があり、それらは有効なトークンを超えることがあります。

理論上のKV(バイト)= 2 × L × Hkv × D × T × B × q;KV(Gio)= KV(バイト)÷ 1 073 741 824

技術的な出典: Hugging Face — キャッシュテンソルの次元 · NIST — 10進単位と2進接頭辞

04 /

計算例:6つのシーケンス、GPU測定なし

40層、KVヘッド8、次元128という架空のアーキテクチャを考えます。値あたり2バイトの均一なキャッシュを仮定します。各シーケンスは最大3 072トークンの入力と、さらに1 024トークン分の予約を受け取るため、上限は4 096位置となります。これらは教育的な仮定であり、実在するモデルの確認済み構成ではありません。

位置あたり・シーケンスあたりの計算コストは2 × 40 × 8 × 128 × 2 = 163,840バイトです。すると4,096位置のシーケンスは671,088,640バイト、つまり0.625 GiBになります。6つのシーケンスでは4,026,531,840バイト、つまりキャッシュだけで3.75 GiBになります。

保持する長さを2倍にすると、この式ではこの項目も2倍になります。8個のKVヘッドを40個に置き換えると、他の前提がすべて同じであれば5倍になります。これらの比率は、実在するモデルの高速化、品質低下、互換性を何ら予告するものではありません。

あくまで理論上の計算です:L = 40、D = 128、q = 2バイト;重みと一時領域は除外。
シーケンス B位置 TKVヘッド計算されたバイト数GiB
14 0968671 088 6400,625
64 09684 026 531 8403,75
68 19288 053 063 6807,5
64 0964020 132 659 20018,75
05 /

キャッシュ戦略に合わせて予算を調整する

動的キャッシュは保持される位置とともに増大します。静的キャッシュは最大容量を予約するため、起動時に観測した短いリクエストだけでなく、この予約を基準にサイズを見積もってください。スライディングウィンドウアテンションでは、一部の層が履歴に上限を設けることがあります。完全アテンションの層は別途計算が必要です。

キャッシュの量子化とCPUへのオフロードも別の戦略であり、モデルとソフトウェアに依存します。これらはストレージ、転送、計算の制約を変化させます。重みを量子化しても、キャッシュが同じ形式である証拠にはなりません。

キャッシュのクラスとその明示的なパラメータを記録してください。また、リクエストの終了やキャンセル後にスロットが解放されることも確認しましょう。短いシーケンスと長いシーケンスが混在する負荷では、均一な最大予約が、有効な内容の合計だけよりも重くのしかかることがあります。

技術的な出典: Hugging Face — 動的・静的・量子化・オフロードキャッシュ

06 /

代表的な負荷を段階的に検証する

3つのケースを用意します。通常の入力、想定される長い入力、許可される同時リクエストの最大数です。モデル、tokenizer、テンプレート、生成上限、終了ルールを固定します。一度に1つの次元だけを変えてください。回答が切り詰められたり文書が切り捨てられたりすると、実行される処理が変わります。

ロード、入力の初期処理、生成を別々に計測します。計時したフェーズを基準にデバイスを同期し、開始時のメモリ水準とピークを保持します。メモリの手法は、allocated と reserved が加算されない理由、およびそれらの最大値の差がキャッシュを単独で分離しない理由を説明します。

検証は、検証済みの上限と許容可能な出力に到達しなければなりません。完了したリクエストの ID、エラー、生成長、品質基準です。短い入力で成功した起動は、最大同時実行を検証したことにはなりません。完了前のメモリエラーは、完全な実行に必要な量の測定にはなりません。

技術的な出典: PyTorch — 選択したデバイスでの処理の同期 · PyTorch — メモリカウンタの対象範囲

07 /

選択を歪めるエラーを排除する

計算したキャッシュを GPU の総容量に置き換えないでください。重み、一時的なアクティベーション、保持される出力、ソフトウェアの分析を加えます。また、この量をカード枚数で自動的に割らないでください。レイヤーやヘッドの配置は、実際に設定し、各デバイスで検証する必要があります。

IteraGPU Lab のノートブックは、アテンションのない小さな密結合ネットワークを対象とした計測の演習です。メモリのフェーズを読み取る助けになりますが、その context パラメータはこの KV 計算を検証するものではありません。負荷シートと推論フォルダを使って、ご自身のモデルの試験を準備してください。

プランの容量を比較するのは、演算量、実際に観測された最大値、まだ検証されていないケースを区別した後に行ってください。レンタルの定額プランは作業時間枠を整えるものであり、コンテキスト長や生成速度を保証するものではありません。

  • Q ヘッドと KV ヘッドを混同する:アーキテクチャの構成を確認し直す。
  • プロンプトだけを予算計上する:生成と実際の予約を組み込む。
  • 「重み 4 ビット」を「キャッシュ 4 ビット」と読む:両方の形式を確認する。
  • 同時実行を忘れる:実際に常駐するシーケンスを数える。
  • カード間でメモリを共有できると約束する:実際の配置を検証する。

実用的な質問

パラメータ数だけから KV キャッシュを知ることはできますか?

いいえ。アテンション層のアーキテクチャ、KVヘッド、その次元、キャッシュ形式、保持される位置、常駐するシーケンスも必要です。サイズが近い2つのモデルでも、異なるKV予算を必要とすることがあります。

静的キャッシュはリクエストの長さしか消費しませんか?

予約する容量を決める必要があります。短いリクエストからこの最大割り当てを推測することはできません。キャッシュのパラメータを確認し、実際に作成された構成を測定してください。

計算式の結果だけでカードを選べますか?

これは公表された前提に基づいてキャッシュのみを見積もります。選択には、他のメモリ項目、ソフトウェアの互換性、そして代表的なコンテキスト、同時実行、品質を備えた完全な試験も考慮する必要があります。