Machine Learning - Compiler Engineer , AWS Neuron, Annapurna Labs

This role is for a software engineer in the Compiler team for AWS Neuron. As part of this role, you will be responsible for building next generation Neuron compiler which transforms ML models written in ML frameworks (e.g, PyTorch, TensorFlow, and JAX) to be deployed AWS Inferentia and Trainium based servers in the Amazon cloud. You will be responsible for solving hard compiler optimization problems to achieve optimum performance for variety of ML model families including massive scale large language models like Llama, Deepseek, and beyond as well as stable diffusion, vision transformers and multi-model models.

Design, implement, test, deploy and maintain innovative software solutions to transform Neuron compiler’s performance, stability and user-interface

Work side by side with chip architects, runtime/OS engineers, scientists and ML Apps teams to seamlessly deploy state of the art ML models on AWS accelerators

Work with open-source software (e.g., StableHLO, OpenXLA, MLIR) to pioneer optimizing advanced ML workloads on AWS software and hardware

Create compiler optimization and verification passes and build features to surface peculiarities of AWS accelerators to developers

Implement tools to analyze numerical errors and resolve the root cause of compiler defects

Participate in design discussions, code review, and communicate with internal and external stakeholders