Title: Spotlight Attention: Towards Efficient LLM Generation via Non-linear Hashing-based KV Cache Retrieval
Venue: NeurIPS 2025
Accepted: yes
Keywords: KV Cache Compression

Efficient Attention
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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).\n\nSTRENGTHS AND WEAKNESSES:\n## 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\nCould the authors provide comparisons with competitive quantization-based baselines to further demonstrate the effectiveness of Spotlight Attention?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n- 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)?\n\nLIMITATIONS:\nYes\n\nThe authors explicitly state that the paper has \"no limitations\" and poses \"no such risks\" regarding negative societal impacts", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n### 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?\n\nQUESTIONS:\nQ1. 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\n\nLIMITATIONS:\nLimitations discussed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-05"}