De CUDA a MLX: el puente que lleva décadas de optimización a Apple Silicon

¿Qué ha ocurrido?

El equipo del Berkeley Sky Lab, en colaboración con IBM Research, ha extendido su framework de búsqueda evolutiva de kernels K-Search con una capa de traducción estructurada de CUDA a MLX. Esta capa permite que el conocimiento experto acumulado en kernels CUDA —como FlashAttention o Mamba— sea transferido automáticamente al ecosistema de Apple Silicon, evitando tener que redescubrir cada optimización manualmente. Los resultados son contundentes: se alcanza el 97% del rendimiento del kernel nativo de MLX para FlashAttention, y se logra una aceleración de hasta 20x en el prefill del modelo de espacio de estados Mamba respecto a la implementación comunitaria mlx-lm.

El trabajo, publicado en el blog de BAIR (Berkeley Artificial Intelligence Research), detalla cómo el framework K-Search, que originalmente se centraba en la búsqueda evolutiva de kernels para CUDA, ha sido adaptado para generar kernels optimizados para MLX, el framework de aprendizaje automático de Apple para su silicio. La clave está en una capa de traducción que mapea las optimizaciones de CUDA a estrategias nativas de MLX, aprovechando las características únicas de la arquitectura unificada de memoria de Apple Silicon. Según los autores, este enfoque no solo acelera el desarrollo, sino que también permite que el ecosistema de Apple se beneficie de décadas de ingeniería de kernels en CUDA.

¿Por qué es importante?

El ecosistema CUDA acumula décadas de ingeniería de kernels optimizados a mano, un conocimiento que no está disponible en plataformas más nuevas como Apple Silicon. Con la creciente adopción de MLX para inferencia local en MacBooks y Mac Studios (cientos de millones de dispositivos), cerrar esta brecha de rendimiento es crucial para que los desarrolladores puedan ejecutar modelos de lenguaje y otros workloads de IA de manera eficiente sin depender de GPUs NVIDIA. Este trabajo demuestra que es posible transferir ese conocimiento de forma automática, lo que podría acelerar significativamente la adopción de hardware alternativo en IA.

Históricamente, la portabilidad de kernels entre ecosistemas ha sido un desafío enorme. Por ejemplo, cuando AMD lanzó ROCm, la comunidad tuvo que reescribir manualmente muchas optimizaciones de CUDA, un proceso que tomó años y aún hoy no alcanza la paridad completa. Del mismo modo, proyectos como OpenAI Triton han intentado abstraer la escritura de kernels, pero aún requieren ajustes específicos por hardware. Este trabajo de Berkeley e IBM representa un salto cualitativo al automatizar la transferencia de conocimiento, lo que podría cambiar las reglas del juego para la adopción de hardware alternativo en IA.

¿Cómo funciona la traducción?

La capa de traducción no copia instrucción por instrucción, sino que mapea las estrategias de optimización de CUDA a técnicas nativas de MLX. Por ejemplo, las optimizaciones de memoria compartida en CUDA se traducen al uso de la memoria unificada de Apple Silicon, y los patrones de paralelización se adaptan a la arquitectura de los cores GPU de Apple. El framework K-Search utiliza búsqueda evolutiva guiada por IA para explorar el espacio de configuraciones de kernel, partiendo de los kernels CUDA traducidos como semilla, y encuentra variantes optimizadas para MLX.

En concreto, el proceso comienza con un kernel CUDA existente, como el de FlashAttention. La capa de traducción lo analiza y extrae las optimizaciones clave, como el uso de memoria compartida para reducir latencia o el tiling para mejorar la localidad de datos. Luego, genera un kernel en MLX que implementa esas mismas estrategias, pero utilizando las primitivas nativas de Metal y la memoria unificada. Este kernel traducido sirve como punto de partida para K-Search, que aplica mutaciones y selección evolutiva para explorar variantes, evaluando su rendimiento en hardware real de Apple Silicon. El resultado es un kernel optimizado que, en el caso de FlashAttention, alcanza el 97% del rendimiento del kernel nativo de MLX desarrollado manualmente por la comunidad.

Consecuencias para el ecosistema

Esta investigación tiene implicaciones profundas:

• Para desarrolladores de IA: podrán ejecutar modelos optimizados en hardware Apple sin esperar a que la comunidad desarrolle kernels nativos. Esto reduce la barrera de entrada para experimentar con modelos de vanguardia en hardware asequible y ampliamente disponible.
• Para Apple: refuerza la posición de MLX como plataforma seria para inferencia de IA, compitiendo directamente con CUDA. Con el respaldo de una herramienta automatizada como K-Search, Apple podría atraer a más desarrolladores a su ecosistema, especialmente aquellos que buscan alternativas a NVIDIA.
• Para la industria: el enfoque es generalizable a otros ecosistemas (como AMD ROCm o Google TPU), lo que podría democratizar el acceso a kernels de alto rendimiento. Empresas como AMD podrían beneficiarse de una herramienta similar para portar automáticamente kernels de CUDA a ROCm, acelerando la adopción de sus GPUs en IA.

Además, este avance podría reducir la dependencia de NVIDIA en el mercado de IA, fomentando una mayor competencia y diversidad de hardware. Para las empresas que invierten en infraestructura de IA, esto se traduce en más opciones y potencialmente menores costos.

Limitaciones y advertencias

Los autores señalan que la transferencia no es perfecta: se alcanza el 97% del rendimiento nativo en FlashAttention, pero hay una pérdida marginal. Además, la técnica se ha evaluado principalmente en kernels de atención y modelos de espacio de estados; su generalización a otros tipos de kernels (convoluciones, normalizaciones) está por verse. El framework K-Search requiere recursos computacionales para la búsqueda evolutiva, aunque es un costo único por kernel. En sus experimentos, los autores utilizaron un clúster de GPUs NVIDIA para ejecutar la búsqueda, lo que podría ser una barrera para algunos desarrolladores, aunque el costo se amortiza si el kernel se utiliza ampliamente.

Otra limitación es que la traducción automática puede no capturar todas las optimizaciones específicas de la arquitectura de Apple Silicon, como el uso de la Neural Engine o las unidades de matriz de Apple. Los kernels resultantes son competitivos, pero aún no superan a los nativos en todos los casos. La comunidad de MLX deberá validar estos resultados en una gama más amplia de modelos y hardware.

Qué deben saber los lectores

Este avance no es un reemplazo inmediato del trabajo manual de optimización, sino una herramienta que reduce drásticamente el tiempo y esfuerzo necesarios para portar kernels entre ecosistemas. Los desarrolladores que trabajen con MLX deberían estar atentos a la integración de K-Search en herramientas como MLX Community o frameworks de Hugging Face. Para empresas que invierten en hardware Apple para inferencia, esto significa que pronto podrán ejecutar modelos con rendimiento cercano al de NVIDIA sin tener que reescribir código desde cero.

En un contexto más amplio, este trabajo se alinea con la tendencia hacia herramientas de compilación y optimización automatizadas, como XLA o TVM, que buscan abstraer las complejidades del hardware subyacente. Sin embargo, K-Search va un paso más allá al transferir conocimiento experto entre ecosistemas, lo que podría acelerar la convergencia hacia un estándar de kernels portables. Como señala Shiyi Cao, autor principal: "Transferir un kernel de un hardware de un proveedor a otro es más difícil de lo que parece, y a menudo significa redescubrir las mismas optimizaciones desde cero. Este trabajo muestra que ese conocimiento puede transferirse automáticamente."

"Transferir un kernel de un hardware de un proveedor a otro es más difícil de lo que parece, y a menudo significa redescubrir las mismas optimizaciones desde cero. Este trabajo muestra que ese conocimiento puede transferirse automáticamente." — Shiyi Cao, autor principal.