cuDNN と TransformerEngine を使用して GPU でアテンションを高速化する

1. はじめに

この Codelab では、NVIDIA GPU でのアテンションの計算に焦点を当てます。まず、単純な実装から始め、それを JAX の組み込みの融合アテンションに置き換えてから、同じオペレーションに対して NVIDIA の cuDNN バックエンドを有効にします。このラボを完了すると、各レベルで何が変化するのか、問題が大きくなるにつれて GPU 最適化パスがどれだけ高速になるのかを理解できるようになります。

演習内容

• 基本的な JAX 演算を使用して、Scaled Dot-Product Attention をゼロから実装する
• JAX の組み込み融合カーネルである jax.nn.dot_product_attention に置き換えます。
• implementation="cudnn" を使用して cuDNN バックエンドを強制し、因果マスキングを追加する
• シーケンス長とバッチサイズをスイープして、融合カーネルが効果を発揮する場所を確認する
• マルチヘッド アテンションの形状 MHA、GQA、MQA を比較する
• NVIDIA TransformerEngine のアテンションをベンチマークし、FP8 コードパスを検査する

必要なもの

• 課金が有効になっている Google Cloud プロジェクト、および GPU の使用量をカバーするワークショップ クレジットまたは予約
• 選択したリージョンで少なくとも 2 つの NVIDIA L4 GPU の割り当て（GPU 割り当てを確認する方法）
• Codelab 1 ～ 4 の完了、または同等の CUDA 対応 JAX GPU 環境
• 行列の乗算とソフトマックスに関する知識。CUDA や cuDNN の経験は不要です。

完了までの推定時間: 60 分。

アテンションの仕組み

Scaled Dot-Product Attention は、クエリ、キー、値の 3 つの入力を受け取り、次のように計算します。

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

スケーリング ファクタ 1 / sqrt(d_k) は、ヘッドの次元が増加するにつれてドット積が大きくなりすぎるのを防ぎます。ドット積が大きくなると、ソフトマックスが勾配が非常に小さい領域に押し込まれます。

4 つのステップがあり、それぞれが次のステップにフィードされます。

1. スコア - Q @ K.T: ドット積は、各クエリ位置が各キー位置にどれだけ注意を払うべきかを測定します。
2. スケール + ソフトマックス - softmax(scores / sqrt(d)): スケーリングにより勾配消失を防ぎ、ソフトマックスによりスコアを合計が 1 になるアテンション重みに変換します。
3. Attend - 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 が作成されるまで待ってから、Pod ログから 1 回限りの 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	Key-Value シーケンスの長さ	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) とデータ型 float32 が表示されます。jax.device_put 呼び出しは、3 つの配列すべてをセットアップ セルで選択した GPU に固定するため、後のベンチマークではホスト転送を測定しません。

4. アテンションをゼロから実装する

融合カーネルに手を伸ばす前に、数式を直接書き出して、カーネルが置き換える内容を正確に把握できるようにします。この実装は、スコアの計算、スケーリング、ソフトマックス、値の乗算という 4 つのステップを順番に実行します。ヘッドとシーケンスのディメンションを転置して、matmul がシーケンス軸で動作するようにします。

これは正しく機能しますが、スコア matmul、softmax、値 matmul に対して 3 つの個別の 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 は、スコア、スケール、ソフトマックス、アテンドのステップを 1 つのオペレーションに融合します。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 融合アテンション カーネルが強制的に使用されます。これらは、アテンション計算全体を 1 つのカーネル起動に融合し、メモリ アクセス パターンを最適化した、手動で最適化された GPU カーネルです。

cuDNN 融合アテンションには、GPU コンピューティング機能が 8.0 以上（Ampere 以降）であることや、入力 dtype が 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 を設定すると、この下三角マスクが融合カーネル内に適用されるため、マスク マトリックスを自分で作成する必要はありません。

以下のコードでは、マスクありとマスクなしでアテンションを 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 がわずかに最速になるはずですが、融合アテンションの利点は通常、完全なアテンション マトリックスを回避することが重要となる、より長いシーケンス長でより明確になります。

シーケンスの長さをスイープする

融合アテンション カーネルの利点は、シーケンス長が長くなるほど大きくなります。単純な実装では、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),
)

スイープには時間がかかります。これは、各シーケンス長で 3 つまたは 4 つの個別の実装をコンパイルしてから、それらの時間を測定するためです。最終的に、64 から 1024 までのシーケンス長ごとに 1 つのテーブル行が作成されます。

このプロットは、シーケンス長に伴うアテンション レイテンシの変化を、ナイーブ実装、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()

実装ごとに 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 ～ 16 のバッチサイズごとに 1 つのエントリを含むテーブルとプロットが表示されます。このスイープでは、単純な実装は除外されているため、どちらの出力も 3 つの融合パスのみを比較しています。

9. MHA、GQA、MQA の比較

これまでのところ、各ヘッドは独自のキーと値を持っていました。しかし、デコーダでは、推論中に Key-Value キャッシュがメモリを占有するため、この対称性が崩れることがよくあります。

マルチヘッド アテンション（MHA）では、各ヘッドに独自の Q、K、V 投影が割り当てられます。Grouped-Query Attention（GQA）と Multi-Query Attention（MQA）は、KV ヘッドの数を減らして、推論時のメモリとコンピューティングを節約します。

jax.nn.dot_product_attention は、K < N の場合に KV ヘッドが自動的にブロードキャストされることで、3 つすべてを処理します。

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

K と V の形状が 8 ヘッドから 2 ヘッド、1 ヘッドに縮小しても、3 つの行すべてで同じ出力形状 (2, 64, 8, 64) が報告されます。重要なのは、アテンションの下流を変更することなく、KV キャッシュを削減できるということです。

10. NVIDIA TransformerEngine と FP8 のベンチマーク

NVIDIA の TransformerEngine は、NVIDIA GPU 向けに最適化された融合アテンション モジュールを提供します。JAX 統合では、Flax Linen スタイルのモジュール（NNX ではない）が使用されるため、モジュールは 1 回初期化されてから、その変数が適用されます。

TransformerEngine を使用してシーケンス長をスイープする

このセルは、NVIDIA TransformerEngine が使用可能かどうかを確認し、複数のシーケンス長にわたって bf16 因果 DotProductAttention をベンチマークします。これは、同じワークロードで XLA バックエンドと cuDNN バックエンドを使用する JAX SDPA と比較されます。

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 のシーケンス長ごとに 1 つのテーブル行と 1 つのプロットポイントが表示されます。この bf16 ベンチマークでは、同じ因果アテンション ワークロードで TransformerEngine と JAX SDPA を比較していますが、TransformerEngine のパフォーマンスの可能性は、通常、FP8 自動キャストが利用可能な Hopper GPU と Blackwell GPU で発揮されます。

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 呼び出しを示す 2 行が表示されます。このスニペットは、FP8 が呼び出しサイトで必要とする唯一の変更です。

11. クリーンアップ

LoadBalancer と永続ボリュームを含む Jupyter ワークロードを削除します。

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 は計算を 1 つのオペレーションに融合し、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: Scale JAX training across multiple GPUs では、両方の L4 で配列をシャーディングし、トレーニング ステップを並列で実行する方法を説明します。
• HEAD_DIM を変更し（32、64、128 を試す）、cuDNN が受け入れる値とタイミングの変化を確認します。
• SEQ_LEN を増やし（256、512、1024、2048 を試す）、メモリ使用量と cuDNN の高速化率を確認します。
• 8 個のクエリヘッドに対して K=2 と K=1 の KV ヘッドを使用して MHA、GQA、MQA の比較を再実行し、推論における KV キャッシュの節約について推論します。

リファレンス ドキュメント

• jax.nn.dot_product_attention
• cuDNN 融合アテンション
• TransformerEngine JAX
• GitHub の TransformerEngine
• Flax NNX MultiHeadAttention
• NGC の NVIDIA JAX コンテナ