Hard End-to-End

Causal LM with KV Cache Generation

Implement causal_lm_generate_with_kv_cache for one prompt. Generate tokens greedily, saving each block's keys and values so later steps can reuse them. Return the prompt followed by the generated tokens.

The cache avoids recomputing projections for previous positions. Attention for a new query still visits all cached keys; its cost grows with the prefix.

Inputs and output

prompt: (P,), a nonempty sequence of integer-valued IDs in [0, V), delivered as floats.
w_emb: (V, D), the token embedding table.
pos_embed: (T_max, D), absolute position embeddings.
blocks_weights: (L, 6, D, D), L post-normalization transformer blocks.
w_head: (D, V), the supplied vocabulary projection.
num_heads: a positive integer dividing D.
max_new_tokens: an integer controlling the number of new tokens.

P, V and D are positive; L may be zero. Shapes are consistent. Floating-point parameters are finite and share a dtype and device; prompt is on that device. Let M=max(0, max_new_tokens). The position table has at least P + max(M-1, 0) rows, enough for every position this function processes.

Return a one-dimensional integer array of length P+M on the input device. Preserve the prompt's token values, then append exactly M IDs. Leave all inputs unchanged. If max_new_tokens<=0, return the prompt converted to integer IDs.

At each choice, select the largest vocabulary logit; break exact ties by choosing the smallest token ID. There is no sampling, temperature, beam search or special end-of-sequence ID. Do not stop early when any particular token is generated.

Block architecture

The six matrices for block i are [w_q, w_k, w_v, w_o, w_mlp1, w_mlp2]. All have shape (D, D); the feed-forward width is D. Each block computes:

q, k, v = x @ w_q, x @ w_k, x @ w_v
a = attention(q, k, v) @ w_o
y = LN(x + a)
x = LN(y + GELU(y @ w_mlp1) @ w_mlp2)

Split projected features into consecutive heads of width D/num_heads. Scale query-key dot products by sqrt(D/num_heads), apply stable softmax over keys, and concatenate head outputs in their original order before multiplying by w_o.

LN normalizes each token across D features using population variance and (token - mean) / sqrt(variance + 1e-5), without learned scale or bias. Use the tanh GELU approximation:

GELU(u) = 0.5 * u * (1 + tanh(sqrt(2/pi) * (u + 0.044715 * u**3)))

There are no linear-layer biases, dropout or additional final normalization. Use w_head independently of w_emb. Build normalization, GELU, attention and cache management yourself; do not use packaged transformer layers or a generation helper. A stable softmax helper is allowed.

1. Process the prompt and save caches

Embed all P prompt tokens and add pos_embed[0:P], obtaining shape (1, P, D). Run the blocks in order. In each block, compute Q, K and V from that block's input, split into heads, and save K and V in separate per-block caches of shape (1, num_heads, P, D/num_heads).

Prompt attention is causal: query position i may use keys 0 through i, including itself. Exclude future scores with negative infinity before softmax, or normalize only over allowed keys. Subtract row maxima for stable softmax.

Project the final prompt states with w_head. The last prompt position's logits select the first new token. If M=1, append that token and return immediately.

2. Process one generated token at a time

For each of the remaining M-1 choices:

Embed the token just generated. If it is the first generated token, add pos_embed[P]; then advance the position by one each time.
At each block, compute this token's new Q, K and V from its current state. Append K and V to that block's existing caches along the sequence axis.
Attend the new query to all keys in the enlarged cache, then apply the output projection, residual normalizations and feed-forward branch above.
Project the final state with w_head, choose the next ID, and append it to the generated sequence.

The enlarged cache includes both earlier tokens and the token being processed. There are no future keys, so this single-query attention needs no triangular mask. Keep caches separate across blocks, and do not append a token twice.

For a prompt of length three and two requested tokens, positions 0, 1 and 2 predict the first token. Process that token at position 3 to predict the second. Return five IDs. The last generated token is returned without another forward pass, so no position embedding is read for it.

With L=0, skip all blocks and caches. Each choice comes directly from the current token-plus-position embedding projected by w_head.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Save earlier keys and values during generation

In a causal transformer, adding a later token does not change how earlier tokens attend to their past. During inference with fixed model behavior, their keys and values can be saved in a KV cache at each layer instead of recomputed.

For a four-token prompt, the first pass stores positions 0 through 3. Its final logits predict the first new token. After sampling that token, the next pass processes it at position 4, appends its key and value, and predicts another token. Sampling a token and processing it are separate steps; confusing them can skip or duplicate tokens.

Keep a separate cache for each layer

Each layer receives different hidden representations, so it needs its own keys and values. Key and value sequence lengths must agree. During single-token decoding, the query length is one, but it can attend to the whole available prefix allowed by the mask.

A cache avoids repeated processing of the prefix. It still consumes memory that grows with stored sequence length, and each new query must read its allowed past keys and values. A sliding-window cache may discard entries according to its window policy.

Compare against a run without a cache

For the same prefix, cached and uncached logits should agree within appropriate floating-point tolerance when inference is otherwise deterministic. Check padding, position offsets, head layout, and window rules if they differ.

Use each new token's actual position for position embeddings or RoPE. Do not restart at position zero. Cached values also belong to the model weights and prefix that produced them; unrelated prompts or changed weights require a fresh cache unless an explicitly valid shared prefix is being reused.

Detecting runtime…