From 85abbe3669149570305d416b0527f768448fbd11 Mon Sep 17 00:00:00 2001 From: Guan Luo <41310872+GuanLuo@users.noreply.github.com> Date: Tue, 14 Jul 2026 18:06:02 -0700 Subject: [PATCH 01/11] fix: expose env var to force engine side conversation dp assignment Signed-off-by: Guan Luo <41310872+GuanLuo@users.noreply.github.com> --- .../src/dynamo/trtllm/request_handlers/handler_base.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/components/src/dynamo/trtllm/request_handlers/handler_base.py b/components/src/dynamo/trtllm/request_handlers/handler_base.py index 5fc3c849988c..be10872df4a9 100644 --- a/components/src/dynamo/trtllm/request_handlers/handler_base.py +++ b/components/src/dynamo/trtllm/request_handlers/handler_base.py @@ -266,6 +266,10 @@ 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 to force engine-side assignment of conversation-affinity + self._engine_conversation_affinity_env_override: bool = ( + os.getenv("DYN_ENGINE_CONV_AFFINITY") == "1" + ) self.encode_client = config.encode_client self.multimodal_processor = config.multimodal_processor self.first_generation = True @@ -1146,7 +1150,7 @@ async def _generate_locally_impl( dp_rank = routing.get("dp_rank") if routing else None conversation_params = None scheduling_params = None - if conv_affinity: + if conv_affinity or self._engine_conversation_affinity_env_override: # Let the engine pick the rank from the conversation id (agent_context.session_id); # do NOT force a rank (an explicit rank is honored before affinity and bypasses it). conversation_params = conversation_params_for( From 266ab73bcad342578b27be800528eac6df09c26e Mon Sep 17 00:00:00 2001 From: Guan Luo <41310872+GuanLuo@users.noreply.github.com> Date: Wed, 15 Jul 2026 00:58:16 -0700 Subject: [PATCH 02/11] chore: expose as cmdline arg as well Signed-off-by: Guan Luo <41310872+GuanLuo@users.noreply.github.com> --- components/src/dynamo/trtllm/backend_args.py | 9 +++++++++ .../dynamo/trtllm/request_handlers/handler_base.py | 11 ++++++----- components/src/dynamo/trtllm/workers/llm_worker.py | 1 + 3 files changed, 16 insertions(+), 5 deletions(-) diff --git a/components/src/dynamo/trtllm/backend_args.py b/components/src/dynamo/trtllm/backend_args.py index 002d9cd98ccd..5ebbd02bb549 100644 --- a/components/src/dynamo/trtllm/backend_args.py +++ b/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 diff --git a/components/src/dynamo/trtllm/request_handlers/handler_base.py b/components/src/dynamo/trtllm/request_handlers/handler_base.py index be10872df4a9..7f6a47ba0ad1 100644 --- a/components/src/dynamo/trtllm/request_handlers/handler_base.py +++ b/components/src/dynamo/trtllm/request_handlers/handler_base.py @@ -242,6 +242,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): @@ -266,10 +268,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 to force engine-side assignment of conversation-affinity - self._engine_conversation_affinity_env_override: bool = ( - os.getenv("DYN_ENGINE_CONV_AFFINITY") == "1" - ) + # 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 @@ -1150,7 +1151,7 @@ async def _generate_locally_impl( dp_rank = routing.get("dp_rank") if routing else None conversation_params = None scheduling_params = None - if conv_affinity or self._engine_conversation_affinity_env_override: + if conv_affinity or self._engine_conversation_affinity_override: # Let the engine pick the rank from the conversation id (agent_context.session_id); # do NOT force a rank (an explicit rank is honored before affinity and bypasses it). conversation_params = conversation_params_for( diff --git a/components/src/dynamo/trtllm/workers/llm_worker.py b/components/src/dynamo/trtllm/workers/llm_worker.py index 5f71ca4ef585..f3b8001a7b35 100644 --- a/components/src/dynamo/trtllm/workers/llm_worker.py +++ b/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 From abaff249eb97c31d8349b3274c9ca88a803ab059 Mon Sep 17 00:00:00 2001 From: Guan Luo <41310872+GuanLuo@users.noreply.github.com> Date: Wed, 15 Jul 2026 09:27:24 -0700 Subject: [PATCH 03/11] fix: fix up Signed-off-by: Guan Luo <41310872+GuanLuo@users.noreply.github.com> --- .../src/dynamo/trtllm/request_handlers/handler_base.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/components/src/dynamo/trtllm/request_handlers/handler_base.py b/components/src/dynamo/trtllm/request_handlers/handler_base.py index 7f6a47ba0ad1..1f443664391f 100644 --- a/components/src/dynamo/trtllm/request_handlers/handler_base.py +++ b/components/src/dynamo/trtllm/request_handlers/handler_base.py @@ -1144,14 +1144,16 @@ async def _generate_locally_impl( "the installed TensorRT-LLM build has no ConversationParams API (requires a " "release newer than 1.3.0rc20)." ) - conv_affinity = self._conversation_affinity + conv_affinity = ( + self._conversation_affinity or self._engine_conversation_affinity_override + ) # 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 conversation_params = None scheduling_params = None - if conv_affinity or self._engine_conversation_affinity_override: + if conv_affinity: # Let the engine pick the rank from the conversation id (agent_context.session_id); # do NOT force a rank (an explicit rank is honored before affinity and bypasses it). conversation_params = conversation_params_for( From 1c06fb2a18f26c53064a05aa0ef4cf4e35a3c4d5 Mon Sep 17 00:00:00 2001 From: Pei Li Date: Wed, 15 Jul 2026 12:44:23 -0700 Subject: [PATCH 04/11] test(trtllm): pin conversation_affinity=False on MagicMock config in TestConversationAffinity MagicMock auto-creates truthy attributes; config.conversation_affinity was returning a MagicMock object instead of False, making _engine_conversation_affinity_override truthy and causing the affinity-off test cases to take the conversation_params branch instead of the scheduling_params branch. Refs: #11705 --- components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py | 1 + 1 file changed, 1 insertion(+) diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py b/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py index 0e61ab44f77a..2c08bfbd965d 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py @@ -1084,6 +1084,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 From 2b41fc8b114b1500367b82a3a8a63220a81cbd47 Mon Sep 17 00:00:00 2001 From: Pei Li Date: Wed, 15 Jul 2026 12:59:52 -0700 Subject: [PATCH 05/11] fix(trtllm): guard override path for ConversationParams availability Raise RuntimeError at startup when --conversation-affinity / DYN_ENGINE_CONV_AFFINITY is set on a TRT-LLM build that lacks ConversationParams (<=1.3.0rc20). The existing guard only covered the auto-detection path. Refs: #11705 --- .../src/dynamo/trtllm/request_handlers/handler_base.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/components/src/dynamo/trtllm/request_handlers/handler_base.py b/components/src/dynamo/trtllm/request_handlers/handler_base.py index 1f443664391f..edcfcdbd17e1 100644 --- a/components/src/dynamo/trtllm/request_handlers/handler_base.py +++ b/components/src/dynamo/trtllm/request_handlers/handler_base.py @@ -271,6 +271,15 @@ def __init__(self, config: RequestHandlerConfig): # 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 + 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.encode_client = config.encode_client self.multimodal_processor = config.multimodal_processor self.first_generation = True From 0ca9721a67498bf70fae47a0d82322f1707d987a Mon Sep 17 00:00:00 2001 From: Pei Li Date: Wed, 15 Jul 2026 14:10:31 -0700 Subject: [PATCH 06/11] fix(trtllm): revert startup guard for ConversationParams availability The __init__ guard fires for any HandlerBase constructed with a bare MagicMock() config (truthy conversation_affinity attribute), breaking existing tests that don't pin it. CONVERSATION_PARAMS_AVAILABLE is also False in the CPU test environment, compounding the breakage. Refs: #11705 --- .../src/dynamo/trtllm/request_handlers/handler_base.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/components/src/dynamo/trtllm/request_handlers/handler_base.py b/components/src/dynamo/trtllm/request_handlers/handler_base.py index edcfcdbd17e1..1f443664391f 100644 --- a/components/src/dynamo/trtllm/request_handlers/handler_base.py +++ b/components/src/dynamo/trtllm/request_handlers/handler_base.py @@ -271,15 +271,6 @@ def __init__(self, config: RequestHandlerConfig): # 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 - 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.encode_client = config.encode_client self.multimodal_processor = config.multimodal_processor self.first_generation = True From 9ec0a0e770330bc7f0275b2c66ccdd9621aa1218 Mon Sep 17 00:00:00 2001 From: Pei Li Date: Wed, 15 Jul 2026 14:21:10 -0700 Subject: [PATCH 07/11] fix(trtllm): guard override path against missing ConversationParams API Raise RuntimeError on first request when DYN_ENGINE_CONV_AFFINITY is set on a TRT-LLM build that lacks ConversationParams (<=1.3.0rc20). The existing guard only covered the engine auto-detection path. Guard goes in _generate_locally_impl (not __init__) because CONVERSATION_PARAMS_AVAILABLE is False in the CPU test environment. Pin config.conversation_affinity=False in all MagicMock-based handler helpers so the guard does not fire spuriously in tests. Refs: #11705 --- .../src/dynamo/trtllm/request_handlers/handler_base.py | 9 +++++++++ .../src/dynamo/trtllm/tests/test_trtllm_handler_base.py | 5 +++++ components/src/dynamo/trtllm/tests/utils.py | 1 + 3 files changed, 15 insertions(+) diff --git a/components/src/dynamo/trtllm/request_handlers/handler_base.py b/components/src/dynamo/trtllm/request_handlers/handler_base.py index 1f443664391f..9bc75f0b37f8 100644 --- a/components/src/dynamo/trtllm/request_handlers/handler_base.py +++ b/components/src/dynamo/trtllm/request_handlers/handler_base.py @@ -1147,6 +1147,15 @@ async def _generate_locally_impl( 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)." + ) # Extract dp_rank from request's routing hints for attention DP routing routing = request.get("routing", {}) diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py b/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py index 2c08bfbd965d..667d38990fb8 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py @@ -467,6 +467,7 @@ class TestDeferredAbortGuard: def _make_handler(self) -> HandlerBase: config = MagicMock() config.shutdown_event = None + config.conversation_affinity = False return _ConcreteHandler(config) @pytest.mark.asyncio @@ -588,6 +589,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): @@ -644,6 +646,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): @@ -714,6 +717,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 @@ -797,6 +801,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 diff --git a/components/src/dynamo/trtllm/tests/utils.py b/components/src/dynamo/trtllm/tests/utils.py index aed2cd663d00..a9157020370a 100644 --- a/components/src/dynamo/trtllm/tests/utils.py +++ b/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 From 610fe2650329b39d68a0eba6e96c9208d44a2cb5 Mon Sep 17 00:00:00 2001 From: Pei Li Date: Wed, 15 Jul 2026 15:54:52 -0700 Subject: [PATCH 08/11] test(trtllm): add coverage for DYN_ENGINE_CONV_AFFINITY override path Three tests in test_trtllm_unit.py verify the env var and CLI flag are read correctly and default to False. Two tests in TestConversationAffinity cover the override path behavior: - override=True suppresses dp_rank and forwards conversation_params (regression case for the conv_kwargs gate) - override=True raises RuntimeError when ConversationParams API is missing (guard added in _generate_locally_impl) Refs: #11705 --- .../trtllm/tests/test_trtllm_handler_base.py | 56 +++++++++++++++++++ .../dynamo/trtllm/tests/test_trtllm_unit.py | 21 +++++++ 2 files changed, 77 insertions(+) diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py b/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py index 27f22b8aafb0..b27d1cb46c56 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py @@ -1290,3 +1290,59 @@ 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.""" + monkeypatch.setattr( + "dynamo.trtllm.conversation_affinity.ConversationParams", + _FakeConversationParams, + ) + # Engine detection returns False (no engine-side affinity config). + handler = self._make_handler(conversation_affinity=False) + # Simulate DYN_ENGINE_CONV_AFFINITY=true wired through RequestHandlerConfig. + 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 (guard in _generate_locally_impl).""" + monkeypatch.setattr( + "dynamo.trtllm.request_handlers.handler_base.CONVERSATION_PARAMS_AVAILABLE", + False, + ) + handler = self._make_handler(conversation_affinity=False) + 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() diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_unit.py b/components/src/dynamo/trtllm/tests/test_trtllm_unit.py index e041efdbb4c2..1a377d48f9be 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_unit.py +++ b/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( From 5feb7a5f2bcf50f131dbab46db64823dab758349 Mon Sep 17 00:00:00 2001 From: Pei Li Date: Wed, 15 Jul 2026 16:27:52 -0700 Subject: [PATCH 09/11] refactor(trtllm): fold override guard into lazy affinity init Per review: consolidate the two runtime checks into the existing `if _conversation_affinity is None` block so both the override guard and the engine-detection guard fire exactly once (on first request), then simplify `conv_affinity = self._conversation_affinity`. Update the two override tests to reset `_conversation_affinity = None` before driving the handler so the lazy init path is exercised. Refs: DYN-3570 --- .../trtllm/request_handlers/handler_base.py | 27 +++++++++---------- .../trtllm/tests/test_trtllm_handler_base.py | 18 ++++++++++--- 2 files changed, 27 insertions(+), 18 deletions(-) diff --git a/components/src/dynamo/trtllm/request_handlers/handler_base.py b/components/src/dynamo/trtllm/request_handlers/handler_base.py index f26a95fa45ca..5996eb6c46d9 100644 --- a/components/src/dynamo/trtllm/request_handlers/handler_base.py +++ b/components/src/dynamo/trtllm/request_handlers/handler_base.py @@ -1115,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( @@ -1124,18 +1134,7 @@ async def _generate_locally_impl( "the installed TensorRT-LLM build has no ConversationParams API (requires a " "release newer than 1.3.0rc20)." ) - 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)." - ) + conv_affinity = self._conversation_affinity # Extract dp_rank from request's routing hints for attention DP routing routing = request.get("routing", {}) diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py b/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py index b27d1cb46c56..030d5a0d0dcd 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py @@ -1296,14 +1296,22 @@ 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.""" + 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, ) - # Engine detection returns False (no engine-side affinity config). + monkeypatch.setattr( + "dynamo.trtllm.request_handlers.handler_base.engine_conversation_affinity_enabled", + lambda _: False, + ) handler = self._make_handler(conversation_affinity=False) - # Simulate DYN_ENGINE_CONV_AFFINITY=true wired through RequestHandlerConfig. + # 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, @@ -1326,12 +1334,14 @@ 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 (guard in _generate_locally_impl).""" + 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() From d298b91eb5f903df0cf534885674e61bf3c9317d Mon Sep 17 00:00:00 2001 From: Yuewei Na <248773860+nv-yna@users.noreply.github.com> Date: Wed, 15 Jul 2026 16:33:00 -0700 Subject: [PATCH 10/11] fix(trtllm): honor conversation-affinity override on unified engine path (#11738) Signed-off-by: Yuewei Na Co-authored-by: Yuewei Na --- components/src/dynamo/trtllm/llm_engine.py | 27 ++++++- ...est_trtllm_engine_conversation_affinity.py | 73 +++++++++++++++++++ 2 files changed, 96 insertions(+), 4 deletions(-) diff --git a/components/src/dynamo/trtllm/llm_engine.py b/components/src/dynamo/trtllm/llm_engine.py index e749d55e1654..d2573274e2dd 100644 --- a/components/src/dynamo/trtllm/llm_engine.py +++ b/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, diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_engine_conversation_affinity.py b/components/src/dynamo/trtllm/tests/test_trtllm_engine_conversation_affinity.py index 0ba1291b785a..a5eb66ea35f7 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_engine_conversation_affinity.py +++ b/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,68 @@ 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.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 From cee306a8781ca7eb1ec418bd1503a0723d1fdb26 Mon Sep 17 00:00:00 2001 From: Pei Li Date: Wed, 15 Jul 2026 18:02:08 -0700 Subject: [PATCH 11/11] test(trtllm): fix two GPU test failures after llm_engine override wiring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit test_trtllm_trace_propagation: _make_engine bypasses __init__ but _generate_started now reads _engine_conversation_affinity_override unconditionally — add the missing default (False) alongside the existing _conversation_affinity = False. test_trtllm_engine_conversation_affinity: the override test patches ConversationParams but not CONVERSATION_PARAMS_AVAILABLE in llm_engine, so the guard fires in a real TRT-LLM 1.3.0rc20 environment. Patch llm_engine.CONVERSATION_PARAMS_AVAILABLE = True to match the intent. Refs: DYN-3570 --- .../trtllm/tests/test_trtllm_engine_conversation_affinity.py | 4 ++++ .../src/dynamo/trtllm/tests/test_trtllm_trace_propagation.py | 3 ++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_engine_conversation_affinity.py b/components/src/dynamo/trtllm/tests/test_trtllm_engine_conversation_affinity.py index a5eb66ea35f7..fb697909612a 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_engine_conversation_affinity.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_engine_conversation_affinity.py @@ -205,6 +205,10 @@ async def test_override_suppresses_dp_rank_and_forwards_conversation_params( """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, diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_trace_propagation.py b/components/src/dynamo/trtllm/tests/test_trtllm_trace_propagation.py index e4334ff31b88..29ba8c0139c7 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_trace_propagation.py +++ b/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