Tutorial 15: Continuous Batching and Scheduling¶
Serving an LLM to multiple users simultaneously requires more than running one request at a time. Continuous batching lets new requests join and finished requests leave the batch between decode steps, maximizing GPU utilization.
1. Static vs Continuous Batching¶
Static batching: All requests in a batch must finish before new ones can start. If one request generates 1000 tokens and another generates 10, the second request's GPU slot sits idle for most of the batch.
Continuous batching: Requests can enter and exit the batch at any step. When request A finishes, request B from the queue takes its slot immediately.
Static (wasteful):
Slot 1: [=======A========================] [idle......]
Slot 2: [==B==] [idle........................................]
Continuous (efficient):
Step 1: [=======A========================] [==B==]
Step 2: [=======A========================] [===C===]
Step 3: [=======A========================] [====D====]
Step 4: [==E==] [====D====]
2. The Scheduler Architecture¶
The BatchedEngine manages the lifecycle of requests:
New Request
│
▼
┌───────┐ ┌───────────┐ ┌──────────┐ ┌──────────────┐
│ Queue │ → │ Prefilling │ → │ Awaiting │ → │ Decode Slots │
│ │ │ (chunked) │ │ Slot │ │ (batch_size) │
└───────┘ └───────────┘ └──────────┘ └──────────────┘
Each state has a specific purpose:
- Queue: Requests waiting to start processing
- Prefilling: Processing the prompt (one request at a time, in chunks)
- Awaiting Slot: Prompt processed, waiting for a free decode slot
- Decode Slots: Actively generating tokens (batch_size slots total)
3. The Step Function¶
The scheduler runs one step() per decode iteration. Each step performs
five operations in order:
def step(self):
"""
Advance the scheduling state by one decode step.
Returns True if there is still work to do, False if idle.
"""
if self.is_idle():
return False
# 1. Cancel any cancelled requests (user aborted, timeout, etc.)
self._sweep_cancellations()
# 2. Admit the next request from the queue into prefilling
self._admit()
# 3. Advance the prefill by one chunk (process more prompt tokens)
self._advance_prefill()
# 4. Move fully-prefilled requests into free decode slots
self._place_ready_requests()
# 5. Run one decode step across all active slots
self._decode_once()
return True
4. Chunked Prefill¶
Processing a 2000-token prompt in one shot may exceed memory limits or cause long stalls. Chunked prefill breaks the prompt into smaller chunks:
def _advance_prefill(self):
req = self.prefilling
if req is None:
return
# How many tokens remain in the prompt?
remaining = len(req.prompt_tokens) - req.offset
# Process at most prefill_step tokens (e.g., 128)
n = min(self.prefill_step, remaining)
chunk = mx.array([req.prompt_tokens[req.offset : req.offset + n]])
# Forward pass on just this chunk
logits = self.model(chunk, offset=[req.offset], cache=req.caches, logits_to_keep=1)
req.offset += n
# If we finished processing the entire prompt, sample the first token
if req.offset == len(req.prompt_tokens):
self.prefilling = None
token = self._sample_row(logits[:, -1, :], req.sampler)
self._handle_new_token(req, token, feed_offset=False)
if not req.finished:
self.awaiting_slot.append(req)
The key insight: during prefill, each chunk processes prefill_step tokens at
once, updating the KV cache for all of them. This is much faster than processing
one token at a time (which is what decode does).
5. Decode Across All Active Slots¶
After all active requests have been prefilled, they share a single batched forward pass:
def _decode_once(self):
active = [(i, req) for i, req in enumerate(self.slots) if req is not None]
if not active:
return
# Collect the next token for each active request
tokens = [req.next_token if req is not None else 0 for req in self.slots]
offsets = [req.offset if req is not None else 0 for req in self.slots]
# Batch them into a single tensor: (batch_size, 1)
y = mx.array(tokens).reshape(-1, 1)
# One forward pass processes all active requests simultaneously
logits = self.model(y, offset=offsets, cache=self.kv_cache, logits_to_keep=1)
# Sample next tokens independently for each request
last = logits[:, -1, :]
logprobs = last - mx.logsumexp(last, axis=-1, keepdims=True)
sampled = [
self._sample_row(logprobs[i : i + 1], req.sampler)
for i, req in active
]
# Handle each sampled token
for row, (i, req) in enumerate(active):
self._handle_new_token(req, sampled[row])
if req.finished:
# Free the slot for the next request
for batch_cache in self.kv_cache:
batch_cache.remove_request(i)
self.slots[i] = None
req.slot = None
6. Request Lifecycle¶
Here is the complete lifecycle of a request through the scheduler:
1. submit() → request enters Queue
2. _admit() → moves to Prefilling (creates per-request KV caches)
3. _advance_prefill() → processes prompt in chunks
4. When prefill complete → moves to Awaiting Slot
5. _place_ready_requests() → moves to a Decode Slot
6. _decode_once() → generates one token per step
7. _handle_new_token() → checks EOS, max length
8. When finished → slot freed, pages returned
Token Handling¶
def _handle_new_token(self, req, token, feed_offset=True):
# Check for end-of-sequence
if token in req.eos_token_ids:
self._finish(req, "stop")
return True
# Record the token
req.output.append(token)
req.events.append(token)
req.generated_count += 1
# Check for max length
if req.generated_count >= req.max_new_tokens:
self._finish(req, "length")
return True
# Update state for next step
req.next_token = token
if feed_offset:
req.offset += 1
if req.offset + 1 >= req.max_seq_len:
self._finish(req, "length")
return True
return False
7. The Poll Interface¶
The server needs to read generated tokens without blocking the scheduler.
The poll() method atomically drains events and reads status:
def poll(self, req_id):
"""
Atomically drain new tokens AND read current status.
This must be atomic: if drain() and status() were separate calls,
the status could change between them, and the final tokens could
be lost.
"""
req = self.requests.get(req_id)
if req is None:
return [], "unknown", None
events, req.events = req.events, [] # atomic swap
if req.finished:
return events, "done", req.finish_reason
return events, "running", None
8. Scheduling Efficiency¶
The scheduler's efficiency depends on the balance between:
- Prefill speed: How fast new requests join the decode batch
- Decode speed: How many requests can be active simultaneously
- Queue depth: How many requests are waiting
If prefill is too slow: queue grows, latency increases
If batch is too small: GPU is underutilized
If batch is too large: decode latency increases (larger matmuls)
The prefill_step parameter controls the chunk size — smaller chunks let
decode requests interleave with prefill, improving latency for active requests.
9. Summary¶
| Phase | What Happens | GPU Usage |
|---|---|---|
| Prefill | Process prompt tokens in chunks | High (large matmuls) |
| Awaiting Slot | Wait for a free slot | None (idle) |
| Decode | One token per step across all slots | Moderate (small matmuls, weight loading) |
Continuous batching achieves high GPU utilization by filling decode slots as quickly as possible and never leaving them idle when there is work to do.