Chunked Prefillとは？長いプロンプトでLLMの生成を止めない仕組み

LLMが回答をstreamingしている途中で、tokenの表示が数秒止まることがある。原因の一つが、別の利用者から届いた長いpromptのPrefillである。prompt全体を一度に処理すると、その間、すでに生成中のrequestが次tokenを計算できない。

Chunked Prefillは、長いpromptを複数の小さなchunkへ分け、Decodeのiterationと混ぜて少しずつ処理する。長文入力がGPUを長時間占有するのを防ぎ、token配信の停止を短くするscheduler技術である。

Prefillを一度に実行すると何が起きるか

LLM推論は、prompt全体を読むPrefillと、回答を1 tokenずつ生成するDecodeに分かれる。Prefillは入力tokenを並列に計算できる一方、promptが長いほど一回の処理時間も大きくなる。

すでにA、B、Cの回答を生成しているGPUへ、数万tokenのrequest Dが届いたとする。schedulerがDのPrefillを優先して一度に実行すると、その間A、B、CのDecodeは待つ。平均throughputが高くても、利用者にはtoken配信が途切れるgeneration stallとして現れる。

OSDI 2024で発表されたSarathi-Serveは、Prefill優先schedulerで数秒に及ぶgeneration stallが起こりうると指摘し、Chunked Prefillとstall-free schedulingを提案した。

長いpromptをtoken chunkへ分ける

Chunked Prefillでは、Dのprompt全体を一度に処理しない。たとえば8,000 tokenのpromptを2,000 tokenずつ4回へ分け、Decodeと交互にscheduleする。

iteration	処理する内容
1	A・B・CのDecode + DのPrefill chunk 1
2	A・B・CのDecode + DのPrefill chunk 2
3	A・B・CのDecode + DのPrefill chunk 3
4	A・B・CのDecode + DのPrefill chunk 4
5以降	A・B・C・DのDecode

Dが最初のtokenを返すまでの総Prefill計算が消えるわけではない。複数iterationへ分けることで、進行中requestのDecodeを間に入れられるようになる。

Continuous Batchingとの違い

Continuous Batchingは、完了したrequestをbatchから外し、空いた枠へ新しいrequestを入れる仕組みである。request単位の入れ替えを扱う。

Chunked Prefillは、一つの長いPrefillをtoken単位の複数chunkへ分ける。まだPrefillが終わっていないrequestを、複数iterationにまたがって進める。

二つは競合する機能ではない。Continuous Batchingで生成中requestを入れ替えながら、空いたtoken budgetへPrefill chunkを詰めることで、GPU utilizationを高めつつDecodeのstallを抑えられる。

Prefill・Decode分離とも目的は似ているが、境界が違う。Chunked Prefillは同じworker内のschedulerで干渉を小さくする。Prefill・Decode分離は別worker poolへ移し、KV Cacheを転送する。後者は独立scaleが可能だがnetworkと運用の費用が増える。

token budgetで一回の仕事量を制限する

Chunked Prefillでは、request数よりiterationごとのtoken budgetが重要である。Decode中のrequestは通常、各iterationで次の1 tokenを必要とする。残ったbudgetへPrefill tokenを入れる。

vLLMのChunked Prefillでは、max_num_batched_tokensが一つのbatchでscheduleするtoken数を制御する。vLLMの設計説明では、Decodeを先にbatchへ入れ、残りのtoken budgetへ待機中Prefillをscheduleし、収まらないPrefillをchunkへ分ける方針が示されている。

考え方を簡略化すると、Prefillへ使えるtoken数は次になる。

\text{Prefill budget} = \text{batch token budget} - \text{Decodeで使うtoken数}

実際にはschedulerの予約、padding、modelやbackendの制約もあるため、この式だけでmemory使用量を正確には求められない。何を調整する値かを理解するための概念式である。

chunkを小さくすればよいわけではない

chunkを小さくすれば、一回のPrefillがDecodeを止める時間は短くなり、Inter-Token Latencyを守りやすい。一方で、長いpromptを多くのiterationへ分けるため、そのrequestのTime to First Tokenが伸びる場合がある。schedule回数やkernel実行のoverheadも増える。

chunkを大きくすると、Prefillを少ないiterationで終えやすく、TTFTやPrefill throughputに有利である。ただし一回の処理が長くなり、進行中Decodeのstallが大きくなる。

chunk・token budget	良くなりやすい指標	悪化しやすい指標
小さい	ITL、短いDecodeの滑らかさ	長文requestのTTFT、scheduler overhead
大きい	Prefill throughput、長文requestのTTFT	進行中requestのITL、tail stall

最適値はGPU、model、tensor parallelism、prompt長、同時request数で変わる。別環境の設定値をそのまま流用しない。

Decode優先にも上限がある

Decodeを常に最優先すると、進行中requestのtokenは滑らかに流れる。しかし、trafficが多くDecode batchが埋まり続ければ、新しいPrefillが残りbudgetを得られず、待機requestのTTFTが伸び続ける。

反対にPrefillを積極的に入れれば新規requestは開始しやすいが、実行中requestのITLが悪化する。schedulerにはfairnessが必要である。

次のような制御を検討する。

Prefill requestの最大queue age
iterationごとのPrefill最小budget
対話とoffline batchのpool分離
長いpromptの同時受付数
tenantごとのtoken量と同時実行上限
deadlineが近いrequestの優先

Chunked Prefillはstallを細かくする仕組みであり、無制限な長文を公平に処理する保証ではない。

Prefix Cacheが効く部分は再計算しない

system promptや共通資料など、prompt前半のKV Cacheを再利用できれば、新しくPrefillするtoken数を減らせる。Chunked Prefillで分割すべきなのは、cache hit後に残る未計算部分である。

ただし、cache hit率だけを見てtoken budgetを決めると危険である。cacheがcoldになるdeployment直後、version変更、tenant分離によるmissでPrefill量が急増する。cache hit時と完全miss時の両方を負荷試験する。

Multimodal入力では境界が複雑になる

画像や動画を含むrequestでは、text tokenだけでなくvision encoderの出力やplaceholder tokenを扱う。任意の位置で単純にchunkを切ると、画像に対応するfeature tokenのまとまりを壊す可能性がある。

vLLMのmultimodal processing設計は、Chunked PrefillやPrefix Cachingのために、placeholder feature tokenと元のmultimodal inputの対応を管理している。text-only modelで動いた設定を、そのままmultimodal modelへ適用せず、対応するarchitectureとruntime versionで検証する。

監視するmetric

tokens/sだけではChunked Prefillの効果を判断できない。throughputが上がっても、長いrequestのTTFTや進行中requestのITLが悪化している場合がある。

少なくとも次を測る。

input token数とoutput token数の分布
Prefill queue lengthと最古requestの待ち時間
requestごとのPrefill chunk数
iterationごとのPrefill token数とDecode sequence数
TTFT、ITL、E2E latencyのp50・p95・p99
generation stallの最大時間
input tokens/s、output tokens/s、total tokens/s
KV Cache使用率、preemption、OOM
deadline超過とcancel後に処理したtoken数

平均ITLだけでなく、単一request内で最も長くtokenが止まった時間を見るとgeneration stallを発見しやすい。

負荷試験は短文と長文を混ぜる

すべてのpromptを同じ長さに揃えると、Chunked Prefillが解決する干渉を再現できない。次のtraceで、無効時と有効時、複数のtoken budgetを比較する。

短いprompt・短いoutputだけを送る
長いpromptを単発で送る
短い回答のstreaming中に長いpromptを投入する
長いpromptをburstで投入する
長いprompt・長いoutputを混ぜる
Prefix Cache hitと完全missを切り替える
clientを途中切断し、残りchunkが停止するか確認する
context上限付近のmultimodal requestを試す

同じmodel、GPU、量子化、runtime version、request traceで比較し、total throughputだけでなくTTFTとITLのSLO達成率を見る。

Chunked Prefillは同じGPU内の干渉を小さくする

Chunked Prefillは、長いpromptのPrefillを複数iterationへ分け、進行中Decodeと同じbatchへ少しずつ入れる。長文requestによるgeneration stallを抑え、GPUの空きtoken budgetを使いやすくする。

代償は、長文requestのPrefillが複数回に分かれ、設定によってTTFTが伸びることである。小さいchunkは常に正解ではない。prompt長とoutput長が混在する本番traceで、TTFT、ITL、tokens/s、KV Cache、fairnessを測り、用途に合うtoken budgetを決める必要がある。