[Bugfix][Model Runner V2][Spec Decode] Fix off-by-one in bad_words draft-prefix matching (#52311)
Signed-off-by: jyan <r_02213@sjtu.edu.cn> Co-authored-by: jyan <r_02213@sjtu.edu.cn> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Jyan-R committed
Aug 16, 2026 at 17:02 UTC
1b079c40ff9d2598d837f4ed1fc08342fca4fd6e
2 files changed
+108
-1
tests/v1/worker/test_gpu_bad_words.py
new
+105
@@ -0,0 +1,105 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import numpy as np
5
+import pytest
6
+import torch
7
+
8
+pytest.importorskip("triton")
9
+if not torch.cuda.is_available():
10
+ pytest.skip(
11
+ "CUDA required for Model Runner V2 bad words tests",
12
+ allow_module_level=True,
13
+ )
14
+
15
+from vllm.sampling_params import SamplingParams
16
+from vllm.v1.worker.gpu.sample.bad_words import BadWordsState
17
+from vllm.v1.worker.gpu.states import RequestState
18
+
19
+DEVICE = torch.device("cuda")
20
+VOCAB_SIZE = 128
21
+
22
+# Committed tokens: prompt [5], output [10, 11]. Draft tokens: [12, 13].
23
+# The sampler passes input_ids gathered at logits_indices, so local position 0
24
+# holds the last committed token (11) and draft tokens start at position 1.
25
+PROMPT_LEN = 1
26
+COMMITTED = [5, 10, 11]
27
+INPUT_IDS = [11, 12, 13]
28
+LOCAL_POS = [0, 1, 2]
29
+
30
+
31
+def _make_state(bad_words_token_ids: list[list[int]]) -> tuple[BadWordsState, int]:
32
+ req_states = RequestState(
33
+ max_num_reqs=4,
34
+ max_model_len=64,
35
+ max_num_batched_tokens=16,
36
+ num_speculative_steps=4,
37
+ vocab_size=VOCAB_SIZE,
38
+ device=DEVICE,
39
+ )
40
+ req_states.add_request(
41
+ req_id="req",
42
+ prompt_len=PROMPT_LEN,
43
+ all_token_ids=COMMITTED,
44
+ num_computed_tokens=len(COMMITTED),
45
+ max_tokens=32,
46
+ )
47
+ req_states.apply_staged_writes()
48
+
49
+ req_idx = req_states.req_id_to_index["req"]
50
+ state = BadWordsState(req_states)
51
+ state.add_request(req_idx, SamplingParams(_bad_words_token_ids=bad_words_token_ids))
52
+ state.apply_staged_writes()
53
+ return state, req_idx
54
+
55
+
56
+def _apply(bad_words_token_ids: list[list[int]]) -> torch.Tensor:
57
+ state, req_idx = _make_state(bad_words_token_ids)
58
+ num_logits = len(INPUT_IDS)
59
+ logits = torch.zeros((num_logits, VOCAB_SIZE), device=DEVICE)
60
+ idx_mapping_np = np.array([req_idx], dtype=np.intp)
61
+ expanded_idx_mapping = torch.tensor(
62
+ [req_idx] * num_logits, dtype=torch.int32, device=DEVICE
63
+ )
64
+ state.apply_bad_words(
65
+ logits,
66
+ expanded_idx_mapping,
67
+ idx_mapping_np,
68
+ torch.tensor(INPUT_IDS, dtype=torch.int32, device=DEVICE),
69
+ torch.tensor(LOCAL_POS, dtype=torch.int32, device=DEVICE),
70
+ )
71
+ return logits.cpu()
72
+
73
+
74
+def test_v2_bad_words_prefix_inside_draft_tokens():
75
+ """A prefix matching entirely within the draft tokens must mask the bad
76
+ word's last token at the draft position that completes the prefix."""
77
+ out = _apply([[12, 13, 40]])
78
+ expected = torch.zeros_like(out)
79
+ expected[2, 40] = -float("inf")
80
+ torch.testing.assert_close(out, expected)
81
+
82
+
83
+def test_v2_bad_words_prefix_spanning_committed_and_draft_tokens():
84
+ """A prefix spanning the committed/draft boundary must mask at the row
85
+ where the prefix completes, not one draft position later."""
86
+ out = _apply([[11, 12, 30]])
87
+ expected = torch.zeros_like(out)
88
+ expected[1, 30] = -float("inf")
89
+ torch.testing.assert_close(out, expected)
90
+
91
+
92
+def test_v2_bad_words_no_spurious_match_from_last_committed_token():
93
+ """The last committed token must not be double-counted as the first draft
94
+ token; [11, 11] never occurs in output [10, 11] + drafts [12, 13]."""
95
+ out = _apply([[11, 11, 50]])
96
+ expected = torch.zeros_like(out)
97
+ torch.testing.assert_close(out, expected)
98
+
99
+
100
+def test_v2_bad_words_committed_prefix():
101
+ """Baseline: a fully committed prefix masks at the first row."""
102
+ out = _apply([[10, 11, 60]])
103
+ expected = torch.zeros_like(out)
104
+ expected[0, 60] = -float("inf")
105
+ torch.testing.assert_close(out, expected)
vllm/v1/worker/gpu/sample/bad_words.py
+3
-1
@@ -153,8 +153,10 @@ def _bad_words_kernel(
153
154
from_spec_input = actual_pos >= output_len
155
if from_spec_input:
156
+ # input_ids at local position 0 is the last committed token;
157
+ # draft tokens start at local position 1.
158
spec_offset = actual_pos - output_len
157
- actual = tl.load(input_ids_ptr + cur_req_first_pos + spec_offset)
159
+ actual = tl.load(input_ids_ptr + cur_req_first_pos + spec_offset + 1)
160
else:
161
actual = tl.load(output_base + actual_pos)
162