Scaling language models to long contexts is constrained by the quadratic cost of self-attention. Block sparse attention is a common efficiency workaround, but it introduces its own bottleneck: choosing which blocks to keep normally requires scoring every query-block pair, which is itself quadratic in sequence length.
To address this, researchers propose PISA, a block-sparse attention mechanism built around a pyramid Top-K selection strategy. Rather than scoring all pairs, PISA narrows candidates progressively across multiple levels, starting from a coarse representation of keys and refining toward finer ones.
The method constructs a coarse-to-fine hierarchy of keys. At each level, LogSumExp scoring is applied to a bounded candidate set to choose which candidates advance to the next, finer level, repeating until the finest level is reached. Pooling is used to build O(log N) levels of keys, giving the overall mechanism a complexity of O(N log N), where N is the sequence length.
The authors also developed hardware-aware Triton kernels for both training and inference. These kernels fuse hierarchical routing and LogSumExp scoring without ever materializing the full query-key score matrix.
On language modeling tasks, the paper reports performance comparable to the baseline on benchmarks such as commonsense reasoning, while delivering better results on retrieval tasks.