🧪 Метод и эксперименты
Основными задачи при разработке кернела были:
• Уменьшить обьем CPU работы и CPU-GPU синхронизации
• Максимально эффективно перекрывать вычисления и коммуникации
Рассматриваются два типа dispatching
Push-based dispatch (традиционный подход)
1. Router для каждого токена определяет его экспертов.
2. Для каждого токена выполняется: Отправить этот токен эксперту №7.
То есть именно токен "толкается" (push) в буфер нужного эксперта.
⚠️ Проблемы
Если тысячи потоков одновременно хотят записать данные эксперту 7, они начинают конкурировать за
• atomic counter,
• место в буфере,
• память.
Получается много случайных записей и плохая эффективность памяти.
Pull-based dispatch
Идея разворачивается наоборот.
Сначала известно, какие токены принадлежат каждому эксперту (например, после сортировки).
После этого сам эксперт приходит за своими токенами.
Теперь данные читаются большими непрерывными кусками.
Никаких scatter-записей больше не требуется.
Pull позволяет организовать данные так, чтобы чтение стало почти непрерывным, что улучшает пропускную способность памяти и снижает накладные расходы на scatter/atomic операции.
Кернел валидируют на разных MoE моделях - GLM-5.2, Qwen-3.5-397B, Kimi-K2.7, DeepSeek V4 Pro.
Удается добиться ускорения ~2x против бейзлайнов на форварде, и ~1.5x на обратном проходе.
📌 Выводы
Выкладывание в открытый доступ данного кернела - хороший подарок компаниям, которым завезли NVL72, и они не знают, как их эффективно использовать.
Основными задачи при разработке кернела были:
• Уменьшить обьем CPU работы и CPU-GPU синхронизации
• Максимально эффективно перекрывать вычисления и коммуникации
Рассматриваются два типа dispatching
Push-based dispatch (традиционный подход)
1. Router для каждого токена определяет его экспертов.
2. Для каждого токена выполняется: Отправить этот токен эксперту №7.
То есть именно токен "толкается" (push) в буфер нужного эксперта.
T0 ---> Expert 2
T1 ---> Expert 7
T2 ---> Expert 2
T3 ---> Expert 5
⚠️ Проблемы
Если тысячи потоков одновременно хотят записать данные эксперту 7, они начинают конкурировать за
• atomic counter,
• место в буфере,
• память.
Получается много случайных записей и плохая эффективность памяти.
Pull-based dispatch
Идея разворачивается наоборот.
Сначала известно, какие токены принадлежат каждому эксперту (например, после сортировки).
После этого сам эксперт приходит за своими токенами.
Expert 2:
беру токены [0,2,9,15]
Expert 5:
беру токены [3,8]
Expert 7:
беру токены [1,4,5,6]
Теперь данные читаются большими непрерывными кусками.
Никаких scatter-записей больше не требуется.
Pull позволяет организовать данные так, чтобы чтение стало почти непрерывным, что улучшает пропускную способность памяти и снижает накладные расходы на scatter/atomic операции.
Кернел валидируют на разных MoE моделях - GLM-5.2, Qwen-3.5-397B, Kimi-K2.7, DeepSeek V4 Pro.
Удается добиться ускорения ~2x против бейзлайнов на форварде, и ~1.5x на обратном проходе.
📌 Выводы
Выкладывание в открытый доступ данного кернела - хороший подарок компаниям, которым завезли NVL72, и они не знают, как их эффективно использовать.