เร่งความเร็ว Attention ใน GPU ด้วย cuDNN และ TransformerEngine

1. บทนำ

เส้นทางการเรียนรู้ Jax บน GPU Lab 5: Attention on GPU

ใน Codelab นี้ คุณจะมุ่งเน้นไปที่การคำนวณ ความสนใจบน GPU ของ NVIDIA โดยจะเริ่มต้นด้วยการใช้งานแบบง่ายๆ จากนั้นแทนที่ด้วยความสนใจแบบผสานในตัวของ JAX แล้วเปิดใช้งานแบ็กเอนด์ cuDNN ของ NVIDIA สำหรับการดำเนินการเดียวกัน เมื่อสิ้นสุด Codelab คุณจะทราบว่ามีการเปลี่ยนแปลงอะไรบ้างในแต่ละระดับ และเส้นทางที่เพิ่มประสิทธิภาพสำหรับ GPU จะเร็วขึ้นมากเพียงใดเมื่อปัญหาใหญ่ขึ้น

สิ่งที่คุณต้องดำเนินการมีดังต่อไปนี้

  • ใช้ความสนใจแบบจุดผลิตภัณฑ์ที่ปรับขนาดตั้งแต่ต้นด้วยการดำเนินการพื้นฐานของ JAX
  • แทนที่ด้วย jax.nn.dot_product_attention ซึ่งเป็นเคอร์เนลแบบผสานในตัวของ JAX
  • บังคับใช้แบ็กเอนด์ cuDNN ด้วย implementation="cudnn" และเพิ่มการมาสก์แบบมีเหตุและผล
  • กวาดความยาวลำดับและขนาดกลุ่มเพื่อดูว่าเคอร์เนลแบบผสานมีประโยชน์เมื่อใด
  • เปรียบเทียบรูปร่างความสนใจแบบหลายหัว (MHA), ความสนใจแบบหลายคิวรี (GQA) และความสนใจแบบหลายคีย์ (MQA)
  • เปรียบเทียบความสนใจของ NVIDIA TransformerEngine และตรวจสอบเส้นทางโค้ด FP8

สิ่งที่คุณต้องมี

  • โปรเจ็กต์ Google Cloud ที่เปิดใช้การเรียกเก็บเงินแล้ว รวมถึงเครดิตของเวิร์กช็อปหรือการจองที่ครอบคลุมการใช้งาน GPU
  • โควต้าสำหรับ GPU ของ NVIDIA L4 อย่างน้อย 2 รายการในภูมิภาคที่เลือก (วิธีตรวจสอบโควต้า GPU)
  • ทำ Codelab 1 ถึง 4 เสร็จสมบูรณ์ หรือมีสภาพแวดล้อม GPU ของ JAX ที่เปิดใช้ CUDA แล้วเทียบเท่า
  • ความคุ้นเคยกับการคูณเมทริกซ์และ Softmax ไม่จำเป็นต้องมีประสบการณ์ด้าน CUDA หรือ cuDNN

เวลาที่ใช้โดยประมาณ: 60 นาที

วิธีการทำงานของความสนใจ

ความสนใจแบบจุดผลิตภัณฑ์ที่ปรับขนาดจะรับอินพุต 3 รายการ ได้แก่ คิวรี คีย์ และค่า แล้วคำนวณดังนี้

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

ปัจจัยการปรับขนาด 1 / sqrt(d_k) จะป้องกันไม่ให้จุดผลิตภัณฑ์มีขนาดใหญ่เกินไปเมื่อมิติข้อมูลของหัวเพิ่มขึ้น ซึ่งจะทำให้ Softmax เข้าสู่ภูมิภาคที่การไล่ระดับสีมีขนาดเล็ก

4 ขั้นตอน โดยแต่ละขั้นตอนจะป้อนข้อมูลให้กับขั้นตอนถัดไป

  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>

จัดเตรียมสภาพแวดล้อม 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 ใช้เวลาประมาณ 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 โค้ดบล็อกทุกบล็อกใน 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 ใช้

3. สร้างอาร์เรย์ทดสอบ Q, K และ V

เคอร์เนลความสนใจแบบผสานจะพิจารณาว่าอาร์เรย์อินพุตจัดวางอย่างไรในหน่วยความจำ ดังนั้นให้สร้าง Q, K และ V ในเลย์เอาต์ที่ JAX คาดหวังไว้ก่อนที่จะเขียนโค้ดความสนใจ jax.nn.dot_product_attention คาดหวังเลย์เอาต์ต่อไปนี้

Dim

ความหมาย

ค่าเริ่มต้นของเรา

B

ขนาดกลุ่ม

4

T

ความยาวลำดับคิวรี

128

S

ความยาวลำดับคีย์/ค่า

128 (เท่ากับ T สำหรับความสนใจตนเอง)

N

จำนวนหัวความสนใจ

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

คุณควรเห็นตารางที่มี 1 แถวต่ออาร์เรย์ โดยแต่ละแถวจะรายงานรูปร่าง (4, 128, 8, 64) และ dtype float32 การเรียก jax.device_put จะตรึงอาร์เรย์ทั้ง 3 รายการไว้กับ GPU ที่คุณเลือกในเซลล์การตั้งค่า ดังนั้นการเปรียบเทียบประสิทธิภาพในภายหลังจึงไม่ได้วัดการโอนโฮสต์

4. ใช้ความสนใจตั้งแต่ต้น

ก่อนที่จะใช้เคอร์เนลแบบผสาน ให้เขียนสูตรออกมาโดยตรงเพื่อให้คุณเห็นสิ่งที่เคอร์เนลแทนที่ได้อย่างชัดเจน การใช้งานนี้ทำตาม 4 ขั้นตอนตามลำดับ ได้แก่ คำนวณคะแนน ปรับขนาด Softmax และคูณด้วยค่า คุณจะเปลี่ยนมิติข้อมูลของหัวและลำดับเพื่อให้ matmul ทำงานผ่านแกนลำดับ

การดำเนินการนี้ทำงานได้อย่างถูกต้อง แต่จะเปิดใช้เคอร์เนล GPU แยกกัน 3 รายการสำหรับการคูณเมทริกซ์คะแนน, 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 ซึ่งเป็นสิ่งที่ขั้นตอนเอาต์พุต ของสูตรสัญญาไว้

5. เปลี่ยนไปใช้ dot_product_attention

jax.nn.dot_product_attention จะผสานขั้นตอนการให้คะแนน การปรับขนาด Softmax และความสนใจเป็นการดำเนินการเดียว จากนั้น JAX และ XLA จะเพิ่มประสิทธิภาพรูปแบบการเข้าถึงหน่วยความจำได้ โดยเฉพาะอย่างยิ่ง จะหลีกเลี่ยงการสร้างเมทริกซ์น้ำหนักความสนใจแบบเต็มเมื่อลำดับยาว

เมื่อ implementation=None เป็นค่าเริ่มต้น JAX จะเลือกแบ็กเอนด์ที่ดีที่สุดที่มีอยู่โดยอัตโนมัติ ใน GPU ที่มี cuDNN และอินพุตที่เข้ากันได้ JAX อาจใช้ cuDNN อยู่แล้ว ในฮาร์ดแวร์อื่นๆ JAX จะกลับไปใช้ 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 เลือกแบ็กเอนด์นั้นสะดวก แต่จากภายนอกคุณจะบอกไม่ได้ว่าเคอร์เนลใดทำงานจริง การตั้งค่า implementation="cudnn" จะบังคับใช้เคอร์เนลความสนใจแบบผสาน cuDNN ของ NVIDIA ซึ่งเป็นเคอร์เนล GPU ที่เพิ่มประสิทธิภาพด้วยตนเองซึ่งผสานการคำนวณความสนใจทั้งหมดเป็นการเปิดใช้เคอร์เนลเดียวด้วยรูปแบบการเข้าถึงหน่วยความจำที่เพิ่มประสิทธิภาพ

โปรดทราบว่าความสนใจแบบผสาน cuDNN มีข้อกำหนดด้านฮาร์ดแวร์และรูปร่าง เช่น ความสามารถในการคำนวณของ GPU ซึ่งควรมีค่า >= 8.0 (Ampere หรือใหม่กว่า) หรือ float16 หรือ bfloat16 เป็น dtype อินพุต หากไม่เป็นไปตามข้อกำหนดและคุณตั้งค่า 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

7. เพิ่มการมาสก์แบบมีเหตุและผล

ในโมเดลแบบถดถอยอัตโนมัติ (ตัวถอดรหัสสไตล์ GPT) แต่ละตำแหน่งจะให้ความสนใจได้เฉพาะตำแหน่งก่อนหน้า การตั้งค่า is_causal=True จะใช้มาสก์สามเหลี่ยมล่างนี้ภายในเคอร์เนลแบบผสาน คุณจึงไม่จำเป็นต้องสร้างเมทริกซ์มาสก์ด้วยตนเอง

โค้ดด้านล่างจะเรียกใช้ความสนใจ 2 ครั้ง โดยมีและไม่มีมาสก์ แล้วเปรียบเทียบ 2 ตำแหน่งเพื่อแสดงสิ่งที่มาสก์เปลี่ยนแปลง

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. การเปรียบเทียบประสิทธิภาพบางอย่าง

ตอนนี้คุณมี 3 วิธีในการคำนวณฟังก์ชันเดียวกัน วิธีที่จะใช้ขึ้นอยู่กับรูปร่างของปัญหาทั้งหมด และคุณสามารถประเมินได้

กำหนดเวลาตัวแปรที่รูปร่างเริ่มต้น

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)

การกวาดด้านล่างจะกำหนดเวลาการใช้งานแต่ละรายการตามความยาวลำดับตั้งแต่ 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),
)

การกวาดใช้เวลาสักครู่ เนื่องจากความยาวลำดับแต่ละรายการจะคอมไพล์การใช้งานแยกกัน 3 หรือ 4 รายการก่อนที่จะกำหนดเวลา คุณควรได้ตาราง 1 แถวต่อความยาวลำดับตั้งแต่ 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()

คุณควรเห็น 1 เส้นต่อการใช้งาน และคุณควรเห็นช่องว่างที่เพิ่มขึ้นเมื่อลำดับยาวขึ้น

กวาดขนาดกลุ่ม

กลุ่มที่ใหญ่ขึ้นจะลดค่าใช้จ่ายในการดำเนินการของการเปิดใช้เคอร์เนลและเพิ่มการใช้งาน 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 รายการต่อขนาดกลุ่มตั้งแต่ 1 ถึง 16 การใช้งานแบบง่ายๆ จะถูกละเว้นจากการกวาดนี้ ดังนั้นเอาต์พุตทั้ง 2 รายการจะเปรียบเทียบเฉพาะเส้นทางแบบผสาน 3 เส้นทาง

9. เปรียบเทียบ MHA, GQA และ MQA

จนถึงตอนนี้ แต่ละหัวมีคีย์และค่าของตัวเอง แต่ตัวถอดรหัสมักจะทำลายความสมมาตรนั้น เนื่องจากระหว่างการอนุมาน แคชคีย์-ค่าคือสิ่งที่ครอบงำหน่วยความจำ

ความสนใจแบบหลายหัว (MHA) จะให้การฉายภาพ Q, K และ V ของตัวเองแก่แต่ละหัว ความสนใจแบบหลายคิวรี (GQA) และความสนใจแบบหลายคีย์ (MQA) จะลดจำนวนหัว KV เพื่อประหยัดหน่วยความจำและการคำนวณระหว่างการอนุมาน

jax.nn.dot_product_attention จัดการทั้ง 3 รายการโดยมีการออกอากาศหัว 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",
)

ทั้ง 3 แถวควรรายงานรูปร่างเอาต์พุตเดียวกัน (2, 64, 8, 64) แม้ว่ารูปร่าง K และ V จะลดลงจาก 8 หัวเป็น 2 และ 1 นั่นคือประเด็นสำคัญ คุณสามารถตัดแคช KV ได้โดยไม่ต้องเปลี่ยนแปลงสิ่งใดก็ตามที่อยู่ปลายน้ำของความสนใจ

10. เปรียบเทียบประสิทธิภาพ NVIDIA TransformerEngine และ FP8

TransformerEngine ของ NVIDIA มีโมดูลความสนใจแบบผสานที่เพิ่มประสิทธิภาพสำหรับ GPU ของ NVIDIA การผสานรวม JAX ใช้โมดูลสไตล์ Flax Linen (ไม่ใช่ NNX) ดังนั้นระบบจะเริ่มต้นโมดูลเพียงครั้งเดียว จากนั้นใช้กับตัวแปรของโมดูล

กวาดความยาวลำดับด้วย TransformerEngine

เซลล์นี้จะตรวจสอบว่า NVIDIA TransformerEngine พร้อมใช้งานหรือไม่ จากนั้นเปรียบเทียบประสิทธิภาพ DotProductAttention แบบมีเหตุและผล bf16 กับความยาวลำดับหลายรายการเทียบกับ 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()

คุณควรได้ตาราง 1 แถวและจุดพล็อต 1 จุดต่อความยาวลำดับตั้งแต่ 128 ถึง 2048 การเปรียบเทียบประสิทธิภาพ bf16 นี้จะเปรียบเทียบ TransformerEngine กับ JAX SDPA ในภาระงานความสนใจแบบมีเหตุและผลเดียวกัน แต่โดยปกติแล้วศักยภาพด้านประสิทธิภาพเต็มรูปแบบของ 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 ตามด้วยบรรทัดที่พิมพ์ 2 บรรทัดที่แสดงการเรียก te.autocast ที่คุณจะใช้ใน GPU ของ Hopper เก็บข้อมูลโค้ดนั้นไว้ เนื่องจากเป็นเพียงการเปลี่ยนแปลงที่ FP8 กำหนดไว้ในตำแหน่งการเรียก

11. ล้างข้อมูล

ลบภาระงาน Jupyter รวมถึง LoadBalancer และโวลุ่มถาวร

kubectl delete -f deploy/jupyter.yaml

ทำลายคลัสเตอร์ Node Pool VPC และบัญชีบริการ

cd terraform
terraform destroy

พิมพ์ yes เมื่อได้รับข้อความแจ้ง จากนั้นยืนยันว่าไม่มีสิ่งใดเหลืออยู่

gcloud container clusters list
gcloud compute instances list

ทั้ง 2 รายการควรว่างเปล่าสำหรับโปรเจ็กต์นี้ หากคุณสร้างโปรเจ็กต์ขึ้นมาเพื่อใช้กับชุดข้อมูลนี้โดยเฉพาะ คุณสามารถลบโปรเจ็กต์ทั้งหมดออกจากคอนโซล 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: ปรับขนาดการฝึก JAX ใน GPU หลายรายการ** จะแสดงวิธีแบ่งอาร์เรย์ใน L4 ทั้ง 2 รายการและการเรียกใช้ขั้นตอนการฝึกแบบขนาน
  • เปลี่ยน HEAD_DIM (ลองใช้ 32, 64, 128) และดูว่า cuDNN ยอมรับค่าใดบ้าง รวมถึงการกำหนดเวลาเปลี่ยนแปลงอย่างไร
  • เพิ่ม SEQ_LEN (ลองใช้ 256, 512, 1024, 2048) และดูการใช้งานหน่วยความจำและอัตราส่วนการเพิ่มความเร็ว cuDNN
  • เรียกใช้การเปรียบเทียบ MHA, GQA และ MQA อีกครั้งโดยใช้หัว KV K=2 และ K=1 เทียบกับหัวคิวรี 8 หัว และพิจารณาการประหยัดแคช KV สำหรับการอนุมาน

เอกสารอ้างอิง