Chunked prefill for long context

Long context is the feature everyone wants and nobody wants to pay for. A 128K-token prompt on a 70B model means billions of FLOPs of prefill compute, all in one shot. If you run that as a single continuous prefill, every decode request in the same batch has to wait for it to finish. Your P99 latency spikes, your throughput tanks, and your users notice.

Chunked prefill is the fix. Instead of processing the entire prompt in one monolithic forward pass, you break it into chunks of, say, 512 or 2048 tokens, and interleave those chunks with decode steps from other requests. It is one of those ideas that sounds obvious in retrospect but requires careful engineering to get right.

The problem: prefill as a bully

Recall from the roofline post that prefill is compute-bound. It does large matrix multiplications across all prompt tokens in parallel, which is great for GPU utilization but terrible for sharing. A single 128K prefill on an H100 can take several seconds. During that time, any request in the decode phase is stuck waiting for its next token.

In a continuous batching system like vLLM, new requests can join the batch at each iteration. But if one request is doing a monster prefill, the iteration itself takes a long time. Everyone else in the batch is hostage to the longest prefill.

The math is straightforward. For a transformer layer, the prefill compute for the attention block scales as O(n^2 * d) where n is the sequence length and d is the head dimension. Double the prompt length and you quadruple the attention compute. At 128K tokens, this dominates everything.

How chunked prefill works

The idea is simple: instead of prefilling all 128K tokens in one iteration, break them into chunks of a fixed size (the "chunk budget") and process one chunk per scheduler iteration. Between chunks, the scheduler can run decode steps for other requests.

# Conceptual scheduler loop with chunked prefill
chunk_budget = 2048  # tokens per iteration for prefill

while requests_pending():
    batch = []

    # 1. Schedule decode tokens first (they're cheap)
    for req in decode_queue:
        batch.append(req.next_decode_token())

    # 2. Fill remaining budget with prefill chunks
    remaining = chunk_budget - len(batch)
    for req in prefill_queue:
        chunk = req.get_next_chunk(max_tokens=remaining)
        batch.append(chunk)
        remaining -= len(chunk)
        if remaining <= 0:
            break

    # 3. Run the batch through the model
    forward_pass(batch)

The key insight is that decode tokens are always prioritized. Each decode step for a request is exactly one token, so it is cheap. The chunk budget controls how much prefill work gets done per iteration, bounding the worst-case iteration time.

What changes inside the model

Chunked prefill introduces a complication: the KV cache for a prefilling request is built incrementally. After processing the first chunk of 2048 tokens, those KV entries are in the cache. The next chunk attends to both the cached keys/values from previous chunks and the new tokens in the current chunk.

This means the attention computation for chunk i looks different from a standard prefill:

• Cross-attention to cached KVs: The new chunk's queries attend to all previously cached keys and values. This is a matrix-vector-like operation (similar to decode) and is memory-bound.
• Self-attention within the chunk: The new tokens attend to each other with a causal mask. This is the standard compute-bound prefill attention.

In practice, most implementations fuse these two components. vLLM's scheduler, for example, handles this transparently. The model sees a batch containing both "partial prefill" and "decode" tokens, and FlashAttention kernels handle the mixed attention patterns efficiently.

Choosing the chunk size

The chunk budget is the main knob, and it involves a real tradeoff:

• Smaller chunks (256 to 512 tokens): Lower iteration latency, better responsiveness for decode requests. But more iterations are needed to finish a prefill, so total prefill time (TTFT) increases due to overhead.
• Larger chunks (4096 to 8192 tokens): Better prefill throughput, fewer iterations to complete. But each iteration takes longer, so decode latency suffers during those iterations.

In vLLM, the default chunk size is 512 tokens for chunked prefill. Sarathi-Serve, the research system that pioneered this approach, found that chunk sizes between 256 and 2048 tokens give good tradeoffs for most workloads. The sweet spot depends on your model size, GPU, and the ratio of long-context to short-context requests in your traffic.

The Sarathi-Serve insight: stall-free batching

The Sarathi-Serve paper (Microsoft Research, 2024) formalized chunked prefill as "stall-free batching." Their key observation was that without chunking, continuous batching systems have a bimodal iteration time distribution: fast iterations when only decodes are running, and slow iterations when a prefill lands. This bimodality makes tail latency unpredictable.

With chunked prefill, every iteration does roughly the same amount of work. The iteration time distribution becomes unimodal. P99 latency drops dramatically, sometimes by 3 to 5x, even though median latency barely changes.

They also showed that chunked prefill actually improves GPU utilization in many cases. The prefill chunks fill the compute capacity that would otherwise be wasted during memory-bound decode iterations. You are essentially piggy-backing prefill compute on decode bandwidth.

Interaction with other techniques

Chunked prefill plays well with the rest of the inference stack:

• Prefix caching: If prefix caching gives you a cache hit, fewer chunks need to be prefilled. The chunk budget goes further.
• Speculative decoding: Both techniques improve decode throughput. Chunked prefill keeps the iteration time stable; speculative decoding produces more tokens per iteration.
• Disaggregated serving: In a disaggregated setup where prefill and decode run on separate GPUs, chunked prefill is less critical because there is no contention. But it is still useful for bounding memory on the prefill nodes.
• Context parallelism: For truly massive contexts (1M+ tokens), context parallelism shards the sequence across GPUs. Chunked prefill operates within each shard.

What I would measure

If you are enabling chunked prefill in production, here is what to watch:

• P99 inter-token latency: This should drop significantly compared to unchunked prefill. If it does not, your chunk size may be too large.
• TTFT for long prompts: This will increase slightly because of the chunking overhead. Track the regression and make sure it is acceptable for your SLOs.
• GPU compute utilization: Should stay high or improve. If it drops, you may be chunking too aggressively and creating too many small iterations.
• Throughput (tokens/second): Should be similar or better than without chunking, because you are filling GPU bubbles.

Chunked prefill is not an optimization. It is a scheduling discipline. It trades a small TTFT regression for dramatically better tail latency, and in production, tail latency is what your users feel.

Next up: small fine-tuned vs large quantized, the real deployment tradeoff nobody talks about enough.