Chunked Prefill

チャンク化プリフィルは、長い入力プロンプトを小さなトークンのグループに分割し、ほかのリクエストのトークン生成と交互に複数ステップで処理するモデルサービングの手法です。この機能がない場合、長いプロンプトを受け付けた推論エンジンは、そのプロンプトだけで1ステップを使い切ります。その間、ほかのユーザーに配信中の応答はすべて、処理が終わるまで一時停止します。2023年のSARATHI論文がチャンク化プリフィルを提案し、後続のSarathi-Serveはこれを基盤に「停止のない」スケジューリングを構築しました。DeepSpeed-FastGenのDynamic SplitFuseも同じ考え方を採用しています。現在、チャンク化プリフィルはvLLMでデフォルト有効になっており、オープンウェイトの大規模言語モデル（LLM）を多数サーブするエンジンであるSGLangとTensorRT-LLMでも利用できます。

プリフィル、デコード、生成の停止#

自己回帰型Transformerの推論には、次の2つのフェーズがあります。

• プリフィル： モデルはプロンプト内のすべてのトークンを並列処理してキーと値を計算し、KVキャッシュに保存します。プリフィルでは、読み込む重みの1バイトあたりに実行する演算が多いため、通常は計算律速になります。
• デコード： モデルはリクエストごとに1ステップで1トークンを生成し、毎回キャッシュ全体を読み込みます。読み込む1バイトあたりの演算が少ないため、通常はメモリ帯域幅に制約されます。

20,000トークンのプリフィルを含むGPUのステップは、デコードだけのステップよりはるかに長くかかるため、バッチ内のすべてのリクエストがその処理を待ちます。Sarathi-Serveはこれらの一時停止を生成停止と呼び、プロンプト自体は短いユーザーのトークン間レイテンシに急増として現れます。

チャンク化プリフィルの仕組み#

スケジューラーは各ステップにトークン予算を割り当てます。vLLMではmax_num_batched_tokensで設定します。まず実行中の各デコードに1トークンを割り当て、残りの予算を待機中のプロンプトのチャンクに使います。予算に収まらないプロンプトは次のステップに持ち越され、プロンプト全体の処理が終わってデコードを開始するまで、KVキャッシュをチャンクごとに拡張します。2つのフェーズを1つのステップで混在させると、GPUの利用効率も上がります。デコードによって計算資源がアイドルになる時間を、プリフィルのチャンクで埋められるためです。

チャンク化で変わるのは、プロンプトのトークンを処理するタイミングだけです。各チャンクは、因果マスクを適用しながら、先行するすべてのチャンクと自身のトークンのキャッシュ済みキーと値にアテンションを適用します。そのため、結果は1回のプリフィルと同じになります。PyTorchのアテンションレイヤーでその様子を確認できます。

import torch
import torch.nn.functional as F

torch.manual_seed(0)
prompt_len, chunk = 1000, 256
q, k, v = torch.randn(3, 1, 8, prompt_len, 64).unbind(0)  # one layer: 8 heads, 64 dims each

# One-shot prefill: all prompt tokens attend causally in a single pass
full = F.scaled_dot_product_attention(q, k, v, is_causal=True)

# Chunked prefill: each chunk attends to the cached keys and values of earlier chunks plus itself
outputs = []
for start in range(0, prompt_len, chunk):
    end = min(start + chunk, prompt_len)
    mask = torch.ones(end - start, end, dtype=torch.bool).tril(diagonal=start)  # causal over absolute positions
    q_chunk, k_cached, v_cached = q[..., start:end, :], k[..., :end, :], v[..., :end, :]
    outputs.append(F.scaled_dot_product_attention(q_chunk, k_cached, v_cached, attn_mask=mask))
chunked = torch.cat(outputs, dim=2)

print(f"{len(outputs)} chunks, same result as one pass: {torch.allclose(full, chunked, atol=1e-5)}")

出力は4 chunks, same result as one pass: Trueです。1,000トークンのプロンプトは256、256、256、232トークンのチャンクとして処理され、単一パスと同じ出力になります。

関連するサービング手法#

• 連続バッチ処理： バッチ全体の処理完了を待たずに、各ステップでリクエストを追加・削除します。チャンク化プリフィルにより、このループ内で長いプロンプトもスケジュールできます。
• PagedAttention： KVキャッシュを固定サイズのページに保存するため、複数ステップにわたってプリフィルするプロンプトも、全長分のメモリを事前に確保せずにキャッシュを拡張できます。
• プロンプトキャッシュ： 繰り返し使われるプロンプトプレフィックスのキーと値を再利用し、その部分のプリフィルを完全に省略します。
• プリフィルとデコードの分離： 2つのフェーズを別々のGPUプールで実行し、その間でKVキャッシュを転送することで、各プールを固有のボトルネックに合わせてスケーリング、調整できます。
• 投機的デコード： 複数のトークン候補を先に生成し、1回のパスで検証することでデコードフェーズを高速化します。

用途#

チャンク化プリフィルは、長いプロンプトと短いプロンプトが同じサーバーを共有する環境で特に有効です。検索拡張生成（RAG）アシスタントが50ページのレポートでコンテキストウィンドウを埋める場合も、ほかのユーザーのチャットを配信し続けながら、そのレポートをチャンク単位でプリフィルできます。視覚言語モデル（VLM）は各プロンプトに数百から数千の画像トークンを追加するため、フレームの説明を送信するカメラパイプラインでは、まさにこうした混在が発生します。Ultralytics LLMインターフェースは、セルフホスト型vLLMサーバーを含むOpenAI互換エンドポイントにテキストと画像を送信します。そこで、人物を検出したフレームの説明を要求するYOLO26パイプラインのリクエストは、ほかのユーザーのトラフィックと並行してチャンク単位でプリフィルされます。

トークン予算の調整#

この予算によって、2つのレイテンシ指標の間でトレードオフが生じます。予算が小さいと各ステップでプリフィルに使える余地が少なくなり、ストリーミング中のユーザーのトークン間レイテンシを低く保つ一方、新しいプロンプトの初回トークンまでの時間は長くなります。予算が大きい場合はその逆です。長いプロンプトでは、後続の各チャンクが先行チャンクの増加したKVキャッシュを再度読み込むため、チャンクが小さすぎるとメモリトラフィックが増加します。運用担当者は、実際のプロンプト長に基づいて、中央値とテールの推論レイテンシ、初回トークンまでの時間、スループットを考慮しながら予算を調整します。vLLM最適化ガイドでは、トークン間レイテンシを優先する予算として2,048トークンを挙げ、大規模GPU上で小型モデルのスループットを最大化するには8,192トークンを超える予算を推奨しています。