SageAttention

环境要求

• 支持 Ampere、Ada 或 Hopper 架构的 NVIDIA GPU
• 已配置 nvcc 的 CUDA 编译环境（源码编译时需要）
• 可用的 Python 与 pip 环境（README 未指定具体版本）
• 已安装 PyTorch（README 未显式列出，按深度学习环境默认准备）

安装与启动步骤

1. 1pip 安装

   安装包含 SageAttention2++ 的 2.2.0 版本，必须加 --no-build-isolation，否则依赖隔离会导致编译失败。

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

   复制

2. 2源码编译安装

   若需要自行编译，克隆仓库后执行 setup.py install；先导出并行编译变量可加快速度，这些变量是可选的。

   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

   复制

3. 3可选安装 FA3

   只有需要与 FlashAttention3 对比速度时才编译，需切到指定 commit 并进入 hopper 目录安装。

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

   复制

4. 4调用注意力算子

   导入 sageattn 后直接调用替换原注意力，q/k/v 需为 FP16 或 BF16；默认 HND 布局，NHD 布局要显式指定。

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

   复制

关键配置

如何确认成功

调用 sageattn 后能正常返回 attn_output，且其形状与输入的 q 保持一致即算成功。

常见问题

注意事项

• SageAttention V1 为 Triton 实现，速度慢于 SageAttention2、2++ 与 3。
• q、k、v 必须为 FP16/BF16，默认形状为 (batch_size, head_num, seq_len, head_dim)。
• 形状为 (batch_size, seq_len, head_num, head_dim) 时需设置 tensor_layout="NHD"。
• 源码编译依赖 CUDA 工具链，编译参数 EXT_PARALLEL、NVCC_APPEND_FLAGS、MAX_JOBS 均为可选。