model_v2.py

"""
model_v2.py -- SpikeWhaleLM v2: optimized base architecture.

Changes vs model.py (v1):

  PERFORMANCE
  - SparseMoEFFN: sort-based expert dispatch (one contiguous slice per expert,
    index_add_ scatter-back) replaces per-expert boolean masking. Far fewer
    kernel launches, torch.compile-friendly (no data-dependent boolean
    indexing in the hot path).
  - Shared experts fused into ONE ExpertFFN with n_shared * intermediate width
    (mathematically equivalent to the averaged sum, 1 matmul set instead of N).

  QUALITY / STABILITY
  - QK-Norm: per-head RMSNorm on Q and K before RoPE (Gemma2/OLMo2-style).
    Stabilizes attention logits, tolerates higher LR. (cfg.use_qk_norm, default ON)
  - z-loss on lm_head logits: zloss_coef * mean(log^2 Z). Prevents logit drift.
    (cfg.zloss_coef, default 1e-4; set 0 to disable)
  - MTP heads REDESIGNED: instead of K independent full H x V matrices (which at
    50M params dwarfed the model), each MTP head is now a small zero-init H x H
    projection feeding the SHARED lm_head. Param cost per head: H^2 instead of
    H*V. MTP loss is down-weighted by cfg.mtp_loss_weight (default 0.3).
  - HC output: learned softmax mix over streams (HCOutputMix) instead of mean().
  - Value-embedding residual (nanoGPT-speedrun style): per-layer learned gate
    (zero-init => exact no-op at init) adds a projection of the token embedding
    into each block's input. (cfg.use_value_embed, default OFF = opt-in)

  All new config keys are read with getattr(cfg, key, default) so your existing
  config.py works unmodified. NOTE: QK-Norm and HCOutputMix add parameters, so
  v1 checkpoints need load_state_dict(strict=False) (new params keep init;
  QK-Norm at init is NOT identity -- prefer training v2 from scratch, or set
  use_qk_norm=False to stay v1-loadable).

  XSA is kept byte-identical to v1 but read the note in MLADerfXSAAttention:
  with num_kv_heads == 1 it removes the SAME rank-1 value subspace from every
  head. A/B it at 50M before keeping it in the final base.
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Tuple, List
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast
from torch.utils.checkpoint import checkpoint as gradient_checkpoint

from config import SpikeWhaleConfig
from fractal import fractal_rope_inv_freq


# ---------------------------------------------------------------------------
# Primitives
# ---------------------------------------------------------------------------

def _attention_blocked(mask, batch, queries, keys, device):
    """Canonical causal mask. 2D: binary keep mask; 4D: SDPA bool/additive."""
    causal = torch.triu(torch.ones(queries, keys, dtype=torch.bool, device=device),
                        diagonal=keys - queries + 1)[None, None]
    if mask is None:
        return causal
    mask = mask.to(device)
    if mask.ndim == 2:
        if tuple(mask.shape) != (batch, keys):
            raise ValueError("2D attention_mask must cover the complete key sequence")
        if not torch.all((mask == 0) | (mask == 1)):
            raise ValueError("2D attention_mask must contain only zero and one")
        blocked = ~mask.bool()[:, None, None, :]
    elif mask.ndim == 4:
        if mask.shape[0] not in (1, batch) or mask.shape[-1] != keys or mask.shape[-2] not in (1, queries):
            raise ValueError("4D attention_mask shape does not match attention")
        if mask.dtype == torch.bool:
            blocked = ~mask
        else:
            if not torch.all((mask == 0) | (mask < -1)):
                raise ValueError("Only zero/negative blocking additive masks are supported")
            blocked = mask < -1
    else:
        raise ValueError("attention_mask must be a 2D keep mask or 4D attention mask")
    return blocked | causal


def _masked_ce(logits, labels):
    valid = labels != -100
    if not valid.any():
        return logits.sum() * 0.0
    return F.cross_entropy(logits[valid], labels[valid])


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight


class RotaryEmbedding(nn.Module):
    """RoPE for the rope partition of Q and K (qk_rope_head_dim dims only)."""

    def __init__(self, dim: int, max_positions: int = 4096, theta: float = 10000.0,
                 inv_freq: Optional[torch.Tensor] = None):
        super().__init__()
        if inv_freq is None:
            inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        t = torch.arange(max_positions).float()
        freqs = torch.outer(t, inv_freq)
        self.register_buffer("cos_cache", freqs.cos())
        self.register_buffer("sin_cache", freqs.sin())

    def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor:
        cos = self.cos_cache[position_ids].unsqueeze(1)   # [B, 1, S, rope_dim//2]
        sin = self.sin_cache[position_ids].unsqueeze(1)
        d = cos.shape[-1]
        x1, x2 = x[..., :d], x[..., d:]
        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)


# ---------------------------------------------------------------------------
# Engram: N-gram hash lookup + DERF gate (unchanged from v1)
# ---------------------------------------------------------------------------

class TokenCompressor(nn.Module):
    """Frozen random projection into the Engram hash space.

    The output is L2-NORMALIZED (v3). The downstream hash is a sign-bit LSH,
    which is scale-invariant on its own, but normalizing here means the hash
    input distribution no longer depends on `hidden_size`, `initializer_range`,
    or how far training has drifted the embedding norms -- so bucket occupancy
    stays uniform for the whole run instead of only at init.
    """
    def __init__(self, embed_dim: int, compress_dim: int):
        super().__init__()
        self.proj = nn.Linear(embed_dim, compress_dim, bias=False)
        # 1/sqrt(fan_in) so the projection preserves scale rather than shrinking
        # it by ~0.02*sqrt(embed_dim); the normalize below makes this cosmetic,
        # but it keeps the pre-norm magnitude O(1) and readable in diagnostics.
        nn.init.normal_(self.proj.weight, std=embed_dim ** -0.5)
        # Frozen LSH-style projection: gradient never reaches it through the
        # integer hash cast, so a fixed random projection is correct (see v1).
        self.proj.weight.requires_grad_(False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        z = self.proj(x)
        return z / (z.norm(dim=-1, keepdim=True) + 1e-6)


class MultiHeadHashLookup(nn.Module):
    """Multi-head n-gram hash memory.

    HASH (v3, fixed). The v1/v2 hash was `h.abs().long() % table_size` where `h`
    was a dot product of magnitude ~0.1. `.long()` truncates toward zero, so
    EVERY token in every n-gram order hashed to bucket 0: the table collapsed to
    a single row, only that row ever received gradient, and Engram degenerated
    into a constant additive bias. Rescaling `h` would paper over it -- the
    truncation threshold is an absolute constant while `h`'s scale tracks
    hidden_size, initializer_range, and training drift, so the same collapse
    returns whenever those change.

    The fix is a real sign-bit LSH: hash by the SIGN of `hash_bits` random
    hyperplane projections, packed into an integer. Signs are invariant to any
    positive rescaling of the input, so the hash cannot silently collapse again
    no matter how the upstream scale moves. Occupancy is verifiable at any time
    via `bucket_occupancy()`.

    Set `legacy_trunc=True` ONLY to load a pre-v3 checkpoint bit-identically.
    """
    def __init__(self, num_heads: int, table_size: int,
                 compress_dim: int, out_dim: int, max_ngram: int = 3,
                 legacy_trunc: bool = False):
        super().__init__()
        self.num_heads = num_heads
        self.table_size = table_size
        self.max_ngram = max_ngram
        self.out_dim = out_dim
        self.legacy_trunc = legacy_trunc
        # Hash width. Addressing the table needs ceil(log2(table_size)) bits,
        # but using EXACTLY that many makes `% table_size` non-uniform whenever
        # table_size is not a power of two: with table_size=300 the 9-bit hash
        # spans 512 values, so buckets 0-211 are reachable twice and 212-299
        # once -- a measured 1.85x occupancy bias. 8 spare bits cut the worst-case
        # bias to 1 + table_size/2^bits (<0.4%), and for a power-of-two table
        # 2^bits divides table_size exactly, so it stays perfectly uniform.
        # Capped at 24: the bits are packed via a float32 accumulation, and
        # integers are only exact up to 2^24 there.
        self.hash_bits = min(24, max(1, math.ceil(math.log2(max(table_size, 2))) + 8))

        self.tables = nn.ModuleList([
            nn.Embedding(table_size, out_dim) for _ in range(num_heads)
        ])
        for t in self.tables:
            nn.init.normal_(t.weight, std=0.01)

        # Bit weights 2^b for packing sign bits into a bucket index.
        self.register_buffer(
            "bit_weights",
            (2.0 ** torch.arange(self.hash_bits, dtype=torch.float32)),
            persistent=False,
        )

        for n in range(1, max_ngram + 1):
            for k in range(n):
                if legacy_trunc:
                    proj = torch.randn(num_heads, compress_dim)
                    proj = proj / (proj.norm(dim=1, keepdim=True) + 1e-8)
                else:
                    # [num_heads, hash_bits, compress_dim] -- one hyperplane per
                    # (head, bit). Distinct per position k so the hash is
                    # ORDER-SENSITIVE: "a b" and "b a" land in different buckets.
                    proj = torch.randn(num_heads, self.hash_bits, compress_dim)
                    proj = proj / (proj.norm(dim=-1, keepdim=True) + 1e-8)
                self.register_buffer(f"hash_proj_n{n}_p{k}", proj)

    def _bucket_ids(self, compressed: torch.Tensor, n: int, valid_len: int) -> torch.Tensor:
        """-> [B, valid_len, num_heads] int64 bucket indices for order-n grams."""
        if self.legacy_trunc:
            h = None
            for k in range(n):
                proj = getattr(self, f"hash_proj_n{n}_p{k}")
                term = torch.matmul(compressed[:, k:k + valid_len, :].float(), proj.t())
                h = term if h is None else h + term
            return h.abs().long() % self.table_size

        # Accumulate the hyperplane projections across the n gram positions
        # BEFORE taking the sign, so the bit pattern depends on the whole gram.
        h = None
        for k in range(n):
            proj = getattr(self, f"hash_proj_n{n}_p{k}")   # [H, bits, D]
            # [B, valid_len, D] x [D, H*bits] -> [B, valid_len, H, bits]
            term = torch.matmul(
                compressed[:, k:k + valid_len, :].float(),
                proj.reshape(self.num_heads * self.hash_bits, -1).t(),
            ).view(-1, valid_len, self.num_heads, self.hash_bits)
            h = term if h is None else h + term
        bits = (h > 0).float()                                   # scale-invariant
        idx = (bits * self.bit_weights).sum(dim=-1).long()
        return idx % self.table_size

    @torch.no_grad()
    def bucket_occupancy(self, compressed: torch.Tensor) -> dict:
        """Diagnostic: fraction of the table actually addressed, per n-gram order.

        A healthy hash gives `distinct / min(tokens, table_size)` near 1.0. The
        v1/v2 collapse showed up here as distinct == 1. Used by the regression
        test in tests/test_engram_hash.py and cheap enough to call during
        training warmup.
        """
        S = compressed.shape[1]
        stats = {}
        for n in range(1, self.max_ngram + 1):
            if S < n:
                continue
            idx = self._bucket_ids(compressed, n, S - n + 1)
            distinct = int(idx.unique().numel())
            stats[f"n{n}"] = {
                "distinct_buckets": distinct,
                "table_size": self.table_size,
                "occupancy": distinct / min(idx.numel(), self.table_size),
            }
        return stats

    def forward(self, compressed: torch.Tensor,
                prefix: Optional[torch.Tensor] = None) -> torch.Tensor:
        """`prefix` is the compressed tail of the preceding tokens (see
        EngramModule.retrieve). Output is only for the `compressed` positions.

        Without it, incremental decoding passes S=1 and every n-gram order with
        n > 1 is skipped by the `S < n` guard -- so a generated token sees only
        its own unigram while the same token in a full forward sees its 2- and
        3-grams. That is a silent train/decode mismatch; it was invisible only
        because the Engram gate used to be a hard zero.
        """
        if prefix is not None and prefix.shape[1] > 0:
            P = prefix.shape[1]
            compressed = torch.cat([prefix.to(compressed.dtype), compressed], dim=1)
        else:
            P = 0

        B, S, _ = compressed.shape
        device = compressed.device
        out = torch.zeros(B, S, self.out_dim, device=device, dtype=compressed.dtype)
        norm = torch.zeros(S, device=device)

        for n in range(1, self.max_ngram + 1):
            if S < n:
                continue
            valid_len = S - n + 1
            start = n - 1

            idx = self._bucket_ids(compressed, n, valid_len)

            for head_idx, table in enumerate(self.tables):
                out[:, start:, :] = out[:, start:, :] + table(idx[:, :, head_idx]).to(out.dtype)

            norm[start:] += self.num_heads

        out = out / norm.view(1, -1, 1).clamp(min=1)
        return out[:, P:, :].to(compressed.dtype)


class DERFContextGate(nn.Module):
    def __init__(self, obs_size: int, init_bias: float = -4.0):
        super().__init__()
        self.proj = nn.Linear(obs_size * 2, obs_size)
        self.alpha = nn.Parameter(torch.ones(obs_size))
        self.bias = nn.Parameter(torch.full((obs_size,), init_bias))
        self.gamma = nn.Parameter(torch.ones(obs_size))

    def forward(self, retrieved: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
        logits = self.proj(torch.cat([retrieved, x], dim=-1))
        gate = self.gamma * ((torch.erf(self.alpha * logits + self.bias) + 1.0) / 2.0)
        return retrieved * gate


class EngramModule(nn.Module):
    def __init__(self, cfg: SpikeWhaleConfig):
        super().__init__()
        self.compressor = TokenCompressor(cfg.hidden_size, cfg.engram_compress_dim)
        self.lookup = MultiHeadHashLookup(
            cfg.engram_num_heads, cfg.engram_table_size,
            cfg.engram_compress_dim, cfg.hidden_size, cfg.engram_max_ngram,
            legacy_trunc=getattr(cfg, "engram_legacy_hash", False),
        )
        self.gate = DERFContextGate(cfg.hidden_size, cfg.engram_gate_init_bias)

    def compress(self, x: torch.Tensor) -> torch.Tensor:
        """Compressed hash-space vectors, one per position.

        These depend ONLY on the token embedding (x is detached and the
        projection is frozen), never on context -- which is what makes it valid
        to cache the tail of them and reuse it as the n-gram prefix on the next
        decode step.
        """
        return self.compressor(x.detach())

    def retrieve(self, x: torch.Tensor,
                 prefix: Optional[torch.Tensor] = None) -> torch.Tensor:
        """Raw (ungated) n-gram memory read. Exposed so mid-network fusion
        layers can reuse the SAME retrieval instead of hashing again."""
        return self.lookup(self.compress(x), prefix=prefix)

    def forward(self, x: torch.Tensor, retrieved: Optional[torch.Tensor] = None) -> torch.Tensor:
        if retrieved is None:
            retrieved = self.retrieve(x)
        return self.gate(retrieved, x)


class EngramLayerFusion(nn.Module):
    """Mid-stack re-injection of the Engram n-gram read (Nanbeige NgramLayerFusion).

    Engram injects n-gram evidence once, at the embedding, where it competes
    with everything else the first layers need to do. Nanbeige re-injects the
    same n-gram features again partway up the stack, gated by the AGREEMENT
    between an n-gram-derived key and the current hidden state -- so a layer
    only pulls surface-form evidence back in when its own representation is
    already pointing that way.

    Bottlenecked to `fusion_dim` so the cost is 3*H*fusion_dim per fusion layer
    rather than 3*H^2. `value_proj` is zero-init => exact no-op at init.
    """
    def __init__(self, hidden_size: int, fusion_dim: int, eps: float = 1e-6):
        super().__init__()
        self.fusion_dim = fusion_dim
        self.hidden_down = nn.Linear(hidden_size, fusion_dim, bias=False)
        self.key_proj = nn.Linear(hidden_size, fusion_dim, bias=False)
        self.value_proj = nn.Linear(hidden_size, fusion_dim, bias=False)
        self.out_proj = nn.Linear(fusion_dim, hidden_size, bias=False)
        self.hidden_norm = RMSNorm(fusion_dim, eps)
        self.ngram_norm = RMSNorm(fusion_dim, eps)
        for m in (self.hidden_down, self.key_proj, self.value_proj):
            nn.init.normal_(m.weight, std=0.02)
        nn.init.zeros_(self.out_proj.weight)

    def forward(self, hidden: torch.Tensor, ngram: torch.Tensor) -> torch.Tensor:
        key = self.ngram_norm(self.key_proj(ngram))
        h = self.hidden_norm(self.hidden_down(hidden))
        gate = (h * key).sum(-1, keepdim=True) / math.sqrt(self.fusion_dim)
        # Signed sqrt compresses the dynamic range of the agreement score before
        # the sigmoid, so the gate does not saturate on a few high-norm tokens.
        gate = torch.sigmoid(gate.abs().clamp_min(1e-6).sqrt() * torch.sign(gate))
        return hidden + self.out_proj(gate * self.value_proj(ngram))


# ---------------------------------------------------------------------------
# Hyper-Connections
# ---------------------------------------------------------------------------

class HyperConnectionLayer(nn.Module):
    """Simplified HC: softmax pre-mix / post-distribute over hc_mult streams.
    Asymmetric init (v1 bugfix) so streams diverge and gradients flow."""
    def __init__(self, hidden_size: int, hc_mult: int,
                 sinkhorn_iters: int = 20, eps: float = 1e-6):
        super().__init__()
        self.hc_mult = hc_mult
        self.pre_weight = nn.Parameter(
            torch.linspace(0.5, -0.5, hc_mult) / max(hc_mult, 1)
        )
        self.post_weight = nn.Parameter(
            torch.linspace(-0.5, 0.5, hc_mult) / max(hc_mult, 1)
        )

    def pre_op(self, copies: torch.Tensor) -> torch.Tensor:
        w = F.softmax(self.pre_weight, dim=0)
        return (copies * w.view(1, -1, 1, 1)).sum(dim=1)

    def post_op(self, copies: torch.Tensor, delta: torch.Tensor) -> torch.Tensor:
        w = F.softmax(self.post_weight, dim=0)
        return copies + delta.unsqueeze(1) * w.view(1, -1, 1, 1)


class HCOutputMix(nn.Module):
    """
    NEW (v2): learned combination of the hc_mult streams at the model output,
    replacing the v1 mean(dim=1). Mean forces the streams toward redundancy at
    exactly the point where you want them specialized. Initialized uniform so
    it starts identical to mean() -- a strict generalization, zero risk.
    """
    def __init__(self, hc_mult: int):
        super().__init__()
        self.weight = nn.Parameter(torch.zeros(hc_mult))   # softmax(0)=uniform=mean

    def forward(self, copies: torch.Tensor) -> torch.Tensor:
        w = F.softmax(self.weight, dim=0)
        return (copies * w.view(1, -1, 1, 1)).sum(dim=1)


# ---------------------------------------------------------------------------
# MLA + (DERF) + XSA Attention, now with QK-Norm
# ---------------------------------------------------------------------------

class MLADerfXSAAttention(nn.Module):
    """
    v2 additions:
      - QK-Norm (cfg.use_qk_norm, default True): per-head RMSNorm applied to Q
        and K BEFORE the rope/nope split. Bounds attention logits, the standard
        modern stability fix; composes cleanly with SDPA and partial RoPE.

    XSA NOTE (unchanged mechanics, important caveat): with num_kv_heads == 1
    (MQA) every query head shares the same value vector, so the self-projection
    subtraction removes the SAME rank-1 value subspace from all heads -- much
    more aggressive than per-head XSA. Ablate use_xsa on/off at 50M before
    locking the base config.
    """

    def __init__(self, cfg: SpikeWhaleConfig, layer_idx: int = 0):
        super().__init__()
        self.layer_idx = layer_idx
        self.num_heads = cfg.num_attention_heads
        self.num_kv_heads = cfg.num_key_value_heads
        self.head_dim = cfg.head_dim
        self.qk_rope_head_dim = cfg.qk_rope_head_dim
        self.nope_head_dim = cfg.nope_head_dim
        self.hidden_size = cfg.hidden_size
        self.use_derf = cfg.use_derf
        self.use_xsa = cfg.use_xsa
        self.use_elo = getattr(cfg, "use_elo", False)
        self.dropout_p = cfg.attention_dropout
        self.kv_groups = self.num_heads // self.num_kv_heads
        self.use_qk_norm = getattr(cfg, "use_qk_norm", True)

        self.q_a_proj = nn.Linear(cfg.hidden_size, cfg.q_lora_rank, bias=False)
        self.q_a_norm = RMSNorm(cfg.q_lora_rank, cfg.rms_norm_eps)
        self.q_b_proj = nn.Linear(cfg.q_lora_rank, self.num_heads * self.head_dim, bias=False)

        self.k_proj = nn.Linear(cfg.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(cfg.hidden_size, self.num_kv_heads * self.head_dim, bias=False)

        self.o_a_proj = nn.Linear(self.num_heads * self.head_dim, cfg.o_lora_rank, bias=False)
        self.o_b_proj = nn.Linear(cfg.o_lora_rank, cfg.hidden_size, bias=False)

        # QK-Norm: one RMSNorm over head_dim, shared across heads (Gemma-2 style).
        if self.use_qk_norm:
            self.q_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps)
            self.k_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps)

        # Fractal RoPE: same endpoints/band count as standard RoPE, but the
        # exponents sit on a Cantor spectrum (fractal.py). gamma=0 == standard.
        rope_inv_freq = None
        if getattr(cfg, "use_fractal_rope", False):
            rope_inv_freq = fractal_rope_inv_freq(
                self.qk_rope_head_dim, cfg.rope_theta, cfg.fractal_rope_gamma)
        self.rope = RotaryEmbedding(
            self.qk_rope_head_dim,
            max_positions=cfg.max_position_embeddings,
            theta=cfg.rope_theta,
            inv_freq=rope_inv_freq,
        )

        if self.use_derf:
            self.derf_alpha = nn.Parameter(torch.ones(self.num_heads))
            self.derf_bias = nn.Parameter(torch.zeros(self.num_heads))
            self.derf_gamma = nn.Parameter(torch.ones(self.num_heads))

        # --- Depth Attention (Nanbeige-4.2 port) ---
        # Mixes this layer's V with the V of earlier checkpointed layers, using
        # softmax over the DEPTH axis of q.k_layer. Parameter-free except the
        # zero-init gate, which makes it an exact no-op at initialization.
        self.use_depth_attention = getattr(cfg, "use_depth_attention", False)
        self.depth_attention_stride = max(1, int(getattr(cfg, "depth_attention_stride", 4)))
        # Layer 0 has no earlier layer to mix from, so its gate could never
        # receive gradient. Don't allocate one rather than ship a parameter that
        # is dead by construction.
        self.depth_gate = (
            nn.Parameter(torch.zeros(1))
            if (self.use_depth_attention and layer_idx > 0) else None
        )

        if self.use_elo:
            # per-head Elo K-factor (softplus>0) and rating->logit gain (tanh,
            # zero-init => no-op at start, so the tournament reduces to plain
            # softmax attention).
            self.elo_log_k = nn.Parameter(torch.full(
                (self.num_heads,), math.log(math.expm1(cfg.elo_k_init))))
            self.elo_gain = nn.Parameter(torch.zeros(self.num_heads))

        for m in (self.q_a_proj, self.q_b_proj, self.k_proj,
                  self.v_proj, self.o_a_proj, self.o_b_proj):
            nn.init.normal_(m.weight, std=cfg.initializer_range)

    def _depth_mix_value(self, q, k, v, depth_kv):
        """Cross-layer value mixing (Nanbeige `_apply_depth_attention`).

        q: [B, num_heads, S, hd] (this layer's queries, pre-cache)
        k, v: [B, num_kv_heads, S, hd] (this layer's fresh K/V, post-RoPE)
        depth_kv: tuple of (k_l, v_l) from earlier checkpointed layers, same shape.

        For each (batch, kv-head, position) we softmax q.k over the LAYER axis
        -- current layer plus each source layer -- and take the corresponding
        convex combination of values. So a token can pull forward the value it
        had built at a shallower depth when that depth's key is a better match:
        a content-addressed skip along depth, complementing the sequence-axis
        attention. Returns the mixed V, shape unchanged.

        Queries are mean-pooled down to the KV-head grouping first, because K/V
        live in num_kv_heads (=2 under GQA) while Q lives in num_heads (=16).
        """
        if not depth_kv:
            return v
        # [B, num_heads, S, hd] -> [B, num_kv_heads, S, hd]
        if self.num_heads != self.num_kv_heads:
            q_g = q.view(q.shape[0], self.num_kv_heads, self.kv_groups,
                         q.shape[2], q.shape[3]).mean(dim=2)
        else:
            q_g = q

        keys = torch.stack([kk for kk, _ in depth_kv] + [k], dim=0)      # [L+1,B,kvH,S,hd]
        vals = torch.stack([vv for _, vv in depth_kv] + [v], dim=0)
        logits = (q_g.unsqueeze(0).float() * keys.float()).sum(-1)       # [L+1,B,kvH,S]
        probs = torch.softmax(logits * (self.head_dim ** -0.5), dim=0).to(vals.dtype)
        mixed = (probs.unsqueeze(-1) * vals).sum(dim=0)
        # Zero-init gate => `mixed` contributes nothing at step 0, so turning
        # depth attention on never perturbs an existing run's starting point.
        return v + torch.tanh(self.depth_gate) * (mixed - v)

    def _elo_attention(self, q, k, v, is_masked, past_rating):
        """Elo/Bradley-Terry tournament attention (ported from wheelerv2).

        Softmax attention already gives each key its win-probability against the
        field; Elo adds a PERSISTENT rating r_j that accumulates across queries
        and biases future logits (winners get boosted), like an AlphaFold-style
        branch-until-winner reinforcement.

          e_ij   = q_i.k_j/sqrt(d)  +  gain * rating_j   (rating = prior + intra-call)
          match  : S_ij = softmax weight key j won at query i ; E_i = 1/n_i baseline
          update : dR_ij = K * (S_ij - E_i)
          rating_j(query i) = sum over earlier queries t<i of dR_tj  (causal prefix sum)

        Two passes keep it feed-forward: pass A scores on PRIOR ratings only to
        define the matches; pass B adds the causal in-call accumulation. The
        summed update is written back to the KV cache so the tournament continues
        through incremental decoding. Zero-init `elo_gain` => pass B == pass A ==
        plain softmax, so enabling Elo is an exact no-op at init.

        Returns (y, rating_out) with rating_out: [B, num_heads, N]."""
        B, H, S, _ = q.shape
        N = k.shape[2]

        device = q.device
        scale = 1.0 / math.sqrt(self.head_dim)
        scores = torch.matmul(q, k.transpose(-2, -1)) * scale        # [B,H,S,N]

        # Take the working dtype from `scores`, NOT from q: under autocast the
        # matmul downcasts to bf16/fp16 while q stays fp32, and fp32's finfo.min
        # (-3.40282e38) overflows bf16's max finite value -- masked_fill then
        # raises "value cannot be converted to BFloat16 without overflow".
        dtype = scores.dtype
        keep = ~is_masked                                            # broadcastable to [.,.,S,N]
        keepf = keep.to(dtype)
        neg = torch.finfo(dtype).min
        K = F.softplus(self.elo_log_k).view(1, H, 1, 1)              # per-head K-factor >0
        gain = torch.tanh(self.elo_gain).view(1, H, 1, 1)            # per-head, 0 at init

        # prior rating of each key = what it accumulated BEFORE this forward
        # (from the cache) + 0 for the keys introduced this call.
        if past_rating is not None:
            prior = torch.cat(
                [past_rating, past_rating.new_zeros(B, H, S)], dim=-1)   # [B,H,N]
        else:
            prior = scores.new_zeros(B, H, N)
        prior_bias = gain * prior.unsqueeze(2)                       # [B,H,1,N]

        masked = scores.masked_fill(is_masked, neg)
        # Match outcome = innate q.k compatibility (the "game result"), NOT the
        # accumulated rating. Keeping matches rating-INDEPENDENT is what makes
        # the parallel prefix-sum (tril @ dR, training) bit-identical to the
        # incremental running sum (decoding); rating is the REPUTATION that
        # biases attention, not what defines the match.
        a_base = torch.softmax(masked, dim=-1)                       # [B,H,S,N]
        n_i = keepf.sum(-1, keepdim=True).clamp(min=1.0)             # valid keys per query
        dR = (K * (a_base - 1.0 / n_i)) * keepf
        # rating seen by query i = prior + causal sum of dR from earlier queries t<i.
        # The wheelerv2 original wrote this as tril(S,S,-1) @ dR, which is O(S^2*N)
        # -- cubic in sequence length, ~1e11 FLOPs/batch-head at S=N=4096 and far
        # too slow to train on. That matmul is exactly an EXCLUSIVE prefix sum over
        # the query axis, so cumsum computes the same tensor in O(S*N).
        intra = torch.cumsum(dR, dim=2) - dR                          # [B,H,S,N]
        a = torch.softmax(masked + prior_bias + gain * intra, dim=-1)
        if self.dropout_p > 0 and self.training:
            a = F.dropout(a, p=self.dropout_p)
        y = torch.matmul(a, v)                                       # [B,H,S,hd]
        # rating carried forward = prior + everything this call's queries contributed
        rating_out = prior + dR.sum(2)                               # [B,H,N]
        return y, rating_out

    def forward(
        self,
        x: torch.Tensor,
        position_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        use_cache: bool = False,
        depth_kv: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[Tuple]]:
        B, S, _ = x.shape

        q = self.q_a_norm(self.q_a_proj(x))
        q = self.q_b_proj(q).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)

        k = self.k_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)

        # QK-Norm before RoPE (v2). Cache stores the NORMALIZED k so prefill and
        # incremental decode agree.
        if self.use_qk_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)

        q_nope = q[..., :self.nope_head_dim]
        q_rope = q[..., self.nope_head_dim:]
        k_nope = k[..., :self.nope_head_dim]
        k_rope = k[..., self.nope_head_dim:]

        q_rope = self.rope(q_rope, position_ids)
        k_rope = self.rope(k_rope, position_ids)

        q = torch.cat([q_nope, q_rope], dim=-1)
        k = torch.cat([k_nope, k_rope], dim=-1)

        # Depth attention operates on the FRESH K/V for this position block,
        # before the KV cache is concatenated -- every layer's contribution is
        # then the same [B, kvH, S, hd] shape, and the mixed V is what gets
        # cached, so incremental decode inherits the mix (Nanbeige does the same
        # ordering: mix, then `past_key_value.update`).
        depth_entry = None
        if self.use_depth_attention:
            # Layer 0 CONSUMES nothing (no earlier layer, hence no gate) but
            # still PRODUCES a source for the layers above it.
            if self.depth_gate is not None:
                v = self._depth_mix_value(q, k, v, depth_kv or ())
            if self.layer_idx % self.depth_attention_stride == 0:
                depth_entry = (k, v)

        past_rating = None
        if past_key_value is not None:
            if len(past_key_value) == 3:          # (k, v, elo_rating) cache
                past_rating = past_key_value[2]
            k = torch.cat([past_key_value[0], k], dim=2)
            v = torch.cat([past_key_value[1], v], dim=2)
        present = (k, v) if use_cache else None
        N = k.shape[2]
        # Preserve causality for explicit padding masks and cached DERF chunks.
        if attention_mask is not None or (self.use_derf and past_key_value is not None):
            blocked = _attention_blocked(attention_mask, B, S, N, q.device)
            attention_mask = torch.zeros_like(blocked, dtype=q.dtype).masked_fill(blocked, float('-inf'))

        if self.kv_groups > 1:
            k = k.unsqueeze(2).expand(-1, -1, self.kv_groups, -1, -1).reshape(
                B, self.num_heads, N, self.head_dim)
            v = v.unsqueeze(2).expand(-1, -1, self.kv_groups, -1, -1).reshape(
                B, self.num_heads, N, self.head_dim)

        if self.use_elo:
            # Elo tournament over keys. Needs explicit scores, so it takes the
            # materialized-mask path rather than SDPA; it precedes DERF.
            if attention_mask is None and past_key_value is None:
                is_masked = torch.triu(
                    torch.ones(S, N, dtype=torch.bool, device=q.device),
                    diagonal=N - S + 1,
                ).unsqueeze(0).unsqueeze(0)
            else:
                is_masked = (attention_mask < -1.0) if attention_mask is not None \
                    else torch.triu(
                        torch.ones(S, N, dtype=torch.bool, device=q.device),
                        diagonal=N - S + 1,
                    ).unsqueeze(0).unsqueeze(0)
            y, rating_out = self._elo_attention(q, k, v, is_masked, past_rating)
            if use_cache:
                present = (present[0], present[1], rating_out)
        elif self.use_derf:
            scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)

            if attention_mask is None and past_key_value is None:
                is_masked = torch.triu(
                    torch.ones(S, N, dtype=torch.bool, device=scores.device),
                    diagonal=N - S + 1,
                ).unsqueeze(0).unsqueeze(0)
            else:
                is_masked = (attention_mask < -1.0) if attention_mask is not None \
                    else torch.zeros_like(scores, dtype=torch.bool)

            safe_scores = scores.masked_fill(is_masked, -10000.0)

            a = self.derf_alpha.view(1, -1, 1, 1)
            b = self.derf_bias.view(1, -1, 1, 1)
            g = self.derf_gamma.view(1, -1, 1, 1)

            attn_weights = g * torch.erf(a * safe_scores + b)
            attn_weights = (attn_weights + g) / 2.0
            attn_weights = attn_weights.masked_fill(is_masked, 0.0)
            attn_weights = attn_weights / (attn_weights.sum(dim=-1, keepdim=True) + 1e-8)

            if self.dropout_p > 0 and self.training:
                attn_weights = F.dropout(attn_weights, p=self.dropout_p)

            y = torch.matmul(attn_weights, v)
        else:
            q = q.contiguous()
            k = k.contiguous()
            v = v.contiguous()
            drop = self.dropout_p if self.training else 0.0
            if past_key_value is None and attention_mask is None:
                y = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=drop)
            else:
                if attention_mask is not None:
                    is_masked = (attention_mask < -1.0)
                else:
                    is_masked = torch.triu(
                        torch.ones(S, N, dtype=torch.bool, device=q.device),
                        diagonal=N - S + 1,
                    ).unsqueeze(0).unsqueeze(0)
                y = F.scaled_dot_product_attention(
                    q, k, v, attn_mask=~is_masked, dropout_p=drop)

        if self.use_xsa:
            past_len = N - S
            v_self = v[:, :, past_len:past_len + S, :]
            vn = v_self / (v_self.norm(dim=-1, keepdim=True) + 1e-8)
            projection = (y * vn).sum(dim=-1, keepdim=True) * vn
            y = y - projection

        y = y.transpose(1, 2).contiguous().view(B, S, self.num_heads * self.head_dim)
        y = self.o_b_proj(self.o_a_proj(y))
        return y, present, depth_entry


# ---------------------------------------------------------------------------
# MoE FFN -- v2: sort-based dispatch + fused shared expert
# ---------------------------------------------------------------------------

class ExpertFFN(nn.Module):
    """Single SwiGLU expert."""
    def __init__(self, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


def sqrtsoftplus(x: torch.Tensor) -> torch.Tensor:
    return torch.sqrt(F.softplus(x) + 1e-8)


class SparseMoEFFN(nn.Module):
    """
    v2 changes:
      - FUSED shared expert: one ExpertFFN with width n_shared * intermediate,
        scaled by 1/n_shared on output -- equivalent to v1's averaged Python
        loop, one fused matmul set. (state-dict key changes: shared_expert.*)
      - SORT-BASED dispatch for routed experts: flatten (token, slot) pairs,
        argsort by expert id, run each expert on ONE contiguous slice, weighted
        index_add_ back. No boolean masks, no nonzero(), no per-expert scatter.
    Routing logic (hash routing, sqrtsoftplus, aux loss) is unchanged.
    """
    def __init__(self, cfg: SpikeWhaleConfig, layer_idx: int = 0):
        super().__init__()
        self.n_routed_experts = cfg.n_routed_experts
        self.n_shared_experts = cfg.n_shared_experts
        self.num_experts_per_tok = cfg.num_experts_per_tok
        self.norm_topk_prob = cfg.norm_topk_prob
        self.scoring_func = cfg.scoring_func
        self.routed_scaling_factor = cfg.routed_scaling_factor
        self.use_hash_routing = layer_idx < cfg.num_hash_layers
        self.aux_loss_coef = cfg.moe_aux_loss_coef

        self.router = nn.Linear(cfg.hidden_size, cfg.n_routed_experts, bias=False)
        if self.use_hash_routing:
            # The first `num_hash_layers` layers route by POSITION, so this
            # Linear is never called and would sit in the optimizer forever with
            # a zero gradient. Kept allocated (state-dict keys stay compatible
            # with existing checkpoints) but frozen, so it is honestly excluded
            # from training rather than silently dead.
            self.router.weight.requires_grad_(False)
        self.experts = nn.ModuleList([
            ExpertFFN(cfg.hidden_size, cfg.moe_intermediate_size)
            for _ in range(cfg.n_routed_experts)
        ])
        # Fused shared expert (v2)
        self.shared_expert = (
            ExpertFFN(cfg.hidden_size,
                      cfg.moe_intermediate_size * cfg.n_shared_experts)
            if cfg.n_shared_experts > 0 else None
        )

        self._last_aux_loss: Optional[torch.Tensor] = None

    def forward(self, x: torch.Tensor,
                position_ids: Optional[torch.Tensor] = None) -> torch.Tensor:
        B, S, H = x.shape
        x_flat = x.view(B * S, H)
        T = B * S
        K = self.num_experts_per_tok

        # Shared expert: always active, single fused pass.
        if self.shared_expert is not None:
            shared_out = self.shared_expert(x_flat)
            if self.n_shared_experts > 1:
                shared_out = shared_out / self.n_shared_experts
        else:
            shared_out = None

        # ---- Routing (unchanged logic) ----
        if self.use_hash_routing:
            if position_ids is not None:
                base = (position_ids.reshape(T, 1) % self.n_routed_experts).long()
            else:
                base = (torch.arange(T, device=x.device) % self.n_routed_experts).unsqueeze(1)
            offsets = torch.arange(K, device=x.device)
            top_k_indices = (base + offsets.unsqueeze(0)) % self.n_routed_experts  # [T, K]
            top_k_weights = torch.full((T, K), 1.0 / K, device=x.device, dtype=x_flat.dtype)
            self._last_aux_loss = None
        else:
            router_logits = self.router(x_flat)
            if self.scoring_func == "sqrtsoftplus":
                routing_scores = sqrtsoftplus(router_logits)
            else:
                routing_scores = F.softmax(router_logits, dim=-1)

            top_k_scores, top_k_indices = torch.topk(routing_scores, K, dim=-1)
            if self.norm_topk_prob:
                top_k_weights = top_k_scores / (top_k_scores.sum(dim=-1, keepdim=True) + 1e-8)
            else:
                top_k_weights = top_k_scores
            top_k_weights = top_k_weights * self.routed_scaling_factor

            softmax_probs = F.softmax(router_logits, dim=-1)
            expert_mask = torch.zeros_like(softmax_probs)
            expert_mask.scatter_(1, top_k_indices, 1.0)
            f_e = expert_mask.mean(0)
            p_e = softmax_probs.mean(0)
            self._last_aux_loss = self.n_routed_experts * (f_e * p_e).sum() * self.aux_loss_coef

        # ---- Sort-based dispatch (v2) ----
        # Flatten the (token, slot) assignment: T*K rows total.
        flat_expert = top_k_indices.reshape(-1)                       # [T*K]
        flat_weight = top_k_weights.reshape(-1, 1)                    # [T*K, 1]
        flat_token = torch.arange(T, device=x.device).repeat_interleave(K)  # [T*K]

        order = torch.argsort(flat_expert, stable=True)               # group by expert
        sorted_expert = flat_expert[order]
        sorted_token = flat_token[order]
        sorted_weight = flat_weight[order]

        counts = torch.bincount(sorted_expert, minlength=self.n_routed_experts)
        # boundaries per expert in the sorted order (CPU sync once per forward;
        # unavoidable without grouped-GEMM, still vastly cheaper than v1's
        # per-expert nonzero/masking)
        counts_list = counts.tolist()

        gathered = x_flat[sorted_token]                               # [T*K, H]
        out_flat = torch.zeros_like(x_flat)

        start = 0
        for expert_idx, cnt in enumerate(counts_list):
            if cnt == 0:
                continue
            end = start + cnt
            seg = gathered[start:end]
            seg_out = self.experts[expert_idx](seg) * sorted_weight[start:end]
            out_flat.index_add_(0, sorted_token[start:end], seg_out.to(out_flat.dtype))
            start = end

        if shared_out is not None:
            out_flat = out_flat + shared_out
        return out_flat.view(B, S, H)

    def get_aux_loss(self) -> Optional[torch.Tensor]:
        return self._last_aux_loss


class DenseFFN(nn.Module):
    def __init__(self, cfg: SpikeWhaleConfig):
        super().__init__()
        self.gate_proj = nn.Linear(cfg.hidden_size, cfg.moe_intermediate_size, bias=False)
        self.up_proj = nn.Linear(cfg.hidden_size, cfg.moe_intermediate_size, bias=False)
        self.down_proj = nn.Linear(cfg.moe_intermediate_size, cfg.hidden_size, bias=False)

    def forward(self, x: torch.Tensor,
                position_ids: Optional[torch.Tensor] = None) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))

    def get_aux_loss(self) -> Optional[torch.Tensor]:
        return None


# ---------------------------------------------------------------------------
# Transformer block
# ---------------------------------------------------------------------------

class MemoryCacheBranch(nn.Module):
    """Memory Caching (arXiv 2602.24281): a linear-attention memory branch with
    growing cross-segment memory, adapted to sit PARALLEL to softmax MLA.

    The sequence is split into N segments of length L. For a token t in segment s
    the Gated Residual Memory readout is

        y_t = gamma_{t,s} * M_t^(s)(q_t)          # own segment, CAUSAL
            + sum_{i<s} gamma_{t,i} * M_L^(i)(q_t) # earlier segments, cached (full)

    where each M^(i) is a linear-attention memory  M^(i) = sum_{m in i} phi(k_m) v_m^T,
    readout is phi(q_t) M / (phi(q_t) z + eps), and the gates are
        gamma_{t,i} = softmax_i( <u_t, meanpool(x over segment i)> ),  i <= s.

    phi = elu(.)+1 (Katharopoulos linear-transformer feature map). Heads are small
    (mc_num_heads x mc_head_dim) so the branch is a cheap add-on. Cross-segment
    readout is O(N^2 L) -- fine at the sequence lengths used here; it is not a
    long-context speed win, it is the faithful gated-residual formulation.
    """
    def __init__(self, cfg: SpikeWhaleConfig):
        super().__init__()
        self.h = cfg.mc_num_heads
        self.d = cfg.mc_head_dim
        self.L = cfg.mc_segment_len
        self.dg = cfg.mc_gate_dim
        self.eps = 1e-6
        H, inner = cfg.hidden_size, self.h * self.d
        self.q_proj = nn.Linear(H, inner, bias=False)
        self.k_proj = nn.Linear(H, inner, bias=False)
        self.v_proj = nn.Linear(H, inner, bias=False)
        self.o_proj = nn.Linear(inner, H, bias=False)
        self.u_proj = nn.Linear(H, self.dg, bias=False)   # per-token gate query
        self.p_proj = nn.Linear(H, self.dg, bias=False)   # per-segment gate key
        self._st = {}   # cache_idx -> incremental segment state

    @staticmethod
    def _phi(x):
        return F.elu(x) + 1.0

    def forward(self, x, position_ids=None, attention_mask=None, state_key=0, use_cache=True):
        # The linear-attention denominator underflows in bf16 as the branch's
        # projections grow, so the branch runs in fp32 with autocast disabled.
        #
        # state_key is the plan slot (cache_idx): unique per (layer, loop pass).
        # A looped stack calls this module once per pass on the same positions with
        # different activations, so each pass needs its own segment state.
        orig = x.dtype
        pos0 = None
        if position_ids is not None and position_ids.numel():
            pos0 = int(position_ids.reshape(position_ids.shape[0], -1)[0, 0])
        if not hasattr(self, "_st") or self._st is None:
            self._st = {}
        st = self._st.get(state_key)
        with torch.autocast(device_type=x.device.type, enabled=False):
            xf = x.float()
            if self.training or not use_cache:
                # Training/recomputed prefixes must neither read nor replace
                # incremental decode state. Capture also duplicates projections.
                y = self._mc_forward(xf, position_ids, attention_mask)
            elif (pos0 is not None and st is not None
                    and st["npos"] == pos0 and st["B"] == x.shape[0]):
                valid = self._valid_mask(xf, attention_mask)
                y = torch.cat([self._step(xf[:, i:i+1], state_key, valid[:, i])
                               for i in range(x.shape[1])], dim=1)
            else:
                y = self._mc_forward(xf, position_ids, attention_mask)
                if pos0 == 0 or st is None or st.get("B") != x.shape[0]:
                    self._capture(xf, state_key, attention_mask)
                else:
                    self._st.pop(state_key, None)   # non-contiguous: state is stale
        return y.to(orig)

    @staticmethod
    def _valid_mask(x, attention_mask):
        if attention_mask is None:
            return x.new_ones(x.shape[:2])
        if attention_mask.ndim == 4:
            keys = attention_mask.shape[-1]
            blocked = _attention_blocked(attention_mask, x.shape[0], x.shape[1], keys, x.device)
            allowed = ~blocked
            valid = allowed.any(dim=1).any(dim=1)
            expected = _attention_blocked(valid, valid.shape[0], x.shape[1], keys, x.device)
            if not torch.equal(blocked.expand_as(allowed), expected.expand_as(allowed)):
                raise ValueError("MC supports causal/key-padding masks, not arbitrary query-specific masks")
            return valid.expand(x.shape[0], -1)[:, -x.shape[1]:].to(x.dtype)
        if attention_mask.ndim != 2 or attention_mask.shape[0] != x.shape[0]:
            raise ValueError("MC requires a [batch, tokens] binary attention mask")
        if attention_mask.shape[1] < x.shape[1]:
            raise ValueError("MC attention mask is shorter than the input")
        return attention_mask[:, -x.shape[1]:].to(device=x.device, dtype=x.dtype)

    def _mc_forward(self, x, position_ids=None, attention_mask=None):
        B, S, H = x.shape
        h, d, L = self.h, self.d, self.L
        # A short prefix contains only one segment. Avoid allocating a full
        # segment-sized attention matrix for a handful of tokens.
        L = min(L, S)
        valid = self._valid_mask(x, attention_mask)
        pad = (L - S % L) % L
        if pad:
            x = F.pad(x, (0, 0, 0, pad))
            valid = F.pad(valid, (0, pad))
        Sp = S + pad
        N = Sp // L

        def heads(proj):
            return proj(x).view(B, N, L, h, d).permute(0, 3, 1, 2, 4)  # [B,h,N,L,d]

        q = self._phi(heads(self.q_proj))
        k = self._phi(heads(self.k_proj))
        vm = valid.view(B, N, L)
        k = k * vm[:, None, :, :, None]
        v = heads(self.v_proj)                                          # [B,h,N,L,d]

        # --- own segment, causal within-segment linear attention (UNnormalized) ---
        A = torch.einsum('bhnld,bhnmd->bhnlm', q, k)                    # [B,h,N,L,L]
        causal = torch.tril(torch.ones(L, L, device=x.device, dtype=A.dtype))
        A = A * causal
        num_intra = torch.einsum('bhnlm,bhnmd->bhnld', A, v)           # [B,h,N,L,d]  (no /Z)
        den_intra = A.sum(-1)                                           # [B,h,N,L]

        # --- cached full-segment memories, cross readout for every (query,key) seg ---
        M = torch.einsum('bhnmd,bhnme->bhnde', k, v)                    # [B,h,N,d,d]
        z = k.sum(3)                                                    # [B,h,N,d]
        num = torch.einsum('bhqld,bhide->bhqile', q, M)               # [B,h,Nq,Ni,L,d] UNnormalized
        den = torch.einsum('bhqld,bhid->bhqil', q, z)                # [B,h,Nq,Ni,L]

        # diagonal (query seg == key seg) uses the CAUSAL running memory, not the
        # full-segment memory. Keep numerators/denominators UNnormalized here.
        idx = torch.arange(N, device=x.device)
        num = num.clone(); den = den.clone()
        num[:, :, idx, idx] = num_intra.to(num.dtype)                 # [B,h,N,L,d]
        den[:, :, idx, idx] = den_intra.to(den.dtype)                 # [B,h,N,L]

        # --- gates: gamma_{t,i} = softmax_i <u_t, meanpool_i>, masked to i<=s ---
        u = self.u_proj(x).view(B, N, L, self.dg)                       # [B,Nq,L,dg]
        pp = self.p_proj(x).view(B, N, L, self.dg)                      # [B,Ni,L,dg]
        pp = pp * vm.unsqueeze(-1)
        pool_full = pp.sum(2) / vm.sum(2).clamp_min(1).unsqueeze(-1)
        score = torch.einsum('bqlg,big->bqil', u, pool_full)           # [B,Nq,Ni,L]
        # CAUSALITY: the current-segment (diagonal) gate must pool ONLY tokens
        # 0..l of the segment -- a full-segment mean leaks that segment's future
        # tokens into token l's gate. Past segments (i<q) are wholly earlier, so
        # their full-segment mean is causal already. Use a running mean on the
        # diagonal.
        counts = vm.cumsum(2).clamp_min(1).unsqueeze(-1)
        cpool = pp.cumsum(2) / counts                                   # [B,N,L,dg] causal running mean
        score_diag = torch.einsum('bqlg,bqlg->bql', u, cpool)          # [B,Nq,L]
        score = score.clone()
        score[:, idx, idx] = score_diag.to(score.dtype)                # diagonal i==q -> causal
        seg_ok = (idx[None, :] <= idx[:, None])                         # [Nq,Ni] i<=q
        score = score.masked_fill(~seg_ok[None, :, :, None], float('-inf'))
        gamma = torch.softmax(score, dim=2)                            # over key seg i  [B,Nq,Ni,L]

        # Gated aggregation with a SINGLE GLOBAL normalizer. The paper's Gated
        # Residual weights the UNnormalized per-segment memories by gamma and does
        # NOT normalize each segment -- normalizing per-segment (the earlier bug)
        # erased the query-key match magnitude, so only the coarse mean-pool gate
        # could pick a segment. Here the numerators keep that magnitude; the gate
        # weights segments; one global Z normalizes the aggregate.
        g = gamma[:, None]                                             # [B,1,Nq,Ni,L]
        y_num = (num * g.unsqueeze(-1)).sum(3)                        # [B,h,Nq,L,d]
        y_den = (den * g).sum(3)                                       # [B,h,Nq,L]
        y = y_num / (y_den.unsqueeze(-1) + self.eps)
        y = y.permute(0, 2, 3, 1, 4).reshape(B, Sp, h * d)
        y = self.o_proj(y)
        return y[:, :S]

    # ---- incremental (KV-cache) support, loop-aware ------------------------
    # State is keyed by the plan slot (cache_idx), because a looped stack calls
    # this same module once per pass with different activations. Sharing one slot
    # across passes folds each token in loop_count times and is worse than no fix.
    #
    # Exact, not approximate: M^(i) = sum_m phi(k_m) v_m^T is a plain sum and the
    # own-segment term is tril-masked, so a running sum over m<=l reproduces the
    # full path's einsum for the current token.
    def reset_state(self):
        self._st = {}

    def _tok(self, proj, x_t):
        B = x_t.shape[0]
        return proj(x_t).view(B, self.h, self.d)

    @torch.no_grad()
    def _capture(self, x, key, attention_mask=None):
        B, S, H = x.shape
        h, d, L = self.h, self.d, self.L
        k = self._phi(self.k_proj(x).view(B, S, h, d)).permute(0, 2, 1, 3)
        valid = self._valid_mask(x, attention_mask)
        k = k * valid[:, None, :, None]
        v = self.v_proj(x).view(B, S, h, d).permute(0, 2, 1, 3)
        pp = self.p_proj(x) * valid.unsqueeze(-1)
        nc = S // L
        dev, dt = x.device, x.dtype
        if nc:
            ks = k[:, :, :nc * L].reshape(B, h, nc, L, d)
            vs = v[:, :, :nc * L].reshape(B, h, nc, L, d)
            Ms = torch.einsum('bhnmd,bhnme->bhnde', ks, vs)
            zs = ks.sum(3)
            counts = valid[:, :nc * L].reshape(B, nc, L).sum(2).clamp_min(1)
            pools = pp[:, :nc * L].reshape(B, nc, L, self.dg).sum(2) / counts.unsqueeze(-1)
        else:
            Ms = torch.zeros(B, h, 0, d, d, device=dev, dtype=dt)
            zs = torch.zeros(B, h, 0, d, device=dev, dtype=dt)
            pools = torch.zeros(B, 0, self.dg, device=dev, dtype=dt)
        r = S - nc * L
        if r:
            kr, vr = k[:, :, nc * L:], v[:, :, nc * L:]
            M_run = torch.einsum('bhmd,bhme->bhde', kr, vr)
            z_run = kr.sum(2)
            pp_sum = pp[:, nc * L:].sum(1)
        else:
            M_run = torch.zeros(B, h, d, d, device=dev, dtype=dt)
            z_run = torch.zeros(B, h, d, device=dev, dtype=dt)
            pp_sum = torch.zeros(B, self.dg, device=dev, dtype=dt)
        self._st[key] = dict(Ms=Ms, zs=zs, pools=pools, M_run=M_run, z_run=z_run,
                             pp_sum=pp_sum, valid_n=valid[:, nc * L:].sum(1),
                             n=r, npos=S, B=B)

    def _step(self, x_t, key, valid=None):
        st = self._st[key]
        B = x_t.shape[0]
        h, d, L = self.h, self.d, self.L
        q = self._phi(self._tok(self.q_proj, x_t))
        k = self._phi(self._tok(self.k_proj, x_t))
        if valid is None:
            valid = x_t.new_ones(B)
        k = k * valid[:, None, None]
        v = self._tok(self.v_proj, x_t)
        u = self.u_proj(x_t).view(B, self.dg)
        pp = self.p_proj(x_t).view(B, self.dg) * valid[:, None]

        st["M_run"] = st["M_run"] + torch.einsum('bhd,bhe->bhde', k, v)
        st["z_run"] = st["z_run"] + k
        st["pp_sum"] = st["pp_sum"] + pp
        st["n"] += 1
        st["valid_n"] = st["valid_n"] + valid

        num_d = torch.einsum('bhd,bhde->bhe', q, st["M_run"])
        den_d = torch.einsum('bhd,bhd->bh', q, st["z_run"])
        nc = st["Ms"].shape[2]
        if nc:
            num_c = torch.einsum('bhd,bhide->bhie', q, st["Ms"])
            den_c = torch.einsum('bhd,bhid->bhi', q, st["zs"])
            score_c = torch.einsum('bg,big->bi', u, st["pools"])
        else:
            num_c = num_d.new_zeros(B, h, 0, d)
            den_c = den_d.new_zeros(B, h, 0)
            score_c = u.new_zeros(B, 0)
        score_d = (u * (st["pp_sum"] / st["valid_n"].clamp_min(1)[:, None])).sum(-1, keepdim=True)
        gamma = torch.softmax(torch.cat([score_c, score_d], -1), dim=-1)

        num_all = torch.cat([num_c, num_d.unsqueeze(2)], dim=2)
        den_all = torch.cat([den_c, den_d.unsqueeze(2)], dim=2)
        g = gamma[:, None]
        y_num = (num_all * g.unsqueeze(-1)).sum(2)
        y_den = (den_all * g).sum(2)
        y = y_num / (y_den.unsqueeze(-1) + self.eps)
        y = self.o_proj(y.reshape(B, 1, h * d))

        st["npos"] += 1
        if st["n"] == L:
            st["Ms"] = torch.cat([st["Ms"], st["M_run"].unsqueeze(2)], dim=2)
            st["zs"] = torch.cat([st["zs"], st["z_run"].unsqueeze(2)], dim=2)
            pool = st["pp_sum"] / st["valid_n"].clamp_min(1)[:, None]
            st["pools"] = torch.cat([st["pools"], pool.unsqueeze(1)], dim=1)
            st["M_run"] = torch.zeros_like(st["M_run"])
            st["z_run"] = torch.zeros_like(st["z_run"])
            st["pp_sum"] = torch.zeros_like(st["pp_sum"])
            st["n"] = 0
            st["valid_n"] = torch.zeros_like(st["valid_n"])
        return y


class TransformerBlock(nn.Module):
    def __init__(self, cfg: SpikeWhaleConfig, layer_idx: int):
        super().__init__()
        self.use_hc = cfg.use_hyper_connections
        self.hidden_dropout = cfg.hidden_dropout
        self.use_value_embed = getattr(cfg, "use_value_embed", False)

        self.attn_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.attn = MLADerfXSAAttention(cfg, layer_idx=layer_idx)
        self.ffn_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)

        # Memory Caching branch (arXiv 2602.24281), parallel to attention. Its own
        # norm + zero-init gate => exact no-op at init.
        self.use_mc = getattr(cfg, "use_memory_cache", False)
        if self.use_mc:
            self.mc_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
            self.mc = MemoryCacheBranch(cfg)
            self.mc_gate = nn.Parameter(torch.zeros(1))

        # Mid-stack Engram re-injection, only on the configured layers.
        fusion_layers = getattr(cfg, "engram_fusion_layers", None) or []
        self.engram_fusion = (
            EngramLayerFusion(cfg.hidden_size,
                              getattr(cfg, "engram_fusion_dim", 256),
                              cfg.rms_norm_eps)
            if (cfg.use_engram and layer_idx in fusion_layers) else None
        )

        if cfg.use_moe and layer_idx in cfg.moe_layers:
            self.ffn = SparseMoEFFN(cfg, layer_idx)
            self.is_moe = True
        else:
            self.ffn = DenseFFN(cfg)
            self.is_moe = False

        if self.use_hc:
            self.hc_attn = HyperConnectionLayer(cfg.hidden_size, cfg.hc_mult,
                                                cfg.hc_sinkhorn_iters, cfg.hc_eps)
            self.hc_ffn = HyperConnectionLayer(cfg.hidden_size, cfg.hc_mult,
                                               cfg.hc_sinkhorn_iters, cfg.hc_eps)

        # NEW (v2, opt-in): value-embedding residual. Zero-init gate -> exact
        # no-op at init; learns to mix raw token-embedding signal into each
        # block's input (nanoGPT-speedrun "value embedding"/U-net skip family;
        # consistent wins at the 50-500M scale).
        if self.use_value_embed:
            self.ve_gate = nn.Parameter(torch.zeros(1))

    def forward(
        self,
        x: torch.Tensor,           # [B, hc_mult, S, H] if HC else [B, S, H]
        position_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple] = None,
        use_cache: bool = False,
        token_embed: Optional[torch.Tensor] = None,   # [B, S, H] (value-embed)
        ngram_embed: Optional[torch.Tensor] = None,   # [B, S, H] (Engram fusion)
        depth_kv: Optional[Tuple] = None,             # ((k,v), ...) earlier layers
        mc_key: int = 0,                              # plan slot: per (layer, pass)
    ) -> Tuple[torch.Tensor, Optional[Tuple], Optional[torch.Tensor], Optional[Tuple]]:

        # --- Attention sub-layer ---
        if self.use_hc:
            h = self.hc_attn.pre_op(x)
        else:
            h = x

        if self.use_value_embed and token_embed is not None:
            h = h + torch.tanh(self.ve_gate) * token_embed

        attn_out, present, depth_entry = self.attn(
            self.attn_norm(h), position_ids, attention_mask, past_key_value,
            use_cache, depth_kv,
        )
        attn_out = F.dropout(attn_out, p=self.hidden_dropout, training=self.training)

        # Memory Caching: fold the gated linear-memory readout into the attention
        # sub-layer contribution (zero-init gate => no-op at start). h is [B,S,H].
        if self.use_mc:
            mc_out = self.mc(self.mc_norm(h), position_ids, attention_mask, mc_key,
                             use_cache=use_cache)
            mc_out = F.dropout(mc_out, p=self.hidden_dropout, training=self.training)
            attn_out = attn_out + torch.tanh(self.mc_gate) * mc_out

        if self.use_hc:
            x = self.hc_attn.post_op(x, attn_out)
            h = self.hc_ffn.pre_op(x)
        else:
            h = h + attn_out

        # --- FFN sub-layer ---
        ffn_out = self.ffn(self.ffn_norm(h), position_ids)
        ffn_out = F.dropout(ffn_out, p=self.hidden_dropout, training=self.training)

        if self.use_hc:
            x = self.hc_ffn.post_op(x, ffn_out)
        else:
            x = h + ffn_out

        # Engram re-injection sits AFTER the block, on the residual stream, so
        # it feeds the next layer rather than competing inside this one. Under
        # hyper-connections x is [B, hc_mult, S, H]; broadcast the fusion across
        # streams by folding the stream axis into the batch.
        if self.engram_fusion is not None and ngram_embed is not None:
            if x.dim() == 4:
                B, M, S, H = x.shape
                fused = self.engram_fusion(
                    x.reshape(B * M, S, H),
                    ngram_embed.unsqueeze(1).expand(B, M, S, H).reshape(B * M, S, H),
                )
                x = fused.view(B, M, S, H)
            else:
                x = self.engram_fusion(x, ngram_embed)

        return x, present, self.ffn.get_aux_loss(), depth_entry


# ---------------------------------------------------------------------------
# HRM refinement (unchanged)
# ---------------------------------------------------------------------------

class HRMRefinementBlock(nn.Module):
    """Iterative refinement of the final hidden state.

    DEAD-MODULE FIX. `gate` and `up.weight` were BOTH zero-initialized, and they
    sit on the same multiplicative path:

        h = h + tanh(gate[t]) * up(silu(down(...)))

    so dL/d(up) is proportional to tanh(gate) = 0, dL/d(gate) is proportional to
    up(...) = 0, and dL/d(down) and dL/d(norm) are proportional to up.weight = 0.
    Every tensor in the block received exactly zero gradient on every step: HRM
    refinement was frozen at the identity for the entire run, and deep
    supervision was supervising `hrm_refine_steps` identical copies of x. No
    error, no NaN, just a headline feature that silently did nothing -- the same
    failure mode as the Engram hash collapse.

    Only ONE endpoint of a multiplicative path may be zero-initialized. `up`
    keeps its zero (that is what makes the block an exact no-op at step 0) and
    `gate` now starts at `gate_init`, so gradient reaches `up` immediately and
    the rest of the block unfreezes as soon as `up` leaves zero.
    """
    def __init__(self, hidden_size: int, refine_dim: int, steps: int, eps: float = 1e-6,
                 gate_init: float = 1.0):
        super().__init__()
        self.steps = steps
        self.norm = RMSNorm(hidden_size, eps)
        self.down = nn.Linear(hidden_size * 2, refine_dim, bias=False)
        self.up = nn.Linear(refine_dim, hidden_size, bias=False)
        self.gate = nn.Parameter(torch.full((steps,), float(gate_init)))
        nn.init.normal_(self.down.weight, std=0.02)
        nn.init.zeros_(self.up.weight)

    def forward(self, x: torch.Tensor, return_all: bool = False):
        """return_all=True additionally returns the per-step states [h_1..h_T],
        so every refinement iteration can be supervised (LDT eq. 1). Default
        False keeps the original single-tensor return for existing callers."""
        anchor = x
        h = x
        states = [] if return_all else None
        for t in range(self.steps):
            inp = torch.cat([self.norm(h), anchor], dim=-1)
            update = self.up(F.silu(self.down(inp)))
            h = h + torch.tanh(self.gate[t]) * update
            if return_all:
                states.append(h)
        return (h, states) if return_all else h


# ---------------------------------------------------------------------------
# Full model
# ---------------------------------------------------------------------------

def loop_layer_plan(num_layers: int, loop_count: int, loop_mode: str = "full",
                    loop_middle_layers: Optional[int] = None):
    """Execution order for the looped stack.

    Returns (plan, num_pass_slots) where plan is a list of (layer_idx, pass_slot)
    and pass_slot indexes the per-pass embedding (None = not in a looped region).
    The list position IS the KV cache slot, which keeps cache indexing correct
    for both modes without either one having to know the other's layout.

    "full"          -- the original Byrne-700M-Looped behaviour: the whole stack,
                       `loop_count` times. Effective depth L*loop_count.
    "middle_split"  -- Nanbeige-4.2 `enable_double_loop_split`: run a fixed
                       prefix unlooped, repeat the middle block `loop_count`
                       times, then a fixed suffix unlooped. Effective depth
                       (L - M) + M*loop_count for a middle of width M, at
                       correspondingly lower compute than "full".
    """
    if loop_count <= 1 or loop_mode == "full":
        return ([(l, p if loop_count > 1 else None)
                 for p in range(max(1, loop_count))
                 for l in range(num_layers)], max(1, loop_count))

    if loop_mode != "middle_split":
        raise ValueError(f"unknown loop_mode {loop_mode!r} (expected 'full' or 'middle_split')")

    M = loop_middle_layers if loop_middle_layers is not None else num_layers // 2
    if M <= 0 or num_layers % M != 0:
        raise ValueError(
            f"loop_middle_layers={M} must be positive and divide "
            f"num_hidden_layers={num_layers}"
        )
    start = (num_layers - M) // 2
    end = start + M
    plan = [(l, None) for l in range(start)]
    plan += [(l, p) for p in range(loop_count) for l in range(start, end)]
    plan += [(l, None) for l in range(end, num_layers)]
    return plan, loop_count


class SpikeWhaleModel(nn.Module):
    """Decoder stack without LM head."""

    def __init__(self, cfg: SpikeWhaleConfig):
        super().__init__()
        self.cfg = cfg
        self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
        nn.init.normal_(self.embed_tokens.weight, std=cfg.initializer_range)

        self.engram = EngramModule(cfg) if cfg.use_engram else None
        self.layers = nn.ModuleList([
            TransformerBlock(cfg, layer_idx=i)
            for i in range(cfg.num_hidden_layers)
        ])
        self.norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.hc_out_mix = (
            HCOutputMix(cfg.hc_mult) if cfg.use_hyper_connections else None
        )
        self.hrm_refine = (
            HRMRefinementBlock(cfg.hidden_size, cfg.hrm_refine_dim, cfg.hrm_refine_steps,
                               cfg.rms_norm_eps,
                               gate_init=getattr(cfg, "hrm_gate_init", 1.0))
            if getattr(cfg, "use_hrm_refine", False) else None
        )
        self.use_value_embed = getattr(cfg, "use_value_embed", False)
        self.deep_supervision = (
            getattr(cfg, "hrm_deep_supervision", False) and self.hrm_refine is not None
        )
        # --- looped transformer ---
        self.loop_count = max(1, int(getattr(cfg, "loop_count", 1)))
        # Per-pass embedding so the shared weights can condition on which
        # iteration they're in. Zero-init => loop_count=1 is bit-identical to
        # the dense baseline, and a fresh looped run starts from the same point.
        self.loop_mode = getattr(cfg, "loop_mode", "full")
        self.loop_plan, n_pass_slots = loop_layer_plan(
            cfg.num_hidden_layers, self.loop_count, self.loop_mode,
            getattr(cfg, "loop_middle_layers", None),
        )
        # How many router invocations the plan makes, relative to one pass over
        # the stack -- used to keep the MoE aux-loss coefficient comparable
        # across loop modes (see the normalization in forward()).
        self.loop_router_scale = len(self.loop_plan) / cfg.num_hidden_layers
        # Cache layout: slots [0, _n_cache_slots) are the per-(pass, layer) K/V
        # entries; if Engram is on, one extra trailing slot carries its n-gram
        # context (see forward()).
        self._n_cache_slots = len(self.loop_plan)
        self.loop_pass_embed = None
        if self.loop_count > 1 and getattr(cfg, "loop_pass_embed", True):
            self.loop_pass_embed = nn.Parameter(
                torch.zeros(n_pass_slots, cfg.hidden_size))
        self.use_depth_attention = getattr(cfg, "use_depth_attention", False)

        # The hc_mult residual streams start as IDENTICAL copies of x, and
        # HyperConnectionLayer.pre_op is a softmax-weighted sum over them -- on
        # identical streams that sum is invariant to the weights, so the very
        # first pre-mix in the network has an exactly-zero gradient. Under
        # loop_mode="full" layer 0 runs again on later passes (streams have
        # diverged by then) and it recovers; under "middle_split" layer 0 runs
        # exactly ONCE, so its pre_weight would be dead for the whole run.
        # Freeze it rather than ship a parameter that cannot train.
        if cfg.use_hyper_connections and self.loop_plan[0][0] == 0:
            runs_layer0_once = sum(1 for l, _ in self.loop_plan if l == 0) == 1
            if runs_layer0_once:
                self.layers[0].hc_attn.pre_weight.requires_grad_(False)

        self.gradient_checkpointing = False

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[Tuple]] = None,
        use_cache: bool = False,
    ) -> Tuple[torch.Tensor, Optional[List[Tuple]], torch.Tensor]:
        B, S = input_ids.shape
        device = input_ids.device

        if position_ids is None:
            past_len = past_key_values[0][0].shape[2] if past_key_values else 0
            position_ids = torch.arange(
                past_len, past_len + S, device=device
            ).unsqueeze(0).expand(B, -1)

        x = self.embed_tokens(input_ids)
        token_embed = x if self.use_value_embed else None

        # Read the n-gram memory ONCE and reuse it for both the embedding-level
        # injection and every mid-stack fusion layer (hashing is not free, and
        # re-hashing a drifted hidden state would address different buckets).
        ngram_embed = None
        engram_tail = None
        if self.engram is not None:
            # Incremental decode passes S=1, which would drop every n-gram order
            # above 1. The cache carries the last (max_ngram - 1) compressed
            # vectors of the prefix so the n-grams straddling the boundary are
            # still computed. See MultiHeadHashLookup.forward.
            prefix = None
            if past_key_values is not None and len(past_key_values) > self._n_cache_slots:
                prefix = past_key_values[self._n_cache_slots]
            compressed = self.engram.compress(x)
            ngram_embed = self.engram.lookup(compressed, prefix=prefix)
            x = x + self.engram(x, retrieved=ngram_embed)
            if use_cache:
                keep = max(0, self.cfg.engram_max_ngram - 1)
                full = compressed if prefix is None else \
                    torch.cat([prefix.to(compressed.dtype), compressed], dim=1)
                # max(0, ...) rather than relying on Python's negative-index
                # clamping: a prompt shorter than the n-gram window would
                # otherwise index with a negative start and only work by accident.
                start = max(0, full.shape[1] - keep)
                engram_tail = full[:, start:, :].detach() if keep else None

        if self.cfg.use_hyper_connections:
            x = x.unsqueeze(1).expand(-1, self.cfg.hc_mult, -1, -1).clone()

        present_key_values = [] if use_cache else None
        total_aux_loss = torch.tensor(0.0, device=device)

        # Gradient checkpointing is incompatible with use_cache (the cache from
        # the discarded forward would be silently wrong on recompute).
        assert not (self.gradient_checkpointing and self.training and use_cache), \
            "use_cache=True is not supported with gradient checkpointing"

        # LOOPED TRANSFORMER: run the whole stack `loop_count` times through the
        # SAME weights. Effective depth = loop_count * num_hidden_layers at the
        # parameter cost of num_hidden_layers. loop_count=1 reduces exactly to
        # the dense baseline (the pass embedding is skipped entirely).
        #
        # KV cache: every (pass, layer) pair produces its own K/V, because the
        # same weights see different activations on each pass. The cache is laid
        # out flat as [pass0_layer0 .. pass0_layerN, pass1_layer0 ..], indexed by
        # `pass_idx * num_layers + layer_idx`. Getting this wrong would silently
        # feed pass 1's keys to pass 0 and corrupt every cached decode.
        #
        # The plan (see loop_layer_plan) also covers loop_mode="middle_split",
        # where only the middle block repeats. The plan's list position is the
        # cache slot, so both modes get correct, non-overlapping cache indexing.
        #
        # Depth attention accumulates (k, v) from checkpointed layers as the
        # plan advances. It is RESET at the start of every looped pass: within a
        # pass, depth means "earlier layer"; carrying the list across passes
        # would grow it by loop_count and make pass 3 attend to a stale pass-1
        # value of the same layer, which is a different (and unvalidated)
        # mechanism. The list is threaded through as an explicit argument rather
        # than mutated inside the layer, so gradient checkpointing's recompute
        # cannot double-append.
        prev_slot = None
        depth_kv: List[Tuple[torch.Tensor, torch.Tensor]] = []
        depth_prefix: List[Tuple[torch.Tensor, torch.Tensor]] = []
        for cache_idx, (layer_idx, pass_slot) in enumerate(self.loop_plan):
            layer = self.layers[layer_idx]

            if pass_slot is not None and pass_slot != prev_slot:
                if self.loop_pass_embed is not None:
                    # [H] broadcasts over both [B,S,H] and the hyper-connection
                    # [B,hc_mult,S,H] layout.
                    x = x + self.loop_pass_embed[pass_slot]
                # Entering a new pass: keep whatever the UNLOOPED prefix
                # contributed (it ran once and is still "earlier depth"), drop
                # the previous pass's own entries.
                if prev_slot is None:
                    depth_prefix = list(depth_kv)
                depth_kv = list(depth_prefix)
            prev_slot = pass_slot

            pkv = past_key_values[cache_idx] if past_key_values else None
            dkv = tuple(depth_kv) if self.use_depth_attention else None

            if self.gradient_checkpointing and self.training:
                x, present, aux_loss, depth_entry = gradient_checkpoint(
                    layer, x, position_ids, attention_mask, None, False,
                    token_embed, ngram_embed, dkv, cache_idx,
                    use_reentrant=False,
                )
            else:
                x, present, aux_loss, depth_entry = layer(
                    x, position_ids, attention_mask, pkv, use_cache,
                    token_embed, ngram_embed, dkv, cache_idx)

            if depth_entry is not None:
                depth_kv.append(depth_entry)
            if use_cache:
                present_key_values.append(present)
            if aux_loss is not None:
                total_aux_loss = total_aux_loss + aux_loss

        # Normalize the MoE aux loss by how many times the plan invoked a
        # router relative to a single pass over the stack: looping produces
        # proportionally more router invocations, so without this the aux
        # coefficient would effectively be scaled up in the looped variant and
        # the two runs would not be comparing like with like. (For loop_mode
        # "middle_split" the factor is not loop_count -- hence the plan-derived
        # ratio rather than a hardcoded divide.)
        if self.loop_router_scale > 1.0:
            total_aux_loss = total_aux_loss / self.loop_router_scale

        # Trailing cache slot for the Engram n-gram context (see above).
        if use_cache and engram_tail is not None:
            present_key_values.append(engram_tail)

        if self.cfg.use_hyper_connections:
            x = self.hc_out_mix(x)          # v2: learned mix (init == mean)

        # Deep supervision (LDT): keep each refinement step's state so the LM
        # head can be applied to all of them. Only collected while training --
        # at inference we read the final iteration only, as the paper does.
        refine_states = None
        if self.hrm_refine is not None:
            if self.deep_supervision and self.training:
                x, states = self.hrm_refine(x, return_all=True)
                refine_states = [self.norm(s) for s in states]
            else:
                x = self.hrm_refine(x)

        x = self.norm(x)
        return x, present_key_values, total_aux_loss, refine_states


class MTPHead(nn.Module):
    """
    v2 MTP head: small zero-init H x H projection feeding the SHARED lm_head.
    Cost per head: H^2 params (e.g. 1M at H=1024) instead of H*V (e.g. 50M+).
    Zero-init means at step 0 the head predicts exactly what lm_head predicts
    for the residual path = 0, i.e. uniform-ish gradient pressure; the residual
    form (x + proj(x)) keeps it anchored to the trunk representation.
    """
    def __init__(self, hidden_size: int):
        super().__init__()
        self.proj = nn.Linear(hidden_size, hidden_size, bias=False)
        nn.init.zeros_(self.proj.weight)

    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
        return hidden + self.proj(hidden)


class SpikeWhaleLM(PreTrainedModel):
    """
    v2 loss = CE + zloss_coef * z-loss
            + mtp_loss_weight * mean(MTP CE)
            + MoE aux loss
    """
    config_class = SpikeWhaleConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["TransformerBlock"]

    def __init__(self, cfg: SpikeWhaleConfig):
        super().__init__(cfg)
        self.model = SpikeWhaleModel(cfg)
        self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
        nn.init.normal_(self.lm_head.weight, std=cfg.initializer_range)

        self.zloss_coef = getattr(cfg, "zloss_coef", 1e-4)
        self.mtp_loss_weight = getattr(cfg, "mtp_loss_weight", 0.3)

        # v2 MTP: H x H residual projections sharing lm_head (see MTPHead).
        self.mtp_heads = nn.ModuleList([
            MTPHead(cfg.hidden_size)
            for _ in range(cfg.num_nextn_predict_layers)
        ]) if cfg.num_nextn_predict_layers > 0 else None

        # LDT-adapted abstention head (see config.py / docs/LDT-notes.md).
        # Zero-init => logit 0 => p(abstain) = 0.5 everywhere at step 0, and it
        # contributes no gradient to the trunk until it learns something.
        self.use_abstain_head = getattr(cfg, "use_abstain_head", False)
        self.abstain_loss_weight = getattr(cfg, "abstain_loss_weight", 0.1)
        self.abstain_pos_weight = getattr(cfg, "abstain_pos_weight", 8.0)
        self.abstain_threshold = getattr(cfg, "abstain_threshold", 0.6)
        self.abstain_head = (
            nn.Linear(cfg.hidden_size, 1, bias=True) if self.use_abstain_head else None
        )

        self.post_init()

        # Zero-init AFTER post_init: HF's _init_weights runs inside post_init and
        # would otherwise overwrite these, leaving the head with a random opinion
        # about its own correctness before it has learned anything.
        if self.abstain_head is not None:
            nn.init.zeros_(self.abstain_head.weight)
            nn.init.zeros_(self.abstain_head.bias)

        # Same reason: EngramLayerFusion.out_proj, MTPHead.proj and
        # HRMRefinementBlock.up are all zero-init BY DESIGN (each makes its
        # module an exact no-op at step 0), and post_init() overwrites all three
        # with normal(0, initializer_range). Restore them here.
        for layer in self.model.layers:
            if getattr(layer, "engram_fusion", None) is not None:
                nn.init.zeros_(layer.engram_fusion.out_proj.weight)
        if self.model.hrm_refine is not None:
            nn.init.zeros_(self.model.hrm_refine.up.weight)
        for head in (self.mtp_heads or ()):     # None when MTP is disabled
            nn.init.zeros_(head.proj.weight)

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def tie_weights(self, **kwargs):
        if self.config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight

    def save_pretrained(self, *args, **kwargs):
        tied = (
            self.config.tie_word_embeddings
            and self.lm_head.weight.data_ptr() == self.model.embed_tokens.weight.data_ptr()
        )
        if tied:
            self.lm_head.weight = nn.Parameter(self.model.embed_tokens.weight.detach().clone())
        try:
            super().save_pretrained(*args, **kwargs)
        finally:
            if tied:
                self.lm_head.weight = self.model.embed_tokens.weight

    def _set_gradient_checkpointing(self, module, value=False):
        if isinstance(module, SpikeWhaleModel):
            module.gradient_checkpointing = value

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[Tuple]] = None,
        labels: Optional[torch.Tensor] = None,
        use_cache: bool = False,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        hidden, present_kvs, aux_loss, refine_states = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
        )

        logits = self.lm_head(hidden)
        loss = None

        # LDT abstention signal: p(this prediction is wrong). Exposed on the
        # output so decoding can gate on it; `abstain_threshold` is deliberately
        # above 0.5 (paper raises theta_CLS at inference vs training).
        abstain_logit = None
        if self.abstain_head is not None:
            abstain_logit = self.abstain_head(hidden).squeeze(-1)   # [B, S]

        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            flat_logits = shift_logits.view(-1, shift_logits.size(-1))
            flat_labels = shift_labels.view(-1)
            loss = _masked_ce(flat_logits, flat_labels)

            # Deep supervision (LDT eq. 1): every refinement iteration gets its
            # own CE, all weighted equally, averaged with the final-state CE.
            # refine_states is populated only while training with the flag on.
            if refine_states:
                step_losses = [loss]
                for s in refine_states[:-1]:        # last state == `hidden`
                    s_logits = self.lm_head(s)[..., :-1, :].contiguous()
                    step_losses.append(_masked_ce(
                        s_logits.view(-1, s_logits.size(-1)),
                        flat_labels))
                loss = torch.stack(step_losses).mean()

            # Abstention head: asymmetric BCE against "top-1 was wrong here".
            # pos_weight>1 makes missing a genuine error costlier than a false
            # alarm -- the paper's w+/w- = 8 asymmetry, pointed at the error
            # direction that actually matters for abstention.
            if abstain_logit is not None and self.abstain_loss_weight > 0:
                valid = flat_labels != -100
                if valid.any():
                    with torch.no_grad():
                        wrong = (flat_logits.argmax(-1) != flat_labels).float()
                    a_flat = abstain_logit[..., :-1].contiguous().view(-1)
                    abstain_loss = F.binary_cross_entropy_with_logits(
                        a_flat[valid].float(), wrong[valid],
                        pos_weight=torch.tensor(
                            self.abstain_pos_weight, device=loss.device))
                    loss = loss + self.abstain_loss_weight * abstain_loss

            # z-loss (v2): penalize log^2 of the partition function on valid
            # positions. Keeps logits from drifting; pairs well with Muon.
            if self.zloss_coef > 0:
                valid = flat_labels != -100
                if valid.any():
                    log_z = torch.logsumexp(flat_logits[valid].float(), dim=-1)
                    loss = loss + self.zloss_coef * (log_z ** 2).mean()

            # MTP (v2): residual H x H head -> shared lm_head, down-weighted.
            if self.mtp_heads is not None and self.mtp_loss_weight > 0:
                mtp_total = torch.tensor(0.0, device=loss.device)
                n_active = 0
                for k, head in enumerate(self.mtp_heads, start=1):
                    offset = k + 1
                    if hidden.size(1) > offset:
                        mtp_hidden = head(hidden[..., :-offset, :])
                        mtp_logits = self.lm_head(mtp_hidden)
                        mtp_labels = labels[..., offset:].contiguous()
                        mtp_total = mtp_total + _masked_ce(
                            mtp_logits.reshape(-1, mtp_logits.size(-1)),
                            mtp_labels.reshape(-1),
                        )
                        n_active += 1
                if n_active > 0:
                    loss = loss + self.mtp_loss_weight * mtp_total / n_active

            loss = loss + aux_loss

        out = CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=present_kvs,
        )
        # Attach the abstention signal without breaking the HF output contract.
        # p_abstain[b, t] = P(the token predicted at position t is wrong).
        if abstain_logit is not None:
            out.abstain_logit = abstain_logit
            out.p_abstain = torch.sigmoid(abstain_logit)
        return out

    def count_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters())


def reset_memory_cache(model):
    """Clear incremental MemoryCacheBranch state (all loop passes) on every layer.
    Call between independent generations that reuse one model object."""
    n = 0
    for m in model.modules():
        if isinstance(m, MemoryCacheBranch):
            m.reset_state(); n += 1
    return n