Staff Machine Learning Engineer (Generative Model Performance & Efficiency)

• The Simulator Team at Waymo builds state-of-the-art simulations of realistic environments for testing, training, and validation of the Waymo Driver. Our team is a diverse, and collaborative group of machine learning (ML) engineers, software engineers, and ML research engineers
• We develop industry-leading simulation solutions using advanced generative and reconstructive ML algorithms, to model the real world, encompassing realistic agents, roads, traffic systems, weather, and the full sensor suite (Camera, Lidar, Radar)
• To accelerate the fidelity, scalability, controllability, and richness of our simulations, we are pushing the frontiers of 3D world modeling. We leverage state-of-the-art ML technologies trained on large-scale datasets to create dynamic and semantically rich virtual worlds, directly impacting the development and validation of the Waymo Driver
• In this role, you will report to a Senior Staff Engineering Manager
• Analyze model architectures and identify bottlenecks in training and inference performance (e.g., memory bandwidth, compute, communication)
• Apply and develop techniques such as quantization (e.g., FP8, INT4), pruning, knowledge distillation, and efficient attention mechanisms
• Optimize model code for specific hardware accelerators (TPUs, GPUs), leveraging compiler features and low-level libraries (e.g., XLA)
• Experiment with different model partitioning and sharding strategies (e.g., data, tensor, pipeline parallelism, expert parallelism) to improve scalability and efficiency
• Design and implement low-latency, high-throughput serving solutions for generative models and optimize training pipelines to reduce training time
• Build and maintain tools for performance analysis, profiling (e.g., xprof), and debugging of ML models

Benefits

• Medical, dental, and vision insurance for employees and dependents
• Employee assistance programs focused on mental health
• Personalized workplace adjustments for diverse needs and abilities, including physical, mental, and neurodivergent considerations
• Access to mental health apps
• Onsite wellness centers
• Support programs including menopause benefit
• Counseling services
• Second medical opinion for you and your loved ones
• Medical advocacy program for transgender employees
• Competitive compensation
• Regular bonus and equity performance grant opportunities
• Generous 401(k) and regional retirement plans
• Annual cross-company compensation review and pay equity analysis
• 1-on-1 financial coaching
• Fertility and growing family assistance
• Parental leave and baby bonding leave
• Survivor income benefit
• Elder care and support
• Backup childcare
• Caregiver leave
• Paid time off, including vacation, bereavement, sick leave, parental leave, disability, and holidays
• Jury duty leave
• Military leave
• Hybrid work model with remote work opportunities also available
• Educational reimbursement
• Peer learning and coaching platform
• Donation matching programs
• Volunteer hours
• Employee resource groups
• Internal community groups and local culture clubs
• Inspiring spaces to work, recharge, and collaborate with fellow Waymonauts
• On-site meals and snacks
• Fitness centers, massage programs, and ergonomic support
• On-demand fitness, wellbeing, and cooking classes
• Commuter benefits- Strong programming skills in Python and potentially C++, with experience in software development best practices
• Proficiency in JAX, Flax, and potentially TensorFlow/PyTorch
• 5+ years of experience with deep learning architectures (especially Transformers, Diffusion Models, MoEs), algorithms, and optimization techniques
• MS or PhD in Computer Science, Machine Learning, Robotics, or a related field
• Hands-on experience with quantization, pruning, distillation, and other model compression methods
• Expertise in using profiling tools (e.g., XProf, Perfetto, NVIDIA Nsight) to diagnose performance issues in ML workloads
• Familiarity with ML compilers like XLA and an understanding of how they translate high-level code to efficient hardware instructions
• Knowledge of TPU and GPU architectures and how to optimize code for them
• Experience contributing to frameworks and libraries that improve training speed and scalability (e.g., JAX, Gemax, XManager)
• Understanding of concepts related to training and serving models across multiple devices and machines