ReplaySSM réduit le trafic d’état récurrent au décodage des modèles de langage

Lorsque les couches à espace d’état d’un modèle réécrivent un grand état récurrent après chaque jeton, peut-on réduire ces écritures sans modifier les sorties ? La question compte parce que ces mises à jour de l’état peuvent représenter une part importante de la latence de décodage, même dans des modèles hybrides où le coût de l’attention augmente avec la longueur du contexte. L’article évalue une autre façon de représenter et de mettre à jour cet état.

ReplaySSM stocke les entrées récentes à côté d’un point de contrôle et calcule les sorties entre ses mises à jour périodiques afin de réduire les transferts de mémoire et de faciliter le retour arrière.

L’article examine si les couches à espace d’état doivent réécrire leur état évolutif complet après chaque jeton, même lorsque chaque mise à jour demande peu de calcul.

ReplaySSM stocke les entrées récentes à côté d’un point de contrôle et le met à jour périodiquement, dans le but de réduire ces écritures tout en conservant les mêmes sorties.

Sur quatre modèles hybrides avec une taille de lot de 256, une étape de décodage du modèle entier était 1,20 à 1,47 fois plus rapide qu’avec le décodage classique vLLM réglé.

Ces résultats appuient l’intérêt d’une autre gestion de l’état dans ces tests ; des expériences distinctes montrent aussi pourquoi les lectures en basse précision ne devraient pas alimenter les mises à jour récurrentes.

Ces résultats proviennent de quatre modèles hybrides testés chacun sur un seul GPU, sans parallélisme tensoriel ; les expériences en basse précision ont utilisé une quantification émulée. Les mesures n’établissent pas que les performances ou la précision seront les mêmes dans d’autres conditions.