Google utilise la technologie IA pour traduire le contenu dans votre langue préférée. Les traductions générées par IA peuvent contenir des erreurs.

Accélérer l'attention sur GPU avec cuDNN et TransformerEngine

1. Introduction

Parcours de formation "Jax sur GPU". Atelier 5 : Attention sur GPU.

Dans cet atelier de programmation, vous allez vous concentrer sur le calcul de l'attention sur les GPU NVIDIA. Vous commencerez par une implémentation naïve, que vous remplacerez par l'attention fusionnée intégrée de JAX, puis vous activerez le backend cuDNN de NVIDIA pour la même opération. À la fin, vous saurez ce qui change à chaque niveau et à quel point les chemins optimisés pour les GPU deviennent plus rapides à mesure que le problème s'aggrave.

Objectifs de l'atelier

Implémenter l'attention par produit scalaire normalisé à partir de zéro avec des opérations JAX de base
Remplacez-la par jax.nn.dot_product_attention, le noyau fusionné intégré de JAX.
Forcer le backend cuDNN avec implementation="cudnn" et ajouter un masquage causal
Balayez la longueur de la séquence et la taille du lot pour voir où les kernels fusionnés sont rentables.
Comparer les formes d'attention multi-têtes MHA, GQA et MQA
Évaluer l'attention de NVIDIA TransformerEngine et inspecter son chemin de code FP8

Prérequis

Un projet Google Cloud avec facturation activée, et des crédits d'atelier ou une réservation couvrant l'utilisation des GPU
Quota d'au moins deux GPU NVIDIA L4 dans la région de votre choix (découvrez comment vérifier le quota de GPU)
Avoir suivi les ateliers de programmation 1 à 4 ou disposer d'un environnement GPU JAX compatible CUDA équivalent
Connaissance de la multiplication matricielle et de Softmax Aucune expérience de CUDA ou cuDNN n'est requise.

Durée estimée : 60 minutes.

Fonctionnement de l'attention

L'attention par produit scalaire à l'échelle prend trois entrées (Query, Key et Value) et calcule :

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

Le facteur de scaling 1 / sqrt(d_k) empêche les produits scalaires de devenir trop grands à mesure que la dimension de la tête augmente, ce qui pousserait la fonction softmax dans des régions où ses gradients sont minuscules.

Quatre étapes, chacune alimentant la suivante :

Score — Q @ K.T : le produit scalaire mesure l'attention que chaque position de requête doit accorder à chaque position de clé.
Mise à l'échelle + Softmax – softmax(scores / sqrt(d)) : la mise à l'échelle empêche la disparition du gradient, et Softmax convertit les scores en pondérations d'attention dont la somme est égale à 1.
Attend — weights @ V : la somme pondérée des vecteurs de valeur produit la sortie pour chaque position de requête.
Sortie : même forme que Q. Chaque position de requête contient désormais des informations provenant des positions auxquelles elle a prêté attention.

2. Avant de commencer

Sélectionner votre projet

Dans la console Google Cloud, sélectionnez ou créez un projet pour lequel la facturation est activée.

Ouvrir Cloud Shell

Cliquez sur Activer Cloud Shell (l'icône de terminal en haut à droite de la console) pour démarrer une session Cloud Shell, puis pointez-la vers votre projet :

gcloud config set project <YOUR_PROJECT_ID>

Provisionner l'environnement GPU

Exécutez la commande suivante dans Cloud Shell. L'atelier de programmation 1 explique chaque commande en détail, y compris les exigences concernant le quota de GPU et la zone.

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

Modifiez terraform.tfvars et définissez project_id = "", puis provisionnez le cluster et déployez JupyterLab :

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

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

terraform apply prend environ 12 minutes. Une fois l'opération terminée, attendez que le pod et l'équilibreur de charge soient prêts, puis lisez le jeton JupyterLab à usage unique dans le journal du 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

Ouvrez http://:8884, collez le jeton et créez un notebook Python 3 dans /workspace. Chaque bloc de code de cet atelier de programmation est placé dans une cellule de ce notebook.

Installer les éléments nécessaires pour cet atelier de programmation

!pip install --quiet matplotlib flax

Configurer et vérifier le GPU

Importez JAX et vérifiez que le backend par défaut est un GPU. Cette cellule définit également block_tree, show_table et show_bars, les assistants que chaque étape ultérieure utilise pour attendre le travail de l'appareil et afficher les résultats.

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

Vous devriez voir une version de JAX, gpu comme backend par défaut, une liste d'appareils CUDA et le GPU utilisé par le reste de l'atelier de programmation.

3. Créer des tableaux de test Q, K et V

Les noyaux d'attention fusionnés se soucient de la disposition des tableaux d'entrée en mémoire. Créez donc Q, K et V dans la disposition attendue par JAX avant d'écrire du code d'attention. jax.nn.dot_product_attention attend la mise en page suivante :

Diminuer la luminosité	Signification	Notre valeur par défaut
B	Taille de lot	4
T	Longueur de la séquence de requête	128
S	Longueur de la séquence clé/valeur	128 (identique à T pour l'auto-attention)
N	Nombre de têtes d'attention	8
H	Dimension par tête	64

Vous commencez dans float32 pour l'implémentation naïve, en initialisant avec des valeurs aléatoires.

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

Un tableau doit s'afficher avec une ligne par tableau, chacune indiquant la forme (4, 128, 8, 64) et le type de données float32. L'appel jax.device_put épingle les trois tableaux au GPU que vous avez sélectionné dans la cellule de configuration. Ainsi, aucun élément des benchmarks ne mesure un transfert d'hôte par la suite.

4. Implémenter l'attention à partir de zéro

Avant d'utiliser un kernel fusionné, écrivez directement la formule pour voir exactement ce que le kernel remplace. Cette implémentation suit les quatre étapes dans l'ordre : calculer les scores, mettre à l'échelle, appliquer la fonction Softmax et multiplier par les valeurs. Vous transposez les dimensions de la tête et de la séquence pour que la multiplication matricielle s'effectue sur l'axe de la séquence.

Il fonctionne correctement, mais il lance trois noyaux GPU distincts pour la multiplication matricielle des scores, la fonction softmax et la multiplication matricielle des valeurs, et il matérialise la matrice de pondération d'attention (B, N, T, S) complète dans la mémoire 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",
)

Vous devriez voir la forme de sortie (4, 128, 8, 64) et le type de données float32, qui est identique à la forme de Q, comme promis par l'étape Sortie de la formule.

5. Passer à dot_product_attention

jax.nn.dot_product_attention fusionne les étapes de score, de mise à l'échelle, de softmax et d'attention en une seule opération. JAX et XLA peuvent ensuite optimiser le modèle d'accès à la mémoire. En particulier, ils peuvent éviter de matérialiser la matrice complète des pondérations d'attention lorsque la séquence est longue.

Avec la valeur par défaut implementation=None, JAX sélectionne automatiquement le meilleur backend disponible. Sur un GPU avec cuDNN disponible et des entrées compatibles, il peut déjà utiliser cuDNN. Sur d'autres matériels, il utilise 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",
)

Vous devriez voir la même forme et le même type de données que la version naïve, une petite différence absolue maximale et True pour la vérification de proximité atol=1e-3.

6. Forcer le backend fusionné cuDNN

Laisser JAX choisir le backend est pratique, mais de l'extérieur, vous ne pouvez pas savoir quel noyau a réellement été exécuté. La définition de implementation="cudnn" force l'utilisation des noyaux d'attention fusionnés cuDNN de NVIDIA. Il s'agit de noyaux GPU optimisés manuellement qui fusionnent l'ensemble du calcul d'attention en un seul lancement de noyau avec des modèles d'accès à la mémoire optimisés.

Notez que l'attention fusionnée cuDNN nécessite des exigences matérielles et de forme, telles que la capacité de calcul du GPU, qui doit être supérieure ou égale à 8.0 (Ampere ou version ultérieure) ou float16 ou bfloat16 comme types de données d'entrée. Si les conditions ne sont pas remplies et que vous définissez implementation="cudnn", JAX génère une erreur au lieu de revenir silencieusement à l'état précédent. C'est pourquoi le code ci-dessous effectue un cast vers bfloat16 et encapsule l'appel dans try/except : il définit HAS_CUDNN_SDPA afin que chaque étape ultérieure sache si le chemin d'accès cuDNN est disponible sur cette machine.

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 votre GPU ne peut pas exécuter cuDNN SDPA, le code affiche la raison et l'atelier de programmation continue sur le backend XLA.

7. Ajouter un masquage causal

Dans les modèles autorégressifs (décodeurs de type GPT), chaque position ne peut s'intéresser qu'aux positions précédentes. La définition de is_causal=True applique ce masque triangulaire inférieur à l'intérieur du noyau fusionné. Vous n'avez donc pas besoin de créer vous-même une matrice de masque.

Le code ci-dessous exécute l'attention deux fois, avec et sans le masque, et compare deux positions pour montrer ce que le masque a modifié.

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

Vous devriez constater une différence importante à la première position et une différence proche de zéro à la dernière. Cette asymétrie est le masque avec la position 0 qui perd toutes les clés sauf la sienne, tandis que la position finale pouvait déjà voir l'intégralité de la séquence.

8. Exécuter des benchmarks

Vous disposez maintenant de trois façons de calculer la même fonction. Le choix de l'une ou l'autre dépend entièrement de la nature du problème et de la façon dont vous pouvez l'évaluer.

Chronométrer les variantes à la forme par défaut

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)

Pour cette petite taille de problème, toutes les implémentations sont très rapides et principalement limitées par les frais généraux. La version naïve est donc compétitive. cuDNN bf16 devrait être légèrement plus rapide, mais l'avantage de l'attention fusionnée devient généralement plus évident pour les séquences plus longues, où il est important d'éviter la matrice d'attention complète.

Balayer la longueur de la séquence

L'avantage des noyaux d'attention fusionnés augmente avec la longueur de la séquence. L'implémentation naïve matérialise une matrice d'attention (B, N, T, S) dans la mémoire GPU, qui est une mémoire O(T²). Les noyaux fusionnés comme cuDNN FlashAttention découpent le calcul en tuiles pour ne jamais matérialiser la matrice complète, ce qui maintient la mémoire à O(T).

Le balayage ci-dessous chronomètre chaque implémentation pour des longueurs de séquence allant de 64 à 1 024.

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

Le balayage prend un certain temps, car chaque longueur de séquence compile trois ou quatre implémentations distinctes avant de les chronométrer. Vous devriez obtenir une ligne de tableau par longueur de séquence, de 64 à 1 024.

Ce graphique compare l'évolution de la latence de l'attention en fonction de la longueur de la séquence pour l'implémentation naïve, JAX/XLA SDPA en float32 et bfloat16, et l'attention fusionnée 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()

Une ligne doit s'afficher pour chaque implémentation. Vous devriez voir l'écart se creuser à mesure que la séquence s'allonge.

Balayer la taille de lot

Les lots plus volumineux amortissent les frais généraux de lancement du kernel et améliorent l'utilisation du GPU jusqu'au point où la mémoire du GPU devient le goulot d'étranglement. Le balayage ci-dessous maintient la longueur de séquence fixe à 256 et fait varier la taille de lot.

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

Vous devriez obtenir un tableau et un graphique avec une entrée par taille de lot, de 1 à 16. L'implémentation naïve est exclue de ce balayage. Les deux résultats ne comparent donc que les trois chemins fusionnés.

9. Comparer MHA, GQA et MQA

Jusqu'à présent, chaque tête comportait ses propres clés et valeurs. Cependant, les décodeurs rompent souvent cette symétrie, car lors de l'inférence, le cache clé-valeur est ce qui domine la mémoire.

L'attention multitête (MHA) donne à chaque tête ses propres projections Q, K et V. L'attention par requêtes groupées (GQA) et l'attention par requêtes multiples (MQA) réduisent le nombre de têtes KV pour économiser de la mémoire et du calcul lors de l'inférence.

jax.nn.dot_product_attention gère les trois, et les têtes KV sont diffusées automatiquement lorsque 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",
)

Les trois lignes doivent indiquer la même forme de sortie, (2, 64, 8, 64), même si les formes K et V passent de 8 têtes à 2, puis à 1. C'est là tout l'intérêt : vous pouvez réduire le cache KV sans rien modifier en aval de l'attention.

10. Effectuer une analyse comparative de NVIDIA TransformerEngine et FP8

TransformerEngine de NVIDIA fournit des modules d'attention fusionnés optimisés pour les GPU NVIDIA. L'intégration JAX utilise des modules de style Flax Linen (et non NNX). Le module est donc initialisé une fois, puis appliqué avec ses variables.

Longueur de la séquence de balayage avec TransformerEngine

Cette cellule vérifie si NVIDIA TransformerEngine est disponible, puis compare son DotProductAttention causal bf16 sur plusieurs longueurs de séquence à JAX SDPA avec les backends XLA et cuDNN sur la même charge de travail.

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

Vous devriez obtenir une ligne de tableau et un point de tracé par longueur de séquence de 128 à 2 048. Ce benchmark bf16 compare TransformerEngine à JAX SDPA sur la même charge de travail d'attention causale, mais le potentiel de performances maximal de TransformerEngine est généralement observé sur les GPU Hopper et Blackwell lorsque l'autocast FP8 est disponible.

Inspecter le chemin d'accès FP8

Sur les GPU Hopper (capacité de calcul >= 9.0, comme le H100), TransformerEngine peut exécuter l'attention en FP8 pour un débit supplémentaire. FP8 utilise la recette DelayedScaling qui suit l'historique du maximum absolu par Tensor pour calculer les facteurs de scaling dynamiques.

E4M3 pour la propagation avant (4 bits d'exposant, 3 bits de mantisse)
E5M2 pour la passe arrière (5 bits d'exposant, 2 bits de mantisse)

Si le GPU n'est pas compatible avec FP8, cette cellule montre à quoi ressemblerait le code sans l'exécuter.

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

Sur le L4, vous devriez voir un tableau nommant votre GPU et indiquant qu'aucune compatibilité FP8 n'a été détectée, suivi des deux lignes imprimées montrant l'appel te.autocast que vous utiliseriez sur un GPU Hopper. Conservez cet extrait : c'est la seule modification requise par FP8 au niveau du site d'appel.

11. Effectuer un nettoyage

Supprimez la charge de travail Jupyter, y compris le LoadBalancer et le volume persistant :

kubectl delete -f deploy/jupyter.yaml

Détruisez le cluster, le pool de nœuds, le VPC et le compte de service :

cd terraform
terraform destroy

Saisissez yes lorsque vous y êtes invité, puis vérifiez qu'il ne reste rien :

gcloud container clusters list
gcloud compute instances list

Les deux doivent être vides pour ce projet. Si vous avez créé un projet uniquement pour cette série, vous pouvez supprimer l'intégralité du projet depuis la console Cloud.

12. Félicitations

Vous êtes passé d'un calcul d'attention écrit à la main à des noyaux fusionnés optimisés pour les GPU, et vous avez mesuré la différence sur du matériel réel.

Connaissances acquises

L'attention naïve fonctionne, mais lance plusieurs kernels GPU et matérialise la matrice d'attention complète en mémoire.
jax.nn.dot_product_attention fusionne le calcul en une seule opération, et avec implementation=None, JAX choisit automatiquement le meilleur backend.
implementation="cudnn" force les noyaux d'attention fusionnés cuDNN de NVIDIA, qui sont les plus rapides sur les séquences longues, et nécessite des entrées bfloat16 ou float16 et une capacité de calcul 8.0 ou ultérieure.
Le masquage causal avec is_causal=True est intégré au kernel fusionné. Aucune matrice de masquage manuelle n'est nécessaire.
GQA et MQA réduisent les têtes KV pour économiser de la mémoire pendant l'inférence, et dot_product_attention gère automatiquement la diffusion.
TransformerEngine fournit une attention fusionnée avec une précision FP8 facultative sur les GPU Hopper.

Étapes suivantes

L'atelier de programmation 6 : Faire évoluer l'entraînement JAX sur plusieurs GPU vous montrera comment segmenter des tableaux sur les deux L4 et exécuter des étapes d'entraînement en parallèle.
Modifiez HEAD_DIM (essayez 32, 64, 128) et observez les valeurs acceptées par cuDNN et l'évolution des temps de traitement.
Augmentez SEQ_LEN (essayez 256, 512, 1024, 2048) et observez l'utilisation de la mémoire et le taux d'accélération de cuDNN.
Réexécutez la comparaison entre MHA, GQA et MQA avec K=2 et K=1 têtes KV par rapport à huit têtes de requête, et réfléchissez aux économies de cache KV pour l'inférence.