Member of Technical Staff — Inference-TPU
Palo Alto, CAFull-time$200–400K/yrPosted todayStill listed today
Most applications go out cold — see where you stand first. No sign-up to start.
You apply on the employer's site, with no Olive account needed.
Watch jobs like this. New roles like this one near Palo Alto, CA, by email.
Don't just apply. Show up ready.
Olive works from this exact posting.
At a glance
Olive lists jobs from US employers, including remote roles you can work from the United States.
Job overview
RadixArk seeks a Member of Technical Staff to build high-performance inference and training systems using JAX, XLA, and Pallas. The role focuses on optimizing large-model workloads on Google TPU hardware, including SGLang-JAX, and improving LLM serving latency and throughput. The engineer will work across compiler, runtime, kernel, and distributed systems layers, collaborating with kernel engineers and compiler teams and contributing to open-source TPU optimization resources.
Skills & qualifications
Skills
Qualifications
Full job description
About the Role RadixArk is looking for a Member of Technical Staff — TPU Systems to build high-performance inference and training systems using JAX, XLA, and Pallas. You'll push model workloads to their limits on TPU hardware, working on SGLang-JAX and other critical infrastructure that enables efficient deployment of frontier models on Google's tensor processing units. Requirements
-
3+ years experience building production ML systems utilizing JAX/Torch, XLA, or TPU-focused frameworks.
-
Bachelor's or Master's degree in Computer Science, Electrical Engineering, or equivalent industry experience
-
Deep understanding of XLA internals preferred: HLO, MLIR, operator fusion, SPMD partitioning, and sharding strategies.
-
Strong performance tuning instincts across compiler and runtime layers
-
Experience with distributed inference systems (e.g. SGLang, vLLM) or training frameworks (e.g. Miles, Alpa, Pathways)
-
Proficiency in Python with demonstrated ability to write high-performance, production-quality code
-
Experience writing custom GPU/TPU/AI Accelerator kernels. Familiarity with Pallas for kernel development is strongly preferred.
Responsibilities
-
Build high-performance inference and training systems using JAX/XLA/Pallas, including SGLang-JAX
-
Push large-model workloads to the limits on the newest TPU hardwares
-
Optimize end-to-end latency and throughput for LLM serving on TPU infrastructure
-
Design and implement SPMD strategies for efficient distributed inference and training
-
Design and implement Pallas kernels for operations that require customized low level control for best performance
-
Profile and optimize XLA compilation pipelines and HLO graph transformations
-
Collaborate with kernel engineers and compiler teams to achieve performance wins across the stack
-
Contribute to open-source projects with TPU optimization guides, benchmarks, and architectural insights
About RadixArk RadixArk is an infrastructure-first company built by engineers who've shipped production AI systems, created SGLang (30K+ GitHub stars, the fastest open LLM serving engine), and developed Miles (our large-scale RL framework). Founded by AI infrastructure veterans from xAI and NVIDIA, we're on a mission to democratize frontier-level AI infrastructure by building world-class open systems for inference and training. Our team has optimized kernels serving billions of tokens daily, designed distributed training systems coordinating 10,000+ GPUs, and contributed to infrastructure that powers leading AI companies and research labs. Compensation Depending on background, skills, and experience, the expected annual salary range for this position is $200,000 - $400,000 USD + equity. Equal Opportunity RadixArk is an Equal Opportunity Employer and is proud to offer equal employment opportunity to everyone regardless of race, color, ancestry, religion, sex, national origin, sexual orientation, age, citizenship, marital status, disability, gender identity, veteran status, and more.
How this job compares
Set against the 6,178 open Machine Learning Engineer jobs on Olive Jobs, counted November 8, 2026.
- Pay
- This posting$200–400K/yr
- California median$180–263K/yr1,442 salary ranges
- US median$151–246K/yr5,016 salary ranges
This range's midpoint, $300K, is 33% above the California median midpoint ($225K) — in the top quarter of Machine Learning Engineer salary ranges in California.
- Openings
- Olive Jobs lists 1,918 open Machine Learning Engineer jobs in California and 6,178 across the US. 8% of them nationwide are listed as remote or flexible.
- Experience
- This posting asks for at least 3 years of experience; Machine Learning Engineer postings that state a minimum ask for 5 years at the median (4,858 postings).
- Education
- 2,039 of the 6,178 open Machine Learning Engineer jobs (33%) ask for a bachelor's degree or higher.
Counted across the open, public US postings on Olive Jobs on November 8, 2026. Pay medians are the middle values of the low ends and of the high ends of postings that list both ends of a range in the same cadence (a yearly salary or an hourly rate); a range's midpoint is halfway between its ends. Pay and experience figures resting on fewer than 20 postings are left out.
Prepare for this role
You've read the whole posting — now see how you match it.
