Staff ML Infra Engineer — TPU Training & Inference

GOOGLE ASIA PACIFIC PTE. LTD.; Singapore; On-site; SGD 250,000 - 350,000; Full time

Job summary

GOOGLE ASIA PACIFIC PTE. LTD. is seeking an experienced ML software engineer to lead onboarding for stable training and inference workloads on the latest TPU NPIs. You will architect onboarding frameworks and work closely with customers to deploy workloads at scale.

The role requires deep expertise in ML infrastructure, large-scale models, and modern training/inference techniques, with opportunities to collaborate across AI and infrastructure teams.

Qualifications

• Bachelor’s degree or equivalent practical experience.
• 8 years of experience in software development.
• 4 years of experience leading ML design and optimizing ML infrastructure.
• 2 years of experience with state of the art training and inference techniques.

Responsibilities

• Act as ML software engineer leading onboarding for stable stack and optimization for training and serving solutions on TPU NPIs.
• Set the technical direction and architect training/inference onboarding and optimization frameworks for customer needs.
• Collaborate with ML research, ML performance, and tooling teams to improve deployment efficiency.
• Ship stable stack to customers for training/inference optimization on TPU.
• Prepare and optimize very large reference models for state-of-the-art inference at scale.

Skills

Software development

ML design

ML infrastructure

Onboarding ML workloads

Education

Bachelor's degree or equivalent practical experience

Tools

Megatron-LM

DeepSpeed

TensorRT-LLM

vLLM

SGLang

Job description

GOOGLE ASIA PACIFIC PTE. LTD. is seeking an experienced ML software engineer to lead onboarding for stable training and inference workloads on the latest TPU NPIs. You will architect onboarding frameworks and work closely with customers to deploy workloads at scale.

The role requires deep expertise in ML infrastructure, large-scale models, and modern training/inference techniques, with opportunities to collaborate across AI and infrastructure teams.