Senior Compiler Engineer - AI
Core
Design and implement program transformations on JAX programs (jaxpr) for large-scale training, including pipeline schedules, task placement, sharding, and communication/computation overlap.
Role type
Senior IC compiler engineer (distributed deep learning)
Builds
JaxPP, NVIDIA's open-source library for pipeline parallelism in JAX, and distributed runtime frameworks for GPU clusters.
Domain
Deep learning, distributed systems, compiler technology
Deliverable
production ML models
Required skills
Compiler design and implementation, Distributed system design and deployment, Deep learning frameworks (JAX, PyTorch), GPU training and inference, C/C++ and Python programming, Parallel computing libraries (MPI)
Preferred skills
JAX internals (jaxpr, tracing, sharding), XLA, StableHLO, MLIR, GPU communication libraries (NCCL, NVSHMEM, RDMA), CUDA programming, Mentoring
Technologies
JAX, XLA, NCCL, CUDA, Python, C/C++, MPI, MaxText, Megatron-LM
Responsibilities
Design and implement program transformations on JAX programs for large-scale training; Develop technologies and products involving optimizing compilers and runtime execution frameworks; Work on compiling a graph intermediate representation (IR) to target specific code and build distributed runtime; Collaborate with software framework and hardware architecture teams; Contribute to the open-source JaxPP project
Seniority
Senior, hands-on IC
