Agentic OS 技術スタックを下から読む 第20回：推論エンジンの中身 ―― KV を「ページ」として配る

モデルは、それだけでは動かない

第1回から第3回で、長文脈の値段は、かなりの部分を KV キャッシュが決める、と見た。KV は過去文脈のための作業場所であり、リクエストごとに固有である。重みのように、多くの利用者で割り勘しにくい。

前回は、推論基盤が立ち上がるまでを早回しで見た。今回は、立ち上がったあとに進む。実際に推論を回すプログラム、つまりエンジンの中身を見る。

「モデルを手に入れた」と言うとき、実際に手にしているのは、重みの詰まったファイルである。重みとは、学習で得られた大量の数値である。だが、ファイルがあるだけでは何も起きない。入力を受け取り、それを何層もの計算に通し、次のトークンを選び、それをまた入力側に戻して、次を出す。この一連の動きを実際に行うのは、エンジンと呼ぶプログラムである。

推論とは、モデルの構造と、重みと、ハードウェアの実行をつなぐ仕事である。そして、その中で一番やっかいな部分の一つが、過去文脈の置き場、つまり KV キャッシュのメモリ管理である。

エンジンがやること、ざっくり

エンジンの仕事は、大きく見ると三つある。

まず、重みを置き場から GPU のメモリに読み込む。GPU は大量の行列計算を速く行うための装置であり、推論では重みを何度も読む。だから、できるだけ近い場所に置いておきたい。

次に、入力されたプロンプトをまとめて計算に通す。この段階をプリフィルと呼ぶ。プリフィルとは、すでに与えられた複数のトークンを一気に処理し、以後の生成に使う内部状態を作る段階である。例えば、入力が二千トークンあるなら、その二千トークンをまとめて計算に流す。多くのトークンを横に並べて処理できるので、計算機を使い切りやすい。

その後、生成に入る。この段階をデコードと呼ぶ。デコードとは、直前までの文脈を使って、次の一トークンを選び、それを文脈に追加して、また次の一トークンを選ぶ段階である。ここでは、一トークンずつしか進まない。次のトークンが決まらないと、その次の計算を始められないからである。

エンジンはさらに、複数のリクエストをまとめて回す。これをバッチと呼ぶ。バッチとは、複数の入力や生成作業を一つのまとまりとして計算機に渡すことだ。うまくまとめられれば、重みの読み出しを複数リクエストで共有できる。これが、一トークンあたりの値段を下げる。

ただし、プリフィルとデコードは性質が違う。プリフィルは大きな入力をまとめて押し込める。デコードは細かく一歩ずつ進む。その間ずっと持ち歩かれるものがある。過去文脈の置き場である。今回の主役は、この KV キャッシュである。

KV キャッシュという、置き場の問題

生成中、エンジンは過去のトークンを毎回ゼロから計算し直さない。そんなことをすれば、文脈が長くなるほど、同じ計算を何度も繰り返すことになる。

そこで、各層の注意機構が使う中間結果を保存しておく。注意機構とは、現在のトークンが過去のどのトークンをどれだけ参照するかを計算する仕組みである。この参照のために、過去の各トークンから「鍵」と「値」と呼ぶ二種類の表現を作る。鍵は、参照先を探すための見出しのようなものだ。値は、見つけた参照先から取り出す中身のようなものだ。

この鍵と値を、英語の頭文字で K と V と呼ぶ。過去トークンぶんの K と V を捨てずに取っておく場所が、KV キャッシュである。

次の一トークンを出すとき、エンジンは新しいトークンのための情報を計算し、保存済みの K と V を参照する。生成したトークンが一つ増えるたびに、そのトークンの K と V も追加される。だから、文脈が伸びるほど KV キャッシュも増える。

ここで重要なのは、KV キャッシュはリクエストごとに違う、という点である。重みは同じモデルを使う全員で共有できる。だが、ある人のプロンプトと生成履歴から作られた KV は、その人のそのリクエストに固有である。別の人の文脈には、そのまま使えない。

問題は、これを GPU のメモリにどう置くかである。GPU のメモリは大きく見えても、同時に大量のリクエストを抱え、長い文脈を扱うと、すぐに厳しくなる。しかも、ただ容量が足りるかどうかだけではない。空き方が悪いと、まだ空き容量があるのに使えない、という状態が起きる。

素朴なやり方は、一枚板を予約すること

いちばん単純な置き方は、リクエストごとに、KV 用の連続したメモリ領域を確保することだ。連続したメモリ領域とは、途中で切れずに一続きになった場所である。ここでは、一枚板と呼ぶ。

例えば、あるリクエストに最大四千トークンまで対応したいとする。すると、エンジンはそのリクエストのために、四千トークンぶんの KV を置ける一枚板を最初から予約する。実際の入力が三百トークンでも、今後どこまで伸びるか分からない。生成が長くなるかもしれない。だから安全側に倒して、最大長まで取っておく。

この方式は、考え方としては分かりやすい。論理的なトークンの並びと、メモリ上の並びが一致する。前から何番目のトークンかが分かれば、メモリ上の位置も簡単に計算できる。実装も単純である。

だが、無駄が大きい。

第一の無駄は、予約したのに使われない場所である。例えば、最大四千トークンぶんを予約したリクエストが、実際には六百トークンで終わるとする。この場合、三千四百トークンぶんの場所は、予約されているが使われない。他のリクエストに貸すこともできない。リクエストの内側に空きが閉じ込められている。これを内部の無駄と呼べる。

第二の無駄は、断片化である。断片化とは、空き容量そのものは残っているのに、それが細かく散らばっていて、必要な形で使えない状態である。

例えば、GPU のメモリの中に、千トークンぶんの空きが十か所あるとする。合計すれば一万トークンぶん空いている。だが、新しいリクエストが四千トークンぶんの連続した一枚板を必要とするなら、どこにも入らない。空きはあるのに、連続していないから使えない。

この状態では、同時にさばけるリクエスト数が早く頭打ちになる。メモリ容量の数字だけを見るとまだ余裕がある。しかし、エンジンから見ると、貸せる形の場所がない。推論基盤では、この差がそのまま処理能力と費用に出る。

KV をページに切って配る

ここで効くのが、古いが強力な発想である。メモリを一枚板として配るのをやめる。固定サイズの小さな区画に切り、それを必要な分だけ配る。

この小さな区画をページと呼ぶ。ページとは、メモリを管理しやすくするための固定サイズの単位である。ここでは、KV キャッシュをページのように扱う。例えば、一ブロック十六トークンぶんの KV を入れられる固定サイズの区画を作る。この区画をブロックと呼ぶ。ブロックとは、KV を一定数のトークンごとに区切った配布単位である。

リクエストは、最初から最大長ぶんを受け取らない。必要になったぶんだけ、空いているブロックを少しずつ受け取る。入力が三百トークンなら、十六トークン単位で十九ブロックほどを使う。生成が進んで三百十七トークンになれば、必要に応じて一ブロック足す。四千トークンまで伸びるかどうかは、実際に伸びてから考えればよい。

重要なのは、これらのブロックが物理的に隣り合っている必要はない、という点である。論理的には、トークンは前から順に並んでいる。だが、メモリ上では、最初の十六トークンぶんのブロックが奥の方にあり、次の十六トークンぶんのブロックが別の場所にあり、その次がまた違う場所にあってもよい。

では、どうやって論理的な順番と物理的な置き場を対応させるのか。そこで使うのがブロック表である。

ブロック表とは、リクエストごとに持つ対応表である。論理的な何番目のブロックが、物理的などのブロックに置かれているかを記録する。例えば、あるリクエストの論理ブロック 0 は物理ブロック 105、論理ブロック 1 は物理ブロック 37、論理ブロック 2 は物理ブロック 912、という具合である。

注意機構が過去トークンの KV を読むとき、エンジンはまず、そのトークンが論理的に何番目のブロックに入っているかを求める。ブロックサイズが十六トークンなら、トークン番号を十六で割れば、どの論理ブロックか分かる。余りを見れば、そのブロック内の何番目かも分かる。次に、ブロック表を見て、対応する物理ブロックを探す。そして、その物理ブロックの中の位置から K と V を読む。

つまり、一枚板の単純さを捨てる代わりに、間接参照を一段入れる。直接「このトークンはこのメモリ位置」と決めるのではなく、「このトークンは論理ブロックのここにあり、その論理ブロックは物理ブロックのここにある」とたどる。

この一段の間接が、メモリ管理を大きく変える。

無駄は最後の一ブロックに閉じ込められる

ページ式にすると、まず最大長を予約しなくてよくなる。

一枚板方式では、最初に四千トークンぶんを取る必要があった。ページ式では、いま使っている長さに合わせてブロックを足すだけでよい。三百トークンなら三百トークンぶんに近い量だけを使う。千トークンまで伸びたら、その時点で千トークンぶんに近い量を使う。

もちろん、無駄が完全にゼロになるわけではない。ブロックが固定サイズだからである。例えば、一ブロック十六トークンで、実際の文脈が三百一トークンだとする。十六トークンずつ詰めると、十八ブロックで二百八十八トークン、十九ブロックで三百四トークンまで入る。最後のブロックには三トークンぶんの空きが残る。

だが、この無駄は小さい。最大でも、最後の一ブロック未満である。一ブロック十六トークンなら、最大でも十五トークンぶんの空きに収まる。四千トークンぶんを予約して、実際には六百トークンしか使わない、という無駄とは桁が違う。

次に、断片化が効きにくくなる。

一枚板方式では、新しいリクエストのために連続した大きな場所を探す必要があった。ページ式では、空いているブロックを必要数だけ集めればよい。それらはばらばらでよい。メモリのあちこちに一ブロックずつ空きが散っていても、それをそのまま貸せる。

リクエストが終われば、そのリクエストが使っていたブロックを空きリストに戻す。空きリストとは、現在使われていないブロックの一覧である。次のリクエストは、そこから必要な数だけ受け取る。連続しているかどうかを気にしないので、空きが使いにくい形で残る問題が大幅に減る。

この効果は、単なるメモリ節約ではない。同じ GPU メモリで、より多くのリクエストを同時に抱えられるようになる。すると、バッチを大きくしやすくなる。バッチが大きいほど、重みの読み出しを複数リクエストで割り勘しやすい。一トークンを出すための費用が下がる。

つまり、KV をページとして配ることは、低レベルの実装上の工夫に見えて、推論の経済性そのものに効いている。

読むときの面倒を、どう引き受けるか

もちろん、ページ式には代償もある。KV が一枚板ではなくなるので、読む側が少し複雑になる。

注意機構は、過去の多くのトークンの K と V を読む。連続したメモリに並んでいれば、読み出しは単純である。ページ式では、ブロック表を見ながら、物理的には散らばったブロックをたどる必要がある。これは余分な処理である。

だから、ブロックサイズは小さすぎてもいけない。例えば、一ブロック一トークンにすれば、無駄はほぼなくなる。しかし、トークンごとに対応表を引くことになり、管理の負担が大きくなる。逆に、一ブロック千二十四トークンにすれば、対応表は小さくなるが、最後のブロック内の無駄が大きくなる。短いリクエストでは、ほとんど使わない大きなブロックを抱えることになる。

そこで、例えば十六トークン、三十二トークン、といった単位が考えられる。小さすぎず、大きすぎない単位を選ぶ。ブロック表の管理費用と、最後のブロックの無駄との釣り合いを見る。

エンジンは、生成中にこの割り当てを繰り返す。デコードで新しいトークンが生まれるたび、そのトークンの K と V を現在の最後のブロックに書く。もし最後のブロックが埋まっていれば、新しい空きブロックを取り、ブロック表に一行足す。こうして、文脈は論理的にはまっすぐ伸びる。物理的には、空いている場所を拾いながら伸びる。

この仕組みの本質は、論理的な連続性と物理的な連続性を切り離すことである。利用者から見える文脈は、前から後ろへ一続きである。モデルの計算から見ても、過去トークンは順番を保っている。だが、メモリの中では、必ずしも一続きでなくてよい。この切り離しが、容量を使い切るための鍵になる。

同じブロックを、複数で共有する

ブロック表という間接の層があると、もう一つ得をする。同じ内容の過去を、複数リクエストで共有できる。

例えば、多くのリクエストが同じ最初の指示文を持つとする。役割の説明、出力形式、禁止事項、共通の手順などである。この前置きが五百トークンあるなら、本来は各リクエストが五百トークンぶんの KV を持つことになる。百件あれば、同じ内容の KV を百個持つ。

だが、前置きが本当に同じなら、その部分の K と V も同じである。そこで、一度作ったブロックを複数のブロック表から指す。各リクエストのブロック表は、自分の論理ブロックとして同じ物理ブロックを参照する。中身をコピーしない。

このとき必要になるのが、参照カウントである。参照カウントとは、あるブロックを何件のリクエストが使っているかを数える仕組みである。ある共有ブロックを三件が指していれば、参照カウントは三である。一件が終われば二になる。最後の一件が離れてゼロになったら、そのブロックを空きリストに戻す。

ただし、共有できるのは同じ内容の部分である。途中から別の生成が始まれば、その先はリクエストごとに違う。共通の前置きブロックは共有し、その後のブロックは各リクエストが自分で持つ。この境目を、ブロック表で自然に表せる。

ここでも、ページ式の意味が出る。一枚板で持っていると、先頭の一部だけを共有し、途中から分ける管理が難しい。ブロック単位なら、前の数ブロックだけ同じ物理ブロックを指し、次のブロックから別々の物理ブロックを指せばよい。

同じ前置きを何度も使う処理では、これはそのままメモリの節約になる。同じ KV を何度も作らないので、計算の節約にもなる。エージェントのように、同じ役割説明や手順を繰り返し使う用途では、この差が大きい。

エンジンはモデルの付属品ではない

ここまで見ると、推論エンジンがモデルの単なる付属品ではないことが分かる。

同じ重みを使っていても、KV を一枚板で持つのか、ページとして配るのかで、同時に抱えられるリクエスト数は変わる。断片化で早く詰まるか、細かい空きを最後まで使えるかが変わる。共通の前置きを毎回コピーするか、同じブロックを共有するかも変わる。

モデルの賢さは、主に重みが決める。どのような知識や能力を持つかは、学習で作られた数値に強く依存する。だが、その賢さをどれだけ安く、どれだけ多くの人に、どれだけ安定して届けられるかは、エンジンの実装が決める。

特に KV キャッシュの管理は、派手ではない。画面に出る機能名でもない。だが、長文脈を扱い、同時リクエストを増やし、一トークンあたりの費用を下げるには、ここを避けて通れない。

推論の現場では、モデルは重みのファイルであり、エンジンはそれを動かす実行系である。そして、実行系の中心には、計算そのものだけでなく、メモリをどう貸し、どう返し、どう共有するかという地味な管理がある。

Agentic OS への含意

Agentic OS の上の層では、エージェントをどう設計するか、どの道具を持たせるか、どの記憶を参照させるか、どの利用者にどう応答するかを考える。だが、その下では、毎回の呼び出しが文脈を作り、KV キャッシュを作り、GPU メモリを消費している。

エージェントは、長い文脈を持ちやすい。役割の指示、手順、過去のやり取り、参照した情報、途中の思考や計画が積み上がる。また、何度も呼ばれる。同じ前置きや同じ作業手順を繰り返し使う。だから、ページ式の KV 管理と、前置きの共有は、エージェント基盤にそのまま効く。

第1回から第3回で見た推論コストは、ここで具体的な姿を持つ。長文脈が高いのは、抽象的に「計算が多い」からだけではない。過去トークンの K と V を、リクエストごとに持ち続ける必要があるからである。その置き方が悪ければ、メモリは早く詰まる。置き方がよければ、同じメモリでより多くの文脈を抱えられる。

上の層が「このエージェントを何人に、いくらで提供できるか」と考えるとき、答えは料金表や設計思想だけでは決まらない。下の層で、KV を一枚板として予約しているのか、ページとして配っているのかまで降りていく。

モデルとは、重みのファイルである。それだけでは動かない。動かすのはエンジンである。そして、エンジンの難所の一つは、過去文脈の置き場をどう管理するかである。鍵は、KV をページとして配ることだ。

次回も、Agentic OS の土台をもう少し掘る。