sustcsonglin/flash-linear-attention

🚀 Эффективные реализации для новых архитектур моделей

Инференс⭐ 5 773Pythonпоследний релиз: v0.5.2
Открыть на GitHub ↗Сайт проекта ↗

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

Библиотека с оптимизированными ядрами для эффективного обучения и инференса современных архитектур последовательностей, таких как линейное внимание, SSM и гибридные LLM.

Зачем нужен

Инструмент решает проблему вычислительной сложности стандартного Attention, предлагая аппаратно-эффективные реализации (kernels) для NVIDIA, AMD и Intel. Он полезен для ускорения обучения и снижения потребления памяти при работе с длинными контекстами и гибридными моделями.

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

  • Интеграция ядер линейного внимания (например, Gated DeltaNet, KDA) в существующие LLM для ускорения инференса.
  • Обучение гибридных моделей, сочетающих традиционный Attention с SSM (State Space Models), с использованием фьюженных слоев.
  • Оптимизация распределенного обучения (Context Parallel) для последовательных моделей на кластерах GPU.
  • Использование готовых реализаций SOTA-архитектур (RWKV, Mamba, RetNet) для экспериментов и бэнчмарков.

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

  • Поддержка широкого спектра архитектур: линейное внимание, SSM, гибридные модели и sparse attention.
  • Аппаратная независимость: верифицировано на NVIDIA, AMD и Intel, с бэкендами Triton, TileLang и Gluon.
  • Наличие фьюзированных модулей (fused modules) и оптимизированных слоев (например, fused linear and cross-entropy) для экономии памяти.
  • Активная экосистема: интеграция в крупные проекты (Qwen3-Next) и наличие фреймворка обучения flame.
  • Регулярные обновления с новыми SOTA-реализациями (GDN-2, Parallax, Wall attention, Mamba3).

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

Релизы

v0.5.2minor
🕐 1 мес назад · релиз на GitHub ↗

Релиз v0.5.2: поддержка кэшированного инференса, исправления переполнения int32, оптимизации для NPU Ascend и фиксы генерации.

  • Added: Добавлена поддержка кэшированного инференса (cached inference) для NSA.
  • Fixed: Исправлено переполнение int32 при работе с длинными последовательностями (T).
  • Fixed: Устранена ошибка инициализации rotary output при left-padded batched generation.
  • Added: Оптимизированы и расширены бэкенды для NPU Ascend (CANN 9.0.0, TileLang, causal_conv1d).
  • Added: Добавлены производительные оптимизации: fused GatedDeltaNet convolutions и ускоренный parallel backward для NSA.