Fix: fuse message not aligned on different processes (NVIDIA#3067)

Signed-off-by: Kaiyu Xie <[email protected]>

Author: kaiyux

tensorrt_llm/executor/postproc_worker.py

@@ -158,7 +158,9 @@ async def _mainloop(self):
 
         async def handle_single_input(inp: PostprocWorker.Input,
                                       batch: List[PostprocWorker.Output]):
-            assert isinstance(inp, PostprocWorker.Input)
+            assert isinstance(
+                inp, PostprocWorker.Input
+            ), f"Expect PostprocWorker.Input, got {type(inp)}."
             client_id = inp.rsp.client_id
             is_final = inp.rsp.result.is_final if isinstance(
                 inp.rsp, tllm.Response) else True

tensorrt_llm/executor/worker.py

@@ -552,7 +552,7 @@ def worker_main(
             # processes, each one is a PAIR zmq socket
             result_queues = [
                 FusedIpcQueue(is_server=True,
-                              fuse_message=True,
+                              fuse_message=not BATCH_RESP_IN_AWAIT,
                               name=f"postprocess_{i}_feedin_queue")
                 for i in range(postproc_worker_config.num_postprocess_workers)
             ]
@@ -803,7 +803,8 @@ def handle_for_ipc_batched(self, responses: List[tllm.Response]) -> None:
 
         if postproc_batches:
             for wid, batch in enumerate(postproc_batches):
-                self.worker.postproc_queues[wid].put(batch)
+                if batch:
+                    self.worker.postproc_queues[wid].put(batch)
 
         if rsp_batch:
             self.worker.result_queue.put(rsp_batch)