Flash-Sparse-Attention cuts the padding waste in NSA's selected attention

A kernel for Native Sparse Attention, not a new attention method

This repository is the official implementation of Flash Sparse Attention, and the README is careful to say what kind of improvement it is. FSA provides an optimized kernel implementation for the NSA selected attention module, and without modifying the NSA algorithm it gives an efficient Triton-based implementation for GQA group sizes smaller than 8. The NSA algorithm itself is unchanged.

That distinction matters for anyone deciding whether to use it. This is not a new sparse attention scheme with different output. It is a faster way to compute the same thing, and the claim is scoped to a specific regime: grouped query attention with a small group size, which the README says is more popular on state-of-the-art large language models.

The paper is on arXiv at 2508.18224, released in August 2025, and the news section records the open sourcing of the repository in the same month, followed by a beta of one-step decoding in a separate directory called fsa_preview, and an announced online profiling module that transitions between NSA and FSA, still to be released at the time of writing.

The audience is therefore a kernel-level audience. Somebody has a working NSA implementation, they have profiled it, and the selected attention module is where the time goes. This is a project for optimising a known workload, not for exploring a new one.

Why NSA's selected attention kernel wastes work

The method section explains the problem precisely, and it is a problem about data layout rather than about mathematics.

In the NSA selected attention module, the kernel loops over query tokens in the outer loop and over key value blocks in the inner loop. To make the computation efficient, NSA batches query heads that share the same key value head. But when the GQA group size is not sufficiently large, the selected attention kernel has to pad data to satisfy hardware requirements on the matrix dimensions used by matrix multiplication.

The hardware constraints are specific. For NVIDIA GPUs, the warp-level matrix multiply-accumulate instructions require that each dimension of a matrix tile executed on a warp be larger than a specified value, at least 8 on Hopper GPUs. In Triton, the same kernel must satisfy that each dimension of a matrix tile executed on a thread block is at least 16.

Put those numbers together and the shape of the problem becomes clear. A GQA group size of 1 means the batch of query heads that could share a key value head is one head, and the tile the hardware wants is at least 8 or 16 wide. The data is padded to fill the tile, and the padded entries are computed and thrown away.

So the waste is proportional to how far the real group size falls below the tile size the hardware requires, and it is why the NSA kernel's efficiency depends so strongly on the group size. That is the problem FSA attacks.

FSA reverses the loop order and splits the work in three

FSA exchanges the kernel loop order of the original NSA kernel design. Instead of query tokens outside and key value blocks inside, FSA loops over key value blocks in the outer loop and query tokens in the inner loop.

That single change breaks the work into three major kernels, and the README describes each one.

The main kernel batches the query tokens that attend to the same key value block and stores the partial results to a buffer. The reduction kernel accumulates the attention results for each query token. The online softmax kernel handles the online softmax statistics computation.

The reason this helps is stated as the key insight behind the arrangement: it reduces unnecessary memory access and computations for the padded data, while avoiding atomic additions for accumulating attention results for each query token across key value blocks. Both halves matter. The padding waste falls because the batch of query tokens is now along the dimension that the tile requirement applies to, and the atomics disappear because the accumulation happens in a separate reduction pass over a buffer rather than through concurrent updates to shared state.

The speedup claim is that it originates from significantly lowered kernel-level memory access volume and computations, and the latency comparison the README shows is normalised so that the FSA execution latency is 1. The conditions are NSA hyperparameters with block size B_K of 64 and a topk value T of 16, a 64K sequence length and 4 key value heads, with the comparison run across varied GQA group sizes. The normalised presentation is worth noting when you read the figure, since a ratio against a normalised baseline is a claim about relative shape, not about absolute throughput.

Where FSA wins and where it deliberately steps aside

The features section draws the line clearly. FSA provides an efficient Triton-based implementation for GQA group sizes smaller than 8, and for GQA group sizes larger than or equal to 8 it usually chooses to fall back to the original NSA implementation for better performance.

That fallback is the most useful thing in the README for a user, because it defines the boundary of the tool. If your model has a group size of 8 or more, the original NSA implementation is the one you want, and installing this project changes nothing for you. If your group size is smaller, and small group sizes are what the modern large language models use, this is where the work pays off.

It also means the integration has a fallback path built in, which matters for a kernel library. A user does not have to reason about which path their configuration will take.

What the README does not give is a threshold for hardware, a minimum supported Triton version in practice beyond the requirement list, or a measurement of how the two paths compare on a given model. The latency figure is presented as a chart across group sizes, which tells you the direction of the effect but leaves the reader to judge the magnitude for their own configuration.

The tested envelope is narrow and worth respecting

FSA is currently well tested with a list of conditions that is short enough to read as a specification.

The hardware is NVIDIA Ampere or Hopper GPUs, with the examples given as A100 SXM, H20, H100 PCIe, H100 NVL, H100 SXM and H200 SXM. The datatype is fp16 or bf16. The head dimension must be the same, less than or equal to 256, across query, key and value. GQA group sizes from 1 to 16 are covered. And the tested modes are training and inference, with the inference covered being prefill.

Two things stand out. The first is that decode is not in the list, which is consistent with the separate one-step decoding beta in fsa_preview: the sparse attention path is validated for the prefill-shaped workload, and the decode-shaped one is where the other work sits.

The second is the same head dimension requirement, which is a constraint inherited from NSA rather than introduced by FSA, but it will still stop some configurations from running. The head dimension limit of 256 is generous for current models, and the group size range covers the interesting cases, so in practice the envelope is wide enough for the workloads the project is aimed at.

What is not in the list is AMD hardware, older architectures, and any CPU path, and the implementation is Triton based, which is an NVIDIA-oriented toolchain in practice.

Installing the dependencies and building the module

The README lists the requirements, and the install is a single command against the requirements file:

pip install -r requirements.txt

The listed requirements are PyTorch 2.4 or later, Triton 3.0 or later, transformers 4.45.0 or later, datasets 3.3.0 or later, accelerate 1.9.0 or later, and flash-attn pinned to 2.6.3. The requirements file in the repository is more specific in places, asking for torch 2.8.0 or later and triton 3.4.0 or later, and it installs flash-attention from Git rather than from an index, both the main v2 branch and the hopper subdirectory for v3, along with nvidia_cudnn_frontend, with a comment about bumping flash-attention from 2.6.3 to 2.8.3. So the README list and the pinned file are not the same set, and the file is the one a build will actually use.

Usage starts by instantiating the module class:

import torch
from fsa.module.fsa import FlashSparseAttention, RopeConfig

FSA = (
    FlashSparseAttention(
        hidden_size=4096,
        num_q_heads=4,
        num_kv_heads=4,
        head_dim=128,
        kernel_size=32,
        kernel_stride=16,
        block_size=64,

The constructor arguments are the model geometry, and the sparse attention hyper-parameters sit alongside them: kernel size and kernel stride for the local window, block size for the key value blocking, and further down topk, init blocks, local blocks and window size. The RopeConfig carries the rotary embedding settings, including the scaling factors and the rope type.

The forward call takes the input and a cumulative sequence length tensor, which is how packed variable length sequences are expressed:

seqlens = torch.LongTensor([65536, 32768]).int().cuda()

cu_seqlens = torch.cat(
    [
        torch.zeros(1, dtype=torch.int32, device="cuda"),
        torch.cumsum(seqlens, dim=0),
    ],
    dim=0,
).to(torch.int32)
x = torch.randn(cu_seqlens[-1], 4096, device="cuda", dtype=torch.bfloat16)

y = FSA(x, cu_seqlens)
loss = (y * torch.randn_like(y)).sum(-1).mean()
loss.backward()

The example is a training step, which is the honest framing for a kernel library: it shows the forward pass and a backward pass, so the module is meant to be dropped into a training loop rather than called as a service.

What the repository holds, and the absence of releases

The top-level layout matches the three things the project is about. The fsa directory holds the implementation, including the module class and the operations, where the README links a topk sparse attention file. The nsa_ref directory holds the reference NSA implementation that FSA is compared against and falls back to. The fsa_preview directory holds the beta one-step decoding code.

Around those are the supporting directories, benchmarks, test, scripts and utils.py, plus a requirements file and a development requirements file.

There are no GitHub releases, and the project is not archived, with the last push on 2026-09-17. That is consistent with a research codebase whose reference point is the arXiv paper rather than a published version. Anyone pinning a dependency is pinning a commit, and should record which one.

The licence is Apache-2.0 with a LICENSE file at the root, and the homepage given for the project is the repository itself on GitHub.

One more detail is worth noticing in the news section. The online profiling module that transitions between NSA and FSA is announced as upcoming, and the paper is the artefact to read for the method. The repository is where the implementation is, and the arXiv paper at 2508.18224 is where the argument and the latency comparison live.

Editorial conclusion

Flash-Sparse-Attention fits a team training or serving a model that uses Native Sparse Attention with a grouped query attention group size below 8, on NVIDIA hardware, and who have measured that the selected attention kernel is the bottleneck. It is the wrong choice for group sizes of 8 or more, where the implementation deliberately falls back to the original NSA kernel, and the wrong choice for anything but NVIDIA GPUs, since the kernel is Triton based and the tested hardware list is Ampere and Hopper. Before adopting it, measure the selected attention kernel against the rest of the model, confirm the group size of your configuration falls in the range where the fallback does not trigger, and read the requirements file rather than the README list, because the two disagree on the torch and Triton versions.

Frequently asked questions

Official sources

1. Issues
2. License: Apache-2.0
3. Project website
4. README
5. Relaxed-System-Lab/Flash-Sparse-Attention on GitHub