何が公開されたか

Ahead of AIに、LLMの推論を高速化する「KVキャッシュ」をゼロから実装するチュートリアルが掲載されました。書籍『Build a Large Language Model (From Scratch)』の第3〜4章で扱われるGPT実装をベースに、KVキャッシュなし版(gpt_ch04.py)とあり版(gpt_with_kv_cache.py)の2本のPythonスクリプトがGitHubで公開されています。

なぜKVキャッシュが効くのか

LLMは1トークンずつ逐次生成します。プロンプト「Time」から「flies」を生成し、次は「Time flies」を再投入して「fast」を出す——という具合に、何もしなければ過去トークンの鍵(K)・値(V)ベクトルを毎ステップ計算し直すことになります。KVキャッシュは、学習済み重み行列W_k・W_vで一度計算したKとVを保持し、次ステップでは新規トークンぶんだけ計算する仕組みです。「Time flies fast」の例なら「Time」は1回計算して2回再利用、「flies」は1回計算して1回再利用。系列が長くなるほど節約は積み上がります。

実装の要点

改修ポイントは限定的です。MultiHeadAttentionのコンストラクタにself.register_bufferでcache_kとcache_vを追加し、forwardにuse_cache引数を新設。キャッシュが存在すればtorch.catでdim=1に沿って新しいK/Vを連結し、Noneなら初期化します。生成呼び出しをまたいで古い鍵にクエリが注意を向けないよう、reset_cacheでバッファをクリアする設計です。

GPTModel側にはself.current_posカウンタを置き、use_cache=True時にこの位置からpos_idsを発番、呼び出しごとにseq_lenぶん進めます(block.att.cache_k.shape[1]を使う実装もあり得ます)。TransformerBlockはuse_cacheを下位に橋渡しし、モデル全体ではreset_kv_cacheで全ブロックのキャッシュとcurrent_posを一括リセット。これを使うgenerate_text_simple_cachedは「リセット→プロンプト一括処理→max_new_tokensまで逐次生成」という素直なフローで動きます。

トレードオフ

速度向上の代償はメモリ消費とコード複雑度です。著者は同書にKVキャッシュを含めなかった主因として、ctx_len=1024のような長コンテキストでKとVを系列長ぶん保持するメモリ負担を挙げています。さらにキャッシュは推論専用で、学習時は使えません。

💼 事業会社視点:これは自社にどう効くか

自社LLMやLLMラッパー型SaaSを運用する事業会社にとって、KVキャッシュはコスト構造に直結する論点です。OpenAIやAnthropicのAPIを叩くだけのプロダクトであれば事業者側は意識する必要はありませんが、自前GPUでオープンウェイトモデル(Llama系・Qwen系など)をホスティングしている企業——たとえばオンプレ要件のある金融・医療・官公庁向け受託開発、社内チャットボットを内製しているメーカー、長文要約系SaaSなど——では、KVキャッシュの有無と容量設計がそのままGPU時間単価と応答レイテンシに跳ね返ります。

経営者・事業責任者が確認すべきは2点。第一に、自社推論基盤がvLLMやTGIなどKVキャッシュ最適化(PagedAttention等)を標準搭載した推論サーバを採用しているか。素のtransformersで動かしている場合、同じGPUでさばける同時接続数が桁違いに少ない可能性があります。第二に、長コンテキスト対応を売りにする場合、メモリ上限がそのまま「同時セッション数×最大トークン長」の天井になる点。値付けやSLAを決める前に、エンジニアにKVキャッシュ容量と同時接続数のトレードオフ表を出させるべきです。

関連リンク