# Member of Technical Staff, Training Performance Engineer

**Company:** [Cohere](https://hotfix.jobs/companies/cohere)
**Location:** London, United Kingdom, New York, NY, Paris, France, Toronto, Canada, Montreal, Canada
**Role:** ML Engineering
**Skills:** Python, JAX, PyTorch, xla, mlir, CUDA, triton, gpu kernels, Distributed Training, Transformers, supercomputing, profiling tools
**Posted:** 2025-02-20

> Optimizes training performance for advanced language models by developing scalable software, GPU kernels, distributed training systems, and profiling tools. The role requires strong software engineering skills, Python and ML framework proficiency, and experience with CUDA or Triton.

## Job Description

## Responsibilities
- Design and write high-performance, scalable software for model training.
- Understand architectural modifications and design choices and their effects on training throughput and quality.
- Write low-level CUDA and Triton kernels to maximize accelerator performance.
- Research, implement, and experiment with ideas on supercomputing and data infrastructure.
- Identify and remove performance bottlenecks.
- Develop training and profiling tools to improve model performance.
- Collaborate with researchers and engineers on large-scale language model systems.

## Requirements
- Extremely strong software engineering skills.
- Proficiency in Python and related machine learning frameworks, including JAX, PyTorch, and XLA/MLIR.
- Experience writing GPU kernels using CUDA, Triton, or similar technologies.
- Experience with large-scale distributed training strategies.
- Familiarity with autoregressive sequence models such as Transformers.

## Nice to Have
- A paper published at a top-tier venue such as NeurIPS, ICML, ICLR, AIStats, MLSys, JMLR, AAAI, Nature, COLING, ACL, or EMNLP.

## Compensation and Benefits
- Weekly lunch stipend of $75/£75 or equivalent in local currency.
- Full health and dental benefits, including a separate mental health budget.
- RRSP matching, 401(k), and pension scheme.
- 100% parental leave top-up for up to six months for either parent.
- Annual enrichment benefits for arts and culture, fitness and wellness, quality time, and workspace improvements.
- Education and learning stipend for conferences, courses, and coaching.
- Six weeks of paid vacation (30 working days).
- Travel budget for remote employees visiting other offices and an annual company offsite.
- Coworking benefit for employees not near an office.
- $500 home-office stipend.

## Similar roles

- [Staff Machine Learning Engineer, Traffic Intelligence](https://hotfix.jobs/jobs/cee28c3a-eb64-42f7-a584-8d2b4a0980b8) - Airbnb - Remote - $212k – $265k/yr
- [Member of Technical Staff, Agentic Environments](https://hotfix.jobs/jobs/f3ceb248-2d3e-4b28-a5d1-10891969916e) - Cohere - Remote
- [Staff Machine Learning Infrastructure Engineer, Embedding Platform](https://hotfix.jobs/jobs/9b4387bf-ac16-48ba-810b-2e6a612e3617) - Reddit - Remote - $253k – $355k/yr
- [Senior Staff Machine Learning Engineer, Feed Relevance](https://hotfix.jobs/jobs/a2f9b178-c6cf-49bf-ba93-00688d1027fb) - Reddit - Remote - $266k – $372k/yr
- [Staff Applied Scientist](https://hotfix.jobs/jobs/d518c7a7-07d4-4f06-af79-9c478df769ca) - Garner Health - New York, NY - $300k – $390k/yr

**Apply:** https://hotfix.jobs/jobs/386cb47a-4bf4-45d3-80da-8fcbe1f811e3
**Canonical:** https://hotfix.jobs/jobs/386cb47a-4bf4-45d3-80da-8fcbe1f811e3