Brayden Zhang

Chunked Prefill

2 min read

Idea to split a long prompt into fixed token chunks and schedule them alongside decode steps under a single per-batch token budget, so one large prefill doesn’t stall every running generation.

Algorithm

  • Without chunking, a 32k-token prefill occupies the GPU for hundreds of milliseconds and every decoding sequence will get increased inter-token latency. Chunked prefill caps each forward pass at a token budget : fill the batch with decode requests first (each costs 1 token), then top up with the next chunk of a prefilling request until the budget is spent.
  • A chunk attends to all previously cached blocks of its own sequence, so correctness needs only the right position offset and a causal mask over the chunk plus the cache.
  • Budget accounting must weight by attention cost, not token count alone, since a chunk late in a long prompt attends over a much longer cache than an early one.

From scratch

def schedule(decodes, prefills, budget=2048, chunk=512):
    """decodes: list of seq ids (1 token each). prefills: list of [id, done, total]."""
    batch, used = [(s, 1) for s in decodes], len(decodes)
    for p in prefills:
        if used >= budget: break
        take = min(chunk, budget - used, p[2] - p[1])
        if take <= 0: continue
        batch.append((p[0], take)); p[1] += take; used += take
    return batch, used
 
def step(model, batch, cache):
    for sid, n_tok in batch:
        pos = cache.length(sid)                      # absolute offset for RoPE
        x = cache.next_tokens(sid, n_tok)            # 1 for decode, `take` for prefill
        logits = model(x, cache[sid], pos=pos)       # causal over chunk + cached prefix
        if cache.length(sid) >= cache.prompt_len(sid):
            cache.emit(sid, logits[:, -1])           # only the last chunk produces a token

Notes

  • Prioritizing decodes protects inter-token latency; prioritizing prefills protects time-to-first-token, and the budget split is the direct TTFT versus TBT knob.
  • Chunking hurts raw prefill throughput slightly (smaller GEMMs, and the KV cache is re-read once per chunk), typically a few percent, in exchange for far tighter tail latency.
  • A piggybacked batch is a healthy mix: decode alone is bandwidth-bound and leaves the tensor cores idle, so adding prefill tokens is close to free until the budget binds.
  • Disaggregated serving is the alternative answer, running prefill and decode on separate hardware pools and shipping the KV cache between them.

Related: Continuous Batching Simulator, Greedy Decoding with KV Cache, Scheduling and Admission Control

Graph View