Block Sparse Attention with Log-Linear Complexity

Hugging Face Daily Papers Papers

Summary

The paper proposes PISA, a block sparse attention mechanism using a pyramid Top-K selection strategy to achieve O(N log N) complexity, enhancing efficiency for long-context language models with comparable performance on benchmarks and better results on retrieval tasks.

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.
Original Article
View Cached Full Text

Cached at: 09/28/26, 04:02 AM

Paper page - Block Sparse Attention with Log-Linear Complexity

Source: https://huggingface.co/papers/2609.31093

Abstract

Scalinglanguagemodelstolongcontextsislimitedbythequadraticcostofself-attention.Blocksparseattentionoffersanefficientalternative,butselectingtheretainedblocksremainsabottleneck.Conventionalblockselectionrequiresscoringallquery-blockpairsandthereforeremainsquadraticinsequencelength.Toaddressthisissue,weproposePISA,ablock-sparseattentionmechanismthatemploysapyramidTop-Kselectionstrategy.Themainideaistograduallynarrowdownthecandidatesacrossdifferentlevels,makingitmoreefficienttofindthemostrelevantkeys.Specifically,weconstructacoarse-to-finehierarchyofkeysandperformselectionfromthecoarsestlevel.Ateachlevel,LogSumExpscoringisappliedtoaboundedcandidatesettoselectcandidatesforthenextfinerlevel,continuinguntilthefinestlevelisreached.Throughpooling,weconstructO(logN)levelsofkeys,yieldinganoverallcomplexityofO(NlogN),whereNdenotesthesequencelength.Wedevelophardware-awareTritonkernelsforbothtrainingandinference,fusinghierarchicalroutingandLogSumExpscoringwithoutmaterializingthequery-keyscorematrix.Wefurtherevaluateourmethodonlanguagemodelingtasks.Comparedwiththebaseline,ourmethodachievescomparableperformanceonbenchmarkssuchascommonsensereasoningwhiledeliveringbetterresultsonretrievaltasks.

View arXiv pageView PDFAdd to collection

Models citing this paper0

No model linking this paper

Cite arxiv.org/abs/2609.31093 in a model README.md to link it from this page.

Datasets citing this paper0

No dataset linking this paper

Cite arxiv.org/abs/2609.31093 in a dataset README.md to link it from this page.

Spaces citing this paper0

No Space linking this paper

Cite arxiv.org/abs/2609.31093 in a Space README.md to link it from this page.

Collections including this paper0

No Collection including this paper

Add this paper to acollectionto link it from this page.

Similar Articles

MiniMax Sparse Attention

Hugging Face Daily Papers

MiniMax Sparse Attention introduces a blockwise sparse attention mechanism that achieves significant speedups for ultra-long-context LLMs, reducing per-token attention compute by 28.4x at 1M context with wall-clock speedups of 14.2x for prefill and 7.6x for decoding on H800 GPUs. The method is accompanied by an open-source inference kernel and a publicly released multimodal model.

COBS: Cumulant Order Block Sparse Attention

arXiv cs.LG

COBS introduces a cumulant-order block sparse attention method that improves block selection by using compressed second-order statistics, achieving near-dense attention accuracy on long-context benchmarks while significantly reducing KV cache read traffic.