Solution Architect (Kernels and Performance, Core ML Frameworks) Full Time

We are seeking a talented Solution Architect (Kernels and Performance, Core ML Frameworks) to drive the performance and optimization of both training and serving, delivering massive impact for customers. In this role, you will join the highly interdisciplinary CoreML team with exposure to the newest Tensor Processing Unit (TPU), Graphics Processing Unit (GPU) hardware, the latest ML models, and the advanced toolchains that bridge them. Your work will directly enable AI research and production deployments across Google Cloud and the broader open-source ecosystem, addressing complex technical issues that directly impact the efficiency and scalability of AI across the industry.

Responsibilities

• Design and optimize high-performance kernels (using languages like Pallas, Mosaic and Triton) targeting Tensor Processing Unit (TPU) and Graphics Processing Unit (GPU) architectures for critical Machine Learning (ML) operations, redefining what’s possible from massive training runs to high-speed inference
• Architecture of infrastructure such as benchmarking suites, autotuning frameworks, performance analysis tools, regression testing and documentation
• Transformation of how the developer community interacts with increasingly critical custom kernels in key Open-Source Software (OSS) libraries
• Track the latest advancements in hardware architectures, compiler technologies and AI models to identify new opportunities for performance optimization through custom kernels
• Engagement with ML researchers, framework developers (Just After eXecution (JAX), PyTorch) and compiler engineers (Accelerated Linear Algebra (XLA)) to enhance adoption
• Identification of new requirements and resolution of bottlenecks by providing appropriate solutions

Requirements

• Bachelor’s degree or equivalent practical experience
• Overall 12+ years of industry experience
• 5+ years of experience with software development in C++ or Python
• 3+ years of experience testing, maintaining or launching software products, and 1 year of experience with software design and architecture
• Background in performance optimization at the kernel level
• Expertise in optimizing TPU/GPU code, using low-level kernel languages like Pallas, Compute Unified Device Architecture (CUDA) or Triton
• Knowledge of ML Frameworks (JAX/PyTorch), common operations like attention and Mixture of Experts (MoEs), including model optimization and low-precision formats
• Understanding of modern accelerators (e.g., data movement, pipelining, heterogeneous compute and scale-out)
• Understanding of compiler principles (optimization, code generation) and toolchains such as MLIR, OpenXLA
• Showcase of building developer infrastructure, including Open-Source Software (OSS) libraries, flexible high-performance APIs and easy-to-consume documentation to empower the community
• Skills in GPGPU/GPU Programming, NVIDIA CUDA, PyTorch and Python
• Excellent investigative and problem-solving capabilities with communication skills across cross-functional teams
• Advanced proficiency in English (B2+/C1)

Nice to have

• Proficiency in C++