Senior Machine Learning Engineer, Runtime and Serving at Waymo UK Ltd

Mountain View, California, United States -

Full Time

Start Date

Immediate

Expiry Date

28 Jul, 26

Salary

263000.0

Posted On

29 Apr, 26

Experience

5 year(s) or above

Remote Job

Yes

Telecommute

Yes

Sponsor Visa

Skills

C++, Python, Machine Learning, JAX, CUDA, Triton, TensorRT, OpenXLA, Distributed systems, Inference engines, Model compression, Performance optimization, Hardware acceleration, Profiling, Benchmarking, Deep learning

Industry

technology;Information and Internet

Description

Waymo is an autonomous driving technology company with the mission to be the world's most trusted driver. Since its start as the Google Self-Driving Car Project in 2009, Waymo has focused on building the Waymo Driver—The World's Most Experienced Driver™—to improve access to mobility while saving thousands of lives now lost to traffic crashes. The Waymo Driver powers Waymo’s fully autonomous ride-hail service and can also be applied to a range of vehicle platforms and product use cases. The Waymo Driver has provided over ten million rider-only trips, enabled by its experience autonomously driving over 100 million miles on public roads and tens of billions in simulation across 15+ U.S. states. The ML Optimization team at Waymo provides a set of tools to support and automate the lifecycle of the machine learning workflow, including feature and experiment management, model development, optimization and monitoring. These efforts have resulted in making machine learning more accessible to teams at Waymo, including Perception, Planner, Research and Simulation. We are looking for engineers with ML software & systems expertise to help build the next generation Waymo onboard ML inference engine for Waymo fundamental model. You'll work across the entire ML stack from the system perspective, from efficient deep learning models, model compression, ML software (e.g. JAX, XLA, Triton, and CUDA), to . You will be pleasantly challenged with deploying Waymo ML models on limited computation resources. In this hybrid role, you will report to the Senior Manager of Runtime and Optimization. You will: Architect and develop an efficient, high-performance ML runtime and serving system tailored for both onboard autonomous vehicle compute and large-scale, offboard data center environments. Lead the integration and feature development for ML inference runtimes across both domains, balancing the strict real-time latency and memory constraints of onboard systems with the high-throughput, highly concurrent demands of offboard serving fleets. Drive the strategic migration of ML workloads toward a JAX-native runtime architecture, which includes extending and modifying underlying ML compilers and runtimes (e.g., OpenXLA/PjRT, TensorRT). Collaborate with world-class Waymo ML practitioners across perception, planner, and research to analyze system-level ML workloads and apply hardware-aware compute optimizations. Design and build robust tooling for profiling, benchmarking, and identifying system-level bottlenecks across the end-to-end ML software stack. You Have: B.S. or M.S. in CS, EE, Deep Learning or a related field 5+ years of professional software engineering experience focused on building, scaling, or maintaining ML systems and infrastructure. 5+ years production programming in C++. 3+ years of production experience in Python and major deep learning frameworks (e.g., PyTorch, JAX). Experience optimizing ML software for hardware accelerators (e.g., GPUs, TPUs, custom silicon). Experience building low-latency, highly concurrent distributed backend systems. We Prefer PhD in CS, EE, Deep Learning or a related field. Experience modifying ML compilers, runtimes, or inference engines (e.g., TensorRT, ONNX Runtime, OpenXLA/PjRT, TVM). Experience building or scaling LLM serving systems, including expertise in distributed inference and performance optimization (e.g., KV/prefix caching, continuous batching). Experience with custom kernel development (e.g., CUDA/CUDA Tile, Triton, JAX/Pallas). Experience architecting unified serving APIs and optimizing tensor buffer management (e.g., zero-copy data transfer, shared memory) for complex, multi-model inference pipelines. The expected base salary range for this full-time position across US locations is listed below. Actual starting pay will be based on job-related factors, including exact work location, experience, relevant training and education, and skill level. Your recruiter can share more about the specific salary range for the role location or, if the role can be performed remote, the specific salary range for your preferred location, during the hiring process. Waymo employees are also eligible to participate in Waymo’s discretionary annual bonus program, equity incentive plan, and generous Company benefits program, subject to eligibility requirements. Salary Range $213,000—$263,000 USD

How To Apply:

Incase you would like to apply to this job directly from the source, please click here

Responsibilities

You will architect and develop high-performance ML runtime and serving systems for both onboard autonomous vehicle compute and large-scale data centers. Additionally, you will lead the migration to JAX-native architectures and collaborate with cross-functional teams to optimize system-level ML workloads.