通过频率波长设计注意力窗口，显著减少计算量且保持模型精度。
ATFlash: Per-RoPE-Wavelength Attention Windows for Compute/Memory-Efficient LLM Inference
- 按旋转位置编码的波长动态裁剪注意力范围，实现高效计算。
- 在长文本上保留96%以上匹配率，输出分布差异小于10^{-3}纳特。
- 兼容现有优化方法，部署简单，长序列推理速度最高提升2.9倍。
旋转位置嵌入（RoPE）的注意力分数可精确分解为二维旋转频率对之和，每个频率对的波长决定了其能区分位置的范围。基于此结构，我们提出每波长距离窗口：按波长比例裁剪超出距离的查询-键内积项。不同于滑动窗口，所有键仍可通过低频对可达。该方法减少率与输入无关，闭式表达对数级依赖于序列长度 $N$，区别于MInference等动态稀疏方法。该词元级选择与频率级裁剪正交，可叠加于其他方法之上。在Qwen2.5-0.5B和Llama-3.2-3B上，窗口在模型原生上下文长度内裁剪37–48%的查询-键内积项。相比全注意力，长文本基准（LongBench-v2）上顶级匹配率维持在96–98%，平均输出分布KL散度处于$10^{-3}$-nat水平。在RULER、OpenAI-MRCR、LongCodeQA和$ fty$Bench等长上下文基准上，绝对得分基本保持。我们将窗口实现为查询-键收缩轴的切片，保持在线softmax递推不变，并以极小代码改动集成至FlashAttention-4预填充与FlashInfer解码中。在RTX PRO 6000上，对Llama的两个版本均优于原版，加速比随上下文增长，128K时达1.29×。端到端在Qwen2.5-7B-1M上，裁剪57%内积项时，1M-token上下文加速比达1.31×。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。