Spotlight Attention: Towards Efficient LLM Generation via Non-linear Hashing-based KV Cache Retrieval

We thank the reviewer for their constructive feedback. We have addressed their concerns below and revised the manuscript accordingly.

Regarding Weaknesses

1. On the Distinction Between Spotlight Attention and MagicPIG

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

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

2. On the Use of Perplexity (PPL) as an Evaluation Metric

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

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

First, we compared Spotlight Attention with other strong baselines on a 7B model (generate 128 tokens):

Second, we evaluated its scaling performance on a 32B model with increasing context lengths:

Regarding Questions

1. On MagicPIG's Performance on TriviaQA and LCC Tasks

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

• For TriviaQA, MagicPIG often prepends the short answer before the full passage (e.g., **55**\nPassage:...), which aligns better with the scoring script.
• 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.

However, to provide a complete picture, we also note that MagicPIG's probabilistic nature can be detrimental on tasks requiring more precise reasoning or summarization:

This analysis provides a more nuanced understanding of the qualitative trade-offs between these attention mechanisms.

2. On Why Github Code Outperforms Arxiv + Books for Training

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

PG19 (Perplexity, Lower is better)

Math (Perplexity, Lower is better)

Code (Perplexity, Lower is better)

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

Thanks for your reply. My concerns are addressed. I recommends for acception, raises score to 5.

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.

Regarding Weaknesses

1. Comparison with Quantization-based Baselines

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.

We 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.)

First, a performance comparison on a 7B model shows competitive performance for our method.

Second, we evaluated how these methods scale on a 32B model with increasing context lengths.

The results indicate that the fidelity of some quantization-based methods can degrade as context length grows, whereas our method's fidelity improves.

2. Evaluation on LongBench v2

Reviewer's Comment: The paper does not report results on LongBench v2, a widely used benchmark for long-context evaluation.

Thank you for this suggestion. We have already begun the evaluation process.

Specifically, 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.

References

[1] Coleman Hopper et al. KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization. arXiv preprint arXiv:2401.18079.

[2] Yuhong Li et al. SnapKV: LLM Knows What You are Looking for Before Generation. arXiv preprint arXiv:2404.14469.

[3] Jay Shah et al. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv preprint arXiv:2407.08608.

[4] Jintao Zhang et al. Sage Attention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration. arXiv preprint arXiv:2410.02367.

We encountered a small issue while adapting Spotlight to vLLM. Consequently, we modified the pred.py script from the original code, replacing the API call with our own custom generate function.

However, for reasons we have not yet fully identified, our reproduced results are significantly worse than those reported by LongBench V2 paper. We suspect this discrepancy may be due to differences in the chat template or prompt compared to the vLLM implementation, but we currently lack the time to investigate this further. For now, I am sharing the results we have obtained:

Current Results on LongBench V2

The parameters we used are rag=0, cot=False, no_context=False. We have placed the relevant evaluation script in an anonymous repository (test_longbench/pred_v2.py).

We thank the reviewer for their constructive feedback. We have addressed their concerns below and revised the manuscript accordingly.

Regarding Weaknesses

1. Ablation on the time cost of encoding, searching, and selecting top-k.

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

Efficiency breakdown on Qwen2.5-7B (batch size=1, time in µs).

2. The fixed hyperparameter for the token budget and its generalization.

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

PG19 (Lower is better), 10 Samples, 8K tokens

Math (Lower is better), 10 Samples, 8K tokens

Code (Lower is better), 10 Samples, 8K tokens

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

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

3. Redundancy in Figure 6.

We agree. We are consolidating the subfigures for batch_size and context_length into a single, clearer set of plots in the revised manuscript.

4. Generalization of the hash MLP.

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

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

Regarding Questions

1. On applying LoRA to the hash MLP.

Yes, 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.

2. On the generalization of the MLP trained with few-shot samples.

Yes, our experiments suggest the learned parameters generalize well from a small training set. Our evidence for this is:

• As shown above, checkpoints trained on three different corpora (Code, C4, Arxiv+Book) perform almost identically across diverse downstream tasks.
• 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.

Thank you for your constructive feedback. We address your concerns below.

Regarding Weaknesses

1. On PDF Truncation

   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?

2. On Evaluation in Sparse Attention Scenarios

   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.

3. On Short-Context Performance

   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.

Regarding Questions

Q1: On Architectural Integration

The following pseudocode illustrates the integration of Spotlight Attention into the attention mechanism.

# Hash the current query and key
q_hash = mlp_hash(q)
k_hash = mlp_hash(k)

# Update the hash cache with the new key's hash
past_key_hash.update(k_hash)

# Perform similarity search to find relevant past keys
score = sim_search(q_hash, past_key_hash)

indices = topk(score)
selected_keys, selected_vals = select_by_indices(past_key, past_val, indices)

# Compute final attention output using the selected KV cache
output = flash_attention(q, selected_keys, selected_vals)

Q2: On Failure Cases and Performance Degradation

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

Q3: On Short-Context Comparisons

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

Q4: On the MagicPIG Comparison Sample Size

We 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:

These expanded results show that Spotlight Attention outperforms MagicPIG across these domains without using a local window or sink tokens.

Q5: On Scaling to Larger Models

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

Thank you for your review. Your feedback has helped us improve the paper. Our responses are detailed below.

Regarding Weaknesses

1. On the Detailed Ablation of the Loss Function

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

2. On the Context for Key Statements

We have provided the specific context for the statements you highlighted and will integrate these clarifications into the manuscript.

• 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:

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

Regarding Questions

Q1. What is the maximum evaluated context length?

In 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:

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

The "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.

Q2. On the computational cost of the ranking loss.

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

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

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