From f3d73b0199d5a43b00977b2b47d9729b27343d04 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Mon, 29 Dec 2025 17:24:26 +0800 Subject: [PATCH] [Diffusion] Refactor qwen_image's rope in a single helper func (#16047) --- .../configs/pipeline_configs/qwen_image.py | 1 + .../runtime/layers/rotary_embedding.py | 68 ++++++++++++++++++- .../runtime/models/dits/flux.py | 14 ++-- .../runtime/models/dits/flux_2.py | 26 ++++--- .../runtime/models/dits/qwen_image.py | 53 +++------------ 5 files changed, 103 insertions(+), 59 deletions(-) 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 415f0ecac..f3fe94188 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py @@ -156,6 +156,7 @@ class QwenImagePipelineConfig(ImagePipelineConfig): # img_shapes: for global entire image img_freqs, txt_freqs = rotary_emb(img_shapes, txt_seq_lens, device=device) + # flashinfer RoPE expects a float32 cos/sin cache concatenated on the last dim 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() diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py index c0a589038..ac5e8ed0e 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py +++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py @@ -27,7 +27,7 @@ """Rotary Positional Embeddings.""" import functools from collections import OrderedDict -from typing import Any +from typing import Any, Optional, Tuple import torch @@ -39,6 +39,72 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) +def apply_flashinfer_rope_qk_inplace( + q: torch.Tensor, + k: torch.Tensor, + cos_sin_cache: torch.Tensor, + *, + head_size: Optional[int] = None, + is_neox: bool = False, + positions: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + if q.dim() != 4 or k.dim() != 4: + raise ValueError( + f"Expected q/k to be 4D [bsz, seqlen, nheads, head_size], " + f"got q:{tuple(q.shape)} k:{tuple(k.shape)}" + ) + if q.shape != k.shape: + raise ValueError( + f"q and k must have the same shape, got {q.shape} vs {k.shape}" + ) + + if not (isinstance(cos_sin_cache, torch.Tensor) and cos_sin_cache.dim() == 2): + raise ValueError("cos_sin_cache must be a 2D torch.Tensor") + + bsz, seqlen, nheads, d = q.shape + if head_size is None: + head_size = d + if head_size != d: + raise ValueError(f"head_size mismatch: inferred {d}, but head_size={head_size}") + + try: + from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace + except Exception as e: + raise RuntimeError( + "flashinfer is required for apply_flashinfer_rope_qk_inplace. " + "Please install flashinfer or disable this optimization." + ) from e + + if positions is None: + pos_1d = torch.arange(seqlen, device="cpu", dtype=torch.long) + positions = pos_1d if bsz == 1 else pos_1d.repeat(bsz) + else: + if not ( + isinstance(positions, torch.Tensor) + and positions.dtype == torch.long + and positions.dim() == 1 + ): + raise ValueError("positions must be a 1D torch.long Tensor") + if positions.numel() != bsz * seqlen: + raise ValueError( + f"positions length must be bsz*seqlen={bsz*seqlen}, got {positions.numel()}" + ) + + positions = positions.to(q.device, non_blocking=True) + + q_flat = q.reshape(bsz * seqlen, nheads * d).contiguous() + k_flat = k.reshape(bsz * seqlen, nheads * d).contiguous() + apply_rope_with_cos_sin_cache_inplace( + positions=positions, + query=q_flat, + key=k_flat, + head_size=d, + cos_sin_cache=cos_sin_cache, + is_neox=is_neox, + ) + return q_flat.view(bsz, seqlen, nheads, d), k_flat.view(bsz, seqlen, nheads, d) + + def _rotate_neox(x: torch.Tensor) -> torch.Tensor: x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index dc88afa50..d7ee2fd6d 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux.py @@ -40,7 +40,7 @@ from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear from sglang.multimodal_gen.runtime.layers.mlp import MLP from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( NDRotaryEmbedding, - _apply_rotary_emb, + apply_flashinfer_rope_qk_inplace, ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import current_platform @@ -182,11 +182,15 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin): if freqs_cis is not None: cos, sin = freqs_cis - query = _apply_rotary_emb( - query, cos, sin, is_neox_style=False, interleaved=False + cos_sin_cache = torch.cat( + [ + cos.to(dtype=torch.float32).contiguous(), + sin.to(dtype=torch.float32).contiguous(), + ], + dim=-1, ) - key = _apply_rotary_emb( - key, cos, sin, is_neox_style=False, interleaved=False + query, key = apply_flashinfer_rope_qk_inplace( + query, key, cos_sin_cache, is_neox=False ) x = self.attn(query, key, value) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py index b99f893b9..232af8fca 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -26,7 +26,7 @@ from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( NDRotaryEmbedding, - _apply_rotary_emb, + apply_flashinfer_rope_qk_inplace, ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import current_platform @@ -187,11 +187,15 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin): if freqs_cis is not None: cos, sin = freqs_cis - query = _apply_rotary_emb( - query, cos, sin, is_neox_style=False, interleaved=True + cos_sin_cache = torch.cat( + [ + cos.to(dtype=torch.float32).contiguous(), + sin.to(dtype=torch.float32).contiguous(), + ], + dim=-1, ) - key = _apply_rotary_emb( - key, cos, sin, is_neox_style=False, interleaved=True + query, key = apply_flashinfer_rope_qk_inplace( + query, key, cos_sin_cache, is_neox=False ) hidden_states = self.attn(query, key, value) @@ -311,11 +315,15 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): if freqs_cis is not None: cos, sin = freqs_cis - query = _apply_rotary_emb( - query, cos, sin, is_neox_style=False, interleaved=True + cos_sin_cache = torch.cat( + [ + cos.to(dtype=torch.float32).contiguous(), + sin.to(dtype=torch.float32).contiguous(), + ], + dim=-1, ) - key = _apply_rotary_emb( - key, cos, sin, is_neox_style=False, interleaved=True + query, key = apply_flashinfer_rope_qk_inplace( + query, key, cos_sin_cache, is_neox=False ) hidden_states = self.attn(query, key, value) hidden_states = hidden_states.flatten(2, 3) 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 b0f61c77d..21bf35481 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -19,6 +19,9 @@ from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConf 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.rotary_embedding import ( + apply_flashinfer_rope_qk_inplace, +) from sglang.multimodal_gen.runtime.layers.triton_ops import ( fuse_scale_shift_gate_select01_kernel, fuse_scale_shift_kernel, @@ -30,12 +33,6 @@ 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 ): @@ -561,9 +558,6 @@ class QwenImageCrossAttention(nn.Module): # Apply RoPE if image_rotary_emb is not None: - 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 @@ -572,41 +566,12 @@ class QwenImageCrossAttention(nn.Module): 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) + 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 + ) # Concatenate for joint attention # Order: [text, image]