arXiv:2506.02572cs.LGcs.AI2025-06ACL被引 6

用哈希编码加速大模型注意力,速度提升7.2倍且不丢精度。

HATA: Trainable and Hardware-Efficient Hash-Aware Top-k Attention for Scalable Large Model Inference

  • 用二进制哈希码替代传统得分计算,低成本获取注意力顺序。
  • 在多个主流模型上实现最高7.2倍加速,精度优于现有方法。
  • 适合需要高效推理的大模型部署场景,开源可用。

大型语言模型(LLMs)虽已取得显著进展,但注意力模块仍是推理阶段的关键瓶颈,即便采用KVCache等技术也难以完全缓解冗余计算问题。尽管已有多种top-k注意力机制通过利用注意力固有的稀疏性来加速推理,但往往难以在效率与精度间取得平衡。本文提出HATA(Hash-Aware Top-k Attention),首次将轻量级学习哈希技术系统性地融入top-k注意力过程。不同于传统方法需精确估计qk得分(成本高昂),HATA将查询与键映射为二进制哈希码,以极低开销获得相对qk得分排序,足以支持top-k注意力。大量实验表明,相较于原生全注意力,HATA可实现高达7.2倍的加速,同时保持模型精度;在多个主流LLM模型和多样化任务中,其精度与效率均超越当前最优top-k注意力方法。代码已开源:https://github.com/gpzlx1/HATA。

原文摘要 · Abstract (English)

Large Language Models (LLMs) have emerged as a pivotal research area, yet the attention module remains a critical bottleneck in LLM inference, even with techniques like KVCache to mitigate redundant computations. While various top-$k$ attention mechanisms have been proposed to accelerate LLM inference by exploiting the inherent sparsity of attention, they often struggled to strike a balance between efficiency and accuracy. In this paper, we introduce HATA (Hash-Aware Top-$k$ Attention), a novel approach that systematically integrates low-overhead learning-to-hash techniques into the Top-$k$ attention process. Different from the existing top-k attention methods which are devoted to seeking an absolute estimation of qk score, typically with a great cost, HATA maps queries and keys into binary hash codes, and acquires the relative qk score order with a quite low cost, which is sufficient for realizing top-k attention. Extensive experiments demonstrate that HATA achieves up to 7.2$\times$ speedup compared to vanilla full attention while maintaining model accuracy. In addition, HATA outperforms the state-of-the-art top-$k$ attention methods in both accuracy and efficiency across multiple mainstream LLM models and diverse tasks. HATA is open source at https://github.com/gpzlx1/HATA.

注意力机制大模型推理哈希编码高效计算

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