Stochastic Sparse Attention for Memory-Bound Inference

Autoregressive decoding becomes bandwidth-limited at long contexts, as generating each token requires reading all $n_k$ key and value vectors from KV cache. We present Stochastic Additive No-mulT Attention (SANTA), a method that sparsifies value-cache access by sampling $S \ll n_k$ indices from the post-softmax distribution and aggregates only those value rows. This yields an unbiased estimator of the post-softmax value aggregation while replacing value-stage multiply-accumulates with gather-and-add. We introduce stratified and systematic sampling to design variance-reduced, GPU-friendly variants. Evaluated on Llama-3.1-8B-Instruct at 32k-token contexts, S$^2$ANTA matches baseline accuracy while achieving up to $1.5\times$ decode-step attention-kernel speedup over FlashInfer and FlashDecoding on an NVIDIA RTX 6000 Ada. In batched long-context generation, these kernel gains translate to up to $1.25\times$ end-to-end decode-latency speedup. Finally, we propose Bernoulli $qK^\mathsf{T}$ sampling as a complementary technique to sparsify the score stage, reducing key-feature access through stochastic ternary queries. Both methods are complementary to upstream quantization, low-rank projection, KV-cache compression, and KV-cache selection methods. Together, they point toward sparse, multiplier-free, and energy-efficient inference. We open-source our kernels at: https://github.com/OPUSLab/SANTA.git

Large language models are increasingly asked to work with long documents and conversations. When generating each new word, they often repeatedly read a large memory of previous words; at long lengths, this data movement can become a speed and energy bottleneck. Our paper focuses on attention, the part of the model that looks back at earlier words to decide what information matters now. We introduce SANTA, a way to make this step cheaper by sampling a small number of stored entries instead of reading and multiplying through all of them. The key idea is like estimating an average from a carefully chosen sample: with enough samples, the model gets nearly the same answer while moving much less data. We design sampling strategies and GPU kernels that make this practical for long-context language model decoding. On Llama-3.1-8B-Instruct, SANTA matches standard attention accuracy on long-context tasks while speeding up the attention step by up to 1.5× and full batched generation by up to 1.25×. Because it changes how the model reads existing memory rather than permanently deleting context, it can work alongside compression and cache-management methods. This points toward faster and more energy-efficient language model inference for long-context applications.