[Bugfix] Sets `is_first_step_output` for TPUModelRunner (#9202)

c6cf9295 · Allen Wang · GitHub · de9fb4be · c6cf9295
Unverified Commit c6cf9295 authored Oct 11, 2024 by Allen Wang Committed by GitHub Oct 11, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

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

No files found.
--- 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]]