engenheiro de software sênior em aws trainium / nki - kernel

Sobre a vaga

Buscamos uma pessoa engenheira de software sênior para atuar com AWS Trainium e NKI (Neuron Kernel Interface). Você será responsável por desenvolver e otimizar kernels de machine learning de alto desempenho, trabalhando em um time que valoriza autonomia e impacto técnico.

Responsabilidades

• Projetar, implementar e otimizar kernels NKI para AWS Trainium e Inferentia.
• Colaborar com times de ML e infraestrutura para melhorar a performance de modelos.
• Analisar e otimizar o uso de memória (SBUF, PSUM, HBM) e de banda (DMA).
• Garantir a correta integração com o AWS Neuron SDK e compilador Neuron.
• Participar de code reviews e contribuir para a melhoria contínua do código.

Requisitos

• Experiência sólida com AWS Trainium, Inferentia e NKI.
• Conhecimento profundo de CUDA, Triton ou linguagens de programação de kernels similares.
• Familiaridade com o ecossistema AWS Neuron (SDK, compilador, runtime).
• Experiência com otimização de kernels para arquiteturas como Trn1 e Trn2.
• Conhecimento de formatos numéricos (FP32, BF16, FP8, INT8).
• Habilidade para trabalhar de forma autônoma e colaborativa.

Diferenciais

• Experiência com XLA e compiladores de ML.
• Conhecimento em RLHF e treinamento de modelos de linguagem.
• Contribuições em projetos open source de ML.

Benefícios

• Remuneração competitiva.
• Trabalho remoto.
• Ambiente flexível e colaborativo.