TPU Kernel Optimization & Performance Engineering Expert (Freelance)

Uber AI Solutions; India; Híbrido; INR 2.480.000 - 4.822.000; Jornada completa

Descripción de la vacante

Uber AI Solutions seeks a skilled ML engineer to accelerate performance on Google TPUs, translating GPU kernels to TPU-friendly forms while ensuring numerical correctness. You will collaborate across ML frameworks, compilers, and distributed systems to optimize workloads.

The role requires strong Python, C++/CUDA, and JAX/Pallas experience, with emphasis on full-stack debugging and frontier-model benchmarking.

Formación

• Bachelor’s degree in CS, EE, or related field.
• Strong programming in Python; proficiency in a systems language (C++, CUDA, or Pallas).
• Experience with ML accelerator code and distributed ML concepts.

Responsabilidades

• Develop and optimize high-performance custom kernels for TPU execution.
• Translate CUDA/Triton kernel behaviors into TPU-friendly forms for PyTorch/JAX.
• Profile distributed training/inference workloads to identify bottlenecks.
• Debug correctness and performance across ML frameworks, compiler/runtime layers, and accelerators.
• Design TPU-friendly tiling, sharding, and mixed-precision strategies.

Conocimientos

Python

C++/CUDA/Pallas

JAX/Pallas

PyTorch-XLA/OpenXLA

Distributed ML

Educación

Bachelor’s degree in Computer Science or Electrical Engineering

Herramientas

JAX/Pallas

PyTorch-XLA

OpenXLA

Descripción del empleo

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.

Minimum Qualifications

• Education: Bachelor’s degree in Computer Science, Electrical Engineering, or a related technical field (or equivalent practical experience).
• 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.

Preferred Qualifications

• Advanced Degree: Master’s degree or PhD in Computer Science, Data Science, or a related technical field.
• TPU Hardware Knowledge: In-depth knowledge of TPU architecture, memory hierarchies, and hardware performance characteristics.
• GPU-to-TPU Translation: 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: 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: 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: 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: Experience with performance-oriented open-source systems such as vLLM, FlashAttention, JAX/Flax, PyTorch/XLA, or equivalent kernel libraries.
• Benchmarking & Tooling: Experience designing reproducible microbenchmarks, performance regression test suites, and profiling tools for accelerator workloads.
• Engagement Type: Independent Freelance Contract