Machine Learning - Compiler Engineer , AWS Neuron, Annapurna Labs

Amazon Inc.; Cupertino (CA); Hybrid; USD 129,300 - 223,600; Full time

Benefits offered by this job

Competitive salary

Health insurance

Career growth opportunities

Job summary

A leading tech company in Seattle is seeking a software engineer for its Compiler team. This role focuses on building the next-generation Neuron compiler, enhancing performance for ML models on AWS accelerators. Candidates should have 3+ years in software development, preferably with C++ or Java, and experience in compiler design is a plus. The position offers competitive compensation and a dynamic work environment.

Qualifications

• 3+ years of non-internship professional software development experience.
• 2+ years of non-internship design or architecture experience.
• Experience programming with at least one software programming language.

Responsibilities

• Design, implement, test, deploy, and maintain software solutions.
• Work closely with various teams to deploy ML models on AWS accelerators.
• Participate in building and improving open-source software.

Skills

C++

Java

Compiler design

Machine Learning

Collaboration

Education

Master’s degree or PhD in Computer Science or a related technical field

Tools

LLVM

MLIR

PyTorch

OpenXLA

Job description

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, TensorFlow, 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, StableHLO, and MLIR is a bonus.

Explore the product and our history: https://awsdocs-neuron.readthedocs-hosted.com/en/latest/neuron-guide/neuron-cc/index.html

https://aws.amazon.com/machine-learning/neuron/ https://github.com/aws/aws-neuron-sdk https://www.amazon.science/how-silicon-innovation-became-the-secret-sauce-behind-awss-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.

Key job responsibilities

• 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., StableHLO, 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.

A day in the life

• 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.

About the team

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.

Basic Qualifications

• 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.

Preferred Qualifications

• Master’s degree or PhD in Computer Science or a related technical field.
• 3+ years of experience writing production‑grade code in object‑oriented languages such as C++/Java.
• Experience in compiler design for CPU/GPU/Vector engines/ML‑accelerators.
• Experience with the open‑source compiler toolset like LLVM/MLIR.
• Experience with technologies such as PyTorch, OpenXLA, StableHLO, JAX, TVM, deep‑learning models, and algorithms.
• Experience with modern build systems like Bazel/CMake.

Amazon is an equal opportunity employer and does not discriminate on the basis of protected veteran status, disability, or other legally protected status.

Our compensation reflects the cost of labor across several U.S. geographic markets. The base pay for this position ranges from $129,300/year in the lowest market up to $223,600/year in the highest market. Pay is based on a number of factors including market location and may vary depending on job‑related knowledge, skills, and experience. Amazon may provide equity, sign‑on payments, and other compensation as part of a total package, along with medical, financial, and other benefits. This position will remain posted until filled. Applicants should apply via our internal or external career site.