Как работает Tensor Core microarchitecture (WGMMA, MMA инструкции) в H100?

Краткий тезис

Tensor Cores — это специализированные вычислительные блоки внутри GPU, оптимизированные для выполнения операции D = A × B + C, где A, B, C, D — матрицы малых размеров. В H100 (четвёртое поколение Tensor Cores) ключевыми инструкциями являются MMA (Matrix Multiply-Accumulate) для работы с warp-уровнем и WGMMA (Warp Group MMA) для масштабирования на warp-группы. Эти инструкции обеспечивают пиковую производительность до 1979 TFLOPS в режиме FP8 sparse и являются основой для ускорения всех GEMM-операций в LLM, включая attention и feed-forward слои.

1. Что такое Tensor Cores и зачем они нужны?

Tensor Cores — это аппаратные блоки, впервые представленные в архитектуре NVIDIA Volta (V100). Они предназначены для выполнения матричного умножения с накоплением за один такт. В отличие от обычных CUDA-ядер, которые выполняют скалярные операции (одно сложение/умножение за раз), Tensor Cores работают с малыми матрицами (например, 4×4, 8×8, 16×16) и выполняют всю операцию D = A×B + C за один шаг.

Зачем это нужно Основная нагрузка в глубоком обучении — это операции GEMM (General Matrix Multiply), которые составляют >90% вычислений в LLM. Без Tensor Cores эти операции выполнялись бы на обычных ядрах с гораздо меньшей пропускной способностью.

Эволюция поколений

2. Архитектура H100: потоковый мультипроцессор (SM) и warp

H100 содержит 132 SM (Streaming Multiprocessor). Каждый SM включает:

• 4 Tensor Core блока (каждый способен выполнять 16×16×16 MMA за такт)
• 64 CUDA cores (обычные ядра)
• 4 warp schedulers (каждый управляет двумя warps)

Warp — это группа из 32 потоков, которые выполняют одну и ту же инструкцию (SIMT-модель). В контексте Tensor Cores:

• MMA работает на уровне одного warp (32 потока).
• WGMMA работает на уровне warp group — объединения нескольких warps (обычно 4 warps = 128 потоков).

Warp group — это новая концепция в H100, позволяющая координировать несколько warps для работы над одной большой матрицей без синхронизации на уровне блоков.

3. Инструкция MMA (Matrix Multiply-Accumulate)

MMA — это инструкция PTX (Parallel Thread Execution), которая выполняет D = A×B + C для малых матриц. В H100 поддерживаются размеры:

• m16n16k16: A (16×16), B (16×16), C/D (16×16) — базовый размер.
• m16n8k16: A (16×16), B (16×8), C/D (16×8) — для несимметричных операций.
• m8n8k32: A (8×8), B (8×8), C/D (8×8) — для меньших матриц.

Как это работает на уровне warp:

1. Каждый из 32 потоков warp'а загружает фрагменты матриц A и B из shared memory или registers.
2. Tensor Core выполняет умножение и накопление за один такт.
3. Результат сохраняется в регистры или shared memory.

Пример кода на CUDA (псевдокод):

// Объявление фрагментов матриц
nvcuda::wmma::fragment<nvcuda::wmma::matrix_a, 16, 16, 16, half, nvcuda::wmma::row_major> a_frag;
nvcuda::wmma::fragment<nvcuda::wmma::matrix_b, 16, 16, 16, half, nvcuda::wmma::col_major> b_frag;
nvcuda::wmma::fragment<nvcuda::wmma::accumulator, 16, 16, 16, float> c_frag;

// Загрузка данных из shared memory
nvcuda::wmma::load_matrix_sync(a_frag, shmem_a, 16);
nvcuda::wmma::load_matrix_sync(b_frag, shmem_b, 16);

// Выполнение MMA
nvcuda::wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);

// Сохранение результата
nvcuda::wmma::store_matrix_sync(shmem_c, c_frag, 16, nvcuda::wmma::mem_row_major);

Ограничения MMA

• Работает только внутри одного warp (32 потока).
• Требует ручного управления shared memory и синхронизации.
• Максимальный размер одной операции — 16×16×16.

4. Инструкция WGMMA (Warp Group MMA)

WGMMA — это новая инструкция в H100, которая расширяет MMA на warp group (обычно 4 warps = 128 потоков). Она позволяет выполнять операцию D = A×B + C для матриц размером 64×64×16 за одну инструкцию.

Ключевые особенности

• Warp-level parallelism: 4 warpa работают параллельно над разными частями матрицы.
• Автоматическое распределение данных: WGMMA сама распределяет фрагменты матриц между warps без явной синхронизации.
• Поддержка асинхронности: WGMMA может выполняться асинхронно, позволяя перекрывать вычисления с загрузкой данных.

Размеры, поддерживаемые WGMMA:

• m64n64k16: A (64×64), B (64×64), C/D (64×64) — базовый размер.
• m64n16k16: A (64×64), B (64×16), C/D (64×16) — для узких матриц.
• m16n64k16: A (16×64), B (64×64), C/D (16×64) — для высоких матриц.

Как это работает

1. Warp group (4 warps) загружает матрицу A размером 64×64 в shared memory.
2. Каждый warp отвечает за 16 строк матрицы A (16×64).
3. WGMMA выполняет умножение A (64×64) на B (64×16) за один такт.
4. Результат (64×16) распределяется между warps.

Пример псевдокода (PTX):

// WGMMA для m64n64k16
wgmma.mma_async.sync.aligned.m64n64k16.f16.f16.f32
    { %r0, %r1, ... },  // фрагменты результата (16 регистров на warp)
    [%ptrA],             // указатель на матрицу A в shared memory
    [%ptrB],             // указатель на матрицу B в shared memory
    %descA, %descB;      // дескрипторы (формат, stride)

Преимущества WGMMA перед MMA

5. Числовые форматы и производительность

Tensor Cores в H100 поддерживают несколько форматов данных:

• FP32: для аккумулятора (C и D), но умножение может быть в FP16.
• FP16: базовый формат для A и B.
• BF16: аналогичен FP16, но с большим диапазоном.
• FP8 (E4M3 и E5M2): новый формат для H100, позволяет удвоить производительность.
• INT8: для инференса с квантизацией.
• Sparse: поддержка разреженных матриц (2:4 sparsity) — удваивает производительность.

Пиковая производительность H100 (на один SM):

Примечание Реальная производительность зависит от размера матриц и эффективности использования памяти.

6. Роль Tensor Cores в LLM

Все основные операции в LLM (Large Language Models) являются GEMM и выполняются через Tensor Cores:

1. Attention: QK^T (Query-Key multiplication):

   • Размер: (batch_size × seq_len × d_model) × (d_model × seq_len)
   • Использует WGMMA для больших batch-размеров.

2. Attention: PV (Probability-Value multiplication):

   • Размер: (batch_size × seq_len × seq_len) × (seq_len × d_model)
   • Аналогично, WGMMA.

3. Feed-Forward Network (FFN):

   • Два последовательных GEMM: (batch_size × seq_len × d_model) × (d_model × 4d_model) и обратно.
   • Доминирующая часть вычислений (≈2/3 FLOPs).

4. Embedding lookup:

   • Хотя это не GEMM, но часто реализуется через матричное умножение.

Пример: для модели LLaMA-70B

• Каждый слой содержит 4 GEMM (Q, K, V, O) + 2 GEMM (FFN).
• При batch_size=1, seq_len=4096, d_model=8192:
  • QK^T: 8192×4096 × 4096×8192 → 2^31 FLOPs.
  • На H100 (989 TFLOPS FP16) это занимает ~2 мкс.

7. Программная модель: как вызывать Tensor Cores

На практике разработчики редко пишут PTX-инструкции напрямую. Вместо этого используются библиотеки:

1. cuBLAS — библиотека BLAS от NVIDIA, автоматически выбирает Tensor Cores для GEMM.

   cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, CUDA_R_16F, lda, B, CUDA_R_16F, ldb, &beta, C, CUDA_R_32F, ldc, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP);
   

2. cuDNN — для свёрточных и рекуррентных сетей.

3. CUTLASS — библиотека с открытым исходным кодом для реализации кастомных GEMM с Tensor Cores.

4. Triton — язык программирования от OpenAI, который автоматически генерирует код для Tensor Cores.

Пример на Triton

import triton
import triton.language as tl

@triton.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
):
    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    pid_m = pid % num_pid_m
    pid_n = pid // num_pid_m

    offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k = tl.arange(0, BLOCK_SIZE_K)

    a_ptrs = a_ptr + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
    b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k in range(0, K, BLOCK_SIZE_K):
        a = tl.load(a_ptrs)
        b = tl.load(b_ptrs)
        accumulator += tl.dot(a, b)
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk

    c = accumulator.to(tl.float16)
    c_ptrs = c_ptr + offs_am[:, None] * stride_cm + offs_bn[None, :] * stride_cn
    tl.store(c_ptrs, c)

Triton автоматически сгенерирует WGMMA-инструкции для H100, если размеры блоков кратны 64.

8. Ограничения и подводные камни

1. Размер матриц: WGMMA эффективен только для больших матриц (M, N ≥ 64). Для малых (например, в batch_size=1) MMA может быть быстрее из-за меньших накладных расходов.

2. Shared memory: WGMMA требует больше shared memory (до 228 KB на SM в H100). Если модель использует много shared memory (например, FlashAttention), может возникнуть конкуренция.

3. Warp divergence: Все потоки в warp должны выполнять одну и ту же инструкцию. Если в коде есть ветвления, Tensor Cores не используются.

4. Precision: FP8 требует осторожной квантизации (calibration). Неправильный scaling может привести к loss of accuracy.

5. Sparse: 2:4 sparsity (каждый второй элемент — ноль) удваивает производительность, но требует специальной структуры матриц.

Пет-проект для закрепления

Задача Написать микро-бенчмарк для сравнения производительности MMA и WGMMA на H100.

Инструменты

• Python 3.10+
• PyTorch 2.0+ (с поддержкой CUDA 12.0)
• NVIDIA H100 (или A100 для MMA-only)
• Triton (опционально)

Шаги:

1. Создайте матрицы разных размеров (16×16, 32×32, 64×64, 128×128, 256×256) с FP16.
2. Реализуйте GEMM через PyTorch (использует cuBLAS, который автоматически выбирает Tensor Cores).
3. Реализуйте кастомный kernel на Triton с размерами блоков 16×16 (MMA) и 64×64 (WGMMA).
4. Замерьте время выполнения для каждого размера с помощью torch.cuda.Event.
5. Постройте график зависимости GFLOPS от размера матрицы.

Ожидаемый результат

• Для малых матриц (≤32) MMA и WGMMA показывают схожую производительность.
• Для больших матриц (≥128) WGMMA даёт прирост в 2-4 раза.
• График должен показать плато на уровне ~900 TFLOPS для FP16.

Код для замера (PyTorch):

import torch
import time

def benchmark_gemm(M, N, K, dtype=torch.float16, device='cuda'):
    A = torch.randn(M, K, dtype=dtype, device=device)
    B = torch.randn(K, N, dtype=dtype, device=device)
    C = torch.zeros(M, N, dtype=torch.float32, device=device)

    # Warmup
    for _ in range(10):
        torch.mm(A, B, out=C)

    # Measure
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    start_event.record()
    for _ in range(100):
        torch.mm(A, B, out=C)
    end_event.record()
    torch.cuda.synchronize()

    elapsed_ms = start_event.elapsed_time(end_event) / 100
    gflops = (2 * M * N * K) / (elapsed_ms * 1e6)
    return gflops

sizes = [16, 32, 64, 128, 256, 512]
for s in sizes:
    gflops = benchmark_gemm(s, s, s)
    print(f"Size {s}x{s}: {gflops:.2f} GFLOPS")

Связь с другими вопросами

Навигация

• Предыдущий: 704
• Следующий: 706
• Индекс: 00. Индекс разборов