Google utiliza tecnología de IA para traducir contenido a tu idioma preferido. Las traducciones realizadas con IA pueden contener errores.

Acelera la atención en la GPU con cuDNN y TransformerEngine

1. Introducción

Ruta de aprendizaje de Jax en GPU. Lab 5: Atención en la GPU

En este codelab, te concentrarás en el cálculo de la atención en las GPUs de NVIDIA. Comenzarás con una implementación simple, la reemplazarás por la atención fusionada integrada de JAX y, luego, activarás el backend de cuDNN de NVIDIA para la misma operación. Al final, sabrás qué cambia en cada nivel y cuánto más rápido se vuelven las rutas optimizadas para GPU a medida que crece el problema.

Actividades

Implementa la atención de producto punto escalado desde cero con operaciones básicas de JAX.
Reemplázala por jax.nn.dot_product_attention, el kernel fusionado integrado de JAX.
Fuerza el backend de cuDNN con implementation="cudnn" y agrega el enmascaramiento causal.
Realiza un barrido de la longitud de la secuencia y el tamaño del lote para ver dónde se amortizan los kernels fusionados.
Compara las formas de atención de varios encabezados MHA, GQA y MQA.
Realiza una comparativa de la atención de TransformerEngine de NVIDIA y, luego, inspecciona sus instrucciones de código FP8.

Requisitos

Un proyecto de Google Cloud con la facturación habilitada y créditos del taller o una reserva que cubra el uso de la GPU.
Cuota para al menos 2 GPUs L4 de NVIDIA en la región elegida (cómo verificar la cuota de GPU).
Haber completado los codelabs del 1 al 4 o tener un entorno de GPU de JAX equivalente habilitado para CUDA.
Conocimientos sobre la multiplicación de matrices y softmax. No se requiere experiencia en CUDA ni cuDNN.

Tiempo estimado para completar esta actividad: 60 minutos.

Cómo funciona la atención

La atención de producto punto escalado toma tres entradas: Query, Key y Value, y calcula lo siguiente:

Attention(Q, K, V) = softmax(Q K^T / sqrt(d_k)) V

El factor de escala 1 / sqrt(d_k) evita que los productos punto crezcan demasiado a medida que aumenta la dimensión del encabezado, lo que llevaría a softmax a regiones en las que sus gradientes son muy pequeños.

Cuatro pasos, cada uno alimentando al siguiente:

Score : Q @ K.T: El producto punto mide cuánto debe atender cada posición de consulta a cada posición de clave.
Scale + Softmax : softmax(scores / sqrt(d)): El escalamiento evita que desaparezca el gradiente, y softmax convierte las puntuaciones en pesos de atención que suman 1.
Attend : weights @ V: La suma ponderada de los vectores de valor produce el resultado para cada posición de consulta.
Output : Tiene la misma forma que Q: Cada posición de consulta ahora contiene información de las posiciones a las que atendió.

2. Antes de comenzar

Selecciona tu proyecto

En la consola de Google Cloud, selecciona o crea un proyecto con la facturación habilitada.

Abre Cloud Shell

Haz clic en Activar Cloud Shell (el ícono de terminal en la parte superior derecha de la consola) para iniciar una sesión de Cloud Shell y, luego, apunta a tu proyecto:

gcloud config set project <YOUR_PROJECT_ID>

Aprovisiona el entorno de GPU

Ejecuta lo siguiente en Cloud Shell. En el codelab 1, se explica cada comando en detalle, incluidos los requisitos de cuota y zona de la GPU.

gcloud services enable \
  container.googleapis.com \
  compute.googleapis.com \
  iam.googleapis.com \
  cloudresourcemanager.googleapis.com \
  logging.googleapis.com \
  monitoring.googleapis.com

git clone https://github.com/Google-Cloud-AI/partner-ai-nvidia.git
cd partner-ai-nvidia/05-workshops/jax-on-gpu/terraform

cp terraform.tfvars.example terraform.tfvars

Edita terraform.tfvars y establece project_id = "". Luego, aprovisiona el clúster e implementa JupyterLab:

terraform init
terraform apply
$(terraform output -raw get_credentials_command)

cd ..
kubectl apply -f deploy/jupyter.yaml

terraform apply tarda unos 12 minutos. Cuando termine, espera el Pod y el LoadBalancer y, luego, lee el token único de JupyterLab del registro del Pod:

kubectl get pod jax-jupyter -w         # wait for Running, then Ctrl+C
kubectl get svc jax-jupyter-svc -w     # wait for EXTERNAL-IP, then Ctrl+C
kubectl logs jax-jupyter | grep -o 'token=[a-z0-9]*' | head -1

Abre http://:8884, pega el token y crea un nuevo notebook de Python 3 en /workspace. Cada bloque de código de este codelab se incluye en una celda de ese notebook.

Instala lo que necesita este codelab

!pip install --quiet matplotlib flax

Configura y verifica la GPU

Importa JAX y confirma que el backend predeterminado es una GPU. Esta celda también define block_tree, show_table y show_bars, los asistentes que usa cada paso posterior para esperar el trabajo del dispositivo y renderizar los resultados.

import os
os.environ["LD_LIBRARY_PATH"] = "/usr/local/nvidia/lib64:" + os.environ.get("LD_LIBRARY_PATH", "")

import html
import math
import time
from functools import partial

from IPython.display import HTML, display
import matplotlib.pyplot as plt
import numpy as np

import jax
import jax.numpy as jnp


devices = jax.devices()
gpu_devices = [d for d in devices if d.platform == "gpu"]
device = gpu_devices[0] if gpu_devices else None

print(f"JAX version:     {jax.__version__}")
print(f"Default backend: {jax.default_backend()}")
print(f"Devices:         {devices}")

assert gpu_devices, f"This lab assumes a GPU backend. Available devices: {devices}"
print(f"Using GPU:       {device}")


def block_tree(tree):
    """Wait until a PyTree of JAX arrays is ready on device."""
    return jax.block_until_ready(tree)


def show_table(headers, rows, title=None, aligns=None):
    """Render rows as an HTML table."""
    aligns = aligns or ["left"] * len(headers)
    parts = ["<div style='font-family: system-ui; max-width: 980px;'>"]
    if title:
        parts.append(f"<h4 style='margin: 0 0 8px 0;'>{html.escape(title)}</h4>")
    parts.append("<table style='border-collapse: collapse; width: 100%; font-size: 13px;'>")
    parts.append("<thead><tr>")
    for h, a in zip(headers, aligns):
        parts.append(
            f"<th style='text-align:{a}; border-bottom:1px solid #d0d7de; padding:6px;'>"
            f"{html.escape(str(h))}</th>"
        )
    parts.append("</tr></thead><tbody>")
    for row in rows:
        parts.append("<tr>")
        for cell, a in zip(row, aligns):
            parts.append(
                f"<td style='text-align:{a}; border-bottom:1px solid #eef1f4; padding:6px;'>"
                f"{html.escape(str(cell))}</td>"
            )
        parts.append("</tr>")
    parts.append("</tbody></table></div>")
    display(HTML("".join(parts)))


def show_bars(rows, title, unit="", lower_is_better=False):
    """Render (label, value) pairs as a horizontal bar chart in HTML."""
    max_value = max(float(value) for _, value in rows) or 1.0
    color = "#1a7f37" if not lower_is_better else "#0969da"
    parts = ["<div style='font-family: Arial, sans-serif; max-width: 760px;'>"]
    parts.append(f"<h4 style='margin: 0 0 8px 0;'>{html.escape(title)}</h4>")
    for label, value in rows:
        width = max(3, 100 * float(value) / max_value)
        parts.append(
            "<div style='display:grid; grid-template-columns: 190px 1fr 130px; gap: 8px; "
            "align-items:center; margin: 6px 0;'>"
            f"<div style='font-size:13px;'>{html.escape(str(label))}</div>"
            "<div style='background:#f6f8fa; border-radius:6px; overflow:hidden; height:22px;'>"
            f"<div style='height:22px; width:{width:.1f}%; background:{color};'></div></div>"
            f"<div style='font-size:13px; font-variant-numeric: tabular-nums;'>{float(value):,.1f} {html.escape(unit)}</div>"
            "</div>"
        )
    parts.append(
        f"<div style='font-size:12px; color:#57606a;'>"
        f"{'Lower' if lower_is_better else 'Higher'} is better.</div></div>"
    )
    display(HTML("".join(parts)))

Deberías ver una versión de JAX, gpu como backend predeterminado, una lista de dispositivos CUDA y la GPU que usa el resto del codelab.

3. Compila arrays de prueba Q, K y V

A los kernels de atención fusionados les importa cómo se disponen los arrays de entrada en la memoria, por lo que debes compilar Q, K y V en el diseño que espera JAX antes de escribir cualquier código de atención. jax.nn.dot_product_attention espera el siguiente diseño:

Atenuar	Significado	Nuestro valor predeterminado
B	Tamaño del lote	4
T	Longitud de la secuencia de consulta	128
S	Longitud de la secuencia de clave/valor	128 (igual que T para la autoatención)
N	Cantidad de encabezados de atención	8
H	Dimensión por encabezado	64

Comienzas en float32 para la implementación simple, inicializando con valores aleatorios.

BATCH = 4
SEQ_LEN = 128
NUM_HEADS = 8
HEAD_DIM = 64

key = jax.random.key(0)
k1, k2, k3 = jax.random.split(key, 3)

q = jax.random.normal(k1, (BATCH, SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
k = jax.random.normal(k2, (BATCH, SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
v = jax.random.normal(k3, (BATCH, SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)

q, k, v = jax.device_put((q, k, v), device)

show_table(
    ["Array", "Shape", "Dtype", "Layout"],
    [
        ("Q (query)", q.shape, q.dtype, "(B, T, N, H)"),
        ("K (key)", k.shape, k.dtype, "(B, S, N, H)"),
        ("V (value)", v.shape, v.dtype, "(B, S, N, H)"),
    ],
    title="Attention inputs on GPU",
)

Deberías ver una tabla con una fila por array, cada una con la forma (4, 128, 8, 64) y el tipo de datos float32. La llamada jax.device_put fija los tres arrays a la GPU que seleccionaste en la celda de configuración, por lo que nada en las comparativas posteriores mide una transferencia de host.

4. Implementa la atención desde cero

Antes de usar un kernel fusionado, escribe la fórmula directamente para que puedas ver exactamente qué reemplaza el kernel. Esta implementación sigue los cuatro pasos en orden: calcula las puntuaciones, la escala, softmax y multiplica por los valores. Transpones las dimensiones del encabezado y la secuencia para que matmul opere sobre el eje de la secuencia.

Funciona correctamente, pero emite tres lanzamientos de kernel de GPU separados para la puntuación matmul, softmax y el valor matmul, y materializa la matriz de peso de atención (B, N, T, S) completa en la memoria de la GPU.

def naive_attention(q, k, v):
    """Scaled dot-product attention from scratch."""
    scale = 1.0 / math.sqrt(q.shape[-1])

    # (B, T, N, H) to (B, N, T, H) so matmul runs over the T/S axis
    q_t = jnp.transpose(q, (0, 2, 1, 3))
    k_t = jnp.transpose(k, (0, 2, 1, 3))
    v_t = jnp.transpose(v, (0, 2, 1, 3))

    # Score: (B, N, T, H) @ (B, N, H, S) to (B, N, T, S)
    scores = jnp.matmul(q_t, jnp.transpose(k_t, (0, 1, 3, 2))) * scale
    weights = jax.nn.softmax(scores, axis=-1)

    # Attend: (B, N, T, S) @ (B, N, S, H) to (B, N, T, H)
    out_t = jnp.matmul(weights, v_t)

    # Back to (B, T, N, H)
    return jnp.transpose(out_t, (0, 2, 1, 3))


naive_out = block_tree(naive_attention(q, k, v))

show_table(
    ["", "Value"],
    [
        ("Output shape", str(naive_out.shape)),
        ("Output dtype", str(naive_out.dtype)),
    ],
    title="Naive attention",
)

Deberías ver la forma de salida (4, 128, 8, 64) y el tipo de datos float32, que es la misma forma que Q, que es lo que prometió el paso Output de la fórmula.

5. Cambia a dot_product_attention

jax.nn.dot_product_attention fusiona los pasos de puntuación, escala, softmax y atención en una sola operación. Luego, JAX y XLA pueden optimizar el patrón de acceso a la memoria. En particular, pueden evitar materializar la matriz de peso de atención completa cuando la secuencia es larga.

Con el valor predeterminado implementation=None, JAX elige automáticamente el mejor backend disponible. En una GPU con cuDNN disponible y entradas compatibles, es posible que ya use cuDNN. En otro hardware, vuelve a XLA.

sdpa_out = block_tree(jax.nn.dot_product_attention(q, k, v))

max_diff = float(jnp.max(jnp.abs(naive_out - sdpa_out)))

show_table(
    ["", "Value"],
    [
        ("Output shape", str(sdpa_out.shape)),
        ("Output dtype", str(sdpa_out.dtype)),
        ("Max |naive − SDPA|", f"{max_diff:.2e}"),
        ("Outputs close (atol=1e-3)", str(bool(jnp.allclose(naive_out, sdpa_out, atol=1e-3)))),
    ],
    title="JAX SDPA vs naive",
)

Deberías ver la misma forma y el mismo tipo de datos que la versión simple, una pequeña diferencia absoluta máxima y True para la verificación de cercanía atol=1e-3.

6. Fuerza el backend fusionado de cuDNN

Dejar que JAX elija el backend es conveniente, pero desde afuera no puedes saber qué kernel se ejecutó realmente. Si configuras implementation="cudnn", se fuerzan los kernels de atención fusionados de cuDNN de NVIDIA. Estos son kernels de GPU optimizados manualmente que fusionan todo el cálculo de atención en un solo lanzamiento de kernel con patrones de acceso a la memoria optimizados.

Ten en cuenta que la atención fusionada de cuDNN tiene requisitos de hardware y forma, como la capacidad de procesamiento de la GPU, que debe ser >= 8.0 (Ampere o posterior) o float16 o bfloat16 como tipos de datos de entrada. Si no se cumplen los requisitos y configuras implementation="cudnn", JAX genera un error en lugar de volver a la versión anterior de forma silenciosa. Por eso, el código que aparece a continuación se convierte en bfloat16 y ajusta la llamada en try/except: establece HAS_CUDNN_SDPA para que cada paso posterior sepa si la ruta de cuDNN está disponible en esta máquina.

q_bf16 = q.astype(jnp.bfloat16)
k_bf16 = k.astype(jnp.bfloat16)
v_bf16 = v.astype(jnp.bfloat16)

HAS_CUDNN_SDPA = False

try:
    cudnn_out = block_tree(
        jax.nn.dot_product_attention(q_bf16, k_bf16, v_bf16, implementation="cudnn")
    )
    HAS_CUDNN_SDPA = True

    xla_bf16_out = block_tree(
        jax.nn.dot_product_attention(q_bf16, k_bf16, v_bf16, implementation="xla")
    )
    max_diff = float(jnp.max(jnp.abs(
        cudnn_out.astype(jnp.float32) - xla_bf16_out.astype(jnp.float32)
    )))

    show_table(
        ["", "Value"],
        [
            ("Output shape", str(cudnn_out.shape)),
            ("Output dtype", str(cudnn_out.dtype)),
            ("Max |cuDNN − XLA| (both bf16)", f"{max_diff:.2e}"),
            ("Outputs close (rtol=1e-2, atol=1e-2)", str(bool(jnp.allclose(cudnn_out, xla_bf16_out, rtol=1e-2, atol=1e-2)))),
        ],
        title="cuDNN fused attention",
    )

except Exception as e:
    print(f"cuDNN SDPA not available on this GPU: {e}")
    print("Continuing with XLA backend only.")

Si tu GPU no puede ejecutar cuDNN SDPA, el código imprime el motivo y el codelab continúa en el backend de XLA.

7. Agrega el enmascaramiento causal

En los modelos autorregresivos (decodificadores de estilo GPT), cada posición solo puede atender a las posiciones anteriores. Si configuras is_causal=True, se aplica esta máscara triangular inferior dentro del kernel fusionado, por lo que no necesitas compilar una matriz de máscara por tu cuenta.

El siguiente código ejecuta la atención dos veces, con y sin la máscara, y compara dos posiciones para mostrar qué cambió la máscara.

causal_out = block_tree(
    jax.nn.dot_product_attention(q, k, v, is_causal=True)
)

# With causal masking, the last position attends to all positions.
# The first position attends only to itself.
nocausal_out = block_tree(
    jax.nn.dot_product_attention(q, k, v, is_causal=False)
)

# First position should differ
first_pos_diff = float(jnp.max(jnp.abs(causal_out[:, 0] - nocausal_out[:, 0])))
# Last position should be the same
last_pos_diff = float(jnp.max(jnp.abs(causal_out[:, -1] - nocausal_out[:, -1])))

show_table(
    ["Position", "Max diff (causal vs full)", "Expected"],
    [
        ("First (t=0)", f"{first_pos_diff:.4f}", "Large — causal restricts to self only"),
        ("Last (t=T-1)", f"{last_pos_diff:.2e}", "~0 — attends to all positions either way"),
    ],
    title="Causal masking effect on attention output",
)

Deberías ver una gran diferencia en la primera posición y una diferencia casi nula en la última. Esa asimetría es la máscara con la posición 0 que pierde todas las claves, excepto la suya, mientras que la posición final ya podía ver toda la secuencia.

8. Ejecuta algunas comparativas

Ahora tienes tres formas de calcular la misma función. Cuál usar depende por completo de la forma del problema y puedes evaluarlo.

Mide el tiempo de las variantes en la forma predeterminada

def benchmark_attention(fn, q, k, v, warmup=3, repeats=50):
    """Time an attention function. Returns median milliseconds per call."""
    jit_fn = jax.jit(fn)

    for _ in range(warmup):
        block_tree(jit_fn(q, k, v))

    times = []
    for _ in range(repeats):
        start = time.perf_counter()
        block_tree(jit_fn(q, k, v))
        times.append((time.perf_counter() - start) * 1000)

    return np.median(times)


t_naive = benchmark_attention(naive_attention, q, k, v)
t_sdpa = benchmark_attention(
    lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
    q, k, v,
)

results = [
    ("Naive (matmul + softmax + matmul)", f"{t_naive:.2f}"),
    ("SDPA (XLA, float32)", f"{t_sdpa:.2f}"),
]
bar_data = [
    ("Naive", t_naive),
    ("SDPA XLA f32", t_sdpa),
]

if HAS_CUDNN_SDPA:
    t_sdpa_bf16 = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
        q_bf16, k_bf16, v_bf16,
    )
    t_cudnn = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="cudnn"),
        q_bf16, k_bf16, v_bf16,
    )
    results.append(("SDPA (XLA, bfloat16)", f"{t_sdpa_bf16:.2f}"))
    results.append(("SDPA (cuDNN, bfloat16)", f"{t_cudnn:.2f}"))
    bar_data.append(("SDPA XLA bf16", t_sdpa_bf16))
    bar_data.append(("SDPA cuDNN bf16", t_cudnn))

show_table(
    ["Implementation", "Median ms/call"],
    results,
    title=f"Attention timing — B={BATCH}, T={SEQ_LEN}, N={NUM_HEADS}, H={HEAD_DIM}",
    aligns=["left", "right"],
)
show_bars(bar_data, "Attention latency (ms per call)", "ms", lower_is_better=True)

Para este tamaño de problema pequeño, todas las implementaciones son muy rápidas y, en su mayoría, están limitadas por la sobrecarga, por lo que la versión simple es competitiva. cuDNN bf16 debería ser un poco más rápido, pero la ventaja de la atención fusionada suele ser más clara en longitudes de secuencia más largas, en las que es importante evitar la matriz de atención completa.

Realiza un barrido de la longitud de la secuencia

La ventaja de los kernels de atención fusionados aumenta con la longitud de la secuencia. La implementación simple materializa una matriz de atención (B, N, T, S) en la memoria de la GPU, que es una memoria O(T²). Los kernels fusionados, como cuDNN FlashAttention, dividen el cálculo en mosaicos para que nunca materialicen la matriz completa, lo que mantiene la memoria en O(T).

El siguiente barrido mide el tiempo de cada implementación en longitudes de secuencia de 64 a 1024.

SEQ_LENS = [64, 128, 256, 512, 1024]

sweep_results = []
for sl in SEQ_LENS:
    rk = jax.random.key(sl)
    rk1, rk2, rk3 = jax.random.split(rk, 3)

    q_s = jax.random.normal(rk1, (BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    k_s = jax.random.normal(rk2, (BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    v_s = jax.random.normal(rk3, (BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    q_s, k_s, v_s = jax.device_put((q_s, k_s, v_s), device)

    q_sb = q_s.astype(jnp.bfloat16)
    k_sb = k_s.astype(jnp.bfloat16)
    v_sb = v_s.astype(jnp.bfloat16)

    row = {"seq_len": sl}

    row["naive_ms"] = benchmark_attention(
        naive_attention,
        q_s, k_s, v_s,
        warmup=2,
        repeats=20,
    )

    row["sdpa_xla_f32_ms"] = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
        q_s, k_s, v_s,
        warmup=2,
        repeats=20,
    )

    row["sdpa_xla_bf16_ms"] = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
        q_sb, k_sb, v_sb,
        warmup=2,
        repeats=20,
    )

    if HAS_CUDNN_SDPA:
        row["cudnn_bf16_ms"] = benchmark_attention(
            lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="cudnn"),
            q_sb, k_sb, v_sb,
            warmup=2,
            repeats=20,
        )

    sweep_results.append(row)


headers = ["Seq len", "Naive (ms)", "SDPA XLA f32 (ms)", "SDPA XLA bf16 (ms)"]
if HAS_CUDNN_SDPA:
    headers.append("cuDNN bf16 (ms)")

table_rows = []
for r in sweep_results:
    row = [
        r["seq_len"],
        f"{r['naive_ms']:.2f}",
        f"{r['sdpa_xla_f32_ms']:.2f}",
        f"{r['sdpa_xla_bf16_ms']:.2f}",
    ]

    if HAS_CUDNN_SDPA:
        row.append(f"{r['cudnn_bf16_ms']:.2f}")

    table_rows.append(row)

show_table(
    headers,
    table_rows,
    title=f"Sequence-length sweep — B={BATCH}, N={NUM_HEADS}, H={HEAD_DIM}",
    aligns=["right"] * len(headers),
)

El barrido tarda un tiempo, ya que cada longitud de secuencia compila tres o cuatro implementaciones separadas antes de medirlas. Deberías terminar con una fila de tabla por longitud de secuencia de 64 a 1024.

En este gráfico, se compara cómo cambia la latencia de atención con la longitud de la secuencia para la implementación simple, JAX/XLA SDPA en float32 y bfloat16, y la atención fusionada de cuDNN en bfloat16.

fig, ax = plt.subplots(figsize=(8, 5))
seq_lens = [r["seq_len"] for r in sweep_results]

ax.plot(
    seq_lens,
    [r["naive_ms"] for r in sweep_results],
    "o-",
    label="Naive",
    color="#d1242f",
)

ax.plot(
    seq_lens,
    [r["sdpa_xla_f32_ms"] for r in sweep_results],
    "s-",
    label="SDPA XLA f32",
    color="#0969da",
)

ax.plot(
    seq_lens,
    [r["sdpa_xla_bf16_ms"] for r in sweep_results],
    "d-",
    label="SDPA XLA bf16",
    color="#8250df",
)

if HAS_CUDNN_SDPA:
    ax.plot(
        seq_lens,
        [r["cudnn_bf16_ms"] for r in sweep_results],
        "^-",
        label="cuDNN bf16",
        color="#1a7f37",
    )

ax.set_xlabel("Sequence length")
ax.set_ylabel("Median ms per call")
ax.set_title("Attention latency vs sequence length")
ax.legend()
ax.grid(True, alpha=0.25)
ax.set_xticks(seq_lens)

fig.tight_layout()
plt.show()

Deberías ver una línea por implementación. Y deberías ver la brecha que crece a medida que la secuencia se alarga.

Realiza un barrido del tamaño del lote

Los lotes más grandes amortizan la sobrecarga de lanzamiento del kernel y mejoran el uso de la GPU hasta el punto en que la memoria de la GPU se convierte en el cuello de botella. El siguiente barrido mantiene la longitud de la secuencia fija en 256 y varía el tamaño del lote.

BATCH_SIZES = [1, 2, 4, 8, 16]
SWEEP_SEQ = 256

batch_results = []
for bs in BATCH_SIZES:
    rk = jax.random.key(bs + 100)
    rk1, rk2, rk3 = jax.random.split(rk, 3)

    q_b = jax.random.normal(rk1, (bs, SWEEP_SEQ, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    k_b = jax.random.normal(rk2, (bs, SWEEP_SEQ, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    v_b = jax.random.normal(rk3, (bs, SWEEP_SEQ, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    q_b, k_b, v_b = jax.device_put((q_b, k_b, v_b), device)

    q_bb = q_b.astype(jnp.bfloat16)
    k_bb = k_b.astype(jnp.bfloat16)
    v_bb = v_b.astype(jnp.bfloat16)

    row = {"batch": bs}

    row["sdpa_xla_f32_ms"] = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
        q_b, k_b, v_b,
        warmup=2,
        repeats=20,
    )

    row["sdpa_xla_bf16_ms"] = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
        q_bb, k_bb, v_bb,
        warmup=2,
        repeats=20,
    )

    if HAS_CUDNN_SDPA:
        row["cudnn_bf16_ms"] = benchmark_attention(
            lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="cudnn"),
            q_bb, k_bb, v_bb,
            warmup=2,
            repeats=20,
        )

    batch_results.append(row)


headers = ["Batch size", "SDPA XLA f32 (ms)", "SDPA XLA bf16 (ms)"]
if HAS_CUDNN_SDPA:
    headers.append("cuDNN bf16 (ms)")

table_rows = []
for r in batch_results:
    row = [
        r["batch"],
        f"{r['sdpa_xla_f32_ms']:.2f}",
        f"{r['sdpa_xla_bf16_ms']:.2f}",
    ]

    if HAS_CUDNN_SDPA:
        row.append(f"{r['cudnn_bf16_ms']:.2f}")

    table_rows.append(row)

show_table(
    headers,
    table_rows,
    title=f"Batch-size sweep — T={SWEEP_SEQ}, N={NUM_HEADS}, H={HEAD_DIM}",
    aligns=["right"] * len(headers),
)


fig, ax = plt.subplots(figsize=(8, 5))
batches = [r["batch"] for r in batch_results]

ax.plot(
    batches,
    [r["sdpa_xla_f32_ms"] for r in batch_results],
    "s-",
    label="SDPA XLA f32",
    color="#0969da",
)

ax.plot(
    batches,
    [r["sdpa_xla_bf16_ms"] for r in batch_results],
    "d-",
    label="SDPA XLA bf16",
    color="#8250df",
)

if HAS_CUDNN_SDPA:
    ax.plot(
        batches,
        [r["cudnn_bf16_ms"] for r in batch_results],
        "^-",
        label="cuDNN bf16",
        color="#1a7f37",
    )

ax.set_xlabel("Batch size")
ax.set_ylabel("Median ms per call")
ax.set_title("Attention latency vs batch size")
ax.legend()
ax.grid(True, alpha=0.25)
ax.set_xticks(batches)

fig.tight_layout()
plt.show()

Deberías obtener una tabla y un gráfico con una entrada por tamaño de lote de 1 a 16. La implementación simple se omite de este barrido, por lo que ambas salidas comparan solo las tres rutas fusionadas.

9. Compara MHA, GQA y MQA

Hasta ahora, cada encabezado ha llevado sus propias claves y valores. Sin embargo, los decodificadores suelen romper esa simetría, ya que durante la inferencia, la caché de par clave-valor es lo que domina la memoria.

La atención de varios encabezados (MHA) le da a cada encabezado sus propias proyecciones Q, K y V. La atención de consulta agrupada (GQA) y la atención de varias consultas (MQA) reducen la cantidad de encabezados KV para ahorrar memoria y procesamiento durante la inferencia.

jax.nn.dot_product_attention controla los tres con los encabezados KV que se transmiten automáticamente cuando K < N.

rk = jax.random.key(42)
rk1, rk2, rk3, rk4, rk5 = jax.random.split(rk, 5)

q_mha = jax.random.normal(rk1, (2, 64, 8, 64), dtype=jnp.float32)

# MHA: 8 KV heads
k_mha = jax.random.normal(rk2, (2, 64, 8, 64), dtype=jnp.float32)
v_mha = jax.random.normal(rk3, (2, 64, 8, 64), dtype=jnp.float32)

# GQA: 2 KV heads (each shared by 4 query heads)
k_gqa = jax.random.normal(rk2, (2, 64, 2, 64), dtype=jnp.float32)
v_gqa = jax.random.normal(rk3, (2, 64, 2, 64), dtype=jnp.float32)

# MQA: 1 KV head (shared by all 8 query heads)
k_mqa = jax.random.normal(rk4, (2, 64, 1, 64), dtype=jnp.float32)
v_mqa = jax.random.normal(rk5, (2, 64, 1, 64), dtype=jnp.float32)

out_mha = block_tree(jax.nn.dot_product_attention(q_mha, k_mha, v_mha))
out_gqa = block_tree(jax.nn.dot_product_attention(q_mha, k_gqa, v_gqa))
out_mqa = block_tree(jax.nn.dot_product_attention(q_mha, k_mqa, v_mqa))

show_table(
    ["Pattern", "Q shape", "K shape", "V shape", "Output shape"],
    [
        ("MHA", q_mha.shape, k_mha.shape, v_mha.shape, out_mha.shape),
        ("GQA", q_mha.shape, k_gqa.shape, v_gqa.shape, out_gqa.shape),
        ("MQA", q_mha.shape, k_mqa.shape, v_mqa.shape, out_mqa.shape),
    ],
    title="Multi-head attention variants — all should produce the same output shape",
)

Las tres filas deben informar la misma forma de salida, (2, 64, 8, 64), aunque las formas K y V se reduzcan de 8 encabezados a 2 a 1. Ese es el punto: puedes cortar la caché de KV sin cambiar nada más adelante de la atención.

10. Realiza una comparativa de TransformerEngine de NVIDIA y FP8

TransformerEngine de NVIDIA proporciona módulos de atención fusionados optimizados para GPUs de NVIDIA. La integración de JAX usa módulos de estilo Flax Linen (no NNX), por lo que el módulo se inicializa una vez y, luego, se aplica con sus variables.

Realiza un barrido de la longitud de la secuencia con TransformerEngine

Esta celda verifica si TransformerEngine de NVIDIA está disponible y, luego, realiza una comparativa de su DotProductAttention causal bf16 en varias longitudes de secuencia con JAX SDPA con los backends de XLA y cuDNN en la misma carga de trabajo.

TE_SEQ_LENS = [128, 256, 512, 1024, 2048]
TE_BATCH = BATCH

HAS_TE = False

try:
    import transformer_engine.jax as te
    import transformer_engine.jax.flax as te_flax
    HAS_TE = True
except ImportError:
    print("TransformerEngine not installed — skipping TE sections.")

if HAS_TE:
    te_results = []

    for sl in TE_SEQ_LENS:
        rk = jax.random.key(sl + 1000)
        rk1, rk2, rk3 = jax.random.split(rk, 3)

        q_te = jax.random.normal(
            rk1, (TE_BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
        )
        k_te = jax.random.normal(
            rk2, (TE_BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
        )
        v_te = jax.random.normal(
            rk3, (TE_BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
        )
        q_te, k_te, v_te = jax.device_put((q_te, k_te, v_te), device)

        te_attention = te_flax.DotProductAttention(
            head_dim=HEAD_DIM,
            num_attention_heads=NUM_HEADS,
            num_gqa_groups=NUM_HEADS,
            attn_mask_type="causal",
            transpose_batch_sequence=False,
        )

        te_vars = te_attention.init(
            jax.random.key(0),
            q_te,
            k_te,
            v_te,
            deterministic=True,
        )

        def te_fn(q, k, v):
            return te_attention.apply(te_vars, q, k, v, deterministic=True)

        row = {"seq_len": sl}

        row["sdpa_xla_bf16_ms"] = benchmark_attention(
            lambda q, k, v: jax.nn.dot_product_attention(
                q, k, v, implementation="xla", is_causal=True
            ),
            q_te, k_te, v_te,
            warmup=2,
            repeats=20,
        )

        if HAS_CUDNN_SDPA:
            row["sdpa_cudnn_bf16_ms"] = benchmark_attention(
                lambda q, k, v: jax.nn.dot_product_attention(
                    q, k, v, implementation="cudnn", is_causal=True
                ),
                q_te, k_te, v_te,
                warmup=2,
                repeats=20,
            )

        row["te_bf16_ms"] = benchmark_attention(
            te_fn,
            q_te, k_te, v_te,
            warmup=2,
            repeats=20,
        )

        te_results.append(row)


    headers = ["Seq len", "SDPA XLA bf16 causal (ms)"]
    if HAS_CUDNN_SDPA:
        headers.append("SDPA cuDNN bf16 causal (ms)")
    headers.append("TE DotProductAttention bf16 causal (ms)")

    table_rows = []
    for r in te_results:
        row = [
            r["seq_len"],
            f"{r['sdpa_xla_bf16_ms']:.2f}",
        ]

        if HAS_CUDNN_SDPA:
            row.append(f"{r['sdpa_cudnn_bf16_ms']:.2f}")

        row.append(f"{r['te_bf16_ms']:.2f}")
        table_rows.append(row)

    show_table(
        headers,
        table_rows,
        title=f"TransformerEngine sequence-length sweep — B={TE_BATCH}, N={NUM_HEADS}, H={HEAD_DIM}",
        aligns=["right"] * len(headers),
    )


    fig, ax = plt.subplots(figsize=(8, 5))
    seq_lens = [r["seq_len"] for r in te_results]

    ax.plot(
        seq_lens,
        [r["sdpa_xla_bf16_ms"] for r in te_results],
        "d-",
        label="SDPA XLA bf16 causal",
        color="#8250df",
    )

    if HAS_CUDNN_SDPA:
        ax.plot(
            seq_lens,
            [r["sdpa_cudnn_bf16_ms"] for r in te_results],
            "^-",
            label="SDPA cuDNN bf16 causal",
            color="#1a7f37",
        )

    ax.plot(
        seq_lens,
        [r["te_bf16_ms"] for r in te_results],
        "o-",
        label="TE DotProductAttention bf16 causal",
        color="#d1242f",
    )

    ax.set_xlabel("Sequence length")
    ax.set_ylabel("Median ms per call")
    ax.set_title("Causal attention latency vs sequence length")
    ax.legend()
    ax.grid(True, alpha=0.25)
    ax.set_xticks(seq_lens)

    fig.tight_layout()
    plt.show()

Deberías obtener una fila de tabla y un punto de gráfico por longitud de secuencia de 128 a 2048. Esta comparativa bf16 compara TransformerEngine con JAX SDPA en la misma carga de trabajo de atención causal, pero el potencial de rendimiento completo de TransformerEngine suele verse en las GPUs Hopper y Blackwell cuando está disponible la conversión automática de FP8.

Inspecciona la ruta FP8

En las GPUs Hopper (capacidad de procesamiento >= 9.0, como H100), TransformerEngine puede ejecutar la atención en FP8 para obtener un rendimiento adicional. FP8 usa la receta DelayedScaling que realiza un seguimiento del historial de valores máximos absolutos por tensor para calcular los factores de escala dinámicos.

Formato E4M3 para el pase directo (4 bits de exponente, 3 de mantisa)
Formato E5M2 para el pase inverso (5 bits de exponente, 2 de mantisa)

Si la GPU no admite FP8, esta celda muestra cómo se vería el código sin ejecutarlo.

if HAS_TE:
    from transformer_engine.common.recipe import DelayedScaling, Format

    gpu_name = f"{device} {getattr(device, 'device_kind', '')}".lower()
    HAS_FP8 = any(
        tag in gpu_name
        for tag in ["h100", "h200", "b100", "b200", "gb200", "blackwell"]
    )

    fp8_recipe = DelayedScaling(
        margin=0,
        fp8_format=Format.HYBRID,
        amax_history_len=1024,
        amax_compute_algo="max",
    )

    if HAS_FP8:
        FP8_SEQ_LEN = 2048
        FP8_BATCH = BATCH

        rk = jax.random.key(9000)
        rk1, rk2, rk3 = jax.random.split(rk, 3)

        q_fp8 = jax.random.normal(
            rk1, (FP8_BATCH, FP8_SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
        )
        k_fp8 = jax.random.normal(
            rk2, (FP8_BATCH, FP8_SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
        )
        v_fp8 = jax.random.normal(
            rk3, (FP8_BATCH, FP8_SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
        )
        q_fp8, k_fp8, v_fp8 = jax.device_put((q_fp8, k_fp8, v_fp8), device)

        fp8_attention = te_flax.DotProductAttention(
            head_dim=HEAD_DIM,
            num_attention_heads=NUM_HEADS,
            num_gqa_groups=NUM_HEADS,
            attn_mask_type="causal",
            transpose_batch_sequence=False,
        )

        bf16_vars = fp8_attention.init(
            jax.random.key(0),
            q_fp8,
            k_fp8,
            v_fp8,
            deterministic=True,
        )

        bf16_out = block_tree(
            fp8_attention.apply(
                bf16_vars,
                q_fp8,
                k_fp8,
                v_fp8,
                deterministic=True,
            )
        )

        with te.autocast(enabled=True, recipe=fp8_recipe):
            fp8_vars = fp8_attention.init(
                jax.random.key(1),
                q_fp8,
                k_fp8,
                v_fp8,
                deterministic=True,
            )
            fp8_out = block_tree(
                fp8_attention.apply(
                    fp8_vars,
                    q_fp8,
                    k_fp8,
                    v_fp8,
                    deterministic=True,
                )
            )

        max_diff_fp8 = float(jnp.max(jnp.abs(
            bf16_out.astype(jnp.float32) - fp8_out.astype(jnp.float32)
        )))

        show_table(
            ["", "Value"],
            [
                ("GPU", getattr(device, "device_kind", str(device))),
                ("Input dtype", str(q_fp8.dtype)),
                ("bf16 output dtype", str(bf16_out.dtype)),
                ("FP8 autocast output dtype", str(fp8_out.dtype)),
                ("Output shape", str(fp8_out.shape)),
                ("Max |TE bf16 - TE FP8 autocast|", f"{max_diff_fp8:.2e}"),
            ],
            title="FP8 attention with TransformerEngine",
        )

    else:
        show_table(
            ["", "Value"],
            [
                ("GPU", getattr(device, "device_kind", str(device))),
                ("FP8 support", "No detected support; requires Hopper/Blackwell-class GPU"),
            ],
            title="FP8 attention — not available on this GPU",
        )

        print()
        print("The FP8 path uses TransformerEngine autocast:")
        print()
        print("  with te.autocast(enabled=True, recipe=fp8_recipe):")
        print("      out = fp8_attention.apply(vars, q, k, v, deterministic=True)")

else:
    print("TransformerEngine not available — FP8 section skipped.")

En la L4, deberías ver una tabla con el nombre de tu GPU y que no informa compatibilidad con FP8 detectada, seguida de las dos líneas impresas que muestran la llamada te.autocast que usarías en una GPU Hopper. Conserva ese fragmento: es el único cambio que requiere FP8 en el sitio de llamada.

11. Limpia

Borra la carga de trabajo de Jupyter, incluido el LoadBalancer y el volumen persistente:

kubectl delete -f deploy/jupyter.yaml

Destruye el clúster, el grupo de nodos, la VPC y la cuenta de servicio:

cd terraform
terraform destroy

Escribe yes cuando se te solicite y, luego, confirma que no quede nada:

gcloud container clusters list
gcloud compute instances list

Ambos deben estar vacíos para este proyecto. Si creaste un proyecto solo para esta serie, puedes borrar todo el proyecto de la consola de Cloud.

12. Felicitaciones

Pasaste de un cálculo de atención escrito a mano a kernels fusionados optimizados para GPU y mediste la diferencia en hardware real.

Qué aprendiste

La atención simple funciona, pero inicia varios kernels de GPU y materializa la matriz de atención completa en la memoria.
jax.nn.dot_product_attention fusiona el cálculo en una sola operación y, con implementation=None, JAX elige automáticamente el mejor backend.
implementation="cudnn" fuerza los kernels de atención fusionados de cuDNN de NVIDIA, que son más rápidos en secuencias largas y requieren entradas bfloat16 o float16 y capacidad de procesamiento 8.0 o posterior.
El enmascaramiento causal con is_causal=True está integrado en el kernel fusionado; no se necesita una matriz de máscara manual.
GQA y MQA reducen los encabezados KV para ahorrar memoria durante la inferencia, y dot_product_attention controla la transmisión automáticamente.
TransformerEngine proporciona atención fusionada con precisión FP8 opcional en las GPUs Hopper.

Próximos pasos

En el **codelab 6: Escala el entrenamiento de JAX en varias GPUs** , se mostrará cómo fragmentar arrays en L4 y ejecutar pasos de entrenamiento en paralelo.
Cambia HEAD_DIM (prueba 32, 64, 128) y observa qué valores acepta cuDNN y cómo se mueven los tiempos.
Aumenta SEQ_LEN (prueba 256, 512, 1024, 2048) y observa el uso de memoria y la proporción de aceleración de cuDNN.
Vuelve a ejecutar la comparación de MHA, GQA y MQA con encabezados KV K=2 y K=1 en comparación con 8 encabezados de consulta, y razona sobre los ahorros de caché de KV para la inferencia.