Vegas: Self-Speculative Decoding with Verification-Guided Sparse Attention

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 against the full KV cache, speeds up inference in a lossless manner. However, they rely on a standalone KV selection algorithm to select the KV entries used for drafting and overlook the fact that the criticality of each KV entry is inherently computed during verification. In this paper, we propose Vegas, a self-speculative decoding method with verification-guided sparse attention. Vegas identifies critical KV cache entries as a byproduct of verification and computes attention only over these entries when drafting subsequent tokens. This not only improves the draft token acceptance rate but also incurs low KV selection overhead, thereby improving decoding throughput. Vegas achieves a 1.25$\times$–2.81$\times$ speedup in decoding throughput over default vLLM and a 1.15$\times$–1.29$\times$ speedup over state-of-the-art sparse attention-based self-speculative decoding methods. Our code is available at https://github.com/platformxlab/vegas.

How fast a large language model (LLM) generates tokens is an important metric for justifying LLM inference efficiency. To generate each new token, the LLM must consider all previous tokens. Because of this, the token generation process has been a well-known bottleneck, especially for long-context scenarios where there would be millions of previous tokens. To speed up the token generation, an LLM can draft a few tokens by only considering the most important previous tokens. After that, it checks these draft tokens in a single pass and drops inaccurate ones. By doing so repeatedly, the LLM can generate tokens at a faster speed by spending less time considering previous tokens. However, to ensure that draft tokens will not be dropped, it is important to find the truly important previous tokens. In our work Vegas, we achieve this by looking at which previous tokens turned out to be important when an LLM is checking the draft tokens. With this insight, Vegas can find important previous tokens with higher accuracy. Therefore, its draft tokens are more likely to pass the check, enabling it to generate tokens faster. Our experimental results show that it can achieve a 2.8 times speedup over standard methods.