mirror attention.py from geolip 0.7.0 @a63a29b (MIRROR law; n_const=1 bit-identical)

attention.py

@@ -61,17 +61,51 @@ class CausalSDPA(nn.Module):
         return self.o(y.transpose(1, 2).reshape(B, 1, d))
 
 
-class CausalSplatHUB(nn.Module):
-    def __init__(self, d: int, K: int = 512, D: int = 32, tau: float = 0.1,
-                 chunk: int = 256):
+class _Constellation(nn.Module):
+    """One codebook with its own routing-owned q/k frames (v2 form).
+
+    The multi-constellation hub is the PRODUCT-CODE form (B2: independent
+    frames compose, .859 -> .955 monotone in members) at lawful supply
+    (ROUND 5e: K <= 2*D per address space — v1's single 512-anchor book in
+    32 dims ran 16x and crowded into 333-646 duplicate pairs)."""
+
+    def __init__(self, d: int, K: int, D: int, tau: float):
         super().__init__()
         self.addr = AlephAddress(K, D, tau)
-        self.chunk = chunk          # 256 measured best at ctx 2048 (bench)
         self.q = nn.Linear(d, D, bias=False)
         self.k = nn.Linear(d, D, bias=False)
+        nn.init.orthogonal_(self.q.weight)
+        nn.init.orthogonal_(self.k.weight)
+
+
+class CausalSplatHUB(nn.Module):
+    def __init__(self, d: int, K: int = 512, D: int = 32, tau: float = 0.1,
+                 chunk: int = 256, n_const: int = 1):
+        super().__init__()
+        if K > 2 * D:
+            import warnings
+            warnings.warn(
+                f"CausalSplatHUB supply K={K} exceeds 2*D={2*D}: anchors on "
+                f"a {D}-dim sphere past ~2x supply CROWD (measured — ROUND "
+                "5e shape ladder + the mini-beatrix-1 hub census: duplicate "
+                "pairs by the hundreds, consumed erank collapse). Provision "
+                "K <= 2*D or raise D.", stacklevel=2)
+        self.n_const = n_const
+        if n_const == 1:
+            # v1 layout, bit-for-bit: state-dict keys addr/q/k unchanged so
+            # every shipped checkpoint and the HF automodel mirror load.
+            self.addr = AlephAddress(K, D, tau)
+            self.q = nn.Linear(d, D, bias=False)
+            self.k = nn.Linear(d, D, bias=False)
+            nn.init.orthogonal_(self.q.weight)
+            nn.init.orthogonal_(self.k.weight)
+        else:
+            self.consts = nn.ModuleList(
+                _Constellation(d, K, D, tau) for _ in range(n_const))
+        self.chunk = chunk          # 256 measured best at ctx 2048 (bench)
         self.v = nn.Linear(d, d, bias=False)
         self.o = nn.Linear(d, d, bias=False)
-        for m in (self.q, self.k, self.v, self.o):
+        for m in (self.v, self.o):
             nn.init.orthogonal_(m.weight)
         self._mask_cache: dict = {}
         self._den_raw = None        # (den tensor, floor) until read
@@ -102,34 +136,27 @@ class CausalSplatHUB(nn.Module):
             self._mask_cache[key] = m
         return m
 
+    # ------------------------------------------------ constellation access
+    def _units(self):
+        """Uniform view: [(addr, q, k)] whether single- or multi-book."""
+        if self.n_const == 1:
+            return [(self.addr, self.q, self.k)]
+        return [(c.addr, c.q, c.k) for c in self.consts]
+
     def _halves(self, x):
-        qp, qn = self.addr.oriented(self.q(x))
-        kp, kn = self.addr.oriented(self.k(x))
-        return qp, qn, kp, kn, self.v(x)
+        """Per-constellation oriented halves + shared values."""
+        outs = []
+        for addr, q, k in self._units():
+            qp, qn = addr.oriented(q(x))
+            kp, kn = addr.oriented(k(x))
+            outs.append((qp, qn, kp, kn))
+        return outs, self.v(x)
 
-    def forward(self, x):
-        """Fast path: the two oriented halves run as ONE 2K-wide pass —
-        every term is a sum of bilinear forms over the halves, so one
-        pass over cat(p, n) is the same arithmetic in half the kernels
-        (equal to forward_naive to fp reorder, ~1.5e-06; speed-harness
-        verdict 2026-08-15: 1.7x eager, 4.0x under torch.compile)."""
-        B, n, d = x.shape
-        qc = self.addr.oriented_cat(self.q(x))            # (B, n, 2K)
-        kc = self.addr.oriented_cat(self.k(x))
-        v = self.v(x)
-        C = min(self.chunk, n)
-        pad = (-n) % C
-        if pad:
-            qc = F.pad(qc, (0, 0, 0, pad))
-            kc = F.pad(kc, (0, 0, 0, pad))
-            v = F.pad(v, (0, 0, 0, pad))
-        nc = (n + pad) // C
+    def _scan_cat(self, qc, kc, v, mask, B, n, nc, C, d):
+        """The exact chunked scan for one 2K-wide constellation."""
         K2 = qc.shape[-1]
         qc = qc.view(B, nc, C, K2)
         kc = kc.view(B, nc, C, K2)
-        v = v.view(B, nc, C, d)
-        mask = self._mask(C, x.device, v.dtype)
-
         S = torch.einsum("bick,bicd->bikd", kc, v)     # per-chunk 2KxD sums
         P = torch.cumsum(S, dim=1) - S                  # exclusive prefix
         zS = kc.sum(dim=2)                              # (B, nc, 2K)
@@ -138,9 +165,36 @@ class CausalSplatHUB(nn.Module):
         num = torch.einsum("bick,bikd->bicd", qc, P) + att @ v
         den = torch.einsum("bick,bik->bic", qc, zP).unsqueeze(-1) \
             + att.sum(dim=-1, keepdim=True)
+        return num.reshape(B, nc * C, d)[:, :n], den.reshape(B, nc * C, 1)[:, :n]
 
-        num = num.reshape(B, nc * C, d)[:, :n]
-        den = den.reshape(B, nc * C, 1)[:, :n]
+    def forward(self, x):
+        """Fast path: the two oriented halves run as ONE 2K-wide pass —
+        every term is a sum of bilinear forms over the halves, so one
+        pass over cat(p, n) is the same arithmetic in half the kernels
+        (equal to forward_naive to fp reorder, ~1.5e-06; speed-harness
+        verdict 2026-08-15: 1.7x eager, 4.0x under torch.compile).
+        Multi-constellation (n_const > 1): each book scans independently
+        and the reads compose BY BUDGET — numerators and agreement masses
+        sum across books before the single divide (never softmax over
+        books; B4 measured comparative composition at -.10)."""
+        B, n, d = x.shape
+        v = self.v(x)
+        C = min(self.chunk, n)
+        pad = (-n) % C
+        vp = F.pad(v, (0, 0, 0, pad)) if pad else v
+        nc = (n + pad) // C
+        vc = vp.view(B, nc, C, d)
+        mask = self._mask(C, x.device, v.dtype)
+        num = den = None
+        for addr, q, k in self._units():
+            qc = addr.oriented_cat(q(x))                # (B, n, 2K)
+            kc = addr.oriented_cat(k(x))
+            if pad:
+                qc = F.pad(qc, (0, 0, 0, pad))
+                kc = F.pad(kc, (0, 0, 0, pad))
+            nu, de = self._scan_cat(qc, kc, vc, mask, B, n, nc, C, d)
+            num = nu if num is None else num + nu
+            den = de if den is None else den + de
         cl = dtype_floor(den)
         self._den_raw = (den.detach(), cl)
         self._den_stats = None
@@ -149,39 +203,54 @@ class CausalSplatHUB(nn.Module):
     # ---------------------------------------------------- incremental decode
     def prefill(self, x):
         """Full causal pass plus the decode cache. The hub's cache is the
-        CONSTANT-SIZE prefix state (Sp, Sn, zp, zn) — O(K·d) regardless of
-        sequence length; this is the linear-attention decode advantage."""
+        CONSTANT-SIZE prefix state (Sp, Sn, zp, zn) per constellation —
+        O(n_const·K·d) regardless of sequence length; this is the
+        linear-attention decode advantage. n_const == 1 keeps the exact
+        v1 cache shape (arms and the Space depend on it)."""
         out = self.forward(x)
-        qp, qn, kp, kn, v = self._halves(x)
-        cache = {"Sp": torch.einsum("bnk,bnd->bkd", kp, v),
-                 "Sn": torch.einsum("bnk,bnd->bkd", kn, v),
-                 "zp": kp.sum(dim=1), "zn": kn.sum(dim=1)}
-        return out, cache
+        halves, v = self._halves(x)
+        caches = [{"Sp": torch.einsum("bnk,bnd->bkd", kp, v),
+                   "Sn": torch.einsum("bnk,bnd->bkd", kn, v),
+                   "zp": kp.sum(dim=1), "zn": kn.sum(dim=1)}
+                  for (qp, qn, kp, kn) in halves]
+        return out, (caches[0] if self.n_const == 1 else {"consts": caches})
 
     def step(self, x_t, cache):
-        """One new position: fold it into the prefix state, read once."""
-        qp, qn, kp, kn, v = self._halves(x_t)          # (B,1,K)/(B,1,d)
-        kp1, kn1, v1 = kp.squeeze(1), kn.squeeze(1), v.squeeze(1)
-        cache["Sp"] = cache["Sp"] + kp1.unsqueeze(-1) * v1.unsqueeze(1)
-        cache["Sn"] = cache["Sn"] + kn1.unsqueeze(-1) * v1.unsqueeze(1)
-        cache["zp"] = cache["zp"] + kp1
-        cache["zn"] = cache["zn"] + kn1
-        qp1, qn1 = qp.squeeze(1), qn.squeeze(1)
-        num = torch.einsum("bk,bkd->bd", qp1, cache["Sp"]) \
-            + torch.einsum("bk,bkd->bd", qn1, cache["Sn"])
-        den = ((qp1 * cache["zp"]).sum(-1)
-               + (qn1 * cache["zn"]).sum(-1)).unsqueeze(-1)
+        """One new position: fold it into each prefix state, read once,
+        compose by budget across constellations."""
+        halves, v = self._halves(x_t)                  # (B,1,K)/(B,1,d)
+        caches = [cache] if self.n_const == 1 else cache["consts"]
+        v1 = v.squeeze(1)
+        num = den = None
+        for (qp, qn, kp, kn), c in zip(halves, caches):
+            kp1, kn1 = kp.squeeze(1), kn.squeeze(1)
+            c["Sp"] = c["Sp"] + kp1.unsqueeze(-1) * v1.unsqueeze(1)
+            c["Sn"] = c["Sn"] + kn1.unsqueeze(-1) * v1.unsqueeze(1)
+            c["zp"] = c["zp"] + kp1
+            c["zn"] = c["zn"] + kn1
+            qp1, qn1 = qp.squeeze(1), qn.squeeze(1)
+            nu = torch.einsum("bk,bkd->bd", qp1, c["Sp"]) \
+                + torch.einsum("bk,bkd->bd", qn1, c["Sn"])
+            de = ((qp1 * c["zp"]).sum(-1)
+                  + (qn1 * c["zn"]).sum(-1)).unsqueeze(-1)
+            num = nu if num is None else num + nu
+            den = de if den is None else den + de
         return self.o((num / den.clamp_min(dtype_floor(den))).unsqueeze(1))
 
     def forward_naive(self, x):
         """Reference cumsum form (the validated probe-bed implementation).
-        O(n·K·d) memory — test oracle only."""
-        qp, qn, kp, kn, v = self._halves(x)
-        Sp = torch.cumsum(torch.einsum("bnk,bnd->bnkd", kp, v), dim=1)
-        Sn = torch.cumsum(torch.einsum("bnk,bnd->bnkd", kn, v), dim=1)
-        zp = torch.cumsum(kp, dim=1)
-        zn = torch.cumsum(kn, dim=1)
-        num = torch.einsum("bnk,bnkd->bnd", qp, Sp) \
-            + torch.einsum("bnk,bnkd->bnd", qn, Sn)
-        den = (qp * zp).sum(-1, keepdim=True) + (qn * zn).sum(-1, keepdim=True)
+        O(n·K·d) memory — test oracle only. Sums constellations by budget,
+        matching forward()."""
+        halves, v = self._halves(x)
+        num = den = None
+        for qp, qn, kp, kn in halves:
+            Sp = torch.cumsum(torch.einsum("bnk,bnd->bnkd", kp, v), dim=1)
+            Sn = torch.cumsum(torch.einsum("bnk,bnd->bnkd", kn, v), dim=1)
+            zp = torch.cumsum(kp, dim=1)
+            zn = torch.cumsum(kn, dim=1)
+            nu = torch.einsum("bnk,bnkd->bnd", qp, Sp) \
+                + torch.einsum("bnk,bnkd->bnd", qn, Sn)
+            de = (qp * zp).sum(-1, keepdim=True) + (qn * zn).sum(-1, keepdim=True)
+            num = nu if num is None else num + nu
+            den = de if den is None else den + de
         return self.o(num / den.clamp_min(dtype_floor(den)))