Elastic Threshold Attention: Learned Contextual Sparsity for Long-Context Decoding

Themistoklis Haris; Henry Li; Maryam Karimzadehgan

cs.LGSep 16, 2026

Elastic Threshold Attention: Learned Contextual Sparsity for Long-Context Decoding

Authors: Themistoklis Haris, Henry Li, Maryam Karimzadehgan

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 $O(1)$ 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, H $_2$ O, NSA) while achieving higher sparsity levels. Our kernel also achieves up to $2.15\times$ end-to-end speedups over FlashAttention-2 at context lengths of up to $512$ 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

Panel of Figure 1 — Figure 1: Overview of Elastic Threshold Attention (ETA). Left (Training): Queries predict a dynamic threshold $\tau_{t}$ that gates attention logits via a differentiable SiLU filter, compressing unselected logits toward zero rather than $-\infty$ . Right (Inference): A custom Triton kernel evaluates geometric-probabilistic block bounds against $\tau_{t}$ in $\mathcal{O}(1)$ time, loading only active KV blocks from HBM.

Panel of Figure 2 — Figure 2: SiLU-gated tail tempering vs. support truncation. (a) A shifted $\operatorname{SiLU}$ in logit space tempers the attention distribution without dropping tokens from its support. (b) Across pruned blocks ( $B_{1}$ ) and co-admitted keys ( $k_{5},k_{6}$ ) in loaded blocks ( $B_{2}$ ), tempering compresses wide raw weight spreads (gray) into a tightly bounded band around $1.0$ while preserving near-threshold order ( $k_{2}<k_{3}$ and $k_{6}<k_{5}$ ).

Panel of Table 1 — Table 1: Language modeling perplexity and zero-shot reasoning accuracy at $L=2048$ on the 1.45B model. Lower sub-text indicates the average fraction of past tokens attended per head. Because ARC-Easy and HellaSwag prompts ( $25\text{--}110$ tokens) are shorter than the $W=256$ local window of the heuristic baselines, those methods degenerate to $100\%$ dense attention on reasoning tasks, whereas fine-grained ETA actively prunes redundant tokens.

$b$	$z$	$D_{\text{head}}$	$D_{\text{union}}$	PPL ( $\downarrow$ )
$4$	$2.0$	$34.77\%$	$51.57\%$	$102.32$
$8$	$2.0$	$37.23\%$	$53.95\%$	$98.39$
$16$	$1.5$	$32.70\%$	$47.51\%$	$99.07$
$32$	$1.0$	$39.00\%$	$50.61\%$	$95.92$
$64$	$1.0$	$\mathbf{44.26\%}$	$51.94\%$	$\mathbf{93.48}$

Table 2: Block size at matched

{\sim}51\%

GQA-union bandwidth and

126

M C4 PPL (

L=2048

Method	Config	$D_{\text{head}}$	$D_{\text{union}}$	Recall ( $\uparrow$ )	C4 PPL ( $\downarrow$ )	Needle ( $\uparrow$ )
Dense Baseline	FA2 / SDPA	$100\%$	$100\%$	$100\%$	$17.49\pm 0.33$	$66.7\%$
ETA ( $b=4$ )	$a=+0.4$	$36.3\%$	$66.6\%$	$98.3\%$	$\mathbf{17.56\pm 0.33}$	$\mathbf{66.7\%}$
	$a=-0.8$	$24.4\%$	$50.6\%$	$97.1\%$	$17.82\pm 0.33$	$\mathbf{66.7\%}$
	$a=-2.5$	$\mathbf{13.1\%}$	$\mathbf{30.1\%}$	$94.8\%$	$19.81\pm 0.34$	$\mathbf{66.7\%}$
ETA ( $b=64$ )	$a=+0.2$	$55.3\%$	$79.0\%$	$\mathbf{98.4\%}$	$18.30\pm 0.29$	$\mathbf{66.7\%}$
	$a=-0.4$	$48.3\%$	$69.7\%$	$97.3\%$	$18.60\pm 0.29$	$\mathbf{66.7\%}$

Table 3: Pareto Frontier on C4 Perplexity (

L=2048

), Token Recall, and Needle Retrieval (

L=4096

) on the 1.45B model across block sizes

b\in\{4,64\}

vs. Quest ( Tang et al., 2024 ) (

b=64

) and

\text{H}_{2}\text{O}

( Zhang et al., 2023 ) . Extended offset sweep in Table 17 .

Panel of Table 4 — Table 4: RULER multi-needle and 2-hop retrieval. Full 16-arm suite with 95% CIs in Subsection J.2 .

$B$	$L$	KV	FA2	ETA	$\times$	FA2	ETA	$\times$	FA2	ETA	$\times$
			$51.2\%$ union ( $38\%$ head)			$35.5\%$ union ( $25\%$ head)			$16.0\%$ union ( $10\%$ head)
$32$	$64$ K	$2$ GB	$0.710$	$0.983$	$0.72$	$0.710$	$0.767$	$0.93$	$0.711$	$0.526$	$\mathbf{1.35}$
$32$	$128$ K	$4$ GB	$1.398$	$1.830$	$0.76$	$1.395$	$1.426$	$0.98$	$1.396$	$0.936$	$\mathbf{1.49}$
$32$	$256$ K	$8$ GB	$2.772$	$3.510$	$0.79$	$2.775$	$2.732$	$1.02$	$2.774$	$1.747$	$\mathbf{1.59}$
$32$	$512$ K	$16$ GB	$5.529$	$6.787$	$0.81$	$5.525$	$5.275$	$1.05$	$5.533$	$3.364$	$\mathbf{1.64}$
$64$	$64$ K	$4$ GB	$1.757$	$1.833$	$0.96$	$1.885$	$1.435$	$\mathbf{1.31}$	$1.897$	$0.937$	$\mathbf{2.02}$

Table 5: Single decode-step attention latency, ETA versus FlashAttention-2. NVIDIA H100 SXM

80

GB, float16 ,

H_{Q}=32

H_{KV}=4

(

G=8

d_{h}=64

b=64

. Speedup (

\times

) is

t_{\text{FA2}}/t_{\text{ETA}}

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

FineWeb (200 batches, 1.62M tokens)
Metric / Benchmark	Normal ETA	Strict ETA
Pretraining Loss (Step 10k)	4.3986	4.4085
Perplexity (PPL)	99.80 $\pm$ 22.59	116.38 $\pm$ 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)

Appendix

Table 6: Downstream perplexity and sparsity comparison between Normal and Strict ETA (126M parameters, 10k pretraining steps). Evaluated using 1% prefill and 99% autoregressive decode under the block-sparse inference kernel (

b=4

, threshold adjustment

a=0.4

z=2.0

Block-sparse kernel decode, FineWeb
Sub-threshold tail metric	Pre-gate (raw)	Normal mode (multiplicative)	Strict mode (additive)
Sub-threshold probability mass	$42.4\%$	$\mathbf{43.8\%}$ (preserved)	$\mathbf{1.5\times 10^{-25}}$ (annihilated)
Normalized tail entropy ( $H/\log n$ )	$0.790$	$\mathbf{0.916}$ (near-uniform)	$\mathbf{0.081}$ (collapsed)
Tail max-to-mean ratio	$15.6$	$\mathbf{2.03}$ (flattened)	$\mathbf{249.7}$ (extreme peakiness)
Perplexity ( $\downarrow$ )	—	$\mathbf{99.80}$	$116.38$ ( $+16.6\%$ )
Compute density ( $\downarrow$ )	—	$\mathbf{34.09\%}$	$62.74\%$ ( $1.84\times$ )

Appendix

Table 7: Mechanistic and downstream comparison of masking formulations (126M parameters, FineWeb). Tail statistics are measured over the sub-threshold key set (

n=24

sequences,

L=1024

). Downstream metrics reflect block-sparse decode (

b=4,z=2.0,a=0.4

). Multiplicative gating flattens sub-threshold keys into a uniform background, achieving superior perplexity and density over additive deletion.

Statistic	$\\|q_{t}\\|$	lse	$s_{\max}$	$s_{\text{std}}$	$H$	$\log t$	$s_{\text{mean}}$	$\\|x_{t}\\|$	Joint (All 8)
$R^{2}$	$0.187$	$0.176$	$0.168$	$0.156$	$0.091$	$0.071$	$0.059$	$0.021$	$\mathbf{0.327}$

Appendix

Table 8: Proportion of variance (

R^{2}

) in learned thresholds

\tau_{t,h}^{\ell}

explained by surface query and score-distribution statistics on held-out C4 across all layers and heads (126M model). Even when combined in a joint multiple regression, all eight statistics account for less than one-third of threshold variance.

Evaluation Regime	Metric	Dynamic $\tau_{t,h}^{\ell}$	Calibrated Static $c^{\ell h}$
In-Distribution ( $L=1024$ )	Perplexity ( $\downarrow$ )	$91.96$	$\mathbf{91.60}$
	Active Head Density ( $\downarrow$ )	$24.14\%$	$\mathbf{22.40\%}$
Length Extrapolation ( $L=2048$ , $z=3.0$ )	Perplexity ( $\downarrow$ )	$100.52$	$\mathbf{100.48}$
	Active Head Density ( $\downarrow$ )	$44.77\%$	$\mathbf{32.83\%}$

Appendix

Table 9: Dynamic per-token thresholds (

\tau_{t,h}^{\ell}

) versus calibrated static per-(layer, head) thresholds (

c^{\ell h}

) , evaluated on held-out C4 via the block-sparse decode kernel (

b=4

, 126M model with

L_{\text{train}}=1024

). Constants are calibrated on a disjoint

16

-sequence split. In-distribution evaluations are paired at the batch level (

\rho=0.9996

t=8.64

Panel of Figure 6 — Figure 6: Inference-time block-sparse selective loading mechanism. Unimportant blocks are pruned in SRAM via lightweight metadata screening, while active blocks are loaded from HBM for fused online softmax accumulation.

	$51.2\%$ union		$35.5\%$ union		$16.0\%$ union
Arm	ms	$\times$	ms	$\times$	ms	$\times$
FlashAttention-2	$15.085$	—	$15.009$	—	$14.790$	—
ETA $b=32$ , $c=64$	$16.617$	$0.91$	$13.498$	$1.11$	$9.207$	$1.61$
ETA $b=64$ , $c=64$	$13.797$	$1.09$	$10.566$	$1.42$	$6.429$	$2.30$
ETA $b=64$ , $c=128$	$\mathbf{13.432}$	$\mathbf{1.12}$	$\mathbf{10.430}$	$\mathbf{1.44}$	$6.520$	$2.27$
ETA $b=128$ , $c=64$	$16.733$	$0.90$	$11.472$	$1.31$	$\mathbf{5.921}$	$\mathbf{2.50}$

Appendix

Table 10: Effect of block size (

b

) and chunk size (

c

) at batch size

B=64

L=512

K (

32.8

GB KV cache), NVIDIA H100 SXM

80

GB, float16 .

c

denotes blocks per chunk. Coarser blocks reduce metadata overhead in sparse regimes (

16\%

), while

b=64

achieves the best balance at

51.2\%

union density.

Pruned Set Size $m=\|\mathcal{P}_{t}\|$	$m=16$	$m=32$	$m=64$	$m=128$	$m=256$	$m=384$
Theoretical Rate ( $1/\sqrt{m}$ )	$0.2500$	$0.1768$	$0.1250$	$0.0884$	$0.0625$	$0.0510$
Empirical 1.45B Norm	$\mathbf{0.2504}$	$\mathbf{0.1750}$	$\mathbf{0.1202}$	$\mathbf{0.0799}$	$\mathbf{0.0448}$	$\mathbf{0.0274}$

Appendix

Table 11: Empirical verification of 1 (

1/\sqrt{m}

background value cancellation) on the 22-layer 1.45B ETA model. Normalized centered norm

\|\bar{v}_{\mathcal{P}_{t}}-\boldsymbol{\mu}_{v}\|_{2}/\sqrt{\operatorname{Tr}(\Sigma_{v})}

measured across pruned set sizes

m=|\mathcal{P}_{t}|

on C4, compared against the theoretical

1/\sqrt{m}

rate.

126M Model (Layers 0–11)			1.45B Model (Layers 0–21, $b=64$ Sub-Block Envelope + $\mathtt{lwb}=2$ )
Layer	$\mathcal{S}_{\cos}$	$\mathcal{E}_{\ell_{2}}$	Layer	$\mathcal{S}_{\cos}$	$\mathcal{E}_{\ell_{2}}$	Layer	$\mathcal{S}_{\cos}$	$\mathcal{E}_{\ell_{2}}$
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%

Appendix

Table 12: Layer-wise residual stream alignment (

\mathcal{S}_{\cos}

) and relative

\ell_{2}

shift (

\mathcal{E}_{\ell_{2}}

) across both the 12-layer 126M and 22-layer 1.45B ETA models on held-out C4 sequences during autoregressive block-sparse decoding.

Hyperparameter	NSA 126M	ETA 126M
Layers ( $L$ ) / Hidden Size ( $d_{\text{model}}$ ) / FFN ( $d_{\text{ffn}}$ )	$12$ / $768$ / $2048$	$12$ / $768$ / $2048$
Query Heads ( $H_{Q}$ ) / KV Heads ( $H_{KV}$ ) / Head Dim ( $d_{k}$ )	$12$ / $4$ / $64$	$12$ / $4$ / $64$
Total Parameters (excl. output embedding)	$125.85\text{M}$ ( $100.66\text{M}$ )	$125.85\text{M}$ ( $100.66\text{M}$ )
Vocabulary Size / RoPE Base ( $\theta$ )	$32{,}768$ / $5\times 10^{5}$ (LLaMA-3 $8\times$ )	$32{,}768$ / $5\times 10^{5}$ (LLaMA-3 $8\times$ )

Appendix

Table 13: Backbone Architectural Hyperparameters. All backbone dimensions and positional encodings are identical between the two models. ETA adds

12\times(768\times 12+12)=0.11

M threshold-predictor parameters (

<0.1\%

of total parameters).

Panel of Figure 8 — Figure 8: Pretraining Loss Convergence and Matched-H100 Throughput. Training loss trajectories across $10{,}000$ steps ( $164$ M tokens, $L=1024$ ) on identical H100 GPUs. ETA achieves $\Delta\text{Loss}=-0.11\text{ to }-0.14$ ( $\Delta\text{PPL}=-9.4$ ) below NSA at matched step latency ( $0.5315$ s/step vs. $0.5311$ s/step).

Dataset	NSA PPL	ETA PPL	$\Delta$ PPL	ETA Sparsity
FineWeb (10BT)	$103.55\pm 26.51$ [ $\pm 1.87$ ]	$\mathbf{99.80\pm 21.71}$ [ $\pm 1.54$ ]	$\mathbf{-3.75}$	53.60%
C4	$\mathbf{100.18\pm 16.08}$ [ $\pm 1.14$ ]	$100.73\pm 17.96$ [ $\pm 1.27$ ]	$+0.55$	53.89%
Wikitext-2	$161.65\pm 22.50$ [ $\pm 1.24$ ]	$\mathbf{158.64\pm 22.30}$ [ $\pm 1.17$ ]	$\mathbf{-3.01}$	51.17%

Appendix

Table 14: Perplexity (PPL) and Sparsity Comparison.

126

M models, evaluated under 1% prompt prefill + 99% autoregressive decode (batch size

4

L=2048

Batch ( $B$ )	Context ( $L$ )	Model	Prefill (TTFT)	Step Latency (ms)	Throughput (tok/s)
16	16,384 (16K)	NSA	0.992 s	$61.319\pm 2.248$	260.93
16	16,384 (16K)	ETA	0.319 s	$\mathbf{24.254\pm 6.086}$	659.68
16	32,768 (32K)	NSA	1.724 s	$60.932\pm 2.927$	262.59
16	32,768 (32K)	ETA	0.649 s	$\mathbf{37.356\pm 12.475}$	428.31
32	8,192 (8K)	NSA	0.174 s	$69.069\pm 1.573$	463.31
32	8,192 (8K)	ETA	0.294 s	$\mathbf{25.353\pm 5.808}$	1,262.16

Appendix

Table 15: Decoding Efficiency & Throughput Sweep. Evaluated under 10% prompt prefill + 90% autoregressive decode across sequence lengths up to 32,768 tokens.

Panel of Figure 9 — Figure 9: GQA Union Loading vs. Compute Sparsity across KV Groups on C4 decoding (126M model).

$b$	False neg.	False pos.	Token recall	Missed mass	PPL
$4$	$\mathbf{3.4\%}$	$30.8\%$	$\mathbf{98.3\%}$	$\mathbf{10.2\%}$	$\mathbf{17.38}$
$16$	$11.7\%$	$16.8\%$	$96.9\%$	$21.5\%$	$18.47$
$64$	$27.5\%$	$\mathbf{5.6\%}$	$95.3\%$	$36.3\%$	$19.15$

Appendix

Table 16: Effect of variance dilution under flat block variance (

z=2.0,a=0.4

) before applying the sub-block envelope (

1.45

B model on FineWeb,

L=2048

). Missed mass is the fraction of softmax probability mass in rejected blocks.

Method Variant	Offset / Budget	Head ( $D_{\text{head}}$ )	Union ( $D_{\text{union}}$ )	Recall / FNR	C4 PPL ( $\downarrow$ )	Needle $4\text{K}$ ( $\uparrow$ )
Dense Baseline	FlashAttention-2	$100.0\%$	$100.0\%$	$100.0\%$ / $0.0\%$	$17.49\pm 0.33$	$66.7\%$ ( $6/9$ )
ETA ( $b=4,z=2.0$ )	$a=+0.4$	$36.3\%$	$66.6\%$	$98.3\%$ / $7.1\%$	$\mathbf{17.56\pm 0.33}$	$\mathbf{66.7\%}$ ( $6/9$ )
	$a=0.0$	$32.0\%$	$61.2\%$	$97.8\%$ / $8.9\%$	$17.61\pm 0.33$	$\mathbf{66.7\%}$ ( $6/9$ )
	$a=-0.8$	$24.4\%$	$50.6\%$	$97.1\%$ / $12.1\%$	$17.82\pm 0.33$	$\mathbf{66.7\%}$ ( $6/9$ )
	$a=-1.5$	$18.9\%$	$41.7\%$	$96.0\%$ / $16.1\%$	$18.22\pm 0.33$	$\mathbf{66.7\%}$ ( $6/9$ )
	$a=-2.0$	$15.7\%$	$35.6\%$	$95.4\%$ / $18.8\%$	$18.81\pm 0.34$	$\mathbf{66.7\%}$ ( $6/9$ )

Appendix

Table 17: Extended C4 Perplexity (

L=2048

) and Needle-in-a-Haystack (

L=4096

) Pareto sweep on the 1.45B model across block size

b

and threshold offset

a

, supplementing Table 3 .

Context	Method	Head ( $\text{Routine}\to\text{Query}$ )	Union ( $\text{Routine}\to\text{Query}$ )	Single	MultiKey	MultiVal	VT (2-Hop)	RULER Avg
$L=2048$	Dense Baseline	$100.0\%$	$100.0\%$	$87.0\%$	$76.9\%$	$88.9\%$	$87.0\%$	$84.95\%\pm 3.19\%$
	ETA ( $b=4,z=1.0,a=0.4$ )	$38.0\%\to 41.0\%$	$66.3\%\to 69.8\%$	$\mathbf{88.0\%}$	$87.0\%$	$\mathbf{88.9\%}$	$\mathbf{88.0\%}$	$\mathbf{87.96\%\pm 2.86\%}$
	ETA ( $b=64,z=1.5,a=0.4$ )	$40.6\%\to 43.5\%$	$54.6\%\to 57.9\%$	$87.0\%$	$\mathbf{88.0\%}$	$\mathbf{88.9\%}$	$77.8\%$	$85.42\%\pm 3.14\%$
	ETA ( $b=64,z=1.5,a=0.0$ )	$25.1\%\to 27.8\%$	$\mathbf{37.4\%\to 39.8\%}$	$87.0\%$	$76.9\%$	$78.7\%$	$67.6\%$	$77.55\%\pm 3.78\%$
	Quest ( $b=64$ , top- $38\%$ )	$39.1\%$	$51.2\%$	$55.6\%$	$23.1\%$	$12.0\%$	$13.0\%$	$25.93\%\pm 4.14\%$
	Quest ( $b=64$ , top- $25\%$ )	$\mathbf{23.1\%}$	$32.7\%$	$34.3\%$	$12.0\%$	$0.9\%$	$0.9\%$	$12.04\%\pm 3.07\%$

Appendix

Table 18: RULER Long-Context Suite on the 1.45B model (

L=2048

Native &

L=4096

with

2\times

NTK RoPE Scaling;

N=432

trials/arm). Reports per-task and average accuracy (

\pm 95\%

CI), plus routine haystack vs. query-step density (

\text{Routine}\to\text{Query}

Panel of Table 19 — Table 19: Single-Needle passkey retrieval accuracy across context lengths up to $L=8192$ ( $4\times L_{\text{train}}$ ). Accuracy averaged over $9$ needle depths ( $0\%\text{--}100\%$ ). Density reports position-averaged head density at $L=1024\to 8192$ .

Context $L$	Method ( $B=64$ )	22-Layer Attn (ms)	Total Step (ms/tok)	E2E Speedup
64K	Dense FlashAttention-2	$42.14$	$53.90$	$1.00\times$
	ETA $51.2\%$ union ( $38\%$ head)	$40.33$	$52.35$	$1.03\times$
	ETA $35.5\%$ union ( $25\%$ head)	$31.57$	$43.59$	$1.24\times$
	ETA $16.0\%$ union ( $10\%$ head)	$\mathbf{20.61}$	$\mathbf{32.63}$	$\mathbf{1.65\times}$
128K	Dense FlashAttention-2	$82.92$	$94.68$	$1.00\times$
	ETA $35.5\%$ union ( $25\%$ head)	$70.93$	$82.95$	$1.14\times$

Appendix

Table 20: Full 22-layer 1.45B model decode latency (

B=64

, float16 , H100 80GB). Total step time sums 22-layer attention, the constant non-attention floor (

11.76\text{ ms}

), and ETA predictor/metadata updates (

0.26\text{ ms}

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

Certified Sparse Attention via Perturbed Retrieval
Oct 18, 2026Runguo LiSparse Attention LLM Inference Acceleration
University of Illinois Urbana-Champaign
Kascade: A Practical Sparse Attention Method for Long-Context LLM Inference
Dec 18, 2025Dhruv Deshmukh, Saurabh Goyal, Nipun Kwatra +1Self-Attention Long-Context Language Model Inference
Microsoft Research India.
SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers
Jun 22, 2026Huzama Ahmad, Se-Young YunSparse Attention Block-Sparse Attention
KAIST

Oct 18, 2026J/K move · Enter open · S save

Certified Sparse Attention via Perturbed Retrieval

Runguo Li

University of Illinois Urbana-Champaign

Sparse Attention LLM Inference Acceleration

Open paper AskPDF arXiv

Sparse attention speeds up long-context decoding by reading few keys of the KV cache, but existing methods do not control the error of a query, and on real KV caches their errors have heavy tails. We show that perturbing the score of every key with an independent Exp(1) variable turns threshold retrieval into sampling keys with probability proportional to their attention weights. Lowering the threshold along a geometric grid gives nested samples, a Horvitz--Thompson estimate of the output at every level, and confidence bounds that hold simultaneously over all levels, so each head stops at the first level whose certificate meets the target error

\varepsilon

. With fresh perturbations at every decoding step, each certificate fails with probability at most

δ

during generation. At a suitable threshold, the estimator reads

O(\varepsilon^{-2}\log(1/δ))

values regardless of the context length, and no adaptive method reads fewer, up to constants. A 4-bit scan with deterministic bounds, shared across a grouped-query-attention group, generates the candidates exactly but reads every key. On three LLMs and ten settings at 16K--128K tokens, no certificate is violated, and at equal keys read our 90th-percentile error is 4.6--7.9

\times

lower than that of the best of 15 baseline families. At equal bytes, the best baselines have 0.6--1.6 times our 90th-percentile error and at least 4.9 times our largest. End to end, our method stays within one point of dense attention on tasks and within 0.1% in PG-19 perplexity, and the strongest baselines at equal bytes score as well. On CPU, decoding steps at 32K--128K tokens run 1.1--1.4

\times

faster than dense attention on the same 16-bit cache.

Abstract

Long-Context Language Model Inference

Explore similar work

Certified Sparse Attention via Perturbed Retrieval

Kascade: A Practical Sparse Attention Method for Long-Context LLM Inference

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