Update NeoLLM model and configuration sources

configuration_neollm.py

@@ -1,5 +1,6 @@
 # ==================== configuration_neollm.py ====================
 import math
+from numbers import Integral
 
 from transformers.configuration_utils import PretrainedConfig
 from transformers.modeling_rope_utils import rope_config_validation
@@ -53,7 +54,8 @@ class NeoLLMConfig(PretrainedConfig):
             Cut Cross Entropy through ``linear_cross_entropy(..., shift=1)``.
             ``"liger"`` uses Liger Kernel's fused linear cross entropy with an
             explicit shift equivalent to ``shift=1``. The two backends are
-            mutually exclusive at runtime.
+            mutually exclusive at runtime. Copy-Paste, when enabled, computes
+            its mixed objective instead of either backend's NTP loss.
         use_liger_kernel (:obj:`bool`, *optional*, defaults to ``False``):
             Convenience flag for ``ntp_loss_backend="liger"``. When enabled,
             the model does not call the CCE loss path for NTP.
@@ -61,15 +63,24 @@ class NeoLLMConfig(PretrainedConfig):
             CCE implementation used when ``ntp_loss_backend="cce"``. The default
             keeps the conservative Kahan/full-classifier-gradient path used for
             pretraining from scratch.
-        use_mile_loss (:obj:`bool`, *optional*, defaults to ``False``):
-            Enable MiLe Loss (Su et al., 2026, arXiv:2310.19531) in the extended
-            CCE backend. MiLe reweights token losses with detached,
-            mean-normalized predictive-entropy weights.
+        use_mile_loss (:obj:`bool`, *optional*, defaults to ``True``):
+            Enable MiLe Loss (Su et al., 2026, arXiv:2310.19531). MiLe
+            reweights token losses with detached, mean-normalized predictive-
+            entropy weights. Copy-Paste derives those weights from the mixed
+            distribution using CCE's sparse statistics on CUDA.
+        use_copy_paste (:obj:`bool`, *optional*, defaults to ``False``):
+            Add the Copy-Paste LM head (Godey and Artzi, arXiv:2610.19247).
+            Its learned copy gate and sink mix vocabulary predictions with a
+            causal distribution over observed input tokens. CUDA training uses
+            CCE's sparse Copy-Paste path; CPU uses a bounded Torch reference.
+            MiLe uses mixture entropy; mu-loss remains a regularizer on the
+            ordinary vocabulary matrix. Inference still returns dense logits.
         mile_loss_gamma (:obj:`float`, *optional*, defaults to ``1.0``):
             Non-negative exponent applied to MiLe's ``1 + entropy`` weight.
-        use_mu_loss (:obj:`bool`, *optional*, defaults to ``False``):
+        use_mu_loss (:obj:`bool`, *optional*, defaults to ``True``):
             Enable the output-embedding centering regularizer from Stollenwerk
-            et al. (2026, arXiv:2601.02031) in the extended CCE backend.
+            et al. (2026, arXiv:2601.02031). CCE uses its extension; the
+            Copy-Paste path computes the regularizer directly in Torch.
         mu_loss_lambda (:obj:`float`, *optional*, defaults to ``1e-4``):
             Non-negative coefficient in ``L_mu = lambda * ||mean(C)||^2``.
         use_meap (:obj:`bool`, *optional*, defaults to ``False``):
@@ -90,7 +101,19 @@ class NeoLLMConfig(PretrainedConfig):
             MEAP implementation, either ``"triton"`` or the reference
             ``"torch"`` path.
         meap_exclude_last (:obj:`bool`, *optional*, defaults to ``True``):
-            Exclude the final eligible input because it has no shifted target.
+            Exclude the last valid input, which has no shifted target, before
+            structural eligibility filters are applied.
+        meap_protect_structural_tokens (:obj:`bool`, *optional*, defaults to ``True``):
+            Preserve configured PAD, BOS, EOS and the dedicated MEAP mask ID
+            during corruption. EOS may be a single ID or a list of IDs.
+        meap_exclude_mask_from_copy (:obj:`bool`, *optional*, defaults to ``True``):
+            During MEAP training, exclude artificial mask positions from the
+            Copy-Paste source distribution. Clean targets remain unchanged.
+        copy_paste_lexical_keys (:obj:`bool`, *optional*, defaults to ``False``):
+            Add a low-rank projection of each observed token embedding to its
+            patch-context key. This has an effect only when Copy-Paste and
+            patch training with K>1 are both active. The rank is derived from
+            the hidden size, with no batch- or sequence-specific setting.
         liger_loss_accum_dtype (:obj:`str`, *optional*, defaults to ``"float32"``):
             Accumulator dtype passed to Liger FLCE for weight/bias gradient
             accumulation. ``"float32"`` is the stability-first setting for BF16
@@ -108,6 +131,15 @@ class NeoLLMConfig(PretrainedConfig):
             Whether to add a bias term to Q, K, V, and output projections.
         attention_dropout (:obj:`float`, *optional*, defaults to 0.1):
             Dropout probability applied to attention weights during training.
+        use_vha_attention (:obj:`bool`, *optional*, defaults to ``False``):
+            Factorize query projections into virtual heads and apply a low-rank
+            residual head mix (Ding et al., arXiv:2610.17959). For Hq query
+            heads and Hkv KV heads, the G=Hq/Hkv physical query/gate heads are
+            shared across KV groups and expanded to Hq virtual heads; the
+            postmix rank is max(1, G//2). The configured KV count is unchanged.
+            Hkv=2 gives the paper's halving of physical query heads; other KV
+            counts use the same group-specific factorization. Enabling this
+            option changes checkpoint parameter shapes.
         use_hola_memory (:obj:`bool`, *optional*, defaults to ``False``):
             Enable the FA2-compatible HOLA-derived episodic memory read path for
             the IHA seq-expand branch. This is not a second attention branch: it
@@ -190,6 +222,19 @@ class NeoLLMConfig(PretrainedConfig):
             Epsilon for the denominator of the XSA projection to prevent
             division by zero: ``‖v_ref‖² + xsa_eps``. Ignored when
             ``use_xsa=False``.
+        use_msar (:obj:`bool`, *optional*, defaults to ``False``):
+            Enable Block Multi-Stream Attention Residuals (arXiv:2610.13630).
+            It is exclusive with the other residual topologies, GPAS, FAN
+            residual and DAGFormer. Single-stream depth summaries are retained;
+            causal views and independent depth reads are reconstructed per branch.
+        msar_block_branches (:obj:`int` or ``None``, optional):
+            Number of attention/MLP branches per block. ``None`` derives a
+            size from model depth aiming for roughly eight blocks.
+        msar_kernel_sizes (:obj:`tuple`, optional):
+            Causal depthwise widths; the paper uses ``(4, 8, 12)``.
+        msar_backend (:obj:`str`, optional):
+            ``auto`` uses Triton on CUDA and Torch on CPU; ``torch`` selects
+            the fully differentiable reference for validation.
         use_stack_memory (:obj:`bool`, *optional*, defaults to ``False``):
             Enable the StackMemory / STACKTRANS hidden-state stack between
             decoder layers (Zhang et al., NeurIPS 2025). The module keeps the
@@ -675,6 +720,7 @@ class NeoLLMConfig(PretrainedConfig):
         ntp_loss_backend="cce",
         use_liger_kernel=False,
         cce_loss_impl="cce_kahan_full_c",
+        use_copy_paste=False,
         # ── Optional extended-CCE objectives and MEAP input corruption ────
         use_mile_loss=True,
         mile_loss_gamma=1.0,
@@ -689,12 +735,16 @@ class NeoLLMConfig(PretrainedConfig):
         meap_seed=0,
         meap_implementation="triton",
         meap_exclude_last=True,
+        meap_protect_structural_tokens=True,
+        meap_exclude_mask_from_copy=True,
+        copy_paste_lexical_keys=False,
         liger_loss_accum_dtype="float32",
         rope_theta=10000.0,
         rope_scaling=None,
         partial_rotary_factor=0.25,
         attention_bias=False,
         attention_dropout=0.1,
+        use_vha_attention=False,
         # ── HOLA-derived episodic memory for FlashAttention-2 ─────────────
         use_hola_memory=False,
         hola_memory_size=64,
@@ -721,6 +771,11 @@ class NeoLLMConfig(PretrainedConfig):
         # ── Attention Residuals (Kimi Team, 2026) ─────────────────────────
         use_attn_res=False,
         attn_res_num_blocks=4,
+        # MSAR: independent residual topology, not combined with SiameseNorm.
+        use_msar=False,
+        msar_block_branches=None,  # auto: approximately eight depth blocks
+        msar_kernel_sizes=(4, 8, 12),
+        msar_backend="auto",  # Triton on CUDA, Torch on CPU; "torch" is a reference
         # ── StackMemory / STACKTRANS (Zhang et al., NeurIPS 2025) ───────
         use_stack_memory=False,
         stack_d_model=24,  # H * d_s = 4 * 6
@@ -868,11 +923,12 @@ class NeoLLMConfig(PretrainedConfig):
             if use_meap:
                 reserved_for_model = {
                     token_id
-                    for token_id in (
+                    for value in (
                         kwargs.get("pad_token_id"),
                         kwargs.get("bos_token_id"),
                         kwargs.get("eos_token_id"),
                     )
+                    for token_id in (value if isinstance(value, (list, tuple)) else (value,))
                     if token_id is not None
                 }
                 if meap_mask_token_id in reserved_for_model:
@@ -1055,6 +1111,7 @@ class NeoLLMConfig(PretrainedConfig):
             "use_attn_res": bool(use_attn_res),
             "use_stack_memory": bool(use_stack_memory),
             "use_siamesenorm": bool(use_siamesenorm),
+            "use_msar": bool(use_msar),
         }
         _active_residual_flow = [
             name for name, enabled in _residual_flow_flags.items() if enabled
@@ -1064,16 +1121,34 @@ class NeoLLMConfig(PretrainedConfig):
             and bool(use_siamesenorm)
             and not bool(use_laurel)
             and not bool(use_attn_res)
+            and not bool(use_msar)
         )
         if len(_active_residual_flow) > 1 and not allow_stack_siamese:
             raise ValueError(
                 "Enable at most one residual-flow mechanism among "
                 "`use_laurel`, `use_attn_res`, `use_stack_memory`, "
-                "and `use_siamesenorm`, except the explicit "
+                "`use_siamesenorm`, and `use_msar`, except the explicit "
                 "`use_stack_memory` + `use_siamesenorm` pair. "
                 f"Active flags: {_active_residual_flow}."
             )
 
+        # MSAR owns the residual path: reject other residual contributors too.
+        if use_msar and (use_gpas or use_fan_residual or use_dagformer):
+            raise ValueError(
+                "`use_msar=True` cannot be combined with `use_gpas`, "
+                "`use_fan_residual`, or `use_dagformer`."
+            )
+        if msar_block_branches is None:
+            # Approximately eight blocks, without tuning for batch or sequence.
+            msar_block_branches = 2 * max(1, math.ceil(num_hidden_layers / 8))
+        if isinstance(msar_block_branches, bool) or not isinstance(msar_block_branches, int) or msar_block_branches < 1:
+            raise ValueError("`msar_block_branches` must be a positive integer.")
+        if (not isinstance(msar_kernel_sizes, (tuple, list)) or not msar_kernel_sizes
+                or any(isinstance(k, bool) or not isinstance(k, int) or k < 1 for k in msar_kernel_sizes)):
+            raise ValueError("`msar_kernel_sizes` must contain positive integer widths.")
+        if msar_backend not in ("auto", "torch"):
+            raise ValueError("`msar_backend` must be 'auto' or 'torch'.")
+
         # ── LAuReL: sub-flag consistency ──────────────────────────────────
         if use_laurel and not use_laurel_rw and not use_laurel_lr:
             raise ValueError(
@@ -1225,6 +1300,13 @@ class NeoLLMConfig(PretrainedConfig):
         self.hidden_size = hidden_size
         self.intermediate_size = intermediate_size
         self.num_hidden_layers = num_hidden_layers
+        if use_vha_attention and (
+            isinstance(num_attention_heads, bool) or not isinstance(num_attention_heads, Integral)
+            or isinstance(num_key_value_heads, bool) or not isinstance(num_key_value_heads, Integral)
+            or num_attention_heads <= 0 or num_key_value_heads <= 0
+            or num_attention_heads % num_key_value_heads != 0
+        ):
+            raise ValueError("VHA requires positive query/KV head counts, with KV heads dividing query heads.")
         self.num_attention_heads = num_attention_heads
         self.num_key_value_heads = num_key_value_heads
         self.hidden_act = hidden_act
@@ -1235,6 +1317,7 @@ class NeoLLMConfig(PretrainedConfig):
         self.ntp_loss_backend = ntp_loss_backend
         self.use_liger_kernel = bool(use_liger_kernel)
         self.cce_loss_impl = str(cce_loss_impl).strip()
+        self.use_copy_paste = bool(use_copy_paste)
         self.use_mile_loss = bool(use_mile_loss)
         self.mile_loss_gamma = float(mile_loss_gamma)
         self.use_mu_loss = bool(use_mu_loss)
@@ -1247,6 +1330,9 @@ class NeoLLMConfig(PretrainedConfig):
         self.meap_seed = int(meap_seed)
         self.meap_implementation = meap_implementation
         self.meap_exclude_last = bool(meap_exclude_last)
+        self.meap_protect_structural_tokens = bool(meap_protect_structural_tokens)
+        self.meap_exclude_mask_from_copy = bool(meap_exclude_mask_from_copy)
+        self.copy_paste_lexical_keys = bool(copy_paste_lexical_keys)
         self.liger_loss_accum_dtype = liger_loss_accum_dtype
 
         # ── Positional encoding ───────────────────────────────────────────
@@ -1257,6 +1343,7 @@ class NeoLLMConfig(PretrainedConfig):
         # ── Attention ─────────────────────────────────────────────────────
         self.attention_bias = attention_bias
         self.attention_dropout = attention_dropout
+        self.use_vha_attention = bool(use_vha_attention)
         # ── HOLA-derived episodic memory for FlashAttention-2 ─────────────
         self.use_hola_memory = bool(use_hola_memory)
         self.hola_memory_size = int(hola_memory_size)
@@ -1294,6 +1381,10 @@ class NeoLLMConfig(PretrainedConfig):
         #   block_size = num_hidden_layers // 4 = 3 layers per block.
         #   Memory cost: O(num_blocks × batch × seq × hidden) instead of
         #   O(num_layers × batch × seq × hidden).
+        self.use_msar = bool(use_msar)
+        self.msar_block_branches = msar_block_branches
+        self.msar_kernel_sizes = tuple(msar_kernel_sizes)
+        self.msar_backend = msar_backend
         self.use_attn_res = use_attn_res
         self.attn_res_num_blocks = attn_res_num_blocks
 

modeling_neollm.py

@@ -103,6 +103,15 @@ except ImportError:
     linear_cross_entropy = None
     _CCE_AVAILABLE = False
 
+try:
+    from cut_cross_entropy import copy_paste_loss as _cce_copy_paste_loss
+except (ImportError, AttributeError):
+    _cce_copy_paste_loss = None
+_CCE_COPY_PASTE_HAS_METRICS = (
+    _cce_copy_paste_loss is not None
+    and "return_copy_metrics" in inspect.signature(_cce_copy_paste_loss).parameters
+)
+
 try:
     from cut_cross_entropy import meap_mask_inputs  # type: ignore[import-not-found]
 except ImportError:
@@ -4155,6 +4164,14 @@ class NeoLLMAttention(nn.Module):
         self.use_learnable_multipliers = getattr(
             config, "use_learnable_multipliers", True
         )
+        self.use_vha_attention = bool(getattr(config, "use_vha_attention", False))
+        # A virtual group shares physical Q and gate projections. The final Q
+        # heads still match the configured GQA geometry expected by every
+        # attention backend and optional head-space operation.
+        self.vha_physical_heads = (
+            config.num_attention_heads // config.num_key_value_heads
+            if self.use_vha_attention else config.num_attention_heads
+        )
 
         self.fan_layer = FANLayer(
             hidden_size=config.hidden_size,
@@ -4166,7 +4183,7 @@ class NeoLLMAttention(nn.Module):
 
         self.q_proj = LinearWithMultipliers(
             fan_output_dim,
-            config.num_attention_heads * self.head_dim * 2,
+            self.vha_physical_heads * self.head_dim * 2,
             bias=config.attention_bias,
             use_row_multiplier=True,
             use_column_multiplier=False,
@@ -4211,6 +4228,18 @@ class NeoLLMAttention(nn.Module):
         self.q_norm = _make_norm(self.head_dim, eps=config.rms_norm_eps)
         self.k_norm = _make_norm(self.head_dim, eps=config.rms_norm_eps)
 
+        if self.use_vha_attention:
+            # Premix is group-specific and close to identity; postmix starts as
+            # the identity because its second low-rank factor is exactly zero.
+            rank = max(1, self.vha_physical_heads // 2)
+            identity = torch.eye(self.head_dim).expand(
+                config.num_key_value_heads, -1, -1).clone()
+            self.vha_premix = nn.Parameter(
+                identity + torch.randn_like(identity) * (0.1 * config.initializer_range))
+            self.vha_postmix_a = nn.Parameter(
+                torch.randn(config.num_attention_heads, rank) * config.initializer_range)
+            self.vha_postmix_b = nn.Parameter(torch.zeros(config.num_attention_heads, rank))
+
         if self.use_mea_attention:
             self.mea_key_mix = nn.Parameter(
                 torch.eye(self.num_mea_component_heads, config.num_key_value_heads)
@@ -5474,7 +5503,12 @@ class NeoLLMAttention(nn.Module):
         kv_shape = (*input_shape, self.num_mea_component_heads, self.head_dim)
 
         q_projected = self.q_proj(h_fan_q).view(
-            *input_shape, self.config.num_attention_heads, self.head_dim * 2)
+            *input_shape, self.vha_physical_heads, self.head_dim * 2)
+        if self.use_vha_attention:
+            q_projected = q_projected.unsqueeze(-3).expand(
+                *input_shape, self.config.num_key_value_heads,
+                self.vha_physical_heads, self.head_dim * 2,
+            ).reshape(*input_shape, self.config.num_attention_heads, self.head_dim * 2)
         q_raw, gate = q_projected.chunk(2, dim=-1)
         gate = gate.reshape(*input_shape, -1)
 
@@ -5486,8 +5520,27 @@ class NeoLLMAttention(nn.Module):
                 features = self.lambda_1 * first_layer_fan.unsqueeze(0) + self.lambda_2 * features
             hq = self.config.num_attention_heads
             hk = self.num_mea_component_heads
-            q_delta = _dag_project(features, dag_coefficients[..., :hq, :], self.q_proj, hq,
-                                   q_projected if self.dag_reuse_latest else None)
+            if self.use_vha_attention:
+                # DAG selects a different depth mixture per virtual head. Apply
+                # those coefficients after expanding each physical projection;
+                # this retains DAG's routing instead of silently sharing it.
+                sources = features.shape[0]
+                projected = self.q_proj(features).reshape(
+                    sources, *input_shape, self.vha_physical_heads, self.head_dim * 2)
+                projected = projected.unsqueeze(-3).expand(
+                    sources, *input_shape, self.config.num_key_value_heads,
+                    self.vha_physical_heads, self.head_dim * 2,
+                ).reshape(sources, *input_shape, hq, self.head_dim * 2)
+                coeff = dag_coefficients[..., :hq, :]
+                earlier_coeff = coeff[..., :-1] if self.dag_reuse_latest else coeff
+                q_delta = torch.einsum(
+                    "sbthd,bths->bthd", projected.float(), earlier_coeff.float())
+                if self.dag_reuse_latest:
+                    q_delta = q_delta + coeff[..., -1, None].float() * q_projected.float()
+                q_delta = q_delta.to(features.dtype)
+            else:
+                q_delta = _dag_project(features, dag_coefficients[..., :hq, :], self.q_proj, hq,
+                                       q_projected if self.dag_reuse_latest else None)
             dq, dg = q_delta.chunk(2, dim=-1)
             q_raw = q_raw + dq
             gate = gate + dg.reshape(*input_shape, -1)
@@ -5495,6 +5548,13 @@ class NeoLLMAttention(nn.Module):
                                          k_raw if self.dag_reuse_latest else None)
             v_raw = v_raw + _dag_project(features, dag_coefficients[..., hq+hk:hq+2*hk, :], self.v_proj, hk,
                                          v_raw if self.dag_reuse_latest else None)
+        if self.use_vha_attention:
+            q_raw = torch.einsum(
+                "btgid,gde->btgie",
+                q_raw.reshape(*input_shape, self.config.num_key_value_heads,
+                              self.vha_physical_heads, self.head_dim),
+                self.vha_premix,
+            ).reshape(query_shape)
         q = self.q_norm(q_raw.view(query_shape)).transpose(1, 2)
         k = self.k_norm(k_raw).transpose(1, 2)
         v = v_raw.transpose(1, 2)
@@ -5963,7 +6023,18 @@ class NeoLLMAttention(nn.Module):
 
         # ── Reshape → o_proj → Gated Attention gate → dropout ────────────
         attn_out_flat = attn_out.reshape(*input_shape, -1).contiguous()
-        attn_out_gated = self.o_proj(attn_out_flat * torch.sigmoid(gate))
+        if self.use_vha_attention:
+            # Keep NeoLLM's nonlinear head gate before postmix. Consequently
+            # the linear postmix can be folded into a dense output projection;
+            # folding into a constrained Hadamard parameterization is not assumed.
+            gated_heads = (attn_out_flat * torch.sigmoid(gate)).reshape(
+                *input_shape, self.config.num_attention_heads, self.head_dim)
+            low_rank = torch.einsum("bthd,hr->btrd", gated_heads, self.vha_postmix_a)
+            mixed_heads = gated_heads + torch.einsum(
+                "btrd,hr->bthd", low_rank, self.vha_postmix_b)
+            attn_out_gated = self.o_proj(mixed_heads.reshape(*input_shape, -1))
+        else:
+            attn_out_gated = self.o_proj(attn_out_flat * torch.sigmoid(gate))
         attn_out_gated = self.dropout(attn_out_gated)
 
         return attn_out_gated, attn_weights, current_layer_fan
@@ -6607,6 +6678,360 @@ def _dag_project(features, coefficients, projection, heads, latest=None):
     return _dag_mix(projected,coefficients.reshape(b*t,heads,size)).reshape(b,t,heads,-1)
 
 
+def _msar_views(sources, banks, valid, position, widths):
+    """Expand sources; positions prevent attention across packed boundaries."""
+    b, t, d = sources[0].shape
+    valid = valid.to(torch.bool)
+    expanded = []
+    for source, bank in zip(sources, banks):
+        x = source.float()
+        streams = [x * valid[..., None]]
+        start = 0
+        for width in widths:
+            view = torch.zeros_like(x)
+            for delta in range(width):
+                if delta >= t:
+                    break
+                if delta:
+                    past = F.pad(x[:, :t-delta, :], (0, 0, delta, 0))
+                    past_valid = F.pad(valid[:, :t-delta], (delta, 0))
+                    past_pos = F.pad(position[:, :t-delta], (delta, 0))
+                else:
+                    past, past_valid, past_pos = x, valid, position
+                allowed = valid & past_valid & ((position - past_pos) == delta)
+                view = view + past * allowed[..., None] * bank[start + delta]
+            streams.append(view)
+            start += width
+        expanded.append(torch.stack(streams, dim=2))
+    return torch.stack(expanded, dim=2)  # B,T,N,M,D
+
+
+def _msar_reference_read(sources, banks, query, valid, position, widths, eps):
+    """Eqs. 9--12: each source has its own causal expansion bank."""
+    v = _msar_views(sources, banks, valid, position, widths)
+    q = query.float()[None, None, None, :, :]
+    inv = torch.rsqrt(v.square().mean(dim=-1, keepdim=True) + eps)
+    score = (v * inv * q).sum(-1)
+    probs = torch.softmax(score, dim=2)
+    return (probs[..., None] * v).sum(dim=2).to(sources[0].dtype)
+
+
+if triton is not None:
+    @triton.jit
+    def _msar_read_kernel(SOURCES, BANKS, QUERY, VALID, POSITION, OUT, LSE,
+                    T: tl.constexpr, D: tl.constexpr, NS: tl.constexpr,
+                    M: tl.constexpr,
+                    WIDTHS: tl.constexpr, EPS: tl.constexpr, BD: tl.constexpr):
+        """Online depth softmax; expanded source views stay in registers."""
+        row = tl.program_id(0)
+        stream = tl.program_id(1)
+        batch = row // T
+        token = row % T
+        c = tl.arange(0, BD)
+        mask_c = c < D
+        is_valid = tl.load(VALID + row).to(tl.int1)
+        position = tl.load(POSITION + row)
+        q = tl.load(QUERY + stream * D + c, mask_c, other=0).to(tl.float32)
+        running_max = tl.full((), -float('inf'), tl.float32)
+        denom = tl.full((), 0., tl.float32)
+        accum = tl.full((BD,), 0., tl.float32)
+        for i in tl.static_range(NS):
+            value = tl.full((BD,), 0., tl.float32)
+            if stream == 0:
+                value = tl.load(SOURCES[i] + row * D + c,
+                                mask_c & is_valid, other=0).to(tl.float32)
+            offset = 0
+            for r in tl.static_range(len(WIDTHS)):
+                if stream == r + 1:
+                    for delta in tl.static_range(WIDTHS[r]):
+                        past = row - delta
+                        in_seq = token >= delta
+                        past_valid = tl.load(VALID + past, in_seq, other=0).to(tl.int1)
+                        past_pos = tl.load(POSITION + past, in_seq, other=0)
+                        allowed = in_seq & is_valid & past_valid & ((position - past_pos) == delta)
+                        x = tl.load(SOURCES[i] + past * D + c,
+                                    mask_c & allowed, other=0).to(tl.float32)
+                        w = tl.load(BANKS[i] + (offset + delta) * D + c,
+                                    mask_c, other=0).to(tl.float32)
+                        value += x * w
+                offset += WIDTHS[r]
+            norm = tl.rsqrt(tl.sum(value * value, 0) / D + EPS)
+            score = tl.sum(value * q, 0) * norm
+            new_max = tl.maximum(running_max, score)
+            old_weight = tl.exp(running_max - new_max)
+            new_weight = tl.exp(score - new_max)
+            accum = accum * old_weight + value * new_weight
+            denom = denom * old_weight + new_weight
+            running_max = new_max
+        tl.store(OUT + row * M * D + stream * D + c, accum / denom, mask_c)
+        tl.store(LSE + row * M + stream, running_max + tl.log(denom))
+
+
+    @triton.jit
+    def _msar_vjp_value(SOURCE, BANK, QUERY, VALID, POSITION, LSE, Z, G, DV, QPART,
+                   ROWS: tl.constexpr, T: tl.constexpr, D: tl.constexpr,
+                   M: tl.constexpr,
+                   WIDTHS: tl.constexpr, EPS: tl.constexpr, R: tl.constexpr,
+                   BD: tl.constexpr, GROUPS: tl.constexpr):
+        """Per-source VJP through softmax and RMS; reduce query partials by rows."""
+        group = tl.program_id(0)
+        stream = tl.program_id(1)
+        row = group * R + tl.arange(0, R)
+        c = tl.arange(0, BD)
+        row_ok = row < ROWS
+        tok = row % T
+        valid = tl.load(VALID + row, row_ok, other=0).to(tl.int1)
+        pos = tl.load(POSITION + row, row_ok, other=0)
+        q = tl.load(QUERY + stream*D+c, c < D, other=0).to(tl.float32)
+        v = tl.full((R, BD), 0., tl.float32)
+        if stream == 0:
+            v = tl.load(SOURCE + row[:,None]*D+c[None,:],
+                        row_ok[:,None] & valid[:,None] & (c[None,:] < D), other=0).to(tl.float32)
+        offset = 0
+        for r in tl.static_range(len(WIDTHS)):
+            if stream == r+1:
+                for delta in tl.static_range(WIDTHS[r]):
+                    past = row - delta
+                    in_seq = row_ok & (tok >= delta)
+                    pv = tl.load(VALID + past, in_seq, other=0).to(tl.int1)
+                    pp = tl.load(POSITION + past, in_seq, other=0)
+                    allowed = in_seq & valid & pv & ((pos-pp) == delta)
+                    x = tl.load(SOURCE + past[:,None]*D+c[None,:],
+                                allowed[:,None] & (c[None,:]<D), other=0).to(tl.float32)
+                    w = tl.load(BANK+(offset+delta)*D+c,
+                                c<D, other=0).to(tl.float32)
+                    v += x*w[None,:]
+            offset += WIDTHS[r]
+        inv = tl.rsqrt(tl.sum(v*v, 1)/D + EPS)
+        qv = tl.sum(v*q[None,:], 1)
+        lse = tl.load(LSE+row*M+stream, row_ok, other=0)
+        p = tl.exp(qv*inv-lse)
+        z = tl.load(Z+row[:,None]*M*D+stream*D+c[None,:],
+                    row_ok[:,None] & (c[None,:]<D), other=0).to(tl.float32)
+        g = tl.load(G+row[:,None]*M*D+stream*D+c[None,:],
+                    row_ok[:,None] & (c[None,:]<D), other=0).to(tl.float32)
+        ds = p*tl.sum(g*(v-z), 1)
+        dv = p[:,None]*g + ds[:,None]*(q[None,:]*inv[:,None] -
+            v*(qv*inv*inv*inv/D)[:,None])
+        tl.store(DV+stream*ROWS*D+row[:,None]*D+c[None,:], dv,
+                 row_ok[:,None] & (c[None,:]<D))
+        qp = tl.sum(ds[:,None]*v*inv[:,None], 0)
+        tl.store(QPART+stream*GROUPS*D+group*D+c, qp, c<D)
+
+
+    @triton.jit
+    def _msar_vjp_source(DV, BANK, VALID, POSITION, DX, ROWS: tl.constexpr,
+                    T: tl.constexpr, D: tl.constexpr,
+                    WIDTHS: tl.constexpr, BLOCK: tl.constexpr):
+        """Transpose the causal depthwise convolutions without atomics."""
+        j = tl.program_id(0)*BLOCK+tl.arange(0,BLOCK)
+        row = j//D
+        c = j%D
+        ok = row<ROWS
+        token = row%T
+        valid = tl.load(VALID+row, ok, other=0).to(tl.int1)
+        pos = tl.load(POSITION+row, ok, other=0)
+        dx = tl.load(DV+j, ok & valid, other=0)
+        offset = 0
+        for r in tl.static_range(len(WIDTHS)):
+            for delta in tl.static_range(WIDTHS[r]):
+                future = row+delta
+                in_seq = ok & ((token+delta)<T)
+                fv = tl.load(VALID+future, in_seq, other=0).to(tl.int1)
+                fp = tl.load(POSITION+future, in_seq, other=0)
+                allowed = in_seq & valid & fv & ((fp-pos)==delta)
+                v = tl.load(DV+(r+1)*ROWS*D+future*D+c, allowed, other=0)
+                w = tl.load(BANK+(offset+delta)*D+c, ok, other=0).to(tl.float32)
+                dx += v*w
+            offset += WIDTHS[r]
+        tl.store(DX+j, dx, ok)
+
+
+    @triton.jit
+    def _msar_vjp_bank(SOURCE, DV, VALID, POSITION, WPART, ROWS: tl.constexpr,
+                  T: tl.constexpr, D: tl.constexpr, WIDTHS: tl.constexpr,
+                  R: tl.constexpr, C: tl.constexpr, GROUPS: tl.constexpr):
+        """Reduce source x view-gradient products into convolution bank partials."""
+        group = tl.program_id(0)
+        c = tl.program_id(1)*C + tl.arange(0,C)
+        row = group*R + tl.arange(0,R)
+        ok = row<ROWS
+        tok = row%T
+        valid = tl.load(VALID+row, ok, other=0).to(tl.int1)
+        pos = tl.load(POSITION+row, ok, other=0)
+        offset = 0
+        for r in tl.static_range(len(WIDTHS)):
+            grad = tl.load(DV+(r+1)*ROWS*D+row[:,None]*D+c[None,:],
+                           ok[:,None] & (c[None,:]<D), other=0).to(tl.float32)
+            for delta in tl.static_range(WIDTHS[r]):
+                past = row-delta
+                in_seq = ok & (tok>=delta)
+                pv = tl.load(VALID+past, in_seq, other=0).to(tl.int1)
+                pp = tl.load(POSITION+past, in_seq, other=0)
+                allowed = in_seq & valid & pv & ((pos-pp)==delta)
+                x = tl.load(SOURCE+past[:,None]*D+c[None,:],
+                            allowed[:,None] & (c[None,:]<D), other=0).to(tl.float32)
+                part = tl.sum(grad*x, 0)
+                tl.store(WPART+(offset+delta)*GROUPS*D+group*D+c, part, c<D)
+            offset += WIDTHS[r]
+
+
+@torch.library.custom_op("neollm_msar::read", mutates_args=())
+def _msar_read_op(sources: list[torch.Tensor], banks: list[torch.Tensor],
+               query: torch.Tensor, valid: torch.Tensor, position: torch.Tensor,
+               widths: list[int], eps: float) -> tuple[torch.Tensor, torch.Tensor]:
+    b, t, d = sources[0].shape
+    sources = [x.contiguous() for x in sources]
+    banks = [x.contiguous() for x in banks]
+    query = query.contiguous()
+    valid = valid.contiguous()
+    position = position.contiguous()
+    result = sources[0].new_empty((b, t, len(widths)+1, d))
+    lse = torch.empty((b,t,len(widths)+1),device=result.device,dtype=torch.float32)
+    _msar_read_kernel[(b*t, len(widths)+1)](tuple(sources), tuple(banks), query, valid, position,
+        result, lse, t, d, len(sources), len(widths)+1,
+        tuple(widths), eps, triton.next_power_of_2(d), num_warps=4)
+    return result, lse
+
+
+@_msar_read_op.register_fake
+def _msar_read_fake(sources, banks, query, valid, position, widths, eps):
+    b, t, d = sources[0].shape
+    return (sources[0].new_empty((b, t, len(widths)+1, d)),
+            torch.empty((b,t,len(widths)+1),device=sources[0].device,dtype=torch.float32))
+
+
+@torch.library.custom_op("neollm_msar::backward", mutates_args=())
+def _msar_backward_op(sources: list[torch.Tensor], banks: list[torch.Tensor],
+                  query: torch.Tensor, valid: torch.Tensor, position: torch.Tensor,
+                  widths: list[int], eps: float,
+                  grad: torch.Tensor, lse: torch.Tensor) -> tuple[list[torch.Tensor], list[torch.Tensor], torch.Tensor]:
+    """Fused VJP with exact FP32 stream recomputation and one-source workspace."""
+    b, t, d = sources[0].shape
+    n, m, ksum = b*t, len(widths)+1, sum(widths)
+    sources = [x.contiguous() for x in sources]
+    banks, query = [x.contiguous() for x in banks], query.contiguous()
+    valid, position, grad = valid.contiguous(), position.contiguous(), grad.contiguous()
+    # Recompute FP32 retrieved streams instead of saving them for every layer.
+    z = torch.empty((b,t,m,d),device=grad.device,dtype=torch.float32)
+    lse_recomputed = torch.empty_like(lse)
+    _msar_read_kernel[(n,m)](tuple(sources), tuple(banks), query, valid, position, z, lse_recomputed,
+        t,d,len(sources),m,tuple(widths),eps,
+        triton.next_power_of_2(d),num_warps=4)
+    # Only one source's view gradients are materialized at a time.
+    dv = torch.empty((m,n,d),device=grad.device,dtype=torch.float32)
+    qrows = max(1, min(8, 8192 // triton.next_power_of_2(d)))
+    qgroups = triton.cdiv(n,qrows)
+    qpart = torch.empty((m,qgroups,d),device=grad.device,dtype=torch.float32)
+    wgroups = triton.cdiv(n,64)
+    wpart = torch.empty((ksum,wgroups,d),device=grad.device,dtype=torch.float32)
+    dq = torch.zeros_like(query, dtype=torch.float32)
+    dws = []
+    dxs = []
+    for source, bank in zip(sources, banks):
+        _msar_vjp_value[(qgroups,m)](source,bank,query,valid,position,lse,z,grad,dv,qpart,
+            n,t,d,m,tuple(widths),eps,qrows,triton.next_power_of_2(d),qgroups,num_warps=4)
+        dx = torch.empty_like(source)
+        _msar_vjp_source[(triton.cdiv(n*d,256),)](dv,bank,valid,position,dx,n,t,d,
+            tuple(widths),256,num_warps=4)
+        _msar_vjp_bank[(wgroups,triton.cdiv(d,32))](source,dv,valid,position,wpart,
+            n,t,d,tuple(widths),64,32,wgroups,num_warps=4)
+        dq.add_(qpart.sum(dim=1))
+        dws.append(wpart.sum(dim=1).to(bank.dtype))
+        dxs.append(dx)
+    return dxs, dws, dq.to(query.dtype).contiguous()
+
+
+@_msar_backward_op.register_fake
+def _msar_backward_fake(sources, banks, query, valid, position, widths, eps, grad, lse):
+    return ([torch.empty_like(x) for x in sources], [torch.empty_like(x) for x in banks],
+            torch.empty_like(query))
+
+
+def _msar_setup(ctx, inputs, output):
+    sources, banks, query, valid, position, widths, eps = inputs
+    ctx.save_for_backward(*sources, *banks, query, valid, position, output[1])
+    ctx.mark_non_differentiable(output[1])
+    ctx.count = len(sources)
+    ctx.widths = widths
+    ctx.eps = eps
+
+
+def _msar_autograd_backward(ctx, grad, grad_lse):
+    saved = ctx.saved_tensors
+    sources = list(saved[:ctx.count])
+    banks = list(saved[ctx.count:2*ctx.count])
+    query, valid, position, lse = saved[2*ctx.count:]
+    dx, dw, dq = _msar_backward_op(sources, banks, query, valid, position,
+                               ctx.widths, ctx.eps, grad.contiguous(), lse)
+    return dx, dw, dq, None, None, None, None
+
+
+_msar_read_op.register_autograd(_msar_autograd_backward, setup_context=_msar_setup)
+
+
+def _msar_read(sources, banks, query, valid, position, widths, eps):
+    return _msar_read_op(sources, banks, query, valid, position, widths, eps)[0]
+
+
+class MSARRead(nn.Module):
+    """One paper-style depth read with independent queries and H-pre mixing.
+
+    Sources are single-stream block summaries; the bank index identifies the
+    embedding, last branch of each completed block, or previous branch for the
+    current partial. RMS normalizations have no learned affine parameters.
+    The paper specifies zero query/convolution initialization and scale 0.01,
+    but not the mixer matrix or bias initialization: here they use the model's
+    normal initializer and zero bias, respectively. The terminal merge uses
+    the last branch input plus its output; no further expansion is performed.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.widths = tuple(config.msar_kernel_sizes)
+        self.eps = float(config.rms_norm_eps)
+        self.backend = config.msar_backend
+        m, d = len(self.widths) + 1, config.hidden_size
+        self.query = nn.Parameter(torch.zeros(m, d))
+        self.mixer_weight = nn.Parameter(torch.empty(m, m * d))
+        nn.init.normal_(self.mixer_weight, std=config.initializer_range)
+        self.mixer_bias = nn.Parameter(torch.zeros(m))
+        self.mixer_scale = nn.Parameter(torch.tensor(0.01))
+
+    def forward(self, sources, partial, banks, branch, block_size, valid, position):
+        source_list = list(sources)
+        bank_ids = [0] + [i * block_size for i in range(1, len(sources))]
+        if partial is not None:
+            source_list.append(partial)
+            bank_ids.append(branch)
+        # Slices remain views of the shared bank parameter. Their pointers
+        # let Triton reuse a specialization across branches with equal source counts.
+        selected_banks = [banks[i] for i in bank_ids]
+        if self.backend == "auto" and source_list[0].is_cuda:
+            if triton is None:
+                raise RuntimeError("MSAR on CUDA requires Triton; set msar_backend='torch' for reference.")
+            streams = _msar_read(source_list, selected_banks, self.query, valid, position,
+                                 list(self.widths), self.eps)
+        else:
+            streams = _msar_reference_read(source_list, selected_banks, self.query, valid,
+                                           position, self.widths, self.eps)
+        flat = streams.float().flatten(-2)
+        normed = flat * torch.rsqrt(flat.square().mean(dim=-1, keepdim=True) + self.eps)
+        gates = torch.sigmoid(self.mixer_scale.float() *
+                              F.linear(normed, self.mixer_weight.float()) +
+                              self.mixer_bias.float())
+        return (gates.unsqueeze(-1) * streams.float()).sum(dim=-2).to(streams.dtype)
+
+
+def _msar_start_branch(sources, partial, branch, block_size):
+    """Publish a completed single-stream block before a new block starts."""
+    if branch and branch % block_size == 0:
+        if partial is None:
+            raise RuntimeError("MSAR block boundary is missing its partial summary.")
+        return sources + (partial,), None
+    return sources, partial
+
+
 class NeoLLMDecoderLayer(GradientCheckpointingLayer):
     """
     Decoder layer with standard residual connections.
@@ -6776,6 +7201,10 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         # Pseudo-queries are initialized to ZERO so AttnRes starts as uniform
         # average (equivalent to standard residual mean) and training volatility
         # is avoided. This is critical per the paper's ablation.
+        self.use_msar = bool(getattr(config, "use_msar", False))
+        self.msar_block_branches = config.msar_block_branches
+        self.msar_attn_read = MSARRead(config) if self.use_msar else None
+        self.msar_mlp_read = MSARRead(config) if self.use_msar else None
         self.use_attn_res = getattr(config, "use_attn_res", False)
         if self.use_attn_res:
             self.attn_res_query_attn = nn.Parameter(torch.zeros(config.hidden_size))
@@ -7208,6 +7637,11 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         first_layer_fan: Optional[torch.Tensor] = None,
         attn_res_sources: Optional[list] = None,
         attn_res_partial: Optional[torch.Tensor] = None,
+        msar_sources=None,
+        msar_partial=None,
+        msar_banks=None,
+        msar_valid=None,
+        msar_position=None,
         output_attentions: Optional[bool] = False,
         repo_rope_args: Optional[Tuple[torch.Tensor, float]] = None,
         position_ids: Optional[torch.LongTensor] = None,
@@ -7233,7 +7667,14 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         # attn_res_partial carries the intra-block standard residual that
         # connects the attention and MLP sublayers within this layer.
         # When inactive, flow is identical to the original.
-        if (
+        if self.use_msar:
+            branch = 2 * self.layer_idx
+            msar_sources, msar_partial = _msar_start_branch(
+                msar_sources, msar_partial, branch, self.msar_block_branches)
+            h_attn = self.msar_attn_read(msar_sources, msar_partial, msar_banks,
+                branch, self.msar_block_branches, msar_valid, msar_position)
+            residual_attn = msar_partial
+        elif (
             self.use_attn_res
             and attn_res_sources is not None
             and attn_res_partial is not None
@@ -7283,7 +7724,10 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         )
 
         # ── Residual junction (attention sublayer) ────────────────────────
-        if self.use_laurel:
+        if self.use_msar:
+            msar_partial = hidden_states if msar_partial is None else msar_partial + hidden_states
+            attn_aug = msar_partial
+        elif self.use_laurel:
             attn_aug = self._laurel_residual(
                 residual_attn,
                 hidden_states,
@@ -7301,7 +7745,14 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         # After attention, the partial sum is updated with h_tilde.
         # The pre-MLP AttnRes attends over the same sources but with h_tilde
         # as the current partial — capturing the within-layer attention output.
-        if self.use_attn_res and attn_res_sources is not None:
+        if self.use_msar:
+            branch = 2 * self.layer_idx + 1
+            msar_sources, msar_partial = _msar_start_branch(
+                msar_sources, msar_partial, branch, self.msar_block_branches)
+            h_mlp = self.msar_mlp_read(msar_sources, msar_partial, msar_banks,
+                branch, self.msar_block_branches, msar_valid, msar_position)
+            residual_mlp = msar_partial
+        elif self.use_attn_res and attn_res_sources is not None:
             h_mlp = self._attn_res(
                 attn_res_sources,
                 h_tilde,
@@ -7334,7 +7785,11 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
 
         # ── Residual junction (MLP sublayer) ─────────────────────────────
         # LAuReL augments the base MLP residual (residual_mlp + delta_m).
-        if self.use_laurel:
+        if self.use_msar:
+            msar_partial = delta_m if msar_partial is None else msar_partial + delta_m
+            # Representative intermediate state for existing layer diagnostics.
+            mlp_aug = h_mlp + delta_m
+        elif self.use_laurel:
             mlp_aug = self._laurel_residual(
                 residual_mlp,
                 delta_m,
@@ -7349,6 +7804,8 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         hidden_states = self.gpas_mlp(mlp_aug) if self.use_gpas else mlp_aug
 
         outputs = (hidden_states,)
+        if self.use_msar:
+            outputs += (msar_sources, msar_partial)
         if jtok_aux_stats is not None:
             outputs += jtok_aux_stats
         if output_attentions:
@@ -8058,6 +8515,11 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                 for layer_idx in range(config.num_hidden_layers)
             ]
         )
+        self.msar_banks = (
+            nn.Parameter(torch.zeros(2 * config.num_hidden_layers,
+                sum(config.msar_kernel_sizes), config.hidden_size))
+            if config.use_msar else None
+        )
         self.dagformer = (DAGFormerTopology(config, self.layers)
                           if bool(getattr(config, "use_dagformer", False)) else None)
         self.use_siamesenorm = bool(getattr(config, "use_siamesenorm", False))
@@ -8191,15 +8653,17 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
             )
         return inputs_embeds, jtok_z_tilde
 
-    def prepare_patch_inputs(self, input_ids, attention_mask, *, patch_size, meap_active=False):
+    def prepare_patch_inputs(self, input_ids, attention_mask, *, patch_size, meap_active=False,
+                             return_token_features=False):
         """Prepare differentiable core inputs without running decoder layers.
 
         Input IDs must already contain the selected MEAP replacements, while
         clean future targets are maintained separately by the trainer. SBE and
         the dedicated MEAP vector are composed before pooling. The returned
         tuple is the pool_patch_inputs contract, including separate attention
-        and geometry validity. The K=1 path shares the ordinary producer;
-        no model parameters, optimizer state or coordinate gradients are copied.
+        and geometry validity. When requested, the final item is the observed
+        raw-token representation for lexical copy keys. The K=1 path shares
+        the ordinary producer; no parameters or coordinate gradients are copied.
         """
         from cut_cross_entropy import pool_patch_inputs
 
@@ -8210,10 +8674,11 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
             raise ValueError("Patch preparation requires an explicit boolean raw-token validity mask.")
         embeddings, coordinates = self._compose_input_embeddings(input_ids, attention_mask, meap_active)
         corruption = input_ids.eq(int(self.config.meap_mask_token_id)) if meap_active else None
-        return pool_patch_inputs(
+        pooled = pool_patch_inputs(
             embeddings, attention_mask, patch_size=patch_size,
             coordinates=coordinates, corruption_mask=corruption,
         )
+        return (*pooled, embeddings) if return_token_features else pooled
 
     def forward(
         self,
@@ -8477,6 +8942,18 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         # In both modes, attn_res_partial is the current intra-block accumulated
         # hidden state that connects the attn and MLP sublayers and flows between
         # decoder layers within a block.
+        use_msar = bool(getattr(self.config, "use_msar", False))
+        if use_msar:
+            msar_sources, msar_partial = (hidden_states,), None
+            if attention_mask is not None and attention_mask.ndim != 2:
+                raise ValueError("MSAR requires a 2-D token validity mask.")
+            msar_valid = (attention_mask.to(dtype=torch.bool).contiguous()
+                          if attention_mask is not None else
+                          torch.ones(hidden_states.shape[:2], dtype=torch.bool,
+                                     device=hidden_states.device))
+            msar_position = position_ids.expand(hidden_states.shape[:2]).contiguous()
+        else:
+            msar_sources = msar_partial = msar_valid = msar_position = None
         use_attn_res = getattr(self.config, "use_attn_res", False)
         attn_res_sources = None
         attn_res_partial = None
@@ -8564,6 +9041,9 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                     first_layer_fan=self.first_layer_fan,
                     attn_res_sources=attn_res_sources,
                     attn_res_partial=attn_res_partial if use_attn_res else None,
+                    msar_sources=msar_sources, msar_partial=msar_partial,
+                    msar_banks=self.msar_banks, msar_valid=msar_valid,
+                    msar_position=msar_position,
                     output_attentions=output_attentions,
                     repo_rope_args=repo_rope_args,
                     position_ids=position_ids,
@@ -8576,7 +9056,13 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                     **kwargs,
                 )
                 hidden_states = layer_outputs[0]
-                extras_start = 1
+                if use_msar:
+                    msar_sources, msar_partial = layer_outputs[1:3]
+                    extras_start = 3
+                    # The last branch's h_mlp + delta_m is the terminal
+                    # residual merge. No new source expansion is applied.
+                else:
+                    extras_start = 1
 
             if collect_residual:
                 packed = residual_dynamics(before_x.detach(), hidden_states.detach(),
@@ -8689,8 +9175,8 @@ def _prepare_lm_labels(labels, device, pad_token_id=None):
 
 
 def _lm_perplexity_sums(labels, ce, mile_delta, *, pad_token_id, patch_targets,
-                        mile_enabled):
-    """Detached sufficient statistics for CCE CE and MiLe pseudo-perplexity.
+                        mile_enabled, vocab_size=None):
+    """Detached sufficient statistics for NTP CE and MiLe pseudo-perplexity.
 
     For N valid supervised targets, return (N*CE, N*MiLe, N).
     MiLe = CE + reweighting_delta, excluding mu and all auxiliary objectives.
@@ -8708,6 +9194,8 @@ def _lm_perplexity_sums(labels, ce, mile_delta, *, pad_token_id, patch_targets,
     valid = targets != -100
     if pad_token_id is not None:
         valid = valid & (targets != pad_token_id)
+    if vocab_size is not None:
+        valid = valid & (targets >= 0) & (targets < vocab_size)
     count = valid.sum().to(dtype=torch.float32)
     ce = ce.detach().float()
     mile = ce + mile_delta.detach().float() if mile_enabled else ce
@@ -10506,6 +10994,235 @@ def compute_nextlat_loss(
     return total, mse_loss, kl_loss
 
 
+_COPY_PASTE_MAX_LOGIT_ELEMENTS = 1 << 24
+
+
+class CopyPasteHead(nn.Module):
+    """Causal vocabulary/pointer mixture from Godey and Artzi (2026).
+
+    A separate gate vector is equivalent to the paper's extra <copy> output
+    row without adding an ID to the tokenizer. The learned sink returns copy
+    probability to the ordinary LM. Input IDs are those actually observed by
+    the backbone, including any MEAP corruption. For pooled patches, all raw
+    tokens become visible together. By default they share the patch key;
+    optional lexical keys use ``k_i = W_k h_patch + U(V e_i)``, where ``e_i``
+    is the actual post-MEAP input embedding. The automatic rank depends only on
+    hidden size. No clean target is used to construct a copy key.
+
+    CUDA training delegates to CCE's sparse Copy-Paste kernels without a
+    [batch, sequence, vocabulary] tensor. The chunked Torch implementation
+    remains the CPU/reference path; inference still returns dense logits.
+    """
+
+    def __init__(self, hidden_size: int, lexical_keys: bool = False):
+        super().__init__()
+        self.query = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.key = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.copy_gate = nn.Linear(hidden_size, 1, bias=False)
+        self.sink_key = nn.Parameter(torch.zeros(hidden_size))
+        self.scale = hidden_size ** -0.5
+        rank = min(hidden_size, max(4, hidden_size // 16))
+        self.lexical_down = nn.Linear(hidden_size, rank, bias=False) if lexical_keys else None
+        self.lexical_up = nn.Linear(rank, hidden_size, bias=False) if lexical_keys else None
+
+    def _parts(self, hidden, query, keys, source_ids, source_mask, ends, weight, bias):
+        logits = F.linear(hidden, weight, bias).float()
+        lse = torch.logsumexp(logits, dim=-1)
+        gate_delta = self.copy_gate(hidden).squeeze(-1).float() - lse
+        scores = torch.matmul(query.float(), keys.float().transpose(-1, -2)) * self.scale
+        source_position = torch.arange(source_ids.shape[1], device=source_ids.device)
+        allowed = source_mask[:, None, :] & (source_position[None, None, :] <= ends[None, :, None])
+        scores = scores.masked_fill(~allowed, -torch.inf)
+        sink = torch.matmul(query.float(), self.sink_key.float()) * self.scale
+        normalizer = torch.logsumexp(torch.cat((sink.unsqueeze(-1), scores), dim=-1), dim=-1)
+        log_gate = -F.softplus(-gate_delta)
+        log_base_weight = torch.logaddexp(-F.softplus(gate_delta), log_gate + sink - normalizer)
+        return logits, lse, scores, normalizer, log_gate, log_base_weight
+
+    @staticmethod
+    def _distribution(logits, lse, scores, normalizer, log_gate, log_base_weight, source_ids):
+        probs = torch.exp(logits - lse.unsqueeze(-1) + log_base_weight.unsqueeze(-1))
+        copy_probs = torch.exp(scores - normalizer.unsqueeze(-1) + log_gate.unsqueeze(-1))
+        indices = source_ids[:, None, :].expand(-1, scores.shape[1], -1)
+        return probs.scatter_add(-1, indices, copy_probs)
+
+    def _loss_chunk(self, hidden, query, keys, source_ids, source_mask, ends,
+                    targets, weight, bias, need_entropy, need_topk):
+        logits, lse, scores, norm, log_gate, log_base = self._parts(
+            hidden, query, keys, source_ids, source_mask, ends, weight, bias)
+        valid = (targets >= 0) & (targets < weight.shape[0])
+        safe_targets = targets.clamp(0, weight.shape[0] - 1)
+        target_logits = logits.gather(-1, safe_targets)
+        matches = ((source_ids[:, None, None, :] == safe_targets[..., None])
+                   & torch.isfinite(scores).unsqueeze(-2))
+        # Avoid logsumexp(-inf, ..., -inf), whose backward is undefined. Rows
+        # without a match use a harmless finite reduction and then receive an
+        # exact -inf log-probability, so FP16 needs no overflowing sentinel.
+        has_match = matches.any(-1)
+        matched_scores = scores.unsqueeze(-2).masked_fill(~matches, -torch.inf)
+        matched_scores = torch.where(has_match.unsqueeze(-1), matched_scores,
+                                     torch.zeros_like(matched_scores))
+        copy_target = torch.where(
+            has_match, torch.logsumexp(matched_scores, dim=-1) - norm.unsqueeze(-1),
+            torch.full_like(target_logits, -torch.inf))
+        log_probability = torch.logaddexp(
+            log_base.unsqueeze(-1) + target_logits - lse.unsqueeze(-1),
+            log_gate.unsqueeze(-1) + copy_target,
+        )
+        nll = torch.where(valid, -log_probability, torch.zeros_like(log_probability))
+        entropy = torch.zeros_like(lse)
+        counts = logits.new_zeros(5)
+        with torch.no_grad():
+            # Sums are weighted by valid target slots, so partial patches and
+            # microbatches can be merged without averaging batch means.
+            c = log_gate.exp()
+            sink_score = torch.matmul(query.float(), self.sink_key.float()) * self.scale
+            s = (sink_score - norm).exp()
+            row_count = valid.sum(-1)
+            responsibility = (log_gate.unsqueeze(-1) + copy_target - log_probability).exp()
+            copy_metrics = torch.stack((
+                (c * row_count).sum(), (s * row_count).sum(),
+                (c * (1.0 - s) * row_count).sum(),
+                torch.where(valid, responsibility, 0.0).sum(),
+                (valid & has_match).sum(), valid.sum(),
+            )).float()
+        if need_entropy or need_topk:
+            with torch.no_grad():
+                probs = self._distribution(
+                    logits, lse, scores, norm, log_gate, log_base, source_ids)
+                if need_entropy:
+                    entropy = -(probs * probs.clamp_min(torch.finfo(torch.float32).tiny).log()).sum(-1)
+                    entropy = entropy.clamp(0.0, math.log(weight.shape[0]))
+                if need_topk:
+                    target = safe_targets[..., 0]
+                    target_prob = probs.gather(-1, target.unsqueeze(-1)).squeeze(-1)
+                    finite = torch.isfinite(target_prob)
+                    keep = valid[..., 0] & finite
+                    ids = torch.arange(weight.shape[0], device=probs.device)
+                    ahead = ((probs > target_prob.unsqueeze(-1)) |
+                             ((probs == target_prob.unsqueeze(-1)) & (ids < target.unsqueeze(-1))))
+                    rank = ahead.sum(-1)
+                    counts = torch.stack(((keep & (rank == 0)).sum(),
+                                          (keep & (rank < 3)).sum(),
+                                          keep.new_zeros((), dtype=torch.long),
+                                          valid[..., 0].sum(),
+                                          (valid[..., 0] & ~finite).sum())).float()
+        return nll, entropy, counts, copy_metrics
+
+    @getattr(torch.compiler, "nested_compile_region", lambda function: function)
+    def _loss_chunk_reusable(self, *args):
+        """Allow reuse of a compiled tile region when PyTorch supports it."""
+        return self._loss_chunk(*args)
+
+    def loss(self, hidden, source_ids, source_mask, targets, lm_head,
+             *, patch_size, group_fraction, mile_gamma, use_mile, use_mu,
+             mu_lambda, need_topk, source_token_features=None):
+        """Mixed loss with optional lexical keys from observed token features."""
+        rows = targets.shape[1]
+        if rows == 0:
+            zero = hidden.sum() * 0.0
+            mu = (float(mu_lambda) * lm_head.weight.float().mean(0).square().sum()
+                  if use_mu else zero)
+            return (zero + mu, zero.detach(), zero.detach(), mu.detach(),
+                    zero.new_zeros(5) if need_topk else None, zero.new_zeros(6))
+        keys = self.key(hidden).repeat_interleave(patch_size, dim=1)
+        if source_token_features is not None:
+            if self.lexical_down is None or source_token_features.shape != keys.shape:
+                raise ValueError("Lexical copy features must match the raw source token domain.")
+            keys = keys + self.lexical_up(self.lexical_down(source_token_features))
+        query = self.query(hidden[:, :rows])
+        if hidden.is_cuda:
+            if not _CCE_COPY_PASTE_HAS_METRICS:
+                raise ImportError(
+                    "CUDA Copy-Paste training requires a cut-cross-entropy installation "
+                    "with Copy-Paste metric support. Update the CCE repository first."
+                )
+            selected = hidden[:, :rows]
+            gate = self.copy_gate(selected).squeeze(-1).float()
+            sink = torch.matmul(query.float(), self.sink_key.float()) * self.scale
+            return _cce_copy_paste_loss(
+                selected, lm_head.weight, lm_head.bias, query, keys, source_ids,
+                source_mask, targets, gate, sink, patch_size=patch_size,
+                group_fraction=group_fraction, mile_gamma=mile_gamma, use_mile=use_mile,
+                use_mu=use_mu, mu_lambda=mu_lambda, need_topk=need_topk,
+                return_copy_metrics=True,
+            )
+        # Bound each temporary FP32 logit tile by ~64 MiB when one row fits.
+        # The value is a memory ceiling, not a model-specific geometry.
+        tile = max(1, min(rows, _COPY_PASTE_MAX_LOGIT_ELEMENTS // max(1, hidden.shape[0] * lm_head.weight.shape[0])))
+        repeated_tiles = tile < rows
+        if repeated_tiles:
+            # invoke_subgraph rejects aliased inputs (labels can be a view of
+            # input_ids); one compact label copy avoids that compiler failure.
+            targets = targets.clone()
+        chunk_function = self._loss_chunk_reusable if repeated_tiles else self._loss_chunk
+        nll_chunks, entropy_chunks, count_chunks, copy_metric_chunks = [], [], [], []
+        for start in range(0, rows, tile):
+            stop = min(start + tile, rows)
+            ends = (torch.arange(start, stop, device=hidden.device) + 1) * patch_size - 1
+            args = (hidden[:, start:stop], query[:, start:stop], keys,
+                    source_ids, source_mask, ends, targets[:, start:stop],
+                    lm_head.weight, lm_head.bias, use_mile, need_topk)
+            if torch.is_grad_enabled() and self.training:
+                values = _activation_checkpoint(chunk_function, *args, use_reentrant=False)
+            else:
+                values = chunk_function(*args)
+            nll_chunks.append(values[0])
+            entropy_chunks.append(values[1])
+            count_chunks.append(values[2])
+            copy_metric_chunks.append(values[3])
+        nll = torch.cat(nll_chunks, dim=1)
+        valid = (targets >= 0) & (targets < lm_head.weight.shape[0])
+        valid_count = valid.sum().to(torch.float32)
+        count = valid_count.clamp_min(1.0)
+        ce = nll.sum() / count
+        weighted = ce
+        if use_mile:
+            entropy = torch.cat(entropy_chunks, dim=1).detach()
+            row_count = valid.sum(-1).to(torch.float32)
+            active = row_count > 0
+            log_base = torch.where(active, torch.log1p(entropy), 0.0)
+            if group_fraction is None:
+                raw = torch.exp(float(mile_gamma) * (log_base - log_base.amax()))
+                scale = valid_count / (raw * row_count).sum().clamp_min(1e-20)
+                weights = raw * scale
+            else:
+                group = group_fraction.float()
+                other = 1.0 - group
+                group_active = active & (group > 0)
+                other_active = active & (other > 0)
+                group_peak = torch.where(group_active, log_base, -torch.inf).amax()
+                other_peak = torch.where(other_active, log_base, -torch.inf).amax()
+                group_peak = torch.where(torch.isfinite(group_peak), group_peak, 0.0)
+                other_peak = torch.where(torch.isfinite(other_peak), other_peak, 0.0)
+                group_offset = torch.where(group_active, log_base - group_peak, 0.0)
+                other_offset = torch.where(other_active, log_base - other_peak, 0.0)
+                group_raw = torch.where(group_active, torch.exp(float(mile_gamma) * group_offset), 0.0)
+                other_raw = torch.where(other_active, torch.exp(float(mile_gamma) * other_offset), 0.0)
+                group_count = (row_count * group).sum()
+                other_count = (row_count * other).sum()
+                group_scale = group_count / (row_count * group * group_raw).sum().clamp_min(1e-20)
+                other_scale = other_count / (row_count * other * other_raw).sum().clamp_min(1e-20)
+                weights = group * group_raw * group_scale + other * other_raw * other_scale
+            weighted = (nll * weights.unsqueeze(-1)).sum() / count
+        mu = (float(mu_lambda) * lm_head.weight.float().mean(0).square().sum()
+              if use_mu else ce.new_zeros(()))
+        topk = torch.stack(count_chunks).sum(0) if need_topk else None
+        return (weighted + mu, ce.detach(), (weighted - ce).detach(), mu.detach(), topk,
+                torch.stack(copy_metric_chunks).sum(0).detach())
+
+    def logits(self, hidden, source_ids, source_mask, positions, lm_head, *, patch_size=1):
+        """Return normalized log-probabilities for selected query positions."""
+        keys = self.key(hidden).repeat_interleave(patch_size, dim=1)
+        selected = hidden[:, positions]
+        query = self.query(selected)
+        ends = (positions + 1) * patch_size - 1
+        parts = self._parts(selected, query, keys, source_ids, source_mask, ends,
+                            lm_head.weight, lm_head.bias)
+        probs = self._distribution(*parts[:3], parts[3], parts[4], parts[5], source_ids)
+        return probs.clamp_min(torch.finfo(torch.float32).tiny).log()
+
+
 class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
     """
     Causal LM with NeoLLM backbone.
@@ -10519,6 +11236,11 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         self.model = NeoLLMModel(config)
         self.vocab_size = config.vocab_size
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.copy_paste_head = (
+            CopyPasteHead(config.hidden_size,
+                          lexical_keys=bool(getattr(config, "copy_paste_lexical_keys", False)))
+            if getattr(config, "use_copy_paste", False) else None
+        )
         self.nitp_projector = NITPProjector(config) if config.use_nitp else None
         self.nitp_temporal_transition = (
             NITPTemporalTransition(config)
@@ -10550,10 +11272,13 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 "`ntp_loss_backend` must be 'cce' or 'liger'. "
                 f"Got {self.ntp_loss_backend!r}."
             )
-        if bool(getattr(config, "use_mile_loss", False)):
-            _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")
+        # Copy-Paste computes MiLe and mu-loss in its own Torch path; those
+        # CCE loss extensions are only needed when CCE computes the NTP loss.
+        if self.copy_paste_head is None:
+            if bool(getattr(config, "use_mile_loss", False)):
+                _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")
         # 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.
@@ -10561,6 +11286,7 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         self._last_ntp_ce_unweighted = None
         self._last_lm_metric_sums = None
         self._last_topk_metric_sums = None
+        self._last_copy_metric_sums = None
         self._last_mile_reweighting_delta = None
         self._last_mu_loss = None
         self._last_meap_mask_fraction = None
@@ -10679,8 +11405,9 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
     def get_dynamics_parameter_groups(self):
         """Existing observers plus residual controllers; never duplicate the backbone.
 
-        Control modules share one sampled group per layer/method. Groups are
-        sampled around real optimizer updates by train.py, not in forward.
+        Control modules share one sampled group per layer/method. VHA factors
+        and Copy-Paste projections each have a separate sampled group. Groups
+        are sampled around real optimizer updates by train.py, not in forward.
         """
         groups = {}
         generator = getattr(self.model, "token_generator", None)
@@ -10688,7 +11415,16 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             for name, parameter in generator.named_parameters():
                 if parameter.requires_grad:
                     groups[f"leviathan/{name}"] = (parameter,)
+        if self.copy_paste_head is not None:
+            for name, parameter in self.copy_paste_head.named_parameters():
+                if parameter.requires_grad:
+                    groups[f"copy_paste/{name}"] = (parameter,)
         for layer in self.model.layers:
+            if layer.self_attn.use_vha_attention:
+                for name in ("vha_premix", "vha_postmix_a", "vha_postmix_b"):
+                    parameter = getattr(layer.self_attn, name)
+                    if parameter.requires_grad:
+                        groups[f"vha/layer_{layer.layer_idx}/{name}"] = (parameter,)
             module = getattr(layer, "jtok", None)
             if module is not None:
                 kind = "jtokm" if module.use_mixture else "jtok"
@@ -10729,7 +11465,12 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         self.lm_head = new_embeddings
 
     def _corrupt_meap_inputs(self, input_ids, attention_mask, meap_seed, *, return_mask):
-        """Shared raw-token corruption for ordinary and prepared training."""
+        """Corrupt raw IDs while preserving the configured structural tokens.
+
+        The last valid position is removed before structural-token filtering so
+        that a final EOS does not additionally exclude the preceding token.
+        The disabled protection path retains the legacy MEAP behavior.
+        """
         if meap_mask_inputs is None:
             raise ImportError("Active MEAP training requires the extended cut-cross-entropy package.")
         if "return_metrics" not in inspect.signature(meap_mask_inputs).parameters:
@@ -10737,20 +11478,42 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         mask_token_id = getattr(self.config, "meap_mask_token_id", None)
         if mask_token_id is None:
             raise ValueError("`meap_mask_token_id` is required when MEAP is enabled.")
-        if meap_seed is None:
-            seed = input_ids.new_tensor(int(getattr(self.config, "meap_seed", 0)), dtype=torch.long)
-        elif isinstance(meap_seed, torch.Tensor):
+        implementation = str(getattr(self.config, "meap_implementation", "triton"))
+        if isinstance(meap_seed, torch.Tensor):
             seed = meap_seed.to(device=input_ids.device, dtype=torch.long)
+            if implementation == "torch":
+                raise TypeError("MEAP's torch reference requires a Python integer seed; use triton for tensor seeds.")
+            kernel_seed = seed
         else:
-            seed = input_ids.new_tensor(int(meap_seed), dtype=torch.long)
+            seed_value = int(getattr(self.config, "meap_seed", 0)) if meap_seed is None else int(meap_seed)
+            seed = input_ids.new_tensor(seed_value, dtype=torch.long)
+            kernel_seed = seed_value if implementation == "torch" else seed
+
+        eligible = (attention_mask.to(device=input_ids.device, dtype=torch.bool)
+                    if attention_mask is not None else None)
+        exclude_last = bool(getattr(self.config, "meap_exclude_last", True))
+        if bool(getattr(self.config, "meap_protect_structural_tokens", True)):
+            if eligible is None:
+                eligible = torch.ones_like(input_ids, dtype=torch.bool)
+            if exclude_last and input_ids.shape[1] > 0:
+                positions = torch.arange(input_ids.shape[1], device=input_ids.device)
+                last_valid = torch.where(eligible, positions, -1).amax(dim=1, keepdim=True)
+                eligible = eligible & positions.ne(last_valid)
+                # The CCE kernel normally removes the last *eligible* token;
+                # doing so after filtering EOS would drop a second position.
+                exclude_last = False
+            for value in (getattr(self.config, "pad_token_id", None),
+                          getattr(self.config, "bos_token_id", None),
+                          getattr(self.config, "eos_token_id", None), mask_token_id):
+                for token_id in value if isinstance(value, (list, tuple)) else (value,):
+                    if token_id is not None:
+                        eligible = eligible & input_ids.ne(int(token_id))
         result = meap_mask_inputs(
             input_ids, int(mask_token_id), enabled=True,
             mask_ratio=float(getattr(self.config, "meap_mask_ratio", 0.15)),
-            eligible_mask=(attention_mask.to(device=input_ids.device, dtype=torch.bool)
-                           if attention_mask is not None else None),
-            seed=seed, exclude_last=bool(getattr(self.config, "meap_exclude_last", True)),
+            eligible_mask=eligible, seed=kernel_seed, exclude_last=exclude_last,
             return_mask=return_mask, return_metrics=True,
-            implementation=str(getattr(self.config, "meap_implementation", "triton")),
+            implementation=implementation,
         )
         if return_mask:
             corrupted, selected, metrics = result
@@ -10792,9 +11555,13 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             input_ids, selected, metrics, seed = self._corrupt_meap_inputs(
                 input_ids, attention_mask, meap_seed, return_mask=True,
             )
-        e, z, core, geometry, fraction = self.model.prepare_patch_inputs(
+        need_lexical = (self.copy_paste_head is not None
+                        and self.copy_paste_head.lexical_down is not None and patch_size > 1)
+        prepared_inputs = self.model.prepare_patch_inputs(
             input_ids, attention_mask, patch_size=patch_size, meap_active=active,
+            return_token_features=need_lexical,
         )
+        e, z, core, geometry, fraction = prepared_inputs[:5]
         if selected is not None:
             grouped_valid = attention_mask.unflatten(-1, (core.size(1), patch_size))
             fraction = ((selected & attention_mask).unflatten(-1, (core.size(1), patch_size))
@@ -10813,6 +11580,13 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             patch_aux_valid_mask=aux_valid,
             prepared_meap_metrics=metrics, prepared_meap_seed=seed,
         )
+        if self.copy_paste_head is not None:
+            # Preserve only the observed raw IDs and their original mask;
+            # externally prepared token features must follow the same contract:
+            # clean labels/features would leak MEAP-hidden content to the pointer.
+            result.update(copy_input_ids=input_ids, copy_attention_mask=attention_mask)
+            if need_lexical:
+                result["copy_token_features"] = prepared_inputs[5]
         if z is not None:
             result.update(prepared_jtok_coordinates=z, prepared_jtok_valid_mask=geometry)
         return result
@@ -10860,11 +11634,23 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         prepared_meap_seed: Optional[torch.Tensor] = None,
         raw_patch_size: Optional[int] = None,
         patch_target_slots: int = 4,
+        copy_input_ids: Optional[torch.LongTensor] = None,
+        copy_attention_mask: Optional[torch.Tensor] = None,
+        copy_token_features: Optional[torch.Tensor] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> CausalLMOutputWithPast:
         self._last_meap_mask_fraction = None
         self._last_meap_masked_tokens = None
         self._last_meap_seed = None
+        copy_source_ids = copy_input_ids
+        copy_source_mask = copy_attention_mask
+        copy_source_features = copy_token_features
+        copy_patch_size = 1
+        if self.copy_paste_head is None and (copy_input_ids is not None or copy_attention_mask is not None
+                                             or copy_token_features is not None):
+            raise ValueError("Copy source IDs, mask and token features require use_copy_paste=True.")
+        if copy_token_features is not None and not patch_training_enabled:
+            raise ValueError("copy_token_features is only accepted for prepared patches.")
 
         if raw_patch_size is not None:
             # Static K/shape signatures are warmed once. The changing optimizer
@@ -10872,7 +11658,9 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             # so embedding/coordinate VJPs are part of the compiled backward.
             if (patch_training_enabled or inputs_embeds is not None
                     or input_ids is None or labels is None or attention_mask is None
-                    or position_ids is not None or any(value is not None for value in (
+                    or position_ids is not None or copy_input_ids is not None
+                    or copy_attention_mask is not None or copy_token_features is not None
+                    or any(value is not None for value in (
                         patch_group_fraction, patch_aux_valid_mask,
                         prepared_meap_metrics, prepared_meap_seed,
                         kwargs.get("prepared_jtok_coordinates"),
@@ -10884,6 +11672,11 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 patch_size=raw_patch_size, target_slots=patch_target_slots,
                 meap_seed=meap_seed,
             )
+            if self.copy_paste_head is not None:
+                copy_source_ids = prepared["copy_input_ids"]
+                copy_source_mask = prepared["copy_attention_mask"]
+                copy_source_features = prepared.get("copy_token_features")
+                copy_patch_size = raw_patch_size
             input_ids = None
             inputs_embeds = prepared["inputs_embeds"]
             attention_mask = prepared["attention_mask"]
@@ -10959,6 +11752,46 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             self._last_meap_masked_tokens = masked_count.detach()
             self._last_meap_seed = effective_meap_seed.detach()
 
+        if self.copy_paste_head is not None:
+            if input_ids is not None:
+                if copy_input_ids is not None or copy_attention_mask is not None:
+                    raise ValueError("Copy-Paste uses the observed input_ids and attention_mask; do not supply separate copy inputs too.")
+                copy_source_ids = input_ids
+                copy_source_mask = attention_mask
+            elif not patch_training_enabled and copy_source_mask is None:
+                copy_source_mask = attention_mask
+            if copy_source_ids is None:
+                raise ValueError("Copy-Paste with inputs_embeds requires the corresponding observed copy_input_ids.")
+            if copy_source_ids.ndim != 2 or copy_source_ids.dtype not in (torch.int32, torch.int64):
+                raise ValueError("Copy-Paste source IDs must be a 2D integer tensor.")
+            if copy_source_mask is None:
+                if patch_training_enabled and raw_patch_size is None:
+                    raise ValueError("Prepared patches require their raw copy_attention_mask.")
+                copy_source_mask = torch.ones_like(copy_source_ids, dtype=torch.bool)
+            if copy_source_mask.shape != copy_source_ids.shape or copy_source_mask.device != copy_source_ids.device:
+                raise ValueError("Copy-Paste source mask must match the raw IDs in shape and device.")
+            copy_source_mask = copy_source_mask.to(dtype=torch.bool)
+            copy_source_ids = copy_source_ids.to(dtype=torch.long)
+            if meap_enabled and bool(getattr(self.config, "meap_exclude_mask_from_copy", True)):
+                # Use observed IDs only: never restore the clean target here.
+                copy_source_mask = copy_source_mask & copy_source_ids.ne(int(self.config.meap_mask_token_id))
+            core = inputs_embeds if inputs_embeds is not None else input_ids
+            core_length = core.shape[1]
+            if copy_source_ids.shape[0] != core.shape[0] or copy_source_ids.device != core.device:
+                raise ValueError("Copy-Paste source IDs must match the model batch and device.")
+            if patch_training_enabled and raw_patch_size is None:
+                if core_length == 0 or copy_source_ids.shape[1] % core_length:
+                    raise ValueError("Copy-Paste raw source length must be divisible by patch core length.")
+                copy_patch_size = copy_source_ids.shape[1] // core_length
+            if copy_patch_size < 1 or copy_source_ids.shape[1] != core_length * copy_patch_size:
+                raise ValueError("Copy-Paste source length does not match the active token/patch geometry.")
+            if patch_training_enabled and copy_patch_size > 1 and self.copy_paste_head.lexical_down is not None:
+                if (copy_source_features is None
+                        or copy_source_features.shape != (*copy_source_ids.shape, self.config.hidden_size)
+                        or copy_source_features.device != copy_source_ids.device
+                        or not copy_source_features.is_floating_point()):
+                    raise ValueError("Lexical Copy-Paste requires observed raw-token embedding features.")
+
         tweo_enabled = (
             bool(getattr(self.config, "use_tweo", False))
             and labels is not None
@@ -11086,6 +11919,7 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         self._last_ntp_ce_unweighted = None
         self._last_lm_metric_sums = None
         self._last_topk_metric_sums = None
+        self._last_copy_metric_sums = None
         self._last_mile_reweighting_delta = None
         self._last_mu_loss = None
         self._last_tweo_loss = None
@@ -11126,7 +11960,33 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         self._last_stack_gate_effective_slots = None
         self._last_total_loss = None
         if labels is not None:
-            if self.ntp_loss_backend == "liger":
+            if self.copy_paste_head is not None:
+                targets = _prepare_lm_labels(labels, hidden_states.device, self.config.pad_token_id)
+                if not patch_training_enabled:
+                    targets = targets[:, 1:].unsqueeze(-1)
+                use_mile = bool(getattr(self.config, "use_mile_loss", False))
+                group_fraction = None
+                if use_mile and patch_training_enabled:
+                    group_fraction = patch_group_fraction[:, :-1]
+                elif use_mile and meap_selected_mask is not None:
+                    group_fraction = meap_selected_mask[:, :-1]
+                collect_topk = (bool(getattr(self.config, "use_topk_metrics", False))
+                                and (not patch_training_enabled or raw_patch_size == 1))
+                ntp_loss, ce, delta, mu, topk, copy_metrics = self.copy_paste_head.loss(
+                    hidden_states, copy_source_ids, copy_source_mask, targets, self.lm_head,
+                    patch_size=copy_patch_size, group_fraction=group_fraction,
+                    mile_gamma=float(getattr(self.config, "mile_loss_gamma", 1.0)),
+                    use_mile=use_mile, use_mu=bool(getattr(self.config, "use_mu_loss", False)),
+                    mu_lambda=float(getattr(self.config, "mu_loss_lambda", 1e-4)),
+                    need_topk=collect_topk,
+                    source_token_features=copy_source_features,
+                )
+                self._last_ntp_ce_unweighted = ce
+                self._last_mile_reweighting_delta = delta
+                self._last_mu_loss = mu
+                self._last_topk_metric_sums = topk
+                self._last_copy_metric_sums = copy_metrics
+            elif self.ntp_loss_backend == "liger":
                 ntp_loss = compute_liger_loss(
                     hidden_states,
                     labels,
@@ -11202,6 +12062,7 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                     patch_targets=patch_training_enabled,
                     mile_enabled=(self.ntp_loss_backend == "cce"
                                   and bool(getattr(self.config, "use_mile_loss", False))),
+                    vocab_size=(self.vocab_size if self.copy_paste_head is not None else None),
                 )
             loss = ntp_loss
 
@@ -11541,7 +12402,13 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 if isinstance(logits_to_keep, int)
                 else logits_to_keep
             )
-            logits = self.lm_head(hidden_states[:, slice_indices, :])
+            if self.copy_paste_head is not None:
+                positions = torch.arange(hidden_states.shape[1], device=hidden_states.device)[slice_indices].reshape(-1)
+                logits = self.copy_paste_head.logits(
+                    hidden_states, copy_source_ids, copy_source_mask, positions, self.lm_head,
+                    patch_size=copy_patch_size)
+            else:
+                logits = self.lm_head(hidden_states[:, slice_indices, :])
 
         return CausalLMOutputWithPast(
             loss=loss,