# Software Engineer - AI Compiler and Distributed Systems

**Company**: NVIDIA
**Location**: Santa Clara, CA
**Experience**: senior
**Job type**: full-time
**Salary**: 184,000 USD – 287,500 USD
**Category**: Engineering
**Industry**: Technology

**Apply**: https://nvidia.wd5.myworkdayjobs.com/en-US/NVIDIAExternalCareerSite/job/US-CA-Santa-Clara/Senior-Compiler-Engineer---AI_JR2029239-1?utm_source=yubhub.co&utm_medium=jobs_feed&utm_campaign=apply
**Canonical**: https://yubhub.co/jobs/job_ca3b7247-c0c

## Description

NVIDIA is seeking a world-class Software Engineer with AI Compiler and distributed systems experience for its Deep Learning Compiler Technology team. The team builds JaxPP, NVIDIA's open-source library for pipeline parallelism in JAX.

Responsibilities:

- Design and implement program transformations on JAX programs (jaxpr) for large-scale training

- Develop technologies and products involving optimizing compilers and runtime execution frameworks for accelerating deep learning training on large scale distributed systems

- Work on compiling a graph intermediate representation (IR) to target specific code

- Contribute to advancing the state of the art in compiler technology for delivering high performance on large-scale training workloads

- Collaborate with members of the software framework teams and the hardware architecture teams to accelerate the next generation of deep learning software

- Be a core contributor to the open-source JaxPP project

Requirements:

- Masters or PhD in Computer Science, Computer Engineering, or related field or equivalent experience

- 6+ years of relevant work or research experience in compilers, distributed systems, or deep learning training

- Compiler design and implementation

- Distributed system design, development, and deployment

- Experience with deep learning frameworks (JAX, PyTorch)

- GPU training and inference of deep learning models

- C/C++ and Python programming, with excellent software design skills

- Experience with parallel computing libraries (MPI)

- Strong communication skills

Nice to Have:

- Experience with JAX internals (jaxpr, tracing, sharding, shard_map), or with XLA, StableHLO, or MLIR

- Experience building or tuning large training jobs with pipeline, tensor, expert, and data parallelism

- Knowledge of GPU communication libraries such as NCCL, NVSHMEM, or RDMA, and of CPU and/or GPU architecture

- Contributions to open-source compiler or machine learning projects

Benefits:

- Competitive salaries

- Generous benefits package

- Equity

The expected base salary range for this position is 184,000 USD – 287,500 USD.

## Skills

### Required
- C/C++
- Python
- compilers
- distributed systems
- deep learning
- JAX
- PyTorch
- GPU training
- parallel computing libraries

### Nice to have
- JAX internals
- XLA
- StableHLO
- MLIR
- GPU communication libraries
- open-source compiler or machine learning projects

---

Source: [Apply at nvidia.wd5.myworkdayjobs.com](https://nvidia.wd5.myworkdayjobs.com/en-US/NVIDIAExternalCareerSite/job/US-CA-Santa-Clara/Senior-Compiler-Engineer---AI_JR2029239-1?utm_source=yubhub.co&utm_medium=jobs_feed&utm_campaign=apply)