AWS montre comment le speculative decoding sur Trainium2 accélère la génération dans vLLM

AWS présente une méthode concrète pour accélérer l'inférence LLM sur Trainium2 et en réduire le coût, dans les cas où le modèle génère bien plus de tokens qu'il n'en reçoit en entrée. Il s'agit du speculative decoding : au lieu de forcer un grand modèle à produire les tokens un par un, le système adjoint un petit modèle draft qui propose rapidement plusieurs tokens suivants d'un coup, puis le modèle principal, dit target, les vérifie en une seule passe. Quand les prédictions concordent, le service économise des étapes séquentielles coûteuses, réduit la latence entre tokens et exploite mieux l'accélérateur.

L'enjeu est particulièrement fort pour les charges decode-heavy : assistants d'écriture, coding agents, génération de rapports, documents types et autres tâches à réponse longue. En génération autorégressive classique, chaque nouveau token est calculé séparément : l'accélérateur relit en permanence le KV-cache en mémoire et n'effectue que relativement peu de travail utile par étape. L'inférence bute donc souvent sur la bande passante mémoire plutôt que sur la puissance de calcul brute.

Le speculative decoding s'attaque précisément à ce goulot d'étranglement : le modèle target exécute moins d'étapes de décodage séquentielles, et la vérification d'un lot de tokens densifie la charge. L'approche a toutefois ses exigences. Les modèles draft et target doivent partager le même tokenizer et le même vocabulaire, et idéalement appartenir à la même famille architecturale, afin que le petit modèle devine plus souvent la suite produite par le modèle principal.

Le paramètre clé est le nombre de speculative tokens. Une fenêtre trop petite rend le gain presque imperceptible ; une fenêtre trop grande voit son effet absorbé par les rejets précoces et les vérifications superflues. Pour son test, AWS a utilisé Qwen3-32B comme modèle target et Qwen3-1.

7B comme modèle draft, servis via vLLM sur une instance Trn2 trn2.48xlarge. Pour le speculative decoding, l'équipe a retenu la fused speculation de NeuronX Distributed Inference, qui compile les deux modèles ensemble pour de meilleures performances.

Les configurations de référence et speculative ont été déployées dans le même cluster Amazon EKS, à paramètres identiques : allocation des accélérateurs, tensor parallelism, longueur de contexte, batch limits et image Neuron. Seuls différaient l'ajout du modèle draft et le paramètre num_speculative_tokens. La charge a été envoyée aux deux services via llmperf, et le TTFT, l'inter-token latency et la latence de bout en bout ont été remontés dans CloudWatch pour comparaison.

AWS a aussi testé le plus compact Qwen3-0.6B, mais son acceptance rate s'est révélé inférieur d'environ 60 %, ce qui a suffi à faire perdre l'essentiel du gain. Sur une plage de 5 à 15 speculative tokens, la configuration à sept tokens s'est imposée comme le meilleur point de ces tests, même si les auteurs soulignent que la valeur optimale dépend fortement de la structure des prompts.

C'est d'ailleurs la structure des requêtes qui a déterminé le résultat. Sur les scénarios prévisibles — texte répétitif, séquences numériques, code simple —, le speculative decoding apporte un bénéfice net. Dans ces cas, le modèle draft devine souvent ce que le modèle target aurait de toute façon produit, et le système saute une part importante des étapes séquentielles.

Dans les tests, l'inter-token latency descendait à environ 15 millisecondes par token, et la courbe de latence de bout en bout restait systématiquement sous la baseline. Sur les requêtes ouvertes, moins déterministes, le tableau change : le modèle draft diverge plus souvent du modèle target, les tokens sont rejetés et le gain potentiel disparaît. Pour ces prompts, l'inter-token latency se maintenait autour de 45 millisecondes par token, et les configurations speculative et baseline affichaient une latence globale quasi identique.

Le TTFT, c'est-à-dire le temps jusqu'au premier token, n'a pratiquement pas bougé, car le speculative decoding n'accélère pas la phase de prefill, durant laquelle le modèle encode le contexte d'entrée. Le bénéfice principal intervient plus tard, en phase de décodage, grâce à la réduction du nombre d'étapes séquentielles coûteuses du modèle target. La conclusion pratique de l'article est simple : sur Trainium2, le speculative decoding n'est pas un bouton d'accélération universel, mais une optimisation ciblée pour un type de charge précis.

Si un produit génère souvent une sortie structurée et prévisible — code, extraction de données, rapports types, fichiers de configuration —, ce mode peut directement réduire le coût du token de sortie et augmenter le débit sans perte de qualité. Si en revanche votre usage relève surtout du chat ouvert à génération libre, l'effet peut être minime. Mieux vaut donc ne déployer ce schéma qu'après des mesures sur vos propres prompts, en choisissant un modèle draft compatible et une fenêtre de speculative tokens adaptée aux scénarios réels, plutôt qu'en se fiant à des benchmarks déconnectés du produit.