PaperlayerView exact revision on arXiv

Source-linked AI summary

KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache

Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, Xia Hu

arXiv:2402.02750v2cs.CLcs.LGcs.PF

TL;DR

KV-cache growth limits batched LLM inference, while limited evidence leaves the quantization behavior of keys and values insufficiently understood. KIVI addresses this with tuning-free asymmetric 2bit quantization, reducing peak memory by 2.6× and increasing throughput by 2.35×∼3.47×.

  • Problem

    KV-cache growth creates memory and speed bottlenecks in batched, long-context LLM inference, while its element distribution and quantization limits remain insufficiently studied.

  • Method

    KIVI analyzes KV-cache distributions and applies tuning-free 2bit asymmetric quantization, using per-channel keys and per-token values.

  • Results

    2.6× peak memory reduction with little to no accuracy drop enables up to 4× larger batch sizes and 2.35×∼3.47× throughput.

  • Takeaways & Limitations

    KIVI provides a hardware-friendly approach for compressing KV caches in real LLM generation workloads without fine-tuning.

  • Takeaways & Limitations

    The evaluation scope excludes closed-end tasks such as MMLU because they involve only one decoding step and do not suit compressed-KV-cache analysis.

Abstract

from arXiv · show

Efficiently serving large language models (LLMs) requires batching of many requests to reduce the cost per request. Yet, with larger batch sizes and longer context lengths, the key-value (KV) cache, which stores attention keys and values to avoid re-computations, significantly increases memory demands and becomes the new bottleneck in speed and memory usage. Additionally, the loading of the KV cache causes the computational core to be idle, which limits the inference speed. A straightforward and effective solution to reduce KV cache size is quantization, which decreases the total bytes taken by KV cache. However, there is a lack of in-depth studies that explore the element distribution of KV cache to understand the hardness and limitation of KV cache quantization. To fill the gap, we conducted a comprehensive study on the element distribution in KV cache of popular LLMs. Our findings indicate that the key cache should be quantized per-channel, i.e., group elements along the channel dimension and quantize them together. In contrast, the value cache should be quantized per-token. From this analysis, we developed a tuning-free 2bit KV cache quantization algorithm named KIVI. With hardware-friendly implementation, KIVI can enable Llama, Falcon, and Mistral models to maintain almost the same quality while using $\mathbf{2.6\times}$ less peak memory (including model weight). This reduction in memory usage enables up to $\mathbf{4\times}$ larger batch size, bringing $\mathbf{2.35\times \sim 3.47\times}$ throughput on real LLM inference workload. The source code is available at https://github.com/jy-yuan/KIVI.

1. Introduction

The introduction identifies KV-cache memory and loading as bottlenecks in batched LLM inference and motivates quantization as a simple way to reduce cache size. It presents KIVI, a tuning-free 2bit method that uses per-channel key quantization and per-token value quantization, achieving substantial memory and throughput benefits.

  • Motivation: KV-cache memory grows with batch size and context length, becoming a bottleneck alongside cache-loading overhead in batched LLM inference.The cache stores attention keys and values to avoid recomputation, but loading it can leave computational cores idle.
  • Motivation: Quantization is a simple, effective strategy for reducing KV-cache bytes, but prior work largely applied vanilla 4bit round-to-nearest quantization without deeply studying cache-element distributions.The introduction attributes this gap to the streaming nature of KV caches and related complications.
  • Quantization design: Key caches contain a few fixed large-magnitude channels, so KIVI quantizes them per-channel to confine quantization error within individual channels.Value caches lack an obvious outlier pattern and are quantized per-token because they mix values when computing attention outputs.
  • KIVI: KIVI is a plug-and-play 2bit KV-cache method that requires no fine-tuning and supports hardware-friendly implementation.It quantizes keys per-channel and values per-token; value tensors can be appended directly during autoregressive streaming, while residual cache portions remain in full precision.
  • Results: 2.6× peak memory usage reduction for Llama-2-7B enables up to 4× larger batch size and brings 2.35× throughput.The reported evaluation covers Llama, Mistral, and Falcon on popular generation tasks, with little to no accuracy drop stated for Llama-2-7B.

2. Background: Attention Inference-Time Workflow

Attention inference proceeds through prefill and autoregressive decoding: prompts generate and cache keys and values, while each new token updates the cache and computes attention. The resulting KV cache grows with batch size and sequence length, creating substantial memory demands, exemplified by OPT-175B requiring 1.2TB under a large workload.

  • Workflow: Attention inference has two phases: prefill generates KV cache from the input prompt, while decoding uses and updates it to generate tokens sequentially.The process repeats until a special token marks sentence completion.
  • Prefill Phase: During prefill, the input tensor is projected through key and value layer weights to produce tensors that are cached for decoding.The input tensor has shape X ∈ R^(b×lprompt×d), with b denoting batch size, lprompt prompt length, and d hidden size.
  • Decoding Phase: During decoding, each current token is projected into key and value outputs, which are first used to update the KV cache.The current token embedding has shape t ∈ R^(b×1×d).
  • Decoding Phase: The updated cache is then used with query-layer weights to calculate the attention output for the next-token computation.The attention output layer and other inference components are omitted for illustration.
  • Memory and Speed Analysis: 1.2TB is the KV-cache requirement for OPT-175B with batch size 512, prompt length 512, and output length 32, equal to 3.8 times the model weights.The KV-cache shape is b × (lprompt + lgen) × d, where lgen is the number of generated tokens.

3. Methodology

KIVI studies dimension-wise KV-cache quantization and finds that INT2 accuracy is best when keys are quantized per-channel and values per-token. It implements this asymmetric scheme for streaming decoding using residual full-precision windows and group-wise quantization.

  • Quantization Configuration: INT2 is most accurate when key cache uses per-channel quantization and value cache uses per-token quantization.Per-token quantization of both caches causes a notable accuracy drop at INT2, while per-channel quantization of values significantly worsens accuracy.
  • Analysis of Key Cache: Key-cache channels contain persistent large-magnitude outliers, so per-channel quantization confines errors within individual channels.Per-token quantization can lead to almost 5× larger attention score error than per-channel quantization.
  • Analysis of Value Cache: Value-cache errors are almost 15× smaller with per-token than per-channel quantization because sparse attention emphasizes only a few important tokens.Value cache lacks a clear channel-wise outlier pattern, and per-token quantization confines errors to individual tokens.
  • Algorithm: KIVI retains recent key and value entries in full precision, then quantizes completed residual groups and appends them to the stored cache.Per-channel key quantization uses group-wise processing with zero-padding when needed, while value quantization is appended along the token dimension.
  • Algorithm: The full-precision sliding window is crucial for desirable performance on hard tasks such as GSM8K.The residual length R is a hyperparameter, and for keys it should be divisible by the group size G.

4. Experiments

Experiments evaluate KIVI across multiple LLM families and normal- and long-context generation tasks, showing that asymmetric 2bit KV-cache quantization preserves quality and retrieval ability. They also examine quantization hyperparameters and demonstrate substantially higher serving capacity and throughput than FP16.

  • Experimental setup: KIVI is evaluated on Llama/Llama-2, Falcon, and Mistral using LM-Eval normal-context tasks and LongBench plus NIAH for long-context evaluation.LM-Eval covers CoQA, TruthfulQA, and GSM8K; NIAH measures long-context retrieval ability.
  • Quantization configuration: 2bit K per-channel, V per-token quantization consistently achieves the best results among tested fake-quantization configurations.This supports the asymmetric quantization design identified by the paper’s earlier KV-cache distribution analysis.
  • LM-Eval results: Up to 2% accuracy drop is observed for Llama and Mistral when moving from 16bit to KIVI’s 2bit KV cache.The comparison includes Llama-2-7B, Llama-2-13B, Falcon-7B, and Mistral-7B on CoQA, TruthfulQA, and GSM8K.
  • Long-context results: KIVI maintains retrieval ability with 2bit KV cache and shows minimal accuracy impact across hard long-context generation tasks in LongBench.LongBench evaluation covers Llama2, Falcon-7B, and Mistral-7B variants, while NIAH results are reported for Llama-3-8B-Instruct and Mistral-7B-Instruct-v0.2.

5. Related Work

Prior work scales LLM inference through quantization, system-level memory optimization, and KV-cache token eviction. KIVI is most closely related to KV-cache quantization methods, but distinguishes itself through per-channel key quantization motivated by key-cache outliers.

  • LLM Quantization: LLM quantization includes weight-only methods such as AWQ, GPTQ, and SqueezeLLM, which reduce model weights to lower precision.AWQ uses activation-aware INT4 and INT3 quantization, while GPTQ uses approximate second-order information.
  • KV-Cache Quantization: SmoothQuant balances activation and weight quantization and compresses KV cache to 8bit with minor performance loss, but accuracy drops significantly at 4bit or below.FlexGen is also described as using 4-bit group-wise quantization.
  • KV-Cache Quantization: KIVI uses per-channel quantization because key caches contain more outliers than value caches, extending observations reported by ATOM.The concurrent KVQuant work independently identified a similar observation and approach.
  • System-Level Optimization: System-level methods such as vLLM and S3 reduce KV-cache memory requirements while increasing throughput through PagedAttention or memory-usage prediction.These optimizations are orthogonal to KIVI and can also be applied with its algorithm.
  • KV-Cache Sparsification: KV-cache sparsification methods compress caches by evicting tokens, including H2O’s attention-based retention, Scissorhands’ persistence hypothesis, and StreamingLLM’s attention sinks.These methods retain only selected tokens to preserve performance.

6. Conclusion and Future Work

The paper analyzes KV cache element distributions and concludes that keys should be quantized per-channel while values should be quantized per-token. Based on these observations, it proposes tuning-free 2bit KIVI, which improves batching and throughput, with further implementation optimization planned.

  • Conclusion: KV cache analysis concludes that key caches should be quantized per-channel and value caches per-token.These quantization granularities are based on systematic analysis of KV cache element distributions in popular LLMs.
  • Conclusion: KIVI is a plug-and-play 2bit KV cache quantization algorithm that requires no tuning.The method is derived directly from the observed key- and value-cache distributions.
  • Conclusion: 4× larger batch sizes and 3.47× throughput are achieved by KIVI in real LLM workloads.The reported gains characterize KIVI’s practical serving benefits.
  • Future Work: Future work will further optimize the implementation to reduce quantization overhead.The supplied conclusion passage states this future direction but ends before specifying additional details.

A. Detailed Implementations

The KIVI implementation provides pseudocode for prefill and decoding, maintaining quantized key/value groups alongside residual caches and updating them as they reach the residual-length threshold.

  • Prefill: KIVI’s prefill computes keys and values, group-quantizes value tokens, key groups, and stores residual key/value caches.Value grouping uses the token dimension, while key quantization is applied through KeyQuant.
  • Decoding: During decoding, new keys and values append to residual caches, while quantized groups are retrieved from the KV cache.The algorithm concatenates tK and tV along the token dimension before checking residual lengths.
  • Decoding: When the key residual reaches R, KIVI quantizes it, appends the result to quantized key groups, and clears the residual tensor.This preserves a residual window while moving completed key blocks into the quantized cache.
  • Decoding: When the value residual exceeds R, KIVI quantizes older residual values, appends them to quantized value groups, and retains the latest R values.Attention output combines contributions from quantized groups and the retained residual values.
  • KeyQuant: KeyQuant partitions older keys along the channel dimension with group size G, leaving the newest residual keys unquantized.It computes r = l%R, quantizes XK[: l −r], and retains XK[l −r :].

B. NIAH Setting

The NIAH setting follows a passkey-retrieval prompt template with a 7-digit passkey embedded in Paul Graham Essays as irrelevant background text. The prompt instructs the model to find, memorize, and answer the hidden passkey.

  • B. NIAH Setting: The experiment largely follows Mohtashami and Jaggi’s passkey-retrieval template, using a 7-digit passkey and Paul Graham Essays as background filler.The setup is described as following Arize-ai and Reid et al. (2024).
  • B. NIAH Setting: The prompt instructs the model to find and memorize important information hidden inside a large amount of irrelevant text.It states that the model will later be quizzed about the important information.
  • B. NIAH Setting: The prompt places the 7-digit passkey within a Paul Graham Essays prefix and later asks, “What is the pass key?”The passkey is presented twice before the final retrieval question.

C. More Ablation Results

The ablation study finds that reducing KIVI’s residual length from 128 to 32 substantially improves memory compression and throughput without a substantial performance gap. Accordingly, the authors use residual length 32 across all benchmark datasets.

  • Residual Length Ablation: Residual length 32 gives KIVI a significantly higher memory compression rate and increased throughput.The efficiency evaluation identifies the shorter residual length as beneficial for both memory use and serving speed.
  • Residual Length Ablation: Changing the residual length from 128 to 32 does not produce a substantial performance gap.The ablation compares the two residual lengths and finds no substantial degradation from the shorter setting.
  • Benchmark Configuration: KIVI uses residual length 32 across all benchmark datasets.Tables 6 and 7 report comparisons involving 16bit, KIVI-2, and KIVI-4 at residual lengths 128 and 32.

D. More Experimental Results

The section reports additional LongBench evaluations of KIVI on Llama3-8B, Mistral-7B-v0.2, and LongChat-7B-v1.5, alongside a Needle-in-a-Haystack test. The experiments use model-specific context lengths and KIVI configurations detailed in Tables 8–10.

  • LongBench evaluations: Additional LongBench results evaluate KIVI with Llama3-8B, Mistral-7B-v0.2, and LongChat-7B-v1.5.The results appear in Tables 8, 9, and 10, respectively.
  • Needle-in-a-Haystack test: A Needle-in-a-Haystack test follows the original passkey-retrieval format while incorporating modifications from Arize-ai and Gemini 1.5.The results are shown in Figure 4.
  • LongBench evaluations: Llama-3-8B-Instruct uses 8K context, grouped-query attention with 8 KV-cache heads, group size 32, and residual length 128 for KIVI-2 and KIVI-4.The baseline is full precision.
  • LongBench evaluations: Mistral-7B-Instruct-v0.2 uses 32K context, grouped-query attention with 8 KV-cache heads, group size 32, and residual length 128 for KIVI-2 and KIVI-4.The baseline is full precision.
  • LongBench evaluations: LongChat-7B-v1.5-32K is evaluated on LongBench with 32K context, group size 32, and residual length 128 for KIVI-2 and KIVI-4.The baseline is full precision.
Loading 2402.02750v2…