From c6cf9295e1dad2aeffbce1d92682971df9f71ddf Mon Sep 17 00:00:00 2001 From: Allen Wang Date: Fri, 11 Oct 2024 15:28:10 -0500 Subject: [PATCH] [Bugfix] Sets `is_first_step_output` for TPUModelRunner (#9202) --- vllm/worker/tpu_model_runner.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/worker/tpu_model_runner.py b/vllm/worker/tpu_model_runner.py index f26d1c8c..c13e95f6 100644 --- a/vllm/worker/tpu_model_runner.py +++ b/vllm/worker/tpu_model_runner.py @@ -541,7 +541,8 @@ class TPUModelRunner(ModelRunnerBase[ModelInputForTPU]): seq_group_metadata_list=ctx.seq_group_metadata_list, scheduler_outputs=ctx.scheduler_outputs, is_async=False, - is_last_step=False) + is_last_step=False, + is_first_step_output=i == 0) model_input.async_callback() if use_async_out_proc: return [sampler_outputs[-1]]