About The Role
RadixArk is looking for a TPU Systems Engineer to build high-performance inference and training systems using JAX, XLA, and Pallas. You'll push large-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 with JAX, XLA, or TPU-focused frameworks
Bachelor's or Master's degree in Computer Science, Electrical Engineering, or equivalent industry experience
Deep understanding of JAX/XLA internals: HLO, 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
Familiarity with Pallas for kernel development, or strong ability to learn quickly
Responsibilities
Build high-performance inference and training systems using JAX/XLA/Pallas, including SGLang-JAX
Push large-model workloads to the limits on TPU v4, v5e, and v5p architectures
Optimize end-to-end latency and throughput for LLM serving on TPU infrastructure
Design and implement SPMD strategies for efficient distributed inference and training
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
Create testing frameworks for numerical correctness and performance regression detection
About RadixArk
RadixArk is an infrastructure-first company built by engineers who've shipped production AI systems at xAI, created SGLang (20K+ GitHub stars, the fastest open LLM serving engine), and developed Miles (our large-scale RL framework). 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. We're backed by well-known investors in the infrastructure field and partner with Google, AWS, and frontier AI labs. Join us in building infrastructure that gives real leverage back to the AI community.
Compensation
We offer competitive compensation with significant founding team equity, comprehensive health benefits, and flexible work arrangements. The US base salary range for this full-time position is: $180,000 - $250,000 + equity + benefits. Our salary ranges are determined by location, level, and role. Individual compensation will be determined by experience, skills, and demonstrated expertise in GPU computing and ML systems.
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.
See other positions
