Source-linked AI summary
Block Sparse Attention with Log-Linear Complexity
Bohao Tang, Zhen Qin, Yuqi Pan, Zheng Li, Pengfei Liu
TL;DR
Conventional block selection remains quadratic because every query must be scored against every key block. PISA uses pooled, coarse-to-fine LogSumExp Top-K routing with O(N log N) selection, and reports comparable commonsense performance and better retrieval results than BSA.
Problem
Conventional Top-K block selection scores every query against every key block, leaving selection with quadratic complexity and a major long-context bottleneck.
Source: §1 · pp. 1–2Method
PISA pools key blocks into a hierarchy and applies LogSumExp Top-K selection from coarse to fine, yielding O(N log N) block-selection complexity.
Source: §1 · p. 2Results
PISA matches BSA on commonsense reasoning and performs better on retrieval, with higher average containment accuracy after pretraining than sparse baselines.
Source: §1 · p. 2, §5 · p. 11Takeaways & Limitations
Within the evaluated settings, pyramid block selection reduces selection complexity while maintaining competitive model performance.
Source: §5 · p. 11, E Limitations · p. 21The experiments explore limited model sizes and pretraining data budgets, both of which can substantially affect performance and relative gains over baselines.
Source: E Limitations · p. 21
Abstract
from arXiv · showhide
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.
1 Introduction
Conventional block selection remains quadratic because every query is scored against every key block. PISA narrows candidates through a pooled hierarchy, achieving log-linear selection complexity and competitive benchmark performance.
1 Introduction
- PISA pools key blocks into a coarse-to-fine hierarchy, retaining Top-K candidates at each level and expanding only those candidates toward finer blocks.Each query evaluates a bounded candidate set with LogSumExp scoring at every level.Source: §1 · p. 2
- With fixed Top-K budget, branching factor, and leaf-block size, PISA reduces overall block-selection complexity to O(N log N).The hierarchy has O(log N) levels, with bounded candidate scoring at each level.Source: §1 · p. 2
- Hardware-aware Triton kernels fuse hierarchical routing and LogSumExp scoring for training and inference without materializing a dense query-key score matrix.The kernels maintain only a bounded candidate set at the current level.Source: §1 · p. 2
- PISA achieves comparable commonsense-reasoning performance and better retrieval results than BSA, while long-sequence benchmarks show more efficient hierarchical routing.Selection benchmarks extend to sequence lengths of 256K tokens.Source: §1 · p. 2
2 Related Work
Related sparse-attention methods differ in when they apply sparsity and whether they train or adapt model weights. The section reviews training-free and trainable approaches and summarizes selected methods’ settings and computational complexity.
2 Related Work
- Table 1 compares selected sparse-attention methods’ training settings and computational complexity, with LLSA’s cost stated per attention pass rather than autoregressive prefill.The listed complexities include selection and attention computation.Source: Table 1 · p. 2
- Sparse attention limits the keys and values each query attends to; selected blocks are computed with tiled kernels, while attention over selected blocks scales linearly with sequence length.Block-sparse attention groups keys and values into contiguous blocks.Source: §2 · p. 2
- Training-free methods reduce inference cost without retraining, using sparse patterns or staged selection during prefill and reducing KV-cache accesses during decoding.Some methods use summaries to select blocks or clusters before scoring tokens within them.Source: §2 · p. 3
- Trainable methods include selected long-range connections, pooled or learned block summaries, learned gates, cross-layer distillation, and continued pretraining.The reviewed approaches include both single-level and hierarchical gist summaries.Source: §2 · p. 3
3 Method
The method replaces exhaustive block scoring with hierarchical Top-K routing over pooled key summaries. Its kernel designs and complexity analysis address training, prefill, and decoding separately.
3.1 Block Sparse Attention
- BSA selects K blocks by scoring all M = ⌈N/C⌉ summaries per query, then attends to original keys and values in at most KC keys.Scoring all block summaries costs O(N^2d/C) across queries.Source: §3.1 · p. 3
3.2 Pyramid Top-K Block Selection
- PISA mean-pools adjacent key blocks into a fine-to-coarse hierarchy, then selects from coarse to fine, expanding retained blocks and scoring bounded candidates with LSE.Intermediate blocks use child summaries, while leaf blocks use exact LSE over original keys.Source: §3.2 · pp. 4–5
3.3 Kernel Implementation
- Training and prefill use two-stage selection, while decoding uses one kernel that continues through leaf-level scoring and updates the cached mean pyramid along the new token’s ancestor path.The two-stage design reuses key blocks across query tiles; the single-stage decoding design avoids multiple kernel launches.Source: §3.3 · pp. 6–7
3.4 Complexity
- With fixed selection and model parameters, PISA requires O(log N) selection per query and O(N log N) for training or prefill; decoding averages O(log N) per token.These bounds include occasional cache-expansion costs during decoding.Source: §3.4 · p. 7
- For the stated configuration, the two-stage design has lower leaf-level Q/K IO cost than the single-stage alternative under the analyzed assumptions.The comparison uses GQ = 16, Qtile = 4, and C = 64.Source: §3.4 · p. 8
4 Experiments
Across three model scales, PISA matches sparse baselines on language modeling and commonsense reasoning and achieves the highest average accuracy among sparse methods on containment tasks. Its long-context retrieval and block-selection results further assess performance and efficiency.
4.1 Experimental setup
- The evaluation compares models at 418M, 1.47B, and 2.67B scales after 100B tokens of 4K pretraining and 10B additional tokens of 16K continued pretraining.Sparse methods use C = 64, with K = 8 and K = 32 for the two stages, respectively.Source: §4.1 · p. 8
4.2 Language Modeling and Downstream Evaluation
- Across all three scales, PISA performs comparably to BSA, NSA, and HiLS on language modeling and commonsense reasoning, and leads sparse methods in average containment accuracy.Full Attention retains a higher average containment accuracy.Source: §4.2 · p. 8, Table 2 · p. 9
4.3 Long-context evaluation
- PISA achieves strong retrieval performance on four RULER task families for 2.67B models evaluated from 1K to 16K contexts after continued pretraining at 16K.Table 3 reports accuracy after 10B tokens of continued pretraining; sparse methods use C = 64 and K = 32.Source: §4.3 · p. 8, Table 3 · p. 10
4.4 Ablation Studies
- Across all three scales, PISA has lower training loss and higher average containment accuracy than PISA-1, PISA-2, and BSA.PISA-1 and PISA-2 truncate the scoring expansion at first and second order, respectively.Source: §4.4 · p. 9, Table 2 · p. 9
- On 100 FDA prompts, PISA achieves the highest average Recall@8 and attention mass ratio; the ratio measures retained mass relative to reference blocks, not total full-attention mass.Figure 2 uses identical Full-Attention queries and keys, K = 8, and the same forced-block policy for all selectors and the reference.Source: §4.4 · p. 10, Figure 2 · p. 10
4.5 Block-selection efficiency
- At 64K, 128K, and 256K, PISA is 2.86×, 5.31×, and 9.95× faster than BSA; Q4 is 1.35×, 1.30×, and 1.33× faster than per-query routing.BSA is faster from 4K to 16K; the timings include selection stages but exclude attention over selected blocks.Source: §4.5 · p. 11, Figure 3 · p. 11
5 Conclusion
PISA combines pyramid Top-K selection with LogSumExp scoring to reduce prefill complexity to O(N log N). The conclusion reports comparable language-modeling and commonsense-reasoning performance to sparse baselines, higher average containment accuracy after pretraining, and lower long-sequence block-selection latency than BSA.
5 Conclusion
- PISA reduces prefill complexity to O(N log N) by scoring bounded candidate sets at each level of pyramid Top-K selection.Its Triton kernels fuse intermediate selection levels and reuse leaf-key blocks across queries.Source: §5 · p. 11
- Experiments across three model scales show comparable language-modeling and commonsense-reasoning performance to sparse baselines, with higher average containment accuracy after pretraining.Source: §5 · p. 11
- Block-selection benchmarks show lower latency than BSA at long sequence lengths.Source: §5 · p. 11
A Analysis of Block Scoring
The analysis distinguishes raw-key scoring from intermediate scoring based on pooled child means, then examines score approximations and grouped-head aggregation. It qualifies score-error comparisons: they do not guarantee correct block rankings or recovery of the globally highest-mass leaf blocks.
A.1 Intermediate block scoring
- For fully visible blocks, raw-key LogSumExp ranks blocks like their full-attention mass, while intermediate scoring uses child means that omit within-child key variation.Normalized LSE over child means is at least as close to normalized raw-key LSE as the mean score; equality holds when logits within each child are equal, but the comparison concerns score error only.Source: A Analysis of Block Scoring · p. 16, §A.1 · pp. 16–17
- The intermediate mean score, PISA score, and raw-key LSE are normalized by their respective input counts so they can be compared on the same scale without changing rankings among equal-count candidates.Source: §A.1 · p. 16
A.2 First- and second-order approximations
- PISA-1 keeps only the mean in its normalized-LSE approximation, PISA-2 adds half the variance, and PISA computes LSE without truncating the expansion.Source: §A.2 · p. 17
A.3 Grouped-head scoring
- Grouped-query attention uses one shared selected block set for query heads sharing a KV head, but summed LSE scores can rank blocks differently from averaged full-attention block masses.Source: §A.3 · p. 18
B Implementation Details
The implementation addresses causal eligibility, fixed selection decisions during backpropagation, and cache growth during decoding. Benchmarks show Q4 has the lowest reported latency from 32K to 256K, while BSA is fastest through 16K.
B Implementation Details
- Causal selection excludes candidates outside the sequence or starting after the query, while retaining eligible first, previous, and current leaf blocks and their ancestor paths.Current-block-path summaries may contain future keys, but are retained without their scores affecting ranking; attention separately masks future tokens.Source: B Implementation Details · p. 18
- Backpropagation holds selected block indices fixed, so gradients flow through attention on selected entries rather than through discrete selection decisions.Source: B Implementation Details · p. 18
- The decode cache stores O((N/C)d) elements per KV head, and occasional buffer expansions preserve O(log N) average decoding cost per generated token at fixed model and selection parameters.Source: B Implementation Details · p. 18
- Q4 reaches 1.35×, 1.30×, and 1.33× speedups over per-query at 64K, 128K, and 256K, respectively, and has the lowest reported latency from 32K to 256K.BSA is fastest through 16K; timings include summary construction and selection but exclude sparse attention.Source: B Implementation Details · pp. 18–19, Table 4 · p. 18
C.1 Model and training configuration
The configuration section describes the shared model architecture and training setup, including matched sparse-method settings and a separate continued-pretraining phase. Continued pretraining extends sequence length to 16K from 4K checkpoints under a new learning-rate schedule.
C.1 Model and training configuration
- All models use GPT-2 BPE with a 50,257-token vocabulary, padded embedding and output matrices, pre-RMSNorm decoder blocks, SiLU-gated feed-forward layers, no dropout, and RoPE on the high-frequency half of each head.Table 5 summarizes architecture dimensions and from-scratch training settings.Source: §C.1 · p. 19, Table 5 · p. 19
- Sparse methods use C = 64 and K = 8 with a shared first, previous, and current eligible-block policy; PISA variants differ only in block score, while BSA uses non-hierarchical selection.BSA also normalizes candidate scores within each query head before summing across heads, unlike PISA and its variants.Source: §C.1 · p. 19
- Continued pretraining uses C = 64 and K = 32 for 10B tokens at 16K, with each run using the same block budget during training and evaluation.Source: §C.1 · p. 19
- Each continued-pretraining run starts from its corresponding 4K checkpoint, extends sequence length to 16K, raises the RoPE base to 80,000, and uses a new warm-up and cosine-decay schedule.Packed inputs preserve document boundaries and use no cross-document attention.Source: §C.1 · p. 19
C.2 Evaluation details
The evaluation uses standard language-modeling and multiple-choice benchmarks, six containment retrieval tasks, and four RULER retrieval families across specified contexts.
C.2 Evaluation details
- WikiText and LAMBADA perplexity and eight multiple-choice tasks are evaluated with the Language Model Evaluation Harness.The reported WikiText–LAMBADA average is the arithmetic mean of their perplexities; Acc-8 weights the eight multiple-choice scores equally.Source: §C.2 · pp. 19–20
- Six containment tasks use zero-shot greedy decoding on validation data, with predictions counted correct when they contain a gold answer as a case-insensitive literal substring.Generation is capped at 48 tokens and stops at a newline or EOS.Source: §C.2 · p. 20
- RULER evaluation covers four needle-in-a-haystack families at context lengths from 1K to 16K.The average is computed across four families and five context lengths before rounding.Source: §C.2 · p. 20
D.1 Downstream evaluation after continued pretraining
After continued pretraining, PISA has comparable language-modeling and multiple-choice performance to BSA but lower containment accuracy; in a separate block-selection diagnostic, it leads on three selection measures.
D.1 Downstream evaluation after continued pretraining
- PISA has slightly lower training loss than BSA and PISA-1 at both scales, comparable perplexity and multiple-choice accuracy to BSA, and lower containment accuracy than BSA.Table 6 reports results after 10B continued-pretraining tokens at 16K for 1.47B- and 2.67B-parameter models.Source: Table 6 · p. 20, §D.1 · p. 20
D.2 Block-selection diagnostic
- The diagnostic replays selectors on fixed queries and keys from a 418M Full Attention model pretrained for 100B tokens, without additional training.It uses all 24 layers, C = 64, K = 8, and 100 FDA prompts; scored positions have more than eight eligible blocks.Source: §D.2 · p. 20
- Recall@8 includes a guaranteed overlap of three forced blocks, and the reference and PISA may rank blocks differently because they use different mass-scoring procedures.Attention mass ratio compares captured mass with the mass captured by the reference set, not total full-attention mass.Source: §D.2 · p. 21
- PISA achieves the highest Recall@8, captured mass, and attention mass ratio in the block-selection diagnostic, with selected blocks capturing nearly as much attention mass as the full-attention reference.The diagnostic uses 100 FDA prompts and shared Full-Attention queries and keys; Recall@8 includes three forced blocks.Source: Table 7 · p. 21, §D.2 · p. 21
E Limitations
Computational constraints bound the model sizes and pretraining data budgets explored, both of which can substantially affect performance and relative gains over baselines.
E Limitations
- Computational constraints limit the model sizes and pretraining data budgets explored in the experiments.Both factors can substantially affect performance and relative gains over baselines.Source: E Limitations · p. 21
- The reported central finding is scoped to evaluated settings: pyramid block selection reduces selection complexity while maintaining competitive model performance.Source: E Limitations · p. 21
- The study does not establish how performance or relative baseline gains would change beyond the explored model sizes and pretraining data budgets.Source: E Limitations · p. 21