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

8b9241be · Cade Daniel · GitHub · f942efb5 · 8b9241be
Unverified Commit 8b9241be authored May 08, 2024 by Cade Daniel Committed by GitHub May 08, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

vllm/spec_decode/ngram_worker.py vllm/spec_decode/ngram_worker.py +1 -1

No files found.
--- 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