PaperlayerView exact revision on arXiv

Source-linked AI summary

KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization

Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W. Mahoney, Yakun Sophia Shao, Kurt Keutzer, Amir Gholami

arXiv:2401.18079v6cs.LG

TL;DR

Long-context inference makes KV-cache activations a dominant memory cost, while prior methods struggle to represent them accurately below 4 bits. KVQuant combines distribution-aware Key and Value quantization with outlier handling and custom kernels. Across LLaMA, Llama-2, Llama-3, and Mistral, it achieves under 0.1 perplexity degradation at 3 bits, supports million-token contexts, and provides kernel speedups.

  • Problem

    Long-context inference makes KV-cache activations the dominant memory contributor, while existing methods represent them inaccurately at sub-4-bit precision.

  • Method

    KVQuant combines pre-RoPE per-channel Key quantization, sensitivity-weighted non-uniform datatypes, per-vector dense-and-sparse outlier handling, and custom CUDA kernels.

  • Results

    Across LLaMA, Llama-2, Llama-3, and Mistral, 3-bit KV-cache quantization achieves under 0.1 perplexity degradation on Wikitext-2 and C4.

  • Takeaways & Limitations

    KVQuant enables LLaMA-7B serving with 1M-token context on one A100 GPU and 10M-token context on an 8-GPU system.

  • Takeaways & Limitations

    The work targets efficient inference rather than training long-context models, and its latency benchmarks focus on generation rather than prompt processing.

Abstract

from arXiv · show

LLMs are seeing growing use for applications which require large context windows, and with these large context windows KV cache activations surface as the dominant contributor to memory consumption during inference. Quantization is a promising approach for compressing KV cache activations; however, existing solutions fail to represent activations accurately in sub-4-bit precision. Our work, KVQuant, facilitates low precision KV cache quantization by incorporating several novel methods: (i) Per-Channel Key Quantization, where we adjust the dimension along which we quantize the Key activations to better match the distribution; (ii) Pre-RoPE Key Quantization, where we quantize Key activations before the rotary positional embedding to mitigate its impact on quantization; (iii) Non-Uniform KV Cache Quantization, where we derive per-layer sensitivity-weighted non-uniform datatypes that better represent the distributions; and (iv) Per-Vector Dense-and-Sparse Quantization, where we isolate outliers separately for each vector to minimize skews in quantization ranges. By applying our method to the LLaMA, Llama-2, Llama-3, and Mistral models, we achieve < 0.1 perplexity degradation with 3-bit quantization on both Wikitext-2 and C4, outperforming existing approaches. Our method enables serving LLaMA-7B with a context length of up to 1 million on a single A100-80GB GPU and up to 10 million on an 8-GPU system. We develop custom CUDA kernels for KVQuant, showing that we can achieve up to ~1.7x speedups, compared to baseline fp16 matrix-vector multiplications, for the LLaMA-7B model.

1 Introduction

Long-context inference shifts the memory bottleneck from model weights to cached Key and Value activations, motivating KV-cache compression. KVQuant combines distribution-aware quantization and outlier handling to support accurate, efficient low-bit inference.

  • Longer sequence lengths make cached Key and Value activations the main memory bottleneck, especially during batched inference.
  • Existing KV-cache quantization methods suffer unacceptable accuracy degradation because of activation outliers and suboptimal bit allocation.
  • Keys contain outlier channels before RoPE, so KVQuant quantizes Keys per-channel before applying RoPE.
  • KVQuant derives sensitivity-weighted non-uniform datatypes offline on calibration data instead of relying only on activation magnitude.
  • Removing only 1% of outliers yields under 0.1 perplexity degradation for 3-bit KV-cache quantization across four model families and enables 4.8× longer context.
  • Custom CUDA kernels achieve up to ∼1.7× speedups for LLaMA-7B Key and Value matrix-vector multiplications at 4-bit precision versus fp16.

2 Background

LLM inference becomes memory-bandwidth bound during autoregressive generation, while long contexts make KV-cache storage increasingly important. KV-cache quantization complements token eviction and retrieval methods but remains difficult at low precision because activation distributions contain structured outliers.

  • During generation, each token depends on previous tokens, making small-batch decoder inference typically memory-bandwidth bound.
  • Weight-only quantization primarily targets weights because they dominate memory consumption and runtime at small sequence lengths and batch sizes.
  • Prior KV-cache activation quantization still shows significant perplexity degradation at low precision, including unacceptable loss for some 3-bit methods.
  • KV-cache compression research also includes evicting less important tokens or retrieving only subsets of tokens, whereas KVQuant explores quantization as an orthogonal direction.
  • KV Cache Compression: Key distributions show distinct outlier channels, while Value distributions contain outlier channels and tokens that are less extreme than Key outliers.

3 Method

KVQuant matches quantization granularity and datatype design to KV-cache activation structure. Its method uses pre-RoPE per-channel Key quantization, offline sensitivity-weighted non-uniform datatypes, per-vector outlier handling, attention-sink awareness, and specialized kernels.

  • Per-Channel Quantization: Per-channel quantization groups values with similar channel magnitudes, better matching KV-cache distributions than per-token grouping across channels.
  • Per-Channel Quantization: Offline calibration enables accurate per-channel quantization without the fine-grained grouping or fp16 residual subset required by a concurrent approach.
  • Pre-RoPE Key Quantization: Pre-RoPE Key quantization stores unrotated Keys and applies the position-dependent RoPE rotation during inference.
  • Pre-RoPE Key Quantization: Post-RoPE rotation mixes channel pairs by position, making channel magnitudes less consistent and harder to quantize.
  • Non-Uniform Quantization: KVQuant derives per-layer sensitivity-weighted non-uniform signposts offline, then rescales the datatype per-channel or per-token for online quantization.
  • Dense-and-Sparse Quantization: Per-vector dense-and-sparse quantization uses separate outlier thresholds at the quantization granularity, calibrating per-channel thresholds offline and computing per-token thresholds online.
  • Implementation: Attention Sink-Aware quantization is used across bit widths, while custom kernels store quantized values as lookup-table indices and sparse outliers in CSR or CSC format.

4 Results

KVQuant maintains low perplexity degradation under aggressive KV-cache compression across standard and long-context evaluations, while also reducing latency and memory requirements.

  • Main Evaluation: 3-bit quantization stays under 0.1 perplexity degradation on Wikitext-2 across LLaMA models, with 4.8× memory savings relative to fp16.The same evaluation reports under 0.02 degradation at 4 bits and under 0.5 at 2 bits, with 3.7× and 6.9× memory savings, respectively.
  • Main Evaluation: KVQuant maintains accuracy with larger input contexts on LLaMA-2-7B-32K and Llama-2-70B-32K LongLoRA models evaluated on Wikitext-2.Perplexity was measured at different sequence lengths to assess long-context performance.
  • Long Context Length Evaluation: KVQuant maintains long-context passkey retrieval performance and achieves higher retrieval rates than KIVI at the same compression level.Passkey success rates were evaluated across context lengths for the LLaMA-2-7B-32K model, with additional long-context-model evaluation.
  • Long Context Length Evaluation: On LongBench, 3-bit KVQuant shows minimal degradation relative to fp16 and outperforms KIVI at a similar compression level.LongBench covers question answering, summarization, and few-shot learning tasks with maximum input context length set to 31,500 tokens.
  • Long Context Length Evaluation: On RULER, 3-bit KVQuant scores 14% better than KIVI at similar average bit-width, while 2-bit KVQuant matches KIVI accuracy at 1.5× smaller bit-width.The evaluation uses a maximum context length of 32K.
  • Performance Analysis and Memory Savings: Custom kernels reduce Key and Value multiplication latency by 1.2–1.6× and 1.3–1.7×, respectively, relative to fp16 on an A6000.The benchmark uses batch size 1 and evaluates 4-bit dense-and-sparse kernels across sequence lengths.
  • Performance Analysis and Memory Savings: With nuq2, KVQuant provides 8× KV-cache compression, supporting LLaMA-7B at 1M tokens on one A100 and 10M tokens on eight GPUs with little fp16 degradation.The result is reported for long-context serving and is compared with baseline fp16 inference.

5 Conclusion

KVQuant targets the KV cache memory bottleneck in long-context inference with low-precision quantization designed around activation outliers, RoPE effects, and sensitivity. The method preserves accuracy while enabling substantially longer contexts and efficient inference.

  • 5 Conclusion: KVQuant combines pre-RoPE per-channel Key quantization, sensitivity-weighted non-uniform quantization, and per-vector dense-and-sparse outlier handling.These choices address inconsistent post-RoPE channel magnitudes, suboptimal signpost placement, and outliers that reduce dense quantization resolution.
  • 5 Conclusion: 2-bit KV cache compression enables 1M-token LLaMA-7B inference on one A100-80GB GPU and 10M-token inference on an 8-GPU system.The reported settings use nuq2 compression for the KV cache.
  • 5 Conclusion: 3.7× KV cache compression is achieved with nuq4-1%, while nuq2 enables LLaMA-65B at 32K tokens on one A100-80GB GPU.The LLaMA-65B configuration uses 4-bit-compressed weights and a nuq2-1% KV cache.
  • 5 Conclusion: Sensitivity-weighted quantization derives non-uniform datatypes offline from calibration data rather than placing signposts using activation magnitude alone.The sensitivity measure is based on perturbation of the model loss under quantization.
  • 5 Conclusion: Dense-and-sparse quantization removes numerical outliers from the dense component and stores them separately, with approximately 99% of values occupying a small dynamic-range portion.Per-vector outlier detection is reported to outperform per-matrix detection without additional memory overhead.

G Per-Channel Key Quantization Ablations

Per-channel quantization benefits Keys but not Values: at 3-bit precision, per-channel Keys with per-token Values improves LLaMA-7B perplexity over per-token quantization for both activations.

  • Per-Channel Key Quantization: 3.82 perplexity improvement results from per-channel Key and per-token Value quantization versus standard per-token quantization at 3-bit precision.The comparison is reported for LLaMA-7B.
  • Per-Channel Key Quantization: Per-channel Key quantization mitigates large outlier channels in Keys.The configuration uses an fp16 zeropoint because some channels are entirely positive or negative.
  • Per-Channel Value Quantization: Per-channel Value quantization performs worse than per-token Value quantization.The paper attributes this to localized error accumulation and perturbation of important tokens in attention-weighted sums.
  • Per-Channel Key Quantization: Table 9 compares per-token and per-channel quantization for KV cache activations in LLaMA-7B.PT denotes per-token quantization and PC denotes per-channel quantization.

H Pre-RoPE Key Quantization Ablations

Quantizing Keys before RoPE improves low-precision accuracy relative to post-RoPE quantization by avoiding channel mixing from rotary positional embeddings.

  • Pre-RoPE Key Quantization: 0.82 perplexity improvement is obtained by pre-RoPE rather than post-RoPE Key quantization at 3-bit precision for LLaMA-7B.The configuration uses per-channel Key and per-token Value quantization.
  • Pre-RoPE Key Quantization: RoPE makes Key quantization more challenging by mixing pairs of channels with different magnitudes.Pre-RoPE quantization therefore permits more accurate low-precision quantization.
  • Pre-RoPE Key Quantization: Table 10 compares perplexity for pre-RoPE and post-RoPE Key quantization in LLaMA-7B.All configurations use per-channel Key quantization and per-token Value quantization.

I Sensitivity-Weighted Non-Uniform Quantization Ablations

Sensitivity-weighted non-uniform datatypes improve low-bit KV cache quantization, with additional gains from accounting for per-channel scaling factors during k-means.

  • Sensitivity-Weighted Non-Uniform Quantization: Sensitivity-weighted non-uniform quantization yields particularly noticeable gains over uniform quantization at 3-bit and 2-bit precision.The evaluation covers different LLaMA, Llama-2, and Mistral models on Wikitext-2.
  • Sensitivity-Weighted Non-Uniform Quantization: Sensitivity weighting is necessary to derive performant non-uniform datatypes with a k-means-based approach.The ablation also evaluates nuq3 without sensitivity weighting and variants without per-channel scaling factors.
  • Sensitivity-Weighted Non-Uniform Quantization: Accounting for per-channel scaling factors during k-means provides distinct additional benefits.The comparison includes uniform int3, non-uniform nf3, unweighted k-means, and Fisher-weighted k-means variants.

J Per-Vector Dense-and-Sparse Quantization Ablations

Isolating outliers in separate sparse representations improves KV cache quantization, with per-vector thresholds outperforming a single per-matrix threshold.

  • Per-Vector Dense-and-Sparse Quantization: Per-vector outlier detection provides greater accuracy improvements than per-matrix detection.The method uses separate thresholds per-channel for Keys and per-token for Values.
  • Per-Vector Dense-and-Sparse Quantization: Different thresholds for Keys and Values provide additional benefits when outliers are isolated and stored sparsely.The reported configuration uses per-channel thresholds for Keys and per-token thresholds for Values.
  • Per-Vector Dense-and-Sparse Quantization: Per-vector outlier thresholds significantly improve accuracy relative to per-tensor outlier detection in LLaMA-7B on Wikitext-2.PV denotes per-vector thresholds and PM denotes a single per-matrix threshold.

K Attention Sink-Aware Quantization Ablations

Attention Sink-Aware quantization consistently improves perplexity across evaluated LLaMA-family models and bit widths, with larger gains at lower precision and without sparsity.

  • Attention Sink-Aware quantization consistently improves perplexity across 4-, 3-, and 2-bit settings for dense-only and dense-and-sparse quantization.
  • 0.25 PPL and 0.13 PPL improvements are observed for nuq2-1% on Llama-3-8B and Llama-3-70B, respectively.
  • Perplexity gains are more pronounced at lower bit widths and without sparsity.

L Calibration Ablations

The calibration ablations show that offline Key calibration can match online per-channel calibration when outlier extraction is used, while calibration and online thresholding remain practically manageable.

  • Offline per-channel scaling factors avoid recomputing factors and updating prior cached Keys for every newly appended Key.The factors are computed from a calibration set rather than online during inference.
  • Offline Key calibration matches online per-channel calibration after removing a small percentage of outliers.Without outlier extraction, 3-bit quantization shows minor accuracy degradation.
  • 45% of each LLaMA-7B projection matrix-vector runtime is required by the topk operation for online outlier-threshold computation.The topk operation can also run on the CPU in parallel with the subsequent linear-layer matrix-vector operation.
  • At 4-bit precision, sequential per-layer Fisher-information computation and calibration for LLaMA-65B take at most 6 hours.Each layer’s calibration takes only a few minutes and can be performed independently.
  • The evaluation derives scaling factors, zero-points, and non-uniform datatypes from 16 Wikitext-2 training samples of length 2K.The calibration process is reported to remain robust to the choice of calibration data.
  • NUQ generally outperforms uniform INT and non-uniform NF datatypes across LLaMA, Llama-2, Llama-3, and Mistral models on Wikitext-2.Extracting 0.1% to 1.0% of outliers further improves performance; 1.0% NUQ has memory requirements similar to INT with group size 64.

O Full Perplexity Evaluation

Full evaluations report the benefits of KVQuant across LLaMA, Llama-2, Llama-3, and Mistral models on Wikitext-2 and C4, with calibration remaining robust across datasets.

  • KVQuant improves KV cache compression results across LLaMA, Llama-2, Llama-3, and Mistral models and multiple bit widths on Wikitext-2 and C4.
  • NUQ with pre-RoPE per-channel Key quantization is the configuration used for KVQuant results in the full evaluations.
  • Post-RoPE Key quantization outperforms pre-RoPE Key quantization when Keys use per-token quantization in LLaMA-7B.The paper therefore uses post-RoPE results for per-token baseline comparisons.
  • Perplexity remains similar when calibration uses Wikitext-2 or C4 and evaluation uses either dataset.

R Kernel Implementation Details

KVQuant uses lookup-table-based dense-and-sparse CUDA kernels that dequantize compressed KV activations on demand, with sparse layouts selected for efficient vector appends and balanced GPU work.

  • The 4-bit kernels load compressed Keys and Values and dequantize them only as needed to reduce memory-bandwidth utilization.The kernels use lookup tables plus per-element indices, with arithmetic performed in fp16.
  • CSC stores Key outliers and CSR stores Value outliers because each format permits efficient appending of newly generated vectors.
  • Balanced sparse matrix-vector kernels assign equal numbers of nonzeros per thread to address unbalanced sparsity distributions on GPUs.This design incurs greater synchronization overhead than assigning one thread per row.
  • A single per-layer datatype rescaled per vector avoids loading a different lookup table for each element along the reduction dimension.
  • Table 22 breaks runtime into vector packing and dense-and-sparse matrix-vector multiplication for LLaMA-2-7B-32K at different sequence lengths.
  • The implementation remains inefficient when sparse-matrix updates copy prior-token data during concatenation.Blocked allocation is proposed to avoid overhead from reallocating memory.
Loading 2401.18079v6…