Dao-AILab/flash-attention

Что это за инструмент

FlashAttention — это библиотека для PyTorch, обеспечивающая быстрое и энергоэффективное вычисление механизма внимания (attention) за счет оптимизации работы с памятью GPU.

Зачем нужен

Инструмент решает проблему высокой нагрузки на память и узкой пропускной способности при обучении и инференсе больших языковых моделей (LLM). Он ускоряет вычисления и снижает потребление VRAM, что позволяет работать с более длинными контекстами и большими моделями на существем оборудовании.

Что можно реализовать

• Обучение и тонкая настройка (fine-tuning) больших языковых моделей (LLM) с ограниченным объемом видеопамяти.
• Ускорение инференса LLM, особенно при использовании длинных контекстных окон.
• Применение в архитектурах, использующих механизм внимания, таких как Transformer, для повышения производительности.
• Бенчмаркинг и тестирование производительности на GPU разных поколений (Ampere, Hopper, Blackwell).

Ключевые возможности

• IO-awareness: оптимизация доступа к памяти для минимизации пересылки данных между HBM и SRAM.
• Поддержка нескольких версий: FlashAttention-2 (универсальная), FlashAttention-3 (для Hopper/H100) и FlashAttention-4 (для Hopper/Blackwell с CuTeDSL).
• Широкая поддержка аппаратного обеспечения: NVIDIA (CUDA 12.0+, Ampere/Ada/Hopper/Blackwell) и AMD (ROCm 6.0+, MI200/MI300/RDNA).
• Работа с форматами данных fp16, bf16 и fp8 (в зависимости от версии и GPU).
• Простая установка через pip и интеграция с PyTorch 2.2+.

👤 Кому подойдёт: ML-инженеры, исследователи в области NLP, разработчики LLM и дата-сайентисты, работающие с обучением и развертыванием больших моделей.