feat(trtllm): plumb conversation_id to engine for conversation-aware ADP routing

components/src/dynamo/common/backend/dp_rank.py

@@ -26,6 +26,17 @@ def forced_dp_rank(request: GenerateRequest) -> Optional[int]:
     return None if rank is None else int(rank)
 
 
+def conversation_id_from_request(request: GenerateRequest) -> Optional[str]:
+    """Pull the conversation/session id the frontend KV router forwards in
+    ``request.routing.conversation_id`` (exp H / gap #4). Populates
+    ``disaggregated_params.conversation_id`` so the engine's conversation-affinity
+    ADP router can pin a conversation to a DP rank. Returns ``None`` when absent
+    or empty."""
+    routing = cast("dict[str, Any]", request.get("routing") or {})
+    conv_id = routing.get("conversation_id")
+    return conv_id if conv_id else None
+
+
 def validate_global_dp_rank(
     dp_rank: Optional[int],
     dp_start: int,

components/src/dynamo/common/backend/tests/test_dp_rank.py

@@ -7,7 +7,11 @@
 
 import pytest
 
-from dynamo.common.backend.dp_rank import forced_dp_rank, validate_global_dp_rank
+from dynamo.common.backend.dp_rank import (
+    conversation_id_from_request,
+    forced_dp_rank,
+    validate_global_dp_rank,
+)
 
 pytestmark = [pytest.mark.unit, pytest.mark.gpu_0, pytest.mark.pre_merge]
 
@@ -23,6 +27,20 @@ def test_forced_dp_rank_coerces_to_int():
     assert forced_dp_rank({"token_ids": [], "routing": {"dp_rank": 3}}) == 3
 
 
+def test_conversation_id_from_request():
+    # exp H (gap #4): the frontend KV router forwards the session id as
+    # routing.conversation_id; the worker reads it to drive engine conv-affinity.
+    assert conversation_id_from_request({"token_ids": [1]}) is None
+    assert conversation_id_from_request({"token_ids": [1], "routing": {}}) is None
+    assert conversation_id_from_request({"routing": {"conversation_id": ""}}) is None
+    assert (
+        conversation_id_from_request(
+            {"token_ids": [], "routing": {"conversation_id": "conv-a:abc123", "dp_rank": 2}}
+        )
+        == "conv-a:abc123"
+    )
+
+
 def test_validate_global_dp_rank_passes_in_range():
     # dp_start=2, dp_size=4 → valid global ranks are [2, 6).
     assert validate_global_dp_rank(2, 2, 4, "test") == 2

components/src/dynamo/trtllm/llm_engine.py

@@ -35,7 +35,11 @@
 from dynamo._core import Context
 from dynamo.common.backend import telemetry
 from dynamo.common.backend.disagg import require_prefill_result
-from dynamo.common.backend.dp_rank import forced_dp_rank, validate_global_dp_rank
+from dynamo.common.backend.dp_rank import (
+    conversation_id_from_request,
+    forced_dp_rank,
+    validate_global_dp_rank,
+)
 from dynamo.common.backend.engine import (
     TEST_LOGITS_PROCESSOR_ENV,
     EngineConfig,
@@ -368,6 +372,16 @@ async def start(self, worker_id: int) -> EngineConfig:
         ) or getattr(self._trtllm_metrics_collector, "log_metrics_dict", None)
 
         self._attention_dp_size = self._engine.get_attention_dp_size()
+        # exp-H (gap #4): engine conversation-affinity mode. When ON, the worker does NOT
+        # force attention_dp_rank — it sets disaggregated_params.conversation_id and lets the
+        # engine's ConversationAwareADPRouter pick the rank (round-robin balanced + sticky).
+        self._engine_conv_affinity = os.environ.get("DYN_ENGINE_CONV_AFFINITY") == "1"
+        logger.info(
+            "exp-H: engine conversation-affinity mode = %s (DYN_ENGINE_CONV_AFFINITY); attention_dp_size=%d",
+            "ON" if self._engine_conv_affinity else "OFF",
+            self._attention_dp_size,
+        )
+        self._convid_first_logged = False
         # Always start the metrics-poll thread: it pushes the latest
         # ComponentSnapshot into the framework's SnapshotPublisher and
         # forwards each snap to `_log_iteration_stats` for `trtllm_kv_cache_*`.
@@ -789,6 +803,31 @@ async def _generate_started(
             )
             disaggregated_params = self._decode_prefill_handoff(prefill_result)
 
+        # exp H (gap #4): forward the conversation/session id (set by the frontend KV
+        # router into RoutingHints.conversation_id) onto disaggregated_params so the
+        # engine's conversation-affinity ADP router can pin conv -> DP rank.
+        conv_id = conversation_id_from_request(request)
+        if disaggregated_params is not None and conv_id is not None:
+            disaggregated_params.conversation_id = conv_id
+        # exp-H (gap #4): confirm propagation at run time. PERF-SAFE: log the FIRST request at
+        # INFO (one-shot) so it's confirmable WITHOUT enabling debug at benchmark time; rest debug.
+        _set_on_disagg = disaggregated_params is not None and conv_id is not None
+        if not self._convid_first_logged:
+            self._convid_first_logged = True
+            logger.info(
+                "exp-H convid FIRST: conv_id=%s set_on_disagg=%s engine_conv_affinity=%s (rest at debug)",
+                conv_id,
+                _set_on_disagg,
+                self._engine_conv_affinity,
+            )
+        else:
+            logger.debug(
+                "exp-H convid: conv_id=%s set_on_disagg=%s engine_conv_affinity=%s",
+                conv_id,
+                _set_on_disagg,
+                self._engine_conv_affinity,
+            )
+
         stop_conditions = request.get("stop_conditions", {})
         if is_prefill:
             # Prefill only needs to populate KV — one token is enough.
@@ -804,16 +843,22 @@ async def _generate_started(
         if ignore_eos:
             sampling_params.ignore_eos = ignore_eos
 
-        # Honour the router's DP rank decision; without it TRT-LLM picks
-        # its own rank and KV events land on the wrong publisher.
-        rank = validate_global_dp_rank(
-            forced_dp_rank(request), 0, self._attention_dp_size, "TRT-LLM"
-        )
-        scheduling_params = (
-            SchedulingParams(attention_dp_rank=rank, attention_dp_relax=False)
-            if rank is not None
-            else None
-        )
+        # exp H (gap #4): in engine conversation-affinity mode (DYN_ENGINE_CONV_AFFINITY=1)
+        # do NOT force the rank — let the engine's ConversationAwareADPRouter pick it from
+        # disaggregated_params.conversation_id (round-robin balanced + sticky), matching
+        # native trtllm-serve. Otherwise honour the frontend KV router's DP rank decision
+        # (without it TRT-LLM picks its own rank and KV events land on the wrong publisher).
+        if self._engine_conv_affinity:
+            scheduling_params = None
+        else:
+            rank = validate_global_dp_rank(
+                forced_dp_rank(request), 0, self._attention_dp_size, "TRT-LLM"
+            )
+            scheduling_params = (
+                SchedulingParams(attention_dp_rank=rank, attention_dp_relax=False)
+                if rank is not None
+                else None
+            )
 
         entries = logits_processors_for_request(
             self._logits_processor_spec,

components/src/dynamo/trtllm/request_handlers/handler_base.py

@@ -34,6 +34,7 @@
 from tensorrt_llm.scheduling_params import SchedulingParams
 
 from dynamo._core import Client, Context
+from dynamo.common.backend.dp_rank import conversation_id_from_request
 from dynamo.common.utils.structural_tag import serialize_structural_tag
 from dynamo.health_check import HEALTH_CHECK_KEY
 from dynamo.llm.exceptions import EngineShutdown
@@ -63,6 +64,11 @@
 logger = logging.getLogger(__name__)
 
 
+# exp-H (gap #4): one-shot runtime-verification flag (perf-safe — logs once at INFO so the
+# conversation_id plumbing can be confirmed without enabling debug at benchmark time).
+_EXPH_FIRST_LOGGED = False
+
+
 class TRTLLMEnginePauseController:
     """Adapts TRT-LLM sleep/wake to the standard pause controller interface.
 
@@ -793,6 +799,13 @@ def _setup_disaggregated_params_for_mode(
             # For full EPD flow, make decoded params available to multimodal processor
             ep_disaggregated_params = disaggregated_params
 
+        # exp-H (gap #4): plumb conversation_id (frontend seeded routing.conversation_id from
+        # nvext.session_control) onto disaggregated_params so the engine's ConversationAwareADPRouter
+        # can pin conv -> DP rank. Covers PREFILL (context_only) / AGGREGATED / DECODE (generation_only).
+        _conv_id = conversation_id_from_request(request)
+        if disaggregated_params is not None and _conv_id is not None:
+            disaggregated_params.conversation_id = _conv_id
+
         return disaggregated_params, ep_disaggregated_params, epd_metadata
 
     async def _prepare_input_for_generation(
@@ -1139,15 +1152,37 @@ async def _generate_locally_impl(
         # Extract dp_rank from request's routing hints for attention DP routing
         routing = request.get("routing", {})
         dp_rank = routing.get("dp_rank") if routing else None
+        # exp-H (gap #4): when DYN_ENGINE_CONV_AFFINITY=1, do NOT force attention_dp_rank — the
+        # engine's ConversationAwareADPRouter must pick the rank from disaggregated_params.conversation_id.
+        # An explicit attention_dp_rank is honored BEFORE conv-affinity in the engine and would bypass it.
+        engine_conv_affinity = os.environ.get("DYN_ENGINE_CONV_AFFINITY") == "1"
         scheduling_params = None
-        if dp_rank is not None:
+        if dp_rank is not None and not engine_conv_affinity:
             scheduling_params = SchedulingParams(
                 attention_dp_rank=dp_rank,
                 attention_dp_relax=False,  # Strict routing - use the rank dynamo router selected
             )
             logging.debug(
                 f"Using dynamo router dp_rank={dp_rank} for TRTLLM attention DP scheduling"
             )
+        # exp-H one-shot runtime verification (perf-safe; visible with debug OFF). Confirms the two
+        # invariants the engine conv-affinity router needs: conversation_id set + rank NOT forced.
+        global _EXPH_FIRST_LOGGED
+        if not _EXPH_FIRST_LOGGED:
+            _EXPH_FIRST_LOGGED = True
+            _cid = (
+                disaggregated_params.conversation_id
+                if disaggregated_params is not None
+                else None
+            )
+            logging.info(
+                "exp-H FIRST req: conversation_id=%s attention_dp_rank_forced=%s "
+                "engine_conv_affinity=%s disagg_mode=%s (subsequent requests silent)",
+                _cid,
+                scheduling_params is not None,
+                engine_conv_affinity,
+                self.disaggregation_mode,
+            )
 
         # Priority is a float in [0.0, 1.0]; health checks use 1.0. Default is 0.5.
         priority = request.get("priority", DEFAULT_REQUEST_PRIORITY)

lib/llm/src/kv_router/push_router.rs

@@ -778,6 +778,21 @@ impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<LLMEngineOutpu
         let route_outcome = self.sticky.on_routed(&request, worker, &context_id).await?;
         let deferred_close = route_outcome.deferred_close;
 
+        // exp-H (gap #4): confirm the conversation_id (seeded by the preprocessor from
+        // session_control) is on the routing we dispatch to the worker. PERF-SAFE: the FIRST
+        // routed request logs at INFO/WARN (one-shot) so propagation is confirmable WITHOUT
+        // enabling debug at benchmark time; subsequent requests are debug-only. grep "exp-H convid".
+        {
+            static FIRST: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
+            let first = !FIRST.swap(true, std::sync::atomic::Ordering::Relaxed);
+            match request.routing.as_ref().and_then(|r| r.conversation_id.as_deref()) {
+                Some(cid) if first => tracing::info!(request_id = %context_id, conversation_id = %cid, dp_rank, "exp-H convid forwarded to worker (first; rest at debug)"),
+                Some(cid) => tracing::debug!(request_id = %context_id, conversation_id = %cid, dp_rank, "exp-H convid forwarded to worker"),
+                None if first => tracing::warn!(request_id = %context_id, dp_rank, "exp-H convid ABSENT on first routed request — conv-affinity will NOT engage (session_control missing?)"),
+                None => tracing::debug!(request_id = %context_id, dp_rank, "exp-H convid ABSENT"),
+            }
+        }
+
         let (mut backend_input, context) = request.into_parts();
         backend_input.routing_mut().dp_rank = Some(dp_rank);
         let updated_request = context.map(|_| backend_input);

lib/llm/src/preprocessor.rs

@@ -916,6 +916,9 @@ impl OpenAIPreprocessor {
                 prefill_worker_id: nvext.prefill_worker_id,
                 decode_worker_id: nvext.decode_worker_id,
                 dp_rank: nvext.dp_rank,
+                // exp H (gap #4): seed conversation_id from session_control so the engine's
+                // conversation-affinity ADP router can pin conv -> DP rank.
+                conversation_id: nvext.session_control.as_ref().map(|sc| sc.session_id.clone()),
                 prefill_dp_rank: nvext.prefill_dp_rank,
                 expected_output_tokens: hints.and_then(|h| h.osl),
                 priority_jump: hints.and_then(|h| {

lib/llm/src/protocols/common/preprocessor.rs

@@ -55,6 +55,12 @@ pub struct RoutingHints {
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub dp_rank: Option<u32>,
 
+    /// Conversation/session id forwarded to the worker so the engine's
+    /// conversation-affinity ADP router can pin conv -> DP rank (exp H / gap #4).
+    /// Sourced from `session_control.session_id`.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub conversation_id: Option<String>,
+
     /// Data parallel rank for the prefill worker in disaggregated serving
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub prefill_dp_rank: Option<u32>,