diff --git a/tests/model_executor/test_seeded_exponential_platform.py b/tests/model_executor/test_seeded_exponential_platform.py new file mode 100644 index 000000000..88e728123 --- /dev/null +++ b/tests/model_executor/test_seeded_exponential_platform.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from unittest.mock import Mock + +import pytest +import torch + +from tests.helpers.mark import hardware_test +from vllm_omni.utils import seeded_exponential + +pytestmark = [pytest.mark.core_model] + + +@pytest.mark.cpu +@pytest.mark.parametrize("is_nvidia", [False, True], ids=["rocm", "cuda"]) +def test_cuda_tensor_requires_nvidia_distribution_contract(monkeypatch, is_nvidia: bool) -> None: + monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: is_nvidia) + # HIP tensors expose the same Torch device predicate as NVIDIA tensors. + q = Mock(spec=torch.Tensor) + q.is_cuda = True + q.dtype = torch.float32 + q.dim.return_value = 2 + q.is_contiguous.return_value = True + generators = {0: object(), 1: object()} + + assert seeded_exponential.batched_seeded_exponential_supported(q, generators) is is_nvidia + + +@pytest.mark.cpu +def test_shared_generator_keeps_ordered_torch_draws(monkeypatch) -> None: + monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: True) + q = Mock(spec=torch.Tensor) + q.is_cuda = True + q.dtype = torch.float32 + q.dim.return_value = 2 + q.is_contiguous.return_value = True + shared = object() + + assert not seeded_exponential.batched_seeded_exponential_supported(q, {0: shared, 1: shared}) + + +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) +@pytest.mark.parametrize("seeded_rows", [(0, 1, 2), (0, 2)], ids=["all_seeded", "mixed"]) +@pytest.mark.parametrize("use_fp64_gumbel", [False, True]) +def test_torch_fallback_preserves_samples_and_generator_state(monkeypatch, seeded_rows, use_fp64_gumbel) -> None: + from vllm.v1.sample.ops import topk_topp_sampler as sampler_ops + + import vllm_omni.patch # noqa: F401 + + original = sampler_ops.random_sample.__wrapped__ + monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: False) + + def reject_kernel(*args, **kwargs): + raise AssertionError("the CUDA distribution kernel must not run on the fallback path") + + monkeypatch.setattr(seeded_exponential, "fill_exponential_rows", reject_kernel) + probs = torch.arange(1, 258, device="cuda", dtype=torch.float32).expand(3, -1).contiguous() + probs /= probs.sum(dim=-1, keepdim=True) + expected_gens = {row: torch.Generator(device="cuda").manual_seed(42 + row) for row in seeded_rows} + actual_gens = {row: torch.Generator(device="cuda").manual_seed(42 + row) for row in seeded_rows} + expected_default = torch.cuda.get_rng_state() + actual_default = expected_default.clone() + try: + for _ in range(5): + torch.cuda.set_rng_state(expected_default) + # FP32 sampling divides probabilities in place. Give both paths + # the same untouched input on every step. + expected = original(probs.clone(), expected_gens, use_fp64_gumbel) + expected_default = torch.cuda.get_rng_state() + torch.cuda.set_rng_state(actual_default) + actual = sampler_ops.random_sample(probs.clone(), actual_gens, use_fp64_gumbel) + actual_default = torch.cuda.get_rng_state() + assert torch.equal(actual, expected) + assert torch.equal(actual_default, expected_default) + for row in seeded_rows: + assert torch.equal(actual_gens[row].get_state(), expected_gens[row].get_state()) + finally: + torch.cuda.set_rng_state(expected_default) + + +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) +@pytest.mark.parametrize("rows", [None, [0, 2, 4]], ids=["grouped_codebooks", "selected_rows"]) +def test_direct_fill_fallback_preserves_ordered_draws(monkeypatch, rows) -> None: + monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: False) + + def reject_kernel(*args, **kwargs): + raise AssertionError("the NVIDIA kernel must not run for direct fallback callers") + + monkeypatch.setattr(seeded_exponential, "torch_exponential_policy", reject_kernel) + expected = torch.ones(6, 257, device="cuda") + actual = expected.clone() + expected_gen = torch.Generator(device="cuda").manual_seed(71) + actual_gen = torch.Generator(device="cuda").manual_seed(71) + # A shared seeded generator and an unseeded request make loop order and + # default-state advancement observable, including grouped codebooks. + expected_gens = [expected_gen, None, expected_gen] + actual_gens = [actual_gen, None, actual_gen] + original_default = torch.cuda.get_rng_state() + expected_default = original_default.clone() + actual_default = original_default.clone() + try: + for _ in range(5): + torch.cuda.set_rng_state(expected_default) + targets = expected.reshape(3, -1) if rows is None else [expected[row] for row in rows] + for target, generator in zip(targets, expected_gens): + target.exponential_(generator=generator) + expected_default = torch.cuda.get_rng_state() + torch.cuda.set_rng_state(actual_default) + assert seeded_exponential.fill_exponential_rows(actual, actual_gens, rows) is actual + actual_default = torch.cuda.get_rng_state() + assert torch.equal(actual, expected) + assert torch.equal(actual_gen.get_state(), expected_gen.get_state()) + assert torch.equal(actual_default, expected_default) + finally: + torch.cuda.set_rng_state(original_default) diff --git a/vllm_omni/utils/seeded_exponential.py b/vllm_omni/utils/seeded_exponential.py index 390c75747..027af7468 100644 --- a/vllm_omni/utils/seeded_exponential.py +++ b/vllm_omni/utils/seeded_exponential.py @@ -12,6 +12,7 @@ every value and every generator state afterwards matches the per-request loop. """ import torch +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -73,10 +74,18 @@ def torch_exponential_policy(numel: int, device: torch.device) -> tuple[int, int def fill_exponential_rows(out: torch.Tensor, generators: list, rows: list[int] | None = None) -> torch.Tensor: """Draw ``out[rows[i]].exponential_(generator=generators[i])`` for all ``i`` in one launch. - ``out`` is a contiguous float32 CUDA matrix; ``rows`` defaults to - ``range(len(generators))``. ``None`` generators use the default CUDA - generator in order. Each generator is advanced exactly as torch would. + ``out`` is a contiguous float32 matrix. Without ``rows``, each generator + owns an equal contiguous slice of the flattened output (which may span + several codebook rows). ``None`` generators use the default generator in + order. NVIDIA uses the batched CUDA kernel; other platforms retain + ordered Torch draws, including direct Higgs Audio callers. """ + if not current_platform.is_cuda(): + targets = out.reshape(len(generators), -1) if rows is None else [out[row] for row in rows] + for target, generator in zip(targets, generators): + target.exponential_(generator=generator) + return out + numel = int(out.shape[-1]) if rows is not None else int(out.numel()) // len(generators) device = out.device threads, increment = torch_exponential_policy(numel, device) @@ -109,9 +118,15 @@ def fill_exponential_rows(out: torch.Tensor, generators: list, rows: list[int] | def batched_seeded_exponential_supported(q: torch.Tensor, generators: dict) -> bool: - """Whether ``fill_exponential_rows`` reproduces the per-row loop for ``q``.""" + """Whether the NVIDIA CUDA kernel reproduces the per-row loop for ``q``. + + ROCm tensors also report ``is_cuda``, but HIP uses a different distribution + implementation and cannot compile the CUDA libdevice transform above. + Keep those draws on Torch's original per-row path, including RNG state. + """ return ( - len(generators) > 1 + current_platform.is_cuda() + and len(generators) > 1 and q.is_cuda and q.dtype == torch.float32 and q.dim() == 2