sgl-project--sglang/python/sglang/srt/mem_cache/unified_cache_components/README.md

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:

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.