fix: preserve Qwen Image rotary frequency precision on CUDA

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py

@@ -1,12 +1,39 @@
 import pytest
 import torch
 
-from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import QwenEmbedRope
+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)

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -50,6 +50,17 @@ 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 _normalize_qwen_image_weight_name(name: str) -> str:
     name = name.removeprefix("transformer.")
     if ".to_out.0." in name:
@@ -642,15 +653,23 @@ class QwenImageCrossAttention(nn.Module):
         txt_query = self.norm_added_q(txt_query)
         txt_key = self.norm_added_k(txt_key)
 
-        img_cos = torch.real(vid_freqs).to(img_query.dtype)
-        img_sin = torch.imag(vid_freqs).to(img_query.dtype)
-        txt_cos = torch.real(txt_freqs).to(txt_query.dtype)
-        txt_sin = torch.imag(txt_freqs).to(txt_query.dtype)
+        if img_query.device.type == "cuda":
+            img_query = _apply_qwen_image_rotary_emb(img_query, vid_freqs)
+            img_key = _apply_qwen_image_rotary_emb(img_key, vid_freqs)
+            txt_query = _apply_qwen_image_rotary_emb(txt_query, txt_freqs)
+            txt_key = _apply_qwen_image_rotary_emb(txt_key, txt_freqs)
+        else:
+            # Retain the platform-specific kernels on other accelerators,
+            # which may not support complex tensors.
+            img_cos = torch.real(vid_freqs).to(img_query.dtype)
+            img_sin = torch.imag(vid_freqs).to(img_query.dtype)
+            txt_cos = torch.real(txt_freqs).to(txt_query.dtype)
+            txt_sin = torch.imag(txt_freqs).to(txt_query.dtype)
 
-        img_query = self.rope(img_query, img_cos, img_sin)
-        img_key = self.rope(img_key, img_cos, img_sin)
-        txt_query = self.rope(txt_query, txt_cos, txt_sin)
-        txt_key = self.rope(txt_key, txt_cos, txt_sin)
+            img_query = self.rope(img_query, img_cos, img_sin)
+            img_key = self.rope(img_key, img_cos, img_sin)
+            txt_query = self.rope(txt_query, txt_cos, txt_sin)
+            txt_key = self.rope(txt_key, txt_cos, txt_sin)
 
         seq_len_txt = encoder_hidden_states.shape[1]
         joint_query = torch.cat([txt_query, img_query], dim=1)