Seeking an experienced PyTorch Engineer / Maintainer to help bring new AI accelerator hardware into the PyTorch ecosystem and contribute changes to the official upstream PyTorch project. The ideal candidate will have strong expertise in PyTorch internals, C++, Python, AI accelerators, and open-source development, with experience in areas such as torch.compile, Triton, backend integration, and performance optimization.

Key Responsibilities

• Develop and upstream support for new AI accelerator hardware within PyTorch.
• Participate in RFCs, pull requests, code reviews, issue triage, testing, and ongoing maintenance.
• Work closely with PyTorch maintainers and open-source working groups to align hardware integration with PyTorch architecture and release requirements.
• Design and implement backend integration using PrivateUse1, dispatch keys, operator registration, memory management, streams/events, and autograd.
• Collaborate with compiler, runtime, kernel, driver, and hardware engineering teams.
• Improve and optimize ATen operators, torch.compile, distributed training, mixed-precision workloads, and profiling.
• Develop and maintain unit tests, CI workflows, benchmarks, and performance regression tests.
• Debug correctness, performance, memory, build, and compatibility issues across the PyTorch stack.
• Create technical documentation and mentor other engineers.

Required Qualifications

• 3+ years of relevant experience; 5+ years in open-source ecosystem development or related roles preferred.
• Strong hands-on experience with PyTorch internals, including:
  • ATen
  • Dispatcher
  • Operator registration
  • Autograd
  • Tensor/device abstractions
  • Python/C++ extensions
• Strong C++ and Python programming skills.
• Experience contributing to large open-source projects, including code reviews, pull requests, design discussions, and CI/debugging.
• Experience integrating ML frameworks with AI accelerators, GPUs, runtimes, compilers, kernels, or device drivers.
• Understanding of accelerator architecture concepts such as:
  • Memory hierarchy
  • Asynchronous execution
  • Streams/events
  • Profiling
• Experience with Linux, CMake, GitHub/Git, and CI/build/test workflows.
• Bachelor's degree in Computer Science, Computer Engineering, Electrical Engineering, or related field, or equivalent experience.

Preferred Qualifications

• Recognized PyTorch contributor or maintainer, with contributions accepted into the main PyTorch project.
• Experience developing or maintaining PyTorch accelerator backends, including PrivateUse1 or OpenReg-style implementations.
• Experience with:
  • torch.compile
  • TorchDynamo
  • TorchInductor
  • Triton
  • MLIR/LLVM
  • XLA
• Experience optimizing large-scale LLM training, including frameworks such as vLLM or SGLang.
• Experience building benchmarking, profiling, and performance regression systems.
• Familiarity with distributed training backends such as NCCL, RCCL, HCCL, oneCCL, or custom ProcessGroup implementations.
• Experience working directly with hardware vendors or open-source foundations.

Candidate Profile

The ideal candidate is a strong PyTorch systems engineer who understands both the software framework and underlying accelerator hardware. Candidates with significant PyTorch internals, open-source contribution, accelerator integration, compiler, runtime, or performance optimization experience are encouraged to apply even if they do not have every technology listed above.