feat: Add KV events for Mamba radix cache (#23678)

Signed-off-by: zhongdaor-nv <220807034+zhongdaor-nv@users.noreply.github.com>
Co-authored-by: zhongdaor-nv <220807034+zhongdaor-nv@users.noreply.github.com>

test/registered/4-gpu-models/test_qwen35_hicache.py

@@ -1,10 +1,14 @@
 import shutil
 import tempfile
+import time
 import unittest
 from types import SimpleNamespace
 
 import requests
+import zmq
+from msgspec.msgpack import Decoder
 
+from sglang.srt.disaggregation.kv_events import BlockStored, KVEventBatch
 from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
 
@@ -72,6 +76,8 @@ class TestQwen35WithHiCache(CustomTestCase):
                 "file",
                 "--hicache-storage-prefetch-policy",
                 "wait_complete",
+                "--kv-events-config",
+                '{"publisher": "zmq", "topic": "kv-events"}',
             ],
         )
 
@@ -125,6 +131,42 @@ class TestQwen35WithHiCache(CustomTestCase):
             f"first={first_metrics['score']}, second={second_metrics['score']}",
         )
 
+    def test_kv_events_smoke(self):
+        decoder = Decoder(type=KVEventBatch)
+        context = zmq.Context()
+        sub = context.socket(zmq.SUB)
+        sub.connect("tcp://localhost:5557")
+        sub.setsockopt_string(zmq.SUBSCRIBE, "kv-events")
+
+        try:
+            time.sleep(1.0)
+            res = requests.post(
+                f"{self.base_url}/generate",
+                json={
+                    "text": "HiCache KV event compatibility check. " * 64,
+                    "sampling_params": {"temperature": 0, "max_new_tokens": 1},
+                },
+                timeout=120,
+            )
+            res.raise_for_status()
+
+            events = []
+            deadline = time.time() + 10
+            while time.time() < deadline and not any(
+                isinstance(event, BlockStored) for event in events
+            ):
+                if sub.poll(timeout=100):
+                    _, _, payload = sub.recv_multipart()
+                    events.extend(decoder.decode(payload).events)
+
+            self.assertTrue(
+                any(isinstance(event, BlockStored) for event in events),
+                "Expected at least one BlockStored event from Qwen3.5 HiCache server",
+            )
+        finally:
+            sub.close()
+            context.term()
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_mamba_unittest.py

@@ -3,6 +3,7 @@ import unittest
 import torch
 
 from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
+from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored
 from sglang.srt.environ import envs
 from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
@@ -305,7 +306,98 @@ class TestMamba(unittest.TestCase):
         print(available_and_evictable_str(tree))
         tree.sanity_check()
 
-    def _setup_tree_and_allocator(self):
+    def test_mamba_radix_cache_kv_events(self):
+        tree, allocator, _, make_dummy_req = self._setup_tree_and_allocator(
+            enable_kv_cache_events=True
+        )
+        tree.take_events()  # Clear the reset event.
+
+        stored_hashes = []
+
+        req1 = make_dummy_req()
+        key1 = RadixKey([1, 2, 3])
+        tree.insert(
+            InsertParams(
+                key=key1,
+                value=allocator.alloc(3)[: len(key1)],
+                mamba_value=req1.mamba_pool_idx.unsqueeze(0),
+            )
+        )
+        events = tree.take_events()
+        stored_events = [e for e in events if isinstance(e, BlockStored)]
+        self.assertEqual(len(stored_events), 3)
+        self.assertEqual([e.token_ids[0] for e in stored_events], [1, 2, 3])
+        stored_hashes.extend(e.block_hashes[0] for e in stored_events)
+
+        req2 = make_dummy_req()
+        key2 = RadixKey([1, 2, 3, 4, 5])
+        tree.insert(
+            InsertParams(
+                key=key2,
+                value=allocator.alloc(5)[: len(key2)],
+                mamba_value=req2.mamba_pool_idx.unsqueeze(0),
+            )
+        )
+        events = tree.take_events()
+        stored_events = [e for e in events if isinstance(e, BlockStored)]
+        self.assertEqual(len(stored_events), 2)
+        self.assertEqual([e.token_ids[0] for e in stored_events], [4, 5])
+        stored_hashes.extend(e.block_hashes[0] for e in stored_events)
+
+        # Evicting an internal mamba state creates a tombstone but does not
+        # remove full-attention KV blocks, so it must not emit BlockRemoved.
+        result = tree.evict(EvictParams(num_tokens=0, mamba_num=1))
+        self.assertEqual(result.num_tokens_evicted, 0)
+        self.assertEqual(result.mamba_num_evicted, 1)
+        events = tree.take_events()
+        self.assertEqual([e for e in events if isinstance(e, BlockRemoved)], [])
+
+        result = tree.evict(EvictParams(num_tokens=1))
+        self.assertGreaterEqual(result.num_tokens_evicted, 1)
+        events = tree.take_events()
+        removed_hashes = [
+            e.block_hashes[0] for e in events if isinstance(e, BlockRemoved)
+        ]
+        self.assertCountEqual(removed_hashes, stored_hashes)
+
+    def test_mamba_radix_cache_kv_events_split_hash(self):
+        tree, allocator, _, make_dummy_req = self._setup_tree_and_allocator(
+            enable_kv_cache_events=True
+        )
+        tree.take_events()  # Clear the reset event.
+
+        req1 = make_dummy_req()
+        key1 = RadixKey([1, 2, 3, 4])
+        tree.insert(
+            InsertParams(
+                key=key1,
+                value=allocator.alloc(4)[: len(key1)],
+                mamba_value=req1.mamba_pool_idx.unsqueeze(0),
+            )
+        )
+        first_insert_events = [
+            e for e in tree.take_events() if isinstance(e, BlockStored)
+        ]
+        self.assertEqual(len(first_insert_events), 4)
+        split_parent_hash = first_insert_events[1].block_hashes[0]
+
+        req2 = make_dummy_req()
+        key2 = RadixKey([1, 2, 5, 6])
+        tree.insert(
+            InsertParams(
+                key=key2,
+                value=allocator.alloc(4)[: len(key2)],
+                mamba_value=req2.mamba_pool_idx.unsqueeze(0),
+            )
+        )
+        second_insert_events = [
+            e for e in tree.take_events() if isinstance(e, BlockStored)
+        ]
+        self.assertEqual(len(second_insert_events), 2)
+        self.assertEqual(second_insert_events[0].token_ids, [5])
+        self.assertEqual(second_insert_events[0].parent_block_hash, split_parent_hash)
+
+    def _setup_tree_and_allocator(self, enable_kv_cache_events=False):
         """Helper to create a MambaRadixCache with allocator for testing."""
         set_global_server_args_for_scheduler(
             ServerArgs(model_path="dummy", page_size=1)
@@ -374,6 +466,7 @@ class TestMamba(unittest.TestCase):
             token_to_kv_pool_allocator=allocator,
             page_size=1,
             disable=False,
+            enable_kv_cache_events=enable_kv_cache_events,
         )
         tree = MambaRadixCache(params=params)