UC Berkeley представила mKernel — persistent CUDA‑ядра для объединённой связи и вычислений в распределённых AI‑задачах

Команда UCCL из UC Berkeley представила mKernel — набор persistent (постоянно запущенных) CUDA‑ядер, в которых коммуникация и вычисления выполняются внутри одного ядра с целью снять нагрузку оркестрации с CPU и сократить микросекундные «пустоты» в конвейере. Такое слияние важно для масштабных распределённых AI‑задач, особенно для MoE и attention‑ориентированных рабочих нагрузок, где задержки связи серьёзно влияют на пропускную способность. Архитектура mKernel опирается на самоорганизающиеся CTA: потоки в пределах мультипроцессоров SM динамически самоназначают роли — вычисления, внутрикластерные передачи, межузловые отправки и редукции. Количество SM, отведённых под каждую роль, конфигурируется под форму задачи, а обмен данными и вычисления перекрываются на уровне тайлов/чанков, что обеспечивает мелкозернистое скрытие задержек.

Для межузловых операций mKernel применяет GPU‑инициированные RDMA‑записи на базе libibverbs и не зависит от NCCL или NVSHMEM; сетевой бэкенд разработан с нуля, чтобы поддерживать гетерогенные устройства и разные NIC‑топологии. Такой подход переводит контроль за передачами с хоста на GPU и ориентирован на снижение накладных расходов CPU при коллективных операциях.

Авторы описывают пять готовых сливаемых шаблонов ядра. AllGather + GEMM: каждый ранг хранит шард A и по мере сбора шардов от партнёров по NVLink/RDMA локальный GEMM потребляет поступающие тайлы. GEMM + AllReduce: вычисление C = A @ B сопровождается немедленной отправкой частичных результатов в дерево редукции. MoE Dispatch + GEMM: маршрутизация токенов MoE к экспертам (внутриузловой NVLink и межузловой all‑to‑all) с выполнением сгруппированных GEMM без промежуточной буферизации. поочерёдная ротация KV‑чанков по кольцу при одновременной обработке локальной FlashAttention предыдущего фрагмента. GEMM + ReduceScatter: вычисление с немедленной редукцией и пересылкой владельцу каждого выходного тайла.

mKernel позиционируется как ответ на ограничения host‑driven модели, где CPU оркестрирует коллективы (NCCL/NVSHMEM) и создаёт микросекундные разрывы в конвейере. Команда ссылается на анализ производительности: коммуникация может занимать до 43.6% времени прямого прохода и 32% полного обучения; в MoE‑моделях межустройственная связь достигает 47% выполнения, что делает такие оптимизации критичными при масштабировании.

Авторы также указывают на усиливающуюся чувствительность к накладным расходам CPU на современных конфигурациях: пример GB300 NVL72-72 Blackwell Ultra GPU, 36 Grace CPU, 720 PFLOP/s (FP8/FP6), 1.44 EFLOP/s (FP4) тензорных ядер и 130 TB/s all‑to‑all NVLink. Валидация mKernel проведена на двух двухузловых конфигурациях с по восемь H200 на узел, которые различались только межузловой тканью; в одном варианте использовали AWS EFA/SRD с 16 × 200 Gb/s EFA на узел.