SpecAttn: Co-Designing Sparse Attention with Self-Speculative Decoding
This addresses the memory demands of KV cache in LLM inference, offering a practical speedup for AI applications, though it is incremental as it builds on existing self-speculative decoding with sparse attention.
The paper tackles the memory bottleneck in long-context LLM inference by proposing SpecAttn, a self-speculative decoding method that uses verification-guided sparse attention to improve draft token acceptance and reduce overhead, achieving 2.81× higher throughput than vanilla decoding and 1.29× improvement over state-of-the-art methods.
Long-context large language model (LLM) inference has become the norm for today's AI applications. However, it is severely bottlenecked by the increasing memory demands of its KV cache. Previous works have shown that self-speculative decoding with sparse attention, where tokens are drafted using a subset of the KV cache and verified in parallel with full KV cache, speeds up inference in a lossless way. However, this approach relies on standalone KV selection algorithms to select the KV entries used for drafting and overlooks that the criticality of each KV entry is inherently computed during verification. In this paper, we propose SpecAttn, a self-speculative decoding method with verification-guided sparse attention. SpecAttn identifies critical KV entries as a byproduct of verification and only loads these entries when drafting subsequent tokens. This not only improves draft token acceptance rate but also incurs low KV selection overhead, thereby improving decoding throughput. SpecAttn achieves 2.81$\times$ higher throughput over vanilla auto-regressive decoding and 1.29$\times$ improvement over state-of-the-art sparsity-based self-speculative decoding methods.