From a46f3eb232b8a74bab0aab02b6a070ebf337125f Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Tue, 23 Jun 2026 14:01:13 +0700 Subject: [PATCH] [Bugfix][Model Runner V2] Preserve all allowed_token_ids in the logit bias kernel (#46245) Signed-off-by: Ting Sun --- tests/v1/sample/test_sampling_params_e2e.py | 8 ++++++++ vllm/v1/worker/gpu/sample/logit_bias.py | 4 ++++ 2 files changed, 12 insertions(+) diff --git a/tests/v1/sample/test_sampling_params_e2e.py b/tests/v1/sample/test_sampling_params_e2e.py index fff953323f9..56b93ea1e01 100644 --- a/tests/v1/sample/test_sampling_params_e2e.py +++ b/tests/v1/sample/test_sampling_params_e2e.py @@ -152,6 +152,14 @@ def test_allowed_token_ids(llm): output = llm.generate(PROMPT, SamplingParams(allowed_token_ids=allowed_token_ids)) assert output[0].outputs[0].token_ids[-1] == TOKEN_ID + # Each single-token allowlist must force that token (kernel used to drop some). + for token_id in (1, 5, 100, 500, 2518, 9834, 31999): + output = llm.generate( + PROMPT, + SamplingParams(temperature=0, max_tokens=1, allowed_token_ids=[token_id]), + ) + assert output[0].outputs[0].token_ids[-1] == token_id + # Reject empty allowed_token_ids. with pytest.raises(ValueError): _ = llm.generate(PROMPT, SamplingParams(allowed_token_ids=[])) diff --git a/vllm/v1/worker/gpu/sample/logit_bias.py b/vllm/v1/worker/gpu/sample/logit_bias.py index 396f9f509c6..f3f7c29b3f3 100644 --- a/vllm/v1/worker/gpu/sample/logit_bias.py +++ b/vllm/v1/worker/gpu/sample/logit_bias.py @@ -189,6 +189,8 @@ def _bias_kernel( logits_ptr + token_idx * logits_stride + allowed_token_ids, mask=mask ) + tl.debug_barrier() # save must read original logits before the -inf overwrite + # Set logits to -inf for all tokens. for i in range(0, vocab_size, LOGITS_BLOCK_SIZE): offset = i + tl.arange(0, LOGITS_BLOCK_SIZE) @@ -198,6 +200,8 @@ def _bias_kernel( mask=offset < vocab_size, ) + tl.debug_barrier() # -inf overwrite must finish before restoring saved logits + # Restore logits for allowed token IDs. tl.store( logits_ptr + token_idx * logits_stride + allowed_token_ids,