The role involves identifying and addressing performance issues across various machine learning systems, including research, training, and inference. A primary responsibility includes designing and optimizing kernels for TPUs. The engineer will provide feedback to researchers regarding the impact of model changes on system performance. Additionally, the position requires implementing low-latency sampling for large language models and adapting models for low-precision inference. The work also entails building quantitative models of system performance and designing custom collective communication algorithms. Debugging kernel performance at the assembly level is a core component of the responsibilities.

Candidates must possess significant experience in optimizing machine learning systems for TPUs, GPUs, or other accelerators. The role requires a track record of solving large-scale systems problems and performing low-level optimization. A minimum of a bachelor's degree or an equivalent combination of education, training, and experience is required. Applicants should demonstrate a field of study relevant to the role through coursework or professional experience.