[Bugfix] Sets is_first_step_output for TPUModelRunner (vllm-project…

…#9202) Signed-off-by: Maxime Fournioux <[email protected]>
mfournioux · Nov 20, 2024 · 7b8e2b1 · 7b8e2b1
1 parent 53a45e1
commit 7b8e2b1
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/vllm/worker/tpu_model_runner.py b/vllm/worker/tpu_model_runner.py
@@ -541,7 +541,8 @@ def execute_model(
                         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]]