fix: repair ROCm test routing, CPU timeouts, and Breeze autotuning (#8342)

Signed-off-by: andyluo7 <andy.luo@amd.com>

.buildkite/amd/test-amd-merge.yml

@@ -41,7 +41,33 @@ 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
-        - "pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
+        # Keep multi-GPU tests out of this one-GPU pool even when their module
+        # also carries the cpu marker for unit tests.
+        - "pytest -sv tests/diffusion -m 'core_model and cpu and not (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)' --ignore=tests/diffusion/cache/test_teacache_extractors.py --ignore=tests/diffusion/models/pi05/test_pi05_units.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
+    # This module takes about 22 minutes on MI300 and was consistently assigned to shard 4,
+    # pushing that shard to the 45-minute timeout. Keep its module-scoped fixture in one lane.
+    - label: "Simple · Pi0.5 CPU Test"
+      agent_pool: mi300_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      timeout_in_minutes: 45
+      commands:
+        - export VLLM_ROCM_USE_AITER=0
+        - "pytest -sv tests/diffusion/models/pi05/test_pi05_units.py -m 'core_model and cpu'"
+
+    - label: "Diffusion · AuK VAE GPU Test"
+      agent_pool: mi300_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      timeout_in_minutes: 40
+      commands:
+        - export VLLM_ROCM_USE_AITER=0
+        - >-
+          timeout --signal=TERM --kill-after=2m 35m
+          pytest -sv tests/diffusion/models/auk/test_auk_vae_fastpath.py
+          -m 'core_model and rocm and cards_1' --run-level core_model
 
     - label: "Simple · Engine&Entrypoints Test"
       agent_pool: mi300_1
@@ -131,6 +157,11 @@ steps:
       grade: Blocking
       commands:
         - pytest -s -v tests/diffusion/distributed/test_tensor_parallel.py
+        - >-
+          timeout 20m pytest -s -v
+          tests/diffusion/attention/test_ulysses_uaa.py::test_ulysses_uaa_2d_mask_layout_matches_baseline
+          -m 'core_model and cards_2'
+          --run-level "core_model"
 
     - label: "Diffusion Model Test"
       agent_pool: mi300_1

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

@@ -41,7 +41,33 @@ 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
-        - "pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
+        # Keep multi-GPU tests out of this one-GPU pool even when their module
+        # also carries the cpu marker for unit tests.
+        - "pytest -sv tests/diffusion -m 'core_model and cpu and not (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)' --ignore=tests/diffusion/cache/test_teacache_extractors.py --ignore=tests/diffusion/models/pi05/test_pi05_units.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
+    # This module takes about 22 minutes on MI300 and was consistently assigned to shard 4,
+    # pushing that shard to the 45-minute timeout. Keep its module-scoped fixture in one lane.
+    - label: "Simple · Pi0.5 CPU Test"
+      agent_pool: mi300_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      timeout_in_minutes: 45
+      commands:
+        - export VLLM_ROCM_USE_AITER=0
+        - "pytest -sv tests/diffusion/models/pi05/test_pi05_units.py -m 'core_model and cpu'"
+
+    - label: "Diffusion · AuK VAE GPU Test"
+      agent_pool: mi300_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      timeout_in_minutes: 40
+      commands:
+        - export VLLM_ROCM_USE_AITER=0
+        - >-
+          timeout --signal=TERM --kill-after=2m 35m
+          pytest -sv tests/diffusion/models/auk/test_auk_vae_fastpath.py
+          -m 'core_model and rocm and cards_1' --run-level core_model
 
     - label: "Simple · Engine&Entrypoints Test"
       agent_pool: mi300_1
@@ -141,6 +167,11 @@ steps:
         - export VLLM_LOGGING_LEVEL=DEBUG
         - export VLLM_WORKER_MULTIPROC_METHOD=spawn
         - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
+        - >-
+          timeout 20m pytest -s -v
+          tests/diffusion/attention/test_ulysses_uaa.py::test_ulysses_uaa_2d_mask_layout_matches_baseline
+          -m 'core_model and cards_2'
+          --run-level "core_model"
         - >-
           timeout 20m pytest -s -v
           tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py
@@ -165,6 +196,7 @@ steps:
         # registered multi-GPU marker from this one-GPU job.
         - >-
           timeout 15m pytest -sv tests/diffusion/models/
+          --ignore=tests/diffusion/models/auk/test_auk_vae_fastpath.py
           -m '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)'
           --run-level "core_model"

tests/diffusion/models/auk/test_auk_vae_fastpath.py

@@ -11,10 +11,28 @@ eager off CUDA) and the tiled decode of long clips.
 import pytest
 import torch
 
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE, LowPass, SnakeBeta, Upsample
 from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph, plan_tiles
+from vllm_omni.platforms import current_omni_platform
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model]
+
+
+@pytest.fixture(autouse=True)
+def _bound_cpu_decode_threads(request):
+    if request.node.get_closest_marker("cpu") is None:
+        yield
+        return
+    # These narrow convolutions run in shared CI pods. A host-sized Torch
+    # pool spends more time coordinating workers than decoding small clips.
+    previous_threads = torch.get_num_threads()
+    torch.set_num_threads(1)
+    print(f"AuK CPU decode threads: {previous_threads} -> {torch.get_num_threads()}")
+    try:
+        yield
+    finally:
+        torch.set_num_threads(previous_threads)
 
 
 def _small_vae() -> AuKVAE:
@@ -45,6 +63,23 @@ def _reference_snake(module: SnakeBeta, x: torch.Tensor) -> torch.Tensor:
     return x + (1.0 / (beta + 1e-9)) * torch.sin(x * alpha).pow(2)
 
 
+def _assert_plain_graph_matches_eager(replay: torch.Tensor, eager: torch.Tensor) -> None:
+    assert replay.shape == eager.shape and replay.dtype == eager.dtype
+    assert torch.isfinite(replay).all() and torch.isfinite(eager).all()
+    error = (replay - eager).abs()
+    rms = error.square().mean().sqrt().item()
+    print(f"AuK plain graph: max_abs={error.max().item():.9g}, rms={rms:.9g}")
+    if current_omni_platform.is_rocm():
+        # MI300 build 13040 measured max_abs <= 1.64e-7 and RMS <= 4.13e-8.
+        # Bound both peak and aggregate FP32 error; near-zero samples must
+        # not hide behind a relative tolerance. NVIDIA retains bit parity.
+        torch.testing.assert_close(replay, eager, atol=1e-6, rtol=0.0)
+        assert rms <= 1e-7, rms
+    else:
+        assert torch.equal(replay, eager)
+
+
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_window_padding_decodes_to_raw_zero() -> None:
     """Bucket padding is the normalized latent whose denormalized value is zero."""
@@ -57,6 +92,7 @@ def test_window_padding_decodes_to_raw_zero() -> None:
     torch.testing.assert_close(raw, torch.zeros_like(raw), atol=1e-6, rtol=0.0)
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_eager_snake_matches_the_reference_formula() -> None:
     module = SnakeBeta(6, alpha_logscale=True)
@@ -69,6 +105,7 @@ def test_eager_snake_matches_the_reference_formula() -> None:
     assert torch.equal(module(x), _reference_snake(module, x))
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_filter_cache_leaves_the_filters_and_output_unchanged() -> None:
     lowpass = LowPass(cutoff=0.25, half_width=0.3, stride=2, kernel_size=12, causal=True)
@@ -88,6 +125,7 @@ def test_filter_cache_leaves_the_filters_and_output_unchanged() -> None:
         assert module._expanded.shape == (3, 1, 12)
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_decode_fast_paths_reproduce_the_plain_decode() -> None:
     vae = _small_vae()
@@ -103,6 +141,7 @@ def test_decode_fast_paths_reproduce_the_plain_decode() -> None:
     assert all(module._cached for module in vae.modules() if isinstance(module, SnakeBeta))
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_graph_wrapper_falls_back_to_eager_off_cuda(mocker) -> None:
     vae = _small_vae()
@@ -119,6 +158,7 @@ def test_graph_wrapper_falls_back_to_eager_off_cuda(mocker) -> None:
     assert not wrapper._cache and not wrapper._compiled
 
 
+@pytest.mark.cpu
 def test_compiled_bucket_is_the_smallest_captured_one_that_fits() -> None:
     wrapper = AuKVAEDecodeGraph(_small_vae(), compile_shapes=(16, 8, 8))
     assert wrapper.compile_shapes == [8, 16]
@@ -131,6 +171,7 @@ def test_compiled_bucket_is_the_smallest_captured_one_that_fits() -> None:
     assert wrapper.compiled_bucket(6) == 16
 
 
+@pytest.mark.cpu
 def test_plan_tiles_covers_the_clip_once_with_full_context() -> None:
     assert plan_tiles(40, 64, 10, 3) == [(0, 40, 0, 40)]
     windows = plan_tiles(1000, 512, 53, 9)
@@ -149,10 +190,16 @@ def test_plan_tiles_covers_the_clip_once_with_full_context() -> None:
     # 512 and a 256 window instead of two 512s.
     assert plan_tiles(600, 512, 53, 9, sizes=(128, 256, 512)) == [(0, 512, 0, 503), (344, 256, 503, 600)]
     assert plan_tiles(1000, 512, 53, 9, sizes=(128, 256, 512))[-1] == (872, 128, 953, 1000)
+    narrow = plan_tiles(200, 64, 53, 9)
+    assert len(narrow) == 69
+    assert [emit for _, _, emit, _ in narrow[:3]] == [0, 55, 57]
+    assert all(a[3] == b[2] for a, b in zip(narrow, narrow[1:]))
+    assert narrow[-1][3] == 200
     with pytest.raises(ValueError, match="must exceed"):
         plan_tiles(100, 60, 53, 9)
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_decode_context_bounds_the_measured_receptive_field() -> None:
     vae = _small_vae()
@@ -170,25 +217,26 @@ def test_decode_context_bounds_the_measured_receptive_field() -> None:
     assert differing and all(index < left or end - start - index <= right for index in differing)
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_tiled_decode_matches_the_whole_decode() -> None:
     vae = _small_vae()
     # Off the accelerator every window decodes eagerly, so this isolates the stitching.
-    wrapper = AuKVAEDecodeGraph(vae, compile_shapes=(64,), tile_frames=64)
-    assert wrapper.tile_frames == 64 and wrapper.context_frames == (53, 9)
-    for frames in (64, 65, 200, 331):
+    wrapper = AuKVAEDecodeGraph(vae, compile_shapes=(128,), tile_frames=128)
+    assert wrapper.tile_frames == 128 and wrapper.context_frames == (53, 9)
+    for frames in (128, 129, 200, 331):
         latents = torch.randn(1, frames, vae.latent_dim)
         whole = vae.decode(latents)
         tiled = wrapper(latents)
         assert tiled.shape == whole.shape
         torch.testing.assert_close(tiled, whole, atol=1e-5, rtol=0.0)
-        assert wrapper.last_mode == ("eager" if frames == 64 else "tiled")
+        assert wrapper.last_mode == ("eager" if frames == 128 else "tiled")
     # The streaming form yields the same audio in order.
     latents = torch.randn(1, 200, vae.latent_dim)
     pieces = [(start, chunk.clone()) for start, chunk in wrapper.decode_tiles(latents)]
-    # With a 64-frame tile and a 62-frame halo each interior tile adds two frames.
-    assert [start for start, _ in pieces] == [emit for _, _, emit, _ in plan_tiles(200, 64, 53, 9)]
-    assert [start for start, _ in pieces][:3] == [0, 55, 57]
+    # A 128-frame tile avoids repeatedly decoding the same 62-frame halo.
+    assert [start for start, _ in pieces] == [emit for _, _, emit, _ in plan_tiles(200, 128, 53, 9)]
+    assert [start for start, _ in pieces][:3] == [0, 119, 185]
     torch.testing.assert_close(
         torch.cat([chunk for _, chunk in pieces], dim=1), vae.decode(latents), atol=1e-5, rtol=0.0
     )
@@ -198,20 +246,24 @@ def test_tiled_decode_matches_the_whole_decode() -> None:
         AuKVAEDecodeGraph(vae, tile_frames=60)
 
 
-@pytest.mark.cuda
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
 @torch.inference_mode()
 def test_graph_replay_matches_eager_per_length() -> None:
     vae = _small_vae().to("cuda")
     wrapper = AuKVAEDecodeGraph(vae, max_graphs=2)
     pools = []
+    saved_outputs: list[tuple[torch.Tensor, torch.Tensor]] = []
     for frames in (6, 9, 6, 12, 6):
         latents = torch.randn(1, frames, vae.latent_dim, device="cuda")
         eager = vae.decode(latents)
         replay = wrapper(latents)
-        # Exact-length graphs replay the very same kernels: bit-identical,
-        # including after the generation was retired.
-        assert torch.equal(replay, eager), frames
+        _assert_plain_graph_matches_eager(replay, eager)
+        # Returned waveforms own their storage and survive later replays
+        # and retirement of the generation that produced them.
+        for previous, snapshot in saved_outputs:
+            assert torch.equal(previous, snapshot)
+        saved_outputs.append((replay, replay.clone()))
         pools.append(wrapper._plain_pool)
     # The third distinct length found the cache full, so the whole generation
     # (6 and 9) was retired together with its pool rather than one graph at a
@@ -220,7 +272,7 @@ def test_graph_replay_matches_eager_per_length() -> None:
     assert pools[0] is pools[2] and pools[3] is pools[4] and pools[2] is not pools[3]
 
 
-@pytest.mark.cuda
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
 @torch.inference_mode()
 def test_bucketed_graph_only_disturbs_the_tail() -> None:
@@ -240,7 +292,7 @@ def test_bucketed_graph_only_disturbs_the_tail() -> None:
     torch.testing.assert_close(replay, eager, atol=0.1, rtol=0.0)
 
 
-@pytest.mark.cuda
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="torch.compile + CUDA graph capture requires CUDA")
 @torch.inference_mode()
 def test_compiled_buckets_replay_within_fusion_tolerance_and_leave_longer_clips_to_plain_graphs() -> None:
@@ -266,11 +318,12 @@ def test_compiled_buckets_replay_within_fusion_tolerance_and_leave_longer_clips_
     # a longer clip gets its own plain graph.
     assert wrapper.tile_frames == 0
     long = torch.randn(1, 12, vae.latent_dim, device="cuda")
-    assert torch.equal(wrapper(long), vae.decode(long))
+    replay, eager = wrapper(long), vae.decode(long)
+    _assert_plain_graph_matches_eager(replay, eager)
     assert wrapper.last_mode == "graph" and list(wrapper._cache) == [12]
 
 
-@pytest.mark.cuda
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="torch.compile + CUDA graph capture requires CUDA")
 @torch.inference_mode()
 def test_tiles_replay_the_compiled_bucket_for_long_clips() -> None:

tests/model_executor/models/test_breeze_tts_2_first_code_sampler.py

@@ -4,6 +4,7 @@
 
 import pytest
 import torch
+from vllm.platforms import current_platform
 
 from tests.helpers.mark import hardware_test
 from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import sample_logits
@@ -16,6 +17,10 @@ VOCAB = 2052
 EOS = 2051
 
 
+def _captured_generator_state_supported() -> bool:
+    return current_platform.is_cuda() and hasattr(torch.cuda.CUDAGraph, "register_generator_state")
+
+
 def _logits(count: int, *, device: str, seed: int = 100, vocab: int = VOCAB) -> list[torch.Tensor]:
     source = torch.Generator(device="cpu").manual_seed(seed)
     values = torch.randn((count, vocab), generator=source, dtype=torch.float32) * 3
@@ -164,6 +169,14 @@ def test_cuda_replay_preserves_tokens_rng_and_owned_outputs(monkeypatch) -> None
                     right[codebook].exponential_(generator=live_expected[index])
                 torch.testing.assert_close(left, right, atol=0, rtol=0)
             _assert_states(actual_generators, expected_generators)
+    if not _captured_generator_state_supported():
+        # ROCm deliberately uses the eager sampler because CUDA graph generator-state registration is an
+        # NVIDIA-only path. The loop above still verifies GPU tokens and every request-owned RNG stream.
+        assert not sampler._graphs
+        for value, expected in retained:
+            torch.testing.assert_close(value, expected, atol=0, rtol=0)
+        torch.testing.assert_close(torch.cuda.get_rng_state(), global_state, atol=0, rtol=0)
+        return
     # B1/3/16/17/31/32 must use exactly buckets 1/4/16/32.
     assert len(sampler._graphs) == 4
     rows = _logits(3, device="cuda", seed=2026)
@@ -189,6 +202,7 @@ def test_cuda_replay_preserves_tokens_rng_and_owned_outputs(monkeypatch) -> None
 @torch.inference_mode()
 def test_cuda_parameters_vocab_and_missing_graph_api(monkeypatch) -> None:
     sampler = BreezeFirstCodeSampler()
+    captures_generator_state = _captured_generator_state_supported()
     actual_generators = _generators([42, 43, 44], "cuda")
     expected_generators = _generators([42, 43, 44], "cuda")
     cases = ((0.9, 50, 1.0), (0.7, 100, 0.8), (1.2, 0, 0.95), (1.0, 1, 1.0), (1.0, 0, 0.5))
@@ -201,12 +215,12 @@ def test_cuda_parameters_vocab_and_missing_graph_api(monkeypatch) -> None:
         expected = _reference(rows, parameters, expected_generators)
         _assert_outputs(sampler.sample(rows, parameters, actual_generators), expected, rows)
         _assert_states(actual_generators, expected_generators)
-    assert len(sampler._graphs) == len(cases)
+    assert len(sampler._graphs) == (len(cases) if captures_generator_state else 0)
     rows = _logits(3, device="cuda", vocab=17)
     expected = _reference(rows, PARAMETERS, expected_generators)
     _assert_outputs(sampler.sample(rows, PARAMETERS, actual_generators), expected, rows)
     _assert_states(actual_generators, expected_generators)
-    assert len(sampler._graphs) == len(cases) + 1
+    assert len(sampler._graphs) == (len(cases) + 1 if captures_generator_state else 0)
 
     class UnavailableGraph:
         def __init__(self, *_args, **_kwargs):
@@ -243,7 +257,7 @@ def test_cuda_shared_generator_uses_sequential_eager_fallback(monkeypatch) -> No
     # The duplicate-generator guard must precede lookup of an existing graph.
     sampler.sample(rows, PARAMETERS, _generators([100, 101, 102], "cuda"))
     cached = dict(sampler._graphs)
-    assert len(cached) == 1
+    assert len(cached) == (1 if _captured_generator_state_supported() else 0)
     actual = torch.Generator(device="cuda").manual_seed(2**64 - 23)
     reference = torch.Generator(device="cuda").manual_seed(2**64 - 23)
     multinomial = torch.multinomial
@@ -270,6 +284,7 @@ def test_cuda_shared_generator_uses_sequential_eager_fallback(monkeypatch) -> No
 @torch.inference_mode()
 def test_cuda_cache_limit_falls_back_without_evicting_warmed_graphs(monkeypatch) -> None:
     sampler = BreezeFirstCodeSampler()
+    captures_generator_state = _captured_generator_state_supported()
     rows = _logits(1, device="cuda")
     actual_generators = _generators([42], "cuda")
     expected_generators = _generators([42], "cuda")
@@ -279,7 +294,7 @@ def test_cuda_cache_limit_falls_back_without_evicting_warmed_graphs(monkeypatch)
         _assert_outputs(sampler.sample(rows, parameters, actual_generators), expected, rows)
         _assert_states(actual_generators, expected_generators)
     cached = dict(sampler._graphs)
-    assert len(cached) == 32
+    assert len(cached) == (32 if captures_generator_state else 0)
     parameters = (0.9, 33, 1.0)
     expected = _reference(rows, parameters, expected_generators)
     multinomial = torch.multinomial
@@ -298,6 +313,11 @@ def test_cuda_cache_limit_falls_back_without_evicting_warmed_graphs(monkeypatch)
     assert all(sampler._graphs[key] is value for key, value in cached.items())
     expected = _reference(rows, (0.9, 1, 1.0), expected_generators)
 
+    if not captures_generator_state:
+        _assert_outputs(sampler.sample(rows, (0.9, 1, 1.0), actual_generators), expected, rows)
+        _assert_states(actual_generators, expected_generators)
+        return
+
     def unexpected_multinomial(*_args, **_kwargs):
         raise AssertionError("Cache exhaustion must retain existing captured entries")
 

tests/model_executor/models/test_breeze_tts_2_graphs.py

@@ -9,7 +9,11 @@ from transformers.models.t5gemma2.modeling_t5gemma2 import T5Gemma2TextEncoder
 from vllm.platforms import current_platform
 
 from tests.helpers.mark import hardware_test
-from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import BreezeDepthDecoder, sample_logits
+from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import (
+    BreezeDepthDecoder,
+    _depth_layer_compile_options,
+    sample_logits,
+)
 from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze import BreezeForConditionalGeneration
 from vllm_omni.model_executor.models.breeze_tts_2.text_encoder_graph import (
     BreezeTextEncoderCompiled,
@@ -19,6 +23,17 @@ from vllm_omni.model_executor.models.breeze_tts_2.text_encoder_graph import (
 pytestmark = [pytest.mark.core_model]
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize(("is_rocm", "expected"), [(True, False), (False, True)])
+def test_depth_layer_compile_options(monkeypatch, is_rocm: bool, expected: bool) -> None:
+    monkeypatch.setattr(current_platform, "is_rocm", lambda: is_rocm)
+
+    assert _depth_layer_compile_options() == {
+        "epilogue_fusion": False,
+        "max_autotune": expected,
+    }
+
+
 @pytest.fixture
 def full_precision_matmul():
     # Different padded shapes can select different TF32 kernels. Compare
@@ -241,6 +256,7 @@ def test_depth_graph_buckets_bound_shapes_and_preserve_requests(full_precision_m
     seeds[0], seeds[-1] = 2**63 + 17, 2**64 - 1
     generators = [torch.Generator(device="cuda").manual_seed(seed) for seed in seeds]
     reference_generators = [torch.Generator(device="cuda").manual_seed(seed) for seed in seeds]
+    captures_generator_state = current_platform.is_cuda() and hasattr(torch.cuda.CUDAGraph, "register_generator_state")
 
     expected_greedy = torch.cat(
         [
@@ -314,7 +330,7 @@ def test_depth_graph_buckets_bound_shapes_and_preserve_requests(full_precision_m
             branches = 2 if guidance_scale != 1 else 1
             graph_key = (request_bucket * branches, branches, temperature == 0)
             with monkeypatch.context() as rng_guard:
-                if temperature > 0 and graph_key in depth._graphs:
+                if temperature > 0 and graph_key in depth._graphs and captures_generator_state:
                     # A warmed replay must generate noise within the CUDA
                     # graph, without dispatching new RNG kernels from Python.
                     rng_guard.setattr(torch.Tensor, "exponential_", reject_eager_noise)
@@ -349,7 +365,7 @@ def test_depth_graph_buckets_bound_shapes_and_preserve_requests(full_precision_m
     expected_keys.update((batch, 2, greedy) for batch in (2, 8, 16, 32) for greedy in (False, True))
     assert set(depth._graphs) == expected_keys
     for (batch_bucket, branches, greedy), entry in depth._graphs.items():
-        if greedy:
+        if greedy or not captures_generator_state:
             assert entry.noise_generators is None
         else:
             assert entry.noise_generators is not None

vllm_omni/model_executor/models/breeze_tts_2/depth_decoder.py

@@ -16,6 +16,14 @@ from vllm.platforms import current_platform
 from vllm_omni.platforms import current_omni_platform
 
 
+def _depth_layer_compile_options() -> dict[str, bool]:
+    """Keep ROCm away from the unstable Triton GEMM autotune path."""
+    return {
+        "epilogue_fusion": False,
+        "max_autotune": not current_platform.is_rocm(),
+    }
+
+
 def sample_logits(
     logits: torch.Tensor,
     temperature: float,
@@ -283,7 +291,7 @@ class BreezeDepthDecoder(nn.Module):
                     BreezeDepthLayer.forward,
                     fullgraph=True,
                     dynamic=True,
-                    options={"epilogue_fusion": False, "max_autotune": True},
+                    options=_depth_layer_compile_options(),
                 )
                 self._compiled_sampler = torch.compile(sample_graph_logits, fullgraph=True, dynamic=False)
             entry = BreezeDepthGraph(