ML Kernel Performance Engineer, AWS Neuron, Annapurna Labs

The Annapurna Labs team at Amazon Web Services (AWS) builds AWS Neuron, the software development kit used to accelerate deep learning and GenAI workloads on Amazon’s custom machine learning accelerators, Inferentia and Trainium. The Acceleration Kernel Library team focuses on maximizing performance for AWS's custom ML accelerators, crafting high-performance kernels for ML functions to deliver optimal performance for customers’ demanding workloads.

The AWS Neuron SDK, developed by the Annapurna Labs team, is the backbone for accelerating deep learning and GenAI workloads on Amazon's Inferentia and Trainium ML accelerators. It includes an ML compiler, runtime, and application framework that seamlessly integrates with popular ML frameworks like PyTorch, providing unparalleled inference and training performance.

As part of the broader Neuron Compiler organization, our team works across multiple technology layers—from frameworks and compilers to runtime and collectives...