KV-cache-aware router never matches LoRA or salted requests: lora_id is dropped and cache_salt is hashed with a different algorithm than the engine

"""CPU repro: KV-cache-aware router vs V2 engine reuse-namespace mismatch.

No GPU, no CUDA, no compiled tensorrt_llm bindings. Uses only the pure-Python
hashing modules that both sides of the real system import.

Two independent partitioning dimensions are dropped/derived differently on the
router side vs the engine/event side:

  (1) lora_id  -- tensorrt_llm/serve/router_utils.py:159-170 (block_key_hasher)
                  and :174-179 / :333 (v2 hashers) never pass lora_task_id,
                  while tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py:653-665
                  publishes event hashes WITH scope.lora_id.

  (2) cache_salt -> salt id
                  router : tensorrt_llm/runtime/kv_cache_hash.py:48  get_cache_salt_id  = blake3(s)[:8]
                  engine : tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py:3972
                           _derive_reuse_salt = sha256(s)[:8]
"""
import hashlib
import sys

import sys, types, pathlib, importlib
_ROOT = pathlib.Path("/home/oneknight/projects/oss/tensorrt-llm")
# Load the two pure-Python modules without executing kv_cache_manager_v2/__init__.py,
# which hard-imports the compiled C++ bindings (see _copy_engine.py:41-73).
sys.path.insert(0, str(_ROOT / "tensorrt_llm" / "runtime"))
_pkg = types.ModuleType("kv_cache_manager_v2")
_pkg.__path__ = [str(_ROOT / "tensorrt_llm" / "runtime" / "kv_cache_manager_v2")]
sys.modules["kv_cache_manager_v2"] = _pkg
from kv_cache_hash import get_cache_salt_id, hash_v1_block_key
from kv_cache_manager_v2._block_radix_tree import ReuseScope, RootBlock, Hasher

TPB = 4
TOKENS = list(range(1, 33))          # 32 prompt tokens
LORA = 7
SALT = "tenant-a"


def derive_reuse_salt(cache_salt):
    """Verbatim copy of KVCacheManagerV2._derive_reuse_salt (kv_cache_manager_v2.py:3962)."""
    if cache_salt is None:
        return None
    return int.from_bytes(hashlib.sha256(cache_salt.encode("utf-8")).digest()[:8], "little")


def blocks(tokens, tpb):
    """Both sides block the prompt the same way: exclude the final token."""
    return [tokens[t:min(t + tpb, len(tokens) - 1)]
            for t in range(0, len(tokens) - 1, tpb)]


def worker_v1_event_hashes(tokens, lora_id, salt_id):
    """_event_manager._v1_hash_from_radix_block chain (lines 606-621)."""
    out, parent = [], 0
    for blk in blocks(tokens, TPB):
        parent = hash_v1_block_key(blk, parent_hash=parent,
                                   lora_task_id=lora_id, cache_salt_id=salt_id)
        out.append(parent)
    return out


def router_v1_hashes(tokens, salt_id):
    """router_utils.block_key_hasher chain (lines 159-170, 346-357)."""
    out, parent = [], None
    for blk in blocks(tokens, TPB):
        parent = hash_v1_block_key(blk, parent_hash=0 if parent is None else parent,
                                   cache_salt_id=salt_id)   # <-- no lora_task_id
        out.append(parent)
    return out


def worker_v2_root(lora_id, salt_id):
    return RootBlock.make_key(ReuseScope(lora_id=lora_id, salt=salt_id))


def router_v2_root(salt_id):
    return RootBlock.make_key(ReuseScope(salt=salt_id))       # <-- router_utils.py:177/333


fails = []

def check(name, condition, detail):
    status = "MATCH  " if condition else "MISMATCH"
    print(f"  [{status}] {name}: {detail}")
    return condition


print("=" * 78)
print("A. baseline: no lora, no salt  (expected: router and worker agree)")
print("=" * 78)
w = worker_v1_event_hashes(TOKENS, None, None)
r = router_v1_hashes(TOKENS, None)
if not check("v1 hashes", w == r, f"worker[0]={w[0]}  router[0]={r[0]}"):
    fails.append("baseline")

print()
print("=" * 78)
print("B. LoRA request (lora_task_id=7, no salt)")
print("   engine partitions by lora_id; router never sends it")
print("=" * 78)
w = worker_v1_event_hashes(TOKENS, LORA, None)
r = router_v1_hashes(TOKENS, None)
overlap = set(w) & set(r)
if check("v1 hashes differ", w != r, f"worker[0]={w[0]}  router[0]={r[0]}"):
    pass
else:
    fails.append("B-differ")
if check("zero overlap -> router can never hit worker blocks",
         len(overlap) == 0, f"|overlap|={len(overlap)} of {len(w)} blocks"):
    pass
else:
    fails.append("B-overlap")

# collision side of the same drop: two DIFFERENT adapters look identical to the router
r_a = router_v1_hashes(TOKENS, None)
r_b = router_v1_hashes(TOKENS, None)
w_a = worker_v1_event_hashes(TOKENS, 7, None)
w_b = worker_v1_event_hashes(TOKENS, 8, None)
if not check("router collapses adapter 7 and adapter 8 into one hash",
             r_a == r_b and w_a != w_b,
             f"router same={r_a == r_b}, worker differs={w_a != w_b}"):
    fails.append("B-collision")

print()
print("=" * 78)
print("C. Salted request (cache_salt='tenant-a')")
print("   router salt id = blake3(s)[:8]   engine salt id = sha256(s)[:8]")
print("=" * 78)
router_salt = get_cache_salt_id(SALT)
engine_salt = derive_reuse_salt(SALT)
print(f"   get_cache_salt_id('{SALT}')  = {router_salt}")
print(f"   _derive_reuse_salt('{SALT}') = {engine_salt}")
if not check("salt ids differ", router_salt != engine_salt,
             f"blake3 {router_salt} != sha256 {engine_salt}"):
    fails.append("C-saltid")

w = worker_v1_event_hashes(TOKENS, None, engine_salt)
r = router_v1_hashes(TOKENS, router_salt)
overlap = set(w) & set(r)
if not check("zero overlap on v1 algo", len(overlap) == 0,
             f"|overlap|={len(overlap)} of {len(w)} blocks"):
    fails.append("C-v1")

wk = worker_v2_root(None, engine_salt)
rk = router_v2_root(router_salt)
if not check("v2 root keys differ -> whole chain diverges", wk != rk,
             f"worker root {wk.hex()[:16]}  router root {rk.hex()[:16]}"):
    fails.append("C-v2root")

print()
print("=" * 78)
print("D. control: if the router used the ENGINE's salt derivation, it matches")
print("=" * 78)
w = worker_v1_event_hashes(TOKENS, None, engine_salt)
r = router_v1_hashes(TOKENS, engine_salt)
if not check("v1 hashes agree once salt derivation is shared", w == r,
             f"worker[0]={w[0]}  router[0]={r[0]}"):
    fails.append("D")

print()
print("=" * 78)
print("RESULT:", "all expectations held" if not fails else f"UNEXPECTED: {fails}")
print("=" * 78)
sys.exit(1 if fails else 0)