From 24326ed63e2ef1d5a15e7607699e551accf134a6 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sat, 12 Sep 2026 21:34:47 +0800 Subject: [PATCH 1/2] feat(gpu): Engram byte-n-gram conditional memory (TODO.impl/10, gate opened) The lexical-memory lever: one embedding table addressed by hashes of the byte n-grams ending at each position (orders 2/3/4), summed into the encoder stream at one layer. Zero-init projection leaves the backbone untouched at step 0 (the PKM rule). int8 table + fp16 projection as the artifact form (64 MiB at 2M x 32). Export contract settled by the probe: hashing is NOT graph-exportable (python data-dependence) and int64 bitwise-and is rejected by ONNX - from_addresses is the pure-graph path (Gather + mean + Linear with remainder addressing, which also allows the report's prime table sizes), and the hash becomes a 20-line portable function per runtime, with addresses as an IMF input. Probe: export -> ORT -> torch match within 1e-4. 6 specs: determinism and order separation, byte-not-token addressing, zero-init identity, context sensitivity, int8 roundtrip, ONNX gather survival. Training run remains owner-gated. --- src/gpu/engram.py | 108 +++++++++++++++++++++++++++++++++++++++++++ tests/test_engram.py | 94 +++++++++++++++++++++++++++++++++++++ 2 files changed, 202 insertions(+) create mode 100644 src/gpu/engram.py create mode 100644 tests/test_engram.py diff --git a/src/gpu/engram.py b/src/gpu/engram.py new file mode 100644 index 0000000..2b9f7f8 --- /dev/null +++ b/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(), + } diff --git a/tests/test_engram.py b/tests/test_engram.py new file mode 100644 index 0000000..8951079 --- /dev/null +++ b/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 From ea6b692349506ee79b7e67029eac0354231fa51a Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sat, 12 Sep 2026 21:35:30 +0800 Subject: [PATCH 2/2] fix(experiment): the validated framing protocol (greedy v3) Two protocol bugs invalidated the first framing runs: a cache-rebuild loop that re-fed window tokens onto a cache already containing them, and a window that dropped the PAD prefix (batch=1 degenerated to an order-1 Markov decode - fp32 'diverging' with itself was the tell). v3 windows over the full consumed sequence [PAD]+tokens and uses the judge call's own presents as the advanced cache. Sanity: fp32 framing 480/480 under Python ORT (the physics baseline). --- scripts/static_int8_experiment.py | 36 ++++++++++++++++--------------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/scripts/static_int8_experiment.py b/scripts/static_int8_experiment.py index d81d8e9..a5f5152 100644 --- a/scripts/static_int8_experiment.py +++ b/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