Seamless Integration: Custom Kernels in the Torch.compile Stack Wi... Kshiteej K, Masaki K & Pawel G

This talk covers the integration of custom kernels into the torch compile stack, focusing on the collaborative efforts of the Transformer Engine team at Nvidia. The speakers explain the advantages of the Transformer Engine, including optimized layers for Nvidia's latest hardware, such as Hopper and Blackwell, and the importance of low precision training. They describe how torch compile can enhance performance through operation fusion and reduced CPU overhead. However, they address significant challenges, particularly graph breaks that prevent efficient operation when integrating custom kernels. The presentation concludes with a solution involving custom operations to streamline the integration process and optimize performance, showcasing the benefits through practical examples and performance metrics.

Hello everyone. Uh, my name is Pavel and I work in Transformer Engine team at Nvidia and today together with Masaki and Kshitij will present our work on integrating uh, custom kernels in the torch compile stack. So, we will talk about our experience of integrating Transformer Engine into torch compile. We'll show you why these two don't really work well together and we will show you the solution how

you can integrate the library providing custom operations into PyTorch. Note that Transformer Engine is something we work on but the content of this presentation is related can be related to any library providing custom kernels or standalone custom kernels will also uh, work that way. So, I will start by introducing Transformer Engine. Uh, so Transformer Engine is library developed by Nvidia which provides drop-in replacement layers to PyTorch

layers but these layers are optimized for newer newest Nvidia hardware like Hopper or Blackwell. you can think about it like PyTorch is general software which supports a lot of hardware and a lot of layers but Transformer Engine is focused only on Hopper and Blackwell and probably future in about Rubin. And it contains only a few layers related to language modeling. And because it's so focused, we can

ship much faster and it's much much more performant than standard PyTorch layers. So, we put into Transformer Engine newest kernels for newest hardware we use for example this type engine kernels for our MLPerf submissions. So, we have state-of-the-art kernels kernels inside our library. So, let's look what makes Transformer engine so fast. So, primarily, we support low precision training. Transformer engine was developed to showcase the possibilities of

FP8 current scaling on Hopper. But, recently, we added support for MXFP8 and NVN before training, uh which showcases the possibilities of Blackwell GPUs. Moreover, we provide we expose a kernels uh to attention. So, we perform automated attention bucket selection. Uh we choice the fastest kernels uh from QDNN or from flash attention from unfused PyTorch attention. And we chose the kernel which is the fastest on given configuration

and on given GPU. And we have a lot of many more custom kernels and features, but let's not talk about the Transformer engine. So, I will leave it as it is. And there is I will do short recap what Torch compile does. So, Torch compile takes as input some Python code. Then, it tries to understand that convert to some intermediate representation, and then generate some which is

optimized. So, what are the benefits of Torch compile we'll focus on? It is fusion of operations, which reduce GPU time. It is CPU time reduction. And there is mode reduce overhead in Torch Inductor, which uses CUDA graphs, uh which can reduce CPU time to almost zero. So, I will say a few words why CPU time is more and more important nowadays. Because few years ago, the word

ML workloads all look like this. Like kernel times were relatively long. And time to issue them on CPU were relatively short. But GPUs are getting much much faster and CPUs are getting only slightly faster. So, it may turn out that current workloads may look like this. And this is especially painful for workloads with which have a lot of light kernels, like MOEs training or small batch size

or inference decode phase. So, that's why we care also about CPU Uh so, I showed you in short ways two ways of speeding up the model. First was torch compile and the second was transformer engine. So, we can see that these features and are somehow orthogonal and we may want to use them together in one model. But they don't really work together and this speed up often

does not really add up. And I will try to show you why it is the case. So, consider some PyTorch model composed of some PyTorch layers. So, we can speed up it on two ways. First is change some layers to transformer engine layers to use fast custom kernels. And the second way I mentioned is to adding using torch compile to compile this model into one graph. But

what will happen if you will try to this two things together? We'll get something like this. What happens is that transform torch compile does not understood well transformer engine layers. So, it will split the graph into sub graphs if it understands and run Transformer engine layers eagerly. And this is in fact not faster. Often it's not faster than just applying torch.compile or just changing some layers to

fast Transformer engine layers. So, Masaki will say why it's not Thanks, Pavel. So, let's get into the main issue we are dealing with, which are graph breaks. When we call torch.compile, Dynamo tries to trace a Python program into a single graph, but if it sees anything it cannot trace, for example, a custom CUDA kernel, it falls back to eager execution. Instead of one graph, we end up

with multiple smaller graphs with Python between. This fallback is what we call a graph break. Visually, you can see how one logical model execution now becomes graph one, graph two, graph three, stitched together by Python. This is not rare as we have custom kernel inside TE layers, so we can easily end up with dozens of graph So, why are graph breaks such a big deal? Because they

degrade performance in three separate ways. First, guard overhead on the CPU. Every compiled graph has guard runtime checks that validate whether we can reuse the cached compiled code. If we have n graph breaks, we get n plus one graphs, and each one has its own guard set. Those checks are evaluated on every iteration. Second, lost fusion opportunities. Inductor can only fuse operators within a single graph. Once

a break exists, that becomes a hard boundary, no fusion across it, which means extra kernel launches and extra memory traffic. Third, poor composability. Features like reduced overhead through CUDA of CUDA graphs work best with a single unified Once a model is fragmented, features uh sorry, once a model is fragmented, many of these optimizations become ineffective or unusable. So, graph breaks don't just cause performance, they block compiler

features. Let's make guards concrete. Each compiled graph comes with guards. Runtime conditions that decide whether we can reuse cached compiled code. At compile time, Dynamo records a snapshot of input properties like D type, device, shape, and strides. On every subsequent call, it checks whether current inputs still match. If they pass, we take the fast path. If they fail, we may re- compile and the cache another variant.

Now, here's the key point. Graph breaks multiply guards. Each graph fragment has its own guard manager, and all of those guard checks run every iteration. So, graph breaks don't just block fusion, they add per-core CPU overhead. So, let's take a look at this simple function example. On the left-hand side, there's a graph break that one break splits execution into two graphs, and we have 17 guard checks

at runtime. On the right-hand side, the same function stays as a single graph. There's no break, so we only have one graph and eight guard checks. The takeaway is simple. Breaks don't just add one cost, they multiply cost But usually, the bigger performance hit is lost fusion. Inductor can fuse ops within a single but the break creates a hard boundary and we can't fuse across the boundary.

That leads to more kernel launches and more intermediate tensors written to memory. So, even with a very fast custom kernel, fragmentation around it can dominate end-to-end time, especially in CPU-bound or latency-sensitive workloads. So, let's connect this back to TE with a simplified example. Here is a TE-like linear and the forward calls two quantizations on weight and input, and we call it gem. And we also set Python

quantizer objects that control behavior. This is intentionally simplified, but this captures the core pattern, opaque kernel calls plus non-tensor metadata, which is exactly where graph breaks tend to appear. So, why doesn't Torch compile just work on this linear? Even though it won't trace the custom CUDA kernels, it still needs enough information to compile correctly. First, output metadata during tracing, such as shape, stride, dtype, and device without

running the real kernel. Second, the correct mutation and alias and semantics, so it knows whether the reordering of operation is safe. Third, a safe representation of non-tensor Python objects, like the quantizers, especially if they are mutable or stateful. If any of these are unclear, Dynamo prioritizes correctness over performance and will insert graph break. And that This is exactly what happens with TE's style calls today. Next, Kush

T show how we address these requirements. Thank you, Masaki. So, Masaki explained what the problem with graph breaks are. Let's have a look at how custom op and opaque objects can help us solve this. So, custom op is a decorator which you can wrap around the custom kernel that you have such that it becomes a first-class PyTorch operator. What this also does is it makes it composable

with the torch compile stack as well. Few things that you need to provide as a author are first, the real implementation for the kernel, the fake kernel which will help map the map the input and the output metadata. Third, declare the in-place semantics for your custom operation. And finally, provide an autograd rule if this is intended for training scenario. How torch compile uses this is that during

the tracing time, Dynamo will actually use the fake kernel such that it can understand the output metadata and doesn't need to graph break on that. And at execution time, it would just actually call the kernel that you've wrapped. Let's have a look at why declaring the in-place semantics is really important. Some passes in the torch compile stack may reorder operations and to do so correctly, they need

to understand the tensor aliasing and also the in-place semantics of the operations. If that isn't the case, then they could reorder incorrectly leading to silently incorrect results, which is bad. So, to help with this, custom op also provides an argument called mutate args. The mutate args itself takes a string or a iterable of string. And let's see how we can use this for different scenarios. So, first

is if you have a custom op which is pure or functional, you can just pass empty tuple or empty string. Second, if you know that our X and Y arguments which are in mutated in place, then you can just pass a list or a tuple with X and Y as strings and uh the Torch compile would know how to deal with this. And finally, if you're wrapping

a third-party kernel and you are unsure about what the semantics of the kernel are, then you can just pass mutates args unknown and it will pessimistically assume that all inputs are updated in place. This could lead to slightly poorer performance, but this is still correct. So, uh it's better that this is done Otherwise, you'll have undefined Now, let's have a look at how we can use custom

op to decorate the example that we just saw. And in this case, we just decorate the custom kernels that we have with custom op. Pass Give it a name T linear in this case and also pass the mutates args. Another thing that we do is we register the fake kernel. The job of the fake kernel is to just provide the output metadata, so it doesn't need to

actually do any computation at all. And this would have worked fine if we did not have quantizers object being passed to custom op. So, Torch compile stack can understand plain tensors or types like ints and floats, but if you have some custom Python class that needs to to be passed to the custom op, then that can lead to a graph break as well. But, we have opaque

opaque objects API which can help teach Torch compile stack how to deal with them. So, to use that, you need to use opaque base and also call register opaque type on it. One thing to note is that this API is subject to change and it is experimental. And this opaque objects API also provides two types of semantics that your opaque object can have. One is reference semantics,

and this is useful when you have objects which are mutable or stateful, which would change across iterations. And in this case, these objects get passed to the graph as inputs such that the they are wired into the custom op as a reference, and that custom op can then make any mutations inside it. The way to implement these are that you just have to inherit from opaque base.

For the value type semantic objects, these are useful when you have immutable configuration like mode flags or settings. And in this case, they would get baked into the graph. And to do so, torch compile also needs to add extra guards on this. So, you also author not only have to inherit from opaque base, but also implement few additional methods like equality, hashing, and FX representation for the

same. So, let's have a look at how we can use opaque objects for the quantizer class that we had. So, we just have to inherit The rest of the body for the class should remain the same. And finally, we can just call register opaque type on And finally, about the autograd support, you need to provide two things. One is setup context, and another is the backward rule

for your custom op. Setup context should be familiar for those who've written torch.autograd function, wherein it just isolates saving the tensors for backward and any other metadata. the backward is just the backward computation for computing the gradients. One important point to note here is that both of these have to be torch compile traceable. So, what that means is that they shouldn't do anything which would lead to

graph breaks. So, they cannot access data pointer or anything like that. And this becomes a problem if you have a if you have a custom kernel for the backward pass as well. In which case you need to actually just wrap it in custom up again so that it becomes traceable and it can be understood by torch compile. And finally you just call register autograd passing these two

Moving ahead I'll pass the pass it on back to Pavel so that he can walk you through the results and the current status. Okay, we would like to show you some numbers. Uh so we'll show you the motivation why we are doing this thing. Uh so we run some very basic example one forward with llama for like hugging face on one GB 200 GPU. And you can

see Uh we have two configuration. One is GPU and CPU bound depending on which line we look at. And the second one is pretty artificial uh very small batch size and very small sequence length. And it shows basically only CPU overhead. we can see that torch compile gives us quite huge benefit. FP8 does not give us that much benefit probably because of CPU overhead. Applying two of

them gives us the benefit but not bigger than applying uh torch compile alone. But if we use custom ops to prevent graph breaks we can get the speed up uh which is the fastest. And it's worth noting that this reduced overhead back end can reduce CPU overhead to almost zero. But if there are graph breaks it can be really really really slow because this CUDA graphs have

got a lot of like setup time. So splitting the graph and running CUDA graphs for parts of the graph is extremely expensive. So, some words about the progress of integration on torch compiling to our library transformer engine. So, it's worth noting that this solution pretty show is in fact under active development. Uh and this support for transformer engine is also under active development. We merge support for

some layers. Some layers are still uh worked on and it is not that easy and we need sometimes to contribute to custom ops API and we'll need to probably to do probably to do this also in the future to enable full transformer engine support. Uh here are some references. If you want to know more about custom operations. And we I want to thank Meta folks who custom

op API so that we are able to start our work toward integrating T into Uh thank you very much for your attention. I think we have 4 minutes so maybe one question if there is any. Okay, thank you very much. Okay, uh there is a microphone. Hey, yeah. Um thanks for the presentation. I wanted to ask um because I've had this problem multiple times making the decision

how to do outputs one thing would be you allocate the in PyTorch itself and then I put a mutable argument to your custom kernel um which would give PyTorch um torch compiler like it would actually see the allocations and could potentially do optimizations with that, right? Like in your example, you actually specify like the the out like the output gets allocated in your code, it's invisible to

torch compile um and you need to annotate uh the fake kernel so that it knows the output. Like do you have any insight which of the two options um is better and why? So, first thing is that we are working into like existing library transformer engine and we cannot really afford to like change everything to uh add this torch compile support because it's already quite big. So,

it sometimes is done that way in transformer engine and we need to like there are some other reasons not related uh to do it that way. And we need we need to just to adjust. Uh so, you are asking what's better? Like allocating outside custom up like buffer for output inside? Yes. Uh I don't have like one decisive answer. Do you have something? I think we can

talk a little bit more about this. After the talk. Okay. Uh we don't have a time for next question. So, thank you very much for attending.