[None][feat] Add conversation params API (#15452)

Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Co-authored-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

Author: jiaganc

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -392,6 +392,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)
 
@@ -438,6 +439,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:

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -713,6 +713,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
 
@@ -1041,6 +1042,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)

tensorrt_llm/_torch/pyexecutor/request_utils.py

@@ -38,13 +38,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.
@@ -54,9 +56,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)
 
 
@@ -70,9 +75,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 build_no_fitting_reqs_diagnostic(active_requests, kv_cache_manager) -> str:
@@ -578,6 +582,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(
@@ -589,6 +596,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,
                 ],
             )

tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py

@@ -749,14 +749,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)."""

tensorrt_llm/conversation_params.py

@@ -0,0 +1,20 @@
+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

tensorrt_llm/disaggregated_params.py

@@ -51,10 +51,6 @@ class DisaggregatedParams:
     ctx_dp_rank: Optional[int] = None
     ctx_info_endpoint: Optional[str] = None
     schedule_style: Optional[DisaggScheduleStyle] = 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]]] = (

tensorrt_llm/executor/base_worker.py

@@ -619,6 +619,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
 

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_id: Optional[int] = None,
         arrival_time: Optional[float] = None,
         priority: float = DEFAULT_REQUEST_PRIORITY,
@@ -162,6 +164,7 @@ def generate_async(
             trace_headers=trace_headers,
             multimodal_params=multimodal_params,
             scheduling_params=scheduling_params,
+            conversation_params=conversation_params,
             cache_salt_id=cache_salt_id,
             arrival_time=arrival_time,
             priority=priority)
@@ -180,6 +183,7 @@ def generate(
         prompt_adapter_request: Optional[Union[
             PromptAdapterRequest, List[PromptAdapterRequest]]] = None,
         disaggregated_params: Optional[DisaggregatedParams] = None,
+        conversation_params: Optional[ConversationParams] = None,
     ) -> Union[GenerationResult, List[GenerationResult]]:
         """Generate output for the given prompt token ids in the synchronous mode.
         Synchronous generation accepts either single prompt or batched prompts.
@@ -212,7 +216,8 @@ def generate(
                 lora_request=lora_req,
                 prompt_adapter_request=pa_req,
                 streaming=False,
-                disaggregated_params=disaggregated_params)
+                disaggregated_params=disaggregated_params,
+                conversation_params=conversation_params)
             futures.append(future)
 
         for future in futures:

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
@@ -105,6 +106,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_id: Optional[int] = None,
         arrival_time: Optional[float] = None,
         priority: float = DEFAULT_REQUEST_PRIORITY,
@@ -134,6 +136,7 @@ def __init__(
         self.disaggregated_params = disaggregated_params
         self.trace_headers = trace_headers
         self.scheduling_params = scheduling_params
+        self.conversation_params = conversation_params
         self.cache_salt_id = cache_salt_id
         self.arrival_time = arrival_time
         if not (0.0 <= priority <= 1.0):

tensorrt_llm/llmapi/__init__.py

@@ -1,29 +1,28 @@
 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
 from ..scheduling_params import SchedulingParams
 from .build_cache import BuildCacheConfig
 from .llm import LLM, RequestOutput
-
 # yapf: disable
 from .llm_args import (AttentionDpConfig, AutoDecodingConfig, BatchingType,
                        CacheTransceiverConfig, CalibConfig,
                        CapacitySchedulerPolicy, ContextChunkingPolicy,
                        CudaGraphConfig, DeepSeekSparseAttentionConfig,
-                       DeepSeekV4SparseAttentionConfig,
-                       DFlashDecodingConfig, DraftTargetDecodingConfig,
-                       DynamicBatchConfig, Eagle3DecodingConfig,
-                       EagleDecodingConfig, ExtendedRuntimePerfKnobConfig,
-                       KvCacheConfig, LlmArgs, LookaheadDecodingConfig,
-                       MedusaDecodingConfig, MoeConfig, MTPDecodingConfig,
-                       NGramDecodingConfig, PARDDecodingConfig,
-                       PrometheusMetricsConfig, ReorderRequestPolicyConfig,
-                       RocketSparseAttentionConfig, SADecodingConfig,
-                       SAEnhancerConfig, SaveHiddenStatesDecodingConfig,
-                       SchedulerConfig, SkipSoftmaxAttentionConfig,
-                       TorchCompileConfig, TorchLlmArgs, TrtLlmArgs,
-                       UserProvidedDecodingConfig)
+                       DeepSeekV4SparseAttentionConfig, DFlashDecodingConfig,
+                       DraftTargetDecodingConfig, DynamicBatchConfig,
+                       Eagle3DecodingConfig, EagleDecodingConfig,
+                       ExtendedRuntimePerfKnobConfig, KvCacheConfig, LlmArgs,
+                       LookaheadDecodingConfig, MedusaDecodingConfig, MoeConfig,
+                       MTPDecodingConfig, NGramDecodingConfig,
+                       PARDDecodingConfig, PrometheusMetricsConfig,
+                       ReorderRequestPolicyConfig, RocketSparseAttentionConfig,
+                       SADecodingConfig, SAEnhancerConfig,
+                       SaveHiddenStatesDecodingConfig, SchedulerConfig,
+                       SkipSoftmaxAttentionConfig, TorchCompileConfig,
+                       TorchLlmArgs, TrtLlmArgs, UserProvidedDecodingConfig)
 from .llm_utils import (BuildConfig, KvCacheRetentionConfig, QuantAlgo,
                         QuantConfig)
 from .mm_encoder import MultimodalEncoder
@@ -38,6 +37,7 @@
     'GuidedDecodingParams',
     'SamplingParams',
     'DisaggregatedParams',
+    'ConversationParams',
     'DisaggScheduleStyle',
     'KvCacheConfig',
     'KvCacheRetentionConfig',