diff --git a/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py b/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py index ee69f7e59..2eaa097bd 100644 --- a/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py +++ b/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( diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py b/tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py index 2f10ea33e..16dc85115 100644 --- a/tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py +++ b/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) 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 4f3a81491..3fd141dc1 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 @@ -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) diff --git a/tests/e2e/accuracy/test_qwen_image.py b/tests/e2e/accuracy/test_qwen_image.py index fb2f8cbc8..1fca04f99 100644 --- a/tests/e2e/accuracy/test_qwen_image.py +++ b/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" 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 11568f62c..cb1adbb61 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/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)