About the job
We are looking for versatile software engineers for our XLA team. NVIDIA is at the center for the AI revolution that's transforming how people live, work, and interact with technology. Come join us to build high-performance, production-grade software that's at the core of next-generation AI systems.
Responsibilities
Develop compiler optimization algorithms for deep learning workloads
Optimize inference and training performance for the JAX framework and the OpenXLA compiler on NVIDIA GPUs at scale
Craft and implement compiler optimization techniques for deep learning network graphs
Design novel graph partitioning and tensor sharding techniques for distributed training and inference
Perform performance tuning and analysis
Generate code for NVIDIA GPU backends using open-source compilers such as MLIR, LLVM and OpenAI Triton
Qualifications
Minimum
Bachelors, Masters or Ph.D. in Computer Science, Computer Engineering, related field (or equivalent experience)
4+ years of relevant work or research experience in performance analysis and compiler optimizations
Ability to work independently, define project goals and scope, and lead your own development effort adopting clean software engineering and testing practices
Excellent C/C++ programming and software design skills, including debugging, performance analysis, and test design
Strong foundation in architecture of CPU, GPUs or other high performance hardware accelerators
Knowledge of high-performance computing and distributed programming
Strong interpersonal skills and ability to work in a dynamic product-oriented team
Preferred
CUDA or OpenCL programming experience
Experience with XLA, TVM, MLIR, LLVM, OpenAI Triton, deep learning models and algorithms, and deep learning framework design
History of mentoring junior engineers and interns
Experience working deep learning frameworks such as JAX, PyTorch or TensorFlow
Extensive experience with CUDA or with GPUs in general
Experience with open-source compilers such as XLA, LLVM, MLIR or TVM