diff --git a/python/sglang/srt/mem_cache/unified_cache/components/README.md b/python/sglang/srt/mem_cache/unified_cache/components/README.md index 336e4d1759ab..14f2f15c1d53 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/README.md +++ b/python/sglang/srt/mem_cache/unified_cache/components/README.md @@ -111,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_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) + - Calls `finalize_match_result_in_tree_core()` per component (tree-side: Full/SWA host-hit sums, Mamba `branching_seqlen`); the cache then routes `finalize_match_result_in_cache()` per component post-walk (Mamba performs copy-on-write: allocates new pool slot, copies SSM state) --- @@ -127,7 +127,7 @@ Insert a key-value pair into the tree. | **Mutation** | Creates new leaf nodes; updates component data on overlapping nodes; frees duplicate KV indices; may split nodes; updates LRU lists and evictable sizes | | **Complexity** | **O(K + D·C)** | -**Algorithm detail** (`_insert_helper`): +**Algorithm detail** (the resumable insert steps: `_insert_walk_step` / `_insert_commit_step` / `_insert_tail_step`): 1. At each existing node, calls `_touch_node` → promotes to MRU via `node_has_component_data()` 2. If key diverges mid-node, calls `_split_node` → `redistribute_on_node_split()` per component 3. For each overlapping node, calls `update_component_on_insert_overlap()` per component — returns `consumed_from` index; the tree frees `value[dup_start:consumed_from]` as duplicate pool indices @@ -267,15 +267,15 @@ Each component implements these hooks. See `tree_component.py` for the ABC and d |------|---------|-----------|----------| | `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_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 | +| `finalize_match_result_in_cache()` | Cache-level finalize after the walk (receives the params + NodeId-based result), dispatched 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 | Hook | Purpose | Called By | Default | |------|---------|-----------|----------| -| `update_component_on_insert_overlap()` | Handle key overlap with an existing node during insert. Returns the index within `value_slice` from which this component consumed (took ownership of) pool slots. Full/Mamba: no consumption (`prefix_len`). SWA: may recover tombstoned nodes within the sliding window boundary. | `_insert_helper` | returns `prefix_len` | -| `recover_after_unevict()` | Rebuild auxiliary component data after `_unevict_node_on_insert()` restores a Full device value from fresh KV indices. SWA uses this to rebuild in-window SWA data. | `_insert_helper` | no-op | -| `commit_insert_component_data()` | Finalize component data on the target node after the insert walk completes. Full: no-op (handled by `_add_new_node`). SWA: checks window boundary, may split node — parent becomes tombstone, child gets SWA data. Mamba: sets mamba pool indices and inserts into Mamba LRU. | `_insert_helper` | no-op | +| `update_component_on_insert_overlap()` | Handle key overlap with an existing node during insert. Returns the index within `value_slice` from which this component consumed (took ownership of) pool slots. Full/Mamba: no consumption (`prefix_len`). SWA: may recover tombstoned nodes within the sliding window boundary. | `_insert_walk_step` | returns `prefix_len` | +| `recover_after_unevict()` | Rebuild auxiliary component data after `_unevict_node_on_insert()` restores a Full device value from fresh KV indices. SWA uses this to rebuild in-window SWA data. | `_insert_walk_step` | no-op | +| `commit_insert_component_data()` | Finalize component data on the target node after the insert walk completes. Full: no-op (handled by `_add_new_node`). SWA: checks window boundary, may split node — parent becomes tombstone, child gets SWA data. Mamba: sets mamba pool indices and inserts into Mamba LRU. | `_insert_commit_step` | no-op | ### Node Split diff --git a/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py b/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py index a4efe3e4a78d..bf46238fcf67 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py @@ -890,12 +890,11 @@ 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. + # Drain even if the walk raises so tombstoned slots are not leaked; + # the walk runs behind the tree-core interface (Rust runs it natively). try: - self._evict_excess_path_states( - self.tree_core.node_by_id(action.tail_node_id), - device_frees, - host_frees, + self.tree_core.evict_excess_path_states( + action.tail_node_id, device_frees, host_frees ) finally: self.cache._free_values(device_frees, host_frees) diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index e0110067cb0e..8627190a58fa 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -372,9 +372,6 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): per-component LRUs, the size/leaf bookkeeping, and the component drivers, plus ``reset()``. - TODO(Jialin): the tree operations still live on ``UnifiedRadixCache`` and - reach this state through its proxy properties; they migrate onto this class - as the TreeCore split completes. """ def __init__( @@ -497,6 +494,14 @@ def get_prefix_hash_values(self, node_id: NodeId) -> list[str]: node = self.node_by_id(node_id) return node.get_prefix_hash_values(node.parent) + def get_hash_values(self, node_id: NodeId) -> list[str]: + """The hash values owned by this node, excluding its ancestors.""" + return self.node_by_id(node_id).hash_value or [] + + def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId: + """The NodeId anchoring matches; the single root serves every namespace.""" + return self.root_node.id + def _new_node(self, priority: int = 0) -> UnifiedTreeNode: """Create and register a tree node in the arena.""" node = UnifiedTreeNode(self.component_types, priority=priority) @@ -1279,6 +1284,16 @@ def drive_host_eviction( ) return result + def evict_excess_path_states( + self, + tail_node_id: NodeId, + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> None: + self.components_by_type[ComponentType.MAMBA]._evict_excess_path_states( + self.node_by_id(tail_node_id), device_frees, host_frees + ) + def _evict_host_leaf( self, node: UnifiedTreeNode, diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py index cbd9910ffbed..2e16983880da 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py @@ -165,6 +165,16 @@ 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 get_hash_values(self, node_id: NodeId) -> list[str]: + """The hash values owned by this node, excluding its ancestors.""" + ... + + @abstractmethod + def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId: + """The NodeId anchoring matches for the namespace.""" + ... + @abstractmethod def inc_lock_ref( self, node_id: NodeId, skip_lock_components: Sequence[ComponentType] = () @@ -345,6 +355,17 @@ def drive_host_eviction( """Evict a component's host-side resources; no-op if the component is absent.""" ... + @abstractmethod + def evict_excess_path_states( + self, + tail_node_id: NodeId, + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> None: + """Evict shallow Mamba device checkpoints beyond the per-path cap on the + tail's root path, collecting freed values into the caller's dicts.""" + ... + # ==== HiCache ==== @abstractmethod diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index c11a02e30f68..d4362382658a 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -213,7 +213,10 @@ def __init__( self.hicache_storage_pass_prefix_keys = False self.reset() - logger.info(f"Init Unified RadixTree with components {self.tree_components}") + logger.info( + f"Init Unified Radix Cache. Components: {self.tree_components}. " + f"Tree Core: {type(self.tree_core).__name__}" + ) def _all_reduce_attn_groups(self, tensor: torch.Tensor, op): reduced = False @@ -2191,4 +2194,4 @@ def resolve_node_handle(self, node_handle): def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId: """The root's NodeId -- URC match results carry NodeIds.""" - return self.tree_core.root_node.id + return self.tree_core.root_node_handle(extra_key)