SLA

This repository provides the implementation of SLA (Sparse–Linear Attention), a trainable attention method that fuses sparse and linear attention to accelerate diffusion models.

SLA: Beyond Sparsity in Diffusion Transformers via Fine-Tunable Sparse–Linear Attention
Jintao Zhang, Haoxu Wang, Kai Jiang, Shuo Yang, Kaiwen Zheng, Haocheng Xi, Ziteng Wang, Hongzhou Zhu, Min Zhao, Ion Stoica, Joseph E. Gonzalez, Jianfei Chen, Jun Zhu
Paper: https://www.arxiv.org/pdf/2509.24006

git clone https://github.com/thu-ml/SLA.git
cd SLA

For Linux/macOS:

# Option 1: Use the install script
./clean_install.sh

# Option 2: Manual installation
pip install -e .[triton]

For Windows:

REM Option 1: Use the install script
clean_install.bat

REM Option 2: Manual installation
pip install -e .[triton-windows]

Note:

• Triton is required for running SLA
• Linux/macOS: Use [triton] extra (installs triton>=3.3.0)
• Windows: Use [triton-windows] extra (installs triton-windows>=3.5.1.post22)

If updating from a previous installation, use the clean install scripts to remove old package metadata, or manually run:

# Remove old metadata and reinstall
pip uninstall -y sparse-linear-attention
rm -rf sparse_linear_attention.egg-info build dist  # Linux/macOS
# or on Windows: rd /s /q sparse_linear_attention.egg-info build dist
pip install -e .[triton-windows]  # or [triton] for Linux/macOS

import torch
from sparse_linear_attention import SparseLinearAttention

attn = SparseLinearAttention(
    head_dim=128,
    topk=0.2,                 # = 1 - sparsity
    feature_map="softmax",    # options: elu, relu, softmax
    BLKQ=64,
    BLKK=64,
).cuda()

B, H, L, D = 2, 4, 4096, 128
q = torch.randn((B, H, L, D), dtype=torch.bfloat16, device='cuda')
k = torch.randn((B, H, L, D), dtype=torch.bfloat16, device='cuda')
v = torch.randn((B, H, L, D), dtype=torch.bfloat16, device='cuda')

o = attn(q, k, v)

We provide SageSLA, a very fast SLA (Sparse-Linear Attention) forward pass based on SageAttention. It uses some code from SpargeAttn. Please refer to the SageSLA/ directory for the usage of SageSLA.

If you find this work useful, please cite:

@article{zhang2025sla,
  title={SLA: Beyond Sparsity in Diffusion Transformers via Fine-Tunable Sparse-Linear Attention},
  author={Zhang, Jintao and Wang, Haoxu and Jiang, Kai and Yang, Shuo and Zheng, Kaiwen and Xi, Haocheng and Wang, Ziteng and Zhu, Hongzhou and Zhao, Min and Stoica, Ion and Gonzalez, Joseph E. and Zhu, Jun and Chen, Jianfei},
  journal={arXiv preprint arXiv:2509.24006},
  year={2025}
}
@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}
}