[None][fix] cold-start warmup for KV-aware ADP router (feat/deepseek_v4) - #14388
Merged
Conversation
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>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49602 [ run ] triggered by Bot. Commit: |
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>
Pre-commit collapsed the multi-line ``KVCacheAwareADPRouter(...)`` helper in ``TestKVCacheAwareADPRouterWarmup._make_router`` to a single line. Applying the formatter's preferred form so the hook doesn't keep re-staging the file. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
1 similar comment
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49619 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49621 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49619 [ run ] completed with state |
Collaborator
|
PR_Github #49602 [ run ] completed with state |
lancelly
added a commit
to lancelly/TensorRT-LLM
that referenced
this pull request
May 21, 2026
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>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49664 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49664 [ run ] completed with state
|
Closed
2 tasks
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>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49824 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49824 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49859 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49859 [ run ] completed with state |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Port of #14307 onto
feat/deepseek_v4. 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.This started as a straight cherry-pick of the original commit (with a trivial
dataclassesimport-line conflict resolution, sincefeat/deepseek_v4hasn't picked up unrelated main-branch extras).Problem
With
match_rate_threshold=0.1the 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: