From c0d4d1897c8228f220ac7342cde068b571d92a7a Mon Sep 17 00:00:00 2001 From: Yan Ru Pei Date: Fri, 26 Jun 2026 01:13:40 -0700 Subject: [PATCH] fix: batch BlockRemoved events per radix node (#29265) --- python/sglang/srt/mem_cache/events.py | 17 ++++++------- test/manual/test_kv_events.py | 8 ++++--- .../unit/mem_cache/test_mamba_unittest.py | 10 +++++--- .../unit/mem_cache/test_radix_cache_unit.py | 24 ++++++++++++------- .../unit/mem_cache/test_swa_unittest.py | 10 +++++--- .../test_unified_radix_cache_unittest.py | 10 +++++--- 6 files changed, 51 insertions(+), 28 deletions(-) diff --git a/python/sglang/srt/mem_cache/events.py b/python/sglang/srt/mem_cache/events.py index d659268d1..b1eded35e 100644 --- a/python/sglang/srt/mem_cache/events.py +++ b/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()) diff --git a/test/manual/test_kv_events.py b/test/manual/test_kv_events.py index 95367cef0..4d8874bca 100644 --- a/test/manual/test_kv_events.py +++ b/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( diff --git a/test/registered/unit/mem_cache/test_mamba_unittest.py b/test/registered/unit/mem_cache/test_mamba_unittest.py index 5dda30a73..96f2cb5ca 100755 --- a/test/registered/unit/mem_cache/test_mamba_unittest.py +++ b/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): diff --git a/test/registered/unit/mem_cache/test_radix_cache_unit.py b/test/registered/unit/mem_cache/test_radix_cache_unit.py index b72476217..32afb8840 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_unit.py +++ b/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.""" diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index 93334fd1f..ec869681a 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/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): diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 8d0edb349..67bcf88f5 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/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)