[fix] pp+mtp bs 1 correctness

df49168d · zhuwenwen · f8cf43ae · df49168d
Commit df49168d authored Sep 26, 2025 by zhuwenwen
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

vllm/v1/worker/gpu_model_runner.py vllm/v1/worker/gpu_model_runner.py +2 -2

No files found.
--- a/vllm/v1/worker/gpu_model_runner.py
+++ b/vllm/v1/worker/gpu_model_runner.py
@@ -642,10 +642,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
            if not is_last_rank:
                # Add new_token_ids to token_ids_cpu.
                start_token_index = num_computed_tokens
-                end_token_index = num_computed_tokens + len(new_token_ids)
+                end_token_index = num_computed_tokens + 1
                self.input_batch.token_ids_cpu[
                    req_index,
-                    start_token_index:end_token_index] = new_token_ids
+                    start_token_index:end_token_index] = new_token_ids[-1]
                self.input_batch.num_tokens_no_spec[
                    req_index] = end_token_index
                self.input_batch.num_tokens[req_index] = end_token_index