Google używa technologii AI do tłumaczenia treści na Twój preferowany język. Tłumaczenia wygenerowane przez AI mogą zawierać błędy.

Przyspieszanie uwagi na GPU za pomocą cuDNN i TransformerEngine

1. Wprowadzenie

Ścieżka szkoleniowa Jax na GPU. Laboratorium 5. Uwaga dotycząca GPU.

W tym ćwiczeniu skupisz się na obliczeniach uwagi na GPU NVIDIA. Zaczniesz od prostej implementacji, zastąpisz ją wbudowaną w JAX fuzją uwagi, a następnie aktywujesz backend cuDNN firmy NVIDIA dla tej samej operacji. Na koniec dowiesz się, co się zmienia na każdym poziomie i jak bardzo przyspieszają ścieżki zoptymalizowane pod kątem GPU wraz ze wzrostem złożoności problemu.

Jakie zadania wykonasz

Zaimplementuj od podstaw skalowaną uwagę z iloczynem skalarnym za pomocą podstawowych operacji JAX
Zastąp ją jax.nn.dot_product_attention, czyli wbudowanym w JAX jądrem fuzji
Wymuś użycie backendu cuDNN za pomocą implementation="cudnn" i dodaj maskowanie przyczynowe
Zmień długość sekwencji i wielkość wsadu, aby zobaczyć, gdzie opłacają się jądra fuzji
Porównaj kształty uwagi wielogłowicowej MHA, GQA i MQA
Przeprowadź test porównawczy uwagi NVIDIA TransformerEngine i sprawdź jej ścieżkę kodu FP8

Czego potrzebujesz

Projekt Google Cloud z włączonymi rozliczeniami oraz środkami na warsztaty lub rezerwacją obejmującą korzystanie z GPU
Limit co najmniej 2 GPU NVIDIA L4 w wybranym regionie (jak sprawdzić limit GPU)
Ukończenie ćwiczeń od 1 do 4 lub równoważne środowisko GPU JAX z obsługą CUDA
Znajomość mnożenia macierzy i funkcji softmax. Nie jest wymagane doświadczenie z CUDA ani cuDNN.

Szacowany czas potrzebny do ukończenia: 60 minut.

Jak działa uwaga

Skalowana uwaga z iloczynem skalarnym przyjmuje 3 dane wejściowe – zapytanie, klucz i wartość – i oblicza:

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

Współczynnik skalowania 1 / sqrt(d_k) zapobiega nadmiernemu wzrostowi iloczynów skalarnych wraz ze wzrostem wymiaru głowy, co spowodowałoby przesunięcie funkcji softmax do obszarów, w których jej gradienty są bardzo małe.

4 kroki, z których każdy jest daną wejściową dla następnego:

Wynik – Q @ K.T: iloczyn skalarny mierzy, jak bardzo każda pozycja zapytania powinna zwracać uwagę na każdą pozycję klucza.
Skalowanie + softmax – softmax(scores / sqrt(d)): skalowanie zapobiega zanikaniu gradientu, a funkcja softmax przekształca wyniki w wagi uwagi, które sumują się do 1.
Uwaga – weights @ V: ważona suma wektorów wartości daje wynik dla każdej pozycji zapytania.
Dane wyjściowe – mają taki sam kształt jak Q: każda pozycja zapytania zawiera teraz informacje z pozycji, na które zwracała uwagę.

2. Zanim zaczniesz

Wybór projektu

W konsoli Google Cloud wybierz lub utwórz projekt z włączonymi rozliczeniami.

Otwórz Cloud Shell

Aby rozpocząć sesję Cloud Shell, kliknij Aktywuj Cloud Shell (ikonę terminala w prawym górnym rogu konsoli), a następnie skieruj ją na swój projekt:

gcloud config set project <YOUR_PROJECT_ID>

Zainicjuj środowisko GPU

Uruchom to polecenie w Cloud Shell. Ćwiczenie 1 szczegółowo wyjaśnia każde polecenie, w tym wymagania dotyczące limitu GPU i strefy.

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

Edytuj terraform.tfvars i ustaw project_id = "", a następnie zainicjuj klaster i wdróż JupyterLab:

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

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

terraform apply trwa około 12 minut. Po zakończeniu poczekaj na utworzenie poda i LoadBalancera, a następnie odczytaj jednorazowy token JupyterLab z logu poda:

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

Otwórz http://:8884, wklej token i utwórz nowy notatnik Pythona 3 w katalogu /workspace. Każdy blok kodu w tym ćwiczeniu znajduje się w komórce tego notatnika.

Zainstaluj to, czego potrzebujesz do tego ćwiczenia

!pip install --quiet matplotlib flax

Skonfiguruj i zweryfikuj GPU

Zaimportuj JAX i sprawdź, czy domyślny backend to GPU. Ta komórka definiuje też block_tree, show_table i show_bars, czyli pomocnicze funkcje, których używa każdy kolejny krok, aby poczekać na wykonanie pracy przez urządzenie i wyrenderować wyniki.

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

Powinna pojawić się wersja JAX, gpu jako domyślny backend, lista urządzeń CUDA i GPU, z którego korzysta reszta ćwiczenia.

3. Utwórz tablice testowe Q, K i V

Jądra fuzji uwagi zwracają uwagę na sposób ułożenia tablic wejściowych w pamięci, dlatego przed napisaniem kodu uwagi utwórz Q, K i V w układzie oczekiwanym przez JAX. jax.nn.dot_product_attention oczekuje następującego układu:

Wymiar	Znaczenie	Domyślne ustawienie
B	Wielkość wsadu	4
T	Długość sekwencji zapytania	128
S	Długość sekwencji klucza/wartości	128 (taka sama jak T w przypadku uwagi własnej)
N	Liczba głów uwagi	8
H	Wymiar na głowę	64

Zaczniesz od float32 w przypadku prostej implementacji, inicjując ją losowymi wartościami.

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

Powinna pojawić się tabela z 1 wierszem na tablicę, z których każdy zawiera kształt (4, 128, 8, 64) i typ danych float32. Wywołanie jax.device_put przypina wszystkie 3 tablice do GPU wybranego w komórce konfiguracji, więc nic w późniejszych testach porównawczych nie mierzy transferu hosta.

4. Zaimplementuj uwagę od podstaw

Zanim sięgniesz po jądro fuzji, napisz bezpośrednio formułę, aby zobaczyć, co dokładnie zastępuje jądro. Ta implementacja wykonuje 4 kroki w kolejności: obliczanie wyników, skalowanie, softmax, mnożenie przez wartości. Transponujesz wymiary głowy i sekwencji, aby matmul działał na osi sekwencji.

Działa poprawnie, ale uruchamia 3 oddzielne jądra GPU dla matmul wyniku, softmax i matmul wartości oraz materializuje pełną macierz wag uwagi (B, N, T, S) w pamięci 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",
)

Powinien pojawić się kształt danych wyjściowych (4, 128, 8, 64) i typ danych float32, czyli taki sam kształt jak Q, co obiecał krok Dane wyjściowe formuły.

5. Przejdź na dot_product_attention

jax.nn.dot_product_attention łączy kroki wyniku, skalowania, softmax i uwagi w jedną operację. JAX i XLA mogą wtedy zoptymalizować wzorzec dostępu do pamięci. W szczególności mogą uniknąć materializowania pełnej macierzy wag uwagi, gdy sekwencja jest długa.

Domyślnie implementation=None JAX automatycznie wybiera najlepszy dostępny backend. Na GPU z dostępnym i zgodnym z danymi wejściowymi cuDNN może już używać cuDNN. Na innym sprzęcie wraca do 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",
)

Powinien pojawić się ten sam kształt i typ danych co w przypadku prostej wersji, niewielka maksymalna różnica bezwzględna i True w przypadku sprawdzenia bliskości atol=1e-3.

6. Wymuś użycie backendu fuzji cuDNN

Pozwolenie JAX na wybór backendu jest wygodne, ale z zewnątrz nie można stwierdzić, które jądro zostało faktycznie uruchomione. Ustawienie implementation="cudnn" wymusza użycie jąder fuzji uwagi cuDNN firmy NVIDIA. Są to ręcznie zoptymalizowane jądra GPU, które łączą całe obliczenia uwagi w jedno uruchomienie jądra ze zoptymalizowanymi wzorcami dostępu do pamięci.

Zwróć uwagę, że fuzja uwagi cuDNN ma wymagania sprzętowe i dotyczące kształtu, takie jak moc obliczeniowa GPU, która powinna być >= 8.0 (Ampere lub nowsza) lub float16 lub bfloat16 jako typy danych wejściowych. Jeśli wymagania nie są spełnione i ustawisz implementation="cudnn", JAX zgłosi błąd, zamiast cicho się wycofać. Dlatego poniższy kod rzutuje na bfloat16 i otacza wywołanie blokiem try/except: ustawia HAS_CUDNN_SDPA, aby każdy kolejny krok wiedział, czy ścieżka cuDNN jest dostępna na tym komputerze.

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

Jeśli GPU nie może uruchomić cuDNN SDPA, kod wyświetli przyczynę, a ćwiczenie będzie kontynuowane na backendzie XLA.

7. Dodaj maskowanie przyczynowe

W modelach autoregresyjnych (dekodery w stylu GPT) każda pozycja może zwracać uwagę tylko na wcześniejsze pozycje. Ustawienie is_causal=True stosuje tę maskę dolnego trójkąta wewnątrz jądra fuzji, dzięki czemu nie musisz samodzielnie tworzyć macierzy maski.

Poniższy kod uruchamia uwagę 2 razy, z maską i bez niej, i porównuje 2 pozycje, aby pokazać, co zmieniła maska.

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

Powinna pojawić się duża różnica na pierwszej pozycji i prawie zerowa na ostatniej. Ta asymetria to maska, w której pozycja 0 traci wszystkie klucze oprócz własnego, a pozycja końcowa może już zobaczyć całą sekwencję.

8. Uruchamianie testów porównawczych

Masz teraz 3 sposoby na obliczenie tej samej funkcji. To, który z nich wybrać, zależy wyłącznie od kształtu problemu, który możesz ocenić.

Zmierz czas wariantów przy domyślnym kształcie

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)

W przypadku tego małego rozmiaru problemu wszystkie implementacje są bardzo szybkie i w większości ograniczone narzutem, więc prosta wersja jest konkurencyjna. cuDNN bf16 powinien być nieco szybszy, ale przewaga fuzji uwagi zwykle staje się wyraźniejsza przy dłuższych sekwencjach, gdzie unikanie pełnej macierzy uwagi ma znaczenie.

Zmień długość sekwencji

Przewaga jąder fuzji uwagi rośnie wraz z długością sekwencji. Prosta implementacja materializuje macierz uwagi (B, N, T, S) w pamięci GPU, która ma rozmiar O(T²). Jądra fuzji, takie jak cuDNN FlashAttention, dzielą obliczenia na części, dzięki czemu nigdy nie materializują pełnej macierzy, utrzymując pamięć na poziomie O(T).

Poniższa zmiana mierzy czas każdej implementacji w przypadku długości sekwencji od 64 do 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),
)

Zmiana trwa dość długo, ponieważ każda długość sekwencji kompiluje 3 lub 4 oddzielne implementacje, zanim je zmierzy. Powinna pojawić się tabela z 1 wierszem na długość sekwencji od 64 do 1024.

Ten wykres porównuje, jak zmienia się opóźnienie uwagi wraz z długością sekwencji w przypadku prostej implementacji, JAX/XLA SDPA w float32 i bfloat16 oraz fuzji uwagi cuDNN w 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()

Powinna pojawić się 1 linia na implementację. Powinna też być widoczna luka, która rośnie wraz z długością sekwencji.

Zmień wielkość wsadu

Większe wsady amortyzują narzut uruchamiania jądra i zwiększają wykorzystanie GPU do momentu, w którym wąskim gardłem staje się pamięć GPU. Poniższa zmiana utrzymuje długość sekwencji na stałym poziomie 256 i zmienia wielkość wsadu.

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

Powinna pojawić się tabela i wykres z 1 wpisem na wielkość wsadu od 1 do 16. Prosta implementacja jest pomijana w tej zmianie, więc oba wyniki porównują tylko 3 ścieżki fuzji.

9. Porównaj MHA, GQA i MQA

Do tej pory każda głowa miała własne klucze i wartości. Dekodery często jednak naruszają tę symetrię, ponieważ podczas wnioskowania to pamięć podręczna klucza-wartości dominuje w pamięci.

Uwaga wielogłowicowa (MHA) daje każdej głowie własne projekcje Q, K i V. Uwaga z grupowanym zapytaniem (GQA) i uwaga z wieloma zapytaniami (MQA) zmniejszają liczbę głów KV, aby zaoszczędzić pamięć i moc obliczeniową podczas wnioskowania.

jax.nn.dot_product_attention obsługuje wszystkie 3 przypadki, a głowy KV są automatycznie rozgłaszane, gdy 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",
)

Wszystkie 3 wiersze powinny mieć ten sam kształt danych wyjściowych, (2, 64, 8, 64), mimo że kształty K i V zmniejszają się z 8 głów do 2 do 1. O to właśnie chodzi: możesz zmniejszyć pamięć podręczną KV bez zmiany niczego poniżej uwagi.

10. Test porównawczy NVIDIA TransformerEngine i FP8

TransformerEngine firmy NVIDIA udostępnia moduły fuzji uwagi zoptymalizowane pod kątem GPU NVIDIA. Integracja z JAX używa modułów w stylu Flax Linen (nie NNX), więc moduł jest inicjowany raz, a następnie stosowany z jego zmiennymi.

Zmień długość sekwencji za pomocą TransformerEngine

Ta komórka sprawdza, czy NVIDIA TransformerEngine jest dostępny, a następnie przeprowadza test porównawczy jego przyczynowej DotProductAttention bf16 w przypadku kilku długości sekwencji w porównaniu z JAX SDPA z backendami XLA i cuDNN w tym samym zbiorze zadań.

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

Powinien pojawić się 1 wiersz tabeli i 1 punkt wykresu na długość sekwencji od 128 do 2048. Ten test porównawczy bf16 porównuje TransformerEngine z JAX SDPA w tym samym zbiorze zadań uwagi przyczynowej, ale pełny potencjał wydajności TransformerEngine jest zwykle widoczny na GPU Hopper i Blackwell, gdy dostępna jest funkcja automatycznego rzutowania FP8.

Sprawdź ścieżkę FP8

Na GPU Hopper (moc obliczeniowa >= 9.0, np. H100) TransformerEngine może uruchamiać uwagę w FP8, aby zwiększyć przepustowość. FP8 używa przepisu DelayedScaling, który śledzi historię maksymalnej wartości bezwzględnej na tensor, aby obliczać dynamiczne współczynniki skalowania.

Format E4M3 dla przejścia do przodu (4 bity wykładnika, 3 bity mantysy)
Format E5M2 dla przejścia do tyłu (5 bitów wykładnika, 2 bity mantysy)

Jeśli GPU nie obsługuje FP8, ta komórka pokazuje, jak wyglądałby kod bez jego uruchamiania.

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

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

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

    if HAS_FP8:
        FP8_SEQ_LEN = 2048
        FP8_BATCH = BATCH

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

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

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

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

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

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

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

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

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

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

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

Na L4 powinna pojawić się tabela z nazwą GPU i informacją o braku wykrytej obsługi FP8, a następnie 2 wydrukowane wiersze pokazujące wywołanie te.autocast, którego używasz na GPU Hopper. Zachowaj ten fragment kodu: to jedyna zmiana, jakiej wymaga FP8 w miejscu wywołania.

11. Zwalnianie miejsca

Usuń zbiór zadań Jupyter, w tym LoadBalancer i wolumin trwały:

kubectl delete -f deploy/jupyter.yaml

Zniszcz klaster, pulę węzłów, VPC i konto usługi:

cd terraform
terraform destroy

Gdy pojawi się prośba, wpisz yes, a następnie sprawdź, czy nic nie zostało:

gcloud container clusters list
gcloud compute instances list

W przypadku tego projektu oba powinny być puste. Jeśli projekt został utworzony tylko na potrzeby tej serii, możesz zamiast tego usunąć cały projekt w konsoli Google Cloud.

12. Gratulacje

Przeszliśmy od ręcznie napisanego obliczania uwagi do zoptymalizowanych pod kątem GPU jąder fuzji i zmierzyliśmy różnicę na prawdziwym sprzęcie.

Czego się nauczysz

Prosta uwaga działa, ale uruchamia wiele jąder GPU i materializuje pełną macierz uwagi w pamięci.
jax.nn.dot_product_attention łączy obliczenia w jedną operację, a w przypadku implementation=None JAX automatycznie wybiera najlepszy backend.
implementation="cudnn" wymusza użycie jąder fuzji uwagi cuDNN firmy NVIDIA, które są najszybsze w przypadku długich sekwencji, i wymaga danych wejściowych bfloat16 lub float16 oraz mocy obliczeniowej 8.0 lub nowszej.
Maskowanie przyczynowe za pomocą is_causal=True jest wbudowane w jądro fuzji – nie jest potrzebna ręczna macierz maski.
GQA i MQA zmniejszają liczbę głów KV, aby zaoszczędzić pamięć podczas wnioskowania, a dot_product_attention automatycznie obsługuje rozgłaszanie.
TransformerEngine udostępnia fuzję uwagi z opcjonalną precyzją FP8 na GPU Hopper.

Dalsze kroki

**Ćwiczenie 6. Skalowanie trenowania JAX na wielu GPU** pokaże, jak dzielić tablice na oba L4 i równolegle wykonywać kroki trenowania.
Zmień HEAD_DIM (wypróbuj 32, 64, 128) i sprawdź, które wartości akceptuje cuDNN i jak zmieniają się czasy.
Zwiększ SEQ_LEN (wypróbuj 256, 512, 1024, 2048) i sprawdź wykorzystanie pamięci oraz współczynnik przyspieszenia cuDNN.
Ponownie uruchom porównanie MHA, GQA i MQA z głowami KV K=2 i K=1 w porównaniu z 8 głowami zapytania i zastanów się nad oszczędnościami w pamięci podręcznej KV w przypadku wnioskowania.