Parameterized CUDA Graph Launch in PyTorch: CUDA Graphs Without the Pain - Daniel Galvez, NVIDIA

In this talk, Daniel Galvez discusses the enhancement of CUDA APIs within PyTorch, focusing specifically on improving the developer experience with CUDA graphs. He explains what CUDA graphs are, emphasizing their role in performance optimization by eliminating CPU overhead during GPU operations. The speaker articulates the complexities and challenges associated with using CUDA graphs, along with the introduction of parameterized CUDA graph launch as a solution to these challenges. Through detailed examples, he illustrates how this new feature allows for dynamic input parameters while launching CUDA graphs, thereby addressing issues like memory usage and the potential for silent errors in large codebases. Galvez also delves into the implementation difficulties, especially in distributed environments, and shares insights into ongoing improvements to the CUDA graph API within PyTorch to better support developers.

For those who don't know me, I'm Daniel Galvez. I've been working on I've been working on a just improving CUDA APIs in PyTorch for I don't know, a year or two now. And I'm very particular about, you know, achieving good performance while also having good developer experience. And this uh what I'm going to talk to you about is basically something focused on good as opposed to performance.

And that's going to be about you making CUDA graphs easier to use. Uh one moment. Have we started recording? Excuse me. All right, thank you. Uh so, yeah. Uh so, let's just get going. Uh so, the agenda is basically uh I'm going to describe what CUDA graphs are. Uh thanks for coming. What the benefits of CUDA graphs are. The answer is performance. What are the challenges of

CUDA graphs? This is like the important part. Uh how parameterized CUDA graph launch helps. And how this feature is implemented and why it's hard. And the last part is because uh you know, I was wondering like, "Ah, how should I give this talk?" And because I made the RFC for this, I was like, "Ah, you know, my feature is almost done. It's going to get in." I

did actually find some issues that uh make this not work in a few cases. And, you know, there's a lot of experts here, so figured I'll just talk through those. So, anyway, uh what are CUDA graphs? CUDA graph is basically a DAG. It's a DAG of CUDA work to do. each node is a CUDA is a CUDA kernel, a mem copy, a mem set. It's basically always

a CUDA kernel, though. Always. like 90% of the time these nodes are going to be CUDA kernels. And basically, you have edges declaring like, "Okay, like in the example you see there, uh A must precede B and X like Y X can only after X is finished." Uh pretty simple. Uh and the way you make these is via something called graph capture. Uh graph capture is a

little a difficult to describe, but basically, you mark a CUDA stream as capturing and once you do that, you don't actually execute any CUDA stream anymore until you call CUDA stream end capture, but instead you just lazily enqueue all the CUDA kernels and CUDA work into this graph essentially. And you will replay that via an API API called CUDA graph launch afterwards. So, that's why it's called

a form of lazy execution. the key thing you need to know about this, kernel arguments used at graph capture time are reused at replay time. Like they cannot be Normally, they cannot be changed and they cannot be introspected. What we're going to talk about is how you do that today. Um so, like you see a CUDA graph launch that API right there, like there's a CUDA graph

exact T, CUDA stream T but there's no extra argument, which is like, "Oh, yeah, for the ith argument of like the jth Can you shut that, please?" Uh for the ith argument of the jth node like change the that 8-byte value to a different 8-byte value. Like there there that's not in that CUDA graph launch API, right? So, anyway, what's the in terms of benefit, right? The

benefit is performance, period. Like that's that's the only reason why people care about this API. Like basically uh in the era of Blackwell and Hopper GPUs like what happened was like CP like everyone knows, okay, Moore's law is dead, parallel blah What kind of happened is that uh CPUs like GPUs caught the CPUs in terms of like, you know, like it's not uncommon these days have a

CUDA kernel that executes for say I don't know, let's say 5 microseconds. Um and you know, GPUs are always really doubling in speed every generation for in speed. Um well, like it's not uncommon for like your CPU code to take more than 5 microseconds to launch the next CUDA kernel. Sometimes much longer depending on weird edge cases. So, what ends up happening is that even if you

speed up your GPU code, you still end up with GPU bubbles where you just get a In fact, if you speed up your GPU code in a situation like this at the top, you actually don't get a speed up from speeding up your GPU code, right? Because that just makes the GPU bubble bigger. CUDA graphs eliminate all CPU overhead because essentially like when you launch a CUDA

graph, the GPU itself will automa- automa- all of the nodes in the order specified by the graph. And there's just no more CPU overhead. Uh I'll since Uh yeah, I'll just mention uh you know, a lot of people think like, "Oh, why is the CPU so slow?" A lot of people think, "Oh, Python's slow." Uh you know, I'll just add some nuance for people who might wonder

about why exactly the CPU's There's like basically three factors. There's the CUDA driver, there is Python, and there is uh torches uh 810 dispatch. I would say just ballpark each of those accounts for about 1/3 of the time. Uh so, if you want to speed up the CPU, you can, but you really need to be You need to focus on all three prongs. But, we're not talking

about that today. Uh we're just talking about how we just get rid of all that junk. Uh So, yeah, uh just give you a sense of the uh quantitative performance benefits of CUDA graphs. Uh what you can see is that it's like there are situations where like you're a full non-autistic. Like, this is like the ML on training results. One of my colleagues ran with and without

CUDA graphs uh for a bunch of the workloads in MLPerf. Uh basically, some of these are like three times faster with CUDA graphs. And some of these are even like 512 GPUs. And like your boss is very upset at you if you're using 512 GPUs and you're not Right? That just that is that is bad. Uh so basically you need them for many workloads these days. And

uh you might ask, "Oh, what workloads do I need these for?" Uh basically I would say highly distributed workloads because you don't want to get hit by a tail worker effect or two say one CPU uh core being slower than everything. Garbage collection or a bad scheduler or something weird, who knows. and uh I would say anything where you have a small amount of work. So generally

smaller models sending more affected uh by uh high CPU overhead and thus more amenable to CUDA graphs. Uh actually uh so I'm kind of curious further room the CUDA graph API in PyTorch. All right, so I got some I guess we'll uh go through this then. Uh so basically uh for those who don't know, this is how you use CUDA graphs today in PyTorch. Uh there's another

way which is via torch.compile, but uh I'm not I'm not I'm going to allude to that, but I'm not going to talk too much about that. That's called CUDA graph trees. I think those who might be here. Uh maybe not in this room, but at the at the conference. Um so anyway, uh the left-hand side's just a training function, forward pass, loss function, backward pass, optimizer, whatever.

Uh the right-hand side's the Uh oops. So yeah, the first part you'll see is that uh basically and for the very first iteration you do it something called a warm-up where you just execute the model as it is, like the normal way, the eager way in PyTorch. Uh you would just, you know, call the train function, get the loss, print the loss, whatever. Uh the warm-up is

important for reasons I just don't have time to talk about today, sorry. Um so then the next capture a CUDA graph here this CUDA graph capture. So if you look at that context manager with torch.cuda.graph, uh that basically sets the current stream in PyTorch, well, it makes a new stream, which it then sets to be in stream capture mode, and it actually captures whatever work you put

on that stream into the graph G. there's an X static and a Y static. X static and Y static are uh basically what people call static input buffers. The thing you know about CUDA graphs is whatever inputs you use at graph capture time, whatever tensors use at graph capture time, those are fixed. I've like already said this before but it's like message. Those are fixed. Uh so

you know, X static and Y static, we'll get into that get to those in the next part. And also the output is fixed, too. Like that's why I call it loss underscore So anyway, um yeah, so after you create the CUDA graph, you see there's one less It's a little hard to see. Uh for all other iterations, you just use the CUDA graph and you don't just

call replay, that's what does CUDA graph launch. You actually have to copy into your static buffers, X static and Y static, X and Y. X and Y, you know, like if this is like image recognition, like X is like the image, Y is the label, You need to do that first. So these input these what I'm calling these static input buffer copies is something that it we're

going to talk about a little bit later. Uh so that's like how people do today. I have in the commented section how you would run with parameterized graph launch. It's not too important and I think maybe I'll not spend too much spend too much time on it but basically specify dynamic inputs, in this case X and Y, new values for them. Uh and you get a also

a quote unquote dynamic loss. I I'm calling that dynamic because they're no longer fixed, like loss underscore static is like the same tensor every time you replay, like the same exact backing memory address. Dynamic loss though, every time you'd replay in the dynamic dynamic uh parameterized CUDA graph launch, it's going to be at a different memory address. Unless by chance the caching allocator allocates an address which

is exactly the same, which can happen, but you know, that's not the point. Uh so uh why CUDA graphs hard to use? So, basically I mentioned it's like CUDA graphs are like and it's very hard to describe why CUDA graphs are hard to use. There's like so many things that can go wrong. I made an iceberg meme because a friend thought it would be funny and I

caved. Don't don't don't bother looking at that too much. Maybe if you're like an expert, which I know some people in the room are, maybe you'll be entertained by it, but Yeah, so why are they hard to use? And it's in the subtitle. Basically, a well-formed PyTorch program may not run properly when you use CUDA So, you can have a perfectly valid PyTorch program that does what

you want without CUDA graphs and when you use CUDA graphs, it will not work. So, there's several reasons for this. on four issues today that parameterized CUDA graph launch, the new API I'm talking about today can kind of help deal with. basically, there's kind of four issues that I've kind of already alluded to. First of all, you have to copy your inputs into the static input memory

addresses you use I already talked about that. That's actually isn't too bad depending upon things. It can get really difficult though when you have what people like to call piecewise CUDA graph where you basically mix CUDA graph and non-CUDA graph workloads in the same in the same workload, I guess. Bad bad bad name, workload. Uh They're really tricky because like essentially, you might It is the user's

responsibility to make sure that your to make sure you do all the appropriate copies into those static input memory addresses. It is your job to know what the inputs to a given CUDA graph are. trust me, people Second thing, a replay CUDA graph always writes outputs to the same memory address. So, let's say you uh like two replays of the same graph and you have two outputs,

out one and out two. Out one will have the same value as out two once the second replay finishes because they're actually of each other. It's the same exact memory buffers. So, what can happen is that in large scale code bases you will see that uh essentially like well, you won't see. It's going to be a silent error. What will happen to you is that if you're

using an output memory address of a CUDA graph uh but I run that value is going to be overwritten and you never want to bug a bug like this. This is like super advanced ninja territory. You don't want to be there. Uh and CUDA graphs inflate memory usage is probably the most important one. I don't know why it's number three. This like should be like the top

one. Uh CUDA graphs inflate memory usage. I'll see you soon. But, I'd say it's the top most important issue. Um and there's no way to CUDA graph nodes inputs and output tensors. I kind of alluded to this in like one and two. It's your responsibility to keep track of that. So, parameterized CUDA graphs uh can kind of tackle one, two, three, and four. As I mentioned, there's

other actually spoke about a few weeks ago BTC. If you're curious, that's the QR code. yeah, okay. So, basically, you know, I've been talking about this parameterized CUDA graph one thing for a while. How does it work? Like, what does this dude talk What is he even talking about? I don't know. Uh so, basically, uh the way it works is like conceptually super simple. torch.add You actually

like the internal torch.add op is going to run a CUDA kernel at some point. It's going to look something like this. This is not the code, by the way. I I looked up the code. It's like some C++ templates. You all would be confused. It confused me. So, this is like a basic like straw man uh demonstration of what this uh kernel might look like. Uh it's

pretty simple. Like, basically to update the parameters of a CUDA graph, you would what what you would need to do mechanic mechanistically is update the current the parameters of each individual kernel inside that CUDA We have a CUDA graph which just one kernel in it, just this add arrays, And we you know you got it from this torch.add function. Like, all we need to do is say,

"Okay, for this input, like we know that like we know that input goes to float star A, other goes to float star B, and whatever the output of torch.add is is going to be in float star Uh, so that's really all there all there is to it. Like, conceptually that's what what you're doing. You're like making these associations of like, "Okay, here are my tensors, here are

the memory related memory buffers." Uh, and you can patch these with a a few APIs. The but the best ones CUDA graph kernel update supply. It's the most Uh, I'll mention you cannot update non-pointer arguments. Uh, I don't have time to describe why, but basically we only care about updating Um, so if you really want to update non-pointer arguments, you're out of luck, basically. Um, so yeah.

Uh, so that conceptually is what happens. Kind of difficult to implement this, and we'll see how much time I have to spend how much time I have to spend on that. But, uh, basically how does parameterized CUDA graph help uh, CUDA graph launch help with uh, the issues I mentioned? Well, basically the key thing is that, the typical typical way you run CUDA graphs, there's a private

memory CUDA graph which is distinct from the default memory pool in PyTorch. Uh, what you can do instead, if you just rewrite the freaking param the pointer parameters every time before you launch a CUDA graph, is you just allocate the memory you need from the and replace those inside a private pool, you're essentially borrowing memory uh from the default memory pool and this reduces the external fragmentation

due to uh having a having a private memory pool. This is the biggest most important one because VLM, SG line, a lot of Megatron, a lot state-of-the-art frameworks today uh they're in a situation where they have no choice but to do partial CUDA graphs where like some of the workload uses CUDA graphs but some can't due to basically because they're not CUDA graph compatible and you know,

what makes a workload not CUDA graph compatible is not something I have time to talk about. Essentially, like today, like some like you just need to accept that some of your operation are going to be CUDA graph compatible. And what can happen is that when you adopt CUDA graphs, you realize, oh, I'm getting an out of memory error. That basically prevents you from running with CUDA graphs

at all and like if that making you lose like a 3x performance boost and that's un- Um and there's like a few other things, but that's really the important one. I'll mention uh Yeah, there's no need to copy inputs into your static memory buffers anymore, So, this allows some cool things. You divide you save device memory copies or it's going to like change those static input buffers

to whatever the dynamic input buffer you want to use right now Uh so, that allows input you know, you can't do input mutations uh today in CUDA graphs, funnily enough. Oh, you can sometimes for the static inputs. It's not important. Um So, there's no chance of overwriting previous CUDA graph replace output uh and uh you can so maybe I won't spend too much time on this, but

it's theoretically the CUDA graph for fast restarts is something VLM and like anyone interested in like a lot of good put so, is interested in because restarting your whole training process takes a long a a long time And uh there is ways that that could work, but also ways that I couldn't, but I think I'll skip that for time. Kind of described uh like the the all

the agenda items except the last one. Uh you all can feel free to tune out at this point. Um so basically why is this CUDA graph launch feature hard to implement? So like for perspective, I got this and uh not Torch Titan, sorry. Not all of Torch Titan, no no no no no. All of Torch Bench, um which is a bunch of single GPU workloads, and that

works perfectly fine. um what I found though is that it really breaks apart in a few cases, in particular with distributed computation. way, uh the thing I want to point out is like, okay, like conceptually I said, "Oh, you just need to like know it's like, okay, input tensor becomes float star A, other becomes float star B, whatever." Uh the trouble is that, you know, like what

the CUDA graph actually sees is very stupid. It sees like a void buffer, that kernel per bottom left of the screen. That kernel parameters of just a void star star, like it do- like each each individual CUDA graph kernel kernel node has no idea what code it's even working on. Uh this is like extremely problematic. Uh if you're like, "Oh, how am I I'm just going to

like find the, you know, uh address inputs and outputs, and then it I'll just know at those offsets in that parameter buffer, I just need to replace those." Uh that doesn't actually work. it doesn't actually work it doesn't straight forwardly work out of the box. Uh we're going to skip this one because it's just too complicated. Um so there's actually a clever algorithm you can do uh

to actually detect pointers sitting in these opaque buffers programmatically. Uh this was Zach DeVito who thought of this. Uh he's I guess somewhat famous in uh PyTorch land. Uh, naively and like this was actually described in a paper called Medusa uh for VLM. You can just check if any 8-byte value in these like opaque buffers uh okay from that CUDA graph's private memory pool. If it does,

then okay, you can say like oh that's totally a pointer. But you can get false positives with this. And like if you override one of these uh fake pointers with like a different pointer value, you're going to be uh that's very bad. You do not want to do that. You're going to get totally incorrect values. So, the solution is really simple. Just capture two CUDA graphs. And

you do that same check twice. And basically you can do a proof by contradiction that uh you're never going to false positive with this. Um so, there's like a few ways where this uh detection algorithm fails. I think I'll skip the first one. I want to mention like maybe like the second and third one where it's just You know, I mentioned this works on off torch bench,

but there are like a few theoretical issues that have kind of prevented me from landing this. you know, like kind of one of the things that we assume is that, you know, the value of a pointer itself does Well, actually it does uh for distributed computation. So, nickel it's what's typically used for torch distributed, uh but also like these are issues with NV SHMEM and like PyTorch

symmetric memory as well. Uh essentially uh uh symmetric memory and registered memory uh assume that they're going to be at fixed memory addresses. And if they are, then you can you can get like latency speed ups and well, latency latency reductions, I should say. Um and that's great, but that actually goes against like what we're trying to do, which is we want to change addresses. This works

only if addresses stay fixed. is that we don't have a way for say nickel or uh PyTorch symmetric memory today to communicate to the parameterized graph launch infrastructure that oh this is static address you cannot change it. That's one of the things I've noticed. The second one is like just terrible. I hope you I I hope this one too much but basically uh you know like if

you had torch compile it has guards right it guards on things like device D type whatever. Uh one thing it does not guard on is alignment of your input tensors. So one of the things that I noticed while working on this is that you know like alignment matters like PyTorch will pick different CUDA kernels at graph capture time depending upon the alignment that you have at graph

capture Uh and you might say okay Daniel just make the like at replay time just demand that the alignment is the same as at capture time right? Yes that that's correct. Uh you know to be concrete like if I request from the memory allocator a four byte aligned pointer it could easily appear as an eight byte aligned pointer. So you know like I know that it needs

to be at least four byte aligned but you know PyTorch might be optimistic and say like this is at graph capture time like I'm going to choose the case that only work with eight byte memory aligned values. Uh so this is something I also don't really know how to handle to be honest. Um and this is something is just a gigantic pain to be honest with you.

Uh so I just thought I'd talk through like all this shenanigans about like okay you like just chase them. you know then you just replace them right? There's also a different technique though. Oh my god okay well I'm so sorry about the visual it it worked much better on my screen. I I thought you kind of hard. The point is is that you don't actually need to

update the pointer parameters. You can actually use a listing. You can have basically virtual pages owned by the CUDA graph and virtual pages and physical pages owned by your application. And the CUDA graph will essentially borrow the physical pages by doing a memory mapping. Uh this is actually being prototyped in OpenXLA right now. uh that pro Uh it's actually somewhat working. Uh this might actually just work

with the nickel registration registered memory problem I mentioned earlier, but I spoke with the colleague this morning and he said it was kind of sketchy. A colleague who works on nickel, so I Anyway, um yeah, I I think uh just due to time we won't spend much time on Um so yeah, conclusion, uh where's this going? I would say that there's a few issues, uh two known

known issues and maybe some unknown unknowns with parameterized CUDA graph launch today. Uh that I kind of just talked through. I'd say I'd say from the user's point of view, like if you're like, "Okay, I do get some of these benefits." I'd say the number one benefit to everybody is reduction in memory usage in the partial CUDA graph's use case. What I expect is that you can

probably expect to see in some weeks or months in the future, a weak informed parameterized CUDA graph launch where essentially only the output but not input parameters are changed on each that would have that Without going into details, that allows you to save the memory uh increase that you would otherwise get. Um because you can still borrow the output memory. It's It's not too It's not too

important. Uh probably exposed via torch compile and CUDA graph trees the fighter guys. As an as like an optional boolean flag you would set for CUDA And I think uh we'll stop there. Um I'm around. Uh I think that AV person just walked out, so uh, hopefully thing is uh, successful. Maybe I can take like one question if anyone uh, it's 12:25. Uh, meet me outside if

you want to speak to me. Sure, sure. Uh, yeah, I mean, I think uh, Li Zhaoxin added something like that. Both like for external ops and mark could autograph on safe and then uh, PyTorch uh, What's it called? PyTorch uh, and Dr. Graph partition will actually like take note of that. Uh, and there's also a context manager as well. So, yes, I think that question is yes.

No, there's no API for that today. Could graph to a low level.