# Research Engineer, Infrastructure, Kernels

**Company:** [Thinking Machines Lab](https://hotfix.jobs/companies/thinking-machines-lab)
**Location:** San Francisco, CA
**Role:** DevOps / SRE
**Salary:** $350k – $475k/yr
**Skills:** CUDA, Cute, Triton, PyTorch, JAX, Gpu Programming, Low-Precision Arithmetic, Tensor Parallelism, Pipeline Parallelism, Xla, Tvm
**Posted:** 2026-05-04

> Designs and optimizes high-performance ML kernels (CUDA, CuTe, Triton) for large-scale LLM training, focusing on GPU efficiency, low-precision formats, and distributed compute. Collaborates with researchers to bridge algorithms and hardware.

## Job Description

## What You’ll Do

- 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.

## Skills and Qualifications

**Minimum 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, different cross-functional partners and subject matter experts.
- A bias for action with a mindset to take initiative to work across different stacks and different teams where you spot the opportunity to make sure something ships.
- Proficiency in CUDA, CuTe, Triton, or other GPU programming frameworks.
- Demonstrated ability to analyze, profile, and optimize compute-intensive workloads.

**Preferred qualifications:**

- Experience training or supporting large-scale language models with tens of billions of parameters or more.
- Track record of improving research productivity through infrastructure design or process improvements.
- Experience developing or tuning kernels for deep learning frameworks such as PyTorch, JAX, or custom accelerators.
- Familiarity with tensor parallelism, pipeline parallelism, or distributed data processing frameworks.
- Experience implementing low-precision formats (FP8, INT8, block floating point) or contributing to related compiler stacks (e.g., XLA, TVM).
- Contributions to open-source GPU, ML systems, or compiler optimization projects.
- Prior research or engineering experience in numerical optimization, communication-efficient training, or scalable AI infrastructure.

## Logistics

**Compensation:** Depending on background, skills and experience, the expected annual salary range for this position is $350,000 - $475,000 USD.

**Benefits:** Generous health, dental, and vision benefits, unlimited PTO, paid parental leave, and relocation support as needed.

## Similar roles

- [Reliability Engineer, Supercomputing](https://hotfix.jobs/jobs/e9263e4b-66f2-4203-b8da-b88ba1855dfe) - Thinking Machines Lab - San Francisco, CA - $350k – $475k/yr
- [Network Engineer, Supercomputing](https://hotfix.jobs/jobs/2a952a8c-f34c-4fdd-a219-c8ac22399145) - Thinking Machines Lab - San Francisco, CA - $350k – $475k/yr
- [Performance Engineer, Inference Systems](https://hotfix.jobs/jobs/c4e8b3e8-25ad-4a83-b11d-bebf411d63f1) - Anthropic - San Francisco, CA - $350k – $850k/yr
- [Site Reliability Engineer (SRE)](https://hotfix.jobs/jobs/893a897e-0421-40b8-a8ac-f61c0b30d5ff) - Thinking Machines Lab - San Francisco, CA - $350k – $475k/yr
- [Research Engineer, Infrastructure, Training Systems](https://hotfix.jobs/jobs/9cc8be27-572e-4c19-97d0-8a7ae54e41a8) - Thinking Machines Lab - San Francisco, CA - $350k – $475k/yr

**Apply:** https://hotfix.jobs/jobs/0a8c8d26-d360-4832-90a8-bfc710e06ff4
**Canonical:** https://hotfix.jobs/jobs/0a8c8d26-d360-4832-90a8-bfc710e06ff4