Accélérer Sonar grâce à la spéculation

Le décodage spéculatif accélère la vitesse de génération des grands modèles de langage (LLM) en utilisant un modèle préliminaire rapide et de petite taille pour produire des candidats de complétion qui sont vérifiés par le grand modèle cible. Dans le cadre de cette méthode, au lieu qu'une exécution du modèle cible (qui est coûteux) produise un seul jeton, plusieurs sont émis en une seule étape. Nous présentons ici les détails d'implémentation de divers types de décodage spéculatif, appliqués chez Perplexity pour réduire la latence entre les jetons sur les modèles Sonar.

Décodage spéculatif

Le décodage spéculatif tire parti de la structure des langues naturelles et de la nature autorégressive des transformateurs pour accélérer la génération de jetons. Même si les modèles plus grands, tels que Llama-70B, possèdent plus de connaissances que les plus petits, tels que Llama-1B, ils affichent des performances similaires sur certaines tâches plus simples. Ce chevauchement suggère que certaines séquences sont mieux générées par les modèles les moins coûteux, confiant les problèmes complexes aux plus grands. La difficulté consiste à déterminer quelles complétions sont les meilleures et si la génération du modèle plus petit est de la même qualité que celle du modèle plus grand.

Heureusement, les LLM sont des transformateurs autorégressifs : lorsqu'on leur fournit une séquence de jetons, ils affichent la distribution de probabilité du jeton suivant. De plus, les logits dérivés des caractéristiques intermédiaires associées aux jetons de la séquence d'entrée indiquent également la probabilité que le modèle émette ces jetons précis. Cette propriété permet la spéculation : si une séquence de jetons est générée par un modèle plus petit à partir d'un préfixe d'entrée, elle peut être exécutée par le modèle plus grand afin de déterminer dans quelle mesure elle s'aligne sur le modèle cible. Chaque préfixe des candidats est évalué avec une probabilité, et le plus long dépassant un seuil d'acceptation est sélectionné. En prime, le modèle cible fournit également un jeton supplémentaire sans frais : si un modèle préliminaire génère n jetons, jusqu'à n + 1 jetons peuvent être émis en une seule étape.

Agrandir

Au moment de l'inférence, le processus d'échantillonnage spéculatif peut être divisé en quatre grandes étapes :

Pré-remplissage : les modèles cible et préliminaire doivent tous deux être exécutés sur la séquence d'entrée pour alimenter les entrées du cache KV. Bien que certaines méthodes, telles que Medusa, utilisent des couches denses plus simples pour la prédiction, cet article se concentre sur les brouillons basés sur des transformateurs qui nécessitent leurs propres caches KV.

Génération préliminaire : le modèle préliminaire effectue des itérations pour produire un certain nombre de jetons fixes. La séquence préliminaire peut être linéaire ou le modèle peut explorer une structure arborescente jusqu'à une profondeur donnée (EAGLE, Medusa). Ici, nous nous concentrons sur les séquences linéaires.

Acceptation : le modèle cible s'exécute sur la séquence préliminaire en construisant des logits correspondant à chaque jeton préliminaire. La longueur de la séquence acceptable la plus longue est alors déterminée.

Génération cible : étant donné que le modèle cible a généré des logits, ceux-ci correspondent à un jeton suivant à la position de discordance ou à la toute fin de la séquence. Ces logits peuvent être échantillonnés pour fournir un jeton robuste issu de la cible, ce qui vient couronner la séquence.

Il existe diverses méthodes pour implémenter le décodage spéculatif. Dans cet article, nous nous intéresserons aux méthodes que nous avons utilisées pour accélérer les modèles Sonar à l'aide d'un modèle 1B interne, ainsi qu'aux mécanismes de prédiction que nous concevons pour accélérer les modèles à l'échelle de DeepSeek.

Cible-Brouillon

Le décodage spéculatif peut être réalisé en associant un petit LLM existant en tant que modèle préliminaire à un modèle cible afin de générer des séquences candidates. En production, nous avons accéléré Sonar à l'aide d'un modèle Llama-1B affiné sur le même ensemble de données que la cible. Bien que cette approche n'ait pas nécessité d'entraîner un brouillon à partir de zéro, le petit modèle utilise tout de même une capacité de cache KV importante et introduit un léger surdébit de pré-remplissage, ce qui augmente le TTFT.

Selon cette méthode, le décodeur ne formule des hypothèses que sur les lots réservés au décodage, en générant des jetons par le biais d'un échantillonnage standard lors du pré-remplissage ou sur des lots mixtes (pré-remplissage et décodage). Durant l'étape de pré-remplissage, les logits de la cible sont immédiatement échantillonnés pour pré-remplir également le jeton nouvellement généré dans le cache KV du brouillon. Le modèle préliminaire n'est pas encore échantillonné, mais les logits qu'il produit sont reportés à l'étape de décodage.

Agrandir

Lors du décodage, le modèle préliminaire progresse en échantillonnant le meilleur jeton à chaque étape. Une fois la longueur préliminaire souhaitée atteinte, les jetons sont traités par le modèle cible pour produire les logits à partir desquels l'échantillonneur identifie la longueur de séquence acceptée. L'acceptation est déterminée en comparant l'ensemble des distributions de probabilité du brouillon et de la cible. Comme la cible génère toujours un ensemble de logits à la suite de la séquence préliminaire acceptée, ce dernier est échantillonné pour produire une sortie supplémentaire. Le modèle préliminaire n'ayant pas encore vu ce jeton accepté, il est réexécuté afin de remplir ses entrées correspondantes dans le cache KV en vue de la prochaine étape de décodage, en reportant à nouveau les logits.

EAGLE

EAGLE est une méthode de décodage spéculatif qui explore de multiples séquences préliminaires, générées par un parcours arboresquent de jetons préliminaires probables. Un arbre fixe (EAGLE) ou à forme dynamique (EAGLE-2) est exploré au moyen d'exécutions consécutives des jetons préliminaires, en considérant les candidats Top-K à chaque nœud au lieu de suivre le jeton ayant le score le plus élevé dans une séquence linéaire. Les séquences reçoivent ensuite un score et la plus longue séquence convenable est sélectionnée pour continuer, tout en ajoutant un jeton supplémentaire provenant de la cible.

Agrandir

Afin d'obtenir une prédiction plus précise, un modèle préliminaire EAGLE effectue ses prédictions non seulement à partir des jetons, mais aussi en utilisant les caractéristiques cibles (les états cachés de la dernière couche) du modèle cible. L'inconvénient d'EAGLE réside dans la nécessité d'entraîner de petits modèles préliminaires personnalisés qui sont suffisamment précis pour générer des candidats adaptés dans le respect d'un budget de latence réduit. En règle générale, un modèle préliminaire est constitué d'une unique couche de transformateur identique à une couche de décodeur du modèle d'origine, laquelle est étroitement liée à la cible en se rattachant à ses représentations vectorielles (embeddings) et à ses projections lm_head. Étant donné que cela nécessite moins de capacité de cache KV, EAGLE présente une empreinte mémoire plus faible.

Pour vérifier des séquences arborescentes dans le modèle cible, des masques d'attention personnalisés doivent être utilisés. Malheureusement, l'utilisation d'un masque d'attention personnalisé pour une séquence entière ralentit considérablement l'attention pour des longueurs d'entrée réalistes (jusqu'à 50 %), ce qui annule une partie de l'accélération obtenue grâce à la spéculation. C'est pourquoi nous n'avons pas encore déployé l'exploration arborescente complète en production, préférant nous concentrer sur le cas particulier de la prédiction à jeton unique par le biais de méthodes de type MTP présentées dans le rapport technique de DeepSeek-V3.

MTP

Cette méthode est similaire au décodage brouillon-cible, à l'exception près que des états cachés sont utilisés en plus des jetons pour la prédiction. Un peu plus de travail doit être effectué à la fois lors des étapes de pré-remplissage et de décodage par rapport à la spéculation classique entre brouillon et cible. Le modèle préliminaire utilise à la fois des jetons et des états cachés : le jeton t_{i+1} est échantillonné à partir des logits L_i correspondant au jeton t_i, lesquels sont à leur tour dérivés des états cachés H_i. Par conséquent, les tampons de jetons d'entrée doivent être décalés d'un cran vers la gauche par rapport aux vecteurs d'états cachés émis par la cible. Le schéma ci-dessous indique les correspondances utilisées pour l'entraînement, ainsi que le décalage effectué lors de l'inférence.

Agrandir

Le déroulement du décodage est assez similaire au décodage brouillon-cible, à l'exception du report à la fois des états cachés et des logits. Notre implémentation partage l'ensemble des opérations d'échantillonnage et de traitement des logits associés, et ne spécialise que les appels de propagation vers l'avant du modèle. Lorsque plusieurs jetons sont prédits, le modèle préliminaire utilise des états cachés préliminaires pour sa prédiction, tout en remplissant les entrées du cache KV sur la base de ses propres caractéristiques. À long terme, cela peut nuire à la précision. Par la suite, lors de l'exécution du modèle préliminaire pour alimenter l'entrée du cache KV dédiée à la prédiction de la cible, nous l'exécutons sur la séquence entière en prenant comme entrées les états cachés de la cible, qui sont plus précis. Comme ces modèles préliminaires sont de petite taille, le coût supplémentaire lié au traitement des jetons additionnels est négligeable.

Entraînement des têtes MTP

Afin de tirer parti du MTP, nous avons mis en place l'infrastructure nécessaire pour entraîner des têtes MTP rattachées à nos modèles affinés sur les ensembles de données de Perplexity, le tout exécuté sur un nœud doté de 8 appareils H100. En une journée environ, nous pouvons construire des têtes pour des modèles allant de Llama-1B à Llama-70B et DeepSeek V2-Lite. Pour les modèles plus volumineux, nous nous appuyons sur des têtes MTP générées au cours du processus d'affinage.

L'objectif de l'entraînement MTP est de faire correspondre les états cachés préliminaires et les logits extrapolés à partir des états cachés de la cible avec les logits du jeton suivant et les états cachés de la cible. L'inférence pour les états cachés étant coûteuse, nous les pré-calculons à l'aide de notre implémentation du modèle cible optimisée pour l'inférence, afin de les utiliser pendant l'entraînement. Toutefois, pour valider l'implémentation MTP de l'inférence et s'assurer que les écarts numériques dus à la quantification ou aux optimisations n'altèrent pas les résultats, nous réutilisons pleinement l'implémentation d'inférence de la cible et du modèle préliminaire pour le calcul de la perte de validation et l'estimation de la précision.

Lors de la transition de l'ensemble de données ShareGPT (utilisé dans l'article d'origine) vers des échantillons plus volumineux, nous avons remarqué que l'architecture de tête MTP décrite et implémentée dans l'article sur EAGLE ne parvenait pas à s'entraîner pour les modèles de la taille 70B. Contrairement à ShareGPT, qui contenait un plus grand nombre de séquences courtes, nous nous entrainons sur un nombre légèrement inférieur de requêtes considérablement plus longues. Les têtes EAGLE d'origine présentant une structure légèrement différente de celle d'un transformateur typique, nous avons réintroduit certaines couches de normalisation RMS qui avaient été supprimées. Nous avons constaté que cela permettait non seulement de faire converger l'entraînement, de surcroît, cela améliorait la précision des têtes de quelques points de pourcentage.

Agrandir

Les normalisations de couches facilitent non seulement l'entraînement, mais leur réintroduction s'avère également intuitive sur le plan mathématique. Les têtes MTP réutilisent les représentations vectorielles (embeddings) et les projections de logits du modèle cible, car celles-ci peuvent être d'une taille importante (environ 2 Go pour Llama 70B). Pendant l'entraînement, ces éléments sont gelés et on s'attend à ce que la couche MTP apprenne à intégrer les prédictions dans le même espace vectoriel que celui qu'a appris la couche de projection du modèle d'origine lors de son entraînement. En supprimant les normalisations, on s'attend à ce qu'un seul perceptron multicouche (MLP) apprenne la même fonction qu'un MLP suivi d'une normalisation, ce qui nuit à l'adéquation entre les états cachés du modèle préliminaire et ceux du modèle cible.

Inférence avec le décodage spéculatif

Dans le moteur d'inférence, pour générer des jetons à partir de séquences d'entrée, celles-ci doivent d'abord être regroupées en lots de taille raisonnable, puis des pages doivent être allouées dans le cache KV pour les jetons suivants. Les jetons d'entrée et les informations sur les pages du cache KV sont ensuite empaquetés dans un tampon diffusé à l'ensemble des rangs parallèles qui font tourner le modèle. Enfin, les métadonnées sont copiées dans la mémoire du GPU et le modèle est exécuté afin de produire les logits à partir desquels le jeton suivant sera échantillonné.

Contrairement à certaines implémentations qui associent de manière lâche un serveur d'inférence préliminaire et cible via une interface de programmation (wrapper) coordonnant les requêtes entre eux, nos paires brouillon-cible sont étroitement couplées et progressent de concert dans la génération. La planification des lots (batch scheduling) et l'allocation des pages du cache KV sont partagées entre les modèles pour toutes les formes de décodage spéculatif : cela unifie la logique qui fait le lien entre un modèle et le serveur d'inférence global, étant donné qu'ils exposent tous la même interface.

Le moteur d'exécution d'inférence de Perplexity repose sur FlashInfer, qui détermine les métadonnées à construire afin de configurer et de planifier le noyau d'attention. Étant donné un ensemble de séquences d'entrée formant un lot pour le pré-remplissage, le décodage ou la vérification, un travail du côté du processeur (CPU) doit être effectué pour allouer des tampons intermédiaires et alimenter certains tampons constants utilisés dans l'attention. Ce travail s'ajoute au coût de la planification des lots et de l'allocation des pages du cache KV, qui engendrent également des temps de latence qu'il convient de masquer afin de maximiser l'utilisation du GPU.

Bien que nous ayons entièrement parallélisé le travail côté CPU et côté GPU pour l'inférence sans spéculation, nous avons constaté que l'équilibre entre CPU et GPU s'avère plus complexe dans le cadre du décodage spéculatif. La principale difficulté tient au fait que le nombre de jetons acceptés détermine la longueur de séquence d'une exécution ultérieure, ce qui introduit un point de synchronisation entre GPU et CPU qu'il est difficile d'éviter. Nous avons testé différentes méthodes de planification afin de masquer au mieux la latence du travail du CPU.

Planification brouillon-cible

Bien qu'il soit plus petit qu'un modèle cible, lorsqu'un LLM complet est utilisé comme modèle préliminaire, il génère tout de même une latence considérable sur le GPU, ce qui procure une certaine marge de manœuvre pour masquer les opérations coûteuses du CPU. Les petits modèles ne tirant aucun avantage du parallélisme tensoriel, il existe un décalage entre le nombre de rangs sur lesquels une cible et un brouillon sont répartis (sharded). Dans notre implémentation, le modèle préliminaire ne s'exécute que sur le rang principal (leader rank) d'un groupe de parallélisme tensoriel.

Agrandir

Comme indiqué précédemment, une étape de décodage reporte les logits dans l'exécution suivante. Cela nous permet de superposer une exécution du modèle préliminaire avec le travail de planification des lots du côté du CPU. Une fois le lot constitué, des appels répétés à l'échantillonneur et au modèle préliminaire génèrent les jetons préliminaires. En parallèle, le lot destiné à la vérification est constitué pour le modèle cible et synchronisé avec les processus parallèles. Les logits de la cible sont vérifiés et échantillonnés afin de déterminer les longueurs de séquence acceptées. À ce stade, une synchronisation entre le GPU et le CPU est nécessaire pour établir les longueurs de séquence ultérieures. Le modèle préliminaire ne s'exécutant que sur le nœud principal, son lot est configuré séquentiellement et son exécution est lancée pour alimenter ses entrées du cache KV avec le jeton supplémentaire produit par la cible. Les logits produits par cette exécution du brouillon serviront à échantillonner le premier jeton préliminaire lors de l'exécution suivante. Plus important encore, pendant que le modèle préliminaire est en cours d'exécution, le lot suivant peut être planifié.

Planification MTP pour un jeton unique

Bien que le moteur d'exécution ne propose pas encore d'exploration d'arbres préliminaires de type Eagle, nous avons implémenté un cas particulier de cette méthode en considérant une séquence linéaire de jetons préliminaires produits par un modèle de la taille d'une seule couche de décodeur de transformateur. Cette méthode peut être employée pour la prédiction préliminaire à l'aide des poids open source de DeepSeek R1. Le sous-cas de la prédiction d'un jeton unique présente un intérêt certain, dans la mesure où les grandes couches MTP atteignent des taux d'acceptation suffisamment élevés pour justifier leur surdébit.

La planification MTP est un peu plus complexe, car le modèle préliminaire est beaucoup plus rapide, ce qui masque une moindre latence du côté du CPU. De plus, le modèle préliminaire est réparti (sharded) parallèlement au modèle cible, ce qui nécessite des transferts en mémoire partagée pour les informations de lots. Une exécution commence par le transfert des informations de lots et l'échantillonnage du premier jeton à partir des logits reportés, de la même manière que pour la méthode précédente. Ensuite, la cible est exécutée pour valider les jetons, en traitant 2 * D jetons, où D représente la taille du lot de décodage. Cette approche est idéale pour le traitement par micro-lots dans les modèles de type mélange d'experts (MoE) sur des interconnexions plus lentes telles qu'InfiniBand, car le lot se divise de manière égale en deux moitiés. Les états cachés de la cible sont reportés à l'exécution suivante du brouillon, tandis que les logits sont transmis à l'échantillonneur pour vérification.

Agrandir

En effectuant une quantité limitée de travail supplémentaire sur le GPU, nous évitons la synchronisation entre le CPU et le GPU après l'acceptation de la séquence préliminaire. Une fois les jetons d'entrée des cibles décalés, un noyau insère les jetons cibles suivants dans leurs emplacements respectifs. Le modèle préliminaire est ensuite réexécuté avec les mêmes informations de lot que la cible, ce qui permet de remplir les entrées du cache KV et de générer les logits et les états cachés pour l'exécution suivante, tout en effectuant un travail redondant sur les jetons qui n'ont pas été acceptés. Dans ces situations, la latence liée au travail non utilisé est à peine mesurable en raison de la petite taille du modèle préliminaire. En parallèle de l'exécution du brouillon, les longueurs de séquence sont déterminées sur le CPU et la planification du lot suivant est lancée, sans qu'il soit nécessaire d'attendre la fin du travail sur le GPU.

Le surdébit de travail supplémentaire dans la couche préliminaire n'est pas perceptible au niveau de l'attention, mais les couches MLP s'avèrent plus problématiques. Les instructions de multiplication matricielle s'alignant sur une limite de 64 le long de la dimension correspondant au nombre de jetons, si le doublement ne nécessite pas un nombre nettement plus élevé de blocs, le surdébit est masqué. Pour les séquences préliminaires plus longues, le surdébit est plus coûteux et la méthode employée pour les modèles classiques brouillon-cible s'avère plus efficace.

Références

Inférence rapide à partir de transformateurs grâce au décodage spéculatif

EAGLE : l'échantillonnage spéculatif nécessite de repenser l'incertitude des caractéristiques

EAGLE-2 : inférence plus rapide de modèles de langage grâce à des arbres préliminaires dynamiques

EAGLE-3 : mise à l'échelle de l'accélération d'inférence des grands modèles de langage via un test au moment de l'entraînement

Medusa : framework d'accélération d'inférence de LLM simple doté de têtes de décodage multiples

Normalisation de couche

FlashInfer : moteur d'attention efficace et personnalisable pour le service d'inférence de LLM

Rapport technique de DeepSeek-V3

Citation

@online{perplexity2025AcceleratingSonarThroughSpeculation,
  author = {{Perplexity} and Nandor Licker and Erin McKnight},
  title = {{Accelerating Sonar through speculation}},
  date = {2025-06-10},
  url = {https://perplexity.coolkeypoint.com/hub/blog/accelerating-sonar-through-speculation}
}