تسريع آلية الانتباه على وحدة معالجة الرسومات باستخدام cuDNN وTransformerEngine

1. مقدمة

المسار التعليمي حول Jax على وحدة معالجة الرسومات التمرين العملي 5: آلية الانتباه على وحدة معالجة الرسومات

في هذا الدرس التطبيقي حول الترميز، سنركّز على احتساب الانتباه على وحدات معالجة الرسومات من NVIDIA. ستبدأ بتنفيذ بسيط، ثم تستبدله بآلية الانتباه المدمجة في JAX، ثم تفعّل واجهة cuDNN الخلفية من NVIDIA للعملية نفسها. في النهاية، ستتعرّف على التغييرات التي تحدث على كل مستوى، ومدى سرعة المسارات المحسّنة لوحدة معالجة الرسومات مع تزايد المشكلة.

الإجراءات التي ستنفذّها

  • تنفيذ آلية الانتباه بالضرب القياسي المقيس من البداية باستخدام عمليات JAX الأساسية
  • استبدِلها بـ jax.nn.dot_product_attention، وهي نواة مدمجة في JAX
  • فرض استخدام واجهة cuDNN الخلفية مع implementation="cudnn" وإضافة قناع سببي
  • يمكنك تجربة طول التسلسل وحجم الدفعة لمعرفة الحالات التي تكون فيها النواة المدمجة مفيدة.
  • مقارنة بين أشكال الانتباه المتعدد الرؤوس (MHA) والانتباه المتعدد المجموعات (GQA) والانتباه المتعدد الاستعلامات (MQA)
  • مقارنة أداء NVIDIA TransformerEngine attention وفحص مسار الرمز البرمجي FP8

المتطلبات

  • مشروع على السحابة الإلكترونية من Google Cloud مع تفعيل الفوترة، وأرصدة ورشة عمل أو حجز يغطي استخدام وحدة معالجة الرسومات
  • حصة لوحدتَي معالجة رسومات NVIDIA L4 على الأقل في المنطقة التي تختارها (كيفية التحقّق من حصة وحدة معالجة الرسومات)
  • إكمال الدروس التطبيقية حول الترميز من 1 إلى 4، أو بيئة JAX GPU مكافئة متوافقة مع CUDA
  • الإلمام بعملية ضرب المصفوفات ودالة softmax لا يلزم توفّر خبرة في CUDA أو cuDNN.

الوقت المقدَّر لإكمال العملية: 60 دقيقة

آلية عمل الانتباه

تتلقّى آلية الانتباه Scaled dot-product ثلاثة مدخلات، وهي الطلب والمفتاح والقيمة، وتحسب ما يلي:

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. التركيز — weights @ V: ينتج عن المجموع المرجّح لمتجهات القيمة الناتج لكل موضع طلب بحث.
  4. الناتج: يكون بالشكل نفسه لـ Q، حيث يحمل كل موضع طلب الآن معلومات من المواضع التي تم التركيز عليها.

2. قبل البدء

اختيار مشروعك

في وحدة تحكّم Google Cloud، اختَر مشروعًا تم تفعيل الفوترة فيه أو أنشِئ مشروعًا جديدًا.

فتح Cloud Shell

انقر على تفعيل Cloud Shell (رمز الوحدة الطرفية في أعلى يسار وحدة التحكّم) لبدء جلسة Cloud Shell، ثم وجِّهها إلى مشروعك:

gcloud config set project <YOUR_PROJECT_ID>

توفير بيئة وحدة معالجة الرسومات

نفِّذ الأمر التالي في Cloud Shell. يشرح الدرس التطبيقي الأول حول الترميز كل أمر بالتفصيل، بما في ذلك حصة وحدة معالجة الرسومات ومتطلبات المنطقة.

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، وألصِق الرمز المميّز، وأنشئ دفتر ملاحظات جديدًا في Python 3 في /workspace. يتم إدخال كل مجموعة رموز برمجية في هذا الدرس التطبيقي حول الترميز في خلية من ورقة الملاحظات هذه.

تثبيت ما يتطلّبه هذا الدرس التطبيقي حول الترميز

!pip install --quiet matplotlib flax

إعداد وحدة معالجة الرسومات والتحقّق منها

استورِد JAX وتأكَّد من أنّ الخلفية التلقائية هي وحدة معالجة رسومات. تحدّد هذه الخلية أيضًا 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، ووحدة معالجة الرسومات التي يستخدمها باقي الدرس التطبيقي حول الترميز.

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) والنوع float32. يُثبِّت استدعاء jax.device_put جميع المصفوفات الثلاث على وحدة معالجة الرسومات التي حددتها في خلية الإعداد، لذلك لا يوجد شيء في المعايير لاحقًا يقيس نقل المضيف.

4. تنفيذ آلية الانتباه من الصفر

قبل استخدام نواة مدمجة، اكتب الصيغة مباشرةً لترى بدقة ما تستبدله النواة. يتبع هذا التنفيذ الخطوات الأربع بالترتيب: حساب النتائج، والتحجيم، وsoftmax، والضرب في القيم. يمكنك تبديل أبعاد الرأس والتسلسل لكي تعمل عملية matmul على محور التسلسل.

تعمل هذه الطريقة بشكل صحيح، ولكنّها تُطلق ثلاث عمليات منفصلة لنواة وحدة معالجة الرسومات من أجل ضرب مصفوفة النقاط، ودالة softmax، وضرب مصفوفة القيم، وتُنشئ مصفوفة أوزان الانتباه الكاملة (B, N, T, S) في ذاكرة وحدة معالجة الرسومات.

def naive_attention(q, k, v):
    """Scaled dot-product attention from scratch."""
    scale = 1.0 / math.sqrt(q.shape[-1])

    # (B, T, N, H) to (B, N, T, H) so matmul runs over the T/S axis
    q_t = jnp.transpose(q, (0, 2, 1, 3))
    k_t = jnp.transpose(k, (0, 2, 1, 3))
    v_t = jnp.transpose(v, (0, 2, 1, 3))

    # Score: (B, N, T, H) @ (B, N, H, S) to (B, N, T, S)
    scores = jnp.matmul(q_t, jnp.transpose(k_t, (0, 1, 3, 2))) * scale
    weights = jax.nn.softmax(scores, axis=-1)

    # Attend: (B, N, T, S) @ (B, N, S, H) to (B, N, T, H)
    out_t = jnp.matmul(weights, v_t)

    # Back to (B, T, N, H)
    return jnp.transpose(out_t, (0, 2, 1, 3))


naive_out = block_tree(naive_attention(q, k, v))

show_table(
    ["", "Value"],
    [
        ("Output shape", str(naive_out.shape)),
        ("Output dtype", str(naive_out.dtype)),
    ],
    title="Naive attention",
)

من المفترض أن يظهر لك شكل الناتج (4, 128, 8, 64) ونوع البيانات float32، وهو الشكل نفسه لـ Q، وهو ما وعدت به خطوة الناتج في الصيغة.

5- التبديل إلى dot_product_attention

يدمج jax.nn.dot_product_attention خطوات النتيجة والمقياس وsoftmax والتركيز في عملية واحدة. يمكن لـ JAX وXLA بعد ذلك تحسين نمط الوصول إلى الذاكرة. على وجه الخصوص، يمكنها تجنُّب تجسيد مصفوفة أوزان الانتباه الكاملة عندما يكون التسلسل طويلاً.

باستخدام الإعداد التلقائي implementation=None، يختار JAX أفضل واجهة خلفية متاحة تلقائيًا. على وحدة معالجة رسومات تتوفّر فيها مكتبة 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",
)

من المفترض أن يظهر لك الشكل ونوع البيانات نفسهما كما في الإصدار البسيط، وفرق مطلق صغير كحد أقصى، و"True" للتحقّق من التقارب في "atol=1e-3".

6. فرض الخلفية المدمجة لـ cuDNN

إنّ السماح لـ JAX باختيار الخلفية هو أمر مريح، ولكن من الخارج لا يمكنك معرفة النواة التي تم تشغيلها فعليًا. يؤدي ضبط الإعداد implementation="cudnn" إلى فرض استخدام نواة الاهتمام المدمجة في cuDNN من NVIDIA. وهي عبارة عن نواة لوحدات معالجة الرسومات تم تحسينها يدويًا لدمج عملية حساب الانتباه بالكامل في عملية تشغيل نواة واحدة مع أنماط محسّنة للوصول إلى الذاكرة.

يُرجى العِلم أنّ آلية الانتباه المدمجة في cuDNN تتطلّب أجهزة وأشكالاً معيّنة، مثل قدرة وحدة معالجة الرسومات على الحوسبة التي يجب أن تكون أكبر من أو تساوي 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.")

إذا لم تتمكّن وحدة معالجة الرسومات من تشغيل cuDNN SDPA، سيطبع الرمز السبب وسيستمر الدرس التطبيقي حول الترميز في الخلفية XLA.

7. إضافة قناع سببي

في نماذج الانحدار الذاتي (برامج فك الترميز بنمط GPT)، يمكن لكل موضع أن يركّز فقط على المواضع السابقة. يؤدي ضبط is_causal=True إلى تطبيق قناع المثلث السفلي هذا داخل النواة المدمجة، لذا لن تحتاج إلى إنشاء مصفوفة قناع بنفسك.

يشغّل الرمز أدناه الانتباه مرتين، مع القناع وبدونه، ويقارن بين موضعين لإظهار ما غيّره القناع.

causal_out = block_tree(
    jax.nn.dot_product_attention(q, k, v, is_causal=True)
)

# With causal masking, the last position attends to all positions.
# The first position attends only to itself.
nocausal_out = block_tree(
    jax.nn.dot_product_attention(q, k, v, is_causal=False)
)

# First position should differ
first_pos_diff = float(jnp.max(jnp.abs(causal_out[:, 0] - nocausal_out[:, 0])))
# Last position should be the same
last_pos_diff = float(jnp.max(jnp.abs(causal_out[:, -1] - nocausal_out[:, -1])))

show_table(
    ["Position", "Max diff (causal vs full)", "Expected"],
    [
        ("First (t=0)", f"{first_pos_diff:.4f}", "Large — causal restricts to self only"),
        ("Last (t=T-1)", f"{last_pos_diff:.2e}", "~0 — attends to all positions either way"),
    ],
    title="Causal masking effect on attention output",
)

من المفترض أن يظهر فرق كبير في الموضع الأول وفرق قريب من الصفر في الموضع الأخير. هذا التفاوت هو أنّ القناع في الموضع 0 يفقد كل مفتاح باستثناء مفتاحه، في حين أنّ الموضع النهائي يمكنه رؤية التسلسل بأكمله.

8. جارٍ إجراء بعض عمليات قياس الأداء

لديك الآن ثلاث طرق لحساب الدالة نفسها. ويعتمد اختيارك على شكل المشكلة وكيفية تقييمها.

ضبط توقيت ظهور الأشكال المختلفة بالشكل التلقائي

def benchmark_attention(fn, q, k, v, warmup=3, repeats=50):
    """Time an attention function. Returns median milliseconds per call."""
    jit_fn = jax.jit(fn)

    for _ in range(warmup):
        block_tree(jit_fn(q, k, v))

    times = []
    for _ in range(repeats):
        start = time.perf_counter()
        block_tree(jit_fn(q, k, v))
        times.append((time.perf_counter() - start) * 1000)

    return np.median(times)


t_naive = benchmark_attention(naive_attention, q, k, v)
t_sdpa = benchmark_attention(
    lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
    q, k, v,
)

results = [
    ("Naive (matmul + softmax + matmul)", f"{t_naive:.2f}"),
    ("SDPA (XLA, float32)", f"{t_sdpa:.2f}"),
]
bar_data = [
    ("Naive", t_naive),
    ("SDPA XLA f32", t_sdpa),
]

if HAS_CUDNN_SDPA:
    t_sdpa_bf16 = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
        q_bf16, k_bf16, v_bf16,
    )
    t_cudnn = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="cudnn"),
        q_bf16, k_bf16, v_bf16,
    )
    results.append(("SDPA (XLA, bfloat16)", f"{t_sdpa_bf16:.2f}"))
    results.append(("SDPA (cuDNN, bfloat16)", f"{t_cudnn:.2f}"))
    bar_data.append(("SDPA XLA bf16", t_sdpa_bf16))
    bar_data.append(("SDPA cuDNN bf16", t_cudnn))

show_table(
    ["Implementation", "Median ms/call"],
    results,
    title=f"Attention timing — B={BATCH}, T={SEQ_LEN}, N={NUM_HEADS}, H={HEAD_DIM}",
    aligns=["left", "right"],
)
show_bars(bar_data, "Attention latency (ms per call)", "ms", lower_is_better=True)

بالنسبة إلى حجم المشكلة الصغير هذا، تكون جميع عمليات التنفيذ سريعة جدًا وتقتصر في الغالب على النفقات العامة، لذا فإنّ الإصدار البسيط تنافسي. من المفترض أن يكون cuDNN bf16 أسرع قليلاً، ولكن ميزة "الانتباه المدمج" تصبح عادةً أكثر وضوحًا في أطوال التسلسلات الأطول، حيث يكون تجنُّب مصفوفة الانتباه الكاملة أمرًا مهمًا.

تغيير طول التسلسل

تزداد ميزة نواة الانتباه المدمجة مع طول التسلسل. يؤدي التنفيذ البسيط إلى إنشاء مصفوفة انتباه (B, N, T, S) في ذاكرة وحدة معالجة الرسومات، وهي ذاكرة O(T²). تُقسّم النواة المدمجة، مثل cuDNN FlashAttention، العمليات الحسابية إلى أجزاء صغيرة، وبالتالي لا تُنشئ المصفوفة الكاملة، ما يحافظ على استخدام الذاكرة عند O(T).

يقيس الاختبار أدناه وقت كل عملية تنفيذ على مستوى تسلسلات يتراوح طولها بين 64 و1024.

SEQ_LENS = [64, 128, 256, 512, 1024]

sweep_results = []
for sl in SEQ_LENS:
    rk = jax.random.key(sl)
    rk1, rk2, rk3 = jax.random.split(rk, 3)

    q_s = jax.random.normal(rk1, (BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    k_s = jax.random.normal(rk2, (BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    v_s = jax.random.normal(rk3, (BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    q_s, k_s, v_s = jax.device_put((q_s, k_s, v_s), device)

    q_sb = q_s.astype(jnp.bfloat16)
    k_sb = k_s.astype(jnp.bfloat16)
    v_sb = v_s.astype(jnp.bfloat16)

    row = {"seq_len": sl}

    row["naive_ms"] = benchmark_attention(
        naive_attention,
        q_s, k_s, v_s,
        warmup=2,
        repeats=20,
    )

    row["sdpa_xla_f32_ms"] = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
        q_s, k_s, v_s,
        warmup=2,
        repeats=20,
    )

    row["sdpa_xla_bf16_ms"] = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
        q_sb, k_sb, v_sb,
        warmup=2,
        repeats=20,
    )

    if HAS_CUDNN_SDPA:
        row["cudnn_bf16_ms"] = benchmark_attention(
            lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="cudnn"),
            q_sb, k_sb, v_sb,
            warmup=2,
            repeats=20,
        )

    sweep_results.append(row)


headers = ["Seq len", "Naive (ms)", "SDPA XLA f32 (ms)", "SDPA XLA bf16 (ms)"]
if HAS_CUDNN_SDPA:
    headers.append("cuDNN bf16 (ms)")

table_rows = []
for r in sweep_results:
    row = [
        r["seq_len"],
        f"{r['naive_ms']:.2f}",
        f"{r['sdpa_xla_f32_ms']:.2f}",
        f"{r['sdpa_xla_bf16_ms']:.2f}",
    ]

    if HAS_CUDNN_SDPA:
        row.append(f"{r['cudnn_bf16_ms']:.2f}")

    table_rows.append(row)

show_table(
    headers,
    table_rows,
    title=f"Sequence-length sweep — B={BATCH}, N={NUM_HEADS}, H={HEAD_DIM}",
    aligns=["right"] * len(headers),
)

يستغرق الفحص بعض الوقت، لأن كل طول تسلسل يجمع ثلاثة أو أربعة تطبيقات منفصلة قبل توقيتها. يجب أن ينتهي بك الأمر بصف جدول واحد لكل طول تسلسل من 64 إلى 1024.

يقارن هذا المخطط بين كيفية تغيُّر وقت استجابة الانتباه مع طول التسلسل في التنفيذ البسيط، و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()

من المفترض أن يظهر سطر واحد لكل عملية تنفيذ. ومن المفترض أن تظهر لك الفجوة التي تتسع مع زيادة طول التسلسل.

تغيير حجم الدفعة

تؤدي المجموعات الأكبر حجمًا إلى تقليل تكلفة تشغيل النواة وتحسين استخدام وحدة معالجة الرسومات إلى أن تصبح ذاكرة وحدة معالجة الرسومات هي نقطة الاختناق. يُثبّت الفحص أدناه طول التسلسل عند 256 ويغيّر حجم الدفعة.

BATCH_SIZES = [1, 2, 4, 8, 16]
SWEEP_SEQ = 256

batch_results = []
for bs in BATCH_SIZES:
    rk = jax.random.key(bs + 100)
    rk1, rk2, rk3 = jax.random.split(rk, 3)

    q_b = jax.random.normal(rk1, (bs, SWEEP_SEQ, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    k_b = jax.random.normal(rk2, (bs, SWEEP_SEQ, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    v_b = jax.random.normal(rk3, (bs, SWEEP_SEQ, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    q_b, k_b, v_b = jax.device_put((q_b, k_b, v_b), device)

    q_bb = q_b.astype(jnp.bfloat16)
    k_bb = k_b.astype(jnp.bfloat16)
    v_bb = v_b.astype(jnp.bfloat16)

    row = {"batch": bs}

    row["sdpa_xla_f32_ms"] = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
        q_b, k_b, v_b,
        warmup=2,
        repeats=20,
    )

    row["sdpa_xla_bf16_ms"] = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
        q_bb, k_bb, v_bb,
        warmup=2,
        repeats=20,
    )

    if HAS_CUDNN_SDPA:
        row["cudnn_bf16_ms"] = benchmark_attention(
            lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="cudnn"),
            q_bb, k_bb, v_bb,
            warmup=2,
            repeats=20,
        )

    batch_results.append(row)


headers = ["Batch size", "SDPA XLA f32 (ms)", "SDPA XLA bf16 (ms)"]
if HAS_CUDNN_SDPA:
    headers.append("cuDNN bf16 (ms)")

table_rows = []
for r in batch_results:
    row = [
        r["batch"],
        f"{r['sdpa_xla_f32_ms']:.2f}",
        f"{r['sdpa_xla_bf16_ms']:.2f}",
    ]

    if HAS_CUDNN_SDPA:
        row.append(f"{r['cudnn_bf16_ms']:.2f}")

    table_rows.append(row)

show_table(
    headers,
    table_rows,
    title=f"Batch-size sweep — T={SWEEP_SEQ}, N={NUM_HEADS}, H={HEAD_DIM}",
    aligns=["right"] * len(headers),
)


fig, ax = plt.subplots(figsize=(8, 5))
batches = [r["batch"] for r in batch_results]

ax.plot(
    batches,
    [r["sdpa_xla_f32_ms"] for r in batch_results],
    "s-",
    label="SDPA XLA f32",
    color="#0969da",
)

ax.plot(
    batches,
    [r["sdpa_xla_bf16_ms"] for r in batch_results],
    "d-",
    label="SDPA XLA bf16",
    color="#8250df",
)

if HAS_CUDNN_SDPA:
    ax.plot(
        batches,
        [r["cudnn_bf16_ms"] for r in batch_results],
        "^-",
        label="cuDNN bf16",
        color="#1a7f37",
    )

ax.set_xlabel("Batch size")
ax.set_ylabel("Median ms per call")
ax.set_title("Attention latency vs batch size")
ax.legend()
ax.grid(True, alpha=0.25)
ax.set_xticks(batches)

fig.tight_layout()
plt.show()

من المفترض أن تحصل على جدول ومخطط بياني يتضمّنان إدخالاً واحدًا لكل حجم دفعة من 1 إلى 16. تم استبعاد التنفيذ البسيط من هذا المسح، لذا فإن كلا المخرجين يقارنان فقط المسارات الثلاثة المدمجة.

9. مقارنة بين MHA وGQA وMQA

حتى الآن، كان كل رأس يحمل مفاتيحه وقيمه الخاصة. ولكن غالبًا ما تكسر برامج فك الترميز هذا التماثل، لأنّ ذاكرة التخزين المؤقت للمفتاح والقيمة هي التي تهيمن على الذاكرة أثناء الاستنتاج.

يمنح الانتباه المتعدد الرؤوس (MHA) كل رأس إسقاطات Q وK وV الخاصة به. يؤدي استخدام آلية الانتباه المجمّع للطلبات (GQA) وآلية الانتباه المتعدّد للطلبات (MQA) إلى تقليل عدد رؤوس المفاتيح والقيم (KV) لتوفير الذاكرة والقدرة الحاسوبية أثناء الاستنتاج.

يتعامل jax.nn.dot_product_attention مع الحالات الثلاث، ويتم بث رؤوس 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. هذا هو الهدف: يمكنك تقليل حجم ذاكرة التخزين المؤقت للمفاتيح والقيم بدون تغيير أي شيء في المراحل اللاحقة من عملية الانتباه.

10. مقارنة أداء NVIDIA TransformerEngine وFP8

توفّر TransformerEngine من NVIDIA وحدات اهتمام مدمجة ومحسّنة لوحدات معالجة الرسومات من NVIDIA. يستخدم تكامل JAX وحدات Flax Linen-style (وليس 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 إلى 2048. يقارن مقياس الأداء bf16 هذا بين TransformerEngine وJAX SDPA في عبء عمل الانتباه السببي نفسه، ولكن عادةً ما تظهر إمكانات الأداء الكاملة لـ TransformerEngine على وحدات معالجة الرسومات Hopper وBlackwell عندما يتوفّر البث التلقائي FP8.

فحص مسار FP8

على وحدات معالجة الرسومات Hopper (إمكانية الحوسبة >= 9.0، مثل H100)، يمكن لـ TransformerEngine تشغيل آلية الانتباه في FP8 لزيادة معدل النقل. يستخدم FP8 وصفة DelayedScaling التي تتتبّع سجلّ الحد الأقصى المطلق لكل موتر لحساب عوامل القياس الديناميكية.

  • تنسيق E4M3 لعملية التمرير الأمامي (4 بتات للأس، و3 بتات للجزء العشري)
  • تنسيق E5M2 لعملية التمرير العكسي (5 بتات للأس وبتان للجزء العشري)

إذا كانت وحدة معالجة الرسومات لا تتيح استخدام 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، من المفترض أن يظهر جدول يسمّي وحدة معالجة الرسومات ويشير إلى عدم رصد أي دعم لـ FP8، متبوعًا بسطرَين مطبوعَين يعرضان طلب te.autocast الذي ستستخدمه على وحدة معالجة رسومات 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 Console.

12. تهانينا

لقد انتقلت من عملية حسابية مكتوبة بخط اليد إلى نواة مدمجة محسّنة لوحدة معالجة الرسومات، وقست الفرق على أجهزة حقيقية.

ما تعلّمته

  • تعمل آلية الانتباه البسيطة، ولكنّها تشغّل عدة نواة لوحدة معالجة الرسومات وتُنشئ مصفوفة الانتباه الكاملة في الذاكرة
  • 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 اختيارية على وحدات معالجة الرسومات Hopper

الخطوات التالية

  • الدرس التطبيقي 6: توسيع نطاق التدريب باستخدام JAX على مستوى وحدات معالجة رسومات متعددة سيوضّح لك كيفية تقسيم المصفوفات على مستوى وحدات L4 وتنفيذ خطوات التدريب بالتوازي
  • غيِّر HEAD_DIM (جرِّب 32 و64 و128) وشاهِد القيم التي تقبلها مكتبة cuDNN وكيفية تغيُّر التوقيتات
  • يمكنك رفع قيمة SEQ_LEN (جرِّب 256 و512 و1024 و2048) ومراقبة استخدام الذاكرة ونسبة تسريع cuDNN
  • أعِد تشغيل مقارنة MHA وGQA وMQA مع رؤوس KV بقيمة K=2 وK=1 مقابل 8 رؤوس طلبات بحث، واشرح وفورات ذاكرة التخزين المؤقت لـ KV في الاستدلال

المستندات المرجعية