Aufmerksamkeit auf GPU mit cuDNN und TransformerEngine beschleunigen

1. Einführung

Lernpfad für JAX auf GPUs. Lab 5: Attention auf der GPU.

In diesem Codelab konzentrieren Sie sich auf die Berechnung der Aufmerksamkeit auf NVIDIA-GPUs. Sie beginnen mit einer naiven Implementierung, ersetzen sie durch die integrierte zusammengeführte Aufmerksamkeit von JAX und aktivieren dann das cuDNN-Back-End von NVIDIA für denselben Vorgang. Am Ende wissen Sie, was sich auf jeder Ebene ändert und wie viel schneller die GPU-optimierten Pfade werden, wenn das Problem größer wird.

Aufgaben

  • Skalierte Dot-Product-Aufmerksamkeit von Grund auf mit grundlegenden JAX-Vorgängen implementieren
  • Durch jax.nn.dot_product_attention ersetzen, den integrierten zusammengeführten Kernel von JAX
  • Das cuDNN-Back-End mit implementation="cudnn" erzwingen und kausale Maskierung hinzufügen
  • Sequenzlänge und Batchgröße durchlaufen, um zu sehen, wo sich die zusammengeführten Kernel lohnen
  • Die Formen der Multi-Head-Aufmerksamkeit MHA, GQA und MQA vergleichen
  • Die Aufmerksamkeit der NVIDIA TransformerEngine benchmarken und den FP8-Codepfad prüfen

Voraussetzungen

  • Ein Google Cloud-Projekt mit aktivierter Abrechnung und Workshop-Guthaben oder einer Reservierung für die GPU-Nutzung
  • Kontingent für mindestens zwei NVIDIA L4-GPUs in der ausgewählten Region (GPU-Kontingent prüfen)
  • Abschluss der Codelabs 1 bis 4 oder eine entsprechende CUDA-fähige JAX-GPU-Umgebung
  • Kenntnisse in Matrixmultiplikation und Softmax. CUDA- oder cuDNN-Kenntnisse sind nicht erforderlich.

Ungefähre Dauer: 60 Minuten

Funktionsweise der Aufmerksamkeit

Die skalierte Dot-Product-Aufmerksamkeit verwendet drei Eingaben: Query, Key und Value und berechnet Folgendes:

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

Der Skalierungsfaktor 1 / sqrt(d_k) verhindert, dass die Dot-Produkte zu groß werden, wenn die Dimension des Heads zunimmt. Dadurch würde die Softmax-Funktion in Bereiche verschoben, in denen ihre Gradienten sehr klein sind.

Vier Schritte, die jeweils in den nächsten übergehen:

  1. Score – Q @ K.T: Das Dot-Produkt misst, wie stark jede Abfrageposition auf jede Schlüsselposition reagieren sollte.
  2. Scale + Softmax – softmax(scores / sqrt(d)): Die Skalierung verhindert das Verschwinden des Gradienten und die Softmax-Funktion wandelt die Scores in Aufmerksamkeitsgewichte um, die sich zu 1 summieren.
  3. Attend – weights @ V: Die gewichtete Summe der Wertvektoren erzeugt die Ausgabe für jede Abfrageposition.
  4. Output – dieselbe Form wie Q: Jede Abfrageposition enthält jetzt Informationen von den Positionen, auf die sie reagiert hat.

2. Hinweis

Projekt auswählen

Wählen Sie in der Google Cloud Console ein Projekt mit aktivierter Abrechnung aus oder erstellen Sie ein solches Projekt.

Cloud Shell öffnen

Klicken Sie auf Cloud Shell aktivieren (das Terminalsymbol rechts oben in der Konsole), um eine Cloud Shell-Sitzung zu starten, und richten Sie sie dann auf Ihr Projekt aus:

gcloud config set project <YOUR_PROJECT_ID>

GPU-Umgebung bereitstellen

Führen Sie in der Cloud Shell den folgenden Befehl aus. In Codelab 1 werden alle Befehle im Detail erläutert, einschließlich der Anforderungen an das GPU-Kontingent und die 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

Bearbeiten Sie terraform.tfvars und legen Sie project_id = "" fest. Stellen Sie dann den Cluster bereit und stellen Sie JupyterLab bereit:

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

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

terraform apply dauert etwa 12 Minuten. Warten Sie nach Abschluss auf den Pod und den Load Balancer und lesen Sie dann das einmalige JupyterLab-Token aus dem Pod-Log:

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

Öffnen Sie http://:8884, fügen Sie das Token ein und erstellen Sie ein neues Python 3-Notebook in /workspace. Jeder Codeblock in diesem Codelab wird in eine Zelle dieses Notebooks eingefügt.

Für dieses Codelab erforderliche Komponenten installieren

!pip install --quiet matplotlib flax

GPU einrichten und überprüfen

Importieren Sie JAX und bestätigen Sie, dass das Standard-Back-End eine GPU ist. In dieser Zelle werden auch block_tree, show_table und show_bars definiert. Diese Hilfsfunktionen werden in jedem späteren Schritt verwendet, um auf die Gerätearbeit zu warten und Ergebnisse zu rendern.

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

Sie sollten eine JAX-Version, gpu als Standard-Back-End, eine Liste der CUDA-Geräte und die GPU sehen, die im Rest des Codelabs verwendet wird.

3. Test-Arrays für Q, K und V erstellen

Zusammengeführte Aufmerksamkeitskernel berücksichtigen, wie die Eingabearrays im Arbeitsspeicher angeordnet sind. Erstellen Sie daher Q, K und V im Layout, das JAX erwartet, bevor Sie Aufmerksamkeits-Code schreiben. jax.nn.dot_product_attention erwartet das folgende Layout:

Dimmen

Bedeutung

Standardwert

B

Batchgröße

4

T

Länge der Abfragesequenz

128

S

Länge der Schlüssel-/Wertsequenz

128 (identisch mit T für die Selbstaufmerksamkeit)

N

Anzahl der Aufmerksamkeits-Heads

8

H

Dimension pro Head

64

Sie beginnen in float32 für die naive Implementierung und initialisieren mit Zufallswerten.

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

Sie sollten eine Tabelle mit einer Zeile pro Array sehen, in der jeweils die Form (4, 128, 8, 64) und der Datentyp float32 angegeben sind. Mit dem Aufruf jax.device_put werden alle drei Arrays an die GPU angepinnt, die Sie in der Einrichtungszelle ausgewählt haben. Bei den Benchmarks später wird also keine Hostübertragung gemessen.

4. Aufmerksamkeit von Grund auf implementieren

Bevor Sie einen zusammengeführten Kernel verwenden, schreiben Sie die Formel direkt aus, damit Sie genau sehen können, was der Kernel ersetzt. Diese Implementierung folgt den vier Schritten in der Reihenfolge: Scores berechnen, skalieren, Softmax-Funktion anwenden, mit Werten multiplizieren. Sie transponieren die Head- und Sequenzdimensionen, damit die Matrixmultiplikation auf der Sequenzachse ausgeführt wird.

Sie funktioniert korrekt, löst aber drei separate GPU-Kernelstarts für die Score-Matrixmultiplikation, die Softmax-Funktion und die Wert-Matrixmultiplikation aus und materialisiert die vollständige Aufmerksamkeitsgewichtsmatrix (B, N, T, S) im GPU-Arbeitsspeicher.

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

Sie sollten die Ausgabegröße (4, 128, 8, 64) und den Datentyp float32 sehen. Das ist dieselbe Form wie Q, was im Schritt Output der Formel versprochen wurde.

5. Zu dot_product_attention wechseln

Mit jax.nn.dot_product_attention werden die Schritte für Score, Skalierung, Softmax und Aufmerksamkeit zu einem einzigen Vorgang zusammengefasst. JAX und XLA können dann das Speichermuster optimieren. Insbesondere können sie die Materialisierung der vollständigen Aufmerksamkeitsgewichtsmatrix vermeiden, wenn die Sequenz lang ist.

Mit der Standardeinstellung implementation=None wählt JAX automatisch das beste verfügbare Back-End aus. Auf einer GPU mit verfügbarem und kompatiblen cuDNN kann es bereits cuDNN verwenden. Auf anderer Hardware wird auf XLA zurückgegriffen.

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

Sie sollten dieselbe Form und denselben Datentyp wie bei der naiven Version, eine kleine maximale absolute Differenz und True für die Näheprüfung atol=1e-3 sehen.

6. Das zusammengeführte cuDNN-Back-End erzwingen

Es ist praktisch, JAX das Back-End auswählen zu lassen, aber von außen können Sie nicht erkennen, welcher Kernel tatsächlich ausgeführt wurde. Wenn Sie implementation="cudnn" festlegen, werden die zusammengeführten Aufmerksamkeitskernel von NVIDIA cuDNN erzwungen. Das sind handoptimierte GPU-Kernel, die die gesamte Aufmerksamkeitsberechnung in einem einzigen Kernelstart mit optimierten Speichermustern zusammenfassen.

Beachten Sie, dass für die zusammengeführte Aufmerksamkeit von cuDNN Hardware- und Formvoraussetzungen gelten, z. B. die GPU-Rechenleistung, die >= 8.0 (Ampere oder neuer) sein sollte, oder float16 oder bfloat16 als Eingabedatentypen. Wenn die Anforderungen nicht erfüllt sind und Sie implementation="cudnn" festlegen, gibt JAX einen Fehler aus, anstatt im Hintergrund zurückzugreifen. Aus diesem Grund wird der Code unten in bfloat16 umgewandelt und der Aufruf in try/except eingeschlossen. Dadurch wird HAS_CUDNN_SDPA festgelegt, damit in jedem späteren Schritt bekannt ist, ob der cuDNN-Pfad auf diesem Computer verfügbar ist.

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

Wenn Ihre GPU cuDNN SDPA nicht ausführen kann, wird der Grund dafür ausgegeben und das Codelab wird mit dem XLA-Back-End fortgesetzt.

7. Kausale Maskierung hinzufügen

In autoregressiven Modellen (Decodern im GPT-Stil) kann jede Position nur auf frühere Positionen reagieren. Wenn Sie is_causal=True festlegen, wird diese untere Dreiecksmaske im zusammengeführten Kernel angewendet. Sie müssen also keine Maskenmatrix selbst erstellen.

Der Code unten führt die Aufmerksamkeit zweimal aus, einmal mit und einmal ohne Maske, und vergleicht zwei Positionen, um zu zeigen, was sich durch die Maske geändert hat.

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

Sie sollten einen großen Unterschied an der ersten Position und einen Unterschied nahe null an der letzten Position sehen. Diese Asymmetrie ist die Maske, da Position 0 alle Schlüssel außer ihrem eigenen verliert, während die letzte Position bereits die gesamte Sequenz sehen konnte.

8. Einige Benchmarks ausführen

Sie haben jetzt drei Möglichkeiten, dieselbe Funktion zu berechnen. Welche Sie verwenden, hängt ganz von der Form des Problems ab und Sie können es bewerten.

Die Varianten in der Standardform zeitlich erfassen

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)

Bei dieser kleinen Problemgröße sind alle Implementierungen sehr schnell und hauptsächlich durch den Overhead begrenzt. Die naive Version ist also konkurrenzfähig. cuDNN bf16 sollte etwas schneller sein, aber der Vorteil der zusammengeführten Aufmerksamkeit wird in der Regel bei längeren Sequenzlängen deutlicher, wenn es wichtig ist, die vollständige Aufmerksamkeitsmatrix zu vermeiden.

Die Sequenzlänge durchlaufen

Der Vorteil der zusammengeführten Aufmerksamkeitskernel nimmt mit der Sequenzlänge zu. Die naive Implementierung materialisiert eine Aufmerksamkeitsmatrix (B, N, T, S) im GPU-Arbeitsspeicher, was O(T²) Arbeitsspeicher entspricht. Zusammengeführte Kernel wie cuDNN FlashAttention teilen die Berechnung auf, sodass sie die vollständige Matrix nie materialisieren und der Arbeitsspeicher bei O(T) bleibt.

Bei der folgenden Durchlaufung wird jede Implementierung für Sequenzlängen von 64 bis 1024 zeitlich erfasst.

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

Die Durchlaufung dauert eine Weile, da für jede Sequenzlänge drei oder vier separate Implementierungen kompiliert werden, bevor sie zeitlich erfasst werden. Sie sollten eine Tabellenzeile pro Sequenzlänge von 64 bis 1024 erhalten.

In dieser Grafik wird verglichen, wie sich die Latenz der Aufmerksamkeit mit der Sequenzlänge für die naive Implementierung, JAX/XLA SDPA in float32 und bfloat16 sowie die zusammengeführte Aufmerksamkeit von cuDNN in bfloat16 ändert.

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

Sie sollten eine Linie pro Implementierung sehen. Und Sie sollten die Lücke sehen, die mit zunehmender Sequenzlänge größer wird.

Die Batchgröße durchlaufen

Bei größeren Batches wird der Overhead beim Kernelstart amortisiert und die GPU-Auslastung verbessert, bis der GPU-Arbeitsspeicher zum Engpass wird. Bei der folgenden Durchlaufung wird die Sequenzlänge auf 256 festgelegt und die Batchgröße variiert.

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

Sie sollten eine Tabelle und eine Grafik mit einem Eintrag pro Batchgröße von 1 bis 16 erhalten. Die naive Implementierung wird bei dieser Durchlaufung nicht berücksichtigt. In beiden Ausgaben werden also nur die drei zusammengeführten Pfade verglichen.

9. MHA, GQA und MQA vergleichen

Bisher hatte jeder Head eigene Schlüssel und Werte. Decoder brechen diese Symmetrie jedoch oft, da während der Inferenz der Schlüssel-Wert-Cache den Arbeitsspeicher dominiert.

Bei der Multi-Head-Aufmerksamkeit (MHA) hat jeder Head eigene Q-, K- und V-Projektionen. Bei der Grouped-Query-Aufmerksamkeit (GQA) und der Multi-Query-Aufmerksamkeit (MQA) wird die Anzahl der KV-Heads reduziert, um Arbeitsspeicher und Rechenleistung während der Inferenz zu sparen.

jax.nn.dot_product_attention verarbeitet alle drei. Die KV-Heads werden automatisch übertragen, wenn K < N ist.

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

In allen drei Zeilen sollte dieselbe Ausgabegröße (2, 64, 8, 64) angegeben sein, obwohl die K- und V-Formen von 8 Heads auf 2 bzw. 1 reduziert werden. Das ist der Punkt: Sie können den KV-Cache reduzieren, ohne etwas nach der Aufmerksamkeit zu ändern.

10. NVIDIA TransformerEngine und FP8 benchmarken

Die TransformerEngine von NVIDIA bietet zusammengeführte Aufmerksamkeitsmodule, die für NVIDIA-GPUs optimiert sind. Bei der JAX-Integration werden Module im Flax Linen-Stil (nicht NNX) verwendet. Das Modul wird also einmal initialisiert und dann mit seinen Variablen angewendet.

Sequenzlänge mit TransformerEngine durchlaufen

In dieser Zelle wird geprüft, ob die NVIDIA TransformerEngine verfügbar ist. Anschließend wird die kausale DotProductAttention in bf16 für mehrere Sequenzlängen mit JAX SDPA mit den XLA- und cuDNN-Back-Ends für dieselbe Arbeitslast verglichen.

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

Sie sollten eine Tabellenzeile und einen Plotpunkt pro Sequenzlänge von 128 bis 2048 erhalten. Bei diesem bf16-Benchmark wird TransformerEngine mit JAX SDPA für dieselbe Arbeitslast mit kausaler Aufmerksamkeit verglichen. Das volle Leistungspotenzial von TransformerEngine wird jedoch in der Regel auf Hopper- und Blackwell-GPUs sichtbar, wenn FP8-Autocast verfügbar ist.

Den FP8-Pfad prüfen

Auf Hopper-GPUs (Rechenleistung >= 9.0, z. B. H100) kann TransformerEngine die Aufmerksamkeit in FP8 ausführen, um den Durchsatz zu erhöhen. FP8 verwendet das Rezept DelayedScaling, mit dem der absolute Maximalwert pro Tensor im Zeitverlauf erfasst wird, um dynamische Skalierungsfaktoren zu berechnen.

  • E4M3 -Format für den Forward-Pass (4 Exponentenbits, 3 Mantissenbits)
  • E5M2 -Format für den Backward-Pass (5 Exponentenbits, 2 Mantissenbits)

Wenn die GPU FP8 nicht unterstützt, wird in dieser Zelle angezeigt, wie der Code aussehen würde, ohne ihn auszuführen.

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

Auf der L4 sollten Sie eine Tabelle mit dem Namen Ihrer GPU sehen, in der keine FP8-Unterstützung erkannt wurde. Darauf folgen die beiden gedruckten Zeilen mit dem Aufruf te.autocast, den Sie auf einer Hopper-GPU verwenden würden. Behalten Sie diesen Code-Snippet bei. Das ist die einzige Änderung, die für FP8 am Aufrufort erforderlich ist.

11. Bereinigen

Löschen Sie die Jupyter-Arbeitslast, einschließlich des Load Balancers und des persistenten Volumes:

kubectl delete -f deploy/jupyter.yaml

Zerstören Sie den Cluster, den Knotenpool, die VPC und das Dienstkonto:

cd terraform
terraform destroy

Geben Sie bei Aufforderung yes ein und bestätigen Sie dann, dass nichts zurückgeblieben ist:

gcloud container clusters list
gcloud compute instances list

Beide sollten für dieses Projekt leer sein. Wenn Sie ein Projekt nur für diese Reihe erstellt haben, können Sie stattdessen das gesamte Projekt in der Cloud Console löschen.

12. Glückwunsch

Sie haben von einer handgeschriebenen Aufmerksamkeitsberechnung zu GPU-optimierten zusammengeführten Kerneln gewechselt und den Unterschied auf echter Hardware gemessen.

Das haben Sie gelernt

  • Die naive Aufmerksamkeit funktioniert, löst aber mehrere GPU-Kernel aus und materialisiert die vollständige Aufmerksamkeitsmatrix im Arbeitsspeicher.
  • jax.nn.dot_product_attention fasst die Berechnung zu einem einzigen Vorgang zusammen und mit implementation=None wählt JAX automatisch das beste Back-End aus.
  • implementation="cudnn" erzwingt die zusammengeführten Aufmerksamkeitskernel von NVIDIA cuDNN, die bei langen Sequenzen am schnellsten sind. Erforderlich sind bfloat16 oder float16 Eingaben und eine Rechenleistung von 8.0 oder höher.
  • Die kausale Maskierung mit is_causal=True ist in den zusammengeführten Kernel integriert. Es ist keine manuelle Maskenmatrix erforderlich.
  • GQA und MQA reduzieren die Anzahl der KV-Heads, um Arbeitsspeicher während der Inferenz zu sparen, und dot_product_attention übernimmt die Übertragung automatisch.
  • TransformerEngine bietet zusammengeführte Aufmerksamkeit mit optionaler FP8-Genauigkeit auf Hopper-GPUs.

Nächste Schritte

  • In **Codelab 6: JAX-Training auf mehrere GPUs skalieren** erfahren Sie, wie Sie Arrays auf beide L4-GPUs aufteilen und Trainingsschritte parallel ausführen.
  • Ändern Sie HEAD_DIM (versuchen Sie 32, 64, 128) und beobachten Sie, welche Werte von cuDNN akzeptiert werden und wie sich die Zeiten ändern.
  • Erhöhen Sie SEQ_LEN (versuchen Sie 256, 512, 1024, 2048) und beobachten Sie die Arbeitsspeichernutzung und das cuDNN-Beschleunigungsverhältnis.
  • Führen Sie den Vergleich von MHA, GQA und MQA mit K=2 und K=1 KV-Heads im Vergleich zu 8 Abfrage-Heads noch einmal aus und überlegen Sie, wie viel Arbeitsspeicher durch den KV-Cache bei der Inferenz gespart werden kann.

Referenzdokumente