🚀 Flash Attention на AMD за 80 строк

AMD показала сверхэффективную реализацию Flash Attention для MI300X с помощью TileLang.

TileLang — это DSL для разработки GPU-ядер, который скрывает низкоуровневые детали HIP и CUDA. На его основе удалось реализовать Flash Attention менее чем за 80 строк кода (вместо 500+ в CUDA), сохранив производительность на уровне ручной оптимизации. Ключевые фишки: автотюнинг 108 конфигураций за секунду, оптимизация под shared memory и регистры, а также встроенная поддержка кэширования и memory coalescing. Ядро автоматически адаптируется под архитектуру MI300X.

Решение ускоряет выполнение на 2.7×: 0.36 мс против 0.97 мс у PyTorch. Ошибка вычислений — в пределах допустимого (rtol=0.01).

#amd_mi300x #flash_attention #tilelang #rocm #ai_kernels #gpu_optimization