PaperlayerView exact revision on arXiv

Source-linked AI summary

STAR-KV: Low-Rank KV Cache Compression via Soft Thresholding for Adaptive Rank Control

Priyansh Bhatnagar, Ashkan Moradifirouzabadi, Se-Hyun Yang, SeungJae Lee, Jungwook Choi, Mingu Kang

arXiv:2606.08382v1cs.LGcs.AI

TL;DR

Growing KV caches create a memory bottleneck, while existing low-rank compression methods face rank-selection and accuracy challenges. STAR-KV learns fine-grained ranks and combines hybrid decomposition with low-rank-aware quantization. Across multiple LLMs and benchmarks, it reports up to 20× overall cache reduction while maintaining model quality.

  • Problem

    KV-cache memory grows with context length, while prior low-rank methods use fixed or heuristic ranks and struggle to compress aggressively without accuracy degradation.

    Source: §1 · p. 1, §4.1.1 · p. 3

  • Method

    STAR-KV learns fine-grained ranks through differentiable singular-value thresholds, pairs key/value projections with different decompositions, and applies low-rank-aware mixed-precision quantization.

    Source: §1 · pp. 1–2

  • Results

    Across multiple LLMs and benchmarks, STAR-KV maintains model quality with up to 75% KV-cache reduction from low-rank compression and up to 20× overall reduction with quantization.

    Source: §7 · p. 9

  • Takeaways & Limitations

    STAR-KV offers a fine-grained compression approach that balances cache reduction with model quality across multiple LLMs and benchmarks.

    Source: §7 · p. 9

Abstract

from arXiv · show

Low-rank projection has emerged as a promising approach for compressing the KV cache by exploiting hidden-dimension redundancy. However, prior methods rely on fixed or heuristic rank selection and struggle to achieve aggressive compression with minimal accuracy degradation. We propose STAR-KV, an adaptive low-rank KV cache compression framework with fine-grained rank control. STAR-KV encompasses 1) a differentiable thresholding mechanism that enables optimal rank selection at both attention-head and block levels, 2) a hybrid decomposition strategy that applies different low-rank factorizations according to the sensitivity of key and value projections, and 3) a low-rank-aware mixed precision quantization that leverages data statistics for near lossless low-bit quantization. Evaluated across multiple LLMs and benchmarks, STAR-KV achieves up to 75% KV cache compression and up to 20x overall KV cache reduction when combined with quantization. Enabled by custom Triton-based GPU kernels, STAR-KV delivers up to 6.9x speedup for the attention module and 3.1x end-to-end generation throughput. Our code is publicly available at: https://github.com/PriyanshBhatnagar/STAR-KV.

1. Introduction

STAR-KV targets the memory and latency costs of growing KV caches and the accuracy losses that limit aggressive low-rank compression. It learns fine-grained ranks, combines decomposition and quantization strategies, and reports improved compression–accuracy trade-offs and faster generation.

1. Introduction

  • KV cache storage and access costs grow with context length and can dominate GPU memory, reaching 81% of an 85GB FP16 footprint for LLaMA-3.1-8B at 128K tokens and batch size four.Low-rank compression also adds reconstruction computation at inference, increasing latency.Source: §1 · p. 1
  • STAR-KV learns ranks by applying differentiable thresholds to singular values, rather than relying on prior methods’ heuristic or sensitivity-estimation rank choices.The approach controls ranks at fine granularity across decoder blocks and attention heads.Source: §1 · pp. 1–2
  • Its hybrid decomposition uses head-wise factorization for keys to reduce reconstruction overhead and joint factorization for values to better preserve fidelity.The choice reflects differing sensitivity of key and value projections under compression.Source: §1 · p. 2
  • Low-rank-aware mixed-precision quantization uses ordered singular values to identify and handle outlier channels efficiently.This contrasts with prior approaches that use explicit outlier modeling, non-uniform quantization, or additional computations.Source: §1 · p. 2
  • Up to 20.0× overall KV cache compression with 3.2-bit average mixed-precision quantization accompanies improved accuracy–compression trade-offs across LLMs and tasks.Low-rank compression alone reaches up to 4.0× with negligible accuracy degradation versus the uncompressed model; STAR-KV reports 2.70% better average zero-shot accuracy and 2.0× higher compression than Palu.Source: §1 · p. 2
  • Custom Triton-based kernels deliver up to 6.9× attention-operator speedup and 3.1× higher end-to-end generation throughput in long-context or high-batch scenarios.The attention speedup is reported against PyTorch SDPA; end-to-end throughput is compared with the PyTorch native implementation.Source: §1 · p. 2

2. Related Work

Prior KV-cache optimization includes architectural changes and lightweight post-training methods. Low-rank approaches compress hidden-dimension redundancy by caching compact latent representations.

2. Related Work

  • Architectural approaches such as GQA and MLA reduce KV-cache memory through sharing or projecting heads, but require architectural changes and additional training.Post-training alternatives include quantization, which lowers key/value precision, and token eviction, which retains only a subset of tokens under a cache budget.Source: §2 · p. 2
  • Low-rank KV-cache methods project key/value tensors into a lower-dimensional latent space and cache compact latent representations rather than using quantization or token eviction.A common approach factorizes projection weights offline, then uses the resulting factors at inference without additional runtime overhead.Source: §2 · p. 2
  • Low-rank decomposition exploits redundancy in the hidden channel dimension of key/value tensors.The offline factorization produces compact latent representations for the cache.Source: §2 · p. 2

3. Background

The background introduces multi-head attention, low-rank factorization of key/value projections, and weight absorption as a way to reduce reconstruction work during decoding.

3.1. Multi-Head Attention (MHA)

  • Multi-head attention computes query, key, and value projections, forms attention weights for each head, then concatenates head outputs and projects them back to the hidden dimension.The output projection follows the concatenation and weighted-sum operations.Source: §3.1 · p. 3

3.2. Low-rank KV Cache

  • SVD-based low-rank KV caching retains the top-r singular values and vectors of key/value projection matrices, storing intermediate representations in a reduced r-dimensional subspace.These compact representations reduce cache memory compared with full-dimensional tensors.Source: §3.2 · p. 3

3.3. Weight Absorption

  • Weight absorption reorders value processing to combine reconstruction with output projection, reducing FLOPs as sequence length grows.RoPE prevents offline fusion for key processing, while value processing permits reordering; with sufficiently large sequence length, savings approach a factor of d_h, typically 128 in widely used LLMs.Source: §3.3 · p. 3

4. The STAR-KV Framework

STAR-KV learns fine-grained ranks to balance KV-cache compression and accuracy, then selects key/value decomposition schemes and quantization according to their different sensitivity and channel statistics.

4.1.1. RANK SELECTION CHALLENGE FOR KV CACHE

  • At 75% compression, static and heuristic rank allocation incur 26.33% and 17.14% accuracy degradation, motivating adaptive rank selection.The challenge is to balance approximation error and memory reduction across blocks and projections.Source: §4.1.1 · pp. 3–4

4.1.2. ADAPTIVE RANK SELECTION VIA SOFT THRESHOLDING

  • STAR-KV learns per-block, per-head key and value thresholds using a differentiable soft surrogate, then applies hard thresholding offline before inference.The threshold controls effective rank; larger sharpness more closely approximates a hard cutoff, while smaller sharpness smooths the transition.Source: §4.1.2 · p. 4

4.2. Adaptive Compression Objective

  • The adaptive compression loss increases compression pressure early and reduces it as thresholds grow, leaving room to recover and stabilize task performance.A separate distillation loss uses KL divergence to encourage the compressed student to match the teacher’s predictive distribution.Source: §4.2 · pp. 4–5

4.3. Decomposition Scheme

  • Joint decomposition (JD) factors a concatenation of heads, whereas head-wise decomposition (HD) factors each head independently.The framework compares these alternatives in accuracy and reconstruction overhead.Source: §4.3 · p. 5

4.3.1. RECONSTRUCTION OVERHEAD ANALYSIS

  • At equal compression rates, JD requires h× the reconstruction FLOPs of HD.The analysis attributes the difference to joint versus per-head reconstruction.Source: §4.3.1 · p. 5

4.3.2. RECONSTRUCTION ERROR ANALYSIS

  • Under the same rank constraint, value projections have higher or equal relative approximation error than keys, while JD has no greater Frobenius-norm error than HD at the same budget.The key/value comparison follows from their normalized singular-value spectra; the JD/HD comparison concerns approximation error under a matched compression budget.Source: §4.3.2 · p. 5

4.3.3. HYBRID KV DECOMPOSITION

  • HD for keys and JD for values gives the best accuracy-overhead balance; reordering substantially reduces its normalized reconstruction latency.JD on both projections yields the best accuracy but the highest reconstruction latency; the hybrid assigns JD to more-sensitive values and HD to keys.Source: §4.3.3 · pp. 5–6, Figure 7 · p. 6

4.4. Integration with Quantization

  • Block-wise Hadamard transforms separate outlier and inlier latent channels for mixed-precision quantization, with higher precision for outliers.The method reserves the top 20% of latent channels for higher precision and quantizes the remaining 80% at lower precision; the Hadamard factors can be fused into projections without overhead.Source: §4.4 · p. 6, Figure 8 · p. 6

5. Experiments

STAR-KV is evaluated across multiple LLM families, standard and long-context benchmarks, and compression settings. Its reported results cover accuracy, quantization, attention speed, and end-to-end generation throughput.

5.1. Experiment Settings

  • Evaluation spans four LLM families, WikiText-2 and C4 perplexity, six zero-shot tasks, LongBench and RULER, at 60% and 75% KV cache compression budgets.Thresholds are learned using 3,000 FineWeb-Edu calibration samples, with reported fine-tuning cost below six GPU hours.Source: §5.1 · pp. 6–7

5.2. Accuracy Analysis

  • At 60% compression, LongChat-7B-v1.5 reaches WikiText-2 perplexity 6.83, matching the uncompressed baseline; STAR-KV also outperforms Palu and ReCalKV in reported zero-shot comparisons.For LLaMA-2-7B and LLaMA-3-8B-Instruct, it maintains minimal degradation at 60% and remains ahead of both methods at 75%.Source: §5.2 · pp. 7–8, Table 1 · p. 7
  • Across LongBench and RULER, STAR-KV preserves stronger long-context accuracy than Palu at higher compression, including 89.03 versus 85.06 RULER average on LLaMA-3.1-8B-Instruct.At 60% compression on LongChat-7B-v1.5, its RULER average is 88.26, close to the dense baseline of 88.80.Source: Table 2 · p. 7, §5.2 · p. 8, Table 3 · p. 8
  • With 3.2-bit average quantization on 75%-compressed LongChat, STAR-KV achieves 20× overall KV cache compression while maintaining 59.18% zero-shot average accuracy.The first 20% of channels are preserved in 4-bit and the remainder quantized in 3-bit; Table 4 reports compression and commonsense accuracy.Source: §5.2 · p. 8, Table 4 · p. 8

5.3. System-level Performance Analysis

  • At 75% compression, STAR-KV reaches 4.0× attention speedup at 32K; adding 4-bit KV quantization enables 64K context and 6.9× speedup against an estimated baseline.The PyTorch baseline runs out of memory at 32K and 64K, so those baseline latencies are extrapolated; Figure 9 distinguishes estimated-baseline points.Source: §5.3 · pp. 8–9, Figure 9 · p. 9
  • STAR-KV achieves a 3.1× average end-to-end generation throughput gain and fits up to 128K context on a single RTX 4090.Table 5 reports maximum throughput by context length, increasing batch size until GPU memory limits the run.Source: §5.3 · p. 9, Table 5 · p. 9

6. Ablation Results

The ablations examine rank-profile stability, calibration and loss-weight sensitivity, and combination with token pruning. The reported profiles are consistent across selected calibration datasets and nearby model variants.

1) Stability of learned rank profiles

  • Learned rank profiles are highly consistent across calibration datasets and transfer across nearby model variants, including models from the same family with similar parameter counts.The study reports this as a stability finding.Source: §6 · p. 9

2) Sensitivity to the calibration dataset

  • Changing the calibration dataset produces at most 0.68% difference in average zero-shot accuracy.The study describes learned thresholds as generally robust to calibration-dataset choice.Source: §6 · p. 9

3) Sensitivity to the compression-loss weight

  • Changing the compression-loss weight yields at most 0.25% average zero-shot accuracy variation.The authors report that results are not caused by fragile hyperparameter settings.Source: §6 · p. 9

4) Combination with token eviction

  • On LongChat-v1.5-32K, combining 60% low-rank compression with 60% H2O-style token pruning achieves 84% KV cache compression and 0.04% higher average zero-shot accuracy than H2O alone.The comparison uses H2O with 60% pruning as the baseline.Source: §6 · p. 9

7. Conclusion

STAR-KV combines adaptive rank control, hybrid decomposition, and low-rank-aware mixed-precision quantization. The conclusion reports high KV cache reduction while maintaining model quality, alongside attention and generation speedups.

Conclusion

  • STAR-KV uses soft thresholding to learn effective ranks across attention heads and decoder blocks, avoiding static or heuristic rank selection.Its hybrid decomposition is guided by sensitivity analysis to balance reconstruction cost and accuracy.Source: §7 · p. 9
  • Low-rank compression alone reduces KV cache by up to 75%, and adding quantization raises overall reduction to up to 20× while maintaining model quality.The quantization is low-rank-aware and mixed-precision.Source: §7 · p. 9
  • Custom Triton kernels deliver 3.5× attention speedup at 16K and up to 6.9× at 64K with an estimated baseline, alongside 3.1× higher end-to-end generation throughput.Source: §7 · p. 9

A.1.1. STATIC VS. ADAPTIVE RANK SELECTION

Under the same compression target and training settings, static uniform rank allocation causes substantial accuracy degradation relative to adaptive rank selection. The learned key and value rank maps show the allocation at their respective granularities.

A.1.1. STATIC VS. ADAPTIVE RANK SELECTION

  • Under the same compression target and training settings, uniform static rank allocation substantially degrades accuracy compared with adaptive per-layer and per-head allocation.Table 6 compares static and adaptive rank selection under the same KV cache compression target.Source: §A.1.1 · p. 14, Table 6 · p. 14
  • At 75% compression, the learned key rank maps are shown for each block and head, while value ranks are shown for each block.Figures 10 and 11 display the key and value rank maps learned through soft-thresholding.Source: §A.1.1 · p. 14, Figure 10 · p. 14, Figure 11 · p. 14

A.1.2. STABILITY OF LEARNED RANK PROFILES

STAR-KV compares learned rank profiles across calibration datasets and nearby model variants using normalized distances. The profiles are highly consistent across datasets, especially for keys across model variants.

A.1.2. STABILITY OF LEARNED RANK PROFILES

  • The profile-distance metric reports normalized distances as percentages and compares entries defined separately for keys and values.For the same model across calibration datasets, raw ranks are directly comparable.Source: §A.1.2 · p. 15
  • For cross-model comparisons, STAR-KV compares normalized learned-rank profiles rather than raw ranks, which may be incomparable across different attention layouts.Key profiles are per layer and head; value profiles are per layer.Source: §A.1.2 · p. 15
  • Across FineWeb-Edu, C4, and RedPajama, key-rank differences are 1.2–1.3% of head dimension and value-rank differences stay below 0.5% of model dimension.Across nearby model variants, key profiles are especially stable, while value profiles show larger but still moderate variation.Source: §A.1.2 · p. 15, Table 7 · p. 15

A.1.3. SENSITIVITY TO THE CALIBRATION DATASET

The appendix reports robustness across calibration settings and benchmarks, plus results for combined cache-compression methods and optimized attention kernels. At longer contexts, STAR-KV’s speedups increase, and quantization adds further speedup.

A.1.4. EFFECT OF HADAMARD ROTATION AND MIXED-PRECISION QUANTIZATION

  • Calibration-source changes only mildly affect zero-shot average at 60% compression; block-wise Hadamard with mixed precision achieves best ablation accuracy, and γ = 0.1 gives the best tested average.Global Hadamard without mixed precision and mixed precision without rotation both degrade accuracy; the reported Hadamard ablation uses 75% low-rank compression.Source: §A.1.3 · p. 16, Table 8 · p. 16, §A.1.4 · p. 16

A.1.6. SENSITIVITY TO THE OUTLIER–INLIER SPLIT

  • The 20:80 outlier–inlier split provides a strong balance among compression, speed-up, and downstream accuracy, and is the default setting.Source: §A.1.6 · p. 17, Table 11 · p. 17

A.2. Proof for Lemma 4.2

  • At equal budget R = hr, every head-wise rank-r candidate belongs to the joint rank-R feasible set, so joint optimization can achieve no worse reconstruction error.The joint feasible set contains all matrices of rank at most R.Source: §A.2 · p. 17

A.3. Training Settings

  • Training uses 3,000 FineWeb-Edu samples, then 1,000 distillation steps; evaluations cover zero-shot reasoning, LongBench long-context tasks, and RULER robustness.RULER is evaluated at 4K tokens by default, with an additional 16K evaluation on LongChat-v1.5-32K.Source: §A.3 · p. 17, §A.4 · pp. 17–18

A.5.1. ADDITIONAL MODELS

  • STAR-KV exceeds ReCalKV on Mistral at 60% compression, beats Palu on LLaMA-2-13B, preserves RULER performance, and scores 40.79% on LLaMA-3.1 LongBench at 50% compression.The LongBench baseline is 42.23%; STAR-KV at 60% compression remains substantially better than Palu at 50%.Source: §A.5.1 · p. 18, Table 12 · p. 18, §A.5.2 · p. 18

A.5.4. COMBINATION WITH TOKEN EVICTION

  • Combining STAR-KV with 60% H2O pruning reaches 84% KV cache compression on LongChat and 90% on LLaMA-2-13B, with only small additional zero-shot degradation on the latter.The combined method is more accurate than H2O at 60% pruning alone on LongChat.Source: §A.5.4 · pp. 18–19, Table 15 · p. 19

A.5.5. LATENCY BREAKDOWN

  • At batch 16, STAR-KV achieves 4.1× speedup at 32K on RTX 4090 and 2.9× at 128K on RTX Pro 6000; quantization raises these to 4.9× and 4.3×.At 32K on RTX 4090, PyTorch SDPA runs out of memory while both STAR-KV variants remain executable.Source: Figure 12 · p. 19, §A.5.5 · p. 19, §A.5.6 · p. 19

A.6. Implementation Optimizations

  • Fused Triton kernels combine key dequantization, RoPE, reconstruction, and attention scoring, while value-path reordering combines per-head operations into a GEMM.These optimizations reduce intermediate materialization and redundant memory traffic, and increase arithmetic intensity and GPU utilization.Source: §A.6 · p. 20
Loading 2606.08382v1…