Can LLMs Write Fast Multi-GPU Kernels? — Simran Arora, Together AI

From single GPU kernels to the network between them

00:12Hi everyone. Uh, sorry it's a bit loud in here. Was not expecting this. Um, I'm Siman. I'm a principal scientist at Together AI. Um, I previously did my PhD in the Hazy Research Lab with Chris Ray at Stanford and I'm an incoming professor at Caltech.

00:27Um I lead the frontier performance research team at together where we develop systems, frameworks and algorithms to extract as much performance as possible out of modern um AI hardware. Today I want to share a little bit about our contributions towards simplifying the development of uh multi-GPU AI kernels.

00:51A few years ago, um, GPU utilization used to be limited by poor intraGPU memory access and single GPU kernels. But with significant investment in better kernels like flash attention, uh, memory efficient architectures like from deepseek, sparse attentions, mambas and so on, um, and better DSLs, we've sort of shifted the bottleneck to multi-GPU communication.

01:21During this talk, I'll start by telling you a little bit about why now, why GPU networking now. Then I'll tell you about

Inside an H100, and why memory distance matters

01:29um the sort of problem space. So what are the challenges in maximizing hardware utilization and development simplicity for multiGPU kernels. Um three, we'll talk a little bit about the fundamentals be behind designing effective multiGGPU kernels.

01:38Four, we'll look at whether frontier AI models can uh leverage these fundamental principles. Do they understand them? Can they reason about them? Um, you know, in theory, these models are very good at reasoning. Um, and then five, we'll talk through the results of these frontier models on a benchmark that we've developed called um parallel kernel bench for multiGPU kernel generation evaluation.

02:12Okay, before we dive into those five parts, just basic preliminaries. So, this is an Nvidia GPU, uh, an H100 GPU that you can see on the screen. Um, I always like to help ground people in GPU kernels via looking at the hardware. So, these rainbow colored dots are processors where actual compute is happening.

02:30um all of the you know parallel threads are operating within one of those colored dots and there's typically you know 100 200 of them on modern AI GPUs. Um around those processors you can see um some of the memory that these processors retrieve data so large weights activations from.

02:55So these rectangles between the colored dots are an L2 cache slightly faster memory. um not a ter like crazy large amount of it. And then these black boxes are high bandwidth memory. So when you Nvidia Smi and see you know 80 gigabytes, 1008 gigabytes, whatever it is on your GPU, that's that memory.

03:19Um a GPU is operating a highly parallel program. So multiple threads are combined together in uh into larger coarser units. and we schedule these threads and and blocks onto these processors to perform our AI compute. Beyond the GPU, we'll have multiple GPUs and we'll also have, you know, CPUs um that have memory as well.

03:49Um so to perform computation um the memory that these threads use is going to be stored in a really fast register memory that's right next to the computation units. um simple physics if I am pulling data from very very close to my proc my compute unit it's really fast to get to it because it you know that data is right next to me but there's not a large radius and not a large volume of space that's close by to my process uh my my compute units and so I don't have very much of it so you can

04:22see that the fastest memory here the registers is sup is you know 130 terabytes per second on an H100 but we don't have very much on of it and as we go to the further away memory we have a lot more of it but it takes longer to reach it. Uh again simple physics.

04:40So in multiGPU systems in particular um there is a hierarchy of interconnects.

The interconnect hierarchy: PCIe, NVLink, NVSwitch

04:46So we will have something called PCIe um as the channel for CPUGGPU communications. We'll have multiGPU or multi- uh node communications over infiniband TCP and then in the we're going to focus mostly on the Nvidia sphere here. Um in uh the intraGPU regime we'll have NVLink um providing point-to-point connections between GPUs and the NV switch.

05:12NV switch connects all NVLink in endpoints into a non-blocking fabric for full GPUGGPU communication and NV switch is exciting because it also provides support for in network offdevice acceleration for oper like communication primitives like multiccast and reductions.

05:37Okay, so diving in with the preliminaries in mind. Why GPU networking now? So as I mentioned at the beginning, we've really put a lot of effort into making AIO uh more efficient over recent

Why GPU networking, and why now

05:52years. Um again, we have uh architectures that use less compute, less memory like Mamba or sparse attentions. We have algorithms that make AI more hardware affair aware like flash attention. We have tools to make it easy to map AI algorithms to the hardware like tileang, mojo, triton, gluon, thunder kittens.

06:08Um, and we have new techniques to overlap uh execution across many AI operators very tightly like mega kernels. And we also have tools to make it easy to run on multiple vendor and silicon platforms like thunder mittens for Apple silicon or or hipkittens for AMD and so on.

06:35Um, at this point we really believe that GPU networking offers many new and exciting opportunities for AI efficiency. Modern AI workloads are getting very big and require kernels that span multiple um, you know, GPUs. So on many production um distributed training and inference workloads, communication is increasingly consuming the majority of the runtime and yields low model flop utilization at scale.

07:08So the uh pace of innovation and diversity of approaches that different hardware providers are taking in their networking stacks is another reason why it's an exciting time to study um you know networking and communication. So we can see here um AMD hardware with uh what's called XGMI interconnects um providing pointto-point links between different GPUs in a scaleup domain.

07:31We can see here TPU and uh interconnect uh as well. So the TPU will use a 3D Taurus and also have optical wraparound links. So yet another diverse form of the topology and links. Um and then again for Nvidia we'll have the INV um switch which integrates um compute capabilities directly into the interconnect fabric like for in network reductions and then we'll have our NV link providing up to you know 900 gigabytes of unidirectional bandwidth between any two remote GPUs high bandwidth memory um on you know particular generation of Nvidia hardware

08:19um Beyond the diversity in the networking stacks, there's also a lot of evolution in um how AI workloads are

Disaggregated workloads and 576 GPU systems

08:27adapting to take advantage of this hardware and the diversity of types of um you know hardware that we're we're using simultaneously for one AI workload. So um KV cache memory in modern inference systems is going to beworked across GPU, CPU, disk and remote machines.

08:41um inference systems increasingly disagregate different steps of inference across different hardware backends. So you could run speculative decoding on some hardware, decode on different hardware, prefill on different hardware. Um hardware is also evolving to have larger scaleup domains than ever before with you know 72 GPUs and a scaleup domain in in the coming um chips and Nvidia planning on a single system in 2027 with 576 GPUs.

09:17Um at the same time to take advantage of these more uh intensive scaleup domains we're getting richer uh primitives for fine grained control and kernel writing over these domains. So we have something called tensor memory acceleration where we can provide uh perform asynchronous network transfers from the device side um on these GPUs.

09:40So all of these changes are opening up new opportunities and challenges in um both AI you know how do we build models that take advantage of these trends and in systems. Okay, so the problems that we're going to go after um how do we get peak hardware utilization and also development simplicity um for these multiGPU kernels.

10:06So um it's been very difficult to write multiGGPU kernels and there's a lot of there are a lot of papers a lot of systems reports that document you know challenges here. Um, one of the things here is it's compounded by the fact that communication hardware around GPUs has progressed a lot more slowly relative to compute and memory.

10:21Um, so comparing

Compute improved 7.2x, the network 3x

10:30NVIDIA A100's in 2020 to B200s in 2024, u BF16 tensor core speeds improved by 7.2x. um while intra node communication by just 3x and inter node communication by just 2x. Um and coming back to my points about how diverse networking is right now things like tensor cores that run map moles and our memory hierarchies are pretty consistent and resemble one another across diverse AI vendors and multisilicon.

10:57Um but again as I I mentioned the networking stack is something that is really different across vendors still um you know a first step as we went about all this work is to just study the baselines um that are out there. So um one of the popular tools for um

Where NCCL breaks down, and the roofline gap

11:21communications is this nickel library or Rickle on AMD um that both you know companies respectively spend a lot of you know engineering investment into releasing to make it easy for people to do multiGPU work. Um but they're not very flexible.

11:36So they're tuned for bulk transfers for large contiguous chunks of data transfers. And the design really breaks down when you care about peak performance, fine grain communication, um, and sort of non-trivial collectives that you want to fuse together.

11:54So as a result, you can achieve much higher performance by writing custom communication kernels that directly address these needs. If we look at a naive baseline that's representative of very popular libraries in machine learning stacking pietorch with nickel um we can we find that across um you know the many uh problems in our parallel kernel bench benchmark that the majority of these um simple baselines will fall below 50% of their communication aware roof line bound.

12:26So there's a lot of room for improvement here. Um the current frameworks beyond um nickel which is popular in like Megatron LM, Flex Flow, Nanoflow um you know all again these systems are primarily orchestrating bulk collectives via nickel um and require synchronization before and after data transfers.

12:50So beyond th those off-the-shelf libraries, we have compilers and DSLs that exist. So there's Triton distributed is one of them. Um and uh you know tile link is another one. Um we have found that it's

Compilers, DSLs, and hand tuned operators

13:07very difficult to support the rapid pace of networking improvements within these frameworks. So our benchmarks and our papers highlight results where Triton distributed originally tuned around 8 uh H800 GPUs fails to adapt efficiently to other architectures like H100s.

13:27And then the third category of how people can proceed here is to really handtune specific AI operators one by one. So there's a lot of popular work um DPP um comet ring attention um flux flashdoe um and then several distributed gem kernels from cutless and these methods achieve peak performance but often um they do not you know some of these methods have been designed in one precision and it takes five or six months to scale it to another precision and just the scalability of this hand tuning and um fine grain kernel writing is not very um effective.

The research question, and ParallelKittens

14:10So um with this landscape in mind, our research question was really about whether there is a small set of principles and fundamentals that really governs multiGPU kernel writing and whether there are methods that can leverage those principles if they exist to simplify the development of these kernels.

14:32Um I'll briefly highlight two works here that um govern like that that represent our approach. So first we think it's important to build our own fundamental understanding and to manually do the work to understand it rather than just throwing say an LLM at the problem.

14:46So we spent the time to build out parallel kittens which is a small set of minimal primitives and patterns for multiGPU kernels. Um we use this to understand the trade-offs of multiGPU kernels and to one write a large collection of um peak performance kernels for a variety of parallelism schemes.

15:09Um and this I will use to hopefully you know educate and bring us all on the same page on what patterns we figured out. Um and then once we found that there is indeed a small set of trade-offs governing this landscape, we were curious whether models, especially these models right now that claim to be very good at kernel writing and also reasoning um could reason about these trade-offs when we provide them in context to actually generate a bunch of net new multiGPU kernels for us.

15:38Unfortunately, we found they were not very good, but we'll dive into more of that at the end. Um so just the fundamental section this is going to be more you know educational what are the trade-offs that go into these kernels. So there are three main ways to do um intraGPU data transfers.

15:56Um there's the per GPU what's called copy engine and this is host or CPU initiated work. It's really good for large message transfers.

Transfer mechanisms: copy engine, TMA, register level

16:12So when your message size, the amount of data being transferred is really big um and it can get to sort of like peak bandwidth um on on the communication side. In contrast, you can use device initiated or GPU initiated transfers via that tensor memory accelerator that I mentioned or via register level instructions called uh in sort of their PTX lingo like LDST red multime.

16:35Um, and the uh T TMA is really nice. These device initiated ones are really nice because they can saturate our NVLink bandwidth using relatively small message sizes. And this means that they can be really nice when we're trying to do fine grain communication rather than sending bulk amounts of data over the links all at once coarsely.

17:02Um, there are some trade-offs here. So the copy engine is really nice because it doesn't take away or you know waste a lot of our precious registers that I mentioned are important for compute on the GPU. Um and it doesn't also use any of those rainbow colored dots the processors on our GPU allowing us to repurpose those for memory or um computation on our uh you know other parts of the AI pipeline.

17:27Um, TMA, the second option here, consumes very few registers, which is why it's nice. Um, and it also can achieve high utilization using very few of our processors. So, it's a nice useful tool for fine grain overlapping. Um, TMA does have limitations.

17:44It can't effectively take advantage of these in network um, computations that I mentioned are feasible with technologies like NV switch. And the register level instructions are really nice for being able to take advantage of um you know those those sort of in network reductions that NV switch offers.

18:08So again there are different tradeoffs different functionalities that these transfer mechanisms offer and they face different trade-offs. Second um beyond transfer mechanism the trade-off is around how to overlap compute memory and uh communication in

Overlapping compute and communication across SMs

18:28GPU kernels. So there's two main categories of schedules. The first is intraSM within one of those rainbow dots um where we'll we'll have different warps or threads within that processor specialized to handle either compute or one specialized for communication concurrently.

18:42Um we can dedicate you know different warps to each of these. The challenge with this intm overlapping is that the communication and computation pattern really need to like align and jive with one another. they need to use the same data as inputs for the computation and communication.

18:59When they don't align, you could use something like interm um schedules that are shown on your right here where we'll now have each of the different rainbow colored dots on our GPUs, those different processors specialized to compute communication and memory.

19:21Um and so this this is nice when the colonel would others wise need to split across resources like the register file or shared memory um across these different steps in misaligned ways. Um this is also really nice when it's hard to maximize NVLink traversal with intram overlapping.

19:50So I just wanted to highlight one quick example here where each of the patterns excels in popular you know AI uh uh like kind of patterns that you'll see. So on a gem plus uh reduced uh scatter here we can see that the um intraSM overlapping scheduleuler schedule is very effective in the gem plus all reduce we can see that the interSM which again leverages the in network reductions of envy switch is very effective.

20:19So we face these trade-offs and you can read more about um the design decisions that go into them in our parallel kittens paper. Um, and then finally, ideally, abstraction should allow the developer flexibility to control how they're buffering and synchronizing between data senders and receivers.

20:41So, we encapsulated these ideas into parallel kittens. Again, a simple set of programming primitives and templates for these multiGPU kernels. Parallel Kittens is used in production at Together AI as well as our partner um you know Cursor and other uh companies in the AI space.

In production at Together AI and Cursor

21:01Um here's some sample code. I won't spend too much time here, but we usually add roughly a dozen lines of code over a single GPU kernel to insert these multiGPU primitives. Um, and you can see here across data sequence and expert parallelism how our parallel kittens kernels are achieving state-of-the-art results um, compared to strong reference baselines.

21:22And you can check out our repo to learn more. Okay, so we understand a little bit about um the trade-offs that underly these multiGPU kernels and there's just a couple main, you know, ones that exist. So can models reason through them um and give us these uh you know kernels models are getting better at reasoning today.

21:47Um do they generalize well to these problems or are we benchmaxed on you know benchmarks of the past which are more single GPU centric uh models right now are showing really promising results on single GPU uh benchmarks. So it's a ripe time to to extend it.

22:07In our benchmark parallel kernel bench, each task presents the model with an unoptimized reference implementation

Inside ParallelKernelBench: 87 problems

22:14written in PyTorch with torch distributed um nickel operations and then a system topology that specifies the number of ranks and intraode hardware configuration. And the model needs to rewrite the reference into a performance CUDA kernel that uh uses unified virtual addressing.

22:37The um multiGPU problem space expands combinatorally beyond single GPU um cases. So a standard transformer layer can be parallelized across data sequence, tensor, context, layer, pipeline and expert dimensions. and each composition um induces a different communication pass uh pattern.

22:51So to make sure that our benchmark has high coverage over the representative types of multiGPU problems um we created this taxonomy um that you can read more about in our paper and then picked representative problems for each part of the taxonomy.

23:15Um these are all patterns that arise in real AI workloads from inference to RL to post-training. There are 87 problems overall um drawn from GitHub repositories that we found to be very informative. Um and we uh and like optimized library implementations and DSLs that people have written multiGPU kernels in.

23:40Um, we wanted to really make sure that solving PK this parallel kernel bench would lead to net new useful production kernels rather than artificial or useless kernels. Okay, so that's parallel kernel bench. How do models perform? So we measured to sorry that this is a bit small.

23:58We measured two main metrics. Um pass at K which is the number of correct kernels generated after K attempts and then fast um one at K which counts solutions that are both correct and outperform the speed of the pi torch plus nickel baseline.

24:21Um so pass K just correctness fast one at K is whether you're getting a 1x or higher speed up over the reference. So performanceoriented um we found that in the zeroot setting the best of the frontier models we tried

Zero shot results, and where scaling stalls

24:37solves 28 out of 87 problems and 22 of those problems are faster than the pietorch plus nickel baseline. If we make multiple samples, you know, standard uh scaling up test time compute, we can uh get that number from say like to to 36 correct solutions, but the fast uh one performance still plateaus out at roughly 31%.

25:01So we don't see much room from continuing to scale there as we increase the number of parallel generations. Um we find that the correct once correctness is established um speedups naturally come from eliminating nickel staging overhead in favor of direct NVLink loads in stores.

25:24Um the success patterns here are really concentrated into familiar patterns. So collective primitives, tensor parallel gems and Ulyses style context parallelism. So in other words, patterns that we see heavily represented um on the internet rather than necessarily patterns that the model has used its reasoning abilities to think through.

25:50Okay, even the best available model that we benchmarked here GPT 5.5 drops off very quickly as the speed up threshold increases. So on the x-axis here, we're increasing the speed up threshold over that pietor torch plus nickel baseline. And then we're showing the number of correct kernels that are faster than that baseline or this much faster than the baseline on the y-axis.

26:16So GPT 5.5 is this orange line here and then DeepSeek V4 Pro is the aqua line at the bottom.

The failures are not CUDA syntax

26:28um we found that there's deeper issues than CUDA syntax. So we found that if you do multiple sampling or have the model kind of look at its errors and correct them, it can often compile the kernels. But the models really struggle to reason through the tradeoffs that we talked about in the prior section.

26:44um collective ordering, data partitioning, thinking about intra versus interm scheduling or deciding between the different transfer mechanisms. Um we find that they often do not use things like the register transfer instructions or tensor memory acceleration when writing the kernels.

An agent harness, and what is still missing

27:07Um we wanted to try a pretty simple instantiation of something like a clawed code um coding agent. So, we took uh the mini sui agent multi-turn harness and one of the best performing models, Gemini 3 Pro, and gave it access to a local bash environment to sort of mimic the standard claude code setup.

27:24We found that this could help the agent um go from solving 24 problems to 35 of 87 problems um with um 26 achieving over a 1x speed up over the reference. But we found that as we scaled the amount of time um the performance plateaued as uh and we find that additional techniques would be required to continue seeing the scaling there.

27:50Again these results are discussed in more detail in our paper. Um we think this is a really exciting you know just to wrap up here we hope that people out here can use both parallel kittens and parallel kernel bench. We think that the kernels generated from solving parallel kernel bench will lead to net new production kernels that are important bottlenecks for inference in RL right now.

28:16um they're you know we tried our best to make them you know non-artificial and we can already see signs of life and exciting results where people have not invested a bunch of time to handw write a multiGPU kernel and we've gotten some net new interesting ones like this Nemo vocab parallel um you know filtering kernel um a hyena architecture context parallelism kernel and uh the SAM 3 video segmentation model um IOU suppression kernel.

28:52Just to conclude here um we're really excited about um how uh well just talking about the lessons. First off, we think there aren't that many patterns that are involved in writing intragpu effective kernels again encapsulated by our small set of programming primitives.

29:11Um, but unfortunately models do not currently understand how to reason through these trade-offs even when we provide them in context. Um, we're really excited about methods that can help attack this benchmark. We're excited about architectures that can grow with the trends of how networking stacks are evolving.

29:27Um, you know, larger scale up domains shift away from scale out um and massive onchip memory structures. And we hope that uh you know we can also extend and you you can feel free to reach out to me at my email. Thanks.