Principal Engineer, Frameworks

We are hiring a Principal Engineer to oversee developing our AI platforms, with a heavy emphasis on large language models (LLMs) such as PAXML. Work with a dedicated group of engineers in a fast-paced production environment to make a tangible impact to Deep Learning Frameworks, particularly JAX. Come join us in these exciting times and make a sizable difference in the exploding world of AI! As an all-stack engineer with a focus on the back-end libraries and compilers, you will play a meaningful role in NVIDIA's efforts in contributing to the Deep Learning ecosystem.What You Will Be Doing:Oversee and lead design and development of Deep Learning platform solutions with our Google partnersDesign and build core Frameworks and Compilers and drive peak performance on NVIDIA GPUsEvaluate and improve the end-to-end performance of generative AI models, especially large-language models (e.g. PAXML), on NVIDIA hardwareWork with AI applied researchers and leaders to build highly-efficient ML modelsBuild tools that will increase the efficiency of teams developing AI-based systems.Work to bridge the gap between the latest in numerical computing, simulation and deep learning research and their applications in real world products.Impact the design of next generation of accelerated computing software and hardwareWhat We Need To See:BS in Computer Science or Computer Engineering or related field (or equivalent experience)15+ years relevant experienceBackground in C++, Python, MLIR, LLVMExperience with at least one machine learning framework, e.g. JAX, PyTorch, TensorFlowFamiliar with such concepts as asynchronous computing, sharding, collectives, various flavors of parallelism, GPU kernels, ML building blocks (e.g. Multi-Head Attention, MLP)Proven technical foundation in GPU architectures, numeric libraries, modular software design.Highly motivated with excellent verbal and written communication skills.Ability to work collaboratively with and lead multi-functional teams. Coordinates effectively across organizational boundaries.Ways To Stand Out From The Crowd:Prior experience in MLIR, XLA, JAX, LLVM, and CUDAFirm grasp of distribution computation such as data, model, ZeRO, and pipeline parallelismBeing intimately familiar with Google C++ Style Guide and having contributed to OSSCollaborative and production-minded demeanorThe base salary range is $268,000 - $414,000. Your base salary will be determined based on your location, experience, and the pay of employees in similar positions.You will also be eligible for equity and benefits.NVIDIA is committed to fostering a diverse work environment and proud to be an equal opportunity employer. As we highly value diversity in our current and future employees, we do not discriminate (including in our hiring and promotion practices) on the basis of race, religion, color, national origin, gender, gender expression, sexual orientation, age, marital status, veteran status, disability status or any other characteristic protected by law.