Software Engineer, DL Kernel Libraries

What you'll do

• Design and implement optimized GPU kernels for high-impact AI workloads.
• Develop new AI systems technologies for efficient deep learning inference and training.
• Build extensible abstractions for serving engines and training frameworks.
• Create efficient just-in-time domain-specific compilers and runtimes.
• Develop libraries and code generators for NVIDIA's hardware architecture.
• Contribute to open source projects including FlashInfer, cuDNN, vLLM, and PyTorch.
• Optimize performance for large language models and other high-impact AI workloads.

What we're looking for

• BS, MS, or PhD degree in Computer Science, Electrical Engineering, or related field (or equivalent experience).
• Strong experience in developing or using deep learning frameworks like PyTorch, JAX, TensorFlow, or ONNX.
• Experience using inference engines and runtimes such as vLLM, SGLang, and MLC.
• Strong Python and C/C++ programming skills.
• Background in domain specific compiler and library solutions for LLM inference and training (preferred).
• Expertise in training frameworks and inference engines (preferred).
• Expertise in machine learning compilers such as Apache TVM or MLIR (preferred).
• Strong experience in GPU kernel development and performance optimizations using CUDA C/C++, CUTLASS DSL, cuTile, Triton, or similar (preferred).