Skip to content

Commit c6cf929

Browse files
author
Allen Wang
authored
[Bugfix] Sets is_first_step_output for TPUModelRunner (#9202)
1 parent de9fb4b commit c6cf929

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

vllm/worker/tpu_model_runner.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -541,7 +541,8 @@ def execute_model(
541541
seq_group_metadata_list=ctx.seq_group_metadata_list,
542542
scheduler_outputs=ctx.scheduler_outputs,
543543
is_async=False,
544-
is_last_step=False)
544+
is_last_step=False,
545+
is_first_step_output=i == 0)
545546
model_input.async_callback()
546547
if use_async_out_proc:
547548
return [sampler_outputs[-1]]

0 commit comments

Comments
 (0)