Inference Batching（推論バッチング）

LLM 推論バッチングは訓練側 (tech-39) と前提が違う。可変長プロンプト・未知の出力長・到着駆動・自己回帰デコードのメモリ帯域律速 — この 4 つが GEMM 直観を崩す。本稿は静的 → 動的 → 連続バッチングを Orca と PagedAttention を軸に、KV キャッシュ管理・prefill-decode 分離・投機的デコード・TTFT/TPOT/Goodput・5 大ピットフォールまで通しで扱う。Transformer は tech-35、KV キャッシュが乗る注意機構は tech-38 を参照。

なぜ推論バッチングは訓練バッチングより難しいか

訓練バッチングは 3 つの仮定の上に立つ: 系列長は max_seq_len に packing 済み (GEMM 形状静的)、データはエポック頭で全置換済み、forward と backward の対称性で計算量安定。推論はこの 3 つすべてを壊す。

可変長。プロンプト長はリクエスト依存 (10 〜 32K)、出力長は <eos> まで未知。最長 padding すると 10 トークン要求と 4K 要求を同 batch に乗せた瞬間 99.75 % の compute が padding に食われる。

デコード相のメモリ帯域律速。Prefill (プロンプト処理) は入力トークン並列で compute-bound、decode (1 トークン/step) は各ステップで全モデル重みを読み直すため memory-bandwidth-bound。A100 80GB は FP16 で peak 312 TFLOP/s・帯域 2 TB/s、7B (FP16 約 14 GB) は重みの一往復で ~7 ms — これが batch=1 の per-token latency 下限。バッチングはこの重み読みを複数リクエスト間で償却し compute-bound 側に倒す手段。

到着駆動。推論リクエストはポアソン到着 (heavy tail) で batch が満ちる時刻が事前不明。満杯待ちは TTFT を伸ばし、即発火は batch 効率を落とす。

Head-of-line blocking。静的バッチングは N 件集め同期実行し全件終了で返すため、短い要求 (出力 10) は長い要求 (出力 500) の終端までスロット解放できず、最長と同じ待ち時間を強いられる。長尾の間、終了済みスロットは無駄。

3 つのバッチング方式 — 静的・動的・連続

静的バッチング (static): N 件 (または timeout) を集め最長 padding で同期実行、全件終了で返す。許容は SLO 無し (オフライン分類)・長さ均一・極小モデル (head-of-line 50 ms 未満) のいずれか。本番 LLM では padding 浪費と head-of-line blocking で破綻。

動的バッチング (dynamic / server-side): キューに積み (a) N 件溜まる or (b) 待ち時間 T 超過で発火しキュー内容を batch 実行。NVIDIA Triton の preferred_batch_size / max_queue_delay_microseconds、TorchServe の batching handler が代表。静的より良いが batch 境界で同期 が残り、最長要求が終わるまで次 batch を開始できず batch 単位の head-of-line blocking は消えない。

連続バッチング (continuous / iteration-level scheduling / in-flight batching): Orca (Yu et al., OSDI 2022)。batch を原子的に扱う発想を捨て、デコード 1 ステップごとに再スケジューリング。系列が <eos> か max_length で終了したら次ステップで即座にスロット解放、待機キューから新規を差し込む。batch は drain されず、トークン粒度で流入流出しながら走り続ける。Orca が導入した用語が iteration-level scheduling。論文は現実的分布で静的に対し大幅なスループット改善を報告 (具体値は原典)。HuggingFace TGI が早期本番採用、vLLM (2023) が PagedAttention と組合せ、NVIDIA TensorRT-LLM は同概念を in-flight batching と呼ぶ。

KV キャッシュメモリ管理 — PagedAttention と RadixAttention

自己回帰デコードは新トークンが過去全トークンに attend するため、過去のキーと値を毎ステップ再計算せず GPU メモリに保持する — これが KV キャッシュ。概算式:

KV_bytes = 2 × layers × heads × head_dim × seq_len × batch × dtype_bytes

7B (32 層・32 ヘッド・head_dim=128・FP16) で batch=16・seq=2048 で ~1 GB、batch=64・seq=4096 で ~16 GB。A100 80GB に 7B 重み 14 GB を載せた残りで活性化を捌くと、本番の実効 batch の天井は KV 容量で決まる。

PagedAttention (vLLM) の対象問題は二重: ① 内部断片化 — 100 トークン生成しか使わないリクエストが max_seq_len 分 (4K 等) を pre-allocate し ~97.5 % 無駄、② 共有不可 — system prompt や few-shot の共通プレフィックスを各リクエストが独立保持。仕組みは OS ページングの類推: KV を固定サイズ「ブロック」(典型 16 トークン/block) に分け、各リクエストが論理 KV 位置から物理ブロックへの「ページテーブル」を持つ。物理ブロックは生成進行に応じて on-demand 割当、終了時に解放し他リクエストに再配分。原典 Kwon et al., SOSP 2023。論文はメモリ浪費が大幅削減され、同 GPU メモリで実効 batch が数倍に増えると報告。vLLM (https://github.com/vllm-project/vllm) が OSS 実装。

RadixAttention (SGLang) は PagedAttention がリクエストスコープでクロスリクエスト共有不能な点を解く。KV ブロックをトークン列でインデックスする radix tree (trie) を保持、新規到着時に最長一致プレフィックスを検索し KV ブロックを再利用 — 新規トークン分だけ KV 計算。eviction は LRU。Zheng et al., 2024。RAG の固定コンテキスト・エージェントの system prompt・few-shot のように共通プレフィックスが長い workload で高ヒット率。GitHub: https://github.com/sgl-project/sglang。

実装システム比較

選択指針: 汎用なら vLLM か TensorRT-LLM、共通プレフィックスが多ければ SGLang、NVIDIA エコシステム内なら Triton + TRT-LLM backend、prefill 支配なら DeepSpeed-FastGen 系の chunked prefill。

Prefill-Decode 分離 — chunked prefill と disaggregation

Prefill と decode は compute プロファイルが正反対 (前節)。両者を同イテレーションに混ぜると prefill が decode を starve する: 長 prefill が入ると 1 イテレーションの GPU 予算の大半を食い、共存 active 系列の decode が遅延、進行中の ITL スパイクとなる。

Chunked prefill (Sarathi-Serve / Dynamic SplitFuse): prefill を固定サイズチャンク (例 512 トークン) に分け、各イテレーションで「1 prefill chunk + 全 active decode」を処理。効果は ① latency スパイク抑制、② compute-bound chunk と memory-bound decode の混在で GPU 利用率向上、③ 長プロンプトが batch をブロックせず TTFT 短縮。DeepSpeed-FastGen の Dynamic SplitFuse は同手法でチャンク予算配分をキュー状態に応じ動的決定。

Prefill-Decode Disaggregation (Splitwise / DistServe): prefill 用 GPU 群と decode 用 GPU 群を 物理分離、KV キャッシュをネットワーク (NVLink / Ethernet) で転送。Prefill-decode 干渉は完全に消えるが GPU 間 KV 転送オーバーヘッドが乗る。TTFT SLO と TPOT SLO の両方が同時に厳しい高 QPS で各相を独立に SLO 最適化したい場面向き。

投機的デコードとの相互作用

小さな「ドラフトモデル」が K 個の候補を並列生成、大きな「ターゲット」が 1 forward で検証 (一致 prefix を accept、最初の不一致で reject)。1 検証パスの実効進捗は 0 〜 K+1 と可変。

問題は 系列ごとに accept 数が違うこと — A は 3 進み B は 1 のように batch が不揃いに前進し、KV キャッシュも個別更新が必要で、「全系列が毎ステップ 1 トークン進む」という連続バッチングの素直な前提が壊れる。vLLM 等は検証後に batch を整列し直す chunked operation として実装する。Medusa (2024) はベースモデルに複数予測ヘッドを追加し tree attention で 1 パス検証、EAGLE (2024) は軽量ドラフトがトークン埋め込み + ターゲットの hidden state を条件にしてドラフト品質を上げる。どちらも単一リクエストのスループットは上がるが、リクエストあたりメモリ圧が増え高 batch では accept 数のばらつきで負荷不均衡が出やすい。

性能指標 — TTFT / TPOT / Goodput

7B・A100 80GB・FP16 程度の概算 (illustrative、実装で変動): 低負荷の 512 トークンプロンプトで TTFT は概ね 50〜200 ms、高負荷で prefill キューが詰まれば秒オーダーまで伸び得る。Batch=1 の TPOT は ~20〜30 ms/token、batch=32 程度で重み読みが償却され ~30〜60 ms/token、compute-bound 領域に近づく。Throughput が peak の 80 % 程度なら goodput はほぼ追従するが、95 % に押し込むと tail が TPOT SLO を割り goodput が崩壊する — これが本番の天井。

5 大ピットフォール

1. Prefill と decode を 1 相として扱う。「tokens/sec」だけ見る設定は prefill に倒れて decode を starve させる。TTFT と TPOT を別々に測り、prefill chunk size と batch 構成を独立チューニング。

2. Max batch size を上げすぎて KV preempt / swap を誘発。max_num_seqs を KV 予算と合わせず増やすと vLLM 系は KV 枯渇時に active 系列を preempt (停止+再計算) か swap (CPU オフロード) し、どちらもペナルティが大きい。定常 KV 使用率 80〜85 % で頭打ちになるよう調整、vLLM の gpu_cache_usage_perc を監視。

3. 高 QPS で padding-heavy な静的バッチング。実トラフィックは heavy tail で長さ 20 倍分散・1000 QPS で padded static は GPU compute の 40〜60 % を padding に食う。本番は連続バッチング + 可変長カーネル (FlashAttention varlen) に切り替える。

4. TTFT 測定でトークナイザ overhead を無視。テキスト → トークン ID 変換は CPU bound、実装 (Rust 製 HuggingFace tokenizers vs Python) で 1〜20 ms の幅。生テキスト入力から end-to-end で測る。

5. 均一長プロンプトでベンチ。全 512 トークン入力のような合成負荷は head-of-line blocking を隠す。ShareGPT 等の実分布か Pareto / log-normal 等の heavy-tail 合成を使い、SLO は median ではなく p99 latency を実分布の下で見る。

参考文献

• Yu et al. (2022). Orca: A Distributed Serving System for Transformer-Based Generative Models. OSDI 2022. https://www.usenix.org/conference/osdi22/presentation/yu
• Kwon et al. (2023). Efficient Memory Management for LLM Serving with PagedAttention. SOSP 2023. https://arxiv.org/abs/2309.06180
• Agrawal et al. Sarathi-Serve (chunked prefill). https://arxiv.org/abs/2403.02310
• Splitwise (Microsoft Research). https://arxiv.org/abs/2311.18677
• DistServe. https://arxiv.org/abs/2401.09670
• Zheng et al. SGLang / RadixAttention. https://arxiv.org/abs/2312.07104
• DeepSpeed-FastGen / Dynamic SplitFuse. https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-fastgen
• vLLM. https://github.com/vllm-project/vllm
• HuggingFace TGI. https://github.com/huggingface/text-generation-inference
• NVIDIA TensorRT-LLM. https://github.com/NVIDIA/TensorRT-LLM