Senior Deep Learning Compiler Engineer - XLA
NVIDIAJob Title
Senior Deep Learning Compiler Engineer - XLA
Role Summary
Develop compiler optimization algorithms and systems that improve performance of deep learning training and inference on NVIDIA GPUs. Work on XLA and JAX, collaborating with deep learning framework teams and GPU hardware architects to deliver production-grade compiler features at scale.
Experience Level
Senior-level. Guideline: 4+ years of relevant work or research experience in compiler optimizations or performance analysis.
Responsibilities
You will design and implement compiler and runtime features to accelerate deep learning workloads.
- Implement compiler optimization techniques for deep learning computation graphs.
- Design graph partitioning and tensor sharding solutions for distributed training and inference.
- Perform performance tuning, profiling, and analysis to identify bottlenecks.
- Generate code for NVIDIA GPU backends using MLIR, LLVM, and Triton.
- Design and implement user-facing features in JAX and related libraries.
- Collaborate with framework teams and GPU hardware architects to align software and hardware capabilities.
Requirements
Must-have technical skills and experience. Education requirements are listed separately below.
- 4+ years of experience in compiler optimizations, performance analysis, or related research.
- Excellent C/C++ programming, software design, testing, and debugging skills.
- Strong foundation in CPU/GPU/hardware accelerator architectures and high-performance computing.
- Experience with performance analysis and distributed programming.
- Ability to work independently, define project scope, and lead development efforts.
- Strong interpersonal skills; experience mentoring junior engineers is a plus.
- Nice-to-have: CUDA or OpenCL experience; familiarity with XLA, TVM, MLIR, LLVM, or OpenAI Triton; experience with JAX, PyTorch, or TensorFlow.
Education Requirements
Bachelor's, Master's, or Ph.D. in Computer Science, Computer Engineering, or a related technical field; or equivalent practical experience.
About the Company
Company: NVIDIA
Headquarters: Santa Clara, California, USA
NVIDIA is a global leader in accelerated computing, renowned for its innovative solutions in AI and digital twins that transform diverse industries. The company specializes in networking technologies, providing end-to-end InfiniBand and Ethernet solutions for servers and storage that optimize performance and scalability. NVIDIA serves sectors such as high-performance computing, enterprise data centers, and cloud computing, constantly reinventing its products and services to stay ahead in the market.
