[ROCm] [CI] Fix ROCm CI for v0.26.0 (#5528)

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

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():

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

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,

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

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.
         """

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