Block Sparse Attention with Log-Linear Complexity
Bohao Tang, Zhen Qin, Yuqi Pan, Zheng Li, Pengfei Liu
cs.LG
2026-09-25
PISA uses pyramid Top-K and LogSumExp so block selection is O(N log N). At 2.67B it matches BSA on commonsense and is ~10x faster at 256K selection.
Quadratic attention still sets the price of long context. Block-sparse attention groups keys and values into tiles of size C and lets each query read only Top-K tiles, so the actual attention pass is linear in N. Choosing those tiles is the remaining tax: scoring every query against every block is O(N²/C). The selection stage eats the linear-attention dividend.
Trainable sparse methods such as NSA, MoBA, and HiLS stay quadratic in prefill. Training-free HiP already reaches O(N log N) but does not train. PISA (Pyramid Sparse Attention) wants both: trainable block sparsity whose selection is log-linear. The authors are at Shanghai Jiao Tong University, Shanghai Innovation Institute, and ByteDance Seed.
Keys are stacked into a pyramid. Level 0 is the raw keys; level 1 forms leaf blocks of size C; coarser levels mean-pool adjacent blocks with branching factor 2, giving O(log N) levels. Selection starts at the coarsest level, scores a bounded candidate set with LogSumExp, keeps Top-K, and expands only those nodes into children. Discarded branches are never scored. With a fixed K and branching factor, each query does O(log N) work, so routing a length-N sequence is O(N log N). Decode routing is O(log N) per step.
LSE is used instead of a mean-pooled dot product because it tracks the hottest keys inside a block, which helps retrieval. PISA-1 keeps only the mean; PISA-2 adds the variance term from the LSE Taylor expansion. Full LSE wins training loss and containment at all three scales. The first, previous, and current leaf blocks are forced in, so the causal local window cannot be pruned away.
Training and prefill use a two-stage Triton kernel that fuses intermediate routing and leaf LSE scoring without materializing a dense score matrix. Decode uses a single-stage kernel to avoid launch overhead. Query heads in one GQA group share selected KV blocks. With Qtile=4, up to four queries reuse the same leaf-key tile.
Models are pretrained from scratch for 100B tokens at 4K, then continued for 10B tokens at 16K. Sparse methods use C=64, with K=8 in pretraining and K=32 in continued pretraining. Baselines are Full Attention, BSA (NSA’s selected branch with mean summaries), NSA, and HiLS, at 418M, 1.47B, and 2.67B.
After 2.67B pretraining:
| Method | Train loss | 8-task commonsense | 6-task containment |
| Full Attention | 2.2209 | 58.23 | 53.01 |
| BSA | 2.2184 | 58.61 | 51.13 |
| PISA | 2.2151 | 58.10 | 52.14 |
Commonsense scores among sparse methods sit on top of each other; PISA does not lead. Containment is the best sparse number and still below full attention. The same pattern holds at 418M and 1.47B (PISA 41.77 / 49.83 vs BSA 38.18 / 48.54). On RULER after 16K continued pretraining, 2.67B averages across four needle families and 1K–16K are 66.24 full, 62.80 PISA, 62.92 PISA-2, and 54.99 BSA. A frozen-Q/K diagnostic at K=8 gives Recall@8 of 90.95 for PISA vs 85.91 for BSA, and attention-mass ratio 99.46 vs 98.42.
Latency is the actual product. Selection kernels at C=64, K=8: BSA is faster from 4K to 16K; PISA is 2.86× / 5.31× / 9.95× faster at 64K / 128K / 256K. Do not pay for the pyramid on short prompts.
If block sparsity is going to be trained, selection has to drop below quadratic or prefill stays stuck. PISA shows that hierarchical routing plus LSE can do O(N log N) without giving away much language modeling, and that retrieval likes LSE more than mean scores. The Triton kernels are the part practitioners can copy.
For 4K–8K serving this paper is mostly irrelevant. For sparse prefill beyond 64K, swapping the selector for a pyramid is the actionable bit. It is a complexity patch on the NSA/BSA line, not a new attention primitive.
The authors note that model size and token budget are limited; relative gains may flip at other scales. On RULER, PISA-2’s 62.92 slightly beats full PISA’s 62.80, so “LSE always wins” is not clean for long-context retrieval. After continued pretraining, 2.67B containment actually favors BSA (68.15 vs PISA 66.40); the main-text retrieval claim is the 4K table. Latency numbers exclude attention over the selected blocks, so end-to-end savings are unmeasured. There is no quality bake-off against training-free O(N log N) methods such as HiP. Forced local blocks inflate Recall@8 by guaranteeing overlap with the reference set.