native/dit_decoder.py

#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           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.

from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, Optional, Union, List

import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers.activations import ACT2FN
from transformers.modeling_layers import GradientCheckpointingLayer


from .time_embed import TimestepEmbedder as TimestepEmbedder
from .packed_mul import mul_fixed as mul
from .packed_add import add_fixed as add

def modulate(seq_data, seq_lens, batch_shift, batch_scale):
    scaled_seq_data = mul(seq_data, seq_lens, batch_scale+1)
    shifted_scaled_seq_data = add(scaled_seq_data, seq_lens, batch_shift)
    return shifted_scaled_seq_data


from .portable_attention import flex_attention, create_block_mask, or_masks, and_masks, BlockMask


def full_attention_mask(seq_lens:torch.Tensor, valid_size:int, max_size:int):
    padded_bids = torch.full((max_size,), -1, device=seq_lens.device)
    valid_bids = torch.arange(seq_lens.shape[0], device=seq_lens.device)
    valid_bids = torch.repeat_interleave(valid_bids, seq_lens, output_size=valid_size)
    padded_bids[:valid_size] = valid_bids

    def full_mask(b, h, q_idx, k_idx):
        return (padded_bids[q_idx] == padded_bids[k_idx])
    def valid_mask(b, h, q_idx, k_idx):
        return (padded_bids[q_idx] != -1) & (padded_bids[k_idx] != -1) # (q_idx < valid_size) & (k_idx < valid_size)
    mask = and_masks(full_mask, valid_mask)
    block_mask = create_block_mask(mask, B=1, H=None, Q_LEN=max_size, KV_LEN=max_size)
    return block_mask

def full_incontext_attention_mask(x_bids: torch.Tensor, seq_lens_x:torch.Tensor, valid_size_x:int, max_size_x:int, seq_lens_y:torch.Tensor, valid_size_y:int, max_size_y:int):
    padded_bids_x = torch.full((max_size_x,), -1, device=seq_lens_x.device)
    valid_bids_x = x_bids
    valid_bids_x = torch.repeat_interleave(valid_bids_x, seq_lens_x, output_size=valid_size_x)
    padded_bids_x[:valid_size_x] = valid_bids_x
    padded_bids_y = torch.full((max_size_y,), -1, device=seq_lens_y.device)
    valid_bids_y = x_bids # torch.arange(seq_lens_y.shape[0], device=seq_lens_y.device)
    valid_bids_y = torch.repeat_interleave(valid_bids_y, seq_lens_y, output_size=valid_size_y)
    padded_bids_y[:valid_size_y] = valid_bids_y
    padded_bids_xy = torch.cat([padded_bids_x, padded_bids_y], dim=0)
    def full_mask(b, h, q_idx, k_idx):
        return padded_bids_x[q_idx] == padded_bids_xy[k_idx]
    def valid_mask(b, h, q_idx, k_idx):
        return (padded_bids_x[q_idx] != -1) & (padded_bids_xy[k_idx] != -1) #(q_idx < valid_size_x) & (k_idx < valid_size_x + max_size_x)
    mask = and_masks(full_mask, valid_mask)
    block_mask = create_block_mask(mask, B=1, H=None, Q_LEN=max_size_x, KV_LEN=max_size_x + max_size_y)
    return block_mask

def _compute_default_rope_parameters(
    rope_theta: float,
    head_dim,
    partial_rotary_factor=1.0,
    device: Optional["torch.device"] = None,
    seq_len: Optional[int] = None,
) -> tuple["torch.Tensor", float]:

    base = rope_theta
    dim = int(head_dim * partial_rotary_factor)
    attention_factor = 1.0  # Unused in this type of RoPE
    # Compute the inverse frequencies
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim))
    return inv_freq, attention_factor

class RMSNorm(nn.Module):
    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return (self.weight * hidden_states).to(input_dtype)

class MLP(nn.Module):
    def __init__(self, hidden_size: int, intermediate_size: int, hidden_act: str):
        super().__init__()
        self.hidden_size = hidden_size
        self.intermediate_size = 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[hidden_act]

    def forward(self, x):
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj


class Qwen3VLVisionPatchEmbed(nn.Module):
    def __init__(self, hidden_size: int, patch_size: int, in_channels: int) -> None:
        super().__init__()
        self.patch_size = patch_size
        self.in_channels = in_channels
        self.embed_dim = hidden_size
        self.proj = nn.Conv2d(self.in_channels, self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size, bias=True)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        target_dtype = self.proj.weight.dtype
        N = hidden_states.shape[0]
        hidden_states = hidden_states.view(
            N, self.in_channels, self.patch_size, self.patch_size
        )
        hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
        return hidden_states



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

    def __init__(self, rope_theta: float, head_dim: int, rope_scaling: Optional[dict] = None, device=None):
        super().__init__()
        if rope_scaling is not None and rope_scaling.get("rope_type", "default") is not None:
            self.rope_type = rope_scaling.get("rope_type", "default")
        else:
            self.rope_type = "default"

        inv_freq, self.attention_scaling = _compute_default_rope_parameters(rope_theta, head_dim)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.original_inv_freq = self.inv_freq

        self.mrope_section = rope_scaling.get("mrope_section", [24, 20, 20])

    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]  # just overwrite the first dimension T
        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()
    def forward(self, x, position_ids):
        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(3, -1, 1)
        position_ids_expanded = position_ids[:, None, :].float()  # shape (3, 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(1, 2)
            freqs = self.apply_interleaved_mrope(freqs, self.mrope_section)
            emb = torch.cat((freqs, freqs), dim=-1)
            cos = emb.cos() * self.attention_scaling
            sin = emb.sin() * self.attention_scaling

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





def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


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.
    """
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


class InContextAttention(nn.Module):
    def __init__(self, hidden_size: int, num_heads: int) -> None:
        super().__init__()
        self.dim = hidden_size
        self.num_heads = 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.q_norm = RMSNorm(self.head_dim)
        self.k_norm = RMSNorm(self.head_dim)
        self.conditon_kv = nn.Linear(
            self.dim,
            self.dim * 2,
            bias=True,
        )
        self.proj = nn.Linear(self.dim, self.dim)
        self.scaling = self.head_dim**-0.5
        self.attention_dropout = 0.0
        self.is_causal = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        condition_hidden_states: torch.Tensor,
        attention_mask: BlockMask,
        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        condition_position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        **kwargs,
    ) -> torch.Tensor:
        seq_length = hidden_states.shape[0]
        cond_seq_length = condition_hidden_states.shape[0]
        query_states, key_states, value_states = (
            self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
        )
        condition_key_states, condition_value_states = (
            self.conditon_kv(condition_hidden_states).reshape(cond_seq_length, 2, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
        )
        cos, sin = position_embeddings
        query_states = self.q_norm(query_states)
        key_states = self.k_norm(key_states)
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        condition_cos, condition_sin = condition_position_embeddings
        condition_key_states = self.k_norm(condition_key_states)
        condition_key_states, _ = apply_rotary_pos_emb(condition_key_states, condition_value_states, condition_cos, condition_sin)

        key_states = torch.cat((key_states, condition_key_states), dim=0)
        value_states = torch.cat((value_states, condition_value_states), dim=0)

        query_states = query_states.transpose(0, 1).unsqueeze(0)
        key_states = key_states.transpose(0, 1).unsqueeze(0)
        value_states = value_states.transpose(0, 1).unsqueeze(0)



        attn_output = flex_attention(query_states, key_states, value_states, block_mask=attention_mask)
        attn_output = attn_output.transpose(1, 2)

        attn_output = attn_output.reshape(seq_length, -1).contiguous()
        attn_output = self.proj(attn_output)
        return attn_output

class ConditionRefinerAttention(nn.Module):
    def __init__(self, hidden_size: int, num_heads: int) -> None:
        super().__init__()
        self.dim = hidden_size
        self.num_heads = 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.q_norm = RMSNorm(self.head_dim)
        self.k_norm = RMSNorm(self.head_dim)
        self.proj = nn.Linear(self.dim, self.dim)
        self.scaling = self.head_dim**-0.5
        self.attention_dropout = 0.0
        self.is_causal = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: BlockMask,
        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        **kwargs,
    ) -> torch.Tensor:
        seq_length = hidden_states.shape[0]
        query_states, key_states, value_states = (
            self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
        )
        cos, sin = position_embeddings
        query_states = self.q_norm(query_states)
        key_states = self.k_norm(query_states)
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        query_states = query_states.transpose(0, 1).unsqueeze(0)
        key_states = key_states.transpose(0, 1).unsqueeze(0)
        value_states = value_states.transpose(0, 1).unsqueeze(0)


        attn_output = flex_attention(query_states, key_states, value_states, block_mask=attention_mask)
        attn_output = attn_output.transpose(1, 2)

        attn_output = attn_output.reshape(seq_length, -1).contiguous()
        attn_output = self.proj(attn_output)
        return attn_output

class ConditionRefinerBlock(GradientCheckpointingLayer):
    def __init__(self, hidden_size: int, num_heads: int, intermediate_size: int, hidden_act: str) -> None:
        super().__init__()
        self.norm1 = RMSNorm(hidden_size, eps=1e-6)
        self.norm2 = RMSNorm(hidden_size, eps=1e-6)
        self.attn = ConditionRefinerAttention(hidden_size=hidden_size, num_heads=num_heads)
        self.mlp = MLP(hidden_size=hidden_size, intermediate_size=intermediate_size, hidden_act=hidden_act)
    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: BlockMask,
        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        **kwargs,
    ) -> torch.Tensor:
        hidden_states = hidden_states + self.attn(
            self.norm1(hidden_states),
            attention_mask=attention_mask,
            position_embeddings=position_embeddings,
            **kwargs,
        )
        hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
        return hidden_states


class Qwen3VLVisionBlock(GradientCheckpointingLayer):
    def __init__(self, hidden_size: int, num_heads: int, intermediate_size: int, hidden_act: str) -> None:
        super().__init__()
        self.norm1 = RMSNorm(hidden_size, eps=1e-6)
        self.norm2 = RMSNorm(hidden_size, eps=1e-6)
        self.attn = InContextAttention(hidden_size=hidden_size, num_heads=num_heads)
        self.mlp = MLP(hidden_size=hidden_size, intermediate_size=intermediate_size, hidden_act=hidden_act)

        self.adaLN_modulation = nn.Sequential(
            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        condition_hidden_states: torch.Tensor,
        t: torch.Tensor,
        seq_lens: torch.Tensor,
        attention_mask: BlockMask,
        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        condition_position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        **kwargs,
    ) -> torch.Tensor:
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(t).chunk(6, dim=-1)

        # attention
        attn_in = modulate(self.norm1(hidden_states), seq_lens, shift_msa, scale_msa)
        attn_out = self.attn(
            attn_in,
            condition_hidden_states=condition_hidden_states,
            attention_mask=attention_mask,
            position_embeddings=position_embeddings,
            condition_position_embeddings=condition_position_embeddings,
            **kwargs,
        )
        gated_attn_out = mul(attn_out, seq_lens, gate_msa)
        hidden_states = hidden_states + gated_attn_out

        # mlp
        mlp_in = modulate(self.norm2(hidden_states), seq_lens, shift_mlp, scale_mlp)
        mlp_out = self.mlp(mlp_in)
        gated_mlp_out = mul(mlp_out, seq_lens, gate_mlp)
        hidden_states = hidden_states + gated_mlp_out
        return hidden_states



class FlattenDiT(nn.Module):
    def __init__(self,
                 spatial_merge_size: int,
                 condition_refiner_depth: int,
                 condition_hidden_size: int,
                 depth: int,
                 hidden_size: int,
                 num_heads: int,
                 intermediate_size: int,
                 hidden_act: str,
                 patch_size: int,
                 in_channels: int,
                 rope_theta: float,
                 rope_scaling: Optional[dict] = None,
                 ckpt_path: Optional[str] = None,
                 ) -> None:
        super().__init__()
        self.spatial_merge_size = spatial_merge_size
        deltas = torch.meshgrid(
            torch.zeros((1,), dtype=torch.float32),
            torch.linspace(0, 1, self.spatial_merge_size+1)[:-1],
            torch.linspace(0, 1, self.spatial_merge_size+1)[:-1],
            indexing='ij'
        )
        deltas = torch.stack(deltas, dim=0).view(3, 1, -1)
        self.register_buffer('deltas', deltas)

        self.patch_size = patch_size
        self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size

        self.patch_embed = Qwen3VLVisionPatchEmbed(
            hidden_size=hidden_size,
            patch_size=patch_size,
            in_channels=in_channels,
        )

        head_dim = hidden_size // num_heads
        self.rotary_pos_emb = Qwen3VLTextRotaryEmbedding(
            rope_theta = rope_theta,
            head_dim = head_dim,
            rope_scaling=rope_scaling,
        )

        self.condition_proj = nn.Sequential(
            nn.Linear(condition_hidden_size, hidden_size),
            # nn.LayerNorm(hidden_size, eps=1e-6)
        )
        self.condition_refiner = nn.ModuleList([
            ConditionRefinerBlock(
                hidden_size=hidden_size,
                num_heads=num_heads,
                intermediate_size=intermediate_size,
                hidden_act=hidden_act)
            for _ in range(condition_refiner_depth)])


        self.blocks = nn.ModuleList([
            Qwen3VLVisionBlock(
                hidden_size=hidden_size,
                num_heads=num_heads,
                intermediate_size=intermediate_size,
                hidden_act=hidden_act)
            for _ in range(depth)])

        self.out_proj = nn.Linear(
            hidden_size, patch_size ** 2 * in_channels,
        )
        # init out_proj with zeros
        nn.init.zeros_(self.out_proj.weight)
        nn.init.zeros_(self.out_proj.bias)

        self.gradient_checkpointing = False

        self.t_embedder = TimestepEmbedder(hidden_size)

        if ckpt_path is not None:
            if ckpt_path.endswith('.safetensors'):
                # safetensors
                from safetensors.torch import load_file
                state_dict = load_file(ckpt_path)
            else:
                state_dict = torch.load(ckpt_path, map_location='cpu')

            for k, v in self.state_dict().items():
                if k in state_dict:
                    try:
                        v.copy_(state_dict[k])
                    except Exception as e:
                        print(f"Failed to copy {k}: {e}")
                else:
                    print(f"Parameter {k} not found in checkpoint")
            # self.load_state_dict(state_dict, strict=False)

    def forward(self,
                hidden_states: torch.Tensor,
                t: torch.Tensor,
                position_ids: torch.Tensor,
                seq_lens: torch.Tensor,
                condition_hidden_states: torch.Tensor,
                condition_position_ids: torch.Tensor,
                condition_seq_lens: torch.Tensor,
                visual_pos_idxs: torch.Tensor,
                batch_id_hidden_states: torch.Tensor
                ) -> torch.Tensor:
        """
        Args:
            hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`):
                The final hidden states of the model.
            position_ids (seq_len, 3) [tid, hid, wid]
            seq_lens (batch_size)

        Returns:
            `torch.Tensor`: hidden_states.
        """
        position_ids = position_ids.transpose(0, 1)

        # refiner
        condition_max_seq_lens = hidden_states.shape[0] // self.spatial_merge_unit
        condition_seq_lens = seq_lens // self.spatial_merge_unit
        condition_valid_seq_lens = condition_seq_lens.sum().item()
        condition_hidden_states_wotext = condition_hidden_states[visual_pos_idxs]
        condition_hidden_states_wotext = torch.nn.functional.pad(
            condition_hidden_states_wotext, (0, 0, 0, condition_max_seq_lens - condition_hidden_states_wotext.shape[0]))
        condition_position_ids = position_ids.view(3, condition_max_seq_lens, self.spatial_merge_unit)
        condition_position_ids = torch.mean(condition_position_ids, dim=2, keepdim=False)
        # base_position_ids = condition_position_ids[:, visual_pos_idxs]
        # condition_position_ids = torch.ones((3, condition_max_seq_lens), dtype=torch.float32, device=hidden_states.device)
        # condition_position_ids[:, :condition_valid_seq_lens] = base_position_ids

        condition_hidden_states_wotext = self.condition_proj(condition_hidden_states_wotext).to(torch.bfloat16)
        condition_hidden_states = condition_hidden_states_wotext

        condition_attn_mask = full_attention_mask(condition_seq_lens, condition_valid_seq_lens, condition_max_seq_lens)
        condition_position_embeddings = self.rotary_pos_emb(condition_hidden_states, condition_position_ids)

        for blk in self.condition_refiner:
            condition_hidden_states = blk(
                condition_hidden_states,
                condition_attn_mask,
                condition_position_embeddings,
            )
        t = self.t_embedder(t)
        t = nn.functional.silu(t)

        valid_seq_lens = seq_lens.sum().item()
        max_seq_len, _ = hidden_states.size()

        hidden_states = self.patch_embed(hidden_states)
        position_embeddings = self.rotary_pos_emb(hidden_states, position_ids)
        attn_mask = full_incontext_attention_mask(batch_id_hidden_states,
                                                  seq_lens,
                                                  valid_seq_lens,
                                                  max_seq_len,
                                                  condition_seq_lens,
                                                  condition_valid_seq_lens,
                                                  condition_max_seq_lens)
        hidden_states = hidden_states.reshape(max_seq_len, -1)

        for layer_num, blk in enumerate(self.blocks):
            hidden_states = blk(
                hidden_states,
                condition_hidden_states,
                t,
                seq_lens,
                attn_mask,
                position_embeddings=position_embeddings,
                condition_position_embeddings=condition_position_embeddings
            )

        hidden_states = self.out_proj(hidden_states)
        return hidden_states