Что такое DeepGEMM и зачем она нужна
DeepSeek представила библиотеку DeepGEMM — унифицированное решение для высокопроизводительных вычислений на графических процессорах NVIDIA. В отличие от тяжелых библиотек вроде CUTLASS, DeepGEMM создана для простоты и понятности, что делает её отличным инструментом для изучения оптимизации CUDA-ядер. Ключевая особенность — компиляция ядровых функций в рантайме через DeepJIT, что исключает необходимость ручной сборки CUDA-кода при установке.
Библиотека объединяет ключевые примитивы современных больших языковых моделей (LLM): GEMM в форматах FP8, FP4, BF16, а также специализированные ядра для MoE-архитектур (Mega MoE) и индексации (MQA scoring).
Производительность и поддержка железа
DeepGEMM демонстрирует результаты, сопоставимые или превосходящие экспертно-настроенные библиотеки. На графическом процессоре NVIDIA H800 (архитектура SM90) зафиксирована пиковая производительность до 1550 TFLOPS. Библиотека поддерживает архитектуры SM90 (Hopper) и SM100 (Blackwell), а также, согласно последним обновлениям, доступна версия DeepGEMM Ascend для чипов Huawei.
| Характеристика | SM90 (Hopper/H800) | SM100 (Blackwell) |
|---|---|---|
| Поддерживаемые форматы | FP8, FP4, BF16 | FP8, FP4, BF16 |
| Макс. производительность | До 1550 TFLOPS | Информация недостаточна |
| FP32 | Packed UE8M0 (4 значения в torch.int) | |
| Поддержка layout матриц | Только NT (Non-transposed A, Transposed B) | NT, TN, NN, TT (все варианты) |
Ключевые функции: Mega MoE и индексация
Особое внимание уделено оптимизации Mixture-of-Experts (MoE) моделей. Функция Mega MoE объединяет в одно ядро этапы диспетчеризации (EP dispatch), линейных слоев (Linear 1 & 2), активации SwiGLU и комбинирования результатов (EP combine). Это позволяет эффективно перекрывать коммуникацию по NVLink с вычислениями на тензорных ядрах.
Также добавлены ядра для индексатора DeepSeek v3.2 (MQA scoring), работающие с форматами FP8/FP4. Эти ядра поддерживают как prefilling (неpaged), так и decoding (paged) фазы генерации, что критично для скорости инференса.
Требования и установка
Для работы требуются GPU с архитектурой SM90 или SM100, Python 3.8+, CUDA Toolkit 12.9+ и PyTorch 2.3+. Библиотека использует C++20 и требует наличия CUTLASS 4.0 (подтягивается как submodule). Установка сводится к клонированию репозитория и запуску скрипта ./install.sh, после чего модуль deep_gemm доступен для импорта в Python-проектах.
Источник: Github ↗