Оптимизация MoE-кернелов: переход от push- к pull-based dispatching

🧪 Метод Основными задачи при разработке кернела были: • Уменьшить обьем CPU работы и CPU-GPU синхронизации • Максимально эффективно перекрывать вычисления и коммуникации Рассматриваются два типа dispatching Push-based dispatch (традиционный подход) 1. Router для каждого токена определяет его экспертов. 2. Для каждого токена выполняется: Отправить этот токен эксперту №7. То есть именно токен "толкается" (push) в буфер нужного эксперта.

T0 ---> Expert 2 T1 ---> Expert 7 T2 ---> Expert 2 T3 ---> Expert 5 :ale Проблемы Если тысячи потоков одновременно хотят записать данные эксперту 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, и они не знают, как их эффективно использовать.