K-Search traduz otimização de kernels CUDA para Apple Silicon

Uma equipe da UC Berkeley e da IBM Research publicou em 29 de julho de 2026, no blog do Berkeley Artificial Intelligence Research (BAIR), um trabalho que tenta resolver um problema caro e pouco visível do ecossistema de IA: a transferência de conhecimento de otimização de kernels de GPU entre arquiteturas de hardware diferentes.

Kernels são os programas de baixo nível executados dentro da GPU. Escrevê-los de forma eficiente exige anos de especialização, e o resultado é fortemente amarrado à arquitetura para a qual foi escrito. O ecossistema CUDA, da NVIDIA, acumulou ao longo de mais de uma década implementações refinadas à mão de operações críticas — atenção, modelos de espaço de estados, roteamento de mistura de especialistas — que representam milhares de horas de engenharia. Ecossistemas mais novos, como o Apple Silicon, crescem rápido, mas não têm essa profundidade.

A pergunta do trabalho é direta: essa expertise pode ser transferida automaticamente? A resposta apresentada pelos autores — Shiyi Cao e Joseph E. Gonzalez, da UC Berkeley, com Gal Bloch, Assaf Toledo, Michael Factor e Gil Vernik, da IBM Research — é que sim, desde que a tradução seja conceitual e não literal.

Contexto: o que é o K-Search e por que o MLX

O K-Search é um framework de busca evolutiva para otimização de kernels, desenvolvido originalmente por Shiyi Cao no Sky Lab da UC Berkeley e descrito no artigo K-Search: LLM Kernel Generation via Co-Evolving Intrinsic World Model (arXiv:2602.19128, de Cao, Mao, Gonzalez e Stoica).

O funcionamento é iterativo. Dado um kernel ingênuo e uma especificação de hardware, um modelo de linguagem raciocina sobre quais otimizações tentar em seguida; um modelo gerador escreve os candidatos; e cada candidato é compilado e medido em hardware real. As medições realimentam a busca, que continua refinando, perseguindo direções promissoras e descartando becos sem saída até o desempenho convergir.

Nas execuções descritas, um único modelo — o Gemini 3.5 Pro Preview — desempenha os dois papéis. A metade responsável pelo raciocínio recebe o papel de “engenheiro de desempenho de kernels de GPU” e é obrigada a percorrer uma análise fixa antes de propor qualquer coisa: classificar o kernel (redução, varredura, atenção/softmax), reescrever a computação de referência em forma canônica, mapear layout e padrões de acesso a dados, e formular hipóteses sobre o gargalo provável — largura de banda, latência, computação ou sincronização.

O estado persistente desse raciocínio é chamado pelos autores de world model. Não é uma lista plana de ideias, mas uma árvore de decisão: cada caminho da raiz até uma folha compõe um plano completo de otimização, e ramos irmãos são alternativas concorrentes. Cada nó recebe notas — uma avaliação geral de 0 a 10, uma confiança de 0 a 1 e impactos estimados sobre largura de banda de memória, pressão de registradores e adequação à computação do hardware. Se a melhor pontuação não melhora por algumas rodadas, a busca recua e explora um ramo alternativo.

Quanto ao MLX: é o framework de aprendizado de máquina da Apple para seus próprios chips, com adoção crescente desde o fim de 2023. Com Apple Silicon presente em centenas de milhões de MacBooks e Mac Studios, o MLX viabiliza inferência local de IA sem custo de nuvem, e a arquitetura de memória unificada o torna especialmente interessante para modelos de porte médio — a faixa de 7 a 70 bilhões de parâmetros em chips da série M.

O problema é que muitos kernels críticos de desempenho que o ecossistema NVIDIA já toma como garantidos — atenção paginada, varredura otimizada de modelos de espaço de estados, roteamento fundido de mistura de especialistas — estão ausentes ou implementados de forma ingênua no MLX. Os modelos rodam corretamente, mas deixam desempenho na mesa.

Detalhes principais: a camada de tradução

Construir um backend MLX para o K-Search foi a parte mecânica do trabalho: adaptador de tarefa para compilação e execução via APIs Metal/C++ do MLX, prompts atualizados para geração de kernels em Metal e integração com os utilitários de medição do mlx.core.

O desafio interessante estava em outro lugar. Segundo os autores, simplesmente entregar um kernel CUDA a um modelo de linguagem e pedir a portabilidade não funciona: sem contexto profundo de hardware, o resultado é código sintaticamente válido mas arquiteturalmente errado — tamanhos de tile inadequados, primitivas inválidas, premissas de memória incompatíveis.

A camada de tradução construída pela equipe tem três componentes.

Tabelas de mapeamento de conceitos. Um glossário estruturado de primitivas CUDA e seus equivalentes em MLX/Metal, com restrições rígidas. Alguns exemplos citados no texto:

Esse último item é o mais instrutivo. Não se trata de tradução de sintaxe: a diferença de quase uma ordem de grandeza em largura de banda muda quais otimizações valem o esforço em cada plataforma.

Dicas e padrões específicos de MLX. Padrões concretos de código para operações sem equivalente direto em CUDA, como reduções de linha baseadas em registradores usando simd_shuffle_xor em um layout de tile MMA 8×8, ou o chamado “truque do exp2”.

Asserções reutilizáveis. Comportamentos de kernels de especialistas reformulados como propriedades que a busca evolutiva deve preservar — em vez de código a copiar.

Aspectos técnicos: os dois resultados

Kernel de atenção

Os autores avaliaram três configurações de um kernel de atenção MLX: uma linha de base ingênua, evolução pura sem contexto adicional e a camada de tradução completa, que fornece ao otimizador conhecimento de implementação extraído de kernels de alto desempenho como o FlashAttention-2.

O salto foi de 0,26× para 0,97× da velocidade do kernel de atenção nativo da Apple. Com contexto completo, o kernel evoluído descobriu de forma independente as otimizações centrais do FlashAttention-2: tiling em memória de threadgroup, softmax online, transposição de K para acesso à memória e o truque do exp2.

Esse último merece explicação. A identidade e^x = 2^(x · log₂ e) é exata. Substituir a exponencial natural pela exponencial de base 2 permite ao kernel usar diretamente a instrução de hardware fast::exp2() da Apple, em vez de pagar por uma conversão de base em tempo de execução. É um exemplo típico de otimização que só faz sentido quando se conhece o conjunto de instruções do alvo.

Kernel SSM do Mamba

Para testar generalização além da atenção, a equipe aplicou o método ao kernel de modelo de espaço de estados usado pelo Mamba. Aqui o gargalo não é a softmax, mas uma atualização recorrente de estado — um desafio de otimização substancialmente diferente.

Resultados medidos com o modelo mamba-370m em f16, em um M1 Max de 64 GB:

O ganho de aproximadamente 20× no prefill se explica por uma diferença única: a implementação da comunidade não usa varredura paralela para o SSM. A recorrência de estado hₜ = āₜ · hₜ₋₁ + b̄ₜ parece inerentemente sequencial, mas cada passo pode ser escrito como um par sob uma combinação associativa que reproduz a recorrência exatamente. Como o operador é associativo, a sequência inteira pode ser avaliada por uma varredura de prefixo paralela em O(log N) passos dependentes, em vez de O(N).

Repare que o ganho aparece apenas no prefill, quando a sequência inteira está disponível para ser varrida em paralelo — e não no decode de token único, onde há apenas um token novo por passo e nada a paralelizar. É por isso que a linha de decode permanece praticamente estável enquanto o prefill salta.

Impactos, possibilidades e limitações

O resultado mais interessante do trabalho talvez não seja o número de aceleração, e sim a conclusão dos próprios autores sobre onde estava o gargalo: não na capacidade do modelo de linguagem de escrever código Metal, mas na qualidade do contexto e das restrições fornecidas a ele.

Isso tem implicação prática ampla. Sugere que, para tarefas de engenharia altamente especializadas, o retorno de investimento está menos em modelos maiores e mais na estruturação do conhecimento de domínio que se entrega ao modelo.

Do ponto de vista de mercado, o trabalho aponta para uma redução da vantagem competitiva que o ecossistema CUDA acumulou. Se conhecimento de otimização pode ser traduzido semiautomaticamente para novos alvos, a barreira de entrada de fabricantes alternativos de aceleradores diminui. A equipe já declara estar estendendo o método para o IBM Spyre AIU.

As limitações precisam ser ditas com clareza, e os próprios autores as reconhecem. Foram avaliados dois kernels, atenção e SSM. Não se sabe até onde o método generaliza. O resultado de atenção é 0,97× o kernel nativo da Apple — ou seja, chega perto do desempenho especializado, mas não o supera. O ganho de 20× no Mamba é medido contra uma implementação da comunidade que simplesmente não implementa varredura paralela; é uma comparação legítima e relevante para quem usa aquela biblioteca, mas não equivale a superar uma implementação otimizada de referência. Todas as medições vêm de chips específicos (M1 Max e M3 Max), e o custo computacional da busca evolutiva em si — número de iterações, chamadas ao modelo — não é o foco do texto publicado.

Recomendações

• Equipes que fazem inferência local em Apple Silicon podem reproduzir os resultados: o backend MLX foi construído sobre o repositório aberto do K-Search, com scripts prontos para os kernels de atenção e de varredura seletiva do Mamba.
• Antes de assumir que uma biblioteca de inferência está otimizada, meça. A diferença de 20× no prefill do Mamba estava disponível para qualquer usuário do mlx-lm sem que houvesse indicação disso.
• Para quem avalia hardware alternativo a NVIDIA, vale acompanhar essa linha de trabalho: a maturidade do software passou a ser um item que pode ser acelerado, e não apenas esperado.

Conclusão

O trabalho mostra que otimizações acumuladas em um ecossistema de GPU podem ser convertidas em orientação acionável para outro, desde que traduzidas em nível conceitual e ancoradas nas restrições reais do hardware de destino. Em dois kernels, a busca evolutiva guiada por essa tradução chegou perto do desempenho de especialistas sem uma equipe dedicada partindo do zero.

O que vale observar a seguir é a generalização: a equipe declara estar trabalhando em atenção paginada, roteamento fundido de mistura de especialistas e suporte a novas arquiteturas. Se os ganhos se mantiverem em uma variedade maior de kernels e de alvos de hardware, a implicação para a diversidade do mercado de aceleradores de IA é significativa.

Fontes e referências

• BAIR Blog — From CUDA to MLX: How K-Search Brings Decades of Kernel Expertise to Apple Silicon, por Shiyi Cao, Gal Bloch, Assaf Toledo, Michael Factor, Gil Vernik e Joseph E. Gonzalez (29 de julho de 2026): bair.berkeley.edu
• Cao, S.; Mao, Z.; Gonzalez, J. E.; Stoica, I. — K-Search: LLM Kernel Generation via Co-Evolving Intrinsic World Model, arXiv:2602.19128 (2026): arxiv.org/abs/2602.19128
• Repositório oficial do K-Search: github.com/caoshiyi/K-Search
• Blog do BAIR — página inicial e arquivo: bair.berkeley.edu/blog