Software Engineer, TPU Compilers

📋 Description

• Optimize model execution on TPUs, improving performance across the complete forward pass and the
• Profile workloads, identify bottlenecks, and implement improvements to latency, throughput, and
• Develop and optimize kernels, using an understanding of the surrounding model to guide
• Explore techniques such as quantization to improve inference efficiency while preserving model
• Build benchmarks and tools to measure performance, validate correctness, and guide optimization
• Partner with researchers and engineers to support new model architectures and bring performance

🎯 Requirements

• Have practical experience optimizing machine learning workloads on TPUs, through training
• Understand modern model architectures and how their computation maps onto accelerator hardware.
• Can connect detailed kernel behavior to the performance of a complete model and work across these
• Bring strong software engineering skills and a systematic approach to profiling, debugging, and
• Take ownership of open-ended technical problems and can identify priorities, make design decisions
• Enjoys collaborating with researchers and engineers, learning unfamiliar parts of a system, and

🎁 Benefits

• Offers Equity