[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

tests/diffusion/test_diffusion_step_pipeline.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Tests for step-level diffusion execution across runner / worker / executor / engine."""
 
+import contextlib
 import os
 import queue
 import threading
@@ -230,9 +231,20 @@ def _make_engine_request(req_id: str = "req-1", num_inference_steps: int = 2) ->
     )
 
 
+def _make_vllm_config():
+    @contextlib.contextmanager
+    def set_priority(*args, **kwargs):
+        yield
+
+    return SimpleNamespace(
+        kernel_config=SimpleNamespace(ir_op_priority=SimpleNamespace(set_priority=set_priority)),
+        compilation_config=SimpleNamespace(ir_enable_torch_wrap=True),
+    )
+
+
 def _make_runner():
     runner = object.__new__(DiffusionModelRunner)
-    runner.vllm_config = object()
+    runner.vllm_config = _make_vllm_config()
     runner.od_config = SimpleNamespace(
         cache_backend=None,
         parallel_config=SimpleNamespace(use_hsdp=False),
@@ -248,7 +260,7 @@ def _make_runner():
 
 def _make_distributed_runner(mode: str, device: torch.device):
     runner = object.__new__(DiffusionModelRunner)
-    runner.vllm_config = object()
+    runner.vllm_config = _make_vllm_config()
     runner.od_config = SimpleNamespace(
         cache_backend=None,
         parallel_config=SimpleNamespace(use_hsdp=False),
@@ -449,7 +461,7 @@ class TestRunner:
                 return False
 
         runner = object.__new__(DiffusionModelRunner)
-        runner.vllm_config = object()
+        runner.vllm_config = _make_vllm_config()
         runner.od_config = SimpleNamespace(
             enable_cpu_offload=False,
             enable_layerwise_offload=False,

vllm_omni/diffusion/forward_context.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from contextlib import contextmanager
 from dataclasses import dataclass
 
+import vllm.ir
 from vllm.config import VllmConfig
 
 from vllm_omni.diffusion.attention.backends.abstract import (
@@ -145,6 +146,7 @@ def set_forward_context(
     )
     # vLLM CustomOp dispatch (e.g. QKVParallelLinear) requires a global
     # vLLM config set via set_current_vllm_config().
+    # Also set priority for vLLM IR ops (e.g. RMSNorm), copied from vllm/forward_context.py
     with override_forward_context(forward_context):
         if vllm_config is None:
             yield
@@ -152,5 +154,9 @@ def set_forward_context(
             # Local import to avoid importing vllm.config.vllm at module import time.
             from vllm.config.vllm import set_current_vllm_config
 
-            with set_current_vllm_config(vllm_config):
+            with (
+                set_current_vllm_config(vllm_config),
+                vllm_config.kernel_config.ir_op_priority.set_priority(),
+                vllm.ir.enable_torch_wrap(vllm_config.compilation_config.ir_enable_torch_wrap),
+            ):
                 yield

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -136,6 +136,12 @@ class DiffusionWorker:
             enable_return_routed_experts=False,
         )
         vllm_config.quant_config = self.od_config.quantization_config
+        # Since vLLM v0.20.0, IR wraps GPU ops. Set IR op priority preference to enforce GPU op fusion during wrapping.
+        # Also need to log, because vLLM internally logs another line in VllmConfig.__post_init__. Avoid confusion.
+        vllm_config.kernel_config.ir_op_priority = current_omni_platform.get_default_ir_op_priority(vllm_config)
+        logger.info(
+            "Final IR op priority after setting vLLM-Omni overrides: %s", vllm_config.kernel_config.ir_op_priority
+        )
         self.vllm_config = vllm_config
 
         # Initialize distributed environment

vllm_omni/platforms/cuda/platform.py

@@ -2,6 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import torch
+import vllm.envs as envs
+from vllm.config import VllmConfig
+from vllm.config.kernel import IrOpPriorityConfig
 from vllm.logger import init_logger
 from vllm.platforms.cuda import CudaPlatformBase
 from vllm.platforms.interface import DeviceCapability
@@ -129,3 +132,17 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
     @classmethod
     def get_device_name(cls, device_id: int = 0) -> str:
         return torch.cuda.get_device_name(device_id)
+
+    @classmethod
+    def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
+        """Copied from vllm/platforms/cuda/platform.py v0.20.0 with force using vllm_c kernels"""
+        default = ["vllm_c", "native"]  # Originally using "native" here when compiling
+
+        # Use oink if enabled for rms_norm
+        # TODO(Laurawly/luka): remove this env var,
+        #  users can just use IR op priority directly
+        rms_norm = default
+        if envs.VLLM_USE_OINK_OPS:
+            rms_norm = ["oink"] + default
+
+        return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm)

vllm_omni/platforms/rocm/platform.py

@@ -2,6 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import torch
+from vllm import envs
+from vllm.config import VllmConfig
+from vllm.config.kernel import IrOpPriorityConfig
 from vllm.logger import init_logger
 from vllm.platforms.rocm import RocmPlatform
 
@@ -147,3 +150,21 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform):
 
         os.environ.pop("HIP_VISIBLE_DEVICES", None)
         os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+
+    @classmethod
+    def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
+        """Copied from vllm/platforms/rocm/platform.py v0.20.0 with force using vllm_c kernels"""
+        # TODO(luka/TJ) use aiter, vllm_c, native by default on ROCm
+        cc = vllm_config.compilation_config
+        default = ["vllm_c", "native"]  # Originally using "native" here when compiling
+
+        # This (mostly) preserves previous CustomOp behavior
+        # Necessary on ROCm because it's common that users
+        # enable rms_norm to use the aiter kernel.
+        # TODO(luka/TJ) remove env vars completely
+        if cc.is_custom_op_enabled("rms_norm") and envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_RMSNORM:
+            rms_norm = ["aiter"] + default
+        else:
+            rms_norm = default
+
+        return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm)

vllm_omni/platforms/xpu/platform.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import torch
+from vllm.config import VllmConfig
+from vllm.config.kernel import IrOpPriorityConfig
 from vllm.logger import init_logger
 from vllm.platforms.xpu import XPUPlatform
 
@@ -79,3 +81,10 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):
     def get_profiler_cls(cls) -> str:
         """Return XPU-specific profiler that handles XPU events."""
         return "vllm_omni.platforms.xpu.profiler.XPUTorchProfilerWrapper"
+
+    @classmethod
+    def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
+        """Copied from vllm/platforms/xpu/platform.py v0.20.0 with force using xpu_kernels kernels"""
+        default = ["xpu_kernels", "native"]  # Originally using "native" here when compiling
+
+        return IrOpPriorityConfig.with_default(default)