Chunked prefill
Trace static and iteration-level schedules and explain chunked prefill, preemption and latency tradeoffs.
THE IDEA
Chunked prefill
Chunked prefill splits a long prompt into pieces and mixes one piece into each iteration under a total token budget. With a budget of 256 tokens and 4 decoding requests, each iteration carries 252 prompt tokens, so the 2000-token prompt needs 8 iterations of about 29 ms each instead of one 204 ms stall. The scheduler also gets a decision point between chunks.
Chunking gives the scheduler a decision point between chunks and keeps every decode iteration short. The whole earlier context still has to be processed; the prompt just finishes later.
Predict first. Does splitting a prompt remove the need to process its earlier context?
Chapter context
Requests arrive at different times and finish after different output lengths. Scheduling decides which work enters each iteration and how waiting prompts compete with ongoing generation.
TRY IT YOURSELF
Part 1. Do it hands-on
Work out chunk sizes and times by hand, then plan the chunks in code.
Chunking gives the scheduler a decision point between chunks and keeps every decode iteration short. The whole earlier context still has to be processed; the prompt just finishes later.
- To do: Each iteration may process at most 256 tokens in total; each of the 4 decoding sequences uses 1. How many prompt tokens can be prefilled per iteration?
- To do: The 2000-token prompt is split into chunks of 252. How many iterations until its prefill is done?
- To do: With the same toy costs (1 ms per decode, 0.1 ms per prefill token), ms for a full-chunk iteration (4 decodes + 252 prefill tokens)?
Part 2. Do it in code
- chunk_size(token_budget, decodes): prompt tokens that fit in an iteration after the decodes: max(0, token_budget − decodes).
- plan_chunks(prompt_tokens, chunk): the list of chunk sizes that adds up to prompt_tokens (all `chunk` except possibly a smaller last one). ValueError if chunk < 1.
- chunked_iteration_time(decodes, chunk): the toy cost of an iteration with decodes decoding sequences and `chunk` prefill tokens (1 ms per decode, 0.1 ms per prefill token).
- worst_stall(decodes, chunk): the longest wait a running request sees: the time of an iteration that contains a full chunk.
- tokens_processed(chunks): the total of all chunk sizes (the earlier context is still fully processed).
Stuck, or something not working? You can come back to it any time.
CHECK YOUR UNDERSTANDING
Does splitting a prompt remove the need to process its earlier context?
Chapter reference: math, worked example & code
This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.
Math and notation
In a simplified equal-cost decode schedule, utilization = active sequence-iterations / available sequence-iterations. This omits prefill costs and assumes uniform iteration time. A real scheduler needs measured compute, memory and latency constraints.
Worked example
Two slots serve requests A,B,C with output lengths 2,1,2. Static groups run A+B for two iterations, then C alone for two: 5 active slots out of 8 available. Continuous admission runs A+B, then A+C, then C: 5 active slots out of 6. Utilization rises from 62.5% to about 83.33% in this toy schedule; real token costs are not equal.
Reference code
The following example uses the Python standard library and does not call a remote service.
from math import isclose
static=[['A','B'],['A',None],['C',None],['C',None]]
continuous=[['A','B'],['A','C'],['C',None]]
util=lambda rows: sum(x is not None for row in rows for x in row)/(2*len(rows))
assert util(static) == .625
assert isclose(util(continuous),5/6)
Python examples are provided for study and local execution; this page does not execute Python.
Common mistakes
Continuous batching does not guarantee a fixed improvement. Chunked prefill changes scheduling granularity, not the model's context dependencies. Include cache movement or recomputation costs when comparing preemption policies.