ML Compiler Engineer, Annapurna Labs

The AWS Neuron Compiler team is seeking skilled compiler engineers to develop a state-of-the-art deep learning compiler stack. This role involves optimizing application models across diverse domains, including Large Language and Vision, from frameworks like PyTorch, TensorFlow, and JAX. You'll work with custom-built Machine Learning accelerators such as Inferentia/Trainium, contributing to AWS innovation in advanced ML capabilities and Generative AI.

As an ML Compiler engineer, you'll design, develop, and optimize compiler features, tackling crucial challenges in compiler technology and deep-learning systems software. You'll collaborate with Runtime, Frameworks, and Hardware teams to ensure system-wide performance optimization.

Key responsibilities include:

• Solving challenging technical problems across all stack layers
• Designing, implementing, testing, deploying, and maintaining innovative software solutions
• Researching implementations for optimal customer experiences
• Building high-impact solutions for a large customer base
• Participating in design discussions, code reviews, and stakeholder communication
• Driving business decisions with technical input
• Working in a startup-like environment on high-priority projects

The role offers opportunities for growth, knowledge-sharing, and mentorship within a diverse team. You'll contribute to AWS's suite of generative AI services and cutting-edge cloud computing offerings.

Required qualifications:

• B.S. or M.S. in computer science or related field
• Proficiency in C++ (preferred) or Python
• 3+ years of non-internship professional software development experience
• 2+ years of experience in compiler optimization, graph theory, hardware bring-up, FPGA placement and routing algorithms, or hardware resource management

Preferred qualifications:

• M.S. or Ph.D. in computer science or related field
• Strong knowledge in compiler design, instruction scheduling, memory allocation, data transfer optimization, graph partitioning, parallel programming, code generation, ISAs, hardware bring-up, and hardware-software co-design
• Experience with LLVM and/or MLIR
• Experience with simulation tools algorithm development
• Experience in TensorFlow, PyTorch, and/or JAX
• Experience in LLM, Vision, or other deep-learning models

This role offers competitive compensation and benefits, with base pay ranging from $129,300 to $223,600 per year depending on the geographic market.