Staff Compiler Engineer

Company Overview

Flux Computing designs and manufactures optical processors to train and run inference on large AI models. Join us in London to be part of a highly motivated and skilled team that thrives on delivering impact and innovation at speed.

The role

We’re searching for Staff Compiler Engineers to architect and build the ML backend (compiler, run-time, and debugger) for our next-generation OPTUs. You will own integration with PyTorch, Tensorflow, JAX, and MXNet down to our low-level kernel drivers. Your mission will be to create seamless support for a broad ecosystem of large AI models, and ensure we are pushing the limits of their performance envelope by closely partnering with hardware and modelling teams to understand device trade-offs.

Responsibilities

• Ownership: Define and deliver the technical vision and roadmap for your team that unlocks key strategic technical and business goals that are essential to the success of Flux.
• Collaboration: Partner closely with hardware engineers to align compiler, runtime, and debugger requirements with the OTPU design; ensure software and hardware are designed together to deliver maximum performance.
• Architect & Build: Design and implement our compiler, runtime, and debugger for PyTorch, TensorFlow, JAX, and MXNet on custom hardware.
• Optimize Performance: Apply advanced techniques (layout, fusion, scheduling, tiling) to eliminate bottlenecks and maximize throughput.
• Mentor & Define Standards: Lead code reviews, coach peers, and define best practices in ML backend and performance engineering for your team.
• Guide Technology Direction: Stay ahead of GPU, AI accelerator, and optical computing trends; propose and prototype innovations.

Skills & Experience

• 7+ years of experience in software engineering with a focus on C/C++ programming.
• Extensive experience in ML framework internals, compilers, low-level programming, and optimisation techniques.
• Extensive experience optimising Tensorflow, PyTorch or JAX deep learning models.
• Extensive experience with multiple toolchains like LLVM, OpenXLA/XLA, MLIR, TVM.
• Practical experience applying machine learning in high-performance computing contexts.
• Strong problem-solving skills and the ability to think critically and creatively.
• Experience in high-pace, dynamic work environments.
• Excellent teamwork and communication skills, with the ability to collaborate effectively with cross-functional teams.
• Bachelor's degree in computer science, electrical engineering, telecoms engineering, mathematics, or a related field.
• Personal projects are a key differentiating factor and hold more weight than other requirements.

Compensation & Benefits

• Highly competitive salary and stock options in a rapidly growing AI company, depending on experience.
• Based in our new 5,000 sq. ft. office in the AI hub of Kings Cross, London.
• To foster collaboration in our high-growth environment, we require all employees to work from our London HQ and live within a 45-minute commute. We offer an extra £24,000/year incentive for those living within 20 minutes.
• Comprehensive healthcare insurance.
• 25 days PTO policy plus bank holidays.
• Private access to our in-house 3D printer.

If you are passionate about pushing the boundaries of what's possible in AI and thrive in a high-energy, fast-paced environment, we want to hear from you. Apply now to join Flux and be a key player in shaping the future of computing.