Article
Accélérer Sonar par la Spéculation
Le décodage spéculatif accélère la vitesse de génération des Modèles de Langage de Grande Ampleur (LLMs) en utilisant un modèle de brouillon rapide et petit pour produire des candidats de complétion vérifiés par le modèle cible plus grand.
Selon ce schéma, au lieu qu'une exécution coûteuse du modèle cible 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 langues naturelles et la nature auto-régressive des transformeurs pour accélérer la génération de jetons. Même si les modèles plus grands, comme le Llama-70B, possèdent plus de connaissances que les plus petits, tels que le Llama-1B, sur certaines tâches plus simples, leur performance est similaire. Ce chevauchement suggère que certaines séquences sont mieux générées par des modèles moins coûteux, laissant les problèmes complexes aux plus grands. Le défi consiste à déterminer quelles complétions sont meilleures et si la génération du modèle plus petit est de la même qualité que celle du plus grand.
Heureusement, les LLMs sont des transformeurs auto-régressifs : lorsqu'on leur donne une séquence de jetons, 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 dans la séquence d'entrée indiquent également dans quelle mesure le modèle est susceptible 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 plus petit modèle à partir d'un préfixe d'entrée, elle peut être exécutée sur le plus grand pour déterminer dans quelle mesure elle s'aligne avec le modèle cible. Chaque préfixe des candidats est noté avec une probabilité et le plus long au-dessus d'un seuil d'acceptation est choisi. En prime, le modèle cible fournit également un jeton supplémentaire gratuitement : si un modèle de brouillon génère n jetons, jusqu'à n + 1 peuvent être émis en une étape.
Au moment de l'inférence, le processus d'échantillonnage spéculatif peut être divisé en environ 4 étapes :
Préremplissage : les modèles cible et brouillon doivent être exécutés sur la séquence d'entrée pour remplir les entrées de cache KV. Bien que certains schémas, tels que Medusa, utilisent des couches denses plus simples pour la prédiction, nous nous concentrons ici sur les brouillons basés sur des transformeurs qui nécessitent leurs propres caches KV.
Génération de brouillon : le modèle de brouillon itère pour produire un certain nombre de jetons fixes. La séquence de brouillon peut être linéaire ou le modèle peut explorer une structure arborescente jusqu'à une certaine profondeur (EAGLE, Medusa). Nous nous concentrons ici sur les séquences linéaires.
Acceptation : le modèle cible fonctionne sur la séquence de brouillon, générant des logits correspondant à chaque jeton de brouillon. La longueur de la plus longue séquence acceptable est déterminée.
Génération cible : puisque le cible a généré des logits, à la position inadaptée ou à la fin de la queue de la séquence, les logits correspondent à un jeton suivant. Ces logits peuvent être échantillonnés pour fournir un jeton robuste à partir de la cible, complétant ainsi la séquence.
Il existe plusieurs méthodes 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 interne de taille 1B, 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-Brouillon
Le décodage spéculatif peut être réalisé en associant un LLM existant de petite taille comme modèle de brouillon à 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 adapté finement au même ensemble de données que le modèle cible. Bien que cette approche n'ait pas nécessité l'entraînement d'un brouillon à partir de zéro, le petit modèle utilise encore une capacité de cache KV significative et introduit un léger surcoût de préremplissage, augmentant le TTFT.
Selon ce schéma, le décodeur ne spécule que sur des lots decode-only, générant des jetons grâce à un échantillonnage standard pendant le préremplissage ou sur des lots préremplissage-décode mixtes. Dans la phase de préremplissage, les logits cibles sont immédiatement échantillonnés pour également préremplir le nouveau jeton généré dans le cache KV du brouillon. Le brouillon n'est pas encore échantillonné, mais les logits qu'il produit sont reportés à l'étape de décodage.
En décodage, le modèle de brouillon est avancé, échantillonnant le jeton supérieur à chaque étape. Une fois la longueur de brouillon souhaitée atteinte, les jetons sont transmis au 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 du brouillon et de la cible. Comme la cible génère toujours un ensemble de logits suivant la séquence de brouillon acceptée, cela est échantillonné pour produire une sortie supplémentaire. Étant donné que le modèle de brouillon 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 pour la prochaine étape de décodage, reportant les logits à nouveau.
EAGLE
EAGLE est un schéma de décodage spéculatif qui explore plusieurs séquences de brouillon, générées par une traversée arborescente des jetons de brouillon probables. Un arbre fixe (EAGLE) ou à forme dynamique (EAGLE-2) est exploré à l'aide d'exécutions consécutives des jetons de brouillon, en considérant les candidats Top-K à chaque nœud au lieu de suivre le jeton le plus élevé dans une séquence linéaire. Les séquences sont ensuite notées et la plus longue adéquate est sélectionnée pour continuer, ajoutant également un jeton supplémentaire à partir de la cible.
Pour obtenir une prédiction plus précise, un modèle de brouillon EAGLE prédit non seulement sur la base de jetons, mais utilise également 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é de former des modèles de brouillon petits et spécifiques suffisamment précis pour générer des candidats adéquats avec un budget de faible latence. Typiquement, un modèle de brouillon est une seule couche de transformeur identique à une couche de décodeur du modèle original, qui est étroitement couplée à la cible en étant liée à ses projections d'embeddings et de 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, nous concentrant plutôt sur le cas particulier 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 brouillon-cible, à la différence que les états cachés sont utilisés avec les jetons pour la prédiction. Un peu plus de travail doit être fait à la fois dans les étapes de préremplissage et de décodage par rapport à la spéculation régulière brouillon-cible. Le modèle de brouillon 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 à leur tour sont 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 produits par la cible. La figure ci-dessous marque les correspondances utilisées pour l'entraînement, ainsi que le décalage pendant l'inférence.
Le flux de décodage est assez semblable au décodage brouillon-cible, à l'exception que les états cachés et les logits sont reportés. Notre implémentation partage tous les logits de traitement et d'échantillonnage associés, en spécialisant seulement les appels de modèle direct. Lorsqu'il y a plusieurs prédictions de jetons, le modèle de brouillon utilise les états cachés de brouillon pour la prédiction, remplissant également les entrées de cache KV sur la base de ses propres fonctionnalités. A long terme, cela peut dégrader la précision. Par la suite, lorsque le modèle de brouillon est exécuté pour remplir l'entrée de cache KV pour la prédiction cible, nous l'exécutons sur l'ensemble de la séquence en prenant les états cachés cibles plus précis comme entrées. Comme ces modèles de brouillon sont petits, le coût supplémentaire de traitement des jetons additionnels est négligeable.
Formation des têtes MTP
Pour tirer parti de MTP, nous avons construit l'infrastructure nécessaire pour former des têtes MTP attachées à nos modèles adaptés finement sur les ensembles de données de Perplexity, fonctionnant sur une seule unité avec 8 dispositifs 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 comptons sur les têtes MTP construites pendant le processus d'adaptation fine.
L'objectif de la formation MTP est d'aligner les états cachés de brouillon et les logits extrapolés à partir des états cachés cibles avec les logits et états cachés du jeton suivant de la cible. Étant donné que l'inférence des états cachés est coûteuse, nous les pré-calculons en utilisant notre mise en œuvre optimisée pour l'inférence du modèle cible, à utiliser pendant l'entraînement. Cependant, pour valider l'implémentation d'inférence MTP et s'assurer que les différences numériques dues à la quantification ou aux optimisations ne gênent pas les résultats, pour la validation des pertes et l'estimation de la précision, nous réutilisons intégralement l'implémentation d'inférence de la cible et des modèles de brouillon.
En passant de l'ensemble de données ShareGPT utilisé dans le document original à des échantillons plus grands, nous avons remarqué que l'architecture de la tête MTP décrite et implémentée dans le document EAGLE échouait à s'entraîner pour des modèles de taille 70B. Contrairement à ShareGPT qui contenait un plus grand nombre de séquences plus courtes, nous nous formons sur un nombre légèrement inférieur de prompts substantiellement plus longs. Comme les têtes originales EAGLE divergeaient légèrement en 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 augmentait aussi la précision des têtes de quelques points de pourcentage.
Non seulement les normes de couche facilitent-elles l'entraînement, mais leur réintroduction est également mathématiquement intuitive. Les têtes MTP réutilisent les embeddings et les projections de logits du modèle cible, car ils peuvent être substantiels en taille (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 celui que la couche de projection du modèle original a appris pendant l'entraînement. En supprimant les normes, on s'attend à ce qu'un seul MLP apprenne la même fonction qu'un MLP suivi d'une norme, ce qui nuit à l'alignement entre les états cachés des modèles de brouillon et cibles.
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 prochains jetons. Les jetons d'entrée et l'information 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 en mémoire GPU et le modèle est exécuté pour produire les logits à partir desquels le prochain jeton est échantillonné.
Contrairement à certaines implémentations qui couplent librement un serveur d'inférence brouillon et cible via un wrapper qui orchestre les requêtes entre eux, nos paires brouillon-cible sont étroitement couplées et avancent en unisson lors de la génération. La planification des lots et l'allocation de 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 pour configurer et planifier le noyau d'attention. Étant donné certaines séquences d'entrée formant un lot, pour le préremplissage, le décodage ou la vérification, un travail côté CPU doit être réalisé 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 à l'allocation des pages KV, qui génèrent également des latences qui doivent être cachées pour maximiser l'utilisation du GPU.
Bien que nous ayons entièrement parallélisé le travail côté CPU et GPU pour l'inférence sans spéculation, nous avons constaté que l'équilibre CPU-GPU pour le décodage spéculatif est plus compliqué. Le principal défi provient du fait que le nombre de jetons acceptés détermine la longueur de séquence pour une exécution ultérieure, introduisant un point de synchronisation GPU-CPU difficile à éviter. Nous avons expérimenté différents schémas de planification afin de mieux cacher la latence du travail CPU.
Calendrier Brouillon-Cible
Malgré être plus petit qu'un modèle cible, lorsqu'un LLM entier est utilisé comme brouillon, cela introduit encore une latence considérable sur le GPU, fournissant une marge de manœuvre pour cacher des opérations CPU coûteuses. Étant donné que les modèles plus petits ne bénéficient pas du parallélisme des tenseurs, il y a un décalage entre le nombre de rangs qu'un cible et un brouillon sont partagés. Dans notre implémentation, le modèle de brouillon s'exécute uniquement sur le rang leader d'un groupe TP.
Comme indiqué précédemment, une étape de décodage reporte les logits au cours de la prochaine exécution. Cela nous permet de chevaucher une exécution du modèle de brouillon avec le travail de planification de lots côté CPU. Après que le lot soit assemblé, des appels répétés à l'échantillonneur et au brouillon produisent les jetons de brouillon. En parallèle, le lot pour la vérification est formé 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, une synchronisation GPU-CPU est nécessaire pour déterminer les longueurs de séquence ultérieures. Étant donné que le modèle de brouillon est uniquement exécuté 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 le cible a produit. Les logits produits par cette exécution de brouillon dans l'exécution actuelle seront utilisés pour échantillonner le premier jeton de brouillon à l'exécution suivante. Plus important encore, pendant que le brouillon fonctionne, le lot suivant peut être planifié.
Calendrier MTP pour un jeton unique
Bien que le runtime ne fournisse pas encore l'exploration arborescente de brouillon de style Aigle, nous avons mis en place un cas particulier de ce schéma, considérant une séquence linéaire de jetons de brouillon produite 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 de brouillon utilisant les poids open-source de DeepSeek R1. Le sous-cas de la prédiction d'un seul jeton est intéressant, car de 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 de brouillon est beaucoup plus rapide, cachant moins de latence côté CPU. De plus, le brouillon est partagé avec le 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, semblable au schéma précédent. Ensuite, le cible s'exécute pour valider les jetons, traitant 2 * D jetons, où D est la taille du lot de décodage. Cela est idéal pour le micro-batching dans les modèles de Mixture-of-Experts (MoE) sur des interconnexions plus lentes telles que InfiniBand, car le lot se divise également en deux moitiés. Les états cachés de la cible sont transférés à la prochaine exécution de brouillon, tandis que les logits sont passés à l'échantillonneur pour la vérification.
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 de brouillon. Une fois que les jetons d'entrée des cibles sont décalés, un noyau y insère les prochains jetons cibles à leurs emplacements correspondants. Le brouillon est ensuite réexécuté avec les mêmes informations de lot que la cible, remplissant les entrées de cache KV et construisant les logits et états cachés pour la prochaine exécution, en faisant un peu de 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 de brouillon. En parallèle à l'exécution de brouillon, 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 d'un travail supplémentaire dans la couche de brouillon n'est pas notable dans l'attention, cependant les couches MLP posent plus de problèmes. Étant donné que les instructions de multiplication matricielle remplissent jusqu'à une frontière 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 caché. Pour les séquences de brouillon plus longues, le surcoût est plus cher et le schéma utilisé pour les modèles brouillon-cible réguliers fonctionne mieux.
Références
Inférence Rapide à partir des Transformeurs via le 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 Arbres de Brouillon Dynamiques
Medusa: Cadre Simple d'Accélération d'Inférence LLM avec Têtes de Décodage Multiples
FlashInfer: Moteur d'Attention Efficace et Personnalisable pour le Service d'Inférence LLM