arXiv:2605.22106cs.AI2026-05被引 3

提出树结构感知的缓存管理方法,显著降低大模型推理内存占用

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

  • 基于搜索动态设计树结构感知的缓存淘汰策略
  • 在保持近全精度的前提下实现最高4倍的缓存内存压缩
  • 适合需要深度广度搜索的大模型推理场景

大模型推理正从单次生成转向显式搜索中间推理状态。树思维(ToT)通过分支与回溯构建树状搜索,但大幅增加键值(KV)缓存压力:保留多条路径的中间状态迅速成为内存瓶颈,限制吞吐量并约束搜索深度与宽度。我们观察到,ToT推理中KV重用受搜索动态影响:短期解码主要依赖当前分支及其祖先,而未激活子树虽需保留以支持回溯,但短期内重用概率极低。据此提出ArborKV,一种结合轻量估值器与树感知分配策略的缓存管理框架,采用纯令牌提取式淘汰与懒加载重激活机制,支持高效回访。在ToT推理基准测试中,ArborKV实现高达约4倍的峰值KV内存减少,同时保持接近全保留的准确率,使在固定硬件预算下可支持更大规模的搜索配置。

原文摘要 · Abstract (English)

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.

大模型推理缓存优化树搜索

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。