Hard Research

KV Cache for Autoregressive Decoding

Implement kv_cache_step for one attention head and one new token. Reuse the stored key and value vectors, append the new token's key and value, and compute attention for its query.

Caching saves the work of projecting past tokens again. The new query still attends to the entire available prefix, including the new token itself. The projections are supplied here; you do not compute embeddings or apply query, key, value or output-projection matrices.

Inputs

cached_K: (L, D), keys for the previous L tokens in sequence order.
cached_V: (L, D), values for those same tokens in the same order.
new_q: (1, D), the new token's query.
new_k: (1, D), the new token's key.
new_v: (1, D), the new token's value.

L may be zero; D is positive. All arrays are finite, share a floating-point dtype and device, and have the stated shapes. Keep the row dimension even for the single new token. Leave every input unchanged.

Append, then attend

Append new_k after the cached keys, and new_v after the cached values:

full_K = concatenate([cached_K, new_k], axis=0)
full_V = concatenate([cached_V, new_v], axis=0)
scores = new_q @ full_K.T / sqrt(D)
weights = softmax(scores, axis=-1)
output = weights @ full_V

Both caches now have shape (L+1, D). Scores and weights have shape (1, L+1). Normalize over all L+1 keys together. Use a stable softmax, subtracting the maximum score before exponentiation; a framework softmax helper is allowed.

There are no future tokens in these caches, so every key is allowed. Do not apply a triangular mask to the 1-by-(L+1) score row: its row index is zero, but the query's sequence position is L. Masking by that local row index would incorrectly hide most of the prefix. Do not recompute outputs for old queries or add the query vector to either cache.

Return value

Return a dictionary with exactly these entries, preserving dtype and device:

"output": (1, D), the attention result for the new token.
"updated_K": (L+1, D), the cached keys followed by new_k.
"updated_V": (L+1, D), the cached values followed by new_v.

For the next token, pass updated_K and updated_V back as the caches, along with that next token's new query, key and value. Preserve matching key/value row order as the sequence grows.

Boundary and example

With an empty cache, the only key is new_k. Its attention weight is one, so output equals new_v for any query.

If new_q is all zeros, every score is zero and output is the average of all cached values plus new_v. For cached_V=[[1,0]] and new_v=[[0,1]], that gives output=[[0.5,0.5]]. Keys still appear in updated_K in their original order even when the attention weights are equal.

In the causal panel below, this operation corresponds to the last query row: it can use every key up to and including its own position.

Loading visualization…

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…