Relaxed-System-Lab/Flash-Sparse-Attention

Flash‑Sparse‑Attention (FSA)

What it is – An open‑source, Triton‑based implementation of Native Sparse Attention (NSA) that rearranges the kernel loops to dramatically cut memory traffic and compute overhead. It ships as a drop‑in replacement for the attention layer in large language models (LLMs) and works on modern NVIDIA GPUs (Ampere, Hopper, etc.).

Why it matters – Sparse attention is a way to keep the quadratic cost of full‑attention manageable for very long sequences (tens of thousands of tokens). The original NSA kernel suffers from padding and atomic‑add overhead when the group‑query‑attention (GQA) head‑group size is small (the common case in today’s LLMs). FSA swaps the outer/inner loops, splits the work into three specialized kernels, and avoids those bottlenecks, delivering up to 2‑3× speed‑up on both training (prefill) and inference while using the same API.

Key Features

Feature	Detail
Optimized Triton kernels	Three kernels – main (batched query‑to‑KV), reduction, and online‑softmax – minimise padded‑data work and remove atomics.
GQA‑aware fallback	For GQA group sizes ≥ 8 it automatically falls back to the original NSA implementation, which is faster in that regime.
Broad hardware support	Tested on NVIDIA A100, H100, H200 (both SXM and PCIe) with fp16 and bf16.
Drop‑in API	`FlashSparseAttention` mirrors the standard `torch.nn.Module` interface; you only need to supply `cu_seqlens` and the input tensor.
Training & inference	Works for both forward‑only prefill and full back‑propagation (loss‑backward).
Compatibility	Built on top of PyTorch ≥ 2.4, Triton ≥ 3.0, HuggingFace transformers ≥ 4.45, and the official Flash‑Attention 2.6.3 kernel.

Quick Start

# 1️⃣ Install dependencies (Python 3.10+ recommended)
pip install -r requirements.txt   # pulls torch, triton, transformers, datasets, accelerate, flash‑attn

# 2️⃣ Import and instantiate the module
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,
    topk=16,
    init_blocks=1,
    local_blocks=2,
    window_size=512,
    rope_config=RopeConfig(
        max_position_embeddings=131072,
        head_dim=128,
        rope_theta=500000,
        rope_scaling={
            "factor": 8.0,
            "high_freq_factor": 4.0,
            "low_freq_factor": 1.0,
            "original_max_position_embeddings": 8192,
            "rope_type": "llama3",
        },
    ),
).cuda().to(torch.bfloat16)

# 3️⃣ Prepare cu_seqlens (cumulative lengths) and input
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)

x = torch.randn(cu_seqlens[-1], 4096, device="cuda", dtype=torch.bfloat16)

# 4️⃣ Forward + backward (training) example
y = fsa(x, cu_seqlens)
loss = (y * torch.randn_like(y)).sum(-1).mean()
loss.backward()

The only extra step compared with a vanilla transformer is the construction of cu_seqlens, which encodes the variable‑length batch.

Benchmarking

The repo includes two scripts:

scripts/run_unit_test.sh – checks forward/backward correctness and measures raw kernel latency.
scripts/run_unit_test_sel_attn.sh – benchmarks the selected‑attention part (the main bottleneck).

The README shows two performance tables:

Kernel‑level – FSA’s latency is normalised to 1, while NSA and full Flash‑Attention are 1.7–2.4× slower depending on block size/top‑k.
End‑to‑end – For LLMs such as LLaMA‑2‑70B, training step latency drops from ~1.9 s (NSA) to ~1.2 s (FSA) and prefill latency improves similarly.

When to use it

You are training or serving LLMs with sequence lengths ≥ 32 k tokens.
Your model uses GQA with ≤ 8 heads per KV group (the most common configuration).
You have an NVIDIA Ampere/Hopper GPU and can install Triton.
You already rely on Flash‑Attention for dense attention and want a sparse‑attention alternative without rewriting the model code.

Limitations / Future work

Currently only supports NVIDIA GPUs; no CPU or AMD path.
The implementation assumes the same head dimension for Q, K, V (≤ 256).
An “online profiling” module that can dynamically switch between NSA and FSA is announced for a future release (Sept 2025).

Citation

If you use FSA in a paper, cite the accompanying arXiv pre‑print:

@misc{yan2026fsaalternativeefficientimplementation,
  title={{FSA}: An Alternative Efficient Implementation of Native Sparse Attention Kernel},
  author={Ran Yan and Youhe Jiang and Zhuoming Chen and Haohui Mai and Beidi Chen and Binhang Yuan},
  year={2026},
  eprint={2508.18224},
  archivePrefix={arXiv},
  primaryClass={cs.DC},
  url={https://arxiv.org/abs/2508.18224},
}

TL;DR

Flash‑Sparse‑Attention is a high‑performance Triton kernel library that makes native sparse attention practical for modern LLMs with long contexts. It drops into existing PyTorch/transformers code, runs on Ampere/Hopper GPUs, and delivers 2‑3× speed‑ups over the original NSA implementation while keeping the same API and numerical correctness.