Skip to content

GitLab

Unverified Commit 2eedede8 authored by Megha Agarwal's avatar Megha Agarwal Committed by GitHub
Browse files

[Core] Asynchronous Output Processor (#7049) 


Co-authored-by: default avatarAlexander Matveev <alexm@neuralmagic.com>
parent 015e6cc2
Show whitespace changes
Inline Side-by-side
Showing with 7 additions and 1 deletion
vllm/worker/worker_base.py
View file @ 2eedede8
...@@ -263,6 +263,12 @@ class LocalOrDistributedWorkerBase(WorkerBase): ...@@ -263,6 +263,12 @@ class LocalOrDistributedWorkerBase(WorkerBase):
broadcast_data.update(kwargs) broadcast_data.update(kwargs)
broadcast_tensor_dict(broadcast_data, src=0) broadcast_tensor_dict(broadcast_data, src=0)
if execute_model_req.output_proc_callback_fn:
model_input = dataclasses.replace( # type: ignore
model_input,
output_proc_callback_fn=execute_model_req.
output_proc_callback_fn)
return model_input, worker_input, kwargs return model_input, worker_input, kwargs
def prepare_input( def prepare_input(
...@@ -289,7 +295,7 @@ class LocalOrDistributedWorkerBase(WorkerBase): ...@@ -289,7 +295,7 @@ class LocalOrDistributedWorkerBase(WorkerBase):
def execute_model( def execute_model(
self, self,
execute_model_req: Optional[ExecuteModelRequest] = None execute_model_req: Optional[ExecuteModelRequest] = None,
) -> Optional[List[SamplerOutput]]: ) -> Optional[List[SamplerOutput]]:
"""Executes at least one model step on the given sequences, unless no """Executes at least one model step on the given sequences, unless no
sequences are provided.""" sequences are provided."""
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment