Spotlight Attention: Towards Efficient LLM Generation via Non-linear Hashing-based KV Cache Retrieval
Authors: Wenhao Li
Yuxin Zhang
Gen Luo
Haiyuan Wan
ZiYang Gong
Fei Chao
Rongrong Ji
Venue: NeurIPS 2025 poster ()
Keywords: KV Cache Compression
Efficient Attention
Primary area: deep_learning
TL;DR: KV cache retrieval for large language models using nonlinear hashing function.

Abstract: Reducing the key-value (KV) cache burden in Large Language Models (LLMs) significantly accelerates inference. Dynamically selecting critical KV caches during decoding helps maintain performance. Existing methods use random linear hashing to identify important tokens, but this approach is inefficient due to the orthogonal distribution of queries and keys within two narrow cones in LLMs. We introduce Spotlight Attention, a novel method that employs non-linear hashing functions to optimize the embedding distribution of queries and keys, enhancing coding efficiency and robustness. We also developed a lightweight, stable training framework using a Bradley-Terry ranking-based loss, enabling optimization of the non-linear hashing module on GPUs with 16GB memory in 8 hours. Experimental results show that Spotlight Attention drastically improves retrieval precision while shortening the length of the hash code at least 5$\times$ compared to traditional linear hashing. Finally, we exploit the computational advantages of bitwise operations by implementing specialized CUDA kernels, achieving hashing retrieval for 512K tokens in under 100$\mu$s on a single A100 GPU, with end-to-end throughput up to 3$\times$ higher than vanilla decoding.

Spotlight Attention: Towards Efficient LLM Generation via Non-linear Hashing-based KV Cache
Retrieval
Wenhao Li Xiamen University
Yuxin Zhang Xiamen University
Gen Luo Shanghai Al Laboratory
Haiyuan Wan
Shanghai Al Laboratory
Ziyang Gong Shanghai Jiao Tong University
Fei Chao Xiamen University
Rongrong Ji Xiamen University
Abstract
Reducing the key-value (KV) cache burden in Large Language Models (LLMs) significantly accelerates inference. Dynamically selecting critical KV caches during decoding helps maintain performance. Existing methods use random linear hashing to identify important tokens, but this approach is inefficient due to the orthogonal distribution of queries and keys within two narrow cones in LLMs. We introduce Spotlight Attention, novel method that employs non-linear hashing functions to optimize the embedding distribution of queries and keys, enhancing coding efficiency and robustness. We also developed a lightweight, stable training framework using a Bradley-Teree ranking-basee loss, enabling optimization of the non-linear hashing module GPUs with 16GB memory in 8 hours. Experimental results show that Spotlight Attention drastically improves retrieval precision while shortening the length of the hash code at least 5x compared to traditional linear hashing. Finally, we exploit the computational advantages of bitwise operations by implementing specialized CUDA kernels, achieving hashing retrieval for 512K tokens under 100μss on a single 1100 GPU, with end-to-end throughput up to 3x higher than vanilla decoding. All the training and evaluation stuff can be found at https://gichub..oo// hww.sh/wwwwwwwaolooooossoooooooo
Introduction
Large Language Models (LLMs) are propelling groundbreaking advancements various natural language tasks, significantly enhancing applications such content creation and chat assistance. Generally, the inference process of I.LMs can divided into (1) the pre-filling phase calculates the key-value (KV) cache for input tokens in the prompt prior to autoregressive generation, and (2) the decoding phase ruto-regressively generates tokens, producing one token per forward pass based on the KV cache. Among them, the decoding phase serves the primary inference bottleneck due to the frequent exchanges between on-board and on-chip memory for model parameters and KV cache, which limits GPU scalability [3] more than the pre-filling phase that processes input prompts in parallel. For example, deploying LLaMA2-T7 [2]] on an A100 GPU for a single request achieves nearly 100% GPU utilization during the pre- pre-filling phase but drops to below 10% on average during decoding, which largely restrains the inference efficiency of LLMs.
*Corresponding author
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

IV Cocha
04. Antwee decopparry from Roon Model
Chode Tap
Drade Top-t
Self- Aftentinn
3.5
Uner Heriing
zagore
Tap KV Cocke
3.65
B0م0 XX/XXX//
3.62
3.57
Vatical Poltorn
Vertzal Slash Poterr
Topri Indiss
Calinated Topek
Esimaded Tepk
Self-Amention
Sie Soorch (Vey Foo)
0.44
Yealin
334
Hash Coies
Slash
220
440 1109
Toleen
Figure Overview. (Left) Architecture. Comparison of our Spotlight Attention versus normal attention. Spotlight Attention adds an additional hash code- based retrieval mechanism for each layer. (Middle) Performance. Spotlight attention achieves the most accurate retrieval and generates the closest response compared to the original model on QA datasets. (Right) Visualization. For arbitrarily complex attention patterns, our method estimates the top-k sequences well, with an average correctness rate of more than half for different models.
To alleviate this inference bottleneck, extensive research has focused on heuristically eliminating the KV cache burden based on attention scores [25]
23 17 While effective for short sequences,
such irreversible removal of KV cache can significantly degrade performance on long-sequence tasks, especially in Needle-in-a- in----aysttacck scenarios [14]. To explain, tokens initially considered
unimportant and removed might later attain higher attention scores during the prolonged decoding phase, which is crucial for output quality (20 15). To overcome this limitation, recent works have turned to retaining all KV cache tokens and dynamically selecting important tokens for computation during decoding (20) 10), which is the focus of this paper.
Despite the convineing performance of such on-the-fi fly KV cache selection, how to effectively pick up those important tokens remains challenging. As pioneering effort, Quest [20] selects tokens by matching queries and keys in a block-wise manner. While effective, such coarse-grained selection hardly guarantees precise localization of important tokens. MagicPIG [10] advances Quest by implementing token-level cache retrieval, specifically utilizing Local Sensitive Hashing (LSH) to encode queries and keys into hash codes and pinpointing the best matches as the selected tokens. However, as depicted in Figure[aa the efficacy of such linear hashing heavily depends on the hash code length, e.g., hash code length of 1,024 bits per query/key is necessitated to achieve promising token retrieval. Considering the already substantial size of the KV cache, storing these lengthy hash codes markedly impairs deployment efficiency. Moreover, employing linear projection with large output dimension for key processing incurs significant computational overhead.
Delving deeper, prior work [10] has discovered that the queries and keys within LLMs typically form nearly orthogonal cones within the embedding space, as depicted in Figure Given the truth that linear hashing function partitions the embedding space using random hyperplanes, such orthogonal distribution of queries and keys barely lead to satisfying encoding quality, which can even result in collapse of the hashing outcomes, i.e., identical codes for all queries and keys, as shown in Fi Therefore, extremely long hash codes are necessary to mine meaningful information and accurately match essential tokens. MagicPIG attempted to mitigate this issue by normalizing the keys before retrieval. However, this approach remains suboptimal as it overlooks the query distribution and introduces bias to the retrieval process.
To address the aforementioned limitation, we propose Spotlight Attention, a novel method that replaces random hyperplanes with curved surfaces for space partitioning via a non-linear MLP hashing function. As depicted in Figure[2d] this non-linearity can better fit skewed distributions, thereby improving code quality. In particular, we utilize the Bradley-Terry ranking objective [7] to optimize the non-linear MLP layer, wherein the learning target involves minimizing the difference between the estimated top-k indices and the ground truth top-k indices obtained via the vanilla attention scores. This learning process exceptionally efficient, with the LLM backbone remaining frozen and requiring only minimal amount of calibration data. As result, the optimized non-linear

---

610
Dinribution
Perplesity
3.50
Our MLP Heshing
Query Diaribution
120
256
512
1024
2043
Number per key
(a)
(b)
(c)
(d)
Figure 2: Motivation. (a) The empirical evaluation shows that upgrading the hashing function from linear to MLP can bring a huge improvement, (b) this because query and key are usually distributed in two small cones the space [T]. (c) In this situation, is difficult for the space be uniformly partitioned by linear boundaries, (d) which can be well solved by using an MLP hashing function.
hashing function can match the performance of linear hashing while using 5x shorter hash codes, achieving higher efficiency than MagicPIG. We further implemented CUDA kernels for the hash code processing, including bit-packing and bitwise NXOR GEMM operators, achieving significant latency reductions in practice. For example, our method achieves up to 3x increase in Qwen2.5-7B [24] inference throughput for both 32K and 128K sequences, with only 2% performance degradation on the LLaMA3 [13 series and no loss Qwen2.5 [24 series.
Our contributions are threefold:
We propose Spotlight Attention for accelerating LLM inference, which employs non-linear hashing function to encode and match queries and key values within LLMss, thereby effi- ciently selecting critical KV cache for model inference.
We develop a lightweight and robust training framework based on the Bradley-Terry ranking objective, which effectively optimizes the non-linear hashing function using only a small amount of calibration data.
Extensive experiments demonstrate that Spotlight Attention can drastically reduce LLM inference latency while maintaining the strongest performance retention in comparison with state-of sitte-of-the-aad the-art art methods.
2
Related Work
This section covers the spectrum of studies on LLM KV cache pruning that are closely related to our work, which we heuristically categorize into static pruning, dynamic pruning with permanent eviction, and dynamic pruning without permanent eviction.
Static KV cache pruning. These methods compress the KV cache once after the pre-filling phase, using the compressed cache for subsequent decoding. For example, FastGen [12] introduces pattern- aware approach by identifying five fundamental attention structures and applying targeted selection strategies. SnapKV [16] further simplifies FastGen by focusing solely on retrieving tokens based on their importance scores, showing that only subset of prompt tokens carry critical information for response generation and retain their significance during the whole decoding phase. However, without pruning during decoding, these methods are primarily suited for scenarios with long prompts and relatively short responses.
Dynamic pruning with permanent eviction. This category of methods performs dynamic KV cache pruning during the decoding phase, permanently removing pruned KV cache tokens from memory. For example, H2O [25] leverages cumulative attention scores retain high-impact tokens. NACL [9] identifies fundamental limitation in H2O, namely their dependence on potentially biased local attention statisties. To overcome this issue, they develop an alternative approach, implementing diversified random eviction strategy. Keyformer (2) highlights that token removal distorts the underlying softmax probability distribution. Considering the pivotal role of softmax distributions in token significance evaluation, they incorporate regularization techniques to mitigate these distributional perturbations. Unlike static KV cache selection, these methods enable dynamic pruning during decoding, making them better suited for tasks requiring extensive generation. However, they assume that critical information is concentrated in small subset KV cache tokens, a condition that does not always hold. As MagicPIG [10] points out, token importance vary significantly

---

KV Hash Codes
KV Hash Codos
Ahush
Subtract
Extimatted Score
R
H
(asl aa.k-.27400
10 1
The námoacc
Estimated
Score
877
23
P
May ala bo.
Ilustration
Ground Truth
Top-K
Top-K
viic' fover opantor
Non-Top
Fours
Σ
KV Cache
KV Coche
Layer-wise Trainable
Supervising Signal
Layer-wise Trainable
Supervising Signal
Capacity Missllcatton
Capacity Misallccttion
Figure 3: Optimization. (Left) Reconstruction Loss. This loss minimizes the MSE between the estimated and ground-truth attention scores. It has two main drawbacks. First, it is highly sensitive to score magnitudes and prone to outliers. Second, wastes most of the hashing function's 's capacity on preserving order within the top-k and non-top-k sets. (Right) Our Ranking Loss. Our loss adopts the BradIcy-Terry ranking objective, which is robust to score magnitude and outliers, and provides supervision focused solely on distinguishing between top-k and non-top-k sets.
across tasks, leading to premature eviction of tokens before they are needed. For example, H2O may fail to answer questions like a b, d, a due to forgetting earlier facts.
Dynamic pruning without permanent eviction. The limited applicability of permanent token eviction methods has led to shift toward non-permanent eviction approaches. These methods assume the importance of KV cache tokens varies with each query, requiring importance estimation every decoding step. Instead of permanently evicting unimportant tokens, they exclude them from attention calculations for that step only. While this improves accuracy, it demands frequent importance estimation, unlike permanent eviction methods that prune tokens in batches after many steps. Research has therefore focused on optimizing the efficiency and accuracy of these estimations. Quest 20 groups KV cache tokens into blocks, estimating block importance via dot product between queries and block representations derived from the minimum and maximum key values. Although efficient, this approach suffers from internal fragmentation, as entire blocks are processed even if only few tokens important. MagicPIG [10] eliminates this issue by mapping queries and keys to hash codes for token-level retrieval via Hamming distance, avoiding fragmentation but reducing efficiency. Building on MagicPIG, our method significantly shortens hash code, drastically reducing computation while preserving accuracy.
3 Methodology
3.1 Preliminary
Attention computing. We first present the basic preliminaries for attention computation and KV cache pruning during the decoding phase of LLMs. define the query, key, and value inputs to the attention module Q, K, Rlxd where d is the embedding dimension. We use to denote concatenation, and Keache Vanche E Raxd represent the key-value cache generated during the fre-filling phase and previous decoding steps. With these definitions, the standard attention is calculated as follows:
f(Q, Konhee Kenche) + QKT) (Vanche + V),
A softmax
d
where xX) is iroducp.
KV cache pruning. As the decoding sequence length increases, the size of the KV cache {K, VT,.V)cnchee cache can grow exceedingly large, creating an L.M inference bottleneck. KV cache pruning that selectively preserves only essential portions of the cache for computation serves efficient way alleviate

---

Table 1: KV retrieval accuracy. Measured by loUt PPL.| changes before and after training. (1) Training is essential for efficient MLP hashing. (2) Limited improvement from training linear hashing highlights the necessity of MLP hashing. See Appendix B.1 for per-head IoU scores.
LSH Top-2%
MLP Hashing Top-2% % (ours)
Method
Original
Oracle Top-2%
Before
After
Before
After
LLLMMA-7B
55558
1.00 /.69
5.87
0.20/ 5.84
0.05 20.31
0.41 / 5.72
LLaMA2-7B-CCr
/7.10
1.00/6.87
0.17/7.34
0.19/
0.05/21.34
0.42/
LLaMA3-8B
-/ 6.45
1.00 / 6.63
0.15/7.12
0.18/ 7.07
0.07 / 148.2
0.34 / 6.69
Qwen2.5-7B
-77.17
7.207
0.13/8881
0.167 8.73
0.09 22.07
0.35/ 7.31
this problem. Given desired cache budget K, it first identifies the indices I of top- K important tokens and then extracts a subset of the KV cache {K, V V. subset for attention computation as
{K, V) subset gather( V) ceacee
As previously discussed, existing methods for pruning the KV cache either permanently eliminate cache entries in the set I [25 [6] or retain all caches but dynamically determine I during the
decoding process 120 10 In this paper, we focus on the latter due to its superior performance
preservation.
3.2 Revisiting Token-level Cache Retrieval
Token-level cache retrieval refers to dynamically selecting cached entries at the granularity of individual tokens [10].
Oracle top-k retrieval. We conducted preliminary experiment assess the upper-bound per- formance, using full-precision attention scores f(Q, Keache) to select key tokens 2 Surprisingly, by pruning layers beyond the first two, we discarded up to 98% of the KV cache with only a 0.1 per- plexity increase on PG19, revealing significant untapped potential. This suggests that top-k retrieval enables near-lossless compression, contrasting sharply with prior findings [10].
LSH top-k retrieval. Although oracle attention scores accurately identify key tokens, computing f is impractical for real realdworld applications. To address this, MagicPIG approximates f with f using Locality-Sensii Sensitive Hashing (LSH) efficiently retrieve critical KV entries. Specifically, a linear hash function H. computes:
x(X,X'= H(X) @ H(X),,
(3)
where denotes matrix multiplication-IIn like operation, substituting floating- point multiplication with NXOR. The indices I of the top-k largest values in f(Q, Kenche) are used to retrieve the most relevant KV entries.
LSH groups similar vectors into the same bucket with high probability, making it ideal for dense vector spaces. A common variant employs random hyperplanes, where distinct hyperplanes create linear decision boundaries, assigning data on either side to bit-O or bit-1. The bits sequence from all hyperplanes forms the hash code. In practice, these steps can be simplified multiplication followed by sign operation. For a vector E Rdd we apply random projection matrix R E gdxdx to obtain a hash code by taking the sign of the resulting product:
single matrix
H(x) R)..
MLP hashing top-k retrieval. As shown in Figure 2, queries and keys typically lie within two cone- shaped regions in high-dimensional space (10). This distribution causes uneven partitioning in LSH, reducing encoding efficiency. To address this, we propose MLP hashing, a learned non-linear hashing network tailored query and key distributions. This approach enhances hash code information density, enabling effective partitioning of skewed data through non-linear decision boundaries.
Specifically, we replace the projection matrix R in Eq. with a two-layer MLP:
MLP(x) = WZZSU + bi)),
(5)
where W1. by and W2 are learnable parameters. Hash codes are then computed as: H(x) sign(MLP(
(6)
"See Appendix A.2 for the oracle top-k attention pseudocode.
5

---

Table 2: Perplexity comparison with Quest. Perplexity evaluation PG19 (#1), ProofPile (#2), and CodeParrot (#3) datasets. All models truncate inputs to their maximum supported token length. Spotlight Attention achieves performance comparable to Quest with 10x smaller token budget.
LLaMA2-78
LLaMAZ-7B-CLL
LLaMAS3-B
Qwen2.5-78
Method
Configuration
Layers
#1
92
65
41
92
43
81
42
43
#2
#3
Vanilla
0% Pruned
NIA
6.879
4,277
3,679
9.212
5,943
4.756
5.604
3,517
5.219
11:112
3,833
4.951
Oracle Top K
98% Pruned
[0.1]
6.941
4.317
3,729
9.224
5,891
4.795
8.881
3.590
5,383
10,435
3,587
4,733
75% Prined
7.116
3,404
3.854
9.282
5,900
4.579
9.912
4.024
5,593
10.485
4.251
5.482
Quest
87.5 Prased
93.7% Prunedd
[0.11
9.746
7.335
5.775
4.754
4.054
4,573
12,058
9:820
7.440
6.226
3.650
5.054
17.320
12.454
6.749
4.927
8,653
6.646
19,552
13,596
5,471
7.274
7,823
6.213
9% Pruned
15.494
8,578
5.996
18,719
11,083
7.345
27.530
11.631
14.205
31.583
13.398
12.526
80% Prused
6.887
4.278
3.682
9.107
5.360
4.767
5.612
3.519
5.224
9.766
3,465
4.615
Spotlihtt (ours)
90% Pruned
95% Pranol Pruned
[0.1]
6,899
6.908
4.285
4.504
3.689
3,703
9.067
9.058
5.796
5.748
4,753
4.754
8.754
5.651
3.529
3.552
5.285
5.239
9.825
9.785
3,467
3.475
4.627
4,621
18%
7.306
4,364
3.768
9.262
5,770
4,806
8,977
3,621
5.434
9.930
3.497
4.645
3.3 Optimization.
Optimizing the MLP hashing function H to capture the query and key distribution is more critical than the hashing network design and forms the core contribution of this work. We next introduce two intuitive training objectives and explain their limitations in this context.
Language modeling loss. A natural approach to optimize H. is to minimize the language modeling loss directly through few hundred post-traine sott-traniig steps. However, this method has significant limitations. First, it requires full forward and backward passes through the entire LLM, which is computationally costly. More critically, differentiating the top-k operator in Eq. requires complex "soft top-k" approximation, which is challenging to implement.
Reconstruction loss. An alternative approach uses MSE to align f with f. enabling layer-wise optimization and reducing computational cost. However, this method has key limitation: the objective is to select which KV cache entries to retain, not to rank their relative importance. Training with this loss misallocates capacity by prioritizing the ranking of excluded entries, deviating from the core goal, and causing significant performance degradation. See Figure (Lefi) for its limitations.
V Our ranking loss. To address this issue, we adopt Bradley-Terry ranking objective, inspired by RankNet [8) As shown in Figure 3 (Right), during training, we identify top-k and non-top-k index sets based on attention scores. These indices split the estimated scores derived from Hamming
distances of query and key hash codes into sets B and An optimizer then updates the hashing function parameters to ensure each score in B exceeds every score in C:
Lrank
log (sigmoid (3(B, C2) - a)),
ki(n k)
3,7
where / and are positive constants used to amplify the separation between and C, facilitating convergence. The core of this loss design lies in filtering out supervising signals related to internal ranking within B and C, effectively addressing the issue of capacity misallocation. We provide the pseudo-code for our ranking loss in Appendix A.3 to aid readers familiar with code.
Make hashing function differentiable. After computing the loss, errors can be backpropagated to the MLP hashing function. However, the sign function's 's non-differentiability blocks gradient flow. To resolve this, we substitute the sign function with a soft sign function during training:
softsign(z) 1 7|20
(8)
where 7 E R is hyperparameter controlling the extent of smoothing. This soft sign function used only during training. In inference, the non-differentiable sign function is reinstated.

---

Vanilla LLAMAJ-SO
MMagePIGG
Spoclight tup-2% (ours)
NIAH results. Using LLaMA3-8B I3 as the base model, we compared the retrieval accuracy of MagiePIG with method. Our approach, which relies solely on hash code-based retrieval without local windows or sink tokens, achieves comparable response accuracy.
Figure 4:
Table 3: Perplexity versus MagicPIG. Compari- of perplexity PG19 (#1). ProofPile (#2), and CodeParrot (#3). Due to the timecconumiing evalua- tion process of MagicPIG, we sampled only 10 data points from each dataset for testing.
QA response fidelity. On Long- Bench, output fidelity (measured by Rouge- between compressed and vanilla model outputs) shows our method achieves perfor- mance closest to the vanilla model.
Table 4:
Method
Frone
Local (54) Siit (4)
Configuraionn
Reviev (29)
LLaMA3-0
45
91
Method
Local
Sink
Configuration
Retieve
Frozen
Similarity
Onuele
(0.1]
2.28
Vanilla
x
1.00
13.85
1.4
3.91
Cracke Top-K
163
0.66
MogoPtG [U16]
43.13
8.
6:27
10:55
4.57
LSH Top-K
Quest
x
1024
163
[0.1]
0.56
0.37
NaN
NoN
NoN
MagicPIC Spotlight (ours)
Dynumic 163
Spetlisht
9.99
2.91
1.30
100
Vanila
Spotlight (1024)
400
50
MagicPI6
1000
Magichic
NarrathoQA
MattFiilt244 wn!
Qeaper
MultFieldDA..
HotootOA
2
Musi
Duflander
GovReport
Qutium
MultiNdwp
VESUM
TREC
Trivia24
SAMSum
PassapeCCunt
LSIT
Panugeflstrieral
Prasspeletr oh
RepoBerchh
LCC
1004
Figure 5: Downstream QA tasks. (Left) Relative score of each method compared the vanilla baseline; each point denotes a subtask. (Right) Absolute score comparison.
Experimentation
Our evaluation spans multiple dimensions. (1) We first assess the errors introduced by retrieval and sparsification.. measured by retrieval accuracy and perplexity score. (2) We then evaluate
long-context key information retrieval using the Needle-in-a-Haystace [14] benchmark. (3) Next, we evaluated downstream QA tasks on LongBench [6. comparing response similarity between compressed and original models using Rouge-L. The key insight from this Rouge-L, comparison straightforward: higher-qualit compression produces more concise outputs. (4) We also measure end-to-end throughput gains and the execution efficiency of our CUDA ops.
To evaluate generalization, we tested the performance of the Qwen2.5 series 24 across various model sizes and LLaMAA-8B [13 across diverse training corpora. Additionally, we conducted ablation studies to assess the impact of loss functions and attention estimation methods, with details
provided in Appendix B.5 and B.6
Experimental setups. We employ LLaMA3-8B 13 LLaMA2-7B. LLaMA2-7B-Chi 21]. and Qwen2.5 models (1.5B, 7B, 14B) 24] as base models. The bascline methods compared include (1) oracle top-k retrieval, (2) linear LSH top-k; (3) Quest, and (4) MagicPIG. Their implementation
details provided in Appendix A.2
Our MLP hashing function employs 128-dimensional input, intermediate, and output layers, with a distinct MLP for each head in every layer, producing 128-bit hash code-much shorter than MagicPIG's minimum of 720 bits. Only the hashing functions are trainable. Training data consists of 8,192 samples, evenly drawn from the Book and Arxiv datasets [22]
7

---

SBK Cuatest
12NK Contest
Hal Cuoc (MBB)
Karad Lasena (ia)
Thrugggu polones
100
(ours)
60
Socdight
70000
KV Dache
Spoiligtt (ours)
3000
Our Sim Search (CUBA)
FashArttenton 2,5.8
Megicht
Nagitii
Megichii
Top 2% Guthering
20
30
35000
1520
10c G4C 250k
10k
646
2004
Botck Site
Barth Size
Cantest Length
Sixe 1
Booch Site =2
Throughus (okens)
16
Bisclins
Symbol Def.
Inforence Sub-Step
Complexity
Spotight (our))
Urtwse encoding
((kd)
In head count
20
Smirntty search
Olykd)
embed
Top-k gathering
6kyh log k) o(thd)
(ours)
51.6
32K
13BE
Contest Lengh
Castral
Figure 6: Efficiency. (Upper Left) End-to-end throughput comparison with fixed context length across varying batch sizes. (Bottom Left) End-to-end throughput comparison with fixed batch size across different context lengths. (Upper Right) Hash code size comparison between MagicPIG and our method, alongside the execution latency of the two most computationally intensive operations in our method. (Bottom Right) Complexity comparison of different computational steps.
To improve efficiency, hidden states for all layers are precomputed and stored, enabling independent layer-wise training without joint fine-tuning. Training uses ? = 64, learning rate of x 10-3,
1, and 3, for one epoch. Additional details are provided in Appendix A.1 A The pruning rate remains fixed at 98% during training, irrespective of evaluation settings.
4.1 Main Results
KV retrieval accuracy. This experiment compares linear LSH and MLP hashing for KV retrieval. We used the first sample from PG19 as test data and assessed KV retrieval accuracy with the average IoU score across all heads and layers. IoU is computed as the intersection of the top-k indices retrieved by the algorithm and the oracle top-k indices, divided by their union.
For LSH, initialized the projection matrix via QR decomposition (see Appendix A.S). For MLP hashing, parameters were initialized with random Gaussian distributions. Results in Table [ show that LSH performs best without training but shows minimal improvement post-training. Conversely, MLP hashing markedly improves retrieval accuracy after training, achieving the highest performance.
Language modeling. We evaluated three language modeling benchmarks PG119 [9] ProofPile (5),
and CodeParrot [1]-with 100, 79, and 100 samples, respectively, using perplexity to detect minor errors from sparsification. Additional experimental details are in Appendix A.5
Results in Table show our method outperforms Quest with tenfold reduction in token budget for most models. Comparisons with MagicPIG (Table indicate that MagicPIG depends heavily on local windows and sink tokens, failing without them, while our method autonomously identifies these elements, demonstrating significantly higher retrieval accuracy.
MagicPIG excels retrieving facts in scenarios like Banall-in-a-Haystee [14], its contribution should not be dismissed based solely on its limitations in this language modeling test.
-aalo-iin---aattckk. We use the offline evaluation version of NIAH [.. which differs from ChaiGPT scoring by using the Rouge score to measure output accuracy. We utilize LLaMA3-8B [3] as the base model, evaluating context lengths ranging from 256 to 8192, with a step size of 256 for testing. The Needle dataset is highly prompt-sensitive, we provide the needle, retricval question, and haystack context in Appendix(A.4 Additionally, to reduce output variability, all evaluated methods use greedy search. As shown in Figure Spotlight Attention achieves performance on par with the original model.
Downstream QA tasks. We evaluated the performance of various compression methods on Long- Bench (6] subtasks. Quest employs a token budget of 1 1.024, while MagicPIG uses 4 sink tokens, a 64-token local window, and I, 1,500-bit hash code per key, following their default configurations. Our method uses a token budget of 163. All experimental results reported in the main text use LLaMA3-8B as the base model, scores of other models are provided Appendix B.4, and more experimental setup details are in Appendix A.6 Figure 5 (Right) compares absolute scores across all
8

---

subtasks. More importantly, Figure 5 (Left) shows relative scores, revealing that our method's 's scores closely align with those of the vanilla model.
We also assessed output fidelity using Rouge-I. to measure similarity between the vanilla model's outputs and those of these methods. Results are presented in T Our method's s outputs the most similar to the vanilla model, with the longest consecutive subsequence match approaching 60%.
In contrast, Quest requires retrieving six times more tokens achieve comparable similarity. These findings hold for nearly all subtasks, with detailed scores of each subtask provided in Appendix B.2
4.2 Efficiency
All efficiency experiments were performed on Qwen2.5-7B [24] using eight A100 GPUs. For enhanced flexibility, experiments utilized the HuggingFace Transformers framework, optimized with pipeline parallelization and KV cache pre-allocation to boost throughput.
End-to-end throughput evaluation. To evaluate model throughput at extended context lengths (e.g.. 2M tokens), expanded positional encoding. disregarding output quality. We selected the first sample from the PG19 test set [19], repeating to reach a 2M-roken context. GPU execution time was measured using CUDA events. We generated eight consecutive tokens, computed throughput by dividing by generation time, and averaged three runs per data point. Results in Figure ((e() show that our approach consistently delivers throughput gains, especially at larger input scales.
Kernel evaluation. We implemented CUDA kernels for both bit-packing and NXOR GEMM. Bit- packing compresses 32 Torch boolean values into a single unsigned 32-bit integer, as detailed in Appendix A.7 For NXOR GEMM, we utilized the standard library's 's popcount to count bit- 1s in each NXOR result. For top-k gathering and sparse attention, we employed Torch and FlashAttention implementations, respectively. As shown in (Upur Right), compared to dense vectors, our hashing-bate based similarity search significantly reduces storage and computation, enabling bit bit-packing and similarity search within 100 for context lengths up to 512K.
Table 5: Ablation on model size. Various Qwen2.5 model sizes, augmented with Spotlight Attention, demonstrated better perplexity across diverse language modeling tasks.
Table 6: Ablation on training tasks. Along- side the standard ArXiv+Books training data, we trained models on the C4 and GitHub Code datasets, achieving comparable perplexity (PPL) results.
Model
Method
PG19
Math
Code
Qwea2.5-1.5B
13.828
5,081
LLaMA3-BB
Training Corpus
PG19
Math
Code
Qwen2.5:7B
Spotlight Top-2% (ours)
Vanilla
13.510
11.112
4.143
3,833
4,951
5:164
Oracle Top-28
8.881
3,590
5.353
Spotlight Top-2% (ours)
9.950
3,497
4.645
Spetlight Top-2% (ours)
Book (detaull)
5,434
Qwc42.5-148
Spotlight Top-29 (ours)
Vanilla
8.416
8.261
3,230
3.196
4.440
4.472
Github Code
8,891
8.958
3.611
3,631
5,393
5.417
4.3
Ablations
In the main text, we present ablation studies on model size and training tasks only. Additional ablation experiments, including loss functions and attention estimation methods, are detailed in Appendix B.5 and B.6 respectively.
All ablation experiments use language modeling perplexity as the evaluation metric. We assessed performance on PG19 (1)) Proof-Pile (Math) (5) and CodeParrot (Code) [18 with sample sizes of 100, 79, and 100, respectively.
Model size. We compare Qwen2.5 24 models of 1.5B, 7B, and 14B parameters, all trained with the standard recipe. As shown in Table Spotlight Attention achieves consistently strong performance across different model sizes.
Training tasks. To validate our method's applicability across diverse tasks, we trained models on the GitHub Code [18] and C4 [18] datasets, in addition to the ArXiv and Books [22) datasets used previously. As presented in Table training with GitHub Code or C4 unexpectedly outperformed the ArXiv+Books combination, demonstrating the robust adaptability of our training framework. However, due to the cropllccoopletion of the main results, we did not re-evaluate all benchmarks with these checkpoints despite their superior performance.

---

Conclusion and Limitation
We introduce Spollight Attention, advancement over Quest and MagicPIG, incorporating non- linear hashing function and an optimized framework. This approach addresses the underfitting of MagicPIG's linear hashing while significantly reducing hash code length. Spotlight Attention performs well on downstream tasks but has limitations. The IoU remains around 40% despite non-linear hashing, indicating potential for improvement.
Acknowledgments
This work supported by the National Science Fund for Distinguished Young Scholars NN..62255033, the National Natural Science Foundation China (No. U21B2037, No. U22B2051, No. U23A20383, No. 62176222, No. 62176223, No. 62176226, No. 62072386, No. 62072387, No. 62072389, No. 62002305 and No. 62272401), the Natural Science Foundation Fujian Province of China (No. 2021J06003, No.2022J06001), the National Natural Science Foundation of China No.624B21 19, and the China Postdoctoral Science Foundation N0.8220250384)
10

---

References
[1] Github repository: http://tsst.eccdeeeiaaaassaaaile 2023.
[2] M. Adnan, A. Arunkumar, G. Jain, P. Nair, I. Soloveychik, and P. Kamath. Keyformer: Kv cache reduction through key tokens selection for efficient generative inference. MISys, 2024.
[3] A. Agrawal, A. Panwar, J. Mohan, N. Kwatra, B. S. Gulavani, and R. Ramjee. Sarathi: Efficient Ilm inference by piggybacking decodes with chunked prefills. arXiv, 2023.
[4] A. Andoni, P. Indyk. T. Laarhoven, I. Razenshteyn, and L. Schmidt. Practical and optimal Ish for angular distance. In NeurIPS, 2015.
[5] Azerbayev, E. Ayers, and B. Piotrowski. Github repository: hoskison-ccenterproofff 2022.
[6] Y. Bai, X. Lv, J. Zhang, H. Lyu, Tang, Z. Huang, Z. Du, X. Liu, A. Zeng, L. Hou, et al. LongBench: A bilingual, multitask benchmark for long context understanding. In ACL, 2024.
[7] R. A. Bradley and M. E. Terry. Rank analysis of incomplete block designs: I. the method of paired comparisons. Biometrika, 1952.
[8] Z. Cao, T. Qin, T.-Y. Liu, M.-F. Tsai, and H. Li. Learning to rank: from pairwise approach to listwise approach. In ICML., 2007.
[9] Y. Chen, G. Wang, J. Shang, Cui, Z. Zhang, T. Liu, S. Wang, Y. Sun, al. NACL: A general and effective KV cache eviction framework for LLM at inference time. In NACL, 2024.
[10] Z. Chen, R. Sadhukhan, Z. Ye, Y. Zhou, J. Zhang. N. Nolte, Y. Tian, M. Douze, L. Bottou, Z. Jia, and B. Chen. MagicPIG: LSH sampling for efficient LLM generation. In NeurIPS, 2024.
[11] L. Gao, Tow, B. Abbasi, S. Biderman, Black, A. DiPofi, Foster, al. Github repository: 2024.
[12] S. Ge, Y. Zhang, L. Liu, M. Zhang, J. Han, and J. Gao. Model tells you what to discard: Adaptive KV cache compression for LLMs. In ICLR, 2024.
[13] A. Grattafiori, A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al-Dahle, A. Letman, et al. The Ilama 3 herd of models. arXiv, 2024.
[14] G. Kamradt. Github repository:
2023.
[15] H. Li, Y. Li, A. Tian, Tang, Z. Xu, X. Chen, N. Hu, W. Dong, Q. Li, and L. Chen. A survey on large language model acceleration based on kv cache management. arXiv, 2025.
[16] Y. Li, Y. Huang, B. Yang, B. Venkitesh, A. Locatelli, H. Ye, T. Cai, P. Lewis, and D. Chen. SnapKV: V: LLM knows what you are looking for before generation. NeuriPS, 2024.
[17] M. Oren, M. Hassid, N. Yarden, Y. Adi, and R. Schwartz, Transformers are multi-state rons. arXiv, 2024.
[18] Z. Peitian. Huggingface dataset: namespace-p//ong-llnnee Ilm--aaa-an 2024.
[19] J. W. Rae, A. Potapenko, S. M. Jayakumar, C. Hillier, and T. P. Lillicrap. Compressive transformers for long-range sequence modelling. In ICLR, 2020.
[20] J. Tang, Y. Zhao, K. Zhu, G. Xiao, B. Kasikci, and S. Han. QUEST: Query-Aware Sparsity for Efficient Long-Context LLM Inference. In ICML, 2024.
[21] H. Touvron, L. Martin, K. Stone, T. Scialom, et al. Llama Open foundation and fine-tuned chat models. arXiv, 2023.
[22] M. Weber, D. Y. Fu, Q. G. Anthony, Y. Oren, S. Adams, A. Alexandrov, X. Lyu, et al. Redpajama: an open dataset for training large language models. In NeurlPS, 2024.
[23] G. Xiao, Y. Tian, B. Chen, S. Han, and M. Lewis. Efficient streaming language models with attention sinks. In ICLR, 2024.

---

[24] A. Yang, B. Yang, B. Zhang. B. Hui, B. Zheng, B. Yu, C. Li, D. Liu, F. Huang, et al. Qwen2.5 technical report. arXiv, 2025.
[25] Z. Zhang, Sheng, T. Zhou, T. Chen, L. Zheng, R. Cai, Z. Song, Y. Tian, C. Re, C. Barrett, Z. Wang, and B. Chen. H2o: Heavy-hitter oracle for efficient generative inference of large language models. NeurlPS, 2023.
12

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justificcton right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]", ", it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: The experimental section validates all ideas introduced in the Introduction, confirming the accuracy of all stated contributions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
13

---

Justification: In the Conclusion and Limitations section, we outline the limitations of this study.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-ten system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification: Our experiments do not contain theoretical results, only some empirical conclusions.
Guidelines:
The answer NA means that the paper does not include theoretical results. the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: The Experimental Setup section provides detailed descriptions of the experi- ments, with comprehensive details included in the Appendix A.
14

---

Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel urchitecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification:
We provide open access to data and code in the Abstract.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
hhipss.c//ublic/uuudds httpee///oooeSeeeeeoooeoooe
www.obiiiiiiioooooioiiooo for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
15

---

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All experimental details are provided in the Experiment Setup chapter, with implementation details included in the Appendix A.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification: For experiments involving randomness (Sec 4.2), such as the efficiency tests, we conducted multiple sampling iterations and reported the mean values.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We give the relevant information in Sec 4.2.
Guidelines:
16

---

The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics http:ppss/eaaraps...ppppp//////s httpppublcc/bhiiiccc..eeeelelll
Answer: [Yes]
Justification: Our research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer:
[NA]
Justification:
Our approach does not involve societal impact.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications,
authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
17

---

Answer: [NA]
Justification:
Our work does not include such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging. and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All assets used are well-known, properly referenced, and their information is readily accessible.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the bataasesssppprsstttt paperswi theode. com/datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is available online, the authors are encouraged reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [NA]
Justification: Our paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
18

---

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: Our study focuses on open-source LLMs, but LLMs are not involved in the research methodology.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our L.LM policy httpp://eeerrpppcccoo/oofeeerere for what should or should not be described.
19

---

A Implementation Details
A.1 Training Setup
Table T summarizes the training configuration. For further details, including low-level specifics and weight files, refer to our open-source code.
Table Detailed training configuration.
Learning
Precision
Up
Annealing
Clipping
K,192
liscar
Optimiaer
Optimizer
Weight
Arsi Sumples
Book Samples
LLMMA2
LLaMA3NQee
Trunc
adamw
0.9
0.96
book
4.095
4,096
4,096
$,191
right
A.2 Baseline Methods
Oracle top-k. As introduced in Section this baseline employs original attention to identify top-k indices, serving as theoretical upper bound on performance. provide pseudo-code below to illustrate concrete implementation of the oracle top-k algorithm.
LSH top-k. This method uses training -eeee angular LSH 4 with the same hash code length as Spotlight Attention, providing reference for MagicPIG's linear hashing under our framework. Initialization details for angular LSH are Appendix(A.8
Quest. 20 We adopt Quest's official implementation with default hyperparameters, modifying only the token budget. Beyond the default 1024-token budget, test ultra-low budgets to align with Spotlight Attention. For instance, for LLaMA2-7B and LLaMA3-8B, we use budgets of 128 and 256 tokens, compared Spotlight Attention's 's 81 and 163 tokens.
MagicPIG. [TO] We adopt MagicPIG's official implementation with default hyperparameters (K 13, L 100), retaining 64 local tokens and 4 initial tokens. For all experiments except the efficiency test, we use the official Python evaluation code to conduct the experiments.
A.3 Ranking Loss
We provide the ranking loss calculation in the following pseudo-code To support longer sequence lengths during training, we employ three optimization techniques: (1) Random query selection, where only queries specified by query_ index are optimized, rather than all queries. (2) Random top-k selection, where max max_top top randomly sampled from the top-k set for optimization. (3) Random non-top-k selection, where max_ oth is randomly sampled from the non-top-k set for optimization. These techniques enhance training efficiency in long-context scenarios.
def ranking_loww
draft_atts, quwry_
mas_
ntt.oth.
maskaut.
beta: float 1.0,
alpha: float 0.03:
10
(0nsor(0, dtype-torch. flost]2) criterion torch m. whw.Biitaaaiitaa.nss...
13
14
15
true_attn.shappe shageel
forch.triuutaree oowx (onun,vv, nux_ky), dtypestorch.ban davineetrru,etttn.dee diagonul-1) Dione,
:]
17
index is Hone:
nack..... indey, :]
19
true_ate torch masked_ fill (true, attn, attn, radk, valuertorch. wtnerehh.inffotorrrr (tttttttttoattt...oo wtt...tttttee dtypo) attpeetatttttt))).. min)
20
21
indices
(ereerttttroooeeee (troe_atttn,
dessonding-irus)
22
eap_ent (we.iiii....eeaaa-11 iopiies indices :top_cnt]
rankout))
24
25
oth_incises indicen
26
20

---

27
sax_top top is
top_rnd, indiows torch, ranpperstppppppcccccm dsypestarch.inte devices device)) [ [eaa__top]
29
top, indiows top_indion top_rnd_itddewe
30
1f oth is not None:
oth_rnd_ indicaa cerrhh.reniprraiiin wwn.eaaammmmiiiinn..n.. max_och)
top_cnt, deypestarch, int64, devicewingg dea. device)
32
oth_ine itd_iedde rth, indices]
33
34
tap_sask tatch.gather (sask, expand (s((tru_aatt)) oth_rank tarch sarch-gather (make expand (a(truu__ttt)) diz--1, indecwoth, iateerotttiiniiinne indicas) E.... Nonn, :]
indez-top, indicea) [..., Xone]
36
37
cop_dreft_actte torch, gatbergttrrrrrrr (draft dige-1, indez-top, indices) [.... Mone]
38
ath_draftt_ttnn tarch.patierre
dine-1, inderwoth.
40
residual top_draft_ oth, acta
41
rasidual maak (top_rank ott_rask) (araask..aag (eppnn_aa(reettdie)) aat.arrrwettilllle.eeeeeeaaaaa
42
logita labala tarch, onax, (koggt)))))
(lacten(-3) ["residul] wask bool()
44
45
lass (lagi buta algha, labela)..put epu
47
count comnt
48
49
diff, loss
A.4
Needle-in-a-Hasstee
NIAH is a prompt-sensitive test focusing on relative performance changes before and after applying Spotlight Attention, rather than absolute performance. The haystack, needle, and question used in our evaluation are depicted in Figure The prompt, selected from a set proven effective in practice, shown in Figure
Haystack, Needle and Question
Haystack
Paul Graham Essays
Needle
The best thing do in San Francisco is sandwich and sit in Dolores Park sunny day.
Question
What the thing
Francisco?
Figure 7: The haystack, needle, and retrieval question for the NIAH.
Prompt
LLaMA2-7B LLaMA3-8B
You are helpful AI bot that answers questions for user. Keep your response short and direct. <context> <retrieval question> Don't give information outside the document or repeat your findings. The document definitely contains the answer, and I'm 100% sure. So try your best to find it.
LLaMA2-7B-Chat
<>>
 You helpful Al bot that answers questions for Keep your response short and direct. <context> <retrieval question> Don't give information outside the document repeat your findings. The document definitely contains the answer, and I'm 100% sure. So try your best to find it. [/ISS]]
Figure 8: Prompts used in NIAH test.
A.5 Language Modeling Perplexity
For Quest, we used the first 128 tokens for pre-filling with full attention. For MagicPIG, we fie-filled the first 1024 tokens with full attention to compute the key's mean value. Our Spotlight Attention method employed sparse KV retrieval throughout without pre-filling. For LLaMA2 models, we evaluated the first 4K tokens per sample; for LLaMA3 and Qwen2.5 models, we used the first 8K tokens. We compared Quest and MagicPIG separately, as Quest and our method use fixed token budget, while MagicPIG dynamically selects tokens and uses local windows and sink tokens.
21

---

In comparison with Quest, we calculated the token budget by multiplying the token sequence length by the pruning rate, with minimum budget 20.
A.6 Downstream QA Tasks
For contexts exceeding 8K tokens, we truncated 8K tokens from right to left. We evaluated all LongBench [6] subdatasets using the official test scripts. The LLaMA2 chat model employs the official chat template, whereas other models do not.
A.7 Bit-Packing CUDA Kernel
Bit-packing ( is crucial because PoTor lacks native bit type, and boolean values stored as full bytes. Without compaction, storage usage would increase substantially. The bit-packing program groups 32 boolean values and iteratively packs them into single uint_ detailed in the pseudocode accompanying Figure
Number bits
bit.packing
I Hi. PRESENT THE LOGIC HEAR, il THE INPESMRPITATOO CUDA,
Copy last bit
Lensor Lenaor chunk dix--1)
ouzput torch, peroa((a, // 32), dtypo-uist52)
tamsor:
10
output <C- 1
11
output 1- Dx01
① Shift logical left
reture output
Figure 9: Bit-packing
A.8 LSH Weight Initialization
For the linear hashing function, employ angular LSH with rotation matrix the initial value, outperforming standard random initialization. To generate d-dimensional rotation matrix, we use QR decomposition. We first create a random matrix R iddd with each element independently sampled from a standard normal distribution:
EE R"*" Rr, 10.
We then perform QR decomposition on R, yielding an orthogonal matrix Q and an upper triangular matrix R:
R = QR. QQ
(10)
The matrix Q is not necessarily in the special orthogonal group SO(d), as determinant can be either +1 or -1 To ensure Q E SD(d), we flip the sign of the first column of Q if det(Q) < 0.
B
More Experimental Results
B.1 Per-Head Retrieval Accuracy
The loU reported in the experiments section is averaged across all heads and layers. Given the varied behavior of LLM heads, individual loU scores differ significantly. Figure pressents detailed per-head IoU comparisons for various models, using Spotlight Attention and linear hashing functions, both before and after training.
B.2 Detailed QA Response Fidelity Scores
In the main text, we report the average Rouge-L score across all subdatasets as the similarity metric. However, scores vary significantly across individual subdatasets. II detailed similarity scores for each method compared to the original LLaMA3-8B model. To avoid confusion, we assign Rouge-L. score of to identical outputs, except in special cases like outputs containing only An characters, where tokenization issues may prevent perfect score.
22

---

Figure 10: Per-head retrieval accuracy. Measured by Intersection over Union (IoU) of top-k sets predicted by oracle and Spotlight, (1) The results of linear hashing reveal that most heads maintain low IoU both pre- and post-training, suggesting their latent distributions are challenging to approximate with linear functions. (2) Random parameter initialization results in low before-training IoU for Spotlight Attention, which improves significantly after training.
B.3 Few-Shot Learning
To evaluate Spotlight Attention with short context lengths, we restrict the KV cache 20 tokens and assess performance 5-shot learning datasets from LMM-val-Harness [[], including GLUE, SuperGLUE, OpenBookQA, HellaSwag, PiQA, Winogrande, ARC-E, ARC-C, MathQA, and MMLU.
As shown in Table Spotlight Attention delivers strong performance. Comparisons with Quest and MagiePIG are omitted, as their large local windows or budgets consume most of the prompt length on most datasets, rendering such comparisons less meaningful. Instead, we emphasize the relative performance between Spotlight Attention and the original model.
Table 8: A 5-shot learning comparison between original model and Spotlight Attention under fixed budget of 20 tokens was conducted across: GLUE (1), SuperGLUE (2). OpenBookQA (3), HellaSwag (4), PiQA (5), Winogrande (6), ARC-E (7), ARC-C (8), MathQA (9), and MMLU (10).
#4
#5
#6
#9
#10
Win Rate
LLaMA2-7B
0.489
0.646
0.342
0.604
0.776
0.733
0.793
0.478
0.262
0.468
56%
Spotlight
0.494
0.632
0.336
0.581
0.786
0.748
0.793
0,469
0.284
0.466
44%
LLaMA2-7B-Chan
0.625
0.717
0.346
0,598
0.763
0.729
0.811
0.474
0.295
0.482
89%
+Spotlight
0.582
0.664
0.35
0.578
0.759
0.703
0.809
0.467
0.275
0.482
11%
LLaMA3-8B
0.618
0.727
0.378
0.631
0.804
0.772
0.85
0.534
0.418
0.663
38%
+Spotlight
0.620
0.736
0.378
0.624
0.804
0.773
0.851
0.533
0.411
0.664
62%
B.4 Downstream QA Tasks
The main text primarily reports LLaMA3-SB scores on LongBench. Here, we additionally present results for LLaMA2-7B and LLaMA2-7B-CCT alongside comparisons of our method and Quest under varying token budgets, as shown Table 9
B.5 Ablation Loss Function
In the previous section, we analyzed the limitations of reconstruction loss. Here, we compare its empirical performance with our proposed ranking loss. As shown Table[TO our ranking loss significantly outperforms reconstruction loss, which provides minimal to benefit.
23

---

LLAMA3-5B
Upper Bound 90% Prined
Lincar Hashing
Spollight 99% Pruned
Quest 1024 Budge
Quest
MagiePIG
90% Pruned
256 Bodget
NarrativeQA
0.55
0,67
0.43
0.59
0.53
0.36
0.58
Qasper
MuliFieldQAAen
0.93
0.64
0.58
0.34
0.47
0.52
0.33
0.33
HotpotQA
1.00
0,76
0.48
0.66
0.70
0.49
0.69
2WikiMuhilopAA
1.00
0.83
0.53
0.75
0.77
0.61
0.69
MuSiQue
0.78
0.50
0.72
0.73
DuRender
GovRepurt
0.51
QMSum
1.00
0.48
0.35
0.44
0.39
0.8
0.88
NuluiNwws
0.28
0.82
0.60
0.79
0.84
0.5
0.03
VCSUM
0.21
0.46
0.39
0.39
0.41
0.06
0.10
TriviQA
SAMSum
LSHT
0.67
0.39
0.02
0.26
0.18
0.01
002
PassageCount
1.00
0.82
0.58
0.80
0.65
0.59
0.70
PussageRtrievallenn
1.00
0.71
0.41
0.55
0.60
0.42
0.41
PassgeRstrievl--ah
LCC
RepoBBenc-PP
0:20
0.71
0.31
0.59
0.47
0.30
041
Average
0.81
0.66
0.37
0.53
0.55
0.34
0.44
Figure 11: Detailed similarity between (i) the outputs of different models (including LLaMA3-8B itself) and (ii) those of LLaMA3-8B.
Table 9: Absolute scores on LongBench. Performance comparison of different methods Long-
Bench's 'ss long-text downstream tasks: NarrativeQA (1-1), Qasper (1-2), MultiFieldQA-en (1-3),
MultiFieldQA-zh (1-4), HotpotQA (2-1), 2WikiMultihopOA (2-2), MuSiQue (2-3), DuReader (2-4),
GovReport (3-1), QMSum (3-2), MultiNews (3-3), VCSUM (3-4), TREC (4-1), TriviaQA (4-2),
SAMSum (4-3), LSHT (4-4), PassageCount (5-1), PassageRetrieval-en (5-2), Passageeetrieval-zh (5-3), LCC (6-1), and RepoBench-P (6-2). (1) With 98% tokens pruned and no local window or global sink tokens, Spotlight Attention outperforms Quest and MagicPIG. (2) Spotlight Attention achieves performance on par with the original model, even in tasks like summarization and few-shot learning, (3) On subsets of Chinese (1-4, 2-4, 3-4), other LLaMA2-7B Chat-based models generated answers in English, while Quest produced responses in Chinese, achieving overwhelmingly high scores.
$11
1.3
12
11.0
520
1138
1231
18D
11400
123
5X0
3411
5248
Fraed
8.14
6.0
14.
13:6
13.00
6.6
194
1126
780
14.05
01.0
11
17.00
714
825
8.18
11.5
11-2
213
12.0
2.38
222
2-3
PA12
228
11/0
634
941
2008
Baton
1834
03%
14400
32:1
2-49
14.03
517
1833
16.0
1.30
431
335
1940
12.5
111
69
11
140
1240
13
1033 120er11 A11
15.4
21/6
2116
1331
KX
11.1
1156
KM
336
360
PMM
008
2111
141
11.)
B.6 Ablation on Attention Estimation Methods
We evaluated two approaches: hashing with Hamming distance (our default choice) and down- projection with inner product. At a 16x compression rate, results in Table show that hashing outperforms down-projection, demonstrating greater efficiency when the dimensionality is low.

---

Table 10: Ablation training loss. Compared to attention reconstruction loss, our proposed ranking loss yields significantly improved training out-
Table 11: Ablation on estimation method. We compared two attention estimation methods: hashing with Hamming distance (our default choice) and down- projection with inner product. Results indicate that hash- ing with Hamming distance performs better.
comes.
Loss
PC19
Math
Code
Estimation Method
Training
PG19
Math
Code
Attn, Recon. Loss
21.341
8.856
11,573
Down Proj. (16x) Inner
14.743
5,425
7.619
Ranking Loss (ours)
8.977
3.621
5.434
Hashing Hamming Dist. (ours)
8.977
3.621
5.444
25


Reviews:
[{"review_id": "h6y2MSuEX7", "reviewer": "Reviewer_syFB", "summary": "This paper studies the problem of efficient decoding with autoregressive LLMs, more specifically it studies the problem of efficient selection of kv cache during decoding. It builds over the dynamic query-aware KV cache selection method MagicPIG which hashes each key to a hash vector using an LSH function and similarly hashes each query and then retrieves top few key hashes for it, the retrieved keys are used in attention computation. The paper begins with arguing that linear hashing used in MagicPIG needs much bigger hashing space to get effective performance, to this end, it proposes to replace the linear hashing with a learned non-linear MLP based hash function. It is optimized using a ranking loss which encourages the top-2% ground truth key values to score higher than the rest of the key values when scored using the MLP based hash function. The proposed approach is evaluated on multiple benchmarks like LongBench, Needle-in-a-haystack, and perplexity evals on PG-19.", "questions": "Q1. What is the maximum evaluated context length in the performed experiments (on which quality evaluations are being done)?\n\nQ2. The proposed ranking loss seems quite computationally expensive, as per my understanding for every token in the example it computes total $\\mathcal{O}(N^2)$ ($N$ being the sequence length) loss terms which makes the whole loss computation $\\mathcal{O}(N^3)$ in nature? Can the authors confirm and in the case it is true why not explore alternate losses, few options: binary classification between top-K set and the rest, sigmoid based soft-topk loss (see Section 4.3 in [1])? \n\n--- \n\n[1] Dual-Encoders for Extreme Multi-label Classification, ICLR 2024", "limitations": "Limitations discussed.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 2, "strengths_and_weaknesses": "### Strengths\n- The proposed approach gets strong empirical performance on multiple benchmarks and seems to apply well to multiple LLMs and scale.\n- A detailed code repository is available for reproducibility\n- Writing is mostly clear and supported with useful visualizations\n\n### Weakness\n- Lacks a detailed ablation on the loss design, the proposed loss function is intuitive but lacks rigorous comparisons against alternate design choices or sensitivity to loss hyperparameters \n- Some statements in the introduction / abstract section lack context, for e.g. in line 28 \"GPU utilization drops to below 10% on average during decoding\" lacks the context of for what context sizes is this the case or how is this number supported? Similarly in the abstract \"enabling optimization of the non-linear hashing module on GPUs with 16GB memory in 8 hours\" lacks context - for what dataset, model?", "was_revised": true, "final_justification": "I am keeping my positive assessment after the rebuttal as it solidified the claims of the paper.", "rebuttal": "Thank you for your review. Your feedback has helped us improve the paper. Our responses are detailed below.\n\n---\n\n### **Regarding Weaknesses**\n\n#### **1. On the Detailed Ablation of the Loss Function**\n\nWe agree that the justification for our ranking loss could be strengthened. In the revision, we will move the ablation study from the appendix (Table 10) to the main text and expand on it. We will also add a discussion on how early explorations with simpler methods led to training instability, which will provide better motivation for our proposed ranking loss.\n\n#### **2. On the Context for Key Statements**\n\nWe have provided the specific context for the statements you highlighted and will integrate these clarifications into the manuscript.\n\n* **For \"GPU utilization drops to below 10%...\" (line 28):** This observation was preliminary. We have replaced it with a formal analysis of GPU FLOPs utilization efficiency, showing the decoding bottleneck more clearly:\n\n| Throughput Comparison (batch size=1) | 2048 Context | 4096 Context | 8192 Context |\n| :--- | :---: | :---: | :---: |\n| **Prefill** | 13,653 tok/s | 13,255 tok/s | 12,880 tok/s |\n| **Decoding** | 54 tok/s | 53 tok/s | 53 tok/s |\n| **Utilization Ratio** | 0.39% | 0.40% | 0.41% |\n\n* **For \"...enabling optimization...on GPUs with 16GB memory in 8 hours\" (abstract):** This figure corresponds to the optimization of our non-linear hashing module for the LLaMA2-7B model on a dataset of 8,192 samples (from Books and Arxiv). We will clarify this in the revised abstract.\n\n---\n\n### **Regarding Questions**\n\n\n#### **Q1. What is the maximum evaluated context length?**\n\nIn the original manuscript, evaluations were performed at an 8K context length. In response to your question, we conducted a new experiment at a **128K context length** by re-training and re-evaluating the Qwen2.5-7B model. The results show our method's relative advantage increases at this larger scale:\n\n| Model | Train Context | Baseline (PPL) | MLP Hashing Top-2% (PPL) |\n| :--- | :---: | :---: | :---: |\n| Qwen2.5-7B | 8K | 7.17 | 7.31 |\n| Qwen2.5-7B | 128K | 6.99 | **6.95** |\n\nThis finding is reinforced by a recent benchmark study we conducted comparing various compression methods. The fidelity results for our Spotlight Attention (top-10%) also show that its advantages grow with the context length.\n\n| Model Fidelity | 8K Context | 16K Context | 24K Context |\n| :--- | :---: | :---: | :---: |\n| Qwen2.5-7B-Instruct | 0.98268 | 0.98425 | 0.98852 |\n| Qwen2.5-14B-Instruct | 0.99055 | 0.97375 | 0.99213 |\n| Qwen2.5-32B-Instruct | 0.97953 | 0.98425 | 0.99475 |\n\nThe \"Model Fidelity\" metric in the table measures the similarity of 128 generated tokens between the compressed and original models for the given context length. This metric is conceptually similar to a normalized edit distance.\n\n#### **Q2. On the computational cost of the ranking loss.**\n\nYour analysis is absolutely correct, and we truly appreciate your critical thinking regarding our work. A naive implementation of the ranking loss has a computational complexity of $O(N^3)$, where $N$ is the sequence length. We omitted our practical solution from the original manuscript.\n\nOur implementation (line 376-378) uses a sampling strategy to mitigate this cost. At each step, we sample a limited number of queries (via `--max-que`), top-K keys (`--max-top`), and non-top-K keys (`--max-oth`) up to 1024 each. This reduces the practical training complexity to a constant, $O(1)$, making it efficient for long sequences. This sampling mechanism will be clarified in the revised manuscript.\n\nThank you for suggesting alternate losses like binary classification or the Soft-TopK loss. The Soft-TopK loss, which focuses on identifying the correct *set* of top-K items without regard to their internal order, aligns well with our objective. We see this as a promising direction for future work."}, {"review_id": "WQl7uxPGoq", "reviewer": "Reviewer_myg1", "summary": "The paper introduces **Spotlight Attention**, a novel method to enhance the efficiency of Large Language Model (LLM) inference by optimizing key-value (KV) cache retrieval using non-linear hashing. It addresses the inference bottleneck caused by frequent memory exchanges in LLMs by dynamically selecting critical KV cache entries during decoding. Unlike traditional linear hashing methods (e.g., Locality-Sensitive Hashing, LSH), which struggle with the non-uniform distribution of queries and keys in LLMs, Spotlight Attention employs a two-layer Multi-Layer Perceptron (MLP) for non-linear hashing, improving retrieval precision and reducing hash code length. The method includes a lightweight training framework using a Bradley-Terry ranking-based loss, enabling efficient optimization on modest hardware (16GB GPU in 8 hours). Experimental results demonstrate superior KV retrieval accuracy, comparable perplexity to existing methods like Quest with a smaller token budget, and significant throughput gains, particularly for long-context tasks. Evaluations span multiple benchmarks, including PG19, ProofPile, CodeParrot, LongBench, and Needle-in-a-Haystack, showcasing the method’s robustness and generalization.", "questions": "- The paper mentions adding a hash code-based retrieval mechanism for each layer (Page 2, Figure 1), but the integration process into existing LLM architectures is unclear due to truncation. Could the authors provide a detailed explanation or pseudocode for how Spotlight Attention is incorporated into the attention mechanism? \n\n- The evaluation focuses on successes but lacks discussion of scenarios where Spotlight Attention may underperform (e.g., tasks with highly sparse attention patterns). Could the authors provide examples of tasks or datasets where performance degrades and explain why? \n\n- The paper emphasizes long-context tasks but provides limited insight into short-context or real-time applications (Page 15, Table 8). Can the authors evaluate Spotlight Attention on additional short-context tasks or compare it with Quest/MagicPIG in these settings? \n\n- The comparison with MagicPIG uses only 10 data points due to its time-consuming evaluation (Page 7, Table 3). Could the authors justify this choice or provide additional comparisons with a larger sample size? \n\n- The lightweight training framework (8 hours on a 16GB GPU) is a strength, but the impact of scaling to larger models or datasets is unclear. Can the authors discuss how training time and resource requirements scale with model size (e.g., beyond LLaMA2-7B)?", "limitations": "Yes\n\nThe authors explicitly state that the paper has \"no limitations\" and poses \"no such risks\" regarding negative societal impacts", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths:**\n\n- The paper provides a rigorous evaluation across multiple dimensions, including KV retrieval accuracy (Table 1), perplexity comparisons (Tables 2 and 3), QA response fidelity (Table 4), and throughput efficiency (Figure 6). The use of established benchmarks like PG19, LongBench, and Needle-in-a-Haystack ensures robust validation. The ablation studies (Tables 10 and 11) further strengthen the empirical foundation by comparing the proposed ranking loss and hashing method against alternatives.\n\n- The method addresses a critical bottleneck in LLM inference—low GPU utilization during decoding (Page 1)—which is a pressing issue for scaling long-context LLMs. By achieving comparable performance to methods like Quest with a smaller token budget and outperforming MagicPIG in hash code efficiency (Page 7, Figure 6), the approach has practical implications for resource-constrained environments and long-sequence tasks.\n\n- The paper clearly articulates the problem, proposed solution, and experimental setup. Figures (e.g., Figure 1, Figure 2) and tables (e.g., Table 1, Table 2) effectively illustrate the architecture, motivation, and results. The mathematical formulation of MLP hashing (Page 5, Equations 167-168) is precise, and the training setup details (Page 7, Page 12) enhance reproducibility.\n\n**Weaknesses:**\n\n-  While the paper is generally clear, the extensive truncation in the provided document (e.g., 851 characters on Page 1, 4189 on Page 2, etc.) obscures some details, particularly in the methodology and experimental results sections. For instance, the exact mechanism for integrating Spotlight Attention into existing LLM architectures is partially unclear due to missing text. This could confuse readers unfamiliar with the context.\n\n- The evaluation, while comprehensive, lacks a detailed discussion of failure cases or scenarios where Spotlight Attention underperforms. For example, the paper notes improved performance on long-context tasks but does not explicitly address potential degradation in specific task types (e.g., highly sparse attention patterns). Additionally, the comparison with MagicPIG is limited to a small sample (10 data points, Page 7), which may weaken claims of superiority.\n\n- The paper’s focus on long-context tasks is significant, but its applicability to short-context or real-time applications is underexplored. The few-shot learning evaluation (Page 15, Table 8) is promising but lacks comparisons with Quest and MagicPIG due to their large token budgets, limiting the ability to assess relative performance in these settings.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback. We address your concerns below.\n\n## Regarding Weaknesses\n\n1.  **On PDF Truncation**\n\n    We were unable to reproduce the text truncation issue when viewing the PDF in Google Chrome. Could you please provide details about your viewing environment so we can investigate further?\n\n2.  **On Evaluation in Sparse Attention Scenarios**\n\n    We acknowledge that we did not evaluate Spotlight Attention on synthetic tasks with highly sparse attention patterns. We are considering potential evaluation methods and welcome any suggestions.\n\n3.  **On Short-Context Performance**\n\n    Our rationale for not focusing on short-context scenarios was the potential overhead of Spotlight Attention, which may not offer a speed-up over standard Flash Attention in such cases. However, we agree that a robust attention method should perform well across all context lengths.\n\n## Regarding Questions\n\n**Q1: On Architectural Integration**\n\nThe following pseudocode illustrates the integration of Spotlight Attention into the attention mechanism.\n\n```python\n# Hash the current query and key\nq_hash = mlp_hash(q)\nk_hash = mlp_hash(k)\n\n# Update the hash cache with the new key's hash\npast_key_hash.update(k_hash)\n\n# Perform similarity search to find relevant past keys\nscore = sim_search(q_hash, past_key_hash)\n\nindices = topk(score)\nselected_keys, selected_vals = select_by_indices(past_key, past_val, indices)\n\n# Compute final attention output using the selected KV cache\noutput = flash_attention(q, selected_keys, selected_vals)\n```\n\n**Q2: On Failure Cases and Performance Degradation**\n\nWe believe the **Needle-in-a-Haystack (NIAH)** task serves as a scenario with highly sparse information. We attribute potential performance degradation primarily to the precision of the retrieval process, not the top-K selection itself. Our experimental data (Table 1) shows an Intersection over Union (IoU) of approximately **0.35**. While an improvement over prior methods, this indicates room for further optimization of the system's architecture and loss function.\n\n**Q3: On Short-Context Comparisons**\n\nWe were unable to evaluate Quest and MagicPIG in very short-context settings as their local window size (64 tokens) exceeds the task's context length. Spotlight Attention does not have this constraint, which is why only its performance is presented in Table 8.\n\n**Q4: On the MagicPIG Comparison Sample Size**\n\nWe agree that the initial validation set of 10 samples was small. We have since conducted a more extensive Perplexity (PPL) evaluation on the **PG19**, **Math**, and **Code** datasets. The new results are as follows:\n\n| Method | Frozen | Local(64) | Sink(4) | Retrieve(2%) | PG19 | Math | Code |\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| Vanilla | [0,1] | | | | 8.604 | 3.517 | 5.219 |\n| Oracle Top-K | [0,1] | | | ✅ | 8.881 | 3.590 | 5.353 |\n| MagicPIG | [0,16] | ✅ | ✅ | ✅ | 11.478 | 4.284 | 6.762 |\n| MagicPIG | [0,16] | | | ✅ | NaN | NaN | NaN |\n| Spotlight | [0,1] | | | ✅ | **8.977** | **3.621** | **5.434** |\n\nThese expanded results show that Spotlight Attention outperforms MagicPIG across these domains without using a local window or sink tokens.\n\n**Q5: On Scaling to Larger Models**\n\nThe computational overhead of Spotlight Attention increases with model size, but not linearly. The fine-tuning cost for a 72B model is not 10 times that of a 7B model because we only fine-tune the attention components, and the relative cost of attention decreases as model size grows. Our calculations show the attention-only computation for a **72B model** is about **twice** that of a **7B model**. Therefore, we estimate the fine-tuning time for a 72B model would be approximately double that of a 7B model, given adequate GPU resources."}, {"review_id": "gP5mL1n5V4", "reviewer": "Reviewer_YZ39", "summary": "This paper focuses on compressing the KV cache to improve long-context LLM inference. The authors propose a non-linear hashing mechanism combined with novel training objectives to optimize the distribution of queries and keys. Specialized CUDA kernels are developed to ensure the efficiency of the proposed method.", "questions": "Could the authors provide comparisons with competitive quantization-based baselines to further demonstrate the effectiveness of Spotlight Attention?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n1. The motivation for using non-linear hashing is clear and well-articulated.\n2. The Spotlight Attention pipeline is simple yet effective.\n3. The custom CUDA implementation ensures efficient hash code processing, which is critical for the practical deployment of Spotlight Attention.\n\n**Weaknesses**\n\n1. Quantization is a key technique for KV cache compression. However, the paper does not discuss relevant recent works in this area, such as KIVI [1], KVQuant [2], QuaRot [3], and DuQuant [4]. Including a comparison or discussion with these baselines would strengthen the positioning of this work.\n2. The paper does not report results on LongBench v2 [5], a widely used benchmark for long-context evaluation. It would be helpful to understand Spotlight Attention’s performance on this dataset.\n\n[1]. KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache. ICML 2024.\n\n[2] KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization. NeurIPS 2024.\n\n[3] QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs. NeurIPS 2024.\n\n[4] DuQuant: Distributing Outliers via Dual Transformation Makes Stronger Quantized LLMs. NeurIPS 2024.\n\n[5] LongBench v2: Towards Deeper Understanding and Reasoning on Realistic Long-context Multitasks. ACL 2025.", "was_revised": true, "final_justification": "The authors provide comparisons with KV quantization methods, which address most of my concerns. However, I did not raise my score since the LongBench v2 results do not show sufficient improvement.", "rebuttal": "We thank the reviewer for their feedback. We have addressed the identified weaknesses below and will incorporate these results and discussions into the revised manuscript.\n\n***\n\n### **Regarding Weaknesses**\n\n**1. Comparison with Quantization-based Baselines**\n\n> *Reviewer's Comment: The paper does not discuss or compare with recent relevant works in KV cache quantization, such as KIVI, KVQuant, QuaRot, and DuQuant.*\n\nWe agree that a comparison with KV cache quantization methods is important for contextualizing our work. Coincidentally, in another benchmark study where we evaluate various model compression methods, we have tested the KVQuant [1] method. In addition, we have also tested other compression methods including SnapKV [2], H2O, Flash Attention FP8 [3], SageAttention [4], SparseGPT, GPTQ, AWQ, and so on. The metric we use is a very strict one, similar to a normalized edit distance. This metric measures the similarity of 128 tokens generated by the model before and after compression within a given context, and it aligns highly with GPT4o evaluation. (This work is still unpublished, so we have not provided the specific metric expression.)\n\n*First, a performance comparison on a 7B model shows competitive performance for our method.*\n\n| Qwen2.5-**7B**-Instruct | KVQuant (4-bit-1%) [1] | SnapKV (10% local + 10% budget) [2] | Flash FP8 [3] | SageAttn (qk fp8 + vo int8) [4] | **Spotlight (Top-10%, ours)** |\n| :--- | :---: | :---: | :---: | :---: | :---: |\n| Knowledge | 0.89 | 0.66 | 0.47 | 0.98 | **0.92** |\n| Reason | 0.93 | 0.69 | 0.50 | 0.98 | **0.95** |\n| CN/FR/JP | 0.91 | 0.59 | 0.62 | 0.98 | **0.95** |\n\n*Second, we evaluated how these methods scale on a 32B model with increasing context lengths.*\n\n| Qwen2.5-**32B**-Instruct | KVQuant (4-bit-1%) [1] | SnapKV (10% local + 10% budget) [2] | Flash FP8 [3] | SageAttn (qk fp8 + vo int8) [4] | **Spotlight (Top-10%, ours)** |\n| :--- | :---: | :---: | :---: | :---: | :---: |\n| **8K Ctx** | 0.945 | 0.315 | 0.993 | 1.000 | 0.980 |\n| **16K Ctx** | 0.945 | 0.136 | 0.974 | 0.989 | 0.984 |\n| **24K Ctx** | 0.923 | 0.027 | 0.995 | 0.986 | **0.995** |\n\nThe results indicate that the fidelity of some quantization-based methods can degrade as context length grows, whereas our method's fidelity improves.\n\n**2. Evaluation on LongBench v2**\n\n> *Reviewer's Comment: The paper does not report results on LongBench v2, a widely used benchmark for long-context evaluation.*\n\nThank you for this suggestion. We have already begun the evaluation process.\n\nSpecifically, we are adapting the vLLM testing environment to run the LongBench v2 suite. We anticipate having the complete results **within the next two days**. We will provide the results as soon as they are available and will include them in the final version of the paper.\n\n***\n\n### **References**\n[1] Coleman Hopper et al. KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization. arXiv preprint arXiv:2401.18079.  \n\n[2] Yuhong Li et al. SnapKV: LLM Knows What You are Looking for Before Generation. arXiv preprint arXiv:2404.14469.\n\n[3] Jay Shah et al. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv preprint arXiv:2407.08608. \n\n[4] Jintao Zhang et al. Sage Attention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration. arXiv preprint arXiv:2410.02367."}, {"review_id": "NQGOkHbUBv", "reviewer": "Reviewer_FNkM", "summary": "This paper proposes a method which dynamically selecting critical KV cache in LLMs via non-linear hashing. It computes hash codes for the cached keys and the current query, the algorithm efficiently determine a top-k subset of the KV cache. This reduces memory and compute overhead during inference. \n\nThe method achieves up to 3× throughput improvement on Qwen2.5-7B with long sequences (32K/128K), with minimal accuracy drop across LLaMA and Qwen families.", "questions": "1. Since the MLP for hashing is trained offline, could low-rank adaptation techniques such as LoRA be applied to further reduce training time and memory?\n2. Given that the hash MLP is pre-trained using few-shot (8k samples) for 1 epoch, does the learned parameters generalize well across datasets and downstream tasks? Would retraining be necessary for task-specific adaptation?", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "**strength**\n\n1. The method applies a Bradley-Terry ranking loss to select important tokens. Retaining only 2% of tokens to continue the following attention computation, it maintains strong performance across many tasks. Appendix B.3 shows that even with a fixed 20-token budget, accuracy remains competitive. \n2. A custom CUDA kernel with bit-packing and NXOR GEMM is implemented to accelerate retrieval, enabling lower latency in real implementation. \n3. Evaluation is comprehensive, covering various context lengths, hash code sizes, throughput benchmarks, and downstream tasks. \n4. Compared to previous method MagicPIG, the hash code is much shorter (128-bit), reducing  hashing code storage and computational overhead. \n\n**weakness**\n\n1. The paper lacks ablation on the time cost of the efficiency of encoding Queries and Keys and selecting the topk. If authors can provide a detailed time spent on each step, it would be helpful. \n2. The top-k or percentile for token budget is a fixed hyperparameter. It’s unclear whether this threshold generalizes across different tasks or requires per-task tuning. And this paper doesn’t provide the model performance and latency performance under different token budgets. \n3. Figure 6 has repeat information by two groups of subfigures. To be clarify,  batch_size=1/2, content_length =32k/128k appear twice, which can be consolidated by uniformly show one kind of plot axes. \n4. Generalization of the hash MLP is not discussed. Its robustness across domains and tasks is unclear.", "was_revised": true, "final_justification": "The authors’ responses have addressed all my concerns, particularly the speed evaluations compared to FlashAttention-2, which I consider highly important in the field of inference acceleration. In addition, the authors have provided extra results on more datasets and models in response to other reviewers’ recommendations, further demonstrating the strong performance of the proposed method. I have decided to raise my score to indicate my support for the acceptance of this paper.", "rebuttal": "We thank the reviewer for their constructive feedback. We have addressed their concerns below and revised the manuscript accordingly.\n\n---\n\n### **Regarding Weaknesses**\n\n**1. Ablation on the time cost of encoding, searching, and selecting top-k.**\n\nWe agree that a detailed breakdown of the time cost for each component is essential. We have benchmarked each step and will add the following results to the paper.\n\n*Efficiency breakdown on Qwen2.5-7B (batch size=1, time in µs).*\n| Prompt Length | MLP Hashing (PyTorch) | Sim Search (CUDA) | Top-K Gathering (PyTorch) | Sparse Attn (FA2) | **Our Method (Overall)** | Full Attention (FA2) |\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| Complexity | $\\mathcal{O} (1)$ | $\\mathcal{O} (N)$ | $\\mathcal{O} (N\\log K)$ | $\\mathcal{O} (K^2)$ | - | $\\mathcal{O} (N^2)$ |\n| **1,048,576** | 76 | 99 | 117 | 97 | **389** | 1281 |\n| **2,097,152** | 77 | 171 | 230 | 119 | **597** | 2475 |\n\n**2. The fixed hyperparameter for the token budget and its generalization.**\n\nThe choice of a fixed budget (e.g., top-2%) was an empirical finding, not a per-task tuned hyperparameter. To further analyze its robustness, we conducted new experiments measuring perplexity across different token budgets and training corpora.\n\n*PG19 (Lower is better), 10 Samples, 8K tokens*\n| Model Trained on | Budget: 95% | 96% | 97% | 98% | 99% | Full Attention |\n|:---|:---:|:---:|:---:|:---:|:---:|:---:|\n| **Code** | 9.686 | 9.724 | 9.781 | 9.878 | 10.075 | 9.561 |\n| **Arxiv+Book** | 9.704 | 9.742 | 9.795 | 9.893 | 10.088 | 9.561 |\n\n*Math (Lower is better), 10 Samples, 8K tokens*\n| Model Trained on | Budget: 95% | 96% | 97% | 98% | 99% | Full Attention |\n|:---|:---:|:---:|:---:|:---:|:---:|:---:|\n| **Code** | 3.163 | 3.170 | 3.182 | 3.202 | 3.250 | 3.137 |\n| **Arxiv+Book** | 3.168 | 3.176 | 3.189 | 3.211 | 3.262 | 3.137 |\n\n*Code (Lower is better), 10 Samples, 8K tokens*\n| Model Trained on | Budget: 95% | 96% | 97% | 98% | 99% | Full Attention |\n|:---|:---:|:---:|:---:|:---:|:---:|:---:|\n| **Code** | 2.276 | 2.279 | 2.284 | 2.294 | 2.317 | 2.267 |\n| **Arxiv+Book** | 2.277 | 2.280 | 2.285 | 2.296 | 2.322 | 2.267 |\n\nThe results suggest that a pruning rate of 97-98% (keeping 2-3% of tokens) offers a robust performance-computation trade-off across tasks. We have added these results and this discussion to the appendix.\n\nRegarding latency, the modest speedup compared to FlashAttention is due to the lack of a fully fused sparse operator, which we identify as a direction for future engineering work.\n\n**3. Redundancy in Figure 6.**\n\nWe agree. We are consolidating the subfigures for `batch_size` and `context_length` into a single, clearer set of plots in the revised manuscript.\n\n**4. Generalization of the hash MLP.**\n\nThe generalization of the hash MLP across different data distributions is a critical point. We provided an initial analysis in Table 6, copied below, which shows that models trained on diverse corpora (Arxiv+Book, C4, Github Code) achieve similar performance on different tasks.\n\n| LLaMA3-8B | Training Corpus | PG19 | Math | Code |\n|:---|:---|:---:|:---:|:---:|\n| Oracle Top-2% | - | 8.881 | 3.590 | 5.353 |\n| **Ours (Top-2%)** | **Arxiv+Book (default)** | **8.977** | **3.621** | **5.434** |\n| **Ours (Top-2%)**| **C4** | **8.958** | **3.631** | **5.417** |\n| **Ours (Top-2%)**| **Github Code** | **8.891** | **3.611** | **5.393** |\n\nTo further test the robustness of the hash function itself, independent of semantic content, we plan as a next step to investigate training the MLP on a dataset of *meaningless random tokens*. We will explore this in our future work.\n\n---\n\n### **Regarding Questions**\n\n**1. On applying LoRA to the hash MLP.**\n\nYes, this is a viable approach. The MLP can be designed with a LoRA-like bottleneck structure to reduce trainable parameters. Additionally, after the hashing components are trained, LoRA could be applied to fine-tune the full LLM to potentially recover performance lost from the approximation.\n\n**2. On the generalization of the MLP trained with few-shot samples.**\n\nYes, our experiments suggest the learned parameters generalize well from a small training set. Our evidence for this is:\n* As shown above, checkpoints trained on three different corpora (Code, C4, Arxiv+Book) perform almost identically across diverse downstream tasks.\n* For instruction-tuned models like LLaMA2-7B-Chat, training our hash MLP on its pre-training corpus did not degrade performance on downstream LongBench tasks."}, {"review_id": "eOcRJL1t8d", "reviewer": "Reviewer_BJjk", "summary": "This work proposes spotlight attention to approximately select top-k attention scores with non-linear hash function. This work also provides a layer-wise training process, together with a rank-based objective function and a differentiable softsign function, to learn the data-dependent non-linear hash function.  Compared with local-sensitive hash (linear hash function) and other baselines, the proposed non-linear hash function performs most closely to oracle top-k. (But with less compute).", "questions": "- In figure 5 right, the green bar (magicpig) is significantly better than others on TrivialQA and LCC tasks. Do you know the reason? \n- It is well-known that trainable non-linear hash-function is data-dependent. This work shows ablation on training tasks in table 6. Do you know the the reason why Github Code outperforms Arxiv + Books?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "## Strengths\n- The layer-wise training process, including the rank-based loss function, is very interesting and efficient.\n- This work implements CUDA kernels for the hash code processing, including bit-packing and bitwise NXOR GEMM operators, achieving significant latency reductions in practice.\n- Compared with linear hash function (LSH top-k), the non-linear attention performs most closely to oracle Top-K. \n\n## Weaknesses\n\n- This paper introduces Spolitlight attention as an advancement over MagicPIG (e.g. line 294). However, there is a fundamental difference between Spolitlight attention and magicpig. That is spolitlight attention selects top-k (a hard threshold), while magicpig samples tokens according to probability. This difference worths a clarification. This difference may also explain the performance difference in Figure 5 (right). \n- This work heavly uses perplexity in comparisons (e.g. figure 2, 3, 5, 6). However, this evaluation metric shows very little about the model performance. For example, in figure 2, pruning 98% tokens (oracle top-k) almost doesn't change the perplexity. It would be interesting to use other practical evaluation metric.", "was_revised": true, "final_justification": "My initial score was 4. Now I raises score to 5 because the author adds evaluation metrics beyond Perplexity, comparisons between MagicPig and Splitlight attention.", "rebuttal": "We thank the reviewer for their constructive feedback. We have addressed their concerns below and revised the manuscript accordingly.\n\n***\n\n### **Regarding Weaknesses**\n\n**1. On the Distinction Between Spotlight Attention and MagicPIG**\n\nWe agree with the reviewer's assessment: Spotlight Attention uses a hard top-k threshold, whereas MagicPIG samples tokens based on probability. We recognize this crucial distinction was not made explicit enough in the original manuscript.\n\nWe have now added a dedicated paragraph to the related work section to clarify this fundamental difference in their mechanisms. Furthermore, we wish to offer a more nuanced perspective on this distinction. While Spotlight's design is deterministic top-k, its practical execution is not perfect. As our results show an Intersection-over-Union (IoU) of approximately 0.35 with the oracle top-K, our method does not retrieve the exact set.\n\n**2. On the Use of Perplexity (PPL) as an Evaluation Metric**\n\nWe agree that PPL can be an insufficient metric for generation quality. To provide more robust evidence, we have conducted new evaluations using an **output fidelity** metric. This metric, conceptually similar to a normalized edit distance, better correlates with human judgment of text quality.\n\nThe new results show that Spotlight Attention's strong performance is not an artifact of PPL and that it maintains high-quality generation, with fidelity improving as context length grows.\n\n*First, we compared Spotlight Attention with other strong baselines on a 7B model (generate 128 tokens):*\n| Qwen2.5-**7B**-Instruct | KVQuant (4-bit) | SnapKV (10% local + 10% budget) | Flash FP8 | SageAttn (qk fp8 + vo int8) | **Spotlight (Top-10%, ours)** |\n| :--- | :---: | :---: | :---: | :---: | :---: |\n| Knowledge | 0.89 | 0.66 | 0.47 | 0.98 | **0.92** |\n| Reasoning | 0.93 | 0.69 | 0.50 | 0.98 | **0.95** |\n| CN/FR/JP | 0.91 | 0.59 | 0.62 | 0.98 | **0.95** |\n\n*Second, we evaluated its scaling performance on a 32B model with increasing context lengths:*\n| Qwen2.5-**32B**-Instruct | KVQuant (4-bit) | SnapKV (10% local + 10% budget) | Flash FP8 | SageAttn (qk fp8 + vo int8) | **Spotlight (Top-10%, ours)** |\n| :--- | :---: | :---: | :---: | :---: | :---: |\n| **8K Ctx** | 0.945 | 0.315 | 0.993 | 1.000 | 0.980 |\n| **16K Ctx** | 0.945 | 0.136 | 0.974 | 0.989 | 0.984 |\n| **24K Ctx** | 0.923 | 0.027 | 0.995 | 0.986 | **0.995** |\n\n\n***\n\n### **Regarding Questions**\n\n**1. On MagicPIG's Performance on TriviaQA and LCC Tasks**\n\nOur qualitative analysis of the raw model outputs indicates that MagicPIG's superior scores on these tasks are often due to differences in **output formatting and style** rather than factual correctness.\n\n-   *For TriviaQA*, MagicPIG often prepends the short answer before the full passage (e.g., `**55**\\nPassage:...`), which aligns better with the scoring script.\n-   *For LCC (code generation)*, MagicPIG consistently generates more coherent code snippets than the baseline, whereas our method can sometimes replicate the baseline's undesirable behavior.\n\nHowever, to provide a complete picture, we also note that MagicPIG's probabilistic nature can be detrimental on tasks requiring more precise reasoning or summarization:\n\n|                     | MultifiledQA-EN (QA) | MultiNews (Summarization) | Passage Count (Counting) |\n| :------------------ | :------------------: | :-----------------------: | :----------------------: |\n| Vanilla             |         21.4         |           3.78            |           2.00           |\n| **MagicPIG** |  **17.51** (Worse)   |     **1.37** (Worse)      |    **1.20** (Worse)    |\n| Spotlight (ours)    |        21.54         |           3.36            |           2.03           |\n\nThis analysis provides a more nuanced understanding of the qualitative trade-offs between these attention mechanisms.\n\n**2. On Why Github Code Outperforms Arxiv + Books for Training**\n\nTo investigate the surprising generalization performance of code-based training, we ran a more comprehensive set of experiments. The results consistently confirm that the model trained on code holds a slight edge in perplexity across all tested domains and sparsity levels.\n\n*PG19 (Perplexity, Lower is better)*\n| Model Trained on | Budget: 95% | 96% | 97% | 98% | 99% |\n|:---|:---:|:---:|:---:|:---:|:---:|\n| **Code** | **9.686** | **9.724** | **9.781** | **9.878** | **10.075** |\n| **Arxiv+Book** | 9.704 | 9.742 | 9.795 | 9.893 | 10.088 |\n\n*Math (Perplexity, Lower is better)*\n| Model Trained on | Budget: 95% | 96% | 97% | 98% | 99% |\n|:---|:---:|:---:|:---:|:---:|:---:|\n| **Code** | **3.163** | **3.170** | **3.182** | **3.202** | **3.250** |\n| **Arxiv+Book** | 3.168 | 3.176 | 3.189 | 3.211 | 3.262 |\n\n*Code (Perplexity, Lower is better)*\n| Model Trained on | Budget: 95% | 96% | 97% | 98% | 99% |\n|:---|:---:|:---:|:---:|:---:|:---:|\n| **Code** | **2.276** | **2.279** | **2.284** | **2.294** | **2.317** |\n| **Arxiv+Book** | 2.277 | 2.280 | 2.285 | 2.296 | 2.322 |\n\nWhile the precise theoretical reason remains an open question, this finding has inspired a new direction for our future work. We now plan to investigate if training on a dataset of *meaningless random tokens* could achieve a similar, or perhaps even superior, generalizing effect for learning the hash function's structure. We will add a discussion of this finding and future direction to the paper."}]

Author rebuttal:

Meta-review: This paper introduces Spotlight Attention, a non-linear hashing-based approach for KV cache retrieval in LLMs. The method is well-motivated, addressing inefficiencies in linear hashing, and combines a Bradley-Terry ranking loss with custom CUDA kernels to achieve strong retrieval accuracy and throughput improvements. Reviewers consistently highlight the solid empirical results across multiple benchmarks, robustness to long contexts, and practical engineering contributions. While some weaknesses were noted (e.g., limited ablations, missing comparisons with quantization baselines, and clarity issues), the authors’ rebuttal provided additional experiments and clarifications that largely addressed these concerns.

Overall, the work is technically sound, impactful for efficient LLM inference, and demonstrates both novelty and practicality. I recommend acceptance.
Decision: Accept (poster)