PaperlayerView exact revision on arXiv

Source-linked AI summary

KVmix: Gradient-Based Layer Importance-Aware Mixed-Precision Quantization for KV Cache

Fei Li, Song Liu, Weiguo Wu, Shiqiang Nie, Jinyu Wang

arXiv:2506.08018v3cs.LGcs.AI

TL;DR

KV Cache memory growth constrains LLM inference, while existing quantizers can lack layer-specific precision allocation or long-context prioritization. KVmix uses gradient-based importance to allocate mixed precision and dynamically retains full-precision pivotal KV pairs; it reports near-lossless accuracy alongside lower memory use and faster inference.

  • Problem

    KV Cache memory grows with sequence length, and existing methods may use fixed precision allocation or fail to prioritize critical KVs in long-context tasks.

    Source: Introduction · p. 1

  • Method

    KVmix uses loss gradients for layer-specific Key and Value importance, assigning mixed precision and dynamically retaining full-precision pairs for recent pivotal tokens.

    Source: KV Importance Analysis · p. 4, Dynamic Pivotal Context Selection · p. 5

  • Results

    KVmix reports nearly lossless model accuracy, a 4.9× reduction in memory usage, and a 5.3× speedup in inference efficiency.

    Source: Introduction · p. 2

  • Takeaways & Limitations

    KVmix offers flexibility to adapt quantization strategies to diverse scenarios while pursuing memory and efficiency gains with minimal accuracy loss.

    Source: Conclusion · p. 9

    The quantization perturbation analysis assumes that quantization introduces small perturbations to the Key and Value matrices.

    Source: KV Importance Analysis · p. 3

Abstract

from arXiv · show

The high memory demands of the Key-Value (KV) Cache during the inference of Large Language Models (LLMs) severely restrict their deployment in resource-constrained platforms. Quantization can effectively alleviate the memory pressure caused by KV Cache. However, existing methods either rely on static one-size-fits-all precision allocation or fail to dynamically prioritize critical KV in long-context tasks, forcing memory-accuracy-throughput tradeoffs. In this work, we propose a novel mixed-precision quantization method for KV Cache named KVmix. KVmix leverages gradient-based importance analysis to evaluate how individual Key and Value projection matrices affect the model loss, enabling layer-specific bit-width allocation for mix-precision quantization. It dynamically prioritizes higher precision for important layers while aggressively quantizing less influential ones, achieving a tunable balance between accuracy and efficiency. KVmix also introduces a dynamic long-context optimization strategy that adaptively keeps full-precision KV pairs for recent pivotal tokens and compresses older ones, achieving high-quality sequence generation with low memory usage. Additionally, KVmix provides efficient low-bit quantization and CUDA kernels to optimize computational overhead. On LLMs such as Llama and Mistral, KVmix achieves near-lossless inference performance with extremely low quantization configuration (Key 2.19bit Value 2.38bit), while delivering a remarkable 4.9x memory compression and a 5.3x speedup in inference throughput.

Introduction

KV Cache memory grows with sequence length, limiting LLM deployment and inference efficiency, while existing quantizers struggle to balance memory, accuracy, and throughput. KVmix combines layer-aware mixed precision with dynamic long-context retention and efficient low-bit CUDA implementations.

Introduction

  • A 70B-parameter model may require over 50GB for the KV Cache of a 20k-token sequence, exceeding typical GPU capacity.Concurrent requests require separate caches; memory depletion can trigger transfers to system memory or disks and sharply increase latency.Source: Introduction · p. 1
  • Existing methods use fixed precision allocation or costly dynamic quantization that does not adaptively prioritize critical KVs in long-context tasks.These limitations leave memory, accuracy, and throughput trade-offs unresolved.Source: Introduction · p. 1
  • KVmix reports nearly lossless accuracy, a 4.9× reduction in memory usage, and a 5.3× speedup in inference efficiency.The method also includes efficient CUDA implementations and a high-compression 3-bit quantization method.Source: Introduction · p. 2
  • KVmix uses layer-wise KV importance to assign higher bit-widths to critical layers and lower bit-widths to less influential layers, while retaining full precision for recent pivotal tokens.It also provides low-bit quantization and CUDA implementations.Source: Introduction · pp. 1–2

Related Work and Motivation

Prior work compresses KV Cache through quantization, sparsification, and sharing, while KVmix presents its approach as lightweight and complementary to some existing methods. The motivation is that KV importance varies across layers, and learned projection weights offer a basis for assessing those differences.

Related Work

  • KVmix is described as orthogonal to weight quantization and sparsification, and its layer importance can guide more accurate KV eviction during sparsification..Source: Related Work · p. 2
  • Existing KV quantizers include per-channel/per-token, nonuniform, online dynamic, and search-based approaches; some incur computational or search overhead.The cited examples include KIVI, KVQuant, QAQ, and KVTuner.Source: Related Work · p. 2

Motivation

  • Two-bit quantization of Keys or Values from different layers has varying effects on generation quality, motivating layer-specific precision allocation.Figure 1 compares Llama 2-7B accuracy on GSM8K and TruthfulQA with FP16 and individual 2-bit layer ablations.Source: Motivation · p. 2, Figure 1 · p. 3
  • Key and Value projection weights determine how those states are extracted, vary across layers, and remain fixed during inference, making them usable for importance analysis.Figure 2 depicts the projection weights across layers; its Norm and Range labels refer to each matrix's L2 norm and value range.Source: Motivation · pp. 2–3, Figure 2 · p. 3

Methodology

KVmix profiles layer-wise Key and Value sensitivity to loss, then uses the resulting importance scores to guide mixed-precision quantization. It also combines asymmetric grouping, dynamic retention of recent pivotal contexts, and fused CUDA kernels.

KV Importance Analysis

  • Layer-wise L2 norms of loss gradients with respect to Key and Value projection weights estimate component importance and guide their mixed-precision bit-widths.Scores can be averaged across prompts; higher-scored layers receive higher precision, and the 20%-80% split can be adjusted to trade accuracy for memory.Source: KV Importance Analysis · p. 4

Asymmetric Quantization Strategy

  • Keys are quantized per channel and Values per token, isolating Key-channel outliers and confining Value quantization errors to individual tokens.The method uses per-channel groups for Keys and per-token groups for Values to reduce quantization errors.Source: Asymmetric Quantization Strategy · p. 4

Group-Wise Low-Bit Quantization

  • Group-wise low-bit quantization packs 3-bit values in blocks of 11, storing ten values at 3 bits and one at 2 bits for 10% greater packing density than uniform 3-bit packing.Quantized values are stored using bit operations in int32; dequantization uses the group scale and minimum value.Source: Group-Wise Low-Bit Quantization · p. 4

Dynamic Pivotal Context Selection

  • KVmix selects a per-layer, per-component number of recent pivotal KV pairs from importance-based ratios, keeps them full precision, and mixed-quantizes older pairs.The ratio can differ for Keys and Values; increasing it improves accuracy but requires more memory.Source: Dynamic Pivotal Context Selection · pp. 4–5

CUDA Implementation

  • CUDA kernels fuse quantization with cache concatenation and dequantization with matrix-vector multiplication, while supporting 1-, 2-, 3-, and 4-bit configurations.The fused operations reduce memory access or memory overhead, respectively.Source: CUDA Implementation · p. 5

Experimental Results

Experiments evaluate KVmix across LongBench, language modeling, mathematical reasoning, and inference efficiency; the profiler selects model-specific layer bit-widths, with k2.19v2.38 used as the default configuration. KVmix reports lower LongBench accuracy loss than the tested alternatives, strong results on GSM8K and Wikitext-2, and reduced memory use with increased throughput.

Experimental Setup

  • Across five evaluated models, LongBench, Wikitext-2, and GSM8K measure long-context performance, perplexity, and mathematical reasoning, respectively.The evaluation used an RTX 4090 with 24GB memory; the LongBench maximum sequence length was 4096.Source: Experimental Setup · p. 5

Profiling Results

  • With 20% of layers assigned higher bit-widths, the profiler selects a default configuration averaging 2.19 bits for Key and 2.38 bits for Value.The profiler analyzed 30 LongBench prompts; 20–30 prompts were reported sufficient for reliable importance analysis, taking 10–15 minutes in the tested setup.Source: Profiling Results · p. 6, Figure 5 · p. 6, Figure 6 · p. 6

Long Context

  • 1.67% average LongBench accuracy loss for KVmix-k2.19v2.38, versus 4.53% for KVmix-2bit and 4.06% for random-k2.19v2.38 across four models.Removing RPC reduces average accuracy by 3.28% compared with KVmix-k2.19v2.38; KVmix-k2.28v2.56 nearly matches full 4-bit performance but uses more memory.Source: Long Context · p. 6, Table 1 · p. 7

GSM8K and Wikitext-2

  • KVmix-k2.19v2.38 has 2.00% accuracy loss on GSM8K and near-baseline perplexity on Wikitext-2, outperforming the reported 2-bit and 4-bit baselines on GSM8K.The passage also reports better GSM8K accuracy than KVmix-2bit and random-k2.19v2.38.Source: GSM8K and Wikitext-2 · p. 7, Table 3 · p. 8

Inference Efficiency and Memory Usage Evaluation

  • 4.9× lower memory use and up to 5.3× higher throughput than the baseline are reported for KVmix-k2.19v2.38.On Llama 2-7B-hf, it reached batch size 30 and 1032 tokens per second; the reported memory and throughput comparisons used different batch-size procedures.Source: Inference Efficiency and Memory Usage… · pp. 7–8

Conclusion

KVmix combines gradient-based layer importance analysis with dynamic long-context optimization to reduce KV Cache memory while maintaining generation quality. The paper reports memory and efficiency gains with minimal accuracy loss, and identifies real-time KV bit adjustment as future work.

Conclusion

  • KVmix integrates gradients of KV weights into quantization and dynamic long-context optimization to cut memory usage while maintaining generation quality.The conclusion describes flexibility to adapt quantization strategies to diverse scenarios.Source: Conclusion · pp. 8–9
  • The conclusion reports significant memory and efficiency gains with minimal accuracy loss.Figures 7 and 8 are identified as dynamic peak-memory and inference-throughput comparisons, respectively.Source: Figure 7 · p. 8, Figure 8 · p. 8, Conclusion · p. 9
  • Future work will explore lightweight mechanisms for real-time KV bit adjustments.Source: Conclusion · p. 9

KV Weight Analysis

KV projection-weight distributions vary across layers, models, and between Key and Value, motivating model- and layer-specific quantization.

KV Weight Analysis

  • Across three models, Key and Value projection weights differ across layers and models, and Key and Value distributions also differ within a layer.The analysis uses normalized heatmaps sampling every sixth matrix element, with L2 norm and value range reported per layer.Source: KV Weight Analysis · p. 10

Appendix B

The KVmix profiler runs offline, scoring Key and Value importance from loss gradients before configuring quantization bits. Its analysis uses sampled prompts, and the resulting configuration can be reused for inference.

Appendix B

  • The offline profiler processes data, analyzes KV importance, and configures quantization parameters for later reuse without recalculating during inference.It assesses layer importance using gradients of model loss with respect to Wk and Wv.Source: KVmix Profiler Algorithm · p. 10
  • Limiting prompt length mitigates backpropagation memory demands, while 20–30 prompts can provide sufficient importance-analysis accuracy.Source: KVmix Profiler Algorithm · p. 10
  • For each prompt and layer, the algorithm computes Key and Value loss gradients, takes their norms, then averages scores across prompts.The algorithm classifies the scores into layers and sets quantization bits before running inference and saving predictions.Source: KVmix Profiler Algorithm · pp. 10, 12

Analysis Results of KVmix Profiler Using Different Prompts

The profiler’s layer-importance distribution remains consistent across prompt sources and sample sizes. Using 20 prompts produces nearly identical results to 30, with negligible impact on final inference accuracy.

Analysis Results of KVmix Profiler Using Different Prompts

Impact of RPC Proportion on Model Accuracy

The RPC analysis finds that a 20%/10% ratio limits accuracy loss, while higher ratios bring marginal accuracy gains and can reduce memory compression, especially in short-context decoding.

Impact of RPC Proportion on Model Accuracy

Additional Performance Evaluation of KVmix on LongBench

On LongBench, KVmix-k2.19v2.38 offers a substantially higher compression ratio than KVmix-4bit while keeping average accuracy loss within 1.30%. Increasing the high-bit layer share improves accuracy toward the 4-bit configuration.

Additional Performance Evaluation of KVmix on LongBench

Related papers

Loading 2506.08018v3…