This presentation introduces Luminal, a novel ML compiler designed to radically simplify the deep learning software stack through a "search-first" methodology. The speaker, Joe Fior, explains how Luminal eschews the complexity of frameworks like PyTorch by building all computations from a minimal set of just twelve primitive operations. Models are represented as static, high-level graphs which, while slow to execute naively, serve as the input for a powerful compiler. This compiler treats kernel generation as a search problem, using equality saturation (via egg/egraphs) and a small set of generic rewrite rules to explore a vast space of logically equivalent, low-level kernels. By profiling candidates from this search space, Luminal automatically discovers highly optimized implementations. A key validation of this approach is the compiler's ability to independently discover the Flash Attention algorithm from a naive attention implementation. The presentation also covers Luminal's architecture, its support for training via a tiny autograd extension, and its future roadmap towards multi-platform support and a cloud-based deployment service.
A Search-First Approach to ML Compilation with Luminal
00:00:51 The Problem: Radical Simplification of ML Frameworks: Joe Fior introduces his project, Luminal, an ML compiler designed to combat the immense complexity of modern ML frameworks. He notes that frameworks like PyTorch have an explosion of complexity due to the multiplicative effect of supporting numerous operations, data types, and hardware devices.
00:02:23 The Solution: A Minimal Set of Primitives: Luminal is built upon a core of only 12 primitive operations (e.g., add, multiply, log2, sum reduce). All complex operations, including matrix multiplication, convolution, and even subtraction, are composed from this tiny set.
00:06:31 Static Graphs as a Foundation: Unlike the dynamic nature of PyTorch, Luminal represents models as static, directed acyclic graphs (DAGs). This ahead-of-time knowledge of the entire computation graph is what enables powerful compilation.
00:13:31 Treating Compilation as a Search Problem: Luminal avoids the pitfalls of complex, rule-based compilers by framing kernel optimization as a search problem, akin to how AlphaGo solved the game of Go. Instead of being programmed with complex heuristics, the compiler searches through a massive space of logically equivalent programs to find the fastest one.
00:14:43 The Search Workflow with egg and egraphs:
The process begins by converting the initial high-level graph into an egg expression.
A small set of ~30 generic rewrite rules (e.g., for associativity) are applied using a technique called equality saturation to build an "e-graph," which efficiently represents millions of potential kernel variations.
Finally, the compiler generates and profiles kernels from this search space to empirically discover the most performant implementation.
00:25:12 Supporting Dynamic Shapes with Symbolic Dimensions: To avoid recompiling for every new batch size or sequence length, dimensions can be represented symbolically. This allows Luminal to generate a single, flexible kernel that can be executed with variable input sizes specified at runtime.
00:42:25 Key Validation: Automatically Discovering Flash Attention: The power of Luminal's search-based approach is demonstrated by its ability to automatically derive the highly optimized Flash Attention algorithm. By starting with a simple, naive attention implementation, the compiler's generic rewrite rules and profiling were sufficient to discover the complex, fused kernel.
00:51:39 An Extensible Design for Training: Training is supported via a tiny, external autograd engine (only ~150 lines of code). It generates a backward pass graph from the forward pass, fuses them, and then feeds the combined graph into the same compiler, reusing the entire optimization pipeline.
00:53:37 Future Vision and Business Model: Luminal is open-source for local use and supports CUDA, Metal, and CPU. The future roadmap includes adding more backends (AMD, Groq, TPUs) and support for distributed computing. The business model is a Vercel-style cloud platform that will take user models, run the compilation process, and provide a fully optimized, serverless inference endpoint.