TPU Kernel Optimization & Performance Engineering

About the Role

This role focuses on accelerating the performance, portability, and numerical correctness of cutting-edge machine learning workloads on Google TPUs. You will work across high-level ML frameworks, compiler/runtime systems, distributed execution environments, and low-level custom kernels to eliminate hardware bottlenecks and implement ultra-efficient workloads. The ideal candidate operates seamlessly between model-level math and hardware-aware execution: inspecting GPU-oriented implementations, analyzing their exact semantics, translating them into TPU-compatible forms, and rigorously validating both numerical correctness and execution speed.

Key Responsibilities

• Kernel Development & Porting: Develop and optimize high-performance custom kernels and critical ML operations for TPU execution. Translate CUDA and Triton kernel behaviors into standard PyTorch or JAX implementations that compile efficiently via OpenXLA/XLA.
• Workload Profiling & Diagnostics: Profile distributed training and inference workloads to identify compute, memory, communication, compilation, and data-movement bottlenecks.
• Full-Stack Debugging: Debug correctness and performance issues across ML frameworks, compiler/runtime layers, distributed execution topologies, and accelerator kernels.
• Adversarial Agent Benchmarking & Task Formulation: Design authentic, production-derived engineering challenges from modern open-weight architectures (e.g., DeepSeek, Qwen, GLM) to stress-test and evaluate baseline frontier models (e.g., Gemini 3.7 Flash, Claude Code), establishing ground-truth solutions and rigorous evaluation rubrics for autonomous agents.
• Hardware-Aware Strategy: Design TPU-friendly tiling, sharding, reduction, data-movement, and mixed-precision strategies for standard and custom ML operations.
• Correctness & Gradient Validation: Validate numerical equivalence, gradient correctness, masking/indexing behavior, and edge cases between source GPU implementations and TPU-targeted implementations.
• Benchmark Architecture: Build and maintain diagnostic workflows and reproducible microbenchmarks to measure latency, throughput, utilization, memory usage, communication overhead, and performance regressions.
• Frontier Model Optimization: Analyze transformer and LLM workloads—including attention variants, MoE, KV-cache, quantized matmuls, and distributed training/serving patterns—to capture performance gains.
• Cross-Functional Collaboration: Collaborate with ML systems, compiler, infrastructure, and model engineers to translate workload requirements into efficient accelerator implementations.

Requirements & Tech Stack

• Systems Programming: Strong programming proficiency in Python and at least one systems language (C++, CUDA, or Pallas).
• Accelerator Frameworks: Hands-on experience developing and optimizing ML accelerator code using JAX/Pallas, PyTorch-XLA, or OpenXLA.
• Distributed ML: Practical understanding of distributed machine learning concepts.

Qualifications

• Education: Bachelor’s degree in Computer Science, Electrical Engineering, or a related technical field (or equivalent practical experience).
• TPU Hardware Knowledge: Preferred: In-depth knowledge of TPU architecture, memory hierarchies, and hardware performance characteristics.
• GPU-to-TPU Translation: Preferred: Proven experience converting CUDA or Triton implementations to standard PyTorch or JAX APIs for TPU execution while maintaining numerical and gradient semantics.
• Distributed Collectives & Parallelism: Preferred: Strong background in distributed ML techniques, including data, tensor, sequence, and expert parallelism; multidimensional sharding; and collective communications (all-reduce, all-gather, reduce-scatter, all-to-all).
• Frontier Model Profiling: Preferred: Experience extracting, profiling, and optimizing kernels from frontier open-weight models (e.g., DeepSeek V3/V4, Qwen 2.5/3.8, Kimi K2.5).
• Numerical Stability & Precision: Preferred: Deep experience debugging numerical correctness issues involving bf16/fp32 accumulation, quantization scaling, softmax/log-sum-exp stability, masking, and backward-pass gradient behavior.
• High-Performance ML Stacks: Preferred: Experience with performance-oriented open-source systems such as vLLM, FlashAttention, JAX/Flax, PyTorch/XLA, or equivalent kernel libraries.
• Benchmarking & Tooling: Preferred: Experience designing reproducible microbenchmarks, performance regression test suites, and profiling tools for accelerator workloads.
• Advanced Degree: Preferred: Master’s degree or PhD in Computer Science, Data Science, or a related technical field.

Engagement Details

• Engagement Type: Independent Freelance Contract
• Location: 100% Remote (India)