Continuous batching

The technique that turns transformers serve from a toy demo into a production-capable inference endpoint. Instead of running each request as a fresh batch, the scheduler interleaves prefill and decode steps from many concurrent requests in the same forward pass, using a paged KV cache.

The feature shipped on 2025-08-26 in PR #40426 (initial implementation) and has been refined since.

Where it lives

src/transformers/generation/continuous_batching/
├── __init__.py
├── scheduler.py           # request queue, prefill/decode interleaving
├── paged_attention.py     # block-table builder
├── batch.py               # constructs the merged forward inputs
└── ...

Companion files:

• src/transformers/integrations/eager_paged.py — pure-PyTorch paged attention.
• src/transformers/integrations/sdpa_paged.py — SDPA paged attention.
• src/transformers/integrations/flash_paged.py — FlashAttention paged attention.
• src/transformers/cache_utils.py — PagedCache and friends.

Why it matters

Without continuous batching:

• Each request occupies a fixed slot until it finishes generating.
• A batch of 8 must wait for the slowest request before freeing GPU memory.
• Throughput collapses when request lengths are heterogeneous.

With continuous batching:

• The scheduler holds a pool of KV blocks and assigns blocks to requests as they grow.
• A single forward pass mixes "prefill" computations (a long prompt being ingested) and "decode" computations (single-token decoding for in-flight requests).
• Throughput stays close to optimal even with variable lengths.

Block-paged KV cache

The KV cache is a flat tensor of shape [num_blocks, block_size, num_heads, head_dim]. Each request owns a logical sequence built from a list of physical block indices (the block table). When a request grows past its current allocation, the scheduler appends a free block.

graph TD
    R1[Req 1, 2 blocks] --> Pool[(KV pool, N blocks)]
    R2[Req 2, 5 blocks] --> Pool
    R3[Req 3, 1 block] --> Pool
    Pool --> PA[Paged attention kernel]
    PA --> Logits[Per-request logits]

The sequence-to-block mapping is the block table, and it is consumed by the paged attention kernels.

Forward pass shape

A continuous-batched forward pass concatenates all active sequences into one flat sequence and uses cumulative sequence lengths plus a block table to keep them logically separate. Conceptually:

• input_ids = [tokens of req1, tokens of req2, ...] (1D)
• cu_seqlens = [0, len1, len1+len2, ...]
• block_table = [...] mapping logical token positions to KV blocks.

Models that support this declare _supports_paged_attention = True (or are routed automatically when SDPA/FlashAttention paged variants are available).

Scheduler

The scheduler in src/transformers/generation/continuous_batching/scheduler.py maintains:

• A pool of free blocks.
• A queue of pending requests with their tokenized prompts.
• A list of in-flight requests with current decoded ids and remaining max length.

Each tick:

1. Admit pending requests until the prefill token budget (cb_max_batch_tokens) is hit.
2. Emit a forward pass mixing prefill chunks and per-request decode tokens.
3. Update each request's KV blocks, and release blocks for finished requests.

Tunables

transformers serve exposes the relevant CLI flags:

Flag	Purpose
--continuous-batching	Enable.
--cb-block-size	Tokens per KV block (default model-dependent).
--cb-num-blocks	Pool size. Determines the maximum total context held in GPU.
--cb-max-batch-tokens	Max tokens admitted into a single forward pass.

When to use it

Continuous batching helps when:

• The service handles many concurrent users.
• Request lengths are variable.
• Throughput matters more than per-request latency.

It is overkill (and adds complexity) for batch-1 inference; greedy model.generate(...) with a DynamicCache is simpler.

Limitations

• Not all model architectures advertise paged-attention support yet. Encoder-decoder, some hybrid (Mamba / SSM) models, and a few VLMs are still being migrated.
• Custom logits processors that depend on the full causal context (e.g., n-gram penalty) may need extra wiring.
• Scheduler heuristics are conservative; production deployments may want to tune cb_* knobs.

Integration points

• Generation — generate is replaced by the scheduler in serve mode.
• Cache — paged caches.
• Attention — paged kernels.
• CLI — transformers serve --continuous-batching.
• Serving — the HTTP serving feature relies on this scheduler.

Entry points for modification

• New paged kernel → add to src/transformers/integrations/<backend>_paged.py and register in the dispatcher.
• Scheduler heuristics → src/transformers/generation/continuous_batching/scheduler.py.
• Tests → tests/generation/test_continuous_batching*.py.