Software Engineer (ML, Performance Engineering), London

Software Engineer, ML - Performance Engineering

We are looking for engineers with different levels of experience - Mid through to Senior, Staff or equivalent levels. Isomorphic Labs is a new Alphabet company that is reimagining drug discovery through a computational- and AI-first approach. 
We are on a mission to accelerate the speed, increase the efficacy and lower the cost of drug discovery. You'll be working at the cutting edge of the new era of 'digital biology' to deliver a transformative social impact for the benefit of millions of people.

Come and be part of a multi-disciplinary team driving groundbreaking innovation and play a meaningful role in contributing towards us achieving our ambitious goals, while being a part of an inspiring, collaborative and entrepreneurial culture.

Your impact 

Working in a highly creative, iterative environment, you will be partnering with leading engineers, scientists and ML researchers to build the critical platform driving that transformation. This is a newly created role and you will need to use your previous experience and show initiative in order to fully carve out your contribution. This is an exciting opportunity for you to shape the new performance and efficiency focused capability of Isomorphic Labs. 

What you will do

• Develop highly efficient GPU kernels maximising utilisation and application performance.
• Implement aggressive optimisations to overcome hardware limitations.
• Design low-precision algorithms balancing high performance with minimal accuracy loss.
• Work with ML engineers to create training and inference-friendly model architectures.
• Use your intuitions of both runtime and hardware on-device profiling to remove model inference as a bottleneck in quick drug design iteration cycles.
• Partner with scaling teams to deploy kernels and manage training uptime.
• Profile jax level code, and implement improvements.
• Shorten the drug design iteration cycles by adapting models to inference heavy workloads.

Skills and qualificationsEssential

• Deep understanding of GPU and other AI accelerator architectures.
• Excellent collaboration skills with infrastructure and ML engineers.
• Experience with modern ML frameworks in production settings
• 3+ years of relevant industry experience.
• Strong understanding of data structures and algorithms.
• Good understanding of maths.
• Ability to code up a quick prototype solution in Python. 

Nice to have

• Experience with XLA, Triton, CUDA, Pallas or similar accelerator programming languages/deep learning compiler stacks
• Interest in chemistry and biology.
• Proven ability to optimise ML accuracy with low-precision formats.
• Experience building, deploying and maintaining production systems on GCP.