1. 簡介
在本程式碼實驗室中，您將著重於 NVIDIA GPU 的注意力運算。您會先從簡單的實作開始，接著改用 JAX 內建的融合注意力機制，然後啟用 NVIDIA 的 cuDNN 後端執行相同作業。完成後，您將瞭解每個層級的變化，以及隨著問題規模擴大，GPU 最佳化路徑的速度會提升多少。
學習內容
- 使用基本 JAX 運算，從頭開始實作縮放點積注意力機制
- 改用
jax.nn.dot_product_attention，這是 JAX 內建的融合核心 - 使用
implementation="cudnn"強制 cuDNN 後端，並新增因果遮罩 - 掃描序列長度和批次大小，查看融合核心的成效
- 比較多頭注意力形狀 MHA、GQA 和 MQA
- 對 NVIDIA TransformerEngine 注意力機制進行基準測試，並檢查其 FP8 程式碼路徑
軟硬體需求
- 已啟用計費功能的 Google Cloud 雲端專案，以及涵蓋 GPU 用量的研討會抵免額或預留項目
- 所選區域中至少有 2 個 NVIDIA L4 GPU 的配額 (如何查看 GPU 配額)
- 完成程式碼實驗室 1 至 4，或具備同等的 CUDA 支援 JAX GPU 環境
- 熟悉矩陣相乘和 Softmax。不需要 CUDA 或 cuDNN 經驗。
預計完成時間：60 分鐘。
注意力機制的運作方式
縮放點積注意力機制會採用三個輸入內容：查詢、鍵和值，並計算：
Attention(Q, K, V) = softmax(Q K^T / sqrt(d_k)) V
縮放比例係數 1 / sqrt(d_k) 可防止點積隨著頭部維度增加而過大，避免 Softmax 進入梯度極小的區域。
四個步驟，每個步驟都會影響下一個步驟：
- 分數 —
Q @ K.T：點積會測量每個查詢位置應注意每個鍵位置的程度。 - 縮放 + Softmax —
softmax(scores / sqrt(d))：縮放可防止梯度消失，Softmax 則會將分數轉換為總和為 1 的注意力權重。 - Attend —
weights @ V：值向量的加權總和會產生每個查詢位置的輸出內容。 - 輸出：形狀與 Q 相同，每個查詢位置現在都包含其關注位置的資訊。
2. 事前準備
選取專案
在 Google Cloud 控制台中，選取或建立已啟用計費功能的專案。
開啟 Cloud Shell
點選控制台右上方的終端機圖示「啟用 Cloud Shell」，啟動 Cloud Shell 工作階段，然後將其指向您的專案：
gcloud config set project <YOUR_PROJECT_ID>
佈建 GPU 環境
在 Cloud Shell 執行下列指令，程式碼實驗室 1 會詳細說明各項指令，包括 GPU 配額和可用區規定。
gcloud services enable \
container.googleapis.com \
compute.googleapis.com \
iam.googleapis.com \
cloudresourcemanager.googleapis.com \
logging.googleapis.com \
monitoring.googleapis.com
git clone https://github.com/Google-Cloud-AI/partner-ai-nvidia.git
cd partner-ai-nvidia/05-workshops/jax-on-gpu/terraform
cp terraform.tfvars.example terraform.tfvars
編輯 terraform.tfvars 並設定 project_id = "，然後佈建叢集並部署 JupyterLab：
terraform init
terraform apply
$(terraform output -raw get_credentials_command)
cd ..
kubectl apply -f deploy/jupyter.yaml
terraform apply約需 12 分鐘。完成後，請等待 Pod 和 LoadBalancer，然後從 Pod 記錄讀取一次性 JupyterLab 權杖：
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
開啟 http://，貼上權杖，然後在 /workspace 中建立新的 Python 3 筆記本。本程式碼實驗室中的每個程式碼區塊，都會放入該筆記本的儲存格。
安裝本程式碼實驗室所需的項目
!pip install --quiet matplotlib flax
設定及驗證 GPU
匯入 JAX，並確認預設後端為 GPU。這個儲存格也會定義 block_tree、show_table 和 show_bars，這些輔助函式會在後續步驟中，用於等待裝置工作和算繪結果。
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)))
您應該會看到 JAX 版本、預設後端 gpu、CUDA 裝置清單，以及本程式碼實驗室其餘部分使用的 GPU。
3. 建構 Q、K 和 V 測試陣列
融合注意力核心會考量輸入陣列在記憶體中的配置方式，因此請先在 JAX 預期的版面配置中建構 Q、K 和 V，再編寫任何注意力程式碼。jax.nn.dot_product_attention 預期採用下列版面配置：
調暗 | 意義 | 預設 |
B | 批量 | 4 |
T | 查詢序列長度 | 128 |
S | 鍵/值序列長度 | 128 (與自注意力機制的 T 相同) |
否 | 注意力頭數量 | 8 |
H | 每個頭部的維度 | 64 |
您會先在 float32 中進行基本實作，並使用隨機值初始化。
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",
)
您應該會看到一個表格，其中每個陣列各有一列，分別回報形狀 (4, 128, 8, 64) 和 dtype float32。jax.device_put 呼叫會將這三個陣列固定到您在設定儲存格中選取的 GPU，因此後續的基準測試不會測量主機傳輸。
4. 從頭開始實作注意力機制
使用融合核心前，請直接寫出公式，以便清楚瞭解核心取代的內容。這個實作會依序執行四個步驟：計算分數、縮放、softmax、乘以值。您會轉置頭部和序列維度，讓 matmul 在序列軸上運作。
雖然能正常運作，但會針對分數 matmul、softmax 和值 matmul 發出三個獨立的 GPU 核心啟動作業，並在 GPU 記憶體中具體化完整的 (B, N, T, S) 注意力權重矩陣。
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",
)
您應該會看到輸出形狀 (4, 128, 8, 64) 和 dtype float32，其形狀與 Q 相同，這正是公式的「輸出」步驟所承諾的。
5. 切換至 dot_product_attention
jax.nn.dot_product_attention 會將分數、縮放、softmax 和注意力步驟融合為單一運算。JAX 和 XLA 就能最佳化記憶體存取模式。特別是當序列很長時，可以避免具體化完整的注意力權重矩陣。
如果使用預設的 implementation=None，JAX 會自動挑選最合適的後端。在具有 cuDNN 和相容輸入的 GPU 上，它可能已經使用 cuDNN。在其他硬體上，則會改用 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",
)
您應該會看到與原始版本相同的形狀和 dtype、較小的最大絕對差異，以及 atol=1e-3 接近度檢查的 True。
6. 強制使用 cuDNN 融合後端
讓 JAX 選擇後端很方便，但從外部無法得知實際執行的核心。設定 implementation="cudnn" 會強制使用 NVIDIA 的 cuDNN 融合注意力核心。這些是手動最佳化的 GPU 核心，可將整個注意力運算融合至單一核心啟動，並採用最佳化的記憶體存取模式。
請注意，cuDNN 融合注意力機制有硬體和形狀要求，例如 GPU 運算能力應 >= 8.0 (Ampere 或更新版本)，或 float16 或 bfloat16 做為輸入 dtype。如果不符合需求，且您設定了 implementation="cudnn"，JAX 會引發錯誤，而不會默默改用其他方法。因此，下列程式碼會轉換為 bfloat16，並將呼叫包裝在 try/except 中：設定 HAS_CUDNN_SDPA，讓後續步驟知道這部機器是否可以使用 cuDNN 路徑。
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.")
如果 GPU 無法執行 cuDNN SDPA，程式碼會輸出原因，並在 XLA 後端繼續執行程式碼實驗室。
7. 新增因果遮罩
在自動迴歸模型 (GPT 樣式解碼器) 中，每個位置只能注意先前的位置。設定 is_causal=True 會在融合核心中套用這個下三角遮罩，因此您不必自行建構遮罩矩陣。
以下程式碼會執行兩次注意力機制 (一次使用遮罩，一次不使用)，並比較兩個位置，顯示遮罩改變了什麼。
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",
)
您應該會看到第一個位置有很大的差異，而最後一個位置的差異接近零。這種不對稱性是指位置 0 的遮罩會失去除了自身以外的所有鍵，而最終位置可能已經能看到整個序列。
8. 執行一些基準測試
現在您有三種方法可以計算同一個函式。要使用哪一種，完全取決於問題的形式，您可以自行評估。
預設形狀的時間變化版本
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)
由於問題規模較小，所有實作方式都非常快速，且主要受限於負載，因此基本版本具有競爭力。cuDNN bf16 應該會稍微快一點，但融合注意力機制的優勢通常在序列長度較長時更明顯，因為避免使用完整注意力矩陣很重要。
掃描序列長度
序列長度越長，融合注意力核心的優勢就越明顯。基本實作會在 GPU 記憶體中具體化 (B, N, T, S) 注意力矩陣，也就是 O(T²) 記憶體。cuDNN FlashAttention 等融合核心會將運算分塊，因此不會具體化完整矩陣，記憶體用量可維持在 O(T)。
下方的掃描會測量各個實作項目的時間，序列長度介於 64 到 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),
)
掃描需要一些時間，因為每個序列長度都會先編譯三或四個獨立實作，然後才計時。最後，每個序列長度 (從 64 到 1024) 應有一個資料表資料列。
這張圖表比較了注意力延遲時間隨序列長度變化的情況，包括原始實作、float32 和 bfloat16 的 JAX/XLA SDPA，以及 bfloat16 的 cuDNN 融合注意力。
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()
每個導入項目應各占一行。您應該會看到差距隨著序列變長而增加。
掃描批量
較大的批次可分攤核心啟動負擔，並提高 GPU 使用率，直到 GPU 記憶體成為瓶頸為止。在下方的掃描中，序列長度固定為 256，批次大小則會變動。
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()
您應該會看到一個表格和一個圖表，其中每個批次大小 (從 1 到 16) 都有一個項目。這項掃描作業不含原始實作，因此兩個輸出內容只會比較三個融合路徑。
9. 比較 MHA、GQA 和 MQA
到目前為止，每個注意力頭都有自己的鍵和值。但解碼器往往會打破這種對稱性，因為在推論期間，鍵/值快取會占據大部分記憶體。
多頭注意力機制 (MHA) 會為每個頭部提供專屬的 Q、K 和 V 投影。分組查詢注意力 (GQA) 和多查詢注意力 (MQA) 機制可減少 KV 頭數量，在推論期間節省記憶體和運算資源。
jax.nn.dot_product_attention 會處理這三種情況，當 K < N 時，KV 頭會自動廣播。
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",
)
雖然 K 和 V 的形狀從 8 個頭縮減為 2 個，再縮減為 1 個，但這三列應回報相同的輸出形狀 (2, 64, 8, 64)。重點是，您可以縮減 KV 快取，而不必變更注意力機制下游的任何內容。
10. 基準測試 NVIDIA TransformerEngine 和 FP8
NVIDIA 的 TransformerEngine 提供針對 NVIDIA GPU 最佳化的融合注意力模組。JAX 整合功能使用 Flax Linen 樣式模組 (而非 NNX)，因此模組會初始化一次，然後套用其變數。
使用 TransformerEngine 掃描序列長度
這個儲存格會檢查 NVIDIA TransformerEngine 是否可用，然後在相同工作負載下，針對多個序列長度的 bf16 因果 DotProductAttention，與採用 XLA 和 cuDNN 後端的 JAX SDPA 進行基準測試。
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()
您應該會看到一個表格列，以及一個序列長度為 128 到 2048 的繪圖點。這項 bf16 基準測試比較 TransformerEngine 與 JAX SDPA 在相同因果注意力工作負載上的表現，但 TransformerEngine 的完整效能潛力通常會在 Hopper 和 Blackwell GPU 上，且啟用 FP8 自動轉換時才會顯現。
檢查 FP8 路徑
在 Hopper GPU (運算能力 >= 9.0，例如 H100) 上，TransformerEngine 可以執行 FP8 注意力機制，提高處理量。FP8 採用 DelayedScaling 方案，追蹤每個張量的絕對最大值記錄，藉此計算動態縮放比例係數。
- 正向傳遞採用 E4M3 格式 (4 個指數，3 個尾數位元)
- 反向傳播採用 E5M2 格式 (5 個指數位元、2 個尾數位元)
如果 GPU 不支援 FP8，這個儲存格會顯示程式碼的樣子，但不會執行。
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.")
在 L4 上，您應該會看到一個表格，其中列出 GPU 名稱，並回報未偵測到 FP8 支援，接著是兩行顯示 te.autocast 呼叫的輸出內容，您會在 Hopper GPU 上使用這個呼叫。請保留該程式碼片段：這是 FP8 在呼叫端唯一需要變更的部分。
11. 清除
刪除 Jupyter 工作負載，包括 LoadBalancer 和永久磁碟區：
kubectl delete -f deploy/jupyter.yaml
刪除叢集、節點集區、虛擬私有雲和服務帳戶：
cd terraform
terraform destroy
出現提示時，請輸入 yes，然後確認沒有遺漏任何內容：
gcloud container clusters list
gcloud compute instances list
這項專案的這兩個欄位都應留空。如果您是為了這個系列課程而建立專案，可以改為從 Cloud 控制台刪除整個專案。
12. 恭喜
您從手寫的注意力運算，改用 GPU 最佳化的融合核心，並在實際硬體上測量差異。
目前所學內容
- Naive attention 雖然可行，但會啟動多個 GPU 核心，並在記憶體中具體化完整的注意力矩陣
jax.nn.dot_product_attention會將運算融合為單一作業，而implementation=NoneJAX 則會自動挑選最佳後端implementation="cudnn"會強制使用 NVIDIA 的 cuDNN 融合注意力核心，這在長序列上速度最快，且需要bfloat16或float16輸入內容和運算能力 8.0 以上版本is_causal=True的因果遮罩內建於融合核心，不需手動建立遮罩矩陣- GQA 和 MQA 會減少 KV 頭，在推論期間節省記憶體，而
dot_product_attention會自動處理廣播 - TransformerEngine 在 Hopper GPU 上提供融合注意力機制，並可選用 FP8 精確度
後續步驟
- 程式碼實驗室 6：在多個 GPU 上擴充 JAX 訓練：說明如何在 L4 上分割陣列，並平行執行訓練步驟
- 變更
HEAD_DIM(嘗試 32、64、128)，觀察 cuDNN 接受的值，以及時間如何變化 - 提高
SEQ_LEN(嘗試 256、512、1024、2048)，並觀察記憶體用量和 cuDNN 加速比率 - 使用
K=2和K=1個 KV 頭對比 8 個查詢頭，重新執行 MHA、GQA 和 MQA 比較，並推論推論作業節省的 KV 快取