Senior Compiler Engineer - AI

NVIDIA • Santa Clara, CA

2d

About The Position

NVIDIA is hiring a Senior Compiler Engineer for its Deep Learning Compiler Technology team, which develops JaxPP, an open-source library for pipeline parallelism in JAX. JaxPP converts JAX training programs 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.

The role focuses on compiler and runtime technologies for large-scale deep learning training on distributed systems. The engineer will develop program transformations and distributed execution capabilities, contribute to JaxPP, and collaborate with software framework and hardware architecture teams to accelerate deep learning software. The position also involves integrating JaxPP into JAX training frameworks, including MaxText, and supporting teams adopting the project.

Requirements

• Master’s or PhD in Computer Science, Computer Engineering, or a related field, or equivalent experience.
• At least 6 years of relevant work or research experience in compilers, distributed systems, or deep learning training.
• Experience in compiler design and implementation.
• Experience designing, developing, and deploying distributed systems.
• Experience with deep learning frameworks such as JAX or PyTorch.
• Experience with GPU training and inference of deep learning models.
• C/C++ and Python programming skills, with strong 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 and ability to work in a dynamic, product-oriented team.

Nice To Haves

• Experience with JAX internals, including jaxpr, tracing, sharding, or shard_map, or with XLA, StableHLO, or MLIR.
• Experience building or tuning large training jobs using pipeline, tensor, expert, and data parallelism, for example in MaxText or Megatron-LM, and debugging performance on large 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.

Responsibilities

• Design and implement transformations of JAX programs (jaxpr) for large-scale training, including pipeline schedules, task placement across stages, sharding, and overlapping communication with computation.
• Develop optimizing compiler and runtime technologies for deep learning training on large-scale distributed systems.
• Compile graph intermediate representations to target-specific code and build distributed runtime systems that dispatch per-stage programs to 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, such as 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 working with teams adopting it for training runs.
• Work directly with engineers at partner companies to debug and tune large training runs.

Benefits

• Equity
• Benefits package

What This Job Offers

• Job Type: Full-time
• Career Level: Senior