Lightning Talk: Coding Agents for Compiler Construction: Beyond the... Reza Rahimi & Stefan Krassin

This talk features Reza, the CTO of YASSP, as he presents on coding agents for compiler construction, focusing on optimizing code generation through large language models (LLMs). The speaker discusses the limitations of traditional AI methods in compiler design and proposes an innovative approach that leverages LLMs as autonomous agents capable of generating optimized kernel code for GPUs. By integrating an evaluation system that tests correctness and performance, the process transitions from a linear prompt-response model to a feedback loop where the LLM continually refines its output. The session highlights practical applications, such as optimizing specific layers for hardware like the NVIDIA RTX 6000 Pro and demonstrates significant performance improvements, achieving speeds up to six times faster than traditional methods. Reza emphasizes the potential of this agentic system to create a more efficient compiler that can adapt to various hardware platforms.

Awesome. Thank you, everyone for joining. I'm Reza. I'm the CTO at YASSP. And today I'm presenting you coding agents for compiler construction beyond the AI assistant paradigm. So, we want to go beyond just having a coding agent as an assistant to the developer and see can it replace our compilers. Why are we solving this problem? Because uh a lot of AI problem is the workload is not

designed for the target hardware that you're targeting to. And we want to optimize your model for your target hardware. But how to use LLMs and how to use agents to do this? Um We have started first by prompting LLMs and asking them to write code for our target hardware. Probably most of you there the first target hardware would be in video. So, we ask it to write

CUDA code or God forbid ask it to write PTX code. Um and then see what the result would be. And then you would use NVIDIA's compiler um whatever hardware you're choosing their own native compiler to get to byte code machine code and run it on the GPU. The next step that you can add it to make it a little more intelligent to evaluate the kernel that you're

writing. Does it have the global uh the creators? Is it compilable? Does it compile? Is it correct or not? So, an eval system can help an LLM understand if it's doing a good job or not. So, how can we just turn it from being a linear path from prompt to to an agentic system. Now that we have the eval tool, the eval compile your code, check the

correctness, benchmark the code, get the latency, also profile the kernel. Now that we have all this information, we're not bound to only deploying to one LLM. We're not bound to only having it A to B uh scenario or like a sequence scenario. It can be a loop. We can get and extract this data, give it back to the LLM, and ask it for writing uh better and

more optimized kernel. And in this case, I'm not talking about just making matmul more optimized in A lot of engineers have been spending many, many hours with knowing all the details of the hardware doing that. What now we can do is one a step further. Like you have, let's say, a mumble layer, and you want to optimize it for what's your target hardware? Let's say Nvidia RTX

6000 Pro. Can we change the math here to make sure this is optimized for this target hardware and this layer that we're using with this type of shape with uh um is it dynamic shape or static? What's the uh bandwidth of the dynamic shape? Is it like batch one or batch five or batch thousand? So you can optimize for each of these layers. And also, you can

provide all this information to your LLM's knowledge base. Like what is this hardware that you're optimizing for? And on the other side, what is the mathematics of this model? Like my model has matrix multiplication, and then reduction. How can I optimize this? our code generator using as an as an agent as an autonomous agent. These days like autonomous agents are becoming a meme. Now that we have

our code generator as an autonomous agent, we can build a compiler. What's a compiler? Front end and back end. And back end is an optimizer and code generator. Can I replace my back end with with an autonomous agent? And also, now that we have uh this whole system, can I apply to different hardware? Yes, you can. You can write Hip kernels for AMD. You can uh Nicki

kernel for AWS Trainium and a lot and many other And also, you don't have to have the input as text.md. It doesn't have need to be a text file. It can be a PyTorch So, now it's kind of turning into become a compiler. It takes a PyTorch file. It turn it into an intermediate representation. And now the intermediate representation goes to our back end. And in the

back end, you do optimization and code generation. And generate code for that piece of uh the model, the sub model that you have. So, we did this for Mumble layer, as I said. Uh in Mumble layer, we were able to replace uh two of the matmul operations with one matmul and reduction. By we, I mean our autonomous agent. And we ended up running it as uh six

6.2 times faster than reference. By reference, I mean Torch compile. And ended up running the whole model three times faster. So, we believe with using agentic systems as compiler, we can build the that optimize your whole model, not looking at layer by layer optimization, looking at sub sub model or even future the whole model, and optimize the whole And yeah, that's uh what we're building at YASPA's,

our agent AI compiler. Happy to answer any question. Yeah. Thank you. That's a That's a great question. Checking the correctness is the most important thing. Um so right now we pick the reference PyTorch and run the reference PyTorch with the optimized version and report them to the user back what's the uh absolute tolerance and what relative tolerance is. For now our cap is 10 to the power

-3 for both. And uh if it's lower than that, you will accept it as a as a correct kernel, but we tell the user what's the tolerance is. Like what kind of input data do you get like a set of input data? That's a good question. So if the user is Sometimes user doesn't want to provide the their input data to us, we generate randomly generate those.

And when the user provides the the data, we use user data. Or like if they give us synthetic data, we try to be in the range like even randomly generated data in the range of the data they they have. And of course we need to know the shape of the tensors as well. And And from your So from your experience like and sorry my mistake, is it

enough for like getting the same It's enough for getting the same output from the LLMs. Uh we tested with LLMs, we tested with CNNs. Yeah. Okay. Thank you. There's something to consider here. when we have a beta model that is equivalent to this. Yes. Suppose that >> Yes. Uh depends. Some for some of the layers CUDA, for some of the layers PTX. Yeah. PTX is NVIDIA's lower

level. Optimize for each of Exactly. Yes. For the whole model, and that specific layer six times faster, yeah. And it's you you can't Let's say if I come and tell you I made matmul six times faster, I'm lying to you. it's the whole combination of the layers together that you optimize and makes it the whole model running fast. So, like you have a bunch of mathematical operations,

and then we look at the GPU What's your occupancy? How much memory you're using? And in GPU, you want to use the memory as little as possible, and you want to offload the work more uh more of the computation to the thread. And that's one of the works that we do. So, we increase the occupancy of your GPU, and we try to shuffle the operations, fuse the

operations, and these are all style of optimization. It generates code and try We call it an evaluation search. Uh so, it generates the code, and then the eval tool runs, it provides some feedback to the model, engineer the context of the model, and then we We're not using even one model. We use uh up to 10 models, and then we generate another 10 kernels, and based on

those one they eval and pick three candidates, bring it back to the system, and this loop continues until we reach to a kernel that is accurate and it's fast. For a small model, for example, what is the time how long does this process take? Good question. For let's say use ResNet-18, it took between half an hour to 1 hour. But we also use the caching system. So

the first time you compile model, it's going to take longer. Second time it's going to be much faster because most of it is already cached in the system. I also have the same question. Yes, please. Can you elaborate on the caching system? Because you mentioned that they are not using layers or you're not really changing anything. In the implementation part. Yeah. And you want to be able

to use the same layers for both training and inference. Yeah. My question is if you compile it once, like maybe you discover one innovation, let's say. Then can you sorry, make a library of innovations and then use it? That's a great question. So between the loop, the genetic code generator and the compiler, there is a database. And all these finding you put it into database, fingerprint, which

is binary of that specific part of the model and it iterates. So next time that we see one model, it's not going to take longer. We already know how much to move the model. But sometimes we instead of using that, just to skip the code generator, we feed the new the optimized one to the code generator to see if it can go one step further or And

we also have the caching on the customer side. So if someone wants to if they don't want my code to get out, they can hold the they can have all the data displayed. And if you train the system or run a very long training job, do you search the baby space? Sometimes, yeah. So our default is eight iterations, but if you let it run in 24 iterations,

you'll get better credit.