Bole: Efficient Tree Speculation for Hybrid-Attention Language Models
This work significantly improves the inference efficiency and throughput of hybrid-attention large language models, which are used for long-context generation, benefiting users and developers by reducing computational costs and latency. It is an incremental improvement on existing tree-speculation systems.
This paper introduces Bole, a kernel-runtime co-design for efficient tree speculation in hybrid-attention language models, which are typically memory-bound during autoregressive decoding. Bole accelerates linear-attention tree verification by 3.4-7.7x and reduces transient state memory by 82-99x, leading to up to 4.72x higher offline decode throughput and up to 2.03x higher throughput compared to the strongest tree-speculative baseline. For online agent workloads, it reduces TTFT and TPOT by up to 67.6% and 49.9% respectively.
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.