minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[mem_cache] Free hybrid SWA pages by one representative per page on page_size > 1 (#38159)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-09-07 20:15:54 -07:00
committed by GitHub
parent b23d835048
commit 28ebede865
12 changed files with 276 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/allocator/hisparse.py
+3
View File
@@ -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
python/sglang/srt/mem_cache/allocator/paged.py
+14 -9
View File
@@ -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
python/sglang/srt/mem_cache/allocator/swa.py
+62 -13
View File
@@ -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)
python/sglang/srt/mem_cache/allocator/token.py
+4
View File
@@ -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,
python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py
+23 -17
View File
@@ -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`)."""
python/sglang/srt/mem_cache/common.py
+2 -14
View File
@@ -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
python/sglang/srt/mem_cache/swa_radix_cache.py
+24 -16
View File
@@ -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)
python/sglang/srt/mem_cache/unified_cache/components/swa_component.py
+2 -1
View File
@@ -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: