Technical Summary: Vegas – Self-Speculative Decoding with Verification-Guided Sparse Attention

1. Problem Statement

Long-context inference for Large Language Models (LLMs) is increasingly constrained by memory bandwidth bottlenecks caused by the linear growth of the Key-Value (KV) cache with context length. While sparse attention techniques reduce memory access by processing only a subset of critical KV entries, they often degrade output quality by discarding context. Conversely, self-speculative decoding with sparse attention attempts to balance speed and quality by drafting tokens with sparse attention and verifying them with full attention. However, existing approaches suffer from a fundamental trade-off:

• Drafting Accuracy vs. Overhead: Simple, query-agnostic methods (e.g., sliding windows) have low selection overhead but poor drafting accuracy. Sophisticated, query-aware methods improve accuracy but incur high computational overhead in identifying critical KV entries, often negating the speedup gains.
• Disconnected Processes: Current methods treat the drafting and verification phases as standalone processes. They utilize the verification phase solely for token acceptance/rejection, overlooking the fact that the full attention computation during verification inherently calculates the exact criticality (attention weights/logits) of every KV entry. This "free oracle" is currently unused for guiding the subsequent drafting phase.

2. Methodology: Vegas

Vegas is a self-speculative decoding framework that co-designs the drafting and verification phases to leverage verification results for guiding sparse attention.

Core Mechanism: Verification-Guided Sparse Attention

Instead of using a standalone algorithm to select KV entries for drafting, Vegas treats the drafting phase as dependent on the verification phase of the previous iteration.

1. Verification Phase: The model performs a standard forward pass with full attention to verify draft tokens. During this process, it computes attention logits for all KV entries.
2. KV Selection: Vegas identifies critical KV entries as a byproduct of this verification. These entries are then used to construct the sparse attention mask for the next drafting phase.
3. Drafting Phase: The model drafts a sequence of tokens using sparse attention restricted to the selected critical KV entries.

Key Algorithmic Innovations

A. Aggregated Attention Logit Selection

To avoid overfitting to a single token's attention pattern, Vegas does not select KV entries based solely on the last accepted token. Instead, it aggregates attention logits across all draft tokens (including both accepted and discarded tokens) from the verification phase.

• Rationale: Discarded tokens often possess high acceptance probabilities and remain semantically aligned with the target generation. Including them in the selection process captures the broader local context required for the subsequent token chain, preventing rapid decay in acceptance probability for later draft tokens.
• Formulation: The selected set of prefix tokens $T$ maximizes the cumulative attention logits across all draft tokens and query heads:
  $$T = \arg\max_{T \subseteq P, |T|=k} \sum_{t=1}^{\gamma+1} \sum_{h \in H} \sum_{i \in T} l^{(h)}_{ti}$$

B. "Collect-2-Query" Mechanism

Collecting attention logits for all draft tokens during verification can introduce significant memory bandwidth overhead, especially in optimized kernels like FlashAttention-3. Vegas introduces a "Collect-2-Query" strategy to minimize this cost:

• Observation: KV entries selected by consecutive tokens exhibit high overlap. The overlap diminishes as the positional distance between tokens increases.
• Strategy: Instead of collecting logits for all $\gamma+1$ tokens, Vegas extracts logits only from the first draft token and the bonus token (the token generated after the last draft token). These two tokens represent the maximum positional distance in the draft chain.
• Impact: This reduces the overhead of logit collection to a minimal, constant level (scaling with 2 tokens rather than $\gamma$) while maintaining drafting accuracy comparable to collecting all tokens.

C. Hyperparameter Tuning

Vegas employs a systematic three-step tuning strategy to balance the sparsity ratio (percentage of KV entries selected) and the number of draft tokens ($\gamma$):

1. Identify the sparsity ratio where the number of accepted tokens stabilizes.
2. Sweep $\gamma$ to maximize decoding throughput given the selected sparsity.
3. Fine-tune the sparsity ratio to balance drafting latency and accuracy.

3. Implementation

Vegas is implemented within the vLLM serving framework:

• Kernel Instrumentation: The authors modified the FlashAttention-3 kernel to write attention logits to global memory (HBM) before softmax normalization, without occupying on-chip shared memory.
• Sparse Execution: A custom hook dispatches the attention kernel. For drafting, it uses a PagedAttention kernel with a page size of 1 token, passing selected token indices as page indices to perform computation only on critical entries.
• Integration: Vegas implements the vLLM speculative decoding proposer API, calling the target model's forward pass with sparse attention multiple times.

4. Experimental Results

Experiments were conducted on diverse models (Qwen3 family, gpt-oss-20b) using benchmarks including AIME25, CodeElo, and LongBench-v2 (up to 128K context) on NVIDIA H100 GPUs.

• Throughput Speedup:
  • vs. Default vLLM: Vegas achieves a 1.25×–2.81× speedup in decoding throughput.
  • vs. State-of-the-Art: Vegas achieves a 1.15×–1.29× speedup over existing sparse attention-based self-speculative decoding methods (MagicDec and SpecExtend).
• Drafting Accuracy: Vegas consistently achieves higher drafting accuracy (average accepted tokens per iteration) compared to baselines. For instance, on Qwen3-8B, it outperforms MagicDec-Stream by 1.15×–1.23×.
• Overhead Efficiency: The "Collect-2-Query" mechanism keeps KV selection overhead low (5.9%–9.4% for Qwen3-8B), significantly lower than query-aware baselines like MagicDec-Quest (21.7%) or SpecExtend (11.2%–29.1%).
• Long-Context Scaling: The performance gains of Vegas scale with context length. On LongBench-v2 (96K–120K tokens), Vegas maintains high acceptance rates where baselines struggle, achieving 18%–29% higher throughput than the most competitive baselines.

5. Significance and Claims

The paper claims that Vegas represents a significant advancement in lossless LLM acceleration by fundamentally rethinking the relationship between drafting and verification.

• Co-Design: It demonstrates that treating drafting and verification as a unified process allows the system to utilize the "free oracle" of verification results to guide sparse attention, resolving the traditional trade-off between drafting accuracy and selection overhead.
• Training-Free: Unlike methods requiring auxiliary draft models or model retraining, Vegas is a training-free approach that works with off-the-shelf models.
• Lossless Quality: By verifying all draft tokens with full attention, Vegas guarantees that the final output distribution matches vanilla full-attention decoding, ensuring no degradation in generation quality (e.g., hallucinations).
• Scalability: The approach is particularly effective for long-context scenarios where memory bandwidth is the primary bottleneck, offering a path to efficient inference as context lengths scale to millions of tokens.