Block Sparse Attention with Log-Linear Complexity

Paper

PISA: Pyramid Top-$K$ Selection Lowers the Cost of Block-Sparse Attention to $O(N\log N)$

TL;DR — Block-sparse attention makes attention computation linear, but block selection still costs quadratic time. PISA solves this by layering keys into a pyramid via average pooling and narrowing candidates at each level with LogSumExp (LSE) scores. Prefill selection cost drops to $O(N\log N)$ and decode to $O(\log N)$, delivering up to $9.95\times$ speedup in selection latency over existing BSA at 256K tokens.

Key Idea

PISA (Pyramid Sparse Attention) redesigns Top-$K$ block selection, the bottleneck among the two stages of block-sparse attention. Existing block-sparse attention (BSA) scores every query–block pair and then picks the top $K$; this stage alone costs $O(N^2/C)$ (source: §1). PISA stratifies key blocks via average pooling to narrow candidates at a coarse level and expand only the selected candidates to a finer level (source: §3.2). Because the number of candidates evaluated at each level is bounded, selection costs $O(\log N)$ per query and total prefill drops to $O(N\log N)$ (source: §3.4, Table 1).

There are three core design choices. ① Pyramid hierarchy: keys are aggregated bottom-up with block size $C=64$ and branching factor $g=2$. ② LSE scoring: instead of approximating candidate blocks with a simple mean or mean+variance, it scores them with the exact LSE, $\log\sum\exp(\cdot)$. ③ Hardware-aware kernels: hierarchical routing and LSE scoring are fused into a single Triton kernel; training/prefill use a two-stage kernel and decode uses a one-stage kernel (source: §3.3, Algorithm 1).

Background: The Problem They Solved

The cost of self-attention grows quadratically with sequence length $N$, blocking scaling to long contexts (source: §1). Block-sparse attention sidesteps this by restricting each query to a selected subset of key blocks rather than all keys. Grouping into blocks enables tiled kernels, which is efficient (source: §1).

Block-sparse attention splits into two stages (source: §3.1):

1. Top-$K$ block selection — each query decides which key blocks to attend to
2. Attention computation over the selected blocks — using the original keys/values

The problem lies in stage 1, not stage 2. Because each query attends to at most $KC$ keys, attention computation is linear in the sequence, but the selection stage still requires every query to scan every key block, leaving it at $O(N^2/C)$ (source: §1, §3.1). In other words, this paper identifies selection as the real bottleneck in long-context modeling—the research gap it targets.

The authors clearly position their work here. Existing hierarchical methods such as HiP are training-free, while trainable methods such as NSA and HiLS still have $O(N^2)$ prefill complexity. PISA is the only method that is trainable while achieving both $O(N\log N)$ prefill and $O(\log N)$ decode (source: Table 1, §2).

New Approach: PISA (Pyramid Sparse Attention)

3.1 Building the Key-Block Pyramid

The key-block hierarchy is built fine-to-coarse (source: §3.2). Level $0$ consists of singleton blocks of the original keys, level $1$ consists of leaf blocks of size $C=64$, and each subsequent level is formed by concatenating $g=2$ adjacent child blocks. The summary vector $\bar{k}_i^{(\ell)}$ of each block is computed recursively by average pooling over its child summaries:

where $\operatorname{Ch}_\ell(i)$ is the set of child indices of block $i$ at level $\ell$ (source: §3.2, Eq. mean_pyramid). Because this is downsampling, the number of levels is $L=O(\log N)$ (source: §3.4).

3.2 Selection from Coarse to Fine Levels

Selection starts at the coarsest level $L$ and descends to leaf level $1$ (source: §3.2). At each level $\ell$:

1. Compute LSE scores $\mathbf{s}_t^{(\ell)}$ for the candidate set $\mathcal{A}_t^{(\ell)}$
2. Keep only the top-$K$ blocks $\mathcal{I}_t^{(\ell)}$
3. Form the candidate set $\mathcal{A}_t^{(\ell-1)}$ for the next level from the children of the selected blocks

Since each selected block has at most $g$ children, the number of candidates at the next level is bounded by $gK$ (source: §3.2, Eq. candidate_bound). This is the key to reducing quadratic cost to log-linear.

3.3 LSE Scoring: Why LSE Instead of a Simple Mean

The score of candidate block $K_i^{(\ell)}$ is defined as the exact LogSumExp over its child summaries (source: §3.2):

The authors justify this choice with a Taylor expansion. For the mean $\bar{z}$ and variance $\operatorname{Var}(z)$ of $z_1,\dots,z_m$, the LSE expands as follows (source: §4.4, Appendix A):

In other words, the two variants the authors compare are truncated versions of this expansion (source: §4.4):

• PISA-1: first-order approximation → uses only the mean $\bar{z}$
• PISA-2: second-order approximation → uses $\bar{z} + \tfrac12\operatorname{Var}(z)$
• PISA (this method): no truncation; uses the full LSE

According to the Jensen-inequality analysis in Appendix A, scoring an intermediate block by the LSE of its child means is always closer to the original-key LSE than the mean of child means (PISA-1) (source: Appendix A).

3.4 Hardware-Aware Kernels

PISA lets query heads within a GQA group share the same key blocks, and handles training/prefill and decode with different kernels (source: §3.3, Algorithm 1):

• Training/prefill — two-stage kernel: stage 1 performs hierarchical selection at intermediate levels, and stage 2 scores the original key blocks. In stage 2, queries are grouped into tiles ($Q_{\mathrm{tile}}=4$) so that multiple queries reuse the same key blocks, reducing IO. With $G_Q=16$ and $C=64$, $G_Q + C/Q_{\mathrm{tile}} = 32 < 64 = C$, so leaf-level IO is lower than with a single stage (source: §3.4).
• Decode — one-stage kernel: decode has only one query, so there is no cross-query reuse; to avoid the overhead of launching multiple kernels, it processes all the way to $\ell=1$ in a single kernel. It caches the average pyramid and updates only the current leaf average and the ancestor path for each new token (source: §3.3, Appendix B).

Throughout the process, the query–key score matrix is never materialized, which is a strength in terms of memory traffic (source: §1).

How It Works: A Concrete Example

The following is an illustrative example for intuition; the numbers are scaled down and unrelated to the paper’s actual experimental settings ($C=64$, $K=8/32$, $g=2$).

Suppose there are keys $k_1,\dots,k_{16}$, block size $C=4$, and branching factor $g=2$. The pyramid has three levels.

  flowchart TB
    subgraph L2["Level 2 (coarsest)"]
        R["root: [k1..k16]"]
    end
    subgraph L1["Level 1 (middle)"]
        P1["P1: [k1..k8]"]
        P2["P2: [k9..k16]"]
    end
    subgraph L0["Level 1 leaves (C=4)"]
        B1["B1: [k1..k4]"]
        B2["B2: [k5..k8]"]
        B3["B3: [k9..k12]"]
        B4["B4: [k13..k16]"]
    end
    R --> P1 & P2
    P1 --> B1 & B2
    P2 --> B3 & B4

With a block budget of $K=1$, selection proceeds as follows (source: §3.2, Eq. hierarchical_topk).

The key point is that at each level, only the candidates that survived the previous level are scored, not the entire set. As we descend, the number of blocks doubles, but the number of blocks scored is fixed at at most $gK$ (source: §3.2). Therefore, even as the sequence grows and the number of levels increases to $O(\log N)$, the number of scores computed per query stays at $O(\log N)$.

An important nuance is that intermediate levels use the LSE of summaries, while the leaf level uses the LSE of the original keys. Scoring an intermediate block against all original keys would itself be quadratic, so it approximates with summaries and reads the original keys exactly only at the final leaves (source: §3.2, Appendix A).

Performance Validation: Main Results

Experiments train decoder-only backbones at three scales—418M, 1.47B, and 2.67B—from scratch on 100B tokens @ 4K, then continue pretraining (CPT) on 10B tokens @ 16K. Sparse methods share block size $C=64$, with block budgets of $K=8$ for pretraining and $K=32$ for CPT (source: §4.1).

5.1 Language Modeling and Downstream Performance

After training from scratch (100B @ 4K), PISA is on par with BSA, NSA, and HiLS in language modeling and commonsense reasoning, and its average accuracy on six containment (retrieval) tasks is the highest among sparse methods at all three scales (source: §4.2, Table 2). However, Full Attention still has a higher average (source: §4.2).

Lower loss and higher accuracy are better. Numbers excerpted from Table 2 (§4.2).

These results support the claim that “hierarchical selection reduces computational complexity while preserving the blocks needed for downstream modeling” (source: §1, §4.2).

5.2 Long-Context Retrieval (RULER)

After CPT (16K), the 2.67B model was evaluated on RULER’s four needle-in-a-haystack task groups. PISA led with an average accuracy of 62.80, ahead of BSA (54.99) and NSA (61.07); the PISA-2 variant was best among sparse methods at 62.92 (source: Table 3, §4.3).

In particular, PISA showed the best accuracy among sparse methods on multi-key (1K/2K) and multi-value (1K/2K/4K), but at longer lengths of 8K/16K, NSA led on single-key and multi-key tasks (source: Table 3). In other words, PISA’s strength is concentrated in specific tasks and lengths rather than across the entire long-context regime.

5.3 Block Selection Quality (Oracle Diagnosis)

All selectors were replayed with the queries and keys from a Full-Attention checkpoint to measure how well they match the actual full-attention mass (source: §4.4, Appendix D). For $K=8$:

PISA is best on all three metrics, and the selected blocks recover 99.46% of the mass contained in the full-attention reference (source: §4.4, Appendix D). This is direct evidence that LSE scoring produces more accurate Top-$K$ selection. Interestingly, PISA-1 (simple mean) is even lower than BSA, showing that “score accuracy” translates into actual performance differences.

5.4 Selection Latency (Kernel Microbenchmark)

Prefill block-selection latency was measured from 4K to 256K (source: §4.5, Table 6/Appendix B). The most important insight here is that BSA is faster at short lengths.

Numbers excerpted from Appendix B (Table, kernel_grouped_benchmark) and the text in §4.5. Batch 1, 32 query heads, 2 KV heads, $d=64$, $C=64$, $K=8$, $g=2$, BF16.

This crossover summarizes the paper’s core message. PISA’s advantage is asymptotic. Because of constant overhead (hierarchy construction, multi-stage expansion), BSA leads at 16K and below, but as length grows, BSA’s $O(N^2)$ selection becomes exponentially slower, while PISA’s $O(N\log N)$ grows gently (source: §4.5).

Our Perspective: Strengths, Limitations, and Why This Work Matters

Strengths

The biggest contribution is reframing the problem precisely. Whereas prior block-sparse literature focused on the quadratic cost of attention computation, PISA points out that “the selection stage itself remains a bottleneck at $O(N^2/C)$” and attacks that bottleneck directly (source: §1). The second strength is the justification for LSE scoring. It does not merely use LSE; it provides theoretical support via a Taylor expansion and Jensen’s inequality for why it is superior to mean and mean+variance approximations (source: §4.4, Appendix A). The third is engineering honesty. The authors explain the split into a two-stage training/prefill kernel and a one-stage decode kernel via an inequality on leaf-level Q/K IO cost, and they do not hide the fact that BSA is faster at short lengths (source: §3.4, §4.5).

Limitations: What the Authors Acknowledge

The authors explicitly acknowledge that model scale and pretraining data budget were limited by compute constraints, and note that this may affect performance and relative gains over baselines (source: Appendix E, Limitations). In practice, experiments stop at a 2.67B model and 100B tokens, leaving open whether the gains hold for production LLMs at tens to hundreds of billions of parameters.

Limitations: A Reviewer’s Reasoning

Several points warrant caution.

1. Constant overhead and the crossover point. PISA is slower than BSA at 16K and below (source: §4.5). Which is better in real serving contexts (typically 4K–128K) is sensitive to the crossover location, and this crossover depends on hardware and kernel implementation.
2. Size of the performance advantage. It is “on par” in language modeling and commonsense reasoning; clear advantages are concentrated in containment average (best among sparse) and selection-quality metrics (source: §4.2). Even in long-context RULER, NSA leads on certain 8K/16K tasks (source: Table 3). In other words, it is an advantage under specific conditions, not across the board.
3. Approximation in GQA group scoring. The authors themselves note that summing query-head LSEs and averaging full-attention mass can yield different rankings (source: Appendix A, gqa_objective). This is a source of theoretical mismatch.
4. Pyramid cache cost. Decode caches the average pyramid and doubles the buffer each time length doubles, with total copy cost $O((N/C)d\log N)$ (source: Appendix B). The average remains $O(\log N)$, but the fact that this cache occupies VRAM separately from the KV cache must be considered in inference-time memory accounting.

Why This Work Matters

Ultimately, PISA is progress that lowers the complexity of the “selection stage of sparse attention” by one level. Among trainable methods, it achieves both $O(N\log N)$ prefill and $O(\log N)$ decode, and that complexity is implemented in real kernels and validated with a $9.95\times$ speedup at 256K (source: Table 1, §4.5)—meaningfully pushing the bottleneck of long-context modeling further back.

What’s Next?: The Road Ahead

The authors do not explicitly propose future directions, but given the limitations, reasonable next steps are as follows.

1. Scale validation: most urgent is to verify whether the gains and quality loss of $O(N\log N)$ selection hold at 7B+ and multi-trillion-token scale (source: connected to the limitations discussion in Appendix E).
2. Move the crossover earlier: kernel optimization to reduce the constant overhead (hierarchy construction, multi-stage expansion) in the 4K–16K range where PISA is slower than BSA, bringing the advantage in practical contexts forward.
3. Inference-serving integration: follow-up work quantifying the VRAM overhead of the decode pyramid cache and the feasibility of combining it with PagedAttention-style systems.
4. Extending the scoring theory: since LSE was shown superior to mean and mean+variance approximations, explore more accurate aggregation rules that resolve the mismatch in GQA group scoring (source: Appendix A).

In short, PISA is a clean study showing that making selection cheaper is exactly what makes long context more practical. The next step is to verify whether this log-linear selection keeps its promise at production-scale LLMs and in serving systems.

Tables from the paper

Tables converted mechanically from the arXiv e-print LaTeX source. The numbers are the paper’s own and did not pass through a model.

Table 1. Training settings and computational complexity of selected sparse attention methods. $N$ denotes the sequence length. Complexities include selection and attention computation. $\checkmark$ denotes trainable, and $\times$ denotes a training-free method. For LLSA , which is designed for diffusion models, the listed cost is per attention pass rather than autoregressive prefill.

Table 2. Language modeling and downstream evaluation after 100B tokens of pretraining at 4K. Baselines include NSA and HiLS . Loss is the final training loss, and accuracies are reported as percentages. Gray columns show group averages; boldface marks the best sparse loss and group averages at each scale.

Table 3. RULER needle-in-a-haystack results after 10B tokens of continued pretraining at 16K for the 2.67B models. Sparse methods use $C=64$ and $K=32$. Scores are accuracies in percent. The gray Avg column averages scores over the four task families at 1K, 2K, 4K, 8K, and 16K, before rounding. Boldface marks the best sparse result in each column.

Table 4. From-scratch training configuration.

Table 5. Training loss and downstream performance for the 1.47B and 2.67B models after 10B tokens of continued pretraining at 16K. Sparse methods, including NSA , use $C=64$ and $K=32$. Loss is the final training loss. Gray columns show group averages; the containment average covers SWDE, SQuAD Completion, and FDA. The block budget does not apply to Full Attention.