opt: cumsum/chunked attention kernels, memory-flat CE, block checkpointing, v2 trainer (proven equivalent, 46 tests)

fractus/nn/attention.py

@@ -12,12 +12,30 @@ Math (Katharopoulos 2020, normalized causal form):
 """
 
 import math
+import os
 import torch
 import torch.nn as nn
 
 from .stats import elu_plus_one, stable_softmax
 
 
+_ACTIVE_IMPL = os.environ.get("FRACTUS_ATTN_IMPL", "cumsum")  # cumsum | chunked
+
+
+def set_attention_impl(name: str) -> None:
+    """Select the production linear-attention kernel.
+
+    'cumsum'   — fastest arithmetic (default; CPU-friendly).
+    'chunked'  — memory-flat O(block²+D²) activations (GPU/VRAM-constrained).
+    Both are proven equal to the einsum reference by
+    tests/test_attention_equivalence.py.
+    """
+    global _ACTIVE_IMPL
+    if name not in ("cumsum", "chunked"):
+        raise ValueError(f"unknown attention impl: {name!r}")
+    _ACTIVE_IMPL = name
+
+
 def _mandelbrot_offsets(n_levels: int) -> torch.Tensor:
     """Offsets ω_level = (φ2)^{-level} for level = 0..n_levels-1.
 
@@ -126,17 +144,17 @@ class FractalLinearAttention(nn.Module):
             outputs.append(y_t)
         return torch.stack(outputs, dim=1)  # (B, L, D)
 
-    def _linear_attention_causal_vectorized(
+    def _linear_attention_causal_einsum(
         self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
         carry: tuple = None,
     ) -> torch.Tensor:
-        """Vectorized version of _linear_attention_causal_one_head.
+        """REFERENCE implementation (einsum-with-mask form).
 
-        Same mathematics, but without a Python loop over L. The trick:
-        precompute the cumulative sums S_t and z_t via a lower-triangular
-        convolution, then compute all y_t in parallel.
-
-        Equivalence guaranteed by test_attention_vectorized.py (atol 1e-5).
+        Kept as numerical ground truth for tests/test_attention_equivalence.py.
+        Computes S_t = Σ_{j<=t} outer[j] via a matmul against a triangular
+        mask: O(L^2·D^2) multiply-adds. The production path
+        (_linear_attention_causal_vectorized → cumsum form) computes the same
+        sums with O(L·D^2) adds.
 
         L8 STATE-CARRY: if `carry = (S0, z0)` is provided (each (B, D, D) and
         (B, D)), the running state is INITIALIZED with (S0, z0) instead of
@@ -186,6 +204,136 @@ class FractalLinearAttention(nn.Module):
             return y, (S_final, z_final)
         return y  # (B, L, D)
 
+    def _linear_attention_causal_cumsum(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+        carry: tuple = None,
+    ) -> torch.Tensor:
+        """PRODUCTION implementation: causal linear attention via cumsum.
+
+        Same mathematics as _linear_attention_causal_einsum — proven equal by
+        tests/test_attention_equivalence.py (forward AND gradients, with and
+        without carry). The difference is purely computational:
+
+            S_t = Σ_{j<=t} outer[j]  IS a cumulative sum along the L axis.
+            The mask-matmul form computed that sum with O(L^2·D^2)
+            multiply-adds; cumsum needs O(L·D^2) adds — at L=128 that is
+            ~128x less arithmetic on the dominant term. Backward is a reverse
+            cumsum instead of another full masked matmul.
+
+        Numerical note: accumulation order differs from the matmul form, so
+        results agree within float32 rounding (~1e-4 relative), far below the
+        training signal scale. Checkpoint compatibility is exact: parameter
+        shapes and semantics untouched (open-heart surgery rule).
+        """
+        B, L, D = q.shape
+
+        outer = torch.einsum("btp,btq->btpq", k, v)      # (B, L, D, D)
+        S = torch.cumsum(outer, dim=1)                   # S_t = Σ_{j<=t} k_j⊗v_j
+        z = torch.cumsum(k, dim=1)                       # (B, L, D)
+
+        if carry is not None:
+            S0, z0 = carry
+            S = S + S0.unsqueeze(1)
+            z = z + z0.unsqueeze(1)
+
+        num = torch.einsum("btp,btpq->btq", q, S)        # (B, L, D)
+        denom = (q * z).sum(dim=-1, keepdim=True)        # (B, L, 1)
+        safe = denom.abs() > 1e-10
+        y = torch.where(safe, num / (denom + 1e-20), torch.zeros_like(num))
+
+        if carry is not None:
+            return y, (S[:, -1], z[:, -1])
+        return y
+
+    def _linear_attention_causal_chunked(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+        carry: tuple = None, block: int = 64,
+    ) -> torch.Tensor:
+        """MEMORY-FLAT implementation: blockwise causal linear attention.
+
+        Same inclusive-causal mathematics as the einsum/cumsum forms (proven by
+        tests/test_attention_equivalence.py), but S_t is never materialized
+        for every t. The sequence is processed in blocks of `block` positions:
+
+            intra-block  : y_t += Σ_{j<=t, same block} (φq_t·φk_j) v_j
+                           via an (block×block) masked score matrix
+            inter-block  : y_t += φq_t · S_run / (φq_t·z_run contribution)
+            state update : S_run += Σ_block φk⊗v ; z_run += Σ_block φk
+
+        Peak activation per group drops from O(L·D²) (the outer/S tensors of
+        the cumsum form) to O(block² + D²). At L=128, D=64, block=64 that is
+        a ~32x reduction of the dominant activation — the lever that should
+        re-enable torch.compile and large batches on VRAM-constrained pods.
+
+        FLOP note: intra work is O(L·block·D) instead of O(L·D²); with
+        block=64 < D=64 it is on par. Choose this impl when MEMORY is the
+        binding constraint (GPU), cumsum when raw speed is (CPU).
+
+        `carry` semantics identical: (S0, z0) added to every position's sums;
+        returned state includes it.
+        """
+        B, L, D = q.shape
+        if L % block != 0:
+            # fall back to the exact cumsum path for ragged lengths
+            return self._linear_attention_causal_cumsum(q, k, v, carry=carry)
+
+        dev, dt = q.device, q.dtype
+        if carry is not None:
+            S_run, z_run = carry[0].clone(), carry[1].clone()
+        else:
+            S_run = torch.zeros(B, D, D, dtype=dt, device=dev)
+            z_run = torch.zeros(B, D, dtype=dt, device=dev)
+
+        causal = torch.tril(
+            torch.ones(block, block, dtype=torch.bool, device=dev))
+
+        outs = []
+        for a in range(0, L, block):
+            qb = q[:, a : a + block]          # (G, b, D)
+            kb = k[:, a : a + block]
+            vb = v[:, a : a + block]
+
+            scores = torch.bmm(qb, kb.transpose(1, 2))   # (G, b, b) φq·φk
+
+            num_intra = torch.bmm(scores.masked_fill(~causal, 0.0), vb)
+            den_intra = scores.masked_fill(~causal, 0.0).sum(dim=-1)  # (G, b)
+
+            num_inter = torch.bmm(qb, S_run)             # (G, b, D)
+            den_inter = torch.bmm(qb, z_run.unsqueeze(-1)).squeeze(-1)
+
+            num = num_intra + num_inter
+            den = den_intra + den_inter
+            safe = den.abs() > 1e-10
+            outs.append(torch.where(
+                safe.unsqueeze(-1), num / (den.unsqueeze(-1) + 1e-20),
+                torch.zeros_like(num)))
+
+            # inclusive update AFTER readout: current token joins state for
+            # the NEXT query positions only (matches S_t including token t).
+            S_run = S_run + torch.einsum("btp,btq->bpq", kb, vb)
+            z_run = z_run + kb.sum(dim=1)
+
+        y = torch.cat(outs, dim=1)
+        if carry is not None:
+            return y, (S_run, z_run)
+        return y
+
+    def _linear_attention_causal_vectorized(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+        carry: tuple = None,
+    ) -> torch.Tensor:
+        """Dispatches to the active production kernel (default: cumsum).
+
+        Signature and return contract are unchanged from the original
+        vectorized implementation, so all call sites (CTEBlock.tick_chunk_core,
+        FractalLinearAttention.forward) work without modification. Switch
+        kernels with set_attention_impl('chunked'|'cumsum') or the
+        FRACTUS_ATTN_IMPL environment variable.
+        """
+        if _ACTIVE_IMPL == "chunked":
+            return self._linear_attention_causal_chunked(q, k, v, carry=carry)
+        return self._linear_attention_causal_cumsum(q, k, v, carry=carry)
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """x: (B, L, d_model) → output (B, L, d_model).