静的バッチングが捨てているもの

一言でいうと

静的バッチ処理は、バッチの中で最も長いリクエストが終わるまで、残りのスロットを遊ばせます。連続バッチ処理は、終わったスロットを、すぐに新しいリクエストで埋めます。その差が、スループット2–5倍です。

なぜ必要なのか

画像分類のような作業では、バッチ処理は単純です。入力のサイズが同じで、処理時間も同じなので、32個を集めて一度に動かせば済みます。

LLMの生成は違います。各リクエストの出力長がまちまちです。あるリクエストは10トークン、あるリクエストは1,000トークンです。静的バッチ処理で32個をまとめると、31個が10トークンで終わっても、1,000トークンのものが終わるまで、その31個のスロットが空いたままになります。GPUは、32個分のリソースを確保して、1個を処理します。

どう動くのか

連続バッチ処理(continuous batching、またはin-flight batching)は、バッチを、固定されたまとまりではなく、流動的なスロットの集合として扱います。デコードの反復ごとに、スケジューラーが確認します。終わったシーケンスがあれば、スロットから外し、待機キューにリクエストがあれば、そのスロットに入れます。

そうすると、GPUは、つねに最大に近いバッチを保ちます。ベンチマークで、静的バッチ処理に比べてスループットが2–5倍向上すると報告されています。

ここで調整できる値は、2つです。max_num_seqsは、同時に処理するシーケンス数の上限で、max_num_batched_tokensは、1回の反復で処理するトークン数の上限です。前者はKVキャッシュのメモリに縛られ、後者は演算能力に縛られます。

トレードオフがあります。同時実行数を上げるとスループットは上がりますが、個々のリクエストのレイテンシが増えます。そして、KVキャッシュが足りなくなると、スケジューラーが一部のシーケンスをプリエンプション(preempt)して外に出し、あとでまた入れます。このとき、そのシーケンスのKVキャッシュを捨てて再計算するか、CPUにスワップしますが、どちらも高くつきます。プリエンプションが頻繁に起きると、スループットがかえって下がります。

もう1つの候補が、プレフィックスキャッシングです。複数のリクエストが同じシステムプロンプトで始まるなら、その部分のKVキャッシュを再利用できます。得になる程度は、共通プレフィックスの長さ、再利用率、実装とメモリのポリシーによって変わり、プロンプトが毎回まったく違えば、効果はほとんどありません。実際のリクエストの分布で、キャッシュのヒット率とTTFTをいっしょに測定する必要があります。

現場での姿

キューの深さと、テールのTTFTとの関係を見ると、運用の感覚が身に付きます。流入が処理能力に近づくと、待機キューが増え、高いパーセンタイルのレイテンシが、平均より先に悪くなります。製品のユーザー体験とトラフィックの規模に合わせて、p95やp99のようなテールのパーセンタイルを、SLOの指標として選びます。

そして、目標を決めるときには、順序があります。まず、製品とワークロードに合ったTTFTのSLOを決め(このラボでは、計算のために200msを例として使います)、それを満たす最大の同時実行数を探し、その同時実行数でのスループットを、キャパシティプランニングの根拠として使います。スループットだけを先に最大化すると、テールレイテンシのSLOを外しやすくなります。

KVキャッシュが本当の限界線

連続バッチ処理の同時実行数を何が妨げるのかを尋ねると、たいてい演算能力だと答えますが、実際に先に尽きるのは、ほとんどいつもメモリです。生成中のシーケンスは、ここまでのすべてのトークンのキーと値を持っている必要があり、そのサイズは、トークン数に比例して増え続けます。リクエスト1つが4,000トークンまで進むと、そのリクエストのキャッシュも、最初の数倍になります。

そのため、バッチサイズを決める計算は、次のように進みます。モデルの重みが占めた残りが、キャッシュに使える取り分で、それを、リクエスト1つが平均的に使うキャッシュのサイズで割った値が、同時に収められるシーケンス数です。ここで、平均ではなく、長いほうを見る必要があります。ほとんどのリクエストが短くても、長いリクエストが数個でキャッシュを占有すれば、その分だけスロットが減ります。

キャッシュが足りなくなると、スケジューラーがシーケンスをプリエンプションして外に出しますが、この瞬間が、運用で最も危険な地点です。外に出されたシーケンスは、あとでまた入ってきて、最初から計算し直されるので、すでにやった作業をまたやることになります。負荷が少し増えただけなのに、スループットが急に折れる崖が、ここで生じます。見ておく指標は、3つです。

指標	何を表すか	悪い兆候
KVキャッシュの使用率	メモリの余裕	90%付近に長くとどまります
プリエンプションの回数	巻き戻した作業の量	0でなければ、すでに限界を超えています
待機キューの長さ	受け取ったのにできていない量	増え続けるなら、流入を止める必要があります

対応は、3つに分かれます。リクエストごとの最大出力長を制限して、キャッシュが無限に増えないようにし、max_num_seqsを、プリエンプションが起きない値まで下げ、それでも足りなければ、インスタンスを増やします。プリエンプションを許容して、同時実行数を高く置く選択は、ほとんどいつも損です。見た目には、より多く受け付けているように見えますが、実際には、同じ計算を2回しているからです。

もう1つ、指摘しておくことがあります。長いリクエストと短いリクエストを、同じインスタンスに混ぜると、長いリクエストがキャッシュを占有している間、短いリクエストのレイテンシもいっしょに悪くなります。対話型の応答と、バッチ的な性格の長い生成を、別のインスタンスに分けることは、リソースの無駄ではなく、テールレイテンシを守る、最も確実な方法です。

次のラボですること

静的バッチ処理と連続バッチ処理のスケジューラーを、それぞれシミュレーションで実装して、スループットを比較します。同時実行数の上限を適用し、キューの深さとp99の待ち時間の関係を観察し、prefillとdecodeのコストを分けてモデリングしたあと、ラボ用の例示のSLOであるTTFT 200msを満たす最大の同時実行数を探索します。