diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 89d4c088e..46e32f878 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -867,6 +867,14 @@ class HiCacheController: self.prefetch_thread.start() self.backup_thread.start() + # PREREQ-2: reset the B1 pooled-L2 allocator so flush_cache reclaims the + # shared pool. The radix tree + host pool are cleared by the caller; without + # this the allocator free list / ranges / committed set stay stale (leak). + # Safe here: the ack queues are cleared above and flush_cache is idle-gated + # (no in-flight backup/reserve referencing dropped ranges). + if self.cp_shared_l2_page_allocator is not None: + self.cp_shared_l2_page_allocator.reset() + def clear_draft_host_pool(self) -> None: if self.draft_mem_pool_host is not None: self.draft_mem_pool_host.clear() diff --git a/python/sglang/srt/mem_cache/cp_shared_l2_pool.py b/python/sglang/srt/mem_cache/cp_shared_l2_pool.py index 1a04ae211..5bc11f40c 100644 --- a/python/sglang/srt/mem_cache/cp_shared_l2_pool.py +++ b/python/sglang/srt/mem_cache/cp_shared_l2_pool.py @@ -607,6 +607,26 @@ class CpSharedL2PageAllocator: "cp_shared_l2_objects_evicted": 0, } + def reset(self) -> None: + """Reset to the freshly-constructed all-free state (flush_cache hook). + + Rebuilds the per-slab free list to fully-free and drops all object ranges + + the committed set, so flush_cache reclaims the pooled L2 (the allocator + metadata; the slab bytes are left as-is -- unreferenced once the radix tree + and ranges are cleared). Rank-uniform: every CP rank runs it at the + idle-gated flush point, so the post-reset placement_digest is empty and + identical on every rank. Cumulative stats are preserved. + """ + self._generation = 0 + self._ranges_by_object = {} + self._committed_objects = set() + self._free_by_payload = { + payload_kind: { + slab.slab_id: [(0, slab.num_pages)] for slab in slabs + } + for payload_kind, slabs in self._slabs_by_payload.items() + } + def reserve( self, object_key: str, payload_kind: str, num_pages: int ) -> CpSharedL2ObjectRange: diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 8f13381e5..ea7f88413 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1015,7 +1015,26 @@ class HiRadixCache(RadixCache): slabs_by_payload=cp_shared_l2_slabs_by_payload, ) - self.tp_group = params.tp_cache_group + # PREREQ-1 (CP8DP2 scoping): the B1 commit/evict consensus collectives + # (writing_check MIN, placement_digest, drain/evict MINs, the prefetch group, + # the flush barrier) MUST reduce over the CP group -- the ranks that share the + # replicated event stream + the MAP_SHARED slab. tp_cache_group equals the CP + # group ONLY at dp_size=1 (CP==TP); under CP8DP2 (DP attention) tp_cache_group + # is the size-1 attn-TP group, so every `tp_world_size > 1` collective would + # silently no-op per rank -> divergent placement -> shared-slab corruption. + # Scope to the CP cpu group (a no-op handle change at dp_size=1). + _cp_hicache_cp_size = ( + int(getattr(params.token_to_kv_pool_allocator, "cp_size", 1) or 1) + if self._uses_cp_hicache + else 1 + ) + if self._uses_cp_hicache and _cp_hicache_cp_size > 1: + from sglang.srt.layers.dp_attention import get_attention_cp_group + + _cp_group = get_attention_cp_group() + self.tp_group = getattr(_cp_group, "cpu_group", _cp_group) + else: + self.tp_group = params.tp_cache_group self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group) self.pp_rank = params.pp_rank self.pp_size = params.pp_size diff --git a/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py b/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py index 20345c08b..2c6d8b246 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py +++ b/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py @@ -1019,6 +1019,28 @@ class TestCpSharedL2PageAllocator(unittest.TestCase): all(src == 0 and group == "fake-cpu-group" for src, group, _ in broadcasts) ) + def test_reset_restores_freshly_constructed_all_free_state(self): + # PREREQ-2: flush_cache must reclaim the pooled L2. reset() rebuilds the + # free list all-free + drops ranges/committed, restoring the freshly-built + # placement_digest exactly (rank-uniform at the idle flush point). + allocator = self.make_allocator(pages=8) + fresh_digest = allocator.placement_digest() + allocator.reserve("n0", PAYLOAD_TARGET_KV, 3) + allocator.reserve("n1", PAYLOAD_TARGET_KV, 2) + allocator.mark_object_committed("n0") + self.assertLess(allocator.free_pages(PAYLOAD_TARGET_KV), 8) + self.assertNotEqual(allocator.placement_digest(), fresh_digest) + + allocator.reset() + + self.assertEqual(allocator.free_pages(PAYLOAD_TARGET_KV), 8) + self.assertEqual(allocator.free_pages(PAYLOAD_DRAFT_KV), 8) + self.assertFalse(allocator.is_committed("n0")) + self.assertEqual(allocator.object_ranges("n0"), {}) + self.assertEqual(allocator.placement_digest(), fresh_digest) + # usable after reset + self.assertIsNotNone(allocator.reserve("n2", PAYLOAD_TARGET_KV, 4)) + def test_release_and_abort_return_capacity_once_without_double_free(self): allocator = self.make_allocator(pages=5) allocator.reserve("release-me", PAYLOAD_TARGET_KV, 2)