Staff ML Infrastructure Engineer for AI Evaluation

Waymo LLC; Mountain View (CA); On-site; USD 251,000 - 310,000; Full time

Benefits offered by this job

Annual bonus program

Equity incentives

Benefits program

Job summary

Waymo LLC in Mountain View, CA seeks researchers and software engineers to build ML infrastructure for the Waymo Driver evaluation stack, taking models from data to training to production serving.

You will design scalable systems for dataset generation, model training, and evaluation across real-world and simulation data, improve offline and online inference, and guide architectural decisions while mentoring engineers.

Qualifications

• B.S. in Computer Science, Machine Learning, or a related technical field, or equivalent practical experience.
• 7+ years of hands-on experience building machine learning systems.
• Experience designing, scaling, training, deploying, and optimizing large-scale ML systems from data to model.
• Strong coding experience in Python and/or C++.
• Experience with ML frameworks such as JAX, TensorFlow, or PyTorch, including both the user-facing APIs and their internals.
• Proven ability to set technical direction and influence partner teams in ambiguous problem spaces.

Responsibilities

• Lead the design and development of ML infrastructure that takes DUE models from data to training to production serving, including event detection models, triage foundation models, and agents.
• Design and scale distributed systems for large-scale dataset generation, model training, and evaluation across Waymo's real-world and simulation data.
• Own the reliability, efficiency, and observability of production ML serving, including accelerator capacity planning and performance optimization.
• Keep offline evaluation and production inference consistent, so the models we evaluate are the models we serve.
• Drive technical direction for the team, guide architectural decisions, and mentor engineers.
• Collaborate with modeling, evaluation, and platform teams across Waymo to turn product and business goals into measurable technical deliverables.

Skills

Python

C++

ML systems

TensorFlow

PyTorch

JAX

Distributed systems

Education

B.S. in Computer Science or related field

M.S. or Ph.D. in Computer Science or ML

Tools

TPU

GPU

Job description

Waymo LLC in Mountain View, CA seeks researchers and software engineers to build ML infrastructure for the Waymo Driver evaluation stack, taking models from data to training to production serving.

You will design scalable systems for dataset generation, model training, and evaluation across real-world and simulation data, improve offline and online inference, and guide architectural decisions while mentoring engineers.