HashAttention: Semantic Sparsity for Faster Inference
Authors: Aditya Desai
Shuo Yang
Alejandro Cuadron
Matei Zaharia
Joseph E. Gonzalez
Ion Stoica
Venue: ICML 2025 poster (2025)
Keywords: hashing
learning to hash
attention
sparsity
time to next token
decoding
generative LLM
Primary area: deep_learning->attention_mechanisms
TL;DR: Sparse attention for faster inference based on learned bit signatures

Abstract: Leveraging long contexts is crucial for advanced AI systems, but attention computation poses a scalability challenge. While scaled dot-product attention (SDPA) exhibits token sparsity, i.e. only a few pivotal tokens significantly contribute to output, exploiting this sparsity remains challenging. Existing methods either suffer from quality degradation or require substantial additional resources. We show that identifying pivotal tokens is a Maximum Inner Product Search (MIPS) problem. However, existing MIPS solutions are not well-suited for SDPA, as they are not GPU-friendly and often underperform due to the separated query and key distributions. This paper introduces HashAttention, framing pivotal token identification as a recommendation problem. Given a query, HashAttention encodes keys and queries in Hamming space, capturing the required semantic similarity, using learned mapping functions. HashAttention efficiently identifies pivotal tokens for a given query using bitwise operations and computes attention using only these tokens, improving the overall attention efficiency. Trained on generic data, HashAttention reduces tokens used by up to $16\times$ with minimal quality loss, requiring only 32 bits of auxiliary memory per token. Sparsity can be further improved to $32\times$ through task-specific fine-tuning. 
On A100 GPU, at $32\times$ sparsity, incorporating HashAttention reduces attention latency by up to $4.3\times$ in GPT-FAST and $2.54\times$ in FlashDecode, and achieves up to $3.12\times$ higher throughput for GPT-FAST.

MlshhAttnntion
Semantic Sparsity for Faster Inference
Aditya Desai
Shuo Yang' Alejandro Cuadron
Matei Zaharia
Joseph E. Gonzalez' lon Stoica
Abstract
Hastt PTPA)
Double Sparsity (32 bits PTPA)
Double Sparsity (6-tbits PTPA) * FULL
Leveraging long contexts is crucial for advanced Al systems, but attention computation poses a scalability challenge. While scaled dot-product attention (SDPA) exhibits token sparsity, i.e. only a few pivotal tokens significantly contribute to output, exploiting this sparsity remains challeng- ing. Existing methods either suffer from qual- ity degradation or require substantial additional resources. We show that identifying pivotal to- kens is a Maximum Inner Product Search (MIPS) problem. However, existing MIPS solutions are not well-suited for SDPA, as they are not GPU- friendly and often underperform due the sepa- rated query and key distributions. This paper intro- duces HashAttention, framing pivotal token iden- tification recommendation problem. Given a query, HashAttention encodes keys and queries in Hamming space, capturing the required seman- tic similarity, using Iearned mapping functions. HashAttention efficiently identifies pivotal tokens for given query using bitwise operations and computes attention using only these tokens, im- proving the overall attention efficiency. Trained on generic data, HashAttention reduces tokens used by up to 16x with minimal quality loss, requiring only 32 bits of auxiliary memory per token. Sparsity can be further improved to 32x through tasc-fpcc fine-tuning. On A100 GPU, at 32x sparsity, incorporating HashAttention duces attention latency by up to x.3x in GPT- FAST and 2.54x in FlashDecode, and achieves to 3.1 higher throughput for GPT-FAST.
82.56
84.28
Average Quality
80
80.04
75
70
70.29 *
Throughput
Figure 1. Quality vs. Throughput (on A100 GPU) Auxiliary memory at different sparsity rates for Double Spansity (16 channels, 2bit and 4 bit quantization) and HashAttention implemented in GPT-FAST (batch=1) framework on RULER 32K using Llama- 3.1-8B Instruct. The impact of auxiliary memory is excluded from throughput computation due to the relatively small scale of the problem (batch=1 and 32K context length). Also, through- put is overestimated for Double sparsity by excluding quantiza- tion/dequantization in throughput computation.
to engaging extended conversations (Touvron et al., 2023; Achiam et al., 2023; Liu al., 2024b). is commonplace for LLM models to preprocess and store huge amounts of text in the form of KV Cache, which is later used to process various prompts. However, the Scaled Dot Product Attention (SDPA), which is fundamental to the transformer architecture that has driven the Generative Al revolution (Vaswani. 2017: Brown et al., 2020), does not scale well with context length. Generating single token requires SDPA to access the entire context KV Cache, which be hundreds of GB. For instance, a KV Cache of 512K tokens is 64GB (0.125 MB per token) for the LLama3.1 model.
1. Introduction
Luckily, sparsity naturally arises in SDPA. Due to the soft- max kernel, only a few pivotal tokens significantly con- tribute the final attention computation(Bricken Pehle- van, 2021; Deng et al., 2024). Efficiently identifying these pivotal tokens provides pathway to achieving efficient at- tention. Various approaches to identifying pivotal tokens have been explored in the literature. Heuristic- ods, such as fixed sparsity patterns (iio et al., 2023), ignore the dynamic nature of contextual sparsity in attention, result- ing in suboptimal attention quality. KV cache discard strate-
The ability refer long contexts efficiently is erucial for modern Al applications, from processing lengthy documents
Department of Electrical Engineering and Computer Sciences, UC Berkeley DDpartment Computer Sciences, ETH Zurich. Correspondence to: Aditya Desai <apdesai @ bed
meth-
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

HashAttention
gies, such as those proposed in (Liu al., 2024c: Zhang al., 2023; Li et al., 2024a), identify the global importance of tokens per attention head, incorporating some degree of dynamic sparsity. These methods discard tokens based on their observed importance in context prefixes. However, to- ken importance is dynamic, and tokens deemed unimportant in prefixes can become critical for future inference, leading these methods to fail in certain scenarios (Xiao al., 2024: Tang et al., 2024). Some approaches to truly dynamic spar- sity have emerged(Xiao et al., 2024; Yang et al., 2024; Tang et al., 2024). However, firstly these methods rely on heuris- ties to reduce the computation of pivotal tokens, leading to low recall rates. Secondly, improving recall with these methods needs substantial auxiliary memory.
HashAttention, trained with generic dataset (i.e. indepen- dent of downstream tasks) can give up 16x reduction in KV cache usage while preserving the model quality. For instance, with 16x sparsity, HashAttention maintains model quality within a drop of 0.78 points on LongBench(Bai et al., 2024b) and 1.13 points on RULER@ 16K using Llama-3. 3- 8B-Instruct model. Our ablations show that sparsity can be further improved if HashAttention is finetuned on task- specific data. For instance, on some Longbench datasets with Llama-3. 1-8-..118--nntuue and Destral-7B-v.03--nnsue we can go up to 32x sparsity without loss of model quality,
Incorporating HashAttention in inference systems can improve the latency and throughput of inference. For instance, HashAttention-322: improves latency of GPT- FAST(GPTFast, 2024) attention upto 4.3x and that of FlashDecode(Dao, 2023) upto 2.54x. Additionally, end- to-end model throughput for GPT-FAST, tested with 1-ma3.1-8BB model, can be improved up to 3.12x with HashAttention-S22x Furthermore, HashAttention only uses 32 bits PTPA auxiliary memory. To achieve the same quality, other sparse-attentions either need around 4x more pivotal tokens (i.e. 4x less sparsity) 32 PTPA or need more aux- iliary memory to operate the same sparsity. The code for HashAttention is here. 1
In this paper, we take a principled approach to identify- ing pivotal tokens. Given query and a large set of key- value pairs, we frame the task of retrieving important to- kens as a recommendation problem (Zhang et al., 2021). We show that identifying pivotal tokens under the indepen- dence of value vectors is exactly a Maximum Inner Product Search(MIPS) problem. Existing solutions to MIPS are unsuitable for pivotal token identification due to two rea- sons (1) SOTA solutions MIPS are based on graph data structures which are GPU unfriendly.
(2) Query and Key distributions are separated in Attention causing off-the-shelf solutions to fail (Liu et al., 2024a). We observe that MIPS can be reduced to a cosine-similarity search (Neyshabur & Srebro, 2015; Shrivastava & Li, 2014) which can be further approximated using bit signatures (Gionis al., 1999), or embeddings in Hamming space.
2. Related Work
2.1. Long Context and Retrieval-Augmented Generation
Recently, there has been rising interest in developing models capable of processing long Contexts(Beltagy et al., 2020; Guo et al., 2022: Han et al., 2023; An et al., 2024; Ma et al., 2024), driven by emerging applications such as multi- document question answering, multi-turn chatbots, and tutor bots (Bai et al., 2024a; Feng et al., 2022; Davies & Murff, 2024). Two main approaches to addressing this challenge are: (1) designing models that natively support long contexts using methods like ROPE scaling (Su et al., 2023), and (2) employing a retriever to extract relevant portions of text to augment the prompt, technique known Retrieval- Augmented Generation (RAG) (Lewis et al., 2020).
HashAttention uses learned independent mappings to embed the key-value pairs and queries in Hamming space. These mappings are trained on the actual key-value and query data from the LLM model to encode inner product-based similarity into bit signatures. Given a query, the tokens with key-value signatures closest to the query signature in Hamming space are identified pivotal tokens. These signatures can be packed in integers enabling their efficient storage and processing. HashAttention is GPU friendly and resolves the issue of separated distribution of queries and keys using learned independent mappings.
A recent study highlights the advantages of building long- context L.MMs over RAG approaches (Li et al., 2024b). No- table progress in this area includes the emergence of long- context models, such as the open-source Llama 3.1 (Dubey et al., 2024) series, which supports a context length of 128K tokens, and proprietary models like GPT-4 Turbo (Achiam et al., 2023), Claude-2, and Gemini 1.5 (Team et al., 2023), offering context lengths of up to 128K, 200K, and 1M to- kens, respectively.
An instance of throughput improvement for GPT-FAST, on RULER @ 32K benchmark with 18B is shown in Figure 1. At 32K context length, HashAttention improves the throughput 1.3x (4x sparsity) at quality drop of 0.8 points and 1.7x (8x sparsity) with quality drop of 1.7 points using auxiliary memory of 32 bits per token per attention head (PTPA). At auxiliary memory of even 64 bits PTPA, Double Sparsity(DS), a SOTA sparse attention, fails preserve model quality and improve throughput. Broader experiments reveal more potential for sparsity and throughput improvements.
http:///whhbboooom///ggggaaaaaad 1.0
2

---

HashAttention
2.2. Post-training Sparse Attention
from key-value and queries to identify pivotal tokens.
We classify the previously proposed approaches to sparsify the attention computation of pre-trained models as follows:
Estimating attention output using Hash tables Learning-to-hass attention (Sun et al., 2021) proposes to learn hash table mappings to obtain efficient attention with sparse retrieval. They design a learning algorithm to obtain a balanced distribution of keys in the hash table buckets. MagicPig(Chen et al., 2024) noted that top-k attention biased and can be problematic in cases when attention scores are relatively balanced. They propose importance sampling- based unbiased attention estimation. They use LSH tables to approximate the importance-samplin distribution. These methods can be potentially combined with HashAttention for further efficiency.
Fixed Sparsity: Approaches such StreamingLLMM (Xiao et al., 2023) adopt fixed sparsity pattern to approxi- mate attention based on the observed importance of attention sinks (e.g., the first few tokens) and recent tokens. However, subsequent studies (Zhang et al., 2023; Xiao et al., 2024) have demonstrated the dynamic nature of sparsity and the resulting quality degradation when using fixed sparsity.
Methods such as H2O(Zhang al., 2023), ScissorHands(Liu et al., 2024c), Fas et al., 2023), and SnapKV(Li et al., 2024a) discard tokens based on specific heuristics. However, once discarded, these tokens are no longer available for future generations. This limitation is particularly problematic in scenarios such as multi-turn chat or multiple QA sessions on a fixed document, where it is essential to access different parts of the document in each turn or question.
KV Cache Discarding:
3. Background
3.1. Scaled Dot Product Attention (SDPA)
The computation of SDPA for key and value embeddings, K, V n1 x d, and Q : n2 d can be written as,
QK vd
SDPA(K, V. Q) softmax
V
Estimating top-k attention scores via partial computa- tion Attention scores one of the critical components that determine the contribution of a token to the output of the attention mechanism. Identifying the top-k tokens with
n2 = 1, then it can be simplified into,
the highest attention scores requires ((nd) computation, as it involves the dot product of the query with each token in the KV cache, where n is the number of tokens in the KV cache and d is the dimensionality.
SDPA(K, V Σ (a)V E=1
where (1
exp [K] .g) exp (K
called attention scores.
Double Sparsity (Yang al., 2024) reduces this compu- tational cost by selecting fewer channels to estimate dot products, which are then used to identify the top-k tokens. The channels are chosen based on offline calibration of channel norms. Loki (Singhania et al., 2024) reduces the di- mension by using traditional dimensionality reduction tech- niques to obtain low-dimensional representation for keys. InfLLM (Xiao et al., 2024) and Quest (Tang al., 2024) reduce computation along the n-dimension by creating page- level representations. These methods include or exclude all tokens on a page at once. While these approaches are ef- fective, they often fail to provide high recall for the top-k tokens with respect to attention scores.
3.2. Learning Recommendation Model
The recommendation problem can be abstracted as: Given set of items I and user E 24 the recommendation aims to select a small subset of I that is relevant to u. The histor- ical relevance of items for users is captured in an interaction matrix with shape [24] x [Z]. Traditionally, the learning of recommendation model has been cast as a matrix factor- ization of the interaction matrix. The two matrices, thus obtained, provide with embeddings of items and users that can be used for inference. Subsequently, auxiliary infor- mation such as description, profile, interaction history, etc., and deep learning models were used to obtain richer user and item representations. Alternatively, one can learn the embeddings so that the embeddings of the relevant items lie close the user embeddings terms of lp distance. To im- prove the efficiency of relevant item retrieval, approximate near-neigg neirhne algorithms and data structures utilized on top of the embedding space. Sparsifying attention is recommendation problem with key-value pairs as items and queries as a user. We want to select key-value pairs that minimize the error of output embedding w.r.t full attention for a particular query.
Retrieval Algorithms for Top-k Recently, RetrievalAt- tention (Liu al., 2024a) proposed using a graph-based nearest neighbor algorithm identify the top-k tokens with the maximum inner product. RetrievalAttention offloads the top-k computation to the CPU due to the sparse compu- tations involved with graphs, leading to additional latency. SqueezedAttenntion (Hooper et al., 2024), a concurrent work with ours, proposed to solve the top-k problem efficiently by clustering keys. In contrast, we use succinct bit signatures
3

---

HashAttention
Hamming space of signatures
key value signature
KC V
Indices of tokens with signatures closest query signature
Φ kv
1011
1001
1000
24,
56,
911.
4
2501
q
query signature
Figure 2. HashAttention working: The key-value pairs are mapped to bit signatures via learned mapping function Okv. The query is mapped bit signature function Qq. The tokens closest to the query signature are chosen as candidates for attention computation.
4. HashAttention
The SCORE function has three important aspects.
The general recipe of sparse attention can be viewed as combination of three subroutines (1) SCORE (2) TOPK and (3) GATHER-ATT.
1. Auxiliary memory: Various methods such as Double Sparsity, Quest, and InfLLM including HashAttention use meta representations of keys/tokens that cached for efficient score computation. Since memory is an important bottleneck with huge sizes of KVCache, the amount of memory used for meta-information is an im- portant consideration.
SCORE (KC, q): Given a query and a set tokens K,, V, SCORE assigns score each token in (K, V). TOPK the standard top-k these scores. GATHER-ATT (K,, q. I): Given a set of tokens IC, V, query q and set of indices I identifying pivotal tokens, GATHER-ATT computes the attention only using tokens indicated by I as follows.
2. Latency is another important factor. Since SCORE generally not latency bottleneck, one can devote more computation to obtain higher-quality retrieved tokens 3. Quality: A SCORE with higher recall, implies we can effectively use fewer tokens in GATHER-ATT which will further improve the overall efficiency of the attention procedure.
GATHER-ATT (AC, V.
[N]
1 E I) exp (qq, k,))v,
Σ
F 1 I) exp (1g. kj))
where Vi. k; refer to 7th embeddings of Y and IC respec- tively, and 1 is the indicator function. In the text, we will overload the function SCORE (k, v, q) to denote the score assigned to single key for methods that work on each key individually.
4.1. The SCORE function for HashAttention
Identifying the top tokens for a query is analogous to user-item recommendation problem. HashAttention uses two learnable mapping functions, Φx : R2dd - [0, 1]ah dg Rd [0, 1]" lift key-value and queries to hamming space of dimension b. In this space, given query the keys using hamming distance H. The SCORE function is defined by
Most sparse-attention methods can be viewed combi- nation of the above three routines with differences in the SCORE function. Following the SCORE routine, sparse- attention methods can be implemented using identical rou- tines for TOPK and GATHER-ATT Furthermore, we find in experiments, that most time is spent GATHER-ATE and TOPK routines. Thus we should devise good quality SCORE function with reasonably small latency. The quality of SCORE be measured in terms of its recall(n, k) de- fined as follows. Let I, be the set of top n indices chosen by SCORE function and Ierue,& be the set of & true pivotal tokens. Then, the recall is computed as
efficiently compare
SCORE (k, v. q) -H(ooo v), & (q))
(I)
For Φku and Φg. we use independent feed-forward network F, followed by a sign function to extract bits. The general function can be written as,
((x) relu (eluusign(F(x)))
The bits are packed into an integer. Let diint be the compete mapping function.
recall(In, Itrue.k) = In O Ierue,k
Implementation: The integer signatures for key-values are computed and cached along with KV Cache. During
Krrue

---

HashAttention
Training dataset We find that HashAttention trained on a generic dataset works well on a variety of tasks. To fur- ther improve the quality, we can finetune Atteniinn on the decoder hot path, we compute the query signature and perform bit operations compute the hamming distance between the query and key-values.
task-specific data. Additionally, we find that Attentinn trained on shorter sequences (say <= 64K) do not naturally scale to longer sequences (e.g. 128K). So, HashAltention needs to be trained for appropriate context-lengths.
Hi( (q)) = bitcount(bitwise http:///wiiiiieeoo((eereoieoe xor(dint (k, v), Pint .4 (q)))
We run the LLM model in inference mode in a chunk-based fashion, i.e. the tokens are processed in chunks. At the end of each inference, all HashAttention modules are indepen- dently and locally trained for one step on queries from this chunk and previously cached KV embeddings.
In our experiments, we focus on dwe that acts only the key vector k and ignore the v vector or its features such as norm. Incorporating v in 6kv is left for future exploration.
4.2. HashAttention for adapting to pre-trained LLMs
4.3. Theoretical underpinnings of HashAttenting
To use HashAttention for existing SDPA attention in pre- trained models, SCORE function of HashAttention should align with the best selection of tokens w.r. the final attention computation. The best ordering for an SDPA attention is presented in the lemma below.
In this section, we briefly explain why HashAttention works. Lemma 4.1 shows that analyzing token importance in isola- tion, the contribution of 7th token is proportional to a: [V: 22 where and V are attention score and value vector respec- tively. Identifying tokens with this true score can be cast as a maximum inner product search problem as shown below,
Lemma 4.1. Consider n. KV embeddings K., V and query q. Let be the attention scores of tokens w.tt query q. Then the contribution of a token i, assuming no other token is used, towards the final output is proportional to
Lemma 4.2. Consider a KV Cache K, V, and query vec- tor q. The ordering imposed on scores aslivilla for 7th token is same as the ordering imposed by inner product where [] denotes concatenation.
(lq, 1],
a.l/vallz
(3)
The proof is presented in appendix B.1. The Maximum inner product search problem can be reduced to a cosine-similarity search problem (Shrivastava & Li, 201 4; Neyshabur & Sre- bro, 2015) using asymmetric transformation on query and key-values. Combining Lemma 4.1 and Lemma 4.2 along with results from (Neyshabur & Srebro, 2015), we can equiv- alently determine ordering based on cosine similarity be- tween the following vectors,
The proof is presented in appendix B.I. We can use top tokens identified by this score to train our HashAttention modules for each attention. We pose HashAttention training as a classification problem where each embeddings gener- ated in HashAttention are used to predict the top-k tokens of the attention head that it is associated with. We choose k to be small number such as 64. We use binary cross- entropy loss in multi-class setting to train our functions okv and dg with a standard Adam optimizer. We mention some important details below,
Ukk (k, v) k, log(lv)l))) v /22 -I[k, (o((v
Class imbalance As the context length increases, the class imbalance for classification increases. For instance at 64,000 context length, while using top-64 tokens to predict, only 0.1% of labels are of class 1. We use class weights to resolve the issue of class imbalance. Since the imablance depends on the context length. We use the following for- mula to compute the class weights, parameterized with a and S:
where M is the maximum norm of all vectors [k, log([/vl)
bh (q) = [q. 1,
Cosine similarity can be approximated using bit signatures obtained by random signed projections (Gionis et al., 1999). HashAttention is motivated by this theoretical understand- ing of similarity search. Furthermore, we use Icarnable functions in place of 1a exploit the patterns in query and key distribution and obtain succinct bit signatures. Using independent encoding functions for key-values and queries helps in solving the MIPS problem.
class1- = a + /eonnext-length
(4)
a and 33 are hyperparameters that can be chosen.
Soft- cottipartiiioning While we use the sign function to perform space-pare sance-partitionin and obtain bits in é mapping computa- tion during inference, use the tanh function in place of the sign function as a softer version of partitioning while training.
4.4. Efficiency of token level sparsity.
As opposed to methods such as Quest or InfLLM that use block sparsity in attention, HashAttention uses token-level sparsity. We want to note that token-level sparsity is un- avoidable when trying to implement sparse attention during inference time when training assumes full attention, because
5

---

HashAttention
Table A head-to-head comparison of sparse-attention basclines on datasets from I LongRench from all the categories using the IJ Lama- 21-8B-Instruct InfLLM, DS and Quest use auxiliary memory for shorter representation of keys noted in Aux-budget as bits token. HashAttention while having smallest auxiliary budget outperforms baselines on average. HashAttention trained with mixture openwebtext dataset and LongBench datasets (25 examples dataset). Evaluation is done on first 175 examples (MFQA was excluded
from training since it only 150 samples)
Category + Aux:bits/token
MQA
SQA
Summ QmSm
FS-Learn TQA
Synthetic PassR
Code
Model
Tokens
HPQA
MFQA
RepoB
Average
Full Model
NA
NA
54.83
55.17
24.91
91.31
100.00
55.07
63.55
Oracle(top)
NA
512
52.10
53.45
25.14
91.39
100.00
58.49
63.43
H2O
NA
512
36.62
26.61
17.85
80.75
43.43
55.55
43.47
St.eamI_LM
NA
512
33.32
27.98
17.93
51.95
11.43
57.07
33.28
InfLLM
256(p==32.bit=16)
512
47.75
51.99
23.14
88.32
33.00
42.99
47.87
InfLLM
15)(
512
48.27
53.09
22.90
88.88
32.81
43.45
48.23
DS
32(ch=16.bit=2)
512
50.39
50.57
23.41
90.32
98.86
57.72
61.88
DS
64(ch=16.bit=4)
512
52.60
53.44
23.75
90.44
99.43
57.01
62.78
DS
bit))))))bit=66)
512
41.50
43.21
21.80
86.83
82.29
52.80
54.74
Quest
32(pg=16.bit=2)
512
53.13
51.31
23.01
90.15
98.29
58.29
62.36
Quest
64(pg=16.bit=4)
512
52.51
54.35
23.75
91.50
98.86
59.03
63.33
Quest
128(p8=32.bit= 16)
512
51.60
53.60
22.94
90.53
97.71
57.06
62,24
HashAttention
32
512
54.08
53.35
25.08
92.41
100.00
59.98
64.15
Table 2. Llama3. 1-BB 1-18---8ssttt Instruct with HashAttention- 16x on LongBench. AVG is average without passage count
passage-retin rettie--etricvel-en
multifieldqa-en
samsum
gasper
musique
gav-ropert
multi-news
troc
hotpolqu
passgge-cout
qmsum
2wiki
triviuqa
narrativeqja
lcc
repobenchup
AVG
AVG
54.68
100
43.75
45.80
29.41
34.90
27.11
71.43
55.40
8.29
25.14
45.62
90.55
30.19
62.87
55.31
48.78
51.48
16x
53.65
100
42.82
43.52
29.16
34.99
26.68
70.29
51.92
1.85
24.75
42.97
92.14
27.93
64.89
60.49
48.00
51.08
Table 3. Llama3. 1-8B-Instrue with HashAtten 16x on RULER@I6K
niah. multikey 1
niah. multikey 2
niah multikey 33
niah multivalue
niah multiquery
niah single_1
niah single
niah single.3
Average
cwe
fwe
qa qu_1
Full
100
100
100
100
100
100
99.25
98.75
99.2
50.5
90
87
54
90.66
16 x
100 100
100
99
90
96
98.5
98.75
93
62.4
87.33
82
57
89.53
important tokens, defined by higher attention scores, may not appear in the same block, and not using some important tokens adhere to inferior top-k approximation. Interestingly, and con- trary to popular belief, token-level sparsity does not lead to system inefficiency.
as contiguous data access meets or exceeds this cacheline size. important to note that many state-of-the infer- ence framework implement attention using paged-attention backbone with page size 1. They find that with correct opti- mization, the efficiency does not depend on page size (Ye et al., 2024).
imposed block size will only lead
HashAttention implementation is built upon the vLLM page attention framework (Kwon et al., 2023), where each token corresponds exactly to one page (i.e., page size 1). Af- their indices without
5. Experiments
In the experiments below, we use two versions HashAtten- tion. HashAttention is trained using openwebtext(Gokasan et al., 2019) dataset with samples concatenated together to obtain the required context length. HashAttention" is further finetuned from HashAttention for specific dataset.
ter identifying top-k tokens, physically moving gathering these tokens into 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
baselines: StreamingLLM (Xiao et al., 2023), H2O (Zhang et al., 2023), InfL.LM (Xiao et al., 2024), DoubleSparsity (DS) (Yang al., 2024), and Quest (Tang et al., 2024).
6

---

HashA ttention
LLAMA3.1-3B Cleei
STRALLO 3-78 DS
DS
Mlss Abw Kan
Hual Alae Ferf
Hest MLonite
RULER (LLAMA3,
11m
D0B
SIL2
ovg
TUC
Aggl drave0))
82
10
82
91 Skarsts
32x
-28c 8dc
Soars
10c D64 52c
81 1.
Figure 3. Pareto Curves for [LongBench] Average quality of gessaggeretreivaal een, triviaqa, multifieldqa.een and hotporqa from Longbench different sparsities auxiliary budget of 32bits PTPA. HashAttention almost reaches the full model quality at 16x for Llama-3.1 Instruct and 32x for BB-val-7---
HashAttention" further improves the Pareto curves to obtain full model quality at 32x for LLAMA-3.1--8 1-MA--.1-8B-BBE 1-8--.-t-B--n Detailed dataset results are deferred Figure 6 in Appendix. [RULER@ 16K] Quality of -.ama-3.1-8--Immm.. on RULER @1 16K benchmark at varying sparsity levels. On average, HashAttentions 16x achieves the same quality as the full model. Interestingly for most datasets, sparsity at some level helps improve over the full model (4/6 datasets considered show this sparsity >= 8x)
We also include Oracle which predicts exact top-k. For all methods, we include the first and most recent 128 tokens to ensure that the important parts of the prompt are always in
offered by HashAttention.
5.1. Quality with HashAttention
context. In cases where methods use auxiliary memory (e.g. label cache in DS and bit signatures in HashAttention), we denote as bits per token per hPTA)) The token budget refers to the number of heavy tokens retrieved. InfLLM and Quest are page-based retrievers. These baselines have two parameters: page size (pg) and the number of representative vectors (rep). In InfLLM, we always set rep=4. While these baselines do not inherently incorporate further quanti- zation, we apply additional quantization to reduce auxiliary memory usage, as is done with DoubteSparsity (DS). DS has two parameters: the number of channels (ch) and the number of quantization bits (bit). Baselines are further explained in Appendix A.
We present the following experiments to evaluate the quality of HashAttention.
A head-to-head comparison of HashAttention at fixed budget For this experiment, we randomly select one dataset from each category in the LongBench and use the first 175 samples (the last 25 samples are added the training set). We budget of 512 pivotal tokens for all baselines. The results are presented in the Table 1.
Pareto curves To get broader comparative picture, we use four datasets (Passage Retrieval (English), TriviaQA, HotpotQA, and MultifieldQA (English) from Longbench and six datasets (niah, vt, owe, fwe, qal, and qa2) from RULER (16K) to perform more elaborate Pareto analysis. We compare against Quest and DS at 32 PTPA auxiliary budget. Unless explicitly mentioned, we use HashAttention (and not HashAttention". ). These results are presented in Figure 3.
prompt-offset: To simulate a scenario of preprocessed long context with relatively smaller prompts, without relying on specific datasets, we apply sparse attention only to the last ofese tokens of the complete prompt and all subsequent generations. We use offset of 128 for Longbench and 32 in RULER.
Full benchmark at 16x sparsity We also present HashAt- tention at 16x sparsity on full LongBench and RULER benchmarks for Llama-3. L.ama-3.1-8B-Immeee --B-Irutttt The results are pre- sented in Table 2 and Table 3.
Auxiliary Memory: Since memory is a significant bottle- neck in deploying LLM models, we aim to evaluate different methods under the same auxiliary memory budget of 32 bits PTPA unless stated otherwise.
Mícro benchmarking the recall rates To show that the improvements in quality come from the improvement in pivotal token retrieval, we show the recall rates of top-32 tokens at different context lengths for the Llama-3, Instruct model. The result is presented in Figure 5 in the appendix.
Models and datasets: We use Iahhsrecy et al., 2024) and Mistral-7D-Instruut-00...30 2023) models and Longbench(Bai et al., 2024b) and RULER(Hsich et al., 2024) datasets in our evaluation. All attention heads are replaced with sparse attention.
et al.,
Effiicincy frameworks: We use GPT-FAST(GPTFast, 2024) and FlashDecode(Dao, 2023) to test improvements
We make the following observations,
Table
presents the landscape of different sparse-
7

---

HashA ttention
* GPT-FLASH * Lig
Sparsty. FleshDecode
FPT Hest-Wentor S2IGPT-
HashAAlardionn32x (F aash Decode)
20
11.29
12.45
09
5
2,50
2:53
6.24
10
1.31
2.27
latencyiis)
56
2.43
Tokens
0.5
10.24
0.75
0.20
0.1
"Oros" 5.11
0.05
.0.04
1k
16K
32K
B0K
131K
262K
524K
11
18c
30%
15K
(a) Latency of various kernels at different context lengths
(b) Throughput (batch=1) in GPT-FAST inference system at dif-
ferent context lengths
Figure 4. Efficiency improvements by incorporating HashAttention in GPT-FAST and FlashDecode. HashAttention- -3xx shows improve- ment in latency and end-to-end throughput on sequenees longer than 8K in GPT-FAST and 65K in FlashDecode
attentions. H2O and StreamingLLM fail due the in- correct fixed sparsity assumption. InfLLMM uses dynamic sparsity, but the heuristic-based scoring does not gen- eralize well over all tasks leading to poor performance. Quest and DS perform relatively well on all tasks. How- ever, it needs more auxiliary memory to obtain good re- sults. Overall, HashAAtention outperforms all the base- lines while using minimum auxiliary memory.
5.2. Efficiency with HashAttention
To evaluate the efficiency that HashAttention has to offer, perform the following experiments. In these experiments, use 3-layered MLP (128x 10--12xx128 (88-12xx2) for map- pings learned in HashAttention. Using smaller MLPs will improve latency results for HashAttention. We use 32x sparsity for these experiments.
Figure 3, Table 2, and Table 3 show that HashAttention significantly reduce the number of tokens in attention computation. HashAttention is almost 4x more token- efficient compared to Quest and DS at an auxiliary budget of 32 bits/token. Overall, HashAttention seems to pre- serve most of the model quality at 16x sparsity (0.78 points less on Longhench, 0.4 points without the worst performing dataset, and 1.13 points less on RULER). The
We incorporate HashAttention in the GPT-FAST attention kernel and FlashDecode attention kernel and measure the latency improvements.
We incorporate HashAttention in the GPT-FAST infer- ence system to measure the end-to-end improvement in throughput.
We make the following observations.
results can be further improved using HashAttention'.
The matrix multiplication mapping functions used to obtain bit-signatures dominates the latency up 8K con- text length in GPT-FAST and up to 65K in Flash De- code. Since the main difference between the two attention implementations is that FlashDecode uses sequence par- allelism, whereas GPT-FAST does not, this behavior is expected. Post 4K context length for GPT-FAST and 65K context length for FlashDecode, HashAttention improves the latency for the attention. We can see up 4.3x improvement in GPT-FAST and 2 improvement in FlashDecode.
Certain levels of sparsity seem to improve quality over the full model in certain tasks. This could be potentially explained as sparse attention avoids attention to distractor tokens.
Attenitt trained on English data also performs rea- sonably well Chinese datasets in LongBench, although with higher drop in quality. (2.66 points average. See table 5 in the appendix.) The results on coding tasks (lcc and repobench-p in Longbench) are especially better than the full model. Coding tasks might be more sensitive to distractors than general English tasks.
The throughput improvements can be seen starting at 8K context length. The throughput improvement ( (ratio of the throughputs) starts declining as sequence length increases. This can be attributed to TOPK operation becoming ex- pensive at higher sequence lengths. In this experiment, we the best throughput gains of 3.12x occur at sequence length of 32K for 32 sparsity.
The recall of pivotal tokens shown Figure (in ap- pendix) gives an insight into why HashAttention outper- forms other baselines. The reason is that HashAttention is better retrieving pivotal tokens.
8

---

HashAttention
5.3. Ablations
opposed to training-free methods this is limitation since every model needs its own HashAttention. This is, however, not very different from the calibration used in DS. To obtain better results, we can use further finetuning on the task- specific data. Practically, this is okay for users who care about specific application of LLM. In terms of evaluation, we restrict our evaluation to situations where KV Cache is on the GPU. For extremely long context situations, KV Cache needs to be stored on the CPU RAM, and evaluation of HashAttention in such a scenario using flash-decode would be most valuable. However, such an evaluation is out of the scope of this paper. Having said that, HashAttention clearly shows superior quality, larger sparsity, and lower auxiliary memory compared to different top-k-based sparse attentions under identical conditions.
In this section, we present ablations. The tables / figures for this section are presented in Appendix.
HashAttention Bit Signatures vs. LSH Bit Signatures: We compare LSH bit signatures (Gionis et al., 1999), ob- tained via random signed projections, with learned bit sig- natures used in HashAttentin.. We find that LSH requires significantly more bits to perform reasonably well on top-k prediction. This inefficiency stems from the agnostic nature LSH. The results are presented in Table 8. While HashAt- tention performs well with 32-bit signatures, LSH fails to achieve comparable performance even with more than 1000 bits.
Note that this is not a direct comparison to MagicPig (Chen et al., 2024), as MagicPig builds LSH tables from bit signa- tures for sparse retrieval and modifies the attention computa- tion to incorporate the sampling view of LSH. Nevertheless, it provides reasonable indication of the auxiliary memory used by MagicPig to achieve sparse attention.
7. Conclusion
HashAttention proposes a principled approach to sparsify attention based on hashing key-value pairs and queries to a compressed semantic space. Near neighbor comparisons can be efficiently performed in this semantic space using bit operations. The quality of tokens retrieved by HashAt- tention, under similar resources, is significantly superior to other sparsity baselines, leading upto 4x further of tokens as compared to leading sparse attention baselines. Overall, we see reduction in KV Cache usage upto 16-32x with minimal drop in quality (within point) across bench- marks and models, leading to latency and throughput gains.
Query and Key distributions under HashA em- bedding While HashAttention does not directly aim to resolve the out-of-distrbutiion issue between queries and keys, the learned mappings bring the distributions closer due to the nature of the training loss. We measure the empir- ical cosine similarity between each query and its top-32 to- kens using the original embeddings, HashAttention (signed) embeddings, and HashAttention (tanh) embeddings used during training. We find that the cosine similarity between key and query embeddings is sigmificantly higher when us- ing HashAttention embeddings compared to the original embeddings. The results are presented in Table 7.
further reduction
A Acknowledgmentss
Sky Computing Lab is supported by gifts from Accenture, AMD, Anyscale, Cisco, Google, IBM, Intel. Intesa San- paolo, Lambda, Lightspeed, Mibura, Microsoft, NVIDIA, Samsung SDS, and SAP. Aditya Desai's position at sky-lab is supported by Mayfield Charitable.
Quality vs. bit-width in HashAttention As expected, we find that using more bits helps with the quality of top- k prediction. The cross-entropy loss of top-32 prediction using different wit- is presented in Table 9.
Impact Statement
This paper presents work whose goal is to advance the effi- ciency of Machine Learning which can lead to wider adop- tion of large language models. There are many potential societal consequences of advancing machine learning and its efficiency, none of which we feel must be specifically highlighted here.
Latency of SCORE computation vs. full inner product. Since increasing bit-widths can improve attention quality, it is reasonable to ask how many bits can be used before the computational advantage of HashAttention disappears. We present this ablation in Table 6. We find that even for bit-widths as large as 512, the latency of SCORE compu- tation using our optimized semel-whichh employs bitwise operations- is still lower than that of full inner product com- putation for 1 128-dimensional vectors with f precision.
References
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint arXhv:2303.08774, 2023.
6. Limitations and Discussion
We discuss the limitations of the proposal and evaluation here. HashAttention needs to be trained at least once. As
9

---

HashAttention
An, C., Huang, F., Zhang. J., Gong, S., Qiu, X., Zhou, C., and Kong, L. Training-free long-context scaling of large language models. arXiv preprint arXiv:2402.17461 2024.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang. A., Fan, A., et al. The llama 3 herd of models. arXiv preprint arXiv:2407.2178. 2024.
Bai, G., Liu, J., Bu, X., He, Y., Liu, J., Zhou, Z., Lin, Z., Su, W. Ge, T., Zheng, B., et al. Mt-bench- 101: A fine- grained benchmark for evaluating large language models in multi-turn dialogues. arXiv preprint arXiv:2402. 14762, 2024a.
Feng. J., Sun, Q. Xu, C., Zhao, P., Yang, Y., Tao, C., Zhao, D., and Lin, Q. Mmdialog: A large-scale multi-turn dialogue dataset towards multi-modal open-domain con- versation, 2022. URL https: /ttp///axvvvv...//ee//eeeee.aev. .rweiiiiorrr.oooooooo... aba/ 2211.05719.
Bai, Y., Lv, X., Zhang, J., Lyu, H., Tang, J., Huang. Z., Du, Z., Liu, X., Zeng, A., Hou, L., Dong, Y., Tang, J., and Li, J. LongBench: bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 31 3119-3137, Bangkok, Thailand, August 2024b. Association for Com- putational Linguistics. doi: 10. 10.18663/w1/2004.cclooozo
Ge, S., Zhang, Y., Liu, L., Zhang, M., Han, J., and Gao, J. Model tells you what to discard: Adaptive kv cache compression for llms. arXiv preprint arXiv:23 10.0018010.0101,, 2023.
Gionis, A., Indyk, P., Motwani, R., et al. Similarity search in high dimensions via hashing. In Vldb, volume 99, pp. 518-529, 1999.
172. URL https: ac http:::://clnnntooooooooooo acl acl-long 1.2.
Gokaslan, A., Cohen, V., Pavlick, E., and Tellex, S. Open- webtext corpus. i ht//ppeneebTexxtCorpussss 2019.
githt..
Beltagy. I., Peters, M. E., and Cohan, A. Long- former: The long-document transformer. arXiv preprint arXiv:2004.00 05150, 2020.
GPTFast. Gptfast, 2024. htthtttp:///gtthiiie com/ Accessed: 2024-12-15.
Bricken, and Pehlevan, Attention approximates sparse distributed memory. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 15301-15315. Curran Associates, Inc., 2021. URL cc
pytorch-labs happlllllssspppleaeeeeeegg
Guo, M., Ainslie, J., Uthus, D., Ontanon, S., Ni, J., Sung, Y.-.,, and Yang, Y. LongT5: Efficient text-to-text trans- former for long sequences. In Carpuat, M., de Marneffe, M.-C., and Meza Ruiz, (eds.), Findings of the Asso- ciation for Computationnll Linguisties: NAACL 2022, pp. Seattle, United States, July 2022. Association for Computational Linguistics. doi: 10. (0.18653/v1/0022 555c andingg--aacl.55. URL https: aaa//claantt aclaathooooooo.. 55.
meuripss..
817 4477a2225544335504ccc0ff8883883 Paper. pdf.
findings
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
org brg/2022 2022
Han, L., Jayaram, R., Karbasi, A., Mirrokni, V., Woodruff, D. P., and Zandich, A. Hyperattention: Long- context attention in near-linear time. arXiv preprint arXiv:2310.05869, 2023.
Chen, Z., Sadhukhan, R., Ye, Z., Zhou, Y., Zhang, J., Nolte, N., Tian, Y., Douze, M., Bottou, L., Jia, Z., et al. Mag- icpig: Lsh sampling for efficient 11m generation. arXiv preprint arXiv:2410. /6179 2024.
Hooper, C., Kim, S., Mohammadzadeh, H., Maheswaran, M., Paik. J., Mahoney, M. W., Keutzer, K., and Gholami, A. Squeezed attention: Accelerating long context length Ilm inference. arXiv preprint arXiv:2411.09688 2024.
Dao. T. Flashattention-2: Faster attention with bet- ter parallelism and work partitioning. arXiv preprint MxXiv:2307.0869) 2023.
Hsieh, C.P.., Sun, S., Kriman, S., Acharya, S., Rekesh, D., Jia, F., Zhang, Y., and Ginsburg, B. Ruler: What's 'ss the real context size of your long-context language models? arXiv preprint arXiv:2404.0666 2024.
Davies, R. and Murff, M. Using generative ai and gpt chat- bots to improve instruction. In INTED2024 Proceedings, pp. 5451-5459. IATED, 2024.
Jiang, A. Q. Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., Casas, D. d. I., Bressand, F., Lengyel, G., Lample, G., Saulnier, L., et al. Mistral 7b. arXiv preprint arXiv:2310 06825,, 2023.
Deng, Y., Song, Z., and Yang, C. Attention is naturally sparse with gaussian distributed input. arXiv preprint mxXiv22444.026990 2024.
10

---

HashAttention
Kwon, W., Li, Z.,, Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. F.,, Zhang, H., and Stoica, I. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023.
Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., and Roformer, Y. L.. Enhanced transformer with rotary position embed- ding.. 2021. DOI: https://doi. org/10.1016/ neucom, 2023.
Sun, Z., Yang, Y., and Yoo, S. Sparse attention with learn- ing hash. In International Conference on Learning Representations, 2021.
Lewis, P., Perez, E.,, Piktus, A., Petroni, F., Karpukhin, V., Goyal, N., Küttler, H., Lewis, M., Yih, W.-t, Rocktäschel, generation for knowledge- intensive nlp tasks. Advances in Neural Information Pro- cessing Systems, 33:9459-944 9474, 2020.
T., et al. netrieval-augmennte
Tang. J., Zhao, Y., Zhu, K., Xiao, G., Kasikei, B., and Han, S. Quest: Query-aware sparsity for efficient long-context Ilm inference. arXiv preprint arXiv:2406. 10774, 2024.
Li, Y., Huang, Y., Yang. B., Venkitesh, B., Locatelli, A., Ye, H., Cai, T., Lewis, P., and Chen, D. Snapkv: Llm knows what you are looking for before generation. arXiv preprint ar arXiv:2404. 14469, 2024a.
Team, G., Anil, R., Borgeaud, S., Alayrac, J..B., Yu, J., Sori- cut, R., Schalkwyk, J., Dai, A. M., Hauth, A., Millican, K., et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312.
2023.
Li, Z., Li, C., Zhang, M., Mei, Q., and Bendersky, M. Re- trieval augmented generation or long-context llms? a comprehensive study and hybrid approach. In Proceed- ings of the 2024 Conference on Empirical Methods in Natural Language Processing: Industry Track, pp. 881 893, 2024b.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation lan- guage models. arXiv preprint arXIv:2302.1391 2023.
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Liu, D., Chen, M., Lu, B., Jiang, H., Han, Z., Zhang, Q., Chen, Q. Zhang, C., Ding, B., Zhang, K., et al. Re- trievalattention: Accelerating long-context Ilm inference via vector retrieval. arXiv preprint arXiv:2409. 10516, 2024a.
Xiao, C.,, Zhang, P., Han, X., Xiao, G., Lin, Y., Zhang, Z., Liu, Z., Han, S., and Sun, M. Infllm: Unveiling the intrinsic capacity of Ilms for understanding extremely long sequences with training-free free memory. arXiv preprint Wwxiv:2402.04611 2024.
Liu, H., Yan, W., Zaharia, M., and Abbeel, P. World model on miltton-length video and language with ringattention. arXiv e-prints, pp. arXiv-2402, 2024b.
Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Ef- ficient streaming language models with attention sinks. arXiv preprint arXiv:2309. 17453, 2023.
Liu, Z., Desai, A., Liao, F., Wang, W., Xie, V., Xu, Z., Kyril- lidis, A., and Shrivastava, A. Seissorhands: Exploiting the persistence of importance hypothesis for 11m kv cache compression test time. Advances in Neural Information Processing Systems, 36, 2024c.
Yang. S., Sheng, Y., Gonzalez, J. E., Stoica, I., and Zheng, L. Post-training sparse attention with double sparsity. arXiv preprint arXiv:2408..00 07090, 2024.
Ye, Z., Chen, L.,, Lai, R., Lin, W. Zhang, Y., Wang, S. Chen, T., Kasikci, B., Grover, V., Krishna- murthy, A., and Ceze, I.. Flashinfer documentation, 2024. URL http::::///aassneeeee.///// o2/introduce- fl www.sffffsshaafaaaaaaaaaaaaaaaea html. Accessed: 2025-05-27.
Ma, X., Yang, X., Xiong, W., Chen, B., Yu, I.,, Zhang, H., May. J., Zettlemoyer, L., Levy, O., and Zhou, C. Mega- Efficient I1m pretraining and inference with un- limited context length. arXiv preprint arXiv:2404.08801, 2024.
lodon:
Neyshabur, B. and Srebro, N. On symmetric and asymmetric Ishs for inner product search. In International Conference on Machine Learning, pp. 1926-1934. PMLR, 2015.
Zhang, A., Lipton, Z. C., Li, M., and Smola, A. J. Dive into deep learning. arXiv preprint 23X1v:2106.1132, 2021.
Zhang, Z., Sheng, Y., Zhou, T., Chen, T. Zheng, L.., Cai,
Shrivastava, A. and Li, P. Asymmetric Ish (alsh) for sublin- time maximum inner product search (mips). Advances in neural information processing systems, 27, 2014.
R., Song, Z., Tian, Y., Ré, C., Barrett, C., et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Informtion Processing Systems, 36:34661 -34710, 2023.
Singhania, P., Singh, S., He, S., Feizi, S., and Bhatele, A. Loki: Low-rank keys for efficient sparse attention. arXiv preprint Xiv:24006.000 02542, 2024.

---

HashA
A. Baselines
The primary purpose of HashAttention is to sparsify attention to improve the memory loading and computation of attention computation. Thus we consider approaches in literature with similar goals.
A.1. SLL
This attention was developed for processing streams of text. There are two different challenges when it comes to ingesting stream of text with existing LLM models. Firstly, most LLM models are trained on restricted context length and thus do not support longer texts. Secondly, the computational cost of decoding grows linearly with context length leading to prohibitive computation. StreamingLLLM proposes to keep two portions of the context the first (also called attention sink and the local context. Thus it sparsifies the attention computation.
Since StreamingLLM categorically ignores a major chunk of context. It does not perform well on long context benchmarks which needs the model to consider all the context.
Hyperparameters: sink size and local size.
A.2. ScissorHands H2O
This attention was developed primarily to reduce the memory storage of KV Cache - goal very well aligned with sparsification of attention and reducing memory loading during decoding. However, the setting used ScissorHands / H2O is more restrictive since, the decisions are made in streaming fashion and tokens evicted can never be retrieved unless recomputed. The idea in ScissorHands/ / H2O is that if some tokens in context have not been useful in recent predictions then they are not going to be useful in the future generations and can be dropped. In our experiments we use H2O since they have easier codebase.
Seissorhands and H2O both heuristically drop the tokens. The tokens dropped at one point are not available subsequently. This is clearly an issue in different settings such multi-turn chat etc. should be noted that the proposal of ScissorHands and H2O are for reducing decoding time monologue of LLM. In that particular setting the proposals are useful. But their effectiveness is also restricted to that setting.
Hyperparameters: sink size, local size and token budget.
A.3. Retrieval Attention
In Retrieval Attention, the attention computation is preceded by top-k computation using approximate near neighbor search algorithms and full attention is computed on estimated top-k. Most graph based algorithms (including the one proposed in Retrieval Attention) need to be run on CPUs due to their irregular computation pattern. Thus, Retrieval Attention by default always stores the KVCache on CPU.
This is a close cousin of HashAttention. The motivation of both methods is identical in the sense that attention can be replaced by approximate near neighbour search. RetrievalAAteenion uses traditional graph based search algorithm to find near nieghbours, whereas HashAttention uses learning hash to create a quantized similarity space for retrieval. A major drawback of Retrieval is that it is not GPU friendly which causes indexing and querying to be slower for large contexts.
Hyperparameters: sink size, local size and ROAR graph hyper parameters.
A.4. InfLLM
InfLLM maintains the attention sink and local context as with streaming LLM. Additionally, it also maintains the tokens in between and retrieve chunks from them. It divides the KVCache into contiguous chunks and from each chunk a few representative keys. These keys are used to compute the importance of the chunks w.r.t a given query. Top few chunks are chosen for final attention computation via full computation. In order to choose top scoring chunks, the paper proposes performing full dot product computation can be performed with representative vectors which can also replaced by off-the-shelf near neighbour computation on CPUs.
This again is similar to the setup of HashAttention. The chunk based treatment for retrieval reduces the computational cost
12

---

HashAttention
of computing the relevant chunks. However, the heuristic way of computing the representative keys lead to issues of missing key chunks while retrieving. Apart from that, the method is identical to RetrievalAtttntion.
Hyperparameters: sink size, local size, token budget, page size, number of representative vectors
A.5. Quest
Quest is similar InfLLM with the difference being the computation of the importance of a chunk. Quest maintains and min vectors for each chunk which maintains co-ordinate wise extremes of the chunk. These are used to estimate the maximum possible inner product within chunk given a query.
clear to see that we restrict the token budget, we might end up retrieving falsely important chunks at the cost of discarding important ones. We see this in our experiments as well.
Hyperparameters: sink size, local size, token budget, page size
A.6. Double Sparsity
Double sparsity chooses 16 coordinates from the 128 dimensions of key embeddings via offline calibration and use them for computing top-k. Then full attention is computed these top-k. This 16 dimensional slice of K Cache is called as label cache and can be further quantized to 4 bits without much impact on accuracy.
This again is similar to the setup of HashAttention and RetrievalAttention the difference being how to compute the top-k. Surprisingly the 16 channels (16x16 16 = 256 bits) identified by Double Sparsity are good top k indicators. In comparison HashAttention uses 32 bit signatures and uses bit operations for manipulating signatures.
Hyperparameters: sink size, local size, token budget, label size.
Quality and computation of sparsity with HashAttention vs. baselines
The choice of completely ignoring parts of context in streamingLLM, heuristic based permanent eviction of Scis- sorhands/H2O, and heuristic based representative key selection of InfLLM causes these approaches to be inferior to HashAttention. Retrieval Attention, Double Sparsity and Attenitt all are based on determining the top-k and using for attention computation. Thus, the quality depends on ANN algorithm used. In terms of computational complexity, HashAttention and Double sparsity can run on GPU and thus are faster for reasonably sized context lengths as compared Retrieval Attention. Additionally, HashAttention only uses an integer signature for computation of top-k which is memory and compute effective as compared to Double sparsity
B. Lemma Proofs
B.1. Lemma 1: Best Sparse solution for adaptation for 1-token selection
Using the following notations, V : n x d value matrix, a : n x 1 attention scores. $ n x n diagonal sampling matrix with exactly one non-zero in the diagonal (selection of a single token). Note that this ignores the cross token corrleation in value vectors. Then under the sampling the final embedding is,
a SV
(5)
The residual is ||a TV - a. T SVI12. Then the best sampling is the one that minimizes the
S argmin||o Tv TV- T SVIs
(6)
||a V - -a SV
(7)
()V(s ( (I- S T
(8)
T - (1-VVS (1-S)a
99)
n (1 1st.af1|Vc.11
Σ
(10)
13

---

HashAttention
To minimize the residual, we have choose, st be which have higher 0g Vill
B.2. Lemma 2: Identifying topk w.r. isolated token score is a MIPS problem
(/q. |0ll22
a:|lv:112
(11)
Σ ela.k)
k +log(lvlll
(12)
Eelakj
x (q. k) + log(IIv:ll)
(13)
= (1q, 1], [k, log([Ivil))
(14)
C. Other ablations
C.1. Recall rates of HashAttention and DS
veerage with sparsity (Layer 17)
Average top-32 with sparsity (Layer
DaubleSpensily
HashAttenticn
0.80
0.75
MM
Neca
0.6
0.55
0.60
DoubleSparsity
0.55
HashAttention
10000
20000
40000
50000
60000
10000
30000
40000
50000
60000
Context length
Context length
Figure 5. Microbenchmark for recall rates of 12p-32 tokens 16x sparsity for HashAttention and DS. We that Hash HashAttention consistently outperforms DS at varyingn context lengths.
C.2. Detailed pareto curves for each dataset averaged the main paper.
The pareto curves are shown in Figure
Table 4, Longbench (non-chinese tasks)
passage, retrievel cn.en sumsum
multifieldqa.am en
qusper
musique
gov report
multi neww
tree
hoppttg
passage-contd
gmsunt
2wiki
triviaga
natrativ eqa
Ice
repobenchh-p
AVG
AVG
IlashAtiention 16x
53.65
100
42.82
43.52
29.16
34.99
26.68
70.29
$1.92
1.85
24.75
42.97
92.14
27.93
64,89
60.49
48.00
$1.08
Full
54.68
100
43.75
45.80
29.41
34.90
27.11
71.43
55.40
8.29
25.14
45.62
90.55
30.19
62.87
55.31
48.78
51.48
Table Longbench (chinese tasks). mulsifieldga.hh dureader
The hash attention is trained on english data only passage-retrieaal
vesum
59.31
61.36
34.51
17.14
96.86
44.57
50.888
C.3. Latency of SCORE computation
The latency comparison SCORE computation and full inner product computation is presented in Table 6
14

---

HashAttention
Table In order understand the number of bits upto which HashAttention would show latency gains, compare the full inner product computation latency with bitwise score computation latency for varying number of bits We can see that upto 512 bits, our bit-wise kemell superior full inner product computation.
Latency of SCORE computation
Inner product
Hamming distance using packed 64 bit integers
#tokens
128 16Float
64bit
128bit
256bit
512bit
262144
0.13
0.085
0.084
0.087
0.112
524288
0.266
0.085
0.086
0.117
0.218
1048576
0.619
0.087
0.129
0.226
0,429
2097152
1.698
0.144
0.247
0.445
0.85
4194304
3,272
0.28
0.485
0.88
1.694
8388608
6.733
0.552
0.96
1.755
3.378
16777216
8.306
1.1
1.911
3.504
6,746
33554432
16.61
2.188
3.817
7.005
13.48
Table Average Cosine similarity between query and key emheddings of top-32 tokens. This computed Llama-3.1 10.1--88-nnntuue
model
Average Cosine Similarity of top-32 tokens
Layer number
original embeddings
HashAttennion (tanh)
HashAttentiinn(sin))
-0.111818
0.1777
0.180198
4
-0. 19951
0.4922
0.179186
8
0.1109657
0.2676
0.1802
12
-0.10329
0.1309
0. 0.178136
-0.1 10954
0,2422
0.182292
-0,0985931
0,2188
0.179659
24
-0. 00.14441
0.293
0.179346
28
-0.101307
0.2441
0.179191
C.4. Does HashAttention resolve the out-of-distribution issue of queries and keys
While HashAttention does not directly aim resolve the out-of-distributie the learned mappings do cause the distributions to be closer. We present some data in Table 7
C.5. Learned bit representations performed in HashAttention data-agnostic bit signatures using LSH (random signed projections) (Gionis et al., 1999)
If LSH instead of learned bit signatures via mapping functions as done in HashAttention. the bit signature sizes needed for LSH are considerably higher. This is expected since LSH does not exploit the data distribution in the selection of projections. The results on LongBenchmark datasets is presented in Table 8. This is consistent with results from MagicPig(Chen et al., 2024).
C.6. Quality vs. bits
The Table 9 shows the cross entropy loss (lower the better) for top-64 prediction while using different bit-widths.
15

---

HashA ttention
Table we use LSH (used in MagicPig(Chen et al., 2024)) instead learned bit signatures mapping functions as done
HashAttention, the bit signature sizes needed for LSH are considerahly higher. This is expected since LSH does not exploit the data distribution selection of projections
Compression 16x
#bits
passage retrieval en
multifieldqa
hotpotqa
triviaqa
Average
HashAttention
32
99.43
55.18
52.34
93.32
75.0675
16x
LSH
32
48.57
32.62
34.38
80.02
48,8975
LSH
256
89.14
47.97
53.97
88.07
69.7875
LSH
512
85.14
50.38
51.42
87.6
68.635
LSH
1024
91,43
49.75
55.25
OOM
8x
LSH
32
69.14
42.71
46.63
81.78
60.065
LSH
256
82.86
50.18
52.06
90.3
68.85
LSH
512
95.43
52.64
54.85
87.51
72.6075
LSH
1024
98.86
51.84
54.67
OOM
Table Cross-entropy loss for different embedding dimensions evaluated on 250 32K samples. Dimension
Cross Entropy Loss @ 250 32K samples
16
0.243
32
0.212
64
0.204
16

---

ttention
HTPTQA
PassR
Ovest
25
1120
134
112
105
244/WV
LLAMA3. 1-B--tt---- MFQA
TQA
Oxeet
D5
55
1028
109
HTPTQA
PassR
Quest
DS
HastAttertion
FUll
HastMWertioo"
HastAttertonn
HashWtentonn
55
100
45
35
1/128
1/84
1/32
1/18
30
Sparsity
MISTRAL-7B--0.3
TQA
MFQA
Ques!
HashAteniion
Full
Hlss WAllefiorr
Hashiiturtion
Hashitannioon
90
86
85
1/128
1164
1/22
1M6
Sparsit
1/126
1/64
102
185
Figure Parcto curves at bits PTPA for HashAttention, DS and Quest
17


Reviews:
[{"review_id": "ZxZz1swE7c", "reviewer": "Reviewer_fb4H", "summary": "This paper introduces HashAttention, framing pivotal token identification as a recommendation problem. Given a query, HashAttention encodes keys and queries in Hamming space, capturing the required semantic similarity, using learned mapping functions. HashAttention efficiently identifies pivotal tokens for a given query using bitwise operations and computes attention using only these tokens, improving the overall attention efficiency. Trained on generic data, HashAttention reduces tokens used by up to 16× with minimal quality loss, requiring only 32 bits of auxiliary memory per token.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.", "claims_and_evidence": "From the theoretical analysis, HashAttention transforms the problem into a maximum inner product search problem and approximates it in Hamming space by learning mapping functions, which has a reasonable theoretical. The experimental results support the effectiveness. On multiple data sets and models, HashAttention performs better than baselines in the same auxiliary memory budget.", "methods_and_evaluation_criteria": "Compared with the existing methods, progress has been made in reducing the use of KV cache and improving the efficiency, and improving the efficiency of attention computation.", "theoretical_claims": "From the theoretical analysis, HashAttention transforms the problem into a maximum inner product search problem and approximates it in Hamming space by learning mapping functions, which has a reasonable theoretical.", "experimental_designs_or_analyses": "The experimental design was comprehensive, comparing multiple baselines, covering different models and data sets, and considering different assessment indicators, such as quality, efficiency, recall rate.", "relation_to_broader_scientific_literature": "HashAttention transforms the problem into a maximum inner product search problem and approximates it in Hamming space by learning mapping functions, which has a reasonable theoretical and good application value.", "essential_references_not_discussed": "Not yet.", "other_strengths_and_weaknesses": "Strengths \n1.The idea of paper is interesting, e.g., transforming key token recognition into recommendation problem and realize efficient attention calculation by coding in Hamming space. \n2.Well-written paper with a clear process, the English writing is easy to follow.\n3.Extensive experiments demonstrate the validity of the model.\n\nLimitations\n1.Supplement the experiment in long context where KV cache is located in CPU RAM, and compare the performance of HashAttention with other baselines.\n2.Expand the experiments of HashAttention on complex inference and multimodal tasks.\n3.The formula derivation process is elaborated, and the key terms are explained in detail.", "other_comments_or_suggestions": "No.", "questions_for_authors": "No.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mFthBWgqMQ", "reviewer": "Reviewer_HrF1", "summary": "The authors propose a method to identify relevant tokens in the attention computation by framing it as a MIPS search, using the relationship between MIPs and cosine similarity plus the approximation of cosine similarity in terms of the hamming distance of the corresponding hamming embeddings. The authors propose a framework, which other sparse attention approaches can be fit into, which consists on a scoring, topk and gather attention steps. Hash attention is first trained offline with generic data. The authors learn independent mappings for key-value pairs and queries and then use the hamming distance between these mappings to identify the pivotal tokens.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n1. **Discussion on RAG** – We will remove this from related work.\n\n2. **Dimension of hamming codes:**\n\nWe 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.\n\n| Latency of SCORE computation ||||||\n|:----------------------------:|:-------------:|:---------------------------------------------:|:------:|:------:|:------:|\n|                              | Inner product | Hamming distance||||\n|            #tokens           | 128 x 16Float |                     64bit                     | 128bit | 256bit | 512bit |\n|            262144            |      0.13     | 0.085                     |  0.084 |  0.087 |  0.112 |\n|            524288            |     0.266     | 0.085                     |  0.086 |  0.117 |  0.218 |\n|            1048576           |     0.619     |0.087                     |  0.129 |  0.226 |  0.429 |\n|            2097152           |     1.698     |0.144                     |  0.247 |  0.445 |  0.85  |\n|            4194304           |     3.272     |0.28                     |  0.485 |  0.88  |  1.694 |\n|            8388608           |     6.733     |0.552                     |  0.96  |  1.755 |  3.378 |\n|           16777216           |     8.306     |1.1                      |  1.911 |  3.504 |  6.746 |\n|           33554432           |     16.61     |2.188                     |  3.817 |  7.005 |  13.48 |\n\nWe 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. \n\n3. **unbalanced buckets issue in LSH**\n\nSince 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.\n\n4. **query OOD issue**\nThe 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.\n\n\n|              | Average Cosine Similarity of top-32 tokens |                      |                     |\n|--------------|:------------------------------------------:|:--------------------:|:-------------------:|\n| Layer number | original embeddings                        | HashAttention (tanh) | HashAttention(sign) |\n|            0 |                                  -0.111818 |               0.1777 |            0.180198 |\n|            4 |                                  -0.119951 |               0.4922 |            0.179186 |\n|            8 |                                  -0.109657 |               0.2676 |              0.1802 |\n|           16 |                                  -0.110954 |               0.2422 |            0.182292 |\n|           24 |                                  -0.114441 |                0.293 |            0.179346 |\n\n5. **K-means based topk vs. vanilla LSH vs. HashAttention:**\n\nHashAttention 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\n\t\n6. **Extra fine-tuning details**\n\nTo 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.", "claims_and_evidence": "The claims made in the submission supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The methods and/or evaluation criteria chosen make sense for the problem or application at hand.", "theoretical_claims": "Lemma 4.2 presents a true score for token relevance that has the same ordering as using the inner product between queries and key-value pairs plus states the equivalent with respect to cosine similarity. These lemma includes a proof in the appendix and the relevant references to support it.", "experimental_designs_or_analyses": "Yes, table 1-2 and figure 3", "relation_to_broader_scientific_literature": "The authors include a wide overview for the works in the area and classify them in fixed and dynamic sparsity plus token eviction methods.", "essential_references_not_discussed": "They cite up to date papers such as Retrieval Attention, Magic Pig and Squeezed attention, among others.", "other_strengths_and_weaknesses": "The paper includes a variety of experiments and evaluations that showcase the different strengths of the method.", "other_comments_or_suggestions": "It is not clear why retrieval augmented generation approaches are discussed in related work but these are not compared against in the experiments section. Even though RAG is a competitive approach to reduce the context length, it doesn't immediately relate to dynamic sparsity attention computation so I suggest to remove this from the related work section.", "questions_for_authors": "What is the dimension of the hamming codes to guarantee that it is indeed cheaper than computing the cosine similarity?\nHow do you solve the unbalanced buckets issue in LSH, given that you use a single hash table?\nCould you provide further details of how the learning of independent mappings for key-value pairs and queries remedies the OOD problem?\n Would it be feasible to include the topk with k-means approach in the baselines to understand how it does wrt LSH partitioning?\nCould you provide further details about the extra fine-tuning on datasets for the experiments?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FJbegG6uqn", "reviewer": "Reviewer_buih", "summary": "Dynamic sparse attention has been widely explored in long-context scenarios. This paper proposes a learned hash function-based token-level dynamic sparse loading method. Specifically, it formulates the sparse attention top-K problem as a recommendation task, utilizing a learnable hash function to predict top-K tokens.\nThe learned hash function is implemented as a three-layer MLP, generating a 32-bit hash score. During decoding, the method computes the query hash score, selects top-K tokens, and performs token-level gather sparse attention.\nThe approach is evaluated on LLaMA-3.1-8B and Mistral-v0.3-7B using LongBench and RULER, with end-to-end and kernel-level latency comparisons against FlashDecoding and GPT-Fast. Results show that in long-context scenarios, the proposed method achieves up to 4.3× speedup over GPT-Fast.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n1. **training-free hashing**\nTo 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.\n\n||| bits | passage_retrieval_en| multifieldqa_en| hotpotqa|triviaqa|Avg|\n|-----|---------------|:----:|:---------------------:|:---------------:|:--------:|:--------:|:-------:|\n| 16x | HashAttention |  32  |99.43 |55.18      |   52.34  |   93.32  | 75.0675 |\n| 16x |      LSH      |  32  |48.57 |32.62|   34.38  |   80.02  | 48.8975 |\n|     |      LSH      |  256 | 89.14  |47.97|   53.97  |   88.07  | 69.7875 |\n|     |      LSH      |  512 | 85.14 |50.38|   51.42  |   87.6   |  68.635 |\n|     |      LSH      | 1024 |91.43 |49.75|   55.25  |    OOM   ||\n|  8x |      LSH      |  32  |69.14| 42.71|   46.63  |   81.78  |  60.065 |\n|     |      LSH      |  256 | 82.86 | 50.18|   52.06  |   90.3   |  68.85  |\n|     |      LSH      |  512 | 95.43 | 52.64|   54.85  |   87.51  | 72.6075 |\n|     |      LSH      | 1024 | 98.86 | 51.84|   54.67  |    OOM   ||\n\n2. **Token level hashing / gathering in sparse attention / performance efficiency.** \n\nOur 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.\n\nMany 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\n*“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.\n\n3. **OOD test sets / generalization of learned mappings across tasks.**\n\n 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.\n\n4. **OOD Query** Kindly refer to Response 4 to reviewer HrF1.\n\n5.  **MagicPig vs. HashAttention:**\n\nMagicPig 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. \n\nSecondly, 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.\n\nWe 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.\n\n6. **Lemma 4.2** : We would be happy to cite the correct reference in prior literature. Kindly direct us to the same.\n\nPlease let us know if there are any additional queries.", "claims_and_evidence": "The paper's claims are fairly accurate, as it is the first work to leverage a learnable hash function to address the sparse attention top-K problem.", "methods_and_evaluation_criteria": "The method design is reasonable. Unlike Quest-based approaches, which rely on block-wise min-max values as centroids, this method learns a hash index, which is likely to provide better representations and mitigate OOD issues. The main concern is the generalization capability of the learned hash index.\n\nAdditionally, the choice of a hash index is justified—the hashing process is GPU-friendly, and storing hash results incurs minimal cost.", "theoretical_claims": "The paper also provides an error analysis for sparse attention during decoding and reformulates the MIPS problem using cos-sin transformations with an addition operation. However, this Lemma 4.2 has already been widely adopted in the vector retrieval community.", "experimental_designs_or_analyses": "1. The paper lacks a comparison with training-free hash function methods, such as MagicPIG[1], which also employs LSH for sparse attention. Evaluating against such baselines would help clarify the advantages and limitations of the proposed approach.\n2. Additionally, the token-level approach may not be GPU-friendly, yet there is no latency breakdown provided for the gather sparse attention stage. An ablation study on the performance gain from token-level hash retrieval is also missing.\n3. The study lacks OOD experiments, as it is trained on a retrieval dataset but is not tested on reasoning or other domain-specific benchmarks. Evaluating generalization across different domains would strengthen the analysis.\n4. Finally, there is no discussion or analysis on whether the learned hash index helps mitigate query-key distribution shift issues in sparse attention top-K retrieval. Would you be able to provide additional experiments or insights on this aspect?\n\n[1] MagicPIG: LSH Sampling for Efficient LLM Generation. ICLR 2025.", "relation_to_broader_scientific_literature": "Prior work has explored sparse attention methods for long-context LLMs, leveraging attention sparsity through approaches such as StreamingLLM and H2O. More recently, RetrievalAttention has introduced vector retrieval to optimize sparse attention computation. Building on these efforts, this paper proposes a learnable hash index to make sparse attention top-K selection more GPU-friendly.", "essential_references_not_discussed": "MagicPIG, a concurrent study, also utilizes LSH for sparse attention. It would be beneficial to include a discussion comparing the advantages and limitations of these two approaches. Could you provide a discussion comparing the advantages and limitations of these two approaches?", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "Additionally, I noticed an issue in #43, where 128GB should be corrected to 32GB, where LLaMA-3.1-8B uses GQA with a group number of 8.", "questions_for_authors": "1. Do you have an ablation study using only a training-free hash index?\n2. Do you provide a latency breakdown for gather sparse attention?\n3. Have you conducted any OOD analysis between the training and test domains?\n4. Have you analyzed whether the learned hash index mitigates OOD issues in sparse attention top-K retrieval?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1Ctld9QUd8", "reviewer": "Reviewer_p8mb", "summary": "This paper proposes a simple, effective and plug-and-play method for accelerating inference in autoregressive transformers via top-k attention. The authors propose to accelerate top-k operation by learning mappings to encode queries and keys in Hamming space in a way that the ranking induced using negative Hamming distance of encoded queries and keys follows the ranking induced by original exp(<q, k_i>) * ||v_i||. By efficiently identifying the top keys using bit operations, HashAttention reduces attention computation. Authors demonstrate impressive sparsity levels (up to 32×) with minimal quality degradation, and significant latency improvements (up to 4.3× in GPT-FAST and 2.54× in FlashDecode) with modest auxiliary memory requirements (32 bits per token).", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n\n1. **Hyperparameters for HashAttention Training / Frozen backbone LLM**\n\n\tYes 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\n\n|                            |             |\n|----------------------------|-------------|\n|    Openwebtext samples   |     ~128K    |\n|     32K length samples     |     3750    |\n|    # queries per sample    |     8192    |\n| # total queries in dataset |   30M  |\n|          optimizer         | Adam(0.001) |\n\n2. **performance overhead of computing the Hamming distance vs. full computation**\n\nKindly refer to the response 2 (Reviewer HrF1) for the table. We make the following observations,\n\nA. (Sequence dimension) Improvement with using hamming distances over inner products stays consistent with increasing tokens in KV Cache as expected.\n\nB. (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.\n\nWith such an efficient distance computation, the Hamming distance computation is far from being the bottleneck of HashAttention.\n\n3. **different sparsity levels for different layers / heads**\n\nWe have not explored this dimension. However, we expect such improvements to orthogonally improve all the methods.  We leave this for future work.\n\n4. **HashAttention performance with quantized KV cache implementations**\n\nQuantization 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. \n\nThe 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.\n\n**Other discussion:**\n\n1. **Code**: The anonymous repository for our code used for running long benchmarks is here: https://anonymous.4open.science/r/HashAttention_ICML2025/README.md\nWe plan to release full code on acceptance\n\n2. **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.\n\n| Dimension | Cross Entropy Loss @ 250 x 32K samples |\n|-----------|--------------------|\n|     16    |        0.243       |\n|     32    |        0.212       |\n|     64    |        0.204       |\n\nWe will add more results and discussion on these hyperparameters in final version.\n\n3. **overhead of signatures for < 8K contexts, Does overhead become more prominent with quantized KV**\n\nWe 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.\n\n4. **Formulation**: we will improve the formulation in the paper.", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "yes", "experimental_designs_or_analyses": "yes", "relation_to_broader_scientific_literature": "The proposed method is sufficiently novel.", "essential_references_not_discussed": "none", "other_strengths_and_weaknesses": "## Strengths\n1. HashAttention consistently outperforms existing sparse attention methods across a wide range of benchmarks. Even when trained on generic data, it achieves 16× sparsity with minimal quality loss on LongBench and RULER benchmarks. With task-specific fine-tuning, sparsity can be pushed to 32× for certain tasks.\n\n2. The use of bit operations for Hamming distance computation bitcount(bitwise_xor(phi(q), phi(k))) is computationally efficient. The approach shows impressive latency improvements in real-world inference systems like GPT-FAST and FlashDecode.\n\n3. At just 32 bits of auxiliary memory per token, HashAttention is significantly more memory-efficient than competitive approaches like Double Sparsity, which require more auxiliary memory to achieve similar quality.\n\n4. The paper presents comprehensive evaluations across multiple datasets, models, and metrics, comparing against a range of state-of-the-art baselines.\n\n## Weaknesses\n\n1. Unlike some heuristic-based approaches, HashAttention requires training on task data. While this enables better performance, it introduces additional complexity for deployment. However, this is a minor point.\n\n2. The paper could more thoroughly discuss how hyperparameters like bits per embedding are selected and their impact on performance. It would help to include code in the appendix.\n\n3. The authors acknowledge that for shorter contexts (< 8K tokens), the overhead of computing bit signatures can outweigh the benefits of sparse attention. This limitation should be more prominently discussed - does this overhead become more prominent with quantized KV cache?\n\n4. The method formulation could be cleaner. For example: z = sigmoid(FF(x)), phi(x) = (z.round() - z).detach() + z.", "other_comments_or_suggestions": ".", "questions_for_authors": "1. Can you include hyperparameters? How much training was required for learning on OpenWebText? Are the LLM weights kept frozen during the learning of hash functions?\n\n2. What's the performance overhead of computing the Hamming distance across all tokens, and at what point does this become a bottleneck compared to full attention computation?\n\n3. Recent papers report that different layers/heads can benefit from different sparsity levels (e.g., middle layers allow higher sparsity). Have you explored variable sparsity across different layers/heads?\n\n4. How does HashAttention perform with quantized KV cache implementations? Does the additional auxiliary memory requirement become more significant in these settings?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper presents a method for long-context attention. It is trained to learn an embedding of the key-values and queries into Hamming space. It then imposes sparsity on the attention coefficients by searching for relevant tokens with a max-inner-product search over the Hamming embeddings and attending only to the top-k search results. The method's design enables a GPU-friendly implementation. Intuition is drawn from cosine similarity search with LSH. 


Reviewers were unanimously positive about the comprehensive experiments and empirical improvement. The rebuttals contained useful expansions and empirical results which the authors are encouraged to include in the final version.
Decision: Accept (poster)