Phase 07: Transformers Deep Dive

Cache KV, atención flash y optimización de la inferencia

Type: Build

Languages: Python

Prerequisites: Phase 7 · 02 (Self-Attention), Phase 7 · 05 (Full Transformer), Phase 7 · 07 (GPT)

Time: ~75 minutes

El problema

Un descifrador autoregresor ingenuo sí .O(N²)trabajo para generar NTokens: en cada paso recalcula la atención sobre el prefijo completo. Para una respuesta de tokens 4K que es 16M operaciones de atención, la mayoría de ellos redundantes. Cada estado oculto de un token prefijo es determinista una vez calculado.

Además, la atención misma mueve muchos datos. La atención estándar materializa una matriz de puntaje N×N, salida de N×d softmax, salida final N×d demasiadas lecturas y escrituras a HBM. Para N≥2K, la atención se vuelve limitada a la memoria antes de que se convierta en FLOP-bound.

Dos optimizaciones, ambas de Dao et al., empujaron la inferencia fronteriza de "lento" a "rápido":

1. KV cache.Almacenar los vectores K y V de cada token prefijo. la atención de cada nuevo token es una consulta contra las claves almacenadas en caché.O(N²)¿ Qué ?O(N)por cada paso de generación.
2. Flash Attention.El cálculo de la atención se realiza con una tela de tela para que la matriz completa de N×N nunca alcance el HBM. Todo el softmax + matmul ocurre en SRAM. 24× velocidad del reloj de pared en A100; 510× en H100 con FP8.

Para 2026 ambos son universales. Cada pila de inferencias de producción (vLLM, TensorRT-LLM, SGLang, llama.cpp) los asume.

El concepto

!KV cache growth and Flash Attention tiling

Matemáticas de caché de KV

Por capa de decodificación, por token, por cabeza:

Copybytes_per_token_per_layer = 2 * d_head * dtype_size
                          ^
                          K and V

Para un modelo 7B con 32 capas, 32 cabezas, d_head=128, fp16:

Copyper token per layer = 2 * 128 * 2 = 512 bytes
per token (32 layers) = 16 KB
per 32K context = 512 MB

Para Llama 3 70B (80 capas, d_head=128, GQA con 8 capas de KV):

Copyper token per layer = 2 * 8 * 128 * 2 = 4096 bytes (4 KB)
per 32K context = 10.4 GB

Ese 10 GB es por lo que Llama 3 70B en 128K contexto necesita la mayoría de un 40 GB A100 sólo para KV caché en el tamaño de lote 1.

GQA is the KV-cache win.MHA con 64 cabezas sería de 32 GB. MLA comprime aún más.

Arraste las dimensiones y observa el tamaño de la caché. Empuje la longitud de la secuencia o lote hacia arriba y vea qué tan rápido se desvanece más allá de una sola GPU:

La atención de flash el truco de la balsa

Atención estándar:

CopyS = Q @ K^T          (HBM read, N×N, HBM write)
P = softmax(S)       (HBM read, HBM write)
O = P @ V            (HBM read, HBM write)

Tres viajes de ida y vuelta de HBM. En H100, el ancho de banda de HBM es de 3 TB/s; SRAM es de 30 TB/s. Cada viaje de HBM es un factor de 10 de desaceleración frente a mantener todo en el chip.

Atención instantánea:

Copyfor each block of Q (tile size ~128 × 128):
    load Q_tile into SRAM
    for each block of K, V:
        load K_tile, V_tile into SRAM
        compute S_tile = Q_tile @ K_tile^T     (SRAM)
        running softmax aggregation             (SRAM)
        accumulate into O_tile                  (SRAM)
    write O_tile to HBM

Un viaje de HBM por baldosas.O(N²)¿ Qué ?O(N). El pase hacia atrás recalcula algunos valores del pase hacia adelante en lugar de almacenarlos otra ganancia de memoria.

Numerical trick.El funcionamiento de softmax se mantiene (max, sum)La atención flash calcula la salida idéntica a la atención estándar (modulo fp16 no asociativa).

Version evolution:

Flash 4 es solo avanzado al lanzamiento. El entrenamiento todavía utiliza Flash 3.

Descodación especulativa la otra latencia gana

El modelo barato propone N tokens. El modelo grande verifica todos los N en paralelo. Si la verificación acepta k tokens, pagó 1 pase adelante de modelo grande para k generaciones.

2026 incumplimientos:

• EAGLE 2 / Medusa.Tías de proyección integradas que comparten los estados ocultos del verificador. 23x aceleración sin pérdida de calidad.
• Speculative decoding with draft model.2×4 veces más rápido en el hardware de consumo.
• Lookahead decoding.Iteración Jacobi, no se necesita un modelo de borrador. Nicho pero gratis.

Participación continua

Inferencia clásica de lote: esperar a que termine la secuencia más lenta, luego iniciar un nuevo lote.

Batchamiento continuo (en Orca, ahora en vLLM, TensorRT-LLM, SGLang): intercambiar nuevas solicitudes en el lote tan pronto como terminen las viejas.

PagedAttention KV caché como memoria virtual

La función principal de vLLM. El caché KV se asigna en bloques de 16 tokens; una tabla de página mapea las posiciones lógicas de los bloques físicos. Permite compartir KV a través de muestras paralelas (busca de haces, muestreo paralelo), prefijos de intercambio caliente para el caché rápido y memoria de desfragmentación. Mejora de rendimiento 4x sobre asignación contiguosa ingenua.

Construye el mismo

¿ Qué ?code/main.pyImplementamos:

1. Un ingenuo .O(N²)decodificador incremental.
2. ¿ Qué es esto ?O(N)Descriptor de caché KV.
3. Una softmax de tejas que simula el algoritmo de ejecución máxima de Flash Attention.

Paso 1: Caché de KV

pythonCopyclass KVCache:
    def __init__(self, n_layers, n_heads, d_head):
        self.K = [[[] for _ in range(n_heads)] for _ in range(n_layers)]
        self.V = [[[] for _ in range(n_heads)] for _ in range(n_layers)]

    def append(self, layer, head, k, v):
        self.K[layer][head].append(k)
        self.V[layer][head].append(v)

    def read(self, layer, head):
        return self.K[layer][head], self.V[layer][head]

Sencillo: sigue creciendo por token K, V vectores en por capa, por lista de cabeza.

Paso 2: softmax de azulejos

pythonCopydef tiled_softmax_dot(q, K, V, tile=4):
    """Flash-attention-style softmax(qK^T)V with running max/sum."""
    m = float("-inf")
    s = 0.0
    out = [0.0] * len(V[0])
    for start in range(0, len(K), tile):
        k_block = K[start:start + tile]
        v_block = V[start:start + tile]
        scores = [sum(qi * ki for qi, ki in zip(q, k)) for k in k_block]
        new_m = max(m, *scores)
        exp_old = math.exp(m - new_m) if m != float("-inf") else 0.0
        exp_new = [math.exp(sc - new_m) for sc in scores]
        s = s * exp_old + sum(exp_new)
        for j in range(len(out)):
            out[j] = out[j] * exp_old + sum(e * v[j] for e, v in zip(exp_new, v_block))
        m = new_m
    return [o / s for o in out]

Una salida idéntica a bit softmax(qK) Ven un solo disparo, pero en cualquier momento el juego de trabajo es un tile × d_headBloqueo, no el completo N × d_head¿ Qué ?

Paso 3: Comparar la descifrado ingenuo vs caché en la generación de 100 tokens

Cuenta las operaciones de atención.O(N²)= 5050. En caché: O(N)El código imprime ambas cosas.

Usalo

pythonCopy# HuggingFace transformers auto-enables KV cache on decoder-only generate().
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-3B",
    attn_implementation="flash_attention_2",  # use FA3 if Hopper
    torch_dtype="bfloat16",
)
# generate() uses KV cache automatically

Producción de VLLM:

bashCopypip install vllm
vllm serve meta-llama/Llama-3.1-70B-Instruct \
    --tensor-parallel-size 4 \
    --max-model-len 32768 \
    --enable-prefix-caching \
    --kv-cache-dtype fp8

El caché de prefijos en las solicitudes es un gran éxito para 2026 el mismo sistema de instrucciones, ejemplos de pocos disparos o documento de contexto largo reutiliza KV en las llamadas. Para las cargas de trabajo de agentes con instrucciones de herramientas repetidas, el caché de prefijos es rutinariamente 5x ganancia de rendimiento.

Envío

¿ Qué ?outputs/skill-inference-optimizer.mdLa habilidad selecciona la implementación de la atención, la estrategia de caché KV, la cuantificación y la descifrado especulativo para una nueva implementación de inferencias.

Los ejercicios

1. Easy.- ¿ Qué ?code/main.py- Confirmar que los descifradores ingenuos y almacenados en caché producen la misma salida; nota la diferencia de op-count.
2. Medium.Implementar el caché de prefijos: dado un prompt P y varias completas, ejecuta un pase hacia adelante sobre P para llenar el caché KV, luego ramifica por completado.
3. Hard.Implementar un juguete PagedAttention: KV cache en bloques fijos de 16 tokens con una lista libre. Cuando una secuencia termine, devuelva sus bloques a la piscina. Simula 1.000 terminaciones de chat con diferentes longitudes. Compara fragmentación de memoria con asignación contiguosa.

Términos clave

Leer más

• Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness Flash 1.
• Dao (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning Flash 2.
• Shah et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision Flash 3.
• FlashAttention-4 release notes (Dao-AILab, 2026) Blackwell 5 etapas de tubería y el truco de software-exp2; lea el repo README para las advertencias de lanzamiento sólo hacia adelante que menciona esta lección.
• Kwon et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention Papel de la VLLM.
• Leviathan et al. (2023). Fast Inference from Transformers via Speculative Decoding Descripción de especificaciones.
• Li et al. (2024). EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty Documento de EAGLE-1/2 para el enfoque de proyecto integrado que cita la lección.
• Cai et al. (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads el enfoque Medusa mencionado junto con EAGLE.
• vLLM docs — PagedAttention el profundo inmersión canónica en el bloque de 16 tokens y diseño de tabla de páginas.

This free lesson is part of the AI Engineering from Scratch curriculum. Read the full explanation, run the lesson code, and verify the result in the interactive reader or from the repository source.

Browse the complete course catalog or open this lesson on GitHub.