האצת תשומת הלב ב-GPU באמצעות cuDNN ו-TransformerEngine

1. מבוא

תוכנית הלימודים "Jax ב-GPU". שיעור Lab מספר 5: מנגנון תשומת הלב ב-GPU.

ב-Codelab הזה נתמקד בחישוב של מנגנון הקשב ביחידות GPU של NVIDIA. מתחילים בהטמעה פשוטה, מחליפים אותה במנגנון הקשב המשולב המובנה של JAX, ואז מפעילים את הקצה העורפי של NVIDIA cuDNN לאותה פעולה. בסוף תדעו מה משתנה בכל רמה, וכמה מהר יותר הנתיבים שעברו אופטימיזציה ל-GPU הופכים ככל שהבעיה גדלה.

הפעולות שתבצעו:

  • הטמעה של מנגנון קשב מסוג scaled dot-product מאפס באמצעות פעולות JAX בסיסיות
  • מחליפים אותו ב-jax.nn.dot_product_attention, הליבה הממוזגת המובנית של JAX
  • אילוץ של קצה העורפי של cuDNN באמצעות implementation="cudnn" והוספת מיסוך סיבתי
  • סריקה של אורך הרצף וגודל האצווה כדי לראות איפה הליבות הממוזגות משתלמות
  • השוואה בין צורות הקשב מרובות הראשים MHA,‏ GQA ו-MQA
  • השוואת ביצועים של מנגנון תשומת הלב של NVIDIA TransformerEngine ובדיקת נתיב הקוד שלו ב-FP8

הדרישות

  • פרויקט בענן של Google עם חיוב מופעל, וקרדיטים לסדנה או מקום שמור שמכסה את השימוש ב-GPU
  • מכסה של לפחות 2 מעבדי NVIDIA L4 GPU באזור שבחרתם (איך בודקים את מכסת ה-GPU)
  • השלמה של שיעורי ה-Codelab‏ 1 עד 4, או סביבת GPU של JAX עם CUDA
  • היכרות עם כפל מטריצות ו-softmax. לא נדרש ניסיון ב-CUDA או ב-cuDNN.

זמן משוער לסיום: 60 דקות.

איך פועל מדד תשומת הלב

מנגנון הקשב Scaled dot-product attention מקבל שלושה קלטים – שאילתה, מפתח וערך – ומחשב:

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

גורם קביעת קנה המידה 1 / sqrt(d_k) מונע מהמכפלות הסקלריות לגדול יותר מדי ככל שמימד הראש גדל, מה שעלול לדחוף את ה-softmax לאזורים שבהם הגרדיאנטים שלו זעירים.

ארבעה שלבים, שכל אחד מהם מוביל לשלב הבא:

  1. ציון — Q @ K.T: מכפלה סקלרית מודדת את מידת תשומת הלב שכל מיקום של שאילתה צריך להקדיש לכל מיקום של מפתח.
  2. שינוי קנה מידה + Softmax — softmax(scores / sqrt(d)): שינוי קנה המידה מונע את היעלמות הגרדיאנט, ו-Softmax ממיר את הציונים למשקלי תשומת לב שהסכום שלהם הוא 1.
  3. Attend (התייחסות) — weights @ V: הסכום המשוקלל של וקטורי הערכים יוצר את הפלט לכל מיקום של שאילתה.
  4. פלט – אותה צורה כמו Q: כל מיקום של שאילתה מכיל עכשיו מידע מהמיקומים שהוא התייחס אליהם.

‫2. לפני שמתחילים

בוחרים את הפרויקט.

במסוף Google Cloud, בוחרים או יוצרים פרויקט שהחיוב מופעל בו.

פתיחת Cloud Shell

לוחצים על Activate Cloud Shell (סמל הטרמינל בפינה הימנית העליונה של המסוף) כדי להתחיל סשן של Cloud Shell, ואז מפנים אותו לפרויקט:

gcloud config set project <YOUR_PROJECT_ID>

הקצאת סביבת ה-GPU

מריצים את הפקודה הבאה ב-Cloud Shell. ב-Codelab הראשון יש הסבר מפורט על כל פקודה, כולל מכסת ה-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, ואז קוראים את הטוקן החד-פעמי של JupyterLab מיומן ה-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

פותחים את http://:8884, מדביקים את הטוקן ויוצרים קובץ Notebook חדש של Python 3 ב-/workspace. כל בלוק קוד ב-Codelab הזה נכנס לתא ב-notebook הזה.

התקנת הרכיבים הנדרשים ל-Codelab הזה

!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 שבו משתמשים בשאר ה-Codelab.

3. יצירת מערכי בדיקה של Q,‏ K ו-V

בקרנלים של תשומת לב משולבת, חשוב איך מערכי הקלט מסודרים בזיכרון, אז צריך לבנות את Q, ‏ K ו-V בפריסה ש-JAX מצפה לה לפני שכותבים קוד תשומת לב. הפריסה הצפויה של 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 יפעל על ציר הרצף.

הוא פועל בצורה תקינה, אבל הוא מפעיל שלוש ליבות GPU נפרדות ל-score matmul, ל-softmax ול-value matmul, והוא יוצר את מטריצת משקל הקשב המלאה של (B, N, T, S) בזיכרון ה-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",
)

הפלט אמור להיות בצורה (4, 128, 8, 64) ובסוג הנתונים float32, שהיא אותה צורה כמו Q, כפי שהובטח בשלב הפלט של הנוסחה.

5. מעבר ל-dot_product_attention

ה-jax.nn.dot_product_attention ממזג את שלבי הניקוד, קנה המידה, ה-softmax וההשתתפות לפעולה אחת. לאחר מכן, JAX ו-XLA יכולים לבצע אופטימיזציה של דפוס הגישה לזיכרון. במיוחד, הם יכולים להימנע ממימוש מטריצת משקל הקשב המלאה כשהרצף ארוך.

כברירת מחדל, implementation=None, JAX בוחר באופן אוטומטי את הקצה העורפי הטוב ביותר שזמין. ב-GPU עם cuDNN זמין וקלט תואם, יכול להיות שהוא כבר משתמש ב-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 כמו בגרסה הנאיבית, הבדל מוחלט מקסימלי קטן ו-True לבדיקת הקרבה של atol=1e-3.

6. אילוץ של הקצה העורפי המשולב של cuDNN

נוח לתת ל-JAX לבחור את הקצה העורפי, אבל מבחוץ אי אפשר לדעת איזה ליבה (kernel) הופעלה בפועל. ההגדרה implementation="cudnn" מאלצת את ליבות תשומת הלב הממוזגות של cuDNN של NVIDIA. אלה ליבות GPU שעברו אופטימיזציה ידנית, שמשלבות את כל חישוב הקשב בהפעלה של ליבה אחת עם דפוסי גישה לזיכרון שעברו אופטימיזציה.

שימו לב של-cuDNN fused attention יש דרישות חומרה וצורה, כמו יכולת חישוב של GPU שצריכה להיות >= 8.0 (Ampere או גרסה חדשה יותר) או float16 או bfloat16 כסוגי נתונים של קלט. אם הדרישות לא מתקיימות והגדרתם את 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, הקוד מדפיס את הסיבה וה-Codelab ממשיך ב-XLA backend.

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 אמור להיות קצת יותר מהיר, אבל היתרון של fused attention בדרך כלל מתבהר ברצפים ארוכים יותר, שבהם חשוב להימנע ממטריצת תשומת הלב המלאה.

סריקה של אורך הרצף

היתרון של ליבות תשומת לב משולבות גדל ככל שהרצף ארוך יותר. ההטמעה הנאיבית יוצרת מטריצת קשב של (B, N, T, S) בזיכרון ה-GPU, שהוא זיכרון 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.

בתרשים הזה מוצגת השוואה בין השינויים בזמן האחזור של הקשב בהתאם לאורך הרצף בהטמעה הנאיבית, ב-JAX/XLA SDPA ב-float32 וב-bfloat16, ובקשב הממוזג של cuDNN ב-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()

אמורה להופיע שורה אחת לכל הטמעה. ואתם אמורים לראות שהפער גדל ככל שהרצף מתארך.

סריקה של גודל האצווה

ככל שהאצוות גדולות יותר, כך הן מפחיתות את התקורה של הפעלת הליבה ומשפרות את ניצול ה-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 משלו. ‫Grouped-query attention‏ (GQA) ו-multi-query attention‏ (MQA) מפחיתים את מספר ראשי ה-KV כדי לחסוך בזיכרון ובמשאבי מחשוב במהלך ההסקה.

‫jax.nn.dot_product_attention מטפל בכל שלושת המקרים, וראשי ה-KV משודרים באופן אוטומטי כש-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",
)

כל שלוש השורות אמורות לדווח על אותה צורת פלט, (2, 64, 8, 64), למרות שהצורות K ו-V מתכווצות מ-8 ראשים ל-2 ול-1. זה העניין: אפשר לחתוך את מטמון ה-KV בלי לשנות שום דבר אחרי שלב ה-attention.

10. בדיקת ביצועים של NVIDIA TransformerEngine ו-FP8

‫TransformerEngine של NVIDIA מספק מודולים של fused attention שעברו אופטימיזציה ל-GPU של NVIDIA. השילוב של JAX משתמש במודולים בסגנון Flax Linen (לא NNX), כך שהמודול מאותחל פעם אחת ואז מיושם עם המשתנים שלו.

אורך רצף הסריקה עם TransformerEngine

התא הזה בודק אם NVIDIA TransformerEngine זמין, ואז משווה את ה-bf16 causal DotProductAttention שלו בכמה אורכי רצף ל-JAX SDPA עם הקצוות העורפיים של XLA ו-cuDNN באותו עומס עבודה.

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 עד 2,048. בבדיקת הביצועים הזו של bf16, משווים בין TransformerEngine לבין JAX SDPA באותו עומס עבודה של causal-attention, אבל בדרך כלל רואים את פוטנציאל הביצועים המלא של TransformerEngine במעבדי GPU של Hopper ו-Blackwell כש-FP8 autocast זמין.

בדיקת הנתיב של FP8

במעבדי GPU מסוג Hopper (יכולת חישוב >= 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 שבה הייתם משתמשים ב-GPU של Hopper. שמרו את קטע הקוד הזה: זה השינוי היחיד שנדרש ב-FP8 באתר הקריאה.

11. הסרת המשאבים

מוחקים את עומס העבודה של Jupyter, כולל LoadBalancer והנפח הקבוע:

kubectl delete -f deploy/jupyter.yaml

משמידים את האשכול, מאגר הצמתים, ה-VPC וחשבון השירות:

cd terraform
terraform destroy

מקלידים yes כשמופיעה בקשה, ואז מוודאים שלא נשאר שום דבר:

gcloud container clusters list
gcloud compute instances list

שני השדות צריכים להיות ריקים בפרויקט הזה. אם יצרתם פרויקט רק בשביל הסדרה הזו, אתם יכולים למחוק את כל הפרויקט ממסוף Cloud.

12. מזל טוב

עברתם מחישוב קשב ידני לקרנלים משולבים שעברו אופטימיזציה ל-GPU, ומדדתם את ההבדל בחומרה אמיתית.

מה למדתם

  • מנגנון תשומת לב נאיבי עובד, אבל מפעיל כמה ליבות GPU ומממש את מטריצת תשומת הלב המלאה בזיכרון
  • ‫jax.nn.dot_product_attention ממזג את החישוב לפעולה אחת, ועם implementation=None, ‏JAX בוחר את הקצה העורפי הטוב ביותר באופן אוטומטי
  • ‫implementation="cudnn" מאלץ את ליבות תשומת הלב הממוזגות של cuDNN של NVIDIA, שהן המהירות ביותר ברצפים ארוכים, ודורש קלט של bfloat16 או float16 ויכולת חישוב של 8.0 ומעלה
  • הסתרת סיבתיות עם is_causal=True מובנית בקרנל הממוזג – אין צורך במטריצת הסתרה ידנית
  • ‫GQA ו-MQA מפחיתים את ראשי ה-KV כדי לחסוך בזיכרון במהלך ההסקה, ו-dot_product_attention מטפל בשידור באופן אוטומטי
  • ‫TransformerEngine מספק מנגנון תשומת לב משולב עם דיוק אופציונלי של FP8 במעבדי GPU של Hopper

השלבים הבאים

  • ב-Codelab 6: Scale JAX training across multiple GPUs מוסבר איך לפצל מערכים בין שני מעבדי L4 ולהריץ שלבי אימון במקביל.
  • משנים את HEAD_DIM (נסו 32, 64, 128) ובודקים אילו ערכים cuDNN מקבל ואיך התזמונים משתנים
  • מגדילים את SEQ_LEN (נסו 256, 512, 1024, 2048) ועוקבים אחרי השימוש בזיכרון ויחס ההאצה של cuDNN
  • מריצים מחדש את ההשוואה בין MHA,‏ GQA ו-MQA עם ראשי KV של K=2 ו-K=1 מול 8 ראשי שאילתות, ומסבירים על החיסכון במטמון KV להסקת מסקנות

מסמכי עזר