FlashAttention развернули поперёк RISC-V-вектора и удвоили темп

У FlashAttention есть неприятная привычка: алгоритм тщательно экономит внешнюю память, но может почти не пользоваться арифметикой самого процессора. На EPAC, RISC-V-чипе с регистрами по 16 384 бита, один вектор содержит 512 значений FP32. Портированный CPU-вариант ядра умножал вероятности attention на V вдоль head dimension d=64. В каждом векторе реально работали лишь 64 позиции из 512.

Вектор заполняют запросами

Авторы решили не укорачивать вектор и не менять математику attention. Они повернули представление промежуточных данных. После вычисления score-тайла QKᵀ его строки соответствуют запросам, а столбцы — ключам. В привычной раскладке векторная ось проходит по 64 компонентам выходного признака. В новой раскладке score-тайл и аккумулятор выходов транспонируются: один вектор собирает значения для множества query. Длинный регистр получает работу по своей естественной непрерывной оси.

Ускорение состоит из двух вмешательств

Перестановка данных затрагивает и softmax, где у исходного кода оказался другой скрытый расход. Полиномиальная аппроксимация exp держала свои коэффициенты в векторных регистрах. Для каждой операции их приходилось загружать и обрабатывать как массивы, хотя коэффициенты одинаковы для всех элементов. Их перенесли в scalar-регистры. В конфигурации с Bq=128, Bkv=512, N=1024 и d=64 число vector-memory instructions внутри softmax упало с 62 416 до 13 264, то есть на 78.7%.

Длинный вектор не спасает decode

На одном ядре EPAC при FP32 и d=64 переориентированное ядро выполнило attention для N=1024 за 28.22 Mcycles против 51.94 Mcycles у исходного варианта. Авторы приводят ускорения 1.84×, 2.02× и 2.08× для N=1024, 2048 и 4096 соответственно; для каждого размера выбирался свой лучший тайлинг. Утилизация пикового FP32 throughput выросла с 16.2% до 29.7%. Это заметный результат именно потому, что EPAC уже располагает восемью FMA-каналами и пиком 32 FP32 FLOP за такт: проблема была в том, как подать им данные.

Приписывать весь выигрыш одному транспонированию было бы неверно. Полная цифра до 2.08× включает и новую ориентацию тайлов, и скалярные коэффициенты exp, а сравниваемые ветви используют оптимальные для себя размеры блоков. В отдельном сравнении с уже оптимизированным fused traversal одна переориентация дала 1.43× при N=1024: 40.23 против 28.22 Mcycles. На FPGA авторы получили до 2.71×, однако это тот же дизайн на 50 MHz с иным соотношением стоимости памяти и вычислений; переносить эту величину на чип с частотой 1 GHz нельзя.

При autoregressive decode новый запрос часто приходит по одному, поэтому ось query слишком коротка, чтобы заполнить 512 FP32-позиций. Авторы прямо ограничивают вывод prefill- и training-подобными режимами с крупными блоками запросов. Результаты также относятся к одному ядру, одной attention head и FP32. Но сам приём гораздо шире конкретного ядра: перед тем как добавлять вычислительные блоки в SIMD-ускоритель, стоит проверить, какая ось тензора действительно течёт по регистру. Иногда достаточно развернуть таблицу, и почти пустой 2-КиБ вектор внезапно начинает работать.