FlashSFA: Sparse Feature Attention Kernel

• FlashSFA is an IO-aware kernel that implements Sparse Feature Attention by converting dense query-key representations into k-sparse codes, reducing both computation and memory traffic.
• It computes attention scores using only overlapping active features, lowering the complexity from O(n²d) to O(n²k²/d) under balanced-support assumptions.
• Empirical results demonstrate up to 2.5× speed improvements and nearly 50% reductions in FLOPs and KV-cache while maintaining accuracy in long-context models.

FlashSFA is an IO-aware kernel introduced with Sparse Feature Attention (SFA) for scaling self-attention along the feature axis rather than the sequence axis. In the formulation of "Scaling Attention via Feature Sparsity," queries and keys are represented as $k$-sparse codes, and FlashSFA extends FlashAttention to operate directly on sparse overlaps without materializing dense score matrices. The method is presented as an exact, online-softmax implementation of $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$ over sparse feature intersections, with the attention cost reduced from $\Theta(n^2 d)$ to $\Theta(n^2 k^2/d)$ under the paper’s balanced-support assumption. Across GPT-2 and Qwen3 pretraining, the reported results are speed improvements of up to $2.5\times$ together with nearly $50\%$ reductions in FLOPs and KV-cache, while preserving retrieval accuracy and robustness at long contexts (Xie et al., 17 Mar 2026).

1. Definition, scope, and motivation

FlashSFA belongs to a line of exact attention kernels derived from FlashAttention, but its target is different from sequence-axis sparsification. The motivating claim is that standard self-attention remains bottlenecked by the dense query-key score computation and the associated memory traffic: for sequence length $n$ and head dimension $d$, standard self-attention costs $O(n^2 d)$. Existing efficient-attention methods typically reduce cost along the sequence axis through local windows, kernel approximations, or token-level sparsity, whereas FlashSFA is coupled to SFA, which instead explores feature sparsity (Xie et al., 17 Mar 2026).

The paper distinguishes two layers of abstraction. SFA is the attention formulation: $Q$ and $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$0 are converted into sparse $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$1-active feature codes, and attention scores are computed from feature-support overlaps. FlashSFA is the systems and kernel implementation: an IO-aware kernel that extends FlashAttention to the sparse-overlap setting. This division is central to the method’s presentation, because the arithmetic reduction of SFA alone would not be sufficient if the implementation still materialized an $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$2 score matrix (Xie et al., 17 Mar 2026).

The motivating contrast is with "short embeddings" and other dense low-dimensional baselines. The paper argues that high-dimensional but sparse features preserve diversity and semantic detail better than uniformly shrinking dimension, and that the principal cost of dense attention can be attacked without removing token-token interactions themselves. This is why FlashSFA is described as complementary to token sparsity rather than a replacement for it (Xie et al., 17 Mar 2026).

2. Sparse Feature Attention formulation

For head $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$3, the dense projections are written as

$$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$4

with dense attention

$$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$5

SFA modifies only the query-key representation. Given dense projected features $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$6, it applies a row-wise Top-$$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$7 operator: $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$8 For a vector $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$9,

$\Theta(n^2 d)$0

Each token therefore activates only $\Theta(n^2 d)$1 feature coordinates (Xie et al., 17 Mar 2026).

If the support of sparse query $\Theta(n^2 d)$2 is $\Theta(n^2 d)$3 with $\Theta(n^2 d)$4, and similarly for sparse key $\Theta(n^2 d)$5, then the score for token pair $\Theta(n^2 d)$6 is

$\Theta(n^2 d)$7

This preserves all token-token interactions in principle: any pair $\Theta(n^2 d)$8 may interact, but only along overlapping active coordinates. The attention semantics remain standard, since FlashSFA is presented as an exact implementation of

$\Theta(n^2 d)$9

The backward pass uses a straight-through estimator, with gradients flowing only through selected coordinates: $\Theta(n^2 k^2/d)$0 The same rule applies to keys (Xie et al., 17 Mar 2026).

The complexity argument is derived from the number of sparse feature overlaps. If supports are approximately balanced across dimensions and feature $\Theta(n^2 k^2/d)$1 is selected by

$\Theta(n^2 k^2/d)$2

tokens, then the total number of overlap contributions is

$\Theta(n^2 k^2/d)$3

Accordingly, score computation drops from $\Theta(n^2 k^2/d)$4 to $\Theta(n^2 k^2/d)$5. The paper states that this is a fraction $\Theta(n^2 k^2/d)$6 of dense attention cost (Xie et al., 17 Mar 2026).

3. FlashSFA kernel design

FlashSFA is introduced to ensure that SFA’s feature sparsity does not lose its benefit to memory traffic. Its design principle is the same one associated with FlashAttention: tile the computation, keep temporary scores on-chip, use online softmax, and never materialize the full score matrix. The difference is that FlashSFA replaces dense tiled GEMMs with sparse feature-intersection kernels (Xie et al., 17 Mar 2026).

The sparse storage layout is asymmetric. Queries are stored in CSR format with Q_indptr, Q_indices, and Q_values. Keys are stored in a feature-oriented CSCfeat format, whose columns correspond to feature IDs and whose rows correspond to key or token IDs. This arrangement is chosen so that, for an active query feature $\Theta(n^2 k^2/d)$7, the kernel can immediately retrieve the posting list of all keys whose sparse code also activates $\Theta(n^2 k^2/d)$8 (Xie et al., 17 Mar 2026).

For a query tile of size $\Theta(n^2 k^2/d)$9 and a key tile of size $2.5\times$0, the forward pass initializes a score buffer

$2.5\times$1

in SRAM. For each query row $2.5\times$2, it iterates over the nonzeros in the CSR row, takes feature $2.5\times$3 and value $2.5\times$4, retrieves the corresponding posting list from CSCfeat,

$2.5\times$5

restricts that list to the current key tile by binary search, and accumulates

$2.5\times$6

for every matching key entry in the tile. The sparse overlap is therefore computed exactly rather than approximately. After score accumulation, the kernel applies masking and FlashAttention-style online softmax, loads the necessary rows of $2.5\times$7, and updates the output without materializing dense score matrices (Xie et al., 17 Mar 2026).

The implementation is described as IO-aware because only tiles, not the full $2.5\times$8 matrix, are materialized in SRAM or registers. Stored tensors include the sparse $2.5\times$9 representations, the dense $50\%$0, and tiled temporary buffers; the full dense score matrix is never written to high-bandwidth memory. The paper also states that FlashSFA inherits $50\%$1 IO complexity in the FlashAttention sense (Xie et al., 17 Mar 2026).

4. Training, adaptation, and systems behavior

SFA is used as a drop-in replacement for dense $50\%$2 scoring in Transformer attention layers, while $50\%$3 remains dense. Sparse codes are produced on the fly in both training and inference by applying row-wise Top-$50\%$4 to the current $50\%$5 and $50\%$6 activations: $50\%$7 For efficient Top-$50\%$8, the implementation uses the RTop-k kernel. The paper reports the sparsification complexity as $50\%$9, but with low practical overhead: at $n$0, RTop-k costs $n$1 ms or $n$2 of the total forward pass, whereas at $n$3 it costs $n$4 ms or $n$5 (Xie et al., 17 Mar 2026).

To adapt pretrained dense models, the paper introduces a regularization term

$n$6

with

$n$7

Because FlashAttention and FlashSFA do not materialize the full attention matrix, the regularization is applied to head outputs rather than to attention matrices directly (Xie et al., 17 Mar 2026).

The memory model for sparse query or key storage is also given explicitly. For CSR,

$n$8

whereas dense storage is

$n$9

The approximate ratio is

$d$0

For fp16 or bf16 values, int8 indices, and int32 pointers, the paper simplifies this to

$d$1

This is the basis of the reported KV-cache reduction (Xie et al., 17 Mar 2026).

The kernel implementation is built on FlashAttention-2 style code in LeetCUDA. The grid is

$d$2

with a typical row tile height $d$3. The paper states that each thread block has 256 threads, or 8 warps, and that no atomic adds are needed because each score position in the tile is owned by exactly one thread (Xie et al., 17 Mar 2026).

5. Empirical performance

The headline empirical claim is that, across GPT-2 and Qwen3 pretraining, SFA matches dense baselines while improving speed by up to $d$4 and reducing FLOPs and KV-cache by nearly $d$5. Figure 1 is summarized in the paper as reporting a $d$6 speedup over original dimensionality in one comparison against short-embedding tradeoffs, a $d$7 better performance than the short-embedding baseline, a $d$8 FLOPs reduction, and a $d$9 KV-cache reduction (Xie et al., 17 Mar 2026).

On GPT2-124M, the dense full model has latency $O(n^2 d)$0 at 128k context, perplexity $O(n^2 d)$1, and average accuracy $O(n^2 d)$2. A dense short-embedding baseline with $O(n^2 d)$3 reduces latency to $O(n^2 d)$4 but worsens perplexity to $O(n^2 d)$5 and average accuracy to $O(n^2 d)$6. SFA with $O(n^2 d)$7 reports latency $O(n^2 d)$8, perplexity $O(n^2 d)$9, and average accuracy $Q$0. On GPT2-350M, dense full gives $Q$1, $Q$2, and $Q$3, while SFA with $Q$4 gives $Q$5, $Q$6, and $Q$7. On Qwen3-0.6B, dense full gives latency $Q$8, perplexity $Q$9, and average accuracy $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$00, while SFA with $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$01 gives $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$02, $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$03, and $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$04 (Xie et al., 17 Mar 2026).

The synthetic long-context retrieval results on NIAH are presented as evidence that SFA preserves long-context behavior. When trained up to 8k context, Dense($$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$05) gives $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$06 at $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$07 with speed $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$08, while SFA($$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$09) gives $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$10 at $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$11, and SFA($$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$12) gives $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$13 at $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$14. When trained up to 32k context, Dense($$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$15) gives $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$16, while SFA($$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$17) gives $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$18 at $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$19. Appendix K further reports zero-shot NIAH on pretrained Qwen3-0.6B: at 4k context, dense gives $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$20, SFA($$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$21) gives $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$22, and SFA($$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$23) gives $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$24 with $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$25 speedup (Xie et al., 17 Mar 2026).

Kernel-level latency results are substantially larger at long sequence length and head dimension. At sequence length 65k, Dense$$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$26 is reported as $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$27 ms, Sparse$$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$28 as $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$29 ms, Sparse$$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$30 as $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$31 ms, Sparse$$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$32 as $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$33 ms, and Sparse$$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$34 as $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$35 ms. At the same sequence length with $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$36, Dense$$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$37 is $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$38 ms, while Sparse$$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$39 is $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$40 ms. The paper notes, however, that end-to-end model speedups are smaller than kernel-only speedups because full-model latency includes non-attention work (Xie et al., 17 Mar 2026).

6. Limitations, scope, and related methods

The paper is explicit that FlashSFA is most effective in the long-context regime. For short contexts, sparse-kernel overhead from index lookups, posting-list traversal, and binary-search operations can dominate. The reported qualitative transition is that dense attention can still be competitive below roughly 4k, while SFA consistently outperforms dense attention beyond about 8k–16k (Xie et al., 17 Mar 2026).

The savings are also bounded by the fact that $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$41 remains dense. This means the practical reductions in FLOPs and KV-cache are much smaller than the raw reduction in query-key overlap count might suggest. The paper further notes that current GPUs and CUDA libraries do not yet provide ideal support for general sparse tensor products, and that very aggressive sparsity can hurt quality. It suggests adaptive sparsity budgets as a future direction (Xie et al., 17 Mar 2026).

FlashSFA occupies a specific position within the broader family of flash-style attention accelerators. It is distinct from the systolic-array work "SystolicAttention: Fusing FlashAttention within a Single Systolic Array," which addresses exact FlashAttention execution inside a single systolic array and uses the architecture name FSA rather than FlashSFA (Lin et al., 15 Jul 2025). It is also distinct from "Block Sparse Flash Attention," which keeps exact $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$42 score computation but introduces block sparsity after score formation, using exact-score pruning of value blocks rather than feature sparsity in $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$43 and $$\mathrm{softmax}\!\left(\frac{\bar Q \bar K^\top}{\sqrt d}\right)V$$44 themselves (Ohayon et al., 7 Dec 2025).

Within the paper’s own framing, FlashSFA is best understood as the systems realization of SFA: a FlashAttention-derived kernel that makes feature-level sparsity practical at scale. Its central claim is not merely that sparse features can reduce arithmetic, but that an exact, online-softmax, IO-aware kernel can exploit those sparse overlaps directly, making feature sparsity a complementary and underexplored axis for efficient attention (Xie et al., 17 Mar 2026).