[Bugfix][Model Runner V2] Preserve all allowed_token_ids in the logit bias kernel (#46245)
Signed-off-by: Ting Sun <suntcrick@gmail.com>
Ting SUN committed
Jun 23, 2026 at 14:01 UTC
a46f3eb232b8a74bab0aab02b6a070ebf337125f
2 files changed
+12
tests/v1/sample/test_sampling_params_e2e.py
+8
@@ -152,6 +152,14 @@ def test_allowed_token_ids(llm):
152
output = llm.generate(PROMPT, SamplingParams(allowed_token_ids=allowed_token_ids))
153
assert output[0].outputs[0].token_ids[-1] == TOKEN_ID
154
155
+ # Each single-token allowlist must force that token (kernel used to drop some).
156
+ for token_id in (1, 5, 100, 500, 2518, 9834, 31999):
157
+ output = llm.generate(
158
+ PROMPT,
159
+ SamplingParams(temperature=0, max_tokens=1, allowed_token_ids=[token_id]),
160
+ )
161
+ assert output[0].outputs[0].token_ids[-1] == token_id
162
+
163
# Reject empty allowed_token_ids.
164
with pytest.raises(ValueError):
165
_ = llm.generate(PROMPT, SamplingParams(allowed_token_ids=[]))
vllm/v1/worker/gpu/sample/logit_bias.py
+4
@@ -189,6 +189,8 @@ def _bias_kernel(
189
logits_ptr + token_idx * logits_stride + allowed_token_ids, mask=mask
190
)
191
192
+ tl.debug_barrier() # save must read original logits before the -inf overwrite
193
+
194
# Set logits to -inf for all tokens.
195
for i in range(0, vocab_size, LOGITS_BLOCK_SIZE):
196
offset = i + tl.arange(0, LOGITS_BLOCK_SIZE)
@@ -198,6 +200,8 @@ def _bias_kernel(
200
mask=offset < vocab_size,
201
)
202
203
+ tl.debug_barrier() # -inf overwrite must finish before restoring saved logits
204
+
205
# Restore logits for allowed token IDs.
206
tl.store(
207
logits_ptr + token_idx * logits_stride + allowed_token_ids,