FLAGSHIP BUILD LAB

Build causal attention with a correct incremental KV cache

Download the CPU build · Download the correctness tests

The problem

A decoder predicts one token, appends it to the sequence, and predicts again. A simple implementation recomputes queries, keys, and values for the whole prefix after every new token. Causal attention makes most of that work unnecessary: an earlier token cannot depend on a token that arrives later.

A key-value cache stores projected keys and values for the unchanged prefix. It does not cache attention weights or final outputs. The new query still scores every allowed cached key and mixes the corresponding values.

The difficult part is correctness. A wrong axis, a triangular mask built from local instead of absolute positions, or a restarted position encoding can preserve valid tensor shapes while changing every later result.

Work through a small example

Use batch size B=1, heads H=2, sequence length T=4, and head width Dh=3. Projected Q, K, and V each have shape (1, 2, 4, 3). Their score matrix has shape (1, 2, 4, 4).

For query position 2, only key positions 0, 1, and 2 are legal. Future position 3 receives zero weight. Padding adds another rule: a key must be both non-future and real.

After prefill with four tokens, the cache has K,V: (1,2,4,3). One new token produces Q: (1,2,1,3) and new K,V: (1,2,1,3). Append on the sequence axis to get K,V: (1,2,5,3). The rectangular score matrix is (1,2,1,5).

The general rule

For each head, scaled dot-product attention is:

A = softmax((QKT / √Dh) + M)

O = AV

The additive mask M is zero for allowed pairs and effectively negative infinity for blocked pairs. The code uses a boolean allowed-mask first, then fills blocked scores with the smallest representable value before softmax and sets their final weights to exact zero.

With cached prefix length L and a new chunk of length C:

• new queries use absolute positions L … L+C−1;
• all keys use positions 0 … L+C−1;
• query position q may use key position k only when k ≤ q;
• the validity mask also requires that key k is a real token.

A local C × (L+C) lower triangle is wrong because its first row starts at local position zero. Compare absolute positions instead.

Why positions belong in the correctness contract

The build adds the same deterministic sinusoidal position signal in full and cached modes. If cached decoding restarts at position zero, the new query and key projections change. The mask can still look correct, so a tensor-equivalence test is needed.

The API treats an explicit query_offset as an assertion. If a cache contains three positions, the next chunk must start at 3. A request for offset 0 raises an error instead of silently returning plausible but inconsistent values.

Implement and inspect

The implementation is a small nn.Module with separate query, key, value, and output projections. It uses no optimized attention shortcut so every shape and mask remains visible.

with torch.no_grad():
    full = layer(inputs, valid_tokens=valid)

    cache = None
    pieces = []
    for start, stop in ((0, 3), (3, 5), (5, 7)):
        piece, cache = layer(
            inputs[:, start:stop],
            valid_tokens=valid[:, start:stop],
            cache=cache,
            use_cache=True,
        )
        pieces.append(piece)

    cached = torch.cat(pieces, dim=1)
    torch.testing.assert_close(cached, full, atol=2e-6, rtol=2e-6)

Run python examples/flagship/causal_attention_kv_cache.py. The example uses B=2, T=6, D=12, and H=3. One batch row contains two masked padding positions.

Expected structure:

output shape: (2, 6, 12)
weight shape: (2, 3, 6, 6)
cache shape: (2, 3, 6, 4)
maximum full-vs-cached error: a small floating-point value near zero
future weight above diagonal: 0.0

The exact error can vary slightly with the PyTorch build and hardware. The test checks a tolerance of 2e−6; it does not require bit-for-bit equality because matrix multiplication order may differ.

Engineering checks

1. Full versus one-token decode. Process every position as a chunk of length one and compare every output tensor with a full pass.
2. Full versus multi-token chunks. Use chunks such as (3,2,2). This catches incorrect within-chunk masking.
3. No future-token leakage. Replace tokens after position 3 with large random values. Outputs through position 3 must remain unchanged.
4. Padding is invisible as a key. Every attention weight aimed at a masked key must be exactly zero.
5. Rows normalize. Allowed attention weights must sum to one for every batch, head, and query.
6. Cache growth is exact. Cache length must increase by the chunk length, never by the full prefix length again.
7. Offsets are monotonic. Reject a continuation whose declared position disagrees with the stored length.

Compare floating tensors, not only argmax token IDs. Two wrong logit vectors can still select the same token. Also run tests in evaluation mode when the real model has dropout.

Go deeper

What the cache saves

Without a cache, generation step t repeats projections for all t tokens and recomputes a t × t attention matrix. Summed across a growing output, this repeated full-prefix attention has cubic growth in sequence length under the simple implementation.

With a cache, each layer projects only the new token or chunk. A one-token query still scores all t stored keys, so total attention work over generation remains quadratic. The cache removes repeated prefix work; it does not make attention constant-time.

Memory is the new cost

For ordinary multi-head attention, cache elements per layer are 2 × B × H × L × Dh. The factor 2 counts keys and values. Multiply by layer count and bytes per element for a first memory estimate.

The teaching code concatenates tensors, which copies storage. Production systems commonly preallocate pages or blocks, track write positions, support mixed request lengths, and may use multi-query or grouped-query attention to reduce KV memory. Those optimizations need separate lifecycle and isolation tests.

Limits of this build

This module handles one layer, one shared cache length per batch, absolute positions starting at zero, and trailing masked tokens. It does not implement sliding windows, cache eviction, beam reordering, speculative decoding, quantized caches, cross-attention caches, distributed heads, or request batching with independent offsets.

A numerical match on this toy module proves the local algorithm. It does not demonstrate production speed, memory savings on a particular device, or end-to-end model quality.

Practice and recap

Question: A cache holds positions 0 through 3. A new chunk contains positions 4 and 5. Write the boolean causal rows for its 2 × 6 query-key mask.

References and limits

• Vaswani et al., Attention Is All You Need, for scaled dot-product attention and decoder masking.
• PyTorch scaled dot-product attention documentation, especially its mask convention.

The runnable source favors inspectable arithmetic over kernel efficiency. Benchmark claims require a real model, realistic sequence lengths, warm-up, controlled hardware, and measured memory.