ZYFA: Быстрое и точное FP8-внимание на базе AMD MI355X

Основные результаты

• Скорость на реальных активациях: На активациях, полученных из Mistral-7B, Qwen3-8B и DeepSeek-V2-Lite при длительности 16K токенов и измеренных на постоянных частотах, ядра FP8 от ZYFA достигают 1.16–1.76× от пропускной способности ядер по умолчанию от AITER, а ядра BF16 от ZYFA достигают 1.01–1.02×. Прирост скорости в FP8 меньше, чем на гауссовских входных данных, поскольку тактовая частота MI355X зависит от данных. AITER позволяет большему количеству вероятностей уменьшаться до нуля (underflow), а нули переключают меньше битов, поэтому ядро AITER работает на более высокой частоте.
• Скорость на бенчмарках ядер: На гауссовских входных данных при 32K токенах ядра FP8 от ZYFA достигают 1.26–1.37× от пропускной способности ядер AITER при размерности головы 128 и 1.99× для формы головы многоголового латентного внимания (MLA) от DeepSeek. При размерности головы 128 ядро FP8 достигает 53% от пикового значения по спецификации FP8.
• Точность ядер: При размерности головы 128 ошибка ядер FP8 от ZYFA находится в пределах нескольких процентов от ошибки ядер Composable Kernel (CK) и Triton от AMD и на 62–78% ниже, чем у AITER.
• Точность полной модели: Большая часть ошибки FP8 в полной модели возникает из-за квантования запросов, ключей и значений. На Mistral-7B вращение Адамара для Q и K сокращает эту ошибку вдвое и примерно вдвое уменьшает дивергенцию Кульбака — Лейблера (KL) по сравнению с моделью BF16, практически без дополнительных затрат времени.

Мы описываем нашу методологию и ее ограничения в конце этого поста. Короче говоря, все результаты получены с одного узла MI355X, ядра охватывают только причинное предварительное заполнение (causal prefill), а как наши ядра, так и базовые решения быстро развиваются, поэтому все результаты актуальны на момент времени.

О чем этот пост

• Ограничения FP8-внимания: стоимость весового конвейера для softmax, порядок выполнения матричных инструкций в программном обеспечении, а также точность и потеря точности (underflow) вероятностей FP8.
• Дизайн ядер: компактная матрица P (lean P), двухволновое расписание типа «пинг-понг» и поблочное масштабирование с нормализацией по округленной сумме строк.
• Пропускная способность и числовая точность: захваченные активации, гауссовские входные данные, зависимость тактовой частоты от значений операндов и ошибка FP8 в полных моделях.
• Оставшиеся пробелы в производительности и рекомендации, будущая работа, а также методология и ограничения.
• Приложение А: таблицы со значениями, отображенными на графиках.

Введение

Предварительное заполнение (prefill) — прямой проход по промпту перед тем, как модель сгенерирует свой первый токен — тратит все большую часть своего времени на вычисление внимания по мере роста контекста. FlashAttention хранит матрицу скоров в памяти на кристалле (on-chip), а не в памяти с высокой пропускной способностью (HBM), поэтому объем данных, которые внимание перемещает в HBM и обратно, растет линейно с длиной последовательности, в то время как объем арифметических операций растет квадратично. Таким образом, внимание при предварительном заполнении в условиях длинного контекста ограничено производительностью вычислений (compute-bound), и его скорость зависит от того, насколько хорошо ядро загружает матричные модули графического процессора.

MI355X выполняет 8-битные матричные инструкции с плавающей запятой (FP8) вдвое быстрее пиковой скорости BF16, поэтому FP8-внимание в принципе могло бы занимать вдвое меньше времени, чем BF16-внимание. На практике выигрыш намного меньше, и ядро FP8 от AITER достигает меньшей доли своего потенциального пика, чем его ядро BF16. FP8 также поднимает вопрос точности, поскольку внимание является очень чувствительным к числам блоком. DeepSeek-V3, например, обучает свои линейные слои в FP8, но сохраняет внимание в более высокой точности.

ZYFA состоит из ядер причинного прямого внимания в форматах BF16 и FP8 для MI355X (CDNA4, gfx950). Ядра поддерживают размерность головы 128 и форму головы MLA в DeepSeek-V2 с размерностью запрос-ключ 192 и размерностью значения 128. В разделе «Ограничения FP8-внимания на MI355X» описаны три препятствия, ограничивающие FP8-внимание, а в разделе «Дизайн ядер» описана методика в ядрах FP8 от ZYFA, которая решает каждую из этих проблем. Мы приводим полный анализ в нашем техническом отчете.

Ограничения FP8-внимания на MI355X

Стоимость весового конвейера для Softmax

FP8 сокращает вдвое количество циклов для двух матричных произведений внимания, QKᵀ и PV, которые выполняются на матричном конвейере. Softmax между ними выполняется на весовом конвейере, который выполняет поэлементную арифметику с одинаковой скоростью в FP8 и BF16, а ядро FP8 также должно преобразовывать каждую вероятность в 8 бит. Оба конвейера выполняют инструкции волнами (waves) — группами по 64 потока, которые выдают инструкции вместе. Для одной волны, вычисляющей 32 запроса для блока из 128 ключей, весовому конвейеру требуется 0.44× от циклов матричного конвейера в BF16, но 0.94× в FP8 (Рисунок 2). Таким образом, ядро FP8 загружает матричный конвейер только в том случае, если softmax перекрывается с матричными произведениями практически идеально.

Порядок выполнения матричных инструкций в программном обеспечении

На MI355X матричная инструкция (MFMA) продолжает выполняться, пока волна, которая ее выдала, переходит к другим инструкциям, и аппаратное обеспечение не отслеживает момент завершения MFMA. Графические процессоры Hopper и Blackwell от NVIDIA отслеживают завершение асинхронных матричных инструкций на аппаратном уровне. На MI355X только порядок самих инструкций должен перекрывать softmax с инструкциями MFMA и предотвращать раннее чтение результатов MFMA любой инструкцией. В трассировках потоков ядра FP8 от AITER матричный конвейер простаивает 31.7% циклов во время выполнения softmax.

Точность и потеря точности (underflow) вероятностей FP8

Формат FP8 для вероятностей, e4m3, сохраняет все 3 бита мантиссы только для значений не менее 2⁻⁶ (нормальных значений). Меньшие значения теряют точность, а значения ниже 2⁻¹⁰ падают до нуля (underflow). Большинство вероятностей длинной строки лежат далеко ниже наибольшей из них. В слое 20 Mistral-7B при 16K токенов наибольшая вероятность блока из 64 ключей лежит более чем на 8 бинад (степеней двойки) ниже максимума строки для 91% пар строки и блока. Потеря точности также искажает выходные данные, если сумма строки суммирует неоткругленные вероятности, как в ядре FP8 от AITER, поскольку каждая вероятность, которая уходит в ноль, все равно учитывается в сумме строки.

Дизайн ядер

Ядра FP8 от ZYFA решают каждую проблему с помощью определенной техники.

Компактная матрица P (Lean P): линейная аппроксимация кода FP8

Стандартное ядро FP8 вычисляет каждую вероятность p = 2ˣ, где x — это скор минус максимум строки в двоичных единицах, с помощью экспоненциальной инструкции, а затем округляет ее до e4m3. Мы называем это точным P (exact P), и это занимает 14 циклов весового конвейера на один скор. Lean P использует тот факт, что 8-битный паттерн положительного числа e4m3, прочитанный как целое число и называемый его кодом, увеличивается на 8 каждый раз, когда значение удваивается. Таким образом, код 2ˣ близок к 8x плюс константа, поэтому lean P вычисляет код с помощью одной операции умножения-сложения и одного преобразования в целое число за 6 циклов на один скор.

Lean P является приближением. На 68 захваченных слоях при размерности головы 128 это увеличивает ошибку вызова внимания до 1.22–1.36× от ошибки точного P, что все равно составляет 0.26× от ошибки ядра AITER в среднегеометрическом выражении. ZYFA по умолчанию использует lean P при размерности головы 128 и exact P для формы MLA.

Двухволновое расписание типа «пинг-понг» на ассемблере

Каждый SIMD, представляющий собой один из четырех вычислительных блоков вычислительного блока, выполняет две волны FP8-ядра. На каждой фазе одна волна выдает матричные инструкции своего блока, в то время как другая вычисляет FP8-вероятности, и после барьера эти две волны меняются ролями (Рисунок 3). FlashAttention-3 называет такую организацию пинг-понг планированием. Каждая фаза длится до тех пор, пока обе волны не достигнут барьера, поэтому фаза занимает столько времени, сколько длится более медленная из двух ее волн. При точном P волна софтмакса прибывает последней, а матричный конвейер занят в течение 60% тактов. При усеченном P волна софтмакса прибывает первой, а матричный конвейер занят на 86%.

Поскольку обе волны исходят из одного и того же SIMD, перекрытие зависит от точного размещения каждой инструкции, поэтому генератор на Python выдает FP8-ядра в виде ассемблерного кода, что фиксирует положение каждой инструкции. На записанных активациях версия того же ядра, написанная на HIP, языке C++ ядер от AMD, требует на 43–46% больше тактов, а в пошаговом сравнении из статьи ассемблерное расписание дает наибольший прирост среди всех шагов проектирования.

Поблочное масштабирование и нормализация по округленной сумме строк

FP8-ядра умножают вероятности каждого блока из 64 ключей на степень двойки, так чтобы наибольшая вероятность блока была близка к 1. Масштабированная FP8-инструкция матричного умножения в CDNA4 делит эту степень двойки обратно в произведении PV без дополнительных затрат. Ядра также выполняют деление на сумму округленных вероятностей, которую одно дополнительное инструкция MFMA на блок вычисляет из тех же значений, которые умножаются на V. Каждый выход представляет собой взвешенное среднее с весами, сумма которых равна единице, а дополнительная инструкция MFMA добавляет 6% к матричным тактам на блок (Рисунок 4).

Пропускная способность

Пропускная способность на записанных активациях при постоянных тактовых частотах

Мы замеряем время выполнения каждого ядра на запросах, ключах и значениях, которые производит по одному слою Mistral-7B, Qwen3-8B и DeepSeek-V2-Lite при 16K токенах; мы называем это записанными активациями. При постоянной нагрузке MI355X снижает тактовую частоту шейдеров (SCLK), поэтому каждый запуск производит последовательный вызов одного ядра в течение нескольких секунд, а замер времени производится только после стабилизации тактовой частоты. FP8-ядра ZYFA достигают пропускной способности в 1.16×, 1.17× и 1.76× от показателей AITER на трех слоях (Рисунок 1). Отношение наиболее велико на DeepSeek-V2-Lite, поскольку CK, стандартное FP8-ядро AITER для формы MLA, выдает более старую матричную инструкцию с половинной пропускной способностью за такт. Внутри полных проходов прямого распространения (forward passes) при 16K токенах медианное соотношение FP8 по всем слоям составляет 1.26× на Mistral-7B и 1.25× на Qwen3-8B.

Пропускная способность на гауссовских входных данных

При оценке времени только для ядер на гауссовских входных данных, как это принято в типичных бенчмарках ядер, FP8-ядра ZYFA достигают более высокой пропускной способности, чем ядра AITER, на любой длине последовательности от 1K до 32K токенов, а BF16-ядра ZYFA достигают 1.00–1.03× от пропускной способности AITER при 16K и 32K токенах (Рисунок 5). В режиме BF16 матричные произведения занимают более чем в два раза больше тактов, чем софтмакс, поэтому они ограничивают время работы ядра для обеих библиотек. В режиме FP8 с точным P волна софтмакса достигает каждого барьера после волны MFMA, поэтому софтмакс, напротив, ограничивает время работы ядра.

Зависимость тактовой частоты шейдеров от значений операндов

Установившаяся тактовая частота MI355X зависит от значений в его арифметических блоках. Каждый бит, переключающийся между 0 и 1, потребляет ток, поэтому тактовая частота падает по мере того, как от одной инструкции к другой меняется все больше битов операндов (Рисунок 6). Цикл FP8 MFMA работает на частоте 2.34 ГГц с постоянными операндами и на частоте 2.07 ГГц с операндами, которые меняются при каждой инструкции.

На записанном слое Mistral-7B FP8-ядро ZYFA требует меньше тактов, чем ядро AITER, но работает на частоте 1,628 МГц против 2,108 МГц у AITER (Рисунок 7). AITER записывает субнормальные значения или нули для блоков, находящихся далеко ниже максимума строки, и эти значения переключают мало битов, в то время как поблочное масштабирование ZYFA поддерживает их нормальными. На гауссовских входных данных, которые делают внимание почти равномерным, оба ядра записывают только нормальные значения, и ускорение ZYFA на этом слое возрастает с 1.16× до 1.30×. Следовательно, пропускная способность на случайных входных данных ненадежно предсказывает пропускную способность на активациях реальных моделей.

Численная точность

Погрешность одного вызова внимания

Чтобы изолировать погрешность, вносимую самим ядром, мы сравниваем его вывод с вниманием в формате fp64 на тех же квантованных входных данных. В режиме FP8 при размерности головы 128 погрешности ZYFA, CK и Triton различаются на несколько процентов, в то время как погрешность ZYFA составляет 0.22× и 0.38× от погрешности FP8-ядра AITER на Mistral-7B и Qwen3-8B (Рисунок 8). Большая часть дополнительной погрешности AITER проистекает из деления на сумму неокругленных вероятностей. В полной модели Mistral-7B дивергенция Кульбака — Лейблера (KL) по сравнению с моделью с fp64-вниманием составляет 3.4 × 10⁻³ для FP8-ядра ZYFA и 5.9 × 10⁻³ для AITER.

Погрешность FP8 при инференсе полноразмерной модели

Большая часть погрешности FP8 в полной модели происходит от квантования входных данных (Рисунок 9). В Mistral-7B квантование Q, K и V добавляет в 8 раз больше погрешности, чем квантование вероятностей с усеченным P, причем Q и K добавляют больше погрешности, чем V. Вращение Адамара для Q и K по размерности головы оставляет QKᵀ неизменным, но распределяет большие каналы по всем каналам, что примерно вдвое сокращает ошибку входных данных. На Mistral-7B это вращение снижает KL-дивергенцию от модели BF16 с 45.7 до 21.9 раз от верхней границы шумовой полосы BF16, а внимание с вращением занимает 1.001× от времени работы всех FP8. Сохранение наиболее чувствительных голов в BF16 — менее эффективное использование времени. На Mistral-7B и Qwen3-8B даже 90% голов в BF16 не достигают шумовой полосы, и такой гибридный режим занимает больше времени, чем полностью BF16-внимание.

Оставшиеся пробелы в производительности и рекомендации

При 32K токенах на гауссовских входных данных FP8-ядро ZYFA при размерности головы 128 достигает 53% от пикового значения по даташиту. Остаются три пробела. Установившаяся тактовая частота остается значительно ниже 2.4 ГГц, предполагаемых пиком из даташита, матричный конвейер простаивает на переходах между тайл-блоками запросов, а при усеченном P волна MFMA также вычисляет максимум строки и масштабирование вывода, из-за чего волна MFMA теперь занимает больше времени, чем волна софтмакса, и ограничивает длину каждой фазы. Разработчикам ядер мы рекомендуем следующее:

• Производите измерения на реальных активациях при постоянных тактовых частотах, так как скорость тактовой частоты зависит от входных данных.
• Нацеливайтесь на квантование входных данных, например с помощью вращения Адамара для Q и К, прежде чем квантовать вероятности.
• Сокращайте объем векторной работы на один скор перед наращиванием матричной пропускной способности. Более быстрый матричный формат сам по себе не помогает, пока узким местом является софтмакс.
• Не допускайте упакованные инструкции FP32 в софтмакс, который выполняется параллельно с инструкциями MFMA. Упакованная инструкция применяет одну операцию к двум значениям на один поток, а упакованная инструкция FP32 простаивает, пока на том же SIMD выполняется MFMA.

Направления будущей работы

Мы продолжим совершенствовать ZYFA, уделяя особое внимание следующим областям:

• Устранение оставшихся вышеупомянутых пробелов.
• Распространение усеченного P на другие форматы, код которых увеличивается на фиксированный шаг при каждом удвоении значения, такие как e5m2 и FP4.
• Поддержка декодирования, кешей KV с постраничной организацией (paged KV caches) и батчей переменной длины.

Исходный код доступен по адресу github.com/Zyphra/zyfa, а статья опубликована на arXiv.

Методология и ограничения

• Все эксперименты запущены на одном узле с восемью GPU MI355X с политикой автоматического изменения тактовой частоты по умолчанию и лимитом общей мощности платы (TBP) 1400 Вт на ROCm 7.2.3. В качестве базовых показателей используются выборы ядер по умолчанию в AITER 0.1.19.
• Измерение времени работы только для ядер на гауссовских входных данных сообщает среднее время устройства за 100 запусков после 500 разогревочных запусков. Устойчивые измерения времени выполняют 1 с разогревочных запусков, а затем измеряют окно продолжительностью не менее 3 с.
• Исследование полной модели охватывает три модели, текст и не содержит нисходящих задач.
• Ядра охватывают каузальный прямой проход для равных длин запросов и ключей. Они не охватывают декoding (декодирование), кэши KV с постраничной организацией (paged KV caches) или батчи переменной длины.
• Тактовая частота, зависящая от данных, может отличаться при других настройках прошивки или питания, а наши ядра и базовые показатели быстро меняются, поэтому все результаты актуальны на момент времени.

Приложение A: Таблицы данных

Таблицы ниже содержат значения, представленные на графиках. Полные результаты можно найти в нашем техническом отчете.

A1. Захваченные активации при постоянных тактовых частотах (Рисунки 1 и 7)

По одному слою на модель при S = 16K (Mistral-7B, слой 20; Qwen3-8B, слой 18; DeepSeek-V2-Lite, слой 13). Коэффициент времени представляет собой время AITER, разделенное на время ZYFA, поэтому значения больше 1 означают, что ZYFA тратит меньше времени. Коэффициент для гауссовских данных — это то же отношение на гауссовских входных данных той же формы.

Время работы ядер и устойчивые тактовые частоты шейдеров (SCLK) для тех же запусков.

A2. Бенчмарки гауссовских ядер (Рисунок 5)

Пропускная способность ZYFA, деленная на пропускную способность ядра AITER по умолчанию, только для ядер на гауссовских входных данных при B = 1. Для формы MLA (размерность запрос-ключ 192) базовым показателем FP8 является CK. В последнем столбце указана пропускная способность ZYFA и AITER при 32K.

A3. Погрешность одного вызова внимания (Рисунок 8)

Относительная среднеквадратичная ошибка по сравнению с вниманием fp64 на тех же квантованных входных данных, среднее значение по слоям каждой модели. В скобках указана погрешность каждого ядра, разделенная на погрешность AITER в том же слое, усредненная по слоям, поэтому отношение может отличаться от отношения двух средних значений. Тритоновое ядро AITER не поддерживает FP8 для формы MLA.

A4. Погрешность одного вызова FP8 по этапам (Рисунок 9)

Относительная среднеквадратичная ошибка ядра FP8 от ZYFA, разделенная на квантование входных данных (Q, K и V), квантование вероятностей (P) и остальную часть арифметики ядра. ZYFA использует облегченное P на Mistral-7B и Qwen3-8B и точное P для формы MLA в DeepSeek-V2-Lite.

Сноски

Mistral-7B-v0.3 (слой 20), Qwen3-8B (слой 18) и DeepSeek-V2-Lite (слой 13) с запросами, ключами и значениями после позиционного эмбеддинга поворота (post-RoPE) на тексте WikiText-2 при 16K токенах. Базовым показателем является ядро, которое AITER 0.1.19 выбирает по умолчанию, что для FP8 при форме MLA соответствует Composable Kernel (CK). ↩

Вращение встраивается в проход, который квантует Q, а K и V квантуются по мере записи в кэш KV. На Mistral-7B внимание с вращением занимает 1,001× времени внимания во всех FP8. ↩

Наша публикация по инференсу «Wave 2» измеряет предельные характеристики (rooflines) MI355X, которые дают максимальную пропускную способность, достижимую ядром при каждой арифметической интенсивности (FLOP на байт, переданный в/из HBM), а также пиковую пропускную способность матриц для каждого типа данных. ↩

MLA кэширует один сжатый скрытый вектор на токен вместо ключей и значений для каждого хэда (головки внимания). Префилл (prefill) расширяет скрытые векторы в ключи и значения для каждого хэда и выполняет стандартное внимание с этими размерностями. Декодирование использует другую форму MLA, которую ZYFA не охватывает. ↩

Матричный конвейер — это матричный блок SIMD, а векторный конвейер — его векторный блок. Каждый вычислительный блок MI355X имеет четыре SIMD. Матричный конвейер выполняет матричные инструкции, а векторный конвейер выполняет покомпонентную арифметику, такую как экспоненты и преобразования FP8 для softmax. ↩

Деление на неокругленную сумму рассматривает каждую вероятность, которая ушла в минус до нуля, как ключ, вектор значений которого равен нулю, поэтому каждый выходной результат немного смещается к нулю. Деление на сумму округленных вероятностей вместо этого дает взвешенное среднее значений, чьи вероятности переживают округление, с весами, в сумме дающими единицу. ↩

Например, коды 1, 2 и 4 равны 56, 64 и 72, а код 60, находящийся на полпути между 56 и 64, соответствует 1,5. ↩

Каждый переключающийся бит заряжает или разряжает небольшую емкость, поэтому ток, потребляемый чипом, растет с увеличением числа битов, изменяющихся за цикл. Общая мощность платы остается ниже лимита в 1400 Вт в этих запусках, а телеметрия указывает на быстродействующее ограничение по току, а не по средней мощности. ↩

Шум BF16 представляет собой диапазон расхождений Кульбака — Лейблера, не превышающий таковой между двумя моделями с одинаково корректными ядрами BF16 (ZYFA и AITER), которые округляются в разном порядке. ↩

Пиковое значение из технического описания предполагает, что матричная инструкция выдается на каждом такте при пиковой частоте 2,4 ГГц. Наша публикация по инференсу Wave 2 сравнивает пиковое значение из технического описания для каждого типа данных с пропускной способностью, которую MI355X достигает на практике. ↩