From daf631719690e18d13f67a20eb513fd48c712327 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 27 Aug 2026 19:34:48 -0700 Subject: [PATCH] [mem_cache] Add `free_full` to release the full side of a tombstoned SWA node (#36637) --- python/sglang/srt/mem_cache/allocator/base.py | 8 ++++ .../srt/mem_cache/allocator/hisparse.py | 22 ++++++++++ python/sglang/srt/mem_cache/allocator/swa.py | 30 ++++++++++++- .../srt/mem_cache/multi_ended_allocator.py | 11 +++++ .../sglang/srt/mem_cache/swa_radix_cache.py | 35 ++++++++++----- .../mem_cache/unified_cache/cache_action.py | 11 ++++- .../unified_cache/components/swa_component.py | 8 ++-- .../unified_cache/unified_tree_core.py | 6 ++- .../srt/mem_cache/unified_radix_cache.py | 4 ++ .../mem_cache/test_multi_ended_allocator.py | 27 ++++++++++++ .../unit/mem_cache/test_swa_unittest.py | 43 +++++++++++++++++++ .../test_unified_radix_cache_unittest.py | 5 ++- 12 files changed, 190 insertions(+), 20 deletions(-) diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index ef1de7121..851d32db6 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -123,6 +123,14 @@ class BaseTokenToKVPoolAllocator(abc.ABC): def free(self, free_index: torch.Tensor): raise NotImplementedError() + def free_full(self, free_index: torch.Tensor): + """Free slots whose SWA peers the caller already released. + + A hybrid SWA allocator pairs each full-attention slot with an SWA slot + that can die first; this releases the full side alone. A single pool has + no peer, so it is a plain free().""" + self.free(free_index) + def free_segment(self, free_index: torch.Tensor, *, start_pos: int): """Free ``kv_row[start_pos : start_pos + n]`` of one request (or a page-aligned copy); subclasses may use ``start_pos`` to skip the diff --git a/python/sglang/srt/mem_cache/allocator/hisparse.py b/python/sglang/srt/mem_cache/allocator/hisparse.py index dc283b617..4a10891f8 100644 --- a/python/sglang/srt/mem_cache/allocator/hisparse.py +++ b/python/sglang/srt/mem_cache/allocator/hisparse.py @@ -316,6 +316,7 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.free_pages = None self.release_pages = None self.free_group = None + self.full_free_group = [] self.clear() self.hisparse_kvcache.register_mapping( @@ -364,6 +365,15 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): def free_swa(self, free_indices: torch.Tensor): self.logical_attn_allocator.free_swa(free_indices) + def free_full(self, free_indices: torch.Tensor): + if free_indices.numel() == 0: + return + + if self.free_group is None: + self.logical_attn_allocator.free_full(free_indices) + else: + self.full_free_group.append(self._copy_for_free_group(free_indices)) + def available_size(self) -> int: return min( self.logical_attn_allocator.available_size(), @@ -567,6 +577,7 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.full_to_hisparse_device_index_mapping[:-1].fill_(0) self.free_group = None + self.full_free_group = [] def free(self, free_index: torch.Tensor): if free_index.numel() == 0: @@ -576,3 +587,14 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.logical_attn_allocator.free(free_index) else: self.free_group.append(self._copy_for_free_group(free_index)) + + def free_group_begin(self): + super().free_group_begin() + self.full_free_group = [] + + def free_group_end(self): + super().free_group_end() + if self.full_free_group: + full_free_group = self.full_free_group + self.full_free_group = [] + self.free_full(torch.cat(full_free_group)) diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index 2deffeeb1..a95201664 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -95,6 +95,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.release_pages = None self.free_group = None self.swa_free_group = [] + self.full_free_group = [] self._kvcache = kvcache self.clear() @@ -370,9 +371,24 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.swa_attn_allocator.free(swa_indices) + def free_full(self, free_index: torch.Tensor): + if free_index.numel() == 0: + return + + if self.free_group is None: + # Full side only: a tombstoned range's mapping entries read as the + # padding slot, so `free` would push slot 0 into the SWA free list. + self.full_attn_allocator.free(free_index) + else: + self.full_free_group.append(self._copy_for_free_group(free_index)) + assert ( + self.full_attn_allocator.available_size() <= self.full_attn_allocator.size + ) + def free_group_begin(self): super().free_group_begin() self.swa_free_group = [] + self.full_free_group = [] def free_group_end(self): super().free_group_end() @@ -380,6 +396,10 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): swa_free_group = self.swa_free_group self.swa_free_group = [] self.swa_attn_allocator.free(torch.cat(swa_free_group)) + if self.full_free_group: + full_free_group = self.full_free_group + self.full_free_group = [] + self.free_full(torch.cat(full_free_group)) def _expand_to_full_pages(self, indices: torch.Tensor) -> torch.Tensor: pages = torch.unique(indices // self.page_size) @@ -409,6 +429,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.full_to_swa_index_mapping[:-1].fill_(0) self.free_group = None self.swa_free_group = [] + self.full_free_group = [] def get_cpu_copy(self, indices, mamba_indices=None): return self._kvcache.get_cpu_copy(indices, mamba_indices=mamba_indices) @@ -516,8 +537,13 @@ class PureSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): else: self.free_group.append(self._copy_for_free_group(free_index)) - # Not inherited: the SWA parent's hooks drive swa_free_group, - # which this pure-SWA variant does not have. + def free_full(self, free_index: torch.Tensor): + # All-SWA models have no full-attention pool, so there is nothing to + # release once the SWA side is gone. + return + + # Not inherited: the SWA parent's hooks drive swa_free_group and + # full_free_group, which this pure-SWA variant does not have. def free_group_begin(self): self.free_group = [] diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index 4e8680a8e..154d1b6bf 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -2448,6 +2448,17 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): self.swa_attn_allocator.free(live) self.swa_attn_allocator.clear_inverse_history() + def free_full(self, free_index: torch.Tensor) -> None: + """Release the full-physical page and the virtual id, leaving the swa + side alone -- the caller already tombstoned it (`swa.v2p_page == -1`).""" + if free_index is None or free_index.numel() == 0: + return + if self.free_group is not None: + self.full_free_group.append(self._copy_for_free_group(free_index)) + return + self.full_attn_allocator.free(free_index.detach().to(torch.int64)) + self.full_attn_allocator.clear_inverse_history() + def set_full_to_swa_mapping( self, full_indices: torch.Tensor, swa_indices: torch.Tensor ) -> None: diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 910551eee..747db25bc 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -591,6 +591,21 @@ class SWARadixCache(BasePrefixCache): def total_size(self) -> Tuple[int, int]: return self._total_size_helper() + def _free_node_value( + self, node: TreeNode, value: Optional[torch.Tensor] = None + ) -> Tuple[int, int]: + if value is None: + value = node.value + num_tokens = len(value) + if node.swa_tombstone: + # SWA peers went back in `dec_swa_lock_only` or an SWA evict, so + # only the full side is still ours; `free` would hand the SWA pool + # mapping entries that read as the padding slot. + self.token_to_kv_pool_allocator.free_full(value) + return num_tokens, 0 + self.token_to_kv_pool_allocator.free(value) + return num_tokens, num_tokens + def evict(self, params: EvictParams) -> EvictResult: if self.disable: return EvictResult() @@ -611,11 +626,9 @@ class SWARadixCache(BasePrefixCache): # 1. free node kv indices, evict full and swa tokens self.kv_events.record_remove(x) - self.token_to_kv_pool_allocator.free(x.value) - full_num_evicted += len(x.value) - # Tombstoned leaves had their SWA freed earlier in `dec_swa_lock_only` - if not x.swa_tombstone: - swa_num_evicted += len(x.value) + node_full_evicted, node_swa_evicted = self._free_node_value(x) + full_num_evicted += node_full_evicted + swa_num_evicted += node_swa_evicted # 2. get the next leaf, update the lru lists x_next = self.full_lru_list.get_prev_leaf_no_lock(x) @@ -676,9 +689,9 @@ class SWARadixCache(BasePrefixCache): ), f"leaf node with full lock must also have swa lock, {x.id=}" # 1. a leaf node, free full and swa tokens self.kv_events.record_remove(x) - self.token_to_kv_pool_allocator.free(x.value) - full_num_evicted += len(x.value) - swa_num_evicted += len(x.value) + node_full_evicted, node_swa_evicted = self._free_node_value(x) + full_num_evicted += node_full_evicted + swa_num_evicted += node_swa_evicted # 2. get the next node, update the lru lists x_next = self.swa_lru_list.get_prev_no_lock(x) @@ -1302,7 +1315,7 @@ class SWARadixCache(BasePrefixCache): swa_value = allocator.translate_loc_from_full_to_swa(incoming_full) allocator.set_full_to_swa_mapping(node.value, swa_value) allocator.clear_full_to_swa_mapping(incoming_full) - allocator.full_attn_allocator.free(incoming_full) + allocator.free_full(incoming_full) node.swa_tombstone = False self.swa_lru_list.insert_mru(node) @@ -1346,8 +1359,8 @@ class SWARadixCache(BasePrefixCache): ), f"tombstone swa_lock_ref should always be 0, {node.parent.full_lock_ref=}, {node.parent.swa_lock_ref=}, {node.parent.id=}" # delete tombstone node evicts full tokens self.kv_events.record_remove(node.parent) - self.token_to_kv_pool_allocator.free(node.parent.value) - full_num_evicted += len(node.parent.value) + node_full_evicted, _ = self._free_node_value(node.parent) + full_num_evicted += node_full_evicted self.full_lru_list.remove_node(node.parent) self._delete_tombstone_leaf(node.parent) node = node.parent diff --git a/python/sglang/srt/mem_cache/unified_cache/cache_action.py b/python/sglang/srt/mem_cache/unified_cache/cache_action.py index efea64b41..e94ebdb92 100644 --- a/python/sglang/srt/mem_cache/unified_cache/cache_action.py +++ b/python/sglang/srt/mem_cache/unified_cache/cache_action.py @@ -36,6 +36,13 @@ class FreeDeviceKV(msgspec.Struct, frozen=True): indices: list[torch.Tensor] +class FreeDeviceKVFullOnly(msgspec.Struct, frozen=True): + """Free the full side only, for a tombstoned node whose SWA peers are gone; + FreeDeviceKV would release the SWA side twice.""" + + indices: list[torch.Tensor] + + class ComponentAction(msgspec.Struct, frozen=True): """Base for component-routed actions; the cache dispatches each one to ``component_type``'s class-level ``apply_component_action``; every subclass @@ -100,4 +107,6 @@ class SWARebuild(ComponentAction, frozen=True): # Cache-owned actions, applied by UnifiedRadixCache itself. -CacheAction = ReplaceWriteThroughOnNodeSplit | FreeDeviceKV | BackupKV +CacheAction = ( + ReplaceWriteThroughOnNodeSplit | FreeDeviceKV | FreeDeviceKVFullOnly | BackupKV +) diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py index 49843b525..c28ef35dc 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py @@ -22,7 +22,7 @@ from sglang.srt.mem_cache.hicache_storage import ( from sglang.srt.mem_cache.unified_cache.cache_action import ( FreeComponentDeviceSlot, FreeComponentHostSlot, - FreeDeviceKV, + FreeDeviceKVFullOnly, RebuildFullToSWAMapping, RecoverSWAWithLockedFull, SWARebuild, @@ -295,7 +295,7 @@ class SWAComponent(TreeComponent): ) return 0 full_cd.value = value_slice.clone() - cache_actions.append(FreeDeviceKV([old_full])) + cache_actions.append(FreeDeviceKVFullOnly([old_full])) cache_actions.append(SWARebuild(node.id, value_slice)) return 0 elif swa_evicted_seqlen < total_prefix_len + prefix_len: @@ -313,7 +313,7 @@ class SWAComponent(TreeComponent): ) return start_idx node.component_data[BASE_COMPONENT_TYPE].value = new_full.clone() - cache_actions.append(FreeDeviceKV([old_full])) + cache_actions.append(FreeDeviceKVFullOnly([old_full])) cache_actions.append(SWARebuild(node.id, new_full)) return start_idx else: @@ -1156,7 +1156,7 @@ class SWAComponent(TreeComponent): swa_value = self._translate_full_to_swa(action.incoming_full) alloc.set_full_to_swa_mapping(action.kept_full, swa_value) alloc.clear_full_to_swa_mapping(action.incoming_full) - alloc.full_attn_allocator.free(action.incoming_full) + alloc.free_full(action.incoming_full) self.tree_core.set_component_device_value( action.node_id, self.component_type, swa_value ) diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index c8f56a78d..9797081ba 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -48,6 +48,7 @@ from sglang.srt.mem_cache.unified_cache.cache_action import ( CacheAction, ComponentAction, FreeDeviceKV, + FreeDeviceKVFullOnly, ReplaceWriteThroughOnNodeSplit, ) from sglang.srt.mem_cache.unified_cache.components import ( @@ -939,7 +940,10 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): @staticmethod def _is_deferrable_action(action: CacheAction | ComponentAction) -> bool: """Fire-and-forget actions safe to batch until the next barrier.""" - return isinstance(action, (FreeDeviceKV, ReplaceWriteThroughOnNodeSplit)) + return isinstance( + action, + (FreeDeviceKV, FreeDeviceKVFullOnly, ReplaceWriteThroughOnNodeSplit), + ) def _insert_walk_step(self, state: _InsertWalkState) -> None: """Process one walked node, appending its barrier actions to the state.""" diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 1a7429e84..2c3298474 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -51,6 +51,7 @@ from sglang.srt.mem_cache.unified_cache.cache_action import ( ComponentAction, FreeComponentDeviceSlot, FreeDeviceKV, + FreeDeviceKVFullOnly, ReplaceWriteThroughOnNodeSplit, ) @@ -1052,6 +1053,9 @@ class UnifiedRadixCache(BasePrefixCache): # tree values are page-aligned copies of a kv row: page-exact segments for indices in action.indices: self.token_to_kv_pool_allocator.free_segment(indices, start_pos=0) + elif isinstance(action, FreeDeviceKVFullOnly): + for indices in action.indices: + self.token_to_kv_pool_allocator.free_full(indices) elif isinstance(action, BackupKV): self._execute_and_commit_kv_backup(action) else: diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index 9b8adef09..06fa88555 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -744,6 +744,33 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase): ) self.assertIn(tgt, free_full) + def test_swa_free_full_defers_inside_a_free_group(self): + """The full-only release joins the barrier, like `free`.""" + _, allocator, kvcache = self._build() + v = self._alloc(allocator, kvcache, 3) + target = v[1:2] + tgt = int(target.item()) + # Tombstone the swa side, erasing each marker before its release + # (compaction runs inside both). + target_swa = allocator.swa_attn_allocator.virtual_to_physical[target] + kvcache.swa_kv_pool.buf[target_swa] = -1 + allocator.free_swa(target) + full_phys = int(allocator.full_attn_allocator.virtual_to_physical[tgt].item()) + kvcache.full_kv_pool.buf[full_phys] = -1 + + allocator.free_group_begin() + allocator.free_full(target) + deferred = set( + int(x) for x in allocator.full_attn_allocator.free_virtual_ids.tolist() + ) + self.assertNotIn(tgt, deferred) + + allocator.free_group_end() + drained = set( + int(x) for x in allocator.full_attn_allocator.free_virtual_ids.tolist() + ) + self.assertIn(tgt, drained) + # 4. Compaction diverges between the two sub-pools (each runs its own). def test_swa_compaction_diverges_physical_layout(self): _, allocator, kvcache = self._build() diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index 3a3a10ba8..38c55f5c0 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -851,5 +851,48 @@ class TestSWASplitLeafOnInsert(CustomTestCase): tree.sanity_check() +class TestFreeFullPartition(CustomTestCase): + """`free_full` releases only the full side of a hybrid SWA allocator.""" + + def setUp(self): + _, self.allocator, _ = _build_swa_tree(is_eagle=False) + self.full_baseline = self.allocator.full_available_size() + self.swa_baseline = self.allocator.swa_available_size() + + def _sizes(self): + return ( + self.allocator.full_available_size(), + self.allocator.swa_available_size(), + ) + + def test_free_full_keeps_the_swa_peers_allocated(self): + indices = _swa_alloc(self.allocator, 4) + self.allocator.free_full(indices) + + full_avail, swa_avail = self._sizes() + self.assertEqual(full_avail, self.full_baseline) + self.assertEqual(swa_avail, self.swa_baseline - 4) + + def test_free_full_leaves_the_mapping_intact(self): + indices = _swa_alloc(self.allocator, 4) + before = self.allocator.full_to_swa_index_mapping[indices].clone() + self.allocator.free_full(indices) + + self.assertTrue(bool((before > 0).all())) + self.assertTrue( + torch.equal(self.allocator.full_to_swa_index_mapping[indices], before) + ) + + def test_free_full_is_deferred_inside_a_free_group(self): + indices = _swa_alloc(self.allocator, 4) + + self.allocator.free_group_begin() + self.allocator.free_full(indices) + self.assertEqual(self.allocator.full_available_size(), self.full_baseline - 4) + self.allocator.free_group_end() + + self.assertEqual(self.allocator.full_available_size(), self.full_baseline) + + if __name__ == "__main__": unittest.main() 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 08834f981..ff6c9ee53 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 @@ -6733,6 +6733,7 @@ class TestUnifiedRadixCacheActionRouting(CustomTestCase): # translate the source full to SWA and store it on the node (no free) alloc.translate_loc_from_full_to_swa.assert_called_once_with(source_value) alloc.free.assert_not_called() + alloc.free_full.assert_not_called() cache.tree_core.set_component_device_value.assert_called_once_with( 5, ComponentType.SWA, swa_value ) @@ -6767,7 +6768,9 @@ class TestUnifiedRadixCacheActionRouting(CustomTestCase): alloc.set_full_to_swa_mapping.assert_called_once_with(kept_full, swa_value) # the incoming full's stale mapping is cleared, then its slot freed (full-only) alloc.clear_full_to_swa_mapping.assert_called_once_with(incoming_full) - alloc.full_attn_allocator.free.assert_called_once_with(incoming_full) + alloc.free_full.assert_called_once_with(incoming_full) + # not the inner allocator (skips the free-group defer) and not both halves + alloc.full_attn_allocator.free.assert_not_called() alloc.free.assert_not_called() cache.tree_core.set_component_device_value.assert_called_once_with( 5, ComponentType.SWA, swa_value