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

[bug2] skip swa recovery on locked full kv (#29352)

Browse Source 
Co-authored-by: Zhangheng <hzh0425@apache.org>
Co-authored-by: Hanming Lu <69857889+hanming-lu@users.noreply.github.com>
This commit is contained in:
Feng Yao
2026-06-30 16:07:06 -07:00
committed by GitHub
co-authored by Zhangheng Hanming Lu
parent 99b8f36cb1
commit 53c61bd5e6
4 changed files with 167 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/allocator/swa.py
+3 -6
View File
@@ -331,12 +331,9 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
if full_indices.numel() == 0:
return
assert full_indices.numel() == swa_indices.numel()
if _is_npu:
self.full_to_swa_index_mapping[full_indices.to(torch.int64)] = (
swa_indices.to(torch.int64)
)
else:
self.full_to_swa_index_mapping[full_indices] = swa_indices
full_indices = full_indices.to(torch.int64)
swa_indices = swa_indices.to(self.full_to_swa_index_mapping.dtype)
self.full_to_swa_index_mapping[full_indices] = swa_indices
def free_swa(self, free_index: torch.Tensor):
if free_index.numel() == 0:
python/sglang/srt/mem_cache/swa_radix_cache.py
+65 -18
View File
@@ -1138,27 +1138,50 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
), f"swa_evicted_seqlen must be page aligned, {swa_evicted_seqlen=}, {self.page_size=}"
if swa_evicted_seqlen <= total_prefix_length:
# Branch 1: all swa tokens of value[:prefix_len] are not evicted, so we can insert it to the tree directly.
# Free full tokens in the original tree node.
self.token_to_kv_pool_allocator.free(node.value[:prefix_len])
# Overwrite the new value in request to the tree node.
node.value = value[:prefix_len].clone()
node.swa_tombstone = False
self.swa_lru_list.insert_mru(node)
self.swa_evictable_size_ += len(node.value)
if node.full_lock_ref > 0:
# Full KV is still locked by a running request. Keep it
# and adopt the incoming SWA instead of freeing in-flight
# Full slots.
self._recover_tombstone_keeping_locked_full(
node, value[:prefix_len]
)
else:
# Free full tokens in the original tree node.
self.token_to_kv_pool_allocator.free(
node.value[:prefix_len]
)
# Overwrite the new value in request to the tree node.
node.value = value[:prefix_len].clone()
node.swa_tombstone = False
self.swa_lru_list.insert_mru(node)
self.swa_evictable_size_ += len(node.value)
elif swa_evicted_seqlen < total_prefix_length + prefix_len:
# Branch 2: part of swa tokens of value[:prefix_len] are evicted, so we need to split the node and insert the value to new node.
start_update_idx = swa_evicted_seqlen - total_prefix_length
self.token_to_kv_pool_allocator.free(
node.value[start_update_idx:prefix_len]
)
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(value[:start_update_idx])
node.swa_tombstone = False
self.swa_lru_list.insert_mru(node)
self.swa_evictable_size_ += len(node.value)
if node.full_lock_ref > 0:
# Split first so the recovered suffix keeps the locked
# Full slots, then adopt the incoming SWA for that suffix.
self._split_node(node.key, node, start_update_idx)
self._recover_tombstone_keeping_locked_full(
node, value[start_update_idx:prefix_len]
)
self.token_to_kv_pool_allocator.free(
value[:start_update_idx]
)
else:
self.token_to_kv_pool_allocator.free(
node.value[start_update_idx:prefix_len]
)
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(
value[:start_update_idx]
)
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(value[:prefix_len])
@@ -1217,6 +1240,30 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
return total_prefix_length
def _recover_tombstone_keeping_locked_full(
self, node: TreeNode, incoming_full: torch.Tensor
) -> None:
"""Recover a tombstoned node whose Full KV is locked by a running request.
Keep node.value, the locked Full slots, and re-point its full->SWA
mapping at the incoming request's fresh SWA. Free only the incoming
redundant Full slots, not their SWA slots.
"""
assert len(node.value) == len(incoming_full), (
f"locked-full recover size mismatch: {len(node.value)=}, "
f"{len(incoming_full)=}"
)
allocator = self.token_to_kv_pool_allocator
swa_value = allocator.translate_loc_from_full_to_swa(incoming_full)
allocator.set_full_to_swa_mapping(node.value, swa_value)
allocator.full_to_swa_index_mapping[incoming_full.to(torch.int64)] = 0
allocator.full_attn_allocator.free(incoming_full)
node.swa_tombstone = False
self.swa_lru_list.insert_mru(node)
self.swa_evictable_size_ += len(node.value)
def _add_new_node(
self,
parent: TreeNode,
python/sglang/srt/mem_cache/unified_cache_components/swa_component.py
+30 -10
View File
@@ -101,6 +101,19 @@ class SWAComponent(TreeComponent):
self.cache.lru_lists[ct].insert_mru(node)
self.cache.component_evictable_size_[ct] += len(value)
def _restore_device_value_with_locked_full(
self,
node: UnifiedTreeNode,
full_value: torch.Tensor,
incoming_full_value: torch.Tensor,
) -> None:
allocator = self.cache.token_to_kv_pool_allocator
swa_value = self._translate_full_to_swa(incoming_full_value)
allocator.set_full_to_swa_mapping(full_value, swa_value)
allocator.full_to_swa_index_mapping[incoming_full_value.to(torch.int64)] = 0
allocator.full_attn_allocator.free(incoming_full_value)
self._restore_device_value(node, swa_value)
def create_match_validator(
self, match_device_only: bool = False
) -> Callable[[UnifiedTreeNode], bool]:
@@ -169,6 +182,7 @@ class SWAComponent(TreeComponent):
if not is_tombstone:
return prefix_len
full_cd = node.component_data[BASE_COMPONENT_TYPE]
swa_evicted_seqlen = params.swa_evicted_seqlen
assert (
node.component_data[self.component_type].lock_ref == 0
@@ -179,21 +193,27 @@ class SWAComponent(TreeComponent):
if swa_evicted_seqlen <= total_prefix_len:
# Branch 1: entire value_slice is within SWA window — recover
self.cache.token_to_kv_pool_allocator.free(
node.component_data[BASE_COMPONENT_TYPE].value
)
node.component_data[BASE_COMPONENT_TYPE].value = value_slice.clone()
swa_value = self._translate_full_to_swa(
node.component_data[BASE_COMPONENT_TYPE].value
)
if full_cd.lock_ref > 0:
self._restore_device_value_with_locked_full(
node, full_cd.value, value_slice
)
return 0
self.cache.token_to_kv_pool_allocator.free(full_cd.value)
full_cd.value = value_slice.clone()
swa_value = self._translate_full_to_swa(full_cd.value)
self._restore_device_value(node, swa_value)
return 0
elif swa_evicted_seqlen < total_prefix_len + prefix_len:
# Branch 2: value_slice[start_idx:] is within SWA window — partial recover
start_idx = swa_evicted_seqlen - total_prefix_len
self.cache.token_to_kv_pool_allocator.free(
node.component_data[BASE_COMPONENT_TYPE].value[start_idx:]
)
if full_cd.lock_ref > 0:
self.cache._split_node(node.key, node, start_idx)
full_cd = node.component_data[BASE_COMPONENT_TYPE]
self._restore_device_value_with_locked_full(
node, full_cd.value, value_slice[start_idx:]
)
return start_idx
self.cache.token_to_kv_pool_allocator.free(full_cd.value[start_idx:])
self.cache._split_node(node.key, node, start_idx)
node.component_data[BASE_COMPONENT_TYPE].value = value_slice[
start_idx: