PaperlayerView exact revision on arXiv

Source-linked AI summary

SANTA++: Sampling Attention through Representative Keys

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

arXiv:2609.35629v1cs.LGcs.CL

TL;DR

At long contexts, attention may concentrate on a query-dependent subset of cached tokens, but finding that subset with exact scores requires reading every key. SANTA++ routes queries through representative-key teams, scores sampled members exactly, and corrects their contributions for unequal selection chances. At 32K, it reports near-dense scores with lower KV access across benchmarks, and its GPU implementation reports an attention speedup on an evaluated workload.

  • Problem

    At long context lengths, attention can concentrate on a small subset of cached tokens, but finding that subset by exact scores still requires reading every key.

    Source: §1 · p. 1

  • Method

    SANTA++ routes queries using actual-key representatives, evaluates selected teams’ members exactly, and corrects their contributions for unequal inclusion probabilities.

    Source: §1 · p. 1, §6 · p. 9

  • Results

    At 32K, SANTA++ retains near-dense scores with lower reported KV access across benchmarks; on LongBench v2, it retains 99.07% of the SDPA score at 16.68% reported KV access.

    Source: §6 · p. 9

  • Takeaways & Limitations

    The reported sparsity after latent compression and head sharing suggests SANTA++’s selective token access may complement MLA-style architectures.

    Source: §6 · p. 9

    Reported logical KV access is not measured DRAM traffic, resident cache size, or latency, and nominal sampling budget does not uniquely determine access.

    Source: §3 · p. 5

Abstract

from arXiv · show

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

1 INTRODUCTION

SANTA++ addresses query-dependent attention sparsity without scanning every cached key, routing through actual-key representatives and evaluating sampled teams exactly. At 32K context, it retains most dense-attention scores with substantially less reported KV access, and its implementation also reports an attention speedup on a specific GPU workload.

1 INTRODUCTION

  • At long contexts, identifying query-relevant cached tokens by exact attention scores still requires reading every key, even when attention concentrates on a small subset.For Qwen2.5-7B-Instruct at 32K context, streaming bf16 keys and values amounts to about 64 MiB per layer for each generated token.Source: §1 · p. 1
  • SANTA++ uses actual-key representatives to route each query to teams, then computes exact scores for selected members and corrects team contributions for unequal inclusion probabilities.Newly generated tokens form a separate suffix evaluated exactly.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • SANTA++ is a grouped generalization of SANTA that replaces its full-key scan with approximate routing to reduce KV access.The paper describes SANTA as sampling from the exact attention distribution after scoring every cached key.Source: §1 · pp. 1–2
  • The nominal sampling budget provides an adjustable accuracy–KV-access trade-off, while actual access also depends on team sizes, routing and suffix reads, and overlap across query heads.For the Qwen2.5-7B-Instruct evaluations, parent size P = 16 and up to R = 4 representatives per parent were fixed heuristically.Source: §1 · p. 2
  • 99.07% of the dense baseline’s LongBench v2 score is retained at 16.68% of its KV reads, while HELMET RAG retains 98.00% at 38.49% access.On RULER, increasing the sampling budget raises retained score from 91.07% to 98.54% as access rises from 19.95% to 46.95%.Source: §1 · p. 2
  • At 32K context, the Triton implementation reports a 1.69× attention-operator speedup over Flash SDPA on a batch-one FP16 workload using 31 sampled teams per query head.The reported workload is layer 14 of Qwen2.5-7B-Instruct on an NVIDIA GeForce RTX 5090 Laptop GPU.Source: §1 · p. 2

2 METHOD

SANTA++ uses representative-guided sampling to reduce full-cache reads while estimating attention contributions from sampled teams. Its method builds teams within parent groups, samples them using query-dependent routing, and evaluates a routing diagnostic against dense attention.

2 METHOD

  • Representatives guide which teams to read, but sampled members’ keys and values supply their exact attention contributions to approximate the full-cache mass and value-weighted numerator.The full-cache output is the value-weighted numerator divided by the unnormalized attention mass; the sampled sums are corrected before combination.Source: §2 · p. 3

2.1 APPROXIMATING GROUP MASS

  • A mean-key score can underestimate group attention mass because exponentiating the average logit suppresses high-scoring members compared with summing their exponentials.For 16 constructed logits, the exact mass is approximately 2.20 × 10^4 versus approximately 29 for the surrogate; equality holds when all within-group logits are equal.Source: §2.1 · p. 3, Figure 2 · p. 3

2.2 PREPARING PARENT GROUPS AND TEAMS

  • After dense prefill, parent groups are formed by k-means or contiguous token spans, then each member is assigned to its nearest actual-key representative within its parent.Representatives are selected to cover the parent: the first is nearest its mean key, and later representatives are farthest from existing representatives; team records are reused during decoding.Source: §2.2 · pp. 3–4

2.3 SAMPLING AND CORRECTING TEAM CONTRIBUTIONS

  • Teams are sampled without replacement using representative scores weighted by team size; selected members are scored exactly, and each team’s mass and weighted value sum are divided by its conditional inclusion probability.The two unnormalized sums are unbiased, but their ratio—the attention output—is generally biased; selecting all teams recovers dense attention.Source: §2.3 · p. 4

2.4 ATTENTION-DISTRIBUTION DIAGNOSTIC

  • The routing diagnostic compares dense attention over prompt tokens with routing probabilities spread equally across group members; the P = 16, R = 4 team variant has lower mean KL divergence than both centroid-guided variants on all three tasks.The comparison uses the same dense SDPA queries and cache states and performs no sampling.Source: §2.4 · p. 5, Table 1 · p. 5

3 ACCURACY AND THEORETICAL KV ACCESS

At 32K context, SANTA++ reduces reported logical KV access while retaining much of dense SDPA’s score across HELMET RAG, LongBench v2, and RULER. Access is a logical counter, not measured DRAM traffic or latency, and accuracy can vary substantially across tasks.

Access accounting

  • Reported KV access is logical key-and-value vector access relative to dense decoding, not measured DRAM traffic, resident cache size, or latency.S is a nominal budget, not a unique-row count; team sizes and overlap across query heads affect access.Source: §3 · p. 5

HELMET RAG

  • 98.00% of the SDPA RAG score is retained on HELMET at 38.49% KV access (S = 1024); at S = 128, retention is 94.43% at 17.49% access.Contiguous grouping at S = 1024 retains 97.56% at 31.23% access, with 56.00% absolute score versus 56.25% for k-means.Source: §3 · p. 5, Table 2 · p. 6

HELMET task variation

  • JSON KV retrieval on HELMET falls from 96.80% under SDPA to 45.60% with k-means at S = 128, then recovers to 87.00% at S = 512.HELMET’s RAG summaries equally weight four tasks and include representative-key routing reads; other HELMET tasks are reported separately because their metrics differ.Source: §3 · pp. 5–6

LongBench v2

  • 99.07% of the SDPA score is retained on LongBench v2 at 16.68% reported KV access (S = 128), with absolute accuracy of 31.61% versus 31.91% for SDPA.Contiguous grouping reaches 98.44% of the SDPA score at 22.15% access (S = 512).Source: §3 · p. 6, Table 3 · p. 7

RULER

  • 91.07% of the SDPA aggregate score is retained on RULER at 19.95% reported KV access (S = 256); at 25.70% access, retention rises to 93.99%.The aggregate covers all 13 tasks; MK-NIAH-3 scores 60.20% versus 89.60% for SDPA at contiguous S = 1024.Source: §3 · pp. 6–7, Table 4 · p. 7

Cross-benchmark grouping comparison

  • At S = 512, contiguous grouping scores 75.37% with 21.37% access, versus 76.05% with 19.95% access for k-means; at S = 1024 it retains 95.07% of SDPA below 30% access.It avoids parent k-means preparation.Source: §3 · p. 6

4 GPU IMPLEMENTATION AND ATTENTION LATENCY

The Triton implementation speeds up complete attention on a captured 32K-context workload. Its timing includes routing, selection, and kernel launches, but excludes preparation.

4 GPU IMPLEMENTATION AND ATTENTION LATENCY

  • 63.52 µs versus 107.33 µs for Flash SDPA yields a 1.69× speedup on the captured 32K-context workload.The comparison uses layer 14 of Qwen2.5-7B-Instruct, batch size one, FP16 inputs, and an NVIDIA GeForce RTX 5090 Laptop GPU.Source: §4 · p. 8, Table 5 · p. 8
  • K = 31 sampled teams per query head (nominal S = 124) gives the Triton top-K stage a power-of-two working set of 32 scores.The sampler keeps the top K teams plus the next-highest Gumbel-perturbed routing score to compute the conditional-inclusion threshold.Source: §4 · p. 8
  • With GQA, representative-key and selected-member loads are reused across query heads sharing a KV head, while each head retains its own team selections and corrections.Representatives guide selection; exact selected-member scores determine contributions.Source: §4 · p. 8
  • The benchmark times routing, selection, and all five kernel launches, while excluding cache packing and temporary-buffer preparation.End-to-end gains also depend on preparation cost and generation length.Source: §4 · p. 8, Table 5 · p. 8
  • The kernel benchmark uses S = 124, whereas task-accuracy sweeps start at S = 128.Source: §4 · p. 8

5 RELATED WORK

SANTA++ frames sparse decoding as query-dependent sampling with corrections toward full-cache attention. Its use of actual-key representatives distinguishes its routing from mean-key approaches.

5 RELATED WORK

  • Reducing cache-entry reads targets memory traffic and can complement quantization, KV sharing, or latent compression that reduces memory occupied by weights or the cache.Source: §5 · p. 8
  • Unlike cache-retention methods that limit which KV entries remain available, SANTA++ does not permanently drop prompt tokens, so skipped entries can be selected later.The two approaches can be combined: retention reduces storage, while SANTA++ reduces reads from the retained cache.Source: §5 · p. 8
  • SANTA++ uses query-dependent sampling and corrects sampled contributions toward full-cache attention sums, rather than simply renormalizing a selected subset.This contrasts with structured, low-rank, hashing, and deterministic top-k sparse-attention approaches.Source: §5 · p. 8
  • Jensen’s inequality implies that scoring an arithmetic-mean key systematically underestimates a group’s unnormalized attention mass, motivating actual-key representatives.The paper notes this observation may also inform centroid-based retrieval methods.Source: §5 · p. 9
  • When the number of teams equals context length, each key forms its own team and SANTA++ recovers SANTA’s exact sampling distribution; larger teams trade exact routing for lower KV access.SANTA samples from the exact post-softmax distribution but scores every cached key before sampling.Source: §5 · p. 9
  • SANTA++ shares MagicPIG’s estimator view of sparse attention, but uses explicit groups and representative estimates to allocate samples across the cache.MagicPIG uses locality-sensitive hashing for self-normalized importance sampling.Source: §5 · p. 9

6 CONCLUSION

SANTA++ selects query-relevant teams using representative keys and importance correction, trading attention accuracy against KV access without retraining. Evaluations show high retained scores at reduced access, and experiments suggest useful token sparsity can remain after latent compression and head sharing.

6 CONCLUSION

  • SANTA++ routes queries with actual-key representatives, computes exact scores within selected teams, and corrects contributions for unequal inclusion probabilities.The method offers an adjustable accuracy–KV-access trade-off without retraining.Source: §6 · p. 9
  • 99.07% of the SDPA score on LongBench v2 at 16.68% logical KV access and 98.00% on HELMET RAG at 38.49% access are retained with k-means grouping.On RULER, a larger budget retains 98.54% of the SDPA score at 46.95% KV access; contiguous grouping retains 97.56% on HELMET RAG at 31.23% access while avoiding k-means clustering.Source: §6 · p. 9
  • The DeepSeek-V2-Lite-Chat experiments suggest useful token sparsity can remain after latent compression and head sharing.The paper motivates combining selective token access with MLA-style architectures to reduce memory traffic in already compressed models.Source: §6 · p. 9

AI USE STATEMENT

The paper reports AI assistance across writing, implementation, methodological development, and analysis, with author review and validation. It also reports kernel release and funding support.

AI USE STATEMENT

  • AI tools assisted with writing and presentation, software implementation and debugging, methodological and mathematical development, and analysis and interpretation of results.The authors state that all AI-assisted outputs were reviewed and validated.Source: AI USE STATEMENT · p. 9

REPRODUCIBILITY STATEMENT

A ALGORITHM AND ESTIMATOR DETAILS

SANTA++ forms teams from cached keys, routes each query using actual-key representatives, and samples teams with an inclusion correction. Its estimator makes the unnormalized attention sums unbiased, while a mean-key surrogate can misestimate group mass when member logits differ.

A ALGORITHM AND ESTIMATOR DETAILS

  • SANTA++ partitions prompt keys into nonempty teams, stores each team’s actual-key representative and members, and reuses those records during decoding.Parent partitions can come from k-means or contiguous token spans; generated suffix rows are attended exactly.Source: A ALGORITHM AND ESTIMATOR DETAILS · pp. 12, 14, 15
  • Each query samples K teams using representative scores adjusted by team size, then reads selected members and evaluates their exact attention scores.The nominal budget S determines K subject to the stated divisibility conditions; the exact suffix is included separately.Source: A ALGORITHM AND ESTIMATOR DETAILS · p. 13
  • For Gumbel-top-K sampling, the selected team’s correction uses its leave-one-out threshold, obtained from the observed (K + 1)st perturbed routing score.The first-draw probability is not generally the inclusion probability.Source: A ALGORITHM AND ESTIMATOR DETAILS · p. 13
  • The corrected unnormalized attention-mass and weighted-value sums are unbiased for the current query and cache, but their normalized ratio is generally biased.Selecting all teams recovers dense attention in exact arithmetic.Source: A ALGORITHM AND ESTIMATOR DETAILS · p. 14

B.1 ADDITIONAL JENSEN EXAMPLES

  • The mean-key mass surrogate equals exact unnormalized mass only when all member logits are equal; constructed 16-logit examples show ratios from approximately 1.03 to 760.Bimodal, heavy-tailed, and sparse logits produce larger gaps because high-scoring members contribute more to the exponential sum.Source: §B.1 · pp. 15–16, Figure 3 · p. 16

B.2 DENSE-TRAJECTORY KL AND TASK-VARIANT EVALUATION

The diagnostics compare routing and accuracy on fixed dense-generation trajectories, then examine grouping compactness and the attention mass retained by contiguous windows. Whole-team sampling leads the tested RULER variants in all-task point estimates, with higher reported KV access.

B.2 DENSE-TRAJECTORY KL AND TASK-VARIANT EVALUATION

  • At S = 128, whole-team sampling reaches 84.16% accuracy versus 87.09% for SDPA while accessing 24.97% of dense KV; it has the highest all-task point estimate at every tested budget.Increasing the team budget narrows the aggregate score gap but requires more access; guided variants use fewer logical reads at the same nominal budget.Source: §B.2 · p. 17, Table 6 · p. 18
  • Task-specific trade-offs remain: guided B = 4 has higher QA-2 point estimates at every budget, while teams lead on CWE and VT; greater sampling does not improve every team result.The KL diagnostic holds queries and cache states fixed across variants, removing differences from divergent generated histories.Source: §B.2 · p. 17

B.3 CENTROID-DISTANCE DIAGNOSTIC

  • Contiguous chunks have intermediate mean key-to-centroid distance between key clusters and random chunks on every plotted RULER task.The diagnostic uses 100 prompts per task at 8,192-token context and nominal group size P = 16.Source: §B.3 · pp. 18–19, Figure 4 · p. 19

B.4 ORACLE PREFILL-WINDOW DIAGNOSTIC

  • For most heads, 32 windows of width 16 retain at least half the attention mass of the top 512 individual tokens, but only 15.1%–23.6% reach 0.90 WER.This oracle diagnostic uses dense attention weights from the final prefill query token.Source: §B.4 · p. 19, Table 7 · p. 20

B.5 KEY VERSUS PROBE-FINGERPRINT CLUSTERING

The comparison tests direct key clustering against clustering of probe-query score fingerprints using the same team construction, estimator, and sampling procedure. Accuracy advantages vary by task and budget rather than consistently favoring probe fingerprints.

B.5 KEY VERSUS PROBE-FINGERPRINT CLUSTERING

  • Key clustering has higher MK-NIAH-3 point estimates at all three budgets, while probe fingerprints lead on CWE at S = 128 and 256.On QA-2, the ordering reverses with budget, both variants remain below the dense point estimate, and the intervals overlap substantially.Source: §B.5 · p. 20
  • There is no uniform accuracy advantage for probe-fingerprint clustering; direct key clustering is the simpler preparation choice.The comparison evaluates Qwen2.5-7B-Instruct at 8K on four RULER tasks and budgets S = 128, 256, and 512.Source: §B.5 · p. 20

C.1 SETUP AND METRICS

The evaluation defines benchmark settings and logical KV-access accounting, then describes the GPU decode-kernel pipeline. Its RAG task results show that performance can vary substantially across tasks and sampling budgets.

C.1 SETUP AND METRICS

  • The reported KV-access measure includes representative routing reads, selected prompt rows, and the exact generated suffix, with aggregation rules varying by benchmark.The access counter excludes prefill and team construction; HELMET RAG access averages the four task-level total-access percentages equally.Source: §C.1 · p. 22

C.2 HELMET

  • At 32K, JSON KV retrieval scores are 45.60% with k-means and 33.40% with contiguous grouping at S = 128, rising to 87.00% and 68.60%, respectively, at S = 512.The comparison is against SDPA’s 96.80% score.Source: §C.2 · p. 22, Table 10 · p. 24

C.3 LONGBENCH V2

  • LongBench v2 uses 503 distinct prompts and a 32K context cap, with prompts exceeding the cap middle-truncated and shorter prompts retained at original length.The prompting requests concise reasoning and a fixed final-answer format.Source: §C.3 · p. 25

C.4 RULER

  • The main RULER evaluation uses 500 completed prompts per task at both 8K and 32K, following the reference task definitions, prompts, and grading conventions.Reported GQA-aware access includes selected KV reads and representative routing reads.Source: §C.4 · p. 25

D TRITON DECODE KERNELS

  • The Triton decode pipeline scores representatives, selects teams, merges candidates and GQA unions, computes exact selected-member attention, then reduces partial softmax results.For the timed configuration, K = 31 and the top-K stage retains K + 1 scores for the conditional-inclusion threshold.Source: D TRITON DECODE KERNELS · p. 28

E KERNEL EVALUATION PROTOCOL AND PROFILING

The section documents a controlled 32K kernel benchmark and profiler protocol, then reports timing variability and resource-use differences across attention stages.

E KERNEL EVALUATION PROTOCOL AND PROFILING

  • The captured workload fixes a 32,768-token cache, batch size one, one query token, 28 query heads, four KV heads, and FP16 inputs.It selects 31 teams per query head; prefill, team construction, cache packing, input preparation, and workspace allocation are excluded from timing.Source: E KERNEL EVALUATION PROTOCOL AND… · p. 28
  • Pre-timing checks compare outputs against an independent reference using identical selected rows and inclusion corrections, with maximum absolute error below 4.30 × 10−6.The checks cover random-number generation, top-K selection and exact merging, inclusion corrections, GQA union construction, and buffer guards.Source: E KERNEL EVALUATION PROTOCOL AND… · p. 29
  • Profiler-off latency uses 20 interleaved paired rounds with 1,000 complete-operator observations per method per round and an upstream 512 MiB eviction buffer.CUDA-event intervals include GPU execution and device idle gaps between host-issued launches, not host wall-clock latency.Source: E KERNEL EVALUATION PROTOCOL AND… · p. 29
  • [1.689, 1.691] is the 95% paired-round bootstrap interval for the Flash-to-SANTA++ speedup ratio, summarizing within-session uncertainty.The estimates use 10,000 resamples of paired rounds; all 20 rounds were retained.Source: E KERNEL EVALUATION PROTOCOL AND… · p. 29
  • Isolated-stage medians have different cache and timing boundaries from the uninterrupted operator, so they must not be summed as shares of its latency.The complete five-launch operator is measured directly.Source: E KERNEL EVALUATION PROTOCOL AND… · p. 29, Table 13 · p. 29
  • Dense split-KV forward reports 83.95% DRAM throughput, versus 63.85% for representative scoring and 19.46% for selected attention; Top-K occupancy is 84.51%.NCU values are per-metric medians from three captures using kernel replay, not replacements for canonical profiler-off latency.Source: E KERNEL EVALUATION PROTOCOL AND… · p. 29, Table 14 · p. 30

F TOKEN SPARSITY AFTER LATENT COMPRESSION

The section evaluates whether token selection remains useful with a shared latent cache. It reports an accuracy–access trade-off, while showing that shared-row overlap, routing reads, and task differences constrain the savings.

F TOKEN SPARSITY AFTER LATENT COMPRESSION

  • SANTA++ selects token rows in DeepSeek-V2-Lite-Chat’s shared latent cache rather than further compressing each row; independently selected heads can still create a large shared-row union.Heads share teams but draw independently; generated suffix rows are evaluated exactly and are not added to the fixed prompt teams.Source: F TOKEN SPARSITY AFTER LATENT… · pp. 30, 32
  • The evaluation covers niah multivalue, fwe, and qa 1 at 8K using 50 matched prompts per task and a dense SDPA reference through the same compressed-cache adapter.The table’s absolute accuracy is the equal-weight task mean; access counts routing and the union of selected shared rows across heads.Source: F TOKEN SPARSITY AFTER LATENT… · pp. 30–31, Table 15 · p. 30
  • At S = 512, SANTA++ reaches 84.56% mean accuracy versus 86.22% for SDPA, retaining 98.07% of the dense mean with 86.18% logical shared-row access.At S = 256, it retains 91.56% of the dense mean at 70.09% logical access.Source: F TOKEN SPARSITY AFTER LATENT… · p. 30, Table 15 · p. 30
  • Reported access is a logical measure, not realized memory-traffic savings or a latent-cache speedup, and each method’s denominator follows its own generation trajectory.The reference implementation stores expanded routing keys and reconstructs full K/V before selection.Source: F TOKEN SPARSITY AFTER LATENT… · p. 31
  • At S = 512, attention union and exact suffix use 63.75% of the dense shared-row count, with routing adding 22.43%; at S = 1024, total logical access reaches 101.89%.The higher budget does not improve the aggregate accuracy point estimate; separate routing and attention passes can produce access above 100%.Source: F TOKEN SPARSITY AFTER LATENT… · p. 31
  • At S = 512, fwe matches SDPA’s 96.67% score with 80.34% logical access, while multivalue retrieval and qa 1 remain respectively one and four percentage points below their dense scores.At S = 1024, the qa 1 point estimate rises while the multivalue-retrieval estimate falls, so the sweep is task-dependent rather than monotonically improving.Source: F TOKEN SPARSITY AFTER LATENT… · p. 31, Table 16 · p. 31, Table 17 · p. 31
  • Prompt scores use the fraction of required answers recovered for fwe and niah multivalue, but check whether any accepted answer appears for qa 1.Decoding is greedy, with task-specific generation limits of 50, 128, and 32 new tokens, respectively.Source: F TOKEN SPARSITY AFTER LATENT… · p. 32

Related papers

Loading 2609.35629v1…