Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Comment thread
coderabbitai[bot] marked this conversation as resolved.

**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.

Expand Down
36 changes: 11 additions & 25 deletions tensorrt_llm/serve/openai_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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):
Expand Down Expand Up @@ -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()
Expand Down
79 changes: 36 additions & 43 deletions tensorrt_llm/serve/openai_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
Comment thread
Tabrizian marked this conversation as resolved.
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.
Expand Down Expand Up @@ -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):
Expand Down
Loading