PaperlayerView exact revision on arXiv

Source-linked AI summary

MosaicKV: Serving Long-Context LLM with Dynamic Two-D KV Cache Compression

Sheng Qiang, Ruiwei Chen, Yinpeng Wu, Jinyu Gu, Zhichao Hua, Yubin Xia, Binyu Zang, Haibo Chen

arXiv:2607.00760v1cs.LGcs.DC

TL;DR

Long-context serving strains GPU memory as KV caches grow, while existing compression approaches leave limited headroom and direct two-dimensional compression can hurt accuracy. MosaicKV uses dynamic two-D compression and compressed-cache management; on an H800 it reports faster attention and decoding, higher throughput, lower memory use, and limited accuracy loss.

  • Problem

    Long-context KV caches scale with sequence length and strain GPU memory, while single-dimension compression leaves limited headroom and directly combining dimensions can degrade accuracy.

    Source: §1 · pp. 1–2

  • Method

    MosaicKV selects important elements per KV vector, tailors two-dimensional compression by cache segment, and manages compressed caches using underutilized GPU and CPU resources.

    Source: §1 · p. 2

  • Results

    Compared with the uncompressed baseline, MosaicKV achieves up to 16× attention speedup, 4.8× lower decode latency, 7.3× higher throughput, 3× lower GPU memory usage, and 1.76% average accuracy loss on LongBench and RULER.

    Source: §1 · p. 3

  • Takeaways & Limitations

    MosaicKV applies two-dimensional KV compression throughout decode while retaining near-baseline accuracy and improving serving performance.

    Source: §1 · p. 3

    MosaicKV currently focuses on decode; prefill acceleration with chunked prefill is left as future work.

    Source: §7 · p. 12

Abstract

from arXiv · show

Long-context LLM services now sustain prompts with hundreds of thousands to millions of tokens, making the key-value (KV) cache a first-order serving cost. Because the cache grows linearly with context length, it can exhaust GPU memory, force smaller batches, and reduce serving throughput. Prior KV cache compression techniques typically target only the sequence dimension or only the channel dimension, which leaves limited headroom as context windows scale. Compressing both dimensions promises higher memory reduction, but applying the two forms of compression directly leads to significant accuracy loss. This paper introduces MosaicKV, a dynamic two-D (dimensional) KV cache compression system for extremely long-context serving. MosaicKV uses dynamic two-D compression to address the accuracy challenge, exploiting the non-uniform importance distribution of elements within the KV cache. Instead of applying one compression pattern globally, MosaicKV identifies important elements for each KV vector and selects compression strategies at the granularity of KV cache segments. To address the performance challenge, where fine-grained sparsity and compression management overhead can offset the gains from compression, MosaicKV introduces compressed KV cache management. This mechanism uses underutilized GPU and CPU resources to maintain compressed KV caches and accelerate attention computation. Evaluation on an H800 GPU with multiple LLMs shows that MosaicKV delivers up to 16x attention speedup, 4.8x lower decode latency, and 7.3x higher throughput than the uncompressed baseline. At the same time, it reduces memory usage by 3x and incurs only 1.76% average accuracy loss on LongBench and RULER.

1 Introduction

Long-context serving makes KV memory a bottleneck, while directly combining sequence and channel compression can sharply reduce accuracy. MosaicKV uses dynamic two-D compression and compressed-cache management to address accuracy and performance challenges.

1 Introduction

  • Prior compression methods target sequence or channel dimensions separately, leaving limited headroom as context windows expand.Source: §1 · p. 2
  • Directly applying channel and sequence compression can sharply reduce accuracy; 70% channel compression atop sequence compression causes 82.8% accuracy loss.Source: §1 · p. 2
  • MosaicKV selects important elements per KV vector and partitions the cache into segments with tailored compression strategies, enabling two-dimensional compression throughout decode.The approach addresses non-uniform element importance across vectors and cache regions.Source: §1 · p. 2
  • Fine-grained sparse access and compression management can hurt performance, so MosaicKV uses underutilized GPU and CPU resources for packed sparse attention and cache management.GPU-side buffering supports interim decode computation; CPU-side operations prepare recompressed segments for a nonblocking switch.Source: §1 · pp. 2–3
  • Compared with the uncompressed baseline, MosaicKV achieves up to 16× attention speedup, 4.8× lower decode latency, and 7.3× higher throughput, while reducing GPU memory usage by 3× with 1.76% average accuracy loss.Accuracy was measured on LongBench and RULER.Source: §1 · p. 3

2 Background and Motivation

Long contexts make the KV cache a primary memory bottleneck, while existing methods compress along sequence or channel dimensions. SVD-based channel compression rotates vectors to concentrate information in leading channels.

2.1 Large KV Cache in Long-Context LLM Inference

  • For LLaMA-3.1-8B serving eight concurrent requests, the KV cache requires 128 GB at 128K tokens and 1 TB at 1M tokens.Attention accounts for 84.7% and 97.9% of total Time-Between-Tokens at those respective context lengths.Source: §2.1 · p. 3

2.2 Existing KV Cache Compression Methods

  • Sequence compression selects important tokens for attention, while channel compression retains selected dimensions of KV vectors.Quest selects tokens using query-aware scores; ThinK uses a uniform channel mask computed from prefill keys.Source: §2.2 · pp. 3–4
  • Single-dimensional compression has scaling limits: sequence compression can degrade accuracy and must retain the full cache, while combining dimensions has not covered the full decode lifecycle.Existing approaches partially combine dimensions, but no cited system applies two-dimensional compression throughout token selection and attention computation.Source: §2.2 · p. 4

2.3 Singular Value Decomposition (SVD)

  • SVD rotates K so high-magnitude information concentrates in leading columns, enabling channel compression by retaining only the top-r channels.The rotated KV cache must be transformed back to the original space during attention.Source: §2.3 · p. 4

3 Overview of MosaicKV

MosaicKV targets long-context KV-cache compression along both channel and sequence dimensions, adapting compression to non-uniform element importance. Its overview pairs that approach with cache-management techniques intended to address accuracy and performance challenges.

3.1 Goals

  • MosaicKV targets low memory use, high performance, and high accuracy by compressing both K and V caches along channel and sequence dimensions.The stated goals include improved decode latency and throughput with minimized accuracy loss.Source: §3.1 · p. 4

3.2 Challenges and Approaches

  • 24.5% and 82.8% accuracy loss versus sequence compression alone (Quest) occurred at 30% and 70% channel compression, respectively.These losses were measured for the strawman combining channel compression with Quest.Source: §3.2 · p. 5
  • Only 62.28% of top-25% important elements fall in globally top-25% channels; MosaicKV therefore selects per-vector elements and adapts strategies by segment.The importance distribution also varies across KV-cache regions, motivating different compression rates and rotation matrices.Source: §3.2 · p. 5
  • Fine-grained sparse accesses and decode-stage cache management can push performance below the uncompressed baseline, motivating management that uses underutilized resources.The overview also identifies imbalanced utilization: GPU memory bandwidth is heavily used while CUDA cores and CPU resources are underutilized.Source: §3.2 · p. 5

3.3 System Overview

  • MosaicKV combines per-vector channel selection with query-based token selection, then applies segment-specific strategies to the KV cache.Channel compression occurs during prefill and decode, while sequence compression selects top-scoring tokens before each attention layer.Source: §3.3 · p. 6
  • Packed sparse attention uses CUDA cores to compute with packed KV data, while CPU-assisted buffering and asynchronous strategy generation manage compression during decode.The management design also uses compressed encoding for prefill and buffering to replace temporary segments with final compressed versions.Source: §3.3 · p. 6

4 Dynamic Two-D Compression

MosaicKV performs two-dimensional KV-cache compression in two stages, then adapts compression strategies to cache segments. Channel compression selects important elements per vector, while sequence compression selects tokens dynamically during decode.

4.1 Basic Method

  • The KV cache is partitioned into segments, each receiving a suitable strategy for two-stage compression of K and V before attention computation.The strategy is assigned dynamically to balance performance and accuracy.Source: §4.1 · p. 6

4.2 Two-Stage Two-D Compression

  • Stage 1 applies SVD rotation and retains each KV vector’s top-r elements; stage 2 selects top-k token blocks using query-to-block-key scores before each decode attention layer.Channel compression occurs during prefill and decode; sequence compression occurs during decode.Source: §4.2 · pp. 6–7

4.3 Dynamic Compression Strategy

  • Input caches use configurable fixed-size segments, while output caches start a new segment when accumulated compression loss exceeds a threshold or the maximum length is reached.The stated maximum segment length is 64K tokens.Source: §4.3 · p. 7
  • Each segment’s strategy specifies channel rotation, retention rate and element-group size, plus sequence retention rate and token-block size.The strategy is generated for each segment rather than applied uniformly across the cache.Source: §4.3 · p. 7

5 Compressed KV Cache Management

MosaicKV stores sparse KV vectors in a packed representation and computes attention directly over them. It also overlaps or offloads cache-management work to reduce the burden of compression during decode.

5.1 Packed KV Format

  • The packed KV format stores retained elements contiguously with a D-bit position bitmap, avoiding per-element indices in conventional COO/CSR formats.Variable-granularity bitmaps use one bit per group of adjacent elements to reduce bitmap overhead.Source: §5.1 · p. 8

5.2 Compressed KV Encoding

  • Prefill encoding rotates K and V, channel-compresses them and per-block keys, then stores the results in packed format; bottom 25% channels are truncated.Encoding overlaps layer i−1’s KV-cache work with layer i’s computation.Source: §5.2 · p. 8

5.3 PackedAttention

  • PackedAttention selects top-k blocks and computes attention on packed vectors, reconstructing positions on the fly and merging inverse rotation into the query.It operates on FlashAttention chunks, skips unselected chunks, and uses an update buffer to batch new-token compression.Source: §5.3 · pp. 8–9

5.4 Heterogeneous Double Compression Buffering

  • A GPU buffer temporarily compresses the current segment with the previous strategy while the CPU evaluates compression loss and asynchronously finalizes and recompresses the segment.After compression, the final CPU-side result replaces the temporary GPU-side segment and the new strategy is transferred to the GPU.Source: §5.4 · p. 9

5.5 Incremental Strategy Generation

  • The CPU incrementally updates segment loss and the covariance matrix, then selects channel parameters under a loss threshold and block sizes under a variance threshold.The sequence compression rate is statically determined in the current implementation.Source: §5.5 · p. 9

6 Evaluation

Evaluation compares MosaicKV with uncompressed and Quest baselines across accuracy, memory, attention latency, and end-to-end serving. The reported results include lower decode latency and higher throughput.

6.1 Setup

  • The evaluation uses an H800 with 80GB HBM, compares against Full and Quest, and defaults to 10% token selection with MosaicKV retaining 25% of channels.LLaMA-3.1-8B is used for performance, while accuracy is also evaluated with Qwen3-4B and Ministral-3-8B.Source: §6.1 · p. 10

6.2 Accuracy

  • LongBench accuracy loss is 0.86% versus the uncompressed baseline and 0.64% versus Quest; on 64K RULER, losses versus uncompressed are 1.86%, 2.58%, and 3.55% across the three models.RULER accuracy losses versus Quest are 0.55%, 1.74%, and 2.69%, respectively.Source: §6.2 · p. 10, Table 1 · p. 10
  • On LongBench with LLaMA-3.1-8B, SVD rotation and per-vector selection improve normalized accuracy by 20.8% and 64.5%, while dynamic compression brings total accuracy to 98%.The incremental comparison starts from 10% token retention and uniform top-r channel selection without SVD rotation.Source: §6.2 · p. 10, Figure 9 · p. 11

6.3 Microbenchmark

  • With 25% channel retention, MosaicKV reduces KV cache size by 3× versus Full and 3.18× versus Quest.Quest’s per-block keys add 6.25% memory overhead at block size 16.Source: §6.3 · p. 10
  • MosaicKV achieves up to 16× speedup in total attention latency over Full and up to 2.4× over Quest.For score computing, it averages 4.2× speedup over Quest; Figure 8 breaks latency down by score computing, token selection, and attention.Source: Figure 8 · p. 11, §6.3 · p. 11

6.4 End-to-End Serving Performance

  • At 1M context, MosaicKV achieves 4.8× lower TBT and 7.3× higher throughput than vLLM.Across context lengths, average speedups are 3.8× and 5.5× over vLLM for TBT and throughput, respectively.Source: §6.4 · p. 11, Figure 10 · p. 12
  • Prefill overhead is 7.6% at 64K tokens and 0.4% at 1M tokens; decoding the full update buffer costs 40.6 μs.Table 2 covers compression overhead to the prefill stage.Source: Table 2 · p. 11, §6.4 · p. 12

7 Discussion

MosaicKV is presented as compatible with varied attention architectures and serving optimizations, while its current focus and kernel design leave stated future-work boundaries.

Compatibility

  • MosaicKV applies to attention-based LLMs, including GQA and MLA models, and can benefit MoE models with attention layers.The evaluated LLaMA-3.1-8B model uses GQA.Source: §7 · p. 12
  • PagedAttention, RadixAttention, and FlashAttention are compatible; retaining the full token sequence leaves page-level scheduling and prefix sharing unaffected.The implementation is built on FlashAttention.Source: §7 · p. 12

KV Cache Offloading

  • KV-cache offloading methods are complementary to MosaicKV, and fetching compressed vectors can further reduce PCIe traffic.The paper notes that offloading can introduce decode-path PCIe traffic that bottlenecks latency and throughput.Source: §7 · p. 12

Combining with Chunked Prefill

  • With chunked prefill, MosaicKV can optimize decode and compress input-token KV vectors by chunk to accelerate prefill.Input chunks’ KV vectors can be compressed using heterogeneous double compression buffering.Source: §7 · p. 12

Combining with Quantization

  • MosaicKV is orthogonal to quantization, which can reduce the compressed cache footprint; importance-aware quantization can also be applied to compressed elements.The paper names DiffKV as an example of importance-aware token quantization.Source: §7 · p. 12

Limitations and Future Work

  • MosaicKV currently focuses on decode; CUDA cores may limit PackedAttention if future HBM bandwidth makes memory no longer the bottleneck.Combining tensor cores and CUDA cores is proposed as a potential solution.Source: §7 · p. 12

8 Related Work

Related work compresses KV caches along token or channel dimensions, while MosaicKV’s cited distinction is dynamic two-dimensional compression throughout decode.

KV Cache Compression

  • Existing work selects important tokens [13] [16] [28] [29] [36] [48] [49] [54] [55] or channels [25] [31] [39] [43] [50]; SparK [31] compresses K but requires decompression before each attention step.Quest uses query-aware token selection, while MoBA [36] and NSA [54] retrain for higher accuracy under sequence compression.Source: §8 · pp. 12–13
  • Unlike those systems, MosaicKV performs per-token dynamic two-dimensional compression throughout decode, including token selection and attention computation.The paper characterizes this as achieving high memory compression, performance, and low accuracy loss.Source: §8 · p. 13

Quantization

  • KVQuant [23] isolates element-level outliers with per-vector dense-and-sparse quantization, while DiffKV [58] applies importance-aware KV-cache quantization.The paper states that MosaicKV can be combined with quantization for further memory reduction.Source: §8 · p. 13

Sparse GPU Kernels

  • MosaicKV uses underutilized CUDA cores and packed KV format for decode-stage sparse-cache attention, unlike cuSPARSE ’s extreme-sparsity focus and NVIDIA sparse tensor operation [37]’s 2:4 requirement.FlashInfer [53] accelerates block-sparse attention; PIT [62] densifies sparse tensors via permutation invariant transformation.Source: §8 · p. 13

9 Conclusion

MosaicKV combines channel and sequence compression to address long-context KV-cache memory constraints, using dynamic compression and cache management to retain accuracy and translate compression into serving gains.

Conclusion

  • On an H800 across multiple models and benchmarks, MosaicKV achieves up to 16× attention speedup, 4.8× lower decode latency, and 7.3× higher throughput, with 3× lower KV-cache memory and 1.76% average accuracy loss.Dynamic two-D compression uses per-vector element selection and segment-specific strategies; cache management uses underutilized CUDA cores and CPU resources for packed sparse attention and asynchronous cache management.Source: §9 · p. 13

Related papers

Loading 2607.00760v1…