SANTA++: Sampling Attention through Representative Keys

Kyle Lee; Christian Z. Pratt; Ruoyu Fang; Heekyung Lee; Avinash Lohitsa; Ryan Modafe; Kerem Y. Camsari

cs.LGSep 28, 2026

SANTA++: Sampling Attention through Representative Keys

Authors: Kyle Lee, Christian Z. Pratt, Ruoyu Fang, Heekyung Lee, Avinash Lohitsa, Ryan Modafe, Kerem Y. Camsari

Organizations: Department of Electrical and Computer Engineering University of California, Santa Barbara, CA, USA · Flucta, San Francisco, CA, USA

Abstract

Attention often concentrates on a small subset of tokens in the context, but which subset matters changes from one query to the next. To exploit this changing structure, we introduce SANTA++, a training-free stochastic attention method that uses representative keys for memory-efficient selection without scanning the entire key-value (KV) cache. Cached keys are organized into teams, and the query scores one representative from each team to decide which teams to sample. We compute exact attention scores within the sampled teams and reweight each team's contribution by the inverse of its inclusion probability. This importance sampling correction estimates attention over the full cache, with a sampling budget that lets us trade memory reads for accuracy. Remarkably, with 32 or 64 sampled teams, SANTA++ uses 16% to 22% of dense attention's KV reads and retains 94% to 99% of the dense-attention baseline's scores on LongBench v2 and HELMET's retrieval-augmented generation subset, and 85% to 91% on RULER, with Qwen2.5-7B-Instruct at 32K context. With 31 sampled teams, our GPU implementation delivers a $1.69\times$ attention speedup over the dense FlashAttention baseline at 32K context. By reducing the number of cache entries read, SANTA++ in principle complements architectures with compressed KV representations, such as multi-head latent attention. Our kernels are available at: https://github.com/OPUSLab/santapp-kernel-demo.git.

Figures & tables

Figure 1: Representative-guided attention sampling. After prefill, parent groups cover the prompt, and actual member keys represent the teams built within each parent. For each decoding query, we form a routing weight for each team from its representative score and team size, and use these weights to sample teams

\mathcal{S}

. Selected teams’ keys and values are read in full, and their members’ exact attention contributions are computed. Each team’s unnormalized attention mass and weighted value sum are divided by its conditional inclusion probability

c_{g}

before the corrected sums are combined and normalized (Section 2 ).

Figure 2: A mean-key mass surrogate for 16 constructed logits. One high score dominates the exact mass, which is approximately

760\times

the surrogate.

Task	Teams (k-means) $P=16,R=4$	Centroid-guided sampling $B=4$	Centroid-guided sampling $B=16$
CWE	1.621	2.806	3.429
S-NIAH-1	1.248	2.076	2.827
QA-1	1.613	2.429	3.495

Table 1: Routing approximation to dense attention. Mean

\mathrm{KL}(p\|q)

in nats, where

p

is dense attention over prompt tokens and

q

spreads each group’s routing probability equally across its members. All variants use the same dense SDPA queries and cache states. No sampling is performed in this diagnostic. Appendix B.2 gives the definitions and averaging procedure.

Method	$S$	KV (%)	Score (norm) (%)	Score (abs) (%)
SDPA	–	100.00	100.00	57.40 $\pm$ 2.00
SANTA	128	50.43	99.13	56.90 $\pm$ 2.05
	256	50.75	99.04	56.85 $\pm$ 2.05
	512	51.30	100.17	57.50 $\pm$ 2.02
	1024	52.17	99.48	57.10 $\pm$ 2.07
	2048	53.53	99.65	57.20 $\pm$ 2.08

Table 2: HELMET RAG: full 32K budget sweep with Qwen2.5-7B-Instruct. Score (abs) is the equal-weight mean substring exact match on Natural Questions, TriviaQA, PopQA, and HotpotQA. Score (norm) is the absolute score divided by the corresponding SDPA score at the same context length, expressed as a percentage. KV is the equal-weight mean of the four RAG task-level total-access percentages, including representative-key routing reads.

S

is the nominal budget. The

\pm

entries are 95% bootstrap confidence interval half-widths.

Method	$S$	KV (%)	Score (norm) (%)	Score (abs) (%)
SDPA	—	100.00	100.00	31.91 $\pm$ 2.04
SANTA	128	50.36	95.95	30.62 $\pm$ 2.06
	256	50.59	102.80	32.80 $\pm$ 1.96
	512	50.97	106.85	34.10 $\pm$ 2.04
	1024	51.56	99.84	31.86 $\pm$ 1.91
SANTA ++ (k-means)	128	16.68	99.07	31.61 $\pm$ 2.01

Table 3: LongBench v2: full budget sweep with a 32K context cap. Qwen2.5-7B-Instruct is evaluated on 503 prompts using the brief-reasoning prompt and middle-truncation procedure in Appendix C.3 . Score (abs) is accuracy. Score (norm) is the absolute score divided by the corresponding SDPA score, expressed as a percentage. KV is reported GQA-aware logical access, and

S

is the nominal budget. The

\pm

entries are the reported 95% bootstrap confidence interval half-widths.

Method	$S$	KV (%)	Score (norm) (%)	Score (abs) (%)
SDPA	—	100.00	100.00	83.51 $\pm$ 0.66
SANTA	128	50.32	99.15	82.80 $\pm$ 0.64
	256	50.54	99.75	83.31 $\pm$ 0.67
	512	50.91	100.11	83.61 $\pm$ 0.65
	1024	51.52	99.92	83.45 $\pm$ 0.67
	2048	52.49	99.99	83.50 $\pm$ 0.66

Table 4: RULER: full 32K budget sweep with Qwen2.5-7B-Instruct. Score (abs) is the equal-weight mean of all 13 task scores over 500 completed prompts per task. Score (norm) is the absolute score divided by the corresponding SDPA score at the same context length, expressed as a percentage. KV is pooled GQA-aware logical access, including representative-key routing reads, and

S

is the nominal budget. Absolute-score

\pm

entries are 95% task-stratified bootstrap confidence interval half-widths.

Method	GPU launches	Median latency ( $\mu$ s)	Speedup vs. Flash
SANTA ++	5	63.52	1.69 $\times$
Forced Flash SDPA	2	107.33	1.00 $\times$

Table 5: Complete attention-operator latency on the captured 32K, batch-one FP16 workload. Measurements include selection and all launches. Preparation is excluded.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Additional Jensen examples. Each row compares 16 raw logits

s_{i}

, their exponentials, and the exact-to-surrogate mass ratio

\sum_{i}e^{s_{i}}/\bigl[n\exp(n^{-1}\sum_{i}s_{i})\bigr]

. The displayed ratios are approximately (a)

1.03

for near-uniform logits, (b)

30

for bimodal logits, (c)

135

for the heavy-tailed example, and (d)

760

for the sparse example. The exponentiated-score axes are logarithmic.

		S-NIAH-1		S-NIAH-2		S-NIAH-3		MK-NIAH-1		MK-NIAH-2
Method	$S$	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)
SDPA	—	100.00	100.00 $\pm$ 0.00	100.00	100.00 $\pm$ 0.00	100.00	100.00 $\pm$ 0.00	100.00	100.00 $\pm$ 0.00	100.00	98.00 $\pm$ 2.50
SANTA ++ (Teams)	128	22.54	100.00 $\pm$ 0.00	25.37	100.00 $\pm$ 0.00	25.50	100.00 $\pm$ 0.00	24.95	99.00 $\pm$ 1.50	25.57	94.00 $\pm$ 4.50
	256	30.33	100.00 $\pm$ 0.00	34.75	99.00 $\pm$ 1.50	34.77	100.00 $\pm$ 0.00	34.05	100.00 $\pm$ 0.00	34.60	98.00 $\pm$ 2.50
	512	42.60	100.00 $\pm$ 0.00	48.49	100.00 $\pm$ 0.00	48.19	100.00 $\pm$ 0.00	47.63	100.00 $\pm$ 0.00	47.78	95.00 $\pm$ 4.50
Centroid-guided ( $B=4$ )	128	20.90	100.00 $\pm$ 0.00	22.06	100.00 $\pm$ 0.00	22.53	98.00 $\pm$ 2.50	21.98	99.00 $\pm$ 1.50	22.87	98.00 $\pm$ 2.50

Appendix

Table 6: Task-variant evaluation on RULER at 8K context. Qwen2.5-7B-Instruct with 100 prompts per task. SANTA ++ (teams) uses whole-team sampling with

P=16,R=4

; Centroid-guided (

B=4

) and (

B=16

) denote the two centroid-guided baselines.

S

is the nominal budget. KV is the reported theoretical GQA-aware access relative to dense SDPA. Accuracy is reported as mean

\pm

95% bootstrap confidence interval half-width. All Tasks accuracy is the equal-weight mean over all 13 tasks, with a stratified bootstrap that resamples prompts within each task. This is a separate evaluation from the 20-prompt dense-trajectory KL diagnostic.

Figure 4: Key-to-centroid distance across RULER tasks. Mean distance for key clusters, contiguous chunks, and random chunks on 100 prompts per task, using Qwen2.5-7B-Instruct at 8,192-token context and nominal group size

P=16

. Contiguous chunks lie between key clusters and random chunks on every plotted task.

Task	Context length	WER threshold	Heads (%)
CWE	8192	$\geq 0.50$	82.1 [79.8, 84.9]
CWE	32768	$\geq 0.50$	80.2 [76.9, 84.8]
CWE	8192	$\geq 0.75$	54.0 [48.4, 60.7]
CWE	32768	$\geq 0.75$	49.5 [41.7, 60.7]
CWE	8192	$\geq 0.90$	22.1 [18.4, 27.0]
CWE	32768	$\geq 0.90$	23.6 [16.2, 36.4]

Appendix

Table 7: Oracle prefill-window effectiveness. Percentage of heads across all layers meeting each WER threshold from Equation 12 . For each prompt, layer, and head, WER compares the attention mass in the top

K_{\mathrm{w}}=32

disjoint windows of width

w=16

with the mass in the

wK_{\mathrm{w}}=512

highest-attention tokens. Results use 10 prompts per task at 8K and 32K context. Brackets give 95% bootstrap confidence intervals from resampling prompts.

Figure 5: Key versus probe-fingerprint clustering. Accuracy on four RULER tasks with 50 prompts per task and nominal budgets $S=128,256,512$ . Bars compare key clustering, probe-fingerprint clustering, SANTA, and dense SDPA; error bars denote the displayed 95% confidence intervals.

Method	$S$	KV (%)	Acc (norm) (%)	Acc (abs) (%)
SDPA	–	100.00	100.00	62.65 $\pm$ 2.02
SANTA	128	51.62	100.24	62.80 $\pm$ 2.02
	256	52.70	99.60	62.40 $\pm$ 2.03
	512	54.39	99.60	62.40 $\pm$ 2.03
	1024	56.77	100.16	62.75 $\pm$ 2.05
	2048	60.14	100.00	62.65 $\pm$ 2.07

Appendix

Table 8: HELMET RAG: full 8K budget sweep with Qwen2.5-7B-Instruct. Acc (abs) is the equal-weight mean substring exact match on Natural Questions, TriviaQA, PopQA, and HotpotQA. Acc (norm) is the absolute score divided by the corresponding SDPA score at the same context length, expressed as a percentage. KV is the equal-weight mean of the four RAG task-level total-access percentages, including representative-key routing reads.

S

is the nominal budget. The

\pm

entries are the reported score-interval half-widths.

Method	$S$	Natural Questions		TriviaQA		PopQA		HotpotQA		MS MARCO
Method	$S$	Total access (%)	Acc (%)	Total access (%)	Acc (%)	Total access (%)	Acc (%)	Total access (%)	Acc (%)	Total access (%)	Acc (%)
SDPA	–	100.00	53.20 $\pm$ 4.30	100.00	82.20 $\pm$ 3.30	100.00	64.40 $\pm$ 4.10	100.00	50.80 $\pm$ 4.40	100.00	58.51 $\pm$ 3.89
SANTA	128	51.62	54.40 $\pm$ 4.40	51.55	82.40 $\pm$ 3.40	51.79	63.00 $\pm$ 4.20	51.51	51.40 $\pm$ 4.30	51.72	59.16 $\pm$ 3.81
	256	52.73	53.80 $\pm$ 4.40	52.58	81.80 $\pm$ 3.40	52.97	63.80 $\pm$ 4.20	52.52	50.20 $\pm$ 4.50	52.75	58.03 $\pm$ 3.75
	512	54.46	53.00 $\pm$ 4.30	54.21	82.20 $\pm$ 3.30	54.79	64.20 $\pm$ 4.20	54.09	50.20 $\pm$ 4.50	54.27	59.06 $\pm$ 3.84
	1024	56.90	54.40 $\pm$ 4.30	56.49	82.20 $\pm$ 3.30	57.30	64.00 $\pm$ 4.10	56.39	50.40 $\pm$ 4.40

Appendix

Table 9: HELMET task results at 8K context length. Each accuracy is reported as the mean

\pm

the half-width of its 95% confidence interval.

Method	$S$	Natural Questions		TriviaQA		PopQA		HotpotQA		MS MARCO
Method	$S$	Total access (%)	Acc (%)	Total access (%)	Acc (%)	Total access (%)	Acc (%)	Total access (%)	Acc (%)	Total access (%)	Acc (%)
SDPA	–	100.00	46.20 $\pm$ 4.30	100.00	82.00 $\pm$ 3.50	100.00	62.80 $\pm$ 4.30	100.00	38.60 $\pm$ 4.20	100.00	34.39 $\pm$ 4.36
SANTA	128	50.43	46.40 $\pm$ 4.30	50.41	81.00 $\pm$ 3.50	50.46	62.20 $\pm$ 4.20	50.41	38.00 $\pm$ 4.30	50.46	37.63 $\pm$ 4.60
	256	50.77	45.60 $\pm$ 4.40	50.71	81.00 $\pm$ 3.40	50.81	62.80 $\pm$ 4.20	50.71	38.00 $\pm$ 4.40	50.77	35.79 $\pm$ 4.58
	512	51.33	46.00 $\pm$ 4.40	51.23	81.40 $\pm$ 3.30	51.40	62.80 $\pm$ 4.20	51.23	39.80 $\pm$ 4.20	51.29	34.49 $\pm$ 4.35
	1024	52.23	46.40 $\pm$ 4.40	52.05	80.40 $\pm$ 3.40	52.32	62.80 $\pm$ 4.20	52.06	38.80 $\pm$ 4.10

Appendix

Table 10: HELMET task results at 32K context length. Each accuracy is reported as the mean

\pm

the half-width of its 95% confidence interval.

		S-NIAH-1		S-NIAH-2		S-NIAH-3		MK-NIAH-1		MK-NIAH-2
Method	$S$	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)
SDPA	—	100.00	100.00 $\pm$ 0.00	100.00	100.00 $\pm$ 0.00	100.00	100.00 $\pm$ 0.00	100.00	100.00 $\pm$ 0.00	100.00	98.80 $\pm$ 0.90
SANTA	128	50.97	100.00 $\pm$ 0.00	51.11	100.00 $\pm$ 0.00	51.06	99.80 $\pm$ 0.30	51.13	100.00 $\pm$ 0.00	51.25	98.80 $\pm$ 0.90
	256	51.56	100.00 $\pm$ 0.00	51.81	100.00 $\pm$ 0.00	51.72	99.80 $\pm$ 0.30	51.83	100.00 $\pm$ 0.00	52.08	99.00 $\pm$ 0.90
	512	52.52	100.00 $\pm$ 0.00	52.93	100.00 $\pm$ 0.00	52.76	100.00 $\pm$ 0.00	52.94	100.00 $\pm$ 0.00	53.41	98.20 $\pm$ 1.10
	1024	53.98	100.00 $\pm$ 0.00	54.62	100.00 $\pm$ 0.00	54.35	100.00 $\pm$ 0.00	54.62	100.00 $\pm$ 0.00	55.43	98.20 $\pm$ 1.10

Appendix

Table 11: Per-task accuracy and KV memory access on RULER (8K context) on Qwen2.5-7B-Instruct. KV is the theoretical GQA-aware KV memory access (%), and

S

denotes the nominal budget. Accuracy values are reported as mean

\pm

95% bootstrap confidence interval half-width.

		S-NIAH-1		S-NIAH-2		S-NIAH-3		MK-NIAH-1		MK-NIAH-2
Method	$S$	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)	KV (%)	Acc (%)
SDPA	—	100.00	100.00 $\pm$ 0.00	100.00	100.00 $\pm$ 0.00	100.00	98.60 $\pm$ 1.00	100.00	99.00 $\pm$ 0.90	100.00	93.40 $\pm$ 2.20
SANTA	128	50.28	100.00 $\pm$ 0.00	50.31	100.00 $\pm$ 0.00	50.30	99.20 $\pm$ 0.70	50.31	99.20 $\pm$ 0.70	50.37	93.20 $\pm$ 2.20
	256	50.47	100.00 $\pm$ 0.00	50.53	100.00 $\pm$ 0.00	50.51	98.80 $\pm$ 0.90	50.53	99.00 $\pm$ 0.90	50.64	92.80 $\pm$ 2.20
	512	50.81	100.00 $\pm$ 0.00	50.91	100.00 $\pm$ 0.00	50.87	98.80 $\pm$ 0.90	50.90	99.00 $\pm$ 0.90	51.12	93.60 $\pm$ 2.10
	1024	51.37	100.00 $\pm$ 0.00	51.53	100.00 $\pm$ 0.00	51.45	98.60 $\pm$ 1.00	51.51	99.00 $\pm$ 0.90	51.92	93.40 $\pm$ 2.10

Appendix

Table 12: Per-task accuracy and KV memory access on RULER (32K context) on Qwen2.5-7B-Instruct. KV is the theoretical GQA-aware KV memory access (%), and

S

denotes the nominal budget. Accuracy values are reported as mean

\pm

95% bootstrap confidence interval half-width.

Directly timed region	Median latency ( $\mu$ s)
Grouped representative scorer	20.32
Parallel per-query top- $K$	15.51
Fused exact merge + GQA union	16.16
Grouped selected attention	19.97
Split-softmax reducer	9.95
Complete five-launch SANTA ++	63.52

Appendix

Table 13: Profiler-off SANTA ++ timings for independently measured regions of the captured 32K workload. Entries are medians of round medians; the complete operator is measured directly.

Kernel	CTAs	NCU dur. ( $\mu$ s)	Occ. (%)	Elig.	DRAM (%)	Reg./thr.
SANTA ++ representative scorer	248	14.14	24.49	0.08	63.85	56
SANTA ++ top- $K$	868	10.78	84.51	2.40	10.61	30
SANTA ++ merge + union	4	10.94	16.62	0.38	1.59	44
SANTA ++ selected attention	128	14.30	12.88	0.20	19.46	144
SANTA ++ reducer	28	4.80	8.33	0.05	11.07	40
Flash split-KV forward	128	89.92	8.33	0.10	83.95	200

Appendix

Table 14: Nsight Compute measurements for the captured 32K, batch-one workload. Entries are per-metric medians of three independent captures of a fixed sampling realization; durations use kernel replay.

Method	$S$	KV (%)	Acc (norm) (%)	Acc (abs) (%)
SDPA	—	100.00	100.00	86.22 $\pm$ 4.56
SANTA ++ (Latent k-means)	128	56.28	80.80	69.67 $\pm$ 5.92
	256	70.09	91.56	78.94 $\pm$ 5.14
	512	86.18	98.07	84.56 $\pm$ 4.72
	1024	101.89	97.81	84.33 $\pm$ 5.00

Appendix

Table 15: Token sparsity in a shared-latent cache at 8K. DeepSeek-V2-Lite-Chat, with 50 matched prompts per task on niah_multivalue , fwe , and qa_1 . Absolute accuracy is the equal-weight task mean; normalized accuracy divides it by the SDPA mean, times 100. KV denotes logical shared-row access, including routing, under Equation 14 . Throughout these tables,

\pm

gives half the span of a 95% percentile-bootstrap interval using 10,000 resamples. Prompts are resampled within each task, retaining task strata for the aggregate. Interval endpoints need not be symmetric about the point estimate.

Method	$S$	KV (%)	Acc (norm) (%)	Acc (abs) (%)
SDPA	—	100.00	100.00	96.67 $\pm$ 2.67
SANTA ++ (Latent k-means)	128	47.57	84.82	82.00 $\pm$ 7.00
	256	62.25	94.48	91.33 $\pm$ 5.33
	512	80.34	100.00	96.67 $\pm$ 3.00
	1024	98.45	99.31	96.00 $\pm$ 3.00

Appendix

Table 16: RULER fwe at 8K. Accuracy on 50 matched prompts per method. Normalized accuracy divides the reported two-decimal task score by SDPA’s score on this task, times 100. Access and uncertainty follow Table 15 .

Method	$S$	KV (%)	Acc (norm) (%)	Acc (abs) (%)
SDPA	—	100.00	100.00	94.00 $\pm$ 4.00
SANTA ++ (Latent k-means)	128	57.88	77.66	73.00 $\pm$ 8.75
	256	71.89	93.09	87.50 $\pm$ 5.50
	512	87.67	98.94	93.00 $\pm$ 4.75
	1024	102.64	92.55	87.00 $\pm$ 7.75

Appendix

Table 17: RULER niah_multivalue at 8K. Accuracy on 50 matched prompts per method. Normalized accuracy divides the reported two-decimal task score by SDPA’s score on this task, times 100. Access and uncertainty follow Table 15 .

Method	$S$	KV (%)	Acc (norm) (%)	Acc (abs) (%)
SDPA	—	100.00	100.00	68.00 $\pm$ 13.00
SANTA ++ (Latent k-means)	128	58.41	79.41	54.00 $\pm$ 14.00
	256	72.81	85.29	58.00 $\pm$ 14.00
	512	88.50	94.12	64.00 $\pm$ 13.00
	1024	103.87	102.94	70.00 $\pm$ 13.00

Appendix

Table 18: RULER qa_1 at 8K. Accuracy on 50 matched prompts per method. Normalized accuracy divides the reported two-decimal task score by SDPA’s score on this task, times 100. Access and uncertainty follow Table 15 .

Explore similar work

Stochastic Sparse Attention for Memory-Bound Inference
May 3, 2026Kyle Lee, Corentin Delacour, Kevin Callahan-Coray +5Dynamic Sparse Attention Autoregressive Decoding
Department of Electrical and Computer Engineering, University of California, Santa Barbara, CA, USA · Department of Electrical and Computer Engineering, University of Michigan, Ann Arbor, MI, USA · Department of Electrical and Computer Engineering, Carnegie Mellon University, Pittsburgh, PA, USA
Keyless Attention: Value-Space Routing and Value-Only Caching for Efficient Transformers
Jun 20, 2026Xin GaoTransformer Architectures Reasoning Benchmark
Department of Mathematics and Statistics York University Toronto, M3J 1P3
Simplified Sparse Attention via Gist Tokens
Apr 22, 2026Yuzhen Mao, Michael Y. Li, Emily B. FoxDynamic Sparse Attention Efficient Long-Context Inference
Stanford University

May 3, 2026, cs.LGJ/K move · Enter open · S save

Stochastic Sparse Attention for Memory-Bound Inference

Kyle Lee, Corentin Delacour, Kevin Callahan-Coray, Kyle Jiang+4

Department of Electrical and Computer Engineering, University of California, Santa Barbara, CA, USA · Department of Electrical and Computer Engineering, University of Michigan, Ann Arbor, MI, USA · Department of Electrical and Computer Engineering, Carnegie Mellon University, Pittsburgh, PA, USA

Dynamic Sparse Attention Autoregressive Decoding Key-Value Cache Key-Value Cache Compression Gumbel-Softmax Relaxation Stochastic

Open paper AskPDF arXiv

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

Abstract

Explore similar work

Stochastic Sparse Attention for Memory-Bound Inference

Keyless Attention: Value-Space Routing and Value-Only Caching for Efficient Transformers

Simplified Sparse Attention via Gist Tokens