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

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.

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.

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.

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.

Algorithm detail: Calls acquire_component_lock() for each component.

dec_lock_ref(node, params?) → DecLockRefResult

Unlock a previously locked node path.

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.

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).

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

Insert Phase

Node Split

Eviction Phase

Lock Phase

Caching Phase

Utility

Component Behavior Summary

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.