biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[V1][Spec Decode] Change Spec Decode Rejection Sampling API (#13729)

Browse Source
This commit is contained in:
Lily Liu
2025-02-25 18:14:48 -08:00
committed by GitHub
parent 9ba28043b5
commit 5629f26df7
8 changed files with 102 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
tests/v1/sample/test_rejection_sampler.py
View File

@@ -29,7 +29,6 @@ def create_sampling_metadata(spec_tokens: List[List[int]]) -> SamplingMetadata:
temperature=torch.tensor([]),
all_greedy=True,
all_random=False,
spec_token_ids=spec_tokens,
top_p=None,
top_k=None,
min_p=torch.empty(batch_size, ),
@@ -55,7 +54,7 @@ def test_perfect_match(sampler):
metadata = create_sampling_metadata(spec_tokens)
logits = create_logits_tensor(output_tokens)
output = sampler(logits, metadata)
output = sampler(spec_tokens, logits, metadata)
expected = torch.tensor([[1, 2, 3, 4]],
dtype=torch.int,
device=logits.device)
@@ -70,7 +69,7 @@ def test_early_mismatch(sampler):
metadata = create_sampling_metadata(spec_tokens)
logits = create_logits_tensor(output_tokens)
output = sampler(logits, metadata)
output = sampler(spec_tokens, logits, metadata)
expected = torch.tensor([[1, 5, INVALID_TOKEN_ID, INVALID_TOKEN_ID]],
dtype=torch.int,
device=logits.device)
@@ -85,7 +84,7 @@ def test_multiple_sequences(sampler):
metadata = create_sampling_metadata(spec_tokens)
logits = create_logits_tensor(output_tokens)
output = sampler(logits, metadata)
output = sampler(spec_tokens, logits, metadata)
expected = torch.tensor([[1, 2, 5], [3, 4, INVALID_TOKEN_ID]],
dtype=torch.int,
device=logits.device)
@@ -100,7 +99,7 @@ def test_single_token_sequence(sampler):
metadata = create_sampling_metadata(spec_tokens)
logits = create_logits_tensor(output_tokens)
output = sampler(logits, metadata)
output = sampler(spec_tokens, logits, metadata)
expected = torch.tensor([[1, 2]], dtype=torch.int, device=logits.device)
assert torch.equal(output.sampled_token_ids, expected)
@@ -113,7 +112,7 @@ def test_empty_sequence(sampler):
metadata = create_sampling_metadata(spec_tokens)
logits = create_logits_tensor(output_tokens)
output = sampler(logits, metadata)
output = sampler(spec_tokens, logits, metadata)
expected = torch.tensor([[5]], dtype=torch.int, device=logits.device)
assert torch.equal(output.sampled_token_ids, expected)
@@ -126,7 +125,7 @@ def test_multiple_mismatches(sampler):
metadata = create_sampling_metadata(spec_tokens)
logits = create_logits_tensor(output_tokens)
output = sampler(logits, metadata)
output = sampler(spec_tokens, logits, metadata)
expected = torch.tensor([[1, 2, 7, INVALID_TOKEN_ID],
[4, 8, INVALID_TOKEN_ID, INVALID_TOKEN_ID]],
dtype=torch.int,
@@ -147,7 +146,7 @@ def test_parametrized_cases(sampler, spec_tokens, output_tokens, expected):
metadata = create_sampling_metadata(spec_tokens)
logits = create_logits_tensor(output_tokens)
output = sampler(logits, metadata)
output = sampler(spec_tokens, logits, metadata)
expected_tensor = torch.tensor(expected,
dtype=torch.int,
device=logits.device)
@@ -163,7 +162,7 @@ def test_logits_shape_handling(sampler):
metadata = create_sampling_metadata(spec_tokens)
logits = create_logits_tensor(output_tokens, vocab_size)
output = sampler(logits, metadata)
output = sampler(spec_tokens, logits, metadata)
expected = torch.tensor([[1, 2, 3]], dtype=torch.int, device=logits.device)
assert torch.equal(output.sampled_token_ids, expected)
assert logits.shape[-1] == vocab_size