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 7, 07:45 UTC
[cohere][fix][spec-decode]: fix crash when allowed_token_ids is set without penalties (#35654)
Signed-off-by: kkt-cohere <komal@cohere.com>
Komal Kumar Teru committed
Mar 4, 2026 at 12:50 UTC
9e0f44bec449df17d30ed9abef7aeedc059ddfde
1 file changed
+1
-1
vllm/v1/sample/rejection_sampler.py
+1
-1
@@ -271,7 +271,7 @@ class RejectionSampler(nn.Module):
271
272
# Calculate indices of target logits.
273
if sampling_metadata.allowed_token_ids_mask is not None or has_penalties:
274
- num_requests = len(sampling_metadata.output_token_ids)
274
+ num_requests = len(metadata.num_draft_tokens)
275
num_draft_tokens = torch.tensor(metadata.num_draft_tokens, device="cpu")
276
original_indices = torch.arange(num_requests, device="cpu")
277
repeat_indices_cpu = original_indices.repeat_interleave(num_draft_tokens)