Member of Technical Staff - RL Inference

About the Role

The RL infrastructure team at xAI is seeking an engineer to contribute to low-precision RL training and inference.

Responsibilities

• Design and optimize the inference stack for various RL workloads, from small-scale ablations to production training runs.
• Analyse, profile, and address performance bottlenecks in large-scale RL systems.
• Collaborate with the modelling team to implement novel RL techniques and algorithms efficiently.

Basic Qualifications

• Experience in building, debugging, and optimizing large-scale distributed systems.
• Experience in LLM inference.
• Proficiency in programming languages such as Python, C++, and/or Rust; frameworks such as PyTorch, Jax, CUDA.
• Willingness to tackle complex problems at all levels of the stack.

Preferred Skills and Experience

• Strong knowledge in quantization and numerics in LLM inference and training.
• Experience in developing inference engines, e.g., SGLang, vLLM.