feat: add exact autoregressive KV caching

README.md

@@ -23,6 +23,8 @@ accuracy, forgetting, or gradient stability.
 - causal scaled dot-product attention, RoPE, QK RMSNorm, and SwiGLU
 - bounded per-token beta gates and scalar write targets
 - optional causal one-token key/value shift reference
+- exact dynamic KV caching for autoregressive decoding, including shifted K/V
+- deterministic or sampled token generation with cached/reference parity tests
 - tied token embedding / language-model head
 - import, forward/backward, configuration, and block-size regression tests
 
@@ -111,9 +113,28 @@ backed by an additive-residual control with matched parameters/FLOPs, identical
 data order, multiple seeds, learning curves, downstream evaluation, and
 activation/gradient diagnostics.
 
-The optional `ShiftLinear` is a zero-padded, one-token causal reference. Its
-streaming-cache interface is intentionally unsupported until cache correctness is
-tested.
+## Exact autoregressive cache
+
+`generate_tokens` offers a dynamic per-layer KV cache and a slower full-prefix
+reference path. RoPE positions, causal masking, and optional `ShiftLinear`
+state are carried across decode calls. Cache use is inference-only and fails
+before exceeding `block_size`.
+
+```python
+model.eval()
+generated = model.generate_tokens(tokens[:, :8], max_new_tokens=16, use_cache=True)
+```
+
+Cached token-by-token logits are tested against full-sequence logits with and
+without shifted keys/values. Measure the benefit on target hardware:
+
+```bash
+python benchmarks/kv_cache_generation.py --device cpu
+```
+
+The interface follows the dynamic-cache pattern used by modern inference
+runtimes while remaining an original, repository-local implementation. It does
+not claim a hardware-independent speed-up.
 
 ## License
 

benchmarks/kv_cache_generation.py

@@ -0,0 +1,53 @@
+"""Compare exact cached and full-prefix autoregressive decoding."""
+
+import argparse
+import time
+
+import torch
+
+from model import OuroborosConfig, OuroborosModel
+
+
+def synchronize(device: torch.device) -> None:
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+
+
+def profile(model, prompt, new_tokens, use_cache, iterations):
+    for _ in range(2):
+        model.generate_tokens(prompt, new_tokens, use_cache=use_cache)
+    synchronize(prompt.device)
+    started = time.perf_counter()
+    for _ in range(iterations):
+        model.generate_tokens(prompt, new_tokens, use_cache=use_cache)
+    synchronize(prompt.device)
+    return (time.perf_counter() - started) * 1000 / iterations
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    parser.add_argument("--prompt-length", type=int, default=64)
+    parser.add_argument("--new-tokens", type=int, default=32)
+    parser.add_argument("--iterations", type=int, default=10)
+    args = parser.parse_args()
+
+    device = torch.device(args.device)
+    config = OuroborosConfig(
+        vocab_size=256,
+        num_hidden_layers=4,
+        num_attention_heads=4,
+        hidden_size=128,
+        head_dim=32,
+        block_size=args.prompt_length + args.new_tokens,
+    )
+    model = OuroborosModel(config).to(device).eval()
+    prompt = torch.randint(0, config.vocab_size, (1, args.prompt_length), device=device)
+    for use_cache in (False, True):
+        latency = profile(model, prompt, args.new_tokens, use_cache, args.iterations)
+        label = "dynamic-cache" if use_cache else "full-prefix"
+        print(f"{label}: {latency:.3f} ms/generation")
+
+
+if __name__ == "__main__":
+    main()

model/__init__.py

@@ -1,5 +1,12 @@
 """Public OUROBOROS model API."""
 
+from .cache import OuroborosCache, OuroborosLayerCache
 from .ouroboros import OuroborosConfig, OuroborosModel, OuroborosResidual
 
-__all__ = ["OuroborosConfig", "OuroborosModel", "OuroborosResidual"]
+__all__ = [
+    "OuroborosCache",
+    "OuroborosConfig",
+    "OuroborosLayerCache",
+    "OuroborosModel",
+    "OuroborosResidual",
+]

model/cache.py

@@ -0,0 +1,38 @@
+"""Typed inference cache primitives for OUROBOROS autoregressive decoding."""
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass(frozen=True)
+class OuroborosLayerCache:
+    """Attention and optional shift state for one decoder layer.
+
+    Keys and values use ``(batch, heads, sequence, head_dim)``. Shift states
+    retain the final pre-projection token as ``(batch, 1, hidden_size)``.
+    """
+
+    keys: torch.Tensor
+    values: torch.Tensor
+    key_shift_state: torch.Tensor | None = None
+    value_shift_state: torch.Tensor | None = None
+
+    def __post_init__(self) -> None:
+        if self.keys.ndim != 4 or self.values.ndim != 4:
+            raise ValueError("cached keys and values must be rank-four tensors")
+        if self.keys.shape != self.values.shape:
+            raise ValueError("cached keys and values must have identical shapes")
+        for name, state in (
+            ("key_shift_state", self.key_shift_state),
+            ("value_shift_state", self.value_shift_state),
+        ):
+            if state is not None and (state.ndim != 3 or state.shape[1] != 1):
+                raise ValueError(f"{name} must have shape (batch, 1, hidden_size)")
+
+    @property
+    def sequence_length(self) -> int:
+        return int(self.keys.shape[-2])
+
+
+OuroborosCache = tuple[OuroborosLayerCache, ...]

model/kv_shift.py

@@ -5,8 +5,8 @@
 class ShiftLinear(nn.Linear):
     """Causal one-token shift followed by a linear projection.
 
-    The optional cache argument is reserved for a future streaming cache. Passing
-    a non-None cache fails explicitly rather than silently producing wrong states.
+    Streaming inference may pass the previous final input token as ``cache``.
+    This preserves exact one-token-shift semantics across decode calls.
     """
 
     def __init__(self, in_features, out_features, num_heads, bias=False):
@@ -15,11 +15,28 @@ def __init__(self, in_features, out_features, num_heads, bias=False):
             raise ValueError("out_features must be divisible by num_heads")
         self.num_heads = int(num_heads)
 
-    def forward(self, x: torch.Tensor, cache=None) -> torch.Tensor:
-        if cache is not None:
-            raise NotImplementedError("streaming ShiftLinear cache is not implemented")
+    def forward(
+        self,
+        x: torch.Tensor,
+        cache: torch.Tensor | None = None,
+        *,
+        use_cache: bool = False,
+    ):
         if x.ndim != 3:
             raise ValueError("ShiftLinear expects (batch, sequence, features)")
-        shifted = torch.zeros_like(x)
-        shifted[:, 1:] = x[:, :-1]
-        return super().forward(shifted)
+        if x.shape[1] == 0:
+            raise ValueError("ShiftLinear requires at least one sequence element")
+        if cache is not None:
+            expected = (x.shape[0], 1, x.shape[2])
+            if tuple(cache.shape) != expected:
+                raise ValueError(f"ShiftLinear cache must have shape {expected}")
+            if cache.device != x.device or cache.dtype != x.dtype:
+                raise ValueError("ShiftLinear cache must match input device and dtype")
+            shifted = torch.cat((cache, x[:, :-1]), dim=1)
+        else:
+            shifted = torch.zeros_like(x)
+            shifted[:, 1:] = x[:, :-1]
+        output = super().forward(shifted)
+        if not use_cache:
+            return output
+        return output, x[:, -1:].detach()