Новое HIP‑ядро MoonMath ускорит AMD MI300X и превзойдёт AITER v3

TL;DR

MoonMath.ai опубликовал открытый bf16‑ядро forward‑attention для AMD MI300X, написанное на HIP (MIT‑лицензия). По результатам тестов оно превышает производительность официального ядра AITER v3 на всех проверенных формах и режимах округления (среднее ускорение ≈ 1.15×‑1.18×, максимум ≈ 1.26×).

Что представляет собой ядро

Ядро — это небольшая программа, исполняемая непосредственно на всех вычислительных блоках GPU. Оно реализует операцию attention: softmax(QKᵀ/√d)·V в полубайтовом формате bf16. Поддерживаются входы в раскладках BSHD и BHSD, фиксированное количество голов = 128, произвольные длины последовательностей (включая cross‑attention). Ограничения: отсутствие causal‑mask, GQA и varlen‑batching; только вывод bf16; работает исключительно на архитектуре gfx942 (CDNA3).

Как достигнут прирост

Основная идея — обёртки asm‑инструкций, содержащие ровно одну инструкцию. Такая функция объявлена как __device__ __forceinline__, а ограничения компилятора привязывают входные и выходные регистры к одному VGPR:

• Код остаётся читаемым на уровне HIP, но разработчики выбирают точный opcode.
• Компилятор всё‑равно управляет распределением регистров, устраняя лишние копии.

Архитектурные особенности

Для CDNA3 выбран блок из восьми волн, разделённых на две группы по четыре волны. Обе группы одновременно выполняют цикл Q·K → softmax → O += P·V, но находятся в разных фазах:

• Пока группа A загружает K в LDS и работает в матричном ядре, группа B выполняет softmax и подготавливает следующий набор данных.
• Два барьера s_barrier синхронизируют переключения фаз и границы итераций.

Размещение данных в памяти

• K‑тайл помещён в LDS (double‑buffered, 32 KiB) и делится между всеми восемью волнами.
• V‑тайл хранится «горячим» в L1‑кэше, что позволяет многократно переиспользовать его без обращения к более медленной памяти.
• Q‑тайл и аккумуляторы находятся в регистрах VGPR, что исключает любые лишние загрузки.
• Выбран MFMA‑тиль 16×16×16 bf16 (вместо 32×32×8) — одинаковая пропускная способность, но меньшее давление на регистры и лучшая энергоэффективность.

Бенчмарки

Тесты проведены на MI300X в bf16, head‑dimension = 128, измерения для трёх режимов округления (RTNE, RTNA, RTZ). Пример:

• Shape = (2, 24, 8192, 128), RTNE: 3.083 мс (MoonMath) vs 3.792 мс (AITER) → 1.23× быстрее.
• Shape = (2, 24, 32768, 128), RTNA: 44.44 мс vs 52.36 мс → 1.18× быстрее.
• Геометрическое среднее ускорения: 1.18× (RTNE), 1.15× (RTNA), 1.08× (RTZ).

Практический эффект

В реальном проекте (Pull‑Request в SGLang) ядро ускорило видеодиффузию Wan2.1 на 1.23× без потери качества. Это демонстрирует, что полученный прирост не ограничивается синтетическими тестами, а полезен в продуктивных нагрузках.

Как получить код

Исходники находятся в открытом репозитории GitHub под лицензией MIT: github.com/moonmath-ai/cdna3attention. В репозитории также присутствует README с инструкциями по сборке, примерами использования и скриптами для воспроизведения бенчмарков.

Краткое резюме

• Ядро написано на HIP, а не на hand‑written assembly, что упрощает поддержку.
• ТехнИка одно‑инструкционных asm‑обёрток позволяет точно выбирать opcode без потери оптимизаций компилятора.
• Продуманный пайплайн из двух групп волн обеспечивает 100 % загрузку матричного ядра.
• Оптимальное размещение K, V, Q и аккумуляторов обеспечивает до 1.26× ускорения над AITER v3.
• Код полностью открыт, лицензирован MIT и готов к использованию в собственных проектах на MI300X.