На AMD Instinct MI355X целый запрос в vLLM стал быстрее в 3,01 раза на MXFP4. Модель та же, железо то же, конфигурация сервинга та же, поменялся только список кандидатов, из которых TorchInductor выбирает матричные ядра. Новым кандидатом стал FlyDSL, бэкенд, который команда PyTorch добавила в стек для GPU AMD.

За цифрой 3,01x прячется более интересная история: PyTorch перестал быть фреймворком с одним способом писать ядра под вендора. Теперь это площадка, где несколько DSL соревнуются за каждый конкретный размер матрицы, и побеждает тот, кто на этом размере оказался быстрее.

Что такое FlyDSL

FlyDSL это язык для GPU-ядер с явной раскладкой данных. Автор ядра сам разбивает тайл по блокам, волнам, потокам и значениям, выбирает атомы MFMA и tiled-copy, задаёт swizzle и staging для локальной памяти LDS, расставляет границы синхронизации. Ядра пишутся на Python и понижаются через MLIR, но при этом сохраняют явную привязку к железу.

Работает это на GPU серии MI350 с архитектурой gfx950, через необязательный runtime пакета flydsl версии 0.3.x.

Почему Triton-модели не хватило

Triton предлагает блочно-тензорную модель: автор описывает тайлы уровня программы, а компилятор сам решает, как разложить их по потокам, волнам и инструкциям. Это удобно и почти всегда достаточно. Проблема начинается там, где перемещение данных становится частью алгоритма.

На CDNA4 часть выигрыша зависит от того, кто именно владеет операндом в конкретный момент, какой банк LDS занят, где стоит барьер между двумя группами волн. Компилятор верхнего уровня такие решения принимает по общим правилам, а на конкретной архитектуре оптимальный шаблон может выглядеть иначе. FlyDSL даёт эти ручки напрямую, и в этом смысле он не конкурент Triton, а второй инструмент для другой задачи: архитектурно-специфичные шаблоны, где раскладка и есть алгоритм.

Именно поэтому в ядрах FlyDSL видны вещи, которых в Triton-коде не бывает. В dense GEMM это чередование полного и половинного тайла (HTI) и ступенчатый конвейер LDS. В MXFP это согласованное время жизни операндов и их скейлов вместе со scaled MFMA. В grouped GEMM это планировщик, который держит экспертов в одном постоянном гриде.

Три пути, которыми FlyDSL попадает в модель

Первый путь, eager. Вызовы F.rms_norm и torch.topk уходят в FlyDSL автоматически, если установлен и включён необязательный runtime, а форма, тип и layout подходят. Всё остальное остаётся на ATen.

Второй путь, torch.compile. TorchInductor фильтрует подходящие кандидаты среди dense, grouped и MXFP GEMM, замеряет FlyDSL рядом с ATen и Triton на конкретной задаче, кеширует победителя и дальше запускает самое быстрое из измеренного. Первый вызов включает компиляцию и автотюнинг, последующие переиспользуют выбор.

Третий путь, внимание. Здесь FlyDSL не подключается со стороны PyTorch: AOTriton компилирует эти ядра в свою библиотеку и сам выбирает между FlyDSL, Triton и AITER ASM по офлайновой базе тюнинга. Пользователь F.scaled_dot_product_attention ничего не устанавливает и не настраивает, ему нужна только сборка PyTorch, слинкованная с AOTriton 0.14.2b или новее.

Общая рамка устроена аккуратно: FlyDSL нигде не подменяет поведение по умолчанию. Неподдерживаемые входы молча уходят в прежнюю реализацию, поэтому включение бэкенда не может сломать корректность, только изменить скорость.

Dense GEMM: 1,10x и честный размен по формам

Плотное матричное умножение это проекции и feed-forward слои трансформера, и его размеры гуляют вместе с числом токенов в батче. На 15 формах BF16 в раскладке NT FlyDSL даёт 1,10x геометрического среднего против более быстрого из пары ATen и Triton на каждой отдельной форме.

Разбивка важнее среднего: 10 побед, три ничьи и два поражения в полосе ±1%. Выигрыш концентрируется на малых и средних задачах, например 64 × 4096 × 4096 даёт 1,33x. Два самых больших случая с квадратным выходом слегка выигрывают у ATen, и это ровно тот случай, ради которого кандидатов держат несколько: PyTorch просто выберет ATen.

Механика рекорда в HTI. Для больших GEMM тюнинг выбирает выходной тайл 256 × 256 и рабочую группу из восьми волн (512 потоков при Wave64). HTI делит A по M и B по N, держа в регистрах четыре аккумулятора 128 × 128: A0 кормит C00 и C01, A1 кормит C10 и C11, B0 отвечает за C00 и C10, B1 за C01 и C11.

Обработка идёт парами последовательных K-тайлов через два шага LDS. В установившемся режиме prefetch переиспользует шаги, где лежали t и t+1, под t+2 и t+3, причём переиспользование идёт половинами: как только чтения старой половины операнда закончены и барьеры пройдены, её область можно заполнять, пока MFMA считает по регистровым фрагментам. Дальше в дело входит сдвиг фаз: две группы по четыре волны работают со смещением, так что в одной идёт счёт, а в другой загрузка и асинхронный prefetch. Никто не ждёт, пока весь тайл досчитается.

Отдельная проверка текущего ядра со всеми раскладками на тех же формах дала 1,12x агрегата, уже без разбивки по формам. Сами формы NN, TN и TT тоже поддерживаются, но их ускорение зависит от нагрузки.

MXFP: 1,58x на FP8 и 1,68x на FP4

Низкая точность это место, где AMD-стек исторически отставал сильнее всего, и там же выигрыш самый заметный. На 17 формах NT в MXFP8 FlyDSL даёт 1,58x геометрического среднего к ATen, в MXFP4 1,68x. По отдельным формам разброс шире: MXFP8 от 1,31x до 2,23x с максимумом на 32 × 4096 × 4096, MXFP4 от 1,31x до 2,96x с пиком на 32 × 14336 × 4096. Ни одного проигрыша ATen в этих прогонах нет.

Технически MXFP надстраивается над теми же реализациями full-tile и HTI, добавляя обработку FP8, упакованное хранение FP4, staging скейлов E8M0 и scaled MFMA архитектуры CDNA4. Схема данных: операнды FP8 E4M3 или упакованные FP4 E2M1, один скейл E8M0 на 32 значения вдоль оси редукции (BlockWise1x32), накопление в FP32, выход в FP16 или BF16.

Тонкость в буферизации скейлов. В HTI буферы A и B сдвигаются на K-тайл, а отдельные слоты скейлов держат куски, перекрывающие несколько K-тайлов: каждый MFMA берёт регистровые операнды и соответствующий его позиции K скейл. На границе куска ожидания и барьеры согласуют переиспользование второго слота, чтобы префетчить следующий кусок, пока текущий ещё расходуется. Длина куска зависит от геометрии тайла и рабочей группы, четыре тайла на кусок в примере из статьи это иллюстрация, а не требование формата.

Ограничения тоже конкретные: K должен быть кратен 128, bias допускается только одномерный длиной N и в типе выхода, быстрое накопление не поддерживается. Под torch.compile ограничение на layout срабатывает раньше выбора бэкенда и требует единичный шаг вдоль K для A и B, то есть row-major A и column-major B. Несовместимые входы получают вставку копии, которая съедает часть выигрыша.

Grouped GEMM: MoE с неравномерной нагрузкой

В mixture-of-experts эксперты получают разное число токенов, иногда ноль. PyTorch представляет такие вычисления как конкатенированные активации A[sum(M_g), K], сгруппированные веса B[G, K, N] и смещения, отмечающие границу строк каждого эксперта.

FlyDSL не запускает по ядру на эксперта. Вместо этого постоянный грид выдаёт каждой рабочей группе последовательность тайлов с шагом по кумулятивным смещениям групп, так что пустые эксперты пропускаются естественно, а рабочая группа переходит через границы экспертов без новых запусков. Раскладка учитывает кристаллы: swizzle распределяет работу по восьми Accelerator Complex Dies, а N-major fallback держит одновременные группы на переиспользуемых B-тайлах.

На 24 случаях BF16 результат такой: 1,20x геометрического среднего к Triton и 1,95x к ATen. Однородные группы из 14 случаев дают 1,21x и 2,12x, причём FlyDSL быстрее всех в 11 из 14. Пять случаев с проекционными размерностями дают 1,23x и 1,73x, пять случаев с рваной загрузкой экспертов, включая пустых и сильно перекошенных, дают 1,15x и 1,73x, и там FlyDSL выигрывает все пять.

Чувствительность к форме видна в деталях: в проекционном свипе есть случай с малым N, где Triton быстрее, и широкий N, где результат равен ATen. Рваная нагрузка, ради которой grouped GEMM и нужен, не подводит ни в одном случае.

RMSNorm и TopK: 3,66x и 4,82x на узких местах

RMSNorm это повторяющаяся нормализация активаций в трансформере. FlyDSL ускоряет forward, оставляя привычный API и прежнее поведение backward: обратный проход по-прежнему считается реализацией PyTorch.

Ядро отдаёт один блок потоков на строку и выбирает 256, 512 или 1024 потока по размеру нормализуемой размерности, что соответствует 4, 8 или 16 Wave64. Векторизуемое тело грузится 128-битными копиями, остаток добирается скалярным хвостом, вход остаётся в регистрах. Сумма квадратов считается в FP32: сначала внутри каждой волны через shuffle, потом по одной частичной сумме на волну через небольшой буфер LDS. Регистровые значения переиспользуются для rsqrt, веса и для rstd, который нужен backward.

На выровненных размерностях выигрыш скромный, примерно 1,2x–1,5x. На размерности на единицу выше выровненной, например 4097 вместо 4096, он доходит до 3,66x, потому что выравненный путь перестаёт быть выравненным у конкурента. Это хороший пример того, как пиковое число в таблице живёт в узкой нише.

TopK устроен как две разные машины под два режима. Для малого фиксированного K блок из двух Wave64 обрабатывает две строки, по волне на строку: каждая дорожка читает 128-битные куски, битонически сортирует группу за группой и складывает результат в регистровый локальный top-K, после чего butterfly-shuffle сливает кандидатов по волне и нулевая дорожка пишет финальные пары. Для большого K работают четыре 8-битных radix-прохода: они находят порог K-го элемента, собирают только выживших и сортируют буфер, округлённый от K до ближайшей степени двойки, вместо полной строки.

Цифры: на малых K из {2, 4, 8, 16} геометрическое среднее 4,82x к ATen при выключенной детерминированности и 4,02x при включённой. На больших полосах K среднее держится в диапазоне 1,40x–1,97x, отдельные случаи почти равны ATen.

Есть важная оговорка про воспроизводимость. При включённом torch.use_deterministic_algorithms(True) radix-путь использует слоты с префиксными суммами и сохраняет порядок равных значений, при выключенном берёт атомарную выдачу слотов, и индексы равных значений могут меняться. Регистровое ядро воспроизводимо в обоих режимах, но порядок равных элементов может отличаться от ATen. NaN канонизируются в один ординал, поэтому точный payload и выбранный индекс тоже могут отличаться даже в детерминированном режиме.

Attention: два тайла в полёте и обход лимита LDS

Внимание в FlyDSL устроено интереснее остальных операторов, потому что здесь пришлось решать физическую задачу. При размерности головы до 256 forward использует двухволновой software pipeline: два KV-тайла в работе на рабочую группу, softmax одного тайла считается на VALU, пока другой тайл гоняет QK и PV на MFMA. Внутри волны каждый шаг матричных ядер перекрывается softmax соседнего KV-тайла, а по рабочей группе две группы волн идут со сдвигом в одну фазу: одна грузит K и V, вторая считает.

На размерности головы 512 два резидентных K/V-тайла потребовали бы около 272 КБ LDS при лимите 160 КБ. Поэтому появляется отдельное тело для широкой головы, которое обрабатывает один тайл за раз и делит размерность дважды. По времени: K и V стейджатся в проходящие через те же буферы LDS слайсы D. По волнам: каждая волна накапливает свой срез выхода, чтобы аккумулятор поместился в регистрах.

Backward не изобретает новые карты дорожек, а переиспользует блоки forward. dK/dV это forward-цикл в транспонированном виде: K и V остаются в регистрах, Q и dO идут через LDS, а транспонированные операнды используют тот же чтение-с-транспонированием ds_read_b64_tr_b16. В dQ каждое из трёх умножений совпадает с тем, что forward уже выдаёт, а K-тайл стейджится один раз и читается в двух ориентациях, что и удерживает размерность 512 в пределах LDS.

Публичные числа AOTriton 0.14.2b на gfx950 против 0.13b: forward от 0.92x до 2.34x, backward от 0.93x до 2.47x в зависимости от размерности головы и причинно-следственного режима. Авторы честно предупреждают, что это сравнение релизов, а не изолированный замер FlyDSL против Triton: в 0.14.2b обновились ещё Triton, AITER и база тюнинга.

Ограничение тоже названо прямо: FP32 не поддерживается. Входы FP32 и нагрузки, которые база тюнинга отдаёт другим бэкендам, продолжают идти через Triton или AITER ASM.

Что даёт end-to-end: цифры vLLM

Операторные ускорения легко переоценить, поэтому в статье есть замер на уровне целого запроса. Стенд: vllm bench serve на одной MI355X с TP=1, модели Llama-3.1-8B-Instruct, Qwen3-32B и Llama-3.3-70B-Instruct, вход 256 токенов, фиксированный выход 512 токенов, параллелизм от 8 до 256. Базовая и опытная конфигурации отличаются только списком GEMM-кандидатов в Inductor.

В BF16 картина скромная и это ожидаемо: Llama-3.1-8B даёт до 1,13x на низком и среднем параллелизме, Qwen3-32B доходит до 1,06x, а Llama-3.3-70B остаётся в пределах измеренного шума 1,79%. Числа это медианы минимум двух ABBA-чередованных повторов.

В MXFP8 A8W8 базой становится ATen (hipBLASLt), и улучшаются все 18 комбинаций модель-параллелизм: 1,40x–2,03x у Qwen3-32B, 1,04x–1,92x у Llama-3.1-8B, 1,38x–2,04x у Llama-3.3-70B.

В MXFP4 A4W4, с неовернутым lm_head, улучшаются тоже все 18: 1,53x–3,01x у Qwen3-32B, 1,25x–1,85x у Llama-3.1-8B, 1,50x–2,82x у Llama-3.3-70B. Здесь важна методическая оговорка: в MXFP8 и MXFP4 на каждую ячейку приходится один прогон, поэтому мелкие различия между ячейками не стоит трактовать как проверенную дисперсию. Общий вывод от этого не меняется: основной выигрыш лежит в низкой точности, а не в BF16.

Как включить FlyDSL в PyTorch

Нужен свежий ROCm-ночной PyTorch для карт серии MI350 и необязательный runtime:

python -m pip install --upgrade "flydsl>=0.3.0,<0.4"

Дальше FlyDSL добавляется как кандидат, а выбор остаётся за Inductor:

import torch
import torch._inductor.config as config

config.max_autotune_gemm_backends = "ATEN,TRITON,FLYDSL"
config.flydsl_enable_autotuning = True

a = torch.randn(128, 4096, device="cuda", dtype=torch.bfloat16)
b = torch.randn(4096, 14336, device="cuda", dtype=torch.bfloat16)

@torch.compile(mode="max-autotune")
def mm(a, b):
    return a @ b

out = mm(a, b)

На ROCm устройства по-прежнему называются cuda, это не опечатка. Включение FlyDSL не гарантирует его выбор: ATen или Triton могут оказаться быстрее на конкретной форме, и тогда Inductor возьмёт их. Если нагрузка несовместима с захватом графов, вместо max-autotune берут max-autotune-no-cudagraphs. Сами параметры конфигурации версионно-зависимы и могут меняться между ночными сборками.

Для RMSNorm и TopK достаточно, чтобы runtime был установлен и включён: подходящие eager-вызовы уходят в FlyDSL сами, а сравнить с ATen можно контекстным менеджером pn.flydsl.disabled(). Для внимания не нужно ничего, кроме сборки PyTorch с AOTriton 0.14.2b.

Ограничения и планы

Текущая поддержка привязана к gfx950, то есть к MI350. В планах расширение интеграции и набора ядер на MI450, MXFP8 grouped GEMM, ядра FlexAttention, перенос attention-ядер из вендорной копии в собственный каталог kernels/attention/ внутри FlyDSL, расписания внимания с учётом длины последовательности, fusion эпилога GEMM, более широкий backward, поддержка AOTInductor и расширенные end-to-end бенчмарки.

Из текущих ограничений, которые стоит держать в голове: RMSNorm backward остаётся в PyTorch, FP32 во внимании не поддерживается, MXFP требует K кратный 128 и канонический layout под torch.compile, а flex_attention это отдельный API, который Inductor компилирует в свои Triton-ядра и через AOTriton не проводит.

Что это значит для рынка AMD-инференса

За последний год PyTorch собрал вокруг себя целый набор вендорских путей: TorchTPU для тензорных ядер Google, TorchNeuron для Trainium, нативное устройство Spyre от IBM, Monarch для распределённого обучения. FlyDSL продолжает ту же линию, но с другой стороны: не «вендор приносит свой бэкенд целиком», а «вендор приносит шаблоны ядер, а фреймворк замеряет и выбирает».

Для команды, которая обслуживает инференс, это меняет экономику. Ускорение приходит без переписывания модели и без форка фреймворка: достаточно обновить ночную сборку, поставить пакет и переключить список кандидатов. Обратная сторона тоже очевидна: выигрыш появляется только там, где поддерживаемые операции занимают заметную долю времени, то есть в LLM-инференсе с плотными слоями, вниманием, нормализацией и top-k. Экзотическая архитектура с кастомными слоями останется на прежних ядрах.

И третий эффект, менее очевидный. Когда фреймворк сам замеряет несколько бэкендов на каждой форме, вендорские команды начинают конкурировать не презентациями, а числом побед в автотюнинге. Проигрыш на двух самых больших квадратных GEMM в статье описан так же подробно, как и победы, и это хороший признак: инфраструктура замеров стала достаточно дешёвой, чтобы публиковать неудобные результаты.

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

FlyDSL заменит Triton в PyTorch?

Нет. Оба остаются кандидатами в автотюнинге, и на части форм Triton по-прежнему быстрее, например в одном случае проекционного свипа grouped GEMM и на двух крупных квадратных GEMM. FlyDSL дополняет стек для архитектурно-специфичных шаблонов, где раскладка данных часть алгоритма.

Нужно ли переписывать модель, чтобы получить ускорение?

Нет. Для dense и grouped GEMM достаточно включить бэкенд в конфигурации Inductor, для RMSNorm и TopK хватает установленного runtime, для F.scaled_dot_product_attention не нужно вообще ничего, кроме сборки PyTorch с AOTriton 0.14.2b. Неподдерживаемые входы автоматически идут по прежнему пути.

Почему ускорение в 3 раза видно только на MXFP4?

Потому что там слабее базовая линия. В BF16 сравнение идёт с ATen и Triton, которые на gfx950 уже хорошо оттюнены, и выигрыш держится в районе 1,13x на целом запросе. В MXFP4 базой выступает hipBLASLt, а ядро использует scaled MFMA и упакованное хранение FP4, отсюда 1,53x–3,01x в зависимости от модели и параллелизма.

Работает ли это на других GPU AMD?

Пока нет. Поддержка заявлена для gfx950, то есть MI350, а MI450 в планах. Отдельное исключение: AOTriton поставляет attention-ядра FlyDSL и для RDNA4 (gfx1201), но это только внимание, не GEMM и не нормализации.

Итог

FlyDSL это не «ещё один способ писать ядра», а изменение правил выбора: PyTorch получил право замерять несколько реализаций на каждой конкретной форме и брать победителя. Для AMD это первый случай, когда оптимизированные под CDNA4 ядра приезжают в обычный torch.compile и в F.scaled_dot_product_attention без форков и без вендорских патчей.

Если у вас есть MI350 и нагрузка на LLM-инференс, самый дешёвый способ проверить тезис на своём железе: поставить flydsl из ветки 0.3.x, добавить FLYDSL в max_autotune_gemm_backends и сравнить время запроса до и после на своём батче. Расхождение с цифрами из статьи будет полезнее любого бенчмарка: оно покажет, где именно ваша модель упирается не в матричные ядра.