[#18156][fix] Carry lora_id and share the cache_salt derivation in the KV-aware router

tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py

@@ -12,7 +12,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import hashlib
 import math
 import os
 import sys
@@ -58,7 +57,10 @@
 from tensorrt_llm.llmapi.llm_args import KvCacheConfig, KVEventsConfig
 from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
-from tensorrt_llm.runtime.kv_cache_hash import get_effective_kv_cache_event_hash_algo
+from tensorrt_llm.runtime.kv_cache_hash import (
+    get_cache_salt_id,
+    get_effective_kv_cache_event_hash_algo,
+)
 from tensorrt_llm.runtime.kv_cache_manager_v2 import (
     _KV_CACHE_ITERATION_STATS_DELTA_FIELDS,
     BAD_PAGE_INDEX,
@@ -6409,14 +6411,16 @@ def _derive_reuse_salt(cache_salt: str | None) -> int | None:
         """Derive ``ReuseScope.salt`` (int|None) from the ``cache_salt`` string.
 
         Deterministic so the same string yields the same reuse namespace across
-        processes (matches C++ blockKey hashing on cacheSalt). Shared by cache
-        creation, prefetch, and the cache-aware router probe so they all hit the
-        same radix-tree namespace.
+        processes. Delegates to ``kv_cache_hash.get_cache_salt_id`` -- the single
+        source of truth shared with the serve-layer KV-cache-aware router -- so
+        cache creation, prefetch, the cache-aware router probe, and the router's
+        recomputed event hashes all key the same radix-tree namespace. (The C++
+        V1 ``BlockKeyHasher`` is separate: it mixes
+        ``std::hash<std::string>(cacheSalt)`` directly, not this id.)
         """
         if cache_salt is None:
             return None
-        digest = hashlib.sha256(cache_salt.encode("utf-8")).digest()
-        return int.from_bytes(digest[:8], "little")
+        return get_cache_salt_id(cache_salt)
 
     def _create_kv_cache(
         self,

tensorrt_llm/inputs/utils.py

@@ -32,8 +32,7 @@
                                           _safe_aiohttp_get, _safe_request_get)
 from tensorrt_llm.inputs.media_io import \
     convert_image_mode as convert_image_mode
-from tensorrt_llm.inputs.multimodal import (MultimodalServerConfig,
-                                            default_hasher)
+from tensorrt_llm.inputs.multimodal import MultimodalServerConfig
 from tensorrt_llm.inputs.multimodal_data import \
     BaseModalityData as BaseModalityData
 from tensorrt_llm.inputs.multimodal_data import VideoData as VideoData
@@ -1181,14 +1180,3 @@ def convert_to_conversation_message(
         inputs.append(input)
 
     return inputs
-
-
-def get_cache_salt_id(cache_salt: str) -> int:
-    b = cache_salt.encode("utf-8")
-    h = default_hasher(b).digest(length=8)
-    cache_salt_id = int.from_bytes(h, "little", signed=False)
-    if cache_salt_id < 0 or cache_salt_id >= (1 << 64):
-        raise ValueError(
-            f"cache_salt_id must be in [0, 2**64 - 1], got {cache_salt_id}.")
-
-    return cache_salt_id

tensorrt_llm/runtime/kv_cache_hash.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import hashlib
 from collections.abc import Sequence
 from typing import Optional
 
@@ -46,11 +47,26 @@ def truncate_sha256_hash_to_int64(block_hash: bytes) -> int:
 
 
 def get_cache_salt_id(cache_salt: str) -> int:
-    """Return the cache salt id used by request handling and cache-aware routing."""
-    from blake3 import blake3
-
-    h = blake3(cache_salt.encode("utf-8")).digest(length=8)
-    return int.from_bytes(h, "little", signed=False)
+    """Derive the salt id that keys the KV cache reuse namespace.
+
+    Single source of truth shared by the engine and the KV-cache-aware router:
+    ``KVCacheManagerV2._derive_reuse_salt`` uses it to build the
+    ``ReuseScope.salt`` that seeds the radix-tree root (and therefore every
+    block hash the worker publishes in KV cache events), and the serve-layer
+    router uses it to recompute those hashes for prefix matching. The two
+    derivations MUST stay byte-identical or salted requests silently score
+    zero cache hits on every worker (the block-match walk breaks on the first
+    mismatching block).
+
+    Note: the V1 C++ ``BlockKeyHasher`` (cpp/tensorrt_llm/batch_manager/
+    blockKey.cpp) does not use this id; it mixes
+    ``std::hash<std::string>(cacheSalt)`` directly into the root block hash,
+    which is standard-library-implementation-defined and not reproducible
+    here. Router-computed v1 hashes for salted requests therefore only match
+    workers running KV cache manager V2.
+    """
+    digest = hashlib.sha256(cache_salt.encode("utf-8")).digest()
+    return int.from_bytes(digest[:8], "little", signed=False)
 
 
 def hash_v1_block_key(

tensorrt_llm/serve/router.py

@@ -1020,6 +1020,7 @@ def _routing_key_sync(self, request: OpenAIRequest) -> dict:
         thread. The returned plain dict can be sent over HTTP unchanged.
         """
         cache_salt_id = self._get_request_cache_salt_id(request)
+        lora_task_id = self._get_request_lora_id(request)
         # Hash for every algo any server might use (usually one).
         algos = ({self._routing_hash_algo}
                  if self._routing_hash_algo is not None else {
@@ -1032,7 +1033,7 @@ def _routing_key_sync(self, request: OpenAIRequest) -> dict:
                 "synchronize it with the coordinator before routing")
         token_lists, block_hashes_by_algo = \
             self._tokenize_and_compute_block_hashes_by_algo(
-                request, algos, cache_salt_id)
+                request, algos, cache_salt_id, lora_task_id)
         return {
             "block_hashes_by_algo": block_hashes_by_algo,
             "conv_key": self._content_affinity_key(request),

tensorrt_llm/serve/router_utils.py

@@ -47,6 +47,14 @@
 hash_v1_block_key = kv_cache_hash.hash_v1_block_key
 truncate_sha256_hash_to_int64 = kv_cache_hash.truncate_sha256_hash_to_int64
 
+_UINT64_MASK = (1 << 64) - 1
+
+
+def _to_uint64(value: Optional[int]) -> Optional[int]:
+    """Reduce an id modulo 2**64 so the native uint64 LoRA id fields accept it."""
+    return None if value is None else value & _UINT64_MASK
+
+
 OpenAIRequest = Union[CompletionRequest, ChatCompletionRequest]
 BlockHash = Union[int, str]
 
@@ -155,16 +163,24 @@ def get_request_num_tokens(request: Optional[OpenAIRequest]) -> int:
 
 
 def block_key_hasher(
-    token_ids: list[int], parent_hash: Optional[int] = None, cache_salt_id: Optional[int] = None
+    token_ids: list[int],
+    parent_hash: Optional[int] = None,
+    cache_salt_id: Optional[int] = None,
+    lora_task_id: Optional[int] = None,
 ) -> int:
     parent = 0 if parent_hash is None else parent_hash
     # Fast path: the native C++ BlockKeyHasher is bit-exact with
-    # hash_v1_block_key and avoids the per-token Python loop. Its hash() binding
-    # takes no cache_salt_id, so fall back to Python only when a salt is set
-    # (rare opt-in; never in the unsalted agent/chat completion path).
+    # hash_v1_block_key and avoids the per-token Python loop. The native
+    # BlockKey carries the LoRA task id (a uint64, masked here the same way
+    # hash_v1_block_key masks it), but its binding exposes no cache salt, so
+    # only salted requests take the Python path.
     if cache_salt_id is None:
-        return _NativeBlockKeyHasher.hash(_NativeBlockKey(token_ids), parent)
-    return hash_v1_block_key(token_ids, parent_hash=parent, cache_salt_id=cache_salt_id)
+        return _NativeBlockKeyHasher.hash(
+            _NativeBlockKey(token_ids, lora_task_id=_to_uint64(lora_task_id)), parent
+        )
+    return hash_v1_block_key(
+        token_ids, parent_hash=parent, lora_task_id=lora_task_id, cache_salt_id=cache_salt_id
+    )
 
 
 class BlockHashMixin:
@@ -310,14 +326,15 @@ def _compute_block_hashes(
         token_lists: list[list[int]],
         hash_algo: str = KV_CACHE_HASH_ALGO_DEFAULT,
         cache_salt_id: Optional[int] = None,
+        lora_task_id: Optional[int] = None,
     ) -> list[list[BlockHash]]:
         if hash_algo == KV_CACHE_HASH_ALGO_V1:
             block_hasher = block_key_hasher
         elif hash_algo in (KV_CACHE_HASH_ALGO_V2, KV_CACHE_HASH_ALGO_V2_SHA256_64):
             # V2 keys chain from a reuse-scope root, which is exactly what
             # sequence_to_blockchain_keys yields; its first pair is the root itself, so
             # skip it. In KvCacheManager the last token is not part of any block key.
-            reuse_scope = ReuseScope(salt=cache_salt_id)
+            reuse_scope = ReuseScope(lora_id=_to_uint64(lora_task_id), salt=cache_salt_id)
             to_hash = truncate_sha256_hash_to_int64
             if hash_algo == KV_CACHE_HASH_ALGO_V2:
 
@@ -343,7 +360,10 @@ def to_hash(key: bytes) -> BlockHash:
                 t_end = min(t + self._tokens_per_block, len(token_list) - 1)
                 hash_list.append(
                     block_hasher(
-                        token_list[t:t_end], None if t == 0 else hash_list[-1], cache_salt_id
+                        token_list[t:t_end],
+                        None if t == 0 else hash_list[-1],
+                        cache_salt_id,
+                        lora_task_id,
                     )
                 )
             block_hashes.append(hash_list)
@@ -367,22 +387,26 @@ def _tokenize_and_compute_block_hashes_with_salt(
         self,
         request: OpenAIRequest,
         cache_salt_id: Optional[int] = None,
+        lora_task_id: Optional[int] = None,
     ) -> tuple[list[list[int]], list[list[int]]]:
         token_lists = self._tokenize(request)
-        block_hashes = self._compute_block_hashes(token_lists, cache_salt_id=cache_salt_id)
+        block_hashes = self._compute_block_hashes(
+            token_lists, cache_salt_id=cache_salt_id, lora_task_id=lora_task_id
+        )
         return token_lists, block_hashes
 
     def _tokenize_and_compute_block_hashes_by_algo(
         self,
         request: OpenAIRequest,
         hash_algos: Iterable[str],
         cache_salt_id: Optional[int] = None,
+        lora_task_id: Optional[int] = None,
     ) -> tuple[list[list[int]], dict[str, list[list[BlockHash]]]]:
         """Synchronous tokenize + per-algorithm block hashes for thread offload."""
         token_lists = self._tokenize(request)
         return token_lists, {
             hash_algo: self._compute_block_hashes(
-                token_lists, hash_algo, cache_salt_id=cache_salt_id
+                token_lists, hash_algo, cache_salt_id=cache_salt_id, lora_task_id=lora_task_id
             )
             for hash_algo in set(hash_algos)
         }
@@ -399,3 +423,17 @@ def _text_to_int_sequences(texts: list[str]) -> list[list[int]]:
     def _get_request_cache_salt_id(request: OpenAIRequest) -> Optional[int]:
         cache_salt = getattr(request, "cache_salt", None)
         return None if cache_salt is None else get_cache_salt_id(cache_salt)
+
+    @staticmethod
+    def _get_request_lora_id(request: OpenAIRequest) -> Optional[int]:
+        """LoRA adapter id keying the worker's KV reuse namespace, or None.
+
+        The worker partitions KV cache reuse by ``LoRARequest.adapter_id``
+        (== ``lora_int_id``): it becomes ``LlmRequest.lora_task_id`` and then
+        ``ReuseScope.lora_id`` (V2) / ``BlockKey.loraTaskId`` (V1 C++), so the
+        router must fold the same id into its recomputed block hashes or a
+        LoRA request can never match a worker's published blocks.
+        """
+        lora_request = getattr(request, "lora_request", None)
+        adapter_id = getattr(lora_request, "adapter_id", None)
+        return None if adapter_id is None else int(adapter_id)