連続バッチ処理

Continuous batching は、各生成ステップごとにバッチを動的に再スケジュールすることで、GPU の利用率を最大化します。リクエストが完了すると、バッチ全体が完了するのを待たずに、新しいリクエストが即座にバッチに参加します。これにより、GPU は常にフル稼働し、スループットも高い水準を維持されます。

本番環境でのデプロイには、transformers serveを使用してください。ContinuousBatchingManagerを基盤としており、OpenAI互換のHTTPエンドポイントを公開します。

generate_batch

generate_batch() を通じて連続バッチ処理がサポートされています。トークン化されたプロンプトのリストを渡すと、すべての処理が完了した時点で結果が返されます。generate_batch は内部でスケジューリングを処理し、すべてのリクエストが完了するまでブロックします。

サービングやストリーミングのユースケースでは、ContinuousBatchingManagerを直接使用してリクエストを管理します。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import ContinuousBatchingConfig, GenerationConfig

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-4B",
    attn_implementation="flash_attention_2",
    device_map="auto",
    dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")

prompts = [
    "What's up?",
    "Name a cat breed.",
    "Write a detailed history of quantum mechanics.",
]
inputs = [tokenizer.encode(p) for p in prompts]

generation_config = GenerationConfig(
    max_new_tokens=64,
    eos_token_id=tokenizer.eos_token_id,
)

outputs = model.generate_batch(inputs=inputs, generation_config=generation_config)

for request_id, output in outputs.items():
    text = tokenizer.decode(output.generated_tokens, skip_special_tokens=True)
    print(f"[{request_id}] {text}")

ContinuousBatchingManager

ContinuousBatchingManager はバックグラウンドスレッドを実行し、リクエストの送信と結果の取得を独立して行えます。各生成ステップで、完了したリクエストを確認し、バッチに参加する新しいリクエストをスケジュールします。ストリーミング、リアルタイムサービング、またはリクエストの到着に合わせて送信する場合に便利です。

continuous_batching_context_manager() を使用して、マネージャーを安全に開始および停止します。以下の例では、長さの異なる入力を使用しています。最も短いプロンプトが完了すると、そのプロンプトはバッチから離脱し、より長いプロンプトは生成を続行します。静的バッチングでは、すべてのプロンプトを同じ長さにパディングする必要がありました。連続バッチングにより、完了したプロンプトのリソースが解放されるため、次のプロンプトの処理をすぐに開始できます。

with model.continuous_batching_context_manager(generation_config=generation_config) as manager:
    manager.add_request(
        input_ids=tokenizer.encode("Write a detailed history of quantum mechanics."),
        request_id="long",
        max_new_tokens=512,
    )
    manager.add_request(
        input_ids=tokenizer.encode("What's up?"),
        request_id="short_0",
        max_new_tokens=32,
    )
    manager.add_request(
        input_ids=tokenizer.encode("Name a cat breed."),
        request_id="short_1",
        max_new_tokens=32,
    )

    for result in manager:
        text = tokenizer.decode(result.generated_tokens, skip_special_tokens=True)
        print(f"[{result.request_id}] {text}")

init_continuous_batching()を呼び出して、ライフサイクルを自分で管理することもできます。

manager = model.init_continuous_batching(generation_config=generation_config)
manager.start()

# submit and retrieve requests...

マネージャーのシャットダウン

マネージャーはバックグラウンドスレッドを実行し、分散リソースを保持します。シャットダウンは2段階で行われるため、進行中の処理に対してどのような処理を行うかを指定できます。

stop() を呼び出してバックグラウンドスレッドを停止します。デフォルトでは、マネージャーは新しいサブミットの受け入れを停止し、キューに登録されているリクエストとアクティブなリクエストの完了を待ってからスレッドを終了します。

manager.stop()

hard_stop=True を渡すと、保留中の処理を即座に中断します。キューに登録されているリクエストと実行中のリクエストは、完了する代わりに RuntimeError で失敗します。

manager.stop(hard_stop=True)

stop が呼び出されると、add_request() および add_requests() は新しい送信を拒否し、警告をログに記録します。同じマネージャーで別の生成セッションを実行するには、start を再度呼び出すことができます。

destroy() を呼び出して、分散リソースを解放します。destroy は、マネージャーがまだ実行中の場合はまずそれを停止します。その後、マネージャーを再起動することはできません。プロセスの生存期間中、continuous batching の使用が完了した際に使用してください。

manager.destroy()

continuous_batching_context_manager() がこの処理を担当します。終了時に stop を呼び出し、persistent_manager=True を渡してマネージャーをモデルにキャッシュして次のセッションで再利用しない限り、destroy も呼び出します。

リクエストの追加

add_request() は単一のリクエストを送信します。request_id を指定するか、マネージャーが自動的に生成するようにします。

manager.add_request(input_ids=input_ids, request_id="my_request")

add_requests() はバッチを一度に送信します。ブロック共有が有効な場合、プレフィックスキャッシュのヒット率を最大化するために、入力を自動的にソートします。

manager.add_requests(inputs=inputs)

cancel_request() を使用してリクエストをキャンセルします。

manager.cancel_request(request_id="my_request")

リクエストごとのサンプリングパラメータ

per_request_processors を有効にすると、同じ forward pass 内でリクエストごとに temperature、top_k、top_p を独立して適用できるようになり、リクエストごとに異なるサンプリングパラメータを使用できます（例: 創造的で高温度の出力と、正確で低温度の出力）。

cb_config = ContinuousBatchingConfig(per_request_processors=True)

# each request gets its own sampling parameters
manager.add_request(input_ids=inputs_a, temperature=0.9, top_p=0.95)
manager.add_request(input_ids=inputs_b, temperature=0.1, top_k=10)

GenerationConfig の各パラメータは、実行時に関連する logits プロセッサを作成するために、デフォルト以外の値である必要があります。例えば、リクエストごとの温度制御を有効にするには、temperature を None や 1 以外の値に設定します。その後でも、温度が 1 のリクエストを作成することは可能です。

結果の取得

マネージャーを反復して、結果が到着するたびに受け取ります。

for result in manager:
    print(tokenizer.decode(result.generated_tokens, skip_special_tokens=True))

get_result() は、出力キューから次の結果を取得します。request_id を渡すことで、特定のリクエストに絞り込むことができます。キューの次の結果が一致しない場合は、その結果は再キューイングされ、メソッドは None を返します。

# next available result
result = manager.get_result()

# filter for a specific request
result = manager.get_result(request_id="my_request")

ストリーミング

リクエストにstreaming=Trueを設定し、request_id_iter()を使用して、トークンが生成されるたびに部分出力を反復処理します。

from transformers.generation.continuous_batching import RequestStatus

manager.add_request(input_ids=input_ids, request_id="streamed", streaming=True)

for chunk in manager.request_id_iter(request_id="streamed"):
    token = tokenizer.decode(chunk.generated_tokens[-1:], skip_special_tokens=True)
    print(token, end="", flush=True)
    if chunk.status == RequestStatus.FINISHED:
        break

ContinuousBatchingConfig

ContinuousBatchingConfig は、KV キャッシュ、スケジューリング、CUDA グラフ、メモリ使用量などを制御します。generate_batch または init_continuous_batching() に、continuous_batching_config 引数として渡します。GenerationConfig とは別々に指定してください。これら 2 つの設定は、異なるものを記述しています。GenerationConfig はモデル中心で、サンプリングや停止に関するパラメータを保持します。一方、ContinuousBatchingConfig はハードウェア中心で、メモリやスケジューリングに関するパラメータを保持します。

GenerationConfig に continuous_batching_config を設定することは非推奨となり、FutureWarning が出力されます。v5.19 で削除予定です。

デフォルトでは、max_batch_tokensは8192で、利用可能なGPUメモリによって制限され、256を下回ることはありません。一方、num_blocksは残りのメモリを埋めるように設定されます。適切な機能を選択する際は、以下の表を参考にしてください。

from transformers.generation import ContinuousBatchingConfig

cb_config = ContinuousBatchingConfig(
    max_memory_percent=0.8,
    page_size=256,           # KV cache page size in tokens
    scheduler_type="fifo",        # "fifo" or "prefill_first"
)

outputs = model.generate_batch(
    inputs=inputs,
    generation_config=generation_config,
    continuous_batching_config=cb_config,
)

KVキャッシュのページサイズ

page_size は、1 つのレイヤーにおけるフルアテンションの 1 ページが保持する KV のトークン数を設定します。この値は 4 以上である必要があり、それ未満の場合には初期化時に ValueError が発生します。デフォルト値の 256 は、ほとんどのワークロードで問題なく動作します。

block_size パラメータは非推奨となり、page_size に置き換えられました。

ページサイズを大きくすると管理コストは削減されますが、シーケンスがページ途中で終了した場合にスペースが浪費されます。一方、ページサイズを小さくするとスペースの浪費は少なくなりますが、管理コストが増加します。通常、page_sizeを調整する必要はありません。ワークロードに短いリクエストが多い場合は、より小さい値を試してください。非常に長いリクエストが少数ある場合は、より大きい値を試してください。ページ、ブロック、チャンクがどのように関連し合うかは、メモリ管理を参照してください。

cb_config = ContinuousBatchingConfig(page_size=128)

プリフィルバッチサイズ

max_batch_tokens は、1回のフォワードパスにおけるトークン予算、すなわちモデルが一度に処理するクエリトークンの最大数を設定します。予算を大きくすると、各 prefill に詰められるプロンプトトークンが増え、プロンプト中心のワークロードにおいて prefill のスループットが向上し、time-to-first-token が短縮されます。スケジューラーは、予算を超えるプロンプトを複数のステップに分割します。サイズ超過のプロンプトの処理方法については、チャンク化 prefill を参照してください。

デフォルトでは、max_batch_tokens は 8192 です。バッチごとの入力テンソルが KV キャッシュを圧迫しないように、利用可能な GPU メモリによって値が制限され、256 未満になることはありません。空きメモリが少ない GPU では、この制限により 8192 未満に低下します。

max_batch_tokens と num_blocks は同じメモリ予算から割り当てられるため、互いにトレードオフの関係にあります。トークン予算を大きくすると、より大きな入力バッファが確保され、並行リクエストや長いリクエストに使える KV キャッシュブロックが少なくなります。メモリに余裕がある場合は max_batch_tokens を引き上げてプリフィルのスループットを高め、メモリを KV ブロック用に確保したり、メモリ不足エラーを回避したりする場合は引き下げてください。

cb_config = ContinuousBatchingConfig(max_batch_tokens=16384)

バッチおよびスケジューリングの制限

max_requests_per_batch は、1回の forward pass で実行されるリクエストの上限を指定します。モデルは、バッチ内のトークン数と語彙サイズによってサイズが決まる logits テンソルを生成し、サンプリングのために fp32 にキャストされるため、メモリフットプリントが2倍になります。各リクエストには次の1トークンの予測のみ必要であるため、上限を小さくすることでこのテンソルを小さく保ち、大規模な語彙を持つ prefill 中心のバッチにおけるメモリ不足エラーを回避できます。デフォルトでは、送信されたプロンプトの数に設定され、フォールバックとして 1024 が使用されます。その後、max_batch_tokens および num_blocks に収まるよう上限が調整されます。

cb_config = ContinuousBatchingConfig(max_requests_per_batch=256)

safety_margin は、アクティブなリクエストのために空きキャッシュ容量の一部を確保します。空き容量が safety_margin（キャッシュ総容量に対する割合）を下回ると、スケジューラは新しい prefill の受け入れを停止し、進行中のデコードリクエストのみを継続します。これにより、新しいリクエストを開始するよりも、アクティブな処理を完了させることを優先し、デコードレイテンシを保護するとともに オフローディング を遅延させます。値は 0 と 1 の間でなければならず、0 を指定するとこのマージンは無効になります。

cb_config = ContinuousBatchingConfig(safety_margin=0.15)

デフォルト値はスケジューラによって異なります。FIFOスケジューラの場合は0.15、prefill_firstの場合は0.0です。マージンがスケジューリングとどのように相互作用するかについては、Admissionを参照してください。

対数確率

return_logprobs=True を指定すると、ContinuousBatchingConfig は各生成トークンの対数確率を返します。これは、対数確率が一部の学習ループの入力となる強化学習（RL）で有用です。

cb_config = ContinuousBatchingConfig(return_logprobs=True)

# generate_batch()

for request_id, output in outputs.items():
    for token_id, log_prob in zip(output.generated_tokens, output.logprobs):
        token = tokenizer.decode([token_id])
        print(f"{token} | logprob: {log_prob}")

CUDA グラフ

CUDA graphs は、GPU 実行グラフを一度記録し、形状が一致するバッチに対してそれを再生することで、CPU のディスパッチオーバーヘッドを排除します。use_cuda_graph=True を指定して明示的に有効にしてください。

cb_config = ContinuousBatchingConfig(use_cuda_graph=True)

有効な場合、マネージャーはクエリとKVの長さを固定の区間にパディングし、形状を繰り返し、グラフの再利用を可能にします。q_padding_interval_sizeとkv_padding_interval_sizeの値を小さくすると、パディングによる計算の無駄を減らすことができますが、その代わり、グラフが記録・保存する必要がある一意の形状が増え、メモリコストが高くなります。

cb_config = ContinuousBatchingConfig(
    use_cuda_graph=True,
    q_padding_interval_size=64,
    kv_padding_interval_size=16384,
)

非同期バッチ処理

Async batchingは、次のバッチのCPUスケジューリングと、現在のバッチのGPU計算をオーバーラップさせます。CUDA graphsが必要であり、入力テンソルに使用されるVRAMは約2倍になります。

cb_config = ContinuousBatchingConfig(
    use_cuda_graph=True,
    use_async_batching=True,
)

コンパイル

default_compile_level は、モデルの forward pass に torch.compile を適用します。コンパイルにより、一度きりのウォームアップコストが発生する代わりに、生成時の forward pass が高速化されます。レベルを高くするとより積極的な最適化が実行され、スループットは向上しますが、ウォームアップ時間が長くなります。テストやベンチマークの反復ではレベルを低く保ち、ウォームアップコストを多数のリクエストにわたって回収できる長時間稼働のサービングワークロードでは、レベルを高く設定してください。

レベルは0から3までの範囲です。レベル0がデフォルトであり、コンパイルを完全にスキップします。

cb_config = ContinuousBatchingConfig(default_compile_level=1)

このレベルは、varlen と decode の実行パスにデフォルトの CompileConfig を提供します。明示的な設定が存在しないパスにのみ適用されるため、varlen_compile_config と decode_compile_config が設定されている場合は、それらが優先されます。FlashAttention の場合、varlen パスでは max_length_k により頻繁な再コンパイルが発生するためコンパイルがスキップされ、このレベルは decode パスにのみ影響します。

デコード高速パス

バッチにデコードリクエストのみが含まれている場合（シーケンスごとにクエリトークンが1つ）、マネージャーは可変長カーネルではなくflash_attn_with_kvcacheカーネルにディスパッチできます。このカーネルは、手動更新を経由するのではなく、ブロックテーブルを介してページ化されたKVキャッシュをインプレースで読み書きするため、varlenパスよりも高速です。カーネルレベルの詳細については、Paged attentionを参照してください。

ファストパスのサイズは、リクエストごとのブロックテーブルの次元を決定するmax_blocks_per_requestによって決まります。デフォルトでは、これは自動推論されます。マネージャーにmax_prompt_lengthとmax_generated_lengthが設定されている場合、ブロックテーブルは最大シーケンス長に収まるようにサイズ調整されます。それ以外の場合は、フォールバックのデフォルト値（リクエストあたり32ブロック）が使用されます。

max_blocks_per_request に特定の値を設定して、ブロックテーブルのサイズを明示的に指定できます。リクエストごとの最大シーケンス長が既知であり、ブロックテーブルのメモリコストを制限したい場合に便利です。

cb_config = ContinuousBatchingConfig(max_blocks_per_request=64)

max_blocks_per_request=0 を設定すると、fast path が無効化され、すべてのバッチが varlen kernel を通るようになります。これにより、デフォルト変更前の動作が復元され、attention 実装で fast path が利用できない場合に便利です（基盤となる kernel が使用できない場合、マネージャーが自動的に無効化します）。

cb_config = ContinuousBatchingConfig(max_blocks_per_request=0)

高速パスは、flash_attn_with_kvcache カーネルに依存しています。このカーネルは、デバイスとアテンション実装の2つの組み合わせで利用可能です。

それ以外の組み合わせ、またはカーネルをインポートできない場合、マネージャーは varlen パスにフォールバックします。max_blocks_per_request を明示的に設定した場合にのみ、警告がログに出力されます。

Full attention はカーネルブロックテーブルをサポートしていますが、sliding attention はサポートしていません。モデルに sliding attention レイヤーが含まれている場合（full attention も併用するハイブリッドモデルを含む）、初期化時に max_blocks_per_request は 0 に強制設定され、設定した非ゼロ値は上書きされます。その後、マネージャーは、モデルの一部の attention 型がカーネルブロックテーブルをサポートしていないため、デコードの高速パスが利用できないことをログに記録します。

CPUオフロード

CPU offloading は、GPU の KV キャッシュが満杯になった際に、退避された KV キャッシュブロックを事前割り当て済みのピン留め済み CPU バッファにコピーします。キャッシュの空き領域が確保されると、マネージャーはブロックを GPU にコピーし、プロンプトや生成済みトークンの再計算を行わずにリクエストを再開します。

cpu_offload_space に、GiB 単位の CPU スワップスペースを指定します。デフォルト値の 0.0 は、CPU オフロードを無効にします。

cb_config = ContinuousBatchingConfig(cpu_offload_space=8.0)

デフォルトでは、psutil がインストールされている場合、cpu_offload_space_safety_threshold=0.8 により、要求される容量は利用可能なシステム RAM の 80% に制限されます。スワッププールのサイズをセーフティスレッシュホールドから決定するには、cpu_offload_space=None を設定してください。

Tensor parallel タイムアウト

テンソル並列化では、マネージャーが CPU 通信グループを作成し、ランク間でリクエストの送信、キャンセル、およびシャットダウンを調整します。cpu_group_timeout は、このグループでの集約操作がプロセスのクラッシュに至る前にブロックされ得る最大時間を制限します。あるランクが停止した場合、このタイムアウトにより、他のランクが永遠に待機することを防ぎます。

集約操作を頻繁に行わないワークロードには、より長いタイムアウトを設定するか、None を渡して無効にします。

cb_config = ContinuousBatchingConfig(cpu_group_timeout=600.0)

プレフィックスキャッシュ

システムプロンプトのような共通プレフィックスを複数のリクエストが共有する場合、マネージャーはそれらを再計算する代わりに、KVキャッシュブロックを再利用します。allow_block_sharingのデフォルト値はTrueです。プレフィックス共有は、すべてのアテンションタイプのアロケーターが共有をサポートしている場合にのみ有効になるため、フルアテンションのみのモデルでのみ使用されます。純粋なスライディングモデルやハイブリッドモデルでは、allow_block_sharingがTrueであっても無効になります。

cb_config = ContinuousBatchingConfig(
    allow_block_sharing=True,  # default
)

アテンションバックエンドを選択

Continuous batching は FlashAttention、SDPA、および "paged|eager" をサポートしています。モデルを読み込む際に attn_implementation を設定してください。

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-4B",
    attn_implementation="flash_attention_2",
    device_map="auto",
    dtype=torch.bfloat16,
)

Continuous batching は、eager や SDPA よりも FlashAttention を使用した方が高速に動作します。主な理由は、FlashAttention では attention mask が不要であるためです。モデルが FlashAttention をサポートしている場合、"eager" または "sdpa" で読み込まれたモデルは、FlashAttention 3 に切り替わります。FlashAttention 3 がインストールされていない場合は FlashAttention 2 に切り替わり、警告がログに出力されます。どちらのバージョンもインストールされていない場合、バックエンドは変更されません。

テスト時など、eager または SDPA を維持する場合は、ContinuousBatchingConfig で auto_switch_to_flash=False を設定します。バックエンドが "eager" のままの場合、continuous batching は代わりに "paged|eager" を実行します。

テンソル並列

1つのGPUに収まらないほど大きなモデルでは、テンソル並列性を用いて重みをデバイス間でシャード化します。デバイス数はDistributedConfig(tp_size=N)で設定します。Continuous batchingはモデルからテンソル並列サイズを読み取り、シャードごとにpaged KVキャッシュのサイズを決定します。サポートされているアーキテクチャのリストとシャード化の仕組みについては、テンソル並列性を参照してください。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DistributedConfig
from transformers.generation import ContinuousBatchingConfig, GenerationConfig

distributed_config = DistributedConfig(tp_size=4)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-32B",
    attn_implementation="flash_attention_2",
    distributed_config=distributed_config,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-32B")

inputs = [tokenizer.encode(p) for p in ["What's up?", "Name a cat breed."]]
generation_config = GenerationConfig(max_new_tokens=64, eos_token_id=tokenizer.eos_token_id)

outputs = model.generate_batch(inputs=inputs, generation_config=generation_config)

torchrunでスクリプトを起動し、--nproc-per-nodeをシャード対象のGPU数に設定します。

torchrun --nproc-per-node 4 cb_tp.py

テンソル並列サイズは、モデルのnum_key_value_headsで割り切れる必要があります(モデルの設定を確認してください)。そうでない場合、ページドキャッシュは起動時にエラーを発生させるため、適切な--nproc-per-nodeを選択してください。

device_map と distributed_config を同時に設定しないでください。device_map はモジュール全体を特定の GPU に配置するのに対し、テンソル並列では同じパラメータをすべての GPU に分割するため、両者は競合します。

スライディングウィンドウアテンション

スライディングウィンドウアテンション（Mistral、Gemma 2）を持つモデルは、continuous batching で動作します。スライディングレイヤーは SlidingAttentionCacheAllocator を使用し、小さなページセットを再利用することで、直近のトークンの固定ウィンドウのみを保持します。アテンションの種類は ContinuousBatchingConfig ではなく、モデルの layer_types または sliding_window から取得されます。

from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("google/gemma-2-2b")
config.sliding_window = 4096

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2-2b",
    config=config,
    attn_implementation="sdpa",
    device_map="auto",
    dtype=torch.bfloat16,
)

Prefix cachingとdecode fast pathは、sliding attentionが存在する場合、フルアテンション層を備えたハイブリッドモデルを含むすべてのモデルで無効になります。

次のステップ

• Continuous batching ブログ記事では、KV キャッシング、チャンク化プリフィル、動的スケジューリングについて、パフォーマンスベンチマークの数値とともに解説しています。
• 連続バッチ処理システムの動作について詳しく知りたい場合は、Continuous batching architectureドキュメントを参照してください。