[None][perf] disagg: opt-in ctx-side base64 int32 transport for prompt_token_ids (#15698)

Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
Co-authored-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

Author: lancelly

tensorrt_llm/llmapi/disagg_utils.py

@@ -90,6 +90,10 @@ class DisaggServerConfig():
     # worker's request body / JSON-parse GIL cost at high concurrency. Enable
     # ONLY for text-only, non-harmony deployments (see _get_gen_request).
     gen_strip_message_history: bool = False
+    # Ask context workers to return prompt_token_ids as a base64 int32 buffer so
+    # the orchestrator relays a string instead of materializing the token-id list
+    # on its event loop. Text-only, non-harmony deployments (see _get_ctx_request).
+    gen_tokids_ctxbytes: bool = False
 
 
 @dataclass
@@ -140,6 +144,7 @@ def extract_disagg_cfg(hostname: str = 'localhost',
                            'context_first',
                            'generation_first'] = 'context_first',
                        gen_strip_message_history: bool = False,
+                       gen_tokids_ctxbytes: bool = False,
                        **kwargs: Any) -> DisaggServerConfig:
     context_servers = context_servers or {}
     generation_servers = generation_servers or {}
@@ -188,6 +193,7 @@ def extract_disagg_cfg(hostname: str = 'localhost',
     if schedule_style:
         config.schedule_style = schedule_style
     config.gen_strip_message_history = gen_strip_message_history
+    config.gen_tokids_ctxbytes = gen_tokids_ctxbytes
     return config
 
 

tensorrt_llm/serve/openai_disagg_service.py

@@ -78,6 +78,8 @@ def __init__(
         self._health_check_interval_secs = health_check_interval_secs
         # Opt-in body-shrink for generation_only requests; see _get_gen_request.
         self._strip_gen_message_history = config.gen_strip_message_history
+        # Opt-in: ask context workers to return prompt_token_ids as base64 int32.
+        self._tokids_ctxbytes = config.gen_tokids_ctxbytes
 
         self._ctx_client = None
         self._gen_client = None
@@ -320,6 +322,7 @@ def _get_ctx_request(
                     request_type="context_only",
                     disagg_request_id=disagg_request_id,
                     schedule_style=self._schedule_style,
+                    return_prompt_token_ids_b64=self._tokids_ctxbytes,
                 ),
                 "stream": False,
                 "stream_options": None,
@@ -342,7 +345,12 @@ def _get_gen_request(
             if isinstance(request, CompletionRequest):
                 request.prompt = ctx_response.prompt_token_ids
             elif isinstance(request, ChatCompletionRequest):
-                request.prompt_token_ids = ctx_response.prompt_token_ids
+                # Relay the base64 token-id string verbatim (no int-list
+                # materialization on the orchestrator loop), else the int array.
+                if ctx_response.prompt_token_ids_b64 is not None:
+                    request.prompt_token_ids_b64 = ctx_response.prompt_token_ids_b64
+                else:
+                    request.prompt_token_ids = ctx_response.prompt_token_ids
                 # Opt-in: drop conversation history so the gen worker doesn't
                 # re-parse the full conversation JSON (dominates its GIL at high
                 # concurrency). It uses prompt_token_ids and only reads the last

tensorrt_llm/serve/openai_protocol.py

@@ -137,6 +137,9 @@ class DisaggregatedParams(OpenAIBaseModel):
     ctx_dp_rank: Optional[int] = None
     ctx_info_endpoint: Optional[str] = None
     schedule_style: Optional[DisaggScheduleStyle] = None
+    # Orchestrator -> context-worker instruction: return prompt_token_ids as a
+    # base64 int32 buffer (prompt_token_ids_b64) instead of a JSON int array.
+    return_prompt_token_ids_b64: bool = False
 
 
 class ConversationParams(OpenAIBaseModel):
@@ -618,6 +621,9 @@ class ChatCompletionResponse(OpenAIBaseModel):
     # Add prompt_tokens_ids to the response to remove the tokenization
     # in the generation server in disaggreated serving
     prompt_token_ids: Optional[List[int]] = None
+    # base64 int32 buffer alternative to prompt_token_ids; set by the context
+    # worker so the orchestrator can relay a string instead of the int list.
+    prompt_token_ids_b64: Optional[str] = None
 
 
 class DeltaMessage(OpenAIBaseModel):
@@ -675,6 +681,9 @@ class ChatCompletionRequest(OpenAIBaseModel):
     # Add prompt_tokens_ids to the request to remove the tokenization
     # in the generation server in disaggreated serving
     prompt_token_ids: Optional[List[int]] = None
+    # base64 int32 buffer relayed by the orchestrator from the ctx response;
+    # decoded back to prompt_token_ids on the generation worker. Not for clients.
+    prompt_token_ids_b64: Optional[str] = None
     model: str
     frequency_penalty: Optional[float] = 0.0
     logit_bias: Optional[Dict[str, float]] = None

tensorrt_llm/serve/openai_server.py

@@ -1181,6 +1181,13 @@ async def chat_stream_generator(
                     self.multimodal_server_config,
                     request_media_io_kwargs=request.media_io_kwargs)
 
+            # Decode base64 int32 prompt_token_ids relayed by the orchestrator.
+            if request.prompt_token_ids is None and request.prompt_token_ids_b64:
+                import numpy as np
+                request.prompt_token_ids = np.frombuffer(
+                    base64.b64decode(request.prompt_token_ids_b64),
+                    dtype=np.int32).tolist()
+
             if request.prompt_token_ids is not None:
                 prompt = request.prompt_token_ids
             else:
@@ -1263,6 +1270,18 @@ async def chat_stream_generator(
             else:
                 response = await self._create_chat_response(
                     promise, postproc_params, raw_request, disaggregated_params)
+                # Context-only: optionally return prompt_token_ids as a base64
+                # int32 buffer (one string) so the disagg orchestrator relays it
+                # without materializing the int list on its event loop. The encode
+                # runs here on the context worker (1-of-N), not the orchestrator.
+                if (request.disaggregated_params is not None and
+                        request.disaggregated_params.return_prompt_token_ids_b64
+                        and response.prompt_token_ids is not None):
+                    import numpy as np
+                    response.prompt_token_ids_b64 = base64.b64encode(
+                        np.asarray(response.prompt_token_ids,
+                                   dtype=np.int32).tobytes()).decode("ascii")
+                    response.prompt_token_ids = None
                 return JSONResponse(content=response.model_dump())
         except CppExecutorError:
             logger.error(traceback.format_exc())