[bugfix][CI] Fix qwen image performance degradation w/ vllm 0.20 & CUDA 13.0 (#3352)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>

tests/diffusion/test_diffusion_model_runner.py

@@ -171,3 +171,137 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
     assert runner.cache_backend is None
     assert runner.od_config.cache_backend is None
     assert dummy_cache_backend.enabled is False
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_set_forward_context_enters_vllm_config_contexts(monkeypatch):
+    """Ensure `with set_forward_context(...):` enters vllm's context managers internally and calls desired vllm functions."""
+    import vllm.config.vllm as vllm_config_module
+    import vllm.ir
+    from vllm.config import CompilationConfig, DeviceConfig, VllmConfig
+
+    from vllm_omni.diffusion.forward_context import (
+        get_forward_context,
+        is_forward_context_available,
+        set_forward_context,
+    )
+
+    vllm_config = VllmConfig(
+        device_config=DeviceConfig(device="cpu"),
+        compilation_config=CompilationConfig(),
+    )
+    calls = []
+
+    @contextmanager
+    def _set_current_vllm_config(cfg):
+        calls.append(("set_current_vllm_config", cfg))
+        yield
+        calls.append(("set_current_vllm_config_exit", cfg))
+
+    @contextmanager
+    def _set_priority(*args, **kwargs):
+        del args, kwargs
+        calls.append(("ir_op_priority", None))
+        yield
+        calls.append(("ir_op_priority_exit", None))
+
+    @contextmanager
+    def _enable_torch_wrap(flag):
+        calls.append(("enable_torch_wrap", flag))
+        yield
+        calls.append(("enable_torch_wrap_exit", flag))
+
+    monkeypatch.setattr(vllm_config_module, "set_current_vllm_config", _set_current_vllm_config)
+    monkeypatch.setattr(vllm_config.kernel_config.ir_op_priority, "set_priority", _set_priority)
+    monkeypatch.setattr(vllm.ir, "enable_torch_wrap", _enable_torch_wrap)
+
+    assert not is_forward_context_available()
+
+    with set_forward_context(vllm_config=vllm_config):
+        assert is_forward_context_available()
+        assert get_forward_context().vllm_config is vllm_config
+
+    assert not is_forward_context_available()
+    assert calls == [
+        ("set_current_vllm_config", vllm_config),
+        ("ir_op_priority", None),
+        ("enable_torch_wrap", vllm_config.compilation_config.ir_enable_torch_wrap),
+        ("enable_torch_wrap_exit", vllm_config.compilation_config.ir_enable_torch_wrap),
+        ("ir_op_priority_exit", None),
+        ("set_current_vllm_config_exit", vllm_config),
+    ]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_vllm_set_forward_context_implementation(monkeypatch):
+    """Regression test: ensure that vLLM's set_forward_context implementation has changed."""
+    import vllm.forward_context as vllm_forward_context
+    import vllm.ir
+    from vllm.config import CompilationConfig, DeviceConfig, VllmConfig
+
+    ERROR_MESSAGE = (
+        "If this test fails, it likely means that vLLM's set_forward_context (vllm/forward_context.py) implementation has changed. "
+        "In this case, we should update our forward_context (vllm_omni/diffusion/forward_context.py) as well. "
+        "We should at least confirm that the `try: with (<what's inside?>): yield` part does not miss any information "
+        "(typically by calling the same or similar stuff as vLLM). See #3352 for an example. "
+        "Then, update this test to reflect the new implementation, and also update test_set_forward_context_enters_vllm_config_contexts."
+    )
+
+    vllm_config = VllmConfig(
+        device_config=DeviceConfig(device="cpu"),
+        compilation_config=CompilationConfig(),
+    )
+    calls = []
+
+    @contextmanager
+    def _set_priority():
+        calls.append(("ir_op_priority", None))
+        yield
+        calls.append(("ir_op_priority_exit", None))
+
+    @contextmanager
+    def _enable_torch_wrap(flag):
+        calls.append(("enable_torch_wrap", flag))
+        yield
+        calls.append(("enable_torch_wrap_exit", flag))
+
+    def _set_additional_forward_context(**kwargs):
+        calls.append(("set_additional_forward_context", tuple(sorted(kwargs.keys()))))
+        return {}
+
+    monkeypatch.setattr(vllm_config.kernel_config.ir_op_priority, "set_priority", _set_priority)
+    monkeypatch.setattr(vllm.ir, "enable_torch_wrap", _enable_torch_wrap)
+    monkeypatch.setattr(
+        vllm_forward_context.current_platform,
+        "set_additional_forward_context",
+        _set_additional_forward_context,
+    )
+
+    assert not vllm_forward_context.is_forward_context_available(), ERROR_MESSAGE
+
+    with vllm_forward_context.set_forward_context(None, vllm_config):
+        assert vllm_forward_context.is_forward_context_available(), ERROR_MESSAGE
+        assert vllm_forward_context.get_forward_context().attn_metadata is None, ERROR_MESSAGE
+
+    assert not vllm_forward_context.is_forward_context_available(), ERROR_MESSAGE
+    assert calls == [
+        (
+            "set_additional_forward_context",
+            (
+                "attn_metadata",
+                "batch_descriptor",
+                "cudagraph_runtime_mode",
+                "dp_metadata",
+                "num_tokens",
+                "num_tokens_across_dp",
+                "ubatch_slices",
+                "vllm_config",
+            ),
+        ),
+        ("ir_op_priority", None),
+        ("enable_torch_wrap", vllm_config.compilation_config.ir_enable_torch_wrap),
+        ("enable_torch_wrap_exit", vllm_config.compilation_config.ir_enable_torch_wrap),
+        ("ir_op_priority_exit", None),
+    ], ERROR_MESSAGE