squall/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix] Sets is_first_step_output for TPUModelRunner (#9202)

Browse Source
This commit is contained in:
Allen Wang 2024-10-11 15:28:10 -05:00 committed by GitHub
parent de9fb4bef8
commit c6cf9295e1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
vllm/worker/tpu_model_runner.py
View File

@ -541,7 +541,8 @@ class TPUModelRunner(ModelRunnerBase[ModelInputForTPU]):
seq_group_metadata_list=ctx.seq_group_metadata_list, seq_group_metadata_list=ctx.seq_group_metadata_list,
scheduler_outputs=ctx.scheduler_outputs, scheduler_outputs=ctx.scheduler_outputs,
is_async=False, is_async=False,
is_last_step=False) is_last_step=False,
is_first_step_output=i == 0)
model_input.async_callback() model_input.async_callback()
if use_async_out_proc: if use_async_out_proc:
return [sampler_outputs[-1]] return [sampler_outputs[-1]]