KVキャッシュサイズの計算
入力
| レイヤー数 | 32 |
|---|---|
| key/valueヘッド数 | 32 |
| ヘッド次元 | 128 |
| コンテキスト長 | 4,096 |
| バッチサイズ | 1 |
| 精度 | FP16 / BF16（2バイト） |
KVキャッシュサイズの計算
Transformerが推論中に保持するKey-Valueキャッシュのメモリを、レイヤー数・key/valueヘッド数・ヘッド次元・コンテキスト長・バッチサイズ・1要素あたりのバイト数から推定します。
入力
モデル
ワークロード
結果
値を入力すると計算結果が表示されます。
KVキャッシュサイズ
Transformerの推論では、モデルはそれまでに見たすべてのトークンにアテンションを向けます。各生成ステップでコンテキスト全体のkeyベクトルとvalueベクトルを再計算するのを避けるため、それらを保存して再利用します。これがKey-Valueキャッシュです。このキャッシュは長コンテキストのサービングにおいてアクセラレータメモリの最大の消費者となることが多く、ときにモデルの重み自体より大きくなります。この計算ツールは、レイヤー数・key/valueヘッド数・ヘッド次元・コンテキスト長・バッチサイズ・格納する1要素あたりのバイト数から、そのサイズを推定します。
なぜキャッシュが存在するのか
テキストを1トークンずつ生成するということは、新しいトークンごとにそれ以前のすべてのトークンにアテンションを向ける必要があるということです。キャッシュがなければ、各ステップでプロンプト全体とそれまでの出力をすべて処理し直すことになり、生成コストがシーケンス長の二乗で増えてしまいます。過去トークンのkeyとvalueを保持することで、モデルは各ステップを新しいトークンだけに比例する作業に変えられます。その速度の代償がメモリです。各レイヤー・各key/valueヘッド・コンテキスト内の各トークンについて、keyとvalueという1対のベクトルを保持します。
計算式
キャッシュはkeyとvalueという2つのテンソルを保持するため、バイト単位のサイズは次のとおりです。
Kbytes=2⋅L⋅H⋅d⋅s⋅B⋅eここで はレイヤー数、 はkey/valueヘッド数、 はヘッド次元、 はトークン単位のコンテキスト長、 はバッチサイズ、 は格納する1要素あたりのバイト数です。 で割るとギガバイトになります。サイズは各因子に比例して増えます。これが核心となる直感で、コンテキスト・バッチ・レイヤー数のいずれかを2倍にすればキャッシュも2倍になります。
グループ化クエリアテンション
key/valueヘッド数 こそ、グループ化クエリアテンション（GQA）が狙う対象です。素朴なマルチヘッドアテンションでは、すべてのクエリヘッドがkey/valueペアを所有します。GQAではクエリヘッドのグループが単一のkey/valueヘッドを共有し、マルチクエリアテンションはレイヤー全体でただ1つを共有します。キャッシュは に比例するため、32個のkey/valueヘッドを8個に減らせば、クエリヘッド（とモデル品質の大部分）はそのままにキャッシュを4分の1に縮小できます。最近の大規模モデルの多くは、まさにこの理由でGQAを採用しています。
計算例
32レイヤー、key/valueヘッド32個、ヘッド次元128のモデルが、4,096トークンのコンテキストを保持し、単一のシーケンスを16ビット精度でサービングする場合を考えます。
Kbytes=2×32×32×128×4096×1×2=2,147,483,648これは約2.15 GBです。このようなシーケンスを4つ同時にサービングすると4倍のおよそ8.59 GBになります。キャッシュを8ビット精度に切り替えれば、いずれの数値も半減します。
結果の読み方
すべての因子が線形に掛け合わさるため、長コンテキスト・高バッチのサービングではそれらが積み重なり、キャッシュが重みを圧倒することがあります。これが、本番の推論エンジンが1要素あたりのバイト数を減らすキャッシュ量子化、key/valueヘッドを減らすGQA、断片化なくキャッシュを詰めるページドアテンションに頼る理由です。同じハードウェア予算の演算側はアテンションメモリの計算で、モデルの重みの部分はLLM推論のVRAM計算で扱っています。
よくある質問 (FAQ)
KVキャッシュとは何ですか？
自己回帰生成の間、Transformerは過去のすべてのトークンにアテンションを向けます。各ステップでコンテキスト全体のkeyベクトルとvalueベクトルを再計算する代わりに、それらを保存して再利用します。この保存場所がKVキャッシュです。メモリと引き換えに速度を得る仕組みで、これがなければ新しいトークンを生成するたびにプロンプト全体とそれまでの出力すべてを処理し直すことになります。キャッシュは、各レイヤー・各key/valueヘッドについてkeyとvalueの2つのテンソルを保持し、コンテキスト内のトークン1つにつき1エントリを持ちます。
グループ化クエリアテンションはどのようにキャッシュを削減しますか？
標準的なマルチヘッドアテンションでは、すべてのクエリヘッドが個別のkey/valueペアを保持します。グループ化クエリアテンション（GQA）では複数のクエリヘッドが1つのkey/valueヘッドを共有し、マルチクエリアテンションはこれを極端化して単一の共有key/valueヘッドにします。キャッシュサイズはkey/valueヘッド数に比例するため、たとえば32から8に減らせば、クエリヘッド（とモデル品質の大部分）はそのままにキャッシュを4分の1に縮小できます。
なぜ長いコンテキストはメモリ消費が大きいのですか？
キャッシュはコンテキスト長に比例して増加します。コンテキストのトークン数を2倍にすればキャッシュも2倍です。さらにバッチサイズにも比例して増えるため、長いコンテキストのシーケンスを多数同時に処理すると両者が掛け合わさります。コンテキスト長が大きくなると、KVキャッシュがモデルの重み自体が占めるメモリを上回ることもあり、これが長コンテキストのサービングでキャッシュ量子化・GQA・ページドアテンションといった手法に頼る理由です。
免責事項
この推定は、指定された精度でのkeyテンソルとvalueテンソルのみを数え、メモリの断片化・ページング用メタデータ・フレームワークの予約といったサービングのオーバーヘッドは無視します。実際の推論エンジンでの使用量は、これより幾分大きくなります。