NEXadept Hiring! Full Time TPU Performance Engineer in - Ricebowl

About the Role

We are looking for a TPU Performance Engineer to optimize large-scale LLM inference performance on Google TPU. You will work across TPU, kernels, compilers, and runtime systems, improving inference latency, throughput, and overall efficiency.

Key Responsibilities

• Optimize LLM inference workloads on Google TPU.
• Develop and optimize TPU backends, kernels, compiler integrations, and runtime components.
• Optimize performance-critical workloads such as Attention, GEMM, KV Cache, Sampling, and fused kernels.
• Work with JAX, XLA, Pallas and related compiler/runtime technologies.
• Build benchmarking and profiling infrastructure and identify bottlenecks across compute, memory, compilation, and runtime.
• Collaborate with model, inference, compiler, and hardware teams to improve production performance.

Requirements

• Bachelor’s degree or equivalent experience in Computer Science, Engineering, Machine Learning, Systems, or a related field.
• Hands-on experience with TPU performance optimization, preferably with JAX, XLA, Pallas, or related technologies.
• Strong understanding of TPU architecture, memory behavior, compilation, and ML workload performance.
• Experience with ML kernel optimization, LLM inference, backend/runtime development, or performance engineering.
• Strong C++ and/or Python programming skills.
• Solid experience with performance profiling and benchmarking.

Preferred Qualifications

• Experience with vLLM, SGLang, TensorRT-LLM or other LLM inference frameworks.
• Familiarity with LLM serving, batching, KV Cache, decoding, and inference optimization.
• Experience with MLIR, LLVM, Pallas, XLA or other compiler technologies.
• Knowledge of FP8, INT8, mixed precision, or quantization.
• Contributions to vLLM, JAX/XLA, Pallas, PyTorch/XLA or other open-source AI infrastructure projects.