Splitting Prefill From Decode Inside a Step

The two phases of a transformer inference request have opposite hardware profiles. Processing a prompt of $T$ tokens performs $T$ tokens of attention work against a growing context, so the arithmetic per byte read from memory is high and the operation saturates the matrix units. Generating a token performs one position of work against a context of length $T$, so the arithmetic per byte is low and the operation waits on memory. Putting both into the same batch is attractive because it fills both kinds of idle resource. It is also the source of the worst tail latency in a serving system, and the scheduling problem is to get the benefit without the tail.

Two kernels with different limiting resources

Prefill on a 2,048-token prompt reads the weights once and applies them to 2,048 positions. If the model’s weights are $N$ parameters, the step reads $Nb$ bytes and performs roughly $2NT$ operations, so arithmetic intensity scales with $T$ and a long prompt is comfortably compute-bound. The time to process the prompt is approximately the total floating-point work divided by peak throughput, and it grows linearly with prompt length.

Decode at batch $B$ reads the same $Nb$ bytes and performs $2NB$ operations, with $B$ typically in the tens to low hundreds. Arithmetic intensity is $B$ times lower than prefill at the same length. Time per step is bounded below by the weight read divided by bandwidth, and adding sequences to the batch costs almost nothing until the compute units saturate.

A mixed step that contains one 2,048-token prompt and 64 single-token decodes therefore contains work that is 32 times the decode work, all of it landing on the compute-bound side of the machine. The step’s duration is set by the prompt, and all 64 generating sequences wait for it. The interference is asymmetric: the prefill dominates the step time, and the decodes contribute latency without contributing throughput.

This is the argument for a separate prefill phase. Running the prompt alone gives it the whole device and completes it sooner; running decodes alone keeps their step time at the bandwidth floor. The cost of separation is idle resource: during prefill, memory bandwidth is underused, and during decode, the matrix units are underused.

Why the split is rarely a clean partition

Two-phase execution is the natural first design, and it fails under load for a reason that has nothing to do with the kernels. Decode steps must happen continuously; a decode sequence that is not stepped for a while is not making progress, and its inter-token latency grows. If prefill occupies the device for a long stretch, every generating sequence stalls for that stretch. The stall is bounded by the longest prompt in the prefill queue, and prompts can be arbitrarily long.

The naive fix is to give prefill a small share of each step, so decodes keep moving. That reintroduces the interference, because a prefill slice inside a decode step still lengthens the step by the amount of prompt work it contains. The design question is therefore not whether to mix but how much prompt work to place in each step.

Chunked prefill as a rate control

Chunked prefill bounds the answer. A long prompt is split into fixed-size chunks of $C$ tokens, and each step processes at most one chunk per admitted prefill, alongside the running decode batch. The step time becomes approximately the decode step time plus the cost of $C$ positions of prefill. With $C$ chosen so that the added work is comparable to the decode work, the step stays close to the decode bandwidth floor and the prompt completes over $T/C$ steps.

The scheduler’s decision per step is which chunks to include, subject to a token budget:

def build_batch(running, prefills, token_budget):
    batch = [(seq, 1) for seq in running]        # one position each
    used = len(running)
    for p in prefills:
        remaining = token_budget - used
        if remaining <= 0:
            break
        take = min(p.chunk_size, p.remaining, remaining)
        batch.append((p, take))
        used += take
    return batch, used

The token budget is the control knob. It is set from the target step latency and the device’s throughput for mixed work, and it is bounded above by whatever keeps the decode latency within its requirement. A budget equal to the batch size means the prompt work is roughly the size of the decode work; a budget much larger means prompts are prioritised and decodes suffer.

The gain is that prompt processing no longer arrives in indivisible lumps. A request with a 30,000-token prompt is admitted as 30,000/$C$ chunks and begins generating after the first chunk plus a full forward pass, rather than after the entire prompt is processed. This matters for time-to-first-token, which is the metric the first chunk determines.

The arithmetic of chunk boundaries

Chunking does not reduce attention work, because causal attention is triangular either way: a prompt of $T$ tokens performs about $T^2/2$ position-pairs of attention whether it runs in one pass or in $T/C$ chunks. What chunking changes is the shape of each kernel invocation and the number of weight reads. A single pass over a $T$-token prompt reads the weights once. $T/C$ chunks read them $T/C$ times, and each chunk’s attention runs over a shorter query dimension, which lowers the arithmetic intensity of that kernel and gives the matrix units less to work with. For a long prompt this is a large increase in memory traffic and a loss of intra-pass parallelism, and it is the reason chunk sizes are not made arbitrarily small. The chunk size trades weight re-reads and kernel efficiency against step-latency interference, and the balance point sits where the added traffic is a small fraction of the weight read the decode batch was going to perform anyway.

There is a second-order effect on the cache. Each chunk writes its keys and values into the sequence’s cache as it completes, so the sequence’s residency grows monotonically across chunks, exactly as if the prompt had been processed in one pass. No extra cache memory is required beyond the final size, but the allocation must be reserved before the first chunk or the sequence will fail partway through prompt processing, which is a worse failure than failing at admission. This is why admission against the declared maximum length exists; see how key-value bytes set the batch ceiling.

Choosing the priority order

With a token budget per step, the remaining decision is ordering. Three policies are common and they optimise different metrics.

Decode-first. Fill the step with the running decode batch, then spend the remainder of the budget on prefill chunks. This protects inter-token latency and time-to-first-token suffers when the queue is deep, because a request may wait many steps before its first chunk runs.

Prefill-first. Spend the budget on prompt chunks and run decodes with whatever capacity remains. Time-to-first-token improves and inter-token latency degrades. For an interactive workload where a user is waiting for the first token this can be the right trade, and it is the wrong trade for a workload already streaming.

Fair share. Allocate the budget proportionally between the two classes, so both make steady progress. This bounds the worst case for both metrics at the cost of being optimal for neither. It is what an SLO-driven scheduler converges to when it is given latency targets for both, which is the subject of deadline-aware batch assembly.

The ordering also interacts with admission. A prefill chunk that is admitted must have its cache reserved, so admitting a prompt is a memory decision as much as a scheduling one. A scheduler that admits many prompts to keep the prefill queue short can exhaust the pool and force eviction of running sequences, which converts a latency problem into a correctness-of-progress problem.

When separation beats sharing

The mixed-step design is not universally correct. Two configurations favour running prefill and decode on distinct devices or distinct pools.

The first is a workload with very long prompts and very short outputs, where the prompt work dominates total cost. Mixing gives the prompt a fraction of the device per step and stretches it over many steps, which is strictly worse than processing it at full width. Batching prompts together, rather than mixing them with decodes, is the better use of the compute-bound regime.

The second is a workload where the two phases are served for different clients with different latency requirements. Physically separating them allows each to be sized and scheduled independently, which is the argument for separating prefill and decode into distinct pools. The cost is that the cache must move between pools, which turns a scheduling problem into a transfer-bandwidth problem.

Between those extremes, chunked prefill with a token budget is the general answer. Its failure modes are worth naming: a budget set too high produces decode latency spikes on every long prompt arrival; a budget set too low re-reads weights so often that aggregate throughput falls; and a chunk size that does not divide the prompt evenly leaves a final partial chunk whose shape differs from every other, which some fused kernels handle poorly.