[ROCm] [CI] Migrate CI to mi300x for v0.27.x (#5886)

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
Signed-off-by: TJian <tunjian.tan@embeddedllm.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/ar_diffusion/test_paged_attention.py

@@ -60,9 +60,18 @@ def _dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
     return torch.einsum("bhqk,bkhd->bqhd", probs, value)
 
 
-def _cuda_flash_attn_usable() -> bool:
+def _gpu_flash_attn_usable() -> bool:
     if not torch.cuda.is_available():
         return False
+    if torch.version.hip is not None:
+        for module in ("aiter", "flash_attn"):
+            try:
+                imported = __import__(module, fromlist=["flash_attn_varlen_func"])
+                if getattr(imported, "flash_attn_varlen_func", None) is not None:
+                    return True
+            except ImportError:
+                pass
+        return False
     try:
         spec = find_spec("vllm.vllm_flash_attn")
         if spec is None or spec.origin is None:
@@ -227,12 +236,11 @@ def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len,
     assert st.adapter(POS).completed_chunks == before + (1 if commit_current else 0)
 
 
-@pytest.mark.skipif(not _cuda_flash_attn_usable(), reason="usable CUDA FlashAttention is required")
+@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
 @pytest.mark.parametrize("history_chunks", [1, 3])
 @pytest.mark.parametrize("action_len", [0, 3])
 @pytest.mark.parametrize("commit_current", [False, True])
 def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len, commit_current):
-    pytest.importorskip("vllm.vllm_flash_attn")
     torch.manual_seed(0)
     device = torch.device("cuda")
     dtype = torch.float16

tests/diffusion/attention/test_flash_attn.py

@@ -265,10 +265,15 @@ def test_fa_vs_sdpa():
     print(f"Max absolute difference (valid region): {max_diff:.6f}")
     print(f"Mean absolute difference (valid region): {mean_diff:.6f}")
 
-    # Assert that outputs are close
-    # Using higher tolerance for bfloat16 and different implementations
-    assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
-    assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
+    # AITER's ROCm BF16 kernel is deterministic, but differs from ROCm SDPA by
+    # about 1.3% relative error. An absolute-only bound misclassifies that
+    # expected rounding difference for larger output values (for example,
+    # -2.65625 vs -2.6875). Keep the tighter historical bounds elsewhere.
+    if current_omni_platform.is_rocm():
+        torch.testing.assert_close(output_fa_valid, output_sdpa_valid, rtol=1e-2, atol=1.5e-2)
+    else:
+        assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
+        assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
 
     print("✓ Case 2 PASSED: FA and SDPA outputs are very close!")
 

tests/diffusion/distributed/test_sequence_parallel.py

@@ -416,8 +416,8 @@ def test_sp_correctness_advanced(model_name: str):
 
 
 @pytest.mark.skipif(
-    not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
-    reason="Only tested on CUDA and XPU",
+    not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()),
+    reason="Only tested on CUDA, ROCM and XPU",
 )
 @pytest.mark.diffusion
 @pytest.mark.parallel
@@ -447,8 +447,8 @@ def test_local_sp_padding_mask(monkeypatch: pytest.MonkeyPatch) -> None:
 
 
 @pytest.mark.skipif(
-    not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
-    reason="Only tested on CUDA and XPU",
+    not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()),
+    reason="Only tested on CUDA, ROCM and XPU",
 )
 @pytest.mark.diffusion
 @pytest.mark.parallel
@@ -468,8 +468,8 @@ def test_local_sp_padding_mask_no_padding(monkeypatch: pytest.MonkeyPatch) -> No
 
 
 @pytest.mark.skipif(
-    not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
-    reason="Only tested on CUDA and XPU",
+    not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()),
+    reason="Only tested on CUDA, ROCM and XPU",
 )
 @pytest.mark.diffusion
 @pytest.mark.parallel

tests/diffusion/lora/test_images_generations_lora.py

@@ -24,11 +24,20 @@ from safetensors.torch import save_file
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServer, OmniServerParams
+from vllm_omni.platforms import current_omni_platform
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 MODEL = "Tongyi-MAI/Z-Image-Turbo"
 
+# ROCm's default FLASH_ATTN diffusion backend resolves to AITER. On a cold CI
+# worker, JIT-building its MHA kernel can take longer than the server readiness
+# timeout. This test validates LoRA switching rather than attention-backend or
+# compiled-execution performance, so use the eager SDPA path on ROCm only.
+ROCM_SERVER_ARGS = (
+    ["--enforce-eager", "--diffusion-attention-backend", "TORCH_SDPA"] if current_omni_platform.is_rocm() else []
+)
+
 
 PROMPT = "a photo of a cat sitting on a laptop keyboard"
 SIZE = "256x256"
@@ -42,7 +51,14 @@ server_params = [
             server_args=[
                 "--num-gpus",
                 "1",
+                *ROCM_SERVER_ARGS,
             ],
+            # A cold-cache Z-Image load can take slightly more than the
+            # orchestrator's 900s default on CI (weights alone have been
+            # observed at ~690s and full orchestrator readiness at ~906s).
+            # Match the large-model offline runner's startup allowance without
+            # weakening timeout detection for every online-serving test.
+            init_timeout=1800,
         ),
         id="zimage_turbo_lora",
     ),