O problema que todo time de ML conhece bem
À medida que os modelos de IA crescem em escala e complexidade, uma tensão antiga volta à tona: o hardware tem um potencial teórico enorme, mas pouquíssimas equipes conseguem extrair esse potencial na prática. O desenvolvimento de kernels customizados sempre foi o caminho para fechar essa lacuna — mas esse caminho exige conhecimento profundo de arquitetura de chips, ciclos intermináveis de perfilamento manual e uma curva de aprendizado que poucos times podem se dar ao luxo de percorrer.
A AWS decidiu atacar esse problema de frente. A empresa anunciou o Neuron Agentic Development: um conjunto de agentes e habilidades de IA que torna o desenvolvimento de kernels de alta performance acessível para qualquer engenheiro de Machine Learning (ML) — não apenas para os especialistas em hardware. As capacidades estão disponíveis para desenvolvedores que trabalham com AWS Trainium e AWS Inferentia.
A proposta é ambiciosa: e se qualquer engenheiro de ML pudesse atuar como engenheiro de performance, escrevendo kernels com consciência de hardware, diagnosticando gargalos e entregando modelos otimizados — sem precisar de anos de experiência em nível de chip? E se um desenvolvedor experiente em outra arquitetura pudesse se adaptar ao Trainium em dias, não em meses?
O que é o Neuron Agentic Development
O pacote Neuron Agentic Development oferece cinco habilidades especializadas que seguem o pipeline natural de desenvolvimento de kernels: escrever → depurar → perfilar → analisar. Cada habilidade pode ser invocada individualmente para tarefas específicas, ou encadeada por meio do neuron-nki-agent, que seleciona automaticamente o fluxo de trabalho adequado com base na solicitação do desenvolvedor.
Para utilizá-las, basta adicionar as habilidades ao diretório de skills do seu Ambiente de Desenvolvimento Integrado (IDE) agêntico. Em IDEs como VS Code, Cursor ou Kiro, isso significa colocar os arquivos nos diretórios .kiro/skills ou .claude/skills. Vale destacar: as habilidades de perfilamento e depuração precisam ser executadas em uma instância do Amazon Elastic Compute Cloud (Amazon EC2) baseada em Trainium. As habilidades de escrita e documentação funcionam em qualquer ambiente.
Autoria de kernels
A habilidade neuron-nki-writing é o ponto de partida para criar kernels usando a Interface de Kernel Neuron (NKI). Ela traduz código PyTorch, NumPy ou descrições em linguagem natural para código NKI correto, já respeitando as restrições de hardware — como a dimensão de partição de 128 elementos e a dimensão livre de 512/4096 elementos no PSUM. A habilidade classifica a tarefa por complexidade e carrega apenas as referências necessárias para cada caso.
Depuração
A habilidade neuron-nki-debugging oferece um fluxo sistemático para resolver erros de compilação e execução de kernels NKI no Trainium e no Inferentia. Ela cobre desde a configuração do ambiente com as flags --target corretas até a resolução de erros do compilador, com um índice categorizado dos 28 códigos de erro do Compilador Neuron (NCC), além de validação numérica contra referências calculadas em CPU.
Perfilamento e análise
A habilidade neuron-nki-profiling captura perfis de execução diretamente no hardware. Ela configura as variáveis de ambiente de inspeção em tempo de execução, executa o kernel, identifica o Formato de Arquivo de Execução Neuron (NEFF) correto e captura o trace com o neuron-explorer — incluindo notificações do Motor de Grafo de Acesso Direto à Memória (DGE) para detalhes em nível de DMA. O resultado são arquivos JSON de métricas e arquivos NEFF que alimentam a próxima habilidade.
A habilidade neuron-nki-profile-querying, por sua vez, ingere os arquivos NEFF e NTFF e executa queries SQL para calcular limites de performance, identificar engines com gargalos e localizar ineficiências até linhas específicas do código-fonte NKI. Ela suporta três abordagens de análise: o servidor de API do neuron-explorer, DuckDB diretamente sobre arquivos parquet, ou pandas para computações customizadas.
Documentação
A habilidade neuron-nki-docs é utilizada em todas as fases do desenvolvimento. Durante a autoria, ela fornece assinaturas de API e tutoriais. Durante a depuração, explica códigos de erro. Durante o perfilamento, esclarece detalhes de arquitetura de hardware. É possível consultar APIs específicas como nisa.* ou nl.*, buscar tutoriais ou navegar pelos guias de arquitetura do Trainium 1, 2 e 3.
Os agentes: autonomia de ponta a ponta
Enquanto as habilidades oferecem blocos de construção para tarefas individuais, os agentes combinam múltiplas habilidades em fluxos de trabalho autônomos. Cada agente é uma persona especializada que gerencia cenários de desenvolvimento de múltiplas etapas do início ao fim.
neuron-nki-agent: ponto de entrada unificado para o desenvolvimento NKI. Seleciona automaticamente o fluxo correto (escrita, depuração, perfilamento ou documentação) e orquestra as habilidades adequadas. É o ponto de partida recomendado.neuron-nki-writing-agent: focado exclusivamente na autoria de kernels. Traduz PyTorch, NumPy ou linguagem natural para código NKI e também lida com modificações em kernels existentes.neuron-nki-debugging-agent: resolve erros do compilador de forma autônoma — analisa o erro, busca soluções na documentação e aplica correções. Rastreia até 10 iterações e simplifica progressivamente quando encontra dificuldades.neuron-nki-docs-agent: navegador leve de documentação para assinaturas de API, explicações de códigos de erro, tutoriais e detalhes de arquitetura.neuron-nki-profile-analysis-agent: executa duas habilidades em sequência para identificar gargalos de performance — primeiro captura o perfil de execução no hardware, depois executa queries SQL contra os arquivos de perfil para calcular limites, identificar engines com problemas e localizar ineficiências no código-fonte.
Na prática: otimizando um kernel softmax customizado
O anúncio da AWS inclui um passo a passo detalhado que demonstra como essas capacidades funcionam juntas. O exemplo explora dois kernels: um kernel softmax e um kernel SwiGLU — este último comum em Grandes Modelos de Linguagem (LLMs).
Passo 0: configurando o ambiente
O primeiro passo é iniciar uma instância trn2.3xlarge por meio dos Blocos de Capacidade de ML da AWS (AWS MLCBs), usando a Imagem de Machine Learning de Aprendizado Profundo Neuron da AWS (DLAMI). As regiões de São Paulo (sa-east-1) e Melbourne (ap-southeast-4) são citadas como exemplos no artigo original. Após conectar via SSH, instala-se o Kiro e as habilidades do Neuron Agentic Development seguindo as instruções no repositório neuron-agentic-development. Para instruções detalhadas de configuração, a AWS disponibiliza o Guia de Configuração do Neuron DLAMI.
A DLAMI já vem com um ambiente virtual pré-instalado. Para ativá-lo e verificar os dispositivos Neuron disponíveis:
# Confirm Neuron devices are visible
neuron-ls
# Confirm neuron-explorer is available
which neuron-explorer && neuron-explorer --versionsource ~opt/aws_neuronx_venv_pytorch_2_9/bin/activateCom o ambiente configurado, inicia-se o desenvolvimento de kernels com:
kiro-cli --agent neuron-nki-agentAtenção: instâncias trn2.3xlarge geram cobranças por hora enquanto estão em execução. Lembre-se de encerrá-las ao concluir o exercício.
Passo 1: escrevendo o kernel
Na sessão interativa do Kiro, o desenvolvedor descreve o que precisa em linguagem natural — por exemplo, solicitar um kernel NKI que compute um softmax com escala ao longo da última dimensão, para uma entrada em bfloat16. O agente produz um kernel completo de três passos (máximo por linha, soma de exponenciais, normalização), usando aceleração de hardware para a operação exponencial, acumulação em float32 para estabilidade numérica e tiling adequado. O agente também explica suas decisões de design, como o uso de P_MAX=128 (respeitando o limite de 128 partições do hardware) e F_MAX=2048 (respeitando o limite da dimensão livre no Trainium).
Passo 2: depurando no hardware
Ao solicitar que o agente execute o kernel e verifique a paridade numérica contra uma referência PyTorch, ele encontra um problema imediato: a operação nisa.tensor_tensor não faz broadcast automático de resultados de redução. O agente consulta os padrões de referência NKI, identifica o mecanismo correto de broadcast e reescreve o kernel. Após sincronizar o kernel corrigido e executá-lo no hardware, todos os casos de teste passam com erro máximo bem dentro da tolerância do bfloat16:
PASS: shape=(2, 128, 512), max_diff=0.000008
PASS: shape=(4, 256, 1024), max_diff=0.000004
PASS: shape=(1, 1, 64), max_diff=0.000061
PASS: shape=(2, 300, 768), max_diff=0.000007
All tests passed.Passos 3 e 4: perfilando e analisando a execução
Para demonstrar o perfilamento em uma carga de trabalho real, o exemplo usa um kernel SwiGLU MLP — módulo comum em LLMs. O agente compila o kernel para um NEFF, captura um trace NTFF via neuron-explorer e conduz uma investigação em duas etapas: primeiro analisa estatísticas em nível de kernel e limites de performance; depois investiga ineficiências específicas consultando o perfil em nível de execução de instruções.
A análise revela que o Motor de Tensores (TE) domina a execução e apresenta grandes lacunas de ociosidade — o que sugere uma dependência do Motor de Acesso Direto à Memória (DMA). O agente encontra que as instruções de matmul já operam próximo ao pico de eficiência, mas as instruções DMA estão bem abaixo do tamanho-alvo (ineficientes) e que todas as entradas estão sendo recarregadas oito vezes (redundante). O agente identifica até as três linhas exatas do código NKI responsáveis pelas transferências subótimas — informação diretamente acionável para a próxima rodada de otimizações.
O que vem a seguir
As habilidades de autoria e perfilamento de kernels são apenas o primeiro passo de uma visão mais ampla. Atualmente, os desenvolvedores usam os insights de perfilamento para guiar a próxima rodada de edições no kernel — um ciclo iterativo de perfilar, diagnosticar, refatorar e re-perfilar que consome a maior parte do tempo. A AWS quer tornar esse loop totalmente agêntico: agentes que iteram autonomamente em um kernel até atingir a meta de performance, sem exigir que o desenvolvedor interprete cada resultado e crie manualmente a próxima correção.
Além disso, a empresa reconhece que kernels customizados são apenas uma parte de um desafio maior. Os desenvolvedores também querem que seus modelos rodem no Trainium sem precisar se preocupar com portabilidade de código, lacunas de operadores, otimizações em nível de modelo e validação de correção em escala. A visão declarada é aplicar a mesma abordagem agêntica a esse problema mais amplo — desde experimentação com novas arquiteturas de modelos até execução de modelos em produção em escala.
Para acompanhar as novidades e começar a usar o que já está disponível, a AWS indica o repositório GitHub do Neuron Agentic Development.
Como começar
As capacidades do Neuron Agentic Development já estão disponíveis. O caminho sugerido pela AWS é:
- Clonar o repositório neuron-agentic-development e escrever o primeiro kernel NKI em minutos.
- Começar com o
neuron-nki-agent, que seleciona o fluxo correto com base na solicitação e oferece a experiência autônoma completa de ponta a ponta. - Executar os exemplos de habilidades individualmente — como
/neuron-nki-writingpara tarefas específicas, ou encadear/neuron-nki-profilinge/neuron-nki-profile-queryingapós o kernel estar produzindo resultados corretos. - Abrir uma issue no GitHub para reportar problemas ou sugerir melhorias — a AWS afirma estar desenvolvendo ativamente em conjunto com a comunidade.
- Contribuir com pull requests e compartilhar kernels construídos pela comunidade.
Limpeza do ambiente
Para evitar cobranças contínuas, é essencial encerrar a instância trn2.3xlarge criada durante o exercício. Isso pode ser feito pelo Console de Gerenciamento da AWS (EC2 > Instâncias, selecionar a instância e escolher Estado da instância > Encerrar) ou via linha de comando:
aws ec2 terminate-instances --instance-ids Confirme que o estado da instância mostra “terminated” antes de fechar o console.
Fonte
Stop hand-tuning kernels: How Neuron Agentic Development accelerates AWS Trainium optimizations (https://aws.amazon.com/blogs/machine-learning/stop-hand-tuning-kernels-how-neuron-agentic-development-accelerates-aws-trainium-optimizations/)
Leave a Reply