Multi-GPU Kernels, Intelligence per Watt, Heterogeneous Inference, and More | YC Paper Club

Five researchers on where the performance is still hiding, at a kernels-and-chips edition of YC Paper Club in Mountain View.

Training and inference are pulling apart — different chips, different data centers, different everything. At a special kernels-and-chips edition of YC Paper Club in Mountain View, host Francois Chaubard convened five researchers on where the remaining performance is hiding: multi-GPU kernels, intelligence per watt, whether AI can write kernels, heterogeneous inference infrastructure, and running an entire game engine on the GPU.

Watch on YouTube

Timestamps

0:00 – Francois Chaubard: The case for chip and kernel specialization
7:16 – Stuart Sul: Parallel Kittens - Systematic and Practical Simplification of Multi-GPU Al Kernels
21:29 – Jon Saad-Falcon: Intelligence per Watt - Measuring the Intelligence Efficiency of Local and Cloud AI
31:05 – Mark Saroufim: When Al Starts Writing Systems Code
47:04 – Misha Smelyanskiy: Why AI Inference Needs Heterogeneous Hardware
1:04:33 – Brennan Shacklett: Building a High-Throughput Game Engine that Runs ENTIRELY on the GPU
1:15:35 – Wrap-up & what's next

Transcript

Francois Chaubard: Welcome to YC Mountain View, a special edition of YC Paper Club. Some of the feedback was about making it theme-related. So today it’s the YC Kernel and Chip Club. The trend I see in the lab with Stu and Jon, which we talk about all the time, is this specialization that’s happening. We saw it with TPUv8, the I and the T version, zebrafish and sunfish—is that what it’s called? To make an ASIC, you need such activation energy, and there wasn’t enough demand for there to be such a split, but we just see that at the chip level, there’s going to be massive specialization. The specs you would need for a training data center are going to be very different than the specs for an inference data center. You don’t need any bandwidth in and out for a training data center.

You can literally send a spaceship to the sun, come back with a weight file, and it’s the same, right? You can’t do that with inference. There are so many things like that that we’ll talk about. And there’s so much juice left to squeeze on the CUDA side, on the kernel side. I’m really excited to have some folks talk about that. And there’s juice to squeeze on the algorithm side. Do we need to go to the big heavy model every single time for every single query? If I’m asking what one plus one is, why do I have to do N FLOPs every single time? So there’s so much juice to squeeze still on the algorithm side, on the software side, on the chip side, and on the data center side. First, we’ll have Stuart Sul out of my lab at Stanford, one of my lab mates.

Which is the strip mall research lab now. Used to be called Hazy, formerly called Hazy, research scientist at Cursor and trains Composer. Then we have Jon, CS PhD lab mate with me as well, also co-advised by Azalia Mirhoseini, focused on LLMs, MLSys, and hardware. He really pioneers this idea of intelligence efficiency—intelligence per watt, intelligence per joule. He’ll be talking about that. Then we have Mark coming up next, a former PyTorch maintainer, co-founded GPU MODE with my friend Casey Aylward, and is co-founding Core Automation with OpenAI VP Research, Jerry Tworek. Then Misha, hardware-software co-design, has been working on that for two decades and ran AI infra at NVIDIA. Also worked on hardware-software co-design at Meta, and is founding a very exciting new company as well. You’ll hear about some of it. And then finally, Brennan, who I also met at Stanford, is not in my lab, but having done a lot of RL working on self-driving cars and things like that, it’s amazing how much of simulators, when you call environment.step, is still run on the CPU.

And so that’s usually the bottleneck for a lot of your on-policy rollouts, which makes—let’s call it RL is now taken over by RLHF terminology—but classic RL very, very difficult. And so using GPUs to do weird shit, as he says. He’s one of the best people at doing weird shit with GPUs. I’ll just give two slides on my high level. I’m not the person that knows about this stuff the best at all. But from my vantage point, we’re going to see this proliferation, and we have sufficient demand now because there’s so much demand for tokens that it makes sense to specialize at the chip level. That will pay out and allow you to go through the full new product introduction on the ASIC for specialization. And I think the Zebrafish, Sunfish, TPU, V8 is the first that we’re seeing of this. But I think it’s going to get much, much more split out.

We already do see it a little bit where a lot of the stack for inference—people will go to NVIDIA for prefill and then they’ll go to Cerebras for decode engine. And so that’s already splitting out. But there’s another one that people really don’t talk about, which is batch size one inference. And I think that when latency really, really, really matters—for example, voice agents, when you get on the phone call, and it says, “Hi, how you doing?” and it waits eight seconds because it has to go through this whole process. And it’s like, “Oh my God, well, they’re prioritizing throughput.” And it’s like, “Well, that’s not good. It’s not going to work.” And if they were prioritizing batch size one, they’d run out of GPUs, and it’d be so expensive. And so that’s a chip issue. And so we can solve those types of things.

And then on the data center side, equal amounts of specialization, especially if you’re going to use different chips for each data center. So the primary metric that you care about if you’re training models is basically time per step. And you want to train up bigger and bigger models. The location—you just don’t really care where it is. It doesn’t really need to be close to users. You can put it literally at the sun, like orbiting the sun and just training, training, training. The types of GPUs, the interconnects, you’re going to have to have all-to-all communication. You have to pass around a lot of gradient. And for inference, you have to have all-to-all just for the MoE, but other than that, you don’t. And so you can shard much more nicely. There are a whole bunch of other things to consider that justify splitting it out.

And then our speaker’s going to go much more into detail on that. The next one on 7/29, we’ll do a robotics-focused one. There are a bunch of YC companies that might actually bring some robots too. Ironically, this is how YC actually got founded. The original story across the street is that Trevor Blackwell was starting a company called Anybots, which is the robots right over there. And he had some extra space, and PG asked him if he could take over some of the space on Tuesday nights to run this thing called Y Combinator. And then went across the street, and then YC took over all of the Anybots stuff. And so it’s pretty cool to bring that full circle. And then for the next one after that, I love your feedback on what we should do. Maybe this is an idea—working on RLVR and some of the latest.

Everyone’s talking about these distillation attacks and things like that. It’s actually not clear to me what the distillation strategies that people are really trying that are successful out there. We were talking about it a lot in lab today. It’s not clear what the optimal strategy is—if you had a weight file, how to distill from it, or if you just had an API, how would you actually do it? This is the way that I was explaining to the lab about what I think is happening right now: Claude motherbirds into Kimi 2, and then now Kimi 2 is post-training Thinking Machines. And so it just keeps going and going. Now I’ll hand it over to Stu. Everyone, welcome Stu.

Stuart Sul: All right. Hello everyone. My name is Stuart. I’m a CS PhD student at Stanford. I’m also a researcher at Cursor, where I train Composer. Today I’m here to talk about ParallelKittens, which is our work on making multi-GPU AI kernels simple and fast. One quick note before I begin: the final deliverable of this paper is a CUDA framework that we refer to as ParallelKittens. But the goal of today’s talk is not to promote my open source library, but rather to convey the set of lessons that I learned while I was doing this work, which is a mental framework for thinking about multi-GPU kernels and GPU networking in general. So I hope you’ll walk away with the same set of insights that I was able to gain. Let’s try to motivate this a little bit. Over the past few years, the ML systems community has devoted a lot of effort into single-GPU efficiency.

For instance, we have IO-aware algorithms like FlashAttention, linear attention models, mega kernels, and many programming libraries and frameworks for writing efficient kernels across many architectures. These works mainly helped us squeeze the most out of the single GPU. Thanks to all of these efforts, what we now believe is that GPU networking is the major bottleneck that’s remaining, and that there’s a lot of exciting opportunities left in that area. For example, networking can still consume up to 50% of total runtime for workloads like Llama prefill. Obviously, this number can vary wildly depending on the workload. But the point I’m trying to make is that networking weights GPU utilization. So the standard practice nowadays is to overlap inter-GPU communication with intra-GPU operations like memory access and compute. The granularity of the scheme has gotten much smaller nowadays. It used to be we overlapped the transfer of entire tensors or matrices.

Nowadays, we do it at tile or even token granularity, meaning it’s a few kilobytes or even less than one kilobyte of data. On the other hand, hardware advancements have been quite exciting for networking. For instance, we have this thing called in-network compute. In the past, if you wanted to implement a collective operation like AllReduce, you needed to implement an algorithm where the GPUs do some arithmetic on local data, send the results to peer GPUs, repeat that N number of times, and then you end up with the results. But with in-network compute, you can offload all of that arithmetic communication to the networking fabric, and the GPUs are free to do whatever useful computation they need to do. You also have asynchronous bulk device-initiated networking, which was not possible a few generations ago. And you also have scale-up architectures like NVL 72, which packs 72 GPUs inside a single NVLink domain.

Soon, this is going to extend to hundreds of GPUs. The problem we are interested in solving under all of these settings is that it remains very difficult to write multi-GPU kernels that, one, do fine-grain overlap of compute and communication; two, exploit all of these exciting recent hardware characteristics; and three, stay simple and easy to maintain for both humans and AI agents. I emphasize that it’s the conjunction of these three characteristics that makes this problem especially difficult, because you currently roughly have three alternatives. One is to use off-the-shelf libraries like NCCL, which only expose coarse-grain communication primitives and leave a lot of bubbles in the hardware execution. Two is to use compiler-based approaches, but we find these compilers to be quite suboptimal. They produce kernels that are sometimes slower than non-overlap baselines. Or you can use low-level primitives like OS interprocess communication calls or PTX assembly.

These often require reverse-engineered understanding of the networking hardware. In general, they make code really complex and hard to maintain. So the main question we are asking here is: is there a small set of principles or trade-offs that clarifies the multi-GPU AI kernel design space? And can we use that understanding to build a minimal set of programming abstractions or primitives that really simplify writing efficient multi-GPU kernels for production? PK, or ParallelKittens, is the answer to this question. Before we get to PK, let’s quickly cover the GPU fundamentals for those who are not familiar. Here’s a very simplified view of a GPU. Inside a GPU, you find this thing called an SM, or streaming multiprocessor. An SM is very similar to a core on a CPU in that it performs arithmetic, logical operations, and sometimes matrix multiplications. They’re the juice of the GPUs.

They have all the FLOPs, all the compute you want. Naturally, you have a lot of them. On Blackwell GPUs, you have 148 SMs. On Hopper, you have 132 SMs. In order for these SMs to perform computation, you need data. So you have the memory system, the crossbar, the L2 cache, and the HBM, which serves as the main memory for modern GPUs. You need something to initiate operations on the GPUs, so you connect the CPU and GPU through PCIe. On data center servers, you usually want multiple GPUs, so you have many GPUs. In the past, if you wanted these GPUs to communicate with each other, you had to make the data go through PCIe. But PCIe is a shared resource. It gets congested. It’s relatively slow. So we have NVLink, which allows direct communication of data between the GPUs.

And that’s essentially what NVIDIA DGX machines are all about. DGX has been around for quite a while. It’s widely used in production AI training and inference. It’s a collection of GPUs individually connected to the CPU or the host through PCIe, and all of the GPUs are interconnected via NVLink. NVL72, a more recent addition from NVIDIA, is also quite similar. It’s a collection of 72 GPUs connected via NVLink. From a programmer’s perspective, that’s essentially all you need to know about NVLs and NVL72s. So now when you hear somebody saying they’re trying to optimize GPU kernels, what they’re usually saying is that they’re trying to make sure that all of the SMs stay busy. Because, as I said, SMs contain all the FLOPs, all the compute. So you want them to be consuming those precious FLOPs all the time. The way you do that is, for single GPU kernels, you pipeline the memory loads and writes such that you overlap computation happening on SMs with the memory operations that read from the HBM to the SMs, so that when the current round of computation is done, the data for the next round of computation is residing on the SMs and they’re ready.

For multi-GPU kernels, a similar idea applies, except that you’re overlapping computation with communication with other GPUs, so that when the current computation is done, the data for the next computation is ready and fetched from remote GPU HBMs. Let’s get back to PK. PK is based off three key trade-offs we identified for multi-GPU kernel design. The first trade-off is a transfer mechanism. It turns out that there are three ways to send data over NVLink. One is copy engine. Two is tensor memory accelerators, or TMAs. Three is plain register instructions like LD or SD. Each has its own pros and cons, and naively relying on one transfer mechanism leaves a lot of performance on the table. For example, copy engine is ideal because it’s a separate DMA module sitting on the GPU and it does not require any SMs to transfer data between the GPUs.

But it turns out that the copy engine requires a really large message size or per-transfer data size for it to saturate the NVLink. For fine-grained communication, you’re usually sending a few kilobytes of data, so the copy engine becomes a suboptimal choice. TMA, on the other hand, is able to maintain relatively good throughput regardless of the message size, and it’s also able to do that with relatively few SMs. On Blackwell, we find that with roughly 15 SMs out of 148, TMA is able to saturate the NVLink. But TMA also has its own downside because it’s not able to utilize all the in-network compute features that we discussed previously. If you want that, you have to use the register instructions, but register instructions are also a big pain because you have to think about coalescing and register pressure. All the technical details aside, the point I’m trying to convey here is that choosing the right transfer mechanism for the given workload really matters.

For example, for a commonly used multi-GPU operator like All Gather JAM, we find TMA to be quite optimal, whereas a lot of previous approaches just rely on the copy engine. The second trade-off is how we schedule inter-GPU communication with intra-GPU work. The strategies generally fall into intra-SM overlapping and inter-SM overlapping. In intra-SM overlapping, you have a certain number of threads within an SM performing computation or memory work, and the rest of the threads are just doing the communication work. This is also known as warp specialization. In inter-SM overlapping, what you do is dedicate a certain number of SMs purely for communication and the rest for computation. Again, each has its own trade-offs, its own pros and cons. For example, inter-SM overlapping, as you can probably tell from the diagram, kind of wastes compute because you’re putting a certain number of SMs purely for communication, and all of the FLOPs on those SMs are not getting utilized.

Because communication and computation are happening on different SMs, you have to go through the memory system, the L2 cache, the HBM in order to signal the completion of communication and computation with each other. But for intra-SM overlapping, it’s required that the computation and communication are aligned. They have to operate on the same data. Otherwise, you run out of registers, run out of on-chip shared memory, so you really cannot do any useful computation. You also have to think about local prefetching because, if you look at this diagram, again, the red arrow shows the data path taken when you’re accessing a remote GPU’s HBM. You can see that for the remote HBM, the data goes through the L2 cache, but for the local GPU, it does not go through the L2 cache. This means that if you’re repeatedly accessing the same data, for example, like KV cache during attention kernels, the data is only cached on the far-side GPU.

So ideally what you want to do is prefetch the data, stage them into local HBM before doing any major computation. And for you to do that, you have to rely on inter-SM overlapping scheduling strategy. So again, choosing the right strategy really matters for performance. For JAM reduced scatter, we find intra-SM to be most optimal, but for JAM all reduce, we find intra-SM to be the most optimal schedule. The third and final trade-off is the design overhead. By design overhead, I’m referring to design choices in networking libraries that cause performance loss for reducing complexity. For example, NCCL’s default mode forces intermediate buffers, which adds extra data movement between the sender and the receiver. And for fine-grain communication, this overhead really accumulates. By stripping it out, you can speed up an operation as simple and basic as All-Reduce by up to 80%.

What this tells us is that it’s important for modern GPU networking libraries to explicitly expose such performance-critical controls instead of encapsulating them. So I think I made most of the points that I want to make for today’s talk. Now I’m going to go over the ParallelKittens framework, but as I said, I’m going to make this really quick and brief. So ParallelKittens builds on all the trade-offs and the principles that we discussed. It’s a highly opinionated set of CUDA programming primitives that extends Thunder Kittens, which is one of our previous works for single GPU kernels.

As we previously discussed, a modern GPU has a memory hierarchy ranging from registers, shared memory, L2 cache, global memory. You can also think about remote GPU HBM as another layer in the memory hierarchy—so peer global memory. PK provides a data structure for each layer in the memory hierarchy. You have RT for registers, ST for shared memory, L2Cache is not programmable on GPUs, GL for global memory, and then PGL for peer global memory. What PK provides is various communication primitives that take in these DR structures and act on them, utilizing the most efficient transfer mechanism for the given functionality. It also provides a program template consisting of four different workers: loader, consumer, communicator, and store to enable different scheduling strategies. But again, for more details, please refer to our paper and our GitHub repo. Okay. To see the effectiveness of PK, we compare it against popular baselines across data tensor sequence and expert parallelism, as well as pure collective communication like All-Reduce.

What we find is that with roughly 50 to 100 lines of device code, PK is able to surpass or match hand-optimized kernels that are often hundreds to thousands of lines of code. What I’m particularly proud of is the adoption. PK has been adopted by major AI companies. For example, Cursor is using it to train Composer on tens of thousands of Blackwell GPUs. Together AI is also using it to optimize its inference workloads. Okay, so that was it for my talk. I hope you learned something about multi-GPU kernels and GPU networking today. If you work on ML systems and efficiency, I hope you take a look at PK and give it a shot. Thank you so much for your attention.

Francois: Okay. Next up we have Jon.

Jon Saad-Falcon: Awesome. Hey everybody. Thanks for coming. Today I’ll be talking about one of our recent papers called Intelligence per Watt: Measuring the Intelligence Efficiency of Local and Cloud AI. This was a joint work with one of my collaborators in lab, Avanika and Orion, as well as our advisors, Professor John Hennessy, Azalia Mirhoseini, and Chris Ré. To situate us today, we like to call, at least between me and Avanika, the current era of LMs and hardware the mainframe era, alluding to the previous era of mainframes in the ‘50s and ‘60s when IBM had these computers that used to fill up entire rooms, sometimes entire floors. These were the kinds of computers that got us to the moon. These were the first commercial computers available for businesses. That’s kind of how we’re still constructing at least the data centers today, where you have these warehouses filled with TPUs or GPUs from Google and NVIDIA, respectively.

They’re both doing our training as well as our inference. It makes sense because demand is growing so substantially. We need to satisfy that demand and scale out computers as fast as possible. At the same time, this is demand that we haven’t seen in a very long time, maybe since the US built its railroads over the entire continent. We’re spending upwards of two to three percent of GDP today. That’s demanding 250 gigawatts of new data centers, which requires new power sources, new GPUs, new land, new water to actually handle the demand for both training as well as inference. Even though the demand is going up so quickly year over year, a lot of these requests don’t necessarily need frontier-level intelligence to handle them. We each have our different LM workloads that we handle per day. At least for me personally, I know most of my workloads don’t necessarily need frontier-level coding or frontier-level math LMs.

A lot of them are going through different kinds of data sets, doing different kinds of plumbing when it comes to organizing folders, creating unit tests. Despite the fact that frontier-level intelligence is what we’re continuing to push, and it’s really important that we continue to push that out, there’s a lot of smaller local open source LMs that can actually address the vast majority of LM traffic today. What we’re thinking is that the same way that we shifted from these mainframe IBM computers that filled up entire rooms to PCs from Apple and HP in the ‘70s and ‘80s, we will see something quite similar happening for the hardware today, where we’re seeing this shift away from NVIDIA GPUs and Google TPUs for every single request more towards Mac as well as NVIDIA and AMD personal consumer GPUs that sit in people’s workstations and laptops.

And this is being driven by two main trends. One is the fact that these smaller LMs are getting better year over year. We’re seeing smaller LMs, open source LMs that are quite capable, and it’s only accelerating, combined with the fact that local accelerators are also getting much better. So these are accelerators from Apple, accelerators from NVIDIA, accelerators from AMD. We’re seeing that these GPUs have more than enough memory to handle some of these larger open source LMs. We imagine this trend will continue. With this study, we wanted to ask what role local inference with open source AI can play in redistributing the inference demand from how it is today, which is mostly in the cloud, toward more distributed inference. To help us frame the study, we proposed intelligence per watt to talk about the numerator and the denominator.

By intelligence, we just mean capabilities. This looks very different depending on the tasks that you care about. Ultimately, we care about chat, reasoning, agentic tasks, coding tasks—basically anything that people are using LMs for today. At the same time, the denominator focuses on efficiency. This isn’t just efficiency in terms of the intelligence you can deliver per parameter or per bit of information used for training. We’re talking about the actual compute and energy required for running these different accelerators for all of the different workloads that people care about. What is the actual running power, wattage, or wattage in terms of power for actually delivering that amount of intelligence? For this study, we wanted to do a broad sweep and see what the trends were. We tested about 20-plus different state-of-the-art local models across Llama, GPT-OSS, Qwen, IBM Granite, all in the range of one to 200 billion parameters, some MoE, some dense.

At the same time, we wanted to test the state-of-the-art accelerators, particularly local accelerators over the past two to three years. This was from Apple, from NVIDIA, from AMD, from SambaNova, to basically understand what the upper bounds were for intelligence per watt and intelligence per joule. We wanted to do this across all the different workloads that people care about—chat, reasoning, agentic, coding tasks—basically anything that people are using these models for today, we wanted to examine. We wanted to examine more than just accuracy. We cared about latency, we cared about energy, we cared about power, we cared about compute, and we wanted to do a broad sweep and see what made sense. What we found was pretty compelling. We found that a lot of these local models are quite good and they’re advancing rather rapidly.

Up to 88.7% of these queries could actually be routed to local accelerators running local open source AI. And that accuracy is only going up. In the past two years alone, we’ve seen about 3X improvement in this intelligence delivered per watt. At the same time, we’re seeing this compound. This is combined between the better local models as well as the better local accelerators, building off of each other. When you actually examine intelligence per joule as opposed to intelligence per watt—in this case, looking at the total energy being spent for the entire workload end to end—that’s compounding even more rapidly. In the past two years alone, or actually in about 16 months, we’ve seen about an 18X improvement in the actual intelligence delivered per Joule.

This is primarily driven by better accelerators, with more memory being placed on consumer GPUs, such as the Apple M4 Max, as well as the DGX Spark from NVIDIA. This is also from better quantization techniques that can better take advantage of the available compute, as well as longer pre-training, more post-training, and more sophisticated techniques for actually making these models better. What does this actually entail for the way that we think about redistributing this inference? This means that if you were to route perfectly, you could route somewhere between 80% to 90% of the queries today to these local accelerators running these local open source LMs. That means energy savings, compute savings, and cost savings. Even imperfect routers can save somewhere between 50% to 70% of your energy, compute, and dollar cost.

This fundamentally changes the economics of how we scale inference and how we think about building new data centers and investing in new computing lines of hardware and models. At the same time, local accelerators aren’t perfect. One of the advantages of data center compute is you can do all sorts of sophisticated techniques around kernel writing, batching, and different kinds of quantization that allow you to amortize the investment across more users and more queries. The Apple M4 Max compared to the NVIDIA B200 is still lagging behind, particularly when it comes to intelligence per watt and intelligence per joule. At the same time, there are even more specialized accelerators for inference, such as the SambaNova SN40L, for which consumer accelerators are lagging even more behind. So there is an investment opportunity to make these consumer accelerators better and better.

And so we’re excited to see what the different companies release. I just want to thank our collaborators on this project from NVIDIA, from Google, from Apple, AMD, OpenRouter, and SambaNova. Without their resources and guidance, it wouldn’t have been possible to do this study. What are we looking for next? I’m really excited to see how inference engines change for this new kind of dynamic where you’re leveraging both local resources as well as cloud resources. At the same time, I think there’s a huge opportunity to think about model architectures and kernels focused specifically on energy-efficient inference. This could lead to different kinds of accelerator architecture co-designs that make this possible and accelerate this trend. On the actual analysis side, I’m excited to break down the intelligence per watt metrics even more and see how we can separate out how this looks across different tasks, across different hardware, across different deployment settings, looking at the inference engines, looking at the kernels, looking at the actual compute and memory on the hardware.

And finally, while we focus on intelligence per watt here, I think it’s really important to consider how this intelligence is actually being used in a useful manner—how we’re actually using this to do meaningful work that people care about and that humans would have to do otherwise. Quantifying how that affects GDP, how that affects wages- that’s another interesting thing that we’ll be releasing quite soon. If you have any questions, please feel free to reach out. I’m available over email, over Twitter, over GitHub. I’d love to chat if anybody’s excited about these same topics. We also have a follow-up project called OpenJarvis that’s focused on operationalizing a lot of these insights and making it possible to run the entire personal AI coding stack on device. So you don’t need to pay Claude or pay OpenAI anymore for your LMs. You can instead just run it on your laptop and on your workstation.

So if you’re excited by that as well, please reach out. Thank you so much.

Francois: Okay, next up we have Mark.

Mark Saroufim: All right. Can folks hear me okay? Francois, thank you so much for inviting me. I heard in Stuart’s really nice talk a bunch of questions around how good AIs are at writing kernels. This is going to be the gist of my talk. We’re going to go very deep into this topic. I worked a lot on systems for about five years, working on PyTorch. I also worked a lot on custom kernels as part of co-founding GPU MODE. This has been one of my main focuses at CoreAuto. I’m hoping to give you both a research and commercial angle to a lot of this research. This is a YC Paper Club, so the main lessons are drawn from these two papers: KernelBot, which is a competitive platform—basically, it’s like LeetCode for GPU programmers.

And then the second one is Kernel Guard, which is our platform for detecting cheating. I’ll explain why these are one and the same problem. I’m a systems researcher, and I’ve had to accept that my opinion matters less than AI researchers. What I want is for people to just take the same data, keep running a matmul on it over and over again, and give me a GI. But what my colleagues in research want is something that’s autoregressive. They want lots of tiny kernels, dynamic control flow, sparsity in a lot of it, and print statements. All of these things are frustrating, and they should stop doing this. Okay, thank you. Unfortunately, I can’t do this, so I have to work with these people and give them tools.

The way we typically do this is by giving them programming languages. One example—you might not think of it as a programming language, but it is—is math. Math is beautiful because it’s forward compatible. For example, the matrix multiplication API has been forward compatible for 250 years. It’s AB and then equals C. Maybe you add some args or kwargs, but it’s the same API. PyTorch is very much in a similar vein. Basically, the forward-facing printed API of PyTorch—it’s ops like MM, Softmax, QRD comp—all these things will stand the test of time. Ultimately, these dispatch down to individual CUDA kernels that may or may not be very fast on modern hardware, which is why we get programming languages like Triton, which are tile-based. They’re Pythonic, they’re fast. Triton took over the world about two and a half years ago.

But recently, especially starting with Blackwell, people are saying, “Well, Triton’s programming model is very restrictive. It’s not letting me write state-of-the-art kernels.” So people have been very interested in programming languages such as ThunderKittens, which Stuart helps maintain, CUTLASS and CuTe DSL, which are closed-source libraries by NVIDIA for really state-of-the-art matmuls. CUDA, which is a more general programming model, is not just for matmul—it’s a thread-based programming model. It’s really great for surgical improvements, but it’s quite tedious to write. And at the lowest end, you have what the hardcore people do, which is PTX or inline SASS. Typically, there are no BC guarantees. So it’s not that it’s hard to do it; it’s just that your code will stop working if you change architecture. That’s less pleasant and very time-consuming. Typically, all the maintainers of these libraries and users of them have very strong opinions about what’s actually the best trade-off between performance and productivity.

It’s just that you have to know where on the spectrum you are, how much you care about performance, and you can typically figure out what your opinion actually is, but there’s no wrong opinion. But if we were to try to benchmark this a bit and say, okay, let’s do a leaderboard. Let’s say we have a problem, like a QRD composition problem. I’ll explain what that is in a second. The goal is for people to submit kernels to this problem. What’s really cool, what I’ve been working on in GP1, is this platform called KernelBot where people just make submissions. They can pick any library they want. Then we can basically do data analysis and see which libraries are actually popular and what people actually use. At least for this archetype of people that just want state-of-the-art perf, as far as I can tell, they really prefer CUDA.

It was only when we were doing lots of GEM-related problems that people loved using CuTe DSL. But Triton shows up often in the top five to ten, rarely shows up in the top two. So this is how you should broadly think about it. What’s funny, though, is that I spent all this year working on these galaxy brain programming languages and compilers. Then starting this January, we did this very prolific, very popular competition on NVFP4 kernels. People who’ve never written GPU kernels before started to get really competitive results. They were in the top four and top five. This is Shirga. He’s a researcher, a grad student in China, and reached out to me. He’s like, “I’ve never written GPU code before. I’ve never written an open a CUDA book before. Everything is LLM-generated, but I’m number four.” This post of mine made it to r/Singularity, which is interesting.

It was my first cameo on that, so I read it. But then I was like, “Okay, this is weird.” People told me, “Well, but this guy is smart. He’s a researcher. Obviously he gets it. He’s computational and numerical.” But then I saw this other tweet by Clark, who’s a high school teacher, and he’s telling me, “Oh, hey, I just wrote my first dual GEM problem.” And I’m like, “What?” Typically there’s an order, right? You have to start with vector sum and vector mean, then you do MatMul, and then you do this other stuff. Here people were just YOLOing it, and they’re getting competitive results. They’re looking at the charts; they’re having a good time. As an educator and as a builder of libraries, this is very, very strange to me. Despite working in this field, I did not expect these results to come so quickly.

So I want to briefly talk about how we typically evaluate how good AIs are at writing these kinds of kernels. A lot of this framework is borrowed from Simon Guo and Anne Ouyang’s work on KernelBench, but the framework is the following. You typically have some sort of PyTorch reference because we all agree in the community that PyTorch is generally more correct than most libraries, but not more performant. You basically pick, you sample some random inputs, you then have another reference submission in your favorite kernel DSL. You run an eval.py where you just pass both inputs through both, see if they’re equivalent. You check for correctness, you check for performance, and then you rank something on a leaderboard. Great. This is the basic framework of how all these evals work. If you’ve been working in AI, you’re like, “Oh, it’s verifiable rewards.”

It’s an easy problem.” And when people tell me this, this is me. I’m like, “Okay.” I actually disagree. There’s a very hidden nuance here and I think we’ll dive in. I’m going to give you an example. If you’ve never had a CUDA kernel before, you’re going to do one with me right now. We’re going to write the world’s fastest vector mean kernel. Given a vector, which is a one-dimensional vector, we’re going to just take the average of it. Great. We’re going to sample a bunch of random inputs from a vector of length a million. This is great. We’re going to take the mean of this vector. Okay. This is basically our reference. What our kernel’s going to do is we’re going to write a kernel.

The hint here is that torch.randn samples values with mean zero and variance one by default. So the world’s fastest vector mean kernel just returns zero. This actually beats the speed of light. This is a wonderful algorithm. NVIDIA, I would love to hire you. Life is good. But this was not one isolated example. I’ve read a lot of reward hacks. One common one is that the AIs will cache the output and then reuse it and do something equivalent. So you’re like, okay, I want to ban this. One version of it might be you use data pointers to figure out what the past data is in Python and then you return that. So you’ll put in your eval suite, you cannot use data pointer. Then it turns out Python is a very interesting language.

Instead of doing the dot, you can say get attribute, and you get data pointer. Now you ban the string data pointer. I’ve seen the AI also split that data pointer into two different chars in appendim, and they’ll do them in different lines. So you can’t catch it easily with the regex. Let’s say you ban data pointer outright. People get upset, and then you can just use id. My point is this is very much a chicken-and-egg problem. This just keeps happening. It’s just because Python is a dynamic language, which makes it a terrible choice for these things. One of my favorite examples was by Natalia, who was a student at Stanford at the time, who showed me the craziest reward hack I’ve ever seen. The essence of this reward hack was it noticed that during our correctness suites, we were checking for correctness 15 times, and then we would do performance testing.

But when we were doing performance testing, we weren’t doing correctness testing again. So what the AI was doing is it was counting how many times it did correctness testing and giving us a correct but slow kernel. But then when it came time to do performance testing, it was giving us an incorrect kernel, but still batching all the results together and basically making everything look good. Turns out there’s a precedent for this reward hack. It’s actually Volkswagen. This is called Dieselgate. What they did was, under emission testing, the car would detect it was under emissions testing and reduce emissions by 40X. But in the real world, it would just emit a lot more, and that worked. Big credit to Tinderization, who pointed this out to me on Twitter. I was on this brutal on-call load where we kept reviewing reward hacks, and people would tell me, “Mark, your eval sucks.”

Why can’t you detect these issues? But it’s a lot, and they’re all different. It’s not at all obvious how to codify it. So I thought, okay, if people are mostly going to submit to our competitions with AI, I’m also going to review it with AI as well. The main idea was, as a human, if I see something that’s really fast, I audit it. I mark it as a reward hack. I then provide that as an example to an AI system. I get this AI system to synthesize a Regex-based detector. You deploy that Regex and you can catch a lot of cheaters. And as new things come up, you retrain the model. This idea was by Sinatra, one of the GPU MODE maintainers. Again, if we had more funding, we’d just use a model, but this is very fast and very snappy.

I want to talk about a recent problem that we had that was interesting. One of my colleagues, Rohan, is very into this optimizer called Shampoo. You might have followed it on Twitter. He was beefing with Keller Jordan about whether Muon is Shampoo or not. It’s very interesting. One of the reasons why Shampoo isn’t more deployed is because it depends on this algorithm called the QR factorization, which is very, very slow in PyTorch. I saw this as an opportunity to monetize the beef and make a very old algorithm much faster. The basic idea of QR decomposition is that you have a matrix A and you want to decompose it into two matrices. One is Q, which is an orthogonal matrix, which means the transpose is equal to the inverse, and an upper triangular matrix. This is very useful to guess the curvature of your space, which is useful for second-order methods.

And the function, the API to use this in PyTorch, basically just calls LAPACK under the hood for CPU or cuSolver for GPU. Both of these libraries, I think, need a lot more love. They need a lot more maintainers. They’ve been getting swallowed up by deep learning ops. It was really cool. We basically turned this into a problem for the community and collectively did this broad search over humans and AIs to get a kernel that’s 60 times faster. And it doesn’t nan, which is great. We can actually use this in real training runs. What’s really weird about these kernels is that they’re quite long. For example, the average submission here is on the order of 15,000 lines of code because it’s a single kernel per shape. No human would ever do this because humans have this bias towards beauty, simplicity, and elegance.

The AIs don’t seem to particularly care. What they do is implement a dispatcher where they’re like, “Okay, well, if the shapes are small, maybe I should do a shared memory QR, but if it’s really big, I should do a global memory QR. I could have custom schedules. For shapes that get tested a lot, I will use lower precision. I will use higher precision. But for shapes that aren’t tested a lot, I can cheat the benchmark hardness and use a lower precision shape.” So then you have to look at this. At the end of this, we have tens of thousands of submissions, each of which is tens of thousands of lines long. So we have a few million tokens of this. We’re looking at this and thinking, can we synthesize this into an actual beautiful and elegant kernel?

Turns out, not so obvious. I haven’t figured that out. That’s actually an open problem. But we did read the three fastest kernels very carefully. Myself and one of my colleagues, Marcus, one of the FA4 co-authors, tried to explain to people how you write these kernels. This is what we came up with. We’re like, look, step one is you understand how the GPU memory hierarchy works. Step two is you write a very fast kernel. This sounds facetious. It’s correct, maybe not very actionable, but this is directionally what we want to do. The last thing I want to talk about here is that a big part of what makes our leaderboards work is that by virtue of people competing more, more reward hacks get found out. As a result, our eval gets more robust.

To me, this flywheel is reminiscent. When people say, oh, PyTorch is more correct. PyTorch wasn’t born correct. PyTorch was slowly made to be correct. The idea was at first it was vibecoded in the sense that it was a POC that ported Torch to Python. So basically it just went from Lua to PyTorch ops. It was verifiable because it’s PyTorch to LuaTorch, which goes back to NumPy, which goes back to LAPACK, which goes back to Fortran. So we have this chain of correctness.

We worked with a lot of researchers. They told us what they wanted. They pointed out bugs, we fixed them. And then it turns out if you repeat this over nine years and you maintain strong BC numerics guarantees, you’ll get a very popular library. This is a very good process. I don’t know what’s the right way to do this purely with AI, but I suspect the way we’ve done things in an adversarial way with KernelGarden and KernelBot is probably a good hint as to how we should approach these problems. So, thank you. If you’re interested in some of these problems, feel free to reach out to me on Twitter or hello@coreauto.com. I can tell you about what I think are some important open problems that I don’t think get nearly enough attention, and that are very different from just asking Codex to go cook.

So basically the Kernel alums are very good at exacerbating Amdahl’s law effects. They will find all bottlenecks in a system, whether it be compilation times. Anything that involves waiting is very bad. As a result, I’m very interested if people have any ideas on ways of speeding up compilation. Faster JITs, more efficient AOT packaging of things like Triton and CuTe DSL are things I’m very interested in. CPU simulators of GPUs are quite good. I think a lot of people have been joking around how Fable is just phenomenal at this kind of work. This is the kind of work I’d love to see because if you can spin up CPU simulators, that means you don’t need a lot of GPUs for doing rollouts. Similarly, if you want to deploy these kernels in a real inference engine—well, inference engines, let’s say out of the box, if you’re to use something like SGLang or vLLM, if it’s just loading DeepSeek before the first inference, maybe you’ll wait about 30 minutes. This really sucks for an AI system.

I’m also very curious to hear takes on how to more robustly verify kernel correctness. The way we do it today is we just take random inputs. Some people tell me formal verification is the way. I’d love to hear more ideas that are cheaper than training a model because that’s kind of prohibitive. I’m also curious what people think about what’s in between. We can read code line by line to verify its correctness and we can test it a lot. None of these answers feel very satisfying to me, so I’m curious if people have intermediate takes. And the last one is people have been joking that GP mode has become pay-to-win. That sucks. I think it’s become pay-to-win because test time scaling is happening over days or basically on the order of one or two weeks. So given that we know how to solve this problem in these longer periods, can we speed run it in a couple of hours or a couple of days?

If you have any ideas, I’d love to hear from you. Thank you.

Francois: Okay. Next up, Misha. Thank you.

Misha Smelyanskiy: Can everyone hear me? Okay, cool. Hi everyone. I am Misha. I joined a startup called Marlowe almost a month ago. The startup is focused on building workload-optimized heterogeneous infrastructure. The main premise is that inference is a very heterogeneous workload. Different phases of inference exercise compute, network, storage, and memory bandwidth differently. When we look at it, it makes sense to co-design systems that use different hardware and different systems for different phases. I’ll give a couple of examples today. I want to caveat that there is no data—I’m not going to show you any data right now—but I will try to drive those arguments from first principles. I’m also happy to be debated on them. This is probably very familiar to everybody: this is the life cycle, the lifetime of an inference request.

Basically, a user enters a prompt. There is a CPU system that orchestrates, schedules, and batches it. Then it goes through the prefix cache, looks up the prefix cache, and whatever is left of the prompt that’s not cached, it does the prefill. Prefills are usually very compute-intensive and are done on accelerators. There is also the KV cache creation, and that can involve CPU or other accelerators’ memory, which can go over the network. Then there is autoregressive decode that does one token at a time for a given stream, which is also very memory-intensive and very latency-sensitive. Then you produce the result. There is also a speculative decode along the way that speculates on the next batch of tokens, which then needs to be verified with the main model. The main point is that different phases trigger different hardware.

The bottlenecks just keep moving. There is no one thing. I’m sorry, the slide is a little bit busy, but I felt compelled to put it in. I want to talk a little bit about the fundamental characteristics of prefill and decode.

The main metric here is what we call arithmetic intensity. Arithmetic intensity is the ratio of the amount of FLOPs that your fundamental algorithm does to the amount of data that it moves from memory. If that ratio is greater than the corresponding machine ratio—which is the ratio between peak machine FLOPs and peak memory bandwidth—then your kernel or algorithm is compute bound. If it’s lower, then it’s memory bandwidth bound. This is a roofline, and I’m sure a lot of you are familiar with it. It was actually created for HPC back in early 2012 or 2013 by Sam Williams from Berkeley, who tried to model the behavior of different HPC applications. The idea is that you have arithmetic intensity along the X axis, and for a given intensity, you have how much FLOPs you can achieve.

And so if your arithmetic intensity is less than your machine ratio, you’re in memory bandwidth bound, and if it’s greater, then you’re in the compute bound region. Of course, the reality is much more complicated than that, but that gives you a first-degree approximation. If you look at prefill and decode, you can actually see that even attention and MHA kernels, MLP, they all have different arithmetic intensity. Prefill is generally very compute-bound because you basically do attention. You do a lot of work per all the tokens that you are fetching. You can fetch the weights once, and you operate on all of them. The MLP is even more intensive because you can actually batch things, so it’s pretty compute bound. The decode is different. In decode, you actually work on one token at a time; it’s regressive.

So attention is really horrible. You’re fetching the weight for every token that you want to process, so it’s really inefficient. The MLP is a little bit better because you can actually batch it, but in reality, batches are not that big. The machines, like modern accelerators, have a lot of compute, so you end up being bandwidth-bound for pretty large batches as well. The main idea here is that even within a prefill and decode, you have quite a bit of diversity of behaviors and how different kernels stress the system. You can uplevel this a little bit. This slide basically shows different use spaces, and you can see that inference spans a very large space of workloads. Each circle here—the size of the circle, the bigger, the larger the concurrency. For each use case, we are showing how much time you spend in prefill versus decode.

For interactive chats, you’re entering some prompt, and then maybe you get some answer. But generally it’s, I don’t know, 30/70, fifty-fifty ratio between the prefill and decode, and it’s very strict latency-sensitive. For long context queries, you basically spend all the time on prefill. Then you get an answer, you do decode, but most of the time is spent on prefill. That’s why the blue dot here has a long blue line. For coding agents and the latency, you probably also have strict latency requirements. For something like long-running agents, you have long input, long outputs, you left your Claude run overnight to build your code base, and so it also can be high concurrency, and latency is relaxed. I don’t want to belabor the point. The main point is that inference spans a very large space of workloads.

Before I continue, I want to do a quick detour and talk a little bit about SRAM machines. Those machines became very popular in the last decade. What’s so good about the SRAM machine? Basically, today when you want to do the GEMM, for example, on GPU, and if weights are large—if they are larger than the size of your on-die caches—you have to go to memory. You have to go to HBM and fetch it over the inter-package interconnect into your GPU. If you have a large weight matrix, you have to do it for every token, so it’s inefficient. The SRAM machine basically keeps the entire weight matrix in SRAM memory on die, so you get a lot more bandwidth because it’s on chip. You can access bytes over cycles. The chip interconnect is also fast; you can go between different units in a couple of cycles.

So everything is kept on die, and they’re very efficient. I call them GMV, matrix-vector multiplication accelerators, which just happen to be really good for decode because decode is very bandwidth bound. They offer significantly more bandwidth and significantly lower latency, but there is a catch: because it’s all on die, there are only so many transistors you can pack. You’re limited by the theoretical limit of your die. Because of that, the amount of capacity you can put is limited. If you look at different SRAM accelerators, you have between hundreds of megabytes to tens of gigabytes of memory, but that’s it.

But it moves them much faster. So the question is, how do you leverage them? How do you make sure that the model, even after the model is sharded, you still see the benefit? After the model is sharded within many chips, you still see a benefit, and how do you characterize these benefits? I want to talk about a couple of examples where now I’m switching gears, talking about different systems. It does not have to be a GPU or an SRAM machine, but I just want to give examples of how specialization could help. People know very well prefill and decode disaggregation. People separate them because it lets each of them scale independently and removes shared bottlenecks. So you can basically ask the question: what if instead of disaggregating prefill and decode on two different systems, you disaggregate decode on a different system? Let’s say system A does prefill.

System B does decode. When is it beneficial? When is it actually beneficial? If you think about it from a TCO point of view and what we talk about is tokens per second per watt, it’s actually beneficial when the additional power that your system B adds is offset by the speedups that it brings. You can imagine if you have really short sequence lengths, you spend a lot of time in prefill, and so your benefits might not be as pronounced. For short output lengths, you actually have TCO loss because you just bundled a bunch of new hardware and you are not using it. As you increase output lengths, you’re actually spending more time doing decode, spending more time on system B. At some point, you may become TCO positive. The other example I want to give is the AFD.

As I talked about earlier, attention and MoE parts have very different characteristics. One way to accelerate this workload is to run attention on system A and run the MoE on system B. What’s going to happen, let’s say in the example of a GPU: we know very well that GPUs are really good at high throughput. They operate with very high concurrency. They deliver amazing results. You’re compute-bound, and your latencies can be lower. But the moment you start going to lower concurrency because you want better interactivity and better latency, the performance—the throughput—drops. It drops very sharply because all of a sudden you have a lot of smaller kernels, you have a lot of overheads, you’re bound by memory bandwidth, and also memory bandwidth is not utilized very well because of all these overheads. At this point, you may say, okay, what happens if I introduce system B that is really good at running those memory bandwidth, latency-bound kernels at the point where, say, GPU interactivity goes down?

And for example, you can consider offloading some of those nasty MoE kernels to SRAM machine. And because everything is on die, they have low latency, they all can do quite well for small batch and all of a sudden you see that your interactivity has been extended. Now what happens to the TCO? The TCO may not be as great as running it at a very high concurrency, but at the same time for use cases where interactivity matters more than TCO, it might be a sensible approach because there is a point beyond which GPU cannot deliver more interactivity and this extends the life. So that’s the idea of the desegregation of attention and MAE. And the last example I’m going to give is basically in the case of speculative decoding. I think everybody knows what speculative decoding is. You have the drafter that speculates on subsets of tokens, either autoregressively or in parallel for block diffusion.

Then the verifier verifies them in parallel and whatever number it accepts, that’s how much you can fast forward. So you can imagine that you can run your drafter on system B, right? Your verifier runs on system A, but you can run your drafter instead of running it also on system A, you run it on system B. And I don’t mean sharing exactly the same hardware. You can run it across different systems, like for system A, you can still run verifier on one system and drafter on another. And for system B, just offload the drafter to system B entirely. And if the drafter allows you a much better acceptance rate because it can run a much larger model at the same rate as system A can run, then you can actually see the wins. You can see the wins in terms of latency. It’ll reduce the latency because maybe system A cannot run such a large model, it takes time.

And you can also bundle multiple verifiers to the same decoder to improve the TCO. Okay. Just to wrap it up, I was talking about end-to-end co-design of the heterogeneous infrastructure. And that’s a really full-stack problem. There are many interesting aspects of it. There is a data center aspect, like when you put heterogeneous systems in the same data center, what does it mean for power density, for cooling, for break and fix, for managing? A lot of your problems that you have in a homogeneous environment become harder. There is the networking problem. I didn’t talk about it much, but obviously anytime you move data between system A and system B, networking is a bottleneck. So how do you build it, co-design it together? Where do you connect them, under what network topology, so that your latency is reasonable and doesn’t kill advantages?

And so obviously one thing, you also need performance modeling infrastructure. You need a really good simulator to allow you to go through a lot of design points and figure out what makes sense. And it needs to be well calibrated, but that’s really how you’re going to guide your co-design. So that’s basically what we are starting to do. Thank you.

Francois: All right. Next up, Brennan.

Brennan Shacklett: Hey, thanks for having me. Really appreciate it. I’m going to be presenting some work that I did at Stanford a couple years ago during my PhD. I’ve since moved on, and I’m doing an AI plus gaming startup, but credit where credit is due—this is firmly Stanford work in Kayvon Fatahalian’s group at Stanford. What I’m going to be talking about is this hypothetical idea: what if we put an entire game engine that can simulate any game you can imagine all on the GPU and run it super fast? I’ll talk about how we did that at Stanford over this project, which ran for a couple years.

The reason we cared about this is that games are a great learning environment for a bunch of different tasks. I did a lot of work with roboticists, people doing self-driving car research, and game developers themselves. They’re just a really great platform for learning skills in simulation in a low-cost way. But honestly, somewhat surprisingly, game engines are really inefficient for the throughput-oriented training workload if you just do the obvious thing, which is run a thousand copies of the engine in parallel, because you’ll have all of these copies fighting with each other. You can’t amortize any costs, and you wind up using both the CPU and the GPU hardware really inefficiently.

The better solution to this is what we call batch simulators, where you have a single game engine that actually simulates a batch of a thousand learning environments simultaneously in a big throughput-oriented batch that runs on the GPU. Hopefully this loads. This is a visualization of OpenAI’s old hide-and-seek environment. It was a multi-agent reinforcement learning environment. That gives you a sense of how many of those environments we were actually able to fit on a single GPU. You can wind up getting throughput of millions of frames per second of experience, which is incredibly useful for data-hungry algorithms like RL.

The challenge with building this that we ran into is that existing GPU programming frameworks are a really poor fit for the gameplay logic you actually need to write to build these environments. When you’re trying to build these environments, you want to be able to create them easily and experiment with different mechanics and ideas. Here’s an example of some of the logic from that OpenAI hide-and-seek environment. Agents need to run around and use the objects to hide from each other. They can lock objects, move them around, and so on.

And what you wind up with is a bunch of super branchy code. Some of these internal functions end up doing dynamic memory allocation. If you know anything about GPU programming, a lot of this is starting to sound like really bad news from a performance standpoint. Interestingly, the game industry has actually already pretty much solved this problem for parallel CPU code. Most gameplay logic for a game that you might play actually runs on the CPU. But because we’re moving from a latency-sensitive workload, where you’re trying to play a video game and you want fast response times, to a throughput-oriented workload for training, we can almost one-for-one adapt those design patterns to GPU execution and get a really efficient end result. These design patterns are referred to as entity component system design patterns in the game industry. Don’t worry, I’ll define what exactly that means.

And I’m going to be talking about how we run it on the GPU. The first concept is entities—that’s the E in ECS. That might be obstacles in the environment or the agents in blue, the little blue guys. Then we also have components. These are the data that’s attached to each entity. This might be something like the position, rotation, and then you might have action and reward from your reinforcement learning system. The ECS basically takes all of this data and puts it in big in-memory column stores. For the GPU, we actually just put these all in GPU memory and then pack a bunch of learning environments together into these unified column stores. You see that green column is what environment each of these individual objects maps back to. This unified table storage, as it turns out, is a really powerful way that we can do throughput-oriented dynamic memory allocation on the GPU while hiding it behind a really easy-to-use interface, which is just create an obstacle, create an agent.

This makes it incredibly easy to do procedural generation tasks where you might have different numbers of obstacles in each environment, which is the case for the hide-and-seek environment that I was showing you. These things are very tricky to do in a tensor-based programming language like PyTorch, where you’re looking at fixed-size arrays. The third concept, the S in ECS, is the systems. These are the actual code that runs over the entities—the rows in these tables—and builds the logic. That might be processing the actions from a neural network, from your agent, implementing collision detection, or computing rewards for the RL system. For example, the process actions component declares, “I want to run on every entity that has a position and an action component.” In this case, the agents have those two columns.

This logic will run over every single row in this table. That’s the same example I was showing earlier. On the GPU, it’s actually incredibly simple. All we have to do is map one invocation of this function to one GPU thread, and we’re able to get massive parallelism. For the programming language nerds out there, another really nice aspect of this system is it actually allows you to do runtime type polymorphism. The collision system just says, “Hey, I need things that have a position and a bounding box.” It turns out both agents and obstacles have all of those components. Interestingly, the collision system doesn’t actually need to know either of those types. Traditionally in CPU code, you use virtual dispatch for this, which is incredibly inefficient on the GPU because it breaks statically knowing register assignments and introduces a ton of issues like that.

To actually make a game out of all of this, we take a bunch of those individual systems and combine them together into a task graph that describes how the overall frame works. How do we go from taking in some actions from agents and then actually updating the state of the world? There’s quite a lot to it. This is actually a pretty small subset that goes through physics, observation generation for the agents, and so on. Unfortunately, I don’t have the time to go into a lot of the low-level implementation details to actually make the memory allocation in those tables efficient. We’re taking advantage of the fact that this is a throughput-oriented workload, and we actually do something akin to garbage collection on the GPU where we append rows and then mark rows as deleted and then use a super high-performance GPU sort to clean that all up after the fact.

The task graph is actually implemented with a persistent mega kernel that’s constantly churning through the work that’s remaining in the frame. As many of the ParallelKittens-type people will be very familiar with, the GPU just has incredibly good fast atomics that let you do all of this synchronization really efficiently. I can actually show you that here. This is a visualization where every row is an SM on the GPU. You can think of it as the amount that each colored bar is full in each row is telling you how active that SM is. So it’s pretty active most of the time. The different colors correspond to different systems in the task graph. Between each set of colors, you’ll see these slight vertical white lines.

Those are synchronization points in the frame where we’re switching between workloads, but you can see it’s under 1% of the overall work. We’re actually pretty much fully utilizing the GPU. In this case, this is shown on an RTX 4090 simulating, I think, 4,000 worlds of that hide-and-seek environment. For the initial version of this engine, we built a small set of little baseline environments. There was the hide-and-seek environment, there was overcooked, which is a cooperative reinforcement learning environment. The reason we chose these in particular was they already had easy-to-set-up CPU baselines, so we could actually do an apples-to-apples comparison. What we see is the green, which you can barely see on the top of these, is those original CPU reference implementations from the machine learning code bases.

And basically the performance is abysmal. All of these tasks were majorly bottlenecked by environment throughput. In blue, it’s what you get if you take the ECS ideas that I just told you about and just run it on the CPU, more like a traditional game engine, but scaled up to run across multiple worlds. And then red is you actually run it all on a 4090 on the GPU. And it’s literally over a hundred times faster in many cases. We actually did multiple projects over the course of my research that showed you can extend this to end-to-end training. It’s not only simulation throughput. We built a bunch of fast different environments on this engine over about two years. We had a lot of success with relatively non-technical people, like machine learning researchers who knew nothing about low-level GPU programming, coming in, building an environment, getting really good performance, and accelerating their workloads.

I want to leave you with a more general comment here, which is I think there’s still a lot of work to be done on high-level scripting languages for the GPU, even in the era of LLMs writing code and stuff like this. I see a lot of programming language systems trending towards, okay, it’s still the CUDA programming model, but it’s Python syntax on top, which is great, simplifies things, easier to look at, but it doesn’t actually help you with dynamic memory allocation on the GPU, super irregular parallelism, all these really fundamental issues. GPUs obviously have massive memory bandwidth and compute, even if you completely ignore the tensor cores. I think there are a lot of interesting workloads where you can leverage this plentiful, super fast hardware these days and get really good speedups. But some of these fundamental issues just mean there’s a lot of friction if you’re trying to do this on a workload that no one’s ever done before, like we were doing.

I think there’s generally an opportunity here for abstractions that have good enough default performance. You just utilize that raw GPU horsepower, but make GPU programming far, far simpler. Kind of like a real scripting language, Python-esque for the GPU. Thank you.

Francois: Thank you, Brennan.

Brennan: Yep.

Francois: Okay. So just to wrap up here, what’d you think? Good? Thumbs up. Okay. So I think these thematic clusters are the way we’re going to go through this in the future. The next one will be on robotics. We have a couple cool robotics people, a lot of people in the industry. If you have folks that you think would be a good fit, please either DM me in the Slack or otherwise. All right, stick around. We have another busy hour to hang out and meet each other. If you want to talk to the speakers, you want to talk to anyone, now’s the time. Thank you so much for coming. All right.