[Diffusion] Refactor qwen_image's rope in a single helper func (#16047)

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()

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 :]

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(
+            cos_sin_cache = torch.cat(
-                query, cos, sin, is_neox_style=False, interleaved=False
+                [
+                    cos.to(dtype=torch.float32).contiguous(),
+                    sin.to(dtype=torch.float32).contiguous(),
+                ],
+                dim=-1,
             )
-            key = _apply_rotary_emb(
+            query, key = apply_flashinfer_rope_qk_inplace(
-                key, cos, sin, is_neox_style=False, interleaved=False
+                query, key, cos_sin_cache, is_neox=False
             )
 
         x = self.attn(query, key, value)

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(
+            cos_sin_cache = torch.cat(
-                query, cos, sin, is_neox_style=False, interleaved=True
+                [
+                    cos.to(dtype=torch.float32).contiguous(),
+                    sin.to(dtype=torch.float32).contiguous(),
+                ],
+                dim=-1,
             )
-            key = _apply_rotary_emb(
+            query, key = apply_flashinfer_rope_qk_inplace(
-                key, cos, sin, is_neox_style=False, interleaved=True
+                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(
+            cos_sin_cache = torch.cat(
-                query, cos, sin, is_neox_style=False, interleaved=True
+                [
+                    cos.to(dtype=torch.float32).contiguous(),
+                    sin.to(dtype=torch.float32).contiguous(),
+                ],
+                dim=-1,
             )
-            key = _apply_rotary_emb(
+            query, key = apply_flashinfer_rope_qk_inplace(
-                key, cos, sin, is_neox_style=False, interleaved=True
+                query, key, cos_sin_cache, is_neox=False
             )
         hidden_states = self.attn(query, key, value)
         hidden_states = hidden_states.flatten(2, 3)

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(
+            img_query, img_key = apply_flashinfer_rope_qk_inplace(
-                q_4d: torch.Tensor, k_4d: torch.Tensor, cache: torch.Tensor
+                img_query, img_key, img_cache, is_neox=False
-            ):
+            )
-                bsz, seqlen, nheads, d = q_4d.shape
+            txt_query, txt_key = apply_flashinfer_rope_qk_inplace(
-
+                txt_query, txt_key, txt_cache, is_neox=False
-                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]