overlay/hydra/model.py

"""PostSemClawModel — full-architecture model assembly.

Extracted from the monolithic train.py (W1 modularization). Semantics
unchanged. Imports `GPUEngram` from `hydra.engram` and `MuonAdamW` from
`hydra.optimizer`.

Triton kernel integration status (Phase 2):
  HYDRA_FUSED_BCNORM — DEFERRED. The bcnorm_fused Triton kernel fuses
    LayerNorm + RoPE on B/C projections. However, mamba-ssm's Mamba3 block
    uses RMSNormGated (not LayerNorm) for B/C, and RoPE is applied inside
    the mamba3_siso_combined CUDA kernel via the Angles parameter. Replacing
    would require either (a) monkey-patching RMSNormGated + intercepting the
    fused CUDA scan — invasive, 50+ lines, high breakage risk — or (b) a
    full custom Mamba3Block reimplementation. Both are out of scope for
    Phase 2. The kernel is validated standalone; integration deferred to
    Phase 3 when HYDRA moves to a custom SSM block.

  HYDRA_FUSED_SSD — DEFERRED. The ssd_exp_trap Triton kernel implements
    exponential-trapezoidal discretization as a sequential scan. mamba-ssm's
    Mamba3 block delegates the entire scan + gating + output projection to
    mamba3_siso_combined (a compiled CUDA kernel with tilelang). Replacing
    it would require decomposing the combined kernel into constituent ops
    and substituting only the scan — not feasible without a custom block.
    Same Phase 3 gate as above.

Both env vars are accepted but currently no-ops (gates read, logged, but
the code path is unchanged). This avoids silent regression if someone
sets them expecting a speedup.
"""

from __future__ import annotations

import os

import torch
import torch.nn as nn
import torch.nn.functional as F

try:
    from mamba_ssm import Mamba3
except ModuleNotFoundError:  # local CPU tests may run outside the HF image wheel stack
    Mamba3 = None

from subsystems.hestia_mini import HestiaQAT
from subsystems.htm import HTMLayer
from subsystems.mhc_mini import ManifoldHyperConnection
from subsystems.sdr_semantic import SemanticFoldingSDR

from hydra.engram import GPUEngram
from hydra.htm_cache import htm_cache_key, htm_cache_matches
from hydra.hyena_block import HyenaBlock
from hydra.reality_bridge import RealityPoincareBridge
# GDNBlock is imported lazily inside __init__ so the `fla` dependency is
# only required when HYDRA_GDN_LAYERS is actually non-empty. Baseline
# pure-Mamba3 runs continue to work without flash-linear-attention installed.
from hydra.optimizer import MuonAdamW
from hydra.sampled_softmax import UnigramSampler, sampled_softmax_loss

try:
    from subsystems.cantor_router import CantorRouter
except ModuleNotFoundError:
    from archive.cantor_router import CantorRouter


def norm(x: torch.Tensor) -> torch.Tensor:
    """RMSNorm over the last dim — stateless, autocast-friendly."""
    return F.rms_norm(x, (x.size(-1),))


def paired_slow_fast_orthogonality(w: torch.Tensor) -> torch.Tensor:
    """Penalty for aligned adjacent slow/fast vector pairs."""
    n = (w.shape[0] // 2) * 2
    if n == 0:
        return w.new_zeros(())
    slow = F.normalize(w[:n:2].float(), dim=-1, eps=1e-8)
    fast = F.normalize(w[1:n:2].float(), dim=-1, eps=1e-8)
    return (slow * fast).sum(dim=-1).square().mean().to(dtype=w.dtype)


def semantic_gaussian_mollify(
    x: torch.Tensor,
    std: float = 0.0,
    training: bool = True,
    eval_enabled: bool = False,
) -> torch.Tensor:
    """Optionally add train-time semantic Gaussian noise; disabled is identity."""
    if std <= 0.0 or (not training and not eval_enabled):
        return x
    return x + torch.randn_like(x) * float(std)


class _LocalMamba3Fallback(nn.Identity):
    """Shape-preserving local fallback used only when mamba_ssm is absent."""
    pass


class PostSemClawModel(nn.Module):
    """Full Post-SEM-Claw model assembly.

    Architecture:
        Token Embedding -> [Mamba3 + residual] x n_layer
        -> SDR + Engram (at configured layer) -> norm -> LM head

    Interface (must match prepare.py evaluate_bpb):
        model(x, y, reduction='none').view(-1)  -> per-token losses
        model(x, y, reduction='mean')           -> scalar loss
    """

    def __init__(self, config):
        super().__init__()
        self.config = config

        # Audit 2026-05-09 issue #18: HTM output caching is only safe when the
        # current SDR pattern matches the cached pattern. Under MDLM each
        # forward can resample a different mask pattern with identical (B, T)
        # shape, so shape-only reuse silently feeds stale HTM features into
        # the next forward. We gate the cache on (a) MDLM-off and (b) an
        # active-index match (htm_cache_matches) below.
        self._mdlm_active = os.environ.get("HYDRA_USE_MDLM", "0") == "1"
        self._htm_cache_key = None  # populated by forward()

        # Token embedding
        self.wte = nn.Embedding(config.vocab_size, config.d_model)

        # Mamba-3 blocks — official mamba-ssm fused CUDA kernel. No fallbacks.
        # RoPE is applied internally by the Mamba3 CUDA kernel via the Angles
        # parameter; external cos/sin buffers are not needed.
        #
        # Hyena supplement: layers whose index appears in `config.hyena_layers`
        # are instantiated as HyenaBlock instead of Mamba3. The config field
        # is populated from HYDRA_HYENA_LAYERS at construction time and then
        # persisted to checkpoints, so resume is safe even when the env var
        # is unset. Empty tuple → all-Mamba3, byte-identical to pre-port.
        _hyena_layer_set = set(getattr(config, "hyena_layers", ()) or ())
        _gdn_layer_set = set(getattr(config, "gdn_layers", ()) or ())
        # Hyena wins on overlap; conflict is logged at construction time.
        _both = _hyena_layer_set & _gdn_layer_set
        if _both:
            print(f"[WARN] layers in both hyena_layers and gdn_layers; using Hyena: {sorted(_both)}", flush=True)
            _gdn_layer_set -= _hyena_layer_set
        if os.environ.get("HYDRA_STRICT_OPTIMAL_COMPONENTS", "0") == "1":
            if _hyena_layer_set or _gdn_layer_set:
                raise RuntimeError(
                    "HYDRA_STRICT_OPTIMAL_COMPONENTS=1 requires all layers to use Mamba3; "
                    f"got hyena_layers={sorted(_hyena_layer_set)} gdn_layers={sorted(_gdn_layer_set)}."
                )
            if Mamba3 is None:
                raise RuntimeError("HYDRA_STRICT_OPTIMAL_COMPONENTS=1 requires mamba_ssm/Mamba3 to be importable.")

        if _gdn_layer_set:
            from hydra.gdn_block import GDNBlock  # requires `fla` package

        def _build_block(i: int) -> nn.Module:
            if i in _hyena_layer_set:
                return HyenaBlock(
                    d_model=config.d_model,
                    seq_len=config.sequence_len,
                    order=int(os.environ.get("HYDRA_HYENA_ORDER", "2")),
                    filter_order=int(os.environ.get("HYDRA_HYENA_FILTER_DIM", "64")),
                )
            if i in _gdn_layer_set:
                return GDNBlock(
                    d_model=config.d_model,
                    n_heads=config.n_heads,
                )
            if Mamba3 is None:
                return _LocalMamba3Fallback()
            block = Mamba3(
                d_model=config.d_model,
                d_state=config.d_state,
                expand=config.expand,
                headdim=config.headdim,
                is_mimo=False,          # SISO path uses stable mamba3_siso_combined kernel
                chunk_size=int(os.environ.get("HYDRA_MAMBA3_CHUNK", "64")),  # 64 is the validated default; 128 tripped a Triton autotune hang (>8min, no progress)
                is_outproj_norm=False,
                dtype=torch.bfloat16,
            )
            return block

        self.blocks = nn.ModuleList([_build_block(i) for i in range(config.n_layer)])

        # Full-architecture SDR: offline semantic retina + STE (no-bypass).
        self.sdr_semantic = SemanticFoldingSDR(
            vocab_size=config.vocab_size,
            n_bits=config.sdr_n_bits,
            target_active=config.sdr_target_active,
            delta_rank=config.sdr_delta_rank,
            som_warmup_steps=config.sdr_som_warmup,
            som_update_interval=config.sdr_som_interval,
        )

        # HTM spatial pooler + temporal memory (Rust, Hebbian).
        self.htm = HTMLayer(
            input_bits=config.sdr_n_bits,
            n_columns=config.htm_n_columns,
            cells_per_column=config.htm_cells_per_column,
            batch_size=1,           # grows lazily to actual B on first forward
            seed=42,
            learn=True,
            reset_each_forward=True,
        )

        # Gradient bridge: (n_columns + anomaly) -> d_model.
        self.htm_proj = nn.Linear(config.htm_n_columns + 1, config.d_model, bias=False)

        # GPU Engram with Hebbian writes — runs EVERY step.
        self.engram = GPUEngram(
            d_model=config.d_model,
            n_columns=config.engram_n_columns,
            max_ngram=3,
        )
        self.reality_bridge = None
        self.cantor = None
        if os.environ.get("HYDRA_REALITY_BRIDGE", "0") == "1":
            d_reality = int(os.environ.get("HYDRA_REALITY_D", "133"))
            self.reality_bridge = RealityPoincareBridge(
                d_model=config.d_model,
                d_reality=d_reality,
                l0_k=int(os.environ.get("HYDRA_REALITY_L0_K", "64")),
            )
            if os.environ.get("HYDRA_CANTOR_DISABLE", "0") != "1":
                self.cantor = CantorRouter(
                    depth=int(os.environ.get("HYDRA_CANTOR_DEPTH", "7")),
                    d_query=d_reality,
                    seed=int(os.environ.get("HYDRA_CANTOR_SEED", "42")),
                    device=self.wte.weight.device,
                )
        self.engram_layer_idx = config.engram_layer_idx

        # Manifold-Constrained Hyper-Connections (one per Mamba-3 block).
        self.mhc = nn.ModuleList([
            ManifoldHyperConnection(d_model=config.d_model, n_streams=2, sinkhorn_iters=3)
            for _ in range(config.n_layer)
        ])

        # Hestia QAT — ternary weight quantization applied post-optimizer-step.
        self.hestia = HestiaQAT(enabled=True, bits=1.58)

        # LM head
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

        # Learnability knob 1: Multi-Token Prediction (Llama-3 style).
        #   MTP_K=1 -> standard next-token. MTP_K>1 -> extra heads predict
        #   tokens at positions t+1, t+2, ..., t+K. Heads are weight-tied to
        #   lm_head (we share Parameters), so the only extra compute is
        #   additional CE losses; no new params. Activated via HYDRA_MTP_K.
        self._mtp_k = max(1, int(os.environ.get("HYDRA_MTP_K", "1")))

        # Audit 2026-05-09 issue #22 - Cluster E. UnigramSampler attached
        # post-init by training setup. Stays None in environments that
        # cannot build the unigram cache (CI, unit tests); the legacy
        # uniform-negative path covers that case.
        self._unigram_sampler = None

        # Learnability knob 3: gradient checkpointing on Mamba3 blocks.
        self._grad_ckpt = os.environ.get("HYDRA_GRAD_CKPT", "0") == "1"

        # Learnability knob 4: doc-separator BOS masking in packed sequences.
        self._doc_sep_mask = os.environ.get("HYDRA_DOC_SEP_MASK", "0") == "1"
        # BOS token id is looked up lazily on first forward (requires tokenizer
        # load); -1 means uninitialized.
        self._bos_token_id = -1

        # Learnability knob 5: explicit stop-grad on HTM tensor (htm_rust
        # outputs already have requires_grad=False; this is defense-in-depth).
        self._htm_stop_grad = os.environ.get("HYDRA_HTM_STOP_GRAD", "0") == "1"

        # Learnability knob 6: entropy penalty coefficient on LM logits.
        self._entropy_penalty = float(os.environ.get("HYDRA_ENTROPY_PENALTY", "0.0"))

        # Residual dropout. Audit 2026-05-09 issue #14(a): default lowered from
        # 0.2 -> 0.05. The 0.2 default combined with 1/sqrt(2L) out_proj rescale
        # (audit issue #14b) and frozen dt_bias (#19) was producing zero
        # block_out_rms after ~200 steps (Mamba inertness, see
        # docs/sweeps/2026-04-20_inert_mamba_finding.md). 0.05 is GPT-2 style
        # and leaves room for the residual stream to actually carry block
        # output. Override via HYDRA_DROPOUT env var.
        self.drop = nn.Dropout(float(os.environ.get("HYDRA_DROPOUT", "0.05")))

        # Logits soft-capping. Default raised 15 -> 30 (audit #20): the lower
        # cap was actively flattening the high end of the logit distribution
        # at modern eff-vocab. Override via HYDRA_LOGIT_SOFTCAP env var.
        self.softcap = float(os.environ.get("HYDRA_LOGIT_SOFTCAP", "30.0"))

        # Secondary metrics storage
        self._metrics = {}

        # Per-layer diagnostic panel. Env-gated; zero overhead when off.
        # Emits residual-contribution (delta_ratio), feature std, effective rank,
        # gradient norm per layer; used to identify minimum viable n_layer + find
        # entropy leakage / dead layers. See docs/depth-sweep.md.
        #
        # Audit 2026-05-09 issue #14(c): the inertness tripwire (training.py)
        # needs `layer_*_block_out_rms` and `layer_*_block_in_rms` metrics
        # available at step <= TRIPWIRE_STEP. Force the forward hooks on
        # whenever the tripwire is enabled (default) so the check is not a
        # no-op on production runs that don't set HYDRA_LAYER_DIAGNOSTICS=1.
        # Backward (grad-norm) hooks remain gated by the original env var.
        self._diag_enabled = os.environ.get("HYDRA_LAYER_DIAGNOSTICS", "0") == "1"
        _tripwire_ratio = float(os.environ.get("HYDRA_INERTNESS_TRIPWIRE_RATIO", "0.01"))
        _tripwire_active = _tripwire_ratio > 0.0
        self._fwd_hooks_enabled = self._diag_enabled or _tripwire_active
        self._diag_step = 0
        self._diag_svd_every = int(os.environ.get("HYDRA_LAYER_DIAG_SVD_EVERY", "100"))
        if self._diag_enabled:
            # Gradient-norm backward hooks on each Mamba3 block output.
            for _i, _block in enumerate(self.blocks):
                def _mk_grad_hook(_layer_idx):
                    def _hook(module, grad_input, grad_output):
                        if grad_output and grad_output[0] is not None:
                            g = grad_output[0].detach()
                            self._metrics[f'layer_{_layer_idx}_grad_norm'] = float(
                                g.pow(2).mean().sqrt().item()
                            )
                    return _hook
                _block.register_full_backward_hook(_mk_grad_hook(_i))

        # Forward hooks on each Mamba3 block capture the block's OUTPUT
        # directly. This is the clean measurement: unlike merge_streams()
        # sampling which sees (streams + M*block_output) in bf16 — where
        # small block contributions round to zero against unit-norm
        # residuals — this captures `block_output` itself as produced.
        # Reports both its absolute RMS norm and its ratio to the block
        # INPUT's RMS norm (contribution magnitude relative to the
        # residual it's added to).
        # Auto-enabled when the inertness tripwire is active; otherwise
        # gated by HYDRA_LAYER_DIAGNOSTICS=1 (see _fwd_hooks_enabled above).
        if self._fwd_hooks_enabled:
            for _i, _block in enumerate(self.blocks):
                def _mk_fwd_hook(_layer_idx):
                    def _hook(module, inputs, output):
                        with torch.no_grad():
                            inp = inputs[0].detach().float() if inputs else None
                            out = output.detach().float() if isinstance(output, torch.Tensor) else None
                            if out is not None:
                                out_rms = out.pow(2).mean().sqrt().item()
                                self._metrics[f'layer_{_layer_idx}_block_out_rms'] = float(out_rms)
                                if inp is not None:
                                    in_rms = inp.pow(2).mean().sqrt().item()
                                    self._metrics[f'layer_{_layer_idx}_block_in_rms'] = float(in_rms)
                                    self._metrics[f'layer_{_layer_idx}_contrib_ratio'] = float(
                                        out_rms / (in_rms + 1e-8)
                                    )
                    return _hook
                _block.register_forward_hook(_mk_fwd_hook(_i))

        # Triton kernel integration gates (Phase 2 — deferred, see module docstring).
        self._fused_bcnorm = os.environ.get("HYDRA_FUSED_BCNORM", "0") == "1"
        self._fused_ssd = os.environ.get("HYDRA_FUSED_SSD", "0") == "1"
        if self._fused_bcnorm or self._fused_ssd:
            import sys
            _active = []
            if self._fused_bcnorm:
                _active.append("HYDRA_FUSED_BCNORM")
            if self._fused_ssd:
                _active.append("HYDRA_FUSED_SSD")
            print(
                f"[HYDRA] Triton kernel gates set: {', '.join(_active)}. "
                f"NOTE: Both are DEFERRED (mamba-ssm Mamba3 uses internal "
                f"CUDA kernels). Gates accepted but currently no-ops.",
                file=sys.stderr,
            )

        # R6 optional torch.compile on the impl forward. Gated (default OFF).
        if os.environ.get("HYDRA_MODEL_COMPILE", "0") == "1":
            self._forward_impl = torch.compile(
                self._forward_impl,
                fullgraph=False,
                dynamic=True,
                mode="default",
            )

    @torch.no_grad()
    def init_weights(self) -> None:
        s = 3 ** 0.5 * self.config.d_model ** -0.5

        # Move SDR retina indices (plain attribute, not buffer) to same device as params.
        # Required because to_empty() only moves params/buffers, and _retina_indices
        # is loaded from numpy (always CPU) by SemanticFoldingSDR.__init__.
        device = self.wte.weight.device
        if hasattr(self.sdr_semantic, '_retina_indices'):
            self.sdr_semantic._retina_indices = self.sdr_semantic._retina_indices.to(device)

        # Embedding init: GPT-2 / LLaMA convention. std=1.0 was chosen for
        # vocab=8192; at larger vocabs, smaller std prevents logit blowup.
        # Use std = 1/sqrt(d_model) which scales sensibly with model width.
        import math as _math
        _d_model = self.wte.weight.shape[1]
        wte_std = float(os.environ.get("HYDRA_WTE_STD", str(1.0 / _math.sqrt(_d_model))))
        nn.init.normal_(self.wte.weight, mean=0.0, std=wte_std)
        # LM head init: was std=0.001 — PATHOLOGICAL at vocab>=32k because
        # logits collapse to zero, loss locks at log(V)~=11, gradient through
        # head ∝ 1/V is too small to escape. GPT-2 uses std=0.02; LLaMA uses
        # std=1/sqrt(d_model). Pick 0.02 as robust default, env-overridable.
        lm_head_std = float(os.environ.get("HYDRA_LM_HEAD_STD", "0.02"))
        nn.init.normal_(self.lm_head.weight, mean=0.0, std=lm_head_std)
        # F8 (NOT APPLIED): Weight tying would save V*D params but current LR
        # groups have embedding_lr=1.0 and unembedding_lr=0.005 × d_model_scale
        # — tying forces the shared tensor under a single LR group and either
        # the embeddings learn 200x too slow (under unembed LR) or the LM head
        # becomes unstable (under embed LR). Short 15-step smoke with tying +
        # embed-group update showed initial loss jump 9 -> 20. Deferred until
        # LR groups are re-tuned; see docs/OPTIMIZATION_PLAN.md Post-plan.

        for li, block in enumerate(self.blocks):
            if hasattr(block, 'in_proj') and hasattr(block.in_proj, 'weight'):
                nn.init.uniform_(block.in_proj.weight, -s, s)
            if hasattr(block, 'out_proj') and hasattr(block.out_proj, 'weight'):
                # Audit 2026-05-09 issue #14(b): out_proj init std is now FLAT
                # 0.02 (GPT-2 default for non-residual heads). The previous
                # default of 0.02 / sqrt(2 * n_layer) — designed for L>=24
                # GPT-2-class stacks — was throwing away gradient signal at
                # the L=1-4 configs feather actually trains. Combined with
                # dropout=0.2 and frozen dt_bias this drove block_out_rms to
                # zero after ~200 steps (Mamba inertness).
                #
                # Default behaviour: flat 0.02. Legacy 1/sqrt(2L) divisor is
                # available via HYDRA_OUT_PROJ_DIVISOR=1 for reproducibility
                # of pre-2026-05-09 runs. HYDRA_OUT_PROJ_STD overrides the
                # absolute std value (applied AFTER the optional divisor).
                n_layer = self.config.n_layer
                _legacy_divisor = os.environ.get("HYDRA_OUT_PROJ_DIVISOR", "0") == "1"
                _base_std = 0.02 / (2 * n_layer) ** 0.5 if _legacy_divisor else 0.02
                out_std = float(os.environ.get("HYDRA_OUT_PROJ_STD", str(_base_std)))
                nn.init.normal_(block.out_proj.weight, mean=0.0, std=out_std)

            # Audit 2026-05-09 issue #19: explicit per-head dt_bias init.
            # Mamba-2 paper convention: softplus(dt_bias) ~ U(dt_min, dt_max)
            # so each head learns a different discretisation timescale. Mamba3
            # CUDA kernel exposes `dt_bias` of shape (n_heads,). The default
            # init in mamba_ssm already follows this distribution, but we
            # re-apply it explicitly here both for documentation and for
            # env-tunable bounds (HYDRA_DT_MIN, HYDRA_DT_MAX).
            if hasattr(block, 'dt_bias') and isinstance(block.dt_bias, nn.Parameter):
                _dt_min = float(os.environ.get("HYDRA_DT_MIN", "0.001"))
                _dt_max = float(os.environ.get("HYDRA_DT_MAX", "0.1"))
                _log_dt_min = _math.log(max(_dt_min, 1e-6))
                _log_dt_max = _math.log(max(_dt_max, _dt_min * 1.0001))
                # softplus^{-1}(t) = log(exp(t) - 1). For small t (< ~0.5),
                # softplus^{-1}(t) ~= log(t), so we initialise log_dt and let
                # the kernel's softplus produce dt ~ exp(log_dt) ~ U(dt_min, dt_max).
                with torch.no_grad():
                    block.dt_bias.uniform_(_log_dt_min, _log_dt_max)

        nn.init.normal_(self.htm_proj.weight, mean=0.0, std=s)

        if hasattr(self.engram, "memory"):
            nn.init.normal_(self.engram.memory, mean=0.0, std=0.01)
        if hasattr(self.engram, "gate"):
            nn.init.zeros_(self.engram.gate.weight)
            nn.init.zeros_(self.engram.gate.bias)
        if self.reality_bridge is not None:
            nn.init.normal_(self.reality_bridge.to_reality.weight, mean=0.0, std=0.02)
            nn.init.normal_(self.reality_bridge.to_tangent2.weight, mean=0.0, std=0.02)
        if self.cantor is not None and hasattr(self.cantor, "branch"):
            bound = (3.0 / float(self.cantor.d_query)) ** 0.5
            nn.init.uniform_(self.cantor.branch, -bound, bound)

        # Cast to bf16 to match Mamba3 dtype; Muon groups by shape so mixed
        # dtypes in the same shape group would break lerp_ dtype checks.
        self.wte.to(dtype=torch.bfloat16)
        self.blocks.to(dtype=torch.bfloat16)
        self.htm_proj.to(dtype=torch.bfloat16)
        self.engram.to(dtype=torch.bfloat16)
        if self.reality_bridge is not None:
            self.reality_bridge.to(dtype=torch.bfloat16)
        if self.cantor is not None:
            self.cantor.to(dtype=torch.bfloat16)

    def set_bos_token_id(self, bos_id: int) -> None:
        """Inform the model of the tokenizer's BOS id so doc-separator
        masking (learnability #4) knows which positions to skip. Called from
        training setup once the tokenizer is loaded."""
        self._bos_token_id = int(bos_id)

    def set_unigram_sampler(self, sampler) -> None:
        """Attach a UnigramSampler used by the sampled-softmax LM-head loss.

        Audit 2026-05-09 issue #22 - Cluster E. With a sampler attached the
        LM-head loss draws negatives from the empirical unigram distribution
        and applies the per-id ``log p_unigram`` correction (Jean et al.
        2015, importance-sampled NCE). Without a sampler the legacy
        uniform-negative path runs (uniform proposal + ``log(V/K)`` constant
        correction) - kept as a fallback for environments that cannot build
        the unigram cache (no tokenizer, CI, etc.).

        Eval always uses full softmax regardless of this setting.
        """
        if sampler is not None and not isinstance(sampler, UnigramSampler):
            raise TypeError(
                f"set_unigram_sampler expects UnigramSampler or None, got {type(sampler)}"
            )
        self._unigram_sampler = sampler

    def invalidate_hyena_caches(self) -> None:
        """Invalidate filter-rfft caches on all Hyena blocks.

        MUST be called after each `optimizer.step()` when
        `HYDRA_HYENA_FILTER_CACHE=1` is set, otherwise cached rfft values
        will be reused with stale filter parameters.

        No-op for blocks that are not HyenaBlock (Mamba3, etc.).
        """
        for block in self.blocks:
            if hasattr(block, "operator") and hasattr(block.operator, "invalidate_filter_cache"):
                block.operator.invalidate_filter_cache()

    def flush_hyena_pending_grads(self) -> None:
        """Push pending train-cache filter gradients into filter params.

        Used ONLY when HYDRA_HYENA_TRAIN_CACHE=1. Must be called exactly once
        per optimizer step, BEFORE `optimizer.step()` and BEFORE
        `invalidate_hyena_caches()`. The lightning_module wires this in
        `optimizer_step` around the existing optimizer.step() call.

        No-op if:
          * No HyenaBlocks are in the model, OR
          * No micro-batch ever ran with grad enabled (e.g. all-eval step).
        """
        for block in self.blocks:
            if hasattr(block, "operator") and hasattr(block.operator, "flush_pending_filter_grads"):
                block.operator.flush_pending_filter_grads()

    def estimate_flops(self) -> int:
        nparams = sum(p.numel() for p in self.parameters())
        embed_params = self.wte.weight.numel()
        return 6 * (nparams - embed_params)

    def num_scaling_params(self) -> dict:
        wte = sum(p.numel() for p in self.wte.parameters())
        lm_head = sum(p.numel() for p in self.lm_head.parameters())
        blocks = sum(p.numel() for p in self.blocks.parameters())
        sdr = sum(p.numel() for p in self.sdr_semantic.parameters())
        htm_proj = sum(p.numel() for p in self.htm_proj.parameters())
        engram = sum(p.numel() for p in self.engram.parameters())
        total = sum(p.numel() for p in self.parameters())
        return {
            'wte': wte, 'lm_head': lm_head, 'blocks': blocks,
            'sdr_semantic': sdr, 'htm_proj': htm_proj,
            'engram': engram, 'total': total,
        }

    def get_secondary_metrics(self) -> dict:
        """Flush any lingering CUDA tensors to host (single sync)."""
        flushed = {}
        for k, v in self._metrics.items():
            if hasattr(v, 'item'):
                try:
                    flushed[k] = float(v.item())
                except Exception:
                    flushed[k] = v
            else:
                flushed[k] = v
        return flushed

    def setup_optimizer(self, unembedding_lr=0.004, embedding_lr=0.6, matrix_lr=0.04,
                        weight_decay=0.2, adam_betas=(0.8, 0.95), scalar_lr=0.5):
        """Setup MuonAdamW optimizer with per-component LR groups."""
        model_dim = self.config.d_model

        embedding_params = list(self.wte.parameters())
        lm_head_params = list(self.lm_head.parameters())

        # Muon routing guard: 2D parameters are NOT automatically matrices.
        # Exclude:
        #   (a) params whose name ends in `.freq` — Sin frequency vectors used
        #       by Hyena's implicit filter MLP. Shape (1, dim) is nominally 2D
        #       but semantically a per-dim scalar. Muon's polar-express
        #       orthogonalization would force it toward an orthogonal matrix,
        #       destroying the learned modulation frequencies.
        #   (b) 2-D params with min(shape) < MUON_MIN_DIM. Tiny projections
        #       (e.g. HyenaFilter.implicit_filter.0.weight of shape (64, 3))
        #       get collapsed toward near-identity by orthogonalization on the
        #       narrow axis, damaging expressivity. These belong in AdamW.
        # These exclusions route the params into the AdamW scalar/vector group.
        MUON_MIN_DIM = 8

        def _muon_eligible(name: str, p: torch.Tensor) -> bool:
            if p.dim() != 2:
                return False
            if name.endswith(".freq"):
                return False
            if min(p.shape) < MUON_MIN_DIM:
                return False
            return True

        # Matrix params -> Muon (2D weight matrices passing the routing guard).
        matrix_params = []
        for name, p in self.blocks.named_parameters():
            if _muon_eligible(name, p):
                matrix_params.append(p)
        # NOTE (W1 audit REG-2): SemanticFoldingSDR.delta_u / delta_v are
        # currently GRADIENT-DEAD. The forward path uses `binary_only(idx)` for
        # HTM and stores it as `self._last_sdr`, but does NOT route the STE
        # output through any downstream op. Including them in the Muon group
        # burns compute (stack + orthogonalize + lerp) on zero-grad params
        # every step. Excluded here; a later W5 pass can reconnect STE via a
        # gated residual if the SDR signal is wanted back in-graph. The
        # parameters still exist, so no state_dict break.
        # for p in self.sdr_semantic.parameters():
        #     if p.dim() == 2:
        #         matrix_params.append(p)
        for name, p in self.htm_proj.named_parameters():
            if _muon_eligible(name, p):
                matrix_params.append(p)
        for name, p in self.engram.named_parameters():
            if _muon_eligible(name, p):
                matrix_params.append(p)

        # ------------------------------------------------------------------
        # Contrastive retina (Retina-D): retina_contrastive gets its OWN AdamW
        # group with a dedicated LR so it can learn word2vec-like similarity
        # structure without interfering with the main model LR schedule.
        #
        # Audit 2026-05-09 issue #21: pre-2026-05-09 this group bound to
        # `retina_logits`, which was the contrastive embedding (the shadowing
        # parameter — see subsystems/sdr_semantic.py). After the rename, the
        # group binds to `retina_contrastive`. The learnable-mode binary
        # retina_logits (shape [V, n_bits]) is now bound to the SDR delta
        # group at HYDRA_SDR_LR (1e-3 default) so it actually learns instead
        # of sitting at lr=0.
        #
        # delta_u / delta_v are still gradient-dead (REG-2), but the audit
        # now routes them to a real lr group so subsequent reconnection of
        # the STE through HTM (deferred work) trains them automatically.
        # retina_contrastive is in-graph via contrastive_loss() every N steps.
        #
        # LR (contrastive): 1e-3 — higher than embedding (0.6 * scale) to
        # compensate for the sparse gradient signal (only pairs from current
        # batch update each row). weight_decay=0.01 provides light row-norm
        # regularisation without collapsing logits toward zero.
        # LR (sdr delta + binary retina_logits): HYDRA_SDR_LR=1e-3.
        # ------------------------------------------------------------------
        _contrastive_enabled = os.environ.get("HYDRA_CONTRASTIVE_RETINA", "0") == "1"
        _contrastive_lr = float(os.environ.get("HYDRA_CONTRASTIVE_LR", "1e-3"))
        _sdr_lr = float(os.environ.get("HYDRA_SDR_LR", "1e-3"))

        # Separate delta params (gradient-dead) and binary retina_logits
        # (learnable mode only) from retina_contrastive (in-graph when enabled).
        sdr_delta_ids = {id(self.sdr_semantic.delta_u), id(self.sdr_semantic.delta_v)}
        sdr_contrastive_params = []
        if self.sdr_semantic.retina_contrastive is not None:
            sdr_contrastive_params.append(self.sdr_semantic.retina_contrastive)
        
        sdr_delta_params = [self.sdr_semantic.delta_u, self.sdr_semantic.delta_v]
        # Binary retina_logits exists only in learnable mode. When present,
        # route it through the SDR delta group with HYDRA_SDR_LR so it can
        # learn alongside delta_u/delta_v rather than sit at lr=0.
        sdr_binary_logits = []
        _binary_logits_param = getattr(self.sdr_semantic, 'retina_logits', None)
        if isinstance(_binary_logits_param, torch.nn.Parameter):
            sdr_binary_logits.append(_binary_logits_param)
            sdr_delta_params.append(_binary_logits_param)

        # All SDR param ids — excluded from scalar group regardless.
        sdr_param_ids = {id(p) for p in self.sdr_semantic.parameters()}

        # Audit 2026-05-09 issue #19: dt_bias gets its own optimizer group.
        # Previously rolled into the scalar AdamW group at lr=scalar_lr*scale
        # (~0.5 * dmodel_lr_scale), which is far too low for the per-head
        # discretization-step bias to escape its initialisation. Mamba-2 paper
        # uses lr=1e-2 for dt_proj/dt_bias; replicate that here. wd=0 because
        # this is a per-head bias, not a regulariseable matrix.
        _dt_bias_lr = float(os.environ.get("HYDRA_DT_BIAS_LR", "1e-2"))
        dt_bias_params = []
        for name, p in self.blocks.named_parameters():
            if name.endswith("dt_bias") or name == "dt_bias":
                dt_bias_params.append(p)
        dt_bias_ids = {id(p) for p in dt_bias_params}

        assigned = set(id(p) for p in embedding_params + lm_head_params + matrix_params)
        scalar_params = [
            p for p in self.parameters()
            if id(p) not in assigned
            and id(p) not in sdr_param_ids
            and id(p) not in dt_bias_ids
        ]

        total_assigned = (
            len(embedding_params) + len(lm_head_params) + len(matrix_params)
            + len(scalar_params) + len(sdr_delta_params) + len(sdr_contrastive_params)
            + len(dt_bias_params)
        )
        total_params = len(list(self.parameters()))
        assert total_assigned == total_params, (
            f"Parameter count mismatch: assigned {total_assigned} vs total {total_params}"
        )

        dmodel_lr_scale = (model_dim / 768) ** -0.5
        print(f"Scaling AdamW LRs by 1/sqrt({model_dim}/768) = {dmodel_lr_scale:.6f}")
        if _contrastive_enabled:
            print(f"[retina-d] Contrastive retina ENABLED. retina_contrastive LR={_contrastive_lr:.2e}")
        else:
            print("[retina-d] Contrastive retina DISABLED (set HYDRA_CONTRASTIVE_RETINA=1 to enable)")
        print(
            f"[sdr] delta_u/delta_v + binary retina_logits routed to SDR group "
            f"(lr={_sdr_lr:.2e}) — audit issue #21."
        )
        if dt_bias_params:
            print(
                f"[dt-bias] {len(dt_bias_params)} dt_bias parameter(s) routed to "
                f"dedicated AdamW group (lr={_dt_bias_lr:.2e}, betas=(0.9, 0.95)) — "
                f"audit issue #19."
            )
        else:
            print("[dt-bias] no dt_bias parameters detected (Hyena/GDN-only stack?)")

        param_groups = [
            dict(kind='adamw', params=lm_head_params,
                 lr=unembedding_lr * dmodel_lr_scale, betas=adam_betas,
                 eps=1e-10, weight_decay=0.0),
            dict(kind='adamw', params=embedding_params,
                 lr=embedding_lr * dmodel_lr_scale, betas=adam_betas,
                 eps=1e-10, weight_decay=0.0),
        ]
        # Contrastive retina embedding: dedicated group, always present
        # if the parameter exists, so optimizer state is consistent.
        if sdr_contrastive_params:
            # kind='retina_contrastive' marks this group so
            # training.py skips cosine-LR scaling.
            param_groups.append(
                dict(kind='retina_contrastive', params=sdr_contrastive_params,
                     lr=_contrastive_lr, betas=adam_betas,
                     eps=1e-10, weight_decay=0.01)
            )

        # Audit 2026-05-09 issue #21: SDR delta_u/delta_v (and the
        # learnable-mode binary retina_logits when present) get a real
        # lr=HYDRA_SDR_LR (1e-3 default). Pre-2026-05-09 this group sat
        # at lr=0 — even when STE was reconnected through HTM the
        # parameters could not learn. retina_logits is gradient-active
        # only in learnable mode where contrastive_loss / binary_only
        # use it; in offline mode binary_only goes through _retina_data
        # and these params don't accumulate updates.
        param_groups.append(
            dict(kind='adamw', params=sdr_delta_params,
                 lr=_sdr_lr, betas=adam_betas,
                 eps=1e-10, weight_decay=0.0)
        )

        # Audit 2026-05-09 issue #19: dedicated AdamW group for dt_bias. Held
        # out of the cosine LR scaling applied to model groups by setting
        # kind='dt_bias' (training.py treats this kind like 'retina_contrastive'
        # and pins lr to its initial value). betas=(0.9, 0.95) follows the
        # Mamba-2 paper convention for dt-related parameters.
        if dt_bias_params:
            param_groups.append(dict(
                kind='dt_bias', params=dt_bias_params,
                lr=_dt_bias_lr, betas=(0.9, 0.95),
                eps=1e-8, weight_decay=0.0,
            ))

        if scalar_params:
            param_groups.append(
                dict(kind='adamw', params=scalar_params,
                     lr=scalar_lr * dmodel_lr_scale, betas=adam_betas,
                     eps=1e-10, weight_decay=0.0)
            )

        for shape in sorted({p.shape for p in matrix_params}):
            group_params = [p for p in matrix_params if p.shape == shape]
            # ns_steps: Muon polar-express inner iterations. Default 5 (paper),
            # but 3 converges on small matrices (d_model ~ 384) with ~40% lower
            # optimizer step cost. Env-tunable for experimentation.
            _ns_steps = int(os.environ.get("HYDRA_MUON_NS_STEPS", "3"))
            param_groups.append(dict(
                kind='muon', params=group_params, lr=matrix_lr,
                momentum=0.95, ns_steps=_ns_steps, beta2=0.95, weight_decay=weight_decay,
            ))

        optimizer = MuonAdamW(param_groups)
        for group in optimizer.param_groups:
            group["initial_lr"] = group["lr"]
        return optimizer

    def forward(self, idx, targets=None, reduction='mean'):
        """idx: (B, T) int64. Returns loss if targets given, else logits.

        Nested bf16 autocast is a no-op when ambient autocast is already on;
        when it's off (e.g. integration tests) we establish the dtype contract.
        """
        if torch.is_autocast_enabled():
            return self._forward_impl(idx, targets=targets, reduction=reduction)
        with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
            return self._forward_impl(idx, targets=targets, reduction=reduction)

    def _forward_impl(self, idx, targets=None, reduction='mean'):
        B, T = idx.shape

        # Diagnostic: per-subsystem CUDA event timing. Env-gated; zero overhead
        # when disabled. Logs one timing line per forward call. Used to isolate
        # which subsystem is the tps bottleneck on paid hardware.
        _profile = os.environ.get("HYDRA_PROFILE_FORWARD", "0") == "1"
        if _profile:
            def _ev():
                e = torch.cuda.Event(enable_timing=True)
                e.record()
                return e
            _t0 = _ev()
        else:
            _t0 = None

        # Compute SDR binary ONCE and reuse for both HTM input and the stash.
        sdr_binary = self.sdr_semantic.binary_only(idx)
        self._last_sdr = sdr_binary  # uint8 stash (not bf16 → 256MB avoidance)

        # HTM subsampling: run HTM on 1 of every N micro-batches within a
        # gradient accumulation step, reuse the cached result for the other
        # N-1 micro-batches. Cooperative launch monopolizes all SMs (grid.sync
        # requires full-grid residency), so HTM and mamba can't overlap via
        # streams. Subsampling removes HTM from most micro-batches' critical
        # path instead.
        #
        # Math: N=8, 64 accum steps → 8 HTM calls (10.6ms each) + 56 fast
        # calls (4ms each). Total = 84.8 + 224 = 309ms → 106k tps.
        #
        # HYDRA_HTM_SUBSAMPLE=N (default 8). Set =1 for every-microbatch HTM.
        _htm_sub = int(os.environ.get("HYDRA_HTM_SUBSAMPLE", "8"))
        if not hasattr(self, '_htm_call_idx'):
            self._htm_call_idx = int(os.environ.get("HYDRA_HTM_INITIAL_OFFSET", "0"))

        _run_htm = (self._htm_call_idx % _htm_sub == 0)
        self._htm_call_idx += 1

        if _run_htm:
            htm_handle = self.htm.forward_async(sdr_binary, output_dtype=self.wte.weight.dtype)
        else:
            htm_handle = None

        if _profile: _t_htm_async = _ev()

        dense_emb = self.wte(idx)  # (B, T, d_model) bf16
        dense_emb = semantic_gaussian_mollify(
            dense_emb,
            std=float(os.environ.get("HYDRA_SEMANTIC_SMOOTH_STD", "0.0")),
            training=self.training,
            eval_enabled=os.environ.get("HYDRA_SEMANTIC_SMOOTH_EVAL", "0") == "1",
        )

        if _profile: _t_wte = _ev()

        # Audit 2026-05-09 issue #18: cache reuse must verify the current
        # SDR active-index pattern matches the one the cache was built from.
        # Under MDLM, mask resampling produces forwards with identical shape
        # but different content; shape-only reuse fed stale HTM features.
        # nonzero() is computed lazily — only when we are about to write the
        # cache or test reuse — so the runtime hit on the no-cache path is
        # zero.
        if _run_htm:
            htm_out = self.htm.forward_await(htm_handle)
            self._htm_cache = htm_out.detach()  # cache for non-HTM micro-batches
            self._htm_cache_key = htm_cache_key(sdr_binary.nonzero())
            self._htm_cache_shape = (B, T)
        elif (
            # Throughput mode for ordinary CE training: real GPU HTM still runs
            # every HYDRA_HTM_SUBSAMPLE forwards and learns there. Between real
            # HTM passes, reuse the previous same-shape HTM features instead of
            # exact SDR-key matching, because streaming Nemotron changes SDR
            # content every batch and exact matching degenerates to HTM every
            # microbatch. Keep exact matching for MDLM/masked forwards.
            self.training
            and not self._mdlm_active
            and os.environ.get("HYDRA_HTM_CACHE_MODE", "exact").lower() == "shape"
            and hasattr(self, '_htm_cache') and self._htm_cache is not None
            and getattr(self, '_htm_cache_shape', None) == (B, T)
        ):
            htm_out = self._htm_cache
        elif (
            hasattr(self, '_htm_cache') and self._htm_cache is not None
            and self._htm_cache.shape[0] == B and self._htm_cache.shape[1] == T
            and not self._mdlm_active
            and htm_cache_matches(self._htm_cache_key, sdr_binary.nonzero())
        ):
            htm_out = self._htm_cache
        elif (
            os.environ.get("HYDRA_HTM_ZERO_CACHE_ON_MISS", "0") == "1"
            and self.training
            and not self._mdlm_active
        ):
            htm_out = torch.zeros((B, T, self.config.htm_n_columns + 1), device=dense_emb.device, dtype=dense_emb.dtype)
            self._htm_cache = htm_out.detach()
            self._htm_cache_key = None
            self._htm_cache_shape = (B, T)
        else:
            # Very first call with subsample > 1, OR MDLM is on, OR the SDR
            # pattern has changed from the cached one under exact mode: run HTM.
            htm_handle = self.htm.forward_async(sdr_binary, output_dtype=self.wte.weight.dtype)
            htm_out = self.htm.forward_await(htm_handle)
            self._htm_cache = htm_out.detach()
            self._htm_cache_key = htm_cache_key(sdr_binary.nonzero())
            self._htm_cache_shape = (B, T)

        if _profile: _t_htm_await = _ev()
        with torch.no_grad():
            sdr_active_bits = float(self.sdr_semantic.target_active)
            htm_anomaly = htm_out[..., -1].mean()

        # Learnability #5: explicit stop-grad on HTM output. htm_rust already
        # produces a detached tensor, but making it explicit here hardens the
        # contract against future refactors that might route HTM through a
        # grad-enabled op.
        if self._htm_stop_grad:
            htm_out = htm_out.detach()

        # Gradient bridge: HTM columns+anomaly -> d_model.
        htm_proj_out = self.htm_proj(htm_out.to(dense_emb.dtype))
        x = dense_emb + htm_proj_out
        x = norm(x)

        if _profile: _t_htm_proj = _ev()

        # mHC-routed Mamba-3 stack with Engram injection at configured layer.
        streams = self.mhc[0].init_streams(x)
        _engram_ev = None

        # Per-layer diagnostic panel. The pre-layer merged state h_pre lets us
        # measure residual contribution of each layer: delta_N = h_post - h_pre.
        # All reads are detached no-grad to avoid autograd graph pollution.
        _diag = self._diag_enabled
        if _diag:
            # Cast to float32 for the diagnostic arithmetic: the layer's
            # residual contribution is small (~0.5 × rms-normed block output),
            # which underflows in bf16 subtraction (3-digit mantissa) and
            # reports delta_ratio=0 at the boundaries. float32 snapshot is
            # ~3.8 MB extra memory per diag sample (B=1, T=2048, d=96) —
            # negligible vs peak VRAM.
            with torch.no_grad():
                h_pre = self.mhc[0].merge_streams(streams).detach().float()
            _run_svd = (self._diag_step % self._diag_svd_every) == 0

        for i, (block, mhc_layer) in enumerate(zip(self.blocks, self.mhc)):
            def _block_fn(h, _block=block):
                return self.drop(_block(norm(h)))

            # Learnability #3: gradient checkpointing. Wrap the block-fn so
            # the mhc layer's internal uses of it re-run the block in backward
            # (trading compute for activation memory). use_reentrant=False is
            # the modern API and works cleanly under autocast.
            if self._grad_ckpt and self.training:
                import torch.utils.checkpoint as _ckpt
                _raw_fn = _block_fn
                def _block_fn(h, _raw=_raw_fn):  # noqa: E731
                    return _ckpt.checkpoint(_raw, h, use_reentrant=False)

            streams = mhc_layer(streams, _block_fn)

            if i == self.engram_layer_idx:
                if _profile: _t_pre_engram = _ev()
                # mHC stream-2 preservation at the engram boundary.
                # Old behavior (HYDRA_ENGRAM_RESET_STREAMS=1) re-initialized
                # both streams from the merged tensor, discarding stream-1's
                # accumulated state from layers 0..engram_layer_idx and
                # collapsing mHC topology to n_streams=1 below the engram.
                # New default: treat the engram as a residual update on
                # stream-0 only, so stream-1 carries its prior state into
                # the post-engram half of the network. `streams` here is a
                # tensor of shape (n_streams, B, T, d_model) — see
                # subsystems/mhc_mini.ManifoldHyperConnection.
                x_mid = mhc_layer.merge_streams(streams)
                if self.reality_bridge is not None and self.cantor is not None:
                    rb = self.reality_bridge(x_mid)
                    cantor_leaf_ids, _ = self.cantor(rb.reality, return_scores=False)
                    x_after_engram, hit_rate = self.engram(
                        x_mid,
                        idx,
                        sdr_active_indices=rb.l0_indices,
                        cantor_leaf_ids=cantor_leaf_ids,
                        cantor_n_leaves=self.cantor.n_leaves,
                    )
                else:
                    x_after_engram, hit_rate = self.engram(x_mid, idx)
                if os.environ.get("HYDRA_ENGRAM_RESET_STREAMS", "0") == "1":
                    streams = mhc_layer.init_streams(x_after_engram)
                else:
                    engram_delta = x_after_engram - x_mid
                    streams = streams.clone()
                    streams[0] = streams[0] + engram_delta
                self._metrics['engram_hit_rate'] = hit_rate
                if _diag:
                    with torch.no_grad():
                        s0 = streams[0].detach().float()
                        s1 = streams[1].detach().float() if streams.shape[0] > 1 else None
                        self._metrics['engram_stream0_rms'] = float(s0.pow(2).mean().sqrt().item())
                        if s1 is not None:
                            self._metrics['engram_stream1_rms'] = float(s1.pow(2).mean().sqrt().item())
                            self._metrics['engram_stream_divergence_rms'] = float(
                                (s0 - s1).pow(2).mean().sqrt().item()
                            )
                if _profile: _engram_ev = _ev()

            if _diag:
                with torch.no_grad():
                    h_post = mhc_layer.merge_streams(streams).detach().float()
                    in_n  = h_pre.pow(2).mean().sqrt()
                    out_n = h_post.pow(2).mean().sqrt()
                    d_n   = (h_post - h_pre).pow(2).mean().sqrt()
                    self._metrics[f'layer_{i}_in_norm']     = float(in_n.item())
                    self._metrics[f'layer_{i}_out_norm']    = float(out_n.item())
                    self._metrics[f'layer_{i}_delta_ratio'] = float((d_n / (in_n + 1e-6)).item())
                    self._metrics[f'layer_{i}_feat_std']    = float(h_post.std(dim=-1).mean().item())
                    if _run_svd:
                        # Effective rank via participation ratio of singular values.
                        # eff_rank = (Σσ)^2 / Σσ² — smooth rank proxy, bounded by d_model.
                        # Sampled to keep overhead low (SVD is O(min(B*T, D)^2·D)).
                        flat = h_post.reshape(-1, h_post.shape[-1])[:512].float()
                        try:
                            s = torch.linalg.svdvals(flat)
                            eff_rank = float(((s.sum() ** 2) / (s.pow(2).sum() + 1e-6)).item())
                            self._metrics[f'layer_{i}_eff_rank'] = eff_rank
                        except Exception:
                            pass
                    h_pre = h_post

        if _diag:
            self._diag_step += 1

        if _profile: _t_blocks = _ev()

        self._metrics['sdr_active_bits'] = sdr_active_bits
        self._metrics['htm_anomaly'] = htm_anomaly

        x = self.mhc[-1].merge_streams(streams)
        x = norm(x)

        if _profile: _t_merge = _ev()

        softcap = self.softcap
        _softcap_clamp = os.environ.get("HYDRA_SOFTCAP_CLAMP", "0") == "1"
        if targets is not None:
            smoothing = self.config.label_smoothing
            V = self.config.vocab_size

            # Learnability #4: doc-separator masking. In packed rows,
            # tokenizer.encode(..., prepend=bos_token) places a BOS at every
            # document boundary. Without masking, the model is penalized for
            # failing to predict "doc B's BOS" from the last tokens of doc A
            # — pure noise. We set targets==bos to -1 (ignore_index). Done
            # BEFORE MTP/entropy/sampled-softmax branches so all downstream
            # losses inherit the mask.
            if self._doc_sep_mask and self._bos_token_id >= 0:
                targets = torch.where(
                    targets == self._bos_token_id,
                    torch.full_like(targets, -1),
                    targets,
                )

            # Sampled softmax: instead of computing logits for ALL V tokens,
            # compute only for the target + K negatives drawn from a proposal
            # distribution. Reduces the lm_head matmul from (B*T, d) x (d, V)
            # to per-row (1, d) x (d, K) matmuls fused as einsum. At V=65536
            # and K=4096: 16x less compute, ~1.3-1.8x tps improvement once
            # the LM-head matmul is the bottleneck.
            #
            # Audit 2026-05-09 issue #22 - Cluster E. Two paths:
            #   * UnigramSampler attached -> per-row independent negatives
            #     drawn from the unigram distribution; per-id `log p_unigram`
            #     correction (Jean et al. 2015 importance-sampled NCE). This
            #     is the new default once training.py builds the cache.
            #   * No sampler -> legacy uniform-negative path, single shared
            #     batch of K negatives, constant `log(V/K)` correction.
            #     Retained as fallback when the unigram cache is unavailable
            #     (CI, unit tests, missing tokenizer).
            # Eval (self.training=False) always uses full softmax.
            #
            # Set HYDRA_SAMPLED_SOFTMAX=0 to force full softmax even in train.
            K_neg = int(os.environ.get("HYDRA_SAMPLED_SOFTMAX", "4096"))
            use_sampled = K_neg > 0 and K_neg < V and self.training
            unigram_sampler = getattr(self, "_unigram_sampler", None)

            # neg_logits / log_correction are populated by the legacy uniform
            # path so MTP heads can reuse them. The unigram path leaves them
            # at None and each MTP head draws its own per-row negatives.
            neg_logits = None
            log_correction = None

            if use_sampled and unigram_sampler is not None:
                # ----- Unigram-sampler path (importance-weighted NCE) -----
                h_flat = x.reshape(-1, x.shape[-1])
                t_flat = targets.reshape(-1)
                valid_mask_flat = (t_flat >= 0)

                if reduction == 'none':
                    per_tok = sampled_softmax_loss(
                        h_flat, t_flat, self.lm_head.weight, unigram_sampler, K_neg,
                        label_smoothing=0.0, softcap=softcap,
                        softcap_clamp=_softcap_clamp,
                        valid_mask=valid_mask_flat, reduction='none',
                    )
                    return per_tok

                out = sampled_softmax_loss(
                    h_flat, t_flat, self.lm_head.weight, unigram_sampler, K_neg,
                    label_smoothing=smoothing, softcap=softcap,
                    softcap_clamp=_softcap_clamp,
                    valid_mask=valid_mask_flat, reduction='mean',
                )
            elif use_sampled:
                # ----- Legacy uniform-negative path (fallback) -----
                # Flatten hidden states + targets
                h_flat = x.reshape(-1, x.shape[-1])            # (B*T, d)
                t_flat = targets.reshape(-1)                    # (B*T,)
                n = h_flat.shape[0]

                # Learnability #4 hardening: sampled-softmax gather crashes on
                # negative ids (-1 from doc-sep mask). Replace -1 with 0 for
                # gather; the actual loss is masked below.
                valid_mask_flat = (t_flat >= 0)
                t_flat_safe = torch.where(valid_mask_flat, t_flat, torch.zeros_like(t_flat))

                # Sample K negatives uniformly from [0, V)
                neg_ids = torch.randint(0, V, (K_neg,), device=x.device)
                # Gather lm_head weights for target + negatives
                all_ids = torch.cat([t_flat_safe, neg_ids])     # (B*T + K,)
                sampled_w = self.lm_head.weight[all_ids]        # (B*T + K, d)

                # Compute sampled logits: for each position, dot with its
                # target weight and all K negative weights.
                # Target logit: dot product of h[i] with w[target[i]]
                target_w = sampled_w[:n]                        # (B*T, d)
                neg_w = sampled_w[n:]                           # (K, d)
                target_logit = (h_flat * target_w).sum(-1)      # (B*T,)
                neg_logits = h_flat @ neg_w.t()                 # (B*T, K)

                if not _softcap_clamp:
                    target_logit = softcap * torch.tanh(target_logit / softcap)
                    neg_logits = softcap * torch.tanh(neg_logits / softcap)

                # Sampled softmax loss: -log(exp(target) / (exp(target) + sum(exp(neg))))
                # With log-sum-exp correction for sampling K of V negatives.
                # Correction: add log(V/K) to negative logits to account for
                # the fact that we're only seeing K of V possible negatives.
                log_correction = torch.tensor(V / K_neg, device=x.device).log()
                all_logits = torch.cat([
                    target_logit.unsqueeze(-1),                 # (B*T, 1)
                    neg_logits + log_correction,                # (B*T, K)
                ], dim=-1).float()                              # (B*T, K+1)

                # CE with target always at index 0
                ce_targets = torch.zeros(n, dtype=torch.long, device=x.device)
                if reduction == 'none':
                    per_tok = F.cross_entropy(all_logits, ce_targets, reduction='none')
                    if self._doc_sep_mask and self._bos_token_id >= 0:
                        per_tok = torch.where(valid_mask_flat, per_tok, torch.zeros_like(per_tok))
                    return per_tok
                per_tok_ce = F.cross_entropy(
                    all_logits, ce_targets, reduction='none',
                    label_smoothing=smoothing,
                )
                # Mask doc-separator positions. valid_mask_flat is always
                # computed; when doc_sep_mask is off every token is valid so
                # this reduces to a plain mean.
                valid_f = valid_mask_flat.float()
                valid_n = valid_f.sum().clamp(min=1)
                out = (per_tok_ce * valid_f).sum() / valid_n
            else:
                # Full softmax path (eval or HYDRA_SAMPLED_SOFTMAX=0)
                chunk_size = int(os.environ.get("HYDRA_CE_CHUNK", "1024"))
                if chunk_size <= 0:
                    MAX_LOGITS_BYTES = 256 * 1024 * 1024
                    tokens_per_chunk = max(V, MAX_LOGITS_BYTES // (V * 4))
                    chunk_size = max(1, tokens_per_chunk // max(1, B))
                chunk_size = min(chunk_size, T)

                if reduction == 'none':
                    loss_parts = []
                    for start in range(0, T, chunk_size):
                        end = min(start + chunk_size, T)
                        chunk_logits = self.lm_head(x[:, start:end, :]).float()
                        if _softcap_clamp:
                            chunk_logits = torch.clamp(chunk_logits, -softcap, softcap)
                        else:
                            chunk_logits = softcap * torch.tanh(chunk_logits / softcap)
                        chunk_targets = targets[:, start:end].reshape(-1)
                        chunk_loss = F.cross_entropy(
                            chunk_logits.view(-1, chunk_logits.size(-1)),
                            chunk_targets, ignore_index=-1, reduction='none',
                        )
                        loss_parts.append(chunk_loss)
                    return torch.cat(loss_parts)

                total_loss = 0.0
                total_tokens = 0
                for start in range(0, T, chunk_size):
                    end = min(start + chunk_size, T)
                    chunk_logits = self.lm_head(x[:, start:end, :]).float()
                    if _softcap_clamp:
                        chunk_logits = torch.clamp(chunk_logits, -softcap, softcap)
                    else:
                        chunk_logits = softcap * torch.tanh(chunk_logits / softcap)
                    chunk_targets = targets[:, start:end].reshape(-1)
                    chunk_loss = F.cross_entropy(
                        chunk_logits.view(-1, chunk_logits.size(-1)),
                        chunk_targets, ignore_index=-1, reduction='sum',
                        label_smoothing=smoothing,
                    )
                    total_loss = total_loss + chunk_loss
                    total_tokens += (chunk_targets != -1).sum()
                out = total_loss / total_tokens

            # -----------------------------------------------------------
            # Learnability #1: Multi-Token Prediction.
            # For k in {2..K}, add a CE loss at position (t) predicting
            # the token at position (t+k), using the SAME lm_head weights
            # (weight-tied). Cost: K-1 extra CEs on a subset of positions.
            # Only triggered in reduction='mean' path, training only.
            # -----------------------------------------------------------
            if reduction == 'mean' and self._mtp_k > 1 and self.training and use_sampled:
                # Audit 2026-05-09 issue #22 - Cluster E. When a UnigramSampler
                # is attached each MTP head draws its OWN per-row negatives
                # (independent of the primary head). This matches the math of
                # importance-sampled NCE (each loss term has its own negative
                # pool drawn from the same proposal q) and removes the bias
                # from sharing one batch of negatives across K heads. Cost:
                # one extra alias-method sample call per head - cheap relative
                # to the einsum.
                # Legacy (uniform) path retains the original neg_logits reuse
                # for backward-compat with environments that have no sampler.
                mtp_loss_sum = out.new_tensor(0.0)
                mtp_terms = 0
                neg_logits_bt = None
                if unigram_sampler is None and neg_logits is not None:
                    neg_logits_bt = neg_logits.view(B, T, K_neg)

                for k in range(2, self._mtp_k + 1):
                    shift = k - 1
                    if T <= shift:
                        continue
                    n_k = B * (T - shift)
                    h_k_flat = x[:, :T - shift, :].reshape(n_k, -1)  # (n_k, d)
                    t_k = targets[:, shift:].reshape(-1)             # (n_k,)
                    mask_k = (t_k >= 0)

                    if unigram_sampler is not None:
                        # Independent unigram-distributed negatives per head.
                        per_tok_ce_k = sampled_softmax_loss(
                            h_k_flat, t_k, self.lm_head.weight,
                            unigram_sampler, K_neg,
                            label_smoothing=smoothing, softcap=softcap,
                            softcap_clamp=_softcap_clamp,
                            valid_mask=mask_k, reduction='none',
                        )
                        # sampled_softmax_loss already zeros invalid rows.
                        n_valid_k = mask_k.float().sum().clamp(min=1)
                        mtp_loss_sum = mtp_loss_sum + per_tok_ce_k.sum() / n_valid_k
                    else:
                        # Legacy path: reuse primary's negatives. neg_logits_bt
                        # is non-None here because the legacy branch runs only
                        # when use_sampled and unigram_sampler is None.
                        t_k_safe = torch.where(mask_k, t_k, torch.zeros_like(t_k))
                        tgt_w_k = self.lm_head.weight[t_k_safe]      # (n_k, d)
                        tgt_logit_k = (h_k_flat * tgt_w_k).sum(-1)   # (n_k,)
                        if not _softcap_clamp:
                            tgt_logit_k = softcap * torch.tanh(tgt_logit_k / softcap)
                        neg_logits_k = neg_logits_bt[:, :T - shift, :].reshape(n_k, K_neg)
                        all_logits_k = torch.cat([
                            tgt_logit_k.unsqueeze(-1),
                            neg_logits_k + log_correction,
                        ], dim=-1).float()
                        ce_targets_k = torch.zeros(n_k, dtype=torch.long, device=x.device)
                        per_tok_ce_k = F.cross_entropy(
                            all_logits_k, ce_targets_k, reduction='none',
                            label_smoothing=smoothing,
                        )
                        per_tok_ce_k = torch.where(mask_k, per_tok_ce_k, torch.zeros_like(per_tok_ce_k))
                        n_valid_k = mask_k.sum().clamp(min=1)
                        mtp_loss_sum = mtp_loss_sum + per_tok_ce_k.sum() / n_valid_k
                    mtp_terms += 1
                if mtp_terms > 0:
                    out = (out + mtp_loss_sum) / float(mtp_terms + 1)

            # -----------------------------------------------------------
            # Learnability #6: output entropy penalty.
            # L += -lambda * H(softmax(logits)). Negative entropy penalizes
            # peaked distributions; encourages diverse predictions and
            # breaks repetition loops. Computed on a small subset of
            # positions to keep V-sized logits cost bounded.
            # -----------------------------------------------------------
            if reduction == 'mean' and self._entropy_penalty > 0.0 and self.training:
                # Sample up to 64 random positions. V-sized logits on 64
                # positions = 64 * V * 4 bytes (~50 MB at V=200k) — fits
                # on the 3060 and adds ~2 ms.
                h_flat = x.reshape(-1, x.shape[-1])
                n_pos = h_flat.shape[0]
                n_sample = min(64, n_pos)
                idx_sample = torch.randint(0, n_pos, (n_sample,), device=x.device)
                h_sample = h_flat[idx_sample]
                logits_s = F.linear(h_sample, self.lm_head.weight).float()
                if _softcap_clamp:
                    logits_s = torch.clamp(logits_s, -softcap, softcap)
                else:
                    logits_s = softcap * torch.tanh(logits_s / softcap)
                log_probs = F.log_softmax(logits_s, dim=-1)
                probs = log_probs.exp()
                entropy = -(probs * log_probs).sum(-1).mean()   # scalar, nats
                out = out - self._entropy_penalty * entropy

            if _profile:
                _t_end = _ev()
                torch.cuda.synchronize()
                def _ms(a, b): return a.elapsed_time(b)
                print(
                    f"[PROFILE B={B} T={T}] "
                    f"htm_launch={_ms(_t0, _t_htm_async):.2f} "
                    f"wte={_ms(_t_htm_async, _t_wte):.2f} "
                    f"htm_await={_ms(_t_wte, _t_htm_await):.2f} "
                    f"htm_proj={_ms(_t_htm_await, _t_htm_proj):.2f} "
                    f"mamba_mhc_engram={_ms(_t_htm_proj, _t_blocks):.2f} "
                    f"merge={_ms(_t_blocks, _t_merge):.2f} "
                    f"lm_head_loss={_ms(_t_merge, _t_end):.2f} "
                    f"total={_ms(_t0, _t_end):.2f} ms",
                    flush=True,
                )
            return out

        logits = self.lm_head(x).float()
        if _softcap_clamp:
            logits = torch.clamp(logits, -softcap, softcap)
        else:
            logits = softcap * torch.tanh(logits / softcap)
        return logits