HashAttention: Semantic Sparsity for Faster Inference

We thank the reviewer for supporting our paper. Please find responses to the questions and comments below. Kindly let us know if you have any additional follow-up questions.

1. Hyperparameters for HashAttention Training / Frozen backbone LLM

   Yes the LLM weights are kept frozen during training only hash attention mappings are trained locally and independently for each attention head. Our training for 32K length HashAttention on openwebtext looks like

2. performance overhead of computing the Hamming distance vs. full computation

Kindly refer to the response 2 (Reviewer HrF1) for the table. We make the following observations,

A. (Sequence dimension) Improvement with using hamming distances over inner products stays consistent with increasing tokens in KV Cache as expected.

B. (Bitwidth dimension) We see that we can go upto as large as 512 bitwidth signatures while maintaining advantageous latency for Hamming distance computation. In our experiments, we use 32 bits for HashAttention.

With such an efficient distance computation, the Hamming distance computation is far from being the bottleneck of HashAttention.

3. different sparsity levels for different layers / heads

We have not explored this dimension. However, we expect such improvements to orthogonally improve all the methods. We leave this for future work.

4. HashAttention performance with quantized KV cache implementations

Quantization of KV Caches is a way to overall improve the memory footprint of the Cache. Since, in such a setup, HashAttention still would act upon full precision vectors (since we will have access to those while computing signatures), the topk accuracy/sparsity tradeoff would remain the same. Quantization of KV Caches adds another layer of approximation to attention computation. We believe it affects all the sparse attention methods in the same manner.

The absolute memory footprint of HashAttention remains the same. Auxiliary memory would, of course be higher in relative terms. However, it is true for all sparse attention baselines.

Other discussion:

1. Code: The anonymous repository for our code used for running long benchmarks is here: https://anonymous.4open.science/r/HashAttention_ICML2025/README.md We plan to release full code on acceptance

2. Choosing #bits. In our experiments, we use 32 bits since it works reasonably. As expected, we find that using larger bits helps with the quality of top-k prediction.

We will add more results and discussion on these hyperparameters in final version.

3. overhead of signatures for < 8K contexts, Does overhead become more prominent with quantized KV

We will highlight this overhead in the limitation section of the paper for clarity. We do not believe the overhead increases with quantized KV -- firstly, nothing changes in the signature computation even with quantized KV since we have access to full KV when signatures are created. Secondly, using quantized KV adds dequantize step in the attention computation, which would reduce the relative overhead of signature computation.

4. Formulation: we will improve the formulation in the paper.

We thank the reviewer for supporting our paper. Please find responses to the questions and comments below. Kindly let us know if you have any additional follow up questions.

1. training-free hashing To compare HashAttention to training free methods, we measure quality of sparse attention while using LSH signatures (random signed projections) in place of learned mapping-based HashAttention signatures. As expected, the data-agnostic LSH needs longer signatures to start giving reasonable quality.

2. Token level hashing / gathering in sparse attention / performance efficiency.

Our implementation does not involve gathering pivotal tokens into a contiguous memory space. The implementation is built upon the vLLM page attention framework, where each token corresponds exactly to one page (i.e., page size = 1). After identifying top-k tokens, we use their indices without physically moving or gathering these tokens into a separate memory buffer. The page attention kernel explicitly utilizes these indices to selectively compute attention only for the specified tokens, efficiently ignoring irrelevant tokens. The page attention kernel is highly optimized for GPU memory access patterns. Due to the GPU cache line size of 128 bytes, optimal memory bandwidth utilization is achieved as long as contiguous data access meets or exceeds this cache line size. Each token's head representation has 128 fp16 elements, equivalent to 256 bytes. This naturally exceeds the GPU cache line size, allowing our attention kernel to leverage GPU memory bandwidth effectively.

Many state-of-the art inference framework implement attention using paged-attention backbone with page size 1. They find that with correct optimization, the efficiency does not depend on page size. Quoting from official flashinfer documentation “Some recent work such as LightLLM and sglang uses a special form of PageAttention where page size equals one, for easy management of KV-Cache in complicated serving scenarios such as structured generation. FlashInfer optimizes PageAttention kernels by pre-fetching page indices in GPU shared memory, so that kernel performance is not affected by the page size.” We will explicitly clarify this and use better naming for the different stages in the paper to avoid potential confusion.

3. OOD test sets / generalization of learned mappings across tasks.

We show two sets of results in our experiments. The HashAttention is trained on completely unrelated data from openwebtext dataset and tested on OOD LongBench and RULER. The good performance of HashAttention on these datasets is a testament to its strong generalization performance. Additionally, finetuning for specific benchmark, which gives us HashAttention*, further improves the quality of results.

4. OOD Query Kindly refer to Response 4 to reviewer HrF1.

5. MagicPig vs. HashAttention:

MagicPig can be understood and compared with HashAttention in two respects. Firstly identifying important tokens. HashAttention uses succinct bit signatures (32 bits) obtained via learned mappings to compute important tokens. MagicPIG uses data-agnostic LSH to obtain bit signatures. As expected MagicPIG needs much longer bit signatures to obtain reasonable results ( from their paper 1500-2200 bits). MagicPig builds index on this signatures and due to irregular bucket sizes, these indices need to be stored on CPUs -- a reasonable solution in case of KVCache being stored on CPUs.

Secondly, MagicPig proposes sampling based estimation instead of top-k estimation. Sampling is performed using LSH tables. The idea sampling based estimation is an interesting direction to be explored w.r.t its application to other sparse attention methods, including HashAttention.

We plan to thoroughly compare HashAttention with many concurrent works such as MagicPig, PQCache, Squeeze Attention in our future work. Comparison against these is out of the scope of the rebuttal.

6. Lemma 4.2 : We would be happy to cite the correct reference in prior literature. Kindly direct us to the same.

Please let us know if there are any additional queries.

We thank the reviewer for supporting our paper. Please find responses to the questions and comments below. Kindly let us know if you have any additional follow up questions.

1. Discussion on RAG – We will remove this from related work.

2. Dimension of hamming codes:

We can answer this question by comparing (1) the latency of inner product computation ( dim=128 which is the standard dimension in most models) and (2) the hamming distance computation based on our kernel implementation.

We can see that up to 512 bits, the hamming distance is cheaper than the inner product computation. In practice, we do not need such large bit widths. This excludes the cost of running MLP on a query vector, which does not scale with #tokens.

3. unbalanced buckets issue in LSH

Since HashAttention compares the query signature with all token signatures, it is not directly affected by the unbalanced bucket issue. However, it can affect top-k selection when many tokens share the same Hamming distance. In such cases, we randomly chose $k$ tokens in our implementation.

4. query OOD issue The mappings learned in hashattention transform query and key-value into a low-dimensional semantic space where the smaller the hamming distance, better is the relevance of key-value to the query. While training the mappings, this hamming distance is used to classify the key-value to be relevant or irrelevant to the query. The setup naturally promotes transformed query distribution to be closer to key-value distribution, as shown in the table below. We can see that average cosine similarity significantly improves after the transformation. We provide results for both soft embeddings and hard embeddings for HashAttention.

5. K-means based topk vs. vanilla LSH vs. HashAttention:

HashAttention uses learned mappings to obtain bit-signatures for retrieval (while motivated by LSH, it does not perform LSH). Due to rising interest in this topic, many concurrent works have explored different approaches from information retrieval. These include MagicPig (which uses LSH tables), SqueezeAttention and PQCache (which use clustering), and RetrievalAttention (which employs graph-based retrieval). We plan to thoroughly compare HashAttention with these works in future research. However, such a comparison is beyond the scope of this rebuttal

6. Extra fine-tuning details

To fine-tune HashAttention mappings for downstream LongBench, we use 25 samples from all LongBench task each to create a fine-tuning dataset. Then we further train hashAttention on this dataset. The evaluation is performed on the LongBench excluding samples included in the fine-tuning dataset.

We thank the reviewer for supporting our paper. We are working on extending Hash Attention to scenarios involving KV cache offloading and reasoning tasks, and we will present these in our future work. Please let us know if you have any other questions; we would be happy to clarify.