# Flash-Sparse-Attention: trainable sparse attention kernels for very long contexts

> Flash-Sparse-Attention is a Python library that implements sparse attention as Triton and CuTe kernels with working backward passes, split-KV and split-QO balancing, and a decode path. It suits teams training or serving Transformer models on extremely long sequences, and it carries a few sharp edges: the fastest backend needs extra packages, arbitrary mask support was moved off main, and the stated Python requirement does not match the packaging metadata.

**HKUSTDial/flash-sparse-attention** — Trainable fast and memory-efficient sparse attention

- Repository: https://github.com/HKUSTDial/flash-sparse-attention
- Website: https://hkustdial.github.io/flash-sparse-attention/
- Stars: 761 · Forks: 53
- Language: Python
- License: BSD-3-Clause
- Published: 2026-10-21 · Updated: 2026-10-21 · Language: en
- Canonical page: https://hysenlabs.com/projects/hkustdial-flash-sparse-attention

## Why trainable sparse attention needs its own kernel

Flash-Sparse-Attention targets one specific pain point: Transformer models whose sequences are long enough that dense attention becomes the bottleneck. The project describes itself as combining the memory efficiency of Flash Attention with sparse computation, so that extremely long sequences stay trainable rather than merely inferable.

Training is the part most sparse attention libraries skip. Inference only kernels let you serve a long context but give you nothing to optimize a model with. This library ships forward and backward passes for dense attention, sparse attention and gated attention, alongside a separate decode entry point for serving. Regular batched inputs and variable length inputs are both accepted.

The implementation is Python on top of Triton, with an optional CuTe backend, and it targets accelerators rather than CPUs. Supported devices are listed as GPU, XPU, NPU and PPU. Head dimensions are supported up to 256, and both Grouped Query Attention and Multi Query Attention are covered. If your workload is short context, most of this machinery buys you nothing and adds a compiler dependency to your stack.

## Split-KV and Split-QO reshape the uneven long-context workload

The structural problem with long attention is that the work is not evenly distributed. Causal masking means early query positions attend to a single token while late ones attend to thousands, so a kernel that gives every query the same amount of parallelism finishes when the slowest block finishes. Most of the remaining compute goes idle.

The library answers this with two separate partitioning switches. Split-KV handles workload balancing in forward and decode workloads, dividing the key and value sequence across worker blocks and combining the partial softmax results. Split-QO handles workload balancing in backward workloads, splitting the query and output side instead, because the backward pass has its own imbalance pattern.

The distinction matters when you tune. The forward and decode examples in the documentation turn on is_split_kv, while the backward example turns on both is_split_kv and is_split_qo, so the two passes are tuned separately rather than sharing one configuration. That granularity is the practical reason this library exposes more flags than a drop-in dense attention call would.

## Flex Local Window and Fused Quant decide what gets computed at all

Two feature families control how much work actually happens. The first is locality. Flex Local Window Attention supports per-head arbitrary window sizes and arbitrary local ranges, which is stronger than a single global window size applied uniformly across heads. Causal attention and local window attention are both listed as supported modes.

The second is the sparse softmax threshold. A threshold controls which low-contribution entries are skipped, which is how sparsity is introduced without training a separate mask predictor. Related to it is the gating mechanism: gated attention accepts gate inputs and exposes configurable gating sparsity, so the amount of retained attention can be data dependent.

Precision control comes from Fused Quant, described as low-precision computation on hardware without native FP8 support, which the README presents as the path for low-precision math on devices that lack FP8 instructions. For serving, the library also offers top-k gather KV cache decode and Paged Attention. This is a broad feature set, but every item in it costs a flag and a tuning decision, so expect to read the API documentation rather than guess parameter names.

## Where FlashAttention is the better default

The project positions FlashAttention as the baseline in its benchmark section, and the difference in approach is conceptual rather than incremental. FlashAttention computes every query and key pair while keeping memory use low by tiling and recomputing in the backward pass. It does not skip pairs, because the whole point is exact attention.

Flash-Sparse-Attention adds a mechanism for not computing pairs at all: a threshold that drops low-contribution entries and a windowing scheme that restricts the range each head attends to. The trade-off is that results are no longer identical to dense attention. If your model depends on exact global attention over a moderate context, or if you need reproducible agreement with a reference dense implementation, the sparse path introduces differences that you would have to characterize yourself.

Dense attention remains inside the same library, which is convenient for a controlled comparison in one codebase. The sensible way to evaluate is to run both paths on the same tensors with the same shapes and measure the gap in output and in latency on your own sequence lengths and head dimensions.

## The CuTe backend leads and the Gluon backend is still on the list

The README states plainly which backend to prefer: the CuTe backend in flash_sparse_attn.ops.cute currently delivers the best performance, while a Gluon backend targeting performance parity with CuTe is still work in progress. Gluon therefore should not appear in your deployment plan yet. A TLE backend is also listed among the features the project aims to support, alongside a KV-Cache Manager.

The CuTe path is not the default path. Its dependencies sit in an optional extra in the packaging metadata, and the list is long: nvidia-cutlass-dsl, einops, typing_extensions, apache-tvm-ffi, torch-c-dlpack-ext and quack-kernels. The base install carries only triton as a runtime dependency. Anyone who plans to use the faster backend should budget for that extra dependency set rather than discovering it at deployment time.

There is also a versioning gap worth noting. The published versions run v2.0.6 on 2026-09-09, v2.0.5 on 2026-06-30 and v2.0.4 on 2026-06-21, so the release line is moving at a slow but steady cadence. Pin the version you validated instead of tracking the branch head.

## Installing on Ubuntu 22.04 with Triton 3.6

The stated requirements are Linux, Ubuntu 22.04 or later, a supported accelerator device, Python 3.9 or later, PyTorch 2.5.1 or later, Triton 3.6.0 or later and Triton Kernels 3.6.0 or later. Install the released package from PyPI like this:

```bash
pip install flash-sparse-attn
```

To work from a checkout instead:

```bash
git clone https://github.com/HKUSTDial/flash-sparse-attention.git
cd flash-sparse-attention
pip install .
```

The packaging metadata tells you a slightly different story than the prose above. The project entry in pyproject.toml sets requires-python to ">=3.10" and lists triton>=3.6.0 as its only dependency, with setuptools-scm deriving the version from the flash_sparse_attn package attribute. So the effective Python floor at install time is 3.10, even though the requirements section says 3.9 or later. Treat 3.10 as the real requirement and check it before your image is built.

The repository also ships a Makefile with test, style, quality and packaging targets, a docs directory built with mkdocs and mkdocstrings, and an examples/modeling directory for integration work.

## A first forward, backward, and CUDA graph decode

The public entry points live in the Triton interface module. The import and the tensor shapes used throughout the examples look like this:

```python
import torch
from flash_sparse_attn.ops.triton.interface import (
    flash_sparse_attn_func,
    flash_sparse_attn_with_kvcache_func,
)
dtype = torch.bfloat16
device = torch.device("cuda")
batch_size, seqlen, num_heads, num_kv_heads, head_dim = 2, 4096, 32, 8, 128
```

A forward pass that turns on causal masking, the sparse softmax threshold, locality, quantization and split-KV is written as:

```python
output = flash_sparse_attn_func(
    query, key, value,
    is_causal=True,
    softmax_threshold=1.0,
    is_local=True,
    is_quant=True,
    is_split_kv=True,
)
```

The backward example adds split-QO and calls backward on the result:

```python
output = flash_sparse_attn_func(
    query, key, value,
    is_causal=True,
    softmax_threshold=1.0,
    is_local=True,
    is_quant=True,
    is_split_kv=True,
    is_split_qo=True,
)

output.sum().backward()
```

For serving, the decode path takes a shorter query and reuses the key and value tensors, and the documentation shows it being captured in a CUDA graph after a warmup loop:

```python
def fsa_decode_fn():
    return flash_sparse_attn_with_kvcache_func(
        query, key, value,
        softmax_threshold=1.0,
        is_local=True,
        is_quant=True,
    )
```

The examples use bfloat16 and a cuda device. Both the threshold value of 1.0 and the head dimension of 128 are choices from the documentation, not recommended settings, so sweep them for your own model.

## Main branch dropped arbitrary masks, and the licence is BSD-3-Clause

One feature moved rather than stayed. Support for arbitrary mask and bias shapes is available in a separate branch, and the current main branch no longer maintains that feature set. If your model needs custom masks, you are looking at a fork or a non-default branch, which changes the upgrade story completely.

On upkeep, the repository is not archived and the last push was on 2026-10-22. The three most recent releases are v2.0.6 on 2026-09-09, v2.0.5 on 2026-06-30 and v2.0.4 on 2026-06-21. Project files include a CITATION.cff, SECURITY.md, CONTRIBUTING.md and a code of conduct.

The licence is BSD-3-Clause, and the packaging metadata classifies it as an OSI approved BSD licence. The practical implication is permissive use with redistribution obligations rather than copyleft, which matters for a dependency you ship inside a commercial model stack. Check the LICENSE file and the dependencies you pull in through the CuTe extra before you commit to a distribution model.

## Conclusion

Adopt flash-sparse-attention if you are training or serving a Transformer on sequence lengths where dense attention stops fitting, and if you have the accelerator and the patience for a Triton-level dependency stack. Do not adopt it for short contexts, for CPU work, or if you need arbitrary mask and bias shapes without cherry-picking a side branch. Before committing, verify three things on your own hardware: that the CuTe backend installs cleanly with its optional dependencies, that the Python version you plan to use satisfies the packaging metadata rather than the looser prose requirement, and that the branch of main you rely on still exposes the mask semantics your model needs.

## FAQ

### What are the hardware and software requirements for flash-sparse-attention?

The stated requirements are Linux with Ubuntu 22.04 or later, a GPU, XPU, NPU or PPU device, Python 3.9 or later, PyTorch 2.5.1 or later, and Triton and Triton Kernels 3.6.0 or later. The packaging metadata is stricter on Python, setting requires-python to ">=3.10".

### How do I install flash-sparse-attention and use the CuTe backend?

Run pip install flash-sparse-attn for the released package, or clone the repository and run pip install . from the checkout. The CuTe backend needs extra packages that sit in an optional extra, including nvidia-cutlass-dsl, einops, typing_extensions, apache-tvm-ffi, torch-c-dlpack-ext and quack-kernels.

### Can flash-sparse-attention be used for training, not just inference?

Yes. The library ships forward and backward passes for dense attention, sparse attention and gated attention, and the backward example enables split-QO alongside split-KV before calling backward on the output.

### Does flash-sparse-attention support arbitrary mask and bias shapes?

Not on the default branch. Support for arbitrary mask and bias shapes lives in a separate branch, and the current main branch no longer maintains that feature set.

### Which backend should I use, CuTe or Triton?

The README states that the CuTe backend currently delivers the best performance, while the Gluon backend targeting parity with CuTe is still work in progress and should not be relied on yet.

## Sources

- [HKUSTDial/flash-sparse-attention on GitHub](https://github.com/HKUSTDial/flash-sparse-attention)
- [License: BSD-3-Clause](https://github.com/HKUSTDial/flash-sparse-attention/blob/main/LICENSE)
- [Project website](https://hkustdial.github.io/flash-sparse-attention/)
- [README](https://github.com/HKUSTDial/flash-sparse-attention/blob/main/README.md)
- [Releases](https://github.com/HKUSTDial/flash-sparse-attention/releases)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/hkustdial-flash-sparse-attention