SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers

As large language models (LLMs) scale to handle massive context windows, the quadratic compute cost of prefill and the linear memory cost of the key-value (KV) cache during decoding become significant bottlenecks. While sparse attention techniques can mitigate these costs by attending only to relevant tokens, the process of selecting those tokens is often computationally expensive itself. This paper investigates whether a lightweight, plug-in selector can effectively identify relevant tokens for pretrained, frozen LLMs without sacrificing accuracy.

SpotAttention introduces a small, trainable selector module that attaches to every full-attention layer of a pretrained transformer. The selector is trained via KL distillation to mimic the attention distribution of the frozen backbone. It operates on blocks of tokens rather than individual tokens, allowing the selection process to be executed as a single, efficient tensor-core matrix multiplication. The authors introduce a dual top-p rule that dynamically adjusts the per-query, per-layer budget based on the selector's estimated distribution, while reserving specific blocks (sink and recency) to ensure stability.

SpotAttention demonstrates that sparse selection can match the accuracy of dense models across various benchmarks (RULER, BABILong, InfiniteBench, and LongBench-v2) at context lengths up to 128K. In terms of performance, decoding at 128K context is 3.9x faster than standard FlashAttention and 1.8x faster than the strongest training-free baseline, Twilight. Furthermore, the authors show that the selector's KV-cache can be quantized to INT4 or FP4 precision without any degradation in model accuracy, providing additional memory savings.

This approach provides a practical, plug-and-play solution for deploying long-context LLMs in resource-constrained environments. By decoupling the selection mechanism from the backbone training, SpotAttention allows developers to retrofit existing, high-performing models with efficient sparse routing, significantly reducing the latency and memory overhead of long-context inference.