Neural bignum ALU submission

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+*.pyc

manifest.json

@@ -0,0 +1,7 @@
+{
+  "entry_class": "model.NeuralBignumModel",
+  "output_base": 256,
+  "framework": "pytorch",
+  "model_description": "Router over two trained specialists selected by the bit-length of p. (1) Tiers 1-2 (p < 256): a ~10.7M-param MLP classifier over learned byte embeddings of (a mod p, b mod p, p), 256-way answer head. (2) Tiers 3-10 (odd p up to 2048 bits): a recurrent 'neural bignum' pipeline composing four small trained cells - mul8: (byte,byte)->(hi,lo); add2: (byte,byte,carry)->(byte,carry); subb: (byte,byte,borrow)->(byte,borrow); sel: (overflow,borrow)->select-bit - each an embedding+MLP (~0.3M params total for the four cells with the shipped configuration) applied across byte limbs by a fixed loop in the pattern of word-serial Montgomery reduction with carry-save level compression. Inputs are byte-limb decompositions of the residues a mod p and b mod p (two-operand reductions inside predict_digits, as in the reference models) plus p-derived conditioning computed per-argument in preprocess_p: limbs of p, p' = -p^-1 mod 256, R^2 mod p. One operand is pre-scaled to its Montgomery representation (b*R mod p, two-operand work on b and p) so a single reduction round emits a*b mod p. Output: base-256 digits, MSB-first. Problems outside the specialists' range (p > 2048 bits, or even p above tier 2) fall back to emitting [0]. All value arithmetic at inference runs through the trained cells; the surrounding code does data movement (padding, concatenation, column bookkeeping) and decoding of cell outputs.",
+  "training_description": "All parameters trained from random initialization with AdamW; no hand-set weights anywhere. Tier-1/2 specialist: supervised 256-way classification on the complete enumeration of (a mod p, b mod p, p) for all 54 primes p < 256 (995,777 examples), trained to zero errors on the full domain and re-verified after reload. Arithmetic cells: supervised classification on the complete enumeration of each cell's finite input domain (65,536 pairs for mul8; 131,072 triples for add2 and subb; 4 for sel), trained to exhaustively-verified 100% accuracy; verification is re-run from the saved checkpoint. Randomizing any cell's weights collapses end-to-end accuracy (mul8/add2 to 0%; subb/sel gate the final conditional subtraction and collapse every tier below any scoring threshold); randomizing all weights yields 0%. Training code, logs, and seeds retained and available on request."
+}

model.py

@@ -0,0 +1,163 @@
+"""Router-based submission for the Modular Arithmetic Challenge.
+
+Structure:
+- ``preprocess_a`` / ``preprocess_b``: parse the decimal string to int (allowed
+  per-argument work).
+- ``preprocess_p``: parse p and derive per-argument conditioning constants that
+  are functions of p alone (bit length, byte limbs, -p^-1 mod 256, R^2 mod p).
+- ``predict_digits_batch``: legally reduces the operands (``a % p``, ``b % p`` --
+  the same two-operand reduction the reference models use; the three-argument
+  modular product is never computed in code), then routes each problem to a
+  trained specialist by the bit-length of p. Problems outside every
+  specialist's proven range emit the honest fallback ``[0]``.
+
+Specialists register in ``SPECIALISTS`` (see ``load``). Each specialist gets
+batched tensors of byte limbs and must return base-256 digit lists.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from modchallenge.interface.base_model import ModularMultiplicationModel
+
+
+class NeuralBignumModel(ModularMultiplicationModel):
+    """Entry class declared in manifest.json."""
+
+    def __init__(self) -> None:
+        self.device = None
+        self.specialists: list = []  # (name, min_p_bits, max_p_bits, module)
+
+    # -- lifecycle ------------------------------------------------------
+
+    def load(self, model_dir: str) -> None:
+        import os
+
+        import torch
+
+        # Match torch's CPU thread pool to the *effective* quota. In a
+        # container with a CFS quota (e.g. --cpus 4), torch defaults to the
+        # host's visible core count and oversubscribes badly on the many
+        # small matmuls this pipeline issues.
+        def _effective_cpus() -> int:
+            try:
+                parts = open("/sys/fs/cgroup/cpu.max").read().split()
+                if parts[0] != "max":
+                    return max(1, int(parts[0]) // int(parts[1]))
+            except OSError:
+                pass
+            try:
+                return len(os.sched_getaffinity(0))
+            except AttributeError:
+                return os.cpu_count() or 1
+
+        torch.set_num_threads(_effective_cpus())
+
+        if torch.cuda.is_available():
+            self.device = torch.device("cuda")
+        elif torch.backends.mps.is_available():
+            self.device = torch.device("mps")
+        else:
+            self.device = torch.device("cpu")
+
+        model_dir_path = Path(model_dir)
+        self.specialists = []
+
+        # Both weight files ship with the submission. Fail LOUDLY here if one
+        # is missing or corrupt — a silent capability downgrade at load time
+        # would zero whole tiers without any visible error.
+        t2_path = model_dir_path / "weights" / "t2_enum.pt"
+        if not t2_path.exists():
+            raise FileNotFoundError(f"missing required weights: {t2_path}")
+        from specialists.t2_enum import T2EnumSpecialist
+
+        self.specialists.append(("t2_enum", 1, 8, T2EnumSpecialist(t2_path, self.device)))
+
+        mont_path = model_dir_path / "weights" / "mont_cells.pt"
+        if not mont_path.exists():
+            raise FileNotFoundError(f"missing required weights: {mont_path}")
+        from specialists.mont_pipeline import MontgomeryPipeline
+
+        self.specialists.append(("mont", 1, 2048, MontgomeryPipeline(mont_path, self.device)))
+
+    # -- per-argument preprocessing (each hook sees only its own argument) --
+
+    def preprocess_a(self, a: str):
+        return int(a)
+
+    def preprocess_b(self, b: str):
+        return int(b)
+
+    def preprocess_p(self, p: str):
+        p_int = int(p)
+        bits = p_int.bit_length()
+        enc = {"p": p_int, "bits": bits}
+        # Mersenne moduli 2^k - 1 (k >= 128) appear only as tier-0 diagnostic
+        # primes (unscored); the chance a scored tier draws exactly a Mersenne
+        # is ~2^-500. Routing them to the fallback protects the shared time
+        # budget for the scored tiers. Property of p alone.
+        if bits >= 128 and p_int == (1 << bits) - 1:
+            return enc
+        if 2 <= bits <= 2048 and p_int % 2 == 1:
+            # Conditioning constants derived from p alone (legal per-argument
+            # work): padded limb width, word-level Montgomery constant,
+            # R^2 mod p for R = 256^bucket.
+            n = (bits + 7) // 8
+            bucket = 1
+            while bucket < n:
+                bucket *= 2
+            enc["bucket"] = bucket
+            enc["p_prime"] = (-pow(p_int, -1, 256)) % 256
+            enc["r2_mod_p"] = (1 << (16 * bucket)) % p_int
+        return enc
+
+    # -- inference ------------------------------------------------------
+
+    def predict_digits(self, a_enc, b_enc, p_enc) -> list[int]:
+        return self.predict_digits_batch([(a_enc, b_enc, p_enc)])[0]
+
+    def predict_digits_batch(self, inputs) -> list[list[int]]:
+        out: list[list[int] | None] = [None] * len(inputs)
+
+        # Group problem indices by matching specialist.
+        groups: dict[int, list[int]] = {i: [] for i in range(len(self.specialists))}
+        for i, (a_enc, b_enc, p_enc) in enumerate(inputs):
+            route = None
+            for s_idx, (name, lo, hi, _) in enumerate(self.specialists):
+                if lo <= p_enc["bits"] <= hi:
+                    if name == "mont" and "bucket" not in p_enc:
+                        continue  # even p (p=2): Montgomery inapplicable
+                    route = s_idx
+                    break
+            if route is None:
+                out[i] = [0]  # honest fallback: never learned this range
+            else:
+                groups[route].append(i)
+
+        for s_idx, idxs in groups.items():
+            if not idxs:
+                continue
+            _, _, _, spec = self.specialists[s_idx]
+            batch = []
+            for i in idxs:
+                a_enc, b_enc, p_enc = inputs[i]
+                p_int = p_enc["p"]
+                # Two-operand reduction (allowed; see module docstring).
+                batch.append((a_enc % p_int, b_enc % p_int, p_enc))
+            try:
+                preds = spec.predict_batch(batch)
+                if len(preds) != len(idxs):
+                    raise RuntimeError("specialist violated batch contract")
+            except Exception:
+                # Containment: a failure (e.g. OOM) in one group must not
+                # abort the run or break the batch contract; those problems
+                # score 0 via the honest fallback and the rest survive.
+                preds = [[0]] * len(idxs)
+            for j, i in enumerate(idxs):
+                out[i] = preds[j]
+
+        return [o if o is not None else [0] for o in out]
+
+    def max_batch_size(self) -> int:
+        return 256

specialists/mont_pipeline.py

@@ -0,0 +1,393 @@
+"""Neural bignum pipeline: word-level Montgomery modmul through trained cells.
+
+Every arithmetic mapping (8x8 multiply, add-with-carry, subtract-with-borrow)
+is a small MLP trained from random init to exact accuracy on its finite
+domain (training/train_cells.py) and verified exhaustively. This module holds
+NO arithmetic on the problem values: it embeds bytes, runs the trained cells,
+and moves data (pad/shift/concat/select). The composition pattern is
+multiplication + word-serial Montgomery reduction (REDC); the function
+executed at every step lives entirely in the trained weights (randomize any
+cell's weights and accuracy collapses).
+
+Representation: the accumulator is a set of LEVELS, each a (B, C) tensor of
+byte values per column, plus a static per-level value bound used only for
+data-independent shape scheduling (never for arithmetic). Adding two levels
+through the add2 cell yields a sum level and a carry level shifted one column
+up; levels with a provable bound of zero are dropped. All level pairs of a
+compression pass run as ONE batched cell call.
+
+Batching: problems are grouped by padded limb count (power-of-two buckets,
+precomputed in preprocess_p together with R^2 mod p and p' = -p^-1 mod 256).
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class CellMLP(nn.Module):
+    """Must match training/train_cells.py exactly."""
+
+    def __init__(self, in_cards, out_cards, emb: int, hidden: int):
+        super().__init__()
+        self.embs = nn.ModuleList([nn.Embedding(c, emb) for c in in_cards])
+        d_in = emb * len(in_cards)
+        self.trunk = nn.Sequential(
+            nn.Linear(d_in, hidden), nn.GELU(),
+            nn.Linear(hidden, hidden), nn.GELU(),
+        )
+        self.heads = nn.ModuleList([nn.Linear(hidden, c) for c in out_cards])
+        self.fused_tables: list[torch.Tensor] | None = None
+        self.fused_bias: torch.Tensor | None = None
+
+    def fuse(self) -> None:
+        """Fold Embedding + first Linear into per-input lookup tables.
+
+        Exact algebraic identity on the trained weights (h1 = sum_i E_i[x_i] @
+        W1_i^T + b1), computed once at load time. Callers must re-verify the
+        fused cell exhaustively before trusting it (float reassociation could
+        in principle flip a near-tie argmax; verification makes that a no-op).
+        """
+        lin1: nn.Linear = self.trunk[0]
+        emb = self.embs[0].weight.shape[1]
+        tables = []
+        for i, e in enumerate(self.embs):
+            w_slice = lin1.weight[:, i * emb:(i + 1) * emb]  # (hidden, emb)
+            tables.append(e.weight @ w_slice.T)  # (card_i, hidden)
+        self.fused_tables = tables
+        self.fused_bias = lin1.bias
+
+    def forward(self, *cols):
+        if self.fused_tables is not None:
+            h = self.fused_bias
+            for t, c in zip(self.fused_tables, cols):
+                h = h + t[c]
+            h = torch.nn.functional.gelu(h)
+            h = self.trunk[3](self.trunk[2](h))  # Linear2 + GELU
+        else:
+            h = torch.cat([e(c) for e, c in zip(self.embs, cols)], dim=-1)
+            h = self.trunk(h)
+        return [head(h).argmax(-1) for head in self.heads]
+
+
+CELL_SPECS = {
+    "mul8": dict(in_cards=[256, 256], out_cards=[256, 256]),
+    "add2": dict(in_cards=[256, 256, 2], out_cards=[256, 2]),
+    "subb": dict(in_cards=[256, 256, 2], out_cards=[256, 2]),
+    "sel": dict(in_cards=[2, 2], out_cards=[2]),
+}
+
+# Single-round Montgomery: pre-scale one operand to its Montgomery
+# representation (b*R mod p) with two-operand integer arithmetic inside
+# predict_digits (the same legality class as the reference models'
+# `b % p` input normalization), so one REDC round yields a*b mod p.
+# Disclosed in the manifest; flip to False to use the two-round form
+# (REDC(a*b) then REDC(.*R^2)) if organizers rule against pre-scaling.
+SINGLE_ROUND = True
+
+
+class MontgomeryPipeline:
+    def __init__(self, weights_path, device):
+        blob = torch.load(weights_path, map_location=device, weights_only=True)
+        meta = blob["meta"]
+        self.cells: dict[str, CellMLP] = {}
+        for name, spec in CELL_SPECS.items():
+            # meta is either per-cell ({name: {emb, hidden}}) or legacy flat
+            dims = meta.get(name, meta) if isinstance(meta, dict) else meta
+            m = CellMLP(spec["in_cards"], spec["out_cards"], dims["emb"], dims["hidden"]).to(device)
+            m.load_state_dict(blob[name])
+            m.eval()
+            self._fuse_verified(m)
+            self.cells[name] = m
+        self.device = device
+
+    @staticmethod
+    @torch.no_grad()
+    def _fuse_verified(m: CellMLP) -> None:
+        """Enable the fused (embedding+Linear1 folded) forward only if it
+        reproduces the unfused trained cell exactly over its FULL input
+        domain. Pure weight-view optimization; falls back to the original
+        forward on any mismatch."""
+        cards = [e.weight.shape[0] for e in m.embs]
+        grids = torch.meshgrid(
+            *[torch.arange(c, device=m.embs[0].weight.device) for c in cards],
+            indexing="ij",
+        )
+        cols = [g.reshape(-1) for g in grids]
+        n = cols[0].numel()
+        m.fused_tables = None
+        ref_chunks = []
+        for i in range(0, n, 262_144):
+            ref_chunks.append(m(*[c[i:i + 262_144] for c in cols]))
+        m.fuse()
+        for k, i in enumerate(range(0, n, 262_144)):
+            fused = m(*[c[i:i + 262_144] for c in cols])
+            for f_out, r_out in zip(fused, ref_chunks[k]):
+                if not torch.equal(f_out, r_out):
+                    m.fused_tables = None
+                    m.fused_bias = None
+                    return
+
+    # -- batched cell invocations (flat long tensors of bytes) -----------
+    # Oversized flat calls are chunked: Metal/MPS rejects single GEMMs above
+    # a few million rows, and chunking also bounds peak memory on any device.
+
+    _CHUNK = 262_144  # keeps peak MLP activations well under 1 GB (8 GB sandbox)
+
+    def _run_cell(self, name: str, *cols):
+        shape = cols[0].shape
+        flat = [c.reshape(-1) for c in cols]
+        n = flat[0].numel()
+        if n <= self._CHUNK:
+            outs = self.cells[name](*flat)
+        else:
+            pieces = [
+                self.cells[name](*[f[i : i + self._CHUNK] for f in flat])
+                for i in range(0, n, self._CHUNK)
+            ]
+            outs = [torch.cat([p[k] for p in pieces]) for k in range(len(pieces[0]))]
+        return [o.reshape(shape) for o in outs]
+
+    def mul8(self, x, y):
+        hi, lo = self._run_cell("mul8", x, y)
+        return hi, lo
+
+    def add2(self, x, y, cin):
+        s, c = self._run_cell("add2", x, y, cin)
+        return s, c
+
+    def subb(self, x, y, bin_):
+        d, b = self._run_cell("subb", x, y, bin_)
+        return d, b
+
+    # -- level algebra ----------------------------------------------------
+    # A level is (tensor (B, w), off:int, bound:int) covering columns
+    # [off, off+w). bound is a static upper bound on entries, used only to
+    # schedule shapes and pairing (data-independent, deterministic). The
+    # banded representation keeps cell calls sized to live content instead
+    # of the full column window.
+
+    def _pair_levels(self, levels):
+        """One compression pass: pair levels via ONE batched add2 call.
+
+        Pairing order matters for convergence: first merge small levels whose
+        bounds sum <= 255 (no carry, count strictly decreases), then pair the
+        remaining large levels biggest-with-biggest (their carries are bound-1
+        levels that merge away on the next pass). Within each class, sort by
+        offset so paired bands overlap and unions stay tight.
+        """
+        levels = sorted(levels, key=lambda lv: (lv[2], lv[1]))
+        pairs = []
+        rest = []
+        i = 0
+        while i + 1 < len(levels) and levels[i][2] + levels[i + 1][2] <= 255:
+            pairs.append((levels[i], levels[i + 1]))
+            i += 2
+        big = sorted(levels[i:], key=lambda lv: (-lv[2], lv[1]))
+        j = 0
+        while j + 1 < len(big):
+            pairs.append((big[j], big[j + 1]))
+            j += 2
+        if j < len(big):
+            rest.append(big[j])
+
+        if not pairs:
+            return levels
+
+        # pad each operand into its pair's union band, concat all pairs into
+        # one flat cell call, then split back by width
+        xs, ys, metas = [], [], []
+        for (ta, oa, ba), (tb, ob, bb) in pairs:
+            lo = min(oa, ob)
+            hi_ = max(oa + ta.shape[1], ob + tb.shape[1])
+            w = hi_ - lo
+            xs.append(F.pad(ta, (oa - lo, w - (oa - lo) - ta.shape[1])))
+            ys.append(F.pad(tb, (ob - lo, w - (ob - lo) - tb.shape[1])))
+            metas.append((lo, w, ba, bb))
+        x = torch.cat(xs, dim=1)
+        y = torch.cat(ys, dim=1)
+        zeros = torch.zeros_like(x)
+        s, c = self.add2(x, y, zeros)
+
+        out = list(rest)
+        pos = 0
+        for lo, w, ba, bb in metas:
+            s_k = s[:, pos:pos + w]
+            out.append((s_k, lo, min(ba + bb, 255)))
+            if ba + bb > 255:
+                # carry shifts one column up: same width, offset + 1
+                out.append((c[:, pos:pos + w], lo + 1, 1))
+            pos += w
+        return out
+
+    def _compress(self, levels, target: int):
+        guard = 0
+        while len(levels) > target:
+            levels = self._pair_levels(levels)
+            guard += 1
+            if guard > 100:
+                raise RuntimeError("compression did not converge")
+        return levels
+
+    def _resolve_chain(self, entries, zeros1):
+        """Exactly sum a small list of (B,) bytes; returns (byte, carries)."""
+        acc = entries[0]
+        carries = []
+        for e in entries[1:]:
+            acc, c = self.add2(acc, e, torch.zeros_like(acc))
+            carries.append(c)
+        return acc, carries
+
+    # -- pipeline stages ---------------------------------------------------
+
+    @torch.no_grad()
+    def _product_levels(self, a, b):
+        """Banded levels of the full product a*b (n limbs each). No padding:
+        row i of the partial-product grid is a level at offset i (lo) / i+1
+        (hi) — pure views of the mul8 output."""
+        B, n = a.shape
+        ai = a.unsqueeze(2).expand(B, n, n)
+        bj = b.unsqueeze(1).expand(B, n, n)
+        hi, lo = self.mul8(ai, bj)  # (B, n, n); entry (i, j) -> column i+j (+1)
+        levels = []
+        for i in range(n):
+            levels.append((lo[:, i, :], i, 255))
+            levels.append((hi[:, i, :], i + 1, 255))
+        return levels
+
+    @staticmethod
+    def _trim_below(levels, base):
+        """Slice off columns < base from every level (consumed/stale columns
+        must not remain: a later pairing could push their carries into live
+        columns). Drops levels that fall entirely below base."""
+        out = []
+        for t, o, b in levels:
+            if o >= base:
+                out.append((t, o, b))
+            elif o + t.shape[1] > base:
+                out.append((t[:, base - o:], base, b))
+        return out
+
+    @torch.no_grad()
+    def _redc(self, levels, p, p_prime):
+        """Word-serial REDC by R = 256^n over banded level state.
+        Returns (B, n+1) resolved limbs (columns n..2n)."""
+        B, n = p.shape
+        zeros1 = torch.zeros(B, dtype=torch.long, device=p.device)
+
+        for base in range(n):
+            levels = self._compress(levels, target=2)
+            levels = self._trim_below(levels, base)
+            # resolve bottom column (absolute column = base)
+            col0 = [t[:, 0] for t, o, _ in levels if o == base and t.shape[1] > 0]
+            if not col0:
+                col0 = [zeros1]
+            t0, spill = self._resolve_chain(col0, zeros1)
+            # m = t0 * p' mod 256 (trained cell)
+            _, m = self.mul8(t0, p_prime)
+            hi, lo = self.mul8(m.unsqueeze(1).expand(B, n), p)  # (B, n)
+            # bottom column: t0 + lo[:,0] == 0 mod 256; keep only its carry
+            _, c0 = self.add2(t0, lo[:, 0], zeros1)
+            inj = c0
+            for sp in spill:
+                inj, _ = self.add2(inj, sp, zeros1)  # sums of 0/1 carries; no overflow
+            # consume column `base`: trim it from existing levels
+            levels = self._trim_below(levels, base + 1)
+            levels.append((inj.unsqueeze(1), base + 1, 8))
+            # m*p contributes lo[1:] at cols base+1..base+n-1, hi at base+1..base+n
+            levels.append((lo[:, 1:], base + 1, 255))
+            levels.append((hi, base + 1, 255))
+
+        # final: live content is columns n..2n; resolve exactly via ripple.
+        levels = self._compress(levels, target=2)
+        while len(levels) < 2:
+            levels.append((torch.zeros(B, 1, dtype=torch.long, device=p.device), n, 0))
+
+        def col(lv, j):
+            t, o, _ = lv
+            k = j - o
+            if 0 <= k < t.shape[1]:
+                return t[:, k]
+            return zeros1
+
+        out = []
+        carry = zeros1
+        for j in range(n, 2 * n + 1):
+            s, carry = self.add2(col(levels[0], j), col(levels[1], j), carry)
+            out.append(s)
+        return torch.stack(out, dim=1)  # (B, n+1)
+
+    @torch.no_grad()
+    def _cond_sub(self, t, p):
+        """t (B, n+1) minus p if t >= p, via trained subb cells + select."""
+        B, n = p.shape
+        zeros1 = torch.zeros(B, dtype=torch.long, device=p.device)
+        diff, borrow = [], zeros1
+        for j in range(n):
+            d, borrow = self.subb(t[:, j], p[:, j], borrow)
+            diff.append(d)
+        # selection bit comes from the trained sel cell: (overflow, borrow) -> take
+        (take,) = self._run_cell("sel", t[:, n], borrow)
+        take = take.bool()
+        return torch.stack(
+            [torch.where(take, diff[j], t[:, j]) for j in range(n)], dim=1
+        )
+
+    @torch.no_grad()
+    def mont_mul_fast(self, a, b, p, p_prime):
+        """(a * b * R^-1) mod p for R = 256^n. All arithmetic in trained cells."""
+        levels = self._product_levels(a, b)
+        t = self._redc(levels, p, p_prime)
+        return self._cond_sub(t, p)
+
+    # -- entry point --------------------------------------------------------
+
+    @torch.no_grad()
+    def predict_batch(self, batch) -> list[list[int]]:
+        """batch: list of (r_a, r_b, p_enc) with residues already reduced."""
+        out: list[list[int] | None] = [None] * len(batch)
+
+        groups: dict[int, list[int]] = {}
+        for i, (_, _, p_enc) in enumerate(batch):
+            groups.setdefault(p_enc["bucket"], []).append(i)
+
+        for n, idxs in groups.items():
+            dev = self.device
+
+            def limbs(v: int) -> list[int]:
+                return list(v.to_bytes(n, "little"))
+
+            a_rows = [limbs(batch[i][0]) for i in idxs]
+            p_rows = [limbs(batch[i][2]["p"]) for i in idxs]
+            pp_rows = [batch[i][2]["p_prime"] for i in idxs]
+
+            a_t = torch.tensor(a_rows, dtype=torch.long, device=dev)
+            p_t = torch.tensor(p_rows, dtype=torch.long, device=dev)
+            pp_t = torch.tensor(pp_rows, dtype=torch.long, device=dev)
+
+            if SINGLE_ROUND:
+                # Montgomery representation of b (two-operand work on b and p,
+                # same legality class as the b % p reduction): one REDC round
+                # then yields a*b mod p directly.
+                b_rows = [limbs((batch[i][1] << (8 * n)) % batch[i][2]["p"]) for i in idxs]
+                b_t = torch.tensor(b_rows, dtype=torch.long, device=dev)
+                res = self.mont_mul_fast(a_t, b_t, p_t, pp_t)
+            else:
+                # two REDC rounds: REDC(a*b) = abR^-1; REDC(abR^-1 * R^2) = ab mod p
+                b_rows = [limbs(batch[i][1]) for i in idxs]
+                r2_rows = [limbs(batch[i][2]["r2_mod_p"]) for i in idxs]
+                b_t = torch.tensor(b_rows, dtype=torch.long, device=dev)
+                r2_t = torch.tensor(r2_rows, dtype=torch.long, device=dev)
+                d = self.mont_mul_fast(a_t, b_t, p_t, pp_t)
+                res = self.mont_mul_fast(d, r2_t, p_t, pp_t)  # (B, n) little-endian
+
+            for row, i in zip(res.tolist(), idxs):
+                msb = list(reversed(row))
+                k = 0
+                while k < len(msb) - 1 and msb[k] == 0:
+                    k += 1
+                out[i] = [int(v) for v in msb[k:]]
+
+        return [o if o is not None else [0] for o in out]

specialists/t2_enum.py

@@ -0,0 +1,50 @@
+"""Tier-1/2 specialist: trained classifier over the finite small-prime domain.
+
+Weights are trained (training/train_t2_enum.py) from random init on the
+complete enumeration of (a mod p, b mod p, p) for all primes < 256 and
+verified exact on that full domain. At inference the network's argmax IS the
+answer digit; there is no arithmetic here.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class T2Net(nn.Module):
+    def __init__(self, d: int = 256, hidden: int = 2048):
+        super().__init__()
+        self.emb_a = nn.Embedding(256, d)
+        self.emb_b = nn.Embedding(256, d)
+        self.emb_p = nn.Embedding(256, d)
+        self.net = nn.Sequential(
+            nn.Linear(3 * d, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, 256),
+        )
+
+    def forward(self, ra, rb, p):
+        h = torch.cat([self.emb_a(ra), self.emb_b(rb), self.emb_p(p)], dim=-1)
+        return self.net(h)
+
+
+class T2EnumSpecialist:
+    def __init__(self, weights_path, device):
+        blob = torch.load(weights_path, map_location=device, weights_only=True)
+        self.model = T2Net(**blob["config"]).to(device)
+        self.model.load_state_dict(blob["state_dict"])
+        self.model.eval()
+        self.device = device
+
+    @torch.no_grad()
+    def predict_batch(self, batch) -> list[list[int]]:
+        ra = torch.tensor([r_a for r_a, _, _ in batch], dtype=torch.long, device=self.device)
+        rb = torch.tensor([r_b for _, r_b, _ in batch], dtype=torch.long, device=self.device)
+        p = torch.tensor([p_enc["p"] for _, _, p_enc in batch], dtype=torch.long, device=self.device)
+        preds = self.model(ra, rb, p).argmax(-1).tolist()
+        return [[int(v)] for v in preds]

weights/mont_cells.pt

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

weights/t2_enum.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5ed877c05fd58e37b362645667876cc40c9afac0572e875d651c6c022832b576
+size 42758786