Aller au contenu

Inférence : cache KV, décodage, quantification

Ce chapitre décrit ce qui se passe après l'entraînement, quand le modèle sert des requêtes. C'est un tiers du rapport Kimi K3 — et la partie la plus souvent ignorée des explications grand public.

Le cache KV

À chaque nouveau jeton généré, l'attention doit comparer la requête courante aux clés de toutes les positions précédentes. Les recalculer à chaque pas serait absurde : elles ne changent pas.

On les stocke donc : c'est le cache KV.

À retenir

Le cache KV est la principale consommation mémoire de l'inférence longue. Il croît linéairement avec la longueur du contexte et linéairement avec le nombre de requêtes simultanées.

Chez Kimi K3, la situation est particulière car l'architecture est hybride :

Type de couche Nombre Ce qui est mis en cache Taille
Gated MLA 24 Vecteur latent compressé par jeton (rang 512) Croît avec \(T\)
KDA 69 État récurrent \(\mathbf{S}\), une copie par requête Fixe

C'est cette asymétrie qui rend le cache de préfixe hybride de Kimi K3 techniquement délicat — et qui justifie une section entière du rapport.

Le cache de préfixe (prefix caching)

Dans un usage agentique, chaque requête d'une session reprend tout l'historique et y ajoute quelques jetons. Recalculer 400 000 jetons de préfixe à chaque tour serait ruineux.

Le cache de préfixe conserve le cache KV entre les requêtes : si une nouvelle requête commence par un préfixe déjà traité, on reprend directement à partir de là.

Intuition

Le rapport donne le chiffre parlant : dans un usage de codage typique à 1 M de contexte, l'entrée porte un préfixe de 400 K jetons mais ne demande qu'un incrément de 4 K. Un succès de cache est plusieurs ordres de grandeur moins cher qu'un échec.

C'est aussi ce que reflète la tarification API : 0,30 $/M jetons pour une entrée en cache contre 3,00 $/M hors cache — un facteur 10.

Le décodage spéculatif

Le décodage est limité par la mémoire, pas par le calcul : pour produire un seul jeton, il faut relire 104 milliards de paramètres. Le GPU passe l'essentiel de son temps à attendre la mémoire.

Idée : produire plusieurs jetons candidats avec un petit modèle rapide (le brouillon), puis les vérifier tous en une seule passe du grand modèle.

brouillon  →  "le chat dort sur le"     (5 jetons proposés, très rapide)
grand modèle →  vérifie les 5 d'un coup  (1 passe au lieu de 5)
                 ✓ ✓ ✓ ✗ —              (3 acceptés, on repart de là)

Propriété essentielle

L'échantillonnage spéculatif est sans perte : la distribution des jetons produits est exactement celle du grand modèle seul. Ce n'est pas une approximation, c'est une réorganisation du calcul.

L'accélération dépend du taux d'acceptation :

\[ \text{taux d'acceptation} = \sum_{x \in \mathcal{V}} \min\!\left(p(x), q(x)\right) \]

où \(p\) est la distribution du grand modèle et \(q\) celle du brouillon.

L'affinement de Kimi K3

La pratique courante est d'entraîner le brouillon en minimisant la divergence KL avec le grand modèle. Mais minimiser la KL ne maximise pas le taux d'acceptation quand le brouillon a une capacité limitée. Kimi K3 optimise donc directement le logarithme négatif du taux d'acceptation (la perte LK), sans terme auxiliaire de vérité terrain.

Le brouillon est obtenu en affinant la couche MTP (Multi-Token Prediction) pré-entraînée, dont la structure coïncide avec celle d'un décodeur EAGLE-3. Son entrée fusionne des caractéristiques de bas, moyen et haut niveau — les sorties des 1er, 4e et dernier blocs AttnRes.

Point de vigilance

Le fichier config.json publié porte num_nextn_predict_layers: 0. La couche MTP décrite dans le rapport ne semble donc pas incluse dans le dépôt de poids public. Le décodage spéculatif tel que décrit n'est pas directement reproductible à partir des seuls poids publiés. Voir Ce qui est public et ce qui ne l'est pas.

La quantification

Réduire le nombre de bits par poids diminue la mémoire et accélère le décodage (moins d'octets à lire).

Format Bits/poids Taille de Kimi K3 (ordre)
FP32 32 ~11 To
BF16 16 ~5,6 To
FP8 8 ~2,8 To
MXFP4 4 (+ échelle/32) ~1,56 To (taille réelle du dépôt)

Deux approches :

  • PTQ (Post-Training Quantization) : quantifier après coup. Simple, mais dégrade la qualité.
  • QAT (Quantization-Aware Training) : entraîner avec la quantification simulée, pour que le modèle s'y adapte.

Kimi K3 fait du QAT sur tout le post-entraînement, du SFT jusqu'au RL inclus. Pendant le RL, déploiement et entraînement utilisent le même schéma — ce qui supprime l'écart entraînement/inférence, source classique d'instabilité.

Sélectivité

Seuls les poids des experts routés sont en MXFP4. Attention, experts partagés, projections latentes, routeurs, tête de sortie et encodeur visuel restent en précision supérieure. C'est vérifiable dans le champ ignore du quantization_config publié.

Prefill/decode disaggregation

Prefill et decode ont des profils opposés (calcul vs mémoire). Les exécuter sur la même machine force un compromis. La désagrégation les sépare sur des nœuds distincts, chacun optimisé pour son régime, le cache KV étant transféré de l'un à l'autre.

Kimi K3 gère un cas supplémentaire : les nœuds de prefill et de decode peuvent avoir des degrés de parallélisme tensoriel différents. Le réagencement des données est alors fait sur le chemin de transfert, sans réorganisation côté GPU. C'est pour cela que les états KDA sont stockés « tête par tête, de façon contiguë » : chaque tête est un flux d'octets autonome, unité minimale de transfert.

L'ordonnancement de flotte

À l'échelle d'un service, deux problèmes émergent, tous deux traités par Kimi K3 (voir Ordonnancement de flotte) :

  • Affinité de cache : router chaque session vers le cluster qui détient déjà son cache. Kimi K3 utilise un hachage cohérent qui épingle chaque session à deux clusters — un principal et un secondaire de secours — de sorte qu'une panne ne concentre pas tout le re-prefill sur une seule machine.
  • Contrôle d'admission par budget : le trafic mêle des requêtes de 2 K jetons et de 1 M. Le coût par requête varie sur trois ordres de grandeur. Sans cloisonnement, une rafale de requêtes longues sature le calcul et dégrade le temps de première réponse pour tout le monde. Kimi K3 alloue donc des budgets de ressources séparés par classe de requêtes.

Vérification de compréhension

Pourquoi le décodage spéculatif complique-t-il KDA en particulier ?

Parce que l'état récurrent est mis à jour en place à chaque pas. Si la vérification rejette une partie des jetons proposés, l'état a déjà avancé au-delà du dernier jeton accepté, et il n'existe pas de moyen trivial de revenir en arrière. Garder un instantané par position de brouillon résoudrait le problème mais multiplierait le trafic mémoire.

La solution de Kimi K3 : ne mettre en cache que les entrées projetées des jetons brouillons (bien plus petites que l'état), puis rejouer la récurrence sur puce pour reconstruire l'état accepté. Détail en Noyaux d'inférence.

Combien coûte une requête de 1 M de jetons chez Kimi K3, en entrée seule ?

Au tarif publié de 3,00 $/M jetons hors cache : environ 3 $ pour un seul prefill complet. Avec cache : 0,30 $. D'où l'importance économique du cache de préfixe dans les usages agentiques, où le même contexte est relu des centaines de fois.


Chapitre précédent : Post-entraînement · Chapitre suivant : Matériel et parallélismes