Agentic OS 技術スタックを下から読む 第4回：KV キャッシュを小さくする ―― 頭を減らすか、圧縮するか

前回までで、推論コストの見え方を下から見てきた。

計算量だけを見ていると、モデルは「大きな行列計算をたくさん回すもの」に見える。けれど実際の推論では、メモリ帯域、バッチサイズ、そして KV キャッシュが効いてくる。特に長い文脈になると、値段を決める主役は KV キャッシュに寄っていく。

しかも KV キャッシュは、扱いにくい費用である。

重みの読み出しは、同じモデルを同時に使うリクエスト同士である程度割り勘できる。バッチを組めば、一度読んだ重みを複数のトークン生成に使える。ところが KV キャッシュは違う。各ユーザー、各会話、各エージェントが持つ過去文脈そのものなので、他のリクエストと簡単には共有できない。

長い文脈を持つほど、保存すべき過去が増える。並列に走るエージェントが増えるほど、それぞれの過去も増える。だから Agentic OS のように、長い記憶と多数の実行単位を前提にする世界では、KV キャッシュは単なる実装上の細部ではなくなる。上の層の設計を、下から静かに制約する。

では、料金の出方を受け入れるだけでよいのか。

自然な次の問いはこうなる。モデルの構造そのものを変えて、そもそも KV キャッシュを小さくできないか。ここ数年の注意機構の設計は、この問いにかなりの力を注いできた。

今回は、その中でも二つだけを見る。

ひとつは、キーとバリューの「頭」を減らす道である。もうひとつは、キーとバリューの「中身」を圧縮する道である。

KV キャッシュは、注意機構の副産物

まず、削る前の形を押さえておく。

Transformer の注意機構では、入力された各トークンから、三つの役割を持つ表現を作る。クエリ、キー、バリューである。

クエリは、いま何を探しているかを表す。キーは、過去の各トークンがどんな手がかりを持っているかを表す。バリューは、実際に取り出して混ぜる中身を表す。

次のトークンを生成するとき、モデルはいまの位置のクエリを作る。そして、過去すべてのトークンのキーと照らし合わせる。どの過去トークンを見るべきかを重みとして決め、その重みにしたがって過去のバリューを混ぜる。

これが注意の基本形である。

問題は、次のトークンを一つ出すたびに、過去すべてのキーとバリューを作り直していたら遅すぎることだ。だから一度作った過去トークンのキーとバリューを、メモリ上に保存しておく。次の生成では、それを再利用する。

この保存された過去のキーとバリューが、KV キャッシュである。

ここでさらに、注意機構には「頭」が複数ある。多頭注意と呼ばれる形では、同じトークン列に対して、複数の視点から注意を計算する。ある頭は直前の語との関係を見やすいかもしれない。別の頭は文の遠い位置にある参照関係を見やすいかもしれない。実際にはそこまで単純に役割分担できるわけではないが、直感としては、複数の見方を並べて持つ仕組みだと考えればよい。

標準的な多頭注意では、それぞれの頭が自分専用のキーとバリューを持つ。つまり、頭が多いほど、過去トークンごとに保存するキーとバリューも厚くなる。

文脈長が二倍になれば、保存する過去もおおむね二倍になる。バッチ内の会話数が増えれば、それぞれが別の KV キャッシュを持つ。頭が多ければ、各トークンに付く保存物も増える。

ここに、長文脈の値段が出てくる。

ならば、削る場所も見えてくる。保存するキーとバリューの量を減らせばよい。

そのための第一の道が、頭を減らすことである。

道その一：頭を減らす

いちばん素直な発想は、キーとバリューの頭を減らすことだ。

ただし、すべての頭を単純に減らすわけではない。クエリの頭は多く残す。つまり、いま何を探すかという視点は細かく保つ。その一方で、過去側に保存しておくキーとバリューの種類を少なくする。

複数のクエリ頭が、同じキーとバリューの組を共有する。これが grouped-query attention、つまりグループ化されたクエリ注意の中心にある考え方である。以下では GQA と呼ぶ。

標準的な多頭注意では、たとえばクエリ頭が 32 個あれば、キーとバリューの頭も 32 個ある。GQA では、クエリ頭は 32 個のままでも、キーとバリューは 8 組だけにする、といった構成ができる。すると、4 個のクエリ頭が 1 組のキーとバリューを共有する。

このとき、モデルは「見る側」の視点を完全には捨てていない。クエリは複数あるので、いまの位置から過去を見る角度はまだ多い。しかし、過去側に保存する表現の種類は減っている。保存する KV キャッシュが薄くなる。

もっと極端にすれば、キーとバリューを 1 組だけにすることもできる。すべてのクエリ頭が同じキーとバリューを見る形である。これは multi-query attention、つまり複数のクエリが一つのキー・バリューを共有する形に近い。

その中間に GQA がある。

大事なのは、これは連続したつまみだということだ。キーとバリューをどこまで共有するかを選べる。共有を強くすれば、KV キャッシュは小さくなる。共有を弱くすれば、標準的な多頭注意に近づき、キーとバリュー側の自由度は増える。

当然、ただ得をするわけではない。

キーとバリューの頭を減らすということは、過去側の表現の種類を減らすということでもある。複数のクエリ頭が同じ過去表現を使うので、頭ごとに別々のキーとバリューを持つ場合に比べると、表現の自由度は下がる。

ただ、その損失がいつも大きいとは限らない。特に小さめから中規模のモデルでは、キーとバリューをある程度共有しても、品質を大きく落とさずに済む場合がある。実装も比較的素直である。保存するテンソルの形が単純に小さくなるので、推論エンジン側でも扱いやすい。学習も安定させやすい。

だから GQA は、いまでも強い選択肢である。

新しいモデルだからといって、必ず複雑な圧縮方式を使うわけではない。あえてこの古典的で素直な方法を選ぶこともある。特に、モデルサイズ、学習の安定性、推論実装、対応するハードウェアやランタイムをまとめて考えると、GQA の分かりやすさはかなり大きい。

ここで削っているのは、キーとバリューの「種類」である。

では、種類はあまり減らさず、一つひとつの中身を小さくすることはできないか。これが第二の道である。

道その二：中身を圧縮する

もうひとつの発想は、保存するキーとバリューをそのまま持たないことだ。

各トークンについて、生のキーとバリューをフルサイズで保存するのではなく、もっと小さな潜在表現として保存する。必要になったときに、その小さな表現から注意計算に使う形へ戻す。

この考え方の代表的な形が multi-head latent attention、つまり複数頭の潜在注意である。以下では MLA と呼ぶ。

ここでいう潜在表現とは、モデル内部で使う圧縮された中間表現だと考えればよい。人間が読める要約ではない。キーやバリューとしてすぐ使える形でもない。だが、あとで必要なキーやバリューを再構成するための情報を、小さめの器に詰めておく。

GQA が「キーとバリューの組数を減らす」方法だとすれば、MLA は「保存する一組あたりの中身を小さくする」方法である。

狙っているボトルネックは同じだ。長い文脈で膨らむ KV キャッシュを小さくしたい。しかし、効かせ方が違う。

GQA では、複数のクエリ頭がキーとバリューを共有する。だから保存するキーとバリューの種類が減る。MLA では、保存するもの自体を圧縮された形に変える。あとで使うときに、必要な形へ展開する。

この違いは、長い文脈になるほど効いてくる。

文脈が短ければ、KV キャッシュの節約はそこまで支配的ではないかもしれない。だが、数万トークン、さらにそれ以上の文脈を扱うと、過去トークンごとに何を保存しているかが重くなる。各トークンに付く保存物を小さくできれば、総量は大きく変わる。

モデルが大きくなる場合も同じである。層数や隠れ次元が大きくなり、注意の表現も太くなるほど、KV キャッシュは重くなる。そこで圧縮がうまく働けば、品質を大きく落とさずに、長い文脈をより現実的なコストで扱える可能性がある。

ただし、こちらも無料ではない。

保存するものを圧縮するには、圧縮するための写像が必要になる。使うときには、戻すための写像も必要になる。実装は GQA より複雑になりやすい。推論エンジンも、その構造をきちんと理解して最適化しなければならない。

また、圧縮した表現にどれだけ情報を残すかという設計も難しい。小さくしすぎれば、注意に必要な情報が足りなくなる。大きく残しすぎれば、節約効果が薄れる。モデルの規模、学習データ、目的とする文脈長、推論環境の制約を合わせて調整する必要がある。

MLA は、GQA より洗練された方法に見えるかもしれない。実際、長文脈や大規模モデルでは魅力がある。だが、それは単純な上位互換という意味ではない。複雑さを引き受ける代わりに、KV キャッシュの圧縮率を取りにいく設計である。

ここで削っているのは、キーとバリューの「中身」である。

同じボトルネック、違うつまみ

GQA と MLA は、どちらも同じ壁を見ている。

壁とは、長い文脈で増え続ける KV キャッシュである。

ただし、触っているつまみが違う。

GQA は、キーとバリューの頭数を減らす。複数のクエリ頭に、少数のキーとバリューを共有させる。つまり、保存する種類を減らす。

MLA は、保存するキーとバリューの中身を圧縮する。生の形ではなく、小さな潜在表現として持ち、必要なときに戻す。つまり、一つあたりを小さくする。

この違いを押さえると、「どちらが上か」という問いが少し雑に見えてくる。

小さめのモデルでは、GQA の素直さが効きやすい。仕組みが分かりやすく、実装しやすく、推論エンジンにも乗せやすい。学習や運用の見通しも立てやすい。多少の自由度を削っても、全体としてよいバランスになることがある。

一方で、規模が上がり、長い文脈を本格的に扱うほど、MLA のような圧縮方式が魅力を持ちやすくなる。各トークンに保存するものを小さくできれば、文脈長に比例して効いてくる。特に、長文脈を売りにするモデルでは、この差が推論コストや同時処理数に響きやすい。

ただし、規模が大きければ必ず MLA、というほど単純でもない。モデルの目的、既存の学習レシピ、推論基盤、対応するカーネル、運用上の安定性によって判断は変わる。新しいモデルでも、あえて GQA のような古典的な共有方式だけを採ることはある。逆に、ある段階から圧縮方式へ切り替えることもある。

ここで見ているのは、モデル設計上の流行語ではない。

長い文脈をどう安く持つかという、かなり地味な設計判断である。

これは性能ではなく、出し方の選択

ここで誤解しやすい点がある。

GQA や MLA は、モデルを急に賢くする魔法ではない。注意機構の形を変えたからといって、未知の推論能力が突然生まれるわけではない。

もちろん、同じ予算の中でより大きなモデルを動かせる、より長い文脈を入れられる、より多くのリクエストを並列にさばける、という間接的な効果はある。結果としてユーザーから見える性能が上がることはある。

しかし、機構として直接狙っているのはそこではない。

狙いは、同じくらいの能力を持つモデルを、長い文脈でどれだけ現実的に出せるかである。長い入力を受けてもメモリが破綻しないか。多数の会話を同時に処理しても KV キャッシュが詰まらないか。推論の単価が許容できる範囲に収まるか。

前回の言い方に戻せば、KV キャッシュは割り勘しにくい費用だった。GQA と MLA は、その費用をモデル構造の側から削る試みである。

GQA は、保存するキーとバリューの種類を減らす。MLA は、保存する中身を圧縮する。どちらも、請求書が来てから工夫するのではなく、請求書の元になる構造そのものを変えている。

この見方をすると、注意機構の設計はずっと実務的に見えてくる。

長文脈対応という言葉だけを見ると、単に「たくさん読めるモデル」の話に見える。だが下から見ると、それは「たくさん読んだ過去を、各層でどう保存し続けるか」の話である。保存の仕方を変えなければ、長い文脈はすぐに高くなる。

Agentic OS への含意

Agentic OS の上の層では、モデルを部品として選ぶ。

どのモデルを使うか。どのくらいの文脈長を前提にするか。エージェントを何本並列に走らせるか。会話履歴や作業履歴をどこまで文脈に残すか。これらは一見、アプリケーション側の設計判断に見える。

しかし、下には注意機構の設計がある。

長い記憶を文脈に持たせたいなら、そのモデルが KV キャッシュをどう持っているかが効いてくる。多数のエージェントを同時に走らせたいなら、それぞれが抱える KV キャッシュの厚みが効いてくる。長いタスクを途中で捨てずに進めたいなら、過去を保存し続けるコストが効いてくる。

上の層がモデルを選ぶとき、実は注意の設計も一緒に選んでいる。

GQA のようにキーとバリューを共有するモデルを選ぶのか。MLA のように中身を圧縮するモデルを選ぶのか。あるいは別の方式を組み合わせたモデルを選ぶのか。その違いは、ベンチマークの点数だけでは見えにくい。だが、長文脈、並列実行、常時稼働のエージェントを考えると、運用コストとして表に出てくる。

今回見たのは、KV キャッシュを小さくする二つの道だった。

ひとつは、頭を減らす。より正確には、キーとバリューを複数のクエリ頭で共有する。

もうひとつは、中身を圧縮する。生のキーとバリューではなく、小さな潜在表現として保存する。

どちらも、長い文脈の値段を下げるための構造上の工夫である。

次回は、同じ KV の壁に対するもう一つの攻め方に進む。一つのトークンが、過去のどこまでを見にいくのか。その範囲そのものを絞る道である。固定された窓で見る方法と、学習によって見る場所を選ぶ方法を、今回とは別のつまみとして見ていく。