SANTA Slashes AI Memory Costs by 90% at Long Contexts Without Losing Accuracy
A new paper introduces SANTA, a stochastic sparse attention method that cuts value-cache memory reads at decode time, delivering up to 1.5× kernel speedup with no accuracy loss at 32k-token contexts.
- What: Researchers introduce SANTA, a stochastic sparse attention method that samples a small subset of value vectors instead of reading all of them during LLM decoding.
- Speed: S²ANTA CUDA kernels achieve up to 1.5× decode-step attention kernel speedup and 1.25× end-to-end decode latency speedup over FlashAttention-2 on Llama-3.1-8B at 32k-token contexts.
- Accuracy: At accuracy-matching operating points, S²ANTA matches SDPA baseline within confidence intervals on RULER, LongBench v2, GSM8K, and MMLU benchmarks.
- Efficiency: The method eliminates post-softmax multiplications entirely, replacing them with gather-and-add; value-stage memory access is reduced by over 90% at 32k contexts.
- Bonus method: Bernoulli qK^T sampling is proposed as a complementary score-stage technique that reduces key feature access by ~27% with minimal accuracy loss.
- Open source: CUDA kernels are released on GitHub and are compatible with KV quantization, cache compression, and selection methods.
Every token an LLM generates at long context requires the GPU to read the entire key-value (KV) cache , a cost that grows linearly with sequence length. For Llama-3.1-8B at a 32k-token context, that means streaming roughly 128 MB of key and value vectors per layer, per token. At that scale, the bottleneck is no longer arithmetic; it's memory bandwidth. A new research paper proposes a fundamentally different approach to this problem: instead of reading all the values, sample a small random subset and average them.
The paper, titled "Stochastic Sparse Attention for Memory-Bound Inference", introduces SANTA (Stochastic Additive No-mulT Attention), a drop-in replacement for the value aggregation step in standard attention. The authors , Kyle Lee, Corentin Delacour, Kevin Callahan-Coray, Kyle Jiang, Can Yaras, Samet Oymak, Tathagata Srimani, and Kerem Y. Camsari , also release their CUDA kernels as open source.
Why the value stage is the right place to attack
Standard scaled dot-product attention (SDPA) has two expensive memory stages during decoding: computing attention scores via qK^T (the score stage), and then aggregating value vectors weighted by those scores (the value stage). Both require reading the full KV cache. Prior work like FlashAttention and FlashDecoding optimized how you read everything; SANTA asks whether you need to read everything at all.
Autoregressive decoding becomes bandwidth-limited at long contexts, as generating each token requires reading all key and value vectors from the KV cache. The key insight in SANTA is that the post-softmax attention weights form a probability distribution over keys. That distribution can be sampled. If you draw S indices from it and average the corresponding value rows, you get an unbiased estimate of the full weighted sum , with no multiplications required after the softmax.
The math: sampling as Monte Carlo attention
In standard attention, the output is AV , the attention weight matrix multiplied by all value rows. SANTA replaces this with:
# SANTA: sample S indices from the softmax distribution
# and average the corresponding value rows
import torch
def santa_attention(q, K, V, S):
# Compute attention scores and softmax
scores = (q @ K.T) / (K.shape[-1] ** 0.5)
A = torch.softmax(scores, dim=-1) # shape: [n_k]
# Sample S indices from the categorical distribution
indices = torch.multinomial(A, num_samples=S, replacement=True)
# Gather sampled value rows and average
sampled_V = V[indices] # shape: [S, d_k]
return sampled_V.mean(dim=0) # unbiased estimate of AVThis story is for Pro members
You've reached the end of the free preview. Upgrade to AlphaSignal Pro to read the full article - and everything else behind the paywall.
Comments
