From cb9375c3b40f493fc679cb2149bd821ec2dacd26 Mon Sep 17 00:00:00 2001 From: Liao Lanyu <108499334+lancelly@users.noreply.github.com> Date: Fri, 22 May 2026 13:50:50 +0800 Subject: [PATCH 01/28] [None][feat] adp_router: salt-aware probe_prefix_match_length on v1/v2 KV cache backends (#14435) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com> --- .../_torch/pyexecutor/resource_manager.py | 16 ++++++++++++++-- .../_torch/pyexecutor/scheduler/adp_router.py | 9 ++++++++- 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index ebb8fe753991..220433d6bc16 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -643,11 +643,19 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], device='cpu') self.blocks_per_window = blocks_per_window - def probe_prefix_match_length(self, input_tokens, lora_task_id=None): + def probe_prefix_match_length(self, + input_tokens, + lora_task_id=None, + cache_salt_id=None): """Probe the KV cache radix tree for prefix match length. Returns the number of prefix tokens already cached on this rank. Used by KVCacheAwareADPRouter for cache-aware routing. + + ``cache_salt_id`` scopes the probe to the same per-request + namespace that ``_create_kv_cache`` uses; without it, salted + requests would be probed against the salt=None namespace and the + router would see the wrong match length. """ if not self.enable_block_reuse: return 0 @@ -664,12 +672,16 @@ def probe_prefix_match_length(self, input_tokens, lora_task_id=None): LlmRequest as CppLlmRequest block_key = BlockKey(tokens=input_tokens, lora_task_id=lora_task_id) unique_tokens = block_key.unique_tokens + # buildBlockKeys() on the C++ side pulls cache_salt_id off the + # request, so dummy_req must carry it for the probe namespace to + # match the real create_kv_cache path. dummy_req = CppLlmRequest(request_id=0, max_new_tokens=0, input_tokens=input_tokens, sampling_config=SamplingConfig(), is_streaming=False, - lora_task_id=lora_task_id) + lora_task_id=lora_task_id, + cache_salt_id=cache_salt_id) summary = self.impl.analyze_prefix_reuse(unique_tokens, dummy_req) return summary.reusable_blocks_all * self.tokens_per_block diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index b717fc38ca1e..ea909b0292c5 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -515,7 +515,14 @@ def gather_prefix_matches( probe_tokens = input_tokens[:-1] if len(input_tokens) > 1 else [] lora_config = getattr(req, "lora_config", None) lora_task_id = lora_config.task_id if lora_config is not None else None - match_len = self.kv_cache_manager.probe_prefix_match_length(probe_tokens, lora_task_id) + # cache_salt_id scopes block reuse on both v1 and v2 backends; + # passing None for non-salted requests is a no-op. + cache_salt_id = getattr(req, "cache_salt_id", None) + match_len = self.kv_cache_manager.probe_prefix_match_length( + probe_tokens, + lora_task_id, + cache_salt_id=cache_salt_id, + ) local_matches.extend([req_item.id, match_len]) all_data = self.dist.tp_allgather(local_matches) From f7e3ab3a4eac55dcbf58e873a7dd0437705572ad Mon Sep 17 00:00:00 2001 From: Shi Xiaowei <39303645+Shixiaowei02@users.noreply.github.com> Date: Tue, 26 May 2026 20:58:59 +0800 Subject: [PATCH 02/28] [TRTLLM-12941][fix] Selective retry for transient TCP races in disagg HTTP client (#14562) Signed-off-by: Fanrong Li Signed-off-by: Shi Xiaowei <39303645+Shixiaowei02@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 5 +- tensorrt_llm/serve/openai_client.py | 26 +++- .../test_disagg_openai_client.py | 143 ++++++++++++++++++ 3 files changed, 168 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 188ab3302dda..5462d37d896d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -4522,8 +4522,9 @@ def flag_if_kv_transfer_timed_out(req: LlmRequest, type: str) -> None: elapsed_time = (current_time - req.py_kv_transfer_start_time) * 1000 if elapsed_time > timeout_ms and not req.py_kv_transfer_timed_out: logger.warning( - f"Terminating {type} request {req.py_request_id} due to KV cache transfer timeout" - ) + f"Terminating {type} request {req.py_request_id} due to KV " + f"cache transfer timeout: elapsed {elapsed_time:.0f}ms > " + f"kv_transfer_timeout_ms={timeout_ms}ms") req.py_kv_transfer_timed_out = True for req in self.async_transfer_manager.requests_in_transfer().values(): diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index 0371f2da7e03..7d8ae8f8ef2c 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -164,7 +164,13 @@ async def _post_with_retry( hooks: Optional[ResponseHooks] = None, ) -> AsyncGenerator[Any, None]: is_stream = request.stream - for attempt in range(self._max_retries + 1): + # Loop range must cover the transient-TCP extended budget (up to 5) + # so the conditional raise inside the except block can actually decide + # to keep retrying. Non-transient errors still raise on the first + # attempt that reaches self._max_retries. + _TRANSIENT_TCP_BUDGET = 5 + loop_max = max(self._max_retries, _TRANSIENT_TCP_BUDGET) + 1 + for attempt in range(loop_max): # Regenerate disagg_request_id on retry to avoid ID collision on workers if attempt > 0 and self._disagg_id_generator is not None: dp = getattr(request, "disaggregated_params", None) @@ -215,15 +221,27 @@ async def _post_with_retry( traceback.format_exc(), ) raise - if attempt == self._max_retries: + # Selective retry budget: ServerDisconnectedError and + # ConnectionResetError are transient TCP races (typically at + # burst start when client keepalive vs server keepalive race). + # Give them an extended retry budget while preserving the + # original fail-fast for genuine upstream errors. + is_transient_tcp = isinstance( + e, + (aiohttp.ServerDisconnectedError, ConnectionResetError), + ) + effective_max = self._max_retries + if is_transient_tcp: + effective_max = max(self._max_retries, _TRANSIENT_TCP_BUDGET) + if attempt >= effective_max: logger.error( - f"Client error to {url}: {e} - last retry {attempt} of {self._max_retries}" + f"Client error to {url}: {e} - last retry {attempt} of {effective_max}" "failed", traceback.format_exc(), ) raise logger.error( - f"{self._role} client error to {url}: {e} - retry {attempt} of {self._max_retries}", + f"{self._role} client error to {url}: {e} - retry {attempt} of {effective_max}", traceback.format_exc(), ) await asyncio.sleep(self._retry_interval_sec) diff --git a/tests/unittest/disaggregated/test_disagg_openai_client.py b/tests/unittest/disaggregated/test_disagg_openai_client.py index ee97510580b3..2538c5c92df0 100644 --- a/tests/unittest/disaggregated/test_disagg_openai_client.py +++ b/tests/unittest/disaggregated/test_disagg_openai_client.py @@ -412,3 +412,146 @@ async def test_no_generator_keeps_original_id(self): await client.send_request(req) assert req.disaggregated_params.disagg_request_id == 42 + + +class TestSelectiveTransientTcpRetry: + """Selective retry budget for transient TCP race symptoms. + + ServerDisconnectedError and ConnectionResetError (which include + aiohttp.ClientConnectionResetError via MRO) get an extended retry budget + of up to 5 attempts; all other client errors keep the original + max_retries fail-fast behaviour. + """ + + def _ok_response(self): + return CompletionResponse( + id="cmpl-1", + object="text_completion", + created=0, + model="m", + choices=[CompletionResponseChoice(index=0, text="ok", finish_reason="stop")], + usage=UsageInfo(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ) + + def _mock_http_ok(self, body): + r = AsyncMock() + r.status = 200 + r.headers = {"Content-Type": "application/json"} + r.json = AsyncMock(return_value=body.model_dump()) + r.__aenter__ = AsyncMock(return_value=r) + r.__aexit__ = AsyncMock() + return r + + def _make_client(self, session, max_retries=1): + from prometheus_client.registry import REGISTRY + + REGISTRY._names_to_collectors = {} + REGISTRY._collector_to_names = {} + router = AsyncMock(spec=Router) + router.servers = ["localhost:8000"] + router.get_next_server = AsyncMock(return_value=("localhost:8000", None)) + router.finish_request = AsyncMock() + return OpenAIHttpClient( + router=router, + role=ServerRole.CONTEXT, + timeout_secs=10, + max_retries=max_retries, + retry_interval_sec=0, + session=session, + ) + + def _make_request(self): + return CompletionRequest( + model="m", + prompt="hi", + stream=False, + disaggregated_params=DisaggregatedParams( + request_type="context_only", disagg_request_id=1 + ), + ) + + @pytest.mark.asyncio + async def test_server_disconnected_gets_extra_retries(self): + """ServerDisconnectedError: even with max_retries=1, retry up to 5.""" + session = AsyncMock(spec=aiohttp.ClientSession) + client = self._make_client(session, max_retries=1) + + # 4 disconnect failures then success on the 5th attempt + session.post.side_effect = [ + aiohttp.ServerDisconnectedError(), + aiohttp.ServerDisconnectedError(), + aiohttp.ServerDisconnectedError(), + aiohttp.ServerDisconnectedError(), + self._mock_http_ok(self._ok_response()), + ] + + with patch("asyncio.sleep", new_callable=AsyncMock): + await client.send_request(self._make_request()) + + # 1 original + 4 retries = 5 total attempts (extra budget kicked in) + assert session.post.call_count == 5 + + @pytest.mark.asyncio + async def test_connection_reset_gets_extra_retries(self): + """ConnectionResetError: same extra budget as ServerDisconnectedError.""" + session = AsyncMock(spec=aiohttp.ClientSession) + client = self._make_client(session, max_retries=1) + + session.post.side_effect = [ + ConnectionResetError(), + ConnectionResetError(), + self._mock_http_ok(self._ok_response()), + ] + + with patch("asyncio.sleep", new_callable=AsyncMock): + await client.send_request(self._make_request()) + + # 1 original + 2 retries = 3 total attempts (within extra budget) + assert session.post.call_count == 3 + + @pytest.mark.asyncio + async def test_other_client_error_keeps_fail_fast(self): + """Generic aiohttp.ClientError still respects max_retries (=1).""" + session = AsyncMock(spec=aiohttp.ClientSession) + client = self._make_client(session, max_retries=1) + + session.post.side_effect = aiohttp.ClientError("transient non-tcp") + + with patch("asyncio.sleep", new_callable=AsyncMock): + with pytest.raises(aiohttp.ClientError): + await client.send_request(self._make_request()) + + # Original + 1 retry = 2 attempts, NOT promoted to 5 + assert session.post.call_count == 2 + + @pytest.mark.asyncio + async def test_max_retries_zero_still_gets_transient_tcp_budget(self): + """Even when max_retries=0, transient TCP races still retry up to 5.""" + session = AsyncMock(spec=aiohttp.ClientSession) + client = self._make_client(session, max_retries=0) + + session.post.side_effect = [ + aiohttp.ServerDisconnectedError(), + self._mock_http_ok(self._ok_response()), + ] + + with patch("asyncio.sleep", new_callable=AsyncMock): + await client.send_request(self._make_request()) + + assert session.post.call_count == 2 + + @pytest.mark.asyncio + async def test_transient_tcp_capped_at_5_when_max_retries_smaller(self): + """If transient TCP keeps failing, give up after the extended budget.""" + session = AsyncMock(spec=aiohttp.ClientSession) + client = self._make_client(session, max_retries=1) + + # Always raise — must give up after extended (1 + 5) = 6 attempts + session.post.side_effect = aiohttp.ServerDisconnectedError() + + with patch("asyncio.sleep", new_callable=AsyncMock): + with pytest.raises(aiohttp.ServerDisconnectedError): + await client.send_request(self._make_request()) + + # 1 original + 5 retries + assert session.post.call_count == 6 From 0a5ec7a1c0639f60c5a3c2f7f945fc872212c4ed Mon Sep 17 00:00:00 2001 From: Shi Xiaowei <39303645+Shixiaowei02@users.noreply.github.com> Date: Tue, 26 May 2026 21:28:26 +0800 Subject: [PATCH 03/28] [None][feature] Gen-first disagg with transceiver v2 (#14407) Signed-off-by: Fanrong Li Signed-off-by: Shi Xiaowei <39303645+Shixiaowei02@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 4 +- tensorrt_llm/serve/openai_client.py | 26 +- tensorrt_llm/serve/openai_disagg_service.py | 15 +- tensorrt_llm/serve/router.py | 336 +++++++++++++----- .../accuracy/test_disaggregated_serving.py | 137 +++++-- ...agg_config_conditional_deepseek_v3_v2.yaml | 42 +++ .../defs/disaggregated/test_disaggregated.py | 39 ++ .../test_lists/qa/llm_function_core.txt | 7 +- .../test_lists/test-db/l0_dgx_h100.yml | 7 +- .../test_disagg_openai_client.py | 15 +- 10 files changed, 495 insertions(+), 133 deletions(-) create mode 100644 tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3_v2.yaml diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 5462d37d896d..f96cacc72600 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -558,9 +558,11 @@ def __init__( self._prefetched_request_ids: set[int] = set() self.enable_kv_cache_events = self.kv_cache_manager is not None and self.kv_cache_manager.event_buffer_max_size > 0 self.enable_kv_cache_reuse = self.kv_cache_manager is not None and self.kv_cache_manager.enable_block_reuse + # AsyncTransferManager pin/unpin path is V1-only; V2 holds blocks via _KVCache refcount. self.enable_partial_reuse_for_disagg = ( self.enable_kv_cache_reuse - and self.kv_cache_manager.enable_partial_reuse) + and self.kv_cache_manager.enable_partial_reuse + and not isinstance(self.kv_cache_manager, KVCacheManagerV2)) self.max_input_len = max_input_len # _executor_loop private data diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index 7d8ae8f8ef2c..a3d23b266d2f 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -85,8 +85,12 @@ async def check_ready(self) -> Tuple[List[str], List[str]]: async def shutdown(self) -> None: ... @abstractmethod - async def _finish_request(self, request: UCompletionRequest) -> None: - """Finish the request in the router.""" + async def _finish_request(self, request: UCompletionRequest, success: bool = True) -> None: + """Finish the request in the router. + + ``success`` lets the router distinguish completed vs failed requests + (e.g. for cache backfill). + """ ... @@ -129,9 +133,10 @@ async def _send_request( if server is None: server, _ = await self._router.get_next_server(request) url = f"http://{server}/{endpoint}" - logger.debug( - f"Sending {self._role} request {request.disaggregated_params.ctx_request_id} to {url}" - ) + # disaggregated_params is None when conditional_disagg bypasses ctx. + _dp = request.disaggregated_params + _ctx_rid = _dp.ctx_request_id if _dp is not None else None + logger.debug(f"Sending {self._role} request {_ctx_rid} to {url}") try: self._metrics_collector.total_requests.inc() resp_generator = self._post_with_retry(server, url, request, hooks) @@ -153,7 +158,7 @@ async def _send_request( except Exception: self._metrics_collector.error_requests.inc() # finish the request upon error - await self._finish_request(request) + await self._finish_request(request, success=False) raise async def _post_with_retry( @@ -264,6 +269,7 @@ async def _response_generator( assert "text/event-stream" in http_response.headers.get("Content-Type", ""), ( "Response is not streaming" ) + success = True try: last_token_time = start_time i = 0 @@ -294,17 +300,19 @@ async def _response_generator( # a client error is expected when the response stream is done if the connector has close=True logger.error(f"{self._role} client {server} error: {e}") self._metrics_collector.error_requests.inc() + success = False raise except Exception: self._metrics_collector.error_requests.inc() + success = False raise finally: # finish the request after streaming response is done or error is raised - await self._finish_request(request) + await self._finish_request(request, success=success) - async def _finish_request(self, request: UCompletionRequest) -> None: + async def _finish_request(self, request: UCompletionRequest, success: bool = True) -> None: self._metrics_collector.completed_requests.inc() - await self._router.finish_request(request) + await self._router.finish_request(request, self._session, success=success) async def collect_metrics(self) -> Dict[str, Any]: metrics = {} diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index 9b97551a8919..1da78967a5a8 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -257,13 +257,20 @@ async def _check_conditional_disagg(self, request: UCompletionRequest) -> bool: gen_server, info = await self._gen_router.get_next_server(request) match_length = sum(info["matches"]) total_length = sum(len(token_list) for token_list in info["token_lists"]) - if ( + need_ctx_decision = ( match_length == 0 or total_length - match_length > self.conditional_disagg_config.max_local_prefill_length - ): - return gen_server, True - return gen_server, False + ) + # Visibility hook for verifying bypass triggers in disagg deployments. + logger.debug( + f"[conditional_disagg] gen={gen_server} match={match_length} " + f"total={total_length} residual={total_length - match_length} " + f"max_local_prefill_length=" + f"{self.conditional_disagg_config.max_local_prefill_length} " + f"→ need_ctx={need_ctx_decision}" + ) + return gen_server, need_ctx_decision return None, True async def _check_gen_only_disagg(self, request: UCompletionRequest) -> bool: diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 3e589cf1e356..8ee01677e06f 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -20,17 +20,32 @@ from typing import Awaitable, Callable, Dict, Iterable, List, Optional, Union import aiohttp +from transformers import AutoTokenizer -from tensorrt_llm.bindings.internal.batch_manager import (BlockKey, - BlockKeyHasher) from tensorrt_llm.llmapi.disagg_utils import (MetadataServerConfig, RouterConfig, ServerRole) from tensorrt_llm.logger import logger +from tensorrt_llm.runtime import kv_cache_hash +from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import \ + Block as V2Block +from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import \ + ReuseScope +from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import \ + RootBlock as V2RootBlock from tensorrt_llm.serve.metadata_server import JsonDictionary from tensorrt_llm.serve.openai_protocol import (ChatCompletionRequest, CompletionRequest) +KV_CACHE_HASH_ALGO_DEFAULT = kv_cache_hash.KV_CACHE_HASH_ALGO_DEFAULT +KV_CACHE_HASH_ALGO_V1 = kv_cache_hash.KV_CACHE_HASH_ALGO_V1 +KV_CACHE_HASH_ALGO_V2 = kv_cache_hash.KV_CACHE_HASH_ALGO_V2 +KV_CACHE_HASH_ALGO_V2_SHA256_64 = kv_cache_hash.KV_CACHE_HASH_ALGO_V2_SHA256_64 +get_cache_salt_id = kv_cache_hash.get_cache_salt_id +hash_v1_block_key = kv_cache_hash.hash_v1_block_key +truncate_sha256_hash_to_int64 = kv_cache_hash.truncate_sha256_hash_to_int64 + OpenAIRequest = Union[CompletionRequest, ChatCompletionRequest] +BlockHash = Union[int, str] def get_request_num_tokens(request: OpenAIRequest) -> int: @@ -111,16 +126,52 @@ def __init__( session_provider: Optional[Callable[[], aiohttp.ClientSession]] = None): super().__init__(server, use_tokens, session_provider) - self._kv_cache_block_table: set[int] = set() + self._kv_cache_block_table: set[BlockHash] = set() + self._kv_cache_block_tables: dict[str, set[BlockHash]] = { + KV_CACHE_HASH_ALGO_V1: self._kv_cache_block_table + } + self._kv_cache_hash_algo = KV_CACHE_HASH_ALGO_DEFAULT self._tokens_per_block = tokens_per_block - def add_blocks(self, block_hashes: Iterable[int]): - for hash in block_hashes: - self._kv_cache_block_table.add(hash) + @property + def hash_algo(self) -> str: + return self._kv_cache_hash_algo + + def _block_table(self, hash_algo: str) -> set[BlockHash]: + if hash_algo not in self._kv_cache_block_tables: + self._kv_cache_block_tables[hash_algo] = set() + return self._kv_cache_block_tables[hash_algo] + + def set_hash_algo(self, hash_algo: str): + self._kv_cache_hash_algo = hash_algo + self._block_table(hash_algo) - def remove_blocks(self, block_hashes: Iterable[int]): - for hash in block_hashes: - self._kv_cache_block_table.discard(hash) + async def get_hash_algo(self, hash_algo: Optional[str] = None) -> str: + async with self._lock: + if hash_algo is not None: + self.set_hash_algo(hash_algo) + return self._kv_cache_hash_algo + + def _resolve_hash_algo(self, hash_algo: Optional[str]) -> str: + return self._kv_cache_hash_algo if hash_algo is None else hash_algo + + def add_blocks(self, + block_hashes: Iterable[BlockHash], + hash_algo: Optional[str] = None): + hash_algo = self._resolve_hash_algo(hash_algo) + self.set_hash_algo(hash_algo) + block_table = self._block_table(hash_algo) + for block_hash in block_hashes: + block_table.add(block_hash) + + def remove_blocks(self, + block_hashes: Iterable[BlockHash], + hash_algo: Optional[str] = None): + hash_algo = self._resolve_hash_algo(hash_algo) + self.set_hash_algo(hash_algo) + block_table = self._block_table(hash_algo) + for block_hash in block_hashes: + block_table.discard(block_hash) def update_with_events(self, events: Iterable[dict]): # event_raw: {"id": , "data": } @@ -130,11 +181,16 @@ def update_with_events(self, events: Iterable[dict]): else: event = event_raw + hash_algo = event_raw.get( + "hash_algo", event.get("hash_algo", KV_CACHE_HASH_ALGO_DEFAULT)) + if event["type"] == "created": + self.set_hash_algo(hash_algo) if event["type"] == "stored": - self.add_blocks(block["block_hash"] - for block in event["blocks"]) + self.add_blocks( + (block["block_hash"] for block in event["blocks"]), + hash_algo=hash_algo) elif event["type"] == "removed": - self.remove_blocks(event["block_hashes"]) + self.remove_blocks(event["block_hashes"], hash_algo=hash_algo) async def poll_events(self, session: aiohttp.ClientSession): async with session.post( @@ -142,13 +198,17 @@ async def poll_events(self, session: aiohttp.ClientSession): events_raw = await response.json() return events_raw - async def matched_tokens(self, block_hashes: list[list[int]]) -> int: + async def matched_tokens( + self, + block_hashes: list[list[BlockHash]], + hash_algo: str = KV_CACHE_HASH_ALGO_DEFAULT) -> int: match_count = 0 async with self._lock: + block_table = self._block_table(hash_algo) for hash_list in block_hashes: - for hash in hash_list: + for block_hash in hash_list: # TODO: 1) parent hash verification, 2) partial matching - if hash in self._kv_cache_block_table: + if block_hash in block_table: match_count += self._tokens_per_block else: break @@ -373,7 +433,10 @@ async def get_next_server( '''Select server by request and return some intermediate information, exclude_server is a server to exclude from the selection''' @abstractmethod - async def finish_request(self, request: OpenAIRequest): + async def finish_request(self, + request: OpenAIRequest, + session: Optional[aiohttp.ClientSession] = None, + success: bool = True): pass @property @@ -618,8 +681,11 @@ async def get_next_server( ) return server, {"server_info": self._server_info.get(server, {})} - async def finish_request(self, request: OpenAIRequest): - pass + async def finish_request(self, + request: OpenAIRequest, + session: Optional[aiohttp.ClientSession] = None, + success: bool = True): + del request, session, success class LoadBalancingRouter(LoadBalancingMixin, Router): @@ -657,16 +723,30 @@ async def get_next_server( return server, {"server_info": self._server_info.get(server, {})} - async def finish_request(self, request: OpenAIRequest): + async def finish_request(self, + request: OpenAIRequest, + session: Optional[aiohttp.ClientSession] = None, + success: bool = True): + del session, success async with self._lock: await self._unregister_request(request) def block_key_hasher(token_ids: list[int], - parent_hash: Optional[int] = None) -> int: - block_key = BlockKey(token_ids) - return BlockKeyHasher.hash(block_key, - 0 if parent_hash is None else parent_hash) + parent_hash: Optional[int] = None, + cache_salt_id: Optional[int] = None) -> int: + return hash_v1_block_key( + token_ids, + parent_hash=0 if parent_hash is None else parent_hash, + cache_salt_id=cache_salt_id) + + +def v2_sha256_block_hasher(token_ids: list[int], + parent_hash: Optional[str] = None, + cache_salt_id: Optional[int] = None) -> str: + parent_key = (V2RootBlock.make_key(ReuseScope(salt=cache_salt_id)) + if parent_hash is None else bytes.fromhex(parent_hash)) + return V2Block.make_key(parent_key, token_ids).hex() class BlockHashMixin: @@ -695,15 +775,15 @@ def _get_tokenizer(self, model: str): self._tokenizers[model] = load_custom_tokenizer( self._custom_tokenizer, model) else: - # Route through TransformersTokenizer so block-hash routing - # inherits the same post-load fixes as the rest of TRT-LLM - # (e.g. maybe_fix_byte_level_tokenizer for DeepSeek-V3 - # Metaspace bug, _fallback_to_fast_tokenizer for DeepSeek-V3.2 - # AttributeError on transformers >= 5.x). Peel off .tokenizer - # to return the raw HF tokenizer used by _tokenize() below. - from tensorrt_llm.tokenizer import TransformersTokenizer - self._tokenizers[model] = TransformersTokenizer.from_pretrained( - model, trust_remote_code=True).tokenizer + from tensorrt_llm.tokenizer import \ + maybe_fix_byte_level_tokenizer + tokenizer = AutoTokenizer.from_pretrained( + model, trust_remote_code=True) + # Work around Transformers 5.x LlamaTokenizer overriding + # tokenizer.json's ByteLevel pre-tokenizer with Metaspace, + # which silently strips spaces from prompts (see tokenizer.py). + self._tokenizers[model] = maybe_fix_byte_level_tokenizer( + tokenizer, model, trust_remote_code=True) return self._tokenizers[model] def _tokenize(self, request: OpenAIRequest) -> list[list[int]]: @@ -749,18 +829,43 @@ def _tokenize(self, request: OpenAIRequest) -> list[list[int]]: if len(token_lists) > 1 else token_lists[0]) return token_lists - def _compute_block_hashes(self, - token_lists: list[list[int]]) -> list[list[int]]: - block_hashes: list[list[int]] = [] + def _compute_block_hashes( + self, + token_lists: list[list[int]], + hash_algo: str = KV_CACHE_HASH_ALGO_DEFAULT, + cache_salt_id: Optional[int] = None, + ) -> list[list[BlockHash]]: + if hash_algo == KV_CACHE_HASH_ALGO_V1: + block_hasher = block_key_hasher + elif hash_algo == KV_CACHE_HASH_ALGO_V2: + block_hasher = v2_sha256_block_hasher + elif hash_algo == KV_CACHE_HASH_ALGO_V2_SHA256_64: + reuse_scope = ReuseScope(salt=cache_salt_id) + block_hashes: list[list[BlockHash]] = [] + for token_list in token_lists: + hash_list = [] + parent_key = V2RootBlock.make_key(reuse_scope) + for t in range(0, len(token_list) - 1, self._tokens_per_block): + t_end = min(t + self._tokens_per_block, len(token_list) - 1) + parent_key = V2Block.make_key(parent_key, + token_list[t:t_end]) + hash_list.append(truncate_sha256_hash_to_int64(parent_key)) + block_hashes.append(hash_list) + return block_hashes + else: + raise ValueError( + f"Unsupported KV cache hash algorithm: {hash_algo}") + + block_hashes: list[list[BlockHash]] = [] for token_list in token_lists: hash_list = [] - # in KvCacheManager, the last token is not included in the - # block key + # in KvCacheManager, the last token is not included in the block key for t in range(0, len(token_list) - 1, self._tokens_per_block): t_end = min(t + self._tokens_per_block, len(token_list) - 1) hash_list.append( - block_key_hasher(token_list[t:t_end], - None if t == 0 else hash_list[-1])) + block_hasher(token_list[t:t_end], + None if t == 0 else hash_list[-1], + cache_salt_id)) block_hashes.append(hash_list) return block_hashes @@ -778,6 +883,22 @@ def _tokenize_and_compute_block_hashes( block_hashes = self._compute_block_hashes(token_lists) return token_lists, block_hashes + def _tokenize_and_compute_block_hashes_by_algo( + self, + request: OpenAIRequest, + hash_algos: Iterable[str], + cache_salt_id: Optional[int] = None, + ) -> tuple[list[list[int]], dict[str, list[list[BlockHash]]]]: + """Synchronous tokenize + per-algorithm block hashes for thread offload.""" + token_lists = self._tokenize(request) + return token_lists, { + hash_algo: + self._compute_block_hashes(token_lists, + hash_algo, + cache_salt_id=cache_salt_id) + for hash_algo in set(hash_algos) + } + @staticmethod def _text_to_int_sequences(texts: list[str]) -> list[list[int]]: """Convert text strings to lists of unicode code points. @@ -786,6 +907,11 @@ def _text_to_int_sequences(texts: list[str]) -> list[list[int]]: """ return [[ord(c) for c in text] for text in texts] + @staticmethod + def _get_request_cache_salt_id(request: OpenAIRequest) -> Optional[int]: + cache_salt = getattr(request, "cache_salt", None) + return None if cache_salt is None else get_cache_salt_id(cache_salt) + class KvCacheAwareRouter(BlockHashMixin, LoadBalancingMixin, Router): @@ -800,6 +926,7 @@ def __init__(self, max_batch_size: int = 64, tokens_per_block: int = 32, custom_tokenizer: Optional[str] = None, + backfill_block_hashes_on_finish: bool = False, **kwargs): super().__init__(server_role, servers, metadata_server_cfg, metadata_server, **kwargs) @@ -807,68 +934,89 @@ def __init__(self, self._init_load_balancing(servers, use_tokens) # TODO: use max_num_tokens? per server? self._max_batch_size = max_batch_size + # Opt-in workaround for the disagg-gen path that doesn't emit + # kv_cache_events. Stash hashes at routing, inject on finish. + self._backfill_block_hashes_on_finish = backfill_block_hashes_on_finish + self._pending_block_hashes: dict[int, tuple[list, str]] = {} def _create_server_state(self, server: str) -> KvCacheAwareServerState: return KvCacheAwareServerState(server, self._use_tokens, self._tokens_per_block, lambda: self.session) + async def _get_server_hash_algo(self, server: str) -> str: + server_info = self._server_info.get(server, {}) + hash_algo = server_info.get("kv_cache_hash_algo") + return await self._server_state[server].get_hash_algo(hash_algo) + async def get_next_server( self, request: OpenAIRequest, exclude_server: Optional[str] = None) -> tuple[str, dict]: async with self._lock: - servers = [ + servers = list([ server for server in self._server_state.keys() if server != exclude_server - ] - block_hashes = [] - token_lists = [] + ]) + if not servers: + raise ValueError( + f"No available servers after excluding {exclude_server}") + cache_salt_id = self._get_request_cache_salt_id(request) + hash_algo_by_server = { + server: await self._get_server_hash_algo(server) + for server in servers + } + # Tokenize + block-hash is CPU-bound (~50 ms p50 for a 40 k-token + # chat request with a Rust-backed tokenizer). Running it directly + # inside the async handler blocks the orchestrator's event loop and + # serializes all concurrent requests through it; with HuggingFace + # tokenizers releasing the GIL, offloading to a thread lets multiple + # tokenize calls run in parallel and frees the event loop to + # dispatch HTTP traffic to the CTX/GEN workers meanwhile. + token_lists, block_hashes_by_algo = await asyncio.to_thread( + self._tokenize_and_compute_block_hashes_by_algo, request, + hash_algo_by_server.values(), cache_salt_id) + padded_tokens_by_algo = { + hash_algo: + max( + sum(len(hash_list) + for hash_list in block_hashes) * self._tokens_per_block, 1) + for hash_algo, block_hashes in block_hashes_by_algo.items() + } + # select the server by (KV match - load) + # TODO: more options + workloads = [ + self._server_state[server].num_active_requests() + for server in servers + ] + scores = [] matches = [] - if len(servers) == 0: - raise RuntimeError( - f"No available servers after excluding {exclude_server}") - elif len(servers) == 1: - server = servers[0] - else: - # Tokenize + block-hash is CPU-bound (~50 ms p50 for a 40 k-token - # chat request with a Rust-backed tokenizer). Running it directly - # inside the async handler blocks the orchestrator's event loop and - # serializes all concurrent requests through it; with HuggingFace - # tokenizers releasing the GIL, offloading to a thread lets multiple - # tokenize calls run in parallel and frees the event loop to - # dispatch HTTP traffic to the CTX/GEN workers meanwhile. - token_lists, block_hashes = await asyncio.to_thread( - self._tokenize_and_compute_block_hashes, request) - padded_tokens = sum( - len(hash_list) - for hash_list in block_hashes) * self._tokens_per_block - # select the server by (KV match - load) - # TODO: more options - workloads = [ - state.num_active_requests() - for state in self._server_state.values() - ] - scores = [] - matches = [] - for i in range(len(servers)): - server = servers[i] - # https://github.com/ai-dynamo/dynamo/blob/main/docs/kv_cache_routing.md#kv-cache-routing-and-load-balancing - matches.append( - await - self._server_state[server].matched_tokens(block_hashes)) - score = matches[-1] / padded_tokens - workloads[ - i] / self._max_batch_size - scores.append(score) - max_score = max(scores) - tied = [i for i, s in enumerate(scores) if s == max_score] - winner = tied[self._rr_counter % len(tied)] - self._rr_counter += 1 - server = servers[winner] + for i in range(len(servers)): + server = servers[i] + hash_algo = hash_algo_by_server[server] + block_hashes = block_hashes_by_algo[hash_algo] + padded_tokens = padded_tokens_by_algo[hash_algo] + # https://github.com/ai-dynamo/dynamo/blob/main/docs/kv_cache_routing.md#kv-cache-routing-and-load-balancing + matches.append(await self._server_state[server].matched_tokens( + block_hashes, hash_algo)) + score = matches[-1] / padded_tokens - workloads[ + i] / self._max_batch_size + scores.append(score) + max_score = max(scores) + tied = [i for i, s in enumerate(scores) if s == max_score] + winner = tied[self._rr_counter % len(tied)] + self._rr_counter += 1 + server = servers[winner] + hash_algo = hash_algo_by_server[server] + block_hashes = block_hashes_by_algo[hash_algo] async with self._lock: await self._register_request(server, request) + if self._backfill_block_hashes_on_finish: + self._pending_block_hashes[id(request)] = (block_hashes, + hash_algo) return server, { - "block_hashes": block_hashes, # list[list[int]] + "block_hashes": block_hashes, # list[list[int | str]] + "hash_algo": hash_algo, "token_lists": token_lists, # list[list[int]] "matches": matches, # list[int] "server_info": self._server_info.get(server, {}), @@ -876,12 +1024,24 @@ async def get_next_server( async def finish_request(self, request: OpenAIRequest, - session: Optional[aiohttp.ClientSession] = None): + session: Optional[aiohttp.ClientSession] = None, + success: bool = True): async with self._lock: server = self._req_routing_table.pop(id(request), None) + # Pop unconditionally to avoid leaks; only use for backfill on success. + pending = self._pending_block_hashes.pop(id(request), None) + if not (self._backfill_block_hashes_on_finish and success): + pending = None if server is not None and server in self._server_state: await self._server_state[server].decrement_load(request) if server is not None and server in self._server_state: + # Inject just-served block_hashes into server state — see __init__ note. + if pending is not None: + block_hashes, hash_algo = pending + flat_hashes = (h for hash_list in block_hashes + for h in hash_list) + self._server_state[server].add_blocks(flat_hashes, + hash_algo=hash_algo) await self._server_state[server].poll_and_update(session) def _on_servers_updated(self, old_servers, new_servers): @@ -1338,7 +1498,11 @@ async def get_next_server( return server, {"server_info": self._server_info.get(server, {})} - async def finish_request(self, request: OpenAIRequest): + async def finish_request(self, + request: OpenAIRequest, + session: Optional[aiohttp.ClientSession] = None, + success: bool = True): + del session, success async with self._lock: server = await self._unregister_request(request) self._remove_content_load(server, request) diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index c1d1e370ea2a..7b0881d547c1 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -1386,6 +1386,70 @@ def test_kv_cache_v2_nixl_python(self): self.MODEL_PATH) as llm: run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(60000) + @skip_pre_hopper + @pytest.mark.parametrize( + "enable_attention_dp,mtp_nextn", + [(False, 0), (True, 2)], + ids=["noadp-mtp0", "adp-mtp2"], + ) + def test_gen_first(self, enable_attention_dp, mtp_nextn): + """Gen-first MLA coverage on KVCacheManagerV2 + NIXL python; diagonal ADP/MTP combos.""" + kv_cache_config = { + "enable_block_reuse": False, + "use_kv_cache_manager_v2": True, + } + ctx_server_config = { + "tensor_parallel_size": 2, + "pipeline_parallel_size": 1, + "enable_attention_dp": enable_attention_dp, + "disable_overlap_scheduler": True, + "cuda_graph_config": None, + "cache_transceiver_config": { + "backend": "NIXL", + "transceiver_runtime": "PYTHON", + }, + "kv_cache_config": kv_cache_config, + } + gen_server_config = { + "tensor_parallel_size": 2, + "pipeline_parallel_size": 1, + "enable_attention_dp": enable_attention_dp, + "disable_overlap_scheduler": True, + "cuda_graph_config": None, + "cache_transceiver_config": { + "backend": "NIXL", + "transceiver_runtime": "PYTHON", + }, + "kv_cache_config": kv_cache_config, + } + if mtp_nextn > 0: + spec_config = { + "decoding_type": "MTP", + "max_draft_len": mtp_nextn, + } + ctx_server_config["speculative_config"] = spec_config + gen_server_config["speculative_config"] = spec_config + disaggregated_server_config = { + "hostname": "localhost", + "port": 8000, + "backend": "pytorch", + "context_servers": { + "num_instances": 1, + "urls": ["localhost:8001"] + }, + "generation_servers": { + "num_instances": 1, + "urls": ["localhost:8002"] + }, + "schedule_style": "generation_first", + } + with launch_disaggregated_llm(disaggregated_server_config, + ctx_server_config, gen_server_config, + self.MODEL_PATH) as llm: + run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) + @pytest.mark.timeout(DEFAULT_TEST_TIMEOUT) class TestGemma3_1BInstruct(LlmapiAccuracyTestHarness): @@ -1998,34 +2062,61 @@ def test_auto_dtype_with_helix(self, comms_medium, cuda_graph_config, self.MODEL_PATH) as llm: run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) - @pytest.mark.parametrize( - "gen_tp_pp", [(1, 1), (1, 2), (2, 1), (2, 2)], - ids=["gen_tp1pp1", "gen_tp1pp2", "gen_tp2pp1", "gen_tp2pp2"]) - @pytest.mark.parametrize( - "ctx_tp_pp", - [(1, 1), (1, 2), (2, 1), (2, 2), (1, 4)], - ids=[ - "ctx_tp1pp1", "ctx_tp1pp2", "ctx_tp2pp1", "ctx_tp2pp2", "ctx_tp1pp4" - ], - ) - @pytest.mark.parametrize("enable_attention_dp", [False, True], - ids=["noadp", "adp"]) - def test_gen_first(self, ctx_tp_pp, gen_tp_pp, enable_attention_dp): - ctx_tp, ctx_pp = ctx_tp_pp - gen_tp, gen_pp = gen_tp_pp - total_gpus = ctx_tp * ctx_pp + gen_tp * gen_pp - if total_gpus > get_device_count(): - pytest.skip(f"Not enough devices for {total_gpus} GPUs") + @pytest.mark.skip_less_device(2) + def test_gen_first(self): + """Gen-first dense-model smoke test on KVCacheManagerV2 + NIXL python.""" + kv_cache_config = { + "enable_block_reuse": False, + "use_kv_cache_manager_v2": True, + } + ctx_server_config = { + "disable_overlap_scheduler": True, + "cuda_graph_config": None, + "cache_transceiver_config": { + "backend": "NIXL", + "transceiver_runtime": "PYTHON", + }, + "kv_cache_config": kv_cache_config, + } + gen_server_config = { + "disable_overlap_scheduler": True, + "cuda_graph_config": None, + "cache_transceiver_config": { + "backend": "NIXL", + "transceiver_runtime": "PYTHON", + }, + "kv_cache_config": kv_cache_config, + } + disaggregated_server_config = { + "hostname": "localhost", + "port": 8000, + "backend": "pytorch", + "context_servers": { + "num_instances": 1, + "urls": ["localhost:8001"] + }, + "generation_servers": { + "num_instances": 1, + "urls": ["localhost:8002"] + }, + "schedule_style": "generation_first", + } + with launch_disaggregated_llm(disaggregated_server_config, + ctx_server_config, gen_server_config, + self.MODEL_PATH) as llm: + run_accuracy_test(llm, self.MODEL_NAME, ["MMLU"]) + + @pytest.mark.skip_less_device(2) + def test_gen_first_kv_cache_v1(self): + """Gen-first smoke test on the legacy V1 KV cache manager with block reuse.""" transceiver_runtime = "PYTHON" transceiver_backend = "NIXL" kv_cache_config = { "enable_block_reuse": True, "enable_partial_reuse": False, + "use_kv_cache_manager_v2": False, } ctx_server_config = { - "tensor_parallel_size": ctx_tp, - "pipeline_parallel_size": ctx_pp, - "enable_attention_dp": enable_attention_dp, "disable_overlap_scheduler": True, "cuda_graph_config": None, "cache_transceiver_config": { @@ -2035,9 +2126,6 @@ def test_gen_first(self, ctx_tp_pp, gen_tp_pp, enable_attention_dp): "kv_cache_config": kv_cache_config, } gen_server_config = { - "tensor_parallel_size": gen_tp, - "pipeline_parallel_size": gen_pp, - "enable_attention_dp": enable_attention_dp, "disable_overlap_scheduler": True, "cuda_graph_config": None, "cache_transceiver_config": { @@ -2063,7 +2151,6 @@ def test_gen_first(self, ctx_tp_pp, gen_tp_pp, enable_attention_dp): with launch_disaggregated_llm(disaggregated_server_config, ctx_server_config, gen_server_config, self.MODEL_PATH) as llm: - # MMLU is good enough for such a sanity test run_accuracy_test(llm, self.MODEL_NAME, ["MMLU"]) diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3_v2.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3_v2.yaml new file mode 100644 index 000000000000..851807eec4da --- /dev/null +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3_v2.yaml @@ -0,0 +1,42 @@ +hostname: localhost +model: DeepSeek-V3-Lite/bf16 +backend: pytorch +cuda_graph_config: null +free_gpu_memory_fraction: 0.15 +conditional_disagg_config: + max_local_prefill_length: 100 +disable_overlap_scheduler: true +enable_autotuner: false +context_servers: + num_instances: 1 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + kv_cache_config: + enable_block_reuse: true + enable_partial_reuse: true + event_buffer_max_size: 1024 + free_gpu_memory_fraction: 0.15 + # V2 cache manager — Python-native KV pool + radix tree. + use_kv_cache_manager_v2: true + cache_transceiver_config: + # V2 transceiver — Python-side NIXL driver, paired with cache manager V2. + backend: NIXL + transceiver_runtime: PYTHON +generation_servers: + num_instances: 1 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + router: + type: kv_cache_aware + # disagg-gen path doesn't emit kv_cache_events; backfill router-side state from + # successfully-routed requests so conditional_disagg can actually trigger. + backfill_block_hashes_on_finish: true + kv_cache_config: + enable_block_reuse: true + enable_partial_reuse: true + event_buffer_max_size: 1024 + free_gpu_memory_fraction: 0.15 + use_kv_cache_manager_v2: true + cache_transceiver_config: + backend: NIXL + transceiver_runtime: PYTHON diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index 7223f1b47fbb..65a98b09a94b 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -281,6 +281,8 @@ def get_test_config(test_desc, example_dir, test_root): f"{test_configs_root}/disagg_config_cache_aware_balance_deepseek_v3.yaml", "deepseek_v3_lite_bf16_conditional": f"{test_configs_root}/disagg_config_conditional_deepseek_v3.yaml", + "deepseek_v3_lite_bf16_conditional_v2": + f"{test_configs_root}/disagg_config_conditional_deepseek_v3_v2.yaml", "deepseek_v3_lite_fp8_tp1_two_mtp": f"{test_configs_root}/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_two_mtp.yaml", "deepseek_v3_lite_fp8_ctxpp2_gentp2_one_mtp": @@ -1728,6 +1730,43 @@ def test_disaggregated_deepseek_v3_lite_bf16_conditional( cwd=llm_venv.get_working_directory()) +# V2 variant of the conditional disagg test: KV manager V2 + Python NIXL transceiver. +@skip_no_hopper +@pytest.mark.parametrize("deepseek_v3_model_root", ['DeepSeek-V3-Lite-bf16'], + indirect=True) +def test_disaggregated_deepseek_v3_lite_bf16_conditional_v2( + disaggregated_test_root, disaggregated_example_root, llm_venv, + deepseek_v3_model_root): + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/bf16") + + # Bypassed requests clear disaggregated_params and skip add_per_request_metrics, + # so /perf_metrics records strictly fewer entries than requests issued. + last_total = [0] + + def _check_bypass_signal(server_url: str): + import requests as http_requests + resp = http_requests.get(f"{server_url}/perf_metrics", timeout=10) + assert resp.status_code == 200, \ + f"perf_metrics fetch failed: {resp.status_code}" + metrics = resp.json() + logger.info( + f"conditional_v2 perf_metrics len={len(metrics)} " + f"(bypassed requests absent — grows sub-linearly if bypass works)") + # At least one iteration's worth of requests should have been routed — + # i.e. perf_metrics is non-empty after the first client run. + assert metrics or last_total[0] > 0, \ + "no per-request metrics recorded; bypass logic may be misconfigured" + last_total[0] = len(metrics) + + run_disaggregated_test(disaggregated_example_root, + "deepseek_v3_lite_bf16_conditional_v2", + env=llm_venv._new_env, + extra_endpoints_test=_check_bypass_signal, + model_path=deepseek_v3_model_root, + cwd=llm_venv.get_working_directory()) + + @skip_no_hopper @pytest.mark.parametrize("deepseek_v3_model_root", ['DeepSeek-V3-Lite-fp8'], indirect=True) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index abf5990611a3..95af7b3bf43d 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -97,8 +97,10 @@ accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype_with_helix accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp1tp2cp2] accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp2tp1cp2] accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_chunked_prefill -accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first[noadp-ctx_tp2pp1-gen_tp1pp1] -accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first[adp-ctx_tp2pp1-gen_tp2pp1] +accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first +accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first_kv_cache_v1 +accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_first[noadp-mtp0] +accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_first[adp-mtp2] accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_nixl_backend accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[qwen3vl_2b_instruct] accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[nemotron_nano_v3_omni_fp8] @@ -860,6 +862,7 @@ disaggregated/test_disaggregated.py::test_disaggregated_ctxtp2pp2_gentp2pp2[Tiny disaggregated/test_disaggregated.py::test_disaggregated_cuda_graph[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_cache_aware_balance[DeepSeek-V3-Lite-bf16] disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_conditional[DeepSeek-V3-Lite-bf16] +disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_conditional_v2[DeepSeek-V3-Lite-bf16] disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_empty_batch[DeepSeek-V3-Lite-bf16] disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_attention_dp[DeepSeek-V3-Lite-fp8] disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_attention_dp_gen_only[DeepSeek-V3-Lite-fp8] diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index d65e8ba0cd70..f110a31303dd 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -137,8 +137,10 @@ l0_dgx_h100: - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_ctx_pp_gen_tp_asymmetric[MMLU-gen_tp=2-ctx_pp=2] - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_multi_instance[GSM8K] - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_multi_instance[MMLU] - - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first[noadp-ctx_tp2pp1-gen_tp1pp1] - - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first[adp-ctx_tp2pp1-gen_tp2pp1] + - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first + - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first_kv_cache_v1 + - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_first[noadp-mtp0] + - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_first[adp-mtp2] - disaggregated/test_auto_scaling.py::test_service_discovery[etcd-round_robin] - disaggregated/test_auto_scaling.py::test_worker_restart[etcd-load_balancing] - disaggregated/test_auto_scaling.py::test_worker_restart[etcd-round_robin] @@ -221,6 +223,7 @@ l0_dgx_h100: - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_overlap_cuda_graph[DeepSeek-V3-Lite-fp8] - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_cache_aware_balance[DeepSeek-V3-Lite-bf16] - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_conditional[DeepSeek-V3-Lite-bf16] + - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_conditional_v2[DeepSeek-V3-Lite-bf16] - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_empty_batch[DeepSeek-V3-Lite-bf16] - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_guided_decoding[xgrammar-mtp_nextn=2] - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_ctxpp2_gentp2_one_mtp[DeepSeek-V3-Lite-fp8] diff --git a/tests/unittest/disaggregated/test_disagg_openai_client.py b/tests/unittest/disaggregated/test_disagg_openai_client.py index 2538c5c92df0..efac516a150d 100644 --- a/tests/unittest/disaggregated/test_disagg_openai_client.py +++ b/tests/unittest/disaggregated/test_disagg_openai_client.py @@ -141,7 +141,9 @@ async def test_non_streaming_completion_request( assert isinstance(response, CompletionResponse) assert response.model == "test-model" mock_session.post.assert_called_once() - mock_router.finish_request.assert_called_once_with(completion_request) + mock_router.finish_request.assert_called_once_with( + completion_request, mock_session, success=True + ) @pytest.mark.asyncio async def test_streaming_completion_request( @@ -183,7 +185,9 @@ async def mock_iter_any(): for i, chunk in enumerate(chunks): assert chunk == dummy_data[i] mock_session.post.assert_called_once() - mock_router.finish_request.assert_called_once_with(streaming_completion_request) + mock_router.finish_request.assert_called_once_with( + streaming_completion_request, mock_session, success=True + ) @pytest.mark.asyncio async def test_request_with_custom_server( @@ -220,8 +224,11 @@ async def test_request_error_handling( with pytest.raises(aiohttp.ClientError): await openai_client.send_request(completion_request) - # Should finish request on error - mock_router.finish_request.assert_called_once_with(completion_request) + # Should finish request on error with success=False so the router + # doesn't backfill cache state for a request that didn't complete. + mock_router.finish_request.assert_called_once_with( + completion_request, mock_session, success=False + ) @pytest.mark.asyncio async def test_request_with_retry( From d4bb51462de12ec9dc6e7c92f589ffa37afb4d06 Mon Sep 17 00:00:00 2001 From: lishicheng1996-nv Date: Thu, 28 May 2026 08:15:04 +0800 Subject: [PATCH 04/28] [None][feat] Support tool-calling in KvCacheAwareRouter for disagg serving (feat/deepseek_v4) (#14611) Signed-off-by: Shicheng Li Signed-off-by: Fanrong Li Signed-off-by: lishicheng1996-nv --- tensorrt_llm/serve/router.py | 13 ++ tests/unittest/disaggregated/test_router.py | 175 +++++++++++++++++++- 2 files changed, 187 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 8ee01677e06f..05a410fb5946 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -792,6 +792,17 @@ def _tokenize(self, request: OpenAIRequest) -> list[list[int]]: if request.prompt_token_ids is not None: return [request.prompt_token_ids] tokenizer = self._get_tokenizer(request.model) + # Forward tools and chat_template_kwargs so custom tokenizers + # (e.g. DeepseekV32Tokenizer) render tool schemas and respect + # template flags like `thinking=true` when computing the prompt + # token ids used for cache-aware routing AND passed downstream + # (prompt_token_ids makes the worker skip re-tokenization). + tool_dicts = (None if getattr(request, "tools", None) is None else [ + tool.model_dump() if hasattr(tool, "model_dump") else tool + for tool in request.tools + ]) + chat_template_kwargs = (request.chat_template_kwargs if getattr( + request, "chat_template_kwargs", None) else {}) result = tokenizer.apply_chat_template( [ msg if isinstance(msg, dict) else dict(msg) @@ -800,6 +811,8 @@ def _tokenize(self, request: OpenAIRequest) -> list[list[int]]: add_generation_prompt=request.add_generation_prompt, tokenize=True, return_dict=False, + tools=tool_dicts, + **chat_template_kwargs, ) # Some custom tokenizers (e.g. DeepseekV32Tokenizer) return a # string from apply_chat_template even with tokenize=True. diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index 1accc2ec4102..769b6cce0e7d 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -8,8 +8,10 @@ from tensorrt_llm.llmapi.disagg_utils import RouterConfig from tensorrt_llm.serve.openai_protocol import (ChatCompletionRequest, + ChatCompletionToolsParam, CompletionRequest, - DisaggregatedParams) + DisaggregatedParams, + FunctionDefinition) from tensorrt_llm.serve.router import (BlockHashMixin, ConversationRouter, KvCacheAwareRouter, LoadBalancingRouter, RoundRobinRouter, create_router) @@ -927,3 +929,174 @@ async def test_finish_request_forwards_explicit_session(servers): call = explicit_session.post.call_args posted_url = call.args[0] if call.args else call.kwargs.get("url", "") assert "kv_cache_events" in posted_url + + +# --------------------------------------------------------------------------- +# _tokenize: tools + chat_template_kwargs forwarding (PR #13232) +# --------------------------------------------------------------------------- + + +def _get_weather_tool() -> ChatCompletionToolsParam: + """Build a sample tool definition matching the OpenAI schema.""" + return ChatCompletionToolsParam(function=FunctionDefinition( + name="get_current_weather", + description="Get the current weather for a city.", + parameters={ + "type": "object", + "properties": { + "city": { + "type": "string", + "description": "City name" + } + }, + "required": ["city"], + }, + )) + + +def _mock_tokenizer(token_ids=None): + """Return a mock tokenizer with a recorded apply_chat_template. + + ``apply_chat_template`` records its kwargs and returns the supplied + token id list. + """ + tok = mock.MagicMock() + tok.apply_chat_template.return_value = token_ids or [1, 2, 3, 4, 5] + return tok + + +@pytest.mark.parametrize("router_class", + [KvCacheAwareRouter, ConversationRouter]) +def test_tokenize_forwards_tools_and_chat_template_kwargs(router_class): + """Regression test for PR #13232. + + ``BlockHashMixin._tokenize`` must forward the request's ``tools`` (as a + list of dicts) and ``chat_template_kwargs`` to + ``tokenizer.apply_chat_template``. Without this, custom tokenizers that + render tool schemas into the prompt (e.g. DeepSeek-V3.2) produce + truncated token ids, breaking cache-aware routing decisions and the + ``prompt_token_ids`` handed to the worker downstream. + """ + router = router_class(server_role=None, + servers=["server1"], + use_tokens=False, + max_batch_size=32, + tokens_per_block=32) + + tok = _mock_tokenizer() + with mock.patch.object(router, "_get_tokenizer", return_value=tok): + req = ChatCompletionRequest( + model="TinyLlama", + messages=[{ + "role": "user", + "content": "what's the weather in Paris?" + }], + tools=[_get_weather_tool()], + chat_template_kwargs={"thinking": True}, + ) + router._tokenize(req) + + tok.apply_chat_template.assert_called_once() + kwargs = tok.apply_chat_template.call_args.kwargs + # tools must be forwarded as a list of dicts (model_dump), not the + # Pydantic objects themselves. + assert isinstance(kwargs["tools"], list) and len(kwargs["tools"]) == 1 + tool_dict = kwargs["tools"][0] + assert isinstance(tool_dict, dict) + assert tool_dict["type"] == "function" + assert tool_dict["function"]["name"] == "get_current_weather" + assert "parameters" in tool_dict["function"] + # chat_template_kwargs must be forwarded as **kwargs (not nested). + assert kwargs.get("thinking") is True + + +@pytest.mark.parametrize("router_class", + [KvCacheAwareRouter, ConversationRouter]) +def test_tokenize_without_tools_passes_none(router_class): + """Bare chat request: no tools, no chat_template_kwargs. + + ``apply_chat_template`` still runs but receives ``tools=None`` and no + extra keyword arguments. + """ + router = router_class(server_role=None, + servers=["server1"], + use_tokens=False, + max_batch_size=32, + tokens_per_block=32) + + tok = _mock_tokenizer() + with mock.patch.object(router, "_get_tokenizer", return_value=tok): + req = ChatCompletionRequest(model="TinyLlama", + messages=[{ + "role": "user", + "content": "hello" + }]) + router._tokenize(req) + + tok.apply_chat_template.assert_called_once() + kwargs = tok.apply_chat_template.call_args.kwargs + assert kwargs["tools"] is None + assert "thinking" not in kwargs + # messages and add_generation_prompt must still flow through unchanged. + assert kwargs["tokenize"] is True + + +def test_tokenize_preserves_empty_tools_list(): + """Preserve empty tools list distinct from ``None``. + + ``tools=[]`` is semantically distinct from ``tools=None``; preserve it + so the router's call matches what the worker's own + ``apply_chat_template`` would pass (see ``serve/openai_server.py`` + tool_dicts assignment). + """ + router = KvCacheAwareRouter(server_role=None, + servers=["server1"], + use_tokens=False, + max_batch_size=32, + tokens_per_block=32) + + tok = _mock_tokenizer() + with mock.patch.object(router, "_get_tokenizer", return_value=tok): + req = ChatCompletionRequest(model="TinyLlama", + messages=[{ + "role": "user", + "content": "hi" + }], + tools=[]) + router._tokenize(req) + + kwargs = tok.apply_chat_template.call_args.kwargs + assert kwargs["tools"] == [] + + +def test_tokenize_skipped_when_prompt_token_ids_already_set(): + """Skip tokenization when ``prompt_token_ids`` is already populated. + + When the caller pre-tokenizes (``prompt_token_ids`` set), the router + must not invoke ``apply_chat_template`` at all — the cached token ids + are returned as-is. + """ + router = KvCacheAwareRouter(server_role=None, + servers=["server1"], + use_tokens=False, + max_batch_size=32, + tokens_per_block=32) + + tok = _mock_tokenizer() + with mock.patch.object(router, "_get_tokenizer", + return_value=tok) as get_tok: + req = ChatCompletionRequest( + model="TinyLlama", + messages=[{ + "role": "user", + "content": "irrelevant" + }], + tools=[_get_weather_tool()], + chat_template_kwargs={"thinking": True}, + prompt_token_ids=[10, 20, 30], + ) + out = router._tokenize(req) + + assert out == [[10, 20, 30]] + get_tok.assert_not_called() + tok.apply_chat_template.assert_not_called() From fd1085a85ef47dccdf8566f898ba11f76871b8f4 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Sun, 31 May 2026 16:12:15 -0700 Subject: [PATCH 05/28] [None][perf] use native C++ BlockKeyHasher in disagg router to improve TTFT (#14767) Signed-off-by: Zheyu Fu Signed-off-by: Fanrong Li --- tensorrt_llm/serve/router.py | 18 +++++++++--- tests/unittest/disaggregated/test_router.py | 31 ++++++++++++++++++++- 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 05a410fb5946..398829c46215 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -22,6 +22,10 @@ import aiohttp from transformers import AutoTokenizer +from tensorrt_llm.bindings.internal.batch_manager import \ + BlockKey as _NativeBlockKey +from tensorrt_llm.bindings.internal.batch_manager import \ + BlockKeyHasher as _NativeBlockKeyHasher from tensorrt_llm.llmapi.disagg_utils import (MetadataServerConfig, RouterConfig, ServerRole) from tensorrt_llm.logger import logger @@ -735,10 +739,16 @@ async def finish_request(self, def block_key_hasher(token_ids: list[int], parent_hash: Optional[int] = None, cache_salt_id: Optional[int] = None) -> int: - return hash_v1_block_key( - token_ids, - parent_hash=0 if parent_hash is None else parent_hash, - cache_salt_id=cache_salt_id) + parent = 0 if parent_hash is None else parent_hash + # Fast path: the native C++ BlockKeyHasher is bit-exact with + # hash_v1_block_key and avoids the per-token Python loop. Its hash() binding + # takes no cache_salt_id, so fall back to Python only when a salt is set + # (rare opt-in; never in the unsalted agent/chat completion path). + if cache_salt_id is None: + return _NativeBlockKeyHasher.hash(_NativeBlockKey(token_ids), parent) + return hash_v1_block_key(token_ids, + parent_hash=parent, + cache_salt_id=cache_salt_id) def v2_sha256_block_hasher(token_ids: list[int], diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index 769b6cce0e7d..a4e2a3caeaca 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -1,5 +1,6 @@ import asyncio import copy +import random import threading from unittest import mock @@ -7,6 +8,7 @@ import pytest from tensorrt_llm.llmapi.disagg_utils import RouterConfig +from tensorrt_llm.runtime.kv_cache_hash import hash_v1_block_key from tensorrt_llm.serve.openai_protocol import (ChatCompletionRequest, ChatCompletionToolsParam, CompletionRequest, @@ -14,7 +16,34 @@ FunctionDefinition) from tensorrt_llm.serve.router import (BlockHashMixin, ConversationRouter, KvCacheAwareRouter, LoadBalancingRouter, - RoundRobinRouter, create_router) + RoundRobinRouter, block_key_hasher, + create_router) + + +def test_native_block_key_hasher_matches_python_v1(): + """Native C++ BlockKeyHasher must be bit-exact with hash_v1_block_key. + + Covers single blocks and the chained per-block pattern used by + BlockHashMixin._compute_block_hashes. + """ + random.seed(0) + for _ in range(500): + n = random.choice([1, 2, 31, 32, 33, 64, 97]) + toks = [random.randint(0, 300000) for _ in range(n)] + parent = random.choice([None, 0, random.randint(1, 2**64 - 1)]) + ref = hash_v1_block_key(toks, + parent_hash=0 if parent is None else parent) + assert block_key_hasher(toks, parent) == ref + + # chained per-block (the real _compute_block_hashes pattern) + toks = [random.randint(0, 300000) for _ in range(2000)] + parent = None + for t in range(0, len(toks) - 1, 32): + t_end = min(t + 32, len(toks) - 1) + h = block_key_hasher(toks[t:t_end], parent) + assert h == hash_v1_block_key( + toks[t:t_end], parent_hash=0 if parent is None else parent) + parent = h def _make_mock_aiohttp_session(return_value=None): From 6b831c1b240381c203cdc4db893d32fca31dd7a5 Mon Sep 17 00:00:00 2001 From: Liao Lanyu <108499334+lancelly@users.noreply.github.com> Date: Wed, 3 Jun 2026 22:50:03 +0800 Subject: [PATCH 06/28] [None][fix] adp router: gate in-transfer load accounting behind a config flag (#14860) Signed-off-by: Lanyu Liao Co-authored-by: Lanyu Liao Signed-off-by: Fanrong Li Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com> --- .../_torch/pyexecutor/scheduler/adp_router.py | 16 +++++++++++++--- tensorrt_llm/llmapi/llm_args.py | 13 +++++++++++++ 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index ea909b0292c5..e3f9601fc12a 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -192,6 +192,7 @@ def create( fair_share_multiplier=attention_dp_config.kv_cache_routing_fair_share_multiplier, cold_start_warmup=attention_dp_config.kv_cache_routing_cold_start_warmup, async_transfer_manager=async_transfer_manager, + account_for_in_transfer=attention_dp_config.kv_cache_routing_account_for_in_transfer, ) return DefaultADPRouter(dist=dist) @@ -451,6 +452,7 @@ def __init__( fair_share_multiplier: float = 2.0, cold_start_warmup: bool = False, async_transfer_manager=None, + account_for_in_transfer: bool = False, ): super().__init__(dist) self.kv_cache_manager = kv_cache_manager @@ -464,7 +466,12 @@ def __init__( set(range(self.dist.mapping.dp_size)) if cold_start_warmup else set() ) # Requests still sending KV to GEN are invisible in active_requests; - # fold them back in via the transfer manager (see create_rank_state). + # when ``account_for_in_transfer`` is True, fold them back in via the + # transfer manager (see create_rank_state). Disabled by default + # because counting them inflates num_active_requests reported to the + # inference loop, which then skips the idle-fetch wait and yields + # empty fetch cycles. + self.account_for_in_transfer = account_for_in_transfer self.async_transfer_manager = async_transfer_manager def create_rank_state( @@ -484,8 +491,11 @@ def _active_tokens(req) -> int: n_active_for_state = len(active_requests) # Fold in requests mid KV-transfer to GEN; they are removed from - # active_requests but still counted as per-rank load. - if self.async_transfer_manager is not None: + # active_requests but still counted as per-rank load. Gated behind + # ``account_for_in_transfer`` because the inflated active count + # makes the inference loop's idle-fetch wait expire even when no + # request is actually runnable, causing empty fetch cycles. + if self.account_for_in_transfer and self.async_transfer_manager is not None: in_transfer = self.async_transfer_manager.requests_in_transfer() num_active_tokens += sum(_active_tokens(r) for r in in_transfer.values()) n_active_for_state += len(in_transfer) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 4290b82bfb76..baf73bcb6c6c 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1275,6 +1275,19 @@ class AttentionDpConfig(StrictBaseModel): "scatter requests that would otherwise consolidate on a single warm " "rank, wasting prefill. Default False preserves pre-warmup routing. " "Only used when enable_kv_cache_aware_routing is True.") + kv_cache_routing_account_for_in_transfer: bool = Field( + default=False, + description= + "In-transfer load accounting in KV cache-aware routing. When True, " + "requests still streaming KV to the GEN worker (tracked by the " + "PyExecutor AsyncTransferManager but no longer in active_requests) " + "are folded back into the per-rank load reported via RankState. " + "This can improve balance under heavy disagg traffic but inflates " + "num_active_requests reported upstream, which lets the inference " + "loop's idle-fetch wait expire even when no requests are runnable " + "and causes empty fetch cycles. Default False preserves the prior " + "behaviour (fetch blocks when truly idle). " + "Only used when enable_kv_cache_aware_routing is True.") @model_validator(mode='after') def validate_attention_dp_config(self) -> 'AttentionDpConfig': From f37f24eb7fa984f3d8761590cf2f6e891d8904c2 Mon Sep 17 00:00:00 2001 From: Shi Xiaowei <39303645+Shixiaowei02@users.noreply.github.com> Date: Fri, 5 Jun 2026 22:27:28 +0800 Subject: [PATCH 07/28] [None][fix] Optimize cache-aware router backfill and validate worker handshake (#14950) Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com> Signed-off-by: Fanrong Li Signed-off-by: Shi Xiaowei <39303645+Shixiaowei02@users.noreply.github.com> --- tensorrt_llm/llmapi/disagg_utils.py | 7 + tensorrt_llm/serve/openai_server.py | 19 +- tensorrt_llm/serve/router.py | 160 +++- .../test_lists/test-db/l0_b200.yml | 2 + .../disaggregated/test_disagg_utils.py | 35 + .../disaggregated/test_openai_server_info.py | 82 ++ tests/unittest/disaggregated/test_router.py | 785 +++++++++++++++++- 7 files changed, 1050 insertions(+), 40 deletions(-) create mode 100644 tests/unittest/disaggregated/test_openai_server_info.py diff --git a/tensorrt_llm/llmapi/disagg_utils.py b/tensorrt_llm/llmapi/disagg_utils.py index b5d1fc8a1187..6afa17c2626e 100644 --- a/tensorrt_llm/llmapi/disagg_utils.py +++ b/tensorrt_llm/llmapi/disagg_utils.py @@ -287,6 +287,13 @@ def extract_router_config(server_cfg: dict) -> RouterConfig: if key in server_cfg: args[key] = server_cfg[key] + # tokens_per_block lives under kv_cache_config; the cache-aware router must + # use the same block size as the worker or block hashes never match. Carry + # the explicit server value over unless the router block already set it. + kv_cache_config = server_cfg.get("kv_cache_config") or {} + if "tokens_per_block" not in args and "tokens_per_block" in kv_cache_config: + args["tokens_per_block"] = kv_cache_config["tokens_per_block"] + return RouterConfig(type=router_type, args=args) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 8795bc90979f..385537f70089 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -1997,10 +1997,21 @@ async def update_weights(self, return JSONResponse(content={"status": "success"}) async def get_server_info(self) -> JSONResponse: - return JSONResponse( - content={ - "disaggregated_params": self.generator.disaggregated_params - }) + content = {"disaggregated_params": self.generator.disaggregated_params} + args = getattr(self.generator, "args", None) + if args is not None: + if args.max_batch_size is not None: + content["max_batch_size"] = args.max_batch_size + kv_cache_config = args.kv_cache_config + if kv_cache_config is not None: + content[ + "kv_cache_hash_algo"] = get_effective_kv_cache_event_hash_algo( + kv_cache_config.kv_cache_event_hash_algo, + kv_cache_config.use_kv_cache_manager_v2) + if kv_cache_config.tokens_per_block is not None: + content[ + "tokens_per_block"] = kv_cache_config.tokens_per_block + return JSONResponse(content=content) async def openai_image_generation(self, request: ImageGenerationRequest, raw_request: Request) -> Response: diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 398829c46215..25ca067b2740 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -136,6 +136,9 @@ def __init__( } self._kv_cache_hash_algo = KV_CACHE_HASH_ALGO_DEFAULT self._tokens_per_block = tokens_per_block + self._poll_task: Optional[asyncio.Task] = None + self._poll_pending: bool = False + self._poll_session = None @property def hash_algo(self) -> str: @@ -165,8 +168,7 @@ def add_blocks(self, hash_algo = self._resolve_hash_algo(hash_algo) self.set_hash_algo(hash_algo) block_table = self._block_table(hash_algo) - for block_hash in block_hashes: - block_table.add(block_hash) + block_table.update(block_hashes) def remove_blocks(self, block_hashes: Iterable[BlockHash], @@ -174,8 +176,7 @@ def remove_blocks(self, hash_algo = self._resolve_hash_algo(hash_algo) self.set_hash_algo(hash_algo) block_table = self._block_table(hash_algo) - for block_hash in block_hashes: - block_table.discard(block_hash) + block_table.difference_update(block_hashes) def update_with_events(self, events: Iterable[dict]): # event_raw: {"id": , "data": } @@ -237,6 +238,36 @@ async def poll_and_update(self, session=None): logger.warning( f"Failed to poll KV cache events from {self._server}: {e}") + def schedule_poll_and_update(self, session=None) -> None: + # Coalesce concurrent polls into one in-flight task, but record that a + # fresh poll was requested so the running task re-polls once more before + # exiting. Without the re-arm, a poll requested by the LAST + # finish_request while a poll is already in flight would be dropped and + # its KV events stranded until the next request (which may never come). + # The strong ref on _poll_task keeps the task alive for the GC. + self._poll_pending = True + self._poll_session = session + if self._poll_task is not None and not self._poll_task.done(): + return + self._poll_task = asyncio.create_task(self._drain_poll_and_update()) + + async def _drain_poll_and_update(self) -> None: + while self._poll_pending: + self._poll_pending = False + await self.poll_and_update(self._poll_session) + + async def cancel_poll_task(self) -> None: + # Cancel and await the background poll so shutdown leaves no orphaned + # task polling a closed session. + self._poll_pending = False + if self._poll_task is not None and not self._poll_task.done(): + self._poll_task.cancel() + try: + await self._poll_task + except asyncio.CancelledError: + pass + self._poll_task = None + def num_active_tokens(self): return self._num_active_tokens @@ -950,6 +981,8 @@ def __init__(self, tokens_per_block: int = 32, custom_tokenizer: Optional[str] = None, backfill_block_hashes_on_finish: bool = False, + load_weight: float = 1.0, + load_cap: float = float("inf"), **kwargs): super().__init__(server_role, servers, metadata_server_cfg, metadata_server, **kwargs) @@ -957,20 +990,86 @@ def __init__(self, self._init_load_balancing(servers, use_tokens) # TODO: use max_num_tokens? per server? self._max_batch_size = max_batch_size + self._load_weight = load_weight + self._load_cap = load_cap # Opt-in workaround for the disagg-gen path that doesn't emit # kv_cache_events. Stash hashes at routing, inject on finish. self._backfill_block_hashes_on_finish = backfill_block_hashes_on_finish - self._pending_block_hashes: dict[int, tuple[list, str]] = {} + self._inflight_backfill_hashes: dict[int, tuple[list[BlockHash], + str]] = {} def _create_server_state(self, server: str) -> KvCacheAwareServerState: return KvCacheAwareServerState(server, self._use_tokens, self._tokens_per_block, lambda: self.session) - async def _get_server_hash_algo(self, server: str) -> str: - server_info = self._server_info.get(server, {}) - hash_algo = server_info.get("kv_cache_hash_algo") - return await self._server_state[server].get_hash_algo(hash_algo) + async def close(self): + for state in self._server_state.values(): + await state.cancel_poll_task() + await super().close() + + def _stash_backfill_on_route(self, request: OpenAIRequest, + block_hashes: list[list[BlockHash]], + hash_algo: str) -> None: + if not self._backfill_block_hashes_on_finish: + return + flat = [h for hl in block_hashes for h in hl] + self._inflight_backfill_hashes[id(request)] = (flat, hash_algo) + + def _apply_backfill_on_finish(self, request: OpenAIRequest, + server: Optional[str], success: bool) -> None: + # Pop unconditionally to avoid leaks; apply only when eligible. + entry = self._inflight_backfill_hashes.pop(id(request), None) + if not (self._backfill_block_hashes_on_finish and success): + return + if entry is None: + return + if server is None or server not in self._server_state: + return + flat_block_hashes, hash_algo = entry + self._server_state[server].add_blocks(flat_block_hashes, + hash_algo=hash_algo) + + def _get_server_hash_algo(self, server: str) -> str: + # Lock-free attribute read; state is seeded at handshake and refreshed + # by update_with_events. + return self._server_state[server].hash_algo + + async def _prepare_server(self, server: str): + await super()._prepare_server(server) + if server not in self._prepared_ready_servers: + return + info = self._server_info.get(server, {}) + worker_tpb = info.get("tokens_per_block") + if worker_tpb is not None and worker_tpb != self._tokens_per_block: + # Block-hash chunking is router-side; a mismatch makes hashes never + # match the worker's KV cache. Fail fast. + self._prepared_ready_servers.discard(server) + self._server_info.pop(server, None) + raise RuntimeError( + f"tokens_per_block mismatch on {server}: " + f"router={self._tokens_per_block} worker={worker_tpb}. " + f"Align kv_cache_config.tokens_per_block to fix.") + worker_algo = info.get("kv_cache_hash_algo") + known_algos = { + KV_CACHE_HASH_ALGO_V1, + KV_CACHE_HASH_ALGO_V2, + KV_CACHE_HASH_ALGO_V2_SHA256_64, + } + if worker_algo is None: + # Silent default would map a V2 worker's hashes to V1. + logger.warning( + f"{server} did not expose kv_cache_hash_algo in /server_info; " + f"router will assume {KV_CACHE_HASH_ALGO_DEFAULT}.") + elif worker_algo not in known_algos: + self._prepared_ready_servers.discard(server) + self._server_info.pop(server, None) + raise RuntimeError( + f"Unknown kv_cache_hash_algo on {server}: {worker_algo!r}. " + f"Router supports {sorted(known_algos)}.") + else: + # Persist once so per-request reads can skip the lock+await. + self._server_state[server].set_hash_algo(worker_algo) async def get_next_server( self, @@ -986,7 +1085,7 @@ async def get_next_server( f"No available servers after excluding {exclude_server}") cache_salt_id = self._get_request_cache_salt_id(request) hash_algo_by_server = { - server: await self._get_server_hash_algo(server) + server: self._get_server_hash_algo(server) for server in servers } # Tokenize + block-hash is CPU-bound (~50 ms p50 for a 40 k-token @@ -1006,12 +1105,14 @@ async def get_next_server( for hash_list in block_hashes) * self._tokens_per_block, 1) for hash_algo, block_hashes in block_hashes_by_algo.items() } - # select the server by (KV match - load) - # TODO: more options + # select the server by (KV match - load), bounded by load_cap workloads = [ self._server_state[server].num_active_requests() for server in servers ] + load_fractions = [ + workloads[i] / self._max_batch_size for i in range(len(servers)) + ] scores = [] matches = [] for i in range(len(servers)): @@ -1022,11 +1123,19 @@ async def get_next_server( # https://github.com/ai-dynamo/dynamo/blob/main/docs/kv_cache_routing.md#kv-cache-routing-and-load-balancing matches.append(await self._server_state[server].matched_tokens( block_hashes, hash_algo)) - score = matches[-1] / padded_tokens - workloads[ - i] / self._max_batch_size + score = matches[-1] / padded_tokens - self._load_weight * \ + load_fractions[i] scores.append(score) - max_score = max(scores) - tied = [i for i, s in enumerate(scores) if s == max_score] + # Optional hard cap: drop servers at/over load_cap; fall back to all if + # none remain. Disabled by default (load_cap=inf) to match the original + # score-only selection. + candidate_idx = [ + i for i, lf in enumerate(load_fractions) if lf < self._load_cap + ] + if not candidate_idx: + candidate_idx = list(range(len(servers))) + max_score = max(scores[i] for i in candidate_idx) + tied = [i for i in candidate_idx if scores[i] == max_score] winner = tied[self._rr_counter % len(tied)] self._rr_counter += 1 server = servers[winner] @@ -1034,9 +1143,7 @@ async def get_next_server( block_hashes = block_hashes_by_algo[hash_algo] async with self._lock: await self._register_request(server, request) - if self._backfill_block_hashes_on_finish: - self._pending_block_hashes[id(request)] = (block_hashes, - hash_algo) + self._stash_backfill_on_route(request, block_hashes, hash_algo) return server, { "block_hashes": block_hashes, # list[list[int | str]] "hash_algo": hash_algo, @@ -1051,21 +1158,12 @@ async def finish_request(self, success: bool = True): async with self._lock: server = self._req_routing_table.pop(id(request), None) - # Pop unconditionally to avoid leaks; only use for backfill on success. - pending = self._pending_block_hashes.pop(id(request), None) - if not (self._backfill_block_hashes_on_finish and success): - pending = None if server is not None and server in self._server_state: await self._server_state[server].decrement_load(request) + self._apply_backfill_on_finish(request, server, success) if server is not None and server in self._server_state: - # Inject just-served block_hashes into server state — see __init__ note. - if pending is not None: - block_hashes, hash_algo = pending - flat_hashes = (h for hash_list in block_hashes - for h in hash_list) - self._server_state[server].add_blocks(flat_hashes, - hash_algo=hash_algo) - await self._server_state[server].poll_and_update(session) + # Fire-and-forget; poll runs in background and coalesces per server. + self._server_state[server].schedule_poll_and_update(session) def _on_servers_updated(self, old_servers, new_servers): new_state = {} diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index a958f8b24951..2e1dd92df995 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -281,6 +281,8 @@ l0_b200: - unittest/trt/attention/test_gpt_attention.py -k "trtllm_gen" - unittest/llmapi/test_llm_quant.py # 3.5 mins on B200 - unittest/trt/functional/test_fp4_gemm.py # 3 mins on B200 + - unittest/disaggregated/test_router.py + - unittest/disaggregated/test_openai_server_info.py - condition: ranges: system_gpu_count: diff --git a/tests/unittest/disaggregated/test_disagg_utils.py b/tests/unittest/disaggregated/test_disagg_utils.py index 4e62d802d30f..a086108187ad 100644 --- a/tests/unittest/disaggregated/test_disagg_utils.py +++ b/tests/unittest/disaggregated/test_disagg_utils.py @@ -142,6 +142,41 @@ def test_extract_router_config(sample_yaml_config): assert "max_num_tokens" not in gen_router_config.args +def test_extract_router_config_propagates_tokens_per_block(): + # tokens_per_block lives under kv_cache_config and must reach the router so + # block hashes match the worker's KV cache. + cfg = { + "router": { + "type": "kv_cache_aware" + }, + "max_batch_size": 8, + "kv_cache_config": { + "tokens_per_block": 128 + }, + } + router_config = extract_router_config(cfg) + assert router_config.args["tokens_per_block"] == 128 + + # An explicit router-level value is not overwritten. + cfg2 = { + "router": { + "type": "kv_cache_aware", + "tokens_per_block": 64 + }, + "kv_cache_config": { + "tokens_per_block": 128 + }, + } + assert extract_router_config(cfg2).args["tokens_per_block"] == 64 + + # No kv_cache_config -> nothing propagated. + assert "tokens_per_block" not in extract_router_config({ + "router": { + "type": "kv_cache_aware" + } + }).args + + def test_get_server_configs_dict(): server_configs = [ CtxGenServerConfig(type="ctx", diff --git a/tests/unittest/disaggregated/test_openai_server_info.py b/tests/unittest/disaggregated/test_openai_server_info.py new file mode 100644 index 000000000000..5f1aa142791f --- /dev/null +++ b/tests/unittest/disaggregated/test_openai_server_info.py @@ -0,0 +1,82 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +from types import SimpleNamespace + +import pytest + +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.runtime.kv_cache_hash import KV_CACHE_HASH_ALGO_V1, KV_CACHE_HASH_ALGO_V2 +from tensorrt_llm.serve.openai_server import OpenAIServer + + +def _make_server(*, max_batch_size=None, kv_cache_config=None): + server = object.__new__(OpenAIServer) + server.generator = SimpleNamespace( + disaggregated_params=None, + args=SimpleNamespace( + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + ), + ) + return server + + +@pytest.mark.asyncio +async def test_server_info_includes_effective_kv_cache_hash_algo(): + server = _make_server(kv_cache_config=KvCacheConfig()) + + response = await server.get_server_info() + content = json.loads(response.body) + assert content["kv_cache_hash_algo"] == KV_CACHE_HASH_ALGO_V1 + + server.generator.args.kv_cache_config = KvCacheConfig(use_kv_cache_manager_v2=True) + response = await server.get_server_info() + content = json.loads(response.body) + assert content["kv_cache_hash_algo"] == KV_CACHE_HASH_ALGO_V1 + + server.generator.args.kv_cache_config = KvCacheConfig( + use_kv_cache_manager_v2=True, kv_cache_event_hash_algo=KV_CACHE_HASH_ALGO_V2 + ) + response = await server.get_server_info() + content = json.loads(response.body) + assert content["kv_cache_hash_algo"] == KV_CACHE_HASH_ALGO_V2 + + +@pytest.mark.asyncio +async def test_server_info_includes_max_batch_size_for_router_normalization(): + server = _make_server(max_batch_size=256, kv_cache_config=KvCacheConfig()) + + response = await server.get_server_info() + content = json.loads(response.body) + assert content["max_batch_size"] == 256 + + +@pytest.mark.asyncio +async def test_server_info_omits_max_batch_size_when_none(): + server = _make_server(max_batch_size=None, kv_cache_config=KvCacheConfig()) + + response = await server.get_server_info() + content = json.loads(response.body) + assert "max_batch_size" not in content + + +@pytest.mark.asyncio +async def test_server_info_includes_tokens_per_block_from_kv_cache_config(): + server = _make_server(kv_cache_config=KvCacheConfig(tokens_per_block=64)) + + response = await server.get_server_info() + content = json.loads(response.body) + assert content["tokens_per_block"] == 64 diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index a4e2a3caeaca..bb9fe807596b 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -8,16 +8,27 @@ import pytest from tensorrt_llm.llmapi.disagg_utils import RouterConfig -from tensorrt_llm.runtime.kv_cache_hash import hash_v1_block_key +from tensorrt_llm.runtime.kv_cache_hash import (get_cache_salt_id, + hash_v1_block_key, + truncate_sha256_hash_to_int64) +from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import ( + ReuseScope, sequence_to_blockchain_keys) from tensorrt_llm.serve.openai_protocol import (ChatCompletionRequest, ChatCompletionToolsParam, CompletionRequest, DisaggregatedParams, FunctionDefinition) -from tensorrt_llm.serve.router import (BlockHashMixin, ConversationRouter, - KvCacheAwareRouter, LoadBalancingRouter, - RoundRobinRouter, block_key_hasher, - create_router) +# yapf: disable +from tensorrt_llm.serve.router import (KV_CACHE_HASH_ALGO_V1, + KV_CACHE_HASH_ALGO_V2, + KV_CACHE_HASH_ALGO_V2_SHA256_64, + BlockHashMixin, ConversationRouter, + KvCacheAwareRouter, + KvCacheAwareServerState, + LoadBalancingRouter, RoundRobinRouter, + block_key_hasher, create_router) + +# yapf: enable def test_native_block_key_hasher_matches_python_v1(): @@ -284,6 +295,770 @@ async def test_gen_tokens_balancing_router(servers, requests_fixture, request): await router.get_next_server(requests[0]) +def test_v2_sha256_block_hashes_match_kv_cache_manager_v2(servers): + tokens_per_block = 4 + token_lists = [[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]] + router = KvCacheAwareRouter(server_role=None, + servers=servers, + tokens_per_block=tokens_per_block) + + block_hashes = router._compute_block_hashes(token_lists, + hash_algo=KV_CACHE_HASH_ALGO_V2) + expected_block_hashes = [ + block_key.hex() + for token_block, block_key in sequence_to_blockchain_keys( + tokens_per_block, ReuseScope(), token_lists[0][:-1]) if token_block + ] + + assert block_hashes == [expected_block_hashes] + + +def test_v2_sha256_64_block_hashes_match_truncated_kv_cache_manager_v2(servers): + tokens_per_block = 4 + token_lists = [[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]] + router = KvCacheAwareRouter(server_role=None, + servers=servers, + tokens_per_block=tokens_per_block) + + block_hashes = router._compute_block_hashes( + token_lists, hash_algo=KV_CACHE_HASH_ALGO_V2_SHA256_64) + expected_block_hashes = [ + truncate_sha256_hash_to_int64(block_key) + for token_block, block_key in sequence_to_blockchain_keys( + tokens_per_block, ReuseScope(), token_lists[0][:-1]) if token_block + ] + + assert block_hashes == [expected_block_hashes] + assert all(isinstance(block_hash, int) for block_hash in block_hashes[0]) + + +def test_cache_aware_router_block_hashes_include_cache_salt_id(servers): + tokens_per_block = 4 + cache_salt_id = 123 + token_lists = [[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]] + router = KvCacheAwareRouter(server_role=None, + servers=servers, + tokens_per_block=tokens_per_block) + + salted_v1_block_hashes = router._compute_block_hashes( + token_lists, + hash_algo=KV_CACHE_HASH_ALGO_V1, + cache_salt_id=cache_salt_id) + expected_v1_block_hashes = [] + parent_hash = None + for t in range(0, len(token_lists[0]) - 1, tokens_per_block): + t_end = min(t + tokens_per_block, len(token_lists[0]) - 1) + parent_hash = hash_v1_block_key( + token_lists[0][t:t_end], + parent_hash=0 if parent_hash is None else parent_hash, + cache_salt_id=cache_salt_id) + expected_v1_block_hashes.append(parent_hash) + + salted_v2_block_hashes = router._compute_block_hashes( + token_lists, + hash_algo=KV_CACHE_HASH_ALGO_V2, + cache_salt_id=cache_salt_id) + reuse_scope = ReuseScope(salt=cache_salt_id) + expected_v2_block_hashes = [ + block_key.hex() + for token_block, block_key in sequence_to_blockchain_keys( + tokens_per_block, reuse_scope, token_lists[0][:-1]) if token_block + ] + + salted_v2_64_block_hashes = router._compute_block_hashes( + token_lists, + hash_algo=KV_CACHE_HASH_ALGO_V2_SHA256_64, + cache_salt_id=cache_salt_id) + expected_v2_64_block_hashes = [ + truncate_sha256_hash_to_int64(block_key) + for token_block, block_key in sequence_to_blockchain_keys( + tokens_per_block, reuse_scope, token_lists[0][:-1]) if token_block + ] + + assert salted_v1_block_hashes == [expected_v1_block_hashes] + assert salted_v2_block_hashes == [expected_v2_block_hashes] + assert salted_v2_64_block_hashes == [expected_v2_64_block_hashes] + assert salted_v1_block_hashes != router._compute_block_hashes( + token_lists, hash_algo=KV_CACHE_HASH_ALGO_V1) + assert salted_v2_block_hashes != router._compute_block_hashes( + token_lists, hash_algo=KV_CACHE_HASH_ALGO_V2) + assert salted_v2_64_block_hashes != router._compute_block_hashes( + token_lists, hash_algo=KV_CACHE_HASH_ALGO_V2_SHA256_64) + + +def test_cache_salt_id_derivation_matches_worker_path(): + from tensorrt_llm.inputs.utils import \ + get_cache_salt_id as worker_get_cache_salt_id + + assert get_cache_salt_id("abc") == 3697813978277427044 + assert get_cache_salt_id("tenant-a") == get_cache_salt_id("tenant-a") + assert get_cache_salt_id("tenant-a") != get_cache_salt_id("tenant-b") + assert get_cache_salt_id("tenant-a") == worker_get_cache_salt_id("tenant-a") + + +@pytest.mark.asyncio +async def test_kv_cache_aware_server_state_uses_hash_algo(): + tokens_per_block = 4 + token_lists = [[1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008]] + router = KvCacheAwareRouter(server_role=None, + servers=["server1"], + tokens_per_block=tokens_per_block) + v1_block_hashes = router._compute_block_hashes( + token_lists, hash_algo=KV_CACHE_HASH_ALGO_V1) + v2_block_hashes = router._compute_block_hashes( + token_lists, hash_algo=KV_CACHE_HASH_ALGO_V2) + v2_64_block_hashes = router._compute_block_hashes( + token_lists, hash_algo=KV_CACHE_HASH_ALGO_V2_SHA256_64) + state = KvCacheAwareServerState("server1", + tokens_per_block=tokens_per_block) + + state.update_with_events([{ + "event_id": 0, + "hash_algo": KV_CACHE_HASH_ALGO_V2, + "data": { + "type": + "stored", + "parent_hash": + None, + "blocks": [{ + "block_hash": block_hash + } for block_hash in v2_block_hashes[0]], + }, + }]) + + assert state.hash_algo == KV_CACHE_HASH_ALGO_V2 + assert await state.matched_tokens(v2_block_hashes, + hash_algo=KV_CACHE_HASH_ALGO_V2) == 8 + assert await state.matched_tokens(v1_block_hashes, + hash_algo=KV_CACHE_HASH_ALGO_V1) == 0 + state.add_blocks(["manual-v2-block"]) + assert state.hash_algo == KV_CACHE_HASH_ALGO_V2 + assert await state.matched_tokens([["manual-v2-block"]], + hash_algo=KV_CACHE_HASH_ALGO_V2) == 4 + assert await state.matched_tokens([["manual-v2-block"]], + hash_algo=KV_CACHE_HASH_ALGO_V1) == 0 + state.remove_blocks(["manual-v2-block"]) + assert await state.matched_tokens([["manual-v2-block"]], + hash_algo=KV_CACHE_HASH_ALGO_V2) == 0 + + state.update_with_events([{ + "event_id": 1, + "hash_algo": KV_CACHE_HASH_ALGO_V2_SHA256_64, + "data": { + "type": + "stored", + "parent_hash": + None, + "blocks": [{ + "block_hash": block_hash + } for block_hash in v2_64_block_hashes[0]], + }, + }]) + + assert state.hash_algo == KV_CACHE_HASH_ALGO_V2_SHA256_64 + assert await state.matched_tokens( + v2_64_block_hashes, hash_algo=KV_CACHE_HASH_ALGO_V2_SHA256_64) == 8 + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_routes_to_v2_server(servers): + tokens_per_block = 4 + token_lists = [[1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008]] + router = KvCacheAwareRouter(server_role=None, + servers=servers, + use_tokens=False, + max_batch_size=32, + tokens_per_block=tokens_per_block) + v2_block_hashes = router._compute_block_hashes( + token_lists, hash_algo=KV_CACHE_HASH_ALGO_V2) + v2_server = servers[1] + router._server_state[v2_server].update_with_events([{ + "event_id": 0, + "hash_algo": KV_CACHE_HASH_ALGO_V2, + "data": { + "type": + "stored", + "parent_hash": + None, + "blocks": [{ + "block_hash": block_hash + } for block_hash in v2_block_hashes[0]], + }, + }]) + + request = CompletionRequest(model="TinyLlama", + prompt=copy.deepcopy(token_lists)) + server, info = await router.get_next_server(request) + await router.finish_request(request) + + assert server == v2_server + assert info["hash_algo"] == KV_CACHE_HASH_ALGO_V2 + assert info["block_hashes"] == v2_block_hashes + assert info["matches"] == [0, 8, 0] + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_routes_with_cache_salt( + servers, monkeypatch): + + class SaltedRequest: + model = "TinyLlama" + cache_salt = "tenant-a" + + def __init__(self, prompt): + self.prompt = prompt + + tokens_per_block = 4 + cache_salt_id = 123 + token_lists = [[1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008]] + router = KvCacheAwareRouter(server_role=None, + servers=servers, + use_tokens=False, + max_batch_size=32, + tokens_per_block=tokens_per_block) + monkeypatch.setattr("tensorrt_llm.serve.router.get_cache_salt_id", + lambda cache_salt: cache_salt_id) + for server in servers: + router._server_info[server] = { + "kv_cache_hash_algo": KV_CACHE_HASH_ALGO_V2 + } + + salted_block_hashes = router._compute_block_hashes( + token_lists, + hash_algo=KV_CACHE_HASH_ALGO_V2, + cache_salt_id=cache_salt_id) + unsalted_block_hashes = router._compute_block_hashes( + token_lists, hash_algo=KV_CACHE_HASH_ALGO_V2) + salted_server = servers[1] + router._server_state[salted_server].update_with_events([{ + "event_id": 0, + "hash_algo": KV_CACHE_HASH_ALGO_V2, + "data": { + "type": + "stored", + "parent_hash": + None, + "blocks": [{ + "block_hash": block_hash + } for block_hash in salted_block_hashes[0]], + }, + }]) + + request = SaltedRequest(prompt=copy.deepcopy(token_lists)) + server, info = await router.get_next_server(request) + await router.finish_request(request) + + assert salted_block_hashes != unsalted_block_hashes + assert server == salted_server + assert info["hash_algo"] == KV_CACHE_HASH_ALGO_V2 + assert info["block_hashes"] == salted_block_hashes + assert info["matches"] == [0, 8, 0] + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_uses_server_info_hash_algo(servers): + tokens_per_block = 4 + token_lists = [[1000, 1001, 1002, 1003, 1004]] + router = KvCacheAwareRouter(server_role=None, + servers=[servers[0]], + tokens_per_block=tokens_per_block) + # _prepare_server seeds server state from the handshake; simulate it here. + router._server_info[servers[0]] = { + "kv_cache_hash_algo": KV_CACHE_HASH_ALGO_V2 + } + router._server_state[servers[0]].set_hash_algo(KV_CACHE_HASH_ALGO_V2) + + request = CompletionRequest(model="TinyLlama", + prompt=copy.deepcopy(token_lists)) + server, info = await router.get_next_server(request) + await router.finish_request(request) + + assert server == servers[0] + assert info["hash_algo"] == KV_CACHE_HASH_ALGO_V2 + assert router._server_state[servers[0]].hash_algo == KV_CACHE_HASH_ALGO_V2 + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_finish_request_polls_events(servers): + + class MockPostResponse: + + def __init__(self, events): + self._events = events + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + pass + + async def json(self): + return self._events + + class MockSession: + + def __init__(self, events): + self._events = events + self.post_url = None + + def post(self, url): + self.post_url = url + return MockPostResponse(self._events) + + tokens_per_block = 4 + token_lists = [[1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008]] + router = KvCacheAwareRouter(server_role=None, + servers=[servers[0]], + use_tokens=False, + max_batch_size=32, + tokens_per_block=tokens_per_block) + block_hashes = router._compute_block_hashes(token_lists, + hash_algo=KV_CACHE_HASH_ALGO_V1) + events = [{ + "event_id": 0, + "hash_algo": KV_CACHE_HASH_ALGO_V1, + "data": { + "type": + "stored", + "parent_hash": + None, + "blocks": [{ + "block_hash": block_hash + } for block_hash in block_hashes[0]], + }, + }] + session = MockSession(events) + + request = CompletionRequest(model="TinyLlama", + prompt=copy.deepcopy(token_lists)) + await router.get_next_server(request) + await router.finish_request(request, session) + # Poll runs in a background task; await it before asserting. + await router._server_state[servers[0]]._poll_task + + assert session.post_url == f"http://{servers[0]}/kv_cache_events" + assert await router._server_state[servers[0]].matched_tokens( + block_hashes, hash_algo=KV_CACHE_HASH_ALGO_V1) == 8 + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_finish_request_decrements_on_poll_error( + servers): + + class FailingSession: + + def post(self, url): + raise RuntimeError(f"failed to post to {url}") + + tokens_per_block = 4 + token_lists = [[1000, 1001, 1002, 1003, 1004]] + router = KvCacheAwareRouter(server_role=None, + servers=[servers[0]], + use_tokens=False, + max_batch_size=32, + tokens_per_block=tokens_per_block) + + request = CompletionRequest(model="TinyLlama", + prompt=copy.deepcopy(token_lists)) + await router.get_next_server(request) + assert router._server_state[servers[0]].num_active_requests() == 1 + + await router.finish_request(request, FailingSession()) + + assert router._server_state[servers[0]].num_active_requests() == 0 + assert id(request) not in router._req_routing_table + + +def test_kv_cache_aware_server_state_add_blocks_set_update_fast_path(): + state = KvCacheAwareServerState("server-x", tokens_per_block=4) + state.set_hash_algo(KV_CACHE_HASH_ALGO_V1) + + state.add_blocks([1, 2, 3], hash_algo=KV_CACHE_HASH_ALGO_V1) + assert state._block_table(KV_CACHE_HASH_ALGO_V1) == {1, 2, 3} + + state.add_blocks((h for h in [4, 5]), hash_algo=KV_CACHE_HASH_ALGO_V1) + assert state._block_table(KV_CACHE_HASH_ALGO_V1) == {1, 2, 3, 4, 5} + + state.add_blocks([1, 1, 2], hash_algo=KV_CACHE_HASH_ALGO_V1) + assert state._block_table(KV_CACHE_HASH_ALGO_V1) == {1, 2, 3, 4, 5} + + state.add_blocks([], hash_algo=KV_CACHE_HASH_ALGO_V1) + state.add_blocks(iter([]), hash_algo=KV_CACHE_HASH_ALGO_V1) + assert state._block_table(KV_CACHE_HASH_ALGO_V1) == {1, 2, 3, 4, 5} + + +@pytest.mark.asyncio +async def test_kv_cache_aware_server_state_schedule_poll_coalesces(): + state = KvCacheAwareServerState("server-x", tokens_per_block=4) + ready = asyncio.Event() + + async def slow_poll(session=None): + await ready.wait() + + with mock.patch.object(state, "poll_and_update", side_effect=slow_poll): + state.schedule_poll_and_update(session=None) + first_task = state._poll_task + state.schedule_poll_and_update(session=None) + assert state._poll_task is first_task + ready.set() + await first_task + + +@pytest.mark.asyncio +async def test_kv_cache_aware_server_state_schedule_poll_relaunches_after_done( +): + state = KvCacheAwareServerState("server-x", tokens_per_block=4) + calls = 0 + + async def fake_poll(session=None): + nonlocal calls + calls += 1 + + with mock.patch.object(state, "poll_and_update", side_effect=fake_poll): + state.schedule_poll_and_update(session=None) + await state._poll_task + first_task = state._poll_task + state.schedule_poll_and_update(session=None) + await state._poll_task + assert state._poll_task is not first_task + + assert calls == 2 + + +@pytest.mark.asyncio +async def test_kv_cache_aware_server_state_schedule_poll_rearms_pending(): + # A poll requested while one is in flight must re-run once more so the last + # finish_request can never strand its events behind a coalesced poll. + state = KvCacheAwareServerState("server-x", tokens_per_block=4) + calls = 0 + started = asyncio.Event() + release = asyncio.Event() + + async def slow_poll(session=None): + nonlocal calls + calls += 1 + started.set() + await release.wait() + + with mock.patch.object(state, "poll_and_update", side_effect=slow_poll): + state.schedule_poll_and_update(session=None) + await started.wait() # first poll is now actually in flight + started.clear() + # Second request arrives while the first poll is still blocked. + state.schedule_poll_and_update(session=None) + release.set() + await state._poll_task + + assert calls == 2 + + +@pytest.mark.asyncio +async def test_kv_cache_aware_server_state_cancel_poll_task(): + state = KvCacheAwareServerState("server-x", tokens_per_block=4) + started = asyncio.Event() + + async def blocking_poll(session=None): + started.set() + await asyncio.Event().wait() + + with mock.patch.object(state, "poll_and_update", side_effect=blocking_poll): + state.schedule_poll_and_update(session=None) + await started.wait() + await state.cancel_poll_task() + + assert state._poll_task is None + assert state._poll_pending is False + + +def test_kv_cache_aware_server_state_remove_blocks_silent_on_missing(): + state = KvCacheAwareServerState("server-x", tokens_per_block=4) + state.add_blocks([10, 20, 30], hash_algo=KV_CACHE_HASH_ALGO_V1) + + state.remove_blocks([20, 99, 100], hash_algo=KV_CACHE_HASH_ALGO_V1) + assert state._block_table(KV_CACHE_HASH_ALGO_V1) == {10, 30} + + state.remove_blocks([777, 888], hash_algo=KV_CACHE_HASH_ALGO_V1) + assert state._block_table(KV_CACHE_HASH_ALGO_V1) == {10, 30} + + state.remove_blocks((h for h in [10]), hash_algo=KV_CACHE_HASH_ALGO_V1) + assert state._block_table(KV_CACHE_HASH_ALGO_V1) == {30} + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_backfill_stashes_flat_list_at_routing( + servers): + tokens_per_block = 4 + token_lists = [[1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008]] + router = KvCacheAwareRouter(server_role=None, + servers=servers, + use_tokens=False, + max_batch_size=32, + tokens_per_block=tokens_per_block, + backfill_block_hashes_on_finish=True) + + request = CompletionRequest(model="TinyLlama", + prompt=copy.deepcopy(token_lists)) + server, info = await router.get_next_server(request) + + assert id(request) in router._inflight_backfill_hashes + stashed, stashed_algo = router._inflight_backfill_hashes[id(request)] + expected_flat = [h for hl in info["block_hashes"] for h in hl] + assert stashed == expected_flat + assert stashed and not isinstance(stashed[0], list) + assert stashed_algo == info["hash_algo"] + + await router.finish_request(request) + assert id(request) not in router._inflight_backfill_hashes + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_backfill_inserts_blocks_on_finish(servers): + tokens_per_block = 4 + token_lists = [[2000, 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008]] + router = KvCacheAwareRouter(server_role=None, + servers=servers, + use_tokens=False, + max_batch_size=32, + tokens_per_block=tokens_per_block, + backfill_block_hashes_on_finish=True) + + request = CompletionRequest(model="TinyLlama", + prompt=copy.deepcopy(token_lists)) + server, info = await router.get_next_server(request) + flat_expected = [h for hl in info["block_hashes"] for h in hl] + hash_algo = info["hash_algo"] + + assert await router._server_state[server].matched_tokens( + info["block_hashes"], hash_algo=hash_algo) == 0 + + await router.finish_request(request) + + assert router._server_state[server]._block_table(hash_algo).issuperset( + flat_expected) + assert await router._server_state[server].matched_tokens( + info["block_hashes"], hash_algo=hash_algo) == 8 + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_backfill_disabled_skips_pending(servers): + tokens_per_block = 4 + token_lists = [[3000, 3001, 3002, 3003, 3004, 3005, 3006, 3007, 3008]] + router = KvCacheAwareRouter(server_role=None, + servers=servers, + use_tokens=False, + max_batch_size=32, + tokens_per_block=tokens_per_block) + + request = CompletionRequest(model="TinyLlama", + prompt=copy.deepcopy(token_lists)) + server, info = await router.get_next_server(request) + + assert router._inflight_backfill_hashes == {} + + await router.finish_request(request) + + assert await router._server_state[server].matched_tokens( + info["block_hashes"], hash_algo=info["hash_algo"]) == 0 + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_backfill_drops_on_failure(servers): + tokens_per_block = 4 + token_lists = [[4000, 4001, 4002, 4003, 4004, 4005, 4006, 4007, 4008]] + router = KvCacheAwareRouter(server_role=None, + servers=servers, + use_tokens=False, + max_batch_size=32, + tokens_per_block=tokens_per_block, + backfill_block_hashes_on_finish=True) + + request = CompletionRequest(model="TinyLlama", + prompt=copy.deepcopy(token_lists)) + server, info = await router.get_next_server(request) + assert id(request) in router._inflight_backfill_hashes + + await router.finish_request(request, success=False) + + assert id(request) not in router._inflight_backfill_hashes + assert await router._server_state[server].matched_tokens( + info["block_hashes"], hash_algo=info["hash_algo"]) == 0 + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_load_cap_excludes_overloaded(servers): + tokens_per_block = 4 + token_lists = [[5000] * 16] + router = KvCacheAwareRouter(server_role=None, + servers=servers, + use_tokens=False, + max_batch_size=10, + tokens_per_block=tokens_per_block, + load_cap=0.8) + router._server_state[servers[0]]._num_active_requests = 9 + router._server_state[servers[1]]._num_active_requests = 9 + router._server_state[servers[2]]._num_active_requests = 1 + + request = CompletionRequest(model="TinyLlama", + prompt=copy.deepcopy(token_lists)) + chosen, _ = await router.get_next_server(request) + assert chosen == servers[2] + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_load_cap_fallback_when_all_overloaded( + servers): + router = KvCacheAwareRouter(server_role=None, + servers=servers, + use_tokens=False, + max_batch_size=10, + tokens_per_block=4, + load_cap=0.5) + router._server_state[servers[0]]._num_active_requests = 8 + router._server_state[servers[1]]._num_active_requests = 6 + router._server_state[servers[2]]._num_active_requests = 9 + + request = CompletionRequest(model="TinyLlama", prompt=[[100] * 16]) + chosen, _ = await router.get_next_server(request) + assert chosen == servers[1] + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_load_weight_scales_load_penalty(servers): + tokens_per_block = 4 + token_lists = [[6000] * 16] + router = KvCacheAwareRouter(server_role=None, + servers=servers, + use_tokens=False, + max_batch_size=10, + tokens_per_block=tokens_per_block, + load_weight=5.0, + load_cap=1.0) + block_hashes = router._compute_block_hashes(token_lists) + router._server_state[servers[0]].add_blocks( + [h for hl in block_hashes for h in hl]) + router._server_state[servers[0]]._num_active_requests = 7 + router._server_state[servers[1]]._num_active_requests = 1 + router._server_state[servers[2]]._num_active_requests = 1 + + request = CompletionRequest(model="TinyLlama", + prompt=copy.deepcopy(token_lists)) + chosen, _ = await router.get_next_server(request) + assert chosen != servers[0] + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_prepare_server_raises_on_tpb_mismatch(): + router = KvCacheAwareRouter(server_role=None, + servers=["server-a"], + tokens_per_block=32) + + async def fake_fetch(server, timeout): + return {"tokens_per_block": 64, "kv_cache_hash_algo": "v1_block_key"} + + with mock.patch.object(router, "_fetch_server_info", + side_effect=fake_fetch): + with pytest.raises(RuntimeError, match="tokens_per_block mismatch"): + await router._prepare_server("server-a") + + assert "server-a" not in router._prepared_ready_servers + assert "server-a" not in router._server_info + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_prepare_server_keeps_on_tpb_match(): + router = KvCacheAwareRouter(server_role=None, + servers=["server-a"], + tokens_per_block=32) + + async def fake_fetch(server, timeout): + return {"tokens_per_block": 32, "kv_cache_hash_algo": "v1_block_key"} + + with mock.patch.object(router, "_fetch_server_info", + side_effect=fake_fetch): + await router._prepare_server("server-a") + + assert "server-a" in router._prepared_ready_servers + assert router._server_info["server-a"]["tokens_per_block"] == 32 + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_prepare_server_keeps_when_worker_omits_tpb( +): + router = KvCacheAwareRouter(server_role=None, + servers=["server-a"], + tokens_per_block=32) + + async def fake_fetch(server, timeout): + return {"kv_cache_hash_algo": "v1_block_key"} + + with mock.patch.object(router, "_fetch_server_info", + side_effect=fake_fetch): + await router._prepare_server("server-a") + + assert "server-a" in router._prepared_ready_servers + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_prepare_server_warns_on_missing_algo(): + router = KvCacheAwareRouter(server_role=None, + servers=["server-a"], + tokens_per_block=32) + + async def fake_fetch(server, timeout): + return {"tokens_per_block": 32} + + with mock.patch("tensorrt_llm.serve.router.logger") as mock_logger: + with mock.patch.object(router, + "_fetch_server_info", + side_effect=fake_fetch): + await router._prepare_server("server-a") + + assert "server-a" in router._prepared_ready_servers + assert any("did not expose kv_cache_hash_algo" in str(call.args[0]) + for call in mock_logger.warning.call_args_list) + assert router._get_server_hash_algo("server-a") == KV_CACHE_HASH_ALGO_V1 + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_prepare_server_raises_on_unknown_algo(): + router = KvCacheAwareRouter(server_role=None, + servers=["server-a"], + tokens_per_block=32) + + async def fake_fetch(server, timeout): + return {"tokens_per_block": 32, "kv_cache_hash_algo": "bogus_algo"} + + with mock.patch.object(router, "_fetch_server_info", + side_effect=fake_fetch): + with pytest.raises(RuntimeError, match="Unknown kv_cache_hash_algo"): + await router._prepare_server("server-a") + + assert "server-a" not in router._prepared_ready_servers + assert "server-a" not in router._server_info + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_prepare_server_persists_algo_for_lock_free_read( +): + router = KvCacheAwareRouter(server_role=None, + servers=["server-a"], + tokens_per_block=32) + + async def fake_fetch(server, timeout): + return { + "tokens_per_block": 32, + "kv_cache_hash_algo": KV_CACHE_HASH_ALGO_V2, + } + + with mock.patch.object(router, "_fetch_server_info", + side_effect=fake_fetch): + await router._prepare_server("server-a") + + # Per-request read is now sync; no lock/await on the hot path. + assert router._get_server_hash_algo("server-a") == KV_CACHE_HASH_ALGO_V2 + assert router._server_state["server-a"].hash_algo == KV_CACHE_HASH_ALGO_V2 + + @pytest.mark.asyncio async def test_kv_cache_aware_router(servers, mock_aiohttp_session): # create tokenized requests to skip tokenization From 22689ce56bc4562a0c12dc357aa21a7001b9bdfa Mon Sep 17 00:00:00 2001 From: Shi Xiaowei <39303645+Shixiaowei02@users.noreply.github.com> Date: Sat, 6 Jun 2026 11:58:59 +0800 Subject: [PATCH 08/28] [None][fix] Relax tokens_per_block mismatch check in dis-agg router (#15005) Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com> Signed-off-by: Fanrong Li Signed-off-by: Shi Xiaowei <39303645+Shixiaowei02@users.noreply.github.com> --- .../native/mixers/attention/peer.py | 21 ++++++++++++- tensorrt_llm/serve/router.py | 23 +++++++++----- tests/unittest/disaggregated/test_peer.py | 30 +++++++++++++++++++ tests/unittest/disaggregated/test_router.py | 25 ++++++++++++++-- 4 files changed, 88 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py index d1458b47c870..ed6a8d84a887 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py @@ -335,6 +335,25 @@ def _mismatch(self, field: str, local, peer) -> bool: local != peer, f"{field} mismatch", field=field, local=local, peer=peer ) + def _tpb_check(self, local: int, peer: int) -> bool: + if local == peer: + return False + larger, smaller = max(local, peer), min(local, peer) + if larger % smaller != 0: + logger.warning( + "AttentionPolicy: incompatible: tokens_per_block not divisible; local=%d peer=%d", + local, + peer, + ) + return True + logger.warning( + "AttentionPolicy: tokens_per_block mismatch (local=%d, peer=%d); " + "KV transfer proceeds — ensure block boundaries align during transfer.", + local, + peer, + ) + return False + def check_peer_compatible(self, peer_ri: RankInfo) -> bool: a = self._ri.attention b = peer_ri.attention @@ -348,7 +367,7 @@ def check_peer_compatible(self, peer_ri: RankInfo) -> bool: peer=peer_ri.cp_size, ) or self._mismatch("element_bytes", a.element_bytes, b.element_bytes) - or self._mismatch("tokens_per_block", a.tokens_per_block, b.tokens_per_block) + or self._tpb_check(a.tokens_per_block, b.tokens_per_block) or self._mismatch("dims_per_head", a.dims_per_head, b.dims_per_head) or self._fail_if( a.is_mla and (a.kv_heads_per_rank != 1 or b.kv_heads_per_rank != 1), diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 25ca067b2740..a694c38c5025 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -1042,14 +1042,21 @@ async def _prepare_server(self, server: str): info = self._server_info.get(server, {}) worker_tpb = info.get("tokens_per_block") if worker_tpb is not None and worker_tpb != self._tokens_per_block: - # Block-hash chunking is router-side; a mismatch makes hashes never - # match the worker's KV cache. Fail fast. - self._prepared_ready_servers.discard(server) - self._server_info.pop(server, None) - raise RuntimeError( - f"tokens_per_block mismatch on {server}: " - f"router={self._tokens_per_block} worker={worker_tpb}. " - f"Align kv_cache_config.tokens_per_block to fix.") + larger = max(worker_tpb, self._tokens_per_block) + smaller = min(worker_tpb, self._tokens_per_block) + if larger % smaller != 0: + self._prepared_ready_servers.discard(server) + self._server_info.pop(server, None) + raise RuntimeError( + f"tokens_per_block mismatch on {server}: " + f"router={self._tokens_per_block} worker={worker_tpb} are not divisible. " + f"Align kv_cache_config.tokens_per_block so that one evenly divides the other." + ) + logger.warning( + "tokens_per_block mismatch on %s: router=%d worker=%d. " + "KV events from worker will not align with router block hashes; " + "use backfill_block_hashes_on_finish=true for best hit rate.", + server, self._tokens_per_block, worker_tpb) worker_algo = info.get("kv_cache_hash_algo") known_algos = { KV_CACHE_HASH_ALGO_V1, diff --git a/tests/unittest/disaggregated/test_peer.py b/tests/unittest/disaggregated/test_peer.py index d7a9f9913f8b..da595cfd093c 100644 --- a/tests/unittest/disaggregated/test_peer.py +++ b/tests/unittest/disaggregated/test_peer.py @@ -403,3 +403,33 @@ def test_peer_registrar_get_kv_map_head_mismatch(): reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) assert isinstance(mapper, HeadMismatchMapper) + + +def test_peer_registrar_tpb_divisible_warns_but_compatible(): + # local=16, peer=32: 32 % 16 == 0 → compatible with warning, register succeeds + self_rankinfo = make_rankinfo(instance_name="local", tokens_per_block=16) + reg = _make_peer_registrar(self_rankinfo) + peer_ri = make_rankinfo( + instance_name="peer", + instance_rank=5, + tokens_per_block=32, + layer_num_per_pp=[2], + page_table=make_page_table(), + ) + reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) + assert reg.get_peer_rank_info(peer_ri.instance_name, peer_ri.instance_rank) is not None + + +def test_peer_registrar_tpb_not_divisible_raises(): + # local=16, peer=24: 24 % 16 != 0 → incompatible, register raises ValueError + self_rankinfo = make_rankinfo(instance_name="local", tokens_per_block=16) + reg = _make_peer_registrar(self_rankinfo) + peer_ri = make_rankinfo( + instance_name="peer", + instance_rank=6, + tokens_per_block=24, + layer_num_per_pp=[2], + page_table=make_page_table(), + ) + with pytest.raises(ValueError): + reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index bb9fe807596b..3d0d82fd6fa0 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -947,7 +947,9 @@ async def test_kv_cache_aware_router_load_weight_scales_load_penalty(servers): @pytest.mark.asyncio -async def test_kv_cache_aware_router_prepare_server_raises_on_tpb_mismatch(): +async def test_kv_cache_aware_router_prepare_server_warns_on_tpb_mismatch_divisible( +): + # router=32, worker=64: 64 % 32 == 0 → warning only, server stays ready router = KvCacheAwareRouter(server_role=None, servers=["server-a"], tokens_per_block=32) @@ -957,7 +959,26 @@ async def fake_fetch(server, timeout): with mock.patch.object(router, "_fetch_server_info", side_effect=fake_fetch): - with pytest.raises(RuntimeError, match="tokens_per_block mismatch"): + await router._prepare_server("server-a") + + assert "server-a" in router._prepared_ready_servers + assert router._server_info["server-a"]["tokens_per_block"] == 64 + + +@pytest.mark.asyncio +async def test_kv_cache_aware_router_prepare_server_raises_on_tpb_not_divisible( +): + # router=32, worker=48: 48 % 32 != 0 → RuntimeError, server removed + router = KvCacheAwareRouter(server_role=None, + servers=["server-a"], + tokens_per_block=32) + + async def fake_fetch(server, timeout): + return {"tokens_per_block": 48, "kv_cache_hash_algo": "v1_block_key"} + + with mock.patch.object(router, "_fetch_server_info", + side_effect=fake_fetch): + with pytest.raises(RuntimeError, match="not divisible"): await router._prepare_server("server-a") assert "server-a" not in router._prepared_ready_servers From 9195548e0a46150d8b0727ebe778e973092ca608 Mon Sep 17 00:00:00 2001 From: lishicheng1996-nv Date: Mon, 8 Jun 2026 09:32:43 +0800 Subject: [PATCH 09/28] [None][perf] ConversationRouter: skip block-hash compute on sticky conversation_id routing (#14957) Signed-off-by: Fanrong Li Signed-off-by: lishicheng1996-nv --- tensorrt_llm/serve/router.py | 126 +++++++++++++++-------------------- 1 file changed, 53 insertions(+), 73 deletions(-) diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index a694c38c5025..c2a28d6c80e2 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -1540,91 +1540,71 @@ async def get_next_server( exclude_server: Optional[str] = None) -> tuple[str, dict]: self._validate_servers_available() - # Pre-compute outside the lock (tokenization + hashing) conv_id = self._get_conversation_id(request) + + # Explicit conversation_id: route by the session table alone and never + # touch request content. Block hashes are consumed only by the implicit + # prefix-match path below (requests with no conversation_id), so for a + # pinned session -- new or already-seen -- the GIL-bound + # _request_to_block_hashes is pure overhead that would serialize the + # single orchestrator event loop and cap dispatch throughput. + if conv_id: + weight = self._estimate_content_weight(request) + async with self._lock: + entry = self._session_table.get(conv_id) + if (entry is not None and entry[0] in self._server_state + and entry[0] != exclude_server): + # Sticky hit: keep the pinned server, refresh LRU only. + server = entry[0] + self._session_table.move_to_end(conv_id) + await self._register_request(server, request) + self._add_content_load(server, request, weight) + else: + # New conversation_id (or its server is gone/excluded): + # pin to the least-loaded server. Store no block hashes -- + # the trie is only read for conversation_id-less requests. + server = self._select_least_loaded(exclude_server) + if server is None: + raise ValueError("No available servers after excluding " + f"{exclude_server}") + await self._register_request(server, request) + self._add_content_load(server, request, weight) + self._update_session(conv_id, server, []) + return server, { + "server_info": self._server_info.get(server, {}) + } + + # No conversation_id: content-based routing. Compute block hashes + # (outside the lock) for implicit prefix matching, else least-loaded. block_hashes = self._request_to_block_hashes(request) weight = self._estimate_content_weight(request, block_hashes) async with self._lock: + matched_id = self._find_matching_session(block_hashes, + exclude_server) + if matched_id is not None: + server, _ = self._session_table[matched_id] + self._update_session(matched_id, server, block_hashes) + await self._register_request(server, request) + self._add_content_load(server, request, weight) + logger.debug( + f"ConversationRouter: IMPLICIT match conv_id={matched_id} " + f"-> server={server}, weight={weight}") + return server, { + "server_info": self._server_info.get(server, {}) + } - # 1. Explicit conversation_id — sticky routing. - # Always honour session affinity when the server is alive - # and not explicitly excluded. No overload gate — the - # server itself provides backpressure. - if conv_id and conv_id in self._session_table: - sticky_server, _ = self._session_table[conv_id] - if sticky_server not in self._server_state: - logger.debug( - f"ConversationRouter: STICKY MISS conv_id={conv_id} " - f"-> server={sticky_server} NOT in server_state, " - f"falling through to FALLBACK") - elif sticky_server == exclude_server: - logger.debug( - f"ConversationRouter: STICKY MISS conv_id={conv_id} " - f"-> server={sticky_server} is exclude_server") - else: - self._update_session(conv_id, sticky_server, block_hashes) - await self._register_request(sticky_server, request) - self._add_content_load(sticky_server, request, weight) - loads = { - s: self._get_content_load(s) - for s in self._servers - } - logger.debug( - f"ConversationRouter: STICKY conv_id={conv_id} " - f"-> server={sticky_server}, " - f"content_loads={loads}, weight={weight}") - return sticky_server, { - "server_info": self._server_info.get(sticky_server, {}) - } - elif conv_id: - logger.debug(f"ConversationRouter: NEW conv_id={conv_id} " - f"not in session_table " - f"(size={len(self._session_table)})") - - # 2. Implicit block-hash prefix matching. - # Always honour match when the server is alive. - matched_id = None - if not conv_id: - matched_id = self._find_matching_session( - block_hashes, exclude_server) - if matched_id is not None: - sticky_server, _ = self._session_table[matched_id] - self._update_session(matched_id, sticky_server, - block_hashes) - await self._register_request(sticky_server, request) - self._add_content_load(sticky_server, request, weight) - loads = { - s: self._get_content_load(s) - for s in self._servers - } - logger.debug( - f"ConversationRouter: IMPLICIT match " - f"conv_id={matched_id} -> server={sticky_server}, " - f"content_loads={loads}, weight={weight}") - return sticky_server, { - "server_info": self._server_info.get(sticky_server, {}) - } - - # 3. Fallback — least-loaded server for new sessions or - # sessions whose sticky server is unavailable. server = self._select_least_loaded(exclude_server) if server is None: raise ValueError( f"No available servers after excluding {exclude_server}") await self._register_request(server, request) self._add_content_load(server, request, weight) - - # Store session mapping. - if not conv_id: - conv_id = self._generate_implicit_id() - self._update_session(conv_id, server, block_hashes) - loads = {s: self._get_content_load(s) for s in self._servers} - logger.debug( - f"ConversationRouter: FALLBACK conv_id={conv_id} " - f"-> server={server}, content_loads={loads}, weight={weight}") - - return server, {"server_info": self._server_info.get(server, {})} + implicit_id = self._generate_implicit_id() + self._update_session(implicit_id, server, block_hashes) + logger.debug(f"ConversationRouter: FALLBACK conv_id={implicit_id} " + f"-> server={server}, weight={weight}") + return server, {"server_info": self._server_info.get(server, {})} async def finish_request(self, request: OpenAIRequest, From 460418e53651237f70c448baf5e22c4a07edcfeb Mon Sep 17 00:00:00 2001 From: Shi Xiaowei <39303645+Shixiaowei02@users.noreply.github.com> Date: Mon, 8 Jun 2026 10:18:07 +0800 Subject: [PATCH 10/28] [None][feat] Introduce tokenizer cache and absolute block score (#15040) Signed-off-by: Fanrong Li Signed-off-by: Shi Xiaowei <39303645+Shixiaowei02@users.noreply.github.com> --- tensorrt_llm/serve/router.py | 57 ++++---- tests/unittest/disaggregated/test_router.py | 139 +++++++++++++++++++- 2 files changed, 171 insertions(+), 25 deletions(-) diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index c2a28d6c80e2..05166e706cd8 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -827,40 +827,57 @@ def _get_tokenizer(self, model: str): tokenizer, model, trust_remote_code=True) return self._tokenizers[model] + def _encode_with_prefix_cache(self, rendered: str, key: int, + tokenizer) -> list[int]: + cache = getattr(self, "_tok_prefix_cache", None) + if cache is None: + cache = self._tok_prefix_cache = OrderedDict() + entry = cache.get(key) + if entry is not None and len(rendered) > len(entry[0]) and \ + rendered.startswith(entry[0]): + ids = entry[1] + tokenizer.encode(rendered[len(entry[0]):], + add_special_tokens=False) + else: + ids = tokenizer.encode(rendered, add_special_tokens=False) + cache[key] = (rendered, ids) + cache.move_to_end(key) + while len(cache) > 1024: + cache.popitem(last=False) + return ids + def _tokenize(self, request: OpenAIRequest) -> list[list[int]]: # Handle ChatCompletionRequest (has messages, not prompt) if isinstance(request, ChatCompletionRequest): if request.prompt_token_ids is not None: return [request.prompt_token_ids] tokenizer = self._get_tokenizer(request.model) - # Forward tools and chat_template_kwargs so custom tokenizers - # (e.g. DeepseekV32Tokenizer) render tool schemas and respect - # template flags like `thinking=true` when computing the prompt - # token ids used for cache-aware routing AND passed downstream - # (prompt_token_ids makes the worker skip re-tokenization). tool_dicts = (None if getattr(request, "tools", None) is None else [ tool.model_dump() if hasattr(tool, "model_dump") else tool for tool in request.tools ]) chat_template_kwargs = (request.chat_template_kwargs if getattr( request, "chat_template_kwargs", None) else {}) - result = tokenizer.apply_chat_template( + rendered = tokenizer.apply_chat_template( [ msg if isinstance(msg, dict) else dict(msg) for msg in request.messages ], add_generation_prompt=request.add_generation_prompt, - tokenize=True, + tokenize=False, return_dict=False, tools=tool_dicts, **chat_template_kwargs, ) - # Some custom tokenizers (e.g. DeepseekV32Tokenizer) return a - # string from apply_chat_template even with tokenize=True. - # Encode to token IDs if needed. - if isinstance(result, str): - result = tokenizer.encode(result, add_special_tokens=False) - # Set prompt_token_ids so the worker server skips re-tokenization + if isinstance(rendered, str): + key = hash("".join( + str( + msg.get("content") if isinstance(msg, dict) else + getattr(msg, "content", "")) + for msg in request.messages[:2])) + result = self._encode_with_prefix_cache(rendered, key, + tokenizer) + else: + result = list(rendered) request.prompt_token_ids = result return [result] @@ -981,7 +998,7 @@ def __init__(self, tokens_per_block: int = 32, custom_tokenizer: Optional[str] = None, backfill_block_hashes_on_finish: bool = False, - load_weight: float = 1.0, + load_weight: float = 0.25, load_cap: float = float("inf"), **kwargs): super().__init__(server_role, servers, metadata_server_cfg, @@ -1105,13 +1122,6 @@ async def get_next_server( token_lists, block_hashes_by_algo = await asyncio.to_thread( self._tokenize_and_compute_block_hashes_by_algo, request, hash_algo_by_server.values(), cache_salt_id) - padded_tokens_by_algo = { - hash_algo: - max( - sum(len(hash_list) - for hash_list in block_hashes) * self._tokens_per_block, 1) - for hash_algo, block_hashes in block_hashes_by_algo.items() - } # select the server by (KV match - load), bounded by load_cap workloads = [ self._server_state[server].num_active_requests() @@ -1126,12 +1136,11 @@ async def get_next_server( server = servers[i] hash_algo = hash_algo_by_server[server] block_hashes = block_hashes_by_algo[hash_algo] - padded_tokens = padded_tokens_by_algo[hash_algo] # https://github.com/ai-dynamo/dynamo/blob/main/docs/kv_cache_routing.md#kv-cache-routing-and-load-balancing matches.append(await self._server_state[server].matched_tokens( block_hashes, hash_algo)) - score = matches[-1] / padded_tokens - self._load_weight * \ - load_fractions[i] + score = matches[-1] / self._tokens_per_block - self._load_weight * \ + workloads[i] scores.append(score) # Optional hard cap: drop servers at/over load_cap; fall back to all if # none remain. Disabled by default (load_cap=inf) to match the original diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index 3d0d82fd6fa0..f5ff24e5dd0d 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -1863,7 +1863,7 @@ def test_tokenize_without_tools_passes_none(router_class): assert kwargs["tools"] is None assert "thinking" not in kwargs # messages and add_generation_prompt must still flow through unchanged. - assert kwargs["tokenize"] is True + assert kwargs["tokenize"] is False def test_tokenize_preserves_empty_tools_list(): @@ -1925,3 +1925,140 @@ def test_tokenize_skipped_when_prompt_token_ids_already_set(): assert out == [[10, 20, 30]] get_tok.assert_not_called() tok.apply_chat_template.assert_not_called() + + +class _PrefixCacheFakeTokenizer: + + _SPECIAL = {"": 1, "": 2, "": 3, "": 4, "": 5} + + def apply_chat_template(self, + messages, + add_generation_prompt=False, + tokenize=False, + return_dict=False, + tools=None, + **kwargs): + tag = {"system": "", "user": "", "assistant": ""} + parts = [""] + for m in messages: + role = m["role"] if isinstance(m, dict) else m.role + content = m["content"] if isinstance(m, dict) else m.content + parts.append(tag.get(role, "") + str(content) + "") + if add_generation_prompt: + parts.append("") + return "".join(parts) + + def encode(self, text, add_special_tokens=False): + ids = [] + i = 0 + n = len(text) + while i < n: + special = None + for token, tid in self._SPECIAL.items(): + if text.startswith(token, i): + special = (token, tid) + break + if special is not None: + ids.append(special[1]) + i += len(special[0]) + continue + j = i + while j < n and not any( + text.startswith(token, j) for token in self._SPECIAL): + j += 1 + run = text[i:j] + ids.append(9000 + len(run)) + ids.extend(ord(c) for c in run) + i = j + return ids + + +def _grow_conversation(): + convo = [{ + "role": "system", + "content": "SYS " * 40 + }, { + "role": "user", + "content": "hello " * 20 + }] + yield [dict(m) for m in convo] + for turn in range(4): + convo.append({"role": "assistant", "content": f"resp{turn} " * 30}) + convo.append({"role": "user", "content": f"q{turn} " * 12}) + yield [dict(m) for m in convo] + + +def test_prefix_cache_tokenize_matches_full_encode(servers): + router = KvCacheAwareRouter(server_role=None, + servers=servers, + tokens_per_block=4) + tok = _PrefixCacheFakeTokenizer() + with mock.patch.object(router, "_get_tokenizer", return_value=tok): + for convo in _grow_conversation(): + req = ChatCompletionRequest(model="mock", messages=convo) + rendered = tok.apply_chat_template( + convo, + add_generation_prompt=req.add_generation_prompt, + tokenize=False) + reference = tok.encode(rendered, add_special_tokens=False) + result = router._tokenize(req)[0] + assert result == reference + assert req.prompt_token_ids == reference + + +def test_prefix_cache_tokenize_encodes_only_delta(servers): + router = KvCacheAwareRouter(server_role=None, + servers=servers, + tokens_per_block=4) + tok = _PrefixCacheFakeTokenizer() + encoded_lengths = [] + real_encode = tok.encode + + def _recording_encode(text, add_special_tokens=False): + encoded_lengths.append(len(text)) + return real_encode(text, add_special_tokens=add_special_tokens) + + tok.encode = _recording_encode + with mock.patch.object(router, "_get_tokenizer", return_value=tok): + for index, convo in enumerate(_grow_conversation()): + encoded_lengths.clear() + req = ChatCompletionRequest(model="mock", messages=convo) + rendered = tok.apply_chat_template( + convo, + add_generation_prompt=req.add_generation_prompt, + tokenize=False) + router._tokenize(req) + assert encoded_lengths + if index >= 1: + assert min(encoded_lengths) < len(rendered) + + +def test_prefix_cache_tokenize_falls_back_on_divergent_prefix(servers): + router = KvCacheAwareRouter(server_role=None, + servers=servers, + tokens_per_block=4) + tok = _PrefixCacheFakeTokenizer() + convo_a = [{ + "role": "system", + "content": "A " * 30 + }, { + "role": "user", + "content": "alpha" + }] + convo_b = [{ + "role": "system", + "content": "B " * 30 + }, { + "role": "user", + "content": "beta" + }] + with mock.patch.object(router, "_get_tokenizer", return_value=tok): + for convo in (convo_a, convo_b, convo_a): + req = ChatCompletionRequest(model="mock", + messages=[dict(m) for m in convo]) + rendered = tok.apply_chat_template( + [dict(m) for m in convo], + add_generation_prompt=req.add_generation_prompt, + tokenize=False) + reference = tok.encode(rendered, add_special_tokens=False) + assert router._tokenize(req)[0] == reference From 5d1176f9d49333e58a2236009f3ad5172182f47a Mon Sep 17 00:00:00 2001 From: Liao Lanyu <108499334+lancelly@users.noreply.github.com> Date: Mon, 8 Jun 2026 14:42:37 +0800 Subject: [PATCH 11/28] [None][perf] disagg gen: prepend new generation-only requests to active list (#14867) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com> --- .../pyexecutor/scheduler/scheduler_v2.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 1575021bab5f..1e063e571d39 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -14,6 +14,7 @@ # limitations under the License. import enum +import os from typing import Optional from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy @@ -204,6 +205,14 @@ def __init__( self._disagg_gen_init_state_value = LlmRequestState.DISAGG_GENERATION_INIT.value self._gen_to_complete_state_value = LlmRequestState.GENERATION_TO_COMPLETE.value + # Opt-in (default off): on the disagg generation server, schedule + # just-arrived generation-only requests (first token not yet emitted) + # ahead of older in-flight decodes to reduce TTFT under budget + # pressure. See _schedule_loop. + self._prioritize_first_token_gen = ( + os.environ.get("TLLM_DISAGG_GEN_PRIORITIZE_FIRST_TOKEN", "0") == "1" + ) + def schedule_request( self, active_requests: RequestList, inflight_request_ids: set[int] ) -> SchedulerOutput: @@ -258,6 +267,18 @@ def _schedule_loop(self, active_requests, inflight_request_ids): # Use indexed iteration (while + req_it_end) so that MAX_UTIL # eviction can shrink the range from the tail. requests_list = list(active_requests) + + # Opt-in: prioritize disagg-gen first-token requests (see __init__). + # py_decoding_iter == 0 covers DISAGG_GENERATION_INIT / + # TRANS_IN_PROGRESS / TRANS_COMPLETE. A stable sort keeps FIFO arrival + # order within each class; no-op without generation-only requests. + if self._prioritize_first_token_gen: + requests_list.sort( + key=lambda req: ( + 0 if (req.is_generation_only_request() and req.py_decoding_iter == 0) else 1 + ) + ) + req_it_end = len(requests_list) req_it = 0 From 165fc5366ea7cbefa139185a5c02c8cc73570447 Mon Sep 17 00:00:00 2001 From: Liao Lanyu <108499334+lancelly@users.noreply.github.com> Date: Mon, 8 Jun 2026 18:52:35 +0800 Subject: [PATCH 12/28] [None][feat] Add ConversationAwareADPRouter: explicit conversation->rank affinity for attention DP (#14983) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com> --- .../_torch/pyexecutor/scheduler/adp_router.py | 275 +++++++++++++++--- tensorrt_llm/disaggregated_params.py | 4 + tensorrt_llm/llmapi/llm_args.py | 22 ++ tensorrt_llm/serve/openai_protocol.py | 2 + .../_torch/executor/test_adp_router.py | 174 +++++++++++ .../test_disaggregated_params.py | 23 ++ 6 files changed, 465 insertions(+), 35 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index e3f9601fc12a..84f33c572bf7 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -21,9 +21,9 @@ import math import random from abc import ABC, abstractmethod -from collections import namedtuple +from collections import OrderedDict, namedtuple from dataclasses import MISSING, astuple, dataclass, field, fields, replace -from typing import TYPE_CHECKING, Dict, List, Set, Tuple +from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple from tensorrt_llm.logger import logger @@ -178,6 +178,20 @@ def create( kv_cache_manager has block reuse enabled; DefaultADPRouter otherwise. """ + if ( + attention_dp_config is not None + and attention_dp_config.kv_cache_routing_conversation_affinity + ): + # Explicit conversation_id -> rank affinity. Independent of the KV + # cache manager (works with or without block reuse, though it is + # most beneficial with reuse on), so it is checked before the + # KV-cache-aware path and takes precedence when both are enabled. + return ConversationAwareADPRouter( + dist=dist, + max_sessions=attention_dp_config.kv_cache_routing_max_sessions, + fair_share_multiplier=attention_dp_config.kv_cache_routing_fair_share_multiplier, + ) + if ( attention_dp_config is not None and attention_dp_config.enable_kv_cache_aware_routing @@ -235,6 +249,58 @@ def gather_all_rank_states( responses = self.dist.tp_allgather(local_state.serialize()) return [RankState.deserialize(data=resp) for resp in responses] + @staticmethod + def _assign_explicit_dp_ranks( + requests: List["RequestQueueItem"], + all_ranks_new_requests: Dict[int, List["RequestQueueItem"]], + all_ranks_num_active_requests: List[int], + max_num_active_requests: int, + ) -> List["RequestQueueItem"]: + """Place requests carrying an explicit ``attention_dp_rank`` on that rank + (while it is under ``max_num_active_requests``) and return the rest for + load-balanced assignment. Mutates the ``all_ranks_*`` accumulators in + place. Shared by Default and ConversationAware routers. + """ + remaining: List["RequestQueueItem"] = [] + for req_item in requests: + scheduling_params = getattr(req_item.request, "py_scheduling_params", None) + target_dp_rank = ( + scheduling_params.attention_dp_rank if scheduling_params is not None else None + ) + if ( + target_dp_rank is not None + and all_ranks_num_active_requests[target_dp_rank] < max_num_active_requests + ): + all_ranks_num_active_requests[target_dp_rank] += 1 + all_ranks_new_requests[target_dp_rank].append(req_item) + else: + remaining.append(req_item) + return remaining + + @staticmethod + def _expected_num_active_requests( + all_ranks_num_active_requests: List[int], + num_new_requests: int, + tp_size: int, + *, + multiplier: float = 1.0, + hard_cap: Optional[int] = None, + ) -> int: + """Loose per-rank target = ``multiplier * ceil(fair-share)``, floored at the + busiest rank's current load and optionally capped at ``hard_cap`` + (``max_num_active_requests``). It is the soft cap that bounds per-rank + assignment so no rank exceeds the returned value -- the ADP padding + invariant asserted in py_executor._pad_attention_dp_dummy_request. + """ + fair_share = ( + sum(all_ranks_num_active_requests) + num_new_requests + tp_size - 1 + ) // tp_size + expected = max( + math.ceil(multiplier * fair_share), + max(all_ranks_num_active_requests), + ) + return min(expected, hard_cap) if hard_cap is not None else expected + @abstractmethod def route_requests( self, @@ -311,33 +377,21 @@ def get_relax_value(req_item): sorted_requests = sorted(new_requests, key=get_relax_value) - remaining_unscheduled = [] - for req_item in sorted_requests: - scheduled = False - scheduling_params = getattr(req_item.request, "py_scheduling_params", None) - if scheduling_params is not None: - target_dp_rank = scheduling_params.attention_dp_rank - if ( - target_dp_rank is not None - and all_ranks_num_active_requests[target_dp_rank] < max_num_active_requests - ): - all_ranks_num_active_requests[target_dp_rank] += 1 - scheduled = True - all_ranks_new_requests[target_dp_rank].append(req_item) - - if not scheduled: - remaining_unscheduled.append(req_item) + remaining_unscheduled = self._assign_explicit_dp_ranks( + sorted_requests, + all_ranks_new_requests, + all_ranks_num_active_requests, + max_num_active_requests, + ) num_new_requests_all_ranks = len(remaining_unscheduled) - total_num_active_requests = sum(all_ranks_num_active_requests) # Cap at max_num_active_requests so the per-rank target never # exceeds what the rank can physically schedule. - expected_num_active_requests = min( - max( - (total_num_active_requests + num_new_requests_all_ranks + tp_size - 1) // tp_size, - max(all_ranks_num_active_requests), - ), - max_num_active_requests, + expected_num_active_requests = self._expected_num_active_requests( + all_ranks_num_active_requests, + num_new_requests_all_ranks, + tp_size, + hard_cap=max_num_active_requests, ) all_ranks_new_requests = self._balance_requests_across_ranks( @@ -660,16 +714,12 @@ def _sort_key(req_item): # max_num_active_requests so the per-rank target never exceeds what # the rank can physically schedule (matches DefaultADPRouter). num_new_requests_all_ranks = len(remaining_unscheduled) - total_num_active_requests = sum(all_ranks_num_active_requests) - fair_share = ( - total_num_active_requests + num_new_requests_all_ranks + tp_size - 1 - ) // tp_size - expected_num_active_requests = min( - max( - math.ceil(self.fair_share_multiplier * fair_share), - max(all_ranks_num_active_requests), - ), - max_num_active_requests, + expected_num_active_requests = self._expected_num_active_requests( + all_ranks_num_active_requests, + num_new_requests_all_ranks, + tp_size, + multiplier=self.fair_share_multiplier, + hard_cap=max_num_active_requests, ) eligible_ranks = [ rank @@ -743,3 +793,158 @@ def _sort_key(req_item): ) return all_ranks_new_requests, expected_num_active_requests + + +class ConversationAwareADPRouter(ADPRouter): + """Pins each conversation to a single attention-DP rank: the first request + of a conversation is round-robined, and every later request with the same + ``conversation_id`` returns to that rank, keeping the conversation's + KV-cache prefix on one rank. Falls back to load-balanced round-robin when no + ``conversation_id`` is present. + """ + + # Default LRU cap on the conversation->rank map (entries are ~tens of + # bytes each). Bounds memory on long-running servers as conversations churn. + DEFAULT_MAX_SESSIONS = 1 << 16 + + def __init__( + self, + dist: "Distributed", + max_sessions: int = DEFAULT_MAX_SESSIONS, + fair_share_multiplier: float = 2.0, + ): + super().__init__(dist) + self._conv_to_rank: "OrderedDict[str, int]" = OrderedDict() + self._max_sessions = max(1, int(max_sessions)) + self._fair_share_multiplier = max(1.0, float(fair_share_multiplier)) + self._round_robin_cursor = 0 + + def create_rank_state( + self, + active_requests: list[LlmRequest], + new_requests: list[RequestQueueItem], + ) -> RankState: + if self.dist.has_cp_helix: + num_active_tokens = sum(req.total_input_len_cp for req in active_requests) + else: + num_active_tokens = sum(req.py_orig_prompt_len for req in active_requests) + return RankState( + rank=self.dist.tp_rank, + num_active_requests=len(active_requests), + num_active_tokens=num_active_tokens, + ) + + @staticmethod + def _conversation_id(req_item) -> "str | None": + req = getattr(req_item, "request", None) + if req is None: + return None + disagg = getattr(req, "py_disaggregated_params", None) + if disagg is None: + return None + conv_id = getattr(disagg, "conversation_id", None) + return conv_id if conv_id else None + + def _record_target_rank(self, conv_id: str, rank: int) -> None: + """Bind/refresh the rank a conversation is pinned to (LRU-touch + evict).""" + self._conv_to_rank[conv_id] = rank + self._conv_to_rank.move_to_end(conv_id) + while len(self._conv_to_rank) > self._max_sessions: + self._conv_to_rank.popitem(last=False) + + def route_requests( + self, + all_rank_states: list[RankState], + new_requests: list[RequestQueueItem], + max_num_active_requests: int, + ) -> Tuple[Dict[int, List[RequestQueueItem]], int]: + tp_size = len(all_rank_states) + all_ranks_new_requests: Dict[int, List[RequestQueueItem]] = { + s.rank: [] for s in all_rank_states + } + all_ranks_num_active_requests = [s.num_active_requests for s in all_rank_states] + + def get_relax_value(req_item): + scheduling_params = getattr(req_item.request, "py_scheduling_params", None) + if scheduling_params is None: + return True + return scheduling_params.attention_dp_relax + + sorted_requests = sorted(new_requests, key=get_relax_value) + + # 1) Honour an explicit attention_dp_rank first (strict placement). + remaining_unscheduled = self._assign_explicit_dp_ranks( + sorted_requests, + all_ranks_new_requests, + all_ranks_num_active_requests, + max_num_active_requests, + ) + + # 2) Loose soft cap for spreading new conversations across ranks; sticky + # returns may exceed it (hard cap), so it is re-bumped after the loop. + expected_num_active_requests = self._expected_num_active_requests( + all_ranks_num_active_requests, + len(remaining_unscheduled), + tp_size, + multiplier=self._fair_share_multiplier, + ) + + def _least_loaded(soft_cap: int) -> int: + """Lowest-active rank under soft_cap (deterministic), else global min.""" + cands = [r for r in range(tp_size) if all_ranks_num_active_requests[r] < soft_cap] + if not cands: + cands = list(range(tp_size)) + return min(cands, key=lambda r: (all_ranks_num_active_requests[r], r)) + + def _next_rr(soft_cap: int) -> int: + """Round-robin to the next rank under soft_cap; advance the cursor.""" + for _ in range(tp_size): + r = self._round_robin_cursor % tp_size + self._round_robin_cursor = (self._round_robin_cursor + 1) % tp_size + if all_ranks_num_active_requests[r] < soft_cap: + return r + return _least_loaded(soft_cap) + + for req_item in remaining_unscheduled: + conv_id = self._conversation_id(req_item) + rank = None + + if conv_id is not None and conv_id in self._conv_to_rank: + target_rank = self._conv_to_rank[conv_id] + # Sticky: return the conversation to its target rank while that + # rank is under the HARD cap (max_num_active_requests), NOT the + # soft fair-share `expected`. Stickiness wins over balance so a + # conversation's KV prefix stays resident on one rank instead of + # migrating when the rank gets moderately busy. + if all_ranks_num_active_requests[target_rank] < max_num_active_requests: + rank = target_rank + self._record_target_rank(conv_id, target_rank) # LRU touch + # else: target saturated this batch -> fall through to overflow + # WITHOUT rebinding, so the conversation returns next batch. + + if rank is None: + # First turn of a new conversation, sticky-overflow, or no + # conversation_id -> round-robin spread under the soft cap. + rank = _next_rr(expected_num_active_requests) + if conv_id is not None and conv_id not in self._conv_to_rank: + # Bind this new conversation to its first-turn rank. + self._record_target_rank(conv_id, rank) + + all_ranks_new_requests[rank].append(req_item) + all_ranks_num_active_requests[rank] += 1 + + # Sticky returns use the hard cap, so a rank may now exceed the pre-loop + # soft `expected`. Re-bump so the returned value covers the actual + # per-rank max -- _pad_attention_dp_dummy_request asserts + # expected >= len(active_requests) on every rank. + expected_num_active_requests = max( + expected_num_active_requests, max(all_ranks_num_active_requests) + ) + + logger.debug( + f"[adp_router][conv] new_reqs_per_rank=" + f"{[len(all_ranks_new_requests[r]) for r in range(tp_size)]} " + f"tracked_convs={len(self._conv_to_rank)}" + ) + + return all_ranks_new_requests, expected_num_active_requests diff --git a/tensorrt_llm/disaggregated_params.py b/tensorrt_llm/disaggregated_params.py index 6d6459a8afda..93b9a38f2215 100644 --- a/tensorrt_llm/disaggregated_params.py +++ b/tensorrt_llm/disaggregated_params.py @@ -55,6 +55,10 @@ class DisaggregatedParams: ctx_info_endpoint: Optional[str] = None schedule_style: Optional[DisaggScheduleStyle] = None ctx_usage: Optional[Dict[str, Any]] = None + # Multi-turn conversation id (from session headers such as X-Session-ID), + # carried through so worker-side consumers (e.g. the ADP router) can see + # the same id the disagg orchestrator routed on. + conversation_id: Optional[str] = None # E-P Disaggregated Params multimodal_embedding_handles: Optional[List[Dict[str, Any]]] = ( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index baf73bcb6c6c..330c93eb8bc9 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1288,6 +1288,28 @@ class AttentionDpConfig(StrictBaseModel): "and causes empty fetch cycles. Default False preserves the prior " "behaviour (fetch blocks when truly idle). " "Only used when enable_kv_cache_aware_routing is True.") + kv_cache_routing_conversation_affinity: bool = Field( + default=False, + description= + "Enable explicit conversation-affinity routing for attention DP. When " + "True, the first request of each conversation is round-robined across " + "ranks and every subsequent request carrying the same conversation_id " + "(read from disaggregated_params, populated from the X-Session-ID " + "header) is pinned to that conversation's first-turn rank. This keeps a " + "multi-turn conversation's KV-cache prefix on one rank (maximizing " + "block reuse, minimizing cross-rank migration). Unlike " + "enable_kv_cache_aware_routing (affinity inferred from prefix-match " + "length, which is lost when blocks are evicted), the conversation->rank " + "map is explicit and survives eviction. Falls back to load-balanced " + "round-robin when no conversation_id is available. Takes precedence " + "over enable_kv_cache_aware_routing when both are set.") + kv_cache_routing_max_sessions: int = Field( + default=65536, + description= + "LRU cap on the conversation->rank map used by conversation-affinity " + "routing. The oldest conversations are evicted once more than this many " + "are tracked, bounding memory on long-running servers. Only used when " + "kv_cache_routing_conversation_affinity is True.") @model_validator(mode='after') def validate_attention_dp_config(self) -> 'AttentionDpConfig': diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 9f1ddd3dead7..e80c494e8ae5 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -1329,6 +1329,7 @@ def to_disaggregated_params( ctx_info_endpoint=tllm_disagg_params.ctx_info_endpoint, schedule_style=tllm_disagg_params.schedule_style, ctx_usage=ctx_usage, + conversation_id=tllm_disagg_params.conversation_id, ) @@ -1353,6 +1354,7 @@ def to_llm_disaggregated_params( ctx_info_endpoint=disaggregated_params.ctx_info_endpoint, schedule_style=disaggregated_params.schedule_style, ctx_usage=None if ctx_usage is None else ctx_usage.model_dump(), + conversation_id=disaggregated_params.conversation_id, ) diff --git a/tests/unittest/_torch/executor/test_adp_router.py b/tests/unittest/_torch/executor/test_adp_router.py index 4ead152409b7..b258e622b848 100644 --- a/tests/unittest/_torch/executor/test_adp_router.py +++ b/tests/unittest/_torch/executor/test_adp_router.py @@ -16,6 +16,7 @@ from tensorrt_llm._torch.pyexecutor.scheduler import FCFSWaitingQueue from tensorrt_llm._torch.pyexecutor.scheduler.adp_router import ( ADPRouter, + ConversationAwareADPRouter, DefaultADPRouter, KVCacheAwareADPRouter, RankIterStatsPayload, @@ -1085,3 +1086,176 @@ def test_fair_share_multiplier_caps_per_rank(self): ) assert len(result_strict[0]) == 2 assert sum(len(result_strict[r]) for r in range(1, tp_size)) == 6 + + +def _make_conv_request_item( + req_id, + conversation_id, + target_dp_rank=None, + attention_dp_relax=True, + num_tokens=10, +): + """Mock RequestQueueItem carrying a conversation_id on disaggregated_params.""" + item = MagicMock() + item.id = req_id + item.child_req_ids = None + scheduling_params = MagicMock() + scheduling_params.attention_dp_rank = target_dp_rank + scheduling_params.attention_dp_relax = attention_dp_relax + item.request = MagicMock() + item.request.py_scheduling_params = scheduling_params + item.request.input_token_ids = list(range(num_tokens)) + item.request.py_orig_prompt_len = num_tokens + if conversation_id is None: + item.request.py_disaggregated_params = None + else: + disagg = MagicMock() + disagg.conversation_id = conversation_id + item.request.py_disaggregated_params = disagg + return item + + +class TestConversationAwareADPRouter: + """Routing behavior of ConversationAwareADPRouter (explicit conv->rank).""" + + @staticmethod + def _router(tp_size=4, max_sessions=1 << 16): + return ConversationAwareADPRouter( + dist=_mock_dist(tp_size=tp_size), max_sessions=max_sessions + ) + + @staticmethod + def _states(tp_size, active=None): + active = active or [0] * tp_size + return [ + RankState(rank=r, num_active_requests=active[r], num_active_tokens=active[r] * 10) + for r in range(tp_size) + ] + + @staticmethod + def _route(router, states, items, cap=1000): + assign, _ = router.route_requests(states, items, max_num_active_requests=cap) + pos = {} + for rank, lst in assign.items(): + for it in lst: + pos[it.id] = rank + return pos + + def test_interface_compliance(self): + assert isinstance(self._router(), ADPRouter) + + def test_first_turn_round_robin(self): + """First request of each conversation spreads one-per-rank (RR).""" + router = self._router(tp_size=4) + items = [_make_conv_request_item(i, f"conv{i}") for i in range(4)] + pos = self._route(router, self._states(4), items) + assert sorted(pos.values()) == [0, 1, 2, 3] + + def test_subsequent_turns_are_sticky(self): + """Later turns of a conversation return to its first-turn rank.""" + router = self._router(tp_size=4) + convs = ["A", "B", "C", "D"] + first = [_make_conv_request_item(i, convs[i]) for i in range(4)] + pos1 = self._route(router, self._states(4), first) + home = {convs[i]: pos1[i] for i in range(4)} + # A new batch of second turns (fresh ids) for the same conversations. + second = [_make_conv_request_item(100 + i, convs[i]) for i in range(4)] + pos2 = self._route(router, self._states(4), second) + for i in range(4): + assert pos2[100 + i] == home[convs[i]] + + def test_no_conversation_id_falls_back_and_is_not_recorded(self): + router = self._router(tp_size=4) + items = [_make_conv_request_item(i, None) for i in range(8)] + pos = self._route(router, self._states(4), items) + assert len(pos) == 8 + # Nothing recorded for conversation-less requests. + assert dict(router._conv_to_rank) == {} + # Round-robin spread keeps ranks within one of each other. + counts = [sum(1 for v in pos.values() if v == r) for r in range(4)] + assert max(counts) - min(counts) <= 1 + + def test_routing_is_deterministic_across_ranks(self): + """Two independent instances (= two TP ranks) must agree exactly, or + the no-broadcast allgather protocol would deadlock.""" + convs = ["A", "B", "C", None, "A", "B", "E", None, "C"] + + def run(): + router = self._router(tp_size=4) + items = [_make_conv_request_item(i, convs[i]) for i in range(len(convs))] + return self._route(router, self._states(4), items) + + assert run() == run() + + def test_lru_eviction_bounds_map(self): + router = self._router(tp_size=2, max_sessions=2) + items = [_make_conv_request_item(i, f"c{i}") for i in range(5)] + self._route(router, self._states(2), items) + assert len(router._conv_to_rank) == 2 + assert set(router._conv_to_rank) == {"c3", "c4"} + + def test_explicit_target_dp_rank_respected(self): + router = self._router(tp_size=4) + item = _make_conv_request_item(1, "A", target_dp_rank=2, attention_dp_relax=False) + pos = self._route(router, self._states(4), [item]) + assert pos[1] == 2 + + def test_sticky_overflow_keeps_mapping(self): + """When the home rank is saturated at the HARD cap (max_num_active_requests), + the turn overflows off home but the conversation stays mapped home (so it + returns home next batch once that rank has capacity). Stickiness uses the + hard cap -- not the soft fair-share -- so a conversation's KV prefix stays + resident on one rank in the common case.""" + router = ConversationAwareADPRouter(dist=_mock_dist(tp_size=2)) + # Seed conv A -> rank 0. + self._route(router, self._states(2), [_make_conv_request_item(1, "A")]) + home = router._conv_to_rank["A"] + # Fill the home rank to the hard cap, route A again -> must overflow off home... + states = self._states(2) + states[home] = RankState(rank=home, num_active_requests=5, num_active_tokens=50) + pos = self._route(router, states, [_make_conv_request_item(2, "A")], cap=5) + assert pos[2] != home + # ...but the mapping is unchanged so it can return home later. + assert router._conv_to_rank["A"] == home + + def test_create_rank_state(self): + router = ConversationAwareADPRouter(dist=_mock_dist(tp_rank=2)) + req1 = Mock(py_orig_prompt_len=100) + req2 = Mock(py_orig_prompt_len=50) + state = router.create_rank_state(active_requests=[req1, req2], new_requests=[]) + assert state.rank == 2 + assert state.num_active_requests == 2 + assert state.num_active_tokens == 150 + + def test_factory_selects_conversation_router(self): + cfg = MagicMock() + cfg.kv_cache_routing_conversation_affinity = True + cfg.kv_cache_routing_max_sessions = 8 + router = ADPRouter.create(dist=_mock_dist(), kv_cache_manager=None, attention_dp_config=cfg) + assert isinstance(router, ConversationAwareADPRouter) + assert router._max_sessions == 8 + + def test_factory_default_when_disabled(self): + cfg = MagicMock() + cfg.kv_cache_routing_conversation_affinity = False + cfg.enable_kv_cache_aware_routing = False + router = ADPRouter.create(dist=_mock_dist(), kv_cache_manager=None, attention_dp_config=cfg) + assert isinstance(router, DefaultADPRouter) + + def test_returned_expected_covers_every_rank(self): + """Regression: returned expected_num_active_requests must be >= every + rank's post-assignment active count, else + py_executor._pad_attention_dp_dummy_request asserts and the executor + hangs. A sticky storm (one conversation hammered) must not push its home + rank past expected.""" + router = self._router(tp_size=8) + states = self._states(8) + # 40 requests for the SAME conversation: hard-cap stickiness concentrates + # all of them on the home rank, pushing it well past the pre-loop soft + # `expected`. The post-loop re-bump must lift `expected` to cover the home + # rank's actual count, else _pad_attention_dp_dummy_request asserts. + items = [_make_conv_request_item(i, "A") for i in range(40)] + assign, expected = router.route_requests(states, items, max_num_active_requests=256) + final = [states[r].num_active_requests + len(assign[r]) for r in range(8)] + assert all(expected >= f for f in final), (expected, final) + assert sum(len(v) for v in assign.values()) == 40 # nothing dropped diff --git a/tests/unittest/disaggregated/test_disaggregated_params.py b/tests/unittest/disaggregated/test_disaggregated_params.py index c20f81fcc30f..655fbcabbd13 100644 --- a/tests/unittest/disaggregated/test_disaggregated_params.py +++ b/tests/unittest/disaggregated/test_disaggregated_params.py @@ -64,6 +64,7 @@ def test_to_disaggregated_params(): "cached_tokens": 4, }, }, + conversation_id="conv-abc", ) openai_params = to_disaggregated_params(llm_params) @@ -74,6 +75,7 @@ def test_to_disaggregated_params(): assert openai_params.ctx_info_endpoint == "tcp://10.0.0.1:5000" assert openai_params.ctx_usage.prompt_tokens == 10 assert openai_params.ctx_usage.prompt_tokens_details.cached_tokens == 4 + assert openai_params.conversation_id == "conv-abc" def test_to_llm_disaggregated_params(): @@ -94,6 +96,7 @@ def test_to_llm_disaggregated_params(): total_tokens=10, prompt_tokens_details=PromptTokensDetails(cached_tokens=4), ), + conversation_id="conv-xyz", ) llm_params = to_llm_disaggregated_params(openai_params) @@ -103,6 +106,26 @@ def test_to_llm_disaggregated_params(): assert llm_params.ctx_info_endpoint == "tcp://10.0.0.1:5000" assert llm_params.ctx_usage["prompt_tokens"] == 10 assert llm_params.ctx_usage["prompt_tokens_details"]["cached_tokens"] == 4 + assert llm_params.conversation_id == "conv-xyz" + + +def test_disaggregated_params_conversation_id(): + """conversation_id defaults to None and survives the serve<->llm round-trip.""" + from tensorrt_llm.serve.openai_protocol import DisaggregatedParams as OpenAIDisaggregatedParams + from tensorrt_llm.serve.openai_protocol import ( + to_disaggregated_params, + to_llm_disaggregated_params, + ) + + assert DisaggregatedParams().conversation_id is None + + # serve -> llm -> serve preserves the conversation id end to end. + openai_params = OpenAIDisaggregatedParams( + request_type="context_only", conversation_id="conv-roundtrip" + ) + llm_params = to_llm_disaggregated_params(openai_params) + assert llm_params.conversation_id == "conv-roundtrip" + assert to_disaggregated_params(llm_params).conversation_id == "conv-roundtrip" @patch("tensorrt_llm.disaggregated_params.tllme") From 7794a360f38336c26fed7212dd241cd237dbf801 Mon Sep 17 00:00:00 2001 From: Shi Xiaowei <39303645+Shixiaowei02@users.noreply.github.com> Date: Mon, 8 Jun 2026 23:25:13 +0800 Subject: [PATCH 13/28] [None][feat] Dis-agg content-derived conversation affinity (#15074) Signed-off-by: Fanrong Li Signed-off-by: Shi Xiaowei <39303645+Shixiaowei02@users.noreply.github.com> --- tensorrt_llm/serve/openai_client.py | 2 +- tensorrt_llm/serve/router.py | 112 ++++++++++++---- ...agg_config_conditional_deepseek_v3_v2.yaml | 4 +- .../test_disagg_openai_client.py | 2 +- tests/unittest/disaggregated/test_router.py | 123 +++++++++++++++--- 5 files changed, 195 insertions(+), 48 deletions(-) diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index a3d23b266d2f..9c84d5065c9a 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -89,7 +89,7 @@ async def _finish_request(self, request: UCompletionRequest, success: bool = Tru """Finish the request in the router. ``success`` lets the router distinguish completed vs failed requests - (e.g. for cache backfill). + (e.g. for routed-block tracking). """ ... diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 05166e706cd8..2545132094e3 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -51,6 +51,12 @@ OpenAIRequest = Union[CompletionRequest, ChatCompletionRequest] BlockHash = Union[int, str] +# Max number of conversations whose home-server pin is retained (LRU). +ROUTE_AFFINITY_CACHE_SIZE = 50000 +# Leading token-id count folded into the affinity key so pre-tokenized +# requests (placeholder message content) still key per conversation. +ROUTE_AFFINITY_TOKEN_PREFIX = 256 + def get_request_num_tokens(request: OpenAIRequest) -> int: if request.disaggregated_params is None or request.disaggregated_params.request_type == "context_only": @@ -797,16 +803,19 @@ class BlockHashMixin: """ def _init_block_hashing(self, - tokens_per_block: int = 32, + tokens_per_block: Optional[int] = None, custom_tokenizer: Optional[str] = None): env_tokens_per_block = os.environ.get( "TRTLLM_KVCACHE_AWARE_ROUTER_HASH_TOKENS_PER_BLOCK") if env_tokens_per_block is not None: tokens_per_block = int(env_tokens_per_block) - self._tokens_per_block = tokens_per_block + self._tpb_auto = tokens_per_block is None + self._tokens_per_block = 32 if tokens_per_block is None \ + else tokens_per_block self._tokenizers: dict = {} self._custom_tokenizer = custom_tokenizer logger.info(f"BlockHashMixin: tokens_per_block={self._tokens_per_block}" + f"{' (auto, adopts worker)' if self._tpb_auto else ''}" f", custom_tokenizer={self._custom_tokenizer}") def _get_tokenizer(self, model: str): @@ -995,9 +1004,9 @@ def __init__(self, metadata_server: JsonDictionary = None, use_tokens: bool = False, max_batch_size: int = 64, - tokens_per_block: int = 32, + tokens_per_block: Optional[int] = None, custom_tokenizer: Optional[str] = None, - backfill_block_hashes_on_finish: bool = False, + track_routed_blocks: bool = True, load_weight: float = 0.25, load_cap: float = float("inf"), **kwargs): @@ -1009,11 +1018,8 @@ def __init__(self, self._max_batch_size = max_batch_size self._load_weight = load_weight self._load_cap = load_cap - # Opt-in workaround for the disagg-gen path that doesn't emit - # kv_cache_events. Stash hashes at routing, inject on finish. - self._backfill_block_hashes_on_finish = backfill_block_hashes_on_finish - self._inflight_backfill_hashes: dict[int, tuple[list[BlockHash], - str]] = {} + self._track_routed_blocks = track_routed_blocks + self._pending_routed_blocks: dict[int, tuple[list[BlockHash], str]] = {} def _create_server_state(self, server: str) -> KvCacheAwareServerState: return KvCacheAwareServerState(server, self._use_tokens, @@ -1025,19 +1031,20 @@ async def close(self): await state.cancel_poll_task() await super().close() - def _stash_backfill_on_route(self, request: OpenAIRequest, - block_hashes: list[list[BlockHash]], - hash_algo: str) -> None: - if not self._backfill_block_hashes_on_finish: + def _stash_routed_blocks_on_route(self, request: OpenAIRequest, + block_hashes: list[list[BlockHash]], + hash_algo: str) -> None: + if not self._track_routed_blocks: return flat = [h for hl in block_hashes for h in hl] - self._inflight_backfill_hashes[id(request)] = (flat, hash_algo) + self._pending_routed_blocks[id(request)] = (flat, hash_algo) - def _apply_backfill_on_finish(self, request: OpenAIRequest, - server: Optional[str], success: bool) -> None: + def _apply_routed_blocks_on_finish(self, request: OpenAIRequest, + server: Optional[str], + success: bool) -> None: # Pop unconditionally to avoid leaks; apply only when eligible. - entry = self._inflight_backfill_hashes.pop(id(request), None) - if not (self._backfill_block_hashes_on_finish and success): + entry = self._pending_routed_blocks.pop(id(request), None) + if not (self._track_routed_blocks and success): return if entry is None: return @@ -1052,13 +1059,24 @@ def _get_server_hash_algo(self, server: str) -> str: # by update_with_events. return self._server_state[server].hash_algo + def _events_aligned(self, server: str) -> bool: + worker_tpb = self._server_info.get(server, {}).get("tokens_per_block") + return worker_tpb is None or worker_tpb == self._tokens_per_block + async def _prepare_server(self, server: str): await super()._prepare_server(server) if server not in self._prepared_ready_servers: return info = self._server_info.get(server, {}) worker_tpb = info.get("tokens_per_block") - if worker_tpb is not None and worker_tpb != self._tokens_per_block: + if worker_tpb is not None and getattr(self, "_tpb_auto", False): + if worker_tpb != self._tokens_per_block: + logger.info( + "router tokens_per_block unset: adopting worker's %d on %s", + worker_tpb, server) + self._tokens_per_block = worker_tpb + self._tpb_auto = False + elif worker_tpb is not None and worker_tpb != self._tokens_per_block: larger = max(worker_tpb, self._tokens_per_block) smaller = min(worker_tpb, self._tokens_per_block) if larger % smaller != 0: @@ -1071,9 +1089,9 @@ async def _prepare_server(self, server: str): ) logger.warning( "tokens_per_block mismatch on %s: router=%d worker=%d. " - "KV events from worker will not align with router block hashes; " - "use backfill_block_hashes_on_finish=true for best hit rate.", - server, self._tokens_per_block, worker_tpb) + "KV events from worker cannot align with router block hashes; " + "skipping event polling and relying on routed-block tracking " + "for hit rate.", server, self._tokens_per_block, worker_tpb) worker_algo = info.get("kv_cache_hash_algo") known_algos = { KV_CACHE_HASH_ALGO_V1, @@ -1095,6 +1113,21 @@ async def _prepare_server(self, server: str): # Persist once so per-request reads can skip the lock+await. self._server_state[server].set_hash_algo(worker_algo) + @staticmethod + def _content_affinity_key(request: OpenAIRequest) -> Optional[int]: + messages = getattr(request, "messages", None) + if not messages: + return None + parts = [] + for message in messages[:2]: + content = (message.get("content") if isinstance(message, dict) else + getattr(message, "content", "")) + parts.append(str(content)) + token_ids = getattr(request, "prompt_token_ids", None) + if token_ids: + parts.append(str(list(token_ids[:ROUTE_AFFINITY_TOKEN_PREFIX]))) + return hash("".join(parts)) + async def get_next_server( self, request: OpenAIRequest, @@ -1150,16 +1183,38 @@ async def get_next_server( ] if not candidate_idx: candidate_idx = list(range(len(servers))) - max_score = max(scores[i] for i in candidate_idx) - tied = [i for i in candidate_idx if scores[i] == max_score] - winner = tied[self._rr_counter % len(tied)] + # Conversation affinity: pin all turns of a conversation (keyed by a + # content-derived prefix hash, no conversation-id header) to the server + # it first landed on, so a worker eviction shrinking the match score + # cannot scatter the conversation off its warm home. New conversations + # (no pin yet) fall through to the score, which balances them by load. + affinity = getattr(self, "_route_affinity", None) + if affinity is None: + affinity = self._route_affinity = OrderedDict() + conv_key = self._content_affinity_key(request) + winner = None + if conv_key is not None: + pinned = affinity.get(conv_key) + if pinned in servers: + pinned_idx = servers.index(pinned) + if pinned_idx in candidate_idx: + winner = pinned_idx + if winner is None: + max_score = max(scores[i] for i in candidate_idx) + tied = [i for i in candidate_idx if scores[i] == max_score] + winner = tied[self._rr_counter % len(tied)] self._rr_counter += 1 server = servers[winner] + if conv_key is not None: + affinity[conv_key] = server + affinity.move_to_end(conv_key) + while len(affinity) > ROUTE_AFFINITY_CACHE_SIZE: + affinity.popitem(last=False) hash_algo = hash_algo_by_server[server] block_hashes = block_hashes_by_algo[hash_algo] async with self._lock: await self._register_request(server, request) - self._stash_backfill_on_route(request, block_hashes, hash_algo) + self._stash_routed_blocks_on_route(request, block_hashes, hash_algo) return server, { "block_hashes": block_hashes, # list[list[int | str]] "hash_algo": hash_algo, @@ -1176,8 +1231,9 @@ async def finish_request(self, server = self._req_routing_table.pop(id(request), None) if server is not None and server in self._server_state: await self._server_state[server].decrement_load(request) - self._apply_backfill_on_finish(request, server, success) - if server is not None and server in self._server_state: + self._apply_routed_blocks_on_finish(request, server, success) + if (server is not None and server in self._server_state + and self._events_aligned(server)): # Fire-and-forget; poll runs in background and coalesces per server. self._server_state[server].schedule_poll_and_update(session) diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3_v2.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3_v2.yaml index 851807eec4da..0be9c8201d3f 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3_v2.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3_v2.yaml @@ -28,9 +28,9 @@ generation_servers: pipeline_parallel_size: 1 router: type: kv_cache_aware - # disagg-gen path doesn't emit kv_cache_events; backfill router-side state from + # disagg-gen path doesn't emit kv_cache_events; track router-side state from # successfully-routed requests so conditional_disagg can actually trigger. - backfill_block_hashes_on_finish: true + track_routed_blocks: true kv_cache_config: enable_block_reuse: true enable_partial_reuse: true diff --git a/tests/unittest/disaggregated/test_disagg_openai_client.py b/tests/unittest/disaggregated/test_disagg_openai_client.py index efac516a150d..edc09809732b 100644 --- a/tests/unittest/disaggregated/test_disagg_openai_client.py +++ b/tests/unittest/disaggregated/test_disagg_openai_client.py @@ -225,7 +225,7 @@ async def test_request_error_handling( await openai_client.send_request(completion_request) # Should finish request on error with success=False so the router - # doesn't backfill cache state for a request that didn't complete. + # doesn't record routed-block cache state for a request that didn't complete. mock_router.finish_request.assert_called_once_with( completion_request, mock_session, success=False ) diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index f5ff24e5dd0d..859e3c88d747 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -785,8 +785,7 @@ def test_kv_cache_aware_server_state_remove_blocks_silent_on_missing(): @pytest.mark.asyncio -async def test_kv_cache_aware_router_backfill_stashes_flat_list_at_routing( - servers): +async def test_kv_cache_aware_router_tracks_routed_blocks_at_routing(servers): tokens_per_block = 4 token_lists = [[1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008]] router = KvCacheAwareRouter(server_role=None, @@ -794,25 +793,25 @@ async def test_kv_cache_aware_router_backfill_stashes_flat_list_at_routing( use_tokens=False, max_batch_size=32, tokens_per_block=tokens_per_block, - backfill_block_hashes_on_finish=True) + track_routed_blocks=True) request = CompletionRequest(model="TinyLlama", prompt=copy.deepcopy(token_lists)) server, info = await router.get_next_server(request) - assert id(request) in router._inflight_backfill_hashes - stashed, stashed_algo = router._inflight_backfill_hashes[id(request)] + assert id(request) in router._pending_routed_blocks + stashed, stashed_algo = router._pending_routed_blocks[id(request)] expected_flat = [h for hl in info["block_hashes"] for h in hl] assert stashed == expected_flat assert stashed and not isinstance(stashed[0], list) assert stashed_algo == info["hash_algo"] await router.finish_request(request) - assert id(request) not in router._inflight_backfill_hashes + assert id(request) not in router._pending_routed_blocks @pytest.mark.asyncio -async def test_kv_cache_aware_router_backfill_inserts_blocks_on_finish(servers): +async def test_kv_cache_aware_router_inserts_routed_blocks_on_finish(servers): tokens_per_block = 4 token_lists = [[2000, 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008]] router = KvCacheAwareRouter(server_role=None, @@ -820,7 +819,7 @@ async def test_kv_cache_aware_router_backfill_inserts_blocks_on_finish(servers): use_tokens=False, max_batch_size=32, tokens_per_block=tokens_per_block, - backfill_block_hashes_on_finish=True) + track_routed_blocks=True) request = CompletionRequest(model="TinyLlama", prompt=copy.deepcopy(token_lists)) @@ -840,20 +839,22 @@ async def test_kv_cache_aware_router_backfill_inserts_blocks_on_finish(servers): @pytest.mark.asyncio -async def test_kv_cache_aware_router_backfill_disabled_skips_pending(servers): +async def test_kv_cache_aware_router_routed_blocks_disabled_skips_pending( + servers): tokens_per_block = 4 token_lists = [[3000, 3001, 3002, 3003, 3004, 3005, 3006, 3007, 3008]] router = KvCacheAwareRouter(server_role=None, servers=servers, use_tokens=False, max_batch_size=32, - tokens_per_block=tokens_per_block) + tokens_per_block=tokens_per_block, + track_routed_blocks=False) request = CompletionRequest(model="TinyLlama", prompt=copy.deepcopy(token_lists)) server, info = await router.get_next_server(request) - assert router._inflight_backfill_hashes == {} + assert router._pending_routed_blocks == {} await router.finish_request(request) @@ -862,7 +863,7 @@ async def test_kv_cache_aware_router_backfill_disabled_skips_pending(servers): @pytest.mark.asyncio -async def test_kv_cache_aware_router_backfill_drops_on_failure(servers): +async def test_kv_cache_aware_router_drops_routed_blocks_on_failure(servers): tokens_per_block = 4 token_lists = [[4000, 4001, 4002, 4003, 4004, 4005, 4006, 4007, 4008]] router = KvCacheAwareRouter(server_role=None, @@ -870,16 +871,16 @@ async def test_kv_cache_aware_router_backfill_drops_on_failure(servers): use_tokens=False, max_batch_size=32, tokens_per_block=tokens_per_block, - backfill_block_hashes_on_finish=True) + track_routed_blocks=True) request = CompletionRequest(model="TinyLlama", prompt=copy.deepcopy(token_lists)) server, info = await router.get_next_server(request) - assert id(request) in router._inflight_backfill_hashes + assert id(request) in router._pending_routed_blocks await router.finish_request(request, success=False) - assert id(request) not in router._inflight_backfill_hashes + assert id(request) not in router._pending_routed_blocks assert await router._server_state[server].matched_tokens( info["block_hashes"], hash_algo=info["hash_algo"]) == 0 @@ -1019,6 +1020,44 @@ async def fake_fetch(server, timeout): assert "server-a" in router._prepared_ready_servers +@pytest.mark.asyncio +async def test_kv_cache_aware_router_prepare_server_adopts_worker_tpb_when_unset( +): + router = KvCacheAwareRouter(server_role=None, servers=["server-a"]) + assert router._tpb_auto is True + assert router._tokens_per_block == 32 + + async def fake_fetch(server, timeout): + return {"tokens_per_block": 128, "kv_cache_hash_algo": "v1_block_key"} + + with mock.patch.object(router, "_fetch_server_info", + side_effect=fake_fetch): + await router._prepare_server("server-a") + + assert "server-a" in router._prepared_ready_servers + assert router._tokens_per_block == 128 + assert router._tpb_auto is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("worker_tpb", [32, 64, 128]) +async def test_kv_cache_aware_router_auto_tpb_equals_worker(worker_tpb): + router = KvCacheAwareRouter(server_role=None, servers=["server-a"]) + + async def fake_fetch(server, timeout): + return { + "tokens_per_block": worker_tpb, + "kv_cache_hash_algo": "v1_block_key" + } + + with mock.patch.object(router, "_fetch_server_info", + side_effect=fake_fetch): + await router._prepare_server("server-a") + + assert router._tokens_per_block == worker_tpb + assert "server-a" in router._prepared_ready_servers + + @pytest.mark.asyncio async def test_kv_cache_aware_router_prepare_server_warns_on_missing_algo(): router = KvCacheAwareRouter(server_role=None, @@ -1095,7 +1134,8 @@ async def test_kv_cache_aware_router(servers, mock_aiohttp_session): servers=servers, use_tokens=False, max_batch_size=32, - tokens_per_block=32) + tokens_per_block=32, + track_routed_blocks=False) results = [await router.get_next_server(req) for req in requests] assigned_servers, infos = zip(*results) # Initial routing (empty caches): all 3 should get distinct servers @@ -2062,3 +2102,54 @@ def test_prefix_cache_tokenize_falls_back_on_divergent_prefix(servers): tokenize=False) reference = tok.encode(rendered, add_special_tokens=False) assert router._tokenize(req)[0] == reference + + +@pytest.mark.asyncio +async def test_conversation_affinity_pins_across_load(servers): + router = KvCacheAwareRouter(server_role=None, + servers=servers, + tokens_per_block=4, + load_weight=1.0) + tok = mock.MagicMock() + tok.apply_chat_template.return_value = [1, 2, 3, 4, 5, 6, 7, 8] + convo = [{ + "role": "system", + "content": "S" + }, { + "role": "user", + "content": "u1" + }] + with mock.patch.object(router, "_get_tokenizer", return_value=tok): + req1 = ChatCompletionRequest(model="m", + messages=[dict(m) for m in convo]) + home, _ = await router.get_next_server(req1) + # Saturate the home server so a fresh request would avoid it. + for _ in range(64): + await router._server_state[home].increment_load( + ChatCompletionRequest(model="m", + messages=[{ + "role": "user", + "content": "x" + }])) + # Same conversation (same first-2-message content) -> pinned to home + # even though the score now favours the idle servers. + req2 = ChatCompletionRequest(model="m", + messages=[dict(m) for m in convo] + + [{ + "role": "assistant", + "content": "a1" + }, { + "role": "user", + "content": "u2" + }]) + assert (await router.get_next_server(req2))[0] == home + # A different conversation is free to avoid the saturated home. + req3 = ChatCompletionRequest(model="m", + messages=[{ + "role": "system", + "content": "OTHER" + }, { + "role": "user", + "content": "v1" + }]) + assert (await router.get_next_server(req3))[0] != home From eb15febf177ff8848696eed21fd1690ab90ccd05 Mon Sep 17 00:00:00 2001 From: Liao Lanyu <108499334+lancelly@users.noreply.github.com> Date: Tue, 9 Jun 2026 07:34:30 +0800 Subject: [PATCH 14/28] [None][perf] disagg: single-pass request serialization on orchestrator forward path (#15041) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com> --- tensorrt_llm/serve/openai_client.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index 9c84d5065c9a..265a771daf6c 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -181,11 +181,19 @@ async def _post_with_retry( dp = getattr(request, "disaggregated_params", None) if dp is not None and getattr(dp, "disagg_request_id", None) is not None: dp.disagg_request_id = self._disagg_id_generator() - json_data = request.model_dump(exclude_unset=True, mode="json") + # Serialize once on the orchestrator's single event-loop thread: + # model_dump_json (pydantic-core) is ~2.3x faster than + # model_dump(mode="json") + aiohttp json= (json.dumps). Decodes to + # identical JSON (pydantic just emits compact UTF-8 vs spaced ASCII). + body = request.model_dump_json(exclude_unset=True) try: lines_yielded = 0 start_time = get_steady_clock_now_in_seconds() - async with self._session.post(url, json=json_data) as http_response: + async with self._session.post( + url, + data=body, + headers={"Content-Type": "application/json"}, + ) as http_response: content_type = http_response.headers.get("Content-Type", "") if not is_stream and "text/event-stream" in content_type: raise ValueError( From c3e0d24fa99cdf8eeeb8295d41321309f12e9feb Mon Sep 17 00:00:00 2001 From: Liao Lanyu <108499334+lancelly@users.noreply.github.com> Date: Tue, 9 Jun 2026 07:34:51 +0800 Subject: [PATCH 15/28] [None][perf] disagg: skip per-chunk JSON parse on streaming usage rewrite (#15042) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com> --- tensorrt_llm/serve/openai_disagg_service.py | 130 ++++++++++++++++++++ 1 file changed, 130 insertions(+) diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index 1da78967a5a8..9c4bf8d5bbcf 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -13,7 +13,9 @@ # limitations under the License. import asyncio +import json import os +from collections.abc import AsyncIterator from typing import Any, Callable, Dict, Optional from tensorrt_llm.llmapi.disagg_utils import ( @@ -34,6 +36,7 @@ CompletionRequest, DisaggregatedParams, DisaggScheduleStyle, + PromptTokensDetails, UCompletionRequest, UCompletionResponse, ) @@ -177,6 +180,133 @@ async def _send_disagg_request_ctx_first( return done_generator() return ctx_response + def _ctx_usage_for_client( + self, ctx_response: Optional[UCompletionResponse] + ) -> tuple[Optional[int], int]: + if ctx_response is None or ctx_response.usage is None: + return None, 0 + + prompt_tokens = ctx_response.usage.prompt_tokens + cached_tokens = 0 + prompt_tokens_details = ctx_response.usage.prompt_tokens_details + if prompt_tokens_details is not None: + cached_tokens = prompt_tokens_details.cached_tokens + return prompt_tokens, cached_tokens + + def _rewrite_usage_payload_from_ctx( + self, + usage: dict[str, Any], + ctx_response: Optional[UCompletionResponse], + ) -> None: + prompt_tokens, cached_tokens = self._ctx_usage_for_client(ctx_response) + if prompt_tokens is None: + return + + usage["prompt_tokens"] = prompt_tokens + usage["total_tokens"] = prompt_tokens + (usage.get("completion_tokens") or 0) + prompt_tokens_details = usage.get("prompt_tokens_details") + if not isinstance(prompt_tokens_details, dict): + prompt_tokens_details = {} + usage["prompt_tokens_details"] = prompt_tokens_details + prompt_tokens_details["cached_tokens"] = cached_tokens + + def _rewrite_usage_response_from_ctx( + self, + response: UCompletionResponse, + ctx_response: Optional[UCompletionResponse], + ) -> UCompletionResponse: + prompt_tokens, cached_tokens = self._ctx_usage_for_client(ctx_response) + if prompt_tokens is None or response.usage is None: + return response + + response.usage.prompt_tokens = prompt_tokens + response.usage.total_tokens = prompt_tokens + (response.usage.completion_tokens or 0) + response.usage.prompt_tokens_details = PromptTokensDetails(cached_tokens=cached_tokens) + return response + + @staticmethod + def _sse_separator_index(data: bytes) -> tuple[int, bytes] | None: + indexes = [(data.find(sep), sep) for sep in (b"\n\n", b"\r\n\r\n")] + indexes = [(idx, sep) for idx, sep in indexes if idx >= 0] + if not indexes: + return None + return min(indexes, key=lambda item: item[0]) + + def _rewrite_usage_sse_event_from_ctx( + self, + event: bytes, + ctx_response: Optional[UCompletionResponse], + ) -> bytes: + # Only the final chunk has a top-level `usage` to rewrite, but this ran + # json.loads on every relayed chunk. A usage field can't exist unless the + # bytes contain `"usage"`, so skip the parse when they don't. + if b'"usage"' not in event: + return event + separator_match = self._sse_separator_index(event) + separator = separator_match[1] if separator_match else b"" + event_body = event[: -len(separator)] if separator else event + data_lines = [ + line.removeprefix(b"data:").strip() + for line in event_body.splitlines() + if line.startswith(b"data:") + ] + if len(data_lines) != 1 or data_lines[0] == b"[DONE]": + return event + + try: + payload = json.loads(data_lines[0]) + except json.JSONDecodeError: + return event + + usage = payload.get("usage") if isinstance(payload, dict) else None + if not isinstance(usage, dict): + return event + + self._rewrite_usage_payload_from_ctx(usage, ctx_response) + return b"data: " + json.dumps(payload, separators=(",", ":")).encode("utf-8") + separator + + async def _rewrite_streaming_usage_from_ctx( + self, + response: AsyncIterator[Any], + ctx_response: Optional[UCompletionResponse], + ) -> AsyncIterator[Any]: + pending = b"" + pending_is_str = False + async for chunk in response: + is_str = isinstance(chunk, str) + chunk_bytes = chunk.encode("utf-8") if is_str else chunk + if not isinstance(chunk_bytes, bytes): + yield chunk + continue + + if not pending: + pending_is_str = is_str + pending += chunk_bytes + while separator_match := self._sse_separator_index(pending): + separator_index, separator = separator_match + event_end = separator_index + len(separator) + event = pending[:event_end] + pending = pending[event_end:] + event = self._rewrite_usage_sse_event_from_ctx(event, ctx_response) + yield event.decode("utf-8") if pending_is_str else event + if not pending: + pending_is_str = False + + if pending: + event = self._rewrite_usage_sse_event_from_ctx(pending, ctx_response) + yield event.decode("utf-8") if pending_is_str else event + + def _rewrite_disagg_usage( + self, + response: UCompletionResponseOrGenerator, + ctx_response: Optional[UCompletionResponse], + ) -> UCompletionResponseOrGenerator: + if ctx_response is None: + return response + if hasattr(response, "__aiter__"): + return self._rewrite_streaming_usage_from_ctx(response, ctx_response) + return self._rewrite_usage_response_from_ctx(response, ctx_response) + def _need_gen(self, response: UCompletionResponse) -> bool: if response and response.choices[0].finish_reason not in _GEN_PENDING_FINISH_REASONS: del response.choices[0].disaggregated_params From dd95b12be7e52a36bc282a1c0efe06c2c83e53f3 Mon Sep 17 00:00:00 2001 From: Liao Lanyu <108499334+lancelly@users.noreply.github.com> Date: Wed, 10 Jun 2026 07:11:26 +0800 Subject: [PATCH 16/28] [None][perf] reduce rank-0 GIL contention in disaggregated generation (#15133) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com> --- .../_torch/distributed/communicator.py | 14 ++++---------- .../_torch/pyexecutor/scheduler/adp_router.py | 18 +++++++++--------- 2 files changed, 13 insertions(+), 19 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/communicator.py b/tensorrt_llm/_torch/distributed/communicator.py index 6fbbf16a19df..545ffd2de075 100644 --- a/tensorrt_llm/_torch/distributed/communicator.py +++ b/tensorrt_llm/_torch/distributed/communicator.py @@ -312,18 +312,12 @@ def safe_broadcast(comm, obj, root=0, chunk_size: int = 4 * 1024 * 1024): offset += cur # ---- Reconstruction and deserialization ---- - # Validate the received byte count and unpickle. if rank == root: - # Root already has `serialized` - if len(serialized) != total_size: - raise RuntimeError( - f"Data size mismatch at root: expected {total_size}, got {len(serialized)}" - ) - try: - return pickle.loads(serialized) # nosec B301 - except Exception as e: - raise RuntimeError(f"Deserialization failed: {str(e)}") from e + # Root already holds `obj`; rebuilding it from its own serialized bytes + # would be a needless deep copy. + return obj else: + # Validate the received byte count and unpickle. if len(dst_buf) != total_size: raise RuntimeError( f"Data size mismatch at rank {rank}: expected {total_size}, got {len(dst_buf)}" diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index 84f33c572bf7..dac1030c2bc0 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -448,17 +448,17 @@ def _balance_requests_across_ranks( heapq.heapify(all_ranks_new_requests_heap) - new_requests = sorted( - new_requests, - key=lambda x: len(getattr(x.request, "input_token_ids", [])) if x.request else 0, - reverse=True, - ) + # input_token_ids is a C++ getter that copies the whole token list on every + # access, so read it once per request instead of in both the sort key and + # the loop below. + counted = [ + (len(req.input_token_ids) if (req := item.request) is not None else 0, item) + for item in new_requests + ] + counted.sort(key=lambda item: item[0], reverse=True) - for req_item in new_requests: + for token_count, req_item in counted: val = heapq.heappop(all_ranks_new_requests_heap) - token_count = ( - len(getattr(req_item.request, "input_token_ids", [])) if req_item.request else 0 - ) val = val._replace( num_tokens=val.num_tokens + token_count, num_requests=val.num_requests + 1, From c9e7642ffcc0ca56e1ca45274b3fb7ee3af9b325 Mon Sep 17 00:00:00 2001 From: Liao Lanyu <108499334+lancelly@users.noreply.github.com> Date: Tue, 16 Jun 2026 15:48:27 +0800 Subject: [PATCH 17/28] [None][perf] avoid full input_token_ids copy in ADP router token count (#15199) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com> --- cpp/include/tensorrt_llm/executor/executor.h | 1 + cpp/tensorrt_llm/executor/request.cpp | 5 +++ cpp/tensorrt_llm/executor/requestImpl.h | 5 +++ .../nanobind/executor/request.cpp | 1 + .../_torch/pyexecutor/scheduler/adp_router.py | 25 +++++------- .../_torch/executor/test_adp_router.py | 39 +++++++++++++++++-- 6 files changed, 58 insertions(+), 18 deletions(-) diff --git a/cpp/include/tensorrt_llm/executor/executor.h b/cpp/include/tensorrt_llm/executor/executor.h index 536849cdd87f..9eda85a4d406 100644 --- a/cpp/include/tensorrt_llm/executor/executor.h +++ b/cpp/include/tensorrt_llm/executor/executor.h @@ -758,6 +758,7 @@ class Request ~Request(); [[nodiscard]] VecTokens getInputTokenIds() const; + [[nodiscard]] SizeType32 getNumInputTokens() const; [[nodiscard]] SizeType32 getMaxTokens() const; [[nodiscard]] bool getStreaming() const; [[nodiscard]] SamplingConfig getSamplingConfig() const; diff --git a/cpp/tensorrt_llm/executor/request.cpp b/cpp/tensorrt_llm/executor/request.cpp index e32045892ba7..4d7b1d909c33 100644 --- a/cpp/tensorrt_llm/executor/request.cpp +++ b/cpp/tensorrt_llm/executor/request.cpp @@ -79,6 +79,11 @@ VecTokens Request::getInputTokenIds() const return mImpl->getInputTokenIds(); } +SizeType32 Request::getNumInputTokens() const +{ + return mImpl->getNumInputTokens(); +} + SizeType32 Request::getMaxTokens() const { return mImpl->getMaxNewTokens(); diff --git a/cpp/tensorrt_llm/executor/requestImpl.h b/cpp/tensorrt_llm/executor/requestImpl.h index 55610885b1ae..89cee8569b97 100644 --- a/cpp/tensorrt_llm/executor/requestImpl.h +++ b/cpp/tensorrt_llm/executor/requestImpl.h @@ -135,6 +135,11 @@ class Request::Impl return mInputTokenIds; } + [[nodiscard]] SizeType32 getNumInputTokens() const + { + return static_cast(mInputTokenIds.size()); + } + [[nodiscard]] SizeType32 getMaxNewTokens() const { return mMaxNewTokens; diff --git a/cpp/tensorrt_llm/nanobind/executor/request.cpp b/cpp/tensorrt_llm/nanobind/executor/request.cpp index b502370504d1..3eca465cb544 100644 --- a/cpp/tensorrt_llm/nanobind/executor/request.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/request.cpp @@ -728,6 +728,7 @@ void initRequestBindings(nb::module_& m) nb::arg("cache_salt") = nb::none() ) // clang-format on .def_prop_ro("input_token_ids", &tle::Request::getInputTokenIds) + .def_prop_ro("num_input_tokens", &tle::Request::getNumInputTokens) .def_prop_ro("max_tokens", &tle::Request::getMaxTokens) .def_prop_rw("streaming", &tle::Request::getStreaming, &tle::Request::setStreaming) .def_prop_rw("sampling_config", &tle::Request::getSamplingConfig, &tle::Request::setSamplingConfig) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index dac1030c2bc0..199f9f2d0235 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -36,6 +36,12 @@ HeapVal = namedtuple("HeapVal", ["num_tokens", "num_requests", "rank", "request_list"]) +def _num_input_tokens(request) -> int: + """Token count via the cheap num_input_tokens accessor (avoids materializing + input_token_ids); 0 when request is None.""" + return request.num_input_tokens if request is not None else 0 + + @dataclass class RankIterStatsPayload: """Per-rank IterationStats payload piggybacked on the ADP allgather.""" @@ -448,13 +454,7 @@ def _balance_requests_across_ranks( heapq.heapify(all_ranks_new_requests_heap) - # input_token_ids is a C++ getter that copies the whole token list on every - # access, so read it once per request instead of in both the sort key and - # the loop below. - counted = [ - (len(req.input_token_ids) if (req := item.request) is not None else 0, item) - for item in new_requests - ] + counted = [(_num_input_tokens(item.request), item) for item in new_requests] counted.sort(key=lambda item: item[0], reverse=True) for token_count, req_item in counted: @@ -633,12 +633,6 @@ def _prefix_fingerprint(token_ids, num_tokens: int = 64) -> tuple: return () return tuple(token_ids[:num_tokens]) - @staticmethod - def _req_tokens(req_item) -> int: - if req_item.request is None: - return 0 - return len(getattr(req_item.request, "input_token_ids", [])) - def _match_len(self, rank: int, req_id: int) -> int: matches = self._all_ranks_prefix_matches return matches[rank].get(req_id, 0) if rank < len(matches) else 0 @@ -686,7 +680,8 @@ def get_relax_value(req_item): all_ranks_num_active_requests[target_dp_rank] += 1 # Keep token tally in sync with the soft-balancing phase. effective = max( - self._req_tokens(req_item) - self._match_len(target_dp_rank, req_item.id), + _num_input_tokens(req_item.request) + - self._match_len(target_dp_rank, req_item.id), 0, ) all_ranks_num_active_tokens[target_dp_rank] += effective @@ -731,7 +726,7 @@ def _sort_key(req_item): if not eligible_ranks: break - req_tokens = self._req_tokens(req_item) + req_tokens = _num_input_tokens(req_item.request) req_id = req_item.id # Shuffle eligible ranks per decision so score ties fall to a diff --git a/tests/unittest/_torch/executor/test_adp_router.py b/tests/unittest/_torch/executor/test_adp_router.py index b258e622b848..7b1f58b1a100 100644 --- a/tests/unittest/_torch/executor/test_adp_router.py +++ b/tests/unittest/_torch/executor/test_adp_router.py @@ -21,10 +21,22 @@ KVCacheAwareADPRouter, RankIterStatsPayload, RankState, + _num_input_tokens, ) from tensorrt_llm.scheduling_params import SchedulingParams +class _MockRequest(MagicMock): + """Mock executor Request whose ``num_input_tokens`` mirrors the real binding + (it equals ``len(input_token_ids)``) instead of an auto-generated child mock, + so the routers' cheap token-count reads work without setting it per call site. + """ + + @property + def num_input_tokens(self): + return len(self.input_token_ids) + + def _mock_dist(tp_rank=0, tp_size=1, has_cp_helix=False): """Create a mock Distributed object for testing.""" dist = MagicMock() @@ -35,7 +47,7 @@ def _mock_dist(tp_rank=0, tp_size=1, has_cp_helix=False): def create_mock_request_with_py_schedule_params(attention_dp_rank=None, attention_dp_relax=False): - mock_request = Mock() + mock_request = _MockRequest() if attention_dp_rank is not None: mock_schedule_params = Mock() mock_schedule_params.attention_dp_rank = attention_dp_rank @@ -59,7 +71,7 @@ def _make_request_item(req_id, num_tokens=10, target_dp_rank=None, attention_dp_ scheduling_params = MagicMock() scheduling_params.attention_dp_rank = target_dp_rank scheduling_params.attention_dp_relax = attention_dp_relax - item.request = MagicMock() + item.request = _MockRequest() item.request.py_scheduling_params = scheduling_params item.request.input_token_ids = list(range(num_tokens)) return item @@ -1102,7 +1114,7 @@ def _make_conv_request_item( scheduling_params = MagicMock() scheduling_params.attention_dp_rank = target_dp_rank scheduling_params.attention_dp_relax = attention_dp_relax - item.request = MagicMock() + item.request = _MockRequest() item.request.py_scheduling_params = scheduling_params item.request.input_token_ids = list(range(num_tokens)) item.request.py_orig_prompt_len = num_tokens @@ -1259,3 +1271,24 @@ def test_returned_expected_covers_every_rank(self): final = [states[r].num_active_requests + len(assign[r]) for r in range(8)] assert all(expected >= f for f in final), (expected, final) assert sum(len(v) for v in assign.values()) == 40 # nothing dropped + + +class TestNumInputTokens: + """The cheap token-count accessor used by the routers (avoids copying the + full input_token_ids list just to count it).""" + + def test_none_request(self): + assert _num_input_tokens(None) == 0 + + def test_prefers_num_input_tokens_without_touching_token_list(self): + """When num_input_tokens is available it must be used, and the heavy + input_token_ids getter must NOT be accessed.""" + + class _Req: + num_input_tokens = 1234 + + @property + def input_token_ids(self): # materializing it would be a regression + raise AssertionError("input_token_ids must not be materialized") + + assert _num_input_tokens(_Req()) == 1234 From 983fe50b7d80e08974958eb598e55ea262b42ab7 Mon Sep 17 00:00:00 2001 From: Liao Lanyu <108499334+lancelly@users.noreply.github.com> Date: Wed, 27 May 2026 04:19:42 +0800 Subject: [PATCH 18/28] [None][fix] py_executor: lift conditional tp_allgather sites out of per-rank-divergent gates (#14560) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 102 ++++++++++++++---- 1 file changed, 79 insertions(+), 23 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index f96cacc72600..572920417a29 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -598,6 +598,12 @@ def __init__( # responses and flushing them at a synchronised point in the executor # loop avoids the mismatch. self._pending_transfer_responses: List[Tuple[int, LlmResponse]] = [] + # Same buffer-then-synced-flush pattern as _pending_transfer_responses + # above: _handle_responses and _append_iter_stats are reached from + # per-rank-divergent gates, so their tp_allgather collectives are + # lifted to _handle_kv_transfer_timeouts_synced / _flush_iter_stats_synced. + self._pending_timed_out_requests: List[LlmRequest] = [] + self._pending_iter_stats_dict: Optional[Dict] = None # ADP dummy role for _pad_attention_dp_dummy_request. Default is gen; # updated from observed request types. self._adp_dummy_is_gen: bool = True @@ -981,6 +987,53 @@ def _flush_pending_transfer_responses(self): # can complete. self._enqueue_responses(responses) + def _handle_kv_transfer_timeouts_synced(self): + """ADP-safe drain of the KV-transfer-timeout consensus collective. + + Lifted out of _handle_responses, which is reached from per-rank- + divergent gates (_process_previous_batch, _handle_executed_batch). + Non-ADP runs handle timeouts inline; the buffer is empty here. + """ + if not (self.enable_attention_dp and self.dist.world_size != 1): + return + timed_out = self._pending_timed_out_requests + self._pending_timed_out_requests = [] + any_timed_out = any(self.dist.tp_allgather(bool(timed_out))) + if any_timed_out: + self._handle_errors(error_msg="Request timed out (KV transfer)", + requests=timed_out, + charge_budget=False) + + def _flush_iter_stats_synced(self): + """ADP-safe drain of the TLLM_METRICS_ALL_RANKS dict-gather collective. + + Lifted out of _append_iter_stats (same divergent-gate problem as + _handle_kv_transfer_timeouts_synced). Under ADP every rank sends + either the local dict or ``None`` every iter so the gather stays in + lockstep with peer collectives. When the gate doesn't hold, + _append_iter_stats takes its legacy path and never populates the + buffer. + """ + tp_size = self.dist.tp_size + gather_all_ranks = os.environ.get("TLLM_METRICS_ALL_RANKS", "0") == "1" + if not (gather_all_ranks and self.enable_iter_perf_stats and tp_size > 1 + and self.enable_attention_dp): + return + local_dict = self._pending_iter_stats_dict + self._pending_iter_stats_dict = None + gathered = self.dist.tp_allgather(local_dict) + if self.dist.tp_rank != 0: + return + rank_dicts = [d for d in gathered if d is not None] + if not rank_dicts: + return + with self.stats_lock: + for d in rank_dicts: + self.stats.append(("per_rank_dict", d)) + cap = self.max_stats_len * tp_size + if len(self.stats) > cap: + del self.stats[:len(self.stats) - cap] + # Performance metrics methods are in PerfMetricsManager (self.perf_manager) def _event_loop_wrapper(self): @@ -1903,23 +1956,9 @@ def _append_iter_stats(self, local_dict["gpuForwardTimeMS"] = gpu_forward_time_ms local_dict["schedulerMode"] = scheduler_mode local_dict["rank"] = self.dist.tp_rank - - gathered = self.dist.tp_allgather(local_dict) - - if self.dist.tp_rank == 0: - with self.stats_lock: - # Wrap as ("per_rank_dict", dict) so the serializer can - # distinguish from the legacy (stats, req_stats, kv) tuple. - # Trim before appending so every rank's entry for the - # same iteration lands in the buffer atomically; evicting - # during the append loop would let partial iterations - # survive at the head of self.stats. - cap = self.max_stats_len * tp_size - overflow = max(0, len(self.stats) + len(gathered) - cap) - if overflow: - del self.stats[:overflow] - for d in gathered: - self.stats.append(("per_rank_dict", d)) + # Buffer for _flush_iter_stats_synced; in-place tp_allgather would + # desync (reached from per-rank-divergent gates). + self._pending_iter_stats_dict = local_dict return # Legacy path: rank-0-only (single-rank or iter stats disabled). @@ -2578,6 +2617,13 @@ def _handle_executed_batch(self, executed_batch: Optional[BatchStatePP]): executed_batch.microbatch_id % self.dist.pp_size, ) + # Drain the synced-collective buffers outside the per-rank-divergent + # ``executed_batch is not None`` branch. handle_executed_batches + # invokes this helper the same number of times on every rank + # (broadcast count), so the flushes are rank-symmetric. + self._handle_kv_transfer_timeouts_synced() + self._flush_iter_stats_synced() + @nvtx_range("wait_on_pp_send_handles") def wait_on_pp_send_handles(self, send_handles, microbatch_id): if send_handles[microbatch_id] is not None: @@ -3309,6 +3355,10 @@ def _executor_loop(self): if self.enable_kv_cache_events: self._add_kv_cache_events() + # Drain timeout buffer outside ``if can_queue`` so the synced + # collective fires every iter regardless of future restructuring. + self._handle_kv_transfer_timeouts_synced() + if self.kv_cache_transceiver and self.async_transfer_manager.has_any_inflight_requests( ): self._check_kv_transfer_timeout() @@ -3330,6 +3380,9 @@ def _executor_loop(self): elif gpu_forward_events_from_perf_pool: self.perf_manager.release_forward_timing_events( gpu_forward_start, gpu_forward_end) + # Same lockstep guarantee for iter-stats; no-op when + # TLLM_METRICS_ALL_RANKS=0. + self._flush_iter_stats_synced() self.iter_counter += 1 @@ -3756,6 +3809,12 @@ def _executor_loop_overlap(self): else: self._enqueue_responses([]) + # Drain buffers from the (per-rank-divergent) + # _process_previous_batch above; rank-symmetric companion to + # the _enqueue_responses([]) call in the else branch. + self._handle_kv_transfer_timeouts_synced() + self._flush_iter_stats_synced() + # Call set_exclude_last_generation_logits after _process_previous_batch. # If set before, the response of a request may be incorrect, as it will # use the wrong indices for generation logits when streaming is enabled. @@ -5678,12 +5737,9 @@ def _handle_responses(self, emit_first_iter: bool = True): self._terminate_request(request) if (self.kv_cache_transceiver is not None and self.enable_attention_dp and self.dist.world_size != 1): - any_timed_out = any(self.dist.tp_allgather( - bool(timed_out_requests))) - if any_timed_out: - self._handle_errors(error_msg="Request timed out (KV transfer)", - requests=timed_out_requests, - charge_budget=False) + # Buffer for _handle_kv_transfer_timeouts_synced; in-place + # tp_allgather would desync (reached from per-rank-divergent gates). + self._pending_timed_out_requests.extend(timed_out_requests) else: for req in timed_out_requests: self._handle_errors( From d667c3180dce0d4a3cdde7c4d1881d7a379a48a3 Mon Sep 17 00:00:00 2001 From: Liao Lanyu <108499334+lancelly@users.noreply.github.com> Date: Tue, 2 Jun 2026 21:53:07 +0800 Subject: [PATCH 19/28] [None][fix] disagg: echo prompt_token_ids on chat context-only response (#14844) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com> --- .../test_openai_disagg_service.py | 92 +++++++++++++++++++ 1 file changed, 92 insertions(+) diff --git a/tests/unittest/disaggregated/test_openai_disagg_service.py b/tests/unittest/disaggregated/test_openai_disagg_service.py index 1eca36001798..fa1618ecfe6e 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_service.py +++ b/tests/unittest/disaggregated/test_openai_disagg_service.py @@ -14,6 +14,7 @@ import asyncio from types import SimpleNamespace +from unittest import mock from unittest.mock import AsyncMock import pytest @@ -31,6 +32,9 @@ from tensorrt_llm.serve.openai_disagg_service import OpenAIDisaggregatedService from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatMessage, CompletionRequest, CompletionResponse, CompletionResponseChoice, @@ -101,12 +105,100 @@ def _make_completion_response( ) +def _make_chat_response( + finish_reason: str, + disagg_request_id: int = 42, + prompt_token_ids=None, +) -> ChatCompletionResponse: + if prompt_token_ids is None: + prompt_token_ids = [1, 2, 3] + return ChatCompletionResponse( + model="test-model", + usage=UsageInfo(prompt_tokens=1, completion_tokens=1), + prompt_token_ids=prompt_token_ids, + choices=[ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content=""), + finish_reason=finish_reason, + disaggregated_params=DisaggregatedParams( + request_type="context_only", + disagg_request_id=disagg_request_id, + ctx_request_id=disagg_request_id, + ), + ) + ], + ) + + async def _mock_streaming_response(chunks): for chunk in chunks: await asyncio.sleep(0) yield chunk +def test_get_gen_request_uses_ctx_response_prompt_token_ids_for_chat(): + service = _make_service("context_first") + ctx_prompt_token_ids = [101, 102, 103] + request = ChatCompletionRequest( + model="test-model", + messages=[ + { + "role": "user", + "content": "hello", + } + ], + prompt_token_ids=[1, 2, 3], + ) + ctx_response = _make_chat_response( + finish_reason="length", + prompt_token_ids=ctx_prompt_token_ids, + ) + + gen_request = service._get_gen_request(request, ctx_response, 42) + + assert gen_request.prompt_token_ids == ctx_prompt_token_ids + assert gen_request.disaggregated_params.request_type == "generation_only" + + +@pytest.mark.asyncio +async def test_create_chat_response_sets_prompt_token_ids_for_context_only(): + from tensorrt_llm.serve.openai_server import OpenAIServer + + prompt_token_ids = [101, 102, 103] + + class FakePromise: + def __init__(self) -> None: + self.outputs = [] + self.prompt_token_ids = prompt_token_ids + + async def aresult(self): + return self + + server = OpenAIServer.__new__(OpenAIServer) + server.generator = mock.MagicMock() + server.generator.args.num_postprocess_workers = 0 + server._extract_metrics = mock.AsyncMock() + + post_processor = mock.Mock( + return_value=_make_chat_response( + finish_reason="length", + prompt_token_ids=None, + ) + ) + postproc_params = SimpleNamespace(post_processor=post_processor, postproc_args=object()) + disaggregated_params = SimpleNamespace(request_type="context_only", disagg_request_id=123) + + response = await server._create_chat_response( + FakePromise(), + postproc_params, + raw_request=None, + disaggregated_params=disaggregated_params, + ) + + assert response.prompt_token_ids == prompt_token_ids + + @pytest.mark.asyncio async def test_is_ready_waits_for_router_preparation(): service = _make_service("context_first") From 299e6824fb19d20ff8eba00fa2c885fb90a4d6b9 Mon Sep 17 00:00:00 2001 From: Jin Li <59594262+liji-nv@users.noreply.github.com> Date: Mon, 22 Jun 2026 19:40:28 -0700 Subject: [PATCH 20/28] [None][fix] disagg: use ctx retry id for gen request (#15309) Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> Signed-off-by: Fanrong Li --- tensorrt_llm/serve/openai_disagg_service.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index 9c4bf8d5bbcf..32fab84d861b 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -150,6 +150,9 @@ async def _send_disagg_request_ctx_first( ctx_req, server=ctx_server, hooks=hooks ) await self._verify_ctx_response(ctx_response) + ctx_response_disagg_params = ctx_response.choices[0].disaggregated_params + if ctx_response_disagg_params.disagg_request_id is not None: + disagg_request_id = ctx_response_disagg_params.disagg_request_id gen_req = self._get_gen_request(request, ctx_response, disagg_request_id) else: # Clear synthetic disaggregated_params that may have been From 7be4388007dc18b7911336951204ebb422e476c1 Mon Sep 17 00:00:00 2001 From: Liao Lanyu <108499334+lancelly@users.noreply.github.com> Date: Tue, 23 Jun 2026 12:25:23 +0800 Subject: [PATCH 21/28] [None][perf] disagg: optionally drop message history from generation request (#15366) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com> --- tensorrt_llm/llmapi/disagg_utils.py | 6 ++++++ tensorrt_llm/serve/openai_disagg_service.py | 13 +++++++++++++ 2 files changed, 19 insertions(+) diff --git a/tensorrt_llm/llmapi/disagg_utils.py b/tensorrt_llm/llmapi/disagg_utils.py index 6afa17c2626e..a83c784b3488 100644 --- a/tensorrt_llm/llmapi/disagg_utils.py +++ b/tensorrt_llm/llmapi/disagg_utils.py @@ -97,6 +97,10 @@ class DisaggServerConfig(): schedule_style: Literal['context_first', 'generation_first'] = 'context_first' allow_request_chat_template: bool = False + # Drop conversation history from generation_only requests to shrink the gen + # 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 @dataclass @@ -181,6 +185,7 @@ def extract_disagg_cfg(hostname: str = 'localhost', 'context_first', 'generation_first'] = 'context_first', allow_request_chat_template: bool = False, + gen_strip_message_history: bool = False, **kwargs: Any) -> DisaggServerConfig: context_servers = context_servers or {} generation_servers = generation_servers or {} @@ -230,6 +235,7 @@ def extract_disagg_cfg(hostname: str = 'localhost', config.schedule_style = schedule_style config.allow_request_chat_template = validate_config_bool( allow_request_chat_template, "allow_request_chat_template") + config.gen_strip_message_history = gen_strip_message_history return config diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index 32fab84d861b..b04689f8682d 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -80,6 +80,8 @@ def __init__( self._perf_metrics_collector = perf_metrics_collector self._cluster_storage = disagg_cluster_storage 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 self._ctx_client = None self._gen_client = None @@ -358,6 +360,17 @@ def _get_gen_request( request.prompt = ctx_response.prompt_token_ids elif isinstance(request, ChatCompletionRequest): 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 + # message; tools are preserved. Config-gated because it's unsafe + # for harmony/multimodal workers (model type is fixed per deploy). + if ( + self._strip_gen_message_history + and request.messages + and len(request.messages) > 1 + ): + request.messages = request.messages[-1:] else: # no ctx response, it's either a generation-only request or a generation-first disagg request request.disaggregated_params = DisaggregatedParams( From db2ac6cd63997e838e5b3434d7f466d1da85e92e Mon Sep 17 00:00:00 2001 From: Fanrong Li Date: Thu, 25 Jun 2026 11:16:16 +0000 Subject: [PATCH 22/28] [None][fix] keep disagg routing PR self-contained Signed-off-by: Fanrong Li --- .../_torch/pyexecutor/scheduler/adp_router.py | 7 ++++++- tensorrt_llm/inputs/utils.py | 14 +++++++++++++- tensorrt_llm/serve/openai_server.py | 2 ++ tensorrt_llm/serve/router.py | 12 +++--------- .../disaggregated/test_openai_disagg_service.py | 1 + 5 files changed, 25 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index 199f9f2d0235..57741cf92c93 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -39,7 +39,12 @@ def _num_input_tokens(request) -> int: """Token count via the cheap num_input_tokens accessor (avoids materializing input_token_ids); 0 when request is None.""" - return request.num_input_tokens if request is not None else 0 + if request is None: + return 0 + num_input_tokens = getattr(request, "num_input_tokens", None) + if isinstance(num_input_tokens, int): + return num_input_tokens + return len(getattr(request, "input_token_ids", [])) @dataclass diff --git a/tensorrt_llm/inputs/utils.py b/tensorrt_llm/inputs/utils.py index fcc53dbc73f1..5dc8c3a343bb 100644 --- a/tensorrt_llm/inputs/utils.py +++ b/tensorrt_llm/inputs/utils.py @@ -27,7 +27,8 @@ _safe_aiohttp_get, _safe_request_get) from tensorrt_llm.inputs.media_io import \ convert_image_mode as convert_image_mode -from tensorrt_llm.inputs.multimodal import MultimodalServerConfig +from tensorrt_llm.inputs.multimodal import (MultimodalServerConfig, + default_hasher) from tensorrt_llm.inputs.multimodal_data import \ BaseModalityData as BaseModalityData from tensorrt_llm.inputs.multimodal_data import VideoData as VideoData @@ -894,3 +895,14 @@ def convert_to_conversation_message( inputs.append(input) return inputs + + +def get_cache_salt_id(cache_salt: str) -> int: + b = cache_salt.encode("utf-8") + h = default_hasher(b).digest(length=8) + cache_salt_id = int.from_bytes(h, "little", signed=False) + if cache_salt_id < 0 or cache_salt_id >= (1 << 64): + raise ValueError( + f"cache_salt_id must be in [0, 2**64 - 1], got {cache_salt_id}.") + + return cache_salt_id diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 385537f70089..3e9c54cb6cba 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -47,6 +47,8 @@ from tensorrt_llm.media.encoding import image_to_bytes from tensorrt_llm.media.tensor_payload import is_tensor_format from tensorrt_llm.metrics.collector import MetricsCollector +from tensorrt_llm.runtime.kv_cache_hash import \ + get_effective_kv_cache_event_hash_algo from tensorrt_llm.sampling_params import GuidedDecodingParams from tensorrt_llm.serve.chat_utils import (load_chat_template, parse_chat_messages_coroutines, diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 2545132094e3..c257d0eaa820 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -20,7 +20,6 @@ from typing import Awaitable, Callable, Dict, Iterable, List, Optional, Union import aiohttp -from transformers import AutoTokenizer from tensorrt_llm.bindings.internal.batch_manager import \ BlockKey as _NativeBlockKey @@ -825,15 +824,10 @@ def _get_tokenizer(self, model: str): self._tokenizers[model] = load_custom_tokenizer( self._custom_tokenizer, model) else: - from tensorrt_llm.tokenizer import \ - maybe_fix_byte_level_tokenizer - tokenizer = AutoTokenizer.from_pretrained( + from tensorrt_llm.tokenizer import TransformersTokenizer + tokenizer = TransformersTokenizer.from_pretrained( model, trust_remote_code=True) - # Work around Transformers 5.x LlamaTokenizer overriding - # tokenizer.json's ByteLevel pre-tokenizer with Metaspace, - # which silently strips spaces from prompts (see tokenizer.py). - self._tokenizers[model] = maybe_fix_byte_level_tokenizer( - tokenizer, model, trust_remote_code=True) + self._tokenizers[model] = tokenizer.tokenizer return self._tokenizers[model] def _encode_with_prefix_cache(self, rendered: str, key: int, diff --git a/tests/unittest/disaggregated/test_openai_disagg_service.py b/tests/unittest/disaggregated/test_openai_disagg_service.py index fa1618ecfe6e..4d7d813304f9 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_service.py +++ b/tests/unittest/disaggregated/test_openai_disagg_service.py @@ -171,6 +171,7 @@ class FakePromise: def __init__(self) -> None: self.outputs = [] self.prompt_token_ids = prompt_token_ids + self.error = None async def aresult(self): return self From 55c9e75990b80fab4574981695d7f35bc80efb99 Mon Sep 17 00:00:00 2001 From: Fanrong Li Date: Fri, 26 Jun 2026 01:27:52 +0000 Subject: [PATCH 23/28] [None][feat] add conversation params routing support Signed-off-by: Fanrong Li --- tensorrt_llm/_torch/pyexecutor/_util.py | 2 + tensorrt_llm/_torch/pyexecutor/llm_request.py | 3 + .../_torch/pyexecutor/request_utils.py | 22 +++- .../_torch/pyexecutor/scheduler/adp_router.py | 9 +- tensorrt_llm/conversation_params.py | 35 +++++ tensorrt_llm/executor/base_worker.py | 4 + tensorrt_llm/executor/executor.py | 5 + tensorrt_llm/executor/request.py | 3 + tensorrt_llm/llmapi/__init__.py | 2 + tensorrt_llm/llmapi/llm.py | 9 ++ tensorrt_llm/llmapi/llm_args.py | 8 +- tensorrt_llm/serve/conversation_id.py | 76 +++++++++++ tensorrt_llm/serve/openai_disagg_server.py | 46 +------ tensorrt_llm/serve/openai_disagg_service.py | 14 +- tensorrt_llm/serve/openai_protocol.py | 37 ++++++ tensorrt_llm/serve/openai_server.py | 19 ++- tensorrt_llm/serve/router.py | 7 +- .../defs/disaggregated/test_workers.py | 123 +++++++++++------- .../_torch/executor/test_adp_router.py | 15 ++- .../_torch/executor/test_request_utils.py | 37 ++++++ .../references/conversation_params.yaml | 8 ++ .../api_stability/references/llm.yaml | 6 + tests/unittest/api_stability/test_llm_api.py | 9 +- .../test_openai_disagg_server.py | 27 ++-- .../test_openai_disagg_service.py | 3 + tests/unittest/disaggregated/test_router.py | 13 +- .../apps/test_openai_conversation_params.py | 61 +++++++++ .../llmapi/test_conversation_params.py | 29 +++++ 28 files changed, 492 insertions(+), 140 deletions(-) create mode 100644 tensorrt_llm/conversation_params.py create mode 100644 tensorrt_llm/serve/conversation_id.py create mode 100644 tests/unittest/api_stability/references/conversation_params.yaml create mode 100644 tests/unittest/llmapi/apps/test_openai_conversation_params.py create mode 100644 tests/unittest/llmapi/test_conversation_params.py diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index ac0feab0347d..19e48ebe8f19 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -483,6 +483,7 @@ def _create_dummy_mm_context_request( torch.zeros(1, 1, dtype=torch.int32) } } + request.py_conversation_params = None remaining_tokens -= len(prompt_token_ids) requests.append(request) @@ -532,6 +533,7 @@ def _create_dummy_context_requests( torch.zeros(1, 1, dtype=torch.int32) } } + request.py_conversation_params = None requests.append(request) remaining_tokens -= input_seq_len if self._mapping.enable_attention_dp: diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 26d91045b1ea..10a635e01f61 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -801,6 +801,7 @@ def __init__( self.py_logprobs_mode = LogprobMode( logprobs_mode) # handle passed a raw string self.py_disaggregated_params = None + self.py_conversation_params = None self.py_num_connector_matched_tokens = 0 @@ -1191,6 +1192,8 @@ def executor_request_to_llm_request( llm_request.py_disaggregated_params = getattr(executor_request, "py_disaggregated_params", None) + llm_request.py_conversation_params = getattr(executor_request, + "py_conversation_params", None) if child_req_ids: for child_id in child_req_ids: llm_request.create_child_request(child_id) diff --git a/tensorrt_llm/_torch/pyexecutor/request_utils.py b/tensorrt_llm/_torch/pyexecutor/request_utils.py index 0a8866fdfc69..c40303dc1548 100644 --- a/tensorrt_llm/_torch/pyexecutor/request_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/request_utils.py @@ -33,13 +33,15 @@ def get_num_child_requests(request: ExecutorRequest) -> int: def collect_py_objects_from_requests( - requests: List, attribute_name: str + requests: List, attribute_name: str, include_none: bool = False ) -> Optional[Tuple[str, Dict]]: """Collect Python-only objects from requests. Args: requests: List of RequestQueueItem objects. attribute_name: Name of the attribute to collect. + include_none: Include requests whose attribute value is None. When + enabled, the source request must have the attribute. Returns: Tuple of (attribute_name, dict mapping request_id to object) or None if empty. @@ -49,9 +51,12 @@ def collect_py_objects_from_requests( if not item.is_normal_request: continue if item.request: - obj = getattr(item.request, attribute_name, None) - if obj is not None: - req_id_to_obj[item.id] = obj + if include_none: + req_id_to_obj[item.id] = getattr(item.request, attribute_name) + else: + obj = getattr(item.request, attribute_name, None) + if obj is not None: + req_id_to_obj[item.id] = obj return None if not req_id_to_obj else (attribute_name, req_id_to_obj) @@ -65,9 +70,8 @@ def attach_py_objects_to_requests(requests: List, py_request_objects: Tuple) -> for attr_name, req_obj_dict in py_request_objects: for item in requests: if item.request: - py_obj = req_obj_dict.get(item.id) - if py_obj is not None: - setattr(item.request, attr_name, py_obj) + if item.id in req_obj_dict: + setattr(item.request, attr_name, req_obj_dict[item.id]) def derive_attention_dp_per_rank_request_cap( @@ -527,6 +531,9 @@ def _collect_py_objects(self, new_requests: List) -> Tuple: py_disaggregated_params = collect_py_objects_from_requests( new_requests, "py_disaggregated_params" ) + py_conversation_params = collect_py_objects_from_requests( + new_requests, "py_conversation_params", include_none=True + ) py_lora_path = collect_py_objects_from_requests(new_requests, "py_lora_path") return tuple( @@ -538,6 +545,7 @@ def _collect_py_objects(self, new_requests: List) -> Tuple: py_scheduling_params, py_num_logprobs, py_disaggregated_params, + py_conversation_params, py_lora_path, ], ) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index 57741cf92c93..31bdf4947154 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -836,14 +836,13 @@ def create_rank_state( @staticmethod def _conversation_id(req_item) -> "str | None": - req = getattr(req_item, "request", None) + req = req_item.request if req is None: return None - disagg = getattr(req, "py_disaggregated_params", None) - if disagg is None: + conversation_params = req.py_conversation_params + if conversation_params is None: return None - conv_id = getattr(disagg, "conversation_id", None) - return conv_id if conv_id else None + return conversation_params.conversation_id def _record_target_rank(self, conv_id: str, rank: int) -> None: """Bind/refresh the rank a conversation is pinned to (LRU-touch + evict).""" diff --git a/tensorrt_llm/conversation_params.py b/tensorrt_llm/conversation_params.py new file mode 100644 index 000000000000..c7e450befca8 --- /dev/null +++ b/tensorrt_llm/conversation_params.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass + + +@dataclass(slots=True, kw_only=True) +class ConversationParams: + """Conversation parameters. + + Args: + conversation_id (str): Stable multi-turn conversation id used for routing. + """ + + conversation_id: str + + def __post_init__(self) -> None: + if self.conversation_id is None: + raise ValueError("conversation_id must be non-empty") + conversation_id = str(self.conversation_id).strip() + if not conversation_id: + raise ValueError("conversation_id must be non-empty") + self.conversation_id = conversation_id diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index a03fbb4a3c56..b2263242a2e5 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -572,6 +572,10 @@ def _deduce_max_tokens(request: GenerationRequest, if self._is_pytorch_backend and request.scheduling_params is not None: executor_request.py_scheduling_params = request.scheduling_params + executor_request.py_conversation_params = None + if self._is_pytorch_backend and request.conversation_params is not None: + executor_request.py_conversation_params = request.conversation_params + if request.arrival_time is not None: executor_request.py_arrival_time = request.arrival_time diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py index fdd3dff44a48..c61deee740d6 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -21,6 +21,7 @@ from .._utils import mpi_world_size from ..bindings import executor as tllm from ..builder import Engine +from ..conversation_params import ConversationParams from ..disaggregated_params import DisaggregatedParams from ..llmapi.llm_args import BaseLlmArgs, TorchLlmArgs from ..llmapi.llm_utils import KvCacheRetentionConfig @@ -134,6 +135,7 @@ def generate_async( postproc_params: Optional[PostprocParams] = None, multimodal_params: Optional[MultimodalParams] = None, scheduling_params: Optional[SchedulingParams] = None, + conversation_params: Optional[ConversationParams] = None, cache_salt: Optional[str] = None, arrival_time: Optional[float] = None, encoder_input_token_ids: Optional[Union[torch.Tensor, np.ndarray, @@ -164,6 +166,7 @@ def generate_async( trace_headers=trace_headers, multimodal_params=multimodal_params, scheduling_params=scheduling_params, + conversation_params=conversation_params, cache_salt=cache_salt, arrival_time=arrival_time, encoder_input_token_ids=encoder_input_token_ids, @@ -183,6 +186,7 @@ def generate( prompt_adapter_request: Optional[Union[ PromptAdapterRequest, List[PromptAdapterRequest]]] = None, disaggregated_params: Optional[DisaggregatedParams] = None, + conversation_params: Optional[ConversationParams] = None, cache_salt: Optional[Union[str, List[Optional[str]]]] = None, ) -> Union[GenerationResult, List[GenerationResult]]: """Generate output for the given prompt token ids in the synchronous mode. @@ -218,6 +222,7 @@ def generate( prompt_adapter_request=pa_req, streaming=False, disaggregated_params=disaggregated_params, + conversation_params=conversation_params, cache_salt=cs) futures.append(future) diff --git a/tensorrt_llm/executor/request.py b/tensorrt_llm/executor/request.py index fdd6ba0627ce..a43e89c0e4f7 100644 --- a/tensorrt_llm/executor/request.py +++ b/tensorrt_llm/executor/request.py @@ -8,6 +8,7 @@ from tensorrt_llm.inputs.multimodal import MultimodalParams +from ..conversation_params import ConversationParams from ..disaggregated_params import DisaggregatedParams from ..llmapi.llm_utils import KvCacheRetentionConfig from ..sampling_params import SamplingParams @@ -107,6 +108,7 @@ def __init__( postproc_params: Optional[PostprocParams] = None, multimodal_params: Optional[MultimodalParams] = None, scheduling_params: Optional[SchedulingParams] = None, + conversation_params: Optional[ConversationParams] = None, cache_salt: Optional[str] = None, arrival_time: Optional[float] = None, encoder_input_token_ids: Optional[Union[torch.Tensor, np.ndarray, @@ -138,6 +140,7 @@ def __init__( self.disaggregated_params = disaggregated_params self.trace_headers = trace_headers self.scheduling_params = scheduling_params + self.conversation_params = conversation_params if cache_salt is not None: if not isinstance(cache_salt, str): raise TypeError( diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 96032c23c227..1682eb51e71c 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -1,4 +1,5 @@ from .._torch.async_llm import AsyncLLM +from ..conversation_params import ConversationParams from ..disaggregated_params import DisaggregatedParams, DisaggScheduleStyle from ..executor import CompletionOutput, LoRARequest, RequestError from ..sampling_params import GuidedDecodingParams, SamplingParams @@ -40,6 +41,7 @@ 'GuidedDecodingParams', 'SamplingParams', 'DisaggregatedParams', + 'ConversationParams', 'DisaggScheduleStyle', 'KvCacheConfig', 'KvCacheRetentionConfig', diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index f2ceb32da9b1..e2674df31c3e 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -27,6 +27,7 @@ from ..bindings import executor as tllm from ..bindings import steady_clock_now from ..builder import EngineConfig +from ..conversation_params import ConversationParams from ..disaggregated_params import DisaggregatedParams from ..executor import (DetokenizedGenerationResultBase, GenerationExecutor, GenerationResult, IterationResult, LoRARequest, @@ -489,6 +490,8 @@ def generate( DisaggregatedParams, Sequence[DisaggregatedParams]]] = None, scheduling_params: Optional[Union[SchedulingParams, List[SchedulingParams]]] = None, + conversation_params: Optional[Union[ConversationParams, + List[ConversationParams]]] = None, cache_salt: Optional[Union[str, Sequence[str]]] = None, priority: Union[float, List[float]] = DEFAULT_REQUEST_PRIORITY, ) -> Union[RequestOutput, List[RequestOutput]]: @@ -511,6 +514,8 @@ def generate( Disaggregated parameters. Defaults to None. scheduling_params (tensorrt_llm.scheduling_params.SchedulingParams, List[tensorrt_llm.scheduling_params.SchedulingParams], optional): Scheduling parameters. Defaults to None. + conversation_params (tensorrt_llm.conversation_params.ConversationParams, List[tensorrt_llm.conversation_params.ConversationParams], optional): + Conversation parameters. Defaults to None. cache_salt (str, Sequence[str], optional): If specified, KV cache will be salted with the provided string to limit the kv cache reuse to the requests with the same string. Defaults to None. priority (float, List[float]): The scheduling priority for the request(s), in the range [0, 1]. Higher values indicate higher priority. Defaults to 0.5. @@ -555,6 +560,7 @@ def generate( kv_cache_retention_config, i), disaggregated_params=self._item_at(disaggregated_params, i), scheduling_params=self._item_at(scheduling_params, i), + conversation_params=self._item_at(conversation_params, i), cache_salt=self._item_at(cache_salt, i), priority=self._item_at(priority, i), streaming=False, @@ -585,6 +591,7 @@ def generate_async( trace_headers: Optional[Mapping[str, str]] = None, _postproc_params: Optional[PostprocParams] = None, scheduling_params: Optional[SchedulingParams] = None, + conversation_params: Optional[ConversationParams] = None, cache_salt: Optional[str] = None, priority: float = DEFAULT_REQUEST_PRIORITY, ) -> RequestOutput: @@ -602,6 +609,7 @@ def generate_async( disaggregated_params (tensorrt_llm.disaggregated_params.DisaggregatedParams, optional): Disaggregated parameters. Defaults to None. trace_headers (Mapping[str, str], optional): Trace headers. Defaults to None. scheduling_params (tensorrt_llm.scheduling_params.SchedulingParams, optional): Scheduling parameters. Defaults to None. + conversation_params (tensorrt_llm.conversation_params.ConversationParams, optional): Conversation parameters. Defaults to None. cache_salt (str, optional): If specified, KV cache will be salted with the provided string to limit the kv cache reuse to the requests with the same string. Defaults to None. priority (float): The scheduling priority for the request, in the range [0, 1]. Higher values indicate higher priority. Defaults to 0.5. @@ -671,6 +679,7 @@ def generate_async( postproc_params=_postproc_params, multimodal_params=multimodal_params, scheduling_params=scheduling_params, + conversation_params=conversation_params, cache_salt=cache_salt, arrival_time=arrival_time, encoder_input_token_ids=encoder_input_token_ids, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 330c93eb8bc9..9d9d47810e81 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1293,9 +1293,11 @@ class AttentionDpConfig(StrictBaseModel): description= "Enable explicit conversation-affinity routing for attention DP. When " "True, the first request of each conversation is round-robined across " - "ranks and every subsequent request carrying the same conversation_id " - "(read from disaggregated_params, populated from the X-Session-ID " - "header) is pinned to that conversation's first-turn rank. This keeps a " + "ranks and every subsequent request carrying the same " + "conversation_params.conversation_id is pinned to that conversation's " + "first-turn rank. OpenAI requests use the body conversation_params as " + "canonical; the serve edge only creates conversation_params from the " + "X-Session-ID header when the body does not provide it. This keeps a " "multi-turn conversation's KV-cache prefix on one rank (maximizing " "block reuse, minimizing cross-rank migration). Unlike " "enable_kv_cache_aware_routing (affinity inferred from prefix-match " diff --git a/tensorrt_llm/serve/conversation_id.py b/tensorrt_llm/serve/conversation_id.py new file mode 100644 index 000000000000..7e0fed21640e --- /dev/null +++ b/tensorrt_llm/serve/conversation_id.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Mapping, Optional, Protocol + +# Supported HTTP header protocol for external clients, gateways, or proxies +# that carry a stable multi-turn identifier outside the JSON body. Body +# ``conversation_params.conversation_id`` is canonical when both body and +# headers are set; the serve edge copies the first non-empty header value into +# ``request.conversation_params`` only when the body omits it. Routers then read +# ``conversation_params.conversation_id`` to keep later turns of the same +# conversation on the same backend when sticky conversation routing is enabled. +CONVERSATION_ID_HEADERS = ( + "x-session-id", + "x-correlation-id", + "x-session-affinity", + "x-multi-turn-session-id", +) + + +class RequestWithConversationParams(Protocol): + conversation_params: Any + + +def get_request_conversation_id(request: RequestWithConversationParams) -> Optional[str]: + conversation_params = request.conversation_params + if conversation_params is None: + return None + return conversation_params.conversation_id + + +def extract_conversation_id_from_headers(headers: Optional[Mapping[str, str]]) -> Optional[str]: + if headers is None: + return None + lower_headers = {str(key).lower(): value for key, value in headers.items()} + for header_name in CONVERSATION_ID_HEADERS: + conversation_id = lower_headers.get(header_name) + if conversation_id is None: + continue + conversation_id = str(conversation_id).strip() + if conversation_id: + return conversation_id + return None + + +def resolve_request_conversation_id( + request: RequestWithConversationParams, + headers: Optional[Mapping[str, str]] = None, +) -> Optional[str]: + """Return conversation_params.conversation_id populated at the serve edge. + + Body ``conversation_params.conversation_id`` is canonical. Headers are used + only when the body does not provide an id. + """ + conversation_params = request.conversation_params + if conversation_params is not None: + return conversation_params.conversation_id + + conversation_id = extract_conversation_id_from_headers(headers) + if conversation_id is not None: + from tensorrt_llm.serve.openai_protocol import ConversationParams + + request.conversation_params = ConversationParams(conversation_id=conversation_id) + return conversation_id diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 2ac3e59282fc..051bb4a1c8cf 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -39,12 +39,13 @@ from tensorrt_llm.serve.cluster_storage import ( HttpClusterStorageServer, create_cluster_storage, validate_http_cluster_storage_scope) +from tensorrt_llm.serve.conversation_id import resolve_request_conversation_id from tensorrt_llm.serve.metadata_server import create_metadata_server from tensorrt_llm.serve.openai_client import OpenAIClient, OpenAIHttpClient from tensorrt_llm.serve.openai_disagg_service import ( OpenAIDisaggregatedService, ResponseHooks) from tensorrt_llm.serve.openai_protocol import ( - DisaggregatedParams, UCompletionRequest, UCompletionResponse, + UCompletionRequest, UCompletionResponse, ensure_request_chat_template_allowed) from tensorrt_llm.serve.perf_metrics import DisaggPerfMetricsCollector from tensorrt_llm.serve.responses_utils import (ServerArrivalTimeMiddleware, @@ -81,13 +82,6 @@ def on_resp_done(self, gen_server: str, request: UCompletionRequest, response: U class OpenAIDisaggServer: - _CONVERSATION_ID_HEADERS = ( - "x-session-id", - "x-correlation-id", - "x-session-affinity", - "x-multi-turn-session-id", - ) - def __init__(self, config: DisaggServerConfig, req_timeout_secs: int = 180, @@ -179,38 +173,12 @@ def register_routes(self): @staticmethod def _extract_conversation_id(req: UCompletionRequest, raw_req: Request): - """Populate conversation_id from supported session headers. - - When not already set in the request body, copies the header value - into ``disaggregated_params.conversation_id``. - - Supported headers are checked in priority order: ``X-Session-ID``, - ``X-Correlation-ID``, ``x-session-affinity``, and - ``x-multi-turn-session-id``. We mirror these conventions so the - ConversationRouter can provide session affinity without requiring - clients to set the body field. - - When ``disaggregated_params`` is ``None`` (standard OpenAI - requests without disagg fields), a minimal instance is created - to carry the conversation_id. The service layer always rebuilds - ``disaggregated_params`` in ``_get_ctx_request`` / - ``_get_gen_request`` before forwarding to workers. + """Populate conversation_params.conversation_id from supported headers. + + Body ``conversation_params.conversation_id`` is canonical. Headers are + used only when the body does not provide an id. """ - header_conv_id = None - for header_name in OpenAIDisaggServer._CONVERSATION_ID_HEADERS: - header_conv_id = raw_req.headers.get(header_name) - if header_conv_id is not None and header_conv_id.strip(): - break - else: - return - - if req.disaggregated_params is None: - req.disaggregated_params = DisaggregatedParams( - request_type="context_only", - conversation_id=header_conv_id, - ) - elif req.disaggregated_params.conversation_id is None: - req.disaggregated_params.conversation_id = header_conv_id + resolve_request_conversation_id(req, raw_req.headers) def _wrap_entry_point(self, entry_point: Callable) -> Callable: async def wrapper(req: UCompletionRequest, raw_req: Request) -> Response: diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index b04689f8682d..3e93a80e8392 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -157,13 +157,10 @@ async def _send_disagg_request_ctx_first( disagg_request_id = ctx_response_disagg_params.disagg_request_id gen_req = self._get_gen_request(request, ctx_response, disagg_request_id) else: - # Clear synthetic disaggregated_params that may have been - # injected by _extract_conversation_id (e.g. from the - # X-Correlation-ID header). When need_ctx=False the gen - # server handles full generation and must not see a stale - # request_type="context_only". + # When need_ctx=False the gen server handles full generation and + # must not see a stale request_type="context_only". # _check_gen_only_disagg already sets proper generation_only - # params when applicable, so only clear the synthetic ones. + # params when applicable. if ( gen_req.disaggregated_params is not None and gen_req.disaggregated_params.request_type == "context_only" @@ -320,6 +317,8 @@ def _need_gen(self, response: UCompletionResponse) -> bool: @staticmethod def _get_conversation_id(request: UCompletionRequest) -> Optional[str]: + if request.conversation_params is not None: + return request.conversation_params.conversation_id if request.disaggregated_params is not None: return request.disaggregated_params.conversation_id return None @@ -327,13 +326,14 @@ def _get_conversation_id(request: UCompletionRequest) -> Optional[str]: def _get_ctx_request( self, request: UCompletionRequest, disagg_request_id: Optional[int] ) -> UCompletionRequest: + conversation_id = self._get_conversation_id(request) ctx_request = request.model_copy( update={ "disaggregated_params": DisaggregatedParams( request_type="context_only", disagg_request_id=disagg_request_id, schedule_style=self._schedule_style, - conversation_id=self._get_conversation_id(request), + conversation_id=conversation_id, ), "stream": False, "stream_options": None, diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index e80c494e8ae5..6cffb280806f 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -38,6 +38,7 @@ from tensorrt_llm.executor.request import LoRARequest from tensorrt_llm.inputs.media_io import MediaModality +from tensorrt_llm.llmapi import ConversationParams as LlmConversationParams from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams from tensorrt_llm.llmapi import (DisaggScheduleStyle, GuidedDecodingParams, SamplingParams) @@ -172,6 +173,25 @@ class DisaggregatedParams(OpenAIBaseModel): # before routing MM encoder -> context -> generation through OpenAI protocol. +class ConversationParams(OpenAIBaseModel): + model_config = ConfigDict(extra="forbid", + populate_by_name=True, + validate_assignment=True) + + conversation_id: str = Field( + description=("Stable multi-turn conversation id used for routing"), ) + + @field_validator("conversation_id", mode="before") + @classmethod + def validate_conversation_id(cls, value: Any) -> str: + if value is None: + raise ValueError("conversation_id must be non-empty") + conversation_id = str(value).strip() + if not conversation_id: + raise ValueError("conversation_id must be non-empty") + return conversation_id + + class ErrorResponse(OpenAIBaseModel): object: str = "error" message: str @@ -429,6 +449,10 @@ class CompletionRequest(OpenAIBaseModel): default=None, description=("Parameters for disaggregated serving"), ) + conversation_params: Optional[ConversationParams] = Field( + default=None, + description=("Parameters for multi-turn conversation routing"), + ) # doc: end-completion-extra-params @@ -807,6 +831,10 @@ class ChatCompletionRequest(OpenAIBaseModel): default=None, description=("Parameters for disaggregated serving"), ) + conversation_params: Optional[ConversationParams] = Field( + default=None, + description=("Parameters for multi-turn conversation routing"), + ) cache_salt: Optional[str] = Field( default=None, @@ -1358,6 +1386,15 @@ def to_llm_disaggregated_params( ) +def to_llm_conversation_params( + conversation_params: Optional[ConversationParams] +) -> Optional[LlmConversationParams]: + if conversation_params is None: + return None + return LlmConversationParams( + conversation_id=conversation_params.conversation_id) + + # ============================================================================ # Diffusion API Protocol Classes # ============================================================================ diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 3e9c54cb6cba..eb45ba110f64 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -54,6 +54,7 @@ parse_chat_messages_coroutines, resolve_top_level_model_type) from tensorrt_llm.serve.cluster_storage import create_cluster_storage_client +from tensorrt_llm.serve.conversation_id import resolve_request_conversation_id from tensorrt_llm.serve.disagg_auto_scaling import DisaggClusterWorker from tensorrt_llm.serve.metadata_server import create_metadata_server from tensorrt_llm.serve.openai_protocol import ( @@ -63,7 +64,8 @@ ImageGenerationResponse, ImageObject, MemoryUpdateRequest, ModelCard, ModelList, PromptTokensDetails, ResponseFormat, ResponsesRequest, ResponsesResponse, UpdateWeightsRequest, UsageInfo, - ensure_request_chat_template_allowed, to_llm_disaggregated_params) + ensure_request_chat_template_allowed, to_llm_conversation_params, + to_llm_disaggregated_params) from tensorrt_llm.serve.openai_video_routes import _VideoRoutesMixin from tensorrt_llm.serve.postprocess_handlers import ( ChatCompletionPostprocArgs, ChatPostprocArgs, CompletionPostprocArgs, @@ -1294,6 +1296,10 @@ async def chat_stream_generator( trace_headers = (None if raw_request is None else tracing.extract_trace_headers(raw_request.headers)) + resolve_request_conversation_id( + request, None if raw_request is None else raw_request.headers) + conversation_params = to_llm_conversation_params( + request.conversation_params) scheduling_params = SchedulingParams( agent_hierarchy=request.agent_hierarchy) @@ -1312,6 +1318,7 @@ async def chat_stream_generator( streaming=request.stream, lora_request=request.lora_request, disaggregated_params=disaggregated_params, + conversation_params=conversation_params, cache_salt=request.cache_salt, trace_headers=trace_headers, scheduling_params=scheduling_params, @@ -1583,6 +1590,10 @@ async def generator_wrapper(generator: AsyncIterator[Any]): sampling_params.return_perf_metrics = True disaggregated_params = to_llm_disaggregated_params( request.disaggregated_params) + resolve_request_conversation_id( + request, None if raw_request is None else raw_request.headers) + conversation_params = to_llm_conversation_params( + request.conversation_params) for idx, prompt in enumerate(prompts): postproc_args = CompletionPostprocArgs.from_request(request) postproc_args.prompt_idx = idx @@ -1618,6 +1629,7 @@ async def generator_wrapper(generator: AsyncIterator[Any]): streaming=request.stream, lora_request=request.lora_request, disaggregated_params=disaggregated_params, + conversation_params=conversation_params, trace_headers=trace_headers) asyncio.create_task( self.await_disconnected(raw_request, promise)) @@ -1760,6 +1772,10 @@ async def create_streaming_generator(promise: RequestOutput, postproc_args=postproc_args, ) + resolve_request_conversation_id( + request, None if raw_request is None else raw_request.headers) + conversation_params = to_llm_conversation_params( + request.conversation_params) scheduling_params = SchedulingParams( agent_hierarchy=request.agent_hierarchy) @@ -1773,6 +1789,7 @@ async def create_streaming_generator(promise: RequestOutput, lora_request=request.lora_request, scheduling_params=scheduling_params, disaggregated_params=disaggregated_params, + conversation_params=conversation_params, trace_headers=trace_headers, ) if not self.postproc_worker_enabled: diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index c257d0eaa820..3490abb7294c 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -35,6 +35,7 @@ ReuseScope from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import \ RootBlock as V2RootBlock +from tensorrt_llm.serve.conversation_id import get_request_conversation_id from tensorrt_llm.serve.metadata_server import JsonDictionary from tensorrt_llm.serve.openai_protocol import (ChatCompletionRequest, CompletionRequest) @@ -1311,7 +1312,7 @@ class ConversationRouter(BlockHashMixin, LoadBalancingMixin, Router): """Router that provides session affinity for multi-turn conversations. Routing priority: - 1. Explicit ``conversation_id`` in ``disaggregated_params`` — sticky + 1. Explicit ``conversation_params.conversation_id`` — sticky routing to the previously assigned server. 2. Implicit block-hash prefix matching — find the session whose stored block hashes share the longest prefix with the new request. @@ -1559,9 +1560,7 @@ def _valid(conv_id: str) -> bool: # ── routing helpers ── def _get_conversation_id(self, request: OpenAIRequest) -> Optional[str]: - if request.disaggregated_params is not None: - return request.disaggregated_params.conversation_id - return None + return get_request_conversation_id(request) def _generate_implicit_id(self) -> str: self._implicit_id_counter += 1 diff --git a/tests/integration/defs/disaggregated/test_workers.py b/tests/integration/defs/disaggregated/test_workers.py index 7e6f0ad1bef7..35149391bf32 100644 --- a/tests/integration/defs/disaggregated/test_workers.py +++ b/tests/integration/defs/disaggregated/test_workers.py @@ -21,6 +21,7 @@ from tensorrt_llm import logger from tensorrt_llm.serve.openai_client import OpenAIHttpClient from tensorrt_llm.serve.openai_protocol import (CompletionRequest, + ConversationParams, DisaggregatedParams) from tensorrt_llm.serve.router import (ConversationRouter, KvCacheAwareRouter, KvCacheAwareServerState, ServerRole, @@ -786,70 +787,92 @@ async def _send_via_disagg(self, response.raise_for_status() return response_dict + @staticmethod + def _apply_conversation_id_transport(request: dict, conv_id: str, + use_body_params: bool) -> dict | None: + if use_body_params: + request["conversation_params"] = {"conversation_id": conv_id} + return None + return {"X-Correlation-ID": conv_id} + async def test_explicit_conversation_id(self): """Same conversation_id sticky-routes; different ids load-balance.""" try: async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout( total=self.req_timeout_secs)) as session: - # 1. Same conversation_id always routes to the same server - conv_id = "test-explicit-conv-id" - prompt = ("Hello, this is a test prompt for " - "conversation routing with explicit id") - first_server = None - for i in range(6): - req = CompletionRequest( - model=self.model_name, - prompt=prompt, - disaggregated_params=DisaggregatedParams( - request_type="context_only", - conversation_id=conv_id)) - server, _ = await self.ctx_router.get_next_server(req) - if first_server is None: - first_server = server - else: - assert server == first_server, ( - f"Round {i}: expected {first_server}, " - f"got {server}") - response = await self._send_via_disagg( - session, { + for transport_name, use_body_params in [ + ("header", False), + ("body", True), + ]: + # 1. Same conversation_id always routes to the same server + conv_id = f"test-explicit-conv-id-{transport_name}" + prompt = ("Hello, this is a test prompt for " + "conversation routing with explicit id") + first_server = None + for i in range(6): + req = CompletionRequest( + model=self.model_name, + prompt=prompt, + conversation_params=ConversationParams( + conversation_id=conv_id), + disaggregated_params=DisaggregatedParams( + request_type="context_only")) + server, _ = await self.ctx_router.get_next_server(req) + if first_server is None: + first_server = server + else: + assert server == first_server, ( + f"{transport_name} round {i}: expected " + f"{first_server}, got {server}") + request = { "model": self.model_name, "prompt": prompt, "max_tokens": 32, "ignore_eos": True, "temperature": 0.0, - }, {"X-Correlation-ID": conv_id}) - assert len(response["choices"]) > 0 - await self.ctx_router.finish_request(req) - prompt = prompt + response["choices"][0]["text"] - logger.info( - f"Sticky routing passed: all 6 rounds -> {first_server}") - - # 2. Different conversation_ids are load-balanced across - # servers (not all pinned to a single one). - servers_seen = set() - for i in range(len(self.ctx_servers) * 2): - cid = f"test-diff-conv-{i}" - req = CompletionRequest( - model=self.model_name, - prompt=f"Unique prompt number {i} for load balancing", - disaggregated_params=DisaggregatedParams( - request_type="context_only", conversation_id=cid)) - server, _ = await self.ctx_router.get_next_server(req) - servers_seen.add(server) - response = await self._send_via_disagg( - session, { + } + headers = self._apply_conversation_id_transport( + request, conv_id, use_body_params) + response = await self._send_via_disagg( + session, request, headers) + assert len(response["choices"]) > 0 + await self.ctx_router.finish_request(req) + prompt = prompt + response["choices"][0]["text"] + logger.info(f"{transport_name} sticky routing passed: " + f"all 6 rounds -> {first_server}") + + # 2. Different conversation_ids are load-balanced across + # servers (not all pinned to a single one). + servers_seen = set() + for i in range(len(self.ctx_servers) * 2): + cid = f"test-diff-conv-{transport_name}-{i}" + req = CompletionRequest( + model=self.model_name, + prompt= + f"Unique prompt number {i} for load balancing", + conversation_params=ConversationParams( + conversation_id=cid), + disaggregated_params=DisaggregatedParams( + request_type="context_only")) + server, _ = await self.ctx_router.get_next_server(req) + servers_seen.add(server) + request = { "model": self.model_name, "prompt": f"Unique prompt number {i}", "max_tokens": 1, "temperature": 0.0, - }, {"X-Correlation-ID": cid}) - assert len(response["choices"]) > 0 - await self.ctx_router.finish_request(req) - assert len(servers_seen) > 1, ( - f"Different conv_ids all routed to same server: " - f"{servers_seen}") - logger.info( - f"Load balancing passed: {len(servers_seen)} servers used") + } + headers = self._apply_conversation_id_transport( + request, cid, use_body_params) + response = await self._send_via_disagg( + session, request, headers) + assert len(response["choices"]) > 0 + await self.ctx_router.finish_request(req) + assert len(servers_seen) > 1, ( + f"{transport_name} different conv_ids all routed to " + f"same server: {servers_seen}") + logger.info(f"{transport_name} load balancing passed: " + f"{len(servers_seen)} servers used") finally: await self.ctx_router.close() diff --git a/tests/unittest/_torch/executor/test_adp_router.py b/tests/unittest/_torch/executor/test_adp_router.py index 7b1f58b1a100..4c3726578e0d 100644 --- a/tests/unittest/_torch/executor/test_adp_router.py +++ b/tests/unittest/_torch/executor/test_adp_router.py @@ -23,6 +23,7 @@ RankState, _num_input_tokens, ) +from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.scheduling_params import SchedulingParams @@ -59,6 +60,7 @@ def create_mock_request_with_py_schedule_params(attention_dp_rank=None, attentio mock_request.py_scheduling_params = mock_schedule_params else: mock_request.py_scheduling_params = None + mock_request.py_conversation_params = None mock_request.input_token_ids = [1, 2, 3] return mock_request @@ -73,6 +75,7 @@ def _make_request_item(req_id, num_tokens=10, target_dp_rank=None, attention_dp_ scheduling_params.attention_dp_relax = attention_dp_relax item.request = _MockRequest() item.request.py_scheduling_params = scheduling_params + item.request.py_conversation_params = None item.request.input_token_ids = list(range(num_tokens)) return item @@ -1107,7 +1110,7 @@ def _make_conv_request_item( attention_dp_relax=True, num_tokens=10, ): - """Mock RequestQueueItem carrying a conversation_id on disaggregated_params.""" + """Mock RequestQueueItem carrying conversation params.""" item = MagicMock() item.id = req_id item.child_req_ids = None @@ -1116,14 +1119,12 @@ def _make_conv_request_item( scheduling_params.attention_dp_relax = attention_dp_relax item.request = _MockRequest() item.request.py_scheduling_params = scheduling_params + item.request.py_conversation_params = None + if conversation_id is not None: + item.request.py_conversation_params = ConversationParams(conversation_id=conversation_id) item.request.input_token_ids = list(range(num_tokens)) item.request.py_orig_prompt_len = num_tokens - if conversation_id is None: - item.request.py_disaggregated_params = None - else: - disagg = MagicMock() - disagg.conversation_id = conversation_id - item.request.py_disaggregated_params = disagg + item.request.py_disaggregated_params = None return item diff --git a/tests/unittest/_torch/executor/test_request_utils.py b/tests/unittest/_torch/executor/test_request_utils.py index 899b1e8851ee..085a17e7ea62 100644 --- a/tests/unittest/_torch/executor/test_request_utils.py +++ b/tests/unittest/_torch/executor/test_request_utils.py @@ -6,12 +6,15 @@ """ +from types import SimpleNamespace from unittest.mock import Mock, patch import pytest from tensorrt_llm._torch.pyexecutor.executor_request_queue import RequestQueueItem from tensorrt_llm._torch.pyexecutor.request_utils import ( + RequestBroadcaster, + attach_py_objects_to_requests, can_process_attention_dp_request, derive_attention_dp_per_rank_request_cap, get_from_waiting_queue, @@ -20,6 +23,7 @@ ) from tensorrt_llm._torch.pyexecutor.scheduler import FCFSWaitingQueue from tensorrt_llm.bindings import executor as trtllm +from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.mapping import CpType @@ -58,6 +62,39 @@ def create_mock_request_with_py_schedule_params(attention_dp_rank=None, attentio return mock_request +def test_request_broadcaster_collects_conversation_params_with_none(): + conversation_params = ConversationParams(conversation_id="conv") + source_items = [ + RequestQueueItem(1, SimpleNamespace(py_conversation_params=conversation_params)), + RequestQueueItem(2, SimpleNamespace(py_conversation_params=None)), + ] + + py_request_objects = RequestBroadcaster._collect_py_objects(None, source_items) + + py_objects = dict(py_request_objects) + assert py_objects["py_conversation_params"] == { + 1: conversation_params, + 2: None, + } + + target_items = [ + RequestQueueItem(1, SimpleNamespace()), + RequestQueueItem(2, SimpleNamespace()), + ] + attach_py_objects_to_requests(target_items, py_request_objects) + + assert target_items[0].request.py_conversation_params is conversation_params + assert hasattr(target_items[1].request, "py_conversation_params") + assert target_items[1].request.py_conversation_params is None + + +def test_request_broadcaster_requires_conversation_params_attr(): + source_items = [RequestQueueItem(1, SimpleNamespace())] + + with pytest.raises(AttributeError): + RequestBroadcaster._collect_py_objects(None, source_items) + + def test_merge_helix_requests_with_padding(): """Test merge_helix_requests with basic valid input.""" diff --git a/tests/unittest/api_stability/references/conversation_params.yaml b/tests/unittest/api_stability/references/conversation_params.yaml new file mode 100644 index 000000000000..65959faa85b6 --- /dev/null +++ b/tests/unittest/api_stability/references/conversation_params.yaml @@ -0,0 +1,8 @@ +methods: + __init__: + parameters: + conversation_id: + annotation: str + default: inspect._empty + return_annotation: None +properties: {} diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 09f9968bd689..711ab7d898b7 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -319,6 +319,9 @@ methods: scheduling_params: annotation: Union[tensorrt_llm.scheduling_params.SchedulingParams, List[tensorrt_llm.scheduling_params.SchedulingParams], NoneType] default: null + conversation_params: + annotation: Union[tensorrt_llm.conversation_params.ConversationParams, List[tensorrt_llm.conversation_params.ConversationParams], NoneType] + default: null cache_salt: annotation: Union[str, Sequence[str], NoneType] default: null @@ -339,6 +342,9 @@ methods: annotation: Optional[tensorrt_llm.scheduling_params.SchedulingParams] default: null status: prototype + conversation_params: + annotation: Optional[tensorrt_llm.conversation_params.ConversationParams] + default: null cache_salt: annotation: Optional[str] default: null diff --git a/tests/unittest/api_stability/test_llm_api.py b/tests/unittest/api_stability/test_llm_api.py index 3edd14ecfabb..259eab8b62cc 100644 --- a/tests/unittest/api_stability/test_llm_api.py +++ b/tests/unittest/api_stability/test_llm_api.py @@ -9,8 +9,8 @@ from tensorrt_llm.bindings import executor as tllme from tensorrt_llm.executor.result import IterationResult from tensorrt_llm.llmapi import (CalibConfig, CompletionOutput, - GuidedDecodingParams, QuantConfig, - RequestOutput) + ConversationParams, GuidedDecodingParams, + QuantConfig, RequestOutput) from tensorrt_llm.sampling_params import (BatchedLogitsProcessor, LogitsProcessor, SamplingParams) @@ -75,6 +75,11 @@ class TestGuidedDecodingParams(ApiStabilityTestHarness): REFERENCE_FILE = "guided_decoding_params.yaml" +class TestConversationParams(ApiStabilityTestHarness): + TEST_CLASS = ConversationParams + REFERENCE_FILE = "conversation_params.yaml" + + class TestLogitsProcessor(ApiStabilityTestHarness): TEST_CLASS = LogitsProcessor REFERENCE_FILE = "logits_processor.yaml" diff --git a/tests/unittest/disaggregated/test_openai_disagg_server.py b/tests/unittest/disaggregated/test_openai_disagg_server.py index a9f270d9d7ff..beff804f9ed8 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_server.py +++ b/tests/unittest/disaggregated/test_openai_disagg_server.py @@ -18,7 +18,11 @@ from tensorrt_llm.llmapi.disagg_utils import extract_disagg_cfg from tensorrt_llm.serve.openai_disagg_server import OpenAIDisaggServer -from tensorrt_llm.serve.openai_protocol import CompletionRequest, DisaggregatedParams +from tensorrt_llm.serve.openai_protocol import ( + CompletionRequest, + ConversationParams, + DisaggregatedParams, +) def _raw_request(headers: dict[str, str]): @@ -61,8 +65,8 @@ def test_extract_conversation_id_from_headers(): OpenAIDisaggServer._extract_conversation_id(request, _raw_request(headers)) - assert request.disaggregated_params is not None - assert request.disaggregated_params.conversation_id == expected_conversation_id + assert request.disaggregated_params is None + assert request.conversation_params.conversation_id == expected_conversation_id def test_extract_conversation_id_ignores_empty_headers(): @@ -81,16 +85,15 @@ def test_extract_conversation_id_ignores_empty_headers(): ) assert request.disaggregated_params is None + assert request.conversation_params is None -def test_extract_conversation_id_preserves_body_conversation_id(): +def test_extract_conversation_id_preserves_body_conversation_params(): request = CompletionRequest( model="test-model", prompt="hello", - disaggregated_params=DisaggregatedParams( - request_type="context_only", - conversation_id="body-id", - ), + conversation_params=ConversationParams(conversation_id="body-id"), + disaggregated_params=DisaggregatedParams(request_type="context_only"), ) OpenAIDisaggServer._extract_conversation_id( @@ -98,10 +101,11 @@ def test_extract_conversation_id_preserves_body_conversation_id(): _raw_request({"X-Session-ID": "header-id"}), ) - assert request.disaggregated_params.conversation_id == "body-id" + assert request.conversation_params.conversation_id == "body-id" + assert request.disaggregated_params.conversation_id is None -def test_extract_conversation_id_populates_existing_disaggregated_params(): +def test_extract_conversation_id_populates_conversation_params_with_existing_disaggregated_params(): request = CompletionRequest( model="test-model", prompt="hello", @@ -113,7 +117,8 @@ def test_extract_conversation_id_populates_existing_disaggregated_params(): _raw_request({"x-multi-turn-session-id": "multi-turn-session-id"}), ) - assert request.disaggregated_params.conversation_id == "multi-turn-session-id" + assert request.conversation_params.conversation_id == "multi-turn-session-id" + assert request.disaggregated_params.conversation_id is None def test_disagg_config_allows_request_chat_template_opt_in(): diff --git a/tests/unittest/disaggregated/test_openai_disagg_service.py b/tests/unittest/disaggregated/test_openai_disagg_service.py index 4d7d813304f9..7cc7fe4a649e 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_service.py +++ b/tests/unittest/disaggregated/test_openai_disagg_service.py @@ -38,6 +38,7 @@ CompletionRequest, CompletionResponse, CompletionResponseChoice, + ConversationParams, DisaggregatedParams, DisaggScheduleStyle, PromptTokensDetails, @@ -149,6 +150,7 @@ def test_get_gen_request_uses_ctx_response_prompt_token_ids_for_chat(): } ], prompt_token_ids=[1, 2, 3], + conversation_params=ConversationParams(conversation_id="conv-chat"), ) ctx_response = _make_chat_response( finish_reason="length", @@ -159,6 +161,7 @@ def test_get_gen_request_uses_ctx_response_prompt_token_ids_for_chat(): assert gen_request.prompt_token_ids == ctx_prompt_token_ids assert gen_request.disaggregated_params.request_type == "generation_only" + assert gen_request.conversation_params.conversation_id == "conv-chat" @pytest.mark.asyncio diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index 859e3c88d747..b8d5980e087b 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -13,12 +13,13 @@ truncate_sha256_hash_to_int64) from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import ( ReuseScope, sequence_to_blockchain_keys) +# yapf: disable from tensorrt_llm.serve.openai_protocol import (ChatCompletionRequest, ChatCompletionToolsParam, CompletionRequest, + ConversationParams, DisaggregatedParams, FunctionDefinition) -# yapf: disable from tensorrt_llm.serve.router import (KV_CACHE_HASH_ALGO_V1, KV_CACHE_HASH_ALGO_V2, KV_CACHE_HASH_ALGO_V2_SHA256_64, @@ -1532,11 +1533,15 @@ async def test_get_next_server_exclude_server_insufficient(router_class): def _make_request(conversation_id=None, prompt="the " * 100): - params = DisaggregatedParams(request_type="context_only", - conversation_id=conversation_id) + params = DisaggregatedParams(request_type="context_only") + conversation_params = None + if conversation_id is not None: + conversation_params = ConversationParams( + conversation_id=conversation_id) return CompletionRequest(model="TinyLlama", prompt=[prompt], - disaggregated_params=params) + disaggregated_params=params, + conversation_params=conversation_params) @pytest.mark.asyncio diff --git a/tests/unittest/llmapi/apps/test_openai_conversation_params.py b/tests/unittest/llmapi/apps/test_openai_conversation_params.py new file mode 100644 index 000000000000..b262c880245c --- /dev/null +++ b/tests/unittest/llmapi/apps/test_openai_conversation_params.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace + +import pytest +from pydantic import ValidationError + +from tensorrt_llm.serve.conversation_id import resolve_request_conversation_id +from tensorrt_llm.serve.openai_protocol import ConversationParams + + +def test_openai_conversation_params_requires_conversation_id(): + with pytest.raises(ValidationError): + ConversationParams() + + +def test_openai_conversation_params_normalizes_conversation_id(): + params = ConversationParams(conversation_id=" body-id ") + + assert params.conversation_id == "body-id" + + params.conversation_id = " next-id " + + assert params.conversation_id == "next-id" + + +def test_openai_conversation_params_rejects_empty_conversation_id(): + with pytest.raises(ValidationError): + ConversationParams(conversation_id=" ") + + +def test_resolve_request_conversation_id_preserves_body_params(): + conversation_params = ConversationParams(conversation_id="body-id") + request = SimpleNamespace(conversation_params=conversation_params) + + conversation_id = resolve_request_conversation_id(request, {"X-Session-ID": "header-id"}) + + assert conversation_id == "body-id" + assert request.conversation_params is conversation_params + + +def test_resolve_request_conversation_id_uses_headers_when_body_params_missing(): + request = SimpleNamespace(conversation_params=None) + + conversation_id = resolve_request_conversation_id(request, {"X-Session-ID": " header-id "}) + + assert conversation_id == "header-id" + assert request.conversation_params.conversation_id == "header-id" diff --git a/tests/unittest/llmapi/test_conversation_params.py b/tests/unittest/llmapi/test_conversation_params.py new file mode 100644 index 000000000000..76cdea35212e --- /dev/null +++ b/tests/unittest/llmapi/test_conversation_params.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from tensorrt_llm.conversation_params import ConversationParams + + +def test_conversation_params_normalizes_conversation_id(): + params = ConversationParams(conversation_id=" body-id ") + + assert params.conversation_id == "body-id" + + +def test_conversation_params_rejects_empty_conversation_id(): + with pytest.raises(ValueError): + ConversationParams(conversation_id=" ") From fddf2ca3ba8372d2a596cfafcf50f05ec8a42772 Mon Sep 17 00:00:00 2001 From: Fanrong Li Date: Fri, 26 Jun 2026 03:30:45 +0000 Subject: [PATCH 24/28] [None][fix] pass cache salt to ADP prefix probe Signed-off-by: Fanrong Li --- .../_torch/pyexecutor/resource_manager.py | 12 ++-- .../_torch/pyexecutor/scheduler/adp_router.py | 8 +-- .../executor/test_kvcache_aware_router.py | 62 ++++++++++++++++--- 3 files changed, 64 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 220433d6bc16..cb048c2d9aaf 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -646,13 +646,13 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], def probe_prefix_match_length(self, input_tokens, lora_task_id=None, - cache_salt_id=None): + cache_salt=None): """Probe the KV cache radix tree for prefix match length. Returns the number of prefix tokens already cached on this rank. Used by KVCacheAwareADPRouter for cache-aware routing. - ``cache_salt_id`` scopes the probe to the same per-request + ``cache_salt`` scopes the probe to the same per-request namespace that ``_create_kv_cache`` uses; without it, salted requests would be probed against the salt=None namespace and the router would see the wrong match length. @@ -672,16 +672,16 @@ def probe_prefix_match_length(self, LlmRequest as CppLlmRequest block_key = BlockKey(tokens=input_tokens, lora_task_id=lora_task_id) unique_tokens = block_key.unique_tokens - # buildBlockKeys() on the C++ side pulls cache_salt_id off the - # request, so dummy_req must carry it for the probe namespace to - # match the real create_kv_cache path. + # buildBlockKeys() on the C++ side derives the salt id from + # cache_salt, so the dummy request must carry the same string used by + # the real create_kv_cache path. dummy_req = CppLlmRequest(request_id=0, max_new_tokens=0, input_tokens=input_tokens, sampling_config=SamplingConfig(), is_streaming=False, lora_task_id=lora_task_id, - cache_salt_id=cache_salt_id) + cache_salt=cache_salt) summary = self.impl.analyze_prefix_reuse(unique_tokens, dummy_req) return summary.reusable_blocks_all * self.tokens_per_block diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index 31bdf4947154..acea0242edcb 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -584,13 +584,13 @@ def gather_prefix_matches( probe_tokens = input_tokens[:-1] if len(input_tokens) > 1 else [] lora_config = getattr(req, "lora_config", None) lora_task_id = lora_config.task_id if lora_config is not None else None - # cache_salt_id scopes block reuse on both v1 and v2 backends; - # passing None for non-salted requests is a no-op. - cache_salt_id = getattr(req, "cache_salt_id", None) + # cache_salt scopes block reuse on the backend; passing None for + # non-salted requests is a no-op. + cache_salt = getattr(req, "cache_salt", None) match_len = self.kv_cache_manager.probe_prefix_match_length( probe_tokens, lora_task_id, - cache_salt_id=cache_salt_id, + cache_salt=cache_salt, ) local_matches.extend([req_item.id, match_len]) diff --git a/tests/unittest/_torch/executor/test_kvcache_aware_router.py b/tests/unittest/_torch/executor/test_kvcache_aware_router.py index 172bcf65e817..900151c7ac1d 100644 --- a/tests/unittest/_torch/executor/test_kvcache_aware_router.py +++ b/tests/unittest/_torch/executor/test_kvcache_aware_router.py @@ -41,7 +41,12 @@ def _mock_dist(tp_rank=0, tp_size=1, has_cp_helix=False): def _make_request_item( - req_id, num_tokens=10, target_dp_rank=None, attention_dp_relax=True, lora_task_id=None + req_id, + num_tokens=10, + target_dp_rank=None, + attention_dp_relax=True, + lora_task_id=None, + cache_salt=None, ): """Create a mock RequestQueueItem for testing.""" item = MagicMock() @@ -53,6 +58,7 @@ def _make_request_item( item.request = MagicMock() item.request.py_scheduling_params = scheduling_params item.request.input_token_ids = list(range(num_tokens)) + item.request.cache_salt = cache_salt if lora_task_id is not None: lora_config = MagicMock() lora_config.task_id = lora_task_id @@ -66,14 +72,15 @@ def _mock_kv_cache_manager(probe_results=None): """Create mock KV cache manager with configurable probe results. Args: - probe_results: dict mapping (tuple(input_tokens), lora_task_id) -> match_length. + probe_results: dict mapping (tuple(input_tokens), lora_task_id, cache_salt) + -> match_length. If None, all probes return 0. """ mgr = MagicMock() probe_results = probe_results or {} - def mock_probe(input_tokens, lora_task_id=None): - key = (tuple(input_tokens), lora_task_id) + def mock_probe(input_tokens, lora_task_id=None, cache_salt=None): + key = (tuple(input_tokens), lora_task_id, cache_salt) return probe_results.get(key, 0) mgr.probe_prefix_match_length = Mock(side_effect=mock_probe) @@ -120,8 +127,8 @@ def test_gather_prefix_matches_single_rank(self): tokens_a = list(range(100)) tokens_b = list(range(50)) probe_results = { - (tuple(tokens_a[:-1]), None): 64, - (tuple(tokens_b[:-1]), None): 0, + (tuple(tokens_a[:-1]), None, None): 64, + (tuple(tokens_b[:-1]), None, None): 0, } dist = _mock_dist(tp_rank=0, tp_size=1) @@ -139,7 +146,7 @@ def test_gather_prefix_matches_single_rank(self): def test_gather_prefix_matches_two_ranks(self): tokens_a = list(range(100)) probe_results = { - (tuple(tokens_a[:-1]), None): 64, + (tuple(tokens_a[:-1]), None, None): 64, } dist = _mock_dist(tp_rank=0, tp_size=2) @@ -159,7 +166,7 @@ def test_gather_prefix_matches_two_ranks(self): def test_gather_prefix_matches_with_lora(self): tokens_a = list(range(100)) probe_results = { - (tuple(tokens_a[:-1]), 42): 64, + (tuple(tokens_a[:-1]), 42, None): 64, } dist = _mock_dist(tp_rank=0, tp_size=1) @@ -171,6 +178,24 @@ def test_gather_prefix_matches_with_lora(self): router.gather_prefix_matches([req_a]) assert router._all_ranks_prefix_matches[0] == {1: 64} + def test_gather_prefix_matches_with_cache_salt(self): + tokens_a = list(range(100)) + probe_results = { + (tuple(tokens_a[:-1]), None, "tenant-a"): 64, + } + + dist = _mock_dist(tp_rank=0, tp_size=1) + dist.tp_allgather = Mock(side_effect=lambda x: [x]) + mgr = _mock_kv_cache_manager(probe_results) + router = KVCacheAwareADPRouter(dist=dist, kv_cache_manager=mgr) + + req_a = _make_request_item(1, num_tokens=100, cache_salt="tenant-a") + router.gather_prefix_matches([req_a]) + assert router._all_ranks_prefix_matches[0] == {1: 64} + mgr.probe_prefix_match_length.assert_called_once_with( + tokens_a[:-1], None, cache_salt="tenant-a" + ) + def test_gather_prefix_matches_empty(self): dist = _mock_dist(tp_rank=0, tp_size=1) dist.tp_allgather = Mock(side_effect=lambda x: [x]) @@ -566,3 +591,24 @@ def test_block_to_token_conversion(self): result = KVCacheManager.probe_prefix_match_length(mgr, input_tokens=list(range(200))) assert result == 192 # 3 blocks * 64 tokens/block mgr.impl.analyze_prefix_reuse.assert_called_once() + + def test_cache_salt_uses_cpp_request_cache_salt_kwarg(self): + """Salted probes must construct the C++ request with cache_salt.""" + from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager + + mgr = Mock(spec=KVCacheManager) + mgr.enable_block_reuse = True + mgr.impl = Mock() + mgr.impl.is_variable_window = False + mock_summary = Mock() + mock_summary.reusable_blocks_all = 1 + mgr.impl.analyze_prefix_reuse = Mock(return_value=mock_summary) + mgr.tokens_per_block = 64 + + result = KVCacheManager.probe_prefix_match_length( + mgr, input_tokens=list(range(128)), cache_salt="tenant-a" + ) + + assert result == 64 + dummy_req = mgr.impl.analyze_prefix_reuse.call_args.args[1] + assert dummy_req.cache_salt == "tenant-a" From 2d85deb4325d5bed8c83324e14160cbbadc15203 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Fri, 26 Jun 2026 00:55:10 -0700 Subject: [PATCH 25/28] [None][fix] kv_cache_manager_v2: restore probe_prefix_match_length on V2 backend KVCacheAwareADPRouter calls self.kv_cache_manager.probe_prefix_match_length() unconditionally, but KVCacheManagerV2 only inherits BaseResourceManager and lacked the method (it lives on the v1 KVCacheManager only, after KVCacheManagerV2 was moved out of resource_manager.py on main). KV-cache-aware ADP routing therefore raised AttributeError on the V2 KV cache backend. Add the V2 adapter, deriving the ReuseScope salt from cache_salt exactly as _create_kv_cache does so the probe queries the same reuse namespace and the router sees a correct match length. Extract the cache_salt -> salt-int derivation into a shared _derive_reuse_salt helper so cache creation, prefetch, and the router probe stay in lockstep. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 56 +++++++++++++------ 1 file changed, 40 insertions(+), 16 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index caf80ecba7b5..7d2a2d6e47b6 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -2515,6 +2515,20 @@ def copy_batch_block_offsets( self._stream.cuda_stream, ) + @staticmethod + def _derive_reuse_salt(cache_salt: str | None) -> int | None: + """Derive ``ReuseScope.salt`` (int|None) from the ``cache_salt`` string. + + Deterministic so the same string yields the same reuse namespace across + processes (matches C++ blockKey hashing on cacheSalt). Shared by cache + creation, prefetch, and the cache-aware router probe so they all hit the + same radix-tree namespace. + """ + if cache_salt is None: + return None + digest = hashlib.sha256(cache_salt.encode("utf-8")).digest() + return int.from_bytes(digest[:8], "little") + def _create_kv_cache( self, request_id: int, @@ -2538,14 +2552,7 @@ def _create_kv_cache( self.index_mapper.size(), ) return None - # ReuseScope.salt is int|None; derive a deterministic int from the - # cache_salt string so the same string yields the same reuse namespace - # across processes (matches C++ blockKey hashing on cacheSalt). - salt_int = ( - int.from_bytes(hashlib.sha256(cache_salt.encode("utf-8")).digest()[:8], "little") - if cache_salt is not None - else None - ) + salt_int = self._derive_reuse_salt(cache_salt) kv_cache = self.impl.create_kv_cache( ReuseScope(lora_id=lora_task_id, salt=salt_int), input_tokens, @@ -2565,6 +2572,28 @@ def _create_kv_cache( kv_cache.set_base_page_index_buf(i, pool_idx, memoryview(buffer.numpy())) return kv_cache + def probe_prefix_match_length(self, input_tokens, lora_task_id=None, cache_salt=None): + """Probe the v2 KV cache radix tree for prefix match length. + + Returns the number of prefix tokens already cached on this rank, + without acquiring page ownership. Mirrors the v1 ``KVCacheManager`` + adapter so ``KVCacheAwareADPRouter`` works on both KV cache backends. + + ``cache_salt`` (and ``lora_task_id``) must match the values used by + ``_create_kv_cache``; the salt is derived from the ``cache_salt`` + string the same way, so the probe queries the same reuse namespace. + Otherwise the router would see an incorrect match length. + """ + if not self.enable_block_reuse: + return 0 + if not input_tokens: + return 0 + salt_int = self._derive_reuse_salt(cache_salt) + return self.impl.probe_reuse( + ReuseScope(lora_id=lora_task_id, salt=salt_int), + input_tokens, + ) + def prefetch_for_context_tokens(self, requests: list) -> bool: """Prefetch radix-tree blocks from disk→host for upcoming context requests. @@ -2580,14 +2609,9 @@ def prefetch_for_context_tokens(self, requests: list) -> bool: for req in requests: all_tokens = req.get_tokens(DEFAULT_BEAM_INDEX) tokens = self._augment_tokens_for_block_reuse(all_tokens, req, end=len(all_tokens) - 1) - # Match the ReuseScope salt derivation used in _create_kv_cache so - # the transient cache hits the same radix-tree blocks. - cache_salt = req.cache_salt - salt_int = ( - int.from_bytes(hashlib.sha256(cache_salt.encode("utf-8")).digest()[:8], "little") - if cache_salt is not None - else None - ) + # Use the same salt derivation as _create_kv_cache so the transient + # cache hits the same radix-tree blocks. + salt_int = self._derive_reuse_salt(req.cache_salt) kv_cache = self.impl.create_kv_cache( ReuseScope(lora_id=req.lora_task_id, salt=salt_int), tokens ) From 23c372bd08f62cf35ed17cd84bfa6dc88ca4d9c2 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Fri, 26 Jun 2026 05:57:31 -0700 Subject: [PATCH 26/28] [None][fix] disagg: remove dead usage-rewrite machinery folded into #15042 #15042 ("skip per-chunk JSON parse on streaming usage rewrite") was a 5-line optimization on feat/deepseek_v4, layered on the method-based usage-rewrite machinery introduced by #13620. When this PR carved #15042 onto main, that machinery did not exist there: main's #14177 ("Fix disaggregated usage propagation") had refactored #13620 into postprocess_handlers (rewrite_usage_info_from_ctx, operating on parsed chunk objects). The carve folded the obsolete pre-refactor machinery back in to make the patch apply. The result is dead code: _rewrite_disagg_usage (the chain root) has no caller, so the whole _ctx_usage_for_client / _rewrite_usage_* / _rewrite_streaming_usage_from_ctx subtree is unreachable, and #15042's json.loads-skip optimization sits inside it (a no-op). main already rewrites disagg usage via postprocess_handlers and relays streaming bytes untouched, so the optimized per-chunk-parse path no longer exists. Remove the dead methods and the now-orphaned imports (json, AsyncIterator, PromptTokensDetails). No behavior change: the live usage path is unaffected. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- tensorrt_llm/serve/openai_disagg_service.py | 130 -------------------- 1 file changed, 130 deletions(-) diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index 3e93a80e8392..2123de29629f 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -13,9 +13,7 @@ # limitations under the License. import asyncio -import json import os -from collections.abc import AsyncIterator from typing import Any, Callable, Dict, Optional from tensorrt_llm.llmapi.disagg_utils import ( @@ -36,7 +34,6 @@ CompletionRequest, DisaggregatedParams, DisaggScheduleStyle, - PromptTokensDetails, UCompletionRequest, UCompletionResponse, ) @@ -182,133 +179,6 @@ async def _send_disagg_request_ctx_first( return done_generator() return ctx_response - def _ctx_usage_for_client( - self, ctx_response: Optional[UCompletionResponse] - ) -> tuple[Optional[int], int]: - if ctx_response is None or ctx_response.usage is None: - return None, 0 - - prompt_tokens = ctx_response.usage.prompt_tokens - cached_tokens = 0 - prompt_tokens_details = ctx_response.usage.prompt_tokens_details - if prompt_tokens_details is not None: - cached_tokens = prompt_tokens_details.cached_tokens - return prompt_tokens, cached_tokens - - def _rewrite_usage_payload_from_ctx( - self, - usage: dict[str, Any], - ctx_response: Optional[UCompletionResponse], - ) -> None: - prompt_tokens, cached_tokens = self._ctx_usage_for_client(ctx_response) - if prompt_tokens is None: - return - - usage["prompt_tokens"] = prompt_tokens - usage["total_tokens"] = prompt_tokens + (usage.get("completion_tokens") or 0) - prompt_tokens_details = usage.get("prompt_tokens_details") - if not isinstance(prompt_tokens_details, dict): - prompt_tokens_details = {} - usage["prompt_tokens_details"] = prompt_tokens_details - prompt_tokens_details["cached_tokens"] = cached_tokens - - def _rewrite_usage_response_from_ctx( - self, - response: UCompletionResponse, - ctx_response: Optional[UCompletionResponse], - ) -> UCompletionResponse: - prompt_tokens, cached_tokens = self._ctx_usage_for_client(ctx_response) - if prompt_tokens is None or response.usage is None: - return response - - response.usage.prompt_tokens = prompt_tokens - response.usage.total_tokens = prompt_tokens + (response.usage.completion_tokens or 0) - response.usage.prompt_tokens_details = PromptTokensDetails(cached_tokens=cached_tokens) - return response - - @staticmethod - def _sse_separator_index(data: bytes) -> tuple[int, bytes] | None: - indexes = [(data.find(sep), sep) for sep in (b"\n\n", b"\r\n\r\n")] - indexes = [(idx, sep) for idx, sep in indexes if idx >= 0] - if not indexes: - return None - return min(indexes, key=lambda item: item[0]) - - def _rewrite_usage_sse_event_from_ctx( - self, - event: bytes, - ctx_response: Optional[UCompletionResponse], - ) -> bytes: - # Only the final chunk has a top-level `usage` to rewrite, but this ran - # json.loads on every relayed chunk. A usage field can't exist unless the - # bytes contain `"usage"`, so skip the parse when they don't. - if b'"usage"' not in event: - return event - separator_match = self._sse_separator_index(event) - separator = separator_match[1] if separator_match else b"" - event_body = event[: -len(separator)] if separator else event - data_lines = [ - line.removeprefix(b"data:").strip() - for line in event_body.splitlines() - if line.startswith(b"data:") - ] - if len(data_lines) != 1 or data_lines[0] == b"[DONE]": - return event - - try: - payload = json.loads(data_lines[0]) - except json.JSONDecodeError: - return event - - usage = payload.get("usage") if isinstance(payload, dict) else None - if not isinstance(usage, dict): - return event - - self._rewrite_usage_payload_from_ctx(usage, ctx_response) - return b"data: " + json.dumps(payload, separators=(",", ":")).encode("utf-8") + separator - - async def _rewrite_streaming_usage_from_ctx( - self, - response: AsyncIterator[Any], - ctx_response: Optional[UCompletionResponse], - ) -> AsyncIterator[Any]: - pending = b"" - pending_is_str = False - async for chunk in response: - is_str = isinstance(chunk, str) - chunk_bytes = chunk.encode("utf-8") if is_str else chunk - if not isinstance(chunk_bytes, bytes): - yield chunk - continue - - if not pending: - pending_is_str = is_str - pending += chunk_bytes - while separator_match := self._sse_separator_index(pending): - separator_index, separator = separator_match - event_end = separator_index + len(separator) - event = pending[:event_end] - pending = pending[event_end:] - event = self._rewrite_usage_sse_event_from_ctx(event, ctx_response) - yield event.decode("utf-8") if pending_is_str else event - if not pending: - pending_is_str = False - - if pending: - event = self._rewrite_usage_sse_event_from_ctx(pending, ctx_response) - yield event.decode("utf-8") if pending_is_str else event - - def _rewrite_disagg_usage( - self, - response: UCompletionResponseOrGenerator, - ctx_response: Optional[UCompletionResponse], - ) -> UCompletionResponseOrGenerator: - if ctx_response is None: - return response - if hasattr(response, "__aiter__"): - return self._rewrite_streaming_usage_from_ctx(response, ctx_response) - return self._rewrite_usage_response_from_ctx(response, ctx_response) - def _need_gen(self, response: UCompletionResponse) -> bool: if response and response.choices[0].finish_reason not in _GEN_PENDING_FINISH_REASONS: del response.choices[0].disaggregated_params From 75c3eda565781b8a2a118ee092b0ae9bc1a23724 Mon Sep 17 00:00:00 2001 From: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com> Date: Mon, 29 Jun 2026 01:22:19 -0700 Subject: [PATCH 27/28] [None][fix] disagg: MTP AuxBuffer max_draft_len + conditional_v2 perf-metrics test KVCacheManagerV2 never exposed max_draft_len (only the V1 KVCacheManager did), so the native (Python NIXL) disagg transceiver's _make_aux_buffer() fell back to getattr(kvm, "max_draft_len", 0) and sized the AuxBuffer draft-token region for 0 tokens. Under MTP / speculative decoding a disagg worker filling a slot with >0 draft tokens then raised "draft_tokens length (N) exceeds max_draft_len (0)", terminating the executor loop and hanging the run on KV-transfer timeouts (e.g. accuracy/test_disaggregated_serving.py TestDeepSeekV3Lite::test_gen_first[adp-mtp2]). Mirror the V1 manager and set self.max_draft_len from the speculative config. Also fix the new conditional_v2 disagg test, whose /perf_metrics-based bypass check could never pass: enable per-request perf metrics in the config (return_perf_metrics / perf_metrics_max_requests), and verify once after all client iterations via post_client_test instead of racing each iteration -- routed-request metrics are recorded asynchronously (add_per_request_metrics via create_task) and /perf_metrics is consume-on-read. Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 3 ++ ...agg_config_conditional_deepseek_v3_v2.yaml | 7 +++ .../defs/disaggregated/test_disaggregated.py | 46 +++++++++++-------- 3 files changed, 38 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 7d2a2d6e47b6..3ec74e1fd72b 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -505,6 +505,9 @@ def __init__( from ..speculative import get_num_extra_kv_tokens self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config) + # Mirror V1: expose max_draft_len so the native disagg AuxBuffer + # (_make_aux_buffer's getattr fallback) is sized for MTP/spec decoding. + self.max_draft_len = spec_config.max_draft_len if spec_config is not None else 0 self.max_total_draft_tokens = ( spec_config.max_total_draft_tokens if spec_config is not None else 0 ) diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3_v2.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3_v2.yaml index 0be9c8201d3f..d33e1541406e 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3_v2.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3_v2.yaml @@ -7,10 +7,15 @@ conditional_disagg_config: max_local_prefill_length: 100 disable_overlap_scheduler: true enable_autotuner: false +# Enable per-request perf metrics so the disagg /perf_metrics endpoint (used by +# the test's _check_bypass_signal) records routed requests; bypassed ones absent. +perf_metrics_max_requests: 1000 context_servers: num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 + return_perf_metrics: true + perf_metrics_max_requests: 1000 kv_cache_config: enable_block_reuse: true enable_partial_reuse: true @@ -26,6 +31,8 @@ generation_servers: num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 + return_perf_metrics: true + perf_metrics_max_requests: 1000 router: type: kv_cache_aware # disagg-gen path doesn't emit kv_cache_events; track router-side state from diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index 65a98b09a94b..ad53792c7aba 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -1740,29 +1740,39 @@ def test_disaggregated_deepseek_v3_lite_bf16_conditional_v2( setup_model_symlink(llm_venv, deepseek_v3_model_root, "DeepSeek-V3-Lite/bf16") - # Bypassed requests clear disaggregated_params and skip add_per_request_metrics, - # so /perf_metrics records strictly fewer entries than requests issued. - last_total = [0] - - def _check_bypass_signal(server_url: str): + # Conditional disagg handles short-prefill requests locally on the gen + # server (bypassing the ctx handoff + add_per_request_metrics), while routed + # requests are recorded in the disagg /perf_metrics. Verify ONCE after all + # client iterations via post_client_test (not per-iteration): routed-request + # metrics are recorded asynchronously (add_per_request_metrics via + # create_task on response completion) and surface only after the client + # traffic settles, so a per-iteration read races that lag; /perf_metrics is + # also consume-on-read, so query it exactly once at the end. + def _check_routed_recorded(server_url: str): import requests as http_requests - resp = http_requests.get(f"{server_url}/perf_metrics", timeout=10) - assert resp.status_code == 200, \ - f"perf_metrics fetch failed: {resp.status_code}" - metrics = resp.json() - logger.info( - f"conditional_v2 perf_metrics len={len(metrics)} " - f"(bypassed requests absent — grows sub-linearly if bypass works)") - # At least one iteration's worth of requests should have been routed — - # i.e. perf_metrics is non-empty after the first client run. - assert metrics or last_total[0] > 0, \ - "no per-request metrics recorded; bypass logic may be misconfigured" - last_total[0] = len(metrics) + metrics = [] + deadline = time.time() + 60 + while True: + resp = http_requests.get(f"{server_url}/perf_metrics", timeout=10) + assert resp.status_code == 200, \ + f"perf_metrics fetch failed: {resp.status_code}" + metrics = resp.json() + if metrics or time.time() >= deadline: + break + time.sleep(2) + logger.info(f"conditional_v2 perf_metrics len={len(metrics)} " + f"(routed requests recorded; bypassed ones absent)") + # With short prompts every prompt's first occurrence routes through the + # context server (match=0 -> need_ctx), so at least one routed request + # must be recorded; an empty result means conditional routing never + # engaged. + assert metrics, \ + "no per-request metrics recorded after client runs; conditional routing may be misconfigured" run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_bf16_conditional_v2", env=llm_venv._new_env, - extra_endpoints_test=_check_bypass_signal, + post_client_test=_check_routed_recorded, model_path=deepseek_v3_model_root, cwd=llm_venv.get_working_directory()) From d017cdea29db9476b489e872aea84c6f6f364353 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 30 Jun 2026 23:10:03 -0700 Subject: [PATCH 28/28] [None][fix] api_stability: track conversation_params on serve request models PR #15625 added the `conversation_params: Optional[ConversationParams]` field to CompletionRequest and ChatCompletionRequest in tensorrt_llm/serve/openai_protocol.py but did not register it in the serve API stability reference. This tripped the test_serve_api.py::test_live_fields_are_in_yaml gate. Add the matching rows (kind: extension, status: prototype, default: null, required: false) so the untracked-field check passes. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../api_stability/references/trtllm_serve_api.yaml | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/unittest/api_stability/references/trtllm_serve_api.yaml b/tests/unittest/api_stability/references/trtllm_serve_api.yaml index f5663391d811..71ce2bb4dde3 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_api.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_api.yaml @@ -295,6 +295,12 @@ models: default: null status: stable required: false + conversation_params: + kind: extension + type: Optional[ConversationParams] + default: null + status: prototype + required: false CompletionResponseChoice: fields: @@ -766,6 +772,12 @@ models: default: null status: stable required: false + conversation_params: + kind: extension + type: Optional[ConversationParams] + default: null + status: prototype + required: false cache_salt: kind: extension type: Optional[str]