Prefill and Decode Perform Different Work

Separate prompt prefill, which processes every prompt position, from one-token decode, which consumes one new input and returns one new logits row. Track their distinct shapes and state transitions.

Inference has two phases with the same model parameters but different sequence shapes.

Prefill processes the known prompt. For prompt length $P$ , each layer can calculate all prompt queries, keys, and values together under a causal mask. It returns prompt logits and fills caches with $P$ keys and $P$ values.

Decode processes one new input token. It creates one query, one key, and one value per head and layer, appends the new key and value, and calculates only the new logits row.

Phase	New input shape	Key length after phase	Output used
prefill	$(B,P)$	$P$	final prompt logits row
one decode step	$(B,1)$	$P+1$	new logits row

Sequential Prefill Is a Correctness Tool

The downloadable implementation fills the prompt cache by calling the same one-token function repeatedly. This makes every append and position visible. It then compares the result with the registered model's parallel full-prompt forward pass.

An optimized implementation normally prefills all prompt positions together. The sequential version is not a performance recommendation; it is a smaller surface for proving the cache logic.

The Handoff Must Agree on Length

If prefill caches $P$ positions, the first generated token that is fed back uses position index $P$ . Appending the prompt twice or starting decode at $P-1$ creates a valid-shaped but different computation.

Q1. Distinguish the two phases

A prompt contains three tokens. After processing it, the cache length is three. The selected token is then fed back. What are the new-input length and resulting cache length for that decode step?

Your answer

Answer it first, then check.

Hint

The prompt is already represented in the cache.

Solution

The decode input has length 1; the cache grows from 3 to 4.

Both Phases Need Evaluation State

Use the same weights, normalization behavior, position convention, dtype, and device across the handoff. A cache produced by a different model state is not a valid prefix representation.