FASA: Frequency-aware Sparse Attention

TL;DR

FASA addresses the KV cache memory bottleneck in long-context LLMs by exploiting a RoPE-induced functional sparsity among frequency chunks (FCs). It identifies a small, head-specific set of dominant FCs offline and uses TIP to estimate token importance with this proxy, followed by FAC to compute attention only over a pruned token set, yielding near-oracle accuracy under strong compression. The approach is hardware-aware with two implementations, FASA-M (memory-optimized) and FASA-C (computation-optimized), achieving substantial memory reductions or speedups while remaining compatible with other KV-cache optimizations like PyramidKV. Across long-context benchmarks (LongBench, MATH500, AIME24), FASA consistently outperforms baselines, offering practical, training-free efficiency gains for long-context reasoning and generation in large language models.

Abstract

The deployment of Large Language Models (LLMs) faces a critical bottleneck when handling lengthy inputs: the prohibitive memory footprint of the Key Value (KV) cache. To address this bottleneck, the token pruning paradigm leverages attention sparsity to selectively retain a small, critical subset of tokens. However, existing approaches fall short, with static methods risking irreversible information loss and dynamic strategies employing heuristics that insufficiently capture the query-dependent nature of token importance. We propose FASA, a novel framework that achieves query-aware token eviction by dynamically predicting token importance. FASA stems from a novel insight into RoPE: the discovery of functional sparsity at the frequency-chunk (FC) level. Our key finding is that a small, identifiable subset of "dominant" FCs consistently exhibits high contextual agreement with the full attention head. This provides a robust and computationally free proxy for identifying salient tokens. %making them a powerful and efficient proxy for token importance. Building on this insight, FASA first identifies a critical set of tokens using dominant FCs, and then performs focused attention computation solely on this pruned subset. % Since accessing only a small fraction of the KV cache, FASA drastically lowers memory bandwidth requirements and computational cost. Across a spectrum of long-context tasks, from sequence modeling to complex CoT reasoning, FASA consistently outperforms all token-eviction baselines and achieves near-oracle accuracy, demonstrating remarkable robustness even under constraint budgets. Notably, on LongBench-V1, FASA reaches nearly 100\% of full-KV performance when only keeping 256 tokens, and achieves 2.56$\times$ speedup using just 18.9\% of the cache on AIME24.

Figures (15)

• Figure 1: Functional sparsity of FCs revealed by Contextual Agreement ($\overline{\text{CA}}$) heatmaps. Each heatmap shows $\overline{\text{CA}}$ per FC ($x$-axis) across all heads ($y$-axis). A few "dominant" FCs (bright vertical bands) consistently capture contextual information across attention heads. Results on Qasper ($\mathcal{K}=256$); see Appendix \ref{['part1: sparse fcs investigations']}.
• Figure 2: Method Overview of FASA. First, the TIP stage leverages only dominant FCs to efficiently estimate token importance and select a critical subset of tokens. Then, the FAC stage performs full-dimensional attention exclusively on this reduced subset to generate the next token. See discussion about design in Appendix \ref{['appendix: design choices']}.
• Figure 3: Decoding latency dominates total latency in auto-regressive generation.
• Figure 4: Perplexity results of FASA in comparison with FKV, Oracle, Stream, and Quest on Wikitext (top), PG19 (middle), and C4 corpus (bottom). Token sparsity indicates the retained ratio of tokens.
• Figure 5: FASA under various token budgets ($N_{tip}=16$).