Google 会使用 AI 技术将内容翻译成您偏好的语言。AI 翻译可能包含错误。

使用 cuDNN 和 TransformerEngine 加快 GPU 上的注意力机制计算

1. 简介

Jax on GPU 学习路线。实验 5：GPU 上的注意力机制。

在此 Codelab 中，您将重点关注 NVIDIA GPU 上的注意力机制计算。您将从一个简单的实现开始，然后将其替换为 JAX 的内置融合注意力机制，最后为同一操作启用 NVIDIA 的 cuDNN 后端。最后，您将了解每个级别发生的变化，以及随着问题规模的扩大，GPU 优化路径的运行速度会提高多少。

您将执行的操作

使用基本 JAX 操作从头开始实现缩放点积注意力机制
将其替换为 jax.nn.dot_product_attention，即 JAX 的内置融合内核
使用 implementation="cudnn" 强制使用 cuDNN 后端，并添加因果掩码
扫描序列长度和批次大小，了解融合核在哪些情况下能发挥作用
比较多头注意力机制形状 MHA、GQA 和 MQA
对 NVIDIA TransformerEngine 注意力机制进行基准测试，并检查其 FP8 代码路径

所需条件

启用了结算功能的 Google Cloud 云项目，以及涵盖 GPU 使用量的研讨会赠金或预留
在所选区域中至少有 2 个 NVIDIA L4 GPU 的配额（如何查看 GPU 配额）
完成 Codelab 1 至 4，或拥有等效的支持 CUDA 的 JAX GPU 环境
熟悉矩阵乘法和 softmax。无需具备 CUDA 或 cuDNN 经验。

预计完成时间：60 分钟。

注意力机制的工作原理

缩放点积注意力机制接受三个输入：Query、Key 和 Value，并计算：

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

缩放比例 1 / sqrt(d_k) 可防止点积随着头维度的增加而变得过大，否则会导致 softmax 进入梯度极小的区域。

四个步骤，环环相扣：

得分 - Q @ K.T：点积衡量每个查询位置应关注每个键位置的程度。
缩放 + Softmax - softmax(scores / sqrt(d))：缩放可防止梯度消失，而 Softmax 可将分数转换为总和为 1 的注意力权重。
Attend - weights @ V：值向量的加权和会为每个查询位置生成输出。
输出 - 形状与 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 准备就绪，然后从 Pod 日志中读取一次性 JupyterLab 令牌：

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，粘贴令牌，然后在 /workspace 中创建一个新的 Python 3 笔记本。本 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 设备列表以及本 Codelab 其余部分使用的 GPU。

3. 构建 Q、K 和 V 测试数组

融合注意力核会关注输入数组在内存中的布局，因此在编写任何注意力代码之前，请先按照 JAX 预期的布局构建 Q、K 和 V。jax.nn.dot_product_attention 预期布局如下：

昏暗	含义	默认
B	批次大小	4
T	查询序列长度	128
S	键值对序列长度	128（与自注意力机制的 T 相同）
否	注意力头数	8
H	每个头的维度	64

对于简单实现，您从 float32 开始，使用随机值进行初始化。

BATCH = 4
SEQ_LEN = 128
NUM_HEADS = 8
HEAD_DIM = 64

key = jax.random.key(0)
k1, k2, k3 = jax.random.split(key, 3)

q = jax.random.normal(k1, (BATCH, SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
k = jax.random.normal(k2, (BATCH, SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
v = jax.random.normal(k3, (BATCH, SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)

q, k, v = jax.device_put((q, k, v), device)

show_table(
    ["Array", "Shape", "Dtype", "Layout"],
    [
        ("Q (query)", q.shape, q.dtype, "(B, T, N, H)"),
        ("K (key)", k.shape, k.dtype, "(B, S, N, H)"),
        ("V (value)", v.shape, v.dtype, "(B, S, N, H)"),
    ],
    title="Attention inputs on GPU",
)

您应该会看到一个表格，其中每个数组占一行，分别报告形状 (4, 128, 8, 64) 和 dtype float32。jax.device_put 调用会将所有三个数组固定到您在设置单元中选择的 GPU，因此后续基准测试中不会测量主机传输。

4. 从头开始实现注意力机制

在寻找融合内核之前，请直接写出公式，以便确切了解内核替换的内容。此实现按顺序遵循以下四个步骤：计算分数、缩放、softmax、乘以值。您需要转置头维度和序列维度，以便 matmul 在序列轴上运行。

它能正常运行，但会针对分数 matmul、softmax 和值 matmul 发出三次单独的 GPU 内核启动，并在 GPU 内存中实现完整的 (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) 和 dtype float32，其形状与 Q 相同，这正是公式的输出步骤所承诺的。

5. 切换到 dot_product_attention

jax.nn.dot_product_attention 将评分、缩放、softmax 和注意力步骤融合为单个操作。然后，JAX 和 XLA 可以优化内存访问模式。具体而言，当序列较长时，它们可以避免具体化完整的注意力权重矩阵。

如果使用默认的 implementation=None，JAX 会自动选择最佳可用后端。在具有 cuDNN 和兼容输入的 GPU 上，它可能已经在使用 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、较小的最大绝对差值，以及 atol=1e-3 接近度检查的 True。

6. 强制使用 cuDNN 融合后端

让 JAX 选择后端很方便，但从外部无法判断实际运行的是哪个内核。设置 implementation="cudnn" 会强制使用 NVIDIA 的 cuDNN 融合注意力内核。这些是经过手动优化的 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 后端上运行。

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 应该稍微快一些，但融合注意力机制的优势通常在序列长度较长时更为明显，因为此时避免计算完整注意力矩阵至关重要。

扫描序列长度

融合注意力机制内核的优势会随着序列长度的增加而增加。简单实现会在 GPU 内存中实例化一个 (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）。

此图比较了注意力延迟时间随序列长度的变化情况，涵盖了朴素实现、float32 和 bfloat16 格式的 JAX/XLA SDPA，以及 bfloat16 格式的 cuDNN 融合注意力机制。

fig, ax = plt.subplots(figsize=(8, 5))
seq_lens = [r["seq_len"] for r in sweep_results]

ax.plot(
    seq_lens,
    [r["naive_ms"] for r in sweep_results],
    "o-",
    label="Naive",
    color="#d1242f",
)

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

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

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

ax.set_xlabel("Sequence length")
ax.set_ylabel("Median ms per call")
ax.set_title("Attention latency vs sequence length")
ax.legend()
ax.grid(True, alpha=0.25)
ax.set_xticks(seq_lens)

fig.tight_layout()
plt.show()

您应该会看到每个实现各占一行。您应该会看到，随着序列变长，差距也会变大。

扫描批次大小

较大的批次可以分摊内核启动开销，并提高 GPU 利用率，直到 GPU 内存成为瓶颈。下面的扫描将序列长度固定为 256，并改变批次大小。

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

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

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

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

    row = {"batch": bs}

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

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

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

    batch_results.append(row)


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

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

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

    table_rows.append(row)

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


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

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

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

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

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

fig.tight_layout()
plt.show()

您应该会看到一个表格和一个图表，其中每个批次大小（从 1 到 16）对应一个条目。此扫描中不包含简单实现，因此两个输出仅比较三个融合路径。

9. 比较 MHA、GQA 和 MQA

到目前为止，每个头都有自己的键和值。但解码器往往会打破这种对称性，因为在推理过程中，键值对缓存会占用大部分内存。

多头注意力机制 (MHA) 为每个头提供各自的 Q、K 和 V 投影。分组查询注意力 (GQA) 和多查询注意力 (MQA) 可减少 KV 头的数量，从而在推理期间节省内存和计算资源。

jax.nn.dot_product_attention 会处理所有这三种情况，当 K < N 时，KV 头会自动广播。

rk = jax.random.key(42)
rk1, rk2, rk3, rk4, rk5 = jax.random.split(rk, 5)

q_mha = jax.random.normal(rk1, (2, 64, 8, 64), dtype=jnp.float32)

# MHA: 8 KV heads
k_mha = jax.random.normal(rk2, (2, 64, 8, 64), dtype=jnp.float32)
v_mha = jax.random.normal(rk3, (2, 64, 8, 64), dtype=jnp.float32)

# GQA: 2 KV heads (each shared by 4 query heads)
k_gqa = jax.random.normal(rk2, (2, 64, 2, 64), dtype=jnp.float32)
v_gqa = jax.random.normal(rk3, (2, 64, 2, 64), dtype=jnp.float32)

# MQA: 1 KV head (shared by all 8 query heads)
k_mqa = jax.random.normal(rk4, (2, 64, 1, 64), dtype=jnp.float32)
v_mqa = jax.random.normal(rk5, (2, 64, 1, 64), dtype=jnp.float32)

out_mha = block_tree(jax.nn.dot_product_attention(q_mha, k_mha, v_mha))
out_gqa = block_tree(jax.nn.dot_product_attention(q_mha, k_gqa, v_gqa))
out_mqa = block_tree(jax.nn.dot_product_attention(q_mha, k_mqa, v_mqa))

show_table(
    ["Pattern", "Q shape", "K shape", "V shape", "Output shape"],
    [
        ("MHA", q_mha.shape, k_mha.shape, v_mha.shape, out_mha.shape),
        ("GQA", q_mha.shape, k_gqa.shape, v_gqa.shape, out_gqa.shape),
        ("MQA", q_mha.shape, k_mqa.shape, v_mqa.shape, out_mqa.shape),
    ],
    title="Multi-head attention variants — all should produce the same output shape",
)

所有三行都应报告相同的输出形状 (2, 64, 8, 64)，即使 K 和 V 形状从 8 个头缩减到 2 个，再缩减到 1 个。这就是关键所在：您可以在不改变注意力机制下游任何内容的情况下，削减 KV 缓存。

10. 对 NVIDIA TransformerEngine 和 FP8 进行基准测试

NVIDIA 的 TransformerEngine 提供了针对 NVIDIA GPU 优化的融合注意力机制模块。JAX 集成使用 Flax Linen 样式的模块（而非 NNX），因此模块会初始化一次，然后应用其变量。

使用 TransformerEngine 扫描序列长度

此单元会检查 NVIDIA TransformerEngine 是否可用，然后针对同一工作负载，在 XLA 和 cuDNN 后端上，使用 JAX SDPA 对其 bf16 因果 DotProductAttention 在多个序列长度上进行基准测试。

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 GPU 上体现，前提是 FP8 自动转换可用。

检查 FP8 路径

在 Hopper GPU（计算能力 >= 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 支持，随后是两行输出，显示您在 Hopper GPU 上使用的 te.autocast 调用。保留该代码段：这是 FP8 在调用端唯一需要进行的更改。

11. 清理

删除 Jupyter 工作负载，包括 LoadBalancer 和永久性卷：

kubectl delete -f deploy/jupyter.yaml

销毁集群、节点池、VPC 和服务账号：

cd terraform
terraform destroy

出现提示时，输入 yes，然后确认没有遗漏任何内容：

gcloud container clusters list
gcloud compute instances list

对于此项目，两者都应为空。如果您专门为此系列创建了一个项目，则可以改为从 Cloud 控制台中删除整个项目。

12. 恭喜

您从手写注意力计算过渡到了 GPU 优化的融合内核，并在实际硬件上测量了差异。

您学到的内容

朴素注意力机制虽然有效，但会启动多个 GPU 内核，并在内存中实例化完整的注意力矩阵
jax.nn.dot_product_attention 会将计算融合到单个操作中，而使用 implementation=None 时，JAX 会自动选择最佳后端
implementation="cudnn" 强制使用 NVIDIA 的 cuDNN 融合注意力机制内核，这些内核在长序列上速度最快，并且需要 bfloat16 或 float16 输入以及 8.0 或更高版本的计算能力
is_causal=True 的因果遮盖内置于融合内核中，无需手动遮盖矩阵
GQA 和 MQA 可减少 KV 头，从而在推理期间节省内存，而 dot_product_attention 会自动处理广播
TransformerEngine 在 Hopper GPU 上提供融合注意力机制，并可选择 FP8 精度

后续步骤

Codelab 6：在多个 GPU 上扩缩 JAX 训练将向您展示如何在两个 L4 上对数组进行分片，并并行运行训练步骤
更改 HEAD_DIM（尝试 32、64、128），观察 cuDNN 接受哪些值以及时间如何变化
提高 SEQ_LEN（尝试 256、512、1024、2048），并观察内存用量和 cuDNN 加速比
使用 K=2 和 K=1 KV 头针对 8 个查询头重新运行 MHA、GQA 和 MQA 比较，并推断推理时 KV 缓存节省情况