1. Giriş
Bu codelab'de, NVIDIA GPU'larda dikkat mekanizmasının hesaplanmasına odaklanacaksınız. Önce basit bir uygulama ile başlar, ardından bunu JAX'ın yerleşik birleştirilmiş dikkat mekanizmasıyla değiştirir ve aynı işlem için NVIDIA'nın cuDNN arka ucunu etkinleştirirsiniz. Kursun sonunda, her seviyede nelerin değiştiğini ve sorun büyüdükçe GPU için optimize edilmiş yolların ne kadar hızlandığını öğreneceksiniz.
Yapacaklarınız
- Temel JAX işlemleriyle ölçeklendirilmiş nokta çarpım dikkatini sıfırdan uygulama
- Bunun yerine JAX'in yerleşik birleştirilmiş çekirdeği olan
jax.nn.dot_product_attentionkullanın. implementation="cudnn"ile cuDNN arka ucunu zorunlu kılın ve nedensel maskeleme ekleyin.- Birleştirilmiş çekirdeklerin nerede işe yaradığını görmek için dizi uzunluğunu ve toplu iş boyutunu tarayın.
- MHA, GQA ve MQA çok kafalı dikkat şekillerini karşılaştırma
- NVIDIA TransformerEngine dikkat mekanizmasını karşılaştırma ve FP8 kod yolunu inceleme
İhtiyacınız olanlar
- Faturalandırmanın etkin olduğu bir Google Cloud projesi ve GPU kullanımını kapsayan atölye kredileri veya rezervasyon
- Seçtiğiniz bölgede en az 2 NVIDIA L4 GPU için kota (GPU kotasını kontrol etme)
- 1 ila 4 numaralı codelab'leri veya eşdeğer bir CUDA özellikli JAX GPU ortamını tamamlamış olmak
- Matris çarpımı ve softmax hakkında bilgi sahibi olmak CUDA veya cuDNN deneyimi gerekmez.
Tahmini tamamlanma süresi: 60 dakika.
Dikkat nasıl çalışır?
Ölçeklendirilmiş nokta çarpım dikkat mekanizması üç giriş alır (Sorgu, Anahtar ve Değer) ve şunları hesaplar:
Attention(Q, K, V) = softmax(Q K^T / sqrt(d_k)) V
Ölçeklendirme faktörü 1 / sqrt(d_k), başlık boyutu arttıkça nokta çarpımlarının çok fazla büyümesini engeller. Bu sayede softmax, gradyanlarının çok küçük olduğu bölgelere itilmez.
Her biri bir sonrakini besleyen dört adım:
- Puan —
Q @ K.T: Nokta çarpım, her bir sorgu konumunun her bir anahtar konuma ne kadar dikkat etmesi gerektiğini ölçer. - Ölçekleme + Softmax —
softmax(scores / sqrt(d)): Ölçekleme, gradyanın kaybolmasını önler ve softmax, puanları toplamı 1 olan dikkat ağırlıklarına dönüştürür. - Attend (Katıl) —
weights @ V: Değer vektörlerinin ağırlıklı toplamı, her sorgu konumu için çıkış üretir. - Çıktı: Q ile aynı şekildedir. Her sorgu konumu artık dikkat ettiği konumlardan gelen bilgileri taşır.
2. Başlamadan önce
Projenizi seçin.
Google Cloud Console'da faturalandırmanın etkin olduğu bir proje seçin veya oluşturun.
Cloud Shell'i açın
Cloud Shell oturumu başlatmak için Cloud Shell'i etkinleştir'i (konsolun sağ üst kısmındaki terminal simgesi) tıklayın, ardından oturumu projenize yönlendirin:
gcloud config set project <YOUR_PROJECT_ID>
GPU ortamını temel hazırlıkla oluşturma
Cloud Shell'de aşağıdaki komutu çalıştırın. Codelab 1'de, GPU kotası ve bölge şartları dahil olmak üzere her komut ayrıntılı olarak açıklanmaktadır.
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 dosyasını düzenleyin ve project_id = " değerini ayarlayın, ardından kümeyi temel hazırlık işlemine tabi tutun ve JupyterLab'i dağıtın:
terraform init
terraform apply
$(terraform output -raw get_credentials_command)
cd ..
kubectl apply -f deploy/jupyter.yaml
terraform apply yaklaşık 12 dakika sürer. İşlem tamamlandığında kapsül ve LoadBalancer'ın hazır olmasını bekleyin, ardından kapsül günlüğünden tek seferlik JupyterLab jetonunu okuyun:
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:// adresini açın, jetonu yapıştırın ve /workspace içinde yeni bir Python 3 not defteri oluşturun. Bu codelab'deki her kod bloğu, ilgili not defterinin bir hücresine eklenir.
Bu codelab için gerekenleri yükleme
!pip install --quiet matplotlib flax
GPU'yu ayarlama ve doğrulama
JAX'i içe aktarın ve varsayılan arka ucun GPU olduğunu onaylayın. Bu hücrede ayrıca, sonraki her adımda cihazın çalışmasını beklemek ve sonuçları oluşturmak için kullanılan yardımcılar olan block_tree, show_table ve show_bars tanımlanır.
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 sürümünü, varsayılan arka uç olarak gpu, CUDA cihazlarının listesini ve codelab'in geri kalanının kullandığı GPU'yu görürsünüz.
3. Q, K ve V test dizileri oluşturma
Birleştirilmiş dikkat çekirdekleri, giriş dizilerinin bellekte nasıl düzenlendiğiyle ilgilenir. Bu nedenle, dikkat kodu yazmadan önce Q, K ve V'yi JAX'ın beklediği düzende oluşturun. jax.nn.dot_product_attention şu düzeni bekler:
Loş | Anlamı | Varsayılanımız |
B | Grup boyutu | 4 |
T | Sorgu dizisi uzunluğu | 128 |
G | Anahtar/değer dizisi uzunluğu | 128 (öz dikkat için T ile aynı) |
H | Dikkat başlığı sayısı | 8 |
H | Baş başına boyut | 64 |
Basit uygulama için float32 ile başlayıp rastgele değerlerle başlatırsınız.
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",
)
Her dizi için bir satır içeren bir tablo görürsünüz. Her satırda şekil (4, 128, 8, 64) ve dtype float32 bilgileri yer alır. jax.device_put çağrısı, üç diziyi de kurulum hücresinde seçtiğiniz GPU'ya sabitler. Bu sayede, karşılaştırmalardaki hiçbir şey daha sonra ana makine aktarımını ölçmez.
4. Dikkat mekanizmasını sıfırdan uygulama
Birleştirilmiş bir çekirdeğe ulaşmadan önce, çekirdeğin tam olarak neyi değiştirdiğini görebilmek için formülü doğrudan yazın. Bu uygulama sırasıyla şu dört adımı izler: puanları hesaplama, ölçeklendirme, softmax, değerlerle çarpma. Matmul'un sıra ekseni üzerinde çalışması için baş ve sıra boyutlarının yerini değiştirirsiniz.
Doğru çalışır ancak puan matris çarpımı, softmax ve değer matris çarpımı için üç ayrı GPU çekirdeği başlatır ve GPU belleğinde tam (B, N, T, S) dikkat ağırlık matrisini oluşturur.
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",
)
Formülün Çıkış adımında belirtildiği gibi, Q ile aynı şekle sahip olan (4, 128, 8, 64) çıkış şeklini ve float32 veri türünü görürsünüz.
5. dot_product_attention'a geçme
jax.nn.dot_product_attention; puan, ölçek, softmax ve dikkat adımlarını tek bir işlemde birleştirir. JAX ve XLA, bellek erişim modelini optimize edebilir. Özellikle, dizinin uzun olduğu durumlarda tam dikkat ağırlığı matrisini oluşturmaktan kaçınabilirler.
Varsayılan implementation=None ile JAX, mevcut en iyi arka ucu otomatik olarak seçer. cuDNN'nin kullanılabildiği ve uyumlu girişlerin olduğu bir GPU'da cuDNN zaten kullanılıyor olabilir. Diğer donanımlarda XLA'ya geri döner.
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",
)
Basit sürümle aynı şekli ve veri türünü, küçük bir maksimum mutlak farkı ve atol=1e-3 yakınlık kontrolü için True değerini görmeniz gerekir.
6. cuDNN birleştirilmiş arka ucunu zorunlu kılma
JAX'in arka ucu seçmesine izin vermek kullanışlıdır ancak dışarıdan bakıldığında hangi çekirdeğin çalıştığını anlayamazsınız. implementation="cudnn" ayarı, NVIDIA'nın cuDNN birleştirilmiş dikkat çekirdeklerini zorunlu kılar. Bunlar, tüm dikkat hesaplamasını optimize edilmiş bellek erişim kalıplarıyla tek bir çekirdek başlatmada birleştiren, manuel olarak optimize edilmiş GPU çekirdekleridir.
cuDNN birleştirilmiş dikkat mekanizmasının, GPU işlem kapasitesinin >= 8.0 (Ampere veya daha yeni) olması ya da giriş veri türleri olarak float16 veya bfloat16 kullanılması gibi donanım ve şekil gereksinimleri olduğunu unutmayın. Şartlar karşılanmazsa ve implementation="cudnn" değerini ayarlarsanız JAX, sessizce geri dönmek yerine hata verir. Bu nedenle, aşağıdaki kod bfloat16 değerine dönüştürülür ve çağrıyı try/except içine sarmalar: HAS_CUDNN_SDPA değerini ayarlar, böylece sonraki her adım cuDNN yolunun bu makinede kullanılabilir olup olmadığını bilir.
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'nuz cuDNN SDPA'yı çalıştıramıyorsa kod, nedeni yazdırır ve codelab, XLA arka ucunda devam eder.
7. Nedensel maskeleme ekleme
Otoregresif modellerde (GPT tarzı kod çözücüler) her konum yalnızca önceki konumlara dikkat edebilir. is_causal=True ayarı, bu alt üçgen maskeyi birleştirilmiş çekirdeğin içine uygular. Böylece kendiniz bir maske matrisi oluşturmanız gerekmez.
Aşağıdaki kod, dikkat mekanizmasını maskeli ve maskesiz olarak iki kez çalıştırır ve maskenin neleri değiştirdiğini göstermek için iki konumu karşılaştırır.
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",
)
İlk konumda büyük bir fark, son konumda ise sıfıra yakın bir fark görmeniz gerekir. Bu asimetri, 0 konumundaki maskenin kendi anahtarı dışındaki tüm anahtarları kaybetmesidir. Son konum ise tüm diziyi zaten görebilir.
8. Bazı karşılaştırmalar yapılıyor
Artık aynı işlevi hesaplamanın üç yolu var. Hangisini kullanacağınız tamamen sorunun şekline bağlıdır ve bunu değerlendirebilirsiniz.
Varyantları varsayılan şekle göre zamanlama
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)
Bu küçük problem boyutu için tüm uygulamalar çok hızlıdır ve çoğunlukla ek yük sınırlıdır. Bu nedenle, basit sürüm rekabetçidir. cuDNN bf16 biraz daha hızlı olmalıdır ancak birleştirilmiş dikkatin avantajı, tam dikkat matrisinden kaçınmanın önemli olduğu daha uzun dizi uzunluklarında genellikle daha belirgin hale gelir.
Dizi uzunluğunu tarama
Birleştirilmiş dikkat çekirdeklerinin avantajı, sıra uzunluğuyla birlikte artar. Basit uygulama, GPU belleğinde O(T²) bellek olan bir (B, N, T, S) dikkat matrisi oluşturur. cuDNN FlashAttention gibi birleştirilmiş çekirdekler, hesaplamayı parçalara ayırarak tam matrisi hiçbir zaman oluşturmaz ve belleği O(T) düzeyinde tutar.
Aşağıdaki tarama, her bir uygulamanın süresini 64 ile 1024 arasındaki dizi uzunluklarında ölçer.
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),
)
Her dizi uzunluğu, zamanlamadan önce üç veya dört ayrı uygulama derlediğinden tarama işlemi biraz zaman alır. 64 ile 1024 arasındaki her dizi uzunluğu için bir tablo satırı elde etmeniz gerekir.
Bu grafikte, basit uygulama, float32 ve bfloat16'da JAX/XLA SDPA ve bfloat16'da cuDNN birleştirilmiş dikkat için dikkat gecikmesinin sıra uzunluğuna göre nasıl değiştiği karşılaştırılmaktadır.
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()
Her uygulama için bir satır görürsünüz. Sıra uzadıkça aradaki farkın da arttığını görebilirsiniz.
Yığın boyutunu tarama
Daha büyük gruplar, çekirdek başlatma ek yükünü amorti eder ve GPU belleğinin darboğaz haline geldiği noktaya kadar GPU kullanımını iyileştirir. Aşağıdaki tarama, sıra uzunluğunu 256'da sabit tutar ve toplu iş boyutunu değiştirir.
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'den 16'ya kadar her grup boyutu için bir giriş içeren bir tablo ve grafik elde edersiniz. Basit uygulama bu taramanın dışında bırakıldığından her iki çıkış da yalnızca üç birleştirilmiş yolu karşılaştırır.
9. MHA, GQA ve MQA'yı karşılaştırma
Şimdiye kadar her başlık kendi anahtarlarını ve değerlerini taşıyordu. Ancak kod çözücüler, çıkarım sırasında anahtar/değer önbelleği belleğe hakim olduğu için genellikle bu simetriyi bozar.
Çok başlı dikkat (MHA), her başlığa kendi Q, K ve V projeksiyonlarını verir. Gruplandırılmış sorgu dikkati (GQA) ve çoklu sorgu dikkati (MQA), çıkarım sırasında bellek ve işlem gücünden tasarruf etmek için KV başlıklarının sayısını azaltır.
jax.nn.dot_product_attention, KV başlıkları K < N olduğunda otomatik olarak yayınlanırken üçünü de işler.
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 ve V şekilleri 8 başlıktan 2'ye ve 1'e küçülse de üç satırın da aynı çıkış şeklini ((2, 64, 8, 64)) bildirmesi gerekir. Önemli olan da bu: Dikkat mekanizmasının altındaki hiçbir şeyi değiştirmeden KV önbelleğini kesebilirsiniz.
10. NVIDIA TransformerEngine ve FP8'i karşılaştırma
NVIDIA'nın TransformerEngine'i, NVIDIA GPU'lar için optimize edilmiş birleştirilmiş dikkat modülleri sağlar. JAX entegrasyonu, Flax Linen tarzı modüller (NNX değil) kullandığından modül bir kez başlatılır ve ardından değişkenleriyle uygulanır.
TransformerEngine ile dizi uzunluğunu tarama
Bu hücre, NVIDIA TransformerEngine'in kullanılabilir olup olmadığını kontrol eder, ardından aynı iş yükünde XLA ve cuDNN arka uçlarıyla JAX SDPA'ya karşı çeşitli dizi uzunluklarında bf16 nedensel DotProductAttention'ini karşılaştırır.
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 ile 2048 arasındaki her dizi uzunluğu için bir tablo satırı ve bir çizim noktası elde etmeniz gerekir. Bu bf16 karşılaştırması, TransformerEngine'i aynı nedensel dikkat iş yükünde JAX SDPA ile karşılaştırır. Ancak TransformerEngine'in tam performans potansiyeli genellikle FP8 otomatik yayınlama özelliği mevcut olduğunda Hopper ve Blackwell GPU'larda görülür.
FP8 yolunu inceleme
Hopper GPU'larda (işlem kapasitesi >= 9.0, örneğin H100), TransformerEngine, ek işleme kapasitesi için FP8'de dikkat mekanizmasını çalıştırabilir. FP8, dinamik ölçeklendirme faktörlerini hesaplamak için tensör başına mutlak maksimum geçmişi izleyen DelayedScaling tarifini kullanır.
- İleriye doğru geçiş için E4M3 biçimi (4 üs, 3 mantis biti)
- Geriye doğru geçiş için E5M2 biçimi (5 üs, 2 mantis biti)
GPU, FP8'i desteklemiyorsa bu hücre, kodun çalıştırılmadan nasıl görüneceğini gösterir.
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'te, GPU'nuzu adlandıran ve algılanan FP8 desteği olmadığını bildiren bir tablo, ardından Hopper GPU'da kullanacağınız te.autocast çağrısını gösteren iki yazılı satır görürsünüz. Bu snippet'i saklayın: FP8'in çağrı sitesinde gerektirdiği tek değişiklik budur.
11. Temizleme
LoadBalancer ve kalıcı birim dahil olmak üzere Jupyter iş yükünü silin:
kubectl delete -f deploy/jupyter.yaml
Kümeyi, düğüm havuzunu, VPC'yi ve hizmet hesabını silin:
cd terraform
terraform destroy
İstendiğinde yes yazın, ardından hiçbir şeyin kalmadığını onaylayın:
gcloud container clusters list
gcloud compute instances list
Bu proje için her ikisi de boş olmalıdır. Yalnızca bu seri için bir proje oluşturduysanız projenin tamamını Cloud Console'dan silebilirsiniz.
12. Tebrikler
Elle yazılmış bir dikkat hesaplamasından GPU için optimize edilmiş birleştirilmiş çekirdeklere geçtiniz ve farkı gerçek donanım üzerinde ölçtünüz.
Öğrendikleriniz
- Naive attention works but launches multiple GPU kernels and materializes the full attention matrix in memory
jax.nn.dot_product_attention, hesaplamayı tek bir işlemde birleştirir veimplementation=NoneJAX ile en iyi arka ucu otomatik olarak seçer.implementation="cudnn", NVIDIA'nın uzun dizilerde en hızlı olan cuDNN birleştirilmiş dikkat çekirdeklerini zorunlu kılar vebfloat16veyafloat16girişleri ile 8.0 ya da daha yeni bir işlem kapasitesi gerektirir.is_causal=Trueile nedensel maskeleme, birleştirilmiş çekirdeğe yerleşiktir. Manuel maske matrisi gerekmez.- GQA ve MQA, çıkarım sırasında bellekten tasarruf etmek için KV başlıklarını azaltır ve
dot_product_attention, yayını otomatik olarak yönetir. - TransformerEngine, Hopper GPU'larda isteğe bağlı FP8 hassasiyetiyle birleştirilmiş dikkat sağlar.
Sonraki adımlar
- Codelab 6: JAX eğitimini birden fazla GPU'da ölçeklendirme başlıklı makalede, dizileri her iki L4'te de parçalama ve eğitim adımlarını paralel olarak çalıştırma hakkında bilgi verilmektedir.
HEAD_DIMdeğerini değiştirin (32, 64, 128'i deneyin) ve cuDNN'nin hangi değerleri kabul ettiğini ve zamanlamaların nasıl değiştiğini izleyin.SEQ_LENdeğerini artırın (256, 512, 1024, 2048 değerlerini deneyin) ve bellek kullanımını ve cuDNN hızlandırma oranını izleyin.- MHA, GQA ve MQA karşılaştırmasını 8 sorgu başlığına karşı
K=2veK=1KV başlıklarıyla yeniden çalıştırın ve çıkarım için KV önbelleği tasarrufları hakkında akıl yürütün.