1. Введение
В этой практической работе вы сосредоточитесь на вычислении внимания на графических процессорах NVIDIA. Вы начнете с простой реализации, замените ее встроенным объединенным вниманием JAX, а затем активируете серверную часть cuDNN от NVIDIA для той же операции. В конце вы узнаете, что меняется на каждом уровне и насколько быстрее становятся пути, оптимизированные для графического процессора, по мере роста проблемы.
Что вам предстоит сделать
- Реализуйте масштабированное скалярное произведение с помощью базовых операций JAX
- Замените его на
jax.nn.dot_product_attention, встроенное объединенное ядро JAX. - Принудительно использовать бэкенд cuDNN с
implementation="cudnn"и добавить каузальную маску - Проверьте длину последовательности и размер пакета, чтобы узнать, где объединенные ядра приносят пользу.
- Сравнение форм многоголового внимания MHA, GQA и MQA
- Оцените производительность NVIDIA TransformerEngine и проверьте путь выполнения FP8
Что вам понадобится
- Проект Google Cloud с включенной оплатой и бонусами за участие в семинаре или бронированием, покрывающим использование графического процессора.
- Квота как минимум на два графических процессора NVIDIA L4 в выбранном регионе (узнайте, как проверить квоту на графические процессоры).
- Выполнение практических работ 1–4 или эквивалентная среда графического процессора JAX с поддержкой CUDA.
- Знакомство с умножением матриц и функцией softmax. Опыт работы с CUDA или cuDNN не требуется.
Примерное время выполнения: 60 минут.
Как работает внимание
Масштабированное скалярное произведение принимает три входных значения – запрос, ключ и значение – и вычисляет:
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. - Внимание –
weights @ V: взвешенная сумма векторов значений дает выходные данные для каждой позиции запроса. - Выходные данные – та же форма, что и у Q: каждая позиция запроса теперь содержит информацию из позиций, на которые она обратила внимание.
2. Подготовка
Выберите проект
В консоли Google Cloud выберите или создайте проект с включенными платежными функциями.
Открыть Cloud Shell
Нажмите Activate Cloud Shell (Активировать Cloud Shell) (значок терминала в правом верхнем углу консоли), чтобы запустить сеанс Cloud Shell, а затем укажите свой проект:
gcloud config set project <YOUR_PROJECT_ID>
Как подготовить среду графического процессора
Выполните в Cloud Shell следующую команду: В практической работе 1 подробно описана каждая команда, в том числе квота на графические процессоры и требования к зонам.
gcloud services enable \
container.googleapis.com \
compute.googleapis.com \
iam.googleapis.com \
cloudresourcemanager.googleapis.com \
logging.googleapis.com \
monitoring.googleapis.com
git clone https://github.com/Google-Cloud-AI/partner-ai-nvidia.git
cd partner-ai-nvidia/05-workshops/jax-on-gpu/terraform
cp terraform.tfvars.example terraform.tfvars
Измените terraform.tfvars и задайте project_id = ", а затем подготовьте кластер и разверните JupyterLab:
terraform init
terraform apply
$(terraform output -raw get_credentials_command)
cd ..
kubectl apply -f deploy/jupyter.yaml
terraform apply занимает около 12 минут. Дождитесь завершения процесса, а затем Pod и LoadBalancer. После этого прочитайте одноразовый токен JupyterLab из журнала Pod:
kubectl get pod jax-jupyter -w # wait for Running, then Ctrl+C
kubectl get svc jax-jupyter-svc -w # wait for EXTERNAL-IP, then Ctrl+C
kubectl logs jax-jupyter | grep -o 'token=[a-z0-9]*' | head -1
Откройте http://, вставьте токен и создайте новый блокнот Python 3 в /workspace. Каждый блок кода в этой практической работе помещается в ячейку блокнота.
Установите необходимые инструменты для этой практической работы
!pip install --quiet matplotlib flax
Как настроить и проверить графический процессор
Импортируйте JAX и убедитесь, что в качестве серверной части по умолчанию используется графический процессор. В этой ячейке также определяются block_tree, show_table и show_bars – вспомогательные элементы, которые используются на каждом последующем шаге для ожидания работы устройства и рендеринга результатов.
import os
os.environ["LD_LIBRARY_PATH"] = "/usr/local/nvidia/lib64:" + os.environ.get("LD_LIBRARY_PATH", "")
import html
import math
import time
from functools import partial
from IPython.display import HTML, display
import matplotlib.pyplot as plt
import numpy as np
import jax
import jax.numpy as jnp
devices = jax.devices()
gpu_devices = [d for d in devices if d.platform == "gpu"]
device = gpu_devices[0] if gpu_devices else None
print(f"JAX version: {jax.__version__}")
print(f"Default backend: {jax.default_backend()}")
print(f"Devices: {devices}")
assert gpu_devices, f"This lab assumes a GPU backend. Available devices: {devices}"
print(f"Using GPU: {device}")
def block_tree(tree):
"""Wait until a PyTree of JAX arrays is ready on device."""
return jax.block_until_ready(tree)
def show_table(headers, rows, title=None, aligns=None):
"""Render rows as an HTML table."""
aligns = aligns or ["left"] * len(headers)
parts = ["<div style='font-family: system-ui; max-width: 980px;'>"]
if title:
parts.append(f"<h4 style='margin: 0 0 8px 0;'>{html.escape(title)}</h4>")
parts.append("<table style='border-collapse: collapse; width: 100%; font-size: 13px;'>")
parts.append("<thead><tr>")
for h, a in zip(headers, aligns):
parts.append(
f"<th style='text-align:{a}; border-bottom:1px solid #d0d7de; padding:6px;'>"
f"{html.escape(str(h))}</th>"
)
parts.append("</tr></thead><tbody>")
for row in rows:
parts.append("<tr>")
for cell, a in zip(row, aligns):
parts.append(
f"<td style='text-align:{a}; border-bottom:1px solid #eef1f4; padding:6px;'>"
f"{html.escape(str(cell))}</td>"
)
parts.append("</tr>")
parts.append("</tbody></table></div>")
display(HTML("".join(parts)))
def show_bars(rows, title, unit="", lower_is_better=False):
"""Render (label, value) pairs as a horizontal bar chart in HTML."""
max_value = max(float(value) for _, value in rows) or 1.0
color = "#1a7f37" if not lower_is_better else "#0969da"
parts = ["<div style='font-family: Arial, sans-serif; max-width: 760px;'>"]
parts.append(f"<h4 style='margin: 0 0 8px 0;'>{html.escape(title)}</h4>")
for label, value in rows:
width = max(3, 100 * float(value) / max_value)
parts.append(
"<div style='display:grid; grid-template-columns: 190px 1fr 130px; gap: 8px; "
"align-items:center; margin: 6px 0;'>"
f"<div style='font-size:13px;'>{html.escape(str(label))}</div>"
"<div style='background:#f6f8fa; border-radius:6px; overflow:hidden; height:22px;'>"
f"<div style='height:22px; width:{width:.1f}%; background:{color};'></div></div>"
f"<div style='font-size:13px; font-variant-numeric: tabular-nums;'>{float(value):,.1f} {html.escape(unit)}</div>"
"</div>"
)
parts.append(
f"<div style='font-size:12px; color:#57606a;'>"
f"{'Lower' if lower_is_better else 'Higher'} is better.</div></div>"
)
display(HTML("".join(parts)))
Вы должны увидеть версию JAX, gpu в качестве бэкенда по умолчанию, список устройств CUDA и графический процессор, который используется в оставшейся части практической работы.
3. Создание тестовых массивов Q, K и V
Для ядер объединенного внимания важно, как входные массивы расположены в памяти, поэтому перед написанием кода внимания создайте Q, K и V в макете, который ожидает JAX. jax.nn.dot_product_attention ожидает следующий макет:
Затемнить | Значение | Наш вариант по умолчанию |
B | Размер пакета | 4 |
T | Длина последовательности запросов | 128 |
С | Длина последовательности ключей и значений | 128 (то же, что и T для самовнимания) |
НЕТ | Количество голов внимания | 8 |
H | Параметр "На человека" | 64 |
Вы начинаете с float32 для наивной реализации, инициализируя случайными значениями.
BATCH = 4
SEQ_LEN = 128
NUM_HEADS = 8
HEAD_DIM = 64
key = jax.random.key(0)
k1, k2, k3 = jax.random.split(key, 3)
q = jax.random.normal(k1, (BATCH, SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
k = jax.random.normal(k2, (BATCH, SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
v = jax.random.normal(k3, (BATCH, SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
q, k, v = jax.device_put((q, k, v), device)
show_table(
["Array", "Shape", "Dtype", "Layout"],
[
("Q (query)", q.shape, q.dtype, "(B, T, N, H)"),
("K (key)", k.shape, k.dtype, "(B, S, N, H)"),
("V (value)", v.shape, v.dtype, "(B, S, N, H)"),
],
title="Attention inputs on GPU",
)
Вы увидите таблицу, в которой для каждого массива будет указана одна строка с формой (4, 128, 8, 64) и типом данных float32. Вызов jax.device_put закрепляет все три массива за графическим процессором, выбранным в ячейке настройки, поэтому в последующих тестах производительности не измеряется передача данных на хост.
4. Реализация механизма внимания с нуля
Прежде чем использовать объединенное ядро, запишите формулу напрямую, чтобы точно видеть, что заменяет ядро. Эта реализация состоит из четырех шагов: вычисление оценок, масштабирование, softmax, умножение на значения. Вы транспонируете размеры головы и последовательности, чтобы matmul работал по оси последовательности.
Она работает правильно, но запускает три отдельных ядра графического процессора для умножения матриц оценок, softmax и умножения матриц значений, а также материализует полную матрицу весов внимания (B, N, T, S) в памяти графического процессора.
def naive_attention(q, k, v):
"""Scaled dot-product attention from scratch."""
scale = 1.0 / math.sqrt(q.shape[-1])
# (B, T, N, H) to (B, N, T, H) so matmul runs over the T/S axis
q_t = jnp.transpose(q, (0, 2, 1, 3))
k_t = jnp.transpose(k, (0, 2, 1, 3))
v_t = jnp.transpose(v, (0, 2, 1, 3))
# Score: (B, N, T, H) @ (B, N, H, S) to (B, N, T, S)
scores = jnp.matmul(q_t, jnp.transpose(k_t, (0, 1, 3, 2))) * scale
weights = jax.nn.softmax(scores, axis=-1)
# Attend: (B, N, T, S) @ (B, N, S, H) to (B, N, T, H)
out_t = jnp.matmul(weights, v_t)
# Back to (B, T, N, H)
return jnp.transpose(out_t, (0, 2, 1, 3))
naive_out = block_tree(naive_attention(q, k, v))
show_table(
["", "Value"],
[
("Output shape", str(naive_out.shape)),
("Output dtype", str(naive_out.dtype)),
],
title="Naive attention",
)
Вы должны увидеть форму выходных данных (4, 128, 8, 64) и тип данных float32, который совпадает с формой Q, как и было обещано на шаге Выходные данные формулы.
5. Перейти на dot_product_attention
jax.nn.dot_product_attention объединяет этапы оценки, масштабирования, softmax и внимания в одну операцию. Затем JAX и XLA могут оптимизировать шаблон доступа к памяти. В частности, они могут избежать материализации полной матрицы весов внимания, когда последовательность длинная.
По умолчанию используется implementation=None, и JAX автоматически выбирает лучший доступный бэкенд. На графическом процессоре с доступными и совместимыми входными данными cuDNN он уже может использовать cuDNN. На другом оборудовании используется XLA.
sdpa_out = block_tree(jax.nn.dot_product_attention(q, k, v))
max_diff = float(jnp.max(jnp.abs(naive_out - sdpa_out)))
show_table(
["", "Value"],
[
("Output shape", str(sdpa_out.shape)),
("Output dtype", str(sdpa_out.dtype)),
("Max |naive − SDPA|", f"{max_diff:.2e}"),
("Outputs close (atol=1e-3)", str(bool(jnp.allclose(naive_out, sdpa_out, atol=1e-3)))),
],
title="JAX SDPA vs naive",
)
Вы должны увидеть ту же форму и тип данных, что и в наивной версии, небольшую максимальную абсолютную разницу и True для проверки близости atol=1e-3.
6. Принудительно использовать объединенный бэкенд cuDNN
Позволить JAX выбрать серверную часть удобно, но со стороны вы не можете сказать, какое ядро на самом деле работало. Значение implementation="cudnn" принудительно включает ядра cuDNN fused attention от NVIDIA. Это оптимизированные вручную ядра графического процессора, которые объединяют все вычисления внимания в один запуск ядра с оптимизированными шаблонами доступа к памяти.
Обратите внимание, что для cuDNN fused attention есть требования к оборудованию и форме, например вычислительная способность графического процессора должна быть не ниже 8.0 (Ampere или более поздняя версия) или float16 или bfloat16 в качестве входных типов данных. Если требования не соблюдены и вы установили implementation="cudnn", JAX выдаст ошибку, а не просто вернется к предыдущему состоянию. Поэтому в коде ниже выполняется приведение к bfloat16 и вызов оборачивается в try/except: задается HAS_CUDNN_SDPA, чтобы на каждом последующем шаге было известно, доступен ли путь cuDNN на этом компьютере.
q_bf16 = q.astype(jnp.bfloat16)
k_bf16 = k.astype(jnp.bfloat16)
v_bf16 = v.astype(jnp.bfloat16)
HAS_CUDNN_SDPA = False
try:
cudnn_out = block_tree(
jax.nn.dot_product_attention(q_bf16, k_bf16, v_bf16, implementation="cudnn")
)
HAS_CUDNN_SDPA = True
xla_bf16_out = block_tree(
jax.nn.dot_product_attention(q_bf16, k_bf16, v_bf16, implementation="xla")
)
max_diff = float(jnp.max(jnp.abs(
cudnn_out.astype(jnp.float32) - xla_bf16_out.astype(jnp.float32)
)))
show_table(
["", "Value"],
[
("Output shape", str(cudnn_out.shape)),
("Output dtype", str(cudnn_out.dtype)),
("Max |cuDNN − XLA| (both bf16)", f"{max_diff:.2e}"),
("Outputs close (rtol=1e-2, atol=1e-2)", str(bool(jnp.allclose(cudnn_out, xla_bf16_out, rtol=1e-2, atol=1e-2)))),
],
title="cuDNN fused attention",
)
except Exception as e:
print(f"cuDNN SDPA not available on this GPU: {e}")
print("Continuing with XLA backend only.")
Если ваш графический процессор не может запустить cuDNN SDPA, код выведет причину, и практическая работа продолжится на бэкенде XLA.
7. Добавление причинно-следственной маски
В авторегрессионных моделях (декодерах в стиле GPT) каждая позиция может обращать внимание только на более ранние позиции. Если задать для параметра is_causal=True значение True, то эта нижнетреугольная маска будет применяться внутри объединенного ядра, и вам не придется создавать матрицу маски самостоятельно.
Приведенный ниже код запускает внимание дважды, с маской и без нее, и сравнивает две позиции, чтобы показать, что изменила маска.
causal_out = block_tree(
jax.nn.dot_product_attention(q, k, v, is_causal=True)
)
# With causal masking, the last position attends to all positions.
# The first position attends only to itself.
nocausal_out = block_tree(
jax.nn.dot_product_attention(q, k, v, is_causal=False)
)
# First position should differ
first_pos_diff = float(jnp.max(jnp.abs(causal_out[:, 0] - nocausal_out[:, 0])))
# Last position should be the same
last_pos_diff = float(jnp.max(jnp.abs(causal_out[:, -1] - nocausal_out[:, -1])))
show_table(
["Position", "Max diff (causal vs full)", "Expected"],
[
("First (t=0)", f"{first_pos_diff:.4f}", "Large — causal restricts to self only"),
("Last (t=T-1)", f"{last_pos_diff:.2e}", "~0 — attends to all positions either way"),
],
title="Causal masking effect on attention output",
)
Вы должны увидеть большую разницу на первой позиции и почти нулевую на последней. Эта асимметрия заключается в том, что маска с позицией 0 теряет все ключи, кроме своего собственного, в то время как конечная позиция уже может видеть всю последовательность.
8. Выполняю тесты
Теперь у вас есть три способа вычисления одной и той же функции. Выбор зависит от характера проблемы, которую вы можете оценить.
Задайте время для вариантов по умолчанию
def benchmark_attention(fn, q, k, v, warmup=3, repeats=50):
"""Time an attention function. Returns median milliseconds per call."""
jit_fn = jax.jit(fn)
for _ in range(warmup):
block_tree(jit_fn(q, k, v))
times = []
for _ in range(repeats):
start = time.perf_counter()
block_tree(jit_fn(q, k, v))
times.append((time.perf_counter() - start) * 1000)
return np.median(times)
t_naive = benchmark_attention(naive_attention, q, k, v)
t_sdpa = benchmark_attention(
lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
q, k, v,
)
results = [
("Naive (matmul + softmax + matmul)", f"{t_naive:.2f}"),
("SDPA (XLA, float32)", f"{t_sdpa:.2f}"),
]
bar_data = [
("Naive", t_naive),
("SDPA XLA f32", t_sdpa),
]
if HAS_CUDNN_SDPA:
t_sdpa_bf16 = benchmark_attention(
lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
q_bf16, k_bf16, v_bf16,
)
t_cudnn = benchmark_attention(
lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="cudnn"),
q_bf16, k_bf16, v_bf16,
)
results.append(("SDPA (XLA, bfloat16)", f"{t_sdpa_bf16:.2f}"))
results.append(("SDPA (cuDNN, bfloat16)", f"{t_cudnn:.2f}"))
bar_data.append(("SDPA XLA bf16", t_sdpa_bf16))
bar_data.append(("SDPA cuDNN bf16", t_cudnn))
show_table(
["Implementation", "Median ms/call"],
results,
title=f"Attention timing — B={BATCH}, T={SEQ_LEN}, N={NUM_HEADS}, H={HEAD_DIM}",
aligns=["left", "right"],
)
show_bars(bar_data, "Attention latency (ms per call)", "ms", lower_is_better=True)
Для этой небольшой задачи все реализации очень быстры и в основном ограничены накладными расходами, поэтому наивная версия является конкурентоспособной. cuDNN bf16 должен быть немного быстрее, но преимущество объединенного внимания обычно становится более очевидным при большей длине последовательности, когда важно избегать полной матрицы внимания.
Изменение длины последовательности
Преимущество ядер объединенного внимания растет с увеличением длины последовательности. Простая реализация материализует матрицу внимания (B, N, T, S) в памяти графического процессора, которая представляет собой память O(T²). Объединенные ядра, такие как cuDNN FlashAttention, разбивают вычисления на фрагменты, поэтому они никогда не материализуют полную матрицу, сохраняя память на уровне O(T).
Ниже приведены результаты измерения времени выполнения для каждой реализации при длине последовательности от 64 до 1024.
SEQ_LENS = [64, 128, 256, 512, 1024]
sweep_results = []
for sl in SEQ_LENS:
rk = jax.random.key(sl)
rk1, rk2, rk3 = jax.random.split(rk, 3)
q_s = jax.random.normal(rk1, (BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
k_s = jax.random.normal(rk2, (BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
v_s = jax.random.normal(rk3, (BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
q_s, k_s, v_s = jax.device_put((q_s, k_s, v_s), device)
q_sb = q_s.astype(jnp.bfloat16)
k_sb = k_s.astype(jnp.bfloat16)
v_sb = v_s.astype(jnp.bfloat16)
row = {"seq_len": sl}
row["naive_ms"] = benchmark_attention(
naive_attention,
q_s, k_s, v_s,
warmup=2,
repeats=20,
)
row["sdpa_xla_f32_ms"] = benchmark_attention(
lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
q_s, k_s, v_s,
warmup=2,
repeats=20,
)
row["sdpa_xla_bf16_ms"] = benchmark_attention(
lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
q_sb, k_sb, v_sb,
warmup=2,
repeats=20,
)
if HAS_CUDNN_SDPA:
row["cudnn_bf16_ms"] = benchmark_attention(
lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="cudnn"),
q_sb, k_sb, v_sb,
warmup=2,
repeats=20,
)
sweep_results.append(row)
headers = ["Seq len", "Naive (ms)", "SDPA XLA f32 (ms)", "SDPA XLA bf16 (ms)"]
if HAS_CUDNN_SDPA:
headers.append("cuDNN bf16 (ms)")
table_rows = []
for r in sweep_results:
row = [
r["seq_len"],
f"{r['naive_ms']:.2f}",
f"{r['sdpa_xla_f32_ms']:.2f}",
f"{r['sdpa_xla_bf16_ms']:.2f}",
]
if HAS_CUDNN_SDPA:
row.append(f"{r['cudnn_bf16_ms']:.2f}")
table_rows.append(row)
show_table(
headers,
table_rows,
title=f"Sequence-length sweep — B={BATCH}, N={NUM_HEADS}, H={HEAD_DIM}",
aligns=["right"] * len(headers),
)
Проверка занимает некоторое время, поскольку для каждой длины последовательности компилируется три или четыре отдельные реализации, прежде чем они будут синхронизированы. У вас должна получиться одна строка таблицы на длину последовательности от 64 до 1024.
На этом графике сравнивается, как задержка внимания меняется в зависимости от длины последовательности для наивной реализации, JAX/XLA SDPA в float32 и bfloat16, а также cuDNN fused attention в bfloat16.
fig, ax = plt.subplots(figsize=(8, 5))
seq_lens = [r["seq_len"] for r in sweep_results]
ax.plot(
seq_lens,
[r["naive_ms"] for r in sweep_results],
"o-",
label="Naive",
color="#d1242f",
)
ax.plot(
seq_lens,
[r["sdpa_xla_f32_ms"] for r in sweep_results],
"s-",
label="SDPA XLA f32",
color="#0969da",
)
ax.plot(
seq_lens,
[r["sdpa_xla_bf16_ms"] for r in sweep_results],
"d-",
label="SDPA XLA bf16",
color="#8250df",
)
if HAS_CUDNN_SDPA:
ax.plot(
seq_lens,
[r["cudnn_bf16_ms"] for r in sweep_results],
"^-",
label="cuDNN bf16",
color="#1a7f37",
)
ax.set_xlabel("Sequence length")
ax.set_ylabel("Median ms per call")
ax.set_title("Attention latency vs sequence length")
ax.legend()
ax.grid(True, alpha=0.25)
ax.set_xticks(seq_lens)
fig.tight_layout()
plt.show()
Каждой реализации должна соответствовать одна строка. И вы должны увидеть разрыв, который растет по мере того, как последовательность становится длиннее.
Перебор размера пакета
Большие пакеты амортизируют накладные расходы на запуск ядра и улучшают использование графического процессора до точки, когда память графического процессора становится узким местом. В приведенном ниже примере длина последовательности фиксирована на уровне 256, а размер пакета варьируется.
BATCH_SIZES = [1, 2, 4, 8, 16]
SWEEP_SEQ = 256
batch_results = []
for bs in BATCH_SIZES:
rk = jax.random.key(bs + 100)
rk1, rk2, rk3 = jax.random.split(rk, 3)
q_b = jax.random.normal(rk1, (bs, SWEEP_SEQ, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
k_b = jax.random.normal(rk2, (bs, SWEEP_SEQ, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
v_b = jax.random.normal(rk3, (bs, SWEEP_SEQ, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
q_b, k_b, v_b = jax.device_put((q_b, k_b, v_b), device)
q_bb = q_b.astype(jnp.bfloat16)
k_bb = k_b.astype(jnp.bfloat16)
v_bb = v_b.astype(jnp.bfloat16)
row = {"batch": bs}
row["sdpa_xla_f32_ms"] = benchmark_attention(
lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
q_b, k_b, v_b,
warmup=2,
repeats=20,
)
row["sdpa_xla_bf16_ms"] = benchmark_attention(
lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
q_bb, k_bb, v_bb,
warmup=2,
repeats=20,
)
if HAS_CUDNN_SDPA:
row["cudnn_bf16_ms"] = benchmark_attention(
lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="cudnn"),
q_bb, k_bb, v_bb,
warmup=2,
repeats=20,
)
batch_results.append(row)
headers = ["Batch size", "SDPA XLA f32 (ms)", "SDPA XLA bf16 (ms)"]
if HAS_CUDNN_SDPA:
headers.append("cuDNN bf16 (ms)")
table_rows = []
for r in batch_results:
row = [
r["batch"],
f"{r['sdpa_xla_f32_ms']:.2f}",
f"{r['sdpa_xla_bf16_ms']:.2f}",
]
if HAS_CUDNN_SDPA:
row.append(f"{r['cudnn_bf16_ms']:.2f}")
table_rows.append(row)
show_table(
headers,
table_rows,
title=f"Batch-size sweep — T={SWEEP_SEQ}, N={NUM_HEADS}, H={HEAD_DIM}",
aligns=["right"] * len(headers),
)
fig, ax = plt.subplots(figsize=(8, 5))
batches = [r["batch"] for r in batch_results]
ax.plot(
batches,
[r["sdpa_xla_f32_ms"] for r in batch_results],
"s-",
label="SDPA XLA f32",
color="#0969da",
)
ax.plot(
batches,
[r["sdpa_xla_bf16_ms"] for r in batch_results],
"d-",
label="SDPA XLA bf16",
color="#8250df",
)
if HAS_CUDNN_SDPA:
ax.plot(
batches,
[r["cudnn_bf16_ms"] for r in batch_results],
"^-",
label="cuDNN bf16",
color="#1a7f37",
)
ax.set_xlabel("Batch size")
ax.set_ylabel("Median ms per call")
ax.set_title("Attention latency vs batch size")
ax.legend()
ax.grid(True, alpha=0.25)
ax.set_xticks(batches)
fig.tight_layout()
plt.show()
Вы должны получить таблицу и график с одной записью на размер пакета от 1 до 16. Простая реализация исключена из этого анализа, поэтому оба выходных файла сравнивают только три объединенных пути.
9. Сравнение MHA, GQA и MQA
До сих пор у каждого заголовка были свои ключи и значения. Но декодеры часто нарушают эту симметрию, потому что во время инференса кеш ключ-значение доминирует в памяти.
В механизме Multi-Head Attention (MHA) у каждого блока внимания есть собственные проекции Q, K и V. Grouped-query attention (GQA) и multi-query attention (MQA) уменьшают количество голов KV, чтобы экономить память и вычислительные ресурсы во время вывода.
jax.nn.dot_product_attention обрабатывает все три варианта, автоматически транслируя заголовки KV, если K < N.
rk = jax.random.key(42)
rk1, rk2, rk3, rk4, rk5 = jax.random.split(rk, 5)
q_mha = jax.random.normal(rk1, (2, 64, 8, 64), dtype=jnp.float32)
# MHA: 8 KV heads
k_mha = jax.random.normal(rk2, (2, 64, 8, 64), dtype=jnp.float32)
v_mha = jax.random.normal(rk3, (2, 64, 8, 64), dtype=jnp.float32)
# GQA: 2 KV heads (each shared by 4 query heads)
k_gqa = jax.random.normal(rk2, (2, 64, 2, 64), dtype=jnp.float32)
v_gqa = jax.random.normal(rk3, (2, 64, 2, 64), dtype=jnp.float32)
# MQA: 1 KV head (shared by all 8 query heads)
k_mqa = jax.random.normal(rk4, (2, 64, 1, 64), dtype=jnp.float32)
v_mqa = jax.random.normal(rk5, (2, 64, 1, 64), dtype=jnp.float32)
out_mha = block_tree(jax.nn.dot_product_attention(q_mha, k_mha, v_mha))
out_gqa = block_tree(jax.nn.dot_product_attention(q_mha, k_gqa, v_gqa))
out_mqa = block_tree(jax.nn.dot_product_attention(q_mha, k_mqa, v_mqa))
show_table(
["Pattern", "Q shape", "K shape", "V shape", "Output shape"],
[
("MHA", q_mha.shape, k_mha.shape, v_mha.shape, out_mha.shape),
("GQA", q_mha.shape, k_gqa.shape, v_gqa.shape, out_gqa.shape),
("MQA", q_mha.shape, k_mqa.shape, v_mqa.shape, out_mqa.shape),
],
title="Multi-head attention variants — all should produce the same output shape",
)
Все три строки должны сообщать об одной и той же форме вывода, (2, 64, 8, 64), даже если формы K и V уменьшаются с 8 голов до 2 и до 1. В этом и заключается суть: вы можете сократить кеш KV, не меняя ничего после внимания.
10. Тестирование NVIDIA TransformerEngine и FP8
TransformerEngine от NVIDIA предоставляет модули внимания, оптимизированные для графических процессоров NVIDIA. Интеграция JAX использует модули в стиле Flax Linen (не NNX), поэтому модуль инициализируется один раз, а затем применяется со своими переменными.
Оптимизация длины последовательности с помощью TransformerEngine
Эта ячейка проверяет, доступен ли NVIDIA TransformerEngine, а затем сравнивает его bf16 causal DotProductAttention для нескольких длин последовательностей с JAX SDPA с бэкендами XLA и cuDNN на той же рабочей нагрузке.
TE_SEQ_LENS = [128, 256, 512, 1024, 2048]
TE_BATCH = BATCH
HAS_TE = False
try:
import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax
HAS_TE = True
except ImportError:
print("TransformerEngine not installed — skipping TE sections.")
if HAS_TE:
te_results = []
for sl in TE_SEQ_LENS:
rk = jax.random.key(sl + 1000)
rk1, rk2, rk3 = jax.random.split(rk, 3)
q_te = jax.random.normal(
rk1, (TE_BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
)
k_te = jax.random.normal(
rk2, (TE_BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
)
v_te = jax.random.normal(
rk3, (TE_BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
)
q_te, k_te, v_te = jax.device_put((q_te, k_te, v_te), device)
te_attention = te_flax.DotProductAttention(
head_dim=HEAD_DIM,
num_attention_heads=NUM_HEADS,
num_gqa_groups=NUM_HEADS,
attn_mask_type="causal",
transpose_batch_sequence=False,
)
te_vars = te_attention.init(
jax.random.key(0),
q_te,
k_te,
v_te,
deterministic=True,
)
def te_fn(q, k, v):
return te_attention.apply(te_vars, q, k, v, deterministic=True)
row = {"seq_len": sl}
row["sdpa_xla_bf16_ms"] = benchmark_attention(
lambda q, k, v: jax.nn.dot_product_attention(
q, k, v, implementation="xla", is_causal=True
),
q_te, k_te, v_te,
warmup=2,
repeats=20,
)
if HAS_CUDNN_SDPA:
row["sdpa_cudnn_bf16_ms"] = benchmark_attention(
lambda q, k, v: jax.nn.dot_product_attention(
q, k, v, implementation="cudnn", is_causal=True
),
q_te, k_te, v_te,
warmup=2,
repeats=20,
)
row["te_bf16_ms"] = benchmark_attention(
te_fn,
q_te, k_te, v_te,
warmup=2,
repeats=20,
)
te_results.append(row)
headers = ["Seq len", "SDPA XLA bf16 causal (ms)"]
if HAS_CUDNN_SDPA:
headers.append("SDPA cuDNN bf16 causal (ms)")
headers.append("TE DotProductAttention bf16 causal (ms)")
table_rows = []
for r in te_results:
row = [
r["seq_len"],
f"{r['sdpa_xla_bf16_ms']:.2f}",
]
if HAS_CUDNN_SDPA:
row.append(f"{r['sdpa_cudnn_bf16_ms']:.2f}")
row.append(f"{r['te_bf16_ms']:.2f}")
table_rows.append(row)
show_table(
headers,
table_rows,
title=f"TransformerEngine sequence-length sweep — B={TE_BATCH}, N={NUM_HEADS}, H={HEAD_DIM}",
aligns=["right"] * len(headers),
)
fig, ax = plt.subplots(figsize=(8, 5))
seq_lens = [r["seq_len"] for r in te_results]
ax.plot(
seq_lens,
[r["sdpa_xla_bf16_ms"] for r in te_results],
"d-",
label="SDPA XLA bf16 causal",
color="#8250df",
)
if HAS_CUDNN_SDPA:
ax.plot(
seq_lens,
[r["sdpa_cudnn_bf16_ms"] for r in te_results],
"^-",
label="SDPA cuDNN bf16 causal",
color="#1a7f37",
)
ax.plot(
seq_lens,
[r["te_bf16_ms"] for r in te_results],
"o-",
label="TE DotProductAttention bf16 causal",
color="#d1242f",
)
ax.set_xlabel("Sequence length")
ax.set_ylabel("Median ms per call")
ax.set_title("Causal attention latency vs sequence length")
ax.legend()
ax.grid(True, alpha=0.25)
ax.set_xticks(seq_lens)
fig.tight_layout()
plt.show()
Вы должны получить одну строку таблицы и одну точку графика для каждой длины последовательности от 128 до 2048. В этом тесте bf16 сравнивается TransformerEngine с JAX SDPA на одной и той же рабочей нагрузке с причинным вниманием, но полный потенциал производительности TransformerEngine обычно виден на графических процессорах Hopper и Blackwell, когда доступно автоматическое приведение FP8.
Проверьте путь FP8
На графических процессорах Hopper (вычислительная способность >= 9.0, например H100) TransformerEngine может запускать внимание в FP8 для дополнительной пропускной способности. В FP8 используется рецепт DelayedScaling, который отслеживает историю абсолютных максимумов для каждого тензора, чтобы вычислять динамические коэффициенты масштабирования.
- Формат E4M3 для прямого прохода (4 бита экспоненты, 3 бита мантиссы).
- Формат E5M2 для обратного прохода (5 бит экспоненты, 2 бита мантиссы).
Если графический процессор не поддерживает FP8, в этой ячейке показано, как будет выглядеть код без запуска.
if HAS_TE:
from transformer_engine.common.recipe import DelayedScaling, Format
gpu_name = f"{device} {getattr(device, 'device_kind', '')}".lower()
HAS_FP8 = any(
tag in gpu_name
for tag in ["h100", "h200", "b100", "b200", "gb200", "blackwell"]
)
fp8_recipe = DelayedScaling(
margin=0,
fp8_format=Format.HYBRID,
amax_history_len=1024,
amax_compute_algo="max",
)
if HAS_FP8:
FP8_SEQ_LEN = 2048
FP8_BATCH = BATCH
rk = jax.random.key(9000)
rk1, rk2, rk3 = jax.random.split(rk, 3)
q_fp8 = jax.random.normal(
rk1, (FP8_BATCH, FP8_SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
)
k_fp8 = jax.random.normal(
rk2, (FP8_BATCH, FP8_SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
)
v_fp8 = jax.random.normal(
rk3, (FP8_BATCH, FP8_SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
)
q_fp8, k_fp8, v_fp8 = jax.device_put((q_fp8, k_fp8, v_fp8), device)
fp8_attention = te_flax.DotProductAttention(
head_dim=HEAD_DIM,
num_attention_heads=NUM_HEADS,
num_gqa_groups=NUM_HEADS,
attn_mask_type="causal",
transpose_batch_sequence=False,
)
bf16_vars = fp8_attention.init(
jax.random.key(0),
q_fp8,
k_fp8,
v_fp8,
deterministic=True,
)
bf16_out = block_tree(
fp8_attention.apply(
bf16_vars,
q_fp8,
k_fp8,
v_fp8,
deterministic=True,
)
)
with te.autocast(enabled=True, recipe=fp8_recipe):
fp8_vars = fp8_attention.init(
jax.random.key(1),
q_fp8,
k_fp8,
v_fp8,
deterministic=True,
)
fp8_out = block_tree(
fp8_attention.apply(
fp8_vars,
q_fp8,
k_fp8,
v_fp8,
deterministic=True,
)
)
max_diff_fp8 = float(jnp.max(jnp.abs(
bf16_out.astype(jnp.float32) - fp8_out.astype(jnp.float32)
)))
show_table(
["", "Value"],
[
("GPU", getattr(device, "device_kind", str(device))),
("Input dtype", str(q_fp8.dtype)),
("bf16 output dtype", str(bf16_out.dtype)),
("FP8 autocast output dtype", str(fp8_out.dtype)),
("Output shape", str(fp8_out.shape)),
("Max |TE bf16 - TE FP8 autocast|", f"{max_diff_fp8:.2e}"),
],
title="FP8 attention with TransformerEngine",
)
else:
show_table(
["", "Value"],
[
("GPU", getattr(device, "device_kind", str(device))),
("FP8 support", "No detected support; requires Hopper/Blackwell-class GPU"),
],
title="FP8 attention — not available on this GPU",
)
print()
print("The FP8 path uses TransformerEngine autocast:")
print()
print(" with te.autocast(enabled=True, recipe=fp8_recipe):")
print(" out = fp8_attention.apply(vars, q, k, v, deterministic=True)")
else:
print("TransformerEngine not available — FP8 section skipped.")
На L4 вы должны увидеть таблицу с названием вашего графического процессора и сообщением о том, что поддержка FP8 не обнаружена, а затем две напечатанные строки, показывающие вызов te.autocast, который вы бы использовали на графическом процессоре Hopper. Сохраните этот фрагмент: это единственное изменение, которое требуется для FP8 на сайте вызова.
11. Организовать
Удалите рабочую нагрузку Jupyter, включая LoadBalancer и постоянный том:
kubectl delete -f deploy/jupyter.yaml
Удалите кластер, пул узлов, VPC и сервисный аккаунт:
cd terraform
terraform destroy
Введите yes, когда появится запрос, и убедитесь, что ничего не осталось.
gcloud container clusters list
gcloud compute instances list
Для этого проекта оба поля должны быть пустыми. Если вы создали проект только для этой серии, вы можете удалить весь проект из Cloud Console.
12. Поздравляю!
Вы перешли от вычисления внимания вручную к оптимизированным для графического процессора объединенным ядрам и измерили разницу на реальном оборудовании.
Что вы узнали
- Простое внимание работает, но запускает несколько ядер графического процессора и материализует полную матрицу внимания в памяти.
jax.nn.dot_product_attentionобъединяет вычисления в одну операцию, а сimplementation=NoneJAX автоматически выбирает лучший бэкенд.implementation="cudnn"принудительно использует ядра cuDNN fused attention от NVIDIA, которые быстрее всего работают с длинными последовательностями. Требуются входные данныеbfloat16илиfloat16и вычислительная мощность 8.0 или выше.- Причинное маскирование с помощью
is_causal=Trueвстроено в объединенное ядро, поэтому матрица маскирования вручную не требуется. - GQA и MQA сокращают количество голов KV, чтобы экономить память во время инференса, а
dot_product_attentionавтоматически обрабатывает трансляцию. - TransformerEngine обеспечивает объединенное внимание с необязательной точностью FP8 на графических процессорах Hopper.
Дальнейшие действия
- Практическая работа 6. Масштабирование обучения JAX на нескольких графических процессорах. Вы узнаете, как разбивать массивы на фрагменты на графических процессорах L4 и параллельно выполнять шаги обучения.
- Измените
HEAD_DIM(попробуйте 32, 64, 128) и посмотрите, какие значения принимает cuDNN и как меняется время. - Увеличьте
SEQ_LEN(попробуйте 256, 512, 1024, 2048) и посмотрите, как изменится использование памяти и коэффициент ускорения cuDNN. - Повторно сравните MHA, GQA и MQA с
K=2иK=1головками KV и 8 головками запросов и объясните, как это позволяет экономить кеш KV при инференсе.