Block Sparse Attention with Log-Linear Complexity
Bohao Tang, Zhen Qin, Yuqi Pan, Zheng Li, Pengfei Liu
Shanghai Jiao Tong University · Shanghai Innovation Institute · ByteDance Seed
Published September 25, 2026 · on the feed September 28, 2026 · ▲ 30 · arXiv 2609.31093
The abstract
Scaling language models to long contexts is limited by the quadratic cost of self-attention. Block sparse attention offers an efficient alternative, but selecting the retained blocks remains a bottleneck. Conventional block selection requires scoring all query-block pairs and therefore remains quadratic in sequence length. To address this issue, we propose PISA, a block-sparse attention mechanism that employs a pyramid Top-K selection strategy. The main idea is to gradually narrow down the candidates across different levels, making it more efficient to find the most relevant keys. Specifically, we construct a coarse-to-fine hierarchy of keys and perform selection from the coarsest level. At each level, LogSumExp scoring is applied to a bounded candidate set to select candidates for the next finer level, continuing until the finest level is reached. Through pooling, we construct O(log N) levels of keys, yielding an overall complexity of O(Nlog N), where N denotes the sequence length. We develop hardware-aware Triton kernels for both training and inference, fusing hierarchical routing and LogSumExp scoring without materializing the query-key score matrix. We further evaluate our method on language modeling tasks. Compared with the baseline, our method achieves comparable performance on benchmarks such as commonsense reasoning while delivering better results on retrieval tasks.
The Brief
2:28A short podcast summary of the paper. Written and voiced by AI.
Transcript
Block Sparse Attention with Log-Linear Complexity. By Bohao Tang and four others, from Shanghai Jiao Tong University, the Shanghai Innovation Institute, and one other institution. Published on the twenty-fifth of September. No code repository is linked to it.
Scaling language models to long contexts runs into the quadratic cost of self-attention. Block sparse attention offers an efficient alternative, in which only some blocks of keys are retained for attention. But choosing which blocks to retain has become the bottleneck. Conventional block selection has to score every pair of query and block, and so the selection step itself remains quadratic in sequence length. Attention over the chosen blocks got cheaper. Finding those blocks did not.
PISA is a block-sparse attention mechanism that replaces that exhaustive scoring with a pyramid Top-K selection strategy. The main idea is to narrow down the candidates gradually, across different levels, which makes finding the most relevant keys more efficient. Selection starts coarse, and each level only has to consider what the level above it passed down.
Concretely, the keys are arranged into a coarse-to-fine hierarchy, and selection begins at the coarsest level. At each level, LogSumExp scoring is applied to a bounded set of candidates, and the ones it selects become the candidates for the next finer level. This continues until the finest level is reached. The hierarchy is built by pooling, which yields a number of levels that grows with the logarithm of the sequence length, N. Because each level works on a bounded set, the overall complexity comes to N log N, rather than quadratic.
The method ships with hardware-aware Triton kernels, written for both training and inference. They fuse the hierarchical routing together with the LogSumExp scoring, and they do so without ever materializing the query-key score matrix. That matrix is exactly the object whose size grows quadratically, so never building it is what keeps the implementation faithful to the complexity argument.
PISA is evaluated on language modeling tasks, against a baseline. On benchmarks such as commonsense reasoning, it achieves comparable performance. On retrieval tasks, it delivers better results than the baseline.
What PISA targets is a quadratic cost that survived inside a technique meant to escape quadratic cost. Block sparsity had already made the attention step itself cheap. The pyramid brings the search for blocks down to log-linear as well, by never asking any single level to look at more than a bounded handful of candidates, and the reported results show that economy coming without a loss on commonsense reasoning.
Paper Quest · 7 chapters · about 40 min
Play this paper
Learn it by doing: make the calls the authors made, guess the result before you see it, and defend it in a viva graded on your own words. Chapter 1 is right here.
Chapter 1 🧩 · The problem
Sparse attention made attending cheap. Choosing what to attend to is still expensive.
You want a language model to read very long contexts. Full attention costs grow with the square of the length, so you let each query attend to only a few chunks of keys.
The road ahead
- 1 · The problem ~7 min Attending to only a few chunks of keys makes the attending cheap, but choosing which blocks to attend to still compares every query with every block.
- 2 · What came before ~3 min
- 3 · The key idea ~8 min Pool keys into a coarse-to-fine pyramid, then search it from the top, expanding only the best candidates at each level.
- 4 · Choose your own research ~3 min
- 5 · Build the method ~8 min Pool keys into a pyramid, start at the top, score bounded candidates with LogSumExp, pass winners down, then attend over the final blocks.
- 6 · The experiment ~6 min You’ll guess: Compared with the baseline on language modeling benchmarks, how did PISA's pyramid Top-K block-sparse attention perform on commonsense reasoning and on retrieval?
- 7 · The viva ~5 min Three questions from the examiners, answered in your own words. Win it for the paper’s trophy.
Free, with Google or an email address. Your stars and progress are kept.