flagos-ai/FlagAttention

A collection of memory efficient attention operators implemented in the Triton language.

FlagAttention – 快速且内存效率高的 Attention 内核

简介 – FlagAttention 是一个 Python 库，提供一系列使用 Triton（一种用于编写高性能 GPU 内核的语言）编写的自定义 Attention 算子。这些内核实现了与流行的 FlashAttention 算法相同的功能，并增加了许多专门的变体，例如稀疏、分页缓存、分段、量化与递归 Attention。它们旨在用于各种加速器（NVIDIA GPU、AMD、Ascend 等）上进行大型语言模型或其他 Transformer 风格网络的训练与推理。

关键算子系列

系列	典型使用场景	API 入口点	梯度支持	成熟度
FlashAttention	标准密集 Attention (scaled-dot-product)	`flag_attn.flash_attention`	前向 + 反向	稳定
Piecewise Attention	根据 Token 距离切换两个 Q/K 对 (例如 NLPE)	`flag_attn.piecewise_attention`	前向 + 反向	稳定
Split-KV FlashAttention	长 KV 解码，低 Query 并行度	`flag_attn.flash_attention_split_kv`	仅前向	Alpha
Paged Attention	通过分页 KV 缓存进行单 Token 查询 (常见于仅解码器模型)	`flag_attn.paged_attention`	仅前向	Alpha
MiniMax M3 Sparse Attention	具有 128-token 分页的块稀疏预填充/解码	`flag_attn.minimax_m3_*`	仅推理	Alpha
Chunk GLA / GDN2 / KDA	用于长序列的递归门控线性或 Delta 规则 Attention	`flag_attn.chunk_*`	前向 (+ GLA 的反向)	Alpha
SageAttention	每个块 INT8 量化的 Q/K，可选遮罩	`flag_attn.sage_attention.forward`	仅前向	Alpha
Parallel NSA	原生稀疏 Attention + 压缩	`flag_attn.parallel_nsa.parallel_nsa`	前向 + 反向	Alpha

安装 (快速指南)

# Clone the repo
git clone https://github.com/flagos-ai/FlagAttention.git
cd FlagAttention

# Install the editable package with Triton and test extras
pip install -e "[triton,test]"

您必须已为您的加速器安装兼容的 PyTorch 构建版本 (例如 torch==2.x+cu118)。 本套件刻意不依赖特定的 PyTorch wheel，因为正确的构建版本会因 GPU 供应商而异。

最小使用示例 – FlashAttention

import torch
from flag_attn import flash_attention

B, Hq, Hkv, M, N, D = 2, 16, 4, 2048, 4096, 128
q = torch.randn(B, Hq, M, D, device='cuda', dtype=torch.float16, requires_grad=True)
k = torch.randn(B, Hkv, N, D, device='cuda', dtype=torch.float16, requires_grad=True)
v = torch.randn(B, Hkv, N, D, device='cuda', dtype=torch.float16, requires_grad=True)

out = flash_attention(q, k, v, causal=True)
out.sum().backward()   # gradient flow works

函数签名反映了原始的 FlashAttention API，并增加了用于 Dropout、返回 Log-normalizer、总 Attention 质量与 Philox RNG 状态的可选标志。

何时使用哪种算子

场景	推荐算子
单一 GPU 上的标准密集训练或推理	FlashAttention (稳定)
需要在距离阈值后切换 Attention 矩阵 (例如位置嵌入技巧)	Piecewise Attention (稳定)
Query 较少的超长 KV 缓存 (解码)	Split-KV FlashAttention (alpha)
维护分页 KV 缓存的仅解码器模型 (例如 LLM 服务)	Paged Attention (alpha)
块稀疏模型，如 vLLM 风格的 Token 稀疏解码	MiniMax M3 (alpha)
用于流式或超长序列的递归线性 Attention	Chunk GLA / GDN2 / KDA (alpha)
希望每个块进行 INT8 量化以节省内存带宽	SageAttention (alpha)
需要带有可选压缩的原生稀疏 Attention	Parallel NSA (alpha)

成熟度与测试

本仓库在 conf/operators.yaml 中追踪算子成熟度。只有 FlashAttention 和 Piecewise Attention 被标记为稳定；所有其他皆为 alpha，API 签名可能会有所变动。
完整的 pytest 套件 (pytest tests) 根据 FP32 PyTorch 参考验证数值准确性。准确度容差约为 FP32 误差的两倍，并带有较小的绝对边际。
基准测试提供于 benchmark/ 下，可通过辅助脚本 tools/run_tests.py 或直接执行 (例如 python benchmark/flash_benchmark.py)。结果以与 FlagGems 平台兼容的 JSON 格式输出。

支持的硬件与后端

主要：NVIDIA CUDA GPU (通过 Triton)。本库也包含针对 AMD、Ascend、Enflame、MetaX、MThreads、Intel 等厂商的特定后端。后端选择会根据检测到的 PyTorch 设备自动进行，但并非每个算子都针对每个厂商实现 – 请参阅 conf/operators.yaml 与 src/flag_attn/runtime/backend 目录以获取详细信息。
CPU：您可以导入套件并检查元数据，但若没有 GPU 兼容的执行环境，内核将无法执行。

总结

FlagAttention 是一个真实、生产级的库，为研究人员与工程师提供用于各种 Transformer 工作负载的快速、节省内存的 Attention 内核。其稳定的 FlashAttention 与 PiecewiseAttention API 可取代训练循环中的默认 PyTorch Attention，而 Alpha 算子则为进阶 LLM 服务与研究提供了分页缓存、块稀疏解码、量化内核与递归线性 Attention 等实验性功能。