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

2025-06-03

AI Summary

In LLM inference, attention computation remains a critical bottleneck; existing Top-k sparsification methods struggle to balance efficiency and accuracy. This paper proposes a hash-aware trainable Top-k attention mechanism: for the first time, it integrates a lightweight learnable hash function into the Top-k selection process, using binary hash codes to efficiently approximate the relative similarity ordering between queries and keys—thereby avoiding costly absolute attention score computations. Concurrently, the method jointly optimizes the KV cache. Evaluated across multiple mainstream LLMs and downstream tasks, it achieves zero accuracy loss relative to the dense baseline while accelerating inference by up to 7.2×—outperforming all prior Top-k approaches. The approach demonstrates strong generalization and effectiveness without architectural modification or fine-tuning.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search
• Computer Vision: Large Vision Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

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

Problem

• Addresses attention module bottleneck in LLM inference
• Balances efficiency and accuracy in top-k attention
• Integrates low-cost learning-to-hash for relative qk scoring

Innovation

• Integrates learning-to-hash into Top-k attention
• Uses binary hash codes for low-cost scoring
• Achieves speedup while maintaining model accuracy