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

AI Engineer · 30 min · 215 sentences · from YouTube's caption track

Each timecode opens YouTube at the start of that sentence. Line anchors (#s42) are the cue ids in the WebVTT, and every line carries its start and end seconds. All transcripts has every talk, and the whole corpus as one file.

00:01[music]
00:12Hi everyone.
00:13Uh, sorry it's a bit loud in here.
00:15Was not expecting this.
00:17Um, I'm Siman.
00:18I'm a principal scientist at Together AI.
00:21Um, I previously did my PhD in the Hazy Research Lab with Chris Ray at Stanford and I'm an incoming professor at Caltech.
00:29Um 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.
00:42Today 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.
01:01But 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
01:15to multi-GPU communication.
01:21During this talk, I'll start by telling you a little bit about why now, why GPU networking now.
01:27Then I'll tell you about um the sort of problem space.
01:30So what are the challenges in maximizing hardware utilization and development simplicity for multiGPU kernels.
01:38Um three, we'll talk a little bit about the fundamentals be behind designing effective multiGGPU kernels.
01:45Four, we'll look at whether frontier AI models can uh leverage these fundamental principles.
01:52Do they understand them?
01:53Can they reason about them?
01:55Um, you know, in theory, these models are very good at reasoning.
01:58Um, 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.
02:16So, this is an Nvidia GPU, uh, an H100 GPU that you can see on the screen.
02:23Um, I always like to help ground people in GPU kernels via looking at the hardware.
02:29So, these rainbow colored dots are processors where actual compute is happening.
02:35um 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.
02:47Um around those processors you can see um some of the memory that these processors retrieve data so large weights activations from.
02:57So these rectangles between the colored dots are an L2 cache slightly faster memory.
03:04um not a ter like crazy large amount of it.
03:07And then these black boxes are high bandwidth memory.
03:10So 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.
03:24So multiple threads are combined together in uh into larger coarser units.
03:30and we schedule these threads and and blocks onto these processors to perform our AI compute.
03:38Beyond 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.
04:01um 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
04:17process uh my my compute units and so I don't have very much of it so you can see that the fastest memory here the registers is sup is you know 130 terabytes per second on an H100
04:29but 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.
04:37Uh again simple physics.
04:40So in multiGPU systems in particular um there is a hierarchy of interconnects.
04:46So we will have something called PCIe um as the channel for CPUGGPU communications.
04:53We'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.
05:04Um in uh the intraGPU regime we'll have NVLink um providing point-to-point connections between GPUs and the NV switch.
05:14NV 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
05:29for oper like communication primitives like multiccast and reductions.
05:37Okay, so diving in with the preliminaries in mind.
05:41Why GPU networking now?
05:44So as I mentioned at the beginning, we've really put a lot of effort into making AIO uh more efficient over recent years.
05:53Um again, we have uh architectures that use less compute, less memory like Mamba or sparse attentions.
06:00We have algorithms that make AI more hardware affair aware like flash attention.
06:06We have tools to make it easy to map AI algorithms to the hardware like tileang, mojo, triton, gluon, thunder kittens.
06:15Um, and we have new techniques to overlap uh execution across many AI operators very tightly like mega kernels.
06:23And 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.
06:45Modern AI workloads are getting very big and require kernels that span multiple um, you know, GPUs.
06:52So 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
07:19um you know networking and communication.
07:22So 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:35We can see here TPU and uh interconnect uh as well.
07:39So the TPU will use a 3D Taurus and also have optical wraparound links.
07:45So yet another diverse form of the topology and links.
07:49Um 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
08:07bandwidth between any two remote GPUs high bandwidth memory um on you know particular generation of Nvidia hardware um Beyond the diversity in the networking stacks, there's also a lot of evolution
08:24in um how AI workloads are adapting 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.
08:37So um KV cache memory in modern inference systems is going to beworked across GPU, CPU, disk and remote machines.
08:46um inference systems increasingly disagregate different steps of inference across different hardware backends.
08:53So you could run speculative decoding on some hardware, decode on different hardware, prefill on different hardware.
09:00Um 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
09:13with 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.
09:28So 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.
09:53Okay, 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.
10:17Um, 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:29Um, so comparing NVIDIA A100's in 2020 to B200s in 2024, u BF16 tensor core speeds improved by 7.2x.
10:40um while intra node communication by just 3x and inter node communication by just 2x.
10:47Um 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.
11:02Um 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
11:15um that are out there.
11:16So um one of the popular tools for um communications is this nickel library or Rickle on AMD um that both you know companies respectively spend a lot of you know engineering investment into
11:31releasing to make it easy for people to do multiGPU work.
11:35Um but they're not very flexible.
11:37So they're tuned for bulk transfers for large contiguous chunks of data transfers.
11:43And 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.
12:03If 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
12:16uh 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:30So there's a lot of room for improvement here.
12:33Um 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
12:48um and require synchronization before and after data transfers.
12:52So beyond th those off-the-shelf libraries, we have compilers and DSLs that exist.
12:58So there's Triton distributed is one of them.
13:01Um and uh you know tile link is another one.
13:05Um we have found that it's very difficult to support the rapid pace of networking improvements within these frameworks.
13:13So 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.
13:36So 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
13:52but 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
14:02and just the scalability of this hand tuning and um fine grain kernel writing is not very um effective.
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
14:27to simplify the development of these kernels.
14:32Um I'll briefly highlight two works here that um govern like that that represent our approach.
14:38So 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:49So we spent the time to build out parallel kittens which is a small set of minimal primitives and patterns for multiGPU kernels.
14:58Um 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:11Um and this I will use to hopefully you know educate and bring us all on the same page on what patterns we figured out.
15:18Um 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
15:33um 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:42Unfortunately, we found they were not very good, but we'll dive into more of that at the end.
15:48Um so just the fundamental section this is going to be more you know educational what are the trade-offs that go into these kernels.
15:56So there are three main ways to do um intraGPU data transfers.
16:01Um there's the per GPU what's called copy engine and this is host or CPU initiated work.
16:09It's really good for large message transfers.
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.
16:23In 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
16:37red multime.
16:39Um, and the uh T TMA is really nice.
16:43These device initiated ones are really nice because they can saturate our NVLink bandwidth using relatively small message sizes.
16:52And 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.
17:04So 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.
17:15Um 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:29Um, TMA, the second option here, consumes very few registers, which is why it's nice.
17:35Um, and it also can achieve high utilization using very few of our processors.
17:41So, it's a nice useful tool for fine grain overlapping.
17:45Um, TMA does have limitations.
17:48It can't effectively take advantage of these in network um, computations that I mentioned are feasible with technologies like NV switch.
17:56And 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.
18:18Second um beyond transfer mechanism the trade-off is around how to overlap compute memory and uh communication in GPU kernels.
18:29So there's two main categories of schedules.
18:33The 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
18:47concurrently.
18:48Um we can dedicate you know different warps to each of these.
18:52The challenge with this intm overlapping is that the communication and computation pattern really need to like align and jive with one another.
19:01they need to use the same data as inputs for the computation and communication.
19:06When 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
19:21compute communication and memory.
19:25Um 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
19:38ways.
19:41Um 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.
20:00So 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
20:16which again leverages the in network reductions of envy switch is very effective.
20:21So we face these trade-offs and you can read more about um the design decisions that go into them in our parallel kittens paper.
20:30Um, 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.
20:45Again, a simple set of programming primitives and templates for these multiGPU kernels.
20:51Parallel 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.
21:01Um here's some sample code.
21:02I 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.
21:13Um, 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:26And you can check out our repo to learn more.
21:30Okay, 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.
21:40So can models reason through them um and give us these uh you know kernels models are getting better at reasoning today.
21:48Um 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
22:02uh benchmarks.
22:03So 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 written in PyTorch with torch distributed um nickel operations and then a system topology that specifies the number of ranks and intraode hardware configuration.
22:25And 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.
22:44So a standard transformer layer can be parallelized across data sequence, tensor, context, layer, pipeline and expert dimensions.
22:53and each composition um induces a different communication pass uh pattern.
22:58So 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
23:11problems 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.
23:23There are 87 problems overall um drawn from GitHub repositories that we found to be very informative.
23:31Um 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.
23:51Okay, so that's parallel kernel bench.
23:54How do models perform?
23:58So we measured to sorry that this is a bit small.
24:01We measured two main metrics.
24:03Um 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
24:17the 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.
24:29So performanceoriented um we found that in the zeroot setting the best of the frontier models we tried solves 28 out of 87 problems and 22 of those problems are faster than the pietorch
24:43plus nickel baseline.
24:45If 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
24:59uh one performance still plateaus out at roughly 31%.
25:04So we don't see much room from continuing to scale there as we increase the number of parallel generations.
25:12Um 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.
25:30So collective primitives, tensor parallel gems and Ulyses style context parallelism.
25:35So 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.
26:00So on the x-axis here, we're increasing the speed up threshold over that pietor torch plus nickel baseline.
26:07And 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.
26:28um we found that there's deeper issues than CUDA syntax.
26:32So 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.
26:42But the models really struggle to reason through the tradeoffs that we talked about in the prior section.
26:48um collective ordering, data partitioning, thinking about intra versus interm scheduling or deciding between the different transfer mechanisms.
26:57Um we find that they often do not use things like the register transfer instructions or tensor memory acceleration when writing the kernels.
27:07Um we wanted to try a pretty simple instantiation of something like a clawed code um coding agent.
27:14So, 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:29We 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.
27:44But 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:56Again these results are discussed in more detail in our paper.
28:00Um 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.
28:08We 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:19um 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
28:30kernel 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
28:48model um IOU suppression kernel.
28:52Just to conclude here um we're really excited about um how uh well just talking about the lessons.
29:00First 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.
29:18Um, we're really excited about methods that can help attack this benchmark.
29:22We're excited about architectures that can grow with the trends of how networking stacks are evolving.
29:29Um, you know, larger scale up domains shift away from scale out um and massive onchip memory structures.
29:36And we hope that uh you know we can also extend and you you can feel free to reach out to me at my email.
29:42Thanks.