Acelerar a atenção na GPU com cuDNN e TransformerEngine

1. Introdução

Programa de aprendizado do Jax em GPU. Laboratório 5: atenção na GPU.

Neste codelab, você vai se concentrar na computação da atenção em GPUs NVIDIA. Comece com uma implementação simples, substitua-a pela atenção integrada do JAX e ative o back-end cuDNN da NVIDIA para a mesma operação. Ao final, você saberá o que muda em cada nível e o quanto os caminhos otimizados para GPU ficam mais rápidos à medida que o problema aumenta.

Atividades deste laboratório

  • Implementar a atenção de produto escalar dimensionado do zero com operações básicas do JAX
  • Substituir por jax.nn.dot_product_attention, o kernel mesclado integrado do JAX
  • Forçar o back-end cuDNN com implementation="cudnn" e adicionar mascaramento causal
  • Varrer o comprimento da sequência e o tamanho do lote para ver onde os kernels mesclados valem a pena
  • Comparar as formas de atenção de várias cabeças MHA, GQA e MQA
  • Comparar a atenção do TransformerEngine da NVIDIA e inspecionar o caminho de código FP8

O que é necessário

  • Um projeto na nuvem do Google Cloud com o faturamento ativado e créditos de workshop ou uma reserva que cubra o uso da GPU
  • Cota para pelo menos duas GPUs NVIDIA L4 na região escolhida (como verificar a cota de GPU)
  • Conclusão dos codelabs de 1 a 4 ou um ambiente de GPU JAX equivalente ativado para CUDA
  • Familiaridade com multiplicação de matrizes e softmax. Não é necessário ter experiência com CUDA ou cuDNN.

Tempo estimado para a conclusão: 60 minutos.

Como a atenção funciona

A atenção de produto escalar dimensionado recebe três entradas: consulta, chave e valor, e calcula:

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

O fator de escalonamento 1 / sqrt(d_k) impede que os produtos escalares cresçam muito à medida que a dimensão da cabeça aumenta, o que empurraria o softmax para regiões em que os gradientes são pequenos.

Quatro etapas, cada uma alimentando a próxima:

  1. Pontuação : Q @ K.T: o produto escalar mede o quanto cada posição de consulta deve atender a cada posição de chave.
  2. Escala + Softmax : softmax(scores / sqrt(d)): o escalonamento impede o desaparecimento do gradiente, e o softmax converte as pontuações em ponderações de atenção que somam 1.
  3. Atender : weights @ V: a soma ponderada dos vetores de valor produz a saída para cada posição de consulta.
  4. Saída : mesma forma que Q: cada posição de consulta agora contém informações das posições que ela atendeu.

2. Antes de começar

Selecionar o projeto

No console do Google Cloud, selecione ou crie um projeto com o faturamento ativado.

Abrir o Cloud Shell

Clique em Ativar o Cloud Shell (o ícone do terminal no canto superior direito do console) para iniciar uma sessão do Cloud Shell e aponte para o projeto:

gcloud config set project <YOUR_PROJECT_ID>

Provisionar o ambiente de GPU

Execute o seguinte no Cloud Shell. O codelab 1 explica cada comando em detalhes, incluindo cota de GPU e requisitos de zona.

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

Edite terraform.tfvars e defina project_id = "". Em seguida, provisione o cluster e implante o JupyterLab:

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

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

terraform apply leva cerca de 12 minutos. Quando terminar, aguarde o pod e o balanceador de carga e leia o token único do JupyterLab no registro do 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

Abra http://:8884, cole o token e crie um novo notebook do Python 3 em /workspace. Cada bloco de código neste codelab entra em uma célula desse notebook.

Instalar o que este codelab precisa

!pip install --quiet matplotlib flax

Configurar e verificar a GPU

Importe o JAX e confirme se o back-end padrão é uma GPU. Essa célula também define block_tree, show_table e show_bars, os auxiliares que cada etapa posterior usa para aguardar o trabalho do dispositivo e renderizar os 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)))

Você verá uma versão do JAX, gpu como o back-end padrão, uma lista de dispositivos CUDA e a GPU que o restante do codelab usa.

3. Criar matrizes de teste Q, K e V

Os kernels de atenção mesclados se preocupam com a forma como as matrizes de entrada são dispostas na memória. Portanto, crie Q, K e V no layout esperado pelo JAX antes de escrever qualquer código de atenção. jax.nn.dot_product_attention espera o seguinte layout:

Escurecer

Significado

Nosso padrão

B

Tamanho do lote

4

T

Comprimento da sequência de consultas

128

S

Comprimento da sequência de chave/valor

128 (o mesmo que T para autoatenção)

N

Número de cabeças de atenção

8

H

Dimensão por cabeça

64

Você começa em float32 para a implementação simples, inicializando com valores aleatórios.

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",
)

Você verá uma tabela com uma linha por matriz, cada uma informando a forma (4, 128, 8, 64) e o dtype float32. A chamada jax.device_put fixa todas as três matrizes na GPU selecionada na célula de configuração. Portanto, nada nos comparativos posteriores está medindo uma transferência de host.

4. Implementar a atenção do zero

Antes de usar um kernel mesclado, escreva a fórmula diretamente para que você possa ver exatamente o que o kernel substitui. Essa implementação segue as quatro etapas em ordem: calcular pontuações, dimensionar, softmax, multiplicar por valores. Você transpõe as dimensões da cabeça e da sequência para que o matmul opere no eixo da sequência.

Ele funciona corretamente, mas emite três inicializações de kernel de GPU separadas para matmul de pontuação, softmax, matmul de valor e materializa a matriz de ponderação de atenção (B, N, T, S) completa na memória da 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",
)

Você verá a forma de saída (4, 128, 8, 64) e o dtype float32, que é a mesma forma de Q, que é o que a etapa Saída da fórmula prometeu.

5. Mudar para dot_product_attention

O jax.nn.dot_product_attention mescla as etapas de pontuação, escala, softmax e atendimento em uma única operação. O JAX e o XLA podem otimizar o padrão de acesso à memória. Em particular, eles podem evitar a materialização da matriz de ponderação de atenção completa quando a sequência é longa.

Com o implementation=None padrão, o JAX escolhe automaticamente o melhor back-end disponível. Em uma GPU com cuDNN disponível e entradas compatíveis, ele já pode usar o cuDNN. Em outros hardwares, ele volta para o 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",
)

Você verá a mesma forma e dtype da versão simples, uma pequena diferença absoluta máxima e True para a verificação de proximidade atol=1e-3.

6. Forçar o back-end mesclado cuDNN

Deixar o JAX escolher o back-end é conveniente, mas, de fora, não é possível saber qual kernel foi executado. Definir implementation="cudnn" força os kernels de atenção mesclados do cuDNN da NVIDIA. Esses são kernels de GPU otimizados manualmente que mesclam toda a computação de atenção em uma única inicialização de kernel com padrões de acesso à memória otimizados.

A atenção mesclada do cuDNN tem requisitos de hardware e forma, como a capacidade de computação da GPU, que precisa ser >= 8.0 (Ampere ou mais recente) ou float16 ou bfloat16 como dtypes de entrada. Se os requisitos não forem atendidos e você definir implementation="cudnn", o JAX vai gerar um erro em vez de voltar silenciosamente. É por isso que o código abaixo transmite para bfloat16 e envolve a chamada em try/except: ele define HAS_CUDNN_SDPA para que cada etapa posterior saiba se o caminho do cuDNN está disponível nessa 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.")

Se a GPU não puder executar o SDPA do cuDNN, o código vai imprimir o motivo e o codelab vai continuar no back-end do XLA.

7. Adicionar mascaramento causal

Em modelos autorregressivos (decodificadores no estilo GPT), cada posição só pode atender a posições anteriores. Definir is_causal=True aplica essa máscara triangular inferior dentro do kernel mesclado para que você não precise criar uma matriz de máscara.

O código abaixo executa a atenção duas vezes, com e sem a máscara, e compara duas posições para mostrar o que a máscara mudou.

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",
)

Você verá uma grande diferença na primeira posição e uma diferença quase zero na última. Essa assimetria é a máscara com a posição 0 que perde todas as chaves, exceto a própria, enquanto a posição final já pode ver toda a sequência.

8. Executar alguns comparativos

Agora você tem três maneiras de calcular a mesma função. Qual delas usar depende inteiramente da forma do problema e de como você pode avaliá-lo.

Cronometrar as variantes na forma padrão

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 esse tamanho de problema pequeno, todas as implementações são muito rápidas e, na maioria das vezes, limitadas pela sobrecarga. Portanto, a versão simples é competitiva. O bf16 do cuDNN precisa ser um pouco mais rápido, mas a vantagem da atenção mesclada geralmente fica mais clara em sequências mais longas, em que evitar a matriz de atenção completa é importante.

Varrer o comprimento da sequência

A vantagem dos kernels de atenção mesclados aumenta com o comprimento da sequência. A implementação simples materializa uma matriz de atenção (B, N, T, S) na memória da GPU, que é uma memória O(T²). Kernels mesclados, como o FlashAttention do cuDNN, dividem a computação para que nunca materializem a matriz completa, mantendo a memória em O(T).

A varredura abaixo cronometra cada implementação em comprimentos de sequência 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),
)

A varredura leva um tempo, porque cada comprimento de sequência compila três ou quatro implementações separadas antes de cronometrá-las. Você vai acabar com uma linha de tabela por comprimento de sequência de 64 a 1024.

Este gráfico compara como a latência de atenção muda com o comprimento da sequência para a implementação simples, o SDPA JAX/XLA em float32 e bfloat16 e a atenção mesclada do cuDNN em 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()

Você verá uma linha por implementação. E você verá a lacuna que cresce à medida que a sequência fica mais longa.

Varrer o tamanho do lote

Lotes maiores amortizam a sobrecarga de inicialização do kernel e melhoram a utilização da GPU até o ponto em que a memória da GPU se torna o gargalo. A varredura abaixo mantém o comprimento da sequência fixo em 256 e varia o tamanho do 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()

Você vai receber uma tabela e um gráfico com uma entrada por tamanho do lote de 1 a 16. A implementação simples é deixada de fora dessa varredura. Portanto, as duas saídas comparam apenas os três caminhos mesclados.

9. Comparar MHA, GQA e MQA

Até agora, cada cabeça carregou as próprias chaves e valores. Mas os decodificadores geralmente quebram essa simetria, porque durante a inferência o cache de chave-valor é o que domina a memória.

A atenção de várias cabeças (MHA) dá a cada cabeça as próprias projeções Q, K e V. A atenção de consulta agrupada (GQA) e a atenção de várias consultas (MQA) reduzem o número de cabeças KV para economizar memória e computação durante a inferência.

jax.nn.dot_product_attention processa todas as três com as cabeças KV transmitidas automaticamente quando 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",
)

Todas as três linhas precisam informar a mesma forma de saída, (2, 64, 8, 64), mesmo que as formas K e V diminuam de 8 cabeças para 2 a 1. Esse é o ponto: você pode cortar o cache KV sem mudar nada a jusante da atenção.

10. Comparar o TransformerEngine da NVIDIA e o FP8

O TransformerEngine da NVIDIA oferece módulos de atenção mesclados otimizados para GPUs NVIDIA. A integração do JAX usa módulos no estilo Flax Linen (não NNX). Portanto, o módulo é inicializado uma vez e aplicado com as variáveis.

Varrer o comprimento da sequência com o TransformerEngine

Essa célula verifica se o TransformerEngine da NVIDIA está disponível e, em seguida, compara o bf16 causal DotProductAttention em vários comprimentos de sequência com o SDPA JAX com os back-ends XLA e cuDNN na mesma carga de trabalho.

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()

Você vai receber uma linha de tabela e um ponto de gráfico por comprimento de sequência de 128 a 2048. Esse comparativo bf16 compara o TransformerEngine com o SDPA JAX na mesma carga de trabalho de atenção causal, mas o potencial de desempenho total do TransformerEngine geralmente é visto em GPUs Hopper e Blackwell quando o autocast FP8 está disponível.

Inspecionar o caminho FP8

Em GPUs Hopper (capacidade de computação >= 9.0, como H100), o TransformerEngine pode executar a atenção em FP8 para aumentar a capacidade de processamento. O FP8 usa a receita DelayedScaling que rastreia o histórico absoluto máximo por tensor para calcular fatores de escalonamento dinâmicos.

  • Formato E4M3 para a passagem direta (4 expoentes, 3 bits de mantissa)
  • Formato E5M2 para a passagem inversa (5 expoentes, 2 bits de mantissa)

Se a GPU não oferecer suporte ao FP8, essa célula vai mostrar como o código seria sem executá-lo.

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.")

Na L4, você verá uma tabela com o nome da GPU e informando que não há suporte para FP8 detectado, seguida pelas duas linhas impressas mostrando a chamada te.autocast que você usaria em uma GPU Hopper. Mantenha esse snippet: é a única mudança que o FP8 exige no local da chamada.

11. Limpar

Exclua a carga de trabalho do Jupyter, incluindo o balanceador de carga e o volume permanente:

kubectl delete -f deploy/jupyter.yaml

Destrua o cluster, o pool de nós, a VPC e a conta de serviço:

cd terraform
terraform destroy

Quando solicitado, digite yes e confirme que nada foi deixado para trás:

gcloud container clusters list
gcloud compute instances list

Os dois precisam estar vazios para este projeto. Se você criou um projeto apenas para esta série, poderá excluir todo o projeto do console do Cloud.

12. Parabéns

Você passou de uma computação de atenção escrita à mão para kernels mesclados otimizados para GPU e mediu a diferença no hardware real.

O que você aprendeu

  • A atenção simples funciona, mas inicia vários kernels de GPU e materializa a matriz de atenção completa na memória
  • jax.nn.dot_product_attention mescla a computação em uma única operação e, com implementation=None, o JAX escolhe automaticamente o melhor back-end
  • implementation="cudnn" força os kernels de atenção mesclados do cuDNN da NVIDIA, que são mais rápidos em sequências longas e exigem entradas bfloat16 ou float16 e capacidade de computação 8.0 ou mais recente
  • O mascaramento causal com is_causal=True é integrado ao kernel mesclado. Não é necessária uma matriz de máscara manual
  • GQA e MQA reduzem as cabeças KV para economizar memória durante a inferência, e dot_product_attention processa a transmissão automaticamente
  • O TransformerEngine oferece atenção mesclada com precisão FP8 opcional em GPUs Hopper

Próximas etapas

  • O **codelab 6: dimensionar o treinamento do JAX em várias GPUs** vai mostrar como fragmentar matrizes em L4s e executar etapas de treinamento em paralelo
  • Mude HEAD_DIM (tente 32, 64, 128) e observe quais valores o cuDNN aceita e como os tempos se movem
  • Aumente SEQ_LEN (tente 256, 512, 1024, 2048) e observe o uso de memória e a taxa de aceleração do cuDNN
  • Execute novamente a comparação MHA, GQA e MQA com cabeças KV K=2 e K=1 em relação a 8 cabeças de consulta e raciocine sobre as economias de cache KV para inferência

Documentos de referência