model/kv_cache.py

# model/kv_cache.py
#
# KV-cache for PyCraft-1 incremental decoding.
#
# Without a cache, generating token N re-runs the full forward pass over all
# N-1 previous tokens — quadratic work for a linear amount of output.
# This module stores each layer's keys and values so every decode step only
# computes the single newest position.
#
# Two design choices matter and are easy to get wrong:
#
#   1. Storage is PRE-expansion — (B, n_kv_heads, ...) not (B, n_heads, ...).
#      With GQA 8Q/2KV that is 4x less memory, which is the entire point of
#      grouped-query attention. _repeat_kv re-expands on read.
#
#   2. Memory is PREALLOCATED, never torch.cat'd. Concatenating per step costs
#      O(N) copy per step and O(N^2) overall — roughly 4.3 GB of pointless
#      memcpy across one full-context generation.
#
# The cache holds POST-RoPE keys. RoPE is relative (q_m . k_n depends only on
# m - n), so a key rotated once at its true absolute position stays valid
# forever. Never re-apply RoPE to cached keys.

import torch


# ------------------------------------------------------------------ #
# KV cache
# ------------------------------------------------------------------ #
class KVCache:
    """
    Pre-allocated per-layer key/value store for incremental decoding.

    Per layer:   k[i], v[i] : (batch, n_kv_heads, max_len, head_dim)
    Live region: [:, :, :seq_len, :]

    Typical use:

        cache = KVCache.from_model(model, batch_size=1)
        logits, _ = model(prompt_ids, past_key_values=cache)   # prefill
        logits, _ = model(next_id,    past_key_values=cache)   # decode
    """

    def __init__(
        self,
        n_layers: int,
        batch_size: int,
        n_kv_heads: int,
        head_dim: int,
        max_len: int,
        device,
        dtype: torch.dtype = torch.float32,
    ):
        self.n_layers = n_layers
        self.batch_size = batch_size
        self.n_kv_heads = n_kv_heads
        self.head_dim = head_dim
        self.max_len = max_len
        self.device = device
        self.dtype = dtype
        self.seq_len = 0

        shape = (batch_size, n_kv_heads, max_len, head_dim)
        self.k = [torch.zeros(shape, device=device, dtype=dtype)
                  for _ in range(n_layers)]
        self.v = [torch.zeros(shape, device=device, dtype=dtype)
                  for _ in range(n_layers)]

    # -------------------------------------------------------------- #
    # Called once per layer, per forward pass
    # -------------------------------------------------------------- #
    def update(
        self,
        layer_idx: int,
        k_new: torch.Tensor,   # (batch, n_kv_heads, T, head_dim)
        v_new: torch.Tensor,   # (batch, n_kv_heads, T, head_dim)
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Write this layer's new K/V into the cache and return views covering
        everything written so far: (batch, n_kv_heads, seq_len + T, head_dim).

        Deliberately does NOT advance seq_len. The model advances it once,
        after all layers have written. Advancing here would make layer i see
        the offset that only layer i+1 should see — a silent corruption that
        produces plausible-looking garbage.
        """
        batch, _, T, _ = k_new.shape
        start = self.seq_len

        if start + T > self.max_len:
            raise ValueError(
                f"KVCache overflow: {start} + {T} > max_len={self.max_len}. "
                f"Allocate a larger cache or stop generating."
            )
        if batch != self.batch_size:
            raise ValueError(
                f"batch {batch} does not match cache batch {self.batch_size}"
            )
        if k_new.dtype != self.dtype:
            # Silent casting here would quietly degrade precision every step.
            raise TypeError(
                f"cache dtype {self.dtype} != incoming dtype {k_new.dtype}"
            )

        self.k[layer_idx][:, :, start:start + T] = k_new
        self.v[layer_idx][:, :, start:start + T] = v_new

        # Narrowed views — the zero-padded tail is never visible, so reset()
        # needs no memset.
        return (
            self.k[layer_idx][:, :, :start + T],
            self.v[layer_idx][:, :, :start + T],
        )

    def advance(self, n: int):
        """Advance the write head. Called once per forward, by the model."""
        self.seq_len += n

    def reset(self):
        """Reuse this cache for a new sequence. No zeroing required."""
        self.seq_len = 0

    def __len__(self) -> int:
        return self.seq_len

    def __repr__(self) -> str:
        return (
            f"KVCache(layers={self.n_layers}, batch={self.batch_size}, "
            f"kv_heads={self.n_kv_heads}, head_dim={self.head_dim}, "
            f"seq_len={self.seq_len}/{self.max_len}, dtype={self.dtype})"
        )

    @property
    def memory_bytes(self) -> int:
        """Total allocated cache size in bytes (K and V, all layers)."""
        per = (self.batch_size * self.n_kv_heads
               * self.max_len * self.head_dim)
        return 2 * self.n_layers * per * torch.empty(
            (), dtype=self.dtype).element_size()

    # -------------------------------------------------------------- #
    @classmethod
    def from_model(
        cls,
        model,
        batch_size: int = 1,
        max_len: int | None = None,
        device=None,
        dtype: torch.dtype | None = None,
    ) -> "KVCache":
        """Build a cache matching a model's config, device, and dtype."""
        cfg = model.config
        # Read device/dtype from the embedding, not a Linear: dynamic
        # quantization replaces nn.Linear with a module whose .weight is a
        # method. Activations stay fp32 under dynamic quant either way.
        ref = model.token_embedding.weight
        return cls(
            n_layers=cfg.n_layers,
            batch_size=batch_size,
            n_kv_heads=cfg.n_kv_heads,
            head_dim=cfg.head_dim,
            max_len=max_len or cfg.max_seq_len,
            device=device if device is not None else ref.device,
            dtype=dtype if dtype is not None else ref.dtype,
        )


# ------------------------------------------------------------------ #
# Causal mask construction
# ------------------------------------------------------------------ #
def build_attn_mask(
    q_len: int,
    offset: int,
    device,
    padding_mask: torch.Tensor | None = None,
) -> torch.Tensor | None:
    """
    Build a bottom-right-aligned causal mask for cached attention.

    Args:
        padding_mask: optional (batch, offset + q_len) bool marking real
            tokens True and padding False. Required for batched generation,
            where shorter prompts are left-padded to a common length.

    With no padding mask, returns None for the two cases SDPA handles on its
    own:

        offset == 0   caller passes is_causal=True (square, so PyTorch's
                      upper-left alignment is already correct)
        q_len  == 1   the single newest query may attend to every cached
                      key, so causality holds structurally — no mask needed

    Otherwise returns bool "may attend", (1, 1, q_len, offset + q_len) or
    (batch, 1, q_len, offset + q_len) when padding is involved:

        mask[i, j] = (j <= offset + i) and not padding[j]

    WHY THIS EXISTS: PyTorch builds is_causal as torch.ones(L, S).tril() —
    UPPER-LEFT aligned. With L=1, S=N that mask contains exactly one True
    (column 0), so a cached decode step with is_causal=True would attend only
    to the first prompt token, at every layer. No error, no NaN — just output
    that ignores the prompt and collapses into repetition. Hence the explicit
    dispatch rather than a blanket is_causal=True.
    """
    if padding_mask is None and (offset == 0 or q_len == 1):
        return None

    kv_len = offset + q_len
    q_pos = torch.arange(offset, kv_len, device=device).unsqueeze(1)  # (T, 1)
    k_pos = torch.arange(0, kv_len, device=device).unsqueeze(0)       # (1, S)
    mask = (k_pos <= q_pos)[None, None, :, :]                        # (1,1,T,S)

    if padding_mask is not None:
        if padding_mask.shape[-1] != kv_len:
            raise ValueError(
                f"padding_mask covers {padding_mask.shape[-1]} positions but "
                f"the key length is {kv_len} (offset {offset} + q_len {q_len})"
            )
        # (B,1,1,S) broadcast against the causal (1,1,T,S) -> (B,1,T,S)
        mask = mask & padding_mask[:, None, None, :].to(torch.bool)

        # A fully-masked query row makes softmax(all -inf) = NaN, which then
        # spreads through the residual stream. Left-padded rows hit this: a
        # pad token's own row can have every key masked. Letting each query
        # attend to at least its own position costs nothing (those outputs are
        # discarded) and keeps the tensor finite.
        qi = torch.arange(q_len, device=device)
        mask = mask.clone()
        mask[:, :, qi, offset + qi] = True

    return mask


# ------------------------------------------------------------------ #
# Quick self-test
# ------------------------------------------------------------------ #
if __name__ == "__main__":
    print("Testing build_attn_mask...")

    # Fast paths return None
    assert build_attn_mask(4, 0, "cpu") is None, "offset=0 should return None"
    assert build_attn_mask(1, 7, "cpu") is None, "q_len=1 should return None"

    # Explicit mask matches a brute-force construction
    for q_len, offset in [(4, 7), (3, 1), (2, 2)]:
        got = build_attn_mask(q_len, offset, "cpu")
        kv_len = offset + q_len
        assert got.shape == (1, 1, q_len, kv_len), f"bad shape {got.shape}"
        for i in range(q_len):
            for j in range(kv_len):
                expected = (j <= offset + i)
                assert bool(got[0, 0, i, j]) == expected, (
                    f"mask[{i}][{j}] wrong for q_len={q_len}, offset={offset}"
                )
    print("  build_attn_mask: OK")

    print("\nTesting KVCache...")
    cache = KVCache(n_layers=2, batch_size=1, n_kv_heads=2,
                    head_dim=64, max_len=16, device="cpu")
    print(f"  {cache}")
    print(f"  allocated: {cache.memory_bytes / 1024:.1f} KiB")

    # Prefill 4 positions across both layers
    k = torch.randn(1, 2, 4, 64)
    v = torch.randn(1, 2, 4, 64)
    for layer in range(2):
        kk, vv = cache.update(layer, k, v)
        assert kk.shape == (1, 2, 4, 64), f"bad view shape {kk.shape}"
    assert cache.seq_len == 0, "update() must not advance seq_len"
    cache.advance(4)
    assert len(cache) == 4

    # One decode step
    k1 = torch.randn(1, 2, 1, 64)
    v1 = torch.randn(1, 2, 1, 64)
    kk, vv = cache.update(0, k1, v1)
    assert kk.shape == (1, 2, 5, 64), f"bad decode view {kk.shape}"
    assert torch.equal(kk[:, :, :4], k), "prefill K was corrupted"
    assert torch.equal(kk[:, :, 4:], k1), "decode K not written"
    cache.advance(1)
    print("  writes and views: OK")

    # Guards
    for bad, exc, label in [
        (lambda: cache.update(0, torch.randn(1, 2, 99, 64),
                              torch.randn(1, 2, 99, 64)), ValueError, "overflow"),
        (lambda: cache.update(0, torch.randn(1, 2, 1, 64).half(),
                              torch.randn(1, 2, 1, 64).half()), TypeError, "dtype"),
        (lambda: cache.update(0, torch.randn(3, 2, 1, 64),
                              torch.randn(3, 2, 1, 64)), ValueError, "batch"),
    ]:
        try:
            bad()
            raise AssertionError(f"{label} guard did not fire")
        except exc:
            pass
    print("  guards: OK")

    cache.reset()
    assert len(cache) == 0
    print("  reset: OK")
    print("\nAll kv_cache tests PASSED.")