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

Authors: Hanshi Sun, Li-Wen Chang, Wenlei Bao, Size Zheng, Ningxin Zheng, Harry Dong, Yuejie Chi, Beidi Chen
Primary area: foundation or frontier models, including LLMs
Keywords: Long-Context LLM Inference, KV Cache Optimization

Abstract:
With the widespread deployment of long-context large language models (LLMs), there has been a growing demand for efficient support of high-throughput inference. However, as the key-value (KV) cache expands with the sequence length, the increasing memory footprint and the need to access it for each token generation both result in low throughput when serving long-context LLMs. While various dynamic sparse attention methods have been proposed to speed up inference while maintaining generation quality, they either fail to sufficiently reduce GPU memory consumption or introduce significant decoding latency by offloading the KV cache to the CPU. We present ShadowKV, a high-throughput long-context LLM inference system that stores the low-rank key cache and offloads the value cache to reduce the memory footprint for larger batch sizes and longer sequences. To minimize decoding latency, ShadowKV employs an accurate KV selection strategy that reconstructs minimal sparse KV pairs on-the-fly. By evaluating ShadowKV on a broad range of benchmarks, including RULER, LongBench, and Needle In A Haystack, and models like Llama-3.1-8B, Llama-3-8B-1M, GLM-4-9B-1M, Yi-9B-200K, Phi-3-Mini-128K, and Qwen2-7B-128K, we demonstrate that it can support up to 6$\times$ larger batch sizes and boost throughput by up to 3.04$\times$ on an A100 GPU without sacrificing accuracy, even surpassing the performance achievable with infinite batch size under the assumption of infinite GPU memory.

Decision: reject
Reviewer scores: 8, 6, 5, 8

Reviews and discussion:
<thread>
creation date:
1729566594314

thread_type:
official_review

<thread_object>
creator: reviewer_irxd

<content>
summary:
The paper introduces SHADOWKV, a novel system that improves inference throughput for long-context LLMs. The key challenge addressed is the increasing memory footprint of KV caches as sequence length increases, which slows down inference. SHADOWKV offers a solution by offloading the value cache to the CPU while keeping a low-rank key cache on the GPU, reducing memory consumption without sacrificing performance. By employing a method of KV selection for sparse attention, it boosts throughput and supports larger batch sizes, showing improvements of up to 3.04× throughput on the A100 GPU.

strengths:
The paper tackles a significant problem in LLM inference by reducing GPU memory usage through a hybrid CPU-GPU cache system, enabling long-context LLMs to operate more efficiently without sacrificing accuracy.

The system shows impressive throughput gains, handling up to six times larger batch sizes compared to baseline methods, which could substantially impact real-world LLM deployment.

SHADOWKV is tested on multiple models and benchmarks (e.g., Llama, GLM) across various tasks, demonstrating consistent performance improvements across all settings.

weaknesses:
The main issue is that, as far as I know, SVD is precision sensitive, but I didn't find any discussion about precision in the paper. My main question is what precision is used for ShadowKV and baselines. If you are using precision like FP16/FP32, my question is how does ShadowKV work on FP8/(FP8 & FP16) precision? If the precision is FP8, how does ShadowKV survive from precision-sensitive SVD?

For the speed evaluation, I only found the throughput (tokens/s), are there any experiments for the time of each operation separately?

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730386715054

thread_type:
official_review

<thread_object>
creator: reviewer_5wsx

<content>
summary:
This paper introduces SHADOWKV, a CPU-offloading-based system for long-context LLM inference. SHADOWKV addresses GPU memory constraints by leveraging the low-rank property of the key cache, storing a low-rank pre-ROPE key cache on the GPU while offloading the value cache to the CPU. Additionally, SHADOWKV stores landmarks—mean values of post-ROPE keys for adjacent tokens—along with low-rank pre-ROPE key cache in GPU memory. During decoding, SHADOWKV utilizes these landmark keys to identify significant tokens, selectively recovering their keys and retrieving their values from the CPU. Evaluations on long-context benchmarks demonstrate that SHADOWKV maintains high accuracy while significantly reducing GPU memory consumption and inference latency.

strengths:
* In-depth analysis on low-rank nature of key cache in comparison to value cache as well as weights.
* Leveraging spatial locality of post-ROPE key cache for dynamic sparse attention appears both novel and effective.
* Empirical results are impressive.

weaknesses:
* Incomplete Descriptions

The term "sparse budget" is not clearly defined in the paper, which may lead to confusion. Additionally, while SHADOWKV claims to leverage the temporal locality of the KV cache to reduce computation and communication (by approximately 60%), it lacks any detailed explanations on what that feature is.

* Handling Newly Generated Tokens

While the paper says that it excludes the handling of newly generated tokens for simplicity, this issue is quite significant and should not be ignored. If not addressed, the KV cache for newly generated tokens could negate SHADOWKV’s key benefits of reduced GPU memory usage and lower inference latency, especially with long output sequences. Incorporating mechanisms to handle these tokens within SHADOWKV is essential, and the authors should evaluate and report on its impact on accuracy.

* Lack of Comparison with Infinigen

The paper does not sufficiently compare SHADOWKV with Infinigen [1], a closely related work that similarly stores low-rank key cache in GPU memory, offloads the value cache to the CPU, and selectively fetches important values based on approximate attention scores. Although the paper briefly discusses Infinigen, given the significant similarities, a more in-depth comparison with Infinigen should be made in order to highlight the main differentiator of SHADOWKV.

[1] Lee et al., "InfiniGen: Efficient Generative Inference of Large Language Models with Dynamic KV Cache Management", OSDI'24

questions:
* What exactly is the “sparse budget”?

* How does SHADOWKV leverage the temporal locality of the KV cache?

* What are the exact GPU memory savings of SHADOWKV? Including a quantitative discussion on GPU memory savings in the paper would be helpful.

* How can SHADOWKV handle the KV cache for newly generated tokens, and what would be the impact of this?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730583705434

thread_type:
official_review

<thread_object>
creator: reviewer_4r5o

<content>
summary:
The increasing KV-cache poses great challenges to long-context LLM inference. This work presents a long-context LLM inference system, SHADOWKV. It decreases the memory footprint by storing the low-rank key cache and offloading value cache to CPU, and reduces the decoding latency by reconstructing the sparse KV pairs on-the-fly. Evaluations show that SHADOWKV supports up to 6x larger batch sizes and improves throughput by up to 3X on A100 GPU while maintaining the accuracy.

strengths:
+ This work presents a very interesting observation on the KV cache: pre-RoPE key cache is exceptionally low-rank compared to post-RoPE key cache, value cache and KV projection weights. Built on this observation, the proposed SHADOWKV significantly reduces the memory footprint of the Key cache.
+ The work also improves the previous sparse attention work including QUEST by introducing outlier KV cache.
+ This work also implements the inference system and shows actual throughput improvement on the real world A100 GPUs.

weaknesses:
- The proposed method is complex, including low-rank K cache, CPU offloaded V cache, outlier KV cache, and dynamic sparse attention. However,  the ablation study on each component is missing.
   - In terms of model accuracy: 
      - it is unclear how much accuracy improvement an extra outlier KV cache will bring.
      - previous work Quest uses Min-Max as landmark cache, ShadowKV adopts Mean as landmark cache. It is unclear how much accuracy improvement this change will bring.
   - In terms of efficiency:
      - Authors only show a rough prefiling latency breakdown in Figure 1(c). It is unclear unclear how long it takes for computing the outlier cache (i.e., reduce, cosine-similarity, top-k, gather) in the profiling stage, how long it takes for KV cache chunk selection (i.e, MatMul, Softmax, Max, TopK) in the decoding stage, how long it takes for recomputing the K cache from low-rank cache. These overheads seem to increase linearly with the context length. It would be better to see the efficiency breakdown on every part of the system under different context lengths (e.g., 128K, 256K, 512K).
      - it is unclear how SHADOWKV performs for extremely long context. For example, the authors evaluated on Llama-3-8B-1M but only with up to 128K context length.
     - this work lacks efficiency comparison against previous work LoKi, Quest and MInference.

questions:
My questions are listed in the weakness section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730692687035

thread_type:
official_review

<thread_object>
creator: reviewer_sej1

<content>
summary:
This paper presents a high throughput LLM inference system for long sentence length. It proposes a compression method to leverage the low rank property of key cache and improve the sparse attention method by accurate KV selection.

strengths:
This paper is well-organized and clearly written. The proposed method is well-motivated, addressing relevant challenges, and is supported by thorough analysis. The evaluation is comprehensive and robust, effectively substantiating the claims and demonstrating thoughtful considerations. Overall, this is a strong submission for ICLR.

weaknesses:
1. As shown in Fig. 8, some downstream tasks, such as 'Frequent Words Extraction,' perform significantly worse with sparse KV enabled. A brief analysis of why this approach underperforms for these types of tasks would be helpful, as well as any potential solutions to address these limitations.
2. The proposed solution is currently evaluated on an 8B model with a 128K sequence length. It would strengthen the paper to include an analysis of whether this approach scales effectively for larger models, such as a 70B model with an extremely long sequence length of 1M.

questions:
As listed in the weakness part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732143726610

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Revised Manuscript updated by Authors

comment:
We thank reviewers [R1(sEj1), R2(4R5o), R3(5wSX), R4(iRxd)] for their thoughtful and highly supportive feedback! We were glad that the reviewers found the problem **significant and interesting [R2, R4]**, the observations and analysis **insightful and highly valuable [R1, R2, R3]**, the method **novel and well-motivated [R1, R3, R4]**, the presentation **well-organized and easy to follow [R1]**, experimental results **strong and impressive [R1, R3, R4]**. 

We have updated the paper to incorporate the constructive suggestions, reflected in the revised version. We use **blue color for all newly added content**. Here is a summary of the major changes:

- [R1, R2, R3] **Scalability to larger models, longer contexts, and longer outputs**:
	- Added experiments for Llama-3-70B-1M on RULER with 512K contexts and Needle In A Haystack with up to 1M contexts, demonstrating ShadowKV’s scalability with larger models (mentioned in $\textcolor{blue}{\text{Section 5.1, Page 7}}$ and detailed in $\textcolor{blue}{\text{Appendix A.5, Page 17}}$).
	- Added experiments for Llama-3-8B-1M on RULER with 1M contexts to show ShadowKV’s ability to scale with longer sequences (mentioned in $\textcolor{blue}{\text{Section 5.1, Page 7}}$ and detailed in $\textcolor{blue}{\text{Appendix A.5, Page 17}}$).
	- Added experiments for handling newly generated tokens to validate ShadowKV’s scalability for longer outputs (mentioned in $\textcolor{blue}{\text{Section 4.1, Page 6}}$ and detailed in  $\textcolor{blue}{\text{Appendix A.1, Page 15}}$).
- [R3] **Comparison with InfiniGen**
	- Added a discussion of the differences between InfiniGen and ShadowKV ($\textcolor{blue}{\text{Section 2, Page 3}}$).
	- Included comparative experiments with InfiniGen, showing that ShadowKV outperforms InfiniGen on various downstream tasks. Notably, on the RULER benchmark, ShadowKV achieves an accuracy improvement of 10% over InfiniGen ($\textcolor{blue}{\text{Section 5.1, Page 7-8}}$).
- [R2, R4] **Latency breakdown**
	- Added a latency breakdown for the prefilling stage, demonstrating that ShadowKV introduces minimal overhead, which decreases as sequence length increases, affirming ShadowKV’s scalability with longer sequences (mentioned in $\textcolor{blue}{\text{Section 5.2, Page 9}}$ and detailed in  $\textcolor{blue}{\text{Appendix A.6, Page 18}}$).
	- Added a latency breakdown for the decoding stage, showing how overlapping the recomputation of the key cache with value cache fetching from the CPU effectively reduces decoding latency (mentioned in $\textcolor{blue}{\text{Section 5.2, Page 9}}$ and detailed in  $\textcolor{blue}{\text{Appendix A.6, Page 18}}$).
- [R4] **Numerical stability**
	- Added experiments to show that ShadowKV can retain its accuracy at lower precision, e.g., FP8 (mentioned in $\textcolor{blue}{\text{Section 5.3, Page 10}}$ and detailed in  $\textcolor{blue}{\text{Section A.4, Page 16}}$).
- [R2] **Extra ablations**
	- Added experiments demonstrating the effectiveness of outliers and mean landmarks (compared to Quest’s min-max landmarks), highlighting the improved performance of our method (mentioned in $\textcolor{blue}{\text{Section 5.3, Page 10}}$ and detailed in  $\textcolor{blue}{\text{Appendix A.8, Page 19}}$).
	- Included an efficiency comparison with Quest, showing that our method outperforms Quest in both accuracy and efficiency. Specifically, ShadowKV is up to 4.85$\times$ faster than Quest while also achieving superior accuracy ($\textcolor{blue}{\text{Appendix A.7, Page 18}}$).
</content>
</thread_object>

</thread>


Metareview:
metareview:
After reviewing the entire discussion and considering input from all reviewers, I recommend the rejection of this paper. While the submission has merits, the unresolved concerns outweigh the strengths.  

**Strengths**:  
- The paper presents a novel observation about efficiently compressing pre-RoPE Key caches, reducing CPU-GPU communication costs. This is a potentially useful insight that has practical implications for offloading methods.  
- The revised methodology demonstrates improved scalability for long sequences while maintaining accuracy.  

**Weaknesses**:  
1. **Incremental Contribution**:  
   The proposed work builds on existing methods of CPU offloading and low-rank approximation. However, the enhancements, while effective, appear to be incremental rather than groundbreaking.   

2. **Presentation and Positioning**:  
   The paper struggles with positioning itself clearly within the existing literature. The unique contributions remain insufficiently distinguished from prior work, leaving ambiguity about its originality. This issue undermines the perceived value of the paper's contributions.  


Given these mixed reviews, the lack of consensus among reviewers, and the remaining issues of novelty and presentation, I believe the paper does not meet the threshold for acceptance in its current form. While the work demonstrates promise, a more substantial contribution and clearer articulation of its impact are necessary to justify publication.