Blog post

PISA replaces the keep-or-drop paradigm of sparse attention with exact-or-approximate: critical blocks get exact attention, and every other block is compressed into a "virtual token" via a block-wise Taylor expansion — all under one shared softmax denominator, which is why it works training-free.

Looking closely at PISA's expansion, its numerator is first-order but its block-mass estimate (denominator) is only zeroth-order — the first-order term cancels exactly inside a block, so the next non-zero correction is second order, and it has a clean closed form (following Su's LSE/Softmax Taylor expansion analysis):

$$ \sum_{i\in\mathcal{B}_j} e^{\boldsymbol{q}\cdot\boldsymbol{k}_i} \approx B,e^{\boldsymbol{q}\cdot\bar{\boldsymbol{k}}_j+w_j},\qquad w_j=\tfrac{1}{2},\overline{\boldsymbol{q}^2}\cdot\boldsymbol{\sigma}_j^2 $$

where $\boldsymbol{\sigma}_j^2$ is the per-dimension key variance of block $j$ and $\overline{\boldsymbol{q}^2}$ is pooled per query block. By Jensen's inequality the zeroth-order estimate always underestimates block mass, and $w_j \ge 0$ corrects it in the right direction — a strictly more accurate, still-closed-form estimate of each approximated block's contribution.

• Negligible extra cost over PISA: $w$ is a precomputed [n_query_blocks, n_kv_blocks] matrix reused for both routing and the virtual-token logits, so the kernel adds only a broadcast add plus a small per-block precompute — no extra tensor-core work.
• More accurate per-block approximation: with the same routed blocks, PWT's single-step output tracks full attention more closely than PISA on real diffusion-transformer activations (cosine similarity 0.9952 vs 0.9932).
• Robust on real activations: real q/k have outlier dimensions that blow up the quadratic term; where $w$ exceeds $\ln B$ (the max possible mass gain over zeroth order) the expansion is invalid, so the term is gated to zero and PWT falls back to PISA's zeroth-order behaviour on exactly those blocks.
• Training-free and trainable: a drop-in op that needs no fine-tuning, yet ships a full backward pass (below) for high-sparsity fine-tuning.

The original PISA release is inference-only. This repo adds a complete backward pass so the operator can be fine-tuned at high sparsity:

• exact branch: block-sparse FlashAttention-style backward;
• approximate branch: fused row-scan + chain kernels, with gradients flowing through all block statistics ($\bar{\boldsymbol{k}}_j$, $\boldsymbol{H}_j$, $\sigma_j^2$); the top-k routing is detached (standard practice, as in SLA);
• correctness is verified, not asserted: fp64 gradcheck, exact agreement with SDPA at density 1.0 (forward and backward), and a from-scratch differentiable reference matched to ~1e-3 (bf16) / ~1e-6 (fp32). See tests/.
• long-sequence friendly: the backward processes query blocks in chunks, so the scratch is O(chunk x n_kv_blocks) instead of O(T x n_kv_blocks) — the full-length buffers would need ~15 GB at T=100k, the chunked path a small constant.

Protocol follows the PISA paper (Appendix E): 87.5% sparsity (density=0.125), block size 64, 50 denoising steps with the first 10 dense, the first and last DiT layers kept dense, plus a ±1 exact diagonal band. 480x832, 81 frames (~33k tokens), Morton (Z-order) tile reordering, same seed/prompt for all variants. Speedup is measured on the DiT forward only (attention + FFN blocks; text encoder and VAE excluded — they are <1% of wall time at this scale).

All four rows below come from a single session against one shared full-attention reference, with the default CuTe backend and, for reference, the same variants pinned to Triton:

The CuTe rows are measured with the FA4-style conditional softmax rescaling that is now the default (~4% faster forward); the small PSNR deltas vs the eager-rescaling measurement (21.32 / 22.20 dB) come from that numerics change (oracle error 2.2e-3 -> 2.3e-3) folded through the diffusion trajectory, and sit within its run-to-run variance.

These supersede the earlier non-reordered numbers (21.55 / 21.72 dB at 1.63x). The Morton reorder makes each 64-token block cover a compact 3D neighbourhood and the diagonal band=1 becomes spatially local, so PWT's second-order term absorbs the higher within-tile variance: +0.48 dB for CuTe PWT and +0.74 dB for Triton PWT at essentially the same wall time (a fused gather+transpose kernel keeps the reorder overhead ≤0.5% of DiT time; ~300 μs per attention call vs ~2.4 ms for the naive path). Reorder is enabled by default; pass reorder="none" for a strict PISA head-to-head. The two backends agree within 0.3 dB, as expected from kernels that match to ~2e-3 in bf16.

full attention | PISA | PWT (left to right), CuTe backend — videos in assets/wan14b_final

At the single-sample end-to-end level PWT and PISA are on par (their difference sits within the run-to-run variance of the diffusion trajectory); both are visually indistinguishable from full attention while computing only 12.5% of attention blocks exactly and approximating the rest as Taylor virtual tokens. PWT's advantage is at the operator level (lower single-step approximation error, as above) and is what makes it a cleaner starting point for the trainable regime. The self-attention operator is

3x faster than dense at this sparsity; the DiT-level number is lower because the dense warmup steps and dense layers are, by design, not sparsified.

Requirements: torch >= 2.7, triton >= 3.4, a CUDA GPU.

Optional, to unlock the faster CuTe backend (used automatically when present on an SM100-class GPU): nvidia-cutlass-dsl, cuda-python. Without them everything runs on the Triton kernels.

git clone https://github.com/HaoyiZhu/Piecewise-Taylor-Attention.git
cd Piecewise-Taylor-Attention
pip install -e .

import torch
from pwt_attn import piecewise_sparse_attention
q = torch.randn(1, 24, 32760, 128, device="cuda", dtype=torch.bfloat16)
k, v = torch.randn_like(q), torch.randn_like(q)
out = piecewise_sparse_attention(q, k, v, density=0.125, variant="pwt")
# trainable: full backward through all block statistics
q, k, v = (x.requires_grad_(True) for x in (q, k, v))
out = piecewise_sparse_attention(q, k, v, density=0.125, variant="pwt")
out.float().square().mean().backward()

Knobs: variant="pisa" reproduces the original PISA (0th+1st order) with the same kernels; n_sink_blocks=n forces the first n KV blocks (e.g. an MMDiT text prefix) into the exact branch; band=b keeps the |i-j|<=b diagonal exact (a cheap locality prior); cov_bias=True adds covariance-aware block selection to the router; backend="auto" (default) picks the faster CuTe SM100 kernels when they can run and falls back to Triton — pin it with "cute" / "triton" (see below).

python examples/wan_t2v_demo.py \
 --model Wan-AI/Wan2.1-T2V-14B-Diffusers \
 --density 0.125 --dense_steps 10 \
 --start_layer_idx 1 --num_last_full_layers 1 \
 --num_frames 81 --height 480 --width 832 --steps 50

Generates the same video with full / PISA / PWT (same seed) and reports PSNR against full plus the DiT-only speedup.

pwt_attn/cute is a CuTe-DSL implementation of the same operator that runs about 1.5x faster forward and 1.8x faster backward than the Triton kernels (measured kernel-for-kernel on identical inputs at 87.5% sparsity; reproduce on your own hardware with python bench/bench_cute.py). Its architecture follows Sol-Attn (NVlabs/Sana): the softmax stays resident in the tensor-core accumulator (TMEM) instead of round-tripping through shared memory, blocks are gathered with TMA, and the mainloop is warp-specialized.

It is used automatically. piecewise_sparse_attention dispatches to CuTe whenever it can run for the given inputs — an SM100-class GPU (compute capability 10.x), bf16, head_dim=128, plus nvidia-cutlass-dsl and cuda-python — and silently falls back to the portable Triton kernels otherwise. Nothing to change in existing code:

from pwt_attn import piecewise_sparse_attention
out = piecewise_sparse_attention(q, k, v, density=0.125, variant="pwt")

Pin a backend when you need determinism across machines, or query support:

from pwt_attn import piecewise_sparse_attention
from pwt_attn.cute import is_available
piecewise_sparse_attention(q, k, v, backend="cute") # force CuTe (raises if unsupported)
piecewise_sparse_attention(q, k, v, backend="triton") # force portable Triton
is_available() # True on a capable GPU

Generation quality is unchanged: the two backends agree to ~2e-3 in bf16 and land within 0.1 dB PSNR of each other on the Wan2.1-14B comparison above — inside the diffusion trajectory's run-to-run variance. The DiT-level gain is smaller than the kernel-level gain because the sparsified attention is only a fraction of DiT time under this protocol (the dense warmup steps and the dense layers are, by design, untouched).

The CuTe kernels are derivative works of Sol-Attn (Apache-2.0), which itself builds on FlashAttention (BSD-3-Clause); a minimal subset of both is vendored under pwt_attn/cute/_vendor/. Full license texts, the vendored file list and the modifications made are documented in third_party/NOTICE.md.

• FlashAttention-4-style bf16 backward (this update): all fp32 [B,H,T,D] gradient buffers are gone — the bf16 dq/dk/dv outputs double as accumulation surfaces with a deterministic fixed-order read-modify-write (fp32 math in registers, one bf16 rounding per store; FA4's precision policy). The backward's own peak memory drops ~2 GB at the Wan-14B shape and now sits below full-attention FlashAttention (4.7 vs 5.2 GB peak in the same harness), with zero test-tolerance changes.
• FA4 conditional softmax rescaling, default on (this update): the online-softmax accumulator correction is skipped unless the running max moved by >8 log2 units (FA4's threshold) — ~4% faster forward, oracle error unchanged regime (2.3e-3 vs 1e-2 tolerance). The FA4 cubic-poly exp2 was also implemented and measured ~11% slower on this kernel's FMA-saturated warps, so it stays available but off (PWT_FWD_POLY_EXP2).
• Forward peak trims (this update): the routing score pipeline is in-place, the ranking spread aliases the raw spread storage, and the kernel's route-bias table is built in the gated spread's own storage (bias_inplace) — the forward's own peak drops ~16% (499 → 419 MiB at the Wan-14B shape).
• Cross-arch dispatch + experimental Hopper kernels (this update): SM100 (Blackwell) stays the default CuTe path; new SM90 (Hopper) WGMMA forward/backward kernels following FlashAttention-4's flash_*_sm90 architecture live behind PWT_CUTE_SM90=1 and pass the full test suite (incl. fp64 oracles) on SM90 hardware — see pwt_attn/cute/SM90_STATUS.md. Every other GPU keeps the portable Triton path, which shares the bf16 low-memory backward.
• Morton tile reorder + fused gather kernel (latest): the ±1 exact band now covers a compact 3D neighbourhood; a single Triton kernel fuses the [B,T,H,D]->[B,H,T,D] transpose with the token gather and its inverse (~8x faster than the eager x[:,:,perm].contiguous() chain), and both directions are wrapped in torch.autograd.Function so the reorder participates cleanly in the backward pass (gradcheck passes bit-exact vs eager, see tests/test_reorder.py). This is what buys the +0.48-0.74 dB in the table above. Enabled by default (reorder="morton").
• hc-free block statistics (latest): the [B,H,NT,D,D] per-block cross tensor (~384 MiB at the Wan-14B shape) is no longer materialised — the router's per-block Frobenius norm is atomic-summed inside the chunk kernel and the global first-order tensor is recovered via the exact identity sum_j (K - kc_j)^T V = K^T V - sum_j kc_j (x) vc_j. Cuts ~200 us from every attention call's preamble.
• Router polish (latest): top_k(sorted=False) (set-identical, ~20% cheaper) and a clamped ranking spread — saturated Taylor blocks get the bounded bonus min(w, ln B) for ranking, while the kernel still uses the gated spread for the virtual-token logits. Cleaner rank on the ~7% of rows where the Taylor term saturates on real Wan activations.
• Added the CuTe SM100 kernels (pwt_attn/cute, ~1.5x forward / ~1.8x backward vs Triton) and made them the default backend, with automatic fallback to Triton and a backend= override; the Wan2.1-14B results above were re-measured with them.
• Added the trainable backward pass (above) and its correctness/robustness test suite; the forward semantics are unchanged.
• Made all shape parameters runtime arguments — no per-sequence-length recompilation, so the kernels slot into variable-resolution video-DiT training loops as-is.
• Exposed router knobs (cov_bias, a ±1 exact diagonal band); on the single-sample video benchmark their net quality effect is within the diffusion trajectory's run-to-run variance, so they are conservative defaults rather than headline numbers.
• Gate (not clamp) the second-order term at ln B — fixes a real-activation failure mode (mechanism in the TL;DR and blog post).

Note: per-layer / per-step sparsity schedules can squeeze out more training-free speedup (attention sharpness varies a lot across layers), but they are inference-only tricks and are intentionally left out of the core operator to keep the training path simple.

pytest tests/ -q

Covers density=1.0 forward/backward equivalence with SDPA (both variants), sparse forward/backward finiteness on non-divisible lengths, and kernel reuse across variable sequence lengths.

tests/test_cute.py additionally validates the optional CuTe kernels against an fp64 oracle (forward), an fp64 gradient decomposition (backward), SDPA at density 1.0, and the Triton path under matched routing. It skips itself automatically when no SM100-class GPU is present.

• Forward kernels are adapted from the official PISA implementation (Haopeng Li et al., MIT license). The original piecewise_attn package is kept intact in this repo as the PISA baseline.
• The exact-branch backward kernels are adapted from the official SLA implementation (SLA team, Apache-2.0).
• The optional CuTe kernels (pwt_attn/cute) follow the SM100 architecture of Sol-Attn (NVlabs/Sana, Apache-2.0) and vendor a minimal subset of it, which itself builds on FlashAttention (Tri Dao, BSD-3-Clause). See third_party/NOTICE.md.
• The second-order correction follows Su Jianlin's LSE/Softmax Taylor expansion analysis.

If you find this useful, please cite the blog post:

@misc{zhu2026pwt,
 title = {Sparse Linear Attention: When Sparsity Meets Linear Attention},
 author = {Zhu, Haoyi},
 year = {2026},
 howpublished = {\url{https://www.haoyizhu.site/blog/sparse-linear-attention/}},
 note = {Blog post}
}

and PISA, which this work builds on:

@article{li2026pisa,
 title = {PISA: Piecewise Sparse Attention Is Wiser for Efficient Diffusion Transformers},
 author = {Li, Haopeng and Shao, Shitong and Zhong, Wenliang and Zhou, Zikai and Bai, Lichen and Xiong, Hui and Xie, Zeke},
 journal = {arXiv preprint arXiv:2602.01077},
 year = {2026}
}