Rehan KhanandBob
fd0216d3b5
Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com> Co-authored-by: Bob <bob@ibm.com>
256 lines
11 KiB
Python
256 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import tracemalloc
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.platforms import current_platform
|
|
|
|
if not current_platform.is_cpu():
|
|
pytest.skip("skipping CPU-only tests", allow_module_level=True)
|
|
|
|
import vllm._C # noqa: F401, E402
|
|
|
|
VOCAB_SIZES = [32000, 49152, 128256]
|
|
BATCH_SIZES = [1, 4, 16]
|
|
|
|
|
|
def reference_greedy_sample(logits: torch.Tensor) -> torch.Tensor:
|
|
return logits.argmax(dim=-1).view(-1)
|
|
|
|
|
|
class TestGreedyArgmax:
|
|
@pytest.mark.parametrize("vocab_size", VOCAB_SIZES)
|
|
@pytest.mark.parametrize("batch_size", BATCH_SIZES)
|
|
def test_exact_match(self, vocab_size: int, batch_size: int):
|
|
logits = torch.randn(batch_size, vocab_size, dtype=torch.float32)
|
|
expected = reference_greedy_sample(logits)
|
|
result = torch.ops._C.greedy_argmax(logits)
|
|
torch.testing.assert_close(result, expected)
|
|
|
|
def test_single_dominant(self):
|
|
logits = torch.full((1, 50000), -1e9, dtype=torch.float32)
|
|
logits[0, 42] = 100.0
|
|
assert torch.ops._C.greedy_argmax(logits).item() == 42
|
|
|
|
def test_negative_logits(self):
|
|
logits = torch.randn(8, 32000, dtype=torch.float32) - 10.0
|
|
expected = reference_greedy_sample(logits)
|
|
result = torch.ops._C.greedy_argmax(logits)
|
|
torch.testing.assert_close(result, expected)
|
|
|
|
|
|
class TestFusedGumbelArgmax:
|
|
def test_distribution_chi_squared(self):
|
|
"""Verify sampling distribution via chi-squared goodness of fit."""
|
|
# Small support is enough for chi-squared power; large-vocab bias
|
|
# is covered by test_noise_not_sliding_window (#59786).
|
|
small_vocab = 100
|
|
logits = torch.randn(1, small_vocab, dtype=torch.float32)
|
|
probs = logits.softmax(dim=-1).squeeze(0)
|
|
|
|
n_samples = 100_000
|
|
seeds = torch.arange(n_samples, dtype=torch.long) * 7 + 13
|
|
idx = torch.ops._C.fused_gumbel_argmax(
|
|
logits.expand(n_samples, -1).contiguous(), seeds
|
|
)
|
|
counts = torch.bincount(idx, minlength=small_vocab).float()
|
|
|
|
expected = probs * n_samples
|
|
mask = expected > 5
|
|
chi2 = ((counts[mask] - expected[mask]) ** 2 / expected[mask]).sum()
|
|
dof = mask.sum().item() - 1
|
|
from scipy.stats import chi2 as chi2_dist
|
|
|
|
p_value = 1.0 - chi2_dist.cdf(chi2.item(), dof)
|
|
assert p_value > 0.001, (
|
|
f"Chi-squared test failed: chi2={chi2.item():.1f}, "
|
|
f"dof={dof}, p={p_value:.6f}"
|
|
)
|
|
|
|
def test_noise_not_sliding_window(self):
|
|
"""Table[(seed+i)&MASK] made consecutive seeds a 1-token shift.
|
|
|
|
On uniform logits that yields P(w(s+1) == w(s)-1) ~ 1-2/V. Independent
|
|
per-token noise yields ~1/V. See #59786.
|
|
"""
|
|
vocab, n = 128, 2048
|
|
logits = torch.zeros(n, vocab, dtype=torch.float32)
|
|
seeds = torch.arange(n, dtype=torch.long)
|
|
winners = torch.ops._C.fused_gumbel_argmax(logits, seeds)
|
|
shifted = (winners[:-1] - 1) % vocab
|
|
frac = (winners[1:] == shifted).float().mean().item()
|
|
assert frac < 0.05, (
|
|
f"consecutive-seed winners follow a table shift (frac={frac:.3f})"
|
|
)
|
|
|
|
def test_zipf_head_matches_softmax(self):
|
|
"""Zipf head probability should match softmax (#59786)."""
|
|
vocab, n_samples, batch = 8192, 16384, 256
|
|
logits = -1.1 * torch.arange(1, vocab + 1, dtype=torch.float32).log()
|
|
p0 = torch.softmax(logits.double(), dim=-1)[0].item()
|
|
row = logits.view(1, -1).expand(batch, vocab).contiguous()
|
|
n_top = 0
|
|
for start in range(0, n_samples, batch):
|
|
seeds = torch.arange(start, start + batch, dtype=torch.long)
|
|
tokens = torch.ops._C.fused_gumbel_argmax(row, seeds)
|
|
n_top += int((tokens == 0).sum())
|
|
q0 = n_top / n_samples
|
|
se = (p0 * (1.0 - p0) / n_samples) ** 0.5
|
|
assert abs(q0 - p0) < 8.0 * se, (
|
|
f"P(top-1) kernel={q0:.4f} exact={p0:.4f} se={se:.4f}"
|
|
)
|
|
|
|
def test_deterministic_same_seed(self):
|
|
"""Same seed produces same result."""
|
|
logits = torch.randn(4, 32000, dtype=torch.float32)
|
|
seeds = torch.tensor([42, 123, 456, 789], dtype=torch.long)
|
|
r1 = torch.ops._C.fused_gumbel_argmax(logits, seeds)
|
|
r2 = torch.ops._C.fused_gumbel_argmax(logits, seeds)
|
|
torch.testing.assert_close(r1, r2)
|
|
|
|
def test_different_seeds_differ(self):
|
|
logits = torch.zeros(16, 50000, dtype=torch.float32)
|
|
seeds_a = torch.arange(16, dtype=torch.long)
|
|
seeds_b = torch.arange(16, dtype=torch.long) + 1_000_000
|
|
r_a = torch.ops._C.fused_gumbel_argmax(logits, seeds_a)
|
|
r_b = torch.ops._C.fused_gumbel_argmax(logits, seeds_b)
|
|
assert not torch.equal(r_a, r_b)
|
|
|
|
@pytest.mark.parametrize("vocab_size", VOCAB_SIZES)
|
|
@pytest.mark.parametrize("batch_size", BATCH_SIZES)
|
|
def test_output_in_range(self, vocab_size: int, batch_size: int):
|
|
logits = torch.randn(batch_size, vocab_size, dtype=torch.float32)
|
|
seeds = torch.arange(batch_size, dtype=torch.long)
|
|
result = torch.ops._C.fused_gumbel_argmax(logits, seeds)
|
|
assert result.min() >= 0
|
|
assert result.max() < vocab_size
|
|
|
|
def test_samples_only_from_finite_logits(self):
|
|
"""top-k/top-p mask with -inf; those IDs must never win."""
|
|
allowed = torch.tensor([3, 17, 99, 1024, 8191], dtype=torch.long)
|
|
logits = torch.full((64, 32000), float("-inf"), dtype=torch.float32)
|
|
logits[:, allowed] = torch.randn(64, allowed.numel())
|
|
tokens = torch.ops._C.fused_gumbel_argmax(
|
|
logits, torch.arange(64, dtype=torch.long)
|
|
)
|
|
assert torch.isin(tokens, allowed).all()
|
|
|
|
def test_single_unmasked_token(self):
|
|
logits = torch.full((8, 4096), float("-inf"), dtype=torch.float32)
|
|
logits[:, 123] = 0.0
|
|
seeds = torch.tensor([0, -1, -(2**63), 2**63 - 1, 42, 99, 1000, 7])
|
|
tokens = torch.ops._C.fused_gumbel_argmax(logits, seeds)
|
|
assert (tokens == 123).all()
|
|
|
|
def test_tokens_beyond_gumbel_max_cannot_win(self):
|
|
"""A gap larger than GUMBEL_MAX is unreachable."""
|
|
logits = torch.zeros(256, 4, dtype=torch.float32)
|
|
logits[:, 1:] = -80.0
|
|
tokens = torch.ops._C.fused_gumbel_argmax(
|
|
logits, torch.arange(256, dtype=torch.long)
|
|
)
|
|
assert (tokens == 0).all()
|
|
|
|
def test_tied_huge_logits_not_index_biased(self):
|
|
"""Noise must still break ties when logits sit at float32 max."""
|
|
logits = torch.full(
|
|
(4096, 2), torch.finfo(torch.float32).max, dtype=torch.float32
|
|
)
|
|
tokens = torch.ops._C.fused_gumbel_argmax(
|
|
logits, torch.arange(4096, dtype=torch.long)
|
|
)
|
|
frac0 = (tokens == 0).float().mean().item()
|
|
assert set(tokens.tolist()) == {0, 1}, tokens.unique().tolist()
|
|
assert 0.35 < frac0 < 0.65, f"P(index 0)={frac0:.3f} expected ~0.5"
|
|
|
|
def test_serial_batch1_matches_batched(self):
|
|
"""OpenMP rows must match the batch=1 serial path."""
|
|
logits = torch.randn(8, 2048, dtype=torch.float32)
|
|
seeds = torch.arange(8, dtype=torch.long)
|
|
batched = torch.ops._C.fused_gumbel_argmax(logits, seeds)
|
|
serial = torch.stack(
|
|
[
|
|
torch.ops._C.fused_gumbel_argmax(logits[i : i + 1], seeds[i : i + 1])[0]
|
|
for i in range(8)
|
|
]
|
|
)
|
|
torch.testing.assert_close(batched, serial)
|
|
|
|
def test_noncontiguous_matches_contiguous(self):
|
|
logits = torch.randn(8, 1024, dtype=torch.float32)[:, ::2]
|
|
seeds = torch.arange(16, dtype=torch.long)[::2]
|
|
assert not logits.is_contiguous() and not seeds.is_contiguous()
|
|
expected = torch.ops._C.fused_gumbel_argmax(
|
|
logits.contiguous(), seeds.contiguous()
|
|
)
|
|
actual = torch.ops._C.fused_gumbel_argmax(logits, seeds)
|
|
torch.testing.assert_close(actual, expected)
|
|
|
|
def test_empty_batch(self):
|
|
tokens = torch.ops._C.fused_gumbel_argmax(
|
|
torch.empty(0, 7), torch.empty(0, dtype=torch.long)
|
|
)
|
|
assert tokens.shape == (0,) and tokens.dtype == torch.int64
|
|
|
|
def test_masked_support_matches_softmax(self):
|
|
"""Finite subset after -inf mask should match softmax."""
|
|
from scipy.stats import chisquare
|
|
|
|
logits = torch.tensor(
|
|
[0.0, -0.2, -0.7, -1.5, -3.0, float("-inf")], dtype=torch.float32
|
|
)
|
|
n_samples = 20_000
|
|
tokens = torch.ops._C.fused_gumbel_argmax(
|
|
logits.expand(n_samples, -1).contiguous(),
|
|
torch.arange(n_samples, dtype=torch.long),
|
|
)
|
|
counts = torch.bincount(tokens, minlength=logits.numel()).double()
|
|
assert counts[-1] == 0
|
|
probs = logits.double().softmax(-1)
|
|
_, p_value = chisquare(counts[:-1].numpy(), (probs[:-1] * n_samples).numpy())
|
|
assert p_value > 1e-7, f"masked-support p-value: {p_value}"
|
|
|
|
def test_zipf_tail_matches_softmax(self):
|
|
"""Independent 53-bit Gumbel must not starve the Zipf tail (#59786)."""
|
|
vocab, n_samples, batch = 8192, 16384, 256
|
|
logits = -1.1 * torch.arange(1, vocab + 1, dtype=torch.float32).log()
|
|
p_tail = torch.softmax(logits.double(), dim=-1)[vocab // 2 :].sum().item()
|
|
row = logits.view(1, -1).expand(batch, vocab).contiguous()
|
|
n_tail = 0
|
|
for start in range(0, n_samples, batch):
|
|
seeds = torch.arange(start, start + batch, dtype=torch.long)
|
|
tokens = torch.ops._C.fused_gumbel_argmax(row, seeds)
|
|
n_tail += int((tokens >= vocab // 2).sum())
|
|
q_tail = n_tail / n_samples
|
|
se = (p_tail * (1.0 - p_tail) / n_samples) ** 0.5
|
|
assert abs(q_tail - p_tail) < 8.0 * se, (
|
|
f"P(tail) kernel={q_tail:.4f} exact={p_tail:.4f} se={se:.4f}"
|
|
)
|
|
|
|
|
|
class TestMemory:
|
|
def test_fused_no_intermediate_allocs(self):
|
|
"""Verify that the fused kernel does not allocate large intermediates."""
|
|
logits = torch.randn(16, 128256, dtype=torch.float32)
|
|
seeds = torch.arange(16, dtype=torch.long)
|
|
|
|
torch.ops._C.fused_gumbel_argmax(logits, seeds)
|
|
|
|
tracemalloc.start()
|
|
snap_before = tracemalloc.take_snapshot()
|
|
for _ in range(50):
|
|
torch.ops._C.fused_gumbel_argmax(logits, seeds)
|
|
snap_after = tracemalloc.take_snapshot()
|
|
tracemalloc.stop()
|
|
|
|
diff = snap_after.compare_to(snap_before, "lineno")
|
|
total_new_bytes = sum(s.size_diff for s in diff if s.size_diff > 0)
|
|
vocab_bytes = 16 * 128256 * 4
|
|
assert total_new_bytes < vocab_bytes, (
|
|
f"Fused kernel allocated {total_new_bytes} bytes, "
|
|
f"expected < {vocab_bytes} (one intermediate tensor)"
|
|
)
|