[Bugfix] Sets is_first_step_output for TPUModelRunner (#9202)

2024-10-11 15:28:10 -05:00 · 2024-10-11 15:28:10 -05:00 · c6cf9295e1
commit c6cf9295e1
parent de9fb4bef8
1 changed files with 2 additions and 1 deletions
--- 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]]