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

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

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