AI Engineer World's Fair 2025
Luminal - Search-Based Deep Learning Compilers - Joe Fioti
Read the talk
Luminal: compiling a small operation vocabulary into fast GPU kernels
Deep learning needs only a small mathematical vocabulary. Luminal builds models from that vocabulary, then searches equivalent implementations to recover efficient GPU execution.
From a talk by Joe Fioti
Before you start: Familiarity with tensor shapes, matrix multiplication and the distinction between CPU and GPU execution will help you follow the compiler examples.
Why does simple linear algebra need such complicated libraries?
Deep learning works with scalars, vectors, matrices and tensors, combining addition, multiplication, matrix multiplication and element-wise functions. Why, then, are the libraries implementing it so complicated? Joe Fioti, a co-creator of Luminal, starts with that mismatch: a small mathematical foundation has produced an enormous software ecosystem. His proposed way back to simplicity is to make compilation—and specifically search—do more of the work.
Fioti describes PyTorch as having more than 1,200 operations and more than 15 data types, across targets including CPUs, CUDA, AMD, TPUs and NPUs. The implementation burden grows across combinations, not merely by adding features:
Supporting another operation potentially requires implementations for every relevant data type and device. A new device creates the corresponding problem across operations and types.
Fioti puts PyTorch at roughly three million lines of code. These figures are his scale estimates for the framework at the time, without a version or counting methodology. The practical concern is not size alone: a large implementation surface creates more opportunities for bugs and makes internals harder to understand, extend and debug.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Compose familiar layers from primitive operations
Luminal starts from the other direction: what is the smallest vocabulary needed to represent useful models? Fioti describes a twelve-operation core and names the following mathematical operations:
| Family | Operations |
|---|---|
| Unary | exp2, log2, sine, reciprocal, square root |
| Binary | Addition, multiplication, modulo, less than |
| Reduction | Sum, maximum |
The spoken list contains eleven entries; Luminal’s primitive-graph introduction includes Contiguous as the twelfth. That is the architecture described here, rather than the expanded primitive set in current project documentation. Fioti says this small vocabulary can represent language models, vision-language models, CNNs, RNNs and diffusion models.
Many apparently missing operations are compositions. Subtraction multiplies the second input by −1 and adds it; division multiplies by its reciprocal:
The distinction is between the operation a model author wants to express and the primitives the compiler must understand. A convenient frontend need not make every convenience a separate backend primitive.
Matrix multiplication becomes broadcasted multiplication followed by sum reduction. For A shaped [M, K] and B shaped [K, N], shape metadata presents them as [M, K, 1] and [1, K, N]. Broadcasting yields products indexed by [M, K, N]; reducing over K yields [M, N]:
This is a representation of the computation, not a requirement to materialize the entire broadcasted intermediate. Convolution follows the same approach: shape trackers describe pooled input windows, then a matrix multiplication combines those windows with the convolution kernel.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Keep the graph static and the changing dimensions bounded
Dynamic frameworks made experimentation convenient when researchers were exploring RNNs, LSTMs and changing model structures. Luminal instead asks how much dynamism a particular workload actually needs. In the transformer workload Fioti describes, the changing quantities are sequence length and KV-cache length; the surrounding computation stays fixed. That makes a static representation useful even when some dimensions vary.
The representation is a directed acyclic graph. In the dense-layer example, one node loads the input tensor and another loads the weights. Their outputs feed an element-wise multiply, followed by a sum reduction—the matrix multiplication just decomposed above. Larger models produce larger graphs, but the same vocabulary still specifies their computations.
Fioti reports that Luminal’s implementation at the time was under 5,000 lines of code, without specifying the counting scope. The design goal is equally concrete: a developer should be able to understand the library’s core structure in an afternoon. A small graph language makes that plausible, but it does not yet make execution fast.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
The primitive graph is the starting point, not the execution plan
Fioti illustrates unoptimized execution with Llama 7B: generating a sentence could take roughly all day, with no hardware, precision or token count specified. Executing each primitive separately is not the intended endpoint. The graph first expresses the model correctly; compilers then transform it into faster graphs. A minimal representation and an efficient implementation are separate responsibilities.
The conventional stack he sketches starts with Hugging Face Transformers, passes through PyTorch, xFormers and handwritten kernels, and may call cuDNN or cuBLAS before reaching CUDA. Those kernels must accommodate many workloads. The resulting dependencies complicate both installation and tracing a bug through the stack.
In the path demonstrated here, Luminal’s graph and compilers directly generate CUDA code. That shortens the route from model representation to device execution:
| Stack | Route to CUDA |
|---|---|
| Conventional example | Model library → framework and kernel libraries → cuDNN/cuBLAS → CUDA |
| Luminal’s demonstrated path | Primitive graph → compilers → generated CUDA |
The simplification does not eliminate optimization work. It concentrates that work in the transformation from a primitive graph to an efficient executable graph.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Simpler hardware asks more of the compiler
Why have compilers not already removed all this complexity? Fioti’s answer is that generating sophisticated kernels can require disproportionately sophisticated compiler logic. He uses square-or-cube growth, and a possible tenfold compiler expansion for a doubling of kernel complexity, as illustrations rather than measured scaling laws. The difficulty is especially acute for hardware startups: a promising device still needs software that can schedule and feed it effectively.
His hardware comparison follows the movement of responsibility from circuitry into software:
| Hardware | Responsibilities emphasized in the talk |
|---|---|
| CPU | Hardware predicts and manages more execution details |
| GPU | Software explicitly manages memory use, dispatch and communication; hardware retains scheduling responsibilities |
| TPU | Software takes greater responsibility for scheduling and memory allocation |
The intended tradeoff is simpler, more uniform hardware with better throughput and performance per watt for suitable workloads. But the work removed from hardware must be handled elsewhere, increasingly by the compiler.
This leads to the familiar VLIW, or Very Long Instruction Word, scheduling problem. Hardware can become simpler when the compiler decides ahead of time which work runs together and how resources are used. Beyond a certain level of scheduling complexity, however, writing a compiler that consistently makes good choices becomes the bottleneck.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Search equivalent implementations instead of prescribing every choice
Fioti turns to AlphaGo as an analogy. Its designers did not encode a perfect answer for every Go position; they explored board states using tree search guided by a neural network. Luminal substitutes logically equivalent GPU kernels for board states. Instead of prescribing one optimization sequence that must always produce fast code, simple rules generate alternatives and search chooses among them. The neural-network guidance belongs to the AlphaGo comparison; Fioti does not describe a corresponding learned guide in Luminal.
The implementation uses egglog to represent and expand those alternatives:
- Convert the operation graph into egglog expressions.
- Represent equivalent intermediate expressions compactly in e-graphs.
- Repeatedly apply small rewrite rules that preserve the logical result.
- Profile candidate kernels and select the fastest measured candidate.
- Narrow exploration when the space becomes too large to profile exhaustively.
Fioti describes approximately 20–25 rewrite rules in this search setup. A rewrite need not make a kernel faster by itself: it may produce a slower candidate while opening a path to a better combination of transformations.
Runtime measurement supplies the performance judgment that the rewrite rules deliberately avoid making. Once profiling every possibility becomes infeasible, Fioti says Luminal uses techniques such as Monte Carlo tree search to narrow exploration. The division of labor is crucial: rules establish which alternatives preserve the computation, while search and profiling determine which implementations are worth pursuing.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Fusion removes an intermediate trip through global memory
Consider a tensor passed through sine and then exp2. With separate kernels, the GPU loads the input, computes sine, and writes the intermediate tensor to global memory. The next kernel reads that same intermediate, computes exp2, and writes the result. The arithmetic is simple; the extra write and read are avoidable.
Fioti characterizes data movement as roughly 99% of GPU time and energy; that is a broad claim, not a measured share for this example. The useful distinction is whether a workload is limited by memory traffic, arithmetic or latency, as described in NVIDIA’s GPU performance guide. For this dependent element-wise chain, kernel fusion removes the global-memory intermediate: load once, compute both functions, then store the final output.
The resulting CUDA computation can be expressed directly:
cuda
#include <cuda_runtime.h>
#include <math.h>
__global__ void sine_then_exp2(const float* input, float* output, int n) {
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
const float intermediate = sinf(input[i]);
output[i] = exp2f(intermediate);
}
}
intermediate is local to the computation; there is no separate global-memory tensor between sine and exp2. The mathematical dependency remains intact while the memory path changes.
The larger graph Fioti next shows in the recording applies the same principle repeatedly. Many separate operations, each with an intermediate memory round trip, collapse into one kernel. Fioti reports that this fused kernel takes little longer than an individual kernel from the original graph. His explanation is that the additional arithmetic costs relatively little compared with repeatedly moving intermediate results through memory; the talk supplies no timing configuration for that comparison.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
From naive attention to a FlashAttention implementation
The more ambitious example is FlashAttention, the IO-aware exact-attention algorithm introduced by Tri Dao and colleagues in 2022. Fioti contrasts that date with the introduction of transformers in 2017: an important optimization took years to emerge after the underlying architecture. He reports that Luminal’s compiler can find the optimization through search rather than requiring the full algorithm to be manually encoded as a compiler pass.
The input is a naive multi-head attention graph. Repeated rewrites construct a large space of equivalent implementations; kernel profiling then selects an implementation that Fioti identifies as FlashAttention. He qualifies the claim of compiler uniqueness as being to the team’s knowledge. This is a reported prototype result, without a supplied search budget, numerical-tolerance specification or reproducible runtime comparison.
Later in the recording, Fioti shows an announcement slide with generated kernel code in green and the intermediate representation in white. It makes the output of the search concrete: an optimized representation becomes CUDA code, rather than stopping at an abstract graph transformation.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Reuse buffers and submit kernels together
Search does not have to make every optimization decision. After selecting kernels, Luminal applies deterministic passes before final execution code. Fioti characterizes these as transformations that will not hurt, although their benefit varies. Buffer reuse is the first example: the complete graph exposes when each intermediate is needed and when its storage becomes available again.
In the demonstration, buffer 1 and buffer 3 have non-overlapping lifetimes. By the time buffer 3 is needed, the last use of buffer 1 has finished. They can therefore share an allocation instead of reserving two independent buffers. This is a change to the memory plan, not to the logical identity of the two intermediate results. The comparison below uses allocation-a and allocation-b as illustrative allocation names.
Two logical buffers share storage when their lifetimes do not overlap
Constructed example: The allocation names allocation-a and allocation-b and the symbolic lifetime notation are teaching details. The comparison depicts a memory plan, not an executed allocation trace.
last_use(buffer 1) < first_use(buffer 3)
Operation: Revise the allocation plan so buffer 3 reuses buffer 1's storage after its final use.
buffer 1
allocation-a; live during the earlier interval
allocation-a; live during the earlier interval
buffer 3
allocation-b; live during the later interval
allocation-a; live during the later interval
Separate storage for buffer 3
allocation-b reserved
Not present
The next pass addresses launch overhead. Fioti’s illustrative baseline alternates CPU dispatch, GPU execution and a return to the CPU before the next dispatch. Luminal instead builds a queue ahead of execution so the GPU can proceed through successive kernels with less launch waiting.
That baseline should not be read as a requirement of ordinary CUDA: kernel launches are asynchronous with respect to the host, and operations in a stream can execute in order without a host wait after each kernel. The optimization target here is avoidable host submission and synchronization overhead; the talk does not identify the queue mechanism as CUDA Graphs.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Attach a backward graph without rebuilding the compiler
Luminal began as an inference library, but its graph representation also provides an extension point for training. An external autograd crate takes the forward graph, derives a backward graph and attaches it. The combined graph then passes through the existing downstream compilers.
The compiler optimizes computations, whether they belong to inference or differentiation. Kernel search and the other graph transformations therefore apply to the backward pass as well. Fioti’s characterization of getting training for free refers to this reuse of the compilation machinery, not to eliminating the work of building autograd or running training.
Fioti describes external training support as unusual, qualifying his comparison with other libraries as being to his knowledge. The architectural opportunity is more important than the priority claim: contributors can supply alternative autograd engines, gradient sharding or specialized training arrangements without making each one part of Luminal’s core.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Extend the graph across hardware, environments and deployment
At the time of the talk, Fioti lists CPU, CUDA and Metal support. AMD, Tenstorrent, Groq and TPUs are proposed targets. The aim is to make more hardware usable through the same model representation and reduce dependence on CUDA.
Two further directions extend how much work the graph describes:
- Distributed inference and training: combine data, pipeline and tensor parallelism into full 3D distributed execution.
- GPU-resident reinforcement-learning environments: remove the back-and-forth between a model running on the GPU and an environment stepping on the CPU.
For the latter, the environment itself becomes a Luminal graph and passes through the same compiler flow as the model. Fioti says the team had done this for very simple environments; handling more complex ones remained an open goal. If both pieces fit that representation, the model’s forward pass and the environment step can run together on the GPU.
After a joke about a Dyson Sphere awaiting a Sequoia fundraise, Fioti turns to Luminal Cloud. The proposed deployment interface follows directly from treating a model as a portable graph:
- Export the model graph using the talk’s
graph.exportinterface. - Upload the resulting file.
- Receive a serverless inference endpoint.
This is the workflow presented in the recording, rather than a current SDK call signature.
The service would handle optimization, batching, queuing and machine provisioning. Fioti describes billing only while the graph executes. The cloud proposition therefore extends the same separation used locally: the user supplies the computation, and the system takes responsibility for turning it into an efficient execution plan. His simplest-and-fastest positioning is a product ambition, not a demonstrated cloud comparison.
The closing invitation is practical: contribute ideas and pull requests to the Luminal repository, or bring an inference workload that can test the approach. Fioti connects the small core to the ability to develop new capabilities without the overhead of much larger frameworks. He invites startups and companies to discuss workloads through the Luminal site.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Resources
From the talk
Source code and current instructions for the Luminal compiler.
An equality saturation and Datalog engine with installation instructions and tutorials.
The original paper on exact attention using tiling to reduce GPU memory traffic.
Further reading
An introduction to twelve primitive operations, matrix multiplication decomposition and compiler passes.
Joe Fioti introduces a prototype that searches equivalent kernels and generates CUDA attention code.
- Understanding GPU performance limitsDocumentation
Explains how arithmetic intensity, memory bandwidth and available parallelism limit GPU workloads.
Read the complete timestamped transcript
- 0:02
Hey everyone, I'm Joe. I'm one of the creators of Luminal, and so today in this talk I'm gonna explain what it is, how it works, and why we really think that it's the future of ML libraries.
- 0:14
So the title of the talk is Radical Simplificat- Simplification Through Search, and that really is the theme of Luminal. We are far simpler than most other ML libraries, and yet we are not taking a hit on performance or capability because we are using, uh, compilers and more specifically, search.
- 0:35
So, deep learning fundamentally is very simple. It's simple linear algebra. And so basically all it really is, is, is, you know, scalars, vectors, matrices, and tensors, and then a couple core ops, a couple core operations.
- 0:52
So we have addition, some multiplies, matmuls, some element-wise ops in between, and that's, that's mostly it. But if deep learning is so simple, why aren't the libraries simple? So the machine learning software ecosystem is very, very complicated.
- 1:10
Um, PyTorch is one of the most famous libraries out there. It has over twelve hundred operations, uh, and has over fifteen different data types. It runs on a ton of different kinds of devices, uh, CPU, CUDA, AMD, TPUs, uh, some NPUs out there.
- 1:28
And the problem with all of this is that complexity does not scale just by adding these together. You don't have to have the number of operations plus the number of data types plus the number of, uh, supported devices.
- 1:42
It's actually multiplicative. So the number-- the complexity scales with ops times data types times number of devices. And so once you raise any of those, you wanna support a new op or a new data type or a new device, your complex- complexity starts exploding.
- 2:02
So what that... what's happened to PyTorch is it's now over like three million lines of code. Uh, TensorFlow is even much worse than this. So it's extremely, extremely complicated pieces of software, and this hurts because there's a lot more bugs obviously if you have a ton more code, but it's also just very hard for people to ever
- 2:21
extend or use or build anything inside of.
- 2:26
So what we do is we take the approach of looking top-down at machine learning and then saying fundamentally, "What are the minimum amount of things we need in order to make ML models run?"
- 2:38
So deep learning, again, like I said, is linear algebra. Linear algebra boils down to simple ops. So what if we just built these very complicated models out of like Lego blocks of very simple operations?
- 2:51
So what we do is we have, uh, twelve operations that are very, very, very simple. Um, and so we have exp2, log2, sine, and reciprocal, and square root. Those are our unary operations.
- 3:04
We have, uh, our binary operations: addition, multiplication, modulo, and less than. And then we have our reductions. So we have sum reduce and max reduce. And with those, just those operations, you can support all of the big models out there, all of the commercially relevant models that everybody cares about.
- 3:21
So you can do language models, vision language models, uh, CNNs, RNNs. You can do, um, uh, diffusion models, all these other very, very popular models.
- 3:34
Is that really it? I mean, that's not a lot of ops. [clears throat] Is that really all it takes to, uh, to represent all of this? Well, it's a lot of these other operations that you would have expected to see on the list are really just,
- 3:47
uh, uh, formable, usable by combining different operations on this list. So subtraction is really simple. It's just addition and multiplication with a negative one. Division is really simple because [clears throat] it's just multiplication and a reciprocal on B.
- 4:04
Uh, matmuls are really simple if you have the ability to manipulate the shape metadata of your tensor. So you can basically just do a broadcast at multiply and then sum reduce it down, and you get the output of matmul, so you don't need a matmul op.
- 4:19
Um, convolution is really simple if you have the ability to do pooling through, uh, shape trackers again, and then you do a matmul, like we just discussed, with the, the convolution kernel, and you get the output of a convolution.
- 4:33
And so another thing we realized when we were building this is that all of these existing libraries are built with dynamism at their core. And they were usually, they were mostly built, uh, you know, five to ten years ago when dynamism was very important.
- 4:48
So people were experimenting with RNNs and LSTMs and all these fancy models, and they needed a whole lot of hackability and dynamism, and they didn't so much care about performance.
- 4:59
Um, so they, you know, PyTorch is very, very dynamic, but because of that, it's a lot more complex. So deep learning fundamentally isn't dynamic. It was just there for convenience.
- 5:11
But the dynamism inherent to models here is very, very small and very, very bounded. So in a transformer model, the only real thing that is dynamic is the KV cache length and the actual sequence length.
- 5:26
And aside from that, the entire model is static. Um, and so what we do is we specify these models as directed acyclic graphs of operations. So right in front of us, we see a really simple example where we're loading in a tensor, we're loading in a weight, and then we are element-wise multiplying them together, and then we
- 5:46
are sum reducing them. And right here, like we talked about, this is a matmul. This is what a matrix multiply is. And so right on screen, this is all that a, a single dense, uh, neural network layer is.
- 6:00
It's just your, your, your matrix multiplying by a weight matrix. So on the right-hand side here, we see a much larger model. And so these graphs, you can kind of see get quite a bit more compli- complex, uh, but they are capable of fully specifying these models.
- 6:17
So the consequence of this is that Luminal is really, really simple. Despite it being able to actually represent and run all of these different models in the world, it's under five thousand lines of code.
- 6:29
Uh, it's very easy to understand, and our goal here is that we are going to, uh, this whole library really should be learnable in an afternoon. So you should be able to sit down and understand the core structure and core concepts of it in, you know, a couple hours here.
- 6:44
So that doesn't necessarily mean it's fast, though. So Llama 7B, uh, right out of the box runs, uh, you know, it takes about all day to generate a single sentence.
- 6:57
Super-duper slow. But the point isn't to run these primitive graphs of operations here. The point is to take these graphs and then run them through some functions to transform them into faster graphs.
- 7:10
And so we do that with compilers. So we compile our way back to high performance. So before we get to talk about compilers here, I just wanna again highlight, um, what this simplification gets us.
- 7:23
A traditional stack is you might be using something like Hugging Face Transformers. It's a great library. It sits on top of PyTorch and xFormers and, uh, you know, some other libraries that provide optimized kernels.
- 7:36
These libraries inside them have a bunch of handwritten kernels, uh, that try to be adaptable for all these different use cases. And then those handwritten kernels may call operations in cuDNN or cuBLAS, which then sit on top of CUDA.
- 7:53
And so all of this creates a very complex dependency story. Anybody who's tried to install, uh, a deep learning setup on a new machine knows that this is a non-trivial thing.
- 8:04
Uh, a lot of the times it's, you know, dependency hell that you're stuck in. Um, and it's just, it's just, you know, having very complex, uh, stacks means that once there's a bug, tracing that down through a really complex stack is, is a huge pain as well.
- 8:18
So generally, we wanna keep our stack as simple and straightforward as possible. Uh, with Luminal, we directly emit CUDA code. We directly generate CUDA code. And so there really is nothing between us and CUDA.
- 8:30
It's just our library, our graph, our compilers, and CUDA right beneath.
- 8:37
So like I said, Luminal is slow by default, but you're never meant to run any of these, uh, primitive op graphs. So you're supposed to... We're gonna take these graphs, and we're gonna feed them through compilers, and they're going to spit out much faster graphs.
- 8:53
So how does that actually work? I mean, we're not the first to think of compilers here, and we're not the first to think of ML compilers. Why haven't other ML compilers taken over the world?
- 9:04
Um, it's really because as your code that you want to generate grows in complexity, your compiler scales with the, the square or the cube of that complexity. It scales really, really fast because your compiler needs to now emit, needs to generate that final code.
- 9:23
And so if you double the, the complexity of your kernels, your compiler might have to 10X in complexity. So at some point, these compilers just become too complex for people to write.
- 9:35
Uh, and so that really has put a bottleneck on the current ecosystem. It's put a bottleneck on a lot of hardware startups with very fancy, uh, hardware that, you know, they, they wanna compile stuff down to.
- 9:47
Um, and, and it's actually getting worse. So as we demand more and more out of our hardware, the hardware we need to keep and make simpler and simpler. Because generally, the simpler your hardware can get, uh, the more uniform it can get and the faster it can get.
- 10:03
So as an example of this, CPUs are very, very complex. They try to predict what the programmer wanted to do. Um, they try to make life easier for the programmer.
- 10:13
But as a consequence, they're very complex pieces of hardware, and they don't run very fast for operations that we care about. GPUs are a lot simpler. GPUs require the programmer to specify things ahead of time, set things up, um, explicitly use certain memories, uh, dispatch kernels, and do certain communications explicitly, and the GPU doesn't handle that for
- 10:37
you. It handles some things for you, like scheduling, but, uh, the software needs to do a lot more. But as a consequence, the hardware is simpler, and it's a lot faster.
- 10:47
You get much better performance per watt. Uh, TPUs are even simpler, so the programmer has to handle everything. The programmer has to schedule things and has to, uh, explicitly allocate and manage all this memory.
- 11:00
Um, but the TPUs are very, very simple. So the TPUs are extremely fast, very good performance per watt, and we wanna basically keep walking down this path of simpler, faster hardware and, uh, more complex software.
- 11:16
So what that means is our, our compiler needs to get more complex. Uh, we actually run into a very typical traditional problem in, uh, CS, which is the VLIW compiler problem.
- 11:28
Uh, VLIW stands for very large instruction width, and it's, it's a very, very standard paradox where companies want to make their hardware simple, so they want to, to basically have the compiler statically schedule and set everything up beforehand.
- 11:46
But the problem is that beyond a certain point, compilers just get way too complex, and humans just can't write those compilers anymore. It's just too hard. So how do we actually get past that in Luminal?
- 11:59
Uh, well, we actually turn to the same solution that the AlphaGo, AlphaGo guys did when they were creating AlphaGo, and they wanted to crack the game of Go. They said that they were not super genius Go players, that Go was a very tough problem, and instead of trying to write, like, the perfect algorithm that would be able
- 12:18
to solve Go in one shot and always find the best move
- 12:22
Instead, what they did is they were-- they turned to search, and they said, "We're going to search through a whole bunch of, uh, boards here, of board states. And so we're gonna do this, like, really fancy tree search, and we're gonna use this guided neural network to, to go through it."
- 12:37
Uh, but the point is that search ate the complexity of Go. And so what we're doing is the exact same thing. We are searching instead through... Instead, instead of Go boards, we are searching through logically equivalent GPU kernels.
- 12:51
So this means that we don't need to hand write out a whole bunch of rules and then hope that they always work to produce fast code. What we can do is we can write a whole bunch of simple rules to build this big search space and then let the search go through and find the fastest kernels.
- 13:09
So what does this actually look like? Well, we take our graphs, uh, the graphs that we were talking about before, and we convert them into expressions in this library called Aglog.
- 13:18
Um, this library basically uses E graphs to represent these, this search space in a very memory efficient way, and then it goes ahead and it does this search through all of these equivalent expressions, this equivalent intermediate representation.
- 13:33
And we specify out, you know, twenty or twenty-five different rewrite rules. These rewrite rules are very, very simple. All they do is make a small alteration to a given GPU kernel, and we know that the output is logically equivalent.
- 13:49
We don't know if it's necessarily faster or slower, but we know it's logically equivalent, so it adds to our search space. And our search space, as we iteratively apply and apply and apply these, these simple rewrite rules as many times as we can, we build up a super, super large, uh, search space.
- 14:07
And then what we do is we just go through and we look through all of these different equivalent kernels, and we test the runtimes, and we see how fast they are, and then we choose the fastest one.
- 14:16
It's really that simple. And beyond a certain point, it becomes infeasible to search, uh, to, to profile the runtime of every p- every possible kernel, and so we start to use things like Monte Carlo tree search to sort of prune down the search space.
- 14:32
But at the end of the day, uh, it is fundamentally a search problem.
- 14:38
So what are these kinds of optimizations that end up getting found through this search space? Um, kernel fusion is a very popular one. Uh, simply, you know, you have operation A, operation B.
- 14:52
Operation B operates on the output of operation A. In this example here, we have sine and then followed by exp2. So exp2 operates on the output of sine. So the naive way is we go and we, uh, we do...
- 15:08
We load the tensor from memory into the compute unit. We do sine. We write the tensor back into global memory, and then we read the same [laughs] stuff back into the compute unit, do exp2 again, and then we write it back into memory.
- 15:23
Uh, but this is really bad. Uh, in fact, data movement in GPUs is usually like ninety-nine percent of the energy spent and the time spent. Very few, very small amounts of time and energy are spent on actual compute.
- 15:36
So instead of doing all this round tripping, what we can do is we can just merge exp and sine into the same kernel, load the data in once, and then write the data back out once we are f- we done, we're done, we have our final results.
- 15:53
So what does this actually look like in practice? Well, on the left, we have an unfused graph. It's a very, very sloppy, naive graph where we're doing a whole bunch of these different operations, and then in between them, we always have to write our result back to memory and then read it back into the compute unit for
- 16:09
the next operation. And what our compiler has done here on the right is been able to merge all of them down into one kernel. And like I said how data movement is ninety-nine percent of, uh, runtime, the, the crazy thing is that this real complex kernel on the right here actually doesn't take much longer than any one
- 16:30
of these kernels on the left here. So this whole kernel in ag-- or this whole graph in aggregate is far, far faster than the graph on the left in aggregate.
- 16:41
So one of the real big achievements that we've had, uh, recently with our search technique is we were able to find FlashAttention. FlashAttention is a very, very, very complicated algorithm that, uh, took about five years for the industry to discover, uh, for, for somebody in the industry to discover.
- 16:59
So, uh, TreeDAO discovered it in, uh, twenty twenty-two. Transformers came out in twenty seventeen, and yet this is, like, a really, really important optimization. And our, our compiler now is able to find this completely by itself.
- 17:14
Uh, so again, what do we do? We take in the naive multi-head attention graph. We run all of these different simple rewrite rules. We build out this huge search space.
- 17:24
We profile a bunch of these kernels and find the fastest one, and the fastest one in this case just happens to be FlashAttention. And to our knowledge, we're the only compiler in the world that can do something like this, uh, and it's because we're able to leverage search.
- 17:39
Again, this is an extremely complex optimization here. It's, it's not at all obvious, um, to program into a compiler.
- 17:48
So we did a little announcement about this. Uh, on the right you can see my a-announcement tweet. Here on the left you can see, uh, the generated FlashAttention kernel in green here, and then in white we have the intermediate representation.
- 18:01
It might be a little bit tough to see. Um, but yeah, this is the generated output code.
- 18:08
And so, okay, once we have this really fast, uh, kernels that are generated out of this search function,
- 18:15
what do we do? Do we just directly, uh, generate the CUDA code from that and then run it? We could do that, but there are a set of optimizations that we know will never be harmful.
- 18:26
We don't know exactly how much they will help, but we know they will always be helpful. And so what we do is we run these deterministic optimizations on the output of our search process.
- 18:36
And these optimizations are things like buffer reuse. So obviously, we wanna minimize the amount of memory, uh, we use, and so we want to optimally reuse all of our memory buffers.
- 18:50
And because we have the entire workload specified as this big graph ahead of time, we can have our compiler go in there and say, like, "Okay, uh, in this example right here, buffer one is never being used at the exact same time as buffer three."
- 19:05
And so anytime we-- once we need buffer three, we know buffer one is done. It's not going to be used again. And so what we can say is, "Oh, buffer one and buffer three should actually just be the same buffers."
- 19:17
And so we see in the bottom here, that's exactly what we're doing. We're just saying that these two are the same exact memory buffer. So this is how we can optimally, uh, uh, reduce our memory usage.
- 19:29
Another way we can optimize our final graph here is we issue the kernels all at once. So in traditional, uh, inference, what you do is you have a CPU dispatch a GPU kernel, and then the GPU runs that kernel.
- 19:44
We wait, we wait for it to finish, it goes back to the CPU, the CPU then dispatches the next kernel. That round trip to the CPU and then waiting on the CPU to dispatch the next kernel takes a lot of time.
- 19:56
And so what if we were to dispatch all of our kernels ahead of time, uh, and then the GPU would just run through them one by one by one?
- 20:03
So we do that in our compiler as well. Uh, and so we, we build this big queue. We can actually see the difference on the left-hand side here. The launch time, we actually have to wait quite a while to launch.
- 20:16
Whereas here, we can launch all of the kernels at once, and we save a whole bunch of time.
- 20:23
So Luminal was from day one always an inference library. Uh, it, it was never really designed with training in mind, but due to the s- extreme flexibility that our graph representation gives us, we were able to actually build an external crate, an external library that is an autograd engine.
- 20:42
And it works, uh, directly in Luminal, and it basically derives, given a forward graph, it derives a backward graph and then attaches that to it. And then we run our downstream compilers, which means we basically get training for free.
- 20:55
Uh, so all of the compilers that we have for inference, the search process, all of that also works for training. So it runs right on the, the backward pass as well.
- 21:05
Um, this is pretty neat too that it was added as an extension because to my knowledge, I don't think any other ML library out there is able to do this.
- 21:13
Any, any library that, uh, supports training has to have it as part of their core, uh, whereas we're able to add it in as an external thing, which means somebody else can come in, external contributors can come in and just write their own autograds or their own gradient sharding or their own really fancy training setups.
- 21:33
So that's sort of a brief overview of where we are today, the features we have today. Uh, what's to come? Well, we're really excited about adding more hardware support in.
- 21:43
So right now we support CPU, uh, CUDA, and Metal. Uh, what we really wanna do is support AMD, uh, Tensorring, uh, Groq, and TPUs, um, because these are all, like, really exciting hardwares out there.
- 21:57
We wanna break the CUDA moat ideally and, uh, sort of democratize ML across all these different hardwares. Um, [lip smacks]
- 22:05
we wanna do distributed inference and training, so we wanna do full 3D distributed, uh, through data parallel, pipeline parallel, tensor parallel. Um, [lip smacks] and, uh, we wanna do RL. So a common bottleneck in RL is we want to-- Basically, we run our model on the GPU, but we run our environment on the CPU.
- 22:25
Uh, and then that back and forth is the huge bottleneck. So if we can codify environments, we've done this for very simple environments, but we wanna see how complex we can go, codify the environment in the Luminal graph, and that gets optimized with the rest of the model through our, our same compiler flow.
- 22:43
And so basically, we run the forward pass of the model and step the environment all on the GPU. Um, so this is, this is super exciting 'cause I think it could dramatically accelerate, uh, reinforcement, reinforcement learning workflows.
- 22:57
Um, our Dyson Sphere unfortunately is pending our Sequoia fundraise, so, you know, reach out to us if you have any info on that. Um, but what we've really been working on recently is the Luminal Cloud.
- 23:09
So what we've done, because we were able to represent these models as graphs, if you're working on a model in Luminal, you can do graph.export, get a file out, upload that file to the cloud, and then get a serverless inference endpoint, and we handle everything else.
- 23:23
So we handle optimization, we handle batching and queuing, we handle turning the, you know, provisioning the machines. Um, it's totally serverless. You only pay for when your graph is actually executing.
- 23:35
So we think we can deliver the simplest, fastest, uh, most straightforward cloud experience out there. Um, so yes, come join us. Uh, there's the link to the, uh, link to the, the repo.
- 23:48
We would love, uh, PRs. If you have any ideas, uh, please join us. And we're, we're really pushing into territory that's only been pr- covered by frameworks that are orders of magnitude more complex here.
- 24:01
So it's, it's a really exciting time. [lip smacks] Uh, simplicity really allows us to do these innovations far faster than frameworks that have so much more overhead. Uh, so it's super exciting.
- 24:13
And then if you're a startup or a company that has an inference workload, uh, reach out. Um, we're building again the simplest, fastest ML cloud in the world. And so please reach out to me.
- 24:23
I'm at [REDACTED:email_address], or you can just go to luminalai.com. Um, we'd love to hear what your workload is and if we could help you out. Thanks, guys.