fix: batch BlockRemoved events per radix node (#29265)

python/sglang/srt/mem_cache/events.py

@@ -84,7 +84,7 @@ class KVCacheEventMixin:
                 page_index += 1
 
     def _record_remove_event(self, node: Any, medium=None):
-        # One BlockRemoved per chunk.
+        # One BlockRemoved per radix node.
         # ``medium`` defaults to StorageMedium.GPU but callers may override for
         # lower-tier removals (e.g. StorageMedium.CPU when evicting from host).
         if self.enable_kv_cache_events:
@@ -95,21 +95,22 @@ class KVCacheEventMixin:
             if node.hash_value is None:
                 node.hash_value = compute_node_hash_values(node, self.page_size)
 
-            page_index = 0
+            block_hashes = []
             logical_len = len(node.key)
+            page_index = 0
             for start in range(0, logical_len, self.page_size):
                 end = min(start + self.page_size, logical_len)
                 if end <= start:
                     continue
 
-                block_hash = hash_str_to_int64(node.hash_value[page_index])
-
-                self.kv_event_queue.append(
-                    BlockRemoved(block_hashes=[block_hash], medium=medium)
-                )
-
+                block_hashes.append(hash_str_to_int64(node.hash_value[page_index]))
                 page_index += 1
 
+            if block_hashes:
+                self.kv_event_queue.append(
+                    BlockRemoved(block_hashes=block_hashes, medium=medium)
+                )
+
     def _record_all_cleared_event(self):
         if self.enable_kv_cache_events:
             self.kv_event_queue.append(AllBlocksCleared())

test/manual/test_kv_events.py

@@ -140,10 +140,12 @@ class TestKvEvents(CustomTestCase):
                 elif isinstance(event, BlockRemoved):
                     # Validate BlockRemoved structure
                     self.assertIsInstance(event.block_hashes, list)
-                    self.assertEqual(
-                        len(event.block_hashes), 1, "Should have one hash per block"
+                    self.assertGreater(
+                        len(event.block_hashes),
+                        0,
+                        "Should have at least one removed block hash",
                     )
-                    removed_hashes.add(event.block_hashes[0])
+                    removed_hashes.update(event.block_hashes)
 
             # Verify we got both BlockStored and BlockRemoved events
             self.assertGreater(

test/registered/unit/mem_cache/test_mamba_unittest.py

@@ -29,6 +29,10 @@ register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
 register_amd_ci(est_time=9, suite="stage-b-test-1-gpu-small-amd")
 
 
+def _event_hashes(events):
+    return [block_hash for event in events for block_hash in event.block_hashes]
+
+
 class TestMamba(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
@@ -374,9 +378,9 @@ class TestMamba(unittest.TestCase):
         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)
-        ]
+        removed_hashes = _event_hashes(
+            [e for e in events if isinstance(e, BlockRemoved)]
+        )
         self.assertCountEqual(removed_hashes, stored_hashes)
 
     def test_mamba_radix_cache_kv_events_split_hash(self):

test/registered/unit/mem_cache/test_radix_cache_unit.py

@@ -484,22 +484,25 @@ class TestRadixCache(unittest.TestCase):
         mock_allocator.device = torch.device("cpu")
 
         cache = RadixCache.create_simulated(
-            mock_allocator=mock_allocator, enable_kv_cache_events=True
+            mock_allocator=mock_allocator,
+            page_size=2,
+            enable_kv_cache_events=True,
         )
 
         # Insert and then evict data
+        seq = [1, 2, 3, 4]
         cache.insert(
             InsertParams(
-                key=RadixKey(array("q", [1, 2, 3])),
-                value=torch.tensor([10, 20, 30], dtype=torch.int64),
+                key=RadixKey(array("q", seq)),
+                value=torch.tensor([10, 20, 30, 40], dtype=torch.int64),
             )
         )
-        result = cache.evict(EvictParams(num_tokens=3))
+        result = cache.evict(EvictParams(num_tokens=len(seq)))
         self.assertIsInstance(result, EvictResult)
         self.assertGreaterEqual(
             result.num_tokens_evicted,
-            3,
-            f"evicted {result.num_tokens_evicted} tokens, expected at least 3",
+            len(seq),
+            f"evicted {result.num_tokens_evicted} tokens, expected at least {len(seq)}",
         )
 
         # Take events - should include both store and remove events
@@ -510,10 +513,15 @@ class TestRadixCache(unittest.TestCase):
         event_types = [type(event).__name__ for event in events]
         self.assertIn("BlockStored", event_types)
 
+        stored_hashes = [
+            event.block_hashes[0] for event in events if isinstance(event, BlockStored)
+        ]
+        self.assertEqual(len(stored_hashes), 2)
+
         # Verify BlockRemoved event content
         remove_events = [e for e in events if isinstance(e, BlockRemoved)]
-        for event in remove_events:
-            self.assertGreater(len(event.block_hashes), 0)
+        self.assertEqual(len(remove_events), 1)
+        self.assertEqual(remove_events[0].block_hashes, stored_hashes)
 
     def test_extra_key_isolation(self):
         """Test that keys with different extra_key values are isolated."""

test/registered/unit/mem_cache/test_swa_unittest.py

@@ -27,6 +27,10 @@ register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-large")
 register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd")
 
 
+def _event_hashes(events):
+    return [block_hash for event in events for block_hash in event.block_hashes]
+
+
 class _DummyReq:
     def __init__(self):
         self._kv_committed_len = 0
@@ -172,9 +176,9 @@ class TestSWA(unittest.TestCase):
 
         result = tree.evict(EvictParams(num_tokens=1, swa_num_tokens=0))
         self.assertGreaterEqual(result.num_tokens_evicted, 1)
-        removed_hashes = [
-            e.block_hashes[0] for e in tree.take_events() if isinstance(e, BlockRemoved)
-        ]
+        removed_hashes = _event_hashes(
+            [e for e in tree.take_events() if isinstance(e, BlockRemoved)]
+        )
         self.assertCountEqual(removed_hashes, stored_hashes)
 
     def test_swa_radix_cache_kv_events_split_hash(self):

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -443,6 +443,9 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase):
             events = [e for e in events if e.medium == medium]
         return events
 
+    def _event_hashes(self, events):
+        return [block_hash for event in events for block_hash in event.block_hashes]
+
     def _leaf_for(self, tree, tokens):
         match = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens))))
         self.assertIsNot(match.last_device_node, tree.root_node)
@@ -512,7 +515,8 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase):
         result = tree.evict(EvictParams(num_tokens=len(seq)))
         self.assertGreaterEqual(result.num_tokens_evicted, len(seq))
         removed = self._removed_events(tree, StorageMedium.GPU)
-        self.assertCountEqual([e.block_hashes[0] for e in removed], stored_hashes)
+        self.assertEqual(len(removed), 1)
+        self.assertEqual(removed[0].block_hashes, stored_hashes)
 
     def test_kv_events_split_preserves_block_hash_parentage(self):
         tree, allocator, _ = build_fixture(self.cfg, enable_kv_cache_events=True)
@@ -554,7 +558,7 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase):
 
         tree.evict(EvictParams(num_tokens=len(seq)))
         removed_gpu = self._removed_events(tree, StorageMedium.GPU)
-        self.assertCountEqual([e.block_hashes[0] for e in removed_gpu], stored_hashes)
+        self.assertCountEqual(self._event_hashes(removed_gpu), stored_hashes)
 
         self._load_back_node(tree, node)
         restored_gpu = self._stored_events(tree, StorageMedium.GPU)
@@ -564,7 +568,7 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase):
         self._removed_events(tree, StorageMedium.GPU)
         tree.evict_host(len(seq))
         removed_cpu = self._removed_events(tree, StorageMedium.CPU)
-        self.assertCountEqual([e.block_hashes[0] for e in removed_cpu], stored_hashes)
+        self.assertCountEqual(self._event_hashes(removed_cpu), stored_hashes)
 
     def test_hicache_split_pending_write_through_publishes_fragments(self):
         tree, allocator, _ = build_fixture(self.cfg, enable_kv_cache_events=True)