Senior Research Engineer, JAX

• Maintain and evolve the JAX training framework for scalability and efficiency in large-scale distributed training runs.
• Optimize production JAX inference systems for speech-to-text models using advanced techniques like continuous batching, model sharding, paged attention, and quantization.
• Refactor and modernize model architectures and infrastructure, translating research prototypes into production-ready systems.
• Investigate and resolve performance bottlenecks across the stack, from low-level kernels (XLA, Pallas) to high-level system design.
• Design and deploy scalable, distributed workloads optimized for TPU and GPU architectures.
• Bridge Research and Engineering teams to ensure seamless knowledge transfer and alignment on technical priorities.

• Expert-level proficiency with JAX and its ecosystem (Flax, Optax, XLA compilation pipeline).
• Strong experience optimizing inference systems for production, ideally with LLMs or speech models.
• Hands-on experience with TPU programming and optimization; GPU/CUDA expertise is also valuable.
• Passion for refactoring and improving existing systems to make code faster, cleaner, and more maintainable.
• Familiarity with modern inference optimization techniques: continuous batching, KV-cache management, sharding strategies, quantization.
• Domain knowledge in Speech-to-Text (ASR architectures, audio processing, streaming inference) is a plus.
• Strong Python skills; C++ or Rust experience for kernel-level work is a plus.
• Deep understanding of distributed training at scale and ML infrastructure best practices.
• Excellent communication skills and a collaborative mindset to clearly explain complex tradeoffs and prioritize high-impact work.

Pay range: $190K - $248K