AMD показала, как выжать 85% пиковой пропускной способности HBM из нового ускорителя Instinct MI450 в задаче attention decode. Для этого написали и оптимизировали ядро на Gluon — DSL над Triton, где каждый тензор обязан иметь явную раскладку по регистрам, SIMD-линиям и workgroup’ам.
Ускоритель MI450 принес сразу несколько новых железок. Главная — TDM (Tensor Data Mover). Он умеет асинхронно перебрасывать структурированные тайлы тензоров между HBM и LDS, не заставляя ядро нарезать данные сотнями маленьких векторных загрузок. Появились кластеры workgroup: несколько вычислительных блоков теперь могут синхронизироваться через кластерные барьеры и расшаривать данные multicast-загрузками. Вдобавок на треть выросли VGPR (до 1024 на SIMD32), удвоился LDS (до 320 КБ на WGP), а пропускная способность HBM подскочила с 8,1 до 19,6 ТБ/с.
Оптимизацию вели в четыре шага. Сначала вручную подобрали WMMA-лейауты. В attention decode используется MQA/GQA, поэтому ядро обрабатывает один KV-заголовок и группу Q-заголовков тайлом (BLOCK_M, D). Чтобы WMMA-инструкции v_wmma_scale_f32_16x16x128_f8f6f4 работали без лишних пересылок, выход умножения QK транспонируют прямо в операции (через параметр transpose=True в AMDWMMALayout). Это подает результат сразу на вход умножения PV без дорогого преобразования через LDS. Волны распределяют по сгруппированному Q-измерению, исключая межволновое общение при online-softmax.
Второй шаг — хардовая загрузка K и V через TDM. Чтобы TDM пошёл прямым путем в LDS в обход кэша, внутреннее измерение тайла должно быть не менее 128 байт (рекомендуют 256). Исходная размерность головы 128 (FP8) этого не даёт, поэтому K и V решейпят, загружают широким тайлом, а потом восстанавливают исходную форму при чтении из LDS в регистры.
Третий шаг — конвейеризация с тройной буферизацией. Простой двухстадийный конвейер (загрузка отдельно, вычисления отдельно) оставлял ~680 тактов видимой задержки. Четырёхстадийный разбивает softmax на два блока (VEC0 и VEC1) и чередует их с WMMA QK и PV. Тройная буферизация LDS позволяет выставить TDM-запросы на две итерации вперёд, снижая непокрытую задержку примерно до 360 тактов.
Четвёртый шаг — Split-k через кластеры. Чтобы загрузить все 256 WGP при маленьком батче, KV-последовательность делят на S частей и раздают отдельным workgroup’ам. Дальше частичные результаты сливаются прямо в том же ядре: workgroup’ы синхронизируются через кластер (CGA), обмениваясь данными через общий L2-кэш без лишних записей в глобальную память и без отдельного kernel launch.
Итоговое ядро на FP8 при batch=64, 64 Q-заголовках, одном или двух KV-заголовках и длинах последовательности от 4096 до 65536 токенов показало эффективную полосу 17,10 ТБ/с (85 % от пика в 20 ТБ/с) для GQA и 16,65 ТБ/с (83 %) для MQA. Замеры сняты на ROCm 7.14.0, PyTorch 2.11.0 и Triton (коммит ecfc626), код выложен в репозиторий MI450 Gluon MXFP Attention Example. AMD планирует применить те же техники к другим типам данных, paged KV cache и вариантам внимания вроде MLA/DSA из семейства DeepSeek.