feat: integrate cache_salt with KV routing for multi-tenant isolation [DYN-3178]

components/src/dynamo/router/__main__.py

@@ -141,7 +141,9 @@ async def generate(self, request):
             }
             yield llm_engine_output
 
-    async def best_worker_id(self, token_ids, router_config_override=None):
+    async def best_worker_id(
+        self, token_ids, router_config_override=None, cache_namespace=None
+    ):
         """
         Get the best worker ID for a given set of tokens without actually routing.
 
@@ -154,7 +156,9 @@ async def best_worker_id(self, token_ids, router_config_override=None):
             raise RuntimeError("Router not initialized")
 
         (worker_id, _dp_rank, _overlap_blocks) = await self.kv_router.best_worker(
-            token_ids, router_config_override
+            token_ids,
+            router_config_override,
+            cache_namespace=cache_namespace,
         )
 
         yield worker_id
@@ -177,6 +181,7 @@ async def get_overlap_scores(self, request):
             request.get("block_mm_infos"),
             request.get("lora_name"),
             request.get("include_shared", True),
+            request.get("cache_namespace"),
         )
 
         yield scores

components/src/dynamo/router/tests/test_standalone_router.py

@@ -0,0 +1,125 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import importlib.util
+import sys
+import types
+from pathlib import Path
+from unittest.mock import AsyncMock
+
+import pytest
+
+pytestmark = [pytest.mark.pre_merge, pytest.mark.unit, pytest.mark.gpu_0]
+
+
+def stub_module(name: str, **attributes: object) -> types.ModuleType:
+    module = types.ModuleType(name)
+    for attribute, value in attributes.items():
+        setattr(module, attribute, value)
+    return module
+
+
+def load_standalone_router_handler():
+    placeholder_type = type("Placeholder", (), {})
+    stubs = {
+        "uvloop": stub_module("uvloop", run=lambda coroutine: coroutine),
+        "dynamo": stub_module("dynamo"),
+        "dynamo.llm": stub_module(
+            "dynamo.llm",
+            AicPerfConfig=placeholder_type,
+            KvRouter=placeholder_type,
+            KvRouterConfig=placeholder_type,
+        ),
+        "dynamo.router": stub_module("dynamo.router"),
+        "dynamo.router.args": stub_module(
+            "dynamo.router.args",
+            DynamoRouterConfig=placeholder_type,
+            build_aic_perf_config=lambda config: config,
+            build_kv_router_config=lambda config: config,
+            parse_args=lambda argv=None: argv,
+        ),
+        "dynamo.runtime": stub_module(
+            "dynamo.runtime",
+            Client=placeholder_type,
+            DistributedRuntime=placeholder_type,
+            dynamo_worker=lambda: lambda function: function,
+        ),
+        "dynamo.runtime.logging": stub_module(
+            "dynamo.runtime.logging", configure_dynamo_logging=lambda: None
+        ),
+    }
+    previous = {name: sys.modules.get(name) for name in stubs}
+    sys.modules.update(stubs)
+    try:
+        module_path = Path(__file__).parents[1] / "__main__.py"
+        spec = importlib.util.spec_from_file_location(
+            "standalone_router_main", module_path
+        )
+        assert spec is not None and spec.loader is not None
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        return module.StandaloneRouterHandler
+    finally:
+        for name, previous_module in previous.items():
+            if previous_module is None:
+                sys.modules.pop(name, None)
+            else:
+                sys.modules[name] = previous_module
+
+
+StandaloneRouterHandler = load_standalone_router_handler()
+
+
+def handler_with_router():
+    handler = StandaloneRouterHandler.__new__(StandaloneRouterHandler)
+    router = AsyncMock()
+    handler.kv_router = router
+    return handler, router
+
+
+@pytest.mark.asyncio
+async def test_best_worker_id_forwards_cache_namespace() -> None:
+    handler, router = handler_with_router()
+    router.best_worker.return_value = (7, 0, 3)
+
+    results = [
+        worker_id
+        async for worker_id in handler.best_worker_id(
+            [1, 2, 3, 4],
+            {"temperature": 0.0},
+            cache_namespace="tenant-a",
+        )
+    ]
+
+    assert results == [7]
+    router.best_worker.assert_awaited_once_with(
+        [1, 2, 3, 4],
+        {"temperature": 0.0},
+        cache_namespace="tenant-a",
+    )
+
+
+@pytest.mark.asyncio
+async def test_get_overlap_scores_forwards_cache_namespace() -> None:
+    handler, router = handler_with_router()
+    router.get_overlap_scores.return_value = {"workers": []}
+    request = {
+        "token_ids": [1, 2, 3, 4],
+        "router_config_override": {"temperature": 0.0},
+        "block_mm_infos": None,
+        "lora_name": "adapter-a",
+        "include_shared": False,
+        "cache_namespace": "tenant-a",
+    }
+
+    results = [scores async for scores in handler.get_overlap_scores(request)]
+
+    assert results == [{"workers": []}]
+    router.get_overlap_scores.assert_awaited_once_with(
+        [1, 2, 3, 4],
+        {"temperature": 0.0},
+        None,
+        "adapter-a",
+        False,
+        "tenant-a",
+    )

components/src/dynamo/trtllm/llm_engine.py

@@ -68,6 +68,10 @@
     DisaggregatedParams,
     DisaggregatedParamsCodec,
 )
+from dynamo.trtllm.utils.request_utils import (
+    request_cache_salt,
+    stored_event_cache_salt,
+)
 from dynamo.trtllm.utils.trtllm_utils import deep_update, warn_override_collisions
 
 if TYPE_CHECKING:
@@ -587,6 +591,7 @@ def _dispatch_kv_event(self, event: dict[str, Any]) -> None:
                 block_hashes,
                 parent_hash,
                 lora_name=data.get("lora_name"),
+                cache_salt=stored_event_cache_salt(data),
             )
         elif kind == "removed":
             partial = self._partial_block_hashes_by_rank.get(rank)
@@ -844,12 +849,14 @@ async def _generate_started(
         # Prefill returns one non-streaming chunk carrying the handoff -
         # matches the legacy disagg wire format.
         streaming = not is_prefill
+        cache_salt = request_cache_salt(request)
         generation_result = self._engine.llm.generate_async(
             inputs=token_ids,
             sampling_params=sampling_params,
             streaming=streaming,
             disaggregated_params=disaggregated_params,
             scheduling_params=scheduling_params,
+            cache_salt=cache_salt,
             **telemetry.engine_trace_kwargs(context),
         )
 

components/src/dynamo/trtllm/publisher.py

@@ -37,6 +37,7 @@
 
 from dynamo.common.utils.prometheus import LLMBackendMetrics
 from dynamo.llm import FpmDirectPublisher, KvEventPublisher, WorkerMetricsPublisher
+from dynamo.trtllm.utils.request_utils import stored_event_cache_salt
 
 logger = logging.getLogger(__name__)
 
@@ -160,6 +161,7 @@ def publish_stored(
         block_mm_infos: Optional[list[dict | None]] = None,
         attention_dp_rank: int = 0,
         lora_name: Optional[str] = None,
+        cache_salt: Optional[str] = None,
     ) -> None:
         """Publish a BlockStored event.
 
@@ -182,6 +184,8 @@ def publish_stored(
         }
         if lora_name is not None:
             event["lora_name"] = lora_name
+        if cache_salt is not None:
+            event["cache_salt"] = cache_salt
 
         # Add multimodal info if present
         if block_mm_infos is not None:
@@ -868,16 +872,28 @@ def _handle_kv_event(self, event):
                     block_mm_infos.append(None)
 
             lora_name = data.get("lora_name")
+            try:
+                cache_salt = stored_event_cache_salt(data)
+            except ValueError as error:
+                logger.warning(
+                    "Dropping stored KV event with invalid cache namespace: "
+                    "engine_event_id=%s attention_dp_rank=%s error=%s",
+                    event_id,
+                    attention_dp_rank,
+                    error,
+                )
+                return
 
             logger.debug(
                 "Publishing stored KV event: engine_event_id=%s "
-                "attention_dp_rank=%s blocks=%s tokens=%s lora_name=%s "
+                "attention_dp_rank=%s blocks=%s tokens=%s lora_name=%s has_cache_salt=%s "
                 "has_parent=%s",
                 event_id,
                 attention_dp_rank,
                 len(block_hashes),
                 len(token_ids),
                 lora_name,
+                cache_salt is not None,
                 parent_hash is not None,
             )
             # Publish to ZMQ if consolidator is enabled, otherwise publish to NATS
@@ -892,6 +908,7 @@ def _handle_kv_event(self, event):
                     block_mm_infos,
                     attention_dp_rank,
                     lora_name,
+                    cache_salt,
                 )
             elif self.kv_event_publishers:
                 # No consolidator: publish to NATS (router subscribes directly)
@@ -905,6 +922,7 @@ def _handle_kv_event(self, event):
                         parent_hash,
                         block_mm_infos,
                         lora_name=lora_name,
+                        cache_salt=cache_salt,
                     )
                 else:
                     logging.warning(

components/src/dynamo/trtllm/request_handlers/handler_base.py

@@ -53,6 +53,7 @@
     DisaggregatedParams,
     DisaggregatedParamsCodec,
 )
+from dynamo.trtllm.utils.request_utils import request_cache_salt
 
 if TYPE_CHECKING:
     # tensorrt_llm may use a different version that doesn't have MetricsCollector,
@@ -1110,6 +1111,7 @@ async def _generate_locally_impl(
 
         # Priority is a float in [0.0, 1.0]; health checks use 1.0. Default is 0.5.
         priority = request.get("priority", DEFAULT_REQUEST_PRIORITY)
+        cache_salt = request_cache_salt(request)
 
         try:
             # NEW: Updated engine call to include multimodal data
@@ -1121,6 +1123,7 @@ async def _generate_locally_impl(
                 trace_headers=trace_headers,
                 scheduling_params=scheduling_params,
                 priority=priority,
+                cache_salt=cache_salt,
             )
 
             # In disagg decode mode, wrap abort() to defer until first token