diff --git a/.buildkite/amd/test-amd-merge.yml b/.buildkite/amd/test-amd-merge.yml index 1e3e852e9..e702e4d9c 100644 --- a/.buildkite/amd/test-amd-merge.yml +++ b/.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 diff --git a/.buildkite/amd/test-amd-ready.yml b/.buildkite/amd/test-amd-ready.yml index ae64e7c4b..3491d59c8 100644 --- a/.buildkite/amd/test-amd-ready.yml +++ b/.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" diff --git a/tests/diffusion/models/auk/test_auk_vae_fastpath.py b/tests/diffusion/models/auk/test_auk_vae_fastpath.py index b25381097..8a1b39aa1 100644 --- a/tests/diffusion/models/auk/test_auk_vae_fastpath.py +++ b/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: diff --git a/tests/model_executor/models/test_breeze_tts_2_first_code_sampler.py b/tests/model_executor/models/test_breeze_tts_2_first_code_sampler.py index 660737765..9d9047df1 100644 --- a/tests/model_executor/models/test_breeze_tts_2_first_code_sampler.py +++ b/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") diff --git a/tests/model_executor/models/test_breeze_tts_2_graphs.py b/tests/model_executor/models/test_breeze_tts_2_graphs.py index 72c2cfeea..b60c8bf46 100644 --- a/tests/model_executor/models/test_breeze_tts_2_graphs.py +++ b/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 diff --git a/vllm_omni/model_executor/models/breeze_tts_2/depth_decoder.py b/vllm_omni/model_executor/models/breeze_tts_2/depth_decoder.py index 70b7262f1..0adc38d9b 100644 --- a/vllm_omni/model_executor/models/breeze_tts_2/depth_decoder.py +++ b/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(