modeling_source_qwen25.py

"""Raw-evidence/question recurrent reasoning on Qwen2.5-VL.

The oracle-selected ``raw_evidence_question`` path keeps the complete frozen
visual state as immutable evidence ``E=V*`` and initializes the recurrent state
with the text-only question scaffold ``R0=Q*``.  One shared Cross-Attention
transition reads ``E`` for four residual steps, then parameter-free trajectory
aggregation preserves every recurrent depth::

    R0 = Q*
    Rk = R{k-1} + CrossAttn(R{k-1}, E)
    R_agg = mean(R1, R2, R3, R4) -> y

During training only, ``R_agg`` is also restored toward the image-conditioned
question rows from the frozen multimodal lower pass.  That target is already
available from the pass that builds ``V*`` and never enters the inference
decoder. Historical final-state, state-replacement, and question-anchor modes
remain loadable for checkpoint comparison.

The multimodal upper continuation and its KV cache are discarded.  The lower
replacement cache contains only text-only ``Q*`` rows; the same positions are
overwritten by the selected decoder latent above ``ell_star``.  Thus both
halves have length ``N_q``:
there are no workspace slots, placeholder tokens, learned aggregation, or
direct ``E``-to-decoder connection in this path.

The cross-attention uses Qwen's native 28-query/4-KV-head GQA projections and
starts from the first upper decoder layer's Q/K/V/O and RMSNorm weights.  It is
still cross-attention, so the source layer's self-attention RoPE is not reused.
The current method has no workspace FFN. In residual mode, disabling a visual
read makes that step an identity; in replacement mode it writes zero on valid
question rows. A compatibility-only FFN can still be allocated for older
checkpoints.

Historical slot-workspace implementations remain below for loading and
evaluating old checkpoints.  Configuration validation makes those modules
mutually exclusive with ``raw_evidence_question``.

Shapes: ``B`` batch, ``N_v`` visual tokens, ``N_q`` question tokens, and ``d``
the native backbone width (3584).
"""

from __future__ import annotations

from contextlib import contextmanager
from dataclasses import dataclass, replace

import math

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint
from transformers.cache_utils import Cache, DynamicCache
from transformers.integrations.sdpa_attention import sdpa_attention_forward
from transformers.modeling_outputs import ModelOutput
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
    Qwen2_5_VLForConditionalGeneration,
    Qwen2_5_VLPreTrainedModel,
    apply_multimodal_rotary_pos_emb,
)
from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm

from .configuration_source_qwen25 import CloseQwen2_5_VLConfig
from .source_perceive import (
    PerceiveDeliberateLayout,
    build_answer_decode_mask,
    build_perceive_deliberate_mask,
    build_perceive_deliberate_positions,
)
from .source_spatial import SpatialVisualRecurrentCell
from .source_splitting import (
    block_attention_edges,
    embed_multimodal,
    final_norm,
    image_token_mask,
    make_split_context,
    run_layer_range,
    select_tokens_padded,
    vision_span_mask,
)


def _replace_selected_rows_from_padded(
    full_state: torch.Tensor,
    selection_mask: torch.Tensor,
    padded_rows: torch.Tensor,
    padding_mask: torch.Tensor,
) -> torch.Tensor:
    """Replace variable-count rows of ``full_state`` from a padded sequence.

    This is the inverse layout operation of :func:`select_tokens_padded` for
    valid rows.  The small per-example loop is intentional: image resolution
    makes the multimodal sequence lengths heterogeneous, while preserving each
    sample's original physical/M-RoPE positions is part of the native-layer
    counterfactual.
    """
    if selection_mask.dtype != torch.bool or padding_mask.dtype != torch.bool:
        raise TypeError("selection and padding masks must be boolean")
    if full_state.ndim != 3 or padded_rows.ndim != 3:
        raise ValueError("full_state and padded_rows must be [B, L, D]")
    if selection_mask.shape != full_state.shape[:2]:
        raise ValueError("selection_mask does not match full_state")
    if padding_mask.shape != padded_rows.shape[:2]:
        raise ValueError("padding_mask does not match padded_rows")
    if full_state.shape[0] != padded_rows.shape[0] or (
        full_state.shape[-1] != padded_rows.shape[-1]
    ):
        raise ValueError("full and padded states have incompatible shapes")

    selected_counts = selection_mask.sum(dim=-1)
    padded_counts = (~padding_mask).sum(dim=-1)
    if not torch.equal(selected_counts, padded_counts):
        raise ValueError(
            "selected-row counts disagree with padded valid counts: "
            f"{selected_counts.tolist()} vs {padded_counts.tolist()}"
        )
    output = full_state.clone()
    for batch_index, count_tensor in enumerate(selected_counts):
        count = int(count_tensor)
        output[batch_index, selection_mask[batch_index]] = padded_rows[
            batch_index, :count
        ]
    return output


def answer_cross_entropy(
    logits: torch.Tensor,
    labels: torch.Tensor,
    *,
    per_example: bool = False,
) -> torch.Tensor:
    """Answer CE with either token- or example-balanced reduction.

    ``per_example=False`` exactly matches PyTorch's historical global valid-token
    mean.  ``True`` gives every non-empty sample equal weight regardless of its
    answer length, preventing long caption answers from overwhelming one-token
    visual-decision examples.
    """
    if not per_example:
        return nn.functional.cross_entropy(
            logits.reshape(-1, logits.shape[-1]).float(),
            labels.reshape(-1),
            ignore_index=-100,
        )
    bsz = labels.shape[0]
    token_loss = nn.functional.cross_entropy(
        logits.reshape(-1, logits.shape[-1]).float(),
        labels.reshape(-1),
        ignore_index=-100,
        reduction="none",
    ).view(bsz, -1)
    valid = labels.reshape(bsz, -1).ne(-100)
    counts = valid.sum(dim=-1)
    nonempty = counts > 0
    if not bool(nonempty.any()):
        return logits.float().sum() * 0.0
    sample_loss = (token_loss * valid).sum(dim=-1) / counts.clamp_min(1)
    return sample_loss[nonempty].mean()


def _answer_nll_per_example(
    logits: torch.Tensor,
    labels: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return length-normalized answer NLL and a non-empty-row mask.

    The CRR reliance intervention compares factual and swapped answers across
    examples with very different target lengths.  A per-token mean keeps that
    score gap in the same units for one-token MC labels and caption answers.
    """
    bsz = labels.shape[0]
    token_loss = nn.functional.cross_entropy(
        logits.reshape(-1, logits.shape[-1]).float(),
        labels.reshape(-1),
        ignore_index=-100,
        reduction="none",
    ).view(bsz, -1)
    valid = labels.reshape(bsz, -1).ne(-100)
    counts = valid.sum(dim=-1)
    nll = (token_loss * valid).sum(dim=-1) / counts.clamp_min(1)
    return nll, counts > 0


def _different_answer_residual_donors(
    labels: torch.Tensor,
    question_attention_mask: torch.Tensor | None = None,
) -> tuple[torch.LongTensor, torch.BoolTensor]:
    """Choose a deterministic batch-local donor with a different answer.

    Equal question lengths are preferred so that right-padded Q-shaped
    residuals align exactly.  Equal answer lengths are the second preference,
    which usually keeps MC examples in the same formatting regime.  Rows for
    which the local batch contains no different answer are marked invalid and
    do not contribute to the reliance loss.
    """
    bsz = labels.shape[0]
    rows = torch.arange(bsz, device=labels.device)
    donors = rows.clone()
    found = torch.zeros(bsz, dtype=torch.bool, device=labels.device)
    if bsz < 2:
        return donors, found

    same_answer = labels[:, None, :].eq(labels[None, :, :]).all(dim=-1)
    answer_lengths = labels.ne(-100).sum(dim=-1)
    if question_attention_mask is None:
        question_lengths = torch.ones_like(answer_lengths)
    else:
        question_lengths = question_attention_mask.gt(0).sum(dim=-1)

    # Cyclic offsets avoid concentrating every negative on batch row zero.
    # Compatibility is relaxed only after all offsets at the stricter level
    # have been exhausted.
    for compatibility in ("question", "answer", "any"):
        for offset in range(1, bsz):
            candidates = (rows + offset) % bsz
            eligible = ~same_answer[rows, candidates]
            if compatibility == "question":
                eligible &= question_lengths.eq(
                    question_lengths.index_select(0, candidates)
                )
            elif compatibility == "answer":
                eligible &= answer_lengths.eq(
                    answer_lengths.index_select(0, candidates)
                )
            choose = (~found) & eligible
            donors[choose] = candidates[choose]
            found |= choose
            if bool(found.all()):
                return donors, found
    return donors, found


def _replacement_prefix_layout(
    question_mask: torch.Tensor, num_slots: int
) -> tuple[torch.Tensor, torch.LongTensor, torch.LongTensor]:
    """Mask and logical positions for the physical ``[q_pad; slots]`` layout.

    Questions are right-padded to a common physical width before the ``S`` slot
    rows are appended. The slots must nevertheless occupy the *logical*
    positions immediately after each item's last real question token. Using a
    batch-wide ``arange(N_q + S)`` makes RoPE positions, and therefore answers,
    depend on the longest question in the batch.

    Returns ``(prefix_mask, position_ids, logical_lengths)`` with shapes
    ``[B,N_q+S]``, ``[3,B,N_q+S]``, and ``[B]``. ``logical_lengths`` is the
    first answer-token position, i.e. ``question_length + S``.
    """
    if question_mask.ndim != 2:
        raise ValueError(
            f"question_mask must have shape [B,N_q], got {tuple(question_mask.shape)}"
        )
    if num_slots < 1:
        raise ValueError("num_slots must be >= 1")

    keep = question_mask > 0
    b = keep.shape[0]
    q_positions = keep.long().cumsum(dim=-1) - 1
    q_positions = q_positions.masked_fill(~keep, 0)
    q_lengths = keep.sum(dim=-1).long()  # [B]
    slot_offsets = torch.arange(num_slots, device=question_mask.device)
    slot_positions = q_lengths[:, None] + slot_offsets[None, :]  # [B,S]
    one_d = torch.cat([q_positions, slot_positions], dim=-1)  # [B,N_q+S]
    position_ids = one_d.unsqueeze(0).expand(3, b, -1)
    slot_mask = torch.ones(
        b, num_slots, dtype=question_mask.dtype, device=question_mask.device
    )
    prefix_mask = torch.cat([question_mask, slot_mask], dim=-1)
    return prefix_mask, position_ids, q_lengths + num_slots

# ---------------------------------------------------------------------------
# workspace blocks
# ---------------------------------------------------------------------------


class _Block(nn.Module):
    """Pre-norm block: optional cross-attention over ``context``, then FFN."""

    def __init__(self, d_w: int, n_heads: int, ffn_mult: int, dropout: float, cross: bool):
        super().__init__()
        self.self_norm = nn.LayerNorm(d_w)
        self.self_attn = nn.MultiheadAttention(
            d_w, n_heads, dropout=dropout, batch_first=True
        )
        self.cross = cross
        if cross:
            self.cross_norm_q = nn.LayerNorm(d_w)
            self.cross_norm_kv = nn.LayerNorm(d_w)
            self.cross_attn = nn.MultiheadAttention(
                d_w, n_heads, dropout=dropout, batch_first=True
            )
        self.ffn_norm = nn.LayerNorm(d_w)
        self.ffn = nn.Sequential(
            nn.Linear(d_w, ffn_mult * d_w),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(ffn_mult * d_w, d_w),
        )

    def forward(
        self,
        z: torch.Tensor,  # [B, S, d_w]
        context: torch.Tensor | None = None,  # [B, N, d_w]
        context_padding_mask: torch.Tensor | None = None,  # [B, N] True == ignore
        cond: torch.Tensor | None = None,  # [1|B, S, d_w] norm 입력 조건 (carrier 미반영)
    ) -> torch.Tensor:
        # persistent slot-ID(advisor): cond는 attention/FFN을 조건화할 뿐
        # residual carrier에는 절대 실리지 않는다 — z' = z + F(LN(z + E)).
        # cond=None이면 비트동일.
        zc = z if cond is None else z + cond
        h = self.self_norm(zc)
        z = z + self.self_attn(h, h, h, need_weights=False)[0]
        if self.cross:
            if context is None:
                raise ValueError("cross-attention block called without context")
            zc = z if cond is None else z + cond
            q = self.cross_norm_q(zc)
            kv = self.cross_norm_kv(context)
            z = z + self.cross_attn(
                q, kv, kv, key_padding_mask=context_padding_mask, need_weights=False
            )[0]
        zc = z if cond is None else z + cond
        return z + self.ffn(self.ffn_norm(zc))


class _LowRankMHA(nn.Module):
    """rank-r 부분공간 attention: Up(MHA(Down(q), Down(kv))) — d 유지, 파라미터 O(d·r)."""

    def __init__(self, d: int, r: int, n_heads: int):
        super().__init__()
        self.dq = nn.Linear(d, r, bias=False)
        self.dkv = nn.Linear(d, r, bias=False)
        self.attn = nn.MultiheadAttention(r, n_heads, dropout=0.0, batch_first=True)
        self.up = nn.Linear(r, d, bias=False)

    def forward(self, q, kv, key_padding_mask=None):
        out, _ = self.attn(self.dq(q), self.dkv(kv), self.dkv(kv),
                           key_padding_mask=key_padding_mask, need_weights=False)
        return self.up(out)


class _LowRankBlock(nn.Module):
    """_Block의 low-rank 동형: pre-norm + rank-r attention/FFN 잔차.

    ER v2 (advisor): d_w=d에서 full dense workspace의 파라미터 폭증을 막는
    "state width matching과 parameter explosion의 분리". cond 인터페이스는
    _Block과 동일 (미사용 시 비트동일 경로).
    """

    def __init__(self, d_w: int, n_heads: int, ffn_mult: int, dropout: float,
                 cross: bool, rank: int):
        super().__init__()
        self.self_norm = nn.LayerNorm(d_w)
        self.self_attn = _LowRankMHA(d_w, rank, n_heads)
        self.cross = cross
        if cross:
            self.cross_norm_q = nn.LayerNorm(d_w)
            self.cross_norm_kv = nn.LayerNorm(d_w)
            self.cross_attn = _LowRankMHA(d_w, rank, n_heads)
        self.ffn_norm = nn.LayerNorm(d_w)
        self.ffn = nn.Sequential(
            nn.Linear(d_w, ffn_mult * rank), nn.GELU(), nn.Linear(ffn_mult * rank, d_w),
        )

    def forward(self, z, context=None, context_padding_mask=None, cond=None):
        zc = z if cond is None else z + cond
        h = self.self_norm(zc)
        z = z + self.self_attn(h, h)
        if self.cross:
            if context is None:
                raise ValueError("cross-attention block called without context")
            zc = z if cond is None else z + cond
            q = self.cross_norm_q(zc)
            kv = self.cross_norm_kv(context)
            z = z + self.cross_attn(q, kv, key_padding_mask=context_padding_mask)
        zc = z if cond is None else z + cond
        return z + self.ffn(self.ffn_norm(zc))


def _make_block(config, cross: bool):
    r = int(getattr(config, "workspace_low_rank", 0) or 0)
    if r > 0:
        return _LowRankBlock(config.workspace_width, config.workspace_num_heads,
                             config.workspace_ffn_mult, config.workspace_dropout,
                             cross, r)
    return _Block(config.workspace_width, config.workspace_num_heads,
                  config.workspace_ffn_mult, config.workspace_dropout, cross)


class WorkspaceRead(nn.Module):
    """``r_theta``: condition learnable slots on ``Q*``, then read ``V*``.

    ``Z^(0) = r_theta(Z_init, Q*, V*)``. ``V*`` is consumed only inside this
    module: by :meth:`forward` (the initial read), and -- when
    ``config.read_gating`` is on -- by :meth:`reread`, which applies the SAME
    cross-attention weights to the evolved state under a scalar gate. With
    gating off the read is single by construction and the caller discards
    ``V*`` immediately after; either way no other module ever sees it.
    """

    def __init__(self, config: CloseQwen2_5_VLConfig):
        super().__init__()
        d, d_w = config.backbone_width, config.workspace_width
        self.slots = nn.Parameter(torch.randn(config.num_workspace_slots, d_w) * 0.02)
        _eye = d == d_w  # ER v2: 등폭이면 입사 projection(d→d) 자체가 불필요
        self.q_in = nn.Identity() if _eye else nn.Linear(d, d_w, bias=False)
        self.v_in = nn.Identity() if _eye else nn.Linear(d, d_w, bias=False)
        self.condition = _make_block(config, cross=True)
        self.read = nn.ModuleList(
            _make_block(config, cross=True) for _ in range(config.read_num_blocks)
        )
        # v2: image-dependent from step 0. A random cross-attention starts out
        # near-uniform, so its read is ~the mean of V* regardless of content and
        # the optimizer can settle on ignoring vision (probe 20143). The
        # shortcut injects an orthogonal projection of the LN'd pooled V* into
        # every slot, so ignoring the image is never the zero-gradient default.
        self.pool_in = None
        if getattr(config, "read_shortcut", False):
            self.pool_in = nn.Linear(d, d_w, bias=False)
            # NOTE: post_init()'s HF sweep has overwritten this orthogonal init
            # with N(0, initializer_range) in every run since v2 -- discovered
            # while diagnosing the gate-init wipe (smoke 20552). Deliberately
            # NOT flagged: v4/v5 trained under the overwritten init, and v6
            # must differ from v5 by the gate alone. Revisit only in a fresh
            # baseline generation.
            nn.init.orthogonal_(self.pool_in.weight)
        self.out_norm = nn.LayerNorm(d_w)
        # v6 candidate (read_gating): the SAME reader may be consulted again
        # between transitions, scaled by a scalar gate per item. b_gamma =
        # logit(eps): near-identity start (training begins ~at the proven
        # read-once regime; opening is earned by the losses). No new reader
        # parameters -- re-reads share v_in / read blocks / out_norm, so "how
        # it reads" is identical and only "when" is learned.
        self.gate = None
        if getattr(config, "read_gating", False):
            self.gate = nn.Linear(2 * d_w, 1)
            nn.init.zeros_(self.gate.weight)
            eps = float(getattr(config, "gate_init_eps", 0.02))
            eps = min(max(eps, 1e-6), 1 - 1e-6)
            nn.init.constant_(self.gate.bias, math.log(eps / (1 - eps)))
            # post_init()'s HF sweep re-inits every unflagged nn.Linear
            # (weight ~ N(0, initializer_range), bias -> 0), which silently
            # destroyed this init: with a random w over the unnormalised q̄
            # scale the gate saturated OPEN from step 0 (smoke 20552: gamma
            # 0.98 at the first log, 1.0 thereafter -- the exact opposite of
            # near-identity). The flag is respected by PreTrainedModel.
            # _initialize_weights and keeps the custom init authoritative.
            self.gate._is_hf_initialized = True
        # v7 (read_residual_pure): always-on PURE visual residual reread.
        #   A^(k) = XAttn(LN(U^(k)) -> queries, V* -> keys/values)
        #   Z^(k) = U^(k) + W_o A^(k)
        # No self-attention, no FFN, no output norm, no biases anywhere in the
        # branch, W_o zero-initialised (exact identity at init -- stronger than
        # v6's near-identity). Every atom of the update is a value vector of
        # V*: the battery on ckpt-5724 showed the gated reader's worth was
        # extra DEPTH, not vision; this branch cannot provide content-free
        # depth, so any benefit is visual by construction.
        self.pure_ln = None
        if (getattr(config, "read_residual_pure", False)
                or getattr(config, "final_verify", False)
                or getattr(config, "mid_read_step", 0)):
            self.pure_ln = nn.LayerNorm(d_w, bias=False)
            self.pure_attn = nn.MultiheadAttention(
                d_w, config.workspace_num_heads, batch_first=True, bias=False,
                kdim=d, vdim=d,
            )
            nn.init.zeros_(self.pure_attn.out_proj.weight)
            if getattr(config, "pure_warm_start", False):
                # copy visual-reading geometry from the initial reader: K/V
                # from v_in ([d_w, d] both), Q from the first read block's
                # query in-proj slice ([d_w, d_w]).
                with torch.no_grad():
                    self.pure_attn.k_proj_weight.copy_(self.v_in.weight)
                    self.pure_attn.v_proj_weight.copy_(self.v_in.weight)
                    qw = self.read[0].cross_attn.in_proj_weight[: config.workspace_width]
                    self.pure_attn.q_proj_weight.copy_(qw)
            # post_init()'s sweep visits every SUBMODULE; a flag on the parent
            # MHA does not shield its out_proj Linear (the gate-wipe lesson,
            # caught again by tests/test_pure_residual.py). Flag the leaf.
            self.pure_attn._is_hf_initialized = True
            self.pure_attn.out_proj._is_hf_initialized = True
            self.pure_ln._is_hf_initialized = True

    def forward(
        self,
        q_star: torch.Tensor,  # [B, N_q, d]
        v_star: torch.Tensor,  # [B, N_v, d]
        q_padding_mask: torch.Tensor | None = None,  # [B, N_q] True == ignore
        v_padding_mask: torch.Tensor | None = None,  # [B, N_v] True == ignore
    ) -> torch.Tensor:
        b = q_star.shape[0]
        z = self.slots.unsqueeze(0).expand(b, -1, -1)  # [B, S, d_w]
        q = self.q_in(q_star)  # [B, N_q, d_w]
        z = self.condition(z, context=q, context_padding_mask=q_padding_mask)

        if self.pool_in is not None:
            if v_padding_mask is None:
                pooled = v_star.mean(1)  # [B, d]
            else:
                keep = (~v_padding_mask).to(v_star.dtype).unsqueeze(-1)  # [B, N_v, 1]
                pooled = (v_star * keep).sum(1) / keep.sum(1).clamp_min(1.0)
            pooled = nn.functional.layer_norm(pooled, (pooled.shape[-1],))
            z = z + self.pool_in(pooled).unsqueeze(1)  # broadcast over slots

        # Dynamic resolution means N_v varies across a batch; V* arrives padded
        # to the batch maximum and the mask hides the filler.
        v = self.v_in(v_star)  # [B, N_v, d_w]
        for block in self.read:
            z = block(z, context=v, context_padding_mask=v_padding_mask)
        return self.out_norm(z)  # [B, S, d_w]

    def visual_residual(
        self,
        query_state: torch.Tensor,  # [B, S, d_w] -- query only; NOT added to
        v_star: torch.Tensor,  # [B, N_v, d]
        v_padding_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """``W_o XAttn(LN(q), V*)`` -- the raw visual residual, base-free.

        Separated from :meth:`reread_pure` so held-out interventions can pair
        an arbitrary query source with an arbitrary decoder base (advisor's
        query-only substitution)."""
        if self.pure_ln is None:
            raise RuntimeError("visual_residual() requires config.read_residual_pure=True")
        a, _ = self.pure_attn(
            self.pure_ln(query_state), v_star, v_star,
            key_padding_mask=v_padding_mask, need_weights=False,
        )
        return a

    def reread_pure(
        self,
        u: torch.Tensor,  # [B, S, d_w] provisional state
        v_star: torch.Tensor,  # [B, N_v, d]
        v_padding_mask: torch.Tensor | None = None,
        query_state: torch.Tensor | None = None,  # persistent slot-ID query
    ) -> torch.Tensor:
        """v7: ``Z = U + W_o XAttn(LN(U), V*)``. Update lives in span(V(V*)).

        ``query_state``가 주어지면 query만 그것으로 대체한다 (residual 기저는
        여전히 ``u``) — persistent slot identity가 ``LN(U+E_slot)``을 query로
        쓰되 state에는 E를 누적하지 않기 위한 통로. ``None``이면 비트동일."""
        q = u if query_state is None else query_state
        return u + self.visual_residual(q, v_star, v_padding_mask)

    def reread(
        self,
        z: torch.Tensor,  # [B, S, d_w] current workspace state
        q_star: torch.Tensor,  # [B, N_q, d]
        v_star: torch.Tensor,  # [B, N_v, d]
        q_padding_mask: torch.Tensor | None = None,
        v_padding_mask: torch.Tensor | None = None,
        gate_override: float | torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Gated consultation of ``V*`` by the EVOLVED workspace state.

        ``z' = z + g * (read(z, V*) - z)`` with scalar ``g`` per item, so
        ``g = 0`` reproduces ``z`` exactly (bit-identical to the no-reread
        path) and ``g = 1`` is a full re-read. The reader weights are shared
        with the initial read; only the gate decides WHEN looking again is
        worth it. Returns ``(z', g)`` with ``g`` of shape ``[B]``.

        ``gate_override`` substitutes the learned gate with a fixed value
        (forced-open 1.0 / forced-closed 0.0 / anything between). Held-out
        causal evaluation ONLY -- §3.2's gate interventions. Never set during
        training; ``None`` (the default) is the learned-gate path, bitwise.
        """
        if self.gate is None:
            raise RuntimeError("reread() requires config.read_gating=True")
        if q_padding_mask is None:
            q_bar = self.q_in(q_star).mean(1)  # [B, d_w]
        else:
            keep = (~q_padding_mask).to(q_star.dtype).unsqueeze(-1)
            q_bar = self.q_in((q_star * keep).sum(1) / keep.sum(1).clamp_min(1.0))
        # LN on the pooled state: the gate should react to the state's
        # direction, not to scale drift across steps/runs.
        z_bar = nn.functional.layer_norm(z.mean(1), (z.shape[-1],))  # [B, d_w]
        g = torch.sigmoid(self.gate(torch.cat([z_bar, q_bar], dim=-1))).squeeze(-1)  # [B]
        if gate_override is not None:
            g = torch.as_tensor(gate_override, dtype=g.dtype, device=g.device).expand_as(g)
            if float(g.max()) == 0.0:
                return z, g  # exact noop, mirrors the closed-gate residual form

        v = self.v_in(v_star)  # [B, N_v, d_w]
        upd = z
        for block in self.read:
            upd = block(upd, context=v, context_padding_mask=v_padding_mask)
        upd = self.out_norm(upd)
        return z + g.view(-1, 1, 1) * (upd - z), g


class WorkspaceTransition(nn.Module):
    """``f_theta``: shared recurrent step. Sees ``Z^(k-1)`` and ``Q*``. Never ``V*``.

    Shared across all ``K`` steps, so the trainable parameter count is
    independent of ``K`` (§3.2).
    """

    def __init__(self, config: CloseQwen2_5_VLConfig):
        super().__init__()
        d, d_w = config.backbone_width, config.workspace_width
        self.q_in = (nn.Identity() if d == d_w
                     else nn.Linear(d, d_w, bias=False))
        self.blocks = nn.ModuleList(
            _make_block(config, cross=True) for _ in range(config.transition_num_blocks)
        )
        self.out_norm = nn.LayerNorm(d_w)

    def forward(
        self,
        z: torch.Tensor,  # [B, S, d_w]
        q_star: torch.Tensor,  # [B, N_q, d]
        q_padding_mask: torch.Tensor | None = None,
        cond: torch.Tensor | None = None,  # persistent slot-ID (norm 입력 전용)
    ) -> torch.Tensor:
        q = self.q_in(q_star)  # [B, N_q, d_w]
        for block in self.blocks:
            z = block(z, context=q, context_padding_mask=q_padding_mask, cond=cond)
        return self.out_norm(z)  # [B, S, d_w]


class RawEvidenceQuestionTransition(nn.Module):
    """Shared recurrent update for the slot-free ``E=V*; R0=Q*`` path.

    The state has the same rows and width as the text-only question scaffold.
    Each step reads the complete immutable visual memory with one pre-norm
    cross-attention.  The residual-compatible mode used by existing runs is::

        R_next = R + CrossAttn(RMSNorm(R), RMSNorm(E), RMSNorm(E))

    The state-replacement candidate removes that recurrent identity path::

        H_next = CrossAttn(RMSNorm(H), RMSNorm(E), RMSNorm(E))

    Its question-anchored form replaces the one-time decoder addition with a
    fixed scaffold at every recurrent step::

        Z_next = Q* + CrossAttn(RMSNorm(Z), RMSNorm(E), RMSNorm(E))

    There is no self-attention or FFN in either current method: ``Q*`` is already
    contextualized by the frozen lower backbone, and the upper decoder retains
    its pretrained attention/MLP stack.  Consequently every recurrent state
    change is caused by a visual-memory read.  The same module instance is
    reused for every recurrent step.
    """

    def __init__(self, config: CloseQwen2_5_VLConfig):
        super().__init__()
        text_cfg = config.text_config
        d = config.backbone_width
        self.num_heads = int(text_cfg.num_attention_heads)
        self.num_key_value_heads = int(text_cfg.num_key_value_heads)
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.head_dim = d // self.num_heads
        self.dropout = float(config.workspace_dropout)
        if self.num_heads % self.num_key_value_heads:
            raise ValueError("query heads must be divisible by KV heads")

        # Native Qwen GQA shapes. Q/K/V/O and the input RMSNorm are copied
        # from L(ell_star+1) after post_init, before upper-layer LoRA injection.
        self.read_norm_q = Qwen2RMSNorm(d, eps=text_cfg.rms_norm_eps)
        self.read_norm_e = Qwen2RMSNorm(d, eps=text_cfg.rms_norm_eps)
        self.q_proj = nn.Linear(d, self.num_heads * self.head_dim, bias=True)
        self.k_proj = nn.Linear(
            d, self.num_key_value_heads * self.head_dim, bias=True
        )
        self.v_proj = nn.Linear(
            d, self.num_key_value_heads * self.head_dim, bias=True
        )
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, d, bias=False)
        self.state_replacement = bool(
            getattr(config, "raw_state_replacement", False)
        )
        self.question_anchored = bool(
            getattr(config, "raw_question_anchor", False)
        )
        # Compatibility only: old raw-E/Q checkpoints used a random d->d->d
        # FFN after every read and did not serialize a feature flag. Their
        # configs resolve raw_transition_ffn=None to True; all new launches
        # explicitly set it False and allocate neither module nor parameters.
        self.use_ffn = bool(getattr(config, "raw_transition_ffn", True))
        if self.use_ffn:
            self.ffn_norm = Qwen2RMSNorm(d, eps=text_cfg.rms_norm_eps)
            self.ffn = nn.Sequential(
                nn.Linear(d, config.workspace_ffn_mult * d),
                nn.GELU(),
                nn.Dropout(config.workspace_dropout),
                nn.Linear(config.workspace_ffn_mult * d, d),
            )

    def forward(
        self,
        r: torch.Tensor,  # [B,N_q,d]
        evidence: torch.Tensor | None = None,  # [B,N_v,d] == raw V*
        evidence_padding_mask: torch.Tensor | None = None,
        question_padding_mask: torch.Tensor | None = None,
        read_scale: float = 1.0,  # held-out intervention only
        projected_evidence: tuple[torch.Tensor, torch.Tensor] | None = None,
        question_anchor: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Apply one recurrent update.

        ``projected_evidence`` contains head-split K/V tensors returned by
        :meth:`project_evidence`.  Since raw E is immutable across recurrent
        steps, the normal rollout computes these once and reuses them.  This is
        algebraically identical to calling ``nn.MultiheadAttention`` four
        times, but avoids repeating the dominant ``N_v x d x d`` K/V
        projections for every step.
        """
        if projected_evidence is None:
            if evidence is None:
                raise ValueError("evidence or projected_evidence is required")
            projected_evidence = self.project_evidence(evidence)
        k, v = projected_evidence  # [B,H,N_v,D_h]

        q_norm = self.read_norm_q(r)
        d = q_norm.shape[-1]
        q = self._split_heads(self.q_proj(q_norm), self.num_heads)

        attn_mask = None
        if evidence_padding_mask is not None:
            # SDPA's additive mask uses -inf for keys that must be ignored.
            attn_mask = torch.zeros(
                evidence_padding_mask.shape[0],
                1,
                1,
                evidence_padding_mask.shape[1],
                dtype=q.dtype,
                device=q.device,
            ).masked_fill(
                evidence_padding_mask[:, None, None, :], float("-inf")
            )
        delta = nn.functional.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=attn_mask,
            dropout_p=(self.dropout if self.training else 0.0),
            is_causal=False,
            enable_gqa=True,
        )
        delta = delta.transpose(1, 2).contiguous().view(r.shape[0], r.shape[1], d)
        delta = self.o_proj(delta)
        if self.state_replacement:
            r_next = float(read_scale) * delta
            if self.question_anchored:
                if question_anchor is None:
                    raise ValueError(
                        "question_anchor is required by raw_question_anchor"
                    )
                # Preserve Q* but replace the visual response.  There is no
                # R_{k-1} residual, so the old identity carrier stays removed.
                r_next = question_anchor + r_next
        else:
            r_next = r + float(read_scale) * delta
        if self.use_ffn:
            r_next = r_next + self.ffn(self.ffn_norm(r_next))
        if question_padding_mask is not None:
            # Padding rows are physical cache plumbing only.  Keeping them
            # unchanged makes batched and per-item recurrent states agree.
            padding_value = (
                question_anchor
                if self.question_anchored and question_anchor is not None
                else r
            )
            r_next = torch.where(
                question_padding_mask.unsqueeze(-1), padding_value, r_next
            )
        return r_next

    def _split_heads(self, x: torch.Tensor, num_heads: int) -> torch.Tensor:
        b, n, d = x.shape
        return x.view(b, n, num_heads, d // num_heads).transpose(1, 2)

    def project_evidence(
        self, evidence: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Normalize and project immutable E once for a recurrent rollout."""
        e = self.read_norm_e(evidence)
        k = self.k_proj(e)
        v = self.v_proj(e)
        return (
            self._split_heads(k, self.num_key_value_heads),
            self._split_heads(v, self.num_key_value_heads),
        )


#: Backbone projections the post-``l*`` adapter attaches to.
ADAPTER_TARGET_MODULES = (
    "q_proj",
    "k_proj",
    "v_proj",
    "o_proj",
    "gate_proj",
    "up_proj",
    "down_proj",
)


def build_adapter_config(config: CloseQwen2_5_VLConfig):
    """LoRA config restricted to the decoder layers above ``l*``.

    peft rather than a hand-rolled adapter: a hook-attached delta holds a
    closure over its target module, which is fragile under ``device_map``
    sharding and FSDP -- exactly the settings this will train in. peft's
    ``layers_to_transform`` expresses the "above ``l*`` only" restriction
    directly, and its save/load path is already Trainer-compatible.
    """
    from peft import LoraConfig

    upper = config.upper_slice
    excl = set(getattr(config, "adapter_exclude_layers", []) or [])
    if getattr(config, "recurrence_only_adapter", False):
        target_layers = [int(config.ell_star) + 1]
    else:
        target_layers = [
            li for li in range(upper.start, upper.stop) if li not in excl
        ]
    return LoraConfig(
        r=config.adapter_rank,
        lora_alpha=config.adapter_alpha,
        lora_dropout=config.adapter_dropout,
        bias="none",
        target_modules=list(ADAPTER_TARGET_MODULES),
        layers_to_transform=target_layers,
        layers_pattern="layers",
        task_type=None,
    )


# ---------------------------------------------------------------------------
# output type
# ---------------------------------------------------------------------------


@dataclass
class CloseCausalLMOutput(ModelOutput):
    """Adds the workspace trajectory, which the Trainer needs for ``L_traj``."""

    loss: torch.FloatTensor | None = None
    logits: torch.FloatTensor | None = None
    past_key_values: Cache | None = None
    workspace_states: tuple[torch.FloatTensor, ...] | None = None  # K+1 x [B, S, d_w]
    #: LN'd masked-mean of V*, detached -- the target side of the Trainer's
    #: read-grounding InfoNCE. Never fed back into the model.
    v_star_pooled: torch.FloatTensor | None = None  # [B, d]
    #: bridge-only interface distillation 항 (advisor): 1-cos(LN h_R, LN h_base).
    #: 학습 시에만 채워짐. gradient는 workspace_out·splice_gain·l*+1층 adapter 한정.
    iface_loss: torch.FloatTensor | None = None
    #: Training-only normalized state-restoration loss.  Its target is the
    #: image-conditioned question state from the frozen multimodal lower pass;
    #: inference never receives that state.
    functional_state_loss: torch.FloatTensor | None = None
    #: Training-only score gap for the matched residual-swap intervention:
    #: normalized log p(y_i | B_i + C_i) - log p(y_i | B_i + C_donor).
    residual_swap_score_gap: torch.FloatTensor | None = None  # [B]
    residual_swap_valid: torch.BoolTensor | None = None  # [B]
    #: 1-indexed CRR transition selected by the transition-aware intervention.
    #: Aggregate-swap checkpoints and ordinary inference leave this unset.
    residual_swap_transition: torch.LongTensor | None = None  # scalar
    #: Training-only decoder NLLs for the selected hard-case corrective
    #: transition.  ``current`` remains differentiable; ``previous`` and
    #: ``swap`` are comparator values detached by the Trainer so the objective
    #: cannot win by degrading either baseline.
    corrective_current_nll: torch.FloatTensor | None = None  # [B]
    corrective_previous_nll: torch.FloatTensor | None = None  # [B]
    corrective_swap_nll: torch.FloatTensor | None = None  # [B]
    corrective_valid: torch.BoolTensor | None = None  # [B]
    #: Training-only gold-answer NLL at each decoded CRR prefix
    #: Z_k = B + mean(C_1..C_k).  Every column is scored by the same actual
    #: upper decoder; the Trainer applies a one-sided no-regression objective.
    step_answer_nll: torch.FloatTensor | None = None  # [B, T]
    step_answer_valid: torch.BoolTensor | None = None  # [B, T]
    #: RL-only greedy generations sampled from a stochastic CRR trajectory.
    #: Ordinary SFT/inference calls leave every field below unset.
    generated_ids: torch.LongTensor | None = None  # [B, <=N_a]
    #: Frozen-decoder, token-normalized gold-answer log likelihood for a
    #: sampled CRR trajectory.  This is an environment score only: it is
    #: computed under no_grad and cannot form a pathwise decoder gradient.
    latent_policy_gold_logprob: torch.FloatTensor | None = None  # [B]
    #: The same frozen gold-answer score at every decoded CRR prefix
    #: Z_k = B + mean(C_1..C_k).  This is populated only when latent RL asks
    #: for transition-aware credit assignment; its last column is exactly
    #: ``latent_policy_gold_logprob``.
    latent_policy_step_gold_logprob: torch.FloatTensor | None = None  # [B, T]
    #: Dimension-normalized Gaussian score log-density for C2..CT.
    latent_policy_log_probs: torch.FloatTensor | None = None  # [B, T-1]
    latent_policy_mean_update_rms: torch.FloatTensor | None = None  # [B, T-1]
    latent_policy_noise_rms: torch.FloatTensor | None = None  # [B, T-1]


@dataclass
class LatentPolicyTrace:
    """Training-only score-function trace for stochastic CRR transitions.

    The sampled recurrent states themselves are detached.  Consequently the
    only differentiable quantity returned to the RL Trainer is ``log_probs``;
    answer generation and reward computation cannot create a pathwise gradient
    through the frozen decoder.
    """

    log_probs: torch.FloatTensor  # [B, T-1]
    mean_update_rms: torch.FloatTensor  # [B, T-1], relative to C1 RMS
    noise_rms: torch.FloatTensor  # [B, T-1], relative to C1 RMS


class CRRLatentPolicyAdapter(nn.Module):
    """Zero-init low-rank correction used only by CRR states ``C2..CT``.

    This is intentionally narrower than adapting the shared native L21: L21
    also constructs B/C1 and processes generated answer tokens.  A separate
    latent-only adapter keeps those competence-critical paths frozen and makes
    the score-function action policy the only behavior changed by RL.
    """

    def __init__(self, width: int, rank: int, alpha: float):
        super().__init__()
        if rank < 1:
            raise ValueError("CRRLatentPolicyAdapter rank must be >= 1")
        self.down = nn.Linear(width, rank, bias=False)
        self.up = nn.Linear(rank, width, bias=False)
        self.scaling = float(alpha) / float(rank)

    def forward(self, state: torch.Tensor) -> torch.Tensor:
        return state + self.up(self.down(state)) * self.scaling


class CRRFeaturewiseAggregator(nn.Module):
    """Native-basis-preserving linear aggregation over recurrent depth.

    This is the block-diagonal special case of ``Linear(T * d, d)``: block
    ``k`` is a learned diagonal matrix instead of a dense ``d x d`` matrix.
    It can therefore select a different recurrent depth for every native
    hidden feature without rotating the pretrained decoder basis or adding a
    51M-parameter all-reduce at ``d=3584, T=4``.
    """

    def __init__(self, steps: int, width: int):
        super().__init__()
        if steps < 1 or width < 1:
            raise ValueError("steps and width must both be positive")
        self.steps = int(steps)
        self.width = int(width)
        self.weight = nn.Parameter(torch.empty(self.steps, self.width))
        self.reset_parameters()

    def reset_parameters(self) -> None:
        # Exactly reproduces the historical mean interface at initialization.
        with torch.no_grad():
            self.weight.fill_(1.0 / self.steps)

    def forward(self, concatenated: torch.Tensor) -> torch.Tensor:
        expected = self.steps * self.width
        if concatenated.shape[-1] != expected:
            raise ValueError(
                f"expected concatenated width {expected}, got "
                f"{concatenated.shape[-1]}"
            )
        states = concatenated.reshape(
            *concatenated.shape[:-1], self.steps, self.width
        )
        return (states * self.weight).sum(dim=-2)


# ---------------------------------------------------------------------------
# the model
# ---------------------------------------------------------------------------


class CloseQwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel):
    def _init_weights(self, module):
        super()._init_weights(module)
        if isinstance(module, SpatialVisualRecurrentCell):
            # Generic HF initialization visits the child Linears first.  Zero
            # the residual writer only after that sweep so T=8 is the exact
            # base visual path at initialization and on missing-key loads.
            module.zero_output()
        if isinstance(module, CRRFeaturewiseAggregator):
            # Also runs for missing-key initialization under from_pretrained.
            module.reset_parameters()
        # ER v2.2 C: aggregation 파라미터 init. post_init 스윕(from_backbone)과
        # missing-key 초기화(from_pretrained) 모두 이 훅을 지나므로 두 로드
        # 경로에서 동일 값이 보장된다 (__init__ 직접 대입은 meta 경유에서 유실).
        if module is self and getattr(
            self.config, "concat_aggregation", False
        ) and hasattr(self, "reasoning_aggregator"):
            # 개정 init (advisor 2026-08-12): W = [I/T, ..., I/T] — step 0에서
            # R_agg = mean(R1..RT) 정확히. (구판 [0,..,0,I]는 폐기; random 금지.)
            # I/4 = 0.25는 bf16 정확 표현 — 반올림 오차 없음. 네 블록은 이후
            # answer CE로 독립 분화.
            _T = int(self.config.num_workspace_steps)
            d_w = int(self.reasoning_aggregator.out_features)
            with torch.no_grad():
                w = self.reasoning_aggregator.weight  # [D, T*D]
                w.zero_()
                _eye = torch.eye(d_w, dtype=w.dtype, device=w.device) / _T
                for _i in range(_T):
                    w[:, _i * d_w:(_i + 1) * d_w].copy_(_eye)
        if module is self and getattr(self.config, "temporal_aggregation", False) \
                and hasattr(self, "temporal_logits"):
            _T = int(self.config.num_workspace_steps)
            _w0 = torch.tensor([0.4, 0.3, 0.2, 0.1][:_T]) if _T <= 4 \
                else torch.full((_T,), 1.0 / _T)
            _w0 = _w0 / _w0.sum()
            with torch.no_grad():
                self.temporal_logits.copy_(_w0.log().to(self.temporal_logits.dtype))
        if module is self and hasattr(self, "crr_policy_adapter"):
            # Exact identity at the SFT -> latent-RL handoff.  Down follows the
            # backbone initializer; zero Up is the standard LoRA-style start.
            with torch.no_grad():
                self.crr_policy_adapter.up.weight.zero_()

    """CLOSE: Qwen2.5-VL with its upper visual path replaced by a read-once workspace."""

    config_class = CloseQwen2_5_VLConfig
    _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"]


    def __init__(self, config: CloseQwen2_5_VLConfig, backbone=None):
        super().__init__(config)
        config.validate()
        spatial_visual = bool(
            getattr(config, "spatial_visual_recurrence", False)
        )
        perceive_deliberate = bool(
            getattr(config, "perceive_deliberate_chain", False)
        )
        local_visual_chain = bool(
            getattr(config, "local_visual_latent_chain", False)
        )
        if not spatial_visual:
            config._require_ell_star()

        # The host model. Pass an already-loaded one via :meth:`from_backbone`;
        # `CloseQwen2_5_VLForConditionalGeneration.from_pretrained(<base repo>)`
        # does NOT work, because base checkpoint keys are "model.*" / "lm_head.*"
        # while this tree needs "backbone.model.*". transformers only warns about
        # that, leaving all 8.3B parameters random and every forward NaN.
        self.backbone = (
            backbone
            if backbone is not None
            else Qwen2_5_VLForConditionalGeneration._from_config(config)
        )

        raw_eq = bool(getattr(config, "raw_evidence_question", False))
        crr = bool(
            getattr(config, "counterfactual_residual_recurrence", False)
        )
        q_state_interface = raw_eq or crr
        if perceive_deliberate:
            # LVLC uses one homogeneous latent type; historical PDLC uses two
            # LOOK/THINK types.  Distinct native M-RoPE positions carry order
            # in both one-pass graphs.
            self.latent_chain_embeddings = nn.Embedding(
                1 if local_visual_chain else 2,
                config.backbone_width,
            )
            self.workspace_out = nn.Identity()
        elif spatial_visual:
            self.spatial_visual_cell = SpatialVisualRecurrentCell(
                config.backbone_width,
                int(config.spatial_recurrence_inner_width),
                int(config.spatial_recurrence_heads),
                rms_norm_eps=float(config.text_config.rms_norm_eps),
                residual_scale=float(config.counterfactual_beta),
            )
            self.workspace_out = nn.Identity()
        elif raw_eq:
            # Replacement, not an addition: the raw-evidence path allocates no
            # slot reader, slot transition, evidence slots, reasoning slots,
            # mid-read module, or learned trajectory aggregator.
            self.raw_eq_transition = RawEvidenceQuestionTransition(config)
            self.workspace_out = nn.Identity()
        elif crr:
            # CRR reuses the native L(ell_star+1) module itself.  No copied
            # transition, visual memory, slots, projection, or aggregation is
            # allocated.  RL checkpoints may add one shared low-rank correction
            # that is called only for C2..CT, never for B/C1/token decoding.
            self.workspace_out = nn.Identity()
            if getattr(config, "crr_concat_aggregation", False):
                width = int(config.backbone_width)
                steps = int(config.num_workspace_steps)
                self.reasoning_aggregator = CRRFeaturewiseAggregator(
                    steps, width
                )
            if int(getattr(config, "crr_policy_rank", 0)) > 0:
                self.crr_policy_adapter = CRRLatentPolicyAdapter(
                    config.backbone_width,
                    int(config.crr_policy_rank),
                    float(config.crr_policy_alpha),
                )
        else:
            self.workspace_read = WorkspaceRead(config)
            self.workspace_transition = WorkspaceTransition(config)
            self.workspace_out = nn.Linear(
                config.workspace_width, config.backbone_width
            )
        # -- ER v2 (explicit evidence read; advisor 2026-08-11) --------------
        # R이 E를 별도 memory로 읽는 pre-norm cross-attn 잔차 + 질문 조건화 R^0.
        if getattr(config, "explicit_evidence_read", False):
            # ER v2 canonical (§5): R이 persistent E를 읽는 명시적 cross-attention.
            # query=LN(R), key/value=LN(E). 전 step 공유 (step별 복제 금지, §14).
            assert config.workspace_width == config.backbone_width, (
                "explicit_evidence_read는 native width(d_w == backbone d) 전제 (§2)")
            d_w = config.workspace_width
            self.evidence_read_norm_q = nn.LayerNorm(d_w)
            self.evidence_read_norm_kv = nn.LayerNorm(d_w)
            self.evidence_read = nn.MultiheadAttention(
                d_w, config.workspace_num_heads, dropout=0.0, batch_first=True)
            # §8: 차원-변환 전용 P_Z는 native width에서 Identity.
            # interface 보정 수신자 = splice_gain + post-l* LoRA (§9.5에서 검증).
            self.workspace_out = nn.Identity()
        if getattr(config, "r_init_question", False):
            d_w = config.workspace_width
            # Legacy config name, current semantics: R^0 is learnable slots
            # only. Q* first enters through the shared transition.
            self.reasoning_slots = nn.Parameter(
                torch.randn(config.num_workspace_slots, d_w) * 0.02
            )
        if getattr(config, "concat_aggregation", False):
            # ER v2.2 C 확정판 (§6-8): R_agg = Linear(Concat(R1..RT, dim=hidden)).
            # W ∈ [D, T·D], bias 없음. 초기 W=[I/T,...,I/T], so
            # R_agg=mean(R1..RT). Temporal softmax/gates are absent.
            assert getattr(config, "er_single_chain", False), (
                "concat_aggregation은 er_single_chain=True 전제 (§12)")
            _T = int(config.num_workspace_steps)
            d_w = config.workspace_width
            self.reasoning_aggregator = nn.Linear(_T * d_w, d_w, bias=False)
            # 값(zero + 마지막 블록 Identity)은 _init_weights(self)에서 설정 —
            # __init__ 직접 대입은 from_pretrained meta 경유에서 유실 (실증됨).
        if getattr(config, "temporal_aggregation", False):
            # (legacy — softmax 변형; C 확정판에서 미사용, §7 금지 목록)
            assert getattr(config, "er_single_chain", False), (
                "temporal_aggregation은 er_single_chain=True 전제 (§12)")
            assert not getattr(config, "concat_aggregation", False), "동시 사용 금지"
            _T = int(config.num_workspace_steps)
            self.temporal_logits = nn.Parameter(torch.zeros(_T))
        # Placeholder embedding for the S slot positions in the text-only prefill.
        # Cache-shape plumbing below ell*: these rows are overwritten by R_agg
        # before the upper decoder. They are a fixed text-only placeholder, not
        # an extra learned prompt outside the declared workspace/LoRA method.
        if not q_state_interface and not spatial_visual and not perceive_deliberate:
            self.slot_embedding = nn.Parameter(
                torch.randn(config.num_workspace_slots, config.backbone_width) * 0.02,
                requires_grad=False,
            )
        # RMS splice gain is retained for historical interface-loss checkpoints.
        # In the current answer-only method it is a fixed 1, not an undeclared
        # trainable decoder-interface scalar (and bf16 updates near 1 would be
        # numerically ineffective at the configured learning rate anyway).
        self.splice_gain = (
            nn.Parameter(torch.ones(1), requires_grad=bool(config.interface_loss))
            if (
                not q_state_interface
                and not spatial_visual
                and not perceive_deliberate
                and getattr(config, "splice_scale_match", False)
            )
            else None
        )

        self.post_init()
        if perceive_deliberate:
            init_id = int(getattr(config, "latent_chain_init_token_id", -1))
            if init_id < 0:
                init_id = 0
            with torch.no_grad():
                source = self.vl.get_input_embeddings().weight[init_id]
                self.latent_chain_embeddings.weight.copy_(
                    source.unsqueeze(0).expand_as(
                        self.latent_chain_embeddings.weight
                    )
                )
        if raw_eq:
            self._initialize_raw_attention_from_upper_layer()

        # peft wraps the targeted Linears in place, so this must run after
        # post_init: the adapters start as an exact no-op (B initialised to 0)
        # and must not be re-initialised by the backbone's weight init.
        self.adapter_config = None
        if config.adapter_rank > 0:
            from peft import inject_adapter_in_model

            self.adapter_config = build_adapter_config(config)
            # Inject at `backbone`, not `text_model`: peft matches layer indices
            # with `.*\.layers\.(\d+)\.`, which needs a dot before "layers".
            # From text_model the keys start "layers.0..." and nothing matches;
            # from backbone they are "model.language_model.layers.0...".
            inject_adapter_in_model(self.adapter_config, self.backbone)
            if getattr(config, "recurrence_only_adapter", False):
                # Execution defaults to the frozen pretrained path.  The
                # recurrent loop temporarily enables these same trainable
                # parameters for k>=2 without mutating requires_grad.
                from peft.tuners.tuners_utils import BaseTunerLayer

                for module in self.backbone.modules():
                    if isinstance(module, BaseTunerLayer):
                        module._disable_adapters = True

    def _initialize_raw_attention_from_upper_layer(self) -> None:
        """Copy the first upper decoder layer's visual-read-compatible basis.

        L(ell_star+1) is pretrained to consume layer-ell_star hidden states,
        exactly the feature space occupied by both Q* and V*.  Cross-attention
        remains position-free, but its Q/K/V/O maps and adjacent RMSNorms start
        from that decoder layer instead of an unrelated random basis.

        This runs after ``post_init`` (so HF initialization cannot overwrite
        the copy) and before LoRA injection (so source projections are plain
        Linears and the raw transition receives only the frozen base weights).
        """
        source_idx = int(getattr(self.config, "raw_attention_source_layer", -1))
        if source_idx < 0:
            source_idx = int(self.config.ell_star) + 1
        source_layer = self.text_model.layers[source_idx]
        source_attn = source_layer.self_attn
        target = self.raw_eq_transition
        pairs = (
            (target.q_proj, source_attn.q_proj),
            (target.k_proj, source_attn.k_proj),
            (target.v_proj, source_attn.v_proj),
            (target.o_proj, source_attn.o_proj),
        )
        with torch.no_grad():
            for dst, src in pairs:
                if dst.weight.shape != src.weight.shape:
                    raise ValueError(
                        f"raw GQA/source shape mismatch at L{source_idx}: "
                        f"{tuple(dst.weight.shape)} vs {tuple(src.weight.shape)}"
                    )
                dst.weight.copy_(src.weight)
                if dst.bias is not None:
                    if src.bias is None:
                        raise ValueError("raw GQA expects the source Q/K/V biases")
                    dst.bias.copy_(src.bias)
            target.read_norm_q.weight.copy_(source_layer.input_layernorm.weight)
            target.read_norm_e.weight.copy_(source_layer.input_layernorm.weight)
            if target.use_ffn:
                target.ffn_norm.weight.copy_(
                    source_layer.post_attention_layernorm.weight
                )
        self.raw_attention_source_layer = source_idx

    @classmethod
    def from_backbone(cls, pretrained: str, config: CloseQwen2_5_VLConfig, **kwargs):
        """Build CLOSE around a real Qwen2.5-VL checkpoint.

        The backbone is loaded first and handed to ``__init__`` so its weights
        are never routed through this class's key namespace. Use this, not
        ``from_pretrained``, when starting from the base model;
        ``from_pretrained`` remains correct for reloading a *saved CLOSE*
        checkpoint, whose keys already match.
        """
        backbone = Qwen2_5_VLForConditionalGeneration.from_pretrained(pretrained, **kwargs)
        model = cls(config, backbone=backbone)
        return model.to(backbone.dtype)

    # -- convenience views -------------------------------------------------

    @property
    def vl(self):
        """``Qwen2_5_VLModel`` (visual + language_model)."""
        return self.backbone.model

    @property
    def text_model(self):
        """``Qwen2_5_VLTextModel``."""
        return self.backbone.model.language_model

    def _uses_question_state_interface(self) -> bool:
        """Whether the replacement prefix has exactly the question rows."""
        return bool(
            getattr(self.config, "raw_evidence_question", False)
            or getattr(
                self.config, "counterfactual_residual_recurrence", False
            )
        )

    def _upper_decoder_start(self) -> int:
        """First layer that consumes the image-dependent decoder latent.

        CRR turns L(ell_star+1) into its shared recurrent operator.  That layer
        still processes answer tokens, but its prefix KV is populated only by
        the text-only scaffold; the image-dependent prefix starts at the next
        layer.  Historical methods continue to splice at ell_star+1.
        """
        return int(self.config.ell_star) + (
            2
            if getattr(
                self.config, "counterfactual_residual_recurrence", False
            )
            else 1
        )

    # -- parameter freezing ------------------------------------------------

    def freeze_backbone(self) -> dict[str, int]:
        """Freeze everything, then open exactly the declared method parameters.

        §3.3: "Vision encoder와 기본 backbone은 동결하고 workspace module,
        projection, post-l* low-rank adapter만 학습".

        LoRA parameters live *inside* ``backbone`` after peft injection, so they
        are re-enabled by name after the blanket freeze.
        """
        for p in self.parameters():
            p.requires_grad_(False)
        local_visual_chain = bool(
            getattr(self.config, "local_visual_latent_chain", False)
        )
        # LVLC executes solely through the native Transformer, its homogeneous
        # latent embedding, and LoRA.  Opening legacy recurrent workspace
        # modules here creates dead trainable parameters that never participate
        # in the one-pass graph.
        method_modules = (
            (getattr(self, "latent_chain_embeddings", None),)
            if local_visual_chain
            else (
                getattr(self, "latent_chain_embeddings", None),
                getattr(self, "spatial_visual_cell", None),
                getattr(self, "workspace_read", None),
                getattr(self, "workspace_transition", None),
                getattr(self, "raw_eq_transition", None),
                self.workspace_out,
            )
        )
        for module in method_modules:
            if module is None:
                continue
            for p in module.parameters():
                p.requires_grad_(True)
        if not local_visual_chain:
            for name in (
                "evidence_read_norm_q",
                "evidence_read_norm_kv",
                "evidence_read",
                "reasoning_aggregator",
            ):
                module = getattr(self, name, None)
                if module is not None:
                    for p in module.parameters():
                        p.requires_grad_(True)
            for name in ("reasoning_slots", "temporal_logits"):
                p = getattr(self, name, None)
                if p is not None:
                    p.requires_grad_(True)
            if self.splice_gain is not None and getattr(
                self.config, "interface_loss", False
            ):
                self.splice_gain.requires_grad_(True)

        n_lora = 0
        for name, p in self.backbone.named_parameters():
            if "lora_" in name:
                p.requires_grad_(True)
                n_lora += p.numel()

        n_train = sum(p.numel() for p in self.parameters() if p.requires_grad)
        n_total = sum(p.numel() for p in self.parameters())
        if local_visual_chain:
            latent_parameter = self.latent_chain_embeddings.weight
            expected_trainable = n_lora + latent_parameter.numel()
            if not latent_parameter.requires_grad or n_train != expected_trainable:
                raise RuntimeError(
                    "LVLC must train exactly its latent embedding and LoRA: "
                    f"latent_requires_grad={latent_parameter.requires_grad}, "
                    f"latent_numel={latent_parameter.numel()}, lora={n_lora}, "
                    f"actual_trainable={n_train}, expected={expected_trainable}"
                )
        return {"trainable": n_train, "total": n_total, "lora": n_lora}

    def _aggregate_crr_residuals(
        self, residual_steps: list[torch.Tensor] | tuple[torch.Tensor, ...]
    ) -> torch.Tensor:
        """Aggregate one CRR residual prefix with a stable initial interface.

        The learned full-horizon path is ``Linear(Concat(C1..CT))``.  A shorter
        evaluation/training prefix zero-pads its missing future blocks and is
        multiplied by ``T/k``; at the prescribed ``[I/T,...,I/T]``
        initialization this remains exactly ``mean(C1..Ck)``.  Historical CRR
        checkpoints keep their parameter-free mean because the flag defaults
        to false.
        """
        steps = list(residual_steps)
        if not steps:
            raise ValueError("CRR aggregation requires at least one residual")
        horizon = int(self.config.num_workspace_steps)
        if len(steps) > horizon:
            raise ValueError(
                f"received {len(steps)} CRR states for horizon {horizon}"
            )
        if not getattr(self.config, "crr_concat_aggregation", False):
            return torch.stack(steps, dim=0).mean(dim=0)

        if not hasattr(self, "reasoning_aggregator"):
            raise RuntimeError(
                "crr_concat_aggregation=True but reasoning_aggregator is missing"
            )
        if len(steps) < horizon:
            steps.extend(torch.zeros_like(steps[0]) for _ in range(horizon - len(steps)))
        aggregate = self.reasoning_aggregator(torch.cat(steps, dim=-1))
        prefix_steps = len(residual_steps)
        if prefix_steps < horizon:
            aggregate = aggregate * (horizon / prefix_steps)
        return aggregate

    def decoder_latent_from_states(
        self, states: list[torch.Tensor] | tuple[torch.Tensor, ...]
    ) -> torch.Tensor:
        """Return the exact pre-interface latent decoded by the current model.

        Keeping this operation on the model prevents evaluation scripts from
        silently targeting a stored state that the selected interface does not
        actually decode. Residual Raw-E/Q returns ``R_T``; state replacement
        returns ``Q* + H_T``; question-anchored replacement returns ``Z_T``;
        historical concat models reconstruct their learned ``R_agg``.
        """
        if not states:
            raise ValueError("states must not be empty")
        if getattr(
            self.config, "counterfactual_residual_recurrence", False
        ):
            # states = [B, C1, ..., CT], where B=F(Q_text) is strictly
            # image-independent and every Ck is an image-conditioned residual.
            if len(states) != self.config.num_workspace_steps + 1:
                raise ValueError(
                    f"expected anchor plus {self.config.num_workspace_steps} "
                    f"counterfactual residuals, got {len(states)} tensors"
                )
            residual_steps = list(states[1:])
            # Evaluation-only prefix sweep.  Step k uses the same interface
            # trained by L_step_keep: Z_k = B + mean(C_1..C_k), rather than
            # decoding C_k alone.  Normal training/inference has no override.
            k_ov = int(getattr(self.config, "er_step_override", 0) or 0)
            if k_ov:
                if not 1 <= k_ov <= self.config.num_workspace_steps:
                    raise ValueError(
                        "er_step_override must be in "
                        f"[1,{self.config.num_workspace_steps}]"
                    )
                residual_steps = residual_steps[:k_ov]
            if getattr(
                self.config, "visual_cumulative_recurrence", False
            ):
                # Cumulative CVRR stores C_k = R_k - B and exposes the actual
                # recurrent state R_k.  There is no temporal mean that can
                # dilute or cancel the accumulated visual corrections.
                return states[0] + residual_steps[-1]
            return states[0] + self._aggregate_crr_residuals(residual_steps)
        if getattr(self.config, "raw_evidence_question", False):
            # states = [R0, R1, ..., RT] in residual mode and
            # [Q*, H1, ..., HT] in state-replacement mode.
            if len(states) != self.config.num_workspace_steps + 1:
                raise ValueError(
                    f"expected R0 plus {self.config.num_workspace_steps} recurrent "
                    f"states, got {len(states)} tensors"
                )
            if getattr(self.config, "raw_mean_aggregation", False):
                return torch.stack(list(states[1:]), dim=0).mean(dim=0)
            if getattr(self.config, "raw_state_replacement", False):
                if getattr(self.config, "raw_question_anchor", False):
                    return states[-1]
                return states[0] + states[-1]
            return states[-1]
        if not getattr(self.config, "evidence_reasoning", False):
            return states[-1]
        r_steps = list(states[1:])  # states[0] is persistent evidence E0
        if getattr(self.config, "concat_aggregation", False):
            if len(r_steps) != self.config.num_workspace_steps:
                raise ValueError(
                    f"expected {self.config.num_workspace_steps} reasoning states, "
                    f"got {len(r_steps)}"
                )
            return self.reasoning_aggregator(torch.cat(r_steps, dim=-1))
        if getattr(self.config, "temporal_aggregation", False):
            a = torch.softmax(self.temporal_logits.float(), dim=0).to(r_steps[0].dtype)
            return torch.stack(r_steps, dim=0).mul(
                a.view(-1, *([1] * r_steps[0].ndim))
            ).sum(dim=0)
        return states[-1]

    def decoder_latent_at_depth(
        self,
        states: list[torch.Tensor] | tuple[torch.Tensor, ...],
        depth: int | torch.Tensor,
    ) -> torch.Tensor:
        """Decode a CRR prefix at a training-selected recurrent depth.

        ``depth`` may be one scalar for the batch or one value per example.
        The latter lets annotated GQA rows follow the program curriculum while
        unannotated rows keep the normal full-horizon answer path.  No new
        readout module is introduced: this is the same native decoder interface
        used by the existing early-stop evaluation.
        """
        if not getattr(
            self.config, "counterfactual_residual_recurrence", False
        ):
            raise ValueError("depth-selected decoding is implemented only for CRR")
        horizon = int(self.config.num_workspace_steps)
        if len(states) != horizon + 1:
            raise ValueError(
                f"expected anchor plus {horizon} CRR states, got {len(states)}"
            )

        def latent_for(prefix_depth: int) -> torch.Tensor:
            if not 1 <= prefix_depth <= horizon:
                raise ValueError(
                    f"curriculum readout depth must be in [1,{horizon}], "
                    f"got {prefix_depth}"
                )
            residual_steps = list(states[1 : prefix_depth + 1])
            if getattr(
                self.config, "visual_cumulative_recurrence", False
            ):
                return states[0] + residual_steps[-1]
            return states[0] + self._aggregate_crr_residuals(residual_steps)

        if not isinstance(depth, torch.Tensor):
            return latent_for(int(depth))
        if depth.ndim == 0:
            return latent_for(int(depth.item()))
        if depth.ndim != 1 or depth.shape[0] != states[0].shape[0]:
            raise ValueError(
                "per-example curriculum depth must have shape [B], got "
                f"{tuple(depth.shape)} for batch {states[0].shape[0]}"
            )
        depth = depth.to(device=states[0].device, dtype=torch.long)
        if bool(((depth < 1) | (depth > horizon)).any()):
            raise ValueError(
                f"curriculum readout depths must lie in [1,{horizon}]"
            )
        candidates = torch.stack(
            [latent_for(prefix_depth) for prefix_depth in range(1, horizon + 1)],
            dim=1,
        )
        rows = torch.arange(candidates.shape[0], device=candidates.device)
        return candidates[rows, depth - 1]

    # -- non-recurrent perceive/deliberate latent chain -------------------

    def _perceive_deliberate_inputs(
        self,
        input_ids: torch.LongTensor,
        pixel_values: torch.Tensor,
        image_grid_thw: torch.LongTensor,
        question_ids: torch.LongTensor,
        attention_mask: torch.Tensor | None,
        question_attention_mask: torch.Tensor | None,
        answer_input_ids: torch.LongTensor | None = None,
        answer_input_attention_mask: torch.Tensor | None = None,
    ) -> tuple[
        torch.Tensor,
        torch.LongTensor,
        torch.BoolTensor,
        PerceiveDeliberateLayout,
        torch.Tensor,
        torch.Tensor,
    ]:
        """Materialize one native block-sparse LOOK/THINK sequence."""

        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids, dtype=torch.long)
        if question_attention_mask is None:
            question_attention_mask = torch.ones_like(
                question_ids, dtype=torch.long
            )
        if input_ids.shape[0] != question_ids.shape[0]:
            raise ValueError("multimodal and clean-question batch sizes differ")

        multimodal_embeddings, multimodal_positions = embed_multimodal(
            self.vl,
            input_ids,
            pixel_values,
            image_grid_thw,
            attention_mask,
        )
        question_embeddings = self.vl.get_input_embeddings()(question_ids)
        batch_size = input_ids.shape[0]
        num_pairs = int(self.config.num_workspace_steps) // 2
        if getattr(self.config, "local_visual_latent_chain", False):
            latent_type_ids = torch.zeros(
                2 * num_pairs,
                device=input_ids.device,
                dtype=torch.long,
            )
        else:
            latent_type_ids = torch.tensor(
                [0, 1] * num_pairs,
                device=input_ids.device,
                dtype=torch.long,
            )
        latent_embeddings = self.latent_chain_embeddings(
            latent_type_ids
        ).unsqueeze(0).expand(batch_size, -1, -1)
        latent_embeddings = latent_embeddings.to(multimodal_embeddings.dtype)

        answer_input_length = (
            0 if answer_input_ids is None else int(answer_input_ids.shape[1])
        )
        layout = PerceiveDeliberateLayout(
            multimodal_length=input_ids.shape[1],
            question_length=question_ids.shape[1],
            num_pairs=num_pairs,
            answer_input_length=answer_input_length,
        )
        pieces = [
            multimodal_embeddings,
            question_embeddings.to(multimodal_embeddings.dtype),
            latent_embeddings,
        ]
        if answer_input_ids is not None:
            pieces.append(
                self.vl.get_input_embeddings()(answer_input_ids).to(
                    multimodal_embeddings.dtype
                )
            )
        embeddings = torch.cat(pieces, dim=1)
        position_ids = build_perceive_deliberate_positions(
            layout,
            multimodal_positions,
            attention_mask,
            question_attention_mask,
            answer_input_attention_mask,
        )
        visibility = build_perceive_deliberate_mask(
            layout,
            attention_mask,
            question_attention_mask,
            answer_input_attention_mask,
            multimodal_visual_mask=(
                image_token_mask(input_ids, int(self.config.image_token_id))
                | image_token_mask(input_ids, int(self.config.video_token_id))
            ),
            local_visual_chain=bool(
                getattr(self.config, "local_visual_latent_chain", False)
            ),
            transition_conditioned=bool(
                getattr(self.config, "pdlc_transition_conditioned", False)
            ),
            question_visible_through_pair=int(
                getattr(
                    self,
                    "_pdlc_question_visible_through_pair_override",
                    getattr(
                        self.config,
                        "pdlc_question_visible_through_pair",
                        1,
                    ),
                )
            ),
            disable_chain_links=bool(
                getattr(self.config, "pdlc_disable_chain_links", False)
            ),
            disable_late_visual=bool(
                getattr(self.config, "pdlc_disable_late_visual", False)
            ),
            final_pair_only=bool(
                getattr(self.config, "pdlc_final_pair_only", False)
            ),
        )
        return (
            embeddings,
            position_ids,
            visibility,
            layout,
            attention_mask,
            question_attention_mask,
        )

    @staticmethod
    def _with_attention_graph(ctx, visibility: torch.BoolTensor):
        """Replace every native attention type with one explicit graph."""

        mappings = {
            attention_type: visibility
            for attention_type in ctx.causal_mask_mapping
        }
        return replace(ctx, causal_mask_mapping=mappings)

    def _run_perceive_deliberate_training_layers(self, ctx) -> torch.Tensor:
        """Run the native stack with explicit per-layer checkpointing.

        Calling decoder layers directly is required for the custom graph, so
        the stock text-model wrapper cannot install its gradient-checkpoint
        closures for us.  Mirror that behavior here; otherwise gradients to
        the LOOK/THINK embeddings retain all 28 layer activations and erase
        most of the B200 batch-size advantage.
        """

        hidden = ctx.hidden_states
        checkpoint_layers = bool(
            self.training
            and getattr(self.text_model, "gradient_checkpointing", False)
        )
        for layer in self.text_model.layers:
            attention_type = getattr(layer, "attention_type", "full_attention")
            layer_mask = ctx.causal_mask_mapping[attention_type]

            def layer_forward(states, *, current_layer=layer, mask=layer_mask):
                output = current_layer(
                    states,
                    attention_mask=mask,
                    position_ids=ctx.text_position_ids,
                    past_key_values=None,
                    use_cache=False,
                    cache_position=ctx.cache_position,
                    position_embeddings=ctx.position_embeddings,
                )
                return output[0] if isinstance(output, tuple) else output

            if checkpoint_layers:
                hidden = checkpoint(
                    layer_forward,
                    hidden,
                    use_reentrant=False,
                )
            else:
                hidden = layer_forward(hidden)
        return hidden

    def _forward_perceive_deliberate_chain(
        self,
        input_ids: torch.LongTensor,
        attention_mask: torch.Tensor | None,
        pixel_values: torch.Tensor,
        image_grid_thw: torch.LongTensor,
        question_ids: torch.LongTensor,
        question_attention_mask: torch.Tensor | None,
        labels: torch.LongTensor,
        answer_ids: torch.LongTensor,
    ) -> CloseCausalLMOutput:
        """Teacher-forced answer CE through one native LOOK/THINK graph."""

        if labels.ndim != 2 or answer_ids.ndim != 2:
            raise ValueError("labels and answer_ids must have shape [B, N_a]")
        if labels.shape != answer_ids.shape or labels.shape[1] < 1:
            raise ValueError(
                "labels and answer_ids must share one non-empty answer shape"
            )
        answer_input_ids = answer_ids[:, :-1]
        answer_input_mask = labels[:, :-1].ne(-100)
        (
            embeddings,
            position_ids,
            visibility,
            layout,
            _,
            _,
        ) = self._perceive_deliberate_inputs(
            input_ids,
            pixel_values,
            image_grid_thw,
            question_ids,
            attention_mask,
            question_attention_mask,
            answer_input_ids,
            answer_input_mask,
        )
        ctx = make_split_context(
            self.text_model,
            embeddings,
            position_ids,
            attention_mask=None,
        )
        ctx = self._with_attention_graph(ctx, visibility)
        hidden = self._run_perceive_deliberate_training_layers(ctx)
        hidden = final_norm(self.text_model, hidden)

        # The final THINK predicts the first answer token.  Teacher-forced
        # answer rows predict all subsequent tokens, matching the repository's
        # already-aligned labels convention without an internal label shift.
        seed = hidden[:, layout.think_indices[-1] : layout.think_indices[-1] + 1]
        if layout.answer_input_length:
            prediction_hidden = torch.cat(
                [seed, hidden[:, layout.answer_slice]], dim=1
            )
        else:
            prediction_hidden = seed
        logits = self.backbone.lm_head(prediction_hidden)
        loss = answer_cross_entropy(
            logits,
            labels,
            per_example=bool(
                getattr(self.config, "per_example_answer_loss", False)
            ),
        )

        latent_hidden = hidden[:, layout.latent_start : layout.answer_start]
        expected_latents = int(self.config.num_workspace_steps)
        if latent_hidden.shape[1] != expected_latents:
            raise RuntimeError(
                "latent layout produced an invalid state count: "
                f"expected {expected_latents}, got {latent_hidden.shape[1]}"
            )
        with torch.no_grad():
            normalized = nn.functional.normalize(latent_hidden.float(), dim=-1)
            adjacent = 1.0 - (normalized[:, 1:] * normalized[:, :-1]).sum(-1)
            look = normalized[:, 0::2]
            think = normalized[:, 1::2]
            pair = 1.0 - (look * think).sum(-1)
            think_progress = (
                1.0 - (think[:, 1:] * think[:, :-1]).sum(-1)
                if think.shape[1] > 1
                else think.new_empty(think.shape[0], 0)
            )
            if getattr(self.config, "local_visual_latent_chain", False):
                self.last_local_visual_latent_telemetry = {
                    "transition_drift": adjacent.mean(dim=0),
                }
            else:
                self.last_perceive_deliberate_telemetry = {
                    "adjacent_drift": adjacent.mean(dim=0),
                    "pair_drift": pair.mean(dim=0),
                    "think_progress": think_progress.mean(dim=0),
                }
        states = tuple(
            latent_hidden[:, index : index + 1]
            for index in range(latent_hidden.shape[1])
        )
        return CloseCausalLMOutput(
            loss=loss,
            logits=logits,
            past_key_values=None,
            workspace_states=states,
            v_star_pooled=None,
        )

    @torch.no_grad()
    def _generate_perceive_deliberate_chain(
        self,
        input_ids: torch.LongTensor,
        pixel_values: torch.Tensor,
        image_grid_thw: torch.LongTensor,
        question_ids: torch.LongTensor,
        attention_mask: torch.Tensor | None,
        question_attention_mask: torch.Tensor | None,
        *,
        max_new_tokens: int,
        eos_token_id: int | None,
    ) -> torch.LongTensor:
        """Greedy generation with the multimodal/LOOK cache rows masked."""

        if max_new_tokens < 1:
            raise ValueError("max_new_tokens must be >= 1")
        (
            embeddings,
            position_ids,
            visibility,
            layout,
            _,
            question_attention_mask,
        ) = self._perceive_deliberate_inputs(
            input_ids,
            pixel_values,
            image_grid_thw,
            question_ids,
            attention_mask,
            question_attention_mask,
        )
        cache = DynamicCache(config=self.text_model.config)
        ctx = make_split_context(
            self.text_model,
            embeddings,
            position_ids,
            attention_mask=None,
            past_key_values=cache,
        )
        ctx = self._with_attention_graph(ctx, visibility)
        hidden = run_layer_range(
            self.text_model, ctx, 0, None, use_cache=True
        )
        hidden = final_norm(self.text_model, hidden)
        seed = hidden[:, layout.think_indices[-1] : layout.think_indices[-1] + 1]
        nxt = self.backbone.lm_head(seed).argmax(dim=-1)
        output = [nxt]
        batch_size = input_ids.shape[0]
        done = torch.zeros(
            batch_size, dtype=torch.bool, device=input_ids.device
        )
        if eos_token_id is not None:
            done |= nxt.squeeze(1).eq(eos_token_id)

        final_prefix_position = position_ids[:, :, -1]
        for step in range(1, max_new_tokens):
            if bool(done.all()):
                break
            token_embeddings = self.vl.get_input_embeddings()(nxt)
            cache_position = torch.tensor(
                [layout.answer_start + step - 1], device=input_ids.device
            )
            token_positions = (
                final_prefix_position + step
            ).unsqueeze(-1)
            generated_mask = torch.ones(
                batch_size,
                step,
                dtype=question_attention_mask.dtype,
                device=input_ids.device,
            )
            answer_visibility = build_answer_decode_mask(
                layout,
                question_attention_mask,
                generated_mask,
                local_visual_chain=bool(
                    getattr(
                        self.config,
                        "local_visual_latent_chain",
                        False,
                    )
                ),
                final_pair_only=bool(
                    getattr(self.config, "pdlc_final_pair_only", False)
                ),
                transition_conditioned=bool(
                    getattr(
                        self.config,
                        "pdlc_transition_conditioned",
                        False,
                    )
                ),
            )
            token_ctx = make_split_context(
                self.text_model,
                token_embeddings,
                token_positions,
                attention_mask=None,
                past_key_values=cache,
                cache_position=cache_position,
            )
            token_ctx = self._with_attention_graph(
                token_ctx, answer_visibility
            )
            token_hidden = run_layer_range(
                self.text_model,
                token_ctx,
                0,
                None,
                use_cache=True,
            )
            token_hidden = final_norm(self.text_model, token_hidden)
            nxt = self.backbone.lm_head(token_hidden).argmax(dim=-1)
            if eos_token_id is not None:
                nxt = torch.where(
                    done.unsqueeze(1),
                    nxt.new_full(nxt.shape, eos_token_id),
                    nxt,
                )
                done |= nxt.squeeze(1).eq(eos_token_id)
            output.append(nxt)
        return torch.cat(output, dim=1)

    # -- spatial recurrent visual field -----------------------------------

    def _spatial_horizon(self) -> int:
        """Return the configured train horizon or an explicit eval prefix."""
        override = int(getattr(self.config, "spatial_step_override", -1))
        if override >= 0:
            return override
        return int(self.config.spatial_recurrence_steps)

    def _spatial_image_to_batch(
        self,
        input_ids: torch.LongTensor,
        attention_mask: torch.Tensor | None,
        num_images: int,
    ) -> torch.LongTensor:
        """Map Qwen's batch-major image feature list back to sample rows."""
        if input_ids.ndim != 2:
            raise ValueError("spatial visual recurrence expects [B,L] input_ids")
        starts = input_ids[:, :-1].eq(self.config.vision_start_token_id)
        images = input_ids[:, 1:].eq(self.config.image_token_id)
        markers = starts & images
        if attention_mask is not None:
            markers &= attention_mask[:, :-1].gt(0)
            markers &= attention_mask[:, 1:].gt(0)
        counts = markers.sum(dim=-1).long()
        if int(counts.sum()) != int(num_images):
            raise ValueError(
                "image_grid_thw rows do not match <vision_start><image> "
                f"segments: grids={num_images}, segments={int(counts.sum())}"
            )
        return torch.repeat_interleave(
            torch.arange(input_ids.shape[0], device=input_ids.device), counts
        )

    def _spatial_visual_prefix_embeddings(
        self,
        input_ids: torch.LongTensor,
        pixel_values: torch.Tensor,
        image_grid_thw: torch.LongTensor,
        question_ids: torch.LongTensor,
        attention_mask: torch.Tensor | None,
        question_attention_mask: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Build the ordinary Qwen prefix with image rows replaced by V_T.

        Vision and token embeddings are frozen inputs.  Gradients start at the
        shared spatial cell, continue through the newly built full decoder
        pass, and never enter a cache from visual-state construction.
        """
        if image_grid_thw is None or image_grid_thw.ndim != 2:
            raise ValueError("image_grid_thw is required for spatial recurrence")
        if question_ids.ndim != 2 or question_ids.shape[0] != input_ids.shape[0]:
            raise ValueError("question_ids must be [B,Q] and match input_ids")
        prefix_mask = attention_mask
        if prefix_mask is None:
            prefix_mask = torch.ones_like(input_ids, dtype=torch.long)

        with torch.no_grad():
            prefix_embeddings = self.vl.get_input_embeddings()(input_ids)
            question_embeddings = self.vl.get_input_embeddings()(question_ids)
            image_parts = self.vl.get_image_features(
                pixel_values, image_grid_thw
            )
            image_state = torch.cat(tuple(image_parts), dim=0).to(
                device=prefix_embeddings.device,
                dtype=prefix_embeddings.dtype,
            )

        merge = int(self.config.vision_config.spatial_merge_size)
        merged_grid = image_grid_thw.to(device=image_state.device).clone()
        if bool((merged_grid[:, 1:] % merge).ne(0).any()):
            raise ValueError(
                "image grid height/width must be divisible by spatial_merge_size"
            )
        merged_grid[:, 1:] //= merge
        image_to_batch = self._spatial_image_to_batch(
            input_ids,
            prefix_mask,
            merged_grid.shape[0],
        ).to(image_state.device)
        refined, telemetry = self.spatial_visual_cell(
            image_state,
            merged_grid,
            image_to_batch,
            question_embeddings,
            question_attention_mask,
            steps=self._spatial_horizon(),
        )
        self.last_spatial_visual_telemetry = telemetry

        image_mask, _ = self.vl.get_placeholder_mask(
            input_ids,
            inputs_embeds=prefix_embeddings,
            image_features=refined,
        )
        prefix_embeddings = prefix_embeddings.masked_scatter(
            image_mask, refined.to(prefix_embeddings.dtype)
        )
        return prefix_embeddings, prefix_mask

    def _forward_spatial_visual_field(
        self,
        input_ids: torch.LongTensor,
        attention_mask: torch.Tensor | None,
        pixel_values: torch.Tensor,
        image_grid_thw: torch.LongTensor,
        question_ids: torch.LongTensor,
        question_attention_mask: torch.Tensor | None,
        labels: torch.LongTensor | None,
        answer_ids: torch.LongTensor | None,
        *,
        use_cache: bool | None,
    ) -> CloseCausalLMOutput:
        """Teacher-forced SRVF pass through the unmodified full decoder."""
        prefix_embeddings, prefix_mask = self._spatial_visual_prefix_embeddings(
            input_ids,
            pixel_values,
            image_grid_thw,
            question_ids,
            attention_mask,
            question_attention_mask,
        )
        prefix_positions, rope_deltas = self.vl.get_rope_index(
            input_ids,
            image_grid_thw=image_grid_thw,
            attention_mask=prefix_mask,
        )
        prefix_length = input_ids.shape[1]
        batch_size = input_ids.shape[0]
        if not bool(prefix_mask.gt(0).any(dim=-1).all()):
            raise ValueError("every sample needs a non-empty multimodal prefix")

        full_embeddings = prefix_embeddings
        full_positions = prefix_positions
        full_mask = prefix_mask
        answer_length = 0
        if answer_ids is not None:
            if answer_ids.ndim != 2 or answer_ids.shape[0] != batch_size:
                raise ValueError("answer_ids must be [B,A] and match input_ids")
            answer_length = answer_ids.shape[1]
            with torch.no_grad():
                answer_embeddings = self.vl.get_input_embeddings()(answer_ids)
            if labels is not None:
                if labels.shape != answer_ids.shape:
                    raise ValueError("labels and answer_ids must have the same shape")
                answer_mask = labels.ne(-100).to(prefix_mask.dtype)
            else:
                answer_mask = torch.ones_like(answer_ids, dtype=prefix_mask.dtype)
            physical = torch.arange(
                prefix_length,
                prefix_length + answer_length,
                device=input_ids.device,
            )[None, :]
            answer_positions_1d = physical + rope_deltas.to(physical.device)
            answer_positions = answer_positions_1d.unsqueeze(0).expand(
                3, -1, -1
            )
            full_embeddings = torch.cat(
                [prefix_embeddings, answer_embeddings], dim=1
            )
            full_positions = torch.cat(
                [prefix_positions, answer_positions], dim=-1
            )
            full_mask = torch.cat([prefix_mask, answer_mask], dim=-1)
        elif labels is not None:
            raise ValueError("labels require teacher-forced answer_ids")

        outputs = self.text_model(
            input_ids=None,
            inputs_embeds=full_embeddings,
            position_ids=full_positions,
            attention_mask=full_mask,
            use_cache=bool(use_cache),
            output_attentions=False,
            output_hidden_states=False,
            return_dict=True,
        )
        last_prefix = (prefix_mask.gt(0).sum(dim=-1).long() - 1).clamp_min(0)
        seed_hidden = outputs.last_hidden_state[
            torch.arange(batch_size, device=input_ids.device), last_prefix
        ].unsqueeze(1)
        hidden_for_logits = seed_hidden
        if answer_length > 1:
            hidden_for_logits = torch.cat(
                [
                    seed_hidden,
                    outputs.last_hidden_state[
                        :, prefix_length : prefix_length + answer_length - 1
                    ],
                ],
                dim=1,
            )
        logits = self.backbone.lm_head(hidden_for_logits)
        loss = None
        if labels is not None:
            loss = answer_cross_entropy(
                logits,
                labels,
                per_example=bool(
                    getattr(self.config, "per_example_answer_loss", False)
                ),
            )
        return CloseCausalLMOutput(
            loss=loss,
            logits=logits,
            past_key_values=(outputs.past_key_values if use_cache else None),
            workspace_states=(),
        )

    @torch.no_grad()
    def _generate_spatial_visual_field(
        self,
        input_ids: torch.LongTensor,
        pixel_values: torch.Tensor,
        image_grid_thw: torch.LongTensor,
        question_ids: torch.LongTensor,
        attention_mask: torch.Tensor | None,
        question_attention_mask: torch.Tensor | None,
        *,
        max_new_tokens: int,
        eos_token_id: int | None,
    ) -> torch.LongTensor:
        """Greedy decode from a cache built only after V_T is complete."""
        if max_new_tokens < 1:
            raise ValueError("max_new_tokens must be >= 1")
        prefix_embeddings, prefix_mask = self._spatial_visual_prefix_embeddings(
            input_ids,
            pixel_values,
            image_grid_thw,
            question_ids,
            attention_mask,
            question_attention_mask,
        )
        prefix = self.vl(
            input_ids=input_ids,
            inputs_embeds=prefix_embeddings,
            image_grid_thw=image_grid_thw,
            attention_mask=prefix_mask,
            use_cache=True,
            output_attentions=False,
            output_hidden_states=False,
            return_dict=True,
        )
        cache = prefix.past_key_values
        rope_deltas = prefix.rope_deltas
        batch_size, prefix_length = input_ids.shape
        last_prefix = (prefix_mask.gt(0).sum(dim=-1).long() - 1).clamp_min(0)
        seed = prefix.last_hidden_state[
            torch.arange(batch_size, device=input_ids.device), last_prefix
        ].unsqueeze(1)
        next_token = self.backbone.lm_head(seed).argmax(dim=-1)
        generated = [next_token]
        done = torch.zeros(
            batch_size, dtype=torch.bool, device=input_ids.device
        )
        if eos_token_id is not None:
            done |= next_token.squeeze(1).eq(eos_token_id)

        for step in range(1, max_new_tokens):
            if bool(done.all()):
                break
            with torch.no_grad():
                token_embeddings = self.vl.get_input_embeddings()(next_token)
            cache_position = torch.tensor(
                [prefix_length + step - 1], device=input_ids.device
            )
            one_d = cache_position.view(1, 1).expand(batch_size, 1)
            one_d = one_d + rope_deltas.to(one_d.device)
            position_ids = one_d.unsqueeze(0).expand(3, -1, -1)
            generated_mask = torch.ones(
                batch_size,
                step,
                dtype=prefix_mask.dtype,
                device=input_ids.device,
            )
            decode_mask = torch.cat([prefix_mask, generated_mask], dim=-1)
            decoded = self.text_model(
                input_ids=None,
                inputs_embeds=token_embeddings,
                position_ids=position_ids,
                attention_mask=decode_mask,
                past_key_values=cache,
                cache_position=cache_position,
                use_cache=True,
                output_attentions=False,
                output_hidden_states=False,
                return_dict=True,
            )
            cache = decoded.past_key_values
            next_token = self.backbone.lm_head(
                decoded.last_hidden_state
            ).argmax(dim=-1)
            if eos_token_id is not None:
                next_token = torch.where(
                    done.unsqueeze(1),
                    next_token.new_full(next_token.shape, eos_token_id),
                    next_token,
                )
                done |= next_token.squeeze(1).eq(eos_token_id)
            generated.append(next_token)
        return torch.cat(generated, dim=1)

    # -- §3.2 pipeline -----------------------------------------------------

    @contextmanager
    def _recurrence_adapter_execution(self, enabled: bool):
        """Temporarily toggle recurrence-only LoRA execution.

        PEFT's public ``enable_adapters`` helper also changes
        ``requires_grad``.  That is appropriate for persistent train/eval
        configuration, but not for routing one shared L21 through a frozen
        base call and a trainable recurrent call in the same autograd graph.
        The tuner layer's execution flag is therefore saved and restored
        directly while its parameters remain trainable throughout.
        """
        if not getattr(self.config, "recurrence_only_adapter", False):
            yield
            return

        from peft.tuners.tuners_utils import BaseTunerLayer

        modules = [
            module
            for module in self.backbone.modules()
            if isinstance(module, BaseTunerLayer)
        ]
        previous = [module.disable_adapters for module in modules]
        for module in modules:
            module._disable_adapters = not enabled
        try:
            yield
        finally:
            for module, was_disabled in zip(modules, previous):
                module._disable_adapters = was_disabled

    @torch.no_grad()
    def _set_adapters_enabled(self, enabled: bool) -> None:
        """post-l* LoRA on/off (iface teacher는 frozen-base 동작점 필요)."""
        from peft.tuners.tuners_utils import BaseTunerLayer
        for m in self.backbone.modules():
            if isinstance(m, BaseTunerLayer):
                m.enable_adapters(enabled)

    def _encode_visual_memory(
        self,
        input_ids: torch.LongTensor,  # [B, L_mm] multimodal sequence
        pixel_values: torch.Tensor,
        image_grid_thw: torch.LongTensor,
        attention_mask: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """``V* = Pi_img(F^mm_{<=l*}(I, q))``, right-padded. ``[B, N_v, d]``.

        Runs under ``no_grad``: the vision encoder and lower backbone are frozen,
        and nothing above needs gradients through them.
        """
        embeds, pos = embed_multimodal(
            self.vl, input_ids, pixel_values, image_grid_thw, attention_mask
        )
        ctx = make_split_context(self.text_model, embeds, pos, attention_mask)
        lower = run_layer_range(self.text_model, ctx, 0, self.config.ell_star + 1)
        crr = bool(
            getattr(
                self.config, "counterfactual_residual_recurrence", False
            )
        )
        capture_functional = bool(
            self.training
            and getattr(self.config, "functional_state_loss", False)
        )
        if crr or capture_functional:
            # Free target from the same lower pass that produced V*: normal
            # multimodal question states at the measured read boundary.  The
            # oracle shows that these rows alone retain base V* competence in a
            # Q-only upper context.  Store detached, padded rows for the
            # training-only split-state restoration objective.
            text_mask = ~vision_span_mask(input_ids, self.config)
            if attention_mask is not None:
                text_mask &= attention_mask > 0
            q_mm, q_mm_pad = select_tokens_padded(lower, text_mask)
            q_mm_ids, _ = select_tokens_padded(input_ids.unsqueeze(-1), text_mask)
            q_mm = q_mm.detach()
            q_mm_ids = q_mm_ids.squeeze(-1)
            if crr:
                # Consume-once factual boundary state.  It is the sole image
                # input to CRR; raw visual rows are not materialized below.
                self._counterfactual_q_state = q_mm
                self._counterfactual_q_pad = q_mm_pad
                self._counterfactual_q_ids = q_mm_ids
                if getattr(
                    self.config, "visual_counterfactual_recurrence", False
                ):
                    visual_mask = vision_span_mask(input_ids, self.config)
                    if attention_mask is not None:
                        visual_mask &= attention_mask > 0
                    # Consume-once native multimodal scaffold.  The lower
                    # state/context never reaches the answer decoder; it is
                    # used only by L(ell_star+1)'s visual-on/off recurrent
                    # evaluations and cleared immediately afterwards.
                    self._visual_counterfactual_bundle = (
                        lower.detach(),
                        ctx,
                        text_mask,
                        visual_mask,
                    )
            if capture_functional:
                self._functional_q_target = q_mm
                self._functional_q_target_pad = q_mm_pad
                self._functional_q_target_ids = q_mm_ids
        if self.training and getattr(self.config, "interface_loss", False):
            # iface teacher: frozen-base 정상(mm) 경로의 q_end hidden @ l*+1.
            # adapter를 꺼 base 동작점을 재현 — 학습 target일 뿐 student
            # forward에 입력되지 않으므로 새 causal bypass 없음 (advisor).
            self._set_adapters_enabled(False)
            try:
                h1 = run_layer_range(
                    self.text_model, ctx, self.config.ell_star + 1,
                    self.config.ell_star + 2, hidden_states=lower)
            finally:
                self._set_adapters_enabled(True)
            if attention_mask is not None:
                idx = attention_mask.sum(1).long() - 1  # [B] 마지막 non-pad
            else:
                idx = torch.full(
                    (lower.shape[0],), lower.shape[1] - 1,
                    dtype=torch.long, device=lower.device)
            self._iface_h_base = h1[
                torch.arange(lower.shape[0], device=lower.device), idx
            ].detach()  # [B, d]
        if crr:
            # The multimodal lower pass is needed for Q_mm, but CRR has no V*
            # memory.  Returning an empty tensor prevents an accidental raw
            # visual path while preserving the historical helper signature.
            return (
                lower.new_empty(lower.shape[0], 0, lower.shape[-1]),
                torch.empty(
                    lower.shape[0], 0, dtype=torch.bool, device=lower.device
                ),
            )
        mask = image_token_mask(input_ids, self.config.image_token_id)
        return select_tokens_padded(lower, mask)  # [B, N_v, d], [B, N_v]

    @torch.no_grad()
    def _encode_hierarchical_visual_memories(
        self,
        input_ids: torch.LongTensor,
        pixel_values: torch.Tensor,
        image_grid_thw: torch.LongTensor,
        attention_mask: torch.Tensor | None,
    ) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]:
        """Extract one frozen visual memory per visual-read step in one pass.

        Layer numbers are 0-indexed outputs: requesting L6 returns the hidden
        state after decoder layer 6.  The multimodal branch stops after the
        deepest requested layer because no later multimodal state is consumed
        by this method; the independent text-only branch still reaches
        ``ell_star`` and supplies the decoder scaffold.
        """
        layers = tuple(
            int(layer)
            for layer in getattr(
                self.config, "hierarchical_visual_layers", ()
            )
        )
        if not layers:
            raise RuntimeError(
                "hierarchical visual extraction requires configured layers"
            )

        embeds, pos = embed_multimodal(
            self.vl, input_ids, pixel_values, image_grid_thw, attention_mask
        )
        ctx = make_split_context(self.text_model, embeds, pos, attention_mask)
        visual_mask = image_token_mask(input_ids, self.config.image_token_id)
        hidden = embeds
        start = 0
        memories: list[tuple[torch.Tensor, torch.Tensor]] = []
        reference_pad = None
        for layer in layers:
            hidden = run_layer_range(
                self.text_model,
                ctx,
                start,
                layer + 1,
                hidden_states=hidden,
            )
            memory, padding_mask = select_tokens_padded(hidden, visual_mask)
            if reference_pad is None:
                reference_pad = padding_mask
            elif not torch.equal(reference_pad, padding_mask):
                raise RuntimeError(
                    "visual-token padding changed across frozen layer reads"
                )
            memories.append((memory, padding_mask))
            start = layer + 1
        return tuple(memories)

    def _prefill_text_branch(
        self,
        question_ids: torch.LongTensor,  # [B, N_q] no image tokens
        question_mask: torch.Tensor | None,
    ):
        """Run the text-only lower branch through ``F_{<=l*}``, caching KV.

        Returns ``(q_star, ctx, cache)`` where ``q_star`` is ``[B, N_q, d]`` and
        ``ctx.hidden_states`` holds the lower-branch state.  The raw-evidence
        path uses exactly the question rows; historical slot paths append their
        fixed cache-shape placeholders.
        """
        cfg = self.config
        b, n_q = question_ids.shape

        embeds = self.vl.get_input_embeddings()(question_ids)  # [B, N_q, d]
        if question_mask is None:
            question_mask = torch.ones(b, n_q, dtype=torch.long, device=embeds.device)

        if self._uses_question_state_interface():
            # No placeholders and no duplicated Q prefix.  Logical text RoPE
            # positions remain invariant to right-padding in a batch.
            keep = question_mask > 0
            one_d = keep.long().cumsum(dim=-1) - 1
            one_d = one_d.masked_fill(~keep, 0)
            position_ids = one_d.unsqueeze(0).expand(3, -1, -1)
            cache = DynamicCache(config=self.text_model.config)
            ctx = make_split_context(
                self.text_model,
                embeds,
                position_ids,
                question_mask,
                past_key_values=cache,
            )
            lower = run_layer_range(
                self.text_model, ctx, 0, cfg.ell_star + 1, use_cache=True
            )
            ctx.hidden_states = lower
            return lower, ctx, cache

        s = cfg.num_workspace_slots
        slots = self.slot_embedding.unsqueeze(0).expand(b, -1, -1).to(embeds.dtype)
        embeds = torch.cat([embeds, slots], dim=1)  # [B, N_q + S, d]

        # Physical rows are [right-padded q; slots], but each item's slots get
        # logical RoPE positions immediately after its own real question.
        full_mask, position_ids, _ = _replacement_prefix_layout(question_mask, s)

        cache = DynamicCache(config=self.text_model.config)
        ctx = make_split_context(
            self.text_model, embeds, position_ids, full_mask, past_key_values=cache
        )
        lower = run_layer_range(
            self.text_model, ctx, 0, cfg.ell_star + 1, use_cache=True
        )  # [B, N_q + S, d]
        ctx.hidden_states = lower
        return lower[:, :n_q], ctx, cache

    def _decode_token_block(
        self,
        embeds: torch.Tensor,  # [B,N_a,d]
        position_ids: torch.LongTensor,  # [3,B,N_a]
        lower_attention_mask: torch.Tensor,  # [B,N_q+S+N_a], slots masked
        upper_attention_mask: torch.Tensor,  # [B,N_q+S+N_a], slots visible
        cache: Cache,
        cache_position: torch.LongTensor,  # [N_a], physical cache positions
        return_lower: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        """Decode answer rows while preserving the strict split interface.

        Lower layers may attend only to the text question and preceding answer
        tokens. The physical placeholder slot rows exist in their cache solely
        so every layer has the same cache length, and are masked here. Upper
        layers see those same positions as the overwritten ``R_agg`` rows.
        """
        tm = self.text_model
        lower_ctx = make_split_context(
            tm,
            embeds,
            position_ids,
            lower_attention_mask,
            past_key_values=cache,
            cache_position=cache_position,
        )
        upper_start = self._upper_decoder_start()
        hidden = run_layer_range(
            tm, lower_ctx, 0, upper_start, use_cache=True
        )
        upper_ctx = make_split_context(
            tm,
            embeds,
            position_ids,
            upper_attention_mask,
            past_key_values=cache,
            cache_position=cache_position,
        )
        upper = run_layer_range(
            tm,
            upper_ctx,
            upper_start,
            None,
            hidden_states=hidden,
            use_cache=True,
        )
        if return_lower:
            return upper, hidden
        return upper

    def _competitive_read(self, q_in: torch.Tensor, kv_in: torch.Tensor) -> torch.Tensor:
        """ER v2.1 (advisor 2026-08-11): competitive normalization evidence read.

        기존 ``evidence_read`` MHA의 W_qkv/W_o를 그대로 사용하되, 표준
        row-softmax(evidence 축) 대신 ① reasoning-slot 축 softmax — 각 E_j를
        두고 R slot들이 경쟁 (Slot Attention식) — ② evidence 축 재정규화.
        새 파라미터/모듈/loss 없음. vanilla read에서는 전 slot이 같은 E_j를
        100% 읽는 해가 무벌점이라 실제로 그 해로 붕괴했다 (smoke A/B §5).

        q_in: [B, S, d] = LN(R), kv_in: [B, N, d] = LN(E) — norm은 호출부에서.
        """
        mha = self.evidence_read
        d, h = mha.embed_dim, mha.num_heads
        dh = d // h
        w, b = mha.in_proj_weight, mha.in_proj_bias
        q = nn.functional.linear(q_in, w[:d], None if b is None else b[:d])
        k = nn.functional.linear(kv_in, w[d:2 * d], None if b is None else b[d:2 * d])
        v = nn.functional.linear(kv_in, w[2 * d:], None if b is None else b[2 * d:])
        B, S, _ = q.shape
        N = k.shape[1]
        q = q.view(B, S, h, dh).transpose(1, 2)  # [B, H, S, dh]
        k = k.view(B, N, h, dh).transpose(1, 2)  # [B, H, N, dh]
        v = v.view(B, N, h, dh).transpose(1, 2)  # [B, H, N, dh]
        logits = q @ k.transpose(-1, -2) / math.sqrt(dh)  # [B, H, S, N]
        # 확정 구현 (advisor): log-domain competitive normalization.
        # A = softmax_j(logsoftmax_s(L)) ≡ C/Σ_j C (수학적 동일; FP64 편차 6e-17)
        # 이되, 확률 공간의 나눗셈이 없어 전패 slot에서도 forward/backward 유계
        # (극단 경쟁 unit test: exact=NaN, logdom max|grad| 0.68, 행합 1 유지).
        # 안정화 이력: ①분모-ε → 1/ε=1e6 증폭, grad 1e10 폭발 (x1, step 2.4k)
        # ②ε-분자 → ∂A/∂à = 1/Σ(Ã+ε) ≈ 1/(Nε) 증폭, 재폭발 (step 3.9k)
        # ③분모 clamp → 유계지만 전패 slot 읽기 차단 (의미론 변경이라 기각)
        # ④log-domain — ε/clamp/hyperparameter 없음. 경쟁 의미론 정확 보존.
        log_comp = nn.functional.log_softmax(logits.float(), dim=-2)  # slot 축 경쟁
        a = torch.softmax(log_comp, dim=-1)  # evidence 축 재정규화 (log 공간)
        a = a.to(v.dtype)
        if getattr(self, "_er_attn_capture", None) is not None:
            self._er_attn_capture.append(a.detach().float().cpu())  # [B, H, S, N]
        out = (a @ v).transpose(1, 2).reshape(B, S, d)  # [B, S, d]
        return mha.out_proj(out)

    def _encode_workspace_er(
        self,
        e0: torch.Tensor,  # [B, S, d_w] E^(0) = InitialRead(V*)
        q_star: torch.Tensor,
        v_star: torch.Tensor,
        q_padding_mask: torch.Tensor | None,
        v_padding_mask: torch.Tensor | None,
        gate_overrides: list | None,
        v_star_reread: torch.Tensor | None,
        v_reread_padding_mask: torch.Tensor | None,
        reread_zero_memory: bool,
        reread_memory_overrides: list | None,
        mid_step: int,
    ) -> list[torch.Tensor]:
        """Evidence–Reasoning 분리 — strict 버전 (advisor 2026-08-06 수정).

        decoder 입력은 **R^(T) 단독** — E 직행 우회(E→decoder)를 구조적으로
        제거해 모든 시각 정보가 reasoning을 통과하게 강제한다:
        I → V* → E → R^(1:T) → answer.

        recurrence 입력은 detach 없는 R + E (answer CE가 f_theta를 경유해
        E·reader까지 학습). gradient routing 표(advisor)는 쌍둥이-그래프로
        구현한다: traj가 소비할 R-체인만 sg[E]로 병렬 전개 —
        값은 비트동일(dropout=0 필수), 그래프만 분리되어
          - L_ans   → 전부 (reader/E/transition/adapter)      [answer 체인]
          - L_anchor(k=0) → initial reader·evidence proj만    [E^(0) 그래프]
          - L_reason(k>=1) → transition·d_psi만               [traj 체인, sg[E]]
        이 되고 trainer/optimizer 변경이 필요 없다.

        반환 states = [E^(0), R_t^(1..T)] (traj 체인). answer 체인의 R^(T)는
        self._er_answer_final(consume-once)로 prefill에 전달. E의 detached
        사본은 self._er_last_evidence에 남겨 probe가 쓴다.
        """
        assert mid_step > 0, "ER 모드는 mid_read_step > 0 필요"
        assert float(getattr(self.config, "workspace_dropout", 0.0)) == 0.0, (
            "ER 쌍둥이-체인은 dropout=0 필요 (두 체인의 값 동일성 전제)")
        self._er_answer_final = None  # stale-state 방지: forward 끝에서만 설정
        E = e0
        v2 = bool(getattr(self.config, "explicit_evidence_read", False))
        if getattr(self, "_debug_er_capture", False):
            self._dbg_R_ans, self._dbg_R_traj = [], []
        if getattr(self.config, "r_init_question", False):
            # ER v2 canonical (§4): R^0 = R_init (learnable slots, batch broadcast).
            # Q* 조건화 블록/Pool 없음 — Q*는 매 transition에서 입력되므로 불필요.
            # traj 체인은 stopgrad(R_init) 값 공유 (§9.3: traj -X-> reasoning_slots).
            r0 = self.reasoning_slots.unsqueeze(0).expand(E.shape[0], -1, -1)
            R, R_t = r0, r0.detach()
        else:
            R = torch.zeros_like(E)     # answer 체인 (E gradient 통과)
            R_t = torch.zeros_like(E)   # traj 체인 (sg[E] — evidence/reader 차단)
        states = [E]

        for k in range(self.config.num_workspace_steps):
            if v2:
                # ER v2 canonical (§5-6): R̄ = R + CrossAttn(LN(R), LN(E), LN(E));
                # R = f_theta(R̄, Q*). transition은 잔차 블록이므로 R̄를 외부에서
                # 재가산하지 않는다 (double-residual 금지, §6).
                # traj 체인: evidence_read(R_t, sg(E)) — E 값은 읽되 E 생성부 차단 (§9.3).
                def _read_E(r, mem):
                    q = self.evidence_read_norm_q(r)
                    kv = self.evidence_read_norm_kv(mem)
                    if getattr(self.config, "competitive_evidence_read", False):
                        return r + self._competitive_read(q, kv)
                    out, _ = self.evidence_read(q, kv, kv, need_weights=False)
                    return r + out
                R = self.workspace_transition(_read_E(R, E), q_star, q_padding_mask)
                if getattr(self.config, "er_single_chain", False):
                    # ER v2.2 C (§12): k>=1 traj 감독이 없으므로 twin 불필요 —
                    # states가 answer 체인을 직접 담아 L_ans가 R1..RT 전체에
                    # (aggregation 경유) gradient를 보낸다 (§11).
                    R_t = R
                else:
                    R_t = self.workspace_transition(
                        _read_E(R_t, E.detach()), q_star, q_padding_mask)
                if getattr(self, "_debug_er_capture", False):
                    self._dbg_R_ans.append(R.detach().clone())
                    self._dbg_R_traj.append(R_t.detach().clone())
            else:
                x = R + E
                total = self.workspace_transition(x, q_star, q_padding_mask)
                R = total - E
                x_t = R_t + E.detach()
                total_t = self.workspace_transition(x_t, q_star, q_padding_mask)
                R_t = total_t - E.detach()
            if (k + 1) == mid_step:
                ov = gate_overrides[k] if gate_overrides is not None else None
                off_mid = ov == 0.0
                if not off_mid:
                    v_r, vp_r = v_star, v_padding_mask
                    if v_star_reread is not None:
                        v_r, vp_r = v_star_reread, v_reread_padding_mask
                    if reread_zero_memory:
                        v_r = torch.zeros_like(v_r)
                    elif (reread_memory_overrides is not None
                          and reread_memory_overrides[k] is not None):
                        v_r, vp_r = reread_memory_overrides[k]
                    # ER v2 (advisor): 재읽기 query는 R-only — "무엇을 다시 볼지"는
                    # reasoning state가 정하고, E는 저장소 역할만. (v1 경로는 E+R 유지.)
                    u_q = R if v2 else (E + R)
                    woa = self.workspace_read.visual_residual(u_q, v_r, vp_r)
                    a_eff = float(getattr(self.config, "mid_read_alpha", 1.0))
                    if ov is not None:
                        a_eff = float(ov)  # eval sweep은 절대값으로 대체
                    E_new = E + a_eff * woa
                    with torch.no_grad():
                        uf = (E + R).float()
                        dv = (a_eff * woa).float()
                        den = uf.norm(dim=(-2, -1)).clamp_min(1e-6)
                        self.last_mid_stats_full = {
                            "u_mid": dv.norm(dim=(-2, -1)) / den,
                            "a_mid": nn.functional.cosine_similarity(
                                uf.flatten(1), dv.flatten(1), dim=1),
                            "r_mid": 1 - nn.functional.cosine_similarity(
                                uf.flatten(1), (E_new + R).float().flatten(1), dim=1),
                            "q_mid": (E_new + R).float().norm(dim=(-2, -1)) / den,
                        }
                        if E.shape[0] > 1:
                            def _div(xx):
                                pm = nn.functional.normalize(
                                    xx.float().mean(1), dim=-1)
                                g = pm @ pm.T
                                b_ = g.shape[0]
                                off = g.masked_select(~torch.eye(
                                    b_, dtype=torch.bool, device=g.device))
                                return 1 - off.mean()
                            self.last_mid_stats_full["div_z0"] = _div(
                                states[0]).expand(E.shape[0]).clone()
                            self.last_mid_stats_full["div_z1"] = _div(
                                states[1] if len(states) > 1 else states[0]
                            ).expand(E.shape[0]).clone()
                            self.last_mid_stats_full["div_umid"] = _div(
                                E + R).expand(E.shape[0]).clone()
                    E = E_new
            states.append(R_t)
        # D_{k*->j}: read를 끈 shadow R-chain(E = E^(0) 유지) 대비 발산 —
        # read 갱신이 최종 reasoning까지 남는지의 학습-중 감시.
        with torch.no_grad():
            Rs = states[mid_step - 1] if mid_step > 1 else torch.zeros_like(E)
            prop = []
            for j in range(mid_step, len(states)):
                xs = Rs + e0.detach()
                Rs = self.workspace_transition(xs, q_star, q_padding_mask) - e0.detach()
                prop.append(1 - nn.functional.cosine_similarity(
                    states[j].float().flatten(1), Rs.float().flatten(1), dim=1))
            self.last_mid_prop_full = torch.stack(prop, dim=1)
        self.last_reread_gates = []
        self.last_reread_update_mags = []
        if getattr(self.config, "concat_aggregation", False):
            # ER v2.2 C 확정판 (§6/§10): R_agg = W_agg · Concat(R1..RT).
            # R0/E/V* 미포함 (§10); states는 single-chain이라 answer 체인 =
            # grad 경로 → L_ans가 aggregator를 거쳐 R1..RT 각각에 직접 도달 (§14).
            r_steps = states[1:]  # R1..RT, 각 [B, S, d_w]
            r_cat = torch.cat(r_steps, dim=-1)  # [B, S, T*d_w]
            R_final = self.reasoning_aggregator(r_cat)  # [B, S, d_w]
            with torch.no_grad():
                # §18: step별 block norm ||W_k|| + cos(R_agg, R_k)
                d_w = R.shape[-1]
                w = self.reasoning_aggregator.weight.float()  # [D, T*D]
                self._er_agg_stats = {
                    **{f"aggW{_i + 1}": float(w[:, _i * d_w:(_i + 1) * d_w].norm())
                       for _i in range(len(r_steps))},
                    "agg_norm": float(R_final.float().norm(dim=(-2, -1)).mean()),
                    **{f"agg_cos_R{_i + 1}": float(nn.functional.cosine_similarity(
                        r_steps[_i].float().flatten(1),
                        R_final.float().flatten(1), dim=1).mean())
                       for _i in range(len(r_steps))},
                }
            self._er_answer_final = R_final
        elif getattr(self.config, "temporal_aggregation", False):
            # (legacy — softmax 변형)
            a_w = torch.softmax(self.temporal_logits.float(), dim=0).to(R.dtype)  # [T]
            r_steps = states[1:]  # R1..RT
            R_final = r_steps[0] * a_w[0]
            for _i in range(1, len(r_steps)):
                R_final = R_final + a_w[_i] * r_steps[_i]
            self._er_answer_final = R_final
        else:
            self._er_answer_final = R
        self._er_last_evidence = E.detach()
        return states

    def _encode_counterfactual_residual_recurrence(
        self,
        q_text: torch.Tensor,
        q_mm: torch.Tensor,
        q_padding_mask: torch.Tensor | None,
        ctx,
        visual_bundle=None,
        component_ablation_mode: str | None = None,
    ) -> list[torch.Tensor]:
        """Run the native-layer counterfactual residual recurrence.

        ``F=L(ell_star+1)`` is called as the one shared recurrent operator.  Its
        text-only call populates F's prefix cache, so answer tokens may traverse
        F without receiving an image-conditioned KV bypass.  Every other F call
        uses a cache-free context and therefore changes only the returned
        question-shaped state.

        Returned states are ``[B, C1, ..., CT]``:

        ``B=F(Q_text)``, ``C1=F(Q_mm)-B``,
        ``G_B(C)=F(B+C)-F(B)``, and
        ``Ck=(1-beta)C{k-1}+beta*G_B(C{k-1})``.
        """
        if getattr(
            self.config, "visual_counterfactual_recurrence", False
        ):
            if visual_bundle is None:
                raise RuntimeError(
                    "visual counterfactual recurrence is missing its multimodal "
                    "boundary scaffold"
                )
            return self._encode_visual_counterfactual_recurrence(
                q_text,
                q_padding_mask,
                ctx,
                visual_bundle,
                component_ablation_mode=component_ablation_mode,
            )
        if component_ablation_mode is not None:
            raise ValueError(
                "component ablations require visual counterfactual recurrence"
            )

        steps = int(self.config.num_workspace_steps)
        beta = float(self.config.counterfactual_beta)
        cell_index = int(self.config.ell_star) + 1
        tm = self.text_model

        # This is the only call that writes L21 prefix KV.  Its input is
        # strictly text-only, preserving the no-image lower-cache invariant.
        base_anchor = run_layer_range(
            tm,
            ctx,
            cell_index,
            cell_index + 1,
            hidden_states=q_text,
            use_cache=True,
        )

        # Recurrent evaluations must neither read nor mutate the replacement
        # cache.  The causal mask and RoPE tables remain the same Q-only ones.
        recurrent_ctx = replace(ctx, past_key_values=None)

        def cell(hidden: torch.Tensor) -> torch.Tensor:
            return run_layer_range(
                tm,
                recurrent_ctx,
                cell_index,
                cell_index + 1,
                hidden_states=hidden,
                use_cache=False,
            )

        factual_t1 = cell(q_mm)
        residual = factual_t1 - base_anchor
        if q_padding_mask is not None:
            residual = residual.masked_fill(q_padding_mask.unsqueeze(-1), 0)
        states = [base_anchor, residual]

        if steps > 1:
            fixed_anchor_next = cell(base_anchor)
            for _ in range(1, steps):
                proposed = cell(base_anchor + residual) - fixed_anchor_next
                if beta == 1.0:
                    residual = proposed
                elif beta == 0.0:
                    residual = residual
                else:
                    residual = torch.lerp(residual, proposed, beta)
                if hasattr(self, "crr_policy_adapter"):
                    residual = self.crr_policy_adapter(residual)
                if q_padding_mask is not None:
                    residual = residual.masked_fill(
                        q_padding_mask.unsqueeze(-1), 0
                    )
                states.append(residual)

        return states

    def _encode_visual_counterfactual_recurrence(
        self,
        q_text: torch.Tensor,
        q_padding_mask: torch.Tensor | None,
        text_ctx,
        visual_bundle,
        *,
        component_ablation_mode: str | None = None,
    ) -> list[torch.Tensor]:
        """Run Causal Visual Residual Recurrence with one native layer.

        The first state is the ordinary multimodal output of
        ``F=L(ell_star+1)``.  For every later step, the exact same physical
        multimodal scaffold is evaluated twice from the current question
        state: once normally and once with only text-query -> visual-key/value
        edges blocked.  Their difference is therefore the update attributable
        to a fresh visual read through ``F``:

        ``R1 = Pi_Q F(H_mm*)``
        Historical CVRR uses the intervention difference
        ``Delta_k^V = Pi_Q(F_on([E;R{k-1}]) - F_off([E;R{k-1}]))``.
        Full-state CVRR instead keeps the complete native computation:
        ``R_tilde_k = Pi_Q F_on([E;R{k-1}])`` and
        ``R_k = R_{k-1} + beta (R_tilde_k - R_{k-1})``.  The latter removes
        the visual-off branch from train/inference; it remains an evaluation
        intervention only.  The visual-preserving adapter-correction subtype
        instead uses
        ``Delta_k=F_{base+LoRA,on}([E;R{k-1}])-F_{base,off}([E;R{k-1}])``:
        this cancels frozen non-visual drift without cancelling the pretrained
        image-caused transition.

        ``E`` is the persistent visual portion of the first native output.
        It is never cached for, concatenated into, or exposed to L22+; only the
        selected question-shaped recurrent state reaches the answer decoder
        through CRR's strict interface.
        """
        steps = int(self.config.num_workspace_steps)
        beta = float(self.config.counterfactual_beta)
        cumulative = bool(
            getattr(self.config, "visual_cumulative_recurrence", False)
        )
        full_state = bool(
            getattr(self.config, "visual_full_state_recurrence", False)
        )
        disable_visual_reread = (
            component_ablation_mode == "without_visual_reread"
        )
        if disable_visual_reread and not full_state:
            raise ValueError(
                "without_visual_reread is defined for full-state CVRR only"
            )
        boundary_reentry = bool(
            getattr(self.config, "visual_boundary_reentry", False)
        )
        source_centered_reentry = bool(
            getattr(
                self.config,
                "visual_source_centered_reentry",
                False,
            )
        )
        recurrence_only_adapter = bool(
            getattr(self.config, "recurrence_only_adapter", False)
        )
        visual_preserving_correction = bool(
            getattr(
                self.config,
                "visual_preserving_adapter_correction",
                False,
            )
        )
        cell_index = int(self.config.ell_star) + 1
        tm = self.text_model
        mm_lower, mm_ctx, text_rows, visual_rows = visual_bundle

        # Populate L21's answer-time prefix cache with text-only states.  This
        # is the sole cached call at the recurrent layer.
        with self._recurrence_adapter_execution(False):
            if recurrence_only_adapter:
                with torch.no_grad():
                    base_anchor = run_layer_range(
                        tm,
                        text_ctx,
                        cell_index,
                        cell_index + 1,
                        hidden_states=q_text,
                        use_cache=True,
                    )
            else:
                base_anchor = run_layer_range(
                    tm,
                    text_ctx,
                    cell_index,
                    cell_index + 1,
                    hidden_states=q_text,
                    use_cache=True,
                )

        normal_ctx = replace(mm_ctx, past_key_values=None)
        def initial_native_read(hidden: torch.Tensor) -> torch.Tensor:
            return run_layer_range(
                tm,
                normal_ctx,
                cell_index,
                cell_index + 1,
                hidden_states=hidden,
                use_cache=False,
            )

        # This is the only CVRR pass that needs outputs for every visual row
        # (to construct persistent E).  Saving its 5x-width MLP activations is
        # the dominant long-image memory cost; non-reentrant checkpointing
        # recomputes exactly this one native layer during backward instead.
        if recurrence_only_adapter:
            # R1 is an exact pretrained multimodal anchor.  Adapter gradients
            # begin only at the R1->R2 transition.
            with self._recurrence_adapter_execution(False), torch.no_grad():
                first_full = initial_native_read(mm_lower)
        elif self.training and torch.is_grad_enabled():
            first_full = checkpoint(
                initial_native_read,
                mm_lower,
                use_reentrant=False,
                preserve_rng_state=True,
            )
        else:
            first_full = initial_native_read(mm_lower)
        r1, r1_padding = select_tokens_padded(first_full, text_rows)
        if q_padding_mask is not None and not torch.equal(
            r1_padding, q_padding_mask
        ):
            raise ValueError(
                "multimodal and text-only padding disagree in visual recurrence"
            )

        if component_ablation_mode == "without_entire_visual_path":
            # Matched T-step text-only recurrence.  The multimodal branch above
            # is computed only to preserve the common evaluation setup; none
            # of its states enter this returned decoder prefix.
            text_recurrent_ctx = replace(text_ctx, past_key_values=None)
            r_current = base_anchor
            residual = torch.zeros_like(base_anchor)
            states = [base_anchor, residual]
            for _ in range(1, steps):
                with self._recurrence_adapter_execution(True):
                    proposed = run_layer_range(
                        tm,
                        text_recurrent_ctx,
                        cell_index,
                        cell_index + 1,
                        hidden_states=r_current,
                        use_cache=False,
                    )
                r_current = r_current + beta * (proposed - r_current)
                residual = r_current - base_anchor
                if q_padding_mask is not None:
                    residual = residual.masked_fill(
                        q_padding_mask.unsqueeze(-1), 0
                    )
                states.append(residual)
            return states

        if component_ablation_mode == "without_native_mm_initialization":
            r_current = base_anchor
            residual = torch.zeros_like(base_anchor)
        else:
            r_current = r1
            residual = r1 - base_anchor
        if q_padding_mask is not None:
            residual = residual.masked_fill(q_padding_mask.unsqueeze(-1), 0)
        states = [base_anchor, residual]
        if component_ablation_mode == "without_recurrence":
            # Keep the serialized state-list contract while decoding the same
            # native R1 at every nominal slot; no recurrent layer is called.
            states.extend(residual for _ in range(1, steps))
            return states
        if component_ablation_mode == "without_visual_rows_keep_h1":
            # Retain the exact native multimodal R1, then physically remove the
            # visual rows for every later transition.  The learned recurrent
            # layer, T, beta, and strict decoder stay unchanged; recurrence is
            # evaluated on the depth-matched text-only sequence geometry.
            text_recurrent_ctx = replace(text_ctx, past_key_values=None)
            for _ in range(1, steps):
                with self._recurrence_adapter_execution(True):
                    proposed = run_layer_range(
                        tm,
                        text_recurrent_ctx,
                        cell_index,
                        cell_index + 1,
                        hidden_states=r_current,
                        use_cache=False,
                    )
                r_current = r_current + beta * (proposed - r_current)
                residual = r_current - base_anchor
                if q_padding_mask is not None:
                    residual = residual.masked_fill(
                        q_padding_mask.unsqueeze(-1), 0
                    )
                states.append(residual)
            return states
        update_rms: list[torch.Tensor] = []
        update_relative: list[torch.Tensor] = []

        boundary_q = None
        if boundary_reentry:
            boundary_q, boundary_padding = select_tokens_padded(
                mm_lower, text_rows
            )
            if not torch.equal(boundary_padding, r1_padding):
                raise RuntimeError(
                    "post-L20 and post-L21 question row layouts disagree"
                )

        # The source-centered control separates the fixed text-only anchor
        # from the native image-dependent initial state:
        #   A = F_text(H20), Delta_1 = R1 - A.
        # It also evaluates the recurrent cell once at zero deviation so every
        # later update can subtract that response.  Consequently G(0)=0 even
        # after the recurrence-only adapter has learned a nonzero mapping.
        source_anchor_q = None
        if source_centered_reentry and steps > 1:
            def source_anchor(
                hidden: torch.Tensor,
            ) -> tuple[torch.Tensor, torch.Tensor]:
                with self._recurrence_adapter_execution(True):
                    output, _, padding = self._visual_on_off_native_text_rows(
                        hidden,
                        normal_ctx,
                        text_rows,
                        visual_rows,
                        cell_index,
                        compute_off=False,
                    )
                if output is None:
                    raise RuntimeError(
                        "source-centered anchor returned no output"
                    )
                return output, padding

            if self.training and torch.is_grad_enabled():
                source_anchor_q, source_anchor_padding = checkpoint(
                    source_anchor,
                    mm_lower,
                    use_reentrant=False,
                    preserve_rng_state=True,
                )
            else:
                source_anchor_q, source_anchor_padding = source_anchor(
                    mm_lower
                )
            if not torch.equal(source_anchor_padding, r1_padding):
                raise RuntimeError(
                    "source-centered anchor row layout changed"
                )

        # Historical full-state recurrence reuses post-L21 visual rows and
        # replaces its text rows with the current R.  The boundary-reentry
        # control instead rebuilds every call from the native post-L20 input
        # and carries only R-R1 into those question rows.  Both preserve the
        # original multimodal ordering, masks, and M-RoPE geometry.
        for _ in range(1, steps):
            state_before = r_current
            if boundary_reentry:
                if boundary_q is None:
                    raise RuntimeError("boundary re-entry state was not initialized")
                if source_centered_reentry:
                    # The initial deviation is the native visual effect
                    # R1-A, not zero.  Re-enter L21 in its native H20
                    # coordinates while carrying only this evolving residual.
                    recurrent_q = boundary_q + (r_current - base_anchor)
                else:
                    recurrent_q = boundary_q + (r_current - r1)
                recurrent_input = _replace_selected_rows_from_padded(
                    mm_lower,
                    text_rows,
                    recurrent_q,
                    r1_padding,
                )
            else:
                recurrent_input = _replace_selected_rows_from_padded(
                    first_full,
                    text_rows,
                    r_current,
                    r1_padding,
                )
            if visual_preserving_correction:
                # The two executions intentionally use different adapter
                # routing.  Their difference retains the pretrained visual
                # effect while cancelling the frozen layer's non-visual drift:
                #   F_{base+LoRA,on} - F_{base,off}
                # = (F_{base,on} - F_{base,off})
                #   + (F_{base+LoRA,on} - F_{base,on}).
                # Checkpoint both branches because retaining two native-layer
                # activation sets would otherwise roughly double the long-image
                # recurrent memory footprint.
                def adapter_on(
                    hidden: torch.Tensor,
                ) -> tuple[torch.Tensor, torch.Tensor]:
                    with self._recurrence_adapter_execution(True):
                        output, _, padding = (
                            self._visual_on_off_native_text_rows(
                                hidden,
                                normal_ctx,
                                text_rows,
                                visual_rows,
                                cell_index,
                                compute_off=False,
                            )
                        )
                    if output is None:
                        raise RuntimeError("visual-on branch returned no output")
                    return output, padding

                def frozen_visual_off(
                    hidden: torch.Tensor,
                ) -> tuple[torch.Tensor, torch.Tensor]:
                    with self._recurrence_adapter_execution(False):
                        _, output, padding = (
                            self._visual_on_off_native_text_rows(
                                hidden,
                                normal_ctx,
                                text_rows,
                                visual_rows,
                                cell_index,
                                compute_on=False,
                            )
                        )
                    if output is None:
                        raise RuntimeError("visual-off branch returned no output")
                    return output, padding

                if self.training and torch.is_grad_enabled():
                    on_q, on_padding = checkpoint(
                        adapter_on,
                        recurrent_input,
                        use_reentrant=False,
                        preserve_rng_state=True,
                    )
                    off_q, off_padding = checkpoint(
                        frozen_visual_off,
                        recurrent_input,
                        use_reentrant=False,
                        preserve_rng_state=True,
                    )
                else:
                    on_q, on_padding = adapter_on(recurrent_input)
                    off_q, off_padding = frozen_visual_off(recurrent_input)
                if not torch.equal(on_padding, off_padding):
                    raise RuntimeError(
                        "visual-on and frozen visual-off row layouts disagree"
                    )
            else:
                with self._recurrence_adapter_execution(
                    component_ablation_mode != "without_learned_transition"
                ):
                    if full_state:
                        if disable_visual_reread:
                            _, off_q, on_padding = (
                                self._visual_on_off_native_text_rows(
                                    recurrent_input,
                                    normal_ctx,
                                    text_rows,
                                    visual_rows,
                                    cell_index,
                                    compute_on=False,
                                    compute_off=True,
                                )
                            )
                            if off_q is None:
                                raise RuntimeError(
                                    "visual-off recurrent branch returned no output"
                                )
                            # Preserve the complete learned recurrent layer and
                            # MLP while removing only Q->visual K/V edges.
                            on_q = off_q
                        else:
                            on_q, off_q, on_padding = (
                                self._visual_on_off_native_text_rows(
                                    recurrent_input,
                                    normal_ctx,
                                    text_rows,
                                    visual_rows,
                                    cell_index,
                                    compute_off=False,
                                )
                            )
                    else:
                        on_q, off_q, on_padding = (
                            self._visual_on_off_native_text_rows(
                                recurrent_input,
                                normal_ctx,
                                text_rows,
                                visual_rows,
                                cell_index,
                            )
                        )
                if on_q is None:
                    raise RuntimeError("visual-on branch returned no output")
            if visual_preserving_correction:
                transition_delta = on_q - off_q
            elif full_state:
                if source_centered_reentry:
                    if source_anchor_q is None:
                        raise RuntimeError(
                            "source-centered anchor was not initialized"
                        )
                    # G(Delta) = F(H20 + Delta) - F(H20), hence G(0)=0.
                    # The state update below gives
                    # Delta_k = Delta_{k-1} + beta * G(Delta_{k-1}).
                    transition_delta = on_q - source_anchor_q
                else:
                    transition_delta = on_q - state_before
            else:
                if off_q is None:
                    raise RuntimeError("visual-off branch was not computed")
                transition_delta = on_q - off_q
            if q_padding_mask is not None:
                transition_delta = transition_delta.masked_fill(
                    q_padding_mask.unsqueeze(-1), 0
                )
            if not torch.equal(on_padding, r1_padding):
                raise RuntimeError("recurrent text-row layout changed")

            if full_state or cumulative:
                r_current = state_before + beta * transition_delta
            else:
                r_current = r1 + beta * transition_delta
            residual = r_current - base_anchor
            if q_padding_mask is not None:
                residual = residual.masked_fill(
                    q_padding_mask.unsqueeze(-1), 0
                )
            states.append(residual)

            with torch.no_grad():
                metric_delta = (
                    beta * transition_delta if full_state else transition_delta
                )
                valid = (
                    torch.ones_like(r1_padding, dtype=torch.bool)
                    if q_padding_mask is None
                    else ~q_padding_mask
                )
                valid_f = valid.unsqueeze(-1).float()
                denom = valid_f.sum(dim=(1, 2)).clamp_min(1.0)
                delta_rms = (
                    (metric_delta.float().square() * valid_f).sum(dim=(1, 2))
                    / (denom * metric_delta.shape[-1])
                ).sqrt()
                state_rms = (
                    (state_before.float().square() * valid_f).sum(dim=(1, 2))
                    / (denom * state_before.shape[-1])
                ).sqrt()
                update_rms.append(delta_rms)
                update_relative.append(delta_rms / state_rms.clamp_min(1e-8))

        update_rms_tensor = torch.stack(update_rms, dim=1) if update_rms else None
        update_relative_tensor = (
            torch.stack(update_relative, dim=1) if update_relative else None
        )
        if full_state:
            self.last_recurrent_update_rms = update_rms_tensor
            self.last_recurrent_update_relative = update_relative_tensor
            self.last_visual_counterfactual_update_rms = None
            self.last_visual_counterfactual_update_relative = None
        else:
            self.last_visual_counterfactual_update_rms = update_rms_tensor
            self.last_visual_counterfactual_update_relative = update_relative_tensor
            self.last_recurrent_update_rms = None
            self.last_recurrent_update_relative = None
        return states

    def _visual_on_off_native_text_rows(
        self,
        full_input: torch.Tensor,
        full_ctx,
        text_rows: torch.Tensor,
        visual_rows: torch.Tensor,
        cell_index: int,
        compute_on: bool = True,
        compute_off: bool = True,
    ) -> tuple[
        torch.Tensor | None,
        torch.Tensor | None,
        torch.Tensor,
    ]:
        """Exact text-row execution of one native visual-on/off layer pair.

        Decoder-layer attention and MLP are pointwise in their query/output
        rows.  Consequently, computing Q and the post-attention MLP only for
        text rows while computing K/V for the complete multimodal sequence is
        algebraically identical to two full layer calls followed by ``Pi_Q``.
        Sharing the common Q/K/V projections between the on/off branches also
        preserves gradients: both masks consume the same projected tensors and
        autograd sums their two uses exactly as it would for tied weights.

        This is an execution optimization only.  It removes six discarded
        visual-row MLP activations per T=4 forward, which otherwise dominate
        memory on long 8k-token images.
        """
        if not compute_on and not compute_off:
            raise ValueError("at least one visual on/off branch must be computed")

        layer = self.text_model.layers[cell_index]
        attention = layer.self_attn
        if getattr(layer, "attention_type", "full_attention") != "full_attention":
            raise NotImplementedError(
                "selective CVRR execution currently requires full attention"
            )

        normalized_full = layer.input_layernorm(full_input)
        normalized_q, q_padding = select_tokens_padded(
            normalized_full, text_rows
        )
        residual_q, residual_padding = select_tokens_padded(full_input, text_rows)
        if not torch.equal(q_padding, residual_padding):
            raise RuntimeError("native selective-query padding changed")

        batch_size, query_length, _ = normalized_q.shape
        full_length = normalized_full.shape[1]
        head_dim = attention.head_dim
        query = attention.q_proj(normalized_q).view(
            batch_size, query_length, -1, head_dim
        ).transpose(1, 2)
        key = attention.k_proj(normalized_full).view(
            batch_size, full_length, -1, head_dim
        ).transpose(1, 2)
        value = attention.v_proj(normalized_full).view(
            batch_size, full_length, -1, head_dim
        ).transpose(1, 2)

        cos_full, sin_full = full_ctx.position_embeddings
        def select_rope_rows(
            rope: torch.Tensor,
        ) -> tuple[torch.Tensor, torch.Tensor]:
            # Qwen2.5 multimodal RoPE keeps the t/h/w axis explicitly as
            # [3,B,L,D]; text-only variants may already be [B,L,D].
            if rope.ndim == 3:
                return select_tokens_padded(rope, text_rows)
            if rope.ndim != 4 or rope.shape[0] != 3:
                raise ValueError(
                    f"unexpected M-RoPE tensor shape {tuple(rope.shape)}"
                )
            selected = []
            selected_padding = None
            for axis in range(3):
                rows, padding = select_tokens_padded(rope[axis], text_rows)
                selected.append(rows)
                if selected_padding is None:
                    selected_padding = padding
                elif not torch.equal(selected_padding, padding):
                    raise RuntimeError("M-RoPE axis padding changed")
            return torch.stack(selected, dim=0), selected_padding

        cos_q, cos_padding = select_rope_rows(cos_full)
        sin_q, sin_padding = select_rope_rows(sin_full)
        if not torch.equal(cos_padding, q_padding) or not torch.equal(
            sin_padding, q_padding
        ):
            raise RuntimeError("native selective-query RoPE padding changed")
        mrope_section = attention.rope_scaling["mrope_section"]
        query, _ = apply_multimodal_rotary_pos_emb(
            query,
            query,
            cos_q,
            sin_q,
            mrope_section,
        )
        _, key = apply_multimodal_rotary_pos_emb(
            key,
            key,
            cos_full,
            sin_full,
            mrope_section,
        )

        # Rectangular causal mask at each text row's original physical index.
        # Padded gathered rows get one harmless fallback key and are zeroed
        # after the layer, avoiding all-masked-row kernel corner cases.
        key_indices = torch.arange(full_length, device=full_input.device)
        query_indices = torch.zeros(
            batch_size,
            query_length,
            dtype=torch.long,
            device=full_input.device,
        )
        for batch_index in range(batch_size):
            physical = text_rows[batch_index].nonzero(as_tuple=False).squeeze(-1)
            query_indices[batch_index, : physical.numel()] = physical
        visible = key_indices.view(1, 1, full_length) <= query_indices.unsqueeze(-1)
        if full_ctx.attention_mask is not None:
            if full_ctx.attention_mask.shape != (batch_size, full_length):
                raise ValueError(
                    "selective native execution requires a 2-D [B,L] mask"
                )
            visible &= full_ctx.attention_mask[:, None, :].bool()
        valid_query = ~q_padding
        visible &= valid_query.unsqueeze(-1)
        for batch_index in range(batch_size):
            if bool(q_padding[batch_index].any()):
                visible[batch_index, q_padding[batch_index], 0] = True
        normal_mask = visible.unsqueeze(1)
        visual_off_mask = (
            normal_mask & ~visual_rows[:, None, None, :]
            if compute_off
            else None
        )

        def finish(mask: torch.Tensor) -> torch.Tensor:
            attended, _ = sdpa_attention_forward(
                attention,
                query,
                key,
                value,
                mask,
                dropout=(
                    attention.attention_dropout if self.training else 0.0
                ),
                scaling=attention.scaling,
            )
            attended = attention.o_proj(
                attended.reshape(batch_size, query_length, -1).contiguous()
            )
            hidden = residual_q + attended
            hidden = hidden + layer.mlp(layer.post_attention_layernorm(hidden))
            return hidden.masked_fill(q_padding.unsqueeze(-1), 0)

        normal_output = finish(normal_mask) if compute_on else None
        off_output = finish(visual_off_mask) if visual_off_mask is not None else None
        return normal_output, off_output, q_padding

    def _counterfactual_transition_rollout_prefix(
        self,
        states: list[torch.Tensor] | tuple[torch.Tensor, ...],
        donors: torch.LongTensor,
        transition_index: int,
        question_attention_mask: torch.Tensor | None,
    ) -> torch.Tensor:
        """Intervene on one CRR transition and roll its causal tail forward.

        ``states`` is the factual ``[B, C1, ..., CT]`` trajectory.  At the
        selected 1-indexed transition ``k``, ``Ck`` is replaced by a detached
        different-answer donor state.  Earlier factual states remain fixed and
        later states are recomputed with the target example's anchor and the
        same shared native recurrent cell.  The returned tensor is the exact
        final decoder prefix ``B + Aggregate(C1..CT)`` under that intervention.

        Detaching the donor and factual history prevents the negative branch
        from making another example adversarial or rewriting pre-intervention
        history.  Gradients still pass through the counterfactual tail, the
        shared aggregator (when present), and the upper decoder replay.
        """
        horizon = int(self.config.num_workspace_steps)
        if len(states) != horizon + 1:
            raise ValueError(
                f"expected B plus {horizon} CRR states, got {len(states)}"
            )
        if not 1 <= int(transition_index) <= horizon:
            raise ValueError(
                f"reliance_transition_index must be in [1,{horizon}], "
                f"got {transition_index}"
            )
        batch_size, question_length = states[0].shape[:2]
        if donors.shape != (batch_size,):
            raise ValueError(
                f"donors must have shape {(batch_size,)}, got {tuple(donors.shape)}"
            )

        if question_attention_mask is None:
            q_mask = torch.ones(
                batch_size,
                question_length,
                dtype=torch.long,
                device=states[0].device,
            )
        else:
            q_mask = question_attention_mask
        q_keep = q_mask.gt(0)
        donor_keep = q_keep.index_select(0, donors)
        aligned = q_keep & donor_keep

        k = int(transition_index)
        base_anchor = states[0].detach()
        residual = states[k].index_select(0, donors).detach()
        residual = residual.masked_fill(~aligned.unsqueeze(-1), 0)
        hybrid_residuals = [state.detach() for state in states[1:k]]
        hybrid_residuals.append(residual)

        if k < horizon:
            q_pos_1d = q_keep.long().cumsum(dim=-1) - 1
            q_pos_1d = q_pos_1d.masked_fill(~q_keep, 0)
            q_positions = q_pos_1d.unsqueeze(0).expand(3, -1, -1)
            recurrent_ctx = make_split_context(
                self.text_model,
                base_anchor,
                q_positions,
                q_mask,
            )
            cell_index = int(self.config.ell_star) + 1

            def cell(hidden: torch.Tensor) -> torch.Tensor:
                return run_layer_range(
                    self.text_model,
                    recurrent_ctx,
                    cell_index,
                    cell_index + 1,
                    hidden_states=hidden,
                    use_cache=False,
                )

            fixed_anchor_next = cell(base_anchor)
            beta = float(self.config.counterfactual_beta)
            q_padding_mask = ~q_keep
            for _ in range(k, horizon):
                proposed = cell(base_anchor + residual) - fixed_anchor_next
                if beta == 1.0:
                    residual = proposed
                elif beta == 0.0:
                    residual = residual
                else:
                    residual = torch.lerp(residual, proposed, beta)
                if hasattr(self, "crr_policy_adapter"):
                    residual = self.crr_policy_adapter(residual)
                residual = residual.masked_fill(
                    q_padding_mask.unsqueeze(-1), 0
                )
                hybrid_residuals.append(residual)

        if len(hybrid_residuals) != horizon:
            raise RuntimeError(
                "transition rollout did not reconstruct the full CRR horizon"
            )
        return base_anchor + self._aggregate_crr_residuals(hybrid_residuals)

    def _encode_counterfactual_residual_policy(
        self,
        q_text: torch.Tensor,
        q_mm: torch.Tensor,
        q_padding_mask: torch.Tensor | None,
        ctx,
        *,
        policy_std: float,
        sample: bool,
        sample_mask: torch.Tensor | None = None,
    ) -> tuple[list[torch.Tensor], LatentPolicyTrace]:
        """Sample ``C2..CT`` as an amortized Gaussian latent policy.

        ``B`` and ``C1`` remain the deterministic CRR visual anchor.  For each
        later step, the shared native cell predicts the Gaussian mean

        ``mu_k = (1-beta) C{k-1} + beta G_B(C{k-1})``.

        The action is the next complete residual state
        ``Ck ~ N(mu_k, (policy_std * RMS(C1))^2 I)``.  Sampled states and the
        state supplied to the following transition are detached, as required by
        score-function policy gradients.  This blocks accidental pathwise
        gradients while retaining ``d log pi(Ck|C{k-1}) / d theta`` through the
        policy mean.  Gaussian constants are omitted and the elementwise score
        is averaged over valid latent dimensions to keep its scale independent
        of question length.

        The normal deterministic CRR path is implemented separately above and
        remains untouched.  ``sample=False`` is a structural test/inference
        control: it executes the exact policy mean while still returning a
        trace.
        """
        if not math.isfinite(policy_std) or policy_std <= 0.0:
            raise ValueError("policy_std must be finite and > 0")
        steps = int(self.config.num_workspace_steps)
        if steps < 2:
            raise ValueError("latent policy requires at least two CRR steps")
        beta = float(self.config.counterfactual_beta)
        cell_index = int(self.config.ell_star) + 1
        tm = self.text_model

        recurrent_ctx = replace(ctx, past_key_values=None)

        def cell(hidden: torch.Tensor) -> torch.Tensor:
            return run_layer_range(
                tm,
                recurrent_ctx,
                cell_index,
                cell_index + 1,
                hidden_states=hidden,
                use_cache=False,
            )

        # C1 defines the policy's local visual region and is deliberately not an
        # RL action.  It also fills L21's prefix cache from the text-only branch,
        # exactly as deterministic CRR does.  No graph is retained for B/C1.
        with torch.no_grad():
            base_anchor = run_layer_range(
                tm,
                ctx,
                cell_index,
                cell_index + 1,
                hidden_states=q_text,
                use_cache=True,
            )
            residual = cell(q_mm) - base_anchor
            if q_padding_mask is not None:
                residual = residual.masked_fill(
                    q_padding_mask.unsqueeze(-1), 0
                )

        base_observation = base_anchor.detach()
        residual = residual.detach()
        states = [base_observation, residual]

        if q_padding_mask is None:
            valid = torch.ones(
                residual.shape[:2],
                dtype=torch.bool,
                device=residual.device,
            )
        else:
            valid = ~q_padding_mask
        valid_3d = valid.unsqueeze(-1)
        n_valid_dims = (
            valid.sum(dim=1).to(torch.float32) * residual.shape[-1]
        ).clamp_min(1.0)  # [B]
        c1_square_sum = (
            residual.float().square() * valid_3d
        ).sum(dim=(1, 2))
        c1_rms = (c1_square_sum / n_valid_dims).sqrt().clamp_min(1e-6)
        sigma = (float(policy_std) * c1_rms).view(-1, 1, 1).detach()

        # F(B) is shared by every conditional policy mean.  B is an observation,
        # but the shared cell weights remain differentiable.
        fixed_anchor_next = cell(base_observation)
        log_probs: list[torch.Tensor] = []
        mean_update_rms: list[torch.Tensor] = []
        noise_rms: list[torch.Tensor] = []

        for _ in range(1, steps):
            observation = residual.detach()
            proposed = cell(base_observation + observation) - fixed_anchor_next
            if beta == 1.0:
                policy_mean = proposed
            elif beta == 0.0:
                policy_mean = observation
            else:
                policy_mean = torch.lerp(observation, proposed, beta)
            if hasattr(self, "crr_policy_adapter"):
                policy_mean = self.crr_policy_adapter(policy_mean)
            if q_padding_mask is not None:
                policy_mean = policy_mean.masked_fill(
                    q_padding_mask.unsqueeze(-1), 0
                )

            if sample:
                noise = torch.randn_like(policy_mean)
                if sample_mask is not None:
                    if sample_mask.shape != (policy_mean.shape[0],):
                        raise ValueError(
                            "latent policy sample mask must have shape "
                            f"{(policy_mean.shape[0],)}, got "
                            f"{tuple(sample_mask.shape)}"
                        )
                    noise = noise * sample_mask.to(
                        device=noise.device, dtype=noise.dtype
                    ).view(-1, 1, 1)
                sampled_state = (
                    policy_mean.detach()
                    + sigma.to(policy_mean.dtype) * noise
                )
            else:
                sampled_state = policy_mean.detach()
            if q_padding_mask is not None:
                sampled_state = sampled_state.masked_fill(
                    q_padding_mask.unsqueeze(-1), 0
                )

            # Evaluate the sampled (detached) action under the differentiable
            # policy mean.  Omitting constants/log sigma is exact for gradients
            # because sigma is detached and not a learned parameter.
            standardized = (
                sampled_state.float() - policy_mean.float()
            ) / sigma
            score = -0.5 * standardized.square()
            log_prob = (score * valid_3d).sum(dim=(1, 2)) / n_valid_dims
            log_probs.append(log_prob)

            with torch.no_grad():
                update = (policy_mean.detach() - observation).float()
                perturbation = (
                    sampled_state - policy_mean.detach()
                ).float()
                mean_update_rms.append(
                    ((update.square() * valid_3d).sum(dim=(1, 2)) / n_valid_dims)
                    .sqrt()
                    .div(c1_rms)
                )
                noise_rms.append(
                    (
                        (perturbation.square() * valid_3d).sum(dim=(1, 2))
                        / n_valid_dims
                    )
                    .sqrt()
                    .div(c1_rms)
                )

            residual = sampled_state.detach()
            states.append(residual)

        return states, LatentPolicyTrace(
            log_probs=torch.stack(log_probs, dim=1),
            mean_update_rms=torch.stack(mean_update_rms, dim=1),
            noise_rms=torch.stack(noise_rms, dim=1),
        )

    def _encode_visual_counterfactual_policy(
        self,
        q_text: torch.Tensor,
        q_padding_mask: torch.Tensor | None,
        text_ctx,
        visual_bundle,
        *,
        policy_std: float,
        sample: bool,
        sample_mask: torch.Tensor | None = None,
    ) -> tuple[list[torch.Tensor], LatentPolicyTrace]:
        """Sample the current full-state visual recurrence itself.

        This is the score-function counterpart of
        :meth:`_encode_visual_counterfactual_recurrence`, not the historical
        text-only CRR policy above.  Persistent visual rows ``E`` come from the
        exact pretrained first multimodal read.  ``R1`` is deterministic, and
        for every later transition the existing recurrence-only L21 LoRA
        parameterizes

        ``mu_k = R{k-1} + beta * (F_on([E; R{k-1}])_Q - R{k-1})``.

        The complete next residual state is sampled around that mean.  Sampled
        states are detached before the next transition, so answer rewards can
        reach the recurrent LoRA only through the Gaussian score.  With
        ``sample=False`` this path is exactly the ordinary deterministic
        full-state recurrence.
        """
        if not math.isfinite(policy_std) or policy_std <= 0.0:
            raise ValueError("policy_std must be finite and > 0")
        if not getattr(self.config, "visual_full_state_recurrence", False):
            raise ValueError(
                "visual latent policy requires "
                "visual_full_state_recurrence=True"
            )
        if not getattr(self.config, "recurrence_only_adapter", False):
            raise ValueError(
                "visual latent policy trains the existing recurrence-only "
                "adapter"
            )

        steps = int(self.config.num_workspace_steps)
        if steps < 2:
            raise ValueError("latent policy requires at least two CRR steps")
        beta = float(self.config.counterfactual_beta)
        cell_index = int(self.config.ell_star) + 1
        tm = self.text_model
        mm_lower, mm_ctx, text_rows, visual_rows = visual_bundle

        # B, R1, and persistent E are fixed observations.  In particular, the
        # recurrence LoRA is disabled for these competence-critical reads.
        with self._recurrence_adapter_execution(False), torch.no_grad():
            base_anchor = run_layer_range(
                tm,
                text_ctx,
                cell_index,
                cell_index + 1,
                hidden_states=q_text,
                use_cache=True,
            )
            first_full = run_layer_range(
                tm,
                replace(mm_ctx, past_key_values=None),
                cell_index,
                cell_index + 1,
                hidden_states=mm_lower,
                use_cache=False,
            )

        r1, r1_padding = select_tokens_padded(first_full, text_rows)
        if q_padding_mask is not None and not torch.equal(
            r1_padding, q_padding_mask
        ):
            raise ValueError(
                "multimodal and text-only padding disagree in visual policy"
            )
        residual = r1 - base_anchor
        if q_padding_mask is not None:
            residual = residual.masked_fill(
                q_padding_mask.unsqueeze(-1), 0
            )
        base_anchor = base_anchor.detach()
        residual = residual.detach()
        states = [base_anchor, residual]

        valid = (
            torch.ones(
                residual.shape[:2],
                dtype=torch.bool,
                device=residual.device,
            )
            if q_padding_mask is None
            else ~q_padding_mask
        )
        valid_3d = valid.unsqueeze(-1)
        n_valid_dims = (
            valid.sum(dim=1).to(torch.float32) * residual.shape[-1]
        ).clamp_min(1.0)
        c1_square_sum = (
            residual.float().square() * valid_3d
        ).sum(dim=(1, 2))
        c1_rms = (c1_square_sum / n_valid_dims).sqrt().clamp_min(1e-6)
        sigma = (float(policy_std) * c1_rms).view(-1, 1, 1).detach()

        normal_ctx = replace(mm_ctx, past_key_values=None)
        log_probs: list[torch.Tensor] = []
        mean_update_rms: list[torch.Tensor] = []
        noise_rms: list[torch.Tensor] = []

        for _ in range(1, steps):
            observation = residual.detach()
            state_before = base_anchor + observation
            recurrent_input = _replace_selected_rows_from_padded(
                first_full,
                text_rows,
                state_before,
                r1_padding,
            )
            with self._recurrence_adapter_execution(True):
                on_q, _, on_padding = self._visual_on_off_native_text_rows(
                    recurrent_input,
                    normal_ctx,
                    text_rows,
                    visual_rows,
                    cell_index,
                    compute_off=False,
                )
            if not torch.equal(on_padding, r1_padding):
                raise RuntimeError("recurrent visual-policy row layout changed")

            policy_mean_full = torch.lerp(state_before, on_q, beta)
            policy_mean = policy_mean_full - base_anchor
            if q_padding_mask is not None:
                policy_mean = policy_mean.masked_fill(
                    q_padding_mask.unsqueeze(-1), 0
                )

            if sample:
                noise = torch.randn_like(policy_mean)
                if sample_mask is not None:
                    if sample_mask.shape != (policy_mean.shape[0],):
                        raise ValueError(
                            "latent policy sample mask must have shape "
                            f"{(policy_mean.shape[0],)}, got "
                            f"{tuple(sample_mask.shape)}"
                        )
                    noise = noise * sample_mask.to(
                        device=noise.device, dtype=noise.dtype
                    ).view(-1, 1, 1)
                sampled_state = (
                    policy_mean.detach()
                    + sigma.to(policy_mean.dtype) * noise
                )
            else:
                sampled_state = policy_mean.detach()
            if q_padding_mask is not None:
                sampled_state = sampled_state.masked_fill(
                    q_padding_mask.unsqueeze(-1), 0
                )

            standardized = (
                sampled_state.float() - policy_mean.float()
            ) / sigma
            score = -0.5 * standardized.square()
            log_probs.append(
                (score * valid_3d).sum(dim=(1, 2)) / n_valid_dims
            )

            with torch.no_grad():
                update = (policy_mean.detach() - observation).float()
                perturbation = (
                    sampled_state - policy_mean.detach()
                ).float()
                mean_update_rms.append(
                    (
                        (update.square() * valid_3d).sum(dim=(1, 2))
                        / n_valid_dims
                    ).sqrt().div(c1_rms)
                )
                noise_rms.append(
                    (
                        (perturbation.square() * valid_3d).sum(dim=(1, 2))
                        / n_valid_dims
                    ).sqrt().div(c1_rms)
                )

            residual = sampled_state.detach()
            states.append(residual)

        return states, LatentPolicyTrace(
            log_probs=torch.stack(log_probs, dim=1),
            mean_update_rms=torch.stack(mean_update_rms, dim=1),
            noise_rms=torch.stack(noise_rms, dim=1),
        )

    def _encode_raw_evidence_question(
        self,
        q_star: torch.Tensor,
        v_star: torch.Tensor,
        q_padding_mask: torch.Tensor | None,
        v_padding_mask: torch.Tensor | None,
        gate_overrides: list | None,
        v_star_reread: torch.Tensor | None,
        v_reread_padding_mask: torch.Tensor | None,
        reread_memory_overrides: list | None,
        reread_zero_memory: bool,
        hierarchical_visual_memories: (
            tuple[tuple[torch.Tensor, torch.Tensor], ...] | None
        ) = None,
    ) -> list[torch.Tensor]:
        """Slot-free recurrent read with immutable raw visual evidence.

        Residual mode stores ``[R0=Q*, R1, ..., RT]``.  State-replacement mode
        stores ``[Q*, H1, ..., HT]`` and the decoder later receives ``Q*+HT``.
        Question-anchored replacement stores ``[Z0=Q*, Z1, ..., ZT]`` where
        ``Zk=Q*+CrossAttn(Z{k-1},V*)`` and directly decodes ``ZT``.
        In hierarchical mode the first T-1 steps read their prescribed frozen
        visual memories and the final step reads Q* with the same CrossAttn:
        ``R_T = R_{T-1} + CA(R_{T-1}, Q*)``.  Thus R_T integrates the visual
        trajectory without another image access or another module.  Visual
        memories are never updated or exposed directly to the decoder.
        """
        r = q_star
        states = [r]
        hierarchy = tuple(hierarchical_visual_memories or ())
        if hierarchy:
            expected_visual_reads = self.config.num_workspace_steps - 1
            if len(hierarchy) != expected_visual_reads:
                raise ValueError(
                    "hierarchical visual memories must cover T-1 reads: "
                    f"got {len(hierarchy)} for T={self.config.num_workspace_steps} "
                    f"(expected {expected_visual_reads})"
                )
            if v_star_reread is not None:
                raise ValueError(
                    "hierarchical swaps must use per-step memory overrides"
                )
            base_e = base_pad = base_projected = None
            question_projected = self.raw_eq_transition.project_evidence(q_star)
        else:
            base_e, base_pad = v_star, v_padding_mask
            if v_star_reread is not None:
                base_e, base_pad = v_star_reread, v_reread_padding_mask
            if reread_zero_memory:
                base_e = torch.zeros_like(base_e)
            base_projected = self.raw_eq_transition.project_evidence(base_e)

        for k in range(self.config.num_workspace_steps):
            if hierarchy and k < len(hierarchy):
                e_k, pad_k = hierarchy[k]
                if reread_zero_memory:
                    e_k = torch.zeros_like(e_k)
                projected_k = self.raw_eq_transition.project_evidence(e_k)
            elif hierarchy:
                # Final image-free integration: the evolved state queries the
                # immutable text-only scaffold through the same shared module.
                e_k, pad_k = q_star, q_padding_mask
                projected_k = question_projected
            else:
                e_k, pad_k = base_e, base_pad
                projected_k = base_projected
            if reread_memory_overrides is not None and reread_memory_overrides[k] is not None:
                e_k, pad_k = reread_memory_overrides[k]
                if reread_zero_memory:
                    e_k = torch.zeros_like(e_k)
                projected_k = self.raw_eq_transition.project_evidence(e_k)
            scale = 1.0
            if gate_overrides is not None and gate_overrides[k] is not None:
                # Evaluation-only read ablation/scale; there is no learned gate.
                scale = float(gate_overrides[k])
            r = self.raw_eq_transition(
                r,
                e_k,
                evidence_padding_mask=pad_k,
                question_padding_mask=q_padding_mask,
                read_scale=scale,
                projected_evidence=projected_k,
                question_anchor=(
                    q_star
                    if getattr(self.config, "raw_question_anchor", False)
                    else None
                ),
            )
            states.append(r)

        self.last_reread_gates = []
        self.last_reread_update_mags = []
        return states

    def encode_workspace(
        self,
        q_star: torch.Tensor,  # [B, N_q, d]
        v_star: torch.Tensor,  # [B, N_v, d]
        q_padding_mask: torch.Tensor | None = None,
        v_padding_mask: torch.Tensor | None = None,
        gate_overrides: list | None = None,  # eval-only: per-step None/float
        v_star_reread: torch.Tensor | None = None,  # eval-only: swap rereads' V*
        v_reread_padding_mask: torch.Tensor | None = None,
        reread_memory_overrides: list | None = None,  # eval-only: per-step (v, pad) | None
        reread_zero_memory: bool = False,  # eval-only: V* -> 0 at EVERY reread (spec 10.3)
        final_read_spec: dict | None = None,  # eval-only: latent-necessity battery
        hierarchical_visual_memories: (
            tuple[tuple[torch.Tensor, torch.Tensor], ...] | None
        ) = None,
    ) -> list[torch.Tensor]:
        """``Z^(0..K)``. ``V*`` is consumed here and nowhere else.

        With ``config.read_gating`` the workspace may re-consult ``V*`` between
        transitions through the shared reader (gated, see
        :meth:`WorkspaceRead.reread`); the decoder-facing contract is unchanged
        either way. Returns ``K + 1`` tensors of shape ``[B, S, d_w]``.

        The three eval-only knobs implement §3.2's held-out gate battery:
        ``gate_overrides[k]`` forces step ``k``'s gate (forced-open/closed/
        step-specific), ``v_star_reread`` feeds the REREADS a different visual
        memory while the initial read keeps the true one (mid-recurrence swap).
        Defaults leave the learned path bitwise intact. They are reachable only
        through :meth:`prefill` / :meth:`generate_answer`, never ``forward``,
        so the training loop cannot consume them even by accident.
        """
        if getattr(self.config, "raw_evidence_question", False):
            if final_read_spec is not None:
                raise ValueError("raw_evidence_question does not use final_read_spec")
            return self._encode_raw_evidence_question(
                q_star,
                v_star,
                q_padding_mask,
                v_padding_mask,
                gate_overrides,
                v_star_reread,
                v_reread_padding_mask,
                reread_memory_overrides,
                reread_zero_memory,
                hierarchical_visual_memories,
            )

        z = self.workspace_read(q_star, v_star, q_padding_mask, v_padding_mask)
        gating = getattr(self.config, "read_gating", False)
        pure = getattr(self.config, "read_residual_pure", False)
        verify = getattr(self.config, "final_verify", False)
        mid_step = int(getattr(self.config, "mid_read_step", 0) or 0)
        sid = None
        if getattr(self.config, "persistent_slot_id", False):
            # E_slot을 매 step 입력에 일시 제공 (누적 금지). 저장 state는
            # f_theta 출력 그대로다.
            sid = self.workspace_read.slots.unsqueeze(0).to(z.dtype)  # [1, S, d_w]
        if getattr(self.config, "evidence_reasoning", False):
            # advisor 최종안: Evidence–Reasoning 분리. 전용 경로로 위임 —
            # flag off면 이 아래 기존 코드가 비트동일하게 돈다.
            assert final_read_spec is None, "ER 모드는 final_read_spec 미지원"
            return self._encode_workspace_er(
                z, q_star, v_star, q_padding_mask, v_padding_mask,
                gate_overrides, v_star_reread, v_reread_padding_mask,
                reread_zero_memory, reread_memory_overrides, mid_step,
            )
        if final_read_spec is not None:
            # Latent-necessity battery (advisor): run the ZERO-residual chain
            # U^(0..T) (U^(0) := Z^(0)), then compose one final visual read
            # with independently chosen query source / decoder base:
            #   Z_out = U^(base_t) + alpha * W_o XAttn(LN(query_src), V*).
            # query_override (a tensor) replaces the query source entirely
            # (matched query swap). final_read=False skips the read
            # (latent-only / initial-read-only).
            if not pure and not verify:
                raise RuntimeError("final_read_spec requires a pure verifier branch")
            spec = final_read_spec
            chain = [z]
            for _ in range(self.config.num_workspace_steps):
                chain.append(self.workspace_transition(
                    chain[-1], q_star, q_padding_mask, cond=sid))
            base = chain[int(spec.get("base_t", self.config.num_workspace_steps))]
            if spec.get("final_read", True):
                q_src = spec.get("query_override")
                if q_src is None:
                    q_src = chain[int(spec.get("query_t", self.config.num_workspace_steps))]
                a = self.workspace_read.visual_residual(q_src, v_star, v_padding_mask)
                z_out = base + float(spec.get("alpha", 1.0)) * a
            else:
                z_out = base
            self.last_reread_gates = []
            self.last_reread_update_mags = []
            self.last_reread_gates_full = None
            self.last_reread_update_mags_full = None
            self.last_chain_states = chain  # U^(0..T), query-swap donor 추출용
            return chain + [z_out]
        if not gating and not pure and not verify and not mid_step:
            del v_star  # make the read-once contract visible at the call site
        states = [z]
        gates: list[torch.Tensor] = []
        upd_mags: list[float] = []
        upd_full: list[torch.Tensor] = []
        for k in range(self.config.num_workspace_steps):
            u = self.workspace_transition(z, q_star, q_padding_mask, cond=sid)
            z = u
            if mid_step and (k + 1) == mid_step:
                # MAIN: single intermediate pure read at LOGICAL step k*
                # (python index k = k*-1; the +1 here is the off-by-one guard
                # the directive calls out). Eval hooks share the step-wise
                # keying: gate_overrides[k]==0.0 -> mid-off,
                # reread_memory_overrides[k] -> blank/swap, zero flag -> zero.
                off_mid = gate_overrides is not None and gate_overrides[k] == 0.0
                if not off_mid:
                    v_r, vp_r = v_star, v_padding_mask
                    if v_star_reread is not None:
                        v_r, vp_r = v_star_reread, v_reread_padding_mask
                    if reread_zero_memory:
                        v_r = torch.zeros_like(v_r)
                    elif (reread_memory_overrides is not None
                          and reread_memory_overrides[k] is not None):
                        v_r, vp_r = reread_memory_overrides[k]
                    z = self.workspace_read.reread_pure(
                        u, v_r, vp_r,
                        query_state=None if sid is None else u + sid)
                    # read 스케일: 학습은 config.mid_read_alpha, eval은
                    # gate_overrides가 절대값으로 대체 (sweep 의미 보존).
                    # reread_pure가 z = u + WoA이므로 (z-u) = WoA 정확.
                    a_eff = float(getattr(self.config, "mid_read_alpha", 1.0))
                    if gate_overrides is not None and gate_overrides[k] is not None:
                        a_eff = float(gate_overrides[k])
                    if a_eff != 1.0:
                        z = u + a_eff * (z - u)
                with torch.no_grad():
                    dv = (z - u).float()
                    uf = u.float()
                    den = uf.norm(dim=(-2, -1)).clamp_min(1e-6)
                    self.last_mid_stats_full = {
                        "u_mid": dv.norm(dim=(-2, -1)) / den,
                        "a_mid": nn.functional.cosine_similarity(
                            uf.flatten(1), dv.flatten(1), dim=1),
                        "r_mid": 1 - nn.functional.cosine_similarity(
                            uf.flatten(1), z.float().flatten(1), dim=1),
                        "q_mid": z.float().norm(dim=(-2, -1)) / den,
                    }
                    # 시간축 diversity 감시 (advisor): 배치 내 항목쌍은 무관
                    # 입력이므로, pooled state의 평균 pairwise cos-dist가
                    # d_unrelated의 학습 중 근사가 된다. pre-read 수축이
                    # 스케줄 교정으로 완화되는지 이 곡선으로 판정.
                    if u.shape[0] > 1:
                        def _div(x):  # [B, S, d_w] -> 평균 pairwise 1-cos
                            pm = nn.functional.normalize(
                                x.float().mean(1), dim=-1)  # [B, d_w]
                            g = pm @ pm.T
                            b = g.shape[0]
                            off = g.masked_select(
                                ~torch.eye(b, dtype=torch.bool, device=g.device))
                            return 1 - off.mean()
                        self.last_mid_stats_full["div_z0"] = _div(
                            states[0]).expand(u.shape[0]).clone()
                        self.last_mid_stats_full["div_z1"] = _div(
                            states[1] if len(states) > 1 else states[0]
                        ).expand(u.shape[0]).clone()
                        self.last_mid_stats_full["div_umid"] = _div(u).expand(
                            u.shape[0]).clone()
            if pure:
                # Eval-only residual scale (advisor alpha-sweep): a non-None
                # gate_overrides[k] acts as alpha in Z = U + alpha*W_o A.
                # alpha = 0.0 skips the branch entirely (== read_once), so the
                # sweep's zero point and reread-off are one code path.
                alpha_k = None
                if gate_overrides is not None and gate_overrides[k] is not None:
                    alpha_k = float(gate_overrides[k])
                if alpha_k == 0.0:
                    states.append(z)
                    continue
                # all-steps substitute (global null/swap) -- the v7-ep1 run
                # exposed that this was wired only into the gated branch:
                # null_mem came back bit-identical to native (0.0% changes),
                # which is impossible with a real substitution.
                v_r, vp_r = v_star, v_padding_mask
                if v_star_reread is not None:
                    v_r, vp_r = v_star_reread, v_reread_padding_mask
                if reread_zero_memory:
                    v_r, vp_r = torch.zeros_like(v_star), v_padding_mask
                elif reread_memory_overrides is not None and reread_memory_overrides[k] is not None:
                    v_r, vp_r = reread_memory_overrides[k]
                z = self.workspace_read.reread_pure(u, v_r, vp_r)
                rho = float(getattr(self.config, "residual_norm_cap", 0.0) or 0.0)
                if rho > 0.0:
                    dv = z - u
                    s_cap = (
                        rho * u.float().pow(2).mean((-2, -1)).sqrt()
                        / dv.float().pow(2).mean((-2, -1)).sqrt().clamp_min(1e-6)
                    ).clamp(max=1.0).to(dv.dtype)  # [B]
                    z = u + s_cap.view(-1, 1, 1) * dv
                if alpha_k is not None and alpha_k != 1.0:
                    z = u + alpha_k * (z - u)
                with torch.no_grad():
                    num = (z - u).float().norm(dim=(-2, -1))
                    den = u.float().norm(dim=(-2, -1)).clamp_min(1e-6)
                    u_item = num / den
                    upd_mags.append(float(u_item.mean()))
                    upd_full.append(u_item)
            if gating:
                # Consultation happens through the same reader and lands in Z;
                # f_theta itself still never receives V* (read-mediation holds,
                # the "once" becomes "as often as the gate earns").
                v_r = v_star if v_star_reread is None else v_star_reread
                vp_r = (
                    v_padding_mask if v_star_reread is None else v_reread_padding_mask
                )
                if reread_memory_overrides is not None and reread_memory_overrides[k] is not None:
                    # step-specific visual ablation: THIS reread consults a
                    # substitute memory (blank / donor), all other steps keep
                    # the true V*. The decisive §3.2 battery test.
                    v_r, vp_r = reread_memory_overrides[k]
                z, g = self.workspace_read.reread(
                    u, q_star, v_r, q_padding_mask, vp_r,
                    gate_override=None if gate_overrides is None else gate_overrides[k],
                )
                gates.append(g)
                # u_k: realized update magnitude ||Z-U||_F / ||U||_F -- the
                # gate says how much the model WANTED to look, u_k what the
                # look actually changed (draft §3.2).
                with torch.no_grad():
                    num = (z - u).float().norm(dim=(-2, -1))
                    den = u.float().norm(dim=(-2, -1)).clamp_min(1e-6)
                    u_item = num / den  # [B]
                    upd_mags.append(float(u_item.mean()))
                    upd_full.append(u_item)
            states.append(z)
        if verify:
            # MAIN: single final visual verification. Mid-loop above ran pure
            # f_theta only (no reader calls -- tests/test_final_verify.py
            # counts them). Eval hooks: gate_overrides[-1]==0.0 -> branch off;
            # reread_memory_overrides[-1] -> substitute verifier memory;
            # reread_zero_memory -> zero tensor.
            off = gate_overrides is not None and gate_overrides[-1] == 0.0
            if not off:
                v_r, vp_r = v_star, v_padding_mask
                if reread_zero_memory:
                    v_r = torch.zeros_like(v_star)
                elif (reread_memory_overrides is not None
                      and reread_memory_overrides[-1] is not None):
                    v_r, vp_r = reread_memory_overrides[-1]
                states.append(self.workspace_read.reread_pure(states[-1], v_r, vp_r))
            else:
                states.append(states[-1])
        if mid_step and len(states) == self.config.num_workspace_steps + 1:
            # D_{k*->j} propagation telemetry (advisor: "step 4 update가 빠르게
            # 소거되면 구조는 아름답지만 답에 기여하지 않는다"). Shadow chain =
            # mid-null (read skipped); identical to clean bitwise before k*, so
            # start it at Z^(k*-1). workspace_dropout=0.0 -> deterministic.
            with torch.no_grad():
                zs = states[mid_step - 1]
                prop = []
                for j in range(mid_step, len(states)):
                    zs = self.workspace_transition(
                        zs, q_star, q_padding_mask, cond=sid)
                    prop.append(1 - nn.functional.cosine_similarity(
                        states[j].float().flatten(1), zs.float().flatten(1), dim=1))
                # [B, T-k*+1]; columns j = k*, ..., T
                self.last_mid_prop_full = torch.stack(prop, dim=1)
        # Telemetry for the trainer/analysis: per-step means, [K] each or [].
        self.last_reread_gates = [float(g.detach().mean()) for g in gates]
        self.last_reread_update_mags = upd_mags
        # Full per-item resolution for offline probes (compact logs only carry
        # the means). [B, K] or None when gating is off.
        self.last_reread_gates_full = (
            torch.stack([g.detach().float() for g in gates], dim=1) if gates else None
        )
        self.last_reread_update_mags_full = (
            torch.stack(upd_full, dim=1) if upd_full else None
        )  # [B, K]
        return states

    # -- prefill -----------------------------------------------------------

    def prefill(
        self,
        input_ids: torch.LongTensor,
        pixel_values: torch.Tensor,
        image_grid_thw: torch.LongTensor,
        question_ids: torch.LongTensor,
        attention_mask: torch.Tensor | None = None,
        question_attention_mask: torch.Tensor | None = None,
        workspace_override: torch.Tensor | None = None,
        gate_overrides: list | None = None,
        reread_inputs: dict | None = None,
        reread_step_inputs: dict | None = None,
        reread_zero_memory: bool = False,
        final_read_spec: dict | None = None,
        return_decoder_latent: bool = False,
        latent_policy_std: float | None = None,
        latent_policy_sample: bool = True,
        latent_policy_sample_mask: torch.Tensor | None = None,
        return_latent_policy_trace: bool = False,
        curriculum_readout_depth: int | torch.Tensor | None = None,
        component_ablation_mode: str | None = None,
    ):
        """Run §3.2 end to end and leave a complete replacement cache behind.

        Shared by training and generation, so the two cannot drift apart.

        Args:
            workspace_override: substitute ``Z^(K)`` (``[B, S, d_w]``). Held-out
                causal evaluation only -- content replacement, matched swap,
                path-leakage. Never set during training.
            gate_overrides: per-step forced gate values (``None`` entries keep
                the learned gate). Held-out gate battery only.
            reread_inputs: multimodal inputs (``input_ids``/``pixel_values``/
                ``image_grid_thw``/``attention_mask``) of a DIFFERENT item; its
                ``V*`` is served to the re-reads while the initial read keeps
                the true one (mid-recurrence swap). Held-out battery only.
            reread_step_inputs: ``{step_k: inputs-dict}`` -- like
                ``reread_inputs`` but applied ONLY at the given steps' rereads
                (step-specific blank/swap ablation). Held-out battery only.

        Returns:
            Normally ``(upper, cache, states, v_pooled)``. With
            ``return_decoder_latent=True``, appends the exact latent supplied to
            the decoder (``R_T`` in the raw-E/Q method). This opt-in fifth
            return value keeps old callers compatible while giving causal
            evaluations a non-ambiguous intervention target.
        """
        cfg = self.config
        allowed_component_modes = {
            None,
            "without_recurrence",
            "without_native_mm_initialization",
            "without_learned_transition",
            "without_entire_visual_path",
            "without_visual_reread",
            "without_visual_rows_keep_h1",
        }
        if component_ablation_mode not in allowed_component_modes:
            raise ValueError(
                f"unknown component_ablation_mode={component_ablation_mode!r}"
            )
        if component_ablation_mode is not None and self.training:
            raise ValueError("component ablations are evaluation-only")
        policy_enabled = latent_policy_std is not None
        if return_latent_policy_trace and not policy_enabled:
            raise ValueError(
                "return_latent_policy_trace requires latent_policy_std"
            )
        if policy_enabled and not getattr(
            cfg, "counterfactual_residual_recurrence", False
        ):
            raise ValueError("latent policy is implemented only for CRR")
        if curriculum_readout_depth is not None and policy_enabled:
            raise ValueError(
                "program curriculum and latent-policy sampling are mutually exclusive"
            )
        if curriculum_readout_depth is not None and workspace_override is not None:
            raise ValueError(
                "program curriculum and workspace_override are mutually exclusive"
            )
        tm = self.text_model
        n_q = question_ids.shape[1]
        # Consume-once training buffers must never survive into a later batch.
        self._functional_q_target = None
        self._functional_q_target_pad = None
        self._functional_q_target_ids = None
        self._functional_state_loss = None
        self._counterfactual_q_state = None
        self._counterfactual_q_pad = None
        self._counterfactual_q_ids = None
        self._visual_counterfactual_bundle = None
        q_pad = (
            None if question_attention_mask is None else (question_attention_mask == 0)
        )  # [B, N_q] True == ignore

        # 1. Frozen visual memory.  The hierarchical raw path captures one
        # memory at each prescribed depth during a single multimodal pass.
        hierarchical_visual_memories = None
        if getattr(cfg, "hierarchical_visual_layers", ()):
            hierarchical_visual_memories = (
                self._encode_hierarchical_visual_memories(
                    input_ids,
                    pixel_values,
                    image_grid_thw,
                    attention_mask,
                )
            )
            # Preserve the historical helper plumbing; the raw encoder below
            # consumes the full tuple rather than this final-memory alias.
            v_star, v_pad = hierarchical_visual_memories[-1]
        else:
            v_star, v_pad = self._encode_visual_memory(
                input_ids, pixel_values, image_grid_thw, attention_mask
            )  # [B, N_v, d], [B, N_v]

        # 2. Text-only lower prefill. Raw-E/Q uses Q rows only; historical slot
        #    methods append fixed cache-shape placeholders.
        q_star, ctx, cache = self._prefill_text_branch(
            question_ids, question_attention_mask
        )  # q_star [B, N_q, d]

        if getattr(cfg, "counterfactual_residual_recurrence", False):
            unsupported = {
                "gate_overrides": gate_overrides is not None,
                "reread_inputs": reread_inputs is not None,
                "reread_step_inputs": reread_step_inputs is not None,
                "reread_zero_memory": bool(reread_zero_memory),
                "final_read_spec": final_read_spec is not None,
            }
            active = [name for name, enabled in unsupported.items() if enabled]
            if active:
                raise ValueError(
                    "CRR has no raw-evidence reread controls; unsupported: "
                    + ", ".join(active)
                )

            q_mm = self._counterfactual_q_state
            q_mm_pad = self._counterfactual_q_pad
            q_mm_ids = self._counterfactual_q_ids
            visual_bundle = self._visual_counterfactual_bundle
            self._counterfactual_q_state = None
            self._counterfactual_q_pad = None
            self._counterfactual_q_ids = None
            self._visual_counterfactual_bundle = None
            if q_mm is None or q_mm_pad is None or q_mm_ids is None:
                raise RuntimeError("CRR factual question state was not captured")
            if q_mm.shape != q_star.shape:
                raise ValueError(
                    "multimodal/text-only question rows disagree after removing "
                    f"the vision span: {tuple(q_mm.shape)} vs {tuple(q_star.shape)}"
                )
            valid = ~q_mm_pad
            if q_pad is not None:
                if q_pad.shape != q_mm_pad.shape or not torch.equal(
                    q_pad, q_mm_pad
                ):
                    raise ValueError(
                        "multimodal and text-only question padding masks disagree"
                    )
                valid &= ~q_pad
            if not torch.equal(q_mm_ids[valid], question_ids[valid]):
                raise ValueError(
                    "text-only token ids do not equal multimodal prompt text rows"
                )

            policy_trace = None
            if policy_enabled:
                if getattr(
                    cfg, "visual_counterfactual_recurrence", False
                ):
                    z_states, policy_trace = (
                        self._encode_visual_counterfactual_policy(
                            q_star,
                            q_pad,
                            ctx,
                            visual_bundle,
                            policy_std=float(latent_policy_std),
                            sample=bool(latent_policy_sample),
                            sample_mask=latent_policy_sample_mask,
                        )
                    )
                else:
                    z_states, policy_trace = (
                        self._encode_counterfactual_residual_policy(
                            q_star,
                            q_mm,
                            q_pad,
                            ctx,
                            policy_std=float(latent_policy_std),
                            sample=bool(latent_policy_sample),
                            sample_mask=latent_policy_sample_mask,
                        )
                    )
            else:
                z_states = self._encode_counterfactual_residual_recurrence(
                    q_star,
                    q_mm,
                    q_pad,
                    ctx,
                    visual_bundle=visual_bundle,
                    component_ablation_mode=component_ablation_mode,
                )
            if curriculum_readout_depth is not None:
                z_final = self.decoder_latent_at_depth(
                    z_states, curriculum_readout_depth
                )
            else:
                z_final = (
                    self.decoder_latent_from_states(z_states)
                    if workspace_override is None
                    else workspace_override
                )
            del v_star
            lower = ctx.hidden_states
            if z_final.shape != lower.shape:
                raise ValueError(
                    "CRR decoder latent must match Q* shape, got "
                    f"{tuple(z_final.shape)} vs {tuple(lower.shape)}"
                )
            # In latent-policy mode the generated answer is a non-differentiable
            # environment outcome.  Detaching here makes the score trace above
            # the sole gradient path even if a caller accidentally leaves an
            # upper-decoder parameter trainable.
            decoder_input = (
                z_final.detach() if policy_enabled else z_final
            ).to(lower.dtype)
            if policy_enabled:
                with torch.no_grad():
                    upper = run_layer_range(
                        tm,
                        ctx,
                        self._upper_decoder_start(),
                        None,
                        hidden_states=decoder_input,
                        use_cache=True,
                    )
            else:
                upper = run_layer_range(
                    tm,
                    ctx,
                    self._upper_decoder_start(),
                    None,
                    hidden_states=decoder_input,
                    use_cache=True,
                )
            if return_latent_policy_trace:
                if return_decoder_latent:
                    return upper, cache, z_states, None, z_final, policy_trace
                return upper, cache, z_states, None, policy_trace
            if return_decoder_latent:
                return upper, cache, z_states, None, z_final
            return upper, cache, z_states, None

        # The removed visual-anchor loss must have zero runtime footprint in
        # the current answer-only method. Historical anchor runs opt in through
        # config.read_anchor.
        v_pooled = None
        if getattr(cfg, "read_anchor", False):
            if v_pad is None:
                v_pooled = v_star.float().mean(1)
            else:
                keep = (~v_pad).to(v_star.dtype).unsqueeze(-1)
                v_pooled = (
                    (v_star * keep).sum(1).float()
                    / keep.sum(1).float().clamp_min(1.0)
                )
            v_pooled = torch.nn.functional.layer_norm(
                v_pooled, (v_pooled.shape[-1],)
            ).detach()

        # 3. Z^(0..K). V* dies here.
        v_swap, v_swap_pad = None, None
        step_overrides = None
        if (
            hierarchical_visual_memories is not None
            and reread_inputs is not None
            and "__v_star__" in reread_inputs
        ):
            raise ValueError(
                "hierarchical recurrence requires one substituted memory per "
                "depth; a single precomputed __v_star__ is ambiguous"
            )
        if reread_inputs is not None and "__v_star__" in reread_inputs:
            # eval-only: precomputed substitute memory tensor (e.g. corpus-mean
            # dictionary), bypassing the image encoder.
            v_swap = reread_inputs["__v_star__"]
            v_swap_pad = reread_inputs.get("__v_pad__")
        elif (
            reread_inputs is not None
            and hierarchical_visual_memories is not None
        ):
            swapped_memories = self._encode_hierarchical_visual_memories(
                reread_inputs["input_ids"],
                reread_inputs["pixel_values"],
                reread_inputs["image_grid_thw"],
                reread_inputs.get("attention_mask"),
            )
            step_overrides = [None] * self.config.num_workspace_steps
            step_overrides[: len(swapped_memories)] = swapped_memories
        elif reread_inputs is not None:
            v_swap, v_swap_pad = self._encode_visual_memory(
                reread_inputs["input_ids"],
                reread_inputs["pixel_values"],
                reread_inputs["image_grid_thw"],
                reread_inputs.get("attention_mask"),
            )
        if reread_step_inputs is not None:
            if step_overrides is None:
                step_overrides = [None] * self.config.num_workspace_steps
            for k_, inp in reread_step_inputs.items():
                step_index = int(k_)
                if not 0 <= step_index < self.config.num_workspace_steps:
                    raise ValueError(
                        "reread step index must be in "
                        f"[0,{self.config.num_workspace_steps - 1}], got "
                        f"{step_index}"
                    )
                if "__v_star__" in inp:
                    # eval-only: direct tensor (e.g. true zero memory) -- the
                    # blank IMAGE still yields nonzero V*, so "zero" must be
                    # injected as a tensor, not an image (advisor §4).
                    step_overrides[step_index] = (
                        inp["__v_star__"], inp.get("__v_pad__")
                    )
                elif hierarchical_visual_memories is not None:
                    if step_index >= len(hierarchical_visual_memories):
                        raise ValueError(
                            "the final hierarchical step reads Q*, not an image; "
                            "use a direct tensor override only for an explicit "
                            "integration-step intervention"
                        )
                    replacement_memories = (
                        self._encode_hierarchical_visual_memories(
                            inp["input_ids"],
                            inp["pixel_values"],
                            inp["image_grid_thw"],
                            inp.get("attention_mask"),
                        )
                    )
                    step_overrides[step_index] = replacement_memories[step_index]
                else:
                    step_overrides[step_index] = self._encode_visual_memory(
                        inp["input_ids"], inp["pixel_values"],
                        inp["image_grid_thw"], inp.get("attention_mask"),
                    )
        z_states = self.encode_workspace(
            q_star, v_star, q_pad, v_pad,
            gate_overrides=gate_overrides,
            v_star_reread=v_swap,
            v_reread_padding_mask=v_swap_pad,
            reread_memory_overrides=step_overrides,
            reread_zero_memory=reread_zero_memory,
            final_read_spec=final_read_spec,
            hierarchical_visual_memories=hierarchical_visual_memories,
        )
        del v_star

        if getattr(cfg, "raw_evidence_question", False):
            # Residual: states=[R0=Q*, R1..RT], decode RT or mean(R1..RT).
            # Replacement: states=[Q*, H1..HT], decode Q*+HT.
            # Anchored replacement: states=[Q*, Z1..ZT], decode ZT.
            # The selected mean aggregation is parameter-free; no learned
            # temporal interface is allocated.
            z_final = (
                self.decoder_latent_from_states(z_states)
                if workspace_override is None
                else workspace_override
            )
            if (
                self.training
                and getattr(cfg, "functional_state_loss", False)
                and workspace_override is None
            ):
                target = self._functional_q_target
                target_pad = self._functional_q_target_pad
                target_ids = self._functional_q_target_ids
                self._functional_q_target = None
                self._functional_q_target_pad = None
                self._functional_q_target_ids = None
                if target is None or target_pad is None or target_ids is None:
                    raise RuntimeError("functional-state target was not captured")
                if target.shape != q_star.shape:
                    raise ValueError(
                        "multimodal/text-only question rows disagree after the "
                        f"vision span is removed: {tuple(target.shape)} vs "
                        f"{tuple(q_star.shape)}"
                    )
                valid = ~target_pad
                if q_pad is not None:
                    if q_pad.shape != target_pad.shape or not torch.equal(q_pad, target_pad):
                        raise ValueError(
                            "multimodal and text-only question padding masks disagree"
                        )
                    valid &= ~q_pad
                if not torch.equal(target_ids[valid], question_ids[valid]):
                    raise ValueError(
                        "text-only token ids do not equal the multimodal prompt "
                        "with its vision span removed"
                    )
                # Scale-free visual-state restoration:
                #   ||R_agg-Q*_mm||^2 / sg(||Q*_mm-Q*_text||^2).
                # The denominator removes the large shared text scaffold and
                # defines 1.0 as the image-independent Q* floor.  Per-example
                # reduction prevents long prompts from dominating.
                mask = valid.unsqueeze(-1).to(torch.float32)
                err = (z_final.float() - target.float()).square() * mask
                visual = (target.float() - q_star.float()).square() * mask
                err_sum = err.sum(dim=(1, 2))
                visual_sum = visual.sum(dim=(1, 2)).detach()
                # The normal multimodal branch always has a nonzero image
                # residual, but retain a finite floor for degenerate fixtures.
                floor = mask.sum(dim=(1, 2)).clamp_min(1.0) * 1e-8
                self._functional_state_loss = (
                    err_sum / visual_sum.clamp_min(floor)
                ).mean()
            k_ov = int(getattr(cfg, "er_step_override", 0) or 0)
            if k_ov and workspace_override is None:
                if not 1 <= k_ov <= cfg.num_workspace_steps:
                    raise ValueError(
                        f"er_step_override must be in [1,{cfg.num_workspace_steps}]"
                    )
                z_final = (
                    z_states[0] + z_states[k_ov]
                    if (
                        getattr(cfg, "raw_state_replacement", False)
                        and not getattr(cfg, "raw_question_anchor", False)
                    )
                    else z_states[k_ov]
                )
            lower = ctx.hidden_states  # [B,N_q,d], lower-cache values stay text-only
            if z_final.shape != lower.shape:
                raise ValueError(
                    "raw-evidence decoder latent must match Q* shape, got "
                    f"{tuple(z_final.shape)} vs {tuple(lower.shape)}"
                )
            upper = run_layer_range(
                tm,
                ctx,
                cfg.ell_star + 1,
                None,
                hidden_states=z_final.to(lower.dtype),
                use_cache=True,
            )
            if return_decoder_latent:
                return upper, cache, z_states, v_pooled, z_final
            return upper, cache, z_states, v_pooled

        # 4. overwrite the slot positions with P_Z(Z^(K)). Everything below l*
        #    was computed on a sequence with no image tokens, so this is the
        #    first and only point at which visual information enters.
        lower = ctx.hidden_states  # [B, N_q + S, d]
        z_final = z_states[-1] if workspace_override is None else workspace_override
        if (getattr(self.config, "evidence_reasoning", False)
                and workspace_override is None):
            # strict ER: decoder 입력 = answer 체인의 R^(T) 단독 —
            # E 직행 우회 없음. consume-once로 stale-state 차단.
            r_fin = self._er_answer_final
            assert r_fin is not None, "ER: encode 없이 prefill 시도 (stale/미설정)"
            self._er_answer_final = None
            z_final = r_fin
            # 진단 전용 (checkpoint eval): R^(k)로 디코드하는 step-sweep.
            # ER states = [E^0, R_t^1..R_t^T]이고 dropout=0에서 R_t == R 비트동일.
            # 학습은 이 config 속성이 없으므로 경로 불변 (advisor: 훈련 무접촉).
            k_ov = int(getattr(self.config, "er_step_override", 0) or 0)
            if k_ov:
                T = self.config.num_workspace_steps
                assert 1 <= k_ov <= T, f"er_step_override must be in [1,{T}]"
                z_final = z_states[k_ov]
        z_proj = self.workspace_out(z_final).to(lower.dtype)  # [B, S, d]
        if self.splice_gain is not None:
            # RMS-match the slot rows to the (non-pad) question rows. Fresh
            # linear output is O(1) while layer-20 hiddens run into the
            # hundreds; unmatched, the decoder's attention treats the slots as
            # near-zero vectors and ignoring them is free (probe 20143).
            if question_attention_mask is not None:
                qm = (question_attention_mask > 0).to(lower.dtype)  # [B, N_q]
            else:
                qm = torch.ones(
                    lower.shape[0], n_q, dtype=lower.dtype, device=lower.device
                )
            q_rms = (
                (lower[:, :n_q].float().pow(2).mean(-1).sqrt() * qm).sum(1)
                / qm.sum(1).clamp_min(1.0)
            )  # [B]
            z_rms = z_proj.float().pow(2).mean(-1).sqrt().mean(1)  # [B]
            scale = (q_rms / z_rms.clamp_min(1e-6)).view(-1, 1, 1).to(z_proj.dtype)
            z_proj = z_proj * scale * self.splice_gain.to(z_proj.dtype)
        spliced = torch.cat([lower[:, :n_q], z_proj], dim=1)  # [B, N_q + S, d]

        if (self.training and getattr(cfg, "interface_loss", False)
                and workspace_override is None
                and getattr(self, "_iface_h_base", None) is not None):
            # iface student: R^(T).detach() → 동일 splice 변환(bridge 재적용)
            # → layer l*+1 한 층만 재계산 → 마지막 slot 위치 hidden.
            # gradient는 workspace_out·splice_gain·해당 층 adapter에만 흐른다
            # (R/f_theta/reader/E 차단 — advisor bridge-only 규정).
            zd = self.workspace_out(z_final.detach()).to(lower.dtype)
            if self.splice_gain is not None:
                zd_rms = zd.float().pow(2).mean(-1).sqrt().mean(1)  # [B]
                sc = (q_rms / zd_rms.clamp_min(1e-6)).view(-1, 1, 1).to(zd.dtype)
                zd = zd * sc * self.splice_gain.to(zd.dtype)
            spliced_d = torch.cat([lower[:, :n_q].detach(), zd], dim=1)
            # cache 없는 ctx 사본 필수: 원본 ctx로 돌리면 이 패스가 l*+1층
            # cache에 slot 키를 저장해 본 upper 실행의 키가 2배로 중복된다
            # (게이트에서 [*, 111, 222] shape 오류로 검출).
            import dataclasses as _dc
            ctx_if = _dc.replace(ctx, past_key_values=None)
            h_if = run_layer_range(
                tm, ctx_if, cfg.ell_star + 1, cfg.ell_star + 2,
                hidden_states=spliced_d)  # [B, N_q+S, d]
            h_r = h_if[:, -1].float()  # 마지막 prefix 위치 (= 마지막 slot)
            hb = self._iface_h_base.float()
            self._iface_h_base = None
            d_ = h_r.shape[-1]
            self._iface_loss = (1 - nn.functional.cosine_similarity(
                nn.functional.layer_norm(h_r, (d_,)),
                nn.functional.layer_norm(hb, (d_,)), dim=-1)).mean()

        # 5. F_{>l*} over M* = [Q*; P_Z(Z^(K))], completing the cache.
        upper = run_layer_range(
            tm, ctx, cfg.ell_star + 1, None, hidden_states=spliced, use_cache=True
        )  # [B, N_q + S, d]
        if return_decoder_latent:
            return upper, cache, z_states, v_pooled, z_final
        return upper, cache, z_states, v_pooled

    @torch.no_grad()
    def _greedy_decode_from_prefill(
        self,
        upper: torch.Tensor,
        cache: Cache,
        question_ids: torch.LongTensor,
        question_attention_mask: torch.Tensor | None,
        *,
        max_new_tokens: int,
        eos_token_id: int | None,
    ) -> torch.LongTensor:
        """Greedily decode from an already-built replacement prefix cache."""
        if max_new_tokens < 1:
            raise ValueError("max_new_tokens must be >= 1")
        tm = self.text_model
        b = question_ids.shape[0]
        q_state_interface = self._uses_question_state_interface()
        n_ctx = question_ids.shape[1] + (
            0 if q_state_interface else self.config.num_workspace_slots
        )
        device = question_ids.device
        q_mask = question_attention_mask
        if q_mask is None:
            q_mask = torch.ones_like(question_ids, dtype=torch.long)
        if q_state_interface:
            prefix_mask = q_mask
            logical_prefix_lengths = (q_mask > 0).sum(dim=-1).long()
            last_q = (logical_prefix_lengths - 1).clamp_min(0)
            seed_hidden = upper[
                torch.arange(b, device=device), last_q
            ].unsqueeze(1)
        else:
            prefix_mask, _, logical_prefix_lengths = _replacement_prefix_layout(
                q_mask, self.config.num_workspace_slots
            )
            seed_hidden = upper[:, -1:]

        nxt = self.backbone.lm_head(final_norm(tm, seed_hidden)).argmax(-1)  # [B, 1]
        out = [nxt]
        done = torch.zeros(b, dtype=torch.bool, device=device)
        if eos_token_id is not None:
            done |= nxt.squeeze(1) == eos_token_id

        for step in range(1, max_new_tokens):
            if bool(done.all()):
                break
            pos = n_ctx + step - 1
            embeds = self.vl.get_input_embeddings()(nxt)  # [B, 1, d]
            cache_position = torch.tensor([pos], device=device)
            logical_pos = logical_prefix_lengths + step - 1  # [B]
            position_ids = logical_pos.view(1, b, 1).expand(3, -1, -1)
            generated_mask = torch.ones(
                b, step, dtype=prefix_mask.dtype, device=device
            )
            upper_mask = torch.cat([prefix_mask, generated_mask], dim=-1)
            if q_state_interface:
                lower_mask = upper_mask
            else:
                lower_prefix_mask = prefix_mask.clone()
                lower_prefix_mask[:, -self.config.num_workspace_slots:] = 0
                lower_mask = torch.cat([lower_prefix_mask, generated_mask], dim=-1)
            h = self._decode_token_block(
                embeds,
                position_ids,
                lower_mask,
                upper_mask,
                cache,
                cache_position,
            )
            nxt = self.backbone.lm_head(final_norm(tm, h)).argmax(-1)  # [B, 1]
            if eos_token_id is not None:
                nxt = torch.where(done.unsqueeze(1), nxt.new_full(nxt.shape, eos_token_id), nxt)
                done |= nxt.squeeze(1) == eos_token_id
            out.append(nxt)

        return torch.cat(out, dim=1)

    @torch.no_grad()
    def _gold_answer_logprob_from_prefill(
        self,
        upper: torch.Tensor,
        cache: Cache,
        question_ids: torch.LongTensor,
        question_attention_mask: torch.Tensor | None,
        answer_ids: torch.LongTensor,
        labels: torch.LongTensor,
    ) -> torch.FloatTensor:
        """Score a gold answer against an already sampled latent prefix.

        The score is the mean teacher-forced log probability over valid answer
        tokens.  Running this helper under ``no_grad`` is intentional: the
        decoder acts only as a frozen reward environment, while policy gradient
        reaches the latent adapter solely through ``LatentPolicyTrace``.
        """
        if answer_ids.ndim != 2 or labels.shape != answer_ids.shape:
            raise ValueError(
                "answer_ids and labels must have the same [B, N_a] shape"
            )
        if answer_ids.shape[0] != question_ids.shape[0]:
            raise ValueError("answer and question batch sizes disagree")
        if answer_ids.shape[1] < 1:
            raise ValueError("gold-answer scoring requires at least one token")

        tm = self.text_model
        batch_size, answer_length = answer_ids.shape
        q_state_interface = self._uses_question_state_interface()
        q_mask = question_attention_mask
        if q_mask is None:
            q_mask = torch.ones_like(question_ids, dtype=torch.long)

        if q_state_interface:
            prefix_mask = q_mask
            logical_prefix_lengths = (q_mask > 0).sum(dim=-1).long()
            last_q = (logical_prefix_lengths - 1).clamp_min(0)
            hidden_for_logits = upper[
                torch.arange(batch_size, device=upper.device), last_q
            ].unsqueeze(1)
        else:
            prefix_mask, _, logical_prefix_lengths = _replacement_prefix_layout(
                q_mask, self.config.num_workspace_slots
            )
            hidden_for_logits = upper[:, -1:]

        # The final prefix state predicts y_0.  Consuming gold y_0..y_{N-2}
        # supplies the remaining N-1 teacher-forced prediction states; there is
        # no reason to decode y_{N-1}, whose output would predict beyond labels.
        if answer_length > 1:
            teacher_ids = answer_ids[:, :-1]
            teacher_mask = labels[:, :-1].ne(-100).to(prefix_mask.dtype)
            n_ctx = question_ids.shape[1] + (
                0 if q_state_interface else self.config.num_workspace_slots
            )
            teacher_embeds = self.vl.get_input_embeddings()(teacher_ids)
            cache_position = torch.arange(
                n_ctx,
                n_ctx + teacher_ids.shape[1],
                device=teacher_embeds.device,
            )
            logical_answer_pos = (
                logical_prefix_lengths[:, None]
                + torch.arange(
                    teacher_ids.shape[1], device=teacher_embeds.device
                )[None, :]
            )
            answer_positions = logical_answer_pos.unsqueeze(0).expand(
                3, -1, -1
            )
            upper_mask = torch.cat([prefix_mask, teacher_mask], dim=-1)
            if q_state_interface:
                lower_mask = upper_mask
            else:
                lower_prefix_mask = prefix_mask.clone()
                lower_prefix_mask[:, -self.config.num_workspace_slots:] = 0
                lower_mask = torch.cat(
                    [lower_prefix_mask, teacher_mask], dim=-1
                )
            answer_hidden = self._decode_token_block(
                teacher_embeds,
                answer_positions,
                lower_mask,
                upper_mask,
                cache,
                cache_position,
            )
            hidden_for_logits = torch.cat(
                [hidden_for_logits, answer_hidden], dim=1
            )

        logits = self.backbone.lm_head(final_norm(tm, hidden_for_logits))
        nll, valid = _answer_nll_per_example(logits, labels)
        if not bool(valid.all()):
            raise ValueError("every latent-policy reward row needs a gold token")
        return -nll

    @torch.no_grad()
    def _gold_answer_logprobs_by_crr_prefix(
        self,
        upper: torch.Tensor,
        cache: Cache,
        z_states: list[torch.Tensor] | tuple[torch.Tensor, ...],
        question_ids: torch.LongTensor,
        question_attention_mask: torch.Tensor | None,
        answer_ids: torch.LongTensor,
        labels: torch.LongTensor,
    ) -> torch.FloatTensor:
        """Score ``Z_1..Z_T`` for transition-aware latent RL.

        The final prefix reuses the factual cached decode.  Earlier prefixes
        are stacked into one frozen L22--L27 replay and share the same
        text-only lower answer states, exactly preserving the strict CRR
        interface.  No reward tensor returned here carries a decoder gradient.
        """
        if not getattr(
            self.config, "counterfactual_residual_recurrence", False
        ):
            raise ValueError("prefix reward scoring requires CRR")
        expected_states = int(self.config.num_workspace_steps) + 1
        if len(z_states) != expected_states:
            raise ValueError(
                f"expected B plus {expected_states - 1} CRR states, "
                f"got {len(z_states)}"
            )
        if answer_ids.ndim != 2 or labels.shape != answer_ids.shape:
            raise ValueError(
                "answer_ids and labels must have the same [B, N_a] shape"
            )
        batch_size, answer_length = answer_ids.shape
        if batch_size != question_ids.shape[0] or answer_length < 1:
            raise ValueError("prefix rewards require a non-empty aligned batch")

        tm = self.text_model
        num_steps = expected_states - 1
        n_q = question_ids.shape[1]
        q_mask = question_attention_mask
        if q_mask is None:
            q_mask = torch.ones_like(question_ids, dtype=torch.long)
        q_keep = q_mask.gt(0)
        logical_prefix_lengths = q_keep.sum(dim=-1).long()
        last_q = (logical_prefix_lengths - 1).clamp_min(0)

        prefixes: list[torch.Tensor] = []
        residual_prefix: list[torch.Tensor] = []
        for residual in z_states[1:]:
            residual_prefix.append(residual)
            if getattr(
                self.config, "visual_cumulative_recurrence", False
            ):
                # Full-state CVRR decodes the current recurrent state, not the
                # historical mean of all residual prefixes.
                prefixes.append(z_states[0] + residual)
            else:
                prefixes.append(
                    z_states[0]
                    + self._aggregate_crr_residuals(residual_prefix)
                )

        final_hidden_for_logits = upper[
            torch.arange(batch_size, device=upper.device), last_q
        ].unsqueeze(1)
        answer_lower_hidden = None
        answer_positions = None
        combined_mask = None
        if answer_length > 1:
            teacher_ids = answer_ids[:, :-1]
            teacher_mask = labels[:, :-1].ne(-100).to(q_mask.dtype)
            teacher_embeds = self.vl.get_input_embeddings()(teacher_ids)
            cache_position = torch.arange(
                n_q,
                n_q + teacher_ids.shape[1],
                device=teacher_embeds.device,
            )
            logical_answer_pos = (
                logical_prefix_lengths[:, None]
                + torch.arange(
                    teacher_ids.shape[1], device=teacher_embeds.device
                )[None, :]
            )
            answer_positions = logical_answer_pos.unsqueeze(0).expand(
                3, -1, -1
            )
            combined_mask = torch.cat([q_mask, teacher_mask], dim=-1)
            final_answer_hidden, answer_lower_hidden = self._decode_token_block(
                teacher_embeds,
                answer_positions,
                combined_mask,
                combined_mask,
                cache,
                cache_position,
                return_lower=True,
            )
            final_hidden_for_logits = torch.cat(
                [final_hidden_for_logits, final_answer_hidden], dim=1
            )

        final_logits = self.backbone.lm_head(
            final_norm(tm, final_hidden_for_logits)
        )
        final_nll, final_valid = _answer_nll_per_example(final_logits, labels)
        if not bool(final_valid.all()):
            raise ValueError("every latent-policy reward row needs a gold token")
        if num_steps == 1:
            return -final_nll.unsqueeze(1)

        num_partial = num_steps - 1
        partial_prefix = torch.cat(prefixes[:-1], dim=0)
        q_pos_1d = q_keep.long().cumsum(dim=-1) - 1
        q_pos_1d = q_pos_1d.masked_fill(~q_keep, 0)
        q_positions = q_pos_1d.unsqueeze(0).expand(3, -1, -1)
        partial_q_positions = q_positions.repeat(1, num_partial, 1)
        partial_q_mask = q_mask.repeat(num_partial, 1)

        if answer_lower_hidden is not None:
            partial_input = torch.cat(
                [
                    partial_prefix.to(answer_lower_hidden.dtype),
                    answer_lower_hidden.repeat(num_partial, 1, 1),
                ],
                dim=1,
            )
            partial_positions = torch.cat(
                [
                    partial_q_positions,
                    answer_positions.repeat(1, num_partial, 1),
                ],
                dim=-1,
            )
            partial_attention_mask = combined_mask.repeat(num_partial, 1)
        else:
            partial_input = partial_prefix
            partial_positions = partial_q_positions
            partial_attention_mask = partial_q_mask

        partial_ctx = make_split_context(
            tm,
            partial_input,
            partial_positions,
            partial_attention_mask,
        )
        partial_upper = run_layer_range(
            tm,
            partial_ctx,
            self._upper_decoder_start(),
            None,
            hidden_states=partial_input,
            use_cache=False,
        )
        partial_last_q = last_q.repeat(num_partial)
        partial_hidden_for_logits = partial_upper[
            torch.arange(partial_upper.shape[0], device=partial_upper.device),
            partial_last_q,
        ].unsqueeze(1)
        if answer_lower_hidden is not None:
            partial_hidden_for_logits = torch.cat(
                [
                    partial_hidden_for_logits,
                    partial_upper[
                        :, n_q : n_q + answer_lower_hidden.shape[1]
                    ],
                ],
                dim=1,
            )
        partial_logits = self.backbone.lm_head(
            final_norm(tm, partial_hidden_for_logits)
        )
        partial_nll, partial_valid = _answer_nll_per_example(
            partial_logits,
            labels.repeat(num_partial, 1),
        )
        if not bool(partial_valid.all()):
            raise ValueError("every CRR prefix needs a valid gold-answer score")
        partial_scores = -partial_nll.view(num_partial, batch_size).transpose(
            0, 1
        )
        return torch.cat([partial_scores, -final_nll.unsqueeze(1)], dim=1)

    @torch.no_grad()
    def generate_answer(
        self,
        input_ids: torch.LongTensor,
        pixel_values: torch.Tensor,
        image_grid_thw: torch.LongTensor,
        question_ids: torch.LongTensor,
        attention_mask: torch.Tensor | None = None,
        question_attention_mask: torch.Tensor | None = None,
        max_new_tokens: int = 8,
        eos_token_id: int | None = None,
        workspace_override: torch.Tensor | None = None,
        gate_overrides: list | None = None,
        reread_inputs: dict | None = None,
        reread_zero_memory: bool = False,
    ) -> torch.LongTensor:
        """Greedy decode against the replacement cache. ``[B, <=N_a]``.

        A purpose-built loop rather than ``generate()``: the prefill is not a
        token sequence the generation utilities can reconstruct, and the targets
        here are short VQA answers. Correctness of the read-once contract is
        easier to guarantee with the loop visible.
        """
        if getattr(self.config, "perceive_deliberate_chain", False):
            if (
                workspace_override is not None
                or gate_overrides is not None
                or reread_inputs is not None
                or reread_zero_memory
            ):
                raise ValueError(
                    "historical workspace/reread interventions do not apply "
                    "to perceive_deliberate_chain"
                )
            return self._generate_perceive_deliberate_chain(
                input_ids,
                pixel_values,
                image_grid_thw,
                question_ids,
                attention_mask,
                question_attention_mask,
                max_new_tokens=max_new_tokens,
                eos_token_id=eos_token_id,
            )
        if getattr(self.config, "spatial_visual_recurrence", False):
            if (
                workspace_override is not None
                or gate_overrides is not None
                or reread_inputs is not None
                or reread_zero_memory
            ):
                raise ValueError(
                    "historical workspace/reread interventions do not apply "
                    "to spatial_visual_recurrence"
                )
            return self._generate_spatial_visual_field(
                input_ids,
                pixel_values,
                image_grid_thw,
                question_ids,
                attention_mask,
                question_attention_mask,
                max_new_tokens=max_new_tokens,
                eos_token_id=eos_token_id,
            )
        upper, cache, _, _ = self.prefill(
            input_ids,
            pixel_values,
            image_grid_thw,
            question_ids,
            attention_mask,
            question_attention_mask,
            workspace_override=workspace_override,
            gate_overrides=gate_overrides,
            reread_inputs=reread_inputs,
            reread_zero_memory=reread_zero_memory,
        )
        return self._greedy_decode_from_prefill(
            upper,
            cache,
            question_ids,
            question_attention_mask,
            max_new_tokens=max_new_tokens,
            eos_token_id=eos_token_id,
        )

    def generate_answer_with_latent_policy(
        self,
        input_ids: torch.LongTensor,
        pixel_values: torch.Tensor,
        image_grid_thw: torch.LongTensor,
        question_ids: torch.LongTensor,
        attention_mask: torch.Tensor | None = None,
        question_attention_mask: torch.Tensor | None = None,
        *,
        policy_std: float,
        max_new_tokens: int = 8,
        eos_token_id: int | None = None,
        sample: bool = True,
        sample_mask: torch.Tensor | None = None,
    ) -> tuple[torch.LongTensor, LatentPolicyTrace]:
        """Sample CRR latent actions, then greedily decode a terminal answer.

        The returned token ids are an environment outcome and carry no graph.
        Only ``trace.log_probs`` is differentiable.  Deterministic inference
        continues to use :meth:`generate_answer` and therefore the policy mean.
        """
        if not getattr(
            self.config, "counterfactual_residual_recurrence", False
        ):
            raise ValueError("latent-policy generation requires a CRR checkpoint")
        upper, cache, _, _, trace = self.prefill(
            input_ids,
            pixel_values,
            image_grid_thw,
            question_ids,
            attention_mask,
            question_attention_mask,
            latent_policy_std=float(policy_std),
            latent_policy_sample=bool(sample),
            latent_policy_sample_mask=sample_mask,
            return_latent_policy_trace=True,
        )
        generated = self._greedy_decode_from_prefill(
            upper,
            cache,
            question_ids,
            question_attention_mask,
            max_new_tokens=max_new_tokens,
            eos_token_id=eos_token_id,
        )
        return generated, trace

    def score_gold_answer_with_latent_policy(
        self,
        input_ids: torch.LongTensor,
        pixel_values: torch.Tensor,
        image_grid_thw: torch.LongTensor,
        question_ids: torch.LongTensor,
        answer_ids: torch.LongTensor,
        labels: torch.LongTensor,
        attention_mask: torch.Tensor | None = None,
        question_attention_mask: torch.Tensor | None = None,
        *,
        policy_std: float,
        sample: bool = True,
        sample_mask: torch.Tensor | None = None,
        score_prefixes: bool = False,
    ) -> tuple[
        torch.FloatTensor,
        torch.FloatTensor | None,
        LatentPolicyTrace,
    ]:
        """Sample latent actions and return frozen gold-answer scores.

        ``score_prefixes=False`` preserves the original terminal-only reward
        path.  The opt-in prefix path scores ``Z_1..Z_T`` with the same frozen
        decoder and gold tokens, while keeping the Gaussian score trace as the
        only differentiable path.
        """
        if not getattr(
            self.config, "counterfactual_residual_recurrence", False
        ):
            raise ValueError("latent-policy scoring requires a CRR checkpoint")
        upper, cache, z_states, _, trace = self.prefill(
            input_ids,
            pixel_values,
            image_grid_thw,
            question_ids,
            attention_mask,
            question_attention_mask,
            latent_policy_std=float(policy_std),
            latent_policy_sample=bool(sample),
            latent_policy_sample_mask=sample_mask,
            return_latent_policy_trace=True,
        )
        if score_prefixes:
            step_scores = self._gold_answer_logprobs_by_crr_prefix(
                upper,
                cache,
                z_states,
                question_ids,
                question_attention_mask,
                answer_ids,
                labels,
            )
            score = step_scores[:, -1]
        else:
            score = self._gold_answer_logprob_from_prefill(
                upper,
                cache,
                question_ids,
                question_attention_mask,
                answer_ids,
                labels,
            )
            step_scores = None
        return score, step_scores, trace

    # -- forward -----------------------------------------------------------

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        pixel_values: torch.Tensor | None = None,
        image_grid_thw: torch.LongTensor | None = None,
        question_ids: torch.LongTensor | None = None,
        question_attention_mask: torch.Tensor | None = None,
        labels: torch.LongTensor | None = None,
        answer_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        use_cache: bool | None = None,
        return_dict: bool | None = None,
        workspace_override: torch.Tensor | None = None,
        latent_policy_std: float | None = None,
        latent_policy_sample: bool = True,
        latent_policy_sample_mask: torch.Tensor | None = None,
        latent_policy_max_new_tokens: int = 8,
        latent_policy_eos_token_id: int | None = None,
        latent_policy_score_gold: bool = False,
        latent_policy_score_prefixes: bool = False,
        reliance_transition_index: int | torch.Tensor | None = None,
        curriculum_readout_depth: int | torch.Tensor | None = None,
        **kwargs,
    ) -> CloseCausalLMOutput:
        """Prefill through the workspace, then decode the answer.

        Args:
            input_ids: the *multimodal* sequence (image placeholders included),
                used only to build ``V*``.
            question_ids: the *text-only* question, used for ``Q*`` and for the
                replacement cache. Required -- there is deliberately no fallback
                that strips the vision span implicitly, because getting that
                wrong would silently reintroduce image tokens below ``l*``.
            answer_ids: teacher-forced answer tokens ``[B, N_a]``.
            labels: ``[B, N_a]`` targets **already aligned** to ``answer_ids`` --
                ``labels[:, i]`` is the target at answer position ``i``, and no
                internal shift is applied. The shift lives on the hidden-state
                side: position ``i``'s prediction comes from
                ``[last replacement-cache position] + answer[:i]``. ``-100`` ignores.
            past_key_values: not accepted yet -- see below.
            return_dict: must be ``True``/``None``; tuple returns are not supported.
        """
        cfg = self.config
        if kwargs:
            raise TypeError(f"unexpected forward kwargs: {sorted(kwargs)}")
        if getattr(cfg, "perceive_deliberate_chain", False):
            if past_key_values is not None:
                raise NotImplementedError(
                    "perceive/deliberate forward constructs one fresh graph"
                )
            if return_dict is False:
                raise ValueError(
                    "perceive/deliberate chain returns CloseCausalLMOutput only"
                )
            if workspace_override is not None:
                raise ValueError(
                    "workspace_override does not apply to latent chain tokens"
                )
            if latent_policy_std is not None or latent_policy_sample_mask is not None:
                raise ValueError(
                    "the historical CRR latent policy is incompatible with "
                    "the perceive/deliberate chain"
                )
            if reliance_transition_index is not None:
                raise ValueError(
                    "CRR transition interventions do not apply to the chain"
                )
            if curriculum_readout_depth is not None:
                raise ValueError(
                    "program-curriculum readout does not apply to the chain"
                )
            if input_ids is None or pixel_values is None or image_grid_thw is None:
                raise ValueError(
                    "perceive/deliberate chain requires multimodal inputs"
                )
            if question_ids is None:
                raise ValueError(
                    "perceive/deliberate chain requires clean question_ids"
                )
            if labels is None or answer_ids is None:
                raise ValueError(
                    "teacher-forced chain forward requires labels and answer_ids"
                )
            return self._forward_perceive_deliberate_chain(
                input_ids,
                attention_mask,
                pixel_values,
                image_grid_thw,
                question_ids,
                question_attention_mask,
                labels,
                answer_ids,
            )
        if getattr(cfg, "spatial_visual_recurrence", False):
            if past_key_values is not None:
                raise NotImplementedError(
                    "spatial forward always constructs a fresh final-field prefix"
                )
            if return_dict is False:
                raise ValueError(
                    "spatial visual recurrence returns CloseCausalLMOutput only"
                )
            if workspace_override is not None:
                raise ValueError(
                    "workspace_override does not apply to the visual field"
                )
            if latent_policy_std is not None or latent_policy_sample_mask is not None:
                raise ValueError(
                    "the historical CRR latent policy is incompatible with "
                    "spatial visual recurrence"
                )
            if reliance_transition_index is not None:
                raise ValueError(
                    "transition reliance interventions are not part of SRVF"
                )
            if curriculum_readout_depth is not None:
                raise ValueError(
                    "program-curriculum readout does not apply to SRVF"
                )
            if input_ids is None or pixel_values is None or image_grid_thw is None:
                raise ValueError(
                    "spatial recurrence requires input_ids, pixel_values, and "
                    "image_grid_thw"
                )
            if question_ids is None:
                raise ValueError(
                    "spatial recurrence requires text-only question_ids for "
                    "visual-to-question cross-attention"
                )
            return self._forward_spatial_visual_field(
                input_ids,
                attention_mask,
                pixel_values,
                image_grid_thw,
                question_ids,
                question_attention_mask,
                labels,
                answer_ids,
                use_cache=use_cache,
            )
        if latent_policy_std is None and latent_policy_sample_mask is not None:
            raise ValueError(
                "latent_policy_sample_mask requires latent_policy_std"
            )
        if latent_policy_std is not None:
            if workspace_override is not None:
                raise ValueError(
                    "latent policy and workspace_override are mutually exclusive"
                )
            if past_key_values is not None:
                raise ValueError("latent policy always builds a fresh prefix cache")
            if return_dict is False:
                raise ValueError("latent policy requires return_dict=True")
            if input_ids is None or pixel_values is None or image_grid_thw is None:
                raise ValueError(
                    "latent policy requires multimodal input_ids, pixel_values, "
                    "and image_grid_thw"
                )
            if question_ids is None:
                raise ValueError("latent policy requires text-only question_ids")
            generated = None
            gold_logprob = None
            step_gold_logprob = None
            if latent_policy_score_gold:
                if answer_ids is None or labels is None:
                    raise ValueError(
                        "latent_policy_score_gold requires answer_ids and labels"
                    )
                gold_logprob, step_gold_logprob, trace = (
                    self.score_gold_answer_with_latent_policy(
                        input_ids,
                        pixel_values,
                        image_grid_thw,
                        question_ids,
                        answer_ids,
                        labels,
                        attention_mask,
                        question_attention_mask,
                        policy_std=float(latent_policy_std),
                        sample=bool(latent_policy_sample),
                        sample_mask=latent_policy_sample_mask,
                        score_prefixes=bool(latent_policy_score_prefixes),
                    )
                )
            else:
                if latent_policy_score_prefixes:
                    raise ValueError(
                        "latent_policy_score_prefixes requires "
                        "latent_policy_score_gold"
                    )
                generated, trace = self.generate_answer_with_latent_policy(
                    input_ids,
                    pixel_values,
                    image_grid_thw,
                    question_ids,
                    attention_mask,
                    question_attention_mask,
                    policy_std=float(latent_policy_std),
                    max_new_tokens=int(latent_policy_max_new_tokens),
                    eos_token_id=latent_policy_eos_token_id,
                    sample=bool(latent_policy_sample),
                    sample_mask=latent_policy_sample_mask,
                )
            return CloseCausalLMOutput(
                generated_ids=generated,
                latent_policy_gold_logprob=gold_logprob,
                latent_policy_step_gold_logprob=step_gold_logprob,
                latent_policy_log_probs=trace.log_probs,
                latent_policy_mean_update_rms=trace.mean_update_rms,
                latent_policy_noise_rms=trace.noise_rms,
            )
        compute_residual_swap = bool(
            self.training
            and getattr(cfg, "counterfactual_reliance_loss", False)
            and labels is not None
            and workspace_override is None
        )
        compute_step_retention = bool(
            self.training
            and getattr(cfg, "counterfactual_step_retention_loss", False)
            and labels is not None
            and workspace_override is None
        )
        need_answer_lower = compute_residual_swap or compute_step_retention
        # Every prefill rebuilds the replacement cache from scratch. Accepting a
        # caller-supplied cache would silently discard it and redo the visual
        # read, so refuse instead of pretending to support incremental decode.
        if past_key_values is not None:
            raise NotImplementedError(
                "past_key_values is not supported: forward() always rebuilds the "
                "replacement cache. Incremental decoding needs prepare_inputs_for_"
                "generation, which is not implemented yet."
            )
        if return_dict is False:
            raise ValueError(
                "return_dict=False is not supported; CLOSE always returns a "
                "CloseCausalLMOutput (the workspace trajectory has no tuple slot)."
            )
        if question_ids is None:
            raise ValueError(
                "question_ids (text-only, no vision span) is required. Strip the "
                "vision span explicitly with utils.splitting.vision_span_mask; an "
                "implicit fallback risks leaking image tokens below l*."
            )
        if pixel_values is None or input_ids is None:
            raise ValueError("input_ids and pixel_values are required to build V*.")

        tm = self.text_model
        n_q = question_ids.shape[1]
        upper, cache, z_states, v_pooled = self.prefill(
            input_ids,
            pixel_values,
            image_grid_thw,
            question_ids,
            attention_mask,
            question_attention_mask,
            workspace_override=workspace_override,
            curriculum_readout_depth=curriculum_readout_depth,
        )

        q_state_interface = self._uses_question_state_interface()
        # The last real prefix position predicts answer token 0.  Historical
        # slot paths always end in a valid slot; raw Q* batches may end in
        # right-padding and therefore require a per-item gather.
        if q_state_interface:
            q_mask_for_seed = question_attention_mask
            if q_mask_for_seed is None:
                q_mask_for_seed = torch.ones_like(question_ids, dtype=torch.long)
            last_q = ((q_mask_for_seed > 0).sum(dim=-1).long() - 1).clamp_min(0)
            hidden_for_logits = upper[
                torch.arange(upper.shape[0], device=upper.device), last_q
            ].unsqueeze(1)
        else:
            hidden_for_logits = upper[:, -1:]  # [B, 1, d]
        answer_lower_hidden = None
        ans_positions = None
        upper_mask = None
        if answer_ids is not None and answer_ids.shape[1] > 1:
            n_ctx = n_q + (
                0 if q_state_interface else cfg.num_workspace_slots
            )
            n_a = answer_ids.shape[1]
            ans_embeds = self.vl.get_input_embeddings()(answer_ids)  # [B, N_a, d]
            cache_position = torch.arange(n_ctx, n_ctx + n_a, device=ans_embeds.device)
            q_mask = question_attention_mask
            if q_mask is None:
                q_mask = torch.ones_like(question_ids, dtype=torch.long)
            if q_state_interface:
                prefix_mask = q_mask
                logical_prefix_lengths = (q_mask > 0).sum(dim=-1).long()
            else:
                prefix_mask, _, logical_prefix_lengths = _replacement_prefix_layout(
                    q_mask, cfg.num_workspace_slots
                )
            logical_answer_pos = (
                logical_prefix_lengths[:, None]
                + torch.arange(n_a, device=ans_embeds.device)[None, :]
            )  # [B,N_a]
            ans_positions = logical_answer_pos.unsqueeze(0).expand(3, -1, -1)
            if labels is not None and labels.shape == answer_ids.shape:
                answer_mask = labels.ne(-100).to(prefix_mask.dtype)
            else:
                answer_mask = torch.ones(
                    answer_ids.shape, dtype=prefix_mask.dtype, device=answer_ids.device
                )
            upper_mask = torch.cat([prefix_mask, answer_mask], dim=-1)
            if q_state_interface:
                lower_mask = upper_mask
            else:
                lower_prefix_mask = prefix_mask.clone()
                lower_prefix_mask[:, -cfg.num_workspace_slots:] = 0
                lower_mask = torch.cat([lower_prefix_mask, answer_mask], dim=-1)
            decoded_answer = self._decode_token_block(
                ans_embeds,
                ans_positions,
                lower_mask,
                upper_mask,
                cache,
                cache_position,
                return_lower=need_answer_lower,
            )  # [B,N_a,d]
            if need_answer_lower:
                h_ans, answer_lower_hidden = decoded_answer
            else:
                h_ans = decoded_answer
            hidden_for_logits = torch.cat([hidden_for_logits, h_ans[:, :-1]], dim=1)

        logits = self.backbone.lm_head(final_norm(tm, hidden_for_logits))  # [B, N_a, V]

        # Both training-only counterfactual objectives replay only L22-L27.
        # Question/answer states below that boundary are shared with the
        # factual pass, so neither branch performs another visual read.
        q_mask = None
        q_keep = None
        q_positions = None
        if compute_residual_swap or compute_step_retention:
            q_mask = question_attention_mask
            if q_mask is None:
                q_mask = torch.ones_like(question_ids, dtype=torch.long)
            q_keep = q_mask.gt(0)
            q_pos_1d = q_keep.long().cumsum(dim=-1) - 1
            q_pos_1d = q_pos_1d.masked_fill(~q_keep, 0)
            q_positions = q_pos_1d.unsqueeze(0).expand(3, -1, -1)

        step_answer_nll = None
        step_answer_valid = None
        if compute_step_retention:
            if not getattr(cfg, "counterfactual_residual_recurrence", False):
                raise RuntimeError(
                    "step retention is implemented only for CRR"
                )
            if len(z_states) < 3:
                raise RuntimeError(
                    "step retention requires at least two CRR residual steps"
                )
            if q_mask is None or q_keep is None or q_positions is None:
                raise RuntimeError("step retention is missing question layout")

            # The factual path already scores Z_T.  Score Z_1..Z_{T-1} in one
            # stacked upper-decoder replay to keep the extra launches small.
            base_anchor = z_states[0]
            partial_prefixes = []
            residual_prefix: list[torch.Tensor] = []
            for residual in z_states[1:-1]:
                residual_prefix.append(residual)
                partial_prefixes.append(
                    base_anchor
                    + self._aggregate_crr_residuals(residual_prefix)
                )
            num_partial = len(partial_prefixes)
            partial_prefix = torch.cat(partial_prefixes, dim=0)
            partial_q_positions = q_positions.repeat(1, num_partial, 1)
            partial_q_mask = q_mask.repeat(num_partial, 1)

            if answer_ids is not None and answer_ids.shape[1] > 1:
                if (
                    answer_lower_hidden is None
                    or ans_positions is None
                    or upper_mask is None
                ):
                    raise RuntimeError(
                        "step retention is missing teacher-forced lower states"
                    )
                partial_input = torch.cat(
                    [
                        partial_prefix.to(answer_lower_hidden.dtype),
                        answer_lower_hidden.repeat(num_partial, 1, 1),
                    ],
                    dim=1,
                )
                partial_positions = torch.cat(
                    [
                        partial_q_positions,
                        ans_positions.repeat(1, num_partial, 1),
                    ],
                    dim=-1,
                )
                partial_attention_mask = upper_mask.repeat(num_partial, 1)
            else:
                partial_input = partial_prefix
                partial_positions = partial_q_positions
                partial_attention_mask = partial_q_mask

            partial_ctx = make_split_context(
                tm,
                partial_input,
                partial_positions,
                partial_attention_mask,
            )
            partial_upper = run_layer_range(
                tm,
                partial_ctx,
                self._upper_decoder_start(),
                None,
                hidden_states=partial_input,
                use_cache=False,
            )
            last_q = (q_keep.sum(dim=-1).long() - 1).clamp_min(0)
            partial_last_q = last_q.repeat(num_partial)
            partial_hidden_for_logits = partial_upper[
                torch.arange(
                    partial_upper.shape[0], device=partial_upper.device
                ),
                partial_last_q,
            ].unsqueeze(1)
            if answer_ids is not None and answer_ids.shape[1] > 1:
                n_a = answer_ids.shape[1]
                partial_hidden_for_logits = torch.cat(
                    [
                        partial_hidden_for_logits,
                        partial_upper[:, n_q : n_q + n_a - 1],
                    ],
                    dim=1,
                )
            partial_logits = self.backbone.lm_head(
                final_norm(tm, partial_hidden_for_logits)
            )
            partial_nll, partial_valid = _answer_nll_per_example(
                partial_logits, labels.repeat(num_partial, 1)
            )
            factual_nll, factual_valid = _answer_nll_per_example(logits, labels)
            batch_size = labels.shape[0]
            step_answer_nll = torch.cat(
                [
                    partial_nll.view(num_partial, batch_size),
                    factual_nll.unsqueeze(0),
                ],
                dim=0,
            ).transpose(0, 1)
            step_answer_valid = torch.cat(
                [
                    partial_valid.view(num_partial, batch_size),
                    factual_valid.unsqueeze(0),
                ],
                dim=0,
            ).transpose(0, 1)

        residual_swap_score_gap = None
        residual_swap_valid = None
        residual_swap_transition = None
        corrective_current_nll = None
        corrective_previous_nll = None
        corrective_swap_nll = None
        corrective_valid = None
        if compute_residual_swap:
            if not getattr(cfg, "counterfactual_residual_recurrence", False):
                raise RuntimeError(
                    "counterfactual reliance is implemented only for CRR"
                )
            donors, residual_swap_valid = _different_answer_residual_donors(
                labels, question_attention_mask
            )
            base_anchor = z_states[0]

            if q_mask is None or q_keep is None or q_positions is None:
                raise RuntimeError("CRR residual swap is missing question layout")

            corrective_transition = bool(
                getattr(cfg, "corrective_transition_reliance_loss", False)
            )
            if corrective_transition:
                if reliance_transition_index is None:
                    raise ValueError(
                        "corrective transition reliance requires "
                        "reliance_transition_index"
                    )
                if isinstance(reliance_transition_index, torch.Tensor):
                    if reliance_transition_index.numel() != 1:
                        raise ValueError(
                            "reliance_transition_index must be a scalar"
                        )
                    transition_index = int(reliance_transition_index.item())
                else:
                    transition_index = int(reliance_transition_index)
                horizon = int(cfg.num_workspace_steps)
                if not 2 <= transition_index <= horizon:
                    raise ValueError(
                        "corrective reliance transition must be in "
                        f"[2,{horizon}], got {transition_index}"
                    )

                # Score the selected *local* transition through the same
                # decoder used at inference.  Previous history is preserved;
                # only C_k is donor-swapped in the negative prefix.  Unlike
                # the historical tail rollout, this directly asks whether
                # transition k corrects Z_{k-1} before any later step can hide
                # or amplify its effect.
                current_prefix = base_anchor + self._aggregate_crr_residuals(
                    z_states[1 : transition_index + 1]
                )
                previous_prefix = (
                    base_anchor
                    + self._aggregate_crr_residuals(
                        z_states[1:transition_index]
                    )
                ).detach()

                aligned = q_mask.gt(0) & q_mask.index_select(0, donors).gt(0)
                donor_residual = z_states[transition_index].index_select(
                    0, donors
                ).detach()
                donor_residual = donor_residual.masked_fill(
                    ~aligned.unsqueeze(-1), 0
                )
                swapped_steps = [
                    state.detach()
                    for state in z_states[1:transition_index]
                ]
                swapped_steps.append(donor_residual)
                swapped_prefix = (
                    base_anchor.detach()
                    + self._aggregate_crr_residuals(swapped_steps)
                ).detach()

                prefixes = torch.cat(
                    [current_prefix, previous_prefix, swapped_prefix], dim=0
                )
                num_branches = 3
                if answer_ids is not None and answer_ids.shape[1] > 1:
                    if (
                        answer_lower_hidden is None
                        or ans_positions is None
                        or upper_mask is None
                    ):
                        raise RuntimeError(
                            "corrective transition scoring is missing "
                            "teacher-forced lower states"
                        )
                    replay_input = torch.cat(
                        [
                            prefixes.to(answer_lower_hidden.dtype),
                            answer_lower_hidden.repeat(num_branches, 1, 1),
                        ],
                        dim=1,
                    )
                    replay_positions = torch.cat(
                        [
                            q_positions.repeat(1, num_branches, 1),
                            ans_positions.repeat(1, num_branches, 1),
                        ],
                        dim=-1,
                    )
                    replay_attention_mask = upper_mask.repeat(num_branches, 1)
                else:
                    replay_input = prefixes
                    replay_positions = q_positions.repeat(1, num_branches, 1)
                    replay_attention_mask = q_mask.repeat(num_branches, 1)

                replay_ctx = make_split_context(
                    tm,
                    replay_input,
                    replay_positions,
                    replay_attention_mask,
                )
                replay_upper = run_layer_range(
                    tm,
                    replay_ctx,
                    self._upper_decoder_start(),
                    None,
                    hidden_states=replay_input,
                    use_cache=False,
                )
                batch_size = labels.shape[0]
                last_q = (q_keep.sum(dim=-1).long() - 1).clamp_min(0)
                replay_last_q = last_q.repeat(num_branches)
                replay_hidden_for_logits = replay_upper[
                    torch.arange(
                        replay_upper.shape[0], device=replay_upper.device
                    ),
                    replay_last_q,
                ].unsqueeze(1)
                if answer_ids is not None and answer_ids.shape[1] > 1:
                    n_a = answer_ids.shape[1]
                    replay_hidden_for_logits = torch.cat(
                        [
                            replay_hidden_for_logits,
                            replay_upper[:, n_q : n_q + n_a - 1],
                        ],
                        dim=1,
                    )
                replay_logits = self.backbone.lm_head(
                    final_norm(tm, replay_hidden_for_logits)
                )
                replay_nll, replay_valid = _answer_nll_per_example(
                    replay_logits, labels.repeat(num_branches, 1)
                )
                replay_nll = replay_nll.view(num_branches, batch_size)
                replay_valid = replay_valid.view(num_branches, batch_size)
                corrective_current_nll = replay_nll[0]
                corrective_previous_nll = replay_nll[1]
                corrective_swap_nll = replay_nll[2]
                corrective_valid = (
                    residual_swap_valid
                    & replay_valid[0]
                    & replay_valid[1]
                    & replay_valid[2]
                )
                residual_swap_score_gap = (
                    corrective_swap_nll - corrective_current_nll
                )
                residual_swap_valid = corrective_valid
                residual_swap_transition = torch.tensor(
                    transition_index,
                    dtype=torch.long,
                    device=labels.device,
                )
            else:
                if getattr(cfg, "transition_aware_reliance_loss", False):
                    if reliance_transition_index is None:
                        raise ValueError(
                            "transition-aware reliance requires "
                            "reliance_transition_index"
                        )
                    if isinstance(reliance_transition_index, torch.Tensor):
                        if reliance_transition_index.numel() != 1:
                            raise ValueError(
                                "reliance_transition_index must be a scalar"
                            )
                        transition_index = int(reliance_transition_index.item())
                    else:
                        transition_index = int(reliance_transition_index)
                    swapped_prefix = self._counterfactual_transition_rollout_prefix(
                        z_states,
                        donors,
                        transition_index,
                        q_mask,
                    )
                    residual_swap_transition = torch.tensor(
                        transition_index,
                        dtype=torch.long,
                        device=labels.device,
                    )
                else:
                    aggregated_residual = self._aggregate_crr_residuals(
                        z_states[1:]
                    )
                    donor_residual = aggregated_residual.index_select(0, donors)
                    # Never copy right-padding as semantic residual content.
                    aligned = (
                        q_mask.gt(0)
                        & q_mask.index_select(0, donors).gt(0)
                    )
                    donor_residual = donor_residual.masked_fill(
                        ~aligned.unsqueeze(-1), 0
                    )
                    swapped_prefix = base_anchor + donor_residual

                if answer_ids is not None and answer_ids.shape[1] > 1:
                    if (
                        answer_lower_hidden is None
                        or ans_positions is None
                        or upper_mask is None
                    ):
                        raise RuntimeError(
                            "CRR residual swap is missing teacher-forced lower states"
                        )
                    swapped_input = torch.cat(
                        [
                            swapped_prefix.to(answer_lower_hidden.dtype),
                            answer_lower_hidden,
                        ],
                        dim=1,
                    )
                    swapped_positions = torch.cat(
                        [q_positions, ans_positions], dim=-1
                    )
                    swapped_attention_mask = upper_mask
                else:
                    swapped_input = swapped_prefix
                    swapped_positions = q_positions
                    swapped_attention_mask = q_mask

                # Only L22-L27 are replayed.  Vision, the lower backbone, L21,
                # and teacher-forced answer states below L22 are shared with
                # the factual pass; no second multimodal read is constructed.
                swapped_ctx = make_split_context(
                    tm,
                    swapped_input,
                    swapped_positions,
                    swapped_attention_mask,
                )
                swapped_upper = run_layer_range(
                    tm,
                    swapped_ctx,
                    self._upper_decoder_start(),
                    None,
                    hidden_states=swapped_input,
                    use_cache=False,
                )
                last_q = (q_keep.sum(dim=-1).long() - 1).clamp_min(0)
                swapped_hidden_for_logits = swapped_upper[
                    torch.arange(
                        swapped_upper.shape[0], device=swapped_upper.device
                    ),
                    last_q,
                ].unsqueeze(1)
                if answer_ids is not None and answer_ids.shape[1] > 1:
                    n_a = answer_ids.shape[1]
                    swapped_hidden_for_logits = torch.cat(
                        [
                            swapped_hidden_for_logits,
                            swapped_upper[:, n_q : n_q + n_a - 1],
                        ],
                        dim=1,
                    )
                swapped_logits = self.backbone.lm_head(
                    final_norm(tm, swapped_hidden_for_logits)
                )
                factual_nll, factual_valid = _answer_nll_per_example(
                    logits, labels
                )
                swapped_nll, swapped_valid = _answer_nll_per_example(
                    swapped_logits, labels
                )
                # score = -NLL, hence factual - swap = NLL_swap - NLL_pos.
                residual_swap_score_gap = swapped_nll - factual_nll
                residual_swap_valid &= factual_valid & swapped_valid

        loss = None
        if labels is not None:
            loss = answer_cross_entropy(
                logits,
                labels,
                per_example=bool(
                    getattr(cfg, "per_example_answer_loss", False)
                ),
            )

        iface = getattr(self, "_iface_loss", None)
        self._iface_loss = None  # consume-once
        functional = getattr(self, "_functional_state_loss", None)
        self._functional_state_loss = None  # consume-once
        return CloseCausalLMOutput(
            loss=loss,
            logits=logits,
            past_key_values=cache if use_cache else None,
            workspace_states=tuple(z_states),
            v_star_pooled=v_pooled,
            iface_loss=iface,
            functional_state_loss=functional,
            residual_swap_score_gap=residual_swap_score_gap,
            residual_swap_valid=residual_swap_valid,
            residual_swap_transition=residual_swap_transition,
            corrective_current_nll=corrective_current_nll,
            corrective_previous_nll=corrective_previous_nll,
            corrective_swap_nll=corrective_swap_nll,
            corrective_valid=corrective_valid,
            step_answer_nll=step_answer_nll,
            step_answer_valid=step_answer_valid,
        )

__all__ = [
    "CloseQwen2_5_VLForConditionalGeneration",
    "WorkspaceRead",
    "WorkspaceTransition",
    "RawEvidenceQuestionTransition",
    "answer_cross_entropy",
    "build_adapter_config",
    "CloseCausalLMOutput",
    "LatentPolicyTrace",
    "CRRLatentPolicyAdapter",
]