AWS muestra cómo el speculative decoding en Trainium2 acelera la generación en vLLM

AWS ha mostrado una forma práctica de acelerar y abaratar la inferencia de LLM en Trainium2 en escenarios donde el modelo genera muchos más tokens de los que recibe como entrada. Se trata del speculative decoding: en lugar de obligar al modelo grande a producir los tokens de uno en uno, el sistema incorpora un pequeño modelo draft que propone rápidamente varios tokens siguientes, y el modelo principal (target) los verifica en una sola pasada. Si las predicciones coinciden, el servicio ejecuta menos pasos secuenciales costosos, reduce la latencia entre tokens y aprovecha mejor el acelerador.

Esto es especialmente relevante para cargas decode-heavy: asistentes de escritura, coding agents, generación de informes, documentos basados en plantillas y otras tareas con respuestas largas. En la generación autorregresiva convencional, cada token nuevo se calcula por separado, de modo que el acelerador lee constantemente la KV-cache de memoria y realiza relativamente poco trabajo útil por paso. Por eso la inferencia suele quedar limitada por el ancho de banda de memoria y no por la capacidad de cómputo.

El speculative decoding ataca justo ese cuello de botella: el modelo target ejecuta menos pasos de decode secuenciales y la verificación de un lote de tokens hace la carga más densa. Pero el enfoque tiene requisitos. Los modelos draft y target deben usar el mismo tokenizador y vocabulario, y conviene además que pertenezcan a la misma familia arquitectónica, para que el modelo pequeño acierte con más frecuencia la continuación del principal.

El parámetro clave es el número de speculative tokens. Si la ventana es demasiado pequeña, la ganancia apenas se nota; si es demasiado grande, los rechazos tempranos y la verificación adicional se comen el beneficio. En su prueba, AWS utilizó Qwen3-32B como modelo target y Qwen3-1.

7B como modelo draft, ejecutados con vLLM en una instancia Trn2 trn2.48xlarge. Para el speculative decoding eligió fused speculation en NeuronX Distributed Inference, una modalidad en la que ambos modelos se compilan juntos para obtener mejor rendimiento.

Las configuraciones base y speculative se desplegaron en el mismo clúster de Amazon EKS, con todo lo demás idéntico: asignación de aceleradores, tensor parallelism, longitud de contexto, límites de batch e imagen de Neuron. La única diferencia fue la incorporación del modelo draft y del parámetro num_speculative_tokens. La carga se generó en ambos servicios con llmperf, y el TTFT, la inter-token latency y la latencia de extremo a extremo se enviaron a CloudWatch para compararlas.

AWS probó además el modelo más compacto Qwen3-0.6B, pero su acceptance rate resultó aproximadamente un 60 por ciento más bajo, suficiente para perder la mayor parte de la ganancia. En el rango de 5 a 15 speculative tokens, el punto óptimo en estas pruebas fue la configuración de siete tokens, aunque los autores subrayan que el valor óptimo depende mucho de la estructura de los prompts.

Fue precisamente la estructura de las peticiones la que acabó determinando el resultado. En escenarios predecibles —texto repetitivo, secuencias numéricas, código sencillo— el speculative decoding mostró un beneficio claro. En esos casos, el modelo draft suele acertar lo que el modelo target habría generado de todos modos, así que el sistema se salta una parte importante de los pasos secuenciales.

En las pruebas, la inter-token latency bajó hasta unos 15 milisegundos por token, y la curva de latencia de extremo a extremo se mantuvo de forma estable por debajo del baseline. Con peticiones abiertas y menos deterministas, el panorama cambia: el modelo draft diverge más a menudo del target, los tokens se rechazan y la ganancia potencial desaparece. Para esos prompts, la inter-token latency se mantuvo en torno a 45 milisegundos por token, y las configuraciones speculative y baseline mostraron una latencia total casi idéntica.

El TTFT, es decir, el tiempo hasta el primer token, apenas cambió, porque el speculative decoding no acelera la fase de prefill, en la que el modelo codifica el contexto de entrada. El beneficio principal llega después, en la fase de decode, gracias a la reducción del número de pasos secuenciales costosos del modelo target. La conclusión práctica del artículo es sencilla: el speculative decoding en Trainium2 no es un botón universal de aceleración, sino una optimización puntual para un tipo concreto de carga.

Si un producto genera con frecuencia salidas estructuradas y predecibles —código, extracción de datos, informes basados en plantillas, archivos de configuración—, este modo puede reducir directamente el coste por token de salida y aumentar el throughput sin pérdida de calidad. En cambio, si lo que tienes es sobre todo un chat abierto con generación libre, el efecto puede ser mínimo. Por eso, conviene implantar este esquema solo después de medir con tus propios prompts, eligiendo un modelo draft compatible y ajustando la ventana de speculative tokens a los escenarios reales, y no a partir de benchmarks desligados del producto.