[sgl] Window-aware LRU refresh for SWA prefix cache in unified cache

python/sglang/srt/mem_cache/unified_cache_components/swa_component.py

@@ -60,6 +60,29 @@ def _translate_full_to_swa(self, full_indices: torch.Tensor) -> torch.Tensor:
             full_indices
         )
 
+    def refresh_lru_on_walkdown(self, node: UnifiedTreeNode) -> None:
+        return
+
+    def _bounded_refresh(
+        self, last_node: UnifiedTreeNode, root_node: UnifiedTreeNode
+    ) -> None:
+        self.cache.lru_lists[self.component_type].reset_node_and_window_ancestors_mru(
+            last_node,
+            root_node,
+            self.sliding_window_size + self.cache.page_size,
+            self.node_has_component_data,
+        )
+
+    def refresh_lru_on_match_end(
+        self, last_node: UnifiedTreeNode, root_node: UnifiedTreeNode
+    ) -> None:
+        self._bounded_refresh(last_node, root_node)
+
+    def refresh_lru_on_insert_end(
+        self, target_node: UnifiedTreeNode, root_node: UnifiedTreeNode
+    ) -> None:
+        self._bounded_refresh(target_node, root_node)
+
     def _restore_device_value(self, node: UnifiedTreeNode, value: torch.Tensor) -> None:
         ct = self.component_type
         node.component_data[ct].value = value

python/sglang/srt/mem_cache/unified_cache_components/tree_component.py

@@ -115,6 +115,24 @@ def value_len(self, node: UnifiedTreeNode) -> int:
         value = node.component_data[self.component_type].value
         return len(value) if value is not None else 0
 
+    def refresh_lru_on_walkdown(self, node: UnifiedTreeNode) -> None:
+        cd = node.component_data[self.component_type]
+        if cd.value is None:
+            return
+        self.cache.lru_lists[self.component_type].reset_node_mru(node)
+
+    def refresh_lru_on_match_end(
+        self, last_node: UnifiedTreeNode, root_node: UnifiedTreeNode
+    ) -> None:
+        self.cache.lru_lists[self.component_type].reset_node_and_parents_mru(
+            last_node, root_node, self.node_has_component_data
+        )
+
+    def refresh_lru_on_insert_end(
+        self, target_node: UnifiedTreeNode, root_node: UnifiedTreeNode
+    ) -> None:
+        return
+
     @abstractmethod
     def create_match_validator(
         self, match_device_only: bool = False

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -179,6 +179,24 @@ def reset_node_and_parents_mru(
                 prev_node = node
             node = node.parent
 
+    def reset_node_and_window_ancestors_mru(
+        self,
+        node: UnifiedTreeNode,
+        root_node: UnifiedTreeNode,
+        window_size: int,
+        should_include,
+    ):
+        prev_node = self.head
+        accumulated = 0
+        while node != root_node and accumulated < window_size:
+            if should_include(node):
+                assert node.id in self.cache
+                self._remove_node(node)
+                self._add_node_after(prev_node, node)
+                prev_node = node
+            accumulated += len(node.key)
+            node = node.parent
+
     def in_list(self, node: Optional[UnifiedTreeNode]):
         return node is not None and node.id in self.cache
 
@@ -779,9 +797,7 @@ def _match_post_processor(
         for comp in self._components_tuple:
             if comp.component_type == BASE_COMPONENT_TYPE:
                 continue  # Full uses last_access_time, not LRU
-            self.lru_lists[comp.component_type].reset_node_and_parents_mru(
-                node_update, self.root_node, comp.node_has_component_data
-            )
+            comp.refresh_lru_on_match_end(node_update, self.root_node)
 
         cur_time = get_and_increase_time_counter()
         while node_update:
@@ -855,7 +871,10 @@ def _split_node(
     def _touch_node(self, node: UnifiedTreeNode):
         node.last_access_time = get_and_increase_time_counter()
         if node != self.root_node:
-            self._for_each_component_lru(node, UnifiedLRUList.reset_node_mru)
+            for comp in self._components_tuple:
+                if comp.component_type == BASE_COMPONENT_TYPE:
+                    continue
+                comp.refresh_lru_on_walkdown(node)
 
     def _add_new_node(
         self,
@@ -985,6 +1004,13 @@ def _insert_helper(
                 params=params,
                 result=result,
             )
+
+        if target_node is not self.root_node:
+            for component in self._components_tuple:
+                if component.component_type == BASE_COMPONENT_TYPE:
+                    continue
+                component.refresh_lru_on_insert_end(target_node, self.root_node)
+
         if is_new_leaf:
             self._inc_hit_count(target_node, params.chunked)
         return result