SpargeAttention

推荐 API

我们强烈推荐使用 spas_sage2_attn_meansim_topk_cuda 和 block_sparse_sage2_attn_cuda 这两个 API。它们即插即用且可定制：

即插即用 API

from spas_sage_attn import spas_sage2_attn_meansim_topk_cuda

attn_output = spas_sage2_attn_meansim_topk_cuda(q, k, v, topk=0.5, is_causal=False)

您可以调整 topk 参数，以在注意力精度（topk 越高精度越高）和稀疏性（topk 越低稀疏性越高）之间取得平衡。

自定义块稀疏掩码

from spas_sage_attn import block_sparse_sage2_attn_cuda

attn_output = block_sparse_sage2_attn_cuda(q, k, v, mask_id=None):

在这个 API 中，我们支持为每个注意力头计算任意块稀疏掩码下的注意力。具体来说，每个头的注意力掩码 mask_id 的形状为 (batch_size, num_heads, ⌈seq_len / 128⌉, ⌈seq_len // 64⌉)，由 0 和 1 组成。目前，块大小为 128×64。

---

SpargeAttn 的官方实现，一种通用的无训练稀疏注意力机制，能够加速语言、图像和视频模型。

项目更新

• 请使用 spas_sage2_attn_meansim_topk_cuda 和 block_sparse_sage2_attn_cuda API。
• [2025-07]：发布 Triton Kernel 示例。
• [2025-06]：基于 SageAttention2++ 的 SpargeAttn 正式发布。
• [2025-05]：新增一种非常简单的用法，无需调优或校准：o = spas_sage2_attn_meansim_topk_cuda(q, k, v)。
• [2025-05]：🎉SpargeAttn 和 SageAttention2 已被 ICML 2025 接受！
• [2025-03] 支持在更多 GPU 上实现高效加速，例如 H100。

安装

基础环境

• python>=3.9 , torch>=2.3.0

• CUDA：
  • >=12.8 用于 Blackwell 架构，>=12.4 用于 Ada 架构上的 fp8 支持，>=12.3 用于 Hopper 架构上的 fp8 支持，>=12.0 用于 Ampere 架构。

安装包

pip install ninja   # 用于并行编译
python setup.py install   # 或 pip install -e .

可用 API

• spas_sage2_attn_meansim_topk_cuda：基于 SageAttention2 的 SpargeAttn，我们推荐使用。

• spas_sage2_attn_meansim_cuda：基于 SageAttention2 的 SpargeAttn，但我们不推荐使用。

• spas_sage_attn_meansim_topk_cuda：基于 SageAttention 的 SpargeAttn，我们推荐使用。

• spas_sage_attn_meansim_cuda：基于 SageAttention 的 SpargeAttn，但我们不推荐使用。

使用方法

即插即用用法

只需将 torch.nn.functional.scaled_dot_product_attention API 替换为 spas_sage2_attn_meansim_topk_cuda：

from spas_sage_attn import spas_sage2_attn_meansim_topk_cuda

- attn_output = torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=False)  # is_causal 可以为 True

+ attn_output = spas_sage2_attn_meansim_topk_cuda(q, k, v, topk=0.5, is_causal=False) # is_causal 可以为 True

即插即用 API

from spas_sage_attn import spas_sage2_attn_meansim_topk_cuda

attn_output = spas_sage2_attn_meansim_topk_cuda(q, k, v, topk=0.5, is_causal=False)

您可以调整 topk 参数，以在注意力精度（topk 越高精度越高）和稀疏性（topk 越低稀疏性越高）之间取得平衡。

自定义块稀疏掩码 API

from spas_sage_attn import block_sparse_sage2_attn_cuda

attn_output = block_sparse_sage2_attn_cuda(q, k, v, mask_id=None):

在这个 API 中，我们支持为每个注意力头计算任意块稀疏掩码下的注意力。具体而言，每个头的注意力掩码 mask_id 的形状为 (batch_size, num_heads, ⌈seq_len / 128⌉, ⌈seq_len // 64⌉)，由 0 和 1 组成。目前，块大小为 128×64。

引用

@inproceedings{zhang2025spargeattn,
  title={Spargeattn: 精确稀疏注意力加速任意模型推理},
  author={Zhang, Jintao 和 Xiang, Chendong 和 Huang, Haofeng 和 Wei, Jia 和 Xi, Haocheng 和 Zhu, Jun 和 Chen, Jianfei},
  booktitle={国际机器学习会议 (ICML)},
  year={2025}
}

@article{zhang2026spargeattention2,
  title={SpargeAttention2: 基于混合 Top-k + Top-p 掩码与蒸馏微调的可训练稀疏注意力},
  author={Zhang, Jintao 和 Jiang, Kai 和 Xiang, Chendong 和 Feng, Weiqi 和 Hu, Yuezhou 和 Xi, Haocheng 和 Chen, Jianfei 和 Zhu, Jun},
  journal={arXiv 预印本 arXiv:2602.13515},
  year={2026}
}

@inproceedings{zhang2025sageattention,
  title={SageAttention: 用于即插即用推理加速的精确 8 位注意力}, 
  author={Zhang, Jintao 和 Wei, Jia 和 Zhang, Pengle 和 Zhu, Jun 和 Chen, Jianfei},
  booktitle={国际表征学习会议 (ICLR)},
  year={2025}
}

@inproceedings{zhang2024sageattention2,
  title={Sageattention2: 具有彻底异常值平滑和线程级 int4 量化技术的高效注意力},
  author={Zhang, Jintao 和 Huang, Haofeng 和 Zhang, Pengle 和 Wei, Jia 和 Zhu, Jun 和 Chen, Jianfei},
  booktitle={国际机器学习会议 (ICML)},
  year={2025}
}

SpargeAttn 快速上手指南

SpargeAttn 是一个通用的、无需训练的稀疏注意力加速工具，适用于语言、图像和视频模型推理。它基于 SageAttention2++ 实现，能够显著加速推理过程并保持高精度。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• Python: >= 3.9
• PyTorch: >= 2.3.0
• CUDA (根据显卡架构选择版本):
  • Blackwell: >= 12.8
  • Ada (fp8 支持): >= 12.4
  • Hopper (fp8 支持): >= 12.3
  • Ampere: >= 12.0
• 编译工具: 需要安装 ninja 以支持并行编译。

提示：国内用户建议使用清华源或阿里源加速 Python 包和 CUDA 工具的下载。

安装步骤

1. 安装编译依赖

   pip install ninja
   

2. 克隆代码并安装 进入项目目录后，执行以下命令进行安装（推荐使用 -e 模式以便开发调试）：

   pip install -e .
   

   或者使用传统安装方式：

   python setup.py install
   

基本使用

SpargeAttn 提供了“即插即用”的 API，只需替换原有的 torch.nn.functional.scaled_dot_product_attention 即可。

推荐 API：spas_sage2_attn_meansim_topk_cuda

这是官方最推荐的接口，支持通过 topk 参数灵活平衡注意力精度与稀疏度（topk 越高越精确，越低越稀疏）。

代码示例：

from spas_sage_attn import spas_sage2_attn_meansim_topk_cuda

# q, k, v 为输入的 Query, Key, Value 张量
# topk=0.5 表示保留前 50% 的注意力权重，可根据需求调整
# is_causal 设为 True 可启用因果掩码（适用于解码器模型）
attn_output = spas_sage2_attn_meansim_topk_cuda(q, k, v, topk=0.5, is_causal=False)

迁移示例：

将原有代码：

attn_output = torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=False)

替换为：

attn_output = spas_sage2_attn_meansim_topk_cuda(q, k, v, topk=0.5, is_causal=False)

进阶用法：自定义块稀疏掩码

如果您需要针对每个注意力头自定义块稀疏掩码，可以使用 block_sparse_sage2_attn_cuda 接口。当前支持的块大小为 128×64。

from spas_sage_attn import block_sparse_sage2_attn_cuda

# mask_id 形状：(batch_size, num_heads, ⌈seq_len / 128⌉, ⌈seq_len // 64⌉)
# 由 0 和 1 组成，1 表示保留该块
attn_output = block_sparse_sage2_attn_cuda(q, k, v, mask_id=None)