DeepGEMM - Эффективная библиотека умножения матриц FP8

Четкая и эффективная библиотека умножения матриц FP8 с тонкой настройкой масштабирования, оптимизированная для архитектуры NVIDIA Hopper

DeepGEMM Matrix Multiplication Visualization

Архитектурный дизайн DeepGEMM

DeepGEMM имеет тщательно разработанную архитектуру, оптимизированную для тензорных ядер NVIDIA Hopper, обеспечивающую эффективное умножение матриц FP8.

DeepGEMM Architecture Diagram

Полностью JIT-дизайн

DeepGEMM использует полностью Just-In-Time (JIT) дизайн без компиляции во время установки. Все ядра компилируются во время выполнения с использованием легковесного JIT-модуля, обрабатывая формы GEMM, размеры блоков и этапы конвейера как константы времени компиляции, экономя регистры и позволяя дополнительные оптимизации компилятора.

Функции Hopper TMA

Полностью использует возможности Tensor Memory Accelerator (TMA) архитектуры Hopper для более быстрого и асинхронного перемещения данных, включая загрузки TMA, сохранения, трансляции и предварительную выборку дескрипторов, значительно повышая производительность.

Основные методы оптимизации

DeepGEMM использует несколько продвинутых методов оптимизации, выходящих за рамки традиционных библиотек GEMM для достижения исключительной производительности.

DeepGEMM Optimization Techniques

Специализация постоянных потоков

Следуя дизайну CUTLASS, перекрывает перемещение данных, инструкции MMA тензорного ядра и продвижение ядра CUDA для улучшения вычислительной эффективности.

Неровные размеры блоков

Поддерживает неровные размеры блоков (например, 112) для лучшего использования SM. Например, при M=256, N=7168, использование BLOCK_M=128, BLOCK_N=112 позволяет работать 128 SM вместо 112 с выровненными размерами блоков.

Переплетение FFMA SASS

Модифицирует инструкции FFMA в скомпилированном бинарном файле, переворачивая биты yield и reuse, создавая больше возможностей для перекрытия инструкций MMA с инструкциями продвижения FFMA, что дает +10% производительности в некоторых случаях.

Унифицированный оптимизированный планировщик блоков

Один планировщик для всех негруппированных и группированных ядер, использующий растеризацию для улучшения повторного использования кэша L2, повышая общую производительность.

Объяснение кода

Основная функция ядра DeepGEMM состоит всего из примерно 300 строк кода, с чистым дизайном, который облегчает изучение умножения матриц Hopper FP8 и методов оптимизации.

// Simplified version of the core GEMM kernel
template <
    int BLOCK_M,
    int BLOCK_N,
    int BLOCK_K,
    int CLUSTER_M,
    int CLUSTER_N,
    int STAGES,
    int WARPGROUP_SIZE,
    bool ENABLE_INTERLEAVE
>
__global__ void gemm_fp8_fp8_bf16_nt_kernel(
    const void* __restrict__ lhs,
    const void* __restrict__ lhs_scale,
    const void* __restrict__ rhs,
    void* __restrict__ output,
    const void* __restrict__ rhs_scale,
    int m, int n, int k,
    int lhs_stride, int rhs_stride, int output_stride) {
    
    // Using TMA to load LHS, LHS scale factors, and RHS
    // Using CUDA cores for two-level accumulation (promotion)
    // Using TMA to store output
    
    // Persistent thread specialization implementation
    // Unaligned block size support
    // FFMA SASS interleaving optimization
}

Особенности структуры кода

  • Шаблонный дизайн, поддерживающий различные размеры блоков и конфигурации оптимизации
  • Полное использование функций TMA архитектуры Hopper для эффективного перемещения данных
  • Двухуровневое накопление для решения проблемы неточного накопления тензорных ядер FP8
  • Реализация специализации постоянных потоков для оптимизированных вычислений и перекрытия перемещения данных
DeepGEMM Code Flow Diagram

Дизайн интерфейса

  • Обычный плотный GEMM: Использование функции deep_gemm.gemm_fp8_fp8_bf16_nt
  • Группированный GEMM (Непрерывный макет): Использование функции m_grouped_gemm_fp8_fp8_bf16_nt_contiguous, подходящей для моделей MoE, где эксперты имеют одинаковую форму
  • Группированный GEMM (Маскированный макет): Использование функции m_grouped_gemm_fp8_fp8_bf16_nt_masked, используемой в фазе декодирования вывода с включенными графами CUDA, где CPU не знает количество токенов, получаемых каждым экспертом

Производительность

Производительность DeepGEMM сопоставима или лучше, чем у библиотек, настроенных экспертами, для различных форм матриц, с ускорением до 2.7x в некоторых случаях.

DeepGEMM Performance Comparison

Основные показатели производительности

  • Обычный GEMM: До 2.7x ускорения на определенных формах
  • Группированный GEMM с непрерывным макетом: До 1.2x ускорения
  • Группированный GEMM с маскированным макетом: Также до 1.2x ускорения
  • Вычислительная производительность до 1358 TFLOPS, протестированная на H800 SXM5
  • Высокое использование пропускной способности памяти, достигающее 2668 ГБ/с

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

Исключительная производительность

DeepGEMM обеспечивает ускорение до 2.7x по сравнению с библиотеками, настроенными экспертами, на архитектуре NVIDIA Hopper, достигая более 1350 TFLOPS

Продвинутые оптимизации

DeepGEMM использует специализацию постоянных потоков, неровные размеры блоков и переплетение FFMA SASS для максимальной вычислительной эффективности

Гибкая интеграция

DeepGEMM легко интегрируется с фреймворками глубокого обучения и библиотеками научных вычислений через четкие, хорошо документированные интерфейсы

Тонкое масштабирование

DeepGEMM эффективно масштабируется от малых матриц до крупномасштабных вычислений с оптимизированным планированием блоков и методами растеризации

Нативная поддержка FP8

DeepGEMM обеспечивает первоклассную поддержку точности FP8, оптимизированную специально для тензорных ядер NVIDIA Hopper

Высокая пропускная способность памяти

DeepGEMM достигает исключительного использования пропускной способности памяти до 2668 ГБ/с, максимизируя возможности оборудования

Чистый дизайн кода

DeepGEMM имеет лаконичное ядро из 300 строк с четкими шаблонами проектирования, что облегчает понимание и расширение

Оптимизированный группированный GEMM

DeepGEMM превосходно работает как с непрерывными, так и с маскированными операциями группированного GEMM, идеально подходя для моделей Mixture of Experts (MoE)

Часто задаваемые вопросы

Что такое DeepGEMM и чем он особенный?
DeepGEMM - это эффективная библиотека умножения матриц FP8, оптимизированная для архитектуры NVIDIA Hopper. Особенным его делает чистый дизайн (всего ~300 строк основного кода), исключительная производительность (до 2.7x ускорения по сравнению с библиотеками, настроенными экспертами) и продвинутые методы оптимизации, такие как специализация постоянных потоков и переплетение FFMA SASS.
Для какого оборудования оптимизирован DeepGEMM?
DeepGEMM специально оптимизирован для архитектуры NVIDIA Hopper, полностью используя возможности Tensor Memory Accelerator (TMA) Hopper и тензорные ядра для умножения матриц FP8. Он был протестирован на GPU H800 SXM5, достигая до 1358 TFLOPS вычислительной производительности.
Как DeepGEMM достигает своих преимуществ в производительности?
DeepGEMM достигает своей производительности через несколько продвинутых техник: полностью JIT-дизайн с константами времени компиляции, неровные размеры блоков для лучшего использования SM, переплетение FFMA SASS, двухуровневое накопление для решения проблем точности FP8, и унифицированный оптимизированный планировщик блоков с растеризацией для улучшения повторного использования кэша L2.
Какие типы операций GEMM поддерживает DeepGEMM?
DeepGEMM поддерживает три основных типа операций GEMM: обычный плотный GEMM (используя функцию deep_gemm.gemm_fp8_fp8_bf16_nt), группированный GEMM с непрерывным макетом и группированный GEMM с маскированным макетом. Последние два особенно полезны для моделей Mixture of Experts (MoE).
Подходит ли DeepGEMM для моделей Mixture of Experts (MoE)?
Да, DeepGEMM отлично подходит для моделей MoE. Он предоставляет специализированные операции группированного GEMM как с непрерывным, так и с маскированным макетом, которые оптимизированы для специфических вычислительных паттернов в архитектурах MoE. Эти операции могут обеспечить ускорение до 1.2x по сравнению с другими библиотеками.
Какую точность поддерживает DeepGEMM?
DeepGEMM в основном фокусируется на точности FP8, которая становится все более важной для эффективного вывода AI-моделей. Он реализует двухуровневый подход к накоплению для решения проблем неточного накопления, которые могут возникать с тензорными ядрами FP8, обеспечивая как производительность, так и точность.
Как использование пропускной способности памяти DeepGEMM сравнивается с другими библиотеками?
DeepGEMM достигает исключительного использования пропускной способности памяти, достигающего 2668 ГБ/с. Это возможно благодаря эффективному использованию функций TMA Hopper для более быстрого и асинхронного перемещения данных, включая загрузки TMA, сохранения, трансляции и предварительную выборку дескрипторов.
Могу ли я учиться на реализации DeepGEMM для оптимизации своего собственного кода CUDA?
Абсолютно! Основное ядро DeepGEMM состоит всего из примерно 300 строк кода с чистым дизайном, что делает его отличным учебным ресурсом. Оно демонстрирует продвинутые методы оптимизации CUDA, эффективное использование функций архитектуры Hopper и умные подходы к умножению матриц, которые могут быть применены к другим задачам высокопроизводительных вычислений.