Elastic Threshold Attention: Learned Contextual Sparsity for Long-Context Decoding
Organizations: Google · Boston University
Abstract
Massive KV caches can cause severe memory-bandwidth bottlenecks during long-context decoding. Sparse attention methods mitigate this problem, but often drop necessary context, leading to quality degradation. We introduce Elastic Threshold Attention (ETA), an end-to-end trainable architecture that rivals dense model quality under hardware-aligned block-sparse decoding. ETA predicts dynamic, contextual thresholds directly from query representations, adjusting context retention depending on the task at hand. To learn this policy from scratch while enabling near lossless KV cache pruning at inference time, we filter attention logits through a shifted SiLU gate during training. We show theoretically and empirically that this creates a smooth, near-uniform attention floor that neutralizes sub-threshold value contributions while simultaneously causing localized attention sinks on initial tokens to disappear. To materialize these advantages, we implement a fused inference-time kernel in Triton that screens KV blocks in time using cached geometric-probabilistic bounds. Across language modeling, reasoning, and RULER benchmarks, our 1.45B ETA model matches or exceeds dense quality, outperforming alternative fast decoding methods (Quest, HO, NSA) while achieving higher sparsity levels. Our kernel also achieves up to end-to-end speedups over FlashAttention-2 at context lengths of up to K tokens. Finally, we introduce an offline calibration algorithm for domain-specific deployments that freezes per-head constant thresholds, cutting attention compute by an additional 27% at no quality cost.
Figures & tables
| PPL ( ) | ||||
|---|---|---|---|---|
| Method | Config | Recall ( ) | C4 PPL ( ) | Needle ( ) | ||
|---|---|---|---|---|---|---|
| Dense Baseline | FA2 / SDPA | |||||
| ETA ( ) | ||||||
| ETA ( ) | ||||||
| union ( head) | union ( head) | union ( head) | |||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|
| KV | FA2 | ETA | FA2 | ETA | FA2 | ETA | |||||
| K | GB | ||||||||||
| K | GB | ||||||||||
| K | GB | ||||||||||
| K | GB | ||||||||||
| K | GB | ||||||||||
Appendix figures & tables18 assets
Supplementary material from the paper’s appendix.
Appendix
| Metric / Benchmark | Normal ETA | Strict ETA |
|---|---|---|
| Pretraining Loss (Step 10k) | 4.3986 | 4.4085 |
| FineWeb (200 batches, 1.62M tokens) | ||
| Perplexity (PPL) | 99.80 22.59 | 116.38 24.38 |
| Compute Sparsity (FLOPs Pruned) | 65.91% (34.09% density) | 37.26% (62.74% density) |
| GQA Union Sparsity (I/O Pruned) | 49.23% (50.77% density) | 23.59% (76.41% density) |
| C4 (200 batches, 1.62M tokens) | ||
| Sub-threshold tail metric | Pre-gate (raw) | Normal mode (multiplicative) | Strict mode (additive) |
|---|---|---|---|
| Sub-threshold probability mass | (preserved) | (annihilated) | |
| Normalized tail entropy ( ) | (near-uniform) | (collapsed) | |
| Tail max-to-mean ratio | (flattened) | (extreme peakiness) | |
| Block-sparse kernel decode, FineWeb | |||
| Perplexity ( ) | — | ( ) | |
| Compute density ( ) | — | ( ) | |
| Statistic | lse | Joint (All 8) | |||||||
|---|---|---|---|---|---|---|---|---|---|
| Evaluation Regime | Metric | Dynamic | Calibrated Static |
|---|---|---|---|
| In-Distribution ( ) | Perplexity ( ) | ||
| Active Head Density ( ) | |||
| Length Extrapolation ( , ) | Perplexity ( ) | ||
| Active Head Density ( ) |
| union | union | union | ||||
|---|---|---|---|---|---|---|
| Arm | ms | ms | ms | |||
| FlashAttention-2 | — | — | — | |||
| ETA , | ||||||
| ETA , | ||||||
| ETA , | ||||||
| ETA , | ||||||
| Pruned Set Size | ||||||
|---|---|---|---|---|---|---|
| Theoretical Rate ( ) | ||||||
| Empirical 1.45B Norm |
| 126M Model (Layers 0–11) | 1.45B Model (Layers 0–21, Sub-Block Envelope + ) | |||||||
|---|---|---|---|---|---|---|---|---|
| Layer | Layer | Layer | ||||||
| Layer 0 | 0.9997 | 0 0.96% | Layer 0 | 1.0000 | 0 0.01% | Layer 11 | 0.9881 | 13.56% |
| Layer 1 | 0.9960 | 0 8.98% | Layer 1 | 1.0000 | 0 0.34% | Layer 12 | 0.9867 | 14.19% |
| Layer 2 | 0.9782 | 21.06% | Layer 2 | 1.0000 | 0 0.41% | Layer 13 | 0.9863 | 14.38% |
| Layer 3 | 0.9678 | 25.37% | Layer 3 | 0.9999 | 0 1.26% | Layer 14 | 0.9860 | 14.53% |
| Layer 4 | 0.9613 | 27.52% | Layer 4 | 0.9997 | 0 2.19% | Layer 15 | 0.9863 | 14.19% |
| Hyperparameter | NSA 126M | ETA 126M |
|---|---|---|
| Layers ( ) / Hidden Size ( ) / FFN ( ) | / / | / / |
| Query Heads ( ) / KV Heads ( ) / Head Dim ( ) | / / | / / |
| Total Parameters (excl. output embedding) | ( ) | ( ) |
| Vocabulary Size / RoPE Base ( ) | / (LLaMA-3 ) | / (LLaMA-3 ) |
| Dataset | NSA PPL | ETA PPL | PPL | ETA Sparsity |
|---|---|---|---|---|
| FineWeb (10BT) | [ ] | [ ] | 53.60% | |
| C4 | [ ] | [ ] | 53.89% | |
| Wikitext-2 | [ ] | [ ] | 51.17% |
| Batch ( ) | Context ( ) | Model | Prefill (TTFT) | Step Latency (ms) | Throughput (tok/s) |
|---|---|---|---|---|---|
| 16 | 16,384 (16K) | NSA | 0.992 s | 260.93 | |
| 16 | 16,384 (16K) | ETA | 0.319 s | 659.68 | |
| 16 | 32,768 (32K) | NSA | 1.724 s | 262.59 | |
| 16 | 32,768 (32K) | ETA | 0.649 s | 428.31 | |
| 32 | 8,192 (8K) | NSA | 0.174 s | 463.31 | |
| 32 | 8,192 (8K) | ETA | 0.294 s | 1,262.16 |
| False neg. | False pos. | Token recall | Missed mass | PPL | |
|---|---|---|---|---|---|
| Method Variant | Offset / Budget | Head ( ) | Union ( ) | Recall / FNR | C4 PPL ( ) | Needle ( ) |
| Dense Baseline | FlashAttention-2 | / | ( ) | |||
| ETA ( ) | / | ( ) | ||||
| / | ( ) | |||||
| / | ( ) | |||||
| / | ( ) | |||||
| / | ( ) |
| Context | Method | Head ( ) | Union ( ) | Single | MultiKey | MultiVal | VT (2-Hop) | RULER Avg |
|---|---|---|---|---|---|---|---|---|
| Dense Baseline | ||||||||
| ETA ( ) | ||||||||
| ETA ( ) | ||||||||
| ETA ( ) | ||||||||
| Quest ( , top- ) | ||||||||
| Quest ( , top- ) |
| Context | Method ( ) | 22-Layer Attn (ms) | Total Step (ms/tok) | E2E Speedup |
|---|---|---|---|---|
| 64K | Dense FlashAttention-2 | |||
| ETA union ( head) | ||||
| ETA union ( head) | ||||
| ETA union ( head) | ||||
| 128K | Dense FlashAttention-2 | |||
| ETA union ( head) |
Long-Context Language Model Inference
Long-context language model inference makes processing and generating over lengthy prompts faster and more memory-efficient.
See all 227 papers