[None][fix] cold-start warmup for KV-aware ADP router by lancelly · Pull Request #14307 · NVIDIA/TensorRT-LLM

lancelly · 2026-05-19T12:20:11Z

Summary

Cold-start mitigation for the KV-aware ADP router: round-robin the first relaxed requests across all ranks so every rank caches the (assumed shared) system prompt before cache-affinity scoring would otherwise pin all traffic to the first warm rank.

Problem

With match_rate_threshold=0.1 the cache-affinity gate stays ON for any system prompt longer than ~10% of a request. Once a single rank caches that prompt, the (req_tokens - match_len) term dominates the load term:

Low-traffic sequential arrivals: each new request pins to the warm rank; other ranks never see the prompt.
Eventual traffic spike: cold ranks pay the full prefill cost re-computing the same shared prompt.

Fix

KVCacheAwareADPRouter tracks _pending_warmup_ranks: Set[int]:

Initialised to {0, ..., tp_size-1} only when cold_start_warmup=True; otherwise an empty set.
In route_requests, any target_dp_rank (explicit strict or synthesised warmup) discards its rank from the set. Relaxed requests with no explicit target synthesise min(_pending_warmup_ranks) while the set is non-empty.
Strict pre-assignments count toward warmup naturally; cap-saturated ranks also discard so the synthesiser never loops on a busy rank.
Once the set is empty, warmup is dormant forever and routing is pure scoring.

Opt-in (default False)

The behaviour is gated by AttentionDpConfig.kv_cache_routing_cold_start_warmup (default False). The warmup is a win for workloads that share a long system prompt, but for diverse-prompt workloads it can scatter requests onto cold ranks and bypass cache-affinity scoring entirely for the first tp_size requests — wasting prefill work the scoring path would otherwise consolidate. Default False preserves pre-warmup routing; users who hit cold-start pinning opt in explicitly.

Test layout:

TestKVCacheAwareADPRouterWarmup constructs its router with cold_start_warmup=True to exercise the new code path, plus a new test_disabled_by_default asserting the empty pending set when the flag is absent.
Pre-existing scoring tests no longer need the _pending_warmup_ranks = set() opt-out (default is already off).

With match_rate_threshold=0.1 the cache-affinity gate stays ON for any system prompt longer than ~10% of a request. Once a single rank caches the prompt, the (req_tokens - match_len) term dominates the load term and subsequent single-request batches keep landing on that warm rank until the fair-share cap finally evicts it. Other ranks never warm up and pay the full prefill cost when traffic eventually spills over. Track _pending_warmup_ranks (Set[int]) on the router, initialised in __init__ to {0, ..., tp_size-1}. In route_requests, any target_dp_rank -- explicit strict or synthesised warmup -- discards its rank from the set; relaxed requests with no explicit target synthesise min(_pending_warmup_ranks) so warmup spreads the shared prompt across every rank within roughly the first tp_size requests. Strict pre-assignments count toward the warmup naturally and cap saturation also discards so the synthesiser never loops on a busy rank. Once the set empties the warmup path is dormant forever and routing reverts to pure scoring. Tests that exercise pure scoring behaviour opt out via router._pending_warmup_ranks = set() -- 3 tests in test_kvcache_aware_router.py and the shared _make_router helper in test_adp_router.py. Scope: single shared system prompt. Multi-prompt warmup is a follow-up. Signed-off-by: Lanyu Liao <lancelly@users.noreply.github.com>

coderabbitai · 2026-05-19T12:44:03Z

📝 Walkthrough

Walkthrough

The PR adds cold-start warmup behavior to the KVCacheAwareADPRouter. It tracks pending tensor-parallel ranks that need initial requests, assigns new unrouted requests to the smallest unwarmed rank before cache-aware scoring, and removes ranks from the warmup set once targeted. Existing tests disable warmup to validate pure scoring. New tests comprehensively validate warmup initialization, dispatch, and fallback behavior.

Changes

Warmup-driven routing for ADP

Layer / File(s)	Summary
Warmup state initialization `tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py`	`Set` type is imported, and `_pending_warmup_ranks` field is initialized on `KVCacheAwareADPRouter` with all tensor-parallel ranks to track which ranks need cold-start requests.
Warmup-aware routing in route_requests `tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py`	The request assignment loop in `route_requests` is updated to honor strictly targeted `attention_dp_rank` requests first, then assign unrouted requests to the smallest pending warmup rank and immediately remove that rank from the warmup set, enabling gradual seeding across all ranks before cache-aware scoring resumes.
Test warmup state isolation `tests/unittest/_torch/executor/test_adp_router.py`, `tests/unittest/_torch/executor/test_kvcache_aware_router.py`	Test setup and three existing routing tests explicitly clear `_pending_warmup_ranks` to disable warmup behavior, isolating test assertions to measure pure cache and load-balance scoring without rank pre-assignment.
Comprehensive warmup test suite `tests/unittest/_torch/executor/test_kvcache_aware_router.py`	New `TestKVCacheAwareADPRouterWarmup` test class validates warmup initialization, round-robin dispatch across pending ranks, rank removal after targeting, interaction with strict and relaxed `target_dp_rank` requests, capacity saturation, and fallback to scoring when warmup ranks are exhausted.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title '[None][fix] cold-start warmup for KV-aware ADP router' clearly summarizes the main change: implementing a cold-start warmup feature for the KV-cache-aware ADP router.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	PR description covers the problem, solution, testing approach, and opt-in design clearly with good technical depth.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

coderabbitai

🧹 Nitpick comments (1)

tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py (1)

26-26: 💤 Low value

Consider using built-in set instead of typing.Set.

Per coding guidelines, prefer built-in types (set) over legacy typing.Set since Python 3.10+.

🔧 Suggested change

-from typing import TYPE_CHECKING, Dict, List, Set, Tuple
+from typing import TYPE_CHECKING, Dict, List, Tuple

And at line 455:

-    self._pending_warmup_ranks: Set[int] = set(range(self.dist.tp_size))
+    self._pending_warmup_ranks: set[int] = set(range(self.dist.tp_size))

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py` at line 26, Replace
usage of typing.Set with the built-in set: remove Set from the typing import
list in adp_router.py and update all annotations that use Set[...] to use the
modern built-in syntax set[...] (e.g., change Set[int] to set[int]). Search for
occurrences in this module (e.g., any functions or variables annotated with Set
like in router logic) and update their annotations accordingly; keep other
typing imports (Dict, List, Tuple, TYPE_CHECKING) as-is.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py`:
- Line 26: Replace usage of typing.Set with the built-in set: remove Set from
the typing import list in adp_router.py and update all annotations that use
Set[...] to use the modern built-in syntax set[...] (e.g., change Set[int] to
set[int]). Search for occurrences in this module (e.g., any functions or
variables annotated with Set like in router logic) and update their annotations
accordingly; keep other typing imports (Dict, List, Tuple, TYPE_CHECKING) as-is.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 9a744c6b-eb4e-43fe-82b5-478a7a0992ba

📥 Commits

Reviewing files that changed from the base of the PR and between 989671b and b3c5632.

📒 Files selected for processing (3)

tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py
tests/unittest/_torch/executor/test_adp_router.py
tests/unittest/_torch/executor/test_kvcache_aware_router.py

lancelly · 2026-05-19T14:50:47Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-19T15:00:37Z

PR_Github #49216 [ run ] triggered by Bot. Commit: b3c5632 Link to invocation

tensorrt-cicd · 2026-05-19T22:33:57Z

PR_Github #49216 [ run ] completed with state SUCCESS. Commit: b3c5632
/LLM/main/L0_MergeRequest_PR pipeline #38891 completed with status: 'SUCCESS'

CI Report

Link to invocation

Cold-start round-robin warmup is now controlled by ``AttentionDpConfig.kv_cache_routing_cold_start_warmup`` (default False). The previous behaviour always seeded ``_pending_warmup_ranks`` from ``dist.tp_size``, which forces the first tp_size relaxed requests across ranks regardless of cache affinity. That helps the assumed "shared system prompt" case but can scatter requests with partially overlapping prompts onto cold ranks, wasting prefill. Default False preserves pre-warmup routing for diverse-prompt workloads; users hitting cold-start pinning opt in explicitly. Tests in TestKVCacheAwareADPRouterWarmup now pass ``cold_start_warmup= True`` and a new ``test_disabled_by_default`` covers the default. The three opt-out hacks in scoring tests (manual ``_pending_warmup_ranks = set()``) are no longer needed and are removed. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>

Pure cosmetic sync with feat/deepseek_v4 NVIDIA#14388: move the inline comment in ``test_empty_pending_disables_warmup`` to a standalone line above the assignment, matching the form there. No behaviour change. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>

The ruff-legacy hook (D200) fired on 24 docstrings in ``tensorrt_llm/llmapi/llm_args.py`` that follow the multi-line form for a single line of text: \"\"\" Configuration for X. \"\"\" These predate this PR but the gate only checks files touched by a change, so this PR was the one to surface them. Collapsing to the one-line form satisfies D200 without touching semantics or the baseline ratchet. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>

lancelly · 2026-05-21T07:19:34Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-21T07:26:14Z

PR_Github #49624 [ run ] triggered by Bot. Commit: 7443271 Link to invocation

tensorrt-cicd · 2026-05-21T12:56:10Z

PR_Github #49624 [ run ] completed with state SUCCESS. Commit: 7443271
/LLM/main/L0_MergeRequest_PR pipeline #39244 completed with status: 'SUCCESS'

CI Report

Link to invocation

A saturated warmup pick was being unconditionally removed from ``_pending_warmup_ranks``, so a rank could be marked as warmed without ever receiving a request and never having its shared prefix seeded. Move the ``discard`` under the scheduled branch so only ranks that actually received a request leave the pending set. Also use ``dist.mapping.dp_size`` instead of ``dist.tp_size`` for the initial pending set; the values match under ``enable_attention_dp=True`` but ``dp_size`` reads more clearly at the call site. Addresses reviewer feedback on PR NVIDIA#14307. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>

lancelly · 2026-05-22T02:35:51Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-22T02:42:18Z

PR_Github #49819 [ run ] triggered by Bot. Commit: 7443271 Link to invocation

A saturated warmup pick was being unconditionally removed from ``_pending_warmup_ranks``, so a rank could be marked as warmed without ever receiving a request and never having its shared prefix seeded. Move the ``discard`` under the scheduled branch so only ranks that actually received a request leave the pending set. Also use ``dist.mapping.dp_size`` instead of ``dist.tp_size`` for the initial pending set; the values match under ``enable_attention_dp=True`` but ``dp_size`` reads more clearly at the call site. Addresses reviewer feedback on PR NVIDIA#14307. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>

lancelly · 2026-05-22T02:42:56Z

/bot run --disable-fail-fast

lancelly · 2026-05-22T02:43:07Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-22T02:48:21Z

PR_Github #49821 [ run ] triggered by Bot. Commit: ded52ac Link to invocation

tensorrt-cicd · 2026-05-22T02:49:08Z

PR_Github #49823 [ run ] triggered by Bot. Commit: ded52ac Link to invocation

tensorrt-cicd · 2026-05-22T02:49:13Z

PR_Github #49821 [ run ] completed with state ABORTED. Commit: ded52ac

Link to invocation

tensorrt-cicd · 2026-05-22T02:49:36Z

PR_Github #49819 [ run ] completed with state ABORTED. Commit: 7443271

Link to invocation

lancelly · 2026-05-22T05:45:15Z

/bot run --disable-fail-fast

QiJune

LGTM

tensorrt-cicd · 2026-05-22T05:50:43Z

PR_Github #49858 [ run ] triggered by Bot. Commit: ded52ac Link to invocation

tensorrt-cicd · 2026-05-22T05:56:15Z

PR_Github #49823 [ run ] completed with state ABORTED. Commit: ded52ac

Link to invocation

lancelly · 2026-05-22T06:00:10Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-22T06:06:44Z

PR_Github #49862 [ run ] triggered by Bot. Commit: ded52ac Link to invocation

tensorrt-cicd · 2026-05-22T06:10:16Z

PR_Github #49858 [ run ] completed with state ABORTED. Commit: ded52ac

Link to invocation

lancelly · 2026-05-22T07:09:57Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-22T07:16:36Z

PR_Github #49879 [ run ] triggered by Bot. Commit: ded52ac Link to invocation

tensorrt-cicd · 2026-05-22T07:20:04Z

PR_Github #49862 [ run ] completed with state ABORTED. Commit: ded52ac

Link to invocation

tensorrt-cicd · 2026-05-22T07:48:09Z

PR_Github #49879 [ run ] completed with state SUCCESS. Commit: ded52ac
/LLM/main/L0_MergeRequest_PR pipeline #39459 completed with status: 'SUCCESS'

CI Report

Link to invocation

Signed-off-by: Lanyu Liao <lancelly@users.noreply.github.com> Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Co-authored-by: Lanyu Liao <lancelly@users.noreply.github.com>

github-actions Bot assigned lancelly May 19, 2026

lancelly force-pushed the router_optimization branch 2 times, most recently from 1c80f10 to 0fc5a32 Compare May 19, 2026 12:35

lancelly marked this pull request as ready for review May 19, 2026 12:40

lancelly requested a review from a team as a code owner May 19, 2026 12:40

lancelly requested a review from Tabrizian May 19, 2026 12:41

lancelly force-pushed the router_optimization branch from 0fc5a32 to b3c5632 Compare May 19, 2026 12:42

coderabbitai Bot reviewed May 19, 2026

View reviewed changes

lishicheng1996-nv approved these changes May 21, 2026

View reviewed changes

lancelly mentioned this pull request May 21, 2026

[None][fix] cold-start warmup for KV-aware ADP router (feat/deepseek_v4) #14388

Merged

lancelly requested a review from a team as a code owner May 21, 2026 06:39

lancelly requested a review from Superjomn May 21, 2026 06:39

lancelly added 2 commits May 21, 2026 00:04

QiJune reviewed May 21, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py Outdated

QiJune reviewed May 21, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py

QiJune approved these changes May 22, 2026

View reviewed changes

lancelly merged commit e796f16 into NVIDIA:main May 22, 2026
7 checks passed

Uh oh!

Conversation

lancelly commented May 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Problem

Fix

Opt-in (default False)

Uh oh!

coderabbitai Bot commented May 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

lancelly commented May 19, 2026

Uh oh!

tensorrt-cicd commented May 19, 2026

Uh oh!

tensorrt-cicd commented May 19, 2026

Uh oh!

lancelly commented May 21, 2026

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

Uh oh!

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

lancelly commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

lancelly commented May 22, 2026

Uh oh!

lancelly commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

lancelly commented May 22, 2026

Uh oh!

QiJune left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

lancelly commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

lancelly commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

lancelly commented May 19, 2026 •

edited

Loading

coderabbitai Bot commented May 19, 2026 •

edited

Loading