Job Description

xAI is seeking a highly motivated engineer to develop and optimize low-level CUDA kernels for state-of-the-art inference and training software. The ideal candidate will have a strong understanding of GPU architecture and experience with CUDA, CUTLASS, and deep learning techniques.

Responsibilities

• Develop and improve CUDA kernel optimizations
• Profile, debug, and optimize GPU operations
• Understand GPU memory hierarchy and computation capabilities
• Implement deep learning methods in CUDA kernels
• Innovate new ideas to maximize GPU performance

Requirements

• Experience building high-performance GeMM CUDA kernels
• Comfort with forward and backward kernel development
• Optimization skills for memory-bound and compute-bound operations
• Familiarity with Nsight and GPU utilization analysis
• Knowledge of latest inference and training optimization techniques
• Experience with pybind for kernel integration (JAX/XLA)

Benefits

• No benefits