Staff ML Engineer
Core
Optimize deep learning model architectures and training/inference pipelines for memory, compute, and communication efficiency, focusing on generative models.
Role type
Staff ML Engineer (Model Optimization & Infrastructure)
Builds
Low-latency, high-throughput serving solutions and optimized training pipelines for generative models.
Domain
AI/ML Infrastructure, Model Compression, High-Performance Computing
Deliverable
production ML models
Required skills
Deep learning architecture design, model quantization, pruning, knowledge distillation, efficient attention techniques, TPU/GPU optimization, compiler features, low-level libraries, parallelism strategies (data/tensor/pipeline/expert), Python, C++, ML workload profiling (XProf, Perfetto, Nsight)
Preferred skills
TPU/GPU architecture knowledge, XLA/ML compiler behavior, multi-device/multi-machine training and serving, contributions to JAX, Gemax, or XManager
Technologies
JAX, Flax, TensorFlow, PyTorch, XLA, TPU, GPU, NVIDIA Nsight, Perfetto, XProf
Responsibilities
Analyze model architectures to identify bottlenecks; Develop and apply model compression techniques; Optimize model code for accelerators; Evaluate parallelism strategies; Design and implement serving solutions; Optimize training pipelines and build analysis tools
Seniority
Staff, hands-on IC with strategic impact
