Velocizza l'attenzione sulla GPU con cuDNN e TransformerEngine

1. Introduzione

Percorso di apprendimento Jax on GPU. Lab 5: Attention on GPU.

In questo codelab ti concentrerai sul calcolo dell'attenzione sulle GPU NVIDIA. Inizierai con un'implementazione semplice, la sostituirai con l'attenzione fusa integrata di JAX, quindi attiverai il backend cuDNN di NVIDIA per la stessa operazione. Alla fine, saprai quali modifiche vengono apportate a ogni livello e quanto più veloci diventano i percorsi ottimizzati per la GPU man mano che il problema aumenta.

In questo lab proverai a:

  • Implementare l'attenzione del prodotto scalato da zero con le operazioni JAX di base
  • Sostituirla con jax.nn.dot_product_attention, il kernel fuso integrato di JAX
  • Forzare il backend cuDNN con implementation="cudnn" e aggiungere la maschera causale
  • Scorrere la lunghezza della sequenza e la dimensione del batch per vedere dove i kernel fusi sono più efficaci
  • Confrontare le forme di attenzione multi-testa MHA, GQA e MQA
  • Eseguire il benchmark dell'attenzione di NVIDIA TransformerEngine e ispezionare il relativo percorso del codice FP8

Che cosa ti serve

  • Un progetto Google Cloud con la fatturazione abilitata e crediti del workshop o una prenotazione che copra l'utilizzo della GPU
  • Quota per almeno 2 GPU NVIDIA L4 nella regione scelta (come controllare la quota di GPU)
  • Completamento dei codelab da 1 a 4 o un ambiente GPU JAX con CUDA equivalente
  • Dimestichezza con la moltiplicazione delle matrici e softmax. Non è richiesta alcuna esperienza con CUDA o cuDNN.

Tempo stimato per completare il corso: 60 minuti.

Come funziona l'attenzione

L'attenzione del prodotto scalato accetta tre input: Query, Key e Value e calcola:

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

Il fattore di scalabilità 1 / sqrt(d_k) impedisce ai prodotti scalari di diventare troppo grandi man mano che la dimensione della testa aumenta, il che spingerebbe la softmax in regioni in cui i gradienti sono minuscoli.

Quattro passaggi, ognuno dei quali alimenta il successivo:

  1. Punteggio : Q @ K.T: il prodotto scalare misura la quantità di attenzione che ogni posizione della query deve prestare a ogni posizione della chiave.
  2. Scala + Softmax : softmax(scores / sqrt(d)): la scalabilità impedisce la scomparsa del gradiente e la softmax converte i punteggi in ponderazioni di attenzione che sommano a 1.
  3. Attenzione : weights @ V: la somma ponderata dei vettori di valori produce l'output per ogni posizione della query.
  4. Output : stessa forma di Q: ogni posizione della query ora contiene informazioni dalle posizioni a cui ha prestato attenzione.

2. Prima di iniziare

Seleziona il progetto

Nella console Google Cloud, seleziona o crea un progetto con la fatturazione abilitata.

Apri Cloud Shell

Fai clic su Attiva Cloud Shell (l'icona del terminale in alto a destra nella console) per avviare una sessione di Cloud Shell, quindi puntala al tuo progetto:

gcloud config set project <YOUR_PROJECT_ID>

Esegui il provisioning dell'ambiente GPU

Esegui il comando riportato di seguito in Cloud Shell. Il codelab 1 spiega ogni comando in dettaglio, inclusi i requisiti di quota e zona della 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

Modifica terraform.tfvars e imposta project_id = "", quindi esegui il provisioning del cluster e il deployment di JupyterLab:

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

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

terraform apply richiede circa 12 minuti. Al termine, attendi il Pod e il LoadBalancer, quindi leggi il token JupyterLab una tantum dal log 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

Apri http://:8884, incolla il token e crea un nuovo notebook Python 3 in /workspace. Ogni blocco di codice in questo codelab va in una cella di quel notebook.

Installa ciò che è necessario per questo codelab

!pip install --quiet matplotlib flax

Configura e verifica la GPU

Importa JAX e conferma che il backend predefinito è una GPU. Questa cella definisce anche block_tree, show_table e show_bars, gli helper utilizzati da ogni passaggio successivo per attendere il lavoro del dispositivo e visualizzare i risultati.

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

Dovresti visualizzare una versione JAX, gpu come backend predefinito, un elenco di dispositivi CUDA e la GPU utilizzata dal resto del codelab.

3. Crea array di test Q, K e V

I kernel di attenzione fusi si preoccupano di come gli array di input sono disposti in memoria, quindi crea Q, K e V nel layout previsto da JAX prima di scrivere qualsiasi codice di attenzione. jax.nn.dot_product_attention prevede il seguente layout:

Dim

Significato

Valore predefinito

B

Dimensione del batch

4

T

Lunghezza della sequenza di query

128

S

Lunghezza della sequenza chiave/valore

128 (uguale a T per l'auto-attenzione)

N

Numero di teste di attenzione

8

H

Dimensione per testa

64

Inizia con float32 per l'implementazione semplice, inizializzando con valori casuali.

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

Dovresti visualizzare una tabella con una riga per array, ognuna delle quali riporta la forma (4, 128, 8, 64) e il tipo di dati float32. La chiamata jax.device_put blocca tutti e tre gli array sulla GPU selezionata nella cella di configurazione, quindi nulla nei benchmark successivi misura un trasferimento host.

4. Implementa l'attenzione da zero

Prima di utilizzare un kernel fuso, scrivi direttamente la formula in modo da vedere esattamente cosa sostituisce il kernel. Questa implementazione segue i quattro passaggi in ordine: calcola i punteggi, scala, softmax, moltiplica per i valori. Trasponi le dimensioni della testa e della sequenza in modo che la moltiplicazione delle matrici operi sull'asse della sequenza.

Funziona correttamente, ma emette tre lanci di kernel GPU separati per la moltiplicazione delle matrici di punteggio, softmax, moltiplicazione delle matrici di valori e materializza la matrice di ponderazione dell'attenzione (B, N, T, S) completa nella memoria 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",
)

Dovresti visualizzare la forma di output (4, 128, 8, 64) e il tipo di dati float32, che è la stessa forma di Q, come promesso dal passaggio Output della formula.

5. Passa a dot_product_attention

jax.nn.dot_product_attention unisce i passaggi di punteggio, scala, softmax e attenzione in una singola operazione. JAX e XLA possono quindi ottimizzare il pattern di accesso alla memoria. In particolare, possono evitare di materializzare la matrice di ponderazione dell'attenzione completa quando la sequenza è lunga.

Con implementation=None predefinito, JAX sceglie automaticamente il backend migliore disponibile. Su una GPU con cuDNN disponibile e input compatibili, potrebbe già utilizzare cuDNN. Su altro hardware, torna 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",
)

Dovresti visualizzare la stessa forma e lo stesso tipo di dati della versione semplice, una piccola differenza assoluta massima e True per il controllo di prossimità atol=1e-3.

6. Forza il backend fuso cuDNN

Lasciare che JAX scelga il backend è comodo, ma dall'esterno non è possibile sapere quale kernel è stato effettivamente eseguito. L'impostazione di implementation="cudnn" forza i kernel di attenzione fusi cuDNN di NVIDIA. Si tratta di kernel GPU ottimizzati manualmente che fondono l'intero calcolo dell'attenzione in un singolo lancio di kernel con pattern di accesso alla memoria ottimizzati.

Tieni presente che l'attenzione fusa cuDNN ha requisiti hardware e di forma, come la capacità di calcolo della GPU, che deve essere >= 8.0 (Ampere o versioni successive) o float16 o bfloat16 come tipi di dati di input. Se i requisiti non vengono soddisfatti e imposti implementation="cudnn", JAX genera un errore anziché tornare silenziosamente. Ecco perché il codice riportato di seguito esegue il cast a bfloat16 e racchiude la chiamata in try/except: imposta HAS_CUDNN_SDPA in modo che ogni passaggio successivo sappia se il percorso cuDNN è disponibile su questa macchina.

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 la GPU non è in grado di eseguire SDPA cuDNN, il codice stampa il motivo e il codelab continua sul backend XLA.

7. Aggiungi la maschera causale

Nei modelli autoregressivi (decoder in stile GPT), ogni posizione può prestare attenzione solo alle posizioni precedenti. L'impostazione di is_causal=True applica questa maschera triangolare inferiore all'interno del kernel fuso, quindi non è necessario creare una matrice di maschere.

Il codice riportato di seguito esegue l'attenzione due volte, con e senza la maschera, e confronta due posizioni per mostrare cosa ha modificato la maschera.

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

Dovresti notare una grande differenza nella prima posizione e una differenza quasi nulla nell'ultima. Questa asimmetria è la maschera con la posizione 0 che perde ogni chiave tranne la propria, mentre la posizione finale poteva già vedere l'intera sequenza.

8. Esecuzione di alcuni benchmark

Ora hai tre modi per calcolare la stessa funzione. La scelta dipende interamente dalla forma del problema e puoi valutarla.

Misura il tempo delle varianti nella forma predefinita

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)

Per questa piccola dimensione del problema, tutte le implementazioni sono molto veloci e per lo più limitate all'overhead, quindi la versione semplice è competitiva. bf16 cuDNN dovrebbe essere leggermente più veloce, ma il vantaggio dell'attenzione fusa di solito diventa più chiaro a lunghezze di sequenza maggiori, dove evitare la matrice di attenzione completa è importante.

Scorri la lunghezza della sequenza

Il vantaggio dei kernel di attenzione fusi aumenta con la lunghezza della sequenza. L'implementazione semplice materializza una matrice di attenzione (B, N, T, S) nella memoria GPU, che è una memoria O(T²). I kernel fusi come cuDNN FlashAttention dividono il calcolo in modo da non materializzare mai la matrice completa, mantenendo la memoria a O(T).

Lo scorrimento riportato di seguito misura il tempo di ogni implementazione per lunghezze di sequenza da 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),
)

Lo scorrimento richiede un po' di tempo, perché ogni lunghezza di sequenza compila tre o quattro implementazioni separate prima di misurarne il tempo. Dovresti ottenere una riga di tabella per ogni lunghezza di sequenza da 64 a 1024.

Questo grafico confronta la variazione della latenza dell'attenzione con la lunghezza della sequenza per l'implementazione semplice, JAX/XLA SDPA in float32 e bfloat16 e l'attenzione fusa cuDNN in 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()

Dovresti visualizzare una riga per implementazione. E dovresti vedere il divario che aumenta man mano che la sequenza si allunga.

Scorri la dimensione del batch

I batch più grandi ammortizzano l'overhead di lancio del kernel e migliorano l'utilizzo della GPU fino al punto in cui la memoria GPU diventa il collo di bottiglia. Lo scorrimento riportato di seguito mantiene la lunghezza della sequenza fissa a 256 e varia la dimensione del batch.

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

Dovresti ottenere una tabella e un grafico con una voce per ogni dimensione del batch da 1 a 16. L'implementazione semplice viene esclusa da questo scorrimento, quindi entrambi gli output confrontano solo i tre percorsi fusi.

9. Confronta MHA, GQA e MQA

Finora ogni testa ha portato le proprie chiavi e i propri valori. Tuttavia, i decoder spesso interrompono questa simmetria, perché durante l'inferenza la cache della coppia chiave-valore è ciò che domina la memoria.

L'attenzione multi-testa (MHA) assegna a ogni testa le proprie proiezioni Q, K e V. L'attenzione alle query raggruppate (GQA) e l'attenzione alle query multiple (MQA) riducono il numero di teste KV per risparmiare memoria e calcolo durante l'inferenza.

jax.nn.dot_product_attention gestisce tutte e tre le teste KV vengono trasmesse 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",
)

Tutte e tre le righe devono riportare la stessa forma di output, (2, 64, 8, 64), anche se le forme K e V si riducono da 8 teste a 2 a 1. Questo è il punto: puoi tagliare la cache KV senza modificare nulla a valle dell'attenzione.

10. Esegui il benchmark di NVIDIA TransformerEngine e FP8

TransformerEngine di NVIDIA fornisce moduli di attenzione fusi ottimizzati per le GPU NVIDIA. L'integrazione JAX utilizza moduli in stile Flax Linen (non NNX), quindi il modulo viene inizializzato una volta e poi applicato con le relative variabili.

Scorri la lunghezza della sequenza con TransformerEngine

Questa cella controlla se NVIDIA TransformerEngine è disponibile, quindi esegue il benchmark del relativo DotProductAttention causale bf16 su diverse lunghezze di sequenza rispetto a JAX SDPA con i backend XLA e cuDNN sullo stesso carico di lavoro.

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

Dovresti ottenere una riga di tabella e un punto del grafico per ogni lunghezza di sequenza da 128 a 2048. Questo benchmark bf16 confronta TransformerEngine con JAX SDPA sullo stesso carico di lavoro di attenzione causale, ma il potenziale di prestazioni completo di TransformerEngine viene in genere visualizzato sulle GPU Hopper e Blackwell quando è disponibile FP8 autocast.

Ispeziona il percorso FP8

Sulle GPU Hopper (capacità di calcolo >= 9.0, ad esempio H100), TransformerEngine può eseguire l'attenzione in FP8 per un throughput aggiuntivo. FP8 utilizza la ricetta DelayedScaling che tiene traccia della cronologia assoluta massima per tensore per calcolare i fattori di scalabilità dinamici.

  • Formato E4M3 per il passaggio in avanti (4 bit di esponente, 3 bit di mantissa)
  • Formato E5M2 per il passaggio all'indietro (5 bit di esponente, 2 bit di mantissa)

Se la GPU non supporta FP8, questa cella mostra l'aspetto del codice senza eseguirlo.

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

Su L4 dovresti visualizzare una tabella che indica la tua GPU e segnala che non è stato rilevato il supporto FP8, seguita dalle due righe stampate che mostrano la chiamata te.autocast che utilizzeresti su una GPU Hopper. Conserva questo snippet: è l'unica modifica richiesta da FP8 nel sito di chiamata.

11. Libera spazio

Elimina il carico di lavoro Jupyter, inclusi il LoadBalancer e il volume permanente:

kubectl delete -f deploy/jupyter.yaml

Elimina il cluster, il node pool, il VPC e l'account di servizio:

cd terraform
terraform destroy

Digita yes quando richiesto, quindi verifica che non sia rimasto nulla:

gcloud container clusters list
gcloud compute instances list

Entrambi devono essere vuoti per questo progetto. Se hai creato un progetto solo per questa serie, puoi invece eliminare l'intero progetto dalla console Cloud.

12. Complimenti

Sei passato da un calcolo dell'attenzione scritto a mano a kernel fusi ottimizzati per la GPU e hai misurato la differenza sull'hardware reale.

Che cosa hai imparato

  • L'attenzione semplice funziona, ma avvia più kernel GPU e materializza la matrice di attenzione completa in memoria
  • jax.nn.dot_product_attention unisce il calcolo in una singola operazione e, con implementation=None, JAX sceglie automaticamente il backend migliore
  • implementation="cudnn" forza i kernel di attenzione fusi cuDNN di NVIDIA, che sono i più veloci su sequenze lunghe e richiedono input bfloat16 o float16 e capacità di calcolo 8.0 o versioni successive
  • La maschera causale con is_causal=True è integrata nel kernel fuso: non è necessaria alcuna matrice di maschere manuale
  • GQA e MQA riducono le teste KV per risparmiare memoria durante l'inferenza e dot_product_attention gestisce automaticamente la trasmissione
  • TransformerEngine fornisce un'attenzione fusa con precisione FP8 facoltativa sulle GPU Hopper

Passaggi successivi

  • Il **codelab 6: Scalare l'addestramento JAX su più GPU** ti mostrerà come partizionare gli array su entrambe le L4 ed eseguire i passaggi di addestramento in parallelo
  • Modifica HEAD_DIM (prova 32, 64, 128) e osserva quali valori accetta cuDNN e come si spostano i tempi
  • Aumenta SEQ_LEN (prova 256, 512, 1024, 2048) e osserva la memoria utilizzata e il rapporto di accelerazione cuDNN
  • Esegui di nuovo il confronto MHA, GQA e MQA con le teste KV K=2 e K=1 rispetto a 8 teste di query e ragiona sui risparmi della cache KV per l'inferenza

Documenti di riferimento