RSA-X: A Reproducible Framework and Benchmark Protocol for Sparse Attention

Abstract

Attention in Transformer language models is routinely assumed to be dense, yet a growing body of work builds efficient sparse attention mechanisms on the premise that real attention is concentrated on a small subset of tokens. Whether that premise holds - and how concentrated attention actually is - is rarely measured under controlled, reproducible conditions. We present RSA-X, an open framework that makes this question answerable. RSA-X provides (i) a decoupled backend contract for benchmarking dense and sparse attention implementations under a strict, synchronized execution protocol, and (ii) an observational analysis pipeline that profiles attention concentration in pretrained models using Shannon entropy, top- $k$ cumulative mass, active density, and near-zero sparsity. We apply the pipeline to the GPT-2 family (gpt2, gpt2-medium, gpt2-large; 12, 24, and 36 layers) on 100 sequence blocks (length 512) of the WikiText-103 test split on an NVIDIA Tesla T4, and report three quantitative findings. First, attention is heavily concentrated: on average 70.0-75.5% of attention weights are near-zero (at most 1e-4) across model scales, while the top-10 keys carry 72.9-77.6% of the attention mass. Second, concentration grows monotonically with depth: mean entropy falls from approximately 4.3 nats in early layers to approximately 1.9 nats in deep layers of GPT-2 Large, and layer-wise sparsity peaks at 80.8% (layer 14), with attention tails following a power law $(\hat{\alpha} = 1.68,\ R^2 = 0.96)$ . Third, the depth profile of concentration is scale-invariant: the layer-wise entropy, sparsity, and top- $k$ curves of gpt2, gpt2-medium, and gpt2-large are highly correlated, and entropy and sparsity are strongly negatively correlated, with $r$ less than $-0.8$ at the head level. These results provide direct quantitative support for the central assumption behind sparse attention, KV-cache eviction, and attention-sink methods, and establish RSA-X as a standardized, reproducible instrument for future work. All code, figures, and artifacts are open source.

1. Introduction

Scaled dot-product attention [1] is the computational core of modern Large Language Models (LLMs) and Transformers. It computes a full $O(N^2)$ query–key interaction matrix, which becomes a prohibitive computational and memory bottleneck as sequence length $N$ scales to tens or hundreds of thousands of tokens, and dominates the memory footprint of inference through the key–value cache.

A wide spectrum of efficient attention mechanisms has been proposed to escape this quadratic cost:

Structured sparse patterns — Sparse Transformer [2], Longformer [3], BigBird [4], Mistral Sliding Window [5].
Linear and low-rank approximations — Linformer [6].
Locality-Sensitive Hashing — Reformer [7].
Kernel and random-feature expansions — Performer [8].
Selective state-space models — Mamba [9].
Dynamic KV-cache eviction and compression — H2O [10], StreamingLLM [11], SnapKV [12].

In parallel, exact IO-aware algorithms such as FlashAttention [13], FlashAttention-2 [14], and FlashAttention-3 [15] have shown that GPU memory traffic and warp scheduling frequently dominate asymptotic FLOP counts in practice; distributed long-context paradigms such as Ring Attention [16] extend this to near-infinite contexts; and compiler frameworks such as FlexAttention [17] and Triton [18] make custom sparse attention kernels increasingly easy to express.

Nearly all of these methods share one empirical premise: pretrained attention is not dense - it is concentrated on a small set of keys. Yet this premise is rarely measured directly. Reported evaluations use disparate hardware, unsynchronized CUDA timers, varying floating-point formats, custom padding, or unverified dense controls, making it difficult to distinguish genuine algorithmic savings from hardware- or measurement-dependent artifacts. As highlighted by the Long Range Arena benchmark [19] and efficiency surveys [20], and by work on benchmark variance [21], the field lacks a standardized, inspectable way to answer two questions:

How concentrated is attention in pretrained models, as a function of layer and model scale?
How should a candidate sparse attention backend be measured so that results transfer across machines?

RSA-X addresses both. It contributes:

A minimal, decoupled Python API contract (AttentionBackend) for registering dense, vendor-optimized, and third-party sparse attention backends in under ten lines of code.
A reproducible microbenchmarking protocol that measures numerical fidelity against an exact dense reference (PyTorch SDPA [22]), synchronized wall-clock latency, token throughput, TFLOP/s, and peak CUDA memory, and exports everything to a versioned, portable JSON artifact (rsa-x-attention-benchmark/v1).
An observational attention profiling suite that quantifies Shannon entropy, top- $k$ cumulative mass, active density, and near-zero sparsity across layers, heads, and samples of pretrained models.
A pre-registered evaluation protocol that strictly separates implemented framework features from empirical claims.
An empirical study (this paper) that uses the suite to measure natural attention concentration in the GPT-2 family, providing the first controlled, single-instrument comparison of sparsity and concentration across model scales.

2. Background and Related Work

2.1 Attention and its cost

For queries $Q$ , keys $K$ , and values $V \in \mathbb{R}^{B \times H \times N \times d}$ (batch $B$ , heads $H$ , sequence length $N$ , head dimension $d$ ), dense attention with optional causal mask $M$ (where $M\_{ij} = 0$ for attended positions and $-\infty$ for masked ones) is:

\operatorname{Attn}(Q,K,V) = \operatorname{softmax}\left(\frac{QK^{\mathsf T}}{\sqrt{d}} + M\right)V

The $QK^\top$ term materializes an $O(N^2)$ matrix per head; both compute and memory scale quadratically with $N$ .

2.2 Taxonomies of efficient attention

Pattern-based sparsity. Fixed strides and local windows [2] reduce computed pairs to $O(N\sqrt{N})$ or $O(N\log N)$ . Longformer [3] and BigBird [4] combine local windows with global and random connections to preserve graph connectivity.
Hardware-aware IO optimization. FlashAttention [13, 14, 15] tiles the computation across GPU SRAM and uses online softmax rescaling, avoiding $O(N^2)$ HBM traffic without changing output semantics.
Dynamic context eviction. H2O [10], StreamingLLM [11], and SnapKV [12] prune or compress KV-cache tokens at generation time, motivated by observed spatial attention concentration and positional "sink" tokens.
Flexible execution APIs. PyTorch FlexAttention [17] compiles custom block-sparse masks into optimized Triton [18] kernels.

2.3 The evaluation gap

Despite architectural progress, published speedups often fail to transfer across hardware generations or driver versions [21]. Common methodological problems include: omitting warm-up effects; using asynchronous host-to-device timers that hide true GPU execution time; conflating allocated vs. reserved memory; and comparing against unverified dense controls. RSA-X is designed to close this gap by fixing the variables that cause benchmark variance and by making the observational basis for sparsity claims measurable and reproducible.

3. Methods: The RSA-X Framework

RSA-X comprises two decoupled execution paths: an Observational Analysis Path and an Interchangeable Microbenchmark Path.

3.1 Backend contract

Any candidate backend satisfies the AttentionBackend protocol by defining a unique name and implementing:

class AttentionBackend(Protocol):
    name: str

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        *,
        is_causal: bool,
    ) -> torch.Tensor: ...

Inputs and output use shape [B, H, N, d]; the output must match the input shape, device, and dtype, and inputs are treated as immutable.

3.2 Observational diagnostics

For each attention row $p\_i \in \mathbb{R}^N$ (weights of query token $i$ over keys, normalized so $\\sum\_{j=1}^{N} p\_{ij} = 1$ ), RSA-X computes:

Shannon entropy (base 2):

H(p_i) = -\sum_{j=1}^{N} p_{ij} \ln p_{ij}

Top- $k$ cumulative mass (where $\mathrm{TopK}(p\_i)$ denotes the indices of the $k$ largest entries of $p\_i$ ):

M_k(p_i) = \sum_{j \in \mathrm{TopK}(p_i)} p_{ij}

Active density ratio (fraction of weights above the uniform-attention baseline $1/N$ ):

D(p_i) = \frac{1}{N}\sum_{j=1}^{N} \mathbf{1}\bigl(p_{ij} > 1/N\bigr)

Near-zero sparsity (percentage of weights at or below $\\epsilon = 10^{-4}$ ):

S(p_i) = \frac{100}{N}\sum_{j=1}^{N} \mathbf{1}\bigl(p_{ij} \leq \epsilon\bigr)

In addition, attention tails are fitted to a power law $\ln w\_r = -\alpha \ln r + C$ over ranks $r$ sorted by descending weight, yielding an exponent $\alpha$ and goodness-of-fit $R^2$ .

3.3 Benchmarking and synchronization protocol

To prevent GPU launch latency and memory caching from skewing measurements, the microbenchmark path enforces strict execution boundaries:

Deterministic inputs: $Q, K, V$ are generated with fixed PRNG seeds (torch.manual_seed(seed)), matching the requested workload [B, H, N, d].
Correctness verification: candidate outputs are compared against PyTorch SDPA [22] using cosine similarity, relative $L\_2$ error, exact-match ratio, and max/mean absolute error (stride-safe, computed on .reshape(-1) views).
Warm-up: $W = 10$ warm-up iterations are executed without timing.
Synchronization and memory reset: torch.cuda.synchronize() is called, then torch.cuda.reset_peak_memory_stats().
Timed loop: $M = 30$ measured iterations, each timed individually with time.perf_counter() and synchronized.
Post-loop synchronization before recording latency distribution (p50 median, min, max, std), throughput (tokens/s), TFLOP/s, and peak allocated/reserved memory.

Results are exported as rsa-x-attention-benchmark/v1 JSON plus Markdown and HTML reports, with full environment metadata (GPU model, PyTorch version, driver) for provenance.

4. Experimental Setup

Hardware. All experiments ran on a single NVIDIA Tesla T4 (14.56 GB VRAM, Linux 6.12.90+, PyTorch 2.10.0+cu128) on Kaggle, with torch.set_num_threads(1) to minimize host overhead.

Models. Pretrained decoder-only GPT-2 models at three scales: gpt2 (12 layers, 12 heads), gpt2-medium (24 layers, 16 heads), and gpt2-large (36 layers, 20 heads), loaded via TransformerLens [23] with weights frozen.

Data. 100 contiguous sequence blocks of length 512 from the test split of WikiText-103 (wikitext-103-raw-v1) [24], tokenized with the model's own tokenizer; batch size 1 per sample.

Protocol. Attention patterns were extracted from every layer and head via TransformerLens hooks under torch.no_grad(); the four diagnostics (Section 3.2) were computed per query position and aggregated per head, per layer, and globally. All measured metrics passed automated research validation (finite, in-range values; strict top- $k$ monotonicity: top-1 mass is at most top-5 mass, which is at most top-10 mass, which is at most top-50 mass).

5. Results

5.1 Attention is heavily concentrated at every scale

Table 1 reports global aggregates for the three model scales. Across all scales, 70.0–75.5% of attention weights are below $10^{-4}$ , and the top-10 keys concentrate 72.9–77.6% of the total attention mass. Even a single key (top-1) carries 44.8–46.8% of the mass.

Table 1. Global attention concentration statistics (WikiText-103 test, 100 blocks × 512 tokens, T4).

Top-k values for gpt2-medium were not persisted in the archived run; the suite computes them by default.

5.2 Concentration increases monotonically with depth

Figure 1 shows the layer-wise profile of GPT-2 Large. Entropy is highest in the earliest layers (peak 4.28 nats at layer 1) and declines steadily to ≈ 1.9 nats in layers 13–17 and 29; sparsity rises from ≈ 56% in layer 1 to a peak of 80.8% at layer 14. The same monotone structure appears in gpt2 (sparsity peaks at 88.1%, layer 3) and gpt2-medium.

Figure 1. Layer-wise mean attention entropy (left axis, blue) and mean sparsity (right axis, orange) for GPT-2 Large on WikiText-103 test. Deeper layers attend to fewer keys.

Table 2. Layer-wise metrics for GPT-2 Large (36 layers). Entropy in nats; sparsity % is the fraction of weights ≤ 1e-4; mass columns are cumulative top- $k$ attention mass.

5.3 Attention mass follows a steep power-law tail

Figure 2 shows top- $k$ cumulative mass across layers of GPT-2 Large: in layers 12–33 the top-1 key alone holds ~50-66% of the mass, and the top-5 keys hold ~67–78% of the mass in the same region. A rank-weight power-law fit on a representative deep-layer row yields an exponent $\hat{\alpha} = 1.68$ with $R^2 = 0.956$ (gpt2, layer 5, head 5, query 256) - consistent with heavy-tailed, concentrated attention.

Figure 2. Cumulative mass captured by the top-1, top-5, top-10, and top-50 keys per layer for GPT-2 Large. Mass concentrates sharply with depth.

5.4 The concentration profile is scale-invariant

Figure 3 compares top- $k$ concentration for gpt2 (12 layers) and gpt2-large (36 layers): the curves are nearly identical. Layer-wise entropy and sparsity curves are highly correlated across model scales, and head-level entropy versus sparsity exhibits a strong negative Pearson correlation, with $r$ less than $-0.8$ at every scale consistent with the hypothesis that attention concentration is an intrinsic structural property of pretrained Transformers rather than an artifact of model size.

Figure 3. Top- $k$ cumulative mass for gpt2 (12 layers) vs. gpt2-large (36 layers) on WikiText-103 test. Concentration profiles are nearly identical across a 3× parameter range.

5.5 Numerical integrity

Every reported metric passed automated validation: entropy in [0.009, 5.157] nats, sparsity in [0.59, 99.80]%, density in [0.002, 0.645], with strict top- $k$ monotonicity (top-1 mass is at most top-5 mass, which is at most top-10 mass, which is at most top-50 mass) enforced per row. The microbenchmark path is covered by a unit-test suite and a CI smoke test that verifies artifact generation and correctness checks against the dense SDPA control on CPU and CUDA.

6. Discussion

Our measurements make the implicit assumption behind sparse attention explicit and quantitative:

Sparsity is real and large. Two-thirds to three-quarters of attention weights are numerically negligible (≤ 1e-4). This is direct evidence that thresholded or top- $k$ sparse attention need not discard much probability mass — the top-10 keys already capture ~73–78% of it.
Concentration is systematic, not idiosyncratic. The monotone deepening of concentration (Section 5.2) and its scale invariance (Section 5.4) suggest that layer- and depth-aware policies — e.g., retaining more keys in early layers, aggressive pruning in middle layers — could exploit structure that generalizes across model sizes.
Heavy tails justify rank-based methods. A power-law exponent near 1.7 means a handful of keys dominate; this supports KV-cache eviction heuristics (H2O [10], SnapKV [12]) and attention-sink retention (StreamingLLM [11]).
The benchmark protocol makes claims transferable. By fixing seeds, dense references, warm-up, synchronization, and artifact schemas, RSA-X removes the measurement variance that has made sparse-attention speedups hard to reproduce across papers [21].

We deliberately refrain from claiming that observational concentration causes sparse-attention success: high top- $k$ mass is necessary but not sufficient — dropped mass may still be task-critical. The pre-registered evaluation protocol (Appendix A) is designed to test sufficiency by pairing microbenchmarks with downstream perplexity checks.

7. Limitations and Threats to Validity

Scope. We measured forward-pass attention in three GPT-2-scale models on one corpus. Results may differ for encoder models, instruction-tuned models, longer contexts, or other corpora (e.g., code, multilingual text).
Observational vs. causal sparsity. Concentration is measured on frozen, natural inputs; it does not by itself establish that pruned attention preserves model capability.
Hardware dependence. Absolute latency numbers depend on the GPU, driver, and PyTorch version; the concentration findings reported here are hardware-independent, but the benchmark matrix's timing claims require the pre-registered protocol executed on the target hardware.
Dataset coverage. A single test corpus was used; the dataset-comparison suite (WikiText-2 vs. Penn Treebank) is implemented and validated but its full comparative results are reserved for the expanded matrix.

8. Conclusion and Roadmap

RSA-X provides a rigorous, reproducible instrument for measuring both how concentrated pretrained attention is and how fast sparse backends actually run. Applied to the GPT-2 family, it yields three concrete findings: attention is heavily concentrated (70–75% near-zero weights; top-10 keys hold ~73–78% of mass), concentration deepens monotonically with layer depth (peak sparsity 80.8%, entropy dropping from 4.28 to ~1.9 nats), and the depth profile is scale-invariant across a 3× parameter range, with heavy-tailed (power-law) structure. These results put the empirical foundation of sparse attention on measurable footing and provide a benchmark artifact that future work can extend.

Immediate roadmap. (i) Execute the pre-registered microbenchmark matrix (Appendix A) on T4/P100/A100; (ii) add Triton-based block-sparse and sliding-window candidate backends via the AttentionBackend contract; (iii) run the dataset-comparison and downstream perplexity suites; (iv) publish a public leaderboard with automated schema validation.

Reproducibility, Ethics, and License

Reproducibility. All source code, tests, configuration files, and figure-generation scripts are open source. Benchmark runs use fixed seeds; artifacts embed full environment metadata. The observational dataset (WikiText-103 test) and models (GPT-2 family) are publicly available.

Co-author. Mayank Chaudhary is included as a co-author of this work.

Ethics. This work uses public models and datasets, involves no human subjects, and performs inference only. Energy use is limited to short T4 sessions; we encourage reporting hardware and carbon context alongside benchmark results [25].

License. This paper and its figures are published under the Creative Commons Attribution 4.0 International (CC-BY 4.0) license: you are free to share and adapt with attribution to the author. The software is distributed under the MIT License.

References

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). Attention is all you need. Advances in Neural Information Processing Systems, 30, 5998–6008.
Child, R., Gray, S., Radford, A., & Sutskever, I. (2019). Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509. https://arxiv.org/abs/1904.10509
Beltagy, I., Peters, M. E., & Cohan, A. (2020). Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150. https://arxiv.org/abs/2004.05150
Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Alberti, C., Ontañón, S., Pham, P., Ravula, A., Wang, Q., Yang, L., & Ahmed, A. (2020). Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems, 33, 17283–17297.
Jiang, A. Q., Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., de las Casas, D., Bressand, F., Lengyel, G., Lample, G., Saulnier, L., Lavaud, L. R., Lachaux, M.-A., Stock, P., Le Scao, T., Lavril, T., Wang, T., Lacroix, T., & El Sayed, W. (2023). Mistral 7B. arXiv preprint arXiv:2310.06825. https://arxiv.org/abs/2310.06825
Wang, S., Li, B. Z., Khabsa, M., Fang, H., & Ma, H. (2020). Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768. https://arxiv.org/abs/2006.04768
Kitaev, N., Kaiser, Ł., & Levskaya, A. (2020). Reformer: The efficient transformer. ICLR. https://arxiv.org/abs/2001.04451
Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L., Belanger, D., Colwell, L., & Weller, A. (2021). Rethinking attention with performers. ICLR. https://arxiv.org/abs/2009.14794
Gu, A., & Dao, T. (2023). Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752. https://arxiv.org/abs/2312.00752
Zhang, Z., Sheng, Y., Zhou, T., Chen, T., Zheng, L., Cai, R., Song, Z., Tian, Y., Ré, C., Barrett, C., Wang, Z., & Chen, B. (2023). H2O: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36. https://arxiv.org/abs/2306.14048
Xiao, G., Tian, Y., Chen, B., Han, S., & Lewis, M. (2024). Efficient streaming language models with attention sinks. ICLR. https://arxiv.org/abs/2309.17453
Li, Y., Huang, Y., Yang, B., Venkitesh, B., Locatelli, A., Ye, H., Cai, T., Lewis, P., & Chen, D. (2024). SnapKV: LLM knows what you are looking for before generation. Advances in Neural Information Processing Systems, 37. https://arxiv.org/abs/2404.14469
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and memory-efficient exact attention with IO-awareness. Advances in Neural Information Processing Systems, 35, 16344–16359.
Dao, T. (2024). FlashAttention-2: Faster attention with better parallelism and work partitioning. Transactions on Machine Learning Research. https://arxiv.org/abs/2307.08691
Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., & Dao, T. (2024). FlashAttention-3: Fast and accurate attention with asynchrony and low-precision. arXiv preprint arXiv:2407.08608. https://arxiv.org/abs/2407.08608
Liu, H., Zaharia, M., & Abbeel, P. (2023). Ring attention with blockwise transformers for near-infinity context. arXiv preprint arXiv:2310.01889. https://arxiv.org/abs/2310.01889
Guessous, D., Liang, Y., Dong, J., et al. (2024). FlexAttention: The flexibility of PyTorch with the performance of FlashAttention. PyTorch Blog. https://pytorch.org/blog/flexattention/
Tillet, P., Kung, H. T., & Cox, D. (2019). Triton: An intermediate language and compiler for tile-based neural network computations. Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages (MAPL), 10–19. https://doi.org/10.1145/3315508.3329973
Tay, Y., Dehghani, M., Abnar, S., Shen, Y., Bahri, D., Pham, P., Rao, J., Yang, L., Ruder, S., & Metzler, D. (2021). Long range arena: A benchmark for efficient transformers. ICLR. https://arxiv.org/abs/2011.04006
Tay, Y., Dehghani, M., Bahri, D., & Metzler, D. (2023). Efficient transformers: A survey. ACM Computing Surveys, 55(6), 1–28. https://doi.org/10.1145/3559550
Bouthillier, X., Delaunay, P., Bronzi, M., Trofimov, A., Nichyporuk, B., Szeto, J., Sepah, N., Raff, E., Madan, K., Voleti, V., Kahou, S. E., Michalski, V., Serdyuk, D., Arbel, T., Pal, C., Varoquaux, G., & Vincent, P. (2021). Accounting for variance in machine learning benchmarks. Proceedings of Machine Learning and Systems, 3. https://arxiv.org/abs/2103.03098
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Köpf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., & Chintala, S. (2019). PyTorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems, 32, 8024–8035.
Nanda, N., & Bloom, J. (2022). TransformerLens: A library for mechanistic interpretability of neural networks. https://github.com/neelnanda-io/TransformerLens
Merity, S., Xiong, C., Bradbury, J., & Socher, R. (2017). Pointer sentinel mixture models. ICLR. https://arxiv.org/abs/1609.07843
Dodge, J., Prewitt, T., Tachet des Combes, R., Odmark, E., Schwartz, R., Strubell, E., Luccioni, A. S., Smith, N. A., DeCario, N., & Buchanan, W. (2022). Measuring the carbon intensity of AI in cloud instances. Proceedings of the 2022 ACM Conference on Fairness, Accountability, and Transparency (FAccT), 1877–1894. https://doi.org/10.1145/3531146.3533194

Appendix A. Pre-Registered Benchmark Matrix

Candidate backends will be evaluated across $\\N \in \{128, 256, 512, 1024, 2048, 4096, 8192\}$ , causal and non-causal, in FP32/FP16/BF16, at $B = 1$ , $H = 8$ , $d = 64$ , using the protocol of Section 3.3. Entries remain unpopulated until the matrix is executed on target hardware; the dense SDPA control defines the reference row.