Block Sparse Attention with Log-Linear Complexity
AuthorsBohao Tang, Zhen Qin, Yuqi Pan, Zheng Li, Pengfei Liu
Resources
PISA makes long-context attention more scalable by hierarchically narrowing relevant key blocks instead of comparing every query with every block.
Key results
PISA block-selection latency speedup over BSA at 256K tokens
PISA block-selection latency speedup over BSA at 64K tokens
PISA block-selection recall against Full Attention reference blocks
PISA average retrieval accuracy for the 2.67B model across RULER task families and context lengths
Model scale used in long-context retrieval evaluation
Tokens used for initial pretraining across model scales
What the paper found
Researchers at ByteDance introduce PISA, or Pyramid Sparse Attention, to remove the quadratic bottleneck in block-sparse attention: conventional methods score every query against every key block, while PISA builds a mean-pooled key pyramid and performs coarse-to-fine Top-K routing with LogSumExp scores. Because each query evaluates only a bounded candidate set across O(log N) levels, selection scales as O(N log N) for prefill and O(log N) per decoding step. Fused Triton kernels avoid materializing the dense query-key score matrix, and a two-stage prefill kernel reuses key blocks across grouped-query attention heads. Across 418M, 1.47B, and 2.67B decoder models trained on 100B tokens, PISA matches block-sparse baselines on language modeling and commonsense reasoning while improving retrieval-oriented block selection: on a diagnostic using Full Attention references, Recall@8 reaches 90.95 percent versus 85.91 percent for BSA. After 10B tokens of continued pretraining at 16K context, PISA records a 62.80 percent average on the RULER needle-in-a-haystack suite for the 2.67B model, above BSA’s 54.99 percent, though below Full Attention’s 66.24 percent. Hardware measurements show the largest gains at long contexts: PISA is 2.86× faster than BSA at 64K tokens and 9.95× faster at 256K, demonstrating that hierarchical routing becomes advantageous as sequences grow.
Original 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(N\log 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.
Read the original paperMore in Attention Mechanisms
Browse all 18 papers →CoWindow Attention: Full Causal Coverage Is a Collective Property
Jingze Shi, Zhangyang Peng, Xianduo Li, Yanlin Qi, Xiaotian Lin, Haoxian Chen, Liangdong Wang, Guang Liu, Yuyu Luo
CoWindow Attention makes long-context transformers faster by letting attention heads collectively cover the past instead of redundantly reading all of it.
HLA: Expressive Hybrid Linear Attention via Chunk-Wise Dynamic Mixing
Zhuokun Chen, Xi Lin, Xiyu Wu, Jiahao He, Jianfei Cai, Bohan Zhuang
HLA makes linear attention more selective by letting each query dynamically choose which compressed chunks of long-context history to access.
MinkowskiPE: Minkowski Positional Encoding for Spatiotemporal Perception
Yuhao Li, Louie Hong Yao, Tianyi Shi, Hanqun Cao, Hongxia Hao, Zhen Zhao, Shengchao Liu
MinkowskiPE gives attention a physics-inspired sense of spacetime, improving both molecular dynamics and video prediction with far fewer parameters.