From 51dbdb22026d2d13b1957a9bc0b8660be1561576 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Fri, 26 Dec 2025 22:25:10 +0800 Subject: [PATCH] [diffusion] improve: improve qwen-image-edit performance to align with LightX2V (#15812) Co-authored-by: Mick --- .../configs/models/dits/qwenimage.py | 19 +- .../configs/pipeline_configs/qwen_image.py | 64 ++++--- .../layers/attention/backends/flash_attn.py | 96 ++++++++++ .../runtime/layers/triton_ops.py | 160 ++++++++++++++++ .../runtime/models/dits/qwen_image.py | 178 +++++++++++++----- 5 files changed, 425 insertions(+), 92 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py b/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py index 8a726d8ca..0738b63b0 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py +++ b/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( - 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"), - ] - ) + stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list) 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", } diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py index 042013495..415f0ecac 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +++ b/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_freqs.real.to(dtype=dtype), - img_freqs.imag.to(dtype=dtype), - ) - txt_cos, txt_sin = ( - txt_freqs.real.to(dtype=dtype), - txt_freqs.imag.to(dtype=dtype), - ) + img_cos_half = img_freqs.real.to(dtype=torch.float32).contiguous() + img_sin_half = img_freqs.imag.to(dtype=torch.float32).contiguous() + txt_cos_half = txt_freqs.real.to(dtype=torch.float32).contiguous() + txt_sin_half = txt_freqs.imag.to(dtype=torch.float32).contiguous() - 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_sin = shard_rotary_emb_for_sp(img_sin) + img_cache, txt_cache = freqs_cis + 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, :]) - noisy_img_sin = shard_rotary_emb_for_sp(img_sin[:noisy_img_seq_len, :]) - - # concat back the img_cos for input image (since it is not sp-shared later) - img_cos = torch.cat([noisy_img_cos, img_cos[noisy_img_seq_len:, :]], dim=0).to( - device=device - ) - img_sin = torch.cat([noisy_img_sin, img_sin[noisy_img_seq_len:, :]], dim=0).to( - device=device - ) - + 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_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) = ( - QwenImageEditPlusPipelineConfig.get_freqs_cis( - img_shapes, txt_seq_lens, rotary_emb, device, dtype - ) + 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, :]) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py index adbe410a9..17a9cd576 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/layers/triton_ops.py b/python/sglang/multimodal_gen/runtime/layers/triton_ops.py index 2a8d96af8..761b13763 100644 --- a/python/sglang/multimodal_gen/runtime/layers/triton_ops.py +++ b/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), diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 169055a7d..b0f61c77d 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/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, img_key, img_value = img_qkv.chunk(3, dim=-1) + 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: - txt_qkv, _ = attn.to_added_qkv(encoder_hidden_states) - txt_query, txt_key, txt_value = txt_qkv.chunk(3, dim=-1) + 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 @@ -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 - qkv_dim = num_heads * head_dim * 3 - self.to_qkv = ReplicatedLinear(dim, qkv_dim, bias=True) + # Use separate Q/K/V projections + self.inner_dim = out_dim if out_dim is not None else head_dim * num_heads + 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.to_added_qkv = ReplicatedLinear(added_kv_proj_dim, qkv_dim, bias=True) + self.add_q_proj = ReplicatedLinear( + 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 - img_query = apply_rotary_embedding( - img_query, img_cos, img_sin, interleaved=True - ) - img_key = apply_rotary_embedding( - img_key, img_cos, img_sin, interleaved=True - ) - txt_query = apply_rotary_embedding( - txt_query, txt_cos, txt_sin, interleaved=True - ) - txt_key = apply_rotary_embedding( - txt_key, txt_cos, txt_sin, interleaved=True - ) + if apply_rope_with_cos_sin_cache_inplace is None: + raise RuntimeError("flashinfer is required") + + 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, + ) + return q2.view(1, seqlen, nheads, d), k2.view(1, seqlen, nheads, d) + + positions = pos_1d.repeat(bsz).to(q_4d.device, non_blocking=True) + q2 = q_4d.reshape(bsz * seqlen, nheads * d).contiguous() + k2 = k_4d.reshape(bsz * 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, + ) + 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,32 +694,66 @@ 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] - scale_result = scale[index] - gate_result = gate[index] - else: - shift_result = shift - scale_result = scale - gate_result = gate[:1].unsqueeze(1) + 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] - return fuse_scale_shift_kernel(x, scale_result, shift_result), gate_result + 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( self, 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] - - if self.zero_cond_t: - temb = torch.chunk(temb, 2, dim=0)[0] - txt_mod_params = self.txt_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] # 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