Software Development Engineer, AI/ML, AWS Neuron, Model Inference

What you'll do

• Build distributed inference support and optimize model performance and efficiency on Trainium and Inferentia.
• Design and implement ML infrastructure, high-performance kernels, and system-level optimizations; profile and test workloads.
• Collaborate with hardware, compiler, runtime, and framework teams and customers to enable and tune models.

What they ask for

• Bachelor's degree in computer science or equivalent.
• 3+ years of non-internship professional software development experience and 3+ years of non-internship design or architecture experience.
• Fundamentals of machine learning and LLMs, including architecture, training and inference lifecycles, and experience with model execution optimizations.
• Software development experience in C++ or Python; strong understanding of system performance, memory management, and parallel computing.

Nice to have

• Familiarity with PyTorch, JIT compilation, and AOT tracing.
• Experience with kernel development or production inference serving platforms such as vLLM, SGLang, or TensorRT.

Pay

165,200.00 - 223,600.00 USD annually (base salary)

Tools

AWS Neuron SDK, Trainium, Inferentia, C++, Python, PyTorch, JAX, vLLM

AI focus

Enable and accelerate distributed LLM inference on custom AWS ML accelerators, including model onboarding, kernel development, and latency and throughput optimization.