Update modeling_neollm.py

modeling_neollm.py

@@ -81,6 +81,7 @@ from typing import Optional, Union, Tuple
 import torch
 import torch.nn.functional as F
 from torch import nn
+from torch.utils.checkpoint import checkpoint as _activation_checkpoint
 
 try:
     from cut_cross_entropy import linear_cross_entropy  # type: ignore[import-not-found]
@@ -95,6 +96,35 @@ try:
 except ImportError:
     meap_mask_inputs = None
 
+try:
+    from cut_cross_entropy.polynorm import (  # type: ignore[import-not-found]
+        polynorm as _cce_polynorm,
+    )
+except (ImportError, ModuleNotFoundError):
+    _cce_polynorm = None
+
+try:
+    from cut_cross_entropy.polynorm import (  # type: ignore[import-not-found]
+        polynorm_uses_cute as _cce_polynorm_uses_cute,
+    )
+except (ImportError, ModuleNotFoundError):
+    _cce_polynorm_uses_cute = None
+
+_REPO_GRAPE_KERNEL_AVAILABLE = False
+_cce_repo_grape = None
+_cce_repo_grape_supported = None
+try:
+    from cut_cross_entropy.repo_grape import (  # type: ignore[import-not-found]
+        repo_grape as _cce_repo_grape,
+        repo_grape_supported as _cce_repo_grape_supported,
+    )
+
+    _REPO_GRAPE_KERNEL_AVAILABLE = True
+except (ImportError, ModuleNotFoundError, AttributeError):
+    # REPO-GRAPE remains fully functional through RepoGrapePositioning when the
+    # optional Triton package (or a new-enough CCE checkout) is unavailable.
+    pass
+
 _LEV_KERNEL_AVAILABLE = False
 _leviathan_embedding_compiler_safe = None
 _leviathan_supports = None
@@ -497,7 +527,13 @@ class LeviathanGenerator(nn.Module):
     # in the compiled graph while Inductor excludes only the unsafe Triton
     # launch from CUDA-graph capture. A CUDA launch error is allowed to
     # propagate instead of falling back to reference work on the same stream.
-    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self,
+        token_ids: torch.Tensor,
+        *,
+        meap_mask_embedding: Optional[torch.Tensor] = None,
+        meap_mask_token_id: Optional[int] = None,
+    ) -> torch.Tensor:
         """
         Generate embeddings from discrete token indices.
 
@@ -508,6 +544,24 @@ class LeviathanGenerator(nn.Module):
         Returns:
             embeddings [*token_ids.shape, hidden_size].
         """
+        has_meap = (
+            meap_mask_embedding is not None or meap_mask_token_id is not None
+        )
+        if has_meap:
+            if meap_mask_embedding is None or meap_mask_token_id is None:
+                raise ValueError(
+                    "meap_mask_embedding and meap_mask_token_id must be "
+                    "provided together"
+                )
+            if (
+                meap_mask_embedding.ndim != 1
+                or meap_mask_embedding.numel() != self.hidden_size
+            ):
+                raise ValueError(
+                    "meap_mask_embedding must have shape "
+                    f"({self.hidden_size},)"
+                )
+
         if (
             self.use_leviathan_triton
             and _LEV_KERNEL_AVAILABLE
@@ -541,6 +595,8 @@ class LeviathanGenerator(nn.Module):
                     params,
                     self.config,
                     self.knot_grid,
+                    mask_embedding=meap_mask_embedding,
+                    mask_token_id=meap_mask_token_id,
                 )
             except (TypeError, ValueError, AttributeError):
                 # Only validation/configuration failures use the reference
@@ -592,7 +648,14 @@ class LeviathanGenerator(nn.Module):
 
             e = e + modes.to(self.head_out_weight.dtype) @ self.head_out_weight[m]
 
-        return e.reshape(*orig_shape, self.hidden_size)
+        e = e.reshape(*orig_shape, self.hidden_size)
+        if has_meap:
+            e = torch.where(
+                token_ids.eq(int(meap_mask_token_id)).unsqueeze(-1),
+                meap_mask_embedding.to(device=e.device, dtype=e.dtype),
+                e,
+            )
+        return e
 
 
 
@@ -870,6 +933,18 @@ class NeoLLMRotaryEmbedding(nn.Module):
         self.register_buffer("inv_freq", inv_freq, persistent=False)
         self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
 
+    def _apply(self, fn):
+        # Positional frequencies are an FP32 numerical contract for both the
+        # PyTorch and fused REPO-GRAPE paths. Preserve that contract across
+        # model.to(bfloat16)/model.to(float16) instead of allocating a cast in
+        # every decoder layer. The buffers remain non-persistent as before.
+        super()._apply(fn)
+        for name in ("inv_freq", "original_inv_freq"):
+            value = self._buffers.get(name)
+            if value is not None:
+                self._buffers[name] = value.float()
+        return self
+
     @staticmethod
     def compute_default_rope_parameters(
         config: NeoLLMConfig = None,
@@ -1697,10 +1772,16 @@ class REPOModule(nn.Module):
     where ``h_i ∈ R^d`` is the hidden state of token ``i`` entering the
     decoder layer (pre-FANLayer), and ``d_p = hidden_size // 8`` by default.
 
-    The resulting positions ``z [B, H, S]`` are real-valued and
-    unconstrained.  They are used to compute per-head ``cos/sin`` embeddings
-    inline from ``inv_freq``, replacing the standard integer-based
-    ``position_embeddings`` for this layer.
+    The resulting assignments ``z [B, H, S]`` are real-valued and
+    unconstrained.  In the original REPO-RoPE path they are used directly to
+    compute per-head ``cos/sin`` embeddings.  When REPO is composed with
+    REPO-GRAPE, ``RepoGrapePositioning`` keeps ``z`` as the raw REPO output and
+    constructs the final coordinate
+
+        u_i^(h) = i + alpha_h (z_i^(h) - i)
+
+    before the GRAPE phase is formed.  This keeps the semantic distinction
+    ``REPO -> z`` versus ``REPO-GRAPE -> u`` explicit in the implementation.
 
     Design notes:
     - ``W_g`` and ``W_c`` are shared across heads (parameter efficiency).
@@ -1870,20 +1951,64 @@ def _apply_repo_rope(
 
 class RepoGrapePositioning(nn.Module):
     r"""
-    Minimal REPO-GRAPE-M positional operator.
+    Minimal REPO-GRAPE-M positional operator with linear-blend coordinates.
+
+    The implementation deliberately has no positional-mode flag.  Whenever
+    REPO-GRAPE is active it uses exactly one coordinate map, the
+    ``linear-blend-repo`` form from the accompanying derivation:
+
+        z_i^(h) = f_phi^(h)(h_i),
+        u_i^(h) = (1 - alpha_h) i + alpha_h z_i^(h)
+                = i + alpha_h (z_i^(h) - i).
 
-    This module intentionally keeps only the specific multiplicative path:
+    ``i`` is the actual runtime ``position_ids`` value, not an independently
+    generated global arange.  This matters for padding and packed/reset
+    position streams.  ``alpha_h`` is one learned scalar per query head.  It
+    is initialized to 1, so a newly constructed model starts exactly at the
+    previous REPO coordinate u=z and can learn a linear anchor without an
+    abrupt initialization change.
 
-        1. REPO predicts contextual coordinates z_i^(h)=f_phi^(h)(h_i).
-        2. REPO-GRAPE uses u_i^(h)=z_i^(h), with no extra position mode.
-        3. GRAPE-M applies canonical commuting SO(d) rotary planes to Q/K.
-        4. Each query head learns its own angular spectrum
+    The parameter is intentionally stored directly rather than through a
+    sigmoid/clamp.  Therefore alpha_h is initialized inside [0,1] but is not
+    constrained to stay there.  This avoids an additional nonlinear operation,
+    preserves the exact alpha_h=1 initialization, and lets the optimizer learn
+    interpolation or extrapolation.  Constraining alpha would be a different
+    parameterization and is not silently introduced here.
 
-              theta_{h,r}=inv_freq_r*exp(s_{h,r}),
+    Backward is standard autograd.  If ubar_i^(h)=dL/du_i^(h), then for this
+    active linear-blend coordinate
 
-           and therefore its own phase
+        dL/dz_i^(h) = alpha_h * ubar_i^(h),
+        dL/dalpha_h = sum_i ubar_i^(h) * (z_i^(h) - i).
 
-              phi_{b,h,i,r}=u_{b,h,i}*theta_{h,r}.
+    For comparison only, the distinct ``linear-plus-repo`` residual map
+
+        u_i^(h) = i + alpha_h z_i^(h)
+
+    would instead give dL/dalpha_h = sum_i ubar_i^(h) z_i^(h).  That formula
+    is documented to avoid confusing the two maps, but it is not another
+    executable mode and no flag is added for it.
+
+    Storage/lifetime contract:
+    - ``position_blend_alpha`` is an FP32 ``nn.Parameter`` of shape [H_q].
+      It is optimized normally and is saved/restored in ``state_dict``.
+    - ``freq_log_scale`` remains an FP32 saved parameter of shape
+      [H_q, head_dim/2].
+    - ``position_ids`` is runtime metadata and is never copied into a buffer or
+      checkpoint.
+    - ``u`` and the optional IHA-expanded ``u_{i,p}`` are transient autograd
+      intermediates only.  The IHA tensor replaces the previous ``P*z+p``
+      intermediate, so this change does not add another expanded coordinate
+      tensor.
+
+    IHA is applied after the coordinate map.  For P pseudo-slots,
+
+        u_{i,p}^(h) = P u_i^(h) + p,    p in {0,...,P-1},
+
+    and GRAPE sees the phase
+
+        phi_{b,h,i,p,r} = u_{b,h,i,p} * theta_{h,r},
+        theta_{h,r} = inv_freq_r * exp(s_{h,r}).
 
     With one key per query head, the multiplicative action obeys the exact
     relative law
@@ -1894,32 +2019,23 @@ class RepoGrapePositioning(nn.Module):
     GQA introduces a separate geometric restriction: several query heads
     share one key, although their phases phi_{b,h,i,r} can differ.  A shared
     key cannot realize every head-specific rotation simultaneously.  For a
-    KV group H_m and rotary plane r, this implementation therefore projects
-    the requested rotations onto the single optimal shared SO(2) rotation:
+    KV group H_m and rotary plane r, this implementation projects the requested
+    rotations onto the single optimal shared SO(2) rotation:
 
         Z_{b,m,i,r} = (1/g) * sum_{h in H_m} exp(i*phi_{b,h,i,r})
                     = C_{b,m,i,r} + i*S_{b,m,i,r},
 
-        psi*_{b,m,i,r} = arg Z_{b,m,i,r},
+        psi*_{b,m,i,r} = arg Z_{b,m,i,r}.
 
-    or, equivalently when |Z| > 0,
+    Equivalently, when |Z| > 0,
 
         cos(psi*) = C/|Z|,
         sin(psi*) = S/|Z|.
 
-    This is the global minimizer of
-
-        sum_{h in H_m} ||R(phi_{b,h,i,r})k - R(psi)k||_2^2,
-
-    so it replaces the previous arithmetic approximation
-
-        mean(u_h) * mean(theta_h),
-
-    which neither averages the effective phase u_h*theta_h nor respects its
-    2*pi periodicity.  If the circular resultant is numerically zero, the
-    objective does not identify a preferred shared direction; the code uses
-    the identity rotation as a deterministic finite convention.  Under IHA,
-    the projection is computed independently inside every pseudo-slot.
+    If the circular resultant is numerically zero, the objective does not
+    identify a preferred shared direction; the code uses the identity rotation
+    as a deterministic finite convention.  Under IHA, the projection is
+    computed independently inside every pseudo-slot.
     """
 
     def __init__(
@@ -1938,8 +2054,21 @@ class RepoGrapePositioning(nn.Module):
         self.head_dim = int(head_dim)
         self.max_rot_half = self.head_dim // 2
 
-        # The single REPO-GRAPE parameter: a learned log-scale over the
-        # base RoPE/GRAPE frequencies, per query head and rotary plane.
+        # Linear-blend coefficient alpha_h, one scalar per query head:
+        #
+        #     u_i^(h) = i + alpha_h * (z_i^(h) - i)
+        #
+        # Initialize at alpha_h=1 so the first forward is exactly the previous
+        # REPO behavior u=z.  This is a real nn.Parameter, therefore it is
+        # optimized by autograd/AdamW and stored in every state_dict checkpoint.
+        # It is deliberately unconstrained: no sigmoid/clamp and no new flag.
+        self.position_blend_alpha = nn.Parameter(
+            torch.ones(self.num_attention_heads, dtype=torch.float32),
+            requires_grad=True,
+        )
+
+        # Learned log-scale over the base RoPE/GRAPE frequencies, per query
+        # head and rotary plane: theta_{h,r}=theta^0_r*exp(s_{h,r}).
         self.freq_log_scale = nn.Parameter(
             torch.zeros(
                 self.num_attention_heads, self.max_rot_half, dtype=torch.float32
@@ -1947,21 +2076,108 @@ class RepoGrapePositioning(nn.Module):
             requires_grad=True,
         )
 
+    def _load_from_state_dict(
+        self,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
+        # Backward-compatible checkpoint migration.  Older REPO-GRAPE
+        # checkpoints do not contain alpha_h because they used u=z.  Injecting
+        # alpha_h=1 reproduces that old forward exactly at load time, while every
+        # checkpoint saved by this version stores the learned parameter normally.
+        alpha_key = prefix + "position_blend_alpha"
+        local_state_dict = state_dict
+        if alpha_key not in state_dict:
+            # Use a local shallow copy so loading does not mutate the caller's
+            # checkpoint mapping.  CPU FP32 also works when the module was
+            # constructed on meta before Hugging Face materializes parameters.
+            local_state_dict = state_dict.copy()
+            local_state_dict[alpha_key] = torch.ones(
+                self.num_attention_heads, dtype=torch.float32
+            )
+        super()._load_from_state_dict(
+            local_state_dict,
+            prefix,
+            local_metadata,
+            strict,
+            missing_keys,
+            unexpected_keys,
+            error_msgs,
+        )
+
     def _apply(self, fn):
-        # Keep the small spectral parameters in fp32 under model.to(dtype).
+        # Keep the tiny positional parameters in FP32 under model.to(dtype).
+        # Both are checkpointed nn.Parameters; only runtime coordinates are
+        # transient.
         super()._apply(fn)
-        p = self.freq_log_scale
-        p.data = p.data.float()
-        if p.grad is not None:
-            p.grad.data = p.grad.data.float()
+        for p in (self.position_blend_alpha, self.freq_log_scale):
+            p.data = p.data.float()
+            if p.grad is not None:
+                p.grad.data = p.grad.data.float()
         return self
 
     def transform_positions(
         self,
         z: torch.Tensor,
+        position_ids: Optional[torch.Tensor],
     ) -> torch.Tensor:
-        """Use the REPO coordinate directly: u_i^(h)=z_i^(h)."""
-        return z
+        r"""Build the active linear-blend REPO-GRAPE coordinate.
+
+        The only executable map is
+
+            u_i^(h) = (1-alpha_h) i + alpha_h z_i^(h)
+                    = i + alpha_h (z_i^(h) - i).
+
+        The residualized second form is used because it needs only one
+        subtraction, one multiplication and one addition.  ``position_ids`` is
+        broadcast as [B,1,S], so it is not materialized once per head.
+
+        ``u`` is evaluated in FP32 because the next operation is the FP32 phase
+        computation.  It is not registered as a buffer and is not saved in the
+        checkpoint.  Standard autograd gives
+
+            dL/dalpha_h = sum_{b,i} ubar_{b,h,i} (z_{b,h,i} - i).
+        """
+        B, H, S = z.shape
+        if H != self.num_attention_heads:
+            raise ValueError(
+                f"REPO-GRAPE expected z heads={self.num_attention_heads}, got {H}."
+            )
+
+        if position_ids is None:
+            # Defensive fallback for direct attention calls. NeoLLMModel normally
+            # constructs position_ids before entering the decoder layers.
+            pid = torch.arange(S, device=z.device, dtype=torch.float32).view(1, S)
+        else:
+            pid = position_ids
+            if pid.dim() == 1:
+                pid = pid.unsqueeze(0)
+            elif pid.dim() > 2:
+                # Match the existing IHA RoPE convention: trainers may attach
+                # metadata in the final axis; column 0 is the sequential index.
+                pid = pid[..., 0]
+            if pid.dim() != 2 or pid.shape[-1] != S:
+                raise ValueError(
+                    "REPO-GRAPE position_ids must have shape [S] or [B,S] "
+                    f"matching z[...,S={S}], got {tuple(pid.shape)}."
+                )
+            if pid.shape[0] == 1 and B != 1:
+                pid = pid.expand(B, -1)
+            elif pid.shape[0] != B:
+                raise ValueError(
+                    f"REPO-GRAPE position_ids batch={pid.shape[0]} does not match z batch={B}."
+                )
+            pid = pid.to(device=z.device, dtype=torch.float32)
+
+        # Views/broadcasting do not allocate [B,H,S] copies of pid or alpha.
+        i = pid.unsqueeze(1)
+        alpha = self.position_blend_alpha.view(1, H, 1)
+        return i + alpha * (z.float() - i)
 
     def _query_freq(
         self,
@@ -1989,16 +2205,18 @@ class RepoGrapePositioning(nn.Module):
         self,
         q: torch.Tensor,
         k: torch.Tensor,
-        z: torch.Tensor,
+        u: torch.Tensor,
         inv_freq: torch.Tensor,
         attention_scaling: float,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         r"""
-        Apply REPO-GRAPE-M rotation to Q/K using contextual REPO coordinates.
+        Apply REPO-GRAPE-M rotation to Q/K using final coordinates ``u``.
 
+        ``u`` is already the linear-blend coordinate produced by
+        ``transform_positions`` (and, for IHA, already expanded as P*u+p).
         Query heads keep their individual phases
 
-            phi_{b,h,i,r} = z_{b,h,i} * theta_{h,r}.
+            phi_{b,h,i,r} = u_{b,h,i} * theta_{h,r}.
 
         For GQA, each KV head m serves a group H_m of query heads.  Its shared
         rotation is the circular projection
@@ -2018,7 +2236,7 @@ class RepoGrapePositioning(nn.Module):
         Args:
             q:                 [B, H_q_eff, S, head_dim]
             k:                 [B, H_k_eff, S, head_dim]
-            z:                 [B, H_q_base, S]
+            u:                 [B, H_q_base, S] final REPO-GRAPE coordinates
             inv_freq:          [rotary_dim/2]
             attention_scaling: scaling from NeoLLMRotaryEmbedding
 
@@ -2026,13 +2244,13 @@ class RepoGrapePositioning(nn.Module):
             Rotated ``(q, k)`` tensors with the same shapes and dtypes as the
             inputs. Channels beyond ``rotary_dim`` pass through unchanged.
         """
-        B, H_repo, S = z.shape
+        B, H_repo, S = u.shape
         H_q_eff = q.shape[1]
         H_k_eff = k.shape[1]
 
         if H_q_eff % H_repo != 0:
             raise ValueError(
-                f"REPO-GRAPE expected q heads divisible by z heads, "
+                f"REPO-GRAPE expected q heads divisible by coordinate heads, "
                 f"got H_q={H_q_eff} vs H_repo={H_repo}."
             )
         P = H_q_eff // H_repo
@@ -2043,7 +2261,7 @@ class RepoGrapePositioning(nn.Module):
                 f"got H_k={H_k_eff} vs P={P}."
             )
 
-        z_q_struct = z.unsqueeze(2).expand(B, H_repo, P, S)  # [B,H_q_base,P,S]
+        u_q_struct = u.unsqueeze(2).expand(B, H_repo, P, S)  # [B,H_q_base,P,S]
         H_k_base = H_k_eff // max(P, 1)
         if H_repo % H_k_base != 0:
             raise ValueError(
@@ -2052,7 +2270,7 @@ class RepoGrapePositioning(nn.Module):
             )
 
         q_per_k_base = H_repo // H_k_base
-        z_q = z_q_struct.reshape(B, H_q_eff, S)
+        u_q = u_q_struct.reshape(B, H_q_eff, S)
 
         rot_half = int(inv_freq.shape[0])
         rotary_dim = rot_half * 2
@@ -2064,7 +2282,7 @@ class RepoGrapePositioning(nn.Module):
         # Individual query phases are the fundamental quantity.  They are
         # calculated in fp32 before trigonometric evaluation, as in the
         # original REPO-GRAPE path.
-        phase_q = z_q.float().unsqueeze(-1) * freq_q.to(device=q.device).view(
+        phase_q = u_q.float().unsqueeze(-1) * freq_q.to(device=q.device).view(
             1, H_q_eff, 1, rot_half
         )
         cos_q_half = phase_q.cos()
@@ -2681,6 +2899,9 @@ class NeoLLMAttention(nn.Module):
         # REPO-GRAPE reuses the REPO coordinate module f_phi and replaces only
         # the positional action applied to Q/K. Enabling use_repo_grape=True
         # therefore activates the REPO coordinate path even when use_repo=False.
+        # There is no extra position-mode flag: active REPO-GRAPE always maps
+        # raw z to u=i+alpha_h*(z-i) with one learned FP32 alpha_h per query
+        # head before GRAPE-M (and before the IHA P*u+p expansion).
         _repo_requested = bool(getattr(config, "use_repo", False))
         _repo_grape_requested = bool(getattr(config, "use_repo_grape", False))
         _repo_start_layer = getattr(
@@ -2709,8 +2930,14 @@ class NeoLLMAttention(nn.Module):
                 num_key_value_heads=config.num_key_value_heads,
                 head_dim=self.head_dim,
             )
+            # Availability is fixed when the module is constructed, so Dynamo
+            # specializes one stable route. This is deliberately not a new
+            # configuration flag: use_repo_grape still owns the feature, while
+            # unsupported inputs retain the exact PyTorch implementation.
+            self.use_repo_grape_triton = False
         else:
             self.repo_grape = None
+            self.use_repo_grape_triton = False
 
         # ── REPO-GOAT factorised log-prior ───────────────────────────────────
         # Independent of use_repo/use_repo_grape: it can be ablated with standard
@@ -3791,6 +4018,7 @@ class NeoLLMAttention(nn.Module):
 
         # ── RoPE / REPO ───────────────────────────────────────────────────────
         cos, sin = position_embeddings
+        _repo_grape_kernel_used = False
         if self.use_repo:
             # REPO path: f_ϕ predicts continuous per-head positions from the
             # residual stream, then cos/sin are built inline from those positions
@@ -3799,34 +4027,107 @@ class NeoLLMAttention(nn.Module):
             # directly from rotary_emb at forward time — no buffer on this module,
             # no meta-tensor issue on lm_eval / to(device) paths.
             # (Li et al., 2026, §3.2 — Eq. 6–7)
-            z = self.repo_module(hidden_states)  # [B, H, S]
-            if self.use_repo_grape:
-                z = self.repo_grape.transform_positions(z)
+            z = self.repo_module(hidden_states)  # [B, H, S], raw REPO assignment
             inv_freq, attn_scaling = repo_rope_args
-            if _iha_seq_expand:
-                # IHA seq-expand + REPO/REPO-GRAPE: each real position s has P virtual slots.
-                # Assign continuous position z[b,h,s]·P + p to slot (s,p), scaling z by P
-                # to create room for P sub-positions while preserving relative ordering.
+
+            # The optional kernel consumes raw REPO assignments and performs
+            # the linear-blend coordinate, IHA P=2 sequence expansion, GRAPE
+            # phases, GQA circular projection, and optional momentum in one
+            # autograd-aware Triton boundary. Q/K are intentionally normalized
+            # before this call: IHA mixes normalized heads, so moving RMSNorm
+            # into the kernel after that mixing would change the model.
+            _repo_grape_sequence_factor = self.iha_P if _iha_seq_expand else 1
+            if (
+                self.use_repo_grape
+                and self.use_repo_grape_triton
+                and _cce_repo_grape is not None
+                and _cce_repo_grape_supported is not None
+                and _repo_grape_sequence_factor in (1, 2)
+                and _cce_repo_grape_supported(
+                    q,
+                    k,
+                    z,
+                    inv_freq,
+                    self.repo_grape.position_blend_alpha,
+                    self.repo_grape.freq_log_scale,
+                    sequence_pseudo_factor=_repo_grape_sequence_factor,
+                )
+            ):
+                _kernel_position_ids = position_ids
+                if (
+                    _kernel_position_ids is not None
+                    and _kernel_position_ids.dim() > 2
+                ):
+                    # Match transform_positions and the IHA helpers: the first
+                    # metadata column is the runtime sequential position.
+                    _kernel_position_ids = _kernel_position_ids[..., 0]
+                _kernel_gamma = (
+                    self.momentum_gamma if self.use_momentum_attention else 0.0
+                )
+                q, k = _cce_repo_grape(
+                    q,
+                    k,
+                    z,
+                    _kernel_position_ids,
+                    inv_freq,
+                    self.repo_grape.position_blend_alpha,
+                    self.repo_grape.freq_log_scale,
+                    attn_scaling,
+                    sequence_pseudo_factor=_repo_grape_sequence_factor,
+                    momentum_gamma=_kernel_gamma,
+                    output_dtype=q.dtype,
+                )
+                _repo_grape_kernel_used = True
+
+            # REPO-GRAPE consumes u, not raw z.  There is intentionally no mode
+            # flag: when REPO-GRAPE is active, transform_positions always applies
+            # the learned linear blend
+            #
+            #     u_i^(h) = i + alpha_h * (z_i^(h) - i).
+            #
+            # Pure REPO (without GRAPE) keeps the original paper behavior and
+            # therefore continues to use z directly.
+            if not _repo_grape_kernel_used:
+                repo_coords = (
+                    self.repo_grape.transform_positions(z, position_ids)
+                    if self.use_repo_grape
+                    else z
+                )
+
+            if not _repo_grape_kernel_used and _iha_seq_expand:
+                # IHA is applied *after* the coordinate map.  Thus REPO-GRAPE
+                # uses u_{i,p}=P*u_i+p, while pure REPO keeps its previous
+                # z_{i,p}=P*z_i+p behavior.  ``repo_coords`` is transient and this
+                # expanded tensor replaces, rather than accompanies, the old
+                # z_expanded allocation.
                 P = self.iha_P
-                p_offsets = torch.arange(P, device=z.device, dtype=z.dtype)
-                # z[B,H,S] → [B,H,S,1]*P + [1,1,1,P] → [B,H,S,P] → [B,H,S·P]
-                z_expanded = (z.unsqueeze(-1) * P + p_offsets.view(1, 1, 1, P)).reshape(
-                    z.shape[0], z.shape[1], _S_orig * P
+                p_offsets = torch.arange(
+                    P, device=repo_coords.device, dtype=repo_coords.dtype
                 )
+                coords_expanded = (
+                    repo_coords.unsqueeze(-1) * P
+                    + p_offsets.view(1, 1, 1, P)
+                ).reshape(repo_coords.shape[0], repo_coords.shape[1], _S_orig * P)
+
                 if self.use_repo_grape:
                     q, k = self.repo_grape.apply_multiplicative(
-                        q, k, z_expanded, inv_freq, attn_scaling
+                        q, k, coords_expanded, inv_freq, attn_scaling
                     )
                 else:
-                    q, k = _apply_repo_rope(q, k, z_expanded, inv_freq, attn_scaling)
-            else:
-                # Standard REPO path (P=1 or no IHA): z covers original S positions.
+                    q, k = _apply_repo_rope(
+                        q, k, coords_expanded, inv_freq, attn_scaling
+                    )
+            elif not _repo_grape_kernel_used:
+                # P=1/no-IHA: ``repo_coords`` is u for REPO-GRAPE and z for
+                # the original REPO-RoPE path.
                 if self.use_repo_grape:
                     q, k = self.repo_grape.apply_multiplicative(
-                        q, k, z, inv_freq, attn_scaling
+                        q, k, repo_coords, inv_freq, attn_scaling
                     )
                 else:
-                    q, k = _apply_repo_rope(q, k, z, inv_freq, attn_scaling)
+                    q, k = _apply_repo_rope(
+                        q, k, repo_coords, inv_freq, attn_scaling
+                    )
         elif _iha_seq_expand:
             # IHA seq-expand + standard integer RoPE:
             # Build interleaved positions  pos_iha[s·P+p] = position_ids[s]·P + p
@@ -3848,7 +4149,8 @@ class NeoLLMAttention(nn.Module):
             # Standard path: integer positions pre-computed by NeoLLMModel.
             q, k = apply_rotary_pos_emb(q, k, cos, sin)
 
-        q, k = self._apply_momentum_attention(q, k)
+        if not _repo_grape_kernel_used:
+            q, k = self._apply_momentum_attention(q, k)
         # ── MEA head mixing ───────────────────────────────────────────────────
         # For IHA seq-expand (P>1): K/V are [B, H_kv, S·P, d] — same number of
         # heads as without IHA, just a longer sequence.  Use the standard
@@ -4187,7 +4489,44 @@ class PolyNorm(nn.Module):
     def forward(
         self,
         x: torch.Tensor,
+        dropout_p: float = 0.0,
     ) -> torch.Tensor:
+        if _cce_polynorm is not None:
+            if (
+                self.training
+                and torch.is_grad_enabled()
+                and torch.compiler.is_compiling()
+                and _cce_polynorm_uses_cute is not None
+                and _cce_polynorm_uses_cute(
+                    x,
+                    self.weight,
+                    self.bias,
+                    eps=float(self.eps),
+                    exclusive_logits=getattr(self, "exclusive_logits", None),
+                    dropout_p=dropout_p,
+                )
+            ):
+                return _activation_checkpoint(
+                    _cce_polynorm,
+                    x,
+                    self.weight,
+                    self.bias,
+                    use_reentrant=False,
+                    eps=float(self.eps),
+                    proj_eps=float(getattr(self, "proj_eps", 1e-6)),
+                    exclusive_logits=getattr(self, "exclusive_logits", None),
+                    dropout_p=dropout_p,
+                )
+            return _cce_polynorm(
+                x,
+                self.weight,
+                self.bias,
+                eps=float(self.eps),
+                proj_eps=float(getattr(self, "proj_eps", 1e-6)),
+                exclusive_logits=getattr(self, "exclusive_logits", None),
+                dropout_p=dropout_p,
+            )
+
         # Caché de potencias: x_sq reutilizado en x1 y x2; x_cu = x·x_sq evita pow(3)
         x_sq = x.pow(2)
         x_cu = x * x_sq
@@ -4213,6 +4552,9 @@ class PolyNorm(nn.Module):
             self.weight[0] * x3 + self.weight[1] * x2 + self.weight[2] * x1 + self.bias
         )
 
+        if dropout_p:
+            output = F.dropout(output, p=dropout_p, training=True)
+
         return output
 
 
@@ -4259,10 +4601,11 @@ class NeoLLMMLP(nn.Module):
         gate_out = self.gate_proj(x_fan)
         up_out = self.up_proj(x_fan)
 
-        act_out = self.act_fn(gate_out)
+        dropout_p = float(self.dropout.p) if self.training else 0.0
+        act_out = self.act_fn(gate_out, dropout_p=dropout_p)
         act_x_up = act_out * up_out
 
-        result = self.down_proj(self.dropout(act_x_up))
+        result = self.down_proj(act_x_up)
 
         return result
 
@@ -4997,6 +5340,9 @@ class SpellingBeeEmbedding(nn.Module):
         self,
         token_ids: torch.Tensor,  # [B, S] or [N]
         token_embeds: torch.Tensor,  # [B, S, d] or [N, d]
+        *,
+        meap_mask_embedding: Optional[torch.Tensor] = None,
+        meap_mask_token_id: Optional[int] = None,
     ) -> torch.Tensor:
         """
         Args:
@@ -5021,7 +5367,22 @@ class SpellingBeeEmbedding(nn.Module):
         ].mean(-2)  # [..., d], fixed 16-position mean
 
         # ── Step 3: mean with token embeddings ──────────────────────────────
-        return (token_embeds + e_chars) * 0.5
+        output = (token_embeds + e_chars) * 0.5
+        has_meap = (
+            meap_mask_embedding is not None or meap_mask_token_id is not None
+        )
+        if not has_meap:
+            return output
+        if meap_mask_embedding is None or meap_mask_token_id is None:
+            raise ValueError(
+                "meap_mask_embedding and meap_mask_token_id must be provided "
+                "together"
+            )
+        return torch.where(
+            token_ids.eq(int(meap_mask_token_id)).unsqueeze(-1),
+            meap_mask_embedding.to(device=output.device, dtype=output.dtype),
+            output,
+        )
 
     # ── Inference utility ─────────────────────────────────────────────────────
 
@@ -5117,11 +5478,33 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
 
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
-        """Load legacy checkpoints without exposing an uninitialized knot grid."""
-        model = super().from_pretrained(*args, **kwargs)
-        for module in model.modules():
+        """Initialize only a legacy checkpoint that lacks ``knot_grid``.
+
+        Current checkpoints persist this buffer and must retain its exact saved
+        values.  The previous compatibility path regenerated the grid after
+        every successful load, silently overwriting a valid checkpoint (and
+        BF16 direct-linspace rounding can differ from FP32-then-BF16 rounding).
+        """
+        return_loading_info = bool(kwargs.pop("output_loading_info", False))
+        requested_dtype = kwargs.get("dtype", kwargs.get("torch_dtype"))
+        model, loading_info = super().from_pretrained(
+            *args,
+            output_loading_info=True,
+            **kwargs,
+        )
+        missing_keys = set(loading_info.get("missing_keys", ()))
+        for module_name, module in model.named_modules():
             if not isinstance(module, LeviathanGenerator):
                 continue
+            checkpoint_key = (
+                f"{module_name}.knot_grid" if module_name else "knot_grid"
+            )
+            needs_legacy_initialization = (
+                module.knot_grid.device.type == "meta"
+                or checkpoint_key in missing_keys
+            )
+            if not needs_legacy_initialization:
+                continue
             if module.knot_grid.device.type == "meta":
                 module._buffers["knot_grid"] = torch.linspace(
                     0.0,
@@ -5132,15 +5515,45 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
                 )
             else:
                 with torch.no_grad():
+                    # Match training construction: create the fixed grid in
+                    # FP32 first, then let the explicit cast perform rounding.
                     module.knot_grid.copy_(
                         torch.linspace(
                             0.0,
                             1.0,
                             module.num_knots,
                             device=module.knot_grid.device,
-                            dtype=module.knot_grid.dtype,
-                        )
+                            dtype=torch.float32,
+                        ).to(dtype=module.knot_grid.dtype)
+                    )
+
+        # These SiameseNorm tensors are constructed explicitly in FP32 to make
+        # direct ``NeoLLMForCausalLM(config)`` initialization stable. Training
+        # then casts them with ``model.to(BF16)``. Transformers' low-memory
+        # loader honors the explicit constructor dtype instead of ``dtype=``
+        # when copying a checkpoint, so without this normalization a BF16 Hub
+        # checkpoint silently resumes those 12 trainable vectors in FP32.
+        if (
+            isinstance(requested_dtype, torch.dtype)
+            and requested_dtype.is_floating_point
+        ):
+            for module in model.modules():
+                siamese_scale = getattr(module, "siamese_attn_x_scale", None)
+                if siamese_scale is not None:
+                    siamese_scale.data = siamese_scale.data.to(
+                        dtype=requested_dtype
                     )
+                stream_scale = getattr(
+                    module,
+                    "_siamese_stream_scale_value",
+                    None,
+                )
+                if stream_scale is not None:
+                    module._buffers["_siamese_stream_scale_value"] = (
+                        stream_scale.to(dtype=requested_dtype)
+                    )
+        if return_loading_info:
+            return model, loading_info
         return model
 
     def _init_weights(self, module):
@@ -5346,6 +5759,26 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         else:
             self.spelling_bee = None
 
+        # MEAP owns one independent final embedding vector.  It deliberately
+        # lives in the model instead of the optional CCE package so checkpoints
+        # remain loadable for inference when that package is absent.  The last
+        # active embedding producer owns the MEAP selection: Leviathan fuses it
+        # into its output GEMM, while Spelling Bee owns it when byte features
+        # follow. This prevents masked positions from updating dense token rows,
+        # Leviathan codebooks, or Spelling Bee bytes.
+        meap_mask_token_id = getattr(config, "meap_mask_token_id", None)
+        if meap_mask_token_id is None:
+            self.register_parameter("meap_mask_embedding", None)
+        else:
+            self.meap_mask_embedding = nn.Parameter(
+                torch.empty(config.hidden_size)
+            )
+            nn.init.normal_(
+                self.meap_mask_embedding,
+                mean=0.0,
+                std=config.initializer_range,
+            )
+
         self.layers = nn.ModuleList(
             [
                 NeoLLMDecoderLayer(config, layer_idx)
@@ -5422,6 +5855,7 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         output_hidden_states: Optional[bool] = None,
         output_attentions: Optional[bool] = None,
         output_tweo_activations: Optional[bool] = None,
+        meap_active: bool = False,
         return_dict: Optional[bool] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> Tuple:
@@ -5447,9 +5881,35 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
             raise ValueError("Specify exactly one of input_ids or inputs_embeds")
 
         # ── Embedding stage ────────────────────────────────────────────────
+        meap_mask_embedding = None
+        meap_mask_token_id = None
+        if meap_active:
+            if input_ids is None:
+                raise ValueError("Active MEAP requires `input_ids`.")
+            if self.meap_mask_embedding is None:
+                raise ValueError(
+                    "Active MEAP requires a configured `meap_mask_token_id`."
+                )
+            meap_mask_embedding = self.meap_mask_embedding
+            meap_mask_token_id = int(self.config.meap_mask_token_id)
+
+        meap_consumed = False
         if inputs_embeds is None:
             if self.config.use_token_generator:
-                inputs_embeds = self.token_generator(input_ids)
+                # Without a later Spelling Bee epilogue, Leviathan owns the
+                # final embedding write and selects the MEAP vector inside its
+                # GEMM store. If SBE is active, its epilogue must own the final
+                # selection so byte features cannot alter masked positions.
+                fuse_meap_in_leviathan = meap_active and self.spelling_bee is None
+                if fuse_meap_in_leviathan:
+                    inputs_embeds = self.token_generator(
+                        input_ids,
+                        meap_mask_embedding=meap_mask_embedding,
+                        meap_mask_token_id=meap_mask_token_id,
+                    )
+                else:
+                    inputs_embeds = self.token_generator(input_ids)
+                meap_consumed = fuse_meap_in_leviathan
             else:
                 inputs_embeds = self.embed_tokens(input_ids)
 
@@ -5458,7 +5918,29 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         # caller; in that case SBE cannot run (no token_ids available) and is
         # silently skipped — consistent with the standard embedding bypass path.
         if self.spelling_bee is not None and input_ids is not None:
-            inputs_embeds = self.spelling_bee(input_ids, inputs_embeds)
+            inputs_embeds = self.spelling_bee(
+                input_ids,
+                inputs_embeds,
+                meap_mask_embedding=(
+                    meap_mask_embedding if meap_active else None
+                ),
+                meap_mask_token_id=(meap_mask_token_id if meap_active else None),
+            )
+            meap_consumed = meap_active
+
+        # This must be the final operation in the input-embedding pipeline.
+        # The clean/inference path is unchanged and does not depend on CCE.
+        if meap_active and not meap_consumed:
+            meap_positions = input_ids.eq(meap_mask_token_id).unsqueeze(-1)
+            mask_embedding = meap_mask_embedding.to(
+                device=inputs_embeds.device,
+                dtype=inputs_embeds.dtype,
+            ).view(1, 1, -1)
+            inputs_embeds = torch.where(
+                meap_positions,
+                mask_embedding,
+                inputs_embeds,
+            )
 
         if position_ids is None:
             if attention_mask is not None:
@@ -5789,6 +6271,7 @@ def compute_cce_loss(
     cce_impl="cce_kahan_full_c",
     use_mile_loss=False,
     mile_loss_gamma=1.0,
+    mile_group_mask=None,
     use_mu_loss=False,
     mu_loss_lambda=1e-4,
     return_loss_metrics=False,
@@ -5817,6 +6300,9 @@ def compute_cce_loss(
             mile_enabled=True,
             mile_gamma=float(mile_loss_gamma),
         )
+        if mile_group_mask is not None:
+            _require_extended_cce_options("mile_group_mask")
+            extension_kwargs["mile_group_mask"] = mile_group_mask
     if use_mu_loss:
         _require_extended_cce_options("mu_loss_enabled", "mu_loss_lambda")
         extension_kwargs.update(
@@ -7502,19 +7988,9 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             _require_extended_cce_options("mile_enabled", "mile_gamma")
         if bool(getattr(config, "use_mu_loss", False)):
             _require_extended_cce_options("mu_loss_enabled", "mu_loss_lambda")
-        if bool(getattr(config, "use_meap", False)) and meap_mask_inputs is None:
-            raise ImportError(
-                "`use_meap=True` requires the extended cut-cross-entropy package. "
-                f"Install it with `pip install --force-reinstall --no-deps "
-                f"\"{_EXTENDED_CCE_INSTALL}\"`, or set `use_meap=False`."
-            )
-        if bool(getattr(config, "use_meap", False)) and (
-            "return_metrics" not in inspect.signature(meap_mask_inputs).parameters
-        ):
-            raise ImportError(
-                "The installed cut-cross-entropy package does not provide kernel-side "
-                "MEAP metrics. Reinstall the extended backend or set `use_meap=False`."
-            )
+        # The optional MEAP backend is checked only when training actually
+        # applies corruption.  Clean training, evaluation, and inference remain
+        # usable without cut-cross-entropy even when loading a MEAP checkpoint.
         self._last_ntp_loss = None
         self._last_ntp_ce_unweighted = None
         self._last_mile_reweighting_delta = None
@@ -7609,6 +8085,7 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             and self.training
             and labels is not None
         )
+        meap_selected_mask = None
         if meap_enabled:
             if input_ids is None or inputs_embeds is not None:
                 raise ValueError(
@@ -7617,8 +8094,13 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 )
             if meap_mask_inputs is None:
                 raise ImportError(
-                    "MEAP is enabled but is unavailable in the installed "
-                    "cut-cross-entropy package."
+                    "Active MEAP training requires the extended cut-cross-entropy "
+                    "package. Evaluation and inference do not require it."
+                )
+            if "return_metrics" not in inspect.signature(meap_mask_inputs).parameters:
+                raise ImportError(
+                    "The installed cut-cross-entropy package lacks MEAP metrics. "
+                    "Install the extended backend before MEAP training."
                 )
             mask_token_id = getattr(self.config, "meap_mask_token_id", None)
             if mask_token_id is None:
@@ -7642,7 +8124,8 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             )
             exclude_last = bool(getattr(self.config, "meap_exclude_last", True))
             mask_ratio = float(getattr(self.config, "meap_mask_ratio", 0.15))
-            input_ids, meap_metrics = meap_mask_inputs(
+            return_meap_mask = bool(getattr(self.config, "use_mile_loss", False))
+            meap_result = meap_mask_inputs(
                 input_ids,
                 int(mask_token_id),
                 enabled=True,
@@ -7650,12 +8133,16 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 eligible_mask=eligible_mask,
                 seed=effective_meap_seed,
                 exclude_last=exclude_last,
-                return_mask=False,
+                return_mask=return_meap_mask,
                 return_metrics=True,
                 implementation=str(
                     getattr(self.config, "meap_implementation", "triton")
                 ),
             )
+            if return_meap_mask:
+                input_ids, meap_selected_mask, meap_metrics = meap_result
+            else:
+                input_ids, meap_metrics = meap_result
             eligible_count = meap_metrics[0]
             masked_count = meap_metrics[1]
             self._last_meap_mask_fraction = (
@@ -7706,6 +8193,7 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             inputs_embeds=inputs_embeds,
             output_hidden_states=model_output_hidden_states,
             output_tweo_activations=tweo_enabled,
+            meap_active=meap_enabled,
             return_dict=return_dict,
             **kwargs,
         )
@@ -7810,6 +8298,12 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                     mile_loss_gamma=float(
                         getattr(self.config, "mile_loss_gamma", 1.0)
                     ),
+                    mile_group_mask=(
+                        meap_selected_mask
+                        if meap_enabled
+                        and bool(getattr(self.config, "use_mile_loss", False))
+                        else None
+                    ),
                     use_mu_loss=bool(getattr(self.config, "use_mu_loss", False)),
                     mu_loss_lambda=float(
                         getattr(self.config, "mu_loss_lambda", 1e-4)