Software Engineer - GPU Kernels

Responsibilities

Core Engineering Responsibilities

• Design and implement high-performance GPU kernels for key ML operations, including matrix multiplications, attention mechanisms, and mixture-of-experts routing
• Write and optimize code using CUDA, PTX assembly, and architecture-specific techniques
• Apply advanced performance optimization methods such as memory coalescing, warp-level programming, tensor core acceleration, and compute/memory overlap

Performance & Innovation

• 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

Impact & Collaboration

• Contribute to internal and open-source GPU libraries
• Present technical contributions at industry conferences (e.g., NVIDIA GTC, AWS re:Invent)

Requirements

• Strong understanding of GPU architecture and programming paradigms:
  • Memory hierarchy (global, shared, registers, L1/L2 cache)
  • Thread/block/grid organization
  • Synchronization techniques and race condition mitigation
• Proficient in C++ and GPU performance profiling tools
• Knowledge of:
  • CUDA C++ API
  • Memory access patterns and bandwidth optimization
  • Numerical precision and quantization strategies
  • Modern GPU features (e.g., 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