PYTORCH · BUILD IT YOURSELF

KV caches, incremental attention, and correctness

Full build path and setup · Download transformer parts · Download training runner · Download correctness tests

What gets repeated during generation?

A decoder-only transformer predicts one new token after a prefix. A simple implementation runs the entire prefix again after appending each token. Earlier positions have not changed. Causal attention prevents them from depending on the newly added token. Their keys and values can therefore be reused during evaluation.

A key-value cache stores those tensors separately for each layer. It does not store the final answer. It does not replace attention. Cached attention still computes scores between the new query and all allowed stored keys, then mixes their values.

Design the storage before optimizing it

Each layer stores K and V shaped (B,H,L,D). B is batch size. H is head count. L is the cached length. D is head width. A new token supplies tensors shaped (B,H,1,D). Append along the length axis, which is the second-to-last axis.

import torch
old_k = torch.zeros(2, 4, 3, 8)
new_k = torch.ones(2, 4, 1, 8)
all_k = torch.cat((old_k, new_k), dim=-2)
assert all_k.shape == (2, 4, 4, 8)
print(all_k[:, :, -1, :].sum().item())  # 64.0

Do the same for values. Keep each layer's tensors separate because each layer learns different projections. Reset the cache for a new independent prompt. Reusing another request's cache changes predictions and can expose information from that request.

Prefill and decode are different stages

Trace two steps: after prefill with three tokens, each layer caches three key rows and three value rows. Decode one new token and the cache length becomes four. The new query has length one, while its keys cover all four positions.

Prefill processes the prompt together. A four-token prompt uses positions 0 through 3 and fills every layer's cache. Decode receives only the next new input token, uses position 4, and returns updated caches of length 5.

The prompt's final score predicts the first generated token. After selecting that token, feed it into the next decode call. Feeding the entire prompt again alongside an existing cache duplicates the prompt. Shapes may still work, so validate cache length after every call.

Use absolute positions for rectangular attention

Without a cache, query length and key length are equal. With one new query after a prefix of four, the attention score matrix has one row and five columns. That row must be allowed to see all five columns. A naive one-row lower-triangular mask would allow only column zero.

past, new = 4, 2
query_positions = torch.arange(past, past + new)
key_positions = torch.arange(past + new)
allowed = key_positions[None, :] <= query_positions[:, None]
print(allowed.tolist())
# [[True, True, True, True, True, False],
#  [True, True, True, True, True, True]]

The first query is position 4. It cannot see new position 5. The second query can see both new tokens. This is why chunked-cache tests matter. Testing only one token at a time can miss a faulty within-chunk mask.

Preserve rotary positions exactly

Rotate new queries and keys using the cached length as their position offset. Store already-rotated keys. Do not rotate the stored prefix again. If a new query restarts at zero, it asks the wrong relative-position question even though the cache holds correct values.

This teaching cache starts positions at zero. Every layer must have a cache for the same prefix, or every layer must start empty. The code rejects mixed or inconsistent states. It uses one shared sequence length per batch and ordinary multi-head attention. It is not a production cache for arbitrary position origins, mixed-length requests, sliding windows, grouped-query attention, or distributed inference. Those require additional bookkeeping and tests.

Prove equality before measuring speed

Compare cached and uncached logits at each decode step with dropout disabled. If they differ, speed measurements are premature. Check absolute positions, cache concatenation order, mask dimensions, and numerical tolerance in that order.

.venv-learning/bin/python -m unittest discover -s tests -p 'test_transformer_builds.py'
.venv-learning/bin/python examples/pytorch/build_transformer.py decoder-only --steps 500

Put the model in evaluation mode. Disable gradient recording. Calculate scores for a complete sequence. Then calculate the same positions using a prompt plus incremental tokens. Compare the tensors with a numerical tolerance. Greedy output equality alone is insufficient.

Also test chunks of two or more tokens, multiple layers, and batch size greater than one. Assert that cache length increases by the number of new tokens. Test a fresh empty cache for a second prompt. These checks catch stale state, layer sharing, and position-offset mistakes.

Calculate the memory cost

For ordinary multi-head attention, approximate cache bytes equal 2 × layers × B × H × L × D × bytes_per_value. The factor 2 counts keys and values. With two layers, batch 1, four heads, length 100, head width 8, and float32, the result is 51,200 bytes, or 50 KiB. This excludes model weights and temporary attention buffers.

A cache saves repeated prefix projections and other prefix calculations. Each new query still attends across the stored context, so attention work grows with L. Concatenating tensors also copies storage. This simple version favors clarity over allocation efficiency. A faster implementation preallocates buffers and tracks a write position.

A useful failure drill

Do not claim a speedup from theory alone. Benchmark identical models, generated lengths, and devices. Warm up the runtime and report memory use. Tiny CPU models can be dominated by Python overhead, so a correct cache can still be slower in a small experiment.

Read the executable source

These source listings are copied from the runnable files by scripts/sync-build-code.mjs. The training runner and module belong in the same folder. Read a section, run the tests, then make one small change.

@dataclass(frozen=True)
class KVCache:
    """Keys/values for one decoder layer, both shaped ``(B, H, T, Dh)``."""

    key: torch.Tensor
    value: torch.Tensor
    key_padding_mask: torch.Tensor  # (B, T), True for a real token

    def __post_init__(self) -> None:
        if self.key.ndim != 4 or self.value.ndim != 4:
            raise ValueError("cached key and value must have shape (B, H, T, Dh)")
        if self.key.shape != self.value.shape:
            raise ValueError("cached key and value shapes must match")
        expected = (self.key.shape[0], self.key.shape[2])
        _check_valid_mask(self.key_padding_mask, expected, "cached key_padding_mask")

    @property
    def length(self) -> int:
        return self.key.shape[2]

class DecoderOnlyBlock(nn.Module):
    """Pre-norm causal block that can append to a per-layer KV cache."""

    def __init__(self, d_model: int, num_heads: int, hidden_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.attention_norm = ManualLayerNorm(d_model)
        self.self_attention = ManualMultiHeadAttention(
            d_model, num_heads, dropout=dropout, use_rope=True
        )
        self.ffn_norm = ManualLayerNorm(d_model)
        self.feed_forward = FeedForward(d_model, hidden_dim, dropout)

    def forward(
        self,
        x: torch.Tensor,
        valid_tokens: torch.Tensor,
        *,
        cache: KVCache | None = None,
        use_cache: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, KVCache]:
        attention_result = self.self_attention(
            self.attention_norm(x),
            key_padding_mask=valid_tokens,
            query_padding_mask=valid_tokens,
            causal=True,
            cache=cache,
            use_cache=use_cache,
        )
        if use_cache:
            attention_output, new_cache = attention_result
        else:
            attention_output = attention_result
        x = x + attention_output
        x = x + self.feed_forward(self.ffn_norm(x))
        x = x * valid_tokens.unsqueeze(-1).to(x.dtype)
        if use_cache:
            return x, new_cache
        return x

class DecoderOnlyTransformer(nn.Module):
    """Decoder-only language model with one :class:`KVCache` per layer."""

    def __init__(
        self,
        vocab_size: int,
        *,
        d_model: int = 64,
        num_heads: int = 4,
        hidden_dim: int = 128,
        num_layers: int = 2,
        dropout: float = 0.0,
        padding_id: int = 0,
    ) -> None:
        super().__init__()
        self.padding_id = padding_id
        self.d_model = d_model
        self.token_embedding = nn.Embedding(vocab_size, d_model, padding_idx=padding_id)
        self.layers = nn.ModuleList(
            [DecoderOnlyBlock(d_model, num_heads, hidden_dim, dropout) for _ in range(num_layers)]
        )
        self.final_norm = ManualLayerNorm(d_model)

    def forward(
        self,
        token_ids: torch.Tensor,
        *,
        valid_tokens: torch.Tensor | None = None,
        caches: Sequence[KVCache | None] | None = None,
        use_cache: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, list[KVCache]]:
        if token_ids.ndim != 2:
            raise ValueError("token_ids must have shape (B, T)")
        if valid_tokens is None:
            valid_tokens = token_ids.ne(self.padding_id)
        _check_valid_mask(valid_tokens, tuple(token_ids.shape), "valid_tokens")
        if not bool(valid_tokens.any(dim=1).all()):
            raise ValueError("each input sequence needs at least one valid token")
        if caches is None:
            caches = [None] * len(self.layers)
        if len(caches) != len(self.layers):
            raise ValueError("provide exactly one cache entry per layer")
        initialized = [cache is not None for cache in caches]
        if any(initialized) and not all(initialized):
            raise ValueError("layer caches must be either all initialized or all empty")
        lengths = {cache.length for cache in caches if cache is not None}
        if len(lengths) > 1:
            raise ValueError("all layer caches must have the same sequence length")

        x = self.token_embedding(token_ids) * math.sqrt(self.d_model)
        new_caches: list[KVCache] = []
        for layer, cache in zip(self.layers, caches):
            result = layer(x, valid_tokens, cache=cache, use_cache=use_cache)
            if use_cache:
                x, new_cache = result
                new_caches.append(new_cache)
            else:
                x = result
        x = self.final_norm(x) * valid_tokens.unsqueeze(-1).to(x.dtype)
        # Tied output weights avoid a separate, hidden language-model projection.
        logits = F.linear(x, self.token_embedding.weight)
        if use_cache:
            return logits, new_caches
        return logits

    @torch.no_grad()
    def generate_cached(self, prefix_ids: torch.Tensor, max_new_tokens: int) -> torch.Tensor:
        """Prefill once, then decode one token at a time using the KV caches."""
        self.eval()
        generated = prefix_ids
        logits, caches = self(prefix_ids, use_cache=True)
        for _ in range(max_new_tokens):
            next_id = logits[:, -1].argmax(dim=-1, keepdim=True)
            generated = torch.cat((generated, next_id), dim=1)
            logits, caches = self(next_id, caches=caches, use_cache=True)
        return generated


__all__ = [
    "DecoderBlock",
    "DecoderOnlyBlock",
    "DecoderOnlyTransformer",
    "EncoderBlock",
    "FeedForward",
    "KVCache",
    "ManualLayerNorm",
    "ManualMultiHeadAttention",
    "RotaryEmbedding",
    "Seq2SeqTransformer",
    "TransformerDecoder",
    "TransformerEncoder",
    "apply_rope",
    "rope_sin_cos",
    "rotate_pairs",
]

Further reading

Primary references: Attention Is All You Need, RoFormer, and PyTorch attention API and mask conventions. This implementation is a small teaching model, not a reproduction of a pretrained model.