ArborKV: Structure-Aware KV Cache Management for Scaling Tree-based LLM Reasoning

TL;DR

ArborKV is a novel cache management method for tree-based LLM reasoning that significantly reduces memory usage while maintaining accuracy, enabling larger searches within fixed hardware constraints.

Contribution

It introduces a structure-aware eviction framework with a lightweight estimator and lazy rehydration, improving memory efficiency for ToT-style reasoning.

Findings

Achieves up to ~4x peak KV-memory reduction.

Maintains near-full-retention accuracy.

Enables larger search configurations under fixed memory budgets.

Abstract

Recent progress in LLM reasoning has increasingly shifted from single-pass generation to explicit search over intermediate reasoning states. Tree-of-Thoughts (ToT) organizes inference to tree-structured search with branching and backtracking, but it substantially amplifies the Key--Value (KV) cache: retaining KV states for a frontier of partial trajectories quickly becomes a memory bottleneck that limits throughput and constrains search depth and width under fixed hardware budgets. We address this challenge by observing that KV reuse in ToT-style inference is governed by search dynamics: near-term decoding depends primarily on the active branch and its ancestors, whereas inactive subtrees have low short-term reuse probability yet must remain recoverable for backtracking. Motivated by this, we propose ArborKV, a structure-aware eviction framework that couples a lightweight value estimator with a tree-aware allocation policy, and performs purely token-extractive eviction with lazy rehydration to support revisits. Experiments on ToT-style reasoning benchmarks show that ArborKV achieves up to ~4x peak KV-memory reduction while preserving near-full-retention accuracy, enabling larger search configurations under fixed device budgets that would otherwise run out of memory.