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 Decode] Use float64 for uniform_probs (#23803)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Aug 28, 2025 at 05:26 UTC
a3432f18fdd85eb18e29fc32327507fe1063ad57
2 files changed
+7
-2
examples/offline_inference/spec_decode.py
+1
-1
@@ -138,7 +138,7 @@ def main():
138
sampling_params = SamplingParams(temperature=args.temp, max_tokens=args.output_len)
139
if not args.custom_mm_prompts:
140
outputs = llm.generate(
141
- TokensPrompt(prompt_token_ids=prompt_ids),
141
+ [TokensPrompt(prompt_token_ids=x) for x in prompt_ids],
142
sampling_params=sampling_params,
143
)
144
else:
vllm/v1/sample/rejection_sampler.py
+6
-1
@@ -365,9 +365,14 @@ def generate_uniform_probs(
365
A tensor of shape `(num_tokens, )` containing uniform
366
random values in the range [0, 1).
367
"""
368
+ # NOTE(woosuk): We deliberately use float64 instead of float32 here
369
+ # because when using float32, there's a non-negligible chance that
370
+ # uniform_prob is sampled to be exact 0.0 as reported in
371
+ # https://github.com/pytorch/pytorch/issues/16706. Using float64
372
+ # mitigates the issue.
373
uniform_probs = torch.rand(
374
(num_tokens, ),
370
- dtype=torch.float32,
375
+ dtype=torch.float64,
376
device=device,
377
)
378
start_idx = 0