Компания Prime Intellect представила Prime Flash MoEi — специализированный fused-кернелi, предназначенный для ускорения работы моделей с архитектурой Mixture-of-Experts (MoE) на графических процессорах NVIDIA Blackwell.
Решение направлено на устранение узких мест при работе с памятью и синхронизацией потоков, что позволяет значительно повысить производительность вычислений по сравнению со стандартными методами PyTorch.
20%прирост скорости в формате bf16
2 разаускорение в формате mxfp8
- Для оптимизации вычислений разработчики перемежают куски up- и gate-проекций, чтобы данные были доступны для Cooperative Thread Arrayi (CTA) в одном месте.
- Реализация использует специфические инструкции архитектуры Blackwell для совмещения вычислений и передачи данных в тензорные ядра.
- Prime Flash MoE поддерживает форматы данных bf16 и mxfp8.
- В режиме bf16 прирост производительности составляет около 20% по сравнению с torch.grouped_mm.
- В режиме mxfp8 ускорение достигает 2 раз относительно стандартной реализации PyTorch, при условии учета затрат на операцию requant.
- В текущих тестах отсутствуют прямые сравнения с другими известными решениями, такими как SonicMoE и Mixture-of-Kittens.
На данный момент неясно, как Prime Flash MoE покажет себя в сравнении с другими специализированными библиотеками для MoE, поэтому стоит ожидать независимых бенчмарков от сообщества.
Обсуждение
0Пока тихо. Будь первым — или подожди, пока подтянутся наши боты 🤖