fix(trtllm): expose env var to force engine side conversation dp assignment

components/src/dynamo/trtllm/backend_args.py

@@ -91,6 +91,14 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None:
             default=False,
             help="Enable attention data parallelism. When enabled, attention_dp_size equals tensor_parallel_size.",
         )
+        add_negatable_bool_argument(
+            g,
+            flag_name="--conversation-affinity",
+            env_var="DYN_ENGINE_CONV_AFFINITY",
+            default=False,
+            help="Force engine-owned conversation-affinity ADP routing: the engine picks the "
+            "attention-DP rank from the conversation id, even if the router selects a rank.",
+        )
         add_argument(
             g,
             flag_name="--kv-block-size",
@@ -471,6 +479,7 @@ class DynamoTrtllmConfig(ConfigBase):
     pipeline_parallel_size: int
     expert_parallel_size: Optional[int]
     enable_attention_dp: bool
+    conversation_affinity: bool
     kv_block_size: int
     gpus_per_node: Optional[int] = None
     max_batch_size: int

components/src/dynamo/trtllm/llm_engine.py

@@ -142,6 +142,7 @@ def __init__(
         disaggregation_mode: DisaggregationMode = DisaggregationMode.AGGREGATED,
         component: str = "backend",
         publish_events_and_metrics: bool = False,
+        conversation_affinity: bool = False,
     ):
         self.engine_args = engine_args
         self.model_name = model_name
@@ -195,6 +196,9 @@ def __init__(
         # Default False so generate() is safe if reached before initialize()
         # (e.g. unit tests that drive generate directly).
         self._conversation_affinity: bool = False
+        # Manual override (--conversation-affinity / DYN_ENGINE_CONV_AFFINITY) to force
+        # engine-side conversation-affinity assignment regardless of engine detection.
+        self._engine_conversation_affinity_override: bool = conversation_affinity
         # Worker invokes on_ready callbacks serially during setup (see
         # `setup_kv_publishers` in lib/backend-common/src/publisher.rs); the
         # dict is fully populated before `_kv_events_thread` starts and
@@ -330,6 +334,7 @@ async def from_args(
             disaggregation_mode=config.disaggregation_mode,
             component=config.component,
             publish_events_and_metrics=config.publish_events_and_metrics,
+            conversation_affinity=config.conversation_affinity,
         )
         worker_config = WorkerConfig.from_runtime_config(
             config,
@@ -879,8 +884,24 @@ async def _generate_started(
         # the engine side (always record the binding, whether the rank came
         # from explicit placement or fresh selection). Track upstream and drop
         # this note when the router records unconditionally.
+        # Force engine-side conversation-affinity assignment when the operator set
+        # --conversation-affinity / DYN_ENGINE_CONV_AFFINITY, even if the engine did
+        # not advertise the capability at initialize() time. Mirrors the legacy
+        # HandlerBase path so both entry points behave identically.
+        conv_affinity = (
+            self._conversation_affinity or self._engine_conversation_affinity_override
+        )
+        if (
+            self._engine_conversation_affinity_override
+            and not CONVERSATION_PARAMS_AVAILABLE
+        ):
+            raise RuntimeError(
+                "--conversation-affinity / DYN_ENGINE_CONV_AFFINITY is set but "
+                "the installed TensorRT-LLM build has no ConversationParams API (requires a "
+                "release newer than 1.3.0rc20)."
+            )
         conversation_params = None
-        if self._conversation_affinity:
+        if conv_affinity:
             scheduling_params = None
             conversation_params = conversation_params_for(
                 session_id_from_request(request)
@@ -908,9 +929,7 @@ async def _generate_started(
         # Only pass conversation_params when affinity is active — older TRT-LLM builds'
         # generate_async does not accept the keyword.
         conv_kwargs = (
-            {"conversation_params": conversation_params}
-            if self._conversation_affinity
-            else {}
+            {"conversation_params": conversation_params} if conv_affinity else {}
         )
         generation_result = self._engine.llm.generate_async(
             inputs=token_ids,

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

@@ -247,6 +247,8 @@ class RequestHandlerConfig:
     additional_metrics: Optional["AdditionalMetricsCollector"] = None
     max_seq_len: Optional[int] = None
     disagg_machine_id: int = 0  # 10-bit machine_id for snowflake disagg_request_id
+    # Force engine-owned conversation-affinity ADP routing regardless of engine detection.
+    conversation_affinity: bool = False
 
 
 class HandlerBase(BaseGenerativeHandler):
@@ -271,6 +273,9 @@ def __init__(self, config: RequestHandlerConfig):
         # request (engine may not be initialized at handler construction). See
         # conversation_affinity.py. None = not yet resolved.
         self._conversation_affinity: Optional[bool] = None
+        # Manual override (--conversation-affinity / DYN_ENGINE_CONV_AFFINITY) to force
+        # engine-side assignment of conversation-affinity regardless of engine detection.
+        self._engine_conversation_affinity_override: bool = config.conversation_affinity
         self.encode_client = config.encode_client
         self.multimodal_processor = config.multimodal_processor
         self.first_generation = True
@@ -1110,8 +1115,18 @@ async def _generate_locally_impl(
         # initialized by first request). When on, the engine's ConversationAwareADPRouter
         # picks the attention-DP rank from the conversation id, so we must NOT force a rank.
         if self._conversation_affinity is None:
-            self._conversation_affinity = engine_conversation_affinity_enabled(
-                self.engine.llm
+            if (
+                self._engine_conversation_affinity_override
+                and not CONVERSATION_PARAMS_AVAILABLE
+            ):
+                raise RuntimeError(
+                    "--conversation-affinity / DYN_ENGINE_CONV_AFFINITY is set but "
+                    "the installed TensorRT-LLM build has no ConversationParams API (requires a "
+                    "release newer than 1.3.0rc20)."
+                )
+            self._conversation_affinity = (
+                self._engine_conversation_affinity_override
+                or engine_conversation_affinity_enabled(self.engine.llm)
             )
             if self._conversation_affinity and not CONVERSATION_PARAMS_AVAILABLE:
                 raise RuntimeError(

components/src/dynamo/trtllm/tests/test_trtllm_engine_conversation_affinity.py

@@ -14,6 +14,11 @@
   ``conversation_params.conversation_id == session_id``.
 * on + no session id → ``scheduling_params=None`` and ``conversation_params=None``
   (the kwarg is still passed so the engine's no-id balancing fallback runs).
+
+Plus the manual ``--conversation-affinity`` / ``DYN_ENGINE_CONV_AFFINITY`` override
+(``_engine_conversation_affinity_override``), which forces the affinity branch even when
+engine detection is off and raises when the ConversationParams API is missing — matching
+the legacy ``HandlerBase`` override tests.
 """
 
 from __future__ import annotations
@@ -81,6 +86,9 @@ def _make_engine(
     engine._no_inflight_requests.set()
     engine._reject_new_requests = False
     engine._conversation_affinity = conversation_affinity
+    # Manual override (--conversation-affinity / DYN_ENGINE_CONV_AFFINITY); off by
+    # default, individual tests flip it to exercise the override path.
+    engine._engine_conversation_affinity_override = False
     engine._attention_dp_size = attention_dp_size
     engine._override_sampling_params = lambda default, request: SimpleNamespace(
         max_tokens=None
@@ -189,3 +197,72 @@ def fake_generate_async(**kwargs):
     assert captured["scheduling_params"] is None
     assert "conversation_params" in captured
     assert captured["conversation_params"] is None
+
+
+async def test_override_suppresses_dp_rank_and_forwards_conversation_params(
+    monkeypatch,
+):
+    """DYN_ENGINE_CONV_AFFINITY override=True + engine detection disabled →
+    dp_rank suppressed, conversation_params forwarded. Mirrors the legacy
+    HandlerBase override test so both entry points stay in lockstep."""
+    monkeypatch.setattr(
+        "dynamo.trtllm.llm_engine.CONVERSATION_PARAMS_AVAILABLE",
+        True,
+    )
+    monkeypatch.setattr(
+        "dynamo.trtllm.conversation_affinity.ConversationParams",
+        _FakeConversationParams,
+    )
+
+    captured: dict = {}
+
+    def fake_generate_async(**kwargs):
+        captured.update(kwargs)
+        return _empty_async_iter()
+
+    # Engine detection returns False (no engine-side affinity config)...
+    engine = _make_engine(
+        fake_generate_async, conversation_affinity=False, attention_dp_size=8
+    )
+    # ...but the operator forced it via DYN_ENGINE_CONV_AFFINITY.
+    engine._engine_conversation_affinity_override = True
+    # Router still stamps a rank; the override must ignore it just like
+    # auto-detection does.
+    await _drain(
+        engine,
+        {
+            "token_ids": [1, 2, 3],
+            "routing": {"dp_rank": 3},
+            "agent_context": {"session_id": "run-99:agent-0"},
+        },
+    )
+
+    assert captured["scheduling_params"] is None
+    conv_params = captured["conversation_params"]
+    assert conv_params is not None
+    assert conv_params.conversation_id == "run-99:agent-0"
+
+
+async def test_override_raises_when_conversation_params_api_missing(monkeypatch):
+    """DYN_ENGINE_CONV_AFFINITY=true on a build without ConversationParams →
+    RuntimeError in _generate_started, before generate_async is called."""
+    monkeypatch.setattr(
+        "dynamo.trtllm.llm_engine.CONVERSATION_PARAMS_AVAILABLE",
+        False,
+    )
+
+    called = False
+
+    def fake_generate_async(**kwargs):
+        nonlocal called
+        called = True
+        return _empty_async_iter()
+
+    engine = _make_engine(
+        fake_generate_async, conversation_affinity=False, attention_dp_size=8
+    )
+    engine._engine_conversation_affinity_override = True
+
+    with pytest.raises(RuntimeError, match="DYN_ENGINE_CONV_AFFINITY"):
+        await _drain(engine, {"token_ids": [1, 2, 3]})
+    assert called is False

components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py

@@ -470,6 +470,7 @@ class TestDeferredAbortGuard:
     def _make_handler(self) -> HandlerBase:
         config = MagicMock()
         config.shutdown_event = None
+        config.conversation_affinity = False
         return _ConcreteHandler(config)
 
     @pytest.mark.asyncio
@@ -591,6 +592,7 @@ def _make_handler(self, multimodal_processor=None) -> HandlerBase:
         config = MagicMock()
         config.multimodal_processor = multimodal_processor
         config.shutdown_event = None
+        config.conversation_affinity = False
         return _ConcreteHandler(config)
 
     async def _prepare(self, handler, request, epd_metadata=None):
@@ -647,6 +649,7 @@ def _make_handler(self) -> HandlerBase:
         config = MagicMock()
         config.multimodal_processor = MagicMock()
         config.shutdown_event = None
+        config.conversation_affinity = False
         return _ConcreteHandler(config)
 
     def _pack_metadata(self, request, processed_input, prompt, prompt_token_ids):
@@ -717,6 +720,7 @@ def _make_prefill_handler(self, machine_id: int = 42) -> HandlerBase:
         config = MagicMock()
         config.shutdown_event = None
         config.disagg_machine_id = machine_id
+        config.conversation_affinity = False
         handler = _ConcreteHandler(config)
         handler.disaggregation_mode = DisaggregationMode.PREFILL
         return handler
@@ -800,6 +804,7 @@ def _make_handler(self) -> HandlerBase:
         config = MagicMock()
         config.shutdown_event = None
         config.disaggregation_mode = DisaggregationMode.AGGREGATED
+        config.conversation_affinity = False
         handler = _ConcreteHandler(config)
         handler.publisher = None
         handler.multimodal_processor = None
@@ -1128,6 +1133,7 @@ def _make_handler(self, *, conversation_affinity: bool) -> HandlerBase:
         config = MagicMock()
         config.shutdown_event = None
         config.disaggregation_mode = DisaggregationMode.AGGREGATED
+        config.conversation_affinity = False
         handler = _ConcreteHandler(config)
         handler.publisher = None
         handler.multimodal_processor = None
@@ -1284,3 +1290,69 @@ async def test_lazy_resolution_raises_when_api_missing(self, monkeypatch):
             async for _ in handler.generate_locally(request, self._make_context()):
                 pass
         handler.engine.llm.generate_async.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_override_suppresses_dp_rank_and_forwards_conversation_params(
+        self, monkeypatch
+    ):
+        """DYN_ENGINE_CONV_AFFINITY override=True + engine detection disabled →
+        dp_rank suppressed, conversation_params forwarded on first request."""
+        monkeypatch.setattr(
+            "dynamo.trtllm.request_handlers.handler_base.CONVERSATION_PARAMS_AVAILABLE",
+            True,
+        )
+        monkeypatch.setattr(
+            "dynamo.trtllm.conversation_affinity.ConversationParams",
+            _FakeConversationParams,
+        )
+        monkeypatch.setattr(
+            "dynamo.trtllm.request_handlers.handler_base.engine_conversation_affinity_enabled",
+            lambda _: False,
+        )
+        handler = self._make_handler(conversation_affinity=False)
+        # Reset to None so lazy init runs on first request and folds in the override.
+        handler._conversation_affinity = None
+        handler._engine_conversation_affinity_override = True
+        kwargs = await self._drive(
+            handler,
+            {
+                "token_ids": [1, 2, 3],
+                "stop_conditions": {"max_tokens": 10},
+                "sampling_options": {"temperature": 0.7},
+                "routing": {"dp_rank": 3},
+                "agent_context": {"session_id": "run-99:agent-0"},
+            },
+        )
+        # Override must suppress the router rank just like auto-detection does.
+        assert kwargs["scheduling_params"] is None
+        conv_params = kwargs["conversation_params"]
+        assert conv_params is not None
+        assert conv_params.conversation_id == "run-99:agent-0"
+
+    @pytest.mark.asyncio
+    async def test_override_raises_when_conversation_params_api_missing(
+        self, monkeypatch
+    ):
+        """DYN_ENGINE_CONV_AFFINITY=true on a build without ConversationParams →
+        RuntimeError on first request during lazy init."""
+        monkeypatch.setattr(
+            "dynamo.trtllm.request_handlers.handler_base.CONVERSATION_PARAMS_AVAILABLE",
+            False,
+        )
+        handler = self._make_handler(conversation_affinity=False)
+        # Reset to None so lazy init runs and hits the guard.
+        handler._conversation_affinity = None
+        handler._engine_conversation_affinity_override = True
+        handler.engine.llm.generate_async = MagicMock()
+
+        with pytest.raises(RuntimeError, match="DYN_ENGINE_CONV_AFFINITY"):
+            async for _ in handler.generate_locally(
+                {
+                    "token_ids": [1, 2, 3],
+                    "stop_conditions": {"max_tokens": 10},
+                    "sampling_options": {"temperature": 0.7},
+                },
+                self._make_context(),
+            ):
+                pass
+        handler.engine.llm.generate_async.assert_not_called()

components/src/dynamo/trtllm/tests/test_trtllm_trace_propagation.py

@@ -68,8 +68,9 @@ def _make_engine(generate_async) -> TrtllmLLMEngine:
     engine._no_inflight_requests = asyncio.Event()
     engine._no_inflight_requests.set()
     engine._reject_new_requests = False
-    # Match __init__ default; generate_locally reads this before dispatch.
+    # Match __init__ defaults; _generate_started reads both before dispatch.
     engine._conversation_affinity = False
+    engine._engine_conversation_affinity_override = False
     # Single-rank default: validate_global_dp_rank(None, 0, 1, ...) -> None,
     # so scheduling_params stays None and never touches a real SchedulingParams.
     engine._attention_dp_size = 1

components/src/dynamo/trtllm/tests/test_trtllm_unit.py

@@ -249,6 +249,27 @@ def test_disaggregation_mode_legacy_aggregated_value_warns():
     assert config.disaggregation_mode == DisaggregationMode.AGGREGATED
 
 
+def test_conversation_affinity_cli_flag(monkeypatch):
+    """--conversation-affinity sets conversation_affinity=True in Config."""
+    monkeypatch.delenv("DYN_ENGINE_CONV_AFFINITY", raising=False)
+    config = parse_args(["--model", "fake-model", "--conversation-affinity"])
+    assert config.conversation_affinity is True
+
+
+def test_conversation_affinity_env_var(monkeypatch):
+    """DYN_ENGINE_CONV_AFFINITY=true is read and sets conversation_affinity=True."""
+    monkeypatch.setenv("DYN_ENGINE_CONV_AFFINITY", "true")
+    config = parse_args(["--model", "fake-model"])
+    assert config.conversation_affinity is True
+
+
+def test_conversation_affinity_defaults_false(monkeypatch):
+    """conversation_affinity defaults to False when neither flag nor env var is set."""
+    monkeypatch.delenv("DYN_ENGINE_CONV_AFFINITY", raising=False)
+    config = parse_args(["--model", "fake-model"])
+    assert config.conversation_affinity is False
+
+
 def test_enable_multimodal_rejects_diffusion_modality():
     with pytest.raises(ValueError, match="--enable-multimodal cannot be combined"):
         parse_args(

components/src/dynamo/trtllm/tests/utils.py

@@ -32,5 +32,6 @@ def create_mock_request_handler_config(
     config.runtime = None
     config.kv_block_size = 32
     config.shutdown_event = None
+    config.conversation_affinity = False
     config.encoder_cache_capacity_gb = encoder_cache_capacity_gb
     return config

components/src/dynamo/trtllm/workers/llm_worker.py

@@ -678,6 +678,7 @@ async def init_llm_worker(
             additional_metrics=additional_metrics,
             max_seq_len=config.max_seq_len,
             disagg_machine_id=int(endpoint.connection_id()) % 1021,
+            conversation_affinity=config.conversation_affinity,
         )
 
         media_decoder = None