Moonshot AI выложила FlashKDA для ускорения механизма Kimi Delta Attention.
Компания Moonshot AI представила FlashKDA, быструю реализацию Kimi Delta Attention на базе CUTLASS, оптимизированную под современные GPU NVIDIA. По сути, это CUDA-ядра для эффективного запуска их архитектур.
Замеры производительности
- H20: ускорение примерно до 2,31× относительно
chunk_kda. - GB200: до 3,27× на последовательностях переменной длины (variable-length sequences).
Среди ключевых возможностей FlashKDA стоит отметить поддержку батчинга для последовательностей разной длины и готовую интеграцию с библиотекой flash-linear-attention (начиная с версии 0.5.0). Благодаря этому FlashKDA может автоматически подхватываться как бэкенд для расчётов.
Пока требования к окружению довольно высокие: архитектура GPU от SM90+, CUDA 12.9+ и PyTorch 2.4+.
Мы видим интересную тенденцию: одновременно с развитием моделей Moonshot постепенно открывает и весь низкоуровневый стек, делая запуск своих архитектур более доступным и производительным.
Исходный код ждёт вас на GitHub: github.com/MoonshotAI/FlashKDA
Ещё новости:
