[TRTLLM-14510][feat] serve: multi-process HTTP frontends on the classic IPC executor path

tensorrt_llm/executor/base_worker.py

@@ -49,6 +49,7 @@
 from .result import (GenerationResult, LogProbsResult, ResponseWrapper,
                      compute_logprobs, get_metrics_dict)
 from .utils import (ErrorResponse, IntraProcessQueue, RequestError,
+                    bucket_responses_by_frontend, frontend_lane_index,
                     is_llm_response)
 
 if TYPE_CHECKING:
@@ -118,6 +119,9 @@ def __init__(
         self.engine = None
         self.result_queue: Optional[IpcQueue] = None
         self.postproc_queues: Optional[List[IpcQueue]] = None
+        # Multi-frontend serving: one result lane per frontend process,
+        # selected by the frontend id in client_id's top bits.
+        self.frontend_result_queues: Optional[List[IpcQueue]] = None
         self.rank = mpi_rank()
         self.global_rank = global_mpi_rank()
         # mapping: client_id -> GenerationResult
@@ -272,13 +276,21 @@ def fetch_kv_cache_events(self) -> list:
     def set_result_queue(self, queue):
         """In multi-gpu mode, result_queue will be set here to communicate between the proxy and the worker 0 process."""
         assert self.postproc_queues is None
+        assert self.frontend_result_queues is None
         self.result_queue = queue
 
     def set_postproc_queues(self, queues: List["IpcQueue"]):
         """ Set the IPC queues for feeding post-processing processes. """
         assert self.result_queue is None
+        assert self.frontend_result_queues is None
         self.postproc_queues = queues
 
+    def set_frontend_result_queues(self, queues: List["IpcQueue"]):
+        """Multi-frontend serving: one result lane per frontend process."""
+        assert self.result_queue is None
+        assert self.postproc_queues is None
+        self.frontend_result_queues = queues
+
     def _set_iteration_result_queue(self, it_result_queue: IterationResultQueue,
                                     queue: Union[Queue, FusedIpcQueue,
                                                  IntraProcessQueue]):
@@ -1087,20 +1099,20 @@ def responses_handler(self, responses: List[tllm.Response]):
         HandlerKind = AwaitResponseHelper.HandlerKind
 
         if self.handler_kind is HandlerKind.unknown:
-            if not (self.worker.result_queue is not None
-                    or self.worker.postproc_queues is not None):
+            has_ipc_queues = (self.worker.result_queue is not None
+                              or self.worker.postproc_queues is not None
+                              or self.worker.frontend_result_queues is not None)
+            if not has_ipc_queues:
                 logger_debug(f"creating await_response helper for Worker\n",
                              color="yellow")
                 # When ExecutorBindingWorker is used in the main process
                 # aka the single process mode
                 self.handler_kind = HandlerKind.single_process_worker
-            elif self.worker.result_queue is not None or self.worker.postproc_queues is not None:
+            else:
                 # The ExecutorBindingProxy is used
                 logger_debug(f"creating await_response helper for IPC\n",
                              color="yellow")
                 self.handler_kind = HandlerKind.ipc_batched
-            else:
-                raise NotImplementedError
 
         match self.handler_kind:
             case HandlerKind.single_process_worker:
@@ -1279,7 +1291,13 @@ def handle_for_ipc_batched(self, responses: List[tllm.Response]) -> None:
                     self.worker.postproc_queues[wid].put(batch)
 
         if rsp_batch:
-            self.worker.result_queue.put(rsp_batch)
+            if (lanes := self.worker.frontend_result_queues) is not None:
+                for frontend_id, sub_batch in enumerate(
+                        bucket_responses_by_frontend(rsp_batch, len(lanes))):
+                    if sub_batch:
+                        lanes[frontend_id].put(sub_batch)
+            else:
+                self.worker.result_queue.put(rsp_batch)
 
 
 def _get_params_for_first_rsp(
@@ -1393,7 +1411,16 @@ def _send_rsp(
         rsp_batch: Optional[List[tllm.Response]] = None):
     # if postproc_batches is set, append to batch instead of putting to IpcQueue
 
-    if worker.result_queue is not None:
+    if worker.frontend_result_queues is not None:
+        # Route to the origin frontend's result lane; None/out-of-range ids
+        # fall back to lane 0 (see frontend_lane_index).
+        if rsp_batch is not None:
+            rsp_batch.append(response)
+        else:
+            lanes = worker.frontend_result_queues
+            lanes[frontend_lane_index(response.client_id,
+                                      len(lanes))].put(response)
+    elif worker.result_queue is not None:
         if rsp_batch is not None:
             rsp_batch.append(response)
         else:

tensorrt_llm/executor/executor.py

@@ -1,6 +1,8 @@
 import atexit
 import faulthandler
+import json
 import multiprocessing
+import os
 import platform
 import signal
 import traceback
@@ -560,6 +562,35 @@ def create(
                 f"Using {postproc_worker_config.num_postprocess_workers} postprocess parallel processes.\n",
                 "green")
 
+        # Multi-frontend serving: attach to an already-running executor
+        # instead of launching one (set by trtllm-serve for attached
+        # frontend processes).
+        attach_env = os.getenv("TLLM_EXECUTOR_ATTACH_INFO")
+        if attach_env:
+            attach_info = json.loads(attach_env)
+            # Consumed: it carries the HMAC keys and must not leak into
+            # descendant processes.
+            os.environ.pop("TLLM_EXECUTOR_ATTACH_INFO", None)
+            if attach_info.get("mode") != "classic":
+                raise ValueError(
+                    "TLLM_EXECUTOR_ATTACH_INFO only supports the classic IPC "
+                    f"executor path, got mode={attach_info.get('mode')!r}")
+            from .proxy import GenerationExecutorFrontendProxy
+            frontend_id_env = os.getenv("TLLM_EXECUTOR_FRONTEND_ID")
+            if frontend_id_env is None:
+                raise ValueError(
+                    "TLLM_EXECUTOR_ATTACH_INFO is set but "
+                    "TLLM_EXECUTOR_FRONTEND_ID is not; both are set together "
+                    "by trtllm-serve when spawning attached frontends.")
+            frontend_id = int(frontend_id_env)
+            logger.info(f"Attaching executor frontend {frontend_id} to the "
+                        "running classic IPC worker")
+            return GenerationExecutorFrontendProxy(
+                attach_info,
+                frontend_id=frontend_id,
+                postproc_worker_config=postproc_worker_config,
+                is_llm_executor=is_llm_executor)
+
         worker_kwargs = {
             "engine": engine,
             "executor_config": executor_config,

tensorrt_llm/executor/postproc_worker.py

@@ -15,7 +15,7 @@
 from ..sampling_params import SamplingParams
 from .ipc import ZeroMqQueue
 from .postprocessor_hook import load_post_processor_hook
-from .utils import ErrorResponse, is_llm_response
+from .utils import ErrorResponse, bucket_responses_by_frontend, is_llm_response
 
 if TYPE_CHECKING:
     from ..disaggregated_params import DisaggregatedParams
@@ -86,7 +86,7 @@ class Output(NamedTuple):
     def __init__(
         self,
         pull_pipe_addr: tuple[str, Optional[bytes]],
-        push_pipe_addr: tuple[str, Optional[bytes]],
+        push_pipe_addrs: List[tuple[str, Optional[bytes]]],
         tokenizer_dir: str,
         record_creator: Callable[
             ["PostprocWorker.Input", TransformersTokenizer], Any],
@@ -95,7 +95,9 @@ def __init__(
         '''
         Args:
             pull_pipe_addr (tuple[str, Optional[bytes]]): The address and HMAC key of the input IPC.
-            push_pipe_addr (tuple[str, Optional[bytes]]): The address and HMAC key of the output IPC.
+            push_pipe_addrs: The addresses and HMAC keys of the output IPC
+                lanes, one per frontend (a single-element list in
+                single-frontend mode).
             tokenizer_dir (str): The directory to load tokenizer.
             record_creator (Callable[["ResponsePostprocessWorker.Input"], Any]): A creator for creating a record for a request.
             result_handler (Optional[Callable[[GenerationResultBase], Any]]): A callback handles the final result.
@@ -108,11 +110,14 @@ def __init__(
                                       is_async=True,
                                       is_server=False,
                                       name="postprocess_pull_pipe")
-        self._push_pipe = ZeroMqQueue(address=push_pipe_addr,
-                                      is_async=True,
-                                      is_server=False,
-                                      socket_type=zmq.PUSH,
-                                      name="postprocess_push_pipe")
+        self._push_pipes = [
+            ZeroMqQueue(address=addr,
+                        is_async=True,
+                        is_server=False,
+                        socket_type=zmq.PUSH,
+                        name=f"postprocess_push_pipe_{i}")
+            for i, addr in enumerate(push_pipe_addrs)
+        ]
         self._to_stop = asyncio.Event()
 
         self._q = deque()
@@ -192,11 +197,19 @@ async def _batched_put(self):
         ''' Batched IPC send. '''
         async for batch in self._mainloop():
             if batch is None:
-                # notify dispatch_result corountine to quit
-                await self._push_pipe.put_async(None)
+                # notify the dispatch_result coroutine in every frontend to
+                # quit
+                for pipe in self._push_pipes:
+                    await pipe.put_async(None)
                 break
             assert isinstance(batch, list)
-            await self._push_pipe.put_async(batch)
+            if len(self._push_pipes) == 1:
+                await self._push_pipes[0].put_async(batch)
+                continue
+            for frontend_id, sub_batch in enumerate(
+                    bucket_responses_by_frontend(batch, len(self._push_pipes))):
+                if sub_batch:
+                    await self._push_pipes[frontend_id].put_async(sub_batch)
 
     async def _mainloop(self):
         ''' The loop for handle_response and keep producing outputs. '''
@@ -289,13 +302,13 @@ async def main():
 
 @print_traceback_on_error
 def postproc_worker_main(feedin_ipc_addr: tuple[str, Optional[bytes]],
-                         feedout_ipc_addr: tuple[str, Optional[bytes]],
+                         feedout_ipc_addrs: List[tuple[str, Optional[bytes]]],
                          tokenizer_dir: str,
                          record_creator: Callable,
                          post_processor_hook: Optional[str] = None):
     # Pass the hook import path; PostprocWorker builds it once.
     worker = PostprocWorker(feedin_ipc_addr,
-                            feedout_ipc_addr,
+                            feedout_ipc_addrs,
                             tokenizer_dir=tokenizer_dir,
                             record_creator=record_creator,
                             post_processor_hook=post_processor_hook)