> Markdown version of [/jobs/ext/3419049-aiml-staff-ml-infrastructure-engineer-ml-platform-technology-pre-training-infrastructure](https://www.wearedevelopers.com/jobs/ext/3419049-aiml-staff-ml-infrastructure-engineer-ml-platform-technology-pre-training-infrastructure). Every page supports `.md` or `Accept: text/markdown`. Links point to the HTML versions so they work for humans too. Agent guide: [/agents.md](https://www.wearedevelopers.com/agents.md).

---

# AIML - Staff ML Infrastructure Engineer, ML Platform & Technology - Pre-training Infrastructure

- **Company:** Apple Inc.
- **Location:** United States
- **Experience:** Expert
- **Salary:** $184,700.0 - $324,800.0
- **Contract:** Permanent contract
- **Skills:** Nvidia CUDA, Distributed Systems, Python (Programming Language), Machine Learning, Performance Tuning, Azure Machine Learning, Pytorch, Parallel Computation, Information Technology, Machine Learning Operations, Artificial Intelligence Markup Language (AIML), Programming Languages
- **Published:** September 2, 2026
- **Apply:** https://www.indeed.com/viewjob?jk=f416d2f2aec59cbb

## About the Role

Advanced degree in Computer Science, Engineering, or a related field

Experience with accelerators such as TPU or GPU and understanding of accelerator architecture and performance characteristics

Experience with JAX, XLA, PyTorch or other ML compiler/runtime stacks

Experience developing or optimizing accelerator kernels using Pallas, Triton, CUDA, or similar technologies

Experience optimizing large-scale foundation model training and distributed communication

Minimum Qualifications

6+ years of experience building or optimizing high-performance ML or distributed systems

Proficient in Python or other relevant programming languages

Strong understanding of distributed systems, parallel computing, and performance optimization

Experience profiling and optimizing compute-, memory-, or communication-intensive workloads

Ability to clearly communicate complex technical problems and collaborate with partners to develop solutions

Bachelor's degree in Computer Science, Engineering, or a related field

## Description

As an engineer on the ML Compute team, your work will include:

 * Drive performance optimization for large-scale foundation model training on TPUs, focusing on efficiency, throughput, and scalability

 * Profile and optimize JAX/XLA workloads across compute, memory, communication, and compilation.

 * Develop and optimize high-performance TPU kernels for critical ML operations such as attention and Mixture-of-Experts (MoE)

 * Optimize distributed training techniques, sharding strategies, and collective communication over TPU interconnects (ICI/Fabric)

 * Research and implement new techniques across the JAX, XLA, and TPU stack to improve end-to-end training performance

 * Develop performance profiling, benchmarking, and automated tuning capabilities for large-scale training workloads.

 * Collaborate with cross-functional engineers to solve large-scale ML training challenges

 * Lead complex technical projects and mentor engineers in areas of your expertise