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 5, 02:15 UTC
[Model Runner V2][BugFix] fix num_sampled dtype for probabilistic rej… (#39951)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
Apr 15, 2026 at 18:09 UTC
343f65234bb2f6a0c42bf398e80a1a79b9739aaa
1 file changed
+1
-1
vllm/v1/worker/gpu/spec_decode/probabilistic_rejection_sampler_utils.py
+1
-1
@@ -527,7 +527,7 @@ def probabilistic_rejection_sample(
527
sampled = draft_sampled.new_empty(
528
num_reqs, num_speculative_steps + 1, dtype=torch.int64
529
)
530
- num_sampled = sampled.new_empty(num_reqs)
530
+ num_sampled = sampled.new_empty(num_reqs, dtype=torch.int32)
531
target_rejected_logsumexp = target_logits.new_empty(num_reqs, dtype=torch.float32)
532
draft_rejected_logsumexp = target_logits.new_empty(num_reqs, dtype=torch.float32)
533
_probabilistic_rejection_kernel[(num_reqs,)](