SpargeAttn

环境要求

• Python 环境与 PyTorch（示例代码使用 torch 张量）
• 支持 CUDA 的 GPU（库为 CUDA 内核实现）
• 先安装 ninja 以并行编译
• 需在项目源码根目录执行安装（依赖 setup.py）

安装与启动步骤

1. 1安装编译依赖

   ninja 用于并行编译 CUDA 扩展，必须先装，否则编译会变慢或失败。

   pip install ninja

   复制

2. 2源码安装本包

   在 SpargeAttn 源码根目录执行，两种方式二选一，均可完成编译安装。

   python setup.py install

   复制

3. 3（备选）可编辑安装

   需要修改源码或反复调试时，用可编辑模式安装替代上一步。

   pip install -e .

   复制

4. 4验证导入

   能成功导入核心 API 即说明编译与安装通过。

   python -c "from spas_sage_attn import spas_sage2_attn_meansim_topk_cuda"

   复制

5. 5替换注意力调用

   把原来的 scaled_dot_product_attention 换成 spas_sage2_attn_meansim_topk_cuda，参数保持一致。

   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)

   复制

6. 6调节稀疏度

   按需调整 topk：数值越高越精确，越低越稀疏、速度越快。

7. 7（可选）自定义块稀疏

   需要按注意力头指定稀疏掩码时，改用 block_sparse_sage2_attn_cuda 并传入 mask_id。

   from spas_sage_attn import block_sparse_sage2_attn_cuda
   attn_output = block_sparse_sage2_attn_cuda(q, k, v, mask_id=None)

   复制

关键配置

如何确认成功

执行导入命令无报错；将模型中的 SDPA 替换为该 API 后推理输出正常且速度提升。

常见问题

注意事项

• 本库为训练无关方案，替换注意力调用即可，无需改动模型权重。
• 核心内核基于 CUDA，需在具备 GPU 的环境中编译与运行。
• 安装命令须在项目源码根目录执行，否则找不到 setup.py。
• 两种接口二选一：追求易用选 spas_sage2_attn_meansim_topk_cuda，需要自定义稀疏结构选 block_sparse_sage2_attn_cuda。