From 9e0f44bec449df17d30ed9abef7aeedc059ddfde Mon Sep 17 00:00:00 2001 From: Komal Kumar Teru <162363718+kkt-cohere@users.noreply.github.com> Date: Wed, 4 Mar 2026 12:50:15 +0530 Subject: [PATCH] [cohere][fix][spec-decode]: fix crash when allowed_token_ids is set without penalties (#35654) Signed-off-by: kkt-cohere --- vllm/v1/sample/rejection_sampler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py index 278d421eb91..d3e8573458b 100644 --- a/vllm/v1/sample/rejection_sampler.py +++ b/vllm/v1/sample/rejection_sampler.py @@ -271,7 +271,7 @@ class RejectionSampler(nn.Module): # Calculate indices of target logits. if sampling_metadata.allowed_token_ids_mask is not None or has_penalties: - num_requests = len(sampling_metadata.output_token_ids) + num_requests = len(metadata.num_draft_tokens) num_draft_tokens = torch.tensor(metadata.num_draft_tokens, device="cpu") original_indices = torch.arange(num_requests, device="cpu") repeat_indices_cpu = original_indices.repeat_interleave(num_draft_tokens)