¿Por qué la arquitectura de atención ahora define la velocidad de inferencia?

¡Hola Devs! Los workloads agénticos y de contexto largo ya no son excepción—son el estándar. A medida que el contexto crece, la atención consume una porción cada vez mayor del tiempo de inferencia. La forma en que diseñas la atención (no solo la implementas) impacta directamente el rendimiento del modelo en GPUs. Esa es la premisa del co-diseño de modelos de IA: moldear la arquitectura del modelo según cómo la GPU realmente ejecuta.

En este deep dive, analizamos cómo tres parámetros—group size (G), dimensión de cabeza (Hsz) y longitud de secuencia (ISL/KVSL)—afectan la atención densa. Todo basado en fórmulas analíticas (aritmética de formas GEMM) y datos medidos de kernels de prefill y decode con FP8. Para una discusión más amplia sobre diseño de LLMs amigable al hardware, revisa nuestro análisis anterior sobre co-diseño de modelos de IA.

Antes de entrar, aclaremos las dos fases:

  • Prefill: Procesa el prompt completo en paralelo → GEMM-M grande (ISL × G), compute-bound.
  • Decode: Genera un token a la vez → GEMM-M pequeño (G), memory-bound por las lecturas del KV cache en HBM.

Decodificación especulativa puede mover el decode hacia compute-bound aumentando el GEMM-M efectivo. Cada parámetro se analiza por separado para prefill y decode porque sus cuellos de botella difieren.

Diagram illustrating attention head grouping (GQA) and its impact on GPU memory bandwidth during LLM inference Coding Session Visual

Las cuatro guías de co-diseño (con código y matemáticas)

Guía 1: Elige el Group Size (G) para eficiencia en decode—y súbelo

Group size (G) es el número de query heads que comparten una KV head. MHA tiene G=1, GQA tiene G=4,8,16…, MQA tiene G=QH.

Prefill está dominado por ISL, no por G. Intensidad aritmética ≈ 2 × ISL / (1 + 1/G). Con ISL=32K, subir G de 8 a 16 mejora la IA en <6%. El runtime varía <1% entre G=1 y 64.

Decode tiene intensidad aritmética ≈ 2 × G (cuando KVSL ≫ G). Duplicar G duplica la IA del decode y prácticamente reduce el runtime a la mitad (hasta que los overheads dominen).

# Cálculo rápido: intensidad aritmética del decode vs group size
# Asume FP8 (1 byte por elemento)
def decode_arithmetic_intensity(G, KVSL):
    # FLOPs = 4 * DB * QH * KVSL * Hsz (constante en G)
    # Bytes = 2 * DB * KH * Hsz * (G + KVSL)
    # AI = 2 * G * KVSL / (G + KVSL)
    return 2 * G * KVSL / (G + KVSL)

for G in [1, 8, 16, 32]:
    ai = decode_arithmetic_intensity(G, 32768)
    print(f"G={G:2d} -> AI ≈ {ai:.1f} FLOPs/byte")
# Salida: G=1 -> 2.0, G=8 -> 16.0, G=16 -> 32.0, G=32 -> 64.0

Moraleja: G más alto mejora la velocidad del decode y la utilización de la GPU. La decodificación especulativa es otra palanca.

Guía 2: Usa dimensión de cabeza (Hsz) = 128 o 256

Hsz no cambia la intensidad aritmética (duplica FLOPs y bytes). Pero el runtime crece con Hsz porque:

  • Matmul crece en pasos alineados; Hsz=64 aún paga por un tile de 128. Hsz≥512 alcanza el límite de TMEM.
  • Memoria se mueve en unidades de 128 bytes; Hsz debe ser múltiplo de 128.
  • Softmax es independiente de Hsz (opera en la matriz de scores).
# Ops de softmax escalan con longitud de secuencia, no con dimensión de cabeza
def softmax_ops_prefill(PB, QH, ISL):
    return PB * QH * ISL * ISL  # independiente de Hsz

def softmax_ops_decode(DB, QH, KVSL):
    return DB * QH * KVSL      # independiente de Hsz

Moraleja: Hsz=128 o 256 es el punto ideal para alineación con GPU y presupuesto de TMEM.

Guía 3: Reduce el estado KV efectivo cuando sea posible

Prefill escala cuadráticamente con ISL (cada token atiende a todos), mientras decode escala linealmente con KVSL (cada paso lee el KV cache completo). El costo es asimétrico.

Estrategias: compresión de KV cache, atención esparsa/sliding-window, arquitecturas híbridas (ej.: solo algunas capas cargan estado KV global).

Guía 4: Deja que el KH determine la estrategia de paralelismo

Tensor parallelism (TP) divide las attention heads entre GPUs. Límite práctico: TP ≤ KH para que cada GPU tenga al menos un grupo completo (una KV head + sus G query heads). Cuando TP > KH, el estado KV se duplica entre ranks, desperdiciando memoria y ancho de banda.

Para modelos con pocas KV heads (MQA con 1, GQA con 2), el TP se agota rápido. Usa Attention Data Parallelism (ADP) o KV Parallelism (KVP) para atención, más Expert Parallelism (EP) para el FFN MoE. TensorRT-LLM lo implementa como Wide EP (ADP+EP) y Helix Parallelism (KVP+EP).

# Chequeo de paralelismo: asegura que TP no exceda el número de KV heads
def validate_tp(KH, TP):
    if TP > KH:
        raise ValueError(f"TP={TP} excede KH={KH}. Estado KV será duplicado.")
    return True

# Ejemplo: GQA con 2 KV heads
validate_tp(KH=2, TP=2)  # OK
# validate_tp(KH=2, TP=4)  # Lanza error

GPU server rack with high-bandwidth memory handling long-context KV cache for AI model decoding Developer Related Image

Limitaciones y cuidados críticos

  • Prefill y decode no son simétricos. Optimizaciones que ayudan al decode (G alto) casi no mueven el prefill. Siempre perfila ambas fases por separado.
  • Decodificación especulativa cambia el juego. Eleva el GEMM-M efectivo a (1+D)×G, donde D es el número de draft tokens. Cuando es suficientemente grande para llenar tiles de compute, el decode tiende a compute-bound—reduciendo el beneficio relativo de G alto.
  • Overheads dominan en secuencias cortas. Con KVSL=32K, el runtime del decode se aplana más allá de G=16 debido a setup fijo, post-procesamiento y reducción del flash-decoding. Contextos más largos (128K) amortiguan mejor estos costos.
  • TP ≤ KH es límite duro. Violarlo duplica el estado KV y anula las ganancias de escala. Modelos con pocas KV heads deben adoptar ADP/KVP + EP.
  • Hsz=64 es ineficiente. Aún paga por un tile de 128—recibes la mitad del compute por el mismo costo. Evítalo a menos que estés restringido.

Próximos pasos de aprendizaje

  • Experimenta con distintos group sizes en tu modelo: mide el throughput de decode con G=1, 4, 8, 16 en tu GPU objetivo.
  • Perfila prefill y decode por separado con NVIDIA Nsight Systems para identificar si estás compute-bound o memory-bound.
  • Explora Wide EP y Helix Parallelism de TensorRT-LLM para escalar modelos con pocas KV heads.
  • Lee el post complementario sobre atención esparsa (próximamente) para entender cómo la esparsidad cambia el cálculo de intensidad aritmética.
  • Para una perspectiva relacionada sobre integración de IA empresarial, revisa nuestro análisis sobre cómo los agentes de IA están redefiniendo la integración empresarial.

Performance graph showing prefill vs decode runtime scaling with sequence length on NVIDIA GPUs

Conclusión: el co-diseño no es opcional

Las cuatro guías anteriores forman un checklist práctico para desarrolladores de modelos que apuntan a GPUs NVIDIA:

  1. Sube G para eficiencia en decode (prefill es insensible).
  2. Usa Hsz=128 o 256 para alinear con tiles de GPU y transferencias de 128 bytes.
  3. Reduce el estado KV efectivo vía compresión, esparsidad o arquitecturas híbridas.
  4. Casa el paralelismo con el KH—mantén TP ≤ KH y usa Wide EP/Helix para modelos con pocas KV heads.

Con workloads agénticos volviéndose estándar, la arquitectura de atención ya no es un detalle—es la principal palanca de rendimiento de inferencia. Co-diseñar la arquitectura del modelo con la ejecución en GPU en mente es cómo entregas inferencia rápida e interactiva a escala.

Mantente seguro: Al integrar agentes de IA en tu stack, ten en cuenta nuevos riesgos de supply chain como ataques de inyección vía AGENTS.md. La seguridad también necesita co-diseño.

Este contenido fue redactado con la asistencia de herramientas de IA, basándose en fuentes confiables, y fue revisado por nuestro equipo editorial antes de su publicación. No reemplaza el asesoramiento de un profesional especializado.