Senior/Staff Software Engineer (ML Performance Optimization)

Join Zoox, a cutting-edge autonomous vehicle company, as a Senior/Staff Software Engineer specializing in ML Performance Optimization. In this role, you will work across all ML teams, lead a team of software engineers, and drive innovation in autonomous driving and rider experience. You will have the opportunity to significantly impact the future of ML at Zoox and enjoy a range of benefits, including paid parental leave, health insurance, and a comprehensive wellbeing program.

• Lead the design, implementation, and operation of cutting-edge ML Training and inference performance optimization techniques.
• Collaborate closely with cross-functional teams, including ML researchers, software engineers, data engineers, and hardware engineers, to define requirements and align on architectural decisions.
• Develop and execute a strategic vision for the ML Performance Optimization team to unlock ML innovation in autonomous driving and rider experience.

• You do not need to match every listed expectation to apply for this position
• Experience with model compression techniques to reduce model size and improve performance
• Experience using profiling tools like NVIDIA's Nsight or PyTorch's Profiler for identifying model training and serving bottlenecks
• Proficient in Python and C++
• Experience with GPU-accelerated inference using TensorRT, Ray Serve, or similar frameworks
• Strong experience with training frameworks like PyTorch, leveraging GPUs efficiently for distributed model training
• Excellent leadership skills with a demonstrated ability to lead high-performing engineering teams
• 10+ years of total experience, including 4+ years of working on large-scale model training or inference platforms