Открытое ядро MoonMath обгоняет AMD Instinct MI300X AITER на 18% в среднем

До сих пор владельцы серверов на AMD Instinct MI300X пользовались ядром AITER v3, которое разрабатывает сама AMD. MoonMath AI предложила альтернативу, написанную на HIP (язык программирования GPU от AMD), а не на ручном ассемблере. Это упрощает чтение и доработку кода. Голое железо для тестов предоставил облачный провайдер HotAisle, специализирующийся на AMD.

Ядро внимания (attention kernel) выполняет ключевую математическую операцию внутри каждого трансформера, архитектуры, на которой построены ChatGPT, YandexGPT и большинство современных нейросетей. Чем быстрее эта операция, тем быстрее модель отвечает пользователю и тем дешевле обходится каждый запрос.

Что быстрее и на сколько?

По замерам MoonMath AI, ядро обгоняет AITER v3 от AMD на каждой протестированной форме тензора и в каждом режиме округления:

• Среднее геометрическое ускорение (RTNE, режим округления к ближайшему чётному): 1,18 раза
• RTNA (округление к ближайшему, при ничьей от нуля): 1,15 раза
• RTZ (усечение к нулю, самый быстрый режим самого AITER): 1,08 раза
• Максимальное ускорение на отдельной конфигурации: до 1,26 раза

В сравнении с ядром Modular MAX средний выигрыш ещё заметнее: от 1,44 до 1,49 раза, а на отдельных формах до 1,59 раза.

Команда уже отправила патч в проект SGLang для генерации видео. На модели Wan2.1-T2V-1.3B-Diffusers замена AITER на ядро MoonMath ускорила сквозную генерацию видео в 1,23 раза без потери качества картинки.

Как MoonMath добилась ускорения?

Основной выигрыш, по описанию команды, даёт не алгоритм, а размещение данных в памяти GPU. Матрица K (ключи) загружается в локальную разделяемую память (LDS) с двойной буферизацией. Матрица V (значения) остаётся «горячей» в кэше первого уровня. Запросы Q и аккумуляторы живут в регистрах, самой быстрой памяти чипа.

Второй приём: вместо стандартных четырёх волн на вычислительный блок MoonMath запускает восемь, разбитых на две группы по четыре. Пока одна группа загружает матричное ядро, вторая считает softmax и подгружает данные, затем они меняются. Матричный блок не простаивает.

Третий приём, однокомандные asm-обёртки: разработчик выбирает конкретную инструкцию GPU, но компилятор по-прежнему сам распределяет регистры. Код остаётся читаемым HIP-кодом, а не ручным ассемблером.

Ядро работает не на любом железе

У решения есть чёткие ограничения, которые команда указывает явно:

• Работает только на AMD Instinct MI300X (архитектура CDNA3, цель gfx942)
• Только формат bf16, размерность головы фиксирована на 128
• Нет каузальной маски (нужна для авторегрессивной генерации текста)
• Нет GQA (Grouped Query Attention, группированное внимание, которое используют многие современные модели для экономии памяти)
• Нет поддержки varlen batching (пакетирование последовательностей разной длины)

Это значит, что ядро сейчас подходит для задач, где маска не нужна: генерация видео, диффузионные модели, кросс-внимание с произвольной длиной KV.

Как попробовать?

1. Убедитесь, что у вас есть доступ к серверу с AMD Instinct MI300X (облачный провайдер HotAisle или другой поставщик AMD GPU)
2. Установите ядро через pip (PyTorch должен быть собран с поддержкой ROCm)
3. Импортируйте moonmath_attention, передайте тензоры Q, K, V и укажите раскладку bshd или bhsd

Код открыт на GitHub под лицензией MIT: можно форкнуть, изменить и использовать в коммерческом продукте без ограничений.

Есть ли аналог для российского пользователя?

Прямого российского аналога этому ядру нет: это низкоуровневый компонент для серверных GPU, а не пользовательский продукт. Но контекст для аудитории важен.

YandexGPT и GigaChat работают на собственной инфраструктуре, и конечный пользователь не выбирает ядро внимания. Новость актуальна для тех, кто сам разворачивает модели: часть российских и СНГ-компаний арендует серверы с AMD Instinct MI300X в зарубежных облаках, потому что они дешевле и доступнее аналогов от NVIDIA в условиях ограничений на поставки.

Что делать с этим прямо сейчас, по ролям?

Автору Дзена. Напрямую ядро вас не затрагивает. Но если вы пользуетесь сервисами генерации видео по тексту, ускорение диффузионных моделей в 1,23 раза означает, что обложки и видеовставки будут генерироваться быстрее и дешевле, когда провайдеры подтянут обновление.

Маркетологу и предпринимателю. Если ваша команда разворачивает собственные модели на AMD GPU, проверьте, подходит ли ваш сценарий (bf16, размерность 128, без каузальной маски). Лицензия MIT не ограничивает коммерческое использование. Экономия на инференсе при среднем ускорении 1,15 раза ощутима на масштабе тысяч запросов в час.

Техническому руководителю в РФ и СНГ. Открытый код на HIP, а не на ассемблере, значит, что ядро можно адаптировать под свои задачи без реверс-инжиниринга. Следите за добавлением каузальной маски и GQA: после этого ядро станет пригодным для текстовых LLM.

Частые вопросы

Можно ли запустить это ядро на обычной видеокарте AMD?

Нет. Ядро работает исключительно на AMD Instinct MI300X, серверном ускорителе архитектуры CDNA3. Потребительские карты Radeon используют другую архитектуру (RDNA) и не поддерживаются.

Заменяет ли это FlashAttention?

Частично. Ядро решает ту же задачу, быстрое вычисление внимания, но пока только для прямого прохода (forward) в формате bf16 и без каузальной маски. Для полноценной замены FlashAttention в текстовых моделях нужны обратный проход, каузальная маска и GQA.

Нужно ли платить за использование?

Нет. Код распространяется под лицензией MIT: можно использовать, изменять и встраивать в коммерческие продукты бесплатно и без ограничений.

Открытый код под MIT и ускорение, подтверждённое хотя бы на одном реальном пайплайне генерации видео, делают это ядро первым кандидатом на тест для всех, кто уже работает с AMD Instinct MI300X и хочет выжать из железа больше без смены платформы.