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 16dc851153e..70d586eff8b 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 @@ -10,7 +10,9 @@ import torch.nn as nn from vllm_omni.diffusion.layers.rope import RotaryEmbedding +from vllm_omni.diffusion.models.qwen_image import qwen_image_transformer as qwen_mod from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( + _FUSED_MIN_TOKENS, _qwen_image_qk_norm_rope, ) @@ -106,6 +108,13 @@ def _run(data: QwenImageQKInput, *, use_fused: bool = True) -> tuple[torch.Tenso ) +@pytest.fixture +def force_always_fuse(monkeypatch): + """Disable the short-seq host gate so fused-kernel correctness tests keep + exercising the Triton path on tiny shapes.""" + monkeypatch.setattr(qwen_mod, "fused_qk_norm_rope_min_tokens", lambda _default: 0) + + @pytest.mark.cuda @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_qwen_image_qk_norm_rope_cuda_fp32_fallback_matches_reference(): @@ -147,6 +156,7 @@ def test_qwen_image_qk_norm_rope_cuda_fp16_fallback_matches_reference(): def test_qwen_image_fused_qk_norm_rope_cuda_matches_fp32_rope_reference( seq_len: int, packed_qkv_view: bool, + force_always_fuse, ): data = _make_input( seq_len=seq_len, @@ -164,7 +174,7 @@ def test_qwen_image_fused_qk_norm_rope_cuda_matches_fp32_rope_reference( @pytest.mark.cuda @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_qwen_image_fused_qk_norm_rope_torch_compile_fullgraph_capture(): +def test_qwen_image_fused_qk_norm_rope_torch_compile_fullgraph_capture(force_always_fuse): data = _make_input( seq_len=257, dtype=torch.bfloat16, @@ -194,3 +204,44 @@ def fn( 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) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qwen_image_short_seq_default_gate_uses_eager_path(monkeypatch): + """Below _FUSED_MIN_TOKENS the helper must stay on the eager chain (#7780).""" + monkeypatch.delenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", raising=False) + # BATCH*seq_len = 2*512 = 1024 < 2048 + data = _make_input( + seq_len=512, + dtype=torch.bfloat16, + device=torch.device("cuda:0"), + packed_qkv_view=False, + ) + assert BATCH * 512 < _FUSED_MIN_TOKENS + + actual_q, actual_k = _run(data, use_fused=True) + expected_q, expected_k = _eager_rotary(data) + + torch.testing.assert_close(actual_q, expected_q, atol=0, rtol=0) + torch.testing.assert_close(actual_k, expected_k, atol=0, rtol=0) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qwen_image_above_gate_uses_fused_path(monkeypatch): + monkeypatch.delenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", raising=False) + # BATCH*seq_len = 2*1024 = 2048 == _FUSED_MIN_TOKENS → fuse + data = _make_input( + seq_len=1024, + dtype=torch.bfloat16, + device=torch.device("cuda:0"), + packed_qkv_view=False, + ) + assert BATCH * 1024 >= _FUSED_MIN_TOKENS + + actual_q, actual_k = _run(data, use_fused=True) + 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/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index cb1adbb61a2..ed7b66bce36 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -48,6 +48,7 @@ from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( _fused_cuda_supported, fused_qk_norm_rope, + fused_qk_norm_rope_min_tokens, ) from vllm_omni.diffusion.layers.qwen_select01_modulation import ( can_use_qwen_select01_triton, @@ -59,6 +60,10 @@ logger = init_logger(__name__) +# Fuse only when B*S >= this; below it host launch overhead dominates (#7780). +# Override: VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS (0 = always fuse). +_FUSED_MIN_TOKENS = 2048 + def _qwen_image_qk_norm_rope( q: torch.Tensor, @@ -73,14 +78,19 @@ def _qwen_image_qk_norm_rope( ) -> tuple[torch.Tensor, torch.Tensor]: head_dim = q.shape[-1] rotary_dim = freqs.shape[-1] * 2 - if use_fused and _fused_cuda_supported(q, k, head_dim, rotary_dim, interleaved=True): - batch, seq_len, num_heads, _ = q.shape + batch, seq_len, num_heads, _ = q.shape + tokens = batch * seq_len + if ( + use_fused + and tokens >= fused_qk_norm_rope_min_tokens(_FUSED_MIN_TOKENS) + and _fused_cuda_supported(q, k, head_dim, rotary_dim, interleaved=True) + ): num_kv_heads = k.shape[2] rope_table = torch.cat((freqs.real, freqs.imag), dim=-1) - rope_table = rope_table.unsqueeze(0).expand(batch, -1, -1).reshape(batch * seq_len, rotary_dim) + rope_table = rope_table.unsqueeze(0).expand(batch, -1, -1).reshape(tokens, rotary_dim) fused_q, fused_k = fused_qk_norm_rope( - q.reshape(batch * seq_len, num_heads, head_dim), - k.reshape(batch * seq_len, num_kv_heads, head_dim), + q.reshape(tokens, num_heads, head_dim), + k.reshape(tokens, num_kv_heads, head_dim), norm_q.weight, norm_k.weight, rope_table,