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

Co-authored-by: Zhangheng <hzh0425@apache.org>
Co-authored-by: Hanming Lu <69857889+hanming-lu@users.noreply.github.com>

python/sglang/srt/mem_cache/allocator/swa.py

@@ -331,11 +331,8 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         if full_indices.numel() == 0:
             return
         assert full_indices.numel() == swa_indices.numel()
-        if _is_npu:
+        full_indices = full_indices.to(torch.int64)
-            self.full_to_swa_index_mapping[full_indices.to(torch.int64)] = (
+        swa_indices = swa_indices.to(self.full_to_swa_index_mapping.dtype)
-                swa_indices.to(torch.int64)
-            )
-        else:
         self.full_to_swa_index_mapping[full_indices] = swa_indices
 
     def free_swa(self, free_index: torch.Tensor):

python/sglang/srt/mem_cache/swa_radix_cache.py

@@ -1138,8 +1138,18 @@ 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.
+                        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])
+                            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
@@ -1148,6 +1158,17 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
                     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
+                        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]
                             )
@@ -1155,7 +1176,9 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
                             # 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])
+                            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)
@@ -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

@@ -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(
+            if full_cd.lock_ref > 0:
-                node.component_data[BASE_COMPONENT_TYPE].value
+                self._restore_device_value_with_locked_full(
-            )
+                    node, full_cd.value, value_slice
-            node.component_data[BASE_COMPONENT_TYPE].value = value_slice.clone()
-            swa_value = self._translate_full_to_swa(
-                node.component_data[BASE_COMPONENT_TYPE].value
                 )
+                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(
+            if full_cd.lock_ref > 0:
-                node.component_data[BASE_COMPONENT_TYPE].value[start_idx:]
+                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:

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