Software Engineer - GPU Kernels

Stack

Responsibilities

• ▹Design and implement high-performance GPU kernels for key ML operations: matrix multiplications, attention mechanisms, mixture-of-experts routing
• ▹Write and optimize code using CUDA, PTX assembly, and architecture-specific techniques
• ▹Apply advanced performance optimization methods: memory coalescing, warp-level programming, tensor core acceleration, compute/memory overlap
• ▹Implement cutting-edge features like quantization (FP8/FP4), sparsity, and compute/communication overlap
• ▹Identify and resolve performance bottlenecks using tools like Nsight Systems, Nsight Compute, and Torch Profiler
• ▹Collaborate with research teams to productionize theoretical advancements
• ▹Contribute to internal and open-source GPU libraries, present at industry conferences

Requirements

• ▹Strong understanding of GPU architecture and programming paradigms: memory hierarchy, thread/block/grid organization, synchronization techniques
• ▹Proficient in C++ and GPU performance profiling tools
• ▹Knowledge of the CUDA C++ API
• ▹Knowledge of memory access patterns and bandwidth optimization
• ▹Knowledge of numerical precision and quantization strategies
• ▹Knowledge of modern GPU features (tensor cores, async operations)

Nice to have

• ▹Experience with Transformer models and attention optimization (e.g. Flash Attention)
• ▹Familiarity with GPU kernel libraries: Cutlass, Triton, Thrust, CUB
• ▹Background in GEMM tuning and distributed/multi-GPU compute
• ▹Contributions to open-source GPU projects
• ▹Research publications or conference presentations on GPU performance

Soft skills

What we offer

• ▹Competitive compensation, including meaningful equity
• ▹100% coverage of medical, dental, and vision insurance for employee and dependents
• ▹Flexible PTO policy including company-wide Winter Break
• ▹Paid parental leave
• ▹Fertility and family-building stipend through Carrot
• ▹Company-facilitated 401(k)

About the company

Baseten powers mission-critical inference for the world's most dynamic AI companies, like Cursor, Notion, OpenEvidence, Abridge, Clay, Gamma, and Writer. The company recently raised a $1.5B Series F.