GPT-2で追うprefill・decodeとKVキャッシュ

はじめに

LLMの処理は、大きく学習と推論の2つに分けられます。

この記事で扱うKVキャッシュは推論側の話です。推論はさらに、prefillとdecodeという2つの段階に分かれます。ローカルLLMを触っていると、どちらもよく目にする言葉だと思います。2つの段階は、次のように性質が異なります。

decodeがmemory boundということは、decodeの速度はGPUの計算性能よりも、メモリからどれだけ速くデータを読み出せるか（メモリ帯域）で決まるということです。そのため、LLMを動かすハードウェアでは、メモリの容量と帯域、そしてメモリをどう配置するかが重要になります。まずGPT-2でprefillとdecodeを順に追い、KVキャッシュが何を使い回しているのかを見たうえで、最後に、KVキャッシュや重みを置くメモリが、実際のハードウェアでどう配置されているかを確認します。

GPT-2型モデルの推論処理を追う

この記事で追う例

GPT-2に The cat sat on という文章を入力し、続きとして the mat and purred が生成される場面を例に、処理の流れを追います（トークンの区切りと生成される文章は説明用の例です）。

prefillでは、プロンプトの4トークンをまとめて処理し、最初の1トークン the を生成します。decodeでは、直前に生成した1トークンだけを入力して、次の1トークンを生成します。このとき、過去のトークンの計算結果をどう使い回すかが、KVキャッシュの話につながります。

推論処理の流れ

次の図は、推論で行われる処理の流れです。

上の図には多くの部品がありますが、この中で過去のトークンを参照するのは、次の図で赤枠をつけたAttentionだけです。ほかの部品は1トークンずつ独立に計算するので、decodeでも新しいトークンの1行だけを計算すれば済みます。そこでここからは、prefillとdecodeのそれぞれでAttentionが何を計算しているのかに絞って追います。

Transformer Blockに \times N 層とあるように、点線のブロックは1層目、2層目、…、N 層目と同じ構造で積み重なります。1層目の出力が2層目の入力になり、最後の N 層目の出力が後ろのLayerNormとLinearに渡ります。Attentionは各層に1つずつあるので、KVキャッシュも層ごとに持つことになります。

この先の説明は、次の動画の場面を切り出しながら進めます。prefillからdecodeまで、Attentionの計算が層ごとにどう進み、KVキャッシュがあると何を使い回せるのかを動画にしました。左がKVキャッシュあり、右がなしです。色の意味は次のとおりです。

https://youtu.be/eBtSIzcETe0

動画のトークンの下にある 1〜8 は位置の番号で、q_i・k_i・v_i の添字 i と対応しています。以降の式も同じ番号を使います。

prefillの流れ

まずはprefillです。プロンプトの4トークン The cat sat on をまとめて入力し、次のトークン the を求めます。先ほどの全体図で赤枠をつけたMulti-Head Attentionの中を見ていきます。多少数式が出てきますが、行列の積と内積の計算だけです。

この記事の設定では、Multi-Head Attentionは6つのヘッドで構成されています。以降の式に出てくる記号の読み方は次のとおりです。大きさはprefill（4トークン）のときの値です。

次の図は、そのうち1ヘッド分の動きです^[1]。

上の図は、トークン数を最大の256として描いたものです。ここからは、prefillで実際に入力する4トークンで見ていきます。

上の図は、prefillの入力 The cat sat on の4トークンで描き直したもので、赤枠は Q^h・K^h・V^h を作るところです。1トークンは384個の数値（1行）で表されるので、入力 \tilde{X} は1行目が The、2行目が cat、3行目が sat、4行目が on の 4 \times 384 の行列です。ここにヘッド h の重み W_Q^h・W_K^h・W_V^h（各 384 \times 64）を掛けると、Q^h・K^h・V^h も4トークン分の4行を持つ 4 \times 64 の行列になります。

1行ずつ見ると、各トークンの q_i・k_i・v_i は、そのトークン自身の行 \tilde{x}_i だけから計算されます。

こうして作った Q^h・K^h・V^h の3つが、Attentionの入力です。1ヘッドの出力 O^h は、次のように書けます。Transformerの論文（Attention Is All You Need）もこの書き方で、\mathrm{Attention} の中身は、このあと見るCausal Attentionの計算です。

上の図の赤枠は、作った Q^h・K^h・V^h から1ヘッドの出力 O^h を求めるところ（Causal Attention）です。計算は次の2つの式で表せます。この式がKVキャッシュを理解するうえで一番大事な部分で、KVキャッシュが必要になるのもここです。

2つをまとめると、先ほどの \mathrm{Attention} の中身になります。

まず Q^h (K^h)^\top で、全トークンの組み合わせの内積を計算します。ここで右肩の \top は転置で、行列の行と列を入れ替える操作です。Q^h も K^h も 4 \times 64 なので、そのままでは掛け算できません（左の列数64と右の行数4が合わない）。K^h を転置して 64 \times 4 にすると、内側の64がそろって掛けられるようになり、結果は 4 \times 4 になります。

転置すると、K^h で1行ずつ並んでいた k_1〜k_4 が、1列ずつ並ぶ形になります。そのため、結果の i 行 j 列目は、Q^h の i 行目 q_i と、K^h の j 行目 k_j の内積（64個の数値を掛けて足したもの）になります。

これにマスク M を足すと、次のようになります（以下の表では \sqrt{64} で割る処理を省略しています）。

マスク M で未来のトークンを -\infty にしてから行ごとにsoftmaxを取ると、-\infty の部分は0になり、各行の合計が1の重み A^h になります。

最後に O^h = A^h V^h で V^h の行を重み付きで足し合わせます。たとえば on の出力は o_4 = a_{41} v_1 + a_{42} v_2 + a_{43} v_3 + a_{44} v_4 です。つまり、あるトークンの出力を求めるには、そのトークン自身と過去すべてのトークンの k と v が必要になります。

ここまでの計算を動画で見たのが、次の図（prefill・層1の場面）です。

上の図では、4トークン分の q・k・v と内積（青）をすべて計算し、causal maskで右上が -\infty になっています。右のKVキャッシュでは、層1に k_1〜k_4・v_1〜v_4 が保存されています。同じ計算を層2・層3でも行い、それぞれの層のKVキャッシュに保存します。

上の図は、3層を通ったあとの出力の場面です。最後の位置（on）の出力から、次のトークン the を選びます。この時点で、KVキャッシュには3層それぞれに4トークン分の k・v がそろっています。

decodeの流れ

続いてdecodeです。decodeの1ステップ目では、直前に生成した the（位置5）だけを入力し、次のトークン mat を求めます。以降の式では、\textcolor{#146B53}{\boxed{k}} がKVキャッシュから読む値、\textcolor{#2D6BD6}{\boxed{\bm{k}}} が今回計算する値です。

上の図は、動画のdecode 1・層1の場面です。計算するのは位置5の1行（青）だけで、k_1〜k_4・v_1〜v_4 はKVキャッシュから読み出しています（緑）。数式で追うと次のとおりです。

新しく計算するのは、入力した the の1行分の q・k・v だけです。

過去4トークン分の k・v はprefillで計算済みなので、KVキャッシュから読み出し、今回の1行を末尾に足します。

内積を取るのも the の1行だけです。最新のトークンは過去をすべて見てよいので、マスクも要りません。

出力 o_5 は、キャッシュから読んだ v_1〜v_4 と、今回計算した v_5 を混ぜて作ります。

つまり、the が過去の文脈を取り込む材料は、計算済みの k・v だけです。層2以降の k_j・v_j には、位置 j までの文脈がすでに計算されて入っています。

この計算結果は後からトークンが増えても変わらないので、KVキャッシュに保存したものをそのまま使えます。

上の図は、同じdecode 1で層3まで進んだ場面です。次のトークンの予測に使うのは、最後の位置の出力 o_5 だけです^[2]。層1・層2でも、次の層に渡す必要があるのは位置5の行だけなので、過去の q_1〜q_4 や o_1〜o_4（灰色）は使いません。過去のトークンについて残しておく必要があるのは、Attentionで参照される k・v だけです。これが、KVキャッシュに q を入れない理由でもあります。

上の図は、同じdecode 1・層1をKVキャッシュなしで見た場面です。KVキャッシュがなければ、過去の k・v はどこにも残っていません。保存しておいたトークンID列 The cat sat on the をもう一度Embedから通し、全層で 5 \times 5 をまるごと計算し直すことになります。赤（\textcolor{#C0392B}{\boxed{k}}）の部分がこの計算し直した値で、prefillで計算した値とまったく同じです。次のトークンを決めるのに必要なのは、最後の1行（青）だけです。

上の図は、decode 3まで進んだ場面です。decodeを進めるたびに、KVキャッシュは1トークン分ずつ伸びていき、毎回計算するのは青の1つだけです。次の表は1層・1ヘッド分なので、実際には V^h も含めて、層の数 × ヘッドの数だけ同じようなキャッシュを持ちます。

KVキャッシュとハードウェアのメモリ構成

ここまで見てきたように、decodeでは1トークンごとに、モデルの重みと、過去のトークンの計算結果であるKVキャッシュを読み出します。KVキャッシュはコンテキスト長に比例して大きくなるため、メモリの容量と帯域の両方を圧迫します。では、実際のハードウェアではメモリがどう配置されているのかを見てみます。ここでは、代表的な4つの構成について、メモリの配置と帯域から見たざっくりとした特性を描きます。

この先の4枚は、メモリの配置と接続の違いを比べるための自作の模式図です。青は計算チップ、オレンジはメモリ、グレーの枠はパッケージや基板を表します。矢印の太さは帯域の違いを概念的に表したもので、実際の比率や配線を再現したものではありません。

Vera Rubin（データセンター向け）

Rubin GPUのパッケージ内に、HBM4（メモリチップを縦に積み、非常に広いバス幅でGPUとつなぐ高帯域メモリ）が配置されています。Vera CPU側には別途LPDDR5Xがあり、NVLink-C2Cで整合性を保ったまま相互にアクセスできますが、物理的には別のメモリです。次の図ではRubin GPUを1基だけ描いています。

構成の出典：NVIDIA公式・Vera Rubinの技術解説。プラットフォーム全体ではなく、CPU・GPUとメモリの関係を抜粋。

GB10 Grace Blackwell（DGX Spark）

DGX Sparkに搭載されるSoCで、Grace CPU（Arm）とBlackwell GPUがNVLink-C2Cでつながり、128GBのLPDDR5Xを共有します（ユニファイドメモリ）。大きなモデルも載せられる一方、帯域は273GB/sと控えめです。

構成・名称の出典：NVIDIA公式・GB10の発表。容量・帯域：DGX Spark仕様。矢印は論理的なアクセス関係。

Apple Silicon（Mac Studio）

GB10と同じく、CPUとGPUが同じ物理メモリを共有するユニファイドメモリ構成です（内部の接続方式まで同じとは限りません）。違うのは帯域の大きさで、Mac StudioのM5 Ultraは最大512GBのユニファイドメモリを1.2TB/sで読み出せます。GB10（128GB・273GB/s）と比べると容量は4倍、帯域は約4.4倍で、グラフィックボードのRTX 5090（VRAM 32GB・約1.8TB/s）に近い帯域を、より大きな容量で持つ形です。

ただし、メモリはOSやほかのアプリとも共有するため、搭載容量のすべてをLLMに使えるわけではありません。

容量・帯域の出典：Apple公式・Mac Studioの技術仕様。配置は模式的なもので、実際のチップ構成を再現したものではない。

グラフィックボード（RTX 5090）

一般的なPCに挿すグラフィックボードでは、基板上にGPU専用のVRAM（GDDR）があり、CPU側のメインメモリとは別です。たとえばRTX 5090は32GBのGDDR7で、帯域は約1.8TB/sです。広帯域ですが、GPUに置けるモデルやKVキャッシュの量はVRAM容量で頭打ちになります。

製品例の仕様：NVIDIA公式・RTX 5090のメモリ容量と帯域。メモリの個数・配置は特定製品を再現したものではない。

では、VRAMが足りなければマザーボード側のメインメモリ（DDR5）を足せばよいのでしょうか。llama.cppなどでは、VRAMに収まらない部分をメインメモリに置くこと（オフロード）ができますが、一般に生成速度は大きく落ちます。ざっくり言うと、メインメモリは計算するチップから見て読み出しの経路が細いためです。

たとえば重みのうち20GBをメインメモリに置くと、1トークンごとにその20GBを約90GB/sで読むだけで約0.2秒かかります。合計容量が足りていても、速く読めるメモリに載っていなければ、decodeは速くなりません。

さいごに

GPT-2型の小さなモデルを例に、prefillとdecodeの流れを追いながら、KVキャッシュが何を使い回しているのかを見てきました。

Attentionは難しそうに見えますが、中身は行列の積と内積、それにsoftmaxくらいで、今回のように4トークン・1ヘッドに絞れば手で追える計算です。そして、causal maskによって過去のトークンの k・v が後から変わらない、という性質さえ押さえれば、KVキャッシュがなぜ成り立つのかも自然に理解できます。

また、prefillは計算性能（compute）、decodeはメモリ帯域が効く、という違いを意識しておくと、ローカルLLMを扱うときの判断がしやすくなります。どのハードウェアを選ぶか、どこまで量子化するか、コンテキスト長をどれくらいにするか、といった判断の多くは、読み出すデータの量とメモリ帯域の話に置き換えて考えられるからです。実際の速度はCUDAカーネルや推論エンジンの実装にも左右されますが、大まかな見当をつけるうえで、この見方はよく効きます。

今回は記事の長さの都合で、説明用の設定を図と数式で追うところまでにとどめました。今後は、実際にこの記事の設定（埋め込み次元384・6ヘッド・3層）程度の小さなモデルを作り、KVキャッシュの有無による速度やメモリ使用量の違いを検証してみたいと思います。

参考文献

この記事を書くにあたって、特にAttentionやTransformerの仕組みの理解では、次の書籍を大いに参考にしました。トークナイザからTransformer、事前学習・事後学習までを実装しながら学べる一冊です。

• 斎藤 康毅『ゼロから作るDeep Learning ❻ ―LLM編』オライリー・ジャパン、2026年

https://www.oreilly.co.jp/books/9784814401611/