Upload model.py

model.py

@@ -0,0 +1,996 @@
+"""
+Foveated Vision-Language Model.
+
+Architecture: DINOv2 encoder + foveated cross-attention + SmolLM2 LLM.
+Each video frame is compressed to ONE visual token via query-guided attention.
+The LLM controls WHERE to look by generating the query for the next frame.
+
+Three forward modes:
+  1. forward_coarse_fine   -- Training (two parallel passes)
+  2. forward_coarse_only   -- Fast eval (single static-query pass)
+  3. forward_autoregressive -- True inference (sequential, KV-cached)
+
+Loss: text cross-entropy only (no reconstruction, no VAE).
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoModelForCausalLM, AutoConfig
+from typing import Dict, Optional
+
+# Optional: Liger Kernel fused CE loss (never materializes [B, S, V] logits)
+try:
+    from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss
+    _HAS_LIGER = True
+except ImportError:
+    _HAS_LIGER = False
+
+
+class FoveatedVLM(nn.Module):
+    """
+    Foveated Vision-Language Model.
+
+    Parameters
+    ----------
+    llm_name : str
+        HuggingFace model id for SmolLM2 (e.g. "HuggingFaceTB/SmolLM2-135M-Instruct").
+    dino_name : str
+        HuggingFace model id for DINOv2 (e.g. "facebook/dinov2-small").
+    query_dim : int
+        Dimension of the foveated query vectors (matches DINO dim by default).
+    visual_scale : float
+        Multiplicative factor applied to projected visual tokens so their
+        magnitude matches the LLM embedding std (~0.14 for SmolLM2).
+    lambda_coarse : float
+        Weight for the optional auxiliary coarse-pass CE loss during training.
+        Set to 0 to disable.
+    """
+
+    def __init__(
+        self,
+        llm_name: str = "HuggingFaceTB/SmolLM2-135M-Instruct",
+        dino_name: str = "facebook/dinov2-small",
+        query_dim: int = 384,
+        visual_scale: float = 0.14,
+        lambda_coarse: float = 0.0,
+        deep_query: bool = True,
+        use_fused_ce: bool = False,
+    ):
+        super().__init__()
+
+        # ---- delayed import so encoder.py can live next to this file ----
+        from encoder import FoveatedEncoder
+
+        # ---- Vision encoder (DINOv2 + query cross-attention) ----
+        self.encoder = FoveatedEncoder(
+            dino_model_name=dino_name,
+            query_dim=query_dim,
+            output_dim=None,  # output_dim = dino_dim by default inside encoder
+        )
+        dino_dim = self.encoder.dino_dim
+
+        # ---- Language model ----
+        self.llm = AutoModelForCausalLM.from_pretrained(
+            llm_name, attn_implementation="sdpa", torch_dtype=torch.bfloat16,
+        )
+        self.llm.config.use_cache = False  # training default; overridden per-method
+        llm_dim = self.llm.config.hidden_size
+
+        # ---- Projections ----
+        self.dino_to_llm = nn.Linear(dino_dim, llm_dim)
+        self.llm_to_query = nn.Linear(llm_dim, query_dim)
+
+        # ---- Learnable queries ----
+        # BUG-001 FIX: init with std=1.0 so queries dominate over projection
+        # bias and produce meaningful (non-uniform) attention patterns.
+        self.q_static = nn.Parameter(torch.randn(1, query_dim))   # std=1.0
+        self.q_init   = nn.Parameter(torch.randn(1, query_dim))   # std=1.0
+
+        # ---- Hyperparams stored as plain Python (not buffers) ----
+        self.visual_scale = visual_scale
+        self.lambda_coarse = lambda_coarse
+        self.query_dim = query_dim
+        self.deep_query = deep_query
+        self.use_fused_ce = use_fused_ce and _HAS_LIGER
+
+        # ---- Dimension bookkeeping (useful for external code) ----
+        self.dino_dim = dino_dim
+        self.llm_dim = llm_dim
+
+    # ------------------------------------------------------------------
+    # helpers
+    # ------------------------------------------------------------------
+
+    def _get_pad_token_id(self) -> int:
+        """Return pad_token_id from the LLM config (never hardcoded)."""
+        pid = getattr(self.llm.config, "pad_token_id", None)
+        if pid is None:
+            pid = getattr(self.llm.config, "eos_token_id", 0)
+        return pid
+
+    def _llm_dtype(self) -> torch.dtype:
+        """Return the dtype of the LLM parameters (e.g. bfloat16)."""
+        return next(self.llm.parameters()).dtype
+
+    def _embed_text(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """[B, S] -> [B, S, llm_dim] via LLM embedding table."""
+        return self.llm.get_input_embeddings()(input_ids)
+
+    def _project_visual(self, z: torch.Tensor) -> torch.Tensor:
+        """
+        Project DINO features to LLM space and rescale.
+
+        z : [B, T, dino_dim]  or  [B, dino_dim]
+        Returns same shape with last dim = llm_dim.
+        """
+        h = self.dino_to_llm(z)                       # -> llm_dim
+        h = h * self.visual_scale                      # match LLM embedding magnitude
+        return h
+
+    # Maximum frames per DINO encode/query call to prevent OOM on large batches.
+    _MAX_ENCODE_CHUNK = 200
+
+    def _encode_all_frames(self, frames: torch.Tensor, frame_mask=None):
+        """
+        Run DINO patch encoding for every frame in the batch.
+
+        frames     : [B, T, 3, 224, 224]
+        frame_mask : [B, T] bool — True for real frames, False for padding.
+
+        Returns (kv_cache, patch_features, mask_flat):
+            kv_cache       : list of (K, V) per layer, each [n_real, N+1, D]
+                             (compact — only real frames, no padding waste).
+            patch_features : [n_real, N+1, D] final DINO embeddings (for shallow mode).
+            mask_flat      : [B*T] bool tensor or None. Used to scatter results back.
+        """
+        B, T, C, H, W = frames.shape
+        BT = B * T
+        frames_flat = frames.reshape(BT, C, H, W)
+
+        if frame_mask is not None:
+            mask_flat = frame_mask.reshape(BT)
+            n_real = mask_flat.sum().item()
+        else:
+            mask_flat = None
+            n_real = BT
+
+        if mask_flat is not None and n_real < BT:
+            real_frames = frames_flat[mask_flat]          # [n_real, C, H, W]
+        else:
+            real_frames = frames_flat
+
+        # Chunked encoding to prevent OOM on batches with many real frames
+        if real_frames.shape[0] <= self._MAX_ENCODE_CHUNK:
+            patch_features, kv_cache = self.encoder.encode_patches(real_frames)
+        else:
+            pf_chunks, kv_chunks = [], []
+            for start in range(0, real_frames.shape[0], self._MAX_ENCODE_CHUNK):
+                pf_chunk, kv_chunk = self.encoder.encode_patches(
+                    real_frames[start:start + self._MAX_ENCODE_CHUNK]
+                )
+                pf_chunks.append(pf_chunk)
+                kv_chunks.append(kv_chunk)
+            patch_features = torch.cat(pf_chunks, dim=0)
+            kv_cache = [
+                (torch.cat([c[li][0] for c in kv_chunks], dim=0),
+                 torch.cat([c[li][1] for c in kv_chunks], dim=0))
+                for li in range(len(kv_chunks[0]))
+            ]
+
+        return kv_cache, patch_features, mask_flat
+
+    def _batched_query_attend(self, queries: torch.Tensor, kv_cache: list,
+                              patch_features: torch.Tensor = None) -> torch.Tensor:
+        """Chunked query_attend (deep) or shallow_query_attend to prevent OOM."""
+        n = queries.shape[0]
+        if not self.deep_query:
+            # Shallow mode: single cross-attention on final features
+            if n <= self._MAX_ENCODE_CHUNK:
+                return self.encoder.shallow_query_attend(queries, patch_features)
+            chunks = []
+            for start in range(0, n, self._MAX_ENCODE_CHUNK):
+                end = min(start + self._MAX_ENCODE_CHUNK, n)
+                chunks.append(self.encoder.shallow_query_attend(
+                    queries[start:end], patch_features[start:end]))
+            return torch.cat(chunks, dim=0)
+        # Deep mode: propagate through all DINO layers
+        if n <= self._MAX_ENCODE_CHUNK:
+            return self.encoder.query_attend(queries, kv_cache)
+        chunks = []
+        for start in range(0, n, self._MAX_ENCODE_CHUNK):
+            end = min(start + self._MAX_ENCODE_CHUNK, n)
+            kv_slice = [(K[start:end], V[start:end]) for K, V in kv_cache]
+            chunks.append(self.encoder.query_attend(queries[start:end], kv_slice))
+        return torch.cat(chunks, dim=0)
+
+    def _query_all_frames(
+        self, query: torch.Tensor, kv_cache: list,
+        B: int, T: int, mask_flat=None, patch_features=None,
+    ) -> torch.Tensor:
+        """
+        Apply a single query to every frame in ONE batched query_attend call.
+
+        query          : [B, query_dim]
+        kv_cache       : list of (K, V) per layer, each [n_real, N+1, D]
+        B, T           : batch and temporal dimensions
+        mask_flat      : [B*T] bool or None
+        patch_features : [n_real, N+1, D] (needed for shallow mode)
+        Returns        : [B, T, dino_dim]
+        """
+        BT = B * T
+        dd = self.encoder.dino_dim
+
+        # Expand: same query for all T frames → [B*T, qd]
+        query_exp = query.unsqueeze(1).expand(B, T, -1).reshape(BT, -1)
+
+        if mask_flat is not None:
+            n_real = mask_flat.sum().item()
+            if n_real == 0:
+                return torch.zeros(B, T, dd, device=query.device, dtype=query.dtype)
+            query_real = query_exp[mask_flat]                     # [n_real, qd]
+            z_real = self._batched_query_attend(query_real, kv_cache, patch_features)
+            z_flat = torch.zeros(BT, dd, device=query.device, dtype=z_real.dtype)
+            z_flat[mask_flat] = z_real
+        else:
+            z_flat = self._batched_query_attend(query_exp, kv_cache, patch_features)
+
+        return z_flat.reshape(B, T, dd)
+
+    def _query_all_frames_batched(
+        self, queries: torch.Tensor, kv_cache: list,
+        B: int, T: int, mask_flat=None, patch_features=None,
+    ) -> torch.Tensor:
+        """
+        Apply per-frame queries in ONE batched query_attend call.
+
+        queries        : [B, T, query_dim]
+        kv_cache       : list of (K, V) per layer, each [n_real, N+1, D]
+        B, T           : batch and temporal dimensions
+        mask_flat      : [B*T] bool or None
+        patch_features : [n_real, N+1, D] (needed for shallow mode)
+        Returns        : [B, T, dino_dim]
+        """
+        BT = B * T
+        dd = self.encoder.dino_dim
+        queries_flat = queries.reshape(BT, -1)
+
+        if mask_flat is not None:
+            n_real = mask_flat.sum().item()
+            if n_real == 0:
+                return torch.zeros(B, T, dd, device=queries.device, dtype=queries.dtype)
+            query_real = queries_flat[mask_flat]                   # [n_real, qd]
+            z_real = self._batched_query_attend(query_real, kv_cache, patch_features)
+            z_flat = torch.zeros(BT, dd, device=queries.device, dtype=z_real.dtype)
+            z_flat[mask_flat] = z_real
+        else:
+            z_flat = self._batched_query_attend(queries_flat, kv_cache, patch_features)
+
+        return z_flat.reshape(B, T, dd)
+
+    def _extract_frame_kv(self, kv_cache: list, mask_flat, B: int, T: int, frame_idx: int):
+        """
+        Extract single-frame KV cache from flat format (for autoregressive/eval).
+
+        Returns list of (K, V) per layer, each [B, N+1, D].
+        """
+        if mask_flat is not None:
+            # Scatter compact caches to full [B*T] then extract frame
+            N1 = kv_cache[0][0].shape[1]
+            D = kv_cache[0][0].shape[2]
+            frame_kv = []
+            for K_real, V_real in kv_cache:
+                K_full = torch.zeros(B * T, N1, D, dtype=K_real.dtype, device=K_real.device)
+                V_full = torch.zeros(B * T, N1, D, dtype=V_real.dtype, device=V_real.device)
+                K_full[mask_flat] = K_real
+                V_full[mask_flat] = V_real
+                K_t = K_full.reshape(B, T, N1, D)[:, frame_idx]  # [B, N+1, D]
+                V_t = V_full.reshape(B, T, N1, D)[:, frame_idx]
+                frame_kv.append((K_t, V_t))
+            return frame_kv
+        else:
+            N1 = kv_cache[0][0].shape[1]
+            D = kv_cache[0][0].shape[2]
+            frame_kv = []
+            for K_all, V_all in kv_cache:
+                K_t = K_all.reshape(B, T, N1, D)[:, frame_idx]
+                V_t = V_all.reshape(B, T, N1, D)[:, frame_idx]
+                frame_kv.append((K_t, V_t))
+            return frame_kv
+
+    def _build_causal_mask(self, seq_len: int, device: torch.device) -> torch.Tensor:
+        """
+        Standard causal attention mask [1, 1, S, S] for the LLM.
+        True = masked (cannot attend), False = allowed.
+        """
+        mask = torch.ones(seq_len, seq_len, dtype=torch.bool, device=device).triu(1)
+        return mask.unsqueeze(0).unsqueeze(0)  # [1, 1, S, S]
+
+    def _ce_loss(
+        self,
+        logits: torch.Tensor,
+        labels: torch.Tensor,
+        loss_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Standard autoregressive CE loss with shift-by-1.
+
+        logits    : [B, S, V]   (full sequence logits)
+        labels    : [B, S]      (token ids; positions without loss use pad)
+        loss_mask : [B, S]      (1 = compute loss, 0 = ignore). Applied BEFORE
+                    the shift so that loss_mask[i] guards label[i].
+
+        Returns scalar loss.
+        """
+        # Shift: predict position i+1 from position i
+        shift_logits = logits[:, :-1, :].contiguous()   # [B, S-1, V]
+        shift_labels = labels[:, 1:].contiguous()        # [B, S-1]
+
+        if loss_mask is not None:
+            shift_mask = loss_mask[:, 1:].contiguous()   # [B, S-1]
+            # Replace masked positions with -100 (standard PyTorch ignore_index)
+            shift_labels = shift_labels.clone()
+            shift_labels[shift_mask == 0] = -100
+
+        V = shift_logits.shape[-1]
+        loss = F.cross_entropy(
+            shift_logits.reshape(-1, V),
+            shift_labels.reshape(-1),
+            ignore_index=-100,
+            reduction="mean",
+        )
+        return loss
+
+    def _fused_ce_loss(
+        self,
+        hidden_states: torch.Tensor,
+        labels: torch.Tensor,
+        loss_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Fused lm_head + CE loss via Liger Kernel.
+
+        Never materializes the [B, S, V] logits tensor — computes CE in chunks
+        inside the fused kernel.  Saves ~2× memory on the loss computation.
+
+        hidden_states : [B, S, ld]  (LLM hidden states, NOT yet projected by lm_head)
+        labels        : [B, S]      (token ids)
+        loss_mask     : [B, S]      (1 = compute loss, 0 = ignore)
+
+        Returns scalar loss.
+        """
+        # Shift: predict position i+1 from position i
+        h_input = hidden_states[:, :-1, :].contiguous()  # [B, S-1, ld]
+        shift_labels = labels[:, 1:].contiguous()         # [B, S-1]
+
+        if loss_mask is not None:
+            shift_mask = loss_mask[:, 1:].contiguous()
+            # Replace masked positions with -100 (standard PyTorch ignore_index)
+            shift_labels = shift_labels.clone()
+            shift_labels[shift_mask == 0] = -100
+
+        # Flatten for Liger: [B*(S-1), ld] and [B*(S-1)]
+        BSminus1 = h_input.shape[0] * h_input.shape[1]
+        return LigerFusedLinearCrossEntropyLoss(
+            ignore_index=-100
+        )(
+            h_input.reshape(BSminus1, -1),
+            self.llm.lm_head.weight,
+            shift_labels.reshape(-1),
+        )
+
+    # ------------------------------------------------------------------
+    # Forward mode 1: Coarse+Fine (TRAINING)
+    # ------------------------------------------------------------------
+
+    def forward_coarse_fine(
+        self,
+        frames: torch.Tensor,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        loss_mask: Optional[torch.Tensor] = None,
+        frame_mask: Optional[torch.Tensor] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Two-pass parallel training forward.
+
+        Pass 1 (coarse): q_static -> all frames -> z_coarse -> LLM(visual only) -> queries
+        Pass 2 (fine):   shifted queries -> all frames -> z_fine -> LLM + text -> loss
+
+        Optimization: the coarse LLM pass processes ONLY visual tokens (not text).
+        Because causal attention means visual positions never see text tokens,
+        removing text produces mathematically identical hidden states at visual
+        positions while reducing sequence length from T+S to T (~10-30x shorter).
+
+        Parameters
+        ----------
+        frames         : [B, T, 3, 224, 224]
+        input_ids      : [B, S]  tokenized text (prompt + answer)
+        attention_mask : [B, S]  text attention mask
+        loss_mask      : [B, S]  which tokens contribute to loss (1=yes, 0=no).
+                         If None, all non-pad tokens have loss.
+
+        Returns
+        -------
+        dict with keys: loss, logits, coarse_loss (optional), fine_loss
+        """
+        B, T = frames.shape[:2]
+        S = input_ids.shape[1]
+
+        # ---- Step 0: Encode all frames (DINO, shared across both passes) ----
+        # Use prefetched DINO results if available (from CUDA stream overlap)
+        prefetched = self._get_prefetched_dino()
+        if prefetched is not None:
+            kv_cache, patch_features, mask_flat = prefetched
+        else:
+            kv_cache, patch_features, mask_flat = self._encode_all_frames(frames, frame_mask)
+
+        # ---- Pass 1: Coarse (visual tokens ONLY — text is invisible to them) ----
+        q_static = self.q_static.expand(B, -1)                     # [B, qd]
+        z_coarse = self._query_all_frames(q_static, kv_cache, B, T, mask_flat, patch_features)  # [B,T,dd]
+        z_coarse_llm = self._project_visual(z_coarse)              # [B,T,ld]
+
+        # Coarse LLM: process ONLY visual tokens (T tokens, not T+S).
+        # Causal attention: visual pos i only sees visual pos 0..i, never text.
+        # This is ~30x faster for typical T=8, S=256 batches.
+        out_coarse = self.llm.model(inputs_embeds=z_coarse_llm)
+        h_coarse = out_coarse.last_hidden_state                    # [B,T,ld]
+
+        # Extract dynamic queries from visual positions
+        queries = self.llm_to_query(h_coarse)                      # [B,T,qd]
+
+        # Shift queries: frame t gets query from frame t-1; frame 0 gets q_init
+        q_init = self.q_init.expand(B, 1, -1)                     # [B,1,qd]
+        shifted_queries = torch.cat([q_init, queries[:, :-1]], dim=1)  # [B,T,qd]
+
+        # ---- Pass 2: Fine ----
+        z_fine = self._query_all_frames_batched(shifted_queries, kv_cache, B, T, mask_flat, patch_features)  # [B,T,dd]
+        z_fine_llm = self._project_visual(z_fine)                  # [B,T,ld]
+
+        # Build fine sequence: [visual_fine, text]
+        text_embeds = self._embed_text(input_ids)                  # [B,S,ld]
+        seq_fine = torch.cat([z_fine_llm, text_embeds], dim=1)     # [B,T+S,ld]
+
+        out_fine = self.llm.model(inputs_embeds=seq_fine)
+        h_fine = out_fine.last_hidden_state                        # [B,T+S,ld]
+
+        # ---- Loss on text portion ----
+        h_text = h_fine[:, T:, :]                                  # [B,S,ld]
+        if loss_mask is None:
+            loss_mask = attention_mask.float()
+
+        if self.use_fused_ce:
+            # Liger Kernel: fused lm_head + CE, never materializes [B,S,V] logits
+            fine_loss = self._fused_ce_loss(h_text, input_ids, loss_mask)
+            logits_text = None  # not available with fused loss
+        else:
+            logits_text = self.llm.lm_head(h_text)                # [B,S,V]
+            fine_loss = self._ce_loss(logits_text, input_ids, loss_mask)
+
+        # ---- Optional auxiliary coarse loss ----
+        coarse_loss = torch.tensor(0.0, device=frames.device)
+        if self.lambda_coarse > 0:
+            seq_coarse_full = torch.cat([z_coarse_llm, text_embeds], dim=1)
+            out_coarse_full = self.llm.model(inputs_embeds=seq_coarse_full)
+            h_coarse_text = out_coarse_full.last_hidden_state[:, T:, :]
+            if self.use_fused_ce:
+                coarse_loss = self._fused_ce_loss(h_coarse_text, input_ids, loss_mask)
+            else:
+                logits_coarse = self.llm.lm_head(h_coarse_text)
+                coarse_loss = self._ce_loss(logits_coarse, input_ids, loss_mask)
+
+        # ---- Combined loss ----
+        loss = fine_loss + self.lambda_coarse * coarse_loss
+
+        return {
+            "loss": loss,
+            "fine_loss": fine_loss,
+            "coarse_loss": coarse_loss,
+            "logits": logits_text,  # [B,S,V] text positions only
+        }
+
+    # ------------------------------------------------------------------
+    # Forward mode: DPO (preference training)
+    # ------------------------------------------------------------------
+
+    def forward_dpo(
+        self,
+        frames: torch.Tensor,
+        chosen_input_ids: torch.Tensor,
+        chosen_attention_mask: torch.Tensor,
+        chosen_loss_mask: torch.Tensor,
+        rejected_input_ids: torch.Tensor,
+        rejected_attention_mask: torch.Tensor,
+        rejected_loss_mask: torch.Tensor,
+        frame_mask: Optional[torch.Tensor] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        DPO forward pass: run coarse+fine on both chosen and rejected sequences.
+
+        Shares DINO encoding across chosen and rejected (same visual input).
+        Returns per-sample sum of log-probabilities for both chosen and rejected,
+        masked by loss_mask (answer-only tokens).
+
+        Parameters
+        ----------
+        frames                  : [B, T, 3, 224, 224]
+        chosen_input_ids        : [B, S_c]
+        chosen_attention_mask   : [B, S_c]
+        chosen_loss_mask        : [B, S_c]  (1 = answer token, 0 = prompt/pad)
+        rejected_input_ids      : [B, S_r]
+        rejected_attention_mask : [B, S_r]
+        rejected_loss_mask      : [B, S_r]
+        frame_mask              : [B, T] bool (optional)
+
+        Returns
+        -------
+        dict with keys:
+          chosen_logps    : [B]  per-sample sum of log-probs on chosen answer tokens
+          rejected_logps  : [B]  per-sample sum of log-probs on rejected answer tokens
+          chosen_logits   : [B, T+S_c, V]  full logits for chosen
+          rejected_logits : [B, T+S_r, V]  full logits for rejected
+        """
+        B, T = frames.shape[:2]
+
+        # ---- Step 0: Encode all frames (DINO, shared across chosen & rejected) ----
+        kv_cache, patch_features, mask_flat = self._encode_all_frames(frames, frame_mask)
+
+        # ---- Coarse pass (visual tokens ONLY — text invisible in causal attn) ----
+        q_static = self.q_static.expand(B, -1)                          # [B, qd]
+        z_coarse = self._query_all_frames(q_static, kv_cache, B, T, mask_flat, patch_features)
+        z_coarse_llm = self._project_visual(z_coarse)                    # [B, T, ld]
+
+        # Coarse LLM: visual tokens only (T, not T+S_c). Causal attention means
+        # visual positions never see text, so this is mathematically identical.
+        out_coarse = self.llm.model(inputs_embeds=z_coarse_llm)
+        h_coarse = out_coarse.last_hidden_state                          # [B, T, ld]
+
+        # Extract dynamic queries from visual positions
+        queries = self.llm_to_query(h_coarse)                            # [B, T, qd]
+
+        q_init = self.q_init.expand(B, 1, -1)
+        shifted_queries = torch.cat([q_init, queries[:, :-1]], dim=1)    # [B, T, qd]
+
+        # ---- Fine pass: shared visual features ----
+        z_fine = self._query_all_frames_batched(shifted_queries, kv_cache, B, T, mask_flat, patch_features)
+        z_fine_llm = self._project_visual(z_fine)                        # [B, T, ld]
+
+        # ---- Forward on CHOSEN (lm_head on text positions only) ----
+        text_embeds_chosen = self._embed_text(chosen_input_ids)          # [B, S_c, ld]
+        seq_chosen = torch.cat([z_fine_llm, text_embeds_chosen], dim=1)  # [B, T+S_c, ld]
+        out_chosen = self.llm.model(inputs_embeds=seq_chosen)
+        chosen_logits = self.llm.lm_head(out_chosen.last_hidden_state[:, T:, :])  # [B, S_c, V]
+
+        # ---- Forward on REJECTED (lm_head on text positions only) ----
+        text_embeds_rejected = self._embed_text(rejected_input_ids)      # [B, S_r, ld]
+        seq_rejected = torch.cat([z_fine_llm, text_embeds_rejected], dim=1)
+        out_rejected = self.llm.model(inputs_embeds=seq_rejected)
+        rejected_logits = self.llm.lm_head(out_rejected.last_hidden_state[:, T:, :])  # [B, S_r, V]
+
+        # ---- Compute per-token log-probs ----
+        chosen_logps = self._sequence_logprobs(
+            chosen_logits, chosen_input_ids, chosen_loss_mask,
+        )
+        rejected_logps = self._sequence_logprobs(
+            rejected_logits, rejected_input_ids, rejected_loss_mask,
+        )
+
+        return {
+            "chosen_logps": chosen_logps,       # [B]
+            "rejected_logps": rejected_logps,   # [B]
+            "chosen_logits": chosen_logits,     # [B, S_c, V]
+            "rejected_logits": rejected_logits, # [B, S_r, V]
+        }
+
+    def _sequence_logprobs(
+        self,
+        logits: torch.Tensor,
+        input_ids: torch.Tensor,
+        loss_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Compute per-sample sum of log-probabilities on answer tokens.
+
+        logits    : [B, S, V]  text-only logits (visual positions excluded)
+        input_ids : [B, S]     text token ids
+        loss_mask : [B, S]     1.0 for answer tokens, 0.0 otherwise
+
+        Returns   : [B]        sum of log-probs per sample
+        """
+        B, S = input_ids.shape
+
+        # Shift for autoregressive prediction
+        shift_logits = logits[:, :-1, :]                              # [B, S-1, V]
+        shift_labels = input_ids[:, 1:]                               # [B, S-1]
+        shift_mask = loss_mask[:, 1:]                                 # [B, S-1]
+
+        # Per-token log-probs: log_softmax then gather the label's prob
+        log_probs = F.log_softmax(shift_logits, dim=-1)              # [B, S-1, V]
+        per_token_logps = log_probs.gather(
+            dim=-1, index=shift_labels.unsqueeze(-1),
+        ).squeeze(-1)                                                 # [B, S-1]
+
+        # Mask and sum per sample
+        per_token_logps = per_token_logps * shift_mask                # zero out non-answer tokens
+        return per_token_logps.sum(dim=-1)                            # [B]
+
+    # ------------------------------------------------------------------
+    # Forward mode 2: Coarse only (FAST EVAL)
+    # ------------------------------------------------------------------
+
+    def forward_coarse_only(
+        self,
+        frames: torch.Tensor,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        loss_mask: Optional[torch.Tensor] = None,
+        frame_mask: Optional[torch.Tensor] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Single-pass coarse forward (q_static only, no fine queries).
+
+        Used for:
+          - Training A6 ablation (coarse-only training)
+          - Fast eval (wrap in torch.no_grad() externally)
+
+        q_static -> all frames -> z_coarse -> LLM -> logits.
+
+        Parameters
+        ----------
+        frames         : [B, T, 3, 224, 224]
+        input_ids      : [B, S]  (optional, for loss computation)
+        attention_mask : [B, S]  (optional)
+        loss_mask      : [B, S]  (optional)
+
+        Returns
+        -------
+        dict with keys: logits, and optionally loss
+        """
+        B, T = frames.shape[:2]
+
+        kv_cache, patch_features, mask_flat = self._encode_all_frames(frames, frame_mask)
+
+        q_static = self.q_static.expand(B, -1)
+        z_coarse = self._query_all_frames(q_static, kv_cache, B, T, mask_flat, patch_features)
+        z_coarse_llm = self._project_visual(z_coarse)
+
+        if input_ids is not None:
+            text_embeds = self._embed_text(input_ids)
+            seq = torch.cat([z_coarse_llm, text_embeds], dim=1)
+        else:
+            seq = z_coarse_llm
+        # dtype handled by autocast on GPU; float32 on CPU
+
+        out = self.llm.model(inputs_embeds=seq)
+        h = out.last_hidden_state                         # [B, T+S, ld]
+
+        if input_ids is not None:
+            S = input_ids.shape[1]
+            pad_id = self._get_pad_token_id()
+            visual_pad = torch.full(
+                (B, T), pad_id, dtype=input_ids.dtype, device=input_ids.device,
+            )
+            full_labels = torch.cat([visual_pad, input_ids], dim=1)
+
+            if loss_mask is not None:
+                visual_no_loss = torch.zeros(
+                    B, T, dtype=loss_mask.dtype, device=loss_mask.device,
+                )
+                full_loss_mask = torch.cat([visual_no_loss, loss_mask], dim=1)
+            elif attention_mask is not None:
+                visual_no_loss = torch.zeros(
+                    B, T, dtype=attention_mask.dtype, device=attention_mask.device,
+                )
+                full_loss_mask = torch.cat([visual_no_loss, attention_mask], dim=1)
+            else:
+                full_loss_mask = None
+
+            if self.use_fused_ce and self.training:
+                # Fused CE: skip lm_head, never materializes [B, T+S, V]
+                loss = self._fused_ce_loss(h, full_labels, full_loss_mask)
+                logits = None
+            else:
+                logits = self.llm.lm_head(h)
+                loss = self._ce_loss(logits, full_labels, full_loss_mask)
+
+            result: Dict[str, torch.Tensor] = {"logits": logits, "loss": loss}
+            result["coarse_loss"] = loss
+            result["fine_loss"] = torch.tensor(0.0, device=frames.device)
+        else:
+            logits = self.llm.lm_head(h)
+            result: Dict[str, torch.Tensor] = {"logits": logits}
+
+        return result
+
+    # ------------------------------------------------------------------
+    # Forward mode 3: Autoregressive (TRUE INFERENCE)
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def forward_autoregressive(
+        self,
+        frames: torch.Tensor,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        loss_mask: Optional[torch.Tensor] = None,
+        frame_mask: Optional[torch.Tensor] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        True autoregressive inference: sequential frame-by-frame with KV cache.
+
+        q_init -> frame_1 -> z_1 -> LLM -> q_1 -> frame_2 -> z_2 -> ...
+
+        No coarse pass. Each query is derived from the LLM hidden state after
+        processing the *previous* fine visual token -- exactly what happens at
+        real inference time.
+
+        Parameters
+        ----------
+        frames         : [B, T, 3, 224, 224]
+        input_ids      : [B, S]  (optional, for loss computation)
+        attention_mask : [B, S]  (optional)
+        loss_mask      : [B, S]  (optional)
+
+        Returns
+        -------
+        dict with keys: logits, and optionally loss
+        """
+        B, T = frames.shape[:2]
+        device = frames.device
+
+        # Encode all frames with DINO up front (this is OK -- DINO encoding
+        # does not depend on the query, only query_attend does).
+        kv_cache, patch_features, mask_flat = self._encode_all_frames(frames, frame_mask)
+
+        # Enable KV cache on the LLM for incremental decoding
+        orig_use_cache = self.llm.config.use_cache
+        self.llm.config.use_cache = True
+
+        query = self.q_init.expand(B, -1)    # [B, qd]
+        llm_past_kv = None
+
+        for t in range(T):
+            # Foveated extraction with current query
+            frame_kv = self._extract_frame_kv(kv_cache, mask_flat, B, T, t)
+            z_t = self.encoder.query_attend(query, frame_kv)  # [B, dd]
+            z_t_llm = self._project_visual(z_t.unsqueeze(1))            # [B,1,ld]
+            # dtype handled by autocast on GPU; float32 on CPU
+
+            # Incremental LLM forward (one visual token at a time)
+            out = self.llm.model(
+                inputs_embeds=z_t_llm,
+                past_key_values=llm_past_kv,
+                use_cache=True,
+            )
+            llm_past_kv = out.past_key_values
+
+            # Derive query for the NEXT frame from the current hidden state
+            if t < T - 1:
+                h_t = out.last_hidden_state[:, -1, :]   # [B, ld]
+                query = self.llm_to_query(h_t)                   # [B, qd]
+
+        # ---- Now process text (if provided) using the accumulated KV cache ----
+        if input_ids is not None:
+            text_embeds = self._embed_text(input_ids)  # [B, S, ld]
+
+            out_text = self.llm.model(
+                inputs_embeds=text_embeds,
+                past_key_values=llm_past_kv,
+                use_cache=False,
+            )
+            # Combine visual hidden states (already in KV cache) with text states
+            # for logit computation. We only need logits over the text portion
+            # (plus the last visual token which predicts the first text token).
+            #
+            # The KV cache holds T visual positions; out_text.last_hidden_state
+            # holds S text positions.  We reconstruct the full logits as
+            # [visual_logits, text_logits] but only compute loss on text.
+            h_text = out_text.last_hidden_state         # [B, S, ld]
+            logits_text = self.llm.lm_head(h_text)      # [B, S, V]
+
+            # For the loss we also need the logit at the last visual position
+            # (it predicts the first text token).  Re-derive it:
+            h_last_visual = out.last_hidden_state[:, -1:, :]   # [B,1,ld]
+            logits_last_v = self.llm.lm_head(h_last_visual)    # [B,1,V]
+
+            # Full logits over [last_visual, text] = [B, 1+S, V]
+            logits = torch.cat([logits_last_v, logits_text], dim=1)
+
+            # Labels: [pad_for_last_visual, input_ids]
+            pad_id = self._get_pad_token_id()
+            lv_pad = torch.full(
+                (B, 1), pad_id, dtype=input_ids.dtype, device=device,
+            )
+            full_labels = torch.cat([lv_pad, input_ids], dim=1)
+
+            # Loss mask
+            if loss_mask is not None:
+                lv_no_loss = torch.zeros(
+                    B, 1, dtype=loss_mask.dtype, device=device,
+                )
+                full_loss_mask = torch.cat([lv_no_loss, loss_mask], dim=1)
+            elif attention_mask is not None:
+                lv_no_loss = torch.zeros(
+                    B, 1, dtype=attention_mask.dtype, device=device,
+                )
+                full_loss_mask = torch.cat([lv_no_loss, attention_mask], dim=1)
+            else:
+                full_loss_mask = None
+
+            loss = self._ce_loss(logits, full_labels, full_loss_mask)
+
+            self.llm.config.use_cache = orig_use_cache
+            return {"loss": loss, "logits": logits}
+
+        else:
+            # No text -- just return logits at the last visual position
+            h_last = out.last_hidden_state   # [B, 1, ld]
+            logits = self.llm.lm_head(h_last)
+            self.llm.config.use_cache = orig_use_cache
+            return {"logits": logits}
+
+    # ------------------------------------------------------------------
+    # Convenience: unified forward dispatching by name
+    # ------------------------------------------------------------------
+
+    def forward(
+        self,
+        frames: torch.Tensor,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        loss_mask: Optional[torch.Tensor] = None,
+        frame_mask: Optional[torch.Tensor] = None,
+        mode: str = "coarse_fine",
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Unified forward entry point.
+
+        Parameters
+        ----------
+        frames : Tensor [B, T, 3, 224, 224]
+            Preprocessed video frames (DINOv2 normalization).
+            For **video**: T = number of sampled frames (1-64).
+            For **images**: replicate the single frame to T=8 to match training
+            distribution (``frame.unsqueeze(0).repeat(8, 1, 1, 1)``).
+            The model was trained with ``replicate_image_frames: 8`` in
+            Stages 2-3, so single-frame image input will produce degraded
+            results.
+        input_ids : Tensor [B, S]
+            Tokenized text (prompt + response).
+        attention_mask : Tensor [B, S]
+            1 for real tokens, 0 for padding.
+        loss_mask : Tensor [B, S], optional
+            1 for tokens that contribute to loss, 0 to skip.
+        frame_mask : Tensor [B, T] bool, optional
+            True for real frames, False for padding (for variable-length batches).
+        mode : str
+            "coarse_fine"    — two-pass parallel forward (recommended, uses foveation)
+            "coarse_only"    — single static-query pass (fastest, no foveation)
+            "autoregressive" — sequential inference with KV cache
+        """
+        if mode == "coarse_fine":
+            return self.forward_coarse_fine(frames, input_ids, attention_mask, loss_mask, frame_mask)
+        elif mode == "coarse_only":
+            return self.forward_coarse_only(frames, input_ids, attention_mask, loss_mask, frame_mask)
+        elif mode == "autoregressive":
+            return self.forward_autoregressive(frames, input_ids, attention_mask, loss_mask, frame_mask)
+        else:
+            raise ValueError(
+                f"Unknown forward mode '{mode}'. "
+                "Expected one of: coarse_fine, coarse_only, autoregressive"
+            )
+
+    # ------------------------------------------------------------------
+    # CUDA Stream Prefetch — overlap DINO encoding with LLM backward
+    # ------------------------------------------------------------------
+
+    def prefetch_dino(self, frames: torch.Tensor, frame_mask=None, stream=None):
+        """
+        Start DINO encoding on a separate CUDA stream.
+
+        Call this while the previous batch's backward pass is running.
+        The DINO encoder is frozen during training, so there's no gradient
+        dependency between the backward pass and this prefetch.
+
+        Args:
+            frames: [B, T, 3, 224, 224] next batch's frames
+            frame_mask: [B, T] bool, optional
+            stream: torch.cuda.Stream to run on (caller manages lifecycle)
+
+        Returns:
+            None — results are stored internally and retrieved via
+            forward_coarse_fine(..., prefetched_dino=True).
+        """
+        if stream is None:
+            stream = torch.cuda.Stream()
+        with torch.cuda.stream(stream):
+            with torch.no_grad():
+                self._prefetched_dino = self._encode_all_frames(frames, frame_mask)
+                self._prefetch_stream = stream
+
+    def _get_prefetched_dino(self):
+        """Retrieve and clear prefetched DINO results, synchronizing the stream."""
+        if hasattr(self, '_prefetched_dino') and self._prefetched_dino is not None:
+            self._prefetch_stream.synchronize()
+            result = self._prefetched_dino
+            self._prefetched_dino = None
+            self._prefetch_stream = None
+            return result
+        return None
+
+    # ------------------------------------------------------------------
+    # Utility methods for external callers (train.py, eval.py)
+    # ------------------------------------------------------------------
+
+    def enable_gradient_checkpointing(
+        self, llm_only: bool = False, use_reentrant: bool = True,
+    ) -> None:
+        """Turn on activation checkpointing for LLM (and optionally DINO).
+
+        Args:
+            llm_only: If True, only enable for LLM backbone. Leave DINO
+                      un-checkpointed so it can be safely torch.compiled.
+                      DINO is small (22M params) so checkpointing saves
+                      little memory there.
+            use_reentrant: If False, use non-reentrant checkpointing which
+                          is compatible with torch.compile (the reentrant
+                          version causes NaN with compile). Default True
+                          for backward compat; set False when using compile.
+        """
+        ckpt_kwargs = {"use_reentrant": use_reentrant}
+        self.llm.gradient_checkpointing_enable(
+            gradient_checkpointing_kwargs=ckpt_kwargs
+        )
+        if not llm_only and hasattr(self.encoder.dino, 'gradient_checkpointing_enable'):
+            self.encoder.dino.gradient_checkpointing_enable(
+                gradient_checkpointing_kwargs=ckpt_kwargs
+            )
+
+    def get_param_groups(
+        self,
+        lr_backbone: float = 1e-5,
+        lr_connector: float = 1e-4,
+    ) -> list:
+        """
+        Return parameter groups with differential learning rates.
+
+        Groups:
+          1. Connector (dino_to_llm, llm_to_query, q_static, q_init) -- highest LR
+          2. DINO encoder -- backbone LR
+          3. LLM -- backbone LR
+
+        This is a suggestion; train.py may override.
+        """
+        connector_params = set()
+        for name, param in self.named_parameters():
+            if any(k in name for k in [
+                "dino_to_llm", "llm_to_query", "q_static", "q_init",
+                "query_input_proj", "query_output_proj",
+            ]):
+                connector_params.add(id(param))
+
+        encoder_params = set()
+        for name, param in self.encoder.named_parameters():
+            if id(param) not in connector_params:
+                encoder_params.add(id(param))
+
+        groups = [
+            {
+                "params": [p for p in self.parameters()
+                           if id(p) in connector_params and p.requires_grad],
+                "lr": lr_connector,
+                "name": "connector",
+            },
+            {
+                "params": [p for n, p in self.encoder.named_parameters()
+                           if id(p) in encoder_params and p.requires_grad],
+                "lr": lr_backbone,
+                "name": "dino",
+            },
+            {
+                "params": [p for p in self.llm.parameters() if p.requires_grad],
+                "lr": lr_backbone,
+                "name": "llm",
+            },
+        ]
+        return [g for g in groups if len(g["params"]) > 0]