AWS montre comment la décodification spéculative sur Trainium2 accélère la génération dans vLLM

AWS a montré un moyen pratique d'accélérer et de réduire le coût de l'inférence LLM sur Trainium2 pour les scénarios où le modèle génère significativement plus de tokens qu'il n'en reçoit en entrée. Il s'agit de speculative decoding : au lieu de forcer un grand modèle à produire séquentiellement un token à la fois, le système connecte un petit draft-model qui propose rapidement plusieurs tokens suivants à la fois, tandis que le main target-model les vérifie en une seule passe. Si les prédictions correspondent, le service dépense moins d'étapes séquentielles coûteuses, réduit la latence entre tokens et utilise mieux l'accélérateur.

C'est particulièrement important pour les charges decode-heavy — assistants d'écriture, coding agents, génération de rapports, documents modèles et autres tâches avec des réponses longues. Dans la génération autorégrессive standard, chaque nouveau token est calculé séparément, donc l'accélérateur lit constamment KV-cache de la mémoire et effectue relativement peu de travail utile par étape. À cause de cela, l'inférence bute souvent sur les limites de bande passante mémoire plutôt que sur le pur calcul.

Speculative decoding cible exactement ce goulot : le target-model exécute les étapes de decode séquentiel moins fréquemment, et la vérification par lot rend la charge plus dense. Cependant, l'approche a des exigences. Les modèles draft et target doivent utiliser le même tokeniseur et vocabulaire, et idéalement appartenir à la même famille architecturale pour que le petit modèle devine plus souvent la continuation du modèle principal.

Un paramètre clé est le nombre de speculative tokens. Si la fenêtre est trop petite, le gain est à peine perceptible ; si elle est trop grande, les rejets précoces et la vérification inutile consomment le bénéfice. Dans leur test, AWS a utilisé le target-model Qwen3-32B et draft-model Qwen3-1.

7B, exécutés via vLLM sur une instance trn2.48xlarge. Pour speculative decoding, ils ont choisi fused speculation dans NeuronX Distributed Inference, où les deux modèles sont compilés ensemble pour de meilleures performances.

Les configurations baseline et speculative ont été déployées dans un seul cluster Amazon EKS en maintenant tout identique : allocation d'accélérateur, tensor parallelism, longueur du contexte, batch limits et image Neuron. La seule différence était l'ajout du draft-model et du paramètre num_speculative_tokens. La charge a été appliquée aux deux services via llmperf, et TTFT, inter-token latency et latence de bout en bout ont été envoyés à CloudWatch pour comparaison.

AWS a également testé le plus compact Qwen3-0.6B, mais son taux d'acceptation était environ 60 pour cent inférieur, ce qui a suffi à perdre l'essentiel du bénéfice. Dans la plage de 5 à 15 speculative tokens, le point optimal dans ces tests était une configuration avec sept tokens, bien que les auteurs soulignent que la valeur optimale dépend fortement de la structure du prompt.

En fin de compte, la structure de la requête a déterminé le résultat. Sur les scénarios prévisibles — texte répété, séquences numériques, code simple — speculative decoding a montré des bénéfices clairs. Dans de tels cas, le draft-model devine souvent ce que le target-model produirait de toute façon, donc le système ignore une portion significative d'étapes séquentielles.

Dans les tests, inter-token latency est tombée à environ 15 millisecondes par token, et la courbe de latence de bout en bout est restée régulièrement en dessous de la baseline. Sur les requêtes ouvertes et moins déterministes, le tableau est différent : le draft-model diverge plus souvent du target-model, les tokens sont rejetés et le bénéfice potentiel disparaît. Pour ces prompts, inter-token latency s'est maintenue autour de 45 millisecondes par token, et les configurations speculative et baseline ont montré une latence de bout en bout presque identique.

TTFT — temps jusqu'au premier token — a à peine changé car speculative decoding n'accélère pas l'étape prefill, où le modèle encode le contexte d'entrée. Le principal bénéfice apparaît plus tard, dans la phase de decode, en réduisant le nombre d'étapes séquentielles coûteuses du target-model. La conclusion pratique de l'article est simple : speculative decoding sur Trainium2 n'est pas un bouton d'accélération universel, mais une optimisation ciblée pour un type de charge spécifique.

Si votre produit génère souvent des résultats structurés et prévisibles — code, extraction de données, rapports modèles, configs — ce mode peut réduire directement le coût du token de sortie et augmenter le débit sans perte de qualité. Si vous avez principalement du chat ouvert avec génération libre, l'effet peut être minimal. Par conséquent, mettre en œuvre ce schéma ne vaut le coup qu'après avoir fait des benchmarks sur vos propres prompts, en sélectionnant un draft-model compatible et une fenêtre de speculative tokens adaptée aux scénarios réels, plutôt que de vous fier à des benchmarks isolés de votre produit.