MEMBER OF TECHNICAL STAFF, JAX & COMPILER

xAI
Full-time
Palo Alto, CA; San Francisco, CA
$180,000 - $440,000
Posted on a month ago

Job Description

xAI is seeking a highly motivated engineer to develop in-house JAX or compiler/runtime features for large-scale LLM training, debug distributed programs, and support cutting-edge LLM research. The company values engineering excellence, curiosity, and strong communication skills.

Responsibilities

  • Develop JAX or compiler/runtime features
  • Develop and use debugging/profiling toolings
  • Understand CUDA runtime interactions
  • Support LLM research/production tasks

Requirements

  • JAX/XLA expertise
  • Compiler optimization understanding
  • Distributed LLM training experience
  • Kernel compiler experience
  • Excellent communication skills
  • Python/C++ proficiency

Benefits

  • No benefits