Skip to content

← Flax step 25 of 100

Hard Primitives

Multi-Head Attention with KV Cache

Jump to the task

Reuse earlier keys and values

During causal decoding, each new token can attend to itself and earlier tokens. With fixed model parameters and deterministic inference, the keys and values already computed for earlier positions can be reused. A KV cache stores them so the model does not have to project the whole prefix again.

Each decoding step projects one new input, writes its key and value into the next cache slot, and attends over the populated slots. The result is an attention output vector. Choosing the next token would require the rest of a language model, including its vocabulary projection and sampling rule.

This exercise runs only the first decoding step. It does not generate text or fill the cache with a prompt.

Allocate the cache, then populate it

Let H = num_heads, D = qkv_features and d = D // H. Construct nn.MultiHeadDotProductAttention with decode=True and initialize it directly on x, whose shape is (T, D_in).

The returned variables contain two collections:

  • params: the learned query, key, value and output projection parameters.
  • cache: cached_key and cached_value, each shaped (T, H, d), plus the scalar cache_index.

Initialization creates zero-filled cache arrays and sets the index to zero. Passing the whole x to init sizes the cache; it does not store the projected contents of x in that cache. The first call to apply must receive just x[0:1], with shape (1, D_in).

After that call, slot zero contains the first token's projected key and value, and the returned index is one. The single query can attend to slot zero; unused slots are masked. Its projected shape is (1, H, d), the cached keys and values have shape (T, H, d), and scores have shape (H, 1, T). The output projection produces (1, D_in).

Keep the returned cache for the next step

mutable=["cache"] allows updates to the cache collection during apply. The call returns (output, updates). It does not change the variables dictionary you passed in. For another step, use updates["cache"] with the same parameters. Reusing the original variables would restart at slot zero and discard the previous step's history.

Here is a complete two-step example. The requested function below only needs its first step:

import jax
import jax.numpy as jnp
import flax.linen as nn

x = jnp.array([[.2, -.3, .1], [.4, .1, -.2]], dtype=jnp.float32)
attn = nn.MultiHeadDotProductAttention(
    num_heads=2, qkv_features=4, decode=True
)
initial = attn.init(jax.random.PRNGKey(0), x)
assert int(initial["cache"]["cache_index"]) == 0
assert jnp.all(initial["cache"]["cached_key"] == 0)

first, changed = attn.apply(initial, x[0:1], mutable=["cache"])
state = {"params": initial["params"], "cache": changed["cache"]}
second, changed_again = attn.apply(state, x[1:2], mutable=["cache"])

assert first.shape == second.shape == (1, 3)
assert int(changed_again["cache"]["cache_index"]) == 2
assert int(initial["cache"]["cache_index"]) == 0

Allocate enough slots for the intended sequence and stop before exceeding that capacity. Once allocated, this decoding path expects one input position per call. A full prompt is not interchangeable with a one-position step.

What caching saves

At fixed feature widths, repeatedly projecting prefixes of lengths 1, 2, ..., T processes O(T²) positions. Caching reduces that to T new positions. Recomputing full dense attention over every prefix would process O(T³) query-key pairs; one new query per step needs O(T²) pairs in total when attending to the valid prefix.

This Linen implementation uses fixed-capacity arrays. A step forms scores against all allocated slots and masks unused ones, so its attention work depends on cache capacity as well as the number of populated positions. The cache stores 2 * T * D scalar values per unbatched layer. These are work and storage counts, not measured speedups. Initializing on the full input also performs a full module forward pass.

Problem

Implement mha_with_cache(seed, x, num_heads, qkv_features):

  1. Build nn.MultiHeadDotProductAttention(num_heads=H, qkv_features=D, decode=True). Keep the other defaults, including projection biases, zero dropout and output width D_in.
  2. Initialize it directly with jax.random.PRNGKey(int(seed)) and x.
  3. Apply the returned variables to x[0:1] with mutable=["cache"].
  4. Return that output flattened with .reshape(-1). This one-step exercise does not return the updated cache.

x is float32 with shape (T, D_in). Dimensions and head count are positive, and D is divisible by H. D_in need not equal D or be divisible by H. The output is a one-dimensional array of length D_in. If JIT-compiling this function, keep seed, num_heads and qkv_features static to match its initialization and shape setup.

With only slot zero available, each head's attention weight on that slot is one. The result matches position zero of causal attention with the same parameters, including its value and output projections. It is not generally equal to the raw first input vector.

In the causal panel below, the first row illustrates this first-step permission pattern. The panels compare allowed connections, not cache state.

Loading visualization…