ShadowKV: KV Cache in Shadows for High-Throughput Long-Context LLM Inference

Q2: Detailed Time Analysis of Operations

Thank you for the valuable suggestion. We present a detailed latency breakdown in tables below to illustrate the efficiency of each operation under various context lengths for both the prefilling and decoding stages (mentioned in $\textcolor{blue}{\text{Section 5.2, Page 9}}$ and detailed in $\textcolor{blue}{\text{Appendix A.6, Page 18}}$).

• Scalability for Longer Sequences. As shown in table, the overhead of SVD, reduce, cosine similarity, topK, and gather computing is very low and tends to decrease as the sequence scales, proving that ShadowKV's scalability to longer sequences.

Latency breakdown (ms) of a Transformer block of Llama-3-8B-1M during prefilling:

• Overlapping Operations for Latency Reduction. In the table below, we demonstrate how overlapping the recomputation of the key cache with value cache fetching from the CPU significantly reduces decoding latency. This concurrent processing approach ensures that ShadowKV minimizes overhead when handling long-context models.

Latency breakdown (ms) of a Transformer block of Llama-3-8B-1M during decoding:

Thank you for your encouraging feedback and insightful remarks. We have thoroughly addressed each of your questions and hope our responses will lead you to consider raising your score.

Q1: Precision Sensitivity in SVD Computations

We appreciate the suggestion to examine the precision sensitivity of ShadowKV. In the main experiments, we used BF16 for both model weights and KV cache. To further investigate the impact of precision on ShadowKV’s performance, we conducted additional experiments using FP8 precision (torch.float8_e5m2), showing that ShadowKV can retain its accuracy at this lower precision, addressing concerns about precision sensitivity, particularly in SVD computations (mentioned in $\textcolor{blue}{\text{Section 5.3, Page 10}}$ and detailed in $\textcolor{blue}{\text{Section A.4, Page 16}}$).

As detailed in the tables below, ShadowKV and baseline methods were evaluated using FP8. Results show that ShadowKV maintains accuracy and achieves consistently high performance even with FP8 precision. This robustness, despite FP8’s reduced numerical range, confirms that ShadowKV can continue to deliver efficiency gains without compromising accuracy.

• Results on RULER

• Results on LongBench

Dear Reviewer iRxd,

We are truly grateful for the time and effort you have taken to review our manuscript and for providing such thoughtful feedback. Your insights have been instrumental in helping us refine our work.

As today marks the final day for revising the PDF, we wanted to kindly follow up to see if you have any further questions or points of clarification regarding our responses. We are happy to address any remaining concerns you might have.

Thank you again for your invaluable support and guidance throughout this process.

With our deepest appreciation,

Authors

Thank you for your thoughtful question and for pointing this out. The SVD operation is conducted in FP32, as PyTorch currently does not support lower-precision kernels for SVD. To evaluate the robustness of ShadowKV, we cast the decomposition results to FP8, optimizing KV cache storage while ensuring performance remains robust. During decoding, the K reconstruction (i.e., matrix multiplication) is performed directly in FP8 or BF16, just like other FP8 or BF16 LLM decoding operations. Moreover, this SVD part can be executed on the CPU during prefilling, which can be overlapped, mitigating the overhead.

Additionally, we explored PyTorch's approximation SVD method, gesvda, and found that it does not degrade ShadowKV's performance, further demonstrating its precision robustness. Once PyTorch introduces lower-precision SVD kernels, we would be happy to evaluate their impact on our method.

If our responses have adequately addressed your concerns, we would be deeply grateful if you could consider reflecting this in your score. Please let me know if you have any further questions or suggestions.

Thank you sincerely for your thoughtful review and invaluable feedback. We have addressed each of your questions with care and hope that our responses will inspire you to consider raising your score.

Q1: Definition of “sparse budget”

Thanks for pointing it out. The term “sparse budget” refers to the number of selected tokens, i.e. the K of TopK. This budget dictates the portions that need to be fetched from the CPU or reconstructed from the low-rank space. We have added this clarification to the revised paper ($\textcolor{blue}{\text{Section 1, Page 2}}$).

Q2: Explanation of “temporal locality of the KV cache”

Temporal locality of the KV cache refers to the observation that, during decoding, the KV cache pairs selected by the queries of two adjacent decoding steps have a repetition rate of approximately 60%. This means that we don’t need to perform low-rank reconstruction or CPU data fetching for the repeated portions, only for the non-repeated parts. This reduces the overall decoding overhead. We have added this clarification to the revised paper ($\textcolor{blue}{\text{Section 3.2, Page 5}}$).

Q3: Handling of Newly Generated Tokens

We appreciate the suggestion to evaluate the impact of handling newly generated tokens. We present extensive experiments on the RULER and LongBench across different long-context models (mentioned in $\textcolor{blue}{\text{Section 4.1, Page 6}}$ and detailed in $\textcolor{blue}{\text{Appendix A.1, Page 15}}$). The results demonstrate that ShadowKV effectively handles newly generated tokens while maintaining accuracy.

To address the handling of newly generated tokens, we project these tokens' key cache into a low-rank space using the same projections applied during the prefilling phase. This approach preserves the benefits of reduced GPU memory usage, particularly for long output sequences.

As shown in tables below, we refer to this extension as ShadowKV+. Our evaluation across various models demonstrates that ShadowKV+ effectively maintains accuracy and manages newly generated tokens as well.

• Results on RULER:

• Results on LongBench:

Q4: Detailed Comparison with InfiniGen

We provide further clarification on the key distinctions and conduct additional experiments between ShadowKV and InfiniGen (mentioned in $\textcolor{blue}{\text{Section 2, Page 3}}$ and detailed in $\textcolor{blue}{\text{Section 5.1, Page 7-8}}$). These experiments show that ShadowKV significantly outperforms InfiniGen across a wide range of downstream tasks.

• Differences in SVD Usage. Infinigen performs an offline SVD to get a projection matrix, which is applied to post-RoPE key and query states for KV selection, while ShadowKV applies an online, prompt-dependent SVD directly to the pre-RoPE key cache for compression, not for KV selection.

• Methodological Differences. While InfiniGen uses SVD for KV selection, it requires fetching selected, exact KV pairs from the CPU. In contrast, ShadowKV only fetches the value cache from the CPU, reconstructing the key cache from its low-rank storage on the GPU. By overlapping these processes, ShadowKV reduces data-fetch overhead and achieves improved efficiency in KV cache management.

• Accuracy Comparison. To empirically validate ShadowKV’s advantages, we conducted accuracy evaluations. Results, presented in the tables below, confirm ShadowKV’s effectiveness in maintaining accuracy while optimizing memory usage. Although InfiniGen performs well on simpler tasks like RULER-N-S1, it shows significant accuracy drops on more complex tasks, such as RULER-N-MK2, RULER-FWE, LongBench-LCC, and others, where ShadowKV maintains consistently high accuracy.

• Results on RULER:

• Results on LongBench:

Thanks for your response. The sparse budget is calculated as the product of the selected chunk budget (K) and the chunk size (C), i.e., $\text{sparse budget} = K \times C$.

For LongBench, when we say “KV cache budget is set to 256,” it means the sparse budget is 256. If the chunk size is 8, then $K = 256/8=32$.

Feel free to reach out if you have any additional questions. Thanks!

Thank you for your feedback! I'm very glad the explanation clarified things for you. Yes, you are correct. In the efficiency test, we demonstrate that ShadowKV achieves 6x larger batch sizes (refer to $\textcolor{blue}{\text{Table 4}}$), which aligns well with our theoretical analysis.

For RULER-128K, we use a budget of $K = 256$ . For LongBench, since the context lengths are shorter ($\le$32K), we opt for a smaller $K$ to better match the benchmark’s requirements.

Feel free to reach out if you have any additional questions. Thanks again for your engagement!

Hello, thank you for your response. Our cache implementation is actually quite straightforward—we avoid complex mechanisms like LRU, ensuring no additional memory consumption. Specifically, during each decoding step, we overwrite the KV pairs from the previous step using the same buffer ($2(K+O)C$ in the formula), effectively reusing memory without allocating extra space.

If more complex cache rules are used, the hit rate can be improved, but it will bring some additional memory overhead.

Thank you for your thoughtful feedback and for pointing out your concerns. Let me clarify the points you raised.

The term "KV buffers" refers to the $2KC$, i.e., sparse budget component in this context (for LongBench, we set $O = 1$ for ShadowKV). To implement a caching mechanism, both InfiniGen and ShadowKV must retain at least this portion of the memory, which reduces the overall memory savings. Therefore, your concern about the effect of $K$ on memory savings applies equally to InfiniGen. Moreover, InfiniGen introduces additional memory overhead by saving the low-rank projection matrix, which is also not accounted for in the $6.67\times$ memory savings.

Our primary goal with ShadowKV is to optimize for long-context scenarios ($S >16K$, benchmarks like RULER), where its scalability becomes more evident. To provide additional clarity, we have included a table showing memory savings for ShadowKV across different sequence lengths $S$ and $K = 32$. It’s worth noting that, for the LongBench dataset, we only evaluate samples with sequence lengths greater than $4K$.

As shown in the table below, ShadowKV matches InfiniGen's memory savings (which is overestimated) at a sequence length of $8K$ and surpasses it as the sequence length increases. We believe that ShadowKV demonstrates better performance in long-context scenarios compared to InfiniGen and offers superior scalability.

We hope this table and explanation address your concerns and provide further insight into ShadowKV's scalability and performance. Please feel free to reach out with any additional questions or clarifications. We are looking forward to your feedback.

Thank you for following up and for your thoughtful comments. We appreciate the opportunity to provide additional details regarding the comparison between ShadowKV and InfiniGen in terms of GPU memory savings and latency. Below, we address your questions and provide a detailed analysis, demonstrating that ShadowKV outperforms InfiniGen in terms of accuracy, memory savings, and latency.

GPU memory savings: InfiniGen does not provide a quantitative analysis of GPU memory savings in its paper, so we performed an estimation to highlight ShadowKV’s advantages. Following the configuration described in the InfiniGen paper, where the partial weight ratio is set to 0.3, the KV cache size is reduced to 15%, equating to a $6.67\times$ memory savings. In comparison, ShadowKV achieves a $7.08\times$ reduction. It should be noted that InfiniGen’s 15% memory savings does not account for additional memory overheads (e.g., KV buffers). If these overheads are included, the memory savings for InfiniGen would be further diminished.

Latency Analysis: From a latency perspective, ShadowKV is more scalable to long sequences than InfiniGen. In our accuracy evaluation, we maintained the same sparse budget for both methods. Below is the latency breakdown (in ms) for a single Transformer block of Llama-3-8B-1M during decoding, using ShadowKV.

ShadowKV overlaps the recomputation of the K cache with the retrieval of the V cache from CPU. In contrast, InfiniGen overlaps the KV cache fetching time with other computations (GEMM, Softmax, Max, TopK, Attention, FFN, and QKV operations). However, as the sequence length increases, InfiniGen's ability to hide data-fetching costs diminishes. For instance, at a sequence length of $12 \times 256K$, the cost of other computations is 1.37 ms (0.65+0.07+0.16+0.19+0.25+0.05):

• For ShadowKV, the latency is $1.75+1.37=3.12\text{ms}$, as K cache recomputation is effectively overlapped.
• For InfiniGen, the KV cache fetching time alone is already $1.75\times 2=3.5\text{ms}$ (InfiniGen fetches both K and V, while ShadowKV only fetches V from CPU). Thus, at longer sequence lengths, the other computations in InfiniGen cannot sufficiently hide the data-fetching cost.

Therefore, under the same sparse budget, ShadowKV not only outperforms InfiniGen in latency but also achieves significantly better accuracy on complex tasks, where InfiniGen suffers from a notable accuracy drop.

We hope this analysis addresses your concerns and provides a comprehensive understanding of ShadowKV's effectiveness. Thank you for your insightful feedback, and we remain open to any further questions or discussions.

Thank you for detailed review and valuable feedback. We appreciate the reviewer’s recognition of the novelty and effectiveness of our method. Below, we address concerns regarding the lack of some ablation results. We hope the reviewer can consider raising your score in light of our response.

Q1: Accuracy Contribution of Outlier KV Cache and Mean Landmarks

Our additional experiments demonstrate that outliers play a critical role in capturing essential information, even in small numbers (0.049%), significantly enhancing the performance of mean-based landmarks and outperforming the min-max landmarks used in Quest.

We conduct experiments using different numbers of outlier chunks for Llama-3-8B-1M on the RULER benchmark with 128K context length (mentioned in $\textcolor{blue}{\text{Section 5.3, Page 10}}$ and detailed in $\textcolor{blue}{\text{Appendix A.8, Page 19}}$). As presented in the table below, our findings indicate that outliers play a crucial role. For instance, the first chunk, a significant outlier, has previously been shown to act as an attention sink [1], underscoring its importance in maintaining model accuracy.

The results demonstrate that increasing the number of outlier chunks has a positive impact on accuracy, especially in complex tasks. This indicates that even a small number of outliers can effectively capture essential information, reducing the need for full attention. Remarkably, with just 8 outliers (0.049%), ShadowKV outperforms the Quest baseline and nearly matches the accuracy achieved by full attention.

However, when outliers are not adequately managed, the performance of the mean-based landmarks in ShadowKV may fall below the min-max approach used by Quest, underscoring the importance of handling outliers properly.

[1] Efficient Streaming Language Models with Attention Sinks

Q2: Detailed Efficiency Breakdown Across System Components

Thank you for the valuable suggestion. We present a detailed latency breakdown in tables below to illustrate the efficiency of each operation under various context lengths for both the prefilling and decoding stages (mentioned in $\textcolor{blue}{\text{Section 5.2, Page 9}}$ and detailed in $\textcolor{blue}{\text{Appendix A.6, Page 18}}$).

• Scalability for Longer Sequences. As shown in table, the overhead of SVD, reduce, cosine similarity, topK, and gather computing is very low and tends to decrease as the sequence scales, proving that ShadowKV's scalability to longer sequences.

Latency breakdown (ms) of a Transformer block of Llama-3-8B-1M during prefilling:

• Overlapping Operations for Latency Reduction. In the table below, we demonstrate how overlapping the recomputation of the key cache with value cache fetching from the CPU significantly reduces decoding latency. This concurrent processing approach ensures that ShadowKV minimizes overhead when handling long-context models.

Latency breakdown (ms) of a Transformer block of Llama-3-8B-1M during decoding:

Q3: Performance on Extremely Long Context Lengths

We clarify that our paper includes some experiments with extremely long contexts. We tested Llama-3-8B-1M on the NIAH dataset with up to 1M tokens ($\textcolor{blue}{\text{Figure 6, Page 8}}$).

Here we present our additional results, showing that ShadowKV matches the performance of full attention while outperforming other sparse methods on 1M contexts with Llama-3-8B-1M and 512K contexts with Llama-3-70B-1M, as demonstrated on the RULER benchmark (detailed in $\textcolor{blue}{\text{Section 5.1, Page 7}}$ and $\textcolor{blue}{\text{Appendix A.5, Page 17}}$). Additionally, ShadowKV achieves perfect needle retrieval accuracy with Llama-3-70B-1M, evaluated across context lengths ranging from 16K to 1M tokens ($\textcolor{blue}{\text{Figure 10, Page 17}}$).

These findings underline ShadowKV’s ability to handle large-scale inputs effectively, offering robust performance across increasing context lengths and model sizes. This scalability ensures its suitability for real-world applications that require extensive sequence handling and larger model capabilities.

Thank you for the supportive comments and recognizing the novelty of our method and the thorough evaluations. We hope our detailed clarifications and additional experimental results will address the concerns regarding our work.

Q1: Analysis and Solutions for Sparse KV Underperformance in Certain Tasks

For certain downstream tasks, attention distributions can exhibit a long tail effect, where a larger portion of tokens receive small yet non-negligible attention, making exact top-K selection less effective. As observed in Differential Transformer [1], this issue arises due to the nature of the softmax function [2], which tends to produce noise and disperse attention scores.

Since Top-K serves as the theoretical upper bound for sparse attention methods aiming to approximate it, poor performance of Top-K inherently limits the performance of sparse attention methods.

We include TopK results for ‘Frequent Words Extraction’ as a reference below. As shown in the table, sparse attention underperforms due to limitations in TopK’s performance.

A possible solution is to promote sparsity during the pre-training or SFT phases, co-designing it with the training process. By encouraging sparsity throughout training, the model may learn to manage long tail effects within the constraints of a sparse budget and reduce noise, allowing for a more reliable sparse attention mechanism across diverse downstream tasks.

[1] Differential Transformer

[2] softmax is not enough (for sharp out-of-distribution)