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.

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.