"""Raw-evidence/question recurrent reasoning on Qwen3-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 32-query/8-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 (4096).
"""

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.modeling_outputs import ModelOutput
from transformers.models.qwen3_vl.modeling_qwen3_vl import (
    Qwen3VLForConditionalGeneration,
    Qwen3VLPreTrainedModel,
    Qwen3VLTextRMSNorm,
)

from .configuration_source_qwen3 import CloseQwen3VLConfig
from .source_splitting import (
    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:
    """Inverse of ``select_tokens_padded`` on each example's valid rows."""
    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")
    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 _linear_qwen3_vision_patch_embed_forward(
    patch_embed,
    hidden_states: torch.Tensor,
) -> torch.Tensor:
    """Evaluate Qwen3-VL's non-overlapping patch Conv3d as one GEMM.

    The official projection has kernel size equal to stride and receives one
    already-flattened patch per row.  It is therefore exactly a linear map over
    each row.  Avoiding cuDNN here is important for dynamic-resolution batches:
    cuDNN otherwise spends tens of seconds planning a new Conv3d shape whenever
    the total number of patches changes.
    """
    projection = patch_embed.proj
    target_dtype = projection.weight.dtype
    flattened_weight = projection.weight.flatten(1)
    flattened_input = hidden_states.reshape(-1, flattened_weight.shape[1])
    return torch.nn.functional.linear(
        flattened_input.to(dtype=target_dtype),
        flattened_weight,
        projection.bias,
    )


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: CloseQwen3VLConfig):
        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: CloseQwen3VLConfig):
        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: CloseQwen3VLConfig):
        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 = Qwen3VLTextRMSNorm(d, eps=text_cfg.rms_norm_eps)
        self.read_norm_e = Qwen3VLTextRMSNorm(d, eps=text_cfg.rms_norm_eps)
        attention_bias = bool(getattr(text_cfg, "attention_bias", False))
        self.q_proj = nn.Linear(
            d, self.num_heads * self.head_dim, bias=attention_bias
        )
        self.k_proj = nn.Linear(
            d, self.num_key_value_heads * self.head_dim, bias=attention_bias
        )
        self.v_proj = nn.Linear(
            d, self.num_key_value_heads * self.head_dim, bias=attention_bias
        )
        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 = Qwen3VLTextRMSNorm(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: CloseQwen3VLConfig):
    """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.
    residual_swap_transition: torch.LongTensor | None = None  # scalar
    #: 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 learns feature-wise depth selection while preserving the pretrained
    hidden basis and keeping cross-node gradient synchronization negligible.
    """

    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:
        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 CloseQwen3VLForConditionalGeneration(Qwen3VLPreTrainedModel):
    def _init_weights(self, module):
        super()._init_weights(module)
        if isinstance(module, CRRFeaturewiseAggregator):
            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: Qwen3-VL with its upper visual path replaced by a read-once workspace."""

    config_class = CloseQwen3VLConfig
    _no_split_modules = ["Qwen3VLTextDecoderLayer", "Qwen3VLVisionBlock"]


    def __init__(self, config: CloseQwen3VLConfig, backbone=None):
        super().__init__(config)
        config.validate()
        config._require_ell_star()

        # The host model. Pass an already-loaded one via :meth:`from_backbone`;
        # `CloseQwen3VLForConditionalGeneration.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 Qwen3VLForConditionalGeneration._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 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:
            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 getattr(config, "splice_scale_match", False)
            )
            else None
        )

        self.post_init()
        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):
                from peft.tuners.tuners_utils import BaseTunerLayer

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

        # Qwen3's dynamic number of patches makes cuDNN repeatedly compile a
        # Conv3d plan.  The kernel covers exactly one non-overlapping patch, so
        # the equivalent linear evaluation is both state-free and faster.
        self.vl.visual.patch_embed.forward = (
            _linear_qwen3_vision_patch_embed_forward.__get__(
                self.vl.visual.patch_embed,
                type(self.vl.visual.patch_embed),
            )
        )

    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/source bias mismatch")
                    dst.bias.copy_(src.bias)
                elif src.bias is not None:
                    raise ValueError("raw GQA/source bias mismatch")
            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: CloseQwen3VLConfig, **kwargs):
        """Build CLOSE around a real Qwen3-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 = Qwen3VLForConditionalGeneration.from_pretrained(pretrained, **kwargs)
        model = cls(config, backbone=backbone)
        return model.to(backbone.dtype)

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

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

    @property
    def text_model(self):
        """``Qwen3VLTextModel``."""
        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)
        for module in (
            getattr(self, "workspace_read", None),
            getattr(self, "workspace_transition", None),
            getattr(self, "raw_eq_transition", None),
            self.workspace_out,
        ):
            if module is None:
                continue
            for p in module.parameters():
                p.requires_grad_(True)
        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())
        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 CRR states with a mean-compatible concat projection."""
        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
            ):
                # Full-state visual CRR stores C_k = R_k - B and decodes the
                # actual final recurrent state without temporal averaging.
                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]

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

    @contextmanager
    def _recurrence_adapter_execution(self, enabled: bool):
        """Toggle the single recurrent LoRA without changing trainability."""
        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, visual_mask, deepstack = embed_multimodal(
            self.vl,
            input_ids,
            pixel_values,
            image_grid_thw,
            attention_mask,
            return_deepstack=True,
        )
        ctx = make_split_context(
            self.text_model,
            embeds,
            pos,
            attention_mask,
            visual_pos_masks=visual_mask,
            deepstack_visual_embeds=deepstack,
        )
        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
                ):
                    self._visual_counterfactual_bundle = (
                        lower.detach(),
                        ctx,
                        text_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]

    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,
    ) -> 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(
                    "Qwen3 full-state recurrence is missing its multimodal "
                    "boundary scaffold"
                )
            return self._encode_visual_full_state_recurrence(
                q_text,
                q_padding_mask,
                ctx,
                visual_bundle,
            )

        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_full_state_recurrence(
        self,
        q_text: torch.Tensor,
        q_padding_mask: torch.Tensor | None,
        text_ctx,
        visual_bundle,
    ) -> list[torch.Tensor]:
        """Run the promoted full-state CRR with one native Qwen3 layer.

        ``R1`` is the frozen pretrained multimodal output.  For ``k>=2``, the
        same physical layer consumes persistent visual rows together with the
        previous question state, and only its recurrence-only LoRA is active.
        The original multimodal cache is never exposed to answer decoding.
        """
        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
        mm_lower, mm_ctx, text_rows = visual_bundle

        # Populate answer-time KV for the recurrent layer from text only.
        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,
            )

        normal_ctx = replace(mm_ctx, past_key_values=None)

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

        # Exact pretrained multimodal anchor; LoRA begins at R1 -> R2.
        with self._recurrence_adapter_execution(False), torch.no_grad():
            first_full = full_cell(mm_lower)
        r1, r_padding = select_tokens_padded(first_full, text_rows)
        if q_padding_mask is not None and not torch.equal(
            r_padding, q_padding_mask
        ):
            raise ValueError(
                "multimodal and text-only padding disagree in Qwen3 recurrence"
            )

        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]
        update_rms = []
        update_relative = []
        r_current = r1

        for _ in range(1, steps):
            state_before = r_current
            recurrent_input = _replace_selected_rows_from_padded(
                first_full,
                text_rows,
                r_current,
                r_padding,
            )

            def recurrent_text_rows(hidden: torch.Tensor):
                # The adapter context must live inside the checkpointed
                # function: non-reentrant checkpoint invokes this callable
                # again during backward after the outer forward scope ended.
                with self._recurrence_adapter_execution(True):
                    full = full_cell(hidden)
                    return select_tokens_padded(full, text_rows)

            if self.training and torch.is_grad_enabled():
                proposed, proposed_padding = checkpoint(
                    recurrent_text_rows,
                    recurrent_input,
                    use_reentrant=False,
                    preserve_rng_state=True,
                )
            else:
                proposed, proposed_padding = recurrent_text_rows(
                    recurrent_input
                )
            if not torch.equal(proposed_padding, r_padding):
                raise RuntimeError("Qwen3 recurrent text-row layout changed")
            transition_delta = proposed - state_before
            if q_padding_mask is not None:
                transition_delta = transition_delta.masked_fill(
                    q_padding_mask.unsqueeze(-1), 0
                )
            r_current = state_before + 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():
                valid = (
                    torch.ones_like(r_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)
                applied = beta * transition_delta
                delta_rms = (
                    (applied.float().square() * valid_f).sum(dim=(1, 2))
                    / (denom * applied.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)
                )

        self.last_recurrent_update_rms = (
            torch.stack(update_rms, dim=1) if update_rms else None
        )
        self.last_recurrent_update_relative = (
            torch.stack(update_relative, dim=1)
            if update_relative
            else None
        )
        return states

    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:
        """Swap one CRR transition and roll the remaining causal tail forward."""
        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)
        aligned = q_keep & q_keep.index_select(0, donors)

        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,
    ) -> 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)
                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_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,
    ) -> 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``.
        Every step uses the same CrossAttn module and directly reads raw ``V*``.
        The visual memory itself is never updated or exposed to the decoder.
        """
        r = q_star
        states = [r]
        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):
            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]
                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
    ) -> 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,
            )

        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,
        return_latent_policy_trace: bool = False,
    ):
        """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
        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")
        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. V* -- the one and only visual read, under no_grad (frozen branch).
        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:
                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),
                    )
                )
            else:
                z_states = self._encode_counterfactual_residual_recurrence(
                    q_star,
                    q_mm,
                    q_pad,
                    ctx,
                    visual_bundle=visual_bundle,
                )
            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
        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:
            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"),
            )
        step_overrides = None
        if reread_step_inputs is not None:
            step_overrides = [None] * self.config.num_workspace_steps
            for k_, inp in reread_step_inputs.items():
                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[int(k_)] = (inp["__v_star__"], inp.get("__v_pad__"))
                else:
                    step_overrides[int(k_)] = 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,
        )
        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)
            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.
        """
        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,
    ) -> 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),
            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,
        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),
            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_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,
        **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 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),
                        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),
                )
            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,
        )

        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
        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")
            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)
                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 or cache 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 score_pos - score_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,
            step_answer_nll=step_answer_nll,
            step_answer_valid=step_answer_valid,
        )

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