Uso de LLM locales para descubrir algoritmos de alto rendimiento

Desde pequeña me ha fascinado el dibujo. Lo que me llamó la atención no sólo fue el acto de dibujar en sí, sino también la idea de que cada dibujo se podía mejorar cada vez más. Recuerdo alcanzar niveles muy altos con mi estilo de dibujo. Sin embargo, una vez que alcanzaba la cima de la perfección, intentaba ver cómo podía mejorar aún más el dibujo, por desgracia, con resultados desastrosos.

A partir de ahí siempre tengo presente el mismo mantra: “refina e itera y alcanzarás la perfección”. En la universidad, mi enfoque fue leer libros muchas veces, ampliando mis conocimientos buscando otras fuentes, encontrando capas ocultas de significado en cada concepto. Hoy aplico esta misma filosofía a la IA/ML y la codificación.

Sabemos que la multiplicación de matrices (aquí matmul para simplificar) es la parte central de cualquier proceso de IA. En el pasado desarrollé LLM.rust, un espejo de Rust del LLM.c de Karpathy. El punto más difícil en la implementación de Rust ha sido la multiplicación de matrices. Dado que tenemos que realizar miles de iteraciones para ajustar un modelo basado en GPT, necesitamos una operación matmul eficiente. Para ello tuve que utilizar la biblioteca BLAS, implementando una estrategia insegura para superar los límites y barreras. El uso de inseguro en Rust va en contra de la filosofía de Rust, es por eso que siempre estoy buscando métodos más seguros para mejorar matmul en este contexto.

Entonces, inspirándome en la declaración de Sam Altman – “pregúntale a GPT cómo crear valor” – decidí pedirle a los LLM locales que generaran, compararan e iteraran sus propios algoritmos para crear una mejor implementación nativa de Rust matmul.

El desafío tiene algunas limitaciones:

Necesitamos utilizar nuestro entorno local. En mi caso, un MacBook Pro, M3, 36GB de RAM; Superar los límites de las fichas; Cronometrar y comparar el código dentro del propio ciclo de generación.

Sé que lograr rendimientos de nivel BLAS con este método es casi imposible, pero quiero resaltar cómo podemos aprovechar la IA para necesidades personalizadas, incluso con nuestras computadoras portátiles “pequeñas”, para que podamos desbloquear ideas y superar los límites en cualquier campo. Esta publicación quiere ser una inspiración para los profesionales y las personas que desean familiarizarse más con Microsoft Autogen y la implementación local de LLM.

Toda la implementación de cod se puede encontrar en este repositorio de Github. Este es un experimento en curso y se realizarán muchos cambios/mejoras.

idea general

La idea general es tener una mesa redonda de agentes. El punto de partida es el modelo local MrAderMacher Mixtral 8x7B modelo Q4 K_M. A partir del modelo creamos 5 entidades:

al proponente se le ocurre un nuevo algoritmo similar a Strassen, para encontrar una manera mejor y más eficiente de realizar matmul; el Verificador revisa la formulación matmul mediante matemáticas simbólicas; el Coder crea el código Rust subyacente; el Probador lo ejecuta y guarda toda la información en la base de datos de vectores; el Manager actúa silenciosamente, controlando el flujo de trabajo general.

Función AgentRoleProposerAnaliza tiempos de referencia y propone nuevos parámetros de ajuste y formulaciones matmul.Verificador (Actualmente deshabilitado en el código). Verifica la formulación matemática del proponente mediante verificación simbólica.CoderToma los parámetros y elabora el código de la plantilla de Rust.TesterEjecuta el código de Rust, guarda el código y calcula el tiempo de referencia.ManagerControl general del flujo de trabajo.
Pestaña. 1: Roles de los agentes.

El flujo de trabajo general se puede organizar a través de Microsoft Autogen como se muestra en la figura 1.

Fig.1: Optimización de Matmul. El usuario tiene una solicitud inicial con un mensaje. A partir de ahí, el administrador orquesta el flujo de trabajo general: 1) El proponente actúa como teórico y genera un algoritmo similar a Strassen; 2) El verificador verifica la corrección matemática del código; 3) El codificador genera un código Rust Neon; 4) El probador ejecuta el punto de referencia. [Image generated with Nano Banana Pro].

Prepare los datos de entrada y la base de datos vectorial.

Los datos de entrada se recopilan de todos los artículos académicos, centrados en la optimización de la multiplicación de matrices. Muchos de estos artículos tienen referencias y están relacionados con el artículo de Strassen de DeepMind. Quiero comenzar de manera simple, así que recopilé 50 artículos, publicados entre 2020 y 2025, que abordan específicamente la multiplicación de matrices.

A continuación, utilicé croma para crear la base de datos vectorial. El aspecto crítico al generar una nueva base de datos vectorial es cómo se fragmentan los archivos PDF. En este contexto, utilicé un fragmento semántico. A diferencia de los métodos de división de texto, el fragmentador semántico utiliza el significado real del texto para determinar dónde cortar. El objetivo es mantener las oraciones relacionadas juntas en un solo fragmento, haciendo que la base de datos vectorial final sea más coherente y precisa. Esto se realiza utilizando el modelo local BAAI/bge-base-en-v1.5. La esencia de Github a continuación muestra la implementación completa.

El código central: modelos autogen-core y GGML

He utilizado Microsoft Autogen, en particular la variante autogen-core (versión 0.7.5). A diferencia del chat de nivel superior, en autogen-core podemos tener acceso a bloques de construcción controlados por eventos de bajo nivel, que son necesarios para crear un flujo de trabajo controlado por una máquina de estados según lo necesitemos. De hecho, el desafío es mantener un flujo de trabajo estricto. Todos los agentes actuantes deben actuar en un orden específico: Proponente –> Verificador –> Codificador –> Probador.

La parte principal es BaseMatMulAgent, que hereda de RoutedAgent de AutoGen. Esta clase base nos permite estandarizar cómo los agentes de LLM participarán en el chat y se comportarán.

En el código anterior, podemos ver que la clase está diseñada para participar en un chat grupal asincrónico, manejando el historial de conversaciones, llamadas a herramientas externas y generando respuestas a través del LLM local.

El componente principal es @message_handler, un decorador que registra un método como oyente o suscriptor, según el tipo de mensaje. El decorador detecta automáticamente la sugerencia de tipo del argumento del primer método; en nuestro caso es mensaje: GroupChatMessage. Luego suscribe al agente para recibir cualquier evento de ese tipo enviado al tema del agente. El método asincrónico handle_message es responsable de actualizar la memoria interna del agente, sin generar una respuesta.

Una vez implementado el mecanismo oyente-suscriptor, podemos centrarnos en la clase Manager. MatMulManager hereda RoutedAgent y organiza el flujo general de agentes.

El código anterior maneja todos los agentes. Nos saltamos la parte del Verificador por el momento. El codificador publica el código final y el probador se encarga de guardar tanto el código como todo el contexto en la base de datos de vectores. De esta forma podremos evitar consumir todos los tokens de nuestro modelo local. En cada nueva ejecución, el modelo se pondrá al día con los últimos algoritmos generados a partir de la base de datos de vectores y propondrá una nueva solución.

Una advertencia muy importante: para asegurarse de que autogen-core pueda funcionar con modelos de llama en MacOS, utilice el siguiente fragmento:

#!/bin/bash CMAKE_ARGS=”-DGGML_METAL=on” FORCE_CMAKE=instalación de 1 pip –upgrade –verbose –force-reinstall llama-cpp-python –no-cache-dir

La figura 2 resume todo el código. Podemos subdividir aproximadamente el código en 3 bloques principales:

El BaseAgent, que maneja mensajes a través de los agentes de LLM, evaluando la formulación matemática y generando código; MatMulManager organiza todo el flujo de agentes; autogen_core.SingleThreadedAgentRuntime nos permite hacer realidad todo el flujo de trabajo.

Fig.2: Flujo de trabajo general en pocas palabras. El agente base ejecuta el LLM a través de agentes, evalúa la formulación matemática, crea el algoritmo en Rust y guarda toda la información en la base de datos vectorial. MatMulManager es el verdadero núcleo del flujo de trabajo general. Finalmente, autogen_core.SingleThreadedAgentRuntime hace que todo esto funcione en nuestra MacBook PRO. [Image created with Nano Banana Pro.]

Resultados y benchmark

Todo el código de Rust se revisó y se volvió a ejecutar manualmente. Si bien el flujo de trabajo es sólido, trabajar con LLM requiere un ojo crítico. Varias veces el modelo confabuló*, generando código que parecía optimizado pero que no lograba realizar el trabajo matmul real.

La primera iteración genera una especie de algoritmo tipo Strassen (código “Ejecutar 0” en la figura 3):

El modelo piensa en mejores implementaciones, más parecidas a Rust-NEON, de modo que después de 4 iteraciones da el siguiente código (“Ejecución 3” en la figura 3):

Podemos ver el uso de funciones como vaddq_f32, instrucción de CPU específica para procesadores ARM, procedente de std::arch::aarch64. El modelo logra usar rayón para dividir el flujo de trabajo en múltiples núcleos de CPU y, dentro de los subprocesos paralelos, usa elementos intrínsecos de NEON. El código en sí no es totalmente correcto; además, he notado que nos encontramos con un error de falta de memoria cuando trabajamos con matrices de 1024 × 1024. Tuve que reelaborar manualmente el código para que funcionara.

Esto nos lleva de vuelta a nuestro lema “iterar a la perfección”, y podemos preguntarnos: ‘¿puede un agente local refinar de forma autónoma el código Rust hasta el punto de dominar los complejos intrínsecos de NEON?’. Los hallazgos muestran que sí, incluso en hardware de consumo, este nivel de optimización es alcanzable.

La Fig.3 muestra los resultados finales que obtuve después de cada iteración.

Fig.3: Gráfico logarítmico de la implementación de Rust-Neon en varias iteraciones. Los cálculos se han realizado en puntos de referencia de multiplicación de matrices de 1024 × 1024. [Image generated by the author].

Los puntos de referencia 0.º y 2.º tienen algunos errores, ya que es físicamente imposible lograr tales resultados en un matmul de 1024×1024 en una CPU:

el primer código sufre de una falacia diagonal, por lo que el código calcula solo bloques diagonales de la matriz e ignora el resto; el segundo código tiene un búfer roto, ya que sobrescribe repetidamente un pequeño búfer de caché 1028 flotantes, en lugar de atravesar el millón de elementos completos.

Sin embargo, el código produjo dos códigos reales, la ejecución 1 y la ejecución 3. La primera iteración alcanza 760 ms y constituye una línea de base real. Sufre errores de caché y falta de vectorización SIMD. La ejecución 3 registra 359 ms, la mejora es la implementación del paralelismo NEON SIMD y Rayon.

*: Escribí “la modelo confabula” sobre propósitos. Desde un punto de vista médico, todos los LLM no son alucinaciones, sino fabulaciones. Las alucinaciones son una situación totalmente diferente a lo que hacen los LLM cuando balbucean y generan respuestas “incorrectas”.

Conclusiones

Este experimento comenzó con una pregunta que parecía un desafío imposible: “¿podemos utilizar LLM locales de consumo para descubrir algoritmos Rust de alto rendimiento que puedan competir con la implementación de BLAS?”.

Podemos decir que sí, o al menos tenemos una experiencia válida y sólida, donde podemos desarrollar un mejor código para lograr un código completo similar a BLAS en Rust.

La publicación mostró cómo interactuar con Microsoft Autogen, autogen-core y cómo crear una mesa redonda de agentes.

El modelo base en uso proviene de GGUF y puede ejecutarse en una MacBook Pro M3 de 36 GB.

Por supuesto, no hemos encontrado (todavía) nada mejor que BLAS en un único código simple. Sin embargo, demostramos que el flujo de trabajo agente local, en una MacBook Pro, puede lograr lo que antes se pensaba que requería un clúster y modelos masivos. Finalmente, el modelo logró encontrar una implementación razonable de Rust-NEON, “Ejecutar 3 arriba”, que tiene una velocidad de más del 50% con respecto a la implementación estándar de Rayon. Debemos resaltar que la implementación del backbone fue generada por IA.

La frontera está abierta. Espero que esta publicación de blog pueda inspirarte a intentar ver qué límites podemos superar con la implementación local de LLM.

Estoy escribiendo esto a título personal; Estas opiniones son mías.