diff --git a/tensorrt_llm/disaggregated_params.py b/tensorrt_llm/disaggregated_params.py index 5c91dae3f5de..1c6970eac175 100644 --- a/tensorrt_llm/disaggregated_params.py +++ b/tensorrt_llm/disaggregated_params.py @@ -1,3 +1,18 @@ +# 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 from enum import IntEnum from typing import Any, Dict, List, Optional @@ -31,7 +46,6 @@ class DisaggregatedParams: Each entry is a torch.Tensor of shape [num_tokens, vocab_size] (one per beam/sequence). ctx_usage (Dict[str, Any]): The context usage payload to preserve exact usage accounting on the generation server. - multimodal_embedding_handles (List[Dict[str, Any]]): The resulting multimodal embedding handles from ViT. multimodal_hashes (List[List[int]]): The multimodal hashes of each multimodal item in the request. """ @@ -50,10 +64,6 @@ 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/serve/conversation_id.py b/tensorrt_llm/serve/conversation_id.py index 2acfcad8012e..c4146d93e5d8 100644 --- a/tensorrt_llm/serve/conversation_id.py +++ b/tensorrt_llm/serve/conversation_id.py @@ -36,10 +36,7 @@ class RequestWithConversationParams(Protocol): def get_request_conversation_id(request: RequestWithConversationParams) -> Optional[str]: conversation_params = request.conversation_params - if conversation_params is not None: - return conversation_params.conversation_id - disaggregated_params = getattr(request, "disaggregated_params", None) - return None if disaggregated_params is None else disaggregated_params.conversation_id + return None if conversation_params is None else conversation_params.conversation_id def extract_conversation_id_from_headers(headers: Optional[Mapping[str, str]]) -> Optional[str]: @@ -62,8 +59,7 @@ def resolve_request_conversation_id( ) -> 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. + Body ``conversation_params.conversation_id`` takes precedence over headers. """ conversation_params = request.conversation_params if conversation_params is not None: diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index 3d7b9e608ec5..f1ea3c45c5c8 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -192,25 +192,15 @@ def _need_gen(self, response: UCompletionResponse) -> bool: return False return True - @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 - 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=conversation_id, return_prompt_token_ids_b64=self._tokids_ctxbytes, ), "stream": False, @@ -226,12 +216,10 @@ def _get_gen_request( disagg_request_id: Optional[int], ctx_server_info: Optional[dict] = None, ) -> UCompletionRequest: - conversation_id = self._get_conversation_id(request) if ctx_response: request.disaggregated_params = ctx_response.choices[0].disaggregated_params request.disaggregated_params.request_type = "generation_only" request.disaggregated_params.schedule_style = self._schedule_style - request.disaggregated_params.conversation_id = conversation_id request.disaggregated_params.ctx_usage = ctx_response.usage # Replace the string prompt with prompt_tokens_ids if isinstance(request, CompletionRequest): @@ -261,7 +249,6 @@ def _get_gen_request( ctx_request_id=disagg_request_id, disagg_request_id=disagg_request_id, schedule_style=self._schedule_style, - conversation_id=conversation_id, ) if ctx_server_info and "server_info" in ctx_server_info: disaggregated_params = ctx_server_info["server_info"].get("disaggregated_params", {}) diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index ef5fd54cc875..a2608d4ae199 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -1,3 +1,18 @@ +# 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. + # Adapted from # https://github.com/vllm-project/vllm/blob/4db5176d9758b720b05460c50ace3c01026eb158/vllm/entrypoints/openai/protocol.py import base64 @@ -222,7 +237,6 @@ class DisaggregatedParams(OpenAIBaseModel): ctx_dp_rank: Optional[int] = None ctx_info_endpoint: Optional[str] = None schedule_style: Optional[DisaggScheduleStyle] = None - conversation_id: Optional[str] = None ctx_usage: Optional[UsageInfo] = None # TODO(TRTLLM-12407): Multimodal E/PD over trtllm-serve needs these protocol fields too: # encoder embedding handles, multimodal hashes, and optional mRoPE handles. @@ -1560,7 +1574,6 @@ 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, ) @@ -1585,7 +1598,6 @@ 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/disaggregated/test_coordinator_worker.py b/tests/unittest/disaggregated/test_coordinator_worker.py index 4ff67c58c465..6f6684bb6921 100644 --- a/tests/unittest/disaggregated/test_coordinator_worker.py +++ b/tests/unittest/disaggregated/test_coordinator_worker.py @@ -60,6 +60,7 @@ from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, CompletionRequest, + ConversationParams, DisaggregatedParams, ) from tensorrt_llm.serve.router import ( @@ -535,10 +536,10 @@ def _req(conv_id, request_id): return CompletionRequest( model="m", prompt="hi", + conversation_params=ConversationParams(conversation_id=conv_id), disaggregated_params=DisaggregatedParams( request_type="generation_only", ctx_request_id=request_id, - conversation_id=conv_id, ), ) @@ -583,11 +584,11 @@ async def drive(): request = CompletionRequest( model="m", prompt="hello", + conversation_params=ConversationParams(conversation_id="conv-A"), disaggregated_params=DisaggregatedParams( request_type="generation_only", ctx_request_id=assigned_id, disagg_request_id=None, - conversation_id="conv-A", ), ) await remote.gen_router.get_next_server(request) diff --git a/tests/unittest/disaggregated/test_disagg_internal_auth.py b/tests/unittest/disaggregated/test_disagg_internal_auth.py index e409269597cd..6a35cd1e4990 100644 --- a/tests/unittest/disaggregated/test_disagg_internal_auth.py +++ b/tests/unittest/disaggregated/test_disagg_internal_auth.py @@ -21,7 +21,11 @@ request_requires_internal_disagg_auth, validate_internal_disagg_request, ) -from tensorrt_llm.serve.openai_protocol import CompletionRequest, DisaggregatedParams +from tensorrt_llm.serve.openai_protocol import ( + CompletionRequest, + ConversationParams, + DisaggregatedParams, +) from tensorrt_llm.serve.openai_server import OpenAIServer @@ -114,7 +118,7 @@ def test_protected_fields_accept_valid_header_after_wire_roundtrip(): encoded_opaque_state="b3BhcXVl", ctx_info_endpoint="tcp://10.0.0.1:5000", ) - request.disaggregated_params.conversation_id = "conversation-1" + request.conversation_params = ConversationParams(conversation_id="conversation-1") headers = build_internal_disagg_auth_headers("secret", request) wire_request = CompletionRequest.model_validate_json( @@ -139,10 +143,10 @@ def test_ctx_info_endpoint_list_sender_matches_validated_string_receiver(): validate_internal_disagg_request("secret", wire_request, headers) -def test_unprotected_disagg_fields_do_not_invalidate_internal_auth_header(): +def test_conversation_params_do_not_invalidate_internal_auth_header(): request = _make_request(ctx_info_endpoint="tcp://10.0.0.1:5000") headers = build_internal_disagg_auth_headers("secret", request) - request.disaggregated_params.conversation_id = "conversation-1" + request.conversation_params = ConversationParams(conversation_id="conversation-1") validate_internal_disagg_request("secret", request, headers) diff --git a/tests/unittest/disaggregated/test_disaggregated_params.py b/tests/unittest/disaggregated/test_disaggregated_params.py index 75bdfdbe254e..ca13037e3af1 100644 --- a/tests/unittest/disaggregated/test_disaggregated_params.py +++ b/tests/unittest/disaggregated/test_disaggregated_params.py @@ -1,3 +1,18 @@ +# 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 unittest.mock import MagicMock, patch import pytest @@ -81,7 +96,6 @@ def test_to_disaggregated_params(): "cached_tokens": 4, }, }, - conversation_id="conv-abc", ) openai_params = to_disaggregated_params(llm_params) @@ -92,7 +106,6 @@ 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(): @@ -113,7 +126,6 @@ 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) @@ -123,26 +135,6 @@ 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" def test_opaque_state_round_trips_through_openai_protocol(): diff --git a/tests/unittest/disaggregated/test_openai_disagg_server.py b/tests/unittest/disaggregated/test_openai_disagg_server.py index 4f91ed302bdf..159060cacf7e 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_server.py +++ b/tests/unittest/disaggregated/test_openai_disagg_server.py @@ -186,7 +186,6 @@ def test_extract_conversation_id_preserves_body_conversation_params(): ) assert request.conversation_params.conversation_id == "body-id" - assert request.disaggregated_params.conversation_id is None def test_extract_conversation_id_populates_conversation_params_with_existing_disaggregated_params(): @@ -202,7 +201,6 @@ def test_extract_conversation_id_populates_conversation_params_with_existing_dis ) 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 5af2db0b4907..3297f547e27d 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_service.py +++ b/tests/unittest/disaggregated/test_openai_disagg_service.py @@ -234,6 +234,20 @@ def test_get_gen_request_uses_ctx_response_prompt_token_ids_for_chat(): assert gen_request.conversation_params.conversation_id == "conv-chat" +def test_get_ctx_request_preserves_conversation_params_on_wire(): + service = _make_service("context_first") + request = CompletionRequest( + model="test-model", + prompt="hello", + conversation_params=ConversationParams(conversation_id="conv-completion"), + ) + + ctx_request = service._get_ctx_request(request, 42) + + wire_request = ctx_request.model_dump(exclude_unset=True) + assert wire_request["conversation_params"]["conversation_id"] == "conv-completion" + + @pytest.mark.asyncio async def test_create_chat_response_sets_prompt_token_ids_for_context_only(): from tensorrt_llm.serve.openai_server import OpenAIServer