[Bugfix] Gate Qwen-Image fused QK-norm+RoPE on short sequences (#7780) (#7965)

Signed-off-by: NumberWan <wantszkin2003@gmail.com>

tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py

@@ -10,7 +10,9 @@ import torch
 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 test_qwen_image_fused_qk_norm_rope_torch_compile_fullgraph_capture():
 
     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)

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -48,6 +48,7 @@ from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
 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 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 
 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,