fix: preserve Torch seeded sampling on ROCm (#8341)

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

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)

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