cosmos_framework/model/tokenizer/models/text_decoder.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

"""Text decoder modules for image-to-text generation in the tokenizer.

Provides:
    - ImagePositionEmbeddings: Learnable 2D spatial position embeddings
    - SpatialPatchMerger: 2x2 spatial-to-channel merger for image token reduction
    - TextDecoderWrapper: causal LM wrapper with vision token injection

Architecture flow (mid-training / ITG path):
    x_no_proj [N, encoder_dim]
        -> SpatialPatchMerger (2x2 concat + MLP) -> [N/4, decoder_hidden]
        -> ImagePositionEmbeddings (add 2D spatial pos) -> [N/4, decoder_hidden]
        -> inject into text token embeddings at image_patch_indices
        -> causal LM -> logits [B, S, vocab_size]
        -> CrossEntropy ITG loss (ignore_index=-100 masks vision/padding tokens)

Label alignment (next-token prediction):
    input_ids and labels have the same length (no pre-shift).
    Vision token targets are -100, so loss only backprops through caption tokens.
    The shift (logits[:-1] vs labels[1:]) happens in the loss function.
"""

import os
from dataclasses import dataclass
from types import SimpleNamespace
from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F
from loguru import logger as logging

from cosmos_framework.model.attention import attention as i4_attention
from cosmos_framework.model.attention.masks import CausalType
from cosmos_framework.model.attention.natten import NATTEN_SUPPORTED, get_natten_version
from cosmos_framework.model.tokenizer.utils.hf import (
    load_auto_tokenizer_from_cache,
    prepare_nemotron_tokenizer_snapshot,
    resolve_hf_snapshot_path,
)
from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs, stack_with_bounded_inputs
from cosmos_framework.model.tokenizer.utils.vlm_prompt_format import densevl_add_vision_id_text

QWEN3_PAD_TOKEN_ID = 151643
QWEN3_IM_START_TOKEN_ID = 151644
QWEN3_IM_END_TOKEN_ID = 151645
QWEN3_VISION_START_TOKEN_ID = 151652
QWEN3_VISION_END_TOKEN_ID = 151653
QWEN3_VISION_PAD_TOKEN_ID = 151655
QWEN3_THINK_START_TOKEN_ID = 151667
QWEN3_THINK_END_TOKEN_ID = 151668
NEMOTRON_2B_PAD_TOKEN_ID = 0
NEMOTRON_2B_EOS_TOKEN_ID = 11
NEMOTRON_2B_VISION_START_TOKEN_ID = 20
NEMOTRON_2B_VISION_END_TOKEN_ID = 21
NEMOTRON_2B_VISION_PAD_TOKEN_ID = 22
NEMOTRON_2B_IM_START_TOKEN = "<|im_start|>"
NEMOTRON_2B_IM_END_TOKEN = "<|im_end|>"
NEMOTRON_2B_THINK_START_TOKEN = "<think>"
NEMOTRON_2B_THINK_END_TOKEN = "</think>"
TEXT_DECODER_ATTN_IMPLEMENTATION_ENV = "TOKENIZER_TEXT_DECODER_ATTN_IMPLEMENTATION"
PACKED_ATTENTION_BACKEND_SDPA = "sdpa"
PACKED_ATTENTION_BACKEND_NATTEN = "natten"
PACKED_ATTENTION_BACKENDS = frozenset({PACKED_ATTENTION_BACKEND_SDPA, PACKED_ATTENTION_BACKEND_NATTEN})
VQA_THINKING_MODE_OFF = "off"
VQA_THINKING_MODE_ON = "on"
VQA_THINKING_MODE_RAW = "raw"
VQA_THINKING_MODES = frozenset({VQA_THINKING_MODE_OFF, VQA_THINKING_MODE_ON, VQA_THINKING_MODE_RAW})
VQA_REASONING_SUFFIX = (
    "\nAnswer the question using the following format:\n\n"
    "<think>\nYour reasoning.\n</think>\n\n"
    "Write your final answer immediately after the </think> tag."
)


def normalize_vqa_thinking_mode(thinking_mode: str | None) -> str:
    """Normalize the VQA generation thinking-mode knob."""
    normalized = VQA_THINKING_MODE_OFF if thinking_mode is None else str(thinking_mode).strip().lower()
    if normalized not in VQA_THINKING_MODES:
        raise ValueError(
            f"Unsupported VQA thinking_mode={thinking_mode!r}; expected one of {sorted(VQA_THINKING_MODES)}."
        )
    return normalized


def _append_vqa_reasoning_suffix(question: str, reasoning_suffix: str) -> str:
    """Append the training-compatible reasoning-format suffix once."""
    if not reasoning_suffix:
        return question
    if reasoning_suffix.strip() in question:
        return question
    return f"{question.rstrip()}{reasoning_suffix}"


def _rotate_half_for_nemotron(hidden_states: torch.Tensor) -> torch.Tensor:
    """Rotate the two RoPE feature halves using Nemotron's convention."""
    midpoint = hidden_states.shape[-1] // 2
    first_half = hidden_states[..., :midpoint]  # [B,H,T,D/2]
    second_half = hidden_states[..., midpoint:]  # [B,H,T,D/2]
    rotated = torch.cat((-second_half, first_half), dim=-1)  # [B,H,T,D]
    return rotated


def _apply_nemotron_rotary_embeddings(
    query_states: torch.Tensor,
    key_states: torch.Tensor,
    cosine: torch.Tensor,
    sine: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Apply the remote Nemotron model's RoPE formulation to Q and K."""
    rotated_query = _rotate_half_for_nemotron(query_states)  # [B,Hq,T,D]
    rotated_key = _rotate_half_for_nemotron(key_states)  # [B,Hkv,T,D]
    query_states_ = query_states * cosine + rotated_query * sine  # [B,Hq,T,D]
    key_states_ = key_states * cosine + rotated_key * sine  # [B,Hkv,T,D]
    return query_states_, key_states_


def _forward_nemotron_natten_attention(
    attention_module: nn.Module,
    hidden_states: torch.Tensor,
    position_ids: torch.Tensor,
    cumulative_seqlen: torch.Tensor,
    max_seqlen: int,
    rotary_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
    """Run one Nemotron self-attention block with NATTEN causal varlen attention."""
    batch_size, sequence_length, _ = hidden_states.shape
    if batch_size != 1:
        raise ValueError(f"NATTEN varlen requires packed batch size 1, got {batch_size}.")

    query_states = attention_module.q_proj(hidden_states)  # [1,T,Hq*D]
    query_states = query_states.view(
        batch_size,
        sequence_length,
        attention_module.num_heads,
        attention_module.head_dim,
    ).transpose(1, 2)  # [1,Hq,T,D]
    key_states = attention_module.k_proj(hidden_states)  # [1,T,Hkv*D]
    key_states = key_states.view(
        batch_size,
        sequence_length,
        attention_module.num_key_value_heads,
        attention_module.head_dim,
    ).transpose(1, 2)  # [1,Hkv,T,D]
    value_states = attention_module.v_proj(hidden_states)  # [1,T,Hkv*D]
    value_states = value_states.view(
        batch_size,
        sequence_length,
        attention_module.num_key_value_heads,
        attention_module.head_dim,
    ).transpose(1, 2)  # [1,Hkv,T,D]

    if rotary_embeddings is None:
        normalized_position_ids = position_ids.reshape(batch_size, sequence_length).long()  # [1,T]
        cosine, sine = attention_module.rotary_emb(
            value_states,
            normalized_position_ids,
        )  # each [1,1,T,D]
    else:
        cosine, sine = rotary_embeddings
    query_states_, key_states_ = _apply_nemotron_rotary_embeddings(
        query_states,
        key_states,
        cosine,
        sine,
    )  # [1,Hq,T,D], [1,Hkv,T,D]

    query_states_ = query_states_.transpose(1, 2).contiguous()  # [1,T,Hq,D]
    key_states_ = key_states_.transpose(1, 2).contiguous()  # [1,T,Hkv,D]
    value_states = value_states.transpose(1, 2).contiguous()  # [1,T,Hkv,D]
    attention_output = i4_attention(
        query=query_states_,
        key=key_states_,
        value=value_states,
        is_causal=True,
        causal_type=CausalType.DontCare,
        cumulative_seqlen_Q=cumulative_seqlen,
        cumulative_seqlen_KV=cumulative_seqlen,
        max_seqlen_Q=max_seqlen,
        max_seqlen_KV=max_seqlen,
        backend=PACKED_ATTENTION_BACKEND_NATTEN,
    )
    if not isinstance(attention_output, torch.Tensor):
        raise TypeError("NATTEN attention unexpectedly returned log-sum-exp output.")
    attention_output = attention_output.reshape(batch_size, sequence_length, attention_module.hidden_size)  # [1,T,D]
    projected_output = attention_module.o_proj(attention_output)  # [1,T,D]
    return projected_output


def _forward_nemotron_natten_decoder_layer(
    decoder_layer: nn.Module,
    hidden_states: torch.Tensor,
    position_ids: torch.Tensor,
    cumulative_seqlen: torch.Tensor,
    max_seqlen: int,
    rotary_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
    """Run one remote-code Nemotron layer while replacing only self-attention."""
    residual = hidden_states  # [1,T,D]
    normalized_hidden_states = decoder_layer.input_layernorm(hidden_states)  # [1,T,D]
    attention_output = _forward_nemotron_natten_attention(
        decoder_layer.self_attn,
        normalized_hidden_states,
        position_ids,
        cumulative_seqlen,
        max_seqlen,
        rotary_embeddings,
    )  # [1,T,D]
    hidden_states = residual + attention_output  # [1,T,D]

    residual = hidden_states  # [1,T,D]
    normalized_hidden_states = decoder_layer.post_attention_layernorm(hidden_states)  # [1,T,D]
    mlp_output = decoder_layer.mlp(normalized_hidden_states)  # [1,T,D]
    return residual + mlp_output  # [1,T,D]


@dataclass(frozen=True)
class TextDecoderFamilySpec:
    """Static configuration for one supported text decoder family."""

    family: str
    default_model_name: str
    pad_token_id: int
    eos_token_ids: tuple[int, ...]
    vision_start_token_id: int
    vision_end_token_id: int
    vision_pad_token_id: int
    suppress_token_ids: tuple[int, ...] = ()
    trust_remote_code: bool = True
    supports_inputs_embeds_forward: bool = True
    supports_inputs_embeds_generate: bool = True
    supports_cache_position: bool = True


QWEN3_SPEC = TextDecoderFamilySpec(
    family="qwen3",
    default_model_name="Qwen/Qwen3-0.6B",
    pad_token_id=QWEN3_PAD_TOKEN_ID,
    eos_token_ids=(QWEN3_IM_END_TOKEN_ID,),
    vision_start_token_id=QWEN3_VISION_START_TOKEN_ID,
    vision_end_token_id=QWEN3_VISION_END_TOKEN_ID,
    vision_pad_token_id=QWEN3_VISION_PAD_TOKEN_ID,
    suppress_token_ids=(QWEN3_THINK_START_TOKEN_ID,),
)

NEMOTRON_2B_SPEC = TextDecoderFamilySpec(
    family="nemotron_2b",
    default_model_name="nvidia/NVIDIA-Nemotron-3-2B-BF16",
    pad_token_id=NEMOTRON_2B_PAD_TOKEN_ID,
    eos_token_ids=(NEMOTRON_2B_EOS_TOKEN_ID,),
    vision_start_token_id=NEMOTRON_2B_VISION_START_TOKEN_ID,
    vision_end_token_id=NEMOTRON_2B_VISION_END_TOKEN_ID,
    vision_pad_token_id=NEMOTRON_2B_VISION_PAD_TOKEN_ID,
    trust_remote_code=True,
    supports_inputs_embeds_forward=False,
    supports_inputs_embeds_generate=False,
    supports_cache_position=False,
)

TEXT_DECODER_FAMILY_SPECS: dict[str, TextDecoderFamilySpec] = {
    QWEN3_SPEC.family: QWEN3_SPEC,
    NEMOTRON_2B_SPEC.family: NEMOTRON_2B_SPEC,
}


def _resolve_attn_implementation(attn_implementation: str) -> str:
    """Resolve the attention backend, allowing an env override for eval flows."""
    env_override = os.environ.get(TEXT_DECODER_ATTN_IMPLEMENTATION_ENV)
    if env_override:
        return env_override
    return attn_implementation


def _is_flash_attention_error(exc: Exception) -> bool:
    """Return whether one model-load failure is caused by FlashAttention selection."""
    error_text = str(exc).lower()
    return "flash attention" in error_text or "flash_attention" in error_text


def _repair_nemotron_rotary_buffers(module: nn.Module) -> int:
    """Reset Nemotron RoPE inv_freq buffers after low-memory HF loading.

    Nemotron's remote-code rotary module registers ``inv_freq`` as a
    non-persistent buffer. With ``low_cpu_mem_usage=True`` those buffers are not
    restored from the checkpoint and can contain uninitialized values.
    """
    repaired_count = 0
    for child_module in module.modules():
        if (
            not hasattr(child_module, "inv_freq")
            or not hasattr(child_module, "dim")
            or not hasattr(child_module, "base")
        ):
            continue
        old_inv_freq = child_module.inv_freq
        if isinstance(old_inv_freq, torch.Tensor) and old_inv_freq.device.type != "meta":
            device = old_inv_freq.device
        else:
            device = torch.device("cpu")
        dim = int(child_module.dim)
        base = float(child_module.base)
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim))  # [D/2]
        child_module.register_buffer("inv_freq", inv_freq, persistent=False)
        repaired_count += 1
    return repaired_count


def infer_text_decoder_family(model_name: str) -> str:
    """Infer the decoder family from a HuggingFace model name."""
    model_name_lower = model_name.lower()
    if "nemotron" in model_name_lower:
        return NEMOTRON_2B_SPEC.family
    return QWEN3_SPEC.family


def _get_required_token_id(tokenizer: Any, token: str) -> int:
    """Resolve one tokenizer token name into an ID and fail fast if missing."""
    token_id = tokenizer.convert_tokens_to_ids(token)
    if token_id is None:
        raise ValueError(f"Tokenizer does not define required token {token!r}.")
    unk_token_id = getattr(tokenizer, "unk_token_id", None)
    if unk_token_id is not None and int(token_id) == int(unk_token_id):
        token_string = getattr(tokenizer, "unk_token", None)
        raise ValueError(f"Tokenizer maps required token {token!r} to unknown token {token_string!r}.")
    return int(token_id)


def get_text_decoder_family_spec(
    family: str | None = None,
    model_name: str | None = None,
) -> TextDecoderFamilySpec:
    """Resolve one supported text decoder family specification."""
    resolved_family = family or infer_text_decoder_family(model_name or QWEN3_SPEC.default_model_name)
    if resolved_family not in TEXT_DECODER_FAMILY_SPECS:
        raise ValueError(f"Unsupported text decoder family: {resolved_family}")
    return TEXT_DECODER_FAMILY_SPECS[resolved_family]


class ImagePositionEmbeddings(nn.Module):
    """Learnable position embeddings added per coordinate dimension.

    For the text decoder, we use coord_dim=2 for (H, W) after spatial merging.
    Each dimension has its own embedding table; embeddings are summed and added to features.

    Args:
        hidden_size: Feature dimension (must match input features).
        max_position: Maximum coordinate value in any dimension.
        coord_dim: Number of coordinate dimensions (default 2 for H, W).
    """

    def __init__(self, hidden_size: int, max_position: int = 128, coord_dim: int = 2):
        super().__init__()
        self.hidden_size = hidden_size
        self.max_position = max_position
        self.coord_dim = coord_dim

        self.position_embeddings = nn.ModuleList([nn.Embedding(max_position, hidden_size) for _ in range(coord_dim)])
        for emb in self.position_embeddings:
            nn.init.normal_(emb.weight, mean=0.0, std=0.02)

    def forward(self, features: torch.Tensor, coords: torch.Tensor) -> torch.Tensor:
        """Add position embeddings to features.

        Args:
            features: [N, hidden_size] feature tensor.
            coords: [N, coord_dim] integer coordinates, each in [0, max_position).

        Returns:
            [N, hidden_size] features with position embeddings added.
        """
        pos_emb = torch.zeros_like(features)
        for dim in range(self.coord_dim):
            positions = coords[:, dim].long().clamp(0, self.max_position - 1)
            pos_emb = pos_emb + self.position_embeddings[dim](positions)
        return features + pos_emb


class SpatialPatchMerger(nn.Module):
    """Qwen3VL-style spatial patch merger for reducing visual tokens before LLM.

    Merges spatial patches by:
    1. Grouping tokens into merge_size x merge_size windows by (H, W) coordinates
    2. Concatenating features within each window (preserves all information)
    3. Projecting through a 2-layer MLP to the LLM hidden size

    Reduces token count by merge_size^2 (e.g., 4x for 2x2 merging).

    Design notes vs reference qwen_text_decoder.py:
    - Vectorized scatter instead of Python for-loop for placing features into windows.
      The reference iterates per-token which is O(N) Python ops; we use advanced indexing.
    - Properly segments by (batch/segment, T) to support video frames.

    Args:
        input_hidden_size: Encoder output dim (e.g., 1152 for SigLIP2 SO400M).
        out_hidden_size: LLM hidden dim (e.g., 896 for Qwen3-0.6B, read from model config).
        spatial_merge_size: Window size (default 2 for 2x2 merging).
    """

    def __init__(
        self,
        input_hidden_size: int = 1152,
        out_hidden_size: int = 896,
        spatial_merge_size: int = 2,
    ):
        super().__init__()
        self.input_hidden_size = input_hidden_size
        self.out_hidden_size = out_hidden_size
        self.spatial_merge_size = spatial_merge_size
        self.merged_hidden_size = input_hidden_size * (spatial_merge_size**2)

        # Pre-merge LayerNorm on input features
        self.norm = nn.LayerNorm(input_hidden_size, eps=1e-6)
        # 2-layer MLP: concat_dim -> concat_dim (GELU) -> out_hidden_size
        self.linear_fc1 = nn.Linear(self.merged_hidden_size, self.merged_hidden_size)
        self.act_fn = nn.GELU()
        self.linear_fc2 = nn.Linear(self.merged_hidden_size, out_hidden_size)

    def forward(
        self,
        feats: torch.Tensor,
        coords: torch.Tensor,
        layout: list[slice] | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, list[slice] | None]:
        """Merge spatial patches and project to LLM hidden size.

        Args:
            feats: [N, input_hidden_size] token features.
            coords: [N, 4] as (T, H, W, Z) or [N, 5] as (batch/seg, T, H, W, Z).
            layout: Optional batch layout slices from SparseTensor.

        Returns:
            merged_feats: [N', out_hidden_size] where N' ~ N / merge_size^2.
            merged_coords: [N', coord_dims] with scaled H, W.
            new_layout: Updated layout slices.
        """
        if feats.shape[0] == 0:
            empty = torch.zeros(0, self.out_hidden_size, device=feats.device, dtype=feats.dtype)
            return empty, coords, layout

        merge_size = self.spatial_merge_size
        device = feats.device
        orig_dtype = feats.dtype

        # LayerNorm in float32 for numerical stability, cast back
        feats = F.layer_norm(
            feats.float(),
            self.norm.normalized_shape,
            self.norm.weight.float() if self.norm.weight is not None else None,
            self.norm.bias.float() if self.norm.bias is not None else None,
            self.norm.eps,
        ).to(orig_dtype)

        coords_int = coords.long()
        has_segment_dim = coords_int.shape[1] == 5

        if has_segment_dim:
            seg_col, t_col, h_col, w_col, z_col = 0, 1, 2, 3, 4
        else:
            # 4-dim: T serves as segment indicator
            seg_col, t_col, h_col, w_col, z_col = 0, 0, 1, 2, 3

        # Iterate over unique segment values. For images (T=0 always), this
        # gives one group per batch element. For video, different T frames
        # within the same segment share the same seg_col value (the batch idx),
        # so they are merged together spatially — which is correct since spatial
        # merging operates on (H, W) within a frame, and all frames of one image
        # in a packed sample share the same segment_id.
        segment_indices = coords_int[:, seg_col].unique(sorted=True)

        merged_feats_list = []
        merged_coords_list = []
        new_layout = [] if layout is not None else None
        current_offset = 0

        for seg_idx in segment_indices:
            seg_mask = coords_int[:, seg_col] == seg_idx
            seg_feats = feats[seg_mask]
            seg_coords = coords_int[seg_mask]

            h_coords = seg_coords[:, h_col]
            w_coords = seg_coords[:, w_col]

            # Compute merged position and local position within window
            merged_h = h_coords // merge_size
            merged_w = w_coords // merge_size
            local_h = h_coords % merge_size
            local_w = w_coords % merge_size
            local_idx = local_h * merge_size + local_w  # 0..merge_size^2-1

            # Include T in merge keys so patches from different frames are never merged.
            t_coords = seg_coords[:, t_col]
            max_merged_w = (w_coords.max().item() // merge_size) + 1
            max_merged_h = (h_coords.max().item() // merge_size) + 1
            merge_keys = t_coords * (max_merged_h * max_merged_w) + merged_h * max_merged_w + merged_w

            unique_keys, inverse_indices = merge_keys.unique(sorted=True, return_inverse=True)
            num_merged = unique_keys.shape[0]

            # Allocate [num_merged, merge_size^2, input_hidden_size]
            merged_seg_feats = torch.zeros(
                num_merged,
                merge_size * merge_size,
                self.input_hidden_size,
                device=device,
                dtype=orig_dtype,
            )
            # Vectorized scatter: O(1) GPU ops instead of O(N) Python loop.
            # Each token at (inverse_indices[i], local_idx[i]) gets seg_feats[i].
            merged_seg_feats[inverse_indices, local_idx] = seg_feats

            # Flatten: [num_merged, merge_size^2 * input_hidden_size]
            merged_seg_feats = merged_seg_feats.view(num_merged, self.merged_hidden_size)
            merged_feats_list.append(merged_seg_feats)

            # Build merged coordinates (recover T, H, W from flattened key)
            hw_size = max_merged_h * max_merged_w
            merged_t_coords = unique_keys // hw_size
            merged_h_coords = (unique_keys % hw_size) // max_merged_w
            merged_w_coords = (unique_keys % hw_size) % max_merged_w
            z_coord = seg_coords[0, z_col]

            if has_segment_dim:
                merged_seg_coords = torch.stack(
                    [
                        seg_idx.expand(num_merged),
                        merged_t_coords,
                        merged_h_coords,
                        merged_w_coords,
                        z_coord.expand(num_merged),
                    ],
                    dim=1,
                )
            else:
                merged_seg_coords = torch.stack(
                    [
                        merged_t_coords,
                        merged_h_coords,
                        merged_w_coords,
                        z_coord.expand(num_merged),
                    ],
                    dim=1,
                )

            merged_coords_list.append(merged_seg_coords)

            if new_layout is not None:
                new_layout.append(slice(current_offset, current_offset + num_merged))
                current_offset += num_merged

        merged_feats = cat_with_bounded_inputs(merged_feats_list, dim=0)
        merged_coords = cat_with_bounded_inputs(merged_coords_list, dim=0)

        # MLP projection: merged_hidden_size -> out_hidden_size
        merged_feats = self.linear_fc2(self.act_fn(self.linear_fc1(merged_feats)))

        return merged_feats, merged_coords, new_layout


class TextDecoderWrapper(nn.Module):
    """Family-aware causal LM wrapper with vision token injection for image captioning.

    The wrapper keeps the Qwen3 path intact while adding a Nemotron-specific
    adapter for models that do not expose ``inputs_embeds`` on the standard
    decoder forward/generation API.
    """

    def __init__(
        self,
        model_name: str | None = None,
        image_hidden_size: int = 1152,
        spatial_pool_size: int = 2,
        gradient_checkpointing: bool = True,
        dtype: torch.dtype | None = None,
        attn_implementation: str = "flash_attention_2",
        family_spec: TextDecoderFamilySpec | None = None,
        packed_attention_backend: str = PACKED_ATTENTION_BACKEND_SDPA,
    ) -> None:
        super().__init__()

        if dtype is None:
            dtype = torch.bfloat16

        from transformers import AutoModelForCausalLM

        self.spec = family_spec or get_text_decoder_family_spec(model_name=model_name)
        if packed_attention_backend not in PACKED_ATTENTION_BACKENDS:
            raise ValueError(
                f"Unsupported packed_attention_backend={packed_attention_backend!r}; "
                f"expected one of {sorted(PACKED_ATTENTION_BACKENDS)}."
            )
        if packed_attention_backend == PACKED_ATTENTION_BACKEND_NATTEN and self.spec.family != NEMOTRON_2B_SPEC.family:
            raise ValueError("NATTEN packed attention currently supports only the Nemotron 2B text decoder.")
        if (
            packed_attention_backend == PACKED_ATTENTION_BACKEND_NATTEN
            and torch.cuda.is_available()
            and not NATTEN_SUPPORTED
        ):
            raise RuntimeError("NATTEN packed attention was requested but NATTEN is unavailable or unsupported.")
        self.packed_attention_backend: str = packed_attention_backend
        self._model_name: str = model_name or self.spec.default_model_name
        resolved_attn_implementation = _resolve_attn_implementation(attn_implementation)
        if packed_attention_backend == PACKED_ATTENTION_BACKEND_NATTEN:
            # Packed training bypasses the remote decoder attention. Keep all
            # non-packed calls on the explicit dependency-free fallback.
            resolved_attn_implementation = "eager"
            if NATTEN_SUPPORTED:
                logging.info(f"Using NATTEN {get_natten_version()} for packed text decoder attention")
        logging.info(f"Loading text decoder: {self._model_name}")
        logging.info(f"Using text decoder attention backend: {resolved_attn_implementation}")
        hf_cache_dir = os.environ.get("HF_HOME")
        local_files_only = hf_cache_dir is not None
        local_snapshot = resolve_hf_snapshot_path(self._model_name, hf_cache_dir, required_files=("config.json",))
        model_source = local_snapshot or self._model_name
        model_cache_dir = None if local_snapshot is not None else hf_cache_dir
        model_local_files_only = True if local_snapshot is not None else local_files_only

        try:
            self.text_decoder = AutoModelForCausalLM.from_pretrained(
                model_source,
                dtype=dtype,
                attn_implementation=resolved_attn_implementation,
                device_map=None,
                low_cpu_mem_usage=True,
                trust_remote_code=self.spec.trust_remote_code,
                cache_dir=model_cache_dir,
                local_files_only=model_local_files_only,
            )
        except Exception as e:
            if _is_flash_attention_error(e):
                logging.warning("Flash attention unavailable, falling back to eager")
                self.text_decoder = AutoModelForCausalLM.from_pretrained(
                    model_source,
                    dtype=dtype,
                    attn_implementation="eager",
                    device_map=None,
                    low_cpu_mem_usage=True,
                    trust_remote_code=self.spec.trust_remote_code,
                    cache_dir=model_cache_dir,
                    local_files_only=model_local_files_only,
                )
            else:
                raise

        if self.spec.family == NEMOTRON_2B_SPEC.family:
            repaired_rotary_buffers = _repair_nemotron_rotary_buffers(self.text_decoder)
            if repaired_rotary_buffers == 0:
                logging.warning(
                    "No Nemotron rotary embedding buffers were reinitialized; "
                    "the remote-code rotary module may have changed its inv_freq/dim/base attributes."
                )
            else:
                logging.info(f"Reinitialized {repaired_rotary_buffers} Nemotron rotary embedding buffers")

        self.lm_config = self.text_decoder.config
        self.lm_config.pad_token_id = self.spec.pad_token_id
        if getattr(self.lm_config, "eos_token_id", None) is None:
            if len(self.spec.eos_token_ids) == 1:
                self.lm_config.eos_token_id = self.spec.eos_token_ids[0]
            else:
                self.lm_config.eos_token_id = list(self.spec.eos_token_ids)
        self.vision_token_id = self.spec.vision_pad_token_id
        self.image_hidden_size = image_hidden_size
        self.spatial_pool_size = spatial_pool_size
        self._caption_tokenizer: Any | None = None
        hidden_size = self.lm_config.hidden_size

        # Spatial merging or simple projection
        if spatial_pool_size > 1:
            self.spatial_merger = SpatialPatchMerger(
                input_hidden_size=image_hidden_size,
                out_hidden_size=hidden_size,
                spatial_merge_size=spatial_pool_size,
            )
            self.image_proj = None
            logging.info(
                f"SpatialPatchMerger enabled: {image_hidden_size} -> {hidden_size}, "
                f"merge_size={spatial_pool_size} ({spatial_pool_size**2}x token reduction)"
            )
        else:
            self.spatial_merger = None
            self.image_proj = nn.Sequential(
                nn.Linear(image_hidden_size, image_hidden_size * 2),
                nn.GELU(),
                nn.Linear(image_hidden_size * 2, hidden_size),
            )

        # 2D position embeddings for merged image features
        self.image_pos_embed = ImagePositionEmbeddings(
            hidden_size=hidden_size,
            max_position=128,  # supports up to ~2048px / patch_size=16
            coord_dim=2,  # (H, W) only
        )

        # Disable KV cache for training forwards. Generation paths opt back in.
        self.text_decoder.config.use_cache = False

        self._manual_gradient_checkpointing_enabled = False
        gradient_checkpointing_enabled = False
        if gradient_checkpointing:
            try:
                self.text_decoder.gradient_checkpointing_enable()
                gradient_checkpointing_enabled = True
            except ValueError as exc:
                if "does not support gradient checkpointing" in str(exc).lower():
                    if not self.spec.supports_inputs_embeds_forward:
                        self._manual_gradient_checkpointing_enabled = True
                        gradient_checkpointing_enabled = True
                        logging.warning(
                            f"Text decoder {self.text_decoder.__class__.__name__} does not support native "
                            "gradient checkpointing; using manual activation checkpointing in the decoder "
                            "layer loop"
                        )
                    else:
                        logging.warning(
                            f"Text decoder {self.text_decoder.__class__.__name__} does not support "
                            "gradient checkpointing; continuing without it"
                        )
                else:
                    raise
        elif hasattr(self.text_decoder, "gradient_checkpointing_disable"):
            self.text_decoder.gradient_checkpointing_disable()

        logging.info(
            f"TextDecoderWrapper ready: family={self.spec.family}, hidden_size={hidden_size}, "
            f"layers={len(self.text_decoder.model.layers)}, "
            f"gradient_checkpointing={gradient_checkpointing_enabled}, "
            f"packed_attention_backend={self.packed_attention_backend}, use_cache=False"
        )

    def _get_eos_token_ids(self) -> tuple[int, ...]:
        """Resolve all EOS token IDs used for generation and bookkeeping."""
        eos_token_id = getattr(self.lm_config, "eos_token_id", None)
        if isinstance(eos_token_id, list):
            return tuple(int(token_id) for token_id in eos_token_id)
        if eos_token_id is None:
            return self.spec.eos_token_ids
        return (int(eos_token_id),)

    def _ensure_caption_tokenizer(self) -> Any:
        """Lazily load the tokenizer used for generation-time decoding."""
        if self._caption_tokenizer is None:
            hf_cache_dir = os.environ.get("HF_HOME")
            if self.spec.family == NEMOTRON_2B_SPEC.family:
                from transformers import AutoTokenizer

                patched_snapshot = prepare_nemotron_tokenizer_snapshot(self._model_name, hf_cache_dir)
                if patched_snapshot is not None:
                    self._caption_tokenizer = AutoTokenizer.from_pretrained(
                        patched_snapshot,
                        local_files_only=True,
                        trust_remote_code=self.spec.trust_remote_code,
                    )
                else:
                    self._caption_tokenizer = load_auto_tokenizer_from_cache(
                        self._model_name,
                        hf_cache_dir,
                        trust_remote_code=self.spec.trust_remote_code,
                    )
            else:
                self._caption_tokenizer = load_auto_tokenizer_from_cache(
                    self._model_name,
                    hf_cache_dir,
                    trust_remote_code=self.spec.trust_remote_code,
                )
        return self._caption_tokenizer

    def _embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        """Embed token IDs using the model's input embedding module."""
        return self.text_decoder.get_input_embeddings()(input_ids)

    def _forward_nemotron_natten_from_embeddings(
        self,
        text_embeds: torch.Tensor,
        position_ids: torch.Tensor,
        cumulative_seqlen: torch.Tensor,
        max_seqlen: int,
    ) -> SimpleNamespace:
        """Run packed Nemotron layers with NATTEN varlen self-attention."""
        if self.spec.family != NEMOTRON_2B_SPEC.family:
            raise ValueError("NATTEN packed attention currently supports only the Nemotron 2B text decoder.")
        if text_embeds.shape[0] != 1:
            raise ValueError(f"NATTEN varlen requires packed batch size 1, got {text_embeds.shape[0]}.")
        if cumulative_seqlen.dtype != torch.int32:
            raise ValueError(f"NATTEN cumulative sequence lengths must be int32, got {cumulative_seqlen.dtype}.")

        hidden_states = text_embeds  # [1,T,D]
        decoder_layers = self.text_decoder.model.layers
        if not decoder_layers:
            raise ValueError("NATTEN packed attention requires at least one Nemotron decoder layer.")
        first_attention = decoder_layers[0].self_attn
        normalized_position_ids = position_ids.reshape(hidden_states.shape[0], hidden_states.shape[1]).long()  # [1,T]
        rotary_reference = hidden_states.new_empty(
            (
                hidden_states.shape[0],
                first_attention.num_key_value_heads,
                0,
                first_attention.head_dim,
            )
        )  # [1,Hkv,0,D]
        rotary_embeddings = first_attention.rotary_emb(
            rotary_reference,
            normalized_position_ids,
        )  # each [1,1,T,D]
        use_manual_gradient_checkpointing = (
            self._manual_gradient_checkpointing_enabled and self.training and hidden_states.requires_grad
        )
        for decoder_layer in decoder_layers:
            if use_manual_gradient_checkpointing:

                def _layer_forward(
                    layer_hidden_states: torch.Tensor,
                    layer_module: nn.Module = decoder_layer,
                ) -> torch.Tensor:
                    return _forward_nemotron_natten_decoder_layer(
                        layer_module,
                        layer_hidden_states,
                        position_ids,
                        cumulative_seqlen,
                        max_seqlen,
                        rotary_embeddings,
                    )

                hidden_states = torch.utils.checkpoint.checkpoint(
                    _layer_forward,
                    hidden_states,
                    preserve_rng_state=True,
                    use_reentrant=False,
                )  # [1,T,D]
            else:
                hidden_states = _forward_nemotron_natten_decoder_layer(
                    decoder_layer,
                    hidden_states,
                    position_ids,
                    cumulative_seqlen,
                    max_seqlen,
                    rotary_embeddings,
                )  # [1,T,D]

        hidden_states = self.text_decoder.model.norm(hidden_states)  # [1,T,D]
        return SimpleNamespace(last_hidden_state=hidden_states, past_key_values=None)

    def _forward_from_embeddings(
        self,
        text_embeds: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.Tensor | None = None,
        cache_position: torch.Tensor | None = None,
        use_cache: bool = False,
        past_key_values: Any | None = None,
        packed_cumulative_seqlen: torch.Tensor | None = None,
        packed_max_seqlen: int | None = None,
    ) -> SimpleNamespace:
        """Run the LM from caller-provided token embeddings."""
        has_packed_cumulative_seqlen = packed_cumulative_seqlen is not None
        has_packed_max_seqlen = packed_max_seqlen is not None
        if has_packed_cumulative_seqlen != has_packed_max_seqlen:
            raise ValueError("packed_cumulative_seqlen and packed_max_seqlen must be provided together.")
        if has_packed_cumulative_seqlen:
            if self.packed_attention_backend != PACKED_ATTENTION_BACKEND_NATTEN:
                raise ValueError("Packed varlen metadata was provided without the NATTEN packed attention backend.")
            if attention_mask is not None or use_cache or past_key_values is not None:
                raise ValueError("NATTEN packed training does not support masks or KV-cache inputs.")
            if position_ids is None:
                raise ValueError("NATTEN packed attention requires per-token position IDs.")
            assert packed_cumulative_seqlen is not None
            assert packed_max_seqlen is not None
            return self._forward_nemotron_natten_from_embeddings(
                text_embeds=text_embeds,
                position_ids=position_ids,
                cumulative_seqlen=packed_cumulative_seqlen,
                max_seqlen=packed_max_seqlen,
            )

        if self.spec.supports_inputs_embeds_forward:
            model_kwargs: dict[str, Any] = {
                "inputs_embeds": text_embeds,
                "attention_mask": attention_mask,
                "position_ids": position_ids,
                "use_cache": use_cache,
                "return_dict": True,
            }
            if past_key_values is not None:
                model_kwargs["past_key_values"] = past_key_values
            if self.spec.supports_cache_position and cache_position is not None:
                model_kwargs["cache_position"] = cache_position
            return self.text_decoder.model(**model_kwargs)

        hidden_states = text_embeds
        next_decoder_cache = () if use_cache else None
        hidden_padding_mask = None
        if (
            attention_mask is not None
            and attention_mask.dim() == 2
            and attention_mask.shape[0] == hidden_states.shape[0]
            and attention_mask.shape[1] == hidden_states.shape[1]
        ):
            # The manual Nemotron path batches unequal packed segments. Keep
            # padded rows finite between layers so masked SDPA outputs cannot
            # poison later key/value projections.
            hidden_padding_mask = attention_mask.to(device=hidden_states.device, dtype=torch.bool)  # [B,T]
        use_manual_gradient_checkpointing = (
            self._manual_gradient_checkpointing_enabled
            and self.training
            and not use_cache
            and hidden_states.requires_grad
        )
        for idx, decoder_layer in enumerate(self.text_decoder.model.layers):
            layer_past = past_key_values[idx] if past_key_values is not None else None
            if use_manual_gradient_checkpointing:

                def _layer_forward(
                    layer_hidden_states: torch.Tensor,
                    layer_module: nn.Module = decoder_layer,
                ) -> torch.Tensor:
                    layer_outputs = layer_module(
                        layer_hidden_states,
                        attention_mask=attention_mask,
                        position_ids=position_ids,
                        past_key_value=None,
                        use_cache=False,
                    )
                    return layer_outputs[0]

                hidden_states = torch.utils.checkpoint.checkpoint(
                    _layer_forward,
                    hidden_states,
                    use_reentrant=False,
                )
                if hidden_padding_mask is not None:
                    hidden_states = hidden_states.masked_fill(~hidden_padding_mask[:, :, None], 0)  # [B,T,D]
            else:
                layer_outputs = decoder_layer(
                    hidden_states,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    past_key_value=layer_past,
                    use_cache=use_cache,
                )
                hidden_states = layer_outputs[0]
                if hidden_padding_mask is not None:
                    hidden_states = hidden_states.masked_fill(~hidden_padding_mask[:, :, None], 0)  # [B,T,D]
                if use_cache:
                    next_decoder_cache += (layer_outputs[1],)

        hidden_states = self.text_decoder.model.norm(hidden_states)
        return SimpleNamespace(last_hidden_state=hidden_states, past_key_values=next_decoder_cache)

    def _lm_head(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """Project decoder hidden states to token logits."""
        return self.text_decoder.lm_head(hidden_states)

    @property
    def output_projection(self) -> nn.Module:
        """Return the LM head used by the chunked training loss."""
        return self.text_decoder.lm_head

    def _sample_next_token(
        self,
        logits: torch.Tensor,
        *,
        do_sample: bool,
        temperature: float,
        top_p: float,
        top_k: int,
        suppress_token_ids: tuple[int, ...],
    ) -> torch.Tensor:
        """Sample or greedily select the next token from one-step logits."""
        next_token_logits = logits[:, -1, :].float()
        if suppress_token_ids:
            next_token_logits[:, list(suppress_token_ids)] = torch.finfo(next_token_logits.dtype).min

        if not do_sample:
            return torch.argmax(next_token_logits, dim=-1)

        safe_temperature = max(float(temperature), 1.0e-5)
        next_token_logits = next_token_logits / safe_temperature

        if top_k > 0:
            top_values, _ = torch.topk(next_token_logits, k=min(top_k, next_token_logits.shape[-1]), dim=-1)
            kth_values = top_values[:, -1].unsqueeze(-1)
            next_token_logits = next_token_logits.masked_fill(next_token_logits < kth_values, float("-inf"))

        if 0.0 < top_p < 1.0:
            sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True, dim=-1)
            sorted_probs = torch.softmax(sorted_logits, dim=-1)
            cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
            sorted_mask = cumulative_probs > top_p
            sorted_mask[:, 1:] = sorted_mask[:, :-1].clone()
            sorted_mask[:, 0] = False
            next_token_logits.scatter_(
                1,
                sorted_indices,
                next_token_logits.gather(1, sorted_indices).masked_fill(sorted_mask, float("-inf")),
            )

        next_token_probs = torch.softmax(next_token_logits, dim=-1)
        return torch.multinomial(next_token_probs, num_samples=1).squeeze(-1)

    def _build_vqa_user_turn(
        self,
        tokenizer: Any,
        question: str,
        num_image_tokens: int,
    ) -> tuple[list[int], int]:
        """Build one multimodal user turn and return the image-pad start offset."""
        user_turn, image_pad_offsets = self._build_vqa_user_turn_with_visual_blocks(
            tokenizer=tokenizer,
            question=question,
            image_token_counts_by_visual=[int(num_image_tokens)],
        )
        return user_turn, image_pad_offsets[0]

    def _build_vqa_user_turn_with_visual_blocks(
        self,
        tokenizer: Any,
        question: str,
        image_token_counts_by_visual: list[int],
    ) -> tuple[list[int], list[int]]:
        """Build one multimodal user turn and return image-pad offsets per visual."""
        if not image_token_counts_by_visual:
            raise ValueError("VQA prompt construction requires at least one visual token block.")
        if any(int(token_count) <= 0 for token_count in image_token_counts_by_visual):
            raise ValueError(f"Visual token counts must be positive, got {image_token_counts_by_visual!r}.")

        def _encode(text: str) -> list[int]:
            return tokenizer.encode(text, add_special_tokens=False)

        image_placeholder = "<image>"
        num_visuals = len(image_token_counts_by_visual)
        normalized_question = str(question).strip()
        placeholder_count = normalized_question.count(image_placeholder)
        if num_visuals == 1:
            if placeholder_count > 1:
                raise ValueError(f"VQA prompt has {placeholder_count} image placeholders for one visual input.")
        else:
            if placeholder_count == 0:
                visual_placeholders = "\n".join(image_placeholder for _ in image_token_counts_by_visual)
                normalized_question = f"{visual_placeholders}\n{normalized_question}".strip()
            elif placeholder_count != num_visuals:
                raise ValueError(
                    f"VQA prompt has {placeholder_count} image placeholders for {num_visuals} visual inputs."
                )

        if self.spec.family == QWEN3_SPEC.family:
            im_start_id = QWEN3_IM_START_TOKEN_ID
            im_end_id = QWEN3_IM_END_TOKEN_ID
            vision_start_id = QWEN3_VISION_START_TOKEN_ID
            vision_pad_id = QWEN3_VISION_PAD_TOKEN_ID
            vision_end_id = QWEN3_VISION_END_TOKEN_ID
        elif self.spec.family == NEMOTRON_2B_SPEC.family:
            im_start_id = _get_required_token_id(tokenizer, NEMOTRON_2B_IM_START_TOKEN)
            im_end_id = _get_required_token_id(tokenizer, NEMOTRON_2B_IM_END_TOKEN)
            vision_start_id = self.spec.vision_start_token_id
            vision_pad_id = self.spec.vision_pad_token_id
            vision_end_id = self.spec.vision_end_token_id
        else:
            raise NotImplementedError(
                f"VQA prompt construction is not implemented for text decoder family {self.spec.family!r}."
            )

        user_turn = [im_start_id] + _encode("user\n")
        image_pad_offsets: list[int] = []
        remaining_text = normalized_question
        visual_index = 0
        while image_placeholder in remaining_text:
            placeholder_index = remaining_text.find(image_placeholder)
            before_text = remaining_text[:placeholder_index]
            remaining_text = remaining_text[placeholder_index + len(image_placeholder) :]
            user_turn.extend(_encode(before_text))
            if num_visuals > 1:
                user_turn.extend(_encode(densevl_add_vision_id_text("image", visual_index + 1)))
            image_pad_offsets.append(len(user_turn) + 1)
            user_turn.extend([vision_start_id])
            user_turn.extend([vision_pad_id] * int(image_token_counts_by_visual[visual_index]))
            user_turn.extend([vision_end_id])
            visual_index += 1
        if image_pad_offsets:
            user_turn.extend(_encode(remaining_text))
        else:
            image_pad_offsets.append(len(user_turn) + 1)
            user_turn.extend([vision_start_id])
            user_turn.extend([vision_pad_id] * int(image_token_counts_by_visual[0]))
            user_turn.extend([vision_end_id])
            user_turn.extend(_encode("\n"))
            user_turn.extend(_encode(normalized_question))

        user_turn.extend([im_end_id])
        user_turn.extend(_encode("\n"))
        return user_turn, image_pad_offsets

    def forward(
        self,
        input_ids: torch.Tensor,
        image_feats_tensor: torch.Tensor,
        image_coords: torch.Tensor,
        image_patch_indices: torch.Tensor,
        image_layout: list[slice] | None = None,
        segment_ids: torch.Tensor | None = None,
        return_hidden_states: bool = False,
    ) -> tuple[torch.Tensor, int]:
        """Forward pass: inject image features into text and run causal LM.

        When segment_ids are provided, SDPA reshapes the packed sequence into a
        padded per-segment batch, while NATTEN retains the packed layout and uses
        cumulative sequence offsets. Both paths provide segment-isolated causal attention.
        Position IDs reset per segment for correct RoPE.

        Args:
            input_ids: [B, S] text token IDs (with <|image_pad|> as placeholders).
            image_feats_tensor: [N, encoder_dim] raw encoder output features.
            image_coords: [N, 4 or 5] spatial coordinates.
            image_patch_indices: [N_pooled] flat indices into [B*S] sequence where
                merged image features should be inserted. The host collator validates
                and removes padding sentinels before device transfer.
            image_layout: Optional batch layout slices for the image features.
            segment_ids: [B, S] segment IDs for packed sequences. Values >= 0 indicate
                valid segments, -1 indicates padding. When provided, enables
                segment-isolated attention and per-segment position IDs.
            return_hidden_states: Return decoder hidden states instead of projecting
                the complete sequence to vocabulary logits. This is used by the
                chunked training loss to bound LM-head memory.

        Returns:
            Dense [B, S, vocab_size] logits, or [B, S, hidden_size] hidden
            states when ``return_hidden_states=True``.
            num_pooled_tokens: Number of image tokens after spatial merging.
        """
        image_features = image_feats_tensor
        current_coords = image_coords

        if len(image_features) > 0:
            # Spatial merging + MLP projection
            if self.spatial_merger is not None:
                image_features, current_coords, _ = self.spatial_merger(
                    feats=image_features,
                    coords=current_coords,
                    layout=image_layout,
                )
            elif self.image_proj is not None:
                image_features = self.image_proj(image_features)

            # Extract (H, W) for 2D position embeddings
            if current_coords.shape[1] == 5:
                coords_2d = current_coords[:, 2:4]  # (seg, T, H, W, Z) -> (H, W)
            else:
                coords_2d = current_coords[:, 1:3]  # (T, H, W, Z) -> (H, W)

            image_features = self.image_pos_embed(features=image_features, coords=coords_2d)

        num_pooled_tokens = len(image_features)

        # Embed text tokens
        text_embeds = self._embed_input_ids(input_ids)  # [B, S, d]
        B, S, d = text_embeds.shape

        # The host collator validates exact placeholder/index coverage and removes
        # padding sentinels before tensors move to the accelerator.
        # Zero out vision placeholder positions, then insert real image features.
        vision_mask = (input_ids != self.vision_token_id).to(dtype=text_embeds.dtype)  # [B,S]
        text_embeds = text_embeds * vision_mask[:, :, None]  # [B,S,d]

        if image_patch_indices.ndim != 1:
            raise ValueError(
                f"Image patch indices must be one-dimensional, got shape {tuple(image_patch_indices.shape)}."
            )
        if (
            image_patch_indices.dtype == torch.bool
            or image_patch_indices.is_floating_point()
            or image_patch_indices.is_complex()
        ):
            raise TypeError(f"Image patch indices must use an integer dtype, got {image_patch_indices.dtype}.")
        if image_patch_indices.numel() != num_pooled_tokens:
            raise ValueError(
                "Image feature/index count mismatch: "
                f"got {num_pooled_tokens} pooled features and {image_patch_indices.numel()} indices."
            )
        if image_patch_indices.numel() > 0:
            insert_idx = image_patch_indices.to(device=input_ids.device, dtype=torch.long)  # [N]
            text_embeds_flat = text_embeds.reshape(B * S, d)  # [B*S,d]
            text_insert_idx = insert_idx.to(device=text_embeds.device)  # [N]
            text_embeds_flat[text_insert_idx] = image_features.to(text_embeds.dtype)  # [N,d]
            text_embeds = text_embeds_flat.reshape(B, S, d)  # [B,S,d]

        # Segment-aware forward pass for packed sequences
        if segment_ids is not None:
            if return_hidden_states:
                decoder_output = self._forward_packed(  # [B,S,D]
                    text_embeds,
                    input_ids,
                    segment_ids,
                    return_hidden_states=True,
                )
            else:
                decoder_output = self._forward_packed(  # [B,S,Vocab]
                    text_embeds,
                    input_ids,
                    segment_ids,
                )
        else:
            if return_hidden_states:
                decoder_output = self._forward_standard(  # [B,S,D]
                    text_embeds,
                    input_ids=input_ids,
                    return_hidden_states=True,
                )
            else:
                decoder_output = self._forward_standard(  # [B,S,Vocab]
                    text_embeds,
                    input_ids=input_ids,
                )

        return decoder_output, num_pooled_tokens

    def _forward_standard(
        self,
        text_embeds: torch.Tensor,
        input_ids: torch.Tensor | None = None,
        return_hidden_states: bool = False,
    ) -> torch.Tensor:
        """Standard forward pass without segment isolation."""
        effective_text_embeds = text_embeds
        original_seq_len = text_embeds.shape[1]
        effective_seq_len = original_seq_len
        attention_mask = None
        if input_ids is not None:
            pad_token_id = self.lm_config.pad_token_id
            if pad_token_id is None:
                attention_mask = torch.ones_like(input_ids, dtype=torch.long)
            else:
                attention_mask = (input_ids != pad_token_id).to(dtype=torch.long)
                valid_lengths = attention_mask.sum(dim=1)
                if valid_lengths.numel() > 0:
                    # Caption/VQA batches are right padded. Trim the common trailing
                    # padding tail so non-packed Nemotron runs do not forward the
                    # entire fixed max_seq_len block through the decoder.
                    effective_seq_len = max(int(valid_lengths.max().item()), 1)
                    if effective_seq_len < original_seq_len:
                        effective_text_embeds = text_embeds[:, :effective_seq_len, :]
                        attention_mask = attention_mask[:, :effective_seq_len]

        cache_position = torch.arange(effective_seq_len, device=text_embeds.device)
        position_ids = cache_position.unsqueeze(0)
        outputs = self._forward_from_embeddings(
            text_embeds=effective_text_embeds,
            attention_mask=attention_mask,
            position_ids=position_ids,
            cache_position=cache_position,
            use_cache=False,
        )
        decoder_hidden_states = outputs.last_hidden_state  # [B,T,D]
        decoder_output = (
            decoder_hidden_states if return_hidden_states else self._lm_head(decoder_hidden_states)
        )  # [B,T,D|Vocab]
        if effective_seq_len == original_seq_len:
            return decoder_output

        pad_len = original_seq_len - effective_seq_len
        padded_output = decoder_output.new_zeros(  # [B,S-T,D|Vocab]
            decoder_output.shape[0],
            pad_len,
            decoder_output.shape[-1],
        )
        return torch.cat([decoder_output, padded_output], dim=1)  # [B,S,D|Vocab]

    def _forward_packed(
        self,
        text_embeds: torch.Tensor,
        input_ids: torch.Tensor,
        segment_ids: torch.Tensor,
        return_hidden_states: bool = False,
    ) -> torch.Tensor:
        """Segment-isolated forward pass for packed sequences.

        SDPA reshapes packed [B, S] into padded per-segment batches. NATTEN keeps
        valid tokens contiguous and uses cumulative sequence offsets. Position IDs
        reset per segment for correct RoPE in both paths.

        Args:
            text_embeds: [B, S, d] embeddings with image features already injected.
            input_ids: [B, S] original token IDs (for reconstructing output shape).
            segment_ids: [B, S] segment IDs (>=0 for valid, -1 for padding).

        Returns:
            Dense [B, S, vocab_size] logits, or [B, S, hidden_size] hidden
            states when ``return_hidden_states=True``.
        """
        B, S, d = text_embeds.shape
        device = text_embeds.device
        dtype = text_embeds.dtype

        # Process each batch element (typically B=1 for packed sequences)
        all_outputs: list[torch.Tensor] = []
        for b in range(B):
            seg_ids = segment_ids[b]  # [S]
            embeds = text_embeds[b]  # [S, d]

            # Find valid segments (>= 0)
            valid_indices = (seg_ids >= 0).nonzero(as_tuple=True)[0]
            if valid_indices.numel() == 0:
                # All padding — return zeros
                output_size = self.lm_config.hidden_size if return_hidden_states else self.lm_config.vocab_size
                all_outputs.append(torch.zeros(S, output_size, device=device, dtype=dtype))  # [S,D|Vocab]
                continue

            valid_seg_ids = seg_ids.index_select(0, valid_indices)
            valid_embeds = embeds.index_select(0, valid_indices)

            # Packed collate emits each segment as one contiguous run, so
            # unique_consecutive gives the per-segment lengths without
            # scalarizing each segment length back to Python.
            segment_values, seg_lengths = torch.unique_consecutive(valid_seg_ids, return_counts=True)
            if __debug__:
                unique_segment_count = int(valid_seg_ids.unique().numel())
                assert segment_values.numel() == unique_segment_count, (
                    "Segment IDs must be contiguous within packed text decoder inputs"
                )
            num_segments = seg_lengths.numel()

            segment_starts = torch.cumsum(seg_lengths, dim=0) - seg_lengths  # [R]
            segment_positions = torch.arange(valid_indices.numel(), device=device, dtype=torch.long)  # [V]
            segment_positions = segment_positions - torch.repeat_interleave(segment_starts, seg_lengths)  # [V]

            if self.packed_attention_backend == PACKED_ATTENTION_BACKEND_NATTEN:
                segment_lengths_int32 = seg_lengths.to(dtype=torch.int32)  # [R]
                zero_offset = torch.zeros(1, device=device, dtype=torch.int32)  # [1]
                cumulative_lengths = segment_lengths_int32.cumsum(dim=0, dtype=torch.int32)  # [R]
                cumulative_seqlen = torch.cat([zero_offset, cumulative_lengths], dim=0)  # [R+1]
                max_seqlen = int(seg_lengths.max().item())
                packed_embeds = valid_embeds.unsqueeze(0)  # [1,V,D]
                packed_position_ids = segment_positions.unsqueeze(0)  # [1,V]
                outputs = self._forward_from_embeddings(
                    text_embeds=packed_embeds,
                    attention_mask=None,
                    position_ids=packed_position_ids,
                    use_cache=False,
                    packed_cumulative_seqlen=cumulative_seqlen,
                    packed_max_seqlen=max_seqlen,
                )
                valid_hidden_states = outputs.last_hidden_state  # [1,V,D]
                valid_output = (
                    valid_hidden_states if return_hidden_states else self._lm_head(valid_hidden_states)
                )  # [1,V,D|Vocab]
                packed_output = valid_output.new_zeros((S, valid_output.shape[-1]))  # [S,D|Vocab]
                packed_output[valid_indices] = valid_output[0]  # [V,D|Vocab]
                all_outputs.append(packed_output)
                continue

            segment_rows = torch.repeat_interleave(
                torch.arange(num_segments, device=device, dtype=torch.long),
                seg_lengths,
            )  # [V]
            max_seg_len = int(seg_lengths.max().item())
            batch_embeds = torch.zeros(num_segments, max_seg_len, d, device=device, dtype=dtype)  # [R,T,D]
            batch_embeds[segment_rows, segment_positions] = valid_embeds  # [V,D]

            position_template = torch.arange(max_seg_len, device=device, dtype=torch.long)  # [T]
            batch_mask = (position_template.unsqueeze(0) < seg_lengths.unsqueeze(1)).to(dtype=torch.long)  # [R,T]
            batch_pos = position_template.unsqueeze(0).expand(num_segments, -1) * batch_mask  # [R,T]
            outputs = self._forward_from_embeddings(
                text_embeds=batch_embeds,
                attention_mask=batch_mask,
                position_ids=batch_pos,
                cache_position=position_template,
                use_cache=False,
            )
            segment_hidden_states = outputs.last_hidden_state  # [R,T,D]
            segment_output = (
                segment_hidden_states if return_hidden_states else self._lm_head(segment_hidden_states)
            )  # [R,T,D|Vocab]
            packed_output = segment_output.new_zeros((S, segment_output.shape[-1]))  # [S,D|Vocab]
            packed_output[valid_indices] = segment_output[segment_rows, segment_positions]  # [V,D|Vocab]
            all_outputs.append(packed_output)

        # Guard: empty batch (all samples skipped by collate) -> empty output tensor.
        if len(all_outputs) == 0:
            output_size = self.lm_config.hidden_size if return_hidden_states else self.lm_config.vocab_size
            return torch.zeros(0, S, output_size, device=device, dtype=dtype)  # [0,S,D|Vocab]
        return stack_with_bounded_inputs(all_outputs, dim=0)  # [B,S,D|Vocab]

    def _decode_generation_result(
        self,
        generated_ids: torch.Tensor,
        input_len: int,
        tokenizer: Any,
        max_new_tokens: int,
        eos_token_ids: tuple[int, ...],
        skip_special_tokens: bool = True,
        decode_prefix_ids: list[int] | None = None,
    ) -> tuple[str, dict[str, bool | int]]:
        """Decode generated token IDs and expose basic generation metadata."""
        generated_only = generated_ids[0, input_len:]  # [T_gen]
        generated_token_count = int(generated_only.shape[0])
        finished_with_eos = generated_token_count > 0 and int(generated_only[-1].item()) in eos_token_ids
        truncated = generated_token_count >= max_new_tokens and not finished_with_eos
        decode_ids: torch.Tensor | list[int] = generated_only
        if decode_prefix_ids is not None:
            decode_ids = list(decode_prefix_ids) + generated_only.tolist()
        text = tokenizer.decode(decode_ids, skip_special_tokens=skip_special_tokens)
        metadata: dict[str, bool | int] = {
            "generated_tokens": generated_token_count,
            "finished_with_eos": finished_with_eos,
            "truncated": truncated,
        }
        return text, metadata

    def _generate_from_prefix_embeddings(
        self,
        *,
        input_ids: torch.Tensor,
        text_embeds: torch.Tensor,
        max_new_tokens: int,
        do_sample: bool,
        temperature: float,
        eos_token_ids: tuple[int, ...],
        suppress_token_ids: tuple[int, ...] | None = None,
    ) -> torch.Tensor:
        """Autoregressively generate from a caller-provided embedding prefix."""
        resolved_suppress_token_ids = self.spec.suppress_token_ids if suppress_token_ids is None else suppress_token_ids
        if self.spec.supports_inputs_embeds_generate:
            eos_token_id: int | list[int] = eos_token_ids[0] if len(eos_token_ids) == 1 else list(eos_token_ids)
            generate_kwargs = dict(
                input_ids=input_ids,
                inputs_embeds=text_embeds,
                max_new_tokens=max_new_tokens,
                do_sample=do_sample,
                use_cache=True,
                pad_token_id=self.spec.pad_token_id,
                eos_token_id=eos_token_id,
            )
            if resolved_suppress_token_ids:
                generate_kwargs["suppress_tokens"] = list(resolved_suppress_token_ids)
            if do_sample:
                generate_kwargs["temperature"] = temperature
                generate_kwargs["top_p"] = 0.8
                generate_kwargs["top_k"] = 20
            else:
                generate_kwargs["top_p"] = 1.0
                generate_kwargs["top_k"] = 0
                generate_kwargs["temperature"] = 1.0
            return self.text_decoder.generate(**generate_kwargs)

        if input_ids.shape[0] != 1:
            raise ValueError("Manual caption generation for the Nemotron text decoder only supports batch size 1.")

        prefix_len = text_embeds.shape[1]
        prefix_position_ids = torch.arange(prefix_len, device=text_embeds.device, dtype=torch.long).unsqueeze(0)
        attention_mask = torch.ones_like(input_ids, dtype=torch.long)
        outputs = self._forward_from_embeddings(
            text_embeds=text_embeds,
            attention_mask=attention_mask,
            position_ids=prefix_position_ids,
            use_cache=True,
        )
        logits = self._lm_head(outputs.last_hidden_state)
        past_key_values = outputs.past_key_values
        generated_tokens: list[torch.Tensor] = []
        eos_token_tensor = torch.tensor(eos_token_ids, device=input_ids.device, dtype=input_ids.dtype)

        for step_idx in range(max_new_tokens):
            next_token = self._sample_next_token(
                logits,
                do_sample=do_sample,
                temperature=temperature,
                top_p=0.8 if do_sample else 1.0,
                top_k=20 if do_sample else 0,
                suppress_token_ids=resolved_suppress_token_ids,
            )
            generated_tokens.append(next_token.unsqueeze(1))

            if bool(torch.isin(next_token, eos_token_tensor).all()):
                break
            if step_idx == max_new_tokens - 1:
                break

            total_len = input_ids.shape[1] + len(generated_tokens)
            attention_mask = torch.ones((input_ids.shape[0], total_len), dtype=torch.long, device=input_ids.device)
            position_ids = torch.full(
                (input_ids.shape[0], 1),
                fill_value=total_len - 1,
                device=input_ids.device,
                dtype=torch.long,
            )
            next_token_embeds = self._embed_input_ids(next_token.unsqueeze(1))
            outputs = self._forward_from_embeddings(
                text_embeds=next_token_embeds,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                use_cache=True,
            )
            logits = self._lm_head(outputs.last_hidden_state)
            past_key_values = outputs.past_key_values

        if len(generated_tokens) == 0:
            return input_ids

        generated_ids = cat_with_bounded_inputs(generated_tokens, dim=1)  # [B,T_generated]
        return torch.cat([input_ids, generated_ids], dim=1)

    @torch.no_grad()
    def generate_caption(
        self,
        image_feats_tensor: torch.Tensor,
        image_coords: torch.Tensor,
        max_new_tokens: int = 256,
        do_sample: bool = True,
        temperature: float = 0.7,
        return_metadata: bool = False,
    ) -> str | tuple[str, dict[str, bool | int]]:
        """Generate a caption from a single image's encoder features.

        Args:
            image_feats_tensor: [N, encoder_dim] features for ONE image.
            image_coords: [N, 4 or 5] coordinates.
            max_new_tokens: Maximum tokens to generate.
            do_sample: Whether to sample (False = greedy for deterministic vis).
            temperature: Sampling temperature if do_sample=True.
            return_metadata: Whether to also return generation metadata.

        Returns:
            Generated caption string, or ``(caption, metadata)`` when requested.
        """
        # Spatial merge + position embeddings
        if self.spatial_merger is not None:
            image_features, coords, _ = self.spatial_merger(image_feats_tensor, image_coords)
        elif self.image_proj is not None:
            image_features = self.image_proj(image_feats_tensor)
            coords = image_coords
        else:
            image_features = image_feats_tensor
            coords = image_coords

        if coords.shape[1] == 5:
            coords_2d = coords[:, 2:4]
        else:
            coords_2d = coords[:, 1:3]
        image_features = self.image_pos_embed(image_features, coords_2d)

        num_image_tokens = len(image_features)

        # Build input: [<|vision_start|>, <|image_pad|>×N, <|vision_end|>]
        device = image_features.device
        input_ids = torch.tensor(
            [self.spec.vision_start_token_id]
            + [self.spec.vision_pad_token_id] * num_image_tokens
            + [self.spec.vision_end_token_id],
            dtype=torch.long,
            device=device,
        ).unsqueeze(0)

        # Embed and inject image features
        text_embeds = self._embed_input_ids(input_ids)
        vision_mask = (input_ids != self.spec.vision_pad_token_id).to(dtype=text_embeds.dtype)
        text_embeds = text_embeds * vision_mask[:, :, None]
        text_embeds[0, 1 : 1 + num_image_tokens] = image_features.to(text_embeds.dtype)

        input_len = input_ids.shape[1]
        eos_token_ids = self._get_eos_token_ids()
        generated_ids = self._generate_from_prefix_embeddings(
            input_ids=input_ids,
            text_embeds=text_embeds,
            max_new_tokens=max_new_tokens,
            do_sample=do_sample,
            temperature=temperature,
            eos_token_ids=eos_token_ids,
        )

        tokenizer = self._ensure_caption_tokenizer()
        caption, metadata = self._decode_generation_result(
            generated_ids=generated_ids,
            input_len=input_len,
            tokenizer=tokenizer,
            max_new_tokens=max_new_tokens,
            eos_token_ids=eos_token_ids,
        )
        if return_metadata:
            return caption, metadata
        return caption

    @torch.no_grad()
    def generate_answer(
        self,
        image_feats_tensor: torch.Tensor,
        image_coords: torch.Tensor,
        question: str,
        max_new_tokens: int = 512,
        do_sample: bool = True,
        temperature: float = 0.7,
        image_token_counts_by_visual: list[int] | None = None,
        thinking_mode: str = VQA_THINKING_MODE_OFF,
        reasoning_suffix: str = VQA_REASONING_SUFFIX,
        decode_skip_special_tokens: bool | None = None,
        return_metadata: bool = False,
    ) -> str | tuple[str, dict[str, bool | int]]:
        """Generate an answer to a question about an image.

        Qwen3 uses its native chat template. Nemotron uses its own native
        chat template with ``<|im_start|>``, ``<|im_end|>``, and an empty
        ``<think></think>`` no-thinking prefix.

        Args:
            image_feats_tensor: [N, encoder_dim] features for ONE image.
            image_coords: [N, 4 or 5] coordinates.
            question: Question text about the image.
            max_new_tokens: Maximum tokens to generate.
            do_sample: Whether to sample.
            temperature: Sampling temperature if do_sample=True.
            image_token_counts_by_visual: Merged visual-token counts for each visual block.
            thinking_mode: ``off`` forces direct-answer mode, ``on`` requests a
                visible thinking block, and ``raw`` leaves the assistant turn unconstrained.
            reasoning_suffix: User-turn suffix appended in Qwen ``on`` mode. Nemotron
                follows Dense VL generation and uses only the assistant think prefix.
            decode_skip_special_tokens: Whether tokenizer special tokens are hidden
                in the returned answer text. When unset, ``off`` keeps the
                legacy skip-special decode and ``on``/``raw`` preserve tags for
                answer postprocessing.
            return_metadata: Whether to also return generation metadata.

        Returns:
            Generated answer string, or ``(answer, metadata)`` when requested.
        """
        resolved_thinking_mode = normalize_vqa_thinking_mode(thinking_mode)
        resolved_decode_skip_special_tokens = (
            resolved_thinking_mode == VQA_THINKING_MODE_OFF
            if decode_skip_special_tokens is None
            else bool(decode_skip_special_tokens)
        )
        prompt_question = question
        if resolved_thinking_mode == VQA_THINKING_MODE_ON and self.spec.family != NEMOTRON_2B_SPEC.family:
            prompt_question = _append_vqa_reasoning_suffix(question, reasoning_suffix)

        # Spatial merge + position embeddings (same as generate_caption)
        if self.spatial_merger is not None:
            image_features, coords, _ = self.spatial_merger(image_feats_tensor, image_coords)
        elif self.image_proj is not None:
            image_features = self.image_proj(image_feats_tensor)
            coords = image_coords
        else:
            image_features = image_feats_tensor
            coords = image_coords

        if coords.shape[1] == 5:
            coords_2d = coords[:, 2:4]
        else:
            coords_2d = coords[:, 1:3]
        image_features = self.image_pos_embed(image_features, coords_2d)

        num_image_tokens = len(image_features)
        if image_token_counts_by_visual is None:
            visual_token_counts = [num_image_tokens]
        else:
            visual_token_counts = [int(token_count) for token_count in image_token_counts_by_visual]
            if sum(visual_token_counts) != num_image_tokens:
                raise ValueError(
                    f"Visual token counts {visual_token_counts!r} do not sum to merged feature count "
                    f"{num_image_tokens}."
                )
        device = image_features.device

        # Lazy-load tokenizer
        tok = self._ensure_caption_tokenizer()

        def _encode(s: str) -> list[int]:
            return tok.encode(s, add_special_tokens=False)

        answer_decode_prefix_ids: list[int] | None = None
        answer_suppress_token_ids: tuple[int, ...] = ()
        if self.spec.family == QWEN3_SPEC.family:
            system_turn = (
                [QWEN3_IM_START_TOKEN_ID]
                + _encode("system\nYou are a helpful assistant.")
                + [QWEN3_IM_END_TOKEN_ID]
                + _encode("\n")
            )
            user_turn, user_image_pad_offsets = self._build_vqa_user_turn_with_visual_blocks(
                tokenizer=tok,
                question=prompt_question,
                image_token_counts_by_visual=visual_token_counts,
            )
            asst_prefix = [QWEN3_IM_START_TOKEN_ID] + _encode("assistant\n")
            if resolved_thinking_mode == VQA_THINKING_MODE_OFF:
                # Empty think block signals Qwen3 non-thinking mode.
                no_think = [QWEN3_THINK_START_TOKEN_ID] + _encode("\n\n") + [QWEN3_THINK_END_TOKEN_ID] + _encode("\n\n")
                asst_prefix += no_think
                answer_suppress_token_ids = self.spec.suppress_token_ids
            eos_token_ids = self._get_eos_token_ids()
        elif self.spec.family == NEMOTRON_2B_SPEC.family:
            im_start_id = _get_required_token_id(tok, NEMOTRON_2B_IM_START_TOKEN)
            im_end_id = _get_required_token_id(tok, NEMOTRON_2B_IM_END_TOKEN)
            think_id = _get_required_token_id(tok, NEMOTRON_2B_THINK_START_TOKEN)
            end_think_id = _get_required_token_id(tok, NEMOTRON_2B_THINK_END_TOKEN)
            system_turn = [im_start_id] + _encode("system\nYou are a helpful assistant.") + [im_end_id] + _encode("\n")
            user_turn, user_image_pad_offsets = self._build_vqa_user_turn_with_visual_blocks(
                tokenizer=tok,
                question=prompt_question,
                image_token_counts_by_visual=visual_token_counts,
            )
            asst_prefix = [im_start_id] + _encode("assistant\n")
            if resolved_thinking_mode == VQA_THINKING_MODE_OFF:
                asst_prefix += [think_id, end_think_id]
                answer_suppress_token_ids = (think_id, end_think_id)
            elif resolved_thinking_mode == VQA_THINKING_MODE_ON:
                answer_decode_prefix_ids = [think_id] + _encode("\n")
                asst_prefix += answer_decode_prefix_ids
            eos_token_ids = self._get_eos_token_ids()
        else:
            raise NotImplementedError(
                f"VQA generation is not implemented for text decoder family {self.spec.family!r}."
            )

        prompt_ids = system_turn + user_turn + asst_prefix
        input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=device)

        # Embed and inject image features
        text_embeds = self._embed_input_ids(input_ids)
        vision_mask = (input_ids != self.spec.vision_pad_token_id).to(dtype=text_embeds.dtype)
        text_embeds = text_embeds * vision_mask[:, :, None]

        # Image features start inside each multimodal user-turn image_pad block.
        feature_start = 0
        for user_image_pad_offset, visual_token_count in zip(user_image_pad_offsets, visual_token_counts, strict=True):
            ip_start = len(system_turn) + user_image_pad_offset
            feature_end = feature_start + visual_token_count
            text_embeds[0, ip_start : ip_start + visual_token_count] = image_features[feature_start:feature_end].to(
                text_embeds.dtype
            )
            feature_start = feature_end

        # Generate
        # Qwen3 best practice: do NOT use greedy decoding — causes repetitions.
        # Nemotron reuses the same defaults for consistency with captioning.
        input_len = input_ids.shape[1]
        generated_ids = self._generate_from_prefix_embeddings(
            input_ids=input_ids,
            max_new_tokens=max_new_tokens,
            do_sample=do_sample,
            text_embeds=text_embeds,
            temperature=temperature,
            eos_token_ids=eos_token_ids,
            suppress_token_ids=answer_suppress_token_ids,
        )

        answer, metadata = self._decode_generation_result(
            generated_ids=generated_ids,
            input_len=input_len,
            tokenizer=tok,
            max_new_tokens=max_new_tokens,
            eos_token_ids=eos_token_ids,
            skip_special_tokens=resolved_decode_skip_special_tokens,
            decode_prefix_ids=answer_decode_prefix_ids,
        )
        if return_metadata:
            return answer, metadata
        return answer