prefillとdecodeは別の仕事だ

一言でいうと

LLM推論は、2つのフェーズに分かれます。プロンプトを一度に処理するprefillは演算がボトルネックで、トークンを1つずつ作るdecodeは、メモリ帯域幅がボトルネックです。

なぜ必要なのか

学習では、1つの指標だけを見れば十分です。1秒あたりに処理するトークン数です。サービングでは、それでは足りません。ユーザーが体感するものは2つあり、2つは異なる原因から生じます。

TTFT(Time To First Token)は、リクエストを送ってから、最初のトークンが届くまでの、エンドツーエンドの時間です。対話型のサービスでは、200msを例示のSLOとして置けますが、実際の目標は、製品の体験、モデルのサイズ、プロンプトの長さ、デプロイ環境を基準に決める必要があります。TTFTには、ネットワーク、リクエストのキュー、スケジューリング、トークン化、prefillがすべて含まれます。長いプロンプトでは、全体のアテンションを計算するprefillが主要な構成要素になりやすいですが、いつもTTFT全体を単独で決めるわけではありません。

ITLまたはTPOT(Inter-Token Latency、Time Per Output Token)は、その後トークンが1つずつ出てくる間隔です。人が読む速度より速ければ十分です。この値は、decodeの段階が決めます。

2つのフェーズでボトルネックが違うというのが核心です。prefillは、行列の積が大きく発生するので、演算能力に縛られます。decodeは、トークン1つごとにモデルの重み全体を1回ずつ読む必要があるので、メモリ帯域幅に縛られます。そのため、decodeは、バッチを大きくしても、時間があまり増えません。どのみち、重みを読む時間が支配的だからです。この性質が、バッチ処理が劇的に効果的な理由です。

どう動くのか

KVキャッシュが、新しく登場するリソースです。Transformerは、各トークンで、それ以前のすべてのトークンのキーと値を参照します。毎回計算し直すとO(n²)になるので、計算しておいたキーと値を、メモリに積んでおきます。これがKVキャッシュです。

問題はサイズです。公式は次のとおりです。

KV 바이트 = 2 x 레이어 수 x 은닉 차원 x 시퀀스 길이 x 배치 x dtype 바이트

このコードブロックの韓国語の式は、KVのバイト数が、2×レイヤー数×隠れ次元×シーケンス長×バッチ×dtypeのバイト数だ、という意味です。

7Bモデル(32層、隠れ次元4096)をfp16で、4Kコンテキスト、バッチ1で動かすと、2 x 32 x 4096 x 4096 x 1 x 2 = 2GiBです。コンテキストを128Kに延ばすと、64GiBです。モデルの重みが14GBなのに、KVキャッシュが64GBを食います。バッチ8なら512GiBで、GPU1枚には絶対に入りません。

そのため、サービングシステムの設計は、大部分がKVキャッシュの管理の話になります。PagedAttentionは、KVキャッシュを固定サイズのブロックに分けて、OSの仮想メモリのように非連続の割り当てを可能にし、内部断片化を減らします。GQAは、モデルのKVヘッド数に応じて、キー・値ヘッドをグループで共有して、キャッシュのサイズを減らします。KVキャッシュの量子化(INT8、INT4)は、dtypeのバイト数を減らしますが、対応しているかどうかと、精度への影響は、サーバーとモデルごとに確認する必要があります。

現場での姿

測定でよくあるミスが、TTFTとITLを平均でひとまとめにすることです。総レイテンシを総トークン数で割ると、2つの指標が混ざって、何も見えなくなります。必ず別々に測って、それぞれのp50、p95、p99を見る必要があります。

そして、GPUメモリの使用率の設定値(例: vLLMのgpu_memory_utilization)が重要です。一部の専用GPU環境では、0.90前後を出発点にしますが、モデルサーバーのバージョン・量子化・CUDAグラフと共存するプロセスによって、安全な値は変わります。高くしすぎるとCUDA OOMが出て、低くしすぎるとKVキャッシュに使う場所が減って、同時処理量が下がるので、実際のピーク負荷で検証する必要があります。

プロンプトの長さがすべてを変える

前の公式で、シーケンス長が掛け算の項として入っているという点が、実務で繰り返し戻ってきます。プロンプトを長くする決定は、精度を買う決定であると同時に、スループットを売る決定であり、その交換比を知らないと、キャパシティプランニングが毎回ずれます。

3つのことが、一緒に悪くなります。

TTFTが長くなります。prefillは、プロンプト全体に対するアテンションを計算するので、長さに応じてコストが急激に増えます。システムプロンプトが数千トークンあれば、ユーザーが何を尋ねても、そのコストを毎回支払います。

同時処理数が減ります。キャッシュが長さに比例するので、プロンプトが2倍なら、同じメモリに半分しか入りません。前のモジュールで見たプリエンプションが、ここから始まります。

コストが増えます。トークン単位で値段を付けるサービスでも、自前の運用でも、読んで計算する量が、そのまま料金や機材になります。

そのため、実務での対応は決まっています。共通のプレフィックスを固定します。システムプロンプトと例を、リクエストごとに少しずつ変えず、まったく同じに保てば、プレフィックスキャッシュが効きます。1文字だけ違っても、その後ろがすべて計算し直されるので、時刻やリクエストIDのようなものをプロンプトの前のほうに入れると、キャッシュがまるごと無効になります。変わるものは後ろへ回すことが、このキャッシュを生かす基本ルールです。

入れる文書を選びます。検索で持ってきた文書をすべて入れる代わりに、上位の数件だけを入れ、各文書も必要な部分だけを切り出します。文書を2倍入れても、答えが2倍よくなるわけではなく、むしろ関連のない内容が混ざると、精度が落ちる場合も多くあります。

出力の長さに上限を置きます。前に見たとおり、キャッシュは生成が進むほど増え続けます。上限がないと、リクエスト1つがメモリを食い続けて、他のリクエストの場所を奪います。本当に長い出力が必要な作業は、対話型とは別のインスタンスに送ることが、ここでも答えです。

次の確認で見ること

まずクイズで、TTFTとITLのボトルネック、KVキャッシュの容量の条件を区別します。次のモジュールから、トークンストリーミングサーバーを自分で作り、バッチスケジューラーをシミュレーションし、KVキャッシュの計算機を作成します。