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

AI summary

Authors: Yikang Yue, Yuqi Xue, Jian Huang (University of Illinois Urbana-Champaign) arXiv: 2602.07223v2 [cs.LG], License CC BY 4.0

Overview

Research area: Efficient LLM inference systems — specifically long-context serving, KV cache optimization, sparse attention, and speculative decoding.

Technical level: Intermediate. The paper assumes working familiarity with the Transformer attention mechanism, KV caching, and the draft/verify loop of speculative decoding, but explains its own design choices from first principles.

Scope: A training-free decoding method that reuses the attention logits already computed during speculative verification to select which KV cache entries the next drafting round should attend to.

What This Paper Is About

Long-context LLM inference is bottlenecked by memory bandwidth: the KV cache grows linearly with context length, and reading it dominates each decoding step. Sparse attention reduces those reads but risks degrading output quality, while self-speculative decoding with sparse attention preserves exact output quality but is limited by a trade-off between how many draft tokens get accepted (drafting accuracy) and how much extra work is spent picking which KV entries to keep (KV selection overhead). Vegas's goal is to break that trade-off by treating the verification pass — which already computes full attention over every KV entry — as a free source of information about which entries actually matter.

Key Contributions

1. Vegas, a verification-guided self-speculative decoding mechanism. Rather than using verification only to accept or reject draft tokens, Vegas identifies critical KV cache entries as a byproduct of the full-attention computation performed during verification, and computes attention only over those entries when drafting subsequent tokens.

2. A low-overhead KV selection algorithm. Selecting KV entries by aggregating attention logits across draft tokens (including discarded ones) avoids the rapid decay in acceptance probability seen when selecting from the last accepted token alone. The paper formalizes this as selecting the k prefix tokens that maximize cumulative attention logits over all draft tokens.

3. The "Collect-2-Query" optimization. Instead of collecting attention logits for every query token, Vegas extracts them only from the first draft token and the bonus token — the pair with maximum positional distance, whose selected KV sets are least redundant with each other. This bounds collection overhead as the number of draft tokens grows.

4. A vLLM implementation and systematic evaluation. Vegas is implemented inside vLLM using a custom proposer, a repurposed PagedAttention kernel with page size 1 token for token-granular sparse attention, and an instrumented FlashAttention-3 kernel that writes BF16 attention logits to HBM after masking and before softmax. A three-step hyperparameter tuning strategy balances the sparsity ratio against the number of draft tokens γ.

Main Findings

• End-to-end speedup over vLLM: Vegas achieves a 1.25×–2.81× speedup in decoding throughput over default vLLM, and a 1.25×–2.70× speedup on the short-input reasoning workloads (AIME25 and CodeElo).

• Speedup over state-of-the-art sparse self-speculative decoding: 1.15×–1.29× over prior sparse-attention-based self-speculative decoding methods overall, and 1.15×–1.23× over MagicDec-Stream on the short-input reasoning benchmarks.

• Motivating trade-off (MagicDec study): On Qwen3-8B and gpt-oss-20b, Quest (query-aware sparsity) achieved a 22% improvement in the average number of tokens decoded per iteration at γ=5 versus StreamingLLM (query-agnostic), but KV selection inflated iteration duration by 23%, producing a 1.7% reduction in total decoding throughput relative to StreamingLLM (204.5 vs. 208.0 tokens/s).

• Multi-token selection beats last-token selection: Selecting the KV entries with highest attention weights for the last accepted token gives high acceptance probability for the first few draft tokens but a rapid decay for later ones; maximizing coverage across all draft tokens maintains consistently high acceptance probabilities. Ignoring discarded tokens decreases average accepted tokens per iteration by 3%–14% at γ=7 for both Qwen3-8B and gpt-oss-20b.

• Collect-2-Query preserves accuracy: In Table 1 (LongBench-v2, 7% of KV entries selected), using all draft tokens versus only the first draft plus bonus token gave: Qwen3-8B at γ=7 — 6.13 vs. 6.11; Qwen3-8B at γ=11 — 8.91 vs. 8.72; gpt-oss-20b at γ=7 — 6.25 vs. 6.24; gpt-oss-20b at γ=11 — 9.27 vs. 9.22.

• Logit collection overhead reduction: Collecting logits from all query tokens with 12 query tokens cost 53% overhead for Qwen3-8B and 189% for gpt-oss-20b. Collect-2-Query reduces this to as low as 5% for Qwen3-8B and 37% for gpt-oss-20b, and the overhead does not grow with γ.

• KV selection overhead versus baselines: For Qwen3-8B, MagicDec-Quest incurred 21.7% overhead and SpecExtend 11.2%; SpecExtend's one-time selection overhead rose to 29.1% on gpt-oss-20b, whose layers have a higher query-to-KV head ratio and smaller head dimensions. Vegas incurred only 5.9%–9.4%.

• Long-context behavior: On LongBench-v2 with input lengths from 96K to 120K tokens, Vegas achieved 18%–29% higher decoding throughput than the most competitive baselines, with the advantage growing with context length.

• gpt-oss-20b caveat: All self-speculative methods showed smaller speedups on gpt-oss-20b because only half of its transformer blocks use full attention, so sparse attention reduces per-drafting-step latency only to 48% of vanilla decoding latency. Vegas applies sparse attention only to the dense attention blocks.

• Memory-pressure motivation: Serving Qwen3-8B on an H100 (batch size 4, 128K context) requires 36.9 ms per decoding step; the 72 GB KV cache dictates a theoretical minimum transfer time of 18.5 ms, more than 50% of the total step time.

• Hyperparameter behavior: As the sparsity ratio increases, the average number of accepted draft tokens rises initially and then stabilizes; γ should be aligned with the expected number of accepted tokens per iteration.

• Not reported: The paper does not report accuracy or quality-degradation numbers for full-attention baseline outputs versus Vegas outputs (it states the method is lossless by construction), and it does not report energy measurements despite discussing energy implications in the impact statement.

Methodology in Plain English

The researchers start from an existing technique: run the same LLM twice per iteration, once with sparse attention to quickly guess several tokens ahead (drafting), then once with full attention to check those guesses in parallel (verification). Because the full-attention pass already computes how strongly each generated token attends to every earlier token, that information is sitting there essentially for free. Vegas captures it.

Concretely, during verification the system records the raw attention scores (logits, not normalized weights, since logits are cheaper to capture) for the generated tokens. It then combines the scores from the draft tokens to score each prefix token in the cache, and keeps only the top-scoring fraction — using that subset as the context for the next round of drafting. The authors found that scoring based on a single token overfits to that token's pattern and causes acceptance to collapse later in the draft chain, whereas aggregating across draft tokens — including ones that were rejected — produces more stable acceptance. They also found that simply using the first draft token and the final bonus token captures nearly all the benefit of using every token, because selected KV sets from adjacent tokens overlap heavily while distant tokens capture more diversity.

To make this fast enough to matter, they modified the FlashAttention-3 kernel to dump BF16 logits to HBM outside the on-chip shared memory, and repurposed vLLM's PagedAttention kernel with a page size of one token to perform token-granular sparse attention (kernel latency is nearly identical from page size 1 to 16). Finally, they tuned two knobs — the sparsity ratio and the number of draft tokens γ — in a three-step procedure: sweep sparsity until accepted tokens stabilize, then sweep γ for throughput, then fine-tune sparsity again.

Evaluation used Qwen3-4B, Qwen3-8B, Qwen3-30B-A3B-Thinking-2507-FP8, and gpt-oss-20b (with MXFP4 quantization for MoE weights) on two NVIDIA H100 NVL GPUs (94 GB), comparing against vanilla vLLM, MagicDec with StreamingLLM and Quest, and SpecExtend.

Why This Matters

This work shows that speculative decoding's verification stage is an underused information source, not just a correctness filter. It suggests a general design principle for inference systems: look for expensive signals you are already computing and reuse them, rather than computing new approximations. For the sparse-attention literature, it offers a path to accuracy without a separate selection algorithm. For the speculative decoding literature, it improves drafting accuracy without training an auxiliary model or fine-tuning.

Real-world applications:

• Long-document analysis in legal, medical, and financial settings, where the paper's impact statement explicitly notes that quality cannot be traded for speed and contexts reach tens of thousands of tokens.
• Code generation and coding assistants, evaluated here via CodeElo, where outputs are long and generation latency is user-visible.
• Chain-of-thought and reasoning assistants, evaluated via AIME25, where a single answer can require roughly 19,680 generated tokens.
• Retrieval-augmented and agentic pipelines that repeatedly query very long contexts, where per-request throughput directly determines serving cost.

Industry relevance: Vegas is training-free and integrates into vLLM, an open-source serving framework widely used in production, without deploying a separate draft model. It targets H100-class hardware with large batches and is compatible with existing kernels. Because gains scale with context length, the approach becomes more attractive as models move toward 128K, 1M, and longer contexts.

Future Directions

• Reducing collection overhead on unusual architectures. gpt-oss-20b still shows 37% collection overhead and a 29.1% selection overhead for the one-time baseline, driven by a high query-to-KV head ratio and smaller head dimensions. Kernel-level or layout-level fixes for such models remain open.
• Handling hybrid attention architectures. Applying sparse attention to only half the blocks limits the achievable speedup to 48% of vanilla per-drafting-step latency on gpt-oss-20b; how to exploit banded/dense hybrid patterns more fully is unresolved.
• Composition with orthogonal acceleration. The paper notes that layer-skipping drafters and quantized drafters could be combined with Vegas's sparse attention; whether these compose cleanly or interact through the acceptance rate is not established.
• Comparing against concurrent verification-guided designs. SparseSpec selects critical KV entries from all draft tokens and incurs rematerialization overhead; the paper claims its two-token variant is substantially cheaper, but the broader design space of which tokens to sample for KV selection is left largely unexplored.

Target Audience

This paper is most useful to systems and inference engineers building or optimizing LLM serving stacks (vLLM, PagedAttention, FlashAttention), to researchers in efficient LLM inference working on KV cache management, sparse attention, or speculative decoding, and to graduate students with some background in Transformer internals who want a concrete example of co-designing two stages of a decoding pipeline. Practitioners mainly interested in model quality, training methods, or non-serving applications will find less of direct relevance, though the lossless-output guarantee and training-free property make it accessible to deployment-focused readers.

Authors’ abstract

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.