[RFC]: Radix Tree Index — Prefix-Aware KV Cache Management and Cascading Eviction

### Changes proposed

## Summary

Introduce a Radix Tree auxiliary index for the Mooncake Store Master Service to enable:
1. **Prefix queries**: O(1) lookup of all objects sharing a blake3 hash prefix
2. **Parent-child-aware eviction**: Evict leaf nodes first; when space is still insufficient, cascade upward to evict parent nodes — preserving RadixAttention inference path integrity
3. **Zero invasion of primary storage**: Implemented as an auxiliary index layer on top of the existing `unordered_map`, with no impact on hot paths

## 1. Background and Motivation

Mooncake Store stores KV Cache objects in 1024-sharded `unordered_map`. Keys are opaque strings — the Master never parses their internal structure. This leads to three core problems:

**Problem 1: No prefix queries**

SGLang's RadixAttention needs to check whether all keys under a given hash prefix exist (batch_is_exist). Currently this requires scanning all 1024 shards — an O(N) operation that becomes prohibitively slow at million-object scale.

**Problem 2: Eviction is parent-child-unaware**

In RadixAttention, a parent node's KV Cache is the prefix of its children. If a parent is evicted while its children remain, the children lose their context (inference requires a complete root-to-leaf path). Current `BatchEvict` sorts purely by lease timeout, completely ignoring parent-child dependencies.

**Problem 3: Diverse key formats but uniform prefixes**

Key formats vary across model configurations (MLA: ``, MHA: `__`, MHA+TP: `___tp_`), but the first 64 characters are always a blake3 hex digest identifying a RadixAttention node.

**Goals**: Prefix queries O(N) → O(1), eviction guarantees "children alive implies parent alive", zero invasion of primary storage, progressive enablement via `enable_radix_tree` flag.

## 2. Core Design Decisions

### Design 1: Auxiliary Index, Not a Replacement

The Radix Tree is an auxiliary index layer on top of primary storage. Primary storage (MetadataShard[1024]) remains unchanged; a new RadixTreeShard[256] stores only relationship metadata (parent_prefix_hash, children, registered_keys) — not ObjectMetadata itself.

Rationale: ObjectMetadata is non-copyable/non-movable; hash table O(1) lookup is a critical hot path; the Radix Tree only stores string references with minimal overhead.

```
struct RadixTreeNode {
// The prefix hash (blake3 hex digest), e.g.
// "13b825898e41332c2beeae39b161d50994fd2ea4af86f4b95fb5dd0e66893882"
std::string prefix_hash;

// Parent node's prefix hash, or "" for root nodes
std::string parent_prefix_hash;

// Child nodes' prefix hashes
std::unordered_set children;

// Full object keys registered under this prefix hash.
// MLA models: {""}
// MHA models: {"_0_k", "_0_v", "_1_k", ...}
// MHA+TP: {"_0_k_tp_0", "_0_k_tp_1", ...}
std::unordered_set registered_keys;

bool IsLeaf() const { return children.empty(); }
};
```

```
┌──────────────────────────────────────────────────────┐
│ Master Service │
│ │
│ Primary Store (unchanged): │
│ MetadataShard[1024] │
│ key → ObjectMetadata (replicas, lease, ...) │
│ │
│ Relationship Index (NEW): │
│ RadixTreeShard[256] │
│ prefix_hash → RadixTreeNode │
│ ├── parent_prefix_hash │
│ ├── children: {child1_hash, child2_hash} │
│ └── registered_keys: {full_key1, full_key2} │
└──────────────────────────────────────────────────────┘
```

### Design 2: Sharding by prefix_hash

Consistent with MetadataShard's 1024-shard pattern, but with fewer shards. Rationale: Radix Tree node count is far smaller than object count (MLA ~1:1, MHA ~1:2×num_layers), so 256 shards provide sufficient granularity while reducing the number of shards to scan during eviction.

### Design 3: String References Instead of Pointers

`RadixTreeNode` uses `string` for parent/children references rather than `shared_ptr/weak_ptr`. Rationale:

| | Pointer approach | String approach (current) |
|---|---|---|
| Cross-shard safety | `weak_ptr::lock()` produces `shared_ptr` without guaranteeing target shard lock is held | `nodes.find(hash)` operates under shard lock protection |
| Lock ordering | Cannot predict which shards need locking during traversal | Can pre-compute shard indices and lock in ascending order |
| CascadeCleanup | Acquiring new shard locks per level cannot guarantee ascending order | Per-level `find` + shard lock acquisition, implemented and safe |
| Memory overhead | 16-byte pointers | 64-byte strings, but node count is far smaller than object count |

### Design 4: Prefix Extraction via blake3 Fixed Length

blake3 hex digest is always 64 characters. If the first `_` in a key falls at position 64, the first 64 characters are the prefix; otherwise the entire key serves as the prefix (MLA or non-standard formats). O(1) time.

### Design 5: Two Registration Paths

- **Path A**: Inline via `ReplicateConfig.parent_block_hash` at `PutStart` time, registered through the Deferred Pattern
- **Path B**: Standalone `RegisterRadixTreeNode` / `BatchRegisterRadixTreeNode` RPCs, for scenarios like Indexer batch registration

### Design 6: Shell Node Mechanism

When a child node is registered before the parent's `PutStart` arrives, a shell node is created — with `children` populated but `registered_keys` empty. When the parent's `PutStart` arrives later, `registered_keys` is filled in. This ensures correct relationship establishment regardless of registration order.

## 3. Multi-Round Leaf-to-Root Cascading Eviction

### Algorithm Overview

Multi-round iteration, each round has three phases:

1. **Phase 1** (radix tree shared lock): Scan all shards, collect current leaf node candidates
2. **Phase 2** (metadata shared lock): Check whether all keys in each leaf are evictable (lease expired + has evictable replica), sort by `lease_timeout`
3. **Phase 3** (metadata write lock + deferred radix tree cleanup): Execute eviction, update radix tree via Defer/Flush pattern

After eviction, `CascadeCleanup` automatically converts empty-shell parent nodes into new leaves. The next round scans these new leaves and evicts them. The loop continues until the watermark is restored or no more leaves exist.

```
RadixTreeAwareEvict():
for round = 0; round < kMaxCascadeRounds(16); round++:
if watermark recovered: break

Phase 1: Collect current leaf candidates (radix tree shared lock)
Phase 2: Check evictability, sort by lease_timeout (metadata shared lock)
Phase 3: Execute eviction + DeferRadixTreeUnregister (metadata write lock)
FlushRadixTreeUnregister → CascadeCleanup
// Parents that lost all children become new leaves, evictable next round

if watermark recovered: break
```

### Key Design Decisions

| Decision | Rationale |
|----------|-----------|
| Watermark as sole exit condition | Exit when `mem_used_ratio <= high_watermark` |
| Skip parents that still have children | Guarantees "children alive → parent alive" — parents only become new leaves after all children are evicted |
| Maximum cascade depth of 16 rounds | Exceeds any realistic RadixAttention tree depth, while preventing infinite loops |
| Re-scan radix tree each round | Previous round's eviction may have changed tree structure |
| Eviction atomicity = leaf node | All keys under a leaf are either all evictable or none are |
| Fallback to traditional eviction | If cascading eviction is insufficient, BatchEvict continues with 3-pass |

### Watermark vs. Object Count as Exit Condition

Traditional `BatchEvict` uses `evicted_count / object_count = evict_ratio_target` as its target, but this is only a proxy metric:
1. Object size varies greatly (MLA: 1 object/prefix vs MHA: 2×num_layers), so evicting the same count can free 10× different bytes
2. Count-based targets lead to over-eviction or under-eviction
3. Checking count targets inside Phase 3 per-leaf leads to partially-evicted subtrees with inconsistent state

Memory watermark is the precise metric — checked twice per round (before and after the loop body), with negligible overhead.

### Example

```
Initial tree: A → {B → {D(leaf), E(leaf)}, C → {F(leased)}}

Round 1: Evict D, E → B loses all children → B becomes a new leaf
Round 2: Evict B → A still has C → A is retained
Round 3: C still has F (non-evictable) → skip
```

## 4. Deferred Radix Tree Cleanup Pattern

### The Problem

`UnregisterRadixTreeNodeInternal` / `RegisterRadixTreeNodeInternal` require acquiring radix tree shard locks. Many metadata-erasing functions (PutRevoke, Remove, BatchEvict, etc.) already hold metadata shard locks at the point of erasure. Directly calling radix tree functions results in **metadata → radix_tree** — a lock order violation that causes deadlocks.

### Solution

While holding the metadata lock, only collect keys into a vector (`DeferRadixTreeUnregister` — just push_back, no lock needed). After releasing the metadata lock, perform the actual radix tree operations (`FlushRadixTreeUnregister`). Registration follows the same pattern.

### Four Variants

| Pattern | Applicable functions | Approach |
|---------|---------------------|----------|
| A: Scoped block + Defer/Flush | PutStart, UpsertStart, ClearInvalidHandles, BatchEvict, etc. | Metadata operations in scoped block, Flush outside |
| B: Result variable + Defer/Flush | CopyEnd, CopyRevoke, MoveEnd, MoveRevoke, etc. | Use result variable to avoid early return skipping Flush |
| C: Direct call after scope | Remove | Direct Unregister call after accessor destruction releases metadata lock |
| D: Multi-shard Defer + Flush | RemoveByRegex, RemoveAll, BatchRemove | Collect across all shards, Flush at the end |

See the detailed design document for the complete list of 18 affected functions.

## 5. Lock Order

```
1. snapshot_mutex_
2. client_mutex_
3. radix_tree_shards_[min] (ascending order for multi-shard acquisition)
4. radix_tree_shards_[max] (when parent and child are in different shards)
5. metadata_shards_[N]
6. segment_mutex_
7. ObjectMetadata.lock (SpinLock, per-object)
```

**Prohibited**: metadata → radix_tree (resolved by Deferred Pattern), radix_tree[max] → radix_tree[min] (resolved by `ScopedRadixTreeDualLock` ascending lock), radix_tree → segment_mutex_.

## 6. Thread Safety

| Scenario | Safety |
|----------|--------|
| Concurrent PutStart (parent-child in different shards) | Safe: both threads acquire locks in ascending order |
| PutStart concurrent with eviction | Safe: eviction is two-phase; keys in processing state are skipped |
| Cascade deletion concurrent with new child registration | Safe: shard write lock serializes both operations |
| Multi-round cascading eviction concurrent with registration | Safe: re-scan each round; new children prevent parents from becoming leaves |
| Cascading eviction concurrent with traditional eviction | Safe: BatchEvict runs after RadixTreeAwareEvict returns |

## 7. Interface Changes Overview

- **Configuration**: All config classes add `enable_radix_tree = false` flag
- **RPC**: New `RegisterRadixTreeNode`, `BatchRegisterRadixTreeNode`, `GetKeysByPrefix`
- **PutStart inline**: `ReplicateConfig` adds `parent_block_hash` field
- **C API**: `mooncake_replicate_config` adds `parent_block_hash`
- **Client**: `MasterClient` adds corresponding three methods

## 8. Future Work

| Item | Description |
|------|-------------|
| HA persistence | MetadataSerializer needs to serialize/deserialize radix tree state; currently lost on restart, degrading to traditional eviction (safe degradation, no data inconsistency), progressively rebuilt as inference requests arrive |
| Batch subtree traversal API | Current `GetKeysByPrefix` returns single-node data, requiring N RPCs for depth N; future `GetSubtreeKeys(prefix_hash, max_depth)` will support server-side one-shot traversal |
| Responsibility boundary: token→hash mapping | Radix Tree provides block_hash→children relationships, not token_ids→block_hash mapping (maintained by inference engine's HiRadixTree) |
| Python/Go/Rust bindings | Need to expose parent_block_hash parameter |
| Unit tests | Radix tree CRUD, cross-shard parent-child, cascade deletion, eviction correctness |
| Traditional BatchEvict SpinLock | Directly reads lease_timeout without acquiring per-object SpinLock; safe under shard write lock but violates annotation convention |
| Code style unification | Some early return paths skip Flush; Remove() uses direct call instead of deferred pattern — currently correct but inconsistent |

ref #1732

### Before submitting a new issue...

- [x] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)