> Markdown version of [/jobs/ext/2685159-aiml-staff-ml-infrastructure-engineer-ml-platform-technology-pre-training-infrastructure](https://www.wearedevelopers.com/jobs/ext/2685159-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:** San Francisco, CA, United States
- **Experience:** Expert
- **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.sanfranciscogigs.com/job.asp?id=3373897079&tx=KK2323FFP&pt=1&aff=0B19D771-A501-4A5E-8338-2A822B784D54&utm_source=Job%20Feed&utm_medium=textkernel&utm_campaign=DE&utm_term=0B19D771-A501-4A5E-8338-2A822B784D54

## About the Role

* 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, * 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

## 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
 * Cultivate a team centered on collaboration, technical excellence, and innovation


## Related Videos

- [Introduction to Azure Machine Learning](https://www.wearedevelopers.com/videos/368-introduction-to-azure-machine-learning)
- [Coffee with Developers - Stephen Jones - NVIDIA](https://www.wearedevelopers.com/videos/1303-coffee-with-developers-stephen-jones-nvidia)
- [Photonic Computing: Programming a New Class of AI Accelerators (incl. Live Coding)](https://www.wearedevelopers.com/videos/100196-photonic-computing-programming-a-new-class-of-ai-accelerators-incl-live-coding)
- [From Model to Metal: An Open Source Stack for Accelerating Intelligence](https://www.wearedevelopers.com/videos/1636-from-model-to-metal-an-open-source-stack-for-accelerating-intelligence)
- [Making neural networks portable with ONNX](https://www.wearedevelopers.com/videos/301-making-neural-networks-portable-with-onnx)
- [Developer Experience, Platform Engineering and AI powered Apps](https://www.wearedevelopers.com/videos/990-developer-experience-platform-engineering-and-ai-powered-apps)

## Related Articles

- [How to Become an AI Engineer](https://www.wearedevelopers.com/magazine/331-how-to-become-an-ai-engineer)
- [MLOps And AI Driven Development](https://www.wearedevelopers.com/magazine/82-mlops-and-ai-driven-development)
- [What Are Large Language Models?](https://www.wearedevelopers.com/magazine/304-what-are-large-language-models)
- [MLOps – What’s the deal behind it?](https://www.wearedevelopers.com/magazine/125-mlops-what-s-the-deal-behind-it)
- [MLops – Deploying, Maintaining And Evolving Machine Learning Models in Production](https://www.wearedevelopers.com/magazine/115-mlops-deploying-maintaining-and-evolving-machine-learning-models-in-production)
- [Stephan Gillich - Bringing AI Everywhere](https://www.wearedevelopers.com/magazine/489-stephan-gillich-bringing-ai-everywhere)