thu-ml/SageAttention

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

SageAttention — это набор оптимизированных ядер внимания (attention kernels) для GPU, обеспечивающих ускорение инференса в 2-5 раз по сравнению с FlashAttention без потери точности.

Зачем нужен

Инструмент решает проблему медленного выполнения операций внимания в нейросетях за счет квантования (INT8, FP8, FP4) и сглаживания выбросов. Он полезен для разработчиков, так как работает как плагин (plug-and-play), не требуя переобучения моделей, и сохраняет метрики качества на уровне оригинальных моделей.

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

• Ускорение инференса языковых моделей (LLM) на GPU архитектур Ampere, Ada и Hopper.
• Оптимизация генерации изображений и видео (например, CogVideoX) с сохранением визуального качества.
• Запуск распределенного инференса с поддержкой torch.compile и non-cudagraphs mode.
• Работа с моделями, использующими group-query attention (GQA) и переменную длину последовательностей.

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

• Поддержка нескольких версий: SageAttention (INT8), SageAttention2 (INT4/FP8) и SageAttention3 (FP4 для Blackwell).
• Совместимость с современными GPU: NVIDIA H100, H20, RTX 4090, RTX 5090 и другими.
• Высокая точность благодаря стратегии двухуровневой аккумуляции и сглаживанию выбросов.
• Гибкая поддержка форматов данных: FP8, INT8, INT4 и FP16 accumulator.

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