Phase 10 - Lesson 25

Decodificación Especulativa y EAGLE

This lesson includes a graded coding exercise that runs in your browser, unlocked with lifetime access.

Un LLM de frontera generando un token requiere un forward pass completo sobre miles de millones de parámetros. Ese forward pass está masivamente sobredimensionado: la mayor parte del tiempo un modelo mucho más pequeño puede adivinar los siguientes 3-5 tokens correctamente, y el modelo grande solo necesita verificar la suposición. Cuando la suposición es correcta, obtienes 5 tokens por el precio de uno. La decodificación especulativa (speculative decoding) (Leviathan et al. 2023) hizo esto exacto, y EAGLE-3 (2025) elevó las tasas de aceptación a ~4.5 tokens por verificación — una aceleración de 4-5x manteniendo la misma distribución de salida.

Tipo: Build Lenguajes: Python (con numpy) Prerrequisitos: Phase 10 Lesson 12 (Inference Optimization), Phase 10 Lesson 04 (Pre-training Mini-GPT) Tiempo: ~75 minutos

El Problema

El throughput de decodificación para un modelo de la clase de 70B en una H100 es típicamente de 40-80 tokens/segundo. Cada token requiere un forward pass completo leyendo todos los pesos del modelo desde HBM. No puedes hacer el modelo más pequeño sin cambiar su salida. No puedes aumentar el tamaño de lote (batch size) más allá de la memoria. Estás atrapado — a menos que puedas permitir que el modelo genere más de un token por forward pass.

La generación autorregresiva parece inherentemente serial: x_{t+1} = sample(p(· | x_{1:t})). But there is a concurrency opportunity. Si tuvieras un predictor económico que dijera "los próximos 4 tokens son probablemente [a, b, c, d]", podrías verificar las 5 posiciones en un único forward pass del modelo grande y aceptar el prefijo coincidente más largo.

Leviathan, Kalai, Matias (2023, "Fast Inference from Transformers via Speculative Decoding") hicieron esto exacto a través de una ingeniosa regla de aceptación/rechazo que preserva la distribución de muestreo del modelo objetivo. La misma distribución de salida, de 2 a 4 veces más rápido.

El Concepto

La Configuración de Dos Modelos

• Modelo objetivo (Target model) M_p: el modelo grande, lento y de alta calidad del cual realmente deseas obtener muestras. Distribución: p(x).
• Modelo de borrador (Draft model) M_q: un modelo pequeño, rápido y de menor calidad. Distribución: q(x). De 5 a 30 veces más pequeño.

Por paso:

1. El modelo de borrador propone K tokens de forma autorregresiva: x_1, x_2, ..., x_K ~ q.
2. El modelo objetivo ejecuta UN forward pass sobre las K+1 posiciones en paralelo, produciendo p(x_k) para cada token propuesto.
3. Acepta/rechaza cada token de izquierda a derecha a través de la regla modificada de muestreo por rechazo a continuación. Acepta el prefijo coincidente más largo.
4. Si se rechaza algún token, se extrae el reemplazo de la distribución corregida y se detiene. De lo contrario, se extrae un token de bonificación de p(· | x_1...x_K).

Si el borrador coincide perfectamente con el objetivo, obtienes K+1 tokens por forward del objetivo. Si el borrador es incorrecto en la posición 1, obtienes solo 1 token.

La Regla de Exactitud

La decodificación especulativa es demostrablemente equivalente en distribución al muestreo de p. La regla de rechazo:

For each drafted token x_t:
    r ~ Uniform(0, 1)
    if r < p(x_t) / q(x_t):
        accept x_t
    else:
        sample replacement from residual: (p - q)+ / ||(p - q)+||_1
        stop

donde (p - q)+ denota la parte positiva de la diferencia punto a punto. Cuando el borrador y el objetivo coinciden (p ≈ q), la aceptación es casi 1. Cuando difieren, la distribución residual se construye de modo que la muestra general siga siendo exactamente p.

Caso codicioso (Greedy case). Para un muestreo con temperatura=0, solo verifica argmax(p) == x_t. Si es así, acepta; si no, genera argmax(p) y detente.

Aceleración Esperada

Si la tasa de aceptación a nivel de token del modelo de borrador es α, la cantidad esperada de tokens producidos por forward pass del objetivo es:

E[tokens] = (1 - α^{K+1}) / (1 - α)        # K = draft length, α in [0, 1]

A α = 0.8, K = 4: (1 - 0.8^5)/(1 - 0.8) = 3.36 tokens por forward. Un solo forward del objetivo cuesta aproximadamente cost_q * K + cost_p (K pasos de borrador más una verificación del objetivo). Si cost_p >> cost_q * K, la relación de aceleración es de 3.36× / 1 = 3.36× en el throughput.

El único parámetro real es α, que depende completamente de la alineación entre el borrador y el objetivo. Un buen borrador lo es todo.

Entrenamiento del Borrador: Destilación

Un modelo pequeño aleatorio hace un mal borrador. La receta estándar es destilar a partir del objetivo:

1. Elige una arquitectura pequeña (~1B para un objetivo de 70B, ~500M para un objetivo de 7B).
2. Ejecuta el modelo objetivo en un corpus de texto grande; almacena sus distribuciones del siguiente token.
3. Entrena el borrador con divergencia KL contra la distribución del objetivo (not against ground-truth tokens).

El resultado: α típicamente es de 0.6-0.8 en programación y de 0.7-0.85 en chat de lenguaje natural. Aceleraciones de 2 a 3 veces en producción.

EAGLE: Tree Drafting + Reutilización de Características (Feature Reuse)

Li, Wei, Zhang, Zhang (2024, "EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty") observaron dos ineficiencias en la decodificación especulativa estándar:

1. El borrador realiza K pasos seriales, cada uno recorriendo toda la red. Pero el borrador podría reutilizar las características (hidden states) del objetivo de la verificación más reciente: el objetivo ya calculó representaciones ricas que el borrador está volviendo a derivar desde cero.
2. El borrador genera una cadena lineal. Si el borrador pudiera generar un árbol de candidatos (cada nodo con múltiples suposiciones), el único forward pass del objetivo podría verificar múltiples rutas candidatas en paralelo a través de una máscara de atención en árbol (tree attention mask) y elegir la rama aceptada más larga.

Cambios en EAGLE-1:

• Entrada del borrador = estado oculto final del objetivo en la posición t, no tokens brutos.
• Arquitectura del borrador = 1 capa de decodificador transformer (no un modelo pequeño separado).
• Salida = árbol de K = 4-8 candidatos por profundidad, profundidad de 4-6.

EAGLE-2 (2024) agrega topología de árbol dinámica: el árbol se ensancha donde el borrador no está seguro y se mantiene estrecho donde tiene confianza. Eleva la α_effective sin aumentar el costo de verificación.

EAGLE-3 (Li et al. 2025, "EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test") elimina la dependencia fija de características de la capa superior y entrena el borrador con una nueva pérdida de "simulación en tiempo de prueba" (test-time simulation loss); el borrador se entrena en salidas que coinciden con la distribución de tiempo de prueba del objetivo en lugar de la distribución de entrenamiento forzada (teacher-forced). La tasa de aceptación aumenta de 0.75 (EAGLE-2) a 0.82 (EAGLE-3), y la media de tokens por verificación de 3.0 a 4.5.

Verificación de Atención en Árvore

Cuando el borrador genera un árbol, el modelo objetivo lo verifica en un solo forward pass usando una máscara de atención en árbol (tree attention mask): una máscara causal que codifica la topología del árbol en lugar de una línea pura. Cada token atiende solo a sus ancestros en el árbol. El paso de verificación sigue siendo un forward, una multiplicación de matrices (matmul); la máscara topológica solo cuesta unas pocas entradas KV adicionales.

        root
       /    \
      a      b
     / \    / \
    c  d   e   f

Si a, b son candidatos competidores para el primer token y c, d, e, f son candidatos para el segundo token, las seis posiciones se verifican en un solo forward pass. La salida es el prefijo más largo a lo largo de cualquier ruta aceptada.

Cuándo Conviene y Cuándo No

Cuándo conviene (Wins):

• Chat / completado con texto predecible (código, inglés común, salida estructurada). α es alto.
• Escenarios con computación de GPU no utilizada durante la decodificación (fase limitada por memoria/memory-bound). El tree drafting utiliza los FLOPs disponibles.

Cuándo no conviene / Sin ganancias (Loses / no win):

• Salidas altamente estocásticas (escritura creativa a alta temperatura). α cae hacia 1/|vocab|.
• Servido por lotes (batch serving) con muy alta concurrencia: el procesamiento por lotes ya llena las FLOPs, queda poco espacio para la verificación en árbol.
• Modelos objetivo muy pequeños donde el borrador no es mucho más pequeño.

Los entornos de producción suelen reportar una aceleración de 2 a 3 veces en tiempo real (wall-clock) para chat, de 3 a 5 veces en generación de código y casi cero en escritura creativa.

Desarrolla

code/main.py:

• Una referencia speculative_decode(target, draft, prompt, K, temperature) que implementa la regla de rechazo exacta y verifica que preserve la distribución del objetivo (KL empírica < 0.01 frente al muestreo simple del objetivo).
• Un generador de borradores en árbol (tree drafter) de estilo EAGLE que construye un árbol de profundidad K con ramificación top-p.
• Un constructor de máscara de atención en árbol que produce el patrón causal correcto para un verificador.
• Un arnés de tasa de aceptación que ejecuta ambos en un LM diminuto (destilando un GPT-2-small a partir de un objetivo GPT-2-medium).

def speculative_step(p_target, q_draft, K, temperature=1.0):
    """One round of speculative decoding. Returns list of accepted tokens."""
    # 1. Draft K tokens
    draft_tokens = []
    q_probs = []
    state = draft_state_init()
    for _ in range(K):
        probs = softmax(q_draft(state) / temperature)
        t = np.random.choice(len(probs), p=probs)
        draft_tokens.append(t)
        q_probs.append(probs[t])
        state = draft_step(state, t)

    # 2. Target computes p at every drafted position + 1 extra
    p_probs_all = target_forward_batched(p_target, draft_tokens, temperature)

    # 3. Accept/reject left-to-right
    accepted = []
    for k, tok in enumerate(draft_tokens):
        r = np.random.uniform()
        if r < p_probs_all[k][tok] / q_probs[k]:
            accepted.append(tok)
        else:
            residual = np.maximum(p_probs_all[k] - q_probs[k], 0)
            residual /= residual.sum()
            accepted.append(np.random.choice(len(residual), p=residual))
            return accepted
    # 4. All K accepted → sample bonus token from target
    accepted.append(np.random.choice(len(p_probs_all[-1]), p=p_probs_all[-1]))
    return accepted

Utilízalo

• vLLM y SGLang incorporan decodificación especulativa nativa. Parámetros: --speculative_model, --num_speculative_tokens. El soporte para EAGLE-2/3 se habilita mediante el parámetro --spec_decoding_algorithm eagle.
• NVIDIA TensorRT-LLM soporta árboles Medusa y EAGLE de forma nativa.
• Modelos de borrador de referencia: Qwen/Qwen3-0.6B-spec (borradores para Qwen3-32B), meta-llama/Llama-3.2-1B-Instruct-spec (borradores para 70B).
• Medusa heads (Cai et al. 2024, "Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads"): en lugar de un modelo de borrador, agrega K cabezas de predicción paralelas al propio objetivo. Más fácil de implementar, tasa de aceptación ligeramente menor que EAGLE.

Ponlo en Producción

Esta lección produce outputs/skill-speculative-tuning.md — una habilidad que perfila la carga de trabajo de un modelo objetivo y elige: modelo de borrador, K (longitud de borrador), ancho del árbol, temperatura y cuándo recurrir a la decodificación simple.

Ejercicios

1. Implementa la regla de rechazo exacta y verifícala empíricamente. Ejecuta 10K muestras a través de speculative_decode y a través de muestreo simple del objetivo; calcula la distancia de variación total (TV distance) entre las dos distribuciones de salida. Debería ser < 0.01.

2. Calcula la fórmula de aceleración. Dado un α y K fijos, grafica los tokens esperados por forward del objetivo. Encuentra el K óptimo para α ∈ {0.5, 0.7, 0.9}.

3. Entrena un borrador diminuto. Toma un objetivo GPT-2 de 124M y destila un borrador GPT-2 de 30M en 100M de tokens con pérdida KL. Mide α en texto de prueba retenido. Esperado: 0.6-0.7.

4. Implementa el tree drafting estilo EAGLE. En lugar de una cadena, haz que el borrador genere las 3 mejores ramas en cada profundidad. Construye la máscara de atención en árbol. Verifica que el objetivo acepte la rama correcta más larga.

5. Mide los modos de falla. Ejecuta la decodificación especulativa a temperatura=1.5 (alta estocasticidad). Muestra que α colapsa y el algoritmo es más lento que la decodificación simple debido a la sobrecarga del borrador.

Términos Clave

Lecturas Adicionales

• Leviathan, Kalai, Matias, 2023 — "Fast Inference from Transformers via Speculative Decoding" — la regla de rechazo exacta y el análisis de la aceleración teórica
• Chen, Borgeaud, Irving et al., 2023 — "Accelerating Large Language Model Decoding with Speculative Sampling" — artículo concurrente sobre muestreo especulativo de DeepMind
• Cai, Li, Geng, Wang, Wang, Zhu, Dao, 2024 — "Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads" — alternativa de cabezas paralelas a un modelo de borrador
• Li, Wei, Zhang, Zhang, 2024 — "EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty" — reutilización de características y tree drafting
• Li et al., 2024 — "EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees" — topología de árbol dinámica
• Li et al., 2025 — "EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test" — emparejamiento entre tiempo de entrenamiento y tiempo de prueba
• Fu, Haotian, Peng et al., 2024 — "Break the Sequential Dependency of LLM Inference Using Lookahead Decoding" — decodificación Jacobi/lookahead, una alternativa sin especulador