From b751b6e62cf5e92dcaca04badfbc6a66016cb73d Mon Sep 17 00:00:00 2001 From: tzhouam Date: Tue, 8 Sep 2026 11:26:57 +0800 Subject: [PATCH] fix: preserve Qwen Image rotary frequency precision on CUDA Signed-off-by: tzhouam --- .../qwen_image/test_qwen_image_rope_utils.py | 29 ++++++++++++++- .../qwen_image/qwen_image_transformer.py | 35 ++++++++++++++----- 2 files changed, 55 insertions(+), 9 deletions(-) diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py b/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py index 52f1b6aa3..4f3a81491 100644 --- a/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py +++ b/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) diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index c82ace9c4..696728b90 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/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)