x86-альянс представил ACE, новый ускоритель матричных вычислений для ИИ

Intel в своё время добавила в x86 расширение AMX (Advanced Matrix Extensions), набор 2D-регистров-«плиток» (tile) и настроечных регистров для ускорения матричного умножения в задачах машинного обучения. Первой реализацией стал блок TMUL (tile matrix multiply unit) в серверных процессорах Sapphire Rapids: он умел работать с типами INT8, FP16 и BF16, а в Granite Rapids-D появилась поддержка комплексных чисел с компонентами FP16.

Теперь отраслевая группа x86 Ecosystem Advisory Group (объединение вендоров, развивающих архитектуру x86) выпустила whitepaper и спецификацию нового ускорителя ACE, второго типа блока в семействе AMX наряду с TMUL. ACE устроен проще: конфигурация «плиток» больше не гибкая, а фиксированная, всегда 64 байта на 16 строк. Поддержка комплексных чисел исчезла, зато добавлена работа с форматом FP8. Главное архитектурное отличие: если TMUL и обычные AVX512-VNNI считают через операции внутреннего (скалярного) произведения векторов, то ACE, как и конкурирующее расширение Arm SME, использует операции внешнего произведения, оно даёт матрицу ранга 1 и, как показывает автор статьи, лежит в основе почти всей линейной алгебры (умножение матриц можно представить как сумму внешних произведений; на нём же построены сингулярное разложение матриц и даже быстрое преобразование Фурье в переработанном виде). Исторически чаще использовали внутреннее произведение, потому что оно экономнее по регистрам, но конвертировать между двумя формами несложно, и ACE делает ставку на внешнее произведение.

Для сравнения: расширение Arm SME (и его развитие SME2) использует не фиксированный размер регистров, а переменную «потоковую» длину вектора (streaming vector length, SVL), от 128 до 2048 бит с шагом вдвое, как и у SVE. Хранилище ZA (аналог tile-регистров AMX) масштабируется вместе с SVL, от 256 байт до 64 КБ. ACE, в отличие от Arm, сохраняет фиксированные 8 КБ tile-регистров, унаследованные от AMX.

Отдельный блок статьи посвящён деквантованию, преобразованию сжатых весов модели в форматы, которые понимает ускоритель. И ACE, и Arm SME (в отличие, например, от TensorCores NVIDIA) выполняют предобработку входных векторов программно, поэтому могут поддерживать произвольные, даже нестандартные форматы, а не только фиксированный набор «нативных». ACE опирается на 512-битные регистры AVX-512/AVX10, которых достаточно, чтобы хранить таблицу соответствия для типов данных шириной до 6 бит через инструкцию VPERMB, а для 7-битных типов, сразу два таких регистра через VPERMI2B. Новая инструкция VUNPACKB (часть ACE/AVX10.3) распаковывает элементы шириной 2-7 бит в побайтовое выравнивание, после чего в дело идут те же инструкции перестановки. Таким образом всего тремя инструкциями ACE закрывает произвольную ширину типа от 2 до 7 бит, авторы x86 EAG надеются, что это пригодится не только для деквантования весов моделей, но и, например, для кодовых книг сжатия данных.

Arm не может полагаться на широкие векторные регистры как таблицы соответствия, поскольку длина векторов SVE/SME переменная. Поэтому в SME2 добавлен отдельный 512-битный регистр ZT0, работающий как таблица 16 на 4 байта, а инструкции LUTI2 и LUTI4 распаковывают 2- или 4-битные индексы, ищут значения в ZT0 и кладут результат в целевой регистр. Это даёт Arm экономию: конвертация в одну инструкцию вместо двух и меньше давления на векторные регистры, но за счёт гибкости, форматы младше 2 или отличные от 4 бит таблицей ZT0 не ускоряются, а сложные схемы с несколькими таблицами соответствия не поддерживаются вовсе.

Ещё одна тема, блочное масштабирование (block scaling) для низкоточных форматов вроде FP8, у которых маленький динамический диапазон. Спецификация Microscaling Formats консорциума Open Compute Project решает это через один коэффициент масштабирования на блок значений. ACE вводит для этого новый 1024-битный регистр BSR0, разделённый на две 512-битные половины (по одной на каждый вход операции внешнего произведения), каждая, на четыре группы по 16 восьмибитных коэффициентов, выбираемых непосредственно операндом инструкции. У Arm SME похожий механизм реализован через 8-битные поля LSCALE и LSCALE2 регистра FPMR: инструкция BF1CVT использует LSCALE, BF2CVT, LSCALE2; масштабирование и собственно умножение у Arm разнесены на две инструкции, тогда как у ACE, на одну. Обе архитектуры перезаписывают такие регистры целиком, а значит рассчитывают, что коэффициенты масштабирования меняются нечасто.

Наконец, автор разбирает эффект от размера тайлов на конкретном примере: перемножение двух матриц 32768x32768 в INT8 (35,2 триллиона операций умножения-сложения). ACE сохраняет те же 8 КБ tile-регистров, что и у AMX, но математические инструкции теперь берут оба входных операнда из обычных векторных регистров AVX-512, а не из tile-регистров, в отличие от TMUL, где все операнды берутся из tile-регистров, из-за чего часть их ёмкости уходит под временные входные данные, а не под аккумуляторы результата.

Ключевые факты

x86 Ecosystem Advisory Group опубликовала whitepaper и спецификацию ACE, второго типа ускорителя матричных вычислений в семействе Intel AMX, наряду с уже существующим TMUL
ACE упрощает конфигурацию tile-регистров (фиксированные 64 байта на 16 строк), убирает поддержку комплексных чисел, добавляет FP8 и переходит на операции внешнего (а не внутреннего) произведения, как у Arm SME/SME2
Для деквантования весов ACE использует 512-битные регистры AVX-512/AVX10 как таблицы соответствия (VPERMB, VPERMI2B, новая инструкция VUNPACKB) и покрывает типы шириной 2-7 бит тремя инструкциями
Arm SME2 решает ту же задачу иначе, через выделенный регистр ZT0 и инструкции LUTI2/LUTI4, что экономнее по числу инструкций, но ограничено типами в 2 и 4 бита
Для FP8 и блочного масштабирования ACE вводит новый 1024-битный регистр BSR0; у Arm похожий, хотя и двухшаговый, механизм реализован через поля LSCALE/LSCALE2 регистра FPMR

Почему это важно

ACE, попытка x86-индустрии (в лице отраслевой группы x86 Ecosystem Advisory Group) закрыть отставание от Arm в аппаратном ускорении матричных вычислений для ИИ прямо на процессоре, без обращения к GPU. Расширение строится поверх уже существующего Intel AMX, но переносит вычисления на операции внешнего произведения, тот же подход, что использует Arm в SME/SME2, и добавляет поддержку низкоточных форматов (FP8, произвольные форматы деквантования), которые нужны для эффективного инференса больших моделей.

Кому это важно

В первую очередь, инженерам, проектирующим x86-процессоры и компиляторы (Intel, AMD и другие участники x86 Ecosystem Advisory Group), а также разработчикам инференс-движков и библиотек линейной алгебры, которым предстоит писать код под новые инструкции. Полезно и тем, кто сравнивает архитектурные пути x86 и Arm в гонке за эффективный ИИ-инференс на CPU.

Как это применить

ACE предлагает три новые инструкции для деквантования произвольных форматов шириной 2-7 бит (VUNPACKB вместе с VPERMB/VPERMI2B) и отдельный регистр BSR0 для блочного масштабирования FP8-значений. Математические инструкции ACE в отличие от TMUL берут оба операнда из обычных векторных регистров AVX-512, а не из tile-регистров AMX, что освобождает tile-регистры под накопители результата и позволяет использовать бóльшие тайлы при матричном умножении.

Можно ли доверять

Материал, детальный технический разбор от chipsandcheese.com, издания, специализирующегося на анализе микроархитектуры процессоров, основанный на официальном whitepaper и спецификации x86 Ecosystem Advisory Group и на прямом сопоставлении с уже задокументированными и частично реализованными инструкциями Arm SME/SME2. Важная оговорка: ACE, пока спецификация, а не реализованное в кремнии железо, и статья не называет ни производителя, который первым выпустит чип с ACE, ни сроков.

Риски и подводные камни

И у ACE (регистр BSR0), и у Arm SME (регистр FPMR) масштабирующие коэффициенты записываются в регистр целиком, если в реальных нагрузках значения масштабирования будут меняться часто, а не редко, как рассчитывают разработчики обеих архитектур, это создаст барьер сериализации для FP8-инструкций на обеих платформах. Кроме того, ACE платит за гибкость и задел на будущее (произвольная ширина форматов через VUNPACKB) потерей части возможностей TMUL, в частности, поддержки комплексных чисел и гибкой конфигурации tile-регистров.