Mempercepat perhatian pada GPU dengan cuDNN dan TransformerEngine

1. Pengantar

Alur pembelajaran Jax di GPU. Lab 5: Attention di GPU.

Dalam codelab ini, Anda akan berfokus pada komputasi attention pada GPU NVIDIA. Anda akan memulai dengan implementasi sederhana, menggantinya dengan fused attention bawaan JAX, lalu mengaktifkan backend cuDNN NVIDIA untuk operasi yang sama. Pada akhirnya, Anda akan mengetahui apa yang berubah di setiap level, dan seberapa cepat jalur yang dioptimalkan GPU seiring bertambahnya masalah.

Yang akan Anda lakukan

  • Mengimplementasikan scaled dot-product attention dari awal dengan operasi JAX dasar
  • Ganti dengan jax.nn.dot_product_attention, kernel gabungan bawaan JAX
  • Memaksa backend cuDNN dengan implementation="cudnn" dan menambahkan masking kausal
  • Menyapu panjang urutan dan ukuran batch untuk melihat di mana kernel gabungan memberikan hasil
  • Membandingkan bentuk multi-head attention MHA, GQA, dan MQA
  • Membandingkan perhatian NVIDIA TransformerEngine dan memeriksa jalur kode FP8-nya

Yang Anda butuhkan

  • Project Google Cloud dengan penagihan diaktifkan, dan kredit workshop atau pemesanan yang mencakup penggunaan GPU
  • Kuota untuk minimal 2 GPU NVIDIA L4 di region yang Anda pilih (cara memeriksa kuota GPU)
  • Penyelesaian codelab 1 hingga 4, atau lingkungan GPU JAX yang mendukung CUDA yang setara
  • Pemahaman tentang perkalian matriks dan softmax. Tidak memerlukan pengalaman CUDA atau cuDNN.

Perkiraan waktu penyelesaian: 60 menit.

Cara kerja attention

Atensi perkalian titik berskala menggunakan tiga input — Query, Key, dan Value — dan menghitung:

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

Faktor penskalaan 1 / sqrt(d_k) mencegah perkalian titik menjadi terlalu besar seiring bertambahnya dimensi head, yang akan mendorong softmax ke wilayah dengan gradien yang sangat kecil.

Empat langkah, masing-masing mengarah ke langkah berikutnya:

  1. Skor — Q @ K.T: perkalian titik mengukur seberapa besar setiap posisi kueri harus memperhatikan setiap posisi kunci.
  2. Skala + Softmax — softmax(scores / sqrt(d)): penskalaan mencegah hilangnya gradien, dan softmax mengonversi skor menjadi bobot atensi yang berjumlah 1.
  3. Attend — weights @ V: jumlah berbobot dari vektor nilai menghasilkan output untuk setiap posisi kueri.
  4. Output — bentuknya sama dengan Q: setiap posisi kueri kini membawa informasi dari posisi yang diperhatikannya.

2. Sebelum memulai

Pilih project Anda

Di Konsol Google Cloud, pilih atau buat project dengan penagihan yang diaktifkan.

Buka Cloud Shell

Klik Activate Cloud Shell (ikon terminal di kanan atas konsol) untuk memulai sesi Cloud Shell, lalu arahkan ke project Anda:

gcloud config set project <YOUR_PROJECT_ID>

Menyediakan lingkungan GPU

Jalankan perintah berikut di Cloud Shell. Codelab 1 menjelaskan setiap perintah secara mendetail, termasuk kuota GPU dan persyaratan zona.

gcloud services enable \
  container.googleapis.com \
  compute.googleapis.com \
  iam.googleapis.com \
  cloudresourcemanager.googleapis.com \
  logging.googleapis.com \
  monitoring.googleapis.com

git clone https://github.com/Google-Cloud-AI/partner-ai-nvidia.git
cd partner-ai-nvidia/05-workshops/jax-on-gpu/terraform

cp terraform.tfvars.example terraform.tfvars

Edit terraform.tfvars dan tetapkan project_id = "", lalu sediakan cluster dan deploy JupyterLab:

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

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

terraform apply memerlukan waktu sekitar 12 menit. Setelah selesai, tunggu Pod dan LoadBalancer, lalu baca token JupyterLab satu kali dari log 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

Buka http://:8884, tempel token, dan buat notebook Python 3 baru di /workspace. Setiap blok kode dalam codelab ini masuk ke dalam sel notebook tersebut.

Menginstal yang diperlukan codelab ini

!pip install --quiet matplotlib flax

Menyiapkan dan memverifikasi GPU

Impor JAX dan pastikan backend default-nya adalah GPU. Sel ini juga menentukan block_tree, show_table, dan show_bars, yaitu helper yang digunakan setiap langkah berikutnya untuk menunggu pekerjaan perangkat dan merender hasil.

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

Anda akan melihat versi JAX, gpu sebagai backend default, daftar perangkat CUDA, dan GPU yang digunakan di bagian codelab lainnya.

3. Membuat array pengujian Q, K, dan V

Kernel fused attention memperhatikan cara array input diletakkan dalam memori, jadi bangun Q, K, dan V dalam tata letak yang diharapkan JAX sebelum menulis kode attention apa pun. jax.nn.dot_product_attention mengharapkan tata letak berikut:

Redup

Arti

Default kami

B

Ukuran batch

4

T

Panjang urutan kueri

128

S

Panjang urutan Kunci/Nilai

128 (sama dengan T untuk self-attention)

T

Jumlah attention head

8

H

Dimensi per head

64

Anda memulai di float32 untuk implementasi sederhana, melakukan inisialisasi dengan nilai acak.

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

Anda akan melihat tabel dengan satu baris per array, yang masing-masing melaporkan bentuk (4, 128, 8, 64) dan dtype float32. Panggilan jax.device_put menyematkan ketiga array ke GPU yang Anda pilih di sel penyiapan, sehingga tidak ada tolok ukur yang mengukur transfer host nantinya.

4. Mengimplementasikan attention dari awal

Sebelum menggunakan kernel gabungan, tulis formulanya secara langsung agar Anda dapat melihat dengan jelas apa yang digantikan oleh kernel tersebut. Implementasi ini mengikuti empat langkah secara berurutan: menghitung skor, menskalakan, softmax, mengalikan dengan nilai. Anda mentransposisikan dimensi head dan urutan sehingga matmul beroperasi di atas sumbu urutan.

Kode ini berfungsi dengan benar, tetapi mengeluarkan tiga peluncuran kernel GPU terpisah untuk matmul skor, softmax, matmul nilai, dan mewujudkan matriks bobot perhatian (B, N, T, S) lengkap dalam memori 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",
)

Anda akan melihat bentuk output (4, 128, 8, 64) dan dtype float32 yang memiliki bentuk yang sama dengan Q, yang merupakan hasil yang dijanjikan oleh langkah Output dari formula.

5. Beralih ke dot_product_attention

jax.nn.dot_product_attention menggabungkan langkah-langkah skor, skala, softmax, dan perhatian menjadi satu operasi. JAX dan XLA kemudian dapat mengoptimalkan pola akses memori. Secara khusus, model ini dapat menghindari materialisasi matriks bobot atensi penuh saat urutannya panjang.

Dengan implementation=None default, JAX akan otomatis memilih backend terbaik yang tersedia. Pada GPU dengan cuDNN yang tersedia dan input yang kompatibel, GPU tersebut mungkin sudah menggunakan cuDNN. Pada hardware lain, model akan kembali ke 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",
)

Anda akan melihat bentuk dan dtype yang sama dengan versi naive, perbedaan absolut maksimum yang kecil, dan True untuk pemeriksaan kedekatan atol=1e-3.

6. Memaksa backend gabungan cuDNN

Membiarkan JAX memilih backend memang praktis, tetapi dari luar Anda tidak dapat mengetahui kernel mana yang benar-benar berjalan. Menyetel implementation="cudnn" akan memaksa kernel fused attention cuDNN NVIDIA. Ini adalah kernel GPU yang dioptimalkan secara manual yang menggabungkan seluruh komputasi atensi ke dalam satu peluncuran kernel dengan pola akses memori yang dioptimalkan.

Perhatikan bahwa cuDNN fused attention memiliki persyaratan hardware dan bentuk seperti kemampuan komputasi GPU yang harus >= 8.0 (Ampere atau yang lebih baru) atau float16 atau bfloat16 sebagai dtype input. Jika persyaratan tidak terpenuhi dan Anda menetapkan implementation="cudnn", JAX akan memunculkan error, bukan melakukan fallback secara otomatis. Itulah sebabnya kode di bawah melakukan transmisi ke bfloat16 dan menggabungkan panggilan dalam try/except: kode ini menetapkan HAS_CUDNN_SDPA sehingga setiap langkah berikutnya mengetahui apakah jalur cuDNN tersedia di mesin ini.

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

Jika GPU Anda tidak dapat menjalankan cuDNN SDPA, kode akan mencetak alasannya dan codelab akan dilanjutkan di backend XLA.

7. Menambahkan masking kausal

Dalam model autoregresif (dekoder bergaya GPT), setiap posisi hanya dapat memperhatikan posisi sebelumnya. Menyetel is_causal=True akan menerapkan mask segitiga bawah ini di dalam kernel gabungan sehingga Anda tidak perlu membuat matriks mask sendiri.

Kode di bawah menjalankan attention dua kali, dengan dan tanpa mask, serta membandingkan dua posisi untuk menunjukkan apa yang diubah oleh mask.

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

Anda akan melihat perbedaan besar di posisi pertama dan perbedaan yang hampir nol di posisi terakhir. Asimetri tersebut adalah mask dengan posisi 0 kehilangan setiap kunci kecuali kuncinya sendiri, sedangkan posisi akhir sudah dapat melihat seluruh urutan.

8. Menjalankan beberapa tolok ukur

Sekarang Anda memiliki tiga cara untuk menghitung fungsi yang sama. Pilihan yang tepat bergantung sepenuhnya pada bentuk masalah dan Anda dapat mengevaluasinya.

Atur waktu varian pada bentuk default

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)

Untuk ukuran masalah yang kecil ini, semua implementasi sangat cepat dan sebagian besar dibatasi oleh overhead, sehingga versi naifnya kompetitif. cuDNN bf16 seharusnya sedikit lebih cepat, tetapi keunggulan fused attention biasanya menjadi lebih jelas pada panjang urutan yang lebih panjang, di mana menghindari matriks perhatian penuh menjadi penting.

Menyapu panjang urutan

Keunggulan kernel atensi gabungan meningkat seiring dengan panjang urutan. Implementasi sederhana mewujudkan matriks perhatian (B, N, T, S) dalam memori GPU yang merupakan memori O(T²). Kernel gabungan seperti FlashAttention cuDNN membagi komputasi menjadi beberapa bagian sehingga tidak pernah mewujudkan matriks lengkap, sehingga memori tetap berada di O(T).

Sweep di bawah ini mengukur waktu setiap implementasi di seluruh panjang urutan dari 64 hingga 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),
)

Penyisiran memerlukan waktu beberapa saat, karena setiap panjang urutan mengompilasi tiga atau empat implementasi terpisah sebelum mengukur waktunya. Anda akan mendapatkan satu baris tabel per panjang urutan dari 64 hingga 1024.

Plot ini membandingkan perubahan latensi perhatian dengan panjang urutan untuk implementasi sederhana, JAX/XLA SDPA dalam float32 dan bfloat16, serta perhatian gabungan cuDNN dalam 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()

Anda akan melihat satu baris per penerapan. Anda akan melihat kesenjangan yang semakin besar seiring dengan semakin panjangnya urutan.

Menyapu ukuran batch

Batch yang lebih besar mengamortisasi overhead peluncuran kernel dan meningkatkan pemanfaatan GPU hingga memori GPU menjadi bottleneck. Sweep di bawah ini mempertahankan panjang urutan tetap pada 256 dan memvariasikan ukuran batch.

BATCH_SIZES = [1, 2, 4, 8, 16]
SWEEP_SEQ = 256

batch_results = []
for bs in BATCH_SIZES:
    rk = jax.random.key(bs + 100)
    rk1, rk2, rk3 = jax.random.split(rk, 3)

    q_b = jax.random.normal(rk1, (bs, SWEEP_SEQ, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    k_b = jax.random.normal(rk2, (bs, SWEEP_SEQ, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    v_b = jax.random.normal(rk3, (bs, SWEEP_SEQ, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    q_b, k_b, v_b = jax.device_put((q_b, k_b, v_b), device)

    q_bb = q_b.astype(jnp.bfloat16)
    k_bb = k_b.astype(jnp.bfloat16)
    v_bb = v_b.astype(jnp.bfloat16)

    row = {"batch": bs}

    row["sdpa_xla_f32_ms"] = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
        q_b, k_b, v_b,
        warmup=2,
        repeats=20,
    )

    row["sdpa_xla_bf16_ms"] = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
        q_bb, k_bb, v_bb,
        warmup=2,
        repeats=20,
    )

    if HAS_CUDNN_SDPA:
        row["cudnn_bf16_ms"] = benchmark_attention(
            lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="cudnn"),
            q_bb, k_bb, v_bb,
            warmup=2,
            repeats=20,
        )

    batch_results.append(row)


headers = ["Batch size", "SDPA XLA f32 (ms)", "SDPA XLA bf16 (ms)"]
if HAS_CUDNN_SDPA:
    headers.append("cuDNN bf16 (ms)")

table_rows = []
for r in batch_results:
    row = [
        r["batch"],
        f"{r['sdpa_xla_f32_ms']:.2f}",
        f"{r['sdpa_xla_bf16_ms']:.2f}",
    ]

    if HAS_CUDNN_SDPA:
        row.append(f"{r['cudnn_bf16_ms']:.2f}")

    table_rows.append(row)

show_table(
    headers,
    table_rows,
    title=f"Batch-size sweep — T={SWEEP_SEQ}, N={NUM_HEADS}, H={HEAD_DIM}",
    aligns=["right"] * len(headers),
)


fig, ax = plt.subplots(figsize=(8, 5))
batches = [r["batch"] for r in batch_results]

ax.plot(
    batches,
    [r["sdpa_xla_f32_ms"] for r in batch_results],
    "s-",
    label="SDPA XLA f32",
    color="#0969da",
)

ax.plot(
    batches,
    [r["sdpa_xla_bf16_ms"] for r in batch_results],
    "d-",
    label="SDPA XLA bf16",
    color="#8250df",
)

if HAS_CUDNN_SDPA:
    ax.plot(
        batches,
        [r["cudnn_bf16_ms"] for r in batch_results],
        "^-",
        label="cuDNN bf16",
        color="#1a7f37",
    )

ax.set_xlabel("Batch size")
ax.set_ylabel("Median ms per call")
ax.set_title("Attention latency vs batch size")
ax.legend()
ax.grid(True, alpha=0.25)
ax.set_xticks(batches)

fig.tight_layout()
plt.show()

Anda akan mendapatkan tabel dan plot dengan satu entri per ukuran batch dari 1 hingga 16. Implementasi sederhana tidak disertakan dalam sweep ini, sehingga kedua output hanya membandingkan tiga jalur gabungan.

9. Membandingkan MHA, GQA, dan MQA

Sejauh ini, setiap head memiliki kunci dan nilainya sendiri. Namun, dekoder sering kali merusak simetri tersebut, karena selama inferensi, cache nilai kunci adalah yang mendominasi memori.

Multi-head attention (MHA) memberikan proyeksi Q, K, dan V-nya sendiri kepada setiap head. Grouped-query attention (GQA) dan multi-query attention (MQA) mengurangi jumlah KV head untuk menghemat memori dan komputasi selama inferensi.

jax.nn.dot_product_attention menangani ketiganya dengan head KV yang disiarkan secara otomatis saat 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",
)

Ketiga baris harus melaporkan bentuk output yang sama, (2, 64, 8, 64), meskipun bentuk K dan V menyusut dari 8 head menjadi 2 menjadi 1. Itulah intinya: Anda dapat memotong cache KV tanpa mengubah apa pun di downstream attention.

10. Tolok ukur NVIDIA TransformerEngine dan FP8

TransformerEngine NVIDIA menyediakan modul fused attention yang dioptimalkan untuk GPU NVIDIA. Integrasi JAX menggunakan modul bergaya Flax Linen (bukan NNX), sehingga modul diinisialisasi sekali, lalu diterapkan dengan variabelnya.

Menyapu panjang urutan dengan TransformerEngine

Sel ini memeriksa apakah NVIDIA TransformerEngine tersedia, lalu melakukan benchmark DotProductAttention kausal bf16-nya di beberapa panjang urutan terhadap JAX SDPA dengan backend XLA dan cuDNN pada workload yang sama.

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

Anda akan mendapatkan satu baris tabel dan satu titik plot per panjang urutan dari 128 hingga 2048. Tolok ukur bf16 ini membandingkan TransformerEngine dengan JAX SDPA pada workload causal-attention yang sama, tetapi potensi performa penuh TransformerEngine biasanya terlihat pada GPU Hopper dan Blackwell saat autocast FP8 tersedia.

Memeriksa jalur FP8

Pada GPU Hopper (kemampuan komputasi >= 9.0, seperti H100), TransformerEngine dapat menjalankan attention di FP8 untuk throughput tambahan. FP8 menggunakan resep DelayedScaling yang melacak histori maksimum absolut per tensor untuk menghitung faktor penskalaan dinamis.

  • Format E4M3 untuk penerusan maju (4 eksponen, 3 bit mantissa)
  • Format E5M2 untuk backward pass (5 eksponen, 2 bit mantissa)

Jika GPU tidak mendukung FP8, sel ini akan menampilkan tampilan kode tanpa menjalankannya.

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

Di L4, Anda akan melihat tabel yang memberi nama GPU Anda dan melaporkan bahwa tidak ada dukungan FP8 yang terdeteksi, diikuti dengan dua baris yang dicetak yang menunjukkan panggilan te.autocast yang akan Anda gunakan di GPU Hopper. Pertahankan cuplikan tersebut: ini adalah satu-satunya perubahan yang diperlukan FP8 di situs panggilan.

11. Pembersihan

Hapus workload Jupyter, termasuk LoadBalancer dan persistent volume:

kubectl delete -f deploy/jupyter.yaml

Hancurkan cluster, node pool, VPC, dan akun layanan:

cd terraform
terraform destroy

Ketik yes saat diminta, lalu pastikan tidak ada yang tertinggal:

gcloud container clusters list
gcloud compute instances list

Keduanya harus kosong untuk project ini. Jika Anda membuat project hanya untuk seri ini, Anda dapat menghapus seluruh project dari Konsol Cloud.

12. Selamat

Anda beralih dari komputasi atensi yang ditulis secara manual ke kernel gabungan yang dioptimalkan untuk GPU, dan mengukur perbedaannya pada hardware sungguhan.

Yang telah Anda pelajari

  • Naive attention berfungsi, tetapi meluncurkan beberapa kernel GPU dan mewujudkan matriks attention penuh dalam memori
  • jax.nn.dot_product_attention menggabungkan komputasi menjadi satu operasi, dan dengan implementation=None, JAX akan otomatis memilih backend terbaik
  • implementation="cudnn" memaksa kernel fused attention cuDNN NVIDIA, yang tercepat pada urutan panjang, dan memerlukan input bfloat16 atau float16 serta kemampuan komputasi 8.0 atau yang lebih baru
  • Penyamaran kausal dengan is_causal=True diintegrasikan ke dalam kernel gabungan, sehingga tidak memerlukan matriks penyamaran manual
  • GQA dan MQA mengurangi KV head untuk menghemat memori selama inferensi, dan dot_product_attention menangani penyiaran secara otomatis
  • TransformerEngine menyediakan fused attention dengan presisi FP8 opsional di GPU Hopper

Langkah berikutnya

  • Codelab 6: Menskalakan pelatihan JAX di beberapa GPU akan menunjukkan cara melakukan sharding array di beberapa L4 dan menjalankan langkah-langkah pelatihan secara paralel
  • Ubah HEAD_DIM (coba 32, 64, 128) dan perhatikan nilai yang diterima cuDNN dan bagaimana pengaturan waktunya berubah
  • Naikkan SEQ_LEN (coba 256, 512, 1024, 2048) dan perhatikan penggunaan memori serta rasio percepatan cuDNN
  • Jalankan ulang perbandingan MHA, GQA, dan MQA dengan K=2 dan K=1 KV head terhadap 8 query head, dan jelaskan penghematan cache KV untuk inferensi

Dokumen referensi