[None][feat] ADP conversation router: configurable least-queued placement for new conversations

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)

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':

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": "<class 'int'>",

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()