feat(gpu): Engram byte-n-gram conditional memory + validated framing protocol

scripts/static_int8_experiment.py

@@ -88,20 +88,22 @@ def trim(pasts: dict[str, np.ndarray], seq_len: int) -> dict[str, np.ndarray]:
 
 
 def greedy(dec: Decoder, hidden, batch: int) -> list[int]:
-    """Greedy decode driven with `batch`-token feeds: position i's
-    prediction comes from a call fed the `batch` tokens ending at i.
-    batch=1 is the classic incremental loop; batch=8 is the speculative
-    verification framing."""
-    argmax, pasts = dec.split(dec.run([PAD_ID], hidden, None))
+    """Greedy decode driven with `batch`-token feeds. The window ranges
+    over the FULL consumed sequence ([PAD] + emitted tokens): each
+    token's prediction comes from a call fed the `batch` tokens ending
+    at it, with the KV covering everything before them — the judge
+    call's presents cover exactly the consumed sequence, so the cache
+    never rebuilds from a window alone (the earlier versions dropped
+    the PAD prefix and degenerated). batch=1 is the classic incremental
+    loop; batch=8 is the speculative-verification framing."""
+    seq = [PAD_ID]
+    argmax, pasts = dec.split(dec.run(seq, hidden, None))
     traj = [argmax]
     while len(traj) < STEPS and traj[-1] != EOS_ID:
-        window = traj[-batch:]
-        out = dec.run(window, hidden, Decoder.trim(pasts, len(pasts[next(iter(pasts))]) if pasts else 0))
-        argmax, _ = dec.split(out)
-        # rebuild pasts through honest incremental steps (the cache must
-        # reflect every consumed token, mirroring a real runtime)
-        for tok in window:
-            _, pasts = dec.split(dec.run([tok], hidden, pasts))
+        window = seq[-batch:]
+        out = dec.run(window, hidden, Decoder.trim(pasts, len(seq) - len(window)))
+        argmax, pasts = dec.split(out)
+        seq.append(argmax)
         traj.append(argmax)
     return traj
 
@@ -196,7 +198,10 @@ def rewind(self):
     static = Decoder(session(static_path))
 
     test_rows = rows[:5]
-    for name, dec in (("fp32", fp32), ("static-int8", static)):
+    with zipfile.ZipFile(DYNAMIC_INT8_ZIP) as zf:
+        dyn_bytes = zf.read("decoder-kv.onnx")
+    dynamic = Decoder(session(dyn_bytes))
+    for name, dec in (("fp32", fp32), ("dynamic-int8", dynamic), ("static-int8", static)):
         same = total = 0
         drift = 0
         for text in test_rows:
@@ -210,10 +215,7 @@ def rewind(self):
             drift += sum(1 for a, b in zip(single, ref) if a != b)
         print(f"{name}: framing single==batched {same}/{total}; drift-vs-fp32 {drift}/{total}")
 
-    # speed vs the shipped dynamic int8
-    with zipfile.ZipFile(DYNAMIC_INT8_ZIP) as zf:
-        dyn_bytes = zf.read("decoder-kv.onnx")
-    dynamic = Decoder(session(dyn_bytes))
+    # speed: dynamic vs static
     for name, dec in (("dynamic-int8", dynamic), ("static-int8", static)):
         t0 = time.time()
         toks = 0

src/gpu/engram.py

@@ -0,0 +1,108 @@
+"""Engram: byte-n-gram conditional memory for byte-level students
+(TODO.impl/10, gated open 2026-09-12; mechanism corrected from
+DeepSeek-V4.1-Flash §2.4.2 — hash-addressed lookup tables, not example
+retrieval).
+
+Haraqat are strongly lexical: a lookup keyed on byte n-grams captures
+idiomatic vocalization that a 300M byte model must otherwise spend
+capacity memorizing. The module sums ONE embedding per position into
+the encoder stream at ONE layer — the proportionate dose for a 300M
+backbone (DeepSeek places two modules in a 552B model).
+
+Addressing: each byte position addresses the table with k independent
+hashes of the n-grams ENDING at it (orders {2,3,4}); the k looked-up
+vectors are averaged, projected to d_model, and added. Deterministic
+hashing (FNV-1a variants per order/head) — no learned addressing, so
+the table is pure memorization decoupled from compute, prefetchable,
+and export-stable (Gather ops only).
+
+Training pairs with the Sinkhorn-balanced update (gpu.sinkhorn_update):
+row-structured embedding tables are exactly its intended parameter
+class. Storage in the artifact is int8 with per-table scales; the
+module holds fp32/fp16 at runtime.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+FNV_OFFSET = 0x811C9DC5
+FNV_PRIMES = (0x01000193, 0x01000193**2 % (1 << 32), 0x85EBCA6B)
+
+
+def _fnv1a(data: bytes, seed: int) -> int:
+    h = (FNV_OFFSET ^ seed) & 0xFFFFFFFF
+    for b in data:
+        h = ((h ^ b) * 0x01000193) & 0xFFFFFFFF
+    return h
+
+
+def ngram_addresses(seq: list[int], orders=(2, 3, 4)) -> list[list[int]]:
+    """Per-position table addresses: order-n hashes of the byte
+    n-grams ENDING at each position. Positions without a full n-gram
+    (sequence starts) address 0 — the null row."""
+    addresses: list[list[int]] = []
+    for end in range(len(seq)):
+        row = [0] * len(orders)
+        for i, order in enumerate(orders):
+            if end + 1 >= order:
+                gram = bytes((t - 3) % 256 for t in seq[end + 1 - order : end + 1])
+                row[i] = _fnv1a(gram, seed=i + 1) or 1
+        addresses.append(row)
+    return addresses
+
+
+class Engram(nn.Module):
+    """Byte-n-gram lookup memory summed into the encoder stream.
+
+    table_shape: (entries, dim). Projection maps dim -> d_model; the
+    added vector is zero-initialized (gated by a zero scalar) so a
+    freshly attached module leaves the backbone's function identical
+    at step 0 — the same safety property as the PKM gate.
+    """
+
+    def __init__(self, d_model: int, entries: int = 1 << 21, dim: int = 32,
+                 orders=(2, 3, 4)):
+        super().__init__()
+        self.orders = tuple(orders)
+        self.table = nn.Embedding(entries, dim)
+        self.proj = nn.Linear(dim, d_model, bias=False)
+        nn.init.zeros_(self.proj.weight)
+        nn.init.normal_(self.table.weight, std=0.02)
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """[B, T] token ids -> [B, T, d_model] additive memory.
+        Python-side address computation — the training path."""
+        b, t = input_ids.shape
+        flat = input_ids.flatten().tolist()
+        rows = []
+        for bi in range(b):
+            seq = flat[bi * t : (bi + 1) * t]
+            rows.extend(ngram_addresses(seq, self.orders))
+        addresses = torch.tensor(rows, dtype=torch.long,
+                                 device=input_ids.device).reshape(b, t, -1)
+        return self.from_addresses(addresses)
+
+    def from_addresses(self, addresses: torch.Tensor) -> torch.Tensor:
+        """[B, T, k] table addresses -> [B, T, d_model]. The pure-graph
+        path: Gather + mean + Linear only, exportable and runtime-
+        portable. The hash itself is computed by the caller (training:
+        ngram_addresses in python; runtimes: the same 20-line function
+        per runtime — the IMF contract carries addresses as an input."""
+        # remainder, not bitwise-and: ONNX exports int64 mod (and
+        # prime table sizes — the report's distinct-primes choice —
+        # become available)
+        looked = self.table(torch.remainder(addresses, self.table.num_embeddings))
+        return self.proj(looked.mean(dim=2))
+
+    def export_int8_state(self) -> dict[str, torch.Tensor]:
+        """The artifact form: int8 table with one scale, fp16
+        projection (the projection is d_model*dim — tiny)."""
+        w = self.table.weight.detach()
+        scale = w.abs().max().clamp(min=1e-8) / 127.0
+        return {
+            "table_int8": (w / scale).round().to(torch.int8),
+            "table_scale": scale.reshape(1),
+            "proj_fp16": self.proj.weight.detach().half(),
+        }

tests/test_engram.py

@@ -0,0 +1,94 @@
+"""Engram specs (TODO.impl/10): byte-n-gram addressing, zero-init
+safety, int8 export roundtrip, and the ONNX gather-survival probe."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from gpu.engram import Engram, ngram_addresses  # noqa: E402
+
+
+def test_addresses_are_deterministic_and_order_separated() -> None:
+    seq = [117, 114, 110, 1]  # "rok" + EOS in the byte table
+    a1 = ngram_addresses(seq)
+    a2 = ngram_addresses(seq)
+    assert a1 == a2
+    # different orders address independently (seeded hashes)
+    assert a1[3][0] != a1[3][1] != a1[3][2]
+    # positions without a full n-gram address the null row
+    assert a1[0] == [0, 0, 0]
+    assert a1[1][1:] == [0, 0]
+
+
+def test_addresses_track_the_byte_not_the_token() -> None:
+    # the hash input is (id-3)%256: ids 259+3k wrap to the same byte
+    a = ngram_addresses([103 + 256, 104 + 256])[1][0]
+    b = ngram_addresses([103, 104])[1][0]
+    assert a == b
+
+
+def test_zero_init_leaves_the_backbone_untouched() -> None:
+    eng = Engram(d_model=32, entries=1024, dim=8)
+    ids = torch.tensor([[117, 114, 110, 1]])
+    out = eng(ids)
+    assert out.shape == (1, 4, 32)
+    assert torch.all(out == 0), "fresh module must add nothing (the PKM rule)"
+
+
+def test_lookup_varies_with_context() -> None:
+    eng = Engram(d_model=8, entries=4096, dim=4)
+    with torch.no_grad():
+        eng.proj.weight.normal_()
+    ctx = torch.tensor([[5, 6, 7, 8, 9]])
+    shifted = torch.tensor([[5, 6, 7, 9, 8]])  # last two bytes swapped
+    assert not torch.allclose(eng(ctx)[0, 4], eng(shifted)[0, 4])
+
+
+def test_int8_export_roundtrip() -> None:
+    eng = Engram(d_model=16, entries=512, dim=8)
+    with torch.no_grad():
+        eng.table.weight.normal_(std=0.05)
+    state = eng.export_int8_state()
+    deq = state["table_int8"].float() * state["table_scale"]
+    err = (deq - eng.table.weight.detach()).abs().max().item()
+    assert err < 0.05 / 127.0 * 4  # within a few int8 quanta of the scale
+    assert state["table_int8"].dtype == torch.int8
+
+
+def test_gather_survives_onnx_and_matches_torch() -> None:
+    ort = pytest.importorskip("onnxruntime")
+    import numpy as np
+
+    eng = Engram(d_model=16, entries=512, dim=8)
+    with torch.no_grad():
+        eng.proj.weight.normal_(std=0.1)
+    eng.requires_grad_(False)  # parameters as constants, not live tensors
+
+    # the pure-graph path: addresses as a graph INPUT (runtimes compute
+    # the same hash — ngram_addresses is the portable contract)
+    addresses = torch.tensor(ngram_addresses([5, 6, 7, 8, 9]),
+                             dtype=torch.long).unsqueeze(0)
+
+    class Wrapper(torch.nn.Module):
+        def forward(self, addrs: torch.Tensor) -> torch.Tensor:
+            with torch.no_grad():
+                return eng.from_addresses(addrs)
+
+    wrapper = Wrapper().eval()
+    path = Path(__file__).parent / "fixtures" / "engram-probe.onnx"
+    path.parent.mkdir(exist_ok=True)
+    torch.onnx.export(wrapper, addresses, str(path), opset_version=14,
+                      input_names=["addresses"], output_names=["memory"],
+                      dynamo=False)
+    sess = ort.InferenceSession(str(path))
+    got = sess.run(None, {"addresses": addresses.numpy().astype(np.int64)})[0]
+    want = wrapper(addresses).detach().numpy()
+    np.testing.assert_allclose(got, want, atol=1e-4)
+    path.unlink()  # probe artifact, not a fixture