Figura 1: Mapa de traducción de optimización de CUDA a MLX. El conocimiento de optimización de CUDA puede traducirse a estrategias nativas de arquitectura MLX en lugar de copiar instrucción por instrucción.
Nos enfrentamos a una nueva época en la computación. El hardware está cambiando rápidamente; no solo con GPUs más veloces, sino con una creciente gama de chips de diferentes fabricantes, cada uno con su propia arquitectura y, a menudo, adaptados a cargas de trabajo de IA específicas. El software evoluciona con la misma rapidez, y las herramientas de codificación de IA generan ahora en minutos lo que hace pocos años requería meses de esfuerzo.
Con gran parte de la computación centrada en la IA, los kernels de GPU son un componente crucial para su éxito. Estos son los programas de bajo nivel que se ejecutan dentro de la GPU, y escribir versiones eficientes está lejos de ser algo obvio: requiere años de experiencia para lograrlo. Transferir un kernel del hardware de un fabricante a otro es aún más difícil, y a menudo significa redescubrir las mismas optimizaciones desde cero. El ecosistema CUDA, por ejemplo, ha acumulado décadas de experiencia en kernels: implementaciones optimizadas a mano de atención, modelos de espacio de estados (SSM) y otras operaciones críticas que representan miles de horas de ingeniería. Los ecosistemas de hardware más nuevos (Apple Silicon, aceleradores de IA personalizados y otros) están creciendo rápido, pero carecen de esta profundidad.
En este trabajo nos preguntamos si esa experiencia puede transferirse automáticamente. Construimos sobre K-Search, un framework evolutivo de búsqueda de kernels presentado por Cao et al. en el Berkeley Sky Lab que utiliza IA para optimizar kernels de GPU, y lo extendimos con un backend para MLX, el framework de machine learning de Apple para sus propios chips Apple Silicon. Desarrollamos una novedosa capa de traducción estructurada de CUDA a MLX que permite a K-Search tomar kernels de CUDA existentes como base de conocimiento y adaptarlos en kernels de GPU de alta calidad para Apple Silicon, en lugar de reconstruirlos desde cero.
Demostramos que nuestro enfoque alcanza un rendimiento de nivel casi experto en Apple Silicon con una aceleración de 0.97x en comparación con el kernel de atención nativo de MLX, y hasta una aceleración de 20x en el prellenado (prefill) sobre la implementación comunitaria mlx-lm en el kernel Mamba SSM; reportamos las cifras, y cuánto de la ganancia proviene de la capa de traducción, en las secciones a continuación. Aunque nos enfocamos en kernels de MLX para Apple Silicon, el método no es específico para MLX y se aplica a cualquier ecosistema donde la experiencia en CUDA sea transferible.
¿Por qué MLX?
El framework MLX de Apple ha experimentado una adopción notable desde finales de 2023. Con Apple Silicon presente en cientos de millones de MacBooks y Mac Studios, MLX permite la inferencia de IA local sin costos de nube. La arquitectura de memoria unificada lo hace especialmente atractivo para modelos de tamaño mediano (7B a 70B de parámetros en chips serie M).
Sin embargo, bajo este impulso subyace una brecha significativa: muchos kernels críticos para el rendimiento que el ecosistema NVIDIA da por sentado (paged attention, kernels de escaneo SSM optimizados, enrutamiento MoE fusionado) están ausentes o son ingenuos sin un ajuste específico para el hardware. MLX ejecuta los modelos correctamente, pero a menudo deja un rendimiento significativo sobre la mesa.
Esta brecha es lo que motiva el resto de esta publicación.
¿Qué es K-Search?
K-Search es un framework de optimización de kernels evolutivo desarrollado originalmente por nuestro autor principal Shiyi Cao en el UC Berkeley Sky Lab. Dado un kernel ingenuo y una especificación de hardware, ejecuta un bucle de optimización iterativo: un LLM razona sobre qué optimizaciones intentar a continuación, un modelo de escritura de código genera kernels candidatos, y esos candidatos se compilan y prueban (benchmark) en hardware real.
Las mediciones alimentan la búsqueda, que continúa refinándose, persiguiendo direcciones prometedoras y descartando callejones sin salida hasta que el rendimiento converge.
Algoritmo 1: K-Search mediante modelos de mundo co-evolutivos. La búsqueda alterna entre seleccionar la acción más prometedora, instanciar y evaluar código hasta que la mejora se estanca, y evolucionar el modelo de mundo a través de operaciones de inserción, actualización y poda. Adaptado de Cao et al. (2026).
La búsqueda se fundamenta en una Especificación: un documento específico del dominio que codifica reglas de hardware, patrones de optimización y restricciones matemáticas que evita que el código generado alucine primitivas inválidas y garantiza que los candidatos realmente se compilen y ejecuten eficientemente.
En nuestras ejecuciones, un solo modelo (Gemini 3.5 Pro Preview) desempeña ambos roles: mantiene el estado de razonamiento y escribe los kernels. La mitad de razonamiento se solicita como un "ingeniero de rendimiento de kernels de GPU" y se le pide que trabaje a través de un análisis fijo antes de proponer nada: clasificar el kernel (reducción, escaneo, atención/softmax, etc.), reescribir el cálculo de referencia en forma canónica, mapear el diseño de datos y patrones de acceso, e hipotetizar el cuello de botella probable (ancho de banda, latencia, cómputo o sincronización) en cada régimen de ejecución. Solo entonces emite optimizaciones candidatas, cada una como un cambio único implementable en una iteración.
Llamamos al estado de razonamiento persistente un modelo de mundo. En lugar de una lista plana de cosas por intentar, es un árbol de decisión (prefijo): cada camino desde la raíz hasta la hoja compone un plan de optimización completo, y las ramas hermanas son alternativas competitivas. Cada nodo es puntuado: un overall_rating en [0, 10], una confidence en [0, 1], e impacts por nodo en el ancho de banda de memoria, presión de registros y ajuste de cómputo/hardware, para que la búsqueda pueda clasificar planes parciales y expandir los más prometedores. El árbol persiste y crece a través de las rondas: refinar una idea agrega un nodo hijo en lugar de sobrescribir a su padre, y si la mejor puntuación no mejora durante unas cuantas rondas (una ventana de estancamiento), la búsqueda retrocede para explorar una rama alternativa. Un solo nodo, tal como aparece a mitad de la ejecución en el kernel de atención, se ve así:
{
"action": "Replace the threadgroup-memory softmax reduction with a register-only reduction: each SIMD group owns 8 query rows and reduces across lanes with simd_shuffle_xor, removing a threadgroup_barrier.",
"difficulty_1_to_5": 4,
"impacts": {
"memory_bandwidth": 8,
"register_pressure": 4,
"compute_hw_fit": 9
},
"overall_rating_0_to_10": 8,
"confidence_0_to_1": 0.7
}Listado 1: Ejemplo de nodo de modelo de mundo de K-Search. Cada optimización candidata registra una acción concreta, impactos de hardware estimados, una calificación de prioridad general y la confianza del modelo.
Figura 2: Descripción general de K-Search. El framework opera sobre un Estado de Búsqueda $S_t$ estructurado como un árbol de búsqueda. El árbol consiste en nodos Cerrados (azul, estados visitados con programa adjunto como $x_{12}$) y una Frontera de nodos Abiertos (naranja, hipótesis pendientes como $u_{13}$). El flujo de trabajo itera a través de tres fases: (1) Selección de Acción, donde el nodo de acción más prometedor se recupera de la frontera basado en la puntuación de prioridad estimada del modelo de mundo $V$; (2) Refinamiento Local, donde una política estocástica $\pi_{\mathrm{code}}$ muestrea implementaciones concretas hasta el estancamiento; y (3) Actualización del Modelo de Mundo, donde el LLM razona sobre la trayectoria para actualizar el árbol de búsqueda a través de Inserción (agregando nuevas acciones), Actualización (ajustando $V$, por ejemplo, $u_{11}$ bajando de 0.9 a 0.6), y Poda (eliminando nodos menos prometedores como $u_{10}$).
El artículo original de K-Search evaluó esta estrategia de búsqueda en kernels de CUDA de FlashInfer. A través de decodificación GQA, decodificación MLA, prellenado MLA y MoE, K-Search mejoró más consistentemente que OpenEvolve y ShinkaEvolve durante el mismo presupuesto de 120 iteraciones.
Vía Berkeley AI Research.
