[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)