TPU Systems Engineer (JAX / XLA / Pallas)

About the Role

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, 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. 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 TPU systems and ML infrastructure.

  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.