91 lines
4.3 KiB
Python
91 lines
4.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
pytestmark = [
|
|
pytest.mark.core_model,
|
|
pytest.mark.cuda,
|
|
pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"),
|
|
]
|
|
|
|
|
|
def _reference(logits, uniforms, top_k, top_p):
|
|
if top_k > 0:
|
|
threshold = logits.topk(top_k, dim=-1).values[:, -1:]
|
|
logits = logits.masked_fill(logits < threshold, -torch.inf)
|
|
if top_p < 1:
|
|
values, indices = logits.sort(dim=-1, descending=True)
|
|
probs = values.softmax(dim=-1, dtype=torch.float32)
|
|
values[(probs.cumsum(-1) - probs) >= top_p] = -torch.inf
|
|
logits = values.scatter(1, indices, values)
|
|
return (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
|
|
|
|
|
|
@pytest.mark.parametrize("batch,vocab", [(1, 128), (7, 2048), (64, 2048)])
|
|
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
|
@pytest.mark.parametrize("top_k,top_p", [(-1, 0.8), (0, 1.0), (50, 0.8), (50, 0.01)])
|
|
def test_fixed_uniform_sampling_and_graph(batch, vocab, dtype, top_k, top_p):
|
|
from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
|
|
|
|
torch.manual_seed(173)
|
|
logits = torch.randn(batch, vocab, device="cuda", dtype=dtype)
|
|
# The MTP caller supplies a noncontiguous row slice of [B, 15, vocab].
|
|
uniforms = torch.rand(batch, 15, vocab, device="cuda")[:, 3].clamp_min_(1e-7)
|
|
expected = _reference(logits, uniforms, top_k, top_p)
|
|
actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
|
|
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
|
graph = torch.cuda.CUDAGraph()
|
|
with torch.cuda.graph(graph):
|
|
captured = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
|
|
logits.add_(torch.randn_like(logits) * 0.1)
|
|
uniforms.uniform_(1e-7, 1 - 1e-7)
|
|
graph.replay()
|
|
torch.testing.assert_close(captured, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.parametrize("top_k,top_p", [(3, 0.5), (3, 1.0), (0, 0.8)])
|
|
def test_threshold_ties_and_degenerate_rows(top_k, top_p):
|
|
from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
|
|
|
|
logits = torch.tensor(
|
|
[[2, 2, 2, 2, 2, 1, 0, -1], [1] * 8, [-torch.inf] * 8, [torch.nan, 1, 2, 3, 4, 5, 6, 7]],
|
|
device="cuda",
|
|
dtype=torch.float32,
|
|
)
|
|
uniforms = torch.tensor([[0.4, 0.7, 0.8, 0.9, 0.2, 0.3, 0.6, 0.5]], device="cuda").expand(4, -1)
|
|
actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
|
|
torch.testing.assert_close(actual, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
|
def test_native_nucleus_boundaries(dtype):
|
|
from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
|
|
|
|
torch.manual_seed(1923)
|
|
logits = torch.randn(1, 2048, device="cuda", dtype=dtype)
|
|
values, indices = logits.sort(descending=True)
|
|
probabilities = values.softmax(-1, dtype=torch.float32)
|
|
before = probabilities.cumsum(-1) - probabilities
|
|
for rank in [1, 19, 49, 1000, 2000]:
|
|
boundary = before[0, rank]
|
|
uniforms = torch.full_like(logits, 0.001, dtype=torch.float32)
|
|
uniforms[0, indices[0, rank]] = 1 - 1e-7
|
|
for direction in [-torch.inf, torch.inf]:
|
|
top_p = torch.nextafter(boundary, boundary.new_tensor(direction)).item()
|
|
actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=0, top_p=top_p)
|
|
torch.testing.assert_close(actual, _reference(logits, uniforms, 0, top_p), rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.parametrize("top_p", [0.8, 1.0])
|
|
def test_strided_logits_and_non_power_of_two_vocabulary(top_p):
|
|
from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
|
|
|
|
for seed in range(16):
|
|
torch.manual_seed(seed)
|
|
logits = torch.randn(128, 2050, device="cuda")[:, ::2]
|
|
uniforms = torch.rand(128, 3, 1025, device="cuda")[:, 1].clamp_min_(1e-7)
|
|
actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=-1, top_p=top_p)
|
|
torch.testing.assert_close(actual, _reference(logits, uniforms, -1, top_p), rtol=0, atol=0)
|