[Bugfix][Qwen-Image] Restore RotaryEmbedding CUDA RoPE for Diffusers e2e (#7513)

Signed-off-by: NumberWan <wantszkin2003@gmail.com>
Co-authored-by: wangyu <53896905+yenuo26@users.noreply.github.com>

benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py

@@ -22,9 +22,7 @@ from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
     _fused_cuda_supported,
     fused_qk_norm_rope,
 )
-from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
-    _apply_qwen_image_rotary_emb,
-)
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 
 
 def _parse_args() -> argparse.Namespace:
@@ -87,10 +85,10 @@ def _native(
 ) -> tuple[torch.Tensor, torch.Tensor]:
     q = F.rms_norm(q, (q.shape[-1],), q_weight, eps)
     k = F.rms_norm(k, (k.shape[-1],), k_weight, eps)
-    return (
-        _apply_qwen_image_rotary_emb(q, freqs),
-        _apply_qwen_image_rotary_emb(k, freqs),
-    )
+    rope = RotaryEmbedding(is_neox_style=False)
+    cos = freqs.real.to(q.dtype)
+    sin = freqs.imag.to(q.dtype)
+    return rope(q, cos, sin), rope(k, cos, sin)
 
 
 def _fused(

tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py

@@ -11,7 +11,6 @@ import torch.nn as nn
 
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
-    _apply_qwen_image_rotary_emb,
     _qwen_image_qk_norm_rope,
 )
 
@@ -73,22 +72,28 @@ def _make_input(
     return QwenImageQKInput(q=q, k=k, norm_q=norm_q, norm_k=norm_k, freqs=freqs)
 
 
-def _reference(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]:
+def _eager_rotary(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]:
     q = data.norm_q(data.q)
     k = data.norm_k(data.k)
-    if data.q.device.type == "cuda":
-        return (
-            _apply_qwen_image_rotary_emb(q, data.freqs),
-            _apply_qwen_image_rotary_emb(k, data.freqs),
-        )
-
     rope = RotaryEmbedding(is_neox_style=False)
     cos = data.freqs.real.to(data.q.dtype)
     sin = data.freqs.imag.to(data.q.dtype)
     return rope(q, cos, sin), rope(k, cos, sin)
 
 
-def _run(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]:
+def _fp32_complex_rotary(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+    """Local copy of the old CUDA helper; fused kernel still tracks this math."""
+    paired = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+    return torch.view_as_real(paired * freqs.unsqueeze(1)).flatten(3).to(x.dtype)
+
+
+def _fused_kernel_reference(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]:
+    q = data.norm_q(data.q)
+    k = data.norm_k(data.k)
+    return _fp32_complex_rotary(q, data.freqs), _fp32_complex_rotary(k, data.freqs)
+
+
+def _run(data: QwenImageQKInput, *, use_fused: bool = True) -> tuple[torch.Tensor, torch.Tensor]:
     return _qwen_image_qk_norm_rope(
         data.q,
         data.k,
@@ -97,6 +102,7 @@ def _run(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]:
         data.freqs,
         RotaryEmbedding(is_neox_style=False),
         EPS,
+        use_fused=use_fused,
     )
 
 
@@ -110,8 +116,8 @@ def test_qwen_image_qk_norm_rope_cuda_fp32_fallback_matches_reference():
         packed_qkv_view=True,
     )
 
-    actual_q, actual_k = _run(data)
-    expected_q, expected_k = _reference(data)
+    actual_q, actual_k = _run(data, use_fused=False)
+    expected_q, expected_k = _eager_rotary(data)
 
     torch.testing.assert_close(actual_q, expected_q, atol=1e-5, rtol=1e-5)
     torch.testing.assert_close(actual_k, expected_k, atol=1e-5, rtol=1e-5)
@@ -127,8 +133,8 @@ def test_qwen_image_qk_norm_rope_cuda_fp16_fallback_matches_reference():
         packed_qkv_view=True,
     )
 
-    actual_q, actual_k = _run(data)
-    expected_q, expected_k = _reference(data)
+    actual_q, actual_k = _run(data, use_fused=False)
+    expected_q, expected_k = _eager_rotary(data)
 
     torch.testing.assert_close(actual_q, expected_q, atol=1e-3, rtol=1e-3)
     torch.testing.assert_close(actual_k, expected_k, atol=1e-3, rtol=1e-3)
@@ -150,7 +156,7 @@ def test_qwen_image_fused_qk_norm_rope_cuda_matches_fp32_rope_reference(
     )
 
     actual_q, actual_k = _run(data)
-    expected_q, expected_k = _reference(data)
+    expected_q, expected_k = _fused_kernel_reference(data)
 
     torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02)
     torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02)

tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py

@@ -3,37 +3,12 @@ import torch
 
 from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
     QwenEmbedRope,
-    _apply_qwen_image_rotary_emb,
 )
 from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("noncontiguous", [False, True])
-def test_qwen_rotary_preserves_frequency_precision(dtype, noncontiguous):
-    from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen
-
-    generator = torch.Generator().manual_seed(42)
-    x = torch.randn(2, 17, 3, 128, generator=generator).to(dtype)
-    if noncontiguous:
-        x = x.transpose(1, 2).contiguous().transpose(1, 2)
-    angles = torch.randn(17, 64, generator=generator)
-    freqs = torch.polar(torch.ones_like(angles), angles)
-    expected = apply_rotary_emb_qwen(x, freqs, use_real=False)
-    actual = _apply_qwen_image_rotary_emb(x, freqs)
-    assert actual.dtype == dtype
-    assert actual.shape == x.shape
-    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
-
-    if dtype == torch.bfloat16:
-        # A regression to early BF16 frequency rounding must be observable.
-        rounded_freqs = torch.complex(freqs.real.to(dtype).float(), freqs.imag.to(dtype).float())
-        rounded = _apply_qwen_image_rotary_emb(x, rounded_freqs)
-        assert not torch.equal(actual, rounded)
-
-
 def test_txt_seq_lens_from_embeds_uses_padded_width_not_valid_token_count():
     prompt_embeds = torch.zeros(2, 64, 8)
     prompt_embeds_mask = torch.zeros(2, 64, dtype=torch.bool)

tests/e2e/accuracy/test_qwen_image.py

@@ -32,9 +32,9 @@ HEIGHT = 512
 NUM_INFERENCE_STEPS = 20
 TRUE_CFG_SCALE = 4.0
 SEED = 42
-SSIM_THRESHOLD = 0.94
-# With --fa-deterministic, H100 nightlies stably land near SSIM 0.958 / PSNR 27.8
-# vs Diffusers on the FA3 hub path (see #5734 / #5963 / #6971).
+# H100 nightly after RotaryEmbedding CUDA RoPE: SSIM 0.978375 / PSNR 31.83
+# vs Diffusers (Buildkite #15339). SSIM 0.97 so a drop is visible.
+SSIM_THRESHOLD = 0.97
 PSNR_THRESHOLD = 27.0
 
 MODEL_2512_ID = "Qwen/Qwen-Image-2512"

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -60,17 +60,6 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 logger = init_logger(__name__)
 
 
-def _apply_qwen_image_rotary_emb(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
-    """Rotate interleaved pairs before rounding back to the activation dtype.
-
-    Qwen-Image's reference uses complex FP32 multiplication. Rounding the
-    frequencies to BF16 before rotation loses positional precision; those
-    errors accumulate across the denoising steps.
-    """
-    paired = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
-    return torch.view_as_real(paired * freqs.unsqueeze(1)).flatten(3).to(x.dtype)
-
-
 def _qwen_image_qk_norm_rope(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -103,16 +92,11 @@ def _qwen_image_qk_norm_rope(
             fused_k.reshape(batch, seq_len, num_kv_heads, head_dim),
         )
 
+    # Eager path: BF16/activation-dtype RotaryEmbedding on every device.
+    # CUDA used to call FP32 complex multiply here; that matches the Diffusers
+    # helper in unit tests but drops Omni vs Diffusers pipeline PSNR (#7494).
     q = norm_q(q)
     k = norm_k(k)
-    if q.device.type == "cuda":
-        return (
-            _apply_qwen_image_rotary_emb(q, freqs),
-            _apply_qwen_image_rotary_emb(k, freqs),
-        )
-
-    # Retain the platform-specific kernels on other accelerators, which may
-    # not support complex tensors.
     cos = torch.real(freqs).to(q.dtype)
     sin = torch.imag(freqs).to(q.dtype)
     return rope(q, cos, sin), rope(k, cos, sin)