Chunked Prefill
A 32k-token prefill blocks every decode behind it for a second. Slice the prefill into chunks and mix them into decode batches.
- inter-token latency spikes
- chunk size tuning
- piggybacking
- TTFT/ITL trade-off
Motivation
The problem
The S10 simulator showed it as a purple spike: one step where the entire token budget goes to a single long prefill, and every decode behind it produces nothing.
Users experience this as the output freezing. On a 32k-token prompt it can approach a second, long enough that people reload the page. The cause is invisible to the users affected: someone else sent a long prompt.
Reordering cannot fix it. The problem is not which request runs, but that prefill is indivisible. 8,192 tokens go through in one pass, or they do not go through at all.
Core idea
The solution
Make it divisible. Nothing requires a prefill to complete in one forward pass; attention is causal, so the first 1,024 tokens can be processed, their KV written to the cache, and the next 1,024 processed in a later step against the cache the first chunk left behind.
Once prefill is splittable, the scheduler fills each step to its token budget with every decode plus as much prefill as fits. Step size becomes nearly constant, and the spike disappears.
def schedule(self):
budget = self.max_num_batched_tokens
batch = []
for req in self.running: # decodes first, 1 token each
batch.append(req.decode_slice()); budget -= 1
for req in self.waiting_or_partially_prefilled:
take = min(req.prompt_len - req.computed, budget, self.max_chunk)
if take <= 0:
break
batch.append(req.prefill_slice(take)) # a SLICE, not the whole prompt
req.computed += take
budget -= take
return batchChunking also removes the starvation bug
Mechanics
How it works
Why mixing prefill and decode is efficient, not just fair
There is a second reason this works, and it is the more interesting one. Prefill is compute-bound; decode is memory-bandwidth-bound. They contend for different resources.
A pure decode step leaves the arithmetic units idle. A pure prefill step saturates them and leaves the memory system relatively free. Mix the two and the decode tokens ride along on weight reads the prefill is doing anyway: adding 24 decodes to a prefill step costs far less at the margin than running those decodes in a step of their own.
So chunked prefill usually increases total throughput while reducing latency variance. SARATHI, the paper that introduced chunked prefill, calls this ride-along effect piggybacking; its follow-up Sarathi-Serve built a scheduler around it and called the result stall-free batching. The name refers to decodes never stalling behind a prefill, not to the throughput gain.
Choosing the chunk size
The chunk size is the same knob as the token budget, viewed from the other side. Smaller chunks mean smoother inter-token latency and more fixed per-step overhead; larger chunks mean better prefill efficiency and the spike creeping back.
A reasonable starting point: choose the largest chunk whose step time still fits your inter-token latency target, then measure. vLLM enables chunked prefill by default and sizes the token budget in the thousands (8,192 in V1 online serving). For most models that is a step of a few tens of milliseconds. Check your version, since the default has moved.
Explore
Try it
- worst stall for decoders
- 24 ms
- TTFT for the big prompt
- 189 ms
- steps over ITL target
- 0
- decode throughput
- 1451 tok/s
- prefill work
- decode work
Every step carries a slice of the prompt and a token for each of the 24 streaming users. No step is much longer than any other, so nobody sees a stall.
smaller chunks
- + smoother inter-token latency
- + new requests start sooner
- − more steps, so more fixed overhead
- − worse TTFT for the long prompt
larger chunks
- + better prefill efficiency
- + lower TTFT for the long prompt
- − latency spikes return
- − at chunk = prompt, you are back where you started
Set chunk size equal to the prompt length and the chunked path becomes the unchunked one: the two toggles produce the same timeline. Chunking is not a free win. It converts one large latency spike into a slightly longer TTFT, and that is almost always the better trade, because the spike is felt by every user and the TTFT by one.
Push the prompt to 32k with chunking off: one bar towers over the ITL target and every decode is frozen behind it. Turn chunking on and the same work becomes thirty-two ordinary steps. Then set the chunk size equal to the prompt length and watch the spike return. Chunking is about step granularity and nothing else.
Build it
Implementation
The subtlety is in the attention mask. A chunk's queries attend to all previously cached keys, plus the keys within their own chunk under causal masking.
def chunk_attention_mask(chunk_len: int, cached_len: int):
"""
Queries: the chunk_len new tokens.
Keys: cached_len already-computed tokens, then the chunk itself.
"""
total_keys = cached_len + chunk_len
mask = np.zeros((chunk_len, total_keys), dtype=bool)
# every new query may attend to ALL cached keys — no causality needed
mask[:, :cached_len] = True
# within the chunk, ordinary causal masking applies
mask[:, cached_len:] = np.tril(np.ones((chunk_len, chunk_len), dtype=bool))
return maskPosition IDs are absolute, not chunk-relative
$ python code/s11_chunked_prefill.pyOnly NumPy is required — setup instructions.
Production notes
In production
SARATHI is the original paper; it introduced chunked prefill and decode piggybacking. Sarathi-Serve, its follow-up, added the stall-free batching scheduler. vLLM exposes it as enable_chunked_prefill, on by default in v1, with prefill and decode sharing one budget and one batch. TensorRT-LLM ships the same idea under the name "chunked context". The alternative is not to mix them at all, but to run prefill and decode on separate machines. That is disaggregation, the subject of S18.
Exercises
- 1Sweep chunk size from 128 to 8,192 and plot P99 inter-token latency against prefill throughput. Identify the knee for a 50 ms latency target.
- 2Combine chunking with prefix caching: a chunk whose blocks are already cached should be skipped entirely. Measure the interaction; the two features compose better than either alone.
- 3Implement the chunk-relative position bug on purpose, then write the test that catches it. Comparing the KV cache against an unchunked prefill is the cheapest regression test available, and it catches an entire class of errors.
Continue
What's next
Scheduling is in good shape. The attention kernel itself still builds an N×N score matrix in memory, and at long context that matrix is both the memory bottleneck and the speed bottleneck. S12 removes it with online softmax.
Check yourself
Questions
Answer before you look. A wrong answer tells you more than a right one, because the explanation points at the part of the chapter to reread.
Besides smoothing latency, what correctness problem does chunked prefill solve?