Aller au contenu principal

NVIDIA réécrit le stockage d’état des modèles hybrides

NVIDIA réécrit le stockage d'état des modèles hybrides

Pourquoi un modèle conçu pour décoder vite se retrouve-t-il freiné par sa propre mémoire ? La question concerne désormais toute une génération de modèles hybrides, et NVIDIA, avec le Dao AI Lab de Princeton, vient d’y répondre avec une technique baptisée ReplaySSM.

Des couches rapides sur le papier

Un modèle hybride intercale beaucoup de couches SSM (state-space model, modèle à espace d’états) et quelques couches d’attention. Cette architecture est celle de NVIDIA Nemotron 3, de Qwen3.5 jusqu’à Qwen3.8, de Kimi-K3 et de GLM-5.3-Flash. Selon les modèles, la couche SSM est de type Mamba-2, Gated DeltaNet (GDN) ou Kimi Delta Attention (KDA).

L’intérêt tient à une différence de nature. L’attention garde un cache clé-valeur (KV cache) qui grossit à chaque token, comme un carnet dont on ajoute une page à chaque mot lu. Une couche SSM, elle, résume tout le passé dans un état de taille fixe : une fiche unique qu’on réécrit sans cesse. Sur le principe, la fiche coûte moins cher que le carnet.

Sur le terrain, NVIDIA constate pourtant que, aux tailles de batch utilisées en production, les couches SSM peuvent devenir le goulot de latence plutôt que l’attention, et cela jusqu’à 100 000 tokens de contexte.

La fiche qu’on réécrit à chaque token

Prenons la couche Mamba-2 de Nemotron 3 Ultra. À chaque token, le décodage lit l’état précédent dans la HBM (high-bandwidth memory, la mémoire rapide du GPU), lui applique un facteur de décroissance et une mise à jour de rang 1, calcule la sortie, puis réécrit l’état complet. GDN et KDA ajoutent une correction de type règle delta, mais leurs mises à jour s’écrivent sous la même forme, donc le raisonnement vaut pour eux aussi.

Le noyau lit et écrit tout l’état pour très peu d’arithmétique : c’est la bande passante de la HBM qui limite, pas la puissance de calcul. Deux facteurs aggravent la note, mesurés sur Nemotron-3-Super-NVFP4 avec un batch de 256 sur un GPU Blackwell Ultra B300 :

  • le modèle compte cinq fois plus de couches SSM que de couches d’attention ;
  • l’état FP32 (nombres sur 32 bits) d’une seule couche pèse 4 MiB par séquence, quand une entrée de KV cache en FP8 (nombres sur 8 bits) pèse 0,5 KiB par token. Une couche SSM transporte donc autant de mémoire que le cache d’une couche d’attention sur 8 000 tokens.

Trois ennuis nés de la même fusion

Le décodage classique fait tout en même temps : il met à jour l’état, lit la sortie dedans et le range pour le token suivant. NVIDIA fait remonter trois difficultés de cette fusion.

La première est le trafic d’état, décrit ci-dessus. La deuxième est le retour en arrière nécessaire au décodage spéculatif, où un petit modèle propose des tokens que le grand modèle valide ou rejette : l’état ayant été écrasé à chaque pas, revenir sur un token refusé devient délicat. La troisième est l’accumulation d’erreurs quand l’état est quantifié, c’est-à-dire stocké avec moins de bits : un état réécrit à chaque token cumule les arrondis.

Or le décodage spéculatif et la quantification sont justement deux des leviers qui accélèrent de façon fiable les modèles à attention. Les modèles hybrides, eux, les héritaient mal.

Tenir le journal plutôt que le solde

ReplaySSM change la comptabilité. Au lieu de réécrire le solde du compte à chaque opération, on tient un journal. L’état est représenté par un point de contrôle (checkpoint) et un tampon (buffer) qui contient les entrées récentes. La sortie se calcule à partir du checkpoint et du tampon, sans toucher à l’état complet, et le checkpoint n’est réécrit que lorsque le tampon est plein.

Cette séparation entre le calcul de la sortie et la mise à jour de l’état s’attaque aux trois difficultés d’un coup. L’état n’est plus lu et écrit à chaque pas, ce qui réduit le trafic mémoire. Les tokens rejetés par le décodage spéculatif se retirent du tampon sans abîmer le checkpoint. Et l’état quantifié n’est réécrit qu’à intervalles espacés, ce qui limite les arrondis cumulés. Le titre du papier le résume : cacher les entrées, pas l’état.

Une mise à jour du moteur d’inférence suffit

Le coût de servir un modèle se joue en grande partie sur la bande passante mémoire. Si les couches SSM cessent de monopoliser la HBM, un même GPU peut tenir davantage de séquences en parallèle ou répondre plus vite à batch égal. La technique est déjà intégrée à NVIDIA TensorRT-LLM, à vLLM et à FlashInfer, et elle tourne en production dans GLM-5.3-Flash et Kimi-K3. Pour une équipe qui héberge ses propres modèles hybrides, l’adoption passe donc par une mise à jour du moteur d’inférence, sans réentraînement.

L’extrait publié par NVIDIA décrit le mécanisme et les goulots, mais le détail des gains chiffrés se trouve dans le papier ReplaySSM: Cache SSM Inputs, Not State, avec son code publié en ligne. Les accélérations selon la longueur de contexte et le format de quantification sont à vérifier avant de dimensionner une flotte de GPU.

Si les modèles hybrides se généralisent, l’état SSM rejoindra probablement le KV cache parmi les postes de réglage que toute équipe d’inférence finit par maîtriser.

Sources

Ils m’ont fait confiance

« Il ne se contente pas de corriger les symptômes, il cherche à comprendre l'origine des problèmes et à sécuriser les modifications effectuées. J'ai réellement le sentiment d'avoir trouvé un développeur qui comprend à la fois la technique et les enjeux globaux du projet. »

Évaluation client · projet WordPress · septembre 2026

5,0/5 sur 17 évaluations

Faire appel à mes services →

Laisser un commentaire

Votre adresse e-mail ne sera pas publiée. Les champs obligatoires sont indiqués avec *