PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Hanshi Sun, Li-Wen Chang, Wenlei Bao, Size Zheng, Ningxin Zheng, Xin Liu, Harry Dong, Yuejie Chi, Beidi Chen

arXiv:2410.21465v3cs.LGcs.CL

TL;DR

As long-context KV caches grow, their memory footprint and per-token access contribute to low inference throughput, while existing methods face accuracy, memory-reduction, or latency trade-offs. ShadowKV stores low-rank pre-RoPE keys, offloads values to the CPU, and uses sparse KV selection with on-the-fly reconstruction. ShadowKV supports up to 6× larger batch sizes and enhances throughput by up to 3.04× on an A100 across various long-context models.

  • Problem

    As long-context KV caches grow, their memory footprint and per-token access contribute to low inference throughput, while existing methods face accuracy, memory-reduction, or latency trade-offs.

    Source: §1 · pp. 1–2

  • Method

    ShadowKV stores low-rank pre-RoPE keys, offloads values to the CPU, and uses sparse KV selection with on-the-fly reconstruction.

    Source: §1 · p. 2, §6 · p. 9

  • Results

    ShadowKV supports up to 6× larger batch sizes and enhances throughput by up to 3.04× on an A100 across various long-context models.

    Source: §6 · p. 10

  • Takeaways & Limitations

    Across the reported long-context models, ShadowKV supports larger batches and higher throughput while maintaining accuracy.

    Source: §6 · pp. 9–10

Abstract

from arXiv · show

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. The code is available at https://github.com/bytedance/ShadowKV.

1 Introduction

ShadowKV targets long-context inference’s GPU-memory and decoding-latency bottlenecks while preserving accuracy. It compresses pre-RoPE keys, offloads values, and uses sparse KV selection; experiments report larger batches and higher throughput.

1 Introduction

  • Long-context KV caches increase memory use and per-token access, while existing approaches face accuracy loss, insufficient memory savings, or decoding latency.Eviction can lose information; dynamic sparse attention keeps KV on the GPU, and naive CPU offloading adds fetch latency.Source: §1 · pp. 1–2

Low-rank Keys and Offloaded Values for Storage:

Accurate KV Selection for Fast Decoding:

  • 6× larger batch sizes and up to 3.04× throughput on an A100 are reported for Llama-3.1-8B; gains also span other long-context models.The reported maxima across models include 2.97× for Llama-3-8B-1M, 2.56× for GLM-4-9B-1M, and 2.66× for Yi-9B-200K.Source: Accurate KV Selection for Fast Decoding · p. 3

2 Related Works

Related approaches reduce KV-cache costs through eviction, dynamic sparse attention, or quantization, each with distinct trade-offs.

Token Eviction

  • Eviction methods retain selected KV pairs, but discarded tokens cannot be recovered and may cause performance degradation or information loss.Examples include StreamingLLM, H2O, LESS, and SnapKV, which use different retention or eviction policies.Source: §2 · pp. 3–4

Dynamic Sparse Attention

  • Dynamic sparse attention retains the full KV cache while selecting entries for attention, so it can reduce inference latency but not memory footprint.Methods include SparQ, Quest, Loki, TriForce, and InfiniGen; ShadowKV instead uses prompt-dependent SVD for key-cache compression.Source: §2 · p. 4

Quantization

  • KV-cache quantization reduces memory consumption by lowering cache bit width, an approach described as orthogonal to ShadowKV.KIVI quantizes keys and values to 2-bit, while Palu decomposes KV weight matrices offline.Source: §2 · p. 4

3 Observations

The observations motivate two design choices: compress the low-rank pre-RoPE key cache, and use chunk landmarks plus outlier retention to select sparse KV pairs accurately.

3.1 Low-Rank Keys and Offloaded Values for Storage

  • Pre-RoPE keys allow 6× compression without performance degradation, and their low-rank subspaces are shared within a sequence but not strongly across sequences.This makes sequence-specific low-rank key storage more suitable than applying low-rank approximations to weights alone.Source: §3.1 · pp. 4–5
  • Because values are not low-rank, storing low-rank keys on the GPU and offloading values to the CPU reduces memory footprint; selected KV pairs can be reconstructed during decoding.Source: §3.1 · p. 5

3.2 Accurate KV Selection for Fast Decoding

  • On 128K contexts, most post-RoPE keys have high cosine similarity to adjacent tokens, so chunk means can serve as landmarks for approximating attention.The analysis uses chunks of eight tokens and identifies poorly approximated tokens as outliers.Source: §3.2 · p. 5
  • A 1.56% sparse KV budget can be selected accurately using landmarks, while 0.2–0.3% outlier chunks are retained to protect accuracy.Outliers may contain dense or critical information and are difficult to approximate.Source: §3.2 · p. 5
  • Temporal locality enables caching selected KV pairs across adjacent decoding steps, reducing computation and data movement by 60%.The reported high repetition rate supports a cache policy for reducing decoding latency.Source: §3.2 · p. 5

4 ShadowKV

ShadowKV combines low-rank key storage, CPU-offloaded values, landmarks, and on-the-fly sparse KV reconstruction. Its equivalent-bandwidth analysis illustrates how these choices support high-throughput decoding.

4.1 Algorithm

  • During pre-filling, ShadowKV compresses each layer’s pre-RoPE key cache, offloads values, stores chunk-mean landmarks, and keeps poorly approximated outliers on the GPU.Source: §4.1 · p. 6
  • During decoding, landmarks rank chunks; values are fetched from the CPU while keys are reconstructed, and temporal locality limits rebuilding to necessary KV pairs.Optimized CUDA kernels reduce computation and data fetching by 60%.Source: §4.1 · p. 6

4.2 Theoretical Equivalent Bandwidth

  • The equivalent-bandwidth accounting combines GPU bandwidth for selection and sparse attention with PCIe bandwidth for value fetching.It counts the selected chunks and predefined outliers in the attention computation.Source: §4.2 · p. 6
  • 7.2 TB/s equivalent bandwidth is calculated for ShadowKV on an A100 under the stated example assumptions, 3.6× the GPU’s 2 TB/s memory bandwidth.The example uses C=8, S=128K, K=256, O=48, and PCIe bandwidth of 31.5 GB/s.Source: §4.2 · p. 6

5 Empirical Evaluation

ShadowKV’s evaluations report reduced KV-cache memory with maintained accuracy, larger serving batches, and higher throughput; ablations examine sparse-budget, chunk-size, and key-rank trade-offs.

5.1 Accuracy Evaluation

  • ShadowKV reduces KV-cache GPU memory 6× while maintaining accuracy; at 128K and a 1.56% sparse budget, it is robust on RULER and consistently outperforms methods on LongBench.Table 1 compares methods on RULER and LongBench.Source: §5.1 · p. 7, Table 1 · p. 7
  • ShadowKV maintains multi-turn accuracy where SnapKV drops from the second round, and remains compatible with MInference pre-filling on RULER contexts from 8K to 256K.The MInference combination sometimes shows slight performance improvement.Source: §5.1 · p. 8, Table 2 · p. 8

5.2 Efficiency Evaluation

  • 6× larger batch sizes and up to 3.04× throughput are reported on an A100, with throughput surpassing the infinite-batch estimate under assumed infinite GPU memory.The comparison is against full attention.Source: §5.2 · p. 8, Table 3 · p. 8

5.3 Ablation Results

  • At the same sparse budgets, ShadowKV surpasses Quest and achieves higher throughput; on most tasks it maintains full-attention accuracy with a 1.56% budget.Figure 8 compares full cache, ShadowKV, and Quest across models and long-context tasks.Source: Figure 8 · p. 9, §5.3 · p. 9

5.2 Efficiency Evaluation

  • As context length grows, full attention runs out of memory while ShadowKV scales to batch size 48 at 60K context and remains functional for longer sequences.Table 4 reports throughput across batch sizes and sequence lengths for Llama-3-8B-1M.Source: Table 4 · p. 8, §5.2 · p. 9

5.3 Ablation Results

  • Increasing chunk size allows larger batches, but accuracy declines above eight while chunk hit rate remains around 60%.Figure 9 shows chunk-size effects on batch size and accuracy.Source: §5.3 · p. 9, Figure 9 · p. 9
  • Accuracy increases with pre-RoPE key rank up to approximately 160, then stabilizes near full-rank performance, with trends varying across tasks.Some tasks show better performance with low-rank approximations.Source: §5.3 · p. 9, Figure 9 · p. 9

6 Conclusion

ShadowKV combines low-rank key storage and value-cache offloading with sparse attention to reduce GPU memory and decoding overhead. The paper reports higher throughput across long-context models and an extension that handles newly generated tokens while maintaining accuracy and optimizing memory.

6 Conclusion

  • ShadowKV reduces GPU memory through low-rank key-cache storage and value-cache offloading, while accurate sparse attention reduces decoding overhead and maintains accuracy.Source: §6 · p. 9
  • Up to 6× larger batch sizes and up to 3.04× throughput are reported on an A100 across several long-context models.The models listed include Llama-3.1-8B, Llama-3-8B-1M, GLM-4-9B-1M, and Yi-9B-200K.Source: §6 · p. 10

A.1 Handling of Newly Generated Tokens

  • ShadowKV+ projects newly generated tokens’ key cache into the low-rank space and maintains accuracy while optimizing memory usage.The extension is evaluated on RULER and LongBench.Source: §A.1 · p. 15, Table 5 · p. 15, Table 6 · p. 15

A.2 Quantitative Analysis of GPU Memory Savings

The appendix gives a quantitative model of ShadowKV’s GPU memory savings and illustrates its estimate for a long-sequence setting.

A.2 Quantitative Analysis of GPU Memory Savings

  • The memory-savings estimate is parameterized by vector size, sequence length, chunk size, selected-chunk budget, outliers, and pre-RoPE key rank.Source: §A.2 · p. 16
  • With M = 1024, C = 8, S = 128K, K = 256, O = 48, and r = 160, the estimated GPU memory savings are 7.08×.Source: §A.2 · p. 16
  • The example presents 7.08× as a theoretical reduction in KV-cache GPU memory footprint for longer sequences and larger batch sizes.Source: §A.2 · p. 16

A.3 Accuracy Results for Yi-9B-200K

The appendix reports accuracy comparisons for Yi-9B-200K and tests ShadowKV’s precision sensitivity using FP8. It reports maintained accuracy with FP8 and comparative performance across the stated benchmarks.

A.3 Accuracy Results for Yi-9B-200K

A.4 Precision Sensitivity

  • The precision-sensitivity experiments use torch.float8_e5m2 after BF16 was used for the main experiments, including to examine sensitivity in SVD computations.Source: §A.4 · p. 16
  • With FP8 precision, ShadowKV maintains accuracy and achieves consistently high performance in the reported evaluations.Table 9 covers RULER at 128K; Table 10 covers LongBench samples exceeding 4K tokens.Source: §A.4 · p. 16, Table 9 · p. 17, Table 10 · p. 17

A.5 Scalability Analysis for Larger Models and Longer Sequences

ShadowKV maintains performance across larger models and longer contexts, while its latency breakdown shows low scaling overhead and gains from overlapping key reconstruction with value fetching. Under memory-constrained conditions, it outperforms Full Attention and Quest at the same sparse budget.

A.5 Scalability Analysis for Larger Models and Longer Sequences

  • ShadowKV maintains robust performance as context length and model size increase, including 1M-context Llama-3-8B-1M and 512K-context Llama-3-70B-1M evaluations.The RULER results are presented in Figure 10 and Table 11.Source: Table 11 · p. 17, Figure 10 · p. 17, §A.5 · p. 17

A.6 Latency Breakdown

  • The reported SVD, reduce, cosine similarity, topK, and gather-computing overhead is low and tends to decrease as sequence length grows.The latency breakdown is for prefilling.Source: Table 12 · p. 18, §A.6 · p. 18
  • Overlapping key-cache recomputation with CPU value-cache fetching significantly reduces decoding latency.Table 13 reports the decoding latency breakdown for a Llama-3-8B-1M Transformer block.Source: Table 13 · p. 18, §A.6 · p. 18

A.7 Efficiency Comparison with Quest

  • ShadowKV’s advantage over Quest comes from fetching only values from the CPU and using a cache mechanism that leverages KV-cache temporal locality.Source: §A.7 · p. 18
  • When GPU memory cannot hold the KV cache, ShadowKV significantly outperforms Full Attention and Quest under the same sparse budget.The comparison is reported for long contexts or high batch sizes.Source: Table 14 · p. 18, §A.7 · p. 18

A.8 Accuracy Contribution of Outlier KV Cache

Outlier chunks improve ShadowKV’s accuracy, with a small outlier budget nearly matching full attention. Compared with InfiniGen, ShadowKV uses a different SVD and cache strategy and maintains accuracy more consistently on complex tasks.

A.8 Accuracy Contribution of Outlier KV Cache

  • When outliers are inadequately managed, ShadowKV’s mean-based landmarks may fall below Quest’s min-max approach.Source: §A.8 · p. 19
  • At 128K context, increasing outlier chunks improves accuracy, especially on complex tasks; 8 outliers (0.049%) outperform Quest and nearly match full attention.The first chunk is described as a significant outlier previously shown to act as an attention sink [52].Source: Table 15 · p. 19, §A.8 · p. 19

A.9 Detailed Comparison with InfiniGen

  • ShadowKV significantly outperforms InfiniGen across downstream tasks and maintains accuracy on complex tasks where InfiniGen shows drops.InfiniGen performs well on simpler tasks such as RULER-N-S1, but drops on RULER-N-MK2, RULER-FWE, and LongBench-LCC.Source: §A.9 · p. 19
  • ShadowKV applies prompt-dependent online SVD to pre-RoPE keys for compression, whereas InfiniGen uses offline SVD on post-RoPE keys and queries for KV selection.Source: §A.9 · p. 19
  • ShadowKV reconstructs selected keys on the GPU and fetches only values from the CPU, overlapping both operations; InfiniGen fetches exact selected KV pairs.Source: §A.9 · p. 19

B Experiment Details

This appendix describes ShadowKV’s implementation and experiment settings, then reports additional benchmark details and results. The Needle In A Haystack evaluation finds minimal impact relative to full attention across several models and context windows.

B.1 System Implementation.

  • The implementation uses PyTorch, dedicated kernels, FlashAttention, and fused CUDA operations, with multi-streams overlapping key-cache reconstruction and value-cache fetching.Most cases use rank 160, chunk size 8, and a sparse KV-cache budget of 1.56%.Source: §B.1 · p. 20

B.2 Dataset Details

  • The benchmark suite covers RULER’s 13 complex tasks, LongBench’s 6 categories and 21 tasks, and Needle In A Haystack retrieval at context lengths up to 1M tokens.Source: §B.2 · p. 20

B.4 InfiniteBench

  • InfiniteBench comprises 10 tasks spanning QA, coding, dialogue, summarization, and retrieval, with an average length of 214K tokens.Source: §B.4 · p. 20

B.3 Needle In A Haystack

  • Across five additional models, ShadowKV has minimal impact versus full attention on semantic understanding across context windows and needle depths, with a slight improvement for Yi-9B-200K.The results are shown in Figure 11.Source: §B.3 · p. 20, Figure 11 · p. 21
Loading 2410.21465v3…