Optimización de Inferencia de LLMs con Decodificación Especulativa en vLLM y GPUs AMD

La inferencia de Large Language Models (LLMs) se ve fundamentalmente limitada por la naturaleza autoregresiva de su generación de tokens, donde cada token se produce secuencialmente. Este proceso, aunque garantiza la coherencia, impone una latencia significativa y limita el throughput, especialmente en escenarios de alta demanda. La decodificación especulativa aborda este problema de rendimiento al desacoplar la propuesta de tokens de su verificación. En lugar de generar un token a la vez con el modelo completo, un componente de borrador más ligero predice una secuencia de tokens candidatos, que luego son validados en bloque por el modelo objetivo. Si la predicción es correcta, se "aceptan" múltiples tokens en una sola pasada del modelo objetivo, lo que reduce drásticamente el número de operaciones costosas del modelo principal y, por ende, mejora el throughput. Este enfoque es una aplicación práctica del principio de "predicción y corrección" en sistemas distribuidos, donde un componente rápido pero potencialmente impreciso propone una solución que un componente más lento pero preciso valida, optimizando la latencia efectiva sin comprometer la exactitud del resultado final.

La relevancia actual de la decodificación especulativa radica en la creciente demanda de servir LLMs a escala, donde cada milisegundo de latencia y cada token por segundo son críticos para la experiencia del usuario y la eficiencia operativa. A medida que los modelos se vuelven más grandes y complejos, la optimización de la inferencia se convierte en un cuello de botella primario. La decodificación especulativa, junto con otras técnicas como la cuantificación y la paralelización, es esencial para hacer que los LLMs sean económicamente viables y escalables en entornos de producción.

Arquitectura del Sistema

El núcleo de la decodificación especulativa se basa en una arquitectura de dos componentes: un "draft component" (modelo borrador) y un "target model" (modelo objetivo). El modelo objetivo es el LLM original y es responsable de la salida final correcta. El modelo borrador es un modelo más pequeño y rápido, diseñado para proponer secuencias de tokens candidatos. El proceso se desarrolla en rondas: el modelo borrador propone N tokens, que el modelo objetivo evalúa en una única pasada de verificación. Los tokens aceptados se añaden a la secuencia de salida; si un token es rechazado, la secuencia propuesta se trunca en ese punto y el modelo objetivo genera el siguiente token correcto, continuando la generación desde allí.

Existen varias implementaciones del "draft component", que se clasifican en tres categorías principales: módulos MTP nativos, drafters MTP separados y redes de borrador dedicadas condicionadas por el objetivo. Los módulos MTP nativos (como en Qwen3.5/3.6) están integrados directamente en la arquitectura del modelo objetivo y reutilizan componentes, generando tokens secuencialmente. Los drafters MTP separados (como Gemma 4 MTP) usan un checkpoint independiente pero comparten activaciones y el KV cache del modelo objetivo, también generando tokens secuencialmente. Las redes de borrador dedicadas (EAGLE-3, DFlash, DSpark) son modelos especuladores entrenados específicamente para un modelo objetivo. EAGLE-3 genera tokens autoregresivamente a partir de estados ocultos del modelo objetivo. DFlash predice un bloque completo de tokens en paralelo, utilizando estados ocultos fusionados del modelo objetivo como información Key y Value en cada capa de su red. DSpark extiende DFlash añadiendo una "lightweight sequential head" para introducir dependencia entre tokens dentro del bloque propuesto y una "confidence-based prefix selection" (aunque esta última no estaba activa en los experimentos de vLLM). La interacción entre estos componentes se gestiona a través de vLLM, que orquesta la generación y verificación de tokens, el manejo del KV cache y la configuración de los parámetros de especulación como num_speculative_tokens.

python -m vllm.entrypoints.api_server --model google/gemma-4-26B-A4B-it --speculative-config '{"method": "mtp", "model": "google/gemma-4-26B-A4B-it-mtp-assistant", "num_speculative_tokens": 5}'

Fundamentos Teóricos

El concepto de decodificación especulativa se basa en principios de predicción y corrección que tienen raíces en la teoría de control y la informática distribuida. Aunque el artículo no cita un paper fundacional específico para la decodificación especulativa en LLMs, la idea de usar un modelo más pequeño y rápido para predecir y un modelo más grande y preciso para verificar se alinea con el patrón de "predictor-corrector" visto en algoritmos numéricos y sistemas de procesamiento de señales. En el contexto de los LLMs, este patrón fue popularizado por trabajos como "Accelerating Large Language Model Decoding with Speculative Sampling" (Chen et al., 2023) o "Speculative Decoding for Accelerating LLM Inference" (Levi et al., 2023), que formalizaron la aplicación de esta técnica para la generación de texto. Estos papers demuestran cómo la probabilidad de aceptación de tokens disminuye a medida que la longitud de la propuesta aumenta, un fenómeno observado en los resultados experimentales del artículo, donde las tasas de aceptación por posición disminuyen en las posiciones posteriores. Esto subraya la importancia de equilibrar la longitud de la propuesta (N) con la precisión del modelo borrador para maximizar el throughput sin introducir demasiadas correcciones costosas.