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

Sandisk; Bengaluru; On-site; INR 3,500,000 - 7,000,000; Full time

Job summary

Sandisk in Bengaluru seeks an AI Software Lead to own the runtime and NN layer of a next-gen accelerator. Lead design and optimization of NN operators and new ops using CUDA and custom runtime APIs for high-performance execution on AI hardware.

You will collaborate with compiler, PyTorch framework, and low-level teams to drive runtime improvements, operator fusion, and scalable performance across architectures at scale. 10+ years experience preferred.

Qualifications

• 8+ years in systems software, runtime or performance engineering.
• Experience with PyTorch, TensorFlow, or JAX.
• NN operator/kernel development and optimization.
• Experience with operator fusion and graph-level optimizations.
• Proficiency in C/C++ and CUDA.
• Understanding memory hierarchy and parallel execution.

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 for efficient hardware utilization.
• Identify and resolve performance bottlenecks across the stack.
• Collaborate with compiler, PyTorch framework, and low-level SW teams.

Skills

NN operator/kernel development

operator fusion

graph-level optimizations

C/C++

CUDA

PyTorch / TensorFlow / JAX

memory hierarchy

Job description

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

• Full-time
• Job Type (exemption status): Exempt position - Please see related compensation & benefits details below
• Business Function: Firmware Engineering

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
• 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

Equal Employment Opportunity Statement

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 to embrace the diversity of our customers and communities is to mirror it from within. We believe the fusion of various perspectives results in the best outcomes for our employees, our company, our customers, and the world around us. We are committed to an inclusive environment where every individual can thrive through a sense of belonging, respect and contribution.

Sandisk is committed to offering opportunities to applicants with disabilities and ensuring all candidates can successfully navigate our careers website and our hiring process. Please contact us at [e-mail] to advise us of your accommodation request. In your email, please include a description of the specific accommodation you are requesting as well as the job title and requisition number of the position for which you are applying.