Staff ML Engineer

Responsibilities

• Analyze model architectures and identify training and inference bottlenecks involving memory bandwidth, compute, and communication.
• Develop and apply quantization, pruning, knowledge distillation, and efficient attention techniques.
• Optimize model code for TPU and GPU accelerators using compiler features and low-level libraries.
• Evaluate data, tensor, pipeline, and expert parallelism strategies to improve scalability and efficiency.
• Design and implement low-latency, high-throughput serving solutions for generative models.
• Optimize training pipelines and build tools for model performance analysis, profiling, and debugging.

Requirements

• MS or PhD in Computer Science, Machine Learning, Robotics, or a related field.
• 5+ years of experience with deep learning architectures, algorithms, and optimization techniques, especially Transformers, Diffusion Models, and MoEs.
• Proficiency in JAX and Flax, with potential experience in TensorFlow or PyTorch.
• Expertise with ML workload profiling tools such as XProf, Perfetto, and NVIDIA Nsight.
• Hands-on experience with quantization, pruning, distillation, and other model-compression methods.
• Strong Python programming skills and potentially C++, with software development best-practice experience.
• Preferred: knowledge of TPU and GPU architectures, XLA and ML compiler behavior, multi-device and multi-machine training and serving, and contributions to JAX, Gemax, or XManager.

Benefits

• Eligibility for Waymo’s discretionary annual bonus program, equity incentive plan, and generous Company benefits program, subject to eligibility requirements.
• Full-time position with salary ranges varying by US location; remote-location salary details may be provided during the hiring process.