افزایش سرعت توجه در GPU با cuDNN و TransformerEngine

۱. مقدمه

مسیر یادگیری Jax در GPU. آزمایشگاه ۵: توجه به GPU.

در این codelab، بر محاسبه توجه در GPUهای NVIDIA تمرکز می‌کنید. با پیاده‌سازی ساده شروع می‌کنید، آن را با توجه ادغام‌شده داخلی JAX جایگزین می‌کنید، سپس پسای NVIDIA cuDNN را برای همان عملیات فعال می‌کنید. در پایان، خواهید دانست که در هر سطح چه چیزی تغییر می‌کند و با بزرگ‌تر شدن مشکل، مسیرهای بهینه‌سازی‌شده GPU چقدر سریع‌تر می‌شوند.

آنچه انجام خواهید داد

  • پیاده‌سازی توجه ضرب نقطه‌ای مقیاس‌بندی‌شده از ابتدا با عملیات JAX پایه
  • آن را با jax.nn.dot_product_attention، هسته ادغام‌شده داخلی JAX، جایگزین کنید
  • اجبار باطن cuDNN با implementation="cudnn" و افزودن پوشش علّی
  • طول توالی و اندازه دسته را جاروب کنید تا ببینید هسته‌های ادغام‌شده کجا نتیجه می‌دهند
  • مقایسه شکل‌های توجه چندسره MHA،‏ GQA، و MQA
  • معیار توجه NVIDIA TransformerEngine و بررسی مسیر کد FP8 آن

آنچه نیاز خواهید داشت

  • یک پروژه Google Cloud با صورت‌حساب فعال، و اعتبار کارگاه یا رزروی که استفاده از GPU را پوشش دهد
  • سهمیه برای حداقل ۲ واحد پردازش گرافیکی NVIDIA L4 در منطقه انتخابی شما (نحوه بررسی سهمیه واحد پردازش گرافیکی)
  • تکمیل codelabs ۱ تا ۴، یا محیط JAX GPU معادل با CUDA فعال
  • آشنایی با ضرب ماتریس و بیشینه هموار. به تجربه CUDA یا cuDNN نیاز نیست.

زمان تخمینی تکمیل شدن: ۶۰ دقیقه.

نحوه عملکرد «توجه»

توجه ضرب نقطه‌ای مقیاس‌بندی‌شده سه ورودی می‌گیرد — پرس‌وجو، کلید، و مقدار — و موارد زیر را محاسبه می‌کند:

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

ضریب مقیاس 1 / sqrt(d_k) از بزرگ شدن بیش‌ازحد حاصل‌ضرب‌های نقطه‌ای با افزایش ابعاد سر جلوگیری می‌کند، که باعث می‌شود softmax به مناطقی منتقل شود که گرادیان‌های آن بسیار کوچک هستند.

چهار مرحله که هرکدام به مرحله بعدی منتهی می‌شود:

  1. امتیاز — Q @ K.T: ضرب داخلی اندازه‌گیری می‌کند که هر جایگاه پُرسمان چقدر باید به هر جایگاه کلید توجه کند.
  2. Scale + Softmax — softmax(scores / sqrt(d)): مقیاس‌بندی از محو شدن گرادیان جلوگیری می‌کند و softmax امتیازها را به وزن‌های توجه تبدیل می‌کند که مجموع آن‌ها ۱ می‌شود.
  3. توجه — weights @ V: مجموع وزنی بردارهای مقدار، خروجی را برای هر موقعیت پُرسمان تولید می‌کند.
  4. برونداد — همان شکل Q: هر موقعیت پُرسمان اکنون اطلاعاتی را از موقعیت‌هایی که به آن‌ها توجه کرده است حمل می‌کند.

۲. قبل‌از شروع

پروژه‌تان را انتخاب کنید

در کنسول Google Cloud، پروژه‌ای را که صورت‌حساب آن فعال است انتخاب یا ایجاد کنید.

باز کردن Cloud Shell

روی فعال کردن Cloud Shell (نماد پایانه در بالا سمت چپ کنسول) کلیک کنید تا جلسه Cloud Shell شروع شود، سپس آن را به پروژه‌تان هدایت کنید:

gcloud config set project <YOUR_PROJECT_ID>

تأمین محیط GPU

دستور زیر را در Cloud Shell اجرا کنید. Codelab 1 هر فرمان را با جزئیات توضیح می‌دهد، ازجمله سهمیه GPU و الزامات منطقه.

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

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

cp terraform.tfvars.example terraform.tfvars

terraform.tfvars را ویرایش کنید و project_id = "" را تنظیم کنید، سپس خوشه را تأمین کنید و JupyterLab را مستقر کنید:

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

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

terraform apply حدود ۱۲ دقیقه طول می‌کشد. وقتی تمام شد، منتظر 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 را باز کنید، کد را جای‌گذاری کنید، و دفتر یادداشت Python 3 جدیدی در /workspace ایجاد کنید. هر بلوک کد در این codelab به سلولی از آن دفتر یادداشت می‌رود.

نصب کردن آنچه این 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 را ببینید.

۳. آرایه‌های آزمایشی Q،‏ K، و V را بسازید

هسته‌های توجه ادغام‌شده به نحوه چیدمان آرایه‌های ورودی در حافظه اهمیت می‌دهند، بنابراین Q،‏ K، و V را در چیدمانی که JAX انتظار دارد بسازید و سپس کد توجه را بنویسید. jax.nn.dot_product_attention انتظار چیدمان زیر را دارد:

کم‌نور

معنی

پیش‌فرض ما

B

اندازه دسته

4

T

طول توالی پُرسمان

128

جنوب

طول توالی کلید/مقدار

‫۱۲۸ (مشابه T برای خودتوجهی)

شمال

تعداد سرهای توجه

8

H

بُعد سرانه

۶۴

برای پیاده‌سازی ساده، با 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 هر سه آرایه را به واحد پردازش گرافیکی که در سلول راه‌اندازی انتخاب کرده‌اید سنجاق می‌کند، بنابراین هیچ‌چیز در محک‌های بعدی انتقال میزبان را اندازه‌گیری نمی‌کند.

۴. پیاده‌سازی توجه از ابتدا

قبل‌از اینکه به‌سراغ هسته ادغام‌شده بروید، فرمول را مستقیماً بنویسید تا بتوانید دقیقاً ببینید هسته چه چیزی را جایگزین می‌کند. این پیاده‌سازی چهار مرحله را به‌ترتیب دنبال می‌کند: محاسبه امتیازها، مقیاس‌بندی، softmax، ضرب در مقادیر. شما ابعاد سر و توالی را جابجا می‌کنید تا matmul روی محور توالی عمل کند.

این روش به‌درستی کار می‌کند، اما سه راه‌اندازی هسته GPU جداگانه برای ضرب ماتریس امتیاز، softmax، و ضرب ماتریس مقدار صادر می‌کند و ماتریس کامل وزن توجه (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) و dtype float32 را ببینید که همان شکل Q است، یعنی همان چیزی که مرحله خروجی فرمول وعده داده بود.

۵. Switch to dot_product_attention

jax.nn.dot_product_attention مراحل امتیاز، مقیاس، softmax، و attend را در یک عملیات واحد ادغام می‌کند. سپس 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 ببینید.

‫۶. اجباری کردن پسای ادغام‌شده cuDNN

اجازه دادن به JAX برای انتخاب پس‌زمینه راحت است، اما از بیرون نمی‌توانید تشخیص دهید که کدام هسته واقعاً اجرا شده است. تنظیم implementation="cudnn" هسته‌های توجه ادغام‌شده cuDNN مربوط به NVIDIA را اجباری می‌کند. این‌ها هسته‌های GPU بهینه‌شده دستی هستند که کل محاسبات توجه را در یک راه‌اندازی هسته واحد با الگوهای دسترسی به حافظه بهینه‌شده ادغام می‌کنند.

توجه داشته باشید که توجه ادغام‌شده cuDNN دارای الزامات سخت‌افزاری و شکلی است، مانند قابلیت محاسباتی 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 ادامه می‌یابد.

‫۷. افزودن ماسک‌گذاری علّی

در مدل‌های خودبازگشتی (رمزگشاهای سبک 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",
)

باید تفاوت زیادی در موقعیت اول و تفاوتی نزدیک به صفر در موقعیت آخر ببینید. این عدم تقارن به این معنی است که ماسک با موقعیت ۰ همه کلیدها را به‌جز کلید خودش ازدست می‌دهد، درحالی‌که موقعیت نهایی می‌تواند کل توالی را ببیند.

‫۸. درحال اجرای چند محک‌زنی

اکنون سه روش برای محاسبه تابع یکسان دارید. اینکه کدام‌یک را انتخاب کنید کاملاً به شکل مشکل بستگی دارد و می‌توانید آن را ارزیابی کنید.

زمان‌بندی گونه‌ها در شکل پیش‌فرض

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 باید کمی سریع‌ترین باشد، اما مزیت توجه ادغام‌شده معمولاً در طول‌های توالی طولانی‌تر، جایی که اجتناب از ماتریس توجه کامل اهمیت دارد، واضح‌تر می‌شود.

طول توالی را جارو کنید

مزیت هسته‌های توجه ادغام‌شده با طول توالی افزایش می‌یابد. پیاده‌سازی ساده یک ماتریس توجه (B, N, T, S) را در حافظه GPU ایجاد می‌کند که حافظه O(T²) است. هسته‌های ادغام‌شده مانند cuDNN FlashAttention محاسبات را کاشی‌بندی می‌کنند تا هرگز ماتریس کامل را محقق نکنند و حافظه را در O(T) نگه دارند.

جاروب زیر زمان هر پیاده‌سازی را در طول‌های توالی از ۶۴ تا ۱۰۲۴ اندازه‌گیری می‌کند.

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

این جاروب کردن مدتی طول می‌کشد، زیرا هر طول توالی سه یا چهار پیاده‌سازی جداگانه را قبل‌از زمان‌بندی آن‌ها کامپایل می‌کند. باید برای هر طول توالی از ۶۴ تا ۱۰۲۴ یک ردیف جدول داشته باشید.

این نمودار مقایسه می‌کند که چگونه تأخیر توجه با طول توالی برای پیاده‌سازی ساده، 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 به گلوگاه تبدیل شود. در جاروب زیر، طول توالی روی ۲۵۶ ثابت نگه داشته شده است و اندازه دسته تغییر می‌کند.

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

باید یک جدول و یک نمودار با یک ورودی برای هر اندازه دسته از ۱ تا ۱۶ دریافت کنید. پیاده‌سازی ساده از این جاروب حذف شده است، بنابراین هر دو خروجی فقط سه مسیر ادغام‌شده را مقایسه می‌کنند.

‫۹. مقایسه MHA،‏ GQA، و MQA

تاکنون هر سربرگ کلیدها و مقادیر خودش را داشته است. اما رمزگشاها اغلب این تقارن را می‌شکنند، زیرا درطول استنتاج، حافظه پنهان کلید-مقدار چیزی است که بر حافظه تسلط دارد.

توجه چندسره (MHA) به هر سر، طرح‌افکنی Q،‏ K، و V خودش را می‌دهد. توجه پرسمان گروه‌بندی‌شده (GQA) و توجه چندپرسمانی (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 از ۸ سر به ۲ و سپس به ۱ کاهش یابند. نکته همین است: می‌توانید حافظه نهان KV را بدون تغییر دادن هیچ‌چیز در پایین‌دست توجه کاهش دهید.

‫۱۰. معیار NVIDIA TransformerEngine و FP8

TransformerEngine متعلق به NVIDIA واحدهای توجه ادغام‌شده‌ای را ارائه می‌دهد که برای واحدهای پردازش گرافیکی NVIDIA بهینه‌سازی شده‌اند. ادغام JAX از ماژول‌های سبک Flax Linen (نه NNX) استفاده می‌کند، بنابراین ماژول یک‌بار مقداردهی اولیه می‌شود و سپس با متغیرهایش اعمال می‌شود.

طول توالی جاروب با TransformerEngine

این سلول بررسی می‌کند که آیا NVIDIA TransformerEngine دردسترس است یا نه، سپس bf16 علّی 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()

باید یک ردیف جدول و یک نقطه نمودار برای هر طول توالی از ۱۲۸ تا ۲۰۴۸ دریافت کنید. این محک bf16،‏ TransformerEngine را با JAX SDPA در همان حجم کاری توجه علّی مقایسه می‌کند، اما پتانسیل عملکرد کامل TransformerEngine معمولاً در پردازنده‌های گرافیکی Hopper و Blackwell زمانی که FP8 autocast دردسترس است دیده می‌شود.

بازرسی مسیر FP8

در پردازنده گرافیکی Hopper (قابلیت محاسبه >= 9.0، مثل H100)،‏ TransformerEngine می‌تواند توجه را در FP8 برای توان عملیاتی بیشتر اجرا کند. ‫FP8 از دستورالعمل DelayedScaling استفاده می‌کند که سابقه بیشینه مطلق هر تانسور را برای محاسبه ضریب‌های مقیاس پویا ردیابی می‌کند.

  • قالب E4M3 برای پاس رو به جلو (۴ بیت توان، ۳ بیت مانتیس)
  • قالب E5M2 برای پاس برگشتی (۵ بیت توان، ۲ بیت مانتیس)

اگر 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 در محل فراخوانی نیاز دارد.

‫۱۱. مرتب‌سازی

حجم کاری Jupyter، ازجمله LoadBalancer و حجم پایدار را حذف کنید:

kubectl delete -f deploy/jupyter.yaml

خوشه، استخر گره، VPC، و حساب سرویس را ازبین ببرید:

cd terraform
terraform destroy

وقتی درخواست شد، yes را تایپ کنید، سپس تأیید کنید که چیزی جا نمانده است:

gcloud container clusters list
gcloud compute instances list

هر دو باید برای این پروژه خالی باشند. اگر پروژه‌ای را فقط برای این مجموعه ایجاد کرده‌اید، می‌توانید کل پروژه را از کنسول Cloud حذف کنید.

‫۱۲. تبریک می‌گوییم

از محاسبات توجه دست‌نویس به هسته‌های ادغام‌شده بهینه‌سازی‌شده با GPU منتقل شدید و تفاوت را روی سخت‌افزار واقعی اندازه‌گیری کردید.

آنچه آموخته‌اید

  • توجه ساده کار می‌کند اما چندین هسته GPU را راه‌اندازی می‌کند و ماتریس توجه کامل را در حافظه محقق می‌کند
  • jax.nn.dot_product_attention محاسبات را در یک عملیات واحد ادغام می‌کند و با implementation=None JAX بهترین پسای را به‌طور خودکار انتخاب می‌کند
  • implementation="cudnn" هسته‌های توجه ادغام‌شده cuDNN از NVIDIA را اجباری می‌کند که سریع‌ترین هسته‌ها در توالی‌های طولانی هستند و به ورودی‌های bfloat16 یا float16 و قابلیت محاسباتی ۸.۰ یا جدیدتر نیاز دارد
  • ماسک‌گذاری علّی با is_causal=True در هسته ادغام‌شده تعبیه شده است — نیازی به ماتریس ماسک دستی نیست
  • GQA و MQA سرهای KV را کاهش می‌دهند تا حافظه را درطول استنتاج ذخیره کنند، و dot_product_attention پخش را به‌طور خودکار مدیریت می‌کند
  • TransformerEngine توجه ادغام‌شده با دقت FP8 اختیاری را در Hopper GPU ارائه می‌دهد

مراحل بعدی

  • Codelab 6: Scale JAX training across multiple GPUs به شما نشان می‌دهد که چگونه آرایه‌ها را در هر دو L4 قطعه‌بندی کنید و مراحل آموزش را به‌صورت موازی اجرا کنید
  • HEAD_DIM را تغییر دهید (۳۲، ۶۴، ۱۲۸ را امتحان کنید) و ببینید cuDNN کدام مقادیر را می‌پذیرد و زمان‌بندی‌ها چگونه تغییر می‌کنند
  • SEQ_LEN را افزایش دهید (۲۵۶، ۵۱۲، ۱۰۲۴، ۲۰۴۸ را امتحان کنید) و استفاده از حافظه و نسبت افزایش سرعت cuDNN را تماشا کنید
  • مقایسه MHA،‏ GQA، و MQA را با K=2 و K=1 سر KV دربرابر ۸ سر پُرسمان دوباره اجرا کنید و درباره صرفه‌جویی در حافظه پنهان KV برای استنتاج استدلال کنید

اسناد مرجع