FSA: An Alternative Efficient Implementation of Native Sparse Attention Kernel
Authors: Ran Yan
Youhe Jiang
Zhuoming Chen
Haohui Mai
Beidi Chen
Binhang Yuan
Venue: ICLR 2026 Poster (2026)
Keywords: Efficient attention
GPUs
Long context LLMs
Sparse attention
Primary area: foundation or frontier models, including LLMs

Abstract: Recent advance in sparse attention mechanisms has demonstrated strong potential for reducing the computational cost of long-context training and inference in large language models (LLMs). Native Sparse Attention (NSA), one state-of-the-art approach, introduces natively trainable, hardware-aligned sparse attention that delivers substantial system-level performance boost while maintaining accuracy comparable to full attention. However, the kernel implementation of NSA forces a loop order that is only efficient with a relatively large number of query heads in each Grouped Query Attention (GQA) group, whereas existing LLMs widely adopt much smaller number of query heads in each GQA group --- such an inconsistency significantly limits the applicability of this sparse algorithmic advance. In this work, we propose **F**lash **S**parse **A**ttention (**FSA**), an alternative kernel implementation that enables efficient NSA computation across a wide range of popular LLMs with varied smaller number of query heads in each GQA group on modern GPUs. Compared to vanilla NSA kernel implementation, our empirical evaluation demonstrates that FSA achieves (i) up to 3.5x and on average 1.6x kernel-level latency reduction, (ii) up to 1.25x and 1.09x on average end-to-end training speedup on state-of-the-art LLMs, and (iii) up to 1.36x and 1.11x on average for prefill-phase speedup in LLM generative inference.

Published as a conference paper at ICLR 2026
FSA: AN ALTERNATIVE EFFICIENT IMPLEMENTATION OF NATIVE SPARSE ATTENTION KERNEL
Ran Yan'*, Youhe Jiang" ' Zhuoming Chen", Haohui Mai', Beidi Chen', Binhang Yuan
1. The Hong Kong University of Science and Technology
2. Carnegie Mellon University
ryana féconnect. ust. hk, yone amaii..cm com (zhuominc, beidi @andrew. emu, edu.edu, {hachui, Bust hk
ABSTRACT
Recent advances in sparse attention mechanisms have demonstrated strong po- tential for reducing the computational cost of long-context training and inference in large language models (LLMs). Native Sparse Attention (NSA), one state-of- the-art approach, introduces natively trainable, hardware-alignee sparse attention that delivers substantial system-level performance boosts while maintaining accu- racy comparable to full attention. However, the kernel implementation of NSA forces a loop order that is only efficient with a relatively large number of query heads in each Grouped Query Attention (GQA) group, whereas existing L.MS widely adopt a much smaller number of query heads in each GQA group such inconsistency significantly limits the applicability of this sparse algorithmic advance. In this work, we propose Flash Sparse Attention (FSA), an alterna- tive kernel implementation that enables efficient NSA computation across a wide range of popular LLMS with varied, smaller number of heads in each GQA group on modern GPUs. Compared to vanilla NSA kernel implementation, our empirical evaluation demonstrates that FSA achieves (i) up to 3.5x and on aver- age 1.6x kernel-level latency reduction, (ii) up to 1.25x and 1.09 on average end-to-end training speedup state-of-the-ar LLMs, and (iii) up to 1.36x and 1.11x on average for prefill-phase speedup in LLM generative inference. The source code is open-sourced and publicly available athttts:://tttttbse Relaxed- Leaxxed-Systten-aaa/aaaaee Sww.aaass--hhhssseeeeeeeee Attppaarreeeteettiioo
INTRODUCTION
Large Language Models (LLMs) with long context windows (OpenAl 2024 Anthropic 2024
Young et al.] 2024 Dubey et al. 2024) face prohibitive computational costs due to the full atten- tion mechanism's quadratic time and memory complexity. As sequence length increases, attention computation becomes critical bottleneck - for instance, attention can account for 70 80% of to- tal decoding latency at a 64k token context (Yuan et al., (2025). In extreme cases, processing a million-token prompt with an 8B model can take up to 30 minutes on single GPU (Jiang et al.) 2024a). These observations underscore the urgent need for more efficient attention mechanisms in long-context LLM training and inference. A recent promising direction is to exploit sparse attention, whereby the query of each token only interacts with a subset of key and value, dramatically reducing the computation load and HBM I/O volumes. However, implementing efficient sparse attention scale is non-trivial - In fact, the challenge of implementing high-performance kernels has become major obstacle deploying state-of-the-art sparse attention techniques in practice. In this paper, want to explore: Can we design and implement efficient sparse attention kernel for wide range of current LLMs fully unleash the potential of this algorithmic advance over modern GPUs?
Addressing the above question is crucial because adopting sparse attention in long-context LLMs could mitigate the quadratic
and enable new applications (Xu al. 2025a Chen et al. 2024
Acharya et al. 2024 Wang et al.) 2024). By leveraging the inherent sparsity of attention patterns, one can significantly cut down computation and memory overhead. Among such methods, one
represents equal contribution. Correspondence Binhang Yuan.

---

Published as a conference paper at ICLR 2026
recently proposed promising example is Natively Sparse Attention (NSA) YYuan al. 2025),
sparse attention framework, which organizes keys/values into blocks and processes them via three parallel attention modules compressed coarse-grained tokens, selected fine- tine-grained tokens, and sliding local windows. By learning which tokens to compress or drop, NSA achieves long-context efficiency without predefined pattern, making it natural choice for long-context L.LM training.
Nevertheless, implementing an efficient sparse attention kernel, i.e, NSA, is challenging. The core difficulty lies in implementing the sparse mechanism in NSA (i.e., computing attention score based on selectively retained fine-grained tokens), where the query of each token needs to dynamically select different set of keys and values. Such computation results in irregular HBM access patterns on modern GPUs, where each query processes distinct selected keys/values, potentially requiring unnecessary padding for query tiles before executing accumulate instructions (e.g., wmma or wgmma), and leading to the underutilization of tensor cores.
level matrix multiply-and-
This scattered access pattern conflicts with the GPU hardware-efficient design principle: GPUs achieve their peak mathematical throughput when the warps execute dense (no-padded) matrix mul- tiply and accumulation instructions. Thus, current sparse attention implementations fail to translate the theoretical foating-point operations (FLOPs) reduction into wall-clock speedups.
Vanilla NSA kernel implements a two-level loop: In the outer loop, NSA kernel loads one token and batches query attention heads that share the same key and value heads; in the inner loop, NSA kernel loads selected KV block iteratively and performs attention computation. This strategy reaches kernel efficiency only when each Grouped Query Attention (GQA) (Ainslie al., 2023) group has sufficient number of query heads, so that no-padding is required to execute PTX instructions (e.g., anuna or wgmmaa on modern GPUs [ However, such an assumption may not hold for wide range of popular LLMs so that the original NSA kernel efficiency could drop considerably. With an insufficient number of query heads in each GQA group. batching query heads is inefficient to satisfy this hardware requirement. Thus, the original NSA kernel implementation must pad query attention heads to meet instruction requirements, resulting in unnecessary data loading and computations.
To resolve this issue, we propose FSA, which implements optimized kernels efficient for NSA under various GQA group settings. We make the following concrete contributions:
We propose an alternative implementation for the NSA kernel, which exchanges the two-level loop order in NSA implementation - FSA A loops over KV blocks in the outer loop and loops over query tokens in the inner loop to accelerate this system bottleneck. Since the number of query tokens that attend to a given KV block is usually much larger than the hardware required value, FSA introduces no padding, significantly reducing unnecessary kernel memory access and FLOPs, thereby facilitating faster token selection kernel execution.
Contribution 1:
Contribution 2: We analyze the rrade-off between vanilla NSA and FSA implementation in terms of kernel efficiency and memory accessing paradigm, which illustrates the effective design and im- plementation of FSA. A. To maximize the performance benefits of FSA A kernel design, we implement dedicated optimizations for query token memory access, which is accessed in the inner loop of FSA kernel, and employ separate optimized kernels for attention result reduction.
We conduct empirical studies to compare FSA with vanilla NSA and full at- tention. Concretely, we benchmark kernel execution latencies, end-to-end training and inference prefill phase latencies for state-of-thh-art LLMs. Compared to NSA, results show that FSA deliv- ers (i up to 3.5x and average 1.6x kernel-level latency reduction, (ii) up to 1.25x and 1.09x on average end-to-end training speedup, and (iii) up to 1.36x and 1.11 x on average inference prefill-phase speedup. Compared to full attention, the performance boost is further amplified.
Contribution 3:
2 PRELIMINARIES AND RELATED WORK
2.1 GPU KERNEL IMPLEMENTNTION
Parallelization in modern GPUs. Modern GPUs utilize massive threads to execute kernels con- currently. Optimized kernel implementations typically employ two-level parallelism: (i) Thread
'Concretely, performance is downgraded due to hardware requirements on matrix shapes for warp- /warpgroup- level matrix muuliply-and-cccumulll instructions (e.g., wmma wgmma) (NVIDIA where each dimension of matrix tile must he larger than specified value (e.g., least 8 on Hopper GPUss.
2

---

Published as a conference paper at ICLR 2026
block-level parallelism: Optimized implementations partition input matrices into multiple tiles, sign them to thread blocks, and execute computations for each thread block in parallel. Common paradigm within single thread block follows three key steps: Load matrix tiles into the GPU's shared memory; perform computations using the loaded tiles; and store computed results to the out- put tensor. (ii) Warp-level parallelism: Within each thread block, optimized kernels further partition assigned matrix tiles to multiple warps - each containing 32 threads on NVIDIA GPUs (NVIDIA) 2024d) to enable sine-grained parallel execution. Warp-level parallelism maximizes hardware efficiency through coalesced memory access and implicit synchronization within warps.
Efficient kernel implementation. Modern GPU architectures impose strict requirements on the shapes of matrix tiles used in low-level computations. Specifically, PTX warp-level matrix multiply- accumulate instructions (NVIDIA (025) require that for matrix multiplication A E Tm and B E R the dimensions m. n. and & must satisfy minimum size requirements for single-warp processing. On NVIDIA Hopper GPUs, m, n, k must be least 8. To achieve higher efficiency, a thread block typically utilizes multiple warps for sufficient warp-level paral- lelism. Additionally, modern GPUs perform optimally with coalesced and contiguous data loading and storing; non-contiguous memory access leads to lower L2 cache hit rate, thereby reducing effective memory bandwidth and degrading overall kernel efficiency.
AB, where
2.2 ATTENTION MECHANISMS
2023 where each query token attends to all previous KV tokens-inn is standard in LLM training and inference. Formally, given sequence length N, query/key head dimension dx. value head dimension dv, h query heads, and h KV heads, attention computation involves query/kee/vvlue tensor Q E R xdx xh, K E R.N XNxxxxbb xbx V [g.xdy For j-th (j {0, 1,.,, h h-1) 1}) query head, [i/hg] --th (ranging from O to hg-1) key and value head, denote involved matrices Q' Kb//k/] R Nxdx, Vli/hg] RN xdv Full attention computation can be formalized as:
Full attention. Full attention with causality (Vaswani et al. 2017 Ainslie et al.
(K(KL//kk))
O Softmax
y VLi/kk]
(L)
vdx
K Kwon et al. 2023) has optimized full attention from various perspectives. Notably, Flash Attention DDo] (023) optimizes full attention with a two-level loop: Each thread block loads block of query tokens and, while KV tokens remain, iteratively pro- cesses a block of KV tokens and accumulates intermediate results with online softmax (Milakov & Gimelshein] 2018). Results are finally written to the output tensor. This design minimizes redundant memory accesses for query and output tensors, thereby reducing attention execution latency.
On the system side, recent research (Dao 2023
Sparse attention. Recent efforts in sparse attention algorithms (Yuan al. 2025 Lu al. 2025
Lee et al., 2023; TTa et al. 2020 Zhao et al. 2019 Tang et al. 2024 Xiao et al. 2024b Zhu
al. 2024; Lai et al., 2025 Xu et al. 2025b; Zhang al. 2023) and system side optimiza- tions efforts /Zhng et al. 2024h a 2025b) represent an emerging trend aimed at reducing atten- tion computation costs in long-contex LLM training and inference, where standard attention per- forms poorly due to its quadratic complexity with respect to sequence length. The most notable efforts in sparse attention include Native Sparse Attention (NSA) (Yuan et al. 2025). Formally, in NSA, for j-th query head, each query token q Rixdx, J E {0, 1,.,...., N 1} attends to N < N KV tokens via three attention mechanisms e E C, where C = {emp, sel, win). repre- senting compression, selection, and sliding window for keys and values. We denote KV tokens as [//// E Nxdx, V/j/kx] E Nxdy, which contains [j/hx] J-th KV head and a subset of KV tokens of attention mechanism c. Given trainable gating scores of E [0. 1] for three attention modules, NSA combines the three attention mechanisms as follows:
(
( (K (K/////kll
of
Σ of Softmax
V Vi/hkl
Vdx
Notably, the NSA kernel that selectively retains fine-grained tokens is a major system bottleneck across three attention mechanisms. This point is validated in $4.4 The NSA kernel allows each
3

---

Published as a conference paper at ICLR 2026
query token across query heads that share the same KV heads to attend to distinct T KV blocks, with Bx contiguous KV tokens. Distinet KV block selection imposes challenges on effec- tively batching query tokens and performing computation with KV blocks within one thread block. Therefore, it is crucial to optimize the batching strategy for efficient NSA kernl execution.
3 FLASH SPARSE ATTENTION
We present FSA design and compare with vanilla NSA (33.1), then introduce FSA implementation and optimizations ($3 (32).. Finally, we provide a thorough analysis of FSA performance (3.3
3.1 FSA KERNEL DESIGN
An efficient sparse attention kernel must translate theoretical FLOPs reduction into concrete savings in memory access and computation during GPU execution. Vanilla NSA kernel is insufficient in achieving this goal. As illustrated in Figure (left), NSA kernel processes query tokens one by one in the outer loop and KV blocks in the inner loop, while batching query heads. However, if the number of query heads is insufficient, this method requires padding to meet the hardware's 's matrix multiplication shape requirements, leading wasteful memory access and computation.
To achieve higher kernel efficiency, FSA exchanges NSA kernel loop order and processes query heads one by one, looping over KV blocks in the outer loop and batches of query tokens in the inner loop. Since the number of such tokens is typically large enough to meet hardware requirements, this strategy requires no padding and eliminates the overhead of processing padded data.
However, due to inversion of kernel loop order, FSA encounters new challenges:
Non-contiguous memory access for query batches. Due to the sparse nature of NSA token selection, for one KV block, only subset of total query tokens is involved for attention compu- tation, and query token indices are typically non-contiguous. When processing query tokens in FSA inner loop, it is critical minimize the negative impact of non-contiguous memory access.
Online softmax statistics and attention results accumulation. Online softmax and attention results reduction for each query token across distinct KV blocks adds another layer of complexity. the NSA token selection logie, computing the final output for a query token requires accumu- lating partial attention results from its distinct selected KV blocks. Since the NSA kernel's outer loop iterates over query tokens, this accumulation process can be handled within one thread block. In contrast, FSA's inverted loop order means that partial results for single query are computed across different thread blocks, each processing a different KV block. This design necessitates a proper management strategy for accumulating attention results distributed across thread blocks.
3.2 FSA KERNEL IMPLEMENTATION AND OPTIMIZATION
To implement an efficient FSA kernel, we employ an optimized token selection kernel that mini- mizes the negative impact of non-contiguous memory access. Additionally, an online softmax and reduction kernel are designed to efficiently handle online softmax and attention result reduction.
FSA token selection kernel. FSA mitigates the impact of non-contiguous memory access by em- ploying index tensors to orchestrate data movement. During forward pass, as illustrated in Figure] [ (right), each thread block in FSA kernel is assigned single (Query Head, KV Block) pair. The KV block is loaded from main memory once per thread block. The kernel then iterates through batches of non-contiguous query tokens, which are loaded and stored using index tensors I, and O, for i E 21,2.... b}, where b is the total number of KV blocks. These index tensors are pre-computed from the NSA sparse selection tensor E Ax X which records selected KV block indices for each query token. Due to the sparse nature of token selection, each KV block is attended by a subset of query tokens. Consequently, index tensor Z,, which contains query token indices attending to current KV block, typically holds fewer than N valid indices, i.e., NNalid ZZ, < N. To minimize the impact of non-contiguous memory access, a thread block terminates early once it has processed all valid query indices in I, avoiding further memory access computation. Concurrently, index mapping tensor O facilitates contiguous storage of intermediate results. Note that outputs from FSA token selection kernel not final attention scores; they partial results that are reduced
4

---

Published as a conference paper at ICLR 2026
Inner Loop
Loop
dg
V
N x.
N'x dy
V
Nix dig X
N'x
Kyteck:
Nob-contiguoiis Batching
Quacc Rg dx
Load
Load
Load
Inner Loop
Grid Loop
door dnd
Valook dy
Crld
Compute SM
Compute
Store KV Black Buffer
Store HBM
Store via index mapping
Output xdy
Oudt
Nealid x.ddy
Grid Loop
Loop
Figure I: Left: Illustration of NSA kernel Yuan al. 2025 which iterates query tokens in outer loop, and processes KV blocks the inner loop. Right: Illustration FSA kernel, which alternatively iterate KV blocks in the outer loop, and processes query tokens in the inner loop partial attention results are stored in output buffer Our for accumulation (sce $3.2 for more details).
for each query across different KV blocks in separate reduction kernel, which introduce next. In the backward pass, FSA kernel follows similar logic, loading query tokens non-contiguously to compute gradients and storing intermediate gradients buffers. The primary difference that index tensors I₁ and O, computed during the forward pass, are retrieved from cache.
FSA handles query attention results and gradients reduction in separate kernels. In forward pass, FSA parallel computation of attention scores where a single query token's 's results are reduced across multiple KV blocks requires careful implementation of online softmax and reduction logic to ensure numerical correctness. In the backward pass, a similar reduction challenge exists for gradients query tokens. FSA achieves efficient and correct accumulation in two kernels:
FSA reduction kernel. Since a query's attention scores or gradients are computed across multiple thread blocks (each processing a different KV block in FSA token selection kernel), direct reduction into the output tensor in FSA kernel necessitates atomic additions (NVIDIA) 2024a) to prevent race conditions. Given the prohibitive overhead of atomic operations, FSA decouples computation from accumulation. It adopts two-stage process:
(i: FSA token selection kernel (see Figure (right)) computes partial query attention results or gradients without reduction with online softmax and writes them intermediate buffer.
(ii): A dedicated reduction kernel efficiently accumulates these partial results into final output tensor with online softmax scaling, which we introduce next.
This two-stage arrangement effectively eliminates atomic operations and achieves efficient attention result accumulation. However, HBM memory overhead is increased due to intermediate buffers. To minimize memory overhead, allocate buffer sized only for Nealid query tokens relevant each KV block, rather than for all N tokens. Index mapping tensor O, facilitates contiguous I/O into this compact buffer, thereby avoiding the significant overhead of allocating full-sized buffer for each KV block. We present detailed analysis of FSA buffer HBM memory overhead in Appendix E
FSA online softmax kernel. the forward pass, ensure numerical correctness, FSA needs to include online softmax statistics in two aspects:
ii) In the FSA token selection kernel, computation results between each query token and key block must be scaled with historical running maximum Miiakovv & Gimelshein 2018)).
(i): In the reduction kernel, partial attention outputs of query tokens regarding selected KV blocks stored in the output buffer must be scaled with online softmax statistics. Additionally, final output for query token must be scaled with log sum exponentials (Milakov & Gimelshein, (018).
Computing online softmax statistics within the FSA token selection kernel produces incorrect atten- tion results. When multiple thread blocks process the query token, each block computes only
5

---

Published as a conference paper at ICLR 2026
partial statistics, leading incorrect maximum values and attention outputs. To address this chal- lenge, FSA introduces a separate online softmax kernel that compueess online softmax statistics using query and key tensor Q and key tensor K and stores them in buffer.
3.3 FSA PERFORMANCE ANALYSIS
We analyze FSA performance by answering two critical questions regarding FSA and NSA kernel performance:
Normittee Voltme
Nemory Access
Compute
Norma tbee FLOPs
11:54
2.66
Do additional auxiliary kernels like online softmax and reduction implemented in FSA incur additional memory access and computation overhead?
Question 1:
1.33
0.00
Hoods
Heads GOA
Figure 2: Comparison on memory access FLOPs, block size is 64, top-k is 16. FSA's 's memory volume or FI.OPs are normalized to 1.
To answer this question, we conduct detailed memory footprint and computation load analy- sis and derive the following theorem:
Theorem:
Across popular GQA group settings. where each GQA group contains E {1, 4,8) 8} query heads, aggregate memory access volume and FLOPs of FSA token selection, online softmax, and reduction kernel are lower than vanilla NSA kernel. Comparisons are presented in Figure Additional memory access introduced by auxiliary kernels, i.e., FSA online softmax and reduction kernels, remains manageable, falling significantly below memory access wasted on padded data in the original NSA kernel (see more details in Appendix
Since FSA introduces non- contiguous memory access on loading query tokens and requires additional auxiliary ker- nels, is FSA generally applicable across vari-
Question 2:
A300
H100 PCIe
H100 NVL
Normalzed Overhea
2.27
1.52
1.19
1.39
GPU types, and does it consistently provide performance improvements over NSA kernels?
4
2.44
2.56
2.57
To answer this question, conduct group of herrharksss and enumerate the follow- ing analysis of empirical results:
Nomatted
Overnoso
1.83
.2
2.68
0.00
0.00
0.00
Group
GGA Group
Group
Empirical analysis: Profiling results (shown in Figure 3 across various GPU types and GQA group settings confirm superior perfor- mance of FSA. Optimized FSA outperforms vanilla NSA across popular GPU architectures and GQA group settings, despite being compro-
Figure 3: Real-time profiling results of the FSA and NSA kernel execution overhead across different GPUs, under block size Bx 64, and top-k value T 16. FSA latency is normalized
mised by non-contiguous memory access, and reducing attention results in a separate kerne.. When each GQA group contains fewer than 8 query heads, FSA usually demonstrates superior perfor- mance to NSA. These empirical results demonstrate that FSA kernel's performance gains from overall reduced unnecessary memory access and FLOPs more than compensate for the overhead of non-contiguous memory access and executing multiple kernels.
EVALUATION
This section presents comprehensive evaluation of FSA across various NSA configurations. We aim to investigate the following research questions:
Ql: What is the kernel-level performance of FSA compared with NSA and full attention across diverse NSA algorithmic configurations?
Q2: What is the impact of FSA on end-to-end eo-end training and inference performance in practice?
Q3: What is the breakdown performance of FSA, and how effective is each part of FSA?

---

Published as a conference paper at ICLR 2026
NSA
Attention
GQA-1,
H20
GQA-2,
GQA-4,
H20
GQA-B, (64,
3446
3446
3446
3446
Labency (rvs))
2585
2585
2585
3,333
1.3x5.55
2585
1723
1723
1723
1723
16
32
16
32
64
8
16
32
64
16
32
64
GOA-1, (128, 8), H20
GQA-2, (128. 81, H20
GOA-4, (128. 81. H20
GQA-B, (128. B1, H20
3446
3446
3446
3446
Latenry (mss)
2585
1723
24x0.X
1723
2585
0.5
2.2x2.6H
2.2%4.99
2585
1723
1.7x3/4x
2585
1723
1.0x0.9x
11.0x
862
862
862
862
GQA-1, (64, 16), H200
GQA-2, (64, 16), H200
GQA-4, (64, 16), H200
GQA-B, (64, 16), H200
1315
1315
1315
1315
Latoncy (mh))
986
986
986
986
16:21
1.000.00
329
329
329
329
16
16
16
GOA-1, (1.28, 81, H200
GOA-2, (128, B1, H200
GOA-4, (128, 8), H200
GOA-B, (128, 8), H200
1315
1315
Latenry (mos)
657
986
1.5% 2
11100,6 6
986
657
2.23
657
986
10000
1.2x
2.78
986
657
1.00 0.70
1.0x
329
329
329
329
16
32
64
8
16
32
64
8
16
32
64
166
32
64
Soal en (K)
Seqien (K)
Seqien (K)
Scalen (K)
Figure 4: Performance comparison of Triton-based FSA, NSA, and full attention (enabled by Flash Attention) kernels under block sizes and lop-k values of (Bx, T) equals to (64, 16) and (128, 8).
4.1 EXPERIMENTAL SETUP
Experimental setups. We use two GPU types for evaluations: NVIDIA H20 GPUs (NVIDIA 2024b), which provide 148 TFLOPS tensor core computational power and 4 TB/s memory band- width; and NVIDIA H200 GPUs (NVIDIA (024G), which deliver 989 TFLOPS tensor core compu- tational power and 4.8 TB/s memory bandwidth. For end-to-end training and inference evaluations, GPUs are interconnected via NVLink, providing 450 GB/s inter-GPU bandwidth. In our evaluations, use BF16 for training and FP16 for inference.
Baselines. We compare FSA with two basclines:
NSA (Native Sparse Attention) (Yuan et al. 2025). Our primary baseline is vanilla NSA imple- mentation, which introduces natively hariware-aligned trainable sparse attention. NSA maintains algorithmic performance comparable to full attention while substantially reducing computational complexity. We utilize Triton-based NSA kernel (Organization) (024) for evaluation.
Full attention (Flash Attention) (Dao, 2023). Due limited hardware resource utilization, theo- retical FLOPS reductions achieved by NSA or FSA may not translate to proportional performance gains. Therefore, the full attention baseline (with causality), which has sparsity constraints, is essential to demonstrate the practical effectiveness of both NSA and FSA. We utilize an efficient Triton-based Flash Attention kernel 2024) for fair comparison.
Experimental configurations. To ensure comprehensive evaluation, we systematically test FSA and two baselines under varying NSA configurations: ( GQA settings g E 21, 4, 8). where g is number of query heads in one GQA group; (ii) NSA hyperparameter block size Bx and top-k T combinations of (Bx, T E {644, 16), (128, 8)): and (iii) sequence lengths of {8K, 16K, 32K, 64K) tokens." For end-to-ens training and inference evaluations, we evaluate performance using Llama3-
8B (Dubey et al. 2024), Qwen3-14B (Yang et al. 2025), and Qwen2.5-32B [Team] 2024) with
2 MMore experiments on ultra long sequence lengths, i.c., 128K and 256K sequence lengths, are presented in
Appendix F
7

---

Published as a conference paper at ICLR 2026
sequence lengths of 32K and 64K. When the entire model is too large fit single GPU for training, we use pipeline parallelism (Shoeybi et al. 2019 to distribute model across multiple GPUs.
Evaluation metries. Following established practices in prior research (Yuan et al. 2025 Lu et al. 2025| Dao 2023), we employ two metrics to evaluate system efficiency: (i) Kernel exccution latency, which measures computational time required for attention operations, and (ii) training and inference latency, which measures end-to-end time required to process a single batch of data during model training and inference. These metries directly assess FSA's 's computational efficiency.
4.2 FSA KERNEL BENCHMARKING RESULTS (Q1)
FSA A kernel performance. We evaluate the kernel performance of FS A across both H20 and H200 GPUs under various configurations. In this section, we evaluate FSA on one single GPU, while we present distributed evaluations of FSA in Appendix I As shown in Figure + the evaluation results demonstrate that FSA outperforms both NSA and full attention across most of the tested scenarios:
Comparison with NSA. FSA outperforms NSA with significantly lowered memory access vol- ume and FLOPs in NSA token selection module, despite introducing non-contiguous memory access and auxiliary kernels (see details in $ FSA achieves up 3.5x speedup and on average x lower kernel latency on H20 GPUs, and up to 2.9x speedup and on average 1.4x lower ker- nel latency on H200 GPUs compared to NSA. Performance gap between FSA and NSA widens with smaller GQA group settings (g E {1, 2 and longer sequence lengths (32K and 64K tokens), with peak performance improvement of 3.5x observed at g 1 (one query head in one GQA group) and sequence length of 32K tokens. Furthermore, FSA maintains consistent performance improvements across different NSA algorithmic configurations, e.g., where (Bx, T) (64, 16) and (Bx, T) = (128, 8), demonstrating robust efficiency gains across diverse parameter settings.
Comparison with full attention. For long sequences, FSA outperforms full attention with an ef- ficient NSA algorithm and even more efficient token selection. FSA achieves up 6.4 speedup and on average 2.4x lower kernel latency on H20 GPUs, and up to 4.9x speedup and on average 2.3x lower kernel latency on H200 GPUs compared to full attention. Performance gap between FSA and full attention increases dramatically with a larger number of query heads in each GQA group. with the most substantial improvement of 6.4x observed at g 8 (8 query heads in one GQA group) and sequence length of 64K tokens. Similarly, FSA maintains superior efficiency across (Bx, T) E (64, 16), (128, 8)) settings, demonstrating consistent and substantial perfor- mance advantages over full attention. On the other hand, vanilla NSA lags behind full attention in many tested cases, even with its sparse attention mechanism. For example, when the sequence length is 32K, one GQA group contains one query head, NSA consistently falls short of full at- tention, while FSA demonstrates superior performance than full attention.
4.3 END-TO-END PERFORMANCE COMPARISON (Q2)
End-to-end training performance. We benchmark end-to-end training performance of FSA against NSA and full attention across various models and hardware setups. As shown in Figure 5 results demonstrate that FSA consistently reduces training latency across all evaluated cases. Specifically, FSA achieves up to 1.25x speedup and on average 1.09x speedup compared to NSA, and delivers up 2.47 speedup and an average of 1.86x speedup compared to full attention. These efficiency gains are pronounced with longer sequences and on higher-performante hardware like the H200, demonstrating FSA's 's effectiveness in accelerating computation-intessii training scenarios.
Inference performance. For prefill latency, we benchmark FSA against NSA and full attention across various models and hardware setups. As shown in Figure our results demonstrate that FS A achieves lower prefill latency across most evaluated configurations. Specifically, FSA achieves up to 1.36x speedup and on average 1.11x speedup compared to NSA. FSA performance advantages even more significant when compared to full attention, where FSA delivers up to 1.69x speedup and an average of 1.39x speedup. Taken together, these results underscore FSA'S 's efficacy in ac- celerating the prefill phase of LLM inference In terms of decoding latency, FSA matches that of NSA, which reduces memory access of the decoding phase by only loading a sparse subset com-
We present more detailed decoding evaluations in Appendix G
8

---

Published as a conference paper at ICLR 2026
P5A
MNN NSA
MSaa Pull Altention
FSA
NSA
VSEA FUE Allendion
H20, Unmm3-65 65
H20. Qwen3-148 140
H20. Owen2 2-3-22 (64,15)
+20. Unimal-6 6-33--3 (b4.16)
H20. Qwen3-14B 140
H20. Owen2.5-322 (c4,16)
(b4.10)
(54, (1))
(D4,
35
88.5
2562
33 22.7
64.1
104.2
192.1
37.0
27
51.0
42.8
60.5
128.1
11.4
18
34.0
0.0
6.0
32
64
52
32 H200, Lamm33 68 (MM166)
32
64
Uamma3
H200, OwerS
+200, 2.
H200,
Qwen225 328
(64.16)
(84,18)
(64
(M1)10)
(64,16)
10.6
Uiverry
21.2
612
15.5
14.2
40.9
110
28.4
0.0
32
32
52
32
64
Latency
88.5
(125,0)
(2200.6) 60
256.2
11201,00 (91
Leten S 22.7 17.0 11.4
(120.8)
8 27.6
(1240,00
67.9
(1
164.
104.2
192.1
51.
82.8
60.5
128.1
18.4
34.0
22-4
34.7
$4.0
17.0
32
64
32
32
64
32
6-1
H200.
Liama 88 (128,8)
H200,
(128,0)
(126.00
1129,88 RI
(128,0)
(129,8)
Laterry
21.2
13.1
01.2
Leawery
8.0
16.5
14.2
22.1
40.6
11.
7.1
28.4
32
54
32
64
64
64
64 Seaion 1K]
52
64
Soglen (K)
Soaian 1K1
Seajen 1k]
Soolon (K)
Sagien(k)
Figure 5: End-to-end training latency of FSA, NSA, full attention.
Figure 6: Inference Prefil latency of FSA, NSA, full attention.
posed of compressed tokens, selected tokens, and recent tokens from sliding window (Yuan al. 2025
4.4 PERFORMANCE BREAKDOWN & ABLATION STUDIES (Q3)
In this section, we evaluate FSA at both kernel and end-to-end (training or inference) levels. At the kernel level, we analyze forward and backward performance separately, and examine each of the three attention mechanisms within NSA: Compression, selection, and sliding window on key/value tokens. We conduct ablation studies to assess the effectiveness of FSA kernel optimizations. We validate the implementation correctness of FSA by comparing training loss across FSA, NSA, and full attention in Appendix [
Forward and backward breakdown. We conduct a detailed breakdown to analyze forward and backward attention computation latencies of FSA, NSA, and full attention across various NSA con- figurations. As shown in Figure 7 FSA demonstrates superior performance in both forward and backward attention computations across all evaluated scenarios. For forward computation, FSA achieves up to 2.36x speedup and average 1.62x lower latency compared to NSA, and up to 3.23 x speedup and on average 1.83 lower latency compared to full attention. Backward compu- tation analysis reveals even more pronounced advantages, since FSA avoids computation costs for index tensors I. O2 for i-th KV block (see details in $3.2). FSA achieves up to 4.32x speedup and on average 2.59x lower latency compared to NSA, and up to 7.45x speedup and on average 6.89x x lower latency compared to full attention. Performance improvements remain consistent across dif- ferent NSA configurations, demonstrating that FSA provides robust efficiency gains.
Compression, selection, and sliding window breakdown. We conduct detailed breakdown exper- iments for the three essential steps in NSA. As demonstrated in Figure g the token selection phase dominates overall attention computation performance, accounting for up to 79% and on average 65% of total attention overhead across all evaluated configurations. And FSA achieves substantial performance improvements in token selection, delivering up to 7.6x speedup and on average 3.4x lower latency compared to NSA in this critical phase. These results highlight that FSA's primary performance advantages stem from its efficient handling of token selection computation.
Ablaioon study on sparse attention performance. We present an ablation study of FSA kernel performance in Figure 2 where we disable each of additional optimizations of FSA we mentioned in 3 Results demonstrate that by disabling the inner loop (one thread block for one query batch), performance of FSA kernel drops by up to 18.9% and average 11.9%, and by disabling early turn optimization, performance drops by up to 25.2% and on average 18.2%. These empirical results demonstrate the importance of each component of our FSA optimization in enhancing performance.
9

---

Published as a conference paper at ICLR 2026
SF Selecled (NSA)
Compressed (FSA) Compresses (NSA)
Slding IFSA]
FSA-Ned PSM-bwd
NSA/wd NSSdww
Full Accention-fed Pull Attentlon-hed
Sliding INSA)
mIAAm1 (64.19)
60 (M1,10
GQA=1
H200, GOA-4
(64.10)
(64-16)
(64,19)
(64,16)
[65.189
[54,18]
L30ne 15
2.6
1.3
0.3
6.2
0.1
1.7
1.7
0.7
0.7
0.1
0.2
0.1
0.9
0.0
0.0
Swd bed
Ned
bad
fed
bad
Literey t4
3.4
H20, G0Am1
(12%)
3.4
H20, G0Am
(128,
1.3
H200, GOA=1
1128,
K200, 50A-
(126.
1.5
1.1
H20, Gi0te1 (221,0)
0.6
H20, Si0kmm
(128,00
H200, GgA-1
(1286.)
0.3
1200, coAm
(128,0]
1.7
177
0.7
0.1
0.0
0.9
0.3
0.3
0.0
0.0
0.0
0.0
0.0
fed
bwd
Figure 7: Experimental breakdown of FSA, NSA, and full attention latencies during forward and back- ward computation.
Figure 8: Experimental breakdown of token compres- sion, selection, and sliding window attention over- head during forward/backward pass.
MIP (FSA)
MLP (NSAI
MLP [Ful Actention) Atlentin iFul Attantion)
Loop
Attandion (FSA)
Altencion (NSA)
H20 G0A-1
H20, G0A-4
H200 G0A-1
P200, SQA
H20 Liama3 48
H20. Owend-eeeee 14B
H20, Owen2. S-328
Seolian
560) an-325
Seglen=SaK
(128.8)
1128,00
(128,8)
0.25
0.16
0.15
0.10
78.4
127A
254.8
58.8
95.5
176.1
Latercy
0.13
0.08
0.08
0.05
14.2
19.4
31.8
63.7
117.4
58.7
0.00
0.00
0.00
0.00
0.0
0.0
64
0.0
64
120. GDA-1
120, GOA-4
1100 GDAM)
1200, SQX -1
M200, Liamas-00
M200. Owerd-148
M200, Owan2.5-33 2.3-118
0.50
Soa
0.31
Soo)
0.50
0.10
So0lon=64K
25-8
19.5
403
304
74.8
13.0
202
0.25
0.16
8.1
0.00
10.1
18:7
0.8
0.07
32
64
32
64
32
64
0.00
0.00
0.00
0.00
Scatih
Sealon
Scalch
Figure 9: Ablation study (with or without FSA opti- mizations) on FSA kernel.
Figure 10: Breakdown of computation time for atten- tion and MLP during end-to-end training.
End-to-end training breakdown. To isolate the source of performance improvements, we conduct a breakdown analysis of the end-to-end training latency. As shown in Figure 10, results demon- strate that FSA's performance improvements originate from attention computation. Within this component, FSA achieves up to 1.4x and on average 1.23 lower latency than NSA, and realizes speedup of up to 3.87 x and on average 2.91 over full attention. This analysis confirms that overall end-to-end speedup is driven by FSA's fundamental optimizations in NSA token selection.
CONCLUSION
We presented Flash Sparse Attention (FSA), kernel design that broadens the applicability of Native Sparse Attention (NSA) to modern LLMs where each GQA group contains a small number of query heads. By inverting kernel loop order and introducing tailored optimizations for non-contiguous memory access, online softmax, and accumulation, FSA eliminates padding inefficiencies that limit NSA on current GPUs. Evaluation demonstrates that FSA achieves substantial improvements in both kernel-level and end-to-end performance, offering consistent speedups in training/inference across state-of-the-art long-context LLMs. These results highlight that algorithm- systemmm--ymmmmm---e--m co-design is critical for translating theoretical efficiency of sparse attention into practical acceleration. We believe FSA provides a foundation for future exploration of hardware-efficient sparse attention.
ACKNOWLEDGMENT
This work is supported by the HKUST startup grant R9895 from CSE; RGC-ECS project 26218024; RGC-NSFC project CRS HKUST601/24.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Shantanu Acharya, Fei Jia, and Boris Ginsburg. Star attention: Efficient Ilm inference over long sequences. arXiv preprint arXiv:2411.17116, 2024.
Saurabh Agarwal, Bilge Acun, Basil Hosmer, Mostafa Elhoushi, Yejin Lee, Shivaram Venkatara- man, Dimitris Papailiopoulos, and Carole-Jean Wu. Chai: clustered head attention for efficient Ilm inference. In Proceedings of the 41st International Conference Machine Learning, pp.
291-312, 2024.
Joshua Ainslie, James Lee-Thorp, Michiel De Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head check- points. arXiv preprint arXiv:2330.. 23055 2023.
Anthropic. The claude 3 model family: Opus, sonnet, haiku, 2024. URL https://wwwccdd anthropic com /eee/deebbaa 9b008babbll00 9b01c9bbbbbb... htt.o:olcc.ccccaaaaa3aa33600.... htwp/////bbbb6666666666///000 Caro_
pdf
Yaofo Chen, Zeng You, Shuhai Zhang, Haokun Li, Yirui Li, Yaowei Wang, and Mingkui Tan. Core context aware transformers for long context language modeling. arXiv preprint arXiv:2412.12465. 2024.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint wwXivv2307.086911 2023.
Tri Dao, Daniel Haziza, Francisco Massa, and Grigory Sizov. Flash-decoding for long-context inference, 2023. URL
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, al. The llama 3 herd of models. arXiv e-prints, pp. arXiv-2407, 2024.
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin-Yev Lin, et al. Minference 1.0: Accelerating pre-filling for long-context Ilms via dynamic sparse attention. Advances in Neural Information Processing Systems, 37:52481-52511 2024a.
YOUHE JIANG, Ran Yan, and Binhang Yuan. Hexgen-2: Disaggregated generative inference of Ilms in heterogeneous environment. In The Thirfeenth International Conference Learning Representations.
Youhe Jiang, Ran Yan, Xiaozhe Yao, Yang Zhou, Beidi Chen, and Binhang Yuan. Hexgen: gener- ative inference of large language model over heterogeneous environment. In Proceedings of the 41 International Conference on Machine Learning. pp. 21946 21961, 2024b.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th symposium on operating systems princi- plex, pp. 611-626, 2023.
Xunhao Lai, Jianqiao Lu, Yao Luo, Yiyuan Ma, and Xun Zhou. Flexprefill: A context-aware sparse attention mechanism for efficient long-sequence inference. 2025.
Heejun Lee, Jina Kim, Jeffrey Willette, and Sung Ju Hwang. Sea: Sparse linear attention with estimated attention mask. arXiv preprint arXiv:2310.01777, 2023.
Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulum Du, Tao Jiang, Chao Hong, Shaowei Liu, Weiran He, Enming Yuan, Yuzhi Wang. et al. Moba: Mixture of block attention for long-context 1lms. arXiv preprint arXiv:2502.13188 2025.
Maxim Milakov and Natalia Gimelshein, Online normalizer calculation for softmax. arXiv preprint arXiv: 1805.8005000 02867, 2018.
11

---

Published as a conference paper at ICLR 2026
NVIDIA. Cuda C++ programming guide.
https://docs wtw.cs..viiiaacccm/uuuuuu
cuda C- programming--uiide 2024a. Section on Atomic Functions.
NVIDIA. Nvidia h20 solution brief. http:p//iiaagess../aaaaaaoo/oooo
dgx-
2024b.
NVIDIA. H200 tensor core gpu. httppp//www..wddaaaacc./rrrrrrrr
h2c0/ 2024c.
NVIDIA. Cuda c++ best practices guide.
http:://oc3.nv.diaa.oom//
httpa--eeestpprattt-e--------ee- 2024d.
NVIDIA. Parallel thread execution isa version 9.0 - warp-level matrix instructions.
execution/index
html htt.lhaaapp-lleaaaanntttiiissita 2025.
OpenAl. Openai gpt-4o. 2024. httppppp//paatto.......aaaco////
gpt-4a
FLA Organization. Native sparse attention.
https::/github..om/fllee--
sative-spaare-atteeettt 2024.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model par- 2019.
allelism. arXiv preprint
Connor Shorten. MI-arxiv- Ml-arxiv-papers, AIL-ArXiv-Pae Papivs 2024.
Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. Quest: query- aware sparsity for efficient long-context Ilm inference. In Proceedings of the 41st International Conference on Machine Learning, pp. 47901-47911, 2024.
Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse sinkhorn attention. arXiv preprint arXiv:2002.11290 2020.
Qwen Team. Qwen2 technical report. arXiv preprint arX1v:2477.1000 2024.
Triton. Fused attention http://ppp///.t......oo..oo/o//.
Eutorial s/06- fused-attent ieniiiinn.tt. html
2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- fion processing systems, 30, 2017.
Xindi Wang, Mahsa Salmani, Parsa Omidi, Xiangyu Ren, Mehdi Rezagholizadeh, and Armaghan Eshaghi. Beyond the limits: A survey of techniques to extend the context length in large language models. arXiv preprint arXiv:2402.02244, 2024.
Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, and Maosong Sun. Infllm: Training-free long-context extrapolation for Ilms with effi- cient context memory, Advances in Neural Information Processing Systems, 37:1 173119638-116611 2024a.
Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Shang Yang. Haotian Tang, Yao Fu, Song Han, et al. Duoattention: Efficient long-context I1m inference with retrieval and streaming heads. In The Thirteenth International Conference Learning Representations, 2024b.
Chejian Xu, Wei Ping, Peng Xu. Zihan Liu, Boxin Wang, Mohammad Shoeybi. Bo Li, and Bryan Catanzaro. From 128k to 4m: Efficient training of ultra-long context large language models. arXiv preprint a 250iv:2504.068 0..04002244 14, 2025a.
Ruyi Xu, Guangxuan Xiao, Haofeng Huang, Junxian Guo, and Song Han. Xattention: Block sparse attention with antidiagonal scoring. In Proceedings of the 42nd International Conference on Machine Learning (ICMI.) 2025b.
12

---

Published as a conference paper at ICLR 2026
Ran Yan, Youhe Jiang, Xiaonan Nie, Fangcheng Fu, Bin Cui, and Binhang Yuan. Hexiscale: Ac- commodating large language model training over heterogeneous environment. arXiv preprint arXiv:2409.01143, 2024.
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 08t::25005003388. 2025.
Alex Young, Bei Chen, Chao Li, Chengen Huang. Ge Zhang, Guanwei Zhang. Heng Li, Jiangcheng Zhu, Jianqun Chen, Jing Chang, et al. Yi: Open foundation models by 01. ai. arXiv preprint arXiv:2403.044 04652, 2024.
Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, YX Wei, Lean Wang, Zhiping Xiao, et al. Native sparse attention: Hardware-aligned and natively trainable sparse attention. arXiv preprint (rXiv:2502.11080) 2025.
Organization Zai. Longbench: A benchmark for long-context language models., 2023. URL
httpps:/hugginggaaee
Jintao Zhang, Haofeng Huang, Pengle Zhang, Jia Wei, Jun Zhu, and Jianfei Chen. Sageatten- tion2: Efficient attention with thorough outlier smoothing and per-thread int4 quantization. arXiv preprint arXiv:2411. 10958, 2024a.
Jintao Zhang, Jia Wei, Haofeng Huang, Pengle Zhang, Jun Zhu, and Jianfei Chen. Sageat- tention: Accurate 8-bit attention for plug-and-play inference acceleration. arXiv preprint arXiv:2410.02367, 2024b.
Jintao Zhang, Rundong Su, Chunyu Liu, Jia Wei, Ziteng Wang, Pengle Zhang, Haoxu Wang, Huiqiang Jiang, Haofeng Huang, Chendong Xiang, Haocheng Xi, Shuo Yang. Xingyang Li, Yuezhou Hu, Tianyu Fu, Tianchen Zhao, Yicheng Zhang. Youhe Jiang, Chang Chen, Kai Jiang, Huayu Chen, Min Zhao, Xiaoming Xu, Jun Zhu, and Jianfci Chen. A survey of efficient attention methods: Hardware-efficient, sparse, compact, and linear attention. 2025a.
Jintao Zhang, Jia Wei, Pengle Zhang, Xiaoming Xu, Haofeng Huang, Haoxu Wang, Kai Jiang, Jun Zhu, and Jianfei Chen. Sageattention3: Microscaling fp4 attention for inference and an exploration 8-bit training. arXiv preprint arXiv:2505. 11594, 2025b.
Li Zhang, Youhe Jiang, Guoliang He, Xin Chen, Han Lv, Qian Yao, Fangcheng Fu, and Kai Chen. Efficient mixed-precision large language model inference with turbomind. arXiv preprint arXiv:2508. 15601, 2025c.
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, al. H2o: Heavy-hitter oracle for efficient gen- erative inference of large language models. Advances in Neural Information Processing Systems, 36:34661-34710, 2023.
Guangxiang Zhao, Junyang Lin, Zhiyuan Zhang, Xuancheng Ren, Qi Su, and Xu Sun. Ex- plicit sparse transformer: Concentrated attention through explicit selection. In arXiv preprint arXiv:1912.11637, 2019.
Qianchao Zhu, Jiangfei Duan, Chang Chen, Siran Liu, Xiuhong Li, Guanyu Feng, Xin Lv, Huanqi Cao, Xiao Chuanfu, Xingcheng Zhang, et al. Sampleattention: Near-lossless acceleration of long context IIm inference with adaptive structured sparse attention. arXiv preprint arXiv:2406. 24446,6,00..566 2024.
13

---

Published as a conference paper at ICLR 2026
A THE USE OF LARGE LANGUAGE MODELS
In this paper, we leverage LLMs to enhance academic writing quality by ensuring grammatical correctness and improving sentence structure.
B NOTATIONS
The notations used in this paper summarized in Table 1
Table Notations and Explanations.
Notation
Explanation
N
Sequence length.
dk
Head dimension for query and key tensor.
dy
Head dimension for value tensor.
d
Uniform head dimension, i.e., dg = dy.
h
Number of Q heads.
hg
Number of KV heads.
g
GQA group size, defined as g = hx
h
Number of selected KV blocks of each query token.
T
(Hyperparameter of the NSA sparse attention module.)
Bg
Block size of each KV block; NSA hyperparameter.
b
Number of KV blocks; b
7g
BQ
Query batch size in FSA; a FSA hyperparameter.
The set of query indices attending to the i-th KV block.
I
(I; contain non-contiguous query indices, usually [Z,] S N.)
The output tensor mapping for the i-th KV block; c.g., O|]]
O,
gives the storage position of token j in the output buffer.
Nyalid
The number of valid query tokens in Z4.
T
Sparse selected KV block indices in NSA.
Q.KV
Full query, key, and value tensor for attention computation.
Non-contiguous Query balches introduced in FSA.
Qhatch
(One thread block processes multiple Q batch-)
V
The i-th KV block with Bx contiguous KV tokens.
Intermediate buffer which holds query attention results
Obuf
without scaling with online softmax in FSA.
C FSA IMPLEMENTATION DETAILS
FSA is implemented using 10K lines of Python and Triton code. To optimize system performance: (i We apply fineggrained control over FSA selected attention kernel and reduction kernel to op- timize warp-level parallelism. FSA usually assigns warps per thread block for FSA selected attention kernel, which contains matrix multiplication operations, to enable sufficient computational resources of given thread block. FSA usually assigns kernel, which mainly consists of elementwise operations. Warp assignment for reduction kernel efficiently utilizes warp-level parallelism, reducing reduction kernel execution latency. (ii) We spec- ulatively compute online softmax statistics once per KV heads. Due to invariant nature of online softmax (Milakov & Gimelshein 2018), correctness of FSA is maintained, while significant cost for computing online softmax statistics amortized.
warps per thread block for reduction
D FSA CORRECTNESS
FSA correctness. To evaluate correctness of FSA kernels, we fine-tune Llama3-8B model using MI Ar---XXy-Paaers dataset (Shorten 2024). We replace attention module of Llama3-8B model with
14

---

Published as a conference paper at ICLR 2026
FSA
NSA
FA
2.0
1.5
Loss
1.0
0.5
0.0
200
400
600
800
1000
1200
1400
1600
1800
2000
Step
Figure 11: Loss comparison of FSA/NSA/full attention in end-to-end Llama3-8B training.
either FSA NSA, while initializing all other components with pretrained model checkpoints pro- vided by Meta. For fair comparison with full attention, we reinitialize the parameters the attention module. Loss comparison among FSA, NSA, and full attention presented in Figure 11 Results demonstrate that all three methods achieve stable and similar convergence, and FSA exhibits a sim- ilar loss curve NSA, validating the correctness of the FSA kernel.
Further analysis. To strengthen our accuracy evaluation, we conduct additional experiments by fine-tuning smaller models across diverse tasks. Specifically, fine-tune Llama-3.2- 3..1BB.B Llama- 3.2-3B, and Llama-3. 1-8B Dubey et al. (024) on three representative LongBench ZZai (2023) tasks: multi-document QA (MQA) on the HotpotQA dataset, single-document QA (SQA) on the Qasper dataset, and synthetic-date QA (Synthetic) on the PassR-EN dataset. For Llama-3.4 3.2- and Llama- 3.2-3B, we report the average loss after convergence over 2K training steps, comparing FSA, NSA, and Full Attention; these results appear in Table To more comprehensively assess accuracy preser- vation, we further evaluate perplexity and QA Fl across three models and tasks. The results, summarized in Tables and consistently show that FSA matches the accuracy of NSA and Full Attention.
Table Converged Loss Across Datasets Attention Modes.
MQA (HQA)
SQA (Qasper)
Synthetic (PassR-EN)
Model Size
FA
FSA
NSA
FA
FSA
NSA
FA
FSA
NSA
1B
0.200
0.182
0.187
0.216
0.191
0.184
0.231
0.224
0.231
3B
0.173
0.153
0.166
0.087
0.082
0.078
0.123
0.119
0.118
Table 3: PPL Across Datasets, Models, and Attention Modes.
MQA (HQA)
SQA (Qasper)
Synthetic (PassR-EN)
Model Size
FA
FSA
NSA
FA
FSA
NSA
FA
FSA
NSA
5.40
6.79
6.82
8.77
9.48
9.45
3.48
2.52
2.49
3B
2.42
1.50
1.48
1.20
2.64
2.62
1.87
1.94
1.96
8B
1.57
1.17
1.16
1.28
1.71
1.70
1.21
1.26
1.27
Table QA F1 Across Datasets, Models, and Attention Modes.
MQA (HQA)
SQA (Qasper)
Synthetic (PassR-EN)
Model Size
FSA
NSA
FA
FSA
NSA
FSA
NSA
IB
0.05
0.10
0.11
0.08
0.07
0.06
0.22
0.32
0.31
3B
0.28
0.35
0.33
0.15
0.11
0.12
0.39
0.47
0.48
8B
0.32
0.38
0.37
0.23
0.20
0.19
0.83
0.86
0.86
E
FSA AND NSA THEORETICAL MEMORY ACCESS AND FLOPS ANALYSIS
To demonstrate how FSA outperforms NSA selected attention, we analyze as follows. For simplic- ity, assume query/key/value have the same head dimension, i.e. d = dx dv.
15

---

Published as a conference paper at ICLR 2026
FSA analytic advantages. Theoretically, FSA introduces lower memory access volume and num- ber of floating-pont operations (FLOPs) for small GQA group sizes. We analyze FSA/NSA as follows:
FSA memory access volume and FLOPs. We analyze the three key components in FSA as follows:
FSA selected attention kernel launches hb thread blocks, where h is the number of query atten- tion heads, and b is the total number of KV blocks. For sequence of N tokens, the number of KV blocks = Bg where Bx is the KV block size. In one thread block, FSA selected attention kernel runs a two-level loop. In the outer loop, it loads 2B d KV tokens; in the inner loop, it iter- atively loads B&d query tokens, performs attention computation with FLOPs of 4B and stores Bod query attention results. We estimate the number of our inner loop as follows. Assume each query token attends to each KV block with equal probability. Therefore, each query token attends to a given KV block with a probability of f. resulting in an average number of tokens attending to a given KV block of NT, and an average number of query batches for one KV block 5
N
of NT. Assuming each data occupies 2 bytes, we can calculate memory accessed in bytes by b.Qg
FSA selected attention kernel as 4dh + TT), and FLOPs as 4dh BxT.
FSA online softmax kernel operates similarly to the FSA selected attention kernel, with three key differences: It is called per KV head, omits V tensor loading and computation, and interme- diate attention scores storage, storing only a single scalar value per (query token, KV block) pair. Following a similar estimation logic as FSA selected attention kernel, the online softmax kernel introduces 2dh k (N +1 T') memory access volume in bytes, and 2dh B&T FLOPs.
FSA reduction kernel introduces negligible FLOPs, but for each query token, it involves loading attention results of T KV blocks and storing the final attention results. Therefore, FSA reduction kernel introduces 2dhN(1 + T)) memory access in bytes.
In total, FSA incurs d.N (6h + 22g) + +T) memory access in bytes, and dN By + 2hx) FLOPs.
NSA Memory access volume and FLOPs. NSA selected attention kernel launches h& N thread blocks, where hx is the number of KV heads. In each thread block, NSA kernel runs a two-level loop. In the outer loop, NSA kernel loads one query token and g hie same KV head. Due to the hardware requirements on matrix multiplication shapes, when GQA < 8, NSA kernels must load 8 query heads (&d elements), perform computation, and mask out the undesired computation results. In the inner loop, NSA kernel iteratively (up to 'T' times) loads one KV block (2Bxd d elements) and performs attention computation with a FLOPs of 32 To maintain the causal property, i.e., avoiding query tokens to attend to future KV tokens, the actual number of KV blocks that need to be loaded and participate in computations within a thread block on average F. Finally, NSA kernel stores the attention results in the output tensor, incurring gd memory access. Therefore, we can estimate the memory access volume (2 bytes per data) for NSA kernel as 2dh N + g + 8). The FLOPs for NSA kernel are 32dh g N BRT.
_ Q heads that share the
FSA selected attention kernels exhibit lower memory access volume and FLOOPs. With (Bx, 7) (64, 16) and sequence length of 64K, which is the same configuration as presented in the NSA paper, we observe that compared to the NSA selected attention kernel, our method incurs lower memory cess volume and FLOPs for GQAS8, detailed comparisons are presented in Figure 2 In particular, for GQA-4, a common configuration in LLMs, our method theoretically reduces memory access volume to 21.3% and FLOPs to 56.2% of those in NSA. Benefits from the more efficient hardware- aligned kernel design, our method substantially outperforms NSA across various GQA group sizes. Additionally, our method demonstrates superior performance as the NSA hyperparameter B& in- creases. This advantage stems from NSA's inherent inefficiency with larger KV blocks. Although NSA can easily skip loading KV blocks that fully violate causal property, to maintain causality con- straints for KV blocks that partially violate causal property, NSA must mask out many KV tokens within the KV block, leading to wasteful memory accesses where loaded data is only partially valid for computation. As the KV block size BR grows larger, this inefficiency becomes increasingly pronounced, as greater proportion of the loaded KV block remains unused due to causal masking. In contrast, our method processes all query tokens that attend to a given KV block within a single thread block, naturally satisfying causal constraints without requiring extensive masking. This ap- proach achieves superior memory efficiency by ensuring that all loaded KV data contributes to the computation, resulting in significantly lower memory access overhead.
16

---

Published as a conference paper at ICLR 2026
FSA trade-offs. FSA trades lowered memory access volume and FLOPs with non-contiguous loading and more buffer overhead. Theoretical advantages of FSA come at the price of involving non-contiguous memory access and more buffers that occupy HBM memory. We analyze how these factors compromise FSA performance and how FSA optimizes memory access and buffer management as follows:
Optimize memory access. The non-contiguous loading on query batches, which is inefficient modemn GPUs, compromises FSA selected attention kernel performance. Modern GPUs usu- ally operate more efficiently under coalesced and contiguous memory access, which can improve the L.2caache hit rate and thereby kernel efficiency (NVIDIA 2024d). Therefore, the theoret- ical advantages of our method cannot be fully reflected in actual hardware, due to inevitably degraded performance of non-contiguous memory access. Nonetheless, to our best effort, FSA optimizes memory access with fine-grained early return mechanisms that filter out unnecessary query batches loading. For example, for i-th KV block, FSA compactly stores query indices in set I, which is computed via full index table. For each query token, the full index table records whether it should attend to i-th KV block, and Ij filters the tokens that do not attend to i-th KV block. Therefore, when all query tokens in I, are exhausted, FSA returns early.
Optimize buffer management. The newly introduced buffers, Obuf appeared in Figure T (right), bring memory overhead. FSA minimizes buffer overhead from two aspects: (i) FSA Token lection kernel processes a subset of query heads at each time, reusing the buffers for subsequent query heads computations. (ii) FSA introduces an output index mapping tensor store results compactly. For each query head, FSA only reserves buffers for maximum query tokens that attend given KV block. On average, this value is By T, combining that Bx . FSA introduces an output buffer with dNT elements. Assume each data in the output buffer occupies 2 bytes, for sequence with 64K tokens, T at 16, and d at 128, Obut occupies 1 GB HBM memory (This also applies for the buffer for intermediate gradients with respect to Q). Compared to the high HBM memory capacity in modern GPUs.e.g.. 96 GB HBM memory on H20 (NVIDIA) (024b) and 141 GB memory on H200 (NVIDIA 2024c), the additional buffer overhead in FSA remains manageable.
Attention Sink Optimizations. The attention sink phenomenon in NSA sparse token selection presents challenge for FSA's buffer management strategy. The initial KV block receives attention from all query tokens, while subsequent KV blocks exhibit more selective attention patterns. This asymmetry creates a buffer allocation dilemma: In practice, FSA allocates uniform buffer sizes based on the maximum number of valid tokens across all KV blocks. However, the attention sink property forces this maximum as full sequence length, thereby negating the memory efficiency gains that FSA's 's sparse buffer management is designed to achieve. To address this inefficiency, we im- plement a dual-buffer allocation strategy. We maintain separate buffer allocations for the attention sink (first KV block) and the remaining KV blocks. The attention sink buffer accommodates the full query sequence, while buffers for subsequent KV blocks are sized according to their maximum valid query tokens, which are usually much smaller than full sequence length. This approach preserves the memory optimization benefits for the majority of KV blocks while handling the attention sink's dense connectivity requirements.
FSA online profiling module. In real-world deployment, FSA dynamically selects kernel config- uration via online profiling, and potentially falls back to original NSA implementation. To ensure optimal performance across diverse NSA configurations, FSA incorporates a one-time online pro- filing mechanism. Upon its first execution with new set of hyperparameters (e.g., sequence length, GQA group size), FSA benchmarks its kernel performance across several candidate query batch sizes (e.g., 1, 64, 128). When GQA group size is sufficiently large, a query batch size of 1 is addi- tionally searched and serves as a potential fallback to original NSA strategy of batching query heads. Once profiling is complete, the fastest configuration is cached. All subsequent calls with the same hyperparameters directly use this optimal configuration, bypassing profiling step until hyperparam- eters change.
Actual memory footprint of FSA buffers. We conduct additional bircobbenchmarks to measure the memory footprint of FSA buffers. Concretely, we set the head dimension at 128 and use the NSA hyperparameters (Bx, T) (64.16) or (128,8), and report the profiled buffer overheads for sequence lengths ranging from 32K to 256K in Table [ Under extreme cases, ie., when the sequence length is 1 128K 256K, FSA introduces G. or T2.36 GB buffer memory overhead, which is still much
17

---

Published as a conference paper at ICLR 2026
smaller than the memory capacity of modern GPUs (e.g., H200 has 141GB memory). These results confirm that the FSA buffer memory overhead remains acceptable.
Table Profiled Buffer Overhead.
(Bx. 7)
Seqlen (K)
Profiled Buffer Overhead (GB)
(64, 16)
32
0.52
(64, 16)
64
1.88
(64, 16)
128
5.01
(64, 16)
256
12.36
(128, 8)
32
0,26
(128. 8)
64
0.91
(128, 8)
128
2.28
(128, 8)
256
6.15
F EVALUATIONS FOR ULTRA LONG SEQUENCE LENGTHS.
We extend our evaluations to 128K and 256K sequence lengths. Fixing the head dimension at 128 and the number of query heads at 64, while varying the number of key and value heads, evaluate configurations where a GQA group contains to 8 query heads. Using the NSA hyperparame- ters with (Bk. T) (64, 16) or (128, 8), benchmark the performance of FSA, NSA, and Full Attention (FA) on both H20 and H200 GPUs.
F.1 INFERENCE PREFILI AND TRAINING EVALUATIONS
The experimental results Table 2 and and show that FSA also outperforms NSA for ultra-long sequence lengths. For inference prefill execution latency, FSA achieves up to 1.47 x speedup and an average of 1.20x lower kernel latency on H20 GPUs, and up 1.86x speedup with an average of 1.23x lower kernel latency H200 GPUs, compared to NSA. For training execution latency measured over one forward and one backward pass FSA achieves up to 1.91 speedup and average of 1.37x lower kernel latency H20 GPUs, and up to 2.55x speedup with an average 1.49 lower kernel latency on H200 GPUs, relative to NSA.
Results discussion:
Table H20 GPU, Inference Prefill and Training Latency for Different (Bx, T)-
Seq Len (K)
FSA
NSA
FA
FSA
NSA
FA
(Bx, T)
GQA
Fwd (s)
Fwd (s)
Fwd (s)
F+B (s)
F (s)
F+ (s)
128
1.42
2.08
2.64
2.36
4.51
12.08
256
6.40
7.18
10.5
8.70
13.29
48.23
2
128
0.87
1.17
2.62
1.79
2.71
12.04
256
3.74
4.07
10.53
6.03
8.43
48.27
(64,16)
128
0.52
0.61
2.65
1.43
1.75
12.07
4
256
2,41
2.44
10.52
4.66
5.98
48.24
8
128
0.45
0.45
2.61
1.38
1.39
12.05
8
256
1.63
1.64
10.51
3.99
4.75
48.26
128
1.24
1.68
2.64
2.15
3.50
12.08
1
256
5.34
7.50
10.5
7.56
13.65
48.23
128
0.90
1.29
2.62
1.66
2.22
12.04
2
256
3.18
4.24
10.53
5.40
8.33
48.27
(128,8)
128
0.45
0.49
2.65
1.35
1.52
12.07
4
256
2.10
2.53
10.52
4.29
5.55
48.24
0.42
0.43
2.61
1.31
1,41
12.05
8
256
1.56
1.68
10.51
3.74
4.17
48.28
18

---

Published as a conference paper at ICLR 2026
Table 7: H200 GPU, Inference Prefill and Training Latency for Different (Bx. T).
Seq Len
FSA
NSA
FA
FSA
NSA
FA
(Bx, 7)
GQA
(K)
Fwd (s)
Fwd (s)
Fwd (s)
F+B
F+ + (8)
F+B
128
0.78
1.01
1.01
1.12
2.17
4.70
256
3.92
3.97
3.96
4.81
6.99
18.75
2
128
0.46
0.57
0.98
0.80
1.24
4.68
2
256
2.14
2.17
3,98
3.04
4.12
18.77
(64,16)
4
128
0.27
0.33
0.99
0.60
0.82
4.71
4
256
1.20
1.29
3,99
2.15
2.70
18.76
8
128
0.18
0.18
0.97
0.50
0.52
4.67
8
256
0.73
0.73
3.97
1.72
2.01
18.78
128
0.65
1.20
1.00
0.97
2.47
4.70
256
3.16
4.10
3,96
3.99
6.79
18.75
2
128
0.38
0.66
0.98
0.70
1.40
4.68
2
256
1.76
2.21
3,99
2.58
3,90
18.77
(128,8)
128
0.21
0.28
0.99
0.53
0.73
4.71
4
256
1.06
1.24
3.97
1.87
2.44
18.76
8
128
0.17
0.20
0.97
0.48
0.58
4.67
8
256
0.71
0.75
3.95
1.52
1.70
18.78
F.2 INFERENCE END-TO-END EVALUATIONS
By further fixing the number of generated tokens at 512, we evaluate the end-to-ent inference exe- cution latency of FSA, NSA, and Full Attention both H20 and H200 GPUs.
Results and discussion:
The experimental results Table and demonstrate that FSA's perfor-
mance scales well for extremely long sequences. For inference execution latency: (i) compared to NSA, FSA achieves up to 1.40x speedup and on average 1.16x lower kernel latency on H20 GPUs, and up to x.59 speedup and average 1.15x lower kernel latency H200 GPUs. (ii) Compared to Full Attention, FSA achieves up to 7.20 speedup and on average 4.61 x lower kernel latency on H20 GPUs, and up 4.71 speedup and average 2.96x lower kernel latency H200 GPUs.
Table 8: H20 Inference Latency (s) for (Bx. T) at (64.16) and (128,8).
GQA
GQA =
GQA
GQA
Method
(Bę, T)
128K
256K
128K
256K
128K
256K
128K
256K
(64,16)
1.67
6.71
1.12
4.05
0.77
2.72
0.70
1.94
FSA
(128,8)
1.49
5.65
1.15
3.49
0.70
2.41
0.67
1.87
(64,16)
2.33
7.49
1.42
4.38
0.86
2.75
0.70
1.95
NSA
(128,8)
1.93
7.81
1.54
4.55
0.74
2.84
0.68
1.99
HA
4.34
13.45
4.32
13.48
4.35
13.47
4.31
13.46
Table 9: H200 Inference Latency (s) for (Bx, T) (64.16) and (128,8).
GQA =
GQA 2
GQA =4
GQA = 8
Method
(Bx, 7)
128K
256K
128K
256K
128K
256K
128K
256K
(64,16)
1.06
4.24
0.74
2.46
0.55
1.52
0.46
1.05
FSA
(128,8)
0.93
3.48
0.66
2.08
0.49
1.38
0.45
1.03
(64,16)
1.29
4.29
0.85
2.49
0.61
1.61
0.46
1.05
NSA
(128,8)
1.48
4.42
0.94
2.53
0.56
1.56
0.48
1.07
1.88
4.86
1.85
4.88
1.86
4.89
1.84
4.87
G COMPARISON WITH FLASHDECODING
To compare with the ofate-of-the-ar FlashDecoding kernel Dao et al. 2023), we conducted ad- ditional experiments measuring decoding execution latency for FlashDecoding and FSA. Given
19

---

Published as a conference paper at ICLR 2026
prefill sequence length (ranging from 32K 256K), we present the average decoding latency across 1K generated tokens. We fixed the number of attention heads at 64 and the head dimension at 128. For FSA, the sparse-attention hyperparameters were set to block size Bx of 64 and TopK Value T of 16.
Experimental results in Table TO demonstrate that FSA achieves superior per- formance to FlashDecoding. Compared to FlashDecoding, FSA demonstrates an average speedup of 5.46x on H20 GPU and 2.16x, on H200 GPU. During the decoding phase, FlashDecoding partitions the key and value tokens and distributes the resulting attention computation tasks across multi- ple thread blocks, thereby increasing kernel-level parallelism and improving decoding throughput. However, due to the sparsity in FSA, the FSA decoding throughput is still superior to FlashDecod- ing-
Result discussions:
Table 10: Decoding Latency on H20 and H200 GPUs.
H20 Latency (ms)
H200 Latency (ms)
Seq Len (K)
FlashDecoding
NSA
FSA
Fecoding
NSA
FSA
32
0.88
0.46
0.45
0.51
0.48
0.47
64
1,71
0.46
0.48
0.88
0.53
0.54
128
3.32
0.50
0.48
1.70
0.55
0.54
256
5.76
0.62
0.61
1.75
0.62
0.63
H COMPLIATION OVERHEAD
To determine the optimal Triton kernel hyperparameters, both FSA and NSA incur compilation overhead. For given NSA hyperparameter combination, this overhead occurs only once. Setting (BК. 7) at (64, 16) or (128, 8), we evaluate the compilation overhead of FSA and NSA for sequence length across diverse sequence lengths. The experimental results are summarized in Table 11
Table 11: Compilation Overhead on H20 and H200 GPUs.
Seglen (K)
Framework
H20 Overhead (s)
H200 Overhead (s)
32
FSA
2.16
1.82
32
NSA
2.12
1.78
64
FSA
2.37
2.01
64
NSA
2.33
1.95
128
FSA
2.59
2.24
128
NSA
2.55
2.19
256
FSA
2.80
2.36
256
NSA
2.76
2.30
EVALUATIONS ON DISTRIBUTED PERFORMANCE.
1.1 DISTRIBUTED INFERENCE EVALUATION OF THE ATTENTION MODULE
Table 12: Distributed inference latency of the attention module on H20 GPU.
(Bx. T)
Seq Len (K)
Framework
TP=1 (ms)
TP=2 (ms)
TP=4 (ms)
TP=8 (ms)
(64, 16)
32
FSA
82.50 45.00
25,94
16.25
(64, 16)
32
NSA
99.53
53.44
28.75
16.56
(64, 16)
64
FSA
195.84
110.63
61.25
38.63
(64, 16)
64
NSA
221.49
122.81
65.31
39.94
(128, 8)
FSA
80.31
43.44
24.69
15.63
(128, 8)
32
NSA
105.10
54.38
28.68
16.75
(128, 8)
FSA
187.50
102.68
56.25
33.75
(128, 8)
64
NSA
243.88
130.31
70.69
40.00
20

---

Published as conference paper at ICLR 2026
Table 13: Distributed inference latency of the attention module H200 GPU.
(Bx. 7)
Seq Len (K)
Framework
TP=1 (ms)
TP=2 (ms)
TP=4 (ms)
TP=8 (ms)
(64, 16)
32
FSA
43.44
25.00
15.63
11.88
(64, 16)
32
NSA
50.31
27.19
17.63
12.69
(64, 16)
64
FSA
110.00
63.13
39.38
26.13
(64, 16)
64
NSA
121.56
66.81
40.75
27.00
(128, 8)
32
FSA
40.63
23.13
14.38
10.00
(128, 8)
32
NSA
59.06
31.25
17.50
11.63
(128, 8)
64
FSA
96.25
53.75
32.50
22.81
(128, 8)
64
NSA
124.38
65.69
37.50
25.56
We conduct additional experiments evaluate the distributed inference performance of the attention module using FSA and NSA on H20 and H200 GPUs. We fix the number of query heads at 32, and the number of key and value heads 8. This setting indicates that one GQA group contains query heads. The results for both methods measured across different NSA hyperparameters, sequence lengths, and tensor-parallel degrees are summarized in the Table T2 and [3 Compared to NSA, FSA achieves an average speedup of 1.16x on H20 GPUs and 1. 17x on H200 GPUS.
1.2 END-TO-END DISTRIBUTED INFERENCE EVALUATION
Following the same configuration as Figure 6 we evaluate the distributed inference performance of the Llama-3-8B model on H20 and H200 GPUs. The results for NSA and FSA, measured under varying NSA hyperparameters, sequence lengths, and consor-paralled degrees, are presented in the Table and [5| Compared NSA, FSA achieves average speedup of 1.13x on H20 GPUs and 1x on H200 GPUs.
Table 14: End-to-end distributed inference latency H20 GPU.
(Bx. 7)
Seqlen (K)
Framework
TP=1 (s)
TP=2 (s)
TP=4 (8)
TP=8 (s)
(64, 16)
32
FSA
5.28
2.88
1.66
1.04
(64, 16)
32
NSA
6.00
3.22
1.84
1.06
(64, 16)
FSA
11.14
7.08
3.92
2.60
(64, 16)
64
NSA
12.04
7.86
4.18
2.66
(128, 8)
32
FSA
5.14
2.78
1.58
1.00
(128, 8)
32
NSA
6.40
3.32
1.80
1.10
(128, 8)
64
FSA
12.00
6.38
3.60
2.16
(128, 8)
64
NSA
13.72
7.10
4.00
2.12
Table End-to-end distributed inference latency on H200 GPU.
(Bx. 7)
Seqlen (K)
Framework
TP=1 (s)
TP=2 (8)
TP=4 (8)
TP=8
(64, 16)
32
FSA
1.95
1.12
0.70
0.46
(64, 16)
32
NSA
1.97
1.22
0.70
0.49
(64, 16)
64
FSA
4.51
2.83
1.48
1.09
(64, 16)
NSA
4.61
2.81
1.51
1.16
(128, 8)
32
FSA
1.82
1.04
0.64
0.45
(128, 8)
32
NSA
2.37
1.33
0.78
0.53
(128, 8)
64
FSA
3.61
2.13
1.34
0.98
(128, 8)
64
NSA
4.62
2.63
1.61
1.15
21


Reviews:
[{"review_id": "bnTuc71GGo", "reviewer": "Reviewer_AuaY", "summary": "The paper proposes Flash Sparse Attention (FSA)—a new kernel implementation for Native Sparse Attention (NSA) that flips the loop order: instead of iterating queries (outer) and KV blocks (inner) as in the vanilla NSA kernel, FSA iterates KV blocks in the outer loop and batches of query tokens in the inner loop. This change targets a practical inefficiency of NSA on modern GPUs when each GQA group has few query heads, which otherwise forces padding and underutilizes tensor cores. FSA adds (i) an indexed, non-contiguous query loader, (ii) a dedicated reduction kernel, and (iii) a separate online-softmax statistics kernel. Experiments on H20/H200 GPUs report up to 3.5× kernel speedups vs. NSA and 1.25× training / 1.36× prefill speedups on Llama3-8B, Qwen3-14B, and Qwen2.5-32B for long contexts.", "questions": "typo: L63: \"implementation fail\" ->  \"implementations fail\"", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well motivated. It identifies padding-driven inefficiency in NSA for common small-g GQA regimes and addresses it directly via loop reordering.\n2. The paper presents a solid kernel design with detailed techniques such as non-contiguous query batching with early termination; decoupled reduction to avoid atomics; precomputed online softmax stats to maintain numerical correctness.\n3. The paper provides comprehensive evaluation with microbenchmarks across GPUs and (BK,T) settings, plus end-to-end training and inference on multiple LLMs and sequence lengths (8k–64k).", "weaknesses": "1. FSA only provides efficiency gains over NSA when each GQA group has few query heads, which limits its impact.", "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": "3eSvierBGm", "reviewer": "Reviewer_8FpK", "summary": "This paper proposes Flash Sparse Attention to enable efficient NSA computation with a varied number of query heads in each GQA group to improve LLM inference.", "questions": "Please see the weaknesses.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. This paper aims to address an interesting and important problem in long-context LLM applications.\n2. The presentation is good with clear writing.\n3. The evaluation results are comprehensive and good.", "weaknesses": "1. What precision is used for evaluation? Is it FP8, FP16, or FP32? \n2. I am wondering how would the proposed kernel could scale beyond 64K length, especially for inference.\n3. I am wondering if the authors could provide any further insights into optimizing the proposed kernel for different GPU architectures, such as Hopper and Blackwell.\n4. Some system-related works on sparse attention are missing [1-3].\n\n[1] InfiniGen: Efficient Generative Inference of Large Language Models with Dynamic KV Cache Management, OSDI 2024.\n\n[2] Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference, MLSys 2024.\n\n[3] ALISA: Accelerating Large Language Model Inference via Sparsity-Aware KV Caching, ISCA 2024.", "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": "jfcZTUKPvV", "reviewer": "Reviewer_W7t5", "summary": "The core issue addressed in this paper: NSA's cyclic ordering is only efficient when the number of query heads in the GQA group is large. However,  current attention mechanisms often employ GQA with a small number of attention heads, requiring NSA to use padding to meet hardware instruction requirements. This leads to redundant memory access and computation, preventing the theoretical reduction in FLOPs from translating into actual speed improvements.\nThe core solutions are loop order reversal, non-contiguous memory access optimization, staged computation and Softmax optimization, and memory-hardware co-adaptation.\nIts greatest value lies in providing a hardware-friendly sparse attention solution for long-context LLM training and inference, which can be directly deployed within the existing GPU ecosystem, achieving an average 1.11x speedup in inference compared to NSA.", "questions": "1. Can FSA adapt to other sparse patterns (e.g., block-sparse, dynamic)? Is the loop inversion strategy compatible beyond NSA's structure?\n2.  Any evaluation beyond 64K length? How does attention sink affect the dual-buffer design?\n3. Beyond loss curves, can you provide multi-model metrics (e.g., perplexity, QA F1) across diverse tasks?\n4. What's the actual memory footprint of intermediate buffers? Any compilation time or distributed inference benchmarks?\n5. What about other related work,like flashdecoding?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Problem Significance: Addresses the key bottleneck in deploying sparse attention: the incompatibility between native NSA kernels and mainstream LLMs' small GQA group sizes (typically 1-4 query heads).\n2. Core Innovation: Inverts the NSA kernel's loop order (to \"outer loop over KV blocks, inner loop over query tokens\"), eliminating padding requirements for small GQA groups and significantly reducing redundant computation and memory access. Enhanced by memory management and specialized kernels.\n3. Evaluation Results: Comprehensive testing shows substantial improvements: up to 3.5× lower kernel latency, 1.25× faster end-to-end training, and 1.36× faster inference prefill, with consistent performance across configurations.\n4. Technical Depth: Theoretically reduces memory access and FLOPs to 21.3% and 56.2% of NSA respectively (GQA=4), complemented by practical engineering solutions for real-world deployment.", "weaknesses": "1. Limited Generalization to Extreme Lengths: Evaluation only up to 64K tokens, lacking analysis of ultra-long contexts (128K/256K). Attention sink effects on FSA's dual-buffer design remain unexamined at extreme scales.\n2. Missing SOTA Comparisons: Only compares FSA with vanilla NSA and full attention. Omits comparisons with recent sparse kernels like flashdecoding, limiting perspective on performance trade-offs.\n3. Insufficient Accuracy Analysis: Accuracy claims rely solely on Llama3-8B results. Lacks validation across diverse tasks and smaller models, with no study of numerical stability impact from optimizations.\n4. Unquantified Deployment Overhead: Intermediate buffer costs are mentioned but not measured for memory-constrained scenarios. No discussion of integration complexity with mainstream frameworks.", "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)