Autoregressive Generation with a KV Cache: prefill then decode

The problem, from first principles

Without caching, generating token t+1 naively means re-running the whole sequence through the model again, recomputing keys and values for every position — O(t) wasted work per step. Since keys and values for already-seen positions never change (they only depend on the frozen input token at that position), they can be computed once and reused. Implement single-head causal self-attention generation with a KV cache: a "prefill" pass over the whole prompt at once (building the initial cache), then generating n_new_tokens one at a time, each time appending only the new token's key/value to the cache instead of recomputing the whole sequence.

From theory to code

# Prefill: compute K, V once for every prompt position and cache them.
K_cache, V_cache = X_prompt @ W_K, X_prompt @ W_V

# Each decode step, compute K/V ONLY for the new token, append to cache.
k_new, v_new = x_new @ W_K, x_new @ W_V;  K_cache <- concat(K_cache, k_new)

# Attend the new token's query against the FULL cache (old + new).
out_new = softmax(q_new @ K_cache^T / sqrt(d_k)) @ V_cache

Constraints

• Prefill must process the full prompt in a single batched causal attention call and populate the cache.
• Each decode step computes Q/K/V for exactly the single new token (not the whole running sequence).
• The new query attends over the entire cache including the just-appended new K/V (self-attention).
• Return the final generated hidden states for the new tokens AND the final cache length.

Hints