Continuous Batching: Optimizing LLM Inference Throughput › Implementation Considerations
Language AI HandbookAll chapters Continuous Batching: Optimizing LLM Inference Throughput › Implementation ConsiderationsOne section · Sections in this chapter9 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Building a production continuous batching system requires careful attention to several practical concerns that go beyond the basic algorithm. The theory is elegant, but the engineering details determine whether the system reaches its theoretical potential.
The main challenges fall into four areas: memory management, preemption handling, tensor operations, and monitoring. Each requires careful design. A system that gets the scheduling logic right but fails at memory management will hit out-of-memory errors under load. A system that handles memory well but lacks proper monitoring will be opaque when performance degrades.
Memory Management Integration
Continuous batching works best when paired with paged attention, as covered in the Paged Attention chapter. Without it, memory fragmentation limits how many requests can be active simultaneously. Traditional contiguous allocation reserves maximum-length KV cache space for each request, even if most requests use far less. Paged attention allocates memory in small blocks on demand, letting much higher concurrency.
The interaction between continuous batching and paged memory is symbiotic. Continuous batching needs fine-grained memory release to quickly fill freed slots. Paged attention gives exactly this: when a request completes, its pages return to the free pool immediately, ready for allocation to incoming requests. The combination reaches both high throughput (through continuous batching) and high memory efficiency (through paged attention).
@dataclass
class MemoryConfig:
"""Configuration for KV cache memory."""
total_blocks: int # Total memory blocks available
block_size: int # Tokens per block
def estimate_blocks_needed(seq_length: int, block_size: int) -> int:
"""Estimate memory blocks needed for a sequence."""
return (seq_length + block_size - 1) // block_size # Ceiling division
class MemoryAwareScheduler:
"""Scheduler that respects memory constraints."""
def __init__(self, config: MemoryConfig, max_batch_size: int):
self.config = config
self.max_batch_size = max_batch_size
self.free_blocks = config.total_blocks
self.active_requests: dict[int, int] = {} # request_id -> blocks_used
def can_admit(self, request_length: int) -> bool:
"""Check if we have memory for a new request."""
needed = estimate_blocks_needed(request_length, self.config.block_size)
return (
needed <= self.free_blocks
and len(self.active_requests) < self.max_batch_size
)
def admit_request(self, request_id: int, initial_length: int) -> bool:
"""Admit a request if resources are available."""
if not self.can_admit(initial_length):
return False
blocks = estimate_blocks_needed(initial_length, self.config.block_size)
self.active_requests[request_id] = blocks
self.free_blocks -= blocks
return True
def grow_request(self, request_id: int, new_tokens: int = 1) -> bool:
"""Allocate additional blocks as sequence grows."""
if request_id not in self.active_requests:
return False
# Check if we need a new block
current_blocks = self.active_requests[request_id]
current_capacity = current_blocks * self.config.block_size
# Simplified: assume sequence length = current_capacity + new_tokens
new_length = current_capacity + new_tokens
needed_blocks = estimate_blocks_needed(
new_length, self.config.block_size
)
additional = needed_blocks - current_blocks
if additional > 0:
if additional > self.free_blocks:
return False # Out of memory
self.free_blocks -= additional
self.active_requests[request_id] = needed_blocks
return True
def release_request(self, request_id: int):
"""Release all blocks for a completed request."""
if request_id in self.active_requests:
self.free_blocks += self.active_requests[request_id]
del self.active_requests[request_id]Handling Preemption
When memory pressure is high, the scheduler might need to preempt active requests to make room for higher-priority work. This involves saving the request's state, including the KV cache and generated tokens, for later resumption. Preemption is a powerful tool for managing resources but adds significant complexity.
The decision of which request to preempt involves tradeoffs. A policy must balance fairness against throughput while remaining practical to implement. Preempting a request that has generated many tokens wastes the computation invested in it. Preempting a request that has been preempted before may seem unfair. Different policies balance these concerns differently. vLLM's original implementation used a "beam search" style where preempted requests had their KV caches swapped to CPU DRAM, letting resumption without full recomputation when memory pressure eased. This swap-based approach preserves work in progress at the cost of CPU memory and PCIe bandwidth.
An important nuance is that preemption should be a last resort, not a routine operation. Systems that preempt frequently can enter a thrashing state where they spend more time managing preemption and resumption than doing useful generation. The right response to chronic memory pressure is usually to reduce the maximum batch size rather than to preempt aggressively.
from enum import Enum
class RequestStatus(Enum):
WAITING = "waiting"
RUNNING = "running"
PREEMPTED = "preempted"
COMPLETED = "completed"
@dataclass
class PreemptibleRequest:
"""Request that can be preempted and resumed."""
id: int
status: RequestStatus = RequestStatus.WAITING
tokens_generated: int = 0
preempt_count: int = 0
class PreemptionPolicy:
"""Policy for choosing which requests to preempt."""
@staticmethod
def longest_first(
requests: List[PreemptibleRequest],
) -> PreemptibleRequest | None:
"""Preempt the request that has generated the most tokens."""
running = [r for r in requests if r.status == RequestStatus.RUNNING]
if not running:
return None
return max(running, key=lambda r: r.tokens_generated)
@staticmethod
def most_preempted_last(
requests: List[PreemptibleRequest],
) -> PreemptibleRequest | None:
"""Prefer preempting requests that haven't been preempted before."""
running = [r for r in requests if r.status == RequestStatus.RUNNING]
if not running:
return None
return min(
running, key=lambda r: (r.preempt_count, -r.tokens_generated)
)Batched Token Operations
For efficiency, token generation for all active requests should happen in a single batched forward pass. This requires careful tensor management to handle requests at different sequence positions. Unlike training, where all sequences in a batch are padded to the same length, continuous batching must handle sequences of varying lengths efficiently.
The key challenge is that different requests have different position embeddings and attention masks. Position embeddings tell the model where each token sits in its sequence. Attention masks prevent tokens from attending to future positions. Both vary by request, but both must be batched together for efficient GPU utilization.
The solution is to process just one token per request per forward pass during the decode phase. Each request contributes a single token to the batch input, and each token carries its own position identifier. The position IDs differ across requests in the batch because each request is at a different point in its generation, but the batch still processes as a single matrix operation.
import torch
def prepare_batch_tensors(
active_requests: List[dict], vocab_size: int = 32000
) -> dict:
"""
Prepare input tensors for a continuous batch.
Each request has:
- 'last_token': The most recently generated token
- 'seq_length': Current sequence length (for position embedding)
"""
batch_size = len(active_requests)
# Each decode step processes just the last token
input_ids = torch.tensor(
[r["last_token"] for r in active_requests], dtype=torch.long
).unsqueeze(1) # Shape: [batch, 1]
# Position IDs vary by request
position_ids = torch.tensor(
[r["seq_length"] for r in active_requests], dtype=torch.long
).unsqueeze(1) # Shape: [batch, 1]
return {
"input_ids": input_ids,
"position_ids": position_ids,
"batch_size": batch_size,
}# Example batch preparation
example_requests = [
{"last_token": 1234, "seq_length": 50}, # Short sequence
{"last_token": 5678, "seq_length": 200}, # Medium sequence
{"last_token": 9012, "seq_length": 15}, # Very short
{"last_token": 3456, "seq_length": 500}, # Long sequence
]
tensors = prepare_batch_tensors(example_requests)Batch Tensor Preparation: Input IDs shape: torch.Size([4, 1]) Input IDs: [1234, 5678, 9012, 3456] Position IDs: [50, 200, 15, 500]
The key insight is that despite requests being at different points in their generation, they can share a single forward pass because attention masks and position embeddings handle the differences. The input is just one token per request, but each token gets its own position embedding corresponding to its place in that request's sequence. The four requests above are at positions 50, 200, 15, and 500 in their respective sequences, yet they process together in the same matrix operations. This batching is what makes continuous batching efficient rather than just correct.