diff --git a/python/sglang/srt/mem_cache/allocator/hisparse.py b/python/sglang/srt/mem_cache/allocator/hisparse.py index 5647154f7..bef9b2a43 100644 --- a/python/sglang/srt/mem_cache/allocator/hisparse.py +++ b/python/sglang/srt/mem_cache/allocator/hisparse.py @@ -358,6 +358,9 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): def free_swa(self, free_indices: torch.Tensor): self.logical_attn_allocator.free_swa(free_indices) + def free_swa_segment(self, free_indices: torch.Tensor, *, start_pos: int): + self.logical_attn_allocator.free_swa_segment(free_indices, start_pos=start_pos) + def free_full(self, free_indices: torch.Tensor): if free_indices.numel() == 0: return diff --git a/python/sglang/srt/mem_cache/allocator/paged.py b/python/sglang/srt/mem_cache/allocator/paged.py index 4d77446c9..ffc479e8c 100755 --- a/python/sglang/srt/mem_cache/allocator/paged.py +++ b/python/sglang/srt/mem_cache/allocator/paged.py @@ -296,12 +296,19 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): torch.unique(free_index.cpu() // ps), ) + self.free_page_ids(reps // ps) + + def free_page_ids(self, page_ids: torch.Tensor): + """Free exactly these pages; no page twice, no dedup.""" + if page_ids.numel() == 0: + return + if self.free_group is None: - self._release_page_ids(reps // ps) + self._release_page_ids(page_ids) if self.debug_mode: self._debug_check_no_duplicate_pages() else: - self.free_page_reps_group.append(self._copy_for_free_group(reps)) + self.free_page_ids_group.append(self._copy_for_free_group(page_ids)) def _debug_check_no_duplicate_pages(self): pages = self.get_all_free_pages() @@ -316,15 +323,13 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): def free_group_begin(self): super().free_group_begin() - self.free_page_reps_group = [] + self.free_page_ids_group = [] def free_group_end(self): super().free_group_end() - if self.free_page_reps_group: - self._release_page_ids( - torch.cat(self.free_page_reps_group) // self.page_size - ) - self.free_page_reps_group = [] + if self.free_page_ids_group: + self._release_page_ids(torch.cat(self.free_page_ids_group)) + self.free_page_ids_group = [] if self.debug_mode: # the no-double-free contract can only break across a group's calls self._debug_check_no_duplicate_pages() @@ -335,7 +340,7 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): 1, self.num_pages + 1, dtype=torch.int64, device=self.device ) self.free_group = None - self.free_page_reps_group = [] + self.free_page_ids_group = [] # need_sort only: freed pages wait here, unsorted, until an alloc runs short. self.staged_pages: list[torch.Tensor] = [] self.num_staged_pages = 0 diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index 9528c96cb..1b9dc6b65 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -111,6 +111,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.release_pages = None self.free_group = None self.swa_free_group = [] + self.swa_page_ids_group = [] self._kvcache = kvcache @@ -429,19 +430,19 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.full_to_swa_index_mapping.index_fill_(0, full_indices, 0) def free_swa(self, free_index: torch.Tensor): + """Release the SWA peers of an arbitrary slot set and clear their mapping. + Synchronizes at page_size > 1; kv-row segments go through free_swa_segment().""" if free_index.numel() == 0: return if self.page_size == 1: - # A filter here would make the output shape data-dependent, - # which costs a device-to-host sync. - mapping_indices = free_index - swa_indices = self.full_to_swa_index_mapping[mapping_indices] - expect(_SWA_PEER_MAPPED, swa_indices > 0, msg="caller wants free_full") - else: - mapping_indices = self._expand_to_full_pages(free_index) - swa_indices = self.full_to_swa_index_mapping[mapping_indices] + self._free_swa_pages(free_index, start_pos=0) + return + # The expansion gathers slots this caller never allocated, which read as + # the padding slot; `_release_swa` filters them once per group. + mapping_indices = self._expand_to_full_pages(free_index) + swa_indices = self.full_to_swa_index_mapping[mapping_indices] self.clear_full_to_swa_mapping(mapping_indices) if self.free_group is not None: @@ -452,11 +453,50 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self._release_swa(swa_indices) + def free_swa_segment(self, free_index: torch.Tensor, *, start_pos: int): + """free_swa() for a kv-row segment; same start-alignment contract as + free_segment(), and fixed-shape at every page size.""" + if free_index.numel() == 0: + return + self._free_swa_pages(free_index, start_pos=start_pos) + + def _free_swa_pages(self, free_index: torch.Tensor, *, start_pos: int): + ps = self.page_size + assert start_pos % ps == 0, f"segment start {start_pos} is not page-aligned" + # First token of every page the segment touches; the caller allocated + # each one, so a dead entry means the caller wanted free_full. + reps = free_index[::ps] + swa_tokens = self.full_to_swa_index_mapping[reps] + expect(_SWA_PEER_MAPPED, swa_tokens > 0, msg="caller wants free_full") + + if ps == 1: + swa_pages = swa_tokens + mapping_indices = free_index + else: + swa_pages = swa_tokens // ps + # Both pools page in step (alloc_extend / alloc_decode drive them + # with one seq_lens), so a rep's peer page is the whole peer page. + mapping_indices = self._expand_to_full_pages(reps) + if self.swa_attn_allocator.debug_mode: + ref = self.full_to_swa_index_mapping[mapping_indices].cpu() + assert torch.equal( + torch.sort(swa_pages.cpu())[0], + torch.unique(ref[ref > 0] // ps), + ), "swa pages do not match the mapped pages" + self.clear_full_to_swa_mapping(mapping_indices) + + if self.free_group is not None: + # Resolve ownership now, as above. + self.swa_page_ids_group.append(swa_pages) + return + + self.swa_attn_allocator.free_page_ids(swa_pages) + assert self.swa_attn_allocator.available_size() <= self.swa_attn_allocator.size + def _release_swa(self, swa_indices: torch.Tensor): if self.page_size > 1: - # HiCache LOAD_BACK re-pairs a page-aligned full chunk with an offset - # SWA one (commit_hicache_transfer advances by raw token count), so a - # page can hold unmapped slots; one filter per group, not per call. + # Set-shaped frees only (see free_swa): drop the padding-slot entries + # the page expansion picked up. swa_indices = swa_indices[swa_indices > 0] self.swa_attn_allocator.free(swa_indices) assert self.swa_attn_allocator.available_size() <= self.swa_attn_allocator.size @@ -482,7 +522,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): return # SWA first, as in free(): it reads the mapping that a later cache # action in this group may re-point. - self.free_swa(free_index) + self.free_swa_segment(free_index, start_pos=start_pos) self.full_attn_allocator.free_segment(free_index, start_pos=start_pos) def free_full_segment(self, free_index: torch.Tensor, *, start_pos: int): @@ -498,11 +538,16 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): def free_group_begin(self): super().free_group_begin() self.swa_free_group = [] + self.swa_page_ids_group = [] # No full-side pile here: the full allocator's own group defers those. self.full_attn_allocator.free_group_begin() def free_group_end(self): super().free_group_end() + if self.swa_page_ids_group: + swa_page_ids_group = self.swa_page_ids_group + self.swa_page_ids_group = [] + self.swa_attn_allocator.free_page_ids(torch.cat(swa_page_ids_group)) if self.swa_free_group: swa_free_group = self.swa_free_group self.swa_free_group = [] @@ -553,6 +598,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.full_to_swa_index_mapping[:-1].fill_(0) self.free_group = None self.swa_free_group = [] + self.swa_page_ids_group = [] def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): return self._kvcache.get_cpu_copy( @@ -682,6 +728,9 @@ class PureSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): else: self.free_group.append(self._copy_for_free_group(free_index)) + def free_swa_segment(self, free_index: torch.Tensor, *, start_pos: int): + self.free_swa(free_index) + 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. @@ -695,7 +744,7 @@ class PureSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): def free_full_segment(self, free_index: torch.Tensor, *, start_pos: int): return - # Not inherited: the SWA parent's hooks drive swa_free_group and the full + # Not inherited: the SWA parent's hooks drive the SWA piles and the full # allocator's group, which this pure-SWA variant does not have. def free_group_begin(self): BaseTokenToKVPoolAllocator.free_group_begin(self) diff --git a/python/sglang/srt/mem_cache/allocator/token.py b/python/sglang/srt/mem_cache/allocator/token.py index 60aa4af5d..4782a1de2 100644 --- a/python/sglang/srt/mem_cache/allocator/token.py +++ b/python/sglang/srt/mem_cache/allocator/token.py @@ -74,6 +74,10 @@ class TokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): else: self.free_group.append(self._copy_for_free_group(free_index)) + def free_page_ids(self, page_ids: torch.Tensor): + # page_size == 1: page ids are token ids. + self.free(page_ids) + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): return self._kvcache.get_cpu_copy( indices, diff --git a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py index 759c22682..1f199fc69 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py @@ -475,29 +475,13 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): self.full_attn_allocator.clear_inverse_history() self.swa_attn_allocator.clear_inverse_history() - def free_swa( - self, free_index: torch.Tensor, *, start_pos: Optional[int] = None - ) -> None: + def free_swa(self, free_index: torch.Tensor) -> None: """SWA tombstone path: release swa-physical, keep the virtual id and full-physical live; `swa.v2p_page[v_page] = -1` IS the tombstone.""" if free_index is None or free_index.numel() == 0: return v = free_index.detach().to(torch.int64) ps = self.page_size - # `start_pos` promises a contiguous ascending range starting at that prefix - # position, so page reps come from stride arithmetic, not `torch.unique`. - if start_pos is not None and ps > 1: - reps = self.swa_attn_allocator._page_reps(v, start_pos) - # Keep only pages still bound on swa; freeing a tombstoned one would - # corrupt the hole list. `> 0` strict: -1 tombstoned, 0 padding sink. - rep_pages = reps // ps - swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[rep_pages] - live_reps = reps[swa_v2p_pages > 0] - if live_reps.numel() == 0: - return - self.swa_attn_allocator.free(live_reps, _pages=live_reps // ps) - self.swa_attn_allocator.clear_inverse_history() - return v_pages = v // ps # `> 0` strict: -1 = tombstoned, page 0 = padding sink (never freeable). swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages] @@ -512,6 +496,28 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): self.swa_attn_allocator.free(live) self.swa_attn_allocator.clear_inverse_history() + def free_swa_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None: + """free_swa() for a kv-row segment: `start_pos` promises a contiguous + ascending range, so page reps come from stride arithmetic, not `torch.unique`.""" + if free_index is None or free_index.numel() == 0: + return + if self.page_size == 1: + self.free_swa(free_index) + return + ps = self.page_size + reps = self.swa_attn_allocator._page_reps( + free_index.detach().to(torch.int64), start_pos + ) + # Keep only pages still bound on swa; freeing a tombstoned one would + # corrupt the hole list. `> 0` strict: -1 tombstoned, 0 padding sink. + rep_pages = reps // ps + swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[rep_pages] + live_reps = reps[swa_v2p_pages > 0] + if live_reps.numel() == 0: + return + self.swa_attn_allocator.free(live_reps, _pages=live_reps // ps) + 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`).""" diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index ce76960fd..8ac5ca73a 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -101,21 +101,9 @@ def free_swa_out_of_window_slots( free_slots = req_to_token_pool.req_to_token[ req.kv.req_pool_idx, req.kv.swa_evicted_seqlen : new_swa_evicted_seqlen ] - # Local import: the unified allocators import this module lazily for - # eviction; a module-level import here would be a cycle hazard. - from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedSWATokenToKVPoolAllocator, + token_to_kv_pool_allocator.free_swa_segment( + free_slots, start_pos=req.kv.swa_evicted_seqlen ) - - if isinstance(token_to_kv_pool_allocator, UnifiedSWATokenToKVPoolAllocator): - # Contiguous range with host-int bounds: hand the composite its - # start position so the free stays host-sync-free (`free_segment` - # derives page reps by stride math instead of `torch.unique`). - token_to_kv_pool_allocator.free_swa( - free_slots, start_pos=req.kv.swa_evicted_seqlen - ) - else: - token_to_kv_pool_allocator.free_swa(free_slots) req.kv.swa_evicted_seqlen = new_swa_evicted_seqlen diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index f37260171..766e761de 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -599,9 +599,9 @@ class SWARadixCache(BasePrefixCache): # 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) + self.token_to_kv_pool_allocator.free_full_segment(value, start_pos=0) return num_tokens, 0 - self.token_to_kv_pool_allocator.free(value) + self.token_to_kv_pool_allocator.free_segment(value, start_pos=0) return num_tokens, num_tokens def evict(self, params: EvictParams) -> EvictResult: @@ -660,7 +660,9 @@ class SWARadixCache(BasePrefixCache): if len(x.children) > 0: # 1. an internal node, free swa tokens. - self.token_to_kv_pool_allocator.free_swa(x.value) + self.token_to_kv_pool_allocator.free_swa_segment( + x.value, start_pos=0 + ) swa_num_evicted += len(x.value) # 2. get the next node, update the lru lists @@ -673,7 +675,9 @@ class SWARadixCache(BasePrefixCache): # Leaf still holds a full-side lock (can happen when the # SWA leaf-lock early-release optimization revived a # tombstoned leaf. Treat it like an internal tombstone. - self.token_to_kv_pool_allocator.free_swa(x.value) + self.token_to_kv_pool_allocator.free_swa_segment( + x.value, start_pos=0 + ) swa_num_evicted += len(x.value) x_next = self.swa_lru_list.get_prev_no_lock(x) @@ -845,7 +849,9 @@ class SWARadixCache(BasePrefixCache): # swa_lru_list so SWA-eviction won't pick this tombstoned # leaf (which still holds full_lock_ref > 0). The full kv # stays alive until the request releases its full lock. - self.token_to_kv_pool_allocator.free_swa(node.value) + self.token_to_kv_pool_allocator.free_swa_segment( + node.value, start_pos=0 + ) self.swa_lru_list.remove_node(node) node.swa_tombstone = True else: @@ -1202,8 +1208,8 @@ class SWARadixCache(BasePrefixCache): ) else: # Free full tokens in the original tree node. - self.token_to_kv_pool_allocator.free_full( - node.value[:prefix_len] + self.token_to_kv_pool_allocator.free_full_segment( + node.value[:prefix_len], start_pos=0 ) # Overwrite the new value in request to the tree node. node.value = value[:prefix_len].clone() @@ -1220,26 +1226,28 @@ class SWARadixCache(BasePrefixCache): self._recover_tombstone_keeping_locked_full( node, value[start_update_idx:prefix_len] ) - self.token_to_kv_pool_allocator.free_full( - value[:start_update_idx] + self.token_to_kv_pool_allocator.free_full_segment( + value[:start_update_idx], start_pos=0 ) else: - self.token_to_kv_pool_allocator.free_full( - node.value[start_update_idx:prefix_len] + self.token_to_kv_pool_allocator.free_full_segment( + node.value[start_update_idx:prefix_len], start_pos=0 ) self._split_node(node.key, node, start_update_idx) # Here node is the new node after split, so we can overwrite the value to the new node. # The old node is still swa tombstone and the full token is not freed. node.value = value[start_update_idx:prefix_len].clone() - self.token_to_kv_pool_allocator.free_full( - value[:start_update_idx] + self.token_to_kv_pool_allocator.free_full_segment( + value[:start_update_idx], start_pos=0 ) node.swa_tombstone = False self.swa_lru_list.insert_mru(node) self.swa_evictable_size_ += len(node.value) else: # Branch 3: all swa tokens of value[:prefix_len] are evicted, so we don't need to update the node. - self.token_to_kv_pool_allocator.free_full(value[:prefix_len]) + self.token_to_kv_pool_allocator.free_full_segment( + value[:prefix_len], start_pos=0 + ) else: # The node is not tombstone, so we don't need to update the node. # The incoming slice can still straddle this request's own @@ -1275,7 +1283,7 @@ class SWARadixCache(BasePrefixCache): # occurring in normal operation. This check is a defensive guard # against unexpected eviction states from other code paths. if swa_evicted_seqlen == total_prefix_length + len(key): - self.token_to_kv_pool_allocator.free_full(value) + self.token_to_kv_pool_allocator.free_full_segment(value, start_pos=0) return total_prefix_length if ( @@ -1319,7 +1327,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.free_full(incoming_full) + allocator.free_full_segment(incoming_full, start_pos=0) node.swa_tombstone = False self.swa_lru_list.insert_mru(node) 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 0372c6041..01ccb6ece 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 @@ -1436,7 +1436,8 @@ class SWAComponent(TreeComponent): alloc = self.cache.token_to_kv_pool_allocator if isinstance(action, FreeComponentDeviceSlot): for indices in action.indices: - alloc.free_swa(indices) + # Component values are page-aligned copies of a kv row. + alloc.free_swa_segment(indices, start_pos=0) return if isinstance(action, FreeComponentHostSlot): for host_indices in action.host_indices: diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index 10cf0487b..3609eb743 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -299,14 +299,20 @@ class TestSWA(unittest.TestCase): ) def test_free_swa_group_owns_deferred_indices(self): + for page_size in (1, 4): + with self.subTest(page_size=page_size): + self._free_swa_group_owns_deferred_indices(page_size) + + def _free_swa_group_owns_deferred_indices(self, page_size): _, allocator, _ = _build_swa_tree( is_eagle=False, - kv_size=32, - kv_size_swa=32, + page_size=page_size, + kv_size=32 * page_size, + kv_size_swa=32 * page_size, ) index_batches = [] for size in (2, 3, 1, 4): - indices = _swa_alloc(allocator, size) + indices = _swa_alloc(allocator, size * page_size) assert indices is not None index_batches.append(indices) original_indices = torch.cat([indices.clone() for indices in index_batches]) @@ -314,9 +320,10 @@ class TestSWA(unittest.TestCase): available_before_free = allocator.swa_available_size() allocator.free_group_begin() for indices in index_batches: - allocator.free_swa(indices) + allocator.free_swa_segment(indices, start_pos=0) - self.assertEqual(len(allocator.swa_free_group), len(index_batches)) + # The reps were gathered at enqueue time, not from these views. + self.assertEqual(len(allocator.swa_page_ids_group), len(index_batches)) self.assertEqual(allocator.swa_available_size(), available_before_free) for indices in index_batches: indices.zero_() @@ -1144,13 +1151,57 @@ class TestSWAPeerMappedContract(CustomTestCase): def _strict(self): return envs.SGLANG_INVARIANT_CHECK.override(int(InvariantCheckLevel.STRICT)) - def _condition_checked_by(self, allocator, indices): + def _condition_checked_by(self, allocator, indices, start_pos=None): """The predicate free_swa hands the async assert, as a python bool.""" with self._strict(): with mock.patch.object(torch, "_assert_async") as assert_async: - allocator.free_swa(indices) + if start_pos is None: + allocator.free_swa(indices) + else: + allocator.free_swa_segment(indices, start_pos=start_pos) return bool(assert_async.call_args.args[0]) + def test_segment_free_flags_a_page_whose_peer_is_already_gone(self): + _, allocator, _ = _build_swa_tree(is_eagle=False, page_size=4) + live = _swa_alloc(allocator, 8) + stale = _swa_alloc(allocator, 8) + allocator.clear_full_to_swa_mapping(stale) + + self.assertTrue(self._condition_checked_by(allocator, live, start_pos=0)) + self.assertFalse(self._condition_checked_by(allocator, stale, start_pos=0)) + + @unittest.skipUnless(torch.cuda.is_available(), "sync detection needs CUDA") + def test_segment_free_does_not_synchronize_on_pages(self): + """page_size > 1: page reps by stride replace the page expansion's + filter and the inner allocator's torch.unique, in and out of a group.""" + ps = 4 + _, allocator, _ = _build_swa_tree(is_eagle=False, page_size=ps) + + def grouped(indices): + allocator.free_group_begin() + allocator.free_swa_segment(indices, start_pos=0) + allocator.free_group_end() + + # Warm up both paths outside the window: a first-time cudaMalloc can + # synchronize on its own, which the detector would blame on this call. + allocator.free_swa_segment(_swa_alloc(allocator, 2 * ps), start_pos=0) + grouped(_swa_alloc(allocator, 2 * ps)) + first = _swa_alloc(allocator, 3 * ps) + second = _swa_alloc(allocator, 2 * ps) + + # Gate on the pre-fix form: a detector blind to this sync class would pass + # the asserts below no matter how free_swa derives the pages. + if _sync_error(lambda: torch.unique(first // ps)) is None: + self.skipTest("sync debug mode does not flag a data-dependent shape here") + + with self._strict(): + self.assertIsNone( + _sync_error( + lambda: allocator.free_swa_segment(first[: 3 * ps - 1], start_pos=0) + ) + ) + self.assertIsNone(_sync_error(lambda: grouped(second[: 2 * ps - 1]))) + def test_free_swa_flags_a_slot_whose_peer_is_already_gone(self): _, allocator, _ = _build_swa_tree(is_eagle=False) live = _swa_alloc(allocator, 4) @@ -1183,6 +1234,69 @@ class TestSWAPeerMappedContract(CustomTestCase): self.assertIsNone(_sync_error(lambda: allocator.free_swa(indices))) +class TestSWAPageRepsFree(CustomTestCase): + """page_size > 1: with a start position the SWA side frees one representative + per page instead of expanding, filtering and dedup'ing through torch.unique.""" + + PS = 4 + + def _allocator(self): + _, allocator, _ = _build_swa_tree(is_eagle=False, page_size=self.PS) + return allocator + + def _sizes(self, allocator): + return allocator.full_available_size(), allocator.swa_available_size() + + def test_segment_free_releases_the_mapped_pages_for_every_tail(self): + ps = self.PS + for num_tokens in (1, ps, ps + 1, 3 * ps - 1, 3 * ps): + with self.subTest(num_tokens=num_tokens): + allocator = self._allocator() + indices = _swa_alloc(allocator, 3 * ps) + mapping = allocator.full_to_swa_index_mapping + expected = torch.unique(mapping[indices[:num_tokens]] // ps) + before = allocator.swa_attn_allocator.free_pages.numel() + + allocator.free_swa_segment(indices[:num_tokens], start_pos=0) + + free_pages = allocator.swa_attn_allocator.free_pages + freed = free_pages[: free_pages.numel() - before] + self.assertTrue(torch.equal(torch.sort(freed)[0], expected)) + # The whole last page goes back, and its mapping with it. + touched = -(num_tokens // -ps) * ps + self.assertTrue(torch.all(mapping[indices[:touched]] == 0)) + self.assertTrue(torch.all(mapping[indices[touched:]] > 0)) + + def test_node_frees_take_the_page_path_through_the_tree(self): + """Tree values are page-aligned copies of a kv row, so SWA eviction and + the full eviction of its tombstones both free by page reps.""" + ps = self.PS + tree, allocator, _ = _build_swa_tree( + is_eagle=False, page_size=ps, sliding_window_size=ps + ) + full_before, swa_before = self._sizes(allocator) + _insert(tree, allocator, list(range(1, 3 * ps + 1))) + + # Either inner `free` is the torch.unique path a caller falls back to + # when it hands no start position. + with ( + patch.object( + allocator.full_attn_allocator, + "free", + side_effect=AssertionError("full side took the unique path"), + ), + patch.object( + allocator.swa_attn_allocator, + "free", + side_effect=AssertionError("swa side took the unique path"), + ), + ): + tree.evict(EvictParams(num_tokens=0, swa_num_tokens=ps)) + tree.evict(EvictParams(num_tokens=3 * ps, swa_num_tokens=0)) + + self.assertEqual(self._sizes(allocator), (full_before, swa_before)) + + class TestCacheUnfinishedReqEvictedPrefix(CustomTestCase): """An unfinished request whose SWA prefix is already gone must insert that prefix as a tombstone, not as live SWA KV.""" diff --git a/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py b/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py index 198acd505..cc9453680 100644 --- a/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py +++ b/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py @@ -373,7 +373,7 @@ class TestUnifiedSwaFullSideGroup(unittest.TestCase): class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase): """The per-decode-step SWA window ratchet frees a CONTIGUOUS row slice with - host-int, page-aligned bounds, so `free_swa(..., start_pos=)` must reach the + host-int, page-aligned bounds, so `free_swa_segment` must reach the swa side with caller-derived page ids: no `torch.unique` and no stale-slot `.item()` on the per-step path. """ @@ -429,15 +429,15 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase): torch.Tensor, "item", side_effect=AssertionError("item = host sync") ), ): - alloc.free_swa(v[: 4 * self.PS], start_pos=0) - alloc.free_swa(v[4 * self.PS :], start_pos=4 * self.PS) + alloc.free_swa_segment(v[: 4 * self.PS], start_pos=0) + alloc.free_swa_segment(v[4 * self.PS :], start_pos=4 * self.PS) def test_full_only_segment_free_never_syncs(self): """Request-finish shape: the swa side is already tombstoned, so the full side must free by page reps rather than `free_full`'s dedup.""" alloc = self._swa_composite(lazy=True) v = alloc.alloc(8 * self.PS) - alloc.free_swa(v, start_pos=0) + alloc.free_swa_segment(v, start_pos=0) before = alloc.full_available_size() with ( mock.patch.object( @@ -456,7 +456,7 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase): alloc = self._swa_composite(lazy=True) v = alloc.alloc(8 * self.PS) with self.assertRaises(AssertionError): - alloc.free_swa(v[1 : 5 * self.PS], start_pos=1) + alloc.free_swa_segment(v[1 : 5 * self.PS], start_pos=1) def test_start_pos_path_matches_the_fallback_end_state(self): """Derived property: the stride-rep path and the dedup fallback leave @@ -468,7 +468,7 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase): v1 = a1.alloc(6 * self.PS) v2 = a2.alloc(6 * self.PS) self.assertTrue(torch.equal(v1, v2)) - a1.free_swa(v1[: 4 * self.PS], start_pos=0) + a1.free_swa_segment(v1[: 4 * self.PS], start_pos=0) a2.free_swa(v2[: 4 * self.PS]) # fallback (radix shape) self.assertTrue( torch.equal( @@ -487,8 +487,8 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase): liveness filter (radix eviction and the ratchet can overlap).""" alloc = self._swa_composite(lazy=True) v = alloc.alloc(4 * self.PS) - alloc.free_swa(v, start_pos=0) - alloc.free_swa(v, start_pos=0) # all tombstoned -> filtered to empty + alloc.free_swa_segment(v, start_pos=0) + alloc.free_swa_segment(v, start_pos=0) # all tombstoned -> filtered to empty @unittest.skipUnless( 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 df4e258f1..0ced1fc84 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 @@ -7860,13 +7860,15 @@ class TestUnifiedRadixCacheActionRouting(CustomTestCase): indices, start_pos=0 ) - def test_apply_component_action_device_kv_swa_uses_free_swa(self): + def test_apply_component_action_device_kv_swa_uses_free_swa_segment(self): cache = mock.MagicMock() indices = torch.tensor([4, 5]) _component_with_cache(ComponentType.SWA, cache).apply_component_action( FreeComponentDeviceSlot([indices], component_type=ComponentType.SWA) ) - cache.token_to_kv_pool_allocator.free_swa.assert_called_once_with(indices) + cache.token_to_kv_pool_allocator.free_swa_segment.assert_called_once_with( + indices, start_pos=0 + ) def test_apply_component_action_device_kv_mamba_uses_mamba_allocator(self): cache = mock.MagicMock() diff --git a/test/registered/unit/mem_cache/test_unified_tri_pool.py b/test/registered/unit/mem_cache/test_unified_tri_pool.py index 40c8100db..371c752bc 100644 --- a/test/registered/unit/mem_cache/test_unified_tri_pool.py +++ b/test/registered/unit/mem_cache/test_unified_tri_pool.py @@ -407,7 +407,7 @@ class TestTriFreeSwaNoHostSync(unittest.TestCase): torch.Tensor, "item", side_effect=AssertionError("item = host sync") ), ): - alloc.free_swa(v[: 4 * self.PS], start_pos=0) + alloc.free_swa_segment(v[: 4 * self.PS], start_pos=0) self.assertEqual(alloc.verify_byte_accounting(), []) def test_fallback_free_swa_still_correct_for_radix_shapes(self): @@ -416,7 +416,7 @@ class TestTriFreeSwaNoHostSync(unittest.TestCase): a1, a2 = self._tri(), self._tri() v1, v2 = a1.alloc(6 * self.PS), a2.alloc(6 * self.PS) self.assertTrue(torch.equal(v1, v2)) - a1.free_swa(v1[: 4 * self.PS], start_pos=0) + a1.free_swa_segment(v1[: 4 * self.PS], start_pos=0) a2.free_swa(v2[: 4 * self.PS]) self.assertTrue( torch.equal( @@ -715,7 +715,7 @@ class TestTriDeferredAbsorption(unittest.TestCase): v = alloc.alloc(8 * self.PS) sa = alloc.swa_attn_allocator span = sa._span_pages() - alloc.free_swa(v[6 * self.PS :], start_pos=6 * self.PS) # high edge + alloc.free_swa_segment(v[6 * self.PS :], start_pos=6 * self.PS) # high edge self.assertGreater(sa._hole_pages(), 0) # deferred self.assertEqual(sa._span_pages(), span) moved = alloc.flush_opportunistic() @@ -729,7 +729,7 @@ class TestTriDeferredAbsorption(unittest.TestCase): alloc = self._tri() v = alloc.alloc(8 * self.PS) sa = alloc.swa_attn_allocator - alloc.free_swa(v[6 * self.PS :], start_pos=6 * self.PS) + alloc.free_swa_segment(v[6 * self.PS :], start_pos=6 * self.PS) self.assertGreater(sa._hole_pages(), 0) moves_before = len(sa._inverse_history) from sglang.srt.mem_cache.allocator.unified_sub_pool import _relieve_for_alloc @@ -743,7 +743,7 @@ class TestTriDeferredAbsorption(unittest.TestCase): value -- under-reporting is safe, over-reporting would over-admit.""" alloc = self._tri() v = alloc.alloc(8 * self.PS) - alloc.free_swa(v[6 * self.PS :], start_pos=6 * self.PS) + alloc.free_swa_segment(v[6 * self.PS :], start_pos=6 * self.PS) deferred = alloc.available_size() alloc.swa_attn_allocator._flush(urgent=False) absorbed = alloc.available_size() @@ -759,7 +759,7 @@ class TestTriDeferredAbsorption(unittest.TestCase): alloc = self._tri() v = alloc.alloc(8 * self.PS) - alloc.free_swa(v[2 * self.PS : 4 * self.PS], start_pos=2 * self.PS) + alloc.free_swa_segment(v[2 * self.PS : 4 * self.PS], start_pos=2 * self.PS) alloc.flush_opportunistic() # consumes the dirty flag sa = alloc.swa_attn_allocator self.assertGreater(sa._hole_pages(), 0) # interior holes remain @@ -775,10 +775,10 @@ class TestTriDeferredAbsorption(unittest.TestCase): alloc = self._tri() v = alloc.alloc(8 * self.PS) sa = alloc.swa_attn_allocator - alloc.free_swa(v[: 2 * self.PS], start_pos=0) # low-edge holes + alloc.free_swa_segment(v[: 2 * self.PS], start_pos=0) # low-edge holes n_after_free = sa._hole_pages() alloc.alloc(2 * self.PS) # drains them back to live - alloc.free_swa(v[6 * self.PS :], start_pos=6 * self.PS) # high edge + alloc.free_swa_segment(v[6 * self.PS :], start_pos=6 * self.PS) # high edge self.assertEqual(sa._hole_pages(), n_after_free) # same COUNT as before span = sa._span_pages() self.assertGreater(alloc.flush_opportunistic(), 0) # still absorbed @@ -797,7 +797,7 @@ class TestTriDeferredAbsorption(unittest.TestCase): with mock.patch.object( torch.Tensor, "tolist", side_effect=AssertionError("tolist = D2H") ): - alloc.free_swa(v, start_pos=0) + alloc.free_swa_segment(v, start_pos=0) self.assertTrue(sa._is_frontier_transparent()) self.assertEqual(sa._hole_pages(), 0)