From 8b9241be3a0020724e145bf600d9710b3d59b167 Mon Sep 17 00:00:00 2001
From: Cade Daniel <edacih@gmail.com>
Date: Wed, 8 May 2024 16:24:46 -0700
Subject: [PATCH] [Speculative decoding] [Bugfix] Fix overallocation in ngram +
 spec logprobs (#4672)

---
 vllm/spec_decode/ngram_worker.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/vllm/spec_decode/ngram_worker.py b/vllm/spec_decode/ngram_worker.py
index fed8be42054a5..f18f9387f5b23 100644
--- 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