117 lines
5.4 KiB
Python
117 lines
5.4 KiB
Python
# 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)
|