huggingface.copaper

Block Sparse Attention with Log-Linear Complexity (PISA)

PISA is a block-sparse attention mechanism that replaces the quadratic block-selection step of conventional sparse attention with a pyramid Top-K strategy: keys are organized into a coarse-to-fine hierarchy of O(log N) levels built via pooling, and LogSumExp scoring narrows candidates level by level down to the finest granularity, giving overall O(N log N) complexity instead of O(N^2). The authors implement hardware-aware Triton kernels for both training and inference that fuse hierarchical routing and scoring without ever materializing the full query-key score matrix, and show on language modeling benchmarks that PISA matches baseline performance on commonsense reasoning while improving results on retrieval tasks.