Luminal - Search-Based Deep Learning Compilers - Joe Fioti

Hey everyone. I'm Joe, I'm one of the creators of Luminal, and so today in this talk I'm going to explain what it is, how it works, and why we really think that it's the future of ML libraries. So, the title of the talk is "Radical 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.

So, deep learning fundamentally is very simple. It's simple linear algebra, and so basically all that really is, is, is, you know, scalars, vectors, matrices, and tensors, and then a couple core ops— a couple core operations. So we have additions, 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. Um, PyTorch is one of the most famous libraries out there. It has over 1,200 operations, uh, and has over 15 different data types.

It runs on a ton of different kinds of devices: CPU, CUDA, AMD, TPUs, some NPUs out there. 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.

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 want to support a new op, or a new data type, or a new device, your complexity starts exploding.

So what that, what's happened to PyTorch is it's now over, like, 3 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 extend, or use, or build anything inside of.

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. 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? So what we do is we have, uh, 12 operations that were very, very, very simple.

Um, and so we have x2, log2, sine, and reciprocal, and square root. Those are our unary operations. 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. 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.

Is that really it? I mean, that's not a lot of ops. 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, 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 1. Division is really simple because it's just multiplication and a reciprocal on b. 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 broadcasted multiply and then sum reduce it down, and you get the output of matmul. So you don't need a matmul op. 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.

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, 5 to 10 years ago when dynamism was very important.

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. 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. 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.

And aside from that, the entire model is static. Um, and so what we do is we specify these models as directed asymptotic graphs of operations. Soright 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 are sum reducing them.

Andright here, like we talked about, this is a matmul. This is what a matrix multiply is. And soright on screen, this is all that a, uh, single dense, uh, neural network layer is. It's just your, your, your matrix multiplying by a weight matrix.

So on theright-hand side here, we see a much larger model. And so these graphs, you can kind of see get quite a bit more complex, uh, but they are capable of fully specifying these models. 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 5,000 lines of code. 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. 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.

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. 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 want to, again, highlight, um, what this simplification gets us. 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. 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 KubLoss which then sit on top of CUDA. 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.

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.

So generally we want to 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. It's just our library, our graph, our compilers, and CUDAright beneath.

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 going to take these graphs and we're going to feed them through compilers, and they're going to spit out much faster graphs.

So how does that actually work? I mean, we're not the first to think of compilers here. We're not the first to think of ML compilers. Why haven't other ML compilers taken over the world? 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. 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.

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 want to compile stuff down to. 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.

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. 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 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. 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.

Um, but the TPUs are very, very simple. So the TPUs are extremely fast, very good performance per watt. And we want to basically keep walking down this path of simpler, faster hardware and, uh, more complex software. 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. Uh, VLIW stands for very large instruction width. And 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. 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? 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 to solve Go in one shot and always find the best move, 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 going to do this, like, really fancy tree search, and we're going to use this guided neural network to, to go through it.

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 of Go boards, we are searching through logically equivalent GPU kernels. So this means that we don't need to handwrite 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. 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 Egglog.

Um, this library basically uses eGraphs 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. And we specify out, you know, 20 or 25 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. 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. 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. It's really that simple. And beyond a certain point, it becomes infeasible to search, uh, to, to profile the runtime of every, every possible kernel. And so we start to use things like Monte Carlo tree search to sort of prune down the search space.

But at the end of the day, uh, it is fundamentally a search problem.

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. 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, 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 stuff back into the compute unit, do exp2 again, and then we write it back into memory. Uh, but this is really bad.

In fact, data movement in GPUs is usually like 99% of the energy spent and the time spent. Very few, very small amounts of time and energy are spent on actual compute. 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, we're done.

We have our final results.

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 the next operation.

And what our compiler has done here on theright is been able to merge all of them down into one kernel. And like I said, how data movement is 99% of, uh, runtime. The, the crazy thing is that this real complex kernel on theright here actually doesn't take much longer than any one 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. So one of the real big achievements that we've had, uh, recently with our search technique is we were able to find flash attention.

Flash attention 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. So, uh, TreeDAO discovered it in, uh, 2022. Transformers came out in 2017.

And yet this is, like, a really, really important optimization. On our, our compiler now is able to find this completely by itself. 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. We profile a bunch of these kernels and find the fastest one. And the fastest one in this case just happens to be flash attention. 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. Again, this is an extremely complex optimization here. It's, it's not at all obvious, um, to program into a compiler. So we did a little announcement about this. Uh, on theright, you can see my, uh, announcement tweet.

Here on the left, you can see, uh, the generated flash attention kernel in green here. And then in white, we have the intermediate representation. It might be a little bit tough to see. Um, but yeah, this is the generated output code.

And so, okay, once we have this really fast, uh, kernels that are generated out of this search function, 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.

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. And these optimizations are things like buffer reuse.

So obviously we want to minimize the amount of memory, uh, we use. And so we want to optimally reuse all of our memory buffers. 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 exampleright here, buffer one is never being used at the exact same time as buffer three.

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.

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. 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. 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. 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.

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.

Whereas here, we can launch all of the kernels at once and we save a whole bunch of time. 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 extreme flexibility that our graph representation gives us, we were able to actually build an external crate, an external library that is an auto grad engine. 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. Uh, so all of the compilers that we have for inference, the search process, all of that also works for training. So it runsright on the backward path as well.

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. 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 auto grads or their own gradient sharding or their own really fancy training setups.

So that's sort of a brief overview of where we are today, the features we have today. Uh, what's to come, we're really excited about adding more hardware support in. Soright now we support CPU, uh, CUDA and metal. Uh, what we really want to do is support AMD, uh, TensorRN, uh, Grok and TPUs, um, because these are all, like, really exciting hardwares out there.

We want to break the CUDA mode ideally and, uh, sort of democratize ML across all these different hardwares. Um, we want to do distributed inference and training. So we want to do full 3D distributed, uh, through data parallel, pipeline parallel, tensor parallel.

Um, and, uh, we want to 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. 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 want to 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.

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 because I think it could dramatically accelerate, uh, reinforcement, reinforcement learning workflows. 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. 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.

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. 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 where you 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 covered by frameworks that are orders of magnitude more complex here.

So it's, it's a really exciting time. 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. 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. I'm at joe@luminalai.com 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.