AI Software Lead – PyTorch & CUDA Runtime (Next-Gen Accelerator)

Role Overview

We are looking for a Software Lead (8+ years’ experience) to own the runtime and neural network (NN) layer of a next-generation AI accelerator platform. This role focuses on designing, optimizing, and implementing NN operators and developing new ops using CUDA/custom runtime APIs to deliver high-performance execution on custom AI hardware.

Key Responsibilities

Design and optimize NN operators for performance-critical workloads

Develop new NN ops using CUDA/custom runtime APIs

Drive runtime-level optimizations across compute, memory, and scheduling

Own runtime ↔ NN layer interfaces and execution model

Implement and optimize operator fusion (e.g., matmul + bias + LayerNorm) for efficient hardware utilization

Identify and resolve performance bottlenecks across the stack

Collaborate with compiler, PyTorch framework, and low-level SW teams

Impact

Own how efficiently AI workloads execute on the platform

Drive performance, scalability, and hardware utilization through optimized runtime and NN ops design

Required Qualifications

8+ years in systems software / runtime / performance engineering

Strong experience with:PyTorch / TensorFlow / JAX or similar frameworks

NN operator/kernel development and optimization

operator fusion and graph-level optimizations

Hands-on expertise in:C/C++ and CUDA (or similar low-level programming)

runtime systems and execution engines

Strong understanding of:memory hierarchy, data movement, and parallel execution

Preferred Qualifications

Experience with GPU/NPU or custom AI accelerators

Familiarity with XLA / MLIR / compiler-runtime interaction

Experience optimizing LLM or large-scale DL workloads

Sandisk is committed to providing equal opportunities to all applicants and employees and will not discriminate against any applicant or employee based on their race, color, ancestry, religion (including religious dress and grooming standards), sex (including pregnancy, childbirth or related medical conditions, breastfeeding or related medical conditions), gender (including a person’s gender identity, gender expression, and gender-related appearance and behavior, whether or not stereotypically associated with the person’s assigned sex at birth), age, national origin, sexual orientation, medical condition, marital status (including domestic partnership status), physical disability, mental disability, medical condition, genetic information, protected medical and family care leave, Civil Air Patrol status, military and veteran status, or other legally protected characteristics. We also prohibit harassment of any individual on any of the characteristics listed above. Our non-discrimination policy applies to all aspects of employment. We comply with the laws and regulations set forth in the "Know Your Rights: Workplace Discrimination is Illegal” poster. Our pay transparency policy is available here.

Sandisk thrives on the power and potential of diversity. As a global company, we believe the most effective way t