A KV cache that actually generates

Look at what happens at the third generated token. The old generator projects keys and values for tokens 1, 2 and 3 all over again. The cached one only projects token 3, and then its new query reads the saved keys for 1 and 2. Step through the diagram and turn the cache off, you'll see all the repeated cells.

Save the past keys, compute only the new pair

Day 2's GQA post worked out the memory cost of these blue cells before any code stored them. Day 4's sampling post used the orange path, recompute everything every step. Day 5 finally builds the blue path, the one that reuses.

The first pass over a prompt computes a query, key and value for every position. And when you append one new token, the past tokens' keys and values don't change. So for $t$ generated positions after a one-token prompt, a simplified count of key and value projections is $1 + 2 + \cdots + t = t(t+1)/2$ without a cache, and just $t$ with one. But attention still reads all the past keys at each step, so the cache doesn't make decoding constant-time. GPU launches, memory traffic and the final vocabulary projection still cost time. The diagram shows saved work, not a measured speed ratio.

One pair of tensors per layer

Decoder.forward_cached takes a new chunk of tokens and an optional list of past (key, value) tensors. Each layer projects the new chunk, appends it to that layer's pair, runs attention against the full pair and returns the updated cache. generate_cached fills the cache from the prompt first, then sends one token per step. The original forward and generate are still there as the baseline to compare against.

@torch.inference_mode()
def forward_cached(
    self, token_ids: Tensor, cache: list[tuple[Tensor, Tensor]] | None = None
) -> tuple[Tensor, list[tuple[Tensor, Tensor]]]:
    if self.training:
        raise ValueError("cached forward requires eval mode")
    if (
        self.config.attention != "sdpa"
        or self.config.attention_window
        or self.config.kv_compress_block
    ):
        raise ValueError("cached forward requires full SDPA attention")
    if token_ids.ndim != 2 or token_ids.shape[1] < 1:
        raise ValueError("token IDs must have shape [batch, nonempty sequence]")
    if cache is not None and len(cache) != len(self.blocks):
        raise ValueError("cache must have one entry per layer")
    past = 0 if cache is None else cache[0][0].shape[2]
    if past + token_ids.shape[1] > self.config.context_length:
        raise ValueError("cached sequence exceeds context_length")
    rope = None
    if self.config.position == "rope":
        cosine, sine = rope_tables(
            past + token_ids.shape[1],
            self.config.head_size,
            self.config.rope_scale,
            token_ids.device,
            self.config.rope_base,
        )
        rope = cosine[past:], sine[past:]
    x = self.dropout(self._embed(token_ids, past))
    updated = []
    for index, block in enumerate(self.blocks):
        x, layer_cache = block.forward_cached(x, rope, None if cache is None else cache[index])
        updated.append(layer_cache)
    return self.lm_head(self.final_norm(x)), updated

RoPE rotates a key by its absolute position before the key gets stored. So on the next step the old keys have to keep their old angles, if you rotate the whole cache again the attention scores change. The cached call slices the position table starting from the number of tokens already cached, so a fresh query and key get the right offsets.

Then there's a masking trap which is less obvious. In a normal full pass the query and key lengths match and a triangular causal mask is enough. With a cache, a three-token update can have three queries and hundreds of keys. And PyTorch's is_causal=True on a non-square attention matrix uses upper-left alignment, which would make the new queries attend to the wrong prefix. So octlm builds an explicit mask from absolute row positions for multi-token updates. A one-token update doesn't need a mask at all, because every saved key is in its past.

def forward_cached(
    self,
    x: Tensor,
    rope: tuple[Tensor, Tensor] | None,
    cache: tuple[Tensor, Tensor] | None,
) -> tuple[Tensor, tuple[Tensor, Tensor]]:
    batch, length, _ = x.shape
    query, key, value = self._project(x, rope)
    past = 0 if cache is None else cache[0].shape[2]
    if cache is not None:
        key = torch.cat((cache[0], key), dim=2)
        value = torch.cat((cache[1], value), dim=2)
    mask = None
    if past and length > 1:
        rows = torch.arange(past, past + length, device=x.device).unsqueeze(1)
        columns = torch.arange(past + length, device=x.device).unsqueeze(0)
        mask = columns <= rows
    attended = F.scaled_dot_product_attention(
        query,
        key,
        value,
        attn_mask=mask,
        is_causal=not past and length > 1,
        enable_gqa=key.shape[1] != self.n_heads,
    )
    attended = attended.transpose(1, 2).contiguous().view(batch, length, -1)
    return self.output(attended), (key, value)

The cache is only for full SDPA attention, on purpose. A sliding-window or compressed-attention cache needs different eviction or pooling rules, so the function just rejects those configs. When generated text fills the trained 512-token context, generate_cached rebuilds the cache from the newest 512 tokens. That keeps the original generator's context behavior, but we lose the reuse at that boundary. The planned 64, 256 and 512-token measurements all stay inside the window anyway.

GQA is what makes this cache small enough

For one sequence the cache stores two tensors per layer, so its payload is

where $L$ is layers, $H_{kv}$ is key-value heads, $d_{head}$ is the width per head, $T$ is cached positions and $b$ is bytes per element. The Day 4 model has 8 layers, 2 KV heads, 64 values per head and float16 values, so that's 4 KiB per cached position, which means 256 KiB at 64, 1 MiB at 256 and 2 MiB at 512. That's tensor payload, before allocator overhead. With eight KV heads every one of these numbers would be four times bigger.

The benchmark adds up the actual key and value tensor sizes and asserts that they equal this formula. All three measured sizes passed.

Match the old tokens first, then time both paths

The matched control is Day 4's trained 26M modern checkpoint, temperature 0, batch size 1 and a BOS-only prompt. After warming up both paths, octlm.day5 cache times 64, 256 and 512 new tokens on the same device and dtype. It synchronizes CUDA around each timer and records naive and cached tokens per second, cache bytes and whether the greedy tokens are exactly equal. If even one token changes, the speed result doesn't count until parity is fixed.

def stage_cache(args: argparse.Namespace) -> None:
    if args.out.exists():
        raise FileExistsError(f"{args.out} already contains a measurement")
    model, tokenizer = load_model(args)
    if model.config.attention != "sdpa" or any(
        length > model.config.context_length for length in args.lengths
    ):
        raise ValueError("cache benchmark requires SDPA and lengths within context")
    device = next(model.parameters()).device
    if device.type == "cuda":
        model.half()
    prompt = torch.tensor([[tokenizer.bos_id]], device=device)
    model.generate(prompt, 2)
    model.generate_cached(prompt, 2)
    for length in args.lengths:
        naive, naive_seconds = timed_generation(model, prompt, length, False)
        cached, cached_seconds = timed_generation(model, prompt, length, True)
        _, cache = model.forward_cached(cached[:, :-1])
        cache_bytes = sum(
            (key.numel() + value.numel()) * key.element_size() for key, value in cache
        )
        expected = kv_cache_bytes(model.config, length, next(model.parameters()).element_size())
        if cache_bytes != expected:
            raise RuntimeError("actual cache bytes differ from the model formula")
        row = {
            "cache_bytes": cache_bytes,
            "cached_tokens_per_second": length / cached_seconds,
            "device": str(device),
            "dtype": str(next(model.parameters()).dtype),
            "greedy_equal": bool(torch.equal(naive, cached)),
            "length": length,
            "naive_tokens_per_second": length / naive_seconds,
            "pytorch": torch.__version__,
            "type": "cache_benchmark",
        }
        _write(args.out, row)
        if not row["greedy_equal"]:
            raise RuntimeError(f"cached generation changed greedy output at length {length}")

The small-model CPU check matched all 12 generated tokens across an eight-token context boundary. A separate cached three-token call matched the full forward logits, which exercises the offset mask. And then the trained-checkpoint test matched every greedy token at all three lengths.

So the cached path was 10.6 to 15.8 percent faster in these single T4 timings. That's way smaller than the projection-count reduction would make you expect, because attention still reads every saved key, and the PyTorch calls and the vocabulary projection still take time. And one timing per length can't give us a stable speed distribution, so don't read too much into the exact numbers. Decision, keep the cached generator for full SDPA attention, and the generic generator stays for the other attention configs.