From 5a8d2b71040f65ce77abc998c8c0dde1c8fce0ce Mon Sep 17 00:00:00 2001 From: Iman Tabrizian Date: Sat, 25 Jul 2026 15:36:33 -0700 Subject: [PATCH] [None][perf] serve: always use msgspec msgpack for disagg orchestrator->worker body Make the msgspec MessagePack transport for the disaggregated orchestrator->worker request body unconditional and drop the TRTLLM_SERVE_ENABLE_MSGSPEC opt-in flag. msgspec is already a hard dependency (requirements.txt -> install_requires), so the guarded import and its ImportError shim are no longer needed. The worker still keys off the X-TRTLLM-Msgpack header to choose the decoder, so external OpenAI-API clients keep the unchanged stdlib-json path. Signed-off-by: Iman Tabrizian --- ..._Workload_Optimizations_in_TensorRT-LLM.md | 2 +- tensorrt_llm/serve/openai_client.py | 36 +++------ tensorrt_llm/serve/openai_server.py | 79 +++++++++---------- 3 files changed, 48 insertions(+), 69 deletions(-) diff --git a/docs/source/blogs/tech_blog/blog26_DeepSeek_V4_on_NVIDIA_Blackwell_Model_Specific_and_Agentic_Workload_Optimizations_in_TensorRT-LLM.md b/docs/source/blogs/tech_blog/blog26_DeepSeek_V4_on_NVIDIA_Blackwell_Model_Specific_and_Agentic_Workload_Optimizations_in_TensorRT-LLM.md index 260b94448f42..30812197a64d 100644 --- a/docs/source/blogs/tech_blog/blog26_DeepSeek_V4_on_NVIDIA_Blackwell_Model_Specific_and_Agentic_Workload_Optimizations_in_TensorRT-LLM.md +++ b/docs/source/blogs/tech_blog/blog26_DeepSeek_V4_on_NVIDIA_Blackwell_Model_Specific_and_Agentic_Workload_Optimizations_in_TensorRT-LLM.md @@ -428,7 +428,7 @@ After reducing host work inside CTX and GEN, latency could still accumulate befo **Scheduling.** Under high concurrency, a GEN worker may have to choose between newly admitted requests waiting for their first token and requests that have already begun decoding. When the batch or token budget is tight, placing both groups in the same queue can leave a new request behind many ongoing conversations for several scheduling iterations. We added an opt-in policy for disaggregated GEN that gives first-token requests priority while preserving FIFO order within both groups. This can move a newly admitted request into an earlier feasible batch and reduce the queueing delay before its first decode step without adding GPU work. The policy is designed to improve first-token responsiveness and tail-latency SLO compliance rather than raw throughput. -**Orchestrator.** Profiling exposed two sources of per-request overhead. First, **native block-key hashing** removed the routing hot spot. The instance-level KV router originally hashed every token ID in a Python loop on its single asyncio thread. This held the Global Interpreter Lock (GIL) and serialized concurrent requests before they could reach CTX. Replacing the loop with a native C++ block-key hasher preserved the cache-key format while moving prompt-length hashing out of the Python event loop. Second, **efficient relay serialization** reduced request-body processing when a request crossed the HTTP boundaries to CTX and GEN. The default path previously materialized a JSON-compatible Python object containing the full prompt-token list, after which the HTTP client traversed that object again to encode the request body. We changed this path to use `pydantic-core` to emit the JSON body directly. For deployments where request-body processing remains a bottleneck, an opt-in `msgspec` MessagePack transport provides an alternative internal orchestrator-to-worker path that removes JSON encoding and parsing from this hop. The external OpenAI API and the default direct-JSON path remain unchanged. +**Orchestrator.** Profiling exposed two sources of per-request overhead. First, **native block-key hashing** removed the routing hot spot. The instance-level KV router originally hashed every token ID in a Python loop on its single asyncio thread. This held the Global Interpreter Lock (GIL) and serialized concurrent requests before they could reach CTX. Replacing the loop with a native C++ block-key hasher preserved the cache-key format while moving prompt-length hashing out of the Python event loop. Second, **efficient relay serialization** reduced request-body processing when a request crossed the HTTP boundaries to CTX and GEN. The default path previously materialized a JSON-compatible Python object containing the full prompt-token list, after which the HTTP client traversed that object again to encode the request body. We changed this path to use `pydantic-core` to emit the JSON body directly. To remove request-body processing from this hop entirely, the internal orchestrator-to-worker path now uses a `msgspec` MessagePack transport that eliminates JSON encoding and parsing. The external OpenAI API remains unchanged. **CTX-to-GEN protocol.** The CTX-to-GEN protocol had two issues: it could carry more request state than GEN needed, and a retried request could leave CTX and GEN with inconsistent identifiers. In supported text-only, non-Harmony deployments, an opt-in path removes earlier chat history from the GEN request while preserving the prompt token IDs, final message, tool definitions, and generation settings that GEN still needs. This reduces copying, serialization, and transfer without changing the request semantics supported by that path. We also hardened the retry path so that known transient connection failures do not leave CTX and GEN with inconsistent request identifiers. Together, these changes reduce protocol overhead while preserving consistency across the disaggregated handoff. diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index 26646a54344c..33a83d53cdd6 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -14,12 +14,12 @@ # yapf: disable import asyncio -import os import traceback from abc import ABC, abstractmethod from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, List, Optional, Tuple, Type import aiohttp +import msgspec from tensorrt_llm.llmapi.disagg_utils import ServerRole from tensorrt_llm.logger import logger @@ -41,21 +41,11 @@ # yapf: enable -# msgspec msgpack is an opt-in transport for the orchestrator->worker request -# body (alternative to the JSON path). Enable with TRTLLM_SERVE_ENABLE_MSGSPEC=1; +# msgspec msgpack is the transport for the orchestrator->worker request body: # the orchestrator encodes the forwarded body as msgpack and flags it with the # X-TRTLLM-Msgpack header (Content-Type stays application/json so FastAPI still -# routes it through Request.json()). Fails loudly at import if msgspec is missing. -_MSGSPEC_ENABLED = os.getenv("TRTLLM_SERVE_ENABLE_MSGSPEC", "0") == "1" -if _MSGSPEC_ENABLED: - try: - import msgspec - except ImportError as exc: - raise ImportError( - "TRTLLM_SERVE_ENABLE_MSGSPEC=1 requires the msgspec package " - "(listed in requirements.txt)." - ) from exc - _msgpack_encoder = msgspec.msgpack.Encoder() +# routes it through Request.json()). +_msgpack_encoder = msgspec.msgpack.Encoder() class OpenAIClient(ABC): @@ -216,17 +206,13 @@ async def _post_with_retry( if dp is not None and getattr(dp, "disagg_request_id", None) is not None: dp.disagg_request_id = await self._disagg_id_generator() # Serialize once on the orchestrator's single event-loop thread. - if _MSGSPEC_ENABLED: - # msgspec msgpack: encode the request dict to msgpack bytes. Keep - # Content-Type application/json so FastAPI still routes the body - # through Request.json() (it only does that for json/+json content - # subtypes); the X-TRTLLM-Msgpack header tells the worker's - # Request.json() to decode with msgspec instead of stdlib json. - body = _msgpack_encoder.encode(request.model_dump(mode="json", exclude_unset=True)) - req_headers = {"Content-Type": "application/json", "X-TRTLLM-Msgpack": "1"} - else: - body = request.model_dump_json(exclude_unset=True) - req_headers = {"Content-Type": "application/json"} + # msgspec msgpack: encode the request dict to msgpack bytes. Keep + # Content-Type application/json so FastAPI still routes the body + # through Request.json() (it only does that for json/+json content + # subtypes); the X-TRTLLM-Msgpack header tells the worker's + # Request.json() to decode with msgspec instead of stdlib json. + body = _msgpack_encoder.encode(request.model_dump(mode="json", exclude_unset=True)) + req_headers = {"Content-Type": "application/json", "X-TRTLLM-Msgpack": "1"} try: lines_yielded = 0 start_time = get_steady_clock_now_in_seconds() diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index ad1070f2d8ea..a58ee9d6ab75 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -20,6 +20,7 @@ from typing import (Annotated, Any, AsyncGenerator, AsyncIterator, List, Optional, Union) +import msgspec import uvicorn from fastapi import Body, FastAPI, Request from fastapi.exceptions import RequestValidationError @@ -105,52 +106,45 @@ # yapf: enable -# msgspec msgpack is an opt-in transport for the disagg orchestrator->worker -# request body: the large agentic chat body otherwise blocks the serving event -# loop on stdlib json.loads. Enable with TRTLLM_SERVE_ENABLE_MSGSPEC=1 (must be -# set on both orchestrator and worker). The worker decodes bodies flagged with -# the X-TRTLLM-Msgpack header via msgspec and falls back to stdlib json for -# everything else, so the JSON path is byte-for-byte unchanged when the flag is off. -_MSGSPEC_ENABLED = os.getenv("TRTLLM_SERVE_ENABLE_MSGSPEC", "0") == "1" -if _MSGSPEC_ENABLED: - try: - import msgspec - except ImportError as exc: - raise ImportError( - "TRTLLM_SERVE_ENABLE_MSGSPEC=1 requires the msgspec package " - "(listed in requirements.txt).") from exc - _msgpack_decoder = msgspec.msgpack.Decoder() - - class _MsgspecRequest(Request): - """Request that decodes msgpack bodies (X-TRTLLM-Msgpack: 1) with msgspec. - - The orchestrator sends Content-Type application/json (so FastAPI still - routes the body through Request.json()) with the X-TRTLLM-Msgpack header - flagging a msgspec-msgpack payload; everything else is stdlib json. - """ +# msgspec msgpack is the transport for the disagg orchestrator->worker request +# body: the large agentic chat body otherwise blocks the serving event loop on +# stdlib json.loads. The worker decodes bodies flagged with the X-TRTLLM-Msgpack +# header via msgspec and falls back to stdlib json for everything else, so +# external OpenAI-API clients keep the byte-for-byte unchanged JSON path. +_msgpack_decoder = msgspec.msgpack.Decoder() + + +class _MsgspecRequest(Request): + """Request that decodes msgpack bodies (X-TRTLLM-Msgpack: 1) with msgspec. + + The orchestrator sends Content-Type application/json (so FastAPI still + routes the body through Request.json()) with the X-TRTLLM-Msgpack header + flagging a msgspec-msgpack payload; everything else is stdlib json. + """ + + async def json(self): + if not hasattr(self, "_json_body"): + body = await self.body() + if not body: + self._json_body = {} + elif self.headers.get("x-trtllm-msgpack") == "1": + self._json_body = _msgpack_decoder.decode(body) + else: + self._json_body = json.loads(body) + return self._json_body - async def json(self): - if not hasattr(self, "_json_body"): - body = await self.body() - if not body: - self._json_body = {} - elif self.headers.get("x-trtllm-msgpack") == "1": - self._json_body = _msgpack_decoder.decode(body) - else: - self._json_body = json.loads(body) - return self._json_body - class _MsgspecRoute(APIRoute): - """APIRoute that parses request bodies via :class:`_MsgspecRequest`.""" +class _MsgspecRoute(APIRoute): + """APIRoute that parses request bodies via :class:`_MsgspecRequest`.""" - def get_route_handler(self): - original_route_handler = super().get_route_handler() + def get_route_handler(self): + original_route_handler = super().get_route_handler() - async def route_handler(request: Request): - return await original_route_handler( - _MsgspecRequest(request.scope, request.receive)) + async def route_handler(request: Request): + return await original_route_handler( + _MsgspecRequest(request.scope, request.receive)) - return route_handler + return route_handler TIMEOUT_KEEP_ALIVE = 5 # seconds. @@ -439,8 +433,7 @@ async def lifespan(app: FastAPI): self.generator.shutdown() self.app = FastAPI(lifespan=lifespan) - if _MSGSPEC_ENABLED: - self.app.router.route_class = _MsgspecRoute + self.app.router.route_class = _MsgspecRoute @self.app.exception_handler(RequestValidationError) async def validation_exception_handler(_, exc):