sgl-project--sglang/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py

# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

# SPDX-License-Identifier: Apache-2.0

import functools
from typing import Any, Dict, List, Optional, Tuple, Union

import diffusers
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
from diffusers.models.modeling_outputs import Transformer2DModelOutput
from diffusers.models.normalization import AdaLayerNormContinuous

from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig
from sglang.multimodal_gen.runtime.distributed import (
    get_local_torch_device,
    get_tp_world_size,
)
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
    get_sp_world_size,
)
from sglang.multimodal_gen.runtime.distributed.sp_shard_utils import (
    build_shard_plan,
    join_seqs,
    shard_like,
    should_shard_text,
    split_seqs,
    tail_attn_meta,
)
from sglang.multimodal_gen.runtime.layers.attention import (
    DynamicVarlenMaskMeta,
    USPAttention,
    build_varlen_mask_meta,
)
from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd
from sglang.multimodal_gen.runtime.layers.fused_scale_shift_gate import (
    FusedLayerNormScaleShiftGateSelect01,
    FusedResidualLayerNormScaleShiftGateSelect01,
)
from sglang.multimodal_gen.runtime.layers.layernorm import (
    LayerNormScaleShift,
    RMSNorm,
    ScaleResidualLayerNormScaleShift,
    apply_qk_norm_with_optional_rope,
)
from sglang.multimodal_gen.runtime.layers.linear import (
    ColumnParallelLinear,
    MergedColumnParallelLinear,
    ReplicatedLinear,
    RowParallelLinear,
)
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
    QuantizationConfig,
)
from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
    NunchakuConfig,
    is_nunchaku_available,
)
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
    apply_flashinfer_rope_qk_inplace,
)
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
    LayerwiseOffloadableModuleMixin,
)
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
    is_in_breakable_cuda_graph,
)

logger = init_logger(__name__)  # pylint: disable=invalid-name


def _attn_mask_meta_local_pad(attn_mask_meta) -> int:
    if attn_mask_meta is None or isinstance(attn_mask_meta, DynamicVarlenMaskMeta):
        return 0
    return attn_mask_meta.get("local_pad", 0)


try:
    from nunchaku.models.attention import NunchakuFeedForward  # type: ignore[import]
except Exception:
    NunchakuFeedForward = None


def _local_seq_len(seq_len: int, sp_world_size: int) -> int:
    """get the local seq len, from seq_len padding to the next multiple of sp_world_size, then shard to local"""
    if sp_world_size <= 1:
        return seq_len
    padded_len = seq_len
    if padded_len % sp_world_size != 0:
        padded_len += sp_world_size - (padded_len % sp_world_size)
    return padded_len // sp_world_size


def _get_qkv_projections(
    attn: "QwenImageCrossAttention", hidden_states, encoder_hidden_states=None
):
    if attn.use_fused_qkv:
        img_qkv, _ = attn.to_qkv(hidden_states)
        img_query, img_key, img_value = [
            x.contiguous() for x in img_qkv.chunk(3, dim=-1)
        ]
    else:
        img_query, _ = attn.to_q(hidden_states)
        img_key, _ = attn.to_k(hidden_states)
        img_value, _ = attn.to_v(hidden_states)

    txt_query = txt_key = txt_value = None
    if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None:
        if attn.use_fused_added_qkv:
            txt_qkv, _ = attn.to_added_qkv(encoder_hidden_states)
            txt_query, txt_key, txt_value = [
                x.contiguous() for x in txt_qkv.chunk(3, dim=-1)
            ]
        else:
            txt_query, _ = attn.add_q_proj(encoder_hidden_states)
            txt_key, _ = attn.add_k_proj(encoder_hidden_states)
            txt_value, _ = attn.add_v_proj(encoder_hidden_states)

    return img_query, img_key, img_value, txt_query, txt_key, txt_value


class QwenTimestepProjEmbeddings(nn.Module):
    def __init__(self, embedding_dim, use_additional_t_cond=False):
        super().__init__()

        self.time_proj = Timesteps(
            num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000
        )
        self.timestep_embedder = TimestepEmbedding(
            in_channels=256, time_embed_dim=embedding_dim
        )
        self.use_additional_t_cond = use_additional_t_cond
        if use_additional_t_cond:
            self.addition_t_embedding = nn.Embedding(2, embedding_dim)

    def forward(self, timestep, hidden_states, addition_t_cond=None):
        timesteps_proj = self.time_proj(timestep)
        timesteps_emb = self.timestep_embedder(
            timesteps_proj.to(dtype=hidden_states.dtype)
        )  # (N, D)

        conditioning = timesteps_emb
        if self.use_additional_t_cond:
            if addition_t_cond is None:
                raise ValueError(
                    "When additional_t_cond is True, addition_t_cond must be provided."
                )
            addition_t_emb = self.addition_t_embedding(addition_t_cond)
            addition_t_emb = addition_t_emb.to(dtype=hidden_states.dtype)
            conditioning = conditioning + addition_t_emb

        return conditioning


class QwenEmbedRope(nn.Module):
    def __init__(self, theta: int, axes_dim: List[int], scale_rope=False):
        super().__init__()
        self.theta = theta
        self.axes_dim = axes_dim
        pos_index = torch.arange(4096)
        neg_index = torch.arange(4096).flip(0) * -1 - 1
        self.pos_freqs = torch.cat(
            [
                self.rope_params(pos_index, self.axes_dim[0], self.theta),
                self.rope_params(pos_index, self.axes_dim[1], self.theta),
                self.rope_params(pos_index, self.axes_dim[2], self.theta),
            ],
            dim=1,
        )
        self.neg_freqs = torch.cat(
            [
                self.rope_params(neg_index, self.axes_dim[0], self.theta),
                self.rope_params(neg_index, self.axes_dim[1], self.theta),
                self.rope_params(neg_index, self.axes_dim[2], self.theta),
            ],
            dim=1,
        )

        # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART
        self.scale_rope = scale_rope

    def rope_params(self, index, dim, theta=10000):
        """
        Args:
            index: [0, 1, 2, 3] 1D Tensor representing the position index of the token
        """
        device = index.device
        assert dim % 2 == 0
        freqs = torch.outer(
            index,
            (
                1.0
                / torch.pow(
                    theta,
                    torch.arange(0, dim, 2, device=device).to(torch.float32).div(dim),
                )
            ).to(device=device),
        )
        freqs = torch.polar(torch.ones_like(freqs), freqs)
        return freqs

    def forward(
        self,
        video_fhw: Union[Tuple[int, int, int], List[Tuple[int, int, int]]],
        txt_seq_lens: List[int],
        device: torch.device,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            video_fhw (`Tuple[int, int, int]` or `List[Tuple[int, int, int]]`):
                A list of 3 integers [frame, height, width] representing the shape of the video.
            txt_seq_lens (`List[int]`):
                A list of integers of length batch_size representing the length of each text prompt.
            device: (`torch.device`):
                The device on which to perform the RoPE computation.
        """
        # When models are initialized under a "meta" device context (e.g. init_empty_weights),
        # tensors created during __init__ become meta tensors. Calling .to(...) on a meta tensor
        # raises "Cannot copy out of meta tensor". Rebuild the frequencies on the target device
        # in that case; otherwise move them if just on a different device.
        if getattr(self.pos_freqs, "device", torch.device("meta")).type == "meta":
            pos_index = torch.arange(4096, device=device)
            neg_index = torch.arange(4096, device=device).flip(0) * -1 - 1
            self.pos_freqs = torch.cat(
                [
                    self.rope_params(pos_index, self.axes_dim[0], self.theta),
                    self.rope_params(pos_index, self.axes_dim[1], self.theta),
                    self.rope_params(pos_index, self.axes_dim[2], self.theta),
                ],
                dim=1,
            ).to(device=device)
            self.neg_freqs = torch.cat(
                [
                    self.rope_params(neg_index, self.axes_dim[0], self.theta),
                    self.rope_params(neg_index, self.axes_dim[1], self.theta),
                    self.rope_params(neg_index, self.axes_dim[2], self.theta),
                ],
                dim=1,
            ).to(device=device)
        elif self.pos_freqs.device != device:
            self.pos_freqs = self.pos_freqs.to(device)
            self.neg_freqs = self.neg_freqs.to(device)

        if isinstance(video_fhw, list):
            video_fhw = video_fhw[0]
        if not isinstance(video_fhw, list):
            video_fhw = [video_fhw]

        vid_freqs = []
        max_vid_index = 0
        for idx, fhw in enumerate(video_fhw):
            frame, height, width = fhw
            # RoPE frequencies are cached via a lru_cache decorator on _compute_video_freqs
            video_freq = self._compute_video_freqs(frame, height, width, idx)
            video_freq = video_freq.to(device)
            vid_freqs.append(video_freq)

            if self.scale_rope:
                max_vid_index = max(height // 2, width // 2, max_vid_index)
            else:
                max_vid_index = max(height, width, max_vid_index)

        max_len = max(txt_seq_lens)
        txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + max_len, ...]
        vid_freqs = torch.cat(vid_freqs, dim=0).to(device=device)
        return vid_freqs, txt_freqs

    @functools.lru_cache(maxsize=128)
    def _compute_video_freqs(
        self, frame: int, height: int, width: int, idx: int = 0
    ) -> torch.Tensor:
        seq_lens = frame * height * width
        freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
        freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1)

        freqs_frame = (
            freqs_pos[0][idx : idx + frame]
            .view(frame, 1, 1, -1)
            .expand(frame, height, width, -1)
        )
        if self.scale_rope:
            freqs_height = torch.cat(
                [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]],
                dim=0,
            )
            freqs_height = freqs_height.view(1, height, 1, -1).expand(
                frame, height, width, -1
            )
            freqs_width = torch.cat(
                [freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]],
                dim=0,
            )
            freqs_width = freqs_width.view(1, 1, width, -1).expand(
                frame, height, width, -1
            )
        else:
            freqs_height = (
                freqs_pos[1][:height]
                .view(1, height, 1, -1)
                .expand(frame, height, width, -1)
            )
            freqs_width = (
                freqs_pos[2][:width]
                .view(1, 1, width, -1)
                .expand(frame, height, width, -1)
            )

        freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(
            seq_lens, -1
        )
        return freqs.clone().contiguous()


class QwenEmbedLayer3DRope(nn.Module):
    def __init__(self, theta: int, axes_dim: List[int], scale_rope=False):
        super().__init__()
        self.theta = theta
        self.axes_dim = axes_dim
        pos_index = torch.arange(4096)
        neg_index = torch.arange(4096).flip(0) * -1 - 1
        self.pos_freqs = torch.cat(
            [
                self.rope_params(pos_index, self.axes_dim[0], self.theta),
                self.rope_params(pos_index, self.axes_dim[1], self.theta),
                self.rope_params(pos_index, self.axes_dim[2], self.theta),
            ],
            dim=1,
        )
        self.neg_freqs = torch.cat(
            [
                self.rope_params(neg_index, self.axes_dim[0], self.theta),
                self.rope_params(neg_index, self.axes_dim[1], self.theta),
                self.rope_params(neg_index, self.axes_dim[2], self.theta),
            ],
            dim=1,
        )

        self.scale_rope = scale_rope

    def rope_params(self, index, dim, theta=10000):
        """
        Args:
            index: [0, 1, 2, 3] 1D Tensor representing the position index of the token
        """
        device = index.device
        assert dim % 2 == 0
        freqs = torch.outer(
            index,
            (
                1.0
                / torch.pow(
                    theta,
                    torch.arange(0, dim, 2, device=device).to(torch.float32).div(dim),
                )
            ).to(device=device),
        )
        freqs = torch.polar(torch.ones_like(freqs), freqs)
        return freqs

    def forward(self, video_fhw, txt_seq_lens, device):
        """
        Args: video_fhw: [frame, height, width] a list of 3 integers representing the shape of the video Args:
        txt_length: [bs] a list of 1 integers representing the length of the text
        """

        # When models are initialized under a "meta" device context (e.g. init_empty_weights),
        # tensors created during __init__ become meta tensors. Calling .to(...) on a meta tensor
        # raises "Cannot copy out of meta tensor". Rebuild the frequencies on the target device
        # in that case; otherwise move them if just on a different device.
        if getattr(self.pos_freqs, "device", torch.device("meta")).type == "meta":
            pos_index = torch.arange(4096, device=device)
            neg_index = torch.arange(4096, device=device).flip(0) * -1 - 1
            self.pos_freqs = torch.cat(
                [
                    self.rope_params(pos_index, self.axes_dim[0], self.theta),
                    self.rope_params(pos_index, self.axes_dim[1], self.theta),
                    self.rope_params(pos_index, self.axes_dim[2], self.theta),
                ],
                dim=1,
            ).to(device=device)
            self.neg_freqs = torch.cat(
                [
                    self.rope_params(neg_index, self.axes_dim[0], self.theta),
                    self.rope_params(neg_index, self.axes_dim[1], self.theta),
                    self.rope_params(neg_index, self.axes_dim[2], self.theta),
                ],
                dim=1,
            ).to(device=device)
        elif self.pos_freqs.device != device:
            self.pos_freqs = self.pos_freqs.to(device)
            self.neg_freqs = self.neg_freqs.to(device)

        if isinstance(video_fhw, list):
            video_fhw = video_fhw[0]
        if not isinstance(video_fhw, list):
            video_fhw = [video_fhw]

        vid_freqs = []
        max_vid_index = 0
        layer_num = len(video_fhw) - 1
        for idx, fhw in enumerate(video_fhw):
            frame, height, width = fhw
            if idx != layer_num:
                video_freq = self._compute_video_freqs(frame, height, width, idx)
            else:
                # For the condition image, we set the layer index to -1
                video_freq = self._compute_condition_freqs(frame, height, width)
            video_freq = video_freq.to(device)
            vid_freqs.append(video_freq)

            if self.scale_rope:
                max_vid_index = max(height // 2, width // 2, max_vid_index)
            else:
                max_vid_index = max(height, width, max_vid_index)

        max_vid_index = max(max_vid_index, layer_num)
        max_len = max(txt_seq_lens)
        txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + max_len, ...]
        vid_freqs = torch.cat(vid_freqs, dim=0)

        return vid_freqs, txt_freqs

    @functools.lru_cache(maxsize=None)
    def _compute_video_freqs(self, frame, height, width, idx=0):
        seq_lens = frame * height * width
        freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
        freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1)

        freqs_frame = (
            freqs_pos[0][idx : idx + frame]
            .view(frame, 1, 1, -1)
            .expand(frame, height, width, -1)
        )
        if self.scale_rope:
            freqs_height = torch.cat(
                [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]],
                dim=0,
            )
            freqs_height = freqs_height.view(1, height, 1, -1).expand(
                frame, height, width, -1
            )
            freqs_width = torch.cat(
                [freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]],
                dim=0,
            )
            freqs_width = freqs_width.view(1, 1, width, -1).expand(
                frame, height, width, -1
            )
        else:
            freqs_height = (
                freqs_pos[1][:height]
                .view(1, height, 1, -1)
                .expand(frame, height, width, -1)
            )
            freqs_width = (
                freqs_pos[2][:width]
                .view(1, 1, width, -1)
                .expand(frame, height, width, -1)
            )

        freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(
            seq_lens, -1
        )
        return freqs.clone().contiguous()

    @functools.lru_cache(maxsize=None)
    def _compute_condition_freqs(self, frame, height, width):
        seq_lens = frame * height * width
        freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
        freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1)

        freqs_frame = (
            freqs_neg[0][-1:].view(frame, 1, 1, -1).expand(frame, height, width, -1)
        )
        if self.scale_rope:
            freqs_height = torch.cat(
                [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]],
                dim=0,
            )
            freqs_height = freqs_height.view(1, height, 1, -1).expand(
                frame, height, width, -1
            )
            freqs_width = torch.cat(
                [freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]],
                dim=0,
            )
            freqs_width = freqs_width.view(1, 1, width, -1).expand(
                frame, height, width, -1
            )
        else:
            freqs_height = (
                freqs_pos[1][:height]
                .view(1, height, 1, -1)
                .expand(frame, height, width, -1)
            )
            freqs_width = (
                freqs_pos[2][:width]
                .view(1, 1, width, -1)
                .expand(frame, height, width, -1)
            )

        freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(
            seq_lens, -1
        )
        return freqs.clone().contiguous()


class QwenImageCrossAttention(nn.Module):
    def __init__(
        self,
        dim: int,  # query_dim
        num_heads: int,
        head_dim: int,
        window_size=(-1, -1),
        added_kv_proj_dim: int = None,
        out_bias: bool = True,
        qk_norm=True,  # rmsnorm
        eps=1e-6,
        pre_only=False,
        context_pre_only: bool = False,
        parallel_attention=False,
        out_dim: int = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        assert dim % num_heads == 0
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.window_size = window_size
        self.qk_norm = qk_norm
        self.eps = eps
        self.parallel_attention = parallel_attention
        self.added_kv_proj_dim = added_kv_proj_dim
        self.prefix = prefix

        self.use_fused_qkv = isinstance(quant_config, NunchakuConfig)

        self.inner_dim = out_dim if out_dim is not None else head_dim * num_heads
        self.inner_kv_dim = self.inner_dim

        tp_size = get_tp_world_size()
        assert (
            self.num_heads % tp_size == 0
        ), f"num_heads ({self.num_heads}) must be divisible by tp_size ({tp_size})"
        self.local_num_heads = self.num_heads // tp_size

        if self.use_fused_qkv:
            # Use fused QKV projection for nunchaku quantization
            self.to_qkv = MergedColumnParallelLinear(
                dim,
                [self.inner_dim] * 3,
                bias=True,
                quant_config=quant_config,
                prefix=f"{prefix}.to_qkv",
            )
        else:
            self.to_q = ColumnParallelLinear(
                dim,
                self.inner_dim,
                bias=True,
                gather_output=False,
                quant_config=quant_config,
                prefix=f"{prefix}.to_q",
            )
            self.to_k = ColumnParallelLinear(
                dim,
                self.inner_dim,
                bias=True,
                gather_output=False,
                quant_config=quant_config,
                prefix=f"{prefix}.to_k",
            )
            self.to_v = ColumnParallelLinear(
                dim,
                self.inner_dim,
                bias=True,
                gather_output=False,
                quant_config=quant_config,
                prefix=f"{prefix}.to_v",
            )

        if self.qk_norm:
            self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
            self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()

        if added_kv_proj_dim is not None:
            self.use_fused_added_qkv = isinstance(quant_config, NunchakuConfig)
            if self.use_fused_added_qkv:
                self.to_added_qkv = MergedColumnParallelLinear(
                    added_kv_proj_dim,
                    [self.inner_dim] * 3,
                    bias=True,
                    quant_config=quant_config,
                    prefix=f"{prefix}.to_added_qkv",
                )
            else:
                self.add_q_proj = ColumnParallelLinear(
                    added_kv_proj_dim,
                    self.inner_dim,
                    bias=True,
                    gather_output=False,
                    quant_config=quant_config,
                    prefix=f"{prefix}.add_q_proj",
                )
                self.add_k_proj = ColumnParallelLinear(
                    added_kv_proj_dim,
                    self.inner_dim,
                    bias=True,
                    gather_output=False,
                    quant_config=quant_config,
                    prefix=f"{prefix}.add_k_proj",
                )
                self.add_v_proj = ColumnParallelLinear(
                    added_kv_proj_dim,
                    self.inner_dim,
                    bias=True,
                    gather_output=False,
                    quant_config=quant_config,
                    prefix=f"{prefix}.add_v_proj",
                )

        if context_pre_only is not None and not context_pre_only:
            self.to_add_out = RowParallelLinear(
                self.inner_dim,
                self.dim,
                bias=out_bias,
                input_is_parallel=True,
                quant_config=quant_config,
                prefix=f"{prefix}.to_add_out",
            )
        else:
            self.to_add_out = None

        if not pre_only:
            self.to_out = nn.ModuleList([])
            self.to_out.append(
                RowParallelLinear(
                    self.inner_dim,
                    self.dim,
                    bias=out_bias,
                    input_is_parallel=True,
                    quant_config=quant_config,
                    prefix=f"{prefix}.to_out.0",
                )
            )
        else:
            self.to_out = None

        self.norm_added_q = RMSNorm(head_dim, eps=eps)
        self.norm_added_k = RMSNorm(head_dim, eps=eps)

        # Scaled dot product attention
        self.attn = USPAttention(
            num_heads=self.local_num_heads,
            head_size=self.head_dim,
            dropout_rate=0,
            softmax_scale=None,
            causal=False,
            supported_attention_backends={
                AttentionBackendEnum.FA,
                AttentionBackendEnum.AITER,
                AttentionBackendEnum.AITER_SAGE,
                AttentionBackendEnum.TORCH_SDPA,
                AttentionBackendEnum.SAGE_ATTN,
                AttentionBackendEnum.SAGE_ATTN_3,
            },
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
        **cross_attention_kwargs,
    ):
        """Run joint text-image attention.

        `attn_mask` or `attention_mask` takes precedence. Otherwise,
        `encoder_hidden_states_mask` keeps valid text tokens in the joint
        text-image sequence.
        """
        seq_len_txt = encoder_hidden_states.shape[1]
        attn_mask = cross_attention_kwargs.get("attn_mask")
        if attn_mask is None:
            attn_mask = cross_attention_kwargs.get("attention_mask")
        encoder_hidden_states_mask = cross_attention_kwargs.get(
            "encoder_hidden_states_mask"
        )
        # Varlen metadata precomputed in QwenImageTransformer2DModel.forward,
        # paired with the same ``attn_mask`` for the USPAttention FA fast path.
        attn_mask_meta = cross_attention_kwargs.get("attn_mask_meta")
        # When the text stream is sharded across SP ranks the joint sequence is
        # fully sequence-parallel, so no leading tokens are replicated.
        sp_text_sharded = cross_attention_kwargs.get("sp_text_sharded", False)
        # Rows of tail padding inside THIS rank's text chunk (sp_shard meta).
        sp_txt_pad = _attn_mask_meta_local_pad(attn_mask_meta)

        (
            img_query,
            img_key,
            img_value,
            txt_query,
            txt_key,
            txt_value,
        ) = _get_qkv_projections(self, hidden_states, encoder_hidden_states)

        # Reshape for multi-head attention
        img_query = img_query.unflatten(-1, (self.local_num_heads, self.head_dim))
        img_key = img_key.unflatten(-1, (self.local_num_heads, self.head_dim))
        img_value = img_value.unflatten(-1, (self.local_num_heads, self.head_dim))

        txt_query = txt_query.unflatten(-1, (self.local_num_heads, self.head_dim))
        txt_key = txt_key.unflatten(-1, (self.local_num_heads, self.head_dim))
        txt_value = txt_value.unflatten(-1, (self.local_num_heads, self.head_dim))

        img_cache = txt_cache = None
        if image_rotary_emb is not None:
            if not (
                isinstance(image_rotary_emb[0], torch.Tensor)
                and image_rotary_emb[0].dim() == 2
            ):
                raise RuntimeError("image_rotary_emb must be cos_sin_cache tensors")

            img_cache, txt_cache = image_rotary_emb

        if self.qk_norm:
            img_query, img_key = apply_qk_norm_with_optional_rope(
                q=img_query,
                k=img_key,
                q_norm=self.norm_q,
                k_norm=self.norm_k,
                head_dim=img_query.shape[-1],
                cos_sin_cache=img_cache,
                is_neox=False,
                allow_inplace=True,
            )
            txt_query, txt_key = apply_qk_norm_with_optional_rope(
                q=txt_query,
                k=txt_key,
                q_norm=self.norm_added_q,
                k_norm=self.norm_added_k,
                head_dim=txt_query.shape[-1],
                cos_sin_cache=txt_cache,
                is_neox=False,
                allow_inplace=True,
            )
        elif img_cache is not None and txt_cache is not None:
            img_query, img_key = apply_flashinfer_rope_qk_inplace(
                img_query, img_key, img_cache, is_neox=False
            )
            txt_query, txt_key = apply_flashinfer_rope_qk_inplace(
                txt_query, txt_key, txt_cache, is_neox=False
            )

        # Joint order [text, image]; join_seqs relocates any SP text tail-pad
        # behind the image (see sp_shard.join_seqs for why).
        joint_query = join_seqs(txt_query, img_query, sp_txt_pad)
        joint_key = join_seqs(txt_key, img_key, sp_txt_pad)
        joint_value = join_seqs(txt_value, img_value, sp_txt_pad)
        if attn_mask is None and encoder_hidden_states_mask is not None:
            image_mask = torch.ones(
                (hidden_states.shape[0], img_query.shape[1]),
                device=encoder_hidden_states_mask.device,
                dtype=torch.bool,
            )
            attn_mask = torch.cat(
                [encoder_hidden_states_mask.to(dtype=torch.bool), image_mask],
                dim=1,
            )

        # Compute joint attention
        joint_hidden_states = self.attn(
            joint_query,
            joint_key,
            joint_value,
            attn_mask=attn_mask,
            attn_mask_meta=attn_mask_meta,
            num_replicated_prefix=0 if sp_text_sharded else seq_len_txt,
        )

        # Reshape back
        joint_hidden_states = joint_hidden_states.flatten(2, 3)
        joint_hidden_states = joint_hidden_states.to(joint_query.dtype)

        # Split attention outputs back
        txt_attn_output, img_attn_output = split_seqs(
            joint_hidden_states, seq_len_txt, sp_txt_pad
        )

        # Apply output projections
        img_attn_output, _ = self.to_out[0](img_attn_output)
        if len(self.to_out) > 1:
            (img_attn_output,) = self.to_out[1](img_attn_output)  # dropout

        txt_attn_output, _ = self.to_add_out(txt_attn_output)

        return img_attn_output, txt_attn_output


class QwenImageGELU(nn.Module):
    def __init__(
        self,
        dim: int,
        inner_dim: int,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        replicated: bool = False,
    ) -> None:
        super().__init__()
        if replicated:
            self.proj = ReplicatedLinear(
                dim,
                inner_dim,
                bias=True,
                quant_config=quant_config,
                prefix=f"{prefix}.proj",
            )
        else:
            self.proj = ColumnParallelLinear(
                dim,
                inner_dim,
                bias=True,
                gather_output=False,
                quant_config=quant_config,
                prefix=f"{prefix}.proj",
            )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states, _ = self.proj(hidden_states)
        return F.gelu(hidden_states, approximate="tanh")


class QwenImageFeedForward(nn.Module):
    def __init__(
        self,
        dim: int,
        dim_out: int,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        mult: int = 4,
        replicated: bool = False,
    ) -> None:
        super().__init__()
        inner_dim = dim * mult
        if replicated:
            # Keep the whole FFN resident on every rank: no per-block
            # all-reduce. Only worth it when the branch's token count is small
            # enough that the duplicated GEMM is cheaper than the all-reduce.
            down = ReplicatedLinear(
                inner_dim,
                dim_out,
                bias=True,
                quant_config=quant_config,
                prefix=f"{prefix}.net.2",
            )
        else:
            down = RowParallelLinear(
                inner_dim,
                dim_out,
                bias=True,
                input_is_parallel=True,
                quant_config=quant_config,
                prefix=f"{prefix}.net.2",
            )
        self.net = nn.ModuleList(
            [
                QwenImageGELU(
                    dim,
                    inner_dim,
                    quant_config=quant_config,
                    prefix=f"{prefix}.net.0",
                    replicated=replicated,
                ),
                nn.Dropout(0.0),
                down,
            ]
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.net[0](hidden_states)
        hidden_states = self.net[1](hidden_states)
        hidden_states, _ = self.net[2](hidden_states)
        return hidden_states


class QwenImageTransformerBlock(nn.Module):
    def __init__(
        self,
        dim: int,
        num_attention_heads: int,
        attention_head_dim: int,
        qk_norm: str = "rms_norm",
        eps: float = 1e-6,
        quant_config: Optional[QuantizationConfig] | NunchakuConfig = None,
        prefix: str = "",
        zero_cond_t: bool = False,
    ):
        super().__init__()
        self.prefix = prefix

        self.dim = dim
        self.num_attention_heads = num_attention_heads
        self.attention_head_dim = attention_head_dim
        self.quant_config = quant_config
        self.zero_cond_t = zero_cond_t

        # Image processing modules
        self.img_mod = nn.Sequential(
            nn.SiLU(),
            ReplicatedLinear(
                dim,
                6 * dim,
                bias=True,
                quant_config=quant_config,
                prefix=f"{prefix}.img_mod",
            ),  # For scale, shift, gate for norm1 and norm2
        )
        self.img_norm1 = LayerNormScaleShift(
            hidden_size=dim, eps=eps, elementwise_affine=False
        )

        self.attn = QwenImageCrossAttention(
            dim=dim,
            num_heads=num_attention_heads,
            added_kv_proj_dim=dim,
            context_pre_only=False,
            head_dim=attention_head_dim,
            quant_config=quant_config,
            prefix=f"{prefix}.attn",
        )
        self.img_norm2 = ScaleResidualLayerNormScaleShift(
            dim, eps=eps, elementwise_affine=False
        )

        # Text processing modules
        self.txt_mod = nn.Sequential(
            nn.SiLU(),
            ReplicatedLinear(
                dim,
                6 * dim,
                bias=True,
                quant_config=quant_config,
                prefix=f"{prefix}.txt_mod",
            ),  # For scale, shift, gate for norm1 and norm2
        )
        self.txt_norm1 = LayerNormScaleShift(
            hidden_size=dim, eps=eps, elementwise_affine=False
        )
        # Text doesn't need separate attention - it's handled by img_attn joint computation
        self.txt_norm2 = ScaleResidualLayerNormScaleShift(
            hidden_size=dim, eps=eps, elementwise_affine=False
        )
        # Utils
        self.fuse_mul_add = MulAdd()
        self.fused_ln_ss_gate_select01 = FusedLayerNormScaleShiftGateSelect01()
        self.fused_res_ln_ss_gate_select01 = (
            FusedResidualLayerNormScaleShiftGateSelect01()
        )

        nunchaku_enabled = (
            quant_config is not None
            and hasattr(quant_config, "get_name")
            and quant_config.get_name() == "svdquant"
            and is_nunchaku_available()
        )
        if nunchaku_enabled:
            ff_class = diffusers.models.attention.FeedForward
            self.img_mlp = ff_class(
                dim=dim,
                dim_out=dim,
                activation_fn="gelu-approximate",
            )
            self.txt_mlp = ff_class(
                dim=dim,
                dim_out=dim,
                activation_fn="gelu-approximate",
            )
        else:
            self.img_mlp = QwenImageFeedForward(
                dim=dim,
                dim_out=dim,
                quant_config=quant_config,
                prefix=f"{prefix}.img_mlp",
            )
            self.txt_mlp = QwenImageFeedForward(
                dim=dim,
                dim_out=dim,
                quant_config=quant_config,
                prefix=f"{prefix}.txt_mlp",
                # The text branch is ~1K tokens regardless of image size, so
                # sharding its FFN saves less GEMM time than the per-block
                # all-reduce it adds. Measured e2e crossover (H100, 1024x1024):
                # tp=2 replication wins (5.58s -> 5.12s, -8%) but at tp=4 the
                # duplicated GEMM outgrows the all-reduce saved (~1% loss), so
                # gate on the TP degree.
                replicated=get_tp_world_size() <= 2,
            )

        if nunchaku_enabled:
            nunchaku_kwargs = {
                "precision": quant_config.precision,
                "rank": quant_config.rank,
                "act_unsigned": quant_config.act_unsigned,
            }
            self.img_mlp = NunchakuFeedForward(self.img_mlp, **nunchaku_kwargs)
            self.txt_mlp = NunchakuFeedForward(self.txt_mlp, **nunchaku_kwargs)

    def _norm_scale_shift(
        self,
        norm_module: LayerNormScaleShift,
        x: torch.Tensor,
        shift: torch.Tensor,
        scale: torch.Tensor,
    ) -> torch.Tensor:
        return norm_module(x=x, shift=shift, scale=scale)

    def _scale_residual_norm_scale_shift(
        self,
        norm_module: ScaleResidualLayerNormScaleShift,
        *,
        residual: torch.Tensor,
        x: torch.Tensor,
        gate: torch.Tensor | int,
        shift: torch.Tensor,
        scale: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        return norm_module(
            residual=residual,
            x=x,
            gate=gate,
            shift=shift,
            scale=scale,
        )

    def _mul_add(
        self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0
    ) -> torch.Tensor:
        return self.fuse_mul_add(a, b, c, k)

    def _modulate(
        self,
        x: torch.Tensor,
        mod_params: torch.Tensor,
        norm_module: Union[LayerNormScaleShift, ScaleResidualLayerNormScaleShift],
        index: Optional[torch.Tensor] = None,
        gate_x: Optional[torch.Tensor] = None,
        residual_x: Optional[torch.Tensor] = None,
        use_bcg_helpers: bool = False,
    ) -> Union[
        Tuple[torch.Tensor, torch.Tensor],
        Tuple[torch.Tensor, torch.Tensor, torch.Tensor],
    ]:
        # Apply attention gates and add residual (like in Megatron)
        #   - residual_out = gate_x * x + residual_x
        # - x = norm(residual_out) * (1 + scale) + shift
        is_scale_residual = isinstance(norm_module, ScaleResidualLayerNormScaleShift)

        shift, scale, gate = mod_params.chunk(3, dim=-1)
        if index is not None:
            actual_batch = x.shape[0]
            shift0, shift1 = (
                shift[:actual_batch],
                shift[actual_batch : 2 * actual_batch],
            )
            scale0, scale1 = (
                scale[:actual_batch],
                scale[actual_batch : 2 * actual_batch],
            )
            gate0, gate1 = (
                gate[:actual_batch],
                gate[actual_batch : 2 * actual_batch],
            )
            if is_scale_residual:
                x, residual_out, gate_result = self.fused_res_ln_ss_gate_select01(
                    x,
                    residual_x,
                    gate_x,
                    getattr(norm_module.norm, "weight", None),
                    getattr(norm_module.norm, "bias", None),
                    scale0,
                    shift0,
                    gate0,
                    scale1,
                    shift1,
                    gate1,
                    index,
                    norm_module.eps,
                )
                return x, residual_out, gate_result
            else:
                x, gate_result = self.fused_ln_ss_gate_select01(
                    x,
                    getattr(norm_module.norm, "weight", None),
                    getattr(norm_module.norm, "bias", None),
                    scale0,
                    shift0,
                    gate0,
                    scale1,
                    shift1,
                    gate1,
                    index,
                    norm_module.eps,
                )
                return x, gate_result
        else:
            shift_result = shift.unsqueeze(1)
            scale_result = scale.unsqueeze(1)
            gate_result = gate.unsqueeze(1)
        if is_scale_residual:
            if use_bcg_helpers:
                modulated, residual_out = self._scale_residual_norm_scale_shift(
                    norm_module,
                    residual=residual_x,
                    x=x,
                    gate=gate_x,
                    shift=shift_result,
                    scale=scale_result,
                )
            else:
                modulated, residual_out = norm_module(
                    residual=residual_x,
                    x=x,
                    gate=gate_x,
                    shift=shift_result,
                    scale=scale_result,
                )
            return modulated, residual_out, gate_result
        else:
            if use_bcg_helpers:
                modulated = self._norm_scale_shift(
                    norm_module, x=x, shift=shift_result, scale=scale_result
                )
            else:
                modulated = norm_module(x=x, shift=shift_result, scale=scale_result)
            return modulated, gate_result

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        encoder_hidden_states_mask: torch.Tensor,
        temb_img_silu: torch.Tensor,
        temb_txt_silu: torch.Tensor,
        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
        modulate_index: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Get modulation parameters for both streams
        img_mod_params, _ = self.img_mod[1](temb_img_silu)  # [B, 6*dim]
        txt_mod_params, _ = self.txt_mod[1](temb_txt_silu)  # [B, 6*dim]

        if (
            self.quant_config is not None
            and hasattr(self.quant_config, "get_name")
            and self.quant_config.get_name() == "svdquant"
        ):
            # When NOT using nunchaku, reshape mod_params from [B, 6*dim] to [B, dim*6]
            # When using nunchaku (svdquant), keep original format
            img_mod_params = (
                img_mod_params.view(img_mod_params.shape[0], -1, 6)
                .transpose(1, 2)
                .reshape(img_mod_params.shape[0], -1)
            )
            txt_mod_params = (
                txt_mod_params.view(txt_mod_params.shape[0], -1, 6)
                .transpose(1, 2)
                .reshape(txt_mod_params.shape[0], -1)
            )

        # Split modulation parameters for norm1 and norm2
        img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
        txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
        use_bcg_helpers = is_in_breakable_cuda_graph()

        # Process image stream - norm1 + modulation
        img_modulated, img_gate1 = self._modulate(
            hidden_states,
            img_mod1,
            self.img_norm1,
            modulate_index,
            use_bcg_helpers=use_bcg_helpers,
        )
        # Process text stream - norm1 + modulation
        txt_shift1, txt_scale1, txt_gate1_raw = txt_mod1.chunk(3, dim=-1)
        if use_bcg_helpers:
            txt_modulated = self._norm_scale_shift(
                self.txt_norm1,
                encoder_hidden_states,
                shift=txt_shift1,
                scale=txt_scale1,
            )
        else:
            txt_modulated = self.txt_norm1(
                encoder_hidden_states, shift=txt_shift1, scale=txt_scale1
            )
        txt_gate1 = txt_gate1_raw.unsqueeze(1)

        # Use QwenAttnProcessor2_0 for joint attention computation
        # This directly implements the DoubleStreamLayerMegatron logic:
        # 1. Computes QKV for both streams
        # 2. Applies QK normalization and RoPE
        # 3. Concatenates and runs joint attention
        # 4. Splits results back to separate streams
        joint_attention_kwargs = joint_attention_kwargs or {}
        attn_output = self.attn(
            # Image stream (will be processed as "sample")
            hidden_states=img_modulated,
            # Text stream (will be processed as "context")
            encoder_hidden_states=txt_modulated,
            encoder_hidden_states_mask=encoder_hidden_states_mask,
            image_rotary_emb=image_rotary_emb,
            **joint_attention_kwargs,
        )

        # QwenAttnProcessor2_0 returns (img_output, txt_output) when encoder_hidden_states is provided
        img_attn_output, txt_attn_output = attn_output
        # Process image stream - norm2 + MLP
        img_modulated2, hidden_states, img_gate2 = self._modulate(
            img_attn_output,
            img_mod2,
            self.img_norm2,
            modulate_index,
            gate_x=img_gate1,
            residual_x=hidden_states,
            use_bcg_helpers=use_bcg_helpers,
        )
        img_mlp_output = self.img_mlp(img_modulated2)

        if img_mlp_output.dim() == 2:
            img_mlp_output = img_mlp_output.unsqueeze(0)
        if use_bcg_helpers:
            hidden_states = self._mul_add(img_mlp_output, img_gate2, hidden_states)
        else:
            hidden_states = self.fuse_mul_add(img_mlp_output, img_gate2, hidden_states)

        # Process text stream - norm2 + MLP
        txt_shift2, txt_scale2, txt_gate2_raw = txt_mod2.chunk(3, dim=-1)
        if use_bcg_helpers:
            (
                txt_modulated2,
                encoder_hidden_states,
            ) = self._scale_residual_norm_scale_shift(
                self.txt_norm2,
                residual=encoder_hidden_states,
                x=txt_attn_output,
                gate=txt_gate1,
                shift=txt_shift2,
                scale=txt_scale2,
            )
        else:
            txt_modulated2, encoder_hidden_states = self.txt_norm2(
                residual=encoder_hidden_states,
                x=txt_attn_output,
                gate=txt_gate1,
                shift=txt_shift2,
                scale=txt_scale2,
            )
        txt_gate2 = txt_gate2_raw.unsqueeze(1)
        txt_mlp_output = self.txt_mlp(txt_modulated2)

        if txt_mlp_output.dim() == 2:
            txt_mlp_output = txt_mlp_output.unsqueeze(0)
        if use_bcg_helpers:
            encoder_hidden_states = self._mul_add(
                txt_mlp_output, txt_gate2, encoder_hidden_states
            )
        else:
            encoder_hidden_states = self.fuse_mul_add(
                txt_mlp_output, txt_gate2, encoder_hidden_states
            )

        # Clip to prevent overflow for fp16
        if encoder_hidden_states.dtype == torch.float16:
            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
        if hidden_states.dtype == torch.float16:
            hidden_states = hidden_states.clip(-65504, 65504)

        return encoder_hidden_states, hidden_states


def to_hashable(obj):
    if isinstance(obj, list):
        return tuple(to_hashable(x) for x in obj)
    return obj


class QwenImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
    """
    The Transformer model introduced in Qwen.

    """

    _supports_gradient_checkpointing = True
    _no_split_modules = ["QwenImageTransformerBlock"]
    _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
    _repeated_blocks = ["QwenImageTransformerBlock"]

    param_names_mapping = QwenImageDitConfig().arch_config.param_names_mapping
    _fsdp_shard_conditions = QwenImageDitConfig().arch_config._fsdp_shard_conditions

    @classmethod
    def get_nunchaku_quant_rules(cls) -> dict[str, list[str]]:
        return {
            "skip": [
                "norm",
                "embed",
                "rotary",
                "pos_embed",
            ],
            "svdq_w4a4": [
                "attn.to_qkv",
                "attn.to_out",
                "attn.add_qkv_proj",
                "attn.to_add_out",
                "img_mlp",
                "txt_mlp",
            ],
            "awq_w4a16": [
                "img_mod",
                "txt_mod",
            ],
        }

    def __init__(
        self,
        config: QwenImageDitConfig,
        hf_config: dict[str, Any],
        quant_config: Optional[QuantizationConfig] = None,
    ):
        super().__init__(config=config, hf_config=hf_config)
        patch_size = config.arch_config.patch_size
        in_channels = config.arch_config.in_channels
        out_channels = config.arch_config.out_channels
        num_layers = config.arch_config.num_layers
        attention_head_dim = config.arch_config.attention_head_dim
        num_attention_heads = config.arch_config.num_attention_heads
        joint_attention_dim = config.arch_config.joint_attention_dim
        axes_dims_rope = config.arch_config.axes_dims_rope
        self.zero_cond_t = getattr(config.arch_config, "zero_cond_t", False)
        self.out_channels = out_channels or in_channels
        self.inner_dim = num_attention_heads * attention_head_dim

        self.use_additional_t_cond: bool = getattr(
            config.arch_config, "use_additional_t_cond", False
        )  # For qwen-image-layered now
        self.use_layer3d_rope: bool = getattr(
            config.arch_config, "use_layer3d_rope", False
        )  # For qwen-image-layered now

        if not self.use_layer3d_rope:
            self.rotary_emb = QwenEmbedRope(
                theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True
            )
        else:
            self.rotary_emb = QwenEmbedLayer3DRope(
                theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True
            )

        self.time_text_embed = QwenTimestepProjEmbeddings(
            embedding_dim=self.inner_dim,
            use_additional_t_cond=self.use_additional_t_cond,
        )

        self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)

        self.img_in = ColumnParallelLinear(
            in_channels,
            self.inner_dim,
            bias=True,
            gather_output=True,
            quant_config=quant_config,
            prefix="img_in",
        )
        self.txt_in = ColumnParallelLinear(
            joint_attention_dim,
            self.inner_dim,
            bias=True,
            gather_output=True,
            quant_config=quant_config,
            prefix="txt_in",
        )

        self.transformer_blocks = nn.ModuleList(
            [
                QwenImageTransformerBlock(
                    dim=self.inner_dim,
                    num_attention_heads=num_attention_heads,
                    attention_head_dim=attention_head_dim,
                    quant_config=quant_config,
                    prefix=f"transformer_blocks.{layer_idx}",
                    zero_cond_t=self.zero_cond_t,
                )
                for layer_idx in range(num_layers)
            ]
        )

        self.norm_out = AdaLayerNormContinuous(
            self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6
        )
        self.proj_out = ColumnParallelLinear(
            self.inner_dim,
            patch_size * patch_size * self.out_channels,
            bias=True,
            gather_output=True,
            quant_config=quant_config,
            prefix="proj_out",
        )

        self.timestep_zero = torch.zeros(
            (1,), dtype=torch.int, device=get_local_torch_device()
        )

        self.layer_names = ["transformer_blocks"]

    @functools.lru_cache(maxsize=50)
    def build_modulate_index(self, img_shapes: tuple[int, int, int], device):
        sp_world_size = get_sp_world_size()

        modulate_index_list = []
        for sample in img_shapes:
            first_size = sample[0][0] * sample[0][1] * sample[0][2]
            total_size = sum(s[0] * s[1] * s[2] for s in sample)
            if sp_world_size > 1:
                first_local_size = _local_seq_len(first_size, sp_world_size)
                tail_local_size = _local_seq_len(total_size - first_size, sp_world_size)
                idx = torch.cat(
                    [
                        torch.zeros(first_local_size, device=device, dtype=torch.int),
                        torch.ones(tail_local_size, device=device, dtype=torch.int),
                    ]
                )
            else:
                idx = (torch.arange(total_size, device=device) >= first_size).int()
            modulate_index_list.append(idx)

        modulate_index = torch.stack(modulate_index_list)
        return modulate_index

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor = None,
        encoder_hidden_states_mask: torch.Tensor = None,
        timestep: torch.LongTensor = None,
        img_shapes: Optional[List[Tuple[int, int, int]]] = None,
        txt_seq_lens: Optional[List[int]] = None,
        freqs_cis: tuple[torch.Tensor, torch.Tensor] = None,
        additional_t_cond: Optional[torch.Tensor] = None,
        guidance: torch.Tensor = None,
        attention_kwargs: Optional[Dict[str, Any]] = None,
        controlnet_block_samples=None,
        return_dict: bool = True,
    ) -> Union[torch.Tensor, Transformer2DModelOutput]:
        """
        The [`QwenTransformer2DModel`] forward method.

        Args:
            hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
                Input `hidden_states`.
            encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
                Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
            encoder_hidden_states_mask (`torch.Tensor` of shape `(batch_size, text_sequence_length)`):
                Valid-token mask of the input conditions, where True keeps a text token.
            timestep ( `torch.LongTensor`):
                Used to indicate denoising step.
            attention_kwargs (`dict`, *optional*):
                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
                `self.processor` in
                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
                tuple.

        Returns:
            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
            `tuple` where the first element is the sample tensor.
        """
        if (
            attention_kwargs is not None
            and attention_kwargs.get("scale", None) is not None
        ):
            logger.warning(
                "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
            )

        if isinstance(encoder_hidden_states, list):
            encoder_hidden_states = encoder_hidden_states[0]
        if isinstance(encoder_hidden_states_mask, list):
            encoder_hidden_states_mask = encoder_hidden_states_mask[0]

        hidden_states, _ = self.img_in(hidden_states)

        timestep = (timestep / 1000).to(hidden_states.dtype)

        if self.zero_cond_t:
            timestep = torch.cat([timestep, self.timestep_zero], dim=0)
            device = timestep.device
            modulate_index = self.build_modulate_index(to_hashable(img_shapes), device)
        else:
            modulate_index = None

        encoder_hidden_states = self.txt_norm(encoder_hidden_states)
        encoder_hidden_states, _ = self.txt_in(encoder_hidden_states)

        block_attention_kwargs = attention_kwargs.copy() if attention_kwargs else {}
        sp_text_sharded = False
        if encoder_hidden_states_mask is not None:
            encoder_hidden_states_mask = encoder_hidden_states_mask.to(
                device=hidden_states.device, dtype=torch.bool
            )
            batch_size, image_seq_len = hidden_states.shape[:2]
            image_mask = torch.ones(
                (batch_size, image_seq_len),
                dtype=torch.bool,
                device=hidden_states.device,
            )
            joint_mask = torch.cat([encoder_hidden_states_mask, image_mask], dim=1)
            block_attention_kwargs["attn_mask"] = joint_mask
            if is_in_breakable_cuda_graph():
                # Qwen/FireRed BCG buckets text inputs so different prompt
                # lengths can share a graph. Attention break kwargs are captured
                # once, so build varlen metadata replay-locally from the current
                # static mask instead of closing over stale cu_seqlens/indices.
                block_attention_kwargs["attn_mask_meta"] = DynamicVarlenMaskMeta()
            else:
                # Precompute varlen metadata once per request so every block
                # reuses the same cu_seqlens / indices instead of rebuilding.
                block_attention_kwargs["attn_mask_meta"] = build_varlen_mask_meta(
                    joint_mask
                )
        elif should_shard_text(encoder_hidden_states.shape[1]):
            # Shard the replicated text stream across SP ranks; non-divisible
            # lengths tail-pad the last rank and attention skips the pad via the
            # per-request tail meta. Otherwise fall through to replicated text.
            txt_shard = build_shard_plan(encoder_hidden_states.shape[1])
            encoder_hidden_states = shard_like(encoder_hidden_states, txt_shard)
            if freqs_cis is not None:
                img_cache, txt_cache = freqs_cis
                freqs_cis = (img_cache, shard_like(txt_cache, txt_shard, dim=0))
            tail_meta = tail_attn_meta(
                txt_shard,
                encoder_hidden_states.shape[0],
                hidden_states.device,
                image_seq_len=hidden_states.shape[1],
            )
            if tail_meta is not None:
                block_attention_kwargs["attn_mask_meta"] = tail_meta
            sp_text_sharded = True
        block_attention_kwargs["sp_text_sharded"] = sp_text_sharded

        temb = self.time_text_embed(timestep, hidden_states, additional_t_cond)

        temb_img_silu = F.silu(temb)
        if self.zero_cond_t:
            temb_txt = temb.chunk(2, dim=0)[0]
            temb_txt_silu = temb_img_silu.chunk(2, dim=0)[0]
        else:
            temb_txt = temb
            temb_txt_silu = temb_img_silu

        image_rotary_emb = freqs_cis
        for index_block, block in enumerate(self.transformer_blocks):
            encoder_hidden_states, hidden_states = block(
                hidden_states=hidden_states,
                encoder_hidden_states=encoder_hidden_states,
                encoder_hidden_states_mask=encoder_hidden_states_mask,
                temb_img_silu=temb_img_silu,
                temb_txt_silu=temb_txt_silu,
                image_rotary_emb=image_rotary_emb,
                joint_attention_kwargs=block_attention_kwargs,
                modulate_index=modulate_index,
            )

            # controlnet residual
            if controlnet_block_samples is not None:
                interval_control = len(self.transformer_blocks) / len(
                    controlnet_block_samples
                )
                interval_control = int(np.ceil(interval_control))
                hidden_states = (
                    hidden_states
                    + controlnet_block_samples[index_block // interval_control]
                )
        # Use only the image part (hidden_states) from the dual-stream blocks
        hidden_states = self.norm_out(hidden_states, temb_txt)

        output, _ = self.proj_out(hidden_states)
        return output


EntryClass = QwenImageTransformer2DModel