Cache de préfixe hybride KDA–MLA¶

La section d'infrastructure la plus dense du rapport, et l'exemple le plus net de ce qu'une innovation architecturale coûte en aval.

Le problème¶

Chaque bloc de Kimi K3 contient trois couches KDA et une couche Gated MLA, dont les caches diffèrent fondamentalement :

	Cache MLA	État KDA
Structure	Une entrée par jeton	Une grande matrice par requête
Croissance	Linéaire en \(T\)	Fixe
Nombre de copies	Une par jeton	Une par requête
Mise à jour	Ajout	En place

La contrainte de couplage

Un préfixe mis en cache n'est réutilisable que si les deux caches peuvent être restaurés ensemble, à la même frontière.

Pourquoi le cache de préfixe classique s'effondre ici¶

Le cache de préfixe standard fonctionne par hachage de blocs : seuls les blocs complets sont hachés, donc seuls les préfixes alignés sur les blocs sont réutilisables.

Ce couplage se brise chez Kimi K3 :

Le hachage de blocs exige une taille de bloc unique partagée par toutes les couches.
Un succès de préfixe n'est réutilisable que si l'état KDA à la frontière du succès a été persisté.
Or une couche KDA maintient un unique grand état récurrent par séquence, pas des entrées par jeton. Les instantanés ne sont donc abordables qu'à des frontières rares.
La taille de bloc partagée est donc forcée à 1 024–6 144 jetons.
Et comme le hachage est lié au bloc de stockage, la granularité de hachage l'est aussi — alors que les entrées par jeton de MLA toléreraient des blocs bien plus fins.

Le résultat

À une granularité aussi grossière, le cache est presque inutile :

toute requête plus courte qu'un bloc ne peut jamais être réutilisée ;
le prefill par morceaux (chunked prefill) n'exporte aucun préfixe cacheable tant qu'il n'a pas franchi une frontière de bloc complète.

La solution : découpler les deux granularités¶

L'idée centrale

Le hachage de préfixe opère sur des blocs de hachage fins (par exemple 512 jetons) à l'intérieur des pages MLA.
Le bloc physique reste l'unité d'allocation grossière.
Pour KDA, l'alignement va dans l'autre sens : les checkpoints de l'état récurrent ne sont sauvegardés qu'à un sous-ensemble épars des extrémités de blocs de hachage MLA — les seules positions qu'une recherche pourra jamais référencer.

MLA KV   ┌────┬────┬────┬────┬────┬────┬────┬────┬────┬────┬────┬────┐
         │ ██ │ ██ │ ██ │ ██ │ ██ │    │    │    │    │    │    │    │
         └────┴────┴────┴────┴────┴────┴────┴────┴────┴────┴────┴────┘
         └──────────────── bloc physique = 6144 jetons ──────────────┘
          └512┘  blocs de hachage
                             ▲
KDA ckpt   ○    ○    ●    ○  ★    ○    ○    ○    ○    ○    ○    ○
                             │
                    frontière de succès B = 2560 = 5 × 512

○ = pas de checkpoint     ● = checkpoint persisté     ★ = succès

La mise en page unifiée¶

Plutôt que de maintenir un gestionnaire distinct par type de cache — ce qui dupliquerait la logique d'allocation, d'éviction et de transfert — Kimi K3 empaquette les états KDA dans le même pool de blocs paginés que le KV MLA, en unifiant les pages à la même taille en octets.

Les deux types de page partagent ainsi une seule implémentation de l'allocation, du comptage de références et de l'éviction.

À l'intérieur d'une page, les états de toutes les têtes sont stockés contigûment, tête par tête, de sorte que le flux d'octets de chaque tête est autonome et constitue l'unité minimale de transfert inter-nœuds.

Le bénéfice pour la désagrégation prefill/decode

Quand les nœuds de prefill et de decode adoptent des degrés de parallélisme tensoriel différents, le réagencement est effectué sur le chemin de transfert, avec zéro remaniement côté GPU.

Une remarque du rapport qui vaut d'être citée

This asymmetry proved useful during development: any type-confused access yields garbage rather than plausible data — a zero-overhead sanity check on the pooled layout.

Comme les deux types de cache ont des structures radicalement différentes, une confusion de type produit du bruit manifeste plutôt que des valeurs plausibles mais fausses. Le bug se signale de lui-même.

C'est une observation de génie logiciel : rendre les erreurs bruyantes plutôt que silencieuses.

Le fonctionnement du prefill¶

Côté MLA. Une page partiellement remplie est enregistrée dans l'index du cache de préfixe sous le hachage chaîné de son dernier bloc de hachage complet — où chaque hachage couvre tous les blocs de hachage précédents, de sorte que faire correspondre une extrémité certifie tout le préfixe jusqu'à elle. L'extrémité enregistrée avance à mesure que la page se remplit.

Côté KDA. Après chaque passe avant, le noyau KDA persiste l'état récurrent à la dernière position alignée sur un bloc de hachage qui a été traitée.

Comme les checkpoints sont volumineux :

les checkpoints intermédiaires rendus obsolètes par la progression de la requête sont recyclés ;
ceux situés aux frontières de tour de conversation sont conservés pour la réutilisation inter-requêtes.

Les checkpoints sont en lecture seule

Un succès restaure l'état en le copiant dans l'état courant privé de la requête avant la prochaine passe avant, et les nouveaux checkpoints sont écrits dans des emplacements neufs.

Ainsi, un checkpoint visible par d'autres requêtes n'est jamais muté en place — condition indispensable au partage concurrent.

La recherche en deux étages¶

Étage MLA : apparier des blocs physiques entiers par hachage chaîné, et, au premier bloc manquant, se rabattre sur les extrémités de hachage à l'intérieur de ce bloc — de sorte que les pages partiellement remplies restent atteignables.
Étage KDA : exiger un checkpoint à la frontière candidate dans chaque groupe de cache KDA, chacun maintenant un état récurrent indépendant.

Le succès est la plus longue frontière satisfaisant les deux étages — toujours un multiple du bloc de hachage, jamais contraint d'être un multiple du bloc physique.

L'exemple du rapport

Une requête dont les 2 800 premiers jetons correspondent au préfixe caché touche à \(B = 2560 = 5 \times 512\), profondément à l'intérieur d'un bloc physique de 6 144 jetons, et reprend le prefill au jeton \(B\) au lieu de recalculer \([0, B)\).

La cohérence sous ordonnancement concurrent¶

Trois mécanismes, chacun dicté par un mode de défaillance concret. Le contexte : un bloc touché est à la fois une entrée de cache partagée et le point de croissance d'une requête privée, et les groupes MLA et KDA doivent s'accorder sur chaque frontière de succès.

Défaillance 1 — éviction d'un bloc que l'on vient de toucher

Tous les groupes de cache tirent leurs blocs d'une seule liste libre partagée. Allouer une copie privée pour un groupe pourrait donc évincer un bloc qu'un autre groupe vient de toucher.

Parade : tout bloc touché est épinglé dans tous les groupes avant qu'aucune allocation ne soit faite.

Défaillance 2 — lire les octets du propriétaire précédent

La copie vers le bloc privé s'exécute sur GPU immédiatement avant la passe avant. Un bloc alloué ou enregistré dans le pas d'ordonnancement courant livrerait encore les octets du propriétaire précédent à un lecteur.

Parade : de tels blocs sont exclus de l'appariement tant que leurs copies ne sont pas terminées.

Défaillance 3 — checkpoint présent dans un groupe seulement

Un checkpoint ne peut restaurer une requête que s'il existe dans chaque groupe KDA.

Parade : évincer le checkpoint d'un groupe invalide atomiquement ses frères — un checkpoint est soit atteignable dans tous les groupes, soit dans aucun.

Le résultat¶

L'énoncé final du rapport

With these mechanisms, every registered state always corresponds to exactly its declared token prefix, and prefix caching for hybrid KDA–MLA models reaches the same generality as for full-attention models: any shared prefix is reusable at any 512-token boundary, independently of request length, chunking, or scheduling interleaving.

Le désavantage structurel de l'architecture hybride est entièrement compensé. C'est un travail d'ingénierie considérable pour un résultat qui, de l'extérieur, se résume à « le cache marche normalement ».

Vérification de compréhension¶

Pourquoi conserver les checkpoints aux frontières de tour de conversation en particulier ?

Parce que c'est exactement là que les requêtes suivantes reprendront. Une session agentique envoie une requête par tour ; le préfixe partagé entre deux requêtes successives se termine à la fin du tour précédent. Placer les checkpoints durables là maximise le taux de succès inter-requêtes pour un coût de stockage minimal.

Qu'est-ce qu'un « hachage chaîné » et pourquoi est-il nécessaire ?

Le hachage du bloc \(n\) inclut le hachage du bloc \(n-1\), qui inclut celui du \(n-2\), etc. Conséquence : deux séquences ayant le même hachage au bloc \(n\) ont nécessairement le même préfixe complet jusque-là. Sans chaînage, deux séquences différentes partageant un bloc identique au milieu produiraient une fausse correspondance — et un cache corrompu.

Pourquoi ne pas simplement sauvegarder l'état KDA à chaque bloc de 512 jetons ?

Parce que l'état KDA de Kimi K3 pèse ~217 Mio par requête (69 couches × 96 têtes × 128 × 128 en BF16). Un checkpoint tous les 512 jetons sur 1 M de contexte ferait 2 048 checkpoints, soit plus de 400 Gio par requête. C'est précisément pourquoi les checkpoints doivent être épars, et pourquoi le découplage des granularités est nécessaire.

Chapitre précédent : Sandboxes AgentENV · Chapitre suivant : Noyaux d'inférence