From d50ea9fcccf6879da44f2092ff7989768c8e3666 Mon Sep 17 00:00:00 2001 From: TJian Date: Thu, 30 Jul 2026 11:27:32 +0800 Subject: [PATCH] [ROCm] [CI] Fix ROCm CI for v0.26.0 (#5528) Signed-off-by: tjtanaa --- tests/config/test_config_factory.py | 11 ++ tests/diffusion/attention/test_selector.py | 138 ++++++++++++++++++ .../models/cosmos3/test_cosmos3_pipeline.py | 5 + vllm_omni/deploy/qwen3_tts.yaml | 7 + vllm_omni/platforms/musa/platform.py | 4 +- vllm_omni/platforms/rocm/platform.py | 14 ++ 6 files changed, 178 insertions(+), 1 deletion(-) create mode 100644 tests/diffusion/attention/test_selector.py diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 5a59dcf32..3e482f48f 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -2164,6 +2164,17 @@ class TestBaseConfigInheritance: class TestPlatformOverrides: """Test platform-specific deploy config overrides.""" + def test_qwen3_tts_rocm_disables_code2wav_outer_cudagraph(self): + deploy_path = Path(get_deploy_config_path("qwen3_tts.yaml")) + + base = load_deploy_config(deploy_path) + assert base.stages[0].enforce_eager is None + assert base.stages[1].enforce_eager is False + + rocm = _apply_platform_overrides(base, platform="rocm") + assert rocm.stages[0].enforce_eager is None + assert rocm.stages[1].enforce_eager is True + def test_npu_overrides(self): deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not deploy_path.exists(): diff --git a/tests/diffusion/attention/test_selector.py b/tests/diffusion/attention/test_selector.py new file mode 100644 index 000000000..79755ac2e --- /dev/null +++ b/tests/diffusion/attention/test_selector.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for diffusion attention backend selection.""" + +from types import SimpleNamespace + +import pytest + +import vllm_omni.diffusion.attention.selector as selector +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +class _ConfiguredBackend: + @classmethod + def get_name(cls) -> str: + return "CONFIGURED" + + +class _PlatformBackend: + @classmethod + def get_name(cls) -> str: + return "PLATFORM" + + +@pytest.fixture(autouse=True) +def clear_selector_caches(): + selector._cached_get_backend_cls.cache_clear() + selector._log_backend_resolution.cache_clear() + yield + selector._cached_get_backend_cls.cache_clear() + selector._log_backend_resolution.cache_clear() + + +@pytest.mark.parametrize( + ("role", "role_category", "expected_backend"), + [ + ("ltx2.audio_to_video", "cross", "EXACT"), + ("ltx2.video_to_audio", "cross", "CATEGORY"), + ("joint", None, "DEFAULT"), + ], +) +def test_configured_selection_precedence( + monkeypatch, + role: str, + role_category: str | None, + expected_backend: str, +): + config = AttentionConfig( + default=AttentionSpec(backend="DEFAULT"), + per_role={ + "cross": AttentionSpec(backend="CATEGORY"), + "ltx2.audio_to_video": AttentionSpec(backend="EXACT"), + }, + ) + calls = [] + + def fake_get_backend(backend_name, head_size, allow_trtllm_default=True): + calls.append((backend_name, head_size, allow_trtllm_default)) + return _ConfiguredBackend + + monkeypatch.setattr(selector, "_cached_get_backend_cls", fake_get_backend) + + backend, spec = selector.get_attn_backend_for_role( + role=role, + role_category=role_category, + head_size=128, + attention_config=config, + ) + + assert backend is _ConfiguredBackend + assert spec is not None + assert spec.backend == expected_backend + assert calls == [(expected_backend, 128, True)] + + +@pytest.mark.parametrize("attention_config", [None, AttentionConfig()]) +def test_platform_default_used_without_resolved_spec(monkeypatch, attention_config): + calls = [] + + def fake_get_backend(backend_name, head_size, allow_trtllm_default=True): + calls.append((backend_name, head_size, allow_trtllm_default)) + return _PlatformBackend + + monkeypatch.setattr(selector, "_cached_get_backend_cls", fake_get_backend) + + backend, spec = selector.get_attn_backend_for_role( + role="self", + head_size=64, + attention_config=attention_config, + allow_trtllm_default=False, + ) + + assert backend is _PlatformBackend + assert spec is None + assert calls == [(None, 64, False)] + + +def test_backend_class_resolution_is_cached(monkeypatch): + platform_calls = [] + load_calls = [] + fake_platform = SimpleNamespace( + get_diffusion_attn_backend_cls=lambda **kwargs: (platform_calls.append(kwargs) or "fake.module.Backend") + ) + + monkeypatch.setattr("vllm_omni.platforms.current_omni_platform", fake_platform) + + def fake_load_backend(path): + load_calls.append(path) + return _ConfiguredBackend + + monkeypatch.setattr(selector, "_load_backend_cls", fake_load_backend) + + first = selector._cached_get_backend_cls("FLASH_ATTN", 128, False) + second = selector._cached_get_backend_cls("FLASH_ATTN", 128, False) + + assert first is second is _ConfiguredBackend + assert platform_calls == [ + { + "selected_backend": "FLASH_ATTN", + "head_size": 128, + "allow_trtllm_default": False, + } + ] + assert load_calls == ["fake.module.Backend"] + + +def test_load_backend_cls_reports_missing_module(): + with pytest.raises(ImportError, match="Failed to import module missing_attention_backend"): + selector._load_backend_cls("missing_attention_backend.Backend") + + +def test_load_backend_cls_reports_missing_class(monkeypatch): + monkeypatch.setattr(selector.importlib, "import_module", lambda _: SimpleNamespace()) + + with pytest.raises(AttributeError, match="Class MissingBackend not found in module"): + selector._load_backend_cls("fake.module.MissingBackend") diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 3093e6ceb..994acd5a2 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -16,6 +16,7 @@ from PIL import Image from torch import nn from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -857,6 +858,10 @@ def test_flow_unipc_reuses_scheduler_and_forwards_each_request_shift(make_cosmos assert all(call["sigmas"] is None for call in scheduler.set_timesteps_calls) +@pytest.mark.skipif( + current_omni_platform.is_rocm(), + reason="Pytorch shipped with vLLM does not built with LAPACK support", +) def test_flow_unipc_reproducible_with_same_seed(make_cosmos3_pipeline) -> None: from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import ( FlowUniPCMultistepScheduler, diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 8f4c6440e..ad154bed2 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -97,6 +97,13 @@ stages: repetition_penalty: 1.0 platforms: + rocm: + stages: + - stage_id: 1 + # MIOpen Conv1d cannot be invoked while vLLM captures the outer + # Code2Wav model graph. Keep CUDA's graph-enabled default, but run + # Stage 1 eager on ROCm so decoder convolutions execute normally. + enforce_eager: true npu: stages: - stage_id: 0 diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index c71b5a70b..0774bca3a 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/vllm_omni/platforms/musa/platform.py @@ -45,6 +45,7 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): cls, selected_backend: str | None, head_size: int, + allow_trtllm_default: bool = False, ) -> str: """Get the diffusion attention backend class path for MUSA platform. @@ -54,7 +55,8 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): selected_backend: User-selected backend name (e.g., "FLASH_ATTN", "TORCH_SDPA"). If None, uses platform default. head_size: Attention head size. - + allow_trtllm_default: Does not support TRTLLM backend; + arg accepted for signature parity but unused. Returns: Fully qualified class path of the selected backend. """ diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py index cf72baccb..6ae65ed2b 100644 --- a/vllm_omni/platforms/rocm/platform.py +++ b/vllm_omni/platforms/rocm/platform.py @@ -75,7 +75,21 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform): cls, selected_backend: str | None, head_size: int, + allow_trtllm_default: bool = False, ) -> str: + """Get the diffusion attention backend class path for ROCm platform. + + ROCm supports FLASH_ATTN via the aiter library, and SDPA as fallback. + + Args: + selected_backend: User-selected backend name (e.g., "FLASH_ATTN", + "TORCH_SDPA"). If None, uses platform default. + head_size: Attention head size. + allow_trtllm_default: Does not support TRTLLM backend; + arg accepted for signature parity but unused. + Returns: + Fully qualified class path of the selected backend. + """ from vllm._aiter_ops import is_aiter_found_and_supported # Check if aiter is available for Flash Attention support