[ROCm] [CI] Bugfix Existing CI cases (#3946)

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

.buildkite/test-amd-merge.yml

@@ -1,15 +1,28 @@
 steps:
 
-- label: "Simple Unit Test"
+- label: "Simple · Diffusion & Model Executor Test"
   agent_pool: mi325_1
   depends_on: amd-build
+  timeout_in_minutes: 50
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
     - export VLLM_ROCM_USE_AITER=0
     # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
     # that is not supported on CPU
-    - "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+
+- label: "Simple · Other Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  timeout_in_minutes: 50
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_ROCM_USE_AITER=0
+    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
+    # that is not supported on CPU
+    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
 
 - label: "Diffusion Images API LoRA E2E"
   agent_pool: mi325_1
@@ -49,16 +62,6 @@ 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: "Diffusion Cache Backend Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
-
 - label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
   agent_pool: mi325_4
   depends_on: amd-build
@@ -123,7 +126,7 @@ steps:
   grade: Blocking
   commands:
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
+    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
 
 - label: "Qwen3-TTS CustomVoice E2E Test"
   agent_pool: mi325_1

.buildkite/test-amd-ready.yaml

@@ -1,6 +1,6 @@
 steps:
 
-- label: "Simple Unit Test"
+- label: "Simple · Diffusion & Model Executor Test"
   agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
@@ -9,7 +9,18 @@ steps:
     - export VLLM_ROCM_USE_AITER=0
     # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
     # that is not supported on CPU
-    - "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+
+- label: "Simple · Other Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_ROCM_USE_AITER=0
+    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
+    # that is not supported on CPU
+    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
 
 - label: "Voxtral TTS CUDA Unit Test"
   agent_pool: mi325_1
@@ -19,13 +30,13 @@ steps:
   commands:
     - timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
 
-- label: "Diffusion Batching Test"
+- label: "Diffusion · Batch Test"
   agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - timeout 20m pytest -s -v tests/diffusion/batching/test_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
+    - timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
 
 - label: "Custom Pipeline Test"
   agent_pool: mi325_1

.buildkite/test-template-amd-omni.j2

@@ -53,6 +53,9 @@
 {% endif %}
             {{ indented_cmd | safe }}
         priority: 100
+        {% if step.timeout_in_minutes %}
+        timeout_in_minutes: {{ step.timeout_in_minutes }}
+        {% endif %}
         {% if step.grade and step.grade == "Blocking" %}
         soft_fail: false
         {% else %}

examples/offline_inference/ming_flash_omni/README.md

@@ -53,7 +53,7 @@ python examples/offline_inference/ming_flash_omni/end2end.py \
     --image-path ./3_0.png
 ```
 
-### Omni-Speech (Thinker + Talker)
+### Omni-Speech (Thinker + Talker) {#omni-speech-thinker--talker}
 
 The default deploy YAML already runs thinker+talker, so spoken output only requires requesting `audio` (or `text,audio`) modalities.
 The thinker processes your multimodal input, generates text, then the talker synthesises the response as speech.

tests/diffusion/test_omni_diffusion_config_master_port.py

@@ -11,26 +11,54 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def _get_available_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("", 0))
+        return sock.getsockname()[1]
+
+
+def _bind_available_port_below(max_port: int, port_inc: int) -> socket.socket:
+    for _ in range(100):
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        try:
+            sock.bind(("", 0))
+            port = sock.getsockname()[1]
+            if port >= max_port - port_inc:
+                sock.close()
+                continue
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as next_sock:
+                next_sock.bind(("", port + port_inc))
+            return sock
+        except OSError:
+            sock.close()
+    raise RuntimeError(f"Failed to reserve an available port below {max_port - port_inc}")
+
+
 class TestOmniDiffusionConfigMasterPort:
     def test_honors_master_port_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
-        monkeypatch.setenv("MASTER_PORT", "43210")
+        port = _get_available_port()
+        monkeypatch.setenv("MASTER_PORT", str(port))
         config = OmniDiffusionConfig(model="test")
-        assert config.master_port == 43210
+        assert config.master_port == port
 
     def test_honors_explicit_master_port(self, monkeypatch: pytest.MonkeyPatch) -> None:
         monkeypatch.delenv("MASTER_PORT", raising=False)
-        config = OmniDiffusionConfig(model="test", master_port=40000)
-        assert config.master_port == 40000
+        port = _get_available_port()
+        config = OmniDiffusionConfig(model="test", master_port=port)
+        assert config.master_port == port
 
     def test_master_port_env_takes_precedence_over_kwarg(self, monkeypatch: pytest.MonkeyPatch) -> None:
-        monkeypatch.setenv("MASTER_PORT", "43210")
-        config = OmniDiffusionConfig(model="test", master_port=40000)
-        assert config.master_port == 43210
+        env_port = _get_available_port()
+        kwarg_port = _get_available_port()
+        monkeypatch.setenv("MASTER_PORT", str(env_port))
+        config = OmniDiffusionConfig(model="test", master_port=kwarg_port)
+        assert config.master_port == env_port
 
     def test_explicit_master_port_is_stable_without_jitter(self, monkeypatch: pytest.MonkeyPatch) -> None:
         monkeypatch.delenv("MASTER_PORT", raising=False)
-        ports = {OmniDiffusionConfig(model="test", master_port=40123).master_port for _ in range(5)}
-        assert ports == {40123}
+        port = _get_available_port()
+        ports = {OmniDiffusionConfig(model="test", master_port=port).master_port for _ in range(5)}
+        assert ports == {port}
 
     def test_default_master_port_is_available(self, monkeypatch: pytest.MonkeyPatch) -> None:
         monkeypatch.delenv("MASTER_PORT", raising=False)
@@ -42,8 +70,7 @@ class TestOmniDiffusionConfigMasterPort:
 
     def test_settle_port_steps_when_port_is_busy(self, monkeypatch: pytest.MonkeyPatch) -> None:
         monkeypatch.delenv("MASTER_PORT", raising=False)
-        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-            sock.bind(("", 0))
+        with _bind_available_port_below(max_port=60000, port_inc=37) as sock:
             busy_port = sock.getsockname()[1]
             config = OmniDiffusionConfig(model="test", master_port=busy_port)
             assert config.master_port == busy_port + 37

vllm_omni/deploy/qwen3_omni_moe.yaml

@@ -7,8 +7,9 @@
 #   * Stages 0/1 (thinker, talker) do not set max_num_batched_tokens —
 #     chat-sized prefill fits in the 32768 default.
 #   * Stages 0/1 do not set enforce_eager — cudagraph runs by default
-#     (false). Stage 2 (code2wav) sets true because its generation loop
-#     is cudagraph-incompatible.
+#     (false). Stage 2 (code2wav) keeps enforce_eager=false so Code2Wav can
+#     use its inner cudagraph on CUDA. ROCm flips stage 2 to eager because
+#     MIOpen conv_transpose1d is not capture-safe in the Code2Wav graph path.
 #   * Platform sections flip enforce_eager per-stage where platform
 #     cudagraph support differs.
 async_chunk: true
@@ -88,6 +89,15 @@ platforms:
         gpu_memory_utilization: 0.3
         devices: "2"
 
+  # Code2Wav decode uses UpSample1d, which calls F.conv_transpose1d.
+  # On ROCm this routes through MIOpen and is not capture-safe during
+  # vLLM's outer graph capture, so stage 2 runs eager to avoid MIOpen/HIP
+  # graph-capture failures.
+  rocm:
+    stages:
+      - stage_id: 2
+        enforce_eager: true
+
   xpu:
     stages:
       - stage_id: 0

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -523,10 +523,11 @@ class DiffusionWorker:
 
             current_omni_platform.synchronize()
             free_before = current_omni_platform.get_free_memory(self.device)
-            self.sleep(level=task.level)
+            allocator_freed = self.sleep(level=task.level)
             current_omni_platform.synchronize()
             free_after = current_omni_platform.get_free_memory(self.device)
-            real_freed = max(0, free_after - free_before)
+            phys_freed = max(0, free_after - free_before)
+            real_freed = max(int(allocator_freed), phys_freed)
             logger.info(f"[Worker {self.rank}] Preparing ACK: freed_bytes={real_freed / GiB_bytes:.2f} GiB.")
 
             # Ensure all ranks have completed sleep before measuring memory and sending ACK
@@ -553,6 +554,8 @@ class DiffusionWorker:
                 metadata={
                     "source": f"Platform_{current_omni_platform.get_device_name()}",
                     "total_freed_gib": f"{real_freed / GiB_bytes:.2f}",
+                    "allocator_freed_gib": f"{allocator_freed / GiB_bytes:.2f}",
+                    "physical_freed_gib": f"{phys_freed / GiB_bytes:.2f}",
                     "rank_residual_gib": f"{residual_gib:.2f}",
                 },
             )