← All roles

Research Engineer, Infrastructure, Kernels

Thinking Machines Lab· Research Infrastructure (ML Infra)
Apply on company site
San Francisco, CA$350k - $475kFull-timePosted 11mo ago

About the role

Infrastructure research engineer to design, optimize, and maintain compute foundations powering large-scale language model training, focusing on high-performance ML kernels and scalable distributed compute.

Responsibilities

  • Design and implement custom ML kernels (e.g., CUDA, CuTe, Triton) for core LLM operations such as attention, matrix multiplication, gating, and normalization, optimized for modern GPU and accelerator architectures.
  • Design and think through compute primitives to reduce memory bandwidth bottlenecks and improve kernel compute efficiency.
  • Collaborate with research teams to align kernel-level optimizations with model architecture and algorithmic goals.
  • Develop and maintain a library of reusable kernels and performance benchmarks that serve as the foundation for internal model training.
  • Contribute to infrastructure stability and scalability, ensuring reproducibility, consistency across precision formats, and high utilization of compute resources.
  • Document and share insights through internal talks, technical papers, or open-source contributions to strengthen the broader ML systems community.

Qualifications

  • Bachelor’s degree or equivalent experience in computer science, electrical engineering, statistics, machine learning, physics, robotics, or similar.
  • Strong engineering skills, ability to contribute performant, maintainable code and debug in complex codebases
  • Understanding of deep learning frameworks (e.g., PyTorch, JAX) and their underlying system architectures.
  • Thrive in a highly collaborative environment involving many cross-functional partners and subject matter experts.
  • Bias for action with initiative to work across stacks and teams to ship features.
  • Proficiency in CUDA, CuTe, Triton, or other GPU programming frameworks.
  • Demonstrated ability to analyze, profile, and optimize compute-intensive workloads.

Skills

Required:CUDACuTeTritonKernel programmingDeep learning frameworksPerformance profiling

Benefits

  • generous health, dental, and vision benefits
  • unlimited PTO
  • paid parental leave
  • relocation support as needed
Last updated Jul 8, 2026. Interested? Apply on the company site.
Thinking Machines Lab logo

About Thinking Machines Lab

Connectionism: Research Blog by Thinking Machines Lab

Thinking Machines Lab is an artificial intelligence research and product company. We’re building a future where everyone has access to the knowledge and tools to make AI work for their unique needs and goals.

AIartificial intelligencebessemer-ai-infrastructure-2026deep learningliquid-ai-recruiting-targetmachine learningproductreinforcement-learning-platforms
Company profileWebsite

More open roles at Thinking Machines Lab