Research Engineer - AI Performance & Kernel Optimization
Zyphra· R&D - Engineering
San FranciscoFull-timePosted 7mo ago
About the role
Research Engineer focused on improving and optimizing the performance of large-scale language model training and inference stacks, with emphasis on kernel optimization, performance tuning, and cross-hardware portability across accelerator platforms.
Responsibilities
- Kernel development and optimization for large-scale ML workloads, using any level of the stack from PTX/assembly to CUDA, HIP, Triton, or other GPU DSLs
- Performance tuning for training and inference stacks across GPUs and other accelerators
- Profiling and eliminating bottlenecks in memory movement, communication, scheduling, and compute utilization
- Optimizing distributed training and inference systems for large MoE models, including large-scale model parallelism
- Portability and optimization across non-NVIDIA hardware, with special interest in AMD hardware such as the MI300x and MI355x
- Collaboration with research and infrastructure teams to turn systems improvements into real-world model training and inference gains
Qualifications
- Experience writing highly performant GPU kernels at any level of abstraction–PTX, CUDA, HIP, Triton, or other kernel DSLs
- Experience optimizing ML workloads for large-scale training, ideally in language model pretraining or inference environments
- Experience with non-NVIDIA accelerator hardware, such as AMD, AWS Trainium, Google TPU, Qualcomm, ARM, Intel, and custom ASICs
- Strong understanding of distributed training systems and parallelism schemes, including data parallelism, tensor/model parallelism, pipeline parallelism, sharding, and communication/computation overlap
- Experience with performance engineering in other demanding parallel computing environments such as HPC, quantitative finance, scientific computing, graphics, compilers, or numerical simulation
- Strong systems intuition around memory hierarchy, bandwidth constraints, kernel fusion, launch overhead, communication overhead, and hardware utilization
- Experience using profiling and debugging tools to drive performance improvements
- Familiarity with infrastructure underlying large-scale training and inference, including collective communication libraries, and runtime performance analysis
- Background in a highly technical field such as physics, mathematics, theoretical computer science, computer science, or electrical engineering
- Any HPC experience is a strong plus
Skills
Required:GPU kernel programmingML workload optimizationNon-NVIDIA accelerator hardwareDistributed training systemsPerformance profiling and debuggingMemory hierarchy and bandwidth optimization
Benefits
- Comprehensive medical, dental, vision, and FSA plans
- Competitive compensation and 401(k) plan
- Relocation and immigration support on a case-by-case basis
- In-office snacks and meals provided
- Unlimited PTO and company holidays
- In-person team in San Francisco with a collaborative, high-energy environment
Last updated Aug 6, 2026. Interested? Apply on the company site.
More open roles at Zyphra
- Machine Learning Systems Administrator / DevOps Engineer: Training Clusters & Inference Cloud5d agoSan Francisco · R&D - Engineering
- Research Engineer -- Vision and Vision-Language Models5d agoSan Francisco · R&D - Engineering
- Principal Solutions Architect7mo agoSan Francisco · GTM
- Developer Relations Lead7mo agoSan Francisco · GTM
- GTM Engineer7mo agoSan Francisco · GTM