Software Engineer III, Machine Learning Compiler, TPU Scaling

Job Description

Develop XLA compiler capabilities for efficient communication and horizontal scaling on TPUs, including collective operation lowering, communication-computation overlap, and compiler scheduling. Analyze TPU architecture and interconnect features, and collaborate with performance and research teams to diagnose scaling bottlenecks and improve workload performance.

Requirements: Requires a bachelor’s degree or equivalent practical experience, at least two years of low-level systems programming experience, and experience in performance optimization and compiler development. Also requires experience in an ML specialization and end-to-end machine learning; a graduate degree and experience with data structures, algorithms, accessibility, or hardware design and architecture are preferred.

Key Skills: C++, C, Machine Learning, ML Infrastructure, End-to-End Machine Learning, Performance Optimization, Compiler Development, XLA, TPU, Collective Operations, Compiler Scheduling, Program Sharding, Data Structures, Algorithms, Hardware Architecture, Hardware Design

Benefits: Bonus, Equity, Benefits