diff --git a/components/src/dynamo/common/backend/dp_rank.py b/components/src/dynamo/common/backend/dp_rank.py index be5a43a55236..f97d7c398c4e 100644 --- a/components/src/dynamo/common/backend/dp_rank.py +++ b/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, diff --git a/components/src/dynamo/common/backend/tests/test_dp_rank.py b/components/src/dynamo/common/backend/tests/test_dp_rank.py index aeb09d0ff7f5..c046d5e3c901 100644 --- a/components/src/dynamo/common/backend/tests/test_dp_rank.py +++ b/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 diff --git a/components/src/dynamo/trtllm/llm_engine.py b/components/src/dynamo/trtllm/llm_engine.py index b02d735f6e69..632481a5cb9d 100644 --- a/components/src/dynamo/trtllm/llm_engine.py +++ b/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, diff --git a/components/src/dynamo/trtllm/request_handlers/handler_base.py b/components/src/dynamo/trtllm/request_handlers/handler_base.py index afa59dba6128..ad077e6b8b8e 100644 --- a/components/src/dynamo/trtllm/request_handlers/handler_base.py +++ b/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,8 +1152,12 @@ 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 @@ -1148,6 +1165,24 @@ async def _generate_locally_impl( 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) diff --git a/lib/llm/src/kv_router/push_router.rs b/lib/llm/src/kv_router/push_router.rs index 2f48c1de4290..728a349847c0 100644 --- a/lib/llm/src/kv_router/push_router.rs +++ b/lib/llm/src/kv_router/push_router.rs @@ -778,6 +778,21 @@ impl AsyncEngine, ManyOut 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); diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index c7894ec178f0..60e44067c289 100644 --- a/lib/llm/src/preprocessor.rs +++ b/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| { diff --git a/lib/llm/src/protocols/common/preprocessor.rs b/lib/llm/src/protocols/common/preprocessor.rs index a04b5f7b2e08..b7463f084e37 100644 --- a/lib/llm/src/protocols/common/preprocessor.rs +++ b/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, + /// 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, + /// Data parallel rank for the prefill worker in disaggregated serving #[serde(default, skip_serializing_if = "Option::is_none")] pub prefill_dp_rank: Option,