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 @@ 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