Acelerar la inferencia LLM con mucha decodificación con decodificación especulativa en AWS Trainium y vLLM

Puntos de referencia prácticos que muestran una latencia entre tokens más rápida al implementar modelos Qwen3 con vLLM, Kubernetes y chips AI de AWS.

La decodificación especulativa en AWS Trainium puede acelerar la generación de tokens hasta 3 veces para cargas de trabajo con gran cantidad de decodificación, lo que ayuda a reducir el costo por token de salida y mejora el rendimiento sin sacrificar la calidad de la salida. Si crea asistentes de escritura de IA, agentes de codificación u otras aplicaciones de IA generativa, es probable que sus cargas de trabajo produzcan muchos más tokens de los que consumen, lo que hace que la etapa de decodificación sea el costo dominante de la inferencia. Durante la decodificación autorregresiva, los tokens se generan secuencialmente, lo que deja a los aceleradores de hardware limitados por el ancho de banda de la memoria y subutilizados. Esto aumenta el costo por token generado. La decodificación especulativa aborda este cuello de botella al permitir que un pequeño modelo borrador proponga múltiples tokens a la vez, que el modelo objetivo verifica en un solo paso hacia adelante. Menos pasos de decodificación en serie significan una menor latencia y una mayor utilización del hardware, lo que ayuda a reducir los costos de inferencia.

En esta publicación aprenderás:

Cómo funciona la decodificación especulativa y por qué ayuda a reducir el costo por token generado en AWS Trainium2 Cómo habilitar la decodificación especulativa con vLLM en Trainium La metodología de evaluación comparativa que usamos para evaluar el rendimiento Cómo ajustar la selección del modelo preliminar y el tamaño de la ventana del token especulativo para sus cargas de trabajo Instrucciones paso a paso para reproducir los resultados usando Qwen3

¿Qué es la decodificación especulativa?

La decodificación especulativa acelera la generación autorregresiva mediante el uso de dos modelos:

Un borrador de modelo propone n tokens candidatos rápidamente. Un modelo objetivo los verifica en un pase hacia adelante.

Para obtener una visión más profunda de los mecanismos subyacentes, incluida la aceptación y el rechazo de tokens, la especulación basada en EAGLE y los conceptos generales de decodificación especulativa, consulte la publicación del blog Inferentia2, este tutorial de SageMaker EAGLE en AWS Inferentia2, este tutorial de SageMaker EAGLE y este manual. Aquí, nos centramos en los dos botones que controlas en la práctica: el modelo borrador y num_speculative_tokens.

Los modelos borrador y de destino deben compartir el mismo tokenizador y vocabulario, porque la decodificación especulativa opera en ID de token verificados directamente por el modelo de destino. Recomendamos elegir modelos de la misma familia arquitectónica porque sus predicciones del siguiente token coinciden con mayor frecuencia. Puede emparejar modelos con diferentes arquitecturas si comparten un tokenizador, pero una menor concordancia entre el modelo borrador y el de destino reduce las tasas de aceptación y elimina la mayor parte de la ganancia de rendimiento.

Cuando el modelo de destino acepta los borradores de tokens, se comprometen sin incurrir en el costo total de los pasos de decodificación secuenciales. El parámetro principal que controlas es num_speculative_tokens, que establece cuántos tokens propone el modelo preliminar a la vez. Aumentar este valor le permite omitir más pasos de decodificación en serie por paso de verificación, lo que reduce directamente la latencia entre tokens cuando las tasas de aceptación son altas.

La ganancia de rendimiento proviene de dos efectos. Primero, la decodificación especulativa reduce la cantidad de pasos de decodificación del modelo de destino, lo que reduce la cantidad de viajes de ida y vuelta de la memoria caché KV. (La caché KV almacena tensores de clave y valor previamente calculados para que el modelo no vuelva a calcular la atención de tokens pasados. Cada paso de decodificación lee la caché completa de la memoria, lo que hace que la decodificación esté limitada al ancho de banda de la memoria). En segundo lugar, la decodificación especulativa mejora la utilización del hardware durante la decodificación. En la decodificación autorregresiva estándar, cada paso de decodificación produce sólo un token nuevo: el acelerador lanza costosos núcleos de multiplicación de matrices para producir solo un token de trabajo, dejando al motor de elementos de procesamiento en gran medida infrautilizado. Durante la verificación, el modelo de destino procesa n tokens a la vez, amortizando el acceso a la memoria y convirtiendo una secuencia de cálculos pequeños e ineficientes de un solo token en una carga de trabajo más densa en computación. Establecer num_speculative_tokens demasiado bajo limita las ganancias de velocidad.

Configurarlo demasiado alto aumenta la probabilidad de rechazos tempranos, desperdiciando borradores de cálculo y aumentando el costo de verificación del modelo objetivo. Puede ajustar este valor equilibrando el cálculo preliminar con el costo de verificación en función de su tasa de aceptación observada.

Figura 1 Compensaciones de configuración de decodificación especulativa

Para ilustrar estas compensaciones, comparamos los modelos preliminares Qwen3-0.6B y Qwen3-1.7B. El modelo más pequeño de 600 millones fue más rápido de ejecutar, pero su tasa de aceptación fue aproximadamente un 60% menor, suficiente para anular los ahorros de computación. Qwen3-1.7B logró un mejor equilibrio entre velocidad y aceptación.

Para num_speculative_tokens, evaluamos valores de 5 a 15. Las configuraciones más pequeñas (por ejemplo, 5) ofrecieron una aceleración limitada. Las ventanas más grandes (por ejemplo, 15) aumentaron los rechazos y degradaron el rendimiento. La mejor configuración dependía en gran medida de la estructura del mensaje. Probamos tanto indicaciones estructuradas (como repetición, secuencias numéricas y código simple) como lenguaje natural abierto. El mejor saldo provino de Qwen3-1.7B con 7 tokens especulativos. Consulte la sección Lecciones aprendidas para obtener detalles completos sobre el ajuste.

Qué admite la inferencia distribuida de NeuronX (inferencia NxD)

AWS Neuron es el SDK para los chips de IA de AWS. NeuronX Distributed Inference (NxDI) es su biblioteca para inferencia LLM escalable y de alto rendimiento en Trainium e Inferentia. NxDI proporciona soporte nativo para decodificación especulativa en Trainium en cuatro modos:

Decodificación especulativa básica: modelos de borrador y de destino separados compilados de forma independiente. La forma más sencilla de empezar. Especulación fusionada: modelos borrador y objetivo compilados juntos para mejorar el rendimiento. Este es el modo que utilizamos en esta publicación. Especulación de EAGLE: el borrador del modelo aprovecha el contexto de estado oculto del modelo objetivo para mejorar las tasas de aceptación. Especulación de Medusa: múltiples cabezales de predicción pequeños se ejecutan en paralelo para proponer tokens, lo que reduce la sobrecarga del modelo preliminar.

Para obtener documentación completa, consulte la guía de decodificación especulativa y la guía de decodificación especulativa EAGLE. Esta publicación utiliza especulación fusionada, donde el modelo borrador (Qwen3-1.7B) y el modelo objetivo (Qwen3-32B) se compilan junto con enable_fused_speculation=true para un rendimiento óptimo en Neuron.

Comenzando con la decodificación especulativa en AWS Trainium

Implementamos dos servicios de inferencia vLLM en instancias de Trainium en el mismo clúster de Amazon Elastic Kubernetes Service (Amazon EKS), manteniendo todo idéntico excepto el método de decodificación para aislar el impacto en el rendimiento. El servicio básico (qwen-vllm) sirve Qwen3-32B con decodificación estándar. El servicio especulativo (qwen-sd-vllm) sirve el mismo modelo objetivo Qwen3-32B, agregando un modelo borrador Qwen3-1.7B con num_speculative_tokens=7.

Ambos servicios ejecutan configuraciones idénticas en Trn2 (trn2.48xlarge), la misma asignación de acelerador, paralelismo tensor (que distribuye los pesos del modelo en múltiples NeuronCores para adaptarse a modelos grandes), longitud de secuencia, límites de procesamiento por lotes e imagen de Neuron DLC. La única diferencia es la adición del modelo borrador Qwen3-1.7B y num_speculative_tokens=7 para el servicio especulativo. Consulte la Figura 2 para obtener detalles completos de la configuración.

Para comparar las dos configuraciones bajo una carga idéntica, utilizamos llmperf para generar los mismos patrones de tráfico en ambos puntos finales. Capturamos la telemetría de la infraestructura con CloudWatch Container Insights y publicamos métricas personalizadas a nivel de solicitud (TTFT, latencia entre tokens y latencia de un extremo a otro) en los paneles de CloudWatch para realizar un análisis en paralelo.

Diagrama de arquitectura de EKS que muestra la implementación de tres niveles con módulos de aplicaciones, módulos de sistemas e infraestructura para servicios de inferencia LLM

Figura 2 Arquitectura del sistema

Configuración de evaluación comparativa

Usamos LLMPerf para ejecutar casos de prueba estructurados y con gran cantidad de decodificación en implementaciones de decodificación tanto de referencia como especulativas. Los puntos de referencia se ejecutaron dentro de un pod de Kubernetes, qwen-llmperf-pod.yaml, emitiendo solicitudes simultáneas a ambos puntos finales y registrando métricas de latencia a nivel de token. Nuestros casos de prueba abarcaron desde indicaciones altamente estructuradas (secuencias repetitivas, continuaciones numéricas, patrones de código simples) hasta terminaciones abiertas en lenguaje natural, cubriendo tanto el mejor como el peor de los casos para la decodificación especulativa. El conjunto completo de mensajes está disponible en el repositorio de ejemplos.

Para mayor claridad, centramos el análisis en dos tipos de mensajes representativos: un mensaje determinista altamente estructurado (generación de texto repetitivo) y un mensaje abierto. Estos dos casos ilustran el comportamiento tanto en el mejor como en el peor de los casos de la decodificación especulativa.

El pod ejecutó llmperf con longitudes de entrada y salida controladas y temperatura = 0,0 para enfatizar las rutas de decodificación deterministas. Registramos y publicamos métricas que incluyen latencia entre tokens, TTFT, rendimiento y latencia de un extremo a otro en CloudWatch.

Resultados

Gráfico de líneas que muestra la latencia de un extremo a otro en segundos para cuatro configuraciones de prueba que comparan las implementaciones SD y Base a lo largo del tiempo.

Figura 3 Latencia E2E de decodificación especulativa

La decodificación especulativa reduce la latencia de forma selectiva: su efectividad depende en gran medida de la estructura del aviso, y esta dependencia aparece consistentemente en todas las métricas medidas. Esto es lo que puede esperar de cada tipo de aviso:

Mensajes estructurados (por ejemplo, "Repita la siguiente línea exactamente 50 veces"). La decodificación especulativa ofrece una reducción mensurable de la latencia de un extremo a otro. Cuando el borrador del modelo predice de manera confiable lo que generaría el modelo objetivo, el sistema omite una fracción sustancial de los pasos de decodificación del modelo objetivo. En nuestras pruebas, la latencia entre tokens se redujo a aproximadamente 15 ms por token (en comparación con aproximadamente 45 ms para las indicaciones abiertas) y la curva de decodificación especulativa se mantuvo consistentemente por debajo de la línea de base durante toda la ejecución. Preguntas abiertas (por ejemplo, “Creo que el significado de la vida es”). La decodificación especulativa no proporciona ningún beneficio constante. El modelo borrador con frecuencia diverge del modelo objetivo, lo que provoca rechazos simbólicos que anulan las ganancias potenciales. Las curvas de latencia de extremo a extremo especulativas y de referencia se superponen en gran medida, y la latencia entre tokens se mantiene cerca de 45 ms por token para ambas configuraciones.

Gráfico de líneas que muestra la latencia entre tokens en segundos para cuatro configuraciones de LLM que comparan las implementaciones SD y Base durante 3 horas

Figura 4 Latencia entre tokens de decodificación especulativa (Decode)

TTFT (Tiempo hasta el primer token) permanece efectivamente sin cambios en todas las configuraciones (Figura 5). TTFT está dominado por la fase de prellenado, donde el modelo codifica el contexto de entrada. La decodificación especulativa no altera esta etapa, por lo que la latencia de precarga no mejora ni se degrada.

Gráfico de líneas que compara las métricas de rendimiento TTFT para dos mensajes de texto (versiones SD vs Base) de 04:40 a 07:40

Figura 5 Decodificación especulativa TTFT (Prefill)

En conjunto, estos resultados muestran que la decodificación especulativa mejora la latencia total al reducir la cantidad de pasos de decodificación del modelo de destino ejecutados, no al acelerar el paso de decodificación en sí o la etapa de precarga. Esto explica por qué aparecen ganancias en la latencia de un extremo a otro para las indicaciones estructuradas, pero están ausentes en la latencia entre tokens y TTFT, y por qué la decodificación especulativa vuelve al comportamiento básico para la generación abierta.

Reproduciendo los resultados

Proporcionamos ejemplos de código de un extremo a otro y configuraciones de Kubernetes en el repositorio de ejemplos de AWS Neuron EKS. El repositorio incluye:

Manifiestos de Kubernetes para implementar vLLM de línea base y servicios vLLM de decodificación especulativa en Trn2 Indicadores de configuración de vLLM de ejemplo para habilitar la decodificación especulativa fusionada Scripts de evaluación comparativa de llmperf de muestra utilizados para generar carga y recopilar métricas Instrucciones para montar puntos de control del modelo y artefactos compilados a través del controlador S3 CSI Guía sobre la configuración de Neuron DRA, paralelismo tensorial y ubicación de NeuronCore

Estos ejemplos le permiten recrear la misma configuración experimental utilizada en esta publicación, desde la implementación del modelo hasta la evaluación comparativa y la recopilación de métricas.

Conclusión

Las cargas de trabajo de LLM con mucha decodificación están limitadas por la naturaleza secuencial de la generación autorregresiva. La decodificación especulativa rompe este cuello de botella en AWS Trainium2 al reducir la cantidad de pasos de decodificación del modelo de destino necesarios para producir el resultado completo, lo que aumenta efectivamente los tokens generados por paso directo. Para cargas de trabajo donde el espacio de salida es predecible, como la generación de código, la extracción de datos estructurados, la generación de informes con plantillas o la síntesis de archivos de configuración, esto puede traducirse directamente en un menor costo por token de salida y un mayor rendimiento, sin sacrificar la calidad. La decodificación especulativa no es una optimización universal. Su eficacia depende de una estructura rápida, la calidad del modelo preliminar y el ajuste especulativo de los parámetros. Cuando se aplica a las cargas de trabajo adecuadas, ofrece latencia significativa y mejoras de costos en los sistemas de inferencia basados ​​en Trainium.

Próximos pasos

Para comenzar con la decodificación especulativa en AWS Trainium, explore estos recursos:

Sobre los autores

Yahav Biran es arquitecto principal en Amazon y se centra en cargas de trabajo de IA a gran escala. Contribuye a proyectos de código abierto y publica en blogs y revistas académicas de AWS, incluidos los blogs de computación e inteligencia artificial de AWS y el Journal of Systems Engineering. Con frecuencia realiza presentaciones técnicas y colabora con clientes para diseñar aplicaciones en la nube. Yahav tiene un doctorado. en Ingeniería de Sistemas de la Universidad Estatal de Colorado.

Truong Pham es ingeniero de software en Annapurna Labs, Amazon. Se especializa en optimizar el rendimiento de inferencia de modelos de lenguaje grande en aceleradores de IA de AWS, como AWS Inferentia y Trainium, y en diseñar API fáciles de usar para los desarrolladores para la pila de software AWS Neuron. Truong tiene un doctorado. en Ingeniería Química de la Universidad de Minnesota.