Accélérer Sonar par 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 léger pour produire des candidats de complétion qui sont vérifiés par le modèle cible plus important. Selon ce schéma, au lieu qu'une exécution du modèle cible coûteux produise un seul jeton, plusieurs sont émis en une seule étape. Nous présentons ici les détails de mise en œuvre de divers types de décodage spéculatif, appliqués chez Perplexity pour réduire la latence inter-jetons sur les modèles Sonar.

Décodage spéculatif

Le décodage spéculatif exploite la structure des langages naturels et la nature autorégressive des transformeurs pour accélérer la génération de jetons. Bien que les modèles plus grands, tels que Llama-70B, détiennent davantage de connaissances que les plus petits, tels que Llama-1B, ils offrent 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 moins coûteux, laissant les problèmes complexes aux plus grands. Le défi réside dans la détermination des complétions les plus pertinentes et de la qualité de la génération du modèle plus petit par rapport à celle du modèle plus grand.

Heureusement, les LLM sont des transformeurs autorégressifs : lorsqu'une séquence de jetons leur est fournie, ils produisent 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é pour le modèle d'émettre ces jetons exacts. 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 traitée par le modèle plus grand afin de déterminer sa pertinence par rapport au 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 bonus, le modèle cible fournit également un jeton subséquent gratuitement : si un modèle préliminaire génère n jetons, jusqu'à n + 1 peuvent être émis en une seule étape.

Agrandir

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

Preremplissage (Prefill) : le modèle cible et le modèle préliminaire doivent tous deux être exécutés sur la séquence d'entrée pour remplir les entrées du cache KV. Bien que certains schémas, tels que Medusa, utilisent des couches denses plus simples pour la prédiction, nous nous concentrons dans cet article sur les modèles préliminaires basés sur les transformeurs qui nécessitent leurs propres caches KV.

Génération préliminaire : le modèle préliminaire itère pour produire un 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, construisant des logits correspondant à chaque jeton préliminaire. La longueur de la plus longue séquence acceptable est déterminée.

Génération cible : puisque la cible a généré des logits, à la position dépareillée ou à la fin de la séquence, les logits correspondent à un jeton subséquent. Ces logits peuvent être échantillonnés pour fournir un jeton robuste provenant de la cible, clôturant la séquence.

Diverses méthodes existent pour mettre en œuvre le décodage spéculatif. Dans cet article, nous nous concentrerons sur les schémas que nous avons utilisés pour accélérer les modèles Sonar en utilisant un modèle 1B interne, ainsi que sur les mécanismes de prédiction que nous développons pour accélérer les modèles à l'échelle de DeepSeek.

Cible-Préliminaire

Le décodage spéculatif peut être réalisé en couplant un petit LLM existant en tant que modèle préliminaire à un modèle cible pour générer des séquences candidates. En production, nous avons accéléré Sonar en utilisant 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é l'entraînement d'un modèle préliminaire à partir de zéro, le petit modèle utilise tout de même une capacité de cache KV significative et introduit un léger surcoût de preremplissage, augmentant le TTFT.

Dans ce schéma, le décodeur ne spécule que sur les lots de décodage uniquement, générant des jetons via un échantillonnage standard pendant le preremplissage ou sur des lots mixtes preremplissage-décodage. Au stade du preremplissage, les logits cibles sont immédiatement échantillonnés pour également préremplir le jeton nouvellement généré dans le cache KV du modèle préliminaire. Le modèle préliminaire n'est pas encore échantillonné, mais les logits qu'il produit sont reportés au stade du décodage.

Agrandir

En décodage, le modèle préliminaire est avancé, é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 sur la base desquels l'échantillonneur identifie la longueur de séquence acceptée. L'acceptation est déterminée en comparant les distributions de probabilité complètes provenant du modèle préliminaire et de la cible. Puisque la cible produit toujours un ensemble de logits suivant la séquence préliminaire acceptée, celui-ci est échantillonné pour produire une sortie supplémentaire. Étant donné que le modèle préliminaire n'a pas encore vu ce jeton accepté, il est réexécuté pour remplir ses entrées de cache KV correspondantes en préparation de la prochaine étape de décodage, reportant à nouveau les logits.

EAGLE

EAGLE est un schéma de décodage spéculatif qui explore plusieurs séquences préliminaires, générées par un parcours arborescent de jetons préliminaires probables. Un arbre fixe (EAGLE) ou dynamiquement formé (EAGLE-2) est exploré à l'aide 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 meilleur score dans une séquence linéaire. Les séquences sont ensuite notées et la plus longue séquence appropriée est sélectionnée pour continuer, en ajoutant également un jeton supplémentaire provenant de la cible.

Agrandir

Afin d'obtenir une prédiction plus précise, un modèle préliminaire EAGLE prédit non seulement sur la base des jetons, mais également en utilisant les caractéristiques cibles (états cachés de la dernière couche) du modèle cible. L'inconvénient d'EAGLE est la nécessité d'entraîner des modèles préliminaires personnalisés et de petite taille qui sont suffisamment précis pour générer des candidats appropriés dans un budget de latence faible. Typiquement, un modèle préliminaire est une couche de transformeur unique identique à une couche de décodeur du modèle original, étroitement couplée à la cible par ses plongements et ses projections lm_head. Comme cela nécessite moins de capacité de cache KV, EAGLE a une empreinte mémoire plus faible.

Pour vérifier les 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 %), annulant une partie de l'accélération réalisable par la spéculation. Nous n'avons pas encore déployé l'exploration arborescente complète en production pour cette raison, préférant nous concentrer sur le cas spécial de la prédiction à jeton unique via des schémas de type MTP présentés dans le rapport technique DeepSeek-V3.

MTP

Ce schéma est similaire au décodage cible-préliminaire, à l'exception des états cachés qui sont utilisés parallèlement aux jetons pour la prédiction. Un peu plus de travail doit être effectué à la fois dans les étapes de preremplissage et de décodage par rapport à la spéculation cible-préliminaire régulière. 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, qui 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 pas vers la gauche par rapport aux vecteurs d'états cachés sortis par la cible. La figure ci-dessous indique les correspondances utilisées pour l'entraînement, ainsi que le décalage pendant l'inférence.

Agrandir

Le flux de décodage est assez similaire au décodage cible-préliminaire, à l'exception des états cachés et des logits qui sont tous deux reportés. Notre implémentation partage tous les logits de traitement et d'échantillonnage associés, spécialisant uniquement les invocations directes du modèle. Lorsque plusieurs jetons sont prédits, le modèle préliminaire utilise des états cachés préliminaires pour la prédiction, remplissant également les entrées du cache KV sur la base de ses propres caractéristiques. À long terme, cela peut dégrader la précision. Par la suite, lors de l'exécution du modèle préliminaire pour remplir l'entrée du cache KV pour la prédiction cible, nous l'exécutons sur la séquence entière en prenant les états cachés cibles plus précis en entrée. Comme ces modèles préliminaires sont petits, le coût ajouté du traitement des jetons supplémentaires est négligeable.

Entraînement des têtes MTP

Afin de bénéficier du MTP, nous avons construit l'infrastructure nécessaire pour entraîner des têtes MTP attachées à nos modèles affinés sur les ensembles de données de Perplexity, s'exécutant sur un nœud avec 8 périphériques H100. En environ une journée, 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 grands, nous nous appuyons sur des têtes MTP construites pendant le 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 cibles avec les logits du prochain jeton et les états cachés de la cible. Comme l'inférence pour les états cachés est coûteuse, nous les pré-calculons en utilisant notre implémentation optimisée pour l'inférence du modèle cible, afin de les utiliser pendant l'entraînement. Cependant, pour valider l'implémentation MTP de l'inférence et s'assurer que les différences numériques dues à la quantification ou aux optimisations n'entravent pas les résultats, nous réutilisons pleinement l'implémentation d'inférence des modèles cible et préliminaire pour l'estimation de la perte et de la précision de validation.

Lors de la montée en charge à partir de l'ensemble de données ShareGPT utilisé dans l'article original vers des échantillons plus grands, nous avons remarqué que l'architecture de tête MTP décrite et mise en œuvre dans l'article EAGLE échouait à entraîner pour les modèles de taille 70B. Contrairement à ShareGPT qui contenait un plus grand nombre de séquences plus courtes, nous entraînons sur un nombre légèrement plus petit de prompts substantiellement plus longs. Étant donné que les têtes EAGLE originales divergeaient légèrement dans leur structure d'un transformeur typique, nous avons réintroduit certaines couches de normalisation RMS qui avaient été supprimées. Nous avons constaté que cela permettait non seulement à l'entraînement de converger, mais cela boostait également la précision des têtes de quelques points de pourcentage.

Agrandir

Les normalisations de couche facilitent non seulement l'entraînement, mais la réintroduction des normes est également mathématiquement intuitive. Les têtes MTP réutilisent les plongements et les projections de logits du modèle cible, car ils peuvent être d'une taille substantielle (environ 2 Go pour Llama 70B). Pendant l'entraînement, ceux-ci sont gelés et l'attente est que la couche MTP apprenne à intégrer les prédictions dans le même espace vectoriel que ce que la couche de projection du modèle original a appris pendant l'entraînement. En supprimant les normes, un seul MLP est censé apprendre la même fonction qu'un MLP suivi d'une norme, ce qui entrave la correspondance entre les états cachés du modèle préliminaire et du modèle cible.

Inférence avec décodage spéculatif

Dans le moteur d'inférence, afin de générer des jetons pour les séquences d'entrée, elles 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 de page KV sont ensuite emballés dans un tampon diffusé à tous les rangs parallèles exécutant le modèle. Enfin, les métadonnées sont copiées dans la mémoire du GPU et le modèle est exécuté pour produire les logits à partir desquels le jeton suivant est échantillonné.

Contrairement à certaines implémentations qui couplent de manière lâche un serveur d'inférence préliminaire et cible via une enveloppe qui orchestre les requêtes entre eux, nos paires cible-préliminaire sont étroitement couplées et progressent dans la génération à l'unisson. La planification des lots et l'allocation des pages KV sont partagées entre les modèles pour toutes les formes de décodage spéculatif : cela unifie la logique qui relie un modèle au serveur d'inférence global, car ils exposent tous la même interface.

Le runtime d'inférence chez Perplexity est structuré autour de FlashInfer, qui détermine les métadonnées à construire afin de configurer et de planifier le noyau d'attention. Étant donné certaines séquences d'entrée formant un lot, pour le preremplissage, le décodage ou la vérification, un travail côté CPU doit être effectué pour allouer des tampons intermédiaires et remplir 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 KV, qui entraînent également des latences devant être masquées 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 CPU-GPU pour le décodage spéculatif est plus complexe. Le défi principal provient du fait que le nombre de jetons acceptés détermine la longueur de séquence pour une exécution subséquente, introduisant un point de synchronisation GPU-à-CPU difficile à éviter. Nous avons expérimenté différents schémas de planification afin de mieux masquer la latence du travail CPU.

Planification Cible-Préliminaire

Bien qu'il soit plus petit qu'un modèle cible, lorsqu'un LLM entier est utilisé comme modèle préliminaire, il introduit toujours une latence considérable sur le GPU, offrant une marge pour masquer les opérations CPU coûteuses. Comme les modèles plus petits ne bénéficient pas du parallélisme tensoriel, il existe une inadéquation entre le nombre de rangs sur lesquels une cible et un modèle préliminaire sont fragmentés. Dans notre implémentation, le modèle préliminaire ne s'exécute que sur le rang leader d'un groupe TP.

Agrandir

Comme indiqué précédemment, une étape de décodage reporte les logits dans l'exécution suivante. Cela nous permet de chevaucher une exécution du modèle préliminaire avec le travail de planification des lots côté CPU. Une fois le lot assemblé, des appels répétés à l'échantillonneur et au modèle préliminaire produisent les jetons préliminaires. En parallèle, le lot pour la vérification est assemblé pour le modèle cible et synchronisé avec les travailleurs parallèles. Les logits cibles sont vérifiés et échantillonnés pour déterminer les longueurs de séquence acceptées. À ce stade, la synchronisation GPU-à-CPU est nécessaire afin de déterminer les longueurs de séquence subséquentes. Puisque le modèle préliminaire ne s'exécute que sur le nœud leader, son lot est configuré séquentiellement et son exécution est lancée pour remplir ses entrées de cache KV avec le jeton supplémentaire que la cible a produit. Les logits produits par cette exécution préliminaire seront utilisés pour échantillonner le premier jeton préliminaire dans l'exécution subséquente. Plus important encore, pendant que le modèle préliminaire s'exécute, le lot suivant peut être planifié.

Planification MTP pour un jeton unique

Bien que le runtime ne fournisse pas encore d'exploration d'arbre préliminaire de style Eagle, nous avons mis en œuvre un cas spécial de ce schéma, 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 transformeur. Ce schéma peut être utilisé pour la prédiction préliminaire en utilisant les poids open-source de DeepSeek R1. Le sous-cas de la prédiction d'un seul jeton est intéressant, car les grandes couches MTP atteignent des taux d'acceptation suffisamment élevés pour justifier leur surcoût.

La planification MTP est quelque peu plus complexe, car le modèle préliminaire est beaucoup plus rapide, masquant moins de latence côté CPU. De plus, le modèle préliminaire est fragmenté aux côtés du modèle cible, nécessitant des transferts de mémoire partagée pour les informations de lot. Une exécution commence par le transfert des informations de lot et l'échantillonnage du premier jeton à partir des logits reportés, de manière similaire au schéma précédent. Ensuite, la cible est exécutée pour valider les jetons, traitant 2 * D jetons, où D est la taille du lot de décodage. Ceci est idéal pour le micro-lotissement dans les modèles Mixture-of-Experts (MoE) sur des interconnexions plus lentes telles qu'InfiniBand, car le lot se divise uniformément en deux moitiés. Les états cachés de la cible sont reportés à l'exécution préliminaire suivante, 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 CPU-à-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 correspondants. Le modèle préliminaire est ensuite réexécuté avec les mêmes informations de lot que la cible, remplissant les entrées du cache KV et construisant les logits et les états cachés pour l'exécution suivante, effectuant un travail redondant sur les jetons qui n'ont pas été acceptés. Dans ces situations, la latence du travail inutilisé est à peine mesurable en raison de la petite taille du modèle préliminaire. En parallèle avec l'exécution préliminaire, les longueurs de séquence sont déterminées sur le CPU et la planification du lot suivant est lancée, sans avoir à attendre la fin du travail GPU.

Le surcoût du travail supplémentaire dans la couche préliminaire n'est pas perceptible dans l'attention, cependant les couches MLP sont plus problématiques. Étant donné que les instructions de multiplication matricielle se complètent jusqu'à une limite de 64 le long de la dimension du nombre de jetons, si le doublement ne nécessite pas significativement plus de blocs, le surcoût est masqué. Pour des séquences préliminaires plus longues, le surcoût est plus coûteux et le schéma utilisé pour les modèles cible-préliminaire réguliers fonctionne mieux.

Références

Inférence rapide à partir de transformeurs via décodage spéculatif

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

EAGLE-2 : Inférence plus rapide des modèles de langage avec des arbres préliminaires dynamiques

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

Medusa : Cadre simple d'accélération de l'inférence LLM avec plusieurs têtes de décodage

Normalisation de couche

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

Rapport technique 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}
}