[CI/Build][ROCm] Stabilize R2-01 platform selection (#8561)

Signed-off-by: haic0 <149741444+haic0@users.noreply.github.com>
Signed-off-by: andyluo7 <andy.luo@amd.com>
Co-authored-by: andyluo7 <andy.luo@amd.com>

.buildkite/amd/test-amd-ready.yml

@@ -260,7 +260,7 @@ steps:
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
 #     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
 
-- label: "ROCm · Engine & Model Executor GPU Coverage (R2-01)"
+- label: "ROCm · Model Executor GPU Coverage (R2-01)"
   agent_pool: mi300_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
@@ -273,18 +273,17 @@ steps:
     # artifacts relative to its checkout directory.
     - export R2_01_ARTIFACT_DIR="$$BUILDKITE_BUILD_CHECKOUT_PATH/artifacts/rocm-r2-01"
     - mkdir -p "$$R2_01_ARTIFACT_DIR"
+    # The dedicated Engine Test below owns the only selected tests/engine case;
+    # do not duplicate its model startup in this time-bounded coverage lane.
     # The CUDA ready job also covers tests/worker/test_batched_omni_output.py.
-    # Keep the initial ROCm burn-in scoped to the requested directories.
+    # Keep the initial ROCm burn-in scoped to tests/model_executor/.
     # Verbose pytest output and JUnit retain node IDs without a second
     # collection-only pass over the same GPU-backed test tree.
-    # The CUDA/L4-only CosyVoice3 ragged-flow parity test cold-JITs multiple
-    # AITER kernels on ROCm and is covered by the CUDA ready lane.
     - >-
       pytest -s -v -ra
-      tests/engine/ tests/model_executor/
-      -m 'core_model and cuda and not
+      tests/model_executor/
+      -m 'core_model and cuda and not nvidia_only and not
       (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)'
-      -k 'not test_code2wav_streaming_batch_matches_ragged_flow_numerics'
       --run-level "core_model"
       --durations=0
       --junitxml="$$R2_01_ARTIFACT_DIR/pytest.xml"

pyproject.toml

@@ -282,6 +282,7 @@ markers = [
     "parallel: Parallelism/distributed tests",
     "sp: Sequence parallelism tests",
     "example: Doc example code tests",
+    "nvidia_only: Tests that require NVIDIA CUDA and must not run on ROCm",
     # platform markers
     # Tag with [hardware-platform] for tests.helpers.mark.get_supported_platforms().
     # Add [gpu] on accelerator platforms that should also get pytest.mark.gpu.

tests/buildkite/test_amd_ready_r2_01.py

@@ -6,15 +6,17 @@ from shlex import split
 
 import pytest
 import yaml
+from _pytest.mark.expression import Expression
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 AMD_READY_PIPELINE = Path(".buildkite/amd/test-amd-ready.yml")
-R2_01_LABEL = "ROCm · Engine & Model Executor GPU Coverage (R2-01)"
+R2_01_LABEL = "ROCm · Model Executor GPU Coverage (R2-01)"
+ENGINE_LABEL = "Engine Test"
 SINGLE_GPU_MARKERS = (
-    "core_model and cuda and not (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)"
+    "core_model and cuda and not nvidia_only and not "
+    "(cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)"
 )
-EXCLUDED_TEST = "test_code2wav_streaming_batch_matches_ragged_flow_numerics"
 
 
 def _find_step(label: str) -> dict:
@@ -35,14 +37,28 @@ def _find_step(label: str) -> dict:
 
 def _assert_single_gpu_selection(command: str) -> None:
     argv = split(command)
-    assert "tests/engine/" in argv
+    assert "tests/engine/" not in argv
     assert "tests/model_executor/" in argv
     assert "tests/worker/test_batched_omni_output.py" not in argv
     assert argv[argv.index("-m") + 1] == SINGLE_GPU_MARKERS
-    assert argv[argv.index("-k") + 1] == f"not {EXCLUDED_TEST}"
+    assert "-k" not in argv
     assert argv[argv.index("--run-level") + 1] == "core_model"
 
 
+@pytest.mark.parametrize(
+    ("marks", "selected"),
+    [
+        ({"core_model", "cuda", "nvidia_only", "L4", "cards_1"}, False),
+        ({"core_model", "cuda"}, True),
+        ({"core_model", "cuda", "L4", "MI325", "cards_1"}, True),
+        ({"core_model", "cuda", "cards_2"}, False),
+    ],
+)
+def test_r2_01_marker_selection_contract(marks: set[str], selected: bool) -> None:
+    expression = Expression.compile(SINGLE_GPU_MARKERS)
+    assert expression.evaluate(marks.__contains__) is selected
+
+
 def test_r2_01_is_nonblocking_single_gpu_coverage() -> None:
     step = _find_step(R2_01_LABEL)
 
@@ -53,6 +69,11 @@ def test_r2_01_is_nonblocking_single_gpu_coverage() -> None:
     assert step["artifact_paths"] == ["artifacts/rocm-r2-01/**/*"]
 
 
+def test_dedicated_engine_job_preserves_r2_01_engine_coverage() -> None:
+    commands = "\n".join(_find_step(ENGINE_LABEL)["commands"])
+    assert "tests/engine/test_async_omni_engine_abort.py" in commands
+
+
 def test_r2_01_runs_once_fails_closed_and_publishes_results() -> None:
     commands = _find_step(R2_01_LABEL)["commands"]
     pytest_commands = [command for command in commands if command.lstrip().startswith("pytest ")]

tests/model_executor/models/cosyvoice3/test_cosyvoice3_components_cuda.py

@@ -17,6 +17,7 @@ pytestmark = [
 ]
 
 
+@pytest.mark.nvidia_only
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
 def test_code2wav_streaming_batch_matches_ragged_flow_numerics(monkeypatch):
     """A padded flow batch must match individual flow calls on valid mels."""

tests/model_executor/models/personaplex/test_mimi_cudagraph.py

@@ -10,6 +10,7 @@ from types import SimpleNamespace
 import pytest
 import torch
 import torch.nn as nn
+from vllm.platforms import current_platform
 
 from vllm_omni.model_executor.models.personaplex.personaplex_code2wav import _MIMI_DECODE_BATCH_FRAMES
 from vllm_omni.model_executor.models.personaplex.personaplex_mimi import (
@@ -72,6 +73,21 @@ def _snapshot(codec: PersonaPlexMimiCodec) -> list[torch.Tensor]:
     return [tensor.clone() for tensor in _state_tensors(codec)]
 
 
+def _assert_waveform_matches(actual: torch.Tensor, expected: torch.Tensor) -> None:
+    if not current_platform.is_rocm():
+        assert torch.equal(actual, expected)
+        return
+
+    # ROCm convolution kernels can accumulate low-order FP32 differences
+    # between eager execution and CUDA graph replay. Keep token outputs exact,
+    # and bound both the peak and aggregate waveform error here.
+    error = (actual - expected).abs()
+    rms = error.square().mean().sqrt().item()
+    print(f"PersonaPlex Mimi ROCm graph parity: max_abs={error.max().item():.9g}, rms={rms:.9g}")
+    torch.testing.assert_close(actual, expected, atol=1e-3, rtol=1e-4)
+    assert rms <= 1e-4, rms
+
+
 @pytest.mark.cpu
 @pytest.mark.parametrize("kind", ["conv", "conv_replicate", "convtr"])
 def test_stream_carries_update_in_place(kind: str) -> None:
@@ -189,14 +205,14 @@ def test_graph_replay_is_bitwise_equal_to_eager_across_recycle() -> None:
     assert codec._cuda_graphs["decode_f1"].replays == 12
     assert graph_inactive_ok
     assert torch.equal(graph_codes, eager_codes)
-    assert torch.equal(graph_pcm, eager_pcm)
+    _assert_waveform_matches(graph_pcm, eager_pcm)
 
     # A full reset in place (a new stream on the same codec) replays the same
     # frames from a fresh state.
     codec.streaming_init(3)
     again_codes, again_pcm, _ = _drive(codec, frames=12, recycle_at=7, device=device)
     assert torch.equal(again_codes, eager_codes)
-    assert torch.equal(again_pcm, eager_pcm)
+    _assert_waveform_matches(again_pcm, eager_pcm)
     assert pointers == [tensor.data_ptr() for tensor in _state_tensors(codec)]
 
 
@@ -235,7 +251,7 @@ def test_multi_frame_decode_graph_matches_eager_chunks(batch_size: int) -> None:
     for (frames, active), expected, actual in zip(steps, eager, graphed):
         assert actual.shape == (batch_size, frames * FRAME_SIZE)
         rows = slice(None) if active is None else active
-        assert torch.equal(actual[rows], expected[rows])
+        _assert_waveform_matches(actual[rows], expected[rows])
 
 
 @pytest.mark.cuda
@@ -277,7 +293,7 @@ def test_failed_capture_warns_with_the_error_and_stays_eager(monkeypatch: pytest
     assert not kwargs
     # The warmup frames are rolled back, so the eager codec starts a fresh stream.
     assert torch.equal(codec.encode_frame(pcm), expected_codes)
-    assert torch.equal(codec.decode_frame(expected_codes), expected_pcm)
+    _assert_waveform_matches(codec.decode_frame(expected_codes), expected_pcm)
 
 
 @pytest.mark.cuda