Block Sparse Attention with Log-Linear Complexity
QUESTION — How can block sparse attention achieve log-linear computational complexity to efficiently scale language models to long contexts?
The researchers propose PISA, a block-sparse attention mechanism that uses a pyramid Top-K selection strategy to reduce computational complexity from quadratic to log-linear. PISA constructs a coarse-to-fine hierarchy of keys through pooling and applies LogSumExp scoring on bounded candidate sets. They develop hardware-aware Triton kernels for both training and inference that fuse hierarchical routing and scoring without materializing the query-key score matrix, achieving comparable results on commonsense reasoning and superior performance on retrieval tasks compared to baselines.
Through pooling, we construct O(log N) levels of keys, yielding an overall complexity of O(Nlog N), where N denotes the sequence length.
Aphelios-Tang · 25 Sept 2026
read the original ↗