modmul router v0: t12 + composed-t3b members (pre-t3g)

composed_encoding.py

@@ -0,0 +1,193 @@
+"""Composed multiply+reduce CoT — the DEPLOYABLE surface.
+
+divcot validates the reduction in isolation by feeding N=x*y as input; a real
+submission only gets (x, y, p) with x,y < p. This encoding chains the two
+validated scaffolds in one generated trace:
+
+  1. schoolbook multiply as a running accumulator: after consuming each y-limb
+     y_j the model emits acc_j = acc_{j-1} + x*y_j*B^j (2W limbs, LSB-first so
+     the product+add carry chain is local). acc_{W-1} = N.
+  2. copy-reverse: re-emit N MSB-first (attention reversal — staging the
+     dividend so the division block looks exactly like validated divcot).
+  3. long division by p: quotient limbs MSB-first with running-remainder
+     install targets, then the answer (= final remainder) LSB-first.
+
+Surface:
+  "A" x(W,LSB) "B" y(W,LSB) "M" p(W,MSB) "="
+      [ "a" acc_j(2W,LSB) ] * W
+      "N" N(2W,MSB)  Q(2W)  "R" ans(W,LSB) "\n"
+
+Tokens/example = 2W^2 + 9W + 7 (W=5: 102; base-100 tier 5 W=10: 297).
+
+Install targets (NTP-aligned):
+  acarry  carry into each emitted acc limb (the multiply-accumulate state)
+  rem{j}  MSB limb j of the running remainder at each quotient position
+  qhat    leading-limbs quotient estimate at each quotient position
+  ans     limb k of (x*y) mod p at each answer position
+
+Compliance: identical status to divcot/karatsuba — fixed trace SHAPE, every
+value generated by trained weights; randomising weights collapses accuracy.
+"""
+from __future__ import annotations
+
+from limbs import limb_char, limb_str, parse_limbs, to_limbs
+from divcot2_encoding import long_division, qhat_estimate
+
+AMARK, BMARK, DIVMARK, EQ, ACC, NMARK, REVMARK, NL = "A", "B", "M", "=", "a", "N", "R", "\n"
+PLSB = "m"
+
+
+def prompt_str(x: int, y: int, p: int, W: int, base: int, subpad: bool = False) -> str:
+    """subpad adds p LSB-first to the prompt: the MSB field serves quotient
+    estimation, the LSB field aligns with the LSB-emitted qd*p / remainder
+    chains (per-step diagnostic: subtraction digits need REVERSED access into
+    the MSB-only field — the aligned acc stage learned, the division didn't)."""
+    s = (AMARK + limb_str(x, W, base) + BMARK + limb_str(y, W, base)
+         + DIVMARK + limb_str(p, W, base, msb_first=True))
+    if subpad:
+        s += PLSB + limb_str(p, W, base)
+    return s + EQ
+
+
+def gen_len(W: int, scratch: bool = False, cursor: bool = False,
+            subpad: bool = False) -> int:
+    """Tokens generated after '=': W acc blocks, N restage, quotient, answer.
+    scratch=True: each quotient limb is followed by the running remainder
+    (W limbs, LSB-first) — externalizes the division state, which the model
+    cannot hold internally past W~3 (composed tier-3 diagnostic: first
+    quotient digit right, then collapse).
+    cursor=True (requires scratch): each step block starts with the consumed
+    dividend limb, copied from the restaged N — an explicit progress anchor.
+    Diagnostic at scratch step-4k: mid-sequence remainder drift + the model
+    losing count of steps (9 blocks emitted instead of 2W); copying gives the
+    same position mechanism that put restage at 1.00.
+    subpad=True (requires cursor): each step also emits qd*p (W+1 limbs, LSB)
+    before the remainder — externalizes the multiply-subtract that the
+    per-step diagnostic showed failing (steps 0-4 = copies, perfect; real
+    division steps 5+ at 0.04-0.15 even teacher-forced). Tier-2's compact
+    success could memorize qd*p over 48 primes; tier 3's 400 primes need the
+    generic circuit, so generate it like the (perfectly learned) acc rows."""
+    per_step = (1 if cursor else 0) + 1 + ((W + 1) if subpad else 0) + (W if scratch else 0)
+    return W * (1 + 2 * W) + (1 + 2 * W) + 2 * W * per_step + 1 + W
+
+
+def build_example(x: int, y: int, p: int, W: int, base: int, scratch: bool = False,
+                  cursor: bool = False, subpad: bool = False):
+    """Return (text, ann). x,y < p < base**W."""
+    assert not cursor or scratch, "cursor requires scratch"
+    assert not subpad or cursor, "subpad requires cursor"
+    N = x * y
+    y_limbs = to_limbs(y, W, base)
+    x_limbs = to_limbs(x, W, base)
+    parts, ann_g = [], []                       # ann_g: (gen_index, var, val)
+    pos = 0
+
+    acc = 0
+    for j, yj in enumerate(y_limbs):            # multiply-accumulate rows
+        prev = acc
+        acc = prev + x * yj * base ** j
+        parts.append(ACC)
+        pos += 1
+        prev_l = to_limbs(prev, 2 * W, base)
+        c = 0
+        for i in range(2 * W):                  # carry into limb i of acc_j
+            ann_g.append((pos + i, "acarry", c))
+            xi = x_limbs[i - j] if 0 <= i - j < W else 0
+            c = (prev_l[i] + xi * yj + c) // base
+        parts.append(limb_str(acc, 2 * W, base))
+        pos += 2 * W
+    assert acc == N
+
+    parts.append(NMARK + limb_str(N, 2 * W, base, msb_first=True))  # restage MSB
+    pos += 1 + 2 * W
+
+    q_limbs, rems, answer = long_division(N, p, W, base)
+    n_msb = to_limbs(N, 2 * W, base, msb_first=True)
+    r_prev = 0
+    for i, (qd, r) in enumerate(zip(q_limbs, rems)):
+        if cursor:                               # copy the consumed dividend limb
+            parts.append(limb_char(n_msb[i]))
+            pos += 1
+        ann_g.append((pos, "qhat", qhat_estimate(r_prev, n_msb[i], p, base)))
+        if not scratch:                          # compact: remainder is internal state
+            rstr = to_limbs(r, W, base, msb_first=True)
+            for j in range(W):
+                ann_g.append((pos, f"rem{j}", rstr[j]))
+        parts.append(limb_char(qd))
+        pos += 1
+        if subpad:                               # emit qd*p LSB-first (generic mult)
+            qdp = qd * p
+            p_l = to_limbs(p, W, base) + [0]
+            c = 0
+            for k in range(W + 1):               # carry into limb k of qd*p
+                ann_g.append((pos + k, "spcarry", c))
+                c = (qd * p_l[k] + c) // base
+            parts.append(limb_str(qdp, W + 1, base))
+            pos += W + 1
+        if scratch:                              # emit remainder LSB-first (local borrow)
+            if subpad:                           # borrow chain of num - qd*p
+                num_l = [n_msb[i]] + to_limbs(r_prev, W, base)
+                qdp_l = to_limbs(qd * p, W + 1, base)
+                b = 0
+                for k in range(W):
+                    ann_g.append((pos + k, "sborrow", b))
+                    b = 1 if num_l[k] - qdp_l[k] - b < 0 else 0
+            parts.append(limb_str(r, W, base))
+            pos += W
+        r_prev = r
+    parts.append(REVMARK)
+    pos += 1
+    ans_limbs = to_limbs(answer, W, base)
+    for k in range(W):
+        ann_g.append((pos + k, "ans", ans_limbs[k]))
+    parts.append(limb_str(answer, W, base) + NL)
+    pos += W + 1
+
+    prompt = prompt_str(x, y, p, W, base, subpad)
+    text = prompt + "".join(parts)
+    assert len(text) == len(prompt) + gen_len(W, scratch, cursor, subpad) + 1, \
+        (len(text), len(prompt), gen_len(W, scratch, cursor, subpad))
+    ann = [dict() for _ in range(len(text))]
+    base_i = len(prompt)
+    for gi, var, val in ann_g:
+        ann[base_i + gi - 1][var] = int(val)    # NTP alignment
+    return text, ann
+
+
+def var_specs(W: int, base: int, scratch: bool = False, subpad: bool = False):
+    specs = [("ans", base), ("qhat", base), ("acarry", base)]
+    if subpad:                                  # multiply-carry + subtract-borrow chains
+        specs += [("spcarry", base), ("sborrow", 2)]
+    if not scratch:                             # scratch emits remainders: no probe
+        specs += [(f"rem{j}", base) for j in range(W)]
+    return specs
+
+
+def decode_answer(gen_chars: str, W: int, base: int, scratch: bool = False,
+                  cursor: bool = False, subpad: bool = False) -> int:
+    off = gen_len(W, scratch, cursor, subpad) - W  # answer = last W generated limbs
+    return parse_limbs(gen_chars[off:off + W], base)
+
+
+if __name__ == "__main__":
+    import random
+    rng = random.Random(3)
+    for base in (10, 100):
+        for scratch, cursor, subpad in ((False, False, False), (True, False, False),
+                                        (True, True, False), (True, True, True)):
+            for _ in range(4000):
+                W = rng.randint(1, 6)
+                p = rng.randrange(max(2, base ** (W - 1)), base ** W)
+                x, y = rng.randrange(p), rng.randrange(p)
+                text, ann = build_example(x, y, p, W, base, scratch, cursor, subpad)
+                plen = len(prompt_str(x, y, p, W, base, subpad))
+                assert decode_answer(text[plen:], W, base, scratch, cursor, subpad) \
+                    == (x * y) % p, (x, y, p, W, base, scratch, cursor, subpad)
+                assert len(ann) == len(text)
+            print(f"composed base={base} scratch={scratch} cursor={cursor} "
+                  f"subpad={subpad}: 4000/4000 decode OK")
+    print("\ntokens/example (prompt+gen+NL):")
+    for tier, Wd in [(3, 5), (4, 10), (5, 20), (6, 39), (7, 78)]:
+        for base, W in ((10, Wd), (100, (Wd + 1) // 2)):
+            print(f"  tier {tier} base {base:>3} W={W:>3}  compact {3 * W + 4 + gen_len(W) + 1:>6}"
+                  f"  scratch {3 * W + 4 + gen_len(W, True) + 1:>6}")

coppola_pretraining.py

@@ -0,0 +1,613 @@
+#!/usr/bin/env python3
+"""Utilities for full Coppola-typed pretraining in nanochat-style trainers.
+
+This module is intentionally trainer-agnostic. It provides:
+
+- basis extraction from current attention-output and MLP down-project weights
+- typed gradient decomposition and projection
+- a depth-aware scale policy
+- a small controller that can refresh bases from a nanochat-style model
+
+The intended integration is:
+
+1. instantiate a controller from the current model weights
+2. before each matrix optimizer step, project the raw gradient by family
+3. apply Muon / Newton-Schulz or another matrix optimizer
+4. re-project the transformed gradient before the final step
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Callable, Dict, List, Optional, Sequence, Tuple
+
+import torch
+
+
+Tensor = torch.Tensor
+
+
+@dataclass(frozen=True)
+class RankPolicy:
+    """Energy-based rank truncation policy."""
+
+    energy: float = 0.99
+    max_rank: Optional[int] = None
+    min_rank: int = 1
+
+
+@dataclass(frozen=True)
+class FamilyScales:
+    """Typed gradient multipliers for a layer."""
+
+    routing: float = 1.0
+    and_transport: float = 1.0
+    or_update: float = 1.0
+    remainder: float = 1.0
+    coupling: Optional[float] = None
+
+    def coupling_scale(self) -> float:
+        if self.coupling is not None:
+            return self.coupling
+        return 0.5 * (self.and_transport + self.or_update)
+
+
+@dataclass
+class LayerBases:
+    """Typed bases for one transformer layer."""
+
+    routing_head_bases: List[Tensor] = field(default_factory=list)
+    q_head_bases: List[Tensor] = field(default_factory=list)
+    k_head_bases: List[Tensor] = field(default_factory=list)
+    and_output_basis: Optional[Tensor] = None
+    or_input_basis: Optional[Tensor] = None
+
+
+@dataclass
+class ComponentNorms:
+    """Norms of the typed gradient components before scaling."""
+
+    routing: float = 0.0
+    and_transport: float = 0.0
+    or_update: float = 0.0
+    coupling: float = 0.0
+    remainder: float = 0.0
+
+
+def _select_rank(singular_values: Tensor, policy: RankPolicy) -> int:
+    if singular_values.numel() == 0:
+        return policy.min_rank
+    total = float(singular_values.square().sum().item())
+    if total <= 0.0:
+        return policy.min_rank
+    cumulative = singular_values.square().cumsum(0) / total
+    rank = int(torch.searchsorted(cumulative, torch.tensor(policy.energy, device=cumulative.device)).item()) + 1
+    rank = max(policy.min_rank, rank)
+    if policy.max_rank is not None:
+        rank = min(rank, policy.max_rank)
+    return min(rank, singular_values.numel())
+
+
+def _orthonormal_rows(basis: Tensor) -> Tensor:
+    """Return an orthonormal row basis spanning the same row space."""
+    if basis.ndim != 2:
+        raise ValueError(f"basis must be rank-2, got {tuple(basis.shape)}")
+    if basis.shape[0] == 0:
+        return basis
+    q, _ = torch.linalg.qr(basis.T, mode="reduced")
+    return q.T.contiguous()
+
+
+def _project_left(grad: Tensor, basis: Tensor) -> Tensor:
+    basis = basis.to(device=grad.device, dtype=grad.dtype)
+    proj = basis.T @ basis
+    return proj @ grad
+
+
+def _project_right(grad: Tensor, basis: Tensor) -> Tensor:
+    basis = basis.to(device=grad.device, dtype=grad.dtype)
+    proj = basis.T @ basis
+    return grad @ proj
+
+
+def _zero_like(grad: Tensor) -> Tensor:
+    return torch.zeros_like(grad)
+
+
+def compute_attention_output_bases(
+    weight: Tensor,
+    n_head: int,
+    policy: RankPolicy = RankPolicy(),
+) -> List[Tensor]:
+    """Build one output-space basis per attention head block.
+
+    Args:
+        weight: output projection weight of shape [hidden, hidden]
+        n_head: number of attention heads
+        policy: rank selection policy inside each head block
+    """
+    if weight.ndim != 2:
+        raise ValueError(f"attention output weight must be rank-2, got {tuple(weight.shape)}")
+    hidden, width = weight.shape
+    if width % n_head != 0:
+        raise ValueError(f"weight width {width} not divisible by n_head={n_head}")
+    head_dim = width // n_head
+    bases: List[Tensor] = []
+    with torch.no_grad():
+        w = weight.detach().float().cpu()
+        for h in range(n_head):
+            cols = w[:, h * head_dim:(h + 1) * head_dim]
+            u, s, _ = torch.linalg.svd(cols, full_matrices=False)
+            rank = _select_rank(s, policy)
+            basis = u[:, :rank].T.contiguous()
+            bases.append(_orthonormal_rows(basis))
+    return bases
+
+
+def compute_attention_score_bases(
+    weight: Tensor,
+    n_head: int,
+    policy: RankPolicy = RankPolicy(),
+) -> List[Tensor]:
+    """Build one row-space basis per attention head for Q/K score-side blocks."""
+    if weight.ndim != 2:
+        raise ValueError(f"attention score weight must be rank-2, got {tuple(weight.shape)}")
+    width, hidden = weight.shape
+    if width % n_head != 0:
+        raise ValueError(f"weight height {width} not divisible by n_head={n_head}")
+    head_dim = width // n_head
+    bases: List[Tensor] = []
+    with torch.no_grad():
+        w = weight.detach().float().cpu()
+        for h in range(n_head):
+            rows = w[h * head_dim:(h + 1) * head_dim, :]
+            u, s, _ = torch.linalg.svd(rows, full_matrices=False)
+            rank = _select_rank(s, policy)
+            basis = u[:, :rank].T.contiguous()
+            bases.append(_orthonormal_rows(basis))
+    return bases
+
+
+def compute_down_proj_bases(
+    weight: Tensor,
+    output_policy: RankPolicy = RankPolicy(),
+    input_policy: RankPolicy = RankPolicy(),
+) -> Tuple[Tensor, Tensor]:
+    """Build output-side AND basis and input-side OR basis for down_proj."""
+    if weight.ndim != 2:
+        raise ValueError(f"down projection weight must be rank-2, got {tuple(weight.shape)}")
+    with torch.no_grad():
+        w = weight.detach().float().cpu()
+        u, s, vh = torch.linalg.svd(w, full_matrices=False)
+        out_rank = _select_rank(s, output_policy)
+        in_rank = _select_rank(s, input_policy)
+        and_basis = _orthonormal_rows(u[:, :out_rank].T.contiguous())
+        or_basis = _orthonormal_rows(vh[:in_rank, :].contiguous())
+    return and_basis, or_basis
+
+
+def decompose_down_proj_gradient(
+    grad: Tensor,
+    and_output_basis: Optional[Tensor],
+    or_input_basis: Optional[Tensor],
+) -> Dict[str, Tensor]:
+    """Decompose an MLP down_proj gradient into typed components.
+
+    The components are:
+    - and_transport_only
+    - or_update_only
+    - coupling: in both the output and input typed subspaces
+    - remainder
+    """
+    if grad.ndim != 2:
+        raise ValueError(f"gradient must be rank-2, got {tuple(grad.shape)}")
+
+    g_and = _zero_like(grad)
+    g_or = _zero_like(grad)
+    g_both = _zero_like(grad)
+
+    if and_output_basis is not None:
+        g_and = _project_left(grad, and_output_basis)
+    if or_input_basis is not None:
+        g_or = _project_right(grad, or_input_basis)
+    if and_output_basis is not None and or_input_basis is not None:
+        g_both = _project_right(g_and, or_input_basis)
+
+    g_and_only = g_and - g_both
+    g_or_only = g_or - g_both
+    g_remainder = grad - g_and_only - g_or_only - g_both
+
+    return {
+        "and_transport_only": g_and_only,
+        "or_update_only": g_or_only,
+        "coupling": g_both,
+        "remainder": g_remainder,
+    }
+
+
+def project_down_proj_gradient(
+    grad: Tensor,
+    and_output_basis: Optional[Tensor],
+    or_input_basis: Optional[Tensor],
+    scales: FamilyScales,
+) -> Tensor:
+    parts = decompose_down_proj_gradient(grad, and_output_basis, or_input_basis)
+    return (
+        scales.and_transport * parts["and_transport_only"]
+        + scales.or_update * parts["or_update_only"]
+        + scales.coupling_scale() * parts["coupling"]
+        + scales.remainder * parts["remainder"]
+    )
+
+
+def support_project_down_proj_gradient(
+    grad: Tensor,
+    and_output_basis: Optional[Tensor],
+    or_input_basis: Optional[Tensor],
+) -> Tensor:
+    """Project a down-proj update onto typed Coppola support, dropping remainder."""
+    parts = decompose_down_proj_gradient(grad, and_output_basis, or_input_basis)
+    return parts["and_transport_only"] + parts["or_update_only"] + parts["coupling"]
+
+
+def down_proj_component_norms(
+    grad: Tensor,
+    and_output_basis: Optional[Tensor],
+    or_input_basis: Optional[Tensor],
+) -> ComponentNorms:
+    parts = decompose_down_proj_gradient(grad, and_output_basis, or_input_basis)
+    return ComponentNorms(
+        and_transport=float(parts["and_transport_only"].norm().item()),
+        or_update=float(parts["or_update_only"].norm().item()),
+        coupling=float(parts["coupling"].norm().item()),
+        remainder=float(parts["remainder"].norm().item()),
+    )
+
+
+def project_attention_output_gradient(
+    grad: Tensor,
+    routing_head_bases: Sequence[Tensor],
+    scales: FamilyScales,
+) -> Tensor:
+    """Project an attention output gradient head-by-head on the left."""
+    if grad.ndim != 2:
+        raise ValueError(f"gradient must be rank-2, got {tuple(grad.shape)}")
+    if not routing_head_bases:
+        return grad
+    hidden, width = grad.shape
+    if width % len(routing_head_bases) != 0:
+        raise ValueError(f"gradient width {width} incompatible with {len(routing_head_bases)} head bases")
+    head_dim = width // len(routing_head_bases)
+    out = torch.empty_like(grad)
+    for h, basis in enumerate(routing_head_bases):
+        cols = grad[:, h * head_dim:(h + 1) * head_dim]
+        routed = _project_left(cols, basis)
+        remainder = cols - routed
+        out[:, h * head_dim:(h + 1) * head_dim] = scales.routing * routed + scales.remainder * remainder
+    return out
+
+
+def project_attention_qk_gradient(
+    grad: Tensor,
+    q_head_bases: Sequence[Tensor],
+    k_head_bases: Sequence[Tensor],
+    scales: FamilyScales,
+) -> Tensor:
+    """Project fused Q/K/V gradient on Q and K row blocks only, leaving V unchanged."""
+    if grad.ndim != 2:
+        raise ValueError(f"gradient must be rank-2, got {tuple(grad.shape)}")
+    if not q_head_bases and not k_head_bases:
+        return grad
+    hidden3, width = grad.shape
+    if hidden3 % 3 != 0:
+        raise ValueError(f"expected fused qkv height divisible by 3, got {hidden3}")
+    hidden = hidden3 // 3
+    if q_head_bases and hidden % len(q_head_bases) != 0:
+        raise ValueError(f"q block height {hidden} incompatible with {len(q_head_bases)} q bases")
+    if k_head_bases and hidden % len(k_head_bases) != 0:
+        raise ValueError(f"k block height {hidden} incompatible with {len(k_head_bases)} k bases")
+    out = grad.clone()
+    if q_head_bases:
+        head_dim = hidden // len(q_head_bases)
+        for h, basis in enumerate(q_head_bases):
+            rows = grad[h * head_dim:(h + 1) * head_dim, :]
+            routed = _project_left(rows, basis)
+            remainder = rows - routed
+            out[h * head_dim:(h + 1) * head_dim, :] = scales.routing * routed + scales.remainder * remainder
+    if k_head_bases:
+        head_dim = hidden // len(k_head_bases)
+        offset = hidden
+        for h, basis in enumerate(k_head_bases):
+            start = offset + h * head_dim
+            stop = offset + (h + 1) * head_dim
+            rows = grad[start:stop, :]
+            routed = _project_left(rows, basis)
+            remainder = rows - routed
+            out[start:stop, :] = scales.routing * routed + scales.remainder * remainder
+    return out
+
+
+def support_project_attention_output_gradient(
+    grad: Tensor,
+    routing_head_bases: Sequence[Tensor],
+) -> Tensor:
+    """Project an attention-output update onto routing support, dropping remainder."""
+    if grad.ndim != 2:
+        raise ValueError(f"gradient must be rank-2, got {tuple(grad.shape)}")
+    if not routing_head_bases:
+        return grad
+    hidden, width = grad.shape
+    if width % len(routing_head_bases) != 0:
+        raise ValueError(f"gradient width {width} incompatible with {len(routing_head_bases)} head bases")
+    head_dim = width // len(routing_head_bases)
+    out = torch.empty_like(grad)
+    for h, basis in enumerate(routing_head_bases):
+        cols = grad[:, h * head_dim:(h + 1) * head_dim]
+        out[:, h * head_dim:(h + 1) * head_dim] = _project_left(cols, basis)
+    return out
+
+
+def support_project_attention_qk_gradient(
+    grad: Tensor,
+    q_head_bases: Sequence[Tensor],
+    k_head_bases: Sequence[Tensor],
+) -> Tensor:
+    """Project fused Q/K/V update on Q and K support only, leaving V unchanged."""
+    if grad.ndim != 2:
+        raise ValueError(f"gradient must be rank-2, got {tuple(grad.shape)}")
+    if not q_head_bases and not k_head_bases:
+        return grad
+    hidden3, width = grad.shape
+    if hidden3 % 3 != 0:
+        raise ValueError(f"expected fused qkv height divisible by 3, got {hidden3}")
+    hidden = hidden3 // 3
+    out = grad.clone()
+    if q_head_bases:
+        head_dim = hidden // len(q_head_bases)
+        for h, basis in enumerate(q_head_bases):
+            rows = grad[h * head_dim:(h + 1) * head_dim, :]
+            out[h * head_dim:(h + 1) * head_dim, :] = _project_left(rows, basis)
+    if k_head_bases:
+        head_dim = hidden // len(k_head_bases)
+        offset = hidden
+        for h, basis in enumerate(k_head_bases):
+            start = offset + h * head_dim
+            stop = offset + (h + 1) * head_dim
+            rows = grad[start:stop, :]
+            out[start:stop, :] = _project_left(rows, basis)
+    return out
+
+
+def attention_component_norms(grad: Tensor, routing_head_bases: Sequence[Tensor]) -> ComponentNorms:
+    if not routing_head_bases:
+        return ComponentNorms(routing=float(grad.norm().item()))
+    hidden, width = grad.shape
+    head_dim = width // len(routing_head_bases)
+    routed_norm_sq = 0.0
+    remainder_norm_sq = 0.0
+    for h, basis in enumerate(routing_head_bases):
+        cols = grad[:, h * head_dim:(h + 1) * head_dim]
+        routed = _project_left(cols, basis)
+        remainder = cols - routed
+        routed_norm_sq += float(routed.square().sum().item())
+        remainder_norm_sq += float(remainder.square().sum().item())
+    return ComponentNorms(
+        routing=routed_norm_sq**0.5,
+        remainder=remainder_norm_sq**0.5,
+    )
+
+
+def attention_qk_component_norms(
+    grad: Tensor,
+    q_head_bases: Sequence[Tensor],
+    k_head_bases: Sequence[Tensor],
+) -> ComponentNorms:
+    if not q_head_bases and not k_head_bases:
+        return ComponentNorms(routing=float(grad.norm().item()))
+    hidden3, width = grad.shape
+    if hidden3 % 3 != 0:
+        raise ValueError(f"expected fused qkv height divisible by 3, got {hidden3}")
+    hidden = hidden3 // 3
+    routed_norm_sq = 0.0
+    remainder_norm_sq = 0.0
+    if q_head_bases:
+        head_dim = hidden // len(q_head_bases)
+        for h, basis in enumerate(q_head_bases):
+            rows = grad[h * head_dim:(h + 1) * head_dim, :]
+            routed = _project_left(rows, basis)
+            remainder = rows - routed
+            routed_norm_sq += float(routed.square().sum().item())
+            remainder_norm_sq += float(remainder.square().sum().item())
+    if k_head_bases:
+        head_dim = hidden // len(k_head_bases)
+        offset = hidden
+        for h, basis in enumerate(k_head_bases):
+            start = offset + h * head_dim
+            stop = offset + (h + 1) * head_dim
+            rows = grad[start:stop, :]
+            routed = _project_left(rows, basis)
+            remainder = rows - routed
+            routed_norm_sq += float(routed.square().sum().item())
+            remainder_norm_sq += float(remainder.square().sum().item())
+    return ComponentNorms(
+        routing=routed_norm_sq**0.5,
+        remainder=remainder_norm_sq**0.5,
+    )
+
+
+def default_zone_scales(layer_idx: int, n_layer: int) -> FamilyScales:
+    """Default full-Coppola depth policy for from-scratch pretraining."""
+    frac = layer_idx / max(n_layer - 1, 1)
+    if frac < 0.20:
+        return FamilyScales(routing=1.40, and_transport=1.10, or_update=0.70, remainder=0.0, coupling=0.90)
+    if frac < 0.65:
+        return FamilyScales(routing=0.90, and_transport=1.40, or_update=0.90, remainder=0.0, coupling=1.15)
+    if frac < 0.90:
+        return FamilyScales(routing=0.75, and_transport=0.95, or_update=1.40, remainder=0.0, coupling=1.15)
+    return FamilyScales(routing=0.60, and_transport=0.85, or_update=1.15, remainder=0.0, coupling=1.00)
+
+
+@dataclass
+class CoppolaPretrainingConfig:
+    n_head: int
+    attn_output_policy: RankPolicy = field(default_factory=RankPolicy)
+    attn_qk_policy: RankPolicy = field(default_factory=RankPolicy)
+    mlp_output_policy: RankPolicy = field(default_factory=RankPolicy)
+    mlp_input_policy: RankPolicy = field(default_factory=RankPolicy)
+    basis_update_interval: int = 250
+    uniform_scales: FamilyScales = field(default_factory=FamilyScales)
+    scale_fn: Callable[[int, int], FamilyScales] = default_zone_scales
+
+
+class CoppolaPretrainingController:
+    """Refreshes typed bases and projects gradients for nanochat-style models."""
+
+    def __init__(self, config: CoppolaPretrainingConfig):
+        self.config = config
+        self.layer_bases: Dict[int, LayerBases] = {}
+
+    def refresh_from_model(self, model) -> None:
+        layers = self._resolve_layers(model)
+        n_layer = len(layers)
+        bases: Dict[int, LayerBases] = {}
+        for layer_idx, block in enumerate(layers):
+            attn_weight = self._resolve_attn_out_weight(block)
+            q_weight, k_weight = self._resolve_attn_qk_weights(block)
+            mlp_weight = self._resolve_mlp_down_weight(block)
+            routing = compute_attention_output_bases(
+                attn_weight, self.config.n_head, self.config.attn_output_policy
+            )
+            q_bases: List[Tensor] = []
+            k_bases: List[Tensor] = []
+            if q_weight is not None and k_weight is not None:
+                q_bases = compute_attention_score_bases(
+                    q_weight, self.config.n_head, self.config.attn_qk_policy
+                )
+                k_bases = compute_attention_score_bases(
+                    k_weight, self.config.n_head, self.config.attn_qk_policy
+                )
+            and_basis, or_basis = compute_down_proj_bases(
+                mlp_weight,
+                output_policy=self.config.mlp_output_policy,
+                input_policy=self.config.mlp_input_policy,
+            )
+            bases[layer_idx] = LayerBases(
+                routing_head_bases=routing,
+                q_head_bases=q_bases,
+                k_head_bases=k_bases,
+                and_output_basis=and_basis,
+                or_input_basis=or_basis,
+            )
+        self.layer_bases = bases
+        self._n_layer = n_layer
+
+    def scales_for_layer(self, layer_idx: int, mode: str = "zoned") -> FamilyScales:
+        if mode == "uniform":
+            return self.config.uniform_scales
+        if not hasattr(self, "_n_layer"):
+            raise RuntimeError("refresh_from_model() must be called before requesting zoned scales")
+        return self.config.scale_fn(layer_idx, self._n_layer)
+
+    def project_attn_out_grad(self, layer_idx: int, grad: Tensor, mode: str = "zoned") -> Tensor:
+        bases = self.layer_bases[layer_idx]
+        scales = self.scales_for_layer(layer_idx, mode=mode)
+        return project_attention_output_gradient(grad, bases.routing_head_bases, scales)
+
+    def project_attn_qk_grad(self, layer_idx: int, grad: Tensor, mode: str = "zoned") -> Tensor:
+        bases = self.layer_bases[layer_idx]
+        scales = self.scales_for_layer(layer_idx, mode=mode)
+        return project_attention_qk_gradient(grad, bases.q_head_bases, bases.k_head_bases, scales)
+
+    def project_mlp_down_grad(self, layer_idx: int, grad: Tensor, mode: str = "zoned") -> Tensor:
+        bases = self.layer_bases[layer_idx]
+        scales = self.scales_for_layer(layer_idx, mode=mode)
+        return project_down_proj_gradient(grad, bases.and_output_basis, bases.or_input_basis, scales)
+
+    def project_attn_out_support(self, layer_idx: int, grad: Tensor) -> Tensor:
+        bases = self.layer_bases[layer_idx]
+        return support_project_attention_output_gradient(grad, bases.routing_head_bases)
+
+    def project_attn_qk_support(self, layer_idx: int, grad: Tensor) -> Tensor:
+        bases = self.layer_bases[layer_idx]
+        return support_project_attention_qk_gradient(grad, bases.q_head_bases, bases.k_head_bases)
+
+    def project_mlp_down_support(self, layer_idx: int, grad: Tensor) -> Tensor:
+        bases = self.layer_bases[layer_idx]
+        return support_project_down_proj_gradient(grad, bases.and_output_basis, bases.or_input_basis)
+
+    def project_after_matrix_transform(
+        self,
+        layer_idx: int,
+        grad: Tensor,
+        param_kind: str,
+        transform: Callable[[Tensor], Tensor],
+        mode: str = "zoned",
+    ) -> Tensor:
+        """Project -> transform -> re-project for Muon-style optimizers."""
+        g = self.project_gradient(layer_idx, grad, param_kind=param_kind, mode=mode)
+        g = transform(g)
+        return self.project_gradient(layer_idx, g, param_kind=param_kind, mode=mode)
+
+    def project_gradient(self, layer_idx: int, grad: Tensor, param_kind: str, mode: str = "zoned") -> Tensor:
+        if param_kind == "attn_out":
+            return self.project_attn_out_grad(layer_idx, grad, mode=mode)
+        if param_kind == "attn_qk":
+            return self.project_attn_qk_grad(layer_idx, grad, mode=mode)
+        if param_kind == "mlp_down":
+            return self.project_mlp_down_grad(layer_idx, grad, mode=mode)
+        raise ValueError(f"unknown param_kind={param_kind!r}")
+
+    def project_update_support(self, layer_idx: int, grad: Tensor, param_kind: str) -> Tensor:
+        if param_kind == "attn_out":
+            return self.project_attn_out_support(layer_idx, grad)
+        if param_kind == "attn_qk":
+            return self.project_attn_qk_support(layer_idx, grad)
+        if param_kind == "mlp_down":
+            return self.project_mlp_down_support(layer_idx, grad)
+        raise ValueError(f"unknown param_kind={param_kind!r}")
+
+    def gradient_component_norms(self, layer_idx: int, grad: Tensor, param_kind: str) -> ComponentNorms:
+        bases = self.layer_bases[layer_idx]
+        if param_kind == "attn_out":
+            return attention_component_norms(grad, bases.routing_head_bases)
+        if param_kind == "attn_qk":
+            return attention_qk_component_norms(grad, bases.q_head_bases, bases.k_head_bases)
+        if param_kind == "mlp_down":
+            return down_proj_component_norms(grad, bases.and_output_basis, bases.or_input_basis)
+        raise ValueError(f"unknown param_kind={param_kind!r}")
+
+    @staticmethod
+    def _resolve_layers(model):
+        if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
+            return list(model.transformer.h)
+        if hasattr(model, "model") and hasattr(model.model, "layers"):
+            return list(model.model.layers)
+        raise ValueError("could not resolve transformer layers from model")
+
+    @staticmethod
+    def _resolve_attn_out_weight(block) -> Tensor:
+        if hasattr(block, "attn") and hasattr(block.attn, "c_proj"):
+            return block.attn.c_proj.weight
+        if hasattr(block, "self_attn") and hasattr(block.self_attn, "o_proj"):
+            return block.self_attn.o_proj.weight
+        raise ValueError("could not resolve attention output weight")
+
+    @staticmethod
+    def _resolve_attn_qk_weights(block) -> Tuple[Optional[Tensor], Optional[Tensor]]:
+        if hasattr(block, "attn") and hasattr(block.attn, "c_attn"):
+            weight = block.attn.c_attn.weight
+            if weight.shape[0] % 3 != 0:
+                raise ValueError(f"expected fused qkv height divisible by 3, got {weight.shape[0]}")
+            hidden = weight.shape[0] // 3
+            return weight[:hidden, :], weight[hidden:2 * hidden, :]
+        if hasattr(block, "self_attn") and hasattr(block.self_attn, "q_proj") and hasattr(block.self_attn, "k_proj"):
+            return block.self_attn.q_proj.weight, block.self_attn.k_proj.weight
+        return None, None
+
+    @staticmethod
+    def _resolve_mlp_down_weight(block) -> Tensor:
+        if hasattr(block, "mlp") and hasattr(block.mlp, "c_proj"):
+            return block.mlp.c_proj.weight
+        if hasattr(block, "mlp") and hasattr(block.mlp, "down_proj"):
+            return block.mlp.down_proj.weight
+        raise ValueError("could not resolve MLP down projection weight")

divcot2_encoding.py

@@ -0,0 +1,143 @@
+"""Generalized long-division reduction CoT — base-B limbs, estimated quotient,
+optional scratchpad. Scales divcot_encoding.py (validated at tier 2, base 10)
+to higher tiers.
+
+Three levers over the validated v1:
+
+1. **Base-B limbs** (B=100 halves W): each limb is one byte (limbs.py), so
+   tier 5's W=20 decimal becomes 10 limbs, tier 7's W=78 becomes 39.
+2. **Estimated-quotient install** (`qhat`): the classic scaling trick. The
+   one-shot quotient digit qd = floor((r*B+d)/p) needs a W-limb comparison —
+   fine at W=3, strained at W>=5. Humans estimate qd from the LEADING limbs
+   of numerator and divisor (a constant-size table independent of W), then
+   correct by ±1. We install that estimate at an early block so the circuit
+   the model learns is estimate-then-correct, not a W-limb divide.
+3. **Scratchpad mode** (`scratch=True`): for large W the model can't hold the
+   running remainder internally — EMIT it (LSB-first, so the subtraction
+   borrow chain is local). Cost: O(W^2) tokens vs O(W) compact; use compact
+   while it trains, scratch when it stops.
+
+Surface (compact):   "N" N(2W,MSB) "M" p(W,MSB) "=" Q(2W) "R" ans(W,LSB) "\n"
+Surface (scratch):   ... "=" [qd r(W,LSB)]*2W "R" ans(W,LSB) "\n"
+
+Install targets (NTP-aligned: label at position t supervises the residual
+predicting byte t+1):
+    rem{j}  compact only — MSB limb j of the running remainder at each
+            quotient-emission position (the division carry/state)
+    qhat    leading-limbs quotient estimate at each quotient position
+    ans     limb k of N mod p at each answer position
+"""
+from __future__ import annotations
+
+from limbs import limb_char, limb_str, parse_limbs, to_limbs
+
+NMARK, DIVMARK, EQ, REVMARK, NL = "N", "M", "=", "R", "\n"
+
+
+def long_division(N: int, p: int, W: int, base: int):
+    """(quotient_limbs[2W] MSB-first, remainders[2W], answer)."""
+    q_limbs, rems, r = [], [], 0
+    for d in to_limbs(N, 2 * W, base, msb_first=True):
+        r = r * base + d
+        qd = r // p          # single limb 0..base-1 (r < base*p)
+        r -= qd * p
+        q_limbs.append(qd)
+        rems.append(r)
+    return q_limbs, rems, r  # r == N % p
+
+
+def qhat_estimate(r_prev: int, d: int, p: int, base: int) -> int:
+    """Leading-limbs estimate of floor((r*B+d)/p): numerator's top two limbs
+    over divisor's top limb — constant-size lookup regardless of W."""
+    num = r_prev * base + d
+    if num < p:
+        return 0
+    nw = 1
+    while base ** nw <= num:
+        nw += 1
+    pw = 1
+    while base ** pw <= p:
+        pw += 1
+    n_top = num // base ** max(0, nw - 2)       # top 2 limbs of numerator
+    p_top = p // base ** (pw - 1)               # top 1 limb of divisor
+    shift = (nw - 2) - (pw - 1)
+    est = (n_top // p_top) * base ** shift if shift >= 0 else n_top // (p_top * base ** -shift)
+    return min(base - 1, max(0, est))
+
+
+def prompt_str(N: int, p: int, W: int, base: int) -> str:
+    return NMARK + limb_str(N, 2 * W, base, msb_first=True) \
+        + DIVMARK + limb_str(p, W, base, msb_first=True) + EQ
+
+
+def gen_len(W: int, scratch: bool = False) -> int:
+    steps = 2 * W * (1 + W) if scratch else 2 * W
+    return steps + 1 + W
+
+
+def build_example(N: int, p: int, W: int, base: int, scratch: bool = False):
+    q_limbs, rems, answer = long_division(N, p, W, base)
+    n_msb = to_limbs(N, 2 * W, base, msb_first=True)
+    prompt = prompt_str(N, p, W, base)
+    parts, ann_q = [], []                      # ann_q: (char_index, var, val)
+    pos = len(prompt)
+    r_prev = 0
+    for i, (qd, r) in enumerate(zip(q_limbs, rems)):
+        ann_q.append((pos, "qhat", qhat_estimate(r_prev, n_msb[i], p, base)))
+        if not scratch:                        # compact: remainder is internal state
+            rstr = to_limbs(r, W, base, msb_first=True)
+            for j in range(W):
+                ann_q.append((pos, f"rem{j}", rstr[j]))
+        parts.append(limb_char(qd))
+        pos += 1
+        if scratch:                            # emit remainder LSB-first (local borrow)
+            parts.append(limb_str(r, W, base))
+            pos += W
+        r_prev = r
+    ans_limbs = to_limbs(answer, W, base)      # LSB-first
+    parts.append(REVMARK + limb_str(answer, W, base) + NL)
+    text = prompt + "".join(parts)
+    ann = [dict() for _ in range(len(text))]
+    for idx, var, val in ann_q:
+        ann[idx - 1][var] = int(val)           # NTP alignment
+    ans_base = pos + 1                         # after 'R'
+    for k in range(W):
+        ann[ans_base + k - 1]["ans"] = ans_limbs[k]
+    return text, ann
+
+
+def var_specs(W: int, base: int, scratch: bool = False):
+    """(name, n_classes) install vars; assign blocks in the trainer."""
+    specs = [("ans", base), ("qhat", base)]
+    if not scratch:
+        specs += [(f"rem{j}", base) for j in range(W)]
+    return specs
+
+
+def decode_answer(gen_chars: str, W: int, base: int, scratch: bool = False) -> int:
+    off = (2 * W * (1 + W) if scratch else 2 * W) + 1   # skip quotient(+rems) and 'R'
+    return parse_limbs(gen_chars[off:off + W], base)
+
+
+if __name__ == "__main__":
+    import random
+    rng = random.Random(1)
+    for base in (10, 100):
+        for scratch in (False, True):
+            for _ in range(4000):
+                W = rng.randint(1, 8)
+                p = rng.randrange(max(2, base ** (W - 1)), base ** W)
+                N = rng.randrange(p * p) if p > 1 else 0
+                text, ann = build_example(N, p, W, base, scratch)
+                plen = len(prompt_str(N, p, W, base))
+                assert len(text) == plen + gen_len(W, scratch) + 1, (len(text), plen, gen_len(W, scratch))  # +1: NL
+                assert decode_answer(text[plen:], W, base, scratch) == N % p
+                assert len(ann) == len(text)
+            print(f"divcot2 base={base} scratch={scratch}: 4000/4000 decode OK")
+    # tier-relevant lengths
+    print("\ntokens/example (prompt+gen):")
+    for tier, Wd in [(3, 5), (4, 10), (5, 20), (6, 39), (7, 78), (8, 155)]:
+        for base, W in ((10, Wd), (100, (Wd + 1) // 2)):
+            plen = 3 * W + 3
+            print(f"  tier {tier} base {base:>3} W={W:>3}  compact {plen + gen_len(W)}"
+                  f"  scratch {plen + gen_len(W, True)}")

encoding.py

@@ -0,0 +1,97 @@
+"""Shared sequence encoding for the modmul BP-install model.
+
+The model learns the field-multiplication map  (p, x, y) -> (x*y) mod p  with
+x, y in [0, p).  At inference the legal two-operand reduction a%p, b%p produces
+x, y (the same step the baselines use); training samples x, y in [0, p) directly,
+which matches the (a mod p) distribution for p << operand range.
+
+Fixed-width, base-10, reverse-LSB answer (the proven arithmetic recipe). The
+prime p is an explicit in-sequence conditioning field so one model generalises
+across primes:
+
+    p007*012,003=R120\n        # W=3 example: p=7, x=12%7=5? no -- x,y already <p
+                               # e.g. p=7,x=5,y=4 -> 20%7=6 -> "p7*5,4=R6\n" (W=1)
+
+Layout (width W = max decimal digits of any prime in scope):
+    "p" P(W) "*" X(W) "," Y(W) "=R" REV_ANS(W) "\n"
+zero-padded fields; REV_ANS is the W-digit answer reversed (least-significant
+digit first), so digit k is emitted at a fixed position and the carry flows
+left-to-right the way reverse-LSB addition does.
+"""
+from __future__ import annotations
+
+MUL, SEP, EQ, REVMARK, NL = "*", ",", "=", "R", "\n"
+NEWLINE_ID = ord("\n")
+
+
+def width_for_primes(primes) -> int:
+    """Decimal digits needed to hold the largest residue (p-1)."""
+    return max(len(str(p - 1)) for p in primes)
+
+
+def prompt_str(p: int, x: int, y: int, W: int) -> str:
+    return f"p{p:0{W}d}*{x:0{W}d},{y:0{W}d}={REVMARK}"
+
+
+def build_example(p: int, x: int, y: int, W: int):
+    """Return (text, ann) where ann[t] holds the install label for the residual
+    at byte position t (which predicts byte t+1), NTP-aligned."""
+    ans = (x * y) % p
+    rev_ans = f"{ans:0{W}d}"[::-1]            # LSB-first, width W
+    prompt = prompt_str(p, x, y, W)
+    text = prompt + rev_ans + NL
+    ann = [dict() for _ in range(len(text))]
+    base = len(prompt)                        # first answer char index
+    for k in range(W):                        # answer digit k (LSB-first)
+        pos = base + k
+        ann[pos - 1]["ans"] = int(rev_ans[k])  # residual at pos-1 predicts it
+    return text, ann
+
+
+def answer_len(W: int) -> int:
+    return W
+
+
+# --- Scratchpad ("chain-of-thought") inference variant -------------------
+# The model GENERATES the product as reverse-LSB digits before the answer:
+#     "p" P(W) "*" X(W) "," Y(W) "=" REV_PROD(2W) "R" REV_ANS(W) "\n"
+# At inference we feed up to '=' and greedily decode the full scratchpad; the
+# answer is the last W generated digits. Python never computes x*y -- the model
+# emits the product digits (a learned circuit). Only a%p, b%p is done in code.
+
+def prompt_str_sp(p: int, x: int, y: int, W: int) -> str:
+    return f"p{p:0{W}d}*{x:0{W}d},{y:0{W}d}{EQ}"
+
+
+def scratchpad_len(W: int) -> int:
+    return 2 * W + 1 + W
+
+
+def decode_answer_sp(gen_chars: str, W: int) -> int:
+    """Generated suffix is 2W product digits, 'R', then W answer digits."""
+    return decode_answer(gen_chars[2 * W + 1 : 2 * W + 1 + W], W)
+
+
+# --- Schoolbook ("partial products") inference variant -------------------
+# Generated suffix: W partial products (W+1 wide), 'P', 2W product digits,
+# 'R', W answer digits. The model emits the partial products and sums them
+# (learned circuit); the answer is the final W digits.
+
+def school_gen_len(W: int) -> int:
+    return W * (W + 1) + 1 + 2 * W + 1 + W
+
+
+def decode_answer_school(gen_chars: str, W: int) -> int:
+    start = W * (W + 1) + 1 + 2 * W + 1
+    return decode_answer(gen_chars[start:start + W], W)
+
+
+def decode_answer(gen_chars: str, W: int) -> int:
+    """Greedy-generated W chars are the reverse-LSB answer digits. Non-digits
+    (untrained junk) count as 0. Returns the integer answer."""
+    val = 0
+    for k in range(min(W, len(gen_chars))):
+        c = gen_chars[k]
+        d = int(c) if "0" <= c <= "9" else 0  # ASCII only ('³'.isdigit() is True but int() fails)
+        val += d * (10 ** k)
+    return val

kvgen.py

@@ -0,0 +1,172 @@
+"""KV-cached batched greedy generation for ByteGPT — the wall-clock fix.
+
+The naive loop re-forwards the full prefix for every generated token
+(O(T·L²) attention work); harmless at tier-2's ~20-token CoT, fatal for
+tier-3+ chains-of-thought (hundreds-thousands of tokens) under the
+challenge's 5-min/1100-problem inference budget.
+
+This module reuses the model's OWN weights and norm functions — it is a
+faster schedule for the same computation, not a different model. Prefill
+captures each block's K/V via a forward hook on `attn.c_attn` (one ordinary
+forward over the prompt, which also yields the first generated token); each
+subsequent token does one single-position pass per block against the cache.
+
+Supported config (asserted): attn_kind=mha, attn_norm in {entmax15, softmax},
+position_encoding in {alibi, learned, none}, no block gates / loop blocks /
+attention residuals. Validated token-identical against the naive loop in
+__main__ (random-init smoke + optional real checkpoint).
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import torch
+
+_HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(_HERE))                      # submission layout (vendored)
+for _anc in Path(__file__).resolve().parents:       # dev layout (repo experiments/)
+    if (_anc / "experiments" / "coppola_pretrain_tiny.py").exists():
+        sys.path.insert(0, str(_anc / "experiments"))
+        break
+
+import coppola_pretrain_tiny as cpt  # noqa: E402
+
+
+def _norm_row(scores: torch.Tensor, mode: str) -> torch.Tensor:
+    """Attention norm over one query row (all cached keys are valid: no mask)."""
+    if mode == "softmax":
+        return torch.softmax(scores, dim=-1)
+    if mode == "entmax15":
+        return (cpt._pkg_entmax15(scores, dim=-1) if cpt._HAVE_ENTMAX_PKG
+                else cpt._entmax15(scores, dim=-1))
+    raise ValueError(f"kvgen: unsupported attn_norm {mode!r}")
+
+
+def _assert_supported(model) -> None:
+    cfg = model.config
+    assert cfg.attn_kind == "mha", f"kvgen: attn_kind={cfg.attn_kind!r}"
+    assert cfg.attn_norm in ("entmax15", "softmax"), f"kvgen: attn_norm={cfg.attn_norm!r}"
+    assert cfg.position_encoding in ("alibi", "learned", "none"), \
+        f"kvgen: position_encoding={cfg.position_encoding!r}"
+    assert cfg.res_attn == "none", "kvgen: res_attn unsupported"
+    assert not cfg.loop_block_indices and not cfg.loop_unit_indices, "kvgen: loop blocks unsupported"
+    assert model._gate_values() is None, "kvgen: block gates unsupported"
+    assert cfg.dropout == 0.0 or not model.training, "kvgen: eval mode required"
+
+
+@torch.no_grad()
+def generate_kv(model, ids: torch.Tensor, n_gen: int) -> torch.Tensor:
+    """Append n_gen greedy tokens to each row of ids ([B, T0] long). Returns
+    the [B, n_gen] generated tokens. Total length must fit cfg.seq_len."""
+    _assert_supported(model)
+    cfg = model.config
+    B, T0 = ids.shape
+    assert T0 + n_gen <= cfg.seq_len, f"kvgen: {T0}+{n_gen} exceeds seq_len {cfg.seq_len}"
+    device = ids.device
+    blocks = model.transformer.h
+    L = len(blocks)
+    nh = [blk.attn.n_head for blk in blocks]
+    hd = [blk.attn.head_dim for blk in blocks]
+
+    # ---- prefill: one ordinary forward over the prompt, K/V captured by hook
+    kv: list[tuple[torch.Tensor, torch.Tensor] | None] = [None] * L
+
+    def _mk_hook(i):
+        def hook(_mod, _inp, qkv):
+            _q, k, v = qkv.chunk(3, dim=-1)
+            kv[i] = (k.view(B, -1, nh[i], hd[i]).transpose(1, 2).contiguous(),
+                     v.view(B, -1, nh[i], hd[i]).transpose(1, 2).contiguous())
+        return hook
+
+    handles = [blk.attn.c_attn.register_forward_hook(_mk_hook(i))
+               for i, blk in enumerate(blocks)]
+    try:
+        logits, _ = model(ids)
+    finally:
+        for h in handles:
+            h.remove()
+    out = [logits[:, -1].argmax(dim=-1)]            # first generated token
+
+    slopes = (model._alibi_slopes.to(device) if cfg.position_encoding == "alibi" else None)
+
+    # ---- incremental steps: one single-position pass per block per token
+    for step in range(1, n_gen):
+        t = T0 + step - 1                            # position of the token we feed in
+        x = model.transformer.wte(out[-1]).unsqueeze(1)         # [B, 1, E]
+        if cfg.position_encoding == "learned":
+            x = x + model.transformer.wpe(torch.tensor([t], device=device))[None, :, :]
+        for i, blk in enumerate(blocks):
+            h = blk.ln_1(x)
+            q, k, v = blk.attn.c_attn(h).chunk(3, dim=-1)
+            q = q.view(B, 1, nh[i], hd[i]).transpose(1, 2)      # [B, nh, 1, hd]
+            k = k.view(B, 1, nh[i], hd[i]).transpose(1, 2)
+            v = v.view(B, 1, nh[i], hd[i]).transpose(1, 2)
+            K = torch.cat([kv[i][0], k], dim=2)                 # [B, nh, t+1, hd]
+            V = torch.cat([kv[i][1], v], dim=2)
+            kv[i] = (K, V)
+            # [B, nh, 1, t+1]: dot of the single query with every cached key
+            scores = torch.einsum("bhqd,bhjd->bhqj", q, K) / (hd[i] ** 0.5)
+            if slopes is not None:                  # ALiBi row: slope * -(t - j)
+                j = torch.arange(K.size(2), device=device, dtype=scores.dtype)
+                scores = scores + slopes.view(1, -1, 1, 1) * (-(t - j).abs()).view(1, 1, 1, -1)
+            attn = _norm_row(scores, cfg.attn_norm)
+            y = torch.einsum("bhqj,bhjd->bhqd", attn, V)
+            y = y.transpose(1, 2).contiguous().view(B, 1, -1)
+            x = x + blk.attn.c_proj(y)
+            x = x + blk.mlp(blk.ln_2(x))
+        logits1 = model.lm_head(model.transformer.ln_f(x))      # [B, 1, vocab]
+        out.append(logits1[:, -1].argmax(dim=-1))
+    return torch.stack(out, dim=1)                  # [B, n_gen]
+
+
+@torch.no_grad()
+def generate_naive(model, ids: torch.Tensor, n_gen: int) -> torch.Tensor:
+    """Reference loop (full re-forward per token), for validation/timing."""
+    cap = model.config.seq_len
+    for _ in range(n_gen):
+        logits, _ = model(ids[:, -cap:])
+        ids = torch.cat([ids, logits[:, -1].argmax(dim=-1, keepdim=True)], dim=1)
+    return ids[:, -n_gen:]
+
+
+if __name__ == "__main__":
+    import argparse
+    import time
+
+    ap = argparse.ArgumentParser(description="validate kvgen vs naive loop")
+    ap.add_argument("--ckpt", default=None, help="optional trained checkpoint (.pt)")
+    ap.add_argument("--n-gen", type=int, default=40)
+    ap.add_argument("--batch", type=int, default=8)
+    args = ap.parse_args()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    if args.ckpt:
+        sys.path.insert(0, str(_HERE))
+        from train_arith_bp_supervised import TrainConfig, build_model
+        ck = torch.load(args.ckpt, map_location=device, weights_only=False)
+        model = build_model(TrainConfig(**ck["config"]), device)
+        model.load_state_dict(ck["state_dict"])
+    else:
+        cfg = cpt.GPTConfig(vocab_size=256, n_layer=4, n_head=4, n_embd=128,
+                            seq_len=256, dropout=0.0, attn_kind="mha",
+                            attn_norm="entmax15", position_encoding="alibi",
+                            loss_kind="bce")
+        torch.manual_seed(0)
+        model = cpt.ByteGPT(cfg).to(device)
+    model.eval()
+
+    torch.manual_seed(1)
+    ids = torch.randint(40, 70, (args.batch, 24), device=device)
+    t0 = time.time()
+    a = generate_naive(model, ids, args.n_gen)
+    t_naive = time.time() - t0
+    t0 = time.time()
+    b = generate_kv(model, ids, args.n_gen)
+    t_kv = time.time() - t0
+    same = (a == b).all().item()
+    n_diff = (a != b).sum().item()
+    print(f"identical={same}  diff_tokens={n_diff}/{a.numel()}  "
+          f"naive={t_naive:.2f}s  kv={t_kv:.2f}s  speedup={t_naive / max(t_kv, 1e-9):.1f}x")
+    if not same:
+        sys.exit(1)

limbs.py

@@ -0,0 +1,66 @@
+"""Limb codec shared by the scaled CoT encodings (divcot2, karatsuba).
+
+A limb is one digit in base B (10 or 100), encoded as a SINGLE byte
+0x80+value (128..227) so sequence length scales with limb count, not decimal
+digits. Base-100 halves W and roughly quarters quadratic CoT lengths.
+Strings are latin1-safe: encode("latin1") round-trips; never use ascii.
+"""
+from __future__ import annotations
+
+LIMB_OFFSET = 0x80  # limb v -> chr(0x80+v); supports base <= 128
+
+
+def limb_char(v: int) -> str:
+    return chr(LIMB_OFFSET + v)
+
+
+def char_limb(c: str) -> int:
+    """Inverse of limb_char; out-of-range chars decode to 0 (eval leniency)."""
+    v = ord(c) - LIMB_OFFSET
+    return v if v >= 0 else 0
+
+
+def to_limbs(n: int, width: int, base: int, msb_first: bool = False) -> list[int]:
+    out = []
+    for _ in range(width):
+        out.append(n % base)
+        n //= base
+    assert n == 0, "width too small for value"
+    return out[::-1] if msb_first else out
+
+
+def from_limbs(limbs, base: int, msb_first: bool = False) -> int:
+    seq = limbs[::-1] if msb_first else limbs
+    val = 0
+    for k, v in enumerate(seq):
+        val += v * base ** k
+    return val
+
+
+def limb_str(n: int, width: int, base: int, msb_first: bool = False) -> str:
+    return "".join(limb_char(v) for v in to_limbs(n, width, base, msb_first))
+
+
+def parse_limbs(s: str, base: int, msb_first: bool = False) -> int:
+    vals = [min(char_limb(c), base - 1) for c in s]
+    return from_limbs(vals, base, msb_first)
+
+
+def width_for(p_max: int, base: int) -> int:
+    w = 1
+    while base ** w <= p_max - 1:
+        w += 1
+    return w
+
+
+if __name__ == "__main__":
+    import random
+    rng = random.Random(0)
+    for base in (10, 100):
+        for _ in range(2000):
+            w = rng.randint(1, 12)
+            n = rng.randrange(base ** w)
+            for msb in (False, True):
+                s = limb_str(n, w, base, msb)
+                assert len(s) == w and parse_limbs(s, base, msb) == n
+    print("limbs.py: all round-trip tests pass")

manifest.json

@@ -0,0 +1,7 @@
+{
+  "entry_class": "model.ModMulRouter",
+  "output_base": 10,
+  "framework": "pytorch",
+  "model_description": "Prime-magnitude router over from-random-init ByteGPT members (entmax15+ALiBi, BCE+reg): tiers 1-2 scratchpad-CoT member; tier-3 composed multiply+long-division CoT member (model generates accumulator rows, restaged dividend, per-step qd*p and remainders, answer). Routing keys on the size of p only; all answers are generated digit sequences from trained weights.",
+  "training_description": "All members trained from random init with BCE(+reg) LM loss plus per-block BP-install probe losses on intermediate variables (carries, borrows, quotient estimates, remainder digits). Synthetic data from the official tier geometry with held-out seeds; tier-3 member trained on ~all ~6445 primes of the 9-16-bit range (full domain coverage), bit-uniform prime sampling, example-aligned batch rows. Randomising weights collapses accuracy. Training code: github.com/area9innovation/modmul-challenge (available to reviewers on request)."
+}

model.py

@@ -0,0 +1,208 @@
+"""Challenge-interface wrapper for the modmul BP-install model.
+
+Implements ModularMultiplicationModel. Inference:
+  1. per-argument: parse a, b, p to ints (each hook sees only its own arg)
+  2. predict_digits: reduce x=a%p, y=b%p (legal two-operand reduction), and if
+     p is within the model's trained width, greedy-decode the reverse-LSB answer
+     digits from the network; otherwise emit [0] (honest out-of-regime fallback).
+
+The answer comes entirely from the trained network on in-regime primes:
+randomising the weights collapses accuracy. output_base = 10.
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+# Self-contained: ByteGPT + build_model + the entmax15 recipe are VENDORED into
+# this directory (coppola_pretrain_tiny.py, coppola_pretraining.py,
+# train_arith_bp_supervised.py), so the submission loads with only `torch`
+# available and read access limited to its own dir (the eval sandbox contract).
+# entmax15 falls back to a local forward-exact impl when the `entmax` pip
+# package is absent.
+from train_arith_bp_supervised import TrainConfig, build_model  # noqa: E402
+import encoding as enc  # noqa: E402
+import composed_encoding as cenc  # noqa: E402  (composed multiply+reduce CoT)
+import kvgen  # noqa: E402  (KV-cached generation; falls back to naive loop)
+
+from modchallenge.interface.base_model import ModularMultiplicationModel  # noqa: E402
+
+
+class ModMulBP(ModularMultiplicationModel):
+    def __init__(self):
+        self.model = None
+        self.W = 1
+        self.device = None
+        self.regime = 10  # max p exclusive = 10**W
+        self.scratchpad = False
+        self.school = False
+
+    def load(self, model_dir: str, weights: str = "weights.pt") -> None:
+        torch.manual_seed(0)  # determinism is the model's responsibility (rules)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        ckpt = torch.load(Path(model_dir) / weights,
+                          map_location=self.device, weights_only=False)
+        tc = TrainConfig(**ckpt["config"])
+        self.model = build_model(tc, self.device)
+        self.model.load_state_dict(ckpt["state_dict"])
+        self.model.eval()
+        self.W = ckpt["W"]
+        self.scratchpad = bool(ckpt.get("scratchpad", False))
+        self.school = bool(ckpt.get("school", False))
+        self.composed = bool(ckpt.get("composed", False))
+        if self.composed:
+            self.base = ckpt["base"]
+            self.scratch = bool(ckpt.get("scratch", False))
+            self.cursor = bool(ckpt.get("cursor", False))
+            self.subpad = bool(ckpt.get("subpad", False))
+            self.regime = self.base ** self.W
+        else:
+            self.regime = 10 ** self.W
+        # Trained prime span (for router dispatch): derived from the ckpt's
+        # tier list via the official tier geometry; fallback = full regime.
+        self.p_lo, self.p_hi = 2, self.regime - 1
+        try:
+            from modchallenge.config import TIERS
+            spans = [(2 ** TIERS[t].min_bits, 2 ** TIERS[t].max_bits - 1)
+                     for t in ckpt.get("tiers", [])]
+            if spans:
+                self.p_lo = min(lo for lo, _ in spans)
+                self.p_hi = min(max(hi for _, hi in spans), self.regime - 1)
+        except Exception:
+            pass
+
+    # per-argument preprocessing (each sees only its own argument)
+    def preprocess_a(self, a: str) -> int:
+        return int(a)
+
+    def preprocess_b(self, b: str) -> int:
+        return int(b)
+
+    def preprocess_p(self, p: str) -> int:
+        return int(p)
+
+    @torch.no_grad()
+    def predict_digits(self, a_enc, b_enc, p_enc):
+        return self.predict_digits_batch([(a_enc, b_enc, p_enc)])[0]
+
+    @torch.no_grad()
+    def predict_digits_batch(self, inputs):
+        out = [[0]] * len(inputs)
+        if self.composed:
+            prompt_fn = lambda p, x, y, W: cenc.prompt_str(  # noqa: E731
+                x, y, p, W, self.base, self.subpad)
+            n_gen = cenc.gen_len(self.W, self.scratch, self.cursor, self.subpad)
+            decode_fn = lambda g, W: cenc.decode_answer(     # noqa: E731
+                g, W, self.base, self.scratch, self.cursor, self.subpad)
+        else:
+            cot = self.scratchpad or self.school
+            prompt_fn = enc.prompt_str_sp if cot else enc.prompt_str
+            n_gen = (enc.school_gen_len(self.W) if self.school
+                     else enc.scratchpad_len(self.W) if self.scratchpad
+                     else enc.answer_len(self.W))
+            decode_fn = (enc.decode_answer_school if self.school
+                         else enc.decode_answer_sp if self.scratchpad
+                         else enc.decode_answer)
+        prompts, idx = [], []
+        for i, (a, b, p) in enumerate(inputs):
+            if p >= self.regime:          # outside trained width -> honest 0
+                continue
+            x, y = a % p, b % p
+            prompts.append(prompt_fn(p, x, y, self.W))
+            idx.append((i, x, y, p))
+        if not prompts:
+            return out
+        # All in-regime prompts share the same length (fixed width) -> batchable.
+        # latin1: composed prompts carry limb bytes >127 (limbs.py codec).
+        ids = torch.tensor([list(s.encode("latin1")) for s in prompts],
+                           dtype=torch.long, device=self.device)
+        plen = len(prompts[0])
+        # KV-cached generation: the naive loop re-forwards the whole prefix per
+        # token, which blows the 5-min/1100-problem budget on long CoTs. kvgen
+        # is the same computation on the same weights, token-identical
+        # (validated); fall back to the naive loop on unsupported configs.
+        try:
+            gens = kvgen.generate_kv(self.model, ids, n_gen).tolist()
+        except AssertionError:
+            seq_cap = self.model.config.seq_len
+            for _k in range(n_gen):
+                logits, _ = self.model(ids[:, -seq_cap:])
+                nxt = logits[:, -1].argmax(dim=-1, keepdim=True)
+                ids = torch.cat([ids, nxt], dim=1)
+            gens = ids[:, plen:].tolist()
+        for row, (i, x, y, p) in zip(gens, idx):
+            gen = bytes(b & 0xFF for b in row).decode("latin1")
+            # No arithmetic touch-up of the model's answer: a decoded value >= p
+            # would be malformed (scored incorrect anyway), so emit the honest
+            # [0] fallback instead of clamping with % p.
+            ans = decode_fn(gen, self.W)
+            if 0 <= ans < p:
+                out[i] = [int(c) for c in str(ans)]
+        return out
+
+    def max_batch_size(self) -> int:
+        return 512
+
+
+class ModMulRouter(ModularMultiplicationModel):
+    """Routes each problem to the most specialized member model by prime
+    magnitude. Members are weights_r*.pt files (sorted name order); each is a
+    full ModMulBP checkpoint with its own trained regime. A problem goes to
+    the FIRST member whose regime covers its p; out-of-regime problems emit
+    the honest [0].
+
+    Compliance: routing keys on the SIZE of p only (per-argument
+    representation work, like base conversion); every answer comes from a
+    trained member's generated digits.
+    """
+
+    def __init__(self):
+        self.members: list[ModMulBP] = []
+
+    def load(self, model_dir: str) -> None:
+        torch.manual_seed(0)
+        for f in sorted(Path(model_dir).glob("weights_r*.pt")):
+            m = ModMulBP()
+            m.load(model_dir, weights=f.name)
+            self.members.append(m)
+        assert self.members, "router needs weights_r*.pt member checkpoints"
+        self.members.sort(key=lambda m: m.regime)  # most specialized first
+
+    def preprocess_a(self, a: str) -> int:
+        return int(a)
+
+    def preprocess_b(self, b: str) -> int:
+        return int(b)
+
+    def preprocess_p(self, p: str) -> int:
+        return int(p)
+
+    @torch.no_grad()
+    def predict_digits(self, a_enc, b_enc, p_enc):
+        return self.predict_digits_batch([(a_enc, b_enc, p_enc)])[0]
+
+    @torch.no_grad()
+    def predict_digits_batch(self, inputs):
+        out = [[0]] * len(inputs)
+        groups: dict[int, list[int]] = {}
+        for i, (_a, _b, p) in enumerate(inputs):
+            # Prefer the member whose TRAINED prime span contains p; fall back
+            # to the most specialized member whose regime merely covers it.
+            mi = next((k for k, m in enumerate(self.members)
+                       if m.p_lo <= p <= m.p_hi), None)
+            if mi is None:
+                mi = next((k for k, m in enumerate(self.members) if p < m.regime), None)
+            if mi is not None:
+                groups.setdefault(mi, []).append(i)
+        for mi, idxs in groups.items():
+            sub = [inputs[i] for i in idxs]
+            res = self.members[mi].predict_digits_batch(sub)
+            for i, r in zip(idxs, res):
+                out[i] = r
+        return out
+
+    def max_batch_size(self) -> int:
+        return 512

train_arith_bp_supervised.py

@@ -0,0 +1,633 @@
+#!/usr/bin/env python3
+"""BP-install supervised training on arithmetic.
+
+Adds per-block probe-loss to the standard NTP training:
+  total = NTP_loss + λ_install · sum_V CE(probe_V(pre_blk_target_V), V_labels)
+
+For arithmetic, BP variables are (carry_after_k, sum_digit_k) at known
+positions. Target blocks default to Phase-A install pattern:
+  - carry_after_k -> block 2
+  - sum_digit_k   -> block 3
+
+The supervised loss pushes the model to install each BP variable at its
+target block's INPUT (i.e., the residual just before block target_block).
+
+Hypothesis (from user): forcing single-block install via supervision
+will concentrate the per-head distributed signal from Phase A into a
+cleaner BP message.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import random
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from coppola_pretrain_tiny import (
+    ByteGPT, GPTConfig,
+    evaluate_arithmetic_accuracy,
+    bits_per_byte,
+)
+
+
+# ----------------------------- BP annotations ------------------------------
+
+def _make_addition_with_annotations(max_digits: int,
+                                      surface: str = "reverse"
+                                      ) -> Tuple[str, List[Dict[str, int]]]:
+    """Generate an addition problem in one of two surfaces, with per-position
+    BP-variable annotations placed at the NTP predicting residual.
+
+    Surfaces:
+      - "reverse": `a+b=R<ans_lsb_first>\n`, e.g. 12+34 -> "12+34=R64\n"
+      - "forward": `a+b=F<ans_msb_first>\n`, e.g. 12+34 -> "12+34=F46\n"
+
+    Annotation IDENTITY (carry_after_k, sum_digit_k) is invariant across
+    surfaces; only the POSITION at which the predicting residual sits
+    differs. This is the test substrate for surface-form invariance of
+    BP-install location.
+    """
+    if surface not in ("reverse", "forward"):
+        raise ValueError(f"surface must be 'reverse' or 'forward', got {surface!r}")
+    n_a = random.randint(1, max_digits)
+    n_b = random.randint(1, max_digits)
+    a = random.randint(0, 10 ** n_a - 1)
+    b = random.randint(0, 10 ** n_b - 1)
+    ans_int = a + b
+    ans_forward = str(ans_int)
+    ans_emit = ans_forward[::-1] if surface == "reverse" else ans_forward
+    marker = "R" if surface == "reverse" else "F"
+    text = f"{a}+{b}={marker}{ans_emit}\n"
+    a_str = str(a)
+    b_str = str(b)
+    plus_pos = len(a_str)
+    eq_pos = plus_pos + 1 + len(b_str)
+    marker_pos = eq_pos + 1
+    ans_start = marker_pos + 1
+    ans_len = len(ans_emit)
+    ann = [{} for _ in range(len(text))]
+
+    # Compute carries: carry_after_k = carry out of column k (binary).
+    max_n = max(n_a, n_b)
+    carry = 0
+    carries = []
+    for k in range(max_n + 1):
+        da = (a // (10 ** k)) % 10 if k < n_a else 0
+        db = (b // (10 ** k)) % 10 if k < n_b else 0
+        tot = da + db + carry
+        carry = tot // 10
+        carries.append(carry)
+
+    # Sum digits and carries placed at predicting residuals (NTP-aligned).
+    # For position-of-sum_digit_k's EMISSION:
+    #   reverse: sum_digit_k emitted at ans_start + k
+    #   forward: sum_digit_k emitted at ans_start + (ans_len - 1 - k)
+    # The annotation goes at emission_pos - 1.
+    # carry_after_k is needed to PREDICT sum_digit_(k+1), so it sits at the
+    # residual that PRECEDES sum_digit_(k+1)'s emission position.
+    def emission_pos_for_sum_k(k: int) -> int:
+        # k=0 is LSB; emission position depends on surface
+        if surface == "reverse":
+            return ans_start + k
+        # forward: most-significant emitted first
+        # max emitted digit index = ans_len - 1 corresponds to LSB if forward
+        # The kth LSB digit is the (ans_len - 1 - k)-th emitted in forward
+        return ans_start + (ans_len - 1 - k)
+
+    for k in range(ans_len):
+        emit = emission_pos_for_sum_k(k)
+        tgt = emit - 1
+        if 0 <= tgt < len(text):
+            # Extract the actual digit value at that emission position
+            digit_char = text[emit]
+            if digit_char.isdigit():
+                ann[tgt][f"sum_digit_{k}"] = int(digit_char)
+
+    # carry_after_k is needed to predict sum_digit_(k+1).
+    for k, c in enumerate(carries):
+        if k + 1 >= ans_len:
+            continue  # no next sum digit to predict
+        emit_next = emission_pos_for_sum_k(k + 1)
+        tgt = emit_next - 1
+        if 0 <= tgt < len(text):
+            ann[tgt][f"carry_after_{k}"] = c
+
+    return text, ann
+
+
+# Variables we supervise and their target blocks.
+# Per Phase A: carries install around block 2, sums around block 3.
+# For 4L model: blocks 0..3.
+def default_var_specs(max_digits: int, target_carry_block: int = 2,
+                        target_sum_block: int = 3
+                        ) -> List[Tuple[str, int, int]]:
+    """Return list of (var_name, n_classes, target_block_idx)."""
+    specs = []
+    # carry_after_k: 2-class (0 or 1) — for digits beyond highest column,
+    # carry could exceed 1 only in pathological cases, so binary is safe
+    # for 1-3 digit ops with values 0..999+0..999=0..1998.
+    for k in range(max_digits + 1):
+        specs.append((f"carry_after_{k}", 2, target_carry_block))
+    for k in range(max_digits + 1):
+        specs.append((f"sum_digit_{k}", 10, target_sum_block))
+    return specs
+
+
+# ----------------------------- Corpus --------------------------------------
+
+class ArithmeticBPCorpus:
+    """Arithmetic corpus emitting (x, y, bp_labels)."""
+
+    def __init__(self, max_digits: int, var_specs, surfaces=("reverse",),
+                 surface_weights=None):
+        """surfaces: tuple of surface names to sample from.
+        surface_weights: matching list of relative weights (default uniform)."""
+        self.max_digits = max_digits
+        self.var_specs = var_specs
+        self.var_names = [name for name, _, _ in var_specs]
+        self.surfaces = tuple(surfaces)
+        if surface_weights is None:
+            surface_weights = [1.0] * len(self.surfaces)
+        if len(surface_weights) != len(self.surfaces):
+            raise ValueError("surface_weights length mismatch")
+        total = sum(surface_weights)
+        self._cum_weights = []
+        run = 0.0
+        for w in surface_weights:
+            run += w / total
+            self._cum_weights.append(run)
+
+    def _pick_surface(self) -> str:
+        r = random.random()
+        for s, c in zip(self.surfaces, self._cum_weights):
+            if r <= c:
+                return s
+        return self.surfaces[-1]
+
+    def sample_batch(self, batch_size: int, seq_len: int, device: str):
+        need = seq_len + 1
+        x_batch = torch.zeros((batch_size, seq_len), dtype=torch.long)
+        y_batch = torch.zeros((batch_size, seq_len), dtype=torch.long)
+        labels = {v: torch.full((batch_size, seq_len), -1, dtype=torch.long)
+                  for v in self.var_names}
+        for b in range(batch_size):
+            buf_bytes: List[int] = []
+            buf_ann: List[Dict[str, int]] = []
+            while len(buf_bytes) < need:
+                surface = self._pick_surface()
+                text, ann = _make_addition_with_annotations(self.max_digits,
+                                                              surface=surface)
+                buf_bytes.extend(text.encode("ascii"))
+                buf_ann.extend(ann)
+            ids = buf_bytes[:need]
+            x_batch[b] = torch.tensor(ids[:-1], dtype=torch.long)
+            y_batch[b] = torch.tensor(ids[1:], dtype=torch.long)
+            for t in range(seq_len):
+                if t < len(buf_ann):
+                    for v, val in buf_ann[t].items():
+                        if v in labels:
+                            labels[v][b, t] = int(val)
+        return (
+            x_batch.to(device),
+            y_batch.to(device),
+            {v: t.to(device) for v, t in labels.items()},
+        )
+
+
+# ----------------------------- Probes --------------------------------------
+
+class BPProbeHeads(nn.Module):
+    """One linear probe per BP variable. Probes are applied to the residual
+    at the variable's target block. Optionally also at other blocks
+    (anti-spread)."""
+
+    def __init__(self, d_model: int, var_specs):
+        super().__init__()
+        self.var_specs = var_specs
+        self.probes = nn.ModuleDict({
+            name: nn.Linear(d_model, n_classes)
+            for name, n_classes, _ in var_specs
+        })
+
+    def loss(self, residuals_per_block: List[torch.Tensor],
+             labels: Dict[str, torch.Tensor],
+             anti_lambda: float = 0.0) -> Tuple[torch.Tensor, Dict[str, float]]:
+        """Install loss: CE at target block. Anti-spread (optional): for
+        other blocks, push the probe to NOT extract V (negative log-likelihood
+        of UNIFORM distribution = encourage low confidence)."""
+        total = torch.zeros((), device=residuals_per_block[0].device)
+        details = {}
+        for name, n_classes, target_block in self.var_specs:
+            res = residuals_per_block[target_block]
+            logits = self.probes[name](res)  # [B, T, n_classes]
+            y = labels[name]                  # [B, T]
+            mask = y >= 0
+            if int(mask.sum()) == 0:
+                continue
+            ce = F.cross_entropy(
+                logits[mask],
+                y[mask],
+                reduction="mean",
+            )
+            total = total + ce
+            details[f"L_{name}_at_blk{target_block}"] = float(ce.item())
+
+            if anti_lambda > 0.0:
+                anti_terms = []
+                for b_idx, res_b in enumerate(residuals_per_block):
+                    if b_idx == target_block:
+                        continue
+                    other_logits = self.probes[name](res_b)
+                    # Want predicted distribution near uniform -> minimize KL
+                    # to uniform = maximize entropy of softmax.
+                    logp = F.log_softmax(other_logits, dim=-1)
+                    p = logp.exp()
+                    # Negative entropy of p (we want to MAXIMIZE entropy ->
+                    # MINIMIZE negative entropy).
+                    neg_ent = (p * logp).sum(dim=-1)
+                    anti_terms.append(neg_ent[mask].mean())
+                if anti_terms:
+                    anti = torch.stack(anti_terms).mean()
+                    total = total + anti_lambda * anti
+                    details[f"anti_{name}"] = float(anti.item())
+        return total, details
+
+    @torch.no_grad()
+    def probe_accs(self, residuals_per_block: List[torch.Tensor],
+                    labels: Dict[str, torch.Tensor]) -> Dict[str, Dict[int, float]]:
+        out: Dict[str, Dict[int, float]] = {}
+        for name, n_classes, _ in self.var_specs:
+            y = labels[name]
+            mask = y >= 0
+            if int(mask.sum()) == 0:
+                continue
+            out[name] = {}
+            for b_idx, res in enumerate(residuals_per_block):
+                logits = self.probes[name](res)
+                pred = logits.argmax(dim=-1)
+                acc = float((pred[mask] == y[mask]).float().mean())
+                out[name][b_idx] = acc
+        return out
+
+
+# ----------------------------- Forward helper ------------------------------
+
+def head_entropy_loss(model: ByteGPT) -> torch.Tensor:
+    """Penalize entropy of head-norm distribution per row of c_proj.
+
+    For each block, c_proj has shape [d_model, n_head * head_dim]. Reshape to
+    [d_model, n_head, head_dim] and compute per-row squared norms across heads.
+    Normalize per-row to get a head-attribution distribution. Penalize its
+    entropy — minimum (0) means each row written by ONE head.
+    """
+    total = 0.0
+    n_blocks = 0
+    eps = 1e-9
+    n_head = model.config.n_head
+    for blk in model.transformer.h:
+        if not hasattr(blk.attn, "c_proj"):
+            continue
+        W = blk.attn.c_proj.weight
+        d_model = W.shape[0]
+        head_dim = W.shape[1] // n_head
+        if head_dim * n_head != W.shape[1]:
+            continue
+        blocks = W.view(d_model, n_head, head_dim)
+        head_norms_sq = (blocks ** 2).sum(dim=-1)  # [d_model, n_head]
+        p = head_norms_sq / (head_norms_sq.sum(dim=-1, keepdim=True) + eps)
+        ent = -(p * (p + eps).log()).sum(dim=-1)  # [d_model]
+        total = total + ent.mean()
+        n_blocks += 1
+    if n_blocks == 0:
+        return torch.zeros((), device=W.device)
+    return total / n_blocks
+
+
+@torch.no_grad()
+def evaluate_mixed_arithmetic(model: ByteGPT, device: str, n_problems: int,
+                                max_digits: int, surfaces=("reverse", "forward"),
+                                seed: int = 42) -> Dict[str, float]:
+    """Greedy-decode `a+b=<marker>` and compare against ground truth for each
+    surface. Returns per-surface accuracies and a combined number."""
+    random.seed(seed)
+    torch.manual_seed(seed)
+    nl_id = ord("\n")
+    results: Dict[str, Dict[str, list]] = {s: {"correct": [], "by_len": {}}
+                                              for s in surfaces}
+    per_surface_n = max(1, n_problems // len(surfaces))
+    for surface in surfaces:
+        marker = "R" if surface == "reverse" else "F"
+        for _ in range(per_surface_n):
+            n_a = random.randint(1, max_digits)
+            n_b = random.randint(1, max_digits)
+            a = random.randint(0, 10 ** n_a - 1)
+            b = random.randint(0, 10 ** n_b - 1)
+            ans_int = a + b
+            ans_forward = str(ans_int)
+            true_emit = ans_forward[::-1] if surface == "reverse" else ans_forward
+            prompt = f"{a}+{b}={marker}"
+            ids = torch.tensor([list(prompt.encode("ascii"))],
+                                 dtype=torch.long, device=device)
+            max_new = len(true_emit) + 2
+            gen = bytearray()
+            for _ in range(max_new):
+                logits, _ = model(ids)
+                nxt = int(logits[0, -1].argmax())
+                if nxt == nl_id:
+                    break
+                gen.append(nxt)
+                ids = torch.cat([ids, torch.tensor([[nxt]], device=device)],
+                                  dim=1)
+            try:
+                got = gen.decode("ascii", errors="ignore")
+            except Exception:
+                got = ""
+            is_correct = (got == true_emit)
+            results[surface]["correct"].append(1 if is_correct else 0)
+            n_digits = max(n_a, n_b)
+            bucket = results[surface]["by_len"].setdefault(n_digits, [0, 0])
+            bucket[1] += 1
+            if is_correct:
+                bucket[0] += 1
+    out = {}
+    total_correct, total_n = 0, 0
+    for s in surfaces:
+        c = sum(results[s]["correct"])
+        n = max(1, len(results[s]["correct"]))
+        out[f"acc_{s}"] = c / n
+        for k, (cc, tt) in sorted(results[s]["by_len"].items()):
+            out[f"acc_{s}_{k}d"] = cc / max(1, tt)
+        total_correct += c
+        total_n += n
+    out["accuracy"] = total_correct / max(1, total_n)
+    return out
+
+
+def forward_with_block_captures(model: ByteGPT, x: torch.Tensor, y: torch.Tensor):
+    """Run model and capture pre-block residuals + post-ln_f.
+    Returns (logits, lm_loss, residuals_per_block).
+    residuals_per_block[i] is the input to block i (= output of block i-1).
+    """
+    captures: List[torch.Tensor] = [None] * model.config.n_layer
+    handles = []
+    for i, blk in enumerate(model.transformer.h):
+        def make_hook(idx):
+            def hook(module, inputs):
+                captures[idx] = inputs[0]
+            return hook
+        handles.append(blk.register_forward_pre_hook(make_hook(i)))
+    try:
+        logits, lm_loss = model(x, y)
+    finally:
+        for h in handles:
+            h.remove()
+    return logits, lm_loss, captures
+
+
+# ----------------------------- Training ------------------------------------
+
+@dataclass
+class TrainConfig:
+    n_layer: int = 4
+    n_head: int = 4
+    n_embd: int = 128
+    mlp_mult: int = 4
+    mlp_activation: str = "silu"
+    seq_len: int = 128
+    attn_kind: str = "mha"
+    attn_norm: str = "entmax15"
+    position_encoding: str = "alibi"
+    loss: str = "bce"
+    logit_rmsnorm_scale: float = 8.0
+    sparsity_lambda: float = 1e-3
+    # A4: Attention Residuals depth-router + entropy pressure, composed
+    # with the BP-install probe-loss. The entropy term is added to
+    # lm_loss inside ByteGPT.forward; install-loss supplies the correct
+    # support. See docs/attnres-bp-research-program.md A2′-result / A4.
+    res_attn: str = "none"
+    res_attn_norm: str = "softmax"
+    res_attn_temp: float = 1.0
+    res_attn_entropy_lambda: float = 0.0
+
+
+def build_model(tc: TrainConfig, device: str) -> ByteGPT:
+    cfg = GPTConfig(
+        vocab_size=256,
+        n_layer=tc.n_layer,
+        n_head=tc.n_head,
+        n_embd=tc.n_embd,
+        seq_len=tc.seq_len,
+        mlp_mult=tc.mlp_mult,
+        mlp_activation=tc.mlp_activation,
+        attn_kind=tc.attn_kind,
+        attn_norm=tc.attn_norm,
+        position_encoding=tc.position_encoding,
+        loss_kind=tc.loss,
+        logit_rmsnorm_scale=tc.logit_rmsnorm_scale,
+        sparsity_lambda=tc.sparsity_lambda,
+        res_attn=tc.res_attn,
+        res_attn_norm=tc.res_attn_norm,
+        res_attn_temp=tc.res_attn_temp,
+        res_attn_entropy_lambda=tc.res_attn_entropy_lambda,
+    )
+    return ByteGPT(cfg).to(device)
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--steps", type=int, default=8000)
+    ap.add_argument("--batch-size", type=int, default=64)
+    ap.add_argument("--seq-len", type=int, default=128)
+    ap.add_argument("--lr", type=float, default=3e-3)
+    ap.add_argument("--lr-min", type=float, default=3e-4)
+    ap.add_argument("--warmup-steps", type=int, default=500)
+    ap.add_argument("--weight-decay", type=float, default=0.01)
+    ap.add_argument("--max-digits", type=int, default=3)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--surfaces", nargs="+", default=["reverse"],
+                    choices=["reverse", "forward"],
+                    help="Surface forms to sample (uniform).")
+    ap.add_argument("--eval-every", type=int, default=1000)
+    ap.add_argument("--n-eval-arith", type=int, default=200)
+    ap.add_argument("--target-carry-block", type=int, default=2)
+    ap.add_argument("--target-sum-block", type=int, default=3)
+    ap.add_argument("--install-lambda", type=float, default=1.0,
+                    help="Coefficient for BP-install probe loss.")
+    ap.add_argument("--anti-lambda", type=float, default=0.0,
+                    help="Coefficient for anti-spread loss (push non-target "
+                         "block probes to uniform).")
+    ap.add_argument("--head-entropy-lambda", type=float, default=0.0,
+                    help="Coefficient for head specialization loss: entropy of "
+                         "head-norm distribution across rows of c_proj. "
+                         "Minimize -> each row written by one head.")
+    ap.add_argument("--res-attn", choices=["none", "full"], default="none",
+                    help="A4: 'full' = block-granular Attention Residuals "
+                         "(depth-router) composed with BP-install supervision.")
+    ap.add_argument("--res-attn-norm",
+                    choices=["softmax", "entmax15", "sparsemax"],
+                    default="softmax",
+                    help="Depth-attention normalization (res-attn only).")
+    ap.add_argument("--res-attn-temp", type=float, default=1.0,
+                    help="A2′ fixed depth-score temperature (res-attn only).")
+    ap.add_argument("--res-attn-entropy-lambda", type=float, default=0.0,
+                    help="A2′ entropy penalty on α_depth (res-attn only); "
+                         "supplies sparsity, install-loss supplies support.")
+    ap.add_argument("--n-layer", type=int, default=4)
+    ap.add_argument("--n-head", type=int, default=4)
+    ap.add_argument("--n-embd", type=int, default=128)
+    ap.add_argument("--checkpoint-dir", type=Path, required=True)
+    ap.add_argument("--output", type=Path, required=True)
+    ap.add_argument("--save-every", type=int, default=2000)
+    args = ap.parse_args()
+
+    random.seed(args.seed)
+    torch.manual_seed(args.seed)
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    tc = TrainConfig(seq_len=args.seq_len, n_layer=args.n_layer,
+                     n_head=args.n_head, n_embd=args.n_embd,
+                     res_attn=args.res_attn, res_attn_norm=args.res_attn_norm,
+                     res_attn_temp=args.res_attn_temp,
+                     res_attn_entropy_lambda=args.res_attn_entropy_lambda)
+    model = build_model(tc, device)
+    if tc.res_attn != "none":
+        print(f"# AttnRes: {tc.res_attn}/{tc.res_attn_norm} "
+              f"temp={tc.res_attn_temp} entropy_lambda={tc.res_attn_entropy_lambda}")
+
+    var_specs = default_var_specs(args.max_digits, args.target_carry_block,
+                                     args.target_sum_block)
+    probes = BPProbeHeads(tc.n_embd, var_specs).to(device)
+    corpus = ArithmeticBPCorpus(args.max_digits, var_specs=var_specs,
+                                  surfaces=tuple(args.surfaces))
+
+    # Optimizer: include both model and probe params
+    all_params = list(model.parameters()) + list(probes.parameters())
+    optimizer = torch.optim.AdamW(all_params, lr=args.lr,
+                                   weight_decay=args.weight_decay,
+                                   betas=(0.9, 0.95))
+
+    def cur_lr(step):
+        if step < args.warmup_steps:
+            return args.lr * step / max(1, args.warmup_steps)
+        progress = (step - args.warmup_steps) / max(1, args.steps - args.warmup_steps)
+        progress = min(1.0, max(0.0, progress))
+        return args.lr_min + 0.5 * (args.lr - args.lr_min) * (1 + math.cos(math.pi * progress))
+
+    args.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    history: List[Dict] = []
+
+    print(f"# device={device} n_layer={tc.n_layer} n_head={tc.n_head} n_embd={tc.n_embd}")
+    print(f"# BP supervision: λ_install={args.install_lambda} λ_anti={args.anti_lambda}")
+    print(f"# Targets: carry→blk{args.target_carry_block}, sum→blk{args.target_sum_block}")
+    print(f"# Variables supervised: {[s[0] for s in var_specs]}")
+
+    last_train_loss = None
+    last_lm_loss = None
+    last_install_loss = None
+    for step in range(1, args.steps + 1):
+        model.train()
+        probes.train()
+        lr = cur_lr(step)
+        for pg in optimizer.param_groups:
+            pg["lr"] = lr
+
+        x, y, bp_labels = corpus.sample_batch(args.batch_size, args.seq_len, device)
+        optimizer.zero_grad(set_to_none=True)
+        _, lm_loss, captures = forward_with_block_captures(model, x, y)
+        install_loss, details = probes.loss(captures, bp_labels,
+                                              anti_lambda=args.anti_lambda)
+        total = lm_loss + args.install_lambda * install_loss
+        if args.head_entropy_lambda > 0:
+            head_ent = head_entropy_loss(model)
+            total = total + args.head_entropy_lambda * head_ent
+            last_head_ent = float(head_ent.item())
+        else:
+            last_head_ent = 0.0
+        total.backward()
+        torch.nn.utils.clip_grad_norm_(all_params, 1.0)
+        optimizer.step()
+        last_train_loss = float(total.item())
+        last_lm_loss = float(lm_loss.item())
+        last_install_loss = float(install_loss.item())
+
+        if step == 1 or step % args.eval_every == 0 or step == args.steps:
+            model.eval()
+            probes.eval()
+            # Compute val loss + bpb on fresh batch
+            with torch.no_grad():
+                vx, vy, vbp = corpus.sample_batch(args.batch_size, args.seq_len, device)
+                vlogits, vlm_loss, vcaps = forward_with_block_captures(model, vx, vy)
+                vce = F.cross_entropy(vlogits.view(-1, vlogits.size(-1)),
+                                       vy.reshape(-1))
+                vbpb = float(bits_per_byte(vce))
+                # Probe accuracies across all blocks (diagnostic)
+                accs = probes.probe_accs(vcaps, vbp)
+            arith = evaluate_mixed_arithmetic(
+                model, device, n_problems=args.n_eval_arith,
+                max_digits=args.max_digits,
+                surfaces=tuple(args.surfaces),
+            )
+            entry = {
+                "step": step,
+                "lr": lr,
+                "train_loss": last_train_loss,
+                "lm_loss": last_lm_loss,
+                "install_loss": last_install_loss,
+                "head_ent": last_head_ent,
+                "val_lm": float(vlm_loss.item()),
+                "val_ce": float(vce.item()),
+                "val_bpb": vbpb,
+                "arith_acc": arith["accuracy"],
+                "probe_accs": accs,
+            }
+            history.append(entry)
+            line = (f"step={step:>5d} lr={lr:.2e} "
+                    f"lm={last_lm_loss:.3f} install={last_install_loss:.3f} "
+                    f"head_ent={last_head_ent:.3f} "
+                    f"val_bpb={vbpb:.3f} arith_acc={arith['accuracy']:.3f}")
+            print(line)
+            # Compact probe summary: install variable at its target block + spread
+            for name, n_cls, tgt in var_specs:
+                if name in accs:
+                    row = accs[name]
+                    target_acc = row.get(tgt, float("nan"))
+                    others = [v for k, v in row.items() if k != tgt]
+                    spread = (sum(others) / len(others)) if others else 0.0
+                    if step == args.steps or step % (args.eval_every * 4) == 1 or step == 1:
+                        print(f"    {name}: blk{tgt}={target_acc:.3f}  "
+                              f"avg-other={spread:.3f}  delta={target_acc-spread:+.3f}")
+
+        if (args.save_every and step % args.save_every == 0) or step == args.steps:
+            ck = args.checkpoint_dir / f"bp_sup_step_{step}.pt"
+            torch.save({
+                "model_state": model.state_dict(),
+                "model_config": vars(model.config),
+                "probe_state": probes.state_dict(),
+                "var_specs": var_specs,
+                "step": step,
+            }, ck)
+            torch.save({"model_state": model.state_dict(),
+                          "model_config": vars(model.config),
+                          "probe_state": probes.state_dict(),
+                          "var_specs": var_specs,
+                          "step": step}, args.checkpoint_dir / "bp_sup_latest.pt")
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps({
+        "config": vars(args) | {"checkpoint_dir": str(args.checkpoint_dir), "output": str(args.output)},
+        "var_specs": var_specs,
+        "history": history,
+    }, indent=2, default=str))
+
+
+if __name__ == "__main__":
+    main()

weights_r0.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9451db167738bd8b84e25ffc1944af16f14c91a721f39f8257e23b0646ac7734
+size 43512279

weights_r1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6a405c8aeb8bc1d5beb451cc619e7e0c55bd0807b636fae9b7a438c6fe2c5dca
+size 19709291