🤖 AI Summary
This work addresses the inefficiency of autoregressive decoding in linear attention models, which is constrained by recursive state dependencies that necessitate token-by-token processing and hinder existing speculative decoding approaches. To overcome this limitation, the authors propose SpecLA, the first speculative decoding runtime system tailored for stateful linear attention. SpecLA introduces a novel topology-aware verification kernel to efficiently handle token chains and tree structures, complemented by compact state factorization, an EAGLE-style drafter aligned with target model dynamics, confidence-based pruning, and recursive state management that updates only accepted trajectories while minimizing redundant computation. Evaluated on the GDN-1.3B model using NVIDIA H100 GPUs, SpecLA achieves up to a 1.70× end-to-end inference speedup.
📝 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.