EngineX-Ascend/enginex-ascend-910-vllm
11
0
Fork 1
Code Issues Pull Requests Actions Projects Releases Wiki Activity

init v0.23.0

Browse Source 
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
This commit is contained in:
Sun Ruoxi
2026-08-27 15:11:51 +08:00
parent b582a8e7d1
commit 7f8a1b1f7a
2849 changed files with 712887 additions and 22001 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
tests/ut/sample/__init__.py Normal file
View File

570
tests/ut/sample/a2/test_gumbel_sampling.py Normal file
View File

@@ -0,0 +1,570 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
# This file is a part of the vllm-ascend project.
#
# Tests for vllm_ascend.worker.v2.sample.gumbel on Ascend NPU.
# Validates gumbel_sample and apply_temperature against PyTorch references.
import pytest
import torch
from vllm_ascend.worker.v2.sample.gumbel import apply_temperature, gumbel_sample
DEVICE = "npu"
def _ref_apply_temperature(
logits: torch.Tensor,
expanded_idx_mapping: torch.Tensor,
temperature: torch.Tensor,
) -> torch.Tensor:
"""Pure-Python reference for temperature scaling."""
out = logits.clone().float()
for tok in range(logits.shape[0]):
req = expanded_idx_mapping[tok].item()
temp = temperature[req].item()
if temp == 0.0 or temp == 1.0:
continue
out[tok] = out[tok] / temp
return out
class TestGumbelSampling:
@pytest.mark.parametrize(
"num_tokens,vocab_size",
[
(1, 32000),
(8, 32000),
(48, 102400),
(64, 151936),
],
)
def test_apply_temperature(self, num_tokens, vocab_size):
"""Temperature kernel matches PyTorch reference for various vocab sizes."""
torch.manual_seed(0)
logits = torch.randn(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
expanded_idx_mapping = torch.randint(0, num_tokens, (num_tokens,), dtype=torch.int32, device=DEVICE)
temperature = torch.rand(num_tokens, dtype=torch.float32, device=DEVICE) * 1.8 + 0.2
# inject edge cases
temperature[0] = 0.0
if num_tokens > 1:
temperature[1] = 1.0
logits_triton = logits.clone()
apply_temperature(logits_triton, expanded_idx_mapping, temperature)
torch.npu.synchronize()
logits_ref = _ref_apply_temperature(logits, expanded_idx_mapping, temperature)
assert torch.allclose(logits_triton.float(), logits_ref, atol=1e-4, rtol=1e-5), (
f"apply_temperature mismatch: max_diff={(logits_triton.float() - logits_ref).abs().max().item():.6f}"
)
def test_apply_temperature_skip_zero_and_one(self):
"""Logits should be unchanged for temp=0.0 and temp=1.0."""
torch.manual_seed(10)
num_tokens = 4
vocab_size = 32000
logits = torch.randn(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
expanded_idx_mapping = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
temperature = torch.tensor([0.0, 1.0, 0.0, 1.0], dtype=torch.float32, device=DEVICE)
original = logits.clone()
apply_temperature(logits, expanded_idx_mapping, temperature)
torch.npu.synchronize()
assert torch.equal(logits, original), "Logits changed for temp=0.0 or temp=1.0"
@pytest.mark.parametrize(
"num_tokens,num_reqs,vocab_size",
[
(1, 1, 32000),
(4, 4, 32000),
(8, 4, 32000), # expanded: multiple tokens per request
(16, 8, 102400),
],
)
def test_gumbel_sample_greedy(self, num_tokens, num_reqs, vocab_size):
"""temperature=0 must return argmax (greedy)."""
torch.manual_seed(42)
logits = torch.randn(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
expanded_idx_mapping = torch.randint(0, num_reqs, (num_tokens,), dtype=torch.int32, device=DEVICE)
temperature = torch.zeros(num_reqs, dtype=torch.float32, device=DEVICE)
seed = torch.randint(0, 2**31, (num_reqs,), dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
sampled = gumbel_sample(logits, expanded_idx_mapping, temperature, seed, pos, apply_temperature=False)
torch.npu.synchronize()
expected = logits.argmax(dim=-1)
assert torch.equal(sampled, expected), (
f"Greedy mismatch: sampled={sampled.tolist()} expected={expected.tolist()}"
)
def test_gumbel_sample_greedy_apply_temp_flag_irrelevant(self):
"""With temp=0, apply_temperature flag should not affect result (both greedy)."""
torch.manual_seed(55)
num_tokens, num_reqs, vocab_size = 4, 4, 32000
logits = torch.randn(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
expanded_idx_mapping = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
temperature = torch.zeros(num_reqs, dtype=torch.float32, device=DEVICE)
seed = torch.randint(0, 2**31, (num_reqs,), dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
s_false = gumbel_sample(logits, expanded_idx_mapping, temperature, seed, pos, apply_temperature=False)
s_true = gumbel_sample(logits, expanded_idx_mapping, temperature, seed, pos, apply_temperature=True)
torch.npu.synchronize()
expected = logits.argmax(dim=-1)
assert torch.equal(s_false, expected)
assert torch.equal(s_true, expected)
@pytest.mark.parametrize(
"num_tokens,num_reqs,vocab_size",
[
(4, 4, 32000),
(8, 4, 32000),
(16, 8, 102400),
],
)
def test_gumbel_sample_deterministic(self, num_tokens, num_reqs, vocab_size):
"""Same seed must produce identical results across runs."""
torch.manual_seed(7)
logits = torch.randn(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
expanded_idx_mapping = torch.randint(0, num_reqs, (num_tokens,), dtype=torch.int32, device=DEVICE)
temperature = torch.rand(num_reqs, dtype=torch.float32, device=DEVICE) * 1.5 + 0.5
seed = torch.randint(0, 2**31, (num_reqs,), dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
r1 = gumbel_sample(logits, expanded_idx_mapping, temperature, seed, pos, apply_temperature=False)
torch.npu.synchronize()
r2 = gumbel_sample(logits, expanded_idx_mapping, temperature, seed, pos, apply_temperature=False)
torch.npu.synchronize()
assert torch.equal(r1, r2), "gumbel_sample is non-deterministic with same seed"
def test_gumbel_sample_different_seeds(self):
"""Different seeds must (almost surely) produce different results."""
torch.manual_seed(8)
num_tokens, num_reqs, vocab_size = 16, 16, 32000
logits = torch.randn(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
expanded_idx_mapping = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
temperature = torch.ones(num_reqs, dtype=torch.float32, device=DEVICE) * 1.0
pos = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
seed1 = torch.randint(0, 2**31, (num_reqs,), dtype=torch.int64, device=DEVICE)
seed2 = torch.randint(0, 2**31, (num_reqs,), dtype=torch.int64, device=DEVICE)
# Ensure seeds differ
seed2[0] = seed1[0] + 1
r1 = gumbel_sample(logits, expanded_idx_mapping, temperature, seed1, pos, apply_temperature=False)
r2 = gumbel_sample(logits, expanded_idx_mapping, temperature, seed2, pos, apply_temperature=False)
torch.npu.synchronize()
# With 16 tokens and vocab 32000 at temp=1.0, identical results are astronomically unlikely
assert not torch.equal(r1, r2), "Different seeds produced identical results"
@pytest.mark.parametrize(
"num_tokens,num_reqs,vocab_size",
[
(4, 4, 32000),
(8, 4, 32000),
(16, 8, 102400),
],
)
def test_gumbel_sample_valid_token_ids(self, num_tokens, num_reqs, vocab_size):
"""Sampled token IDs must be in [0, vocab_size)."""
torch.manual_seed(3)
logits = torch.randn(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
expanded_idx_mapping = torch.randint(0, num_reqs, (num_tokens,), dtype=torch.int32, device=DEVICE)
temperature = torch.rand(num_reqs, dtype=torch.float32, device=DEVICE) + 0.1
seed = torch.randint(0, 2**31, (num_reqs,), dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
sampled = gumbel_sample(logits, expanded_idx_mapping, temperature, seed, pos, apply_temperature=False)
torch.npu.synchronize()
assert sampled.shape == (num_tokens,)
assert (sampled >= 0).all() and (sampled < vocab_size).all(), (
f"Out-of-range token IDs: min={sampled.min()}, max={sampled.max()}"
)
def test_gumbel_sample_temperature_affects_distribution(self):
"""Higher temperature should increase sampling entropy (less concentrated).
Strategy: create logits with a clear winner. At low temp the winner should
be sampled most often. At high temp other tokens get more probability.
"""
vocab_size = 100
num_trials = 256
logits_base = torch.zeros(1, vocab_size, dtype=torch.float32, device=DEVICE)
logits_base[0, 0] = 10.0 # strong signal at token 0
expanded_idx_mapping = torch.zeros(1, dtype=torch.int32, device=DEVICE)
low_temp = torch.tensor([0.1], dtype=torch.float32, device=DEVICE)
high_temp = torch.tensor([5.0], dtype=torch.float32, device=DEVICE)
low_temp_winner_count = 0
high_temp_winner_count = 0
for i in range(num_trials):
seed = torch.tensor([i * 1000 + 42], dtype=torch.int64, device=DEVICE)
pos = torch.tensor([i], dtype=torch.int32, device=DEVICE)
s_low = gumbel_sample(
logits_base.clone(), expanded_idx_mapping, low_temp, seed, pos, apply_temperature=True
)
s_high = gumbel_sample(
logits_base.clone(), expanded_idx_mapping, high_temp, seed, pos, apply_temperature=True
)
if s_low.item() == 0:
low_temp_winner_count += 1
if s_high.item() == 0:
high_temp_winner_count += 1
torch.npu.synchronize()
# Low temp should pick the winner much more often than high temp
assert low_temp_winner_count > high_temp_winner_count, (
f"Low temp winner count ({low_temp_winner_count}) should be > "
f"high temp winner count ({high_temp_winner_count})"
)
# Low temp with such a strong signal should almost always pick token 0
assert low_temp_winner_count > num_trials * 0.9, (
f"Low temp winner count ({low_temp_winner_count}/{num_trials}) should be >90%"
)
@pytest.mark.parametrize(
"num_tokens,num_reqs,vocab_size",
[
(4, 4, 32000),
(8, 4, 32000),
],
)
def test_gumbel_sample_mixed_temperature(self, num_tokens, num_reqs, vocab_size):
"""Mix of temp=0 and temp>0: temp=0 tokens must be greedy."""
torch.manual_seed(11)
logits = torch.randn(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
# identity mapping: token i -> request i (for simplicity)
expanded_idx_mapping = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
temperature = torch.rand(num_tokens, dtype=torch.float32, device=DEVICE) + 0.5
# force first half to greedy
temperature[: num_tokens // 2] = 0.0
seed = torch.randint(0, 2**31, (num_tokens,), dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
sampled = gumbel_sample(logits, expanded_idx_mapping, temperature, seed, pos, apply_temperature=False)
torch.npu.synchronize()
greedy = logits.argmax(dim=-1)
for tok in range(num_tokens // 2):
assert sampled[tok].item() == greedy[tok].item(), (
f"Token {tok} (temp=0) should be greedy: got {sampled[tok].item()}, expected {greedy[tok].item()}"
)
def test_gumbel_sample_expanded_idx_mapping(self):
"""Multiple tokens mapping to the same request must work correctly."""
torch.manual_seed(99)
num_tokens = 6
num_reqs = 2
vocab_size = 32000
logits = torch.randn(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
# tokens 0,1,2 -> req 0; tokens 3,4,5 -> req 1
expanded_idx_mapping = torch.tensor([0, 0, 0, 1, 1, 1], dtype=torch.int32, device=DEVICE)
temperature = torch.zeros(num_reqs, dtype=torch.float32, device=DEVICE)
seed = torch.randint(0, 2**31, (num_reqs,), dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
sampled = gumbel_sample(logits, expanded_idx_mapping, temperature, seed, pos, apply_temperature=False)
torch.npu.synchronize()
expected = logits.argmax(dim=-1)
assert torch.equal(sampled, expected), (
f"Expanded mapping greedy mismatch: {sampled.tolist()} vs {expected.tolist()}"
)
def test_gumbel_sample_shared_seed_same_request(self):
"""Tokens mapping to the same request share seed, so with same pos they
should produce the same Gumbel noise and therefore the same sample (given
same logits)."""
torch.manual_seed(42)
vocab_size = 32000
num_reqs = 1
# Two tokens with identical logits, same request, same position
logits_row = torch.randn(1, vocab_size, dtype=torch.float32, device=DEVICE)
logits = logits_row.repeat(2, 1)
expanded_idx_mapping = torch.tensor([0, 0], dtype=torch.int32, device=DEVICE)
temperature = torch.tensor([0.8], dtype=torch.float32, device=DEVICE)
seed = torch.randint(0, 2**31, (num_reqs,), dtype=torch.int64, device=DEVICE)
# Same pos -> same Gumbel noise
pos = torch.tensor([5, 5], dtype=torch.int32, device=DEVICE)
sampled = gumbel_sample(logits, expanded_idx_mapping, temperature, seed, pos, apply_temperature=True)
torch.npu.synchronize()
assert sampled[0].item() == sampled[1].item(), (
f"Tokens with same logits, seed, and pos should sample the same token: "
f"got {sampled[0].item()} vs {sampled[1].item()}"
)
def test_gumbel_sample_apply_temperature_true_nonzero(self):
"""apply_temperature=True with temp>0 must divide logits by temperature
before adding Gumbel noise. Verify via processed_logits output."""
torch.manual_seed(77)
num_tokens, num_reqs, vocab_size = 4, 4, 32000
logits = torch.randn(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
expanded_idx_mapping = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
temperature = torch.rand(num_reqs, dtype=torch.float32, device=DEVICE) * 1.5 + 0.5
seed = torch.randint(0, 2**31, (num_reqs,), dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
# Use processed_logits to verify temperature was applied
out_logits = torch.zeros(num_reqs, vocab_size, dtype=torch.float32, device=DEVICE)
gumbel_sample(
logits,
expanded_idx_mapping,
temperature,
seed,
pos,
apply_temperature=True,
output_processed_logits=out_logits,
)
torch.npu.synchronize()
for tok in range(num_tokens):
req = expanded_idx_mapping[tok].item()
temp = temperature[req].item()
expected = logits[tok].float() / temp
assert torch.allclose(out_logits[req].float(), expected, atol=1e-4, rtol=1e-4), (
f"processed_logits mismatch at token {tok} (req {req}, temp={temp:.3f}): "
f"max_diff={(out_logits[req].float() - expected).abs().max().item():.6f}"
)
def test_gumbel_sample_apply_temperature_false_nonzero(self):
"""apply_temperature=False with temp>0: processed_logits must contain
raw logits (no temperature division), but Gumbel noise is still added
to sampling."""
torch.manual_seed(78)
num_tokens, num_reqs, vocab_size = 4, 4, 32000
logits = torch.randn(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
expanded_idx_mapping = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
temperature = torch.rand(num_reqs, dtype=torch.float32, device=DEVICE) * 1.5 + 0.5
seed = torch.randint(0, 2**31, (num_reqs,), dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
out_logits = torch.zeros(num_reqs, vocab_size, dtype=torch.float32, device=DEVICE)
gumbel_sample(
logits,
expanded_idx_mapping,
temperature,
seed,
pos,
apply_temperature=False,
output_processed_logits=out_logits,
)
torch.npu.synchronize()
for tok in range(num_tokens):
req = expanded_idx_mapping[tok].item()
# Without temperature application, stored logits should match raw logits
expected = logits[tok].float()
assert torch.allclose(out_logits[req].float(), expected, atol=1e-4, rtol=1e-4), (
f"processed_logits should be raw logits when apply_temperature=False: "
f"max_diff={(out_logits[req].float() - expected).abs().max().item():.6f}"
)
def test_gumbel_sample_processed_logits_req_state_idx(self):
"""Processed logits must be stored at req_state_idx position, not token_idx.
This tests the EAGLE speculative decoding scenario where the idx_mapping
is non-contiguous (e.g., active requests [2,5,7,0] out of 8 slots).
The buffer is shaped [max_num_reqs, vocab_size] and the kernel must store
at the correct request slot.
"""
torch.manual_seed(200)
num_tokens = 4
max_num_reqs = 8
vocab_size = 4096
logits = torch.randn(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
# Non-contiguous mapping: tokens 0-3 map to requests 2,5,7,0
expanded_idx_mapping = torch.tensor([2, 5, 7, 0], dtype=torch.int32, device=DEVICE)
temperature = torch.ones(max_num_reqs, dtype=torch.float32, device=DEVICE) * 0.8
seed = torch.randint(0, 2**31, (max_num_reqs,), dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
out_logits = torch.zeros(max_num_reqs, vocab_size, dtype=torch.float32, device=DEVICE)
gumbel_sample(
logits,
expanded_idx_mapping,
temperature,
seed,
pos,
apply_temperature=True,
output_processed_logits=out_logits,
)
torch.npu.synchronize()
for tok in range(num_tokens):
req = expanded_idx_mapping[tok].item()
temp = temperature[req].item()
expected = logits[tok].float() / temp
actual = out_logits[req]
assert torch.allclose(actual.float(), expected, atol=1e-4, rtol=1e-4), (
f"Req {req} (tok={tok}, temp={temp:.3f}): max_diff={(actual.float() - expected).abs().max().item():.6f}"
)
# Also verify that unused request slots remain zero
used_reqs = set(expanded_idx_mapping.tolist())
for req in range(max_num_reqs):
if req not in used_reqs:
assert (out_logits[req] == 0).all(), f"Unused request slot {req} should be all zeros"
def test_gumbel_sample_processed_logits_col(self):
"""output_processed_logits_col selects which column (draft step) to write.
Simulates EAGLE with buffer [max_num_reqs, num_steps, vocab_size].
"""
torch.manual_seed(201)
num_tokens = 3
max_num_reqs = 4
vocab_size = 2048
num_steps = 3
logits = torch.randn(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
expanded_idx_mapping = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
temperature = torch.ones(max_num_reqs, dtype=torch.float32, device=DEVICE) * 0.9
seed = torch.randint(0, 2**31, (max_num_reqs,), dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
# Buffer: [max_num_reqs, num_steps, vocab_size]
draft_logits = torch.zeros(max_num_reqs, num_steps, vocab_size, dtype=torch.float32, device=DEVICE)
# Write to column (step) 1
col_tensor = torch.tensor(1, dtype=torch.int32, device=DEVICE)
gumbel_sample(
logits,
expanded_idx_mapping,
temperature,
seed,
pos,
apply_temperature=True,
output_processed_logits=draft_logits,
output_processed_logits_col=col_tensor,
)
torch.npu.synchronize()
for tok in range(num_tokens):
req = expanded_idx_mapping[tok].item()
temp = temperature[req].item()
expected = logits[tok].float() / temp
# Data should be at draft_logits[req, 1, :] (column 1)
actual = draft_logits[req, 1, :]
assert torch.allclose(actual.float(), expected, atol=1e-4, rtol=1e-4), (
f"Token {tok} at col=1: mismatch, max_diff={(actual.float() - expected).abs().max().item():.6f}"
)
# Column 0 and 2 should be untouched (zeros)
assert (draft_logits[req, 0, :] == 0).all(), f"Col 0 should be zeros for req {req}"
assert (draft_logits[req, 2, :] == 0).all(), f"Col 2 should be zeros for req {req}"
def test_gumbel_sample_processed_logits_mixed_temp(self):
"""Processed logits with mixed temperature (1:1 token-to-request mapping):
- temp=0: stored logits should be raw (no scaling)
- temp>0 with apply_temperature=True: stored logits should be logits/temp
Note: In practice, output_processed_logits is only used by EAGLE
speculative decoding, which always has 1:1 token-to-request mapping.
Multiple tokens per request would cause a write race (undefined order).
"""
torch.manual_seed(88)
num_tokens = 4
num_reqs = 4
vocab_size = 4096
logits = torch.randn(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
# 1:1 mapping: token i -> request i (matches EAGLE usage)
expanded_idx_mapping = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
temperature = torch.tensor([0.0, 0.8, 1.5, 0.0], dtype=torch.float32, device=DEVICE)
seed = torch.randint(0, 2**31, (num_reqs,), dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
out_logits = torch.zeros(num_reqs, vocab_size, dtype=torch.float32, device=DEVICE)
gumbel_sample(
logits,
expanded_idx_mapping,
temperature,
seed,
pos,
apply_temperature=True,
output_processed_logits=out_logits,
)
torch.npu.synchronize()
for tok in range(num_tokens):
req = expanded_idx_mapping[tok].item()
temp = temperature[req].item()
if temp == 0.0:
expected = logits[tok].float()
else:
expected = logits[tok].float() / temp
actual = out_logits[req]
assert torch.allclose(actual.float(), expected, atol=1e-4, rtol=1e-4), (
f"Req {req} (tok={tok}, temp={temp:.3f}): max_diff={(actual.float() - expected).abs().max().item():.6f}"
)
def test_gumbel_sample_single_token(self):
"""Single token with temperature > 0 should work."""
torch.manual_seed(42)
logits = torch.randn(1, 32000, dtype=torch.float32, device=DEVICE)
expanded_idx_mapping = torch.tensor([0], dtype=torch.int32, device=DEVICE)
temperature = torch.tensor([0.7], dtype=torch.float32, device=DEVICE)
seed = torch.tensor([12345], dtype=torch.int64, device=DEVICE)
pos = torch.tensor([0], dtype=torch.int32, device=DEVICE)
sampled = gumbel_sample(logits, expanded_idx_mapping, temperature, seed, pos, apply_temperature=True)
torch.npu.synchronize()
assert sampled.shape == (1,)
assert 0 <= sampled.item() < 32000
def test_gumbel_sample_large_vocab(self):
"""Large vocabulary (151936 = Qwen2) should work correctly."""
torch.manual_seed(401)
vocab_size = 151936
num_tokens = 4
logits = torch.randn(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
expanded_idx_mapping = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
temperature = torch.zeros(num_tokens, dtype=torch.float32, device=DEVICE)
seed = torch.randint(0, 2**31, (num_tokens,), dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
sampled = gumbel_sample(logits, expanded_idx_mapping, temperature, seed, pos, apply_temperature=False)
torch.npu.synchronize()
expected = logits.argmax(dim=-1)
assert torch.equal(sampled, expected), "Large vocab greedy mismatch"
def test_gumbel_sample_extreme_temperatures(self):
"""Very low and very high temperatures should not crash."""
torch.manual_seed(42)
num_tokens, vocab_size = 4, 32000
logits = torch.randn(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
expanded_idx_mapping = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
seed = torch.randint(0, 2**31, (num_tokens,), dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_tokens, dtype=torch.int32, device=DEVICE)
# Very low temperature (near-greedy)
low_temp = torch.tensor([0.01, 0.01, 0.01, 0.01], dtype=torch.float32, device=DEVICE)
s1 = gumbel_sample(logits, expanded_idx_mapping, low_temp, seed, pos, apply_temperature=True)
torch.npu.synchronize()
assert (s1 >= 0).all() and (s1 < vocab_size).all()
# Very high temperature (near-uniform)
high_temp = torch.tensor([100.0, 100.0, 100.0, 100.0], dtype=torch.float32, device=DEVICE)
s2 = gumbel_sample(logits, expanded_idx_mapping, high_temp, seed, pos, apply_temperature=True)
torch.npu.synchronize()
assert (s2 >= 0).all() and (s2 < vocab_size).all()

862
tests/ut/sample/test_rejection_sampler.py
View File

@@ -1,25 +1,17 @@
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# This file is a part of the vllm-ascend project.
#
from unittest.mock import patch
import torch
from tests.ut.base import TestBase
from vllm_ascend.sample.rejection_sampler import (
expand_batch_to_tokens, expand_pytorch, rejection_greedy_sample_pytorch,
rejection_random_sample_pytorch, sample_recovered_tokens_pytorch)
expand_batch_to_tokens,
expand_pytorch,
rejection_greedy_sample_pytorch,
rejection_random_sample_block_verify_pytorch,
rejection_random_sample_pytorch,
sample_recovered_tokens_blockwise_pytorch,
sample_recovered_tokens_pytorch,
)
# Global constants
PLACEHOLDER_TOKEN_ID = -1
@@ -27,14 +19,25 @@ GREEDY_TEMPERATURE = 0.0
MAX_SPEC_LEN = 8 # Used as MAX_NUM_TOKENS in expand_batch_to_tokens
class TestAscendRejectionSampler(TestBase):
def mock_pin_memory(original_func):
def func_wo_pin_memory(*args, **kwargs):
if kwargs.get("pin_memory", False):
kwargs["pin_memory"] = False
return original_func(*args, **kwargs)
return func_wo_pin_memory
class TestAscendRejectionSampler(TestBase):
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_rejection_greedy_sample_pytorch(self):
"""Test greedy rejection sampling: stop when draft doesn't match, otherwise append bonus token"""
batch_size = 2
max_spec_len = 2
output_token_ids = torch.full((batch_size, max_spec_len + 1),
PLACEHOLDER_TOKEN_ID)
output_token_ids = torch.full((batch_size, max_spec_len + 1), PLACEHOLDER_TOKEN_ID)
cu_num_draft_tokens = torch.tensor([2, 4])
num_draft_tokens = [2, 2]
@@ -60,25 +63,32 @@ class TestAscendRejectionSampler(TestBase):
assert output_token_ids[1, 0].item() == 20
assert output_token_ids[1, 2].item() == PLACEHOLDER_TOKEN_ID
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_rejection_random_sample_pytorch(self):
"""Test random rejection sampling: accept based on uniform probability"""
batch_size = 2
max_spec_len = 3
output_token_ids = torch.full((batch_size, max_spec_len + 1),
PLACEHOLDER_TOKEN_ID)
output_token_ids = torch.full((batch_size, max_spec_len + 1), PLACEHOLDER_TOKEN_ID)
cu_num_draft_tokens = torch.tensor([2, 1])
draft_token_ids = torch.tensor([1, 0, 2])
draft_probs = torch.tensor([
[0.0, 0.6, 0.0, 0.4], # vocab_size=4
[0.1, 0.2, 0.3, 0.4],
[0.5, 0.5, 0.0, 0.0],
])
target_probs = torch.tensor([
[0.0, 0.8, 0.0, 0.2],
[0.2, 0.1, 0.3, 0.4],
[0.9, 0.1, 0.0, 0.0],
])
draft_probs = torch.tensor(
[
[0.0, 0.6, 0.0, 0.4], # vocab_size=4
[0.1, 0.2, 0.3, 0.4],
[0.5, 0.5, 0.0, 0.0],
]
)
target_probs = torch.tensor(
[
[0.0, 0.8, 0.0, 0.2],
[0.2, 0.1, 0.3, 0.4],
[0.9, 0.1, 0.0, 0.0],
]
)
bonus_token_ids = torch.tensor([[100], [200]])
recovered_token_ids = torch.tensor([1, 2, 3])
uniform_probs = torch.tensor([0.7, 0.6, 0.5])
@@ -104,6 +114,104 @@ class TestAscendRejectionSampler(TestBase):
assert output_token_ids[0, 1].item() == 0
assert output_token_ids[0, 2].item() == 100
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_rejection_random_sample_pytorch_rejects_placeholder(self):
batch_size = 1
max_spec_len = 1
output_token_ids = torch.full((batch_size, max_spec_len + 1), PLACEHOLDER_TOKEN_ID)
cu_num_draft_tokens = torch.tensor([1])
draft_token_ids = torch.tensor([PLACEHOLDER_TOKEN_ID])
target_probs = torch.tensor([[0.0, 0.0, 1.0]])
bonus_token_ids = torch.tensor([[100]])
recovered_token_ids = torch.tensor([2])
uniform_probs = torch.tensor([0.0])
is_greedy = torch.tensor([False])
rejection_random_sample_pytorch(
output_token_ids,
cu_num_draft_tokens,
draft_token_ids,
None,
target_probs,
bonus_token_ids,
recovered_token_ids,
uniform_probs,
is_greedy,
max_spec_len,
vocab_size=3,
IS_NGRAM=True,
)
assert output_token_ids.tolist() == [[2, PLACEHOLDER_TOKEN_ID]]
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_rejection_random_sample_pytorch_rejects_all_placeholder_mtp3(self):
batch_size = 1
max_spec_len = 3
output_token_ids = torch.full((batch_size, max_spec_len + 1), PLACEHOLDER_TOKEN_ID)
cu_num_draft_tokens = torch.tensor([3])
draft_token_ids = torch.tensor([PLACEHOLDER_TOKEN_ID, PLACEHOLDER_TOKEN_ID, PLACEHOLDER_TOKEN_ID])
# Placeholder draft tokens must reject regardless of target probability.
# The recovered token is passed in after recovery sampling.
target_probs = torch.zeros((max_spec_len, 3))
bonus_token_ids = torch.tensor([[100]])
recovered_token_ids = torch.tensor([2, 1, 0])
uniform_probs = torch.tensor([0.0, 0.0, 0.0])
is_greedy = torch.tensor([False])
rejection_random_sample_pytorch(
output_token_ids,
cu_num_draft_tokens,
draft_token_ids,
None,
target_probs,
bonus_token_ids,
recovered_token_ids,
uniform_probs,
is_greedy,
max_spec_len,
vocab_size=3,
IS_NGRAM=True,
)
assert output_token_ids.tolist() == [[2, PLACEHOLDER_TOKEN_ID, PLACEHOLDER_TOKEN_ID, PLACEHOLDER_TOKEN_ID]]
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_sample_recovered_tokens_pytorch_keeps_placeholder_distribution(self):
output_token_ids = torch.empty(1, dtype=torch.int32)
cu_num_draft_tokens = torch.tensor([1])
draft_token_ids = torch.tensor([PLACEHOLDER_TOKEN_ID])
target_probs = torch.tensor([[0.1, 0.2, 0.7]])
q = torch.ones((1, 3), dtype=torch.float32)
sample_recovered_tokens_pytorch(
output_token_ids,
cu_num_draft_tokens,
draft_token_ids,
None,
target_probs,
q,
vocab_size=3,
IS_NGRAM=True,
)
assert output_token_ids.tolist() == [2]
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_expand_pytorch(self):
"""Test expand_pytorch functionality"""
input_ptr = torch.tensor([10, 20, 30], dtype=torch.int32)
@@ -122,39 +230,65 @@ class TestAscendRejectionSampler(TestBase):
expected = torch.tensor([10, 10, 20, 20, 20, 30, 30])
assert torch.equal(output_ptr, expected)
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_expand_batch_to_tokens(self):
"""Test expand_batch_to_tokens wrapper"""
x = torch.tensor([10, 20, 30])
cu_num_tokens = torch.tensor([2, 5, 7])
num_tokens = 7
with patch("vllm_ascend.sample.rejection_sampler.expand_pytorch"
) as mock_kernel:
# Test PyTorch path
with (
patch("vllm_ascend.sample.rejection_sampler.HAS_TRITON", False),
patch("vllm_ascend.sample.rejection_sampler.expand_pytorch") as mock_pytorch,
):
expand_batch_to_tokens(x, cu_num_tokens, num_tokens)
mock_kernel.assert_called_once()
args = mock_kernel.call_args[0]
mock_pytorch.assert_called_once()
args = mock_pytorch.call_args[0]
assert (args[1] == x).all()
assert (args[2] == cu_num_tokens).all()
# Run actual function
result = expand_batch_to_tokens(x, cu_num_tokens, num_tokens)
expected = torch.tensor([10, 10, 20, 20, 20, 30, 30])
assert torch.equal(result, expected)
# Test Triton kernel path
with (
patch("vllm_ascend.sample.rejection_sampler.HAS_TRITON", True),
patch("vllm_ascend.sample.rejection_sampler.expand_triton") as mock_triton,
):
expand_batch_to_tokens(x, cu_num_tokens, num_tokens)
mock_triton.assert_called_once()
call_args = mock_triton.call_args[0]
assert (call_args[2] == x).all()
assert (call_args[3] == cu_num_tokens).all()
# Run actual function
with patch("vllm_ascend.sample.rejection_sampler.HAS_TRITON", False):
result = expand_batch_to_tokens(x, cu_num_tokens, num_tokens)
expected = torch.tensor([10, 10, 20, 20, 20, 30, 30])
assert torch.equal(result, expected)
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_sample_recovered_tokens_pytorch_ngram(self):
"""Test recovered token sampling under n-gram mode"""
output_token_ids = torch.empty(2, dtype=torch.int32)
cu_num_draft_tokens = torch.tensor([1, 2])
draft_token_ids = torch.tensor([1, 2])
draft_probs = None
target_probs = torch.tensor([
[0.1, 0.2, 0.7],
[0.3, 0.3, 0.4],
])
q = torch.tensor([
[0.1, 0.2, 0.7],
[0.5, 0.4, 0.1],
])
target_probs = torch.tensor(
[
[0.1, 0.2, 0.7],
[0.3, 0.3, 0.4],
]
)
q = torch.tensor(
[
[0.1, 0.2, 0.7],
[0.5, 0.4, 0.1],
]
)
vocab_size = 3
sample_recovered_tokens_pytorch(
@@ -171,25 +305,36 @@ class TestAscendRejectionSampler(TestBase):
assert output_token_ids[0].item() == 0
assert output_token_ids[1].item() == 1
def test_sample_recovered_tokens_pytorch_autoregressive(self):
"""Test recovered token sampling for autoregressive models"""
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_reduce_sample_recovered_tokens_pytorch_ngram(self):
"""Test recovered token sampling under n-gram mode"""
output_token_ids = torch.empty(2, dtype=torch.int32)
cu_num_draft_tokens = torch.tensor([1, 1])
draft_token_ids = torch.tensor([0, 1])
draft_probs = torch.tensor([
[0.6, 0.1, 0.3],
[0.2, 0.7, 0.1],
])
target_probs = torch.tensor([
[0.8, 0.1, 0.1],
[0.3, 0.6, 0.1],
])
q = torch.tensor([
[0.5, 0.3, 0.2],
[0.1, 0.8, 0.1],
])
cu_num_draft_tokens = torch.tensor([1, 2])
draft_token_ids = torch.tensor([1, 2])
draft_probs = None
target_probs = torch.tensor(
[
[0.1, 0.2, 0.7],
[0.3, 0.3, 0.4],
]
)
q = torch.tensor(
[
[0.1, 0.2, 0.7],
[0.5, 0.4, 0.1],
]
)
vocab_size = 3
target_indices = torch.tensor(
[
[0, 1, 2],
[0, 1, 2],
]
)
enable_reduce_sampling = True
sample_recovered_tokens_pytorch(
output_token_ids,
cu_num_draft_tokens,
@@ -198,6 +343,593 @@ class TestAscendRejectionSampler(TestBase):
target_probs,
q,
vocab_size,
IS_NGRAM=True,
target_indices=target_indices,
enable_reduce_sampling=enable_reduce_sampling,
)
assert output_token_ids[0].item() == 0
assert output_token_ids[1].item() == 1
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_rejection_random_reduce_sample_block_verify_pytorch(self):
"""Test random rejection sampling for block verify: accept based on uniform probability"""
batch_size = 2
max_spec_len = 3
output_token_ids = torch.full((batch_size, max_spec_len + 1), PLACEHOLDER_TOKEN_ID)
cu_num_draft_tokens = torch.tensor([2, 1])
draft_token_ids = torch.tensor([1, 0, 2])
draft_probs = torch.tensor(
[
[0.0, 0.6, 0.0, 0.4, 0.0],
[0.1, 0.2, 0.3, 0.4, 0.0],
[0.5, 0.5, 0.0, 0.0, 0.0],
]
)
target_probs = torch.tensor(
[
[0.0, 0.8, 0.0, 0.2],
[0.2, 0.1, 0.3, 0.4],
[0.9, 0.1, 0.0, 0.0],
]
)
bonus_token_ids = torch.tensor([[100], [200]])
recovered_token_ids = torch.tensor([1, 2, 3])
uniform_probs = torch.tensor([0.7, 0.6, 0.5])
is_greedy = torch.tensor([False, False])
vocab_size = 5
target_indices = torch.tensor(
[
[0, 1, 2, 3],
[0, 1, 2, 3],
[0, 1, 2, 3],
]
)
enable_reduce_sampling = True
rejection_random_sample_block_verify_pytorch(
output_token_ids,
cu_num_draft_tokens,
draft_token_ids,
draft_probs,
target_probs,
bonus_token_ids,
recovered_token_ids,
uniform_probs,
is_greedy,
max_spec_len,
vocab_size,
IS_NGRAM=False,
target_indices=target_indices,
enable_reduce_sampling=enable_reduce_sampling,
)
assert output_token_ids[0, 0].item() == 1
assert output_token_ids[0, 1].item() == 0
assert output_token_ids[0, 2].item() == 100
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_reduce_sample_recovered_tokens_blockwise_pytorch_ngram(self):
"""Test recovered token sampling for blockwise speculative decoding with n-gram."""
output_token_ids = torch.empty(2, dtype=torch.int32)
cu_num_draft_tokens = torch.tensor([1, 2])
draft_token_ids = torch.tensor([1, 2])
draft_probs = None
target_probs = torch.tensor(
[
[0.1, 0.2, 0.7],
[0.3, 0.3, 0.4],
]
)
q = torch.tensor(
[
[0.1, 0.2, 0.7],
[0.5, 0.4, 0.1],
]
)
vocab_size = 3
target_indices = torch.tensor(
[
[0, 1, 2],
[0, 1, 2],
]
)
enable_reduce_sampling = True
sample_recovered_tokens_blockwise_pytorch(
output_token_ids,
cu_num_draft_tokens,
draft_token_ids,
draft_probs,
target_probs,
q,
vocab_size,
IS_NGRAM=True,
target_indices=target_indices,
enable_reduce_sampling=enable_reduce_sampling,
)
assert output_token_ids[0].item() == 0
assert output_token_ids[1].item() == 1
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_reduce_sample_recovered_tokens_blockwise_pytorch(self):
"""Test recovered token sampling for blockwise speculative decoding."""
output_token_ids = torch.empty(2, dtype=torch.int32)
cu_num_draft_tokens = torch.tensor([1, 2])
draft_token_ids = torch.tensor([0, 1])
draft_probs = torch.tensor(
[
[0.6, 0.1, 0.3],
[0.2, 0.7, 0.1],
]
)
target_probs = torch.tensor(
[
[0.8, 0.1, 0.1],
[0.3, 0.6, 0.1],
]
)
q = torch.tensor(
[
[0.5, 0.3, 0.2],
[0.1, 0.8, 0.1],
]
)
vocab_size = 3
target_indices = torch.tensor(
[
[0, 1, 2],
[0, 1, 2],
]
)
enable_reduce_sampling = True
sample_recovered_tokens_blockwise_pytorch(
output_token_ids,
cu_num_draft_tokens,
draft_token_ids,
draft_probs,
target_probs,
q,
vocab_size,
IS_NGRAM=False,
target_indices=target_indices,
enable_reduce_sampling=enable_reduce_sampling,
)
assert output_token_ids[0].item() == 0
assert output_token_ids[1].item() == 0
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_rejection_random_sample_block_verify_pytorch_standard(self):
"""Test block verify without reduce_sampling: standard full-vocab path."""
batch_size = 2
max_spec_len = 3
output_token_ids = torch.full((batch_size, max_spec_len + 1), PLACEHOLDER_TOKEN_ID)
cu_num_draft_tokens = torch.tensor([2, 3])
draft_token_ids = torch.tensor([1, 0, 2])
draft_probs = torch.tensor(
[
[0.0, 0.6, 0.0, 0.4],
[0.2, 0.0, 0.3, 0.5],
[0.0, 0.0, 0.5, 0.5],
]
)
target_probs = torch.tensor(
[
[0.0, 0.8, 0.0, 0.2],
[0.1, 0.0, 0.3, 0.6],
[0.0, 0.0, 0.9, 0.1],
]
)
bonus_token_ids = torch.tensor([[100], [200]])
recovered_token_ids = torch.tensor([99, 88, 77])
uniform_probs = torch.tensor([0.7, 0.6, 0.5])
is_greedy = torch.tensor([False, False])
vocab_size = 4
rejection_random_sample_block_verify_pytorch(
output_token_ids,
cu_num_draft_tokens,
draft_token_ids,
draft_probs,
target_probs,
bonus_token_ids,
recovered_token_ids,
uniform_probs,
is_greedy,
max_spec_len,
vocab_size,
IS_NGRAM=False,
)
assert output_token_ids[0, 0].item() == 1
assert output_token_ids[0, 1].item() == 0
assert output_token_ids[0, 2].item() == 100
assert output_token_ids[1, 0].item() == 2
assert output_token_ids[1, 1].item() == 200
class TestEntropyVerify(TestBase):
"""Test ENTROPY_VERIFY mode in rejection sampling.
Entropy verify modifies the acceptance threshold based on the entropy
of the original target distribution:
- High entropy (uncertain) → lower effective threshold → more accepting
- Low entropy (certain) → higher effective threshold → stricter
"""
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_entropy_verify_standard_high_entropy_accepts_more(self):
"""High entropy (uniform-like) makes acceptance easier via lower threshold."""
batch_size = 2
max_spec_len = 2
output_token_ids = torch.full((batch_size, max_spec_len + 1), PLACEHOLDER_TOKEN_ID)
cu_num_draft_tokens = torch.tensor([2, 1])
draft_token_ids = torch.tensor([1, 0, 2])
draft_probs = torch.tensor(
[
[0.6, 0.4, 0.0],
[0.2, 0.8, 0.0],
[0.5, 0.5, 0.0],
]
)
target_probs = torch.tensor(
[
[0.8, 0.2, 0.0],
[0.1, 0.9, 0.0],
[0.9, 0.1, 0.0],
]
)
bonus_token_ids = torch.tensor([[100], [200]])
recovered_token_ids = torch.tensor([99, 88, 77])
uniform_probs = torch.tensor([0.7, 0.6, 0.5])
is_greedy = torch.tensor([False, False])
vocab_size = 3
ori_target_probs = torch.tensor(
[
[0.8, 0.19, 0.01],
[0.09, 0.9, 0.01],
[0.9, 0.09, 0.01],
]
)
rejection_random_sample_pytorch(
output_token_ids,
cu_num_draft_tokens,
draft_token_ids,
draft_probs,
target_probs,
bonus_token_ids,
recovered_token_ids,
uniform_probs,
is_greedy,
max_spec_len,
vocab_size,
IS_NGRAM=False,
ENTROPY_VERIFY=True,
POSTERIOR_THRESHOLD=0.95,
POSTERIOR_ALPHA=0.4,
EPSILON=1e-10,
ori_target_probs=ori_target_probs,
)
assert output_token_ids[0, 0].item() == 99
assert output_token_ids[0, 1].item() == -1
assert output_token_ids[0, 2].item() == -1
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_entropy_verify_standard_low_entropy_stricter(self):
"""Low entropy (peaked distribution) keeps threshold near POSTERIOR_THRESHOLD."""
batch_size = 1
max_spec_len = 2
output_token_ids = torch.full((batch_size, max_spec_len + 1), PLACEHOLDER_TOKEN_ID)
cu_num_draft_tokens = torch.tensor([2])
draft_token_ids = torch.tensor([1, 0])
draft_probs = torch.tensor(
[
[0.6, 0.4, 0.0],
[0.8, 0.2, 0.0],
]
)
target_probs = torch.tensor(
[
[0.8, 0.2, 0.0],
[0.1, 0.9, 0.0],
]
)
bonus_token_ids = torch.tensor([[100]])
recovered_token_ids = torch.tensor([99, 88])
uniform_probs = torch.tensor([0.7, 0.6])
is_greedy = torch.tensor([False])
vocab_size = 3
ori_target_probs = torch.tensor(
[
[0.8, 0.19, 0.01],
[0.09, 0.9, 0.01],
]
)
rejection_random_sample_pytorch(
output_token_ids,
cu_num_draft_tokens,
draft_token_ids,
draft_probs,
target_probs,
bonus_token_ids,
recovered_token_ids,
uniform_probs,
is_greedy,
max_spec_len,
vocab_size,
IS_NGRAM=False,
ENTROPY_VERIFY=True,
POSTERIOR_THRESHOLD=0.95,
POSTERIOR_ALPHA=0.4,
EPSILON=1e-10,
ori_target_probs=ori_target_probs,
)
assert output_token_ids[0, 0].item() == 99
assert output_token_ids[0, 1].item() == -1
assert output_token_ids[0, 2].item() == -1
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_entropy_verify_block_verify(self):
"""Entropy verify with block verify mode."""
batch_size = 2
max_spec_len = 3
output_token_ids = torch.full((batch_size, max_spec_len + 1), PLACEHOLDER_TOKEN_ID)
cu_num_draft_tokens = torch.tensor([2, 1])
draft_token_ids = torch.tensor([1, 0, 2])
draft_probs = torch.tensor(
[
[0.6, 0.4, 0.0, 0.0],
[0.2, 0.8, 0.0, 0.0],
[0.5, 0.5, 0.0, 0.0],
]
)
target_probs = torch.tensor(
[
[0.8, 0.2, 0.0, 0.0],
[0.1, 0.9, 0.0, 0.0],
[0.9, 0.1, 0.0, 0.0],
]
)
bonus_token_ids = torch.tensor([[100], [200]])
recovered_token_ids = torch.tensor([99, 88, 77])
uniform_probs = torch.tensor([0.7, 0.6, 0.5])
is_greedy = torch.tensor([False, False])
vocab_size = 4
ori_target_probs = torch.tensor(
[
[0.8, 0.18, 0.01, 0.01],
[0.88, 0.9, 0.01, 0.01],
[0.9, 0.08, 0.01, 0.01],
]
)
rejection_random_sample_block_verify_pytorch(
output_token_ids,
cu_num_draft_tokens,
draft_token_ids,
draft_probs,
target_probs,
bonus_token_ids,
recovered_token_ids,
uniform_probs,
is_greedy,
max_spec_len,
vocab_size,
IS_NGRAM=False,
ENTROPY_VERIFY=True,
POSTERIOR_THRESHOLD=0.95,
POSTERIOR_ALPHA=0.4,
EPSILON=1e-10,
ori_target_probs=ori_target_probs,
)
assert output_token_ids[0, 0].item() == 99
assert output_token_ids[0, 1].item() == -1
assert output_token_ids[0, 2].item() == -1
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_entropy_verify_ngram(self):
"""ENTROPY_VERIFY with IS_NGRAM: draft_probs=None, draft_token_probs=1.0.
In NGRAM mode, acceptance depends on target_prob alone (since
draft_prob=1.0). Entropy verify lowers the threshold for high-entropy
tokens, making acceptance easier when the target distribution is
uncertain.
"""
batch_size = 1
max_spec_len = 2
output_token_ids = torch.full((batch_size, max_spec_len + 1), PLACEHOLDER_TOKEN_ID)
cu_num_draft_tokens = torch.tensor([2])
draft_token_ids = torch.tensor([0, 1])
draft_probs = None
target_probs = torch.tensor(
[
[0.6, 0.2, 0.2],
[0.1, 0.1, 0.8],
]
)
bonus_token_ids = torch.tensor([[100]])
recovered_token_ids = torch.tensor([99, 88])
uniform_probs = torch.tensor([0.7, 0.6])
is_greedy = torch.tensor([False])
vocab_size = 3
ori_target_probs = torch.tensor(
[
[0.6, 0.2, 0.2],
[0.1, 0.1, 0.8],
]
)
rejection_random_sample_pytorch(
output_token_ids,
cu_num_draft_tokens,
draft_token_ids,
draft_probs,
target_probs,
bonus_token_ids,
recovered_token_ids,
uniform_probs,
is_greedy,
max_spec_len,
vocab_size,
IS_NGRAM=True,
ENTROPY_VERIFY=True,
POSTERIOR_THRESHOLD=0.95,
POSTERIOR_ALPHA=0.4,
EPSILON=1e-10,
ori_target_probs=ori_target_probs,
)
assert output_token_ids[0, 0].item() == 0
assert output_token_ids[0, 1].item() == 88
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_entropy_verify_block_verify_ngram(self):
"""ENTROPY_VERIFY + IS_NGRAM + block_verify combined.
Tests the interaction of all three modes: NGRAM (draft_probs=None),
block verify (cumulative acceptance), and entropy-based threshold
adjustment.
"""
batch_size = 1
max_spec_len = 3
output_token_ids = torch.full((batch_size, max_spec_len + 1), PLACEHOLDER_TOKEN_ID)
cu_num_draft_tokens = torch.tensor([2])
draft_token_ids = torch.tensor([0, 1])
draft_probs = None
target_probs = torch.tensor(
[
[0.6, 0.2, 0.2, 0.0],
[0.1, 0.1, 0.8, 0.0],
]
)
bonus_token_ids = torch.tensor([[100]])
recovered_token_ids = torch.tensor([99, 88])
uniform_probs = torch.tensor([0.7, 0.6])
is_greedy = torch.tensor([False])
vocab_size = 4
ori_target_probs = torch.tensor(
[
[0.6, 0.2, 0.2, 0.0],
[0.1, 0.1, 0.8, 0.0],
]
)
rejection_random_sample_block_verify_pytorch(
output_token_ids,
cu_num_draft_tokens,
draft_token_ids,
draft_probs,
target_probs,
bonus_token_ids,
recovered_token_ids,
uniform_probs,
is_greedy,
max_spec_len,
vocab_size,
IS_NGRAM=True,
ENTROPY_VERIFY=True,
POSTERIOR_THRESHOLD=0.95,
POSTERIOR_ALPHA=0.4,
EPSILON=1e-10,
ori_target_probs=ori_target_probs,
)
assert output_token_ids[0, 0].item() == 0
assert output_token_ids[0, 1].item() == 88
@patch("torch.arange", new=mock_pin_memory(torch.arange))
@patch("torch.ones", new=mock_pin_memory(torch.ones))
@patch("torch.full", new=mock_pin_memory(torch.full))
@patch("torch.tensor", new=mock_pin_memory(torch.tensor))
def test_entropy_verify_no_ori_probs_fallback(self):
"""When ori_target_probs is None, fallback to target_probs for entropy."""
batch_size = 1
max_spec_len = 2
output_token_ids = torch.full((batch_size, max_spec_len + 1), PLACEHOLDER_TOKEN_ID)
cu_num_draft_tokens = torch.tensor([2])
draft_token_ids = torch.tensor([1, 0])
draft_probs = torch.tensor(
[
[0.6, 0.4, 0.0],
[0.8, 0.2, 0.0],
]
)
target_probs = torch.tensor(
[
[0.35, 0.33, 0.32],
[0.34, 0.34, 0.32],
]
)
bonus_token_ids = torch.tensor([[100]])
recovered_token_ids = torch.tensor([99, 88])
uniform_probs = torch.tensor([0.7, 0.6])
is_greedy = torch.tensor([False])
vocab_size = 3
rejection_random_sample_pytorch(
output_token_ids,
cu_num_draft_tokens,
draft_token_ids,
draft_probs,
target_probs,
bonus_token_ids,
recovered_token_ids,
uniform_probs,
is_greedy,
max_spec_len,
vocab_size,
IS_NGRAM=False,
ENTROPY_VERIFY=True,
POSTERIOR_THRESHOLD=0.95,
POSTERIOR_ALPHA=0.4,
EPSILON=1e-10,
ori_target_probs=None,
)
assert output_token_ids[0, 0].item() == 1
assert output_token_ids[0, 1].item() in (0, 88)
assert output_token_ids[0, 2].item() == 100

24
tests/ut/sample/test_sampler.py
View File

@@ -1,32 +1,10 @@
from unittest import mock
import torch
from tests.ut.base import TestBase
from vllm_ascend.sample.sampler import AscendSampler, AscendTopKTopPSampler
class TestAscendSampler(TestBase):
def test_init_with_raw_logprobs(self):
sampler = AscendSampler(logprobs_mode="raw_logprobs")
self.assertEqual(sampler.logprobs_mode, "raw_logprobs")
self.assertTrue(hasattr(sampler, 'topk_topp_sampler'))
self.assertTrue(hasattr(sampler, "topk_topp_sampler"))
self.assertIsInstance(sampler.topk_topp_sampler, AscendTopKTopPSampler)
class TestAscendTopKTopPSampler(TestBase):
@mock.patch("torch_npu.npu_top_k_top_p")
def test_npu_topk_topp_called_when_optimized(self, mock_npu_op):
mock_npu_op.return_value = (torch.randn(1, 3))
sampler = AscendTopKTopPSampler()
logits = torch.tensor([[1.0, 2.0, 3.0]])
k = torch.tensor([2])
p = torch.tensor([0.9])
generators = {0: torch.Generator()}
generators[0].manual_seed(42)
sampler.forward_native(logits, generators, k, p)
mock_npu_op.assert_called_once_with(logits, p, k)