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

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 \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.

## introduction recently, kvcache has become a paradigm for the inference of large language model (llm) @xcite , due to its benefit of mitigating redundant computation in the decoding stage. in this situation, massive kv states loading becomes the bottleneck, especially for long sequences and large batch sizes @xcite . top-k attention @xcite has emerged as a promising approach to accelerate llm inference by leveraging the inherent sparsity in attention. by selectively retaining only the top-k most relevant tokens in the kvcache, top-k attention significantly reduces the kvcache loading overhead. however, existing top-k attention algorithms face notable challenges in achieving an optimal trade-off between efficiency and accuracy. low-rank methods, such as loki @xcite and infinigen @xcite , reduce overhead by computing dot-products over a subset of projected dimensions, but they introduce significant computational costs due to the extensive requirements for channel extraction. on the other hand, block-wise methods like quest @xcite and infllm @xcite improve efficiency by grouping contiguous key-value pairs into blocks, but they often compromise accuracy as critical keys may be excluded based on their coarse-grained estimation of query-key (qk) scores. in this paper, we introduce hash-aware top-k attention (hata), a novel approach that systematically integrates low-overhead learning-to-hash techniques into the top-k attention process. unlike existing methods that focus on precise numerical estimation of qk scores, hata maps queries and keys into binary hash codes, acquiring the relative qk score order with minimal computational cost. this approach eliminates costly high-fidelity score approximations, enabling significant speedup while preserving the quality of top-k selection. as illustrated in figure 1 , hata shows superiority in balancing the accuracy and efficiency, compared to state-of-the-art methods. hata leverages the success of learning-tohash @xcite , which has been widely used in similarity-based retrieval tasks such as image search and machine learning. by training hash functions based on the query-key pairs of llm attention, hata is able to encode any query and key vector into a binary code, which further enable hata to achieve low-overhead but precise token selection, making it a hardware-efficient solution for accelerating llm inference. extensive experiments demonstrate that hata achieves up to 7.2× speedup compared to vanilla full attention while maintaining model accuracy. furthermore, hata outperforms state-of-the-art top-k attention methods in both accuracy and efficiency across multiple mainstream llm models and diverse tasks. in summary, our contributions are as follows: ## hata's design to address the aforementioned three challenges, we propose hash-aware top-k attention (hata), a trainable and hardware-efficient approach based on learning-to-hash. in sec 3.1, we formally define the query-keybased learning-to-hash problem and design a training loss function to learn hash codes while preserving similarity. we also incorporate bits balance and uncorrelation constraints @xcite to enhance hash bit quality. in sec 3.2, we introduce hata's workflow, leveraging the learned hash function to significantly accelerate llm inference. ## hardware-efficient optimizations hata is implemented in pytorch @xcite and flashinfer @xcite , comprising 1,470 lines of c++/cuda code (for custom gpu kernels) and 940 lines of python code (for highlevel orchestration). to bridge the gap between theoretical efficiency and practical performance, we introduce three hardware-efficient optimizations, as illustrated in figure 3 , targeting compute and memory bottlenecks in attention with long contexts and large batches. notably, while hata employs extensive low-level optimizations, it remains pluggable and integrates seamlessly with existing inference frameworks. to leverage hata, users need only replace standard attention with hata's attention. kernel fusion for hash encoding. the encode & cache update phase involves a chain of operations such as linear projection, sign function, bitpack, and cache updates. although each operation takes only a few microseconds on the gpu, the cpu requires tens of microseconds to dispatch them, starving gpu compute units. by fusing these into a single cuda kernel, we significantly reduce cpu-gpu synchronization, consequently cutting end-toend inference latency. ## empirical evaluation in this section, we evaluate hata's performance in terms of both accuracy and efficiency. ## conclusion we introduced hash-aware top-k attention (hata), a hardware-efficient method for faster llm inference. hata offers a systematic exploration and validation of the integration of learningto-hash techniques into top-k attention mechanisms, achieving up to 7.2× speedup over dense attention and outperforming sota methods in accuracy and performance, establishing it as an effective solution for llm inference acceleration.