[None][fix] adp_router: salt-aware probe_prefix_match_length on v1/v2 KV cache backends - #14417
Closed
lancelly wants to merge 16 commits into
Closed
[None][fix] adp_router: salt-aware probe_prefix_match_length on v1/v2 KV cache backends#14417lancelly wants to merge 16 commits into
lancelly wants to merge 16 commits into
Conversation
Signed-off-by: Yueh-Ting Chen <yueh.ting.chen@gmail.com> (cherry picked from commit 4baeec3) Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
Build #39108 of PR NVIDIA#14351 failed 7 tests in tests/unittest/disaggregated/ test_router.py with: TypeError: RootBlock.make_key() takes 1 positional argument but 2 were given This PR collapses the prior (lora_task_id, cache_salt_id) pair into a single TreeTaskId seed for the radix tree, but two callers on the base feat/deepseek_v4 branch were not migrated: - tensorrt_llm/serve/router.py — two call sites still pass (None, cache_salt_id) to V2RootBlock.make_key - tests/unittest/disaggregated/test_router.py — two call sites still pass cache_salt_id as the 4th positional argument to sequence_to_blockchain_keys Fix by computing the seed via _KVCache._make_tree_task_id(lora_task_id, cache_salt_id) and passing it as the single TreeTaskId argument. The helper returns lora_task_id unchanged when cache_salt_id is None so the unsalted fast path is unaffected; with a salt it returns the SHA-256 digest that binds the two fields together. Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
…4140) Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
After the ReuseScope refactor cherry-picked into this stack, ``RootBlock.make_key`` takes a single ``ReuseScope`` instead of the legacy ``(lora_task_id, cache_salt_id)`` pair, and ``sequence_to_blockchain_keys`` likewise drops the trailing ``cache_salt_id`` argument in favour of a ``ReuseScope`` positional. Two callers on the base feat/deepseek_v4 branch still used the legacy shapes and would trip the seven disaggregated/test_router.py tests that exercise the V2 hash path: - tensorrt_llm/serve/router.py — two call sites passed ``(None, cache_salt_id)`` to V2RootBlock.make_key - tests/unittest/disaggregated/test_router.py — four call sites passed raw ``None`` or a trailing ``cache_salt_id`` to sequence_to_blockchain_keys Migrate every call to ``ReuseScope(salt=cache_salt_id)`` (or ``ReuseScope()`` for the unsalted path) so the V2 prefix-hash flow matches what the radix tree now expects. Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
After NVIDIA#14140 (cherry-picked in NVIDIA#14353) RootBlock no longer exposes ``lora_task_id`` / ``cache_salt_id`` directly — both fields are folded into a ``ReuseScope`` NamedTuple attached as ``root.reuse_scope``. ``KVCacheEventManager._root_attrs_from_root_block`` was still reading the legacy attribute names via ``getattr(..., None)``, so every emitted event silently received ``(None, None)`` and the V1-compat hash collapsed to the same value for any LoRA/salt request. Downstream Dynamo routing depends on those hashes, so this regression would materially degrade prefix cache hit rate for any request carrying a LoRA task id or a cache salt. Fix ``_root_attrs_from_root_block`` to prefer ``root.reuse_scope`` and fall back to the legacy attributes (keeps any in-flight pre-refactor RootBlock instances working). Also: - Update the now-stale ``test_v2_root_key_distinguishes_lora_from_cache_salt_id`` to use the new single-arg ``RootBlock.make_key(ReuseScope(...))`` API. - Extend ``_FakeRootBlock`` with an optional ``reuse_scope`` kwarg so tests can mimic both the pre- and post-refactor RootBlock shape. - Add a regression test ``test_v2_kv_cache_event_manager_v1_hash_reads_root_reuse_scope`` that asserts (a) a ReuseScope-shaped root and a legacy root with the same lora/salt produce identical V1 event hashes, (b) different scopes still produce different hashes (no silent collapse), and (c) an empty ReuseScope matches a legacy root with no salt/lora. Addresses reviewer feedback on NVIDIA#14351 (peihu-nv). Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
# Conflicts: # tensorrt_llm/_torch/pyexecutor/resource_manager.py
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>
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>
Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com> (cherry picked from commit f9968c6) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
… KV cache backends Make KVCacheAwareADPRouter work on KVCacheManagerV2 and become salt-aware on KVCacheManager (v1). Before this change: - v2 KVCacheManagerV2 had no probe_prefix_match_length, so any v2 backend (including DSV4) would AttributeError on adp_router.gather_prefix_matches with enable_kv_cache_aware_routing=true. - v1 had a probe but never forwarded cache_salt_id onto the dummy LlmRequest, so salted requests were probed against the salt=None namespace while _create_kv_cache wrote to the salted one -- router saw the wrong match length. Add a v2 KVCacheManagerV2.probe_prefix_match_length that calls the new KVCacheManager.probe_reuse (added in PR NVIDIA#14395) with ReuseScope(lora_id, salt). Extend the v1 probe to take cache_salt_id and pass it through to CppLlmRequest so the C++ buildBlockKeys() pulls the right namespace. Update adp_router.gather_prefix_matches to read req.cache_salt_id and forward it. Both backends now probe the same namespace their _create_kv_cache paths write to: (lora_task_id, cache_salt_id). Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
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
The fix commit (
adp_router: salt-aware probe_prefix_match_length on v1/v2 KV cache backends) makesKVCacheAwareADPRouterwork onKVCacheManagerV2and become salt-aware onKVCacheManager(v1):KVCacheManagerV2had noprobe_prefix_match_length, so any v2 backend (including DSV4) wouldAttributeErroronadp_router.gather_prefix_matcheswithenable_kv_cache_aware_routing=true. Added an adapter that calls the newKVCacheManager.probe_reuse()API from [None][feat] KV reuse probe #14395 withReuseScope(lora_id, salt).cache_salt_idonto the dummyLlmRequest, so salted requests were probed against thesalt=Nonenamespace while_create_kv_cachewrote to the salted one — router saw the wrong match length. v1 probe now takescache_salt_idand forwards it ontoCppLlmRequestso the C++buildBlockKeys()picks the right reuse namespace.KVCacheAwareADPRouter.gather_prefix_matchesnow readsreq.cache_salt_idand forwards it.Both backends now probe the same
(lora_task_id, cache_salt_id)namespace that their_create_kv_cachepaths write to.Dependencies
This branch carries merge commits from four in-flight PRs, in this order:
Support cache_salt_id in KV cache v2 manager(introducesReuseScope's salt field)cold-start warmup for KV-aware ADP router(not load-bearing for this fix; only present because this is a test branch)Refactor salting support for KVCacheManagerV2(finalReuseScope(lora_id=…, salt=…)shape)Add KV cache reuse probe(the newprobe_reuse()API this fix wraps)The fix commit logically depends on #14351, #14353, #14395; it is independent of #14388.
Test plan
enable_kv_cache_aware_routing: trueon GEN attn_dp (previously had to befalsebecause v2AttributeError: probe_prefix_match_length)tests/unittest/_torch/executor/test_kvcache_aware_router.pypasses on both v1 and v2 backends with and withoutcache_salt_id