From 343f65234bb2f6a0c42bf398e80a1a79b9739aaa Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Wed, 15 Apr 2026 18:09:11 -0700 Subject: [PATCH] =?UTF-8?q?[Model=20Runner=20V2][BugFix]=20fix=20num=5Fsam?= =?UTF-8?q?pled=20dtype=20for=20probabilistic=20rej=E2=80=A6=20(#39951)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Giancarlo Delfin --- .../gpu/spec_decode/probabilistic_rejection_sampler_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/spec_decode/probabilistic_rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/probabilistic_rejection_sampler_utils.py index 22f5d9df0d..2feaae2453 100644 --- a/vllm/v1/worker/gpu/spec_decode/probabilistic_rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/probabilistic_rejection_sampler_utils.py @@ -527,7 +527,7 @@ def probabilistic_rejection_sample( sampled = draft_sampled.new_empty( num_reqs, num_speculative_steps + 1, dtype=torch.int64 ) - num_sampled = sampled.new_empty(num_reqs) + num_sampled = sampled.new_empty(num_reqs, dtype=torch.int32) target_rejected_logsumexp = target_logits.new_empty(num_reqs, dtype=torch.float32) draft_rejected_logsumexp = target_logits.new_empty(num_reqs, dtype=torch.float32) _probabilistic_rejection_kernel[(num_reqs,)](