SLA makes sparse attention trainable in diffusion transformers

What the method claims to add to sparse attention

The paper title states the ambition directly: going beyond sparsity in diffusion transformers through fine-tunable sparse-linear attention. The README describes SLA as a trainable attention method that fuses sparse and linear attention to accelerate diffusion models. The framing that matters is trainable. Sparsity in attention is normally a hard decision made once, either at design time or as a fixed inference-time budget, and the top-k of tokens a block attends to is a number you pick and live with. Making the sparse-linear mix something that can be adjusted during training turns the sparsity ratio into a parameter of the method rather than a compromise you accept before training starts. The audience is therefore people running diffusion models where attention cost dominates, and who want the speedup to be a tunable rather than a fixed trade.

The knobs in the constructor tell you the mechanism

The usage example constructs the attention module with a small set of arguments, and reading them is the quickest way to understand the design. head_dim sets the per-head dimension, topk selects how much of the sequence is attended to, feature_map picks the nonlinearity applied to the linear attention branch with elu, relu and softmax listed as the options, and BLKQ and BLKK set the block sizes for the query and key sides. The comment on topk is the one to pay attention to, because it reads as one minus the sparsity, which means a value of 0.2 corresponds to keeping 20 percent of the sequence rather than skipping 20 percent. That inversion is a small thing to get wrong and an expensive thing to debug, since a reversed ratio still runs and still returns plausible tensors.

Getting it installed and running the first forward pass

Installation is a clone and an editable install:

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

The package itself is imported as sparse_linear_attention, and the example runs on CUDA with bfloat16 tensors:

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)

The sample shapes are worth noting: four heads, a sequence of 4096, and a head dimension of 128, which is a diffusion transformer scale rather than a language model one. Every tensor is explicitly moved to the cuda device, and there is no CPU branch in the example, which is the practical hint that this is a GPU-only path.

The packaging tells you what the runtime assumes

setup.py is small and specific. The distribution is named sparse_linear_attention at version 0.1.0, packages are found with find_packages, and python_requires is 3.12 or newer, so an older interpreter is refused at install time rather than failing later. The runtime dependencies are torch 2.7.0 or newer and triton 3.3.0 or newer, and Triton being mandatory rather than optional tells you the attention path is a compiled kernel rather than a PyTorch composition of existing ops. There is one extra, benchmark, which pulls in flash-attn, so the comparison baseline is an optional install rather than a hard dependency. The version number is worth noticing: 0.1.0 on a project whose paper is from 2025 tells you the packaging was not versioned with care, which is a signal about how much to rely on the distribution metadata.

SageSLA is the path for production-shaped throughput

A second implementation sits in the SageSLA/ directory, described as a very fast forward pass based on SageAttention, and it uses some code from SpargeAttn. The README points readers to that directory for usage rather than giving a snippet, so the details of how it is configured are in the code rather than in the documentation. The reason to care is that the two paths solve different problems. The main implementation is the one the paper describes, with a trainable sparse-linear mix, and it is what you want when the sparsity ratio is something you are tuning. SageSLA is a fixed fast forward pass, which is what you want in a serving path where the configuration has already been settled. Using the paper implementation for inference because it is the one you found first is a common and avoidable mistake.

Two branches, and a citation list that dates the work

The README states that two versions of SLA are provided and that you can use either the SLA-1 branch or the main branch, which is an unusual thing for a research repository to say and worth taking at face value. Choosing a branch is a decision you have to make before reading any code, and the two are not interchangeable. The citation list also tells you where the work stands. The primary paper is arXiv 2509.24006, SLA: Beyond Sparsity in Diffusion Transformers via Fine-Tunable Sparse-Linear Attention. A second entry cites SLA2: Sparse-Linear Attention with Learnable Routing and Quantization-aware Training as arXiv 2602.12675, and a third cites SageAttention, Accurate 8-Bit Attention for Plug-and-play Inference Acceleration, at ICLR 2025. The presence of the SLA2 citation alongside the implementation suggests the method has already been extended.

Where this codebase is the wrong choice

Three cases, and they are the usual research-code cases. The first is CPU or edge inference: the example is CUDA-only, the dependency on Triton assumes a GPU toolchain, and there is no CPU path documented. The second is a pinned environment: Python 3.12 minimum, torch 2.7 minimum, triton 3.3 minimum, which is a recent stack, and the dependency is unpinned at the top end so a torch upgrade can change kernel behaviour under you. The third is a team that needs a maintained library, because there are no GitHub releases, the last push was on 2026-10-18, and the version in the packaging metadata is still 0.1.0. Against the alternative of plain PyTorch scaled_dot_product_attention, SLA is faster on long sequences but adds a custom kernel, a Triton dependency and a second branch to maintain. If your sequence length is short enough that full attention is not your bottleneck, the simpler path wins.

Editorial conclusion

SLA is a good fit for a diffusion model team that has already moved past full attention and now wants a sparsity mechanism it can tune during training rather than bake in, and that can afford to run on CUDA with a recent PyTorch and Triton toolchain. It is a poor fit for a CPU inference path, for a project pinned to an older Python or PyTorch, and for anyone who needs a supported, versioned library, since the repository publishes no releases and the packaging metadata still says version 0.1.0. Check three things before you integrate it: that your environment satisfies python_requires of 3.12 or newer and torch 2.7.0 with triton 3.3.0 or newer, that the topk value you pass means what you think it means, since the example comments it as one minus the sparsity, and that you pick the right branch, because the README points at both a main branch and an SLA-1 branch and they are not the same code.

Frequently asked questions

Official sources

1. Issues
2. License: Apache-2.0
3. Project website
4. README
5. thu-ml/SLA on GitHub