diff --git a/README.md b/README.md index 7d53e6f..ca3444f 100644 --- a/README.md +++ b/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 diff --git a/benchmarks/kv_cache_generation.py b/benchmarks/kv_cache_generation.py new file mode 100644 index 0000000..2996112 --- /dev/null +++ b/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() diff --git a/model/__init__.py b/model/__init__.py index 4feefcc..4b27748 100644 --- a/model/__init__.py +++ b/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", +] diff --git a/model/cache.py b/model/cache.py new file mode 100644 index 0000000..735de42 --- /dev/null +++ b/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, ...] diff --git a/model/kv_shift.py b/model/kv_shift.py index 742efd8..2e4323e 100644 --- a/model/kv_shift.py +++ b/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() diff --git a/model/ouroboros.py b/model/ouroboros.py index d451233..9c9bf75 100644 --- a/model/ouroboros.py +++ b/model/ouroboros.py @@ -9,6 +9,7 @@ from .rmsnorm import RMSNorm from .kv_shift import ShiftLinear +from .cache import OuroborosCache, OuroborosLayerCache from .init_utils import init_gpt_weights from .pydantic_config import validate_pretrained_config_kwargs from .rotary import Rotary, apply_rotary_emb @@ -54,30 +55,87 @@ def __init__(self, config): # Apply RMSNorm to each head's output dimension self.subln = RMSNorm(self.head_dim, eps=1e-5, elementwise_affine=True) - def forward(self, x): + def forward( + self, + x: torch.Tensor, + *, + cache: OuroborosLayerCache | None = None, + use_cache: bool = False, + ): B, T, C = x.size() # batch size, sequence length, embedding dimensionality (hidden_size) + past_length = cache.sequence_length if cache is not None else 0 q = self.c_q(x).view(B, T, self.n_head, self.head_dim) if self.use_k_shift: - k = self.c_k(x, None).view(B, T, self.n_head, self.head_dim) + key_result = self.c_k( + x, + None if cache is None else cache.key_shift_state, + use_cache=use_cache, + ) + if use_cache: + k, key_shift_state = key_result + else: + k, key_shift_state = key_result, None + k = k.view(B, T, self.n_head, self.head_dim) else: k = self.c_k(x).view(B, T, self.n_head, self.head_dim) + key_shift_state = None if self.use_v_shift: - v = self.c_v(x, None).view(B, T, self.n_head, self.head_dim) + value_result = self.c_v( + x, + None if cache is None else cache.value_shift_state, + use_cache=use_cache, + ) + if use_cache: + v, value_shift_state = value_result + else: + v, value_shift_state = value_result, None + v = v.view(B, T, self.n_head, self.head_dim) else: v = self.c_v(x).view(B, T, self.n_head, self.head_dim) - cos, sin = self.rotary(q) + value_shift_state = None + cos, sin = self.rotary(q, position_offset=past_length) q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin) if self.use_qk_rmsnorm: q = self.q_rms(q) k = self.k_rms(k) - y = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=True) + current_keys = k.transpose(1, 2) + current_values = v.transpose(1, 2) + if cache is not None: + if cache.keys.shape[:2] != current_keys.shape[:2] or cache.keys.shape[-1] != self.head_dim: + raise ValueError("attention cache shape does not match this layer") + if cache.keys.device != x.device or cache.keys.dtype != current_keys.dtype: + raise ValueError("attention cache must match input device and projected dtype") + keys = torch.cat((cache.keys, current_keys), dim=-2) + values = torch.cat((cache.values, current_values), dim=-2) + query_positions = torch.arange(past_length, past_length + T, device=x.device) + key_positions = torch.arange(keys.shape[-2], device=x.device) + causal_mask = key_positions.unsqueeze(0) <= query_positions.unsqueeze(1) + y = F.scaled_dot_product_attention( + q.transpose(1, 2), + keys, + values, + attn_mask=causal_mask, + is_causal=False, + ) + else: + keys, values = current_keys, current_values + y = F.scaled_dot_product_attention( + q.transpose(1, 2), keys, values, is_causal=True + ) if self.using_groupnorm: y = self.subln(y) y = y.transpose(1, 2).contiguous().reshape(B, T, self.n_head * self.head_dim) y = self.c_proj(y) - return y + if not use_cache: + return y + return y, OuroborosLayerCache( + keys=keys.detach(), + values=values.detach(), + key_shift_state=key_shift_state, + value_shift_state=value_shift_state, + ) class MLP(nn.Module): @@ -211,9 +269,19 @@ def __init__(self, config): ) self.last_diagnostics = {} - def forward(self, x): + def forward( + self, + x: torch.Tensor, + *, + cache: OuroborosLayerCache | None = None, + use_cache: bool = False, + ): x_norm = self.ln_1(x) - k_attn = self.attn(x_norm) + attention_result = self.attn(x_norm, cache=cache, use_cache=use_cache) + if use_cache: + k_attn, next_cache = attention_result + else: + k_attn, next_cache = attention_result, None if self.collect_diagnostics: x, attn_diagnostics = self.ouroboros_attn( x, k_in=k_attn, context=x_norm, return_diagnostics=True @@ -233,7 +301,9 @@ def forward(self, x): } else: x = self.ouroboros_mlp(x, k_in=k_mlp, context=x_norm) - return x + if not use_cache: + return x + return x, next_cache @dataclass class OuroborosConfig(PretrainedConfig): @@ -297,16 +367,48 @@ def __init__(self, config): self.ln_f = RMSNorm(config.hidden_size) init_gpt_weights(self, config) - def forward(self, idx, targets=None, return_logits=True, output_all_seq=False): + def forward( + self, + idx, + targets=None, + return_logits=True, + output_all_seq=False, + *, + past_key_values: OuroborosCache | None = None, + use_cache: bool = False, + ): if idx.ndim != 2: raise ValueError("idx must have shape (batch, sequence)") - if idx.size(1) > self.config.block_size: + if idx.size(1) == 0: + raise ValueError("idx must contain at least one token") + if use_cache and self.training: + raise RuntimeError("KV caching is inference-only; call model.eval() first") + if use_cache and targets is not None: + raise ValueError("targets are not supported when use_cache=True") + if past_key_values is not None and not use_cache: + raise ValueError("past_key_values requires use_cache=True") + if past_key_values is not None and len(past_key_values) != len(self.transformer.h): + raise ValueError("past_key_values must contain one entry per decoder layer") + past_length = 0 if past_key_values is None else past_key_values[0].sequence_length + if past_key_values is not None and any( + layer.sequence_length != past_length for layer in past_key_values + ): + raise ValueError("all layer caches must have the same sequence length") + total_length = past_length + idx.size(1) + if total_length > self.config.block_size: raise ValueError( - f"sequence length {idx.size(1)} exceeds block_size {self.config.block_size}" + f"cached sequence length {total_length} exceeds block_size {self.config.block_size}" ) x = self.transformer.wte(idx) # token embeddings of shape (b, t, hidden_size) - for block in self.transformer.h: - x = block(x) + next_key_values = [] + for layer_index, block in enumerate(self.transformer.h): + layer_cache = None if past_key_values is None else past_key_values[layer_index] + block_result = block(x, cache=layer_cache, use_cache=use_cache) + if use_cache: + x, next_cache = block_result + next_key_values.append(next_cache) + else: + x = block_result x = self.ln_f(x) logits_scale = 1.0 @@ -331,7 +433,65 @@ def forward(self, idx, targets=None, return_logits=True, output_all_seq=False): if not return_logits: logits = None - return logits, loss + if not use_cache: + return logits, loss + return logits, loss, tuple(next_key_values) + + @torch.inference_mode() + def generate_tokens( + self, + idx: torch.Tensor, + max_new_tokens: int, + *, + temperature: float = 0.0, + top_k: int | None = None, + use_cache: bool = True, + ) -> torch.Tensor: + """Generate tokens with an exact dynamic KV cache or a reference path.""" + + if max_new_tokens < 0: + raise ValueError("max_new_tokens must be non-negative") + if temperature < 0: + raise ValueError("temperature must be non-negative") + if top_k is not None and top_k <= 0: + raise ValueError("top_k must be positive") + if idx.ndim != 2 or idx.shape[1] == 0: + raise ValueError("idx must have shape (batch, non-empty sequence)") + if idx.shape[1] + max_new_tokens > self.config.block_size: + raise ValueError("prompt plus generated tokens exceeds block_size") + + was_training = self.training + self.eval() + try: + output = idx + cache = None + logits = None + if use_cache: + logits, _, cache = self(output, use_cache=True) + for _ in range(max_new_tokens): + if not use_cache: + logits, _ = self(output) + next_logits = logits[:, -1, :] if logits.shape[1] > 1 else logits[:, 0, :] + if top_k is not None: + values, _ = torch.topk(next_logits, min(top_k, next_logits.shape[-1])) + next_logits = next_logits.masked_fill( + next_logits < values[:, [-1]], float("-inf") + ) + if temperature == 0: + next_token = torch.argmax(next_logits, dim=-1, keepdim=True) + else: + probabilities = F.softmax(next_logits / temperature, dim=-1) + next_token = torch.multinomial(probabilities, num_samples=1) + output = torch.cat((output, next_token), dim=1) + if use_cache: + logits, _, cache = self( + next_token, + past_key_values=cache, + use_cache=True, + ) + return output + finally: + self.train(was_training) def get_residual_diagnostics(self): """Return per-layer metrics from the most recent instrumented forward pass.""" diff --git a/model/rotary.py b/model/rotary.py index 175016e..22e0b5d 100644 --- a/model/rotary.py +++ b/model/rotary.py @@ -25,8 +25,15 @@ def __init__(self, dim: int, base: float = 10000.0, rope_ratio: float = 1.0): inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) self.register_buffer("inv_freq", inv_freq / float(rope_ratio), persistent=False) - def forward(self, x: torch.Tensor): - positions = torch.arange(x.shape[1], device=x.device, dtype=self.inv_freq.dtype) + def forward(self, x: torch.Tensor, *, position_offset: int = 0): + if position_offset < 0: + raise ValueError("position_offset must be non-negative") + positions = torch.arange( + position_offset, + position_offset + x.shape[1], + device=x.device, + dtype=self.inv_freq.dtype, + ) frequencies = torch.outer(positions, self.inv_freq.to(device=x.device)) embedding = torch.cat((frequencies, frequencies), dim=-1) cos = embedding.cos().to(dtype=x.dtype)[None, :, None, :] diff --git a/tests/test_cache.py b/tests/test_cache.py new file mode 100644 index 0000000..20e178e --- /dev/null +++ b/tests/test_cache.py @@ -0,0 +1,74 @@ +import pytest +import torch + +from model import OuroborosConfig, OuroborosModel + + +def tiny_config(**overrides): + values = dict( + vocab_size=32, + num_hidden_layers=2, + num_attention_heads=2, + hidden_size=16, + head_dim=8, + block_size=16, + using_groupnorm=False, + ) + values.update(overrides) + return OuroborosConfig(**values) + + +@pytest.mark.parametrize("use_shift", [False, True]) +def test_incremental_cache_matches_full_sequence_logits(use_shift): + torch.manual_seed(17) + model = OuroborosModel( + tiny_config(use_k_shift=use_shift, use_v_shift=use_shift) + ).eval() + tokens = torch.randint(0, 32, (2, 7)) + + full_logits, _ = model(tokens, output_all_seq=True) + cache = None + decoded = [] + for position in range(tokens.shape[1]): + logits, _, cache = model( + tokens[:, position : position + 1], + past_key_values=cache, + use_cache=True, + ) + decoded.append(logits) + + torch.testing.assert_close(torch.cat(decoded, dim=1), full_logits) + assert len(cache) == model.config.num_hidden_layers + assert all(layer.sequence_length == tokens.shape[1] for layer in cache) + + +def test_cached_and_reference_greedy_generation_match(): + torch.manual_seed(23) + model = OuroborosModel(tiny_config()).eval() + prompt = torch.tensor([[1, 2, 3]]) + + cached = model.generate_tokens(prompt, 5, use_cache=True) + reference = model.generate_tokens(prompt, 5, use_cache=False) + + torch.testing.assert_close(cached, reference) + + +def test_cache_is_inference_only_and_honors_block_size(): + model = OuroborosModel(tiny_config(block_size=4)) + with pytest.raises(RuntimeError, match="inference-only"): + model(torch.ones(1, 1, dtype=torch.long), use_cache=True) + model.eval() + with pytest.raises(ValueError, match="exceeds block_size"): + model.generate_tokens(torch.ones(1, 3, dtype=torch.long), 2) + + +def test_shift_cache_rejects_wrong_shape(): + model = OuroborosModel(tiny_config(use_k_shift=True, use_v_shift=True)).eval() + _, _, cache = model(torch.ones(1, 1, dtype=torch.long), use_cache=True) + with pytest.raises(ValueError, match="key_shift_state"): + type(cache[0])( + keys=cache[0].keys, + values=cache[0].values, + key_shift_state=torch.zeros(1, 2, model.config.hidden_size), + value_shift_state=cache[0].value_shift_state, + )