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 14, 05:30 UTC
[bugfix] [spec-decoding] fix data race in sample_recovered_tokens_kernel (vLLM v1) (#23829)
Signed-off-by: He-Jingkai <he-jingkai@outlook.com>
Jingkai He committed
Aug 29, 2025 at 03:05 UTC
57d4ede520b6071341ebd310c0ddd4c6f4d54917
1 file changed
+2
-15
vllm/v1/sample/rejection_sampler.py
+2
-15
@@ -598,17 +598,10 @@ def sample_recovered_tokens_kernel(
598
vocab_offset = tl.arange(0, PADDED_VOCAB_SIZE)
599
if NO_DRAFT_PROBS:
600
draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
601
- orig_prob = tl.load(target_probs_ptr + (start_idx + pos) * vocab_size +
602
- draft_token_id)
603
- # Temporarily zero out the probability of the draft token.
604
- # This is essentially the same as target_prob - draft_prob, except that
605
- # n-gram does not have draft_prob. We regard it as 1.
606
- tl.store(
607
- target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id,
608
- 0)
601
prob = tl.load(target_probs_ptr + (start_idx + pos) * vocab_size +
602
vocab_offset,
611
- mask=vocab_offset < vocab_size,
603
+ mask=((vocab_offset < vocab_size) &
604
+ (vocab_offset != draft_token_id)),
605
other=0)
606
else:
607
draft_prob = tl.load(draft_probs_ptr + (start_idx + pos) * vocab_size +
@@ -628,9 +621,3 @@ def sample_recovered_tokens_kernel(
621
other=float("-inf"))
622
recovered_id = tl.argmax(prob / q, axis=-1)
623
tl.store(output_token_ids_ptr + start_idx + pos, recovered_id)
631
-
632
- if NO_DRAFT_PROBS:
633
- # Restore the original probability.
634
- tl.store(
635
- target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id,
636
- orig_prob)