Cómo K-Search aporta décadas de experiencia en kernel a Apple Silicon – The Berkeley Artificial Intelligence Research Blog

Figura 1: Mapa de traducción de optimización de CUDA a MLX. El conocimiento de optimización de CUDA se puede traducir en estrategias MLX nativas de la arquitectura en lugar de copiar instrucción por instrucción.

Nos enfrentamos a una nueva época en la informática. El hardware está cambiando rápidamente: no solo GPU más rápidas, sino una gama cada vez mayor de chips de diferentes proveedores, cada uno con su propia arquitectura y, a menudo, adaptados a cargas de trabajo de IA específicas. El software está cambiando con la misma rapidez y las herramientas de codificación de IA ahora generan en minutos lo que hace unos años requería meses de esfuerzo.

Ahora que gran parte de la informática se centra en la IA, los núcleos de GPU son un componente crucial de su éxito. Estos son los programas de bajo nivel que se ejecutan dentro de la GPU, y escribir programas eficientes no es nada obvio: se necesitan años de experiencia para hacerlo bien. Transferir un kernel del hardware de un proveedor 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 kernel ganada con esfuerzo: implementaciones de atención ajustadas manualmente, modelos de espacio de estados 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ápidamente, pero carecen de esta profundidad.

En este trabajo nos preguntamos si esa experiencia se puede transferir automáticamente. Nos basamos en K-Search, un marco de búsqueda de kernel evolutivo introducido por Cao et al. en Berkeley Sky Lab que utiliza IA para optimizar los núcleos de GPU y lo amplió con un backend para MLX, el marco de aprendizaje automático de Apple para sus propios chips Apple Silicon. Desarrollamos una nueva capa de traducción estructurada de CUDA a MLX que permite a K-Search tomar los núcleos CUDA existentes como base de conocimientos y adaptarlos a núcleos de GPU de alta calidad para Apple Silicon, en lugar de reconstruirlos desde cero.

Mostramos que nuestro enfoque alcanza un rendimiento de nivel casi experto en Apple Silicon con una aceleración de 0,97 veces en comparación con el kernel nativo MLX Attention y una aceleración de precarga de hasta 20 veces sobre la implementación comunitaria mlx-lm en el kernel Mamba SSM; Informamos los números y la cantidad de ganancia que proviene de la capa de traducción en las secciones siguientes. Aunque nos centramos en los kernels MLX para Apple Silicon, el método no es específico de MLX y se aplica a cualquier ecosistema donde la experiencia CUDA sea transferible.

¿Por qué MLX?

El marco MLX de Apple ha experimentado una adopción notable desde finales de 2023. Con Apple Silicon 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 la hace especialmente atractiva para modelos de tamaño mediano (parámetros 7B–70B en chips de la serie M).

Sin embargo, debajo de este impulso se esconde una brecha significativa: muchos núcleos críticos para el rendimiento que el ecosistema NVIDIA da por sentado: atención paginada, núcleos de escaneo SSM optimizados, enrutamiento MoE fusionado están ausentes o son ingenuos sin un ajuste específico del hardware. MLX ejecuta 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 marco de optimización del kernel evolutivo desarrollado originalmente por nuestro primer autor, Shiyi Cao, en UC Berkeley Sky Lab. Dado un kernel ingenuo y una especificación de hardware, ejecuta un ciclo de optimización iterativo: un LLM razona qué optimizaciones probar a continuación, un modelo de escritura de código genera kernels candidatos, y esos candidatos se compilan y comparan en hardware real.

Las mediciones retroalimentan la búsqueda, que sigue refinándose, siguiendo direcciones prometedoras y abandonando callejones sin salida hasta que el rendimiento converge.

Algoritmo 1: K-Search a través de modelos mundiales en coevolución. La búsqueda alterna entre seleccionar la acción más prometedora, crear instancias y evaluar el código hasta que la mejora se estanque y hacer evolucionar el modelo mundial mediante operaciones de inserción, actualización y poda. Adaptado de Cao et al. (2026).

La búsqueda se basa en una especificación: un documento específico de dominio que codifica reglas de hardware, patrones de optimización y restricciones matemáticas que evita que el código generado alucine con primitivas no válidas y garantiza que los candidatos realmente se compilarán y ejecutarán de manera eficiente.

En nuestras ejecuciones, un solo modelo (Gemini 3.5 Pro Preview) desempeña ambas funciones: mantiene el estado de razonamiento y escribe los núcleos. A la mitad razonadora se le solicita que sea un “ingeniero de rendimiento del kernel de GPU” y se le pide que trabaje en un análisis fijo antes de proponer cualquier cosa: clasificar el kernel (reducción, escaneo, atención/softmax,…), reescribir el cálculo de referencia en forma canónica, mapear el diseño de los datos y los patrones de acceso, y formular hipótesis sobre el probable cuello de botella (ancho de banda, latencia, computación o sincronización) en cada régimen de tiempo de ejecución. Sólo entonces emite optimizaciones candidatas, cada una como un único cambio implementable en una iteración.

Al estado de razonamiento persistente lo llamamos modelo mundial. En lugar de una lista plana de cosas para probar, es un árbol de decisión (prefijo): cada ruta raíz→hoja compone un plan de optimización completo, y las ramas hermanas son alternativas competitivas. Se califica cada nodo (una calificación general en [0, 10], una confianza en [0, 1] y los impactos por nodo en el ancho de banda de la memoria, la presión del registro y el ajuste de computación/hardware) para que la búsqueda pueda clasificar los planes parciales y expandir los más prometedores. El árbol persiste y crece a lo largo de las rondas: refinar una idea agrega un nodo hijo en lugar de sobrescribir su padre, y si la mejor puntuación no mejora durante algunas rondas (una ventana de estancamiento), la búsqueda retrocede para explorar una rama alternativa. Un solo nodo, tal como aparece en mitad de la ejecución en el núcleo de atención, se ve así:

Listado 1: Ejemplo de nodo de modelo mundial K-Search. Cada optimización candidata registra una acción concreta, impactos estimados en el hardware, una calificación de prioridad general y la confianza del modelo.

Figura 2: Descripción general de K-Search. El marco opera en un estado de búsqueda $S_t$ estructurado como un árbol de búsqueda. El árbol consta de nodos cerrados (azules, estados visitados con un programa adjunto como $x_{12}$) y una frontera de nodos abiertos (naranja, hipótesis pendientes como $u_{13}$). El flujo de trabajo recorre tres fases: (1) Selección de acción, donde el nodo de acción más prometedor se recupera de la frontera según el puntaje de prioridad estimado del modelo mundial $V$; (2) Refinamiento local, donde una política estocástica $pi_{mathrm{code}}$ muestra implementaciones concretas hasta el estancamiento; y (3) Actualización del modelo mundial, donde el LLM razona sobre la trayectoria para actualizar el árbol de búsqueda mediante Insertar (agregando nuevas acciones), Actualizar (ajustando $V$, por ejemplo, $u_{11}$ cayendo de 0,9 a 0,6) y Podar (eliminando nodos menos prometedores como $u_{10}$).

El artículo original de K-Search evaluó esta estrategia de búsqueda en núcleos CUDA de FlashInfer. En términos de decodificación GQA, decodificación MLA, precarga MLA y MoE, K-Search mejoró de manera más consistente que OpenEvolve y ShinkaEvolve con el mismo presupuesto de 120 iteraciones. Estos resultados establecen el marco de búsqueda que construimos aquí; El resto de esta publicación pregunta si su conocimiento de optimización puede transferirse más allá de CUDA.

Figura 3: Principales resultados del artículo original de K-Search. En tres ejecuciones, K-Search logra los mejores puntajes de búsqueda hasta ahora, rendimiento del kernel por carga de trabajo y distribuciones más rápidas que OpenEvolve y ShinkaEvolve en cuatro kernels FlashInfer CUDA. Reproducido exactamente de Cao et al. (2026).

Construyendo un backend MLX

Para llevar K-Search a Apple Silicon, primero creamos un backend MLX nativo. Implementamos un adaptador de tareas completo específico de MLX para K-Search, que incluye:

Un backend de tareas de MLX en k_search/tasks/ que maneja la compilación y ejecución del kernel en Apple Silicon a través de las API Metal/C++ de MLX. Solicitudes actualizadas del generador de kernel para escribir y modificar kernels Metal/MLX. Integración de evaluaciones comparativas específicas de MLX utilizando las utilidades de medición mlx.core.

Traduciendo la experiencia de CUDA a MLX

Sin embargo, el desafío más interesante no fue simplemente ejecutar K-Search en MLX. La idea clave es que los núcleos CUDA expertos codifican décadas de conocimiento de optimización que es transferible a la GPU de Apple si se puede cerrar la brecha conceptual. Simplemente entregarle a un LLM un kernel CUDA y pedirle que lo porte no es suficiente: sin un contexto de hardware profundo, produce código que es sintácticamente válido pero arquitectónicamente incorrecto (tamaños de mosaico incorrectos, primitivas no válidas, suposiciones de memoria que no coinciden).

Nuestra capa de traducción consta de:

Tablas de mapas conceptuales: un glosario estructurado de primitivas CUDA y sus equivalentes MLX/Metal con restricciones estrictas. Por ejemplo: __shared__ se asigna a la memoria del grupo de subprocesos Metal pero con un límite estricto de 32 KB (frente a los 48 KB de NVIDIA) warp_reduce se asigna a MMA (preferido) __syncthreads() se convierte en threadgroup_barrier(mem_flags::mem_tg) HBM3 de ~3,35 TB/s de H100 se asigna a la DRAM unificada de ~400 GB/s de M3 Max, una diferencia de ancho de banda eso remodela qué optimizaciones vale la pena realizar. Sugerencias y patrones específicos de MLX: Patrones concretos a nivel de código para operaciones sin equivalente CUDA directo, como reducciones de filas basadas en registros usando simd_shuffle_xor en un diseño de mosaico MMA de 8×8, o el “truco exp2” (reemplazando $exp(x)$ con $exp_2(x log_2 e)$) para un softmax más rápido en la rápida instrucción de hardware $exp_2$ de Apple. Afirmaciones reutilizables: comportamientos expertos del kernel replanteados como propiedades que la búsqueda evolutiva debe preservar, en lugar de código para copiar.

Rendimiento del kernel experto coincidente: el kernel de atención

Evaluamos tres configuraciones de un núcleo de atención MLX para Apple Silicon: (1) una línea de base ingenua, (2) evolución pura sin contexto adicional proporcionado y (3) una capa de traducción de contexto completo, que proporciona al optimizador conocimiento de implementación específico de la arquitectura extraído de núcleos de alto rendimiento (por ejemplo, FlashAttention-2), permitiendo que la búsqueda evolutiva razone sobre las estrategias de implementación en lugar de comenzar desde un núcleo ingenuo. Juntas, estas tres configuraciones nos permiten aislar el impacto exacto de la capa de traducción.

Figura 4: Escalado del rendimiento de Attention Kernel mediante optimizaciones apiladas. La configuración de "Contexto completo" descubre e implementa con éxito estrategias avanzadas como el doble almacenamiento en búfer y el desenrollado de bucles, logrando un rendimiento casi experto.

El salto de 0,26× a 0,97× la velocidad del núcleo de atención de última generación de Apple ilustra lo importante que es la capa de traducción. Con contexto completo, el kernel evolucionado descubre de forma independiente las optimizaciones clave en FlashAttention 2: mosaico de memoria de grupo de subprocesos, softmax en línea, transposición K para acceso a memoria y el truco exp2. El último de ellos reemplaza cada exponencial softmax con un exponencial de base 2,

[e^x = 2^{x log_2 e},]

lo cual es exacto y permite que el kernel use la instrucción de hardware fast fast::exp2() de Apple directamente en lugar de pagar por una conversión base en tiempo de ejecución.

Un precarga 20 veces más rápido: el kernel Mamba SSM

Para evaluar si K-Search se generaliza más allá de los núcleos de atención, lo aplicamos al núcleo del modelo de espacio de estados (SSM) utilizado por Mamba. A diferencia de la atención, el cuello de botella computacional es una actualización de estado recurrente en lugar de un softmax, lo que genera un desafío de optimización sustancialmente diferente. Comparamos la implementación evolucionada con la implementación comunitaria de MLX (mlx-lm) y la implementación de referencia de PyTorch (mamba.py) en un M1 Max.

Evaluado en mamba-370m f16, M1 Max 64GB:

Tabla 1: Rendimiento de precarga y decodificación en mamba-370m (f16, M1 Max 64 GB). mlx-mamba (el nuestro) alcanza un rendimiento de precarga ~20 veces mayor que la línea base de mlx-lm de la comunidad, mientras que la decodificación sigue siendo comparable.

La aceleración de ~20× de precarga sobre mlx-lm se reduce a una diferencia: mlx-lm no implementa un escaneo paralelo para el SSM. La recurrencia del estado

[h_t = bar{a}_t h_{t-1} + bar{b}_t]

parece inherentemente secuencial, pero cada paso se puede escribir como un par $(bar{a}_t, bar{b}_t)$ bajo la combinación asociativa

[(a_2, b_2) circ (a_1, b_1) = left(a_2 a_1, a_2 b_1 + b_2right),]

que reproduce la recurrencia exactamente. Debido a que el operador es asociativo, toda la secuencia se puede evaluar con un escaneo paralelo (prefijo) en $O(log N)$ pasos dependientes en lugar de $O(N)$. mlx-lm omite esto y procesa los tokens uno a la vez, dejando la mayor parte de la computación de Apple Silicon inactiva; Nuestro kernel Metal evolucionado aplica el escaneo y hace un uso mucho más completo del rendimiento de la GPU. La ganancia aparece en el prerrelleno, donde la secuencia completa está disponible para escanear en paralelo, y no en la decodificación de un solo token, donde solo hay un token nuevo por paso y no hay escaneo para paralelizar, razón por la cual la fila de decodificación es más o menos plana mientras que el prellenado es ~20×.

mamba.py es lento tanto en el llenado previo como en la decodificación porque es una implementación de referencia de PyTorch que recurre a la CPU o MPS en Apple Silicon, renunciando a las optimizaciones específicas de hardware que hace posible el backend Metal de MLX.

¿Qué sigue?

En los dos núcleos que estudiamos, la búsqueda de núcleos evolutivos impulsada por IA basada en conocimientos estructurados de traducción multiplataforma alcanzó un rendimiento casi experto en Apple Silicon sin un equipo de expertos en GPU que comenzara desde cero. Aún no sabemos hasta qué punto esto se generaliza, pero el resultado es alentador.

Para nosotros, la principal conclusión es que el cuello de botella no fue la capacidad del LLM para escribir código Metal, sino la calidad del contexto y las limitaciones que le asignamos. Nuestra capa de traducción CUDA convierte la experiencia existente en el kernel de NVIDIA en una guía práctica para Apple Silicon y permite que la búsqueda evolutiva de K-Search haga el resto.

Estamos ampliando activamente este trabajo en varias direcciones: soporte de nuevas arquitecturas, con esfuerzos actuales centrados en desarrollar nuevos núcleos para IBM Spyre AIU y objetivos de hardware más amplios; agregar más núcleos, como atención paginada y enrutamiento MoE fusionado; y mejorar la integración con el ciclo de evolución de K-Search para hacer que el contexto de traducción sea aún más automático.

Expresiones de gratitud

Este trabajo fue realizado por IBM Research y se basa en K-Search del UC Berkeley Sky Lab (Cao et al., 2026). Agradecemos la colaboración y los comentarios de MLX y de comunidades más amplias de sistemas de IA. Si está trabajando en la optimización del kernel para hardware que no es CUDA, nos encantaría saber de usted.

Citación