Machine Learning Compiler Engineer
Jane Street1 day ago
London, United KingdomSenior
Responsibilities
- Develop and improve compiler technologies that convert ML research ideas and architectures into high-throughput training jobs and low-latency inference.
- Shape the development of Jane Street’s ML compiler stack across frameworks, graph compilers, kernel DSLs, and compiler infrastructure.
- Work on optimization techniques including algebraic rewriting, fusion, equality saturation, scheduling, loop-level optimization, and autotuning.
- Develop or improve compiler frontends involving graph capture, symbolic and dynamic shapes, control flow, and automatic differentiation.
- Program and optimize GPU kernels and investigate ML inference and training performance.
- Debug and optimize GPU workloads using profiling and debugging tools.
Requirements
- Experience with an ML compiler is ideal, while non-ML compiler experience is also valuable.
- Understanding of ML compiler stack components such as PyTorch, Jax, TensorFlow, Inductor, XLA, TVM, Triton, Helion, or cuTile and how they fit together.
- Knowledge of compiler optimization, compiler frontends, intermediate representations, static analysis, instruction selection, or related concepts is valuable.
- Experience or interest in GPU kernel programming, CUDA, tensor cores, ML inference, ML training, GPU debugging, formal methods, or programming-language and DSL design is valuable.
- Fluency in English is required.
Tech Stack
PyTorchTensorFlow
Categories
About Jane Street
Jane Street is a global quantitative trading firm and liquidity provider that builds in-house software and research platforms to trade across asset classes. It makes markets and executes proprietary strategies on exchanges and electronic venues, serving institutional markets rather than individual investors. Founded in 2000 and headquartered in New York, it is privately held with offices in London, Hong Kong, Singapore, and Amsterdam.