[Speculative decoding] [Bugfix] Fix overallocation in ngram + spec lo…

…gprobs (#4672)
vllm-project · May 8, 2024 · 8b9241b · 8b9241b
1 parent f942efb
commit 8b9241b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vllm/spec_decode/ngram_worker.py b/vllm/spec_decode/ngram_worker.py
@@ -138,7 +138,7 @@ def sampler_output(
                 SamplerOutput(
                     outputs=None,
                     sampled_token_probs=token_probs[i],
-                    logprobs=token_logprobs,
+                    logprobs=token_logprobs[i],
                     sampled_token_ids=token_ids[i],
                 ))
         return outputs, False