スクラッチから理解するLLMのKVキャッシュの仕組みと実装

LLM の本番環境における効率的な推論において、KV キャッシュは最も重要な技術の一つです。KV キャッシュは、本番環境での計算効率の高い LLM 推論を実現するための重要なコンポーネントです。この記事では、概念とコードの両面から KV キャッシュがどのように動作するかを、人間が読みやすい形式でゼロから実装した例を用いて解説します。 LLM の基本概念を説明する技術チュートリアルを提供してから随分経ちました。現在、怪我からの回復中であり、より大規模な LLM 研究に焦点を当てた記事に取り組んでいるため、読者の方々からいくつかご要望いただいたトピック（私の『ゼロから作る Large Language Model』の書籍には含まれていなかったため）についてチュートリアル記事を共有しようと思いました。 お楽しみください！ 概要 簡単に言えば、KV キャッシュは推論時（トレーニング後）に再利用するために、中間的なキー (K) とバリュー (V) の計算結果を保存するものであり、これによりテキスト生成時の速度が大幅に向上します。KV キャッシュの欠点は、コードの複雑さが増し、メモリ要件が高まること（これが当初書籍に含めなかった主な理由です）、そしてトレーニング時には使用できないことです。しかし、本番環境で LLM を使用する際、コードの複雑さとメモリのトレードオフに対して、推論速度の向上は十分に価値があると言えます。 KV キャッシュとは何か？ LLM がテキストを生成している様子を想像してみてください。具体的には、LLM に以下のプロンプトが与えられていると仮定します："Time"。ご存知の通り、LLM は一度に一つの単語（またはトークン）を生成するため、続く二つのテキスト生成ステップは、下の図に示すようなものになります： この図は、LLM が1トークンずつテキストを生成する様子を示しています。プロンプト「Time」から始めると、モデルは次のトークン「flies」を生成します。次のステップでは、「Time flies」という完全なシーケンスが再処理され、トークン「fast」が生成されます。 生成された LLM のテキスト出力には冗長性があることに注意してください。これは次の図で強調されています： この図は、各生成ステップで LLM が再処理しなければならない繰り返しコンテキスト（「Time flies」）を強調しています。LLM は中間のキー/バリュー状態をキャッシュしないため、新しいトークン（例えば「fast」）が生成されるたびに完全なシーケンスを再エンコードします。 LLM テキスト生成関数を実装する際、通常は各ステップから最後に生成されたトークンのみを使用します。しかし、上記の可視化は概念的なレベルでの主要な非効率性の一つを浮き彫りにしています。この非効率性（または冗長性）は、アテンションメカニズムそのものに焦点を当てて拡大して見ると、より明確になります（アテンションメカニズムについて詳しく知りたい場合は、私の著書『Build a Large Language Model (From Scratch)』の第 3 章、あるいは「LLM における自己アテンション、マルチヘッドアテンション、因果的アテンション、クロスアテンションの理解と実装」の記事をお読みください）。 以下の図は、LLM の中核をなすアテンションメカニズム計算の一部を示しています。ここでは、入力トークン（"Time" と "flies"）が 3 次元ベクトルとして符号化されています（実際にはこれらのベクトルははるかに大きくなりますが、そうすると小さな図に収めるのが難しくなるためです）。行列 W は、アテンションメカニズムの重み行列であり、これら入力を変換してキー、バリュー、クエリベクトルを生成します。 以下の図は、キーおよびバリューベクトルが強調表示された、基礎となるアテンションスコア計算の一部を示しています： ![image](https://substackcdn.com/image/fetch/$s_!4BB1!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F3748127c-532e-4169-8e12-1fb48e263dbd_945x445.png) この図は、LLM がアテンション計算中にトークン埋め込みからキー (k) ベクトルとバリュー (v) ベクトルをどのように導出するかを示しています。各入力トークン（例："Time" と "flies"）は、学習された行列 W_k および W_v を用いて投影され、それぞれに対応するキーおよびバリューベクトルを取得します。 前述の通り、LLM は一度に 1 つの単語（またはトークン）を生成します。LLM が単語 "fast" を生成し、次のラウンドのプロンプトが "Time flies fast" になったと仮定しましょう。これは以下の図で示されています： ![image](https://substackcdn.com/image/fetch/$s_!gBu0!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F06c2f011-ce16-4832-a3aa-4927703fb752_1259x877.png) この図は、各生成ステップにおいて LLM が既知のトークン（"Time" と "flies"）に対するキーおよびバリューベクトルを再計算する様子を示しています。3 番目のトークン（"fast"）を生成する際、モデルは k(1)/v(1) および k(2)/v(2) ベクトルを再利用するのではなく、再度同じものを再計算します。この繰り返される計算は、自己回帰的デコーディング中に KV キャッシュを使用しないことの非効率性を浮き彫りにしています。 前述の 2 つの図を比較すると分かる通り、最初の 2 つのトークンに対するキーおよびバリューベクトルは完全に同一であり、次のトークンのテキスト生成ラウンドごとにそれらを再計算するのは無駄です。 さて、KV キャッシュの考え方は、以前に生成されたキーベクトルとバリューベクトルを保存して再利用するキャッシュ機構を実装し、これにより不要な再計算を防ぐことです。 LLM によるテキスト生成（KV キャッシュなしおよびあり） 前節で基本概念を確認した後、具体的なコード実装を見る前に、もう少し詳細を見ていきましょう。"Time flies fast"というテキストを KV キャッシュなしで生成する場合、以下のように考えることができます： ![image](https://substackcdn.com/image/fetch/$s_!z-sX!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F5fdd4b41-96e6-40c2-baf3-aedfeee8d1de_741x194.png) 冗長性に注目してください：トークン"Time"と"flies"は、新しい生成ステップごとに再計算されています。KV キャッシュは、以前に計算されたキーベクトルとバリューベクトルを保存して再利用することで、この非効率性を解消します。 まず、モデルが入力トークンのキーベクトルとバリューベクトルを計算し、キャッシュに格納します。 次に、生成される各新しいトークンに対して、モデルはその特定のトークンに対するキーベクトルとバリューベクトルのみを計算します。 以前に計算されたベクトルはキャッシュから取得され、冗長な計算が回避されます。 以下の表は、計算およびキャッシュのステップと状態を要約しています： ここでの利点は、「Time」は一度計算されて二度再利用され、「flies」は一度計算されて一度再利用されることです。（これは単純化のための短いテキスト例ですが、テキストが長くなるほど、すでに計算済みのキーと値をより多く再利用できることが直感的に理解でき、それによって生成速度が向上します。） 以下の図は、KV キャッシュありなしの両方を並べて比較した生成ステップ 3 を示しています。 ![image](https://substackcdn.com/image/fetch/$s_!PjfC!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F78382a83-f634-4cfa-92b9-bbea30c61a60_841x926.png) KV キャッシュありなしのテキスト生成を比較。上部のパネル（キャッシュなし）では、各トークンステップごとにキーと値ベクトルが再計算され、冗長な演算が生じます。下部のパネル（キャッシュあり）では、高速化のために以前に計算されたキーと値を KV キャッシュから取得し、再計算を回避します。 したがって、コードで KV キャッシュを実装したい場合は、通常通りキーと値を計算した上で、次回ラウンドでそれらを取得できるように保存するだけです。次のセクションでは、具体的なコード例を用いてこれを説明します。 スクラッチから実装する KV キャッシュ KV キャッシュを実装する方法は多数あり、その主な考え方は、各生成ステップにおいて新たに生成されたトークンのキーおよびバリューテンソルのみを計算することです。 私はコードの可読性を重視したシンプルな実装を選びました。実装方法を理解するには、コードの変更点をスクロールして確認するのが最も簡単だと考えています。 GitHub で公開している 2 つのファイルは、KV キャッシュあり・なしの両方でゼロから LLM を実装する自己完結型の Python スクリプトです: gpt_ch04.py: 私の『Build a Large Language Model (From Scratch)』書籍の第 3 章および第 4 章から抜粋した自己完結型コードで、LLM の実装と単純なテキスト生成関数の実行を行います。 gpt_with_kv_cache.py: 上記と同じ内容ですが、KV キャッシュを実装するための必要な変更が加えられています。 KV キャッシュ関連のコード変更を確認するには、以下のいずれかの方法をとることができます: a. gpt_with_kv_cache.py ファイルを開き、# NEW セクションとしてマークされた新しい変更箇所を探します: ![image](https://substackcdn.com/image/fetch/$s_!KVf4!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F3338a3b0-c3ad-4d37-9d3f-15db18db51ff_982x881.png) b. 任意のファイル差分ツールを使用して、2 つのコードファイルを比較し変更点を確認します: ![image](https://substackcdn.com/image/fetch/$s_!qPKJ!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fccfd3b14-df6e-4988-b986-851b4b207eae_1468x861.png) さらに、実装の詳細を要約するために、以下の小節で簡単なウォークスルーを行います。 1. キャッシュバッファの登録 MultiHeadAttention のコンストラクタ内では、ステップを超えて結合されたキーと値を保持する 2 つのバッファ、cache_k と cache_v を追加します: self.register_buffer("cache_k", None) self.register_buffer("cache_v", None) (バッファについてさらに学びたい場合は、YouTube で「Understanding PyTorch Buffers」という動画を公開しています。) 2. use_cache フラグを使用したフォワードパス 次に、MultiHeadAttention クラスの forward メソッドを拡張し、use_cache 引数を受け付けるようにします: def forward(self, x, use_cache=False): b, num_tokens, d_in = x.shape keys_new = self.W_key(x) # Shape: (b, num_tokens, d_out) values_new = self.W_value(x) queries = self.W_query(x) #... if use_cache: if self.cache_k is None: self.cache_k, self.cache_v = keys_new, values_new else: self.cache_k = torch.cat([self.cache_k, keys_new], dim=1) self.cache_v = torch.cat([self.cache_v, values_new], dim=1) keys, values = self.cache_k, self.cache_v else: keys, values = keys_new, values_new ここでキーと値の保存および取得は、KV キャッシュ (Key-Value Cache) の核心的なアイデアを実装しています。 保存 具体的には、`if self.cache_k is None:` を通じてキャッシュを初期化した後、`self.cache_k = torch.cat(...)` および `self.cache_v = torch.cat(...)` を用いて、新たに生成されたキー（keys）と値（values）をそれぞれキャッシュに追加します。 取得 次に、`keys, values = self.cache_k, self.cache_v` によって、キャッシュから保存されていた値とキーを取得します。 これが KV キャッシュの核心的な格納・取得メカニズムです。続くセクション 3 と 4 では、実装における細部の処理を担当します。 3. キャッシュのクリア テキスト生成を行う際、2 つの別々のテキスト生成呼び出しの間には、必ずキーと値のバッファをリセットする必要があります。これを怠ると、新しいプロンプトのクエリ（queries）が前のシーケンスから残された古くなったキーに注意を向けてしまい、モデルが無関係な文脈に依存して一貫性のない出力を生み出す原因となります。これを防ぐため、後でテキスト生成呼び出しの間で使用できるように、`MultiHeadAttention` クラスに `reset_kv_cache` メソッドを追加します： ```python def reset_cache(self): self.cache_k, self.cache_v = None, None ``` 4. 全体モデル内での use_cache の伝播 `MultiHeadAttention` クラスへの変更を適用した上で、今度は `GPTModel` クラスを変更します。まず、トークンインデックスの位置追跡をインストラクターに追加します： ```python self.current_pos = 0 ``` これは、増分的な生成セッション中にモデルがすでにキャッシュしているトークンの数を記憶する単純なカウンターです。 次に、1 ラインのブロック呼び出しを明示的なループに置き換え、各トランスフォーマーブロックに use_cache を渡します: def forward(self, in_idx, use_cache=False): # ... if use_cache: pos_ids = torch.arange( self.current_pos, self.current_pos + seq_len, device=in_idx.device, dtype=torch.long ) self.current_pos += seq_len else: pos_ids = torch.arange( 0, seq_len, device=in_idx.device, dtype=torch.long ) pos_embeds = self.pos_emb(pos_ids).unsqueeze(0) x = tok_embeds + pos_embeds # ... for blk in self.trf_blocks: x = blk(x, use_cache=use_cache)use_cache=True を設定した場合、上記の動作は、self.current_pos から開始して seq_len 分だけステップを進めることを意味します。その後、カウンターをインクリメントして、次のデコーディング呼び出しが中断した場所から継続できるようにします。 self.current_pos の追跡が必要な理由は、新しいクエリがすでに保存されているキーと値の直後に正しく配置される必要があるからです。カウンターを使用しない場合、各新しいステップは位置 0 から再び開始することになり、モデルは新しいトークンを以前のトークンと重複しているとみなしてしまいます。(あるいは、offset = block.att.cache_k.shape[1] を通じて追跡することも可能です。) 上記の変更により、TransformerBlock クラスにも use_cache 引数を受け取るための小さな修正が必要になります: def forward(self, x, use_cache=False): # ... self.att(x, use_cache=use_cache)最後に、利便性のために GPTModel にモデルレベルのリセットを追加し、すべてのブロックキャッシュを一度にクリアします: def reset_kv_cache(self): for blk in self.trf_blocks: blk.att.reset_cache() self.current_pos = 05. 生成におけるキャッシュの使用 GPTModel、TransformerBlock、および MultiHeadAttention への変更により、ついに KV キャッシュを単純なテキスト生成関数でどのように使用するかを示します: def generate_text_simple_cached( model, idx, max_new_tokens, use_cache=True ): model.eval() ctx_len = model.pos_emb.num_embeddings # max sup. len., e.g. 1024 if use_cache: # Init cache with full prompt model.reset_kv_cache() with torch.no_grad(): logits = model(idx[:, -ctx_len:], use_cache=True) for _ in range(max_new_tokens): # a) pick the token with the highest log-probability next_idx = logits[:, -1].argmax(dim=-1, keepdim=True) # b) append it to the running sequence idx = torch.cat([idx, next_idx], dim=1) # c) feed model only the new token with torch.no_grad(): logits = model(next_idx, use_cache=True) else: for _ in range(max_new_tokens): with torch.no_grad(): logits = model(idx[:, -ctx_len:], use_cache=False) next_idx = logits[:, -1].argmax(dim=-1, keepdim=True) idx = torch.cat([idx, next_idx], dim=1) return idx このコードでは、c) のステップで `logits = model(next_idx, use_cache=True)` として新しいトークンだけをモデルに渡しています。キャッシュを使用しない場合、保存されたキー（keys）や値（values）を再利用できないため、`logits = model(idx[:, -ctx_len:], use_cache=False)` として入力全体をモデルに渡すことになります。 簡単なパフォーマンス比較 KV キャッシュについて概念的に解説した後、大きな疑問は、この仕組みが実際には小規模な例においてどれほど効果的に機能するかという点です。実装を試すために、前述の 2 つのコードファイルを Python スクリプトとして実行できます。これにより、4 トークンのプロンプト「Hello, I am」から開始して、1.24 億パラメータの小規模な大規模言語モデル（LLM）が 200 個の新しいトークンを生成します: pip install -r https://raw.githubusercontent.com/rasbt/LLMs-from-scratch/refs/heads/main/requirements.txt python gpt_ch04.py python gpt_with_kv_cache.py M4 チップ（CPU）を搭載した Mac Mini での結果は以下の通りです: ![image](https://substackcdn.com/image/fetch/$s_!LWIV!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F8d75a600-22e3-4edf-b642-44700dc7a2db_743x160.png) ご覧の通り、わずか 1.24 億パラメータの小規模モデルと短い 200 トークンのシーケンス長であっても、すでに約 5 倍の高速化が実現されています。（なお、この実装はコードの可読性を最適化したものであり、CUDA や MPS ランタイムでの速度を最適化したものではありません。後者の場合、テンソルを再構築して結合するのではなく、事前にテンソルを確保する必要があります。） 注記：どちらの場合もモデルは「意味不明なテキスト」を生成します。具体的には以下のようなテキストです: 出力テキスト: Hello, I am Featureiman Byeswickattribute argue logger Normandy Compton analogous bore ITVEGIN ministriesysics Kle functional recountrictionchangingVirgin embarrassedgl ... これは、まだモデルを訓練していないからです。次の章ではモデルの訓練を行い、訓練済みモデルで KV キャッシュを使用して一貫性のあるテキストを生成できます（ただし、KV キャッシュは推論時のみ使用することを意図しています）。ここではコードを単純に保つために、未訓練のモデルを使用しています。 さらに重要なのは、gpt_ch04.py と gpt_with_kv_cache.py の両方の実装が全く同じテキストを生成することです。これは、KV キャッシュが正しく実装されていることを示しています——インデックス付けのミスは結果の乖離につながる可能性があり、非常に簡単に起こり得ます。 Ahead of AI をお読みいただきありがとうございます！新しい投稿を受け取り、私の活動を支援するために無料で購読してください。 KV キャッシュの利点と欠点 シーケンス長が増加するにつれて、KV キャッシュのメリットとデメリットは以下の通りより顕著になります： [良い] 計算効率の向上：キャッシュを使用しない場合、t ステップでのアテンションは新しいクエリを t 個の以前のキーと比較する必要があり、累積的な作業量は二次関数的に O(n²) でスケーリングします。キャッシュを使用すると、各キーと値は一度計算されその後再利用されるため、ステップごとの総複雑度が線形 O(n) に削減されます。 [悪い] メモリ使用量の線形的増加：新しいトークンごとに KV キャッシュに追加されます。長いシーケンスや大規模な LLM の場合、累積的な KV キャッシュは大きくなり、(GPU) メモリの大幅な、あるいは実質的に許容できない量を消費する可能性があります。回避策として KV キャッシュを切り捨てることもできますが、これによりさらに複雑さが増します（ただし、LLM をデプロイする際には、その価値がある場合も十分にあります）。 KV キャッシュの実装の最適化 上記の概念実装は明確さを保つために役立ち、主にコードの可読性と教育目的に焦点を当てていますが、実際の運用環境（特に大規模モデルや長いシーケンス長の場合）でデプロイする場合は、より慎重な最適化が必要です。 キャッシュのスケーリング時の一般的な落とし穴 メモリの断片化と繰り返しの割り当て：前述のように torch.cat を使用してテンソルを連続的に結合すると、頻繁なメモリ割り当てと再割り当てによりパフォーマンスのボトルネックが生じます。 メモリ使用量の線形的増加：適切に処理されない場合、KV キャッシュのサイズは非常に長いシーケンスに対して実用的ではなくなります。 ヒント 1: メモリの事前割り当て テンソルを繰り返し結合するのではなく、予想される最大シーケンス長に基づいて十分に大きなテンソルを事前に割り当てることもできます。これによりメモリ使用量が一定に保たれ、オーバーヘッドが削減されます。擬似コードでは以下のように記述できます。 # キーと値の事前割り当て例 max_seq_len = 1024 # 予想される最大シーケンス長 cache_k = torch.zeros( (batch_size, num_heads, max_seq_len, head_dim), device=device ) cache_v = torch.zeros( (batch_size, num_heads, max_seq_len, head_dim), device=device )推論時には、これらの事前に割り当てられたテンソルのスライスに単に書き込むだけで済みます。 ヒント 2: スライディングウィンドウによるキャッシュの切り捨て GPU メモリを爆発させないために、動的な切り捨てを伴うスライディングウィンドウアプローチを実装できます。スライディングウィンドウを通じて、キャッシュには最後の window_size トークンだけを保持します: # スライディングウィンドウキャッシュの実装 window_size = 512 cache_k = cache_k[:, :, -window_size:, :] cache_v = cache_v[:, :, -window_size:, :] Optimiz

スクラッチから理解するLLMのKVキャッシュの仕組みと実装

背景や根拠まで確認しますか？

調べる

選ぶ

サイト