cuDNN 및 TransformerEngine을 사용하여 GPU에서 어텐션 속도 높이기

1. 소개

이 Codelab에서는 NVIDIA GPU에서 어텐션 계산 에 집중합니다. 기본 구현으로 시작하여 JAX의 기본 융합 어텐션으로 대체한 다음 동일한 작업에 NVIDIA의 cuDNN 백엔드를 활성화합니다. 결국 각 수준에서 변경되는 사항과 문제가 커짐에 따라 GPU 최적화 경로가 얼마나 빨라지는지 알게 됩니다.

실습할 내용

• 기본 JAX 작업을 사용하여 처음부터 확장된 점곱 어텐션을 구현합니다.
• 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분

어텐션 작동 방식

확장된 점곱 어텐션은 세 가지 입력(쿼리, 키, 값)을 가져와 다음을 계산합니다.

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. 어텐션 — 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분 정도 걸립니다. 완료되면 포드와 LoadBalancer를 기다린 후 포드 로그에서 일회성 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와 동일)
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",
)

배열당 하나의 행이 있는 표가 표시되며 각 행은 도형 (4, 128, 8, 64) 및 dtype float32를 보고합니다. jax.device_put 호출은 설정 셀에서 선택한 GPU에 세 배열을 모두 고정하므로 나중에 벤치마크에서 호스트 전송을 측정하지 않습니다.

4. 처음부터 어텐션 구현

융합 커널에 도달하기 전에 커널이 대체하는 항목을 정확히 확인할 수 있도록 수식을 직접 작성합니다. 이 구현은 점수 계산, 확장, 소프트맥스, 값 곱셈의 4단계를 순서대로 따릅니다. matmul이 시퀀스 축에서 작동하도록 헤드 및 시퀀스 크기를 전치합니다.

올바르게 작동하지만 점수 matmul, 소프트맥스, 값 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은 점수, 확장, 소프트맥스, 어텐션 단계를 단일 작업으로 융합합니다. 그러면 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을 입력 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를 설정하면 융합 커널 내에 이 하위 삼각 마스크가 적용되므로 마스크 행렬을 직접 빌드할 필요가 없습니다.

아래 코드는 마스크가 있는 경우와 없는 경우에 어텐션을 두 번 실행하고 두 위치를 비교하여 마스크가 변경한 항목을 보여줍니다.

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에 가까운 차이가 표시됩니다. 이 비대칭은 위치 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",
)

K 및 V 도형이 8개의 헤드에서 2개에서 1개로 줄어들더라도 세 행 모두 동일한 출력 도형 (2, 64, 8, 64)를 보고해야 합니다. 이것이 요점입니다. 어텐션의 다운스트림을 변경하지 않고 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의 전체 성능 잠재력은 일반적으로 FP8 자동 캐스트를 사용할 수 있는 Hopper 및 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 호출을 보여주는 두 개의 인쇄된 줄이 표시됩니다. 이 스니펫을 유지합니다. 이는 호출 사이트에서 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 은 계산을 단일 작업으로 융합하고 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 속도 향상 비율을 확인합니다.
• 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 컨테이너