A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 14:00 UTC
[BugFix] Fix bad words with speculative decoding (#31908)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jan 7, 2026 at 12:46 UTC
10ef65eded8187df92c370d6ffd7fd2b8a3c1d3c
2 files changed
+26
-22
tests/v1/sample/test_rejection_sampler.py
+13
-14
@@ -691,9 +691,13 @@ def test_frequency_penalties(rejection_sampler):
691
692
693
def test_bad_words(rejection_sampler):
694
- """Test rejection sampling with bad words constraints"""
694
+ """Test rejection sampling with bad words constraints.
695
+
696
+ This test applies bad words to non-consecutive requests (0 and 2, but not 1)
697
+ to verify correct logit indexing when iterating over requests with bad words.
698
+ """
699
spec_tokens = [[1, 2, 3], [1, 15, 3], [1, 2, 3]]
696
- output_tokens = [[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]]
700
+ output_tokens = [[1, 2, 3, 4], [1, 15, 3, 4], [1, 2, 3, 4]]
701
702
logits = create_logits_tensor(output_tokens, token_idx_to_override=15)
703
metadata = create_sampling_metadata(
@@ -701,17 +705,9 @@ def test_bad_words(rejection_sampler):
705
output_token_ids=[[2], [3], [4]],
706
spec_token_ids=spec_tokens,
707
bad_words_token_ids={
704
- 0: [
705
- [
706
- 2,
707
- ]
708
- ],
709
- 1: [
710
- [
711
- 2,
712
- ]
713
- ],
714
- # Do not apply bad words to the last request
708
+ 0: [[2]],
709
+ # Request 1 has no bad words (to test non-consecutive request handling)
710
+ 2: [[2]],
711
},
712
)
713
bonus_token_tensor = torch.tensor(
@@ -726,8 +722,11 @@ def test_bad_words(rejection_sampler):
722
sampling_metadata=metadata,
723
)
724
725
+ # Request 0: bad word [2] matches prefix, so token 2 is rejected -> 15
726
+ # Request 1: no bad words, all tokens match -> [1, 15, 3, 4]
727
+ # Request 2: bad word [2] matches prefix, so token 2 is rejected -> 15
728
expected = torch.tensor(
730
- [[1, 15, -1, -1], [1, 15, 3, 4], [1, 2, 3, 4]],
729
+ [[1, 15, -1, -1], [1, 15, 3, 4], [1, 15, -1, -1]],
730
dtype=torch.int,
731
device=logits.device,
732
)
vllm/v1/sample/ops/bad_words.py
+13
-8
@@ -42,11 +42,16 @@ def apply_bad_words_with_drafts(
42
num_draft_tokens: list[int],
43
) -> None:
44
start_idx = 0
45
- for i, bad_words_ids in bad_words_token_ids.items():
46
- for draft_idx in range(num_draft_tokens[i]):
47
- _apply_bad_words_single_batch(
48
- logits[start_idx + draft_idx],
49
- bad_words_ids,
50
- past_tokens_ids[start_idx + draft_idx],
51
- )
52
- start_idx += num_draft_tokens[i]
45
+ remaining = len(bad_words_token_ids)
46
+ for i, n in enumerate(num_draft_tokens):
47
+ if (bad_words_ids := bad_words_token_ids.get(i)) is not None:
48
+ for draft_idx in range(start_idx, start_idx + n):
49
+ _apply_bad_words_single_batch(
50
+ logits[draft_idx],
51
+ bad_words_ids,
52
+ past_tokens_ids[draft_idx],
53
+ )
54
+ remaining -= 1
55
+ if not remaining:
56
+ break
57
+ start_idx += n