[Unified Tree] Port SWA Branching-Point Caching to the Rust TreeCore (#37584)

python/sglang/srt/mem_cache/rust_tree_core/adapter.py

@@ -228,6 +228,7 @@ def _insert_step_from_binding(step) -> InsertStepResult:
             prefix_len=step.result.prefix_len,
             last_device_node=step.result.last_device_node,
             mamba_exist=step.result.mamba_exist,
+            swa_branch_inserted=step.result.swa_branch_inserted,
             host_insert_dropped=step.result.host_insert_dropped,
             adopted_ranges=(
                 {
@@ -252,6 +253,7 @@ def _match_result_from_binding(result) -> MatchResult:
         best_match_node=result.best_match_node_id,
         host_hit_length=result.host_hit_length,
         swa_host_hit_length=result.swa_host_hit_length,
+        swa_branching_seqlen=result.swa_branching_seqlen,
         mamba_host_hit_length=result.mamba_host_hit_length,
         mamba_branching_seqlen=result.mamba_branching_seqlen,
         full_kv_hit_length=result.full_kv_hit_length,
@@ -589,6 +591,7 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
                 mamba_value=params.mamba_value,
                 prev_prefix_len=params.prev_prefix_len,
                 swa_evicted_seqlen=params.swa_evicted_seqlen,
+                swa_branching_seqlen=params.swa_branching_seqlen,
                 chunked=params.chunked,
                 priority=0 if params.priority is None else params.priority,
                 track_adopted_ranges=params.track_adopted_ranges,
@@ -630,6 +633,13 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
     def set_hicache_enabled(self) -> None:
         self._binding.set_hicache_enabled()
 
+    def set_host_memory_buffer_only(self) -> None:
+        self._binding.set_host_memory_buffer_only()
+
+    @property
+    def is_host_memory_buffer_only(self) -> bool:
+        return self._binding.is_host_memory_buffer_only()
+
     @property
     def page_size(self) -> int:
         # Read-only: the Rust core freezes it at construction.

python/sglang/srt/mem_cache/unified_cache/components/swa_component.py

@@ -89,13 +89,27 @@ class SWAComponent(TreeComponent):
 
     component_type = ComponentType.SWA
 
-    def _dirty_backup_window(self, node: UnifiedTreeNode) -> list[UnifiedTreeNode]:
+    def _collect_unbacked_swa_nodes(
+        self, node: UnifiedTreeNode
+    ) -> list[UnifiedTreeNode]:
+        """Nodes whose SWA data needs a host backup, deepest first.
+
+        Buffer mode stages one node per FIFO backup intent; cache mode backs
+        up every device-only node within one sliding window of ``node``.
+        """
         if not self.tree_core.has_swa_host_pool:
             return []
+        if self.tree_core.is_host_memory_buffer_only:
+            cd = node.component_data[self.component_type]
+            return [node] if cd.value is not None else []
+        return self._collect_unbacked_swa_nodes_in_window(node)
 
+    def _collect_unbacked_swa_nodes_in_window(
+        self, node: UnifiedTreeNode
+    ) -> list[UnifiedTreeNode]:
         ct = self.component_type
         covered = 0
-        dirty: list[UnifiedTreeNode] = []
+        unbacked: list[UnifiedTreeNode] = []
         cur = node
         while (
             cur is not self.tree_core.root_node and covered < self.sliding_window_size
@@ -109,12 +123,12 @@ class SWAComponent(TreeComponent):
                 break
             covered += len(value)
             if cd.value is not None and cd.host_value is None:
-                dirty.append(cur)
+                unbacked.append(cur)
             cur = cur.parent
-        return dirty
+        return unbacked
 
     def needs_incremental_backup(self, node: UnifiedTreeNode) -> bool:
-        return bool(self._dirty_backup_window(node))
+        return bool(self._collect_unbacked_swa_nodes(node))
 
     def reset_session_state(self) -> None:
         super().reset_session_state()
@@ -919,6 +933,11 @@ class SWAComponent(TreeComponent):
         # that boundary so insertion creates a tombstone instead of live SWA KV.
         insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen
 
+        # A recurrent checkpoint must stay attached to its exact token prefix.
+        # Let MambaComponent select the insertion length for hybrid caches.
+        if self.cache.is_mamba_enabled:
+            return None
+
         branching_seqlen = req.swa_branching_seqlen
         if branching_seqlen is None or branching_seqlen <= req.kv.cache_protected_len:
             return None
@@ -991,8 +1010,7 @@ class SWAComponent(TreeComponent):
         elif prefetch_pages <= 0:
             return PreparePrefetchResult()
         elif (
-            self.tree_core.is_root(node_id)
-            or self.cache.host_memory_mode == "buffer_only"
+            self.tree_core.is_root(node_id) or self.tree_core.is_host_memory_buffer_only
         ):
             # Sub-window fetch: at root the sequence IS its window; mid-tree
             # (buffer mode) the window head is the device prefix's own ring
@@ -1030,22 +1048,17 @@ class SWAComponent(TreeComponent):
             return None
 
         if phase == CacheTransferPhase.BACKUP_HOST:
-            if self.cache.host_memory_mode == "buffer_only":
-                # Buffer mode stages one node/hash span per FIFO backup intent.
-                cd = node.component_data[ct]
-                dirty = [node] if cd.value is not None else []
-            else:
-                dirty = self._dirty_backup_window(node)
-            if not dirty:
+            unbacked_swa_nodes = self._collect_unbacked_swa_nodes(node)
+            if not unbacked_swa_nodes:
                 return None
-            dirty.reverse()
+            unbacked_swa_nodes.reverse()
             return [
                 PoolTransfer(
                     name=PoolName.SWA,
                     device_indices=torch.cat(
-                        [n.component_data[ct].value for n in dirty]
+                        [n.component_data[ct].value for n in unbacked_swa_nodes]
                     ).to(torch.int64),
-                    nodes_to_load=[n.id for n in dirty],
+                    nodes_to_load=[n.id for n in unbacked_swa_nodes],
                 )
             ]
 

python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py

@@ -415,6 +415,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
         self.page_size = params.page_size
         self.is_eagle = params.is_eagle and ComponentType.MAMBA not in components
         self.enable_hicache = False
+        self.is_host_memory_buffer_only = False
         self.enable_storage = False
         self.enable_external_cache_linker = False
         self.write_through_threshold = 256
@@ -2045,6 +2046,9 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
     def set_hicache_enabled(self) -> None:
         self.enable_hicache = True
 
+    def set_host_memory_buffer_only(self) -> None:
+        self.is_host_memory_buffer_only = True
+
     def insert_host(
         self,
         node_id: NodeId,

python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py

@@ -152,6 +152,8 @@ class UnifiedTreeCoreInterface(ABC):
     write_through_threshold: int
     is_write_back: bool
     has_swa_host_pool: bool
+    # Whether the host tier stages one node per FIFO backup intent.
+    is_host_memory_buffer_only: bool
     kv_events: KVCacheEventRecorder
 
     # ==== Tree API ====
@@ -456,6 +458,11 @@ class UnifiedTreeCoreInterface(ABC):
         """Mark the host tier (HiCache) as wired."""
         ...
 
+    @abstractmethod
+    def set_host_memory_buffer_only(self) -> None:
+        """Mark the host tier as buffer-only: one node staged per backup intent."""
+        ...
+
     @abstractmethod
     def insert_host(
         self,

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -456,6 +456,7 @@ class UnifiedRadixCache(BasePrefixCache):
                 self.tree_core.has_swa_host_pool = swa._swa_kv_pool_host is not None
 
         if self.host_memory_mode == "buffer_only":
+            self.tree_core.set_host_memory_buffer_only()
             swa = self.components.get(ComponentType.SWA)
             validate_buffer_only_stack(
                 sidecar_pool_specs=self.sidecar_pool_specs,