#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from src/transformers/models/qwen3_vl/modular_qwen3_vl.py.
#               Do NOT edit this file manually as any edits will be overwritten by the generation of
#             the file from the modular. If any change should be done, please apply the change to the
#                          modular_qwen3_vl.py file directly. One of our CI enforces this.
#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# coding=utf-8
# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from dataclasses import dataclass
from typing import Any, Optional, Union
import math 

import torch
import torch.nn as nn
import torch.nn.functional as F
USE_BF16_ROPE = os.environ.get("USE_BF16_ROPE", "0")
from shared.attention import pay_attention
from transformers.activations import ACT2FN
from transformers.cache_utils import Cache, DynamicCache
from transformers.generation import GenerationMixin
from transformers.integrations import use_kernel_forward_from_hub
from transformers.masking_utils import create_causal_mask
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
from transformers.modeling_utils import PreTrainedModel
from transformers.processing_utils import Unpack
from transformers.utils import TransformersKwargs, is_torchdynamo_compiling
from transformers.utils.generic import check_model_inputs
from .qwen3_vl_configuration import Qwen3VLConfig, Qwen3VLTextConfig, Qwen3VLVisionConfig


# Python-level benchmark toggles for the static AR prefix cache. The cache
# stores activations on CPU; MMGP remains the only owner of weight movement.
HIDREAM_AR_CACHE = False
HIDREAM_AR_CACHE_PINNED = False


def auto_docstring(*args, **kwargs):
    if args and len(args) == 1 and callable(args[0]) and not kwargs:
        return args[0]

    def decorator(obj):
        return obj

    return decorator


def _ffn_chunk_size(seq_len, hidden_size, intermediate_size):
    if seq_len <= 1024:
        return 0
    chunk_size = seq_len * hidden_size // max(intermediate_size, 1)
    return max(128, min(seq_len, chunk_size))


def _take_tensor(value):
    if isinstance(value, list):
        tensor = value[0]
        value.clear()
        return tensor
    return value


def _cat_tensors(tensors, dim):
    out = torch.cat(tensors, dim=dim)
    if hasattr(tensors, "clear"):
        tensors.clear()
    return out


def _cache_tensor_cpu(tensor, pinned):
    tensor = tensor.detach()
    if pinned and torch.cuda.is_available():
        out = torch.empty(tensor.shape, dtype=tensor.dtype, device="cpu", pin_memory=True)
        out.copy_(tensor, non_blocking=False)
        return out
    return tensor.to(device="cpu", copy=True)


def _copy_cached_prefix(dst, src, pinned):
    dst.copy_(src, non_blocking=bool(pinned and dst.device.type == "cuda"))


def _static_ar_prefix_len(token_types):
    token_types = token_types[0] if token_types.dim() > 1 else token_types
    gen_positions = token_types.bool()
    ar_len = int((~gen_positions).sum().item())
    if ar_len == 0:
        return 0
    if gen_positions[:ar_len].any() or (~gen_positions[ar_len:]).any():
        return -1
    return ar_len


def _attention_mask_for_pay_attention(attention_mask, q_len, num_heads):
    if attention_mask is None:
        return None
    if attention_mask.dim() == 4:
        if attention_mask.shape[1] == num_heads and attention_mask.shape[2] == q_len:
            return attention_mask.transpose(1, 2)
        if attention_mask.shape[2] in (1, num_heads) and attention_mask.shape[1] == q_len:
            return attention_mask
        if attention_mask.shape[1] == 1 and attention_mask.shape[2] == q_len:
            return attention_mask.transpose(1, 2)
    if attention_mask.dim() == 3:
        return attention_mask[:, :, None, :]
    if attention_mask.dim() == 2:
        return attention_mask[:, None, None, :]
    return attention_mask


def _build_explicit_text_mask(attention_mask, query_length, device):
    causal = torch.ones(query_length, query_length, dtype=torch.bool, device=device)
    causal.tril_()
    causal = causal.view(1, query_length, 1, query_length)
    key_mask = attention_mask[:, None, None, :query_length].to(device=device, dtype=torch.bool)
    return causal & key_mask


def _repeat_kv_for_pay_attention(hidden_states, n_rep):
    hidden_states = _take_tensor(hidden_states)
    if n_rep == 1:
        return hidden_states
    batch, seq_len, num_key_value_heads, head_dim = hidden_states.shape
    hidden_states = hidden_states[:, :, :, None, :].expand(batch, seq_len, num_key_value_heads, n_rep, head_dim)
    return hidden_states.reshape(batch, seq_len, num_key_value_heads * n_rep, head_dim)


def _transpose_contiguous(x_list, dim0, dim1):
    x = _take_tensor(x_list)
    out = x.transpose(dim0, dim1).contiguous()
    del x
    return out


def _apply_rotary_inplace(x, cos, sin, unsqueeze_dim):
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    half = x.shape[-1] // 2
    x1 = x[..., :half]
    x2 = x[..., half:]
    scratch = x1.clone()
    x1.mul_(cos).addcmul_(x2, sin, value=-1)
    x2.mul_(cos).addcmul_(scratch, sin)
    del scratch
    return x


def _pay_attention_token_types(qkv_list, token_types, softmax_scale):
    q, k, v = qkv_list
    qkv_list.clear()
    batch, seq_len = q.shape[:2]
    token_types = token_types.to(q.device)
    if token_types.dim() == 1:
        token_types = token_types.unsqueeze(0)
    if token_types.shape[0] == 1 and batch > 1:
        token_types = token_types.expand(batch, -1)

    ar_len = _static_ar_prefix_len(token_types)
    gen_len = seq_len - ar_len
    if gen_len == 0:
        attn_qkv = [q, k, v]
        q = k = v = None
        return pay_attention(attn_qkv, softmax_scale=softmax_scale, causal=True, force_attention="sdpa", recycle_q=True)

    if ar_len == 0:
        attn_qkv = [q, k, v]
        q = k = v = None
        return pay_attention(attn_qkv, softmax_scale=softmax_scale, causal=False, recycle_q=True)

    out = q
    q = None
    q_attn = out[:, :ar_len]
    k_attn = k[:, :ar_len]
    v_attn = v[:, :ar_len]
    attn_qkv = [q_attn, k_attn, v_attn]
    q_attn = k_attn = v_attn = None
    out_ar = pay_attention(attn_qkv, softmax_scale=softmax_scale, causal=True, force_attention="sdpa", recycle_q=True)
    out[:, :ar_len].copy_(out_ar)
    del attn_qkv, out_ar

    q_attn = out[:, ar_len:]
    k_attn = k
    v_attn = v
    attn_qkv = [q_attn, k_attn, v_attn]
    q_attn = k_attn = v_attn = None
    k = v = None
    out_gen = pay_attention(attn_qkv, softmax_scale=softmax_scale, causal=False, recycle_q=True)
    out[:, ar_len:].copy_(out_gen)
    del attn_qkv, out_gen

    return out


class Qwen3VLVisionMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.linear_fc1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True)
        self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, hidden_state):
        hidden_state = _take_tensor(hidden_state)
        seq_len = hidden_state.shape[-2] if hidden_state.ndim > 1 else hidden_state.shape[0]
        chunk_size = _ffn_chunk_size(seq_len, self.hidden_size, self.intermediate_size)
        if chunk_size == 0:
            hidden = self.linear_fc1(hidden_state)
            del hidden_state
            hidden = self.act_fn(hidden)
            return self.linear_fc2(hidden)
        out = hidden_state.new_empty(*hidden_state.shape[:-1], self.hidden_size)
        for start in range(0, seq_len, chunk_size):
            chunk = hidden_state.narrow(-2, start, min(chunk_size, seq_len - start))
            chunk_out = self.linear_fc1(chunk)
            del chunk
            chunk_out = self.act_fn(chunk_out)
            chunk_out = self.linear_fc2(chunk_out)
            out.narrow(-2, start, chunk_out.shape[-2]).copy_(chunk_out)
            del chunk_out
        return out


class Qwen3VLVisionPatchEmbed(nn.Module):
    def __init__(self, config) -> None:
        super().__init__()
        self.patch_size = config.patch_size
        self.temporal_patch_size = config.temporal_patch_size
        self.in_channels = config.in_channels
        self.embed_dim = config.hidden_size

        kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]
        self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = _take_tensor(hidden_states)
        target_dtype = self.proj.weight.dtype
        hidden_states = hidden_states.view(
            -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
        )
        hidden_states = hidden_states.to(dtype=target_dtype)
        projected = self.proj(hidden_states)
        del hidden_states
        return projected.view(-1, self.embed_dim)


class Qwen3VLVisionRotaryEmbedding(nn.Module):
    inv_freq: torch.Tensor  # fix linting for `register_buffer`

    def __init__(self, dim: int, theta: float = 10000.0) -> None:
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.dim = dim
        self.theta = theta

    def reset_inv_freq(self) -> None:
        inv_freq = 1.0 / (self.theta ** (torch.arange(0, self.dim, 2, dtype=torch.float) / self.dim))
        self.inv_freq = inv_freq

    def forward(self, seqlen: int) -> torch.Tensor:
        seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(seq, self.inv_freq)
        return freqs


class Qwen3VLVisionPatchMerger(nn.Module):
    def __init__(self, config: Qwen3VLVisionConfig, use_postshuffle_norm=False) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size * (config.spatial_merge_size**2)
        self.use_postshuffle_norm = use_postshuffle_norm
        self.norm = nn.LayerNorm(self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6)
        self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size)
        self.act_fn = nn.GELU()
        self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = _take_tensor(x)
        x = self.norm(x.view(-1, self.hidden_size) if self.use_postshuffle_norm else x).view(-1, self.hidden_size)
        hidden = self.linear_fc1(x)
        del x
        hidden = self.act_fn(hidden)
        return self.linear_fc2(hidden)


def apply_rotary_pos_emb_vision(
    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    q = _take_tensor(q)
    k = _take_tensor(k)
    orig_q_dtype = q.dtype
    orig_k_dtype = k.dtype
    q, k = q.float(), k.float()
    cos, sin = cos.float(), sin.float()
    q = _apply_rotary_inplace(q, cos, sin, -2).to(orig_q_dtype)
    k = _apply_rotary_inplace(k, cos, sin, -2).to(orig_k_dtype)
    return q, k


class Qwen3VLVisionAttention(nn.Module):
    def __init__(self, config: Qwen3VLVisionConfig) -> None:
        super().__init__()
        self.dim = config.hidden_size
        self.num_heads = config.num_heads
        self.head_dim = self.dim // self.num_heads
        self.num_key_value_groups = 1  # needed for eager attention
        self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True)
        self.proj = nn.Linear(self.dim, self.dim)
        self.scaling = self.head_dim**-0.5
        self.config = config
        self.attention_dropout = 0.0
        self.is_causal = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        cu_seqlens: torch.Tensor,
        rotary_pos_emb: Optional[torch.Tensor] = None,
        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        **kwargs,
    ) -> torch.Tensor:
        hidden_states = _take_tensor(hidden_states)
        seq_length = hidden_states.shape[0]
        qkv = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3)
        del hidden_states
        query_states, key_states, value_states = qkv.unbind(0)
        del qkv
        cos, sin = position_embeddings
        query_list = [query_states]
        key_list = [key_states]
        query_states = key_states = None
        query_states, key_states = apply_rotary_pos_emb_vision(query_list, key_list, cos, sin)

        query_states = query_states.unsqueeze(0)
        key_states = key_states.unsqueeze(0)
        value_states = value_states.unsqueeze(0)
        lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
        cursor = 0
        for length in lengths:
            end = cursor + length
            qkv_list = [
                query_states[:, cursor:end].contiguous(),
                key_states[:, cursor:end].contiguous(),
                value_states[:, cursor:end].contiguous(),
            ]
            chunk_output = pay_attention(qkv_list, softmax_scale=self.scaling, causal=False, recycle_q=True)
            query_states[:, cursor:end].copy_(chunk_output)
            del chunk_output
            cursor = end
        del key_states, value_states
        attn_output = query_states.reshape(seq_length, -1).contiguous()
        del query_states
        attn_output = self.proj(attn_output)
        return attn_output


class Qwen3VLVisionBlock(nn.Module):
    def __init__(self, config, attn_implementation: str = "sdpa") -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6)
        self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6)
        self.attn = Qwen3VLVisionAttention(config=config)
        self.mlp = Qwen3VLVisionMLP(config=config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        cu_seqlens: torch.Tensor,
        rotary_pos_emb: Optional[torch.Tensor] = None,
        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        **kwargs,
    ) -> torch.Tensor:
        hidden_states = _take_tensor(hidden_states)
        normed_states = self.norm1(hidden_states)
        normed_states_list = [normed_states]
        normed_states = None
        attn_output = self.attn(
            normed_states_list,
            cu_seqlens=cu_seqlens,
            rotary_pos_emb=rotary_pos_emb,
            position_embeddings=position_embeddings,
            **kwargs,
        )
        hidden_states.add_(attn_output)
        del attn_output
        normed_states = self.norm2(hidden_states)
        normed_states_list = [normed_states]
        normed_states = None
        mlp_output = self.mlp(normed_states_list)
        hidden_states.add_(mlp_output)
        del mlp_output
        return hidden_states


class Qwen3VLTextRotaryEmbedding(nn.Module):
    inv_freq: torch.Tensor  # fix linting for `register_buffer`

    def __init__(self, config: Qwen3VLTextConfig, device=None):
        super().__init__()
        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
            self.rope_type = config.rope_scaling.get("rope_type", "default")
        else:
            self.rope_type = "default"
        self.max_seq_len_cached = config.max_position_embeddings
        self.original_max_seq_len = config.max_position_embeddings

        self.config = config
        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]

        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.original_inv_freq = self.inv_freq
        # self.original_inv_freq, _ = self.rope_init_fn(self.config, device)

        self.mrope_section = (config.rope_scaling or {}).get("mrope_section", [24, 20, 20])

    def reset_inv_freq(self) -> None:
        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, None)
        self.inv_freq = inv_freq
        self.original_inv_freq = inv_freq

    def apply_interleaved_mrope(self, freqs, mrope_section):
        """Apply interleaved MRoPE to 3D rotary embeddings.
        Reorganizes frequency layout from chunked [TTT...HHH...WWW] to
        interleaved [THTHWHTHW...TT], preserving frequency continuity.
        args:
            x: (3, bs, seq_len, head_dim // 2)
            mrope_section: (3,)
        returns:
            x_t: (bs, seq_len, head_dim // 2)
        """
        freqs_t = freqs[0].clone()
        for dim, offset in enumerate((1, 2), start=1):  # H, W
            length = mrope_section[dim] * 3
            idx = slice(offset, length, 3)
            freqs_t[..., idx] = freqs[dim, ..., idx]
        return freqs_t

    @torch.no_grad()
    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
    def forward(self, x, position_ids):
        # In contrast to other models, Qwen3VL has different position ids for the grids
        # So we expand the inv_freq to shape (3, ...)
        if position_ids.ndim == 2:
            position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
        if USE_BF16_ROPE == "1":
            inv_freq_expanded = self.inv_freq[None, None, :, None].float().to(device=x.device).expand(3, position_ids.shape[1], -1, 1)
        else:
            inv_freq_expanded = self.original_inv_freq[None, None, :, None].float().to(device=x.device).expand(3, position_ids.shape[1], -1, 1)
        # inv_freq_expanded = self.inv_freq[None, None, :, None].float().to(device=x.device).expand(3, position_ids.shape[1], -1, 1)
        position_ids_expanded = position_ids[:, :, None, :].float()  # shape (3, bs, 1, positions)

        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
            del inv_freq_expanded, position_ids_expanded
            freqs = self.apply_interleaved_mrope(freqs, self.mrope_section)
            cos = freqs.cos()
            freqs.sin_()
            cos.mul_(self.attention_scaling)
            freqs.mul_(self.attention_scaling)

        cos = cos.to(dtype=x.dtype)
        sin = freqs.to(dtype=x.dtype)
        return cos, sin


@use_kernel_forward_from_hub("RMSNorm")
class Qwen3VLTextRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
        """
        Qwen3VLTextRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = _take_tensor(hidden_states)
        return F.rms_norm(hidden_states, self.weight.shape, self.weight, self.variance_epsilon)

    def extra_repr(self):
        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"


def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    """Applies Rotary Position Embedding to the query and key tensors.

    Args:
        q (`torch.Tensor`): The query tensor.
        k (`torch.Tensor`): The key tensor.
        cos (`torch.Tensor`): The cosine part of the rotary embedding.
        sin (`torch.Tensor`): The sine part of the rotary embedding.
        position_ids (`torch.Tensor`, *optional*):
            Deprecated and unused.
        unsqueeze_dim (`int`, *optional*, defaults to 1):
            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
    Returns:
        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
    """
    q = _take_tensor(q)
    k = _take_tensor(k)
    q = _apply_rotary_inplace(q, cos, sin, unsqueeze_dim)
    k = _apply_rotary_inplace(k, cos, sin, unsqueeze_dim)
    return q, k


class Qwen3VLTextAttention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(self, config: Qwen3VLTextConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
        self.scaling = self.head_dim**-0.5
        self.attention_dropout = config.attention_dropout
        self.is_causal = True

        self.q_proj = nn.Linear(
            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
        )
        self.k_proj = nn.Linear(
            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
        )
        self.v_proj = nn.Linear(
            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
        )
        self.o_proj = nn.Linear(
            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
        )
        self.q_norm = Qwen3VLTextRMSNorm(self.head_dim, eps=config.rms_norm_eps)  # unlike olmo, only on the head dim!
        self.k_norm = Qwen3VLTextRMSNorm(
            self.head_dim, eps=config.rms_norm_eps
        )  # thus post q_norm does not need reshape

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor],
        past_key_values: Optional[Cache] = None,
        cache_position: Optional[torch.LongTensor] = None,
        token_types: Optional[torch.Tensor] = None,
        ar_cache: Optional[dict] = None,
        ar_cache_pinned: bool = False,
        ar_len: int = 0,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        hidden_states = _take_tensor(hidden_states)
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)

        query_states = self.q_proj(hidden_states).view(hidden_shape)
        query_list = [query_states]
        query_states = None
        query_states = self.q_norm(query_list).transpose(1, 2)
        key_states = self.k_proj(hidden_states).view(hidden_shape)
        key_list = [key_states]
        key_states = None
        key_states = self.k_norm(key_list).transpose(1, 2)
        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        del hidden_states

        cos, sin = position_embeddings
        query_list = [query_states]
        key_list = [key_states]
        query_states = key_states = None
        query_states, key_states = apply_rotary_pos_emb(query_list, key_list, cos, sin)

        if past_key_values is not None:
            # sin and cos are specific to RoPE models; cache_position needed for the static cache
            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
            del cache_kwargs

        query_list = [query_states]
        query_states = None
        query_states = _transpose_contiguous(query_list, 1, 2)
        key_list = [key_states]
        key_states = None
        key_states = _transpose_contiguous(key_list, 1, 2)
        value_list = [value_states]
        value_states = None
        value_states = _transpose_contiguous(value_list, 1, 2)
        key_list = [key_states]
        key_states = None
        key_states = _repeat_kv_for_pay_attention(key_list, self.num_key_value_groups)
        value_list = [value_states]
        value_states = None
        value_states = _repeat_kv_for_pay_attention(value_list, self.num_key_value_groups)
        if ar_cache is not None and ar_len > 0 and not ar_cache.get("built", False):
            ar_cache["k"] = _cache_tensor_cpu(key_states[:, :ar_len], ar_cache_pinned)
            ar_cache["v"] = _cache_tensor_cpu(value_states[:, :ar_len], ar_cache_pinned)
        qkv_list = [query_states, key_states, value_states]
        del query_states, key_states, value_states
        if token_types is not None:
            attn_output = _pay_attention_token_types(qkv_list, token_types, self.scaling)
        else:
            causal = attention_mask is None
            attention_mask = _attention_mask_for_pay_attention(attention_mask, input_shape[-1], self.config.num_attention_heads)
            attn_output = pay_attention(qkv_list, attention_mask=attention_mask, softmax_scale=self.scaling, causal=causal, recycle_q=True)

        attn_weights = None
        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        attn_output = self.o_proj(attn_output)
        return attn_output, attn_weights

    def forward_with_ar_cache(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        ar_cache: dict,
        ar_len: int,
        ar_cache_pinned: bool = False,
    ) -> torch.Tensor:
        hidden_states = _take_tensor(hidden_states)
        suffix_len = hidden_states.shape[1]
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)

        query_states = self.q_proj(hidden_states).view(hidden_shape)
        query_list = [query_states]
        query_states = None
        query_states = self.q_norm(query_list).transpose(1, 2)
        key_states = self.k_proj(hidden_states).view(hidden_shape)
        key_list = [key_states]
        key_states = None
        key_states = self.k_norm(key_list).transpose(1, 2)
        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        del hidden_states
        cos, sin = position_embeddings
        cos = cos[:, ar_len:]
        sin = sin[:, ar_len:]
        query_list = [query_states]
        key_list = [key_states]
        query_states = key_states = None
        query_states, key_states = apply_rotary_pos_emb(query_list, key_list, cos, sin)

        query_list = [query_states]
        query_states = None
        query_states = _transpose_contiguous(query_list, 1, 2)
        key_list = [key_states]
        key_states = None
        key_states = _transpose_contiguous(key_list, 1, 2)
        value_list = [value_states]
        value_states = None
        value_states = _transpose_contiguous(value_list, 1, 2)
        key_list = [key_states]
        key_states = None
        key_states = _repeat_kv_for_pay_attention(key_list, self.num_key_value_groups)
        value_list = [value_states]
        value_states = None
        value_states = _repeat_kv_for_pay_attention(value_list, self.num_key_value_groups)

        total_len = ar_len + suffix_len
        key_full = key_states.new_empty(key_states.shape[0], total_len, key_states.shape[2], key_states.shape[3])
        value_full = value_states.new_empty(value_states.shape[0], total_len, value_states.shape[2], value_states.shape[3])
        _copy_cached_prefix(key_full[:, :ar_len], ar_cache["k"], ar_cache_pinned)
        _copy_cached_prefix(value_full[:, :ar_len], ar_cache["v"], ar_cache_pinned)
        key_full[:, ar_len:].copy_(key_states)
        value_full[:, ar_len:].copy_(value_states)
        del key_states, value_states

        attn_qkv = [query_states, key_full, value_full]
        del query_states, key_full, value_full
        attn_output = pay_attention(attn_qkv, softmax_scale=self.scaling, causal=False, recycle_q=True)
        del attn_qkv
        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        return self.o_proj(attn_output)


class Qwen3VLTextMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def _forward_chunk(self, x):
        x = _take_tensor(x)
        hidden = self.act_fn(self.gate_proj(x))
        up = self.up_proj(x)
        del x
        hidden.mul_(up)
        del up
        return self.down_proj(hidden)

    def forward(self, x):
        x = _take_tensor(x)
        seq_len = x.shape[-2] if x.ndim > 1 else x.shape[0]
        chunk_size = _ffn_chunk_size(seq_len, self.hidden_size, self.intermediate_size)
        if chunk_size == 0:
            x_list = [x]
            x = None
            return self._forward_chunk(x_list)
        out = x.new_empty(*x.shape[:-1], self.hidden_size)
        for start in range(0, seq_len, chunk_size):
            chunk = x.narrow(-2, start, min(chunk_size, seq_len - start))
            chunk_out = self._forward_chunk(chunk)
            out.narrow(-2, start, chunk_out.shape[-2]).copy_(chunk_out)
            del chunk, chunk_out
        return out


class Qwen3VLTextDecoderLayer(nn.Module):
    def __init__(self, config: Qwen3VLTextConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size

        self.self_attn = Qwen3VLTextAttention(config=config, layer_idx=layer_idx)

        self.mlp = Qwen3VLTextMLP(config)
        self.input_layernorm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        use_cache: Optional[bool] = False,
        cache_position: Optional[torch.LongTensor] = None,
        ar_cache: Optional[dict] = None,
        ar_cache_pinned: bool = False,
        ar_len: int = 0,
        **kwargs: Unpack[TransformersKwargs],
    ) -> torch.Tensor:
        hidden_states = _take_tensor(hidden_states)
        if ar_cache is not None and ar_len > 0 and ar_cache.get("built", False):
            if ar_len >= hidden_states.shape[1]:
                _copy_cached_prefix(hidden_states, ar_cache["hidden"], ar_cache_pinned)
                return hidden_states

            suffix = hidden_states[:, ar_len:]
            residual = suffix
            suffix = self.input_layernorm([suffix])
            suffix_list = [suffix]
            suffix = None
            suffix = self.self_attn.forward_with_ar_cache(
                hidden_states=suffix_list,
                position_embeddings=position_embeddings,
                ar_cache=ar_cache,
                ar_len=ar_len,
                ar_cache_pinned=ar_cache_pinned,
            )
            residual.add_(suffix)
            del residual, suffix

            suffix = hidden_states[:, ar_len:]
            residual = suffix
            suffix = self.post_attention_layernorm([suffix])
            suffix_list = [suffix]
            suffix = None
            suffix = self.mlp(suffix_list)
            residual.add_(suffix)
            del residual, suffix
            _copy_cached_prefix(hidden_states[:, :ar_len], ar_cache["hidden"], ar_cache_pinned)
            return hidden_states

        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        # Self Attention
        hidden_states_list = [hidden_states]
        hidden_states = None
        hidden_states, _ = self.self_attn(
            hidden_states=hidden_states_list,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            cache_position=cache_position,
            position_embeddings=position_embeddings,
            ar_cache=ar_cache,
            ar_cache_pinned=ar_cache_pinned,
            ar_len=ar_len,
            **kwargs,
        )
        hidden_states = residual.add_(hidden_states)
        del residual

        # Fully Connected
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states_list = [hidden_states]
        hidden_states = None
        hidden_states = self.mlp(hidden_states_list)
        hidden_states = residual.add_(hidden_states)
        if ar_cache is not None and ar_len > 0 and not ar_cache.get("built", False):
            ar_cache["hidden"] = _cache_tensor_cpu(hidden_states[:, :ar_len], ar_cache_pinned)
            ar_cache["built"] = True
        return hidden_states


@dataclass
@auto_docstring(
    custom_intro="""
    Base class for Llava outputs, with hidden states and attentions.
    """
)
class Qwen3VLModelOutputWithPast(ModelOutput):
    r"""
    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).

        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
        `past_key_values` input) to speed up sequential decoding.
    rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
        The rope index difference between sequence length and multimodal rope.
    """

    last_hidden_state: Optional[torch.FloatTensor] = None
    past_key_values: Optional[Cache] = None
    hidden_states: Optional[tuple[torch.FloatTensor]] = None
    attentions: Optional[tuple[torch.FloatTensor]] = None
    rope_deltas: Optional[torch.LongTensor] = None
    x_pred: Optional[torch.FloatTensor] = None
    mid_results: Optional[list] = None


@auto_docstring
class Qwen3VLPreTrainedModel(PreTrainedModel):
    config: Qwen3VLConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False
    _no_split_modules = ["Qwen3VLTextDecoderLayer", "Qwen3VLVisionBlock"]
    _skip_keys_device_placement = "past_key_values"
    _supports_flash_attn = False
    _supports_sdpa = True

    _can_compile_fullgraph = True
    _supports_attention_backend = True
    _can_record_outputs = {
        "hidden_states": Qwen3VLTextDecoderLayer,
        "attentions": Qwen3VLTextAttention,
    }


class Qwen3VLVisionModel(Qwen3VLPreTrainedModel):
    config: Qwen3VLVisionConfig
    _no_split_modules = ["Qwen3VLVisionBlock"]

    def __init__(self, config, *inputs, **kwargs) -> None:
        super().__init__(config, *inputs, **kwargs)
        self.spatial_merge_size = config.spatial_merge_size
        self.patch_size = config.patch_size
        self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size

        self.patch_embed = Qwen3VLVisionPatchEmbed(
            config=config,
        )

        self.pos_embed = nn.Embedding(config.num_position_embeddings, config.hidden_size)
        self.num_grid_per_side = int(config.num_position_embeddings**0.5)

        head_dim = config.hidden_size // config.num_heads
        self.rotary_pos_emb = Qwen3VLVisionRotaryEmbedding(head_dim // 2)

        self.blocks = nn.ModuleList([Qwen3VLVisionBlock(config) for _ in range(config.depth)])
        self.merger = Qwen3VLVisionPatchMerger(
            config=config,
            use_postshuffle_norm=False,
        )

        self.deepstack_visual_indexes = config.deepstack_visual_indexes
        self.deepstack_merger_list = nn.ModuleList(
            [
                Qwen3VLVisionPatchMerger(
                    config=config,
                    use_postshuffle_norm=True,
                )
                for _ in range(len(config.deepstack_visual_indexes))
            ]
        )

        self.gradient_checkpointing = False

    def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
        merge_size = self.spatial_merge_size

        max_hw = int(grid_thw[:, 1:].max().item())
        freq_table = self.rotary_pos_emb(max_hw)  # (max_hw, dim // 2)
        device = freq_table.device

        total_tokens = int(torch.prod(grid_thw, dim=1).sum().item())
        pos_ids = torch.empty((total_tokens, 2), dtype=torch.long, device=device)

        offset = 0
        for num_frames, height, width in grid_thw:
            merged_h, merged_w = height // merge_size, width // merge_size

            block_rows = torch.arange(merged_h, device=device)  # block row indices
            block_cols = torch.arange(merged_w, device=device)  # block col indices
            intra_row = torch.arange(merge_size, device=device)  # intra-block row offsets
            intra_col = torch.arange(merge_size, device=device)  # intra-block col offsets

            # Compute full-resolution positions
            row_idx = block_rows[:, None, None, None] * merge_size + intra_row[None, None, :, None]
            col_idx = block_cols[None, :, None, None] * merge_size + intra_col[None, None, None, :]

            row_idx = row_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1)
            col_idx = col_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1)

            coords = torch.stack((row_idx, col_idx), dim=-1)

            if num_frames > 1:
                coords = coords.repeat(num_frames, 1)

            num_tokens = coords.shape[0]
            pos_ids[offset : offset + num_tokens] = coords
            offset += num_tokens
            del block_rows, block_cols, intra_row, intra_col, row_idx, col_idx, coords

        embeddings = freq_table[pos_ids]  # lookup rotary embeddings
        del freq_table, pos_ids
        return embeddings.flatten(1)

    def fast_pos_embed_interpolate(self, grid_thw):
        grid_ts, grid_hs, grid_ws = grid_thw[:, 0], grid_thw[:, 1], grid_thw[:, 2]

        idx_list = [[] for _ in range(4)]
        weight_list = [[] for _ in range(4)]

        for t, h, w in zip(grid_ts, grid_hs, grid_ws):
            h_idxs = torch.linspace(0, self.num_grid_per_side - 1, h)
            w_idxs = torch.linspace(0, self.num_grid_per_side - 1, w)

            h_idxs_floor = h_idxs.int()
            w_idxs_floor = w_idxs.int()
            h_idxs_ceil = (h_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)
            w_idxs_ceil = (w_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)

            dh = h_idxs - h_idxs_floor
            dw = w_idxs - w_idxs_floor

            base_h = h_idxs_floor * self.num_grid_per_side
            base_h_ceil = h_idxs_ceil * self.num_grid_per_side

            indices = [
                (base_h[None].T + w_idxs_floor[None]).flatten(),
                (base_h[None].T + w_idxs_ceil[None]).flatten(),
                (base_h_ceil[None].T + w_idxs_floor[None]).flatten(),
                (base_h_ceil[None].T + w_idxs_ceil[None]).flatten(),
            ]

            weights = [
                ((1 - dh)[None].T * (1 - dw)[None]).flatten(),
                ((1 - dh)[None].T * dw[None]).flatten(),
                (dh[None].T * (1 - dw)[None]).flatten(),
                (dh[None].T * dw[None]).flatten(),
            ]

            for i in range(4):
                idx_list[i].extend(indices[i].tolist())
                weight_list[i].extend(weights[i].tolist())

        idx_tensor = torch.tensor(idx_list, dtype=torch.long, device=self.pos_embed.weight.device)
        weight_tensor = torch.tensor(
            weight_list, dtype=self.pos_embed.weight.dtype, device=self.pos_embed.weight.device
        )
        pos_embeds = self.pos_embed(idx_tensor)
        del idx_tensor
        pos_embeds.mul_(weight_tensor[:, :, None])
        del weight_tensor
        patch_pos_embeds = pos_embeds[0].clone()
        patch_pos_embeds.add_(pos_embeds[1]).add_(pos_embeds[2]).add_(pos_embeds[3])
        del pos_embeds

        patch_pos_embeds = patch_pos_embeds.split([h * w for h, w in zip(grid_hs, grid_ws)])

        patch_pos_embeds_permute = []
        merge_size = self.config.spatial_merge_size
        for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws):
            pos_embed = pos_embed.repeat(t, 1)
            pos_embed = (
                pos_embed.view(t, h // merge_size, merge_size, w // merge_size, merge_size, -1)
                .permute(0, 1, 3, 2, 4, 5)
                .flatten(0, 4)
            )
            patch_pos_embeds_permute.append(pos_embed)
        del patch_pos_embeds
        return _cat_tensors(patch_pos_embeds_permute, dim=0)

    def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor:
        """
        Args:
            hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`):
                The final hidden states of the model.
            grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`):
                The temporal, height and width of feature shape of each image in LLM.

        Returns:
            `torch.Tensor`: hidden_states.
        """
        hidden_states = _take_tensor(hidden_states)
        hidden_states_list = [hidden_states]
        hidden_states = None
        hidden_states = self.patch_embed(hidden_states_list)

        pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
        hidden_states.add_(pos_embeds)
        del pos_embeds

        rotary_pos_emb = self.rot_pos_emb(grid_thw)

        seq_len, _ = hidden_states.size()
        hidden_states = hidden_states.reshape(seq_len, -1)
        rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
        rotary_cos = rotary_pos_emb.cos()
        rotary_pos_emb.sin_()
        position_embeddings = (rotary_cos, rotary_pos_emb)
        del rotary_cos, rotary_pos_emb

        cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0])
        # FA2 requires int32, while ONNX export requires the same dtype as grid_thw.
        cu_seqlens = cu_seqlens.to(dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32)
        cu_seqlens.cumsum_(dim=0)
        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)

        deepstack_feature_lists = []
        for layer_num, blk in enumerate(self.blocks):
            hidden_states_list = [hidden_states]
            hidden_states = None
            hidden_states = blk(
                hidden_states_list,
                cu_seqlens=cu_seqlens,
                position_embeddings=position_embeddings,
                **kwargs,
            )
            if getattr(self, "_interrupt", False):
                return None, None
            if layer_num in self.deepstack_visual_indexes:
                deepstack_feature = self.deepstack_merger_list[self.deepstack_visual_indexes.index(layer_num)](
                    hidden_states
                )
                deepstack_feature_lists.append(deepstack_feature)

        del cu_seqlens, position_embeddings
        hidden_states_list = [hidden_states]
        hidden_states = None
        hidden_states = self.merger(hidden_states_list)

        return hidden_states, deepstack_feature_lists


@auto_docstring(
    custom_intro=(
        "Text part of Qwen3VL, "
        "not a pure text-only model, as DeepStack integrates visual features into the early hidden states."
    )
)
class Qwen3VLTextModel(Qwen3VLPreTrainedModel):
    config: Qwen3VLTextConfig
    _no_split_modules = ["Qwen3VLTextDecoderLayer"]

    def __init__(self, config: Qwen3VLTextConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList(
            [Qwen3VLTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        self.norm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = Qwen3VLTextRotaryEmbedding(config=config)
        self.gradient_checkpointing = False

        # Initialize weights and apply final processing
        self.post_init()

    def forward_joint(
        self,
        inputs_embeds: list[torch.FloatTensor],
        position_ids: list[torch.LongTensor],
        token_types: list[torch.Tensor],
        return_mid_results_layers: Optional[list] = None,
        ar_cache: Optional[dict] = None,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> BaseModelOutputWithPast:
        hidden_states_list = list(inputs_embeds)
        if hasattr(inputs_embeds, "clear"):
            inputs_embeds.clear()
        position_ids_list = list(position_ids)
        if hasattr(position_ids, "clear"):
            position_ids.clear()
        token_types_list = list(token_types)
        if hasattr(token_types, "clear"):
            token_types.clear()
        text_position_ids_list = []
        position_embeddings_list = []
        mid_results = [] if return_mid_results_layers else None
        ar_cache_layers = ar_cache.get("layers") if ar_cache is not None else None
        ar_lens = ar_cache.get("ar_lens") if ar_cache is not None else None
        ar_cache_pinned = bool(ar_cache.get("pinned", False)) if ar_cache is not None else False

        for idx, hidden_states in enumerate(hidden_states_list):
            cur_position_ids = position_ids_list[idx]
            if cur_position_ids.ndim == 2:
                cur_position_ids = cur_position_ids[None, ...].expand(3, cur_position_ids.shape[0], -1)
            if cur_position_ids.ndim == 3 and cur_position_ids.shape[0] == 4:
                text_position_ids = cur_position_ids[0]
                cur_position_ids = cur_position_ids[1:]
            else:
                text_position_ids = cur_position_ids[0]
            text_position_ids_list.append(text_position_ids)
            position_embeddings_list.append(self.rotary_emb(hidden_states, cur_position_ids))
            token_types_list[idx] = token_types_list[idx].to(hidden_states.device)

        for layer_idx, decoder_layer in enumerate(self.layers):
            for idx, hidden_states in enumerate(hidden_states_list):
                layer_cache = ar_cache_layers[idx][layer_idx] if ar_cache_layers is not None else None
                ar_len = ar_lens[idx] if ar_lens is not None else 0
                hidden_states_list[idx] = None
                if hasattr(decoder_layer.forward, "__wrapped__"):
                    layer_input = hidden_states
                else:
                    layer_input = [hidden_states]
                    hidden_states = None
                hidden_states_list[idx] = decoder_layer(
                    layer_input,
                    attention_mask=None,
                    position_ids=text_position_ids_list[idx],
                    past_key_values=None,
                    cache_position=None,
                    position_embeddings=position_embeddings_list[idx],
                    token_types=token_types_list[idx],
                    ar_cache=layer_cache,
                    ar_cache_pinned=ar_cache_pinned,
                    ar_len=ar_len,
                    **kwargs,
                )
                hidden_states = None
                if getattr(self, "_interrupt", False):
                    return BaseModelOutputWithPast(last_hidden_state=None, past_key_values=None)

            if return_mid_results_layers is not None and layer_idx in return_mid_results_layers:
                mid_results.append([hidden_states.clone() for hidden_states in hidden_states_list])

        del position_embeddings_list, text_position_ids_list, token_types_list
        for idx, hidden_states in enumerate(hidden_states_list):
            hidden_states_list[idx] = None
            norm_list = [hidden_states]
            hidden_states = None
            hidden_states_list[idx] = self.norm(norm_list)

        output = BaseModelOutputWithPast(last_hidden_state=hidden_states_list, past_key_values=None)
        output.mid_results = mid_results
        return output

    @check_model_inputs
    @auto_docstring
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        # args for deepstack
        visual_pos_masks: Optional[torch.Tensor] = None,
        deepstack_visual_embeds: Optional[list[torch.Tensor]] = None,
        return_mid_results_layers: Optional[list] = None,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> Union[tuple, BaseModelOutputWithPast]:
        r"""
        visual_pos_masks (`torch.Tensor` of shape `(batch_size, seqlen)`, *optional*):
            The mask of the visual positions.
        deepstack_visual_embeds (`list[torch.Tensor]`, *optional*):
            The deepstack visual embeddings. The shape is (num_layers, visual_seqlen, embed_dim).
            The feature is extracted from the different visual encoder layers, and fed to the decoder
            hidden states. It's from the paper DeepStack(https://arxiv.org/abs/2406.04334).
        """
        token_types = kwargs.pop("token_types", None)
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        # torch.jit.trace() doesn't support cache objects in the output
        if use_cache and past_key_values is None and not torch.jit.is_tracing():
            past_key_values = DynamicCache(config=self.config)

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)
        else:
            inputs_embeds = _take_tensor(inputs_embeds)
        input_ids = None

        if cache_position is None:
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
            cache_position = torch.arange(
                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
            )

        # the hard coded `3` is for temporal, height and width.
        if position_ids is None:
            position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
        elif position_ids.ndim == 2:
            position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)

        if position_ids.ndim == 3 and position_ids.shape[0] == 4:
            text_position_ids = position_ids[0]
            position_ids = position_ids[1:]
        else:
            text_position_ids = position_ids[0]

        if token_types is None:
            if past_key_values is None:
                if attention_mask is not None:
                    attention_mask = _build_explicit_text_mask(attention_mask, inputs_embeds.shape[1], inputs_embeds.device)
            else:
                attention_mask = create_causal_mask(
                    config=self.config,
                    input_embeds=inputs_embeds,
                    attention_mask=attention_mask,
                    cache_position=cache_position,
                    past_key_values=past_key_values,
                    position_ids=text_position_ids,
                )
        else:
            token_types = token_types.to(inputs_embeds.device)
            attention_mask = None

        hidden_states = inputs_embeds
        inputs_embeds = None
        mid_results = [] if return_mid_results_layers else None

        # create position embeddings to be shared across the decoder layers
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        # decoder layers
        for layer_idx, decoder_layer in enumerate(self.layers):
            if hasattr(decoder_layer.forward, "__wrapped__"):
                layer_input = hidden_states
            else:
                layer_input = [hidden_states]
                hidden_states = None
            hidden_states = decoder_layer(
                layer_input,
                attention_mask=attention_mask,
                position_ids=text_position_ids,
                past_key_values=past_key_values,
                cache_position=cache_position,
                position_embeddings=position_embeddings,
                token_types=token_types,
                **kwargs,
            )

            # add visual features to the hidden states of first several layers
            if deepstack_visual_embeds is not None and layer_idx in range(len(deepstack_visual_embeds)):
                visual_embeds = deepstack_visual_embeds[layer_idx]
                deepstack_visual_embeds[layer_idx] = None
                hidden_states = self._deepstack_process(
                    hidden_states,
                    visual_pos_masks,
                    visual_embeds,
                )

            if return_mid_results_layers is not None and layer_idx in return_mid_results_layers:
                mid_results.append(hidden_states.clone())

            if getattr(self, "_interrupt", False):
                return BaseModelOutputWithPast(last_hidden_state=None, past_key_values=past_key_values)

        if hasattr(deepstack_visual_embeds, "clear"):
            deepstack_visual_embeds.clear()
        del attention_mask, position_embeddings
        hidden_states_list = [hidden_states]
        hidden_states = None
        hidden_states = self.norm(hidden_states_list)

        output = BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values,
        )
        output.mid_results = mid_results
        return output

    def _deepstack_process(
        self, hidden_states: torch.Tensor, visual_pos_masks: torch.Tensor, visual_embeds: torch.Tensor
    ):
        visual_pos_masks = visual_pos_masks.to(hidden_states.device)
        visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype)
        local_this = hidden_states[visual_pos_masks, :].clone()
        local_this.add_(visual_embeds)
        del visual_embeds
        hidden_states[visual_pos_masks, :] = local_this
        del local_this
        return hidden_states

class BottleneckPatchEmbed(nn.Module):
    def __init__(self, config, patch_size=16, in_chans=3, pca_dim=768, embed_dim=768, bias=True):
        super().__init__()
        self.proj1 = nn.Linear(patch_size*patch_size*in_chans, pca_dim, bias=False)
        self.proj2 = nn.Linear(pca_dim, embed_dim, bias=bias)
        self.initialize_weights()

    def initialize_weights(self):
        w1 = self.proj1.weight.data
        nn.init.xavier_uniform_(w1.view([w1.shape[0], -1]))
        w2 = self.proj2.weight.data
        nn.init.xavier_uniform_(w2.view([w2.shape[0], -1]))
        nn.init.constant_(self.proj2.bias, 0)

    def forward(self, x):
        x = _take_tensor(x)
        hidden = self.proj1(x)
        del x
        return self.proj2(hidden)
    
class FinalLayer(nn.Module):
    def __init__(self, config, hidden_size, patch_size, out_channels):
        super().__init__()
        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            nn.init.zeros_(m.weight)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)

    def forward(self, x, adaln_input=None):
        x = _take_tensor(x)
        return self.linear(x)



class TimestepEmbedder(nn.Module):
    """
    Embeds scalar timesteps into vector representations.
    """
    def __init__(self, config, hidden_size, frequency_embedding_size=256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size, bias=True),
        )
        nn.init.normal_(self.mlp[0].weight, std=0.02)
        nn.init.normal_(self.mlp[2].weight, std=0.02)
        self.frequency_embedding_size = frequency_embedding_size

    @staticmethod
    def timestep_embedding(t, dim, max_period=10000):
        """
        Create sinusoidal timestep embeddings.
        :param t: a 1-D Tensor of N indices, one per batch element.
                          These may be fractional.
        :param dim: the dimension of the output.
        :param max_period: controls the minimum frequency of the embeddings.
        :return: an (N, D) Tensor of positional embeddings.
        """
        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
        half = dim // 2
        freqs = torch.exp(
            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
        ).to(device=t.device)
        args = t[:, None].float() * freqs[None]
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        if dim % 2:
            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
        return embedding

    def forward(self, t):
        t_freq = self.timestep_embedding(t * 1000, self.frequency_embedding_size)
        t_emb = self.mlp(t_freq.to(self.mlp[0].weight.dtype))
        return t_emb

@auto_docstring
class Qwen3VLModel(Qwen3VLPreTrainedModel):
    base_model_prefix = ""
    _checkpoint_conversion_mapping = {}
    # Reference: fix gemma3 grad acc #37208
    accepts_loss_kwargs = False
    config: Qwen3VLConfig
    _no_split_modules = ["Qwen3VLTextDecoderLayer", "Qwen3VLVisionBlock"]

    def __init__(self, config):
        super().__init__(config)
        self.visual = Qwen3VLVisionModel._from_config(config.vision_config)
        self.language_model = Qwen3VLTextModel._from_config(config.text_config)
        self.rope_deltas = None  # cache rope_deltas here
        self._ar_cache = None

        self.patch_size = 32
        self.in_channels = 3
        hidden_size = config.text_config.hidden_size
        bottleneck_dim = hidden_size // 4

        self.t_embedder1 = TimestepEmbedder(self.config, hidden_size)
        self.x_embedder = BottleneckPatchEmbed(self.config,  patch_size = self.patch_size, in_chans = self.in_channels, pca_dim = bottleneck_dim, embed_dim = hidden_size, bias=True)
    
        # self.t_embedder2 = TimestepEmbedder(self.config, hidden_size)
        self.t_embedder2 = None
        self.final_layer2 = FinalLayer(self.config, hidden_size = hidden_size, patch_size = self.patch_size, out_channels = self.in_channels)
        self.tms_token_id = 151673

        # Initialize weights and apply final processing
        self.post_init()

    def clear_runtime_caches(self):
        self.rope_deltas = None
        self._ar_cache = None

    def _get_ar_cache(self, sample_count, layer_count, ar_lens, pinned):
        signature = (sample_count, layer_count, tuple(ar_lens), bool(pinned))
        if self._ar_cache is None or self._ar_cache.get("signature") != signature:
            self._ar_cache = {
                "signature": signature,
                "ar_lens": list(ar_lens),
                "pinned": bool(pinned),
                "layers": [[{} for _ in range(layer_count)] for _ in range(sample_count)],
            }
        return self._ar_cache

    def get_input_embeddings(self):
        return self.language_model.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.language_model.set_input_embeddings(value)

    def set_decoder(self, decoder):
        self.language_model = decoder

    def get_decoder(self):
        return self.language_model

    def get_rope_index(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        image_grid_thw: Optional[torch.LongTensor] = None,
        video_grid_thw: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Different from the original implementation, Qwen3VL use timestamps rather than absolute time position ids."""

        # Since we use timestamps to seperate videos, like <t1> <vision_start> <frame1> <vision_end> <t2> <vision_start> <frame2> <vision_end>, the video_grid_thw should also be split
        if video_grid_thw is not None:
            video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0)
            video_grid_thw[:, 0] = 1

        spatial_merge_size = self.config.vision_config.spatial_merge_size
        image_token_id = self.config.image_token_id
        video_token_id = self.config.video_token_id
        vision_start_token_id = self.config.vision_start_token_id
        mrope_position_deltas = []
        if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
            total_input_ids = input_ids
            if attention_mask is None:
                attention_mask = torch.ones_like(total_input_ids)
            position_ids = torch.ones(
                3,
                input_ids.shape[0],
                input_ids.shape[1],
                dtype=input_ids.dtype,
                device=input_ids.device,
            )
            image_index, video_index = 0, 0
            attention_mask = attention_mask.to(total_input_ids.device)
            for i, input_ids in enumerate(total_input_ids):
                input_ids = input_ids[attention_mask[i] == 1]
                image_nums, video_nums = 0, 0
                vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
                vision_tokens = input_ids[vision_start_indices + 1]
                image_nums = (vision_tokens == image_token_id).sum()
                video_nums = (vision_tokens == video_token_id).sum()
                input_tokens = input_ids.tolist()
                llm_pos_ids_list: list = []
                st = 0
                remain_images, remain_videos = image_nums, video_nums
                for _ in range(image_nums + video_nums):
                    if image_token_id in input_tokens and remain_images > 0:
                        ed_image = input_tokens.index(image_token_id, st)
                    else:
                        ed_image = len(input_tokens) + 1
                    if video_token_id in input_tokens and remain_videos > 0:
                        ed_video = input_tokens.index(video_token_id, st)
                    else:
                        ed_video = len(input_tokens) + 1
                    if ed_image < ed_video:
                        t, h, w = (
                            image_grid_thw[image_index][0],
                            image_grid_thw[image_index][1],
                            image_grid_thw[image_index][2],
                        )
                        image_index += 1
                        remain_images -= 1
                        ed = ed_image

                    else:
                        t, h, w = (
                            video_grid_thw[video_index][0],
                            video_grid_thw[video_index][1],
                            video_grid_thw[video_index][2],
                        )
                        video_index += 1
                        remain_videos -= 1
                        ed = ed_video
                    llm_grid_t, llm_grid_h, llm_grid_w = (
                        t.item(),
                        h.item() // spatial_merge_size,
                        w.item() // spatial_merge_size,
                    )
                    text_len = ed - st

                    st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
                    llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)

                    # t_index is always 0 because llm_grid_t is always 1 (we use timestamps to encode the temporal information for videos)
                    t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()
                    h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
                    w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
                    llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
                    st = ed + llm_grid_t * llm_grid_h * llm_grid_w

                if st < len(input_tokens):
                    st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
                    text_len = len(input_tokens) - st
                    llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)

                llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
                position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
                mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))
            mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
            return position_ids, mrope_position_deltas
        else:
            if attention_mask is not None:
                position_ids = attention_mask.long().cumsum(-1)
                position_ids.sub_(1)
                position_ids.masked_fill_(attention_mask == 0, 1)
                position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
                max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
                mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
            else:
                position_ids = (
                    torch.arange(input_ids.shape[1], device=input_ids.device)
                    .view(1, 1, -1)
                    .expand(3, input_ids.shape[0], -1)
                )
                mrope_position_deltas = torch.zeros(
                    [input_ids.shape[0], 1],
                    device=input_ids.device,
                    dtype=input_ids.dtype,
                )

            return position_ids, mrope_position_deltas

    def get_video_features(
        self, pixel_values_videos: torch.FloatTensor, video_grid_thw: Optional[torch.LongTensor] = None
    ):
        """
        Encodes videos into continuous embeddings that can be forwarded to the language model. The deepstack visual features are also returned.

        Args:
            pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
                The tensors corresponding to the input videos.
            video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
                The temporal, height and width of feature shape of each video in LLM.
        """
        # Same implementation as for images
        return self.get_image_features(pixel_values_videos, video_grid_thw)

    def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor] = None):
        """
        Encodes images into continuous embeddings that can be forwarded to the language model. The deepstack visual features are also returned.

        Args:
            pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
                The tensors corresponding to the input images.
            image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
                The temporal, height and width of feature shape of each image in LLM.
        """
        pixel_values = pixel_values.type(self.visual.dtype)
        pixel_values_list = [pixel_values]
        pixel_values = None
        image_embeds, deepstack_image_embeds = self.visual(pixel_values_list, grid_thw=image_grid_thw)
        if image_embeds is None:
            return None, None
        split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist()
        image_embeds = torch.split(image_embeds, split_sizes)
        return image_embeds, deepstack_image_embeds

    def get_placeholder_mask(
        self,
        input_ids: torch.LongTensor,
        inputs_embeds: torch.FloatTensor,
        image_features: Optional[torch.FloatTensor] = None,
        video_features: Optional[torch.FloatTensor] = None,
    ):
        """
        Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
        equal to the length of multimodal features. If the lengths are different, an error is raised.
        """
        if input_ids is None:
            special_image_mask = inputs_embeds == self.get_input_embeddings()(
                torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
            )
            special_image_mask = special_image_mask.all(-1)
            special_video_mask = inputs_embeds == self.get_input_embeddings()(
                torch.tensor(self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device)
            )
            special_video_mask = special_video_mask.all(-1)
        else:
            special_image_mask = input_ids == self.config.image_token_id
            special_video_mask = input_ids == self.config.video_token_id

        n_image_tokens = special_image_mask.sum()
        special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
        if image_features is not None and n_image_tokens * inputs_embeds.shape[-1] != image_features.numel():
            raise ValueError(
                f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {image_features.shape[0]}"
            )

        n_video_tokens = special_video_mask.sum()
        special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
        if video_features is not None and n_video_tokens * inputs_embeds.shape[-1] != video_features.numel():
            raise ValueError(
                f"Videos features and video tokens do not match: tokens: {n_video_tokens}, features {video_features.shape[0]}"
            )

        return special_image_mask, special_video_mask

    def _forward_generation(self, input_ids, position_ids, vinputs, timestep, token_types,
                             attention_mask=None, pixel_values=None, pixel_values_videos=None,
                             image_grid_thw=None, video_grid_thw=None, image_embeds=None, video_embeds=None,
                             prediction_mask=None,
                             return_mid_results_layers=None,
                             **kwargs):
        """Forward pass for image generation (denoising step).

        Args:
            input_ids: [batch, txt_seq_len] - text token IDs (without image tokens)
            position_ids: [3, batch, total_seq_len] - 3D RoPE positions covering text+image
            vinputs: [batch, img_tokens, patch_dim] - patchified noise input
            timestep: [batch] - timestep values (scalar per sample)
            token_types: [batch, total_seq_len] or [total_seq_len, 1] - 0=AR, >0=gen
            attention_mask: ignored (created internally for non-flash path)
            pixel_values: optional image pixel values for conditioned generation
            pixel_values_videos: optional video pixel values
            image_grid_thw: optional image grid info
            video_grid_thw: optional video grid info

        Returns:
            Qwen3VLModelOutputWithPast with x_pred field set.
        """
        is_joint = isinstance(input_ids, (list, tuple))
        sample_count = len(input_ids) if is_joint else 1

        def pick(value, idx, consume=False):
            if not isinstance(value, (list, tuple)):
                return value
            item = value[idx]
            if consume and isinstance(value, list):
                value[idx] = None
            return item

        def prepare_sample(idx):
            cur_input_ids = pick(input_ids, idx)
            cur_pixel_values = pick(pixel_values, idx, consume=True)
            cur_pixel_values_videos = pick(pixel_values_videos, idx, consume=True)
            cur_image_grid_thw = pick(image_grid_thw, idx)
            cur_video_grid_thw = pick(video_grid_thw, idx)
            cur_image_embeds = pick(image_embeds, idx, consume=True)
            cur_video_embeds = pick(video_embeds, idx, consume=True)
            cur_timestep = pick(timestep, idx)
            cur_vinputs = pick(vinputs, idx, consume=True)
            cur_token_types = pick(token_types, idx)
            cur_prediction_mask = pick(prediction_mask, idx)

            inputs_embeds = self.get_input_embeddings()(cur_input_ids)
            if cur_image_embeds is None and cur_pixel_values is not None:
                cur_image_embeds, _ = self.get_image_features(cur_pixel_values, cur_image_grid_thw)
                cur_pixel_values = None
                if cur_image_embeds is None:
                    return None
                image_embed_list = list(cur_image_embeds)
                cur_image_embeds = None
                cur_image_embeds = _cat_tensors(image_embed_list, dim=0)
            if cur_image_embeds is not None:
                cur_image_embeds = cur_image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
                image_mask, _ = self.get_placeholder_mask(cur_input_ids, inputs_embeds=inputs_embeds, image_features=cur_image_embeds)
                inputs_embeds.masked_scatter_(image_mask, cur_image_embeds)
                del image_mask, cur_image_embeds

            if cur_video_embeds is None and cur_pixel_values_videos is not None:
                cur_video_embeds, _ = self.get_video_features(cur_pixel_values_videos, cur_video_grid_thw)
                cur_pixel_values_videos = None
                if cur_video_embeds is None:
                    return None
                video_embed_list = list(cur_video_embeds)
                cur_video_embeds = None
                cur_video_embeds = _cat_tensors(video_embed_list, dim=0)
            if cur_video_embeds is not None:
                cur_video_embeds = cur_video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
                _, video_mask = self.get_placeholder_mask(cur_input_ids, inputs_embeds=inputs_embeds, video_features=cur_video_embeds)
                inputs_embeds.masked_scatter_(video_mask, cur_video_embeds)
                del video_mask, cur_video_embeds

            batch_size = cur_vinputs.shape[0]
            if inputs_embeds.shape[0] == 1 and batch_size > 1:
                inputs_embeds = inputs_embeds.expand(batch_size, -1, -1).clone()
            cur_timestep = cur_timestep.to(inputs_embeds.device)
            t_emb = self.t_embedder1(cur_timestep)
            cur_timestep = None
            tms_mask = cur_input_ids == self.tms_token_id
            cur_input_ids = None
            if tms_mask.shape[0] == 1 and batch_size > 1:
                tms_mask = tms_mask.expand(batch_size, -1)
            t_emb_expanded = t_emb.unsqueeze(1).expand_as(inputs_embeds)
            inputs_embeds[tms_mask] = t_emb_expanded[tms_mask]
            del t_emb, tms_mask, t_emb_expanded

            cur_vinputs = cur_vinputs.to(inputs_embeds.device)
            cur_vinputs_list = [cur_vinputs]
            cur_vinputs = None
            vinputs_embedded = self.x_embedder(cur_vinputs_list).to(inputs_embeds.dtype)
            embed_list = [inputs_embeds, vinputs_embedded]
            inputs_embeds = vinputs_embedded = None
            inputs_embeds = _cat_tensors(embed_list, dim=1)
            batch_size, total_seq_len, _ = inputs_embeds.shape

            cur_token_types = cur_token_types.to(inputs_embeds.device)
            if cur_token_types.dim() == 1:
                cur_token_types = cur_token_types.unsqueeze(0)
            elif cur_token_types.dim() == 2 and cur_token_types.shape[-1] == 1 and cur_token_types.shape[0] == total_seq_len:
                cur_token_types = cur_token_types.squeeze(-1).unsqueeze(0)
            if cur_token_types.shape[0] == 1 and batch_size > 1:
                cur_token_types = cur_token_types.expand(batch_size, -1)
            return inputs_embeds, pick(position_ids, idx), cur_token_types, cur_prediction_mask

        prepared = [prepare_sample(i) for i in range(sample_count)]
        if any(item is None for item in prepared):
            return Qwen3VLModelOutputWithPast(last_hidden_state=None, x_pred=None, mid_results=None)

        ar_cache = None
        if HIDREAM_AR_CACHE and return_mid_results_layers is None:
            ar_lens = [_static_ar_prefix_len(item[2]) for item in prepared]
            if any(ar_len > 0 for ar_len in ar_lens) and all(ar_len >= 0 for ar_len in ar_lens):
                ar_cache = self._get_ar_cache(sample_count, len(self.language_model.layers), ar_lens, HIDREAM_AR_CACHE_PINNED)

        if is_joint or ar_cache is not None:
            inputs_embeds_list = [item[0] for item in prepared]
            position_ids_list = [item[1] for item in prepared]
            token_types_list = [item[2] for item in prepared]
            prediction_mask_list = [item[3] for item in prepared]
            prepared.clear()
            prepared = None
            outputs = self.language_model.forward_joint(
                inputs_embeds=inputs_embeds_list,
                position_ids=position_ids_list,
                return_mid_results_layers=return_mid_results_layers,
                token_types=token_types_list,
                ar_cache=ar_cache,
                **kwargs,
            )
            hidden_states_list = outputs.last_hidden_state
            mid_results = getattr(outputs, "mid_results", None)
            outputs = None
            if hidden_states_list is None:
                return Qwen3VLModelOutputWithPast(last_hidden_state=None, x_pred=None, mid_results=mid_results)
            x_pred_list = []
            for idx, (hidden_states, cur_prediction_mask) in enumerate(zip(hidden_states_list, prediction_mask_list)):
                hidden_states_list[idx] = None
                prediction_mask_list[idx] = None
                batch_size = hidden_states.shape[0]
                if cur_prediction_mask is not None:
                    cur_prediction_mask = cur_prediction_mask.to(hidden_states.device)
                    if cur_prediction_mask.dim() == 1:
                        cur_prediction_mask = cur_prediction_mask.unsqueeze(0)
                    if cur_prediction_mask.shape[0] == 1 and batch_size > 1:
                        cur_prediction_mask = cur_prediction_mask.expand(batch_size, -1)
                    hidden_states = hidden_states[cur_prediction_mask].view(batch_size, -1, hidden_states.shape[-1])
                hidden_states_list_for_head = [hidden_states]
                hidden_states = None
                x_pred_list.append(self.final_layer2(hidden_states_list_for_head))
            if not is_joint:
                return Qwen3VLModelOutputWithPast(last_hidden_state=None, x_pred=x_pred_list[0], mid_results=mid_results)
            return Qwen3VLModelOutputWithPast(last_hidden_state=None, x_pred=x_pred_list, mid_results=mid_results)

        inputs_embeds, position_ids, token_types, prediction_mask = prepared.pop()
        prepared = None
        inputs_embeds_list = [inputs_embeds]
        inputs_embeds = None
        outputs = self.language_model(
            input_ids=None,
            position_ids=position_ids,
            attention_mask=None,
            inputs_embeds=inputs_embeds_list,
            use_cache=False,
            return_mid_results_layers=return_mid_results_layers,
            token_types=token_types,
        )
        hidden_states = outputs.last_hidden_state
        mid_results = getattr(outputs, "mid_results", None)
        outputs = None
        if hidden_states is None:
            return Qwen3VLModelOutputWithPast(last_hidden_state=None, x_pred=None, mid_results=mid_results)

        batch_size = hidden_states.shape[0]
        if prediction_mask is not None:
            prediction_mask = prediction_mask.to(hidden_states.device)
            if prediction_mask.dim() == 1:
                prediction_mask = prediction_mask.unsqueeze(0)
            if prediction_mask.shape[0] == 1 and batch_size > 1:
                prediction_mask = prediction_mask.expand(batch_size, -1)
            hidden_states = hidden_states[prediction_mask].view(batch_size, -1, hidden_states.shape[-1])
        hidden_states_list = [hidden_states]
        hidden_states = None
        x_pred = self.final_layer2(hidden_states_list)

        return Qwen3VLModelOutputWithPast(last_hidden_state=None, x_pred=x_pred, mid_results=mid_results)

    @auto_docstring
    @check_model_inputs
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        pixel_values: Optional[torch.Tensor] = None,
        pixel_values_videos: Optional[torch.FloatTensor] = None,
        image_grid_thw: Optional[torch.LongTensor] = None,
        video_grid_thw: Optional[torch.LongTensor] = None,
        cache_position: Optional[torch.LongTensor] = None,
        vinputs: Optional[torch.Tensor] = None,
        timestep: Optional[torch.Tensor] = None,
        token_types: Optional[torch.Tensor] = None,
        image_embeds: Optional[torch.Tensor] = None,
        video_embeds: Optional[torch.Tensor] = None,
        prediction_mask: Optional[torch.Tensor] = None,
        return_mid_results_layers: Optional[list] = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> Union[tuple, Qwen3VLModelOutputWithPast]:
        r"""
        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
            The temporal, height and width of feature shape of each image in LLM.
        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
            The temporal, height and width of feature shape of each video in LLM.
        """
        # Dispatch to generation forward if vinputs is provided
        if vinputs is not None:
            return self._forward_generation(
                input_ids=input_ids, position_ids=position_ids,
                vinputs=vinputs, timestep=timestep, token_types=token_types,
                attention_mask=attention_mask,
                pixel_values=pixel_values, pixel_values_videos=pixel_values_videos,
                image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw,
                image_embeds=image_embeds, video_embeds=video_embeds,
                prediction_mask=prediction_mask,
                return_mid_results_layers=return_mid_results_layers,
                **kwargs)

        pixel_values = _take_tensor(pixel_values)
        pixel_values_videos = _take_tensor(pixel_values_videos)
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if inputs_embeds is None:
            inputs_embeds = self.get_input_embeddings()(input_ids)
        else:
            inputs_embeds = _take_tensor(inputs_embeds)

        image_mask = None
        video_mask = None

        if pixel_values is not None:
            image_embeds, deepstack_image_embeds = self.get_image_features(pixel_values, image_grid_thw)
            pixel_values = None
            image_embed_list = list(image_embeds)
            image_embeds = None
            image_embeds = _cat_tensors(image_embed_list, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
            image_mask, _ = self.get_placeholder_mask(
                input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
            )
            inputs_embeds.masked_scatter_(image_mask, image_embeds)
            image_embeds = None

        if pixel_values_videos is not None:
            video_embeds, deepstack_video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw)
            pixel_values_videos = None
            video_embed_list = list(video_embeds)
            video_embeds = None
            video_embeds = _cat_tensors(video_embed_list, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
            _, video_mask = self.get_placeholder_mask(
                input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds
            )
            inputs_embeds.masked_scatter_(video_mask, video_embeds)
            video_embeds = None

        visual_pos_masks = None
        deepstack_visual_embeds = None
        if image_mask is not None and video_mask is not None:
            # aggregate visual_pos_masks and deepstack_visual_embeds
            image_mask = image_mask[..., 0]
            video_mask = video_mask[..., 0]
            visual_pos_masks = image_mask | video_mask
            deepstack_visual_embeds = []
            image_mask_joint = image_mask[visual_pos_masks]
            video_mask_joint = video_mask[visual_pos_masks]
            for img_embed, vid_embed in zip(deepstack_image_embeds, deepstack_video_embeds):
                embed_joint = img_embed.new_zeros(visual_pos_masks.sum(), img_embed.shape[-1])
                embed_joint[image_mask_joint, :] = img_embed
                embed_joint[video_mask_joint, :] = vid_embed
                deepstack_visual_embeds.append(embed_joint)
            del deepstack_image_embeds, deepstack_video_embeds, image_mask_joint, video_mask_joint
        elif image_mask is not None:
            image_mask = image_mask[..., 0]
            visual_pos_masks = image_mask
            deepstack_visual_embeds = deepstack_image_embeds
        elif video_mask is not None:
            video_mask = video_mask[..., 0]
            visual_pos_masks = video_mask
            deepstack_visual_embeds = deepstack_video_embeds
        image_mask = video_mask = None

        if position_ids is None:
            attention_mask_tensor = (
                attention_mask if not isinstance(attention_mask, dict) else attention_mask["full_attention"]
            )
            if attention_mask_tensor is not None and attention_mask_tensor.ndim == 4:
                attention_mask_tensor = torch.diagonal(attention_mask_tensor[:, 0], dim1=1, dim2=2)
                # Only apply conversion for floating point tensors (inverted masks)
                if attention_mask_tensor.dtype.is_floating_point:
                    attention_mask_tensor = attention_mask_tensor.clone()
                    attention_mask_tensor.div_(torch.finfo(attention_mask_tensor.dtype).min).neg_().add_(1.0)
                    attention_mask_tensor = attention_mask_tensor.int()

            # Calculate RoPE index once per generation in the pre-fill stage only.
            # When compiling, we can't check tensor values thus we check only input length
            # It is safe to assume that `length!=1` means we're in pre-fill because compiled
            # models currently cannot do asssisted decoding
            prefill_compiled_stage = is_torchdynamo_compiling() and (
                (input_ids is not None and input_ids.shape[1] != 1)
                or (inputs_embeds is not None and inputs_embeds.shape[1] != 1)
            )
            prefill_noncompiled_stage = not is_torchdynamo_compiling() and (
                (cache_position is not None and cache_position[0] == 0)
                or (past_key_values is None or past_key_values.get_seq_length() == 0)
            )
            if (prefill_compiled_stage or prefill_noncompiled_stage) or self.rope_deltas is None:
                position_ids, rope_deltas = self.get_rope_index(
                    input_ids,
                    image_grid_thw,
                    video_grid_thw,
                    attention_mask=attention_mask_tensor,
                )
                self.rope_deltas = rope_deltas
            # then use the prev pre-calculated rope-deltas to get the correct position ids
            else:
                batch_size, seq_length, _ = inputs_embeds.shape
                delta = (
                    (cache_position[0] + self.rope_deltas).to(inputs_embeds.device)
                    if cache_position is not None
                    else 0
                )
                position_ids = torch.arange(seq_length, device=inputs_embeds.device)
                position_ids = position_ids.view(1, -1).expand(batch_size, -1)
                if cache_position is not None:  # otherwise `deltas` is an int `0`
                    delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0)
                position_ids = position_ids.add(delta)
                position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
            attention_mask_tensor = None

        input_ids = None
        inputs_embeds_list = [inputs_embeds]
        inputs_embeds = None
        outputs = self.language_model(
            input_ids=None,
            position_ids=position_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds_list,
            cache_position=cache_position,
            visual_pos_masks=visual_pos_masks,
            deepstack_visual_embeds=deepstack_visual_embeds,
            **kwargs,
        )

        last_hidden_state = outputs.last_hidden_state
        output_past_key_values = outputs.past_key_values
        outputs = None
        return Qwen3VLModelOutputWithPast(
            last_hidden_state=last_hidden_state,
            past_key_values=output_past_key_values,
            rope_deltas=self.rope_deltas,
        )


@dataclass
@auto_docstring(
    custom_intro="""
    Base class for Qwen3VL causal language model (or autoregressive) outputs.
    """
)
class Qwen3VLCausalLMOutputWithPast(ModelOutput):
    r"""
    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
        Language modeling loss (for next-token prediction).
    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).

        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
        `past_key_values` input) to speed up sequential decoding.
    rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
        The rope index difference between sequence length and multimodal rope.
    """

    loss: Optional[torch.FloatTensor] = None
    logits: Optional[torch.FloatTensor] = None
    past_key_values: Optional[Cache] = None
    hidden_states: Optional[tuple[torch.FloatTensor]] = None
    attentions: Optional[tuple[torch.FloatTensor]] = None
    rope_deltas: Optional[torch.LongTensor] = None
    x_pred: Optional[torch.FloatTensor] = None
    mid_results: Optional[list] = None

class Qwen3VLForConditionalGeneration(Qwen3VLPreTrainedModel, GenerationMixin):
    _checkpoint_conversion_mapping = {}
    _tied_weights_keys = ["lm_head.weight"]
    # Reference: fix gemma3 grad acc #37208
    accepts_loss_kwargs = False
    config: Qwen3VLConfig

    def __init__(self, config):
        super().__init__(config)
        self.model = Qwen3VLModel(config)
        self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)

        self.post_init()

    def get_input_embeddings(self):
        return self.model.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.model.set_input_embeddings(value)

    def set_decoder(self, decoder):
        self.model.set_decoder(decoder)

    def get_decoder(self):
        return self.model.get_decoder()

    def get_video_features(
        self, pixel_values_videos: torch.FloatTensor, video_grid_thw: Optional[torch.LongTensor] = None
    ):
        return self.model.get_video_features(pixel_values_videos, video_grid_thw)

    def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor] = None):
        return self.model.get_image_features(pixel_values, image_grid_thw)

    # Make modules available through conditional class for BC
    @property
    def language_model(self):
        return self.model.language_model

    @property
    def visual(self):
        return self.model.visual

    def clear_runtime_caches(self):
        self.model.clear_runtime_caches()

    @check_model_inputs
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        pixel_values: Optional[torch.Tensor] = None,
        pixel_values_videos: Optional[torch.FloatTensor] = None,
        image_grid_thw: Optional[torch.LongTensor] = None,
        video_grid_thw: Optional[torch.LongTensor] = None,
        cache_position: Optional[torch.LongTensor] = None,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        vinputs: Optional[torch.Tensor] = None,
        timestep: Optional[torch.Tensor] = None,
        token_types: Optional[torch.Tensor] = None,
        image_embeds: Optional[torch.Tensor] = None,
        video_embeds: Optional[torch.Tensor] = None,
        prediction_mask: Optional[torch.Tensor] = None,
        return_mid_results_layers: Optional[list] = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> Union[tuple, Qwen3VLCausalLMOutputWithPast]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
            The temporal, height and width of feature shape of each image in LLM.
        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
            The temporal, height and width of feature shape of each video in LLM.

        Example:
            TODO: Add example
        """
        is_generation = vinputs is not None
        if not is_generation:
            inputs_embeds_for_model = [inputs_embeds] if inputs_embeds is not None else None
            pixel_values_for_model = [pixel_values] if pixel_values is not None else None
            pixel_values_videos_for_model = [pixel_values_videos] if pixel_values_videos is not None else None
            inputs_embeds = pixel_values = pixel_values_videos = None
            vinputs_for_model = None
        else:
            inputs_embeds_for_model = inputs_embeds
            pixel_values_for_model = [pixel_values] if isinstance(pixel_values, torch.Tensor) else pixel_values
            pixel_values_videos_for_model = (
                [pixel_values_videos] if isinstance(pixel_values_videos, torch.Tensor) else pixel_values_videos
            )
            vinputs_for_model = [vinputs] if isinstance(vinputs, torch.Tensor) else vinputs
            image_embeds = [image_embeds] if isinstance(image_embeds, torch.Tensor) else image_embeds
            video_embeds = [video_embeds] if isinstance(video_embeds, torch.Tensor) else video_embeds
            pixel_values = pixel_values_videos = vinputs = None
        outputs = self.model(
            input_ids=input_ids,
            pixel_values=pixel_values_for_model,
            pixel_values_videos=pixel_values_videos_for_model,
            image_grid_thw=image_grid_thw,
            video_grid_thw=video_grid_thw,
            position_ids=position_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds_for_model,
            cache_position=cache_position,
            vinputs=vinputs_for_model,
            timestep=timestep,
            token_types=token_types,
            image_embeds=image_embeds,
            video_embeds=video_embeds,
            prediction_mask=prediction_mask,
            return_mid_results_layers=return_mid_results_layers,
            **kwargs,
        )

        # Generation path: return x_pred directly
        if is_generation:
            return Qwen3VLCausalLMOutputWithPast(
                x_pred=outputs.x_pred if outputs is not None else None,
                mid_results=outputs.mid_results if hasattr(outputs, 'mid_results') else None,
            )

        hidden_states = outputs[0]
        output_past_key_values = outputs.past_key_values
        output_rope_deltas = outputs.rope_deltas
        outputs = None

        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        selected_hidden_states = hidden_states[:, slice_indices, :]
        logits = self.lm_head(selected_hidden_states)
        del selected_hidden_states, hidden_states

        loss = None
        if labels is not None:
            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size)

        return Qwen3VLCausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=output_past_key_values,
            rope_deltas=output_rope_deltas,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        cache_position=None,
        position_ids=None,
        use_cache=True,
        pixel_values=None,
        pixel_values_videos=None,
        image_grid_thw=None,
        video_grid_thw=None,
        **kwargs,
    ):
        # Overwritten -- in specific circumstances we don't want to forward image inputs to the model

        model_inputs = super().prepare_inputs_for_generation(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            cache_position=cache_position,
            position_ids=position_ids,
            pixel_values=pixel_values,
            pixel_values_videos=pixel_values_videos,
            image_grid_thw=image_grid_thw,
            video_grid_thw=video_grid_thw,
            use_cache=use_cache,
            **kwargs,
        )

        # Qwen3VL position_ids are prepareed with rope_deltas in forward
        model_inputs["position_ids"] = None

        if cache_position[0] != 0:
            model_inputs["pixel_values"] = None
            model_inputs["pixel_values_videos"] = None

        return model_inputs

    def _get_image_nums_and_video_nums(
        self,
        input_ids: Optional[torch.LongTensor],
        inputs_embeds: Optional[torch.Tensor] = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Get the number of images and videos for each sample to calculate the separation length of the sample tensor.
        These parameters are not passed through the processor to avoid unpredictable impacts from interface modifications.

        Args:
            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
                Indices of input sequence tokens in the vocabulary.

        Returns:
            image_nums (`torch.LongTensor` of shape `(batch_size, num_images_sample)`)
            video_nums (`torch.LongTensor` of shape `(batch_size, num_videos_sample)`)
        """
        image_token_id = self.config.image_token_id
        video_token_id = self.config.video_token_id
        vision_start_token_id = self.config.vision_start_token_id

        if inputs_embeds is not None:
            first_channel = inputs_embeds[..., 0]
            vision_start_value = self.get_input_embeddings()(
                torch.tensor(vision_start_token_id, dtype=torch.long, device=inputs_embeds.device)
            )[0]
            image_value = self.get_input_embeddings()(
                torch.tensor(image_token_id, dtype=torch.long, device=inputs_embeds.device)
            )[0]
            video_value = self.get_input_embeddings()(
                torch.tensor(video_token_id, dtype=torch.long, device=inputs_embeds.device)
            )[0]
            vision_start_mask = first_channel == vision_start_value
            image_mask = first_channel == image_value
            video_mask = first_channel == video_value
        else:
            vision_start_mask = input_ids == vision_start_token_id
            image_mask = input_ids == image_token_id
            video_mask = input_ids == video_token_id

        vision_first_mask = torch.roll(vision_start_mask, shifts=1, dims=1)
        image_nums = torch.sum(vision_first_mask & image_mask, dim=1)
        video_nums = torch.sum(vision_first_mask & video_mask, dim=1)

        return image_nums, video_nums

    def _expand_inputs_for_generation(
        self,
        expand_size: int = 1,
        is_encoder_decoder: bool = False,
        input_ids: Optional[torch.LongTensor] = None,
        **model_kwargs,
    ) -> tuple[torch.LongTensor, dict[str, Any]]:
        # Overwritten -- Support for expanding tensors without a batch size dimension
        # e.g., pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw, second_per_grid_t
        # pixel_values.shape[0] is sum(seqlen_images for samples)
        # image_grid_thw.shape[0] is sum(num_images for samples)

        if expand_size == 1:
            return input_ids, model_kwargs

        visual_keys = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw", "second_per_grid_ts"]

        def _expand_dict_for_generation_visual(dict_to_expand):
            image_grid_thw = model_kwargs.get("image_grid_thw", None)
            video_grid_thw = model_kwargs.get("video_grid_thw", None)
            image_nums, video_nums = self._get_image_nums_and_video_nums(
                input_ids, inputs_embeds=model_kwargs.get("inputs_embeds", None)
            )

            def _repeat_interleave_samples(x, lengths, repeat_times):
                samples = torch.split(x, lengths)
                repeat_args = [repeat_times] + [1] * (x.dim() - 1)
                result = torch.cat([sample.repeat(*repeat_args) for sample in samples], dim=0)
                return result

            for key in dict_to_expand:
                if key == "pixel_values":
                    # split images into samples
                    samples = torch.split(image_grid_thw, list(image_nums))
                    # compute the sequence length of images for each sample
                    lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
                    dict_to_expand[key] = _repeat_interleave_samples(
                        dict_to_expand[key], lengths=lengths, repeat_times=expand_size
                    )
                elif key == "image_grid_thw":
                    # get the num of images for each sample
                    lengths = list(image_nums)
                    dict_to_expand[key] = _repeat_interleave_samples(
                        dict_to_expand[key], lengths=lengths, repeat_times=expand_size
                    )
                elif key == "pixel_values_videos":
                    samples = torch.split(video_grid_thw, list(video_nums))
                    lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
                    dict_to_expand[key] = _repeat_interleave_samples(
                        dict_to_expand[key], lengths=lengths, repeat_times=expand_size
                    )
                elif key == "video_grid_thw":
                    lengths = list(video_nums)
                    dict_to_expand[key] = _repeat_interleave_samples(
                        dict_to_expand[key], lengths=lengths, repeat_times=expand_size
                    )
                elif key == "second_per_grid_ts":
                    dict_to_expand[key] = _repeat_interleave_samples(
                        dict_to_expand[key], lengths=list(video_nums), repeat_times=expand_size
                    )
            return dict_to_expand

        def _expand_dict_for_generation(dict_to_expand):
            for key in dict_to_expand:
                if (
                    key != "cache_position"
                    and dict_to_expand[key] is not None
                    and isinstance(dict_to_expand[key], torch.Tensor)
                    and key not in visual_keys
                ):
                    dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0)
            return dict_to_expand

        model_kwargs = _expand_dict_for_generation_visual(model_kwargs)

        if input_ids is not None:
            input_ids = input_ids.repeat_interleave(expand_size, dim=0)

        model_kwargs = _expand_dict_for_generation(model_kwargs)

        if is_encoder_decoder:
            if model_kwargs.get("encoder_outputs") is None:
                raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.")
            model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"])

        return input_ids, model_kwargs