"""Attention blocks: CausalSDPA (the workhorse) and CausalSplatHUB (the instrumented aleph read). CausalSplatHUB is causal linear attention through the oriented address: prefix-sum memories over the two K-wide halves of the 2K softmax, read by the query's halves and normalized by the scalar agreement mass. O(n·K·d) compute, no softmax over positions, no selection event anywhere. The naive cumsum form materializes (B, n, K, d) — fine on probe beds, fatal at mission scale. forward() therefore uses an exact chunked scan: within-chunk causal affinity (B, C, C) + cross-chunk carried states (B, K, d). `forward_naive()` is kept verbatim as the equivalence oracle for the test array. """ from __future__ import annotations import torch import torch.nn as nn import torch.nn.functional as F from .address import AlephAddress, dtype_floor class CausalSDPA(nn.Module): def __init__(self, d: int, heads: int = 8): super().__init__() assert d % heads == 0 self.h = heads self.qkv = nn.Linear(d, 3 * d, bias=False) self.o = nn.Linear(d, d, bias=False) nn.init.orthogonal_(self.qkv.weight) nn.init.orthogonal_(self.o.weight) def forward(self, x): B, n, d = x.shape q, k, v = self.qkv(x).chunk(3, dim=-1) q, k, v = (t.view(B, n, self.h, d // self.h).transpose(1, 2) for t in (q, k, v)) y = F.scaled_dot_product_attention(q, k, v, is_causal=True) return self.o(y.transpose(1, 2).reshape(B, n, d)) # ---------------------------------------------------- incremental decode def prefill(self, x): """Full causal pass that also returns the decode cache (K/V).""" B, n, d = x.shape q, k, v = self.qkv(x).chunk(3, dim=-1) q, k, v = (t.view(B, n, self.h, d // self.h).transpose(1, 2) for t in (q, k, v)) y = F.scaled_dot_product_attention(q, k, v, is_causal=True) return self.o(y.transpose(1, 2).reshape(B, n, d)), {"k": k, "v": v} def step(self, x_t, cache): """One new position attending over everything cached (KV cache).""" B, _, d = x_t.shape q, k, v = self.qkv(x_t).chunk(3, dim=-1) q, k, v = (t.view(B, 1, self.h, d // self.h).transpose(1, 2) for t in (q, k, v)) cache["k"] = torch.cat([cache["k"], k], dim=2) cache["v"] = torch.cat([cache["v"], v], dim=2) y = F.scaled_dot_product_attention(q, cache["k"], cache["v"]) return self.o(y.transpose(1, 2).reshape(B, 1, d)) 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.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.v, self.o): nn.init.orthogonal_(m.weight) self._mask_cache: dict = {} self._den_raw = None # (den tensor, floor) until read self._den_stats = None # cached floats after first read # den stats are LAZY: the reference forward paid three .item() GPU # syncs per call just to keep this attribute warm; instruments read # it at most once per health interval. Property keeps the tuple API. @property def last_den_stats(self): if self._den_stats is None and self._den_raw is not None: den, cl = self._den_raw with torch.no_grad(): self._den_stats = (den.min().item(), den.mean().item(), (den <= cl).float().mean().item()) return self._den_stats @last_den_stats.setter def last_den_stats(self, value): self._den_stats = value self._den_raw = None def _mask(self, C: int, device, dtype): key = (C, device, dtype) m = self._mask_cache.get(key) if m is None: m = torch.tril(torch.ones(C, C, device=device, dtype=dtype)) self._mask_cache[key] = m return m def _prefix(self, nc: int, device, dtype): """Strictly-lower-triangular ones (nc, nc): the exclusive prefix sum as ONE tensor-core GEMM. The cumsum scan kernel ran ~6x off its memory roofline on the (B, nc, 2K·H, d) layout (C2d, Blackwell 2026-08-26) and its backward is flip+cumsum+flip; matmul accumulates fp32 inside the GEMM — strictly MORE precise than a bf16 cumsum.""" key = ("prefix", nc, device, dtype) m = self._mask_cache.get(key) if m is None: m = torch.tril(torch.ones(nc, nc, device=device, dtype=dtype), diagonal=-1) self._mask_cache[key] = m return m # ------------------------------------------------ constellation access def _code_cat_qk(self, x): """BOTH oriented codes (q and k, every book) in one batched pass. Stack all 2H frame weights, one projection einsum, one address einsum, ONE fused softmax. The oriented address IS softmax over the 2K half-axes — exp(cat[u−m, −u−m])/Σ with m = max|u| is bit-the-same quantity as F.softmax(cat[u, −u]) (softmax subtracts its own max, which is exactly m). This is a KERNEL substitution, not a mechanism change: no softmax over positions, no softmax across books — composition stays budget. The old chain was ~8 unfused GB-scale elementwise passes per call, twice per forward (C2d: 24.6 ms). AUTOCAST TRAP (measured, Blackwell 2026-08-26): torch.einsum is in autocast's PROMOTE category, and F.normalize / exp / softmax are on its fp32 list — one fp32 operand drags the whole downstream scan to fp32. Operands are cast to the autocast dtype explicitly; the softmax accumulates fp32 inside the kernel (standard attention practice) and the returned CODE is in the compute dtype so the num/S/P scan and its backward run bf16. DELIBERATE exception: den's reductions (kc.sum, att.sum) stay fp32 by autocast policy — the agreement mass keeps v1's fp32 dtype_floor semantics at ~5% of scan traffic (dtype audit 2026-08-26). Returns (qc, kc), each (B, n, H*2K), per-book layout [K pos | K neg] matching oriented()/forward_naive.""" units = self._units() dt = (torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() and x.is_cuda else x.dtype) W = torch.stack([q.weight for _, q, _ in units] + [k.weight for _, _, k in units]).to(dt) # (2H, D, d) tau = units[0][0].tau # tau folds into the codebook (a few-MB fp32 tensor op, MORE precise # than dividing bf16 u afterwards), and the query-side row # normalization folds into ONE post-GEMM scale: a per-row scalar # commutes through the linear map, so (xh/||xh||) @ A^T / tau == # (xh @ (A/tau)^T) * (1/||xh||) exactly (fp reorder). Kills the # fp32 normalize-div + cast + separate tau-div passes (perf audit # 2026-08-26). Same 1e-12 floor as F.normalize. A = (F.normalize(torch.stack([a.codebook for a, _, _ in units]), dim=-1) / tau).to(dt) # (H, K, D) A = torch.cat([A, A]) # (2H, K, D) xh = torch.einsum("bnd,hkd->bnhk", x.to(dt), W) # (B,n,2H,D) inv = torch.linalg.vector_norm( # fp32 by autocast xh, dim=-1, keepdim=True).clamp_min(1e-12) \ .reciprocal().to(dt) # policy; cast back u = torch.einsum("bnhd,hkd->bnhk", xh, A) * inv # (B,n,2H,K) B, n = x.shape[:2] H = len(units) # Split-axis-first WITHOUT a copy: the permute is a view, and the # cat (which must write a fresh tensor anyway) absorbs it — so the # q/k split below is a pure view instead of two GB-scale reshape # copies. Layout per book stays [K pos | K neg], book-major. u = u.view(B, n, 2, H, -1).permute(2, 0, 1, 3, 4) # softmax: the explicit dtype arg opts out of autocast's fp32 # override (fp32_set_opt_dtype policy) while the CUDA kernel still # accumulates fp32 internally — bf16-in/bf16-out, no fp32 e pass, # and the softmax BACKWARD chain halves too. e = F.softmax(torch.cat([u, -u], dim=-1), dim=-1, dtype=dt) return e[0].reshape(B, n, -1), e[1].reshape(B, n, -1) 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): """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 _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) S = torch.einsum("bick,bicd->bikd", kc, v) # per-chunk 2KxD sums L = self._prefix(nc, qc.device, qc.dtype) P = torch.matmul(L, S.reshape(B, nc, -1)).view_as(S) # excl. prefix zS = kc.sum(dim=2) # (B, nc, 2K) zP = torch.matmul(L, zS) att = torch.einsum("bick,bijk->bicj", qc, kc) * mask # (B,nc,C,C) 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] 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) # BATCHED path, both n_const cases (2026-08-26 Blackwell verdicts: # the per-book Python loop was 512 sequential little scans — # launch-bound; then the split q/k exp chains were ~8 unfused # GB-scale passes each). Budget composition is algebraically ONE # scan over the concatenated code: num and den are sums of per-book # bilinear forms, so scanning cat_h(qc_h) against cat_h(kc_h) # equals summing H separate scans (fp reorder). qc, kc = self._code_cat_qk(x) # (B, n, H*2K) if qc.dtype != v.dtype: # einsum-promote guard (belt-and-braces; qc = qc.to(v.dtype) # _code_cat_qk already returns the kc = kc.to(v.dtype) # compute dtype) if pad: qc = F.pad(qc, (0, 0, 0, pad)) kc = F.pad(kc, (0, 0, 0, pad)) num, den = self._scan_cat(qc, kc, vc, mask, B, n, nc, C, d) cl = dtype_floor(den) self._den_raw = (den.detach(), cl) self._den_stats = None return self.o(num / den.clamp_min(cl)) # ---------------------------------------------------- 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) 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) 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 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. 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)))