# Speculative Decoding and Blockwise Sparse Attention

## Introduction
Scaling the context length of LLMs continues to be crucial for improving their capabilities and unlocking new categories of applications. However, this leaves a [greater proportion](https://arxiv.org/pdf/2403.14123) of newer ML accelerators underutilized, since it decreases the [operational intensity](https://en.wikipedia.org/wiki/Roofline_model#Arithmetic_intensity) of generating tokens.

New attention mechanisms such as [DeepSeek’s Native Sparse Attention (NSA)](https://arxiv.org/abs/2502.11089v2) and [Moonshot AI’s Mixture of Block Attention (MoBA)](https://arxiv.org/abs/2502.13189) divide the context into blocks and allow each token to dynamically attend to just a subset of these blocks. By reducing the volume of data that needs to be moved from high bandwidth memory (HBM) into compute cores, these “blockwise sparse” methods aim to increase the operational intensity of decoding.

## Operational Intensity of Verification
Speculative decoding (SD) uses a cheap draft model to generate k candidate tokens that are verified in parallel by a target model. SD reduces the latency or increases the throughput of decoding by the formula:

\begin{align*}
\text{Speedup} &= \frac{\mathbb{E}[\text{\# Tokens Generated per Round of Speculation}]}{(k \times t_\text{draft} + t_\text{verify}) / t_\text{decode}}
\end{align*}

Where a round of speculation is the generation and subsequent verification of each block of k draft tokens.

| Full Attention | 8192 | 16384 | 32768 | 65536 |
| ---------------|------|-------|-------|-------|
| NSA (verification best case) | 2048 | 2560 | 3584 | 5632 |
| NSA (verification worst case) | 5120 | 5632 | 6656 | 8704 |
| NSA (reduction in memory access volume in best case compared to worst case) | **2.5×** | **2.2×** | **1.9×** | **1.5×** |

Table 1: Memory access volume (in equivalent number of tokens) during a forward pass. The size of the KV cache is the same when decoding one token or verifying k draft tokens.

## Forcing a Block of Tokens to Attend to the Same Subset of Context
To always achieve the best case operational intensity during verification, we modify the KV cache so that all tokens in a block attend to the same subset of keys and values. This method allows for increased efficiency when loading from HBM during the verification phase.

Figure 1: Operational intensity of verification as a function of the number of draft tokens for target models with 3B active parameters, showing significant variations in performance based on the subset of context attended to.

## Ablations
We investigate various training methods and the effect on model quality when forcing blocks of tokens to attend to the same context subset. Our tests reveal that while training methods can affect model performance, maintaining selected attention paths is important for preserving quality across a range of draft token configurations.

| Config | Cross-Entropy |
| -------|---------------|
| Baseline | **2.128** |
| No selection | **2.147** |
| Free selection | **2.132** |
| Test-time only | **2.128** |

Table 3: Cross-entropy results for models with 184 million parameters.

## Conclusion
We have shown that we can force a block of tokens to attend to the same subset of the context while preserving model quality. Our findings emphasize the significance of operational intensity in speculative decoding mechanisms, suggesting further exploration in this line of research.