feat: add cache salt support to KV cache events (#30827)

Signed-off-by: jthomson04 <jwillthomson19@gmail.com>

test/registered/unit/disaggregation/test_encode_receiver.py

@@ -0,0 +1,61 @@
+"""Unit tests for request construction in the encode-disaggregation path."""
+
+import unittest
+from array import array
+from types import SimpleNamespace
+
+from sglang.srt.disaggregation.encode_receiver import MMReceiverBase
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.sampling.sampling_params import SamplingParams
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+class TestEncodeReceiverRequestConstruction(CustomTestCase):
+    def test_extra_key_and_cache_salt_are_forwarded(self):
+        scheduler = SimpleNamespace(
+            model_config=SimpleNamespace(hf_eos_token_id={2}, vocab_size=128),
+            disaggregation_mode=DisaggregationMode.NULL,
+            metrics_reporter=SimpleNamespace(enable_metrics=False),
+            metrics_collector=None,
+            dllm_config=None,
+            tokenizer=object(),
+        )
+        receiver = SimpleNamespace(scheduler=scheduler)
+        recv_req = SimpleNamespace(
+            rid="request-1",
+            input_text="hello",
+            input_ids=array("q", [1, 2]),
+            sampling_params=SamplingParams(max_new_tokens=1),
+            return_logprob=False,
+            top_logprobs_num=0,
+            token_ids_logprob=None,
+            stream=False,
+            lora_id=None,
+            input_embeds=None,
+            custom_logit_processor=None,
+            require_reasoning=False,
+            return_hidden_states=False,
+            return_routed_experts=False,
+            routed_experts_start_len=0,
+            bootstrap_host=None,
+            bootstrap_port=None,
+            bootstrap_room=None,
+            routed_dp_rank=None,
+            disagg_prefill_dp_rank=None,
+            priority=None,
+            extra_key="classification",
+            cache_salt="tenant-a",
+            http_worker_ipc=None,
+        )
+
+        req = MMReceiverBase.create_req(receiver, recv_req)
+
+        self.assertEqual(req.extra_key, "classification")
+        self.assertEqual(req.cache_salt, "tenant-a")
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/disaggregation/test_kv_events.py

@@ -9,7 +9,14 @@ the router can subscribe per replica (the `dp_size` it reads from
 
 import unittest
 
+import msgspec
+
 from sglang.srt.disaggregation.kv_events import (
+    BlockStored,
+    BlockStoredMetadata,
+    BlockStoredWithMetadata,
+    KVEventBatch,
+    StorageMedium,
     ZmqEventPublisher,
     select_kv_publisher_dp_rank,
 )
@@ -94,5 +101,44 @@ class TestSelectKvPublisherDpRank(CustomTestCase):
                 self.assertEqual(len(ranks), dp_size)
 
 
+class TestBlockStoredWireFormat(CustomTestCase):
+    def _event(self, metadata=None):
+        event_type = BlockStored if metadata is None else BlockStoredWithMetadata
+        kwargs = dict(
+            block_hashes=[123],
+            parent_block_hash=None,
+            token_ids=[1, 2],
+            block_size=2,
+            lora_id=None,
+            medium=StorageMedium.GPU,
+        )
+        if metadata is not None:
+            kwargs["metadata"] = metadata
+        return event_type(**kwargs)
+
+    def test_unsalted_event_keeps_legacy_array_shape(self):
+        decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(self._event()))
+        self.assertEqual(len(decoded), 7)
+
+    def test_salted_event_appends_typed_metadata(self):
+        event = self._event(BlockStoredMetadata(cache_salt="tenant-a"))
+        encoded = msgspec.msgpack.encode(event)
+        decoded = msgspec.msgpack.decode(encoded)
+        round_tripped = msgspec.msgpack.decode(encoded, type=BlockStoredWithMetadata)
+        self.assertEqual(len(decoded), 8)
+        self.assertEqual(decoded[7], {"cache_salt": "tenant-a"})
+        self.assertEqual(round_tripped.metadata.cache_salt, "tenant-a")
+
+    def test_salted_event_remains_compatible_with_typed_batch_consumers(self):
+        batch = KVEventBatch(
+            ts=1.0,
+            events=[self._event(BlockStoredMetadata(cache_salt="tenant-a"))],
+        )
+        round_tripped = msgspec.msgpack.decode(
+            msgspec.msgpack.encode(batch), type=KVEventBatch
+        )
+        self.assertEqual(round_tripped.events[0].block_hashes, [123])
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/entrypoints/openai/test_serving_chat.py

@@ -316,6 +316,8 @@ class ServingChatTestCase(unittest.TestCase):
             input_ids=[101, 102, 103],
             stop=["STOP"],
             return_prompt_token_ids=True,
+            cache_salt="tenant-a",
+            extra_key="classification",
         )
 
         with patch(
@@ -329,6 +331,8 @@ class ServingChatTestCase(unittest.TestCase):
         self.assertEqual(adapted.input_ids, [101, 102, 103])
         self.assertTrue(adapted.return_prompt_token_ids)
         self.assertEqual(adapted.sampling_params["stop"], ["STOP"])
+        self.assertEqual(adapted.cache_salt, "tenant-a")
+        self.assertEqual(adapted.extra_key, "classification")
         conv_mock.assert_not_called()
 
     def test_kimi_k3_usage_excludes_assistant_generation_stub(self):

test/registered/unit/entrypoints/openai/test_serving_completions.py

@@ -65,6 +65,29 @@ class ServingCompletionTestCase(unittest.TestCase):
         internal, _ = self.sc._convert_to_internal_request(req)
         self.assertEqual(internal.input_ids, [1, 2, 3, 4])
 
+    def test_cache_salt_and_extra_key_remain_distinct(self):
+        req = CompletionRequest(
+            model="x",
+            prompt=[1, 2, 3, 4],
+            max_tokens=1,
+            cache_salt="tenant-a",
+            extra_key="classification",
+        )
+        internal, _ = self.sc._convert_to_internal_request(req)
+        self.assertEqual(internal.cache_salt, "tenant-a")
+        self.assertEqual(internal.extra_key, "classification")
+
+    def test_single_request_rejects_batched_cache_salt(self):
+        req = CompletionRequest(
+            model="x",
+            prompt=[1, 2, 3, 4],
+            max_tokens=1,
+            cache_salt=["tenant-a"],
+        )
+        internal, _ = self.sc._convert_to_internal_request(req)
+        with self.assertRaisesRegex(ValueError, "single request"):
+            internal.normalize_batch_and_arguments()
+
     # ---------- echo-handling ----------
     def test_echo_with_list_of_strings_streaming(self):
         req = CompletionRequest(

test/registered/unit/managers/test_io_struct.py

@@ -519,6 +519,53 @@ class TestGenerateReqInputNormalization(CustomTestCase):
         req.normalize_batch_and_arguments()
         self.assertEqual(req.extra_key, "solo")
 
+    def test_cache_salt_normalization(self):
+        req = GenerateReqInput(
+            text=["Hello", "World"],
+            cache_salt=["tenant-A", ""],
+            sampling_params=[{}, {}],
+        )
+        req.normalize_batch_and_arguments()
+        self.assertEqual(req.cache_salt, ["tenant-A", None])
+        self.assertEqual(req[0].cache_salt, "tenant-A")
+        self.assertIsNone(req[1].cache_salt)
+
+        req = GenerateReqInput(
+            text=["Hello", "World"],
+            cache_salt="shared",
+            sampling_params={"n": 2},
+        )
+        req.normalize_batch_and_arguments()
+        self.assertEqual(req.cache_salt, ["shared", "shared"] * 2)
+
+        req = GenerateReqInput(text="Hello", cache_salt="")
+        req.normalize_batch_and_arguments()
+        self.assertIsNone(req.cache_salt)
+
+        req = GenerateReqInput(
+            text=["Hello", "World"],
+            cache_salt=["only-one"],
+            sampling_params=[{}, {}],
+        )
+        with self.assertRaisesRegex(ValueError, "batch size"):
+            req.normalize_batch_and_arguments()
+
+    def test_cache_key_normalization_rejects_invalid_types(self):
+        for field_name in ("extra_key", "cache_salt"):
+            with self.subTest(field_name=field_name, mode="single"):
+                req = GenerateReqInput(text="Hello", **{field_name: ["value"]})
+                with self.assertRaisesRegex(ValueError, "single request"):
+                    req.normalize_batch_and_arguments()
+
+            with self.subTest(field_name=field_name, mode="batch"):
+                req = GenerateReqInput(
+                    text=["Hello", "World"],
+                    sampling_params=[{}, {}],
+                    **{field_name: ["value", 1]},
+                )
+                with self.assertRaisesRegex(ValueError, "should be a string"):
+                    req.normalize_batch_and_arguments()
+
     def test_logprob_parameters_normalization(self):
         """Test normalization of logprob-related parameters."""
         # Test single example

test/registered/unit/managers/test_priority_scheduling_disaggregation.py

@@ -244,6 +244,7 @@ class TestDecodePreallocQueueRebootstrapPayload(unittest.TestCase):
             bootstrap_room=7,
             priority=10,
             extra_key=None,
+            cache_salt=None,
             routing_key=None,
             disagg_prefill_dp_rank=None,
         )
@@ -260,6 +261,7 @@ class TestDecodePreallocQueueRebootstrapPayload(unittest.TestCase):
         self.assertTrue(all(type(x) is int for x in payload["input_ids"]))
         self.assertEqual(payload["sampling_params"]["max_new_tokens"], 1)
         self.assertEqual(payload["bootstrap_room"], 7)
+        self.assertIsNone(payload["cache_salt"])
         # The prefill /generate URL is derived from bootstrap info on the decode
         # side, not sent in the payload; and the boundary token is replayed via
         # the decode-side override, so neither belongs in the payload.

test/registered/unit/managers/test_scheduler_chunked_req_gate.py

@@ -47,6 +47,7 @@ def _make_req(
     req.logprob_start_len = -1
     req.positional_embed_overrides = None
     req.extra_key = None
+    req.cache_salt = None
     req.mamba_pool_idx = None
     req.sampling_params = SimpleNamespace(max_new_tokens=128, ignore_eos=False)
     return req

test/registered/unit/mem_cache/test_decode_radix_lock_ref.py

@@ -78,6 +78,7 @@ class MockReq:
         self.cache_protected_len = cache_protected_len
         self.last_node = last_node
         self.extra_key = None
+        self.cache_salt = None
         self.prefix_indices = torch.empty(0, dtype=torch.int64)
         self.priority = 0
         self.kv_committed_len = len(fill_ids)

test/registered/unit/mem_cache/test_mem_cache_utils.py

@@ -5,6 +5,7 @@ import sys
 import types
 import unittest
 from array import array
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 from sglang.srt.mem_cache.evict_policy import (
@@ -17,6 +18,7 @@ from sglang.srt.mem_cache.evict_policy import (
     SLRUStrategy,
 )
 from sglang.srt.mem_cache.utils import (
+    compute_node_event_hash_values,
     compute_node_hash_values,
     get_eviction_strategy,
     get_hash_str,
@@ -54,9 +56,10 @@ def _legacy_page_hashes(key, page_size, prior_hash=None):
 
 
 class _HashKey:
-    def __init__(self, token_ids, is_bigram=False):
+    def __init__(self, token_ids, is_bigram=False, cache_salt=None):
         self.token_ids = token_ids
         self.is_bigram = is_bigram
+        self.cache_salt = cache_salt
 
     def __len__(self):
         if self.is_bigram:
@@ -68,8 +71,12 @@ class _HashKey:
             start = index.start or 0
             stop = index.stop if index.stop is not None else len(self)
             if self.is_bigram:
-                return _HashKey(self.token_ids[start : stop + 1], is_bigram=True)
-            return _HashKey(self.token_ids[start:stop])
+                return _HashKey(
+                    self.token_ids[start : stop + 1],
+                    is_bigram=True,
+                    cache_salt=self.cache_salt,
+                )
+            return _HashKey(self.token_ids[start:stop], cache_salt=self.cache_salt)
         if self.is_bigram:
             return (self.token_ids[index], self.token_ids[index + 1])
         return self.token_ids[index]
@@ -298,6 +305,7 @@ class TestComputeNodeHashValues(unittest.TestCase):
         node = MagicMock()
         node.key = key
         node.parent = parent
+        node.event_hash_value = None
         if parent is not None:
             parent.hash_value = parent_hash_values
         return node
@@ -318,6 +326,60 @@ class TestComputeNodeHashValues(unittest.TestCase):
                     _legacy_page_hashes(key, page_size=page_size),
                 )
 
+    def test_cache_salt_seeds_root_hash_chain(self):
+        key = _HashKey(array("q", range(1, 17)), cache_salt="tenant-a")
+        seed = hashlib.sha256(b"sglang-cache-salt-v1\0tenant-a").hexdigest()
+        self.assertEqual(
+            compute_node_event_hash_values(self._make_node(key), page_size=8),
+            _legacy_page_hashes(key, page_size=8, prior_hash=seed),
+        )
+        self.assertEqual(
+            compute_node_hash_values(self._make_node(key), page_size=8),
+            _legacy_page_hashes(key, page_size=8),
+        )
+
+        other = _HashKey(array("q", range(1, 17)), cache_salt="tenant-b")
+        self.assertNotEqual(
+            compute_node_event_hash_values(self._make_node(key), page_size=8),
+            compute_node_event_hash_values(self._make_node(other), page_size=8),
+        )
+
+    def test_cache_salt_event_hashes_are_memoized(self):
+        node = self._make_node(
+            _HashKey(array("q", range(1, 17)), cache_salt="tenant-a")
+        )
+        with patch(
+            "sglang.srt.mem_cache.utils.get_hash_str", wraps=get_hash_str
+        ) as mock_get_hash_str:
+            first = compute_node_event_hash_values(node, page_size=8)
+            second = compute_node_event_hash_values(node, page_size=8)
+
+        self.assertIs(first, second)
+        mock_get_hash_str.assert_called_once()
+
+    def test_cache_salt_event_hash_walk_is_iterative(self):
+        root = SimpleNamespace(
+            key=_HashKey(array("q")),
+            parent=None,
+            hash_value=[],
+            event_hash_value=None,
+        )
+        node = root
+        path = []
+        for token_id in range(1, 1102):
+            node = SimpleNamespace(
+                key=_HashKey(array("q", [token_id]), cache_salt="tenant-a"),
+                parent=node,
+                hash_value=None,
+                event_hash_value=None,
+            )
+            path.append(node)
+
+        result = compute_node_event_hash_values(node, page_size=1)
+
+        self.assertEqual(len(result), 1)
+        self.assertTrue(all(item.event_hash_value is not None for item in path))
+
     def test_parent_hash_is_used_only_when_parent_has_nonempty_key_and_hash(self):
         parent = MagicMock()
         parent.key = _HashKey(array("q", range(1, 17)))

test/registered/unit/mem_cache/test_radix_cache_cpp_unit.py

@@ -0,0 +1,38 @@
+"""Unit tests for fail-closed C++ radix-cache request validation."""
+
+import importlib
+import sys
+import types
+import unittest
+from unittest.mock import patch
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+class TestRadixCacheCppCacheSalt(CustomTestCase):
+    def test_cache_salt_is_rejected_without_loading_cpp_extension(self):
+        extension_name = "sglang.srt.mem_cache.cpp_radix_tree.radix_tree"
+        module_name = "sglang.srt.mem_cache.radix_cache_cpp"
+        fake_extension = types.ModuleType(extension_name)
+        fake_extension.IOHandle = object
+        fake_extension.RadixTreeCpp = object
+        fake_extension.TreeNodeCpp = object
+
+        original_module = sys.modules.pop(module_name, None)
+        try:
+            with patch.dict(sys.modules, {extension_name: fake_extension}):
+                module = importlib.import_module(module_name)
+                module.RadixCacheCpp._reject_cache_salt(None)
+                with self.assertRaisesRegex(ValueError, "experimental C\\+\\+"):
+                    module.RadixCacheCpp._reject_cache_salt("tenant-a")
+        finally:
+            sys.modules.pop(module_name, None)
+            if original_module is not None:
+                sys.modules[module_name] = original_module
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_radix_cache_unit.py

@@ -100,6 +100,16 @@ class TestRadixKey(unittest.TestCase):
         self.assertEqual(list(key[2:2].token_ids), [])  # Empty slice
         self.assertEqual(list(key[:].token_ids), [1, 2, 3, 4, 5])  # Full slice
 
+    def test_cache_salt_is_preserved_by_slicing(self):
+        key = RadixKey(
+            array("q", [1, 2, 3, 4]),
+            extra_key="classification",
+            cache_salt="tenant-a",
+        )
+        sliced = key[1:3]
+        self.assertEqual(sliced.extra_key, "classification")
+        self.assertEqual(sliced.cache_salt, "tenant-a")
+
     def test_getitem_invalid_index(self):
         """Test __getitem__ with invalid indices."""
         key = RadixKey(array("q", [1, 2, 3]))
@@ -424,6 +434,7 @@ class TestRadixCache(unittest.TestCase):
             req_pool_idx=0,
             cache_protected_len=0,
             extra_key=None,
+            cache_salt=None,
             priority=0,
             last_node=cache.root_node,
         )
@@ -579,6 +590,107 @@ class TestRadixCache(unittest.TestCase):
         # Non-existent extra_key should not match
         self.assertEqual(len(result4.device_indices), 0)
 
+    def test_cache_salt_isolation_is_independent_of_extra_key(self):
+        cache = RadixCache.create_simulated()
+        tokens = array("q", [1, 2, 3])
+
+        cache.insert(
+            InsertParams(
+                key=RadixKey(tokens, extra_key="bc", cache_salt="a"),
+                value=torch.tensor([10, 20, 30], dtype=torch.int64),
+            )
+        )
+        cache.insert(
+            InsertParams(
+                key=RadixKey(tokens, extra_key="c", cache_salt="ab"),
+                value=torch.tensor([40, 50, 60], dtype=torch.int64),
+            )
+        )
+
+        first = cache.match_prefix(
+            MatchPrefixParams(key=RadixKey(tokens, extra_key="bc", cache_salt="a"))
+        )
+        second = cache.match_prefix(
+            MatchPrefixParams(key=RadixKey(tokens, extra_key="c", cache_salt="ab"))
+        )
+        torch.testing.assert_close(
+            first.device_indices, torch.tensor([10, 20, 30], dtype=torch.int64)
+        )
+        torch.testing.assert_close(
+            second.device_indices, torch.tensor([40, 50, 60], dtype=torch.int64)
+        )
+
+    def test_cache_salt_is_included_in_store_and_remove_events(self):
+        mock_allocator = unittest.mock.Mock()
+        mock_allocator.device = torch.device("cpu")
+        cache = RadixCache.create_simulated(
+            mock_allocator=mock_allocator,
+            page_size=2,
+            enable_kv_cache_events=True,
+        )
+        tokens = array("q", [1, 2, 3, 4])
+        cache.insert(
+            InsertParams(
+                key=RadixKey(tokens, cache_salt="tenant-a"),
+                value=torch.tensor([10, 20, 30, 40], dtype=torch.int64),
+            )
+        )
+        cache.evict(EvictParams(num_tokens=len(tokens)))
+        events = cache.take_events()
+        stored = [event for event in events if isinstance(event, BlockStored)]
+        removed = [event for event in events if isinstance(event, BlockRemoved)]
+
+        self.assertEqual(len(stored), 2)
+        self.assertTrue(
+            all(event.metadata.cache_salt == "tenant-a" for event in stored)
+        )
+        self.assertEqual(stored[1].parent_block_hash, stored[0].block_hashes[0])
+        self.assertEqual(
+            removed[0].block_hashes,
+            [event.block_hashes[0] for event in stored],
+        )
+
+        unsalted = RadixCache.create_simulated(page_size=2, enable_kv_cache_events=True)
+        unsalted.insert(InsertParams(key=RadixKey(tokens), value=None))
+        unsalted_hashes = [
+            event.block_hashes[0]
+            for event in unsalted.take_events()
+            if isinstance(event, BlockStored)
+        ]
+        self.assertNotEqual(
+            unsalted_hashes, [event.block_hashes[0] for event in stored]
+        )
+
+    def test_cache_salt_event_hashes_are_preserved_across_node_split(self):
+        cache = RadixCache.create_simulated(page_size=2, enable_kv_cache_events=True)
+        original = RadixKey(array("q", [1, 2, 3, 4]), cache_salt="tenant-a")
+        cache.insert(
+            InsertParams(
+                key=original,
+                value=torch.tensor([10, 20, 30, 40], dtype=torch.int64),
+            )
+        )
+        original_node = cache.match_prefix(
+            MatchPrefixParams(key=original)
+        ).last_device_node
+        original_hashes = list(original_node.event_hash_value)
+
+        cache.insert(
+            InsertParams(
+                key=RadixKey(array("q", [1, 2, 9, 10]), cache_salt="tenant-a"),
+                value=torch.tensor([10, 20, 90, 100], dtype=torch.int64),
+            )
+        )
+        split_child = cache.match_prefix(
+            MatchPrefixParams(key=original)
+        ).last_device_node
+        split_parent = split_child.parent
+
+        self.assertEqual(
+            split_parent.event_hash_value + split_child.event_hash_value,
+            original_hashes,
+        )
+
     def test_lock_ref_operations(self):
         """Test lock reference counting operations."""
         cache = RadixCache.create_simulated()

test/registered/unit/mem_cache/test_radix_force_miss.py

@@ -31,6 +31,7 @@ class _StubReq:
         self.origin_input_ids = array("q", token_ids)
         self.output_ids = array("q")
         self.extra_key = None
+        self.cache_salt = None
         self.prefix_indices = None
         self.last_node = None
         self.last_host_node = None

test/registered/unit/mem_cache/test_session_token_share_unit.py

@@ -46,6 +46,7 @@ def _recv(rid, input_ids, max_new_tokens=8):
         priority=None,
         routing_key=None,
         extra_key=None,
+        cache_salt=None,
         http_worker_ipc=None,
         time_stats=None,
     )

test/registered/unit/mem_cache/test_streaming_session_unit.py

@@ -66,6 +66,7 @@ class _FakeReq:
         self.origin_input_ids = list(range(committed))
         self.output_ids = []
         self.extra_key = None
+        self.cache_salt = None
         self.last_node = None
         self.cache_protected_len = 0
         self.swa_uuid_for_lock = None

test/registered/unit/mem_cache/test_swa_eviction_boundary.py

@@ -110,6 +110,7 @@ def _make_req(req_pool_idx, token_ids, cache_protected_len, tree):
             swa_evicted_seqlen=0,
         ),
         extra_key=None,
+        cache_salt=None,
         last_node=tree.root_node,
         swa_uuid_for_lock=None,
         swa_prefix_lock_released=False,

test/registered/unit/mem_cache/test_swa_unittest.py

@@ -607,6 +607,7 @@ class TestSWA(unittest.TestCase):
             (req.req_pool_idx, slice(0, req._kv_committed_len)), kv_indices
         )
         req.extra_key = None
+        req.cache_salt = None
         req.last_node = tree.root_node
         req.swa_uuid_for_lock = None
         req.kv.swa_evicted_seqlen = 0
@@ -644,6 +645,7 @@ class TestSWA(unittest.TestCase):
             (req2.req_pool_idx, slice(0, req2._kv_committed_len)), kv_indices2
         )
         req2.extra_key = None
+        req2.cache_salt = None
         req2.last_node = tree.root_node
         req2.swa_uuid_for_lock = None
         req2.kv.swa_evicted_seqlen = 0