iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Model Runner V2][BugFix] fix num_sampled dtype for probabilistic rej… (#39951)

Browse Source 
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
This commit is contained in:
Giancarlo Delfin
2026-04-15 18:09:11 -07:00
committed by GitHub
parent 19fa90ed0d
commit 343f65234b
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/v1/worker/gpu/spec_decode/probabilistic_rejection_sampler_utils.py
+1 -1
View File
@@ -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,)](