Skip to content

Commit 483a60c

Browse files
Allen Wangsumitd2
authored andcommitted
[Bugfix] Sets is_first_step_output for TPUModelRunner (vllm-project#9202)
Signed-off-by: Sumit Dubey <sumit.dubey2@ibm.com>
1 parent c76b33f commit 483a60c

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)