Что такое TMA (Tensor Memory Accelerator) в H100 и как он ускоряет FlashAttention-3?

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

TMA (Tensor Memory Accelerator) — это асинхронный аппаратный движок в GPU NVIDIA H100, предназначенный для копирования данных между global memory (HBM) и shared memory без участия CUDA cores. В FlashAttention-3 TMA используется для параллельной загрузки блоков матриц K и V во время вычисления attention, что устраняет простои ядер и даёт ускорение до 2x по сравнению с Attention 2|FlashAttention-2. Это снижает требования к occupancy и позволяет эффективнее использовать память.

1. Что такое TMA (Tensor Memory Accelerator)

TMA — это специализированный блок внутри SM (Streaming Multiprocessor) GPU H100, который берёт на себя задачу перемещения данных между уровнями памяти. В отличие от обычного копирования через LD/ST (load/store) инструкции, которые выполняются CUDA-ядрами, TMA работает полностью асинхронно и не занимает вычислительные ресурсы.

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

• Работает на уровне тензоров (многомерных массивов), а не отдельных элементов.
• Поддерживает 2D/3D-блоки данных с автоматическим вычислением адресов.
• Может выполнять copy with padding (копирование с выравниванием) и swizzle (перестановку данных для улучшения кэш-локальности).
• Интегрирован с тензорными ядрами (Tensor Cores) — данные могут быть направлены напрямую в регистры тензорных ядер.

Термин «global memory» — основная память GPU (HBM), большая (~80 GB), но медленная (~1.5 TB/s). «Shared memory» — быстрая память внутри SM (~228 KB на SM), доступная всем потокам блока.

2. Архитектура H100 и место TMA

NVIDIA H100 (Hopper) содержит 132 SM, каждый SM имеет:

• 4 блока тензорных ядер (четвёртого поколения).
• TMA unit — один на SM (или на пару SM, в зависимости от конфигурации).
• Shared memory — 228 KB (настраивается).
• L1 cache — 192 KB.

TMA подключён к Memory Partition Unit (MPU) и может инициировать async copy из global memory в shared memory (и обратно) без участия CUDA-ядер. Это освобождает CUDA cores для вычислений.

Сравнение с предыдущими поколениями

3. Проблема в FlashAttention-2

FlashAttention-2 (FA2) — алгоритм точного attention, который разбивает матрицы Q, K, V на блоки и обрабатывает их в on-chip памяти (shared memory), избегая записи промежуточных матриц в HBM. Однако узким местом остаётся загрузка блоков K и V из HBM в shared memory.

В FA2 каждый шаг:

1. CUDA-ядра загружают блок K и V из global memory в shared memory (через LD/ST).
2. Пока идёт загрузка, тензорные ядра простаивают (или используются для других блоков, но это снижает occupancy).
3. После загрузки вычисляется attention (QK^T, softmax, weighted sum).
4. Результат записывается обратно.

Загрузка занимает значительное время (latency ~100-200 циклов), и CUDA-ядра вынуждены ждать, если нет других блоков для переключения. Это ограничивает occupancy — количество активных блоков на SM.

4. Как TMA решает проблему в FlashAttention-3

FlashAttention-3 (FA3) использует TMA для асинхронной загрузки блоков K и V. Ключевая идея — перекрытие (overlap) загрузки данных с вычислениями.

Pipeline в FA3

1. TMA инициирует копирование следующего блока K/V из global memory в shared memory (асинхронно).
2. Пока TMA копирует, тензорные ядра вычисляют attention для текущего блока (уже загруженного ранее).
3. После завершения вычислений и загрузки следующего блока, данные переключаются (double buffering).
4. TMA копирует следующий блок, а тензорные ядра обрабатывают только что загруженный.

Двойная буферизация (double buffering): в shared memory выделяется два буфера: один используется для вычислений, второй — для загрузки. TMA заполняет второй буфер, пока первый обрабатывается.

Результат CUDA-ядра никогда не ждут данные — latency загрузки полностью скрыта. Это даёт:

• Ускорение ~2x по сравнению с FA2.
• Снижение требований к occupancy (можно запускать меньше блоков, так как TMA не блокирует ядра).
• Возможность использовать более крупные блоки (например, 128x128 вместо 64x64), что улучшает эффективность тензорных ядер.

5. Детали реализации FlashAttention-3 с TMA

В коде FA3 на CUDA используется PTX-инструкция cp.async.bulk (или tma.async.copy в более новых версиях). Пример псевдокода:

// Инициализация TMA дескриптора для 2D-блока
TmaDescriptor desc = tma_create_2d(src_ptr, dst_ptr, block_width, block_height, stride);

// Асинхронное копирование блока K
tma_async_copy(desc, k_buffer, /* sync flag */);

// Пока копируется K, вычисляем attention для предыдущего блока
compute_attention(q, prev_k_buffer, prev_v_buffer);

// Ожидание завершения копирования
tma_wait();

// Переключение буферов
swap(k_buffer, prev_k_buffer);

Важно TMA требует, чтобы данные в global memory были выровнены (обычно 16 байт). Для attention это выполняется, так как матрицы хранятся в формате row-major с padding.

Сравнение с программным __pipeline_memcpy_async (A100):

6. Сравнение производительности

Данные приблизительные, основаны на докладах NVIDIA и бенчмарках FA3.

7. Ограничения и требования

• Только H100 и новее (Hopper и Blackwell). На A100 TMA отсутствует.
• Поддержка в CUDA 11.8+ (PTX cp.async.bulk).
• Выравнивание данных — матрицы должны быть выровнены по 16 байт.
• Размер блока — TMA эффективен для блоков размером от 64x64 элементов (FP16).
• Программирование — требуется низкоуровневый код (PTX или библиотека CUTLASS 3.x). High-level фреймворки (PyTorch, JAX) пока не предоставляют прямого API, но FA3 уже интегрирован в xformers и flash-attn.

8. Влияние на LLM inference и training

TMA в FlashAttention-3 позволяет:

• Увеличить длину контекста — attention на 128K токенов становится в 2 раза быстрее.
• Снизить latency — особенно важно для real-time RAG-систем, где каждый миллисекунд на счету.
• Уменьшить энергопотребление — меньше простоев ядер, эффективнее использование HBM.

Для Agentic RAG (где агенты делают много вызовов LLM с длинными контекстами) ускорение attention напрямую сокращает время ответа.

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

Задача Написать простой CUDA-код, демонстрирующий асинхронное копирование через TMA (симуляция на H100 или эмуляция на A100 через pipeline).

Инструменты CUDA 12.x, компилятор nvcc, профилировщик nsys.

Шаги:

1. Создать две матрицы в global memory (K и V) размером 1024x1024 (FP16).
2. Выделить два буфера в shared memory (double buffering).
3. Реализовать цикл по блокам:
   • Использовать cp.async.bulk для копирования блока K/V в буфер B.
   • Пока копируется, обработать буфер A (например, просто сложить элементы).
   • После завершения копирования поменять буферы местами.
4. Сравнить время выполнения с версией без TMA (обычное копирование через memcpy в ядре).

Ожидаемый результат Версия с TMA должна быть быстрее на ~30-50% (зависит от размера блока). Профилировщик покажет, что TMA-копирование перекрывается с вычислениями.

Примечание Для реального TMA нужен H100. На A100 можно использовать __pipeline_memcpy_async для имитации, но эффект будет слабее.

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

Навигация

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