AI/ML Software Engineer - TPU Efficiency for YouTube

Google LLC; Mountain View (CA); On-site; USD 147,000 - 210,000; Full time

Benefits offered by this job

Equity

Job summary

Google is seeking a Software Engineer III, AI/ML, TPU Efficiency for YouTube in Mountain View, CA. You will own optimization of models powering the YouTube algorithm and work on model efficiency with low-precision quantization, distillation, and hardware-friendly architectures.

You will build and optimize models in the Recommendation System stack, adapt to next-gen TPU hardware, and collaborate across the stack from end-user ML models to Hardware/Software architecture.

Qualifications

• Bachelor's degree or equivalent practical experience.
• 2 years of experience with ML infrastructure (e.g., model deployment, model evaluation, data processing, debugging).
• 2 years of experience with accelerator hardware architectures (TPUs/GPUs).
• 2 years of experience programming in C++ and Python.
• Experience with ML frameworks such as TensorFlow, JAX, and PyTorch, or ML compilers (e.g., accelerated linear algebra (XLA)).
• Experience with ML model optimization.

Responsibilities

• Analyze ML workloads, identify compute and memory bandwidth bottlenecks, and optimize accelerator utilization to maximize compute efficiency.
• Explore, implement, and productionize algorithmic efficiency techniques, including low-precision quantization, knowledge distillation, parameter sharing, and attention optimizations.
• Optimize auxiliary serving and distributed data pipelines, including data ingestion, feature transformation, embedding lookups, and memory caching to support real-time training and inference.
• Collaborate with ML model developers and researchers to co-design hardware-friendly model architectures and deploy universal efficiency libraries.

Skills

ML infra

C++

Python

TPU/GPUs

ML frameworks

Model optimization

Education

Bachelor's degree

Tools

TensorFlow

JAX

PyTorch

XLA

Job description

Google is seeking a Software Engineer III, AI/ML, TPU Efficiency for YouTube in Mountain View, CA. You will own optimization of models powering the YouTube algorithm and work on model efficiency with low-precision quantization, distillation, and hardware-friendly architectures.

You will build and optimize models in the Recommendation System stack, adapt to next-gen TPU hardware, and collaborate across the stack from end-user ML models to Hardware/Software architecture.