chore: import upstream snapshot with attribution

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

@@ -0,0 +1,325 @@
+# Unified Radix Cache
+
+A component-based, pluggable prefix cache framework for SGLang that unifies Full-attention, Sliding-Window-Attention (SWA), and Mamba/SSM caching into a single radix tree.
+
+## Design Goals
+
+1. **Unified tree structure** — One radix tree manages all KV cache types instead of separate specialized implementations (`SWARadixCache`, `MambaRadixCache`, etc.).
+2. **Pluggable components** — Each attention/state type (Full, SWA, Mamba) is a `TreeComponent` that implements hook interfaces. Adding a new cache type only requires adding a new component.
+3. **Per-component resource isolation** — Each component has its own lock reference counting, evictable/protected size tracking, and eviction driver. Auxiliary components use per-component LRUs; Full uses device/host leaf sets.
+4. **Cascade eviction with priority** — When a component evicts a node, lower-or-equal-priority components on the same node are evicted together, maintaining cross-component consistency.
+5. **Zero special-casing in the main tree** — The tree operates purely on keys (logical). All physical resource management (allocation, freeing, copy-on-write) is handled by components through hooks.
+
+## Architecture
+
+```
+┌───────────────────────────────────────────────┐
+│              UnifiedRadixCache                │
+│            (unified_radix_cache.py)           │
+│                                               │
+│  root_node ──► UnifiedTreeNode (radix tree)   │
+│  components ► {ComponentType → TreeComponent} │
+│  lru_lists ─► {ComponentType → UnifiedLRUList}│
+└──────────┬───────────┬───────────┬────────────┘
+           │           │           │
+           ▼           ▼           ▼
+   ┌────────────┐ ┌──────────┐ ┌─────────────┐
+   │    Full    │ │   SWA    │ │    Mamba    │
+   │ Component  │ │Component │ │  Component  │
+   └─────┬──────┘ └────┬─────┘ └──────┬──────┘
+         │             │              │
+         └─────────────┼──────────────┘
+                       ▼
+               ┌──────────────┐
+               │TreeComponent │
+               │    (ABC)     │
+               └──────────────┘
+```
+
+### Key Data Structures
+
+**`UnifiedTreeNode`** — Each node stores per-component data independently:
+
+```python
+node.component_data[ComponentType.FULL]   # FullComponent data
+node.component_data[ComponentType.SWA]    # SWAComponent data
+node.component_data[ComponentType.MAMBA]  # MambaComponent data
+```
+
+**`UnifiedLRUList`** — One doubly-linked list per auxiliary component, threaded through the same tree nodes via `lru_prev[component_type]`/`lru_next[component_type]`. Host LRUs use a separate pointer slot range so device and host lists do not collide. Supports O(1) insert/remove/promote and O(L) scan for eviction (L = locked nodes skipped). Full eviction is driven by `evictable_device_leaves` / `evictable_host_leaves`, not a Full LRU.
+
+**`ComponentData`** — Per-component data stored on each node:
+- `value: Tensor | None` — Device indices into the component's memory pool (`TokenToKVPool` for Full, `SWAKVPool` for SWA, `MambaPool` for Mamba). `None` means tombstone (data evicted but node structure retained).
+- `lock_ref: int` — Reference count of active requests using this node's component data. `lock_ref > 0` protects the node from eviction.
+- `metadata: dict` — Component-specific state (e.g., SWA stores `component_uuid` for window-lock boundary tracking).
+- `host_value: Tensor | None` — Host-side indices when HiCache has backed up this component.
+- `host_lock_ref: int` — Reference count protecting host-side component data from host eviction.
+
+---
+
+## File Layout
+
+| File | Contents |
+|------|----------|
+| `../unified_radix_cache.py` | `UnifiedRadixCache`, `UnifiedTreeNode`, `UnifiedLRUList` |
+| `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 |
+| `mamba_component.py` | `MambaComponent` — Mamba/SSM state component with copy-on-write |
+| `../hybrid_cache/hybrid_cache_controller.py` | `HybridCacheController` — HiCache multi-pool controller (L1 GPU → L2 CPU, optional L3 storage) |
+| `__init__.py` | Re-exports: `ComponentType`, `ComponentData`, `TreeComponent`, `FullComponent`, `SWAComponent`, `MambaComponent` |
+
+---
+
+## Public API Reference
+
+All public APIs are on `UnifiedRadixCache`, which implements `BasePrefixCache`.
+
+**Notation**: K = key length (tokens), D = matched path depth in tree (D ≤ K/P), P = page_size, C = number of components (≤ 3, treated as constant).
+
+All tree traversal operations have two cost components: **O(K)** for data operations (key comparison, tensor clone/concat) + **O(D·C)** for component overhead (C hooks per node). Since D ≤ K/P and C is constant, overall **O(K)**.
+
+### `match_prefix(params: MatchPrefixParams) → MatchResult`
+
+Find the longest cached prefix for a token sequence.
+
+| Aspect | Detail |
+|--------|--------|
+| **Purpose** | Walk the radix tree to find the longest prefix where **all** component validators pass |
+| **Inputs** | `params.key: RadixKey` — token IDs + optional extra key for namespace isolation |
+| **Output** | `MatchResult(device_indices, last_device_node, last_host_node, best_match_node, host_hit_length, mamba_branching_seqlen, ...)` |
+| **Mutation** | Updates `last_access_time` on matched path; promotes matched nodes to MRU in all component LRU lists; may trigger `_split_node` if match ends mid-node |
+| **Complexity** | **O(K + D·C)** |
+
+**Algorithm detail:**
+1. Calls `create_match_validator(match_device_only=...)` once per component — returns a stateful closure (e.g., SWA tracks accumulated window length). In HiCache mode, matching tracks both the best device-only node and the best device-or-host node.
+2. Walks tree edges via `RadixKey.match()`; at each node, calls all validator closures — the match boundary is only advanced when **all** validators return `True`
+3. If match ends mid-node, calls `_split_node` → triggers `redistribute_on_node_split()` per component
+4. Post-match (`_match_post_processor`):
+   - 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)
+
+---
+
+### `insert(params: InsertParams) → InsertResult`
+
+Insert a key-value pair into the tree.
+
+| Aspect | Detail |
+|--------|--------|
+| **Purpose** | Insert token sequence + KV indices, reusing existing prefix and freeing duplicate KV slots |
+| **Inputs** | `params.key: RadixKey`, `params.value: Tensor` (KV pool indices), plus component-specific fields (`mamba_value`, `swa_evicted_seqlen`, `prev_prefix_len`) |
+| **Output** | `InsertResult(prefix_len, mamba_exist)` — `prefix_len` = length of reused prefix |
+| **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`):
+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
+   - Full: returns `prefix_len` (no consumption, default behavior)
+   - SWA: checks if the overlapping node is a tombstone (SWA value = None) within the SWA window boundary (`swa_evicted_seqlen`):
+     - If entirely within window: **recovers tombstone** — frees old `full_value`, clones `value_slice`, translates to SWA indices, inserts into SWA LRU (returns `0` = all consumed)
+     - If partially within window: **splits node** at boundary, recovers SWA on the window portion (returns `start_idx`)
+     - If entirely outside window: returns `prefix_len` (no consumption)
+   - Mamba: returns `prefix_len` (no consumption, default behavior)
+4. Creates leaf via `_add_new_node` (clones value tensor, updates Full leaf-set tracking). A leaf survives on its Full value alone, so it is materialized even when an auxiliary component holds only a tombstone for the span (e.g. the whole leaf is outside the SWA window)
+5. Calls `commit_insert_component_data()` per component on the final target node (SWA may trigger a secondary split for window boundary; Mamba sets mamba pool indices and inserts into Mamba LRU)
+
+---
+
+### `evict(params: EvictParams) → EvictResult`
+
+Free cached tokens to reclaim memory.
+
+| Aspect | Detail |
+|--------|--------|
+| **Purpose** | Each component drives eviction from its own LRU list until its target is met |
+| **Inputs** | `params.num_tokens` (full), `params.swa_num_tokens` (SWA), `params.mamba_num` (Mamba) |
+| **Output** | `EvictResult(num_tokens_evicted, swa_num_tokens_evicted, mamba_num_evicted)` |
+| **Mutation** | Frees pool indices; removes nodes from LRU lists; deletes leaf nodes from tree; cascades to lower-priority components; walks up parent chain to delete tombstone ancestors |
+| **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:
+   - 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
+2. After each node eviction, calls `_cascade_evict`:
+   - Queries `eviction_priority()` per component; evicts all with priority ≤ trigger's
+   - Calls `evict_component()` + `node_has_component_data()` for cascaded components
+   - For leaf: removes from parent, then `_iteratively_delete_tombstone_leaf` walks up **O(H)** ancestors
+
+**Cascade eviction rules:**
+- **Leaf nodes**: all priorities = 0 → evicting any cascades to all (node deleted)
+- **Internal nodes**: Full(2) > SWA(1) > Mamba(0)
+  - Evicting Mamba: no cascade
+  - Evicting SWA: cascades to Mamba
+  - Evicting Full: cascades to SWA + Mamba
+
+---
+
+### `inc_lock_ref(node: UnifiedTreeNode) → IncLockRefResult`
+
+Lock a node to protect it (and its ancestors) from eviction.
+
+| Aspect | Detail |
+|--------|--------|
+| **Purpose** | Called when a request begins using a cached prefix — prevents eviction of nodes it depends on |
+| **Inputs** | `node` — the last matched node (deepest) |
+| **Output** | `IncLockRefResult(swa_uuid_for_lock)` |
+| **Mutation** | Increments `lock_ref` per component along the path; moves tokens from evictable to protected size counters |
+| **Complexity** | **O(D)** — Full: node to root; SWA: up to window boundary O(min(D, W)); Mamba: O(1).|
+
+**Algorithm detail:** Calls `acquire_component_lock()` for each component.
+
+| Component | Strategy |
+|-----------|----------|
+| Full | **Path-lock**: walks from node to root, `lock_ref += 1` on every ancestor. On first lock (`lock_ref: 0→1`), moves tokens from `component_evictable_size_` to `component_protected_size_`. |
+| SWA | **Window-lock**: walks upward, accumulating SWA value lengths until `sliding_window_size` is filled. Records a `component_uuid` at the boundary node for `dec_lock_ref` to know where to stop. |
+| Mamba | **Single-node lock**: only `lock_ref += 1` on the node itself (mamba state is per-leaf, not per-path). |
+
+---
+
+### `dec_lock_ref(node, params?) → DecLockRefResult`
+
+Unlock a previously locked node path.
+
+| Aspect | Detail |
+|--------|--------|
+| **Purpose** | Called when a request finishes — releases eviction protection |
+| **Inputs** | `node`, optional `params.swa_uuid_for_lock` for SWA boundary detection |
+| **Output** | `DecLockRefResult()` |
+| **Mutation** | Decrements `lock_ref` per component; moves tokens from protected back to evictable when `lock_ref` reaches 0 |
+| **Complexity** | **O(D)** — symmetric to `inc_lock_ref` |
+
+**Algorithm detail:** Calls `release_component_lock()` for each component. Full walks to root; SWA walks up until matching `component_uuid`; Mamba decrements single node.
+
+---
+
+### `cache_finished_req(req: Req, is_insert: bool = True)`
+
+Cache a completed request's KV data into the tree.
+
+| Aspect | Detail |
+|--------|--------|
+| **Purpose** | After a request finishes, insert its token/KV data into the tree for future reuse |
+| **Inputs** | `req` — the finished request; `is_insert` — whether to insert (True) or just release locks (False) |
+| **Output** | `None` |
+| **Mutation** | Calls component hooks → `insert` → `dec_lock_ref` → component cleanup. Frees unaligned tail KV indices; frees non-inserted KV indices when `is_insert=False`. |
+| **Complexity** | **O(K + D·C)** — insert O(K + D·C) + lock release O(D). Simplifies to **O(K)**. |
+
+**Algorithm detail:**
+1. `prepare_for_caching_req()` per component — sets component-specific insert params, returns effective cache length (SWA: sets `swa_evicted_seqlen`; Mamba: prepares `mamba_value` from ping-pong buffer, returns `mamba_last_track_seqlen` as truncation hint)
+2. Truncates if `effective_cache_len < len(token_ids)`: frees excess pool indices
+3. Converts token IDs (bigram if EAGLE), page-aligns keys, then calls `insert()`
+4. Frees unaligned tail KV indices beyond page boundary
+5. Calls `dec_lock_ref()` on the previous `req.last_node`
+6. `cleanup_after_caching_req()` per component (Mamba: frees forked mamba_value based on `mamba_exist`, handles ping-pong buffer cleanup)
+
+---
+
+### `cache_unfinished_req(req: Req, chunked=False)`
+
+Cache an in-progress request's partial KV data (chunked prefill).
+
+| Aspect | Detail |
+|--------|--------|
+| **Purpose** | During chunked prefill, insert partial results so the next chunk can match the prefix |
+| **Inputs** | `req` — the in-progress request |
+| **Output** | `None` |
+| **Mutation** | Inserts partial KV → re-matches prefix → updates `req.prefix_indices`, `req.cache_protected_len`, `req.last_node`; transfers lock from old node to new node |
+| **Complexity** | **O(K + D·C)** — two tree traversals: insert O(K + D·C) + re-match O(K + D·C) + lock transfer O(D). Simplifies to **O(K)**. |
+
+**Algorithm detail:**
+1. `prepare_for_caching_req()` per component
+2. `insert()` — first tree traversal
+3. `match_prefix()` — **second** tree traversal to get updated indices
+4. Writes new prefix indices into `req_to_token_pool`
+5. `dec_lock_ref()` on old `req.last_node`
+6. `inc_lock_ref()` on new matched node
+7. Updates `req.prefix_indices`, `req.cache_protected_len`, `req.last_node`
+8. `cleanup_after_caching_req()` per component
+
+---
+
+## TreeComponent Hook Reference
+
+Each component implements these hooks. See `tree_component.py` for the ABC and docstrings.
+
+### Match Phase
+
+| 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 |
+
+### 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 |
+
+### Node Split
+
+| Hook | Purpose | Called By | Default |
+|------|---------|-----------|----------|
+| `redistribute_on_node_split()` | Redistribute component data between new parent (prefix) and child (suffix) when a node is split. Full: copies `lock_ref` to parent. SWA: slices SWA value, copies `lock_ref` and `component_uuid`. Mamba: parent gets `None`/`lock_ref=0` (mamba stays on leaf). | `_split_node` | *abstract* |
+
+### Eviction Phase
+
+| Hook | Purpose | Called By | Default |
+|------|---------|-----------|----------|
+| `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 |
+
+### Lock Phase
+
+| Hook | Purpose | Called By | Default |
+|------|---------|-----------|----------|
+| `acquire_component_lock(lock_host=False)` | Increment device or host lock refs; moves device tokens from evictable to protected. Full: path-lock for device, single-node host lock. SWA: window-lock with UUID boundary. Mamba: single-node lock. | `inc_lock_ref`, `inc_host_lock_ref` | *abstract* |
+| `release_component_lock(lock_host=False)` | Decrement device or host lock refs; moves device tokens from protected to evictable when `lock_ref` → 0. Full path-unlocks device; SWA walks up to UUID boundary; Mamba unlocks a single node. | `dec_lock_ref`, `dec_host_lock_ref` | *abstract* |
+
+### Caching Phase
+
+| Hook | Purpose | Called By | Default |
+|------|---------|-----------|----------|
+| `prepare_for_caching_req()` | Prepare component-specific data before insert, fill fields in `InsertParams`, return effective cache length. Full: no-op. SWA: sets `swa_evicted_seqlen`. Mamba: prepares `mamba_value` from ping-pong buffer, returns `mamba_last_track_seqlen`. | `cache_finished/unfinished_req` | returns `None` |
+| `cleanup_after_caching_req()` | Post-cache cleanup. Full/SWA: no-op. Mamba: frees forked `mamba_value` based on `mamba_exist`, handles ping-pong buffer `keep_idx`, resets `mamba_last_track_seqlen` on unfinished. | `cache_finished/unfinished_req` | no-op |
+
+### Utility
+
+| Hook | Purpose | Called By | Default |
+|------|---------|-----------|----------|
+| `build_hicache_transfers()` | Build component-specific `PoolTransfer` descriptors for D→H, H→D, H→Storage, or Storage→H transfers. | HiCache paths | `None` |
+| `commit_hicache_transfer()` | Commit component state after a HiCache transfer finishes. | HiCache paths | no-op |
+| `node_has_component_data(target=EvictLayer.DEVICE)` | Check if a node has this component's device or host data. Used as filter for LRU operations and cascade checks. | multiple | `value is not None` / `host_value is not None` |
+
+### Component Behavior Summary
+
+| Behavior | FullComponent | SWAComponent | MambaComponent |
+|----------|--------------|-------------|----------------|
+| **Validator** | Full device data, or host backup in HiCache match | Tracks accumulated window; `True` when ≥ `sliding_window_size` | Mamba device data, or host backup in HiCache match |
+| **Lock strategy** | Path-lock (root → node) | Window-lock (up to window boundary, UUID-tagged) | Single-node lock |
+| **Internal eviction priority** | 2 (last) | 1 (middle) | 0 (first) |
+| **Split behavior** | Copy `lock_ref` to parent | Slice SWA value + copy UUID | Parent gets `None` (mamba stays on leaf) |
+| **Match finalize** | No-op | No-op | Copy-on-write: allocate new mamba slot, copy state |
+| **Drive eviction** | Full leaf sets → cascade all | SWA LRU → tombstone internal, cascade leaf | Mamba LRU → tombstone internal, cascade leaf |
+
+---
+
+## Construction
+
+`UnifiedRadixCache` is constructed directly from `mem_cache/registry.py` when `SGLANG_ENABLE_UNIFIED_RADIX_TREE` is enabled. The registry sets `params.tree_components` before construction:
+
+- Regular full-attention models → `(ComponentType.FULL,)`
+- Hybrid SWA models → `(ComponentType.FULL, ComponentType.SWA)`
+- Hybrid SSM/Mamba models → `(ComponentType.FULL, ComponentType.MAMBA)`
+
+When hierarchical cache is enabled, the registry calls `cache.init_hicache(server_args, params)` after construction.

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

@@ -0,0 +1,31 @@
+from sglang.srt.mem_cache.unified_cache_components.full_component import FullComponent
+from sglang.srt.mem_cache.unified_cache_components.mamba_component import MambaComponent
+from sglang.srt.mem_cache.unified_cache_components.swa_component import SWAComponent
+from sglang.srt.mem_cache.unified_cache_components.tree_component import (
+    _NUM_COMPONENT_TYPES,
+    BASE_COMPONENT_TYPE,
+    CacheTransferPhase,
+    ComponentData,
+    ComponentType,
+    EvictLayer,
+    LRURefreshPhase,
+    TreeComponent,
+    get_and_increase_time_counter,
+    next_component_uuid,
+)
+
+__all__ = [
+    "BASE_COMPONENT_TYPE",
+    "ComponentData",
+    "ComponentType",
+    "EvictLayer",
+    "FullComponent",
+    "CacheTransferPhase",
+    "LRURefreshPhase",
+    "MambaComponent",
+    "SWAComponent",
+    "TreeComponent",
+    "_NUM_COMPONENT_TYPES",
+    "next_component_uuid",
+    "get_and_increase_time_counter",
+]

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

@@ -0,0 +1,333 @@
+from __future__ import annotations
+
+import heapq
+from typing import TYPE_CHECKING, Callable, Optional, Sequence
+
+import torch
+
+from sglang.srt.mem_cache.base_prefix_cache import (
+    DecLockRefParams,
+    EvictParams,
+    IncLockRefResult,
+    InsertResult,
+    MatchPrefixParams,
+    MatchResult,
+)
+from sglang.srt.mem_cache.hicache_storage import (
+    PoolName,
+    PoolTransfer,
+    PoolTransferResult,
+)
+from sglang.srt.mem_cache.unified_cache_components.tree_component import (
+    CacheTransferPhase,
+    ComponentType,
+    EvictLayer,
+    TreeComponent,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+    from sglang.srt.mem_cache.unified_radix_cache import (
+        UnifiedTreeNode,
+    )
+
+
+class FullComponent(TreeComponent):
+    component_type = ComponentType.FULL
+
+    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
+
+    def create_match_validator(
+        self, match_device_only: bool = False
+    ) -> Callable[[UnifiedTreeNode], bool]:
+        if match_device_only:
+            return (
+                lambda node: node.component_data[self.component_type].value is not None
+            )
+
+        # HiCache: evicted + backuped nodes are valid match boundaries.
+        return lambda node: (
+            node.component_data[self.component_type].value is not None or node.backuped
+        )
+
+    def finalize_match_result(
+        self,
+        result: MatchResult,
+        params: MatchPrefixParams,
+        value_chunks: list[torch.Tensor],
+        best_value_len: int,
+    ) -> MatchResult:
+        # Compute Full KV host hit length: walk from last_host_node up to
+        # last_device_node, summing host_value lengths of evicted nodes.
+        ct = self.component_type
+        kv_host_hit = 0
+        node = result.best_match_node
+        root_node = self.cache.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:
+                kv_host_hit += len(full_host)
+            node = node.parent
+        if kv_host_hit > 0:
+            return result._replace(
+                host_hit_length=max(result.host_hit_length, kv_host_hit)
+            )
+        return result
+
+    def redistribute_on_node_split(
+        self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode
+    ):
+        ct = self.component_type
+        new_parent.component_data[ct].lock_ref = child.component_data[ct].lock_ref
+        child_cd = child.component_data[ct]
+        split_len = len(new_parent.key)
+        if child_cd.value is not None:
+            new_parent.component_data[ct].value = child_cd.value[:split_len].clone()
+            child_cd.value = child_cd.value[split_len:].clone()
+        if child_cd.host_value is not None:
+            new_parent.component_data[ct].host_value = child_cd.host_value[
+                :split_len
+            ].clone()
+            child_cd.host_value = child_cd.host_value[split_len:].clone()
+
+    def evict_component(
+        self,
+        node: UnifiedTreeNode,
+        target: EvictLayer = EvictLayer.DEVICE,
+    ) -> tuple[int, int]:
+        cd = node.component_data[self.component_type]
+        freed = 0
+        host_freed = 0
+
+        # Device layer
+        if EvictLayer.DEVICE in target and cd.value is not None:
+            self._free_full(cd.value)
+            freed = len(cd.value)
+            self.cache.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
+
+        # 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)
+            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
+        ]
+        heapq.heapify(heap)
+        ct = self.component_type
+        while tracker[ct] < request and heap:
+            _, x = heapq.heappop(heap)
+            if x not in self.cache.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),
+                )
+
+    def drive_host_eviction(
+        self, num_tokens: int, tracker: dict[ComponentType, int]
+    ) -> 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
+        ]
+        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:
+                continue
+            self.cache._evict_host_leaf(x, tracker)
+            if x.parent is not None and x.parent in self.cache.evictable_host_leaves:
+                heapq.heappush(
+                    heap,
+                    (self.cache.eviction_strategy.get_priority(x.parent), x.parent),
+                )
+
+    def acquire_component_lock(
+        self,
+        node: UnifiedTreeNode,
+        result: IncLockRefResult,
+        lock_host: bool = False,
+    ) -> IncLockRefResult:
+        ct = self.component_type
+
+        # Only the last host node needs to be protected.
+        if lock_host:
+            cd = node.component_data[ct]
+            if cd.host_value is None:
+                return result
+            cd.host_lock_ref += 1
+            self.cache._update_evictable_leaf_sets(node)
+            return result
+
+        root = self.cache.root_node
+        cur = node
+
+        # Skip the bottom evicted segment
+        while cur is not root and cur.component_data[ct].value is None:
+            result.skip_lock_node_ids.setdefault(ct, set()).add(cur.id)
+            cur = cur.parent
+
+        # Lock the device-on segment up to root
+        delta = 0
+        while cur is not root:
+            cd = cur.component_data[ct]
+            assert (
+                cd.value is not None
+            ), 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
+                delta += key_len
+            cd.lock_ref += 1
+            self.cache.evictable_device_leaves.discard(cur)
+            cur = cur.parent
+        result.delta = delta
+        return result
+
+    def release_component_lock(
+        self,
+        node: UnifiedTreeNode,
+        params: Optional[DecLockRefParams],
+        lock_host: bool = False,
+    ) -> None:
+        ct = self.component_type
+        if lock_host:
+            cd = node.component_data[ct]
+            if cd.host_value is None or cd.host_lock_ref == 0:
+                return
+            cd.host_lock_ref -= 1
+            self.cache._update_evictable_leaf_sets(node)
+            return
+
+        root = self.cache.root_node
+        skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
+        cur = node
+        while cur != root:
+            if cur.id in skip_lock_node_ids:
+                cur = cur.parent
+                continue
+            cd = cur.component_data[ct]
+            assert cd.value is not None
+            assert cd.lock_ref > 0
+
+            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
+            cd.lock_ref -= 1
+            if cd.lock_ref == 0:
+                self.cache._update_evictable_leaf_sets(cur)
+            cur = cur.parent
+
+    # ---- HiCache Hooks ----
+
+    def build_hicache_transfers(
+        self,
+        node: UnifiedTreeNode,
+        phase: CacheTransferPhase,
+        *,
+        req: Optional[Req] = None,
+        token_ids: Optional[Sequence[int]] = None,
+        prefetch_tokens: int = 0,
+        last_hash: Optional[str] = None,
+    ) -> Optional[list[PoolTransfer]]:
+        ct = self.component_type
+
+        if phase == CacheTransferPhase.BACKUP_HOST:
+            # Full KV backup is handled by the main flow
+            # (write_backup → cache_controller.write on host_value directly).
+            # No extra PoolTransfer needed.
+            return None
+
+        if phase == CacheTransferPhase.LOAD_BACK:
+            # `node` is best_match_node. FULL device evict only from leaves,
+            # so once we hit a device-on node, everything above is also device-on
+            backed_up: list[torch.Tensor] = []
+            nodes: list = []
+            cur = node
+            while cur.evicted:
+                cd = cur.component_data[ct]
+                assert cd.host_value is not None
+                backed_up.append(cd.host_value)
+                nodes.append(cur)
+                cur = cur.parent
+            backed_up.reverse()
+            nodes.reverse()
+            return [
+                PoolTransfer(
+                    name=PoolName.KV,
+                    host_indices=(
+                        torch.cat(backed_up)
+                        if backed_up
+                        else torch.empty((0,), dtype=torch.int64, device="cpu")
+                    ),
+                    device_indices=None,
+                    nodes_to_load=nodes,
+                )
+            ]
+
+        return None
+
+    def commit_hicache_transfer(
+        self,
+        node: UnifiedTreeNode,
+        phase: CacheTransferPhase,
+        transfers: list[PoolTransfer] = (),
+        *,
+        insert_result: Optional[InsertResult] = None,
+        pool_storage_result: Optional[PoolTransferResult] = None,
+    ) -> None:
+        ct = self.component_type
+
+        if phase == CacheTransferPhase.BACKUP_HOST:
+            if transfers and transfers[0].host_indices is not None:
+                node.component_data[ct].host_value = transfers[0].host_indices.clone()
+
+        elif phase == CacheTransferPhase.LOAD_BACK:
+            if not transfers or transfers[0].device_indices is None:
+                self.cache._update_evictable_leaf_sets(node)
+                return
+
+            xfer = transfers[0]
+            device_indices = xfer.device_indices
+            offset = 0
+            for n in xfer.nodes_to_load or []:
+                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.cache._update_evictable_leaf_sets(node)

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

@@ -0,0 +1,628 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable, Optional, Sequence
+
+import torch
+
+from sglang.srt.mem_cache.base_prefix_cache import (
+    DecLockRefParams,
+    EvictParams,
+    IncLockRefResult,
+    InsertParams,
+    InsertResult,
+    MatchPrefixParams,
+    MatchResult,
+)
+from sglang.srt.mem_cache.hicache_storage import (
+    PoolHitPolicy,
+    PoolName,
+    PoolTransfer,
+    PoolTransferResult,
+)
+from sglang.srt.mem_cache.unified_cache_components.tree_component import (
+    CacheTransferPhase,
+    ComponentType,
+    EvictLayer,
+    TreeComponent,
+    get_and_increase_time_counter,
+)
+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_radix_cache import (
+        UnifiedRadixCache,
+        UnifiedTreeNode,
+    )
+
+
+class MambaComponent(TreeComponent):
+    component_type = ComponentType.MAMBA
+
+    def __init__(self, cache: UnifiedRadixCache, params: CacheInitParams):
+        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)}"
+        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}"
+        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
+        # HiCache state
+        self._mamba_pool_host = None  # set to host mamba pool when HiCache enabled
+
+    def create_match_validator(
+        self, match_device_only: bool = False
+    ) -> Callable[[UnifiedTreeNode], bool]:
+        ct = self.component_type
+        if match_device_only:
+            return lambda node: node.component_data[ct].value is not None
+
+        # HiCache: evicted + backuped (host_value present) is also a valid match
+        return lambda node: (
+            node.component_data[ct].value is not None
+            or node.component_data[ct].host_value is not None
+        )
+
+    def finalize_match_result(
+        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
+
+        # HiCache can still use prefix matches and load back host-backed Mamba
+        # states. We temporarily skip branching-state fill in that mode and can
+        # add a HiCache-aware branching policy later.
+        if self.cache.cache_controller is None and len(value_chunks) > best_value_len:
+            chunk_size = get_server_args().mamba_cache_chunk_size
+            aligned_seqlen = (
+                sum(len(v) for v in value_chunks) // chunk_size
+            ) * chunk_size
+            branching_seqlen = aligned_seqlen if aligned_seqlen > 0 else None
+        else:
+            branching_seqlen = 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:
+            result = result._replace(
+                mamba_host_hit_length=max(result.mamba_host_hit_length, 1)
+            )
+
+        return result._replace(mamba_branching_seqlen=branching_seqlen)
+
+    def commit_insert_component_data(
+        self,
+        node: UnifiedTreeNode,
+        is_new_leaf: bool,
+        params: InsertParams,
+        result: InsertResult,
+    ) -> 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(
+                params.mamba_value
+            )
+            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]
+            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(
+                params.mamba_value
+            )
+            node.last_access_time = get_and_increase_time_counter()
+            return
+        self.cache.lru_lists[self.component_type].reset_node_mru(node)
+        node.last_access_time = get_and_increase_time_counter()
+        result.mamba_exist = True
+
+    def redistribute_on_node_split(
+        self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode
+    ):
+        ct = self.component_type
+        new_parent.component_data[ct].value = None
+        new_parent.component_data[ct].lock_ref = 0
+        # HiCache: mamba host_value stays on child (mamba = leaf-only data)
+        new_parent.component_data[ct].host_value = None
+        new_parent.component_data[ct].host_lock_ref = 0
+
+    def evict_component(
+        self,
+        node: UnifiedTreeNode,
+        target: EvictLayer = EvictLayer.DEVICE,
+    ) -> tuple[int, int]:
+        cd = node.component_data[self.component_type]
+        freed = 0
+        host_freed = 0
+
+        # Device layer
+        if EvictLayer.DEVICE in target and cd.value is not None:
+            self._free_mamba_value(cd.value)
+            freed = len(cd.value)
+            self.cache.component_evictable_size_[self.component_type] -= freed
+            cd.value = None
+
+        # Host layer
+        host_lru = self.cache.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)
+            cd.host_value = None
+            if host_lru.in_list(node):
+                host_lru.remove_node(node)
+
+        # After device tombstone: if only host_value remains, insert into host LRU
+        if (
+            target is EvictLayer.DEVICE
+            and cd.value is None
+            and cd.host_value is not None
+        ):
+            if not host_lru.in_list(node):
+                host_lru.insert_mru(node)
+
+        return freed, host_freed
+
+    def drive_eviction(
+        self, params: EvictParams, tracker: dict[ComponentType, int]
+    ) -> None:
+        request = params.mamba_num
+        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):
+            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
+
+    def acquire_component_lock(
+        self,
+        node: UnifiedTreeNode,
+        result: IncLockRefResult,
+        lock_host: bool = False,
+    ) -> IncLockRefResult:
+        ct = self.component_type
+        if node is self.cache.root_node:
+            return result
+        cd = node.component_data[ct]
+        value = cd.host_value if lock_host else cd.value
+        # A node in skip_lock_node_ids was a tombstone when this lock was acquired.
+        if value is None:
+            result.skip_lock_node_ids.setdefault(ct, set()).add(node.id)
+            return result
+
+        if lock_host:
+            if cd.host_lock_ref == 0:
+                host_lru = self.cache.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
+            cd.lock_ref += 1
+        return result
+
+    def release_component_lock(
+        self,
+        node: UnifiedTreeNode,
+        params: Optional[DecLockRefParams],
+        lock_host: bool = False,
+    ) -> None:
+        ct = self.component_type
+        if node is self.cache.root_node:
+            return
+        cd = node.component_data[ct]
+        skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
+        if node.id in skip_lock_node_ids:
+            return
+
+        value = cd.host_value if lock_host else cd.value
+        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]
+                if not host_lru.in_list(node):
+                    host_lru.insert_mru(node)
+            return
+
+        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
+            cd.lock_ref -= 1
+
+    def _alloc_mamba_slot(self) -> torch.Tensor:
+        """Allocate one mamba pool slot, evicting if necessary."""
+        slot = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
+        if slot is None:
+            self.cache.evict(EvictParams(num_tokens=0, mamba_num=1))
+            slot = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
+            assert slot is not None, "Can not alloc mamba cache"
+        return slot
+
+    @property
+    def int8_ckpt_pool(self):
+        return getattr(self.cache.req_to_token_pool, "mamba_ckpt_pool", None)
+
+    def _alloc_int8_ckpt_slot(self) -> torch.Tensor:
+        slot = self.int8_ckpt_pool.alloc(1)
+        if slot is None:
+            self.cache.evict(EvictParams(num_tokens=0, mamba_num=1))
+            slot = self.int8_ckpt_pool.alloc(1)
+            assert slot is not None, "Can not alloc int8 mamba checkpoint slot"
+        return slot
+
+    def _commit_int8_checkpoint(self, active_slots: torch.Tensor) -> torch.Tensor:
+        ckpt_slot = self._alloc_int8_ckpt_slot()
+        self.int8_ckpt_pool.store_from_active(
+            self.cache.req_to_token_pool.mamba_pool,
+            active_slots.view(-1),
+            ckpt_slot,
+        )
+        return ckpt_slot
+
+    def _free_mamba_value(self, mamba_value: torch.Tensor) -> None:
+        if self.int8_ckpt_pool is not None:
+            self.int8_ckpt_pool.free(mamba_value)
+        else:
+            self.cache.req_to_token_pool.mamba_allocator.free(mamba_value)
+
+    def prepare_for_caching_req(
+        self,
+        req: Req,
+        insert_params: InsertParams,
+        token_ids_len: int,
+        is_finished: bool,
+    ) -> Optional[int]:
+        if self.enable_mamba_extra_buffer:
+            cache_len = req.mamba_last_track_seqlen
+        else:
+            cache_len = token_ids_len
+            # ReplaySSM (no_buffer): `temporal[slot]` lags the live state by the
+            # slot's unflushed ring depth (`write_pos`), so on request finish cap
+            # the donate to the last flush boundary (where temporal is current)
+            # and reset the cursor, keeping the donated checkpoint consistent with
+            # its key length. page_size is asserted == 1, so no realign. Mirrors
+            # MambaRadixCache.cache_finished_req.
+            if is_finished:
+                write_pos_buf = (
+                    self.cache.req_to_token_pool.mamba_pool.replayssm_write_pos
+                )
+                if write_pos_buf is not None:
+                    cache_len -= int(write_pos_buf[req.mamba_pool_idx].item())
+                    write_pos_buf[req.mamba_pool_idx] = 0
+
+        if is_finished:
+            if cache_len is None:
+                cache_len = 0
+            if self.enable_mamba_extra_buffer:
+                keep_idx = self.cache.req_to_token_pool.get_mamba_ping_pong_keep_idx(
+                    req
+                )
+                active_value = (
+                    req.mamba_ping_pong_track_buffer[keep_idx].unsqueeze(-1).clone()
+                )
+            else:
+                active_value = req.mamba_pool_idx.unsqueeze(-1).clone()
+            if self.int8_ckpt_pool is not None:
+                insert_params.mamba_value = self._commit_int8_checkpoint(active_value)
+            else:
+                insert_params.mamba_value = active_value
+            return cache_len
+        else:
+            if cache_len is None:
+                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:
+                    new_slot = self._alloc_mamba_slot()
+                    src_active = (
+                        self.cache.req_to_token_pool.donate_mamba_ping_pong_slot(
+                            req, new_slot
+                        )
+                    )
+                    mamba_value_donated = self._commit_int8_checkpoint(src_active)
+                    self.cache.req_to_token_pool.mamba_allocator.free(src_active)
+                else:
+                    mamba_value_donated = self._commit_int8_checkpoint(
+                        req.mamba_pool_idx.view(-1)
+                    )
+            elif self.enable_mamba_extra_buffer:
+                new_slot = self._alloc_mamba_slot()
+                mamba_value_donated = (
+                    self.cache.req_to_token_pool.donate_mamba_ping_pong_slot(
+                        req, new_slot
+                    )
+                )
+            else:
+                mamba_value_donated = self._alloc_mamba_slot()
+                # mamba_pool is a pure PHYSICAL store; translate both slot ids
+                # virtual->physical (identity for the non-unified memory pool) first.
+                translate = self.cache.req_to_token_pool.translate_mamba_indices
+                self.cache.req_to_token_pool.mamba_pool.copy_from(
+                    translate(req.mamba_pool_idx.unsqueeze(0)),
+                    translate(mamba_value_donated),
+                )
+            insert_params.mamba_value = mamba_value_donated
+            return cache_len
+
+    def cleanup_after_caching_req(
+        self,
+        req: Req,
+        is_finished: bool,
+        insert_result: Optional[InsertResult] = None,
+        insert_params: Optional[InsertParams] = None,
+    ) -> None:
+        if is_finished:
+            mamba_value_inserted = (
+                insert_result is not None and not insert_result.mamba_exist
+            )
+            pool = self.cache.req_to_token_pool
+
+            if self.int8_ckpt_pool is not None:
+                insert_value_unused = (
+                    not mamba_value_inserted
+                    and insert_params is not None
+                    and insert_params.mamba_value is not None
+                )
+                if insert_value_unused:
+                    self._free_mamba_value(insert_params.mamba_value)
+                pool.free_mamba_cache(req)
+                return
+
+            if self.enable_mamba_extra_buffer:
+                keep_idx = (
+                    pool.get_mamba_ping_pong_keep_idx(req)
+                    if mamba_value_inserted
+                    else None
+                )
+                pool.free_mamba_cache(
+                    req, mamba_ping_pong_track_buffer_to_keep=keep_idx
+                )
+                return
+
+            if not mamba_value_inserted:
+                pool.free_mamba_cache(req)
+        else:
+            if insert_params.mamba_value is not None and (
+                insert_result is None or insert_result.mamba_exist
+            ):
+                self._free_mamba_value(insert_params.mamba_value)
+            req.mamba_last_track_seqlen = None
+
+    # ---- HiCache Hooks ----
+
+    def build_hicache_transfers(
+        self,
+        node: UnifiedTreeNode,
+        phase: CacheTransferPhase,
+        *,
+        req: Optional[Req] = None,
+        token_ids: Optional[Sequence[int]] = None,
+        prefetch_tokens: int = 0,
+        last_hash: Optional[str] = None,
+    ) -> Optional[list[PoolTransfer]]:
+        ct = self.component_type
+
+        if phase == CacheTransferPhase.BACKUP_HOST:
+            cd = node.component_data[ct]
+            if cd.value is None:
+                return None
+            return [
+                PoolTransfer(
+                    name=PoolName.MAMBA,
+                    device_indices=cd.value,
+                )
+            ]
+
+        if phase == CacheTransferPhase.LOAD_BACK:
+            transfers: list[PoolTransfer] = []
+
+            cd = node.component_data[ct]
+            if cd.value is not None:
+                return None
+
+            # restore single node if host_value exists
+            if cd.host_value is not None and cd.value is None:
+                transfers.append(
+                    PoolTransfer(
+                        name=PoolName.MAMBA,
+                        host_indices=cd.host_value,
+                        nodes_to_load=[node],
+                    )
+                )
+
+            # Per-request mamba CoW (H→D copy into request's device slot)
+            cd = node.component_data[ct]
+            if req is not None and cd.host_value is not None:
+                if req.mamba_pool_idx is None:
+                    dst = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
+                    if dst is None:
+                        self.cache.evict(EvictParams(num_tokens=0, mamba_num=1))
+                        dst = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
+                        assert dst is not None, "Cannot alloc mamba for load_back"
+                    req.mamba_pool_idx = dst[0]
+                transfers.append(
+                    PoolTransfer(
+                        name=PoolName.MAMBA,
+                        host_indices=cd.host_value,
+                        device_indices=req.mamba_pool_idx.unsqueeze(0),
+                    )
+                )
+
+            return transfers if transfers else None
+
+        if phase == CacheTransferPhase.BACKUP_STORAGE:
+            cd = node.component_data[ct]
+            if cd.host_value is None or not node.hash_value:
+                return None
+            return [
+                PoolTransfer(
+                    name=PoolName.MAMBA,
+                    host_indices=cd.host_value,
+                    keys=[node.hash_value[-1]],
+                    hit_policy=PoolHitPolicy.TRAILING_PAGES,
+                )
+            ]
+
+        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 []
+            return [
+                PoolTransfer(
+                    name=PoolName.MAMBA,
+                    host_indices=host_indices,
+                    keys=["__placeholder__"],
+                    hit_policy=PoolHitPolicy.TRAILING_PAGES,
+                )
+            ]
+
+        return None
+
+    def commit_hicache_transfer(
+        self,
+        node: UnifiedTreeNode,
+        phase: CacheTransferPhase,
+        transfers: list[PoolTransfer] = (),
+        *,
+        insert_result: Optional[InsertResult] = None,
+        pool_storage_result: Optional[PoolTransferResult] = None,
+    ) -> None:
+        ct = self.component_type
+
+        if phase == CacheTransferPhase.BACKUP_HOST:
+            if transfers and transfers[0].host_indices is not None:
+                cd = node.component_data[ct]
+                if cd.host_value is None:
+                    cd.host_value = transfers[0].host_indices.clone()
+
+        elif phase == CacheTransferPhase.LOAD_BACK:
+            if not transfers:
+                return
+            transfer = transfers[0]
+            if transfer.device_indices is not None:
+                cd = node.component_data[ct]
+                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]
+                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
+
+        elif phase == CacheTransferPhase.PREFETCH:
+            if not transfers:
+                return
+            transfer = transfers[0]
+            host_indices = transfer.host_indices
+            loaded = (
+                pool_storage_result is not None
+                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
+            )
+            if (
+                host_indices is None
+                or target_node is None
+                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]
+                )
+                if insert_result is not None:
+                    insert_result.mamba_exist = True
+                return
+
+            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]
+                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]
+    ) -> 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]
+        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:
+                # Host leaf: atomic eviction (all components host + delete)
+                self.cache._evict_host_leaf(x, tracker)
+            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.cache._cascade_evict(x, self, tracker, target=EvictLayer.HOST)
+                self.cache._update_evictable_leaf_sets(x)
+            x = x_next

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

@@ -0,0 +1,871 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable, Optional, Sequence
+
+import torch
+
+from sglang.srt.mem_cache.base_prefix_cache import (
+    DecLockRefParams,
+    EvictParams,
+    IncLockRefResult,
+    InsertParams,
+    InsertResult,
+    MatchPrefixParams,
+    MatchResult,
+)
+from sglang.srt.mem_cache.common import free_swa_out_of_window_slots
+from sglang.srt.mem_cache.hicache_storage import (
+    PoolHitPolicy,
+    PoolName,
+    PoolTransfer,
+    PoolTransferResult,
+)
+from sglang.srt.mem_cache.unified_cache_components.tree_component import (
+    BASE_COMPONENT_TYPE,
+    CacheTransferPhase,
+    ComponentType,
+    EvictLayer,
+    LRURefreshPhase,
+    TreeComponent,
+    next_component_uuid,
+)
+
+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_radix_cache import (
+        UnifiedRadixCache,
+        UnifiedTreeNode,
+    )
+
+
+class SWAComponent(TreeComponent):
+    """Sliding window attention component.
+
+    Each SWA node stores translated SWA pool indices as its component
+    value, independent of the full attention indices on the same tree node.
+    When SWA data is evicted from an internal node the node is tombstoned
+    — its SWA component value becomes None while the full attention
+    value stays intact.
+    """
+
+    def __init__(self, cache: UnifiedRadixCache, params: CacheInitParams):
+        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)}"
+        super().__init__(cache, params)
+        self.sliding_window_size = params.sliding_window_size
+        # HiCache state: set to host SWA pool when HiCache enabled
+        self._swa_kv_pool_host = None
+
+    component_type = ComponentType.SWA
+
+    def _translate_full_to_swa(self, full_indices: torch.Tensor) -> torch.Tensor:
+        return self.cache.token_to_kv_pool_allocator.translate_loc_from_full_to_swa(
+            full_indices
+        )
+
+    def refresh_lru(
+        self,
+        phase: LRURefreshPhase,
+        node: UnifiedTreeNode,
+        root_node: UnifiedTreeNode,
+    ) -> None:
+        match phase:
+            case LRURefreshPhase.WALKDOWN:
+                # Walk-down would refresh every visited ancestor to MRU,
+                # but most are outside the active sliding window and must
+                # stay evictable. Window-bounded refresh runs at
+                # MATCH_END / INSERT_END instead.
+                return
+            case LRURefreshPhase.MATCH_END | LRURefreshPhase.INSERT_END:
+                self.cache.lru_lists[
+                    self.component_type
+                ].reset_node_and_window_ancestors_mru(
+                    node,
+                    root_node,
+                    self.sliding_window_size + self.cache.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]:
+        sliding_window_size = self.sliding_window_size
+        ct = self.component_type
+        state = {"len": float("inf")}
+
+        # 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
+        )
+
+        def validator(node: UnifiedTreeNode) -> bool:
+            cd = node.component_data[ct]
+            # HiCache: a host-only tombstone is a valid match boundary too
+            # — load_back will restore SWA from host before use.
+            if cd.value is None and (match_device_only or cd.host_value is None):
+                state["len"] = 0
+                if swa_device_only_hicache and (node.backuped or not node.evicted):
+                    return True
+                return False
+            state["len"] += len(node.key)
+            return state["len"] >= sliding_window_size
+
+        return validator
+
+    def finalize_match_result(
+        self,
+        result: MatchResult,
+        params: MatchPrefixParams,
+        value_chunks: list[torch.Tensor],
+        best_value_len: int,
+    ) -> MatchResult:
+        ct = self.component_type
+        n_swa = 0
+        swa_host_hit = 0
+        node = result.best_match_node
+        root = self.cache.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:
+                n_swa += len(cd.value)
+            elif cd.host_value is not None:
+                # TODO(hzh): load_back may currently restore a full host-tombstone
+                # segment whose length exceeds sliding_window_size. Once
+                # load_back is constrained to fetch only one sliding window
+                # worth of pages, cap swa_host_hit at sliding_window_size
+                # here so the scheduler budget matches the actual device-pool
+                # consumption.
+                swa_host_hit += len(cd.host_value)
+                n_swa += len(cd.host_value)
+            else:
+                break
+            node = node.parent
+        if swa_host_hit > 0:
+            return result._replace(
+                swa_host_hit_length=max(result.swa_host_hit_length, swa_host_hit)
+            )
+        return result
+
+    def update_component_on_insert_overlap(
+        self,
+        node: UnifiedTreeNode,
+        prefix_len: int,
+        total_prefix_len: int,
+        value_slice: torch.Tensor,
+        params: InsertParams,
+    ) -> int:
+        if params.prev_prefix_len >= total_prefix_len + prefix_len:
+            return prefix_len
+
+        is_tombstone = node.component_data[self.component_type].value is None
+        if not is_tombstone:
+            return prefix_len
+
+        full_cd = node.component_data[BASE_COMPONENT_TYPE]
+        swa_evicted_seqlen = params.swa_evicted_seqlen
+        assert (
+            node.component_data[self.component_type].lock_ref == 0
+        ), f"tombstone {self.component_type} lock_ref should be 0, node {node.id}"
+        assert (
+            swa_evicted_seqlen % self.cache.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
+            if full_cd.lock_ref > 0:
+                self._restore_device_value_with_locked_full(
+                    node, full_cd.value, 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)
+            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:]
+                )
+                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)
+            return start_idx
+        else:
+            # Branch 3: entire value_slice is outside SWA window — not consumed
+            return prefix_len
+
+    def recover_after_unevict(
+        self,
+        node: UnifiedTreeNode,
+        prefix_len: int,
+        total_prefix_len: int,
+        params: InsertParams,
+    ) -> 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
+        # the in-window portion. There is no old SWA slot to free here.
+        ct = self.component_type
+        if node.component_data[ct].value is not None:
+            return
+        assert (
+            node.component_data[ct].lock_ref == 0
+        ), 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
+        ), 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)
+        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)
+        else:
+            return
+        self._restore_device_value(node, swa_value)
+
+    def commit_insert_component_data(
+        self,
+        node: UnifiedTreeNode,
+        is_new_leaf: bool,
+        params: InsertParams,
+        result: InsertResult,
+    ) -> 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:
+            # Entire leaf is outside the SWA window — left as a tombstone.
+            return
+
+        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.
+        """
+        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
+
+        page_size = self.cache.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
+        split_at = leaf_len - tail_size
+        if page_size > 1 and (split_at % page_size != 0 or leaf_len % page_size != 0):
+            return
+
+        self.cache._split_node(leaf.key, leaf, split_at)
+
+    def redistribute_on_node_split(
+        self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode
+    ):
+        new_parent.component_data[self.component_type].lock_ref = child.component_data[
+            self.component_type
+        ].lock_ref
+
+        child_swa_value = child.component_data[self.component_type].value
+        if child_swa_value is not None:
+            split_len = len(new_parent.key)
+            new_parent.component_data[self.component_type].value = child_swa_value[
+                :split_len
+            ].clone()
+            child.component_data[self.component_type].value = child_swa_value[
+                split_len:
+            ].clone()
+        else:
+            new_parent.component_data[self.component_type].value = None
+
+        child_swa_host_value = child.component_data[self.component_type].host_value
+        if child_swa_host_value is not None:
+            split_len = len(new_parent.key)
+            new_parent.component_data[self.component_type].host_value = (
+                child_swa_host_value[:split_len].clone()
+            )
+            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]
+            if new_parent.component_data[self.component_type].value is None:
+                host_lru.insert_mru(new_parent)
+            if child.component_data[
+                self.component_type
+            ].value is None and not host_lru.in_list(child):
+                host_lru.insert_mru(child)
+
+        # parent inherits the swa_uuid from child for swa lock ref
+        new_parent.component_data[self.component_type].metadata["uuid"] = (
+            child.component_data[self.component_type].metadata.get("uuid")
+        )
+        child.component_data[self.component_type].metadata.pop("uuid", None)
+
+    def evict_component(
+        self,
+        node: UnifiedTreeNode,
+        target: EvictLayer = EvictLayer.DEVICE,
+    ) -> tuple[int, int]:
+        ct = self.component_type
+        cd = node.component_data[ct]
+        freed = 0
+        host_freed = 0
+
+        # Device layer
+        if EvictLayer.DEVICE in target and cd.value is not None:
+            # 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(
+                node.component_data[BASE_COMPONENT_TYPE].value
+            )
+            freed = len(cd.value)
+            self.cache.component_evictable_size_[ct] -= freed
+            cd.value = None
+
+        # Host layer
+        host_lru = self.cache.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)
+            cd.host_value = None
+            if host_lru.in_list(node):
+                host_lru.remove_node(node)
+
+        # After device tombstone: if host_value remains, move into host LRU
+        if (
+            target is EvictLayer.DEVICE
+            and cd.value is None
+            and cd.host_value is not None
+        ):
+            if not host_lru.in_list(node):
+                host_lru.insert_mru(node)
+
+        return freed, host_freed
+
+    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
+        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):
+            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
+
+    def acquire_component_lock(
+        self,
+        node: UnifiedTreeNode,
+        result: IncLockRefResult,
+        lock_host: bool = False,
+    ) -> IncLockRefResult:
+        ct = self.component_type
+        root = self.cache.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]
+
+        # Tombstoned nodes (cd.value is None) have no SWA chunk to protect
+        # skip them and keep walking up. This path is hit when HiCache
+        # backs up a FULL present internal node whose SWA was already evicted.
+        cur = node
+        while cur != root and swa_lock_size < sliding_window_size:
+            comp = cur.component_data[ct]
+            value = comp.host_value if lock_host else comp.value
+            if value is None:
+                result.skip_lock_node_ids.setdefault(ct, set()).add(cur.id)
+                cur = cur.parent
+                continue
+
+            ref = comp.host_lock_ref if lock_host else comp.lock_ref
+            if ref == 0:
+                if lock_host:
+                    if lru.in_list(cur):
+                        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
+            if lock_host:
+                comp.host_lock_ref = ref + 1
+            else:
+                comp.lock_ref = ref + 1
+            swa_lock_size += len(value)
+            if swa_lock_size >= sliding_window_size:
+                if comp.metadata.get(uuid_key) is None:
+                    comp.metadata[uuid_key] = next_component_uuid()
+                swa_uuid = comp.metadata[uuid_key]
+            cur = cur.parent
+
+        if lock_host:
+            result.swa_uuid_for_host_lock = swa_uuid
+        else:
+            result.swa_uuid_for_lock = swa_uuid
+        return result
+
+    def release_component_lock(
+        self,
+        node: UnifiedTreeNode,
+        params: Optional[DecLockRefParams],
+        lock_host: bool = False,
+    ) -> None:
+        ct = self.component_type
+        root = self.cache.root_node
+        swa_uuid_for_lock = (
+            (params.swa_uuid_for_host_lock if lock_host else params.swa_uuid_for_lock)
+            if params
+            else None
+        )
+        skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
+        dec_swa = True
+        uuid_key = "host_uuid" if lock_host else "uuid"
+
+        # A node in skip_lock_node_ids was a tombstone when this lock was acquired.
+        cur = node
+        while cur != root and dec_swa:
+            comp = cur.component_data[ct]
+            if cur.id in skip_lock_node_ids:
+                cur = cur.parent
+                continue
+            ref = comp.host_lock_ref if lock_host else comp.lock_ref
+            if ref == 0:
+                cur = cur.parent
+                continue
+            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]
+                        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
+            if lock_host:
+                comp.host_lock_ref = ref - 1
+            else:
+                comp.lock_ref = ref - 1
+            if swa_uuid_for_lock and comp.metadata.get(uuid_key) == swa_uuid_for_lock:
+                dec_swa = False
+            cur = cur.parent
+
+    def release_window_lock(
+        self,
+        node: UnifiedTreeNode,
+        swa_uuid_for_lock: Optional[int] = None,
+    ) -> None:
+        """Early-release the SWA lock along [node, swa_uuid_for_lock] while
+        leaving Full and Mamba locks intact.
+
+        Called when a request's decode position has advanced past the sliding
+        window — the SWA portion of the tree lock is no longer needed but the
+        Full lock must stay so the request's prefix is protected.
+
+        Caller (UnifiedRadixCache.dec_swa_lock_only) must ensure this is
+        invoked at most once per (node, swa_uuid_for_lock) pair.
+        """
+        ct = self.component_type
+        root = self.cache.root_node
+
+        cur = node
+        while cur is not root:
+            cd = cur.component_data[ct]
+            # Acquire skips tombstoned nodes; release must skip them too. Same
+            # for nodes with lock_ref == 0 — acquire never credited them.
+            if cd.value is None or cd.lock_ref == 0:
+                if swa_uuid_for_lock and cd.metadata.get("uuid") == swa_uuid_for_lock:
+                    break
+                cur = cur.parent
+                continue
+
+            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
+                    )
+
+            if swa_uuid_for_lock and cd.metadata.get("uuid") == swa_uuid_for_lock:
+                break
+            cur = cur.parent
+
+    def prepare_for_caching_req(
+        self,
+        req: Req,
+        insert_params: InsertParams,
+        token_ids_len: int,
+        is_finished: bool,
+    ) -> Optional[int]:
+        # Unfinished requests can already have an SWA-evicted prefix; preserve
+        # that boundary so insertion creates a tombstone instead of live SWA KV.
+        insert_params.swa_evicted_seqlen = req.swa_evicted_seqlen
+        return None
+
+    def free_out_of_window_slots(
+        self, req: Req, pre_len: int, insert_params: InsertParams
+    ) -> None:
+        if self.sliding_window_size is not None:
+            free_swa_out_of_window_slots(
+                req,
+                pre_len,
+                sliding_window_size=self.sliding_window_size,
+                page_size=self.cache.page_size,
+                req_to_token_pool=self.cache.req_to_token_pool,
+                token_to_kv_pool_allocator=self.cache.token_to_kv_pool_allocator,
+            )
+        insert_params.swa_evicted_seqlen = req.swa_evicted_seqlen
+
+    # ---- HiCache Hooks ----
+
+    def build_hicache_transfers(
+        self,
+        node: UnifiedTreeNode,
+        phase: CacheTransferPhase,
+        *,
+        req: Optional[Req] = None,
+        token_ids: Optional[Sequence[int]] = None,
+        prefetch_tokens: int = 0,
+        last_hash: Optional[str] = None,
+    ) -> Optional[list[PoolTransfer]]:
+        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:
+            return None
+
+        if phase == CacheTransferPhase.BACKUP_HOST:
+            cd = node.component_data[ct]
+            if cd.value is None:
+                return None
+            # cd.value already holds SWA-pool indices (translated at insert time).
+            # Host pool indexing wants int64.
+            return [
+                PoolTransfer(
+                    name=PoolName.SWA,
+                    device_indices=cd.value.to(torch.int64),
+                )
+            ]
+
+        if phase == CacheTransferPhase.LOAD_BACK:
+            # `node` is best_match_node; the SWA validator guarantees every
+            # ancestor within `sliding_window_size` has value or host_value.
+            n_swa = 0
+            backed_up: list[torch.Tensor] = []
+            nodes: list = []
+            cur = node
+            while cur is not self.cache.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:
+                    # device exists, skip it
+                    n_swa += len(cd.value)
+                else:
+                    # host only, collect it
+                    backed_up.append(cd.host_value)
+                    nodes.append(cur)
+                    n_swa += len(cd.host_value)
+                cur = cur.parent
+
+            if not backed_up:
+                return None
+
+            backed_up.reverse()
+            nodes.reverse()
+
+            return [
+                PoolTransfer(
+                    name=PoolName.SWA,
+                    host_indices=torch.cat(backed_up),
+                    device_indices=None,
+                    nodes_to_load=nodes,
+                )
+            ]
+
+        if phase == CacheTransferPhase.BACKUP_STORAGE:
+            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
+            if num_pages == 0:
+                return None
+            return [
+                PoolTransfer(
+                    name=PoolName.SWA,
+                    host_indices=cd.host_value[-num_pages * self.cache.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 []
+            return [
+                PoolTransfer(
+                    name=PoolName.SWA,
+                    host_indices=host_indices,
+                    keys=["__placeholder__"] * sw_pages,
+                    hit_policy=PoolHitPolicy.TRAILING_PAGES,
+                )
+            ]
+
+        return None
+
+    def commit_hicache_transfer(
+        self,
+        node: UnifiedTreeNode,
+        phase: CacheTransferPhase,
+        transfers: list[PoolTransfer] = (),
+        *,
+        insert_result: Optional[InsertResult] = None,
+        pool_storage_result: Optional[PoolTransferResult] = None,
+    ) -> None:
+        ct = self.component_type
+
+        if phase == CacheTransferPhase.BACKUP_HOST:
+            if transfers and transfers[0].host_indices is not None:
+                cd = node.component_data[ct]
+                if cd.host_value is None:
+                    cd.host_value = transfers[0].host_indices.clone()
+            return
+
+        if phase == CacheTransferPhase.LOAD_BACK:
+            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
+
+            offset = 0
+            for n in xfer.nodes_to_load or []:
+                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)
+                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)
+                offset += n_tokens
+            assert offset == len(xfer.host_indices)
+            return
+
+        if phase == CacheTransferPhase.PREFETCH:
+            self._commit_prefetch(
+                node,
+                transfers,
+                insert_result=insert_result,
+                pool_storage_result=pool_storage_result,
+            )
+            return
+
+    def _release_swa_host(self, host_indices: torch.Tensor) -> 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)]
+            )
+
+    def _attach_swa_host_value(
+        self, node: UnifiedTreeNode, host_indices: torch.Tensor
+    ) -> None:
+        """Write host_indices into node's SWA host_value and refresh tree state."""
+        ct = self.component_type
+        cd = node.component_data[ct]
+        cd.host_value = host_indices.clone()
+        host_lru = self.cache.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)
+        if node.parent:
+            self.cache._update_evictable_leaf_sets(node.parent)
+
+    def _commit_prefetch(
+        self,
+        anchor,
+        transfers: list[PoolTransfer],
+        *,
+        insert_result: Optional[InsertResult] = None,
+        pool_storage_result: Optional[PoolTransferResult] = None,
+    ) -> None:
+        """Fill the prefetched SWA window onto the leaf→anchor path.
+
+        All-or-nothing over one full window: ``loaded_pages`` is the cross-rank
+        MIN, so ``loaded_pages < window_pages`` drops the whole window (keeps the
+        tree identical across TP ranks). Otherwise map the buffer to token range
+        ``[loaded_start, total_len)`` and walk leaf→anchor, filling SWA
+        tombstones and releasing slices that already have host_value.
+        """
+        if not transfers:
+            return
+        ct = self.component_type
+        page_size = self.cache.page_size
+        host_indices = transfers[0].host_indices
+        window_require_pages = (
+            host_indices.numel() // page_size if host_indices is not None else 0
+        )
+        loaded_pages = (
+            pool_storage_result.extra_pool_hit_pages.get(PoolName.SWA, 0)
+            if pool_storage_result
+            else 0
+        )
+        target = insert_result.inserted_host_node if insert_result else None
+        if (
+            target is None
+            or window_require_pages == 0
+            or loaded_pages < window_require_pages
+        ):
+            self._release_swa_host(host_indices)
+            return
+
+        # Buffer covers token range [loaded_start, total_len).
+        loaded_start = insert_result.total_len - window_require_pages * page_size
+
+        # Walk leaf → anchor; ``pos`` is the right edge of ``cur`` in tokens.
+        pos, cur = insert_result.total_len, target
+        while cur is not anchor and pos > loaded_start:
+            node_start = pos - len(cur.key)
+            # Intersection of cur's range and the buffer.
+            fill_start = max(node_start, loaded_start)
+            fill_len = pos - fill_start
+            buf_off = fill_start - loaded_start
+            slice_ = host_indices[buf_off : buf_off + fill_len]
+
+            cd = cur.component_data[ct]
+            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)
+                self._attach_swa_host_value(cur, slice_)
+            else:
+                # Already has SWA (or empty overlap): drop this slice.
+                self._release_swa_host(slice_)
+
+            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])
+
+    def drive_host_eviction(
+        self, num_tokens: int, tracker: dict[ComponentType, int]
+    ) -> 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]
+        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)
+            else:
+                assert cd.host_value is not None
+                self.cache._evict_component_and_detach_lru(
+                    x, self, target=EvictLayer.HOST, tracker=tracker
+                )
+                self.cache._cascade_evict(x, self, tracker, target=EvictLayer.HOST)
+            x = x_next

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

@@ -0,0 +1,415 @@
+from __future__ import annotations
+
+import dataclasses
+from abc import ABC, abstractmethod
+from enum import Enum, IntFlag
+from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence
+
+import torch
+from numpy import float64
+
+from sglang.srt.mem_cache.base_prefix_cache import (
+    DecLockRefParams,
+    EvictParams,
+    IncLockRefResult,
+    InsertParams,
+    InsertResult,
+    MatchPrefixParams,
+    MatchResult,
+)
+from sglang.srt.mem_cache.hicache_storage import PoolTransfer, PoolTransferResult
+
+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_radix_cache import (
+        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)
+_COMPONENT_UUID_COUNTER = 1
+
+
+@dataclasses.dataclass
+class ComponentData:
+    value: Optional[torch.Tensor] = None
+    lock_ref: int = 0
+    metadata: dict[str, Any] = dataclasses.field(default_factory=dict)
+    host_value: Optional[torch.Tensor] = None
+    host_lock_ref: int = 0
+
+
+class EvictLayer(IntFlag):
+    """Which storage layer(s) to evict.  Combinable via bitwise OR."""
+
+    DEVICE = 1
+    HOST = 2
+    ALL = DEVICE | HOST
+
+
+class CacheTransferPhase(str, Enum):
+
+    BACKUP_HOST = "backup_host"  # D→H
+    LOAD_BACK = "load_back"  # H→D
+    BACKUP_STORAGE = "backup_storage"  # H→Storage
+    PREFETCH = "prefetch"  # Storage→H
+
+
+class LRURefreshPhase(str, Enum):
+
+    WALKDOWN = "walkdown"  # touching a node while walking through the tree
+    MATCH_END = "match_end"  # end of a successful prefix match
+    INSERT_END = "insert_end"  # after a new/updated leaf is committed
+
+
+def get_and_increase_time_counter() -> float64:
+    global _LAST_ACCESS_TIME_COUNTER_FLOAT
+    ret = _LAST_ACCESS_TIME_COUNTER_FLOAT
+    _LAST_ACCESS_TIME_COUNTER_FLOAT += 1.0
+    return ret
+
+
+def next_component_uuid() -> int:
+    global _COMPONENT_UUID_COUNTER
+    _COMPONENT_UUID_COUNTER += 1
+    return _COMPONENT_UUID_COUNTER
+
+
+class TreeComponent(ABC):
+    def __init__(self, cache: UnifiedRadixCache, params: CacheInitParams):
+        self.cache = cache
+
+    # Subclasses MUST set this as a class attribute (not @property)
+    component_type: ComponentType
+
+    def node_has_component_data(
+        self, node: UnifiedTreeNode, target: EvictLayer = EvictLayer.DEVICE
+    ) -> bool:
+        cd = node.component_data[self.component_type]
+        if target is EvictLayer.DEVICE:
+            return cd.value is not None
+        return cd.host_value is not None
+
+    def value_len(self, node: UnifiedTreeNode) -> int:
+        value = node.component_data[self.component_type].value
+        return len(value) if value is not None else 0
+
+    def refresh_lru(
+        self,
+        phase: LRURefreshPhase,
+        node: UnifiedTreeNode,
+        root_node: UnifiedTreeNode,
+    ) -> None:
+        ct = self.component_type
+        match phase:
+            case LRURefreshPhase.WALKDOWN:
+                if node.component_data[ct].value is None:
+                    return
+                self.cache.lru_lists[ct].reset_node_mru(node)
+            case LRURefreshPhase.MATCH_END:
+                self.cache.lru_lists[ct].reset_node_and_parents_mru(
+                    node, root_node, self.node_has_component_data
+                )
+            case LRURefreshPhase.INSERT_END:
+                # WALKDOWN already refreshed every node on the insert path
+                # (including the new leaf), so there is nothing more to do.
+                return
+            case _:
+                raise ValueError(f"Unknown LRURefreshPhase: {phase}")
+
+    @abstractmethod
+    def create_match_validator(
+        self, match_device_only: bool = False
+    ) -> Callable[[UnifiedTreeNode], bool]:
+        """Return a per-match stateful predicate that decides whether a node
+        is a valid match boundary for this component.
+        Called once per match_prefix; the returned closure may carry state.
+        When match_device_only is true, host-backed nodes must not be accepted
+        as valid match boundaries.
+        - Full: returns True if the node has full component data.
+        - SWA: tracks accumulated length since last gap; returns True only
+          when the contiguous window reaches swa_sliding_window_size.
+        - Mamba: returns True iff the node has mamba component data."""
+        ...
+
+    def finalize_match_result(
+        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."""
+        return result
+
+    def update_component_on_insert_overlap(
+        self,
+        node: UnifiedTreeNode,
+        prefix_len: int,
+        total_prefix_len: int,
+        value_slice: torch.Tensor,
+        params: InsertParams,
+    ) -> int:
+        """Called per-node when an insert's key overlaps an existing node.
+        Returns the index within value_slice from which this component
+        consumed (took ownership of) the underlying KV pool slots.
+        Returns prefix_len if nothing was consumed (default).
+        _insert_helper uses this to free only the non-consumed duplicate
+        portion: value_slice[dup_start:consumed_from]."""
+        return prefix_len
+
+    def recover_after_unevict(
+        self,
+        node: UnifiedTreeNode,
+        prefix_len: int,
+        total_prefix_len: int,
+        params: InsertParams,
+    ) -> 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
+        their own data from the freshly assigned base value when their entry
+        is still tombstoned. Default no-op."""
+        return None
+
+    def commit_insert_component_data(
+        self,
+        node: UnifiedTreeNode,
+        is_new_leaf: bool,
+        params: InsertParams,
+        result: InsertResult,
+    ) -> None:
+        """Finalize component data on the target (leaf) node after the insert
+        walk completes. Called once per insert.
+        - Full: no-op (full data is handled by _add_new_node).
+        - SWA: for new leaves, checks whether the node straddles the SWA
+          eviction boundary (swa_evicted_seqlen). If so, splits the node
+          via _split_node — the parent becomes a tombstone (no SWA) and the
+          child (the deeper portion) receives SWA data. If the entire node
+          is within the window, sets SWA directly. If entirely outside,
+          leaves SWA as None (tombstone).
+        - Mamba: sets the mamba component value from params, inserts into
+          mamba LRU list, and increments evictable size. If the node already
+          has mamba data, resets its LRU position instead."""
+        pass
+
+    @abstractmethod
+    def redistribute_on_node_split(
+        self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode
+    ):
+        """Redistribute component data between new_parent and child when a
+        node is split. new_parent is the newly created prefix node.
+        - Full: copies child's lock_ref to new_parent.
+        - SWA: slices (or clones) the swa value for new_parent, copies
+          lock_ref and component_uuid metadata, then syncs child's swa
+          value with its (now-trimmed) full_value.
+        - Mamba: sets new_parent's mamba value to None and lock_ref to 0
+          (mamba data stays on the original leaf, not on prefix nodes)."""
+        ...
+
+    @abstractmethod
+    def evict_component(
+        self,
+        node: UnifiedTreeNode,
+        target: EvictLayer = EvictLayer.DEVICE,
+    ) -> tuple[int, int]:
+        """Free this component's KV resources on a node being evicted.
+
+        *target* controls which layer(s) to evict:
+          - DEVICE: free device memory and tombstone (value = None).
+                    Host data is untouched.
+          - HOST:   free host memory (host_value = None).
+                    Device data is untouched.
+          - ALL:    free both device and host memory.
+                    No tombstone — caller will delete the node.
+
+        Returns (device_freed, host_freed) token counts."""
+        ...
+
+    def eviction_priority(self, is_leaf: bool) -> int:
+        """Eviction priority on this node type. Higher = evicted later.
+        When a component is evicted, all other components with equal or
+        lower priority on the same node are also cascade-evicted.
+
+        Leaf: all components equal (0) — evicting any cascades to all,
+        because the node will be deleted.
+
+        Internal: full=2 > swa=1 > mamba=0.
+        Why swa > mamba: SWA data on internal nodes is *path data* —
+        the sliding window needs continuous SWA coverage along the path
+        from root to the match boundary. E.g. A->B->C->D->E where C
+        and E both have mamba and the window covers C->E: if C's mamba
+        is evicted, C's SWA must stay so E remains reachable.
+        Mamba data, by contrast, is only meaningful at the match
+        boundary node; on internal nodes it
+        contributes nothing to the path. So SWA is more valuable to
+        keep and should be evicted later.
+
+        Cascade consequences:
+        - Mamba evict internal: no cascade.
+        - SWA evict internal: cascades to Mamba. SWA gone -> SWA
+          validator fails -> mamba data is useless (match requires all
+          validators to pass).
+        - Full evict internal: cascades to SWA + Mamba."""
+        return 0
+
+    @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."""
+        ...
+
+    @abstractmethod
+    def acquire_component_lock(
+        self,
+        node: UnifiedTreeNode,
+        result: IncLockRefResult,
+        lock_host: bool = False,
+    ) -> IncLockRefResult:
+        """Increment component lock refs, protecting nodes from
+        eviction. Updates evictable → protected size on first lock.
+        - Full: path-lock — walks from node up to root, incrementing
+          lock_ref on every ancestor.
+        - SWA: path-lock — walks upward collecting swa values until the
+          sliding window is filled; records a component_uuid at the
+          boundary for release_component_lock to know where to stop.
+        - Mamba: single-node lock — only increments lock_ref on the
+          node itself (mamba state is per-leaf, not per-path).
+
+        When ``lock_host`` is True, the lock applies to host-side state:
+        - Full: single-node host lock.
+        - SWA: host window-lock with a dedicated host UUID boundary.
+        - Mamba: single-node host lock with host LRU detach."""
+        ...
+
+    @abstractmethod
+    def release_component_lock(
+        self,
+        node: UnifiedTreeNode,
+        params: Optional[DecLockRefParams],
+        lock_host: bool = False,
+    ) -> None:
+        """Decrement component lock refs, un-protecting nodes.
+        Updates protected → evictable size when lock_ref drops to 0.
+        - Full: path-unlock — walks from node up to root, decrementing
+          lock_ref on every ancestor.
+        - SWA: path-unlock — walks upward, stopping at the node whose
+          component_uuid matches the one recorded during acquire.
+        - Mamba: single-node unlock — only decrements lock_ref on the
+          node itself.
+
+        When ``lock_host`` is True, the inverse host-side semantics apply."""
+        ...
+
+    def prepare_for_caching_req(
+        self,
+        req: Req,
+        insert_params: InsertParams,
+        token_ids_len: int,
+        is_finished: bool,
+    ) -> Optional[int]:
+        """Prepare component-specific data before insert, fill component
+        fields in insert_params, return effective cache_len.
+        Return None for no truncation opinion (use full length);
+        return int >= 0 for effective cache length.
+        - Full: no-op, returns None.
+        - SWA: sets insert_params.swa_evicted_seqlen on finished; returns None.
+        - Mamba: prepares mamba_value (finished from ping-pong buffer,
+          unfinished fork from req); returns mamba_last_track_seqlen."""
+        return None
+
+    def cleanup_after_caching_req(
+        self,
+        req: Req,
+        is_finished: bool,
+        insert_result: Optional[InsertResult] = None,
+        insert_params: Optional[InsertParams] = None,
+    ) -> None:
+        """Post-cache cleanup for component-specific resources.
+
+        ``is_finished`` — whether the request has finished generation.
+        True means the request is complete and its resources can be released;
+        ``insert_result`` is None when insert was skipped (cache disabled
+        or effective_cache_len <= 0); treat as "no insert happened".
+        ``insert_params`` is None only on the disabled path; on early-return
+        paths it is still provided so components can free their resources."""
+        pass
+
+    def free_out_of_window_slots(
+        self, req: Req, pre_len: int, insert_params: InsertParams
+    ) -> None:
+        pass
+
+    # ---- HiCache Hooks ----
+
+    def build_hicache_transfers(
+        self,
+        node: UnifiedTreeNode,
+        phase: CacheTransferPhase,
+        *,
+        req: Optional[Req] = None,
+        token_ids: Optional[Sequence[int]] = None,
+        prefetch_tokens: int = 0,
+        last_hash: Optional[str] = None,
+    ) -> Optional[list[PoolTransfer]]:
+        """Build transfer descriptors for this component in the given phase.
+        Returns None if the component has nothing to transfer."""
+        return None
+
+    def commit_hicache_transfer(
+        self,
+        node: UnifiedTreeNode,
+        phase: CacheTransferPhase,
+        transfers: list[PoolTransfer] = (),
+        *,
+        insert_result: Optional[InsertResult] = None,
+        pool_storage_result: Optional[PoolTransferResult] = None,
+    ) -> None:
+        """Post-transfer bookkeeping: store host indices, update LRU, etc."""
+        pass
+
+    def drive_host_eviction(
+        self, num_tokens: int, tracker: dict[ComponentType, int]
+    ) -> None:
+        """Evict from this component's host-side resources.
+        Called by HostPoolGroup when the host pool is full.
+        Default no-op for components without host storage."""
+        pass