[diffusion] improve: improve qwen-image-edit performance to align with LightX2V (#15812)

Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/multimodal_gen/configs/models/dits/qwenimage.py

@@ -22,27 +22,10 @@ class QwenImageArchConfig(DiTArchConfig):
     axes_dims_rope: Tuple[int, int, int] = (16, 56, 56)
     zero_cond_t: bool = False
 
-    stacked_params_mapping: list[tuple[str, str, str]] = field(
+    stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list)
-        default_factory=lambda: [
-            # (param_name, shard_name, shard_id)
-            (".to_qkv", ".to_q", "q"),
-            (".to_qkv", ".to_k", "k"),
-            (".to_qkv", ".to_v", "v"),
-            (".to_added_qkv", ".add_q_proj", "q"),
-            (".to_added_qkv", ".add_k_proj", "k"),
-            (".to_added_qkv", ".add_v_proj", "v"),
-        ]
-    )
 
     param_names_mapping: dict = field(
         default_factory=lambda: {
-            # QKV fusion mappings
-            r"(.*)\.to_q\.(weight|bias)$": (r"\1.to_qkv.\2", 0, 3),
-            r"(.*)\.to_k\.(weight|bias)$": (r"\1.to_qkv.\2", 1, 3),
-            r"(.*)\.to_v\.(weight|bias)$": (r"\1.to_qkv.\2", 2, 3),
-            r"(.*)\.add_q_proj\.(weight|bias)$": (r"\1.to_added_qkv.\2", 0, 3),
-            r"(.*)\.add_k_proj\.(weight|bias)$": (r"\1.to_added_qkv.\2", 1, 3),
-            r"(.*)\.add_v_proj\.(weight|bias)$": (r"\1.to_added_qkv.\2", 2, 3),
             # LoRA mappings
             r"^(transformer_blocks\.\d+\.attn\..*\.lora_[AB])\.default$": r"\1",
         }

python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py

@@ -156,16 +156,14 @@ class QwenImagePipelineConfig(ImagePipelineConfig):
         # img_shapes: for global entire image
         img_freqs, txt_freqs = rotary_emb(img_shapes, txt_seq_lens, device=device)
 
-        img_cos, img_sin = (
+        img_cos_half = img_freqs.real.to(dtype=torch.float32).contiguous()
-            img_freqs.real.to(dtype=dtype),
+        img_sin_half = img_freqs.imag.to(dtype=torch.float32).contiguous()
-            img_freqs.imag.to(dtype=dtype),
+        txt_cos_half = txt_freqs.real.to(dtype=torch.float32).contiguous()
-        )
+        txt_sin_half = txt_freqs.imag.to(dtype=torch.float32).contiguous()
-        txt_cos, txt_sin = (
-            txt_freqs.real.to(dtype=dtype),
-            txt_freqs.imag.to(dtype=dtype),
-        )
 
-        return (img_cos, img_sin), (txt_cos, txt_sin)
+        img_cos_sin_cache = torch.cat([img_cos_half, img_sin_half], dim=-1)
+        txt_cos_sin_cache = torch.cat([txt_cos_half, txt_sin_half], dim=-1)
+        return img_cos_sin_cache, txt_cos_sin_cache
 
     def _prepare_cond_kwargs(self, batch, prompt_embeds, rotary_emb, device, dtype):
         batch_size = prompt_embeds[0].shape[0]
@@ -184,15 +182,15 @@ class QwenImagePipelineConfig(ImagePipelineConfig):
         ] * batch_size
         txt_seq_lens = [prompt_embeds[0].shape[1]]
 
-        (img_cos, img_sin), (txt_cos, txt_sin) = self.get_freqs_cis(
+        freqs_cis = self.get_freqs_cis(
             img_shapes, txt_seq_lens, rotary_emb, device, dtype
         )
 
-        img_cos = shard_rotary_emb_for_sp(img_cos)
+        img_cache, txt_cache = freqs_cis
-        img_sin = shard_rotary_emb_for_sp(img_sin)
+        img_cache = shard_rotary_emb_for_sp(img_cache)
         return {
             "txt_seq_lens": txt_seq_lens,
-            "freqs_cis": ((img_cos, img_sin), (txt_cos, txt_sin)),
+            "freqs_cis": (img_cache, txt_cache),
         }
 
     def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
@@ -252,7 +250,7 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
             ],
         ] * batch_size
         txt_seq_lens = [prompt_embeds[0].shape[1]]
-        (img_cos, img_sin), (txt_cos, txt_sin) = QwenImagePipelineConfig.get_freqs_cis(
+        freqs_cis = QwenImagePipelineConfig.get_freqs_cis(
             img_shapes, txt_seq_lens, rotary_emb, device, dtype
         )
 
@@ -261,20 +259,14 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
             1 * (height // vae_scale_factor // 2) * (width // vae_scale_factor // 2)
         )
 
-        noisy_img_cos = shard_rotary_emb_for_sp(img_cos[:noisy_img_seq_len, :])
+        img_cache, txt_cache = freqs_cis
-        noisy_img_sin = shard_rotary_emb_for_sp(img_sin[:noisy_img_seq_len, :])
+        noisy_img_cache = shard_rotary_emb_for_sp(img_cache[:noisy_img_seq_len, :])
-
+        img_cache = torch.cat(
-        # concat back the img_cos for input image (since it is not sp-shared later)
+            [noisy_img_cache, img_cache[noisy_img_seq_len:, :]], dim=0
-        img_cos = torch.cat([noisy_img_cos, img_cos[noisy_img_seq_len:, :]], dim=0).to(
+        ).to(device=device)
-            device=device
-        )
-        img_sin = torch.cat([noisy_img_sin, img_sin[noisy_img_seq_len:, :]], dim=0).to(
-            device=device
-        )
-
         return {
             "txt_seq_lens": txt_seq_lens,
-            "freqs_cis": ((img_cos, img_sin), (txt_cos, txt_sin)),
+            "freqs_cis": (img_cache, txt_cache),
         }
 
     def preprocess_condition_image(
@@ -441,16 +433,28 @@ class QwenImageEditPlusPipelineConfig(QwenImageEditPipelineConfig):
         ] * batch_size
         txt_seq_lens = [prompt_embeds[0].shape[1]]
 
-        (img_cos, img_sin), (txt_cos, txt_sin) = (
+        freqs_cis = QwenImageEditPlusPipelineConfig.get_freqs_cis(
-            QwenImageEditPlusPipelineConfig.get_freqs_cis(
             img_shapes, txt_seq_lens, rotary_emb, device, dtype
         )
-        )
 
         # perform sp shard on noisy image tokens
         noisy_img_seq_len = (
             1 * (height // vae_scale_factor // 2) * (width // vae_scale_factor // 2)
         )
+
+        if isinstance(freqs_cis[0], torch.Tensor) and freqs_cis[0].dim() == 2:
+            img_cache, txt_cache = freqs_cis
+            noisy_img_cache = shard_rotary_emb_for_sp(img_cache[:noisy_img_seq_len, :])
+            img_cache = torch.cat(
+                [noisy_img_cache, img_cache[noisy_img_seq_len:, :]], dim=0
+            ).to(device=device)
+            return {
+                "txt_seq_lens": txt_seq_lens,
+                "freqs_cis": (img_cache, txt_cache),
+                "img_shapes": img_shapes,
+            }
+
+        (img_cos, img_sin), (txt_cos, txt_sin) = freqs_cis
         noisy_img_cos = shard_rotary_emb_for_sp(img_cos[:noisy_img_seq_len, :])
         noisy_img_sin = shard_rotary_emb_for_sp(img_sin[:noisy_img_seq_len, :])
 

python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py

@@ -2,6 +2,7 @@
 
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass
+from functools import lru_cache
 from typing import Any
 
 import torch
@@ -18,6 +19,14 @@ try:
 except ImportError as e:
     raise e
 
+
+try:
+    from flash_attn_interface import (
+        flash_attn_varlen_func as flash_attn_varlen_func_upstream,
+    )
+except Exception:
+    flash_attn_varlen_func_upstream = None
+
 from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
     AttentionBackend,
     AttentionImpl,
@@ -31,6 +40,51 @@ logger = init_logger(__name__)
 fa_ver = 3
 
 
+@lru_cache(maxsize=128)
+def _get_cu_seqlens(device_index: int, bsz: int, seqlen: int) -> torch.Tensor:
+    return torch.arange(
+        0,
+        (bsz + 1) * seqlen,
+        step=seqlen,
+        device=torch.device("cuda", device_index),
+        dtype=torch.int32,
+    )
+
+
+@lru_cache(maxsize=256)
+def _should_use_upstream_flash_attention(
+    upstream_available: bool,
+    upstream_heads_ok: bool,
+    q_shape: tuple[int, ...],
+    k_shape: tuple[int, ...],
+    v_shape: tuple[int, ...],
+) -> bool:
+    if not upstream_available or not upstream_heads_ok:
+        return False
+
+    if len(q_shape) != 4 or len(k_shape) != 4 or len(v_shape) != 4:
+        return False
+
+    bsz, seqlen, nheads_q, d = q_shape
+    bsz_k, seqlen_k, nheads_k, d_k = k_shape
+    bsz_v, seqlen_v, nheads_v, d_v = v_shape
+
+    if (
+        bsz != bsz_k
+        or bsz != bsz_v
+        or seqlen != seqlen_k
+        or seqlen != seqlen_v
+        or d != d_k
+        or d != d_v
+    ):
+        return False
+    if nheads_k != nheads_v:
+        return False
+    if nheads_k == 0 or (nheads_q % nheads_k) != 0:
+        return False
+    return True
+
+
 def set_fa_ver(ver: int):
     global fa_ver
     fa_ver = ver
@@ -102,9 +156,19 @@ class FlashAttentionImpl(AttentionImpl):
         prefix: str = "",
         **extra_impl_args,
     ) -> None:
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_size = head_size
         self.causal = causal
         self.softmax_scale = softmax_scale
         self.attention_metadata = FlashAttentionMetadata()
+        if self.num_kv_heads is None:
+            self._upstream_heads_ok = True
+        else:
+            # For gqa, the num_heads must be a multiple of num_kv_heads
+            self._upstream_heads_ok = (
+                self.num_kv_heads > 0 and (self.num_heads % self.num_kv_heads) == 0
+            )
 
     def forward(
         self,
@@ -124,6 +188,38 @@ class FlashAttentionImpl(AttentionImpl):
         else:
             max_seqlen_q = query.shape[1]
             max_seqlen_k = key.shape[1]
+        q_shape = tuple(query.shape)
+        k_shape = tuple(key.shape)
+        v_shape = tuple(value.shape)
+        use_upstream = _should_use_upstream_flash_attention(
+            flash_attn_varlen_func_upstream is not None,
+            self._upstream_heads_ok,
+            q_shape,
+            k_shape,
+            v_shape,
+        )
+
+        if use_upstream:
+            bsz, seqlen, nheads_q, d = q_shape
+            bsz_k, seqlen_k, nheads_k, d_k = k_shape
+            bsz_v, seqlen_v, nheads_v, d_v = v_shape
+            q_ = query.contiguous().reshape(bsz * seqlen, nheads_q, d)
+            k_ = key.contiguous().reshape(bsz * seqlen, nheads_k, d)
+            v_ = value.contiguous().reshape(bsz * seqlen, nheads_v, d)
+            cu = _get_cu_seqlens(q_.device.index, bsz, seqlen)
+            out = flash_attn_varlen_func_upstream(
+                q_,
+                k_,
+                v_,
+                cu,
+                cu,
+                seqlen,
+                seqlen,
+                softmax_scale=self.softmax_scale,
+                causal=self.causal,
+            )
+            return out.reshape(bsz, seqlen, nheads_q, d)
+
         output = flash_attn_func(
             q=query,  # type: ignore[no-untyped-call]
             k=key,

python/sglang/multimodal_gen/runtime/layers/triton_ops.py

@@ -129,6 +129,94 @@ def fuse_scale_shift_kernel_blc_opt(
     tl.store(y_ptr + x_off, y, mask=mask)
 
 
+@triton.jit
+def fuse_scale_shift_gate_select01_kernel_blc_opt(
+    x_ptr,
+    shift0_ptr,
+    scale0_ptr,
+    gate0_ptr,
+    shift1_ptr,
+    scale1_ptr,
+    gate1_ptr,
+    index_ptr,
+    y_ptr,
+    gate_out_ptr,
+    B,
+    L,
+    C,
+    stride_x_b,
+    stride_x_l,
+    stride_x_c,
+    stride_s0_b,
+    stride_s0_c,
+    stride_sc0_b,
+    stride_sc0_c,
+    stride_g0_b,
+    stride_g0_c,
+    stride_s1_b,
+    stride_s1_c,
+    stride_sc1_b,
+    stride_sc1_c,
+    stride_g1_b,
+    stride_g1_c,
+    stride_i_b,
+    stride_i_l,
+    stride_go_b,
+    stride_go_l,
+    stride_go_c,
+    BLOCK_L: tl.constexpr,
+    BLOCK_C: tl.constexpr,
+):
+    pid_l = tl.program_id(0)
+    pid_c = tl.program_id(1)
+    pid_b = tl.program_id(2)
+
+    l_offsets = pid_l * BLOCK_L + tl.arange(0, BLOCK_L)
+    c_offsets = pid_c * BLOCK_C + tl.arange(0, BLOCK_C)
+
+    mask_l = l_offsets < L
+    mask_c = c_offsets < C
+    mask = mask_l[:, None] & mask_c[None, :]
+
+    x_off = (
+        pid_b * stride_x_b
+        + l_offsets[:, None] * stride_x_l
+        + c_offsets[None, :] * stride_x_c
+    )
+    x = tl.load(x_ptr + x_off, mask=mask, other=0)
+
+    idx_off = pid_b * stride_i_b + l_offsets * stride_i_l
+    idx = tl.load(index_ptr + idx_off, mask=mask_l, other=0).to(tl.int1)[:, None]
+
+    s0_off = pid_b * stride_s0_b + c_offsets[None, :] * stride_s0_c
+    sc0_off = pid_b * stride_sc0_b + c_offsets[None, :] * stride_sc0_c
+    g0_off = pid_b * stride_g0_b + c_offsets[None, :] * stride_g0_c
+    s1_off = pid_b * stride_s1_b + c_offsets[None, :] * stride_s1_c
+    sc1_off = pid_b * stride_sc1_b + c_offsets[None, :] * stride_sc1_c
+    g1_off = pid_b * stride_g1_b + c_offsets[None, :] * stride_g1_c
+
+    shift0 = tl.load(shift0_ptr + s0_off, mask=mask_c[None, :], other=0)
+    scale0 = tl.load(scale0_ptr + sc0_off, mask=mask_c[None, :], other=0)
+    gate0 = tl.load(gate0_ptr + g0_off, mask=mask_c[None, :], other=0)
+    shift1 = tl.load(shift1_ptr + s1_off, mask=mask_c[None, :], other=0)
+    scale1 = tl.load(scale1_ptr + sc1_off, mask=mask_c[None, :], other=0)
+    gate1 = tl.load(gate1_ptr + g1_off, mask=mask_c[None, :], other=0)
+
+    shift = tl.where(idx, shift1, shift0)
+    scale = tl.where(idx, scale1, scale0)
+    gate = tl.where(idx, gate1, gate0)
+
+    y = x * (1 + scale) + shift
+    tl.store(y_ptr + x_off, y, mask=mask)
+
+    go_off = (
+        pid_b * stride_go_b
+        + l_offsets[:, None] * stride_go_l
+        + c_offsets[None, :] * stride_go_c
+    )
+    tl.store(gate_out_ptr + go_off, gate, mask=mask)
+
+
 def fuse_scale_shift_kernel(
     x: torch.Tensor,
     scale: torch.Tensor,
@@ -241,6 +329,78 @@ def fuse_scale_shift_kernel(
     return output
 
 
+def fuse_scale_shift_gate_select01_kernel(
+    x: torch.Tensor,
+    scale0: torch.Tensor,
+    shift0: torch.Tensor,
+    gate0: torch.Tensor,
+    scale1: torch.Tensor,
+    shift1: torch.Tensor,
+    gate1: torch.Tensor,
+    index: torch.Tensor,
+    block_l: int = 128,
+    block_c: int = 128,
+):
+    assert x.is_contiguous()
+    B, L, C = x.shape
+    output = torch.empty_like(x)
+    gate_out = torch.empty_like(x)
+
+    if (
+        scale0.dim() != 2
+        or shift0.dim() != 2
+        or gate0.dim() != 2
+        or scale1.dim() != 2
+        or shift1.dim() != 2
+        or gate1.dim() != 2
+    ):
+        raise ValueError("scale0/shift0/gate0/scale1/shift1/gate1 must be 2D [B, C]")
+    if index.dim() != 2:
+        raise ValueError("index must be 2D [B, L]")
+
+    grid = (triton.cdiv(L, block_l), triton.cdiv(C, block_c), B)
+    fuse_scale_shift_gate_select01_kernel_blc_opt[grid](
+        x,
+        shift0,
+        scale0,
+        gate0,
+        shift1,
+        scale1,
+        gate1,
+        index,
+        output,
+        gate_out,
+        B,
+        L,
+        C,
+        x.stride(0),
+        x.stride(1),
+        x.stride(2),
+        shift0.stride(0),
+        shift0.stride(1),
+        scale0.stride(0),
+        scale0.stride(1),
+        gate0.stride(0),
+        gate0.stride(1),
+        shift1.stride(0),
+        shift1.stride(1),
+        scale1.stride(0),
+        scale1.stride(1),
+        gate1.stride(0),
+        gate1.stride(1),
+        index.stride(0),
+        index.stride(1),
+        gate_out.stride(0),
+        gate_out.stride(1),
+        gate_out.stride(2),
+        BLOCK_L=block_l,
+        BLOCK_C=block_c,
+        num_warps=4,
+        num_stages=2,
+    )
+    return output, gate_out
+
+
 @triton.autotune(
     configs=[
         triton.Config({"BLOCK_HS_HALF": 32}, num_warps=2),

python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py

@@ -9,6 +9,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import numpy as np
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from diffusers.models.attention import FeedForward
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
@@ -19,7 +20,7 @@ from sglang.multimodal_gen.runtime.layers.attention import USPAttention
 from sglang.multimodal_gen.runtime.layers.layernorm import LayerNorm, RMSNorm
 from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
 from sglang.multimodal_gen.runtime.layers.triton_ops import (
-    apply_rotary_embedding,
+    fuse_scale_shift_gate_select01_kernel,
     fuse_scale_shift_kernel,
 )
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
@@ -29,16 +30,24 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 logger = init_logger(__name__)  # pylint: disable=invalid-name
 
 
+try:
+    from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
+except Exception:
+    apply_rope_with_cos_sin_cache_inplace = None
+
+
 def _get_qkv_projections(
     attn: "QwenImageCrossAttention", hidden_states, encoder_hidden_states=None
 ):
-    img_qkv, _ = attn.to_qkv(hidden_states)
+    img_query, _ = attn.to_q(hidden_states)
-    img_query, img_key, img_value = img_qkv.chunk(3, dim=-1)
+    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:
-        txt_qkv, _ = attn.to_added_qkv(encoder_hidden_states)
+        txt_query, _ = attn.add_q_proj(encoder_hidden_states)
-        txt_query, txt_key, txt_value = txt_qkv.chunk(3, dim=-1)
+        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
 
@@ -464,20 +473,27 @@ class QwenImageCrossAttention(nn.Module):
         self.parallel_attention = parallel_attention
         self.added_kv_proj_dim = added_kv_proj_dim
 
-        # Use ReplicatedLinear for fused QKV projections
+        # Use separate Q/K/V projections
-        qkv_dim = num_heads * head_dim * 3
+        self.inner_dim = out_dim if out_dim is not None else head_dim * num_heads
-        self.to_qkv = ReplicatedLinear(dim, qkv_dim, bias=True)
+        self.inner_kv_dim = self.inner_dim
+        self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=True)
 
         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()
 
-        self.inner_dim = out_dim if out_dim is not None else head_dim * num_heads
-        self.inner_kv_dim = self.inner_dim
-
         if added_kv_proj_dim is not None:
-            # Use ReplicatedLinear for added (encoder) QKV projections
+            self.add_q_proj = ReplicatedLinear(
-            self.to_added_qkv = ReplicatedLinear(added_kv_proj_dim, qkv_dim, bias=True)
+                added_kv_proj_dim, self.inner_dim, bias=True
+            )
+            self.add_k_proj = ReplicatedLinear(
+                added_kv_proj_dim, self.inner_dim, bias=True
+            )
+            self.add_v_proj = ReplicatedLinear(
+                added_kv_proj_dim, self.inner_dim, bias=True
+            )
 
         if context_pre_only is not None and not context_pre_only:
             self.to_add_out = ReplicatedLinear(self.inner_dim, self.dim, bias=out_bias)
@@ -545,19 +561,52 @@ class QwenImageCrossAttention(nn.Module):
 
         # Apply RoPE
         if image_rotary_emb is not None:
-            (img_cos, img_sin), (txt_cos, txt_sin) = image_rotary_emb
+            if apply_rope_with_cos_sin_cache_inplace is None:
-            img_query = apply_rotary_embedding(
+                raise RuntimeError("flashinfer is required")
-                img_query, img_cos, img_sin, interleaved=True
+
+            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
+
+            def _apply_flashinfer_rope(
+                q_4d: torch.Tensor, k_4d: torch.Tensor, cache: torch.Tensor
+            ):
+                bsz, seqlen, nheads, d = q_4d.shape
+
+                pos_1d = torch.arange(seqlen, device="cpu", dtype=torch.long)
+                if bsz == 1:
+                    positions = pos_1d.to(q_4d.device, non_blocking=True)
+                    q2 = q_4d.squeeze(0).reshape(seqlen, nheads * d).contiguous()
+                    k2 = k_4d.squeeze(0).reshape(seqlen, nheads * d).contiguous()
+                    apply_rope_with_cos_sin_cache_inplace(
+                        positions=positions,
+                        query=q2,
+                        key=k2,
+                        head_size=d,
+                        cos_sin_cache=cache,
+                        is_neox=False,
                     )
-            img_key = apply_rotary_embedding(
+                    return q2.view(1, seqlen, nheads, d), k2.view(1, seqlen, nheads, d)
-                img_key, img_cos, img_sin, interleaved=True
+
-            )
+                positions = pos_1d.repeat(bsz).to(q_4d.device, non_blocking=True)
-            txt_query = apply_rotary_embedding(
+                q2 = q_4d.reshape(bsz * seqlen, nheads * d).contiguous()
-                txt_query, txt_cos, txt_sin, interleaved=True
+                k2 = k_4d.reshape(bsz * seqlen, nheads * d).contiguous()
-            )
+                apply_rope_with_cos_sin_cache_inplace(
-            txt_key = apply_rotary_embedding(
+                    positions=positions,
-                txt_key, txt_cos, txt_sin, interleaved=True
+                    query=q2,
+                    key=k2,
+                    head_size=d,
+                    cos_sin_cache=cache,
+                    is_neox=False,
                 )
+                return q2.view(bsz, seqlen, nheads, d), k2.view(bsz, seqlen, nheads, d)
+
+            img_query, img_key = _apply_flashinfer_rope(img_query, img_key, img_cache)
+            txt_query, txt_key = _apply_flashinfer_rope(txt_query, txt_key, txt_cache)
 
         # Concatenate for joint attention
         # Order: [text, image]
@@ -645,14 +694,50 @@ class QwenImageTransformerBlock(nn.Module):
     def _modulate(self, x, mod_params, index=None):
         shift, scale, gate = mod_params.chunk(3, dim=-1)
         if index is not None:
-            shift_result = shift[index]
+            actual_batch = x.shape[0]
-            scale_result = scale[index]
+            shift0, shift1 = (
-            gate_result = gate[index]
+                shift[:actual_batch],
-        else:
+                shift[actual_batch : 2 * actual_batch],
-            shift_result = shift
+            )
-            scale_result = scale
+            scale0, scale1 = (
-            gate_result = gate[:1].unsqueeze(1)
+                scale[:actual_batch],
+                scale[actual_batch : 2 * actual_batch],
+            )
+            gate0, gate1 = gate[:actual_batch], gate[actual_batch : 2 * actual_batch]
 
+            if x.is_cuda:
+                if not x.is_contiguous():
+                    x = x.contiguous()
+                if not index.is_contiguous():
+                    index = index.contiguous()
+                x, gate_result = fuse_scale_shift_gate_select01_kernel(
+                    x,
+                    scale0=scale0.contiguous(),
+                    shift0=shift0.contiguous(),
+                    gate0=gate0.contiguous(),
+                    scale1=scale1.contiguous(),
+                    shift1=shift1.contiguous(),
+                    gate1=gate1.contiguous(),
+                    index=index,
+                )
+                return x, gate_result
+            else:
+                mask = (index == 0).unsqueeze(-1)
+                shift_result = torch.where(
+                    mask, shift0.unsqueeze(1), shift1.unsqueeze(1)
+                )
+                scale_result = torch.where(
+                    mask, scale0.unsqueeze(1), scale1.unsqueeze(1)
+                )
+                gate_result = torch.where(mask, gate0.unsqueeze(1), gate1.unsqueeze(1))
+                return (
+                    fuse_scale_shift_kernel(x, scale_result, shift_result),
+                    gate_result,
+                )
+        else:
+            shift_result = shift.unsqueeze(1)
+            scale_result = scale.unsqueeze(1)
+            gate_result = gate.unsqueeze(1)
             return fuse_scale_shift_kernel(x, scale_result, shift_result), gate_result
 
     def forward(
@@ -660,17 +745,15 @@ class QwenImageTransformerBlock(nn.Module):
         hidden_states: torch.Tensor,
         encoder_hidden_states: torch.Tensor,
         encoder_hidden_states_mask: torch.Tensor,
-        temb: 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[List[int]] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         # Get modulation parameters for both streams
-        img_mod_params = self.img_mod(temb)  # [B, 6*dim]
+        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.zero_cond_t:
-            temb = torch.chunk(temb, 2, dim=0)[0]
-        txt_mod_params = self.txt_mod(temb)  # [B, 6*dim]
 
         # Split modulation parameters for norm1 and norm2
         img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
@@ -880,13 +963,22 @@ class QwenImageTransformer2DModel(CachableDiT):
 
         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=temb,
+                temb_img_silu=temb_img_silu,
+                temb_txt_silu=temb_txt_silu,
                 image_rotary_emb=image_rotary_emb,
                 joint_attention_kwargs=attention_kwargs,
                 modulate_index=modulate_index,
@@ -902,10 +994,8 @@ class QwenImageTransformer2DModel(CachableDiT):
                     hidden_states
                     + controlnet_block_samples[index_block // interval_control]
                 )
-        if self.zero_cond_t:
-            temb = temb.chunk(2, dim=0)[0]
         # Use only the image part (hidden_states) from the dual-stream blocks
-        hidden_states = self.norm_out(hidden_states, temb)
+        hidden_states = self.norm_out(hidden_states, temb_txt)
 
         output = self.proj_out(hidden_states)
         return output