Compute-communication overlap в MoE снижает latency за счёт тайлового планирования и сигналов. Это напрямую влияет на throughput и загрузку GPU.
Современные Mixture-of-Experts (MoE) системы упираются не в compute, а в коммуникацию. В распределённом исполнении каждый слой требует двух all-to-all операций, и вторая — возврат результатов — попадает в критический путь. Классическая схема запускает её только после завершения GEMM, из-за чего GPU простаивает, ожидая передачи данных. Это усугубляется дисбалансом: вычислительная мощность GPU растёт быстрее, чем пропускная способность interconnect, поэтому доля времени на communication увеличивается.
Попытки перекрыть вычисления и коммуникацию (compute-communication overlap) обычно используют крупнозернистую декомпозицию. GEMM и all-to-all разбиваются на чанки и выполняются в пайплайне. Такой подход прост, но даёт побочные эффекты: ухудшается эффективность tensor cores на малых чанках, появляется синхронизация на стороне хоста и возрастает overhead kernel launch. Альтернатива — fusion, где compute и communication объединяются в один kernel. Это снижает накладные расходы, но требует сложной инженерии и кастомных примитивов.
В рассматриваемом подходе выбран компромисс: сохранить раздельные kernel’ы, но синхронизировать их на уровне тайлов (tile-level signaling). Ключевая идея — начать передачу данных сразу после готовности части GEMM, не дожидаясь завершения всей операции.
Решение строится как producer-consumer модель внутри GPU. Producer — это persistent GEMM kernel, который обрабатывает все эксперты на ранге без повторных запусков. Consumer — отдельный persistent kernel, который занимается передачей данных через NVSHMEM. Они работают параллельно на разных SM (streaming multiprocessors), что устраняет конкуренцию за ресурсы.
Критический элемент — сигнализация на уровне тайлов. После вычисления каждого tile в epilogue GEMM kernel публикуется флаг готовности. Consumer отслеживает эти флаги и инициирует передачу данных сразу, как только накоплен достаточный объём. Это устраняет зависимость от host-side orchestration и сокращает latency между compute и communication.
Чтобы сделать такую схему возможной, вводится специальная организация данных — remote-owner-aligned layout. Каждому tile заранее соответствует один destination rank. Это устраняет необходимость маршрутизации на уровне строк и позволяет выполнять передачу как один contiguous write. Дополнительно используется стратегия remote-first scheduling: сначала вычисляются данные, которые нужно отправить другим GPU. Это увеличивает окно overlap.
Гранулярность передачи тоже становится параметром. Передавать каждый tile отдельно неэффективно из-за низкой загрузки bandwidth. Поэтому вводится сегментация: несколько row bands объединяются в один transfer. Малые сегменты уменьшают latency старта, большие — повышают throughput. Баланс подбирается эмпирически.
На уровне исполнения система разделяет SM между compute и communication. Небольшая часть выделяется под consumer kernel. Это создаёт trade-off: слишком мало ресурсов — communication не успевает за compute; слишком много — замедляется GEMM. Оптимальный диапазон оказался в средней зоне, где достигается устойчивый overlap без деградации вычислений.
Результаты показывают, что такой compute-communication overlap даёт до 2.64x ускорения end-to-end и до 2.74x на уровне MoE слоя по сравнению с базовой реализацией без overlap. При сравнении с другими MoE системами прирост зависит от модели. В сценариях с большим объёмом коммуникации выигрыш наиболее заметен. Если размер скрытого слоя меньше, эффект снижается, так как уменьшается объём данных, который можно «спрятать» за вычислениями.
Отдельно важно, что решение не требует модификации базовых примитивов (GEMM, all-to-all) и не вводит жёсткой зависимости от кастомных kernel’ов. Это делает его практичным для интеграции в существующие стеки.
С инженерной точки зрения это эволюционное улучшение: вместо изменения алгоритма или модели оптимизируется критический путь исполнения. Подход показывает, что дальнейшее масштабирование MoE будет зависеть не только от архитектуры моделей, но и от того, насколько эффективно система умеет скрывать коммуникационные издержки.
Источник информации
arXiv — крупнейший открытый репозиторий препринтов (с 1991 года, под эгидой Корнелла), где исследователи оперативно размещают рабочие версии статей; материалы общедоступны, но не проходят полное рецензирование, поэтому результаты следует считать предварительными и, по возможности, сверять с обновленными версиями или рецензируемыми журналами. arxiv.org