Upload 6 files

config.py

@@ -0,0 +1,121 @@
+
+use_liger=False
+xla_on_gpu=False
+# ---------------------------------------------------------------------------
+# Core shape -- chosen for GPU efficiency, not just "a number that fits"
+# ---------------------------------------------------------------------------
+CONTEXT = 1024
+max_gen_tokens = 1024
+gen_headroom = 128
+vocab_size = 8192+4096
+# ── Engram (static n-gram memory) ──────────────────────────────────
+use_engram              = True
+engram_layer_ids        = (1, 8)         # two layers, one early, one mid
+engram_max_ngram        = 3              # 2-gram + 3-gram tables
+engram_vocab_size       = (4096, 4096)   # halved from default
+engram_embed_per_ngram  = 128
+engram_n_heads          = 4
+engram_kernel_size      = 4
+engram_seed             = 0
+engram_compressed_vocab = 0              # 0 → set from tokenizer lookup
+bias = False
+d_rope = None
+top_k = 64
+moe_top_k = 2
+use_triton_fusions = False
+
+
+gradient_checkpointing = False
+num_kv_heads = 4
+pattern = "dense"
+sliding_window_size = 512
+
+tensor_parallel_size = 1
+
+# ─────────────────────────────────────────────────────────────────────
+# Curriculum learning (token-level mixing across sources)
+# ─────────────────────────────────────────────────────────────────────
+# Master switch. False = old behavior (uniform random over the whole
+# flat token array). True = per-source anchor schedule below.
+CURRICULUM_ENABLED = False
+
+# Number of training steps the anchor schedule spans. Set equal to
+# MAX_STEPS in train.py so the ramp completes at run end.
+CURRICULUM_TOTAL_STEPS = 20000
+
+# Sources, in curriculum order (easiest first). Each entry:
+#   (name, [glob_patterns], w@0, w@T/3, w@2T/3, w@T)
+# Between anchors, weights interpolate linearly. Columns should sum
+# to ~1.0 (the code renormalizes if not).
+# Files matching none of the globs are EXCLUDED from training.
+CURRICULUM_SOURCES = [
+    ("tinystories", ["tinystories*.txt", "tiny_*.txt"],
+     0.70, 0.10, 0.03, 0.02),
+    ("simplewiki",  ["simplewiki*.txt",  "wiki_*.txt"],
+     0.10, 0.70, 0.10, 0.03),
+    ("cosmopedia",  ["cosmopedia*.txt",  "cosmo_*.txt"],
+     0.10, 0.10, 0.70, 0.10),
+    ("finephrase",  ["finephrase*.txt",  "fine_*.txt"],
+     0.10, 0.10, 0.17, 0.85),
+]
+
+# Hard floor: no source ever drops below this fraction at any step.
+# 0.0 disables. 0.02 = 2% minimum, prevents catastrophic forgetting.
+CURRICULUM_FLOOR = 0.02
+
+mod_capacity = 0.25
+use_mtp    = False
+mtp_depth  = 1      
+mtp_lambda = 0.1 
+use_nvfp4 = False
+# config.py
+use_fp8 = False
+fp8_recipe = "rowwise_with_gw_hp"        # default
+fp8_bf16_tail_frac = 0.15                # default
+use_int8 = False
+int8_group_size = 128
+mod_budget_coeff = 0.05
+mod_gate_entropy_coeff = 0.0
+mod_hard_eval = False
+use_flexattention = True
+numberoflayers = 16     
+numberofheads = 12
+D_MODEL = 384        
+ALPHA_COMPILE_MODE="reduce-overhead"  
+num_kv_heads = 6 
+sliding_window_size = 512  
+use_sdpa=False
+optimizer_8bit = True
+use_wsd = True
+wsd_decay_steps = 3500 
+ffn_mult = 8 // 3           
+use_moe = False            
+n_experts = 8
+n_shared = 1
+# T4 path
+
+PRECISION = 'bf16'         
+GLOBAL_DTYPE = PRECISION
+USE_ASYNC_LAYER_PREFETCH = True
+intra_op_threads = 2
+num_workers = 2
+GPU_SIDE_BATCHING = True
+DISABLE_MMAP = False
+
+batch_size = 320
+weight_decay = 0.01        
+use_static_lr = False
+static_lr = 1e-3        
+learning_rate = 1e-3
+
+
+attn_gate_window = 16
+max_steps = 19000
+data_parallel_only = True
+use_fsdp = False
+dropout = 0.0                   
+use_muon = True
+DATA_PATH = './training_data'
+data_path = DATA_PATH
+TOKENIZER_MODEL_PATH = './tokenizer'
+use_streaming_batcher = False

kernel.py

@@ -0,0 +1,88 @@
+"""FlexAttention-only attention kernel.
+
+One backend, one code path: no device probing, no backend dispatch, no try/except fallbacks.
+Every attention layer calls `fused_attention(q, k, v, block_mask)`; what each layer may attend to
+(causal / sliding window / same-document) lives entirely in the BlockMask, which GPT.forward builds
+once per forward and shares across layers.
+
+Requires CUDA + torch >= 2.5 (`create_block_mask(..., _compile=True)`). Attention-weight dropout is not
+supported (DenseMHA refuses dropout > 0 at construction).
+"""
+import functools
+
+import torch
+from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+
+_USE_COMPILED_FLEX = torch.cuda.is_available() and torch.version.cuda is not None
+if _USE_COMPILED_FLEX:
+    _flex = torch.compile(flex_attention, dynamic=False)
+else:
+    _flex = flex_attention
+
+def fused_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, block_mask) -> torch.Tensor:
+    """q: (B, Hq, T, D); k, v: (B, Hkv, T, D). GQA is handled natively (Hq may be a multiple of Hkv)."""
+    return _flex(q, k, v, block_mask=block_mask, enable_gqa=q.shape[1] != k.shape[1])
+
+
+# ---------------------------------------------------------------------------
+# BlockMasks. One uniform mask_mod for every layer: `window = window_size or seq_len`, and since
+# q - kv <= seq_len - 1 < seq_len, "no window" is the same expression as "window = seq_len".
+# Both builders are compiler-disabled (BlockMask construction is not traced) and are called once per
+# forward from GPT.forward -- that is the only graph boundary in the attention path.
+# ---------------------------------------------------------------------------
+
+@functools.lru_cache(maxsize=None)
+def _causal_mask(seq_len: int, window_size, device_key: str):
+    window = window_size or seq_len
+
+    def mask_mod(b, h, q_idx, kv_idx):
+        return (kv_idx <= q_idx) & (q_idx - kv_idx < window)
+
+    return create_block_mask(mask_mod, B=None, H=None, Q_LEN=seq_len, KV_LEN=seq_len, device=device_key)
+
+
+@torch.compiler.disable
+def causal_masks(seq_len: int, windows, device) -> dict:
+    """{window_size: BlockMask} for plain causal attention; cached per (seq_len, window, device)."""
+    return {w: _causal_mask(seq_len, w, str(device)) for w in windows}
+
+
+@torch.compiler.disable
+def doc_masks(doc_ids: torch.Tensor, windows) -> dict:
+    """{window_size: BlockMask} for causal + same-document attention. doc_ids: (B, T) int, per batch.
+    _compile=True fuses mask_mod, so the (B, T, T) boolean grid is never materialised."""
+    b_, t_ = doc_ids.shape
+
+    def build(window_size):
+        window = window_size or t_
+
+        def mask_mod(b, h, q_idx, kv_idx):
+            return (doc_ids[b, q_idx] == doc_ids[b, kv_idx]) & (kv_idx <= q_idx) & (q_idx - kv_idx < window)
+
+        return create_block_mask(mask_mod, B=b_, H=None, Q_LEN=t_, KV_LEN=t_,
+                                 device=doc_ids.device, _compile=True)
+
+    return {w: build(w) for w in windows}
+
+
+# ---------------------------------------------------------------------------
+# Telemetry counters (model.py records its Liger/Triton success/fallback events here).
+# ---------------------------------------------------------------------------
+
+_backend_stats: dict[str, int] = {}
+
+
+def _record_backend(name: str) -> None:
+    # Inside a compiled region this body runs only at trace time, so counting here under-reports (verified:
+    # 5 compiled calls -> 3 counts). Record nothing while compiling rather than report wrong numbers.
+    if torch.compiler.is_compiling():
+        return
+    _backend_stats[name] = _backend_stats.get(name, 0) + 1
+
+
+def get_backend_stats() -> dict[str, int]:
+    return dict(_backend_stats)
+
+
+def reset_backend_stats() -> None:
+    _backend_stats.clear()

model.safetensors

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

modeling_pragya.py

@@ -0,0 +1,95 @@
+import dataclasses
+
+import torch
+from transformers import GenerationMixin, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_pragya import PragyaConfig
+from .model import GPT, GPTConfig
+
+# Config field names that differ between PragyaConfig and GPTConfig.
+_RENAMED = {"top_k": "gpt_top_k"}
+# Fields GPTConfig stores as tuples, which PragyaConfig stores as lists.
+_TUPLE_FIELDS = ("engram_layer_ids", "engram_vocab_size")
+
+
+def to_gpt_config(cfg: PragyaConfig) -> GPTConfig:
+    """Build a GPTConfig from every field, so no value comes from config.py defaults."""
+    kwargs = {}
+    for f in dataclasses.fields(GPTConfig):
+        value = getattr(cfg, _RENAMED.get(f.name, f.name))
+        if f.name in _TUPLE_FIELDS:
+            value = tuple(value)
+        kwargs[f.name] = value
+    return GPTConfig(**kwargs)
+
+
+def _non_persistent_buffers(module):
+    """{name: tensor} for every non-persistent buffer in a torch module tree."""
+    out = {}
+    for prefix, sub in module.named_modules():
+        for buf_name in sub._non_persistent_buffers_set:
+            full = f"{prefix}.{buf_name}" if prefix else buf_name
+            out[full] = getattr(sub, buf_name)
+    return out
+
+
+class PragyaForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = PragyaConfig
+    base_model_prefix = "model"
+    main_input_name = "input_ids"
+    supports_gradient_checkpointing = False
+    # wte.weight and lm_head.weight are one tensor; save only lm_head.weight.
+    _tied_weights_keys = {"model.transformer.wte.weight": "model.lm_head.weight"}
+
+    def __init__(self, config: PragyaConfig):
+        super().__init__(config)
+        if config.use_engram and config.engram_compressed_vocab <= 0:
+            raise ValueError(
+                "engram_compressed_vocab must be set before loading. It is computed from the "
+                "tokenizer when the model is first built, so the export step must record it."
+            )
+        self.model = GPT(to_gpt_config(config))
+        self.post_init()
+
+    def _move_missing_keys_from_meta_to_device(self, missing_keys, device_map, device_mesh, hf_quantizer):
+        # Transformers replaces every non-persistent buffer (RoPE tables, XSA masks, Engram
+        # `mods`) with uninitialised memory when loading. They are derived from the config,
+        # so rebuild them from a freshly constructed GPT.
+        super()._move_missing_keys_from_meta_to_device(missing_keys, device_map, device_mesh, hf_quantizer)
+        fresh = GPT(to_gpt_config(self.config))
+        own = _non_persistent_buffers(self.model)
+        with torch.no_grad():
+            for name, buf in _non_persistent_buffers(fresh).items():
+                own[name].copy_(buf)
+        del fresh
+        # wte and lm_head are one tensor (checkpoint stores only lm_head). Re-tie after load.
+        self.model.transformer.wte.weight = self.model.lm_head.weight
+
+    def _init_weights(self, module):
+        # Weights come from GPT.__init__ (Engram tables, zero-init attention gate, tied
+        # embeddings) or from the checkpoint. Transformers' default would overwrite them.
+        pass
+
+    def get_input_embeddings(self):
+        return self.model.transformer.wte
+
+    def get_output_embeddings(self):
+        return self.model.lm_head
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
+        if attention_mask is not None and not bool(torch.all(attention_mask == 1)):
+            # The model has no attention-mask support. Padded batches would silently
+            # attend to pad tokens, so refuse rather than return wrong results.
+            raise NotImplementedError(
+                "attention_mask with padding is not supported. Run one sequence at a time "
+                "or pass unpadded batches of equal length."
+            )
+        # With labels: full logits (B, T, V) and loss.
+        # Without labels: logits for the last position only, (B, 1, V).
+        logits, loss = self.model(input_ids, targets=labels)
+        return CausalLMOutputWithPast(loss=loss, logits=logits)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        # No KV cache: GPT keeps its own cache format, so every step runs on the full sequence.
+        return {"input_ids": input_ids}

ngram.py

@@ -0,0 +1,459 @@
+"""
+Engram -- static n-gram memory (arXiv:2601.07372, DeepSeek engram_demo_v1.py),
+adapted for Veylon/Arya's standard transformer (hc_mult = 1, tiktoken).
+
+Fixes vs. the drafted port (all verified in test_engram.py):
+
+  1. Hashing ran on CPU/numpy INSIDE forward(): a GPU->CPU sync + torch.compile
+     graph break on every Engram layer, every step (and fatal on XLA). Now the
+     hash is pure int64 torch ops on the model device, compile-friendly.
+  2. Every Engram layer built its own NgramHashMapping (re-decoding the whole
+     vocab) and each call hashed ALL layers, then kept one. Now ONE NgramHasher
+     lives on GPT, hashes all layers once per forward.
+  3. pad_id=None crashed np.pad; pad_id=-1 silently indexed the LAST lookup
+     entry. Left-context padding is now a dedicated reserved id (never
+     collides with a real token class).
+  4. nn.RMSNorm (torch>=2.4 only, eps=finfo(fp16)=1e-3 default) replaced by a local
+     _RMSNorm (eps=1e-5, weight cast to the input dtype) matching model.RMSNorm's eager
+     path. It does NOT upcast to fp32.
+  5. Gate dot-product computed in fp32 (sum over D in fp16 can overflow/lose
+     precision), cast back after the sigmoid.
+  6. The compressed-vocab lookup is now a persistent buffer, so a checkpoint
+     carries it and inference needs no tokenizer object.
+  7. sympy dependency dropped (trial-division primality is plenty for table sizes).
+  8. Special-id discovery no longer assumes TokenizerWrapper._special_ids is a
+     set of ints (dict / missing attr both handled).
+  9. embed_per_ngram % n_heads != 0 silently produced a wrong-width table ->
+     hard assert. Same for len(engram_vocab_size) < max_ngram-1.
+"""
+
+from __future__ import annotations
+
+import math
+import re
+import unicodedata
+import warnings
+from typing import List, Sequence
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Primes
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _is_prime(n: int) -> bool:
+    if n < 2:
+        return False
+    if n < 4:
+        return True
+    if n % 2 == 0 or n % 3 == 0:
+        return False
+    i = 5
+    while i * i <= n:
+        if n % i == 0 or n % (i + 2) == 0:
+            return False
+        i += 6
+    return True
+
+
+def find_next_prime(start: int, seen_primes: set) -> int:
+    candidate = start + 1
+    while True:
+        if _is_prime(candidate) and candidate not in seen_primes:
+            return candidate
+        candidate += 1
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# CompressedTokenizer -- build-time only (numpy). Output is a lookup table.
+# ─────────────────────────────────────────────────────────────────────────────
+
+class CompressedTokenizer:
+    """Maps raw token IDs -> compressed IDs (NFKC / strip-accents / lowercase /
+    whitespace-collapse of each token's surface string). Build-time only: the
+    result (`lookup_table`) is handed to the model as a persistent buffer."""
+
+    _SENTINEL = "\uE000"
+    _WS_RE = re.compile(r"[ \t\r\n]+")
+
+    def __init__(self, tokenizer_wrapper):
+        self.tokenizer = tokenizer_wrapper
+        self._special_ids = self._collect_special_ids(tokenizer_wrapper)
+        self.lookup_table, self.num_new_token = self._build_lookup_table()
+
+    def __len__(self):
+        return self.num_new_token
+
+    @staticmethod
+    def _collect_special_ids(tw) -> set:
+        ids = set()
+        raw = getattr(tw, "_special_ids", None)
+        if raw is not None:
+            if isinstance(raw, dict):
+                for k, v in raw.items():
+                    for cand in (k, v):
+                        if isinstance(cand, (int, np.integer)):
+                            ids.add(int(cand))
+            else:
+                for v in raw:
+                    if isinstance(v, (int, np.integer)):
+                        ids.add(int(v))
+        vs = int(tw.vocab_size)
+        for name in ("bos_id", "eos_id", "pad_id"):
+            v = getattr(tw, name, None)
+            if isinstance(v, (int, np.integer)) and 0 <= int(v) < vs:
+                ids.add(int(v))
+        return ids
+
+    @classmethod
+    def _normalize(cls, text: str) -> str:
+        text = unicodedata.normalize("NFKC", text)
+        text = unicodedata.normalize("NFD", text)
+        text = "".join(c for c in text if unicodedata.category(c) != "Mn")
+        text = text.lower()
+        text = cls._WS_RE.sub(" ", text)
+        if text == " ":                      # lone space survives strip()
+            text = cls._SENTINEL
+        text = text.strip()
+        return text.replace(cls._SENTINEL, " ")
+
+    def _decode_id(self, tid: int) -> str:
+        tw = self.tokenizer
+        try:
+            enc = getattr(tw, "enc", None)
+            if enc is not None:
+                return enc.decode([tid])
+            return tw.decode([tid])
+        except Exception:
+            return ""
+
+    def _build_lookup_table(self):
+        key2new, new_tokens = {}, []
+        vocab_size = int(self.tokenizer.vocab_size)
+        lookup = np.empty(vocab_size, dtype=np.int64)
+        for tid in range(vocab_size):
+            if tid in self._special_ids:
+                key = f"__special_{tid}"
+            else:
+                text = self._decode_id(tid)
+                if not text or "\ufffd" in text:
+                    key = f"__raw_{tid}"
+                else:
+                    norm = self._normalize(text)
+                    key = norm if norm else f"__raw_{tid}"
+            nid = key2new.get(key)
+            if nid is None:
+                nid = len(new_tokens)
+                key2new[key] = nid
+                new_tokens.append(key)
+            lookup[tid] = nid
+        return lookup, len(new_tokens)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# NgramHasher -- static, deterministic, layer-specific, pure torch int64
+# ─────────────────────────────────────────────────────────────────────────────
+
+class NgramHasher(nn.Module):
+    """(B, T) raw ids -> per-layer embedding indices of shape (B, T, n_cols).
+
+    No parameters, no gradients. `lookup` and the hash multipliers `mult` are
+    persistent (state_dict); the prime moduli `mods` are re-derived from config.
+    `mult` comes from NumPy's Generator, which is not guaranteed stable across
+    NumPy releases, so it is saved with the checkpoint instead of re-derived.
+    """
+
+    def __init__(self, layer_ids: Sequence[int], max_ngram: int,
+                 vocab_size_per_ngram: Sequence[int], n_heads: int,
+                 compressed_vocab: int, raw_vocab_size: int, seed: int,
+                 lookup: torch.Tensor | None = None):
+        super().__init__()
+        assert max_ngram >= 2, "engram_max_ngram must be >= 2"
+        assert len(vocab_size_per_ngram) >= max_ngram - 1, (
+            f"engram_vocab_size needs {max_ngram - 1} entries (ngram 2..{max_ngram}), "
+            f"got {len(vocab_size_per_ngram)}"
+        )
+        assert compressed_vocab > 0, "engram_compressed_vocab must be set (>0)"
+
+        self.layer_ids = tuple(int(l) for l in layer_ids)
+        self.max_ngram = int(max_ngram)
+        self.n_heads = int(n_heads)
+        self.n_cols = (self.max_ngram - 1) * self.n_heads
+        self.compressed_vocab = int(compressed_vocab)
+        # Reserved id for left-context padding: distinct from every token class.
+        self.pad_cid = self.compressed_vocab
+
+        if lookup is None:
+            # Every token would hash to class 0 and Engram would silently ignore its input. This is only
+            # valid if a checkpoint (which carries `lookup`) is loaded right after construction.
+            warnings.warn(
+                "NgramHasher built without a lookup table: all tokens map to class 0 until a checkpoint "
+                "containing `lookup` is loaded. Pass engram_lookup from CompressedTokenizer for training.",
+                stacklevel=2)
+            lookup = torch.zeros(int(raw_vocab_size), dtype=torch.int64)
+        else:
+            lookup = lookup.detach().to(torch.int64).clone()
+            assert lookup.numel() == int(raw_vocab_size), (
+                f"lookup has {lookup.numel()} entries, vocab_size={raw_vocab_size}")
+            assert int(lookup.max()) < self.compressed_vocab
+        self.register_buffer("lookup", lookup, persistent=True)
+
+        # Multipliers: r*2+1, bounded so tok*mult can never overflow int64.
+        # (+1 vocab slot for the reserved pad id.)
+        max_long = int(np.iinfo(np.int64).max)
+        M_max = max_long // (self.compressed_vocab + 1)
+        half_bound = max(1, M_max // 2)
+        PRIME_1 = 10007
+        mult = np.empty((len(self.layer_ids), self.max_ngram), dtype=np.int64)
+        for li, layer_id in enumerate(self.layer_ids):
+            g = np.random.default_rng(int(seed + PRIME_1 * int(layer_id)))
+            r = g.integers(low=0, high=half_bound, size=(self.max_ngram,), dtype=np.int64)
+            mult[li] = r * 2 + 1
+        self.register_buffer("mult", torch.from_numpy(mult), persistent=True)
+
+        # Distinct prime moduli per (layer, ngram order, head); unique globally.
+        seen: set = set()
+        self.head_sizes: List[List[int]] = []      # per layer: flat list, len n_cols
+        for _ in self.layer_ids:
+            flat = []
+            for ngram in range(2, self.max_ngram + 1):
+                search_start = int(vocab_size_per_ngram[ngram - 2]) - 1
+                for _h in range(self.n_heads):
+                    found = find_next_prime(search_start, seen)
+                    seen.add(found)
+                    flat.append(found)
+                    search_start = found
+            self.head_sizes.append(flat)
+        self.register_buffer(
+            "mods", torch.tensor(self.head_sizes, dtype=torch.int64), persistent=False)
+        head_offsets = []
+        for sizes in self.head_sizes:
+            offsets = []
+            offset = 0
+            for size in sizes:
+                offsets.append(offset)
+                offset += size
+            head_offsets.append(offsets)
+        self.register_buffer(
+            "head_offsets", torch.tensor(head_offsets, dtype=torch.int64), persistent=False)
+
+    def compress(self, idx: torch.Tensor) -> torch.Tensor:
+        return self.lookup[idx.long()]            # idx is never negative (the -1 ignore index lives in targets)
+
+    def forward(self, idx: torch.Tensor, doc_ids: torch.Tensor | None = None,
+                cache: dict | None = None):
+        """Hash tokens, optionally using cached history for incremental decoding.
+
+        doc_ids (B, T), when supplied, replaces cross-document history with the
+        left-padding id so packed documents hash as if processed separately.
+        """
+        ids = self.compress(idx)                                 # (B, T_new)
+        B, T_new = ids.shape
+        K = self.max_ngram - 1
+
+        if cache is not None:
+            hist_ids = cache.get("hash_hist_ids")
+            hist_doc = cache.get("hash_hist_doc")
+            if hist_ids is None:
+                hist_ids = ids.new_full((B, K), self.pad_cid)
+                hist_doc = (
+                    torch.full((B, K), -1, dtype=doc_ids.dtype, device=doc_ids.device)
+                    if doc_ids is not None else None
+                )
+            else:
+                hist_ids = hist_ids.to(device=ids.device)
+                if hist_doc is not None:
+                    hist_doc = hist_doc.to(device=ids.device)
+            full_ids = torch.cat([hist_ids, ids], dim=1)          # (B, K + T_new)
+            if doc_ids is not None and hist_doc is not None:
+                full_doc = torch.cat([hist_doc, doc_ids], dim=1)
+            else:
+                full_doc = doc_ids
+            cache["hash_hist_ids"] = full_ids[:, -K:].detach()
+            if full_doc is not None:
+                cache["hash_hist_doc"] = full_doc[:, -K:].detach()
+        else:
+            full_ids, full_doc = ids, doc_ids
+
+        T_full = full_ids.shape[1]
+        # Keep only one running hash per layer and one shifted-id tensor per order;
+        # broadcasting all layers and orders at once creates a large int64 product.
+        mixes = [full_ids * self.mult[li, 0] for li in range(len(self.layer_ids))]
+        out = {
+            layer_id: torch.empty(
+                (B, T_full, self.n_cols), dtype=torch.int64, device=full_ids.device)
+            for layer_id in self.layer_ids
+        }
+        for n in range(2, self.max_ngram + 1):
+            shifted = F.pad(full_ids, (n - 1, 0), value=self.pad_cid)[:, :T_full]
+            if full_doc is not None:
+                same = F.pad(full_doc, (n - 1, 0), value=-1)[:, :T_full] == full_doc
+                shifted = torch.where(same, shifted, torch.full_like(shifted, self.pad_cid))
+            lo = (n - 2) * self.n_heads
+            hi = lo + self.n_heads
+            for li, layer_id in enumerate(self.layer_ids):
+                mix = torch.bitwise_xor(mixes[li], shifted * self.mult[li, n - 1])
+                mixes[li] = mix
+                col = mix.unsqueeze(-1) % self.mods[li, lo:hi]
+                col += self.head_offsets[li, lo:hi]
+                out[layer_id][..., lo:hi] = col
+        if T_full > T_new:
+            out = {layer_id: values[:, -T_new:].contiguous() for layer_id, values in out.items()}
+        return out
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Modules
+# ─────────────────────────────────────────────────────────────────────────────
+
+class _RMSNorm(nn.Module):
+    """RMSNorm with fused accumulation for the input dtype."""
+
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.rms_norm(x, (x.shape[-1],), self.weight.to(x.dtype), self.eps)
+
+
+class ShortConv(nn.Module):
+    """(B, T, D) -> (B, T, D). Depthwise causal conv, dilation = max_ngram."""
+
+    def __init__(self, hidden_size: int, kernel_size: int = 4,
+                 dilation: int = 1, activation: bool = True):
+        super().__init__()
+        self.activation = activation
+        self.kernel_size = kernel_size
+        self.dilation = dilation
+        self.conv = nn.Conv1d(
+            hidden_size, hidden_size, kernel_size=kernel_size,
+            groups=hidden_size, bias=False,
+            padding=(kernel_size - 1) * dilation, dilation=dilation,
+        )
+        self.norm = _RMSNorm(hidden_size)
+        self.act_fn = nn.SiLU() if activation else None
+
+    def forward(self, x: torch.Tensor, doc_ids: torch.Tensor | None = None,
+                cache: dict | None = None) -> torch.Tensor:
+        T = x.shape[1]
+        K, d = self.kernel_size, self.dilation
+        rf = (K - 1) * d
+        B, _, D = x.shape
+        xn = self.norm(x)                                        # (B, T, D)
+
+        if cache is not None:
+            hist_xn = cache.get("sc_hist_x")
+            hist_doc = cache.get("sc_hist_doc")
+            if hist_xn is None:
+                hist_xn = xn.new_zeros(B, rf, D)
+                hist_doc = (
+                    torch.full((B, rf), -1, dtype=doc_ids.dtype, device=doc_ids.device)
+                    if doc_ids is not None else None
+                )
+            else:
+                hist_xn = hist_xn.to(dtype=xn.dtype, device=xn.device)
+                if hist_doc is not None:
+                    hist_doc = hist_doc.to(device=xn.device)
+            full_xn = torch.cat([hist_xn, xn], dim=1) if rf > 0 else xn
+            if doc_ids is not None and hist_doc is not None:
+                full_doc = torch.cat([hist_doc, doc_ids], dim=1) if rf > 0 else doc_ids
+            else:
+                full_doc = doc_ids
+            if rf > 0:
+                cache["sc_hist_x"] = full_xn[:, -rf:].detach()
+                if full_doc is not None:
+                    cache["sc_hist_doc"] = full_doc[:, -rf:].detach()
+        else:
+            full_xn, full_doc = xn, doc_ids
+
+        T_full = full_xn.shape[1]
+        w = self.conv.weight[:, 0, :].to(full_xn.dtype)          # (D, K)
+        xp = F.pad(full_xn, (0, 0, rf, 0))
+
+        tap = xp[:, :T_full]
+        if full_doc is not None and K > 1:
+            s = (K - 1) * d
+            same = F.pad(full_doc, (s, 0), value=-1)[:, :T_full] == full_doc
+            tap = tap * same.unsqueeze(-1)
+        y = tap * w[:, 0]
+        for j in range(1, K):
+            tap = xp[:, j * d:j * d + T_full]
+            if full_doc is not None:
+                s = (K - 1 - j) * d
+                same = F.pad(full_doc, (s, 0), value=-1)[:, :T_full] == full_doc
+                tap = tap * same.unsqueeze(-1)
+            y.addcmul_(tap, w[:, j])
+        if T_full > T:
+            y = y[:, -T:].contiguous()
+        if self.activation:
+            y = self.act_fn(y)
+        return y
+
+
+class MultiHeadEmbedding(nn.Module):
+    def __init__(self, list_of_N: List[int], D: int):
+        super().__init__()
+        self.embedding = nn.Embedding(sum(list_of_N), D)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.normal_(self.embedding.weight, std=0.01)
+
+    def forward(self, hash_ids: torch.Tensor) -> torch.Tensor:
+        return self.embedding(hash_ids)
+
+
+class Engram(nn.Module):
+    """Static n-gram memory with conditional gated injection (hc_mult = 1).
+
+    forward() returns the DELTA to add to the residual stream:
+        x = x + engram(x, hash_ids)
+    """
+
+    def __init__(self, hidden_size: int, head_sizes: List[int], max_ngram: int,
+                 embed_per_ngram: int, n_heads: int, kernel_size: int):
+        super().__init__()
+        assert embed_per_ngram % n_heads == 0, (
+            f"engram_embed_per_ngram={embed_per_ngram} must be divisible by "
+            f"engram_n_heads={n_heads}")
+        self.hidden_size = hidden_size
+        d_head = embed_per_ngram // n_heads
+        engram_hidden = (max_ngram - 1) * embed_per_ngram
+
+        self.multi_head_embedding = MultiHeadEmbedding(list(head_sizes), d_head)
+        self.value_proj = nn.Linear(engram_hidden, hidden_size, bias=False)
+        self.key_proj = nn.Linear(engram_hidden, hidden_size, bias=False)
+        self.norm1 = _RMSNorm(hidden_size)
+        self.norm2 = _RMSNorm(hidden_size)
+        self.short_conv = ShortConv(hidden_size, kernel_size=kernel_size,
+                                    dilation=max_ngram)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.multi_head_embedding.reset_parameters()
+        nn.init.normal_(self.value_proj.weight, std=0.02)
+        nn.init.normal_(self.key_proj.weight, std=0.02)
+        # Zero-init the conv: ShortConv RMS-normalises its input, so without this the gated value is
+        # blown back up to ~1 rms and the initial delta is ~0.13 rms (larger than the residual stream at
+        # D<=256). With zeros the branch starts at ~0 and the conv weight gets gradient on step 1.
+        nn.init.zeros_(self.short_conv.conv.weight)
+
+    def forward(self, hidden_states: torch.Tensor, hash_ids: torch.Tensor,
+                doc_ids: torch.Tensor | None = None, cache: dict | None = None):
+        emb = self.multi_head_embedding(hash_ids).flatten(start_dim=-2)  # (B,T,engram_hidden)
+        key = self.key_proj(emb)
+        nk = self.norm1(key)
+        nq = self.norm2(hidden_states)
+        # Product AND sum in fp32: with bf16 inputs, rounding nk*nq to bf16 first leaves gate errors up to
+        # ~0.02 (sqrt(|g|) amplifies small pre-activations); the casts fuse under torch.compile.
+        gate = (nk.float() * nq.float()).sum(dim=-1) / math.sqrt(self.hidden_size)
+        gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign()
+        gate = gate.sigmoid().unsqueeze(-1).to(hidden_states.dtype)      # (B,T,1)
+        value = gate * self.value_proj(emb)
+        return value + self.short_conv(value, doc_ids, cache=cache)