SageAttention

本仓库提供了 SageAttention、SageAttention2 和 SageAttention2++ 的官方实现，这些方法在大多数 GPU 上实现了显著加速，且在所有模型中均无精度损失，并支持即插即用（plug-and-play）。

SageAttention: 面向即插即用推理加速的高精度 8 位注意力机制
张金涛、魏佳、黄浩峰、张鹏乐、朱军、陈健飞
论文：https://arxiv.org/abs/2410.02367

SageAttention2: 基于彻底异常值平滑与每线程 INT4 量化的高效注意力机制
张金涛、黄浩峰、张鹏乐、魏佳、朱军、陈健飞
论文：https://arxiv.org/abs/2411.10958

SageAttention3: 面向推理的微缩放 FP4 注意力机制及对 8 位训练的探索
张金涛、魏佳、王浩旭、张鹏乐、徐晓明、黄浩峰、姜凯、陈健飞、朱军
论文：https://arxiv.org/abs/2505.11594

注：SageAttention2++ 在保持相同精度性能的同时实现了更高的速度。

当前特性

• 针对 Ampere、Ada 和 Hopper GPU 的优化内核。
• 对 $QK^\top$ 支持多种粒度的 INT8 量化和平滑处理。
• 对 $PV$ 支持 FP8 量化，并为 FP8/FP16 $PV$ 提供 FP16 累加器。
• 采用两级累加策略（two-level accumulation strategy）提升 FP8 MMA 和 WGMMA 中 $PV$ 的精度。
• 支持 torch.compile（非 cudagraphs 模式）和分布式推理。

项目更新

• [2025-09-27]：🎉 SageAttention3 被 NeurIPS 2025 接收为 Spotlight 论文！
• [2025-09-27]：SageAttention3 的代码已在本仓库的 sageattention3_blackwell 目录中发布。若您能花一点时间填写 Huggingface 上的表单，我们将不胜感激。请注意，由于 SageAttention2 精度更高，我们仍推荐在对精度敏感的应用中使用 SageAttention2。
• [2025-07-01]：SageAttention2++ 的代码已在本仓库中发布。若您能花一点时间填写 Huggingface 上的表单，我们将不胜感激。非常感谢！

• [2025-06-19]：Sparse SageAttention1 API 和 Sparse SageAttention2 API 可以极快地计算任意块稀疏模式（block sparse pattern）下的注意力。
• [2025-05-02]：🎉 SageAttention2 和 SpargeAttn 被 ICML 2025 接收！
• [2025-02-25]：🔥 我们发布了 SpargeAttn，这是一种基于 SageAttention2 的稀疏注意力机制，无需训练即可加速任意模型。
• [2025-02-15]：🔥 编译代码已更新，支持 RTX5090！在 RTX5090 上，SageAttention 达到 560T，比 FlashAttention2 快 2.7 倍！
• [2025-01-28]：🔥⚡ SageAttention 现已支持 Hopper GPU（H100、H800、H20）！其速度与 FlashAttention3-FP8 相当，但精度显著更高！

FlashAttention2	FlashAttention3	FlashAttention3-FP8	SageAttention
25'34''	17'32''	12'14''	12'07''

NVIDIA H20 GPU 上 CogVideoX1.5-5B 的结果

• [2025-01-24]：🎉 SageAttention 被 ICLR 2025 接收！

• [2024-12-20]：🔥 更新了 SageAttention2 论文。

• [2024-12-20]：🔥 发布 SageAttention 2.0.1 Beta！此版本引入了新特性：每线程量化（per-thread quantization），在保持硬件效率的同时提供更细粒度的量化。

• [2024-11-21]：🔥 SageAttention 2.0.0 beta 发布！目前 SageAttention 已在 L20、L40、A100、A800、A6000、RTX3090 和 RTX4090 上测得加速效果。

• [2024-11-12]：现已支持 sageattn_varlen。

• [2024-11-11]：现已支持 q 与 k, v 具有不同的序列长度、输入形状为 (batch_size, head_num, seq_len, head_dim) 或 (batch_size, seq_len, head_num, head_dim)，以及分组查询注意力（group-query attention）。

安装

基础环境

• python>=3.9、torch>=2.3.0、triton>=3.0.0

• CUDA:
  • Blackwell 或 SageAttention2++ 需要 >=12.8
  • Ada 架构上支持 fp8 需要 >=12.4
  • Hopper 架构上支持 fp8 需要 >=12.3
  • Ampere 架构需要 >=12.0

• flash-attn（用于基准测试）

安装包

对于基于 Triton 的 SageAttention V1（比 SageAttention V2/V2++/V3 慢），请参考 SageAttention-1 分支，并通过 pip 安装：

pip install sageattention==1.0.6

若要使用包含 SageAttention2++ 的 SageAttention 2.2.0，可通过 pip 安装：

pip install sageattention==2.2.0 --no-build-isolation

或者 从源码编译：

git clone https://github.com/thu-ml/SageAttention.git
cd SageAttention 
export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32 # 可选
python setup.py install

若要与 FlashAttention3 进行速度对比，请从源码编译 FlashAttention3：

git clone https://github.com/Dao-AILab/flash-attention.git --recursive
git checkout b7d29fb3b79f0b78b1c369a52aaa6628dabfb0d7 # 2.7.2 版本
cd hopper
python setup.py install

使用方法

from sageattention import sageattn
attn_output = sageattn(q, k, v, tensor_layout="HND", is_causal=False)

• q, k, v 为 FP16/BF16 类型，默认 tensor_layout="HND" 时形状为 (batch_size, head_num, seq_len, head_dim)。若形状为 (batch_size, seq_len, head_num, head_dim)，请设置 tensor_layout="NHD"。
• is_causal 控制是否使用因果掩码（causal mask）。

可用 API：

• sageattn：根据 GPU 自动选择最优 kernel，在性能与精度之间取得良好平衡。
• sageattn_qk_int8_pv_fp16_triton：使用 Triton 后端，对 $QK^\top$ 进行 INT8 量化（8 位整数量化），对 $PV$ 使用 FP16（半精度浮点）。
• sageattn_qk_int8_pv_fp16_cuda：使用 CUDA 后端，对 $QK^\top$ 进行 INT8 量化，对 $PV$ 使用 FP16。
• sageattn_qk_int8_pv_fp8_cuda：使用 CUDA 后端，对 $QK^\top$ 进行 INT8 量化，对 $PV$ 使用 FP8（8 位浮点）。（注意：设置 pv_accum_dtype=fp32+fp16 对应 SageAttention2++。）
• sageattn_qk_int8_pv_fp8_cuda_sm90：使用 CUDA 后端，对 $QK^\top$ 进行 INT8 量化，对 $PV$ 使用 FP8，并针对 Hopper 架构 GPU 进行专门优化。
• sageattn_varlen：使用 Triton 后端，对 $QK^\top$ 进行 INT8 量化，对 $PV$ 使用 FP16，支持同一批次内不同序列长度。

为了在自定义设备和模型上获得最佳的速度与精度表现，我们强烈建议参考 此文件 获取详细指导。

注意：
支持 q 与 k,v 之间的不同序列长度以及 分组查询注意力（group-query attention）。

即插即用示例

我们可以轻松替换 scaled_dot_product_attention。
以 CogvideoX 为例：

添加以下代码并运行：

import torch.nn.functional as F

+ from sageattention import sageattn
+ F.scaled_dot_product_attention = sageattn

具体操作如下：

cd example
python cogvideox_infer.py --model cogvideox-2b --compile --attention_type sage

你将得到一个无损视频，保存在 ./example/videos/<model>/<attention_type>/ 目录下，且速度比使用 --attention_type sdpa 更快。更多示例和指南请参见 example/ 目录。

注意： 并非所有模型都适用于 F.scaled_dot_product_attention = sageattn。从技术上讲，你应该通过修改目标模型的 Attention Class 来替换原始注意力机制。对于图像和视频模型，我们建议仅替换 DiT 中的注意力模块（详见 example/modify_mochi.py）。

Kernel 基准测试

我们提供了一个基准测试脚本，用于比较 SageAttention、FlashAttention2 和 FlashAttention3 等不同 kernel 的速度。详情请参阅 benchmark/ 目录。

性能

Kernel 速度

8+8 表示对 $QK^\top$ 使用 INT8 量化，对 $PV$ 使用 FP8 量化。8+16 表示对 $PV$ 使用 FP16 并搭配 FP16 累加器。

注意： TOPS 结果仅指注意力 Kernel 本身，不包括量化和平滑（smoothing）操作。

端到端性能

端到端精度：

端到端加速比：

注：SageAttention2++ 实现了更高的速度。

引用

如果你使用了本代码或认为我们的工作有价值，请引用以下论文：

@inproceedings{zhang2025sageattention,
  title={SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration}, 
  author={Zhang, Jintao and Wei, Jia and Zhang, Pengle and Zhu, Jun and Chen, Jianfei},
  booktitle={International Conference on Learning Representations (ICLR)},
  year={2025}
}
@inproceedings{zhang2024sageattention2,
  title={Sageattention2: Efficient attention with thorough outlier smoothing and per-thread int4 quantization},
  author={Zhang, Jintao and Huang, Haofeng and Zhang, Pengle and Wei, Jia and Zhu, Jun and Chen, Jianfei},
  booktitle={International Conference on Machine Learning (ICML)},
  year={2025}
}
@article{zhang2025sageattention2++,
  title={Sageattention2++: A more efficient implementation of sageattention2},
  author={Zhang, Jintao and Xu, Xiaoming and Wei, Jia and Huang, Haofeng and Zhang, Pengle and Xiang, Chendong and Zhu, Jun and Chen, Jianfei},
  journal={arXiv preprint arXiv:2505.21136},
  year={2025}
}
@article{zhang2025sageattention3,
  title={SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Training},
  author={Zhang, Jintao and Wei, Jia and Zhang, Pengle and Xu, Xiaoming and Huang, Haofeng and Wang, Haoxu and Jiang, Kai and Zhu, Jun and Chen, Jianfei},
  journal={arXiv preprint arXiv:2505.11594},
  year={2025}
}

SageAttention 快速上手指南

环境准备

系统与依赖要求

• Python ≥ 3.9
• PyTorch ≥ 2.3.0
• Triton ≥ 3.0.0
• CUDA 版本要求（根据 GPU 架构）：
  • Blackwell 或使用 SageAttention2++：≥ 12.8
  • Ada 架构（如 RTX 4090）启用 FP8：≥ 12.4
  • Hopper 架构（如 H100）启用 FP8：≥ 12.3
  • Ampere 架构（如 A100、RTX 3090）：≥ 12.0

建议使用国内 PyPI 镜像加速安装，例如：pip install -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

推荐方式：pip 安装（含 SageAttention2++）

pip install sageattention==2.2.0 --no-build-isolation

或从源码编译安装

git clone https://github.com/thu-ml/SageAttention.git
cd SageAttention
export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32  # 可选，加速编译
python setup.py install

注意：若仅需旧版 SageAttention V1（性能较低），请切换至 sageattention-1 分支并安装 sageattention==1.0.6。

基本使用

最简调用示例

from sageattention import sageattn

# q, k, v 为 FP16/BF16 张量，形状 (batch_size, head_num, seq_len, head_dim)
attn_output = sageattn(q, k, v, tensor_layout="HND", is_causal=False)

参数说明

• tensor_layout：
  • "HND"：输入形状为 (B, H, S, D)
  • "NHD"：输入形状为 (B, S, H, D)
• is_causal：是否使用因果掩码（用于生成任务）

即插即用替换（以 CogVideoX 为例）

在模型推理脚本开头添加：

import torch.nn.functional as F
from sageattention import sageattn
F.scaled_dot_product_attention = sageattn

然后运行示例脚本：

cd example
python cogvideox_infer.py --model cogvideox-2b --compile --attention_type sage

注意：并非所有模型都支持直接替换 F.scaled_dot_product_attention。对于图像/视频生成模型（如 DiT 架构），建议手动替换模型中的 Attention 模块（参考 example/modify_mochi.py）。