English | 简体中文

Flash-Sparse-Attention is a high-performance trainable sparse attention implementation that combines Flash Attention's memory efficiency with sparse computation for handling extremely long sequences in Transformer models.

Key Features

Note

Support for arbitrary mask and bias shapes is available in this branch. The current main branch no longer maintains that feature set.

Supported Features

• Forward and backward passes for dense attention, sparse attention, and gated attention
• Regular batched inputs and varlen inputs
• Causal attention and local window attention
• Arbitrary combinations of Q and KV sequence lengths, with head dimensions up to 256
• Grouped Query Attention and Multi Query Attention
• Sparse softmax threshold control
• Gated attention with gate inputs and configurable gating sparsity
• Flex Local Window Attention with per-head arbitrary window sizes and local ranges
• Split-KV for workload balancing in forward and decode workloads
• Split-QO for workload balancing in backward workloads
• Fused Quant for low-precision computation on hardware without native FP8 support
• Top-k gather KV-cache decode
• Paged Attention

For complete API documentation, please refer to here

Features We Aim to Support

• KV-Cache Manager
• TLE backend support
• Gluon backend support

Installation

Requirements

• Linux: Ubuntu 22.04 or later
• Device: GPU, XPU, NPU, or PPU
• Python: 3.9 or later
• PyTorch: 2.5.1 or later
• Triton: 3.6.0 or later
• Triton Kernels: 3.6.0 or later

Install

Install from PyPI:

pip install flash-sparse-attn

Additionally, install triton_kernels:

pip install "triton_kernels @ git+https://github.com/triton-lang/triton.git@v3.6.0#subdirectory=python/triton_kernels"

To install from source (includes all dependencies automatically):

git clone https://github.com/flash-algo/flash-sparse-attn.git
cd flash-sparse-attn
pip install .

Install via HuggingFace Kernel

You can also load the kernels directly from HuggingFace Kernel without installing the package:

from kernels import get_kernel

fsa = get_kernel("JingzeShi/flash-sparse-attn", version=1, trust_remote_code=True)

# Forward
out = fsa.flash_sparse_attn_func(q, k, v, is_causal=True)
# Backward
out.sum().backward()
# Decode
out = fsa.flash_sparse_attn_with_kvcache_func(q, k_cache, v_cache)

Quick Start

Basic Usage

Below are examples for forward, backward, and decode.

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

Forward

Combine flex window, split-KV, fused quant, and sparse softmax for maximum performance.

query = torch.randn(batch_size, seqlen, num_heads, head_dim, dtype=dtype, device=device)
key = torch.randn(batch_size, seqlen, num_kv_heads, head_dim, dtype=dtype, device=device)
value = torch.randn(batch_size, seqlen, num_kv_heads, head_dim, dtype=dtype, device=device)

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,
)

Backward

Combine flex window, split-QO, fused quant, and low-contribution skipping for maximum backward performance.

query = torch.randn(batch_size, seqlen, num_heads, head_dim, dtype=dtype, device=device, requires_grad=True)
key = torch.randn(batch_size, seqlen, num_kv_heads, head_dim, dtype=dtype, device=device, requires_grad=True)
value = torch.randn(batch_size, seqlen, num_kv_heads, head_dim, dtype=dtype, device=device, requires_grad=True)

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()

Decode

Combine flex window, split-KV, fused quant, sparse softmax, packed GQA, and Graph for maximum decode performance.

query = torch.randn(batch_size, num_heads, head_dim, dtype=dtype, device=device)
key = torch.randn(batch_size, seqlen, num_kv_heads, head_dim, dtype=dtype, device=device)
value = torch.randn(batch_size, seqlen, num_kv_heads, head_dim, dtype=dtype, device=device)

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

# Warmup
for _ in range(3):
    fsa_decode_fn()
torch.cuda.synchronize()

# Capture Graph
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
    output = fsa_decode_fn()

# Replay
graph.replay()

Performance

The following benchmarks cover forward, backward, and decode workloads, using FlashAttention as the baseline.

NVIDIA GPU

A100

Forward Performance

Backward Performance

Decode Performance

H20

Forward Performance

Backward Performance

Decode Performance

RTX PRO 6000

Forward Performance

Backward Performance

Decode Performance

Benchmarking

Benchmark scripts are located under tests, covering forward, backward, and decoding performance.

Forward Performance

python tests/benchmark_forward.py

Backward Performance

python tests/benchmark_backward.py

Decode Performance

python tests/benchmark_decode.py

Citation

If you use FSA in your research, please cite:

@misc{shi2025trainabledynamicmasksparse,
      title={Trainable Dynamic Mask Sparse Attention},
      author={Jingze Shi and Yifan Wu and Bingheng Wu and Yiran Peng and Liangdong Wang and Guang Liu and Yuyu Luo},
      year={2025},
      eprint={2508.02124},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2508.02124},
}

Acknowledgments

This project builds upon and integrates several excellent works:

• OpenSeek - Kernel development support
• Flash-Attention - Memory-efficient attention computation
• NVIDIA CUTLASS - High-performance matrix operations library

We thank the open-source community for its contributions to efficient Transformer implementations. 🤗