Cursor открыл код ядра Mixture of Experts: обучение MoE ускоряется до 2,37 раза

Cursor, компания, известная одноимённым ИИ-редактором кода, развивает собственную линейку моделей Composer. Для их обучения на десятках тысяч GPU команда написала специализированное ядро MoK, которое решает главную проблему: при обучении моделей вроде DeepSeek-V3 слой mixture of experts может занимать больше половины общего времени, причём не из-за математики, а из-за ожидания, пока GPU обменяются данными. Публикация кода на GitHub даёт исследователям и компаниям с доступом к кластерам Blackwell готовый инструмент, а не только описание метода.

Что именно тормозило обучение MoE-моделей?

В архитектуре mixture of experts каждый слой маршрутизирует токены (минимальные единицы текста, которые «видит» модель) к разным «экспертам», разбросанным по десяткам GPU. Сами вычисления быстрые, но перед ними GPU должны переслать друг другу нужные токены, а после собрать результаты обратно. Эта пересылка и становится бутылочным горлышком.

На стойках GB300 NVL72, где 72 GPU связаны скоростной шиной NVLink, проблема меняет характер. Полоса пропускания огромная, но процессоры Grace, встроенные в эти стойки, медленнее GPU, и любая синхронизация через CPU съедает выигрыш. MoK убирает CPU из цикла целиком.

Три решения, которые дали результат

Pull вместо push при отправке токенов. Существующие подходы (например, DeepEP) «толкают» данные на чужие GPU. Cursor замерила: при таком способе обратная полоса NVLink простаивает, а сигнализация о завершении занимает 103 микросекунды. Если GPU сам «тянет» нужные токены, утилизация канала растёт на 29%, а задержка сигнала падает до 18 микросекунд. MoK комбинирует оба режима: pull для отправки, push для сборки результатов.

Средняя степень перекрытия. Конкуренты выбирают либо очень мелкое дробление операций (Comet), либо очень крупное (DeepEP). Cursor показала, что оптимум зависит от нагрузки, и для моделей класса Kimi 2.5 (базовая модель Composer 2.5) минимальный порог составляет 2 368 токенов на пакет вычислений.

Кольцевой буфер токенов. Вместо того чтобы выбрасывать «лишние» токены или просить CPU рассчитать размер буфера, MoK использует фиксированный кольцевой буфер размером в несколько сотен мегабайт. Буфер обходится в обратном порядке при обратном проходе, что сокращает повторное воспроизведение активаций.

Всё это упаковано в единое мегаядро: один вызов GPU выполняет и пересылку, и вычисления, и получает детерминированный (воспроизводимый) результат. Поддерживаются форматы BF16 и MXFP8 (сжатое представление чисел, экономящее память).

Что показали замеры?

Тесты слоя MoE проводились внутри одной стойки NVL72 с 64 GPU. Каждый GPU обрабатывал 2 048 токенов до маршрутизации. Сравнивались четыре комбинации: NCCL+PyTorch, DeepEP+PyTorch, DeepEP+TransformerEngine и HybridEP+Megatron.

• Формат MXFP8, прямой проход: до 2,37 раза быстрее лучшего конкурента.
• Формат MXFP8, обратный проход: 1,78 раза.
• Формат BF16, прямой проход: 1,92 раза.
• Формат BF16, обратный проход: 1,58 раза.

Сквозной тест на 512 GPU (несколько стоек GB300 NVL72) дал рост с 760,9 до 1 070,2 токенов в секунду на GPU, то есть 1,41 раза.

Тестировались конфигурации моделей Kimi K2.7 Code, GLM-5.2, Qwen3.5-397B-A17B и DeepSeek-V4-Pro.

Что делать с этим прямо сейчас?

Исследователям и инженерам с доступом к Blackwell-кластерам. Репозиторий на GitHub даёт рабочий код и конкретные архитектурные решения для ускорения обучения MoE-моделей типа DeepSeek. Три ключевых приёма (pull-dispatch, кольцевой буфер, средняя гранулярность перекрытия) можно изучить и адаптировать, даже если вы строите собственное ядро.

Компаниям в РФ и СНГ, работающим с GPU-кластерами. Стойки GB200/GB300 NVL72 пока редкость, но национальные вычислительные центры и крупные облака (например, Yandex Cloud, SberCloud) наращивают GPU-мощности. Когда Blackwell появится в доступе, MoK станет готовым инструментом, а не исследовательской статьёй, которую нужно реализовывать с нуля.

Авторам и маркетологам. Прямого применения в контенте нет, но понимание, что ускорение обучения MoE-моделей в 2,37 раза реально и выложено в открытый доступ, помогает оценить скорость, с которой следующие поколения моделей будут появляться.

Открытие кода MoK под Apache-2.0 даёт индустрии не готовый продукт для массового использования, а детальный чертёж: как именно снять главный тормоз при обучении MoE-моделей на масштабе тысяч GPU. Для тех, у кого есть железо, это ускорение в 1,41 раза на полном цикле обучения, для остальных это три конкретных инженерных приёма, которые можно адаптировать уже сейчас.

По данным Marktechpost