Bole: Efficient Tree Speculation for Hybrid-Attention Language Models

📄 arXiv: 2608.01651v1 📥 PDF

作者: Li Wang, Yi Su, Xiabao Wu, Chiran You, Yongchao Liu, Zhan Qiu, Juelu Zhang, Jiajun Zheng, Fangxin Liu, Jie Zhang, Chen Tian, Chengying Huan

分类: cs.DC, cs.CL, cs.LG

发布日期: 2026-08-03

备注: 14 pages, 12 figures, 7 tables

💡 一句话要点

提出Bole以解决混合注意力语言模型的树状推测效率问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合注意力 树状推测 语言模型 GPU加速 内存优化 深度学习 推理效率

📋 核心要点

  1. 现有的树状推测系统在混合注意力模型上性能不佳，导致验证延迟和内存使用效率低下。
  2. Bole通过将线性注意力递归转化为树状结构，并实现并行验证，显著提升了推测效率。
  3. 在多个模型和数据集上，Bole的解码吞吐量提高了4.72倍，在线工作负载下TTFT和TPOT分别降低了67.6%和49.9%。

📝 摘要（中文）

混合注意力的大型语言模型结合了全注意力与递归线性注意力，以降低长上下文推理成本，但其自回归解码仍受限于内存。树状推测解码提供了加速的可能性，但现有系统在混合模型上表现不佳。本文提出Bole，一个内核-运行时协同设计，能够高效实现混合注意力LLM的树状推测。Bole将线性注意力递归转化为树状结构，并通过资源高效的GPU内核并行验证所有提议节点，显著提升了线性注意力树验证的速度，并减少了瞬态状态内存。

🔬 方法详解

问题定义：本文旨在解决混合注意力语言模型在自回归解码中存在的内存瓶颈和效率低下的问题。现有的树状推测系统在处理混合模型时，因逐层遍历递归层而导致验证延迟和瞬态内存使用不当。

核心思路：Bole的核心思路是将线性注意力的递归结构转化为树状闭合形式，并通过高效的GPU内核实现并行验证，从而加速推测过程，降低内存占用。

技术框架：Bole的整体架构包括将线性注意力递归转化为树状结构的模块，以及一个资源高效的GPU内核，用于并行验证提议节点。该框架还集成了状态管理与批量验证预算的调节，确保高效的推测过程。

关键创新：Bole的主要创新在于其将线性注意力的递归转化为树状结构的能力，以及通过并行验证显著提高了线性注意力树验证的速度，减少了瞬态状态内存的使用。

关键设计：Bole在设计中采用了损失函数和网络结构的优化，以确保在推测过程中状态更新的无损编码，并仅重建采样后选择的状态，从而实现了高效的内存管理。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Bole在多个模型和GPU平台上实现了自回归解码吞吐量的最高提升达4.72倍，相较于最强的树状推测基线，在线工作负载下TTFT和TPOT分别降低了67.6%和49.9%。这些结果表明Bole在推测效率和内存管理方面的显著优势。

🎯 应用场景

Bole的研究成果在大型语言模型的推理加速中具有广泛的应用潜力，尤其是在需要处理长上下文的自然语言处理任务中。其高效的状态管理和内存使用优化可以为实时应用提供更快的响应时间，提升用户体验。未来，Bole的技术可以扩展到其他需要高效推测的深度学习模型中。

📄 摘要（原文）

Hybrid-attention large language models combine full attention with recurrent linear attention to reduce long-context inference costs, yet their autoregressive decoding remains memory-bound. Tree speculative decoding offers an attractive acceleration path, but existing tree-speculation systems are designed around the key--value caches of full-attention models. On hybrid models, they traverse recurrent layers branch by branch and materialize a full state for every proposal node, causing verification latency and transient memory to scale poorly with tree and batch sizes. We present Bole, a kernel--runtime co-design that enables efficient tree speculation for hybrid-attention LLMs. Bole transforms the linear-attention recurrence into a tree-structured closed form and realizes it with a resource-efficient GPU kernel, verifying all proposal nodes in parallel and accelerating linear-attention tree verification by 3.4--7.7$\times$. It losslessly encodes speculative state updates as token-level factors and reconstructs only the state selected after sampling, reducing transient state memory by 82--99$\times$ and freeing GPU capacity for KV caches. Its integration into SGLang, a widely deployed production LLM serving engine, couples efficient state management with a batch-wide verification budget calibrated to the complete hybrid forward. Across four models, two GPU platforms, and diverse datasets, Bole delivers up to $4.72\times$ the offline decode throughput of autoregressive decoding and up to $2.03\times$ that of the strongest tree-speculative baseline. Under online agent workloads, it reduces TTFT and TPOT by up to $67.6%$ and $49.9%$, respectively, over the strongest tree-speculative baseline.