Основными задачи при разработке кернела были:
• Уменьшить обьем 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, и они не знают, как их эффективно использовать.