Ускоряем LLM по максимуму. Как я создал кросс-платформенный Flash Attention с поддержкой Turing+ архитектур и не только

На сегодняшний день трансформеры правят балом хайпа в мире машинного обучения, особенно после появления ChatGPT и ему подобных языковых моделей. Это стало возможным благодаря лежащему в основе их архитектуры механизму внимания (attention), однако он же и является слабым местом с точки зрения производительности и потребления памяти. Хотя в связи с этим и была разработана изящная концепция Flash Attention (Tri Dao), её существующие реализации имеют ряд ограничений.

Поэтому представляю вашему вниманию первую и единственную open-source реализацию Flash Attention 2 на Triton с поддержкой Linux и Windows, Turing-Blackwell архитектур (теперь можно работать в Google Colab и Kaggle), гомо и гетерогенных кластеров, опциональным детерминизмом, а также возможностью ручной кастомизации ядер (kernels) для более гибкой настройки под каждую GPU архитектуру отдельно. Более подробно о том как это устроено и не только — далее в статье.