Member of Technical Staff, TPU Kernel Engineering

The role

This person will design, write, and optimize kernels for TPU in Pallas and JAX. They will profile production-relevant models on TPU to remove bottlenecks through kernel or compiler changes.

What you will do

• Design, write, and optimize kernels for TPU in Pallas and JAX
• Profile production-relevant models on TPU
• Find bottlenecks and write kernel or compiler changes to remove them
• Provide clear correctness, latency, and throughput benchmarks
• Make TPU support in vLLM usable, fast, benchmarked, and maintainable
• Write Pallas attention or MoE kernels
• Implement low-latency, high-throughput sampling for large language models on TPU
• Adapt existing models for low-precision inference on TPU
• Debug kernel performance at the compiler-artifact or assembly level

Required qualifications

• Bachelor's degree or equivalent experience in computer science, engineering, systems, machine learning, or similar
• Hands-on experience designing and implementing kernels for TPUs or other ML accelerators, in Pallas, Triton, CUDA, or a similar kernel language
• Hands-on experience building or optimizing TPU workloads using JAX, XLA, Pallas, or related compiler and runtime tooling
• Deep understanding of TPU execution, memory behavior, compilation, and performance constraints for ML workloads
• Experience optimizing ML kernels such as attention, GEMM, sampling, KV cache, or fused kernels
• Strong performance profiling and benchmarking skills

Nice to have

• Experience with vLLM, SGLang, TensorRT-LLM, XLA-based serving, or other LLM inference systems
• Familiarity with batching, KV cache, decoding, serving tradeoffs, and backend performance constraints in production inference systems
• Experience with compiler technologies such as XLA, MLIR, LLVM, Pallas, or other kernel DSLs, including lowering, fusion, and backend code generation
• Knowledge of quantization methods such as INT8, FP8, mixed precision, or TPU-specific numeric formats, including accuracy and performance tradeoffs
• Experience building quantitative models of system performance, such as roofline analysis, simulation, or hardware microbenchmarking
• Contributed to vLLM, JAX/XLA, Pallas, PyTorch/XLA, compiler projects, or other open-source ML infrastructure
• Built TPU benchmarking infrastructure or automated performance regression detection for accelerator workloads
• Worked directly with Google TPU ecosystem stakeholders, accelerator platform teams, or early-access programs

Benefits

• Health benefits
• Dental benefits
• Vision benefits
• 401(k) company match

Before you apply

• We sponsor visas on a case-by-case basis

About the company

Inferact's mission is to grow vLLM as the world's AI inference engine and accelerate AI progress by making inference cheaper and faster. Founded by the creators and core maintainers of vLLM, we sit at the intersection of models and hardware.