Triton FP8: реализация и автотюнинг GEMM/Attention под RTX 40xx/Blackwell

Привет, Хабр!
Я тут на досуге решил разобраться с 8-битными числами с плавающей запятой (FP8) и попробовать написать под них свои GPU-ядра на Triton. Зачем? Ну, новые ускорители от NVIDIA обещают невиданную ранее производительность на FP8 — вдвое больше операций, чем на FP16. Для тренировки огромных нейросетей это прям манна небесная: меньше памяти, больше скорость. Но, конечно, всё не так просто.
Сам по себе FP8 формат капризный, требует тщательного подхода: два разных варианта чисел (E4M3 и E5M2), масштабирование (скейлинг) тензоров, аккуратное накопление сумм в FP16/FP32, разбиение вычислений на тайлы под размер быстрого SRAM (shared memory) и даже специальные меры против переполнения. В общем, нюансов хватает. Расскажу, что узнал и как реализовал GEMM (умножение матриц) и внимание (Attention) в FP8 на языке Triton, да ещё и с автотюнингом.