PaperlayerView exact revision on arXiv

Source-linked AI summary

FASA: Frequency-aware Sparse Attention

Yifei Wang, Yueqi Wang, Zhenrui Yue, Huimin Zeng, Yong Wang, Ismini Lourentzou, Zhengzhong Tu, Xiangxiang Chu, Julian McAuley

arXiv:2602.03152v3cs.CL

TL;DR

Long-context LLM inference is constrained by the memory and computational costs of the KV cache. FASA uses dominant RoPE frequency chunks for query-aware token selection and focused attention, achieving near-full-KV performance across long-context and generation tasks while surpassing baselines.

  • Problem

    Long-context LLM applications face memory and computational challenges from the KV cache’s linear growth and repeated full-cache access.

  • Method

    FASA uses sparse, task-agnostic dominant RoPE frequency chunks to predict important tokens dynamically, then computes focused attention over the retained subset.

  • Results

    Less than 0.7% performance reduction from full-KV cache, while consistently surpassing baselines across long-context benchmarks, sequence modeling, and LongCoT reasoning.

  • Takeaways & Limitations

    FASA’s memory- and speed-optimized variants provide a practical approach to efficient long-context inference under constrained KV-cache budgets.

  • Takeaways & Limitations

    The evaluation uses open-source models whose societal biases FASA does not address directly.

Abstract

from arXiv · show

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. 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. 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.

1 INTRODUCTION

FASA addresses the linear memory and computational costs of long-context KV caches through training-free, query-aware token importance prediction based on RoPE-induced frequency-chunk sparsity. Its two-stage design achieves near-full-KV performance while reducing memory or computation across long-context, sequence-modeling, and LongCoT tasks.

  • Motivation: Long-context processing makes the KV cache grow linearly, creating memory and computational challenges for tasks such as repository-level code analysis and document summarization.
  • Motivation: Existing token-eviction strategies remove tokens statically or adaptively, risking irreversible information loss or inadequately modeling query-dependent importance.
  • Method: FASA uses RoPE-induced functional sparsity among frequency chunks, where dominant FCs provide a computationally free proxy for salient tokens.
  • Method: FASA first predicts token importance with dominant FCs, then performs focused attention computation on the resulting critical-token subset with minimal overhead.Dominant-FC identification is one-time and task-invariant.
  • Results: Less than 0.7% performance reduction preserves full-KV-comparable results, while FASA-M provides 8× KV-cache compression and FASA-C delivers 2.6× speedups across three task paradigms.The evaluated paradigms are long-context benchmarking, long-sequence modeling, and LongCoT reasoning.

2 RELATED WORKS

Prior KV-cache compression work primarily exploits query-dependent attention sparsity through token eviction or low-rank approximation, but existing methods incur information loss, rely on heuristics, or add memory overhead. FASA addresses these limitations by operating in-place on the KV cache without auxiliary memory overhead.

  • Token Eviction: Token-eviction methods exploit query-dependent attention sparsity, but Stream’s rigid retention of initial and recent tokens can discard crucial intermediate information.SnapKV improves on Stream through one-time, prefill-stage filtering.
  • Low-rank Compression: Low-rank compression assumes KV-cache information lies in a low-dimensional subspace, while SparQ’s query-magnitude heuristic is suboptimal because it is head-agnostic.SparQ selects key dimensions based on high query-vector magnitudes.
  • Low-rank Compression: LoKi uses PCA to project key states into a compact subspace but requires substantial memory for projection matrices, whereas FASA operates in-place without auxiliary memory overhead.FASA’s in-place operation circumvents the projection-matrix memory cost.

3 OBSERVATION

FASA’s observation is that RoPE’s frequency chunks have heterogeneous roles, with a sparse, broadly recurring subset of dominant FCs closely matching full-head contextual selection. The Contextual Agreement metric identifies these FCs, whose collective attention can predict salient tokens and outperform token-eviction baselines under tight budgets.

  • Frequency-Chunk Perspective on RoPE: RoPE partitions each d-dimensional vector into d/2 orthogonal 2D frequency chunks, each associated with a distinct base angular frequency.The resulting frequency-domain representation provides the basis for analyzing FC-specific functionality.
  • Position vs. Semantics: Different Roles of FCs: High-frequency FCs primarily construct positional patterns, whereas low-frequency FCs specialize in carrying contextual information.This functional division motivates focusing token-importance prediction on contextual FCs.
  • Contextual Agreement: Contextual Agreement measures alignment between a single FC’s attention pattern and the full attention head using normalized overlap of their top-K token sets.Mean CA is computed across samples to assess FC importance robustly.
  • Sparse and Universal Idom: Less than 1% of FCs are dominant, while approximately 90% or more are non-dominant with typically < 0.15 CA scores across architectures and scales.This establishes both sparsity and universality of FC functionality.
  • Reconstructing Functionality from Idom: 43% accuracy using 1/8 of components under budget 64 surpasses SnapKV by an average of 10.3% across all budget levels.The result supports reconstructing full-head functionality from a small dominant-FC subset for token-importance prediction.

4 METHOD

FASA uses a training-free, coarse-to-fine method that first estimates token importance from calibrated dominant frequency chunks, then applies full-dimensional attention only to the selected tokens. This design reduces attention computation and KV-cache memory movement while preserving token positions for high-fidelity generation.

  • FASA overview: FASA’s TIP stage uses dominant frequency chunks to select contextually salient tokens, followed by FAC attention over only that reduced subset.The two stages provide a computationally frugal importance proxy and full-fidelity attention on selected tokens.
  • Token Importance Predictor: Dominant frequency indices are identified once offline per attention head by maximizing expected average contextual-agreement scores over a calibration dataset.The resulting set is task-agnostic, robustly identified from few samples, and has negligible calibration cost.
  • Token Importance Predictor: At each decoding step, TIP aggregates only dominant-frequency contributions to score tokens and selects the top-Nfac indices for FAC.This training-free online prediction bypasses computation for non-dominant frequencies.
  • Focused Attention Computation: FAC gathers keys and values at the selected indices, computes attention on them, and preserves their original absolute positions.Preserving positions maintains positional embeddings and avoids degradation from positional distortion.
  • Hardware-aware variants: FASA-M reduces GPU memory by offloading value caches and non-dominant key components, whereas FASA-C prioritizes inference speed.FASA-M can pair with prefetching to mitigate CPU-GPU transfer latency.
  • Computational analysis: FASA’s TIP complexity is O(2tNtip), while FAC complexity is O(Nfacd), with dominant-frequency detection performed offline once.The low-dimensional TIP stage and reduced-token FAC stage avoid full attention’s O(td) per-head operations.

5 EXPERIMENTS

FASA is evaluated against token-eviction baselines and full-KV oracles across long-context understanding, sequence modeling, and reasoning tasks. It preserves long-range dependencies and evolving reasoning traces while remaining robust to compression, calibration choices, and implementation trade-offs.

  • Evaluation Setup: Experiments compare FASA with Stream, SnapKV, RKV, Quest, and H2O, using FKV and Oracle as full-KV and look-ahead upper bounds.LongBench-V1 comparisons retain a constant 256-token budget and 25% of FCs for FASA.
  • Evaluation Benchmarks: Evaluations cover LongBench long-context understanding and perplexity-based modeling on PG-19, WikiText, and C4.The paper frames these benchmarks as tests of critical-information identification and long-sequence modeling.
  • Long-Sequence Modeling: FASA captures long-term dependencies more effectively than Stream and Quest, whose attention-sink and page-level heuristics discard critical noncontiguous context and increase perplexity.The comparison uses token-by-token decoding with eviction applied iteratively before prediction.
  • Long-CoT Reasoning: On R1-Llama, SnapKV reaches 21.6 accuracy versus 72.4 for FKV, while FASA surpasses standard baselines and specialized R-KV for long-CoT reasoning.Static compression heuristics fail to preserve dynamically shifting thought traces and the logical dependencies required for reasoning.
  • Efficiency Analysis: FASA-M delivers pronounced memory savings for long sequences, while its CPU-GPU transfer overhead can be mitigated through asynchronous prefetching.FASA-C is implemented with Triton, and the analysis evaluates both variants.
  • Robustness and Trade-offs: FASA performance is largely insensitive to calibration-window size K and calibration dataset, while Ntip and Nfac trade token-selection precision against retained-context volume.Smaller K values can be slightly better, and low coefficient of variation indicates stable FC detection across calibration sources.

6 EXTENDING FASA TO NON-ROPE MODELS

Section 6 examines whether FASA generalizes beyond full-RoPE architectures by testing functional sparsity and FASA performance under ALiBi and Partial-RoPE. Across these position-encoding variants, FASA matches or surpasses FKV without significant performance trade-offs.

  • Motivation: FASA’s extension to non-RoPE models depends on whether functional sparsity emerges under alternative position-encoding schemes.The section investigates this property before evaluating FASA in those frameworks.
  • Functional Sparsity in ALiBi and Partial-RoPE (MLA): The analysis covers ALiBi, which adds head-specific linear distance biases to attention logits, and Partial-RoPE in MLA, which applies RoPE to only part of the head dimensions.These variants probe sparsity behavior beyond standard full-RoPE attention.
  • FASA Evaluation on Other PEs: Across diverse position-encoding architectures, FASA matches or surpasses FKV without significant performance trade-offs.The reported results establish broad applicability beyond RoPE.

7 CONCLUSION

The conclusion identifies functional sparsity among frequency chunks (FCs) and presents FASA as a coarse-to-fine framework for reducing LLM KV-cache memory and bandwidth costs. FASA uses dominant FCs for dynamic, query-aware token selection without costly training before focused attention computation.

  • FASA addresses the memory footprint and bandwidth introduced by the KV cache in LLMs.
  • Functional sparsity enables a subset of dominant FCs to show high contextual awareness.
  • FASA uses a coarse-to-fine two-stage framework for token selection and focused attention computation.
  • Its first stage uses dominant FCs for dynamic, query-aware token selection without costly training.

ETHICS STATEMENT … A.2 TASK-INVARIANCE PROPERTY OF FUNCTIONAL SPARSITY

The paper frames FASA as an efficiency-focused method with broad accessibility and sustainability benefits, while acknowledging dual-use and model-bias limitations. Additional analyses show functional sparsity is stable across architectures and scales, and dominant frequency chunks have task-agnostic saliency.

  • ETHICS STATEMENT: FASA reduces LLM inference memory and computational overhead, potentially improving accessibility, affordability, and environmental sustainability.The authors emphasize benefits for researchers and institutions with limited resources deploying long-context models.
  • ETHICS STATEMENT: Efficiency improvements may lower barriers for malicious actors to deploy existing models for misinformation or spam at scale.The work does not create new harmful-content capabilities; it optimizes existing-model performance.
  • ETHICS STATEMENT: Experiments used public LongBench, MATH, and AIME benchmarks with open-source models, excluding private, sensitive, and user-generated data.The authors note that evaluated foundation models may inherit societal biases, which FASA does not directly address.
  • A.1 FURTHER GENERALIZATION ON MODEL SCALES AND ARCHITECHTURES: Functional sparsity remains similar between Qwen2.5-14B-Instruct and its 1M-token long-context variant.Figure 10 compares Mean Contextual Agreement heatmaps across frequency chunks and attention heads, calibrated on Qasper.
  • A.1 FURTHER GENERALIZATION ON MODEL SCALES AND ARCHITECHTURES: Functional sparsity persists across 3B and 32B models, with stable dominant frequency-chunk patterns.The observed stability supports functional sparsity as a scalable characteristic of RoPE.
  • A.1 FURTHER GENERALIZATION ON MODEL SCALES AND ARCHITECHTURES: Cross-architectural and cross-scale analyses characterize functional sparsity as a universal, stable property intrinsic to RoPE rather than model training dynamics or size.The authors describe the hierarchy of frequency roles as fundamental and predetermined.
  • A.2 TASK-INVARIANCE PROPERTY OF FUNCTIONAL SPARSITY: Dominant frequency-chunk saliency is largely task-agnostic, with highly consistent importance rankings across question answering and summarization.Strong alignment between task-specific saliency maps suggests these chunks serve a fundamental architectural role.

A.3 MORE ANALYSIS RESULTS · A.4 QUANTITATIVE EVIDENCE ON SPARSITY & UNIVERSALITY & TASK- INVARIANCE · B EXPERIMENTS DETAILS

FASA’s functional sparsity principle is universal, but dominant frequency-chunks (FCs) specialize dynamically across model depth and attention heads. The supplied evidence visualizes agreement patterns across heads and layers and documents cross-task and score-range analyses of dominant FCs.

  • A.3 MORE ANALYSIS RESULTS: Dominant FCs are not static across layers, despite functional sparsity remaining universal.Their changing identities indicate layer-dependent specialization.
  • A.3 MORE ANALYSIS RESULTS: Agreement-score heatmaps compare attention heads on Qasper and GovReport from LongBench-V1 using K = 256.The comparisons use Mistral-7B-Instruct-v0.3.
  • A.3 MORE ANALYSIS RESULTS: Agreement-score heatmaps also examine dominant-FC behavior across different layers.Figure 13 provides the layer-wise analysis underlying the reported cross-layer specialization.
  • A.3 MORE ANALYSIS RESULTS: Dominant FCs exhibit specialization across model depth and individual attention heads.This dynamic behavior reflects division of labor within the transformer architecture.
  • A.4 QUANTITATIVE EVIDENCE ON SPARSITY & UNIVERSALITY & TASK- INVARIANCE: The quantitative sparsity analysis reports the ratio of dominant FCs to non-dominant FCs.The supplied passage identifies this analysis but does not provide the table’s numerical ratios.
  • A.4 QUANTITATIVE EVIDENCE ON SPARSITY & UNIVERSALITY & TASK- INVARIANCE: The cross-task analysis measures dominant-FC overlap as percentages between row and column datasets.Each sub-table reports the intersection percentage for a dataset pair.
  • A.4 QUANTITATIVE EVIDENCE ON SPARSITY & UNIVERSALITY & TASK- INVARIANCE: The supplied evidence also analyzes the predictive distribution of dominant FCs across attention-score ranges.No separate experiment-detail passage is included in the supplied input for section B.

B.1 EXPERIMENT CONFIGURATIONS. … C ADDITIONAL EXPERIMENTAL RESULTS

The paper evaluates FASA with decode-only comparisons, broad long-context and reasoning benchmarks, task-specific metrics, and a FlashAttention2-compatible implementation. Its experimental setup uses calibrated dominant frequency chunks, standard benchmark protocols, and sparse decoding integration.

  • B.1 EXPERIMENT CONFIGURATIONS.: All methods omit KV-cache optimization during prefilling to isolate decode-stage acceleration under direct, fair comparisons.Baseline configurations follow original-paper standards or fair, strong comparison setups.
  • B.1 EXPERIMENT CONFIGURATIONS.: FASA identifies dominant FC indices through a one-time, task-agnostic offline calibration using a single Qasper sample for LongBench.The passage describes this minimal calibration as robust because the generated response provides sufficient identification signal.
  • B.2 BENCHMARK DETAILS: LongBench averages performance across diverse long-context tasks, including question answering, summarization, few-shot learning, synthetic tasks, and code completion.The benchmark covers both single-document and multi-document question answering.
  • B.2 BENCHMARK DETAILS: The evaluation spans mathematical reasoning with MATH500 and AIME, plus long-sequence modeling on C4, PG19, and WikiText.PG19 examples are full books for testing very long dependencies, while C4 is a cleaned general-domain corpus and WikiText contains formatted Wikipedia articles.
  • B.3 EVALUATION PROTOCOLS: LongBench reports f1 score for question answering, rouge_score for summarization, and code_sim_score for code completion, averaged across tasks.These metrics follow the official LongBench evaluation protocol.
  • B.3 EVALUATION PROTOCOLS: Long-sequence modeling uses perplexity (PPL), while long CoT reasoning on MATH500 and AIME2024 uses pass@1 in extended autoregressive generation.Lower PPL indicates better prediction; AIME2024 determines pass@1 from k = 16 independent generations.
  • B.4 IMPLEMENT DETAILS: FASA is implemented with HuggingFace Transformers by monkey-patching FlashAttention2’s forward pass and storing pre-computed dominant-FC indices globally.Figure 14 describes a two-stage, FlashAttention-compatible pipeline whose FAC stage enables sparse computation through the standard FlashAttention API.

C.1 PERFORMANCE ANALYSIS ON DIFFERENT BUDGETS · D DISCUSSION ON FASA · D.1 VARIANTS OF FASA

FASA outperforms the query-magnitude heuristic used by SparQ under constrained budgets while reducing per-token inference overhead through offline calibration. Its memory-optimized variant, FASA-M, minimizes GPU KV-cache usage by offloading non-dominant keys and values to CPU memory and transferring only required subsets just in time.

  • C.1 PERFORMANCE ANALYSIS ON DIFFERENT BUDGETS: FASA’s budget sensitivity is evaluated on Qwen2.5-7B-Instruct across various token budgets with Ntip = 16.Figure 15 reports this budget analysis.
  • C.1 PERFORMANCE ANALYSIS ON DIFFERENT BUDGETS: FASA’s budget sensitivity is also evaluated on Meta-3.1-Llama-8B-Instruct across various token budgets with Ntip = 16.Figure 16 reports this model-specific budget analysis.
  • C.1 PERFORMANCE ANALYSIS ON DIFFERENT BUDGETS: Under a constrained budget of 256 tokens, SparQ’s performance collapses because query magnitudes cannot reliably identify critical tokens.This comparison is reported on LongBench.
  • D DISCUSSION ON FASA: FASA uses one-time offline calibration, giving it substantially lower per-token inference cost than SparQ’s query-wise dimension re-evaluation.SparQ must re-evaluate high-magnitude dimensions for every new query, whereas FASA avoids that repeated overhead.
  • D.1 VARIANTS OF FASA: FASA-M retains only dominant Key-cache parts on the GPU for initial token-importance prediction, while offloading non-dominant keys and all values to CPU memory.This design targets constrained-GPU-memory scenarios such as consumer-grade hardware.
  • D.1 VARIANTS OF FASA: During focused attention, FASA-M transfers only the required non-dominant key and value subsets for the identified critical tokens, using just-in-time data movement.The GPU therefore remains occupied primarily by critical components.
  • D.1 VARIANTS OF FASA: FASA-M can approach an 8× memory reduction when dominant FCs occupy 25% of d and the token budget b is 10% of L.The comparison is against a full KV cache occupying Nlayers × L × 2d × bytes_per_param.

D.2 DESIGN CHOICES • · D.3 ALGORITHM ON FASA · E LLM USAGE

FASA’s design treats FC-scores as selectors rather than attention substitutes and preserves frequency chunks as indivisible units. Its workflow calibrates dominant FCs offline, predicts token importance from them during inference, computes focused full attention on selected tokens, and uses CPU/GPU cache management; ChatGPT use was limited to language refinement.

  • D.2 DESIGN CHOICES •: FC-scores accurately rank token importance but cannot replace attention probabilities, whose direct substitution causes catastrophic performance degradation.Their validated role is selecting salient tokens rather than approximating the final attention distribution.
  • D.2 DESIGN CHOICES •: Selecting individual dimensions fails catastrophically, so RoPE optimization must treat each Frequency Chunk as an indivisible functional unit.The design respects the inherent coupling of dimension pairs.
  • D.3 ALGORITHM ON FASA: Algorithm 1 calibrates dominant FCs from dataset Ω by collecting contextual-agreement scores across examples and token-generation steps.It computes full attention scores, calculates each FC’s CA score using Eq. 4, and stores the results by layer, head, and FC index.
  • D.3 ALGORITHM ON FASA: Algorithm 1 averages stored CA scores for each layer-head-FC triplet and returns the top-k FC indices as Idom.The output is the set of dominant FC indices selected by mean CA score.
  • E LLM USAGE: During manuscript preparation, ChatGPT was used only for grammar correction, stylistic enhancement, and rephrasing for clarity.The authors state that the scientific concepts, data analyses, and conclusions were original work without substantive language-model contribution.
  • D.3 ALGORITHM ON FASA: FASA-M first splits keys by dominant FCs and predicts token importance, then selects dominant key parts for focused attention computation.The inference algorithm takes the current query, key, value, dominant FC indices, token budget, and past CPU KV caches as inputs.
  • D.3 ALGORITHM ON FASA: FASA-M updates non-dominant key and value caches on the CPU, transfers required parts to the GPU, reconstructs full keys for selected tokens, and computes full attention on that subset.The algorithm returns the next hidden state and updated caches.
Loading 2602.03152v3…