From 28558cc37471da8258c95aa515363a4a05fce601 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Zeyu=20Huang=20=7C=20=E9=BB=83=E6=BE=A4=E5=AE=87?= <11222265+fhfuih@users.noreply.github.com> Date: Wed, 6 May 2026 21:11:35 +0800 Subject: [PATCH] [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> --- .../diffusion/test_diffusion_model_runner.py | 134 ++++++++++++++++++ .../diffusion/test_diffusion_step_pipeline.py | 18 ++- vllm_omni/diffusion/forward_context.py | 8 +- .../diffusion/worker/diffusion_worker.py | 6 + vllm_omni/platforms/cuda/platform.py | 17 +++ vllm_omni/platforms/rocm/platform.py | 21 +++ vllm_omni/platforms/xpu/platform.py | 9 ++ 7 files changed, 209 insertions(+), 4 deletions(-) diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index b63f6d888..cb12019bb 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/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 (): 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 diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index 9417d540b..a443363a9 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/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, diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index 46b2b4a5c..f6df4730a 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/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 diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 887ad27d3..5c8cde3e2 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/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 diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index 6d955dcfb..6afdfcb66 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/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) diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py index 84aca4629..87cc08418 100644 --- a/vllm_omni/platforms/rocm/platform.py +++ b/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) diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index bd50163dc..b2982d263 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/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)