Senior Neuron Runtime Software Engineer for ML Accelerators

Amazon; Cupertino (CA); On-site; USD 150,000 - 210,000; Full time

Benefits offered by this job

Health insurance (medical, dental, and

Vision and prescription coverage

Job summary

Amazon is seeking a Software Development Engineer for the Neuron Runtime Team to develop and maintain high-performance runtime libraries and drivers for ML applications and AI accelerators. You will design, deploy Neuron Runtime components and advance profiling to optimize workloads on Trainium and Inferentia devices.

You will also work on improving ML kernel and framework performance, manage the runtime life cycle, and collaborate with teams to expose key performance data through the C++

Qualifications

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

Responsibilities

• Develop and maintain high-performance runtime libraries and drivers for machine learning applications and AI accelerators.
• Design, develop, and deploy Neuron Runtime and other Neuron components.
• Enhance profiling capabilities to help internal and external customers optimize AI workloads across Trainium and Inferentia devices.
• Improve performance of ML kernels and ML frameworks.
• Manage the full development life cycle of the Neuron Runtime, ensuring scalability, reliability, and usability.
• Collaborate with cross-functional teams to ensure the C++ compiler exposes key performance information for customers to understand and optimize performance.
• Drive innovations to support multiple frameworks (e.g., PyTorch, JAX, XLA).

Skills

Software development experience

Design/architecture

Programming languages

Education

Bachelor’s degree in computer science

Job description

Amazon is seeking a Software Development Engineer for the Neuron Runtime Team to develop and maintain high-performance runtime libraries and drivers for ML applications and AI accelerators. You will design, deploy Neuron Runtime components and advance profiling to optimize workloads on Trainium and Inferentia devices.

You will also work on improving ML kernel and framework performance, manage the runtime life cycle, and collaborate with teams to expose key performance data through the C++