[TRTLLM-17273][chore] Choose the trtllm-serve post-processing path once per server

tensorrt_llm/executor/result.py

@@ -128,6 +128,7 @@ class CompletionOutput:
         logprobs_diff (TokenLogprobs | SimpleTokenLogprobs): Logprobs of newly generated tokens.
         text_diff (str): Newly generated tokens.
         routed_experts (Optional[torch.Tensor]): Per-token pre-EPLB logical top-k MoE expert ids (Router Replay / R3), or None when not requested.
+        postprocess_result (Any): What the post-processor returned for this output when post-processing ran in a postproc worker; None otherwise.
     """
     index: int
     text: str = ""
@@ -216,6 +217,14 @@ def token_ids_diff(self) -> List[int]:
     def logprobs_diff(self) -> TokenLogprobs | SimpleTokenLogprobs:
         return self.logprobs[self._last_logprobs_len:]
 
+    @property
+    def postprocess_result(self) -> Any:
+        """The post-processor's return value stored by a postproc worker.
+
+        None when post-processing did not run in a postproc worker.
+        """
+        return self._postprocess_result
+
 
 class GenerationResultBase:
     """This holds the core logic of the GenerationResult class."""

tensorrt_llm/serve/openai_server.py

@@ -118,11 +118,11 @@
                                              build_request_metrics_record)
 from tensorrt_llm.serve.postprocess_handlers import (
     ChatCompletionPostprocArgs, ChatPostprocArgs, CompletionPostprocArgs,
-    ResponsesAPIPostprocArgs, chat_harmony_post_processor,
+    PostprocStrategy, ResponsesAPIPostprocArgs, chat_harmony_post_processor,
     chat_harmony_streaming_post_processor, chat_response_post_processor,
     chat_stream_post_processor, completion_response_post_processor,
     completion_stream_post_processor, responses_api_post_processor,
-    responses_api_streaming_post_processor)
+    responses_api_streaming_post_processor, select_postproc_strategy)
 from tensorrt_llm.serve.render import (HarmonyConversionError, RenderResources,
                                        UnsupportedRenderError,
                                        mount_render_endpoints,
@@ -648,6 +648,7 @@ def __init__(
         load_builtin_extensions()
         self.generator = generator
         self._is_visual_gen = _is_visual_gen_instance(generator)
+        self._postproc = select_postproc_strategy(self.postproc_worker_enabled)
         self._embedding_max_queue_delay = embedding_max_queue_delay
         self._embedding_max_queue_size = embedding_max_queue_size
         self.embedding_batcher: Optional[EncodeBatcher] = None
@@ -1146,6 +1147,13 @@ def postproc_worker_enabled(self) -> bool:
 
         return True if self.generator.args.num_postprocess_workers > 0 else False
 
+    @functools.cached_property
+    def _postproc(self) -> PostprocStrategy:
+        # __init__ pre-caches this; computing on demand keeps instances
+        # built without __init__ (unit tests use OpenAIServer.__new__)
+        # working.
+        return select_postproc_strategy(self.postproc_worker_enabled)
+
     @property
     def _vocab_size(self) -> Optional[int]:
         if self.tokenizer is not None and self.tokenizer.tokenizer is not None:
@@ -1960,11 +1968,7 @@ async def _create_chat_response(
         await promise.aresult()
         if promise.error is not None:
             raise RuntimeError(f"Generation failed: {promise.error}")
-        if self.postproc_worker_enabled:
-            chat_response = promise.outputs[0]._postprocess_result
-        else:
-            post_processor, args = postproc_params.post_processor, postproc_params.postproc_args
-            chat_response = post_processor(promise, args)
+        chat_response = self._postproc.apply(promise, postproc_params)
 
         if disaggregated_params is not None and disaggregated_params.request_type == "context_only":
             chat_response.prompt_token_ids = promise.prompt_token_ids
@@ -2077,22 +2081,17 @@ async def chat_stream_generator(
                 promise: RequestOutput,
                 postproc_params: PostprocParams) -> AsyncGenerator[str, None]:
             try:
-                if not self.postproc_worker_enabled:
-                    post_processor, args = postproc_params.post_processor, postproc_params.postproc_args
                 first_response = await anext(promise)
                 raw_request.state.server_first_token_time = self._adjusted_steady_clock.now(
                 )
-                pp_results = first_response.outputs[
-                    0]._postprocess_result if self.postproc_worker_enabled else post_processor(
-                        first_response, args)
+                pp_results = self._postproc.apply(first_response,
+                                                  postproc_params)
                 for pp_res in pp_results:
                     yield pp_res
                 # Making sure we can handle the situation where there is only one response
                 res = first_response
                 async for res in promise:
-                    pp_results = res.outputs[
-                        0]._postprocess_result if self.postproc_worker_enabled else post_processor(
-                            res, args)
+                    pp_results = self._postproc.apply(res, postproc_params)
                     for pp_res in pp_results:
                         yield pp_res
                 yield "data: [DONE]\n\n"
@@ -2958,11 +2957,7 @@ async def completion_response(
             response = await promise
             if response.error is not None:
                 raise RuntimeError(f"Generation failed: {response.error}")
-            if not self.postproc_worker_enabled:
-                post_processor, args = postproc_params.post_processor, postproc_params.postproc_args
-                pp_result = post_processor(response, args)
-            else:
-                pp_result = response.outputs[0]._postprocess_result
+            pp_result = self._postproc.apply(response, postproc_params)
             if disaggregated_params and disaggregated_params.request_type and disaggregated_params.request_type == "context_only":
                 # Include prompt token ids for context-only requests
                 pp_result.prompt_token_ids = response.prompt_token_ids
@@ -3003,11 +2998,7 @@ async def completion_generator(promise: RequestOutput,
                                        params: Optional[PostprocParams]):
             try:
                 async for output in promise:
-                    if not self.postproc_worker_enabled:
-                        post_processor, args = params.post_processor, params.postproc_args
-                        pp_result = post_processor(output, args)
-                    else:
-                        pp_result = output.outputs[0]._postprocess_result
+                    pp_result = self._postproc.apply(output, params)
                     for pp_res in pp_result:
                         yield pp_res
                 await self._extract_metrics(output, raw_request)
@@ -3182,21 +3173,16 @@ async def chat_harmony(self, request: ChatCompletionRequest,
         async def create_streaming_generator(promise: RequestOutput,
                                              postproc_params: PostprocParams):
             try:
-                if not self.postproc_worker_enabled:
-                    post_processor, args = postproc_params.post_processor, postproc_params.postproc_args
                 first_response = await anext(promise)
                 raw_request.state.server_first_token_time = (
                     self._adjusted_steady_clock.now())
-                pp_results = (first_response.outputs[0]._postprocess_result if
-                              self.postproc_worker_enabled else post_processor(
-                                  first_response, args))
+                pp_results = self._postproc.apply(first_response,
+                                                  postproc_params)
                 for pp_res in pp_results:
                     yield pp_res
                 res = first_response
                 async for res in promise:
-                    pp_results = (res.outputs[0]._postprocess_result
-                                  if self.postproc_worker_enabled else
-                                  post_processor(res, args))
+                    pp_results = self._postproc.apply(res, postproc_params)
                     for pp_res in pp_results:
                         yield pp_res
                 yield "data: [DONE]\n\n"
@@ -3343,11 +3329,10 @@ async def create_response(
                 promise: RequestOutput,
                 postproc_params: PostprocParams) -> ResponsesResponse:
             await promise.aresult()
-            if self.postproc_worker_enabled:
-                response = promise.outputs[0]._postprocess_result
-            else:
-                args = postproc_params.postproc_args
-                response = await responses_api_create_response(
+            args = postproc_params.postproc_args
+            response = await self._postproc.apply_async(
+                promise,
+                lambda: responses_api_create_response(
                     generator=promise,
                     request=request,
                     sampling_params=args.sampling_params,
@@ -3361,7 +3346,7 @@ async def create_response(
                     tool_parser=args.tool_parser,
                     num_prompt_tokens=args.num_prompt_tokens,
                     tokenizer=self.tokenizer,
-                )
+                ))
             # A postprocessing worker's result has no prompt token ids, and the
             # context-only handoff needs them, so take them from the promise as
             # the chat path does.
@@ -3375,16 +3360,14 @@ async def create_response(
 
         async def create_streaming_generator(promise: RequestOutput,
                                              postproc_params: PostprocParams):
-            post_processor, args = postproc_params.post_processor, postproc_params.postproc_args
+            args = postproc_params.postproc_args
             streaming_processor = args.streaming_processor
             initial_responses = streaming_processor.get_initial_responses()
             for initial_response in initial_responses:
                 yield initial_response
 
             async for res in promise:
-                pp_results = res.outputs[
-                    0]._postprocess_result if self.postproc_worker_enabled else post_processor(
-                        res, args)
+                pp_results = self._postproc.apply(res, postproc_params)
                 for pp_res in pp_results:
                     yield pp_res
             await self._extract_metrics(res, raw_request)

tensorrt_llm/serve/postprocess_handlers.py

@@ -14,8 +14,10 @@
 # limitations under the License.
 import json
 import time
+from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import Any, List, Literal, Optional, Tuple, Union
+from typing import (Any, Awaitable, Callable, List, Literal, Optional, Tuple,
+                    Union)
 
 from tensorrt_llm.logger import logger
 from tensorrt_llm.serve.responses_utils import ResponsesStreamingProcessor
@@ -25,7 +27,7 @@
 from .._utils import nvtx_range_debug
 from ..executor import (DetokenizedGenerationResultBase, GenerationResult,
                         GenerationResultBase)
-from ..executor.postproc_worker import PostprocArgs
+from ..executor.postproc_worker import PostprocArgs, PostprocParams
 from ..executor.result import Logprob, TokenLogprobs
 from ..llmapi import SamplingParams
 from ..llmapi.disagg_utils import (get_usage_tokens_from_ctx,
@@ -1106,3 +1108,53 @@ def responses_api_streaming_post_processor(
             args.streaming_processor.get_final_response_non_store(
                 rsp, args.num_prompt_tokens, tokenizer=args.tokenizer))
     return outputs
+
+
+class PostprocStrategy(ABC):
+    """How a handler gets the post-processed result of one engine response.
+
+    A server chooses one strategy at construction from
+    num_postprocess_workers. WorkerPostprocStrategy: a postproc worker
+    already ran the post-processor and the frontend stored its return value
+    on outputs[0]. InlinePostprocStrategy: run the post-processor in the
+    server process.
+
+    With n > 1 the two are not equivalent in what the frontend records: in
+    worker mode only outputs[0] carries the result and the metrics arrive
+    merged, while inline mode records metrics per candidate. That difference
+    is known and deliberately not changed here.
+    """
+
+    @abstractmethod
+    def apply(self, result: GenerationResultBase,
+              postproc_params: Optional[PostprocParams]) -> Any:
+        """Return the post-processed value for one engine response."""
+
+    @abstractmethod
+    async def apply_async(self, result: GenerationResultBase,
+                          inline: Callable[[], Awaitable[Any]]) -> Any:
+        """Like apply, for in-server post-processing that is a coroutine."""
+
+
+class WorkerPostprocStrategy(PostprocStrategy):
+
+    def apply(self, result, postproc_params):
+        return result.outputs[0].postprocess_result
+
+    async def apply_async(self, result, inline):
+        return result.outputs[0].postprocess_result
+
+
+class InlinePostprocStrategy(PostprocStrategy):
+
+    def apply(self, result, postproc_params):
+        return postproc_params.post_processor(result,
+                                              postproc_params.postproc_args)
+
+    async def apply_async(self, result, inline):
+        return await inline()
+
+
+def select_postproc_strategy(postproc_worker_enabled: bool) -> PostprocStrategy:
+    return (WorkerPostprocStrategy()
+            if postproc_worker_enabled else InlinePostprocStrategy())

tests/integration/test_lists/test-db/l0_cpu.yml

@@ -334,6 +334,7 @@ l0_cpu:
   - unittest/quantization/test_mode.py
   - unittest/scaffolding
   - unittest/scripts
+  - unittest/serve/test_postproc_strategy.py
   - unittest/tools
   - unittest/usage
   - unittest/utils/test_logger.py

tests/unittest/api_stability/references/completion_output.yaml

@@ -35,3 +35,6 @@ properties:
   routed_experts:
     annotation: Optional[torch.Tensor]
     default: inspect._empty
+  postprocess_result:
+    annotation: Any
+    default: inspect._empty

tests/unittest/disaggregated/test_disagg_responses_api.py

@@ -327,7 +327,7 @@ async def nothing(*args, **kwargs):
         return None
 
     promise = SimpleNamespace(
-        outputs=[SimpleNamespace(_postprocess_result=postprocessed)],
+        outputs=[SimpleNamespace(postprocess_result=postprocessed)],
         prompt_token_ids=[5, 6, 7],
         aresult=nothing,
     )

tests/unittest/grpc/openengine/test_metrics.py

@@ -80,6 +80,7 @@ def _args() -> SimpleNamespace:
         enable_iter_perf_stats=True,
         guided_decoding_backend=None,
         reasoning_parser=None,
+        num_postprocess_workers=0,
         max_seq_len=100,
         max_batch_size=8,
         max_num_tokens=100,

tests/unittest/llmapi/apps/test_anthropic_routes.py

@@ -310,7 +310,7 @@ async def __anext__(self):
             frames = self._steps.pop(0)
             return SimpleNamespace(
                 finished=not self._steps,
-                outputs=[SimpleNamespace(_postprocess_result=frames)],
+                outputs=[SimpleNamespace(postprocess_result=frames)],
             )
         if self._then is not None:
             raise self._then

tests/unittest/llmapi/run_llm_with_postproc.py

@@ -56,7 +56,7 @@ async def generate_async():
                                                _postproc_params=postproc_params,
                                                streaming=True):
             print(output)
-            outputs.append(output.outputs[0]._postprocess_result)
+            outputs.append(output.outputs[0].postprocess_result)
 
     asyncio.run(generate_async())
 

tests/unittest/llmapi/test_executor.py

@@ -429,6 +429,7 @@ def test_handle_response_propagates_should_abort():
     result._handle_response(output)
     assert result.aborted()
     assert result._outputs[0]._postprocess_result == "mock_sse_data"
+    assert result._outputs[0].postprocess_result == "mock_sse_data"
 
 
 def test_PostprocWorker_Output_tracing_fields():