Senior Compiler Engineer - AI

NVIDIA•Redmond, WA

3d•$184,000 - $287,500

About The Position

NVIDIA is hiring a Senior Compiler Engineer for its Deep Learning Compiler Technology team. The team develops JaxPP, NVIDIA’s open-source library for pipeline parallelism in JAX, which transforms JAX training programs into per-stage tasks and runs them across multiple GPUs. JaxPP enables large models to be trained on NVIDIA GPU clusters using unmodified JAX and XLA. The role focuses on compiler and runtime technology for large-scale deep learning training. The engineer will design transformations for JAX programs, compile graph intermediate representations to target code, and develop distributed runtime systems that dispatch programs across GPUs and move data between stages. The work includes advancing performance for large training workloads, including mixture-of-experts models, and collaborating with JAX, XLA, and hardware architecture teams. The engineer will also contribute to the open-source JaxPP project and its integration with JAX training frameworks such as MaxText. The expected base salary range is $184,000–$287,500 USD. The employee will also be eligible for equity and benefits. Applications will be accepted at least until November 6, 2026. This posting is for an existing vacancy.

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 with compiler design and implementation.
Experience with distributed system design, development, and deployment.
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.
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 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 communication/computation overlap.
Develop optimizing compiler and runtime technologies for accelerating deep learning training on large 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 with hardware architecture teams.
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.
Define project goals and scope and lead independent development efforts.