Update config and release files

README.md

@@ -88,12 +88,16 @@ inputs = model.prepare_inputs(
     image,
     "What is the dominant color? A. Red B. Blue C. Green D. Yellow. Answer with the option letter.",
 )
-logits = model.next_token_logits(**inputs)
-token_id = logits.argmax(dim=-1).item()
-print(model.tokenizer.decode([token_id]))
+answer_ids = model.generate(
+    **inputs,
+    do_sample=False,
+    max_new_tokens=64,
+    eos_token_id=model.tokenizer.eos_token_id,
+)
+print(model.tokenizer.decode(answer_ids[0], skip_special_tokens=True))
 ```
 
-This model exposes the strict first-answer-token readout used in our analyses. Multi-token generation is not currently supported by this loader.
+This model supports strict first-answer-token readout and deterministic greedy answer generation. The default recurrent update coefficient is `beta=0.5`.
 
 Use one complete model replica per GPU. This custom Transformers loader does
 not support automatic model sharding, generic `save_pretrained()` reserialization,

config.json

@@ -15,16 +15,7 @@
     "recurrent_layer": 36,
     "upper_decoder_start": 37,
     "inference_T": 4,
-    "inference_beta": 0.33,
-    "lora_rank": 32,
-    "lora_alpha": 12.0,
-    "lora_dropout": 0.01,
-    "checkpoint_step": 500,
-    "format": "cvrr_native_plus_merged_transition_v1",
-    "merged_weight_dtype": "float32",
-    "native_weight_bytes": 18277586944,
-    "status": "gpu_vstar_comparison_completed",
-    "upload_ready": true,
-    "hub_model_id": "dmis-lab/InternVL3-9B-CVRR"
+    "inference_beta": 0.5,
+    "format": "cvrr_native_plus_merged_transition_v1"
   }
 }

configuration_source_qwen25.py

@@ -1,1142 +0,0 @@
-"""Configuration for CLOSE latent reasoning on Qwen2.5-VL.
-
-The current ``raw_evidence_question`` path keeps the complete layer-``ell_star``
-visual memory as immutable evidence and recurrently updates the text-only
-question state by cross-attending to it. Its selected interface averages all
-four residual recurrent states and may use a training-only restoration target
-from the normal multimodal question state. Historical slot-workspace fields are
-retained so old checkpoints remain loadable, but they are mutually exclusive
-with the raw-evidence path.
-
-Every architectural hyperparameter of Method §3 lives here -- nothing is
-hardcoded in ``modeling_close_qwen2_5_vl.py``.
-
-Layer-index convention (used consistently across the repo)
-----------------------------------------------------------
-``ell_star`` is **0-indexed and inclusive**: it is the index of the last decoder
-layer belonging to the lower branch. With ``num_hidden_layers == 28``::
-
-    F_{<=l*} = language_model.layers[: ell_star + 1]     # lower_slice
-    F_{>l*}  = language_model.layers[ell_star + 1 :]     # upper_slice
-
-so ``ell_star=18`` means layers 0..18 run before the split and 19..27 after it.
-``ell_star`` is measured once by ``scripts/localize_read_point.py`` (§3.1) and
-frozen before any workspace training.
-
-Careful with ``Qwen2_5_VLConfig``: its ``__setattr__`` forwards any attribute
-whose name already exists in ``text_config.__dict__`` down to the sub-config, and
-its ``__init__`` builds ``text_config`` from ``**kwargs``. Workspace fields are
-therefore declared as explicit named parameters so they never enter ``kwargs``
-and never shadow a text-config field. ``tests/test_config_roundtrip.py`` guards
-this.
-"""
-
-from __future__ import annotations
-
-import math
-
-from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig
-
-#: How the ``S`` workspace slots are assigned M-RoPE positions once they are
-#: spliced into the replacement cache above ``ell_star``. Qwen2.5-VL uses 3D
-#: M-RoPE (t, h, w) with ``mrope_section=[16, 24, 24]``, so slots -- which have
-#: no spatial extent -- need an explicit convention.
-WORKSPACE_ROPE_MODES = (
-    # Slots continue the 1-D text position sequence directly after Q*, with
-    # t == h == w. Treats the workspace as "more text".
-    "continue",
-    # Slots reuse the (t, h, w) span the image tokens occupied in the original
-    # multimodal prefill, subsampled to S positions. Preserves the read point's
-    # positional geometry but reintroduces image-derived indices.
-    "image_span",
-    # All slots share one position (the first position after Q*). Makes the
-    # workspace order-free, matching the set semantics of Perceiver-style slots.
-    "shared",
-)
-
-
-class CloseQwen2_5_VLConfig(Qwen2_5_VLConfig):
-    r"""Configuration for :class:`CloseQwen2_5_VLForConditionalGeneration`.
-
-    Args:
-        ell_star: 0-indexed inclusive last layer of the lower branch. ``None``
-            until §3.1 localization has run; the model refuses to build with
-            ``None`` rather than falling back to a guessed depth.
-        num_workspace_steps: the fixed recurrent horizon (paper notation ``T``;
-            per-item segment counts ``n`` exist only in the losses, never here).
-            The trainable parameter count is independent of it because
-            ``f_theta`` is shared across steps, so it is a free experimental
-            axis. ``0`` is legal and is the no-recurrence floor: ``Z^(0)`` goes
-            straight to the decoder.
-        num_workspace_slots: ``S``, slots in ``Z``. Free capacity knob.
-        workspace_width: ``d_w``, the slot width. Free capacity knob -- the
-            method's claim is path exclusivity (all visual information routed
-            through ``Z``), NOT an information bottleneck, so ``S`` and ``d_w``
-            carry no methodological commitment. Kept below ``d`` (3584 for the
-            7B) since ``P_Z`` projects up at the splice.
-        workspace_num_heads: attention heads inside ``r_theta`` / ``f_theta``.
-        read_num_blocks: cross-attention blocks in ``r_theta``. The *read* is
-            once regardless of depth: all blocks see the same ``V*`` in a single
-            forward, and ``V*`` is discarded afterwards.
-        transition_num_blocks: blocks in the shared recurrent ``f_theta``.
-        workspace_ffn_mult: FFN expansion inside workspace blocks.
-        workspace_dropout: dropout inside workspace blocks.
-        workspace_rope_mode: one of :data:`WORKSPACE_ROPE_MODES`.
-        adapter_rank: rank of the low-rank adapter on the post-``ell_star``
-            backbone layers. ``0`` disables it (frozen backbone only).
-        adapter_alpha: LoRA-style scaling; effective scale is
-            ``adapter_alpha / adapter_rank``.
-        adapter_dropout: dropout probability applied only to the LoRA branch.
-        lambda_traj: legacy trajectory-loss weight. The current method sets it
-            to zero and trains with final-answer CE only; it remains serialized
-            solely so historical checkpoints can still be inspected.
-    """
-
-    model_type = "close_qwen2_5_vl"
-
-    def __init__(
-        self,
-        ell_star: int | None = None,
-        num_workspace_steps: int = 4,
-        num_workspace_slots: int = 32,
-        workspace_width: int = 512,
-        workspace_num_heads: int = 8,
-        read_num_blocks: int = 1,
-        transition_num_blocks: int = 2,
-        workspace_ffn_mult: int = 4,
-        workspace_dropout: float = 0.0,
-        workspace_rope_mode: str = "continue",
-        adapter_rank: int = 16,
-        adapter_alpha: int = 32,
-        adapter_dropout: float = 0.0,
-        adapter_exclude_layers: tuple = (),
-        lambda_traj: float = 0.0,
-        read_shortcut: bool = False,
-        splice_scale_match: bool = False,
-        read_gating: bool = False,
-        gate_init_eps: float = 0.02,
-        read_residual_pure: bool = False,
-        residual_norm_cap: float = 0.0,
-        pure_warm_start: bool = False,
-        final_verify: bool = False,
-        mid_read_step: int = 0,
-        persistent_slot_id: bool = False,
-        mid_read_alpha: float = 1.0,
-        evidence_reasoning: bool = False,
-        raw_evidence_question: bool = False,
-        raw_attention_source_layer: int = -1,
-        raw_transition_ffn: bool | None = None,
-        raw_state_replacement: bool = False,
-        raw_question_anchor: bool = False,
-        raw_mean_aggregation: bool = False,
-        hierarchical_visual_layers: tuple = (),
-        counterfactual_residual_recurrence: bool = False,
-        visual_counterfactual_recurrence: bool = False,
-        visual_cumulative_recurrence: bool = False,
-        visual_full_state_recurrence: bool = False,
-        visual_boundary_reentry: bool = False,
-        visual_source_centered_reentry: bool = False,
-        recurrence_only_adapter: bool = False,
-        visual_preserving_adapter_correction: bool = False,
-        perceive_deliberate_chain: bool = False,
-        local_visual_latent_chain: bool = False,
-        pdlc_transition_conditioned: bool = False,
-        pdlc_question_visible_through_pair: int = 1,
-        pdlc_question_curriculum: bool = False,
-        latent_chain_init_token_id: int = -1,
-        spatial_visual_recurrence: bool = False,
-        spatial_recurrence_steps: int = 8,
-        spatial_recurrence_inner_width: int = 256,
-        spatial_recurrence_heads: int = 8,
-        spatial_step_override: int = -1,
-        counterfactual_beta: float = 1.0,
-        crr_concat_aggregation: bool = False,
-        crr_policy_rank: int = 0,
-        crr_policy_alpha: float = 32.0,
-        counterfactual_reliance_loss: bool = False,
-        transition_aware_reliance_loss: bool = False,
-        corrective_transition_reliance_loss: bool = False,
-        counterfactual_step_retention_loss: bool = False,
-        functional_state_loss: bool = False,
-        per_example_answer_loss: bool = False,
-        interface_loss: bool = False,
-        read_anchor: bool = False,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.ell_star = ell_star
-        self.num_workspace_steps = num_workspace_steps
-        self.num_workspace_slots = num_workspace_slots
-        self.workspace_width = workspace_width
-        self.workspace_num_heads = workspace_num_heads
-        self.read_num_blocks = read_num_blocks
-        self.transition_num_blocks = transition_num_blocks
-        self.workspace_ffn_mult = workspace_ffn_mult
-        self.workspace_dropout = workspace_dropout
-        self.workspace_rope_mode = workspace_rope_mode
-        self.adapter_rank = adapter_rank
-        self.adapter_alpha = adapter_alpha
-        self.adapter_dropout = adapter_dropout
-        # E run: dense full-FT 층은 LoRA 대상에서 제외 (이중 파라미터화 방지)
-        self.adapter_exclude_layers = list(adapter_exclude_layers or ())
-        self.lambda_traj = lambda_traj
-        # v2 anti-collapse switches (defaults False so v1 checkpoints reload).
-        # read_shortcut: r_theta adds an orthogonally-projected pooled-V* term to
-        # every slot, so Z^(0) is image-dependent from step 0 instead of relying
-        # on a random cross-attention finding signal.
-        # splice_scale_match: P_Z(Z) rows are RMS-matched to the question rows at
-        # the splice. Unmatched, slot states are ~O(1) against layer-20 hiddens
-        # in the hundreds -- the decoder could ignore them for free, and did
-        # (probe 20143: visual sensitivity 0.94 -> 1.0000 over training).
-        self.read_shortcut = read_shortcut
-        self.splice_scale_match = splice_scale_match
-        # v6 candidate: gated re-reads DURING the recurrence. Every re-read
-        # goes through r_theta into Z, so the mediation invariant (decoder sees
-        # only [Q*; P_Z(Z)]) is untouched; what changes is that the workspace
-        # may consult V* again mid-reasoning, scaled by a learned scalar gate.
-        # The gate is free at EVERY step -- no hand-coded schedule. Degeneration
-        # into answer-time retrieval is disciplined by the objective itself:
-        # L_traj pins each intermediate state to its teacher segment, so a
-        # wait-then-fetch trajectory misses its intermediate targets. The gate
-        # bias initialises negative, i.e. training starts from the proven
-        # no-reread regime and departs only where gradients demand it.
-        # Default False: v4/v5 checkpoints reload bit-identically.
-        self.read_gating = read_gating
-        # Initial gate open-rate epsilon_gamma: b_gamma = logit(eps). Sigmoid
-        # never reaches exactly 0 at finite logits, so this is near-identity
-        # (not bitwise) initialisation w.r.t. the read-once path.
-        self.gate_init_eps = gate_init_eps
-        # v7: always-on PURE visual residual rereads. Every step adds
-        # Z = U + W_o XAttn(LN(U) -> queries, V* -> keys/values): no self-attn,
-        # no FFN, no output norm, no biases in the branch, W_o zero-init.
-        # Motivation (battery, ep1 ckpt 20752): read_once hurt badly while
-        # null_step (blank memory at one step) was harmless -- the gated
-        # reader's value was extra recurrent DEPTH, not vision. This branch
-        # removes every path that can change Z without V* content, so any
-        # benefit it shows IS visual by construction.
-        self.read_residual_pure = read_residual_pure
-        # v7.1 (advisor): relative-norm trust region on the visual residual.
-        #   s_k = min(1, rho * RMS(U) / (RMS(Delta) + eps));  Z = U + s_k * Delta
-        # Unlike a fixed alpha, the optimizer cannot undo this by rescaling
-        # W_o (s_k adapts); ||Delta_used|| <~ rho * ||U|| structurally, and
-        # V* = 0 => Delta = 0 keeps the zero-memory identity. 0.0 disables.
-        self.residual_norm_cap = residual_norm_cap
-        # v7.1: initialise pure_attn K/V projections from the initial visual
-        # reader's v_in (same [d_w, d] shape) and Q from its first read
-        # block's in-proj slice, instead of random xavier -- the ep1 finding
-        # was that Q/K/V stayed ~random while only W_o trained.
-        self.pure_warm_start = pure_warm_start
-        # MAIN METHOD (advisor-confirmed): Final-Verified Recurrent Workspace.
-        # Initial read -> PURE latent recurrence (no mid-step visual anything)
-        # -> exactly ONE pure visual verification at the end:
-        #   Z_out = U^(T) + W_o XAttn(LN(U^(T)), V*),  alpha = 1, W_o = 0 init.
-        # Uses the same bias-free verifier module as read_residual_pure; the
-        # two flags differ only in WHERE the branch fires (every step vs once).
-        self.final_verify = final_verify
-        # MAIN (advisor 최종 확정): Mid-Read Recurrent Workspace. 고정 상수
-        # k* = floor(T/2) = 4 (스텝 sweep 산물이 아니라 판독 전후 동수 전이의
-        # midpoint)에서 pure visual read 1회:
-        #   Z^(4) = U^(4) + W_o XAttn(LN(U^(4)), V*),  alpha=1, W_o=0 init.
-        # 논리적 step 4 == python 루프 index 3 (0-based) -- off-by-one은
-        # tests/test_mid_read.py가 상태 분기 위치로 고정한다. 0이면 비활성.
-        self.mid_read_step = mid_read_step
-        # advisor(2026-08-06): E_slot(=workspace_read.slots)을 매 recurrent
-        # step 입력과 mid-read query에 일시 제공. state에는 누적하지 않는다.
-        self.persistent_slot_id = persistent_slot_id
-        # 고정 read 스케일 Z = U + alpha*W_o A (advisor scale-mismatch 검사).
-        # 1.0 = 기존과 비트동일. eval의 gate_overrides alpha는 절대값으로 이를 대체.
-        self.mid_read_alpha = mid_read_alpha
-        # advisor 최종안: E(persistent evidence) / R(recurrent reasoning) 분리.
-        # states = [E^(0), R^(1..T)], 결합 Z^(T)=E+R^(T)는 prefill에서.
-        self.evidence_reasoning = evidence_reasoning
-        # Current replacement candidate: no learned evidence/reasoning slots.
-        # E is the complete raw V* sequence and R0 is the text-only Q* sequence.
-        # A single shared CrossAttn transition produces R1..RT. Residual mode
-        # decodes RT; the replacement candidate decodes Q*+HT. Cross-attention
-        # uses the backbone's native GQA shape and
-        # is initialized from the first upper layer by default. Checkpoints
-        # created before the CrossAttn-only decision did not serialize
-        # raw_transition_ffn; None therefore preserves their historical FFN
-        # for faithful evaluation, while every new training launch passes False.
-        self.raw_evidence_question = bool(raw_evidence_question)
-        self.raw_attention_source_layer = int(raw_attention_source_layer)
-        self.raw_transition_ffn = (
-            True if raw_transition_ffn is None else bool(raw_transition_ffn)
-        )
-        # Experimental simplification of the same raw-E/Q path.  The residual
-        # carrier is replaced, not augmented:
-        #   H0=Q*, Hk=CrossAttn(H{k-1}, E), decoder_latent=Q*+HT.
-        # False preserves all residual raw-E/Q checkpoints and active runs.
-        self.raw_state_replacement = bool(raw_state_replacement)
-        # Question-anchored replacement keeps the fixed text scaffold in every
-        # recurrent query without restoring the old recurrent identity path:
-        #   Z0=Q*, Zk=Q*+CrossAttn(Z{k-1}, E), decoder_latent=ZT.
-        # This replaces the one-time Q*+HT decoder interface above.
-        self.raw_question_anchor = bool(raw_question_anchor)
-        # Q-shaped recurrent candidate selected by the V* interface oracle:
-        #   R0=Q*, Rk=R{k-1}+CrossAttn(R{k-1}, V*)
-        #   R_agg=mean(R1..RT)
-        # This is parameter-free and replaces both question-anchored state
-        # replacement and final-state-only decoding.
-        self.raw_mean_aggregation = bool(raw_mean_aggregation)
-        # A non-empty sequence maps the first T-1 recurrent steps to frozen
-        # visual-token states from the listed language layers.  The final step
-        # uses Q* as memory to integrate the staged reads without another image
-        # access.  One CrossAttention module is shared by all T steps.
-        self.hierarchical_visual_layers = [
-            int(layer) for layer in (hierarchical_visual_layers or ())
-        ]
-        # Counterfactual Residual Recurrence (CRR): the image-conditioned and
-        # text-only question states are propagated by the same native layer
-        # F=L(ell_star+1).  Only their counterfactual residual is recurrently
-        # refined; there are no slots, V* memory, custom attention modules, or
-        # inference-time auxiliary modules.  beta is the fixed relaxed-update
-        # coefficient; an optional batch-swap reliance loss is training-only.
-        self.counterfactual_residual_recurrence = bool(
-            counterfactual_residual_recurrence
-        )
-        # Causal Visual Residual Recurrence (CVRR), implemented as the
-        # vision-native CRR transition.  At every recurrent step the same
-        # native layer is evaluated with and without only Q->vision attention;
-        # their difference is the image-caused update.  It deliberately
-        # reuses CRR's strict Q-shaped decoder/cache interface and mean
-        # aggregation, hence the parent CRR flag remains explicit.
-        self.visual_counterfactual_recurrence = bool(
-            visual_counterfactual_recurrence
-        )
-        # Minimal cumulative CVRR variant.  Unlike the historical anchored
-        # mean path, visual updates are accumulated into the running state and
-        # the final state is decoded directly:
-        #   R_k = R_{k-1} + beta * Delta_k^V,  decode R_T.
-        # One explicit flag controls both choices so a checkpoint cannot
-        # silently combine cumulative dynamics with the old mean interface.
-        self.visual_cumulative_recurrence = bool(
-            visual_cumulative_recurrence
-        )
-        # Full-state CVRR keeps the native recurrent computation instead of
-        # treating the visual on/off intervention difference as the next
-        # hidden state:
-        #   R_tilde_k = Pi_Q F([E; R_{k-1}])
-        #   R_k = R_{k-1} + beta * (R_tilde_k - R_{k-1}).
-        # The visual-off branch remains available to evaluation code, but is
-        # not part of this model's train/inference transition.
-        self.visual_full_state_recurrence = bool(
-            visual_full_state_recurrence
-        )
-        # Input-grounded closure control.  Later recurrent calls rebuild the
-        # native post-L20 scaffold and inject only the carried question-state
-        # deviation, rather than feeding post-L21 rows back into L21.
-        self.visual_boundary_reentry = bool(visual_boundary_reentry)
-        # Source-centered boundary recurrence.  The text-only L21 state is a
-        # fixed anchor, while the native multimodal residual initializes a
-        # nonzero, image-dependent deviation.  The recurrent cell evolves
-        # only that deviation after subtracting its zero-deviation response.
-        self.visual_source_centered_reentry = bool(
-            visual_source_centered_reentry
-        )
-        # Restrict LoRA to the single shared recurrent cell L(ell_star+1).
-        # Its execution is disabled for R1, the text-only cache, answer-token
-        # decoding, and L22+; only recurrent transitions k>=2 enable it.
-        self.recurrence_only_adapter = bool(recurrence_only_adapter)
-        # Preserve the pretrained image-caused transition while removing the
-        # frozen layer's non-visual recurrent drift:
-        #   Delta_k = F_{base+LoRA,on}(E,R_{k-1}) - F_{base,off}(E,R_{k-1})
-        #           = (F_{base,on} - F_{base,off})
-        #             + (F_{base+LoRA,on} - F_{base,on}).
-        # This is a full-state recurrence subtype and introduces no new
-        # parameters or inference-time module.
-        self.visual_preserving_adapter_correction = bool(
-            visual_preserving_adapter_correction
-        )
-        # Local Visual Latent Chain (LVLC): K homogeneous latent-token
-        # positions are evaluated in one native Transformer forward.  The
-        # first latent sees clean Q plus visual source rows; every later latent
-        # sees only the immediately preceding latent plus the same visual
-        # rows; answer rows see clean Q plus the final latent.  This is a local
-        # attention graph, not a recurrent forward or temporal aggregation.
-        self.local_visual_latent_chain = bool(local_visual_latent_chain)
-        # Perceive--Deliberate Latent Chain (PDLC).  K must be even and is
-        # interpreted as K/2 native LOOK/THINK pairs.  The complete VLM stack
-        # runs once over a block-sparse graph: LOOK rows may inspect the
-        # multimodal prompt, THINK rows may inspect only clean question rows
-        # and their paired LOOK, and answer rows may inspect only clean
-        # question rows and THINK states.  There is no recurrent cell, visual
-        # state update, split-state splice, or temporal aggregator.
-        # LVLC reuses the same one-pass latent-sequence execution plumbing but
-        # replaces the LOOK/THINK token types and visibility graph entirely.
-        self.perceive_deliberate_chain = bool(
-            perceive_deliberate_chain or self.local_visual_latent_chain
-        )
-        # Tight PDLC graph: pair 1 bootstraps from prompt/question, later
-        # LOOK rows receive visual placeholders plus the preceding latent
-        # prefix, later THINK rows receive only that causal latent prefix, and
-        # the answer consumes final THINK.  It changes visibility only and
-        # allocates no additional trainable module.
-        self.pdlc_transition_conditioned = bool(
-            pdlc_transition_conditioned
-        )
-        # Strict inference uses 1: only the bootstrap pair reads the clean
-        # question.  Larger values are an explicit evaluation intervention.
-        # The training-only curriculum changes a runtime override and always
-        # returns to 1 for validation/inference.
-        self.pdlc_question_visible_through_pair = int(
-            pdlc_question_visible_through_pair
-        )
-        self.pdlc_question_curriculum = bool(pdlc_question_curriculum)
-        # Both learned latent token types start from one native text embedding
-        # (the launcher supplies the tokenizer's newline id).  -1 falls back
-        # to token 0 for synthetic/unit-test configs only.
-        self.latent_chain_init_token_id = int(latent_chain_init_token_id)
-        # Spatial Recurrent Visual Field (SRVF).  Unlike all historical
-        # ell_star methods, the recurrent state is the complete native visual
-        # token grid immediately after the vision merger.  The final field is
-        # inserted into the ordinary Qwen prefix and the full decoder runs from
-        # scratch, making the zero-output initialization exactly base-equivalent.
-        self.spatial_visual_recurrence = bool(spatial_visual_recurrence)
-        self.spatial_recurrence_steps = int(spatial_recurrence_steps)
-        self.spatial_recurrence_inner_width = int(
-            spatial_recurrence_inner_width
-        )
-        self.spatial_recurrence_heads = int(spatial_recurrence_heads)
-        # -1 uses the trained/default horizon.  Non-negative values are an
-        # evaluation-only prefix sweep, with 0 being the exact base path.
-        self.spatial_step_override = int(spatial_step_override)
-        self.counterfactual_beta = float(counterfactual_beta)
-        # Learn one feature-wise mixture of the complete recurrent residual
-        # trajectory.  The bias-free T*d -> d map starts from
-        # [I/T, ..., I/T], so enabling it is exactly the historical CRR mean
-        # interface before the first optimizer step.  A dedicated flag keeps
-        # old CRR checkpoints and the historical slot-ER aggregator distinct.
-        self.crr_concat_aggregation = bool(crr_concat_aggregation)
-        # Optional latent-only policy adapter.  It corrects the deterministic
-        # C2..CT means but never touches B, C1, or answer-token decoding.  Rank 0
-        # is the exact historical CRR path; RL checkpoints serialize rank > 0.
-        self.crr_policy_rank = int(crr_policy_rank)
-        self.crr_policy_alpha = float(crr_policy_alpha)
-        # Training-only CRR intervention.  A batch-matched residual from a
-        # different-answer example replaces the factual residual before the
-        # upper decoder.  The model returns the factual-minus-swapped answer
-        # score gap; the Trainer applies the margin loss.  Evaluation and
-        # generation never execute the extra branch.
-        self.counterfactual_reliance_loss = bool(
-            counterfactual_reliance_loss
-        )
-        # Replace the aggregate residual swap with one intervention inside the
-        # recurrent chain.  C_k is swapped and C_{k+1:T} is rolled out again,
-        # so the answer gap is attributable to a selected transition rather
-        # than to an undifferentiated full-trajectory replacement.
-        self.transition_aware_reliance_loss = bool(
-            transition_aware_reliance_loss
-        )
-        # Training-only hard-case corrective ranking.  At one selected
-        # transition k>=2, the actual decoder scores Z_{k-1}, Z_k, and a
-        # donor-swapped Z_k.  The Trainer requires the factual transition to
-        # improve over both detached comparators.  This replaces, rather than
-        # stacks with, the historical transition-aware reliance objective.
-        self.corrective_transition_reliance_loss = bool(
-            corrective_transition_reliance_loss
-        )
-        # Training-only functional no-regression objective.  Gold-answer NLL
-        # is measured through the actual upper decoder at every recurrent CRR
-        # prefix; inference remains the unchanged T-step mean interface.
-        self.counterfactual_step_retention_loss = bool(
-            counterfactual_step_retention_loss
-        )
-        # Training-only self-distillation target.  The multimodal lower pass
-        # already computed for V* also provides its image-conditioned question
-        # rows.  They supervise R_agg at the split boundary, but are never
-        # exposed to the inference decoder.
-        self.functional_state_loss = bool(functional_state_loss)
-        # Token CE lets long caption answers dominate one-token visual-choice
-        # examples.  The per-example variant first averages valid answer tokens
-        # within each sample, then averages samples.  It changes training/eval
-        # loss reduction only and has no inference-time effect.
-        self.per_example_answer_loss = bool(per_example_answer_loss)
-        # ER v2 (advisor 2026-08-11): explicit evidence read —
-        # R^(k) = f_theta(R^(k-1), Q*; E), E는 별도 memory로 cross-attn 읽기.
-        # (E+R 합산-감산 제거. evidence_reasoning=True 전제.)
-        self.explicit_evidence_read = bool(kwargs.pop("explicit_evidence_read", False))
-        # Historical field name retained for checkpoint compatibility. In ER
-        # v2.2 this means R^(0) is a batch-broadcast set of learnable slots; it
-        # is deliberately NOT initialized from Q*. Q* conditions every shared
-        # transition instead.
-        self.r_init_question = bool(kwargs.pop("r_init_question", False))
-        # ER v2.1 (advisor 2026-08-11): evidence read의 competitive normalization —
-        # slot 축 softmax 후 evidence 축 재정규화 (Slot Attention식 경쟁).
-        # 파라미터/모듈/loss 무추가; evidence_read의 W_qkv/W_o 재사용.
-        self.competitive_evidence_read = bool(
-            kwargs.pop("competitive_evidence_read", False))
-        # ER v2.2 C: k>=1 traj 감독 부재 시 twin 체인 제거 (states = answer 체인).
-        self.er_single_chain = bool(kwargs.pop("er_single_chain", False))
-        # ER v2.2 C (§5-7): decoder latent = softmax(temporal_logits)로 가중한
-        # R1..RT 결합 (전역 스칼라 T개). er_single_chain 전제. (legacy)
-        self.temporal_aggregation = bool(kwargs.pop("temporal_aggregation", False))
-        # ER v2.2 C 확정판: R_agg = Linear(Concat(R1..RT), bias=False),
-        # 초기 W=[I/T,...,I/T] (R_agg=mean(R1..RT)). er_single_chain 전제.
-        self.concat_aggregation = bool(kwargs.pop("concat_aggregation", False))
-        # ER v2: 워크스페이스 블록의 rank-r 부분공간 attention/FFN (0 = dense).
-        # d_w=d에서 state width 유지와 파라미터 폭증을 분리 (advisor).
-        self.workspace_low_rank = int(kwargs.pop("workspace_low_rank", 0))
-        # advisor: single-layer text-side interface alignment (bridge-only).
-        # teacher = frozen-base mm 경로 q_end@l*+1 (adapter off, sg);
-        # student = detach(R^T) splice의 l*+1 한 층 재계산. 추론 시 무존재.
-        self.interface_loss = interface_loss
-        # Training-only visual anchor switch. False in the current answer-only
-        # method; when false the model does not even materialize pooled V* for
-        # the Trainer.
-        self.read_anchor = bool(read_anchor)
-        if evidence_reasoning and persistent_slot_id:
-            raise ValueError("evidence_reasoning은 persistent_slot_id와 동시 사용 불가")
-        if evidence_reasoning and not mid_read_step:
-            raise ValueError("evidence_reasoning은 mid_read_step > 0 필요")
-        self.validate()
-
-    # -- derived views -----------------------------------------------------
-
-    @property
-    def num_decoder_layers(self) -> int:
-        return self.text_config.num_hidden_layers
-
-    @property
-    def backbone_width(self) -> int:
-        """``d`` -- backbone hidden width (3584 for the 7B)."""
-        return self.text_config.hidden_size
-
-    @property
-    def lower_slice(self) -> slice:
-        """Layers forming ``F_{<=l*}``."""
-        self._require_ell_star()
-        return slice(0, self.ell_star + 1)
-
-    @property
-    def upper_slice(self) -> slice:
-        """Layers forming ``F_{>l*}``."""
-        self._require_ell_star()
-        return slice(self.ell_star + 1, self.num_decoder_layers)
-
-    # -- validation --------------------------------------------------------
-
-    def _require_ell_star(self) -> None:
-        if self.ell_star is None:
-            raise ValueError(
-                "ell_star is unset. Run scripts/localize_read_point.py (Method 3.1) "
-                "and pass the measured layer explicitly; it must not be guessed."
-            )
-
-    def validate(self) -> None:
-        """Reject configurations that silently break the §3.2 contract."""
-        # ``PretrainedConfig.from_pretrained`` may apply user overrides with
-        # setattr *after* __init__.  In that path, setting LVLC=True would not
-        # re-run the constructor's promotion of the shared one-pass execution
-        # flag.  Re-establish the invariant before any model reads the config.
-        if getattr(self, "local_visual_latent_chain", False):
-            self.perceive_deliberate_chain = True
-        if self.ell_star is not None:
-            n = self.num_decoder_layers
-            # Both branches must be non-empty: an empty lower branch means there
-            # is no V* to read, an empty upper branch means the workspace never
-            # reaches the decoder.
-            if not 0 <= self.ell_star <= n - 2:
-                raise ValueError(
-                    f"ell_star={self.ell_star} out of range for {n} decoder layers; "
-                    f"expected 0 <= ell_star <= {n - 2} so both branches are non-empty."
-                )
-        # K=0 is the no-recurrence control (read only), not a misconfiguration.
-        if self.num_workspace_steps < 0:
-            raise ValueError("num_workspace_steps (K) must be >= 0.")
-        if self.num_workspace_slots < 1:
-            raise ValueError("num_workspace_slots (S) must be >= 1.")
-        if self.workspace_num_heads < 1:
-            raise ValueError("workspace_num_heads must be >= 1.")
-        if self.workspace_width % self.workspace_num_heads != 0:
-            raise ValueError(
-                f"workspace_width={self.workspace_width} must be divisible by "
-                f"workspace_num_heads={self.workspace_num_heads}."
-            )
-        if self.workspace_width > self.backbone_width:
-            raise ValueError(
-                f"workspace_width (d_w={self.workspace_width}) cannot exceed "
-                f"backbone width (d={self.backbone_width}). Native-width d_w=d is "
-                "the current method and uses an Identity decoder interface."
-            )
-        if self.workspace_rope_mode not in WORKSPACE_ROPE_MODES:
-            raise ValueError(
-                f"workspace_rope_mode={self.workspace_rope_mode!r} not in "
-                f"{WORKSPACE_ROPE_MODES}."
-            )
-        if self.adapter_rank < 0:
-            raise ValueError("adapter_rank must be >= 0 (0 disables the adapter).")
-        if not math.isfinite(self.adapter_dropout) or not 0.0 <= self.adapter_dropout < 1.0:
-            raise ValueError("adapter_dropout must be finite and in [0, 1).")
-        if getattr(self, "crr_policy_rank", 0) < 0:
-            raise ValueError("crr_policy_rank must be >= 0.")
-        if not math.isfinite(getattr(self, "crr_policy_alpha", 0.0)) or getattr(
-            self, "crr_policy_alpha", 0.0
-        ) <= 0.0:
-            raise ValueError("crr_policy_alpha must be finite and > 0.")
-        if self.read_num_blocks < 1 or self.transition_num_blocks < 1:
-            raise ValueError("read_num_blocks and transition_num_blocks must be >= 1.")
-        if self.workspace_ffn_mult < 1:
-            raise ValueError("workspace_ffn_mult must be >= 1.")
-        mid = int(getattr(self, "mid_read_step", 0) or 0)
-        if mid < 0 or mid > self.num_workspace_steps:
-            raise ValueError(
-                f"mid_read_step={mid} must be in [0, num_workspace_steps="
-                f"{self.num_workspace_steps}]."
-            )
-        low_rank = int(getattr(self, "workspace_low_rank", 0) or 0)
-        if low_rank < 0:
-            raise ValueError("workspace_low_rank must be >= 0.")
-        if low_rank and low_rank % self.workspace_num_heads != 0:
-            raise ValueError(
-                f"workspace_low_rank={low_rank} must be divisible by "
-                f"workspace_num_heads={self.workspace_num_heads}."
-            )
-        if getattr(self, "competitive_evidence_read", False) and not getattr(
-            self, "explicit_evidence_read", False
-        ):
-            raise ValueError(
-                "competitive_evidence_read requires explicit_evidence_read=True; "
-                "otherwise the flag has no execution path."
-            )
-        if getattr(self, "explicit_evidence_read", False) and not getattr(
-            self, "evidence_reasoning", False
-        ):
-            raise ValueError(
-                "explicit_evidence_read requires evidence_reasoning=True; otherwise "
-                "the evidence-read module is allocated but never called."
-            )
-        if getattr(self, "r_init_question", False) and not getattr(
-            self, "evidence_reasoning", False
-        ):
-            raise ValueError(
-                "r_init_question (legacy name for learnable R0 slots) requires "
-                "evidence_reasoning=True; otherwise reasoning_slots are unused."
-            )
-        if (getattr(self, "concat_aggregation", False)
-                or getattr(self, "temporal_aggregation", False)):
-            if self.num_workspace_steps < 1:
-                raise ValueError("reasoning aggregation requires num_workspace_steps >= 1.")
-            if not getattr(self, "evidence_reasoning", False):
-                raise ValueError(
-                    "reasoning aggregation requires evidence_reasoning=True; "
-                    "otherwise the aggregation module is unused."
-                )
-        if getattr(self, "concat_aggregation", False) and not getattr(
-            self, "er_single_chain", False
-        ):
-            raise ValueError("concat_aggregation requires er_single_chain=True.")
-        if getattr(self, "raw_evidence_question", False):
-            if self.workspace_width != self.backbone_width:
-                raise ValueError(
-                    "raw_evidence_question requires native workspace_width == "
-                    "backbone_width: E=V* and R0=Q* are used without projections."
-                )
-            if self.num_workspace_steps < 1:
-                raise ValueError(
-                    "raw_evidence_question requires num_workspace_steps >= 1."
-                )
-            source_layer = (
-                self.ell_star + 1
-                if self.raw_attention_source_layer < 0
-                else self.raw_attention_source_layer
-            )
-            if not self.ell_star < source_layer < self.text_config.num_hidden_layers:
-                raise ValueError(
-                    "raw_attention_source_layer must be an upper-backbone layer, "
-                    f"got ell_star={self.ell_star}, source={source_layer}."
-                )
-            incompatible = {
-                "evidence_reasoning": bool(getattr(self, "evidence_reasoning", False)),
-                "explicit_evidence_read": bool(getattr(self, "explicit_evidence_read", False)),
-                "competitive_evidence_read": bool(getattr(self, "competitive_evidence_read", False)),
-                "r_init_question": bool(getattr(self, "r_init_question", False)),
-                "er_single_chain": bool(getattr(self, "er_single_chain", False)),
-                "temporal_aggregation": bool(getattr(self, "temporal_aggregation", False)),
-                "concat_aggregation": bool(getattr(self, "concat_aggregation", False)),
-                "read_shortcut": bool(getattr(self, "read_shortcut", False)),
-                "splice_scale_match": bool(getattr(self, "splice_scale_match", False)),
-                "read_gating": bool(getattr(self, "read_gating", False)),
-                "read_residual_pure": bool(getattr(self, "read_residual_pure", False)),
-                "final_verify": bool(getattr(self, "final_verify", False)),
-                "mid_read_step": bool(getattr(self, "mid_read_step", 0)),
-                "persistent_slot_id": bool(getattr(self, "persistent_slot_id", False)),
-                "interface_loss": bool(getattr(self, "interface_loss", False)),
-                "read_anchor": bool(getattr(self, "read_anchor", False)),
-            }
-            active = [name for name, enabled in incompatible.items() if enabled]
-            if active:
-                raise ValueError(
-                    "raw_evidence_question replaces the slot workspace and is "
-                    f"incompatible with: {', '.join(active)}"
-                )
-            if self.raw_state_replacement and self.raw_transition_ffn:
-                raise ValueError(
-                    "raw_state_replacement is CrossAttn-only and is incompatible "
-                    "with raw_transition_ffn=True."
-                )
-            if self.raw_question_anchor and not self.raw_state_replacement:
-                raise ValueError(
-                    "raw_question_anchor requires raw_state_replacement=True."
-                )
-            if self.raw_mean_aggregation and (
-                self.raw_state_replacement or self.raw_question_anchor
-            ):
-                raise ValueError(
-                    "raw_mean_aggregation replaces raw_state_replacement and "
-                    "raw_question_anchor; use the residual recurrence."
-                )
-            hierarchy = list(
-                getattr(self, "hierarchical_visual_layers", ()) or ()
-            )
-            if hierarchy:
-                if not self.raw_mean_aggregation:
-                    raise ValueError(
-                        "hierarchical_visual_layers requires "
-                        "raw_mean_aggregation=True."
-                    )
-                if self.raw_transition_ffn:
-                    raise ValueError(
-                        "hierarchical visual recurrence is CrossAttn-only and "
-                        "requires raw_transition_ffn=False."
-                    )
-                expected_visual_reads = self.num_workspace_steps - 1
-                if len(hierarchy) != expected_visual_reads:
-                    raise ValueError(
-                        "hierarchical_visual_layers must contain exactly one "
-                        "layer per visual-read step, followed by one Q* "
-                        "integration step: "
-                        f"got {len(hierarchy)} layers for "
-                        f"T={self.num_workspace_steps} "
-                        f"(expected {expected_visual_reads})."
-                    )
-                if hierarchy != sorted(set(hierarchy)):
-                    raise ValueError(
-                        "hierarchical_visual_layers must be strictly increasing."
-                    )
-                invalid = [
-                    layer
-                    for layer in hierarchy
-                    if not 0 <= layer <= int(self.ell_star)
-                ]
-                if invalid:
-                    raise ValueError(
-                        "hierarchical_visual_layers must lie in the frozen lower "
-                        f"branch [0,{self.ell_star}], got {invalid}."
-                    )
-                if self.functional_state_loss:
-                    raise ValueError(
-                        "functional_state_loss is not defined for hierarchical "
-                        "visual recurrence; use final-answer CE only."
-                    )
-            if self.functional_state_loss and not self.raw_mean_aggregation:
-                raise ValueError(
-                    "functional_state_loss requires raw_mean_aggregation=True "
-                    "so its target is the exact decoder-facing recurrence."
-                )
-        elif getattr(self, "raw_state_replacement", False):
-            raise ValueError(
-                "raw_state_replacement requires raw_evidence_question=True."
-            )
-        elif getattr(self, "raw_question_anchor", False):
-            raise ValueError(
-                "raw_question_anchor requires raw_evidence_question=True."
-            )
-        elif getattr(self, "raw_mean_aggregation", False):
-            raise ValueError(
-                "raw_mean_aggregation requires raw_evidence_question=True."
-            )
-        elif getattr(self, "hierarchical_visual_layers", ()):
-            raise ValueError(
-                "hierarchical_visual_layers requires "
-                "raw_evidence_question=True."
-            )
-        elif getattr(self, "functional_state_loss", False):
-            raise ValueError(
-                "functional_state_loss requires raw_evidence_question=True."
-            )
-        if (
-            getattr(self, "counterfactual_reliance_loss", False)
-            and not getattr(self, "counterfactual_residual_recurrence", False)
-        ):
-            raise ValueError(
-                "counterfactual_reliance_loss requires "
-                "counterfactual_residual_recurrence=True."
-            )
-        if (
-            getattr(self, "transition_aware_reliance_loss", False)
-            and not getattr(self, "counterfactual_reliance_loss", False)
-        ):
-            raise ValueError(
-                "transition_aware_reliance_loss requires "
-                "counterfactual_reliance_loss=True."
-            )
-        if (
-            getattr(self, "transition_aware_reliance_loss", False)
-            and self.num_workspace_steps < 2
-        ):
-            raise ValueError(
-                "transition_aware_reliance_loss requires at least two steps."
-            )
-        if (
-            getattr(self, "corrective_transition_reliance_loss", False)
-            and not getattr(self, "counterfactual_reliance_loss", False)
-        ):
-            raise ValueError(
-                "corrective_transition_reliance_loss requires "
-                "counterfactual_reliance_loss=True."
-            )
-        if (
-            getattr(self, "corrective_transition_reliance_loss", False)
-            and self.num_workspace_steps < 2
-        ):
-            raise ValueError(
-                "corrective_transition_reliance_loss requires at least two steps."
-            )
-        if (
-            getattr(self, "corrective_transition_reliance_loss", False)
-            and getattr(self, "transition_aware_reliance_loss", False)
-        ):
-            raise ValueError(
-                "corrective_transition_reliance_loss replaces, and cannot be "
-                "combined with, transition_aware_reliance_loss."
-            )
-        if (
-            getattr(self, "corrective_transition_reliance_loss", False)
-            and getattr(self, "counterfactual_step_retention_loss", False)
-        ):
-            raise ValueError(
-                "corrective_transition_reliance_loss replaces, and cannot be "
-                "combined with, counterfactual_step_retention_loss."
-            )
-        if (
-            getattr(self, "counterfactual_step_retention_loss", False)
-            and not getattr(self, "counterfactual_residual_recurrence", False)
-        ):
-            raise ValueError(
-                "counterfactual_step_retention_loss requires "
-                "counterfactual_residual_recurrence=True."
-            )
-        if (
-            getattr(self, "counterfactual_step_retention_loss", False)
-            and self.num_workspace_steps < 2
-        ):
-            raise ValueError(
-                "counterfactual_step_retention_loss requires at least two steps."
-            )
-        if getattr(self, "crr_concat_aggregation", False) and not getattr(
-            self, "counterfactual_residual_recurrence", False
-        ):
-            raise ValueError(
-                "crr_concat_aggregation requires "
-                "counterfactual_residual_recurrence=True."
-            )
-        if getattr(self, "counterfactual_residual_recurrence", False):
-            if getattr(self, "raw_evidence_question", False):
-                raise ValueError(
-                    "counterfactual_residual_recurrence replaces "
-                    "raw_evidence_question; enable exactly one q-state method."
-                )
-            if self.num_workspace_steps < 1:
-                raise ValueError(
-                    "counterfactual_residual_recurrence requires "
-                    "num_workspace_steps >= 1."
-                )
-            if self.crr_policy_rank > 0 and self.num_workspace_steps < 2:
-                raise ValueError(
-                    "crr_policy_rank > 0 requires num_workspace_steps >= 2."
-                )
-            if not 0.0 <= self.counterfactual_beta <= 1.0:
-                raise ValueError(
-                    "counterfactual_beta must lie in [0, 1], got "
-                    f"{self.counterfactual_beta}."
-                )
-            incompatible = {
-                "evidence_reasoning": bool(getattr(self, "evidence_reasoning", False)),
-                "explicit_evidence_read": bool(getattr(self, "explicit_evidence_read", False)),
-                "competitive_evidence_read": bool(getattr(self, "competitive_evidence_read", False)),
-                "r_init_question": bool(getattr(self, "r_init_question", False)),
-                "er_single_chain": bool(getattr(self, "er_single_chain", False)),
-                "temporal_aggregation": bool(getattr(self, "temporal_aggregation", False)),
-                "concat_aggregation": bool(getattr(self, "concat_aggregation", False)),
-                "read_shortcut": bool(getattr(self, "read_shortcut", False)),
-                "splice_scale_match": bool(getattr(self, "splice_scale_match", False)),
-                "read_gating": bool(getattr(self, "read_gating", False)),
-                "read_residual_pure": bool(getattr(self, "read_residual_pure", False)),
-                "final_verify": bool(getattr(self, "final_verify", False)),
-                "mid_read_step": bool(getattr(self, "mid_read_step", 0)),
-                "persistent_slot_id": bool(getattr(self, "persistent_slot_id", False)),
-                "interface_loss": bool(getattr(self, "interface_loss", False)),
-                "read_anchor": bool(getattr(self, "read_anchor", False)),
-                "functional_state_loss": bool(getattr(self, "functional_state_loss", False)),
-            }
-            active = [name for name, enabled in incompatible.items() if enabled]
-            if active:
-                raise ValueError(
-                    "counterfactual_residual_recurrence replaces the slot/raw "
-                    f"workspace and is incompatible with: {', '.join(active)}"
-                )
-        elif getattr(self, "crr_policy_rank", 0) > 0:
-            raise ValueError(
-                "crr_policy_rank > 0 requires "
-                "counterfactual_residual_recurrence=True."
-            )
-        if getattr(self, "visual_counterfactual_recurrence", False):
-            if not getattr(
-                self, "counterfactual_residual_recurrence", False
-            ):
-                raise ValueError(
-                    "visual_counterfactual_recurrence requires "
-                    "counterfactual_residual_recurrence=True."
-                )
-            # CVRR is the intentionally minimal CE-only candidate.  The old
-            # CRR tail-rollout losses/policy implement a different transition
-            # equation and must not be silently applied to this trajectory.
-            incompatible = {
-                "counterfactual_reliance_loss": bool(
-                    getattr(self, "counterfactual_reliance_loss", False)
-                ),
-                "transition_aware_reliance_loss": bool(
-                    getattr(self, "transition_aware_reliance_loss", False)
-                ),
-                "corrective_transition_reliance_loss": bool(
-                    getattr(self, "corrective_transition_reliance_loss", False)
-                ),
-                "counterfactual_step_retention_loss": bool(
-                    getattr(self, "counterfactual_step_retention_loss", False)
-                ),
-                "crr_concat_aggregation": bool(
-                    getattr(self, "crr_concat_aggregation", False)
-                ),
-                "crr_policy_rank": bool(getattr(self, "crr_policy_rank", 0)),
-            }
-            active = [name for name, enabled in incompatible.items() if enabled]
-            if active:
-                raise ValueError(
-                    "visual_counterfactual_recurrence uses final-answer CE only "
-                    "and a parameter-free decoder interface; incompatible with: "
-                    + ", ".join(active)
-                )
-        elif getattr(self, "visual_cumulative_recurrence", False):
-            raise ValueError(
-                "visual_cumulative_recurrence requires "
-                "visual_counterfactual_recurrence=True."
-            )
-        if getattr(self, "visual_full_state_recurrence", False):
-            if not getattr(self, "visual_counterfactual_recurrence", False):
-                raise ValueError(
-                    "visual_full_state_recurrence requires "
-                    "visual_counterfactual_recurrence=True."
-                )
-            if not getattr(self, "visual_cumulative_recurrence", False):
-                raise ValueError(
-                    "visual_full_state_recurrence requires "
-                    "visual_cumulative_recurrence=True."
-                )
-        if getattr(self, "visual_boundary_reentry", False) and not getattr(
-            self, "visual_full_state_recurrence", False
-        ):
-            raise ValueError(
-                "visual_boundary_reentry requires "
-                "visual_full_state_recurrence=True."
-            )
-        if getattr(self, "visual_source_centered_reentry", False):
-            if not getattr(self, "visual_boundary_reentry", False):
-                raise ValueError(
-                    "visual_source_centered_reentry requires "
-                    "visual_boundary_reentry=True."
-                )
-            if getattr(self, "visual_preserving_adapter_correction", False):
-                raise ValueError(
-                    "visual_source_centered_reentry cannot be combined with "
-                    "visual_preserving_adapter_correction."
-                )
-        if getattr(self, "recurrence_only_adapter", False):
-            if not getattr(self, "visual_full_state_recurrence", False):
-                raise ValueError(
-                    "recurrence_only_adapter requires "
-                    "visual_full_state_recurrence=True."
-                )
-            if self.adapter_rank <= 0:
-                raise ValueError(
-                    "recurrence_only_adapter requires adapter_rank > 0."
-                )
-            if self.ell_star is None:
-                raise ValueError(
-                    "recurrence_only_adapter requires an explicit ell_star."
-                )
-            recurrent_layer = int(self.ell_star) + 1
-            if recurrent_layer in set(self.adapter_exclude_layers):
-                raise ValueError(
-                    "recurrence_only_adapter cannot exclude its recurrent "
-                    f"layer {recurrent_layer} from LoRA."
-                )
-        if getattr(self, "visual_preserving_adapter_correction", False):
-            if not getattr(self, "visual_full_state_recurrence", False):
-                raise ValueError(
-                    "visual_preserving_adapter_correction requires "
-                    "visual_full_state_recurrence=True."
-                )
-            if not getattr(self, "recurrence_only_adapter", False):
-                raise ValueError(
-                    "visual_preserving_adapter_correction requires "
-                    "recurrence_only_adapter=True."
-                )
-        if getattr(self, "perceive_deliberate_chain", False):
-            if self.num_workspace_steps < 2 or self.num_workspace_steps % 2:
-                raise ValueError(
-                    "one-pass latent chains require an even "
-                    "num_workspace_steps >= 2."
-                )
-            vocab_size = int(self.text_config.vocab_size)
-            if not -1 <= self.latent_chain_init_token_id < vocab_size:
-                raise ValueError(
-                    "latent_chain_init_token_id must be -1 or a valid text "
-                    f"token id below {vocab_size}."
-                )
-            num_pairs = self.num_workspace_steps // 2
-            if self.local_visual_latent_chain and (
-                self.pdlc_transition_conditioned
-                or self.pdlc_question_curriculum
-                or self.pdlc_question_visible_through_pair != 1
-            ):
-                raise ValueError(
-                    "local_visual_latent_chain replaces PDLC transition and "
-                    "question-visibility controls"
-                )
-            if not 1 <= self.pdlc_question_visible_through_pair <= num_pairs:
-                raise ValueError(
-                    "pdlc_question_visible_through_pair must lie in "
-                    f"[1,{num_pairs}]"
-                )
-            if self.pdlc_question_curriculum and not (
-                getattr(self, "pdlc_transition_conditioned", False)
-            ):
-                raise ValueError(
-                    "pdlc_question_curriculum requires "
-                    "pdlc_transition_conditioned=True"
-                )
-            incompatible = {
-                "raw_evidence_question": bool(
-                    getattr(self, "raw_evidence_question", False)
-                ),
-                "counterfactual_residual_recurrence": bool(
-                    getattr(self, "counterfactual_residual_recurrence", False)
-                ),
-                "spatial_visual_recurrence": bool(
-                    getattr(self, "spatial_visual_recurrence", False)
-                ),
-                "evidence_reasoning": bool(
-                    getattr(self, "evidence_reasoning", False)
-                ),
-                "read_gating": bool(getattr(self, "read_gating", False)),
-                "read_residual_pure": bool(
-                    getattr(self, "read_residual_pure", False)
-                ),
-                "final_verify": bool(getattr(self, "final_verify", False)),
-                "mid_read_step": bool(getattr(self, "mid_read_step", 0)),
-                "interface_loss": bool(getattr(self, "interface_loss", False)),
-                "read_anchor": bool(getattr(self, "read_anchor", False)),
-                "counterfactual_reliance_loss": bool(
-                    getattr(self, "counterfactual_reliance_loss", False)
-                ),
-                "counterfactual_step_retention_loss": bool(
-                    getattr(self, "counterfactual_step_retention_loss", False)
-                ),
-                "functional_state_loss": bool(
-                    getattr(self, "functional_state_loss", False)
-                ),
-            }
-            active = [name for name, enabled in incompatible.items() if enabled]
-            if active:
-                raise ValueError(
-                    "perceive_deliberate_chain replaces every recurrent/"
-                    "workspace path; incompatible with: " + ", ".join(active)
-                )
-        elif getattr(self, "pdlc_transition_conditioned", False):
-            raise ValueError(
-                "pdlc_transition_conditioned requires "
-                "perceive_deliberate_chain=True."
-            )
-        elif getattr(self, "pdlc_question_curriculum", False) or (
-            getattr(self, "pdlc_question_visible_through_pair", 1) != 1
-        ):
-            raise ValueError(
-                "PDLC question visibility controls require "
-                "perceive_deliberate_chain=True."
-            )
-        if getattr(self, "spatial_visual_recurrence", False):
-            if not 0.0 <= self.counterfactual_beta <= 1.0:
-                raise ValueError(
-                    "counterfactual_beta must lie in [0, 1], got "
-                    f"{self.counterfactual_beta}."
-                )
-            if self.spatial_recurrence_steps < 1:
-                raise ValueError(
-                    "spatial_recurrence_steps must be >= 1 for training"
-                )
-            if self.spatial_recurrence_inner_width < 1:
-                raise ValueError(
-                    "spatial_recurrence_inner_width must be positive"
-                )
-            if self.spatial_recurrence_heads < 1 or (
-                self.spatial_recurrence_inner_width
-                % self.spatial_recurrence_heads
-            ):
-                raise ValueError(
-                    "spatial_recurrence_inner_width must be divisible by "
-                    "spatial_recurrence_heads"
-                )
-            if self.spatial_step_override < -1:
-                raise ValueError("spatial_step_override must be >= -1")
-            if self.adapter_rank != 0:
-                raise ValueError(
-                    "spatial_visual_recurrence replaces decoder LoRA; set "
-                    "adapter_rank=0 so the T=0 path remains the frozen base model"
-                )
-            incompatible = {
-                "raw_evidence_question": bool(
-                    getattr(self, "raw_evidence_question", False)
-                ),
-                "counterfactual_residual_recurrence": bool(
-                    getattr(self, "counterfactual_residual_recurrence", False)
-                ),
-                "evidence_reasoning": bool(
-                    getattr(self, "evidence_reasoning", False)
-                ),
-                "read_gating": bool(getattr(self, "read_gating", False)),
-                "read_residual_pure": bool(
-                    getattr(self, "read_residual_pure", False)
-                ),
-                "final_verify": bool(getattr(self, "final_verify", False)),
-                "mid_read_step": bool(getattr(self, "mid_read_step", 0)),
-                "interface_loss": bool(getattr(self, "interface_loss", False)),
-                "read_anchor": bool(getattr(self, "read_anchor", False)),
-                "counterfactual_reliance_loss": bool(
-                    getattr(self, "counterfactual_reliance_loss", False)
-                ),
-                "counterfactual_step_retention_loss": bool(
-                    getattr(self, "counterfactual_step_retention_loss", False)
-                ),
-                "functional_state_loss": bool(
-                    getattr(self, "functional_state_loss", False)
-                ),
-            }
-            active = [name for name, enabled in incompatible.items() if enabled]
-            if active:
-                raise ValueError(
-                    "spatial_visual_recurrence replaces every historical "
-                    "workspace/reliance path; incompatible with: "
-                    + ", ".join(active)
-                )
-        modes = [bool(getattr(self, m, False)) for m in
-                 ("read_gating", "read_residual_pure", "final_verify")]
-        modes.append(bool(getattr(self, "mid_read_step", 0)))
-        if sum(modes) > 1:
-            raise ValueError(
-                "read_gating / read_residual_pure / final_verify are mutually "
-                "exclusive reread modes."
-            )
-
-
-__all__ = ["CloseQwen2_5_VLConfig", "WORKSPACE_ROPE_MODES"]

configuration_source_qwen3.py

@@ -1,740 +0,0 @@
-"""Configuration for CLOSE latent reasoning on Qwen3-VL.
-
-The current ``raw_evidence_question`` path keeps the complete layer-``ell_star``
-visual memory as immutable evidence and recurrently updates the text-only
-question state by cross-attending to it. Its selected interface averages all
-four residual recurrent states and may use a training-only restoration target
-from the normal multimodal question state. Historical slot-workspace fields are
-retained so old checkpoints remain loadable, but they are mutually exclusive
-with the raw-evidence path.
-
-Every architectural hyperparameter of Method §3 lives here -- nothing is
-hardcoded in ``modeling_close_qwen3_vl.py``.
-
-Layer-index convention (used consistently across the repo)
-----------------------------------------------------------
-``ell_star`` is **0-indexed and inclusive**: it is the index of the last decoder
-layer belonging to the lower branch. With ``num_hidden_layers == 28``::
-
-    F_{<=l*} = language_model.layers[: ell_star + 1]     # lower_slice
-    F_{>l*}  = language_model.layers[ell_star + 1 :]     # upper_slice
-
-so ``ell_star=18`` means layers 0..18 run before the split and 19..27 after it.
-``ell_star`` is measured once by ``scripts/localize_read_point.py`` (§3.1) and
-frozen before any workspace training.
-
-Careful with ``Qwen3VLConfig``: its ``__setattr__`` forwards any attribute
-whose name already exists in ``text_config.__dict__`` down to the sub-config, and
-its ``__init__`` builds ``text_config`` from ``**kwargs``. Workspace fields are
-therefore declared as explicit named parameters so they never enter ``kwargs``
-and never shadow a text-config field. ``tests/test_config_roundtrip.py`` guards
-this.
-"""
-
-from __future__ import annotations
-
-import math
-
-from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig
-
-#: How the ``S`` workspace slots are assigned M-RoPE positions once they are
-#: spliced into the replacement cache above ``ell_star``. Qwen3-VL uses 3D
-#: M-RoPE (t, h, w) with ``mrope_section=[16, 24, 24]``, so slots -- which have
-#: no spatial extent -- need an explicit convention.
-WORKSPACE_ROPE_MODES = (
-    # Slots continue the 1-D text position sequence directly after Q*, with
-    # t == h == w. Treats the workspace as "more text".
-    "continue",
-    # Slots reuse the (t, h, w) span the image tokens occupied in the original
-    # multimodal prefill, subsampled to S positions. Preserves the read point's
-    # positional geometry but reintroduces image-derived indices.
-    "image_span",
-    # All slots share one position (the first position after Q*). Makes the
-    # workspace order-free, matching the set semantics of Perceiver-style slots.
-    "shared",
-)
-
-
-class CloseQwen3VLConfig(Qwen3VLConfig):
-    r"""Configuration for :class:`CloseQwen3VLForConditionalGeneration`.
-
-    Args:
-        ell_star: 0-indexed inclusive last layer of the lower branch. ``None``
-            until §3.1 localization has run; the model refuses to build with
-            ``None`` rather than falling back to a guessed depth.
-        num_workspace_steps: the fixed recurrent horizon (paper notation ``T``;
-            per-item segment counts ``n`` exist only in the losses, never here).
-            The trainable parameter count is independent of it because
-            ``f_theta`` is shared across steps, so it is a free experimental
-            axis. ``0`` is legal and is the no-recurrence floor: ``Z^(0)`` goes
-            straight to the decoder.
-        num_workspace_slots: ``S``, slots in ``Z``. Free capacity knob.
-        workspace_width: ``d_w``, the slot width. Free capacity knob -- the
-            method's claim is path exclusivity (all visual information routed
-            through ``Z``), NOT an information bottleneck, so ``S`` and ``d_w``
-            carry no methodological commitment. Kept below ``d`` (4096 for the
-            7B) since ``P_Z`` projects up at the splice.
-        workspace_num_heads: attention heads inside ``r_theta`` / ``f_theta``.
-        read_num_blocks: cross-attention blocks in ``r_theta``. The *read* is
-            once regardless of depth: all blocks see the same ``V*`` in a single
-            forward, and ``V*`` is discarded afterwards.
-        transition_num_blocks: blocks in the shared recurrent ``f_theta``.
-        workspace_ffn_mult: FFN expansion inside workspace blocks.
-        workspace_dropout: dropout inside workspace blocks.
-        workspace_rope_mode: one of :data:`WORKSPACE_ROPE_MODES`.
-        adapter_rank: rank of the low-rank adapter on the post-``ell_star``
-            backbone layers. ``0`` disables it (frozen backbone only).
-        adapter_alpha: LoRA-style scaling; effective scale is
-            ``adapter_alpha / adapter_rank``.
-        adapter_dropout: dropout probability applied only to the LoRA branch.
-        lambda_traj: legacy trajectory-loss weight. The current method sets it
-            to zero and trains with final-answer CE only; it remains serialized
-            solely so historical checkpoints can still be inspected.
-    """
-
-    model_type = "close_qwen3_vl"
-
-    def __init__(
-        self,
-        ell_star: int | None = None,
-        num_workspace_steps: int = 4,
-        num_workspace_slots: int = 32,
-        workspace_width: int = 512,
-        workspace_num_heads: int = 8,
-        read_num_blocks: int = 1,
-        transition_num_blocks: int = 2,
-        workspace_ffn_mult: int = 4,
-        workspace_dropout: float = 0.0,
-        workspace_rope_mode: str = "continue",
-        adapter_rank: int = 16,
-        adapter_alpha: int = 32,
-        adapter_dropout: float = 0.0,
-        adapter_exclude_layers: tuple = (),
-        lambda_traj: float = 0.0,
-        read_shortcut: bool = False,
-        splice_scale_match: bool = False,
-        read_gating: bool = False,
-        gate_init_eps: float = 0.02,
-        read_residual_pure: bool = False,
-        residual_norm_cap: float = 0.0,
-        pure_warm_start: bool = False,
-        final_verify: bool = False,
-        mid_read_step: int = 0,
-        persistent_slot_id: bool = False,
-        mid_read_alpha: float = 1.0,
-        evidence_reasoning: bool = False,
-        raw_evidence_question: bool = False,
-        raw_attention_source_layer: int = -1,
-        raw_transition_ffn: bool | None = None,
-        raw_state_replacement: bool = False,
-        raw_question_anchor: bool = False,
-        raw_mean_aggregation: bool = False,
-        counterfactual_residual_recurrence: bool = False,
-        visual_counterfactual_recurrence: bool = False,
-        visual_cumulative_recurrence: bool = False,
-        visual_full_state_recurrence: bool = False,
-        recurrence_only_adapter: bool = False,
-        counterfactual_beta: float = 1.0,
-        crr_concat_aggregation: bool = False,
-        crr_policy_rank: int = 0,
-        crr_policy_alpha: float = 32.0,
-        counterfactual_reliance_loss: bool = False,
-        transition_aware_reliance_loss: bool = False,
-        counterfactual_step_retention_loss: bool = False,
-        functional_state_loss: bool = False,
-        per_example_answer_loss: bool = False,
-        interface_loss: bool = False,
-        read_anchor: bool = False,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.ell_star = ell_star
-        self.num_workspace_steps = num_workspace_steps
-        self.num_workspace_slots = num_workspace_slots
-        self.workspace_width = workspace_width
-        self.workspace_num_heads = workspace_num_heads
-        self.read_num_blocks = read_num_blocks
-        self.transition_num_blocks = transition_num_blocks
-        self.workspace_ffn_mult = workspace_ffn_mult
-        self.workspace_dropout = workspace_dropout
-        self.workspace_rope_mode = workspace_rope_mode
-        self.adapter_rank = adapter_rank
-        self.adapter_alpha = adapter_alpha
-        self.adapter_dropout = adapter_dropout
-        # E run: dense full-FT 층은 LoRA 대상에서 제외 (이중 파라미터화 방지)
-        self.adapter_exclude_layers = list(adapter_exclude_layers or ())
-        self.lambda_traj = lambda_traj
-        # v2 anti-collapse switches (defaults False so v1 checkpoints reload).
-        # read_shortcut: r_theta adds an orthogonally-projected pooled-V* term to
-        # every slot, so Z^(0) is image-dependent from step 0 instead of relying
-        # on a random cross-attention finding signal.
-        # splice_scale_match: P_Z(Z) rows are RMS-matched to the question rows at
-        # the splice. Unmatched, slot states are ~O(1) against layer-20 hiddens
-        # in the hundreds -- the decoder could ignore them for free, and did
-        # (probe 20143: visual sensitivity 0.94 -> 1.0000 over training).
-        self.read_shortcut = read_shortcut
-        self.splice_scale_match = splice_scale_match
-        # v6 candidate: gated re-reads DURING the recurrence. Every re-read
-        # goes through r_theta into Z, so the mediation invariant (decoder sees
-        # only [Q*; P_Z(Z)]) is untouched; what changes is that the workspace
-        # may consult V* again mid-reasoning, scaled by a learned scalar gate.
-        # The gate is free at EVERY step -- no hand-coded schedule. Degeneration
-        # into answer-time retrieval is disciplined by the objective itself:
-        # L_traj pins each intermediate state to its teacher segment, so a
-        # wait-then-fetch trajectory misses its intermediate targets. The gate
-        # bias initialises negative, i.e. training starts from the proven
-        # no-reread regime and departs only where gradients demand it.
-        # Default False: v4/v5 checkpoints reload bit-identically.
-        self.read_gating = read_gating
-        # Initial gate open-rate epsilon_gamma: b_gamma = logit(eps). Sigmoid
-        # never reaches exactly 0 at finite logits, so this is near-identity
-        # (not bitwise) initialisation w.r.t. the read-once path.
-        self.gate_init_eps = gate_init_eps
-        # v7: always-on PURE visual residual rereads. Every step adds
-        # Z = U + W_o XAttn(LN(U) -> queries, V* -> keys/values): no self-attn,
-        # no FFN, no output norm, no biases in the branch, W_o zero-init.
-        # Motivation (battery, ep1 ckpt 20752): read_once hurt badly while
-        # null_step (blank memory at one step) was harmless -- the gated
-        # reader's value was extra recurrent DEPTH, not vision. This branch
-        # removes every path that can change Z without V* content, so any
-        # benefit it shows IS visual by construction.
-        self.read_residual_pure = read_residual_pure
-        # v7.1 (advisor): relative-norm trust region on the visual residual.
-        #   s_k = min(1, rho * RMS(U) / (RMS(Delta) + eps));  Z = U + s_k * Delta
-        # Unlike a fixed alpha, the optimizer cannot undo this by rescaling
-        # W_o (s_k adapts); ||Delta_used|| <~ rho * ||U|| structurally, and
-        # V* = 0 => Delta = 0 keeps the zero-memory identity. 0.0 disables.
-        self.residual_norm_cap = residual_norm_cap
-        # v7.1: initialise pure_attn K/V projections from the initial visual
-        # reader's v_in (same [d_w, d] shape) and Q from its first read
-        # block's in-proj slice, instead of random xavier -- the ep1 finding
-        # was that Q/K/V stayed ~random while only W_o trained.
-        self.pure_warm_start = pure_warm_start
-        # MAIN METHOD (advisor-confirmed): Final-Verified Recurrent Workspace.
-        # Initial read -> PURE latent recurrence (no mid-step visual anything)
-        # -> exactly ONE pure visual verification at the end:
-        #   Z_out = U^(T) + W_o XAttn(LN(U^(T)), V*),  alpha = 1, W_o = 0 init.
-        # Uses the same bias-free verifier module as read_residual_pure; the
-        # two flags differ only in WHERE the branch fires (every step vs once).
-        self.final_verify = final_verify
-        # MAIN (advisor 최종 확정): Mid-Read Recurrent Workspace. 고정 상수
-        # k* = floor(T/2) = 4 (스텝 sweep 산물이 아니라 판독 전후 동수 전이의
-        # midpoint)에서 pure visual read 1회:
-        #   Z^(4) = U^(4) + W_o XAttn(LN(U^(4)), V*),  alpha=1, W_o=0 init.
-        # 논리적 step 4 == python 루프 index 3 (0-based) -- off-by-one은
-        # tests/test_mid_read.py가 상태 분기 위치로 고정한다. 0이면 비활성.
-        self.mid_read_step = mid_read_step
-        # advisor(2026-08-06): E_slot(=workspace_read.slots)을 매 recurrent
-        # step 입력과 mid-read query에 일시 제공. state에는 누적하지 않는다.
-        self.persistent_slot_id = persistent_slot_id
-        # 고정 read 스케일 Z = U + alpha*W_o A (advisor scale-mismatch 검사).
-        # 1.0 = 기존과 비트동일. eval의 gate_overrides alpha는 절대값으로 이를 대체.
-        self.mid_read_alpha = mid_read_alpha
-        # advisor 최종안: E(persistent evidence) / R(recurrent reasoning) 분리.
-        # states = [E^(0), R^(1..T)], 결합 Z^(T)=E+R^(T)는 prefill에서.
-        self.evidence_reasoning = evidence_reasoning
-        # Current replacement candidate: no learned evidence/reasoning slots.
-        # E is the complete raw V* sequence and R0 is the text-only Q* sequence.
-        # A single shared CrossAttn transition produces R1..RT. Residual mode
-        # decodes RT; the replacement candidate decodes Q*+HT. Cross-attention
-        # uses the backbone's native GQA shape and
-        # is initialized from the first upper layer by default. Checkpoints
-        # created before the CrossAttn-only decision did not serialize
-        # raw_transition_ffn; None therefore preserves their historical FFN
-        # for faithful evaluation, while every new training launch passes False.
-        self.raw_evidence_question = bool(raw_evidence_question)
-        self.raw_attention_source_layer = int(raw_attention_source_layer)
-        self.raw_transition_ffn = (
-            True if raw_transition_ffn is None else bool(raw_transition_ffn)
-        )
-        # Experimental simplification of the same raw-E/Q path.  The residual
-        # carrier is replaced, not augmented:
-        #   H0=Q*, Hk=CrossAttn(H{k-1}, E), decoder_latent=Q*+HT.
-        # False preserves all residual raw-E/Q checkpoints and active runs.
-        self.raw_state_replacement = bool(raw_state_replacement)
-        # Question-anchored replacement keeps the fixed text scaffold in every
-        # recurrent query without restoring the old recurrent identity path:
-        #   Z0=Q*, Zk=Q*+CrossAttn(Z{k-1}, E), decoder_latent=ZT.
-        # This replaces the one-time Q*+HT decoder interface above.
-        self.raw_question_anchor = bool(raw_question_anchor)
-        # Q-shaped recurrent candidate selected by the V* interface oracle:
-        #   R0=Q*, Rk=R{k-1}+CrossAttn(R{k-1}, V*)
-        #   R_agg=mean(R1..RT)
-        # This is parameter-free and replaces both question-anchored state
-        # replacement and final-state-only decoding.
-        self.raw_mean_aggregation = bool(raw_mean_aggregation)
-        # Counterfactual Residual Recurrence (CRR): the image-conditioned and
-        # text-only question states are propagated by the same native layer
-        # F=L(ell_star+1).  Only their counterfactual residual is recurrently
-        # refined; there are no slots, V* memory, custom attention modules, or
-        # inference-time auxiliary modules.  beta is the fixed relaxed-update
-        # coefficient; an optional batch-swap reliance loss is training-only.
-        self.counterfactual_residual_recurrence = bool(
-            counterfactual_residual_recurrence
-        )
-        # Qwen3 counterpart of the promoted Qwen2.5 full-state CRR path.  The
-        # multimodal scaffold is used only inside the shared recurrent cell;
-        # answer decoding still receives a question-shaped state and a
-        # text-only prefix cache.
-        self.visual_counterfactual_recurrence = bool(
-            visual_counterfactual_recurrence
-        )
-        self.visual_cumulative_recurrence = bool(
-            visual_cumulative_recurrence
-        )
-        self.visual_full_state_recurrence = bool(
-            visual_full_state_recurrence
-        )
-        self.recurrence_only_adapter = bool(recurrence_only_adapter)
-        self.counterfactual_beta = float(counterfactual_beta)
-        # Bias-free Concat(C1..CT) -> C_agg, initialized to the exact
-        # historical mean interface for checkpoint-compatible warm starts.
-        self.crr_concat_aggregation = bool(crr_concat_aggregation)
-        # Optional latent-only policy adapter.  It corrects the deterministic
-        # C2..CT means but never touches B, C1, or answer-token decoding.  Rank 0
-        # is the exact historical CRR path; RL checkpoints serialize rank > 0.
-        self.crr_policy_rank = int(crr_policy_rank)
-        self.crr_policy_alpha = float(crr_policy_alpha)
-        # Training-only CRR intervention.  A batch-matched residual from a
-        # different-answer example replaces the factual residual before the
-        # upper decoder.  The model returns the factual-minus-swapped answer
-        # score gap; the Trainer applies the margin loss.  Evaluation and
-        # generation never execute the extra branch.
-        self.counterfactual_reliance_loss = bool(
-            counterfactual_reliance_loss
-        )
-        self.transition_aware_reliance_loss = bool(
-            transition_aware_reliance_loss
-        )
-        # Training-only functional no-regression objective.  Gold-answer NLL
-        # is measured through the actual upper decoder at every recurrent CRR
-        # prefix; inference remains the unchanged T-step mean interface.
-        self.counterfactual_step_retention_loss = bool(
-            counterfactual_step_retention_loss
-        )
-        # Training-only self-distillation target.  The multimodal lower pass
-        # already computed for V* also provides its image-conditioned question
-        # rows.  They supervise R_agg at the split boundary, but are never
-        # exposed to the inference decoder.
-        self.functional_state_loss = bool(functional_state_loss)
-        # Token CE lets long caption answers dominate one-token visual-choice
-        # examples.  The per-example variant first averages valid answer tokens
-        # within each sample, then averages samples.  It changes training/eval
-        # loss reduction only and has no inference-time effect.
-        self.per_example_answer_loss = bool(per_example_answer_loss)
-        # ER v2 (advisor 2026-08-11): explicit evidence read —
-        # R^(k) = f_theta(R^(k-1), Q*; E), E는 별도 memory로 cross-attn 읽기.
-        # (E+R 합산-감산 제거. evidence_reasoning=True 전제.)
-        self.explicit_evidence_read = bool(kwargs.pop("explicit_evidence_read", False))
-        # Historical field name retained for checkpoint compatibility. In ER
-        # v2.2 this means R^(0) is a batch-broadcast set of learnable slots; it
-        # is deliberately NOT initialized from Q*. Q* conditions every shared
-        # transition instead.
-        self.r_init_question = bool(kwargs.pop("r_init_question", False))
-        # ER v2.1 (advisor 2026-08-11): evidence read의 competitive normalization —
-        # slot 축 softmax 후 evidence 축 재정규화 (Slot Attention식 경쟁).
-        # 파라미터/모듈/loss 무추가; evidence_read의 W_qkv/W_o 재사용.
-        self.competitive_evidence_read = bool(
-            kwargs.pop("competitive_evidence_read", False))
-        # ER v2.2 C: k>=1 traj 감독 부재 시 twin 체인 제거 (states = answer 체인).
-        self.er_single_chain = bool(kwargs.pop("er_single_chain", False))
-        # ER v2.2 C (§5-7): decoder latent = softmax(temporal_logits)로 가중한
-        # R1..RT 결합 (전역 스칼라 T개). er_single_chain 전제. (legacy)
-        self.temporal_aggregation = bool(kwargs.pop("temporal_aggregation", False))
-        # ER v2.2 C 확정판: R_agg = Linear(Concat(R1..RT), bias=False),
-        # 초기 W=[I/T,...,I/T] (R_agg=mean(R1..RT)). er_single_chain 전제.
-        self.concat_aggregation = bool(kwargs.pop("concat_aggregation", False))
-        # ER v2: 워크스페이스 블록의 rank-r 부분공간 attention/FFN (0 = dense).
-        # d_w=d에서 state width 유지와 파라미터 폭증을 분리 (advisor).
-        self.workspace_low_rank = int(kwargs.pop("workspace_low_rank", 0))
-        # advisor: single-layer text-side interface alignment (bridge-only).
-        # teacher = frozen-base mm 경로 q_end@l*+1 (adapter off, sg);
-        # student = detach(R^T) splice의 l*+1 한 층 재계산. 추론 시 무존재.
-        self.interface_loss = interface_loss
-        # Training-only visual anchor switch. False in the current answer-only
-        # method; when false the model does not even materialize pooled V* for
-        # the Trainer.
-        self.read_anchor = bool(read_anchor)
-        if evidence_reasoning and persistent_slot_id:
-            raise ValueError("evidence_reasoning은 persistent_slot_id와 동시 사용 불가")
-        if evidence_reasoning and not mid_read_step:
-            raise ValueError("evidence_reasoning은 mid_read_step > 0 필요")
-        self.validate()
-
-    # -- derived views -----------------------------------------------------
-
-    @property
-    def num_decoder_layers(self) -> int:
-        return self.text_config.num_hidden_layers
-
-    @property
-    def backbone_width(self) -> int:
-        """``d`` -- backbone hidden width (4096 for the 7B)."""
-        return self.text_config.hidden_size
-
-    @property
-    def lower_slice(self) -> slice:
-        """Layers forming ``F_{<=l*}``."""
-        self._require_ell_star()
-        return slice(0, self.ell_star + 1)
-
-    @property
-    def upper_slice(self) -> slice:
-        """Layers forming ``F_{>l*}``."""
-        self._require_ell_star()
-        return slice(self.ell_star + 1, self.num_decoder_layers)
-
-    # -- validation --------------------------------------------------------
-
-    def _require_ell_star(self) -> None:
-        if self.ell_star is None:
-            raise ValueError(
-                "ell_star is unset. Run scripts/localize_read_point.py (Method 3.1) "
-                "and pass the measured layer explicitly; it must not be guessed."
-            )
-
-    def validate(self) -> None:
-        """Reject configurations that silently break the §3.2 contract."""
-        if self.ell_star is not None:
-            n = self.num_decoder_layers
-            # Both branches must be non-empty: an empty lower branch means there
-            # is no V* to read, an empty upper branch means the workspace never
-            # reaches the decoder.
-            if not 0 <= self.ell_star <= n - 2:
-                raise ValueError(
-                    f"ell_star={self.ell_star} out of range for {n} decoder layers; "
-                    f"expected 0 <= ell_star <= {n - 2} so both branches are non-empty."
-                )
-        # K=0 is the no-recurrence control (read only), not a misconfiguration.
-        if self.num_workspace_steps < 0:
-            raise ValueError("num_workspace_steps (K) must be >= 0.")
-        if self.num_workspace_slots < 1:
-            raise ValueError("num_workspace_slots (S) must be >= 1.")
-        if self.workspace_num_heads < 1:
-            raise ValueError("workspace_num_heads must be >= 1.")
-        if self.workspace_width % self.workspace_num_heads != 0:
-            raise ValueError(
-                f"workspace_width={self.workspace_width} must be divisible by "
-                f"workspace_num_heads={self.workspace_num_heads}."
-            )
-        if self.workspace_width > self.backbone_width:
-            raise ValueError(
-                f"workspace_width (d_w={self.workspace_width}) cannot exceed "
-                f"backbone width (d={self.backbone_width}). Native-width d_w=d is "
-                "the current method and uses an Identity decoder interface."
-            )
-        if self.workspace_rope_mode not in WORKSPACE_ROPE_MODES:
-            raise ValueError(
-                f"workspace_rope_mode={self.workspace_rope_mode!r} not in "
-                f"{WORKSPACE_ROPE_MODES}."
-            )
-        if self.adapter_rank < 0:
-            raise ValueError("adapter_rank must be >= 0 (0 disables the adapter).")
-        if not math.isfinite(self.adapter_dropout) or not 0.0 <= self.adapter_dropout < 1.0:
-            raise ValueError("adapter_dropout must be finite and in [0, 1).")
-        if getattr(self, "crr_policy_rank", 0) < 0:
-            raise ValueError("crr_policy_rank must be >= 0.")
-        if not math.isfinite(getattr(self, "crr_policy_alpha", 0.0)) or getattr(
-            self, "crr_policy_alpha", 0.0
-        ) <= 0.0:
-            raise ValueError("crr_policy_alpha must be finite and > 0.")
-        if self.read_num_blocks < 1 or self.transition_num_blocks < 1:
-            raise ValueError("read_num_blocks and transition_num_blocks must be >= 1.")
-        if self.workspace_ffn_mult < 1:
-            raise ValueError("workspace_ffn_mult must be >= 1.")
-        mid = int(getattr(self, "mid_read_step", 0) or 0)
-        if mid < 0 or mid > self.num_workspace_steps:
-            raise ValueError(
-                f"mid_read_step={mid} must be in [0, num_workspace_steps="
-                f"{self.num_workspace_steps}]."
-            )
-        low_rank = int(getattr(self, "workspace_low_rank", 0) or 0)
-        if low_rank < 0:
-            raise ValueError("workspace_low_rank must be >= 0.")
-        if low_rank and low_rank % self.workspace_num_heads != 0:
-            raise ValueError(
-                f"workspace_low_rank={low_rank} must be divisible by "
-                f"workspace_num_heads={self.workspace_num_heads}."
-            )
-        if getattr(self, "competitive_evidence_read", False) and not getattr(
-            self, "explicit_evidence_read", False
-        ):
-            raise ValueError(
-                "competitive_evidence_read requires explicit_evidence_read=True; "
-                "otherwise the flag has no execution path."
-            )
-        if getattr(self, "explicit_evidence_read", False) and not getattr(
-            self, "evidence_reasoning", False
-        ):
-            raise ValueError(
-                "explicit_evidence_read requires evidence_reasoning=True; otherwise "
-                "the evidence-read module is allocated but never called."
-            )
-        if getattr(self, "r_init_question", False) and not getattr(
-            self, "evidence_reasoning", False
-        ):
-            raise ValueError(
-                "r_init_question (legacy name for learnable R0 slots) requires "
-                "evidence_reasoning=True; otherwise reasoning_slots are unused."
-            )
-        if (getattr(self, "concat_aggregation", False)
-                or getattr(self, "temporal_aggregation", False)):
-            if self.num_workspace_steps < 1:
-                raise ValueError("reasoning aggregation requires num_workspace_steps >= 1.")
-            if not getattr(self, "evidence_reasoning", False):
-                raise ValueError(
-                    "reasoning aggregation requires evidence_reasoning=True; "
-                    "otherwise the aggregation module is unused."
-                )
-        if getattr(self, "concat_aggregation", False) and not getattr(
-            self, "er_single_chain", False
-        ):
-            raise ValueError("concat_aggregation requires er_single_chain=True.")
-        if getattr(self, "raw_evidence_question", False):
-            if self.workspace_width != self.backbone_width:
-                raise ValueError(
-                    "raw_evidence_question requires native workspace_width == "
-                    "backbone_width: E=V* and R0=Q* are used without projections."
-                )
-            if self.num_workspace_steps < 1:
-                raise ValueError(
-                    "raw_evidence_question requires num_workspace_steps >= 1."
-                )
-            source_layer = (
-                self.ell_star + 1
-                if self.raw_attention_source_layer < 0
-                else self.raw_attention_source_layer
-            )
-            if not self.ell_star < source_layer < self.text_config.num_hidden_layers:
-                raise ValueError(
-                    "raw_attention_source_layer must be an upper-backbone layer, "
-                    f"got ell_star={self.ell_star}, source={source_layer}."
-                )
-            incompatible = {
-                "evidence_reasoning": bool(getattr(self, "evidence_reasoning", False)),
-                "explicit_evidence_read": bool(getattr(self, "explicit_evidence_read", False)),
-                "competitive_evidence_read": bool(getattr(self, "competitive_evidence_read", False)),
-                "r_init_question": bool(getattr(self, "r_init_question", False)),
-                "er_single_chain": bool(getattr(self, "er_single_chain", False)),
-                "temporal_aggregation": bool(getattr(self, "temporal_aggregation", False)),
-                "concat_aggregation": bool(getattr(self, "concat_aggregation", False)),
-                "read_shortcut": bool(getattr(self, "read_shortcut", False)),
-                "splice_scale_match": bool(getattr(self, "splice_scale_match", False)),
-                "read_gating": bool(getattr(self, "read_gating", False)),
-                "read_residual_pure": bool(getattr(self, "read_residual_pure", False)),
-                "final_verify": bool(getattr(self, "final_verify", False)),
-                "mid_read_step": bool(getattr(self, "mid_read_step", 0)),
-                "persistent_slot_id": bool(getattr(self, "persistent_slot_id", False)),
-                "interface_loss": bool(getattr(self, "interface_loss", False)),
-                "read_anchor": bool(getattr(self, "read_anchor", False)),
-            }
-            active = [name for name, enabled in incompatible.items() if enabled]
-            if active:
-                raise ValueError(
-                    "raw_evidence_question replaces the slot workspace and is "
-                    f"incompatible with: {', '.join(active)}"
-                )
-            if self.raw_state_replacement and self.raw_transition_ffn:
-                raise ValueError(
-                    "raw_state_replacement is CrossAttn-only and is incompatible "
-                    "with raw_transition_ffn=True."
-                )
-            if self.raw_question_anchor and not self.raw_state_replacement:
-                raise ValueError(
-                    "raw_question_anchor requires raw_state_replacement=True."
-                )
-            if self.raw_mean_aggregation and (
-                self.raw_state_replacement or self.raw_question_anchor
-            ):
-                raise ValueError(
-                    "raw_mean_aggregation replaces raw_state_replacement and "
-                    "raw_question_anchor; use the residual recurrence."
-                )
-            if self.functional_state_loss and not self.raw_mean_aggregation:
-                raise ValueError(
-                    "functional_state_loss requires raw_mean_aggregation=True "
-                    "so its target is the exact decoder-facing recurrence."
-                )
-        elif getattr(self, "raw_state_replacement", False):
-            raise ValueError(
-                "raw_state_replacement requires raw_evidence_question=True."
-            )
-        elif getattr(self, "raw_question_anchor", False):
-            raise ValueError(
-                "raw_question_anchor requires raw_evidence_question=True."
-            )
-        elif getattr(self, "raw_mean_aggregation", False):
-            raise ValueError(
-                "raw_mean_aggregation requires raw_evidence_question=True."
-            )
-        elif getattr(self, "functional_state_loss", False):
-            raise ValueError(
-                "functional_state_loss requires raw_evidence_question=True."
-            )
-        if (
-            getattr(self, "counterfactual_reliance_loss", False)
-            and not getattr(self, "counterfactual_residual_recurrence", False)
-        ):
-            raise ValueError(
-                "counterfactual_reliance_loss requires "
-                "counterfactual_residual_recurrence=True."
-            )
-        if (
-            getattr(self, "transition_aware_reliance_loss", False)
-            and not getattr(self, "counterfactual_reliance_loss", False)
-        ):
-            raise ValueError(
-                "transition_aware_reliance_loss requires "
-                "counterfactual_reliance_loss=True."
-            )
-        if (
-            getattr(self, "transition_aware_reliance_loss", False)
-            and self.num_workspace_steps < 2
-        ):
-            raise ValueError(
-                "transition_aware_reliance_loss requires at least two steps."
-            )
-        if (
-            getattr(self, "counterfactual_step_retention_loss", False)
-            and not getattr(self, "counterfactual_residual_recurrence", False)
-        ):
-            raise ValueError(
-                "counterfactual_step_retention_loss requires "
-                "counterfactual_residual_recurrence=True."
-            )
-        if (
-            getattr(self, "counterfactual_step_retention_loss", False)
-            and self.num_workspace_steps < 2
-        ):
-            raise ValueError(
-                "counterfactual_step_retention_loss requires at least two steps."
-            )
-        if getattr(self, "crr_concat_aggregation", False) and not getattr(
-            self, "counterfactual_residual_recurrence", False
-        ):
-            raise ValueError(
-                "crr_concat_aggregation requires "
-                "counterfactual_residual_recurrence=True."
-            )
-        if getattr(self, "counterfactual_residual_recurrence", False):
-            if getattr(self, "raw_evidence_question", False):
-                raise ValueError(
-                    "counterfactual_residual_recurrence replaces "
-                    "raw_evidence_question; enable exactly one q-state method."
-                )
-            if self.num_workspace_steps < 1:
-                raise ValueError(
-                    "counterfactual_residual_recurrence requires "
-                    "num_workspace_steps >= 1."
-                )
-            if self.crr_policy_rank > 0 and self.num_workspace_steps < 2:
-                raise ValueError(
-                    "crr_policy_rank > 0 requires num_workspace_steps >= 2."
-                )
-            if not 0.0 <= self.counterfactual_beta <= 1.0:
-                raise ValueError(
-                    "counterfactual_beta must lie in [0, 1], got "
-                    f"{self.counterfactual_beta}."
-                )
-            incompatible = {
-                "evidence_reasoning": bool(getattr(self, "evidence_reasoning", False)),
-                "explicit_evidence_read": bool(getattr(self, "explicit_evidence_read", False)),
-                "competitive_evidence_read": bool(getattr(self, "competitive_evidence_read", False)),
-                "r_init_question": bool(getattr(self, "r_init_question", False)),
-                "er_single_chain": bool(getattr(self, "er_single_chain", False)),
-                "temporal_aggregation": bool(getattr(self, "temporal_aggregation", False)),
-                "concat_aggregation": bool(getattr(self, "concat_aggregation", False)),
-                "read_shortcut": bool(getattr(self, "read_shortcut", False)),
-                "splice_scale_match": bool(getattr(self, "splice_scale_match", False)),
-                "read_gating": bool(getattr(self, "read_gating", False)),
-                "read_residual_pure": bool(getattr(self, "read_residual_pure", False)),
-                "final_verify": bool(getattr(self, "final_verify", False)),
-                "mid_read_step": bool(getattr(self, "mid_read_step", 0)),
-                "persistent_slot_id": bool(getattr(self, "persistent_slot_id", False)),
-                "interface_loss": bool(getattr(self, "interface_loss", False)),
-                "read_anchor": bool(getattr(self, "read_anchor", False)),
-                "functional_state_loss": bool(getattr(self, "functional_state_loss", False)),
-            }
-            active = [name for name, enabled in incompatible.items() if enabled]
-            if active:
-                raise ValueError(
-                    "counterfactual_residual_recurrence replaces the slot/raw "
-                    f"workspace and is incompatible with: {', '.join(active)}"
-                )
-        elif getattr(self, "crr_policy_rank", 0) > 0:
-            raise ValueError(
-                "crr_policy_rank > 0 requires "
-                "counterfactual_residual_recurrence=True."
-            )
-        if getattr(self, "visual_counterfactual_recurrence", False):
-            if not getattr(self, "counterfactual_residual_recurrence", False):
-                raise ValueError(
-                    "visual_counterfactual_recurrence requires "
-                    "counterfactual_residual_recurrence=True."
-                )
-            if not getattr(self, "visual_cumulative_recurrence", False):
-                raise ValueError(
-                    "Qwen3 visual recurrence currently requires the promoted "
-                    "cumulative full-state path."
-                )
-            if not getattr(self, "visual_full_state_recurrence", False):
-                raise ValueError(
-                    "Qwen3 visual recurrence currently requires "
-                    "visual_full_state_recurrence=True."
-                )
-            incompatible = {
-                "counterfactual_reliance_loss": bool(
-                    getattr(self, "counterfactual_reliance_loss", False)
-                ),
-                "transition_aware_reliance_loss": bool(
-                    getattr(self, "transition_aware_reliance_loss", False)
-                ),
-                "counterfactual_step_retention_loss": bool(
-                    getattr(self, "counterfactual_step_retention_loss", False)
-                ),
-                "crr_concat_aggregation": bool(
-                    getattr(self, "crr_concat_aggregation", False)
-                ),
-                "crr_policy_rank": bool(getattr(self, "crr_policy_rank", 0)),
-            }
-            active = [name for name, enabled in incompatible.items() if enabled]
-            if active:
-                raise ValueError(
-                    "Qwen3 full-state visual recurrence uses final-answer CE "
-                    "and final-state decoding only; incompatible with: "
-                    + ", ".join(active)
-                )
-        elif getattr(self, "visual_cumulative_recurrence", False) or getattr(
-            self, "visual_full_state_recurrence", False
-        ):
-            raise ValueError(
-                "visual cumulative/full-state flags require "
-                "visual_counterfactual_recurrence=True."
-            )
-        if getattr(self, "recurrence_only_adapter", False):
-            if not getattr(self, "visual_full_state_recurrence", False):
-                raise ValueError(
-                    "recurrence_only_adapter requires "
-                    "visual_full_state_recurrence=True."
-                )
-            if self.adapter_rank <= 0:
-                raise ValueError(
-                    "recurrence_only_adapter requires adapter_rank > 0."
-                )
-            recurrent_layer = int(self.ell_star) + 1
-            if recurrent_layer in set(self.adapter_exclude_layers):
-                raise ValueError(
-                    "recurrence_only_adapter cannot exclude its recurrent "
-                    f"layer {recurrent_layer} from LoRA."
-                )
-        modes = [bool(getattr(self, m, False)) for m in
-                 ("read_gating", "read_residual_pure", "final_verify")]
-        modes.append(bool(getattr(self, "mid_read_step", 0)))
-        if sum(modes) > 1:
-            raise ValueError(
-                "read_gating / read_residual_pure / final_verify are mutually "
-                "exclusive reread modes."
-            )
-
-
-__all__ = ["CloseQwen3VLConfig", "WORKSPACE_ROPE_MODES"]

cvrr_release_config.json

@@ -1,19 +0,0 @@
-{
-  "name": "CVRR-InternVL3-9B",
-  "base_model": "OpenGVLab/InternVL3-9B",
-  "ell_star": 35,
-  "recurrent_layer": 36,
-  "upper_decoder_start": 37,
-  "inference_T": 4,
-  "inference_beta": 0.33,
-  "lora_rank": 32,
-  "lora_alpha": 12.0,
-  "lora_dropout": 0.01,
-  "checkpoint_step": 500,
-  "format": "cvrr_native_plus_merged_transition_v1",
-  "merged_weight_dtype": "float32",
-  "native_weight_bytes": 18277586944,
-  "status": "gpu_vstar_comparison_completed",
-  "upload_ready": true,
-  "hub_model_id": "dmis-lab/InternVL3-9B-CVRR"
-}

modeling_cvrr_merged.py

@@ -6,7 +6,6 @@ tensors or external backbone download is required after loading the package.
 from __future__ import annotations
 
 import contextlib
-import json
 import types
 from pathlib import Path
 from typing import TYPE_CHECKING
@@ -21,16 +20,16 @@ from .configuration_cvrr_merged import CVRRMergedConfig
 # Transformers 4.57's local-directory loader copies direct relative imports;
 # enumerate transitive helpers without eagerly importing every backbone.
 if TYPE_CHECKING:
-    from .source_splitting import SplitContext
-    from .source_perceive import PerceiveDeliberateLayout
-    from .source_spatial import SpatialVisualRecurrentCell
+    raise RuntimeError("This release does not bundle source_splitting")
+    raise RuntimeError("This release does not bundle source_perceive")
+    raise RuntimeError("This release does not bundle source_spatial")
     from .source_helpers import _layer_hidden
     from .source_gemma import GemmaCVRR
     from .source_internvl import InternVLCVRR
-    from .configuration_source_qwen25 import CloseQwen2_5_VLConfig
-    from .configuration_source_qwen3 import CloseQwen3VLConfig
-    from .modeling_source_qwen25 import CloseQwen2_5_VLForConditionalGeneration
-    from .modeling_source_qwen3 import CloseQwen3VLForConditionalGeneration
+    raise RuntimeError("This release does not bundle configuration_source_qwen25")
+    raise RuntimeError("This release does not bundle configuration_source_qwen3")
+    raise RuntimeError("This release does not bundle modeling_source_qwen25")
+    raise RuntimeError("This release does not bundle modeling_source_qwen3")
 
 
 class NativeMergedLinear(nn.Module):
@@ -106,6 +105,108 @@ def install_merged(runtime, weights, *, generic):
     return runtime.eval()
 
 
+def _generic_text_model(runtime):
+    return (runtime.base_model.language_model.model if runtime.model_type == 'internvl_chat'
+            else runtime.base_model.model)
+
+
+def _generic_last_logits(runtime, hidden):
+    if runtime.model_type == 'internvl_chat':
+        logits = runtime.base_model.language_model.get_output_embeddings()(hidden[:, -1:])
+    else:
+        logits = runtime.base_model.lm_head(hidden[:, -1:])
+        if runtime.model_type in {'gemma4', 'gemma4_unified'}:
+            cap = runtime.base_model.config.text_config.final_logit_softcapping
+            if cap is not None:
+                logits = torch.tanh(logits / cap) * cap
+    return logits[:, -1].float()
+
+
+def _replace_layer_hidden(output, hidden):
+    if hasattr(output, 'to_tuple'):
+        raise TypeError('decoder layer unexpectedly returned a ModelOutput')
+    if isinstance(output, tuple):
+        return (hidden, *output[1:])
+    if isinstance(output, list):
+        return [hidden, *output[1:]]
+    return hidden
+
+
+def _generic_cache_length(cache):
+    """Read text length from either Transformers Cache or a legacy KV tuple."""
+    if hasattr(cache, 'get_seq_length'):
+        return int(cache.get_seq_length())
+    if not isinstance(cache, (tuple, list)) or not cache:
+        raise TypeError('strict generation received an unsupported KV cache')
+    lengths = []
+    for layer in cache:
+        if not isinstance(layer, (tuple, list)) or len(layer) < 2:
+            raise TypeError('strict generation received a malformed legacy KV cache')
+        for tensor in layer[:2]:
+            if not torch.is_tensor(tensor) or tensor.ndim < 3:
+                raise TypeError('strict generation received a malformed legacy KV tensor')
+            lengths.append(int(tensor.shape[-2]))
+    if len(set(lengths)) != 1:
+        raise RuntimeError('strict generation found inconsistent legacy KV lengths')
+    return lengths[0]
+
+
+def _generic_strict_prefill(runtime, inputs, trace, recurrent):
+    """Cache only the image-free prompt, splicing R_T above the recurrent cell."""
+    if inputs['input_ids'].shape[0] != 1:
+        raise ValueError('strict generation currently supports one item at a time')
+    ids = inputs['input_ids'][trace.text_rows].unsqueeze(0)
+    called = []
+
+    def splice(_module, _args, output):
+        called.append(True)
+        return _replace_layer_hidden(output, recurrent)
+
+    hook = runtime.layers[runtime.cell_index].register_forward_hook(splice)
+    try:
+        with runtime.adapters(False):
+            result = _generic_text_model(runtime)(
+                input_ids=ids, attention_mask=torch.ones_like(ids),
+                use_cache=True, return_dict=True,
+            )
+    finally:
+        hook.remove()
+    if len(called) != 1 or result.past_key_values is None:
+        raise RuntimeError('strict text-only prefill did not produce one splice and a KV cache')
+    if _generic_cache_length(result.past_key_values) != ids.shape[1]:
+        raise RuntimeError('strict text-only KV cache length differs from the image-free prompt')
+    return _generic_last_logits(runtime, result.last_hidden_state), result.past_key_values
+
+
+def _generic_greedy_tokens(runtime, logits, cache, question_length, *, max_new_tokens, eos_token_ids):
+    generated = []
+    with runtime.adapters(False):
+        for step in range(max_new_tokens):
+            if not bool(torch.isfinite(logits).all()):
+                raise FloatingPointError('non-finite greedy-generation logits')
+            token = logits.argmax(-1).reshape(1, 1)
+            generated.append(int(token.item()))
+            if generated[-1] in eos_token_ids or step + 1 == max_new_tokens:
+                break
+            mask = torch.ones((1, question_length + len(generated)), dtype=torch.long, device=token.device)
+            output = _generic_text_model(runtime)(
+                input_ids=token, attention_mask=mask,
+                past_key_values=cache, use_cache=True, return_dict=True,
+            )
+            cache = output.past_key_values
+            if _generic_cache_length(cache) != question_length + len(generated):
+                raise RuntimeError('strict text-only KV cache did not advance by one text token')
+            logits = _generic_last_logits(runtime, output.last_hidden_state)
+    return generated
+
+
+def _add_eos_ids(target, value):
+    if value is None:
+        return
+    values = value if isinstance(value, (tuple, list, set)) else (value,)
+    target.update(int(token) for token in values)
+
+
 class CVRRMergedModel(PreTrainedModel):
     config_class = CVRRMergedConfig
     main_input_name = 'input_ids'
@@ -147,17 +248,16 @@ class CVRRMergedModel(PreTrainedModel):
                                           token=token, local_files_only=offline))
         if config is None:
             config = CVRRMergedConfig.from_pretrained(path, local_files_only=True)
-        settings = json.loads((path / 'cvrr_release_config.json').read_text())
-        if config.release != settings:
-            raise ValueError('Root configuration and release metadata disagree')
+        settings = config.release
         if settings['format'] != 'cvrr_native_plus_merged_transition_v1':
             raise ValueError('Unsupported release layout')
         native_path = str(path / 'native_backbone')
         generic = 'Qwen' not in settings['name']
+        # The generic constructors retain their internal LoRA defaults only to
+        # build temporary projection wrappers. install_merged() replaces those
+        # wrappers with the exported dense weights before inference.
         common = dict(ell_star=settings['ell_star'], steps=settings['inference_T'],
-                      beta=settings['inference_beta'], rank=settings['lora_rank'],
-                      alpha=settings['lora_alpha'], dropout=settings['lora_dropout'],
-                      device=device, offline=True)
+                      beta=settings['inference_beta'], device=device, offline=True)
         if generic:
             if 'InternVL' in settings['name']:
                 from .source_internvl import InternVLCVRR
@@ -168,12 +268,12 @@ class CVRRMergedModel(PreTrainedModel):
         else:
             if 'Qwen2.5' in settings['name']:
                 from transformers import Qwen2_5_VLForConditionalGeneration as Native
-                from .configuration_source_qwen25 import CloseQwen2_5_VLConfig as SourceConfig
-                from .modeling_source_qwen25 import CloseQwen2_5_VLForConditionalGeneration as Source
+                raise RuntimeError("This release does not bundle configuration_source_qwen25")
+                raise RuntimeError("This release does not bundle modeling_source_qwen25")
             else:
                 from transformers import Qwen3VLForConditionalGeneration as Native
-                from .configuration_source_qwen3 import CloseQwen3VLConfig as SourceConfig
-                from .modeling_source_qwen3 import CloseQwen3VLForConditionalGeneration as Source
+                raise RuntimeError("This release does not bundle configuration_source_qwen3")
+                raise RuntimeError("This release does not bundle modeling_source_qwen3")
             source_config = SourceConfig.from_pretrained(path / 'source_config', local_files_only=True)
             source_config.num_workspace_steps = settings['inference_T']
             source_config.counterfactual_beta = settings['inference_beta']
@@ -272,7 +372,25 @@ class CVRRMergedModel(PreTrainedModel):
         if do_sample:
             raise ValueError('Only deterministic greedy decoding is supported by this release')
         if self.is_generic:
-            raise NotImplementedError('Gemma/InternVL release validation currently covers strict next-token readout only')
+            if max_new_tokens < 1:
+                raise ValueError('max_new_tokens must be positive')
+            explicit_eos = inputs.pop('eos_token_id', None)
+            tokenizer = getattr(self, 'tokenizer', getattr(self.runtime, 'tokenizer', None))
+            eos = set()
+            for source in (getattr(self.runtime.base_model, 'generation_config', None), tokenizer):
+                _add_eos_ids(eos, getattr(source, 'eos_token_id', None))
+            _add_eos_ids(eos, explicit_eos)
+            if self.runtime.model_type == 'internvl_chat':
+                sep = self.runtime.base_model.conv_template.sep.strip()
+                eos.add(int(self.runtime.tokenizer.convert_tokens_to_ids(sep)))
+            trace = self.runtime.extract(inputs)
+            recurrent = self.runtime.rollout(trace)[-1]
+            logits, cache = _generic_strict_prefill(self.runtime, inputs, trace, recurrent)
+            tokens = _generic_greedy_tokens(
+                self.runtime, logits, cache, int(trace.question_lengths[0].item()),
+                max_new_tokens=max_new_tokens, eos_token_ids=eos,
+            )
+            return torch.tensor([tokens], dtype=torch.long, device=logits.device)
         return self.runtime.generate_answer(**inputs, max_new_tokens=max_new_tokens)
 
 

source_perceive.py

@@ -1,414 +0,0 @@
-"""Attention layouts for non-recurrent visual latent reasoning chains.
-
-The sequence is laid out as::
-
-    [multimodal prompt ; clean question ; LOOK_1 ; THINK_1 ; ... ; answer]
-
-All rows are processed once by the native VLM decoder.  A block-sparse causal
-graph makes LOOK rows the only latent rows with access to the multimodal
-prefix, while THINK rows integrate one visual read without directly seeing
-that prefix.  Answer rows can see the clean question and THINK rows, but never
-the multimodal prefix or LOOK rows.  Consequently the image-dependent answer
-path is
-
-    image -> LOOK_k -> THINK_k -> answer,
-
-without a tied recurrent cell, a visual-state update, or an aggregation
-module.
-
-The transition-conditioned variant tightens the graph without adding a
-module.  The first LOOK/THINK pair bootstraps from the complete prompt and the
-clean question.  Every later pair sees the causal latent prefix, while its
-LOOK row receives only visual placeholder rows from the multimodal prefix.
-Thus later pairs cannot independently re-solve the original image/question
-prompt; their task semantics must arrive through earlier latent states.  The
-last THINK is the sole latent exposed to answer rows and therefore acts as the
-native-transformer aggregation state.
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-import torch
-
-
-@dataclass(frozen=True)
-class PerceiveDeliberateLayout:
-    """Physical row layout of one padded training/generation sequence."""
-
-    multimodal_length: int
-    question_length: int
-    num_pairs: int
-    answer_input_length: int = 0
-
-    def __post_init__(self) -> None:
-        if self.multimodal_length < 1:
-            raise ValueError("multimodal_length must be positive")
-        if self.question_length < 1:
-            raise ValueError("question_length must be positive")
-        if self.num_pairs < 1:
-            raise ValueError("num_pairs must be positive")
-        if self.answer_input_length < 0:
-            raise ValueError("answer_input_length must be non-negative")
-
-    @property
-    def num_latent_tokens(self) -> int:
-        return 2 * self.num_pairs
-
-    @property
-    def question_start(self) -> int:
-        return self.multimodal_length
-
-    @property
-    def latent_start(self) -> int:
-        return self.multimodal_length + self.question_length
-
-    @property
-    def answer_start(self) -> int:
-        return self.latent_start + self.num_latent_tokens
-
-    @property
-    def sequence_length(self) -> int:
-        return self.answer_start + self.answer_input_length
-
-    @property
-    def question_slice(self) -> slice:
-        return slice(self.question_start, self.latent_start)
-
-    @property
-    def answer_slice(self) -> slice:
-        return slice(self.answer_start, self.sequence_length)
-
-    @property
-    def look_indices(self) -> tuple[int, ...]:
-        return tuple(self.latent_start + 2 * index for index in range(self.num_pairs))
-
-    @property
-    def think_indices(self) -> tuple[int, ...]:
-        return tuple(index + 1 for index in self.look_indices)
-
-
-def build_perceive_deliberate_mask(
-    layout: PerceiveDeliberateLayout,
-    multimodal_attention_mask: torch.Tensor,
-    question_attention_mask: torch.Tensor,
-    answer_input_attention_mask: torch.Tensor | None = None,
-    *,
-    multimodal_visual_mask: torch.Tensor | None = None,
-    local_visual_chain: bool = False,
-    transition_conditioned: bool = False,
-    question_visible_through_pair: int = 1,
-    disable_chain_links: bool = False,
-    disable_late_visual: bool = False,
-    final_pair_only: bool = False,
-) -> torch.BoolTensor:
-    """Build the strict LOOK/THINK/answer visibility graph.
-
-    The returned boolean mask has shape ``[B, 1, L, L]`` and follows PyTorch
-    SDPA semantics: ``True`` means that the query may attend to the key.
-
-    In the original graph, LOOK_k sees the valid multimodal prompt, the clean
-    question, and (for k>1) only THINK_{k-1}; THINK_k sees the clean question
-    and LOOK_k.  In the transition-conditioned graph, pair 1 keeps those
-    bootstrap inputs, but every later latent sees its complete causal latent
-    prefix, later LOOK rows see only visual placeholder keys, and later THINK
-    rows normally receive no independent question shortcut.  For a controlled
-    curriculum/intervention, ``question_visible_through_pair`` may temporarily
-    retain the clean-question edge through a later pair; ``1`` is the strict
-    inference graph.  Answer rows see only the final THINK in that variant.
-    Every row also sees itself so padded query
-    rows remain numerically defined; padded rows are never exposed as keys to
-    semantic rows.
-    """
-
-    if multimodal_attention_mask.ndim != 2:
-        raise ValueError("multimodal_attention_mask must have shape [B, L_mm]")
-    if question_attention_mask.ndim != 2:
-        raise ValueError("question_attention_mask must have shape [B, L_q]")
-    batch_size = multimodal_attention_mask.shape[0]
-    expected_mm = (batch_size, layout.multimodal_length)
-    expected_q = (batch_size, layout.question_length)
-    if tuple(multimodal_attention_mask.shape) != expected_mm:
-        raise ValueError(
-            "multimodal mask/layout mismatch: "
-            f"expected {expected_mm}, got {tuple(multimodal_attention_mask.shape)}"
-        )
-    if tuple(question_attention_mask.shape) != expected_q:
-        raise ValueError(
-            "question mask/layout mismatch: "
-            f"expected {expected_q}, got {tuple(question_attention_mask.shape)}"
-        )
-    if layout.answer_input_length:
-        if answer_input_attention_mask is None:
-            raise ValueError("answer_input_attention_mask is required")
-        expected_answer = (batch_size, layout.answer_input_length)
-        if tuple(answer_input_attention_mask.shape) != expected_answer:
-            raise ValueError(
-                "answer mask/layout mismatch: "
-                f"expected {expected_answer}, got "
-                f"{tuple(answer_input_attention_mask.shape)}"
-            )
-    elif answer_input_attention_mask is not None and answer_input_attention_mask.numel():
-        raise ValueError("received answer mask for an empty answer-input segment")
-
-    device = multimodal_attention_mask.device
-    length = layout.sequence_length
-    visible = torch.zeros(
-        batch_size, length, length, dtype=torch.bool, device=device
-    )
-    mm_valid = multimodal_attention_mask.bool()
-    q_valid = question_attention_mask.bool()
-    visual_valid = None
-    if transition_conditioned or local_visual_chain:
-        if not 1 <= question_visible_through_pair <= layout.num_pairs:
-            raise ValueError(
-                "question_visible_through_pair must lie in "
-                f"[1,{layout.num_pairs}]"
-            )
-        if multimodal_visual_mask is None:
-            raise ValueError(
-                "multimodal_visual_mask is required for the "
-                "transition-conditioned graph"
-            )
-        if tuple(multimodal_visual_mask.shape) != expected_mm:
-            raise ValueError(
-                "visual mask/layout mismatch: "
-                f"expected {expected_mm}, got "
-                f"{tuple(multimodal_visual_mask.shape)}"
-            )
-        visual_valid = multimodal_visual_mask.bool()
-        if bool((visual_valid & ~mm_valid).any()):
-            raise ValueError("visual keys must be a subset of valid multimodal keys")
-
-    # The source multimodal prompt keeps the native causal graph.  Padded query
-    # rows are harmless and receive a self edge below.
-    mm_causal = torch.ones(
-        layout.multimodal_length,
-        layout.multimodal_length,
-        dtype=torch.bool,
-        device=device,
-    ).tril()
-    visible[:, : layout.multimodal_length, : layout.multimodal_length] = (
-        mm_causal.unsqueeze(0) & mm_valid[:, None, :]
-    )
-
-    # The duplicated question is deliberately text-only: it is causally
-    # connected only to preceding valid rows of its own segment.
-    q_causal = torch.ones(
-        layout.question_length,
-        layout.question_length,
-        dtype=torch.bool,
-        device=device,
-    ).tril()
-    visible[
-        :, layout.question_slice, layout.question_slice
-    ] = q_causal.unsqueeze(0) & q_valid[:, None, :]
-
-    if local_visual_chain:
-        # Homogeneous one-pass latent chain.  Visual memory is persistent and
-        # available at every position, while task state has exactly one local
-        # predecessor edge.  No latent may consume the complete causal prefix.
-        latent_indices = range(layout.latent_start, layout.answer_start)
-        for latent_offset, latent in enumerate(latent_indices):
-            visible[:, latent, : layout.multimodal_length] = visual_valid
-            if latent_offset == 0:
-                visible[:, latent, layout.question_slice] = q_valid
-            else:
-                visible[:, latent, latent - 1] = True
-
-        if layout.answer_input_length:
-            answer_valid = answer_input_attention_mask.bool()
-            answer_causal = torch.ones(
-                layout.answer_input_length,
-                layout.answer_input_length,
-                dtype=torch.bool,
-                device=device,
-            ).tril()
-            visible[:, layout.answer_slice, layout.question_slice] = q_valid[:, None, :]
-            visible[:, layout.answer_slice, layout.answer_start - 1] = True
-            visible[:, layout.answer_slice, layout.answer_slice] = (
-                answer_causal.unsqueeze(0) & answer_valid[:, None, :]
-            )
-
-        diagonal = torch.arange(length, device=device)
-        visible[:, diagonal, diagonal] = True
-        return visible.unsqueeze(1)
-
-    for pair_index, (look, think) in enumerate(
-        zip(layout.look_indices, layout.think_indices)
-    ):
-        if final_pair_only and pair_index != layout.num_pairs - 1:
-            continue
-
-        if transition_conditioned and not final_pair_only:
-            # Pair 1 bootstraps the causal latent prefix.  Later pairs cannot
-            # recover the task from the original prompt: only image placeholder
-            # rows remain visible outside the latent prefix.
-            if pair_index == 0:
-                visible[:, look, : layout.multimodal_length] = mm_valid
-                visible[:, look, layout.question_slice] = q_valid
-                visible[:, think, layout.question_slice] = q_valid
-            else:
-                if not disable_late_visual:
-                    visible[:, look, : layout.multimodal_length] = visual_valid
-                if not disable_chain_links:
-                    visible[:, look, layout.latent_start:look] = True
-                    visible[:, think, layout.latent_start:look] = True
-                if pair_index < question_visible_through_pair:
-                    visible[:, look, layout.question_slice] = q_valid
-                    visible[:, think, layout.question_slice] = q_valid
-
-            # The local LOOK -> THINK edge is never an inter-pair chain-link
-            # intervention and therefore remains present in every arm.
-            visible[:, think, look] = True
-        else:
-            # Original graph, also used by the final-pair-only capacity
-            # control: each active pair may independently consume prompt + Q.
-            if not disable_late_visual or pair_index == 0 or final_pair_only:
-                visible[:, look, : layout.multimodal_length] = mm_valid
-            visible[:, look, layout.question_slice] = q_valid
-            if pair_index and not disable_chain_links and not final_pair_only:
-                visible[:, look, layout.think_indices[pair_index - 1]] = True
-
-            # Deliberation: no original multimodal key can be consumed here.
-            visible[:, think, layout.question_slice] = q_valid
-            visible[:, think, look] = True
-
-    if layout.answer_input_length:
-        answer_valid = answer_input_attention_mask.bool()
-        answer_causal = torch.ones(
-            layout.answer_input_length,
-            layout.answer_input_length,
-            dtype=torch.bool,
-            device=device,
-        ).tril()
-        visible[:, layout.answer_slice, layout.question_slice] = q_valid[:, None, :]
-        answer_thinks = (
-            [layout.think_indices[-1]]
-            if final_pair_only or transition_conditioned
-            else list(layout.think_indices)
-        )
-        visible[:, layout.answer_slice, answer_thinks] = True
-        visible[:, layout.answer_slice, layout.answer_slice] = (
-            answer_causal.unsqueeze(0) & answer_valid[:, None, :]
-        )
-
-    # Avoid all-masked softmax rows for physical padding.  Semantic queries do
-    # not receive padded keys because every segment assignment above uses its
-    # validity mask.
-    diagonal = torch.arange(length, device=device)
-    visible[:, diagonal, diagonal] = True
-    return visible.unsqueeze(1)
-
-
-def build_perceive_deliberate_positions(
-    layout: PerceiveDeliberateLayout,
-    multimodal_position_ids: torch.LongTensor,
-    multimodal_attention_mask: torch.Tensor,
-    question_attention_mask: torch.Tensor,
-    answer_input_attention_mask: torch.Tensor | None = None,
-) -> torch.LongTensor:
-    """Extend native multimodal M-RoPE with logical text positions.
-
-    The clean question keeps ordinary relative text positions but is shifted
-    after the largest valid multimodal M-RoPE coordinate.  Latent and answer
-    rows then continue from each item's *logical* question length, independent
-    of right padding.  All three M-RoPE coordinates are equal for these
-    non-spatial rows.
-    """
-
-    if multimodal_position_ids.ndim != 3 or multimodal_position_ids.shape[0] != 3:
-        raise ValueError("multimodal_position_ids must have shape [3, B, L_mm]")
-    batch_size = multimodal_attention_mask.shape[0]
-    if tuple(multimodal_position_ids.shape[1:]) != (
-        batch_size,
-        layout.multimodal_length,
-    ):
-        raise ValueError("multimodal position/layout mismatch")
-
-    mm_valid = multimodal_attention_mask.bool()
-    masked_mm = multimodal_position_ids.masked_fill(~mm_valid.unsqueeze(0), -1)
-    continuation = masked_mm.amax(dim=(0, 2)).clamp_min(-1) + 1  # [B]
-
-    q_valid = question_attention_mask.bool()
-    q_relative = q_valid.long().cumsum(dim=-1) - 1
-    q_relative = q_relative.masked_fill(~q_valid, 0)
-    q_positions = continuation[:, None] + q_relative
-    q_lengths = q_valid.sum(dim=-1)
-
-    latent_relative = torch.arange(
-        layout.num_latent_tokens,
-        device=multimodal_position_ids.device,
-    )
-    latent_positions = (
-        continuation[:, None] + q_lengths[:, None] + latent_relative[None, :]
-    )
-
-    segments = [multimodal_position_ids, q_positions.unsqueeze(0).expand(3, -1, -1)]
-    segments.append(latent_positions.unsqueeze(0).expand(3, -1, -1))
-    if layout.answer_input_length:
-        if answer_input_attention_mask is None:
-            raise ValueError("answer_input_attention_mask is required")
-        answer_relative = torch.arange(
-            layout.answer_input_length,
-            device=multimodal_position_ids.device,
-        )
-        answer_positions = (
-            continuation[:, None]
-            + q_lengths[:, None]
-            + layout.num_latent_tokens
-            + answer_relative[None, :]
-        )
-        segments.append(answer_positions.unsqueeze(0).expand(3, -1, -1))
-    return torch.cat(segments, dim=-1)
-
-
-def build_answer_decode_mask(
-    layout: PerceiveDeliberateLayout,
-    question_attention_mask: torch.Tensor,
-    generated_attention_mask: torch.Tensor,
-    *,
-    local_visual_chain: bool = False,
-    final_pair_only: bool = False,
-    transition_conditioned: bool = False,
-) -> torch.BoolTensor:
-    """Visibility for one cached answer query after the latent prefill.
-
-    ``generated_attention_mask`` includes the current answer token.  Physical
-    multimodal and LOOK cache rows remain present but are invisible.
-    """
-
-    if generated_attention_mask.ndim != 2:
-        raise ValueError("generated_attention_mask must have shape [B, N]")
-    batch_size, generated_length = generated_attention_mask.shape
-    if tuple(question_attention_mask.shape) != (
-        batch_size,
-        layout.question_length,
-    ):
-        raise ValueError("question mask/layout mismatch")
-    total_keys = layout.answer_start + generated_length
-    visible = torch.zeros(
-        batch_size, 1, 1, total_keys, dtype=torch.bool,
-        device=question_attention_mask.device,
-    )
-    visible[:, 0, 0, layout.question_slice] = question_attention_mask.bool()
-    if local_visual_chain:
-        visible[:, 0, 0, layout.answer_start - 1] = True
-    else:
-        answer_thinks = (
-            [layout.think_indices[-1]]
-            if final_pair_only or transition_conditioned
-            else list(layout.think_indices)
-        )
-        visible[:, 0, 0, answer_thinks] = True
-    visible[:, 0, 0, layout.answer_start:] = generated_attention_mask.bool()
-    return visible
-
-
-__all__ = [
-    "PerceiveDeliberateLayout",
-    "build_answer_decode_mask",
-    "build_perceive_deliberate_mask",
-    "build_perceive_deliberate_positions",
-]

source_spatial.py

@@ -1,349 +0,0 @@
-"""Question-conditioned recurrence over the complete native visual field.
-
-The state is the dynamic-resolution sequence emitted by Qwen2.5-VL's vision
-merger, before those image embeddings enter the language model.  Tokens keep
-their native two-dimensional grid throughout the recurrence.  One shared cell
-combines directional local messages with visual-to-question cross-attention::
-
-    V_0 = VisionMerger(image)
-    V_{k+1} = V_k + Cell(V_k, question, grid),  k = 0, ..., T - 1
-
-Only ``V_T`` is inserted into the ordinary Qwen multimodal prefix.  The cell's
-output projection is exactly zero at initialization, so every horizon,
-including the default T=8, is initially bitwise identical to the base visual
-embedding path.  The full language model is then run from scratch; no cache
-created while constructing the visual field is available to answer decoding.
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
-
-
-@dataclass
-class SpatialVisualTelemetry:
-    """Detached recurrence diagnostics; no visual trajectory is retained."""
-
-    update_rms: torch.Tensor  # [T]
-    update_relative: torch.Tensor  # [T]
-    step_drift: torch.Tensor  # [T], 1 - cos(V_k, V_{k+1})
-    final_relative: torch.Tensor  # scalar, RMS(V_T - V_0) / RMS(V_0)
-
-
-class SpatialVisualRecurrentCell(nn.Module):
-    """Shared O(N) spatial cell for packed, variable-resolution image tokens.
-
-    Global N-by-N visual attention is deliberately absent: a sample may carry
-    8192 merged visual tokens.  Four directional grid messages preserve local
-    geometry in linear time, while low-rank visual-to-question attention makes
-    every patch update task dependent.  The recurrent *state* remains at the
-    native language-model width; ``inner_width`` limits only the update
-    operator's compute and is not a state bottleneck.
-    """
-
-    def __init__(
-        self,
-        width: int,
-        inner_width: int,
-        num_heads: int,
-        *,
-        rms_norm_eps: float,
-        residual_scale: float,
-    ) -> None:
-        super().__init__()
-        if width < 1 or inner_width < 1:
-            raise ValueError("width and inner_width must be positive")
-        if num_heads < 1 or inner_width % num_heads:
-            raise ValueError(
-                "inner_width must be divisible by the positive num_heads"
-            )
-        if not 0.0 <= residual_scale <= 1.0:
-            raise ValueError("residual_scale must lie in [0, 1]")
-        self.width = int(width)
-        self.inner_width = int(inner_width)
-        self.num_heads = int(num_heads)
-        self.head_dim = self.inner_width // self.num_heads
-        self.residual_scale = float(residual_scale)
-
-        self.state_norm = Qwen2RMSNorm(self.width, eps=rms_norm_eps)
-        self.question_norm = Qwen2RMSNorm(self.width, eps=rms_norm_eps)
-
-        # One projection contains distinct north/south/west/east maps.  This
-        # keeps orientation identifiable without allocating a dense 3x3 dxd
-        # convolution over the 3584-wide native visual state.
-        self.center_proj = nn.Linear(self.width, self.inner_width, bias=False)
-        self.neighbor_proj = nn.Linear(
-            self.width, 4 * self.inner_width, bias=False
-        )
-
-        self.query_proj = nn.Linear(self.width, self.inner_width, bias=False)
-        self.key_proj = nn.Linear(self.width, self.inner_width, bias=False)
-        self.value_proj = nn.Linear(self.width, self.inner_width, bias=False)
-        self.output_proj = nn.Linear(self.inner_width, self.width, bias=False)
-
-    def zero_output(self) -> None:
-        """Restore the exact base-model identity after generic HF init."""
-        with torch.no_grad():
-            self.output_proj.weight.zero_()
-
-    @staticmethod
-    def _token_layout(
-        merged_grid_thw: torch.LongTensor,
-        image_to_batch: torch.LongTensor,
-        batch_size: int,
-        expected_tokens: int,
-    ) -> tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor]:
-        """Return token-to-batch ids, within-sample positions and counts."""
-        if merged_grid_thw.ndim != 2 or merged_grid_thw.shape[-1] != 3:
-            raise ValueError("merged_grid_thw must have shape [num_images, 3]")
-        if image_to_batch.ndim != 1 or image_to_batch.shape[0] != len(
-            merged_grid_thw
-        ):
-            raise ValueError("image_to_batch must contain one id per image")
-        per_image = merged_grid_thw.prod(dim=-1).long()
-        if int(per_image.sum()) != int(expected_tokens):
-            raise ValueError(
-                "merged grids and visual field disagree: "
-                f"grid tokens={int(per_image.sum())}, state tokens={expected_tokens}"
-            )
-        token_to_batch = torch.repeat_interleave(image_to_batch, per_image)
-        if token_to_batch.numel() and (
-            int(token_to_batch.min()) < 0
-            or int(token_to_batch.max()) >= batch_size
-        ):
-            raise ValueError("image_to_batch contains an out-of-range batch id")
-        if token_to_batch.numel() > 1 and bool(
-            (token_to_batch[1:] < token_to_batch[:-1]).any()
-        ):
-            raise ValueError(
-                "images must follow the same batch-major order as Qwen inputs"
-            )
-        counts = torch.bincount(token_to_batch, minlength=batch_size)
-        starts = counts.cumsum(0) - counts
-        within = torch.arange(
-            expected_tokens, device=token_to_batch.device
-        ) - torch.repeat_interleave(starts, counts)
-        return token_to_batch, within, counts
-
-    @staticmethod
-    def _neighbor_edges(
-        merged_grid_thw: torch.LongTensor,
-        expected_tokens: int,
-    ) -> tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor]:
-        """Build packed directed 4-neighbour edges once for all recurrent steps.
-
-        Direction ids are 0/1/2/3 = north/south/west/east *as seen by the
-        destination token*.  Temporal slices are separate 2-D fields.
-        """
-        device = merged_grid_thw.device
-        destinations: list[torch.Tensor] = []
-        sources: list[torch.Tensor] = []
-        directions: list[torch.Tensor] = []
-        offset = 0
-        for t_value, h_value, w_value in merged_grid_thw.detach().cpu().tolist():
-            t, h, w = int(t_value), int(h_value), int(w_value)
-            count = t * h * w
-            if min(t, h, w) < 1:
-                raise ValueError(f"invalid merged visual grid {(t, h, w)}")
-            index = torch.arange(
-                offset, offset + count, device=device, dtype=torch.long
-            ).view(t, h, w)
-
-            def append(dst: torch.Tensor, src: torch.Tensor, direction: int) -> None:
-                if dst.numel() == 0:
-                    return
-                destinations.append(dst.reshape(-1))
-                sources.append(src.reshape(-1))
-                directions.append(
-                    torch.full(
-                        (dst.numel(),), direction, device=device, dtype=torch.long
-                    )
-                )
-
-            append(index[:, 1:, :], index[:, :-1, :], 0)  # north
-            append(index[:, :-1, :], index[:, 1:, :], 1)  # south
-            append(index[:, :, 1:], index[:, :, :-1], 2)  # west
-            append(index[:, :, :-1], index[:, :, 1:], 3)  # east
-            offset += count
-        if offset != expected_tokens:
-            raise ValueError(
-                f"edge grids contain {offset} tokens, expected {expected_tokens}"
-            )
-        if not destinations:
-            empty = torch.empty(0, device=device, dtype=torch.long)
-            return empty, empty, empty
-        return (
-            torch.cat(destinations),
-            torch.cat(sources),
-            torch.cat(directions),
-        )
-
-    def _question_attention(
-        self,
-        normalized_state: torch.Tensor,
-        question_keys: torch.Tensor,
-        question_values: torch.Tensor,
-        question_valid: torch.BoolTensor,
-        token_to_batch: torch.LongTensor,
-        within_sample: torch.LongTensor,
-        token_counts: torch.LongTensor,
-    ) -> torch.Tensor:
-        """Visual-query/text-memory SDPA for a packed visual sequence."""
-        batch_size = question_keys.shape[0]
-        max_tokens = int(token_counts.max()) if token_counts.numel() else 0
-        if max_tokens == 0:
-            return normalized_state.new_empty(0, self.inner_width)
-
-        query_flat = self.query_proj(normalized_state)
-        query = query_flat.new_zeros(
-            batch_size, max_tokens, self.inner_width
-        )
-        query[token_to_batch, within_sample] = query_flat
-
-        bsz, question_length, _ = question_keys.shape
-        query = query.view(
-            bsz, max_tokens, self.num_heads, self.head_dim
-        ).transpose(1, 2)
-        key = question_keys.view(
-            bsz, question_length, self.num_heads, self.head_dim
-        ).transpose(1, 2)
-        value = question_values.view(
-            bsz, question_length, self.num_heads, self.head_dim
-        ).transpose(1, 2)
-        # Boolean SDPA masks use True for entries that are allowed to attend.
-        allowed = question_valid[:, None, None, :]
-        attended = F.scaled_dot_product_attention(
-            query,
-            key,
-            value,
-            attn_mask=allowed,
-            dropout_p=0.0,
-            is_causal=False,
-        )
-        attended = attended.transpose(1, 2).reshape(
-            bsz, max_tokens, self.inner_width
-        )
-        return attended[token_to_batch, within_sample]
-
-    def forward(
-        self,
-        visual_state: torch.Tensor,
-        merged_grid_thw: torch.LongTensor,
-        image_to_batch: torch.LongTensor,
-        question_embeddings: torch.Tensor,
-        question_attention_mask: torch.Tensor | None,
-        *,
-        steps: int,
-    ) -> tuple[torch.Tensor, SpatialVisualTelemetry]:
-        if visual_state.ndim != 2 or visual_state.shape[-1] != self.width:
-            raise ValueError(
-                f"visual_state must be [N,{self.width}], got "
-                f"{tuple(visual_state.shape)}"
-            )
-        if question_embeddings.ndim != 3 or question_embeddings.shape[-1] != self.width:
-            raise ValueError(
-                f"question_embeddings must be [B,Q,{self.width}]"
-            )
-        if steps < 0:
-            raise ValueError("steps must be non-negative")
-        batch_size = question_embeddings.shape[0]
-        token_to_batch, within_sample, token_counts = self._token_layout(
-            merged_grid_thw,
-            image_to_batch,
-            batch_size,
-            visual_state.shape[0],
-        )
-        edge_dst, edge_src, edge_direction = self._neighbor_edges(
-            merged_grid_thw, visual_state.shape[0]
-        )
-
-        if question_attention_mask is None:
-            question_valid = torch.ones(
-                question_embeddings.shape[:2],
-                dtype=torch.bool,
-                device=question_embeddings.device,
-            )
-        else:
-            if question_attention_mask.shape != question_embeddings.shape[:2]:
-                raise ValueError("question attention mask shape mismatch")
-            question_valid = question_attention_mask > 0
-        if not bool(question_valid.any(dim=-1).all()):
-            raise ValueError("every image-bearing sample needs a question token")
-
-        normalized_question = self.question_norm(question_embeddings)
-        question_keys = self.key_proj(normalized_question)
-        question_values = self.value_proj(normalized_question)
-
-        initial = visual_state
-        state = visual_state
-        update_rms: list[torch.Tensor] = []
-        update_relative: list[torch.Tensor] = []
-        step_drift: list[torch.Tensor] = []
-        for _ in range(steps):
-            normalized = self.state_norm(state)
-            center = self.center_proj(normalized)
-            directional = self.neighbor_proj(normalized).view(
-                state.shape[0], 4, self.inner_width
-            )
-            spatial = center.new_zeros(center.shape)
-            degree = center.new_zeros(center.shape[0], 1)
-            if edge_dst.numel():
-                messages = directional[edge_src, edge_direction]
-                spatial.index_add_(0, edge_dst, messages)
-                degree.index_add_(
-                    0,
-                    edge_dst,
-                    torch.ones(
-                        edge_dst.shape[0], 1, dtype=center.dtype, device=center.device
-                    ),
-                )
-                spatial = spatial / degree.clamp_min(1.0)
-            question = self._question_attention(
-                normalized,
-                question_keys,
-                question_values,
-                question_valid,
-                token_to_batch,
-                within_sample,
-                token_counts,
-            )
-            mixed = F.silu((center + spatial + question) / (3.0**0.5))
-            update = self.output_proj(mixed) * self.residual_scale
-            next_state = state + update
-
-            with torch.no_grad():
-                state_rms = state.float().square().mean().sqrt().clamp_min(1e-8)
-                update_norm = update.float().square().mean().sqrt()
-                cosine = F.cosine_similarity(
-                    state.float(), next_state.float(), dim=-1
-                ).mean()
-                update_rms.append(update_norm.detach())
-                update_relative.append((update_norm / state_rms).detach())
-                step_drift.append((1.0 - cosine).detach())
-            state = next_state
-
-        with torch.no_grad():
-            initial_rms = initial.float().square().mean().sqrt().clamp_min(1e-8)
-            final_relative = (
-                (state.float() - initial.float()).square().mean().sqrt()
-                / initial_rms
-            ).detach()
-            empty = initial.new_empty(0, dtype=torch.float32)
-            telemetry = SpatialVisualTelemetry(
-                update_rms=(torch.stack(update_rms) if update_rms else empty),
-                update_relative=(
-                    torch.stack(update_relative) if update_relative else empty
-                ),
-                step_drift=(torch.stack(step_drift) if step_drift else empty),
-                final_relative=final_relative,
-            )
-        return state, telemetry
-
-
-__all__ = ["SpatialVisualRecurrentCell", "SpatialVisualTelemetry"]

source_splitting.py

@@ -1,473 +0,0 @@
-"""Layer-split forward for Qwen2.5-VL and Qwen3-VL.
-
-Method §3.2 needs ``F = F_{>l*} o F_{<=l*}`` as two separately runnable halves so
-that the multimodal branch can be cut off at ``l*`` and replaced by the
-workspace. §3.1 needs the same split to read and patch activations at a chosen
-depth.
-
-This module reimplements the prologue of the native Qwen VL text-model forward
--- embedding merge, M-RoPE index, causal mask, rotary embeddings -- as a
-reusable :class:`SplitContext`, then exposes the decoder-layer loop as a range
-you can run piecewise.  Qwen3-VL additionally injects three ``DeepStack``
-vision features after language layers 0--2; those tensors are carried in the
-split context and applied at the identical layer boundaries.  It deliberately
-mirrors transformers 4.57.6 rather than monkeypatching it;
-``tests/test_split_equivalence.py`` asserts the composed halves reproduce the
-stock forward bit-for-bit, which is what makes the mirroring safe to rely on.
-
-Shapes use ``B`` batch, ``L`` sequence, ``d`` backbone width (3584 on the 7B),
-``N_v`` visual tokens, ``N_q`` question tokens.
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass, replace
-from typing import Any
-
-import torch
-from transformers.cache_utils import Cache
-from transformers.masking_utils import (
-    create_causal_mask,
-    create_sliding_window_causal_mask,
-)
-
-
-@dataclass
-class SplitContext:
-    """Per-forward state shared by every decoder layer.
-
-    Computed once by :func:`make_split_context` so that layer ranges can be run
-    independently without recomputing masks or rotary tables.
-    """
-
-    hidden_states: torch.Tensor  # [B, L, d] - mutated as layers run
-    position_ids: torch.Tensor  # [3, B, L] - M-RoPE (t, h, w)
-    position_embeddings: tuple[torch.Tensor, torch.Tensor]  # (cos, sin) [B, L, head_dim]
-    causal_mask_mapping: dict[str, torch.Tensor | None]
-    cache_position: torch.Tensor  # [L]
-    text_position_ids: torch.Tensor | None  # [B, L] only when packed
-    past_key_values: Cache | None
-    # Original 2-D key-padding mask.  ``create_causal_mask`` is allowed to
-    # return ``None`` for SDPA and delegate causality to ``is_causal``; keeping
-    # this tensor lets counterfactual branches materialize the equivalent mask
-    # before removing a precisely selected set of attention edges.
-    attention_mask: torch.Tensor | None = None  # [B, L_kv]
-    # Qwen3-VL only.  DeepStack adds one visual feature tensor after each of
-    # the first three language layers.  They stay ``None`` for Qwen2.5-VL and
-    # for every text-only branch.
-    visual_pos_masks: torch.Tensor | None = None  # [B, L] bool
-    deepstack_visual_embeds: list[torch.Tensor] | None = None
-
-    def clone_at(self, hidden_states: torch.Tensor) -> "SplitContext":
-        """Same context, different hidden states (for patched re-runs)."""
-        return SplitContext(
-            hidden_states=hidden_states,
-            position_ids=self.position_ids,
-            position_embeddings=self.position_embeddings,
-            causal_mask_mapping=self.causal_mask_mapping,
-            cache_position=self.cache_position,
-            text_position_ids=self.text_position_ids,
-            past_key_values=self.past_key_values,
-            attention_mask=self.attention_mask,
-            visual_pos_masks=self.visual_pos_masks,
-            deepstack_visual_embeds=self.deepstack_visual_embeds,
-        )
-
-
-# ---------------------------------------------------------------------------
-# embedding / position construction
-# ---------------------------------------------------------------------------
-
-
-def embed_multimodal(
-    vl_model,
-    input_ids: torch.LongTensor,  # [B, L]
-    pixel_values: torch.Tensor | None = None,
-    image_grid_thw: torch.LongTensor | None = None,
-    attention_mask: torch.Tensor | None = None,
-    *,
-    return_deepstack: bool = False,
-) -> (
-    tuple[torch.Tensor, torch.Tensor]
-    | tuple[
-        torch.Tensor,
-        torch.Tensor,
-        torch.Tensor | None,
-        list[torch.Tensor] | None,
-    ]
-):
-    """Token embeddings with image features scattered in, plus M-RoPE indices.
-
-    Mirrors the prefill path of ``Qwen2_5_VLModel.forward``. Pass
-    ``pixel_values=None`` to get the text-only branch used for ``Q*``.
-
-    Args:
-        vl_model: a native ``Qwen2_5_VLModel`` or ``Qwen3VLModel`` (i.e.
-            ``model.model``, not the ``...ForConditionalGeneration`` wrapper).
-        return_deepstack: also return Qwen3-VL's visual-position mask and
-            DeepStack features.  The default two-value return keeps all
-            Qwen2.5 callers backward compatible.
-
-    Returns:
-        ``(inputs_embeds [B, L, d], position_ids [3, B, L])`` and optionally
-        ``(visual_pos_masks, deepstack_visual_embeds)``.
-    """
-    inputs_embeds = vl_model.get_input_embeddings()(input_ids)  # [B, L, d]
-    model_type = str(getattr(vl_model.config, "model_type", ""))
-    # A freshly wrapped model exposes the native qwen3_vl config here.
-    # Reloading a fully saved CLOSE checkpoint reconstructs the nested native
-    # backbone from the wrapper config, so its model view carries
-    # close_qwen3_vl instead. Both use the tuple-returning Qwen3 feature API.
-    is_qwen3_vl = model_type in {"qwen3_vl", "close_qwen3_vl"}
-    visual_pos_masks = None
-    deepstack_visual_embeds = None
-
-    if pixel_values is not None:
-        image_features = vl_model.get_image_features(pixel_values, image_grid_thw)
-        if is_qwen3_vl:
-            image_embeds, deepstack_visual_embeds = image_features
-        else:
-            image_embeds = image_features
-        image_embeds = torch.cat(image_embeds, dim=0).to(
-            inputs_embeds.device, inputs_embeds.dtype
-        )  # [N_v_total, d]
-        image_mask, _ = vl_model.get_placeholder_mask(
-            input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
-        )
-        inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
-        if is_qwen3_vl:
-            visual_pos_masks = image_mask[..., 0]
-
-    if is_qwen3_vl:
-        position_ids, _ = vl_model.get_rope_index(
-            input_ids,
-            image_grid_thw,
-            None,  # video_grid_thw
-            attention_mask=attention_mask,
-        )
-    else:
-        position_ids, _ = vl_model.get_rope_index(
-            input_ids,
-            image_grid_thw,
-            None,  # video_grid_thw
-            second_per_grid_ts=None,
-            attention_mask=attention_mask,
-        )
-    if return_deepstack:
-        return (
-            inputs_embeds,
-            position_ids,
-            visual_pos_masks,
-            deepstack_visual_embeds,
-        )
-    return inputs_embeds, position_ids
-
-
-def make_split_context(
-    text_model,
-    inputs_embeds: torch.Tensor,  # [B, L, d]
-    position_ids: torch.Tensor,  # [3, B, L]
-    attention_mask: torch.Tensor | None = None,
-    past_key_values: Cache | None = None,
-    cache_position: torch.Tensor | None = None,
-    visual_pos_masks: torch.Tensor | None = None,
-    deepstack_visual_embeds: list[torch.Tensor] | None = None,
-) -> SplitContext:
-    """Build masks and rotary embeddings once, as the stock forward does.
-
-    Args:
-        text_model: ``Qwen2_5_VLTextModel`` (``vl_model.language_model``).
-    """
-    if cache_position is None:
-        past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
-        cache_position = torch.arange(
-            past_seen, past_seen + inputs_embeds.shape[1], device=inputs_embeds.device
-        )  # [L]
-
-    if position_ids.ndim == 2:
-        position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
-
-    model_type = str(getattr(text_model.config, "model_type", ""))
-
-    # Packed-sequence convention: a leading text-only row makes it [4, B, L].
-    if position_ids.ndim == 3 and position_ids.shape[0] == 4:
-        text_position_ids = position_ids[0]  # [B, L]
-        position_ids = position_ids[1:]  # [3, B, L]
-    elif model_type == "qwen3_vl_text":
-        # Qwen3-VL always passes the temporal M-RoPE row to both the causal-mask
-        # builder and decoder layers, even for ordinary (non-packed) inputs.
-        text_position_ids = position_ids[0]
-    else:
-        text_position_ids = None
-
-    mask_kwargs: dict[str, Any] = {
-        "config": text_model.config,
-        "input_embeds": inputs_embeds,
-        "attention_mask": attention_mask,
-        "cache_position": cache_position,
-        "past_key_values": past_key_values,
-        "position_ids": text_position_ids,
-    }
-    causal_mask_mapping = {"full_attention": create_causal_mask(**mask_kwargs)}
-    if getattr(text_model, "has_sliding_layers", False):
-        causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(
-            **mask_kwargs
-        )
-
-    position_embeddings = text_model.rotary_emb(inputs_embeds, position_ids)
-
-    return SplitContext(
-        hidden_states=inputs_embeds,
-        position_ids=position_ids,
-        position_embeddings=position_embeddings,
-        causal_mask_mapping=causal_mask_mapping,
-        cache_position=cache_position,
-        text_position_ids=text_position_ids,
-        past_key_values=past_key_values,
-        attention_mask=attention_mask,
-        visual_pos_masks=visual_pos_masks,
-        deepstack_visual_embeds=deepstack_visual_embeds,
-    )
-
-
-def block_attention_edges(
-    ctx: SplitContext,
-    query_mask: torch.Tensor,
-    key_mask: torch.Tensor,
-) -> SplitContext:
-    """Return ``ctx`` with selected query-to-key attention edges removed.
-
-    ``query_mask`` and ``key_mask`` are boolean ``[B, L]`` supports in the
-    current (cache-free) sequence.  Every ordinary causal/padding constraint is
-    preserved; only their Cartesian product is additionally masked.  Both
-    boolean SDPA masks (``True`` means visible) and additive eager masks
-    (``0``/negative infinity) are supported.
-
-    The helper deliberately rejects cached contexts.  Its intended use is a
-    counterfactual recurrent layer evaluation, never autoregressive decoding,
-    and silently guessing the key offset of a populated cache would invalidate
-    the causal comparison.
-    """
-    if ctx.past_key_values is not None and ctx.past_key_values.get_seq_length() > 0:
-        raise ValueError("block_attention_edges requires a cache-free context")
-    if query_mask.dtype != torch.bool or key_mask.dtype != torch.bool:
-        raise TypeError("query_mask and key_mask must be boolean tensors")
-    if query_mask.shape != key_mask.shape or query_mask.ndim != 2:
-        raise ValueError(
-            "query_mask and key_mask must have the same [B, L] shape, got "
-            f"{tuple(query_mask.shape)} and {tuple(key_mask.shape)}"
-        )
-    batch_size, seq_len = query_mask.shape
-    if ctx.hidden_states.shape[:2] != (batch_size, seq_len):
-        raise ValueError(
-            "edge masks must match the SplitContext sequence, got "
-            f"{tuple(query_mask.shape)} for {tuple(ctx.hidden_states.shape[:2])}"
-        )
-
-    blocked = query_mask[:, None, :, None] & key_mask[:, None, None, :]
-    updated: dict[str, torch.Tensor] = {}
-    for attention_type, base_mask in ctx.causal_mask_mapping.items():
-        if base_mask is None:
-            # SDPA may omit an all-valid causal mask.  Materialize exactly that
-            # lower triangle, then reapply key padding before deleting edges.
-            q_positions = ctx.cache_position
-            if q_positions.numel() != seq_len:
-                raise ValueError(
-                    "cache-free context must have one cache position per row"
-                )
-            key_positions = torch.arange(seq_len, device=query_mask.device)
-            visible = key_positions[None, :] <= q_positions[:, None]
-            visible = visible[None, None].expand(batch_size, 1, -1, -1)
-            if ctx.attention_mask is not None:
-                if ctx.attention_mask.shape != (batch_size, seq_len):
-                    raise ValueError(
-                        "counterfactual edge masking expects a 2-D [B, L] "
-                        "attention mask"
-                    )
-                visible = visible & ctx.attention_mask[:, None, None, :].bool()
-            updated[attention_type] = visible & ~blocked
-            continue
-
-        if not isinstance(base_mask, torch.Tensor) or base_mask.ndim != 4:
-            raise TypeError(
-                "counterfactual edge masking supports tensor 4-D attention "
-                f"masks, got {type(base_mask)!r}"
-            )
-        if base_mask.shape[0] not in (1, batch_size):
-            raise ValueError("attention-mask batch dimension is incompatible")
-        if base_mask.shape[-2:] != (seq_len, seq_len):
-            raise ValueError(
-                "counterfactual edge masking expects a square cache-free mask, "
-                f"got {tuple(base_mask.shape)}"
-            )
-        if base_mask.dtype == torch.bool:
-            updated[attention_type] = base_mask & ~blocked
-        elif base_mask.is_floating_point():
-            updated[attention_type] = base_mask.masked_fill(
-                blocked, torch.finfo(base_mask.dtype).min
-            )
-        else:
-            raise TypeError(
-                f"unsupported attention mask dtype {base_mask.dtype}"
-            )
-
-    return replace(ctx, causal_mask_mapping=updated)
-
-
-# ---------------------------------------------------------------------------
-# running layer ranges
-# ---------------------------------------------------------------------------
-
-
-def run_layer_range(
-    text_model,
-    ctx: SplitContext,
-    start: int,
-    stop: int | None = None,
-    use_cache: bool = False,
-    hidden_states: torch.Tensor | None = None,
-    collect: bool = False,
-) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
-    """Run ``text_model.layers[start:stop]`` on ``ctx``.
-
-    ``self.norm`` is *not* applied -- it belongs to the very top of the stack.
-    Call :func:`final_norm` after the last range.
-
-    Args:
-        hidden_states: override the context's states (leave ``None`` to chain).
-        collect: also return the input hidden states of every layer in the range
-            plus the range output, i.e. ``stop - start + 1`` tensors.
-
-    Returns:
-        ``[B, L, d]``, or ``(output, collected)`` when ``collect``.
-    """
-    layers = text_model.layers
-    stop = len(layers) if stop is None else stop
-    h = ctx.hidden_states if hidden_states is None else hidden_states
-
-    collected: list[torch.Tensor] = []
-    for layer_index, layer in enumerate(layers[start:stop], start=start):
-        if collect:
-            collected.append(h)
-        attention_type = getattr(layer, "attention_type", "full_attention")
-        h = layer(
-            h,
-            attention_mask=ctx.causal_mask_mapping[attention_type],
-            position_ids=ctx.text_position_ids,
-            past_key_values=ctx.past_key_values,
-            use_cache=use_cache,
-            cache_position=ctx.cache_position,
-            position_embeddings=ctx.position_embeddings,
-        )
-        # 4.57 decoder layers return a bare tensor; older ones returned a tuple.
-        if isinstance(h, tuple):
-            h = h[0]
-        if (
-            ctx.deepstack_visual_embeds is not None
-            and layer_index < len(ctx.deepstack_visual_embeds)
-        ):
-            if ctx.visual_pos_masks is None:
-                raise ValueError("DeepStack features require visual_pos_masks")
-            h = text_model._deepstack_process(
-                h,
-                ctx.visual_pos_masks,
-                ctx.deepstack_visual_embeds[layer_index],
-            )
-
-    if collect:
-        collected.append(h)
-        return h, collected
-    return h
-
-
-def final_norm(text_model, hidden_states: torch.Tensor) -> torch.Tensor:
-    """Apply the stack's final RMSNorm. ``[B, L, d] -> [B, L, d]``."""
-    return text_model.norm(hidden_states)
-
-
-# ---------------------------------------------------------------------------
-# token selection operators (Pi_img / Pi_q in §3.2)
-# ---------------------------------------------------------------------------
-
-
-def image_token_mask(input_ids: torch.LongTensor, image_token_id: int) -> torch.Tensor:
-    """``Pi_img`` support: ``[B, L]`` bool, True at image placeholder positions."""
-    return input_ids == image_token_id
-
-
-def vision_span_mask(input_ids: torch.LongTensor, config) -> torch.Tensor:
-    """``[B, L]`` bool covering ``<|vision_start|>``, image pads, ``<|vision_end|>``.
-
-    Use this (not :func:`image_token_mask`) when *removing* the visual segment to
-    build the text-only branch, so the delimiters do not survive as orphans.
-    """
-    ids = {
-        config.vision_start_token_id,
-        config.vision_end_token_id,
-        config.image_token_id,
-        config.video_token_id,
-    }
-    mask = torch.zeros_like(input_ids, dtype=torch.bool)
-    for tid in ids:
-        mask |= input_ids == tid
-    return mask
-
-
-def select_tokens(
-    hidden_states: torch.Tensor,  # [B, L, d]
-    mask: torch.Tensor,  # [B, L] bool
-) -> torch.Tensor:
-    """Gather masked positions. Requires an equal count per batch element.
-
-    Returns ``[B, N, d]`` where ``N`` is that per-element count.
-    """
-    counts = mask.sum(dim=1)
-    if counts.numel() > 1 and not bool((counts == counts[0]).all()):
-        raise ValueError(
-            f"select_tokens needs the same number of selected tokens per batch "
-            f"element, got {counts.tolist()}. Bucket by visual-token count or "
-            f"gather per-example instead."
-        )
-    n = int(counts[0])
-    b, _, d = hidden_states.shape
-    return hidden_states[mask].view(b, n, d)
-
-
-def select_tokens_padded(
-    hidden_states: torch.Tensor,  # [B, L, d]
-    mask: torch.Tensor,  # [B, L] bool
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Gather masked positions, right-padded to the batch maximum.
-
-    Qwen2.5-VL uses dynamic resolution, so ``N_v`` differs across a batch. §3.1's
-    patching genuinely needs equal counts (it transplants position by position),
-    but ``r_theta`` only cross-attends over ``V*`` -- a variable-length memory is
-    exactly what a key-padding mask is for.
-
-    Returns ``(padded [B, N_max, d], key_padding_mask [B, N_max])`` where the
-    mask is ``True`` at padding, matching ``nn.MultiheadAttention``.
-    """
-    counts = mask.sum(dim=1)
-    n_max = int(counts.max())
-    b, _, d = hidden_states.shape
-    out = hidden_states.new_zeros((b, n_max, d))
-    pad = torch.ones((b, n_max), dtype=torch.bool, device=hidden_states.device)
-    for i in range(b):
-        n = int(counts[i])
-        out[i, :n] = hidden_states[i][mask[i]]
-        pad[i, :n] = False
-    return out, pad
-
-
-__all__ = [
-    "SplitContext",
-    "embed_multimodal",
-    "make_split_context",
-    "block_attention_edges",
-    "run_layer_range",
-    "final_norm",
-    "image_token_mask",
-    "vision_span_mask",
-    "select_tokens",
-    "select_tokens_padded",
-]