Quick overview
Job description
About the Role
Join an AI team building efficient training and inference systems for image, video, and world-model workloads. You will improve performance across GPU kernels, distributed execution, and multi-node serving.
What You'll Do
Optimize system and GPU performance for training and inference across image, video, and world-model workloads.
Profile and address bottlenecks at the kernel, memory, system, and cluster levels using Nsight and related tools.
Implement low-level optimizations with CUDA and Triton.
Design and improve distributed training and inference engines for diffusion models.
Improve communication performance across GPUs and nodes using NCCL, RDMA, InfiniBand, RoCE, and disaggregated serving approaches.
Build benchmarking and regression harnesses to preserve performance gains in production.
Collaborate on hardware-aware kernel, runtime, and model design.
What We're Looking For
Two or more years of experience in deep learning training or inference systems, distributed systems, or related ML systems work.
At least one year of hands-on ML systems or GPU performance work in a current or recent role.
Experience authoring core features in a training or inference framework, rather than focusing only on deployment or integration.
Direct GPU kernel development experience with CUDA, CUTLASS, Triton, or PTX on NVIDIA GPUs.
Experience improving inference or training performance across GPU kernels, runtime systems, or distributed execution.
Strong machine learning and computer science fundamentals, with a degree in computer science or a related quantitative field.
Experience optimizing diffusion, video, image, or other multimodal workloads is relevant.
Compensation & Benefits
Base salary: $180,000 to $260,000 annually. Visa sponsorship is available.
Location
On-site in Menlo Park, California, United States.
Haystack's AI agent can fill in and submit this application for you. How AI auto apply works