Senior ML Software Engineer, Data Plane

The MLIL DataPlane team is looking for a Senior Software Development Engineer to

own the design and implementation of our inference data plane. We build the

software that makes large models run efficiently on custom hardware - spanning

model execution, memory management, data movement, and serving integration.

Our work covers the full inference path: integrating serving engines with custom

hardware, developing high-performance compute kernels, enabling efficient data

movement, and driving models from early validation through production. We

operate at frontier scale with large distributed models.

This is a ground-up effort with rapidly evolving hardware and software. We need

a senior IC who can write and optimize low-level code for custom hardware,

validate model architectures end-to-end, build test and profiling

infrastructure, and drive performance across the stack.

Key job responsibilities

- Develop and optimize compute kernels for a custom ML accelerator architecture,

targeting production-level performance for large language model inference.

- Implement and validate LLM architectures end-to-end - from PyTorch model

definition through distributed execution on custom hardware.

- Integrate custom accelerator backends into open-source ML serving frameworks

(vLLM, PyTorch), including scheduler extensions, memory management, and model

parallelism.

- Build and maintain test infrastructure for model correctness validation across

CPU, GPU, simulator, and hardware targets.

- Profile and optimize inference workloads - identify bottlenecks, instrument

critical paths, and drive latency and throughput improvements from simulation

through hardware bringup.

- Own features end-to-end: from design through implementation, testing, and

integration into the broader software stack.

- Contribute to CI/CD pipelines that gate model and kernel changes on

correctness and performance regressions.

- Mentor engineers, drive design reviews, and raise the engineering bar across

the team. Basic

Qualifications

• - Bachelor's degree in computer science or
• equivalent
• - 7+ years of full software development life cycle, including coding standards,
• code reviews, source control management, build processes, testing, and
• operations experience
• - Strong proficiency in C/C++
• - Strong Linux systems knowledge
• - Experience developing compute kernels for GPUs, DSPs, or custom accelerators
• - Proven track record of owning and delivering complex software features
• end-to-end

Preferred Qualifications

• - Knowledge of Machine Learning and LLM
• fundamentals, including transformer architecture, training/inference lifecycles,
• and optimization techniques
• - Knowledge of ML frameworks including JAX, PyTorch, vLLM, SGLang, Dynamo,
• TorchXLA, and TensorRT
• - Familiarity with speculative decoding, KV cache optimization, or other LLM
• serving optimizations
• - Experience with distributed systems - collective communication, RDMA, or
• high-speed interconnect programming
• - Experience with hardware simulation environments and model validation
• workflows
• - Demonstrated early adopter of AI-assisted development tools - uses LLMs or
• code-generation agents as part of daily workflow
• Amazon is an equal opportunity employer and does not discriminate on the basis
• of protected veteran status, disability, or other legally protected status.
• Los Angeles County applicants: Job duties for this position include: work safely
• and cooperatively with other employees, supervisors, and staff; adhere to
• standards of excellence despite stressful conditions; communicate effectively
• and respectfully with employees, supervisors, and staff to ensure exceptional
• customer service; and follow all federal, state, and local laws and Company
• policies. Criminal history may have a direct, adverse, and negative relationship
• with some of the material job duties of this position. These include the duties
• and responsibilities listed above, as well as the abilities to adhere to company
• policies, exercise sound judgment, effectively manage stress and work safely and
• respectfully with others, exhibit trustworthiness and professionalism, and
• safeguard business operations and the Company’s reputation. Pursuant to the Los
• Angeles County Fair Chance Ordinance, we will consider for employment qualified
• applicants with arrest and conviction records.
• 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
• [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.
• The base salary range for this position is listed below. Your Amazon package
• will include sign-on payments and restricted stock units (RSUs). Final
• compensation will be determined based on factors including experience,
• qualifications, and location. Amazon also offers comprehensive benefits
• including health insurance (medical, dental, vision, prescription, Basic Life &
• AD&D insurance and option for Supplemental life plans, EAP, Mental Health
• Support, Medical Advice Line, Flexible Spending Accounts, Adoption and Surrogacy
• Reimbursement coverage), 401(k) matching, paid time off, and parental leave.
• Learn more about our benefits at https://amazon.jobs/en/benefits
• [https://amazon.jobs/en/benefits].
• USA, CA, Cupertino - 193,300.00 - 261,500.00 USD annually