UC Berkeley presenta mKernel para fusionar comunicación multi‑GPU y multi‑nodo en un kernel persistente

El equipo UCCL de UC Berkeley presentó mKernel, una librería de kernels CUDA persistentes que traslada la coordinación de comunicaciones del CPU a la GPU, integrando NVLink intra‑nodo, RDMA entre nodos y cómputo con solapamiento fino a nivel de mosaicos.

UC Berkeley anunció mKernel, una librería de kernels CUDA persistentes diseñada para mover la orquestación de comunicaciones del CPU a la GPU. El proyecto parte de la premisa de que la comunicación pesa de forma significativa en el entrenamiento: citan que la comunicación puede consumir 43.6% del forward pass y 32% del tiempo de entrenamiento end‑to‑end, y que en modelos Mixture‑of‑Experts (MoE) la comunicación inter‑dispositivo llega hasta 47% del tiempo total.

mKernel fusiona en un único kernel persistente la comunicación intra‑nodo por NVLink, la comunicación inter‑nodo mediante RDMA y el cómputo denso. Los CTAs (bloques de hilos) se autoasignan roles — cómputo, comunicación intra, envío inter y reducción inter— y el número de SMs dedicado a cada rol puede configurarse según la forma del tensor. El diseño soporta además solapamiento fino a nivel de tile/chunk para maximizar la concurrencia entre transmisión y cálculo.

El proyecto se propone superar las limitaciones del modelo host‑driven actual, en el que la CPU orquesta llamadas a bibliotecas como NCCL o NVSHMEM y lanza operaciones colectivas; esa orquestación introduce burbujas de pipeline a pesar de ocurrir en microsegundos. mKernel se orienta específicamente al escenario multi‑nodo y ofrece como referencia un rack GB300 NVL72: 72 GPUs Blackwell Ultra, 36 CPUs Grace, 720 PFLOP/s FP8/FP6, 1.44 EFLOP/s FP4 y 130 TB/s de NVLink intra‑rack.

Para lograr independencia del host, mKernel usa RDMA iniciado por GPU a través de libibverbs y un backend de comunicaciones implementado desde cero para admitir dispositivos heterogéneos, evitando dependencia directa de NCCL o NVSHMEM. Si la solución se valida en despliegues reales, podría reducir latencias de orquestación y permitir un solapamiento más fino entre comunicación y cómputo, con implicaciones para el entrenamiento y servicio de modelos a gran escala.