diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index 8e27b731a..3dce815f3 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -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: diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index c690e9636..ea898c759 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -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, 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 ff161d68d..35801c706 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 @@ -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: 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 47638f54e..ae1d48df0 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 @@ -1287,6 +1287,75 @@ class UnifiedRadixCacheSuite: self.assertEqual(len(m.device_indices), len(seq)) tree.sanity_check() + def test_swa_unfinished_recovery_preserves_locked_full_value(self): + if not self.cfg.has_swa or self.cfg.has_mamba: + self.skipTest("requires SWA without Mamba") + if self.cfg.page_size != 1 or self.cfg.sliding_window_size != 4: + self.skipTest("requires page_size=1, sliding_window_size=4") + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + + tokens = self._make_seq(1, 4) + self._insert(tree, allocator, req_to_token_pool, tokens) + self._insert( + tree, allocator, req_to_token_pool, tokens + self._make_seq(100, 1) + ) + + node = tree.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens))) + ).last_device_node + old_full_value = node.component_data[ComponentType.FULL].value.clone() + swa_component = tree.components[ComponentType.SWA] + tracker = {ct: 0 for ct in tree.tree_components} + tree._evict_component_and_detach_lru(node, swa_component, tracker=tracker) + self.assertIsNone(node.component_data[ComponentType.SWA].value) + + lock_result = tree.inc_lock_ref(node) + req = self._make_req(req_to_token_pool) + req.origin_input_ids = array("q", tokens) + req.output_ids = [] + req.full_untruncated_fill_ids = array("q", tokens) + req.set_extend_range(0, len(req.full_untruncated_fill_ids)) + kv_len = len(tokens) + fresh_value = self._alloc(allocator, kv_len) + req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), fresh_value) + req.kv_committed_len = kv_len + req.last_node = tree.root_node + req.cache_protected_len = 0 + req.swa_uuid_for_lock = None + req.extra_key = None + req.swa_evicted_seqlen = 0 + + full_available_before_insert = allocator.full_attn_allocator.available_size() + + tree.cache_unfinished_req(req) + + self.assertEqual( + allocator.full_attn_allocator.available_size(), + full_available_before_insert + len(tokens), + ) + self.assertTrue( + torch.equal( + node.component_data[ComponentType.FULL].value, + old_full_value, + ) + ) + swa_value = node.component_data[ComponentType.SWA].value + self.assertIsNotNone(swa_value) + self.assertTrue( + torch.equal( + allocator.translate_loc_from_full_to_swa(old_full_value), + swa_value, + ) + ) + self.assertEqual(req.cache_protected_len, len(tokens)) + + tree.dec_lock_ref( + req.last_node, + DecLockRefParams(swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None)), + ) + tree.dec_lock_ref(node, lock_result.to_dec_params()) + tree.sanity_check() + def test_swa_insert_keeps_full_leaf_when_entire_span_is_outside_window(self): # A leaf survives on its Full value alone: even when the whole span is # past the SWA window (swa_evicted_seqlen >= total), the Full leaf must