Source-linked AI summary
Prism: Spectral-Aware Block-Sparse Attention
Xinghao Wang, Pengyu Wang, Xiaoran Liu, Fangxu Liu, Jason Chu, Kai Song, Xipeng Qiu
TL;DR
Long-context block-sparse attention needs efficient block-importance estimation, but mean-pooled proxies can lose local positional information and require costly token-level corrections. Prism uses dual-band spectral scoring with energy-based calibration to restore these signals using block-level operations, maintaining full-attention accuracy parity while delivering up to 5.1× speedup at 128K tokens.
Problem
Block-sparse attention requires efficient identification of relevant blocks, while mean-pooled coarse-grained proxies are inaccurate and prior methods incur costly search or token-level verification.
Method
Prism is a training-free framework that separates high- and low-frequency block scoring and applies energy-based calibration to preserve attenuated positional information using block-level operations.
Results
5.1× speedup at 128K tokens is reported while Prism maintains performance parity with full attention.
Takeaways & Limitations
Prism provides a block-level approach to precise importance estimation without the token-level overhead common in prior methods.
Takeaways & Limitations
The theoretical mean-pooling analysis assumes semantic content remains relatively stable within each local block.
Abstract
from arXiv · showhide
Block-sparse attention is promising for accelerating long-context LLM pre-filling, yet identifying relevant blocks efficiently remains a bottleneck. Existing methods typically employ coarse-grained attention as a proxy for block importance estimation, but often resort to expensive token-level searching or scoring, resulting in significant selection overhead. In this work, we trace the inaccuracy of standard coarse-grained attention via mean pooling to a theoretical root cause: the interaction between mean pooling and Rotary Positional Embeddings (RoPE). We prove that mean pooling acts as a low-pass filter that induces destructive interference in high-frequency dimensions, effectively creating a "blind spot" for local positional information (e.g., slash patterns). To address this, we introduce Prism, a training-free spectral-aware approach that decomposes block selection into high-frequency and low-frequency branches. By applying energy-based temperature calibration, Prism restores the attenuated positional signals directly from pooled representations, enabling block importance estimation using purely block-level operations, thereby improving efficiency. Extensive evaluations confirm that Prism maintains accuracy parity with full attention while delivering up to $\mathbf{5.1\times}$ speedup.
1 Introduction
Prism identifies why mean-pooled block importance estimates miss local positional structure and replaces costly corrections with spectral-aware, block-level selection. It preserves accuracy parity with full attention while achieving up to 5.1× speedup at 128K tokens.
- Motivation: Mean pooling under RoPE acts as a low-pass filter, attenuating high-frequency positional signals and creating a blind spot for local patterns such as slashes.Destructive interference during aggregation can erase fine-grained positional information from pooled representations.
- Method: Prism separates block importance estimation into high-frequency and low-frequency branches instead of treating embeddings as monolithic vectors.The framework applies energy-based temperature calibration to restore attenuated positional signals from pooled representations.
- Method: Prism estimates block importance using exclusively block-level operations, eliminating the token-level selection bottleneck common in prior methods.This design targets the overhead caused by heuristic search and token-level verification.
- Results: Prism closely matches full-attention accuracy across language modeling, long-context understanding, retrieval, and video-understanding evaluations.The evaluation covers PG19, LongBench, RULER, VideoMME, and LongVideoBench.
- Results: 5.1× speedup at 128K tokens is reported for Prism, alongside latency improvements over baselines across all sequence lengths.The reported trade-off combines state-of-the-art accuracy and speedup.
2 Related Work
Related work motivates dynamic block-sparse attention as a way to reduce long-context computation while adapting to arbitrary dependencies. Its central difficulty is estimating relevant blocks efficiently without incurring full-attention costs.
- Block-Sparse Attention: Block-sparse attention exploits the concentration of attention mass on a small subset of tokens to reduce computational overhead.This sparsity offers an alternative to the quadratic cost of full self-attention.
- Block-Sparse Attention: Static sparse patterns use fixed windows or sink tokens, but can miss long-range dependencies scattered arbitrarily across a sequence.This limitation is described as the “needle in a haystack” problem.
- Block Importance Estimation: Dynamic sparse attention adaptively determines attention patterns, shifting the research focus toward efficient block importance estimation.The goal is to identify relevant Key blocks for each Query block without computing the full attention matrix.
3 Method
The method explains how RoPE makes mean pooling frequency-dependent, attenuating high-frequency positional information while preserving low-frequency semantics. Prism addresses this blind spot through spectral decomposition and calibrated block-level importance estimation.
- Coarse-grained Attention: Mean pooling compresses each query and key block into one representative vector, whose coarse-grained scores generate a binary top-k or top-p block mask.The sequence is divided into N = ⌈L/B⌉ blocks before pooling and mask generation.
- Spectral Structure of RoPE: RoPE’s high-frequency dimensions encode local relative positions, whereas low-frequency dimensions primarily capture global semantic content.These divergent spectral roles correspond to slash-like local patterns and block-sparse semantic patterns.
- Mean Pooling as a Low-Pass Filter: Mean pooling acts as a low-pass filter because geometric summation causes destructive interference at high frequencies and preserves low-frequency signals.The attenuation factor λj(B) measures the pooled-to-original magnitude ratio; negligible rotations yield λj →1.
- Mean Pooling as a Low-Pass Filter: For block size B = 128 and head dimension d = 128, the cutoff is approximately 2j ≈28, creating a Dead Zone through roughly the first 30 dimensions.The Dead Zone has near-zero signal, the Transition Zone remains attenuated, and the Semantic Zone preserves magnitude.
- Empirical Verification: Token-level query representations retain robust Dead Zone energy around RMS ≈1.0, but block pooling collapses that energy while preserving the Semantic Zone.This observed spectral divergence supports the theoretical attenuation analysis.
- Prism: Spectral-Aware Block Selection: Prism separately pools high- and low-frequency bands, then uses branch-specific energy-based temperature calibration to estimate block importance without token-level operations.The high-frequency and low-frequency query/key slices are pooled independently before computing calibrated branch scores.
4 Experiments
Prism is evaluated across long-context language, understanding, retrieval, video, and video-generation workloads, using block-level spectral selection to reduce estimation overhead. Across these settings, it preserves accuracy while improving efficiency relative to full and sparse-attention baselines.
- Experimental Setup: Prism is evaluated on language modeling, long-context understanding, retrieval, video understanding, and video generation.The comparisons include FlashAttention-2, MInference, FlexPrefill, XAttention, and PBS-Attn.
- Language Modeling: 5.1× speedup at 128K accompanies virtually zero perplexity degradation relative to full attention on PG19.XAttention reaches 3.0× at 128K, while Prism maintains the full-attention perplexity level.
- Long-Context Understanding: 41.08 and 39.12 average LongBench scores show negligible degradation below 0.4% versus full attention on Llama-3.1-8B-Instruct and Qwen-3-8B.Prism also outperforms FlexPrefill and XAttention on average for both models.
- Long-Context Retrieval: Prism matches RULER performance using exclusively block-level semantic retrieval, unlike baselines that use token-level estimation heuristics.The reported parity holds across the configured threshold parameters.
- Video Understanding: 64.00 versus 63.11 on Long VideoMME shows Prism surpassing full attention for videos spanning 54K to 107K tokens.Prism also outperforms existing approaches on VideoMME and LongVideoBench while remaining comparable to full attention overall.
- Video Generation: 1.60× speedup with PSNR/SSIM/LPIPS of 21.6/0.748/0.224 exceeds XAttention’s 1.54× and 21.4/0.725/0.228 at comparable video-generation quality.At a higher-quality setting, Prism also improves fidelity and speedup over XAttention.
- Efficiency: Prism reaches 5.1× speedup at 128K and the lowest estimation latency across sequence lengths through purely block-level scoring.At 128K, it uses approximately 20% of FlexPrefill’s memory, while XAttention’s estimation latency rises to approximately 85 ms.
- Ablations: Energy-based calibration consistently dominates fixed-temperature scoring, while overlapping spectral bands stabilize transition-zone coverage.The ablations identify high-frequency recovery beyond the dead zone and spectral overlap as important design choices.
5 Conclusion
The paper identifies mean-pooling spectral attenuation under RoPE as the bottleneck in block-importance estimation and introduces Prism to preserve high-frequency information. Prism combines dual-band scoring with energy-based calibration and achieves near-full-attention performance with up to 5× speedup at 128K.
- Conclusion: Prism addresses mean-pooling spectral attenuation under RoPE with dual-band scoring and energy-based calibration.The framework performs block importance estimation using exclusively block-level operations.
- Conclusion: 5× speedup at 128K is achieved while maintaining performance parity with full attention.The conclusion presents Prism as a scalable solution for long-context and multimodal LLMs.
A Derivation of Spectral Attenuation Factor
The derivation models mean pooling of RoPE features as a geometric sum and defines a spectral attenuation factor for the pooled representation. For small frequencies, this factor approaches a normalized sinc response, establishing mean pooling as a rectangular-window low-pass filter.
- Geometric summation: Mean pooling a RoPE frequency component over block size B produces a geometric-series representation whose magnitude determines retained signal strength.The derivation assumes locally stationary semantic content within the pooling window.
- Magnitude calculation: The attenuation factor λ_j(B) is the ratio of the pooled-vector magnitude to the original content magnitude.The block-start phase has unit magnitude and therefore does not affect this ratio.
- Magnitude calculation: The resulting magnitude expression is the normalized Dirichlet kernel, obtained by applying the half-angle identity to the complex fraction.The kernel describes the spectral response of the discrete periodic summation.
- Sinc approximation: For RoPE dimensions with sufficiently small frequencies, sin(x) ≈ x converts the attenuation expression into a normalized sinc form.RoPE frequencies decay exponentially across dimension index, supporting this approximation for most of the spectrum.
- Sinc approximation: Mean pooling therefore acts as a rectangular-window filter with a sinc-shaped spectral response.This response explains why frequency-dependent attenuation appears in pooled representations.
A.4 Relaxing the Locally Stationary Assumption
The analysis relaxes the locally stationary semantic-content assumption by decomposing pooled representations into mean content and within-block variation. The mean component remains attenuated in the Dead Zone, while arbitrary variation cannot restore the vanished signal.
- Relaxing the assumption: When semantic content varies within a block, the pooled representation decomposes into a mean-content term and a residual variation term.The locally stationary assumption is used only to isolate the RoPE effect and expose the attenuation factor.
- Dead-Zone attenuation: In the Dead Zone, λ_j(B) ≈ 0 removes the mean-content term through destructive interference.This is the same attenuation mechanism derived under locally stationary content.
- Dead-Zone attenuation: Arbitrary intra-block semantic variation cannot restore the mean signal eliminated by destructive interference.Restoration would require implausible oscillation aligned with the opposite RoPE phase across many Dead Zone dimensions.
- Approximation scope: For Qwen3, the small-angle sinc approximation is already highly accurate by the 10th dimension pair and is analytically exact for most higher-index dimensions.At θ_10 ≈ 0.11, the stated relative error is below 0.2%.
- Variation energy: Cross terms in the variation energy are suppressed by weak positional correlations or rapidly rotating Dead Zone phases.The resulting typical-case analysis retains the diagonal variation terms.
B Compatibility with YaRN
The YaRN analysis shows that its frequency interpolation leaves the Dead Zone unchanged and only mildly scales part of the transition region. Consequently, Prism’s spectral split remains effective after Qwen3-8B is extended to 128K context.
- YaRN frequency mapping: YaRN modifies RoPE frequencies through NTK-by-parts interpolation, with coefficients determined by r_j = Lθ_j/(2π).Its additional attention scaling factor does not change the frequency boundaries analyzed here.
- Spectral zones: The spectral zones exhibit distinct YaRN behavior: the Dead Zone and part of the transition remain unchanged, while the transition tail and semantic region are scaled.The zone table links these changes to Prism’s attenuation behavior.
- Compatibility with Prism: YaRN does not modify the Dead Zone frequencies motivating Prism’s high-frequency branch.Within d_high = 64, dimensions up to approximately 47 remain unchanged.
- Compatibility with Prism: Only the tail of the transition region receives mild scaling, which reduces rather than increases attenuation.The low-frequency branch covers the fully scaled long-wavelength region and adjacent transition region.
- Compatibility with Prism: The same spectral split remains effective for Qwen3-8B after YaRN extends the native context from 32K to 128K.The conclusion follows from preserved Dead Zone frequencies and near-lossless low-frequency behavior.
D Experimental Setup Details
Figure 10 presents the PyTorch-style implementation of Prism’s Top-P block-selection process.
- Top-P selection: The Top-P implementation takes block-level probabilities and sorts key blocks for each query block by relevance.It then selects the minimal cumulative-probability set above threshold p and restores original spatial order with scatter.
D.1 Datasets
The evaluation covers language, long-context retrieval, video understanding, and video generation benchmarks, alongside comparisons with several training-free sparse-attention baselines. Video experiments use approximately 107K tokens per hour, while HunyuanVideo uses axis-wise RoPE attenuation criteria and specialized frequency branches.
- D.1 Datasets: PG19 evaluates long-range dependency modeling through perplexity on full-length books.
- D.1 Datasets: LongBench measures broader understanding across 21 bilingual datasets spanning six English and Chinese task categories.
- D.1 Datasets: RULER evaluates the retrieval capability of long-context language models using synthetic tasks.
- D.1 Datasets: VideoMME and LongVideoBench use 327680 maximum pixels per frame and one frame per second, corresponding to approximately 107K tokens per hour.
- D.1 Datasets: HunyuanVideo evaluation samples prompts from VBench, while the model uses 3D-RoPE with temporal, height, and width subspaces.
- Baselines: Prism is compared with MInference, FlexPrefill, XAttention, and PBS-Attn, while SpargeAttention is excluded from the main quantitative comparison to isolate estimation mechanisms.
E Qualitative Video Generation Results
The qualitative video-generation comparison uses identical prompts and sampling settings for full attention, XAttention, and Prism. The reported results indicate that Prism preserves visual content and temporal consistency while improving efficiency, and Figure 12 examines block-size trade-offs.
- E Qualitative Video Generation Results: Prism preserves visual content and temporal consistency relative to full attention and XAttention while improving efficiency in HunyuanVideo generation.
- Block Size Trade-off: Smaller block sizes improve Signal-to-Noise Ratio by reducing spectral attenuation but quadratically increase estimation overhead as the number of blocks N = L/B grows.
- Block Size Trade-off: B = 64 consistently performs better across densities in the accuracy panel, while B = 128 closely follows and matches full attention.