Member of Technical Staff — Inference-TPU
RadixArkPalo Alto, California, United States · Posted 8 months agoDescription
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.
Similar jobs
Est. 300,000 USD
About the Role RadixArk is seeking a Member of Technical Staff — Inference to push the limits of large-scale AI inference. You will work on the core systems that serve frontier models at scale, optimizing performance, la…
Est. 300,000 USD
About the Role RadixArk is seeking a Member of Technical Staff - Inference-Multi-Hardware to push the limits of performance for frontier AI systems. Most performance engineering assumes a single vendor's stack. This role…
Est. 300,000 USD
About the Role RadixArk is seeking a Member of Technical Staff, Developer Technology (DevTech) to make LLM inference and training dramatically faster, cheaper, and more accessible on modern GPU hardware. Our systems sit…
Est. 300,000 USD
About the Role RadixArk is looking for a Member of Technical Staff Cluster Infrastructure to architect and scale the core compute platform that powers frontier-level AI training and inference. You will design and operate…
Est. 300,000 USD
About the Role RadixArk is hiring a Member of Technical Staff — Performance in Palo Alto, CA — someone who can push LLM inference and training systems to the limit across real production workloads. You’ll work on the per…
Est. 300,000 USD
About the Role RadixArk is seeking a Member of Technical Staff — Kernel / Compiler / Communication to push the limits of performance for frontier AI systems. You will work at the lowest layers of the stack — kernels, run…
Est. 300,000 USD
About the Role RadixArk is looking for a Member of Technical Staff — Backend/API Platform Engineer to build the API layer, control plane, and platform services that power SGLang and Miles in production. You'll design and…
Est. 300,000 USD
About the Role RadixArk is seeking a Member of Technical Staff — Inference-Multimodal & Diffusion to advance the frontier of generative modeling. You will work on cutting-edge diffusion and flow-based models for imag…
Est. 165,000 USD
About the Role As a Technical Program Manager at RadixArk, you'll drive the execution of complex, cross-functional programs across our inference and training infrastructure. You'll partner closely with Product Management…
Est. 300,000 USD
About the Role RadixArk is seeking experienced product-focused engineers to join our team in building the developer-facing surfaces of our inference and training infrastructure. As a Member of Technical Staff — Product,…
Est. 300,000 USD
About the Role As a Member of Technical Staff, Training, you will design, build, and operate the distributed systems behind large-scale model post-training — spanning training, inference, and orchestration, with a focus…
Est. 300,000 USD
About the Role RadixArk is seeking a Developer Advocate to build and engage our technical community around SGLang, Miles, and our open source infrastructure. SGLang already has 30K+ GitHub stars and serves billions of to…
Est. 300,000 USD
About the Role RadixArk is hiring a Member of Technical Staff — CI / Infrastructure to own the infrastructure that keeps SGLang moving. Our CI system runs 300+ GPU tests across NVIDIA, AMD, Intel, and Ascend hardware poo…
Est. 144,000 USD
Key Responsibilities Product Strategy & Roadmap Define, prioritize, and drive the product roadmap for inference and training infrastructure. Stay ahead of AI trends, including new model architectures, hardware optimi…
Est. 140,000 USD
About The Role RadixArk is launching a full-time, paid, 1-year residency program for aspiring AI infrastructure engineers. You'll rotate across inference, training, kernels, compilers, and cluster infrastructure, working…
Est. 155,000 USD
About the Role RadixArk is seeking a Product Marketing Manager to own how SGLang, Miles, and our open source infrastructure are positioned and perceived across the market. SGLang already has 20K+ GitHub stars and serves…
Est. 150,000 USD
About the Role We're looking for a Head of Business Development to build the BD function at RadixArk from the ground up. The BD team is the institutional memory of this company — maintaining active relationships across e…
Est. 159,000 USD
About the Role We're looking for a hands-on Talent Operations Specialist to build and run the machinery behind talent and people ops as we scale. This isn't a traditional HR generalist role - it's for someone who treats…
Est. 289,800 USD
XPENG is a leading smart technology company at the forefront of innovation, integrating advanced AI and autonomous driving technologies into its vehicles, including electric vehicles (EVs), electric vertical take-off and…
Est. 235,200 USD
XPENG is a leading smart technology company at the forefront of innovation, integrating advanced AI and autonomous driving technologies into its vehicles, including electric vehicles (EVs), electric vertical take-off and…
Est. 310,000 USD
SpaceXAI’s mission is to create AI systems that can accurately understand the universe and aid humanity in its pursuit of knowledge. Our team is small, highly motivated, and focused on engineering excellence. This organi…
Est. 310,000 USD
At River AI, our mission is to create personal AI owned and shaped by each individual. To achieve this, we are rewriting the entire stack from scratch: personal hardware for local inference, custom training infrastructur…
Est. 310,000 USD
At River AI, our mission is to create personal AI owned and shaped by each individual. To achieve this, we are rewriting the entire stack from scratch: personal hardware for local inference, bespoke training infrastructu…
Est. 115,000 USD
About the Role RadixArk builds the open-source AI infrastructure behind SGLang and Miles, used by developers and enterprises around the world. We're looking for a visual designer to join our design team and to give our b…
Est. 120,000 GBP
The Firm XTX Markets is a leading algorithmic trading firm which uses state-of-the-art machine learning technology to produce price forecasts for over 50,000 financial instruments across equities, fixed income, currencie…
Est. 289,800 USD
XPENG is a leading smart technology company at the forefront of innovation, integrating advanced AI and autonomous driving technologies into its vehicles, including electric vehicles (EVs), electric vertical take-off and…
Est. 237,500 USD
Who We Are Lightning AI is the company behind PyTorch Lightning. Founded in 2019, we build an end-to-end platform for developing, training, and deploying AI systems—designed to take ideas from research to production with…
Est. 235,200 USD
XPENG is a leading smart technology company at the forefront of innovation, integrating advanced AI and autonomous driving technologies into its vehicles, including electric vehicles (EVs), electric vertical take-off and…
Est. 222,500 USD
About Turing Turing’s mission is to accelerate superintelligence to drive real economic progress. Headquartered in San Francisco, Turing works with frontier AI labs to generate high-quality datasets, reinforcement learni…
Est. 310,000 USD
SpaceXAI’s mission is to create AI systems that can accurately understand the universe and aid humanity in its pursuit of knowledge. Our team is small, highly motivated, and focused on engineering excellence. This organi…