arXiv CSopen access

Stochastic Sparse Attention for Memory-Bound Inference

2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

cloud distributed-computing parallel-computing

distributed computing, parallel computing, cloud

Kyle Lee 1 Corentin Delacour 1 Kevin Callahan-Coray 1 Kyle Jiang 1 Can Yaras 2 Samet Oymak 2 Tathagata Srimani 3 Kerem Y. Camsari 1

arXiv:2605.01910v1 [cs.LG] 3 May 2026

Abstract

generating each new token repeatedly streams the Key– Value (KV) cache for all prior tokens. For example, for Llama-3.1-8B-Instruct at a 32k-token context, streaming bf16 keys and values is on the order of ∼128 MB per layer per generated token, and KV memory demands scale linearly with context length and batch size.

Autoregressive decoding becomes bandwidthlimited at long contexts, as generating each token requires reading all nk 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 ≪ nk 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 sampling to design variancereduced, GPU-friendly variants, demonstrating 1.5× decode-step attention kernel speedup over FlashInfer and FlashDecoding on an NVIDIA RTX 6000 Ada while matching baseline accuracy at 32k-token contexts. Finally, we propose Bernoulli qK T sampling as a complementary technique to sparsify the score stage, reducing key-feature access through stochastic ternary queries. Both methods are orthogonal to upstream techniques such as ternary quantization, low-rank projections, and KV-cache compression. Together, they point toward sparse, multiplier-free, and energy-efficient inference. We open-source our kernels at: https:// github.com/OPUSLab/SANTA.git

A broad literature mitigates this bottleneck from several angles. KV-cache quantization/compression reduces bytes per element (Liu et al., 2024; Hooper et al., 2024); cachemanagement policies reduce the number of retained tokens (Tang et al., 2024; Zhang et al., 2023); and structured sparsity or candidate selection reduces attention cost (Child et al., 2019; Zaheer et al., 2020; Beltagy et al., 2020; Wang et al., 2020; Kitaev et al., 2020; Chen et al., 2025). Architectural changes such as grouped-query attention (GQA) reduce KV footprint by sharing keys/values across heads (Ainslie et al., 2023), and decoding kernels improve IO locality for exact attention (Dao et al., 2023; Ye et al., 2025). However, even with optimized exact kernels, long-context decoding still incurs a persistent cost: each step touches essentially the full KV state. We study a complementary inference-time lever: reducing the amount of KV data accessed per decode step without permanently discarding cache contents. Our primary focus is the post-softmax value aggregation, where dense attention computes AV by multiplying attention weights by all nk value rows. We introduce Stochastic Additive No-mulT Attention (SANTA), which samples S ≪ nk indices from the post-softmax distribution and aggregates only those sampled rows of V via gather-and-add. This yields an unbiased estimator of the post-softmax value aggregation while reducing value-cache row reads per query from nk to S (and eliminating value-stage multiplications in favor of additive accumulation). We further propose variance-reduced variants (S2 ANTA) based on stratified and systematic sampling that bolster accuracy and admit GPU-friendly implementations. Since sparsifying V reads yields the clearest bandwidth benefit in decoding, our systems evaluation targets decode-step acceleration at long contexts.

1. Introduction Transformers (Vaswani et al., 2017) underpin modern language models (OpenAI, 2023; Touvron et al., 2023; Gemini Team et al., 2023). As deployments increasingly rely on long contexts, autoregressive decoding enters a regime where throughput is often limited by memory bandwidth: 1 Department of Electrical and Computer Engineering, University of California, Santa Barbara, CA, USA 2 Department of Electrical and Computer Engineering, University of Michigan, Ann Arbor, MI, USA 3 Department of Electrical and Computer Engineering, Carnegie Mellon University, Pittsburgh, PA, USA. Correspondence to: Kyle Lee <[email protected]>.

On modern GPUs, these changes translate into measurable kernel-level gains: our S 2 ANTA CUDA kernels achieve up to 1.5× decode-step attention-kernel speedup over

Preprint. May 5, 2026.

Stochastic Sparse Attention for Memory-Bound Inference

value matrix, and A ∈ R1×nk represents the attention scores as a probability distribution.

FlashInfer and FlashDecoding at 32k-token contexts while matching baseline accuracy on long-context prompts (Section 3).

To approximate AV , we treat A as a categorical distribution over keys and sample S indices i.i.d. from it (with replacement). For query vector q, we fetch the corresponding S rows of V and average them. This yields an estimate of the attention output using only row indexing and addition.

Finally, we propose Bernoulli qK T sampling as a complementary technique for the score stage: it represents query elements as Bernoulli variables to form sparse ternary queries in {−1, 0, +1}, yielding an unbiased estimator of qK T and inducing feature-sparse key access during decoding. In this work, our implemented kernels and measured speedups are driven primarily by value-stage sparsification; Bernoulli qK T is presented as a complementary mechanism that can be combined with SANTA to sparsify both branches.

As an illustrative example, consider number of keys nk = 3. We sample S = 3 one-hot attention vectors from the categorical distribution A: d = 1 1 0 0 + 1 0 0 + 0 0 1 V, (2) AV 3 which simplifies by distributivity of matrix multiplication over addition: d=1 1 0 0 V + 1 0 0 V + 0 0 1 V AV 3 1 = V1 + V 1 + V 3 . 3 (3)

Contributions. Concretely, we: (i) introduce SANTA and variance-reduced S 2 ANTA as unbiased estimators of the post-softmax value aggregation that reduce value-cache reads from nk to S; (ii) design GPU kernels that parallelize sampling and sparse accumulation, demonstrating up to 1.5× decode-step attention-kernel speedup over FlashInfer/FlashDecoding at 32k-token contexts while matching baseline accuracy; (iii) empirically characterize tunable accuracy–efficiency trade-offs on GSM8K, MMLU, and long-context benchmarks; and (iv) propose Bernoulli qK T sampling as a complementary score-stage estimator that can reduce key-feature access during decoding.

where each Vi ∈ R1×dk is a sampled row of V . The onehot attention vectors in Eqs. 2, 3 are only for illustrative purposes; they represent an indexing and gather operation, and are not materialized as vectors in memory. Finally, SANTA can be summarized as:

2. Stochastic Additive No-mulT Attention (SANTA)

S X d= 1 AV Vi ≈ AV, S s=1 s

Stochastic Additive No-mulT Attention (SANTA) Sample

Gather

KT ∈

(4)

where is is a sampled index and each Vis denotes a gathered row from V . Multiplications in the value stage have therefore been eliminated in favor of sampling, indexing and addition. If S is chosen as a power of 2 (and the model has a fixed point representation), normalization may be implemented as a bit-shift. Fig. 1 illustrates SANTA’s memory access benefits for autoregressive decoding: SANTA sparsely reads only S ≪ nk sampled rows of V and computes an additive average without multiplications after the softmax.

dk × nk

Sum

Attn distribution

Figure 1. SANTA requires S ≪ nk memory accesses to sampled rows of V , eliminating multiplications following the softmax operation. Normalization by S is a bit shift if S is a power of 2.

We provide a formal mathematical description of SANTA in Appendix A. SANTA is an unbiased estimator of the attention function whose variance scales as 1/S (Appendix A.1, A.2). A dimension-free high-probability concentration bound for SANTA is deferred to Appendix B.

In principle, SANTA can be employed for the prefill or decode stages. Since the benefit of sparse memory access becomes apparent in decode, we present SANTA as a decoding algorithm. In autoregressive decoding, we consider a single query vector q ∈ R1×dk . Scaled dot-product attention (SDPA) is expressed as a function of the key matrix and value matrix: qK T Attn(q, K, V ) = softmax √ V = AV, (1) dk

2.1. S 2 ANTA: stratified and systematic sampling S 2 ANTA reduces variance via stratified sampling: dividing the CDF into S equal probability mass intervals and sampling once per interval for better coverage of the attention distribution. We devise two variants: independent stratified sampling (S 2 ANTA-strat) and systematic sampling (S 2 ANTA-sys).

where K ∈ Rnk ×dk is the key matrix, V ∈ Rnk ×dk is the 2

Stochastic Sparse Attention for Memory-Bound Inference

The mechanics of S 2 ANTA-strat remain identical to default SANTA, except the S samples are drawn per stratum instead of iid from the full distribution. Systematic sampling (S 2 ANTA-sys) similarly splits the attention distribution CDF into S equal probability mass partitions. However, whereas independent stratified sampling draws S random offsets for each of S strata, systematic sampling draws a single offset U and applies the same offset to all strata, effectively drawing S samples with a single random number.

requiring ≫ S rows of V when the union across all queries is considered (see Appendices F, G for prefill analysis).

3. GPU kernels for variance-reduced SANTA Section 2 introduced SANTA and variance-reduced sampling rules (stratified/systematic), which specify what to sample to form an unbiased estimator of the value-stage aggregation AV . In this section we focus on how to realize these samplers efficiently at decode time on GPUs: we introduce two kernel strategies, S 2 ANTA-prop and S 2 ANTAflash, which differ in how they allocate samples across tiles (global proportional allocation vs. speculative per-tile sampling with deferred normalization).

Construction. Fix a query q. Let Aq = (pq1 , . . . , pqnk ) be the attention weights and Fq the associated CDF on [0, 1). Partition [0, 1) into S equal intervals Im := [m/S, (m + 1)/S),

m = 0, . . . , S − 1.

3.1. Parallelizing stochastic attention on GPUs

Define two schemes:

Efficient implementation on GPUs presents a unique systems challenge: parallelizing the sampling control flow. Standard attention (SDPA) is amenable to “split-KV” strategies (e.g., FlashDecoding (Dao et al., 2023)) because the reduction of values is associative: partial sums can be computed locally and merged later. Sampling, however, creates a sequential dependency: determining which rows of V to access generally requires a global Cumulative Distribution Function (CDF). A naive implementation would require a sequential pass over the global probability distribution to draw sample indices, leaving the GPU compute units idle and preventing parallel access to the value matrix.

• Independent stratified (S 2 ANTA-strat). Draw Tm ∼ Unif(Im ) independently for m = 0, . . . , S − ind dq 1, set Jm := Fq−1 (Tm ), and output AV := PS−1 1 m=0 VJm . S • Systematic (S 2 ANTA-sys). Draw a single U ∼ Unif([0, 1/S)) and take thresholds Tm := U + sys d m/S. With Jm := Fq−1 (Tm ), output AV := q PS−1 1 m=0 VJm . S S 2 ANTA-strat and S 2 ANTA-sys are unbiased estimators of attention (see Appendix A.4). Only S 2 ANTA-strat has a theoretical guarantee of variance reduction compared to default SANTA (see Appendix A.5). Despite lacking a variance reduction guarantee, we validate that S 2 ANTA-sys exhibits comparable performance to S 2 ANTA-strat on reasoning benchmarks and long-context benchmarks. S 2 ANTA-sys demonstrates similar variance reduction to S 2 ANTA-strat in practice (see Appendix D). Systematic sampling with S 2 ANTA-sys is appealing because it draws samples using one random number instead of S random numbers, which may be easier to implement in hardware.

To unlock parallel V -access, we must determine the number of samples per tile before or independently of the full global reduction. We propose two strategies to break this dependency: 1. Proportional Sample Allocation (S 2 ANTA-prop): We pre-calculate the probability mass of each tile to deterministically assign sample budgets. This requires a global barrier but ensures exact load balancing. 2. Speculative Local Sampling (S 2 ANTA-flash): Tiles sample independently with uniform budgets, speculating that all tiles are equally important, followed by a post-hoc re-weighting. This removes the barrier but introduces sample waste.

2.2. Theoretical operations and memory accesses By virtue of SANTA’s sparsity, in autoregressive decoding, SANTA demonstrates a S/nk reduction in post-softmax additions and reads to rows of V , where S is the sample budget and nk is the number of keys. SANTA removes postsoftmax multiplications in favor of additive accumulation (see Appendix E for analysis of decode-step operations and memory accesses).

We implement both strategies and evaluate their decodestep kernel latency and long-context accuracy at 32k-token prompts in Section 3.4.

In prefill, while similar arguments hold for SANTA’s reduced additions and removed multiplies, there is no longer sparse memory access to the value matrix V . Each of nq = nk queries may each sample S rows of V , therefore

S 2 ANTA-prop (Figure 2) breaks the sequential sampling dependency by resolving the global partition function Z in a lightweight initial pass.

3.2. S 2 ANTA-prop: exact budget allocation via global sync

Stochastic Sparse Attention for Memory-Bound Inference (1) Tiled score matrix computation

(3) Per-tile systematic sampling, gather

To eliminate the latency of the global synchronization barrier, S 2 ANTA-flash (Figure 3) adopts a speculative approach structurally similar to FlashDecoding. Instead of pre-calculating the global budget, each tile is allocated a fixed, uniform budget Stile ≈ S/T , where T is the number of tiles. Each tile samples locally as if it contained the entire probability mass. A second reduction kernel then computes the true global partition function Z and down-weights the partial outputs of tiles that had low probability mass.

(2) Allocate tile samples (global sync) Attn distribution

samples

This method maximizes parallelism by removing the global barrier but suffers from sample waste: tiles with low probability mass still consume compute and bandwidth to draw samples that are effectively down-weighted during the final merge. In our 32k-token evaluation (Section 3.4), matching SDPA accuracy requires substantially larger budgets for S 2 ANTA-flash; the accuracy-matching operating point is S=2048 for S 2 ANTA-flash versus S=128 for S 2 ANTAprop (Table 1). See Appendix P for full pseudocode.

Figure 2. S 2 ANTA-prop kernel overview.

The algorithm proceeds in three stages. The key dimension nk is partitioned into tiles of length Btile . 1. Pass 1 (score stash & tile stats): We compute attention scores exactly. The exponentiated scores (or a normalized equivalent) are written to global memory. Crucially, since the scores are scalars (1×nk ) while the values are vectors (nk × dk ), this stash consumes negligible bandwidth (1/dk ) compared to a full value fetch. We also compute and write local partition functions Ztile .

3.4. Kernel results at long contexts (32k) We evaluate decode-step attention kernel latency together with end-task accuracy in the long-context decoding regime. Within this section, SDPA is used for prefill and only the decode step uses S 2 ANTA. Kernel timings are reported for a single decoding step on an NVIDIA RTX 6000 Ada GPU (batch size 1) using Llama-3.1-8B-Instruct (Meta AI, 2024) tensor shapes; we compare against optimized FlashInfer and FlashDecoding backends. The measurement protocol and tensor layouts are provided in Appendix Q.

2. Global budget allocation: A separate kernel sums the tile statistics to find the global Z. It then assigns a specific sample count to each tile: Stile ∝ S ·(Ztile /Z). 3. Pass 2 (parallel gather): Using the stashed scores and the assigned Stile , each tile systematically samples indices and accumulates rows from V . Tiles with low probability mass are assigned Stile = 0, allowing them to skip the expensive V -read entirely.

To make latency comparisons meaningful, we select one operating point per kernel family as the smallest sample budget S that matches SDPA accuracy on 32k-token long-context tasks. This yields S=128 for S 2 ANTA-prop and S=2048 for S 2 ANTA-flash. These are exactly the configurations that deliver ≈ 1.5× decode-step attention-kernel speedup in Fig. 4 while matching SDPA within confidence intervals in Table 1.

Comprehensive pseudocode is in Appendix O. 3.3. S 2 ANTA-flash: Speculative Sampling with Deferred Normalization (1) Per-tile systematic sampling, equal sample budgets

We next verify that these speedup points preserve longcontext accuracy. We evaluate on tasks from RULER (Hsieh et al., 2024) using 32,768-token prompts: frequent-word extraction (FWE), needle-in-a-haystack (NIAH), single-hop QA (QA1 , derived from SQuAD (Rajpurkar et al., 2018)), and multi-hop QA (QA2 , derived from HotpotQA (Yang et al., 2018)). Results are shown in Table 1.

Partial tile outputs

No global sync

At these operating points, the speedups are driven primarily by reduced memory traffic from sparsifying reads of the value matrix V . For S 2 ANTA-prop, S=128 corresponds to extremely sparse value access at 32k contexts (e.g., ≲ 1.56% V access when accounting for GQA). S 2 ANTAflash attains comparable speedup, but requires a larger total

samples per tile

(2) Reweight partial tile outputs by probability mass:

Figure 3. S 2 ANTA-flash kernel overview.

Stochastic Sparse Attention for Memory-Bound Inference

(a)

0.20

S ANTA-prop (S = 128) S2ANTA-prop (S = 256)

S2ANTA-prop (S = 1024)

Table 1. Accuracy of Llama 8B on 4 long-context tasks with a 32,768-token prompt. SDPA is used in prefill. Tasks: frequent word extraction (FWE), needle-in-a-haystack (NIAH), single-hop QA (QA1 ), multi-hop QA (QA2 ). S: sample budget. Accuracy shows 95% bootstrap confidence intervals.

FlashInfer FlashDecoding

1.50 ×

Kernel time (ms)

0.25

0.15

Kernel 2

0.10 0.05 4096 8192

16384

Sequence length nk

32768

FWE

NIAH

QA1

QA2

S ANTA-prop S 2 ANTA-prop S 2 ANTA-prop

128 95.40±1.10 98.25±0.62 64.40±4.50 60.20±4.20 256 95.47±1.10 98.50±0.57 63.40±4.10 60.60±4.00 1024 95.40±1.07 98.40±0.60 64.20±4.00 59.60±4.40

S 2 ANTA-flash S 2 ANTA-flash S 2 ANTA-flash

256 66.20±2.40 88.95±1.43 63.00±4.20 57.20±4.30 1024 91.60±1.43 98.10±0.62 62.20±4.10 61.00±4.30 2048 94.13±1.13 98.25±0.62 64.60±4.20 60.00±4.40

SDPA (baseline)

— 95.60±1.00 98.35±0.62 64.00±4.10 58.80±4.10

(b)

0.20

S2ANTA-flash (S = 256)

S2ANTA-flash (S = 1024) S ANTA-flash (S = 2048) 2

FlashInfer FlashDecoding

>90% reductions in value-stage memory access at 32k contexts. Second, replacing dense multiply-accumulate with gather-and-add eliminates multiplication operations entirely. While our current GPU kernels primarily exploit the first benefit, hardware architectures optimized for sparse, addercentric computation could realize both. These characteristics make S 2 ANTA a natural fit for emerging AI accelerators where memory bandwidth and multiplier energy are primary constraints.

1.51 ×

Kernel time (ms)

0.25

0.15 0.10 0.05 4096 8192

16384

Sequence length nk

32768

Figure 4. S 2 ANTA kernel latency on Llama 8B tensor shapes for 1 decoding step. (a) S 2 ANTA-prop demonstrates a 1.50× speedup relative to FlashInfer at a 32k-token context. (b) S 2 ANTA-flash similarly exhibits a 1.51× speedup. Both operating points (S = 128 for prop, S = 2048 for flash) correspond to configurations that recover baseline accuracy (Table 1).

4. General reasoning & accuracy verification In this section, we consider SANTA, S 2 ANTA-strat (stratified sampling), and S 2 ANTA-sys (systematic sampling) as mathematical constructions (as in Section 2), validating the methods’ accuracy on GSM8K, MMLU, and longcontext benchmarks using PyTorch reference implementations. While demonstrated GPU kernel speedups are for autoregressive decoding, evaluations in this section use SANTA in both prefill and decoding. Since SANTA in principle reduces additions and V memory access by S/nk and removes post-softmax multiplies, future additioncentric hardware may enable gains in energy efficiency in the compute-heavy prefill phase; for instance, on modern CMOS, a FP16 multiply typically costs 1.1 pJ vs. 0.4 pJ for an add (Horowitz, 2014).

sample budget because uniform per-tile sampling expends samples on low-mass tiles before the deferred renormalization step (“sample waste”). These timings are kernel-level improvements; realized endto-end speedups are subject to Amdahl’s Law and are largest in the memory-bound long-context regime (e.g., nk ≳ 8k) where KV-cache streaming dominates decode time. While S 2 ANTA is multiplier-free in principle, current GPUs are highly optimized for dense fused multiply-accumulate instructions; we therefore view multiplier-free arithmetic as an additional energy-oriented benefit that may become more pronounced on future hardware.

4.1. GSM8K We consider Llama-3.1-8B-Instruct (“Llama 8B”) (Meta AI, 2024) on the GSM8K dataset as a mathematical reasoning benchmark (Cobbe et al., 2021; OpenAI, 2023). Table 2 provides both accuracy and the average number of tokens (prompt + answer). Because the compute and memory access costs of SANTA are contextualized relative to the sequence length nk , these sequence lengths justify the regimes where S ≪ nk and SANTA is, in principle, cheaper than SDPA in attention’s value stage.

Additional benchmark results for 8k-token contexts are available in Appendix R. 3.5. Energy considerations Beyond latency, S 2 ANTA’s design offers two complementary paths to energy reduction. First, sparse memory access reduces value-cache traffic from nk to S rows per query, at accuracy-matching operating points, this corresponds to

Default SANTA exhibits respectable performance, but 5

Stochastic Sparse Attention for Memory-Bound Inference Table 2. GSM8K accuracy and average context length (prompt + answer) for Llama-3.1-8B-Instruct. S: SANTA sample budget. Accuracy shows 95% bootstrap confidence intervals. S 2 ANTA-sys

S 2 ANTA-strat

SANTA

Acc. (%)

Tok.

Acc. (%)

Tok.

Acc. (%)

Tok.

2 4 8 16 32 64 128 256

1.11±0.34 1.67±0.40 2.10±0.45 44.63±1.58 68.59±1.42 76.42±1.34 77.33±1.34 77.56±1.34

1071 569 340 349 339 341 342 343

1.01±0.32 1.54±0.38 1.74±0.40 39.12±1.50 67.00±1.48 74.43±1.31 75.64±1.33 78.17±1.28

1071 611 325 352 343 343 341 343

1.26±0.34 1.57±0.37 1.44±0.38 5.51±0.72 38.26±1.43 63.63±1.49 70.23±1.42 75.61±1.42

1075 825 207 350 348 346 341 342

78.06±1.33

344

SDPA (baseline)

which is significantly shorter than the average sequence length nk = 417. S 2 ANTA variants outperform default SANTA across S budgets of 64–256, recovering baseline SDPA accuracy (within 1%) with multiplier-free arithmetic and sparse memory access to rows of V since S ≪ nk . Extended results with DeepSeek 7B and top-k comparisons are provided in Appendix H. 4.3. Long context benchmarks Table 4. Accuracy of Llama 8B on 4 long-context tasks with an 8k-token prompt. Tasks: frequent word extraction (FWE), needle-in-a-haystack (NIAH), single-hop QA (QA1 ), multi-hop QA (QA2 ). S: SANTA sample budget. Accuracy shows 95% bootstrap confidence intervals.

variance-reduced S 2 ANTA variants demonstrate superior performance across every sample budget S. S 2 ANTA implementations approach the accuracy of full SDPA (within 1%) while S remains shorter than the sequence length nk , translating to memory-access savings and FLOP energy savings. In deployment settings where LLM users do not require maximum model performance, the sample budget S can be tuned to save computation cost. For instance, S 2 ANTA-sys achieves 76.42% accuracy at S = 64, which is about 19% of the average sequence length of nk = 341 tokens. For a modest accuracy trade-off compared to SDPA’s 78.06%, S 2 ANTA-sys in principle costs 19% of the additions, 19% of the V -matrix accesses, and no multiplies.

Kernel 2

4.2. MMLU Table 3. MMLU accuracy and average context length (prompt + answer) for Llama-3.1-8B-Instruct. S: SANTA sample budget. Accuracy shows 95% bootstrap confidence intervals. S 2 ANTA-strat

Acc. (%)

Tok.

Acc. (%)

Tok.

Acc. (%)

Tok.

2 4 8 16 32 64 128 256

24.52±1.22 24.49±1.18 24.47±1.23 34.14±1.39 45.48±1.45 48.70±1.39 50.60±1.37 51.12±1.46

1141 774 292 353 403 418 421 421

24.28±1.21 24.65±1.26 25.63±1.23 32.46±1.32 43.78±1.50 49.25±1.48 49.92±1.47 50.90±1.49

1141 807 286 350 398 415 421 422

24.89±1.27 25.13±1.20 25.06±1.23 26.24±1.26 33.81±1.34 41.11±1.40 47.46±1.45 49.75±1.51

1144 961 284 313 354 397 412 417

49.86±1.47

424

SDPA (baseline)

NIAH

QA1

QA2

S ANTA-sys S 2 ANTA-sys S 2 ANTA-sys

64 96.27±0.97 94.05±1.07 71.80±3.90 67.20±4.20 128 97.80±0.70 94.15±1.02 68.20±4.10 66.80±4.00 256 97.87±0.70 95.00±0.95 71.20±3.90 69.40±4.10

S 2 ANTA-strat S 2 ANTA-strat S 2 ANTA-strat

64 95.93±0.94 94.35±1.15 71.40±4.00 65.40±4.20 128 97.73±0.80 94.30±1.10 70.80±4.20 67.40±4.20 256 98.33±0.60 94.55±1.05 71.20±4.00 67.00±4.20

SANTA SANTA SANTA

64 92.53±1.27 93.95±1.30 64.60±4.20 63.40±4.20 128 94.53±1.13 91.45±1.38 67.80±4.10 63.40±4.30 256 96.20±0.93 93.15±1.18 68.00±4.10 66.40±4.10

SDPA (baseline)

— 98.53±0.60 94.55±1.00 71.80±3.90 68.60±4.20

For full clarity, when Llama 8B’s grouped-query attention (Ainslie et al., 2023) with 4 grouped queries is fully accounted for, then memory access to unique rows of V can in the worst case reach 12.5%, which remains extremely sparse. Extended results with top-k comparisons are provided in Appendix H.

SANTA

FWE

We benchmark SANTA on long-context tasks borrowed from RULER (Hsieh et al., 2024) with a prompt length of 8192 tokens. The theoretical benefit of S 2 ANTA’s sparse computation is most stark in this long-context regime. With a budget of S = 256 in an 8192-token sequence, S 2 ANTA uses just S/nk = 3.125% of the value-stage additions and memory accesses compared to full attention (following Table 8). It achieves this while using no multiplications and recovering baseline SDPA accuracy (within ±1%) across every task.

We provide additional DeepSeek-R1-Distill-Qwen7B (DeepSeek-AI, 2025) (DeepSeek 7B) results and comparisons to sparse top-k attention in Appendix H.

S 2 ANTA-sys

5. Bernoulli qK T for sparse key access SANTA reduces memory access for the value matrix by sampling from the softmax distribution; however, it leaves the score computation (qK T ) intact. As illustrated in Fig. 5, we demonstrate that Bernoulli sampling of the query naturally zeroes out query-irrelevant features, thereby effectively reducing memory access for the key matrix.

We further evaluate SANTA on the MMLU benchmark (Hendrycks et al., 2021; Center for AI Safety, 2024), which tests factual recall and general reasoning. MMLU benchmarking observes similar trends as GSM8K, where SANTA virtually recovers baseline accuracy at S = 256,

We propose representing each query element as an independent Bernoulli variable, allowing the corresponding key 6

Stochastic Sparse Attention for Memory-Bound Inference

Bernoulli Gather

Sample

Table 5. Performance of Bernoulli qK T (stratified) on GSM8K (BitNet 2B) with the test set evaluated three times, and confidence intervals obtained by bootstrapping.

Sum

Normalized query

Figure 5. Bernoulli qK T represents elements of the query vector as Bernoulli probabilities, sparsifying along the feature dimension dk with multiplier-free arithmetic.

feature to be pruned when the sampling probability is low. While this method is applicable to the prefill phase, it induces sparse key access specifically during decoding; therefore, we focus on a single query q. First, the query vector is normalized as q ← q/norm, where norm ≥ maxi |qi |. This ensures that each element qi ∈ [−1, 1] can be interpreted as a probability. Next, we estimate the product p = qK T as: norm X (n) T qb K , B n=1

Acc. (%) (95% CI)

Avg. Tok.

SDPA

–

65.7 ± 1.5

306.1

100

Bernoulli qK T

1 2 4 8 16

1.9 ± 0.43 1478.6 48.7 ± 1.4 345.9 64.5 ± 1.5 309.4 65.1 ± 1.5 316.9 65.8 ± 1.5 317.7

25.6 46.3 67.5 82.6 91.1

65.8 88.0 97.9 99.8 100

2 4 8 16

48.6 ± 1.5 63.8 ± 1.5 65.6 ± 1.6 66.2 ± 1.5

57.9 84.7 97.1 99.7

Mean Group Query Bernoulli qK T

318.6 311.4 311.4 310.2

effect: with B = 4 samples, K T access jumps from 67.5% per head to 97.9% for the group. To preserve sparsity in this scenario, we propose using a representative query m for the entire group, defined as the mean of the queries’ absolute values. We then construct the estimator m b for the mean query m via Monte Carlo sampling (detailed in Appendix C.2), which stochastically prunes the low-amplitude mean features.

pb =

K T Access (%) Head Group

Method

(5)

where B is the sample count and qb(n) is a sparse ternary vector constructed via Bernoulli sampling (detailed in Appendix C). The linearity of expectation ensures this estimator is unbiased. Because qb(n) contains only values in {−1, 0, +1}, pb can be computed without multiplications by gathering and accumulating selected rows of K T , followed by normalizing bit shifts for fixed-point representations.

With B = 4 samples, m b contains approximately 15% of zero elements, which allows fetching only 85% of key features for the group (Table 5). The product p = qK T for each query is finally estimated as:

Table 5 shows GSM8K results for Bernoulli qK T applied to BitNet 2B (Ma et al., 2025) for both prefill and decoding. Using B = 4 stratified samples nearly recovers the accuracy of standard SDPA for the GSM8K benchmark, although the approximated scores typically have a 30% error (Appendix C.1). This suggests that the qK T product within the highly quantized BitNet 2B architecture can tolerate numerical errors while maintaining strong overall performance.

pb = m b⊙

q T K , m

(6)

where ⊙ denotes element-wise multiplication. In this scenario, calculating pb requires normalization by m to compensate for the bias introduced by the mean group query. As shown in Table 5, the resulting estimator achieves similar accuracy to the single-query configuration while saving 15% of key memory accesses. We employ the mean group query estimator (Eq. 6) in all the following experiments.

Crucially, Bernoulli qK T enables sparse access to cached keys, a benefit that is uniquely effective during decoding (the prefill phase requires nearly all features). This is shown in Table 5: with B = 4, the model needs only 67.5% of key features per head (vs. ∼ 100% for prefill), while maintaining accuracy within 2% of the SDPA baseline.

5.2. Long context benchmarks We evaluate our method on long-context tasks using Llama 8B, applying Bernoulli qK T strictly during decoding since prefill requires full feature access. As detailed in Table 6, using B = 8 samples reduces key feature access to 72.8% for Llama 8B, with minimal accuracy loss (< 2% on NIAH, QA1 , and QA2 ). Furthermore, this experiment highlights that key memory access patterns vary significantly across models. In contrast to Llama 8B, BitNet 2B requires accessing 97.1% of features for the same budget (B = 8, Table 5), a difference we attribute to distinct query distributions.

5.1. Mean group query While Bernoulli qK T effectively prunes irrelevant features for individual heads, Grouped Query Attention (GQA) (Ainslie et al., 2023) complicates this by sharing keys and values across multiple heads (e.g., groups of 4 in BitNet). Consequently, the set of accessed features becomes the union of those required by all queries in the group, significantly increasing memory traffic. Table 5 illustrates this 7

Stochastic Sparse Attention for Memory-Bound Inference Table 6. Accuracy of Llama 8B on 4 long-context tasks with a 4,096-token prompt. Tasks: frequent word extraction (FWE), needle-in-a-haystack (NIAH), single-hop QA (QA1 ), multi-hop QA (QA2 ). B: Bernoulli qK T (mean group query) sample budget. Accuracy shows 95% bootstrap confidence intervals. Kernel (qK T )

KT Access (%)

FWE

NIAH

QA1

QA2

Bernoulli Bernoulli Bernoulli

4 8 16

47.3 72.8 92.1

70.4±2.0 89.4±1.4 92.7±1.2

87.2±1.9 93.8±1.4 94.5±1.2

72.0±3.8 73.6±3.9 74.8±3.5

66.2±4.1 70.2±3.9 71.6±4.2

SDPA

–

94.0±1.1

95.0±1.2

73.8±3.8

70.8±4.0

evicting or pruning tokens (Tang et al., 2024; Zhang et al., 2023), introducing approximation by permanently discarding state. SANTA and Bernoulli qK T are orthogonal to these approaches: they reduce the number of elements accessed over whatever cache is retained, and therefore can stack with compression, GQA, and eviction. Approximating attention and sampling vs. truncation. Structured sparsity and low-rank approximations reduce attention cost (often in prefill) (Child et al., 2019; Zaheer et al., 2020; Beltagy et al., 2020; Wang et al., 2020). In decoding, candidate selection and query-pruning methods reduce the cost of identifying relevant tokens or approximating qK T (Kitaev et al., 2020; Chen et al., 2025; Ribar et al., 2024). This is where our two methods differ in scope: SANTA operates in the value stage and is orthogonal to score-stage approximations, while Bernoulli qK T is an alternative score-stage estimator that induces feature-sparse key access via stochastic ternary queries. Sampling-based attention approximations have also been explored (Kim & Ko, 2022). A direct deterministic baseline is top-k attention (Gupta et al., 2021), which reduces value reads by truncation but is biased. In contrast, SANTA provides an unbiased estimator of the post-softmax value aggregation with controllable variance, and we emphasize variance reduction (stratified/systematic sampling) together with GPU kernels that realize measured speedups in the memory-bound decoding regime.

5.3. Combining Bernoulli qK T and SANTA Combining SANTA with Bernoulli qK T enables sparse key– value accesses during decoding while replacing the majority of multiplications with additions. Table 7 shows the accuracy for the GSM8K benchmark using BitNet 2B and B = 4 Bernoulli samples for the qK T computation, as the number of SANTA samples S is varied for the value stage. Using B = 4 for queries and S = 64 for attention weights almost recovers the SDPA baseline with less than 2% accuracy drop. Critically, this combined approach reduces key access by 15.3% and value access by 89.2%. Table 7. Performance of Bernoulli qK T (mean group query) + SANTA on GSM8K (BitNet 2B) with the test set evaluated three times (CI computed by bootstrapping). Acc. (%)

Avg.

(K,V) Access (%)

(95% CI)

Tok.

Group

SDPA

65.7 ± 1.5

306.1

(100, 100)

B = 4, S = 8 B = 4, S = 16 B = 4, S = 32 B = 4, S = 64 B = 4, S = 128

55.0 ± 1.6 60.7 ± 1.5 62.3 ± 1.5 63.8 ± 1.5 64.1 ± 1.5

319.2 313.3 308.4 304.6 303.2

(84.7, 2.4) (84.7, 3.7) (84.7, 6.2) (84.7, 10.8) (84.7, 14.3)

Config

7. Conclusion Long-context autoregressive decoding is frequently constrained by memory bandwidth due to repeated KV-cache access. We introduced SANTA, a stochastic method that avoids streaming the full value cache by sampling S ≪ nk value rows from the post-softmax distribution and aggregating them with sparse gather-and-add. SANTA is unbiased for the post-softmax value aggregation, admits practical variance reduction via S 2 ANTA stratified/systematic sampling, and provides a tunable accuracy–latency knob in the sample budget S. On modern GPUs, our observed speedups primarily come from reducing value-cache bandwidth: our S 2 ANTA kernels achieve up to a 1.5× decodestep attention-kernel speedup over FlashInfer/FlashDecoding at 32k-token contexts while retaining baseline performance at modest budgets.

6. Related work Exact attention kernels and memory-bound decoding. FlashAttention-style kernels optimize IO locality and utilization for exact attention, including decoding-focused variants (Dao et al., 2023; Ye et al., 2025). However, exact decoding still streams essentially the full KV state each step at long contexts. Our approach is complementary: rather than further optimizing a full KV scan, we reduce the amount of KV data accessed at runtime by sparsifying value-cache reads (and, for Bernoulli qK T , key-feature reads).

We also proposed Bernoulli qK T sampling as a complementary, forward-looking mechanism for the score stage that reduces key-feature access during decoding via stochastic ternary queries, which can be combined with SANTA to sparsify both branches of attention. Beyond reduced memory bandwidth, the predominantly additive structure of both estimators may also be attractive for future adder-centric hardware and/or prefill, which we leave to future work.

Reducing KV footprint and retention. KV quantization/compression reduces bytes per KV element (Liu et al., 2024; Hooper et al., 2024), and weight-only quantization targets model parameters (Frantar et al., 2023; Dettmers et al., 2023). GQA reduces KV size via head sharing (Ainslie et al., 2023). Cache-management methods bound KV growth by 8

Stochastic Sparse Attention for Memory-Bound Inference

h i d q − µq ∥2 = 1 tr(Σq ). Equivalently, E ∥AV 2 S

Impact Statement This paper presents work whose goal is to advance the field of machine learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.

Proof. Unbiasedness was shown in Proposition A.2. For d q − µq = 1 PS Xs with the variance, write AV s=1 S E[X s ] = 0 and Cov(X s ) = Σq . The Xs are i.i.d., hence PS PS Cov S1 s=1 Xs = S12 s=1 Cov(Xs ) = S1 Σq . Taking the trace gives the mean-squared error identity.

A. Formal description of SANTA For query q, the SANTA estimator can be written as S X dq = 1 AV Vi , S s=1 q,s

iq,s ∼ Categorical(Aq ).

A.3. S2 ANTA estimator Proposition A.4 (Unbiasedness of S 2 ANTA). Let µq := ind sys P d d j pqj Vj = (AV )q . Then E[AV q ] = E[AV q ] = µq .

(7)

The proof is shown in Section A.4. Remark A.5 (Systematic vs. independent). S 2 ANTA-strat is unbiased and admits a variance-dominance theorem, guaranteeing lower variance than that of i.i.d. sampling with default SANTA (Section A.5). S 2 ANTA-sys is also unbiased, but does not provide tractable variance guarantees; empirically it performs comparably.

Definition A.1 (SANTA estimator).√ Let Rnq ×dk , Q ∈ nk ×dk T nq ×nk K∈R , A = softmax QK / dk ∈ R , and V ∈ Rnk ×dk . For each query index r ∈ {1, . . . , nq }, sami.i.d. dr = ple ir,1 , . . . , ir,S ∼ Categorical(Ar ) and define AV P S 1 d ∈ Rnq ×dk . Vi . Stacking all queries gives AV S

s=1

r,s

A.1. Unbiasedness of SANTA

A.4. Unbiasedness of S 2 ANTA

d ] = AV. Proposition A.2 (Unbiasedness). E[AV

Proof of Proposition A.4. For S 2 ANTA-strat,

Proof of Proposition A.2. Fix any q. Since Pr(iq,s = j) = Aqj ,

ind

nk S nk X 1 XX d E[AV q ] = Aqj Vj = Aqj Vj = (AV )q . S s=1 j=1 j=1

Linearity of expectation completes the result.

pqj Vj

A.2. Variance of SANTA

= µq .

Proposition A.3 (Variance scaling of SANTA). The covariance of the SANTA estimator scales as 1/S.

For S 2 ANTA-sys, observe that marginally each Tm is uniform on Im (though dependent across m). By linearity of expectation, the sum of expectations is unchanged by this dependence, so the same integral argument applies; see Cochran (1977, Ch. 8) and Owen (2013).

Throughout this subsection, fix a query index q ∈ {1, . . . , nq }. Let Aq ∈ ∆nk −1 denote the attention weights Pnk over keys, write pqj ≜ Aqj , and let µq ≜ j=1 pqj Vj = (AV )q be the exact attention output for that query. Default

A.5. Variance reduction for S 2 ANTA-strat

i.i.d.

SANTA draws iq,1 , . . . , iq,S ∼ Categorical(Aq ) and returns S X dq = 1 AV Vi ∈ Rdk . S s=1 q,s

TheoremPA.6 (Variance reduction for S 2 ANTA-strat). Let ⊤ Σq := and define withinj pqj (Vj − µq )(Vj − µq ) stratum covariances Σ(m) := Cov VFq−1 (T ) T ∼ Unif(Im ) . q

Define centered summands Xs ≜ Viq,s − µq and the (queryPnk specific) covariance Σq ≜ Cov(Vi ) = j=1 pqj (Vj − ⊤ µq )(Vj − µq ) .

Then ind

We now aim to show that SANTA’s covariance scales as 1/S (Proposition A.3): d q ] = µq , E[AV

S−1 i 1 X h E VFq−1 (Tm ) S m=0 Z 1X = 1{Fq (j−1) ≤ t < Fq (j)} Vj dt

dq ] = E[AV

dq ) = Cov(AV

dq Cov AV

S−1 1 X (m) 1 d q ), Σ ≤Loewner Σq = Cov(AV S 2 m=0 q S

with strict improvement unless all stratum means coincide. See Lohr (2010, Ch. 4) and Cochran (1977, Ch. 5).

1 Σq . S 9

Stochastic Sparse Attention for Memory-Bound Inference ind

C.1. Bernoulli qK T error

dq ) = Proof. Independence across strata gives Cov(AV P 1 m Cov VFq−1 (T ) | T ∈ Im . By the law of total coS2

Bernoulli qKT L2 norm error

variance with T ∼ Unif(0, 1) and partition {Im }, h i Σq = E Cov VFq−1 (T ) | T ∈ Im h i + Cov E VFq−1 (T ) | T ∈ Im h i ≥Loewner E Cov VFq−1 (T ) | T ∈ Im .

and averaging then dividing by S yields the claim. Strictness holds unless the stratum means are all equal.

𝟏

𝐁 𝟏 𝐁

Corollary A.7 i (MSE ordering). i For any q, h h ind 2 d d E ∥AV − µq ∥ ≤ E ∥AV q − µq ∥2 = 1 tr(Σq ). q

B. Vector Bernstein tail for SANTA Theorem B.1 (Vector Bernstein tail for SANTA). Assume a uniform almost-sure bound ∥Vj − µq ∥2 ≤ Vmax for all j (hence ∥Xs ∥2 ≤ Vmax ). Let vq ≜ E ∥Vi − µq ∥22 = tr(Σq ). Then for all t > 0, S t2 d Pr AV q − µq 2 ≥ t ≤ 2 exp − . 2 (vq + Vmax t/3) Equivalently, with probability at least 1 − δ, r 2 vq log(2/δ) 2Vmax log(2/δ) d q − µq AV ≤ + . 2 S 3S Proof sketch. This is a Hilbert-space (vector) Bernstein inequality: apply Bernstein to the mean of independent, meanzero, L-bounded vectors with variance proxy vq ; see, e.g., Boucheron et al. (2013, Thm. 2.10; see also Cor. 2.11) and Pinelis (1994).

To maintain sparse key access when keys share multiple queries, as in GQA, we define a common representative query for the group, such as the mean of the absolute values of the queries:

We first normalize the query vector as q ← q/norm, where norm ≥ maxi |qi | so that each element qi ∈ [−1, 1] is interpreted as a probability. Then, we introduce the unbiased query estimator qb where each element qbi is an independent ternary Bernoulli variable:

1 X |qg |, |G|

(10)

g∈G

where G denotes the query group. We then estimate the mean query using Monte Carlo sampling, after normalizing by norm ≥ maxi mi :

(8)

qbi = bi × sign(qi ) ∈ {−1, 0, +1}.

bi ∼ Bernoulli(mi /norm)

Sampling from qb provides an unbiased estimator for the product p = qK T :

(11)

m b =

norm X (n) T pb = qb K , B n=1

Fig. 6 shows the scaling of the L2 norm error with the number of samples B, defined as e = ||b p − p||2 /||p||2 . In similar fashion as SANTA, stratified Bernoulli sampling reduces the variance of each independent element pbi . The variance scales as O(1/B 2 ), resulting in an overall L2 norm scaling as O(1/B). For this numerical experiment, which uses a head dimension dk = 128 and a context length nk = 1024, using only B = 4 samples induces a mean L2 norm error of approximately 60% for standard Bernoulli sampling, compared to a much lower 30% for stratified sampling. C.2. qK T for Grouped Query Attention

C. Estimating qK T via Bernoulli sampling

bi ∼ Bernoulli(|qi |),

Figure 6. Mean L2 norm error for Bernoulli qK T . The score error is calculated during decoding with dk = 128, a context length nk = 1024, √ and averaged over 100 instances q ∼ N (0, 1), K ∼ N (0, 1)/ dk , V ∼ N (0, 1).

(9)

norm X (n) b B n=1

(12)

where b(n) are Bernoulli query vectors with independent elements bi . After fetching the corresponding K T rows corresponding to non-zero elements of m, b the product p =

(n)

where qb is a Bernoulli query vector, and B is the number of samples. 10

Stochastic Sparse Attention for Memory-Bound Inference

qK T for each query is estimated as: pb = m b⊙

q T K , m

Table 8 displays the theoretical additions, multiplications, and memory accesses incurred by SANTA in the value stage of attention, where the attention scores A multiply V . The score computation qK T is similar for both SANTA and SDPA, and is therefore omitted in this analysis (see Appendix G). We count individual scalar elements which are added together, multiplied together, or accessed by memory. Architecture-dependent effects such as multiply-accumulate (MAC) instructions, caching effects, or operation parallelism are not accounted for; arguments are fundamental and hardware-agnostic.

(13)

where ⊙ denotes element-wise multiplication. In this mean group query scenario, calculating pb requires multiplications by the original query q and normalization by m to compensate for the bias introduced by the mean query approximation.

D. Empirical variance

By virtue of SANTA’s sparsity, SANTA demonstrates a S/nk reduction in post-softmax additions and reads to rows of V , where S is the sample budget and nk is the number of keys (sequence length). A key distinction is that SANTA costs fewer additions than SDPA; SANTA does not seek to replace the multiplications in the AV computation, but rather, leverages sparse sampling to decrease the total number of additions while eliminating multiplications. So long as S ≪ nk , the advantage of attention sparsity holds true. Sampling overhead is not included; it is highly dependent on implementation, but drawing S ≪ nk samples should be lightweight relative to memory accesses, multiplies, and adds (indeed, this proves to be the case with GPU kernel latency demonstrations in Section 3).

Figure 7. Empirical variance of SANTA on single-hop QA prompts (8k tokens).

To further validate our results, we empirically measure the variance of SANTA, S 2 ANTA-strat, and S 2 ANTA-sys. We sweep the sample budget S on 8k-token single-hop QA prompts, averaged across model layers and 30 prompts (normalized per head). S 2 ANTA variants exhibit significantly lower variance than SANTA. Though systematic sampling admits no tractable theoretical variance guarantees, it empirically demonstrates similar variance to independent stratified sampling (and requires 1 random number per query instead of S random numbers). On a log-log scale, slopes for SANTA, S 2 ANTA-strat, and S 2 ANTA-sys are −1.1, −1.27, and −1.29, respectively, matching the 1/S scaling prescribed by theory (see Appendix M for details regarding the measurement protocol).

Whereas SDPA incurs nk dk multiplications from computing dot products, SANTA only costs dk normalizing multiplies by 1/S. dk is a constant model parameter and does not increase with the sequence length nk . Furthermore, if S is chosen as a power of 2 and the model has a fixed point representation, this normalization may be implemented as a bit-shift. In prefill, while similar arguments hold for SANTA’s reduced additions and multiplies, there is no longer sparse memory access to the value matrix V . Each of nq queries may each sample S rows of V , therefore requiring ≫ S rows of V when the union across all queries is considered (see Appendices F, G).

F. Unique V -row coverage for SANTA prefill

E. Theoretical operations and memory accesses in autoregressive decoding

Table 9. Unique V rows read (union over all query positions in a causal prefill pass), reported as a percentage of all nk rows in the value matrix. Each entry is the mean over 10 trials and all attention heads, with Gaussian inputs Q, K, V ∼ N (0, 1) (FP16), H=32 heads, and head dimension dk =128.

Table 8. Per-query, per-head value-stage compute and memory for dense scaled dot-product attention (SDPA) and SANTA. nk is the number of keys, dk is the head dimension, and S is the SANTA sample budget per query. Value reads counts V elements accessed per query; Writes counts output elements. Metric

SDPA

SANTA

Adds Mults / Divs Value reads Writes

nk dk nk dk nk dk dk

Sdk dk Sdk dk

Unique V rows read (% of nk ) nk =1024 2048 4096 8192 16384

1 4 8 16

49.90 79.95 88.92 94.07

50.04 80.00 88.83 94.11

49.96 79.96 88.86 94.10

49.95 79.96 88.86 94.09

49.95 79.95 88.85 94.08

Stochastic Sparse Attention for Memory-Bound Inference

Table 9 shows the percentage of unique V rows read by SANTA during prefill passes of differing sequence lengths nk . Standard Gaussian distributed inputs Q, K, V are passed into SANTA with dk = 128 and H = 32 heads for multi-headed attention. Even in the long-context, S ≪ nk regime, any appreciable number of samples S results in SANTA sampling most rows of V when the union across all queries is considered. For S = 16, selected sequence lengths from nk = 1k − 16k show SANTA reading almost all (> 94%) rows of V , because each of nq = nk queries stochastically attend to S sampled keys. Therefore, while SANTA features sparse V access during autoregressive decoding, SANTA does not have apparent benefit from sparse memory access in prefill. Arguments concerning SANTA’s reduced addition and multiplication operations are still valid in prefill (Table 10).

tic estimators like SANTA provide an unbiased estimate of full attention and retain robustness in such cases. For both top-k and SANTA, we define a notion of compute budget, where k is the number of keys selected by top-k for each query of each attention head, while S is a sample budget representing the number of samples that SANTA draws for each query of each attention head. By default, we employ uniform S across model layers, though ablations suggest that S can be optimized per layer (Appendix I, J). We omit softmax computation, top-k sorting overhead, and sampling from this analysis, as they are lightweight relative to memory accesses and the V matrix multiply, and are highly implementation-dependent. For example, many implementations use partial sorting (Samaga B L et al., 2024; Xie et al., 2025). In both FLOPs and memory reads, SDPA scales as nq nk , reflecting the cost of full quadratic attention computation. Both top-k and SANTA reduce this to nq k and nq S, respectively, assuming fixed budgets with k, S ≪ nk . This shifts the overall complexity of the value stage from quadratic to linear in sequence length, given a fixed budget k or S.

G. Theoretical operations and memory accesses in prefill Table 10. Per-head FLOPs and memory for scaled dot-product attention (SDPA), top-k, and SANTA. nq : number of queries, nk : number of keys, dk : head dimension, k: top-k budget, S: SANTA sample budget per query. “Reads” count logical value/key elements accessed; “Writes” count output or score elements. Softmax, top-k selection, and sampling overheads are omitted. Stage / Variant

Adds

Mults / Divs

Reads

Writes

Score stage (QK ) SDPA / top-k / SANTA

nq nk dk

(nq + nk )dk

nq nk

Value stage SDPA Top-k SANTA

nq nk dk nq kdk nq Sdk

nq nk dk nq kdk nq dk

nq nk dk nq kdk nq Sdk

nq d k nq d k nq d k

Top-k still incurs multiplication costs while SANTA eliminates multiplies. Top-k performs nq kdk multiplications (per head), while SANTA requires only nq dk divisions for normalization (if S is chosen as a power of 2, even these few divisions become simple bit-shifts).

H. Additional DeepSeek and top-k results Table 11. GSM8K accuracy and average context length (prompt + answer) for DeepSeek-R1-Distill-Qwen-7B. k: number of keys in top-k, S: SANTA sample budget. Accuracy shows 95% bootstrap confidence intervals.

While SANTA’s apparent memory sparsity benefits are most viable in autoregressive decoding, for full clarity, this section provides generalized analysis for prefill. We consider the computational cost of SANTA, SDPA, and top-k attention, focusing on FLOPs and memory traffic in the value stage. Table 10 summarizes the theoretical costs per attention head. For the score stage (QK T ), the “Reads” entry (nq + nk )dk counts each query and key vector once per head, assuming they are loaded once from DRAM and then reused in on-chip memory during the matmul. In contrast, for the value stage we report nq nk dk , nq kdk , and nq Sdk for dense SDPA, top-k, and SANTA, respectively, which count the per-query accesses rather than unique loads of elements from memory.

Top-k

SANTA

k|S

Acc. (%)

Tok.

Acc. (%)

Tok.

2 4 8 16 32 64 128 256

0.23±0.15 23.88±1.31 72.25±1.35 82.51±1.21 86.23±1.02 86.76±1.03 87.54±1.02 88.20±0.97

4049 3298 2332 1860 1694 1624 1626 1600

0.00±0.00 0.08±0.09 52.46±1.57 79.00±1.28 83.19±1.23 84.99±1.09 85.42±1.00 86.93±1.01

4217 4208 2799 1778 1667 1587 1594 1627

88.75±1.00

1532

SDPA (baseline)

Tables 11, 12, 13, and 14 feature additional GSM8K and MMLU results with comparisons using DeepSeek-R1Distill-Qwen-7B and top-k attention. Table 15 exhibits long-context benchmark results identical to Table 4, except with additional top-k comparisons.

While top-k attention is memory-efficient and generally performs well, it introduces bias and can degrade significantly when the attention distribution is heavy-tailed. For example, recent work (Chen et al., 2025) shows that tasks such as common word extraction suffer when relevant information is not concentrated in the top-k entries. In contrast, stochas-

Comparing SANTA and top-k requires care, as they have different computational profiles. We present results for sample 12

Stochastic Sparse Attention for Memory-Bound Inference Table 12. GSM8K accuracy and average context length (prompt + answer) for Llama-3.1-8B-Instruct. Top-k k|S

Acc. (%)

SANTA Tok.

Acc. (%)

Tok.

S 2 ANTA-strat

S 2 ANTA-sys

Acc. (%)

Tok.

Table 14. MMLU accuracy and average context length (prompt + answer) for Llama-3.1-8B-Instruct. Top-k k|S

Tok.

2 1.21±0.32 430 1.26±0.34 1075 1.01±0.32 1071 1.11±0.34 1071 4 6.27±0.72 583 1.57±0.37 825 1.54±0.38 611 1.67±0.40 569 8 41.32±1.50 549 1.44±0.38 207 1.74±0.40 325 2.10±0.45 340 16 62.88±1.57 430 5.51±0.72 350 39.12±1.50 352 44.63±1.58 349 32 72.50±1.34 384 38.26±1.43 348 67.00±1.48 343 68.59±1.42 339 64 76.12±1.32 358 63.63±1.49 346 74.43±1.31 343 76.42±1.34 341 128 77.13±1.30 349 70.23±1.42 341 75.64±1.33 341 77.33±1.34 342 256 78.19±1.33 340 75.61±1.42 342 78.17±1.28 343 77.56±1.34 343

Top-k

2 4 8 16 32 64 128 256

Acc. (%) 24.97±1.24 38.93±1.32 55.08±1.47 59.07±1.36 60.77±1.40 62.03±1.38 61.68±1.49 62.05±1.46

SDPA (baseline)

Acc. (%)

Tok.

S 2 ANTA-sys

Acc. (%)

Tok.

49.86±1.47 424

SDPA (baseline)

Table 13. MMLU accuracy and average context length (prompt + answer) for DeepSeek-R1-Distill-Qwen-7B. k: number of keys in top-k, S: SANTA sample budget. Accuracy shows 95% bootstrap confidence intervals.

k|S

Tok.

S 2 ANTA-strat

2 23.56±1.20 588 24.89±1.27 1144 24.28±1.21 1141 24.52±1.22 1141 4 28.37±1.27 359 25.13±1.20 961 24.65±1.26 807 24.49±1.18 774 8 39.54±1.43 532 25.06±1.23 284 25.63±1.23 286 24.47±1.23 292 16 45.26±1.37 484 26.24±1.26 313 32.46±1.32 350 34.14±1.39 353 32 49.49±1.44 448 33.81±1.34 354 43.78±1.50 398 45.48±1.45 403 64 48.33±1.37 433 41.11±1.40 397 49.25±1.48 415 48.70±1.39 418 128 49.47±1.44 418 47.46±1.45 412 49.92±1.47 421 50.60±1.37 421 256 49.75±1.43 421 49.75±1.51 417 50.90±1.49 422 51.12±1.46 421

78.06±1.33 344

SDPA (baseline)

Acc. (%)

SANTA

Table 15. Accuracy of Llama 8B on 4 long context tasks with an 8k-token prompt. Tasks: frequent word extraction (FWE), needle-in-a-haystack (NIAH), single-hop QA (QA1 ), multi-hop QA (QA2 ). k: number of keys for top-k. S: SANTA sample budget. Accuracy shows 95% bootstrap CIs.

SANTA Tok. 3598 2236 1659 1301 1173 1103 1061 1041

Acc. (%)

Kernel

Tok.

4.33±0.58 23.10±1.20 44.55±1.47 58.81±1.42 61.05±1.45 60.59±1.43 62.38±1.32 62.27±1.39

4178 3666 2080 1506 1228 1104 1075 1074

63.16±1.45

1020

budget S equal to the top-k budget k not as an iso-compute benchmark, but to demonstrate that SANTA achieves comparable or superior accuracy while using a fundamentally cheaper set of operations. At S = k: top-k costs nq kdk multiplications + nq kdk additions, whereas SANTA costs 0 multiplications + nq Sdk additions (Table 10).

k|S

FWE

NIAH

QA1

QA2

S ANTA-sys S 2 ANTA-sys S 2 ANTA-sys

64 96.27±0.97 94.05±1.07 71.80±3.90 67.20±4.20 128 97.80±0.70 94.15±1.02 68.20±4.10 66.80±4.00 256 97.87±0.70 95.00±0.95 71.20±3.90 69.40±4.10

S 2 ANTA-strat S 2 ANTA-strat S 2 ANTA-strat

64 95.93±0.94 94.35±1.15 71.40±4.00 65.40±4.20 128 97.73±0.80 94.30±1.10 70.80±4.20 67.40±4.20 256 98.33±0.60 94.55±1.05 71.20±4.00 67.00±4.20

SANTA SANTA SANTA

64 92.53±1.27 93.95±1.30 64.60±4.20 63.40±4.20 128 94.53±1.13 91.45±1.38 67.80±4.10 63.40±4.30 256 96.20±0.93 93.15±1.18 68.00±4.10 66.40±4.10

SDPA (baseline)

— 98.53±0.60 94.55±1.00 71.80±3.90 68.60±4.20

top-k top-k top-k

64 93.07±1.20 92.55±1.10 65.60±4.10 60.80±4.10 128 95.40±0.96 93.60±1.15 69.20±3.90 62.20±4.50 256 97.27±0.80 94.55±1.05 70.60±3.90 64.20±4.10

For benchmarking, we apply SANTA and top-k to both prefill and autoregressive decoding. However, in practice, these sparse attention implementations may see greater benefit during decoding, when only a single query sparsely selects rows of V .

A 32-bit FP multiply costs 3.7 pJ versus 0.9 pJ for addition (Horowitz, 2014). For equal budgets k = S, the valuestage energy consumption for the two methods are approximately: k × (3.7 + 0.9) = 4.6k pJ per element for top-k, vs. S × 0.9 = 0.9S pJ per element for SANTA. Relative to topk, SANTA yields an ≈ 5× energy reduction for value-stage FLOPs. We acknowledge that these are approximate energy estimates; we aim to provide hardware-agnostic comparisons, as modern GPU hardware is optimized for matmuls and contiguous memory access, while SANTA involves adds and sparse memory access.

SANTA variants demonstrate strong performance relative to top-k on GSM8K and MMLU benchmarks, despite SANTA requiring no multiplications and a similar number of additions and memory accesses for k = S. However, the most remarkable result is shown in Table 15, where S 2 ANTA variants outperform top-k across virtually every long-context task for all k = S. S 2 ANTA’s lead over top-k is most apparent for multi-hop question answering (QA2 ), where S 2 ANTA variants lead in performance by ≈ 5% in several cases. It is plausible that long context tasks — multi-hop QA in particular — have information that is not concentrated in the top-k keys, and therefore SANTA’s unbiased estimate of attention outperforms top-k, which drops information and cannot attend to keys outside of the k largest attention scores.

Regarding memory accesses: SANTA’s accesses to elements of the V matrix scale as nq Sdk , which in the worst case matches top-k’s nq kdk memory accesses at k = S. However, since SANTA samples with replacement, the number of unique keys is ≤ S. In practice, as we quantify in Appendix L, the number of unique keys sampled is < S, which may offer significant caching advantages. 13

Stochastic Sparse Attention for Memory-Bound Inference

I. Ablation study: one-hot attention samples

baseline 1-hot schedule 5

100

1-hot layer index

baseline 1-hot schedule 10

1-hot layer index

50 baseline 1-hot schedule 5

100

(b) Deepseek 7B - MMLU

Accuracy (%)

100

(a) Deepseek 7B - GSM

Accuracy (%)

100

Algorithm 1 Layer-wise sample-budget optimization via REINFORCE Require: initial logits θ ∈ R28 (θi = 0), total budget N =224, learning-rate α=0.02 1: while training do 2: Workers (in parallel): 3: for e = 1, . . . , E do {E=10 episodes / worker run} P 4: Compute pi ← exp(θi ) j exp(θj ) 5: Sample schedule S ∼ Multinomial(N, p) 6: Evaluate 16 GSM8K questions using schedule s; obtain reward r ∈ [0, 1] 7: Append (s, r) to shared folder 8: end for

1-hot layer index

(d) Llama 8B - MMLU

50 baseline 1-hot schedule 0

1-hot layer index

Figure 8. Ablations with one-hot stochastic attention. Baseline scores employ a fixed sampling budget of 16 and 256 for DeepSeek 7B and Llama 8B models, respectively. The horizontal axis shows the index of the model layer which is reduced to stochastic hard attention.

9: 10: 11: 12:

Aggregator: Collect all new (s(k) , r(k) )P tuples {k = 1 . . . K} 1 (k) Compute baseline b ← K r k 1 X (k) (r − b) s(k) − N p 13: Gradient: g ← K k 14: Update logits θ ← θ + α g 15: Write back updated θ 16: end while

This section probes a fundamental question: how much does each attention layer matter? We find that one-hot stochastic attention, a limiting case of SANTA with S = 1, acts as a surprisingly sharp diagnostic. It reveals which transformer layers are robust to severe approximation and which are not, offering a practical lens into where attention precision matters most. We perform these ablations by setting S = 1 in a single layer while keeping all other layers at fixed sample budgets: S = 16 for DeepSeek 7B (28 layers) and S = 256 for Llama 8B (32 layers). In this setting, each query attends to only one randomly sampled key, making that layer a hard attention bottleneck.

We now consider a global sample budget shared across all 28 transformer layers of DeepSeek 7B, allowing each layer to have a distinct sample count. We optimize this allocation using a REINFORCE-style algorithm (Williams, 1992), detailed in Algorithm 1.

Models are evaluated on GSM8K and MMLU with the same protocol as Tables 2 and 3. The effect of this localized bottleneck on accuracy is summarized in Fig. 8. Some layers exhibit almost no degradation when ablated in this way, while others collapse entirely, most notably middle layers in DeepSeek 7B and the first layer in Llama 8B.

We train the policy on the GSM8K training set. The policy logits θ ∈ R28 define a softmax distribution over layers. At each iteration, we sample schedules using a multinomial draw over this distribution, with a total budget of 224 samples (i.e., 8 × 28). Each episode evaluates 16 questions; the reward is the fraction answered correctly. We intentionally starve the baseline with S = 8 samples per layer to create room for the learned schedule to improve.

In DeepSeek 7B, layers 12–18 are highly sensitive to approximation, while early and final layers are much more tolerant. In contrast, Llama 8B exhibits extreme sensitivity in just the first layer. These patterns suggest that attention importance is neither uniform nor trivially architectural.

Fig. 9 shows the final learned schedule after 501 iterations. The allocation aligns strikingly with the ablation trends from Fig. 8: more samples are assigned to the first and middle layers, where one-hot attention was most detrimental. That REINFORCE independently rediscovers this structure from only reward signals confirms that these patterns are not artifacts of the ablation method but instead are intrinsic to the models.

We observe similar qualitative trends regardless of the baseline sample budget; the S = 16 and S = 256 settings are arbitrary. The early-layer sensitivity in Llama 8B is consistent with prior findings on high-entropy attention in shallow layers (Clark et al., 2019), and with prior work that avoids approximating early layers (Tang et al., 2024; Sun et al., 2024). The middle-layer fragility in DeepSeek 7B remains an open and intriguing observation.

This experiment is not tuned for peak accuracy, but demonstrates that adaptive per-layer budgets meaningfully outperform fixed ones. Please see Appendix K for implementation details of distributed reinforcement learning using a Kubernetes compute cluster.

J. Layer-wise sample budget through RL 14

Stochastic Sparse Attention for Memory-Bound Inference

(a) Learned sample schedule

read (schedule, reward)

write

Shared filesystem

)

rd wa

sa m

Accuracy (%)

Aggregator

, nd le pe du ap che (s

# Samples

GPU 1

Layer (b) Reinforcement learning curve

Aggregator loop. Whenever at least one new result file appears, the aggregator moves all current files into a private folder for the next iteration, computes the REINFORCE update and writes a new policy snapshot. It never pauses or signals the workers; it simply processes whatever has arrived. In practice, workers are unlikely to finish evaluation simultaneously, thus nearly every policy update comes from one worker output (10 × 16 questions). If two land together, the aggregator just uses a double-sized batch, which lowers the gradient’s variance.

50 40 30 100

200

300

Iteration

400

500

Asynchrony and fault-tolerance.

Figure 9. Learned schedule with a total budget of 224 samples across all 28 transformer blocks (DeepSeek 7B). (b) RL iterations over time.

• Workers can be terminated at any time; partial output vanishes harmlessly.

Table 16. RL sample-schedule performance. Accuracy shows 95% bootstrap confidence intervals. Schedule

Task

Accuracy (%)

Baseline (8 samples) Baseline (8 samples) Learned schedule Learned schedule

GSM8K MMLU GSM8K MMLU

52.46 ± 1.47 44.55 ± 1.47 66.06 ± 1.47 48.86 ± 1.53

... GPU 50

Figure 10. Kubernetes deployment schematic.

20 0

GPU 2

• Workers always proceed with the newest policy that has reached disk; no locking is needed. • If the aggregator restarts it scans the history directory, resumes from the last completed iteration, and continues. Scaling and resources. The worker Deployment’s replica count can be scaled up or down at any time during learning. For the 501 iterations in Fig. 9 we ran 50 GPU workers on a mix of RTX 3090, L4, and A10 cards (24 GB Ampere generation cards) for roughly 20 hours, totaling ≈ 1000 GPU-hours. The aggregator requests just one CPU core and a few hundred MiB of RAM.

K. Distributed reinforcement learning on Kubernetes Topology. All pods mount a single read–write volume that serves as the rendezvous point.

Terminology. We keep the names aggregator and worker to highlight the simple “parameter-server” flavor of the design, but these roles align exactly with the common learner/actor split in distributed reinforcement learning.

• Aggregator pod. Holds the current policy and records a running history of baselines and gradients. • Worker pods. Operate in a tight loop: (i) read the latest policy; (ii) sample a 28-dimensional schedule; (iii) evaluate ten batches of 16 questions; (iv) append a small JSON-lines result file to a designated “inbox” directory on the volume.

L. Unique key-count Since all versions of SANTA sample keys with replacement, the number of unique keys may be less than the sample budget S. In Table 17, we empirically measure the average 15

Stochastic Sparse Attention for Memory-Bound Inference Table 17. Unique-key count vs. S (8192 tokens, single-query). 2

(2) Variance trace of the estimator. We report the trace d, of the covariance (expected squared ℓ2 error) of AV d ) := E ∥AV d − µ∥2 = tr Cov(AV d) , VarTrace(AV 2

SANTA

S ANTA-strat

S ANTA-sys

32 64 128 256

11.0 16.7 25.9 38.3

12.7 21.2 35.8 60.3

13.0 21.7 36.4 61.6

specialized to each sampling scheme: Pnk • Multinomial (closed form). Let Σ := j=1 pj (Vj − P 2 2 µ)(Vj − µ)⊤ so that tr(Σ) = p ∥V ∥ j 2 − ∥µ∥2 . j j Then

number of unique keys sampled by the last query position for 8192-token single-hop QA prompts (normalized per head and averaged across model layers). This is indicative of the key diversity observed in autoregressive decoding. For a given S, S 2 ANTA variants achieve higher key diversity than default SANTA. Furthermore, the number of unique keys tends to be ≪ S, implying that fewer than S unique rows of V require memory accesses. See Appendix M for details about the measurement protocol.

d ) = 1 tr(Σ). VarTracemulti (AV S • Independent equal-mass stratified (closed form). Partition [0, 1) into S equal strata; for each stratum m, form the within-stratum discrete distribution induced by p and let Σm be the corresponding value covariance. Because draws across strata are independent,

M. Implementation & measurement protocol for empirical variance and unique keys sampled

S−1 X d) = 1 VarTracestrat (AV tr(Σm ). S 2 m=0

Setting. For a single layer/head at the last prefill position q ⋆ , let post-softmax attention over nk keys be pj ≥ 0 Pthe nk pj = 1, and let Vj ∈ Rdk denote the correspondwith j=1 ing value vectors. The dense-attention mean is µ =

nk X

• Systematic (replicate-based estimate). There is no general closed form that holds uniformly for systematic sampling, so we estimate variance by repeated random starts. With R ≥ 2 independent offsets, we compute (r) d for each replicate, then use AV

pj Vj ∈ Rdk .

j=1

d) = \ sys (AV VarTrace

R 1 X d (r) AV −V R − 1 r=1

2 , 2

SANTA estimator. Given a per-head sampling budget S, all SANTA variants estimate µ with

V =

X d = 1 VJ , AV S s=1 s

1 X d (r) AV . R r=1

This estimator is applied per head at q ⋆ . M.2. How we aggregate and report

where the key indices {Js } are drawn by one of three schemes: (i) multinomial; (ii) independent equal-mass stratified; or (iii) systematic (random-start, fixed-stride). These are the same variants used in the main paper; the code paths that implement them also record the per-head statistics described below at q ⋆ .

For each prompt, we compute U and the variance-trace scalar per head at q ⋆ , then average over heads within a layer, average over layers, and finally average over prompts to obtain the quantities reported in the variance plots: U (S, variant)

d ). and T (S, variant) = VarTrace(AV

This yields the curves summarized in the main text as a function of the budget S and the SANTA variant.

M.1. What we measure (per head at q ⋆ ) (1) Unique keys. The number of distinct keys touched by the sampler,

Protocol summary. All measurements use evaluationmode forward passes, prefill-only, and record statistics at the last token per head; no additional dense attention pass is performed to form empirical errors. For multinomial and stratified we use the exact formulas above; for systematic we use the replicate-based estimator.

U := {J1 , . . . , JS } , upper-bounded by S (can be < S if high-probability keys repeat). This is a proxy for how many memory locations each head actually reads. 16

Stochastic Sparse Attention for Memory-Bound Inference

N. Prompting, Parsing, and Grading

Llama 8B + GSM8K

In Section 4.1, we evaluate the GSM8K test split (1319 prompts) 3 times. Evaluation code uses straightforward answer parsing and borrows grading code from PRM800K (Lightman et al., 2024). For all runs, temperature = 0.6, top-p = 0.95, and repetition penalty = 1.1. We use left-padded batches of 8–16 prompts. In Section 4.2, we average results over three runs on the MMLU validation set (1531 questions), using identical decoding parameters as GSM8K.

Aspect

Procedure

Prompt

Meta chat template. System: “You are a precise mathematician. Explain your reasoning step by step, then output ONLY the final number on a new line.” User: the GSM8K question. Inspect the model reply in order: (i) the last non-empty line of the response; (ii) if that fails, the final 25 whitespaceseparated tokens. From the selected chunk, take the first integer-looking token (commas allowed, optional leading minus). Numeric equality judged by the MITlicensed PRM800K grader; every question is counted, and blank or unparsable answers score 0.

Answer extraction

In Section 4.3, long-context benchmark results are averaged over 500 prompts and greedy decoding is employed with no repetition penalty applied.

Grading

DeepSeek 7B + GSM8K Aspect

Procedure

Prompt

Question, blank line, then Please reason step by step, and put your final numeric answer within \boxed{}. Perform a brace-balanced search for the last occurrence of \boxed{...} (robust to nested braces). If none is found, fall back to the final ∼200 characters of the model output. Numeric equality judged by the MITlicensed PRM800K grader; every prompt is counted, and blank or unparsable predictions score 0.

Answer extraction

Grading

Table 20. Prompting and grading for Llama 8B on GSM8K.

Llama 8B + MMLU

Table 18. Prompting and grading for DeepSeek 7B on GSM8K. Aspect

Procedure

Prompt

Meta chat template. System: “You are a knowledgeable and concise subjectmatter expert. Work through the problem step by step. Finally, on a new line, output ONLY the single capital letter (A, B, C, or D) that corresponds to the correct choice.” User: the question followed by the four labeled choices A–D. Take the substring after the last newline and scan it left-to-right for the first capital A–D (case-insensitive). If none is found, inspect the final 10 whitespace-separated tokens of the whole response, scanning them right-to-left for A–D. Predicted letter (upper-cased) is compared with the gold key. Every question is counted; if no valid letter is found or it does not match the gold answer, the item scores 0.

DeepSeek 7B + MMLU Aspect

Procedure

Prompt

Question, four labeled choices, then Answer briefly. Put your final answer as a single letter inside \boxed{}. Brace-balanced search for the last occurrence of \boxed{...}; if none is found, inspect the final ∼300 characters of the output. Scan this region backwards and take the first occurrence of A–D (case-insensitive), then convert it to uppercase. The predicted letter is compared with the gold key (case-insensitive). If no valid letter is found—or the output is blank— the item scores 0; otherwise it scores 1 when the letters match.

Answer extraction

Grading

Answer extraction

Grading

Table 21. Prompting and grading for Llama 8B on MMLU.

Table 19. Prompting and grading for DeepSeek 7B on MMLU.

Llama 8B + RULER prompts 17

Stochastic Sparse Attention for Memory-Bound Inference Aspect

Procedure

Prompt

Prompts are generated for the following RULER sub-tasks with a length of 4096– 32768 tokens: fwe, niah multivalue, qa 1, and qa 2. Greedy decoding is used. The model generates a maximum of 64 new tokens. fwe: partial credit = (# of gold words present in the model output) / 3 (caseinsensitive set match). niah multivalue: take the first four unique integers from the model output; partial credit = (# of these that appear in the gold set) / 4. qa 1 and qa 2: correct if the gold answer string appears as a case-insensitive substring of the model output.

Grading

Algorithm 3 S2 ANTA-prop, Kernel 1: pass-1 tile statistics (and optional u-stash) Require: Q ∈ RH×dk ; K ∈ Rnk ×Hkv ×dk ; tile length Btile . Ensure: m ∈ RH×T , ℓ ∈ RH×T , and optional u ∈ RH×nk . 1: T ← ⌈nk /Btile ⌉; G ← H/Hkv . 2: for each KV head k ∈ {0, . . . , Hkv −1} and tile t ∈ {0, . . . , T −1} in parallel do 3: Tt ← [tBtile , min((t+1)Btile , nk )). 4: for each grouped head g ∈ {0, . . . , G−1} √ do 5: h ← kG + g; sh,n ← Q⊤ dk for all n ∈ Tt . h Kn,k / P 6: mh,t ← maxn∈Tt sh,n ; ℓh,t ← n∈Tt exp(sh,n − mh,t ). 7: (Optional) uh,n ← exp(sh,n − mh,t ) for all n ∈ Tt . 8: end for 9: end for

Table 22. Prompting and grading for Llama 8B on RULER prompts.

Algorithm 4 S2 ANTA-prop, Kernel 2: proportional per-tile budgets (largest remainder) Require: Tile stats mh,t , ℓh,t ; samples/head S; seed. P Ensure: Integer budgets Sh,t with t Sh,t = S, plus invδh,t and offsets a0,h,t . 1: for each head h ∈ {0, . . . , H−1} in parallel do 2: m⋆h ← maxt mh,t . P −1 3: Wh,t ← exp(mh,t − m⋆h ) ℓh,t ; Zh ← Tt=0 Wh,t . 4: qt ← S · Wh,t /Zh ; P Sh,t ← ⌊qt ⌋ for all t. 5: Give remaining S − t Sh,t samples to tiles with largest fractional part of qt (largest remainder). 6: invδh,t ← (Sh,t > 0) ? Sh,t /ℓh,t : 0; a0,h,t ∼ Unif(0, 1). 7: end for

Licenses. • Models. deepseek-ai/DeepSeek-R1-Distill-Qwen-7B (MIT) and meta-llama/Meta-Llama-3.1-8B-Instruct (Llama 3.1 license) on Hugging Face. • Datasets. GSM8K openai/gsm8k (MIT) and MMLU cais/mmlu (MIT) on Hugging Face. RULER prompts (Apache). • Grader. PRM800K numeric grader (MIT).

O. S 2 ANTA-prop pseudocode Algorithm 2 provides a high-level overview of S 2 ANTAprop. Algorithms 3, 4, and 5 show breakdowns of the kernels that perform score computation, sample budget allocation, and the V gather, respectively. Algorithm 2 S2 ANTA-prop (overview, decode/GQA): proportional tile budgets + systematic sparse-V Require: Q ∈ RH×dk ; K, V ∈ Rnk ×Hkv ×dk (GQA); tile length Btile ; samples/head S. Ensure: O ∈ RH×dk approximating softmax(QK T )V . 1: T ← ⌈nk /Btile ⌉, G ← H/Hkv , k(h) ← ⌊h/G⌋, Tt = [tBtile , min((t+1)Btile , nk )). 2: Pass P 1: compute mh,t = maxn∈Tt sh,n and ℓh,t = optionally stash uh,n = n∈Tt exp(sh,n − mh,t ); √ exp(sh,n − mh,t ), where sh,n = Q⊤ h Kn,k(h) / dk . ⋆ ⋆ 3: Budgets: mh = maxt mh,t ; W Ph,t = exp(mh,t − mh ) ℓh,t ; allocatePintegers {Sh,t }t with t Sh,t = S and Sh,t ≈ S · Wh,t / t Wh,t (largest remainder); set invδh,t = Sh,t /ℓh,t and a0,h,t ∼ U [0, 1). 4: Pass 2: systematicPcounts ch,n from xh,n = invδh,t uh,n , then Oh ← Oh + n ch,n Vn,k(h) . 5: return O/S.

Stochastic Sparse Attention for Memory-Bound Inference

Algorithm 5 S2 ANTA-prop, Kernel 3: systematic resampling and sparse-V accumulation (GQA grouped)

Algorithm 7 S2 ANTA-flash, Kernel 1: fused per-tile partials (stats + systematic sampling)

Require: V ∈ Rnk ×Hkv ×dk ; uh,n ; invδh,t , a0,h,t ; tile length Btile . Ensure: Accumulator O ∈ RH×dk (fp32). 1: G ← H/Hkv ; O ← 0. 2: for each KV head k ∈ {0, . . . , Hkv −1} and tile t ∈ {0, . . . , T −1} in parallel do 3: Tt ← [tBtile , min((t+1)Btile , nk )); heads Hk = {kG, . . . , (k+1)G−1}. 4: for each head h ∈ Hk do 5: p ← 0. 6: for each n ∈ Tt (increasing) do 7: x ← invδh,t uh,n . 8: ch,n ← ⌊a0,h,t + p + x⌋ − ⌊a0,h,t + p⌋; p ← p + x. 9: end for 10: end for 11: (Optional) compact emitted rows Et = {n ∈ Tt : ∃h ∈ Hk , ch,n > 0}. 12: for each n ∈ Et and each h ∈ Hk do 13: Oh ← Oh + ch,n Vn,k . 14: end for 15: end for 16: return O (apply final scaling by 1/S outside or in a final cast/scale kernel).

Require: Q ∈ RH×dk ; K, V ∈ Rnk ×Hkv ×dk ; tile length Btile ; samples/head S; seed. eh,t ) for all heads and tiles. Ensure: Per-tile outputs (mh,t , ℓh,t , O 1: T ← ⌈nk /Btile ⌉; G ← H/Hkv ; Stile ≈ S/T (rounded). 2: for each KV head k ∈ {0, . . . , Hkv −1} and tile t ∈ {0, . . . , T −1} in parallel do 3: Tt ← [tBtile , min((t+1)Btile , nk )); heads Hk = {kG, . . . , (k+1)G−1}. 4: for each head h ∈ Hk√do 5: sh,n ← Q⊤ h Kn,k / dk for n ∈ Tt . 6: mh,t ←Pmaxn∈Tt sh,n ; uh,n ← exp(sh,n − mh,t ); ℓh,t ← n∈Tt uh,n . 7: invδh,t ← (Stile /ℓh,t ) (or 0 if ℓh,t is tiny); a0,h,t ∼ U [0, 1). 8: Use systematic counts ch,n from xh,n = invδh,t uh,n eh,t ← P and accumulate O n∈Tt ch,n Vn,k . 9: end for 10: end for

Algorithm 8 S2 ANTA-flash, Kernel 2: FlashAttention-style merge across tiles eh,t ) from Alg. 7; Stile . Require: Per-tile (mh,t , ℓh,t , O H×dk . Ensure: Final O ∈ R 1: for each head h ∈ {0, . . . , H−1} in parallel do 2: m⋆h ← maxt mh,t . P −1 3: Wh,t ← exp(mh,t − m⋆h ) ℓh,t ; Zh← Tt=0 Wh,t . P T −1 eh,t /Stile . 4: Oh ← 1 Wh,t O

P. S 2 ANTA-flash pseudocode Algorithm 6 provides a high-level overview of S 2 ANTAflash. Algorithms 7 and 8 describe the 2 kernels that constitute S 2 ANTA-flash in detail.

t=0

5: end for 6: return O.

Algorithm 6 S2 ANTA-flash (overview, decode/GQA): uniform per-tile sampling + FlashDecoding-style merge Require: Q ∈ RH×dk ; K, V ∈ Rnk ×Hkv ×dk (GQA); tile length Btile ; samples/head S; seed. Ensure: O ∈ RH×dk approximating softmax(QK T )V . 1: T ← ⌈nk /Btile ⌉; G ← H/Hkv ; k(h) ← ⌊h/G⌋; uniform per-tile budget Stile ≈ S/T (rounded). 2: Kernel 1 (tile partials): for each tile (k, t), compute eh,t = P (mh,t , ℓh,t ) and O n∈Tt ch,n Vn,k via systematic sampling with budget Stile . ⋆ 3: Kernel 2 (merge): for each head h, set Pmh = maxt mh,t , ⋆ Wh,t = exp(m − m )ℓ , Z = W h h,t , and Oh = h h,t t h,t P 1 e t Wh,t Oh,t /Stile . Zh 4: return O.

Q. Decode-step kernel microbenchmark details This appendix describes the microbenchmark used to measure decode-step attention latency for FlashAttention2 decoding (“FlashDecoding”), FlashInfer decoding, and our custom S 2 ANTA CUDA kernels. The intent is to be explicit about tensor shapes, software/hardware environment, cache state, and what is (and is not) included in the reported runtimes, in order to support a fair comparison. What is being measured. We measure GPU kernel time for a single decoding step (one new token) as a function of the KV-cache length nk . Timing is performed with CUDA events, recorded on the active CUDA stream and synchronized on the end event each iteration. The reported milliseconds therefore reflect elapsed time on the GPU timeline for the CUDA work launched by a single backend call (which may internally launch multiple kernels). All Python/CPU overhead and input construction are excluded by allocating and populating tensors outside the timed region and recording events immediately around the backend call. 19

Stochastic Sparse Attention for Memory-Bound Inference

Hardware and driver. All measurements were run on a single NVIDIA RTX 6000 Ada Generation GPU (48 GB). The installed driver reported by nvidia-smi was 580.82.07 (CUDA compatibility reported as 13.0).

Timing protocol and trial structure. For each cache length nk and each backend, we run a fixed number of warmup and timed iterations: • Warmup: 10 untimed iterations to amortize one-time effects (e.g., kernel selection/caching inside libraries).

Software environment (versions). The benchmark environment used:

• Timing: 40 timed iterations. Each iteration performs an optional cache-thrash (not timed), then records a start CUDA event, launches the backend once, records an end CUDA event, synchronizes on the end event, and reads elapsed milliseconds from the event pair.

• PyTorch 2.6.0+cu124 • flash-attn 2.7.3 (FlashAttention2) • flashinfer-python 0.5.3

• Aggregation: we report the mean latency over the 40 timed iterations.

Our S 2 ANTA kernels are invoked through a PyTorch CUDA extension (custom CUDA code callable from PyTorch); the timed region includes all CUDA kernels launched by this extension entrypoint.

All tensor allocations and concatenations/appends used to construct inputs are performed outside the timed region.

Decode geometry and tensor shapes. We benchmark a single-batch decode step with head geometry matching a Llama-style GQA configuration: H = 32,

Hkv = 8,

dk = 128,

Note on KV-cache appends in FlashDecoding. FlashAttention2 decoding performs a fused KV-cache append/update as part of its decode call, whereas FlashInfer decoding and our S 2 ANTA decode kernels do not perform an internal append in this benchmark. To isolate the attention-like decode computation for FlashInfer and S 2 ANTA, we provide a cache tensor that already includes the current token in the preallocated [nk +1, Hkv , dk ] storage.

G = H/Hkv = 4,

where H is the number of query heads, Hkv is the number of KV heads, and dk is the head dimension. The query length is one token (decode), batch size is one, and attention is causal.

This asymmetry is not material in the long-context regime: the fused append cost is O(Hkv dk ) per step and does not scale with nk , so its fractional contribution decreases as nk grows. Our speedup comparisons focus on large KV-cache lengths (e.g., 8k–32k tokens), where the attention computation dominates. Including the fused append inside FlashDecoding therefore does not change qualitative conclusions in this regime (and makes the FlashDecoding measurement slightly more conservative relative to an “attention-only” timing).

For each cache length nk , the following tensors are constructed once and reused across timing iterations: • Query tensor with shape [H, dk ]. • Key/value cache tensors stored in a sequence-major layout with shape [nk +1, Hkv , dk ]. (The extra +1 slot corresponds to the current token.) All backends are run with the √ same causal setting and the standard softmax scaling 1/ dk , with dropout disabled.

R. Additional S 2 ANTA kernel results

Input distribution and dtype. All inputs are i.i.d. Gaussian (torch.randn) generated directly on the GPU. Unless otherwise stated, the query, keys, and values are provided in bfloat16 (the harness also supports fp16). A fixed RNG seed is used so that, for a given nk , each backend sees the same input tensors.

Table 23. Accuracy of Llama 8B on 4 long-context tasks with an 8k-token prompt. Tasks: frequent word extraction (FWE), needle-in-a-haystack (NIAH), single-hop QA (QA1 ), multi-hop QA (QA2 ). S: sample budget. Accuracy shows 95% bootstrap confidence intervals. Kernel 2

Cold-cache protocol (L2 thrashing). To reduce sensitivity to favorable cache residency effects: before each timed iteration, we touch a large GPU buffer by performing an in-place update on a 512 MB float32 tensor. This operation is intended to thrash GPU caches (including L2) via a large read+write stream. The cache-thrash operation is executed outside the timed region (before the start event is recorded), and therefore is not included in the reported kernel time.

FWE

NIAH

QA1

QA2

S ANTA-prop S 2 ANTA-prop S 2 ANTA-prop

64 96.60±0.52 93.23±0.69 70.73±2.43 67.33±2.46 128 98.04±0.40 93.78±0.60 70.33±2.40 66.87±2.43 256 98.07±0.40 94.05±0.62 70.67±2.33 66.40±2.30

S 2 ANTA-flash S 2 ANTA-flash S 2 ANTA-flash S 2 ANTA-flash

64 128 256 1024

SDPA (baseline)

66.47±1.12 86.11±0.91 94.80±0.62 98.11±0.40

83.57±1.05 91.70±0.73 92.77±0.63 94.22±0.58

65.67±2.44 69.47±2.30 70.67±2.20 70.73±2.30

62.27±2.40 66.27±2.50 66.80±2.50 66.33±2.40

— 98.53±0.60 94.55±1.00 71.80±3.90 68.60±4.20

Stochastic Sparse Attention for Memory-Bound Inference

References

DeepSeek-AI. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning, 2025. URL https://arxiv.org/abs/2501.12948.

Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebron, F., and Sanghai, S. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. In The 2023 Conference on Empirical Methods in Natural Language Processing, 2023.

Dettmers, T., Pagnoni, A., Holtzman, A., and Zettlemoyer, L. Qlora: Efficient finetuning of quantized llms. Advances in neural information processing systems, 36:10088–10115, 2023.

Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020.

Frantar, E., Ashkboos, S., Hoefler, T., and Alistarh, D. OPTQ: Accurate quantization for generative pre-trained transformers. In The Eleventh International Conference on Learning Representations, 2023. URL https: //openreview.net/forum?id=tcbBPnfwxS.

Boucheron, S., Lugosi, G., and Massart, P. Concentration Inequalities: A Nonasymptotic Theory of Independence. Oxford University Press, Oxford, 2013. Center for AI Safety. MMLU: Massive multitask language understanding dataset. https://huggingface. co/datasets/cais/mmlu, 2024. MIT license. Commit c30699e (uploaded 2024-03-08). Accessed 202505-14. Chen, Z., Sadhukhan, R., Ye, Z., Zhou, Y., Zhang, J., Nolte, N., Tian, Y., Douze, M., Bottou, L., Jia, Z., and Chen, B. MagicPIG: LSH sampling for efficient LLM generation. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreview. net/forum?id=ALzTQUgW8a. Child, R., Gray, S., Radford, A., and Sutskever, I. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.

Gemini Team, Anil, R., Borgeaud, S., Alayrac, J.-B., Yu, J., Soricut, R., Schalkwyk, J., Dai, A. M., Hauth, A., Millican, K., et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805, 2023. Gupta, A., Dar, G., Goodman, S., Ciprut, D., and Berant, J. Memory-efficient transformers via top-k attention. In Moosavi, N. S., Gurevych, I., Fan, A., Wolf, T., Hou, Y., Marasović, A., and Ravi, S. (eds.), Proceedings of the Second Workshop on Simple and Efficient Natural Language Processing, pp. 39–52, Virtual, November 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.sustainlp-1.5. URL https:// aclanthology.org/2021.sustainlp-1.5/. Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Measuring massive multitask language understanding. In International Conference on Learning Representations, 2021. URL https:// openreview.net/forum?id=d7KBjmI3GmQ.

Clark, K., Khandelwal, U., Levy, O., and Manning, C. D. What does BERT look at? an analysis of BERT‘s attention. In Linzen, T., Chrupała, G., Belinkov, Y., and Hupkes, D. (eds.), Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 276–286, Florence, Italy, August 2019. Association for Computational Linguistics. doi: 10.18653/v1/W19-4828. URL https:// aclanthology.org/W19-4828/.

Hooper, C., Kim, S., Mohammadzadeh, H., Mahoney, M. W., Shao, Y. S., Keutzer, K., and Gholami, A. Kvquant: Towards 10 million context length llm inference with kv cache quantization. Advances in Neural Information Processing Systems, 37:1270–1303, 2024.

Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021.

Horowitz, M. 1.1 computing’s energy problem (and what we can do about it). In 2014 IEEE International Solid-State Circuits Conference Digest of Technical Papers (ISSCC), pp. 10–14, 2014. doi: 10.1109/ISSCC.2014.6757323.

Cochran, W. G. Sampling Techniques. Wiley, New York, 3rd edition, 1977.

Hsieh, C.-P., Sun, S., Kriman, S., Acharya, S., Rekesh, D., Jia, F., and Ginsburg, B. Ruler: What’s the real context size of your long-context language models? In First Conference on Language Modeling, 2024.

Dao, T., Haziza, D., Massa, F., and Sizov, G. Flash-Decoding for long-context inference. https://crfm.stanford.edu/2023/10/ 12/flashdecoding.html, October 2023. Stanford Center for Research on Foundation Models (CRFM) blog post, accessed 2026-01-20.

Kim, H. and Ko, J. Fast monte-carlo approximation of the attention mechanism. In Proceedings of the AAAI conference on artificial intelligence, volume 36, pp. 7185– 7193, 2022. 21

Stochastic Sparse Attention for Memory-Bound Inference

Kitaev, N., Kaiser, L., and Levskaya, A. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020. URL https:// openreview.net/forum?id=rkgNKkHtvB.

Samaga B L, Y., Yerram, V., You, C., Bhojanapalli, S., Kumar, S., Jain, P., and Netrapalli, P. Hire: High recall approximate top-k estimation for efficient llm inference, 2024. URL https://arxiv.org/abs/ 2402.09360.

Lightman, H., Kosaraju, V., Burda, Y., Edwards, H., Baker, B., Lee, T., Leike, J., Schulman, J., Sutskever, I., and Cobbe, K. Let’s verify step by step. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum? id=v8L0pN6EOi.

Sun, H., Chen, Z., Yang, X., Tian, Y., and Chen, B. Triforce: Lossless acceleration of long sequence generation with hierarchical speculative decoding. In First Conference on Language Modeling, 2024. Tang, J., Zhao, Y., Zhu, K., Xiao, G., Kasikci, B., and Han, S. Quest: query-aware sparsity for efficient long-context llm inference. In Proceedings of the 41st International Conference on Machine Learning, pp. 47901–47911, 2024.

Liu, Z., Yuan, J., Jin, H., Zhong, S., Xu, Z., Braverman, V., Chen, B., and Hu, X. Kivi: A tuning-free asymmetric 2bit quantization for kv cache. In International Conference on Machine Learning, pp. 32332–32344. PMLR, 2024. Lohr, S. L. Sampling: Design and Analysis. Brooks/Cole, Boston, 2nd edition, 2010. Ma, S., Wang, H., Huang, S., Zhang, X., Hu, Y., Song, T., Xia, Y., and Wei, F. Bitnet b1.58 2b4t technical report, 2025. URL https://arxiv.org/abs/2504. 12285. Meta AI. Meta Llama 3.1 8B Instruct: Version 1. https://huggingface.co/meta-llama/ Llama-3.1-8B-Instruct, 2024. Llama 3.1 Community License. Model release 2024-07-23. Accessed 2025-05-14.

Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lample, G. Llama: Open and efficient foundation language models, 2023. URL https://arxiv.org/ abs/2302.13971. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. Advances in neural information processing systems, 30, 2017. Wang, S., Li, B. Z., Khabsa, M., Fang, H., and Ma, H. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020.

OpenAI. Gpt-4 technical report, 2023. URL https:// arxiv.org/abs/2303.08774.

Williams, R. J. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning, 8(3-4):229–256, 1992.

OpenAI. GSM8K: Grade school math word problems (v1.0). https://huggingface.co/datasets/ openai/gsm8k, 2023. MIT License, commit e53f048. Accessed 2025-05-14.

Xie, X., Luo, Y., Peng, H., and Ding, C. RTop-k: Ultra-fast row-wise top-k selection for neural network acceleration on GPUs. In The Thirteenth International Conference on Learning Representations, 2025. URL https:// openreview.net/forum?id=PHg4rAXFVH.

Owen, A. B. Monte Carlo Theory, Methods and Examples. 2013. Monograph. Pinelis, I. Optimum bounds for the distributions of martingales in banach spaces. The Annals of Probability, 22(4): 1679–1706, 1994. Rajpurkar, P., Jia, R., and Liang, P. Know what you don’t know: Unanswerable questions for SQuAD. In Gurevych, I. and Miyao, Y. (eds.), Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pp. 784–789, Melbourne, Australia, July 2018. Association for Computational Linguistics. doi: 10.18653/v1/P18-2124. URL https://aclanthology.org/P18-2124/.

Yang, Z., Qi, P., Zhang, S., Bengio, Y., Cohen, W., Salakhutdinov, R., and Manning, C. D. HotpotQA: A dataset for diverse, explainable multi-hop question answering. In Riloff, E., Chiang, D., Hockenmaier, J., and Tsujii, J. (eds.), Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 2369–2380, Brussels, Belgium, OctoberNovember 2018. Association for Computational Linguistics. doi: 10.18653/v1/D18-1259. URL https: //aclanthology.org/D18-1259/.

Ribar, L., Chelombiev, I., Hudlass-Galley, L., Blake, C., Luschi, C., and Orr, D. Sparq attention: Bandwidthefficient llm inference. In International Conference on Machine Learning, pp. 42558–42583. PMLR, 2024.

Ye, Z., Chen, L., Lai, R., Lin, W., Zhang, Y., Wang, S., Chen, T., Kasikci, B., Grover, V., Krishnamurthy, A., and Ceze, L. Flashinfer: Efficient and customizable attention engine for LLM inference serving. In Eighth Conference 22

Stochastic Sparse Attention for Memory-Bound Inference

on Machine Learning and Systems, 2025. URL https: //openreview.net/forum?id=RXPofAsL8F. Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Alberti, C., Ontanon, S., Pham, P., Ravula, A., Wang, Q., Yang, L., et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283–17297, 2020. Zhang, Z., Sheng, Y., Zhou, T., Chen, T., Zheng, L., Cai, R., Song, Z., Tian, Y., Ré, C., Barrett, C., et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36:34661–34710, 2023.

Record · ID 155250 · SHA-256 66286dba5e223dbd
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.