De CUDA para Apple Silicon: como o K-Search transfere décadas de conhecimento em kernels de GPU automaticamente

De CUDA para Apple Silicon: como o K-Search transfere décadas de conhecimento em kernels de GPU automaticamente

Escrever kernels eficientes de GPU é uma das tarefas mais difíceis da engenharia de software — exige anos de expertise e um conhecimento profundo do hardware. O ecossistema CUDA da NVIDIA acumulou décadas de otimizações manuais para operações críticas como atenção, modelos de espaço de estados (SSMs) e roteamento MoE. Agora, pesquisadores da UC Berkeley mostram que esse conhecimento pode ser transferido automaticamente para o Apple Silicon usando o framework K-Search.

O artigo, publicado no BAIR (Berkeley Artificial Intelligence Research), descreve uma extensão do K-Search — um framework evolutivo de otimização de kernels que usa IA — com um backend para o MLX, o framework de machine learning da Apple para seus chips próprios. A inovação está em uma camada de tradução estruturada CUDA→MLX que permite reaproveitar kernels CUDA existentes como base de conhecimento, adaptando-os para Apple Silicon em vez de recriar tudo do zero.

Por que isso importa agora

O MLX teve adoção notável desde o final de 2023. Com Apple Silicon em centenas de milhões de MacBooks e Mac Studios, ele permite inferência local de IA sem custos de nuvem. A arquitetura de memória unificada é especialmente atraente para modelos médios (7B–70B parâmetros nos chips série M).

Mas há uma lacuna importante: muitos kernels de desempenho crítico que o ecossistema NVIDIA trata como commodity — atenção paginada, kernels de scan SSM otimizados, roteamento MoE fundido — ou estão ausentes no MLX ou são implementações ingênuas sem tuning específico para hardware. O MLX roda modelos corretamente, mas frequentemente deixa desempenho significativo na mesa.

Como o K-Search funciona

O K-Search é um framework de otimização evolutiva: dado um kernel simples e uma especificação de hardware, ele roda um loop iterativo onde um LLM raciocina sobre quais otimizações tentar, um modelo de geração de código produz kernels candidatos, e os candidatos são compilados e testados em hardware real.

As medições alimentam a busca, que continua refinando — perseguindo direções promissoras e descartando becos sem saída até a convergência de desempenho. Nos experimentos, um único modelo (Gemini 3.5 Pro Preview) atua tanto no raciocínio quanto na geração de código.

O raciocínio é estruturado em um “modelo de mundo” — uma árvore de decisão onde cada caminho raiz→folha compõe um plano completo de otimização. O LLM classifica o kernel (redução, scan, atenção/softmax), reescreve a computação de referência em forma canônica, mapeia padrões de acesso a dados e identifica o gargalo provável (largura de banda, latência, computação ou sincronização). Só então propõe otimizações candidatas.

A busca é guiada por uma Spec: um documento específico do domínio que codifica regras de hardware, padrões de otimização e restrições matemáticas — evitando que o código gerado alucine primitivas inválidas.

Resultados: desempenho próximo ao de especialistas humanos

Os números impressionam. A abordagem alcançou:

• 0,97x de speedup comparado ao kernel nativo de Atenção do MLX — ou seja, desempenho praticamente equivalente ao código escrito por especialistas humanos
• Até 20x de speedup no prefill sobre a implementação da comunidade mlx-lm para o kernel SSM Mamba

Os ganhos vêm majoritariamente da camada de tradução CUDA→MLX, que não faz tradução literal instrução-por-instrução, mas sim um mapeamento conceitual: identifica qual padrão de otimização CUDA corresponde a qual estratégia nativa no MLX (veja a figura de tradução no artigo original).

Além do Apple Silicon

Embora o foco seja MLX para Apple Silicon, o método não é específico do MLX. Ele se aplica a qualquer ecossistema onde o conhecimento CUDA possa ser transferido — aceleradores customizados de IA, novos chips, até mesmo GPUs de outros fabricantes.

Com ferramentas de código de IA gerando em minutos o que antes levava meses, e hardware mudando rapidamente com uma variedade crescente de chips de diferentes fornecedores, a capacidade de transferir expertise de otimização entre arquiteturas se torna cada vez mais estratégica.

Para desenvolvedores brasileiros que trabalham com Apple Silicon — uma base que cresce com a popularidade dos MacBooks com chips M — o K-Search + MLX representa um caminho para obter desempenho de ponta em inferência local sem depender de GPUs NVIDIA na nuvem.