SpecLA: Efficient Speculative Decoding for Linear-Attention Models
Computation and Language
Summary
The authors introduce SpecLA, a method to speed up a type of language model called stateful linear-attention models during text generation. Normally, these models process one word at a time, which can be slow. SpecLA uses a special way to check multiple possible next words at once and carefully manages the model's internal memory to avoid extra work. Their tests showed that SpecLA can make the model run up to 1.7 times faster compared to the usual method. This work focuses on handling the unique challenges of memory and verification in linear-attention models.
Authors
Zhibin Wang, Xuying Han, Zhaohua Yang, Fuliang Liu, Xue Li, Rong Gu, Sheng Zhong, Chen Tian
Abstract
Linear-attention models replace the growing KV cache with recurrent states, but autoregressive decoding still reads, updates, and writes these states one token at a time. Speculative decoding can reduce this cost by verifying several draft tokens in one target pass, yet existing speculative systems are designed for Transformer KV caches. For stateful linear-attention targets, verification must follow recurrent dependencies across chains and branches, acceptance must update only the accepted state trajectory, and the drafter must avoid submitting candidates that waste stateful verification work. This paper presents SpecLA, a speculative decoding runtime for stateful linear-attention models. SpecLA verifies chains and trees with topology-aware kernels, stores compact factors produced during verification to recover accepted states, and uses confidence pruning plus a target-aligned EAGLE-style drafter to feed useful candidates to the verifier. On an NVIDIA H100 with a public GDN-1.3B target, SpecLA achieves up to 1.70x end-to-end speedup over autoregressive decoding.