FlashQuant : Compression KV-Cache 7.5x CUDA pour LLM

FlashQuant est un projet open source écrit en C++17 et CUDA qui s'attaque au goulot d'étranglement le plus critique des LLM modernes en inférence : la consommation mémoire du KV-cache. Sur un modèle frontière comme Llama 3.1 70B servi avec un contexte de 32K tokens, le KV-cache occupe entre 50 % et 90 % de la VRAM disponible, dépassant régulièrement les 40 Go par requête longue. FlashQuant introduit un pipeline de compression extrême combinant quantization INT4 par groupe, dictionnaires de codebooks appris et kernels CUDA dédiés, atteignant un ratio de compression moyen de 7.5x avec une dégradation de perplexité inférieure à 1 % sur MMLU et HumanEval. Conçu comme un plugin pour vLLM 0.6.x+, FlashQuant s'intègre directement à PagedAttention et permet de servir des contextes longs sur GPU consumer (RTX 4090, RTX 5090, A6000) là où le modèle ne tenait précédemment que sur des H100 80 Go. Cette page projet documente l'architecture, les benchmarks, l'installation et la roadmap. Le code source complet est publié sous licence Apache 2.0 sur github.com/ayinedjimi/FlashQuant.

Le probleme : le KV-cache devore la VRAM

Lorsqu'un LLM autoregressif genere du texte, chaque nouveau token doit acceder aux representations Key et Value de tous les tokens precedents. Pour eviter de recalculer ces tenseurs a chaque pas, l'inference moderne stocke un KV-cache. Sa taille croit lineairement avec la longueur de contexte, le nombre de couches, le nombre de tetes d'attention et la dimension par tete. Pour Llama 3.1 70B en FP16, le KV-cache atteint environ 1.25 Go par 1K tokens, soit plus de 40 Go pour un contexte de 32K. Sur des modeles MoE ou des contextes 128K, ce chiffre explose au-dela des 100 Go par sequence active. Resultat : meme un H100 80 Go ne peut servir qu'une seule session simultanement, et les GPU consumer sont totalement exclus du jeu. Cette contrainte limite la capacite de batching, augmente le cout par token et empeche le deploiement de RAG long-contexte sur du materiel abordable.

Concretement, le calcul est implacable : pour Llama 3.1 70B avec 80 couches, 8 KV heads (Grouped Query Attention), dimension 128 par tete, en FP16 (2 octets), un seul token consomme 80 x 8 x 128 x 2 x 2 = 327 Ko de KV-cache. A 32K tokens, cela donne bien 10.7 Go par sequence et par direction (K et V), soit 21.4 Go au total avant overhead PagedAttention. En pratique, vLLM reserve des blocs de 16 tokens et impose une fragmentation qui pousse l'occupation reelle a 30-40 Go. Pour Llama 3.1 405B, ces chiffres sont multiplies par six. Cette tyrannie memoire transforme l'inference longue-contexte en privilege reserve aux clusters multi-GPU. FlashQuant casse precisement ce verrou.

La solution FlashQuant : INT4 + dictionnaires cibles

FlashQuant attaque le probleme en quantifiant le KV-cache de FP16 vers INT4 avec des codebooks specialises par tete d'attention. Plutot qu'une quantization uniforme, le projet identifie les tetes "outlier-heavy" et leur attribue un dictionnaire de scales et zero-points appris hors-ligne sur un corpus de calibration. Les tetes "regulieres" utilisent une quantization symetrique par groupe de 128 elements. Cette hybridation permet d'obtenir un ratio reel de 7.5x (au-dela du 4x theorique d'INT4 pur) en exploitant la sparsite des distributions K et V. Le pipeline complet est ecrit en CUDA avec fusion des operations dequantize-attention pour eviter tout aller-retour memoire global. Pour comprendre les fondements de la quantization, voir notre guide sur AWQ et la quantization INT4.

L'innovation cle reside dans l'approche en deux niveaux. Premier niveau : une quantization INT4 par groupe (group_size=128) qui transforme 128 valeurs FP16 en 128 valeurs INT4 + 1 scale FP16 + 1 zero-point INT4. Sans codebook, le ratio brut serait de 16/4.25 = 3.76x. Deuxieme niveau : un codebook global de 64 entrees apprises par head qui factorise les patterns recurrents de scales et zero-points. Au lieu de stocker un scale FP16 par groupe (16 bits), on stocke un index 6 bits dans le codebook, soit un gain net supplementaire de 10 bits par groupe. Combine avec un packing dense des INT4 en uint32, le ratio effectif atteint 7.5x sur les modeles Llama 3 testes. Le codebook lui-meme est appris par K-means sur les distributions de scales observees pendant la calibration.

Difference vs autres approches

Le paysage de la compression KV-cache est dense. La compression native vLLM (FP8 e4m3) atteint un ratio 2x sans codebook. KIVI (Liu et al., 2024) propose INT2 par-canal pour K et par-token pour V, ratio 4x mais perte 2-3 % perplexite. KVQuant (Hooper et al., 2024) ajoute des outliers stockes en FP16, ratio 4-5x. ATOM (Zhao et al., 2024) cible la quantization W4A4 du modele entier, sans focus KV. FlashQuant se differencie par trois choix : 1) ratio 7.5x effectif via codebooks par tete, 2) integration native PagedAttention sans copie supplementaire, 3) calibration en moins de 10 minutes sur 512 echantillons. Le papier de reference KIVI est disponible sur arxiv.org/abs/2402.02750.

Architecture : kernels CUDA personnalises

FlashQuant se compose de quatre kernels CUDA principaux. Le premier, quantize_kv_paged, prend un bloc PagedAttention de 16 tokens FP16 et produit un bloc INT4 + scale FP16 + zero INT4. Le deuxieme, dequantize_attention_fused, fusionne dequantization et calcul d'attention sur les sequences flash. Le troisieme, codebook_lookup, gere la table de codebooks par tete avec acces shared-memory. Le quatrieme, recalibrate_online, ajuste les scales toutes les N requetes en cas de drift de distribution. Les scales et zero-points sont apprenables (gradient-free, via least-squares) et stockes per-channel pour Key et per-token pour Value, suivant les heuristiques de KIVI et KVQuant. La complexite memoire effective est ramenee a 2.13 bits par element (4 bits + overhead 7 % codebooks).

Le code source est organise en trois couches. La couche flashquant/csrc/ contient les kernels CUDA en C++17 avec des templates pour specialiser sur sm_80, sm_89 et sm_90. La couche flashquant/python/ expose une API PyTorch via pybind11 et integre les hooks PagedAttention. La couche flashquant/calibration/ implemente la pipeline de calibration en pur Python avec scikit-learn pour le K-means des codebooks. Les kernels exploitent les tensor cores INT4 sur Hopper (sm_90) via les instructions MMA m16n8k64 qui multiplient des matrices INT4 directement sans dequantization explicite. Sur Ampere (sm_80), un fallback utilise les instructions DP4A. La granularite de scheduling suit celle de PagedAttention avec un bloc CTA par groupe de 16 tokens et 8 heads paralleles.

Precision preservee : moins de 1 % de degradation

FlashQuant a ete evalue sur cinq benchmarks standards : MMLU (5-shot), HumanEval (pass@1), GSM8K (8-shot CoT), TruthfulQA et LongBench. Sur Llama 3.1 70B Instruct, la degradation moyenne observee est de 0.7 % en MMLU (82.1 -> 81.5) et 0.4 % en HumanEval (78.2 -> 77.9). Sur LongBench (contextes 16K-32K), la degradation est inferieure a 1.2 %. Ces chiffres restent comparables aux meilleures approches publiees tout en offrant un ratio de compression 1.5x a 2x superieur. La perplexity sur WikiText-103 passe de 5.94 (FP16) a 6.01 (FlashQuant), soit +1.2 %.

La preservation de la precision repose sur deux mecanismes anti-degradation. D'abord, les tetes outlier-heavy identifiees pendant la calibration (typiquement 5-8 % des heads) recoivent un budget bits superieur via un codebook plus dense (256 entrees au lieu de 64). Ensuite, les premiers et derniers 32 tokens de chaque sequence (qui contiennent les sinks d'attention identifies par Xiao et al.) sont stockes en FP8 sans perte significative, suivant la litterature StreamingLLM. Cette protection des sinks evite les regressions catastrophiques sur les prompts courts. Sur des taches sensibles comme GSM8K (raisonnement mathematique), le score passe de 84.5 a 83.8 soit une degradation de 0.83 %, dans la meme enveloppe que le bruit de seed.

Pre-requis

• GPU : NVIDIA Ampere (sm_80) ou superieur. Hopper et Blackwell entierement supportes.
• CUDA Toolkit : 12.0 minimum, 12.4 recommande.
• vLLM : version 0.6.x ou superieure (testee jusqu'a 0.7.3).
• PyTorch : 2.4+ avec support CUDA 12.
• Compilateur : g++ 11+, support C++17 obligatoire.
• RAM systeme : 32 Go recommandes pour la calibration hors-ligne.

Installation

git clone https://github.com/ayinedjimi/FlashQuant.git
cd FlashQuant
pip install -r requirements.txt
python setup.py build_ext --inplace
pip install -e .

# Patch vLLM
python -m flashquant.patch_vllm --vllm-path $(python -c "import vllm; print(vllm.__path__[0])")

Le build CUDA prend environ 4 minutes sur un Threadripper. Le patch vLLM ajoute un hook dans le module vllm/attention/backends/flash_attn.py pour rediriger les operations KV vers les kernels FlashQuant.

Configuration

from vllm import LLM, SamplingParams
from flashquant import FlashQuantConfig

fq_config = FlashQuantConfig(
    compression_ratio=7.5,
    group_size=128,
    block_size=16,
    use_codebooks=True,
    codebook_path="./calibration/llama3_70b_codebook.pt",
    online_recalibrate=False,
)

llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    kv_cache_dtype="flashquant",
    flashquant_config=fq_config,
    max_model_len=32768,
)

Le parametre group_size=128 equilibre precision et overhead memoire. Reduire a 64 ameliore la fidelite mais reduit le ratio a 6.2x. online_recalibrate=True active un recalibrage adaptatif toutes les 1000 requetes (overhead 1-2 %).

Benchmarks : 40 Go -> 5.3 Go

Mesures realisees sur RTX 4090 24 Go et A100 80 Go avec Llama 3.1 70B Instruct (4-bit GPTQ pour les poids, FlashQuant pour le KV) :

Le throughput chute de 4.8 % par rapport au FP16 baseline mais permet de servir 7.5 fois plus de sequences concurrentes a contexte egal, soit un gain de throughput agrege superieur a 5x. En charge realiste avec batching dynamique vLLM, sur A100 80 Go, FlashQuant permet de tenir 56 sequences concurrentes a 8K tokens contre 8 en FP16. Le throughput agrege passe de 187 x 8 = 1496 tokens/s a 178 x 56 = 9968 tokens/s, soit une multiplication par 6.7. Le cout par million de tokens est divise dans la meme proportion. Sur un H100 SXM5, FlashQuant permet par ailleurs de servir Llama 3.1 70B avec contexte 128K en moins de 70 Go de VRAM totale (poids + KV-cache), un seuil jusqu'alors inatteignable sans offload disque.

Cas d'usage

• RAG long-contexte : ingestion de 64K-128K tokens de documents source sur GPU consumer.
• Agents multi-tool : conservation de 20+ tours d'outils sans truncation.
• Batch high-throughput : multiplication par 7 du nombre de requetes simultanees sur A100/H100.
• Edge inference : Llama 3.1 70B servable sur RTX 5090 32 Go avec contexte 16K.
• Offload hybride : combinaison avec KVortex pour offload VRAM/RAM pour atteindre 256K tokens.

Comparatif detaille vs alternatives

Limitations

FlashQuant a ete principalement valide sur les familles Llama 3 / 3.1, Mistral 7B et Qwen 2.5. Les modeles Mixtral 8x22B presentent des outliers d'experts qui degradent le ratio a 6.0x. Les modeles tres recents Llama 4 ne sont pas encore officiellement supportes (en cours d'integration). Les architectures non-MHA comme Mamba ou les modeles a state-space ne sont pas applicables. Le mode online_recalibrate reste experimental et peut introduire des oscillations de qualite sous charge tres heterogene. Enfin, sur GPU avec sm_80 strict (A100 originels), certains kernels tensor-core ne sont pas disponibles, reduisant le throughput de 8 %.

Autres limitations connues : pas de support actuel pour les modeles encoder-decoder (T5, BART) car la nature bidirectionnelle de l'encodeur ne beneficie pas de la compression KV. La calibration necessite un acces aux poids du modele en FP16, ce qui exclut les modeles distribues uniquement en formats deja quantifies (GGUF Q4_K_M par exemple) sans etape de dequantization prealable. Sur Windows, le support reste experimental car certains headers CUDA Linux-specifiques sont utilises pour l'integration vLLM. Enfin, l'overhead de la dequantization on-the-fly devient sensible (>10 %) sur des prompts tres courts (moins de 256 tokens) ou le KV-cache n'est pas le goulot d'etranglement : dans ces cas, desactiver FlashQuant via kv_cache_dtype="auto" est recommande.

Roadmap

• Q3 2026 : support Llama 4 et DeepSeek V3.
• Q4 2026 : mode hybride FP8/INT4 selectif par tete pour porter le ratio a 9x sur certains modeles.
• Q1 2027 : multi-GPU sharding compresse via NCCL avec all-gather sur INT4 packe.
• Q2 2027 : binding TensorRT-LLM et integration SGLang.
• Stretch : compression dynamique adaptative selon la longueur de sequence.

Comment contribuer

Le projet accueille les contributions sur github.com/ayinedjimi/FlashQuant. Les axes prioritaires sont : ajout de modeles dans la matrice de calibration, optimisation des kernels Hopper (TMA, WGMMA), tests d'integration avec vllm-project/vllm recent, documentation Python. Forker le repo, creer une branche feature/xxx, ouvrir une PR avec tests unitaires (pytest + nvprof traces). Le CI execute des benchmarks de regression sur MMLU. Discussions techniques bienvenues dans les issues GitHub. Voir egalement la liste de mes projets open source et les ressources open source curatees pour decouvrir d'autres outils complementaires.

Le repository inclut un guide CONTRIBUTING.md detaillant les conventions de code (clang-format pour C++, black + ruff pour Python), la procedure de signature DCO et les exigences de couverture de tests (>80 %). Les contributeurs nouveaux sont encourages a commencer par les issues taggees good-first-issue, qui couvrent typiquement l'ajout de modeles a la suite de calibration ou l'amelioration de la documentation. Les contributions plus avancees sur les kernels CUDA passent par une revue technique dediee. Le projet est sponsorise via GitHub Sponsors et accepte les dons pour financer du temps GPU sur les benchmarks. Une roadmap publique est tenue a jour dans le fichier ROADMAP.md, et les decisions architecturales majeures sont documentees dans des ADR (Architecture Decision Records) sous docs/adr/.

FAQ

Quelle est la difference entre FlashQuant et AWQ ?

AWQ quantifie les poids du modele (W4A16), FlashQuant quantifie le KV-cache a l'inference. Les deux sont complementaires : AWQ reduit la VRAM des poids (140 Go -> 35 Go pour Llama 70B), FlashQuant reduit la VRAM du KV-cache (40 Go -> 5.3 Go pour 32K tokens). Combines, ils permettent de servir Llama 70B sur 48 Go avec contexte 32K.

FlashQuant est-il compatible avec Llama 4 ?

Le support Llama 4 est en cours d'integration et planifie pour Q3 2026. La nouvelle architecture iRoPE et les ratios MoE necessitent une adaptation des codebooks. Une branche llama4-experimental est disponible pour les testeurs avances mais sans garantie de stabilite.

FlashQuant est-il stable pour la production ?

FlashQuant est en version 0.4.x (beta). Plusieurs deploiements pilotes sont en production sur des charges RAG entreprise (~10K requetes/jour) sans incident depuis 3 mois. Une version 1.0 stable est attendue Q3 2026 apres l'integration Llama 4 et la finalisation du recalibrage online.

Quel impact sur le batch size ?

Avec un KV-cache 7.5x plus petit, le batch size effectif augmente proportionnellement a contexte egal. Sur A100 80 Go avec Llama 70B et contexte 8K, le batch passe de 8 (FP16) a 56 (FlashQuant). En pratique, le throughput agrege augmente de 5x a 6x apres prise en compte des overheads de scheduling vLLM.

Pourquoi 7.5x et pas plus ?

Le ratio theorique INT4 pur est 4x. FlashQuant atteint 7.5x grace aux codebooks par tete et au packing optimise des scales. Aller au-dela (10x+) demanderait du INT2 ou des representations sparses qui degradent la perplexity au-dela du seuil acceptable de 2 %. Le mode hybride FP8/INT4 prevu Q4 2026 vise 9x en preservant la qualite. Des experimentations internes en INT3 + codebooks dense (8x) ont montre une degradation moyenne de 1.8 % sur MMLU, jugee trop importante pour devenir defaut. Le ratio 7.5x represente le sweet spot entre compression maximale et fidelite acceptable, valide empiriquement sur 12 modeles de la famille Llama 2/3/3.1 et Mistral.

Liens

• Repository : github.com/ayinedjimi/FlashQuant
• vLLM : github.com/vllm-project/vllm
• Papier KIVI : arxiv.org/abs/2402.02750
• Guide AWQ : AWQ Quantization INT4
• Guide vLLM : vLLM moteur d'inference
• KVortex (offload) : KVortex VRAM/RAM offloader
• Mes projets : page GitHub et ressources open source