From 01da19e22333136d29c09a4237adb92b6d3ea8c0 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sun, 12 Jul 2026 19:23:03 -0700 Subject: [PATCH] [None][feat] ADP conversation router: configurable least-queued placement for new conversations kv_cache_routing_conversation_affinity pins conversations to ranks and round-robins first turns, equalizing per-rank conversation counts. But conversation weight (turn rate, agentic fan-out, prefill length) is not uniform: on an agentic multi-turn workload at concurrency 1000 we measured near-uniform per-rank conversation counts (max/min 1.14x) yet 1.8-3.4x per-rank queue-depth skew, with the coldest ranks dummy-padded in up to 15.7% of iterations while hot ranks queued 24 deep. Add attention_dp_config.kv_cache_routing_new_conv_placement: - round_robin (default): existing behavior, unchanged. - least_queued: place unpinned requests (first turns, conversation-less requests, sticky overflow) on the rank with the fewest live requests, reusing the existing _least_loaded helper. Heavy pinned conversations keep their rank's count high, so new conversations steer away from hot ranks; the in-batch count accumulator spreads bursts valley-first. Sticky returns, soft/hard caps, and the padding invariant are unchanged; routing stays a deterministic pure function of allgathered state, as required by the replicated no-broadcast protocol. Measured on a 4x(TP8, attention-DP) prefill disagg deployment at concurrency 1000 (agentic multi-turn replay, 95% KV block reuse): per-rank queue skew 3.4x -> 2.0x, TTFT p90 -15.5%, p99 -17.2%, KV reuse unchanged, throughput neutral. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/pyexecutor/scheduler/adp_router.py | 21 +++++-- tensorrt_llm/llmapi/llm_args.py | 15 +++++ .../usage/llm_args_golden_manifest.json | 10 ++++ .../_torch/executor/test_adp_router.py | 56 +++++++++++++++++++ 4 files changed, 97 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index acea0242edcb..21fd55971bd6 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -201,6 +201,7 @@ def create( dist=dist, max_sessions=attention_dp_config.kv_cache_routing_max_sessions, fair_share_multiplier=attention_dp_config.kv_cache_routing_fair_share_multiplier, + new_conv_placement=attention_dp_config.kv_cache_routing_new_conv_placement, ) if ( @@ -797,10 +798,11 @@ def _sort_key(req_item): class ConversationAwareADPRouter(ADPRouter): """Pins each conversation to a single attention-DP rank: the first request - of a conversation is round-robined, and every later request with the same + of a conversation is placed by ``new_conv_placement`` (round-robin by + default, or least-queued), and every later request with the same ``conversation_id`` returns to that rank, keeping the conversation's - KV-cache prefix on one rank. Falls back to load-balanced round-robin when no - ``conversation_id`` is present. + KV-cache prefix on one rank. Requests without a ``conversation_id`` fall + back to the same ``new_conv_placement`` policy. """ # Default LRU cap on the conversation->rank map (entries are ~tens of @@ -812,11 +814,15 @@ def __init__( dist: "Distributed", max_sessions: int = DEFAULT_MAX_SESSIONS, fair_share_multiplier: float = 2.0, + new_conv_placement: str = "round_robin", ): super().__init__(dist) self._conv_to_rank: "OrderedDict[str, int]" = OrderedDict() self._max_sessions = max(1, int(max_sessions)) self._fair_share_multiplier = max(1.0, float(fair_share_multiplier)) + self._new_conv_placement = ( + "least_queued" if new_conv_placement == "least_queued" else "round_robin" + ) self._round_robin_cursor = 0 def create_rank_state( @@ -923,8 +929,13 @@ def _next_rr(soft_cap: int) -> int: if rank is None: # First turn of a new conversation, sticky-overflow, or no - # conversation_id -> round-robin spread under the soft cap. - rank = _next_rr(expected_num_active_requests) + # conversation_id -> spread under the soft cap: round-robin + # (count-uniform) or least-queued (steers away from ranks kept + # busy by heavy pinned conversations). + if self._new_conv_placement == "least_queued": + rank = _least_loaded(expected_num_active_requests) + else: + rank = _next_rr(expected_num_active_requests) if conv_id is not None and conv_id not in self._conv_to_rank: # Bind this new conversation to its first-turn rank. self._record_target_rank(conv_id, rank) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 6a63ad9c8c49..98294d7566ca 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1412,6 +1412,21 @@ class AttentionDpConfig(StrictBaseModel): "routing. The oldest conversations are evicted once more than this many " "are tracked, bounding memory on long-running servers. Only used when " "kv_cache_routing_conversation_affinity is True.") + kv_cache_routing_new_conv_placement: Literal[ + "round_robin", "least_queued"] = Field( + default="round_robin", + description= + "Placement policy in conversation-affinity routing for requests " + "with no pinned rank yet (first turn of a conversation, requests " + "without a conversation_id, sticky overflow). 'round_robin' " + "(default) equalizes per-rank conversation counts. 'least_queued' " + "places them on the rank with the fewest live requests instead: " + "per-conversation load (turn rate, fan-out, prefill length) is " + "not uniform, so count-uniform round-robin can leave some ranks " + "with deep queues while others idle; steering new conversations " + "by queue depth evens that out and cuts tail TTFT. Existing " + "conversation->rank pins are unaffected. Only used when " + "kv_cache_routing_conversation_affinity is True.") @model_validator(mode='after') def validate_attention_dp_config(self) -> 'AttentionDpConfig': diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 023d4f07f4c1..41abe20dcbdb 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -87,6 +87,16 @@ "kind": "value", "path": "attention_dp_config.kv_cache_routing_max_sessions" }, + { + "allowed_values": [ + "round_robin", + "least_queued" + ], + "annotation": "Literal['round_robin', 'least_queued']", + "converter": "", + "kind": "categorical", + "path": "attention_dp_config.kv_cache_routing_new_conv_placement" + }, { "allowed_values": [], "annotation": "", diff --git a/tests/unittest/_torch/executor/test_adp_router.py b/tests/unittest/_torch/executor/test_adp_router.py index 4c3726578e0d..3d2c418beec8 100644 --- a/tests/unittest/_torch/executor/test_adp_router.py +++ b/tests/unittest/_torch/executor/test_adp_router.py @@ -1247,6 +1247,62 @@ def test_factory_selects_conversation_router(self): router = ADPRouter.create(dist=_mock_dist(), kv_cache_manager=None, attention_dp_config=cfg) assert isinstance(router, ConversationAwareADPRouter) assert router._max_sessions == 8 + # A mocked (non-string) placement value must fall back to round_robin. + assert router._new_conv_placement == "round_robin" + + @staticmethod + def _lq_router(tp_size=4): + return ConversationAwareADPRouter( + dist=_mock_dist(tp_size=tp_size), new_conv_placement="least_queued" + ) + + def test_least_queued_places_new_conversation_on_least_loaded_rank(self): + router = self._lq_router() + pos = self._route( + router, self._states(4, active=[5, 1, 3, 4]), [_make_conv_request_item(1, "A")] + ) + assert pos[1] == 1 + assert router._conv_to_rank["A"] == 1 + + def test_least_queued_fills_valleys_within_one_batch(self): + """The shared count accumulator spreads a burst valley-first: + [3, 0, 2, 3] + 4 new conversations ends level at [3, 3, 3, 3].""" + items = [_make_conv_request_item(i, f"c{i}") for i in range(4)] + pos = self._route(self._lq_router(), self._states(4, active=[3, 0, 2, 3]), items) + placed = [sum(1 for v in pos.values() if v == r) for r in range(4)] + assert placed == [0, 3, 1, 0] + + def test_least_queued_sticky_returns_unaffected(self): + """A later turn returns to its pinned rank even when it is the busiest.""" + router = self._lq_router() + home = self._route( + router, self._states(4, active=[2, 0, 1, 1]), [_make_conv_request_item(1, "A")] + )[1] + assert home == 1 + pos = self._route( + router, self._states(4, active=[0, 9, 0, 0]), [_make_conv_request_item(2, "A")], cap=100 + ) + assert pos[2] == home + + def test_least_queued_routing_is_deterministic_across_ranks(self): + convs = ["A", "B", None, "A", "C", None, "B"] + + def run(): + items = [_make_conv_request_item(i, convs[i]) for i in range(len(convs))] + return self._route(self._lq_router(), self._states(4, active=[2, 5, 0, 1]), items) + + assert run() == run() + + def test_new_conv_placement_config(self): + """Factory forwards the knob; unknown values fall back to round_robin.""" + cfg = MagicMock() + cfg.kv_cache_routing_conversation_affinity = True + cfg.kv_cache_routing_max_sessions = 8 + cfg.kv_cache_routing_new_conv_placement = "least_queued" + router = ADPRouter.create(dist=_mock_dist(), kv_cache_manager=None, attention_dp_config=cfg) + assert router._new_conv_placement == "least_queued" + bad = ConversationAwareADPRouter(dist=_mock_dist(tp_size=4), new_conv_placement="banana") + assert bad._new_conv_placement == "round_robin" def test_factory_default_when_disabled(self): cfg = MagicMock()