Software Engineer, GPU Kernels

River AI is building personal AI systems through custom hardware, training infrastructure, user interfaces, and deep learning research. The Software Engineer, GPU Kernels will develop and optimize GPU compute primitives for model training and inference, including attention, matrix multiplication, mixture-of-experts execution, low-precision computation, benchmarking, and runtime integration.

Responsibilities:

• Build fast GPU kernels for attention, matrix multiplication, expert routing, and related operations
• Optimize memory access, tiling, and synchronization to make efficient use of GPU hardware
• Develop FP8, FP4, and mixed-precision kernels while preserving numerical correctness
• Accelerate fine-tuning and RL through optimized adapters, backward passes, and fused operations
• Profile real workloads and integrate improvements into training and inference runtimes
• Build reproducible benchmarks that verify correctness, gradients, and performance