Parallelize output socket IO on client side

Signed-off-by: Nick Hill <nhill@redhat.com>

vllm/v1/engine/async_llm.py

@@ -244,9 +244,10 @@ async def _run_output_handler(self):
                 # event loop for too long.
                 num_outputs = len(outputs.outputs)
                 if num_outputs <= OUTPUT_PROCESSING_CHUNK_SIZE:
-                    slices = (outputs.outputs,)
+                    slices = (outputs.outputs, )
                 else:
-                    slices = np.array_split(outputs.outputs,
+                    slices = np.array_split(
+                        outputs.outputs,
                         math.ceil(num_outputs / OUTPUT_PROCESSING_CHUNK_SIZE))
 
                 iteration_stats = None

vllm/v1/engine/core_client.py

@@ -1,8 +1,9 @@
+import asyncio
 import os
 import signal
 import weakref
 from abc import ABC, abstractmethod
-from typing import List, Type
+from typing import List, Optional, Type
 
 import msgspec
 import zmq
@@ -242,10 +243,23 @@ def __init__(self, vllm_config: VllmConfig,
             log_stats=True,
         )
 
+        self.outputs_queue: Optional[asyncio.Queue[bytes]] = None
+        self.queue_task: Optional[asyncio.Task] = None
+
     async def get_output_async(self) -> EngineCoreOutputs:
+        if self.outputs_queue is None:
+            # Perform IO in separate task to parallelize as much as possible
+            self.outputs_queue: asyncio.Queue[bytes] = asyncio.Queue()
+
+            async def process_outputs_socket():
+                while True:
+                    (frame, ) = await self.output_socket.recv_multipart(
+                        copy=False)
+                    self.outputs_queue.put_nowait(frame.buffer)
+
+            self.queue_task = asyncio.create_task(process_outputs_socket())
 
-        frames = await self.output_socket.recv_multipart(copy=False)
-        return self.decoder.decode(frames[0].buffer)
+        return self.decoder.decode(await self.outputs_queue.get())
 
     async def _send_input(self, request_type: EngineCoreRequestType,
                           request: EngineCoreRequestUnion) -> None: