Senior Compiler Engineer - AI

Details

NVIDIA's Deep Learning Compiler Technology team is hiring a software engineer with AI compiler and distributed systems experience. The team builds JaxPP, NVIDIA's open-source library for pipeline parallelism in JAX. JaxPP transforms a JAX training program into per-stage tasks and runs them across multiple GPUs, enabling large models to be trained on NVIDIA GPU clusters using unmodified JAX and XLA.

Responsibilities

• Design and implement program transformations on JAX programs (jaxpr) for large-scale training, including pipeline schedules, task placement across stages, sharding—including JAX explicit sharding—and overlapping communication with computation.
• Develop optimizing compiler and runtime execution technologies for accelerating deep learning training on large-scale distributed systems.
• Compile graph intermediate representations (IR) to target-specific code and build distributed runtimes that dispatch per-stage programs to multiple GPUs and move data between stages using communication libraries such as NCCL.
• Advance compiler technology for high-performance large-scale training workloads, including mixture-of-experts models.
• Collaborate with software framework teams, including JAX and XLA, and hardware architecture teams to accelerate deep learning software.
• Integrate JaxPP into JAX training frameworks, including MaxText.
• Contribute to the open-source JaxPP project by maintaining compatibility with new JAX releases, testing, tracking performance, and supporting teams adopting it for training runs.

Requirements

• Master's or PhD in Computer Science, Computer Engineering, or a related field, or equivalent experience.
• 6+ years of relevant work or research experience in compilers, distributed systems, or deep learning training.
• Experience designing and implementing compilers.
• Experience designing, developing, and deploying distributed systems.
• Experience with deep learning frameworks such as JAX and PyTorch.
• Experience with GPU training and inference of deep learning models.
• Proficiency in C/C++ and Python, with excellent software design skills, including debugging, performance analysis, and test design.
• Experience with parallel computing libraries such as MPI.
• Ability to work independently, define project goals and scope, and lead development efforts.
• Strong communication skills, including the ability to work directly with engineers at partner companies to debug and tune large training runs.
• Ability to work effectively in a dynamic, product-oriented team.

Preferred Qualifications

• Experience with JAX internals, including jaxpr, tracing, sharding, and shard_map, or with XLA, StableHLO, or MLIR.
• Experience building or tuning large training jobs with pipeline, tensor, expert, and data parallelism, including work with MaxText or Megatron-LM.
• Experience debugging the performance of large training jobs on GPU clusters.
• Knowledge of GPU communication libraries such as NCCL, NVSHMEM, or RDMA, and of CPU and/or GPU architecture.
• CUDA programming experience.
• Contributions to open-source compiler or machine learning projects.
• Experience mentoring early-career engineers and interns.

Benefits

The position offers equity and benefits in addition to the base salary. The expected base salary range is USD 184,000–287,500, determined by location, experience, and other objective job-related factors.