Anatomía de un kernel de paralelismo de expertos de alto rendimiento

Fuentes: Anatomy of a high-performance EP kernel
Imagen generada por IA con el prompt: Editorial illustration of GPU server racks connected by glowing fiber links, with packets of data flowing between nodes, dark blue and cyan tones, technical schematic style
Imagen generada con IA

Los modelos de lenguaje de gran tamaño (LLM) requieren coordinar muchas GPUs para funcionar. Una de las técnicas clave es el paralelismo de expertos (EP), esencial en los modelos MoE a gran escala. A diferencia de otras formas de paralelismo, cuyas comunicaciones siguen patrones fijos, en EP el router decide en tiempo de ejecución qué tokens envía a qué expertos, repartidos por distintos nodos y GPUs. Esa variabilidad obliga a diseñar kernels de comunicación a medida para repartir las activaciones y reunirlas tras los cálculos. La biblioteca DeepEP, creada por DeepSeek, marcó la forma moderna de estos kernels. Este artículo descompone paso a paso un kernel estilo DeepEP: cómo se prepara el envío (fase de coordinación para conocer cuántos tokens recibirá cada experto), cómo se reagrupan las activaciones en un buffer denso apto para un GEMM agrupado, y las dos variantes del kernel: una optimizada para throughput (fase de prefill) y otra optimizada para latencia (fase de decode). La primera intercambia primero los conteos entre ranks para reservar solo la memoria HBM necesaria, mientras que la segunda reserva un rectángulo fijo dimensionado al peor caso para minimizar latencia, a costa de gastar más HBM. Se discute además la comunicación intra-nodo por NVLink e inter-nodo por RDMA, y cómo se solapa el envío con el cómputo de GEMM mediante técnicas como el microbatching, habituales en stacks como vLLM o SGLang. En conjunto, el texto explica la ingeniería de las comunicaciones que permite servir modelos MoE como DeepSeek a gran escala.