Machine Learning Performance Engineer, Annapurna Labs

Amazon is hiring a Machine Learning Performance Engineer, Annapurna Labs in Tel Aviv, Israel. Level rates it ; you can apply on Amazon.

Level's read on this job

Builds AI ●●●● is how Level rates this role: the job is building AI systems.
Tools and skills the posting cites include PyTorch, TensorFlow and JAX.
Level classifies it as a mid-level role.
Level lists 20,000+ other open roles at Amazon.
The role was posted within the past two weeks.

AI in this role

Profile and optimize deep learning and generative AI workloads on custom AWS machine learning accelerators.

PyTorchTensorFlowJAXAWS NeuronInferentiaTrainium

performance-engineeringmachine-learningkernel-developmentprofilingoptimization

The Annapurna Labs team at Amazon Web Services (AWS) builds AWS Neuron, the software development kit used to accelerate deep learning and generative AI workloads on Amazon's custom machine learning accelerators — Inferentia and Trainium. These chips power workloads for thousands of AWS customers, from large language model training runs to real-time inference serving billions of daily predictions.

We are building the first Neuron performance engineering team in Tel Aviv. As a Machine Learning Performance Engineer, you'll help shape the direction of this team from the ground up — profiling and optimizing workloads across the full ML software stack, writing high-performance kernels, and improving the Neuron SDK that external developers depend on. You'll work at the boundary between software and hardware, collaborating directly with compiler, runtime, and chip design engineers to close performance gaps customers care about.

The team is new and small, which means broad scope, direct ownership, and real influence over the technical direction we take. If you enjoy digging into performance bottlenecks and turning analysis into measurable wins, this role is for you.

Key job responsibilities
Design and implement high-performance compute kernels for ML operations, leveraging the Neuron architecture and programming models.
Profile ML workloads end-to-end to identify bottlenecks — memory, compute, or communication — and drive optimizations through to a measured improvement.
Enhance the programming model and tooling that kernel and model developers rely on, improving usability and debugging workflows.
Identify and drive optimization opportunities across the Neuron software stack (compiler, runtime, frameworks).
Document software designs, operational runbooks, and performance findings so the broader team can build on your work.

A day in the life
You might start your morning reviewing profiling data from a customer's large diffusion model training job, tracing a utilization gap back to a specific kernel. After a design discussion with compiler engineers about a new operator fusion strategy, you spend the afternoon writing and benchmarking a kernel prototype. Later, you review a teammate's pull request for a runtime optimization and share your findings in a short write-up for the broader Neuron organization. Your work directly translates into faster model execution and lower cost for AWS customers running ML workloads at scale.

About the team
The Neuron Performance Engineering team in Tel Aviv is part of Annapurna Labs within AWS. Our mission is to make sure every ML workload running on Inferentia and Trainium chips reaches its full performance potential. We partner closely with compiler, framework, and hardware teams across Annapurna Labs, and we work directly with AWS customers to understand their models and unblock their adoption.

We are a newly formed group which is part of the larger Neuron organization. If you want to shape a team's technical culture from its earliest days while working on problems that matter to the future of AI infrastructure, we'd love to hear from you.

Basic qualifications

- 3+ years of non-internship professional software development experience
- Knowledge of Python and/or C++ programming
- Knowledge of computer architecture, operating systems, and parallel computing
- Experience with PyTorch, TensorFlow, and/or JAX

Preferred qualifications

- Master's degree in Computer Science, Engineering, Mathematics, or a related field
- Experience optimizing performance for LLM, Vision, or other deep-learning models
- Experience with kernel writing or parallel programming (CUDA, Triton, CUTLASS, Pallas, Mojo, SIMD, MPI)
- Experience with compiler optimization or hardware-software co-design

Our inclusive culture empowers Amazonians to deliver the best results for our customers. If you have a disability and need a workplace accommodation or adjustment during the application and hiring process, including support for the interview or onboarding process, please visit https://amazon.jobs/content/en/how-we-hire/accommodations for more information. If the country/region you’re applying in isn’t listed, please contact your Recruiting Partner.

How we rate this

Machine Learning Performance Engineer, Annapurna Labs at Amazon rates 90 out of 100 for how much of the daily work is AI. That makes it Builds AI (AI Level 4 of 4). The level is about AI in the job, not seniority.

Classification

Builds AI. The job is building AI systems.

●●●● Builds AI80 to 100
●●●○ Works on AI60 to 79
●●○○ Uses AI40 to 59
●○○○ Little AI0 to 39

Levels come from how often the tools, models and workflows of the role are named in the posting itself. Read the description and count the AI tools, models and workflows it names.

Key facts

Machine Learning Performance Engineer, Annapurna Labs is an open job at Amazon.
Level rates this role AI Level 4 of 4: Builds AI.
The role is based in IL, Tel Aviv.
The job was posted on October 26, 2026.
Seniority: mid.
Applications are handled by Amazon. Apply on Amazon

Prepare for this job

Built from this posting and from similar jobs on Level.

Skills and AI tools this role asks for

Performance EngineeringMachine learningKernel DevelopmentProfilingOptimizationPyTorchTensorFlowJAX

Questions you could be asked

Tell me about a project where performance engineering was part of your work. What did you do?
Tell me about a project where machine learning was part of your work. What did you do?
Tell me about a project where kernel development was part of your work. What did you do?
Tell me about a project where profiling was part of your work. What did you do?
Tell me about a project where optimization was part of your work. What did you do?

Adapt your resume

List these exact terms on your resume: Performance Engineering, Machine learning, Kernel Development, Profiling, and Optimization. An applicant tracking system matches the wording, not the idea.
Attach one line of real, concrete experience to at least one of them. A tool named with nothing behind it rarely survives a human read.
Lead with what you built, trained or shipped. This role is judged on the AI system itself, not the tools around it.

Get your CV reviewed before you apply

Send your CV and the role you want next. Your review arrives by email within 3 days.

Get a free CV review Check your resume with the ATS resume checker