diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index f29073a10561..90ec6573cc8d 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -677,11 +677,19 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], device='cpu') self.blocks_per_window = blocks_per_window - def probe_prefix_match_length(self, input_tokens, lora_task_id=None): + def probe_prefix_match_length(self, + input_tokens, + lora_task_id=None, + cache_salt_id=None): """Probe the KV cache radix tree for prefix match length. Returns the number of prefix tokens already cached on this rank. Used by KVCacheAwareADPRouter for cache-aware routing. + + ``cache_salt_id`` scopes the probe to the same per-request + namespace that ``_create_kv_cache`` uses; without it, salted + requests would be probed against the salt=None namespace and the + router would see the wrong match length. """ if not self.enable_block_reuse: return 0 @@ -698,12 +706,16 @@ def probe_prefix_match_length(self, input_tokens, lora_task_id=None): LlmRequest as CppLlmRequest block_key = BlockKey(tokens=input_tokens, lora_task_id=lora_task_id) unique_tokens = block_key.unique_tokens + # buildBlockKeys() on the C++ side pulls cache_salt_id off the + # request, so dummy_req must carry it for the probe namespace to + # match the real create_kv_cache path. dummy_req = CppLlmRequest(request_id=0, max_new_tokens=0, input_tokens=input_tokens, sampling_config=SamplingConfig(), is_streaming=False, - lora_task_id=lora_task_id) + lora_task_id=lora_task_id, + cache_salt_id=cache_salt_id) summary = self.impl.analyze_prefix_reuse(unique_tokens, dummy_req) return summary.reusable_blocks_all * self.tokens_per_block @@ -2585,6 +2597,33 @@ def is_request_active(self, request_id: int) -> bool: kv_cache = self.kv_cache_map.get(request_id) return kv_cache is not None and kv_cache.is_active + def probe_prefix_match_length(self, + input_tokens, + lora_task_id=None, + cache_salt_id=None): + """Probe the v2 KV cache radix tree for prefix match length. + + Returns the number of prefix tokens already cached on this rank, + without acquiring page ownership. Mirrors the v1 KVCacheManager + adapter so KVCacheAwareADPRouter works on both backends. + + ``cache_salt_id`` (and ``lora_task_id``) must match the values + used by ``_create_kv_cache`` (which constructs + ``ReuseScope(lora_id=lora_task_id, salt=cache_salt_id)``). + Otherwise the probe queries the wrong reuse namespace and the + router sees an incorrect match length. + """ + if not self.enable_block_reuse: + return 0 + if not input_tokens: + return 0 + from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import \ + ReuseScope + return self.impl.probe_reuse( + ReuseScope(lora_id=lora_task_id, salt=cache_salt_id), + input_tokens, + ) + def _effective_draft_len(self, req: LlmRequest) -> int: """Draft token length to use for next-step KV capacity calculation. diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index 48494af7a79d..96a2a484da4c 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -424,7 +424,14 @@ def gather_prefix_matches( probe_tokens = input_tokens[:-1] if len(input_tokens) > 1 else [] lora_config = getattr(req, "lora_config", None) lora_task_id = lora_config.task_id if lora_config is not None else None - match_len = self.kv_cache_manager.probe_prefix_match_length(probe_tokens, lora_task_id) + # cache_salt_id scopes block reuse on both v1 and v2 backends; + # passing None for non-salted requests is a no-op. + cache_salt_id = getattr(req, "cache_salt_id", None) + match_len = self.kv_cache_manager.probe_prefix_match_length( + probe_tokens, + lora_task_id, + cache_salt_id=cache_salt_id, + ) local_matches.extend([req_item.id, match_len]) all_data = self.dist.tp_allgather(local_matches)