thu-ml/SpargeAttn

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

SpargeAttn — это метод разреженного внимания (sparse attention) без дообучения, который ускоряет инференс языковых, визуальных и видео-моделей.

Зачем нужен

Инструмент решает проблему вычислительной сложности стандартного механизма внимания, заменяя его на оптимизированные CUDA-ядра. Это позволяет значительно ускорить генерацию токенов и обработку видео без потери качества и необходимости переобучения моделей.

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

• Замена torch.nn.functional.scaled_dot_product_attention в существующих LLM для ускорения инференса
• Оптимизация работы видео-моделей (например, CogVideoX) за счет настройки параметров sparsity
• Применение к любым моделям с архитектурой Transformer без необходимости калибровки или дообучения
• Использование кастомных блочных масок для специфических паттернов внимания в разных головках (heads)

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

• Полностью training-free подход: работает сразу после установки, без этапа калибровки
• Гибкое управление точностью и скоростью через параметр topk (баланс между точностью и разреженностью)
• Поддержка кастомных блочных масок (block-sparse mask) размером 128x64 для каждой головы внимания
• Оптимизированные CUDA-ядра (SageAttention2), поддерживающие fp8 на GPU Ada/Hopper и Ampere
• Plug-and-play интеграция: простая замена стандартного API attention

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