Member of Technical Staff, TPU Kernel Engineering
Inferact · San Francisco
Extracted from the posting
Design and optimize Pallas and JAX kernels for Google TPUs for LLM inference operations such as attention, GEMM, sampling, and KV cache. Requires kernel experience in Pallas, Triton, or CUDA, plus JAX/XLA workload optimization and TPU performance profiling. Work includes benchmarking correctness, latency, and throughput to make vLLM TPU support usable and maintainable.
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, a position that took years to build.
About the Role
We're looking for a TPU kernel engineer to make vLLM a first-class inference engine on Google TPUs. A large part of this role is designing, writing, and optimizing kernels for TPU in Pallas and JAX, for the operations that dominate LLM inference: attention, GEMM, sampling, KV cache, and fused operations.
You'll work at the boundary of kernels, compilers, inference systems, and hardware architecture. You'll profile production-relevant models on TPU, find the bottleneck, and write the kernel or compiler change that removes it, with clear correctness, latency, and throughput benchmarks. Your work will help make TPU support in vLLM usable, fast, benchmarked, and maintainable.
Inferact has three kernel roles. If your kernel work is mostly on Google TPUs (Pallas, JAX), apply to TPU Kernel Engineering. If it is mostly on AMD GPUs (HIP, CK, AITER), apply to AMD GPU Kernel Engineering. For NVIDIA GPUs and cross-platform kernel work, apply to Kernel Engineering. Apply to one; we route you to the best fit.
Skills and Qualifications
Minimum 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, with the ability to use measurements, compiler artifacts, correctness tests, and reproducible benchmarks to guide optimization work.
Preferred qualifications:
- 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, and using them to guide kernel work.
Bonus points if you have:
- 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 to ship backend, compiler, or inference performance improvements.
Representative projects
- Write a Pallas attention or MoE kernel that beats the current XLA lowering for a production model.
- 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.
Logistics
- Location: This role is based in San Francisco, California. Will consider remote in the US for exceptional candidates.
- Compensation: Depending on background, skills, and experience, the expected annual salary range for this position is $200,000 - $400,000 USD + equity.
- Visa sponsorship: We sponsor visas on a case-by-case basis.
- Benefits: Inferact offers generous health, dental, and vision benefits as well as 401(k) company match.
Inferact, Inc. is an E-Verify company, and you may review the Notice of E-Verify Participation and the Right to Work posters in English and Spanish.
Compensation: $200K - $400K
