Update model.py

model.py

@@ -26,6 +26,7 @@ _USE_LIGER              = bool(getattr(alpha_config, 'use_liger', False))
 
 _LigerRMSNormFn = None
 try:
+    # pyrefly: ignore [missing-import]
     from liger_kernel.ops.rms_norm import LigerRMSNormFunction as _LigerRMSNormFn
 except ImportError:
     pass
@@ -53,6 +54,15 @@ except ImportError:  # pragma: no cover - fallback for package-style imports
 
 _SIMPLE_BLOCK_PATH = not _USE_FLASH_OPS
 
+try:
+    from ngram import Engram, NgramHasher
+except ImportError:  # pragma: no cover - fallback for package-style imports
+    try:
+        from .ngram import Engram, NgramHasher
+    except ImportError:
+        Engram = None
+        NgramHasher = None
+
 
 
 _alpha_config_warned: set[str] = set()
@@ -98,7 +108,7 @@ _NVFP4_DISABLED_REASON: "str | None" = None
 
 _NVFP4_SKIP_KEYWORDS: tuple[str, ...] = (
     "wte", "lm_head", "ln_", "ln_f", "norm",
-    "router", "gate", "expert_bias",
+    "router", "gate", "expert_bias", "engram",
 )
 
 
@@ -235,7 +245,9 @@ def apply_int8_to_model(model: nn.Module, group_size: int = 128) -> int:
         return True
 
     
-    before_ids = {fqn: id(m) for fqn, m in model.named_modules() if _filter_with_gs(m, fqn)}
+    # torchao replaces module.weight in place and keeps the nn.Linear object, so
+    # identity must be tracked on the weight (same test apply_nvfp4_to_model uses).
+    before_ids = {fqn: id(m.weight) for fqn, m in model.named_modules() if _filter_with_gs(m, fqn)}
 
     
     if _INT8_AVAILABLE and _Int8Config is not None:
@@ -253,13 +265,16 @@ def apply_int8_to_model(model: nn.Module, group_size: int = 128) -> int:
     _torchao_quantize(model, config_obj, filter_fn=_filter_with_gs)
 
     
-    count = sum(
-        1 for fqn, old_id in before_ids.items()
-        if id(dict(model.named_modules()).get(fqn)) != old_id
-    )
+    current = dict(model.named_modules())
+    count = 0
+    for fqn, old_id in before_ids.items():
+        mod = current.get(fqn)
+        if mod is not None and getattr(mod, "weight", None) is not None and id(mod.weight) != old_id:
+            count += 1
     model._quant_mode = "int8"
     model._use_int8 = True
     model._use_nvfp4 = False
+    return count
 
 def _alpha_config_attr(name: str, default):
     if not hasattr(alpha_config, name):
@@ -361,8 +376,9 @@ class RMSNorm(nn.Module):
         self._use_liger = _USE_LIGER and _LigerRMSNormFn is not None
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        
-        if self._use_liger and _liger_cuda_gate(x):
+        # Hoist the liger CUDA guard: x.is_cuda is a tensor bool (no string
+        # comparison), already specialized by dynamo when _use_liger is False.
+        if self._use_liger and x.is_cuda:
             out = _liger_rmsnorm_attempt(x, self.weight, self.eps)
             if out is not None:
                 return out
@@ -554,7 +570,8 @@ def _liger_fused_ce_attempt(hidden: torch.Tensor, weight: torch.Tensor, targets:
     try:
         loss_fn = _LigerFusedLinearCrossEntropyLoss(ignore_index=-1)
         
-        loss = loss_fn(hidden.reshape(-1, hidden.size(-1)), weight, targets.reshape(-1))
+        # LigerFusedLinearCrossEntropyLoss.forward(lin_weight, _input, target): weight FIRST.
+        loss = loss_fn(weight, hidden.reshape(-1, hidden.size(-1)), targets.reshape(-1))
         kernel._record_backend("liger_fused_ce_success")
         return loss
     except Exception as e:  # noqa: BLE001 -- must never crash training
@@ -581,14 +598,26 @@ def chunked_cross_entropy(
     total_loss = torch.zeros((), device=h.device, dtype=torch.float32)
     n_valid = (t != ignore_index).sum().float()
     if n_valid == 0:
-        return total_loss
+        # A bare zeros(()) has no grad_fn, so loss.backward() raised on an all-ignored batch.
+        # Multiplying by 0 keeps a graph edge to `hidden` while contributing zero gradient.
+        return h.float().sum() * 0.0
+
+    def _chunk_loss(h_chunk: torch.Tensor, t_chunk: torch.Tensor) -> torch.Tensor:
+        # Recomputed during backward via checkpoint -- only the scalar loss is
+        # retained, not the (chunk_size, vocab) fp32 logit tensor.
+        logits_chunk = F.linear(h_chunk, weight).float()
+        return F.cross_entropy(logits_chunk, t_chunk, ignore_index=ignore_index, reduction="sum")
+
     for start in range(0, N, chunk_size):
         end = min(start + chunk_size, N)
         h_chunk = h[start:end]                         # (C, D)
         t_chunk = t[start:end]                         # (C,)
-        logits_chunk = F.linear(h_chunk, weight).float()  # (C, vocab) -- freed after backward through this chunk
-        # use sum reduction then normalise manually so partial chunks average correctly
-        loss_chunk = F.cross_entropy(logits_chunk, t_chunk, ignore_index=ignore_index, reduction="sum")
+        # checkpoint recomputes the F.linear in backward; only the scalar
+        # loss_chunk is kept alive in the autograd graph between forward and
+        # backward, so the (C, vocab) fp32 logit tensor is never retained.
+        loss_chunk = torch.utils.checkpoint.checkpoint(
+            _chunk_loss, h_chunk, t_chunk, use_reentrant=False
+        )
         total_loss = total_loss + loss_chunk
     return total_loss / n_valid.clamp(min=1)
 
@@ -700,7 +729,19 @@ class DenseMHA(nn.Module):
         self.dropout = nn.Dropout(dropout)
 
         self.use_xsa = False   # set by GPT.__init__ after reading config
-        
+        self.xsa_alpha = None  # nn.Parameter, per-head gate; set by GPT.__init__ when use_xsa
+
+        self.use_qk_norm = False  # set by GPT.__init__ after reading config
+        self.q_norm = None         # RMSNorm(head_dim); set by GPT.__init__ when use_qk_norm
+        self.k_norm = None         # RMSNorm(head_dim); set by GPT.__init__ when use_qk_norm
+
+        # Gated Attention (Qiu et al., NeurIPS 2025): per-head sigmoid gate
+        # on the SDPA output, before out_proj.  Wired by GPT.__init__.
+        # Zero-init weight -> 2*sigmoid(0) = 1.0 -> gate is identity at step 0.
+        self.use_attn_gate = False
+        self.attn_gate_window = None   # int: how many dims of block input to read
+        self.attn_gate_proj = None     # nn.Linear(gate_window, num_heads, bias=False)
+
         self._xsa_v_reps = local_num_heads // local_num_kv_heads
 
         if use_rope:
@@ -708,6 +749,17 @@ class DenseMHA(nn.Module):
             self.register_buffer("rope_cos", rope_cos, persistent=False)
             self.register_buffer("rope_sin", rope_sin, persistent=False)
 
+        # XSA pos-0 guard mask: registered once at init so forward() never
+        # allocates a fresh (T,) tensor.  Non-persistent: re-created from
+        # max_seq_len on load, costs nothing in the state_dict.
+        self.register_buffer(
+            "_xsa_pos_mask",
+            torch.ones(max_seq_len),
+            persistent=False,
+        )
+        # Slot 0 is always masked off (position-0 has no context beyond itself).
+        self._xsa_pos_mask[0] = 0.0
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -724,6 +776,14 @@ class DenseMHA(nn.Module):
         k = reshape_heads(k_flat, self.num_kv_heads, self.head_dim)
         v = reshape_heads(v_flat, self.num_kv_heads, self.head_dim)
 
+        # QK-Norm: independent of XSA, applied before the rope split. Normalizes
+        # q/k per head (over head_dim) to stabilize attention-logit scale, which
+        # keeps softmax from saturating on a few positions. Off by default --
+        # ablate separately from XSA so gains are attributable.
+        if self.use_qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
         if self.use_rope:
             q, k = apply_rope_qk(
                 q,
@@ -766,15 +826,59 @@ class DenseMHA(nn.Module):
             )
             weights = None
 
+        # ── Gated Attention (Qiu et al., NeurIPS 2025 Best Paper) ───────────────
+        # Per-head sigmoid gate modulates the SDPA output before XSA/out_proj.
+        # Gate input: first `attn_gate_window` dims of the post-norm block input
+        # (hidden_states).  Sparse window keeps the gate head lightweight.
+        # 2*sigmoid(z) is used so the zero-init projection gives gate=1 (identity)
+        # at step 0; the head learns to suppress or amplify over training.
+        if self.use_attn_gate:
+            gw = self.attn_gate_window
+            gate_in = hidden_states[..., :gw]                 # (B, T, gw)
+            gate = self.attn_gate_proj(gate_in)               # (B, T, num_heads)
+            gate = 2.0 * torch.sigmoid(gate)                  # (0, 2), =1.0 at init
+            gate = gate.transpose(1, 2).unsqueeze(-1)         # (B, H, T, 1)
+            context = context * gate
+        # ────────────────────────────────────────────────────────────────────────
+
         if self.use_xsa:
             v_xsa = v
             if self._xsa_v_reps > 1:
-                v_xsa = v.repeat_interleave(self._xsa_v_reps, dim=1)
+                # GQA/MQA-aware broadcast: mathematically identical to
+                # v.repeat_interleave(self._xsa_v_reps, dim=1) (verified via
+                # torch.equal), but avoids materializing a full copy of v.
+                Bv, Hkv, Tv, Dv = v.shape
+                v_xsa = (
+                    v.reshape(Bv, Hkv, 1, Tv, Dv)
+                    .expand(Bv, Hkv, self._xsa_v_reps, Tv, Dv)
+                    .reshape(Bv, Hkv * self._xsa_v_reps, Tv, Dv)
+                )
 
             v_norm = v_xsa.norm(dim=-1, keepdim=True).clamp(min=1e-6)
             v_unit = v_xsa / v_norm
             proj_coeff = (context * v_unit).sum(dim=-1, keepdim=True)
-            context = context - proj_coeff * v_unit
+            # Learnable per-head gate (modded-nanoGPT style): tanh(alpha) in (-1, 1).
+            # alpha starts at 0 -> gate=0 -> forward pass is numerically identical
+            # to vanilla SA at step 0. Each head learns how much (if any) self-value
+            # component to remove, rather than a hard, deterministic subtraction.
+            gate = torch.tanh(self.xsa_alpha).view(1, -1, 1, 1)
+            correction = gate * proj_coeff * v_unit
+
+            # Position-0 guard: causal attention at position 0 has exactly one
+            # reachable key (itself), so context[0] == v[0] identically, at every
+            # step of training. XSA's projection therefore removes 100% of the
+            # context for that head, for any gate > 0. This is structural, not
+            # a training artifact -- there is no "context minus self" at
+            # position 0 because there is no context other than self there.
+            # Mask the correction off so the first token always keeps a
+            # full-strength value contribution.
+            T = context.shape[2]
+            # Slice the pre-built buffer; no allocation per forward. Applied for T == 1 as well:
+            # a 1-token input IS position 0 and must be masked exactly as it is in training.
+            mask = self._xsa_pos_mask[:T].to(dtype=context.dtype).view(1, 1, T, 1)
+            correction = correction * mask
+
+            context = context - correction
 
         merged_context = merge_heads(context)
         # ─────────────────────────────────────────────────────────────────────
@@ -928,13 +1032,17 @@ class AryaSparseMoE(nn.Module):
             return out, self.last_aux_loss
 
         routed_out = torch.zeros_like(x_flat)
+        # Build all expert membership masks at once on the GPU with one_hot:
+        # shape (B*T, n_routed), dtype bool.  No .any() host syncs at all.
+        # one_hot produces (B*T, top_k, n_routed); .any(dim=1) collapses top_k.
+        expert_masks = F.one_hot(idx_flat, num_classes=self.n_routed).bool().any(dim=1)  # (B*T, n_routed)
         for expert_id in range(self.n_routed):
             expert = self.routed_experts[expert_id]
-            token_mask = (idx_flat == expert_id).any(dim=-1)
-            if not token_mask.any():
-                continue
+            token_mask = expert_masks[:, expert_id]   # (B*T,) — already on GPU, no sync
 
             selected_tokens = x_flat[token_mask]
+            if selected_tokens.shape[0] == 0:
+                continue
             expert_out = expert(selected_tokens)  # Returns delta only
             weight = (
                 (idx_flat[token_mask] == expert_id).float() * gate_flat[token_mask]
@@ -1128,7 +1236,10 @@ class AryaBlock(nn.Module):
         
         if _SIMPLE_BLOCK_PATH:
             normed = self.ln_1(x)
-            attn_delta = module_output_tensor(self.attn(normed))
+            # In the simple path output_attentions/return_kv_cache_estimate/
+            # skip_out_proj are all False, so DenseMHA returns a raw tensor --
+            # no isinstance check needed, no dynamo guard on AttentionOutput.
+            attn_delta = self.attn(normed)
             x = x + attn_delta
             ffn_in = self.ln_2(x)
             ffn_result = self.ffn(ffn_in)
@@ -1281,9 +1392,14 @@ class MTPHead(nn.Module):
         self.eps = 1e-5
 
     def forward(self, hidden: torch.Tensor) -> torch.Tensor:
-        h = hidden.float()
-        h_norm = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + self.eps)
-        h_norm = h_norm.to(hidden.dtype) * self.weight
+        # F.rms_norm handles the dtype internally -- no manual fp32 upcast
+        # allocation, no (B, T, D) intermediate tensor, free for any dtype.
+        if _HAS_NATIVE_RMSNORM:
+            h_norm = _native_rms_norm(hidden, (self.weight.shape[0],), self.weight, self.eps)
+        else:  # torch < 2.4: same fp32-upcast fallback as RMSNorm
+            hf = hidden.float()
+            h_norm = (hf / torch.sqrt(hf.pow(2).mean(-1, keepdim=True) + self.eps)).to(hidden.dtype) \
+                * self.weight.to(hidden.dtype)
         return h_norm + self.proj(h_norm)
 
 
@@ -1348,6 +1464,10 @@ class GPTConfig:
     
     use_liger: bool = _alpha_config_attr('use_liger', False)
     use_xsa: bool = _alpha_config_attr('use_xsa', False)   # Exclusive Self-Attention: strips self-referential component from attention output
+    use_qk_norm: bool = _alpha_config_attr('use_qk_norm', False)   # RMSNorm on q/k per head, before rope; independent of use_xsa
+    # Gated Attention (Qiu et al., NeurIPS 2025 Best Paper, arXiv:2505.06708)
+    use_attn_gate: bool = _alpha_config_attr('use_attn_gate', False)  # per-head sigmoid gate after SDPA, before out_proj
+    attn_gate_window: int = _alpha_config_attr('attn_gate_window', 0)  # 0 = full n_embd; 12-32 = sparse speedrun window
     
     moe_top_k: int = _alpha_config_attr('moe_top_k', 2)
 
@@ -1382,10 +1502,24 @@ class GPTConfig:
 
     mod_gate_entropy_coeff: float = _alpha_config_attr('mod_gate_entropy_coeff', 0.01)
 
+    # Engram static n-gram memory (arXiv:2601.07372). Tuples, not lists:
+    # dataclass rejects mutable defaults.
+    use_engram: bool = _alpha_config_attr('use_engram', False)
+    engram_layer_ids: tuple = tuple(_alpha_config_attr('engram_layer_ids', (1,)))
+    engram_max_ngram: int = _alpha_config_attr('engram_max_ngram', 3)
+    engram_vocab_size: tuple = tuple(_alpha_config_attr('engram_vocab_size', (8192, 8192)))
+    engram_embed_per_ngram: int = _alpha_config_attr('engram_embed_per_ngram', 128)
+    engram_n_heads: int = _alpha_config_attr('engram_n_heads', 4)
+    engram_kernel_size: int = _alpha_config_attr('engram_kernel_size', 4)
+    engram_seed: int = _alpha_config_attr('engram_seed', 0)
+    # Set automatically from the tokenizer lookup at build time; saved in
+    # checkpoints so GPT(config) can be rebuilt without a tokenizer.
+    engram_compressed_vocab: int = _alpha_config_attr('engram_compressed_vocab', 0)
+
 class GPT(nn.Module):
     """Full GPT language model"""
 
-    def __init__(self, config):
+    def __init__(self, config, engram_lookup=None):
         super().__init__()
         if config.vocab_size is None:
             raise ValueError("config.vocab_size must be set")
@@ -1441,6 +1575,46 @@ class GPT(nn.Module):
         ))
         self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
 
+        # ── Engram (optional) ────────────────────────────────────────────
+        self.engram_hasher = None
+        self.engram_layers = nn.ModuleDict()
+        self._engram_layer_ids = frozenset()
+        if bool(getattr(config, 'use_engram', False)):
+            if Engram is None:
+                raise ImportError("use_engram=True but engram.py could not be imported")
+            ids = tuple(int(l) for l in config.engram_layer_ids)
+            if not ids or any(not (0 <= l < config.n_layer) for l in ids) or len(set(ids)) != len(ids):
+                raise ValueError(
+                    f"engram_layer_ids={ids} must be unique and within [0, {config.n_layer})")
+            if engram_lookup is not None:
+                config.engram_compressed_vocab = int(engram_lookup.max()) + 1
+            if int(getattr(config, 'engram_compressed_vocab', 0)) <= 0:
+                raise ValueError(
+                    "use_engram=True needs engram_lookup (built from the tokenizer via "
+                    "engram.CompressedTokenizer) or a config with engram_compressed_vocab set")
+            self.engram_hasher = NgramHasher(
+                layer_ids=ids,
+                max_ngram=config.engram_max_ngram,
+                vocab_size_per_ngram=config.engram_vocab_size,
+                n_heads=config.engram_n_heads,
+                compressed_vocab=config.engram_compressed_vocab,
+                raw_vocab_size=config.vocab_size,
+                seed=config.engram_seed,
+                lookup=engram_lookup,
+            )
+            self.engram_layers = nn.ModuleDict({
+                str(l): Engram(
+                    hidden_size=config.n_embd,
+                    head_sizes=self.engram_hasher.head_sizes[i],
+                    max_ngram=config.engram_max_ngram,
+                    embed_per_ngram=config.engram_embed_per_ngram,
+                    n_heads=config.engram_n_heads,
+                    kernel_size=config.engram_kernel_size,
+                )
+                for i, l in enumerate(ids)
+            })
+            self._engram_layer_ids = frozenset(ids)
+
         
         use_mtp = bool(getattr(config, 'use_mtp', False))
         mtp_depth = int(getattr(config, 'mtp_depth', 1))
@@ -1450,6 +1624,8 @@ class GPT(nn.Module):
 
         
         self.apply(self._init_weights)
+        for _eg in self.engram_layers.values():
+            _eg.reset_parameters()   # tables std=0.01 (apply() above overwrote it with 0.02)
 
         
         if config.use_xsa:
@@ -1457,6 +1633,40 @@ class GPT(nn.Module):
                 real_block = block.inner if isinstance(block, MoDBlock) else block
                 if isinstance(real_block.attn, DenseMHA):
                     real_block.attn.use_xsa = True
+                    # Zero-init -> tanh(0) = 0 -> gate is a no-op at step 0, so the
+                    # forward pass starts numerically identical to vanilla SA and
+                    # each head learns its own correction during training.
+                    real_block.attn.xsa_alpha = nn.Parameter(
+                        torch.zeros(real_block.attn.num_heads)
+                    )
+
+        # Separate from use_xsa on purpose: QK-Norm changes attention-logit scale
+        # (and therefore what XSA's projection has to work with), so it needs its
+        # own on/off switch to keep the two ablatable independently.
+        if config.use_qk_norm:
+            for block in self.transformer.h:
+                real_block = block.inner if isinstance(block, MoDBlock) else block
+                if isinstance(real_block.attn, DenseMHA):
+                    real_block.attn.use_qk_norm = True
+                    real_block.attn.q_norm = RMSNorm(real_block.attn.head_dim, eps=1e-6)
+                    real_block.attn.k_norm = RMSNorm(real_block.attn.head_dim, eps=1e-6)
+
+        # Gated Attention (Qiu et al., NeurIPS 2025 Best Paper, arXiv:2505.06708).
+        # Independent switch: keep ablatable from XSA and QK-Norm.
+        # Gate reads the first `attn_gate_window` dims of the post-norm block input.
+        # 0 (or unset) -> use full n_embd (dense gate).
+        if config.use_attn_gate:
+            gw = int(config.attn_gate_window) or int(config.n_embd)
+            gw = min(gw, config.n_embd)
+            for block in self.transformer.h:
+                real_block = block.inner if isinstance(block, MoDBlock) else block
+                if isinstance(real_block.attn, DenseMHA):
+                    real_block.attn.use_attn_gate = True
+                    real_block.attn.attn_gate_window = gw
+                    proj = nn.Linear(gw, real_block.attn.num_heads, bias=False)
+                    # Zero-init: 2*sigmoid(0) = 1.0 -> gate is identity at step 0.
+                    nn.init.zeros_(proj.weight)
+                    real_block.attn.attn_gate_proj = proj
 
         for block in self.transformer.h:
             if isinstance(block, MoDBlock):
@@ -1587,7 +1797,12 @@ class GPT(nn.Module):
         mod_gate_loss = torch.zeros((), device=x.device, dtype=torch.float32)
         moe_aux_loss = torch.zeros((), device=x.device, dtype=torch.float32)
 
-        for block, (is_mod, has_moe) in zip(self.transformer.h, self._block_flags):
+        engram_hashes = self.engram_hasher(idx) if self.engram_hasher is not None else None
+
+        for layer_i, (block, (is_mod, has_moe)) in enumerate(zip(self.transformer.h, self._block_flags)):
+            if engram_hashes is not None and layer_i in self._engram_layer_ids:
+                # Official placement: h = h + Engram(h, ids) BEFORE the block's attention.
+                x = x + self.engram_layers[str(layer_i)](x, engram_hashes[layer_i])
             if self.config.gradient_checkpointing and self.training:
                 if is_mod:
                     (x, moe_aux), gate_aux = torch.utils.checkpoint.checkpoint(