arXiv:2608.01651cs.DCcs.CL2026-08被引 1

让混合注意力大模型的树形推测加速3倍以上,显著降低内存占用。

Bole: Efficient Tree Speculation for Hybrid-Attention Language Models

  • 将线性注意力递归转为树结构闭式,用高效GPU核并行验证所有候选节点
  • 验证延迟降低3.4–7.7倍,临时状态内存减少82–99倍,释放更多显存用于缓存
  • 适合需要高吞吐、低延迟的在线推理场景,尤其在代理任务中表现突出

混合注意力大语言模型结合全量注意力与循环线性注意力以降低长上下文推理成本,但其自回归解码仍受内存限制。树形推测解码虽具潜力,但现有系统针对全量注意力设计,需逐分支遍历循环层并为每个候选节点生成完整状态,导致验证延迟和瞬态内存随树与批量规模增长而急剧上升。本文提出Bole,一种内核-运行时协同设计,实现混合注意力模型的高效树形推测。Bole将线性注意力递归转化为树状闭式表达,并通过资源高效的GPU内核实现,可并行验证所有候选节点,使线性注意力树验证速度提升3.4–7.7倍。它以令牌级因子无损编码推测状态更新,仅重构采样选定状态,使瞬态状态内存减少82–99倍,释放大量显存供KV缓存使用。集成至SGLang(广泛应用的生产级大模型服务引擎)后,结合全局批量验证预算,覆盖完整混合前向计算。在四个模型、两个GPU平台及多样数据集上,Bole相比自回归解码最高提升4.72倍离线吞吐,相比最强基线提升2.03倍;在在线代理负载下,相比最强基线,首次生成时间(TTFT)和每输出令牌时间(TPOT)分别降低67.6%和49.9%。

原文摘要 · Abstract (English)

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.

推理加速注意力机制内存优化树形推测

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