なぜホームAIのコンテキスト長が増えるとKVキャッシュも増加するのか？

KVキャッシュは、ランタイムが保持しているすべてのトークンについて、複数のモデルレイヤーにわたるアテンションのキーと値を保存するため、コンテキスト長に応じて増加します。

短い家庭用AIプロンプトなら、複数のユーザー分のメモリが残る場合があります。一方、長いドキュメント、長時間のチャット履歴、エージェントの実行トレースなどは、モデルファイル自体を変更することなく、はるかに多くのワーキングメモリを消費する可能性があります。キャッシュはプロンプト処理中に作成され始め、モデルが新しいトークンを生成するたびに増え続けます。そのサイズは、レイヤー数、アテンションアーキテクチャ、数値精度、同時リクエスト数にも左右されます。以下では、1つのトークンからサーバー全体のメモリ上限に至るまで、その増加の仕組みを説明します。

保持されるトークンごとにアテンション状態が追加される

Transformer推論では、各レイヤーが処理済みのトークンからキーと値のテンソルを生成します。ランタイムはこれらのテンソルを保持することで、次のトークンがコンテキスト全体を再計算せずに過去のコンテキストへアテンションを向けられるようにします。

vLLMの研究では、効率的なLLM提供における主要なサービングメモリ要件として、トークンごとのKV状態が特定されています。新しいトークンが追加されると新たなキャッシュエントリが作成され、保持済みのエントリは後続のアテンション処理で引き続き利用されます。

したがって、キャッシュは静的なモデル重みの一部ではなく、変更されるリクエスト状態です。同じモデルでも、アクティブな会話が長くなればメモリフットプリントは大きくなります。

キャッシュの増加は、保持されるシーケンス長に対しておおむね線形になる

モデルアーキテクチャとキャッシュ精度が同じであれば、保持するトークン数を2倍にすると、そのリクエストで保持されるKVエントリもおおむね2倍になります。プロンプトと生成された回答は、どちらもアクティブなシーケンスに含まれます。

H2Oは、KVキャッシュがシーケンス長とバッチサイズに応じてどのように拡大するかを説明しています。追加される各トークンが、キャッシュを生成するすべてのレイヤーでキーと値を生み出すため、この関係はおおむね線形です。

そのため、小さなコンテキストからはるかに大きな最大値へランタイム設定を変更すると、モデルが同じ重みを使用していても、実際のメモリ上限が変わる可能性があります。

最大設定値と実際の使用量は異なります。ランタイムによっては必要に応じてキャッシュブロックを割り当てますが、将来の拡張を保証するために、早い段階で大きな領域を予約するものもあります。

モデルアーキテクチャによってトークンあたりのバイト数が変わる

同じパラメータ数のモデルでも、レイヤー数、ヘッドの次元数、アテンションヘッド数、グループ化クエリアテンション、マルチクエリアテンションなどが異なるため、必要なKVメモリ量が異なる場合があります。

KIVIはKVキャッシュの精度を研究し、キーと値をより少ないビット数で保存することで、ピークメモリを大幅に削減できることを示しています。この効果は基盤となるモデル重みを小さくするものではなく、リクエスト状態に適用されます。

グループ化クエリアテンションとマルチクエリアテンションでは、より多くのクエリヘッド間でキー・バリューヘッドを共有するため、完全なマルチヘッドアテンションと比べて、トークンあたりのキャッシュ容量を削減できます。ただし、レイヤー数とヘッド幅は保持される状態の量を引き続き増加させます。

したがって、実用的な見積もりには、別のモデルから流用した一般的な「トークンあたりのバイト数」ではなく、正確なモデルアーキテクチャとランタイムのキャッシュ形式を使用する必要があります。

Featured

ZimaCube 2 パーソナルクラウドホームNAS

ZimaCube2

$799.00 - $2,499.00

Buy Now

プリフィル後も生成トークンによってキャッシュは拡張し続ける

プロンプト処理によって、入力コンテキスト用の初期キャッシュが作成されます。その後の自己回帰デコードでは、受理された出力トークンごとに状態が追加され、後続のトークンが会話全体を参照できるようになります。

vAttentionは、リクエスト開始時には最終的な出力長が分からないため、動的なキャッシュ拡張を割り当て問題として扱います。多く予約しすぎるとメモリを浪費し、少なすぎるとプリエンプションや拡張処理が必要になる可能性があります。

余裕をもって収まるプロンプトでも、長い回答の生成中にメモリ上限を超える場合があります。そのため、出力制限は回答の長さや生成時間だけでなく、メモリも保護します。

同時利用ユーザーが増えると個別のコンテキスト状態も増加する

モデル重みは複数のリクエストで共有できますが、通常、アクティブな各会話は独自のトークン履歴とKVキャッシュを保持します。同じモデルを使っているからといって、長いコンテキストを持つ5人のユーザーが1つの共通キャッシュを共有するわけではありません。

近年のKV管理に関する研究では、リクエストごとの予約を、メモリ効率とプリエンプションリスクの中心的なトレードオフとして捉えています。出力長が不明なため、単純なユーザー数よりも全体のピーク使用量を予測するのが難しくなります。

ランタイムが正確なプレフィックス一致に対応していれば、共有プロンプトのプレフィックスでキャッシュ状態を再利用できる場合があります。しかし、個別のチャット履歴や異なる出力によって、別々の分岐が作られます。

ZimaSpaceのハードウェアガイドでは、コンテキストと同時実行性を、モデルファイル以外にも必要となるメモリ要件として扱っています。そのため、1人でのテストでは、家庭用アシスタントに必要なRAMやVRAMを過小評価する可能性があります。

ページング、量子化、エビクションは上限を変えるが、原因は変えない

ページ単位の割り当てでは、キャッシュ状態を小さなブロックに分割することで断片化を抑えられます。そのため、ランタイムはあらゆるシーケンスに備えて、サイズの大きい連続領域を1つ予約する必要がありません。

PagedAttentionはブロック単位の割り当てを提供し、キャッシュ量子化は保存する値あたりのバイト数を削減します。また、エビクションポリシーは選択した古いエントリを破棄します。これらの方法によって収められるコンテキスト量は変わりますが、保持されるアテンション状態がトークンの蓄積に伴って増加する点は変わりません。

エビクションやスライディングウィンドウによって古いトークンを削除すればメモリを上限内に抑えられますが、モデルは通常のフルコンテキスト経路を通じて、破棄された状態を参照できなくなります。圧縮や選択的な保持によって、品質やワークロード固有のトレードオフが生じる場合もあります。

実際のモデル、コンテキスト、キャッシュ精度、バッチサイズ、ユーザー数を使ってキャッシュ使用量を測定してください。別のトークンやリクエストを、エビクション、オフロード、再計算、または失敗なしには受け入れられなくなった時点が、実用上の上限です。

よくある質問

コンテキスト長を増やすとモデルファイルは大きくなりますか？

いいえ。モデルの重みは変わりません。増加するメモリは、アクティブなプロンプトと生成されたトークンのためにランタイムが作成する状態です。

最大コンテキストを大きく設定すると、KVメモリは必ずすぐに全量割り当てられますか？

いいえ。割り当ての動作はランタイムによって異なります。早い段階で容量を予約するものもあれば、ページングシステムではトークンが受け入れられるたびにブロックを割り当てます。

VRAMがいっぱいになった場合、システムRAMでKVキャッシュを保持できますか？

キャッシュ状態をオフロードまたは移動できるランタイムもありますが、転送によってレイテンシが増加し、ソフトウェアの対応状況、帯域幅、アクティブなアテンション経路に左右されます。