format

vllm-project · Aug 28, 2024 · e26b18b · e26b18b
1 parent c29e4da
commit e26b18b
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/vllm/worker/multi_step_model_runner.py b/vllm/worker/multi_step_model_runner.py
@@ -313,7 +313,6 @@ def execute_model(
             model_input = self._advance_step(
                 model_input, model_input.cached_outputs[-1].sampler_output)
 
-        output_proc_callback = None
         if (frozen_model_input.use_async_and_multi_step
                 and model_input.is_first_multi_step):
             output_proc_callback = frozen_model_input.async_callback
@@ -376,8 +375,10 @@ def execute_model(
 
         # Pythonize the output and block if needed since it is the last step
         if model_input.is_last_step:
-            outputs = self._final_process_outputs(model_input,
-                                                  output_proc_callback)
+            outputs = self._final_process_outputs(
+                model_input,
+                model_input.frozen_model_input.async_callback.  # type: ignore
+                keywords["output_proc_callback"])  # type: ignore
             return outputs
 
         # should be [SamplerOutput]