[BugFix][Spec Decode] Use float64 for uniform_probs (#23803)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
2025-08-28 05:26:45 -07:00
parent 67cee40da0
commit a3432f18fd
2 changed files with 7 additions and 2 deletions
--- a/examples/offline_inference/spec_decode.py
+++ b/examples/offline_inference/spec_decode.py
@@ -138,7 +138,7 @@ def main():
    sampling_params = SamplingParams(temperature=args.temp, max_tokens=args.output_len)
    if not args.custom_mm_prompts:
        outputs = llm.generate(
-            TokensPrompt(prompt_token_ids=prompt_ids),
+            [TokensPrompt(prompt_token_ids=x) for x in prompt_ids],
            sampling_params=sampling_params,
        )
    else: