MoonMath.ai выложила в открытый доступ HIP-ядро для AMD MI300X, которое обходит родной AITER v3
MoonMath.ai опубликовала bf16-ядро прямого внимания (forward attention) для ускорителя AMD MI300X. Код написан на HIP, а не на ассемблере, и доступен под лицензией MIT. По утверждению разработчиков, ядро обходит оптимизированное решение AMD AITER v3 на каждой протестированной форме матриц и при всех режимах округления. Средний геометрический выигрыш составил 1,18x, 1,15x и 1,08x в зависимости от сценария, максимальный — 1,26x. Доступ к «голому железу» предоставил облачный провайдер AMD HotAisle.
Основной трюк — однокомандные обёртки на ассемблере: разработчики помещают ровно одну инструкцию в __device__ __forceinline__ функцию, а компилятору оставляют выделение регистров и отслеживание потока данных. За счёт связывания входа и выхода с тем же VGPR (+v) не генерируется лишняя копирующая инструкция. Большая часть ускорения — за счёт размещения данных: K через HBM в LDS с двойной буферизацией, V горячая остаётся в L1, Q и аккумуляторы — в регистрах. Выбранный размер тайла 16×16×16 для MFMA (вместо 32×32×8) даёт пропускную способность, равную альтернативе, при меньшем тайле.
Архитектура: восемь волн (wavefronts) на блок, разделённых на две группы по четыре. Группы выполняют последовательность Q*K, softmax, O += P*V со смещением по фазе — пока одна группа насыщает матричное ядро, другая занимается softmax и загрузками. Синхронизация — через два s_barrier на итерацию. Это напоминает FlashAttention-3, но без выделенных producer-волн: на CDNA3 все перемещения памяти уже асинхронны. Ядро работает только с bf16 на gfx942, фиксированная головная размерность 128, без causal-маски и GQA, но поддерживает любую длину последовательности (включая cross-attention). Численные результаты детерминированы, каждый конечный выход не выходит за 1 bf16 ULP от AITER; NaN и Inf обрабатываются побитово идентично.
Практическая проверка: реальный PR для SGLang уже использует это ядро и ускоряет генерацию видео Wan2.1 в 1,23 раза без потери качества.