NNX MHA With KV Cache

Reuse earlier keys and values

When processing a sequence one token at a time, the current query needs the keys and values of earlier tokens. A KV cache stores those projections so later steps can reuse them. Each step projects only the new token, appends its K/V, and attends over the populated cache including itself.

This exercise implements one attention layer. It returns a feature vector for the new token, not a sampled token or a complete language-model output.

Separate parameters from decoding state

The four linear projections are learned parameters. The cache holds temporary state for one sequence:

• k_cache and v_cache: float32 buffers shaped (max_len, H, Dh).
• pos: an int32 scalar giving the next unused slot, initially zero.

Wrap these three arrays in nnx.Cache. Its type lets you select decoding state separately from parameters with nnx.state(model, nnx.Cache). It does not automatically manage capacity, reset state or enforce masking.

Read a variable's array with variable[...]. JAX's .at[idx].set(value) produces an updated array; assigning it back to the NNX variable updates the module state. The array expression itself is not a Python in-place write.

Process one token

Read idx = pos[...], project the token to Q/K/V shaped (H, Dh), then write K and V at slot idx. Set pos to idx + 1. Attention may use slots zero through idx, including the slot just written.

The buffers keep a fixed shape, which is convenient for JIT. Unfilled slots still occupy space, so exclude them from softmax. Replace their scores with negative infinity. A finite penalty such as -1e9 can fail when valid scores are even more negative.

At least the newly written slot is valid, so no row is fully masked. This assumes finite scores at valid positions. Softmax over a row of only negative infinities would be undefined.

A complete cache update example

The following example starts after projection, with one head of width one. Its zero query gives equal scores to every valid key.

import jax
import jax.numpy as jnp
from flax import nnx

keys = nnx.Cache(jnp.zeros((3, 1, 1), dtype=jnp.float32))
values = nnx.Cache(jnp.zeros((3, 1, 1), dtype=jnp.float32))
pos = nnx.Cache(jnp.array(0, dtype=jnp.int32))

def step(value):
    idx = pos[...]
    keys[...] = keys[...].at[idx].set(jnp.array([[1.0]]))
    values[...] = values[...].at[idx].set(jnp.array([[value]]))
    pos[...] = idx + 1
    scores = jnp.zeros((1, 3))
    valid = jnp.arange(3) <= idx
    weights = jax.nn.softmax(jnp.where(valid, scores, -jnp.inf), axis=-1)
    return (weights @ values[...].reshape(3, 1))[0, 0]

assert jnp.allclose(step(2.0), 2.0)
assert jnp.allclose(step(6.0), 4.0)
assert int(pos[...]) == 2
assert jnp.allclose(values[...].reshape(3), jnp.array([2.0, 6.0, 0.0]))

The second output averages 2 and 6. The unused zero slot contributes nothing because it is masked. Without masking, that average would be (2 + 6 + 0) / 3.

What caching saves

For dense projections of width D, computing K/V once per token costs O(T * D²) across T tokens. Reprojecting every prefix costs O(T² * D²). Attention still compares each query with its prefix, totaling O(T² * D) work. Caching does not make the whole attention computation linear in T.

This fixed-buffer implementation scores all max_len slots on every step, even those later masked out. Its attention work is O(T * max_len * D), and its K/V storage is O(max_len * D). The entrypoint also rebuilds the cache from history on every call. To reuse it across actual decoding steps, retain the same model and cache for the same sequence and weights.

A retained cache has a fixed capacity. Call step at most max_len times; this exercise does not implement overflow handling. A separate sequence needs separate or reset cache state. Changed weights or changed historical inputs invalidate cached projections.

The diagram compares attention patterns. The final row of causal self-attention corresponds to the new query reading its full prefix; the diagram does not depict cache writes or memory allocation.