From 443e33114ea3dd65b7ef1344ed329fe4e68d155c Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Mon, 8 Jun 2026 14:19:43 -0700 Subject: [PATCH 01/10] Port cache salt routing support Signed-off-by: jthomson04 --- components/src/dynamo/trtllm/llm_engine.py | 22 +- components/src/dynamo/trtllm/publisher.py | 9 +- .../trtllm/request_handlers/handler_base.py | 21 +- .../trtllm/tests/test_trtllm_fpm_publisher.py | 57 +++++ .../trtllm/tests/test_trtllm_handler_base.py | 22 ++ .../tests/test_trtllm_trace_propagation.py | 22 +- deploy/inference-gateway/ext-proc/src/epp.rs | 59 ++++- lib/bindings/c/src/lib.rs | 33 ++- lib/bindings/python/rust/llm/kv.rs | 25 ++- lib/bindings/python/src/dynamo/_core.pyi | 6 + lib/kv-router/src/indexer/branch_sharded.rs | 2 + lib/kv-router/src/indexer/kv_indexer.rs | 2 + lib/kv-router/src/indexer/local.rs | 3 +- lib/kv-router/src/indexer/tests.rs | 78 ++++++- lib/kv-router/src/indexer/thread_pool.rs | 2 + lib/kv-router/src/indexer/traits.rs | 6 + lib/kv-router/src/protocols.rs | 208 +++++++++++++++++- lib/kv-router/src/services/indexer/server.rs | 8 +- lib/kv-router/src/services/selection/input.rs | 1 + lib/kv-router/src/zmq_wire/convert.rs | 57 +++-- lib/kv-router/src/zmq_wire/deserialize.rs | 19 +- lib/kv-router/src/zmq_wire/extra_keys.rs | 38 ++++ lib/kv-router/src/zmq_wire/mod.rs | 53 ++++- lib/kv-router/src/zmq_wire/tests.rs | 124 ++++++++++- lib/kv-router/src/zmq_wire/types.rs | 8 + .../tests/standalone_indexer_http.rs | 1 + lib/kvbm-consolidator/tests/dedup.rs | 1 + lib/kvbm-consolidator/tests/e2e.rs | 2 + lib/kvbm-consolidator/tests/kvbm_bridge.rs | 1 + lib/kvbm-consolidator/tests/lifecycle.rs | 1 + .../tests/output_contract.rs | 2 + lib/kvbm-consolidator/tests/zmq_ingress.rs | 1 + lib/llm/src/kv_router.rs | 50 ++++- lib/llm/src/kv_router/prefill_router/query.rs | 2 + lib/llm/src/kv_router/publisher/tests.rs | 44 +++- .../src/kv_router/push_router/selection.rs | 6 + lib/llm/src/kv_router/route_lookup.rs | 45 +++- lib/llm/src/kv_router/shared_cache.rs | 32 ++- lib/llm/src/preprocessor.rs | 1 + lib/llm/src/protocols/common/extensions.rs | 30 +++ lib/llm/src/protocols/common/preprocessor.rs | 24 ++ lib/mocker/src/replay/online/router.rs | 4 +- 42 files changed, 1049 insertions(+), 83 deletions(-) diff --git a/components/src/dynamo/trtllm/llm_engine.py b/components/src/dynamo/trtllm/llm_engine.py index a4d517f653be..971a9c07cfce 100644 --- a/components/src/dynamo/trtllm/llm_engine.py +++ b/components/src/dynamo/trtllm/llm_engine.py @@ -18,7 +18,7 @@ import sys import threading import time -from collections.abc import AsyncGenerator, Callable +from collections.abc import AsyncGenerator, Callable, Mapping from dataclasses import asdict from typing import TYPE_CHECKING, Any, Optional @@ -107,6 +107,23 @@ } +def _request_cache_salt(request: Mapping[str, Any]) -> Optional[str]: + routing = request.get("routing") or {} + if isinstance(routing, dict): + cache_salt = routing.get("cache_salt") + if cache_salt is not None: + return cache_salt + + extra_args = request.get("extra_args") or {} + nvext = extra_args.get("nvext") if isinstance(extra_args, dict) else None + if isinstance(nvext, dict): + cache_salt = nvext.get("cache_salt") + if cache_salt is not None: + return cache_salt + + return None + + def _to_signed_i64(value: int | None) -> int | None: """Two's-complement cast of a Python int into the signed 64-bit range.""" if value is None: @@ -587,6 +604,7 @@ def _dispatch_kv_event(self, event: dict[str, Any]) -> None: block_hashes, parent_hash, lora_name=data.get("lora_name"), + cache_salt=data.get("cache_salt"), ) elif kind == "removed": partial = self._partial_block_hashes_by_rank.get(rank) @@ -844,12 +862,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), ) diff --git a/components/src/dynamo/trtllm/publisher.py b/components/src/dynamo/trtllm/publisher.py index bf6f0438e54c..a95fee016a3a 100644 --- a/components/src/dynamo/trtllm/publisher.py +++ b/components/src/dynamo/trtllm/publisher.py @@ -160,6 +160,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 +183,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 +871,18 @@ def _handle_kv_event(self, event): block_mm_infos.append(None) lora_name = data.get("lora_name") + cache_salt = data.get("cache_salt") 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 cache_salt=%s " "has_parent=%s", event_id, attention_dp_rank, len(block_hashes), len(token_ids), lora_name, + cache_salt, parent_hash is not None, ) # Publish to ZMQ if consolidator is enabled, otherwise publish to NATS @@ -892,6 +897,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 +911,7 @@ def _handle_kv_event(self, event): parent_hash, block_mm_infos, lora_name=lora_name, + cache_salt=cache_salt, ) else: logging.warning( diff --git a/components/src/dynamo/trtllm/request_handlers/handler_base.py b/components/src/dynamo/trtllm/request_handlers/handler_base.py index 625ed0f11ef6..e3343b16ad20 100644 --- a/components/src/dynamo/trtllm/request_handlers/handler_base.py +++ b/components/src/dynamo/trtllm/request_handlers/handler_base.py @@ -18,7 +18,7 @@ import logging import os import re -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Mapping from contextlib import asynccontextmanager from dataclasses import asdict, dataclass from typing import TYPE_CHECKING, Any, Optional, Protocol, Union @@ -64,6 +64,23 @@ logger = logging.getLogger(__name__) +def _request_cache_salt(request: Mapping[str, Any]) -> Optional[str]: + routing = request.get("routing") or {} + if isinstance(routing, dict): + cache_salt = routing.get("cache_salt") + if cache_salt is not None: + return cache_salt + + extra_args = request.get("extra_args") or {} + nvext = extra_args.get("nvext") if isinstance(extra_args, dict) else None + if isinstance(nvext, dict): + cache_salt = nvext.get("cache_salt") + if cache_salt is not None: + return cache_salt + + return None + + class TRTLLMEnginePauseController: """Adapts TRT-LLM sleep/wake to the standard pause controller interface. @@ -1110,6 +1127,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 +1139,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 diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py index 859f1dcd25ed..88728f89ab2f 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py @@ -368,6 +368,63 @@ def test_publisher_initialize_constructs_fpm_direct_publisher_when_fpm_enabled( assert pub.fpm_publisher is not None +def _publisher_for_kv_event_test(): + pub = publisher_mod.Publisher.__new__(publisher_mod.Publisher) + pub.additional_metrics = None + pub._last_engine_event_id_by_rank = {} + pub.processing_initial_created_events = False + pub.partial_block_hashes = set() + pub.kv_block_size = 4 + pub.max_window_size = None + return pub + + +def _stored_kv_event(cache_salt="tenant-a"): + return { + "event_id": 1, + "attention_dp_rank": 0, + "data": { + "type": "stored", + "parent_hash": None, + "cache_salt": cache_salt, + "blocks": [ + { + "block_hash": 123, + "tokens": [ + {"token_id": 1}, + {"token_id": 2}, + {"token_id": 3}, + {"token_id": 4}, + ], + } + ], + }, + } + + +def test_handle_kv_event_forwards_cache_salt_to_direct_publisher(): + pub = _publisher_for_kv_event_test() + publisher = MagicMock() + pub.zmq_kv_event_publisher = None + pub.kv_event_publishers = {0: publisher} + + pub._handle_kv_event(_stored_kv_event()) + + publisher.publish_stored.assert_called_once() + assert publisher.publish_stored.call_args.kwargs["cache_salt"] == "tenant-a" + + +def test_handle_kv_event_forwards_cache_salt_to_zmq_publisher(): + pub = _publisher_for_kv_event_test() + pub.zmq_kv_event_publisher = MagicMock() + pub.kv_event_publishers = None + + pub._handle_kv_event(_stored_kv_event()) + + pub.zmq_kv_event_publisher.publish_stored.assert_called_once() + assert pub.zmq_kv_event_publisher.publish_stored.call_args.args[-1] == "tenant-a" + + def test_publisher_initializes_fpm_publisher_under_attention_dp(monkeypatch): """Under attention-DP (attention_dp_size > 1), Publisher.initialize() constructs one FpmDirectPublisher channel per attention-DP rank.""" diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py b/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py index 32ed43c6fe24..2ca0257246ea 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py @@ -731,3 +731,25 @@ async def test_regular_request_gets_default_priority(self): handler.engine.llm.generate_async.assert_called_once() _, kwargs = handler.engine.llm.generate_async.call_args assert kwargs["priority"] == DEFAULT_REQUEST_PRIORITY + + @pytest.mark.asyncio + async def test_routing_cache_salt_forwarded_to_generate_async(self): + handler = self._make_handler() + generation_result = self._make_mock_generation_result() + handler.engine.llm.generate_async = MagicMock(return_value=generation_result) + + request = { + "token_ids": [1, 2, 3], + "stop_conditions": {"max_tokens": 10}, + "sampling_options": {"temperature": 0.7}, + "routing": {"cache_salt": "tenant-a"}, + } + + chunks = [ + c async for c in handler.generate_locally(request, self._make_context()) + ] + assert len(chunks) > 0 + + handler.engine.llm.generate_async.assert_called_once() + _, kwargs = handler.engine.llm.generate_async.call_args + assert kwargs["cache_salt"] == "tenant-a" diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_trace_propagation.py b/components/src/dynamo/trtllm/tests/test_trtllm_trace_propagation.py index 0804f3fd1a7e..4313b17ed2ca 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_trace_propagation.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_trace_propagation.py @@ -77,8 +77,10 @@ def _make_engine(generate_async) -> TrtllmLLMEngine: return engine -async def _drain(engine: TrtllmLLMEngine, ctx: _FakeContext) -> None: - async for _ in engine.generate({"token_ids": [1, 2, 3]}, ctx): +async def _drain( + engine: TrtllmLLMEngine, ctx: _FakeContext, request: dict | None = None +) -> None: + async for _ in engine.generate(request or {"token_ids": [1, 2, 3]}, ctx): pass @@ -111,3 +113,19 @@ def fake_generate_async(**kwargs): # kwarg omitted (engine_trace_kwargs returns {}). assert "trace_headers" not in captured + + +async def test_forwards_routing_cache_salt(): + captured: dict = {} + + def fake_generate_async(**kwargs): + captured.update(kwargs) + return _empty_async_iter() + + await _drain( + _make_engine(fake_generate_async), + _FakeContext(), + {"token_ids": [1, 2, 3], "routing": {"cache_salt": "tenant-a"}}, + ) + + assert captured["cache_salt"] == "tenant-a" diff --git a/deploy/inference-gateway/ext-proc/src/epp.rs b/deploy/inference-gateway/ext-proc/src/epp.rs index 2ea901b2385e..32bea1ce05a3 100644 --- a/deploy/inference-gateway/ext-proc/src/epp.rs +++ b/deploy/inference-gateway/ext-proc/src/epp.rs @@ -197,10 +197,10 @@ impl Router { /// Tokenize a JSON request body and extract router queue priorities. /// - /// Returns `(token_ids, priority_jump, strict_priority)`. Priorities default - /// to zero when absent. Mirrors the standalone Dynamo preprocessor lift in - /// `lib/llm/src/preprocessor.rs`. - pub fn tokenize(&self, request_json: &str) -> Result<(Vec, f64, u32)> { + /// Returns `(token_ids, cache_namespace, priority_jump, strict_priority)`. + /// Priorities default to zero when absent. Mirrors the standalone Dynamo + /// preprocessor lift in `lib/llm/src/preprocessor.rs`. + pub fn tokenize(&self, request_json: &str) -> Result<(Vec, Option, f64, u32)> { // TODO(epp-request-routing): Reuse shared preprocessing so expected output // length, LoRA, pins, sessions, topology constraints, additional protocols, // and multimodal routing hashes are preserved. @@ -209,6 +209,7 @@ impl Router { let priority_jump = extract_priority_jump(&request); let strict_priority = extract_strict_priority(&request); + let cache_namespace = extract_cache_namespace(&request); let formatted_prompt = self .preprocessor @@ -218,6 +219,7 @@ impl Router { let encoding = self.preprocessor.tokenize(&formatted_prompt)?; Ok(( encoding.token_ids().to_vec(), + cache_namespace, priority_jump, strict_priority, )) @@ -297,6 +299,7 @@ impl Router { pub async fn route_prefill( &self, tokens: &[u32], + cache_namespace: Option, priority_jump: f64, strict_priority: u32, allowed_worker_ids: Option>, @@ -313,6 +316,7 @@ impl Router { tokens, None, None, + cache_namespace, priority_jump, strict_priority, allowed_worker_ids, @@ -342,6 +346,7 @@ impl Router { &self, tokens: &[u32], is_disaggregated: bool, + cache_namespace: Option, priority_jump: f64, strict_priority: u32, allowed_worker_ids: Option>, @@ -360,6 +365,7 @@ impl Router { config_override.as_ref(), false, None, + cache_namespace, priority_jump, strict_priority, None, @@ -378,6 +384,7 @@ impl Router { worker_id: u64, dp_rank: u32, is_disaggregated: bool, + cache_namespace: Option, ) -> Result<()> { let decode_router = self.decode_router.clone(); let request_id = request_id.to_owned(); @@ -388,7 +395,7 @@ impl Router { let router_config_override = decode_router_config_override(is_disaggregated); let overlap_blocks = decode_router - .get_overlap_blocks(&tokens, None, worker, None) + .get_overlap_blocks(&tokens, None, worker, None, cache_namespace.as_deref()) .await .map_err(|e| anyhow::anyhow!("get_overlap_blocks failed: {e:?}"))?; @@ -403,6 +410,7 @@ impl Router { None, worker, None, + cache_namespace, router_config_override.as_ref(), ) .await; @@ -494,6 +502,15 @@ fn extract_strict_priority( .unwrap_or(0) } +fn extract_cache_namespace( + request: &dynamo_llm::types::openai::chat_completions::NvCreateChatCompletionRequest, +) -> Option { + request + .nvext + .as_ref() + .and_then(|nvext| nvext.cache_salt.clone()) +} + struct DiscoveredModelBootstrap { preprocessor: Arc, card: ModelDeploymentCard, @@ -911,7 +928,7 @@ impl EndpointPicker for Router { let body_str = std::str::from_utf8(&req.body) .map_err(|e| PickError::TokenizationFailed(format!("Invalid UTF-8: {e}")))?; - let (tokens, priority_jump, strict_priority) = self + let (tokens, cache_namespace, priority_jump, strict_priority) = self .tokenize(body_str) .map_err(|e| PickError::TokenizationFailed(e.to_string()))?; @@ -922,6 +939,7 @@ impl EndpointPicker for Router { let prefill_result = self .route_prefill( &tokens, + cache_namespace.clone(), priority_jump, strict_priority, allowed_worker_ids.clone(), @@ -947,6 +965,7 @@ impl EndpointPicker for Router { .route_decode( &tokens, is_disaggregated, + cache_namespace.clone(), priority_jump, strict_priority, allowed_worker_ids, @@ -989,6 +1008,7 @@ impl EndpointPicker for Router { decode_worker.worker_id, decode_worker.dp_rank, is_disaggregated, + cache_namespace, ) .await { @@ -1137,4 +1157,31 @@ mod tests { .unwrap(); assert_eq!(extract_strict_priority(&without_nvext), 0); } + + #[test] + fn cache_namespace_lifted_from_nvext_cache_salt() { + let with_cache_salt: dynamo_llm::types::openai::chat_completions::NvCreateChatCompletionRequest = + serde_json::from_str( + r#"{ + "model": "test", + "messages": [{"role": "user", "content": "hi"}], + "nvext": {"cache_salt": "tenant-a"} + }"#, + ) + .unwrap(); + assert_eq!( + extract_cache_namespace(&with_cache_salt).as_deref(), + Some("tenant-a") + ); + + let without_nvext: dynamo_llm::types::openai::chat_completions::NvCreateChatCompletionRequest = + serde_json::from_str( + r#"{ + "model": "test", + "messages": [{"role": "user", "content": "hi"}] + }"#, + ) + .unwrap(); + assert_eq!(extract_cache_namespace(&without_nvext), None); + } } diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index 26b36300ac88..e02219330931 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -18,7 +18,9 @@ use dynamo_kv_router::{ use dynamo_llm::kv_router::publisher::KvEventPublisher; use dynamo_llm::model_card::ModelDeploymentCard; use dynamo_llm::preprocessor::OpenAIPreprocessor; -use dynamo_llm::protocols::common::extensions::{NvExt, routing_constraints_to_kv}; +use dynamo_llm::protocols::common::extensions::{ + NvExt, NvExtProvider, routing_constraints_to_kv, +}; use dynamo_llm::types::openai::chat_completions::NvCreateChatCompletionRequest; use dynamo_llm::types::openai::completions::NvCreateCompletionRequest; use dynamo_runtime::discovery::{DiscoveryQuery, hash_pod_name}; @@ -234,6 +236,7 @@ fn kv_event_create_stored_block_from_parts( kv_block_size, BlockHashOptions { lora_name, + cache_namespace: None, ..Default::default() }, )[0]; @@ -458,6 +461,7 @@ impl RouterHandles { tokens: &[u32], block_mm_infos: Option<&[Option]>, lora_name: Option, + cache_namespace: Option, priority_jump: f64, strict_priority: u32, allowed_worker_ids: Option>, @@ -473,6 +477,7 @@ impl RouterHandles { tokens, block_mm_infos, lora_name, + cache_namespace, priority_jump, strict_priority, allowed_worker_ids, @@ -513,10 +518,12 @@ impl RouterHandles { /// selection. State updates require a `context_id` (request id) and are managed via the /// explicit bookkeeping APIs (`add_request`, `mark_prefill_complete`, `free_request`). /// Returns (worker, overlap_blocks) on success. + #[expect(clippy::too_many_arguments)] async fn query_decode_worker( &self, tokens: &[u32], is_disaggregated: bool, + cache_namespace: Option, priority_jump: f64, strict_priority: u32, allowed_worker_ids: Option>, @@ -549,6 +556,7 @@ impl RouterHandles { false, false, None, + cache_namespace, priority_jump, strict_priority, None, @@ -610,6 +618,12 @@ fn extract_routing_constraints(nvext: Option<&NvExt>) -> RoutingConstraints { .unwrap_or_default() } +fn extract_cache_namespace(request: &R) -> Option { + request + .nvext() + .and_then(|nvext| nvext.cache_salt.clone()) +} + /// Opaque handle for the router pair pub type RouterHandlesPtr = *mut RouterHandles; @@ -911,7 +925,7 @@ pub unsafe extern "C" fn add_request( // Compute overlap_blocks using the public method let overlap_blocks = match decode_router - .get_overlap_blocks(&tokens, None, worker, None) + .get_overlap_blocks(&tokens, None, worker, None, None) .await { Ok(overlap) => overlap, @@ -931,6 +945,7 @@ pub unsafe extern "C" fn add_request( None, worker, None, // lora_name + None, // cache_namespace Some(&router_config_override), ) .await; @@ -1120,10 +1135,12 @@ pub unsafe extern "C" fn free_routing_result(result: *mut CRoutingResult) { /// absent. This mirrors the standalone Dynamo preprocessor lift in /// `lib/llm/src/preprocessor.rs` so the GAIE/EPP path produces the same queue /// ordering as a non-EPP deployment. +type PreprocessedRequest = (Vec, Option, f64, u32, RoutingConstraints); + unsafe fn preprocess_request( handles: &RouterHandles, request_json: *const c_char, -) -> Result<(Vec, f64, u32, RoutingConstraints), QueryRouterResult> { +) -> Result { let preprocessor = match &handles.preprocessor { Some(p) => p, None => { @@ -1155,6 +1172,7 @@ unsafe fn preprocess_request( }; let priority_jump = extract_priority_jump(request.nvext.as_ref()); let strict_priority = extract_strict_priority(request.nvext.as_ref()); + let cache_namespace = extract_cache_namespace(&request); let routing_constraints = extract_routing_constraints(request.nvext.as_ref()); let (token_ids, _) = match handles .runtime @@ -1178,6 +1196,7 @@ unsafe fn preprocess_request( return Ok(( token_ids, + cache_namespace, priority_jump, strict_priority, routing_constraints, @@ -1194,6 +1213,7 @@ unsafe fn preprocess_request( let priority_jump = extract_priority_jump(request.nvext.as_ref()); let strict_priority = extract_strict_priority(request.nvext.as_ref()); + let cache_namespace = extract_cache_namespace(&request); let routing_constraints = extract_routing_constraints(request.nvext.as_ref()); let formatted_prompt = match preprocessor.apply_template(&request) { @@ -1224,6 +1244,7 @@ unsafe fn preprocess_request( Ok(( token_ids, + cache_namespace, priority_jump, strict_priority, routing_constraints, @@ -1313,7 +1334,7 @@ pub unsafe extern "C" fn route_prefill_request( let handles = unsafe { &*handle }; - let (tokens, priority_jump, strict_priority, routing_constraints) = + let (tokens, cache_namespace, priority_jump, strict_priority, routing_constraints) = match unsafe { preprocess_request(handles, request_json) } { Ok(t) => t, Err(code) => return code, @@ -1327,6 +1348,7 @@ pub unsafe extern "C" fn route_prefill_request( &tokens, None, None, + cache_namespace, priority_jump, strict_priority, allowed_worker_ids, @@ -1394,7 +1416,7 @@ pub unsafe extern "C" fn route_decode_request( let handles = unsafe { &*handle }; - let (tokens, priority_jump, strict_priority, routing_constraints) = + let (tokens, cache_namespace, priority_jump, strict_priority, routing_constraints) = match unsafe { preprocess_request(handles, request_json) } { Ok(t) => t, Err(code) => return code, @@ -1407,6 +1429,7 @@ pub unsafe extern "C" fn route_decode_request( .query_decode_worker( &tokens, is_disaggregated, + cache_namespace, priority_jump, strict_priority, allowed_worker_ids, diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index b4bcefb958ca..5d5f6d003273 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -729,7 +729,7 @@ fn init_standalone_logging() { } #[pyfunction] -#[pyo3(name = "compute_block_hash_for_seq", signature = (tokens, kv_block_size, block_mm_infos=None, lora_name=None, is_eagle=None))] +#[pyo3(name = "compute_block_hash_for_seq", signature = (tokens, kv_block_size, block_mm_infos=None, lora_name=None, is_eagle=None, cache_namespace=None))] pub fn compute_block_hash_for_seq_py( _py: Python, tokens: Vec, @@ -737,6 +737,7 @@ pub fn compute_block_hash_for_seq_py( block_mm_infos: Option>, lora_name: Option, is_eagle: Option, + cache_namespace: Option, ) -> PyResult> { if kv_block_size == 0 { return Err(PyErr::new::( @@ -755,6 +756,7 @@ pub fn compute_block_hash_for_seq_py( BlockHashOptions { block_mm_infos: mm_infos.as_deref(), lora_name: lora_name.as_deref(), + cache_namespace: cache_namespace.as_deref(), is_eagle, }, ); @@ -952,7 +954,7 @@ impl KvEventPublisher { } #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (token_ids, num_block_tokens, block_hashes, parent_hash=None, block_mm_infos=None, lora_name=None, is_eagle=None))] + #[pyo3(signature = (token_ids, num_block_tokens, block_hashes, parent_hash=None, block_mm_infos=None, lora_name=None, is_eagle=None, cache_salt=None))] fn publish_stored( &self, py: Python, @@ -963,6 +965,7 @@ impl KvEventPublisher { block_mm_infos: Option>, lora_name: Option, is_eagle: Option, + cache_salt: Option, ) -> PyResult<()> { let kv_block_size = self.kv_block_size as u32; let dp_rank = self.dp_rank; @@ -989,6 +992,7 @@ impl KvEventPublisher { &num_block_tokens, &block_hashes_u64, lora_name.as_deref(), + cache_salt.as_deref(), &warning_count, mm_infos.as_deref(), is_eagle, @@ -1808,7 +1812,7 @@ impl KvRouter { } #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (token_ids, router_config_override=None, request_id=None, update_indexer=false, block_mm_infos=None, lora_name=None, routing_constraints=None, strict_priority=0, policy_class=None))] + #[pyo3(signature = (token_ids, router_config_override=None, request_id=None, update_indexer=false, block_mm_infos=None, lora_name=None, routing_constraints=None, strict_priority=0, policy_class=None, cache_namespace=None))] fn best_worker<'p>( &self, py: Python<'p>, @@ -1821,6 +1825,7 @@ impl KvRouter { routing_constraints: Option, strict_priority: u32, policy_class: Option, + cache_namespace: Option, ) -> PyResult> { let router_config_override = if let Some(obj) = router_config_override { let override_config: RouterConfigOverride = @@ -1847,6 +1852,7 @@ impl KvRouter { update_states, false, lora_name.clone(), + cache_namespace.clone(), 0.0, strict_priority, policy_class, @@ -1881,6 +1887,10 @@ impl KvRouter { if let Some(lora_name) = lora_name.as_ref() { tokens_with_hashes = tokens_with_hashes.with_lora_name(lora_name.clone()); } + if let Some(cache_namespace) = cache_namespace.as_ref() { + tokens_with_hashes = + tokens_with_hashes.with_cache_namespace(cache_namespace.clone()); + } chooser .record_routing_decision(tokens_with_hashes, best_worker) .await @@ -1919,13 +1929,14 @@ impl KvRouter { }) } - #[pyo3(signature = (token_ids, block_mm_infos=None, lora_name=None))] + #[pyo3(signature = (token_ids, block_mm_infos=None, lora_name=None, cache_namespace=None))] fn get_potential_loads<'p>( &self, py: Python<'p>, token_ids: Vec, block_mm_infos: Option, lora_name: Option, + cache_namespace: Option, ) -> PyResult> { let block_mm_infos = block_mm_infos .map(|obj| depythonize_block_mm_infos(obj.bind(py))) @@ -1939,6 +1950,7 @@ impl KvRouter { None, block_mm_infos.as_deref(), lora_name.as_deref(), + cache_namespace.as_deref(), ) .await .map_err(to_pyerr)?; @@ -1953,7 +1965,8 @@ impl KvRouter { }) } - #[pyo3(signature = (token_ids, router_config_override=None, block_mm_infos=None, lora_name=None, include_shared=true))] + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = (token_ids, router_config_override=None, block_mm_infos=None, lora_name=None, include_shared=true, cache_namespace=None))] fn get_overlap_scores<'p>( &self, py: Python<'p>, @@ -1962,6 +1975,7 @@ impl KvRouter { block_mm_infos: Option, lora_name: Option, include_shared: bool, + cache_namespace: Option, ) -> PyResult> { let router_config_override = if let Some(obj) = router_config_override { let override_config: RouterConfigOverride = @@ -1982,6 +1996,7 @@ impl KvRouter { router_config_override.as_ref(), block_mm_infos.as_deref(), lora_name.as_deref(), + cache_namespace.as_deref(), include_shared, ) .await diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index 5a4010ef9451..54f94b52315d 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -346,6 +346,7 @@ def compute_block_hash_for_seq( block_mm_infos: Optional[List[Optional[Dict[str, Any]]]] = None, lora_name: Optional[str] = None, is_eagle: Optional[bool] = None, + cache_namespace: Optional[str] = None, ) -> List[int]: """ Compute block hashes for a sequence of tokens, optionally including multimodal metadata. @@ -1097,6 +1098,7 @@ class KvEventPublisher: block_mm_infos: Optional[List[Optional[Dict[str, Any]]]] = None, lora_name: Optional[str] = None, is_eagle: Optional[bool] = None, + cache_salt: Optional[str] = None, ) -> None: """ Publish a KV stored event. @@ -2845,6 +2847,7 @@ class KvRouter: routing_constraints: Optional[RoutingConstraints] = None, strict_priority: int = 0, policy_class: Optional[str] = None, + cache_namespace: Optional[str] = None, ) -> Tuple[int, int, int]: """ Find the best matching worker for the given tokens. @@ -2862,6 +2865,7 @@ class KvRouter: block_mm_infos: Optional block-level multimodal metadata aligned to request blocks. When provided, this is used in block hash computation to enable MM-aware worker selection. + cache_namespace: Optional cache namespace used in block hash computation. policy_class: Requested policy family, or an exact explicit class. Missing, unknown, and ordinary physical-class names use the configured default family before cache-bucket resolution. @@ -2879,6 +2883,7 @@ class KvRouter: token_ids: List[int], block_mm_infos: Optional[List[Optional[Dict[str, Any]]]] = None, lora_name: Optional[str] = None, + cache_namespace: Optional[str] = None, ) -> List[Dict[str, int]]: """ Get potential prefill and decode loads for all workers. @@ -2911,6 +2916,7 @@ class KvRouter: block_mm_infos: Optional[List[Optional[Dict[str, Any]]]] = None, lora_name: Optional[str] = None, include_shared: bool = True, + cache_namespace: Optional[str] = None, ) -> Dict[str, Any]: """ Get per-worker KV overlap by storage tier. diff --git a/lib/kv-router/src/indexer/branch_sharded.rs b/lib/kv-router/src/indexer/branch_sharded.rs index 22d836a81534..18276c1dda09 100644 --- a/lib/kv-router/src/indexer/branch_sharded.rs +++ b/lib/kv-router/src/indexer/branch_sharded.rs @@ -801,6 +801,7 @@ impl KvIndexerInterface for BranchShardedIndexer { &self, tokens: &[u32], lora_name: Option<&str>, + cache_namespace: Option<&str>, is_eagle: Option, ) -> Result { let sequence = compute_block_hash_for_seq( @@ -808,6 +809,7 @@ impl KvIndexerInterface for BranchShardedIndexer { self.kv_block_size, BlockHashOptions { lora_name, + cache_namespace, is_eagle, block_mm_infos: None, }, diff --git a/lib/kv-router/src/indexer/kv_indexer.rs b/lib/kv-router/src/indexer/kv_indexer.rs index d3f6f4eadef7..153538a234bb 100644 --- a/lib/kv-router/src/indexer/kv_indexer.rs +++ b/lib/kv-router/src/indexer/kv_indexer.rs @@ -511,6 +511,7 @@ impl KvIndexerInterface for KvIndexer { &self, tokens: &[u32], lora_name: Option<&str>, + cache_namespace: Option<&str>, is_eagle: Option, ) -> Result { tracing::debug!( @@ -523,6 +524,7 @@ impl KvIndexerInterface for KvIndexer { self.kv_block_size, BlockHashOptions { lora_name, + cache_namespace, is_eagle, ..Default::default() }, diff --git a/lib/kv-router/src/indexer/local.rs b/lib/kv-router/src/indexer/local.rs index 67ad120d78af..633555ad0040 100644 --- a/lib/kv-router/src/indexer/local.rs +++ b/lib/kv-router/src/indexer/local.rs @@ -718,10 +718,11 @@ impl KvIndexerInterface for LocalKvIndexer { &self, tokens: &[u32], lora_name: Option<&str>, + cache_namespace: Option<&str>, is_eagle: Option, ) -> Result { self.indexer - .find_matches_for_request(tokens, lora_name, is_eagle) + .find_matches_for_request(tokens, lora_name, cache_namespace, is_eagle) .await } diff --git a/lib/kv-router/src/indexer/tests.rs b/lib/kv-router/src/indexer/tests.rs index 7640732cd141..07793dc8e89d 100644 --- a/lib/kv-router/src/indexer/tests.rs +++ b/lib/kv-router/src/indexer/tests.rs @@ -284,7 +284,7 @@ async fn route_approx_tokens( async fn request_scores(index: &dyn KvIndexerInterface, tokens: &[u32]) -> OverlapScores { index - .find_matches_for_request(tokens, None, None) + .find_matches_for_request(tokens, None, None, None) .await .unwrap() } @@ -834,7 +834,7 @@ mod interface_tests { let tokens: Vec = (1..=96).collect(); let scores = index - .find_matches_for_request(&tokens, None, None) + .find_matches_for_request(&tokens, None, None, None) .await .unwrap(); assert!(scores.scores.is_empty()); @@ -856,7 +856,7 @@ mod interface_tests { flush_and_settle(index.as_ref()).await; let scores = index - .find_matches_for_request(&tokens, None, None) + .find_matches_for_request(&tokens, None, None, None) .await .unwrap(); assert_eq!(scores.scores.get(&WorkerWithDpRank::new(0, 0)), Some(&3)); @@ -1534,6 +1534,78 @@ mod lora_tests { assert!(scores_b.scores.contains_key(&WorkerWithDpRank::new(1, 0))); assert!(!scores_b.scores.contains_key(&WorkerWithDpRank::new(0, 0))); } + + #[tokio::test] + #[apply(indexer_template)] + async fn test_different_cache_namespaces_do_not_conflict(variant: &str) { + let index = make_indexer(variant); + let kv_block_size: u32 = 32; + + let tokens: Vec = (0..kv_block_size * 2).collect(); + + let hashes_a = compute_block_hash_for_seq( + &tokens, + kv_block_size, + BlockHashOptions { + cache_namespace: Some("tenant-a"), + ..Default::default() + }, + ); + let hashes_b = compute_block_hash_for_seq( + &tokens, + kv_block_size, + BlockHashOptions { + cache_namespace: Some("tenant-b"), + ..Default::default() + }, + ); + + assert_ne!( + hashes_a, hashes_b, + "Different cache namespaces must produce different hashes" + ); + + let seq_a = compute_seq_hash_for_block(&hashes_a); + let seq_b = compute_seq_hash_for_block(&hashes_b); + + index + .apply_event(router_event( + 0, + 0, + 0, + KvCacheEventData::Stored(KvCacheStoreData { + parent_hash: None, + start_position: None, + blocks: stored_blocks_with_sequence_hashes(&hashes_a, &seq_a), + }), + )) + .await; + + index + .apply_event(router_event( + 1, + 0, + 0, + KvCacheEventData::Stored(KvCacheStoreData { + parent_hash: None, + start_position: None, + blocks: stored_blocks_with_sequence_hashes(&hashes_b, &seq_b), + }), + )) + .await; + + flush_and_settle(index.as_ref()).await; + + let scores_a = index.find_matches(hashes_a.clone()).await.unwrap(); + assert_eq!(scores_a.scores.len(), 1); + assert!(scores_a.scores.contains_key(&WorkerWithDpRank::new(0, 0))); + assert!(!scores_a.scores.contains_key(&WorkerWithDpRank::new(1, 0))); + + let scores_b = index.find_matches(hashes_b.clone()).await.unwrap(); + assert_eq!(scores_b.scores.len(), 1); + assert!(scores_b.scores.contains_key(&WorkerWithDpRank::new(1, 0))); + assert!(!scores_b.scores.contains_key(&WorkerWithDpRank::new(0, 0))); + } } // ============================================================================ diff --git a/lib/kv-router/src/indexer/thread_pool.rs b/lib/kv-router/src/indexer/thread_pool.rs index ace763f6070e..d0f5782d659f 100644 --- a/lib/kv-router/src/indexer/thread_pool.rs +++ b/lib/kv-router/src/indexer/thread_pool.rs @@ -599,6 +599,7 @@ impl KvIndexerInterface for ThreadPoolIndexer { &self, tokens: &[u32], lora_name: Option<&str>, + cache_namespace: Option<&str>, is_eagle: Option, ) -> Result { let sequence = compute_block_hash_for_seq( @@ -606,6 +607,7 @@ impl KvIndexerInterface for ThreadPoolIndexer { self.kv_block_size, BlockHashOptions { lora_name, + cache_namespace, is_eagle, ..Default::default() }, diff --git a/lib/kv-router/src/indexer/traits.rs b/lib/kv-router/src/indexer/traits.rs index 8efc9809e60f..c108bdd8b0c3 100644 --- a/lib/kv-router/src/indexer/traits.rs +++ b/lib/kv-router/src/indexer/traits.rs @@ -18,10 +18,14 @@ use crate::protocols::*; #[async_trait] pub trait SharedKvCache: Send + Sync { /// Query which blocks exist in the shared cache for the given token sequence. + /// + /// `cache_namespace` must be honored by implementations to avoid scoring + /// shared-cache hits across isolated cache namespaces. async fn check_blocks( &self, tokens: &[u32], block_size: u32, + cache_namespace: Option<&str>, ) -> Result; } @@ -72,6 +76,7 @@ pub trait KvIndexerInterface { /// /// * `tokens` - A vector of `u32` tokens. /// * `lora_name` - Optional LoRA adapter name to include in block hash computation. + /// * `cache_namespace` - Optional cache namespace to include in block hash computation. /// /// ### Returns /// @@ -80,6 +85,7 @@ pub trait KvIndexerInterface { &self, tokens: &[u32], lora_name: Option<&str>, + cache_namespace: Option<&str>, is_eagle: Option, ) -> Result; diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index 36187c702c47..df6b4a3590ec 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -21,6 +21,8 @@ pub const KV_EVENT_SUBJECT: &str = "kv-events"; /// Seed for XXH3 hashing, consistent with indexer.rs pub const XXH3_SEED: u64 = 1337; +const LORA_HASH_SEED: u64 = XXH3_SEED ^ 0x9e37_79b1_85eb_ca87; +const CACHE_NAMESPACE_HASH_SEED: u64 = XXH3_SEED ^ 0xc2b2_ae3d_27d4_eb4f; /// Compute the hash of a local block. pub fn compute_block_hash(data: &[u8]) -> LocalBlockHash { @@ -31,9 +33,24 @@ pub fn compute_block_hash(data: &[u8]) -> LocalBlockHash { pub struct BlockHashOptions<'a> { pub block_mm_infos: Option<&'a [Option]>, pub lora_name: Option<&'a str>, + pub cache_namespace: Option<&'a str>, pub is_eagle: Option, } +fn block_hash_seed(options: BlockHashOptions<'_>) -> u64 { + let mut seed = XXH3_SEED; + if let Some(name) = options.lora_name.filter(|n| !n.is_empty()) { + seed = seed.wrapping_add(xxh3::xxh3_64_with_seed(name.as_bytes(), LORA_HASH_SEED)); + } + if let Some(namespace) = options.cache_namespace.filter(|n| !n.is_empty()) { + seed = seed.wrapping_add(xxh3::xxh3_64_with_seed( + namespace.as_bytes(), + CACHE_NAMESPACE_HASH_SEED, + )); + } + seed +} + #[inline] fn hash_block_no_mm(chunk: &[u32], seed: u64, scratch_bytes: &mut Vec) -> LocalBlockHash { #[cfg(target_endian = "little")] @@ -70,18 +87,16 @@ pub fn pad_value_for_mm_hash(mm_hash: u64) -> u32 { (MM_PAD_SHIFT_VALUE + (mm_hash & MM_PAD_HASH_MASK)) as u32 } -/// Compute the hash for a sequence of tokens, optionally including multimodal metadata -/// and LoRA adapter identity. +/// Compute the hash for a sequence of tokens, optionally including multimodal metadata, +/// LoRA adapter identity, and cache namespace. /// /// When multimodal extra info is provided, the mm_hashes are included in the hash computation /// to ensure that blocks with identical tokens but different multimodal objects produce /// different hashes. /// -/// When `lora_name` is provided, the adapter name is mixed into the XXH3 seed so that -/// blocks cached under different LoRA adapters (or the base model) produce distinct hashes. -/// Because LoRA identity applies uniformly to every block in a sequence, encoding it in the -/// seed is more efficient than appending per-block bytes and matches the approach used by -/// KVBM's `SaltHash`. +/// When `lora_name` or `cache_namespace` is provided, those request-wide identities are +/// mixed into the XXH3 seed so blocks cached under different adapters or namespaces produce +/// distinct hashes. Empty strings are treated as absent. pub fn compute_block_hash_for_seq( tokens: &[u32], kv_block_size: u32, @@ -91,10 +106,7 @@ pub fn compute_block_hash_for_seq( return Vec::new(); } - let seed = match options.lora_name.filter(|n| !n.is_empty()) { - Some(name) => XXH3_SEED.wrapping_add(xxh3::xxh3_64(name.as_bytes())), - None => XXH3_SEED, - }; + let seed = block_hash_seed(options); let is_eagle_flag = options.is_eagle.unwrap_or(false); let stride = kv_block_size as usize; let window_size = if is_eagle_flag { stride + 1 } else { stride }; @@ -444,6 +456,8 @@ pub enum RouterRequest { strict_priority: u32, #[serde(default, skip_serializing_if = "Option::is_none")] lora_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + cache_namespace: Option, }, PotentialLoads { tokens: Vec, @@ -451,6 +465,8 @@ pub enum RouterRequest { block_mm_infos: Option>>, #[serde(default, skip_serializing_if = "Option::is_none")] lora_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + cache_namespace: Option, }, MarkPrefill { // once prefill completes, the frontend might not be allowed to send a @@ -475,6 +491,7 @@ impl Default for RouterRequest { priority_jump: 0.0, strict_priority: 0, lora_name: None, + cache_namespace: None, } } } @@ -1033,6 +1050,7 @@ pub struct TokensWithHashes { block_size: u32, block_mm_infos: Option>>, lora_name: Option, + cache_namespace: Option, block_hashes: Option>, seq_hashes: Option>, is_eagle: Option, @@ -1046,6 +1064,7 @@ impl TokensWithHashes { block_size, block_mm_infos: None, lora_name: None, + cache_namespace: None, block_hashes: None, seq_hashes: None, is_eagle: None, @@ -1066,6 +1085,13 @@ impl TokensWithHashes { self } + /// Sets the cache namespace for hash computation. + pub fn with_cache_namespace(mut self, namespace: String) -> Self { + self.cache_namespace = Some(namespace); + self.invalidate_hashes(); + self + } + /// Sets Eagle hashing semantics for this token sequence. pub fn with_is_eagle(mut self, is_eagle: bool) -> Self { self.set_is_eagle(is_eagle); @@ -1122,6 +1148,7 @@ impl TokensWithHashes { BlockHashOptions { block_mm_infos: self.block_mm_infos.as_deref(), lora_name: self.lora_name.as_deref(), + cache_namespace: self.cache_namespace.as_deref(), is_eagle: self.is_eagle, }, )); @@ -1295,6 +1322,62 @@ mod tests { ); } + #[test] + fn test_cache_namespace_produces_different_hash() { + let tokens: Vec = (0..4).collect(); + let base = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()); + let namespace_a = compute_block_hash_for_seq( + &tokens, + 4, + BlockHashOptions { + cache_namespace: Some("tenant-a"), + ..Default::default() + }, + ); + let namespace_b = compute_block_hash_for_seq( + &tokens, + 4, + BlockHashOptions { + cache_namespace: Some("tenant-b"), + ..Default::default() + }, + ); + let lora_a = compute_block_hash_for_seq( + &tokens, + 4, + BlockHashOptions { + lora_name: Some("tenant-a"), + ..Default::default() + }, + ); + + assert_ne!(base[0], namespace_a[0]); + assert_ne!(base[0], namespace_b[0]); + assert_ne!(namespace_a[0], namespace_b[0]); + assert_ne!( + namespace_a[0], lora_a[0], + "namespace and lora salts must use independent seed domains" + ); + } + + #[test] + fn test_cache_namespace_empty_string_normalized_to_none() { + let tokens: Vec = (0..4).collect(); + let base = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()); + let empty = compute_block_hash_for_seq( + &tokens, + 4, + BlockHashOptions { + cache_namespace: Some(""), + ..Default::default() + }, + ); + assert_eq!( + base, empty, + "empty cache_namespace should be treated as absent" + ); + } + #[test] fn test_tokens_with_hashes_lora() { let tokens: Vec = (0..8).collect(); @@ -1369,6 +1452,47 @@ mod tests { assert_ne!(actual_sequence_hashes, text_sequence_hashes); } + #[test] + fn test_tokens_with_hashes_cache_namespace() { + let tokens: Vec = (0..8).collect(); + + let mut base = TokensWithHashes::new(tokens.clone(), 4); + let base_hashes = base.get_or_compute_block_hashes().to_vec(); + + let mut with_namespace = + TokensWithHashes::new(tokens, 4).with_cache_namespace("tenant-a".to_string()); + let namespace_hashes = with_namespace.get_or_compute_block_hashes().to_vec(); + + assert_eq!(base_hashes.len(), namespace_hashes.len()); + for (base, namespaced) in base_hashes.iter().zip(namespace_hashes.iter()) { + assert_ne!(base, namespaced); + } + } + + #[test] + fn test_tokens_with_hashes_cache_namespace_change_recomputes_cached_hashes() { + let tokens: Vec = (0..8).collect(); + let mut with_hashes = TokensWithHashes::new(tokens.clone(), 4); + let base_sequence_hashes = with_hashes.get_or_compute_seq_hashes().to_vec(); + + let mut with_hashes = with_hashes.with_cache_namespace("tenant-a".to_string()); + let actual_block_hashes = with_hashes.get_or_compute_block_hashes().to_vec(); + let actual_sequence_hashes = with_hashes.get_or_compute_seq_hashes().to_vec(); + let expected_block_hashes = compute_block_hash_for_seq( + &tokens, + 4, + BlockHashOptions { + cache_namespace: Some("tenant-a"), + ..Default::default() + }, + ); + let expected_sequence_hashes = compute_seq_hash_for_block(&expected_block_hashes); + + assert_eq!(actual_block_hashes, expected_block_hashes); + assert_eq!(actual_sequence_hashes, expected_sequence_hashes); + assert_ne!(actual_sequence_hashes, base_sequence_hashes); + } + #[test] fn test_compute_block_hash_for_seq_eagle_windows() { let tokens: Vec = (0..6).collect(); @@ -1607,6 +1731,7 @@ mod tests { priority_jump: 5.0, strict_priority: 0, lora_name: None, + cache_namespace: None, }; let serialized = serde_json::to_string(&request).unwrap(); @@ -1634,6 +1759,7 @@ mod tests { priority_jump: 0.0, strict_priority: 0, lora_name: Some("adapter-a".to_string()), + cache_namespace: None, }; let serialized = serde_json::to_string(&request).unwrap(); @@ -1677,6 +1803,7 @@ mod tests { priority_jump: 0.0, strict_priority: 4, lora_name: None, + cache_namespace: None, }; let serialized = serde_json::to_string(&request).unwrap(); @@ -1702,6 +1829,7 @@ mod tests { priority_jump: 0.0, strict_priority: 0, lora_name: None, + cache_namespace: None, }; assert_eq!( serde_json::to_string(&zero).unwrap(), @@ -1709,12 +1837,42 @@ mod tests { ); } + #[test] + fn test_router_request_new_serialization_with_cache_namespace() { + let request = RouterRequest::New { + tokens: vec![1, 2, 3], + block_mm_infos: None, + routing_constraints: RoutingConstraints::default(), + priority_jump: 0.0, + strict_priority: 0, + lora_name: None, + cache_namespace: Some("tenant-a".to_string()), + }; + + let serialized = serde_json::to_string(&request).unwrap(); + let deserialized: RouterRequest = serde_json::from_str(&serialized).unwrap(); + + assert_eq!( + serialized, + r#"{"method":"new","tokens":[1,2,3],"priority_jump":0.0,"cache_namespace":"tenant-a"}"# + ); + assert!(matches!( + deserialized, + RouterRequest::New { + tokens, + cache_namespace: Some(ref cache_namespace), + .. + } if tokens == vec![1, 2, 3] && cache_namespace == "tenant-a" + )); + } + #[test] fn test_router_request_potential_loads_serialization_with_lora_name() { let request = RouterRequest::PotentialLoads { tokens: vec![1, 2, 3], block_mm_infos: None, lora_name: Some("adapter-a".to_string()), + cache_namespace: None, }; let serialized = serde_json::to_string(&request).unwrap(); @@ -1730,10 +1888,37 @@ mod tests { tokens, block_mm_infos: None, lora_name: Some(ref lora_name), + cache_namespace: None, } if tokens == vec![1, 2, 3] && lora_name == "adapter-a" )); } + #[test] + fn test_router_request_potential_loads_serialization_with_cache_namespace() { + let request = RouterRequest::PotentialLoads { + tokens: vec![1, 2, 3], + block_mm_infos: None, + lora_name: None, + cache_namespace: Some("tenant-a".to_string()), + }; + + let serialized = serde_json::to_string(&request).unwrap(); + let deserialized: RouterRequest = serde_json::from_str(&serialized).unwrap(); + + assert_eq!( + serialized, + r#"{"method":"potential_loads","tokens":[1,2,3],"cache_namespace":"tenant-a"}"# + ); + assert!(matches!( + deserialized, + RouterRequest::PotentialLoads { + tokens, + cache_namespace: Some(ref cache_namespace), + .. + } if tokens == vec![1, 2, 3] && cache_namespace == "tenant-a" + )); + } + #[test] fn test_router_request_potential_loads_defaults_lora_name() { let deserialized: RouterRequest = @@ -1745,6 +1930,7 @@ mod tests { tokens, block_mm_infos: None, lora_name: None, + cache_namespace: None, } if tokens == vec![1, 2, 3] )); } diff --git a/lib/kv-router/src/services/indexer/server.rs b/lib/kv-router/src/services/indexer/server.rs index a62fa6d9a67e..a29fdc4b03d2 100644 --- a/lib/kv-router/src/services/indexer/server.rs +++ b/lib/kv-router/src/services/indexer/server.rs @@ -127,8 +127,7 @@ pub struct QueryRequest { pub tenant_id: String, #[serde(default)] pub lora_name: Option, - /// Optional per-request cache salt (Mooncake RFC #1403). Currently accepted - /// but not yet mixed into hashes — engines apply their own internally. + /// Optional per-request cache salt (Mooncake RFC #1403), mixed into `/query` hashes. #[serde(default)] pub cache_salt: Option, } @@ -139,8 +138,8 @@ pub struct QueryByHashRequest { pub model_name: String, #[serde(default = "default_tenant")] pub tenant_id: String, - /// Optional per-request cache salt (Mooncake RFC #1403). Currently accepted - /// but not yet mixed into hashes — engines apply their own internally. + /// Optional per-request cache salt (Mooncake RFC #1403). For `/query_by_hash`, callers + /// must precompute `block_hashes` with the same salt. #[serde(default)] pub cache_salt: Option, } @@ -339,6 +338,7 @@ async fn query(State(state): State>, Json(req): Json block_size, BlockHashOptions { lora_name: req.lora_name.as_deref(), + cache_namespace: req.cache_salt.as_deref(), ..Default::default() }, ); diff --git a/lib/kv-router/src/services/selection/input.rs b/lib/kv-router/src/services/selection/input.rs index 0d30d416252e..923dca203285 100644 --- a/lib/kv-router/src/services/selection/input.rs +++ b/lib/kv-router/src/services/selection/input.rs @@ -130,6 +130,7 @@ fn normalize_tokens( BlockHashOptions { block_mm_infos, lora_name, + cache_namespace: None, is_eagle: Some(is_eagle), }, ); diff --git a/lib/kv-router/src/zmq_wire/convert.rs b/lib/kv-router/src/zmq_wire/convert.rs index e7cf6be000fe..45ca44e4be2f 100644 --- a/lib/kv-router/src/zmq_wire/convert.rs +++ b/lib/kv-router/src/zmq_wire/convert.rs @@ -37,6 +37,7 @@ pub fn convert_event( token_ids, block_size, lora_name, + cache_namespace, block_mm_infos, medium: _, is_eagle, @@ -90,6 +91,7 @@ pub fn convert_event( &num_block_tokens, &block_hashes_u64, lora_name.as_deref(), + cache_namespace.as_deref(), warning_count, block_mm_infos.as_deref(), is_eagle, @@ -175,15 +177,29 @@ fn substitute_pad_values(token_ids: &[u32], image_token_id: u32, mm_objects: &[u out } +#[derive(Default)] +pub struct StoredBlockOptions<'a> { + pub lora_name: Option<&'a str>, + pub cache_namespace: Option<&'a str>, + pub mm_extra_info: Option, + pub is_eagle: Option, + pub image_token_id: Option, +} + pub fn create_stored_block_from_parts( kv_block_size: u32, block_hash: u64, token_ids: &[u32], - lora_name: Option<&str>, - mm_extra_info: Option, - is_eagle: Option, - image_token_id: Option, + options: StoredBlockOptions<'_>, ) -> KvCacheStoredBlockData { + let StoredBlockOptions { + lora_name, + cache_namespace, + mm_extra_info, + is_eagle, + image_token_id, + } = options; + // When the model has a routing image token and this block carries mm // objects (vLLM events), normalize to the canonical pad_value scheme: // substitute pad_value over the image_token_id runs and hash WITHOUT @@ -200,6 +216,7 @@ pub fn create_stored_block_from_parts( BlockHashOptions { block_mm_infos: None, lora_name, + cache_namespace, is_eagle, }, )[0] @@ -212,6 +229,7 @@ pub fn create_stored_block_from_parts( BlockHashOptions { block_mm_infos: block_mm_infos.as_deref(), lora_name, + cache_namespace, is_eagle, }, )[0] @@ -240,6 +258,7 @@ pub fn create_stored_blocks( num_block_tokens: &[u64], block_hashes: &[u64], lora_name: Option<&str>, + cache_namespace: Option<&str>, warning_count: &Arc, block_mm_infos: Option<&[Option]>, is_eagle: Option, @@ -285,10 +304,13 @@ pub fn create_stored_blocks( kv_block_size, *block_hash_it, tokens, - lora_name, - mm_extra_info, - is_eagle, - image_token_id, + StoredBlockOptions { + lora_name, + cache_namespace, + mm_extra_info, + is_eagle, + image_token_id, + }, )); token_offset += *num_tokens_it as usize; } @@ -322,10 +344,11 @@ mod normalize_tests { block_size, 0xabcd, &vllm_tokens, - None, - Some(mm_info), - None, - Some(image_token_id), + StoredBlockOptions { + mm_extra_info: Some(mm_info), + image_token_id: Some(image_token_id), + ..Default::default() + }, ); // Frontend side: same tokens but image positions already pad_value, @@ -354,13 +377,13 @@ mod normalize_tests { block_size, 0x1, &tokens, - None, - None, - None, - Some(151655), + StoredBlockOptions { + image_token_id: Some(151655), + ..Default::default() + }, ); let without = - create_stored_block_from_parts(block_size, 0x1, &tokens, None, None, None, None); + create_stored_block_from_parts(block_size, 0x1, &tokens, StoredBlockOptions::default()); assert_eq!(with_img.tokens_hash, without.tokens_hash); } } diff --git a/lib/kv-router/src/zmq_wire/deserialize.rs b/lib/kv-router/src/zmq_wire/deserialize.rs index d10a53a7f451..92cc8784c4de 100644 --- a/lib/kv-router/src/zmq_wire/deserialize.rs +++ b/lib/kv-router/src/zmq_wire/deserialize.rs @@ -9,7 +9,7 @@ use serde::de::{self, IgnoredAny, MapAccess, SeqAccess, Visitor}; use crate::protocols::BlockExtraInfo; -use super::extra_keys::extra_keys_to_block_mm_infos; +use super::extra_keys::{extra_keys_to_block_mm_infos, extra_keys_to_cache_namespace}; use super::filter::{BlockStoredTrailingField, KvCacheEventMetadata, KvCacheEventTrailingField}; use super::types::{BlockHashValue, ExtraKeyItem, KvTokenIds, RawKvEvent}; @@ -49,6 +49,7 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { let mut block_size: Option = None; let mut medium: Option> = None; let mut lora_name: Option> = None; + let mut cache_namespace: Option> = None; let mut extra_keys: Option>>>> = None; let mut block_mm_infos: Option>>> = None; let mut metadata = KvCacheEventMetadata::default(); @@ -76,6 +77,9 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { "lora_name" => { lora_name = Some(map.next_value()?); } + "cache_salt" => { + cache_namespace = Some(map.next_value()?); + } "extra_keys" => { extra_keys = Some(map.next_value()?); } @@ -106,16 +110,22 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { let block_size = block_size.ok_or_else(|| de::Error::missing_field("block_size"))?; let medium = medium.unwrap_or(None); + let lora_name = lora_name.unwrap_or(None); + let extra_keys = extra_keys.unwrap_or(None); + let cache_namespace = cache_namespace.unwrap_or(None).or_else(|| { + extra_keys_to_cache_namespace(extra_keys.as_deref(), lora_name.as_deref()) + }); let block_mm_infos = block_mm_infos .unwrap_or(None) - .or_else(|| extra_keys_to_block_mm_infos(extra_keys.unwrap_or(None))); + .or_else(|| extra_keys_to_block_mm_infos(extra_keys)); Ok(RawKvEvent::BlockStored { block_hashes, parent_block_hash: parent_block_hash.unwrap_or(None), token_ids: raw_token_ids, block_size, medium, - lora_name: lora_name.unwrap_or(None), + lora_name, + cache_namespace, block_mm_infos, is_eagle: Some(is_eagle), group_idx: metadata.group_idx, @@ -194,6 +204,8 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { while seq.next_element::()?.is_some() {} + let cache_namespace = + extra_keys_to_cache_namespace(extra_keys.as_deref(), lora_name.as_deref()); let block_mm_infos = block_mm_infos.or_else(|| extra_keys_to_block_mm_infos(extra_keys)); let (raw_token_ids, is_eagle) = normalize_token_ids(token_ids); @@ -205,6 +217,7 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { block_size, medium, lora_name, + cache_namespace, block_mm_infos, is_eagle: Some(is_eagle), group_idx: metadata.group_idx, diff --git a/lib/kv-router/src/zmq_wire/extra_keys.rs b/lib/kv-router/src/zmq_wire/extra_keys.rs index df33469a5660..ff26aa373d45 100644 --- a/lib/kv-router/src/zmq_wire/extra_keys.rs +++ b/lib/kv-router/src/zmq_wire/extra_keys.rs @@ -17,6 +17,44 @@ pub fn parse_mm_hash_from_extra_key(s: &str) -> Option { None } +fn cache_namespace_candidate<'a>(value: &'a str, lora_name: Option<&str>) -> Option<&'a str> { + if value.is_empty() { + return None; + } + if lora_name.is_some_and(|name| name == value) { + return None; + } + if parse_mm_hash_from_extra_key(value).is_some() { + return None; + } + Some(value) +} + +/// Extract a vLLM cache salt from `extra_keys` when a producer does not emit +/// top-level `cache_salt`. vLLM aligns `extra_keys` with blocks and includes +/// request-wide extras in each block, so the first block is enough. +pub fn extra_keys_to_cache_namespace( + extra_keys: Option<&[Option>]>, + lora_name: Option<&str>, +) -> Option { + let first_block = extra_keys?.first()?.as_ref()?; + first_block.iter().find_map(|key| match key { + ExtraKeyItem::Hash(hash) + | ExtraKeyItem::HashWithSignedOffset((hash, _)) + | ExtraKeyItem::HashWithUnsignedOffset((hash, _)) => { + cache_namespace_candidate(hash, lora_name).map(str::to_owned) + } + ExtraKeyItem::Bytes(bytes) => std::str::from_utf8(bytes) + .ok() + .and_then(|value| cache_namespace_candidate(value, lora_name)) + .map(str::to_owned), + ExtraKeyItem::Signed(_) + | ExtraKeyItem::Unsigned(_) + | ExtraKeyItem::Float(_) + | ExtraKeyItem::Bool(_) => None, + }) +} + /// Convert vLLM BlockStored extra_keys to block-level MM infos. /// extra_keys is a list aligned with blocks: /// - None => no MM content in that block diff --git a/lib/kv-router/src/zmq_wire/mod.rs b/lib/kv-router/src/zmq_wire/mod.rs index 55251265ec78..4cb83b524516 100644 --- a/lib/kv-router/src/zmq_wire/mod.rs +++ b/lib/kv-router/src/zmq_wire/mod.rs @@ -23,8 +23,12 @@ mod filter; mod tests; mod types; -pub use convert::{convert_event, create_stored_block_from_parts, create_stored_blocks}; -pub use extra_keys::{extra_keys_to_block_mm_infos, parse_mm_hash_from_extra_key}; +pub use convert::{ + StoredBlockOptions, convert_event, create_stored_block_from_parts, create_stored_blocks, +}; +pub use extra_keys::{ + extra_keys_to_block_mm_infos, extra_keys_to_cache_namespace, parse_mm_hash_from_extra_key, +}; pub use filter::KvCacheSpecKind; pub use types::{BlockHashValue, ExtraKeyItem, KvEventBatch, KvTokenIds, RawKvEvent}; @@ -43,6 +47,7 @@ pub struct ZmqEventNormalizer { image_token_id: Option, warning_count: Arc, group_metadata: FxHashMap<(DpRank, u32), KvCacheGroupMetadata>, + cache_namespaces: FxHashMap, } #[derive(Debug, Clone, Copy)] @@ -79,6 +84,7 @@ impl ZmqEventNormalizer { image_token_id: None, warning_count: Arc::new(AtomicU32::new(0)), group_metadata: FxHashMap::default(), + cache_namespaces: FxHashMap::default(), } } @@ -88,6 +94,7 @@ impl ZmqEventNormalizer { image_token_id: None, warning_count, group_metadata: FxHashMap::default(), + cache_namespaces: FxHashMap::default(), } } @@ -105,7 +112,7 @@ impl ZmqEventNormalizer { pub fn preprocess_with_reason( &mut self, - raw: RawKvEvent, + mut raw: RawKvEvent, worker: WorkerWithDpRank, ) -> Result { if raw.is_ignored() { @@ -119,6 +126,7 @@ impl ZmqEventNormalizer { if let Some(reason) = self.filter_reason(metadata, worker.dp_rank) { return Err(reason); } + self.propagate_cache_namespace(&mut raw); Ok(raw) } @@ -163,6 +171,45 @@ impl ZmqEventNormalizer { ); } + fn propagate_cache_namespace(&mut self, raw: &mut RawKvEvent) { + match raw { + RawKvEvent::BlockStored { + block_hashes, + parent_block_hash, + cache_namespace, + .. + } => { + if cache_namespace.is_none() + && let Some(parent) = parent_block_hash.as_ref() + && let Some(namespace) = + self.cache_namespaces.get(&(*parent).into_u64()).cloned() + { + *cache_namespace = Some(namespace); + } + + if let Some(namespace) = cache_namespace.as_ref().filter(|ns| !ns.is_empty()) { + for block_hash in block_hashes.iter() { + self.cache_namespaces + .insert((*block_hash).into_u64(), namespace.clone()); + } + } else { + for block_hash in block_hashes.iter() { + self.cache_namespaces.remove(&(*block_hash).into_u64()); + } + } + } + RawKvEvent::BlockRemoved { block_hashes, .. } => { + for block_hash in block_hashes.iter() { + self.cache_namespaces.remove(&(*block_hash).into_u64()); + } + } + RawKvEvent::AllBlocksCleared => { + self.cache_namespaces.clear(); + } + RawKvEvent::Ignored => {} + } + } + fn filter_reason( &self, metadata: KvCacheEventMetadata, diff --git a/lib/kv-router/src/zmq_wire/tests.rs b/lib/kv-router/src/zmq_wire/tests.rs index 82a495d8b09c..ee0d9df34c24 100644 --- a/lib/kv-router/src/zmq_wire/tests.rs +++ b/lib/kv-router/src/zmq_wire/tests.rs @@ -4,7 +4,8 @@ use std::sync::Arc; use std::sync::atomic::{AtomicU32, Ordering}; -use rmp_serde::{from_slice, to_vec}; +use rmp_serde::{from_slice, to_vec, to_vec_named}; +use serde::Serialize; use crate::protocols::{ BlockExtraInfo, BlockHashOptions, BlockMmObjectInfo, ExternalSequenceBlockHash, @@ -50,6 +51,80 @@ fn test_deserialize_bigram_block_stored_sequence() { } } +#[derive(Serialize)] +struct MapBlockStoredFixture { + #[serde(rename = "type")] + event_type: &'static str, + block_hashes: Vec, + parent_block_hash: Option, + token_ids: Vec, + block_size: usize, + medium: Option, + lora_name: Option, + #[serde(skip_serializing_if = "Option::is_none")] + cache_salt: Option, + #[serde(skip_serializing_if = "Option::is_none")] + extra_keys: Option>>>, +} + +impl Default for MapBlockStoredFixture { + fn default() -> Self { + Self { + event_type: "BlockStored", + block_hashes: vec![BlockHashValue::Unsigned(11)], + parent_block_hash: None, + token_ids: vec![10, 11], + block_size: 2, + medium: None, + lora_name: None, + cache_salt: None, + extra_keys: None, + } + } +} + +#[test] +fn test_deserialize_map_block_stored_cache_salt() { + let encoded = to_vec_named(&MapBlockStoredFixture { + cache_salt: Some("tenant-a".to_string()), + ..Default::default() + }) + .unwrap(); + let event: RawKvEvent = from_slice(&encoded).unwrap(); + + let RawKvEvent::BlockStored { + cache_namespace, .. + } = event + else { + panic!("expected BlockStored"); + }; + assert_eq!(cache_namespace.as_deref(), Some("tenant-a")); +} + +#[test] +fn test_deserialize_extra_keys_cache_namespace_fallback() { + let mm_hash = "0123456789abcdef00112233445566778899aabbccddeefffedcba9876543210"; + let encoded = to_vec_named(&MapBlockStoredFixture { + lora_name: Some("adapter-a".to_string()), + extra_keys: Some(vec![Some(vec![ + mm_hash.to_string(), + "adapter-a".to_string(), + "tenant-a".to_string(), + ])]), + ..Default::default() + }) + .unwrap(); + let event: RawKvEvent = from_slice(&encoded).unwrap(); + + let RawKvEvent::BlockStored { + cache_namespace, .. + } = event + else { + panic!("expected BlockStored"); + }; + assert_eq!(cache_namespace.as_deref(), Some("tenant-a")); +} + fn block_stored_sequence( group_idx: Option, kv_cache_spec_kind: Option<&'static str>, @@ -431,6 +506,51 @@ fn test_normalizer_does_not_learn_metadata_from_remove_events() { assert!(normalizer.preprocess(bare_remove, worker).is_none()); } +#[test] +fn test_normalizer_propagates_cache_namespace_from_parent() { + let worker = WorkerWithDpRank::new(7, 0); + let mut normalizer = ZmqEventNormalizer::new(2); + let parent = RawKvEvent::BlockStored { + block_hashes: vec![BlockHashValue::Unsigned(1)], + parent_block_hash: None, + token_ids: vec![10, 11], + block_size: 2, + medium: None, + lora_name: None, + cache_namespace: Some("tenant-a".to_string()), + block_mm_infos: None, + is_eagle: Some(false), + group_idx: None, + kv_cache_spec_kind: None, + kv_cache_spec_sliding_window: None, + }; + let child = RawKvEvent::BlockStored { + block_hashes: vec![BlockHashValue::Unsigned(2)], + parent_block_hash: Some(BlockHashValue::Unsigned(1)), + token_ids: vec![12, 13], + block_size: 2, + medium: None, + lora_name: None, + cache_namespace: None, + block_mm_infos: None, + is_eagle: Some(false), + group_idx: None, + kv_cache_spec_kind: None, + kv_cache_spec_sliding_window: None, + }; + + assert!(normalizer.preprocess(parent, worker).is_some()); + let child = normalizer.preprocess(child, worker).unwrap(); + + let RawKvEvent::BlockStored { + cache_namespace, .. + } = child + else { + panic!("expected BlockStored"); + }; + assert_eq!(cache_namespace.as_deref(), Some("tenant-a")); +} + #[test] fn test_normalizer_ignores_non_main_attention_kind_with_group_idx_zero() { let raw_event: RawKvEvent = from_slice(&sequence_with_cache_spec_kind( @@ -457,6 +577,7 @@ fn test_convert_event_bigram_emits_eagle_windows() { block_size: 2, medium: None, lora_name: None, + cache_namespace: None, block_mm_infos: None, is_eagle: Some(true), group_idx: None, @@ -529,6 +650,7 @@ fn cpu_block_stored(fixture: CpuBlockStoredFixture<'_>) -> RawKvEvent { block_size: fixture.block_size, medium: Some("CPU".to_string()), lora_name: None, + cache_namespace: None, block_mm_infos: None, is_eagle: None, group_idx: None, diff --git a/lib/kv-router/src/zmq_wire/types.rs b/lib/kv-router/src/zmq_wire/types.rs index a0d7ff7a3975..79a700ccad0c 100644 --- a/lib/kv-router/src/zmq_wire/types.rs +++ b/lib/kv-router/src/zmq_wire/types.rs @@ -69,6 +69,14 @@ pub enum RawKvEvent { /// LoRA adapter name for adapter-aware block hashing #[serde(default, skip_serializing_if = "Option::is_none")] lora_name: Option, + /// Cache namespace for salted block hashing. The wire field remains `cache_salt` + /// to match backend event schemas. + #[serde( + default, + rename = "cache_salt", + skip_serializing_if = "Option::is_none" + )] + cache_namespace: Option, /// Multimodal extra info for each block (length should match block_hashes) #[serde(default, skip_serializing_if = "Option::is_none")] block_mm_infos: Option>>, diff --git a/lib/kv-router/tests/standalone_indexer_http.rs b/lib/kv-router/tests/standalone_indexer_http.rs index 487c29a691c2..2b86ff5b4cdf 100644 --- a/lib/kv-router/tests/standalone_indexer_http.rs +++ b/lib/kv-router/tests/standalone_indexer_http.rs @@ -379,6 +379,7 @@ fn raw_block_stored( block_size, medium: Some(medium.to_string()), lora_name: None, + cache_namespace: None, block_mm_infos: None, is_eagle: None, group_idx: None, diff --git a/lib/kvbm-consolidator/tests/dedup.rs b/lib/kvbm-consolidator/tests/dedup.rs index 024d9f79ddfc..cd9b40ab99bf 100644 --- a/lib/kvbm-consolidator/tests/dedup.rs +++ b/lib/kvbm-consolidator/tests/dedup.rs @@ -49,6 +49,7 @@ fn bs_event( block_size, lora_name, medium: None, + cache_namespace: None, block_mm_infos: None, is_eagle: None, group_idx: None, diff --git a/lib/kvbm-consolidator/tests/e2e.rs b/lib/kvbm-consolidator/tests/e2e.rs index 33dbb1030482..79bf827b870c 100644 --- a/lib/kvbm-consolidator/tests/e2e.rs +++ b/lib/kvbm-consolidator/tests/e2e.rs @@ -62,6 +62,7 @@ fn make_synthetic_payload_blobs() -> Vec> { block_size: 4, lora_name: None, medium: None, + cache_namespace: None, block_mm_infos: None, is_eagle: None, group_idx: None, @@ -80,6 +81,7 @@ fn make_synthetic_payload_blobs() -> Vec> { block_size: 4, lora_name: None, medium: None, + cache_namespace: None, block_mm_infos: None, is_eagle: None, group_idx: None, diff --git a/lib/kvbm-consolidator/tests/kvbm_bridge.rs b/lib/kvbm-consolidator/tests/kvbm_bridge.rs index b6050dabfade..3cff1d646ac4 100644 --- a/lib/kvbm-consolidator/tests/kvbm_bridge.rs +++ b/lib/kvbm-consolidator/tests/kvbm_bridge.rs @@ -51,6 +51,7 @@ fn bs_event(hashes: Vec, tokens: Vec, block_size: usize) -> RawKvEvent block_size, lora_name: None, medium: None, + cache_namespace: None, block_mm_infos: None, is_eagle: None, group_idx: None, diff --git a/lib/kvbm-consolidator/tests/lifecycle.rs b/lib/kvbm-consolidator/tests/lifecycle.rs index 25614e29b6fd..a8cc0fa974e9 100644 --- a/lib/kvbm-consolidator/tests/lifecycle.rs +++ b/lib/kvbm-consolidator/tests/lifecycle.rs @@ -21,6 +21,7 @@ fn bs(hash: u64, tokens: Vec, block_size: usize) -> RawKvEvent { block_size, lora_name: None, medium: None, + cache_namespace: None, block_mm_infos: None, is_eagle: None, group_idx: None, diff --git a/lib/kvbm-consolidator/tests/output_contract.rs b/lib/kvbm-consolidator/tests/output_contract.rs index 33782150f5ed..fbf613917f40 100644 --- a/lib/kvbm-consolidator/tests/output_contract.rs +++ b/lib/kvbm-consolidator/tests/output_contract.rs @@ -24,6 +24,7 @@ fn bs(hash: u64, parent: Option, tokens: Vec, block_size: usize) -> Ra block_size, lora_name: None, medium: None, + cache_namespace: None, block_mm_infos: None, is_eagle: None, group_idx: None, @@ -40,6 +41,7 @@ fn bs_lora(hash: u64, tokens: Vec, lora_name: String) -> RawKvEvent { block_size: 4, lora_name: Some(lora_name), medium: None, + cache_namespace: None, block_mm_infos: None, is_eagle: None, group_idx: None, diff --git a/lib/kvbm-consolidator/tests/zmq_ingress.rs b/lib/kvbm-consolidator/tests/zmq_ingress.rs index 7666490ac1fc..ca1549e486a6 100644 --- a/lib/kvbm-consolidator/tests/zmq_ingress.rs +++ b/lib/kvbm-consolidator/tests/zmq_ingress.rs @@ -32,6 +32,7 @@ fn bs_event( block_size, lora_name, medium: None, + cache_namespace: None, block_mm_infos: None, is_eagle: None, group_idx: None, diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index fed1dd7bdde2..a40873b6fbd7 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -478,6 +478,7 @@ where update_states: bool, return_routing_hashes: bool, lora_name: Option, + cache_namespace: Option, priority_jump: f64, strict_priority: u32, expected_output_tokens: Option, @@ -493,6 +494,7 @@ where update_states, return_routing_hashes, lora_name, + cache_namespace, priority_jump, strict_priority, None, @@ -515,6 +517,7 @@ where update_states: bool, return_routing_hashes: bool, lora_name: Option, + cache_namespace: Option, priority_jump: f64, strict_priority: u32, policy_class: Option, @@ -543,6 +546,7 @@ where let hash_options = BlockHashOptions { block_mm_infos, lora_name: lora_name.as_deref(), + cache_namespace: cache_namespace.as_deref(), is_eagle: Some(self.is_eagle), }; @@ -577,6 +581,7 @@ where tokens, self.block_size, block_hashes, + cache_namespace.as_deref(), retain_block_hashes, ) .await?; @@ -698,6 +703,7 @@ where router_config_override: Option<&RouterConfigOverride>, update_states: bool, lora_name: Option, + cache_namespace: Option, priority_jump: f64, strict_priority: u32, expected_output_tokens: Option, @@ -713,6 +719,7 @@ where update_states, false, lora_name, + cache_namespace, priority_jump, strict_priority, expected_output_tokens, @@ -746,12 +753,14 @@ where expected_output_tokens: Option, worker: WorkerWithDpRank, lora_name: Option, + cache_namespace: Option, router_config_override: Option<&RouterConfigOverride>, ) { let isl_tokens = tokens.len(); let hash_options = BlockHashOptions { block_mm_infos, lora_name: lora_name.as_deref(), + cache_namespace: cache_namespace.as_deref(), is_eagle: Some(self.is_eagle), }; @@ -873,9 +882,10 @@ where block_mm_infos: Option<&[Option]>, worker: WorkerWithDpRank, lora_name: Option<&str>, + cache_namespace: Option<&str>, ) -> Result { Ok(self - .get_cache_hit_estimate(tokens, block_mm_infos, worker, lora_name) + .get_cache_hit_estimate(tokens, block_mm_infos, worker, lora_name, cache_namespace) .await? .rounded_overlap_blocks()) } @@ -886,10 +896,18 @@ where block_mm_infos: Option<&[Option]>, worker: WorkerWithDpRank, lora_name: Option<&str>, + cache_namespace: Option<&str>, ) -> Result { - self.get_cache_hit_estimate_with_hashes(tokens, block_mm_infos, worker, lora_name, false) - .await - .map(|(estimate, _)| estimate) + self.get_cache_hit_estimate_with_hashes( + tokens, + block_mm_infos, + worker, + lora_name, + cache_namespace, + false, + ) + .await + .map(|(estimate, _)| estimate) } pub(crate) async fn get_cache_hit_estimate_with_hashes( @@ -898,6 +916,7 @@ where block_mm_infos: Option<&[Option]>, worker: WorkerWithDpRank, lora_name: Option<&str>, + cache_namespace: Option<&str>, return_routing_hashes: bool, ) -> Result<(WorkerCacheHitEstimate, Option), KvRouterError> { let block_hashes = compute_block_hash_for_seq( @@ -906,6 +925,7 @@ where BlockHashOptions { block_mm_infos, lora_name, + cache_namespace, is_eagle: Some(self.is_eagle), }, ); @@ -932,11 +952,13 @@ where router_config_override: Option<&RouterConfigOverride>, block_mm_infos: Option<&[Option]>, lora_name: Option<&str>, + cache_namespace: Option<&str>, ) -> Result> { let isl_tokens = tokens.len(); let hash_options = BlockHashOptions { block_mm_infos, lora_name, + cache_namespace, is_eagle: Some(self.is_eagle), }; let block_hashes = compute_block_hash_for_seq(tokens, self.block_size, hash_options); @@ -973,11 +995,13 @@ where router_config_override: Option<&RouterConfigOverride>, block_mm_infos: Option<&[Option]>, lora_name: Option<&str>, + cache_namespace: Option<&str>, include_shared: bool, ) -> Result { let hash_options = BlockHashOptions { block_mm_infos, lora_name, + cache_namespace, is_eagle: Some(self.is_eagle), }; let block_hashes = compute_block_hash_for_seq(tokens, self.block_size, hash_options); @@ -987,7 +1011,10 @@ where let (shared_hits, shared_error) = if include_shared { if let Some(shared_cache) = self.shared_cache.as_ref() { - match shared_cache.check_blocks(tokens, self.block_size).await { + match shared_cache + .check_blocks(tokens, self.block_size, cache_namespace) + .await + { Ok(hits) => (Some(hits), None), Err(err) => { tracing::warn!(error = %err, "Shared cache overlap query failed"); @@ -1056,6 +1083,7 @@ where priority_jump, strict_priority, lora_name, + cache_namespace, } => { let request_context = ctx.context(); let mut schedule = Box::pin(self.find_best_match_details_with_policy_class( @@ -1066,6 +1094,7 @@ where true, false, lora_name, + cache_namespace, priority_jump, strict_priority, policy_class, @@ -1113,6 +1142,7 @@ where tokens, block_mm_infos, lora_name, + cache_namespace, } => RouterResponse::PotentialLoads { loads: self .get_potential_loads( @@ -1120,6 +1150,7 @@ where None, block_mm_infos.as_deref(), lora_name.as_deref(), + cache_namespace.as_deref(), ) .await?, pending_count: self.pending_count(), @@ -1229,6 +1260,7 @@ mod tests { &self, _tokens: &[u32], _block_size: u32, + _cache_namespace: Option<&str>, ) -> Result { if self.should_error { Err(KvRouterError::IndexerOffline) @@ -1362,6 +1394,7 @@ mod tests { None, false, None, + None, 0.0, 0, None, @@ -1397,6 +1430,7 @@ mod tests { None, false, None, + None, 0.0, 0, None, @@ -1422,6 +1456,7 @@ mod tests { None, false, None, + None, 0.0, 0, None, @@ -1463,6 +1498,7 @@ mod tests { false, true, None, + None, 0.0, 0, None, @@ -1486,6 +1522,7 @@ mod tests { BlockHashOptions { block_mm_infos: None, lora_name: None, + cache_namespace: None, is_eagle: Some(false), }, ); @@ -1515,6 +1552,7 @@ mod tests { false, false, None, + None, 0.0, 0, None, @@ -1549,7 +1587,7 @@ mod tests { .await; let scores = router - .get_overlap_scores(&[11, 12, 21, 22], None, None, None, true) + .get_overlap_scores(&[11, 12, 21, 22], None, None, None, None, true) .await .unwrap(); diff --git a/lib/llm/src/kv_router/prefill_router/query.rs b/lib/llm/src/kv_router/prefill_router/query.rs index 90536ffd18f4..388fb62f9399 100644 --- a/lib/llm/src/kv_router/prefill_router/query.rs +++ b/lib/llm/src/kv_router/prefill_router/query.rs @@ -21,6 +21,7 @@ impl PrefillRouter { token_ids: &[u32], block_mm_infos: Option<&[Option]>, lora_name: Option, + cache_namespace: Option, priority_jump: f64, strict_priority: u32, allowed_worker_ids: Option>, @@ -46,6 +47,7 @@ impl PrefillRouter { false, false, lora_name, + cache_namespace, priority_jump, strict_priority, None, diff --git a/lib/llm/src/kv_router/publisher/tests.rs b/lib/llm/src/kv_router/publisher/tests.rs index dfa1f0a998a7..04f666eb33c1 100644 --- a/lib/llm/src/kv_router/publisher/tests.rs +++ b/lib/llm/src/kv_router/publisher/tests.rs @@ -18,6 +18,7 @@ use std::time::Duration; mod test_event_processing { use super::*; use dynamo_kv_router::protocols::{BlockHashOptions, compute_block_hash_for_seq}; + use dynamo_kv_router::zmq_wire::StoredBlockOptions; // --------------------------------------------------------------------- // create_stored_block_from_parts -------------------------------------- @@ -32,10 +33,7 @@ mod test_event_processing { kv_block_size, blk_hash, &token_ids, - None, - None, - None, - None, + StoredBlockOptions::default(), ); assert_eq!(stored.block_hash.0, blk_hash); @@ -45,6 +43,36 @@ mod test_event_processing { assert!(stored.mm_extra_info.is_none()); } + #[test] + fn test_create_stored_block_from_parts_with_cache_salt() { + let kv_block_size = 4; + let token_ids = vec![10, 20, 30, 40]; + + let stored = create_stored_block_from_parts( + kv_block_size, + 0xdead_beef, + &token_ids, + StoredBlockOptions { + cache_namespace: Some("tenant-a"), + ..Default::default() + }, + ); + + let expected_hash = compute_block_hash_for_seq( + &token_ids, + kv_block_size, + BlockHashOptions { + cache_namespace: Some("tenant-a"), + ..Default::default() + }, + )[0]; + let base_hash = + compute_block_hash_for_seq(&token_ids, kv_block_size, BlockHashOptions::default())[0]; + + assert_eq!(stored.tokens_hash, expected_hash); + assert_ne!(stored.tokens_hash, base_hash); + } + // --------------------------------------------------------------------- // create_stored_blocks ------------------------------------------------- // --------------------------------------------------------------------- @@ -62,6 +90,7 @@ mod test_event_processing { &num_block_tokens, &block_hashes, None, + None, &Arc::new(AtomicU32::new(0)), None, None, @@ -87,6 +116,7 @@ mod test_event_processing { &num_block_tokens, &block_hashes, None, + None, &warning_count, None, None, @@ -111,6 +141,7 @@ mod test_event_processing { block_size: 4, medium: None, lora_name: None, + cache_namespace: None, block_mm_infos: None, is_eagle: None, group_idx: None, @@ -142,6 +173,7 @@ mod test_event_processing { block_size: 4, medium: None, lora_name: None, + cache_namespace: None, block_mm_infos: None, is_eagle: None, group_idx: None, @@ -155,6 +187,7 @@ mod test_event_processing { block_size: 4, medium: None, lora_name: Some("my-lora".to_string()), + cache_namespace: None, block_mm_infos: None, is_eagle: None, group_idx: None, @@ -209,6 +242,7 @@ mod test_event_processing { block_size: 4, medium: None, lora_name: None, + cache_namespace: None, block_mm_infos: None, is_eagle: None, group_idx: None, @@ -222,6 +256,7 @@ mod test_event_processing { block_size: 4, medium: None, lora_name: None, + cache_namespace: None, block_mm_infos: None, is_eagle: None, group_idx: None, @@ -1086,6 +1121,7 @@ mod tests_startup_helpers { block_size: 4, medium: None, lora_name: None, + cache_namespace: None, block_mm_infos: None, is_eagle: None, group_idx: None, diff --git a/lib/llm/src/kv_router/push_router/selection.rs b/lib/llm/src/kv_router/push_router/selection.rs index 485d7e6a1c56..059a42c691f7 100644 --- a/lib/llm/src/kv_router/push_router/selection.rs +++ b/lib/llm/src/kv_router/push_router/selection.rs @@ -59,6 +59,7 @@ struct BestMatchArgs<'a> { update_states: bool, return_routing_hashes: bool, lora_name: Option, + cache_namespace: Option, priority_jump: f64, strict_priority: u32, policy_class: Option, @@ -82,6 +83,7 @@ impl KvPushRouter { args.update_states, args.return_routing_hashes, args.lora_name, + args.cache_namespace, args.priority_jump, args.strict_priority, args.policy_class, @@ -126,6 +128,7 @@ impl KvPushRouter { let _nvtx_select = dynamo_nvtx_range!("route.select_worker"); let routing = request.routing.as_ref(); let lora_name = routing.and_then(|routing| routing.lora_name.clone()); + let cache_namespace = routing.and_then(|routing| routing.cache_namespace.clone()); let priority_jump = routing .and_then(|routing| routing.priority_jump) .unwrap_or(0.0); @@ -158,6 +161,7 @@ impl KvPushRouter { update_states: !is_query_only, return_routing_hashes, lora_name, + cache_namespace, priority_jump, strict_priority, policy_class, @@ -192,6 +196,7 @@ impl KvPushRouter { return Ok(selection); }; + let cache_namespace = routing.and_then(|routing| routing.cache_namespace.clone()); let pinned_worker = resolve_pinned_worker_rank( pinned_worker_id, @@ -229,6 +234,7 @@ impl KvPushRouter { update_states: !is_query_only, return_routing_hashes, lora_name, + cache_namespace, priority_jump, strict_priority, policy_class, diff --git a/lib/llm/src/kv_router/route_lookup.rs b/lib/llm/src/kv_router/route_lookup.rs index 582327e5d594..5a20ea1167d9 100644 --- a/lib/llm/src/kv_router/route_lookup.rs +++ b/lib/llm/src/kv_router/route_lookup.rs @@ -47,11 +47,20 @@ pub(super) async fn query_tiered_matches( tokens: &[u32], block_size: u32, block_hashes: Vec, + cache_namespace: Option<&str>, retain_block_hashes: bool, ) -> Result { if retain_block_hashes { let (tiered_matches, shared_cache_hits, indexer_duration, shared_cache_duration) = - query_retained(indexer, shared_cache, tokens, block_size, &block_hashes).await?; + query_retained( + indexer, + shared_cache, + tokens, + block_size, + &block_hashes, + cache_namespace, + ) + .await?; return Ok(TieredLookupResult { tiered_matches, @@ -62,8 +71,15 @@ pub(super) async fn query_tiered_matches( }); } - let (tiered_matches, shared_cache_hits, indexer_duration, shared_cache_duration) = - query_owned(indexer, shared_cache, tokens, block_size, block_hashes).await?; + let (tiered_matches, shared_cache_hits, indexer_duration, shared_cache_duration) = query_owned( + indexer, + shared_cache, + tokens, + block_size, + block_hashes, + cache_namespace, + ) + .await?; Ok(TieredLookupResult { tiered_matches, @@ -80,6 +96,7 @@ async fn query_retained( tokens: &[u32], block_size: u32, block_hashes: &[LocalBlockHash], + cache_namespace: Option<&str>, ) -> Result< ( TieredMatchDetails, @@ -101,7 +118,14 @@ async fn query_retained( let indexer_fut = indexer .find_matches_by_tier_ref(block_hashes) .instrument(tracing::info_span!("kv_router.find_matches")); - join_indexer_and_shared_cache(indexer_fut, shared_cache, tokens, block_size).await + join_indexer_and_shared_cache( + indexer_fut, + shared_cache, + tokens, + block_size, + cache_namespace, + ) + .await } async fn query_owned( @@ -110,6 +134,7 @@ async fn query_owned( tokens: &[u32], block_size: u32, block_hashes: Vec, + cache_namespace: Option<&str>, ) -> Result< ( TieredMatchDetails, @@ -131,7 +156,14 @@ async fn query_owned( let indexer_fut = indexer .find_matches_by_tier(block_hashes) .instrument(tracing::info_span!("kv_router.find_matches")); - join_indexer_and_shared_cache(indexer_fut, shared_cache, tokens, block_size).await + join_indexer_and_shared_cache( + indexer_fut, + shared_cache, + tokens, + block_size, + cache_namespace, + ) + .await } async fn join_indexer_and_shared_cache( @@ -139,6 +171,7 @@ async fn join_indexer_and_shared_cache( shared_cache: &dyn SharedKvCache, tokens: &[u32], block_size: u32, + cache_namespace: Option<&str>, ) -> Result< ( TieredMatchDetails, @@ -152,7 +185,7 @@ where I: Future>, { let shared_fut = shared_cache - .check_blocks(tokens, block_size) + .check_blocks(tokens, block_size, cache_namespace) .instrument(tracing::info_span!("kv_router.shared_cache_check")); let indexer_timed = async { diff --git a/lib/llm/src/kv_router/shared_cache.rs b/lib/llm/src/kv_router/shared_cache.rs index 3b2e87cf9fac..482605c5cfa8 100644 --- a/lib/llm/src/kv_router/shared_cache.rs +++ b/lib/llm/src/kv_router/shared_cache.rs @@ -167,7 +167,16 @@ impl SharedKvCache for HicacheSharedKvCache { &self, tokens: &[u32], block_size: u32, + cache_namespace: Option<&str>, ) -> Result { + if cache_namespace + .filter(|namespace| !namespace.is_empty()) + .is_some() + { + tracing::debug!("Skipping SGLang Mooncake HiCache lookup for cache-namespaced request"); + return Ok(SharedCacheHits::default()); + } + let Some(config) = self.resolve_mooncake_config() else { tracing::debug!("No SGLang Mooncake HiCache runtime config available"); return Ok(SharedCacheHits::default()); @@ -559,7 +568,7 @@ mod tests { let cache = HicacheSharedKvCache::new(runtime_watch_with_config(config)); let hits = cache - .check_blocks(&[1, 2, 3, 4, 5, 6, 7, 8], 4) + .check_blocks(&[1, 2, 3, 4, 5, 6, 7, 8], 4, None) .await .unwrap(); @@ -568,4 +577,25 @@ mod tests { mock.assert_async().await; } + + #[tokio::test] + async fn test_check_blocks_skips_mooncake_for_cache_namespace() { + let server = Server::new_async().await; + let server_url = Url::parse(&server.url()).unwrap(); + + let config = SglangHicacheMooncakeConfig { + master_server_address: Some(format!("{}:50051", server_url.host_str().unwrap())), + master_metrics_port: server_url.port().unwrap(), + ..mooncake_config() + }; + + let cache = HicacheSharedKvCache::new(runtime_watch_with_config(config)); + let hits = cache + .check_blocks(&[1, 2, 3, 4, 5, 6, 7, 8], 4, Some("tenant-a")) + .await + .unwrap(); + + assert!(hits.ranges.is_empty()); + assert_eq!(hits.total_hits, 0); + } } diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index 2be53d16c19a..3313353ef5ae 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -856,6 +856,7 @@ impl OpenAIPreprocessor { strict_priority, priority, lora_name, + cache_namespace: nvext.cache_salt.clone(), allowed_worker_ids: None, routing_constraints: nvext .routing_constraints diff --git a/lib/llm/src/protocols/common/extensions.rs b/lib/llm/src/protocols/common/extensions.rs index 12f7c6464e1c..c003c06b8e1f 100644 --- a/lib/llm/src/protocols/common/extensions.rs +++ b/lib/llm/src/protocols/common/extensions.rs @@ -242,6 +242,7 @@ pub const HEADER_DP_RANK: &str = "x-dynamo-dp-rank"; pub const HEADER_PREFILL_DP_RANK: &str = "x-dynamo-prefill-dp-rank"; pub const HEADER_REQUEST_PRIORITY: &str = "x-dynamo-request-priority"; pub const HEADER_REQUEST_STRICT_PRIORITY: &str = "x-dynamo-request-strict-priority"; +pub const HEADER_TENANT_ID: &str = "x-tenant-id"; // Compatibility aliases for the original unprefixed names. Future agents may remove these after // the deprecation window. pub const HEADER_WORKER_INSTANCE_ID_ALIAS: &str = "x-worker-instance-id"; @@ -299,6 +300,7 @@ pub fn session_affinity_from_headers(headers: &HeaderMap) -> Option `prefill_dp_rank` /// - `x-dynamo-request-priority` -> `agent_hints.priority` /// - `x-dynamo-request-strict-priority` -> `agent_hints.strict_priority` +/// - `x-tenant-id` -> `cache_salt` /// /// Routing headers take priority over existing nvext values when present. /// If no headers are present, returns the original nvext unchanged. @@ -338,6 +340,11 @@ pub fn apply_header_routing_overrides(nvext: Option, headers: &HeaderMap) .get(HEADER_REQUEST_STRICT_PRIORITY) .and_then(|v| v.to_str().ok()) .and_then(|s| s.parse::().ok()); + let tenant_id = headers + .get(HEADER_TENANT_ID) + .and_then(|v| v.to_str().ok()) + .filter(|s| !s.is_empty()) + .map(str::to_owned); if worker_id.is_none() && prefill_id.is_none() @@ -345,6 +352,7 @@ pub fn apply_header_routing_overrides(nvext: Option, headers: &HeaderMap) && prefill_dp_rank.is_none() && priority.is_none() && strict_priority.is_none() + && tenant_id.is_none() { return nvext; } @@ -372,6 +380,9 @@ pub fn apply_header_routing_overrides(nvext: Option, headers: &HeaderMap) hints.strict_priority = Some(strict_priority); } } + if let Some(salt) = tenant_id { + ext.cache_salt = Some(salt); + } Some(ext) } @@ -806,6 +817,25 @@ mod tests { assert_eq!(hints.osl, Some(99)); } + #[test] + fn apply_header_routing_overrides_sets_cache_salt_from_tenant_header() { + let mut headers = HeaderMap::new(); + headers.insert(HEADER_TENANT_ID, "tenant-a".parse().unwrap()); + + let nvext = apply_header_routing_overrides(None, &headers).unwrap(); + assert_eq!(nvext.cache_salt.as_deref(), Some("tenant-a")); + + let mut headers = HeaderMap::new(); + headers.insert(HEADER_TENANT_ID, "tenant-header".parse().unwrap()); + let nvext = NvExt { + cache_salt: Some("tenant-body".to_string()), + ..Default::default() + }; + + let nvext = apply_header_routing_overrides(Some(nvext), &headers).unwrap(); + assert_eq!(nvext.cache_salt.as_deref(), Some("tenant-header")); + } + #[test] fn apply_header_routing_overrides_supports_unprefixed_aliases() { let mut headers = HeaderMap::new(); diff --git a/lib/llm/src/protocols/common/preprocessor.rs b/lib/llm/src/protocols/common/preprocessor.rs index 0afad99a50b6..61eee68aac3f 100644 --- a/lib/llm/src/protocols/common/preprocessor.rs +++ b/lib/llm/src/protocols/common/preprocessor.rs @@ -54,6 +54,14 @@ pub struct RoutingHints { #[serde(default, skip_serializing_if = "Option::is_none")] pub lora_name: Option, + /// Cache namespace for request-scoped KV cache isolation. + #[serde( + default, + rename = "cache_salt", + skip_serializing_if = "Option::is_none" + )] + pub cache_namespace: Option, + /// Priority jump in seconds for queue ordering. /// A positive value decreases the effective arrival time, moving the request /// ahead in the scheduler queue. @@ -526,4 +534,20 @@ mod tests { "encoder_result must be absent from wire when None; got {json}" ); } + + #[test] + fn routing_hints_cache_namespace_serializes_as_cache_salt() { + let hints = RoutingHints { + cache_namespace: Some("tenant-a".to_string()), + ..Default::default() + }; + + let value = serde_json::to_value(&hints).unwrap(); + + assert_eq!(value["cache_salt"], "tenant-a"); + assert!(value.get("cache_namespace").is_none()); + + let decoded: RoutingHints = serde_json::from_value(value).unwrap(); + assert_eq!(decoded.cache_namespace.as_deref(), Some("tenant-a")); + } } diff --git a/lib/mocker/src/replay/online/router.rs b/lib/mocker/src/replay/online/router.rs index 836506836d0d..eefc6be25ac2 100644 --- a/lib/mocker/src/replay/online/router.rs +++ b/lib/mocker/src/replay/online/router.rs @@ -53,11 +53,11 @@ impl ReplayIndexer { ) -> Result { match self { Self::Single(indexer) => indexer - .find_matches_for_request(tokens, lora_name, None) + .find_matches_for_request(tokens, lora_name, None, None) .await .map_err(Into::into), Self::Concurrent(indexer) => indexer - .find_matches_for_request(tokens, lora_name, None) + .find_matches_for_request(tokens, lora_name, None, None) .await .map_err(Into::into), } From 029a596dbf1a6736655e3e7acb2abcb8bd3c5cc0 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Mon, 29 Jun 2026 11:14:39 -0700 Subject: [PATCH 02/10] test(trtllm): cover cache salt routing isolation Signed-off-by: jthomson04 --- tests/router/common.py | 109 +++++++++++++++++++ tests/router/e2e_harness.py | 31 ++++++ tests/router/test_router_e2e_with_unified.py | 37 ++++++- 3 files changed, 176 insertions(+), 1 deletion(-) diff --git a/tests/router/common.py b/tests/router/common.py index c123d41abb09..5b4834280d9c 100644 --- a/tests/router/common.py +++ b/tests/router/common.py @@ -3151,6 +3151,115 @@ async def _verify_selection_service_scores(): asyncio.run(_verify_selection_service_scores()) +def _test_router_cache_salt_isolation( + engine_workers, + endpoint, + model_name: str, + block_size: int, +): + """Verify cache-salted TRT-LLM events remain isolated in the router index.""" + + async def test_sync(): + expected_num_instances = engine_workers.num_workers + kv_router = _create_kv_router_with_timeout( + router_factory=lambda: KvRouter( + endpoint=endpoint, + block_size=block_size, + kv_router_config=KvRouterConfig( + router_snapshot_threshold=20, + use_kv_events=True, + router_event_threads=4, + ), + ), + num_workers=expected_num_instances, + engine_workers=engine_workers, + ) + + worker_ids = await wait_for_workers_ready( + endpoint, + kv_router, + expected_num_workers=expected_num_instances, + model_name=model_name, + ) + assert len(worker_ids) >= 2, "cache-salt isolation requires two workers" + + worker_a = (worker_ids[0], 0) + worker_b = (worker_ids[1], 0) + token_ids = [random.randint(1, 10_000) for _ in range(block_size * 2)] + expected_blocks = len(token_ids) // block_size + + async def generate(cache_salt: str, worker_id: int) -> None: + request = { + "model": model_name, + "token_ids": token_ids, + "stop_conditions": {"ignore_eos": True, "max_tokens": 2}, + "sampling_options": {}, + "output_options": {}, + "eos_token_ids": [], + "extra_args": {"nvext": {"cache_salt": cache_salt}}, + "routing": { + "backend_instance_id": worker_id, + "cache_salt": cache_salt, + }, + } + stream = await kv_router.generate_from_request(request) + terminal = None + async for response in stream: + if ( + isinstance(response, dict) + and response.get("finish_reason") is not None + ): + terminal = response + assert terminal is not None, f"tenant {cache_salt} request did not finish" + + async def device_blocks( + cache_salt: Optional[str], + ) -> dict[tuple[int, int], int]: + scores = await kv_router.get_overlap_scores( + token_ids, + include_shared=False, + cache_namespace=cache_salt, + ) + assert scores["block_size"] == block_size + assert scores["num_blocks"] == expected_blocks + assert scores["shared_cache"]["enabled"] is False + return { + (row["worker_id"], row["dp_rank"]): row["device_blocks"] + for row in scores["workers"] + } + + async def wait_for_scores( + cache_salt: Optional[str], + expected: dict[tuple[int, int], int], + ) -> None: + deadline = time.monotonic() + 10 + last_scores: dict[tuple[int, int], int] = {} + expected_nonzero = {key: value for key, value in expected.items() if value} + while time.monotonic() < deadline: + last_scores = await device_blocks(cache_salt) + actual_nonzero = { + key: value for key, value in last_scores.items() if value + } + if actual_nonzero == expected_nonzero: + return + await asyncio.sleep(0.25) + raise AssertionError( + f"cache_salt={cache_salt!r}: expected {expected}, got {last_scores}" + ) + + await generate("tenant-a", worker_a[0]) + await wait_for_scores("tenant-a", {worker_a: expected_blocks}) + await wait_for_scores("tenant-b", {}) + await wait_for_scores(None, {}) + + await generate("tenant-b", worker_b[0]) + await wait_for_scores("tenant-b", {worker_b: expected_blocks}) + await wait_for_scores("tenant-a", {worker_a: expected_blocks}) + await wait_for_scores(None, {}) + + asyncio.run(test_sync()) + + def _test_busy_threshold_endpoint( engine_workers, block_size: int, diff --git a/tests/router/e2e_harness.py b/tests/router/e2e_harness.py index 97ed6566a8f9..d00cc48bdfa9 100644 --- a/tests/router/e2e_harness.py +++ b/tests/router/e2e_harness.py @@ -8,6 +8,7 @@ from tests.router.common import ( _test_router_basic, + _test_router_cache_salt_isolation, _test_router_decisions, _test_router_decisions_disagg, _test_router_indexers_sync, @@ -285,6 +286,36 @@ def run_router_decisions_test( ) +def run_cache_salt_isolation_test( + *, + engine_process_cls, + engine_args_name: str, + engine_args: dict[str, Any], + request, + request_plane: str, + model_name: str, + block_size: int, + component_name: str, +): + process = _create_engine_process( + engine_process_cls=engine_process_cls, + engine_args_name=engine_args_name, + engine_args=engine_args, + request=request, + request_plane=request_plane, + default_process_kwargs={"num_workers": 2, "single_gpu": True}, + engine_process_kwargs=None, + ) + with process as engine_workers: + endpoint = get_engine_endpoint(engine_workers, request_plane, component_name) + _test_router_cache_salt_isolation( + engine_workers, + endpoint, + model_name, + block_size, + ) + + def run_disagg_router_decisions_test( *, engine_process_cls, diff --git a/tests/router/test_router_e2e_with_unified.py b/tests/router/test_router_e2e_with_unified.py index f0cd6f029ae6..9e2fae719b67 100644 --- a/tests/router/test_router_e2e_with_unified.py +++ b/tests/router/test_router_e2e_with_unified.py @@ -25,7 +25,11 @@ import pytest -from tests.router.e2e_harness import run_basic_router_test, run_router_decisions_test +from tests.router.e2e_harness import ( + run_basic_router_test, + run_cache_salt_isolation_test, + run_router_decisions_test, +) from tests.router.test_router_e2e_with_sglang import MODEL_NAME as SGLANG_MODEL_NAME from tests.router.test_router_e2e_with_sglang import SGLANG_ARGS, SGLangProcess from tests.router.test_router_e2e_with_trtllm import MODEL_NAME as TRTLLM_MODEL_NAME @@ -421,6 +425,37 @@ def test_unified_trtllm_router_decisions_multiple_workers( ) +@pytest.mark.pre_merge +@pytest.mark.gpu_1 +@pytest.mark.trtllm +@pytest.mark.model(TRTLLM_MODEL_NAME) +@pytest.mark.profiled_vram_gib(7.8) +@pytest.mark.requested_trtllm_kv_tokens(2592) +@pytest.mark.timeout(600) +@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True) +def test_unified_trtllm_cache_salt_isolation( + request, + runtime_services_dynamic_ports, + predownload_models, + request_plane, +) -> None: + """Identical prompts under different cache salts never share KV entries. + + This crosses the real unified TRT-LLM engine and its KV-event publisher, + then queries the router index independently for each tenant namespace. + """ + run_cache_salt_isolation_test( + engine_process_cls=UnifiedTRTLLMProcess, + engine_args_name="trtllm_args", + engine_args=TRTLLM_ARGS, + request=request, + request_plane=request_plane, + model_name=TRTLLM_MODEL_NAME, + block_size=TRTLLM_BLOCK_SIZE, + component_name="backend", + ) + + @pytest.mark.gpu_2 @pytest.mark.nightly @pytest.mark.trtllm From 002b7911e92dfffad13806472512d25b43d31297 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Mon, 29 Jun 2026 11:50:20 -0700 Subject: [PATCH 03/10] fix(router): address cache namespace review gaps Signed-off-by: jthomson04 --- .github/workflows/pr.yaml | 2 + components/src/dynamo/trtllm/llm_engine.py | 22 +-- components/src/dynamo/trtllm/publisher.py | 4 +- .../trtllm/request_handlers/handler_base.py | 22 +-- .../src/dynamo/trtllm/utils/request_utils.py | 23 ++++ .../epp/pkg/plugins/disagg/decode_scorer.go | 18 ++- .../pkg/plugins/dynamo_kv_scorer/plugin.go | 71 ++++++++-- .../plugins/dynamo_kv_scorer/plugin_test.go | 26 ++++ deploy/inference-gateway/ext-proc/src/epp.rs | 21 +++ lib/bindings/c/src/lib.rs | 129 +++++++++++++++++- lib/kv-router/src/indexer/tests.rs | 32 +++++ .../src/services/selection/core/mod.rs | 3 + lib/kv-router/src/services/selection/input.rs | 7 +- lib/kv-router/src/services/selection/tests.rs | 25 ++++ lib/kv-router/src/zmq_wire/extra_keys.rs | 22 ++- lib/kv-router/src/zmq_wire/mod.rs | 66 +++++++-- lib/kv-router/src/zmq_wire/tests.rs | 38 ++++++ lib/llm/src/kv_router/publisher/tests.rs | 27 ++++ lib/llm/src/preprocessor.rs | 59 ++++++-- 19 files changed, 525 insertions(+), 92 deletions(-) create mode 100644 components/src/dynamo/trtllm/utils/request_utils.py diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 0692bd16ace1..89fba5de9d82 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -131,6 +131,8 @@ jobs: operator: needs: changed-files + # Keep this trigger set a superset of every deploy-operator* job guarded by + # needs.operator.result == 'success'. if: | needs.changed-files.outputs.operator == 'true' || needs.changed-files.outputs.vllm == 'true' || diff --git a/components/src/dynamo/trtllm/llm_engine.py b/components/src/dynamo/trtllm/llm_engine.py index 971a9c07cfce..dc77b313d22e 100644 --- a/components/src/dynamo/trtllm/llm_engine.py +++ b/components/src/dynamo/trtllm/llm_engine.py @@ -18,7 +18,7 @@ import sys import threading import time -from collections.abc import AsyncGenerator, Callable, Mapping +from collections.abc import AsyncGenerator, Callable from dataclasses import asdict from typing import TYPE_CHECKING, Any, Optional @@ -68,6 +68,7 @@ DisaggregatedParams, DisaggregatedParamsCodec, ) +from dynamo.trtllm.utils.request_utils import request_cache_salt from dynamo.trtllm.utils.trtllm_utils import deep_update, warn_override_collisions if TYPE_CHECKING: @@ -107,23 +108,6 @@ } -def _request_cache_salt(request: Mapping[str, Any]) -> Optional[str]: - routing = request.get("routing") or {} - if isinstance(routing, dict): - cache_salt = routing.get("cache_salt") - if cache_salt is not None: - return cache_salt - - extra_args = request.get("extra_args") or {} - nvext = extra_args.get("nvext") if isinstance(extra_args, dict) else None - if isinstance(nvext, dict): - cache_salt = nvext.get("cache_salt") - if cache_salt is not None: - return cache_salt - - return None - - def _to_signed_i64(value: int | None) -> int | None: """Two's-complement cast of a Python int into the signed 64-bit range.""" if value is None: @@ -862,7 +846,7 @@ 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) + cache_salt = request_cache_salt(request) generation_result = self._engine.llm.generate_async( inputs=token_ids, sampling_params=sampling_params, diff --git a/components/src/dynamo/trtllm/publisher.py b/components/src/dynamo/trtllm/publisher.py index a95fee016a3a..96ca03b4876b 100644 --- a/components/src/dynamo/trtllm/publisher.py +++ b/components/src/dynamo/trtllm/publisher.py @@ -875,14 +875,14 @@ def _handle_kv_event(self, event): logger.debug( "Publishing stored KV event: engine_event_id=%s " - "attention_dp_rank=%s blocks=%s tokens=%s lora_name=%s cache_salt=%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, + cache_salt is not None, parent_hash is not None, ) # Publish to ZMQ if consolidator is enabled, otherwise publish to NATS diff --git a/components/src/dynamo/trtllm/request_handlers/handler_base.py b/components/src/dynamo/trtllm/request_handlers/handler_base.py index e3343b16ad20..564f72a0a02d 100644 --- a/components/src/dynamo/trtllm/request_handlers/handler_base.py +++ b/components/src/dynamo/trtllm/request_handlers/handler_base.py @@ -18,7 +18,7 @@ import logging import os import re -from collections.abc import AsyncGenerator, Mapping +from collections.abc import AsyncGenerator from contextlib import asynccontextmanager from dataclasses import asdict, dataclass from typing import TYPE_CHECKING, Any, Optional, Protocol, Union @@ -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, @@ -64,23 +65,6 @@ logger = logging.getLogger(__name__) -def _request_cache_salt(request: Mapping[str, Any]) -> Optional[str]: - routing = request.get("routing") or {} - if isinstance(routing, dict): - cache_salt = routing.get("cache_salt") - if cache_salt is not None: - return cache_salt - - extra_args = request.get("extra_args") or {} - nvext = extra_args.get("nvext") if isinstance(extra_args, dict) else None - if isinstance(nvext, dict): - cache_salt = nvext.get("cache_salt") - if cache_salt is not None: - return cache_salt - - return None - - class TRTLLMEnginePauseController: """Adapts TRT-LLM sleep/wake to the standard pause controller interface. @@ -1127,7 +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) + cache_salt = request_cache_salt(request) try: # NEW: Updated engine call to include multimodal data diff --git a/components/src/dynamo/trtllm/utils/request_utils.py b/components/src/dynamo/trtllm/utils/request_utils.py new file mode 100644 index 000000000000..cbef9333ece8 --- /dev/null +++ b/components/src/dynamo/trtllm/utils/request_utils.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Mapping +from typing import Any, Optional + + +def request_cache_salt(request: Mapping[str, Any]) -> Optional[str]: + """Return cache_salt using routing hints before legacy extra_args.""" + routing = request.get("routing") or {} + if isinstance(routing, dict): + cache_salt = routing.get("cache_salt") + if cache_salt is not None: + return cache_salt + + extra_args = request.get("extra_args") or {} + nvext = extra_args.get("nvext") if isinstance(extra_args, dict) else None + if isinstance(nvext, dict): + cache_salt = nvext.get("cache_salt") + if cache_salt is not None: + return cache_salt + + return None diff --git a/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go index d867a48da9fe..ca89eadd65b0 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go @@ -58,6 +58,7 @@ type DecodeRoutingState struct { DpRank uint32 PrefillWorkerID string TokenData []int64 + CacheNamespace string } // Clone implements plugins.StateData. @@ -69,6 +70,7 @@ func (s *DecodeRoutingState) Clone() plugins.StateData { WorkerID: s.WorkerID, DpRank: s.DpRank, PrefillWorkerID: s.PrefillWorkerID, + CacheNamespace: s.CacheNamespace, } if s.TokenData != nil { clone.TokenData = make([]int64, len(s.TokenData)) @@ -181,9 +183,10 @@ func (s *DynDecodeScorer) Score(ctx context.Context, cycleState *schedtypes.Cycl // Store routing state for PreRequest bookkeeping if req.RequestId != "" { routingState := &DecodeRoutingState{ - WorkerID: workerIDStr, - DpRank: result.DpRank, - TokenData: result.TokenData, + WorkerID: workerIDStr, + DpRank: result.DpRank, + TokenData: result.TokenData, + CacheNamespace: result.CacheNamespace, } s.pluginState.Write(req.RequestId, plugins.StateKey(decodeStateKey), routingState) } @@ -222,7 +225,13 @@ func (s *DynDecodeScorer) PreRequest(ctx context.Context, request *schedtypes.In return } - if addErr := dynscorer.CallAddRequest(request.RequestId, state.TokenData, workerIDUint, state.DpRank); addErr != nil { + if addErr := dynscorer.CallAddRequest( + request.RequestId, + state.TokenData, + workerIDUint, + state.DpRank, + state.CacheNamespace, + ); addErr != nil { logger.V(logutil.DEFAULT).Error(addErr, "DynDecodeScorer PreRequest: failed to add request", "requestID", request.RequestId) return @@ -232,6 +241,7 @@ func (s *DynDecodeScorer) PreRequest(ctx context.Context, request *schedtypes.In "requestID", request.RequestId, "workerID", state.WorkerID, "dpRank", state.DpRank, + "cacheNamespace", state.CacheNamespace, "tokenCount", len(state.TokenData)) } diff --git a/deploy/inference-gateway/epp/pkg/plugins/dynamo_kv_scorer/plugin.go b/deploy/inference-gateway/epp/pkg/plugins/dynamo_kv_scorer/plugin.go index 93ef628964e8..f3320f2bd14d 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/dynamo_kv_scorer/plugin.go +++ b/deploy/inference-gateway/epp/pkg/plugins/dynamo_kv_scorer/plugin.go @@ -56,6 +56,8 @@ typedef struct { uint32_t decode_dp_rank; uint32_t *token_ids; size_t token_count; + uint8_t *cache_namespace; + size_t cache_namespace_len; } CRoutingResult; // Router bindings API @@ -75,12 +77,14 @@ query_router_result_t route_decode_request(RouterHandles *handle, bool is_disaggregated, CRoutingResult *out_result); -query_router_result_t add_request(RouterHandles *handle, - const char *request_id, - const uint32_t *token_ids, - size_t token_count, - uint64_t worker_id, - uint32_t dp_rank); +query_router_result_t add_request_with_cache_namespace(RouterHandles *handle, + const char *request_id, + const uint32_t *token_ids, + size_t token_count, + uint64_t worker_id, + uint32_t dp_rank, + const uint8_t *cache_namespace, + size_t cache_namespace_len); query_router_result_t mark_prefill_complete(RouterHandles *handle, const char *request_id); @@ -319,6 +323,9 @@ func BuildOpenAIRequest(req *schedtypes.InferenceRequest) (map[string]any, error if nvext := extractNvext(req.Body.Payload); nvext != nil { requestBody["nvext"] = nvext } + if cacheSalt := extractTopLevelCacheSalt(req.Body.Payload); cacheSalt != "" { + requestBody["cache_salt"] = cacheSalt + } return requestBody, nil } @@ -354,8 +361,17 @@ func extractNvext(payload fwkrh.RequestPayload) map[string]any { return nvext } +func extractTopLevelCacheSalt(payload fwkrh.RequestPayload) string { + pm, ok := payload.(fwkrh.PayloadMap) + if !ok { + return "" + } + cacheSalt, _ := pm["cache_salt"].(string) + return cacheSalt +} + // CallAddRequest registers a request with the router's bookkeeping. -func CallAddRequest(requestID string, tokenData []int64, workerID uint64, dpRank uint32) error { +func CallAddRequest(requestID string, tokenData []int64, workerID uint64, dpRank uint32, cacheNamespace string) error { if !routerInitialized { return fmt.Errorf("dynamo router not initialized") } @@ -375,19 +391,26 @@ func CallAddRequest(requestID string, tokenData []int64, workerID uint64, dpRank cRequestID := C.CString(requestID) defer C.free(unsafe.Pointer(cRequestID)) + var cCacheNamespace *C.uint8_t + if cacheNamespace != "" { + cCacheNamespace = (*C.uint8_t)(C.CBytes([]byte(cacheNamespace))) + defer C.free(unsafe.Pointer(cCacheNamespace)) + } var cTokens *C.uint32_t if len(tokens) > 0 { cTokens = (*C.uint32_t)(unsafe.Pointer(&tokens[0])) } - rc := C.add_request( + rc := C.add_request_with_cache_namespace( router, cRequestID, cTokens, C.size_t(len(tokens)), C.uint64_t(workerID), C.uint32_t(dpRank), + cCacheNamespace, + C.size_t(len(cacheNamespace)), ) if rc != C.QUERY_ROUTER_OK { @@ -446,9 +469,10 @@ func CallFreeRequest(requestID string) error { // RoutingResult holds the result of a prefill or decode routing call. type RoutingResult struct { - WorkerID uint64 - DpRank uint32 - TokenData []int64 + WorkerID uint64 + DpRank uint32 + TokenData []int64 + CacheNamespace string } // extractTokenData copies token IDs from a C result into Go memory. @@ -465,6 +489,15 @@ func extractTokenData(result *C.CRoutingResult) []int64 { return nil } +// extractCacheNamespace copies the namespace bytes from a C result into Go memory. +func extractCacheNamespace(result *C.CRoutingResult) string { + count := int(result.cache_namespace_len) + if count > 0 && result.cache_namespace != nil { + return string(unsafe.Slice((*byte)(unsafe.Pointer(result.cache_namespace)), count)) + } + return "" +} + // CallRoutePrefillRequest routes a request to the best prefill worker. // It tokenizes the request and queries only the prefill router. func CallRoutePrefillRequest(requestJSON string, podsJSON string) (*RoutingResult, error) { @@ -495,11 +528,17 @@ func CallRoutePrefillRequest(requestJSON string, podsJSON string) (*RoutingResul } tokens := extractTokenData(&result) + cacheNamespace := extractCacheNamespace(&result) workerID := uint64(result.prefill_worker_id) dpRank := uint32(result.prefill_dp_rank) C.free_routing_result(&result) - return &RoutingResult{WorkerID: workerID, DpRank: dpRank, TokenData: tokens}, nil + return &RoutingResult{ + WorkerID: workerID, + DpRank: dpRank, + TokenData: tokens, + CacheNamespace: cacheNamespace, + }, nil } // CallRouteDecodeRequest routes a request to the best decode worker. @@ -532,9 +571,15 @@ func CallRouteDecodeRequest(requestJSON string, podsJSON string, isDisaggregated } tokens := extractTokenData(&result) + cacheNamespace := extractCacheNamespace(&result) workerID := uint64(result.decode_worker_id) dpRank := uint32(result.decode_dp_rank) C.free_routing_result(&result) - return &RoutingResult{WorkerID: workerID, DpRank: dpRank, TokenData: tokens}, nil + return &RoutingResult{ + WorkerID: workerID, + DpRank: dpRank, + TokenData: tokens, + CacheNamespace: cacheNamespace, + }, nil } diff --git a/deploy/inference-gateway/epp/pkg/plugins/dynamo_kv_scorer/plugin_test.go b/deploy/inference-gateway/epp/pkg/plugins/dynamo_kv_scorer/plugin_test.go index f9c1f2883824..887d0a8bb982 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/dynamo_kv_scorer/plugin_test.go +++ b/deploy/inference-gateway/epp/pkg/plugins/dynamo_kv_scorer/plugin_test.go @@ -64,6 +64,32 @@ func TestBuildOpenAIRequest_ForwardsAgentHintsPriority(t *testing.T) { } } +func TestBuildOpenAIRequest_ForwardsLegacyTopLevelCacheSalt(t *testing.T) { + req := &schedtypes.InferenceRequest{ + TargetModel: "test-model", + Body: &fwkrh.InferenceRequestBody{ + ChatCompletions: &fwkrh.ChatCompletionsRequest{ + Messages: []fwkrh.Message{ + {Role: "user", Content: fwkrh.Content{Raw: "hi"}}, + }, + }, + Payload: fwkrh.PayloadMap{ + "messages": []any{map[string]any{"role": "user", "content": "hi"}}, + "model": "test-model", + "cache_salt": "tenant-legacy", + }, + }, + } + + body, err := BuildOpenAIRequest(req) + if err != nil { + t.Fatalf("BuildOpenAIRequest returned error: %v", err) + } + if got := body["cache_salt"]; got != "tenant-legacy" { + t.Fatalf("expected legacy cache_salt forwarded to FFI body, got %v", got) + } +} + func TestBuildOpenAIRequest_CompletionsTokenPromptUsesPromptIDs(t *testing.T) { req := &schedtypes.InferenceRequest{ TargetModel: "test-model", diff --git a/deploy/inference-gateway/ext-proc/src/epp.rs b/deploy/inference-gateway/ext-proc/src/epp.rs index 32bea1ce05a3..6ee8789ba621 100644 --- a/deploy/inference-gateway/ext-proc/src/epp.rs +++ b/deploy/inference-gateway/ext-proc/src/epp.rs @@ -509,6 +509,13 @@ fn extract_cache_namespace( .nvext .as_ref() .and_then(|nvext| nvext.cache_salt.clone()) + .or_else(|| { + request + .unsupported_fields + .get("cache_salt") + .and_then(|value| value.as_str()) + .map(str::to_owned) + }) } struct DiscoveredModelBootstrap { @@ -1183,5 +1190,19 @@ mod tests { ) .unwrap(); assert_eq!(extract_cache_namespace(&without_nvext), None); + + let legacy_top_level: dynamo_llm::types::openai::chat_completions::NvCreateChatCompletionRequest = + serde_json::from_str( + r#"{ + "model": "test", + "messages": [{"role": "user", "content": "hi"}], + "cache_salt": "tenant-legacy" + }"#, + ) + .unwrap(); + assert_eq!( + extract_cache_namespace(&legacy_top_level).as_deref(), + Some("tenant-legacy") + ); } } diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index e02219330931..e69a6fff314d 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -420,6 +420,10 @@ pub struct CRoutingResult { pub token_ids: *mut u32, /// Number of tokens in the request pub token_count: usize, + /// UTF-8 cache namespace bytes (needed for add_request callback) + pub cache_namespace: *mut u8, + /// Number of bytes in the cache namespace + pub cache_namespace_len: usize, } impl Default for CRoutingResult { @@ -432,6 +436,8 @@ impl Default for CRoutingResult { decode_dp_rank: 0, token_ids: ptr::null_mut(), token_count: 0, + cache_namespace: ptr::null_mut(), + cache_namespace_len: 0, } } } @@ -622,6 +628,13 @@ fn extract_cache_namespace(request: &R) -> Option { request .nvext() .and_then(|nvext| nvext.cache_salt.clone()) + .or_else(|| { + request + .unsupported_fields + .get("cache_salt") + .and_then(|value| value.as_str()) + .map(str::to_owned) + }) } /// Opaque handle for the router pair @@ -892,6 +905,42 @@ pub unsafe extern "C" fn add_request( token_count: usize, worker_id: u64, dp_rank: u32, +) -> QueryRouterResult { + unsafe { + add_request_with_cache_namespace( + handle, + request_id, + token_ids, + token_count, + worker_id, + dp_rank, + ptr::null(), + 0, + ) + } +} + +/// Add a cache-namespaced request to the router's bookkeeping after worker selection. +/// +/// This preserves the original `add_request` ABI for unsalted callers while allowing callers +/// that routed with a namespace to use the same hashes for scheduler bookkeeping. +/// +/// # Safety +/// - `handle` must be a valid RouterHandles handle +/// - `request_id` must be a valid null-terminated C string +/// - `token_ids` must point to at least `token_count` valid u32 values +/// - `cache_namespace` must point to at least `cache_namespace_len` valid UTF-8 bytes when the +/// length is non-zero +#[unsafe(no_mangle)] +pub unsafe extern "C" fn add_request_with_cache_namespace( + handle: RouterHandlesPtr, + request_id: *const c_char, + token_ids: *const u32, + token_count: usize, + worker_id: u64, + dp_rank: u32, + cache_namespace: *const u8, + cache_namespace_len: usize, ) -> QueryRouterResult { if handle.is_null() || request_id.is_null() { return QueryRouterResult::ErrInvalidParam; @@ -902,6 +951,19 @@ pub unsafe extern "C" fn add_request( Ok(s) => s.to_owned(), Err(_) => return QueryRouterResult::ErrInvalidParam, }; + let cache_namespace = if cache_namespace_len == 0 { + None + } else if cache_namespace.is_null() { + return QueryRouterResult::ErrInvalidParam; + } else { + match std::str::from_utf8(unsafe { + std::slice::from_raw_parts(cache_namespace, cache_namespace_len) + }) { + Ok("") => None, + Ok(namespace) => Some(namespace.to_owned()), + Err(_) => return QueryRouterResult::ErrInvalidParam, + } + }; let tokens: Vec = if token_count > 0 && !token_ids.is_null() { unsafe { std::slice::from_raw_parts(token_ids, token_count) }.to_vec() @@ -925,7 +987,7 @@ pub unsafe extern "C" fn add_request( // Compute overlap_blocks using the public method let overlap_blocks = match decode_router - .get_overlap_blocks(&tokens, None, worker, None, None) + .get_overlap_blocks(&tokens, None, worker, None, cache_namespace.as_deref()) .await { Ok(overlap) => overlap, @@ -945,7 +1007,7 @@ pub unsafe extern "C" fn add_request( None, worker, None, // lora_name - None, // cache_namespace + cache_namespace.clone(), Some(&router_config_override), ) .await; @@ -954,6 +1016,7 @@ pub unsafe extern "C" fn add_request( request_id = %request_id_str, worker_id = worker_id, dp_rank = dp_rank, + cache_namespace = cache_namespace.as_deref(), overlap_blocks = overlap_blocks, token_count = tokens.len(), "add_request completed" @@ -1124,6 +1187,18 @@ pub unsafe extern "C" fn free_routing_result(result: *mut CRoutingResult) { res.token_ids = ptr::null_mut(); res.token_count = 0; } + + // Free cache namespace bytes + if !res.cache_namespace.is_null() && res.cache_namespace_len > 0 { + drop(unsafe { + Box::from_raw(std::ptr::slice_from_raw_parts_mut( + res.cache_namespace, + res.cache_namespace_len, + )) + }); + res.cache_namespace = ptr::null_mut(); + res.cache_namespace_len = 0; + } } /// Parse a JSON request string, collect completion prompts directly or apply @@ -1305,6 +1380,17 @@ fn write_tokens_to_result(tokens: &[u32], out: &mut CRoutingResult) { std::mem::forget(tokens_boxed); } +/// Write cache namespace bytes into a `CRoutingResult`, transferring ownership to the caller. +fn write_cache_namespace_to_result(cache_namespace: Option<&str>, out: &mut CRoutingResult) { + let Some(cache_namespace) = cache_namespace.filter(|namespace| !namespace.is_empty()) else { + return; + }; + let mut namespace_boxed = cache_namespace.as_bytes().to_vec().into_boxed_slice(); + out.cache_namespace = namespace_boxed.as_mut_ptr(); + out.cache_namespace_len = namespace_boxed.len(); + std::mem::forget(namespace_boxed); +} + /// Route a request to select the best **prefill** worker only. /// /// This is used in disaggregated mode where the EPP runs separate prefill and decode @@ -1348,7 +1434,7 @@ pub unsafe extern "C" fn route_prefill_request( &tokens, None, None, - cache_namespace, + cache_namespace.clone(), priority_jump, strict_priority, allowed_worker_ids, @@ -1378,6 +1464,7 @@ pub unsafe extern "C" fn route_prefill_request( out.prefill_worker_id = prefill_worker_id; out.prefill_dp_rank = prefill_dp_rank; write_tokens_to_result(&tokens, out); + write_cache_namespace_to_result(cache_namespace.as_deref(), out); QueryRouterResult::Ok } Err(code) => code, @@ -1429,7 +1516,7 @@ pub unsafe extern "C" fn route_decode_request( .query_decode_worker( &tokens, is_disaggregated, - cache_namespace, + cache_namespace.clone(), priority_jump, strict_priority, allowed_worker_ids, @@ -1458,6 +1545,7 @@ pub unsafe extern "C" fn route_decode_request( out.decode_worker_id = decode_worker.worker_id; out.decode_dp_rank = decode_worker.dp_rank; write_tokens_to_result(&tokens, out); + write_cache_namespace_to_result(cache_namespace.as_deref(), out); QueryRouterResult::Ok } Err(code) => code, @@ -1713,4 +1801,37 @@ mod tests { .expect("test request must parse as completion"); assert_eq!(extract_priority_jump(req.nvext.as_ref()), 5.0); } + + #[test] + fn cache_namespace_supports_legacy_top_level_field() { + let request: dynamo_llm::types::openai::chat_completions::NvCreateChatCompletionRequest = + serde_json::from_str( + r#"{ + "model": "test", + "messages": [{"role": "user", "content": "hi"}], + "cache_salt": "tenant-legacy" + }"#, + ) + .expect("test request must parse as chat completion"); + assert_eq!( + extract_cache_namespace(&request).as_deref(), + Some("tenant-legacy") + ); + } + + #[test] + fn routing_result_round_trips_cache_namespace_bytes() { + let mut result = CRoutingResult::default(); + write_cache_namespace_to_result(Some("tenant\0a"), &mut result); + + assert_eq!(result.cache_namespace_len, 8); + let namespace = unsafe { + std::slice::from_raw_parts(result.cache_namespace, result.cache_namespace_len) + }; + assert_eq!(namespace, b"tenant\0a"); + + unsafe { free_routing_result(&mut result) }; + assert!(result.cache_namespace.is_null()); + assert_eq!(result.cache_namespace_len, 0); + } } diff --git a/lib/kv-router/src/indexer/tests.rs b/lib/kv-router/src/indexer/tests.rs index 07793dc8e89d..38d5faee34a7 100644 --- a/lib/kv-router/src/indexer/tests.rs +++ b/lib/kv-router/src/indexer/tests.rs @@ -1605,6 +1605,38 @@ mod lora_tests { assert_eq!(scores_b.scores.len(), 1); assert!(scores_b.scores.contains_key(&WorkerWithDpRank::new(1, 0))); assert!(!scores_b.scores.contains_key(&WorkerWithDpRank::new(0, 0))); + + let request_scores_a = index + .find_matches_for_request(&tokens, None, Some("tenant-a"), None) + .await + .unwrap(); + assert_eq!(request_scores_a.scores.len(), 1); + assert!( + request_scores_a + .scores + .contains_key(&WorkerWithDpRank::new(0, 0)) + ); + assert!( + !request_scores_a + .scores + .contains_key(&WorkerWithDpRank::new(1, 0)) + ); + + let request_scores_b = index + .find_matches_for_request(&tokens, None, Some("tenant-b"), None) + .await + .unwrap(); + assert_eq!(request_scores_b.scores.len(), 1); + assert!( + request_scores_b + .scores + .contains_key(&WorkerWithDpRank::new(1, 0)) + ); + assert!( + !request_scores_b + .scores + .contains_key(&WorkerWithDpRank::new(0, 0)) + ); } } diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 2f37bea7196b..9bffffac3dfd 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -1014,6 +1014,7 @@ mod tests { sequence_hashes: None, isl_tokens: None, lora_name: None, + cache_namespace: None, is_eagle: None, } } @@ -1241,6 +1242,7 @@ mod tests { sequence_hashes: Some(vec![1, 2]), isl_tokens: Some(8), lora_name: None, + cache_namespace: None, is_eagle: None, }, router_config_override: None, @@ -1267,6 +1269,7 @@ mod tests { sequence_hashes: Some(vec![101, 102]), isl_tokens: Some(8), lora_name: None, + cache_namespace: None, is_eagle: None, }, router_config_override: None, diff --git a/lib/kv-router/src/services/selection/input.rs b/lib/kv-router/src/services/selection/input.rs index 923dca203285..35fd32b63103 100644 --- a/lib/kv-router/src/services/selection/input.rs +++ b/lib/kv-router/src/services/selection/input.rs @@ -36,6 +36,8 @@ pub struct PromptRequest { pub isl_tokens: Option, #[serde(default)] pub lora_name: Option, + #[serde(default, rename = "cache_salt")] + pub cache_namespace: Option, #[serde(default)] pub is_eagle: Option, } @@ -51,6 +53,7 @@ impl PromptRequest { token_ids, block_size, self.lora_name.as_deref(), + self.cache_namespace.as_deref(), block_mm_infos, self.is_eagle.unwrap_or(default_is_eagle), )); @@ -78,6 +81,7 @@ impl PromptRequest { token_ids, block_size, self.lora_name.as_deref(), + self.cache_namespace.as_deref(), block_mm_infos, self.is_eagle.unwrap_or(default_is_eagle), ); @@ -121,6 +125,7 @@ fn normalize_tokens( token_ids: &[u32], block_size: u32, lora_name: Option<&str>, + cache_namespace: Option<&str>, block_mm_infos: Option<&[Option]>, is_eagle: bool, ) -> NormalizedPrompt { @@ -130,7 +135,7 @@ fn normalize_tokens( BlockHashOptions { block_mm_infos, lora_name, - cache_namespace: None, + cache_namespace, is_eagle: Some(is_eagle), }, ); diff --git a/lib/kv-router/src/services/selection/tests.rs b/lib/kv-router/src/services/selection/tests.rs index f48c73a4add7..dc7e8887102f 100644 --- a/lib/kv-router/src/services/selection/tests.rs +++ b/lib/kv-router/src/services/selection/tests.rs @@ -114,6 +114,7 @@ fn prompt_normalization_uses_mm_routing_info_and_eagle_hashing() { sequence_hashes: None, isl_tokens: None, lora_name: Some("adapter".to_string()), + cache_namespace: Some("tenant-a".to_string()), is_eagle: Some(true), }; @@ -126,6 +127,7 @@ fn prompt_normalization_uses_mm_routing_info_and_eagle_hashing() { BlockHashOptions { block_mm_infos: Some(&mm_infos), lora_name: Some("adapter"), + cache_namespace: Some("tenant-a"), is_eagle: Some(true), }, ); @@ -137,6 +139,29 @@ fn prompt_normalization_uses_mm_routing_info_and_eagle_hashing() { assert_eq!(normalized.isl_tokens, 8); } +#[test] +fn prompt_request_cache_salt_changes_normalized_hashes() { + let salted: PromptRequest = serde_json::from_value(serde_json::json!({ + "token_ids": [1, 2, 3, 4], + "cache_salt": "tenant-a" + })) + .expect("deserialize cache_salt"); + let unsalted: PromptRequest = serde_json::from_value(serde_json::json!({ + "token_ids": [1, 2, 3, 4] + })) + .expect("deserialize unsalted prompt"); + + let salted = salted + .normalize_for_selection(4, false) + .expect("normalize salted prompt"); + let unsalted = unsalted + .normalize_for_selection(4, false) + .expect("normalize unsalted prompt"); + + assert_ne!(salted.block_hashes, unsalted.block_hashes); + assert_ne!(salted.sequence_hashes, unsalted.sequence_hashes); +} + #[test] fn overlap_scores_response_honors_override_and_includes_python_shape_fields() { let worker = WorkerWithDpRank::new(1, 0); diff --git a/lib/kv-router/src/zmq_wire/extra_keys.rs b/lib/kv-router/src/zmq_wire/extra_keys.rs index ff26aa373d45..f6882a77e54f 100644 --- a/lib/kv-router/src/zmq_wire/extra_keys.rs +++ b/lib/kv-router/src/zmq_wire/extra_keys.rs @@ -32,7 +32,9 @@ fn cache_namespace_candidate<'a>(value: &'a str, lora_name: Option<&str>) -> Opt /// Extract a vLLM cache salt from `extra_keys` when a producer does not emit /// top-level `cache_salt`. vLLM aligns `extra_keys` with blocks and includes -/// request-wide extras in each block, so the first block is enough. +/// cache salt only in the first block. Only MessagePack string values are +/// candidates; byte values such as prompt-embedding hashes must never become +/// cache namespaces. pub fn extra_keys_to_cache_namespace( extra_keys: Option<&[Option>]>, lora_name: Option<&str>, @@ -44,11 +46,8 @@ pub fn extra_keys_to_cache_namespace( | ExtraKeyItem::HashWithUnsignedOffset((hash, _)) => { cache_namespace_candidate(hash, lora_name).map(str::to_owned) } - ExtraKeyItem::Bytes(bytes) => std::str::from_utf8(bytes) - .ok() - .and_then(|value| cache_namespace_candidate(value, lora_name)) - .map(str::to_owned), - ExtraKeyItem::Signed(_) + ExtraKeyItem::Bytes(_) + | ExtraKeyItem::Signed(_) | ExtraKeyItem::Unsigned(_) | ExtraKeyItem::Float(_) | ExtraKeyItem::Bool(_) => None, @@ -110,3 +109,14 @@ pub fn extra_keys_to_block_mm_infos( Some(infos) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn prompt_embedding_bytes_are_not_cache_namespace() { + let extra_keys = [Some(vec![ExtraKeyItem::Bytes(b"prompt-embed".to_vec())])]; + assert_eq!(extra_keys_to_cache_namespace(Some(&extra_keys), None), None); + } +} diff --git a/lib/kv-router/src/zmq_wire/mod.rs b/lib/kv-router/src/zmq_wire/mod.rs index 4cb83b524516..194a374cfb86 100644 --- a/lib/kv-router/src/zmq_wire/mod.rs +++ b/lib/kv-router/src/zmq_wire/mod.rs @@ -47,7 +47,13 @@ pub struct ZmqEventNormalizer { image_token_id: Option, warning_count: Arc, group_metadata: FxHashMap<(DpRank, u32), KvCacheGroupMetadata>, - cache_namespaces: FxHashMap, + cache_namespaces: FxHashMap<(WorkerWithDpRank, u64), CacheNamespaceState>, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum CacheNamespaceState { + Namespaced(String), + Ambiguous, } #[derive(Debug, Clone, Copy)] @@ -59,6 +65,7 @@ struct KvCacheGroupMetadata { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ZmqEventFilterReason { IgnoredEvent, + AmbiguousCacheNamespace, NonMainAttentionKind, UnknownKind, NonMainAttentionGroup, @@ -69,6 +76,7 @@ impl ZmqEventFilterReason { pub fn as_label(self) -> &'static str { match self { Self::IgnoredEvent => "ignored_event", + Self::AmbiguousCacheNamespace => "ambiguous_cache_namespace", Self::NonMainAttentionKind => "non_main_attention_kind", Self::UnknownKind => "unknown_kind", Self::NonMainAttentionGroup => "non_main_attention_group", @@ -126,7 +134,7 @@ impl ZmqEventNormalizer { if let Some(reason) = self.filter_reason(metadata, worker.dp_rank) { return Err(reason); } - self.propagate_cache_namespace(&mut raw); + self.propagate_cache_namespace(&mut raw, worker)?; Ok(raw) } @@ -171,7 +179,11 @@ impl ZmqEventNormalizer { ); } - fn propagate_cache_namespace(&mut self, raw: &mut RawKvEvent) { + fn propagate_cache_namespace( + &mut self, + raw: &mut RawKvEvent, + worker: WorkerWithDpRank, + ) -> Result<(), ZmqEventFilterReason> { match raw { RawKvEvent::BlockStored { block_hashes, @@ -181,33 +193,65 @@ impl ZmqEventNormalizer { } => { if cache_namespace.is_none() && let Some(parent) = parent_block_hash.as_ref() - && let Some(namespace) = - self.cache_namespaces.get(&(*parent).into_u64()).cloned() { - *cache_namespace = Some(namespace); + match self.cache_namespaces.get(&(worker, (*parent).into_u64())) { + Some(CacheNamespaceState::Namespaced(namespace)) => { + *cache_namespace = Some(namespace.clone()); + } + Some(CacheNamespaceState::Ambiguous) => { + return Err(ZmqEventFilterReason::AmbiguousCacheNamespace); + } + None => {} + } } - if let Some(namespace) = cache_namespace.as_ref().filter(|ns| !ns.is_empty()) { + if let Some(namespace) = cache_namespace + .as_ref() + .filter(|namespace| !namespace.is_empty()) + { + let state = CacheNamespaceState::Namespaced(namespace.clone()); for block_hash in block_hashes.iter() { self.cache_namespaces - .insert((*block_hash).into_u64(), namespace.clone()); + .entry((worker, (*block_hash).into_u64())) + .and_modify(|existing| { + if *existing != state { + *existing = CacheNamespaceState::Ambiguous; + } + }) + .or_insert_with(|| state.clone()); } } else { + // Do not retain every unsalted block. If this hash already + // belongs to a namespace, however, fail closed on future + // propagation because the external hash is now ambiguous. for block_hash in block_hashes.iter() { - self.cache_namespaces.remove(&(*block_hash).into_u64()); + if let Some(existing) = self + .cache_namespaces + .get_mut(&(worker, (*block_hash).into_u64())) + { + *existing = CacheNamespaceState::Ambiguous; + } } } } RawKvEvent::BlockRemoved { block_hashes, .. } => { for block_hash in block_hashes.iter() { - self.cache_namespaces.remove(&(*block_hash).into_u64()); + let key = (worker, (*block_hash).into_u64()); + if !matches!( + self.cache_namespaces.get(&key), + Some(CacheNamespaceState::Ambiguous) + ) { + self.cache_namespaces.remove(&key); + } } } RawKvEvent::AllBlocksCleared => { - self.cache_namespaces.clear(); + self.cache_namespaces + .retain(|(known_worker, _), _| *known_worker != worker); } RawKvEvent::Ignored => {} } + Ok(()) } fn filter_reason( diff --git a/lib/kv-router/src/zmq_wire/tests.rs b/lib/kv-router/src/zmq_wire/tests.rs index ee0d9df34c24..ea85e74c82c9 100644 --- a/lib/kv-router/src/zmq_wire/tests.rs +++ b/lib/kv-router/src/zmq_wire/tests.rs @@ -551,6 +551,44 @@ fn test_normalizer_propagates_cache_namespace_from_parent() { assert_eq!(cache_namespace.as_deref(), Some("tenant-a")); } +#[test] +fn test_normalizer_rejects_ambiguous_parent_cache_namespace() { + let worker = WorkerWithDpRank::new(7, 0); + let mut normalizer = ZmqEventNormalizer::new(2); + let stored = + |cache_namespace: Option<&str>, block_hashes, parent_block_hash| RawKvEvent::BlockStored { + block_hashes, + parent_block_hash, + token_ids: vec![10, 11], + block_size: 2, + medium: None, + lora_name: None, + cache_namespace: cache_namespace.map(str::to_owned), + block_mm_infos: None, + is_eagle: Some(false), + group_idx: None, + kv_cache_spec_kind: None, + kv_cache_spec_sliding_window: None, + }; + + let parent_a = stored(Some("tenant-a"), vec![BlockHashValue::Unsigned(1)], None); + let parent_b = stored(Some("tenant-b"), vec![BlockHashValue::Unsigned(1)], None); + let child = stored( + None, + vec![BlockHashValue::Unsigned(2)], + Some(BlockHashValue::Unsigned(1)), + ); + + assert!(normalizer.preprocess(parent_a, worker).is_some()); + assert!(normalizer.preprocess(parent_b, worker).is_some()); + assert_eq!( + normalizer + .preprocess_with_reason(child, worker) + .expect_err("ambiguous parent must be rejected"), + ZmqEventFilterReason::AmbiguousCacheNamespace + ); +} + #[test] fn test_normalizer_ignores_non_main_attention_kind_with_group_idx_zero() { let raw_event: RawKvEvent = from_slice(&sequence_with_cache_spec_kind( diff --git a/lib/llm/src/kv_router/publisher/tests.rs b/lib/llm/src/kv_router/publisher/tests.rs index 04f666eb33c1..e8009245a4e1 100644 --- a/lib/llm/src/kv_router/publisher/tests.rs +++ b/lib/llm/src/kv_router/publisher/tests.rs @@ -100,6 +100,33 @@ mod test_event_processing { assert_eq!(blocks.len(), 2); assert_eq!(blocks[0].block_hash.0, 111); assert_eq!(blocks[1].block_hash.0, 222); + + let salted_blocks = create_stored_blocks( + kv_block_size, + &token_ids, + &num_block_tokens, + &block_hashes, + None, + Some("tenant-a"), + &Arc::new(AtomicU32::new(0)), + None, + None, + None, + ); + for (block, tokens) in salted_blocks + .iter() + .zip(token_ids.chunks(kv_block_size as usize)) + { + let expected = compute_block_hash_for_seq( + tokens, + kv_block_size, + BlockHashOptions { + cache_namespace: Some("tenant-a"), + ..Default::default() + }, + )[0]; + assert_eq!(block.tokens_hash, expected); + } } #[test] diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index 3313353ef5ae..0954f3ce2d2c 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -304,6 +304,19 @@ impl OpenAIPreprocessor { } } + fn request_cache_namespace(request: &R) -> Option { + request + .nvext() + .and_then(|nvext| nvext.cache_salt.clone()) + .or_else(|| { + request + .unsupported_fields() + .and_then(|fields| fields.get("cache_salt")) + .and_then(|value| value.as_str()) + .map(str::to_owned) + }) + } + fn nvext_passthrough_args( request: &R, ) -> Option> { @@ -313,9 +326,6 @@ impl OpenAIPreprocessor { if let Some(ref fields) = nvext.extra_fields { nvext_passthrough.insert("extra_fields".to_string(), serde_json::json!(fields)); } - if let Some(ref salt) = nvext.cache_salt { - nvext_passthrough.insert("cache_salt".to_string(), serde_json::json!(salt)); - } if let Some(ref metadata_upload) = nvext.metadata_upload { nvext_passthrough.insert( "metadata_upload".to_string(), @@ -327,12 +337,7 @@ impl OpenAIPreprocessor { } } - if !nvext_passthrough.contains_key("cache_salt") - && let Some(salt) = request - .unsupported_fields() - .and_then(|fields| fields.get("cache_salt")) - .and_then(|value| value.as_str()) - { + if let Some(salt) = Self::request_cache_namespace(request) { nvext_passthrough.insert("cache_salt".to_string(), serde_json::json!(salt)); } @@ -838,6 +843,7 @@ impl OpenAIPreprocessor { builder.annotations(request.annotations().unwrap_or_default()); builder.mdc_sum(Some(self.mdcsum.clone())); let lora_name = self.lora_name.clone(); + let cache_namespace = Self::request_cache_namespace(request); // Extract routing hints from nvext if present if let Some(nvext) = request.nvext() { @@ -856,7 +862,7 @@ impl OpenAIPreprocessor { strict_priority, priority, lora_name, - cache_namespace: nvext.cache_salt.clone(), + cache_namespace: cache_namespace.clone(), allowed_worker_ids: None, routing_constraints: nvext .routing_constraints @@ -864,11 +870,12 @@ impl OpenAIPreprocessor { .map(routing_constraints_to_kv), }; builder.routing(Some(routing)); - } else if lora_name.is_some() { - // Ensure routing hints exist when we have LoRA, - // even when nvext is absent. + } else if lora_name.is_some() || cache_namespace.is_some() { + // Ensure routing hints exist when we have LoRA or a legacy + // top-level cache_salt, even when nvext is absent. builder.routing(Some(RoutingHints { lora_name, + cache_namespace, ..Default::default() })); } @@ -3695,6 +3702,32 @@ mod tests { ); } + #[test] + fn test_request_cache_namespace_supports_legacy_top_level_field() { + let legacy: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "hi"}], + "cache_salt": "tenant-legacy" + })) + .unwrap(); + assert_eq!( + OpenAIPreprocessor::request_cache_namespace(&legacy).as_deref(), + Some("tenant-legacy") + ); + + let nvext_wins: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "hi"}], + "cache_salt": "tenant-legacy", + "nvext": {"cache_salt": "tenant-nvext"} + })) + .unwrap(); + assert_eq!( + OpenAIPreprocessor::request_cache_namespace(&nvext_wins).as_deref(), + Some("tenant-nvext") + ); + } + #[test] fn test_internal_preserve_omitted_max_tokens_option() { assert_eq!( From e2b014ad28252461719f7e6b0f3cccc2545deab9 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Tue, 30 Jun 2026 12:39:14 -0700 Subject: [PATCH 04/10] fix(trtllm): preserve per-block cache salt events Signed-off-by: jthomson04 --- components/src/dynamo/trtllm/llm_engine.py | 7 ++- components/src/dynamo/trtllm/publisher.py | 3 +- .../trtllm/tests/test_trtllm_fpm_publisher.py | 2 +- .../tests/test_trtllm_kv_event_adapter.py | 60 +++++++++++++++++++ .../trtllm/tests/test_trtllm_request_utils.py | 46 ++++++++++++++ .../src/dynamo/trtllm/utils/request_utils.py | 27 +++++++++ 6 files changed, 141 insertions(+), 4 deletions(-) create mode 100644 components/src/dynamo/trtllm/tests/test_trtllm_kv_event_adapter.py create mode 100644 components/src/dynamo/trtllm/tests/test_trtllm_request_utils.py diff --git a/components/src/dynamo/trtllm/llm_engine.py b/components/src/dynamo/trtllm/llm_engine.py index dc77b313d22e..774a32ce73d6 100644 --- a/components/src/dynamo/trtllm/llm_engine.py +++ b/components/src/dynamo/trtllm/llm_engine.py @@ -68,7 +68,10 @@ DisaggregatedParams, DisaggregatedParamsCodec, ) -from dynamo.trtllm.utils.request_utils import request_cache_salt +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: @@ -588,7 +591,7 @@ def _dispatch_kv_event(self, event: dict[str, Any]) -> None: block_hashes, parent_hash, lora_name=data.get("lora_name"), - cache_salt=data.get("cache_salt"), + cache_salt=stored_event_cache_salt(data), ) elif kind == "removed": partial = self._partial_block_hashes_by_rank.get(rank) diff --git a/components/src/dynamo/trtllm/publisher.py b/components/src/dynamo/trtllm/publisher.py index 96ca03b4876b..4608dc8bbdf4 100644 --- a/components/src/dynamo/trtllm/publisher.py +++ b/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__) @@ -871,7 +872,7 @@ def _handle_kv_event(self, event): block_mm_infos.append(None) lora_name = data.get("lora_name") - cache_salt = data.get("cache_salt") + cache_salt = stored_event_cache_salt(data) logger.debug( "Publishing stored KV event: engine_event_id=%s " diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py index 88728f89ab2f..a6d77b9c5426 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py @@ -386,10 +386,10 @@ def _stored_kv_event(cache_salt="tenant-a"): "data": { "type": "stored", "parent_hash": None, - "cache_salt": cache_salt, "blocks": [ { "block_hash": 123, + "cache_salt": cache_salt, "tokens": [ {"token_id": 1}, {"token_id": 2}, diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_kv_event_adapter.py b/components/src/dynamo/trtllm/tests/test_trtllm_kv_event_adapter.py new file mode 100644 index 000000000000..6012445ce01c --- /dev/null +++ b/components/src/dynamo/trtllm/tests/test_trtllm_kv_event_adapter.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +try: + from dynamo.trtllm.llm_engine import TrtllmLLMEngine +except ImportError: + pytest.skip("tensorrt_llm backend not available", allow_module_level=True) + +pytestmark = [ + pytest.mark.unit, + pytest.mark.trtllm, + pytest.mark.gpu_1, + pytest.mark.pre_merge, +] + + +def _stored_kv_event(cache_salt: str | None = "tenant-a") -> dict: + return { + "event_id": 1, + "attention_dp_rank": 0, + "data": { + "type": "stored", + "parent_hash": None, + "blocks": [ + { + "type": "stored_block", + "block_hash": 123, + "cache_salt": cache_salt, + "tokens": [ + {"token_id": 1}, + {"token_id": 2}, + {"token_id": 3}, + {"token_id": 4}, + ], + } + ], + }, + } + + +def test_dispatch_kv_event_forwards_per_block_cache_salt() -> None: + engine = TrtllmLLMEngine.__new__(TrtllmLLMEngine) + publisher = MagicMock() + engine._kv_publishers = {0: publisher} + engine._last_event_id_by_rank = {} + engine._warned_unknown_dp_rank = False + engine._additional_metrics = None + engine._partial_block_hashes_by_rank = {} + engine.kv_block_size = 4 + + engine._dispatch_kv_event(_stored_kv_event()) + + publisher.publish_stored.assert_called_once() + assert publisher.publish_stored.call_args.kwargs["cache_salt"] == "tenant-a" diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_request_utils.py b/components/src/dynamo/trtllm/tests/test_trtllm_request_utils.py new file mode 100644 index 000000000000..471a04131ed9 --- /dev/null +++ b/components/src/dynamo/trtllm/tests/test_trtllm_request_utils.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from dynamo.trtllm.utils.request_utils import stored_event_cache_salt + +pytestmark = [ + pytest.mark.unit, + pytest.mark.trtllm, + pytest.mark.gpu_0, + pytest.mark.pre_merge, +] + + +def test_stored_event_cache_salt_uses_per_block_schema() -> None: + data = { + "blocks": [ + {"cache_salt": "tenant-a"}, + {"cache_salt": "tenant-a"}, + ] + } + + assert stored_event_cache_salt(data) == "tenant-a" + + +def test_stored_event_cache_salt_supports_parent_fallback() -> None: + assert stored_event_cache_salt({"cache_salt": "tenant-a", "blocks": []}) == ( + "tenant-a" + ) + + +def test_stored_event_cache_salt_allows_unsalted_blocks() -> None: + assert stored_event_cache_salt({"blocks": [{}, {}]}) is None + + +def test_stored_event_cache_salt_rejects_conflicting_blocks() -> None: + data = { + "blocks": [ + {"cache_salt": "tenant-a"}, + {"cache_salt": "tenant-b"}, + ] + } + + with pytest.raises(ValueError, match="conflicting cache_salt"): + stored_event_cache_salt(data) diff --git a/components/src/dynamo/trtllm/utils/request_utils.py b/components/src/dynamo/trtllm/utils/request_utils.py index cbef9333ece8..a7e8aea1ddf9 100644 --- a/components/src/dynamo/trtllm/utils/request_utils.py +++ b/components/src/dynamo/trtllm/utils/request_utils.py @@ -21,3 +21,30 @@ def request_cache_salt(request: Mapping[str, Any]) -> Optional[str]: return cache_salt return None + + +def stored_event_cache_salt(data: Mapping[str, Any]) -> Optional[str]: + """Extract one cache salt from a TRT-LLM stored-event payload. + + TRT-LLM 1.3 serializes ``cache_salt`` on each item in ``data["blocks"]``. + Keep the parent-level lookup as a compatibility fallback, but fail closed + if an event combines blocks from different cache namespaces. + """ + cache_salts: set[str] = set() + + parent_cache_salt = data.get("cache_salt") + if parent_cache_salt is not None: + cache_salts.add(parent_cache_salt) + + blocks = data.get("blocks") or [] + for block in blocks: + if not isinstance(block, Mapping): + continue + block_cache_salt = block.get("cache_salt") + if block_cache_salt is not None: + cache_salts.add(block_cache_salt) + + if len(cache_salts) > 1: + raise ValueError("stored KV event contains conflicting cache_salt values") + + return next(iter(cache_salts), None) From 1cb79baddc9c40b7a6efcc5f957087ae42c2b411 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Tue, 30 Jun 2026 17:22:04 -0700 Subject: [PATCH 05/10] fix(router): address cache salt review feedback Signed-off-by: jthomson04 --- .github/workflows/pr.yaml | 2 - .../epp/pkg/plugins/disagg/decode_scorer.go | 2 +- deploy/inference-gateway/ext-proc/src/epp.rs | 60 +-------- docs/components/frontend/nvext.md | 32 +++++ docs/components/router/standalone-indexer.md | 9 +- lib/bindings/c/src/lib.rs | 39 +----- lib/kv-router/src/services/indexer/server.rs | 16 ++- .../tests/standalone_indexer_http.rs | 123 +++++++++++++++++- lib/llm/src/preprocessor.rs | 45 +------ lib/llm/src/protocols/common/extensions.rs | 65 +++++++++ 10 files changed, 249 insertions(+), 144 deletions(-) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 89fba5de9d82..0692bd16ace1 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -131,8 +131,6 @@ jobs: operator: needs: changed-files - # Keep this trigger set a superset of every deploy-operator* job guarded by - # needs.operator.result == 'success'. if: | needs.changed-files.outputs.operator == 'true' || needs.changed-files.outputs.vllm == 'true' || diff --git a/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go index ca89eadd65b0..040c7fcb7ebe 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go @@ -241,7 +241,7 @@ func (s *DynDecodeScorer) PreRequest(ctx context.Context, request *schedtypes.In "requestID", request.RequestId, "workerID", state.WorkerID, "dpRank", state.DpRank, - "cacheNamespace", state.CacheNamespace, + "hasCacheNamespace", state.CacheNamespace != "", "tokenCount", len(state.TokenData)) } diff --git a/deploy/inference-gateway/ext-proc/src/epp.rs b/deploy/inference-gateway/ext-proc/src/epp.rs index 6ee8789ba621..5e5d3ae1dd1d 100644 --- a/deploy/inference-gateway/ext-proc/src/epp.rs +++ b/deploy/inference-gateway/ext-proc/src/epp.rs @@ -20,6 +20,7 @@ use dynamo_llm::kv_router::prefill_router::PrefillQueryOutcome; use dynamo_llm::kv_router::{KvRouter, PrefillRouter}; use dynamo_llm::model_card::ModelDeploymentCard; use dynamo_llm::preprocessor::OpenAIPreprocessor; +use dynamo_llm::protocols::common::extensions::request_cache_salt; use dynamo_runtime::discovery::{DiscoveryInstance, DiscoveryQuery, hash_pod_name}; use dynamo_runtime::pipeline::RouterMode; use dynamo_runtime::{DistributedRuntime, Runtime}; @@ -209,7 +210,7 @@ impl Router { let priority_jump = extract_priority_jump(&request); let strict_priority = extract_strict_priority(&request); - let cache_namespace = extract_cache_namespace(&request); + let cache_namespace = request_cache_salt(&request).map(str::to_owned); let formatted_prompt = self .preprocessor @@ -502,22 +503,6 @@ fn extract_strict_priority( .unwrap_or(0) } -fn extract_cache_namespace( - request: &dynamo_llm::types::openai::chat_completions::NvCreateChatCompletionRequest, -) -> Option { - request - .nvext - .as_ref() - .and_then(|nvext| nvext.cache_salt.clone()) - .or_else(|| { - request - .unsupported_fields - .get("cache_salt") - .and_then(|value| value.as_str()) - .map(str::to_owned) - }) -} - struct DiscoveredModelBootstrap { preprocessor: Arc, card: ModelDeploymentCard, @@ -1164,45 +1149,4 @@ mod tests { .unwrap(); assert_eq!(extract_strict_priority(&without_nvext), 0); } - - #[test] - fn cache_namespace_lifted_from_nvext_cache_salt() { - let with_cache_salt: dynamo_llm::types::openai::chat_completions::NvCreateChatCompletionRequest = - serde_json::from_str( - r#"{ - "model": "test", - "messages": [{"role": "user", "content": "hi"}], - "nvext": {"cache_salt": "tenant-a"} - }"#, - ) - .unwrap(); - assert_eq!( - extract_cache_namespace(&with_cache_salt).as_deref(), - Some("tenant-a") - ); - - let without_nvext: dynamo_llm::types::openai::chat_completions::NvCreateChatCompletionRequest = - serde_json::from_str( - r#"{ - "model": "test", - "messages": [{"role": "user", "content": "hi"}] - }"#, - ) - .unwrap(); - assert_eq!(extract_cache_namespace(&without_nvext), None); - - let legacy_top_level: dynamo_llm::types::openai::chat_completions::NvCreateChatCompletionRequest = - serde_json::from_str( - r#"{ - "model": "test", - "messages": [{"role": "user", "content": "hi"}], - "cache_salt": "tenant-legacy" - }"#, - ) - .unwrap(); - assert_eq!( - extract_cache_namespace(&legacy_top_level).as_deref(), - Some("tenant-legacy") - ); - } } diff --git a/docs/components/frontend/nvext.md b/docs/components/frontend/nvext.md index d0a10b5cc02b..285fd0ae8d18 100644 --- a/docs/components/frontend/nvext.md +++ b/docs/components/frontend/nvext.md @@ -36,6 +36,7 @@ Include `nvext` as a top-level field alongside standard OpenAI-compatible fields | `backend_instance_id` | `u64` | `None` | Router | Routes the request to a specific backend instance. | | `token_data` | `u32[]` | `None` | Preprocessor | Pre-tokenized prompt tokens. When provided with `backend_instance_id`, tokenization is skipped. | | `max_thinking_tokens` | `u32` | `None` | Backend | Maximum thinking tokens allowed (passed through to backends). | +| `cache_salt` | `string` | `None` | Router/backend | Isolates KV-cache routing and reuse to requests carrying the same non-empty salt. This is the recommended cache-isolation input. | | `extra_fields` | `string[]` | `None` | Response builder | Fields to include in the response `nvext`. Supported: `"worker_id"`, `"timing"`, `"routed_experts"`, `"engine_data"`, `"stop_reason"`. | | `prefill_worker_id` | `u64` | `None` | Router | Routes the request to a specific prefill worker (disaggregated serving). | | `decode_worker_id` | `u64` | `None` | Router | Routes the request to a specific decode worker (disaggregated serving). | @@ -64,12 +65,43 @@ Routing fields can also be set via HTTP headers, which take priority over `nvext | `x-dynamo-prefill-instance-id` | `prefill_worker_id` | | `x-dynamo-dp-rank` | `dp_rank` | | `x-dynamo-prefill-dp-rank` | `prefill_dp_rank` | +| `x-tenant-id` | `cache_salt` | > [!WARNING] > The unprefixed forms (`x-worker-instance-id`, `x-prefill-instance-id`, `x-dp-rank`, > `x-data-parallel-rank`, and `x-prefill-dp-rank`) are compatibility aliases planned for future > deprecation. Use the `x-dynamo-*` headers for new integrations. +### Cache salt and tenant isolation + +Use `nvext.cache_salt` to prevent requests in different cache namespaces from matching or reusing +each other's KV-cache blocks: + +```json +{ + "model": "my-model", + "messages": [{"role": "user", "content": "Hello"}], + "nvext": { + "cache_salt": "tenant-a" + } +} +``` + +Dynamo accepts three inputs, in descending precedence: + +1. The non-empty `x-tenant-id` HTTP header, intended for gateway-controlled tenant identity. +2. The recommended `nvext.cache_salt` request field. +3. The compatibility top-level `cache_salt` field on chat and completion requests. + +Empty strings are treated as absent. In particular, an empty `nvext.cache_salt` falls back to a +non-empty top-level compatibility value. Requests without a salt retain the unsalted hashing and +cache-reuse behavior. + +`DYN_ENABLE_FRONTEND_NVEXT=false` disables both the `nvext` form and routing-header overrides, +including `x-tenant-id`. The top-level backend-compatibility field is not part of the NvExt +protocol. Cache salt is an isolation key, not an authentication or authorization mechanism; +gateways must still authenticate the tenant identity they place in `x-tenant-id`. + Session identity is header-only. Use the coding-agent headers or Dynamo session headers described in [Session IDs](../../agents/session-ids.md); `nvext` does not accept session identity fields. diff --git a/docs/components/router/standalone-indexer.md b/docs/components/router/standalone-indexer.md index 9b2be07a97c7..e27821709bfc 100644 --- a/docs/components/router/standalone-indexer.md +++ b/docs/components/router/standalone-indexer.md @@ -343,7 +343,7 @@ All counts are in **matched tokens** (block overlap count × block size). | `model_name` | yes | — | Model name (selects the indexer) | | `tenant_id` | no | `"default"` | Tenant identifier | | `lora_name` | no | — | LoRA adapter (overrides indexer-level lora_name for this query) | -| `cache_salt` | no | — | Per-request cache salt (Mooncake RFC #1403). Currently parsed for forward compatibility — engines apply their own salting today. | +| `cache_salt` | no | — | Per-request cache salt (Mooncake RFC #1403). The indexer mixes it into hashes computed from `token_ids`; equal tokens with different salts do not match. | ### `POST /query_by_hash` — Query overlap for pre-computed hashes @@ -360,7 +360,12 @@ Same response format as `/query`, including the per-instance `instances` map. Sc | `block_hashes` | yes | — | Pre-computed block hash array | | `model_name` | yes | — | Model name (selects the indexer) | | `tenant_id` | no | `"default"` | Tenant identifier | -| `cache_salt` | no | — | Per-request cache salt (Mooncake RFC #1403). Currently parsed for forward compatibility — engines apply their own salting today. | +| `cache_salt` | no | — | Must be omitted or `null`. Any string value, including an empty string, returns `400 Bad Request`. | + +`block_hashes` are opaque outputs of token hashing, so the indexer cannot apply or verify a salt +after they have been computed. Callers must precompute these hashes with the intended cache salt +and omit `cache_salt` from `/query_by_hash`. Use `/query` when the indexer should compute salted +hashes from tokens server-side. ### Per-instance tier breakdown diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index e69a6fff314d..9783e74a7ad0 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -19,7 +19,7 @@ use dynamo_llm::kv_router::publisher::KvEventPublisher; use dynamo_llm::model_card::ModelDeploymentCard; use dynamo_llm::preprocessor::OpenAIPreprocessor; use dynamo_llm::protocols::common::extensions::{ - NvExt, NvExtProvider, routing_constraints_to_kv, + NvExt, request_cache_salt, routing_constraints_to_kv, }; use dynamo_llm::types::openai::chat_completions::NvCreateChatCompletionRequest; use dynamo_llm::types::openai::completions::NvCreateCompletionRequest; @@ -623,20 +623,6 @@ fn extract_routing_constraints(nvext: Option<&NvExt>) -> RoutingConstraints { .map(routing_constraints_to_kv) .unwrap_or_default() } - -fn extract_cache_namespace(request: &R) -> Option { - request - .nvext() - .and_then(|nvext| nvext.cache_salt.clone()) - .or_else(|| { - request - .unsupported_fields - .get("cache_salt") - .and_then(|value| value.as_str()) - .map(str::to_owned) - }) -} - /// Opaque handle for the router pair pub type RouterHandlesPtr = *mut RouterHandles; @@ -1016,7 +1002,7 @@ pub unsafe extern "C" fn add_request_with_cache_namespace( request_id = %request_id_str, worker_id = worker_id, dp_rank = dp_rank, - cache_namespace = cache_namespace.as_deref(), + has_cache_namespace = cache_namespace.is_some(), overlap_blocks = overlap_blocks, token_count = tokens.len(), "add_request completed" @@ -1247,7 +1233,7 @@ unsafe fn preprocess_request( }; let priority_jump = extract_priority_jump(request.nvext.as_ref()); let strict_priority = extract_strict_priority(request.nvext.as_ref()); - let cache_namespace = extract_cache_namespace(&request); + let cache_namespace = request_cache_salt(&request).map(str::to_owned); let routing_constraints = extract_routing_constraints(request.nvext.as_ref()); let (token_ids, _) = match handles .runtime @@ -1288,7 +1274,7 @@ unsafe fn preprocess_request( let priority_jump = extract_priority_jump(request.nvext.as_ref()); let strict_priority = extract_strict_priority(request.nvext.as_ref()); - let cache_namespace = extract_cache_namespace(&request); + let cache_namespace = request_cache_salt(&request).map(str::to_owned); let routing_constraints = extract_routing_constraints(request.nvext.as_ref()); let formatted_prompt = match preprocessor.apply_template(&request) { @@ -1802,23 +1788,6 @@ mod tests { assert_eq!(extract_priority_jump(req.nvext.as_ref()), 5.0); } - #[test] - fn cache_namespace_supports_legacy_top_level_field() { - let request: dynamo_llm::types::openai::chat_completions::NvCreateChatCompletionRequest = - serde_json::from_str( - r#"{ - "model": "test", - "messages": [{"role": "user", "content": "hi"}], - "cache_salt": "tenant-legacy" - }"#, - ) - .expect("test request must parse as chat completion"); - assert_eq!( - extract_cache_namespace(&request).as_deref(), - Some("tenant-legacy") - ); - } - #[test] fn routing_result_round_trips_cache_namespace_bytes() { let mut result = CRoutingResult::default(); diff --git a/lib/kv-router/src/services/indexer/server.rs b/lib/kv-router/src/services/indexer/server.rs index a29fdc4b03d2..73d6004696b4 100644 --- a/lib/kv-router/src/services/indexer/server.rs +++ b/lib/kv-router/src/services/indexer/server.rs @@ -138,8 +138,8 @@ pub struct QueryByHashRequest { pub model_name: String, #[serde(default = "default_tenant")] pub tenant_id: String, - /// Optional per-request cache salt (Mooncake RFC #1403). For `/query_by_hash`, callers - /// must precompute `block_hashes` with the same salt. + /// Invalid for `/query_by_hash`. Callers must precompute `block_hashes` with the intended + /// cache salt and omit this field; a non-null value is rejected. #[serde(default)] pub cache_salt: Option, } @@ -353,6 +353,18 @@ async fn query_by_hash( Json(req): Json, ) -> Response { let model = req.model_name.clone(); + if req.cache_salt.is_some() { + let mut resp = ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ + "error": "cache_salt is not accepted by /query_by_hash; block_hashes must already include the intended cache salt" + })), + ) + .into_response(); + resp.extensions_mut().insert(AccessLogModel(model)); + return resp; + } + let key = IndexerKey { model_name: req.model_name, tenant_id: req.tenant_id, diff --git a/lib/kv-router/tests/standalone_indexer_http.rs b/lib/kv-router/tests/standalone_indexer_http.rs index 2b86ff5b4cdf..1bbc4022021a 100644 --- a/lib/kv-router/tests/standalone_indexer_http.rs +++ b/lib/kv-router/tests/standalone_indexer_http.rs @@ -16,8 +16,9 @@ use std::sync::Arc; use std::time::Duration; use dynamo_kv_router::protocols::{ - ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheStoreData, - KvCacheStoredBlockData, LocalBlockHash, RouterEvent, StorageTier, compute_seq_hash_for_block, + BlockHashOptions, ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheStoreData, + KvCacheStoredBlockData, LocalBlockHash, RouterEvent, StorageTier, compute_block_hash_for_seq, + compute_seq_hash_for_block, }; use dynamo_kv_router::services::indexer::registry::{IndexerKey, WorkerRegistry}; use dynamo_kv_router::services::indexer::server::{AppState, create_router}; @@ -256,6 +257,124 @@ async fn query_by_hash_returns_per_instance_tier_breakdown() { ); assert_eq!(inst8["longest_matched"], (2 * BLOCK_SIZE) as u64); + let null_salt_resp = client + .post(format!("{base_url}/query_by_hash")) + .json(&json!({ + "block_hashes": [11_i64, 12, 13], + "model_name": MODEL, + "tenant_id": TENANT, + "cache_salt": null, + })) + .send() + .await + .expect("POST /query_by_hash with null cache_salt"); + assert_eq!(null_salt_resp.status(), reqwest::StatusCode::OK); + + cancel.cancel(); + task.await.expect("server task join"); +} + +/// `/query` owns token hashing, so equal tokens under different cache salts must match only the +/// worker whose stored blocks used the same salt. +#[tokio::test] +async fn query_isolates_cache_salts() { + const BLOCK_SIZE: u32 = 4; + const MODEL: &str = "test-model"; + const TENANT: &str = "default"; + let tokens = vec![1_u32, 2, 3, 4, 5, 6, 7, 8]; + + let hashes_a = compute_block_hash_for_seq( + &tokens, + BLOCK_SIZE, + BlockHashOptions { + cache_namespace: Some("tenant-a"), + ..Default::default() + }, + ); + let hashes_b = compute_block_hash_for_seq( + &tokens, + BLOCK_SIZE, + BlockHashOptions { + cache_namespace: Some("tenant-b"), + ..Default::default() + }, + ); + let events = vec![ + store_event( + 7, + 0, + 1, + &[], + &hashes_a.iter().map(|hash| hash.0).collect::>(), + StorageTier::Device, + ), + store_event( + 8, + 0, + 1, + &[], + &hashes_b.iter().map(|hash| hash.0).collect::>(), + StorageTier::Device, + ), + ]; + let registry = registry_with_events(MODEL, TENANT, BLOCK_SIZE, events).await; + let state = make_app_state(registry); + let (base_url, cancel, task) = spawn_indexer_http(state).await; + let client = reqwest::Client::new(); + + for (salt, expected_worker, other_worker) in [("tenant-a", "7", "8"), ("tenant-b", "8", "7")] { + let resp = client + .post(format!("{base_url}/query")) + .json(&json!({ + "token_ids": tokens.clone(), + "model_name": MODEL, + "tenant_id": TENANT, + "cache_salt": salt, + })) + .send() + .await + .expect("POST /query with cache_salt"); + assert_eq!(resp.status(), reqwest::StatusCode::OK); + let body: serde_json::Value = resp.json().await.expect("parse /query body"); + assert_eq!(body["scores"][expected_worker]["0"], tokens.len() as u64); + assert!(body["scores"].get(other_worker).is_none()); + } + + cancel.cancel(); + task.await.expect("server task join"); +} + +/// `/query_by_hash` cannot apply a cache salt after token hashes have already been produced. +#[tokio::test] +async fn query_by_hash_rejects_cache_salt() { + const MODEL: &str = "test-model"; + const TENANT: &str = "default"; + let registry = registry_with_events(MODEL, TENANT, 4, Vec::new()).await; + let state = make_app_state(registry); + let (base_url, cancel, task) = spawn_indexer_http(state).await; + let client = reqwest::Client::new(); + + for cache_salt in ["tenant-a", ""] { + let resp = client + .post(format!("{base_url}/query_by_hash")) + .json(&json!({ + "block_hashes": [11_i64, 12], + "model_name": MODEL, + "tenant_id": TENANT, + "cache_salt": cache_salt, + })) + .send() + .await + .expect("POST /query_by_hash with cache_salt"); + assert_eq!(resp.status(), reqwest::StatusCode::BAD_REQUEST); + let body: serde_json::Value = resp.json().await.expect("parse rejection body"); + assert!( + body["error"] + .as_str() + .is_some_and(|error| error.contains("block_hashes must already include")) + ); + } + cancel.cancel(); task.await.expect("server task join"); } diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index 0954f3ce2d2c..c08444bda6f7 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -64,7 +64,7 @@ use crate::protocols::{ TokenIdType, common::{ OutputOptionsProvider, SamplingOptionsProvider, StopConditionsProvider, - extensions::{AgentHints, NvExtProvider, routing_constraints_to_kv}, + extensions::{AgentHints, NvExtProvider, request_cache_salt, routing_constraints_to_kv}, }, openai::{ DeltaGeneratorExt, @@ -304,19 +304,6 @@ impl OpenAIPreprocessor { } } - fn request_cache_namespace(request: &R) -> Option { - request - .nvext() - .and_then(|nvext| nvext.cache_salt.clone()) - .or_else(|| { - request - .unsupported_fields() - .and_then(|fields| fields.get("cache_salt")) - .and_then(|value| value.as_str()) - .map(str::to_owned) - }) - } - fn nvext_passthrough_args( request: &R, ) -> Option> { @@ -337,7 +324,7 @@ impl OpenAIPreprocessor { } } - if let Some(salt) = Self::request_cache_namespace(request) { + if let Some(salt) = request_cache_salt(request) { nvext_passthrough.insert("cache_salt".to_string(), serde_json::json!(salt)); } @@ -843,7 +830,7 @@ impl OpenAIPreprocessor { builder.annotations(request.annotations().unwrap_or_default()); builder.mdc_sum(Some(self.mdcsum.clone())); let lora_name = self.lora_name.clone(); - let cache_namespace = Self::request_cache_namespace(request); + let cache_namespace = request_cache_salt(request).map(str::to_owned); // Extract routing hints from nvext if present if let Some(nvext) = request.nvext() { @@ -3702,32 +3689,6 @@ mod tests { ); } - #[test] - fn test_request_cache_namespace_supports_legacy_top_level_field() { - let legacy: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({ - "model": "test-model", - "messages": [{"role": "user", "content": "hi"}], - "cache_salt": "tenant-legacy" - })) - .unwrap(); - assert_eq!( - OpenAIPreprocessor::request_cache_namespace(&legacy).as_deref(), - Some("tenant-legacy") - ); - - let nvext_wins: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({ - "model": "test-model", - "messages": [{"role": "user", "content": "hi"}], - "cache_salt": "tenant-legacy", - "nvext": {"cache_salt": "tenant-nvext"} - })) - .unwrap(); - assert_eq!( - OpenAIPreprocessor::request_cache_namespace(&nvext_wins).as_deref(), - Some("tenant-nvext") - ); - } - #[test] fn test_internal_preserve_omitted_max_tokens_option() { assert_eq!( diff --git a/lib/llm/src/protocols/common/extensions.rs b/lib/llm/src/protocols/common/extensions.rs index c003c06b8e1f..4c6e1f2685e0 100644 --- a/lib/llm/src/protocols/common/extensions.rs +++ b/lib/llm/src/protocols/common/extensions.rs @@ -394,6 +394,26 @@ pub trait NvExtProvider { } } +/// Return the request's non-empty cache salt using Dynamo's public precedence rules. +/// +/// `nvext.cache_salt` is the canonical input. The top-level `cache_salt` field remains a +/// compatibility fallback for request types that retain unsupported OpenAI fields. Empty strings +/// are treated as absent so an empty canonical value can still fall back to a non-empty legacy +/// value. +pub fn request_cache_salt(request: &R) -> Option<&str> { + request + .nvext() + .and_then(|nvext| nvext.cache_salt.as_deref()) + .filter(|salt| !salt.is_empty()) + .or_else(|| { + request + .unsupported_fields() + .and_then(|fields| fields.get("cache_salt")) + .and_then(|value| value.as_str()) + .filter(|salt| !salt.is_empty()) + }) +} + pub fn routing_constraints_to_kv( constraints: RoutingConstraints, ) -> dynamo_kv_router::protocols::RoutingConstraints { @@ -638,6 +658,51 @@ mod tests { HEADER_OPENCODE_SESSION_ID, }; + #[derive(Default)] + struct CacheSaltRequest { + nvext: Option, + unsupported_fields: HashMap, + } + + impl NvExtProvider for CacheSaltRequest { + fn nvext(&self) -> Option<&NvExt> { + self.nvext.as_ref() + } + + fn raw_prompt(&self) -> Option { + None + } + + fn unsupported_fields(&self) -> Option<&HashMap> { + Some(&self.unsupported_fields) + } + } + + #[test] + fn request_cache_salt_uses_canonical_precedence_and_empty_fallbacks() { + let mut request = CacheSaltRequest::default(); + assert_eq!(request_cache_salt(&request), None); + + request + .unsupported_fields + .insert("cache_salt".to_string(), serde_json::json!("tenant-legacy")); + assert_eq!(request_cache_salt(&request), Some("tenant-legacy")); + + request.nvext = Some(NvExt { + cache_salt: Some("tenant-nvext".to_string()), + ..Default::default() + }); + assert_eq!(request_cache_salt(&request), Some("tenant-nvext")); + + request.nvext.as_mut().unwrap().cache_salt = Some(String::new()); + assert_eq!(request_cache_salt(&request), Some("tenant-legacy")); + + request + .unsupported_fields + .insert("cache_salt".to_string(), serde_json::json!("")); + assert_eq!(request_cache_salt(&request), None); + } + #[test] fn shared_nvext_builder_default() { let nv_ext = NvExt::builder().build().unwrap(); From 27d60e9cc2bad7b632c70377494d99c5002132a5 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Fri, 3 Jul 2026 14:33:30 -0700 Subject: [PATCH 06/10] fix(vllm): propagate cache salt in unified engine Signed-off-by: jthomson04 --- components/src/dynamo/vllm/llm_engine.py | 2 ++ .../vllm/tests/test_vllm_delta_streaming.py | 27 +++++++++++++++ tests/router/common.py | 2 +- tests/router/test_router_e2e_with_unified.py | 34 ++++++++++++++++++- 4 files changed, 63 insertions(+), 2 deletions(-) diff --git a/components/src/dynamo/vllm/llm_engine.py b/components/src/dynamo/vllm/llm_engine.py index 25688e2155af..ffae0839f820 100644 --- a/components/src/dynamo/vllm/llm_engine.py +++ b/components/src/dynamo/vllm/llm_engine.py @@ -72,6 +72,7 @@ from .handlers import ( VllmEnginePauseController, + _apply_nvext_cache_salt, build_sampling_params, get_dp_range_for_worker, ) @@ -400,6 +401,7 @@ async def generate( self.disaggregation_mode, ) prompt = prepared_prompt.prompt + _apply_nvext_cache_salt(prepared_prompt.request, prompt) # Multimodal decode may replace token_ids with the expanded prefill # sequence. Sampling limits must use that same effective request. diff --git a/components/src/dynamo/vllm/tests/test_vllm_delta_streaming.py b/components/src/dynamo/vllm/tests/test_vllm_delta_streaming.py index de9cffc33b5d..448638b300ec 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_delta_streaming.py +++ b/components/src/dynamo/vllm/tests/test_vllm_delta_streaming.py @@ -289,3 +289,30 @@ async def test_unified_llm_engine_passes_delta_chunks_and_counts_usage(): "completion_tokens": 3, "total_tokens": 5, } + + +@pytest.mark.asyncio +async def test_unified_llm_engine_forwards_cache_salt_to_prompt(): + pytest.importorskip("vllm.usage.usage_lib") + from dynamo.vllm.llm_engine import VllmLLMEngine + + engine = VllmLLMEngine.__new__(VllmLLMEngine) + engine.engine_client = _FakeEngineClient([]) + engine._default_sampling_params = {} + engine._model_max_len = None + engine.disaggregation_mode = DisaggregationMode.AGGREGATED + engine.enable_rl = False + engine._dp_range = None + + request = { + "token_ids": [10, 11], + "sampling_options": {}, + "stop_conditions": {}, + "output_options": {}, + "extra_args": {"nvext": {"cache_salt": "tenant-a"}}, + } + async for _ in VllmLLMEngine.generate(engine, request, _FakeContext()): + pass + + prompt = engine.engine_client.calls[0][0][0] + assert prompt["cache_salt"] == "tenant-a" diff --git a/tests/router/common.py b/tests/router/common.py index 5b4834280d9c..302384ade0b2 100644 --- a/tests/router/common.py +++ b/tests/router/common.py @@ -3157,7 +3157,7 @@ def _test_router_cache_salt_isolation( model_name: str, block_size: int, ): - """Verify cache-salted TRT-LLM events remain isolated in the router index.""" + """Verify cache-salted engine events remain isolated in the router index.""" async def test_sync(): expected_num_instances = engine_workers.num_workers diff --git a/tests/router/test_router_e2e_with_unified.py b/tests/router/test_router_e2e_with_unified.py index 9e2fae719b67..e72beea1de28 100644 --- a/tests/router/test_router_e2e_with_unified.py +++ b/tests/router/test_router_e2e_with_unified.py @@ -213,6 +213,38 @@ def test_unified_vllm_router_decisions_multiple_workers( ) +@pytest.mark.pre_merge +@pytest.mark.gpu_1 +@pytest.mark.vllm +@pytest.mark.model(VLLM_MODEL_NAME) +@pytest.mark.profiled_vram_gib(6.9) +@pytest.mark.requested_vllm_kv_cache_bytes(331_801_000) +@pytest.mark.timeout(360) +@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True) +def test_unified_vllm_cache_salt_isolation( + request, + runtime_services_dynamic_ports, + predownload_models, + set_ucx_tls_no_mm, + request_plane, +) -> None: + """Cache-salted vLLM events remain isolated in the router index. + + This crosses the real unified vLLM engine and its KV-event publisher, then + queries the router index independently for each tenant namespace. + """ + run_cache_salt_isolation_test( + engine_process_cls=UnifiedVLLMProcess, + engine_args_name="vllm_args", + engine_args=VLLM_ARGS, + request=request, + request_plane=request_plane, + model_name=VLLM_MODEL_NAME, + block_size=VLLM_BLOCK_SIZE, + component_name="backend", + ) + + @pytest.mark.gpu_2 @pytest.mark.nightly @pytest.mark.vllm @@ -439,7 +471,7 @@ def test_unified_trtllm_cache_salt_isolation( predownload_models, request_plane, ) -> None: - """Identical prompts under different cache salts never share KV entries. + """Cache-salted TRT-LLM events remain isolated in the router index. This crosses the real unified TRT-LLM engine and its KV-event publisher, then queries the router index independently for each tenant namespace. From e320c03521d91c68c2e435d45261cfe36f10198d Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Fri, 3 Jul 2026 15:45:29 -0700 Subject: [PATCH 07/10] fix(kv-router): address cache salt review feedback Signed-off-by: jthomson04 --- Cargo.lock | 1 + components/src/dynamo/trtllm/publisher.py | 12 ++- .../trtllm/tests/test_trtllm_fpm_publisher.py | 43 +++++++++++ .../trtllm/tests/test_trtllm_request_utils.py | 38 +++++++++- .../src/dynamo/trtllm/utils/request_utils.py | 6 +- components/src/dynamo/vllm/handlers.py | 16 +++- .../vllm/tests/test_vllm_delta_streaming.py | 2 +- .../vllm/tests/test_vllm_tito_parity.py | 30 ++++++-- deploy/inference-gateway/ext-proc/src/epp.rs | 44 ++++++++++- docs/components/frontend/nvext.md | 15 +++- lib/kv-router/Cargo.toml | 1 + lib/kv-router/src/protocols.rs | 26 ++++++- lib/kv-router/src/zmq_wire/extra_keys.rs | 74 ++++++++++++------- lib/kv-router/src/zmq_wire/mod.rs | 13 +++- lib/kv-router/src/zmq_wire/tests.rs | 72 +++++++++++++++++- 15 files changed, 343 insertions(+), 50 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2bfffbb3062a..250f5683cf19 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2493,6 +2493,7 @@ dependencies = [ "chrono", "dashmap", "derive_builder", + "dynamo-kv-hashing", "dynamo-runtime", "dynamo-tokens", "flume 0.12.0", diff --git a/components/src/dynamo/trtllm/publisher.py b/components/src/dynamo/trtllm/publisher.py index 4608dc8bbdf4..22d0f7cf18ad 100644 --- a/components/src/dynamo/trtllm/publisher.py +++ b/components/src/dynamo/trtllm/publisher.py @@ -872,7 +872,17 @@ def _handle_kv_event(self, event): block_mm_infos.append(None) lora_name = data.get("lora_name") - cache_salt = stored_event_cache_salt(data) + 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 " diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py index a6d77b9c5426..cc089b02116d 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py @@ -30,6 +30,7 @@ from __future__ import annotations import asyncio +import logging import queue import threading from unittest.mock import MagicMock @@ -425,6 +426,48 @@ def test_handle_kv_event_forwards_cache_salt_to_zmq_publisher(): assert pub.zmq_kv_event_publisher.publish_stored.call_args.args[-1] == "tenant-a" +@pytest.mark.asyncio +async def test_polling_loop_drops_conflicting_salts_and_processes_next_event(caplog): + pub = _publisher_for_kv_event_test() + pub._stop_event = threading.Event() + pub.zmq_kv_event_publisher = None + publisher = MagicMock() + pub.kv_event_publishers = {0: publisher} + + conflicting = _stored_kv_event("tenant-a") + conflicting["data"]["blocks"].append( + { + **conflicting["data"]["blocks"][0], + "block_hash": 456, + "cache_salt": "tenant-b", + } + ) + valid = _stored_kv_event("tenant-c") + valid["event_id"] = 2 + + async def fetch_events(): + yield conflicting + yield valid + + def handle_event(event): + pub._handle_kv_event(event) + if event["event_id"] == 2: + pub._stop_event.set() + + with caplog.at_level(logging.WARNING): + await pub._polling_loop( + fetch_events, + handle_event, + min_sleep=0.001, + max_sleep=0.001, + backoff_factor=1.0, + ) + + publisher.publish_stored.assert_called_once() + assert publisher.publish_stored.call_args.kwargs["cache_salt"] == "tenant-c" + assert "Dropping stored KV event with invalid cache namespace" in caplog.text + + def test_publisher_initializes_fpm_publisher_under_attention_dp(monkeypatch): """Under attention-DP (attention_dp_size > 1), Publisher.initialize() constructs one FpmDirectPublisher channel per attention-DP rank.""" diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_request_utils.py b/components/src/dynamo/trtllm/tests/test_trtllm_request_utils.py index 471a04131ed9..f27d5a2643ff 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_request_utils.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_request_utils.py @@ -3,7 +3,10 @@ import pytest -from dynamo.trtllm.utils.request_utils import stored_event_cache_salt +from dynamo.trtllm.utils.request_utils import ( + request_cache_salt, + stored_event_cache_salt, +) pytestmark = [ pytest.mark.unit, @@ -13,6 +16,39 @@ ] +@pytest.mark.parametrize( + ("request_body", "expected"), + [ + ( + { + "routing": {"cache_salt": "tenant-routing"}, + "extra_args": {"nvext": {"cache_salt": "tenant-nvext"}}, + }, + "tenant-routing", + ), + ( + { + "routing": {"cache_salt": ""}, + "extra_args": {"nvext": {"cache_salt": "tenant-nvext"}}, + }, + "tenant-nvext", + ), + ( + { + "routing": {"cache_salt": ""}, + "extra_args": {"nvext": {"cache_salt": ""}}, + }, + None, + ), + ({}, None), + ], +) +def test_request_cache_salt_precedence_and_empty_fallback( + request_body, expected +) -> None: + assert request_cache_salt(request_body) == expected + + def test_stored_event_cache_salt_uses_per_block_schema() -> None: data = { "blocks": [ diff --git a/components/src/dynamo/trtllm/utils/request_utils.py b/components/src/dynamo/trtllm/utils/request_utils.py index a7e8aea1ddf9..8930da12023d 100644 --- a/components/src/dynamo/trtllm/utils/request_utils.py +++ b/components/src/dynamo/trtllm/utils/request_utils.py @@ -6,18 +6,18 @@ def request_cache_salt(request: Mapping[str, Any]) -> Optional[str]: - """Return cache_salt using routing hints before legacy extra_args.""" + """Return the first non-empty cache_salt, preferring routing hints.""" routing = request.get("routing") or {} if isinstance(routing, dict): cache_salt = routing.get("cache_salt") - if cache_salt is not None: + if cache_salt: return cache_salt extra_args = request.get("extra_args") or {} nvext = extra_args.get("nvext") if isinstance(extra_args, dict) else None if isinstance(nvext, dict): cache_salt = nvext.get("cache_salt") - if cache_salt is not None: + if cache_salt: return cache_salt return None diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index 263027dacfcd..323c21791c7f 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -463,13 +463,25 @@ def _nvext_extra_field_requested(request: Dict[str, Any], field: str) -> bool: ) +# Must match DYNAMO_CACHE_SALT_PREFIX in lib/kv-router/src/zmq_wire/extra_keys.rs. +_DYNAMO_CACHE_SALT_PREFIX = "dynamo-cache-salt:" + + def _apply_nvext_cache_salt(request: Dict[str, Any], prompt: Any) -> None: + """Pass an internally tagged cache salt to vLLM. + + vLLM publishes cache salts as otherwise-untyped strings in ``extra_keys`` + alongside LoRA and multimodal metadata. The tag lets Dynamo recover the + namespace without guessing from the user-controlled value. It is removed + again by the Rust KV-event decoder, so Dynamo's public namespace is + unchanged. + """ if not isinstance(prompt, dict): return for source in _iter_nvext_sources(request): cache_salt = source.get("cache_salt") - if cache_salt is not None: - prompt["cache_salt"] = cache_salt + if cache_salt: + prompt["cache_salt"] = f"{_DYNAMO_CACHE_SALT_PREFIX}{cache_salt}" return diff --git a/components/src/dynamo/vllm/tests/test_vllm_delta_streaming.py b/components/src/dynamo/vllm/tests/test_vllm_delta_streaming.py index 448638b300ec..923d7d93bbe9 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_delta_streaming.py +++ b/components/src/dynamo/vllm/tests/test_vllm_delta_streaming.py @@ -315,4 +315,4 @@ async def test_unified_llm_engine_forwards_cache_salt_to_prompt(): pass prompt = engine.engine_client.calls[0][0][0] - assert prompt["cache_salt"] == "tenant-a" + assert prompt["cache_salt"] == "dynamo-cache-salt:tenant-a" diff --git a/components/src/dynamo/vllm/tests/test_vllm_tito_parity.py b/components/src/dynamo/vllm/tests/test_vllm_tito_parity.py index 80c9fd50dae3..d835e205e111 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_tito_parity.py +++ b/components/src/dynamo/vllm/tests/test_vllm_tito_parity.py @@ -84,7 +84,7 @@ def test_multiple_tokens_per_position(self): class TestCacheSaltWiring: - """Verify cache_salt is extracted from extra_args and placed on the prompt.""" + """Verify cache_salt is extracted and tagged for vLLM's prompt.""" @staticmethod def _build_token_mode_request(cache_salt=None, token_ids=None): @@ -100,7 +100,7 @@ def _build_token_mode_request(cache_salt=None, token_ids=None): return req def test_cache_salt_attached_to_prompt(self): - """When extra_args.nvext.cache_salt is set, the prompt dict gets it.""" + """The prompt receives an internal tag around the public cache salt.""" from vllm.inputs import TokensPrompt from dynamo.vllm.handlers import _apply_nvext_cache_salt @@ -109,7 +109,7 @@ def test_cache_salt_attached_to_prompt(self): prompt = TokensPrompt(prompt_token_ids=req["token_ids"]) _apply_nvext_cache_salt(req, prompt) - assert prompt.get("cache_salt") == "step_42" + assert prompt.get("cache_salt") == "dynamo-cache-salt:step_42" def test_no_cache_salt_when_absent(self): """When extra_args has no cache_salt, prompt should not gain the key.""" @@ -134,8 +134,8 @@ def test_prefill_and_decode_share_cache_salt_helper(self): _apply_nvext_cache_salt(req, prefill_prompt) _apply_nvext_cache_salt(req, decode_prompt) - assert prefill_prompt["cache_salt"] == "step_43" - assert decode_prompt["cache_salt"] == "step_43" + assert prefill_prompt["cache_salt"] == "dynamo-cache-salt:step_43" + assert decode_prompt["cache_salt"] == "dynamo-cache-salt:step_43" def test_cache_salt_from_top_level_nvext(self): """cache_salt under the raw request["nvext"] shape is also honored, @@ -146,7 +146,25 @@ def test_cache_salt_from_top_level_nvext(self): prompt = {"prompt_token_ids": req["token_ids"]} _apply_nvext_cache_salt(req, prompt) - assert prompt["cache_salt"] == "top_level" + assert prompt["cache_salt"] == "dynamo-cache-salt:top_level" + + def test_empty_cache_salt_is_absent(self): + from dynamo.vllm.handlers import _apply_nvext_cache_salt + + req = self._build_token_mode_request(cache_salt="") + prompt = {"prompt_token_ids": req["token_ids"]} + _apply_nvext_cache_salt(req, prompt) + + assert "cache_salt" not in prompt + + def test_cache_salt_prefix_is_escaped_by_reprefixing(self): + from dynamo.vllm.handlers import _apply_nvext_cache_salt + + req = self._build_token_mode_request(cache_salt="dynamo-cache-salt:tenant-a") + prompt = {"prompt_token_ids": req["token_ids"]} + _apply_nvext_cache_salt(req, prompt) + + assert prompt["cache_salt"] == ("dynamo-cache-salt:dynamo-cache-salt:tenant-a") class TestTokenInSamplingDefaults: diff --git a/deploy/inference-gateway/ext-proc/src/epp.rs b/deploy/inference-gateway/ext-proc/src/epp.rs index 5e5d3ae1dd1d..97fd9c4447e2 100644 --- a/deploy/inference-gateway/ext-proc/src/epp.rs +++ b/deploy/inference-gateway/ext-proc/src/epp.rs @@ -20,11 +20,12 @@ use dynamo_llm::kv_router::prefill_router::PrefillQueryOutcome; use dynamo_llm::kv_router::{KvRouter, PrefillRouter}; use dynamo_llm::model_card::ModelDeploymentCard; use dynamo_llm::preprocessor::OpenAIPreprocessor; -use dynamo_llm::protocols::common::extensions::request_cache_salt; +use dynamo_llm::protocols::common::extensions::{HEADER_TENANT_ID, request_cache_salt}; use dynamo_runtime::discovery::{DiscoveryInstance, DiscoveryQuery, hash_pod_name}; use dynamo_runtime::pipeline::RouterMode; use dynamo_runtime::{DistributedRuntime, Runtime}; +use crate::envoy_helpers::find_header; use crate::picker::{Endpoint, EndpointPicker, PickError, PickResult, RequestInfo}; const BOOKKEEPING_TIMEOUT: Duration = Duration::from_secs(5); @@ -66,6 +67,16 @@ fn decode_router_config_override(is_disaggregated: bool) -> Option, +) -> Option { + find_header(headers, HEADER_TENANT_ID) + .filter(|tenant_id| !tenant_id.is_empty()) + .map(str::to_owned) + .or(body_cache_namespace) +} + /// Name of the inference-serving HTTP port on a Dynamo worker pod. /// /// Mirrors `commonconsts.DynamoContainerPortName` in @@ -920,9 +931,11 @@ impl EndpointPicker for Router { let body_str = std::str::from_utf8(&req.body) .map_err(|e| PickError::TokenizationFailed(format!("Invalid UTF-8: {e}")))?; - let (tokens, cache_namespace, priority_jump, strict_priority) = self + let (tokens, body_cache_namespace, priority_jump, strict_priority) = self .tokenize(body_str) .map_err(|e| PickError::TokenizationFailed(e.to_string()))?; + let cache_namespace = + cache_namespace_with_header_override(&req.headers, body_cache_namespace); // Try prefill routing first (disaggregated mode). // @@ -1099,6 +1112,33 @@ impl EndpointPicker for Router { mod tests { use super::*; + #[test] + fn tenant_header_overrides_body_cache_namespace() { + let headers = vec![("X-Tenant-ID".to_string(), "tenant-header".to_string())]; + + assert_eq!( + cache_namespace_with_header_override(&headers, Some("tenant-body".to_string())) + .as_deref(), + Some("tenant-header") + ); + } + + #[test] + fn empty_tenant_header_falls_back_to_body_cache_namespace() { + let headers = vec![(HEADER_TENANT_ID.to_string(), String::new())]; + + assert_eq!( + cache_namespace_with_header_override(&headers, Some("tenant-body".to_string())) + .as_deref(), + Some("tenant-body") + ); + } + + #[test] + fn absent_cache_namespace_stays_absent() { + assert_eq!(cache_namespace_with_header_override(&[], None), None); + } + /// Proves the core feature: `nvext.agent_hints.priority` lifts into a /// non-zero `priority_jump`, and absence collapses to `0.0`. If this /// regresses, the GAIE ext-proc path is back to ignoring priority. diff --git a/docs/components/frontend/nvext.md b/docs/components/frontend/nvext.md index 285fd0ae8d18..561376c88445 100644 --- a/docs/components/frontend/nvext.md +++ b/docs/components/frontend/nvext.md @@ -36,7 +36,7 @@ Include `nvext` as a top-level field alongside standard OpenAI-compatible fields | `backend_instance_id` | `u64` | `None` | Router | Routes the request to a specific backend instance. | | `token_data` | `u32[]` | `None` | Preprocessor | Pre-tokenized prompt tokens. When provided with `backend_instance_id`, tokenization is skipped. | | `max_thinking_tokens` | `u32` | `None` | Backend | Maximum thinking tokens allowed (passed through to backends). | -| `cache_salt` | `string` | `None` | Router/backend | Isolates KV-cache routing and reuse to requests carrying the same non-empty salt. This is the recommended cache-isolation input. | +| `cache_salt` | `string` | `None` | Router / supported backends | Namespaces Dynamo KV routing. vLLM and TensorRT-LLM also isolate backend KV-cache reuse; see [Backend support](#backend-support). This is the recommended cache-isolation input. | | `extra_fields` | `string[]` | `None` | Response builder | Fields to include in the response `nvext`. Supported: `"worker_id"`, `"timing"`, `"routed_experts"`, `"engine_data"`, `"stop_reason"`. | | `prefill_worker_id` | `u64` | `None` | Router | Routes the request to a specific prefill worker (disaggregated serving). | | `decode_worker_id` | `u64` | `None` | Router | Routes the request to a specific decode worker (disaggregated serving). | @@ -74,8 +74,9 @@ Routing fields can also be set via HTTP headers, which take priority over `nvext ### Cache salt and tenant isolation -Use `nvext.cache_salt` to prevent requests in different cache namespaces from matching or reusing -each other's KV-cache blocks: +Use `nvext.cache_salt` to namespace KV-cache routing. Dynamo also forwards the salt to supported +backend engines so identical prompts in different namespaces cannot reuse the same backend +KV-cache entries: ```json { @@ -87,6 +88,14 @@ each other's KV-cache blocks: } ``` +#### Backend support + +| Backend | Support | Behavior | +|---------|---------|----------| +| vLLM | Supported | Router matching and backend KV-cache reuse are isolated by salt. | +| TensorRT-LLM | Supported | Router matching and backend KV-cache reuse are isolated by salt. | +| SGLang | Not supported end to end | Dynamo request hashes are namespaced, but the embedded SGLang engine does not receive the salt. SGLang KV events and radix-cache reuse remain unsalted. Do not rely on `cache_salt` for tenant cache isolation with SGLang. | + Dynamo accepts three inputs, in descending precedence: 1. The non-empty `x-tenant-id` HTTP header, intended for gateway-controlled tenant identity. diff --git a/lib/kv-router/Cargo.toml b/lib/kv-router/Cargo.toml index d0a977b6ba72..d3167702a228 100644 --- a/lib/kv-router/Cargo.toml +++ b/lib/kv-router/Cargo.toml @@ -64,6 +64,7 @@ reqwest = { workspace = true, optional = true } tracing-appender = { version = "0.2", optional = true } [dev-dependencies] +dynamo-kv-hashing = { workspace = true } rstest = "0.18.2" rstest_reuse = "0.7.0" serde_json = { workspace = true } diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index df6b4a3590ec..b1410b191430 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -21,7 +21,6 @@ pub const KV_EVENT_SUBJECT: &str = "kv-events"; /// Seed for XXH3 hashing, consistent with indexer.rs pub const XXH3_SEED: u64 = 1337; -const LORA_HASH_SEED: u64 = XXH3_SEED ^ 0x9e37_79b1_85eb_ca87; const CACHE_NAMESPACE_HASH_SEED: u64 = XXH3_SEED ^ 0xc2b2_ae3d_27d4_eb4f; /// Compute the hash of a local block. @@ -40,7 +39,8 @@ pub struct BlockHashOptions<'a> { fn block_hash_seed(options: BlockHashOptions<'_>) -> u64 { let mut seed = XXH3_SEED; if let Some(name) = options.lora_name.filter(|n| !n.is_empty()) { - seed = seed.wrapping_add(xxh3::xxh3_64_with_seed(name.as_bytes(), LORA_HASH_SEED)); + // Preserve the established LoRA formula shared with dynamo-kv-hashing. + seed = seed.wrapping_add(xxh3::xxh3_64(name.as_bytes())); } if let Some(namespace) = options.cache_namespace.filter(|n| !n.is_empty()) { seed = seed.wrapping_add(xxh3::xxh3_64_with_seed( @@ -1304,6 +1304,28 @@ mod tests { assert_ne!(lora_a[0], lora_b[0]); } + #[test] + fn test_lora_hash_matches_kv_hashing_contract() { + let tokens: Vec = (0..4).collect(); + let lora_name = "adapter-a"; + let actual = compute_block_hash_for_seq( + &tokens, + 4, + BlockHashOptions { + lora_name: Some(lora_name), + ..Default::default() + }, + ); + let token_bytes = tokens + .iter() + .flat_map(|token| token.to_le_bytes()) + .collect::>(); + let salt_hash = dynamo_kv_hashing::compute_salt_hash(None, Some(lora_name)).unwrap(); + let expected = LocalBlockHash(dynamo_kv_hashing::compute_hash_v2(&token_bytes, salt_hash)); + + assert_eq!(actual, vec![expected]); + } + #[test] fn test_lora_name_empty_string_normalized_to_none() { let tokens: Vec = (0..4).collect(); diff --git a/lib/kv-router/src/zmq_wire/extra_keys.rs b/lib/kv-router/src/zmq_wire/extra_keys.rs index f6882a77e54f..f702c35c7171 100644 --- a/lib/kv-router/src/zmq_wire/extra_keys.rs +++ b/lib/kv-router/src/zmq_wire/extra_keys.rs @@ -5,6 +5,9 @@ use crate::protocols::{BlockExtraInfo, BlockMmObjectInfo}; use super::types::ExtraKeyItem; +// Must match _DYNAMO_CACHE_SALT_PREFIX in components/src/dynamo/vllm/handlers.py. +const DYNAMO_CACHE_SALT_PREFIX: &str = "dynamo-cache-salt:"; + /// Parse MM hash from extra_keys string: /// - Only accept canonical vLLM MM identifiers (64-char hex digest) /// - Convert by taking the first 16 hex chars as u64 @@ -17,40 +20,30 @@ pub fn parse_mm_hash_from_extra_key(s: &str) -> Option { None } -fn cache_namespace_candidate<'a>(value: &'a str, lora_name: Option<&str>) -> Option<&'a str> { - if value.is_empty() { - return None; - } - if lora_name.is_some_and(|name| name == value) { - return None; - } - if parse_mm_hash_from_extra_key(value).is_some() { - return None; - } - Some(value) -} - /// Extract a vLLM cache salt from `extra_keys` when a producer does not emit /// top-level `cache_salt`. vLLM aligns `extra_keys` with blocks and includes -/// cache salt only in the first block. Only MessagePack string values are -/// candidates; byte values such as prompt-embedding hashes must never become -/// cache namespaces. +/// cache salt only in the first block. Dynamo tags the opaque value before +/// passing it to vLLM because bare strings are otherwise ambiguous with LoRA +/// names and legacy multimodal hashes. The first matching LoRA item is skipped +/// before looking for the tag so a salt equal to the LoRA name still works. pub fn extra_keys_to_cache_namespace( extra_keys: Option<&[Option>]>, lora_name: Option<&str>, ) -> Option { let first_block = extra_keys?.first()?.as_ref()?; - first_block.iter().find_map(|key| match key { - ExtraKeyItem::Hash(hash) - | ExtraKeyItem::HashWithSignedOffset((hash, _)) - | ExtraKeyItem::HashWithUnsignedOffset((hash, _)) => { - cache_namespace_candidate(hash, lora_name).map(str::to_owned) + let mut unmatched_lora = lora_name.filter(|name| !name.is_empty()); + first_block.iter().find_map(|key| { + let ExtraKeyItem::Hash(value) = key else { + return None; + }; + if unmatched_lora.is_some_and(|name| name == value) { + unmatched_lora = None; + return None; } - ExtraKeyItem::Bytes(_) - | ExtraKeyItem::Signed(_) - | ExtraKeyItem::Unsigned(_) - | ExtraKeyItem::Float(_) - | ExtraKeyItem::Bool(_) => None, + value + .strip_prefix(DYNAMO_CACHE_SALT_PREFIX) + .filter(|namespace| !namespace.is_empty()) + .map(str::to_owned) }) } @@ -119,4 +112,33 @@ mod tests { let extra_keys = [Some(vec![ExtraKeyItem::Bytes(b"prompt-embed".to_vec())])]; assert_eq!(extra_keys_to_cache_namespace(Some(&extra_keys), None), None); } + + #[test] + fn untagged_strings_are_not_cache_namespaces() { + let extra_keys = [Some(vec![ExtraKeyItem::Hash("tenant-a".to_string())])]; + assert_eq!(extra_keys_to_cache_namespace(Some(&extra_keys), None), None); + } + + #[test] + fn tagged_cache_namespace_is_decoded() { + let extra_keys = [Some(vec![ExtraKeyItem::Hash( + "dynamo-cache-salt:tenant-a".to_string(), + )])]; + assert_eq!( + extra_keys_to_cache_namespace(Some(&extra_keys), None).as_deref(), + Some("tenant-a") + ); + } + + #[test] + fn cache_namespace_equal_to_lora_name_is_decoded() { + let extra_keys = [Some(vec![ + ExtraKeyItem::Hash("adapter-a".to_string()), + ExtraKeyItem::Hash("dynamo-cache-salt:adapter-a".to_string()), + ])]; + assert_eq!( + extra_keys_to_cache_namespace(Some(&extra_keys), Some("adapter-a")).as_deref(), + Some("adapter-a") + ); + } } diff --git a/lib/kv-router/src/zmq_wire/mod.rs b/lib/kv-router/src/zmq_wire/mod.rs index 194a374cfb86..1f0af44df4cf 100644 --- a/lib/kv-router/src/zmq_wire/mod.rs +++ b/lib/kv-router/src/zmq_wire/mod.rs @@ -191,6 +191,9 @@ impl ZmqEventNormalizer { cache_namespace, .. } => { + if cache_namespace.as_deref() == Some("") { + *cache_namespace = None; + } if cache_namespace.is_none() && let Some(parent) = parent_block_hash.as_ref() { @@ -201,7 +204,15 @@ impl ZmqEventNormalizer { Some(CacheNamespaceState::Ambiguous) => { return Err(ZmqEventFilterReason::AmbiguousCacheNamespace); } - None => {} + None => { + // Deliberately preserve the unsalted interpretation when a + // listener joins in the middle of a chain. The vLLM wire + // format cannot distinguish an unknown salted parent from a + // genuinely unsalted one, and retaining every unsalted block + // here would duplicate the index on the event hot path. The + // backend still enforces its own cache isolation; this narrow + // fail-open case can only pollute the router overlap score. + } } } diff --git a/lib/kv-router/src/zmq_wire/tests.rs b/lib/kv-router/src/zmq_wire/tests.rs index ea85e74c82c9..8fcba7929550 100644 --- a/lib/kv-router/src/zmq_wire/tests.rs +++ b/lib/kv-router/src/zmq_wire/tests.rs @@ -107,9 +107,9 @@ fn test_deserialize_extra_keys_cache_namespace_fallback() { let encoded = to_vec_named(&MapBlockStoredFixture { lora_name: Some("adapter-a".to_string()), extra_keys: Some(vec![Some(vec![ - mm_hash.to_string(), "adapter-a".to_string(), - "tenant-a".to_string(), + mm_hash.to_string(), + "dynamo-cache-salt:tenant-a".to_string(), ])]), ..Default::default() }) @@ -125,6 +125,30 @@ fn test_deserialize_extra_keys_cache_namespace_fallback() { assert_eq!(cache_namespace.as_deref(), Some("tenant-a")); } +#[test] +fn test_deserialize_hex_cache_namespace_is_not_multimodal() { + let cache_namespace = "0123456789abcdef00112233445566778899aabbccddeefffedcba9876543210"; + let encoded = to_vec_named(&MapBlockStoredFixture { + extra_keys: Some(vec![Some(vec![format!( + "dynamo-cache-salt:{cache_namespace}" + )])]), + ..Default::default() + }) + .unwrap(); + let event: RawKvEvent = from_slice(&encoded).unwrap(); + + let RawKvEvent::BlockStored { + cache_namespace: decoded_namespace, + block_mm_infos, + .. + } = event + else { + panic!("expected BlockStored"); + }; + assert_eq!(decoded_namespace.as_deref(), Some(cache_namespace)); + assert!(block_mm_infos.is_none()); +} + fn block_stored_sequence( group_idx: Option, kv_cache_spec_kind: Option<&'static str>, @@ -589,6 +613,50 @@ fn test_normalizer_rejects_ambiguous_parent_cache_namespace() { ); } +#[test] +fn test_normalizer_treats_empty_namespace_as_absent() { + let worker = WorkerWithDpRank::new(7, 0); + let mut normalizer = ZmqEventNormalizer::new(2); + let parent = RawKvEvent::BlockStored { + block_hashes: vec![BlockHashValue::Unsigned(1)], + parent_block_hash: None, + token_ids: vec![10, 11], + block_size: 2, + medium: None, + lora_name: None, + cache_namespace: Some("tenant-a".to_string()), + block_mm_infos: None, + is_eagle: Some(false), + group_idx: None, + kv_cache_spec_kind: None, + kv_cache_spec_sliding_window: None, + }; + let child = RawKvEvent::BlockStored { + block_hashes: vec![BlockHashValue::Unsigned(2)], + parent_block_hash: Some(BlockHashValue::Unsigned(1)), + token_ids: vec![12, 13], + block_size: 2, + medium: None, + lora_name: None, + cache_namespace: Some(String::new()), + block_mm_infos: None, + is_eagle: Some(false), + group_idx: None, + kv_cache_spec_kind: None, + kv_cache_spec_sliding_window: None, + }; + + assert!(normalizer.preprocess(parent, worker).is_some()); + let child = normalizer.preprocess(child, worker).unwrap(); + let RawKvEvent::BlockStored { + cache_namespace, .. + } = child + else { + panic!("expected BlockStored"); + }; + assert_eq!(cache_namespace.as_deref(), Some("tenant-a")); +} + #[test] fn test_normalizer_ignores_non_main_attention_kind_with_group_idx_zero() { let raw_event: RawKvEvent = from_slice(&sequence_with_cache_spec_kind( From 488319b36ff81e4dc14d021484108cf7a4597182 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Mon, 6 Jul 2026 13:58:26 -0700 Subject: [PATCH 08/10] fix(kv-router): preserve cache salt through consolidator Signed-off-by: jthomson04 --- components/src/dynamo/router/__main__.py | 9 +- .../router/tests/test_standalone_router.py | 125 +++++++++++++ lib/kv-hashing/src/salt.rs | 15 +- lib/kv-router/Cargo.toml | 2 +- lib/kv-router/src/protocols.rs | 40 +++-- lib/kv-router/src/zmq_wire/mod.rs | 36 ++-- lib/kv-router/src/zmq_wire/tests.rs | 42 +++++ .../src/egress/zmq_publisher.rs | 9 +- .../src/ingress/zmq_subscriber.rs | 7 +- lib/kvbm-consolidator/src/tracker.rs | 166 +++++++++++++++++- lib/kvbm-consolidator/src/wire/router_out.rs | 7 + lib/kvbm-consolidator/tests/common/mod.rs | 2 + lib/kvbm-consolidator/tests/e2e.rs | 8 + lib/kvbm-consolidator/tests/zmq_ingress.rs | 113 +++++++++++- 14 files changed, 535 insertions(+), 46 deletions(-) create mode 100644 components/src/dynamo/router/tests/test_standalone_router.py diff --git a/components/src/dynamo/router/__main__.py b/components/src/dynamo/router/__main__.py index 311fa87d567c..74a7d192a49a 100644 --- a/components/src/dynamo/router/__main__.py +++ b/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 diff --git a/components/src/dynamo/router/tests/test_standalone_router.py b/components/src/dynamo/router/tests/test_standalone_router.py new file mode 100644 index 000000000000..e1a166ca6843 --- /dev/null +++ b/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", + ) diff --git a/lib/kv-hashing/src/salt.rs b/lib/kv-hashing/src/salt.rs index f409e1995b14..2de92809bcdb 100644 --- a/lib/kv-hashing/src/salt.rs +++ b/lib/kv-hashing/src/salt.rs @@ -9,18 +9,19 @@ //! //! # Router parity //! -//! For requests with no extra `salt`, this function reproduces the seed used by -//! `dynamo_kv_router::protocols::compute_block_hash_for_seq` -//! (`lib/kv-router/src/protocols.rs:79-82`): +//! This function is the shared seed derivation used by +//! `dynamo_kv_router::protocols::compute_block_hash_for_seq` and canonical +//! pre-hashed producers: //! //! ```text //! (salt=None, lora=None) → CHAIN_XXH3_SEED //! (salt=None, lora=Some(name)) → CHAIN_XXH3_SEED.wrapping_add(xxh3_64(name)) +//! (salt=Some(value), ...) → seed.wrapping_add(xxh3_64_with_seed(value, 1)) //! ``` //! //! Producer events whose `block_hash` is `compute_block_hash(tokens, salt_hash)` therefore -//! match the router's `compute_block_hash_for_seq(tokens, _, BlockHashOptions { lora_name })` -//! byte-for-byte on the no-salt path — required for kv-router's indexers (which key on +//! match the router's `compute_block_hash_for_seq(tokens, _, BlockHashOptions { .. })` +//! byte-for-byte — required for kv-router's indexers (which key on //! both `tokens_hash` and the `seq_hash` chain) to find matches against consolidator-emitted //! events. //! @@ -57,8 +58,8 @@ pub fn compute_salt_hash( seed = seed.wrapping_add(compute_hash_v2(name.as_bytes(), 0)); } if let Some(s) = salt { - // Router has no concept of caller-supplied salt; mix orthogonally to lora so - // salt-isolated requests stay distinct from both no-salt and lora-only requests. + // Mix salt orthogonally to lora so salt-isolated requests stay distinct from + // both no-salt and lora-only requests. // The 1 vs 0 inner seed (and outer wrapping_add) keeps every (salt, lora) pair // separable: same lora + different salts diverge, and a future `(salt=lora_bytes, // lora=None)` request will not collide with `(salt=None, lora=lora_bytes)`. diff --git a/lib/kv-router/Cargo.toml b/lib/kv-router/Cargo.toml index d3167702a228..ce9e19d76c29 100644 --- a/lib/kv-router/Cargo.toml +++ b/lib/kv-router/Cargo.toml @@ -26,6 +26,7 @@ standalone-selection = ["standalone-indexer"] [dependencies] # repo dynamo-runtime = { workspace = true, optional = true } +dynamo-kv-hashing = { workspace = true } dynamo-tokens = { workspace = true } # workspace @@ -64,7 +65,6 @@ reqwest = { workspace = true, optional = true } tracing-appender = { version = "0.2", optional = true } [dev-dependencies] -dynamo-kv-hashing = { workspace = true } rstest = "0.18.2" rstest_reuse = "0.7.0" serde_json = { workspace = true } diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index b1410b191430..4a37c5025fd6 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -21,7 +21,6 @@ pub const KV_EVENT_SUBJECT: &str = "kv-events"; /// Seed for XXH3 hashing, consistent with indexer.rs pub const XXH3_SEED: u64 = 1337; -const CACHE_NAMESPACE_HASH_SEED: u64 = XXH3_SEED ^ 0xc2b2_ae3d_27d4_eb4f; /// Compute the hash of a local block. pub fn compute_block_hash(data: &[u8]) -> LocalBlockHash { @@ -37,18 +36,8 @@ pub struct BlockHashOptions<'a> { } fn block_hash_seed(options: BlockHashOptions<'_>) -> u64 { - let mut seed = XXH3_SEED; - if let Some(name) = options.lora_name.filter(|n| !n.is_empty()) { - // Preserve the established LoRA formula shared with dynamo-kv-hashing. - seed = seed.wrapping_add(xxh3::xxh3_64(name.as_bytes())); - } - if let Some(namespace) = options.cache_namespace.filter(|n| !n.is_empty()) { - seed = seed.wrapping_add(xxh3::xxh3_64_with_seed( - namespace.as_bytes(), - CACHE_NAMESPACE_HASH_SEED, - )); - } - seed + dynamo_kv_hashing::compute_salt_hash(options.cache_namespace, options.lora_name) + .expect("string salt derivation is infallible") } #[inline] @@ -1382,6 +1371,31 @@ mod tests { ); } + #[test] + fn test_cache_namespace_hash_matches_kv_hashing_contract() { + let tokens: Vec = (0..4).collect(); + let cache_namespace = "tenant-a"; + let lora_name = "adapter-a"; + let actual = compute_block_hash_for_seq( + &tokens, + 4, + BlockHashOptions { + lora_name: Some(lora_name), + cache_namespace: Some(cache_namespace), + ..Default::default() + }, + ); + let token_bytes = tokens + .iter() + .flat_map(|token| token.to_le_bytes()) + .collect::>(); + let salt_hash = + dynamo_kv_hashing::compute_salt_hash(Some(cache_namespace), Some(lora_name)).unwrap(); + let expected = LocalBlockHash(dynamo_kv_hashing::compute_hash_v2(&token_bytes, salt_hash)); + + assert_eq!(actual, vec![expected]); + } + #[test] fn test_cache_namespace_empty_string_normalized_to_none() { let tokens: Vec = (0..4).collect(); diff --git a/lib/kv-router/src/zmq_wire/mod.rs b/lib/kv-router/src/zmq_wire/mod.rs index 1f0af44df4cf..47c23ca163f7 100644 --- a/lib/kv-router/src/zmq_wire/mod.rs +++ b/lib/kv-router/src/zmq_wire/mod.rs @@ -52,7 +52,7 @@ pub struct ZmqEventNormalizer { #[derive(Debug, Clone, PartialEq, Eq)] enum CacheNamespaceState { - Namespaced(String), + Namespaced(Arc), Ambiguous, } @@ -194,12 +194,26 @@ impl ZmqEventNormalizer { if cache_namespace.as_deref() == Some("") { *cache_namespace = None; } - if cache_namespace.is_none() - && let Some(parent) = parent_block_hash.as_ref() - { + let namespace = if let Some(namespace) = cache_namespace.as_deref() { + parent_block_hash + .as_ref() + .and_then(|parent| { + self.cache_namespaces.get(&(worker, (*parent).into_u64())) + }) + .and_then(|state| match state { + CacheNamespaceState::Namespaced(parent_namespace) + if parent_namespace.as_ref() == namespace => + { + Some(Arc::clone(parent_namespace)) + } + _ => None, + }) + .or_else(|| Some(Arc::from(namespace))) + } else if let Some(parent) = parent_block_hash.as_ref() { match self.cache_namespaces.get(&(worker, (*parent).into_u64())) { Some(CacheNamespaceState::Namespaced(namespace)) => { - *cache_namespace = Some(namespace.clone()); + *cache_namespace = Some(namespace.to_string()); + Some(Arc::clone(namespace)) } Some(CacheNamespaceState::Ambiguous) => { return Err(ZmqEventFilterReason::AmbiguousCacheNamespace); @@ -212,15 +226,15 @@ impl ZmqEventNormalizer { // here would duplicate the index on the event hot path. The // backend still enforces its own cache isolation; this narrow // fail-open case can only pollute the router overlap score. + None } } - } + } else { + None + }; - if let Some(namespace) = cache_namespace - .as_ref() - .filter(|namespace| !namespace.is_empty()) - { - let state = CacheNamespaceState::Namespaced(namespace.clone()); + if let Some(namespace) = namespace { + let state = CacheNamespaceState::Namespaced(namespace); for block_hash in block_hashes.iter() { self.cache_namespaces .entry((worker, (*block_hash).into_u64())) diff --git a/lib/kv-router/src/zmq_wire/tests.rs b/lib/kv-router/src/zmq_wire/tests.rs index 8fcba7929550..4d6128ce295a 100644 --- a/lib/kv-router/src/zmq_wire/tests.rs +++ b/lib/kv-router/src/zmq_wire/tests.rs @@ -566,6 +566,18 @@ fn test_normalizer_propagates_cache_namespace_from_parent() { assert!(normalizer.preprocess(parent, worker).is_some()); let child = normalizer.preprocess(child, worker).unwrap(); + let CacheNamespaceState::Namespaced(parent_namespace) = + &normalizer.cache_namespaces[&(worker, 1)] + else { + panic!("expected namespaced parent"); + }; + let CacheNamespaceState::Namespaced(child_namespace) = + &normalizer.cache_namespaces[&(worker, 2)] + else { + panic!("expected namespaced child"); + }; + assert!(Arc::ptr_eq(parent_namespace, child_namespace)); + let RawKvEvent::BlockStored { cache_namespace, .. } = child @@ -575,6 +587,36 @@ fn test_normalizer_propagates_cache_namespace_from_parent() { assert_eq!(cache_namespace.as_deref(), Some("tenant-a")); } +#[test] +fn test_normalizer_shares_cache_namespace_across_blocks() { + let worker = WorkerWithDpRank::new(7, 0); + let mut normalizer = ZmqEventNormalizer::new(2); + let event = RawKvEvent::BlockStored { + block_hashes: vec![BlockHashValue::Unsigned(1), BlockHashValue::Unsigned(2)], + parent_block_hash: None, + token_ids: vec![10, 11, 12, 13], + block_size: 2, + medium: None, + lora_name: None, + cache_namespace: Some("tenant-a".to_string()), + block_mm_infos: None, + is_eagle: Some(false), + group_idx: None, + kv_cache_spec_kind: None, + kv_cache_spec_sliding_window: None, + }; + + assert!(normalizer.preprocess(event, worker).is_some()); + + let CacheNamespaceState::Namespaced(first) = &normalizer.cache_namespaces[&(worker, 1)] else { + panic!("expected first block namespace"); + }; + let CacheNamespaceState::Namespaced(second) = &normalizer.cache_namespaces[&(worker, 2)] else { + panic!("expected second block namespace"); + }; + assert!(Arc::ptr_eq(first, second)); +} + #[test] fn test_normalizer_rejects_ambiguous_parent_cache_namespace() { let worker = WorkerWithDpRank::new(7, 0); diff --git a/lib/kvbm-consolidator/src/egress/zmq_publisher.rs b/lib/kvbm-consolidator/src/egress/zmq_publisher.rs index 06a4b74725db..2821b090a484 100644 --- a/lib/kvbm-consolidator/src/egress/zmq_publisher.rs +++ b/lib/kvbm-consolidator/src/egress/zmq_publisher.rs @@ -95,6 +95,7 @@ fn consolidated_to_event(ev: ConsolidatedEvent) -> anyhow::Result { token_ids, block_size, lora_name, + cache_namespace, source: _, } => { let token_ids_i32: Vec = token_ids @@ -116,6 +117,7 @@ fn consolidated_to_event(ev: ConsolidatedEvent) -> anyhow::Result { token_ids: token_ids_i32, block_size: block_size_i32, lora_name, + cache_namespace, medium: None, }) } @@ -178,7 +180,11 @@ pub async fn spawn( let batch = EventBatch(now_f64, events, Some(0)); let mut buf = Vec::new(); - if let Err(e) = batch.serialize(&mut rmp_serde::Serializer::new(&mut buf)) { + // Store events have optional named fields such as cache_salt. Encode + // structs as maps so omitted fields cannot shift positional values. + let mut serializer = + rmp_serde::Serializer::new(&mut buf).with_struct_map(); + if let Err(e) = batch.serialize(&mut serializer) { tracing::warn!("Failed to publish batch: {e}"); continue; } @@ -210,6 +216,7 @@ mod sort_tests { token_ids: vec![], block_size: 0, lora_name: None, + cache_namespace: None, source: EventSource::Vllm, } } diff --git a/lib/kvbm-consolidator/src/ingress/zmq_subscriber.rs b/lib/kvbm-consolidator/src/ingress/zmq_subscriber.rs index c3b837396dbd..82e0921dfb43 100644 --- a/lib/kvbm-consolidator/src/ingress/zmq_subscriber.rs +++ b/lib/kvbm-consolidator/src/ingress/zmq_subscriber.rs @@ -96,6 +96,7 @@ fn process_event(tracker: &mut Tracker, event: RawKvEvent, engine_source: EventS token_ids, block_size, lora_name, + cache_namespace, is_eagle, block_mm_infos, .. @@ -149,16 +150,20 @@ fn process_event(tracker: &mut Tracker, event: RawKvEvent, engine_source: EventS } let mut current_parent = parent_block_hash.map(|h| h.into_u64().to_string()); + let cache_namespace = cache_namespace + .filter(|namespace| !namespace.is_empty()) + .map(Arc::::from); for (i, block_hash) in block_hashes.into_iter().enumerate() { let hash_str = block_hash.into_u64().to_string(); - tracker.handle_store( + tracker.handle_store_with_cache_namespace( engine_source, hash_str.clone(), current_parent.clone(), token_chunks[i].clone(), block_size, lora_name.clone(), + cache_namespace.clone(), ); current_parent = Some(hash_str); } diff --git a/lib/kvbm-consolidator/src/tracker.rs b/lib/kvbm-consolidator/src/tracker.rs index 68c1b4fb9077..ae01b1f6c710 100644 --- a/lib/kvbm-consolidator/src/tracker.rs +++ b/lib/kvbm-consolidator/src/tracker.rs @@ -22,6 +22,7 @@ //! PLH already computed by the upstream registry. use std::collections::{HashMap, HashSet, VecDeque}; +use std::sync::Arc; use dynamo_kv_hashing::Request; use dynamo_tokens::PositionalLineageHash; @@ -38,6 +39,7 @@ pub enum ConsolidatedEvent { token_ids: Vec, block_size: usize, lora_name: Option, + cache_namespace: Option>, source: EventSource, }, Remove { @@ -59,6 +61,7 @@ pub struct BlockState { pub sources: HashSet, pub registry_handle: Option, pub published: bool, + pub cache_namespace: Option>, } impl std::fmt::Debug for BlockState { @@ -67,6 +70,7 @@ impl std::fmt::Debug for BlockState { .field("sources", &self.sources) .field("registry_handle", &self.registry_handle.is_some()) .field("published", &self.published) + .field("has_cache_namespace", &self.cache_namespace.is_some()) .finish() } } @@ -101,10 +105,12 @@ impl Tracker { token_ids: &[u32], block_size: usize, lora_name: Option<&str>, + cache_namespace: Option<&str>, ) -> Option { let request = Request::builder() .tokens(token_ids.to_vec()) .lora_name(lora_name.map(str::to_string)) + .salt(cache_namespace.map(str::to_string)) .build() .ok()?; let blocks = request.into_blocks(block_size as u32).ok()?; @@ -128,14 +134,39 @@ impl Tracker { block_size: usize, lora_name: Option, ) -> bool { - let parent_plh = match parent_external_hash { + self.handle_store_with_cache_namespace( + source, + external_hash, + parent_external_hash, + token_ids, + block_size, + lora_name, + None, + ) + } + + /// Handle a STORE event with an optional cache namespace. + pub fn handle_store_with_cache_namespace( + &mut self, + source: EventSource, + external_hash: String, + parent_external_hash: Option, + token_ids: Vec, + block_size: usize, + lora_name: Option, + cache_namespace: Option>, + ) -> bool { + let parent_key = parent_external_hash + .as_ref() + .map(|external_hash| (source, external_hash.clone())); + let parent_plh = match parent_key.as_ref() { None => None, - Some(ref peh) => match self.external_to_seq.get(&(source, peh.clone())) { + Some(key) => match self.external_to_seq.get(key) { Some(&p) => Some(p), None => { tracing::warn!( "Unresolved parent external hash {:?} for source {:?}; treating as root", - peh, + key.1, source ); None @@ -143,8 +174,22 @@ impl Tracker { }, }; - let plh = match Self::compute_plh(parent_plh, &token_ids, block_size, lora_name.as_deref()) - { + let cache_namespace = cache_namespace + .filter(|namespace| !namespace.is_empty()) + .or_else(|| { + parent_plh + .and_then(|parent_plh| self.blocks.get(&parent_plh)) + .and_then(|state| state.cache_namespace.as_ref()) + .cloned() + }); + + let plh = match Self::compute_plh( + parent_plh, + &token_ids, + block_size, + lora_name.as_deref(), + cache_namespace.as_deref(), + ) { Some(h) => h, None => { tracing::warn!( @@ -163,10 +208,14 @@ impl Tracker { return false; } + self.external_to_seq.insert((source, external_hash), plh); + match self.blocks.get_mut(&plh) { Some(state) => { state.sources.insert(source); - self.external_to_seq.insert((source, external_hash), plh); + if state.cache_namespace.is_none() { + state.cache_namespace = cache_namespace.clone(); + } // A prior source registered the block without publishable metadata // (KVBM-bridge create with empty tokens / block_size 0). This is the // first real source — publish now. @@ -177,6 +226,7 @@ impl Tracker { token_ids, block_size, lora_name, + cache_namespace, source, }); return true; @@ -196,14 +246,15 @@ impl Tracker { sources, registry_handle, published: true, + cache_namespace: cache_namespace.clone(), }, ); - self.external_to_seq.insert((source, external_hash), plh); self.event_queue.push_back(ConsolidatedEvent::Store { seq_hash: plh, token_ids, block_size, lora_name, + cache_namespace, source, }); true @@ -228,7 +279,6 @@ impl Tracker { return false; } }; - let (empty, published) = match self.blocks.get_mut(&plh) { Some(state) => { state.sources.remove(&source); @@ -277,6 +327,7 @@ impl Tracker { token_ids, block_size, lora_name, + cache_namespace: None, source: EventSource::Kvbm, }); return true; @@ -296,6 +347,7 @@ impl Tracker { sources, registry_handle, published: publishable, + cache_namespace: None, }, ); if publishable { @@ -304,6 +356,7 @@ impl Tracker { token_ids, block_size, lora_name, + cache_namespace: None, source: EventSource::Kvbm, }); true @@ -790,6 +843,103 @@ mod tests { other => panic!("expected Store, got {:?}", other), } } + + #[test] + fn tracker_cache_namespace_isolates_identical_tokens() { + let mut t = tracker(); + let tokens = vec![10, 20, 30, 40]; + + t.handle_store_with_cache_namespace( + EventSource::Vllm, + "tenant-a-block".into(), + None, + tokens.clone(), + 4, + None, + Some(Arc::from("tenant-a")), + ); + t.handle_store_with_cache_namespace( + EventSource::Vllm, + "tenant-b-block".into(), + None, + tokens.clone(), + 4, + None, + Some(Arc::from("tenant-b")), + ); + + let events = t.drain_events(); + assert_eq!(events.len(), 2); + let stores = events + .iter() + .map(|event| match event { + ConsolidatedEvent::Store { + seq_hash, + cache_namespace, + .. + } => (*seq_hash, cache_namespace.as_deref()), + other => panic!("expected Store, got {other:?}"), + }) + .collect::>(); + assert_ne!(stores[0].0, stores[1].0); + assert_eq!(stores[0].1, Some("tenant-a")); + assert_eq!(stores[1].1, Some("tenant-b")); + + for (namespace, actual) in [("tenant-a", stores[0].0), ("tenant-b", stores[1].0)] { + let expected = Request::builder() + .tokens(tokens.clone()) + .salt(Some(namespace.to_string())) + .build() + .unwrap() + .into_blocks(4) + .unwrap()[0] + .plh; + assert_eq!(actual, expected); + } + } + + #[test] + fn tracker_inherits_cache_namespace_from_parent() { + let mut t = tracker(); + let namespace = Arc::::from("tenant-a"); + + t.handle_store_with_cache_namespace( + EventSource::Vllm, + "parent".into(), + None, + vec![1, 2, 3, 4], + 4, + None, + Some(Arc::clone(&namespace)), + ); + t.handle_store_with_cache_namespace( + EventSource::Vllm, + "child".into(), + Some("parent".into()), + vec![5, 6, 7, 8], + 4, + None, + None, + ); + + let events = t.drain_events(); + let namespaces = events + .iter() + .map(|event| match event { + ConsolidatedEvent::Store { + cache_namespace: Some(cache_namespace), + .. + } => cache_namespace, + other => panic!("expected namespaced Store, got {other:?}"), + }) + .collect::>(); + assert!(Arc::ptr_eq(namespaces[0], namespaces[1])); + assert!(Arc::ptr_eq(namespaces[0], &namespace)); + + t.handle_remove(EventSource::Vllm, "child"); + t.handle_remove(EventSource::Vllm, "parent"); + assert!(t.blocks.is_empty()); + } } #[cfg(test)] diff --git a/lib/kvbm-consolidator/src/wire/router_out.rs b/lib/kvbm-consolidator/src/wire/router_out.rs index 14db63368e8e..dcdea404ee2d 100644 --- a/lib/kvbm-consolidator/src/wire/router_out.rs +++ b/lib/kvbm-consolidator/src/wire/router_out.rs @@ -7,6 +7,7 @@ //! (matching vLLM's `msgspec(array_like=True)` envelope). use serde::Serialize; +use std::sync::Arc; /// Batch envelope: `(timestamp, events, data_parallel_rank)`. #[derive(Debug, Serialize)] @@ -25,6 +26,12 @@ pub enum Event { block_size: i32, #[serde(default, skip_serializing_if = "Option::is_none")] lora_name: Option, + #[serde( + default, + rename = "cache_salt", + skip_serializing_if = "Option::is_none" + )] + cache_namespace: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] medium: Option, }, diff --git a/lib/kvbm-consolidator/tests/common/mod.rs b/lib/kvbm-consolidator/tests/common/mod.rs index 8607278a2bb7..db58192b74f3 100644 --- a/lib/kvbm-consolidator/tests/common/mod.rs +++ b/lib/kvbm-consolidator/tests/common/mod.rs @@ -57,6 +57,8 @@ pub enum EventMirror { block_size: i32, #[serde(default)] lora_name: Option, + #[serde(default, rename = "cache_salt")] + cache_namespace: Option, #[serde(default)] medium: Option, }, diff --git a/lib/kvbm-consolidator/tests/e2e.rs b/lib/kvbm-consolidator/tests/e2e.rs index 79bf827b870c..3c07465e8200 100644 --- a/lib/kvbm-consolidator/tests/e2e.rs +++ b/lib/kvbm-consolidator/tests/e2e.rs @@ -34,6 +34,12 @@ enum SnapEvent { block_size: i32, #[serde(default, skip_serializing_if = "Option::is_none")] lora_name: Option, + #[serde( + default, + rename = "cache_salt", + skip_serializing_if = "Option::is_none" + )] + cache_namespace: Option, #[serde(default, skip_serializing_if = "Option::is_none")] medium: Option, }, @@ -168,6 +174,7 @@ fn convert_event(e: common::EventMirror) -> SnapEvent { token_ids, block_size, lora_name, + cache_namespace, medium, } => SnapEvent::BlockStored { block_hashes, @@ -175,6 +182,7 @@ fn convert_event(e: common::EventMirror) -> SnapEvent { token_ids, block_size, lora_name, + cache_namespace, medium, }, common::EventMirror::BlockRemoved { diff --git a/lib/kvbm-consolidator/tests/zmq_ingress.rs b/lib/kvbm-consolidator/tests/zmq_ingress.rs index ca1549e486a6..9f6885bb0d66 100644 --- a/lib/kvbm-consolidator/tests/zmq_ingress.rs +++ b/lib/kvbm-consolidator/tests/zmq_ingress.rs @@ -10,7 +10,7 @@ mod common; use std::time::Duration; use common::{EventMirror, TestBatch, ZmqPubHandle, ZmqSubHandle, init_tracing, sync_pulse}; -use kvbm_consolidator::wire::vllm_in::{BlockHashValue, RawKvEvent}; +use kvbm_consolidator::wire::vllm_in::{BlockHashValue, KvEventBatch, RawKvEvent}; use kvbm_consolidator::{ConsolidatorBuilder, EventSource}; // ─── helpers ───────────────────────────────────────────────────────────────── @@ -21,6 +21,17 @@ fn bs_event( tokens: Vec, block_size: usize, lora_name: Option, +) -> RawKvEvent { + bs_event_with_cache_namespace(block_hashes, parent, tokens, block_size, lora_name, None) +} + +fn bs_event_with_cache_namespace( + block_hashes: Vec, + parent: Option, + tokens: Vec, + block_size: usize, + lora_name: Option, + cache_namespace: Option, ) -> RawKvEvent { RawKvEvent::BlockStored { block_hashes: block_hashes @@ -32,7 +43,7 @@ fn bs_event( block_size, lora_name, medium: None, - cache_namespace: None, + cache_namespace, block_mm_infos: None, is_eagle: None, group_idx: None, @@ -41,6 +52,104 @@ fn bs_event( } } +#[tokio::test] +async fn zmq_cache_namespaces_remain_isolated() { + init_tracing(); + + tokio::time::timeout(Duration::from_secs(5), async { + let pub_handle = ZmqPubHandle::spawn().await; + let egress_port = common::pick_port(); + let egress_ep = common::make_endpoint(egress_port); + + let consolidator = ConsolidatorBuilder::new(&egress_ep, EventSource::Vllm) + .zmq_in(&pub_handle.endpoint) + .poll_interval(Duration::from_millis(20)) + .build() + .await + .expect("build consolidator"); + let mut sub = ZmqSubHandle::spawn(&egress_ep).await.expect("spawn sub"); + assert!( + sync_pulse(&pub_handle, &mut sub, Duration::from_secs(4)).await, + "sync_pulse timed out" + ); + + let tokens = vec![1, 2, 3, 4]; + let batch = TestBatch( + 1.0, + vec![ + bs_event_with_cache_namespace( + vec![101], + None, + tokens.clone(), + 4, + None, + Some("tenant-a".to_string()), + ), + bs_event_with_cache_namespace( + vec![202], + None, + tokens, + 4, + None, + Some("tenant-b".to_string()), + ), + ], + None, + ); + let payload = rmp_serde::to_vec_named(&batch).expect("encode named batch"); + let decoded: KvEventBatch = rmp_serde::from_slice(&payload).expect("decode named batch"); + let decoded_namespaces = decoded + .events + .iter() + .map(|event| match event { + RawKvEvent::BlockStored { + cache_namespace, .. + } => cache_namespace.as_deref(), + other => panic!("expected BlockStored, got {other:?}"), + }) + .collect::>(); + assert_eq!(decoded_namespaces, [Some("tenant-a"), Some("tenant-b")]); + pub_handle + .send_frames(vec![vec![], vec![0u8; 8], payload]) + .await + .expect("send_batch"); + + let msgs = sub.recv_n(2, Duration::from_secs(3)).await.expect("recv_n"); + let stores = msgs + .iter() + .flat_map(|(_, batch)| batch.1.iter()) + .filter_map(|event| match event { + EventMirror::BlockStored { + block_hashes, + lora_name, + cache_namespace, + .. + } => Some(( + block_hashes[0], + lora_name.as_deref(), + cache_namespace.as_deref(), + )), + _ => None, + }) + .collect::>(); + + assert_eq!( + stores.len(), + 2, + "expected one store per namespace: {stores:?}" + ); + assert_ne!(stores[0].0, stores[1].0); + assert_eq!(stores[0].1, None); + assert_eq!(stores[1].1, None); + assert_eq!(stores[0].2, Some("tenant-a")); + assert_eq!(stores[1].2, Some("tenant-b")); + + consolidator.shutdown().await; + }) + .await + .expect("timed out"); +} + // ─── test 1: zmq_ingress_roundtrip ─────────────────────────────────────────── /// A single vLLM batch with 3 chained blocks propagates to egress as 3 STOREs. From 9cd076d7766ae637700dd9b0d7b477da0dba5f45 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Mon, 6 Jul 2026 14:29:14 -0700 Subject: [PATCH 09/10] fix(ci): unblock pre-merge clippy checks Signed-off-by: jthomson04 --- lib/bindings/kvbm/Cargo.lock | 1 + lib/bindings/python/Cargo.lock | 1 + .../src/ingress/zmq_subscriber.rs | 6 +- lib/kvbm-consolidator/src/tracker.rs | 73 ++++++++++++++----- 4 files changed, 58 insertions(+), 23 deletions(-) diff --git a/lib/bindings/kvbm/Cargo.lock b/lib/bindings/kvbm/Cargo.lock index e643664d6139..86ac737dd2eb 100644 --- a/lib/bindings/kvbm/Cargo.lock +++ b/lib/bindings/kvbm/Cargo.lock @@ -1543,6 +1543,7 @@ dependencies = [ "async-trait", "dashmap", "derive_builder", + "dynamo-kv-hashing", "dynamo-runtime", "dynamo-tokens", "flume", diff --git a/lib/bindings/python/Cargo.lock b/lib/bindings/python/Cargo.lock index 5d998a46a00c..8a1a6a560b13 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -2142,6 +2142,7 @@ dependencies = [ "chrono", "dashmap", "derive_builder", + "dynamo-kv-hashing", "dynamo-runtime", "dynamo-tokens", "flume", diff --git a/lib/kvbm-consolidator/src/ingress/zmq_subscriber.rs b/lib/kvbm-consolidator/src/ingress/zmq_subscriber.rs index 82e0921dfb43..3e897059ae79 100644 --- a/lib/kvbm-consolidator/src/ingress/zmq_subscriber.rs +++ b/lib/kvbm-consolidator/src/ingress/zmq_subscriber.rs @@ -24,7 +24,7 @@ use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; use crate::source::EventSource; -use crate::tracker::Tracker; +use crate::tracker::{StoreInput, Tracker}; use crate::wire::vllm_in::{KvEventBatch, RawKvEvent}; use crate::zmq_util::{connect_sub_socket, multipart_message}; @@ -156,7 +156,7 @@ fn process_event(tracker: &mut Tracker, event: RawKvEvent, engine_source: EventS for (i, block_hash) in block_hashes.into_iter().enumerate() { let hash_str = block_hash.into_u64().to_string(); - tracker.handle_store_with_cache_namespace( + tracker.handle_store_input(StoreInput::new( engine_source, hash_str.clone(), current_parent.clone(), @@ -164,7 +164,7 @@ fn process_event(tracker: &mut Tracker, event: RawKvEvent, engine_source: EventS block_size, lora_name.clone(), cache_namespace.clone(), - ); + )); current_parent = Some(hash_str); } } diff --git a/lib/kvbm-consolidator/src/tracker.rs b/lib/kvbm-consolidator/src/tracker.rs index ae01b1f6c710..e4854e25a129 100644 --- a/lib/kvbm-consolidator/src/tracker.rs +++ b/lib/kvbm-consolidator/src/tracker.rs @@ -49,6 +49,39 @@ pub enum ConsolidatedEvent { ClearAll, } +/// Inputs for a STORE event received from a string-hashed source. +pub(crate) struct StoreInput { + source: EventSource, + external_hash: String, + parent_external_hash: Option, + token_ids: Vec, + block_size: usize, + lora_name: Option, + cache_namespace: Option>, +} + +impl StoreInput { + pub(crate) fn new( + source: EventSource, + external_hash: String, + parent_external_hash: Option, + token_ids: Vec, + block_size: usize, + lora_name: Option, + cache_namespace: Option>, + ) -> Self { + Self { + source, + external_hash, + parent_external_hash, + token_ids, + block_size, + lora_name, + cache_namespace, + } + } +} + /// Per-block state: which sources have it, an optional registry handle keeping the /// block present in kvbm-logical's shared radix tree, and whether a publishable /// `ConsolidatedEvent::Store` has been emitted downstream. @@ -134,7 +167,7 @@ impl Tracker { block_size: usize, lora_name: Option, ) -> bool { - self.handle_store_with_cache_namespace( + self.handle_store_input(StoreInput::new( source, external_hash, parent_external_hash, @@ -142,20 +175,20 @@ impl Tracker { block_size, lora_name, None, - ) + )) } /// Handle a STORE event with an optional cache namespace. - pub fn handle_store_with_cache_namespace( - &mut self, - source: EventSource, - external_hash: String, - parent_external_hash: Option, - token_ids: Vec, - block_size: usize, - lora_name: Option, - cache_namespace: Option>, - ) -> bool { + pub(crate) fn handle_store_input(&mut self, input: StoreInput) -> bool { + let StoreInput { + source, + external_hash, + parent_external_hash, + token_ids, + block_size, + lora_name, + cache_namespace, + } = input; let parent_key = parent_external_hash .as_ref() .map(|external_hash| (source, external_hash.clone())); @@ -849,7 +882,7 @@ mod tests { let mut t = tracker(); let tokens = vec![10, 20, 30, 40]; - t.handle_store_with_cache_namespace( + t.handle_store_input(StoreInput::new( EventSource::Vllm, "tenant-a-block".into(), None, @@ -857,8 +890,8 @@ mod tests { 4, None, Some(Arc::from("tenant-a")), - ); - t.handle_store_with_cache_namespace( + )); + t.handle_store_input(StoreInput::new( EventSource::Vllm, "tenant-b-block".into(), None, @@ -866,7 +899,7 @@ mod tests { 4, None, Some(Arc::from("tenant-b")), - ); + )); let events = t.drain_events(); assert_eq!(events.len(), 2); @@ -903,7 +936,7 @@ mod tests { let mut t = tracker(); let namespace = Arc::::from("tenant-a"); - t.handle_store_with_cache_namespace( + t.handle_store_input(StoreInput::new( EventSource::Vllm, "parent".into(), None, @@ -911,8 +944,8 @@ mod tests { 4, None, Some(Arc::clone(&namespace)), - ); - t.handle_store_with_cache_namespace( + )); + t.handle_store_input(StoreInput::new( EventSource::Vllm, "child".into(), Some("parent".into()), @@ -920,7 +953,7 @@ mod tests { 4, None, None, - ); + )); let events = t.drain_events(); let namespaces = events From 76ec167f4290fbd7d582d4208981c173c238a0f5 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Mon, 6 Jul 2026 20:27:56 -0700 Subject: [PATCH 10/10] test(vllm): expect tagged cache salt Signed-off-by: jthomson04 --- .../src/dynamo/vllm/tests/test_vllm_delta_streaming.py | 7 +++++++ components/src/dynamo/vllm/tests/test_vllm_unit.py | 2 +- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/components/src/dynamo/vllm/tests/test_vllm_delta_streaming.py b/components/src/dynamo/vllm/tests/test_vllm_delta_streaming.py index 923d7d93bbe9..813e7d85c92e 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_delta_streaming.py +++ b/components/src/dynamo/vllm/tests/test_vllm_delta_streaming.py @@ -295,6 +295,9 @@ async def test_unified_llm_engine_passes_delta_chunks_and_counts_usage(): async def test_unified_llm_engine_forwards_cache_salt_to_prompt(): pytest.importorskip("vllm.usage.usage_lib") from dynamo.vllm.llm_engine import VllmLLMEngine + from dynamo.vllm.multimodal_utils.request_processor import ( + VllmMultimodalRequestProcessor, + ) engine = VllmLLMEngine.__new__(VllmLLMEngine) engine.engine_client = _FakeEngineClient([]) @@ -302,6 +305,10 @@ async def test_unified_llm_engine_forwards_cache_salt_to_prompt(): engine._model_max_len = None engine.disaggregation_mode = DisaggregationMode.AGGREGATED engine.enable_rl = False + engine._multimodal_request_processor = VllmMultimodalRequestProcessor( + model="test-model", + enable_multimodal=False, + ) engine._dp_range = None request = { diff --git a/components/src/dynamo/vllm/tests/test_vllm_unit.py b/components/src/dynamo/vllm/tests/test_vllm_unit.py index d14025c31c07..2901232ced0e 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_unit.py +++ b/components/src/dynamo/vllm/tests/test_vllm_unit.py @@ -1503,4 +1503,4 @@ async def abort_monitor(*args, **kwargs): ] assert chunks - assert captured["prompt"]["cache_salt"] == "tenant-a" + assert captured["prompt"]["cache_salt"] == "dynamo-cache-salt:tenant-a"