源码精读：RadixCache 实现细节

mem_cache/
├── base_prefix_cache.py        ← 抽象基类
├── radix_cache.py              ← 我们的主角（826 行）
├── hiradix_cache.py            ← HiCache 扩展（CPU 镜像）
├── chunk_cache.py              ← chunk-only 模式（关闭 radix）
├── memory_pool.py              ← KV 物理存储
├── allocator.py                ← slot 分配
├── evict_policy.py             ← 淘汰策略接口
├── cache_init_params.py        ← 初始化参数
└── cpp_radix_tree/             ← C++ 加速版（可选）

@dataclass
class RadixKey:
    token_ids: List[int]
    extra_key: Optional[str] = None       # LoRA id / 命名空间 / cache version

    def page_aligned(self, page_size):
        ...

    def maybe_to_bigram_view(self, is_eagle):
        # EAGLE 投机解码用：把 token 序列看作 bigram
        ...

class TreeNode:
    counter = 0

    def __init__(self, id=None, priority=0):
        self.children = defaultdict(TreeNode)      # 按首 token id 索引
        self.parent: TreeNode = None
        self.key: RadixKey = None                  # 这条边上的 token 序列
        self.value: Optional[torch.Tensor] = None  # device KV slot 索引
        self.lock_ref = 0                          # device 引用计数
        self.last_access_time = time.monotonic()
        self.creation_time = time.monotonic()

        self.hit_count = 0                         # 用于 LFU
        self.host_ref_counter = 0                  # CPU 镜像引用计数
        self.host_value: Optional[torch.Tensor] = None  # CPU 镜像
        self.hash_value: Optional[List[str]] = None     # HiCache 用，每 page 一个 hash

        self.priority = priority                   # priority-aware eviction
        self.id = TreeNode.counter if id is None else id
        TreeNode.counter += 1

class RadixCache(KVCacheEventMixin, BasePrefixCache):
    def __init__(self, params: CacheInitParams):
        self.token_to_kv_pool_allocator = params.allocator
        self.page_size = params.page_size
        self.disable = params.disable
        self.is_eagle = params.is_eagle
        self.root_node = TreeNode()                  # 根
        self._empty_match_result = MatchResult(...)
        # ... 锁 / events / mixin 状态

def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
    key = params.key
    key, _ = key.maybe_to_bigram_view(self.is_eagle)
    if self.disable or len(key) == 0:
        return self._empty_match_result
    key = key.page_aligned(self.page_size)
    if len(key) == 0:
        return self._empty_match_result

    value, last_node = self._match_prefix_helper(self.root_node, key)
    if value:
        value = torch.cat(value)
    else:
        value = self._empty_match_result.device_indices
    return MatchResult(
        device_indices=value,
        last_device_node=last_node,
        last_host_node=last_node,
        best_match_node=last_node,
    )

def _match_prefix_helper(self, node, key):
    """从 node 出发，匹配 key 的最长前缀。"""
    if len(key) == 0:
        return [], node

    first_token = key.first_token()
    if first_token not in node.children:
        return [], node                            # 当前节点子树里没匹配

    child = node.children[first_token]
    common_prefix_len = longest_common_prefix(child.key, key)

    if common_prefix_len == len(child.key):
        # 完全消耗 child 的边，递归下钻
        remaining_key = key.slice_from(common_prefix_len)
        child.last_access_time = time.monotonic()
        child.hit_count += 1
        sub_values, last_node = self._match_prefix_helper(child, remaining_key)
        return [child.value] + sub_values, last_node
    else:
        # 部分匹配 → 分裂 child
        new_node = self._split_node(child, common_prefix_len)
        return [new_node.value], new_node

def _split_node(self, child, split_pos):
    """把 child 在 split_pos 处一分为二。
    Returns 上半部分（也是新的中间节点）。"""
    head_key = child.key.slice_to(split_pos)
    tail_key = child.key.slice_from(split_pos)

    head_node = TreeNode()
    head_node.key = head_key
    head_node.value = child.value[:split_pos]      # KV 切片（只是改 tensor view，不复制）
    head_node.parent = child.parent

    # 改 child
    child.key = tail_key
    child.value = child.value[split_pos:]
    child.parent = head_node
    head_node.children[tail_key.first_token()] = child

    # 父节点的指针改向新 head
    child.parent_old = child.parent  # 暂存
    head_node.parent.children[head_key.first_token()] = head_node

    return head_node

def insert(self, params: InsertParams) -> InsertResult:
    """把 (key, value=device_indices) 挂到树上。"""
    # 1. 先尝试 match，命中部分跳过
    matched_node, depth = self._find_attach_point(params.key)

    # 2. 剩余部分建新节点
    remaining_key = params.key.slice_from(depth)
    remaining_value = params.value[depth:]
    new_node = TreeNode()
    new_node.key = remaining_key
    new_node.value = remaining_value
    new_node.parent = matched_node
    matched_node.children[remaining_key.first_token()] = new_node

    # 3. 引用计数（让插入的节点至少有一个 ref，防止立刻被 evict）
    new_node.lock_ref += 1

    return InsertResult(node=new_node, ...)

def evict(self, params: EvictParams) -> EvictResult:
    """释放 params.num_tokens 个 token 的 KV slot."""
    leaves = self._collect_evictable_leaves()       # lock_ref == 0 的叶子
    leaves.sort(key=lambda n: (n.priority, n.last_access_time))  # LRU + priority

    freed = 0
    freed_slots = []
    while freed < params.num_tokens and leaves:
        node = leaves.pop(0)
        freed_slots.append(node.value)
        freed += node.value.shape[0]
        self._detach(node)                           # 从父节点移除
        # 父节点是否变成新叶子？是的话加进 queue
        parent = node.parent
        if parent is not self.root_node and len(parent.children) == 0 and parent.lock_ref == 0:
            insort(leaves, parent, key=lambda n: (n.priority, n.last_access_time))

    self.token_to_kv_pool_allocator.free(torch.cat(freed_slots))
    return EvictResult(num_freed_tokens=freed)