KV Cache Quantization (KIVI)#

KIVI est une méthode de quantification 2-bit du KV cache (key-value cache) qui réduit la mémoire de 2.6× et permet des batch sizes 4× plus grands, sans tuning et sans dégradation de qualité. Son innovation clé : une approche asymétrique — quantification per-channel pour les keys, per-token pour les values.

Le problème : le KV cache est un bottleneck mémoire caché#

Pendant la génération autoregressive d'un LLM, chaque token produit doit « se souvenir » de tous les tokens précédents. Ce souvenir est stocké dans le KV cache : pour chaque couche d'attention et chaque tête, on garde les vecteurs Key et Value de chaque token passé.

Croissance explosive du KV cache#

Modèle: LLaMA-7B (32 layers × 32 heads × 128 dim), précision FP16. En multi-batch (32 requêtes × 4K tokens), le KV cache peut dépasser 100 GB !

Pourquoi le KV cache devient dominant#

Le KV cache DOMINE avec contexte long + batch ! Quantifier les poids (GPTQ/AWQ) ne suffit pas : le KV cache reste en FP16 et explose.

Niveau néophyte : l'analogie#

Imaginez une réunion où chaque participant prend des notes. Plus la réunion dure, plus les notes s'accumulent. Au bout de 2 heures, vous avez 50 pages de notes qui encombrent votre bureau. Vous ne pouvez plus travailler efficacement.

Le KV cache, c'est ces notes. KIVI, c'est un système de sténographie ultra-condensée : au lieu d'écrire chaque mot en entier, vous utilisez des abréviations à 2 niveaux seulement (2-bit). Au début, ça semble trop extrême — mais KIVI a découvert que les notes des Keys et les notes des Values ont des structures différentes : il faut les abréger différemment. C'est cette asymétrie qui rend le 2-bit viable.

Niveau technique : l'algorithme KIVI#

Découverte clé : distributions asymétriques des Keys et Values#

KIVI part d'une observation empirique fondamentale : les distributions du Key cache et du Value cache sont structurellement différentes et nécessitent des stratégies de quantification différentes.

Architecture : attention avec KV cache quantifié#

KEY INSIGHT : le KV cache est stocké en 2-bit en mémoire mais déquantifié en FP16 pendant le calcul d'attention. Le gain = mémoire (2.6×), pas compute.

Configuration 2-bit tuning-free#

KIVI n'a aucun hyperparamètre à tuner. La configuration par défaut fonctionne universellement :

Pourquoi le 2-bit fonctionne (sans dégradation)#

Caractéristiques#

Résultats#

Qualité préservée#

Conclusion : perte de qualité négligeable (< 0.05 perplexity) avec une compression 8× du KV cache.

Mémoire et throughput#

Réductions : mémoire peak 2.6×, batch size 4× plus grand, throughput 2.9× supérieur.

Autres approches de quantification du KV cache#

KIVI n'est pas la seule méthode. Voici un panorama :

1. Q-Hitter / KVQuant / autres méthodes#

2. Sensitivity analysis (KVQuant)#

Une approche alternative consiste à analyser la sensibilité de chaque channel du KV cache et à allouer plus de bits aux channels critiques :

→ Mixed-precision adaptatif, plus précis que KIVI mais plus complexe (nécessite analyse + calibration).

3. Quantification + low-rank compensation (GEAR)#

GEAR combine quantification INT2/INT4 avec une correction low-rank de l'erreur de quantification :

  KV_quantifié = Quant(KV) + LowRank(erreur)
                 ↑                ↑
            INT2/INT4         SVD de l'erreur de quantification
                             (quelques vecteurs capturent l'essentiel)

Exemple pratique#

Installation et utilisation de KIVI#

# Cloner le repository KIVI
git clone https://github.com/jy-yuan/KIVI.git
cd KIVI

# Installer les dépendances
pip install -e .

Quantification du KV cache avec KIVI#

from kivi import KIVIConfig, KIVIKVCache
from transformers import AutoModelForCausalLM, AutoTokenizer

# 1. Charger le modèle (poids en FP16, INT4, ou autre — KIVI est orthogonal)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    torch_dtype="float16",
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# 2. Activer KIVI — configuration 2-bit par défaut (tuning-free)
kivi_config = KIVIConfig(
    bits=2,              # 2-bit quantization
    group_size=32,       # group size par défaut
    residual_length=128, # longueur de fenêtre résiduelle en FP16
)

model.enable_kivi(kivi_config)

# 3. Génération — le KV cache est automatiquement quantifié en 2-bit
input_ids = tokenizer("Expliquez la relativité générale :", return_tensors="pt").input_ids
output = model.generate(input_ids, max_new_tokens=512)

print(tokenizer.decode(output[0]))

KIVI avec vLLM#

# vLLM supporte la quantification du KV cache
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-2-7b-hf \
    --kv-cache-dtype fp8 \    # FP8 KV cache (alternative à KIVI INT2)
    --dtype float16

Note : vLLM supporte nativement le KV cache en FP8 (--kv-cache-dtype fp8). Pour INT2 KIVI, utiliser le repository officiel KIVI.

Mesurer l'impact mémoire#

import torch

# Avant KIVI
torch.cuda.reset_peak_memory_stats()
# ... génération de 2048 tokens sans KIVI ...
mem_before = torch.cuda.max_memory_allocated() / 1e9
print(f"Mémoire sans KIVI : {mem_before:.1f} GB")

# Avec KIVI
model.enable_kivi(KIVIConfig(bits=2))
torch.cuda.reset_peak_memory_stats()
# ... génération de 2048 tokens avec KIVI ...
mem_after = torch.cuda.max_memory_allocated() / 1e9
print(f"Mémoire avec KIVI : {mem_after:.1f} GB")
print(f"Réduction : {mem_before/mem_after:.1f}×")

Combinaison avec d'autres méthodes#

KIVI est orthogonal à la quantification des poids — il peut être combiné avec n'importe quelle méthode :

EXEMPLE : GPTQ INT4 poids + KIVI INT2 KV cache
- Poids : 14 GB → 3.5 GB (GPTQ INT4, 4× compression)
- KV cache : 16 GB → 2 GB (KIVI INT2, 8× compression)
- Total : 30 GB → 5.5 GB (5.5× compression globale !)

Limitations#

1. Quantifie uniquement le KV cache — les poids et activations ne sont pas affectés (combiner avec GPTQ/AWQ pour une compression complète).
2. Overhead de déquantification — le KV cache doit être déquantifié à la volée pendant le calcul d'attention, ajoutant un léger coût compute.
3. Efficacité dépend du contexte — plus pertinent pour les longs contextes (≥ 2048 tokens). Pour des contextes très courts, le KV cache est négligeable.
4. Fenêtre résiduelle FP16 — KIVI garde les derniers tokens en FP16 (residual length) pour éviter les erreurs sur les tokens récents, ce qui ajoute une petite surcharge mémoire.
5. Implémentation hardware-aware — les gains de throughput nécessitent une implémentation optimisée du kernel de déquantification.

Références#

• Papier arXiv : KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache (Liu et al., ICML 2024)
• Code source : github.com/jy-yuan/KIVI
• KVQuant : arXiv:2401.18079 — sensitivity-based KV quantization
• GEAR : arXiv:2403.05527 — KV quantization + low-rank compensation
• Voir aussi : QoQ / QServe · GPTQ · AWQ · SmoothQuant · Index quantification