Software Engineer, Systems ML - Frameworks/Compilers/Kernels (PhD)

In this role, you will be a member of the PyTorch Accelerator Enablement(AE) team and part of the bigger industry-leading PyTorch AI framework organization. Our team focus on optimizing GenAI model inference on GPUs. As a key member of our Research & Development team, you will work on developing and implementing various optimization techniques at the algorithm, kernel/compiler, and runtime levels. This is an exciting opportunity to contribute to the development of state-of-the-art GenAI models and collaborate with product and research groups across Meta. Our work involves a blend of software systems and ML, and you'll have the opportunity to learn about how PyTorch works at a more in-depth level as well. Come join us!

Software Engineer, Systems ML - Frameworks/Compilers/Kernels (PhD) Responsibilities:

• Research and development of SW stack with one of the following core focus areas: AI frameworks, compiler stack, high performance kernel development and acceleration onto next generation of AI Accelerators architectures.
• Develop custom/novel architectures, define use cases, and develop methodology & benchmarks to evaluate different approaches.
• Collaborate with AI research scientists to accelerate the next generation of deep learning models on the next generation of hardwares.
• Performance tuning and optimizations of deep learning framework & software components.

• Currently has, or is in the process of obtaining a Bachelor's degree in Computer Science, Computer Engineering, relevant technical field, or equivalent practical experience. Degree must be completed prior to joining Meta.
• Currently has, or is in the process of obtaining, a PhD degree in Computer Science, Computer Engineering, relevant technical field, or equivalent practical experience. Degree expected to be completed prior to joining Meta.
• Python and C/C++ programming skills.
• Research experience in AI framework development or accelerating deep learning models on hardware architectures.
• Must obtain work authorization in country of employment at the time of hire and maintain ongoing work authorization during employment.

• Knowledge of GPU, CPU, or AI hardware accelerator architectures.
• Experience working with kernel frameworks like Triton, or CUDA.
• OR AI Compiler: Experience with compiler optimizations such as loop optimizations, vectorization, parallelization, hardware specific optimizations such as SIMD. Experience with MLIR, LLVM, IREE, XLA, TVM, Halide.
• OR AI frameworks: Experience in developing training and inference framework components. Experience in system performance optimizations such as runtime analysis of latency, memory bandwidth, I/O access, compute utilization analysis and associated tooling development.
• Experience developing Large Language Models (LLM) serving frameworks like vLLM or SGLang.
• Demonstrated software engineer experience via an internship, work experience, coding competitions, or widely used contributions in open source repositories (e.g. GitHub).
• Experience working and communicating cross-functionally in a team environment.
• Proven track record of achieving significant results as demonstrated by grants, fellowships, patents, as well as first-authored publications at leading workshops or conferences.