Machine Learning - Compiler Engineer , AWS Neuron, Annapurna Labs

Do you want to be part of the AI revolution? At AWS our vision is to make deep learning pervasive for everyday developers and to democratize access to AI hardware and software infrastructure. In order to deliver on that vision, we’ve created innovative software and hardware solutions that make it possible. AWS Neuron is the SDK that optimizes the performance of complex ML models executed on AWS Inferentia and Trainium, our custom chips designed to accelerate deep‑learning workloads.

This role is for a software engineer in the Compiler team for AWS Neuron. You will be responsible for building the next‑generation Neuron compiler that transforms ML models written in frameworks such as PyTorch, Tensor Flow, and JAX to run on AWS Inferentia and Trainium servers in the Amazon cloud. Your work will solve difficult compiler‑optimization problems to achieve optimum performance for a wide variety of ML model families, including large‑scale language models like Llama and Deepseek, stable diffusion models, vision transformers, and multi‑model architectures.

You will need to understand the internals of these models deeply to make informed compiler‑generation decisions and will partner in technical communication with internal and external stakeholders. The role also includes involvement in pre‑silicon design, bringing new products and features to market, and ultimately making the Neuron compiler both highly performant and easy to use.

Required experience includes proficiency in object‑oriented languages such as C++ or Java.

Experience with compilers, building ML models on accelerators (e.g., GPUs), or technologies like OpenXLA, Stable

HLO, and MLIR is a bonus.

Explore the product and our history: -

 https://(Use the "Apply for this Job" box below).-success

AWS Utility Computing (UC) provides product innovations—from foundational services such as Amazon S3 and EC2 to new product innovations that set AWS’s services apart in the industry. As a member of the UC organization, you’ll support the development and management of Compute, Database, Storage, IoT, Platform, and Productivity Apps services in AWS, including support for customers who require specialized security solutions.

• Design, implement, test, deploy, and maintain innovative software solutions to enhance Neuron compiler performance, stability, and user interface.
• Work closely with chip architects, runtime/OS engineers, scientists, and ML Apps teams to deploy state‑of‑the‑art ML models on AWS accelerators with optimal cost/performance.
• Participate in building and improving open‑source software (e.g., Stable
  
  HLO, OpenXLA, MLIR) to pioneer optimization of advanced ML workloads on AWS software and hardware.
• Build innovative features that deliver the best possible experience for developers worldwide.

• Design and code solutions to drive efficiencies in compiler architecture, creating optimization and verification passes.
• Build features that surface AWS accelerator capabilities to developers and implement tools to analyze numerical errors.
• Resolve root causes of compiler defects and participate in design discussions and code reviews.
• Communicate with internal and external stakeholders, including open‑source communities.
• Work in a startup‑like development environment focused on high‑impact work.

Our team is dedicated to supporting new members. We have a broad mix of experience levels and tenures and are building an environment that celebrates knowledge‑sharing and mentorship. Senior members provide one‑on‑one mentoring and thorough, but kind, code reviews. We care about career growth and strive to assign projects that help team members develop expertise and feel empowered to take on more complex tasks in the future.

• 3+ years of non‑internship professional software development experience.
• 2+ years of non‑internship design or architecture experience (design patterns, reliability, and scaling) of new and existing systems.
• Experience programming with at least one software programming language.

• Master’s degree or PhD in Computer…