Become a pivotal member of Cohere’s technical team, specializing in AI performance engineering. Your expertise will help in advancing our language models through optimizations and robust system designs.

Cohere is looking for a Technical Staff member focusing on improving the performance of AI models. This role involves working within our Pre-Training team to enhance model training metrics while utilizing your software engineering expertise. You will contribute to the cutting-edge landscape of natural language processing from anywhere in the world, preferably aligning with ET to CET time zones for collaboration.

Key Responsibilities
• Create scalable and robust software for model training
• Analyze design choices and their impacts on throughput




• Write performance-enhancing CUDA and triton kernels
• Explore solutions on advanced data infrastructures
• Engage with leading researchers to drive innovation

Requirements:
• Extensive software engineering skills
• Strong knowledge of Python and ML frameworks, especially JAX
• Proficient in writing GPU kernels with CUDA and triton
• Experience with distributed training methods
• Familiarity with autoregressive models, particularly Transformers

Contribute your technical acumen to elevate AI capabilities at Cohere.

📌 Technical Staff in AI Performance Engineering (Canada)
🏢 Cohere
📍 Canada