1. Giới thiệu
Trong lớp học lập trình này, bạn sẽ tập trung vào việc tính toán cơ chế chú ý trên GPU NVIDIA. Bạn bắt đầu bằng một phương thức triển khai đơn giản, thay thế bằng cơ chế chú ý hợp nhất tích hợp sẵn của JAX, sau đó kích hoạt phần phụ trợ cuDNN của NVIDIA cho cùng một thao tác. Cuối cùng, bạn sẽ biết những thay đổi ở mỗi cấp và tốc độ của các đường dẫn được tối ưu hoá cho GPU khi vấn đề tăng lên.
Bạn sẽ thực hiện
- Triển khai cơ chế chú ý tích vô hướng theo tỷ lệ từ đầu bằng các phép toán JAX cơ bản
- Thay thế bằng
jax.nn.dot_product_attention, hạt nhân hợp nhất tích hợp sẵn của JAX - Buộc phần phụ trợ cuDNN sử dụng
implementation="cudnn"và thêm mặt nạ nhân quả - Quét độ dài trình tự và kích thước lô để xem các hạt nhân hợp nhất có hiệu quả ở đâu
- So sánh các hình dạng cơ chế chú ý nhiều đầu MHA, GQA và MQA
- Đo điểm chuẩn cho cơ chế chú ý của NVIDIA TransformerEngine và kiểm tra đường dẫn mã FP8 của cơ chế này
Bạn cần có
- Một dự án trên đám mây của Google Cloud đã bật tính năng thanh toán và có tín dụng hội thảo hoặc một gói đặt trước bao gồm mức sử dụng GPU
- Hạn mức cho ít nhất 2 GPU NVIDIA L4 ở khu vực bạn chọn (cách kiểm tra hạn mức GPU)
- Hoàn thành lớp học lập trình từ 1 đến 4 hoặc môi trường GPU JAX tương đương có hỗ trợ CUDA
- Quen thuộc với phép nhân ma trận và softmax. Không cần có kinh nghiệm về CUDA hoặc cuDNN.
Thời gian hoàn thành ước tính: 60 phút.
Cách hoạt động của cơ chế chú ý
Cơ chế chú ý tích vô hướng theo tỷ lệ nhận ba dữ liệu đầu vào — Truy vấn, Khoá và Giá trị — và tính toán:
Attention(Q, K, V) = softmax(Q K^T / sqrt(d_k)) V
Hệ số tỷ lệ 1 / sqrt(d_k) giúp tích vô hướng không tăng quá lớn khi kích thước đầu tăng lên, điều này sẽ đẩy softmax vào các vùng có độ dốc rất nhỏ.
Bốn bước, mỗi bước là tiền đề cho bước tiếp theo:
- Điểm số —
Q @ K.T: tích vô hướng đo lường mức độ mà mỗi vị trí truy vấn cần chú ý đến từng vị trí khoá. - Tỷ lệ + Softmax —
softmax(scores / sqrt(d)): việc điều chỉnh tỷ lệ giúp ngăn chặn hiện tượng biến mất độ dốc, còn softmax chuyển đổi điểm số thành trọng số chú ý có tổng bằng 1. - Attend —
weights @ V: tổng có trọng số của các vectơ giá trị tạo ra kết quả cho mỗi vị trí truy vấn. - Đầu ra — cùng hình dạng với Q: mỗi vị trí truy vấn hiện mang thông tin từ các vị trí mà nó tham dự.
2. Trước khi bắt đầu
Chọn dự án của bạn
Trong Google Cloud Console, hãy chọn hoặc tạo một dự án đã bật tính năng thanh toán.
Mở Cloud Shell
Nhấp vào Activate Cloud Shell (Kích hoạt Cloud Shell – biểu tượng dòng lệnh ở trên cùng bên phải của bảng điều khiển) để bắt đầu một phiên Cloud Shell, sau đó trỏ phiên đó vào dự án của bạn:
gcloud config set project <YOUR_PROJECT_ID>
Cung cấp môi trường GPU
Chạy lệnh sau trong Cloud Shell. Lớp học lập trình 1 giải thích chi tiết từng lệnh, bao gồm cả hạn mức GPU và yêu cầu về vùng.
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
Chỉnh sửa terraform.tfvars và đặt project_id = ", sau đó cấp phép cụm và triển khai JupyterLab:
terraform init
terraform apply
$(terraform output -raw get_credentials_command)
cd ..
kubectl apply -f deploy/jupyter.yaml
terraform apply mất khoảng 12 phút. Khi quá trình này hoàn tất, hãy đợi Pod và LoadBalancer, sau đó đọc mã thông báo JupyterLab một lần từ nhật ký 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
Mở http://, dán mã thông báo và tạo một sổ ghi chú Python 3 mới trong /workspace. Mỗi khối mã trong lớp học lập trình này sẽ được đưa vào một ô của sổ tay đó.
Cài đặt những gì lớp học lập trình này cần
!pip install --quiet matplotlib flax
Thiết lập và xác minh GPU
Nhập JAX và xác nhận rằng phần phụ trợ mặc định là GPU. Ô này cũng xác định block_tree, show_table và show_bars, các trình trợ giúp mà mọi bước sau này đều sử dụng để chờ thiết bị hoạt động và kết xuất kết quả.
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)))
Bạn sẽ thấy phiên bản JAX, gpu làm phần phụ trợ mặc định, danh sách thiết bị CUDA và GPU mà phần còn lại của lớp học lập trình sử dụng.
3. Tạo mảng kiểm thử Q, K và V
Các hạt nhân chú ý hợp nhất quan tâm đến cách bố trí mảng đầu vào trong bộ nhớ, vì vậy, hãy tạo Q, K và V theo bố cục mà JAX mong đợi trước khi viết bất kỳ mã chú ý nào. jax.nn.dot_product_attention yêu cầu bố cục sau:
Giảm độ sáng | Ý nghĩa | Mặc định |
B | Kích thước lô | 4 |
T | Độ dài chuỗi truy vấn | 128 |
N | Độ dài chuỗi khoá/giá trị | 128 (giống như T cho cơ chế tự chú ý) |
Không | Số lượng đầu chú ý | 8 |
Cao | Phương diện theo đầu người | 64 |
Bạn bắt đầu bằng float32 để triển khai đơn giản, khởi chạy bằng các giá trị ngẫu nhiên.
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",
)
Bạn sẽ thấy một bảng có một hàng cho mỗi mảng, mỗi hàng báo cáo hình dạng (4, 128, 8, 64) và dtype float32. Lệnh gọi jax.device_put ghim cả ba mảng vào GPU mà bạn đã chọn trong ô thiết lập, vì vậy, không có điểm chuẩn nào sau này đo lường quá trình chuyển máy chủ.
4. Triển khai cơ chế chú ý từ đầu
Trước khi sử dụng hạt nhân hợp nhất, hãy viết trực tiếp công thức để bạn có thể thấy chính xác nội dung mà hạt nhân thay thế. Quá trình triển khai này tuân theo 4 bước theo thứ tự: tính điểm, điều chỉnh tỷ lệ, softmax, nhân với giá trị. Bạn chuyển vị các chiều đầu và trình tự để phép nhân ma trận hoạt động trên trục trình tự.
Mã này hoạt động chính xác nhưng đưa ra ba lần khởi chạy hạt nhân GPU riêng biệt cho phép nhân ma trận điểm, softmax, phép nhân ma trận giá trị và mã này hiện thực hoá toàn bộ ma trận trọng số chú ý (B, N, T, S) trong bộ nhớ 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",
)
Bạn sẽ thấy hình dạng đầu ra (4, 128, 8, 64) và dtype float32 có cùng hình dạng với Q, đây là điều mà bước Đầu ra của công thức đã hứa hẹn.
5. Chuyển sang dot_product_attention
jax.nn.dot_product_attention kết hợp các bước tính điểm, điều chỉnh tỷ lệ, softmax và chú ý thành một thao tác duy nhất. Sau đó, JAX và XLA có thể tối ưu hoá mẫu truy cập bộ nhớ. Cụ thể, các mô hình này có thể tránh việc tạo ra ma trận trọng số chú ý đầy đủ khi trình tự dài.
Với implementation=None mặc định, JAX sẽ tự động chọn phần phụ trợ phù hợp nhất hiện có. Trên GPU có cuDNN và dữ liệu đầu vào tương thích, có thể GPU đã sử dụng cuDNN. Trên phần cứng khác, nó sẽ quay lại 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",
)
Bạn sẽ thấy hình dạng và dtype giống như phiên bản đơn giản, chênh lệch tuyệt đối tối đa nhỏ và True cho quy trình kiểm tra độ gần atol=1e-3.
6. Buộc sử dụng phần phụ trợ hợp nhất cuDNN
Việc cho phép JAX chọn phần phụ trợ là rất thuận tiện, nhưng từ bên ngoài, bạn không thể biết nhân nào thực sự chạy. Việc đặt implementation="cudnn" sẽ buộc các hạt nhân chú ý hợp nhất cuDNN của NVIDIA. Đây là các nhân GPU được tối ưu hoá thủ công, kết hợp toàn bộ phép tính chú ý vào một lần khởi chạy nhân duy nhất với các mẫu truy cập bộ nhớ được tối ưu hoá.
Xin lưu ý rằng cơ chế chú ý kết hợp cuDNN có các yêu cầu về phần cứng và hình dạng, chẳng hạn như khả năng tính toán của GPU phải >= 8.0 (Ampere trở lên) hoặc float16 hoặc bfloat16 làm dtype đầu vào. Nếu các yêu cầu không được đáp ứng và bạn đặt implementation="cudnn", JAX sẽ báo lỗi thay vì tự động quay lại. Đó là lý do mã bên dưới truyền đến bfloat16 và gói lệnh gọi trong try/except: mã này đặt HAS_CUDNN_SDPA để mọi bước sau đó đều biết liệu đường dẫn cuDNN có sẵn trên máy này hay không.
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.")
Nếu GPU của bạn không thể chạy cuDNN SDPA, thì mã sẽ in ra lý do và lớp học lập trình sẽ tiếp tục trên phần phụ trợ XLA.
7. Thêm mặt nạ nhân quả
Trong các mô hình tự hồi quy (bộ giải mã kiểu GPT), mỗi vị trí chỉ có thể chú ý đến các vị trí trước đó. Việc đặt is_causal=True sẽ áp dụng mặt nạ tam giác dưới này bên trong hạt nhân hợp nhất, vì vậy, bạn không cần phải tự tạo ma trận mặt nạ.
Đoạn mã dưới đây chạy cơ chế chú ý hai lần, có và không có mặt nạ, đồng thời so sánh hai vị trí để cho thấy những thay đổi do mặt nạ mang lại.
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",
)
Bạn sẽ thấy sự khác biệt lớn ở vị trí đầu tiên và sự khác biệt gần như bằng không ở vị trí cuối cùng. Sự bất đối xứng đó là mặt nạ có vị trí 0 sẽ mất mọi khoá ngoại trừ khoá của chính nó, trong khi vị trí cuối cùng đã có thể nhìn thấy toàn bộ trình tự.
8. Đang chạy một số phép đo điểm chuẩn
Giờ đây, bạn có 3 cách để tính toán cùng một hàm. Việc chọn công cụ nào hoàn toàn phụ thuộc vào dạng vấn đề và bạn có thể đánh giá vấn đề đó.
Thời gian các biến thể ở hình dạng mặc định
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)
Đối với kích thước bài toán nhỏ này, tất cả các phương thức triển khai đều rất nhanh và chủ yếu bị giới hạn về mức hao tổn, vì vậy phiên bản đơn giản có tính cạnh tranh. cuDNN bf16 sẽ nhanh nhất một chút, nhưng lợi thế của cơ chế chú ý hợp nhất thường trở nên rõ ràng hơn ở độ dài chuỗi dài hơn, trong đó việc tránh ma trận chú ý đầy đủ là quan trọng.
Quét độ dài trình tự
Lợi thế của nhân chú ý hợp nhất tăng lên theo độ dài trình tự. Cách triển khai đơn giản sẽ tạo ra một ma trận chú ý (B, N, T, S) trong bộ nhớ GPU, tức là bộ nhớ O(T²). Các hạt nhân hợp nhất như cuDNN FlashAttention phân chia quá trình tính toán thành các ô để không bao giờ tạo ra ma trận đầy đủ, giữ bộ nhớ ở mức O(T).
Quá trình quét bên dưới tính thời gian cho mỗi lần triển khai trên các độ dài trình tự từ 64 đến 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),
)
Quá trình quét mất một lúc vì mỗi độ dài trình tự sẽ biên dịch 3 hoặc 4 cách triển khai riêng biệt trước khi tính thời gian. Bạn sẽ nhận được một hàng trong bảng cho mỗi độ dài trình tự từ 64 đến 1024.
Biểu đồ này so sánh sự thay đổi của độ trễ chú ý theo độ dài trình tự đối với cách triển khai đơn giản, JAX/XLA SDPA ở định dạng float32 và bfloat16, cũng như cơ chế chú ý kết hợp cuDNN ở định dạng 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()
Bạn sẽ thấy một dòng cho mỗi lần triển khai. Và bạn sẽ thấy khoảng cách ngày càng lớn khi chuỗi ngày càng dài.
Quét kích thước lô
Các lô lớn hơn sẽ phân bổ mức hao tổn khi khởi chạy hạt nhân và cải thiện mức sử dụng GPU cho đến khi bộ nhớ GPU trở thành điểm tắc nghẽn. Quá trình quét bên dưới giữ độ dài trình tự cố định ở mức 256 và thay đổi kích thước lô.
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()
Bạn sẽ nhận được một bảng và một biểu đồ với một mục nhập cho mỗi kích thước lô từ 1 đến 16. Cách triển khai đơn giản không được đưa vào lần quét này, vì vậy, cả hai kết quả đầu ra chỉ so sánh ba đường dẫn hợp nhất.
9. So sánh MHA, GQA và MQA
Cho đến nay, mỗi phần đầu đều có khoá và giá trị riêng. Tuy nhiên, bộ giải mã thường phá vỡ tính đối xứng đó, vì trong quá trình suy luận, bộ nhớ đệm khoá-giá trị là yếu tố chiếm ưu thế trong bộ nhớ.
Cơ chế chú ý nhiều đầu (MHA) cung cấp cho mỗi đầu các phép chiếu Q, K và V riêng. Cơ chế chú ý truy vấn theo nhóm (GQA) và cơ chế chú ý nhiều truy vấn (MQA) giúp giảm số lượng đầu KV để tiết kiệm bộ nhớ và tài nguyên điện toán trong quá trình suy luận.
jax.nn.dot_product_attention xử lý cả ba trường hợp này bằng cách tự động truyền phát các đầu KV khi 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",
)
Cả ba hàng đều phải báo cáo cùng một hình dạng đầu ra, (2, 64, 8, 64), ngay cả khi hình dạng K và V thu nhỏ từ 8 đầu xuống còn 2 rồi còn 1. Đó chính là điểm mấu chốt: bạn có thể cắt giảm bộ nhớ đệm KV mà không cần thay đổi bất kỳ điều gì sau cơ chế chú ý.
10. Đo điểm chuẩn NVIDIA TransformerEngine và FP8
TransformerEngine của NVIDIA cung cấp các mô-đun chú ý hợp nhất được tối ưu hoá cho GPU NVIDIA. Tính năng tích hợp JAX sử dụng các mô-đun kiểu Flax Linen (không phải NNX), vì vậy, mô-đun này được khởi chạy một lần rồi áp dụng với các biến của mô-đun.
Độ dài chuỗi quét bằng TransformerEngine
Ô này kiểm tra xem NVIDIA TransformerEngine có sẵn hay không, sau đó đo điểm chuẩn cho DotProductAttention nhân quả bf16 của nó trên một số độ dài trình tự so với JAX SDPA với phần phụ trợ XLA và cuDNN trên cùng một khối lượng công việc.
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()
Bạn sẽ nhận được một hàng trong bảng và một điểm biểu đồ cho mỗi độ dài trình tự từ 128 đến 2048. Điểm chuẩn bf16 này so sánh TransformerEngine với JAX SDPA trên cùng một khối lượng công việc chú ý nhân quả, nhưng tiềm năng hiệu suất đầy đủ của TransformerEngine thường được thấy trên GPU Hopper và Blackwell khi có tính năng tự động truyền FP8.
Kiểm tra đường dẫn FP8
Trên GPU Hopper (khả năng điện toán >= 9.0, chẳng hạn như H100), TransformerEngine có thể chạy cơ chế chú ý trong FP8 để tăng thêm công suất. FP8 sử dụng công thức DelayedScaling theo dõi nhật ký giá trị tối đa tuyệt đối trên mỗi tensor để tính toán hệ số tỷ lệ động.
- Định dạng E4M3 cho lượt chuyển tiếp (4 bit số mũ, 3 bit phần định trị)
- Định dạng E5M2 cho lượt truyền ngược (5 bit số mũ, 2 bit phần định trị)
Nếu GPU không hỗ trợ FP8, thì ô này sẽ cho biết mã sẽ trông như thế nào mà không cần chạy mã.
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.")
Trên L4, bạn sẽ thấy một bảng đặt tên cho GPU và báo cáo rằng không phát hiện thấy tính năng hỗ trợ FP8, theo sau là hai dòng in cho thấy lệnh gọi te.autocast mà bạn sẽ sử dụng trên GPU Hopper. Hãy giữ lại đoạn mã đó: đây là thay đổi duy nhất mà FP8 yêu cầu tại vị trí gọi.
11. Dọn dẹp
Xoá khối lượng công việc Jupyter, bao gồm cả LoadBalancer và ổ đĩa ổn định:
kubectl delete -f deploy/jupyter.yaml
Huỷ cụm, bộ nút, VPC và tài khoản dịch vụ:
cd terraform
terraform destroy
Nhập yes khi được nhắc, sau đó xác nhận rằng không còn gì:
gcloud container clusters list
gcloud compute instances list
Cả hai đều phải trống đối với dự án này. Nếu đã tạo một dự án chỉ cho loạt bài này, thì bạn có thể xoá toàn bộ dự án khỏi Cloud Console.
12. Xin chúc mừng
Bạn đã chuyển từ tính toán cơ chế chú ý viết tay sang nhân hợp nhất được tối ưu hoá cho GPU và đo lường sự khác biệt trên phần cứng thực tế.
Kiến thức bạn học được
- Cơ chế chú ý đơn giản hoạt động nhưng khởi chạy nhiều nhân GPU và hiện thực hoá ma trận chú ý đầy đủ trong bộ nhớ
jax.nn.dot_product_attentionkết hợp phép tính thành một phép toán duy nhất và vớiimplementation=None, JAX sẽ tự động chọn phần phụ trợ tốt nhấtimplementation="cudnn"buộc các hạt nhân chú ý hợp nhất cuDNN của NVIDIA (nhanh nhất trên các chuỗi dài) và yêu cầu đầu vàobfloat16hoặcfloat16và khả năng tính toán 8.0 trở lên- Mặt nạ nhân quả với
is_causal=Trueđược tích hợp vào nhân hợp nhất — không cần ma trận mặt nạ thủ công - GQA và MQA giảm số lượng đầu KV để tiết kiệm bộ nhớ trong quá trình suy luận, còn
dot_product_attentiontự động xử lý việc truyền tin - TransformerEngine cung cấp cơ chế chú ý kết hợp với độ chính xác FP8 tuỳ chọn trên GPU Hopper
Các bước tiếp theo
- Lớp học lập trình 6: Mở rộng quy mô huấn luyện JAX trên nhiều GPU sẽ hướng dẫn bạn cách phân đoạn mảng trên cả L4 và chạy các bước huấn luyện song song
- Thay đổi
HEAD_DIM(thử 32, 64, 128) và xem giá trị nào được cuDNN chấp nhận và thời gian di chuyển như thế nào - Tăng
SEQ_LEN(thử 256, 512, 1024, 2048) và xem mức sử dụng bộ nhớ cũng như tỷ lệ tăng tốc cuDNN - Chạy lại phép so sánh MHA, GQA và MQA với
K=2vàK=1đầu KV so với 8 đầu truy vấn, đồng thời suy luận về mức tiết kiệm bộ nhớ đệm KV cho suy luận