arXiv:2506.05300cs.LG2025-06

用动态阈值筛选注意力,减少显存搬运,提升大模型推理效率。

Power Law Guided Dynamic Sifting for Efficient Attention

  • 基于幂律规律动态设阈值,替代耗时的top-k筛选。
  • 在保持模型质量前提下,降低显存带宽使用量。
  • 特别适合需高效推理的大语言模型部署场景。

大型语言模型在GPU上的高效推理仍受制于内存带宽瓶颈,尤其在注意力计算中高带宽内存(HBM)与SRAM间的数据传输。近似注意力方法虽能降低计算和内存开销,但常依赖昂贵的top-$k$操作,在GPU上表现不佳。本文提出SiftAttention,以基于阈值的逐元素过滤操作替代top-$k$步骤。我们发现注意力分数的τ-分位数在生成序列中遵循可预测的幂律分布,据此在每步生成时动态估算阈值,仅加载并使用高于阈值的注意力分数及其对应值向量,从而减少HBM与SRAM间的数据移动。实验表明,相比现有近似注意力方法,SiftAttention在保持模型质量的同时显著降低值向量加载的内存带宽消耗。

原文摘要 · Abstract (English)

Efficient inference on GPUs using large language models remains challenging due to memory bandwidth limitations, particularly during data transfers between High Bandwidth Memory (HBM) and SRAM in attention computations. Approximate attention methods address this issue by reducing computational and memory overhead but often rely on expensive top-$k$ operations, which perform poorly on GPUs. We propose SiftAttention, a novel approximate attention method that replaces the top-$k$ step with a computationally efficient element-wise filtering operation based on a threshold value. Our intuition for doing this is based on our empirical observation that the $τ$-th quantile of attention scores follows a predictable power-law over sequential generation steps. Exploiting this insight, our approach dynamically estimates a threshold value per prompt at each generation step. Only attention scores above this threshold and their corresponding value vectors are loaded/used to compute the attention output, reducing data movement between HBM and SRAM. Our evaluation demonstrates that SiftAttention preserves model quality better than existing approximate attention methods while reducing memory bandwidth usage when loading value vectors.

注意力机制大模型推理内存优化幂律

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。