ATFlash: Per-RoPE-Wavelength Attention Windows for Compute/Memory-Efficient LLM Inference

2026-08-03

AI Summary

This work addresses the computational and memory bottlenecks of large language models in long-context reasoning by introducing a novel, input-agnostic, closed-form logarithmic compression method for attention sparsification. Leveraging the frequency-domain wavelength properties of RoPE positional encoding, the approach defines distance windows based on wavelength and prunes query–key inner product terms that exceed a proportional threshold, while preserving full key reachability. The mechanism is compatible with FlashAttention-4 and FlashInfer and requires no modification to online softmax recurrence. Experiments on Qwen2.5 and Llama-3.2 demonstrate that pruning 37–48% of inner products yields Top-1 match rates of 96–98% and maintains output KL divergence at the 10⁻³ level. End-to-end speedups of 1.29× and 1.31× are achieved on 128K and 1M context lengths, respectively.

Abstract

The attention score with rotary position embeddings (RoPE) decomposes exactly into a sum over its 2D-rotation frequency pairs, and each pair's wavelength limits how far it can discriminate position. Aligned with this structure, we propose the per-RoPE-wavelength distance window: it prunes the query--key inner-product terms beyond a wavelength-proportional distance. Unlike a sliding window, every key remains reachable, at least through the low-frequency pairs. The reduction rate is input-independent, with a closed form logarithmic in the sequence length $N$, in contrast to dynamic-sparse methods like MInference. Such token-level selection is orthogonal to our frequency-level pruning. The window can therefore be applied on top of those methods. On Qwen2.5-0.5B and Llama-3.2-3B, the window prunes 37--48\% of the query--key inner-product terms within each model's native context length. Relative to full attention, the top-1 match rate stays at 96--98\% and the mean output-distribution KL at the $10^{-3}$-nat level on LongBench-v2 contexts. We examine absolute scores on long-context benchmarks such as RULER, OpenAI-MRCR, LongCodeQA, and $\infty$Bench: they are broadly preserved. We implement the window as a slice of the query--key contraction axis, leaving the online-softmax recurrences untouched, and port it with minimal diffs into the released FlashAttention-4 prefill and FlashInfer decode. On RTX PRO 6000 with Llama, both ports outpace stock with gains growing with context length, up to $1.29\times$ at 128K. End to end on Qwen2.5-7B-1M, with 57\% of the inner-product terms pruned, the speedup reaches $1.31\times$ at a 1M-token context.

Problem

• efficient attention
• long-context inference
• RoPE
• compute efficiency
• memory efficiency

Innovation

• RoPE
• attention pruning
• wavelength-aware window
• compute-efficient LLM inference
• FlashAttention