---

本仓库提供了 Flash Sparse Attention (FSA) 的官方实现，其中包含一种新颖的内核设计，能够在现代 GPU 上高效地支持多种主流大模型中的原生稀疏注意力（NSA）。

• 新闻
• 方法
• 优势
• 特性
• 安装
• 使用
  • 实例化 FSA 模块
  • 使用 FSA 进行训练
• 评估
  • FSA 模块基准测试
  • FSA 稀疏注意力模块基准测试
• 性能
  • 内核性能
  • 端到端性能
• 引用
• 致谢

新闻

• $\texttt{[2025-09, 即将发布]}$: 🚀 可在 NSA 和 FSA 之间无缝切换的在线剖析模块即将发布。
• $\texttt{[2025-08]}$: 💥 我们的 Arxiv 论文 已正式发表。
• $\texttt{[2025-08]}$: 🎈 一步解码的 Beta 版本已发布，请查看位于 fsa_preview 中的代码。
• $\texttt{[2025-08]}$: 🎉 开源了 Flash-Sparse-Attention，提供了一种针对 NSA 的优化实现，进一步拓展了这种可原生训练的新型稀疏注意力技术的应用范围。

方法

对于 NSA 稀疏注意力模块而言，其主要的系统瓶颈在于：外层循环遍历查询 token，而内层循环则遍历 KV 块。为了优化性能，NSA 会将共享相同键值头的查询头进行批处理，以提升计算效率。然而，当 GQA 分组大小不够大时，NSA 稀疏注意力内核必须对数据进行填充，以满足硬件对矩阵乘法维度的要求。具体来说，在 NVIDIA GPU 上，warp 级矩阵乘累加指令要求在 warp 上执行的矩阵分块的每个维度都必须大于指定值（Hopper 架构的 GPU 至少为 8）。而在 Triton 中，NSA 稀疏注意力内核还必须确保线程块上执行的矩阵分块的每个维度至少为 16。

相比之下，FSA 改变了原始 NSA 内核的循环顺序，即外层循环遍历 KV 块，内层循环则遍历查询 token。为了优化性能，FSA 将计算过程拆分为三个主要内核：(i) 主内核将关注同一 KV 块的查询 token 批量处理，并将部分结果存储到缓冲区中；(ii) 归约内核负责累积每个查询 token 的注意力结果；(iii) 在线 softmax 内核则用于处理在线 softmax 统计信息的计算。这一设计的关键在于有效减少了对填充数据的不必要内存访问和计算，同时避免了跨 KV 块对每个查询 token 的注意力结果进行 atomic 加法操作。

NSA（左）与 FSA 主内核（右）的具体计算流程对比可以可视化如下：

优势

🚀 FSA 的速度提升源于显著降低的内核级内存访问量和计算量。

在不同的 GQA 分组大小下，NSA 的超参数设置为：KV 块大小 $B_K=64$，topk 值 $T=16$，序列长度为 64K，KV 头数为 4。NSA 与我们方法的执行延迟对比结果如下（我们方法的执行延迟被归一化为 1）：

特性

FSA 提供了一种针对 NSA 稀疏注意力模块的优化内核实现。在不改变 NSA 算法的前提下，FSA 基于 Triton 提供了一种高效的实现方案，适用于当前最先进的大型语言模型（LLMs）中更为常见的 GQA 分组大小小于 8 的情况，且运行于现代高性能 NVIDIA GPU 上。而对于 GQA 分组大小大于或等于 8 的情况，FSA 通常会选择回退到原始的 NSA 实现，以获得更好的性能。

目前，FSA 已在以下环境中经过充分测试：

• NVIDIA Ampere 或 Hopper 架构的 GPU（例如 A100 SXM、H20、H100 PCIe、H100 NVL、H100 SXM、H200 SXM）；
• 数据类型为 fp16 和 bf16；
• 查询、键和值的头维度相同（不超过 256）；
• GQA 分组大小从 1 到 16 不等；
• 支持训练和推理（预填充阶段）。

安装

请确保满足以下依赖条件：

• PyTorch ≥ 2.4
• Triton ≥ 3.0
• transformers ≥ 4.45.0
• datasets ≥ 3.3.0
• accelerate ≥ 1.9.0
• flash-attn == 2.6.3

您可以通过以下命令安装 FSA 的依赖项：

pip install -r requirements.txt

使用

实例化 FSA 模块

我们提供了 FlashSparseAttention 类供您使用，用法示例如下：

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

# 随机输入
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()

在底层，调用了 FSATopkSparseAttention 类，该类提供了优化的内核，用于加速 NSA 选择性注意力模块。

使用 FSA 进行训练

使用 FSA 进行训练可以通过简单地替换注意力模块来实现。你可能需要做的唯一事情就是实例化 FSA 模块，并为 FSA 计算 cu_seqlens。我们在 SparseLlamaAttention 中提供了一个示例，说明如何将 FSA 插入到 LLM 中。

评估

FSA 模块基准测试

我们在 scripts/run_unit_test.sh 中提供了详细的命令，以便于对 FSA 模块进行便捷的基准测试。该基准测试包括前向和反向输出的正确性比较、性能比较以及内存使用情况的比较。

FSA 选择性注意力模块基准测试

作为系统主要瓶颈的优化 NSA 选择性注意力模块，可以通过 scripts/run_unit_test_sel_attn.sh 中的命令进行基准测试。

[!Tip] 尝试在提供的脚本中调整不同的 gqa、seqlen、block_size、topk 参数，以在你的机器上进行更全面的基准测试！与基准测试 FSA 注意力模块相比，基准测试 FSA 选择性注意力模块通常能带来更高的加速效果。

性能

内核性能

在各种配置下，基于 Triton 的 FSA、NSA 和全注意力（由 Flash Attention 启用）内核的性能比较。元组 ($64$, $16$) / ($128$, $8$) 分别表示块大小 $BK$ 和 top-k 值 $Topk$。对于 FSA 和 NSA，执行延迟由压缩、选择和滑动注意力组成；而对于全注意力，则是 Flash Attention 内核的执行延迟。

端到端性能

具有 FSA、NSA 或全注意力的最先进 LLM 的端到端训练（右）和预填充（左）延迟。

引用

@article{yan2025flashsparseattentionalternative,
  title={Flash Sparse Attention: 更高效的原生可训练稀疏注意力},
  author={Yan, Ran and Jiang, Youhe and Yuan, Binhang},
  journal={arXiv 预印本 arXiv:2508.18224},
  year={2025}
}

致谢

NSA 论文： 原生稀疏注意力

NSA 参考实现： 原生稀疏注意力 Triton

Flash-Sparse-Attention 快速上手指南

Flash-Sparse-Attention (FSA) 是一种针对原生稀疏注意力（NSA）的高效内核实现，旨在现代 NVIDIA GPU 上加速大语言模型（LLM）的训练与推理，特别优化了小 GQA 组大小场景下的性能。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 硬件：NVIDIA Ampere 或 Hopper 架构 GPU（如 A100, H100, H20 等）。
• 软件依赖：
  • PyTorch >= 2.4
  • Triton >= 3.0
  • transformers >= 4.45.0
  • datasets >= 3.3.0
  • accelerate >= 1.9.0
  • flash-attn == 2.6.3
• 数据类型：支持 fp16 和 bf16。
• 头维度：Query、Key、Value 的头维度需一致且 <= 256。

提示：国内开发者建议使用清华源或阿里源加速 Python 包安装。

安装步骤

1. 克隆仓库（如果尚未克隆）：

   git clone https://github.com/your-repo/Flash-Sparse-Attention.git
   cd Flash-Sparse-Attention
   

2. 安装依赖项：

   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   如果 requirements.txt 未包含所有特定版本约束，也可手动安装核心依赖：

   pip install torch>=2.4 triton>=3.0 transformers>=4.45.0 datasets>=3.3.0 accelerate>=1.9.0 flash-attn==2.6.3 -i https://pypi.tuna.tsinghua.edu.cn/simple
   

基本使用

以下是实例化 FlashSparseAttention 模块并执行一次前向传播与反向传播的最小示例。

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

# 实例化 FSA 模块
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)
)

# 准备输入数据
# 模拟两个序列的长度：65536 和 32768
seqlens = torch.LongTensor([65536, 32768]).int().cuda()

# 计算累积序列长度 (cu_seqlens)
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()

集成到训练流程： 若需在现有 LLM 中使用 FSA，只需将原有的注意力模块替换为 FlashSparseAttention，并确保在数据加载阶段正确计算 cu_seqlens 传入模型即可。详细集成示例可参考仓库中的 test/train.py。