perf: express geometric residual as one addcmul operator

README.md

@@ -61,6 +61,23 @@ print(logits.shape, float(loss))
 `hidden_size` must equal `num_attention_heads * head_dim`. Inputs longer than
 `block_size` fail explicitly. `crop_block_size(n)` only decreases that limit.
 
+## Allocation-aware update
+
+The default geometric update uses `torch.addcmul(x, k, delta)` rather than
+materializing `delta * k` as a Python-level intermediate before addition. An
+explicit reference implementation remains in `model.geometric`, with output and
+backward-gradient parity covered by tests.
+
+Measure the two paths on the target hardware:
+
+```bash
+python benchmarks/geometric_update.py --device cpu --iterations 200
+```
+
+On CUDA the script also reports peak allocated memory. The measurable acceptance
+target is equal outputs/gradients and no increase in peak allocation; this
+repository does not claim a hardware-independent latency improvement.
+
 ## Research status and benchmark policy
 
 The spectral statement above is an exact fixed-`k`, fixed-`beta` property.

benchmarks/geometric_update.py

@@ -0,0 +1,65 @@
+"""Profile literal and allocation-aware geometric residual updates."""
+
+import argparse
+import time
+
+import torch
+
+from model.geometric import geometric_update, reference_geometric_update
+
+
+def synchronize(device):
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+
+
+def profile(function, values, iterations):
+    device = values[0].device
+    with torch.inference_mode():
+        for _ in range(20):
+            function(*values)
+    synchronize(device)
+    if device.type == "cuda":
+        torch.cuda.reset_peak_memory_stats(device)
+
+    started = time.perf_counter()
+    with torch.inference_mode():
+        for _ in range(iterations):
+            function(*values)
+    synchronize(device)
+    latency = (time.perf_counter() - started) * 1000 / iterations
+    peak = (
+        torch.cuda.max_memory_allocated(device) / 1024**2
+        if device.type == "cuda"
+        else None
+    )
+    return latency, peak
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument("--sequence-length", type=int, default=1024)
+    parser.add_argument("--hidden-size", type=int, default=768)
+    parser.add_argument("--iterations", type=int, default=200)
+    args = parser.parse_args()
+
+    device = torch.device(args.device)
+    shape = (args.batch_size, args.sequence_length, args.hidden_size)
+    values = (
+        torch.randn(shape, device=device),
+        torch.randn(shape, device=device),
+        torch.randn(shape[:-1] + (1,), device=device),
+    )
+    for name, function in (
+        ("literal", reference_geometric_update),
+        ("addcmul", geometric_update),
+    ):
+        latency, peak = profile(function, values, args.iterations)
+        peak_text = f", peak_allocated={peak:.1f} MiB" if peak is not None else ""
+        print(f"{name}: {latency:.3f} ms/iteration{peak_text}")
+
+
+if __name__ == "__main__":
+    main()

model/geometric.py

@@ -0,0 +1,19 @@
+"""Auditable implementations of the OUROBOROS elementwise update."""
+
+import torch
+
+
+def reference_geometric_update(
+    x: torch.Tensor, k: torch.Tensor, delta: torch.Tensor
+) -> torch.Tensor:
+    """Literal reference: x + delta * k."""
+
+    return x + delta * k
+
+
+def geometric_update(
+    x: torch.Tensor, k: torch.Tensor, delta: torch.Tensor
+) -> torch.Tensor:
+    """Use addcmul so the update is expressed as one elementwise operator."""
+
+    return torch.addcmul(x, k, delta)

model/ouroboros.py

@@ -12,6 +12,7 @@
 from .init_utils import init_gpt_weights
 from .pydantic_config import validate_pretrained_config_kwargs
 from .rotary import Rotary, apply_rotary_emb
+from .geometric import geometric_update
 
 
 def _logit(p: float) -> float:
@@ -163,8 +164,7 @@ def forward(self, x: torch.Tensor, *, k_in: torch.Tensor, context: torch.Tensor)
                 v = torch.sigmoid(v) * self.v_sigmoid_scale
 
         delta = (beta * (v - proj)).to(dtype=x.dtype)  # (B, T, 1)
-        update = delta * k
-        return x + update
+        return geometric_update(x, k, delta)
 
 
 class OuroborosBlock(nn.Module):

tests/test_geometric.py

@@ -0,0 +1,25 @@
+import torch
+
+from model.geometric import geometric_update, reference_geometric_update
+
+
+def run_and_grad(function, values):
+    x, k, delta = [value.detach().clone().requires_grad_(True) for value in values]
+    output = function(x, k, delta)
+    gradients = torch.autograd.grad(output.square().mean(), (x, k, delta))
+    return output, gradients
+
+
+def test_addcmul_matches_reference_output_and_gradients():
+    torch.manual_seed(11)
+    values = (
+        torch.randn(2, 7, 16),
+        torch.randn(2, 7, 16),
+        torch.randn(2, 7, 1),
+    )
+    reference_output, reference_gradients = run_and_grad(reference_geometric_update, values)
+    optimized_output, optimized_gradients = run_and_grad(geometric_update, values)
+
+    torch.testing.assert_close(optimized_output, reference_output)
+    for optimized, reference in zip(optimized_gradients, reference_gradients):
+        torch.testing.assert_close(optimized, reference)