Senior Staff Engineer PyTorch Accelerator

Job Title: Senior Staff Engineer – PyTorch Accelerator

Location: San Jose, CA (Hybrid)

Duration: 6+ Months

Xoriant reasonably expects the pay rate for this position to be within the following range: $70/hr-$90/hr

Summary:

We are seeking an experienced PyTorch maintainer to bring new AI accelerators to PyTorch and get that work accepted into the official upstream project.

Responsibilities:

• Upstream new hardware support to PyTorch: RFCs, pull requests, code reviews, issue triage, and ongoing maintenance.
• Work with PyTorch maintainers and working groups to align integration with PyTorch architecture, testing standards, and release requirements.
• Design and implement backend integration (PrivateUse1, dispatch keys, operator registration, memory management, streams/events, autograd), partnering with compiler, runtime, kernel, driver, and hardware teams.
• Improve ATen operators, torch.compile, distributed training, mixed precision, and profiling; build and maintain tests, CI workflows, and benchmarks for correctness and performance.
• Debug correctness, performance, memory, build, and compatibility issues across the PyTorch stack; write clear documentation and mentor other engineers.

Basic Qualifications:

• 5+ years of experience in open-source ecosystem development or related roles.
• Strong hands-on experience with PyTorch internals: ATen, the dispatcher, operator registration, autograd, tensor/device abstractions, and Python/C++ extensions.
• Strong C++ and Python skills, with a track record of contributing clean, well-tested code to large open-source projects (design discussions, code reviews, pull requests, CI debugging).
• Experience integrating ML frameworks with hardware accelerators, runtimes, compilers, kernels, or device drivers, and a good understanding of GPU/AI accelerator architecture (memory hierarchy, asynchronous execution, streams/events, profiling).
• Comfortable with PyTorch build/test workflows (Linux, CMake, GitHub, CI) and debugging across Python, C++, runtime, and hardware layers.
• Bachelor's degree in Computer Science, Computer Engineering, Electrical Engineering, or a related field, or equivalent experience.

Preferred Qualifications:

• Recognized PyTorch contributor or maintainer, with changes accepted into the main project.
• Experience with PyTorch accelerator integration: PrivateUse1 backend development, OpenReg-style reference implementations, or out-of-tree backend maintenance.
• Experience with torch.compile (TorchDynamo, TorchInductor), MLIR/LLVM, Triton, or XLA, or optimizing large-scale LLM training (e.g. VeRL) and inference (e.g., vLLM, SGLang).
• Experience building benchmarking, profiling, and performance regression systems; familiarity with distributed backends such as NCCL, RCCL, HCCL, oneCCL, or custom implementations.

Experience working with hardware vendors or open-source foundations.

• Bachelor's degree in Computer Science, Computer Engineering,
• Electrical Engineering, or a related field, or equivalent experience.

Regards,

Akangsha Mohite