GPU Kernel Developer - AI Trainer

Obsidian; Chicago (IL); On-site; USD 120,000 - 180,000; Full time

Job summary

Obsidian in Chicago seeks a specialist to evaluate GPU/accelerator kernel development tasks used to train frontier AI models. You will judge numerical correctness, performance benchmarks, task scoping, and compile/run validity across diverse kernel task types.

This role demands rubric-based written feedback and clear guidance to improve kernels, with emphasis on cross-framework translation and hardware-target migrations.

Qualifications

• 3+ years hands-on experience developing, optimizing, or verifying GPU kernels in at least two of CUDA, Triton, NKI, or Pallas (JAX).
• Strong understanding of numerical-correctness criteria for kernels (absolute/relative/ULP tolerances, reference-implementation selection).
• Experience with performance profiling and benchmarking (nsight, ncu, roofline analysis, or framework-native profilers).
• Familiarity with common compilation/runtime failure modes (driver mismatches, OOM, launch-configuration errors, shape/stride mismatches, autotuning failures).
• Experience with at least three kernel task types: generation from specification, translation/lowering across frameworks, migration between hardware targets, debugging, performance optimization, or operator fusion.

Skills

CUDA

Triton

NKI

Pallas (JAX)

Performance profiling

Tools

Nsight

NCU

Job description

Evaluate the quality, correctness, and completeness of GPU/accelerator kernel development tasks used to train and evaluate a frontier AI lab's models. You'll assess numerical correctness, performance-benchmarking fairness, task scoping, and compilation/runtime validity across diverse kernel task types — and provide clear, rubric-based written feedback.

Basic Qualifications

• 3+ years of hands-on experience developing, optimizing, or verifying GPU/accelerator kernels in at least two of: CUDA, Triton, NKI, or Pallas (JAX)
• Strong understanding of numerical-correctness criteria for kernels (absolute/relative/ULP tolerances, reference-implementation selection)
• Demonstrated experience with performance profiling and benchmarking (nsight, ncu, roofline analysis, or framework-native profilers)
• Familiarity with common compilation and runtime failure modes (driver mismatches, OOM, launch-configuration errors, shape/stride mismatches, autotuning failures)
• Experience with at least three kernel task types: generation from specification, translation/lowering across frameworks, migration between hardware targets, debugging, performance optimization, or operator fusion

Preferred Qualifications

• Experience across both NVIDIA GPU (CUDA/Triton) and custom-accelerator (NKI/Pallas/TPU) ecosystems
• Background in compiler engineering, MLIR, or intermediate-representation lowering
• Understanding of memory-hierarchy optimization (shared-memory tiling, register pressure, bank conflicts, coalescing patterns)
• Contributions to kernel libraries (cuBLAS, cuDNN, Triton community kernels, JAX/XLA custom calls)