ReplaySSM : Repenser le stockage d’état pour une décodage plus rapide des modèles hybrides

ReplaySSM : Repenser le stockage d’état pour une décodage plus rapide de modèles hybrides

Les modèles linguistiques hybrides — des architectures qui mêlent des couches d’attention à des couches d’espace d’état ou récurrentes linéaires — promettaient le meilleur des deux mondes : un rappel exact grâce à l’attention et un comportement de flux à mémoire constante propre aux structures récurrentes. En pratique, de nombreuses équipes constatent que cette promesse tient pendant la phase de préremplissage, mais s’érode progressivement lors du décodage. Le débit par token de sortie diminue, il devient difficile d’analyser l’utilisation de la mémoire, et des fonctionnalités habituelles pour les transformateurs purs — recherche par faisceau, décodage spéculatif, plusieurs échantillons par requête — deviennent complexes.

L’article du blog de NVIDIA intitulé « ReplaySSM : Repenser le stockage d’état pour une décodage plus rapide de modèles hybrides » (daté de 12 octobre 2026) examine cette tension et propose de considérer le stockage d’état comme une question de conception première plutôt que comme un détail d’implémentation. Cet article explique l’espace des problèmes, présente une configuration pratique pour expérimenter avec le décodage hybride sur son propre matériel, et indique ce qu’il convient de mesurer. Les explications conceptuelles ici sont interprétatives ; les détails de la méthode proposée doivent être tirés de l’article original.

Pourquoi les modèles hybrides sont attrayants — et difficiles à décoder

Une couche transformateur conserve un cache clé/valeur qui augmente linéairement avec la longueur de la séquence. Une couche SSM conserve un état récurrent de taille fixe qui est réécrit à chaque token. En les superposant, on obtient un modèle capable d’appliquer l’attention précisément là où c’est important et de traiter les autres parties de manière efficace.

Le problème, c’est que ces deux systèmes de mémoire n’ont presque rien en commun :

• Le cache KV est uniquement en écriture, par token, et adressé par position. On peut le découper, le paginer, partager des préfixes et supprimer un suffixe sans toucher au reste.
• L’état SSM est un résumé en cours d’exécution. Il ne possède pas d’index de position. Le token t génère un nouvel état qui efface l’ancien, et cet état est une fonction compressée de tout ce qui l’a précédé.

Cette asymétrie est à l’origine du problème de décodage. Un cache KV est une structure de données que l’on peut manipuler. Un état récurrent est une valeur qu’il faut soit conserver, soit rec calculer.

Pourquoi le décodage, et non le préremplissage, est la partie difficile

Le pré-remplissage traite toute la requête en une seule passe. Chaque token est disponible en même temps, la GPU est saturée par les opérations matmul, et le coût de l’écriture de l’état est amorti sur une longue séquence. La décodification est l’inverse : un token par étape, de petites opérations matmul, et un noyau limité en mémoire dont le temps d’exécution est dominé par la lecture et la mise à jour de l’état.

Pour un modèle pur d’attention, les performances de décodification dépendent principalement de la taille du cache KV et de l’efficacité avec laquelle on peut y accéder. La quantification, le paginage et le partage de préfixes affectent tous ces aspects. Pour un modèle hybride, on hérite de toutes ces contraintes plus une deuxième famille d’états structuralement différente qui ne peut pas être paginée ou partagée de la même manière.

Le goulot d’étranglement lié au stockage de l’état dans les charges de travail ramifiées

Le problème devient aigu chaque fois qu’un préfixe génère plusieurs continuations :

• La recherche par faisceau conserve k hypothèses actives, chacune avec son propre état.
• La décodification spéculative propose plusieurs tokens provisoires, les vérifie en parallèle et doit annuler ceux qui sont rejetés.
• L’échantillonnage du meilleur résultat génère n complétions indépendantes à partir d’une requête partagée.
• Le raisonnement structuré en arbre explore des branches qui se divisent et se rejoignent.

Pour les couches d’attention, tous ces cas sont gérables. Annuler un draft rejeté consiste simplement à tronquer le cache KV. Partager un préfixe de requête signifie partager ses clés et valeurs mémorisées. Pour les couches SSM, aucune de ces opérations n’est gratuite. Un point de branchement implique que l’état doit exister indépendamment pour chaque branche, et une annulation exige que l’état soit restauré à une valeur antérieure qui n’existe plus sauf si on l’a enregistrée.

Il n’y a que trois solutions, chacune ayant un coût :

1. Prendre des snapshots de manière agressive. Conserver une copie de l’état récurrent à chaque point où il pourrait être nécessaire de revenir. La mémoire augmente en fonction du facteur de ramification multiplié par la fréquence des points de contrôle.
2. Récomputer à partir d’un point de contrôle. Stocker l’état rarement, et relancer les couches à partir du dernier point de contrôle lorsque l’on a besoin d’un état intermédiaire. La mémoire est faible ; la latence est élevée.
3. Restructurer le stockage lui-même. Changer ce qui est stocké, quand et sous quelle forme, afin que la reconstruction soit suffisamment peu coûteuse pour être la méthode par défaut.

Ce que implique une conception axée sur la rejouabilité

Le nom de cette approche fait référence à la troisième option. « Rejouabilité » indique que la reconstruction n’est plus une solution d’urgence mais le mode normal, et que la couche de stockage est conçue pour rendre la rejouabilité peu coûteuse.

Interprété de cette manière, la question de conception passe de « combien d’état devons-nous conserver ? » à « quel est le minimum que nous devons conserver afin que tout état nécessaire puisse être reconstruit rapidement ? » Cette formulation a plusieurs conséquences méritant réflexion :

• Le stockage devient un paramètre ajustable plutôt qu’un coût fixe. Vous pouvez échanger de la mémoire contre des recalculs sur une échelle continue, et choisir un point en fonction de la charge de travail.
• La boucle de décodage acquiert une nouvelle primitive : restaurer l’état à la position p. Que cela soit réalisé en chargeant un snapshot ou en rejouant un court suffixe modifie la latence, mais pas la sémantique.
• Les systèmes de service peuvent allouer des budgets d’état de la même manière qu’ils allouent des budgets de cache KV — en tant que ressource par requête avec des politiques d’éviction, plutôt qu’en allocation inévitable par séquence.
• La correction devient plus facile à démontrer, car un état rejoué doit être numériquement identique (ou proche) de l’état stocké. Il s’agit d’une propriété testable, contrairement aux compromis flous liés à la quantification.

La méthode publiée correspond-elle en détail à cette description ? La source y répond directement ; le raisonnement ci-dessus est un cadre pour l’interpréter, et non un substitut.

• Une GPU compatible CUDA disposant de suffisamment de mémoire pour les poids ainsi que pour la plus grande taille d’état que vous souhaitez autoriser. Commencez avec un modèle qui tient confortablement, puis réduisez les ressources.
• Un interpréteur Python 3.x à jour pris en charge par la version du framework que vous avez choisie. Mélanger un Python de pointe avec des paquets CUDA fixes est la cause la plus fréquente d’échec de configuration.
• PyTorch conçu pour votre pilote CUDA. Un désaccord de version entre le pilote et le paquet génère des erreurs en temps de exécution confuses plutôt que des échecs clairs.
• Un framework qui comprend les architectures hybrides. Les versions récentes transformers prennent en charge des objets de cache hybrides ; des moteurs de service tels que vLLM et TensorRT-LLM gèrent les modèles hybrides avec différents niveaux de maturité. Vérifiez la compatibilité avec votre modèle spécifique avant de réaliser des tests de performance.
• Assez de RAM et d’espace disque sur l’hôte si vous prévoyez de créer des points de contrôle ou de transférer l’état. Un tenseur d’état qui est petit sur la GPU peut devenir volumineux lorsqu’il est multiplié à travers des séquences et des points de contrôle.

Les exigences de version exactes pour l’implémentation de ReplaySSM, le cas échéant, appartiennent à la publication originale et à ses artefacts associés.

Installation étape par étape

Les commandes ci-dessous créent un environnement générique pour l’inférence de modèles hybrides. Elles ne installent pas de package ReplaySSM spécifique — consultez la source pour obtenir le chemin officiel.

Tout d’abord, vérifiez que la GPU est visible et que le pilote fonctionne correctement.

nvidia-smi

Créez un environnement isolé afin que les versions du framework ne entrèrent pas en conflit avec d’autres éléments sur la machine.

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip wheel

Installez PyTorch. Sélectionnez l’URL d’index correspondant à votre runtime CUDA ; le paquet PyPI par défaut peut cibler une version CUDA différente de celle de votre pilote.

pip install torch --index-url https://download.pytorch.org/whl/cu124

Vérifiez que PyTorch détecte la GPU avant d’installer quoi que ce soit d’autre, car chaque échec ultérieur est plus facile à diagnostiquer à partir d’une base fonctionnelle connue.

python -c "import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0))"

Installez les bibliothèques de modélisation et de service que vous comptez comparer.

pip install transformers accelerate
pip install vllm

Si votre modèle hybride utilise des couches de type Mamba, des bibliothèques de noyaux qui accélèrent les couches linéaires récurrentes sont couramment utilisées en parallèle du framework principal. Installez-les uniquement si le modèle que vous exécutez les indique comme dépendance.

pip install flash-linear-attention causal-conv1d

Enfin, enregistrez les versions exactes dans un lockfile afin de pouvoir reproduire plus tard un résultat de benchmark.

pip freeze > requirements-lock.txt

Configuration

La configuration pour le décodage hybride se divise en deux catégories : les paramètres déjà exposés par votre framework, et tous les paramètres ajoutés par l’implémentation ReplaySSM. La première catégorie est utilisable dès aujourd’hui.

Définissez l’allocateur afin de réduire la fragmentation lors de l’allouement et du libération des tenseurs d’état de tailles variables au cours des différentes requêtes.

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Fixez l’exécution à une seule GPU afin que les mesures ne soient pas influencées par les décisions de placement.

export CUDA_VISIBLE_DEVICES=0

Lors du service, activez le cache de préfixes et le préremplissage par blocs si votre moteur les prend en charge. Ces deux fonctionnalités modifient la manière dont l’état est partagé et planifié, et interagissent avec toute stratégie de stockage d’état.

vllm serve YOUR_HYBRID_MODEL_ID \
  --enable-prefix-caching \
  --enable-chunked-prefill \
  --max-model-len 8192 \
  --gpu-memory-utilization 0.90

Pour la deuxième catégorie, consultez le code source. Tout bloc de configuration relatif à ReplaySSM doit être lu dans la documentation officielle plutôt que déduit, car les noms des champs et leurs valeurs par défaut déterminent si vous mesurez bien la caractéristique en question.

Exemples d’utilisation

Exemple 1 : Débit de décodage de référence

Établissez une référence avant de modifier quoi que ce soit. Ce script décode un nombre fixe de tokens et indique le nombre de tokens par seconde, ce qui est la métrique la plus sensible aux coûts liés au traitement de l’état.

import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "YOUR_HYBRID_MODEL_ID"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, torch_dtype=torch.bfloat16, device_map="cuda"
).eval()

prompt = "Explain why recurrent state is expensive to branch."
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

steps = 128
with torch.no_grad():
    out = model(**inputs, use_cache=True)
    next_id = out.logits[:, -1].argmax(-1, keepdim=True)

    torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(steps):
        out = model(
            input_ids=next_id,
            past_key_values=out.past_key_values,
            use_cache=True,
        )
        next_id = out.logits[:, -1].argmax(-1, keepdim=True)
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - start

print(f"{steps / elapsed:.1f} tokens/sec")

Exemple 2 : Mesure de l’empreinte mémoire de l’état

La mémoire maximale utilisée vous indique combien d’espace reste disponible pour des lots plus importants ou davantage de branches. Réinitialisez les statistiques avant la région à mesurer afin de ne capturer que ce que le cycle de décodage alloue.

import torch

torch.cuda.reset_peak_memory_stats()
# ... run the decode loop from Example 1 ...
peak_gib = torch.cuda.max_memory_allocated() / 2**30
print(f"peak allocated: {peak_gib:.2f} GiB")

Soustrayez les poids de cette valeur pour estimer l’empreinte mémoire de l’état, puis divisez par le nombre de séquences simultanées pour obtenir une valeur par séquence. Cette valeur détermine la taille maximale de lot possible dans le cadre d’une stratégie de stockage d’état.

Exemple 3 : Décodage spéculatif et retrait

C’est le volume de travail où la politique de stockage est la plus importante. La structure ci-dessous est schématique ; adaptez les appels à la mémoire tampon à l’API de votre framework.

# 1. Save a reference to the state after the accepted prefix.
state_checkpoint = cache.snapshot()

# 2. Propose k draft tokens and run them through the model in one pass.
draft_ids = drafter.generate(prefix_ids, k=4)
candidate_logits = target_model.forward_with_cache(draft_ids, cache)

# 3. Accept the longest prefix that matches the target distribution.
accepted = longest_matching_prefix(candidate_logits, draft_ids)

# 4. If nothing was accepted, restore rather than re-run the whole prefix.
if accepted == 0:
    cache.restore(state_checkpoint)
else:
    cache.truncate_to(accepted)

Les appels snapshot et restore marquent le point où les deux stratégies divergent. Une implémentation basée sur des captures d’écran entraîne des coûts en mémoire à l’étape 1. ; une implémentation basée sur la rejouissance entraîne des coûts de recomputation à l’étape 4.. Le choix approprié dépend du taux d’acceptation : avec un taux élevé, les retraits sont rares et la rejouissance est avantageuse ; avec un taux faible, les captures d’écran peuvent se rentabiliser.

Ce qu’il faut mesurer avant et après

Un changement de stockage d’état ne vaut la peine d’être adopté que s’il améliore une métrique qui vous importe. Suivez-les de manière cohérente :

• Temps par token de sortie, et non seulement le délai total de la requête. Le bruit lié au préremplissage cache les régressions de décodage.
• Mémoire maximale du dispositif pour une taille de lot fixe, ce qui détermine si vous pouvez augmenter la taille du lot.
• Concurrence maximale soutenable, mesurée en augmentant progressivement la taille du lot jusqu’à ce qu’une erreur de mémoire ou un objectif de délai soit atteint.
• Taux d’acceptation, si le décodage spéculatif est utilisé. Il contrôle directement la fréquence à laquelle la rejouissance est déclenchée.
• Coût des retraits, mesuré comme le délai supplémentaire causé par un brouillon rejeté. C’est ce chiffre que les conceptions basées sur la rejouissance cherchent à réduire.

Exécutez chaque configuration au moins trois fois, en excluant une passe d’échauffement. Les mesures de décodage hybride sont sujettes au bruit, et les comparaisons sur une seule exécution aboutissent fréquemment à des conclusions trompeuses.

Limites et questions ouvertes

Quelques mises en garde honnêtes. L’avantage de toute optimisation du stockage d’état dépend du rapport entre les couches d’attention et les couches récurrentes dans votre modèle ; un modèle légèrement hybride peut ne montrer que peu de changement, tandis qu’un modèle fortement récurrent peut en voir beaucoup. La rejouissance entraîne des coûts de recomputation, donc elle n’est avantageuse que lorsque la mémoire ou le délai de retrait est la contrainte limitante, et non les FLOPs bruts. De plus, la courbe d’équilibre est spécifique au volume de travail : l’intervalle optimal pour les points de contrôle dans le cas du meilleur échantillonnage n’est probablement pas optimal pour la recherche par faisceau.

Où la méthode publiée se situe sur cette courbe, et quels sont les charges de travail qu’elle vise, ce sont des questions auxquelles l’article original répond directement. Considérez cette présentation comme un cadre d’aide à la lecture, et non comme un substitut.

Conclusion

Les modèles hybrides ne sont pas plus difficiles à exécuter que les transformateurs ; ils sont plus difficiles à planifier. La raison en est qu’ils comportent deux types de mémoire incompatibles : une mémoire cache KV adressable et un état récurrent non adressable. Tout ce qui rend l’exécution moderne efficace — pagination, partage de préfixes, rollback spéculatif — a été conçu autour du premier type.

La contribution de ReplaySSM, telle que définie par son titre, consiste à cesser de considérer le second type de mémoire comme un coût inévitable et à commencer à le traiter comme une variable de conception. Si l’état peut être reconstitué à un coût suffisamment faible, le stockage devient un paramètre ajustable plutôt qu’un obstacle, et les charges de travail à branchement complexes que les modèles hybrides gèrent mal deviennent à nouveau ordinaires.

La voie pratique est simple : créer un environnement propre, mesurer le débit de décodage et la mémoire maximale utilisée avant de modifier quoi que ce soit, isoler le coût du rollback dans une boucle de décodage spéculative, puis décider seulement où, sur la courbe mémoire versus recomputation, votre charge de travail souhaite se situer.