AWS muestra cómo la decodificación especulativa en Trainium2 acelera la generación en vLLM

AWS mostró una forma práctica de acelerar y reducir el costo de la inferencia de LLM en Trainium2 para escenarios donde el modelo genera significativamente más tokens de los que recibe como entrada. Se trata de speculative decoding: en lugar de obligar a un modelo grande a producir secuencialmente un token a la vez, el sistema conecta un pequeño draft-model que rápidamente propone varios tokens siguientes a la vez, mientras que el main target-model los verifica en un único paso. Si las predicciones coinciden, el servicio gasta menos pasos secuenciales costosos, reduce la latencia entre tokens y aprovecha mejor el acelerador.

Esto es especialmente importante para cargas decode-heavy — asistentes de escritura, coding agents, generación de reportes, documentos plantilla y otras tareas con respuestas largas. En la generación autorregresiva estándar, cada nuevo token se calcula por separado, por lo que el acelerador constantemente lee KV-cache de la memoria y realiza relativamente poco trabajo útil por paso. Debido a esto, la inferencia a menudo se ve limitada por el ancho de banda de la memoria en lugar del puro cálculo.

Speculative decoding ataca exactamente este cuello de botella: el target-model ejecuta pasos de decode secuencial con menos frecuencia, y la verificación por lotes hace la carga más densa. Sin embargo, el enfoque tiene requisitos. Los modelos draft y target deben usar el mismo tokenizador y vocabulario, e idealmente pertenecer a la misma familia arquitectónica para que el modelo pequeño adivine más frecuentemente la continuación del principal.

Un parámetro clave es el número de speculative tokens. Si la ventana es demasiado pequeña, la ganancia es apenas perceptible; si es demasiado grande, los rechazos tempranos y la verificación innecesaria consumen el beneficio. En su prueba, AWS utilizó el target-model Qwen3-32B y draft-model Qwen3-1.

7B, ejecutados a través de vLLM en una instancia trn2.48xlarge. Para speculative decoding, eligieron fused speculation en NeuronX Distributed Inference, donde ambos modelos se compilan juntos para mejor rendimiento.

Las configuraciones baseline y speculative se desplegaron en un único clúster Amazon EKS manteniendo todo idéntico: asignación de acelerador, tensor parallelism, longitud del contexto, batch limits e imagen Neuron. La única diferencia fue la adición del draft-model y el parámetro num_speculative_tokens. La carga se aplicó a ambos servicios a través de llmperf, y TTFT, inter-token latency y latencia de extremo a extremo se enviaron a CloudWatch para comparación.

AWS también probó el más compacto Qwen3-0.6B, pero su tasa de aceptación fue aproximadamente 60 por ciento menor, lo que fue suficiente para perder la mayor parte del beneficio. En el rango de 5 a 15 speculative tokens, el punto óptimo en estas pruebas fue una configuración con siete tokens, aunque los autores enfatizan que el valor óptimo depende fuertemente de la estructura del prompt.

En última instancia, la estructura de la solicitud determinó el resultado. En escenarios predecibles — texto repetido, secuencias numéricas, código simple — speculative decoding mostró beneficios claros. En tales casos, el draft-model frecuentemente adivina lo que el target-model produciría de todos modos, por lo que el sistema salta una porción significativa de pasos secuenciales.

En las pruebas, inter-token latency cayó a aproximadamente 15 milisegundos por token, y la curva de latencia de extremo a extremo se mantuvo consistentemente por debajo de la baseline. En solicitudes abiertas y menos deterministas, el cuadro es diferente: el draft-model diverge más frecuentemente del target-model, los tokens se rechazan y la ganancia potencial desaparece. Para estos prompts, inter-token latency se mantuvo alrededor de 45 milisegundos por token, y las configuraciones speculative y baseline mostraron latencia de extremo a extremo casi idéntica.

TTFT — tiempo hasta el primer token — cambió poco porque speculative decoding no acelera la etapa prefill, donde el modelo codifica el contexto de entrada. El principal beneficio aparece después, en la fase de decode, al reducir el número de pasos secuenciales costosos del target-model. La conclusión práctica del artículo es simple: speculative decoding en Trainium2 no es un botón de aceleración universal, sino una optimización dirigida para un tipo de carga específico.

Si su producto genera frecuentemente salida estructurada y predecible — código, extracción de datos, reportes plantilla, configs — este modo puede reducir directamente el costo del token de salida y aumentar el rendimiento sin pérdida de calidad. Si principalmente tiene chat abierto con generación de forma libre, el efecto puede ser mínimo. Por lo tanto, implementar este esquema es recomendable solo después de hacer benchmarks en sus propios prompts, seleccionando un draft-model compatible y una ventana de speculative tokens adecuada para escenarios reales, en lugar de confiar en benchmarks aislados de su producto.