[V1][Spec Decode] Ngram Spec Decode (#12193)

Signed-off-by: LiuXiaoxuanPKU <lilyliupku@gmail.com>

[V1][Spec Decode] Ngram Spec Decode (#12193)
Signed-off-by: LiuXiaoxuanPKU <lilyliupku@gmail.com>
80f63a39 · Lily Liu · GitHub · 367cb8ce · 80f63a39
Unverified Commit 80f63a39 authored Feb 15, 2025 by Lily Liu Committed by GitHub Feb 15, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

vllm/v1/worker/tpu_model_runner.py vllm/v1/worker/tpu_model_runner.py +1 -1

No files found.
--- a/vllm/v1/worker/tpu_model_runner.py
+++ b/vllm/v1/worker/tpu_model_runner.py
@@ -695,7 +695,7 @@ class TPUModelRunner:
        model_runner_output = ModelRunnerOutput(
            req_ids=all_req_ids,
            req_id_to_index=self.input_batch.req_id_to_index,
-            sampled_token_ids=sampled_token_ids,
+            sampled_token_ids=[[token_id] for token_id in sampled_token_ids],
            logprobs=None,
            prompt_logprobs_dict=prompt_logprobs_dict,  # type: ignore[arg-type]
        )