TPU Kernel Engineer

About Anthropic

• Anthropic’s mission is to create reliable, interpretable, and steerable AI systems. We want AI to be safe and beneficial for our users and for society as a whole. Our team is a quickly growing group of committed researchers, engineers, policy experts, and business leaders working together to build beneficial AI systems.

About the Role

• As a TPU Kernel Engineer, you'll be responsible for identifying and addressing performance issues across many different ML systems, including research, training, and inference. A significant portion of this work will involve designing and optimizing kernels for the TPU. You will also provide feedback to researchers about how model changes impact performance. Strong candidates will have a track record of solving large-scale systems problems and low-level optimization.

You may be a good fit if you

• Have significant experience optimizing ML systems for TPUs, GPUs, or other accelerators
• Are results-oriented, with a bias towards flexibility and impact
• Pick up slack, even if it goes outside your job description
• Enjoy pair programming (we love to pair!)
• Want to learn more about machine learning research
• Care about the societal impacts of your work

Strong candidates may also have experience with:

• High performance, large-scale ML systems
• Designing and implementing kernels for TPUs or other ML accelerators
• Understanding accelerators at a deep level, e.g. a background in computer architecture
• ML framework internals
• Language modeling with transformers

Representative projects

• Implement low-latency, high-throughput sampling for large language models
• Adapt existing models for low-precision inference
• Build quantitative models of system performance
• Design and implement custom collective communication algorithms
• Debug kernel performance at the assembly level

• The annual compensation range for this role is listed below.
• For sales roles, the range provided is the role’s On Target Earnings ("OTE") range, meaning that the range includes both the sales commissions/sales bonuses target and annual base salary for the role.

How we're different

• We believe that the highest-impact AI research will be big science. At Anthropic we work as a single cohesive team on just a few large-scale research efforts. And we value impact — advancing our long-term goals of steerable, trustworthy AI — rather than work on smaller and more specific puzzles. We view AI research as an empirical science, which has as much in common with physics and biology as with traditional efforts in computer science. We're an extremely collaborative group, and we host frequent research discussions to ensure that we are pursuing the highest-impact work at any given time. As such, we greatly value communication skills.
• The easiest way to understand our research directions is to read our recent research. This research continues many of the directions our team worked on prior to Anthropic, including: GPT-3, Circuit-Based Interpretability, Multimodal Neurons, Scaling Laws, AI & Compute, Concrete Problems in AI Safety, and Learning from Human Preferences.