[Speculative decoding] [Bugfix] Fix overallocation in ngram + spec logprobs (#4672)

2024-05-08 16:24:46 -07:00 · 2024-05-08 16:24:46 -07:00 · 8b9241be3a
commit 8b9241be3a
parent f942efb5a3
1 changed files with 1 additions and 1 deletions
--- a/vllm/spec_decode/ngram_worker.py
+++ b/vllm/spec_decode/ngram_worker.py
@ -138,7 +138,7 @@ class NGramWorker(LoraNotSupportedWorkerBase):
                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