Scaling Attention via Feature Sparsity
Authors: Yan Xie
Tiansheng Wen
Tang Da Huang
Bo Chen
Chenyu You
Stefanie Jegelka
Yifei Wang
Venue: ICLR 2026 Poster (2026)
Keywords: Self-Attention
Sparse Representation
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose Sparse Feature Attention (SFA), which converts dense Q/K into k-sparse codes and computes attention via FlashSFA kernel, preserving near-dense quality while significantly reducing compute, latency, and KV-cache.

Abstract: Scaling Transformers to ultra-long contexts is bottlenecked by the $O(n^2 d)$ cost of self-attention. Existing methods reduce this cost along the sequence axis through local windows, kernel approximations, or token-level sparsity, but these approaches consistently degrade accuracy. In this paper, we instead explore an orthogonal axis: \emph{feature sparsity}. We propose \textbf{Sparse Feature Attention (SFA)}, where queries and keys are represented as $k$-sparse codes that preserve high-dimensional expressivity while reducing the cost of attention from $\Theta(n^2 d)$ to $\Theta(n^2 k^2/d)$. To make this efficient at scale, we introduce \textbf{FlashSFA}, an IO-aware kernel that extends FlashAttention to operate directly on sparse overlaps without materializing dense score matrices. Across GPT-2 and Qwen3 pretraining, SFA matches dense baselines while improving speed by up to $2.5\times$ and reducing FLOPs and KV-cache by nearly 50\%. On synthetic and downstream benchmarks, SFA preserves retrieval accuracy and robustness at long contexts, outperforming short-embedding baselines that collapse feature diversity. These results establish feature-level sparsity as a complementary and underexplored axis for efficient attention, enabling Transformers to scale to orders-of-magnitude longer contexts with minimal quality loss.

Published as a conference paper at ICLR 2026
SCALING ATTENTION VIA FEATURE SPARSITY
Yan Xie'* Tiansheng Wen 1.2- Tangda Huang'
Bo Chen't
Chenyu You
Stefanie Jegelkaš.4 Yifei Wang
School of Electronic Engineering, Xidian University
2 Stony Brook University
3 TUM MIIT "Amazon AGI SF Lab!
ABSTRACT
Scaling Transformers to ultra-long contexts is bottlenecked by the O(n²d) cost of self-attention. Existing methods reduce this cost along the sequence axis through local windows, kernel approximations, or token-level sparsity, but these approaches consistently degrade accuracy. this paper, we instead explore an or- thogonal axis: feature sparsity. We propose Sparse Feature Attention (SFA), where queries and keys are represented as k-sparse codes that preserve high- dimensional expressivity while reducing the cost of attention from (nn2d) @n?* /d). To make this efficient at scale, we introduce FlashSFA, an IO- aware kernel that extends FlashAttention operate directly on sparse overlaps without materializing dense score matrices. Across GPT-2 and Qwen3 pre- training, SFA matches dense baselines while improving speed by up 2.5x and reducing FLOPs and KV-cache by nearly 50%. On synthetic and down- stream benchmarks, SFA preserves retrieval accuracy and robustness long contexts, outperforming short-embedding baselines that collapse feature diver- sity. These results establish feature-level sparsity as complementary and un- derexplored axis for efficient attention, enabling Transformers to scale to orders- of-magnitude longer contexts with minimal quality loss. Code is available at wwt..eowooooeoeeee/sseeeeee VammxlleSaarse-Faauureeeeeee Saaaauss--Faatur--ttttiioo
https://githu..co//
1.0
140
FLOPs
KV-Cache
20
120
Full Attention
0.9
Val Perplexity
Full Altention SFA (ours)
100
(ours)
19.
FLOPs
0.8
Cache
18H
0..
17
40-
0.5
16
10
15
20
25
30
35
40
Base
32
16
0.5
Context Length (128K)
Aclivate Dim (K)
(a) Latency comparison
(b) FLOPs & KV-cache comparison
Figure Overview of proposed method. (a) Trade-off between performance and speed. Compared to directly reducing dimensionality with short embeddings, our method achieves a more favorable balance, delivering 259% speedup over the original dimensionality while improving performance by 21.4% relative to the short-embedding baseline. (b) Computational and memory efficiency comparison. Our method reduces KV-cache memory usage by 41 % and FLOPs by 49%.
* "qqual Contribution.
and Yifei Wang (yifeiwg @amazon.com) 1 This work done MIT prior Yifei Wang joining Amazon.
sporesdn Authors: Bo Chen (bchen

---

Published as a conference paper at ICLR 2026
INTRODUCTION
Scaling language models to ever longer contexts is fundamentally limited by the On cost of self-attention, where is the sequence length and d the feature dimension. Most existing approaches attempt to reduce this cost along the sequence axis. Windowed or low-rank attention variants con- strain interactions achieve linear complexity, while token-level sparsity prunes which tokens teract (Child et al. 2019 Beltagy et al. 2020 (Zaheer et al. 2020 Choromanski et al. [022]] Wang al. 2020; Xiong et al. 2021). Yet large-scale benchmarks consistently show that these approx- imations sacrifice accuracy. leaving dense attention the most reliable option at long ranges. This raises a natural question: rather than reducing the set of tokens, can we explore feature diversity as an orthogonal axis for scaling attention?
This question is motivated by findings in representation learning, where sparse embeddings (For- mal et al. 2021: Wen al. 2025 Guo al. 2026; You et al.] 2025 Duan al. 2024 Xie al.) 2025) show that high-dimensionas spaces encode rich features and that selective activation can pre- serve expressivity while yielding large efficiency gains. If attention itself can be viewed retrieval over feature coordinates, then sparsifying queries and keys by activating only their most salient di- mensions could reduce computation without collapsing representational capacity. The challenge is to realize this idea in practice: how to preserve expressivity while sparsifying, how to implement kernels that benefit from sparsity without materializing the score matrix, and how to adapt pretrained dense models without eroding their quality.
We address these challenges with Sparse Feature Attention (SFA). Instead of dense d-dimensional queries and keys, SFA learns k-sparse codes in which each token activates only handful of coor- dinates. Attention scores are computed solely from overlaps between these supports, reducing the arithmetic of QK from (nndd) to (n22 /d) a fraction (k/d)² of the dense cost while stor- ing only (nkk) nonzeros. To make this efficient at scale, we introduce FlashSFA, new IO-aware kernel that extends FlashAttention by operating directly on sparse overlaps with online softmax. This design avoids materializing any dense n xn scores, retains exactness, and brings compute and memory scaling in line with feature sparsity.
The benefits of this shift are demonstrated in Figure I Compared to simply shrinking hidden size ('short embeddings"), SFA achieves much better trade-off: it improves perplexity by more than 20% while delivering over 2.5x speedup, and reduces FLOPPs by nearly half together with a 41 % drop in KV-cache memory. Experiments confirm that these benefits extend broadly. On GPT-2 and Qwen3 pretraining, SFA matches dense baselines in perplexity and downstream accuracy. On syn- thetic long-context benchmarks such Needle-in-a-Haasttne unseen lengths, while providing consistent latency gains. Crucially, the method is orthogonal to teken-level sparsity and paging, multiplying their benefits by lowering per-interaction cost.
sustains retrieval accuracy across
This work thus establishes feature-level sparsity as a powerful and previously underexplored axis for efficient attention. By leveraging feature diversity rather than compressing it away, SFA pre- serves high-dimensional expressivity while unlocking substantial efficiency gains. Together with FlashSFA, it makes exact long-context attention practical at scale, and paves the way for context windows extended by orders of magnitude without compromising model quality.
2 PRELIMINARIES
Transformers and multi-head attention. Let a sequence of tokens be represented by hidden states X E In x ddeata For each head à E {1....,H) H} with head dimension d, standard scaled
dot- dot-product attention computes:
Qn = XW E =
KA XW K Boxd
Vn. XWX E Baxd
(1)
Q
S₁
E Rnxx
=ostmaax O M) E R"*"
= = E Boxd
(2)
Vd
where M encodes causal or padding masks, and the head outputs are concatenated and projected. The principal cost arises from the dense K, and materialization of Pı; 10-aware kernels (e.g., FlashAttention) compute O, in tiles without forming P, explicitly, minimizing HBM traffic while remaining exact (Dao al. 2022
Dao 2024 Shah al. 2024).
2

---

Published as a conference paper at ICLR 2026
Standard Attention
Sparse Attention
Feature Attention (ours)
X
Key
Suandy
Sparse Feelure Attantion
Query
Attention
N
Altention
Attention
Sparse
Query
Figure 2: Three paradigms of attention. Left: Standard attention computes all N x N query-key interactions in the full feature dimension d. Middle: Sparse attention reduces cost by selecting, for each query each retained interaction still spans all features. Right: Sparse Feature Attention (ours) keeps all tokens but sparsifies along the feature axis by selecting the top-k channels in Q and K (Q Topk, (Q), K = Topk, (K)). Attention is then computed only over overlapping selected features with sparse matrix multiplication. This shifts sparsity from the token axis (N x N) the feature axis, achieving efficiency while preserving token coverage.
small subset of keys 52 and masking the remaining logits before softmax, but
Sparse formats for efficient storage. Sparse matrices that contain only a few non-zero elements can be stored efficiently in sparse formats. Consider a matrix A E Ruxd with anz(A) nonzero elements. In the Compressed Sparse Row (CSR) format, we store three arrays: (i) data E IR containing the values all nonzero entries, (ii) indices E {0.... .......d' d }} |)nz(A), recording the column index of each nonzero, and (iii) indptr E {0, ..., (mzz()))++) where indptr[i] marks the offset in data/indices where row begins. Thus, the nonzeros of row can be read quickly from data data(indptrri]:inn (initrrii:indpr((++++ I]]- The Compressed Sparse Column (CSC) format is analogous, but compresses by columns instead of rows, with an indptr array of length d + (Saad 2003 Davis 2006
sparse matrices, the cost is not proportional to the dense size n x d but rather to the number of structural intersections between the nonzero patterns of rows and columns. This operation, called Sparse General Matrix Multi- plication (SpGEMM), is typically implemented by Gustavson's 's row-wise accumulation algorithm methods (Buluc & Gilbert (00T). The efficiency of SpGEMM therefore depends how many row-column index sets overlap, making CSR and CSC natural formats for storing query and key matrices in our method.
Efficient multiplication with spare formats. When multiplying
(Gustavson] 1978) or by
SPARSE FEATURE ATTENTION
This section introduces Sparse Feature Attention (SFA). a drop-in modification of self- attention that operates along the feature axis. Each query/key vector is converted into a k-sparse code; attention scores are then computed only on overlapping active coordinates. This preserves the probabilistic semantics of exact softmax attention over learned supports while reducing arithmetic, memory traffic, and KV-eache growth.
3.1 ATTENTION VIA SPARSE MATRIX MULTIPLICATION
The key idea of SFA is sparsify the query and key features before attention computation, that only their most salient coordinates contribute to similarity scores. As illustrated Figure (right), (right), given dense projections Q, K, V gaa we apply row-wise Top-k operator to both Q and K:
Q
Topky
K (K),
(3)
where for E 12d
Ius u arg (0, otherwise.
Topkg(x),
(4)
Thus each query and key vector is converted into a k-sparse representation, preserving only its largest-magnitude entries. These Q, K serve as sparse query and key features for attention.
Sparse attention via sparse matrix multiplication. Attention scores then computed as QKT. Instead of full dense multiplication, exploit sparsity: each nonzero in d interacts only
3

---

Published as a conference paper at ICLR 2026
with keys that share the same active coordinate. For query i with support St
Sij
Fi,u kuuu NES,nS,
(5)
which corresponds to sparse matrix multiplication between Q (CSR format) and K (CSC format). Traversing active coordinates yields only the nonzero attention edges. The resulting scores are then passed through the usual softmax and value aggregation steps.
Backward computation. Leveraging the sparse structure, we can also skip computing the gradient for the full query and key matrices at backward computation. Specifically, estimator: gradients flow back only through the selected coordinates. For query i with support S.
straight-through
ac
BC
u EE S.S,
(6)
0q.a
0,
u St,
and similarly for k Both forward and backward passes scale only with the sparse edge set.
Efficiency analysis. Dense attention requires O(n²d) computation and (n22) memory, since every query interacts with every key across all d feature dimensions. In contrast, SFA only forms scores along feature coordinates selected by both queries and keys. Each token activates i: features, giving nk nonzeros in total. Assuming supports are balanced across dimensions, each coordinate is chosen by about deg(u) s nk/d tokens. The number of query-key overlaps contributed by coordinate is then deg(u)", and summing over all d coordinates yields:
deg(u)² dl (nnk ((=))²= n²k²
E
Thus the total cost for attention shrinks from @(n²d) (dense) to ((n" /d) (sparse), which is only fraction k- /d= of the dense cost. Both forward and backward passes then cost O(E + Ed,) FLOPs, and memory for storing query and key drops from ((nd) to O(nk) with the sparse formats. For concreteness, with d 128 and & 16 (default setting considered in this work), the ratio is k2/d²= = 1/64, i.c. about a 64x reduction in theory. As the dimension d increases in larger models, the gain could be even higher. For d 1024 and l = 32 (shown to have very similar retrieval performance in Wen et al. (2025)), the ratio is 32 322/1024* /1024 1//022, i.e., reduction of more than 1000x. This means sparse feature attention can potentially extend context length by one to three orders of magnitude similar compute cost. For example, turning a IM context window into 64M or even 1G, opening up substantial improvements for long-context applications.
3.2 FLASHSFA: FAST SPARSE FEATURE ATTENTION WITHOUT MATERIALIZATION
A key challenge in Sparse Feature Attention (SFA) is that, although we reduce the number of pair- wise interactions from dd to n² n² /d, a naive implementation would still require materializing an score matrix apply the softmax. This would destroy the memory advantage, as the ((n)) storage is often the real bottleneck at long sequence lengths.
Attentonn addressed exactly this issue in the dense case: it avoids storing QKT by processing queries and keys in small tiles, keeping only temporary tile buffer of partial scores on-chip. An online softmax update maintains numerical stability and exactness without ever writing the full x matrix to memory (Dao et al.) 2022). FlashAttention-2 and -3 extend this idea with more parallelism and precision refinements (Dao) 2024; Shah et al., 2024). Our proposed FlashSFA extends this machinery of FlashAttention, but replace dense tile multiplications with sparse feature-intersection kernels. For tile of queries (rows i E fio, + B,)) and keys (columns j E [jo. these tokens, intersects their supports, and performs scatter-adds into a compact B, x B. score buffer. This buffer is immediately consumed by the online softmax update, large score matrix is ever written memory. The result is mathematically identical to computing softmax( QK /VdVV, but with both compute and memory scaling as in SFA.
principle to SFA. We retain the 10- aware tiling and
the kernel iterates over the active features of
Efficiency and design. FlashSFA inherits the same (nn) I0 complexity of FlashAttention, since only tiles (not the full matrix) touch bigh-bandwidth memory. Within each tile, the work is pro- portional to the number of overlapping features rather than d, yielding the O /d) complexity
4

---

Published as a conference paper at ICLR 2026
Table 1: Perplexity and Accuracy results. Dense baselines use full hidden size and uncompressed KV cache; "Dense (d=X)" denotes short-embedding baselines with reduced feature dimension. PPL. evaluated on OpenWebTexx for GPT-2 and Pile for Qwen3. Note that "Dense (full)" serves as reference upper bound; we highlights the best results among the sparse/compressed baselines.
Latency 4
PPI.
Acc
Model
128k context
OWT/Pile
PIQA
LAMBADA
ARC-e
ARC-c
HellaS
Avg-Ace
Dense (full)
16.86
17.29
56.34
22.78
28.35
14.32
19.61
28.28
Dense (d 32)
7.86
20.88
51.30
19.39
25.72
12.47
14.26
24.63
GPT2-124M
SFA (k = 8)
9.41
18.27
54.92
21.03
28.41
7.39
19.26
27.40
Dense (full)
46.78
15.03
59.79
24.74
30.19
15.78
22.04
30.51
Dense (d 32)
20.58
19.89
55.17
19.96
28.15
11.83
18.43
26.71
GPT2-350M
SFA (A = 8)
23.67
16.78
58.02
23.83
30.22
13.66
22.13
29.57
Dense (full)
77.65
4.66
62.47
34.82
45.41
20.35
33.95
39.40
Dense (d = 64)
30.84
6.03
58.43
31.27
41.58
15.83
28.29
36.68
Qwen3-0.6B
SFA (k:
34,20
4.81
61.73
34.05
45.62
19.27
34.03
38.94
analyzed in $3.1 The online softmax logic, masking for causality, and streaming of are un- changed. Indices for sparse features add modest overhead (O(nk)), and can be stored efficiently with 16-bit integers for typical d S 65,535.
By marrying the sparsity of SFA with the memory-efficient tiling of FlashAttention, FlashSFA achieves the best of both worlds: avoids Onn²) materialization while preserving the z reduc- tion in arithmetic and memory cost. This enables exact attention with dramatically lower compute and memory footprints, making long-context training and inference practical at scale. We defer full description of the FlashSFA algorithm to Appendix C
EXPERIMENTS
4.1 PRETRAINING EXPERIMENTS
Having introduced Sparse Feature Attention (SFA) and the FlashSFA kernel, we next examine whether autoregressive LMs trained from scratch maintain modeling quality under feature spar- sification. We evaluate GPT-2 and Qwen3 models against dense and short-embedding baselines, measuring both modeling quality and efficiency.
Models and baselines. We study GPT-2 Small/Medium (Radford et al., 2019) and Qwen3-0.6B Yang et al. 2025), replacing dense QKT scoring with SFA while keeping V dense. Sparsity budgets l E {8, 16} are tested. Baselines include standard dense attention and short-embedding variants (halving the hidden size of Q/K). Note that "Dense (full)" serves reference upper bound; we highlights the best results among the sparse/compressed baselines. We use the RTopK kernel (Xie et al. 2024) for efficient topk operations. Additional implementation details, including model contigurations and handling of RoPE dimensions in Qwen3, are deferred to Appendix A.
Datasets and benchmarks. GPT-2 models are trained on Open WebText (Gokaslan & Cohen 2019 Qwen3 on The Pile (Gao et al. 2020 Biderman et al.
2022). We report validation perplexity (PPL), zero-shot accuracy on PIQA Bisk al.] (020), LAMBADA (Paperno al. 2016)) ARC-e/ARC-
e (Clark et al., 2018), and HellaSwag (Zellers et al. 2019). as well as decoding throughput at 128k tokens (Speed3 28k) to assess long-context efficiency.
GPT-2 results. Table I shows that SFA with k = 16 (not shown here but consistent with h 8 trends) closely tracks dense baselines, with negligible differences in perplexity and accuracy. SFA with A 8 incurs slightly higher PPL and minor accuracy drops, but these remain within accept- able bounds. This demonstrates that sparsified features preserve most of the model's expressive capacity. By contrast, short-embedding baselines degrade more substantially: they reduce perplex- ity efficiency and underperform on challenging tasks such as ARC-c, especially for GPT-2 Small. While such baselines deliver higher throughput, their quality-efficiency balance is skewed toward
5

---

Published as a conference paper at ICLR 2026
Table 2: Long context pretraining results. Comparison of NIAH accuracy rates for different lengths under various training lengths. (a) Models are trained on 8k synthetic NIAH data, and the aceuracy rate on test lengths from 1k to 8k. (b) Models are trained on 32k synthetie NIAH data, and the accuracy rate test lengths from 1k to 32k.
(a) NIAH accuracy (%) within 8k Sequence Length.
(b) NIAH accuracy (%) within 32k Sequence Length.
Context Length
Context Length
Method
1k
2k
4k
8k
Speed 8k
Method
1k
4k
16k
32k
Speed @ 32k
Dense (d 64)
94%
93%
90%
95%
1.0x
Dense (d 64)
92%
94%
83%
80%
1.0x
SFA (k: 2)
95%
95%
97%
98%
1.9x
SFA (k 8)
95%
94%
83%
82%
1.3x
SFA (A= 8)
98%
100%
99%
98%
1.3x
SFA (k 16)
97%
96%
83%
83%
1.0x
200
Daros out (Sum
Doras stecorn
3k
Denue Hilt) jurrd
100-
Demue
Latency (ms)
150
Latency (ms)
400
80-
2k
100
Latency
Latency
60
200
1k:
20-
0
32
16
32
16
0k
32
16
0
32
16
8 Sparsity(k)
Sparsity(k)
Sparsity(k)
Sparsity(k)
(a) Dot-Product only
(b) Attention Block
(c) FlashAttention Block
(d) Entire Transformer
Figure 3: Latency vs. feature sparsity. Latency Comparison of dense attention and SFA (ours) at different modular levels in Transformers under 16k context length. Higher sparsity brings substantial decrease in latency.
speed, making them less appealing. On retrieval-like tasks (LAMBADA, HellaSwag), sparse models underperform relative their PPL, motivating further Eetrieval-focusee experiments (Section 4.2).
Qwen3 results. For Qwen3-0.6B. 0.6.....BB.BB. also in Table II SFA with k = 8 maintains perplexity nearly identical to dense (4.81 vs. 4.66) and preserves accuracy across PiQA, ARC-e, and HellaSwag. The small differences on ARC-c (19.27 vs. 20.35) and average accuracy (38.94 vs. 39.40) suggest only a marginal quality cost. Short-embedding baselines again degrade more severely, with higher PPL. (6.03) and lower accuracy (Avg-Ace 36.68). This confirms that even in modern architectures with RoPE and normalization refinements, sparsified features remain competitive with dense attention, while offering clear efficiency benefits at long context.
Efficiency results. Across GPT-2 and Qwen3, short-embedding variants provide the largest raw speedups due to narrower hidden size, but their accuracy loss makes them less practical. Sparse models present a more balanced trade-off: k = 16 maintains baseline-level quality, and k: = 8 provides moderate speedups while remaining close in accuracy. In practice, k emerges as the most attractive setting, balancing efficiency and modeling quality. This setting is therefore used in subsequent scaling and efficiency benchmarks (Section 4.3)
4.2 SYNTHETIC NIAH EXPERIMENTS
The synthetic Needle-in-a-Hayeer (NIAH) benchmark provides controlled way to examine how models handle extremely long contexts and retrieval-style reasoning. To further examine whether sparse attention preserves retrieval capacity over long contexts, we conduct experiments on the syn- thetic NIAH task. Following the RULER methodology, haystacks are constructed by repeating the character "4" and inserting a single target "needle" token that the model must recover. We train GPT-2 models (124M) from scratch on synthetic NIAH QA data under two training regimes: one restricted to 8k contexts and one extended to 32k contexts. In both cases, we then evaluate test accu- racy across multiple held-out lengths, measuring how well models generalize beyond their training window. Speed is also measured at the maximum training length to capture efficiency.
Results within 8k. Table 2a reports results when models trained up to 8k tokens. Dense base- lines perform well across all lengths but incur standard compute costs. Sparse models not only match but slightly exceed dense accuracy, achieving near-perfect recovery at all test lengths. At the

---

Published as a conference paper at ICLR 2026
Exres DuC
Boras tNd Stecaro
60k
Samte Hull
Latency (ms
30
Latancy (mS
200
7k
Latency (ms)
20:
150-
Latency
5k-
40k
20k-
50
2k
0
32
16
8 Sparsity(k)
32
16
32 0k
16
8 Sparsity(k)
0k:
32
Sparsity(k)
Sparsity(K)
(a) 128 dim (4k context) (b) 256 dim (4k context) (c) 128 dim (65k context) (d) 256 dim (65k context)
Figure 4: Latency vs. feature sparsity with various config, Latency Comparison of dense attention and SFA (ours) at different head dimensions and context lengths. Notably, the latency of SFA can be much lower than dense attention under high dimension per head and long context, e.g., Figure 4d
8k
Dense (ful))
125
SFA (km4)
(full)
(full)
Latencymi
6K-
SFA (K-8)
SFA (km4)
Latency(ms
SFA (k-8)
KVCache((Mie)
SFA (Kc B]
SFA (kmm)
4k
SFA (K-16)
75
SFA (6))
1.ck
(Ic-16)
50
25
500-
1.k
32k
Leq
SeqLen
Sequen
(a) Latency vs context length
(h) FLOPs vs context length
(c) KV-cache vs context length
Figure 5: Scaling dense attention and SFA with context length. SFA can consistently reduce both the computatin cost and KV cache size by a constant factor of at least 2.
same time, SFA delivers 1.9x decoding speedup at 8k for k 2, confirming that sparse scoring reduces computation without sacrificing reliability.
Results within 32k. Table 2b extends training to 32k tokens. Dense baselines degrade as length grows, dropping to 80% accuracy at 32k. SFA models maintain higher accuracy: h = holds steady at 82% and A = 16 at 83%. Notably, k 8 delivers 1.3x faster generation at 32k, while = 16 matches dense throughput. These results show that sparse attention generalizes robustly across unseen lengths while simultaneously reducing long-context latency.
The NIAH task isolates retrieval in a controlled setting, making it possible to compare dense and sparse features without confounding factors. Across both 8k and 32k training regimes, SFA preserves or improves accuracy while achieving consistent speedups. This complements the pretraining results in Section 4.1 sparse attention does not erode retrieval ability, and under syn- thetic stress tests it can even provide stronger length generalization than dense attention.
Discussion.
4.3 BENCHMARKING COMPUTATION AND MEMORY EFFICIENCY OF SFA
We benchmark Sparse Feature Attention (SFA) in both training and inference scenarios, since they stress different system bottlenecks. Training- -tmm attention is dominated by quadratic computation, while inference-ting run on an A800 GPU with CUDA 12.4, using INT32 for i indptr and INT8 for ind ces, FP16 for va 1 ves, and pinned batches in HBM. Timing excludes dataloader overhead. All kernels are compiled with CUDA and Ninja, and report medians over 50 warm runs. We built our FlashSFA kernel upon LecICUDA
attention with KV cache is dominated by memory traffic. Experiments are
Influence of SFA in Transformers. Figure 3 compares latency of SFA and dense attention across different modular levels of Transformer: from the raw dot-product to the full model. As sparsity increases (smaller k), latency drops significantly. Importantly, the benefit compounds with complex- ity: while doo-product alone shows modest gains, the full Transformer achieves over 2x reduction. This demonstrates that sparsity scales well when applied throughout the network stack.
7

---

Published as a conference paper at ICLR 2026
Table 3: Evaluation on general reasoning tasks and synthetic retrieval (NIAH). Accuracy is in %.
General Tasks
NIAH Acc
Model
Variant
GSM-8K
Arxiv
PubMed
4096
8192
16384
32768
Base
59.59
13.65
10.48
90
87
77
52
Qwen3-0.6B
Fine-tune
63.42
41.17
40.54
SFA (k: 16)
61,46
39.14
39.03
95
93
77
53
Base
75.44
31.52
29.19
Qwen3-4B
Fine-tune
76.18
49.31
49.05
99
96
92
84
SFA (k: = 16)
75.56
46.28
47.91
99
93
91
84
Base
87.62
40,13
37.22
100
100
97
92
Qwen3-8B
Fine-tune
89.11
54.26
55.07
SFA (k: 16)
87.99
52.74
52.61
100
100
100
97
Influence of Dimension and Context Length. Figure examines latency under varying head di- mensions (128 vs. 256) and context lengths (4k vs. 65k). At shorter contexts (4k tokens), SFA offers consistent but moderate gains. However, under long contexts (65k tokens) and larger head sizes (256 dim), the improvement is dramatic: SFA reduces latency by more than an order of magnitude. This confirms that sparsity is most effective in the large-scale regime, where dense attention becomes prohibitively expensive.
Latency and Memory Scaling at Inference. Figure benchmarks autoregressive inference with KV cache. For short contexts (<4k), dense attention remains competitive because sparse kernels incur lookup overhead. Beyond 8k-16k tokens, however, SFA consistently outperforms dense atten- tion. Moreover, SFA reduces VV-cache size proportionally to sparsity, saving up to 440% memory at 4. This makes sparse features especially valuable for long-context inference, where memory footprint is often the limiting factor.
Together, these results show that SFA addresses both compute and memory bottlenecks. During training, it accelerates high-dimension, long-context workloads by cutting FLOOs; during inference, reduces both latency and K V-cache usage for long sequences. These complementary benefits make SEA well-suited for scaling LLMs ultra-long contexts. More results are shown in Appendix E
EXPLORING SFA ADAPTATION WITH PRETRAINED LLMS
In addition to incorporating SFA during the pretraining stage, we also attempted to adjust models with dense pretraining to a sparse feature attention pattern through fine-tuning. In this section, we explore the use of SFA in fine-tuning.
Regularized Sparse Finetuning. During finetuning, we keep SFA consistent with our strategy in the pre-training phase (Eqs.) 3 (). Nevertheless, the sparsification of pretrained dense features in- troduces a severe distribution shift for the pretrained model. Therefore, we regularize the finetuning with additional MSE loss such that SFA's attention scores approximate that of dense features. Since FlashAttention and FlashSPA do not materialize the full attention matrix, in practice we ap- proximate the denese attention output On (with stop gradient) with SFA's attention output Ön at each head h, leading to the final finetuning objective:
H El stopgrad(O:) (8) HI
1
C LLM + Creg Ezz,) log poly xz; S,V)
Datasets. To comprehensively evaluate the performance of SFA during fine-tuning, we conduct
experiments using mathematical tasks, document question answering, and long-context retrieval tasks. We use GSM-8K (Cobbe et al. 202I). Sci-papers (Arxiv and PubMed (Cohan et al., 2018 ) and NIAH data constructed from real texts, respectively. Because applying TopK to the features almost resets the pattern of the previous dense features, we first restore the model's 's language ability by training similar reasoning dataset, MWP-200k (Mitra et al. 2024), before GSM-8K. For the
8

---

Published as a conference paper at ICLR 2026
NIAH data, we use the Pile dataset haystack for random filling. The size of the training set is set to 100k, and 100 test data entries for each length in the test set.
Training Settings. We fine-tune Qwen3-0.6B and Qwen3-4B using Llama-Factory (Zheng et al,, 2024) with = 16 for SFA. For mathematical reasoning and science QA tasks, the training con- text length is set to 16,384 tokens, while for long-context retrieval tasks it is set to 32,768 tokens, with evaluation spanning 4k-32k contexts. All models are trained for three epochs with identical hyperparameters. Detailed experiment setting can be found in Section A.2
Result Analysis. Table 3 compares the base model, dense fine-tuning, and our Top-16 variant. On general tasks, dense fine-tuning yields large gains Arxiv and PubMed by adapting to the evaluation format, and SFA closely tracks these improvements, showing that sparsified features preserve document-comppreeeniion signals even under hard k. On GSM-8K, Top-16 lags slightly behind dense fine-tuning, indicating that arithmetic reasoning is more sensitive to pruning. For long-context retrieval (NIAH), Top-16 performs nearly identically to dense fine-tuning, consistent with Section 4.2] suggesting that sparse supports provide an effective inductive bias for locality. At the 4B scale, Top- 16 remains within 1-3 points of dense on general tasks and holds parity on NIAH, confirming its robustness and compatibility with larger backbones.
RELATED WORK
Token-level sparsity. Many approaches reduce the quadratic cost by pruning which tokens interact.
Structured patterns (loca//strdded/llba)) and learned routing yield strong long-context performance:
Sparse Transformers (Child et al. 2019), Longformer and BigBird (Beltagy et al.) 2020 Zaheer
al. 2020), Routing Transformers (Roy et al. 2021), and Reformer (Kitaev et al.] 2020). Recent
inference systems dynamically select salient tokens or pages (H₂O, SnapK V, Quest) (Zhang et al.)
2023 Li et al.) 2024] Tang et al., 2024). These methods are orthogonal to ours: they sparsify the of tokens, while we sparsify the feature coordinates used to score any retained token pair. In practice, SFA composes with token sparsity and paging by shrinking the per-interaction cost.
Low-rank/kernel approximations vs. feature sparsity. A parallel line alters the operator to achieve linear or near-linear time via low-rank or kernel approximations: Linformer projects K, V (Wang et al. 2020); Performer approximates softmax with random features (Choromanski et al.) 2021] F Nyströmformer uses landmark decompositions (Xiong et al. (022). These compress infor- mation into dense r < d space, often trading expressivity for speed. By contrast, SFA keeps the high-dimennsional feature space but activates only k dl learned coordinates per token; attention scores are computed exactly over the overlap of active supports (no kernel surrogates). This is closer in spirit to sparse coding and sparse embeddings (e.g., SPLADE; CSR) that preserve semantic detail
while enabling unverted-inder efficiency (Formal et al. 2021 Wen et al. 2025).
Efficient attention kernels and sparse representations. FlashAttention reorders computation and I0 to keep attention exact while minimizing off-chip traffic (Dao et al. [022] Dao 2024] Shah
et al. 2024); systems like xFormers and flashinfer expose page/block sparsity primitives (Lefaudeux
et al. 2022 Ye et al. 2025). Some works use feature cues to drive token selection atop such kernels (c.g., SPAR-Q: LoKI) (Ribar al. 2024 Singhania et al. 2024). SFA differr by learning sparse Q/K codes as first-class representations and introducing an 10-aware kernel (FlashSFA) that iterates intersections of active coordinates rather than dense d-dimensional products, yielding arithmetic and bandwidth savings proportional to and composing naturally with loken-sparse routing. Our focus is thus complementary: we open the underexplored axis of feature-level sparsity inside attention while remaining compatible with token-level sparsity and paging.
CONCLUSION AND LIMITATIONS
We presented Sparse Feature Attention (SFA), a new approach to scaling long-context Transform- ers through dimension-level sparsity. By learning sparse query/key codes and computing attention via feature overlaps, SFA preserves high-dimensional expressivity while reducing both memory and compute. We introduced two adaptation strategies (end-to-end Top-k k finetuning and udapter-based training) and an 10-aware FlashSFA kernel that integrates sparsity directly into the online- online-softmax softmax pipeline. Experiments across synthetic and real tasks show that SFA achieves comparable qual-

---

Published as a conference paper at ICLR 2026
ity to dense attention with growing efficiency gains at longer contexts, and complements existing token-level sparsity methods.
While promising, several aspects remain open. Sparse tensor products require stronger support from GPU hardware and CUDA libraries to fully unlock their efficiency, though these system-level challenges are likely be resolved over time. Very sparse query/key codes can lead occasional quality degradation, suggesting the need for adaptive sparsity budgets. Finally, how to best com- bine token-level and dimension-level sparsity remains an exciting direction, offering the possibility of compounding gains in both compute and memory. We view SFA as first step toward explor- ing this new axis of sparsity attention, and hope it motivates further work the intersection of representation learning, attention design, and efficient systems.
Ethics Statement. This work complies with the ICLR Code of Ethics. Our research primarily utilizes publicly available datasets and pretrained models, and we do not foresee any direct negative societal impacts or ethical concerns arising from our methodology.
Reproducibility. We provide detailed descriptions of our methodology, datasets, model configu- rations, and evaluation metrics in the main text and Appendix. Upon acceptance, we will release source code and scripts to enable full replication of our experiments.
ACKNOWLEDGEMENTS
Yan Xie, Tiansheng Wen, Tangda Huang, and Bo Chen were supported in part by the National Natural Science Foundation of China under Grant 62576266: in part by the Fundamental Research Funds for the Central Universities QTZX24003 and QTZX23018; in part by the 111 Project under Grant B18039. Yifei Wang and Stefanie Jegelka were supported part by the NSF Al Institute TILOS (NSF CCF-2112665), and an Alexander von Humboldt Professorship.
REFERENCES
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint urXiv:2004.05/54 2020.
Stella Biderman, Kieran Bicheno, and Leo Gao. Datashect for the pile. arXiv preprint 2rXIv:2201.07311 2022.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical com- monsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Aydin Buluc and John R. Gilbert. The combinatorial blas: design, implementation, and applications. IJHPCA, 25(4):99-5509 2011.
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904. 10509, 2019.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, al. Rethinking attention with performers. In International Conference on Learning Representations, 2021.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint
2018.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jemy Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint 1rXiv:2110.14166 14168, 2021.
10

---

Published as a conference paper at ICLR 2026
Arman Cohan, Franck Dernoncourt, Doo Soon Kim, Trung Bui, Seokhwan Kim, Walter Chang, and Nazli Goharian, A discourse-awan attention model for abstractive summarization of long docu- ments. Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Techmologies, Volume 2 (Short Papers), 2018. doi:
10.18653/v1/118- 2097. URLhttp: /x doi. org. 111111..66.100
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In The Twelfth International Conference on Learning Representations, 2024.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances neural information processing systems, 11116344-63399 2022.
Timothy A Davis. Direct methods for sparse linear systems. SIAM, 2006.
DefTruth and Many Others. Leetcuda: A modern cuda learn notes with pytorch for beginners, 2025. URL htttts::///ttiubbbbb.bb wwtp://.oooocccmmo CUDD, available at https///ttttb..com///xcccccceeet
Open-source software
Zhibin Duan, Tiansheng Wen, Yifei Wang, Chen Zhu, Bo Chen, and Mingyuan Zhou. Contrastive factor analysis, 2024. httts:tt::///ttt wrg. abs/ 2407. 21740
Thibault Formal, Carlos Lassance, Benjamin Piwowarski, and Stéphane Clinchant. SPLADE v2: Sparse lexical and expansion model for information retrieval. arXiv preprint arXIv:2109.10000 2021.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, al. The Pile: An 800GB dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00022 2020.
Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. http::/skk/Lio0077..gg....
io/
@ppenWebTextCorpuu 2019.
Lixuan Guo, Yifei Wang, Tiansheng Wen, Yifan Wang, Aosong Feng, Bo Chen, Stefanie Jegelka, and Chenyu You. Csrv2: Unlocking ultra-sparse embeddings. 2026. URL https://arxivv
org/abs/2602. 05735
Fred Gustavson. Two fast algorithms for sparse matrices: Multiplication and permuted transposi- tion. ACM Transactions on Mathematical Software (TOMS), 4(3):250-269, 1978.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020.
Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, Daniel Haziza, Luca Wehrstedt, Jeremy Reizenstein, and Grigory Sizov. xformers: modular and hackable trans- former modelling library. 2022.
Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Lim knows what you are looking for before generation. Advances in Neural Information Processing Systems, 37:22947-229707 2024.
Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, al. Deepseek-v2: strong, economical, and efficient mixture- of-experts language model. arXiv preprint arXiv:2405.04434 2024a.
Zechun Liu, Barlas Oguz, Changsheng Zhao, Ernie Chang, Pierre Stock, Yashar Mehdad, Yangyang Shi, Raghuraman Krishnamoorthi, and Vikas Chandra. LIm-qat: Data-free quantization aware training for large language models. In Findings of the Association for Computational Linguistics: ACL 2024, pp. 467-484 2024b.
Arindam Mitra, Hamed Khanpour, Corby Rosset, and Ahmed Awadallah. Orca-math: Unlocking the potential of slms in grade school math. arXiv preprint arXiv:2402. 14830, 2024.
11

---

Published as a conference paper at ICLR 2026
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc-Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th annual meeting of the association for computational linguistics (volume I: Long papers), pp. 1525-1534, 2016.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models are unsupervised multitask learners. 2019,
Luka Ribar, Ivan Chelombiev, Luke Hudlass- Gallsy, Charlie Blake, Carlo Luschi, and Douglas Orr. Sparq attention: bandwidth-efficien IIm inference. In Proceedings of the 41st International Conference on Machine Learning, pp. 42558-42583, 2024.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguis- lics, 9:53-68, 2021.
Yousef Saad. Iterative methods for sparse linear systems. SIAM, 2003.
Jay Shah, Ganesh Bikshandi. Ying Zhang, Vijay Thakkar. Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. Advances in Neural Information Processing Systems, (1666688-66888 2024.
Prajwal Singhania, Siddharth Singh, Shwai Hc, Soheil Feizi, and Abhinav Bhatele. Loki: Low- rank keys for efficient sparse attention. Advances in Neural Information Processing Systems, 37: 16692-16723, 2024.
Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. Quest: query- aware sparsity for efficient long-context I1m inference. Proceedings of the 41st International Conference on Machine Learning, pp. 47901-47911 2024.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04766 2020.
Tiansheng Wen, Yifei Wang, Zequn Zeng, Zhong Peng, Yudi Su, Xinyang Liu, Bo Chen, Hongwei Liu, Stefanie Jegelka, and Chenyu You. Beyond matryoshka: Revisiting sparse coding for adap- tive representation. In International Conference on Machine Learning, pp. 16520-665536 PMLR, 2025.
Xi Xie, Yuebo Luo, Hongwu Peng, and Caiwen Ding. Rtop-k: Ultra-fast row-wise top-k selection for neural network acceleration on gpus. In The Thirleenth International Conference Learning Representations, 2024.
Yan Xie, Zequn Zeng. Hao Zhang, Yucheng Ding, Yi Wang, Zhengjue Wang, Bo Chen, and Hong- wei Liu. Discovering fine-grained visual-concept relations by disentangled optimal transport concept bottleneck models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 30199-30209, 2025.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström- based algorithm for approximating self-attention. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pp. 14138-14148, 2021.
An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao. Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint M.Xiv:2505.09388. 2025.
Zihao Ye, Lequn Chen, Ruihang Lai, Wuwei Lin, Yineng Zhang, Stephanie Wang, Tianqi Chen, Baris Kasikci, Vinod Grover, Arvind Krishnamurthy, et al. Flashinfer: Efficient and customizable attention engine for Ilm inference serving. Proceedings of Machine Learning and Systems, 7, 2025.
Chenyu You, Haocheng Dai. Yifei Min, Jasjeet S Sekhon, Sarang Joshi, and James Duncan. Un- covering memorization effect in the presence of spurious correlations. Nature Communications, 16(1):5424, 2025.
12

---

Published as a conference paper at ICLR 2026
Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, Yuxing Wei, Lean Wang. Zhiping Xiao, et al. Native sparse attention: hardware-aligned and natively trainable sparse attention. In Proceedings of the 63rd Anmual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 23078-23097, 2025.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang. et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283-17299 2020.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? Proceedings of the 57th annual meeting of the association for computational linguistics, pp. 4791-4800, 2019.
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient gen- erative inference of large language models. Advances Neural Information Processing Systems, 36:34661-34710. 2023.
Yaowei Zheng, Richong Zhang, Junhao Zhang, Yanhan Ye, and Zheyan Luo. Llamafactory: Unified efficient fine-tuning of 100+ language models. In Proceedings of the 62nd annual meeting of the association for computational linguistics (volume 3: system demonstrations), pp. 400-410, 2024.
13

---

Published as a conference paper at ICLR 2026
A
ADDITIONAL EXPERIMENTAL DETAILS
A.1
PRETRAINING SETUP
Model configurations. Table lists detailed configurations of GPT-2 and Qwen3 models, including
parameter counts, hidden dimensions, number of layers/heads, and short-embedding baselines.
Size
iParameters
hidden size
num layers
num heads
short hidden
position embedding
Small
124M
768
384
APE
Medium
350M
1024
512
APE
Large
596M
1024
512
RoPE
Table 4: Base model configurations. "Short" refers to halving the hidden size for Q/K.
Implementation notes. For fairness, short-embedding baselines insert only linear projections be- fore and after attention. For Qwen3, we add an extra linear transformation after RoPE to isolate positional dimensions from sparsification. FlashSFA kernels are used for tiled execution.
Training. GPT-2 models trained OpenWebText and Qwen3 The Pile with standard LM objectives. Validation PPL reported on held-out splits. Zero-shot evaluations follow PiQA, LAM- BADA, ARC-e/ARC-e, and HellaSwag. Long-context efficiency is measured decoding through- put at 128k tokens.
A.2 FINE-TUNING SETUP
Table 5: Configurations for fine-tuning MoE models
Model
Dataset
Epoch
Batch_Sizz Size
Lr
Warmup_ Ratio Gradient Gradient_Ckeepoinne Ckecpointing
GSM8K
256
6e-4
0.1
False
Arxiv
256
le-5
0.05
False
Qwen3-0.6B
PubMed
2
256
2e-5
0.05
False
NIAH
256
2e-5
0.05
False
GSM8K
256
6e-6
0.1
True
Arxiv
2
256
2e-6
0.1
True
Qwen3-4B
PubMed
256
2e-6
0.1
True
NIAH
256
2e-6
0.1
True
B ADDITIONAL EXPERIMENTS
B.1 LATENCY
We benchmarked the latency of the attention module on three feature dimensions: 256, 128, and 64, respectively.
Prefilling Latency The computational complexity of the full attention module can be expressed (n)dd). So we can express Latency Latencyattn 22d. To better analyze the impact of the feature dimension on computational complexity, we conduct the analysis in the logarithmic space while fix Batch and Heaads =
Log(Leten)).tn)) 2log logd
(9)
The results in the logarithmic coordinate system shown in Figure 6 we can observe that the latency generally exhibits a linear relationship with the sequence length. Furthermore, the latency gap between different compression ratios are close to constant value in the logarithmic space, which also indicates that the absolute efficiency improvement achieved by compressing the feature dimension increases exponentially with the sequence length.
As shown in the figure,
14

---

Published as a conference paper at ICLR 2026
iTaeh 128
Sparno, 16/121
Sparne 8/64
Spares 111331
sparte_4/64
Beese. 32/120 Sespppree1///// 1ttp:e11111---------- Spw.ama
Desse 16/01e116/61
SSparsc./oo
4000
Denox
12
Faster
Faster
Deron, 1211
Denoi
02
lond
4096
8192
16984
32768
65536
(a) Latency in full attention scenario.
(b) Latency in kv-cache attention scenario.
Figure 6: Comparison of latencies in different attention scenarios: (a) full attention(Time-to-fiiss
token) and (b) kv-cache attention Caattion(Timmtueee fimee-o--cxt-tokee
KV-cache Latency KV cache, which has been widely used in LLM decoding, is known as memory bound task. Therefore, benchmarked the inference latency and KV cache memory usage of sparse features and dense features in the KV cache decoding scenario, while keeping Batch = and Heads 8 unchanged. Since setting the length of the Query N, complexity of the decoding attention can be expressed as O(Nd), which means that as the sequence length increases, the computational complexity grows linearly. Our experimental results confirm this. As shown in Table. 2, sparse attention becomes increasingly advantageous the context length grows. At short sequences (e.g., 4k tokens), dense attention still competitive or even faster because sparse kernels pay overhead for index lookups and binary searches. However, once the context exceeds about 8k-16k tokens, the sparse variants consistently overtake the dense baselines.
the computational
B.2 FLOPS
To further analyze the operation of full-attention, we separately counted the number of floating-poin operations (FLOPs) and integer operations (INOPs) under different settings.
Table 6: Operation counts for standard flash attention and flash attention sparse. The number of floating-poit operations (FLOPs) and integer operations (INOPs) were counted separately for feature dimensions of 64 and 128 under different context lengths.
8192
16384
32768
65536
Config
TFLOPs
INOPs
TFLOPs
INOPs
TFLOPs
INOPs
TFLOPs
INOPs
Dense_ _28
2.23
8.92
/
35.67
142.67
Sparse_32/128
32/128
1.20
28.31
4.79
58.72
19.17
121.63
76.70
251.67
Sparse 16/128
1.15
18.87
4.59
39.85
18.35
83.89
73.40
176.17
Sparse_80 8/128
1.13
13.63
4.54
29.36
18.14
62.91
72.57
134.22
Dense 64
1.12
4.48
17.94
71.75
$pa6se_16/64
0.61
15.16
2.42
29.36
9.69
60.82
38.76
125. 83
Sparse 8/64
0.58
9.44
2.32
19.92
9.27
41.94
37.11
88,08
Sparse 4/64
0.57
6.82
2.30
14.68
9.17
31.46
36.70
67.11
As shown in Tab(6] because we directly reduced the number of non-zero elements in the feature vectors, the number of floating-poin operations has significantly decreased, and a large proportion of the doating-poin operations in the sparse version come from matrix multiplication in the P@ stage. The reason is that sparse feature attention converts large number of FLOPs into the process of finding overlapping non zero elements in sparse matrix multiplication, which corresponds to the INOPs in the table.
15

---

Published as a conference paper at ICLR 2026
C IMPLEMENTATION DETAILS OF FLASHSFA KERNEL
In this section, provide comprehensive breakdown of the FlashSFA CUDA kernel implemen- tation, focusing on the parallelism hierarchy, core data structures, memory access patterns, and detailed complexity analysis of the sparsification overhead. Since current GPUs do support gen- eral sparse matrix multiplication well, for fair comparison, we compared Flash Attention Sparse with FMA-based Dense Flash Attention on the code base of Flash Attention 2 in the LeetCUDA open-source library (DefTrrth & Others 2025).
C.I ALGORITHM PROCEDURE
Algorithm FlashSFA (forward with tile (B, x Be))
Require: CSR(():
Q_indptr, Q Q_indices, values; CSCrea(()) Kf Kf_indptr, Kf if_indices, Kf_ values; V (dense, row-major in HBM); tile offsets (on; tile sizes (B,, B.).
1: Init score tile storage: scores zeros B,, Be) in SRAM.
2: Init CSR(P) row pointers: P indptr(fol <- current nnz counter tp. D b in + r is the global query index
3: for 0 to B, do
4:
f in +
5: Register accumulator: row ww_sccores(0:AA.
kept in registers per thread/warp
6:
t, + Q_indptr(i]] fR - Q_inptrrii+t]
7:
for t = tr. lR do
iterate nonzeros of query row i
8:
f f Q_ QLindices[t]: qu +- Q [aaluesstt] Po - Kfffifinda [ff_indptr[fii P1 & Kf indptr[f+1] (PL.PR) for p Pr. to PR do
9:
posting list for feature indices/ppppi ), [jo: to B.)) >> only keys that fall inside the key tile
BINARY
11:
j K Kffiindicespp
c jo
13:
kv + KT Kf-values(p)
14:
row_scores(c] += (qv ku)/V
> feature-overlap accumulation in registers
15:
end for 16: end for
17: for c= to B. do
18:
scores/r, t- row sco_sccss
store to SRAM after register accumulation
19: end for
20: end for
21: Mask (optional): apply causal mask in-place to scores.
Online softmax per row (as in FA)
23: for r 0 to B, do ig rr o, + zeros(d.) ) in registers
25:
accumulator for output row i iterate nonzeros Psj in row i row-vector, contiguous load from HBM
t P [iid LR +- P P_indptr [i+1]
27:
for t = t to fR - do
28:
j + P_indices[{]: p P P_values[t]
29:
v, V:d,
30:
0, + v.
31:
end for
32: Write back: add o; to the corresponding row of O.
end for
Binary Search for Tiling, To restrict computations to the current processing tile [io. + Be), we employ a BINAARINSSSSS SEARCH SEAARH RANGE routine. for a feature f. we search the sorted array Kf iiniices to find the sub-range (Pt., Phs) such that all indices fall within the current key tile. Since this operation runs in registers with a fixed number of iterations, highly efficient and branch-regular.
16

---

Published as a conference paper at ICLR 2026
C.2 PARALLELISM MODEL
The FlashSFA kernel adopts tiling strategy similar to standard FlashAttention but optimized for sparse operations. The mapping CUDA threads, warps, and blocks to the computation is designed to maximize occupancy and eliminate the need for atomic operations during score accumulation.
Grid and Block Mapping.
The computation grid is defined as grid (TN/B., B, B H), where N is the sequence length, B, the row tile height, is the batch size, and H is the number of heads.
Grid: blockIdx selects row tile of height B, (typically 128), and blockIdx.y selects the specific batch-head pair.
Block: Each thread block consists of 256 threads (8 warps). A single block is responsible for computing the attention output for a specific row tile (ro, ro + B,).
Warp and Thread Hierarchy. Within a block, the workload is distributed as follows:
Warps: The 8 warps in block process disjoint stripes of rows. Each warp is assigned 16-row stripe within the B, rows handled by the block.
Within warp, threads are mapped to 2D grid to process the score matrix. Each thread is responsible for 2 x patch the score tile (two rows x two columns). Across all 32 threads in warp, these patches perfectly tile the warp's assigned stripe.
Threads:
Loop Parallelization and Atomicity. The kernel iterates through column tiles (Key/Value blocks) in the outer loop. Inside the loop:
For the Query (Q), lines are distributed across warps.
For the non-zero elements (nnz) within Q row, threads iterate sequentially over the CSR segment.
No Atomic Operations. A critical design choice is the absence of atomic operations (e.g.- at comiAdd)) for score accumulation. Since each output score position SIr. c] in the tile is owned by exactly one thread (via the fixed 2 x 2 patch mapping), each thread accumulates all partial con- tributions for assigned scores in registers. This ensures thread safety and maximizes throughput.
C.3 CORE DATA STRUCTURES
FlashSFA relies on specialized sparse storage format to efficiently handle the intersection of active queries and keys.
Feature-wise CSC Format (CSC)eett). Standard sparse matrices typically store data Com- pressed Sparse Row (CSR) Column (CSC) format where dimensions correspond to token indices. However, to efficiently retrieve relevant keys for a given active feature in the query, we utilize a transposed view for the Key matrix K, denoted CSCCfeal (K).
Columms: Correspond Feature IDs (f (0,
Rows: Correspond to Key IDs (j E N)).
This layout allows the kernel to quickly access the posting list (list of token indices) for any specific feature of activated by the query.
C.4 MEMORY ACCESS STRATEGY
Efficient memory access is paramount for GPU performance. FlashSFA employs specific strategies ensure coalesced access despite the irregularity of sparse data.
17

---

Published as a conference paper at ICLR 2026
Accessing Q (CSR). For the Query matrix stored in CSR format, each block processes tiguous row range [ro. ro + B,). We calculate the span of non-zero elements [g qhs) = (Qindptr [ro] Qiddpr (ro B,]), which corresponds to a single contiguous segment in HBM. This allows for streamlined streaming of Q indices and values.
Accessing K (Coalesced Sparse Reads). For given key tile and active feature f, once the sub- range Pioo Phi) is identified via binary search:
Warps cooperatively load the indices and values from CSCseat-
We utilize lane-strided access patterns, e.g., loading index p = Pto + Lane_id + 32 x h.
This ensures that even though the semantic access is sparse (random features), the physical memory transactions are coalesced into contiguous slices of HBM, which are then staged into shared memory.
Accessing V (Sparse-Dense Multiplication). The computation of O P&V involves sparse attention matrix P and a dense value matrix V. Since is stored in row-major dense layout, the access any specific row Vj is contiguous:
The non-zero structure P determines which rows of to access.
Threads in warp cooperatively load row Vj using vectorized instructions (e.g., float4 or half2), ensuring high bandwidth utilization.
Memory Bandwidth Analysis Empirical profiling results in Table confirms that the memory system delivers close to peak bandwidth (approx. 919 GB/s in "memory-only" benchmarks), indi- cating that memory access V is not bottleneck.
Table HBM Bandwidth Comparison. "w'o compute" denotes measuring memory throughput with computation logic disabled.
Kernel
Dense
Dense w/o compute
FlashSFA
FlashSFA w/o compute
HBM Bandwidth (GB/s)
14.22
1194.34
17.14
919.38
C.5
REVISED COMPLEXITY ANALYSIS
potential concern with sparse attention the overhead introduced by the sparsification process (Top-k selection). SFA utilizes RTop-k kernel(Xie et al.) 2024) to sparsify Q and K with a com- putational complexity of ((Nd). This kernel employs GPU-parallclized binary search where each warp processes feature row. As shown in Table the latency of RTop-k is negligible compared to the full attention computation.
Table Latency comparison (ms) between standard torch. topk and RTop-k kernel across different context lengths. The Ratio indicates the percentage of time RTop-k consumes relative to the total attention forward pass.
Context Length (N)
1024
4096
8192
16384
32768
65536
torch. took
0.730
2.701
5.336
10.684
21.205
42.374
RTop-k (Ours)
0.221
0.589
1.089
2.080
4.057
8.080
Ratio of RTop-k (%)
10.51
2.10
1.96
1.90
1.03
0.51
Complexity P&V. The multiplication P@V effectively becomes a Sparse Matrix-Matrix Mul- tiplication (SpMM). Since the sparsity of P is induced by the sparsity of Q and K, the number of non-zero elements in P is significantly reduced. While a standard dense attention requires O( operations, SFA reduces the effective FLOPs count significantly. The memory access pattern de- scribed in Section ensures that the theoretical FLOPs reduction translates into actual wall-clock speedup by maintaining high memory bandwidth efficiency.
18

---

Published as a conference paper at ICLR 2026
D LATENCY TIMING RESULTS
Table 9: Latency (ms) versus context length.
Context Length
Variant
1024
4096
8192
16384
32768
65536
Dense_250
10.98
176.20
712.98
2894.46
11 772.47
49 197.70
Sparse 32/256
3.21
45.89
180.38
715.02
2886.21
11 529.74 74
Sparse 24/256
2.77
33.96
154.79
612.81
2488.04
8309.51
Sparse 16/256
2.31
29.15
128.52
510.18
2079.82
7388.78
Sparse 12/256
1.97
21.28
109.04
431.50
1769.73
5063.70
Sparse 10/256
1.93
20.06
106.93
422.95
1734.99
4877.05
Sparse_ 8/256
1.86
17.41
102.63
405.52
1665.69
4235.00
9p2rse_6/256 6/256
1.68
15.62
96.25
365.26
1505.59
3841.40
Sparse_4/22 4/256
1.51
13.94
82.48
324.97
1345.26
3412.10
Sparse_2/2 2/256
1.43
13.32
77.66
305.81
1273.37
2999.38
Dense 128
2.10
28.05
112.88
449.61
1981.92
7879.33
Sparse 32/128
2.17
31.01
120.58
465.68
1802.03
7101.95
Sparse 28/128
1.80
25.06
98.13
387.02
1535.46
6103.98
Sparse 24/128
1.70
23.84
94.10
373.42
1486.14
5909.66
Sparse_16// 16/128
1.56
17.94
70.62
279.72
1108.96
4412.02
Sparse_12/122 12/128
1.11
16.25
64.39
255.24
1017.04
4047.53
Sparse110/122 10/128
1.00
15.03
60.41
239.77
954.21
3814.06
Sparse_8/10
0.92
13.17
54.25
215.37
777.15
3323.53
Sparse 6/128
0.79
11.17
42.21
161.43
681.73
2738.84
Sparse_4/128
4/128
0.67
10.08
40.09
157.88
579.75
2576.93
Sparse 2/128
22122
0.58
9.90
38.92
154.71
539.49
2423.82
Dense 64
0.77
13.51
50.62
202.56
801.50
3137.78
Sparse_166n 16/64
0.90
12.51
39.41
195.53
779.19
2963.94
Sparse_12/d 12/64
0.70
9.71
38.23
151.60
603.18
2400.37
Sparse_10ke 10/64
0.67
9.23
36.36
144.17
573.31
2282.26
Sparse_8/66 8/64
0.59
8.14
32.00
126.99
504.23
2014.14
Sparse_6/0
0.51
7.05
27.64
109.36
434.58
1727.43
Sparse 4/64
0.41
5.41
21.07
83.12
328.83
1311.59
Sparse 2/64
0.39
5.15
19.75
77.96
309.13
1233.64
19

---

Published as a conference paper at ICLR 2026
E COMPARISON OF EFFICIENT ATTENTION BY TRAINING
Table 10: Latency, Perplexity and Accuracy results comparison with various compression and ac- celeration techniques, categorized into Token- and Feature-Level Operations. For token-level operations. "Longforemer" (Beltagy et al. 2020 denotes fixed token sparsity pattern, "NSA" Yuaan ett al. 2025) denotes dynamic token sparsity pattern. "Dense (full)" baselines use full hidden size and uncompressed KV cache: "Short (d X)" denotes baselines with half feature dimensions; "Quant" denotes 8-bit quantization aware training (QAT (Liu et al. (024b)) on weights and activa- tions; "Low-Rank' denotes PCA-based projection matrix tuning, "MLA" denotes multi head latent attention (Liu et al., 2024a), and "MLA + SFA" combines SFA with latent key/value. "La- tency @ 128k" is measured by "Decoding with KV cache (TTNT) (ms)" and "Prefilling with full attention (TTFT) (s)". PPL. is evaluated on Open WexText for GPT-2 and Pile for Qwen3.
Latency @ 128k
PPL J
Ace t
Model
Variant
Decode
Forward
OWT/Pile
PIQA
LAMBDA
ARC-e
ARC-c
HellaS
Avg
Dense (full)
17.08
16.86
17.29
56.34
22.78
28.35
14.32
19.61
28.28
Token-Level Operation
Longformer
6.75
21.27
13.01
18.92
28.10
GPT2
+SFA (k 8)
5.23
6.18
19.30
52.81
20.54
26.39
12.59
17.24
25.91
124M
Feature-1 evel Operation
Short (d 32)
8.37
7.86
20.70
51.30
19.39
25.72
12.47
14,26
24.63
Low-Rank
8,93
7.99
19.89
51,79
20.04
26.47
12.92
14,99
25.24
MLA
5.04
15.39
17.38
57.83
22.29
28.37
13.92
19.66
28.41
MLAA SFA
3.98
15.05
21.92
13.10
19.01
27.25
Quant
14.26
12,97
17.64
56.18
21.03
28.09
13.58
19.05
27.59
SFA (k- 8)
14.12
9.41
18.17
54,92
21.03
28.41
13.41
19.26
27.40
SFA (quant)
12.28
8.72
18.54
54,53
20.81
28.39
13.27
18.97
27.12
Dense(full)
80.84
77.65
4.66
62.47
34.82
45.41
20.35
33.95
39.40
Token-1 Lvel Operation
NSA
9.73
20.32
4.57
62.69
35.01
45.10
20.47
34.42
39.54
Qwen3
+SFA
16)
17.17
33,58
42.74
18.31
32.48
37.43
0.68
Feature- Level Operation
(d 64)
31,27
15.83
28.29
35.08
Low-Rank
40.58
32.46
59.19
31.49
41.77
15.80
30.65
35.78
MLA
8.74
68.92
4.69
62.39
34.71
45.41
20.17
34,21
39.38
MLA SFA
6.72
65.29
4.9
61.22
33.94
43.36
19.25
33.94
38.34
Quant
72.23
59.73
4.71
62.29
34.33
45.39
20.02
33.91
39.19
SFA (k- 16)
66.29
34.20
4.81
61.73
34.05
45.62
19.27
34.03
38.94
SFA (quant)
$7.47
30.74
5.16
$9.63
33.10
44.93
15.98
33.64
37.46
Table I0 compares SFA with variety of token-level and feature-level compression / accelcration techniques on GPT-2 124M and Qwen3-0.6B. We report both prefill ("Forward") and decoding ("Decode") latency at 128K context, together with perplexity and downstream accuracy.
Orthogonality to token-level methods. For token-level operations, SFA is applied on top of Long- former and NSA as drop-in replacement for their dense attention blocks. In both models, adding SFA consistently reduces both Decode and Forward latency while achieving comparable perfor- mance. This shows that SFA is orthogonal to token-level sparsification: can be combined with existing token-level sparse attention methods further accelerate long-context inference.
Feature-level speed-accuracy trade-off. Among feature-level methods, SFA can also be combined with MLA(on the compressed latent vector) and quantization. Pure SFA reduces latency compared to the dense baseline while keeping PPL and average accuracy close. Compared with Short and Low-Rank feature compression, which suffer larger accuracy drops, SEA and SFA (quant) maintain much higher accuracy at similar or better speed. Overall, SFA and its combinations deliver the strongest performance among feature-level approaches while still providing significant end-to-end speedups.
20

---

Published as a conference paper at ICLR 2026
F
LOAD BALANCE
1
0.9
Head index
-0.8
Normllzed entrooy
7
10
13:
10
13
15
16
18
20
23
25
26
27
Layer index
(a) Entropy of Q per layer per head
1-
Head index
5-
0.6
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
Layer index
(b) Entropy of K per layer per head
Figure 7: Entropy of Top-K feature selection across layers and heads. We plot the normalized entropy of the TopK index distribution for each attention head and layer of Qwen3- 3-0.6B when apply- ing SFA. (a) Entropy of the TopK positions of query vectors Q (16 heads, due to GQA). (b) Entropy of the TopK positions of key vectors (8 heads). Each cell corresponds to one (layer, head) pair, and brighter colors indicate higher entropy (more balanced use of feature dimensions).
natural concern for TopK sparsification on Q and K' that some heads layers might collapse to using only few feature dimensions, leading to poor load balance. To study this, we measure the normalized entropy of the TopK index distribution for every head and layer a small but diverse evaluation set: sample 50 samples from each of the Arxiv, Github, FreeLaw, and PubMed domains the Pile validation split (200 samples in total) (Figure For the 16 Q-heads in Qwen3 (due to GQA), the entropy ranges from 0.88 to 0.98 with an average of 0.94. For the K- heads, the entropy ranges from 0.85 0.97 with an average of 0.93. These values are close to the maximum possible entropy (1.0) and show only mild variation across layers and heads, indicating that the selected TopK dimensions remain well distributed rather than concentrating on a few indices.
Although SFA does not introduce any explicit load-balance loss, the feature activations remain nearly balanced. We hypothesize that, unlike TopK applied to weights (as in MoE routing), applying TopK directly feature vectors during end-to-end training encourages the model exploit its full expressive capacity: different dimensions are naturally used whenever they help reduce the training objective. As result, the model tends to learn near-uniform utilization of features.
21

---

Published as a conference paper at ICLR 2026
G ORTHOGONAL BASELINES
Table 11: Comparison results with token-sparse, KV-pruning, low-rank, and kernel base- lines on GPT-2. Token-sparse training methods include Routing(Roy et al. 2021) and Long- (ellar al. 2020 KV-pruning (training-free methods include H2OZZang
al.) 2023),
Quest Tang et al. 2024). and SnapK Vt (022)); Loki Loki(Singaania et al. 2024) is low-rank key compression method (training-frees
al.) 2021 is kemel-based approx- imation. Rows marked "+SFA (k = 8)" apply our feature-spars SFA to Longformer and SnapK showing that SFA is orthogonal to these approaches and can be combined with them.
Performer(Choromane
Latency @ 128k
PPL
Ace
Model
Decode
Forward
OWT/Pill
PIQA
LAMBDA
ARC-e
ARC-e
HellaS
Avg
Dense (full)
17.08
16.86
17.29
42.74
22.78
28.35
8.12
19.61
24.32
14.12
18.17
41.62
21.03
28.41
19.26
23.54
Token Spurse (Training)
Routing Longformer
7.92
18.64
41.39
21.08
28.31
18.89
23.35
6.75
7.93
18.73
41.28
21.27
28.02
7.01
18.92
23.30
+SFA (A 8)
5.23
6.18
19.30
40.75
20.54
26.39
6.63
17.24
22.31
GPT2 124M
VV-punning (Training- free)
H₂O
13.32
16.86
18.02
41.81
20.55
27.04
7.38
18.75
23.11
Quest
10.84
17.95
42.34
20.79
28.3
7.82
18.83
23.62
SnapKV
9.88
16.86
17.91
42.49
21.92
28.43
8.01
19.38
24.05
+SFA (k 8)
6.92
9.41
19.44
39.99
20.24
27.13
6.83
17.74
22.39
Low-rank keys (Training-five frree)
Loki
11.39
16.86
17.82
42.1
21.29
28.01
7.99
19.24
23.73
+SFA
9.09
19.29
40.83
20.04
27.85
7.13
18.03
22.78
Kemel Method
Performer
9.43
19.72
39.83
19.11
26.72
6.77
15.38
21.56
Orthogonality and composability with existing token sparse methods. Table 11 compares SFA with representative cong-context techniques on GPT-2 124M and efficient attention methods. As a standalone replacement of dense attention, SFA already improves efficiency over the dense baseline while perplexity and average accuracy remain close. More impor- tantly, SFA orthogonal to existing methods and can be combined with them for additional gains.
SFA with other
Token-sparse training methods. When applied on top of Longformer, we sparsify selected tokens. SFA further reduces latency from 6.75/7.97 to 5.23/6.18 (xx 1.3x faster decode and prefill), with only a modest change in quality. This shows that feature-level sparsification SFA complements token-level sparsity patterns.
KV-pruning and Low-rank keys methods. KV-pruning methods such as H₂O, Quest, and SnapKV improve speed by compressing the number of tokens in the KV cache, so they only accelerate the Decode stage and leave Forward latency unchanged. When we combine SFA with SnapKV, we obtain additional acceleration in both stages. Similar behavior holds relative H₂O and Quest. This shows that SEA is complementary to KV-pruning: KV-pruning reduces the number of cached tokens for decoding, while SFA sparsifies feature dimensions and brings additional gains.
22

---

Published as a conference paper at ICLR 2026
H ABLATION
30
1.5
Dense (full)
Dense (full)
SFA (ours)
Latency@32k
SFA (ours)
Perplexity
25
1.0
20
0.5
15
10
2
6
12
16
6
8
12
16
Sparsity(k)
Sparsity(k)
(a) Perplexity sparsity
(b) Latency vs sparsity
Figure 8: Ablation of sparsity k on GPT-2 124M with fixed head dimension d = 64. Perplexity on WeenTe (left) and latency at 32K context (right) as function of the Top-k sparsity level used by SFA. The dashed gray line denotes the dense (full) attention baseline; the red curve shows SEA with different
30
2.5
Dense (full)
Dense (full)
SFA (ours)
2.0
SFA (ours)
Perplexity
25
Latency@325
1.5
20
1.0
15
0.5
10
0.0
32
64
128
256
32
64
128
256
Head Dimension
Head Dimension
(a) Perplexity sparsity
(b) Latency sparsity
Figure 9: Ablation of head dimension dnead on GPT-2 124M with fixed sparsity A: = 8. Perplexity on WpenWetText (left) and latency at 32K context (right) as function of the head dimension dhend used by SFA. The dashed gray line denotes the dense (full) attention baseline; the red curve shows SEA with different dhead-
Sensitivity to sparsity k. Figure studies how the Top-k sparsity level affects performance. As increases from very sparse settings (e.g., k = 2) to denser ones (e.g., k 16), perplexity monotoni- cally decreases and quickly approaches the dense baseline; for k > 8, the SFA curve is very close to dense attention. In contrast, latency at 32K grows smoothly with k: very small k yields the largest speedup, while moderate & (around ko 8) still keeps a substantial latency advantage over the dense model with only small perplexity gap. Overall. SFA exhibits stable speed- accuracy trade-off and is not overly sensitive to the exact choice of k. allowing practitioners to pick k to match a desired latency budget.
Sensitivity to head dimension dheud- Figure varies the head dimension while keeping SFA en- abled. When the heads are extremely small (e.g., dhead = 32), perplexity degrades noticeably. As we increase the dimension, perplexity quickly improves, and at dhead 64 it is already very close to the dense baseline while latency remains substantially lower. Further increasing dhead beyond 64 brings only marginal perplexity gains but steadily increases latency. Thus dhess 64 emerges the sweet spot of the speed speed-accuracy trade-off: it recovers most of the dense-model performance while preserving most of the acceleration provided by SFA.
23

---

Published as a conference paper at ICLR 2026
I TRAINING STABILITY ANALYSIS
In this section, we investigate the training process of SFA. Figure 10 illustrates the validation loss trajectories across varying sparsity levels {2, 8,8, 16}) of GPT2-124M. We observe that the loss curves exhibit smooth, monotonic convergence devoid of divergent spikes or chaotic oscilla- tions. Notably, even under the most aggressive sparsity constraint (k red line), the model converges steadily. These empirical results suggest that SFA can intrinsically maintain training sta- bility without suffering from excessive variance or optimization instability.
SFA(K=16)
SFA(k=8)
SFA(K=4)
SFA(k=2)
Vall Loss
3.0-
20000
40000
60000
80000
100000
Steps
Figure 10: Validation loss curves of SFA on GPT-2 (124M) pre-training. We compare varying sparsity levels & E [2, 8, 16). The curves decrease smoothly and monotonically without divergent spikes, demonstrating that SFA maintains training stability even under aggressive sparsity (k 2).
MEMORY SAVING
In our implementation, memory gain can be achieved when & < dd. The memory savings of SFA compared to the dense model depend on the data precision used store the col_indices array and row_poi nter array the CSR matrix.
For CSR matrix with shape (N, d) where each row has a fixed number of l non-zero values, the required bytes for each component are calculated as follows:
value array memory:
Mem (N k) Sval
(10)
indices array memory:
MeManiicss = (N x k)) x Sixx
(11)
indptr array memory (length is N + 1):
MeMinipir (N + 1) x Sper
(12)
Where S denotes the number of bytes for the data format.
Therefore, the total memory consumption of the CSR format the sum of these three parts:
Mem Mem MeMandices
MeMandper
(13)
24

---

Published as a conference paper at ICLR 2026
Substituting the above formulas, obtain the final memory consumption formula:
Memer (N Stal (N k: Sixx) ((N + 1) Spi)
N x k (Stal + S + (N + x Spa
(14)
Consequently, compared to dense matrix of the same shape, the ratio of memory consumption is:
MeMaense
N d x Skal
x Syal
Ratio
(15)
McMas
N x 1 x (Sval Sjax) + (N + x Spir k X (Sval + Sjax) + Spr
As the Q/K feature dimension in Transformers is generally small, indices are typically stored in int format and indptr in int32 format. When we use fp1 6/bf1 store the value array:
MeMdense
d x 2
2d
2d
Ratio
2
(16)
Mem
k
+4
3k 4
3k
K ADDITIONAL NIAH EXPERIMENT
To verify that SFA functions effectively general-purpose mechanism without requiring task- specific supervision, we evaluated the retrieval capabilities of SFA in zero-shot setting. We trained the Qwen3-0.6B model equipped with SFA solely on general language corpora (standard pre-training) and evaluated it on the NIAH task.
As presented in Table 12 SFA consistently outperforms the dense attention baseline across all tested context lengths (Ik to 4k), despite lacking specific training for retrieval tasks.
At context length of 4k, SFA (k = 16) achieves an accuracy of 71 71%, significantly surpassing the dense baseline (62%). Even with aggressive sparsity (66%).
8), SFA maintains superior performance
In addition to improved accuracy, SFA provides substantial speedups. Specifically, SFA (A: 8) achieves a 1.5x speedup at 4k context length compared to the dense baseline.
These findings indicate that feature-level sparsification does not introduce an information bottle- neck. On the contrary, the results suggest that SFA preserves essential semantic information while potentially filtering out noise in long-contes scenarios, allowing it to function effectively within general- purpose foundation model paradigm.
Table 12: NIAH accuracy (%) within 4k Context Length. Qwen3-0.6B trained on Pile dataset with 4k window, and the accuracy rate on NIAH test lengths from 1 to 4k.
Context Length
lk
2k
3k
4k
Speedup @ 4k
Dense(full)
93
87
79
62
1.0x
SFA(k 8)
95
90
80
66
1.5x
SFA(k 16)
96
90
83
71
1.2x
25

---

Published as a conference paper at ICLR 2026
L SVD ANALYSIS
120
120
80
Rankg0,3
60
60
40
Layer
(a) Q matrix
(b) K matrix
Figure 11:
Eigenvalue spectrum analysis for Qwen3-0.6B model. Layer-wise effective dimension
of (a) query and (b) key activation with normalized cumulative eigenvalue of 0.9, evaluated on the same sampled subset of the Pile validation set Appendix F
To better understand why Top-k feature sparsification can preserve semantic information in atten- tion, we analyze the intrinsic dimensionality of the query and key representations pretrained dense model.
We use the pretrained Qwen3-0.6B model and run it on the same sampled subset of the Pile valida- tion set in Appendix [ For each transformer layer and attention head, we collect the corresponding query and key vectors Q, K IR (with head dimension = 128). We then perform singular value decomposition (SVD) the stacked feature matrices and compute the effective rank given energy threshold 7 0.9.
As shown in Figure [1] despite the nominal head dimension d 128, both Q and K exhibit sistently low effective rank, typically around 50-60 across layers. This confirms that the attention features lie on a low-dimensional manifold and are therefore highly compressible. The key matrices tend to have slightly lower effective rank than queries, but both are far from full rank, indicating substantial redundancy in the dense representations.
M LLM USAGE STATEMENT.
In line with the ICLR policy, we disclose the use Large Language Models during the preparation of this manuscript. Our use of these tools was strictly limited to assistance with language and formatting. Specifically, we employed an LLM correct grammatical errors and improve the clarity and readability of sentences. The LLM had no role in the core scientific aspects of this work, including research ideation, methodological design, experimental analysis, or the generation any results or conclusions. All intellectual contributions and the core content of this paper are solely the work of the authors.
26


Reviews:
[{"review_id": "S0syYTVeX7", "reviewer": "Reviewer_Q1QB", "summary": "The authors propose a feature sparse attention to improve the prefill time efficiency. They provide a FlashSFA kernel which integrates their idea into flash attention kernel. They show that models can be trained from scratch to get reasonably close perplexity and downstream accuracy as compared to dense models.", "questions": "1. Does FlashSFA also provide improvement in time to next token (TTNT)? From what i understand it can only provide improvement in time to first token(TTFT). In table 2 the speedups refer to TTFT or TTNT?\n\n2.Related to above, given a $k$, what is the memory footprint of storing $k$ indices for one row in CSR format isnt it $3k$? So if we had d=64, $k < 64/3$ is required to get any memory gains.  is that correct understanding?\n\n3. While the paper dismisses sparsity at sequence level as degrading accuracy, we do not see any comparison with such methods in experiments. sparsity in training (NSA / DSA) would be good baselines in my understanding. Can Authors elaborate more on the what are good baslines for FlashSFA in sparse attention techniques. \n\n4. Related to above, the FlashSFA also shows degradation across the board compared to baseline. So it is reasonable to compare against other sparse attention techniques at training time. \n\n5. The fact that we had to train for retreival tasks indepdnently from scratch is a bit worrying. It might have something to do with the information bottleneck in SFA. In principle we are looking for a model that can do all of the tasks. training for specific task is against the foundational model paradigm. \n\n6. How does top-k coordinate distribution look like in practice. Equation 7 assumes that it is balanced. But in practice (atleast in the case of dense models), some channels can show outlier distribution (See AWQ quantization).\n\n7. Also, authors mention that the benefits are multiplicative and their technique can be combined with sequence-sparse attention methods.  Is there some evidence to this effect in authors' experiments / literature.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. Feature sparsity is not well explored in literature to the best of my knowledge. This paper does that\n2. Opensource FlashSFA kernel that be used by the community.", "weaknesses": "1. Lack of baselines.\n2. experiments are not convincing. There is degradation in quality when compared to baseline. Requires separate training for retrieval task.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "orOiFWZOb1", "reviewer": "Reviewer_FnGK", "summary": "The paper addresses the $O(n^2d)$ computational bottleneck of self-attention in Transformers. While existing methods focus on sparsity along the sequence axis, this paper explores an \"orthogonal axis: feature sparsity\". The authors propose Sparse Feature Attention (SFA). SFA represents queries (Q) and keys (K) as \"k-sparse codes\", claiming to reduce the attention cost from $\\Theta(n^2d)$ to $\\Theta(n^2k^2/d)$. To implement this efficiently, the paper introduces FlashSFA, an \"IO-aware kernel that extends FlashAttention\"  to operate on sparse overlaps without materializing dense score matrices. Experiments on GPT-2 and Qwen3 pretraining show SFA matches dense baselines while improving speed and reducing FLOPs and KV-cache.", "questions": "1. Can the authors provide the corresponding implementation details for the FlashSFA kernel, wrt the above questions in W1?\n\n2. Can the authors provide a revised complexity analysis for the entire SFA block, including the $O(ndk)$ or $O(nd \\log d)$ cost of Top-k and the $O(n^2 d_v)$ cost of the P@V multiplication? How does this revised, complete analysis affect the scaling claims?\n\n3. Given that P@V is a \"large proportion\" of the FLOPs, why was V not sparsified? Did the authors experiment with this, and if so, what was the impact on performance?\n\n4. The $k^2/d^2$ scaling analysis (Eq. 7) relies on an assumption that \"supports are balanced across dimensions\". Is there empirical evidence for this? What happens to performance and cost if the feature selection is unbalanced?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. **Originality**: The paper's primary strength is shifting the optimization focus from the token axis to the \"feature axis\". This approach is \"orthogonal to token-level sparsity and paging\".\n\n2. **Theoretical Idea**: The concept of using k-sparse codes to reduce the computational complexity of the $QK^\\top$ interaction is novel.\n\n3. **Empirical Quality** (as reported): The reported results, if verifiable, suggest that SFA can maintain \"comparable quality to dense attention\" (Table 1 , Table 3 ) in pretraining and fine-tuning scenarios.", "weaknesses": "* 1. **Reproducibility**: The paper's central empirical results (e.g., \"up to 2.5×\" speedup , Figure 3 , Figure 5 ) are entirely dependent on the \"IO-aware kernel\" FlashSFA. This kernel is positioned as a primary contribution. However, the paper provides only high-level pseudocode (Algorithm 1 ) and a brief description (Appendix C). The kernel code itself, which is essential for reproducing any of the speedup benchmarks, is not provided.\n  * Parallelism Model: How are CUDA threads, warps, and blocks mapped to the tile computations? How are the nested loops (Line 2, Line 4) parallelized?\n  * Atomic Operations: How are the \"scatter-adds\" to the `scores` buffer (Line 11) handled? Multiple threads processing different features will inevitably write to the same `scores[r, c]` location, requiring atomic operations which are not mentioned but are critical for correctness and performance.\n  * Memory Management: The algorithm hides all memory access patterns. How are reads from the sparse `Q` and `K` structures (e.g., `Q_indices`, `Kf_indices`) coalesced for efficient HBM bandwidth usage? How is the `scores` buffer  truly managed in SRAM (shared memory) across threads?\n  * Core Data Structures: The algorithm relies on a non-standard $CSC_{feat}(\\overline{K})$ format. The precise memory layout of this \"feature-wise\" CSC matrix is not defined.\n  * Search Implementation: The efficiency of `BINARY_SEARCH_RANGE` (Line 7) is crucial. Its implementation details in a parallel CUDA environment are non-trivial and un-discussed.\n\n* 2. **Complexity Analysis that Ignores the P@V Bottleneck**: The paper's headline cost analysis, $\\Theta(n^2k^2/d)$ (Abstract, Section 3.1), is confusing. This analysis only accounts for the $QK^\\top$ computation. The authors explicitly state they are \"keeping V dense\". The authors admit in their own Appendix B.2 that \"a large proportion of the floating-point operations in the sparse version come from matrix multiplication in the P@V stage\". This non-sparsified $O(n^2 d_v)$ operation remains a quadratic bottleneck. The paper's main analysis in Section 3.1 ignores this, and the resulting speculative claim of a \"reduction of more than 1000x\" is unsubstantiated as it omits a dominant, quadratic term that the authors themselves acknowledge.\n\n* 3. **Core Operational Cost (Top-k) in Analysis**: The paper's efficiency analysis in Section 3.1 omits the computational cost of the $Topk_k$ operation (Eq. 3). This operation must be performed on both Q and K matrices (size $n \\times d$) for every forward and backward pass. While the authors mention using an \"RTopK kernel\" and \"lookup overhead\", they never incorporate this non-trivial (e.g., $O(ndk)$ or $O(nd \\log d)$) cost into their formal complexity model, making that analysis incomplete.\n\n* 4. **Limited Applicability and Inherent Distribution Shift**: The authors' own experiments in Section 5 demonstrate that applying SFA to pretrained models \"introduces a severe distribution shift\". The authors' solution is a multi-stage, regularized fine-tuning process, which requires \"an additional MSE loss\" and even pre-finetuning on another dataset. This implies the method's high adaptation cost.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9ZXFfmSuqU", "reviewer": "Reviewer_gwpA", "summary": "This paper proposes a novel method of introducing sparsity along the feature dimension of attention computation. Specifically, instead of directly introducing sparsity on the attention matrix, the paper proposes to instead introduce in the key and query matrices by retaining only the top-$k$ largest magnitude entries along the feature dimension. Naively, the full attention matrix requires $O(n^2)$ memory to store and process, but the authors combined this sparsity idea with FlashAttention, which processes attention weights by a tiling mechanism that only requires a compact score buffer.\n\nExperiments show that the proposed method achieves significant latency reduction while retaining most of the performance of dense models with only minimal degradations, measured both in terms of PPL (pre-training setup) and NIAH/Math/QA performance (post-training setup). Additionally, the paper also shows that sparsity patterns can be easily introduced by fine-tuning, thus opening the door of converting dense pre-trained models into sparse models for more efficient inference.", "questions": "1. For Table 1, how is the \"Best\" in \"Best results highlighted\" defined? I'm confused because SFA numbers are highlighted on the OWT/Pile scale even though dense models achieved lower numbers.\n2. For Table 2, your dense model has a dimension size of 2. Is that really the case? That looks like a typo to me because there is an experiment with $k=8$ below.\n3. Can you clarify on the comment on L376 -- \"sparse kernels incur lookup overhead\"? What overhead are you referring to?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 4, "strengths": "1. The simplicity of the method is quite impressive considering how significant of an efficiency gain it was able to achieve.\n2. Performance preservation has been verified thoroughly, covering both LLM pre-training and post-training measures.", "weaknesses": "1. One main drawback of the paper is the lack of discussion and comparison with the DeepSeek Multi-Head Latent Attention, which computes a low-rank compression of the key/value pairs. To me, there is a pretty significant resemblance between those two ideas, even though MLA did not explicitly explore sparsity. Some form of detailed discussion on the differences of those two ideas and some empirical comparison would further highlight the benefit of the proposal.\n2. Another minor complaint is that the experiments are mostly on toy tasks and are of pretty small model scale. I know the authors might be limited by the compute they can access, but it would be great if there's at least one experiment on a medium-scale model (7B-14B) which verifies that the latency benefit is able to extrapolate over model scale.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9Yvpuyju4b", "reviewer": "Reviewer_5rhv", "summary": "The paper introduces Sparse Feature Attention (SFA), which enforces k-sparse query/key vectors and computes attention only on overlapping active feature coordinates, reducing the arithmetic of QKᵀ from Θ(n²d) to Θ(n²k²/d). To make this practical, the authors propose FlashSFA, an IO-aware kernel that integrates sparse overlaps into the FlashAttention tiling/online-softmax pipeline without materializing n×n score matrices. Experiments on GPT-2 and Qwen3 pretraining, synthetic long-context retrieval (NIAH), and several downstream tasks claim to match dense-attention quality while yielding up to ~2.5× speedups, ~49% FLOPs, and ~41% KV-cache savings, particularly at long contexts.", "questions": "- Can you add head-to-head comparisons with H2O, Quest, SnapKV, LoKI, Routing/BigBird/Longformer, Performer/Nyström at the same training/inference configs (same context windows, same datasets), and report both quality and speed/KV?\n- GSM-8K drops relative to dense FT—can adaptive k or hybrid (dense for a subset of heads/layers) recover this?", "limitations": "", "rating": 4, "confidence": 2, "soundness": 1, "presentation": 3, "contribution": 3, "strengths": "- Clear, orthogonal idea: Moves the sparsity lever from tokens to feature coordinates, preserving token coverage while lowering per-interaction cost (Θ(n²k²/d)). The complexity argument is clean and intuitive.\n- Kernel contribution: FlashSFA marries sparse overlaps with FlashAttention’s online softmax/tiling, avoiding n² materialization and targeting the actual memory bottleneck.\n- Evidence of efficiency/quality trade-off: At fixed models, SFA often matches dense perplexity/accuracy and beats “short embeddings” that shrink d, with up to ~2.5× speed, ~49% FLOPs, and ~41% KV-cache reductions reported.\n- Long-context relevance: On synthetic NIAH and long-sequence latency/KV scaling, SFA’s advantages grow with context length/head dimension, aligning with the paper’s motivation.\n- Compatibility: Conceptually composes with token sparsity/paging; could multiply gains in real systems.", "weaknesses": "- Missing or thin comparisons to modern long-context/efficient baselines beyond “short embeddings”: e.g., token-sparse (Longformer/BigBird/Routing), KV-pruning/paging (H2O, Quest, SnapKV), low-rank keys (LoKI), and kernel methods (Performer/Nyström). SFA should be combined and compared head-to-head under the same training/inference regimes.\n- The k-sparsification uses row-wise Top-k with a straight-through estimator; stability/variance, gradient bias, and sensitivity to k (and to head dim d) are not adequately analyzed.\n- “Short-embedding” baselines may be overly weak (simple linear downsizing). Stronger compression baselines (e.g., low-rank Q/K, grouped-query attention variants, learned projections) should be included.\n- Reported GSM-8K lags for SFA vs dense after finetuning suggest arithmetic reasoning sensitivity to feature pruning; this deserves deeper diagnosis and mitigation (e.g., adaptive k, soft-top-k).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)