[ICML2025] SpargeAttention: A training-free sparse attention that accelerates any model inference.
SpargeAttention
Recommended API
We highly recommend using thespassage2attnmeansimtopkcuda and blocksparsesage2attn_cuda APIs. They are plug-and-play and customizable:
Plug-and-Play API
from spassageattn import spassage2attnmeansimtopk_cuda
attnoutput = spassage2attnmeansimtopkcuda(q, k, v, topk=0.5, is_causal=False)
topk to balance between attention accuracy (higher topk is more accurate) and sparsity (lower topk is more sparse).
Customize your Block-Sparse Mask
from spassageattn import blocksparsesage2attncuda
attnoutput = blocksparsesage2attncuda(q, k, v, maskid=None):
In this API, we support computing attention in any block sparse mask per attention head. Specifically, the per-head attention mask maskid has shape (batchsize, numheads, โseqlen / 128โ, โseq_len // 64โ) and consists of 0 and 1. Currently, the block size is 128ร64.
The official implementation of SpargeAttn, a universal training-free sparse attention accelerating language, image, and video models.
SpargeAttention: Accurate and Training-free Sparse Attention
Accelerating Any Model Inference
Project Updates
- Please use the
spassage2attnmeansimtopkcudaandblocksparsesage2attn_cudaAPIs. - [2025-07]: Release a Triton Kernel example.
- [2025-06]: SpargeAttn based on SageAttention2++ is released.
- [2025-05]: Add a very simple usage without tuning or calibration:
o = spassage2attnmeansimtopk_cuda(q, k, v). - [2025-05]: ๐SpargeAttn and SageAttention2 are accepted by ICML 2025!
- [2025-03] Support high acceleration on more GPUs, e.g., H100.
Installation
Base environment
python>=3.9,torch>=2.3.0CUDA:
>=12.8 for Blackwell, >=12.4 for fp8 support on Ada, >=12.3 for fp8 support on Hopper, >=12.0 for Ampere
Install Package
pip install ninja # for parallel compilation
python setup.py install # or pip install -e .Available API
spassage2attnmeansimtopkcuda: SpargeAttn based on SageAttention2 that we recommend using.
spassage2attnmeansimcuda: SpargeAttn based on SageAttention2 that we do not recommend.
spassageattnmeansimtopkcuda: SpargeAttn based on SageAttention that we recommend using.
spassageattnmeansimcuda: SpargeAttn based on SageAttention that we do not recommend.
Usage
Plug-and-Play Usage
Just replacetorch.nn.functional.scaleddotproductattention API using spassage2attnmeansimtopkcuda:
from spassageattn import spassage2attnmeansimtopk_cuda
- attnoutput = torch.nn.functional.scaleddotproductattention(q, k, v, iscausal=False) # iscausal can be True
- attnoutput = spassage2attnmeansimtopkcuda(q, k, v, topk=0.5, iscausal=False) # iscausal can be True
Plug-and-Play API
from spassageattn import spassage2attnmeansimtopk_cuda
attnoutput = spassage2attnmeansimtopkcuda(q, k, v, topk=0.5, is_causal=False)
topk to balance between attention accuracy (higher topk is more accurate) and sparsity (lower topk is more sparse).
Customize your Block-Sparse Mask API
from spassageattn import blocksparsesage2attncuda
attnoutput = blocksparsesage2attncuda(q, k, v, maskid=None):
In this API, we support computing attention for any block-sparse mask per attention head. Specifically, the per-head attention mask maskid has shape (batchsize, numheads, โseqlen / 128โ, โseq_len // 64โ) and consists of 0 and 1. Currently, the block size is 128ร64.
Citation
@inproceedings{zhang2025spargeattn,
title={Spargeattn: Accurate sparse attention accelerating any model inference},
author={Zhang, Jintao and Xiang, Chendong and Huang, Haofeng and Wei, Jia and Xi, Haocheng and Zhu, Jun and Chen, Jianfei},
booktitle={International Conference on Machine Learning (ICML)},
year={2025}
}
@article{zhang2026spargeattention2, title={SpargeAttention2: Trainable Sparse Attention via Hybrid Top-k+ Top-p Masking and Distillation Fine-Tuning}, author={Zhang, Jintao and Jiang, Kai and Xiang, Chendong and Feng, Weiqi and Hu, Yuezhou and Xi, Haocheng and Chen, Jianfei and Zhu, Jun}, journal={arXiv preprint arXiv:2602.13515}, year={2026} }
@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} }
