ReplaySSM accélère le décodage des modèles hybrides en mettant en cache les entrées plutôt que l'état

NVIDIA et le Dao AI Lab de l'université de Princeton ont présenté ReplaySSM, une nouvelle méthode pour optimiser la phase de décodage des grands modèles de langage hybrides. Publiée en octobre 2026, cette approche repense la manière dont les modèles à espace d'états (SSM) stockent et mettent à jour leur état interne lors de l'inférence. En mettant en cache les jetons d'entrée récents plutôt qu'en écrasant continuellement la matrice d'état complète, la technique réduit considérablement l'utilisation de la bande passante mémoire et permet un décodage prédictif plus efficace.

Ce qui s'est passé

Les modèles hybrides, qui alternent couches SSM et couches d'attention traditionnelles, deviennent courants dans les architectures modernes telles que NVIDIA Nemotron 3, Qwen3.5 et Kimi-K3. Ces modèles utilisent des SSM pour gérer efficacement les longs contextes, remplaçant le cache clé-valeur croissant des mécanismes d'attention par un état de taille fixe. Cependant, avec des tailles de lot de service élevées, les couches SSM constituent souvent le principal goulot d'étranglement de latence car elles nécessitent des lectures et écritures fréquentes de grandes matrices d'état depuis une mémoire à haute bande passante.

L'équipe de recherche a identifié que le décodage conventionnel fusionne le calcul de sortie et la mise à jour de l'état en une seule étape, entraînant un trafic mémoire excessif. ReplaySSM découple ces opérations en représentant l'état comme un point de contrôle (checkpoint) accompagné d'un tampon d'entrées récentes. Cette modification permet au système de calculer les sorties sans écrire constamment l'état complet en mémoire, ne mettant à jour le point de contrôle que lorsque le tampon d'entrées est plein. La méthode est désormais intégrée aux principaux frameworks d'inférence, notamment TensorRT-LLM, vLLM et FlashInfer.

Comment cela fonctionne

Dans le décodage SSM standard, chaque étape de génération de jeton charge l'état précédent, applique une décroissance et une mise à jour, calcule la sortie et stocke le nouvel état. Ce processus est limité par la mémoire car l'intensité arithmétique est faible par rapport au mouvement des données. ReplaySSM modifie cela en déroulant la récurrence sur une petite fenêtre de jetons. Au lieu de mettre à jour l'état à chaque étape, il ajoute les vecteurs d'entrée du jeton actuel à un tampon. La sortie est calculée en utilisant le point de contrôle stocké et les entrées mises en tampon.

Lorsque le tampon atteint une taille prédéfinie, une opération de vidange (flush) intègre toutes les entrées mises en tampon dans le point de contrôle et l'écrit en mémoire une seule fois. Cela réduit le nombre d'écritures d'état d'environ moitié. Pour le décodage prédictif, où plusieurs jetons brouillons sont générés et vérifiés, ReplaySSM simplifie le retour arrière. Au lieu de sauvegarder des instantanés de l'état entier pour chaque jeton brouillon, le système ajoute simplement les entrées brouillon au tampon. Si un brouillon est rejeté, le pointeur du tampon se réinitialise, écartant les entrées non validées sans restauration coûteuse de l'état.

Détails clés

• ReplaySSM réduit le trafic mémoire lié à l'état d'environ 50 % en n'écrivant les points de contrôle qu'après remplissage d'un tampon d'entrées récentes.
• Sur les GPU NVIDIA Blackwell Ultra B300 avec une taille de lot de 256, les noyaux de décodage SSM s'exécutent 1,46x à 1,80x plus vite que les méthodes conventionnelles.
• Le débit du décodage prédictif augmente de 1,85x à 2,10x pour les modèles de classe 120B car le retour arrière ne nécessite plus de sauvegarder ni de restaurer des instantanés d'état.
• La technique prend en charge le stockage d'état en basse précision en conservant un point de contrôle en haute précision pour la récurrence tout en utilisant des copies en basse précision pour les lectures, limitant ainsi l'accumulation d'erreurs.
• Les modèles évalués incluent les variantes Nemotron-3 et Qwen3.5, montrant des gains constants à travers différentes architectures SSM comme Mamba-2 et Gated DeltaNet.
• La méthode est déjà déployée dans des systèmes de production tels que GLM-5.3-Flash et Kimi-K3.

Pourquoi c'est important

Pour les ingénieurs construisant et servant de grands modèles hybrides, la bande passante mémoire est souvent le facteur limitant du débit. En réduisant le trafic mémoire dominant dans les couches SSM, ReplaySSM permet des tailles de lot plus élevées sans augmenter la latence par jeton. Ceci est crucial pour un service rentable, car il permet une meilleure utilisation du matériel sur des GPU coûteux comme le H100 et le B300. Le gain de performance évolue avec la taille du lot, ce qui le rend particulièrement précieux pour les environnements de production à forte concurrence.

De plus, ReplaySSM rend le décodage prédictif viable pour les modèles basés sur SSM. Auparavant, le coût de sauvegarde des instantanés d'état pour le retour arrière annulait les bénéfices de la spéculation dans ces architectures. En traitant le tampon d'entrées comme un pointeur de cache clé-valeur, la méthode restaure l'efficacité de la prédiction multi-jetons. Cela permet aux développeurs d'utiliser des stratégies de spéculation agressives pour réduire davantage la latence, rapprochant les performances SSM des limites théoriques du matériel.

Ce que vous pouvez faire

• Lire l'article complet sur ReplaySSM sur arXiv pour comprendre la dérivation mathématique et les métriques d'évaluation détaillées.
• Explorer l'implémentation open source sur GitHub pour expérimenter le mécanisme de mise en cache dans vos propres projets.
• Mettre à niveau vers les dernières versions de TensorRT-LLM ou vLLM pour activer la prise en charge de ReplaySSM pour les modèles hybrides compatibles.
• Effectuer des benchmarks de vos charges de travail SSM actuelles avec des tailles de lot plus élevées pour quantifier les améliorations potentielles de débit dues à la réduction du trafic mémoire.
• Étudier les configurations de décodage prédictif avec des longueurs de brouillon plus importantes, car ReplaySSM supprime les pénalités mémoire précédemment associées au retour arrière.
• Surveiller les erreurs de quantification si vous utilisez des états en basse précision, notant que la séparation point de contrôle-tampon aide à maintenir une précision proche des niveaux FP32.