Conozca mKernel: una biblioteca de núcleo fusionado con múltiples GPU y múltiples nodos para comunicación impulsada por GPU

La sobrecarga de comunicación de la GPU es un cuello de botella mensurable en las cargas de trabajo de producción de IA. Según los datos citados por el proyecto mKernel, la comunicación puede consumir el 43,6% del pase directo y el 32% del tiempo de entrenamiento de un extremo a otro. En los modelos populares de Mezcla de Expertos (MoE), la comunicación entre dispositivos puede representar hasta el 47% del tiempo total de ejecución. Investigadores del proyecto UCCL de UC Berkeley han lanzado mKernel, una biblioteca de núcleos CUDA persistentes que fusionan la comunicación NVLink entre nodos, RDMA entre nodos y la computación en un solo núcleo.

El problema: comunicación impulsada por el host

El modelo estándar para la comunicación multi-GPU está controlado por el host: la CPU ejecuta la ruta de control y llama a una biblioteca como NCCL o NVSHMEM. La biblioteca emite la operación colectiva (AllReduce, AllGather, etc.) en todas las GPU. La computación y la comunicación se ejecutan en flujos CUDA separados y se superponen en los límites del núcleo.

El equipo de investigación identifica dos problemas con este enfoque:

(1) Las CPU no se escalan con el cómputo de GPU. Un bastidor GB300 NVL72 integra 72 GPU Blackwell Ultra y 36 CPU Grace, lo que ofrece un rendimiento de 720 PFLOP/s FP8/FP6, 1,44 EFLOP/s FP4 Tensor Core y 130 TB/s de ancho de banda NVLink intra-rack completo. A esas velocidades, la sobrecarga de orquestación del host a escala de microsegundos (una llamada cudaLaunchKernel, una verificación de “todas las escrituras realizadas” en el lado de la CPU, un evento entre transmisiones) se muestra directamente como burbujas de canalización.

(2) Los sistemas controlados por host se superponen a la computación y la comunicación en los límites aproximados del núcleo. No es posible una superposición más fina a nivel de mosaico o fragmento desde el lado anfitrión.

La alternativa es la comunicación impulsada por GPU: la propia GPU activa las transferencias, con la comunicación fusionada en el mismo núcleo que el proceso. La mayoría de las bibliotecas de kernel fusionadas existentes operan dentro de un solo nodo o una sola GPU. mKernel apunta al caso de múltiples nodos.

¿Qué hace mKernel?

mKernel es una biblioteca de núcleos CUDA persistentes. Cada kernel fusiona la comunicación NVLink dentro del nodo, la RDMA entre nodos y la computación densa en un solo kernel.

Multi-GPU + multi-nodo, en un kernel: Tanto NVLink intranodo como RDMA entre nodos viven dentro del mismo kernel persistente.

Superposición intranúcleo detallada: la computación y la comunicación se superponen en granularidad de mosaico/fragmento, cubriendo la comunicación de GPU tanto dentro como entre nodos.

Núcleo persistente con especialización en SM: los CTA autoasignan roles: computación, comunicación intracomunitaria, envío entre envíos y reducción entre ellos. La cantidad de SM dedicados a cada función se puede ajustar por forma.

Red impulsada por GPU basada en libibverbs: mKernel utiliza escrituras RDMA iniciadas por GPU sin depender de NCCL o NVSHMEM. El backend de comunicación está escrito desde cero para maximizar el rendimiento y admitir dispositivos de red heterogéneos.

Los cinco núcleos fusionados

KernelQué fusionaDescripciónAllGather + GEMMAllGather → GEMMCada rango contiene un fragmento de A. Mientras que los rangos reúnen los fragmentos de sus pares a través de NVLink/RDMA, el GEMM local consume mosaicos tan pronto como llegan.GEMM + AllReduceGEMM → AllReduceCalcula C = A @ B y reduce las salidas parciales en todos los rangos en un solo lanzamiento. Los mosaicos de salida se insertan en el árbol de reducción en el instante en que se producen. Envío MoE + envío GEMMTodo a todos → GEMM agrupado Enruta los tokens MoE a sus rangos de expertos (NVLink intranodo + todos a todos entre nodos) y ejecuta el GEMM agrupado por experto en el mismo núcleo. Los tokens se procesan tan pronto como aterrizan, sin búfer de preparación de ida y vuelta. Ring Attention Intercambio Ring KV → FlashAttention Secuencia de atención paralela entre rangos. Cada paso rota un fragmento de KV alrededor del anillo mientras FlashAttention local consume el fragmento recibido previamente. Computar y el envío/recepción del anillo se ejecutan simultáneamente dentro de un único núcleo persistente. GEMM + ReduceScatterGEMM → ReduceScatterCalcula C = A @ B y reduce-dispersa la salida. Cada ficha de salida se reduce y se envía a su rango de propiedad tan pronto como se produce.

Configuración de evaluación

El equipo de investigación evaluó mKernel en dos clústeres de 2 nodos × 8-H200 que difieren solo en su tejido entre nodos:

Banco de pruebasNodos × GPUIntranodoTransporte entre nodosNICAWS EFA2 × 8 H200NVLinkAWS EFA / SRD16 × 200 Gb/s EFA por nodoConnectX-72 × 8 H200NVLinkInfiniBand8 × 400 Gb/s NVIDIA ConnectX-7 por nodo

mKernel se comparó con NCCL, Triton-distribuido, Flux, Mercury, MagiAttention, Transformer-Engine y ring-flash-attention. El equipo señala que aún se están realizando más evaluaciones comparativas a mayor escala.

Backends y requisitos

mKernel admite dos backends de red:

BackendMacroTransportDónde se ejecutaCX7-DINTERNODE_BACKEND_IBVERBSlibibverbs RCConnectX-7 / InfiniBand / RoCEEFA-DINTERNODE_BACKEND_EFAlibibverbs + efadv (SRD)AWS p5/p5e (H200, EFA)

Ambos backends comparten la misma API del lado del host y el mismo kernel en la GPU. Solo difiere la implementación de proxy/sesión (session.h para CX7, session_efa.h para EFA). Requisitos: GPU NVIDIA Hopper (destinos de compilación predeterminados sm_90a), CUDA 12.9, Python con PyTorch. El backend de CX7 requiere bibliotecas y encabezados de desarrollo de libibverbs. El backend de EFA requiere la instalación de AWS EFA con encabezados libfabric, libibverbs, efadv y EFA en EFA_HOME=/opt/amazon/efa de forma predeterminada.

Explicador visual de Marktechpost

01 / 07 — Descripción general

Qué es mKernel?

mKernel es una biblioteca de código abierto de núcleos CUDA persistentes del proyecto UCCL de UC Berkeley. Fusiona comunicación NVLink entre nodos, RDMA entre nodos y computación densa en un solo núcleo.

La mayoría de las bibliotecas de kernel fusionadas existentes operan dentro de un solo nodo o una sola GPU. mKernel está diseñado desde el principio para abarcar los límites de los nodos.

43,6%

de pase directo consumido por la comunicación en producción

47%

del tiempo total de ejecución en modelos MoE populares

32%

del tiempo de capacitación de un extremo a otro consumido por la comunicación

02 / 07 — El problema

Por qué Impulsado por el host La comunicación se queda corta

El modelo estándar está controlado por el host: la CPU llama a NCCL o NVSHMEM, que emite operaciones colectivas entre las GPU. El equipo de UCCL identifica dos problemas.

Las CPU no se escalan con las GPU. Un bastidor GB300 NVL72 ofrece 720 PFLOP/s FP8/FP6 y 1,44 EFLOP/s FP4. A esas velocidades, la sobrecarga a escala de microsegundos de cudaLaunchKernel, las comprobaciones de sincronización del lado de la CPU y los eventos entre transmisiones aparecen directamente como burbujas de canalización.

🔲

La superposición es demasiado gruesa. Los sistemas controlados por host se superponen a la computación y la comunicación sólo en los límites del núcleo. No es posible una superposición más fina a nivel de mosaico o fragmento desde el lado anfitrión.

🔀

La respuesta: comunicación impulsada por GPU. La propia GPU activa transferencias detalladas, fusionadas en el mismo núcleo que el cómputo.

03 / 07 — Diseño

Diseño de cuatro núcleos Propiedades

🖧

Multi-GPU + multi-nodo, en un solo kernel. NVLink dentro del nodo y RDMA entre nodos viven dentro del mismo núcleo persistente.

🔬

Superposición intrakernel de grano fino. La computación y la comunicación se superponen en la granularidad de mosaico/fragmento, cubriendo tanto la comunicación dentro como entre nodos.

⚙️

Kernel persistente con especialización SM. Las llamadas a la acción autoasignan roles: computación, intracomunicación, interenvío, interreducción. La división SM se puede ajustar por forma.

📡

Redes impulsadas por GPU a través de libibverbs. Utiliza escrituras RDMA iniciadas por GPU. Sin dependencia de NCCL o NVSHMEM. El backend de comunicación está escrito desde cero.

04 / 07 — Granos

los cinco Núcleos fusionados

Todos reunidos + GEMM

Todos reunidos —> GEMM

Cada rango contiene un fragmento de A. El GEMM local consume mosaicos a través de NVLink/RDMA a medida que llegan: matmul comienza antes de que finalice el colectivo.

GEMM + TodoReducir

GEMM —> TodoReducir

Calcula C = A @ B y reduce las salidas parciales en todos los rangos en un solo lanzamiento. Los mosaicos de salida ingresan al árbol de reducción en el instante en que se producen.

Despacho del Ministerio de Educación + GEMM

Despacho todos a todos —> GEMM agrupado

Enruta tokens MoE a rangos de expertos a través de NVLink + entre nodos de forma integral, luego ejecuta GEMM agrupados por expertos en el mismo núcleo. Sin viaje de ida y vuelta del buffer de preparación.

Anillo de atención

Intercambio de KV en anillo —> FlashAttention

Atención secuencial-paralela entre rangos. Cada paso rota un fragmento de KV alrededor del anillo mientras FlashAttention local consume el fragmento recibido previamente.

GEMM + Reducir Dispersión

GEMM —> Reducir dispersión

Calcula C = A @ B y reduce-dispersa la salida. Cada ficha se reduce y se envía a su rango de propiedad tan pronto como se produce.

05 / 07 — Evaluación

Evaluación Configuración

Probado en dos clústeres de 2 nodos × 8-H200 que se diferencian solo en la estructura entre nodos.

Banco de pruebasNodos × GPUInternodoNIC AWS EFA2 × 8 H200AWS EFA / SRD16 × 200 Gb/s EFA por nodo ConnectX-72 × 8 H200InfiniBand8 × 400 Gb/s CX7 por nodo

Ambos bancos de pruebas utilizan NVLink dentro del nodo. Comparado con: NCCL, Triton-distribuido, Flux, Mercury, MagiAttention, Transformer-Engine y ring-flash-attention. Todavía se están realizando evaluaciones comparativas a mayor escala.

06/07 — Backends y requisitos

Motores y Requisitos

BackendTransportDónde se ejecuta CX7libibverbs RCConnectX-7 / InfiniBand / RoCE EFAlibibverbs + efadv (SRD)AWS p5/p5e (H200, EFA)

📋

Requisitos: GPU NVIDIA Hopper (sm_90a predeterminado), CUDA 12.9, Python con PyTorch. CX7 necesita encabezados libibverbs. EFA necesita libfabric, libibverbs, efadv en EFA_HOME=/opt/amazon/efa.

📝

Licencia y atribución: licencia MIT. MMA/código informático adaptado de ThunderKittens (HazyResearch).

07 / 07 — Hoja de ruta y conclusiones clave

Hoja de ruta y Conclusiones clave

Núcleos multinodo fusionados impulsados ​​por GPU (AG+GEMM, GEMM+AR, MoE Dispatch+GEMM, Ring Attention, GEMM+RS)

Backends de ConnectX-7 y AWS EFA

🚧

Compatibilidad total con acelerador/NIC heterogéneo con descubrimiento, ubicación y enrutamiento con reconocimiento de topología

🚧

Meganúcleos entre nodos: colapso de varios pasos fusionados en un solo meganúcleo que abarca una capa transformadora

🚧

Compatibilidad con GPU Blackwell

Fusiona NVLink, RDMA entre nodos y computación en un único núcleo CUDA persistente

Cinco núcleos: AllGather+GEMM, GEMM+AllReduce, MoE Dispatch+GEMM, Ring Attention, GEMM+ReduceScatter

RDMA iniciado por GPU a través de libibverbs: sin dependencia de NCCL o NVSHMEM

Requiere GPU Hopper (sm_90a) y redes ConnectX-7 o AWS EFA

Conclusiones clave

mKernel fusiona NVLink entre nodos, RDMA entre nodos y computación en un único kernel CUDA persistente. Los gastos generales de comunicación representan hasta el 47 % del tiempo de ejecución en los modelos MoE según los datos de producción citados. Se incluyen cinco núcleos: AllGather+GEMM, GEMM+AllReduce, MoE Dispatch+GEMM, Ring Attention y GEMM+ReduceScatter. La RDMA iniciada por GPU se implementa directamente a través de libibverbs, sin dependencia de NCCL o NVSHMEM. Actualmente requiere GPU Hopper (sm_90a) y redes ConnectX-7 o AWS EFA; El soporte de Blackwell está en la hoja de ruta.

Consulte el repositorio y los detalles técnicos. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 150.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.

¿Necesita asociarse con nosotros para promocionar su repositorio de GitHub O su página principal de Hugging O su lanzamiento de producto O seminario web, etc.? Conéctate con nosotros