[None][feat] adp_router: salt-aware probe_prefix_match_length on v1/v2 KV cache backends - #14435
Merged
lancelly merged 1 commit intoMay 22, 2026
Conversation
… 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> (cherry picked from commit a5ee47b) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
lancelly
changed the title
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49829 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49829 [ run ] completed with state |
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 10, 2026
…2 KV cache backends (NVIDIA#14435) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 11, 2026
…2 KV cache backends (NVIDIA#14435) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 12, 2026
…2 KV cache backends (NVIDIA#14435) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 26, 2026
…2 KV cache backends (NVIDIA#14435) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
Shixiaowei02
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 29, 2026
…2 KV cache backends (NVIDIA#14435) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
lancelly
added a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 29, 2026
…2 KV cache backends (NVIDIA#14435) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com> Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com>
Shixiaowei02
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 30, 2026
…2 KV cache backends (NVIDIA#14435) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com> Signed-off-by: Liao Lanyu <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
Make
KVCacheAwareADPRouterwork onKVCacheManagerV2(currentlyAttributeErrors) and become salt-aware onKVCacheManagerv1 (currently probes the wrong namespace for salted requests).Before this change:
probe_prefix_match_length, so any v2 backend (DSV4 in particular) crashes withAttributeErrorinadp_router.gather_prefix_matchesas soon asattention_dp_config.enable_kv_cache_aware_routing=true.cache_salt_idonto the dummyLlmRequest, so salted requests were probed against thesalt=Nonenamespace while_create_kv_cachewrote to the salted one. Router saw match_len=0 for everything that was actually warm.After this change, both backends probe the same namespace their
_create_kv_cachepaths write to:(lora_task_id, cache_salt_id).What's in this PR
tensorrt_llm/_torch/pyexecutor/resource_manager.py:KVCacheManagerV2.probe_prefix_match_length(probe_tokens, lora_task_id, *, cache_salt_id=None)that delegates toKVCacheManager.probe_reuse(ReuseScope(lora_task_id, cache_salt_id))(the probe API added in [None][feat] KV reuse probe #14395).probe_prefix_match_lengthto acceptcache_salt_idand pass it through to the C++LlmRequestsobuildBlockKeys()resolves the correct namespace.tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py:gather_prefix_matches, readreq.cache_salt_id(defaulting toNone) and forward it to the probe call.