← All roles

Member of Technical Staff — Inference-TPU

RadixArk· Member of Technical Staff
Apply on company site
Palo Alto, CA$200k - $400kfull-timePosted 5mo ago

About the role

Role focused on building high-performance inference and training systems for TPU hardware using JAX, XLA, and Pallas, pushing model workloads to limits and contributing to TPU-focused infrastructure and open-source projects.

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

Qualifications

  • 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

Skills

Required:JAXXLAPythonDistributed systems
Nice to have:Pallas

Benefits

  • Not specified in description
Last updated Aug 18, 2026. Interested? Apply on the company site.
RadixArk logo

About RadixArk

SHIP AI FOR ALL

RadixArk is an infrastructure-first, deep-tech company building large-scale inference and training systems for the entire AI community.

AI systemsMilesSGLangbessemer-ai-infrastructure-2026inferenceinference-inflection-pointinfrastructureopen infrastructure
Company profileWebsite

More open roles at RadixArk