Radix Cache Split: Spin off TreeCore (#29901)

python/sglang/srt/disaggregation/decode_hicache_mixin.py

@@ -71,7 +71,7 @@ class DecodeHiCachePreallocMixin:
         l3_storage_hit_length = 0
         last_host_node = None
         if self.scheduler.enable_decode_hicache:
-            last_host_node = result.last_host_node
+            last_host_node = self.tree_cache.resolve_node_handle(result.last_host_node)
             if last_host_node.backuped or last_host_node is self.tree_cache.root_node:
                 matched_len = l1_prefix_len + l2_host_hit_length
                 suffix_tokens = req.origin_input_ids[matched_len:]
@@ -82,7 +82,7 @@ class DecodeHiCachePreallocMixin:
                     else None
                 )
                 l3_storage_hit_length = self.tree_cache.query_storage_hit_length(
-                    last_host_node,
+                    result.last_host_node,
                     suffix_tokens,
                     last_hash,
                     prefix_keys,
@@ -93,7 +93,9 @@ class DecodeHiCachePreallocMixin:
             l2_host_hit_length=l2_host_hit_length,
             l3_storage_hit_length=l3_storage_hit_length,
             last_device_node=result.last_device_node,
-            last_host_node=last_host_node if l3_storage_hit_length > 0 else None,
+            last_host_node=(
+                result.last_host_node if l3_storage_hit_length > 0 else None
+            ),
         )
 
     def _start_hicache_prefetch(
@@ -110,7 +112,7 @@ class DecodeHiCachePreallocMixin:
         ):
             return
         try:
-            node = prefix_match.last_host_node
+            node = self.tree_cache.resolve_node_handle(prefix_match.last_host_node)
             matched_len = prefix_match.l1_prefix_len + prefix_match.l2_host_hit_length
             suffix = req.origin_input_ids[
                 matched_len : matched_len + prefix_match.l3_storage_hit_length
@@ -122,7 +124,7 @@ class DecodeHiCachePreallocMixin:
                 else None
             )
             self.tree_cache.prefetch_from_storage(
-                req.rid, node, suffix, last_hash, prefix_keys
+                req.rid, prefix_match.last_host_node, suffix, last_hash, prefix_keys
             )
             prefix_match.prefetch_registered = (
                 req.rid in self.tree_cache.ongoing_prefetch

python/sglang/srt/environ.py

@@ -847,6 +847,8 @@ class Envs:
 
     # Unified Radix Tree
     SGLANG_ENABLE_UNIFIED_RADIX_TREE = EnvBool(False)
+    # Registered TreeCore backend serving the unified radix cache.
+    SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND = EnvStr("python")
 
     # CUDA Graph
     SGLANG_USE_BREAKABLE_CUDA_GRAPH = EnvBool(False)

python/sglang/srt/kv_canary/radix_cache_walker.py

@@ -1,29 +1,19 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
 import torch
 
 from sglang.srt.mem_cache.radix_cache import RadixCache
 from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
-from sglang.srt.mem_cache.unified_cache_components import (
-    BASE_COMPONENT_TYPE,
-    ComponentType,
+from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import (
+    RadixCacheWalkResult,
 )
 from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
 
 if TYPE_CHECKING:
     from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
     from sglang.srt.mem_cache.radix_cache import TreeNode
-    from sglang.srt.mem_cache.unified_radix_cache import UnifiedTreeNode
-
-
-@dataclass(frozen=True, slots=True, kw_only=True)
-class RadixCacheWalkResult:
-    slot_indices: torch.Tensor
-    positions: torch.Tensor
-    prev_slot_indices: torch.Tensor
 
 
 def walk_radix_cache_for_canary(
@@ -40,11 +30,11 @@ def walk_radix_cache_for_canary(
     req. ``swa_resident_only=True`` skips SWA-tombstoned nodes (slots evicted from the SWA
     window)."""
     cache_type = type(radix_cache)
-    if (
-        cache_type is not RadixCache
-        and cache_type is not SWARadixCache
-        and cache_type is not UnifiedRadixCache
-    ):
+    if cache_type is UnifiedRadixCache:
+        return radix_cache.tree_core.walk_for_kv_canary(
+            unlocked_only=unlocked_only, swa_resident_only=swa_resident_only
+        )
+    if cache_type is not RadixCache and cache_type is not SWARadixCache:
         raise NotImplementedError(
             f"walk_radix_cache_for_canary does not support {cache_type.__name__}"
         )
@@ -78,7 +68,7 @@ def walk_radix_cache_for_canary(
 
 def _walk_radix_subtree(
     *,
-    node: TreeNode | UnifiedTreeNode,
+    node: TreeNode,
     radix_cache: BasePrefixCache,
     depth: int,
     parent_last_slot: int,
@@ -89,13 +79,11 @@ def _walk_radix_subtree(
     unlocked_only: bool,
     swa_resident_only: bool,
 ) -> None:
-    node_slots = _node_slots_for_canary(node=node, radix_cache=radix_cache)
+    node_slots = _node_slots_for_canary(node=node)
 
     if unlocked_only:
         emit_slots = not is_root and _node_is_unlocked_for_canary(
-            node=node,
-            radix_cache=radix_cache,
-            swa_resident_only=swa_resident_only,
+            node=node, radix_cache=radix_cache
         )
     else:
         emit_slots = not is_root
@@ -114,12 +102,7 @@ def _walk_radix_subtree(
             prev_slot_buf.append(prev)
         chain_last_slot = slot
 
-    child_depth = depth + _node_len_for_canary(
-        node=node,
-        radix_cache=radix_cache,
-        node_slots=node_slots,
-        is_root=is_root,
-    )
+    child_depth = depth + len(node_slots)
     for child in node.children.values():
         _walk_radix_subtree(
             node=child,
@@ -135,42 +118,17 @@ def _walk_radix_subtree(
         )
 
 
-def _node_slots_for_canary(
-    *,
-    node: TreeNode | UnifiedTreeNode,
-    radix_cache: BasePrefixCache,
-) -> list[int]:
-    value: Any
-    if type(radix_cache) is UnifiedRadixCache:
-        value = node.component_data[BASE_COMPONENT_TYPE].value
-    else:
-        value = node.value
-
+def _node_slots_for_canary(*, node: TreeNode) -> list[int]:
+    value: Any = node.value
     if isinstance(value, torch.Tensor):
         return [int(s) for s in value.tolist()]
     return []
 
 
-def _node_len_for_canary(
-    *,
-    node: TreeNode | UnifiedTreeNode,
-    radix_cache: BasePrefixCache,
-    node_slots: list[int],
-    is_root: bool,
-) -> int:
-    if type(radix_cache) is not UnifiedRadixCache:
-        return len(node_slots)
-
-    if is_root or node.key is None:
-        return len(node_slots)
-    return len(node.key)
-
-
 def _node_is_unlocked_for_canary(
     *,
-    node: TreeNode | UnifiedTreeNode,
+    node: TreeNode,
     radix_cache: BasePrefixCache,
-    swa_resident_only: bool,
 ) -> bool:
     if type(radix_cache) is RadixCache:
         return node.lock_ref == 0
@@ -178,13 +136,6 @@ def _node_is_unlocked_for_canary(
     if type(radix_cache) is SWARadixCache:
         return node.full_lock_ref == 0
 
-    if type(radix_cache) is UnifiedRadixCache:
-        if swa_resident_only and radix_cache.supports_swa():
-            # Unified SWA owns an independent component lock. A node can still
-            # hold Full KV for a running request while its SWA slots are unused.
-            return node.component_data[ComponentType.SWA].lock_ref == 0
-        return node.component_data[BASE_COMPONENT_TYPE].lock_ref == 0
-
     raise NotImplementedError(
         f"walk_radix_cache_for_canary does not support {type(radix_cache).__name__}"
     )
@@ -192,15 +143,10 @@ def _node_is_unlocked_for_canary(
 
 def _node_is_swa_resident_for_canary(
     *,
-    node: TreeNode | UnifiedTreeNode,
+    node: TreeNode,
     radix_cache: BasePrefixCache,
 ) -> bool:
     if type(radix_cache) is SWARadixCache:
         return not node.swa_tombstone
 
-    if type(radix_cache) is UnifiedRadixCache:
-        if not radix_cache.supports_swa():
-            return True
-        return node.component_data[ComponentType.SWA].value is not None
-
     return True

python/sglang/srt/managers/schedule_batch.py

@@ -1249,7 +1249,9 @@ class Req(ReqDllmMixin):
                 )
             )
             if envs.SGLANG_RADIX_FORCE_MISS.get():
-                match_result = zero_match_result(tree_cache, match_result)
+                match_result = zero_match_result(
+                    tree_cache, match_result, extra_key=self.extra_key
+                )
             (
                 self.prefix_indices,
                 self.last_node,

python/sglang/srt/managers/schedule_policy.py

@@ -115,7 +115,9 @@ def match_prefix_for_req(
         )
     )
     if envs.SGLANG_RADIX_FORCE_MISS.get():
-        match_result = zero_match_result(tree_cache, match_result)
+        match_result = zero_match_result(
+            tree_cache, match_result, extra_key=req.extra_key
+        )
     (
         req.prefix_indices,
         req.last_node,
@@ -290,7 +292,7 @@ class SchedulePolicy:
                 )
                 if envs.SGLANG_RADIX_FORCE_MISS.get():
                     match_result = zero_match_result(
-                        self.waiting_queue_radix_tree, match_result
+                        self.waiting_queue_radix_tree, match_result, extra_key=extra_key
                     )
                 in_batch_matching_prefixes = match_result.device_indices
                 if (
@@ -328,7 +330,8 @@ class SchedulePolicy:
         """Sorts the waiting queue based on a depth-first search weighting."""
         last_node_to_reqs = defaultdict(list)
         for req in waiting_queue:
-            last_node_to_reqs[req.last_node].append(req)
+            last_node = tree_cache.resolve_node_handle(req.last_node)
+            last_node_to_reqs[last_node].append(req)
 
         node_to_weight = defaultdict(int)
         for node in last_node_to_reqs:

python/sglang/srt/managers/scheduler.py

@@ -2418,25 +2418,25 @@ class Scheduler(
     def _prefetch_kvcache(self, req: Req):
         if self.enable_hicache_storage:
             req.init_next_round_input(self.tree_cache, cow_mamba=False)
-            last_host_node = req.last_host_node
-            if last_host_node.backuped or last_host_node is self.tree_cache.root_node:
-                last_hash = last_host_node.get_last_hash_value()
+            tree_cache = self.tree_cache
+            if tree_cache.is_backuped(req.last_host_node) or tree_cache.is_root(
+                req.last_host_node
+            ):
                 matched_len = len(req.prefix_indices) + req.host_hit_length
                 match_end = req._compute_max_prefix_len(
                     len(req.full_untruncated_fill_ids)
                 )
                 new_input_tokens = req.full_untruncated_fill_ids[matched_len:match_end]
-
                 prefix_keys = (
-                    last_host_node.get_prefix_hash_values(last_host_node.parent)
-                    if self.tree_cache.hicache_storage_pass_prefix_keys
+                    tree_cache.get_prefix_hash_values(req.last_host_node)
+                    if tree_cache.hicache_storage_pass_prefix_keys
                     else None
                 )
-                self.tree_cache.prefetch_from_storage(
+                tree_cache.prefetch_from_storage(
                     req.rid,
-                    last_host_node,
+                    req.last_host_node,
                     new_input_tokens,
-                    last_hash,
+                    tree_cache.get_last_hash_value(req.last_host_node),
                     prefix_keys,
                 )
 

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -9,6 +9,7 @@ from typing import (
     NamedTuple,
     Optional,
     Protocol,
+    Sequence,
     Tuple,
     runtime_checkable,
 )
@@ -26,6 +27,10 @@ from sglang.srt.observability.metrics_collector import (
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
     from sglang.srt.mem_cache.radix_cache import RadixKey
+    from sglang.srt.mem_cache.unified_cache.cache_action import (
+        CacheAction,
+        ComponentAction,
+    )
     from sglang.srt.mem_cache.unified_cache_components.tree_component import (
         ComponentType,
     )
@@ -78,6 +83,10 @@ class InsertResult:
     last_device_node: Any = None
     mamba_exist: bool = False
     inserted_host_node: Any = None
+    # Controller-applied actions from the non-stepped channels (e.g. insert_host); the stepped insert emits via InsertStepResult.actions.
+    cache_actions: list[CacheAction | ComponentAction] = dataclasses.field(
+        default_factory=list
+    )
 
 
 @dataclasses.dataclass
@@ -191,13 +200,17 @@ class MatchResult(NamedTuple):
     mamba_branching_seqlen: Optional[int] = None
     cache_protected_len: Optional[int] = None
     full_kv_hit_length: int = 0
+    # Actions the Controller applies: CacheActions itself, ComponentActions routed to the owning component.
+    cache_actions: Sequence[CacheAction | ComponentAction] = ()
 
 
-def zero_match_result(tree_cache, match_result: MatchResult) -> MatchResult:
+def zero_match_result(
+    tree_cache, match_result: MatchResult, extra_key: Optional[str] = None
+) -> MatchResult:
     if tree_cache.is_chunk_cache():
         # Chunk caches' match_prefix already returns a miss; no root_node to walk back to.
         return match_result
-    root = tree_cache.root_node
+    root = tree_cache.root_node_handle(extra_key=extra_key)
     return match_result._replace(
         # [:0] keeps dtype and device of the original tensor (e.g. CUDA int64)
         # without allocating a fresh empty tensor.
@@ -258,6 +271,37 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
     def supports_fast_match_prefix(self) -> bool:
         return False
 
+    def resolve_node_handle(self, node_handle: Any) -> Any:
+        """Map a node handle to its node -- e.g. UnifiedRadixCache looks up the
+        node object from its NodeId. Temporary API for the Unified Radix Cache
+        split migration.
+
+        TODO(Jialin): Remove after the Unified Radix Cache split.
+        """
+        return node_handle
+
+    def root_node_handle(self, extra_key: Optional[str] = None) -> Any:
+        """The root handle as match results carry it -- the raw node by default,
+        the root's NodeId for UnifiedRadixCache. extra_key scopes the root for
+        implementations that shard trees per cache namespace."""
+        return self.root_node
+
+    def is_backuped(self, node: Any) -> bool:
+        """Whether the node's Full KV is present on host."""
+        return node.backuped
+
+    def is_root(self, node: Any) -> bool:
+        """Whether the node is a tree root."""
+        return node is self.root_node
+
+    def get_last_hash_value(self, node: Any) -> Optional[str]:
+        """The node's last page hash, or None when it was never hashed."""
+        return node.get_last_hash_value()
+
+    def get_prefix_hash_values(self, node: Any) -> list[str]:
+        """The hash chain of the node's ancestors, in root-to-parent order."""
+        return node.get_prefix_hash_values(node.parent)
+
     @abstractmethod
     def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs):
         pass

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

@@ -0,0 +1,18 @@
+"""Components of Unified Radix Cache.
+
+- **TreeCore** (``unified_tree_core.UnifiedTreeCore``): the pure tree mechanism. Owns
+  the tree structure, per-node values, the LRU(s), bookkeeping, and KV-cache
+  events -- the only object that touches that state; no scheduling policy or IO,
+  and never touches the cache. References the cache-owned component drivers to
+  drive their tree-level hooks.
+- **Components** (``unified_cache_components.TreeComponent`` subclasses): the
+  per-component (FULL/SWA/MAMBA) drivers. Built and owned by the Controller;
+  hold the cache for cache-level logic and the TreeCore for tree state.
+- **Controller** (``unified_radix_cache.UnifiedRadixCache``): the scheduling brain
+  and scheduler-facing facade. Owns the scheduling policy (when to evict / backup /
+  load / prefetch), the IO (HiCache controller, device/host allocators, and the
+  device<->host transfer), the async ack machinery, prefetch/storage, the
+  distributed groups, and the component drivers; holds and drives the TreeCore.
+
+TODO(Jialin): move all unified radix cache logic into this folder.
+"""

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

@@ -0,0 +1,103 @@
+"""A TreeCore emits CacheActions through the TreeCoreInterface to guide
+Controller behavior.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import msgspec
+
+from sglang.srt.mem_cache.unified_cache.component_type import ComponentType
+
+if TYPE_CHECKING:
+    import torch
+
+    from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import NodeId
+
+
+class ReplaceWriteThroughOnNodeSplit(msgspec.Struct, frozen=True):
+    """Replace the pending write-through node on a node split:
+
+        parent -> node    =>    parent -> new_node -> new_child
+
+    old_node_id (the pre-split node) is replaced by new_node_id + new_child_node_id.
+    """
+
+    ack_id: int
+    old_node_id: NodeId
+    new_node_id: NodeId
+    new_child_node_id: NodeId
+
+
+class FreeDeviceKV(msgspec.Struct, frozen=True):
+    """Free unreferenced device KV slots (a SWA-aware combined free)."""
+
+    indices: list[torch.Tensor]
+
+
+class ComponentAction(msgspec.Struct, frozen=True):
+    """Base for component-routed actions; the cache dispatches each one to
+    ``component_type``'s class-level ``apply_component_action``; every subclass
+    declares its ``component_type`` field."""
+
+
+class FreeComponentDeviceSlot(ComponentAction, frozen=True):
+    """Free only the given ``component_type``'s device KV slots."""
+
+    indices: list[torch.Tensor]
+    component_type: ComponentType
+
+
+class FreeComponentHostSlot(ComponentAction, frozen=True):
+    """Free the given ``component_type``'s host KV pages."""
+
+    host_indices: list[torch.Tensor]
+    component_type: ComponentType
+
+
+class BackupKV(msgspec.Struct, frozen=True):
+    """Back up node_ids device->host in order, stopping at the first failure; write-through
+    ids form a contiguous root-first parent chain (each id's parent precedes it), write-back
+    actions carry a single eviction victim."""
+
+    node_ids: list[NodeId]
+
+
+class MambaEvictExcessPathStates(ComponentAction, frozen=True):
+    """Per-path Mamba state-cap eviction from the tail's root path; applied at
+    the insert's commit barrier, after the walk-time backups whose
+    write-through locks shield the backed-up chain."""
+
+    tail_node_id: NodeId
+    component_type: ComponentType = ComponentType.MAMBA
+
+
+class RebuildFullToSWAMapping(ComponentAction, frozen=True):
+    """Rebuild the SWA allocator's full->swa index mapping for loaded chunks."""
+
+    full_indices: list[torch.Tensor]
+    swa_indices: list[torch.Tensor]
+    component_type: ComponentType = ComponentType.SWA
+
+
+class RecoverSWAWithLockedFull(ComponentAction, frozen=True):
+    """Recover an SWA tombstone whose full is locked: keep the locked full, remap it
+    onto the incoming full's SWA translation, and free only the incoming full."""
+
+    node_id: NodeId
+    kept_full: torch.Tensor
+    incoming_full: torch.Tensor
+    component_type: ComponentType = ComponentType.SWA
+
+
+class SWARebuild(ComponentAction, frozen=True):
+    """Rebuild a node's SWA value by translating its source full value, then store it."""
+
+    node_id: NodeId
+    source_value: torch.Tensor
+    component_type: ComponentType = ComponentType.SWA
+
+
+# Cache-owned actions, applied by UnifiedRadixCache itself.
+CacheAction = ReplaceWriteThroughOnNodeSplit | FreeDeviceKV | BackupKV

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

@@ -0,0 +1,29 @@
+"""The per-attention component identity; a leaf module importable from anywhere."""
+
+from enum import Enum
+
+
+class ComponentType(int, Enum):
+    """Integer enum so that per-node list/tuple storage can be indexed directly."""
+
+    FULL = 0
+    SWA = 1
+    MAMBA = 2
+
+    def __str__(self) -> str:  # keep human-readable logging
+        return self.name.lower()
+
+    @property
+    def is_full(self) -> bool:
+        return self == ComponentType.FULL
+
+    @property
+    def is_swa(self) -> bool:
+        return self == ComponentType.SWA
+
+    @property
+    def is_mamba(self) -> bool:
+        return self == ComponentType.MAMBA
+
+
+BASE_COMPONENT_TYPE = ComponentType.FULL

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

@@ -0,0 +1,73 @@
+"""Registry for pluggable TreeCore implementations.
+
+The unified cache constructs its TreeCore through `create_tree_core`, selected
+by SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND (default "python"). To plug in a custom
+implementation, register it under a string name via
+`register_tree_core_backend(name, factory)`.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable, Optional
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.cache_init_params import CacheInitParams
+    from sglang.srt.mem_cache.unified_cache.component_type import ComponentType
+    from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import (
+        UnifiedTreeCoreInterface,
+    )
+    from sglang.srt.mem_cache.unified_cache_components import TreeComponent
+
+TreeCoreFactory = Callable[
+    ["CacheInitParams", "dict[ComponentType, TreeComponent]"],
+    "UnifiedTreeCoreInterface",
+]
+
+_TREE_CORE_REGISTRY: dict[str, TreeCoreFactory] = {}
+
+
+def register_tree_core_backend(name: str, factory: TreeCoreFactory) -> None:
+    """Register a TreeCore factory under `name`."""
+    if not name.strip():
+        raise ValueError(
+            f"register_tree_core_backend: name must be non-empty, got {name!r}"
+        )
+    if name in _TREE_CORE_REGISTRY:
+        raise ValueError(f"register_tree_core_backend: {name!r} is already registered")
+    _TREE_CORE_REGISTRY[name] = factory
+
+
+def get_tree_core_factory(name: str) -> Optional[TreeCoreFactory]:
+    return _TREE_CORE_REGISTRY.get(name)
+
+
+def registered_tree_core_backends() -> list[str]:
+    return list(_TREE_CORE_REGISTRY.keys())
+
+
+def _python_tree_core_factory(
+    params: CacheInitParams, components: dict[ComponentType, TreeComponent]
+) -> UnifiedTreeCoreInterface:
+    """The pure-Python TreeCore."""
+    from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore
+
+    return UnifiedTreeCore(params, components)
+
+
+register_tree_core_backend("python", _python_tree_core_factory)
+
+
+def create_tree_core(
+    name: str,
+    params: CacheInitParams,
+    components: dict[ComponentType, TreeComponent],
+) -> UnifiedTreeCoreInterface:
+    """Construct the TreeCore registered under `name`."""
+    factory = get_tree_core_factory(name)
+    if factory is None:
+        raise ValueError(
+            f"SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND={name!r} is not registered. "
+            f"Registered backends: {registered_tree_core_backends()}. "
+            "External backends must call register_tree_core_backend(...) at import time."
+        )
+    return factory(params, components)

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

@@ -0,0 +1,485 @@
+"""Interface of the TreeCore: the radix tree mechanism that owns the tree
+structure, per-node values, the LRU(s), and bookkeeping. The Controller drives a
+TreeCore exclusively through this surface, so an alternative implementation (e.g.
+a Rust TreeCore) can satisfy it without subclassing the Python tree.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from typing import TYPE_CHECKING, Optional, Sequence
+
+import msgspec
+
+from sglang.srt.mem_cache.events import KVCacheEventMixin
+
+# Tree node id -- the node handle used outside the TreeCore. The concrete tree
+# node is a TreeCore-internal type.
+NodeId = int
+
+
+class BaseEvictionResult(msgspec.Struct):
+    """Base of the eviction step results: the component-keyed device/host
+    values the step freed, for the Controller to drain right after the call
+    returns, plus `tracker` -- the step's per-component evicted counts for the
+    Controller to accumulate."""
+
+    device_frees: dict[ComponentType, list[torch.Tensor]] = msgspec.field(
+        default_factory=lambda: defaultdict(list)
+    )
+    host_frees: dict[ComponentType, list[torch.Tensor]] = msgspec.field(
+        default_factory=lambda: defaultdict(list)
+    )
+    tracker: dict[ComponentType, int] = msgspec.field(
+        default_factory=lambda: defaultdict(int)
+    )
+
+    def __del__(self) -> None:
+        # Drop tripwire: every returned value must be drained before disposal.
+        assert (
+            not self.device_frees and not self.host_frees
+        ), "BaseEvictionResult dropped with undrained values"
+
+
+class EvictDeviceNextNodeResult(BaseEvictionResult):
+    node_id: Optional[NodeId] = None
+
+
+class EvictDeviceLeafResult(BaseEvictionResult):
+    backup_kv: Optional[BackupKV] = None
+
+
+class DemoteResult(BaseEvictionResult):
+    pass
+
+
+class DropSubtreeNoHostResult(BaseEvictionResult):
+    is_dropped: bool = False
+
+
+class DriveHostEvictionResult(BaseEvictionResult):
+    pass
+
+
+class DecSwaLockOnlyResult(BaseEvictionResult):
+    pass
+
+
+class RadixCacheWalkResult(msgspec.Struct, frozen=True, kw_only=True):
+    """Flat (slot, position, prev-slot) rows emitted by the KV-canary walk."""
+
+    slot_indices: torch.Tensor
+    positions: torch.Tensor
+    prev_slot_indices: torch.Tensor
+
+
+class InsertStepResult(msgspec.Struct, frozen=True):
+    """One step of a resumable insert: the Controller executes ``actions``, then
+    resumes while ``result`` is None; ``result`` is set on the final step."""
+
+    actions: list[CacheAction | ComponentAction]
+    result: Optional[InsertResult] = None
+
+
+if TYPE_CHECKING:
+    import torch
+
+    from sglang.srt.managers.schedule_batch import Req
+    from sglang.srt.mem_cache.base_prefix_cache import (
+        DecLockRefParams,
+        DecLockRefResult,
+        IncLockRefResult,
+        InsertParams,
+        InsertResult,
+        MatchPrefixParams,
+        MatchResult,
+    )
+    from sglang.srt.mem_cache.hicache_storage import PoolTransfer, PoolTransferResult
+    from sglang.srt.mem_cache.radix_cache import RadixKey
+    from sglang.srt.mem_cache.unified_cache.cache_action import (
+        BackupKV,
+        CacheAction,
+        ComponentAction,
+    )
+    from sglang.srt.mem_cache.unified_cache.unified_tree_core import (
+        StorageBackupSpec,
+        UnifiedTreeNode,
+    )
+    from sglang.srt.mem_cache.unified_cache_components import (
+        CacheTransferPhase,
+        ComponentType,
+    )
+
+
+class UnifiedTreeCoreInterface(KVCacheEventMixin, ABC):
+    """Methods the Controller invokes on the Tree Core. The Controller treats the
+    Tree Core as opaque behind this surface, which grows as tree operations
+    migrate onto the TreeCore. Inherits KVCacheEventMixin for the KV-event API
+    (take_events, _record_* recorders)."""
+
+    # ==== Tree-owned state the Controller reads (or, via its facade setters, writes) ====
+    page_size: int
+    is_eagle: bool
+    device: torch.device
+    enable_hicache: bool
+    enable_storage: bool
+    write_through_threshold: int
+    is_write_back: bool
+    has_swa_host_pool: bool
+
+    # ==== Tree API ====
+
+    @abstractmethod
+    def reset(self) -> None:
+        """Drop the entire tree and reinitialize empty state."""
+        ...
+
+    @abstractmethod
+    def node_by_id(self, node_id: NodeId) -> UnifiedTreeNode:
+        """Resolve a NodeId -- the tree-node identity the Controller passes
+        across the boundary (e.g. MatchResult fields, lock-ref args) -- back to
+        its tree node.
+
+        TODO(Jialin): Remove after the Unified Radix Cache split.
+        """
+        ...
+
+    @abstractmethod
+    def is_backuped(self, node_id: NodeId) -> bool:
+        """Whether the node's KV is already backed up to host."""
+        ...
+
+    @abstractmethod
+    def is_root(self, node_id: NodeId) -> bool:
+        """Whether the node is the tree root."""
+        ...
+
+    @abstractmethod
+    def get_last_hash_value(self, node_id: NodeId) -> Optional[str]:
+        """The node's last page hash, or None when it was never hashed."""
+        ...
+
+    @abstractmethod
+    def get_prefix_hash_values(self, node_id: NodeId) -> list[str]:
+        """The hash chain of the node's ancestors, in root-to-parent order."""
+        ...
+
+    @abstractmethod
+    def inc_lock_ref(self, node_id: NodeId) -> IncLockRefResult:
+        """Bump the reference count on a node's component locks."""
+        ...
+
+    @abstractmethod
+    def dec_lock_ref(
+        self,
+        node_id: NodeId,
+        params: Optional[DecLockRefParams] = None,
+        skip_swa: bool = False,
+    ) -> DecLockRefResult:
+        """Decrease the reference count on a node's component locks."""
+        ...
+
+    @abstractmethod
+    def dec_swa_lock_only(
+        self, node_id: NodeId, swa_uuid_for_lock: Optional[int]
+    ) -> DecSwaLockOnlyResult:
+        """Decrease only the SWA (and lower-priority co-located) reference
+        counts; the result carries the freed slots."""
+        ...
+
+    # ==== Device eviction (driven step-wise by the Controller's evict()) ====
+
+    @abstractmethod
+    def evict_device_start(
+        self, component_type: ComponentType, request_cnt: int
+    ) -> None:
+        """Begin a device-eviction walk for one component."""
+        ...
+
+    @abstractmethod
+    def evict_device_next_node(
+        self, component_type: ComponentType, tracker: dict[ComponentType, int]
+    ) -> EvictDeviceNextNodeResult:
+        """The next evictable node (None node_id when the walk is exhausted);
+        tracker is the caller's running totals, read for the doneness check."""
+        ...
+
+    @abstractmethod
+    def evict_device_leaf(
+        self, node_id: NodeId, is_write_back: bool
+    ) -> EvictDeviceLeafResult:
+        """Evict a leaf's device value; the result carries a BackupKV when a
+        D->H backup must run (write_back) before the node can be demoted."""
+        ...
+
+    @abstractmethod
+    def drop_subtree_no_host(self, node_id: NodeId) -> DropSubtreeNoHostResult:
+        """Drop an unbacked D-leaf's subtree when its write-back backup failed
+        under host pressure; declines (is_dropped=False) if any node is locked."""
+        ...
+
+    @abstractmethod
+    def demote(self, node_id: NodeId) -> DemoteResult:
+        """Demote a backed-up node: drop its device value after a successful backup."""
+        ...
+
+    @abstractmethod
+    def evict_device_end(self, component_type: ComponentType) -> None:
+        """Finish the component's device-eviction walk."""
+        ...
+
+    @abstractmethod
+    def inc_host_lock_ref(self, node_id: NodeId) -> IncLockRefResult:
+        """Bump the reference count on a node's host-side component locks."""
+        ...
+
+    @abstractmethod
+    def dec_host_lock_ref(
+        self, node_id: NodeId, params: Optional[DecLockRefParams] = None
+    ) -> DecLockRefResult:
+        """Decrease the reference count on a node's host-side component locks."""
+        ...
+
+    @abstractmethod
+    def evictable_size(self) -> int: ...
+
+    @abstractmethod
+    def protected_size(self) -> int: ...
+
+    @abstractmethod
+    def component_evictable_size(self, component_type: ComponentType) -> int:
+        """Evictable token count for one component (0 if the component is absent)."""
+        ...
+
+    @abstractmethod
+    def full_evictable_size(self) -> int: ...
+
+    @abstractmethod
+    def full_protected_size(self) -> int: ...
+
+    @abstractmethod
+    def swa_evictable_size(self) -> int: ...
+
+    @abstractmethod
+    def mamba_evictable_size(self) -> int: ...
+
+    @abstractmethod
+    def swa_protected_size(self) -> int: ...
+
+    @abstractmethod
+    def mamba_protected_size(self) -> int: ...
+
+    @abstractmethod
+    def total_size(self) -> tuple[int, int]:
+        """(full_tokens, aux_tokens) summed across the whole tree."""
+        ...
+
+    @abstractmethod
+    def all_values_flatten(self) -> torch.Tensor: ...
+
+    @abstractmethod
+    def all_mamba_values_flatten(self) -> torch.Tensor: ...
+
+    @abstractmethod
+    def walk_for_kv_canary(
+        self, unlocked_only: bool, swa_resident_only: bool
+    ) -> RadixCacheWalkResult:
+        """Flatten every FULL device slot into (slot, position, prev-slot) rows
+        for the KV-canary sweep."""
+        ...
+
+    @abstractmethod
+    def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
+        """Match a key against the tree; returns device indices + boundary NodeIds."""
+        ...
+
+    @property
+    @abstractmethod
+    def empty_match_result(self) -> MatchResult:
+        """A shared empty MatchResult (empty device indices + boundary NodeIds)."""
+        ...
+
+    @abstractmethod
+    def is_full_device_evicted(self, node_id: NodeId) -> bool:
+        """Whether the node's FULL device value has been evicted."""
+        ...
+
+    @abstractmethod
+    def collect_full_device_indices(
+        self, from_node_id: NodeId, until_node_id: NodeId
+    ) -> torch.Tensor:
+        """Concatenate FULL device values from from_node up to (exclusive) until_node."""
+        ...
+
+    @abstractmethod
+    def begin_insert(self, params: InsertParams) -> InsertStepResult:
+        """Start the (single-flight) insert, running to its first barrier or completion."""
+        ...
+
+    @abstractmethod
+    def resume_insert(self) -> InsertStepResult:
+        """Continue the suspended insert after its step actions were executed."""
+        ...
+
+    @abstractmethod
+    def has_ongoing_insert(self) -> bool:
+        """Whether an insert walk is suspended at a barrier."""
+        ...
+
+    @abstractmethod
+    def end_insert(self) -> list[CacheAction | ComponentAction]:
+        """Finish the insert (idempotent); returns still-pending actions to drain."""
+        ...
+
+    @abstractmethod
+    def drive_host_eviction(
+        self, component_type: ComponentType, num_tokens: int
+    ) -> DriveHostEvictionResult:
+        """Evict a component's host-side resources; no-op if the component is absent."""
+        ...
+
+    # ==== HiCache ====
+
+    @abstractmethod
+    def set_hicache_enabled(self) -> None:
+        """Mark the host tier (HiCache) as wired."""
+        ...
+
+    @abstractmethod
+    def insert_host(
+        self,
+        node_id: NodeId,
+        key: RadixKey,
+        host_value: torch.Tensor,
+        hash_value: list[str],
+    ) -> InsertResult:
+        """Insert a host-side (backuped) tree path descending from the given node."""
+        ...
+
+    @abstractmethod
+    def build_backup_spec(
+        self, node_id: NodeId
+    ) -> tuple[torch.Tensor, dict[ComponentType, list[PoolTransfer]]]:
+        """Read a node's device->host backup spec (device value + transfers) now."""
+        ...
+
+    @abstractmethod
+    def build_storage_backup_spec(
+        self, node_id: NodeId, pass_prefix_keys: bool
+    ) -> Optional[StorageBackupSpec]:
+        """Gather a node's device->storage backup spec; None if not backuped."""
+        ...
+
+    @abstractmethod
+    def build_hicache_transfers(
+        self,
+        component_type: ComponentType,
+        node_id: NodeId,
+        phase: CacheTransferPhase,
+        *,
+        host_indices: Optional[torch.Tensor] = None,
+        token_ids: Optional[Sequence[int]] = None,
+        prefetch_tokens: int = 0,
+        last_hash: Optional[str] = None,
+    ) -> Optional[list[PoolTransfer]]:
+        """Build a component's HiCache transfers for the given node and phase."""
+        ...
+
+    @abstractmethod
+    def build_load_back_spec(
+        self, node_id: NodeId, req: Optional[Req] = None
+    ) -> tuple[PoolTransfer, dict[ComponentType, list[PoolTransfer]]]:
+        """Build the H->D load-back KV transfer plus per-component aux transfers."""
+        ...
+
+    @abstractmethod
+    def prefetch_anchor_info(self, node_id: NodeId) -> Optional[str]:
+        """The anchor node's key extra_key."""
+        ...
+
+    @abstractmethod
+    def commit_hicache_transfers(
+        self,
+        node_id: NodeId,
+        phase: CacheTransferPhase,
+        comp_xfers: dict[ComponentType, list[PoolTransfer]],
+        *,
+        cache_actions: list[CacheAction | ComponentAction],
+        insert_result: Optional[InsertResult] = None,
+        pool_storage_result: Optional[PoolTransferResult] = None,
+    ) -> None:
+        """Commit each component's HiCache transfers onto the node."""
+        ...
+
+    @abstractmethod
+    def commit_backup(
+        self,
+        node_id: NodeId,
+        host_indices: torch.Tensor,
+        comp_xfers: dict[ComponentType, list[PoolTransfer]],
+    ) -> None:
+        """Commit a successful backup to the node."""
+        ...
+
+    @abstractmethod
+    def commit_load_back(
+        self,
+        node_id: NodeId,
+        device_indices: torch.Tensor,
+        kv_xfer: PoolTransfer,
+        comp_xfers: dict[ComponentType, list[PoolTransfer]],
+    ) -> list[CacheAction | ComponentAction]:
+        """Commit a successful H->D load-back onto the node; returns any cache actions."""
+        ...
+
+    @abstractmethod
+    def mark_write_through_pending(self, node_id: NodeId) -> None:
+        """Mark a node as having an in-flight write-through backup."""
+        ...
+
+    @abstractmethod
+    def finish_write_through(self, node_ids: list[NodeId], ack_id: int) -> None:
+        """Clear the write-through-pending mark (when it matches ack_id) and record the
+        host store event for each acked node."""
+        ...
+
+    @abstractmethod
+    def set_component_device_value(
+        self, node_id: NodeId, component_type: ComponentType, value: torch.Tensor
+    ) -> None:
+        """Store an auxiliary (non-Full) component's device value onto a node."""
+        ...
+
+    @abstractmethod
+    def get_component_device_value(
+        self, node_id: NodeId, component_type: ComponentType
+    ) -> Optional[torch.Tensor]:
+        """The component's device value on the node, or None if evicted."""
+        ...
+
+    @abstractmethod
+    def component_has_host_value_only(
+        self, node_id: NodeId, component_type: ComponentType
+    ) -> bool:
+        """Whether the component's data is device-evicted but host-backed."""
+        ...
+
+    # ==== Others ====
+
+    @abstractmethod
+    def sanity_check(
+        self,
+        ongoing_write_through: list[tuple[int, NodeId]],
+        ongoing_load_back: list[tuple[int, NodeId]],
+    ) -> None:
+        """Verify tree invariants and raise AssertionError on any violation.
+
+        ongoing_write_through/ongoing_load_back are (id, node_id) pairs for in-flight ops.
+        """
+        ...
+
+    @abstractmethod
+    def pretty_print(self) -> None:
+        """Print the tree structure for debugging."""
+        ...

python/sglang/srt/mem_cache/unified_cache_components/README.md

@@ -17,6 +17,15 @@ A component-based, pluggable prefix cache framework for SGLang that unifies Full
 │              UnifiedRadixCache                │
 │            (unified_radix_cache.py)           │
 │                                               │
+│  controller: executes all pool/host I/O and   │
+│  drains the tree's deferred Cache/Component   │
+│  Actions ("tree decides, cache executes")     │
+└──────────────────────┬────────────────────────┘
+                       ▼
+┌───────────────────────────────────────────────┐
+│               UnifiedTreeCore                 │
+│      (unified_cache/unified_tree_core.py)     │
+│                                               │
 │  root_node ──► UnifiedTreeNode (radix tree)   │
 │  components ► {ComponentType → TreeComponent} │
 │  lru_lists ─► {ComponentType → UnifiedLRUList}│
@@ -61,7 +70,10 @@ node.component_data[ComponentType.MAMBA]  # MambaComponent data
 
 | File | Contents |
 |------|----------|
-| `../unified_radix_cache.py` | `UnifiedRadixCache`, `UnifiedTreeNode`, `UnifiedLRUList` |
+| `../unified_radix_cache.py` | `UnifiedRadixCache` — the controller: pool/host I/O, deferred-action draining |
+| `../unified_cache/unified_tree_core.py` | `UnifiedTreeCore` — the tree, LRUs, and size counters; `UnifiedTreeNode`, `UnifiedLRUList` |
+| `../unified_cache/unified_tree_core_interface.py` | `UnifiedTreeCoreInterface`, `NodeId` — the tree/cache boundary contract |
+| `../unified_cache/cache_action.py` | Deferred `CacheAction`/`ComponentAction` types emitted by the tree |
 | `tree_component.py` | `TreeComponent` ABC, `ComponentType`, `ComponentData`, `get_and_increase_time_counter`, `next_component_uuid` |
 | `full_component.py` | `FullComponent` — standard full-attention KV cache component |
 | `swa_component.py` | `SWAComponent` — sliding-window attention component with tombstone/window tracking |
@@ -99,7 +111,7 @@ Find the longest cached prefix for a token sequence.
    - Promotes matched path to MRU in each component's LRU via `node_has_component_data()` as filter
    - Updates `last_access_time` with decreasing timestamps up the path (parent < child)
    - Concatenates matched device indices via `torch.cat` (concat length ≤ K, subsumed by O(K))
-   - Calls `finalize_match_result()` per component (Mamba performs copy-on-write: allocates new pool slot, copies SSM state)
+   - Calls `finalize_match_result_in_tree_core()` per component (tree-side: Full/SWA host-hit sums, Mamba `branching_seqlen`); the cache then routes the static `finalize_match_result_in_cache()` per component post-walk (Mamba performs copy-on-write: allocates new pool slot, copies SSM state)
 
 ---
 
@@ -143,7 +155,7 @@ Free cached tokens to reclaim memory.
 | **Complexity** | **O(E·H + L)** — E = nodes evicted, H = tombstone chain height, L = locked nodes skipped in LRU scan. |
 
 **Algorithm detail:**
-1. Calls `drive_eviction()` for each component:
+1. Drives each component's walk via `evict_device_start()` / `evict_device_next_node()` / `evict_device_end()`:
    - Full: drives eviction from `evictable_device_leaves` using `last_access_time`; only device leaves are evicted atomically
    - SWA: scans SWA LRU from tail; **internal** nodes are tombstoned (evict SWA data, keep node), **leaf** nodes are fully deleted; both trigger cascade
    - Mamba: scans Mamba LRU from tail; **internal** nodes are tombstoned, **leaf** nodes are fully deleted; both trigger cascade
@@ -254,7 +266,8 @@ Each component implements these hooks. See `tree_component.py` for the ABC and d
 | Hook | Purpose | Called By | Default |
 |------|---------|-----------|----------|
 | `create_match_validator(match_device_only=False)` | Return a per-match stateful predicate that decides whether a node is a valid match boundary. Full: requires Full device data, or host backup when `match_device_only=False`. SWA: tracks accumulated window length across device/host data. Mamba: requires Mamba device data, or host backup when `match_device_only=False`. | `_match_prefix_helper` | *abstract* |
-| `finalize_match_result()` | Post-process the match result after prefix matching completes. Full/SWA: pass-through. Mamba: copy-on-write — allocates a new mamba pool slot, copies SSM state into the request pool, records `branching_seqlen`. | `_match_post_processor` | pass-through |
+| `finalize_match_result_in_tree_core()` | Tree-side post-processing inside the match walk. Full/SWA: host-hit sums. Mamba: records `branching_seqlen` + the host-hit bump. | `_match_post_processor` | pass-through |
+| `finalize_match_result_in_cache()` | Static, cache-level finalize after the walk (receives the cache + NodeId-based result), dispatched class-level by `UnifiedRadixCache.match_prefix`. Mamba: copy-on-write — allocates a new mamba pool slot, copies SSM state into the request pool. | `UnifiedRadixCache.match_prefix` | pass-through |
 
 ### Insert Phase
 
@@ -276,8 +289,8 @@ Each component implements these hooks. See `tree_component.py` for the ABC and d
 |------|---------|-----------|----------|
 | `evict_component(target=EvictLayer.DEVICE)` | Free this component's device, host, or both resources on a node being evicted. Internal device eviction tombstones (`value = None`); host eviction clears `host_value`. Returns `(device_freed, host_freed)`. | `_evict_component_and_detach_lru` | *abstract* |
 | `eviction_priority()` | Return cascade eviction priority (higher = evicted later). Leaf: all 0. Internal: Full(2) > SWA(1) > Mamba(0). When evicting, all components with ≤ priority on the same node are cascade-evicted. | `_cascade_evict` | `0` |
-| `drive_eviction()` | Drive device eviction until the target amount is freed. Full: leaf-set heap. SWA/Mamba: component LRUs with internal tombstones and atomic leaf deletion. | `evict` | *abstract* |
-| `drive_host_eviction()` | Drive host eviction for this component. Full uses host leaves; SWA/Mamba use host LRUs. | `evict_host` | no-op |
+| `evict_device_start()` / `evict_device_next_node()` / `evict_device_end()` | Step-wise device eviction walk the Controller drives: build the cursor/heap, return the next evictable leaf (freed values collected for the Controller to drain), clear the walk. Full: leaf-set heap. SWA/Mamba: component LRUs with internal tombstones and atomic leaf deletion. | `UnifiedRadixCache._evict_components` | *abstract* |
+| `drive_host_eviction()` | Drive host eviction for this component, collecting freed values for the Controller to drain. Full uses host leaves; SWA/Mamba use host LRUs. | `evict_host` | no-op |
 
 ### Lock Phase
 

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

@@ -10,6 +10,7 @@ from sglang.srt.mem_cache.unified_cache_components.tree_component import (
     EvictLayer,
     LRURefreshPhase,
     PrepareLoadBackResult,
+    PreparePrefetchResult,
     TreeComponent,
     get_and_increase_time_counter,
     next_component_uuid,
@@ -25,6 +26,7 @@ __all__ = [
     "LRURefreshPhase",
     "MambaComponent",
     "PrepareLoadBackResult",
+    "PreparePrefetchResult",
     "SWAComponent",
     "TreeComponent",
     "_NUM_COMPONENT_TYPES",

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

@@ -7,7 +7,6 @@ import torch
 
 from sglang.srt.mem_cache.base_prefix_cache import (
     DecLockRefParams,
-    EvictParams,
     IncLockRefResult,
     InsertResult,
     MatchPrefixParams,
@@ -18,6 +17,7 @@ from sglang.srt.mem_cache.hicache_storage import (
     PoolTransfer,
     PoolTransferResult,
 )
+from sglang.srt.mem_cache.unified_cache.cache_action import FreeComponentDeviceSlot
 from sglang.srt.mem_cache.unified_cache_components.tree_component import (
     CacheTransferPhase,
     ComponentType,
@@ -26,8 +26,12 @@ from sglang.srt.mem_cache.unified_cache_components.tree_component import (
 )
 
 if TYPE_CHECKING:
-    from sglang.srt.managers.schedule_batch import Req
+    from sglang.srt.mem_cache.unified_cache.cache_action import (
+        CacheAction,
+        ComponentAction,
+    )
     from sglang.srt.mem_cache.unified_radix_cache import (
+        NodeId,
         UnifiedTreeNode,
     )
 
@@ -37,13 +41,6 @@ class FullComponent(TreeComponent):
 
     def __init__(self, cache, params):
         super().__init__(cache, params)
-        allocator = cache.token_to_kv_pool_allocator
-        # When SWA is present, only free full-attention KV here;
-        # SWA KV will be freed by cascade via SWAComponent.evict_component.
-        if ComponentType.SWA in cache.tree_components:
-            self._free_full = allocator.full_attn_allocator.free
-        else:
-            self._free_full = allocator.free
         # HiCache state: set to host KV pool when HiCache enabled
         self._full_kv_pool_host = None
 
@@ -60,7 +57,7 @@ class FullComponent(TreeComponent):
             node.component_data[self.component_type].value is not None or node.backuped
         )
 
-    def finalize_match_result(
+    def finalize_match_result_in_tree_core(
         self,
         result: MatchResult,
         params: MatchPrefixParams,
@@ -72,7 +69,7 @@ class FullComponent(TreeComponent):
         ct = self.component_type
         kv_host_hit = 0
         node = result.best_match_node
-        root_node = self.cache.root_node
+        root_node = self.tree_core.root_node
         while node is not result.last_device_node and node is not root_node:
             full_host = node.component_data[ct].host_value
             if full_host is not None:
@@ -103,6 +100,8 @@ class FullComponent(TreeComponent):
     def evict_component(
         self,
         node: UnifiedTreeNode,
+        device_frees: dict[ComponentType, list[torch.Tensor]],
+        host_frees: dict[ComponentType, list[torch.Tensor]],
         target: EvictLayer = EvictLayer.DEVICE,
     ) -> tuple[int, int]:
         cd = node.component_data[self.component_type]
@@ -111,9 +110,9 @@ class FullComponent(TreeComponent):
 
         # Device layer
         if EvictLayer.DEVICE in target and cd.value is not None:
-            self._free_full(cd.value)
+            device_frees[self.component_type].append(cd.value)
             freed = len(cd.value)
-            self.cache.component_evictable_size_[self.component_type] -= freed
+            self.tree_core.component_evictable_size_[self.component_type] -= freed
             # NOTE: cd.value = None is deferred to _cascade_evict (Full as trigger)
             # because SWA's free_swa still needs to read Full.value.
             # cd.value = None
@@ -121,54 +120,78 @@ class FullComponent(TreeComponent):
         # Host layer
         if EvictLayer.HOST in target and cd.host_value is not None:
             host_freed = len(cd.host_value)
-            if self._full_kv_pool_host is not None:
-                self._full_kv_pool_host.free(cd.host_value)
+            host_frees[self.component_type].append(cd.host_value)
             cd.host_value = None
         return freed, host_freed
 
     def eviction_priority(self, is_leaf: bool) -> int:
         return 0 if is_leaf else 2
 
-    def drive_eviction(
-        self, params: EvictParams, tracker: dict[ComponentType, int]
-    ) -> None:
-        request = params.num_tokens
-        heap = [
-            (self.cache.eviction_strategy.get_priority(n), n)
-            for n in self.cache.evictable_device_leaves
+    def _evict_device_start(self, request_cnt: int) -> None:
+        self._evict_device_request_cnt = request_cnt
+        self._evict_device_last_node = None
+        self._evict_device_heap = [
+            (self.tree_core.eviction_strategy.get_priority(n), n)
+            for n in self.tree_core.evictable_device_leaves
         ]
-        heapq.heapify(heap)
+        heapq.heapify(self._evict_device_heap)
+
+    def _evict_device_next_node(
+        self,
+        tracker: dict[ComponentType, int],
+        device_frees: dict[ComponentType, list[torch.Tensor]],
+        host_frees: dict[ComponentType, list[torch.Tensor]],
+    ) -> Optional[NodeId]:
         ct = self.component_type
-        while tracker[ct] < request and heap:
-            _, x = heapq.heappop(heap)
-            if x not in self.cache.evictable_device_leaves:
+        lv = self._evict_device_last_node
+        if (
+            lv is not None
+            and lv.parent is not None
+            and lv.parent in self.tree_core.evictable_device_leaves
+        ):
+            heapq.heappush(
+                self._evict_device_heap,
+                (self.tree_core.eviction_strategy.get_priority(lv.parent), lv.parent),
+            )
+        self._evict_device_last_node = None
+        while tracker[ct] < self._evict_device_request_cnt and self._evict_device_heap:
+            _, x = heapq.heappop(self._evict_device_heap)
+            if x not in self.tree_core.evictable_device_leaves:
                 continue
-            self.cache._evict_device_leaf(x, tracker)
-            if x.parent is not None and x.parent in self.cache.evictable_device_leaves:
-                heapq.heappush(
-                    heap,
-                    (self.cache.eviction_strategy.get_priority(x.parent), x.parent),
-                )
+            self._evict_device_last_node = x
+            return x.id
+        return None
+
+    def _evict_device_end(self) -> None:
+        self._evict_device_heap = []
+        self._evict_device_last_node = None
 
     def drive_host_eviction(
-        self, num_tokens: int, tracker: dict[ComponentType, int]
+        self,
+        num_tokens: int,
+        tracker: dict[ComponentType, int],
+        device_frees: dict[ComponentType, list[torch.Tensor]],
+        host_frees: dict[ComponentType, list[torch.Tensor]],
     ) -> None:
         """Evict host leaves to free KV host pool space."""
         heap = [
-            (self.cache.eviction_strategy.get_priority(n), n)
-            for n in self.cache.evictable_host_leaves
+            (self.tree_core.eviction_strategy.get_priority(n), n)
+            for n in self.tree_core.evictable_host_leaves
         ]
         heapq.heapify(heap)
         ct = self.component_type
         while tracker[ct] < num_tokens and heap:
             _, x = heapq.heappop(heap)
-            if x not in self.cache.evictable_host_leaves:
+            if x not in self.tree_core.evictable_host_leaves:
                 continue
-            self.cache._evict_host_leaf(x, tracker)
-            if x.parent is not None and x.parent in self.cache.evictable_host_leaves:
+            self.tree_core._evict_host_leaf(x, tracker, device_frees, host_frees)
+            if (
+                x.parent is not None
+                and x.parent in self.tree_core.evictable_host_leaves
+            ):
                 heapq.heappush(
                     heap,
-                    (self.cache.eviction_strategy.get_priority(x.parent), x.parent),
+                    (self.tree_core.eviction_strategy.get_priority(x.parent), x.parent),
                 )
 
     def acquire_component_lock(
@@ -182,15 +205,14 @@ class FullComponent(TreeComponent):
         # Only the last host node needs to be protected.
         if lock_host:
             cd = node.component_data[ct]
-            # write_back mode: the anchor may be device-only (no host_value);
-            # pin it anyway so the host-pressure drop fallback cannot delete it.
-            if cd.host_value is None and not self.cache.is_write_back:
+            # write_back mode: the anchor may be device-only (no host_value); pin it anyway.
+            if cd.host_value is None and not self.tree_core.is_write_back:
                 return result
             cd.host_lock_ref += 1
-            self.cache._update_evictable_leaf_sets(node)
+            self.tree_core._update_evictable_leaf_sets(node)
             return result
 
-        root = self.cache.root_node
+        root = self.tree_core.root_node
         cur = node
 
         # Skip the bottom evicted segment
@@ -207,11 +229,11 @@ class FullComponent(TreeComponent):
             ), f"FULL invariant broken: evicted ancestor {cur.id} above device-on segment"
             if cd.lock_ref == 0:
                 key_len = len(cd.value)
-                self.cache.component_evictable_size_[ct] -= key_len
-                self.cache.component_protected_size_[ct] += key_len
+                self.tree_core.component_evictable_size_[ct] -= key_len
+                self.tree_core.component_protected_size_[ct] += key_len
                 delta += key_len
             cd.lock_ref += 1
-            self.cache.evictable_device_leaves.discard(cur)
+            self.tree_core.evictable_device_leaves.discard(cur)
             cur = cur.parent
         result.delta = delta
         return result
@@ -228,13 +250,13 @@ class FullComponent(TreeComponent):
             if cd.host_lock_ref == 0:
                 return
             # Mirror of `acquire`. write_back uses a pure counter.
-            if cd.host_value is None and not self.cache.is_write_back:
+            if cd.host_value is None and not self.tree_core.is_write_back:
                 return
             cd.host_lock_ref -= 1
-            self.cache._update_evictable_leaf_sets(node)
+            self.tree_core._update_evictable_leaf_sets(node)
             return
 
-        root = self.cache.root_node
+        root = self.tree_core.root_node
         skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
         cur = node
         while cur != root:
@@ -247,11 +269,11 @@ class FullComponent(TreeComponent):
 
             if cd.lock_ref == 1:
                 key_len = len(cd.value)
-                self.cache.component_evictable_size_[ct] += key_len
-                self.cache.component_protected_size_[ct] -= key_len
+                self.tree_core.component_evictable_size_[ct] += key_len
+                self.tree_core.component_protected_size_[ct] -= key_len
             cd.lock_ref -= 1
             if cd.lock_ref == 0:
-                self.cache._update_evictable_leaf_sets(cur)
+                self.tree_core._update_evictable_leaf_sets(cur)
             cur = cur.parent
 
     # ---- HiCache Hooks ----
@@ -261,7 +283,8 @@ class FullComponent(TreeComponent):
         node: UnifiedTreeNode,
         phase: CacheTransferPhase,
         *,
-        req: Optional[Req] = None,
+        mamba_pool_idx: Optional[torch.Tensor] = None,
+        host_indices: Optional[torch.Tensor] = None,
         token_ids: Optional[Sequence[int]] = None,
         prefetch_tokens: int = 0,
         last_hash: Optional[str] = None,
@@ -270,7 +293,7 @@ class FullComponent(TreeComponent):
 
         if phase == CacheTransferPhase.BACKUP_HOST:
             # Full KV backup is handled by the main flow
-            # (write_backup → cache_controller.write on host_value directly).
+            # (cache_controller.write on host_value directly).
             # No extra PoolTransfer needed.
             return None
 
@@ -297,7 +320,7 @@ class FullComponent(TreeComponent):
                         else torch.empty((0,), dtype=torch.int64, device="cpu")
                     ),
                     device_indices=None,
-                    nodes_to_load=nodes,
+                    nodes_to_load=[n.id for n in nodes],
                 )
             ]
 
@@ -309,6 +332,7 @@ class FullComponent(TreeComponent):
         phase: CacheTransferPhase,
         transfers: list[PoolTransfer] = (),
         *,
+        cache_actions: list[CacheAction | ComponentAction],
         insert_result: Optional[InsertResult] = None,
         pool_storage_result: Optional[PoolTransferResult] = None,
     ) -> None:
@@ -320,19 +344,39 @@ class FullComponent(TreeComponent):
 
         elif phase == CacheTransferPhase.LOAD_BACK:
             if not transfers or transfers[0].device_indices is None:
-                self.cache._update_evictable_leaf_sets(node)
+                self.tree_core._update_evictable_leaf_sets(node)
                 return
 
             xfer = transfers[0]
             device_indices = xfer.device_indices
             offset = 0
-            for n in xfer.nodes_to_load or []:
+            for nid in xfer.nodes_to_load or []:
+                n = self.tree_core.node_by_id(nid)
                 cd = n.component_data[ct]
                 n_len = len(cd.host_value)
                 cd.value = device_indices[offset : offset + n_len].clone()
                 offset += n_len
                 # Full uses leaf sets, not LRU
-                self.cache.component_evictable_size_[ct] += n_len
-                self.cache._update_evictable_leaf_sets(n)
+                self.tree_core.component_evictable_size_[ct] += n_len
+                self.tree_core._update_evictable_leaf_sets(n)
 
-            self.cache._update_evictable_leaf_sets(node)
+            self.tree_core._update_evictable_leaf_sets(node)
+
+    def free_host_values(self, host_values: list[torch.Tensor]) -> None:
+        if self._full_kv_pool_host is None:
+            return
+        for host_value in host_values:
+            self._full_kv_pool_host.free(host_value)
+
+    def apply_component_action(self, action: ComponentAction) -> None:
+        if isinstance(action, FreeComponentDeviceSlot):
+            alloc = self.cache.token_to_kv_pool_allocator
+            for indices in action.indices:
+                if self.cache.is_swa_enabled:
+                    alloc.full_attn_allocator.free(indices)
+                else:
+                    alloc.free(indices)
+            return
+        raise AssertionError(
+            f"FullComponent: unhandled ComponentAction {type(action).__name__}"
+        )

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

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections import defaultdict
 from typing import TYPE_CHECKING, Callable, Optional, Sequence
 
 import torch
@@ -19,12 +20,18 @@ from sglang.srt.mem_cache.hicache_storage import (
     PoolTransfer,
     PoolTransferResult,
 )
+from sglang.srt.mem_cache.unified_cache.cache_action import (
+    FreeComponentDeviceSlot,
+    FreeComponentHostSlot,
+    MambaEvictExcessPathStates,
+)
 from sglang.srt.mem_cache.unified_cache_components.tree_component import (
     CacheTransferPhase,
     ComponentType,
     EvictLayer,
     LRURefreshPhase,
     PrepareLoadBackResult,
+    PreparePrefetchResult,
     TreeComponent,
     get_and_increase_time_counter,
 )
@@ -33,7 +40,12 @@ from sglang.srt.runtime_context import get_server_args
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
     from sglang.srt.mem_cache.cache_init_params import CacheInitParams
+    from sglang.srt.mem_cache.unified_cache.cache_action import (
+        CacheAction,
+        ComponentAction,
+    )
     from sglang.srt.mem_cache.unified_radix_cache import (
+        NodeId,
         UnifiedRadixCache,
         UnifiedTreeNode,
     )
@@ -46,15 +58,13 @@ class MambaComponent(TreeComponent):
         from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool
 
         assert isinstance(
-            cache.req_to_token_pool, HybridReqToTokenPool
-        ), f"MambaComponent requires HybridReqToTokenPool, got {type(cache.req_to_token_pool)}"
+            params.req_to_token_pool, HybridReqToTokenPool
+        ), f"MambaComponent requires HybridReqToTokenPool, got {type(params.req_to_token_pool)}"
         if not params.enable_mamba_extra_buffer:
             assert (
-                cache.page_size == 1
-            ), f"MambaComponent requires page_size=1 when mamba_extra_buffer is disabled, got {cache.page_size}"
+                params.page_size == 1
+            ), f"MambaComponent requires page_size=1 when mamba_extra_buffer is disabled, got {params.page_size}"
         super().__init__(cache, params)
-        self.enable_mamba_extra_buffer = params.enable_mamba_extra_buffer
-        self.enable_mamba_extra_buffer_lazy = params.enable_mamba_extra_buffer_lazy
         self.mamba_cache_chunk_size = get_server_args().mamba_cache_chunk_size
         self.mamba_max_states_per_path = get_server_args().mamba_max_states_per_path
         # HiCache state
@@ -78,7 +88,7 @@ class MambaComponent(TreeComponent):
                 return
             case LRURefreshPhase.MATCH_END:
                 if node.component_data[ct].value is not None:
-                    self.cache.lru_lists[ct].reset_node_mru(node)
+                    self.tree_core.lru_lists[ct].reset_node_mru(node)
             case LRURefreshPhase.INSERT_END:
                 return
             case _:
@@ -97,15 +107,13 @@ class MambaComponent(TreeComponent):
             or node.component_data[ct].host_value is not None
         )
 
-    def finalize_match_result(
+    def finalize_match_result_in_tree_core(
         self,
         result: MatchResult,
         params: MatchPrefixParams,
         value_chunks: list[torch.Tensor],
         best_value_len: int,
     ) -> MatchResult:
-        cow_mamba = params.cow_mamba
-        req = params.req
         last_node = result.best_match_node
 
         mamba_boundary_len = len(result.device_indices) + result.host_hit_length
@@ -120,75 +128,102 @@ class MambaComponent(TreeComponent):
             aligned_seqlen if aligned_seqlen > mamba_boundary_len else None
         )
 
-        mamba_value = last_node.component_data[self.component_type].value
-        if cow_mamba and mamba_value is not None:
-            assert req is not None
-            if req.mamba_pool_idx is None:
-                dst_index = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
-                if dst_index is None:
-                    # Capture the inc result and thread swa_uuid_for_lock back
-                    # into dec. Without it, SWA's release walks past this
-                    # request's window boundary all the way to root and
-                    # over-decrements SWA locks held by other resident requests
-                    # on ancestor nodes.
-                    lock_result = self.cache.inc_lock_ref(last_node)
-                    self.cache.evict(EvictParams(num_tokens=0, mamba_num=1))
-                    dst_index = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
-                    self.cache.dec_lock_ref(last_node, lock_result.to_dec_params())
-                    assert dst_index is not None, "Can not alloc mamba cache"
-                req.mamba_pool_idx = dst_index[0]
-            req.mamba_cow_src_index = mamba_value
-            req.mamba_needs_clear = False
-
         # HiCache: if mamba was evicted from device but has host backup,
         # ensure mamba_host_hit_length >= 1 so load_back is triggered.
-        cd = last_node.component_data[self.component_type]
-        if cd.value is None and cd.host_value is not None:
+        if self.has_host_value_only(last_node):
             result = result._replace(
                 mamba_host_hit_length=max(result.mamba_host_hit_length, 1)
             )
 
         return result._replace(mamba_branching_seqlen=branching_seqlen)
 
+    def finalize_match_result_in_cache(
+        self, params: MatchPrefixParams, result: MatchResult
+    ) -> MatchResult:
+        # Copy-on-write the matched device mamba state into a per-request slot.
+        if not params.cow_mamba:
+            return result
+        src_index = self.tree_core.get_component_device_value(
+            result.best_match_node, self.component_type
+        )
+        if src_index is None:
+            return result
+        req = params.req
+        assert req is not None
+        if req.mamba_pool_idx is None:
+            dst_index = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
+            if dst_index is None:
+                # Pin the window via inc/dec_lock_ref so evict's SWA release
+                # stops at this request's window boundary instead of walking to
+                # root and over-decrementing locks held by other requests.
+                lock_result = self.cache.inc_lock_ref(result.best_match_node)
+                self.cache.evict(EvictParams(num_tokens=0, mamba_num=1))
+                dst_index = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
+                self.cache.dec_lock_ref(
+                    result.best_match_node, lock_result.to_dec_params()
+                )
+                assert dst_index is not None, "Can not alloc mamba cache"
+            req.mamba_pool_idx = dst_index[0]
+        req.mamba_cow_src_index = src_index
+        req.mamba_needs_clear = False
+        return result
+
     def commit_insert_component_data(
         self,
         node: UnifiedTreeNode,
         is_new_leaf: bool,
         params: InsertParams,
         result: InsertResult,
+        cache_actions: list[CacheAction | ComponentAction],
     ) -> None:
         assert params.mamba_value is not None
         if is_new_leaf:
             node.component_data[self.component_type].value = params.mamba_value
-            self.cache.lru_lists[self.component_type].insert_mru(node)
-            self.cache.component_evictable_size_[self.component_type] += len(
+            self.tree_core.lru_lists[self.component_type].insert_mru(node)
+            self.tree_core.component_evictable_size_[self.component_type] += len(
                 params.mamba_value
             )
-            self._evict_excess_path_states(node)
+            self._emit_excess_path_states_eviction(node, cache_actions)
             return
         if node.component_data[self.component_type].value is None:
             node.component_data[self.component_type].value = params.mamba_value
             # move from host LRU to device LRU
-            host_lru = self.cache.host_lru_lists[self.component_type]
+            host_lru = self.tree_core.host_lru_lists[self.component_type]
             if host_lru.in_list(node):
                 host_lru.remove_node(node)
-            self.cache.lru_lists[self.component_type].insert_mru(node)
-            self.cache.component_evictable_size_[self.component_type] += len(
+            self.tree_core.lru_lists[self.component_type].insert_mru(node)
+            self.tree_core.component_evictable_size_[self.component_type] += len(
                 params.mamba_value
             )
             node.last_access_time = get_and_increase_time_counter()
-            self._evict_excess_path_states(node)
+            self._emit_excess_path_states_eviction(node, cache_actions)
             return
-        self.cache.lru_lists[self.component_type].reset_node_mru(node)
+        self.tree_core.lru_lists[self.component_type].reset_node_mru(node)
         node.last_access_time = get_and_increase_time_counter()
         result.mamba_exist = True
 
-    def _evict_excess_path_states(self, tail: UnifiedTreeNode) -> None:
+    def _emit_excess_path_states_eviction(
+        self,
+        tail: UnifiedTreeNode,
+        cache_actions: list[CacheAction | ComponentAction],
+    ) -> None:
+        """Defer the path-cap eviction so it runs after the insert's BackupKV."""
+        if self.mamba_max_states_per_path < 0:
+            return
+        cache_actions.append(MambaEvictExcessPathStates(tail.id))
+
+    def _evict_excess_path_states(
+        self,
+        tail: UnifiedTreeNode,
+        device_frees: dict[ComponentType, list[torch.Tensor]],
+        host_frees: dict[ComponentType, list[torch.Tensor]],
+    ) -> None:
         """Evict shallow eligible device checkpoints beyond the path cap.
 
         Full KV and any existing host backup are retained. The tail, forks,
-        locked nodes, and device leaves are preserved, so the cap is a
-        best-effort soft limit.
+        locked nodes (including a pending backup chain's write-through locks),
+        and device leaves are preserved, so the cap is a best-effort soft
+        limit. Freed slots are collected into the caller's dicts.
         """
         cap = self.mamba_max_states_per_path
         if cap < 0:
@@ -197,7 +232,7 @@ class MambaComponent(TreeComponent):
         ct = self.component_type
         holders = []
         node = tail
-        while node is not None and node is not self.cache.root_node:
+        while node is not None and node is not self.tree_core.root_node:
             if node.component_data[ct].value is not None:
                 holders.append(node)
             node = node.parent
@@ -212,12 +247,17 @@ class MambaComponent(TreeComponent):
                 break
             if node.component_data[ct].lock_ref > 0 or len(node.children) != 1:
                 continue
-            if node in self.cache.evictable_device_leaves:
+            if node in self.tree_core.evictable_device_leaves:
                 continue
-            self.cache._evict_component_and_detach_lru(
-                node, self, target=EvictLayer.DEVICE, tracker=tracker
+            self.tree_core._evict_component_and_detach_lru(
+                node,
+                self,
+                device_frees,
+                host_frees,
+                target=EvictLayer.DEVICE,
+                tracker=tracker,
             )
-            self.cache._cascade_evict(node, self, tracker)
+            self.tree_core._cascade_evict(node, self, tracker, device_frees, host_frees)
             excess -= 1
 
     def redistribute_on_node_split(
@@ -233,6 +273,8 @@ class MambaComponent(TreeComponent):
     def evict_component(
         self,
         node: UnifiedTreeNode,
+        device_frees: dict[ComponentType, list[torch.Tensor]],
+        host_frees: dict[ComponentType, list[torch.Tensor]],
         target: EvictLayer = EvictLayer.DEVICE,
     ) -> tuple[int, int]:
         cd = node.component_data[self.component_type]
@@ -241,17 +283,16 @@ class MambaComponent(TreeComponent):
 
         # Device layer
         if EvictLayer.DEVICE in target and cd.value is not None:
-            self._free_mamba_value(cd.value)
+            device_frees[self.component_type].append(cd.value)
             freed = len(cd.value)
-            self.cache.component_evictable_size_[self.component_type] -= freed
+            self.tree_core.component_evictable_size_[self.component_type] -= freed
             cd.value = None
 
         # Host layer
-        host_lru = self.cache.host_lru_lists[self.component_type]
+        host_lru = self.tree_core.host_lru_lists[self.component_type]
         if EvictLayer.HOST in target and cd.host_value is not None:
             host_freed = len(cd.host_value)
-            if self._mamba_pool_host is not None:
-                self._mamba_pool_host.free(cd.host_value)
+            host_frees[self.component_type].append(cd.host_value)
             cd.host_value = None
             if host_lru.in_list(node):
                 host_lru.remove_node(node)
@@ -267,30 +308,56 @@ class MambaComponent(TreeComponent):
 
         return freed, host_freed
 
-    def drive_eviction(
-        self, params: EvictParams, tracker: dict[ComponentType, int]
-    ) -> None:
-        request = params.mamba_num
+    def _evict_device_start(self, request_cnt: int) -> None:
+        """Begin the device-eviction walk from this component's LRU cursor."""
+        self._evict_device_request_cnt = request_cnt
+        self._evict_device_cursor = self.tree_core.lru_lists[
+            self.component_type
+        ].get_lru_no_lock()
+
+    def _evict_device_next_node(
+        self,
+        tracker: dict[ComponentType, int],
+        device_frees: dict[ComponentType, list[torch.Tensor]],
+        host_frees: dict[ComponentType, list[torch.Tensor]],
+    ) -> Optional[NodeId]:
+        """Return the next device-leaf node for the driver to evict, or None.
+        Internal nodes are tombstoned inline (no IO); the cursor is re-validated
+        (reset to LRU head) if the previous node's eviction removed it."""
         ct = self.component_type
-        lru = self.cache.lru_lists[ct]
-        x = lru.get_lru_no_lock()
-        while tracker[ct] < request and x is not None and lru.in_list(x):
+        lru = self.tree_core.lru_lists[ct]
+        if self._evict_device_cursor is not None and not lru.in_list(
+            self._evict_device_cursor
+        ):
+            self._evict_device_cursor = lru.get_lru_no_lock()
+        while (
+            tracker[ct] < self._evict_device_request_cnt
+            and self._evict_device_cursor is not None
+            and lru.in_list(self._evict_device_cursor)
+        ):
+            x = self._evict_device_cursor
             assert x.component_data[ct].value is not None
-            if x in self.cache.evictable_device_leaves:
-                # D-leaf: atomic eviction of all components
-                x_next = lru.get_prev_no_lock(x)
-                self.cache._evict_device_leaf(x, tracker)
-                if not lru.in_list(x_next):
-                    x_next = lru.get_lru_no_lock()
-                x = x_next
-            else:
-                # Internal: tombstone Mamba + cascade
-                x_next = lru.get_prev_no_lock(x)
-                self.cache._evict_component_and_detach_lru(
-                    x, self, target=EvictLayer.DEVICE, tracker=tracker
-                )
-                self.cache._cascade_evict(x, self, tracker)
-                x = x_next
+            if x in self.tree_core.evictable_device_leaves:
+                self._evict_device_cursor = lru.get_prev_no_lock(x)
+                return x.id
+            x_next = lru.get_prev_no_lock(x)
+            self.tree_core._evict_component_and_detach_lru(
+                x,
+                self,
+                target=EvictLayer.DEVICE,
+                tracker=tracker,
+                device_frees=device_frees,
+                host_frees=host_frees,
+            )
+            self.tree_core._cascade_evict(
+                x, self, tracker, device_frees=device_frees, host_frees=host_frees
+            )
+            self._evict_device_cursor = x_next
+        return None
+
+    def _evict_device_end(self) -> None:
+        """Clear the device-eviction walk cursor state."""
+        self._evict_device_cursor = None
 
     def acquire_component_lock(
         self,
@@ -299,7 +366,7 @@ class MambaComponent(TreeComponent):
         lock_host: bool = False,
     ) -> IncLockRefResult:
         ct = self.component_type
-        if node is self.cache.root_node:
+        if node is self.tree_core.root_node:
             return result
         cd = node.component_data[ct]
         value = cd.host_value if lock_host else cd.value
@@ -310,15 +377,15 @@ class MambaComponent(TreeComponent):
 
         if lock_host:
             if cd.host_lock_ref == 0:
-                host_lru = self.cache.host_lru_lists[ct]
+                host_lru = self.tree_core.host_lru_lists[ct]
                 if host_lru.in_list(node):
                     host_lru.remove_node(node)
             cd.host_lock_ref += 1
         else:
             if cd.lock_ref == 0:
                 vlen = len(value)
-                self.cache.component_evictable_size_[ct] -= vlen
-                self.cache.component_protected_size_[ct] += vlen
+                self.tree_core.component_evictable_size_[ct] -= vlen
+                self.tree_core.component_protected_size_[ct] += vlen
             cd.lock_ref += 1
         return result
 
@@ -329,7 +396,7 @@ class MambaComponent(TreeComponent):
         lock_host: bool = False,
     ) -> None:
         ct = self.component_type
-        if node is self.cache.root_node:
+        if node is self.tree_core.root_node:
             return
         cd = node.component_data[ct]
         skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
@@ -340,7 +407,7 @@ class MambaComponent(TreeComponent):
         if lock_host:
             cd.host_lock_ref -= 1
             if cd.host_lock_ref == 0 and cd.value is None and cd.host_value is not None:
-                host_lru = self.cache.host_lru_lists[ct]
+                host_lru = self.tree_core.host_lru_lists[ct]
                 if not host_lru.in_list(node):
                     host_lru.insert_mru(node)
             return
@@ -348,8 +415,8 @@ class MambaComponent(TreeComponent):
         if cd.lock_ref > 0:
             if cd.lock_ref == 1:
                 vlen = len(value)
-                self.cache.component_evictable_size_[ct] += vlen
-                self.cache.component_protected_size_[ct] -= vlen
+                self.tree_core.component_evictable_size_[ct] += vlen
+                self.tree_core.component_protected_size_[ct] -= vlen
             cd.lock_ref -= 1
 
     def _alloc_mamba_slot(self) -> torch.Tensor:
@@ -395,7 +462,7 @@ class MambaComponent(TreeComponent):
         token_ids_len: int,
         is_finished: bool,
     ) -> Optional[int]:
-        if self.enable_mamba_extra_buffer:
+        if self.cache.enable_mamba_extra_buffer:
             cache_len = req.mamba_last_track_seqlen
         else:
             cache_len = token_ids_len
@@ -416,7 +483,7 @@ class MambaComponent(TreeComponent):
         if is_finished:
             if cache_len is None:
                 cache_len = 0
-            if self.enable_mamba_extra_buffer:
+            if self.cache.enable_mamba_extra_buffer:
                 keep_idx = self.cache.req_to_token_pool.get_mamba_ping_pong_keep_idx(
                     req
                 )
@@ -435,7 +502,7 @@ class MambaComponent(TreeComponent):
                 return 0
             # Donate the mamba index to the radix cache instead of copying.
             if self.int8_ckpt_pool is not None:
-                if self.enable_mamba_extra_buffer:
+                if self.cache.enable_mamba_extra_buffer:
                     new_slot = self._alloc_mamba_slot()
                     src_active = (
                         self.cache.req_to_token_pool.donate_mamba_ping_pong_slot(
@@ -448,7 +515,7 @@ class MambaComponent(TreeComponent):
                     mamba_value_donated = self._commit_int8_checkpoint(
                         req.mamba_pool_idx.view(-1)
                     )
-            elif self.enable_mamba_extra_buffer:
+            elif self.cache.enable_mamba_extra_buffer:
                 new_slot = self._alloc_mamba_slot()
                 mamba_value_donated = (
                     self.cache.req_to_token_pool.donate_mamba_ping_pong_slot(
@@ -491,7 +558,7 @@ class MambaComponent(TreeComponent):
                 pool.free_mamba_cache(req)
                 return
 
-            if self.enable_mamba_extra_buffer:
+            if self.cache.enable_mamba_extra_buffer:
                 keep_idx = (
                     pool.get_mamba_ping_pong_keep_idx(req)
                     if mamba_value_inserted
@@ -515,17 +582,16 @@ class MambaComponent(TreeComponent):
 
     def prepare_load_back(
         self,
-        node: UnifiedTreeNode,
+        node_id: NodeId,
         *,
         req: Optional[Req] = None,
     ) -> PrepareLoadBackResult:
-        cd = node.component_data[self.component_type]
-        # skip unless the node needs a load-back (device value absent), like build_hicache_transfers
         if (
             req is None
             or req.mamba_pool_idx is not None
-            or cd.host_value is None
-            or cd.value is not None
+            or not self.tree_core.component_has_host_value_only(
+                node_id, self.component_type
+            )
         ):
             return PrepareLoadBackResult()
         dst = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
@@ -544,12 +610,27 @@ class MambaComponent(TreeComponent):
             self.cache.req_to_token_pool.mamba_allocator.free(prep.allocated_mamba_slot)
             req.mamba_pool_idx = None
 
+    def prepare_prefetch(
+        self,
+        node_id: NodeId,
+        *,
+        prefetch_tokens: int = 0,
+    ) -> PreparePrefetchResult:
+        host_indices = self._mamba_pool_host.alloc(1)
+        if host_indices is None:
+            self.cache.evict_host(1, ComponentType.MAMBA)
+            host_indices = self._mamba_pool_host.alloc(1)
+        if host_indices is None:
+            return PreparePrefetchResult(alloc_failed=True)
+        return PreparePrefetchResult(host_indices=host_indices)
+
     def build_hicache_transfers(
         self,
         node: UnifiedTreeNode,
         phase: CacheTransferPhase,
         *,
-        req: Optional[Req] = None,
+        mamba_pool_idx: Optional[torch.Tensor] = None,
+        host_indices: Optional[torch.Tensor] = None,
         token_ids: Optional[Sequence[int]] = None,
         prefetch_tokens: int = 0,
         last_hash: Optional[str] = None,
@@ -580,19 +661,18 @@ class MambaComponent(TreeComponent):
                     PoolTransfer(
                         name=PoolName.MAMBA,
                         host_indices=cd.host_value,
-                        nodes_to_load=[node],
+                        nodes_to_load=[node.id],
                     )
                 )
 
-            # Per-request mamba CoW: H→D copy into the request's device slot allocated by prepare_load_back.
+            # Per-request mamba CoW: H→D copy into the request's device slot pre-allocated on the caller side.
             cd = node.component_data[ct]
-            if req is not None and cd.host_value is not None:
-                assert req.mamba_pool_idx is not None
+            if mamba_pool_idx is not None and cd.host_value is not None:
                 transfers.append(
                     PoolTransfer(
                         name=PoolName.MAMBA,
                         host_indices=cd.host_value,
-                        device_indices=req.mamba_pool_idx.unsqueeze(0),
+                        device_indices=mamba_pool_idx.unsqueeze(0),
                     )
                 )
 
@@ -612,12 +692,7 @@ class MambaComponent(TreeComponent):
             ]
 
         if phase == CacheTransferPhase.PREFETCH:
-            host_indices = self._mamba_pool_host.alloc(1)
-            if host_indices is None:
-                self.cache.evict_host(1, ComponentType.MAMBA)
-                host_indices = self._mamba_pool_host.alloc(1)
-            if host_indices is None:
-                return []
+            assert host_indices is not None
             return [
                 PoolTransfer(
                     name=PoolName.MAMBA,
@@ -635,6 +710,7 @@ class MambaComponent(TreeComponent):
         phase: CacheTransferPhase,
         transfers: list[PoolTransfer] = (),
         *,
+        cache_actions: list[CacheAction | ComponentAction],
         insert_result: Optional[InsertResult] = None,
         pool_storage_result: Optional[PoolTransferResult] = None,
     ) -> None:
@@ -655,11 +731,11 @@ class MambaComponent(TreeComponent):
                 cd.value = transfer.device_indices.clone()
                 count = len(cd.value)
                 # Move from host LRU to device LRU
-                host_lru = self.cache.host_lru_lists[ct]
+                host_lru = self.tree_core.host_lru_lists[ct]
                 if host_lru.in_list(node):
                     host_lru.remove_node(node)
-                self.cache.lru_lists[ct].insert_mru(node)
-                self.cache.component_evictable_size_[ct] += count
+                self.tree_core.lru_lists[ct].insert_mru(node)
+                self.tree_core.component_evictable_size_[ct] += count
 
         elif phase == CacheTransferPhase.PREFETCH:
             if not transfers:
@@ -671,7 +747,10 @@ class MambaComponent(TreeComponent):
                 and pool_storage_result.extra_pool_hit_pages.get(PoolName.MAMBA, 0) >= 1
             )
             target_node = (
-                insert_result.inserted_host_node if insert_result is not None else None
+                self.tree_core.node_by_id(insert_result.inserted_host_node)
+                if insert_result is not None
+                and insert_result.inserted_host_node is not None
+                else None
             )
             if (
                 host_indices is None
@@ -679,8 +758,10 @@ class MambaComponent(TreeComponent):
                 or not loaded
                 or target_node.component_data[ct].host_value is not None
             ):
-                self.cache.cache_controller.append_host_mem_release(
-                    extra_pools=[transfer]
+                cache_actions.append(
+                    FreeComponentHostSlot(
+                        [host_indices], component_type=ComponentType.MAMBA
+                    )
                 )
                 if insert_result is not None:
                     insert_result.mamba_exist = True
@@ -688,33 +769,86 @@ class MambaComponent(TreeComponent):
 
             target_node.component_data[ct].host_value = host_indices.clone()
             if target_node.component_data[ct].value is None:
-                host_lru = self.cache.host_lru_lists[ct]
+                host_lru = self.tree_core.host_lru_lists[ct]
                 if not host_lru.in_list(target_node):
                     host_lru.insert_mru(target_node)
             if insert_result is not None:
                 insert_result.mamba_exist = False
 
     def drive_host_eviction(
-        self, num_tokens: int, tracker: dict[ComponentType, int]
+        self,
+        num_tokens: int,
+        tracker: dict[ComponentType, int],
+        device_frees: dict[ComponentType, list[torch.Tensor]],
+        host_frees: dict[ComponentType, list[torch.Tensor]],
     ) -> None:
         """Evict mamba host resources.
         Internal nodes: private tombstone (free host mamba only).
         Host leaves: atomic eviction via _evict_host_leaf."""
         ct = self.component_type
-        host_lru = self.cache.host_lru_lists[ct]
+        host_lru = self.tree_core.host_lru_lists[ct]
         x = host_lru.get_lru_no_host_lock()
         while tracker[ct] < num_tokens and x is not None and host_lru.in_list(x):
             x_next = host_lru.get_prev_no_host_lock(x)
             cd = x.component_data[ct]
-            if x in self.cache.evictable_host_leaves:
+            if x in self.tree_core.evictable_host_leaves:
                 # Host leaf: atomic eviction (all components host + delete)
-                self.cache._evict_host_leaf(x, tracker)
+                self.tree_core._evict_host_leaf(x, tracker, device_frees, host_frees)
             else:
                 # Internal: tombstone Mamba + cascade
                 assert cd.host_value is not None
-                self.cache._evict_component_and_detach_lru(
-                    x, self, target=EvictLayer.HOST, tracker=tracker
+                self.tree_core._evict_component_and_detach_lru(
+                    x,
+                    self,
+                    target=EvictLayer.HOST,
+                    tracker=tracker,
+                    device_frees=device_frees,
+                    host_frees=host_frees,
                 )
-                self.cache._cascade_evict(x, self, tracker, target=EvictLayer.HOST)
-                self.cache._update_evictable_leaf_sets(x)
+                self.tree_core._cascade_evict(
+                    x,
+                    self,
+                    tracker,
+                    device_frees=device_frees,
+                    host_frees=host_frees,
+                    target=EvictLayer.HOST,
+                )
+                self.tree_core._update_evictable_leaf_sets(x)
             x = x_next
+
+    def free_host_values(self, host_values: list[torch.Tensor]) -> None:
+        if self._mamba_pool_host is None:
+            return
+        for host_value in host_values:
+            self._mamba_pool_host.free(host_value)
+
+    def apply_component_action(self, action: ComponentAction) -> None:
+        if isinstance(action, MambaEvictExcessPathStates):
+            device_frees: dict[ComponentType, list[torch.Tensor]] = defaultdict(list)
+            host_frees: dict[ComponentType, list[torch.Tensor]] = defaultdict(list)
+            # Drain even if the walk raises so tombstoned slots are not leaked.
+            try:
+                self._evict_excess_path_states(
+                    self.tree_core.node_by_id(action.tail_node_id),
+                    device_frees,
+                    host_frees,
+                )
+            finally:
+                self.cache._free_values(device_frees, host_frees)
+            return
+        if isinstance(action, FreeComponentDeviceSlot):
+            for indices in action.indices:
+                self._free_mamba_value(indices)
+            return
+        if isinstance(action, FreeComponentHostSlot):
+            for host_indices in action.host_indices:
+                if host_indices is not None and host_indices.numel() > 0:
+                    self.cache.cache_controller.append_host_mem_release(
+                        extra_pools=[
+                            PoolTransfer(name=PoolName.MAMBA, host_indices=host_indices)
+                        ]
+                    )
+            return
+        raise AssertionError(
+            f"MambaComponent: unhandled ComponentAction {type(action).__name__}"
+        )

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

@@ -6,7 +6,6 @@ import torch
 
 from sglang.srt.mem_cache.base_prefix_cache import (
     DecLockRefParams,
-    EvictParams,
     IncLockRefResult,
     InsertParams,
     InsertResult,
@@ -20,12 +19,21 @@ from sglang.srt.mem_cache.hicache_storage import (
     PoolTransfer,
     PoolTransferResult,
 )
+from sglang.srt.mem_cache.unified_cache.cache_action import (
+    FreeComponentDeviceSlot,
+    FreeComponentHostSlot,
+    FreeDeviceKV,
+    RebuildFullToSWAMapping,
+    RecoverSWAWithLockedFull,
+    SWARebuild,
+)
 from sglang.srt.mem_cache.unified_cache_components.tree_component import (
     BASE_COMPONENT_TYPE,
     CacheTransferPhase,
     ComponentType,
     EvictLayer,
     LRURefreshPhase,
+    PreparePrefetchResult,
     TreeComponent,
     next_component_uuid,
 )
@@ -33,7 +41,12 @@ from sglang.srt.mem_cache.unified_cache_components.tree_component import (
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
     from sglang.srt.mem_cache.cache_init_params import CacheInitParams
+    from sglang.srt.mem_cache.unified_cache.cache_action import (
+        CacheAction,
+        ComponentAction,
+    )
     from sglang.srt.mem_cache.unified_radix_cache import (
+        NodeId,
         UnifiedRadixCache,
         UnifiedTreeNode,
     )
@@ -53,8 +66,8 @@ class SWAComponent(TreeComponent):
         from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
 
         assert isinstance(
-            cache.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator
-        ), f"SWAComponent requires SWATokenToKVPoolAllocator, got {type(cache.token_to_kv_pool_allocator)}"
+            params.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator
+        ), f"SWAComponent requires SWATokenToKVPoolAllocator, got {type(params.token_to_kv_pool_allocator)}"
         super().__init__(cache, params)
         self.sliding_window_size = params.sliding_window_size
         # HiCache state: set to host SWA pool when HiCache enabled
@@ -81,39 +94,17 @@ class SWAComponent(TreeComponent):
                 # MATCH_END / INSERT_END instead.
                 return
             case LRURefreshPhase.MATCH_END | LRURefreshPhase.INSERT_END:
-                self.cache.lru_lists[
+                self.tree_core.lru_lists[
                     self.component_type
                 ].reset_node_and_window_ancestors_mru(
                     node,
                     root_node,
-                    self.sliding_window_size + self.cache.page_size,
+                    self.sliding_window_size + self.tree_core.page_size,
                     self.node_has_component_data,
                 )
             case _:
                 raise ValueError(f"Unknown LRURefreshPhase: {phase}")
 
-    def _restore_device_value(self, node: UnifiedTreeNode, value: torch.Tensor) -> None:
-        ct = self.component_type
-        node.component_data[ct].value = value
-        host_lru = self.cache.host_lru_lists[ct]
-        if host_lru.in_list(node):
-            host_lru.remove_node(node)
-        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]:
@@ -124,7 +115,7 @@ class SWAComponent(TreeComponent):
         # unified_kv never caches the SWA ring (per-request, not content-stable),
         # so SWA bookkeeping must not gate the match here.
         swa_device_only_hicache = (
-            self._swa_kv_pool_host is None and self.cache.cache_controller is not None
+            not self.tree_core.has_swa_host_pool and self.tree_core.enable_hicache
         )
 
         def validator(node: UnifiedTreeNode) -> bool:
@@ -141,7 +132,7 @@ class SWAComponent(TreeComponent):
 
         return validator
 
-    def finalize_match_result(
+    def finalize_match_result_in_tree_core(
         self,
         result: MatchResult,
         params: MatchPrefixParams,
@@ -152,7 +143,7 @@ class SWAComponent(TreeComponent):
         n_swa = 0
         swa_host_hit = 0
         node = result.best_match_node
-        root = self.cache.root_node
+        root = self.tree_core.root_node
         while node is not root and n_swa < self.sliding_window_size:
             cd = node.component_data[ct]
             if cd.value is not None:
@@ -182,6 +173,7 @@ class SWAComponent(TreeComponent):
         total_prefix_len: int,
         value_slice: torch.Tensor,
         params: InsertParams,
+        cache_actions: list[CacheAction | ComponentAction],
     ) -> int:
         if params.prev_prefix_len >= total_prefix_len + prefix_len:
             return prefix_len
@@ -196,40 +188,38 @@ class SWAComponent(TreeComponent):
             node.component_data[self.component_type].lock_ref == 0
         ), f"tombstone {self.component_type} lock_ref should be 0, node {node.id}"
         assert (
-            swa_evicted_seqlen % self.cache.page_size == 0
+            swa_evicted_seqlen % self.tree_core.page_size == 0
         ), f"{self.component_type}: swa_evicted_seqlen must be page-aligned, {swa_evicted_seqlen=}"
 
         if swa_evicted_seqlen <= total_prefix_len:
             # Branch 1: entire value_slice is within SWA window — recover
+            old_full = full_cd.value
             if full_cd.lock_ref > 0:
-                self._restore_device_value_with_locked_full(
-                    node, full_cd.value, value_slice
+                cache_actions.append(
+                    RecoverSWAWithLockedFull(node.id, old_full, 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)
+            cache_actions.append(FreeDeviceKV([old_full]))
+            cache_actions.append(SWARebuild(node.id, value_slice))
             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
-            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:]
+            is_locked = full_cd.lock_ref > 0
+            old_full = full_cd.value[start_idx:]
+            _, action = self.tree_core._split_node(node.key, node, start_idx)
+            if action is not None:
+                cache_actions.append(action)
+            new_full = value_slice[start_idx:]
+            if is_locked:
+                cache_actions.append(
+                    RecoverSWAWithLockedFull(node.id, old_full, new_full)
                 )
                 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:
-            ].clone()
-            swa_value = self._translate_full_to_swa(
-                node.component_data[BASE_COMPONENT_TYPE].value
-            )
-            self._restore_device_value(node, swa_value)
+            node.component_data[BASE_COMPONENT_TYPE].value = new_full.clone()
+            cache_actions.append(FreeDeviceKV([old_full]))
+            cache_actions.append(SWARebuild(node.id, new_full))
             return start_idx
         else:
             # Branch 3: entire value_slice is outside SWA window — not consumed
@@ -241,6 +231,7 @@ class SWAComponent(TreeComponent):
         prefix_len: int,
         total_prefix_len: int,
         params: InsertParams,
+        cache_actions: list[CacheAction | ComponentAction],
     ) -> None:
         # _unevict_node_on_insert already wrote the request's fresh KV slice
         # into the base value. We just need to rebuild SWA from that slice for
@@ -253,20 +244,24 @@ class SWAComponent(TreeComponent):
         ), f"tombstone {ct} lock_ref should be 0 on unevict, node {node.id}"
         swa_evicted_seqlen = params.swa_evicted_seqlen
         assert (
-            swa_evicted_seqlen % self.cache.page_size == 0
+            swa_evicted_seqlen % self.tree_core.page_size == 0
         ), f"{ct}: swa_evicted_seqlen must be page-aligned, {swa_evicted_seqlen=}"
 
-        full_value = node.component_data[BASE_COMPONENT_TYPE].value
         if swa_evicted_seqlen <= total_prefix_len:
-            swa_value = self._translate_full_to_swa(full_value)
+            pass  # entire node is within the SWA window
         elif swa_evicted_seqlen < total_prefix_len + prefix_len:
             start_idx = swa_evicted_seqlen - total_prefix_len
-            self.cache._split_node(node.key, node, start_idx)
-            full_value = node.component_data[BASE_COMPONENT_TYPE].value
-            swa_value = self._translate_full_to_swa(full_value)
+            _, action = self.tree_core._split_node(node.key, node, start_idx)
+            if action is not None:
+                cache_actions.append(action)
         else:
             return
-        self._restore_device_value(node, swa_value)
+        cache_actions.append(
+            SWARebuild(
+                node.id,
+                node.component_data[BASE_COMPONENT_TYPE].value,
+            )
+        )
 
     def commit_insert_component_data(
         self,
@@ -274,57 +269,65 @@ class SWAComponent(TreeComponent):
         is_new_leaf: bool,
         params: InsertParams,
         result: InsertResult,
+        cache_actions: list[CacheAction | ComponentAction],
     ) -> None:
         if not is_new_leaf:
             return
 
         node_start = result.prefix_len
         split_pos = params.swa_evicted_seqlen - node_start
-
-        if split_pos <= 0:
-            swa_value = self._translate_full_to_swa(
-                node.component_data[BASE_COMPONENT_TYPE].value
-            )
-            node.component_data[self.component_type].value = swa_value
-            self.cache.lru_lists[self.component_type].insert_mru(node)
-            self.cache.component_evictable_size_[self.component_type] += len(swa_value)
-        elif split_pos < len(node.key):
-            # Node straddles the SWA eviction boundary
-            # Split into parent (tombstone, no SWA) and child (with SWA)
-            # After _split_node, `node` becomes the child
-            self.cache._split_node(node.key, node, split_pos)
-            swa_value = self._translate_full_to_swa(
-                node.component_data[BASE_COMPONENT_TYPE].value
-            )
-            node.component_data[self.component_type].value = swa_value
-            self.cache.lru_lists[self.component_type].insert_mru(node)
-            self.cache.component_evictable_size_[self.component_type] += len(swa_value)
-        else:
+        if split_pos >= len(node.key):
             # Entire leaf is outside the SWA window — left as a tombstone.
             return
+        if split_pos > 0:
+            # Node straddles the boundary: split into an out-of-window parent
+            # (tombstone) and an in-window child; `node` becomes the child.
+            _, action = self.tree_core._split_node(node.key, node, split_pos)
+            assert action is None, "new leaf cannot be write-through-pending"
+        # Cap the in-window leaf at one window for lock granularity, then rebuild SWA
+        # onto the in-window node(s) at apply time; rebuild the older prefix first so
+        # the in-window tail lands more-MRU.
+        capped_parent = self._maybe_split_leaf_for_swa_lock(node)
+        if capped_parent is not None:
+            cache_actions.append(
+                SWARebuild(
+                    capped_parent.id,
+                    capped_parent.component_data[BASE_COMPONENT_TYPE].value,
+                )
+            )
+        cache_actions.append(
+            SWARebuild(
+                node.id,
+                node.component_data[BASE_COMPONENT_TYPE].value,
+            )
+        )
 
-        self._maybe_split_leaf_for_swa_lock(node)
-
-    def _maybe_split_leaf_for_swa_lock(self, leaf: UnifiedTreeNode) -> None:
-        """Cap a fresh SWA leaf at one page-aligned window so locking it pins
-        only one window of SWA pool, not the whole (long chunked-prefill) leaf.
+    def _maybe_split_leaf_for_swa_lock(
+        self, leaf: UnifiedTreeNode
+    ) -> Optional[UnifiedTreeNode]:
+        """Cap a fresh in-window SWA leaf at one page-aligned window so locking it pins
+        only one window of SWA pool, not the whole (long chunked-prefill) leaf; return
+        the split-off parent (older window) or None. The SWA value is stamped later, so
+        this runs on the tombstone leaf.
         """
         ct = self.component_type
         cd = leaf.component_data[ct]
-        if leaf is self.cache.root_node or cd.value is None or cd.lock_ref > 0:
-            return
+        if leaf is self.tree_core.root_node or cd.lock_ref > 0:
+            return None
 
-        page_size = self.cache.page_size
+        page_size = self.tree_core.page_size
         # Smallest page-aligned size that still covers the sliding window.
         tail_size = (self.sliding_window_size + page_size - 1) // page_size * page_size
         leaf_len = len(leaf.key)
         if leaf_len <= tail_size:
-            return
+            return None
         split_at = leaf_len - tail_size
         if page_size > 1 and (split_at % page_size != 0 or leaf_len % page_size != 0):
-            return
+            return None
 
-        self.cache._split_node(leaf.key, leaf, split_at)
+        new_parent, action = self.tree_core._split_node(leaf.key, leaf, split_at)
+        assert action is None, "fresh SWA leaf cannot be write-through-pending"
+        return new_parent
 
     def redistribute_on_node_split(
         self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode
@@ -354,7 +357,7 @@ class SWAComponent(TreeComponent):
             child.component_data[self.component_type].host_value = child_swa_host_value[
                 split_len:
             ].clone()
-            host_lru = self.cache.host_lru_lists[self.component_type]
+            host_lru = self.tree_core.host_lru_lists[self.component_type]
             if new_parent.component_data[self.component_type].value is None:
                 host_lru.insert_mru(new_parent)
             if child.component_data[
@@ -371,6 +374,8 @@ class SWAComponent(TreeComponent):
     def evict_component(
         self,
         node: UnifiedTreeNode,
+        device_frees: dict[ComponentType, list[torch.Tensor]],
+        host_frees: dict[ComponentType, list[torch.Tensor]],
         target: EvictLayer = EvictLayer.DEVICE,
     ) -> tuple[int, int]:
         ct = self.component_type
@@ -383,19 +388,18 @@ class SWAComponent(TreeComponent):
             # Pass full indices to free_swa so slots with no SWA pair are
             # skipped. Freeing swa_value directly would double free those
             # entries since they all map to the same sentinel slot.
-            self.cache.token_to_kv_pool_allocator.free_swa(
+            device_frees[self.component_type].append(
                 node.component_data[BASE_COMPONENT_TYPE].value
             )
             freed = len(cd.value)
-            self.cache.component_evictable_size_[ct] -= freed
+            self.tree_core.component_evictable_size_[ct] -= freed
             cd.value = None
 
         # Host layer
-        host_lru = self.cache.host_lru_lists[ct]
+        host_lru = self.tree_core.host_lru_lists[ct]
         if EvictLayer.HOST in target and cd.host_value is not None:
             host_freed = len(cd.host_value)
-            if self._swa_kv_pool_host is not None:
-                self._swa_kv_pool_host.free(cd.host_value)
+            host_frees[ct].append(cd.host_value)
             cd.host_value = None
             if host_lru.in_list(node):
                 host_lru.remove_node(node)
@@ -414,30 +418,56 @@ class SWAComponent(TreeComponent):
     def eviction_priority(self, is_leaf: bool) -> int:
         return 0 if is_leaf else 1
 
-    def drive_eviction(
-        self, params: EvictParams, tracker: dict[ComponentType, int]
-    ) -> None:
-        request = params.swa_num_tokens
+    def _evict_device_start(self, request_cnt: int) -> None:
+        """Begin the device-eviction walk from this component's LRU cursor."""
+        self._evict_device_request_cnt = request_cnt
+        self._evict_device_cursor = self.tree_core.lru_lists[
+            self.component_type
+        ].get_lru_no_lock()
+
+    def _evict_device_next_node(
+        self,
+        tracker: dict[ComponentType, int],
+        device_frees: dict[ComponentType, list[torch.Tensor]],
+        host_frees: dict[ComponentType, list[torch.Tensor]],
+    ) -> Optional[NodeId]:
+        """Return the next device-leaf node for the driver to evict, or None.
+        Internal nodes are tombstoned inline (no IO); the cursor is re-validated
+        (reset to LRU head) if the previous node's eviction removed it."""
         ct = self.component_type
-        lru = self.cache.lru_lists[ct]
-        x = lru.get_lru_no_lock()
-        while tracker[ct] < request and x is not None and lru.in_list(x):
+        lru = self.tree_core.lru_lists[ct]
+        if self._evict_device_cursor is not None and not lru.in_list(
+            self._evict_device_cursor
+        ):
+            self._evict_device_cursor = lru.get_lru_no_lock()
+        while (
+            tracker[ct] < self._evict_device_request_cnt
+            and self._evict_device_cursor is not None
+            and lru.in_list(self._evict_device_cursor)
+        ):
+            x = self._evict_device_cursor
             assert x.component_data[ct].value is not None
-            if x in self.cache.evictable_device_leaves:
-                # D-leaf: atomic eviction of all components
-                x_next = lru.get_prev_no_lock(x)
-                self.cache._evict_device_leaf(x, tracker)
-                if not lru.in_list(x_next):
-                    x_next = lru.get_lru_no_lock()
-                x = x_next
-            else:
-                # Internal: tombstone SWA + cascade
-                x_next = lru.get_prev_no_lock(x)
-                self.cache._evict_component_and_detach_lru(
-                    x, self, target=EvictLayer.DEVICE, tracker=tracker
-                )
-                self.cache._cascade_evict(x, self, tracker)
-                x = x_next
+            if x in self.tree_core.evictable_device_leaves:
+                self._evict_device_cursor = lru.get_prev_no_lock(x)
+                return x.id
+            x_next = lru.get_prev_no_lock(x)
+            self.tree_core._evict_component_and_detach_lru(
+                x,
+                self,
+                target=EvictLayer.DEVICE,
+                tracker=tracker,
+                device_frees=device_frees,
+                host_frees=host_frees,
+            )
+            self.tree_core._cascade_evict(
+                x, self, tracker, device_frees=device_frees, host_frees=host_frees
+            )
+            self._evict_device_cursor = x_next
+        return None
+
+    def _evict_device_end(self) -> None:
+        """Clear the device-eviction walk cursor state."""
+        self._evict_device_cursor = None
 
     def acquire_component_lock(
         self,
@@ -446,12 +476,16 @@ class SWAComponent(TreeComponent):
         lock_host: bool = False,
     ) -> IncLockRefResult:
         ct = self.component_type
-        root = self.cache.root_node
+        root = self.tree_core.root_node
         sliding_window_size = self.sliding_window_size
         swa_lock_size = 0
         swa_uuid = None
         uuid_key = "host_uuid" if lock_host else "uuid"
-        lru = self.cache.host_lru_lists[ct] if lock_host else self.cache.lru_lists[ct]
+        lru = (
+            self.tree_core.host_lru_lists[ct]
+            if lock_host
+            else self.tree_core.lru_lists[ct]
+        )
 
         # Tombstoned nodes (cd.value is None) have no SWA chunk to protect
         # skip them and keep walking up. This path is hit when HiCache
@@ -472,8 +506,8 @@ class SWAComponent(TreeComponent):
                         lru.remove_node(cur)
                 else:
                     key_len = len(cur.key)
-                    self.cache.component_evictable_size_[ct] -= key_len
-                    self.cache.component_protected_size_[ct] += key_len
+                    self.tree_core.component_evictable_size_[ct] -= key_len
+                    self.tree_core.component_protected_size_[ct] += key_len
             if lock_host:
                 comp.host_lock_ref = ref + 1
             else:
@@ -498,7 +532,7 @@ class SWAComponent(TreeComponent):
         lock_host: bool = False,
     ) -> None:
         ct = self.component_type
-        root = self.cache.root_node
+        root = self.tree_core.root_node
         swa_uuid_for_lock = (
             (params.swa_uuid_for_host_lock if lock_host else params.swa_uuid_for_lock)
             if params
@@ -522,13 +556,13 @@ class SWAComponent(TreeComponent):
             if ref == 1:
                 if lock_host:
                     if comp.value is None and comp.host_value is not None:
-                        host_lru = self.cache.host_lru_lists[ct]
+                        host_lru = self.tree_core.host_lru_lists[ct]
                         if not host_lru.in_list(cur):
                             host_lru.insert_mru(cur)
                 else:
                     key_len = len(comp.value)
-                    self.cache.component_evictable_size_[ct] += key_len
-                    self.cache.component_protected_size_[ct] -= key_len
+                    self.tree_core.component_evictable_size_[ct] += key_len
+                    self.tree_core.component_protected_size_[ct] -= key_len
             if lock_host:
                 comp.host_lock_ref = ref - 1
             else:
@@ -540,7 +574,9 @@ class SWAComponent(TreeComponent):
     def release_window_lock(
         self,
         node: UnifiedTreeNode,
-        swa_uuid_for_lock: Optional[int] = None,
+        swa_uuid_for_lock: Optional[int],
+        device_frees: dict[ComponentType, list[torch.Tensor]],
+        host_frees: dict[ComponentType, list[torch.Tensor]],
     ) -> None:
         """Early-release the SWA lock along [node, swa_uuid_for_lock] while
         leaving Full and Mamba locks intact.
@@ -553,7 +589,7 @@ class SWAComponent(TreeComponent):
         invoked at most once per (node, swa_uuid_for_lock) pair.
         """
         ct = self.component_type
-        root = self.cache.root_node
+        root = self.tree_core.root_node
 
         cur = node
         while cur is not root:
@@ -569,11 +605,15 @@ class SWAComponent(TreeComponent):
             cd.lock_ref -= 1
             if cd.lock_ref == 0:
                 key_len = len(cur.key)
-                self.cache.component_protected_size_[ct] -= key_len
-                self.cache.component_evictable_size_[ct] += key_len
-                if self.cache._is_device_leaf(cur):
-                    self.cache._evict_component_and_detach_lru(
-                        cur, self, target=EvictLayer.DEVICE
+                self.tree_core.component_protected_size_[ct] -= key_len
+                self.tree_core.component_evictable_size_[ct] += key_len
+                if self.tree_core._is_device_leaf(cur):
+                    self.tree_core._evict_component_and_detach_lru(
+                        cur,
+                        self,
+                        target=EvictLayer.DEVICE,
+                        device_frees=device_frees,
+                        host_frees=host_frees,
                     )
 
             if swa_uuid_for_lock and cd.metadata.get("uuid") == swa_uuid_for_lock:
@@ -608,12 +648,36 @@ class SWAComponent(TreeComponent):
 
     # ---- HiCache Hooks ----
 
+    def prepare_prefetch(
+        self,
+        node_id: NodeId,
+        *,
+        prefetch_tokens: int = 0,
+    ) -> PreparePrefetchResult:
+        # unified_kv keeps SWA as a device-only ring -- nothing to prefetch into.
+        if self._swa_kv_pool_host is None:
+            return PreparePrefetchResult()
+        sw_pages = (
+            self.cache.sliding_window_size + self.cache.page_size - 1
+        ) // self.cache.page_size
+        if sw_pages == 0 or prefetch_tokens // self.cache.page_size < sw_pages:
+            return PreparePrefetchResult()
+        num_tokens = sw_pages * self.cache.page_size
+        host_indices = self._swa_kv_pool_host.alloc(num_tokens)
+        if host_indices is None:
+            self.cache.evict_host(num_tokens, ComponentType.SWA)
+            host_indices = self._swa_kv_pool_host.alloc(num_tokens)
+        if host_indices is None:
+            return PreparePrefetchResult(alloc_failed=True)
+        return PreparePrefetchResult(host_indices=host_indices)
+
     def build_hicache_transfers(
         self,
         node: UnifiedTreeNode,
         phase: CacheTransferPhase,
         *,
-        req: Optional[Req] = None,
+        mamba_pool_idx: Optional[torch.Tensor] = None,
+        host_indices: Optional[torch.Tensor] = None,
         token_ids: Optional[Sequence[int]] = None,
         prefetch_tokens: int = 0,
         last_hash: Optional[str] = None,
@@ -621,7 +685,7 @@ class SWAComponent(TreeComponent):
         ct = self.component_type
 
         # unified_kv keeps SWA as a device-only ring.
-        if self._swa_kv_pool_host is None and self.cache.cache_controller is not None:
+        if not self.tree_core.has_swa_host_pool and self.tree_core.enable_hicache:
             return None
 
         if phase == CacheTransferPhase.BACKUP_HOST:
@@ -644,7 +708,9 @@ class SWAComponent(TreeComponent):
             backed_up: list[torch.Tensor] = []
             nodes: list = []
             cur = node
-            while cur is not self.cache.root_node and n_swa < self.sliding_window_size:
+            while (
+                cur is not self.tree_core.root_node and n_swa < self.sliding_window_size
+            ):
                 cd = cur.component_data[ct]
                 assert cd.host_value is not None or cd.value is not None
                 if cd.value is not None:
@@ -668,7 +734,7 @@ class SWAComponent(TreeComponent):
                     name=PoolName.SWA,
                     host_indices=torch.cat(backed_up),
                     device_indices=None,
-                    nodes_to_load=nodes,
+                    nodes_to_load=[n.id for n in nodes],
                 )
             ]
 
@@ -676,32 +742,21 @@ class SWAComponent(TreeComponent):
             cd = node.component_data[ct]
             if cd.host_value is None or not node.hash_value:
                 return None
-            num_pages = len(cd.host_value) // self.cache.page_size
+            num_pages = len(cd.host_value) // self.tree_core.page_size
             if num_pages == 0:
                 return None
             return [
                 PoolTransfer(
                     name=PoolName.SWA,
-                    host_indices=cd.host_value[-num_pages * self.cache.page_size :],
+                    host_indices=cd.host_value[-num_pages * self.tree_core.page_size :],
                     keys=node.hash_value[-num_pages:],
                     hit_policy=PoolHitPolicy.TRAILING_PAGES,
                 )
             ]
 
         if phase == CacheTransferPhase.PREFETCH:
-            # Require a full sliding window.
-            sw_pages = (
-                self.sliding_window_size + self.cache.page_size - 1
-            ) // self.cache.page_size
-            if sw_pages == 0 or prefetch_tokens // self.cache.page_size < sw_pages:
-                return None
-            num_tokens = sw_pages * self.cache.page_size
-            host_indices = self._swa_kv_pool_host.alloc(num_tokens)
-            if host_indices is None:
-                self.cache.evict_host(num_tokens, ComponentType.SWA)
-                host_indices = self._swa_kv_pool_host.alloc(num_tokens)
-            if host_indices is None:
-                return []
+            assert host_indices is not None
+            sw_pages = host_indices.numel() // self.tree_core.page_size
             return [
                 PoolTransfer(
                     name=PoolName.SWA,
@@ -719,6 +774,7 @@ class SWAComponent(TreeComponent):
         phase: CacheTransferPhase,
         transfers: list[PoolTransfer] = (),
         *,
+        cache_actions: list[CacheAction | ComponentAction],
         insert_result: Optional[InsertResult] = None,
         pool_storage_result: Optional[PoolTransferResult] = None,
     ) -> None:
@@ -735,35 +791,47 @@ class SWAComponent(TreeComponent):
             assert transfers and transfers[0].device_indices is not None
             xfer = transfers[0]
             device_indices = xfer.device_indices
-            allocator = self.cache.token_to_kv_pool_allocator
 
+            full_chunks: list[torch.Tensor] = []
+            swa_chunks: list[torch.Tensor] = []
             offset = 0
-            for n in xfer.nodes_to_load or []:
+            for nid in xfer.nodes_to_load or []:
+                n = self.tree_core.node_by_id(nid)
                 cd_n = n.component_data[ct]
                 cd_full_n = n.component_data[BASE_COMPONENT_TYPE]
                 n_tokens = len(cd_n.host_value)
                 swa_chunk = device_indices[offset : offset + n_tokens].clone()
-                self._restore_device_value(n, swa_chunk)
+                self.tree_core.set_component_device_value(
+                    n.id, self.component_type, swa_chunk
+                )
                 assert cd_full_n.value is not None and len(cd_full_n.value) == n_tokens
-                # rebuild the mapping for the loaded SWA chunk
-                allocator.set_full_to_swa_mapping(cd_full_n.value, swa_chunk)
+                full_chunks.append(cd_full_n.value)
+                swa_chunks.append(swa_chunk)
                 offset += n_tokens
             assert offset == len(xfer.host_indices)
+            # rebuild the mapping for the loaded SWA chunk, defer to orchestrator level
+            if full_chunks:
+                cache_actions.append(RebuildFullToSWAMapping(full_chunks, swa_chunks))
             return
 
         if phase == CacheTransferPhase.PREFETCH:
             self._commit_prefetch(
                 node,
                 transfers,
+                cache_actions=cache_actions,
                 insert_result=insert_result,
                 pool_storage_result=pool_storage_result,
             )
             return
 
-    def _release_swa_host(self, host_indices: torch.Tensor) -> None:
+    def _release_swa_host(
+        self,
+        host_indices: torch.Tensor,
+        cache_actions: list[CacheAction | ComponentAction],
+    ) -> None:
         if host_indices is not None and host_indices.numel() > 0:
-            self.cache.cache_controller.append_host_mem_release(
-                extra_pools=[PoolTransfer(name=PoolName.SWA, host_indices=host_indices)]
+            cache_actions.append(
+                FreeComponentHostSlot([host_indices], component_type=ComponentType.SWA)
             )
 
     def _attach_swa_host_value(
@@ -773,18 +841,19 @@ class SWAComponent(TreeComponent):
         ct = self.component_type
         cd = node.component_data[ct]
         cd.host_value = host_indices.clone()
-        host_lru = self.cache.host_lru_lists[ct]
+        host_lru = self.tree_core.host_lru_lists[ct]
         if cd.value is None and not host_lru.in_list(node):
             host_lru.insert_mru(node)
-        self.cache._update_evictable_leaf_sets(node)
+        self.tree_core._update_evictable_leaf_sets(node)
         if node.parent:
-            self.cache._update_evictable_leaf_sets(node.parent)
+            self.tree_core._update_evictable_leaf_sets(node.parent)
 
     def _commit_prefetch(
         self,
         anchor,
         transfers: list[PoolTransfer],
         *,
+        cache_actions: list[CacheAction | ComponentAction],
         insert_result: Optional[InsertResult] = None,
         pool_storage_result: Optional[PoolTransferResult] = None,
     ) -> None:
@@ -799,7 +868,7 @@ class SWAComponent(TreeComponent):
         if not transfers:
             return
         ct = self.component_type
-        page_size = self.cache.page_size
+        page_size = self.tree_core.page_size
         host_indices = transfers[0].host_indices
         window_require_pages = (
             host_indices.numel() // page_size if host_indices is not None else 0
@@ -809,13 +878,18 @@ class SWAComponent(TreeComponent):
             if pool_storage_result
             else 0
         )
-        target = insert_result.inserted_host_node if insert_result else None
+        target = (
+            self.tree_core.node_by_id(insert_result.inserted_host_node)
+            if insert_result is not None
+            and insert_result.inserted_host_node is not None
+            else None
+        )
         if (
             target is None
             or window_require_pages == 0
             or loaded_pages < window_require_pages
         ):
-            self._release_swa_host(host_indices)
+            self._release_swa_host(host_indices, cache_actions)
             return
 
         # Buffer covers token range [loaded_start, total_len).
@@ -835,37 +909,105 @@ class SWAComponent(TreeComponent):
             if cd.host_value is None and fill_len > 0:
                 # Tombstone: split off the in-buffer tail if needed, then fill.
                 if fill_start > node_start:
-                    self.cache._split_node(cur.key, cur, fill_start - node_start)
+                    _, action = self.tree_core._split_node(
+                        cur.key, cur, fill_start - node_start
+                    )
+                    if action is not None:
+                        cache_actions.append(action)
                 self._attach_swa_host_value(cur, slice_)
             else:
                 # Already has SWA (or empty overlap): drop this slice.
-                self._release_swa_host(slice_)
+                self._release_swa_host(slice_, cache_actions)
 
             pos = node_start
             cur = cur.parent
 
         # Buffer prefix that fell outside the anchor→leaf path.
         if pos > loaded_start:
-            self._release_swa_host(host_indices[: pos - loaded_start])
+            self._release_swa_host(host_indices[: pos - loaded_start], cache_actions)
 
     def drive_host_eviction(
-        self, num_tokens: int, tracker: dict[ComponentType, int]
+        self,
+        num_tokens: int,
+        tracker: dict[ComponentType, int],
+        device_frees: dict[ComponentType, list[torch.Tensor]],
+        host_frees: dict[ComponentType, list[torch.Tensor]],
     ) -> None:
         """Evict SWA host resources.
         Internal nodes: private tombstone (free SWA host only).
         Host leaves: atomic eviction via _evict_host_leaf."""
         ct = self.component_type
-        host_lru = self.cache.host_lru_lists[ct]
+        host_lru = self.tree_core.host_lru_lists[ct]
         x = host_lru.get_lru_no_host_lock()
         while tracker[ct] < num_tokens and x is not None and host_lru.in_list(x):
             x_next = host_lru.get_prev_no_host_lock(x)
             cd = x.component_data[ct]
-            if x in self.cache.evictable_host_leaves:
-                self.cache._evict_host_leaf(x, tracker)
+            if x in self.tree_core.evictable_host_leaves:
+                self.tree_core._evict_host_leaf(x, tracker, device_frees, host_frees)
             else:
                 assert cd.host_value is not None
-                self.cache._evict_component_and_detach_lru(
-                    x, self, target=EvictLayer.HOST, tracker=tracker
+                self.tree_core._evict_component_and_detach_lru(
+                    x,
+                    self,
+                    target=EvictLayer.HOST,
+                    tracker=tracker,
+                    device_frees=device_frees,
+                    host_frees=host_frees,
+                )
+                self.tree_core._cascade_evict(
+                    x,
+                    self,
+                    tracker,
+                    device_frees=device_frees,
+                    host_frees=host_frees,
+                    target=EvictLayer.HOST,
                 )
-                self.cache._cascade_evict(x, self, tracker, target=EvictLayer.HOST)
             x = x_next
+
+    def free_host_values(self, host_values: list[torch.Tensor]) -> None:
+        if self._swa_kv_pool_host is None:
+            return
+        for host_value in host_values:
+            self._swa_kv_pool_host.free(host_value)
+
+    def apply_component_action(self, action: ComponentAction) -> None:
+        alloc = self.cache.token_to_kv_pool_allocator
+        if isinstance(action, FreeComponentDeviceSlot):
+            for indices in action.indices:
+                alloc.free_swa(indices)
+            return
+        if isinstance(action, FreeComponentHostSlot):
+            for host_indices in action.host_indices:
+                if host_indices is not None and host_indices.numel() > 0:
+                    self.cache.cache_controller.append_host_mem_release(
+                        extra_pools=[
+                            PoolTransfer(name=PoolName.SWA, host_indices=host_indices)
+                        ]
+                    )
+            return
+        if isinstance(action, RebuildFullToSWAMapping):
+            assert len(action.full_indices) == len(action.swa_indices)
+            for full, swa in zip(action.full_indices, action.swa_indices):
+                alloc.set_full_to_swa_mapping(full, swa)
+            return
+        if isinstance(action, RecoverSWAWithLockedFull):
+            # Keep the locked full; remap it onto the incoming full's SWA translation,
+            # freeing only the incoming full, then store the swa on the node.
+            swa_value = self._translate_full_to_swa(action.incoming_full)
+            alloc.set_full_to_swa_mapping(action.kept_full, swa_value)
+            alloc.full_to_swa_index_mapping[action.incoming_full.to(torch.int64)] = 0
+            alloc.full_attn_allocator.free(action.incoming_full)
+            self.tree_core.set_component_device_value(
+                action.node_id, self.component_type, swa_value
+            )
+            return
+        if isinstance(action, SWARebuild):
+            # Translate the node's source full value to SWA and store it on the node.
+            swa_value = self._translate_full_to_swa(action.source_value)
+            self.tree_core.set_component_device_value(
+                action.node_id, self.component_type, swa_value
+            )
+            return
+        raise AssertionError(
+            f"SWAComponent: unhandled ComponentAction {type(action).__name__}"
+        )

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

@@ -10,7 +10,6 @@ from numpy import float64
 
 from sglang.srt.mem_cache.base_prefix_cache import (
     DecLockRefParams,
-    EvictParams,
     IncLockRefResult,
     InsertParams,
     InsertResult,
@@ -18,40 +17,26 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     MatchResult,
 )
 from sglang.srt.mem_cache.hicache_storage import PoolTransfer, PoolTransferResult
+from sglang.srt.mem_cache.unified_cache.component_type import (  # noqa: F401
+    BASE_COMPONENT_TYPE,
+    ComponentType,
+)
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
     from sglang.srt.mem_cache.cache_init_params import CacheInitParams
+    from sglang.srt.mem_cache.unified_cache.cache_action import (
+        CacheAction,
+        ComponentAction,
+    )
+    from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore
     from sglang.srt.mem_cache.unified_radix_cache import (
+        NodeId,
         UnifiedRadixCache,
         UnifiedTreeNode,
     )
 
 
-class ComponentType(int, Enum):
-    """Integer enum so that per-node list/tuple storage can be indexed directly."""
-
-    FULL = 0
-    SWA = 1
-    MAMBA = 2
-
-    def __str__(self) -> str:  # keep human-readable logging
-        return self.name.lower()
-
-    @property
-    def is_full(self) -> bool:
-        return self == ComponentType.FULL
-
-    @property
-    def is_swa(self) -> bool:
-        return self == ComponentType.SWA
-
-    @property
-    def is_mamba(self) -> bool:
-        return self == ComponentType.MAMBA
-
-
-BASE_COMPONENT_TYPE = ComponentType.FULL
 _NUM_COMPONENT_TYPES = len(ComponentType)
 
 _LAST_ACCESS_TIME_COUNTER_FLOAT = float64(1.0)
@@ -83,6 +68,16 @@ class PrepareLoadBackResult:
     allocated_mamba_slot: Optional[torch.Tensor] = None
 
 
+@dataclasses.dataclass(frozen=True)
+class PreparePrefetchResult:
+    """Outcome of prepare_prefetch; default = nothing to prepare."""
+
+    # Host pool exhausted; the caller aborts the prefetch.
+    alloc_failed: bool = False
+    # The component's pre-allocated host buffer (None = skip the build).
+    host_indices: Optional[torch.Tensor] = None
+
+
 class CacheTransferPhase(str, Enum):
 
     BACKUP_HOST = "backup_host"  # D→H
@@ -114,6 +109,9 @@ def next_component_uuid() -> int:
 class TreeComponent(ABC):
     def __init__(self, cache: UnifiedRadixCache, params: CacheInitParams):
         self.cache = cache
+        # Populated when the component passed to TreeCore constructor.
+        self.tree_core: Optional[UnifiedTreeCore] = None
+        self.is_evict_device_ongoing = False
 
     # Subclasses MUST set this as a class attribute (not @property)
     component_type: ComponentType
@@ -130,6 +128,11 @@ class TreeComponent(ABC):
         value = node.component_data[self.component_type].value
         return len(value) if value is not None else 0
 
+    def has_host_value_only(self, node: UnifiedTreeNode) -> bool:
+        """Whether this component's data is evicted from device but host-backed."""
+        cd = node.component_data[self.component_type]
+        return cd.value is None and cd.host_value is not None
+
     def refresh_lru(
         self,
         phase: LRURefreshPhase,
@@ -141,9 +144,9 @@ class TreeComponent(ABC):
             case LRURefreshPhase.WALKDOWN:
                 if node.component_data[ct].value is None:
                     return
-                self.cache.lru_lists[ct].reset_node_mru(node)
+                self.tree_core.lru_lists[ct].reset_node_mru(node)
             case LRURefreshPhase.MATCH_END:
-                self.cache.lru_lists[ct].reset_node_and_parents_mru(
+                self.tree_core.lru_lists[ct].reset_node_and_parents_mru(
                     node, root_node, self.node_has_component_data
                 )
             case LRURefreshPhase.INSERT_END:
@@ -168,18 +171,22 @@ class TreeComponent(ABC):
         - Mamba: returns True iff the node has mamba component data."""
         ...
 
-    def finalize_match_result(
+    def finalize_match_result_in_tree_core(
         self,
         result: MatchResult,
         params: MatchPrefixParams,
         value_chunks: list[torch.Tensor],
         best_value_len: int,
     ) -> MatchResult:
-        """Post-process the match result after prefix matching completes.
-        - Full & SWA: pass through unchanged.
-        - Mamba: performs copy-on-write — allocates a new mamba slot, copies
-          the matched node's mamba state into the request pool, and records
-          branching_seqlen in result."""
+        """Tree-side post-processing inside the match walk (no cache access)."""
+        return result
+
+    def finalize_match_result_in_cache(
+        self, params: MatchPrefixParams, result: MatchResult
+    ) -> MatchResult:
+        """Cache-level finalize after the match walk, dispatched by
+        `UnifiedRadixCache.match_prefix`; receives the NodeId-based result.
+        - Mamba: performs the copy-on-write into a per-request slot."""
         return result
 
     def update_component_on_insert_overlap(
@@ -189,6 +196,7 @@ class TreeComponent(ABC):
         total_prefix_len: int,
         value_slice: torch.Tensor,
         params: InsertParams,
+        cache_actions: list[CacheAction | ComponentAction],
     ) -> int:
         """Called per-node when an insert's key overlaps an existing node.
         Returns the index within value_slice from which this component
@@ -204,6 +212,7 @@ class TreeComponent(ABC):
         prefix_len: int,
         total_prefix_len: int,
         params: InsertParams,
+        cache_actions: list[CacheAction | ComponentAction],
     ) -> None:
         """Called after _unevict_node_on_insert restores the base (Full) value
         on an evicted node. Aux components (e.g. SWA) override this to rebuild
@@ -217,6 +226,7 @@ class TreeComponent(ABC):
         is_new_leaf: bool,
         params: InsertParams,
         result: InsertResult,
+        cache_actions: list[CacheAction | ComponentAction],
     ) -> None:
         """Finalize component data on the target (leaf) node after the insert
         walk completes. Called once per insert.
@@ -250,6 +260,8 @@ class TreeComponent(ABC):
     def evict_component(
         self,
         node: UnifiedTreeNode,
+        device_frees: dict[ComponentType, list[torch.Tensor]],
+        host_frees: dict[ComponentType, list[torch.Tensor]],
         target: EvictLayer = EvictLayer.DEVICE,
     ) -> tuple[int, int]:
         """Free this component's KV resources on a node being evicted.
@@ -262,6 +274,9 @@ class TreeComponent(ABC):
           - ALL:    free both device and host memory.
                     No tombstone — caller will delete the node.
 
+        Freed indices are collected into *device_frees*/*host_frees* for the
+        Controller to drain, never freed inline.
+
         Returns (device_freed, host_freed) token counts."""
         ...
 
@@ -292,17 +307,53 @@ class TreeComponent(ABC):
         - Full evict internal: cascades to SWA + Mamba."""
         return 0
 
+    def evict_device_start(self, request_cnt: int) -> None:
+        """Begin this component's device-eviction walk (build its cursor/heap)."""
+        assert (
+            not self.is_evict_device_ongoing
+        ), f"{self.component_type} device eviction already in progress"
+        self._evict_device_start(request_cnt)
+        self.is_evict_device_ongoing = True
+
+    def evict_device_next_node(
+        self,
+        tracker: dict[ComponentType, int],
+        device_frees: dict[ComponentType, list[torch.Tensor]],
+        host_frees: dict[ComponentType, list[torch.Tensor]],
+    ) -> Optional[NodeId]:
+        """Return the next device-leaf node for the driver to evict, or None.
+        Internal nodes are tombstoned inline (no IO)."""
+        assert (
+            self.is_evict_device_ongoing
+        ), f"{self.component_type} device eviction not started"
+        return self._evict_device_next_node(tracker, device_frees, host_frees)
+
+    def evict_device_end(self) -> None:
+        """Clear this component's device-eviction walk state."""
+        assert (
+            self.is_evict_device_ongoing
+        ), f"{self.component_type} device eviction not started"
+        self._evict_device_end()
+        self.is_evict_device_ongoing = False
+
     @abstractmethod
-    def drive_eviction(
-        self, params: EvictParams, tracker: dict[ComponentType, int]
-    ) -> None:
-        """Drive eviction from this component's LRU list.
-        Each component extracts its own request from params, walks its own
-        LRU, evicts, and calls cache._cascade_evict for priority cascade.
-        Updates the shared tracker with freed amounts for all components.
-        - Full: walks leaf LRU, evicts full then cascades entire leaf.
-        - Mamba: walks full LRU; tombstones internal nodes (with cascade
-          to equal-priority components like swa), cascades leaves to all."""
+    def _evict_device_start(self, request_cnt: int) -> None:
+        """Build this component's eviction cursor/heap."""
+        ...
+
+    @abstractmethod
+    def _evict_device_next_node(
+        self,
+        tracker: dict[ComponentType, int],
+        device_frees: dict[ComponentType, list[torch.Tensor]],
+        host_frees: dict[ComponentType, list[torch.Tensor]],
+    ) -> Optional[NodeId]:
+        """Advance the walk; return the next device leaf or None."""
+        ...
+
+    @abstractmethod
+    def _evict_device_end(self) -> None:
+        """Clear this component's eviction cursor/heap."""
         ...
 
     @abstractmethod
@@ -386,11 +437,15 @@ class TreeComponent(ABC):
     ) -> None:
         pass
 
+    def free_host_values(self, host_values: list[torch.Tensor]) -> None:
+        """Free evicted host-tier values back to this component's host pool."""
+        raise NotImplementedError(f"{self.component_type} must free its host values")
+
     # ---- HiCache Hooks ----
 
     def prepare_load_back(
         self,
-        node: UnifiedTreeNode,
+        node_id: NodeId,
         *,
         req: Optional[Req] = None,
     ) -> PrepareLoadBackResult:
@@ -404,12 +459,22 @@ class TreeComponent(ABC):
         not go through."""
         pass
 
+    def prepare_prefetch(
+        self,
+        node_id: NodeId,
+        *,
+        prefetch_tokens: int = 0,
+    ) -> PreparePrefetchResult:
+        """Cache-level host pre-allocation before a prefetch builds its transfers."""
+        return PreparePrefetchResult()
+
     def build_hicache_transfers(
         self,
         node: UnifiedTreeNode,
         phase: CacheTransferPhase,
         *,
-        req: Optional[Req] = None,
+        mamba_pool_idx: Optional[torch.Tensor] = None,
+        host_indices: Optional[torch.Tensor] = None,
         token_ids: Optional[Sequence[int]] = None,
         prefetch_tokens: int = 0,
         last_hash: Optional[str] = None,
@@ -424,6 +489,7 @@ class TreeComponent(ABC):
         phase: CacheTransferPhase,
         transfers: list[PoolTransfer] = (),
         *,
+        cache_actions: list[CacheAction | ComponentAction],
         insert_result: Optional[InsertResult] = None,
         pool_storage_result: Optional[PoolTransferResult] = None,
     ) -> None:
@@ -431,9 +497,20 @@ class TreeComponent(ABC):
         pass
 
     def drive_host_eviction(
-        self, num_tokens: int, tracker: dict[ComponentType, int]
+        self,
+        num_tokens: int,
+        tracker: dict[ComponentType, int],
+        device_frees: dict[ComponentType, list[torch.Tensor]],
+        host_frees: dict[ComponentType, list[torch.Tensor]],
     ) -> None:
-        """Evict from this component's host-side resources.
+        """Evict from this component's host-side resources, collecting freed
+        values into *device_frees*/*host_frees* for the Controller to drain.
         Called by HostPoolGroup when the host pool is full.
         Default no-op for components without host storage."""
         pass
+
+    def apply_component_action(self, action: ComponentAction) -> None:
+        """Apply a component-routed cache action; dispatched by the cache."""
+        raise NotImplementedError(
+            f"{self.component_type} cannot apply {type(action).__name__}"
+        )