src/manas_gpt/kv_cache.py

"""Request-local inference cache; overflow preserves crop-and-recompute semantics."""
import torch
from torch.nn import functional as F
from .rotary import RotaryAttention, rotate_pairs


class CacheSession:
    def __init__(self, model):
        if model.training:
            raise ValueError("KV caching requires model.eval(); training stays uncached")
        self.rotary = model.config.position_encoding == "rope"
        if not self.rotary and not hasattr(model, "position_embedding"):
            raise ValueError("A learned-position model requires its position table")
        for block in model.blocks:
            attention = block.attention
            if self.rotary != isinstance(attention, RotaryAttention):
                raise ValueError("Cache position encoding does not match attention modules")
            kv_heads = getattr(attention, "n_kv_head", attention.n_head)
            width = model.config.n_embd + 2 * kv_heads * (model.config.n_embd // attention.n_head)
            if kv_heads < 1 or attention.n_head % kv_heads or attention.qkv.out_features != width:
                raise ValueError("Cache projection/head shapes do not match the model")
        self.model = model
        self.layers = None
        self.length = 0
        self.batch_size = None

    @torch.no_grad()
    def _forward(self, token_ids, past, offset):
        model = self.model
        if model.training:
            raise ValueError("The cached model must remain in evaluation mode")
        batch, time = token_ids.shape
        positions = torch.arange(offset, offset + time, device=token_ids.device)
        x = model.token_embedding(token_ids)
        if not self.rotary:
            x = x + model.position_embedding(positions)[None]
        x = model.embedding_dropout(x)
        layers = []
        for index, block in enumerate(model.blocks):
            attention = block.attention
            dim = model.config.n_embd // attention.n_head
            kv_heads = getattr(attention, "n_kv_head", attention.n_head)
            q, k, v = attention.qkv(block.ln_attention(x)).split(
                (model.config.n_embd, kv_heads * dim, kv_heads * dim), dim=-1)
            q = q.view(batch, time, attention.n_head, dim).transpose(1, 2)
            k, v = (value.view(batch, time, kv_heads, dim).transpose(1, 2) for value in (k, v))
            if self.rotary:
                cos, sin = attention.cos[offset:offset + time], attention.sin[offset:offset + time]
                q, k = rotate_pairs(q, cos, sin), rotate_pairs(k, cos, sin)
            if past is not None:
                k = torch.cat((past[index][0], k), dim=-2)
                v = torch.cat((past[index][1], v), dim=-2)
            else:
                # Do not retain Q's storage through views of the combined projection.
                k = k.clone(memory_format=torch.contiguous_format)
                v = v.clone(memory_format=torch.contiguous_format)
            layers.append((k, v))
            attention_k, attention_v = k, v
            if kv_heads != attention.n_head:
                repeats = attention.n_head // kv_heads
                attention_k = k.repeat_interleave(repeats, dim=1)
                attention_v = v.repeat_interleave(repeats, dim=1)
            # A single decode query has no future keys. A rectangular causal
            # mask aligned at the upper left would wrongly hide its cached past.
            attended = F.scaled_dot_product_attention(
                q, attention_k, attention_v, dropout_p=0.0, is_causal=past is None)
            attended = attended.transpose(1, 2).contiguous().view(batch, time, model.config.n_embd)
            x = x + attention.output_dropout(attention.output(attended))
            x = x + block.ffn(block.ln_ffn(x))
        return model.lm_head(model.final_norm(x)[:, -1:, :]), layers

    def prefill(self, token_ids):
        if token_ids.ndim != 2 or not 0 < token_ids.shape[1] <= self.model.config.block_size:
            raise ValueError("Prefill requires a nonempty (B,T) prompt within the context window")
        self.layers, self.length, self.batch_size = None, 0, None
        logits, layers = self._forward(token_ids, None, 0)
        self.layers, self.length, self.batch_size = layers, token_ids.shape[1], token_ids.shape[0]
        return logits

    def decode(self, token_ids):
        if self.layers is None:
            raise ValueError("Call prefill before decode")
        if token_ids.shape != (self.batch_size, 1):
            raise ValueError("Decode accepts exactly one new token per existing batch row")
        if self.length >= self.model.config.block_size:
            raise ValueError("Full cache: explicitly prefill the cropped window instead")
        logits, layers = self._forward(token_ids, self.layers, self.length)
        self.layers = layers
        self.length += 1
        return logits

    def storage(self):
        tensors = [] if self.layers is None else [value for pair in self.layers for value in pair]
        unique = {value.untyped_storage().data_ptr(): value.untyped_storage().nbytes() for value in tensors}
        return {"positions": self.length,
                "logical_bytes": sum(value.numel() * value.element_size() for value in tensors),
                "storage_bytes": sum(unique.values()),
                "dtype": str(tensors[0].dtype) if tensors else None,
                "allocation": "dynamic concatenation; no preallocated unused capacity"}