Member of Technical Staff — Inference-TPU
RadixArk
Salary undisclosed
Need a reasonable accommodation to apply or interview? Contact us.
JobMinglr uses automated technology to recommend jobs based on profile information and job preferences. Match Score does not determine eligibility for a position, prevent a user from viewing or applying to a job, or make hiring decisions on behalf of an employer.
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.
About this role
RadixArk is building high-performance inference and training systems on Google's TPU hardware, and they're looking for someone with deep expertise in JAX, XLA, and low-level compiler optimization to make it happen. You'd work on SGLang-JAX and related infrastructure, tackling the hard problems of pushing large language models to their limits—designing SPMD strategies for distributed inference, writing custom Pallas kernels, and optimizing XLA compilation pipelines. This is hands-on systems work that spans the full stack from compiler internals to runtime performance.
The role expects 3+ years building production ML systems with JAX or Torch, solid understanding of XLA internals (HLO, MLIR, operator fusion, sharding), and experience with distributed inference or training frameworks. You should be comfortable writing high-performance Python and custom accelerator kernels; Pallas experience is strongly preferred. A CS or EE degree (or equivalent industry track record) is required.
RadixArk was founded by infrastructure engineers from xAI and NVIDIA who created SGLang and Miles. The team has optimized kernels at scale and designed systems coordinating thousands of GPUs. You'd be in Palo Alto, in-office, with a salary range of $200K–$400K plus equity depending on background.
How this employer is doing
average
- Sponsors Visas
This role's local market on JobMinglr
Pay for this role
The employer didn't post a pay range for this role. That usually means pay is set in negotiation, which favors whoever arrives with numbers. Check ranges on comparable Member of Technical Staff — Inference-TPU postings in Palo Alto, and analyze any offer before you accept it.
Before you apply, worth reading
How JobMinglr reads this job
Every listing here is scored against your profile before you apply: skills overlap, experience level, location and work arrangement, each weighted and explained. You see the score and the reasons, not just a list. How the matching works.