Lightning Talk: Faster Than SOTA Kernels in Torch.compile With Subgrap... Elias Ellison & Paul Zhang

In this talk, Paul Zheng from the PyTorch team at Meta discusses advancements in subgraph fusions and custom operation auto-tuning with Torch Compile. He introduces the Inductor Triton gem backend, which enhances the way matrix multiplications are processed by using a mode called max auto tune, allowing for better optimization between different implementations. The speaker explains the challenges with traditional Triton kernels when dealing with large dimensions and highlights the approach of decomposing K to improve performance. He presents benchmark results demonstrating that this method can yield significant speedups compared to standard implementations. Additionally, Zheng shares insights on the dynamic tuning of hyperparameters based on varying batch sizes and future plans for optimizing fused operations through automation.

Hi everyone. Thanks for coming to my talk. My name is Paul Zheng. I work on the PyTorch team at Meta. And today I'll be talking about sub graph fusions and custom op auto tuning with torch compile. Um, so to start off with um I'll kind of take you guys through the journey in terms of like how this work evolved from just standard torch compile. So here I

kind of introduce the inductor Triton gem backend and essentially what this is is like when you torch compile your model traditionally the the matrix multiplications are just routed to A10 just like with eager. But if you specify this mode called max auto tune for torch compile, what it does is it takes extra time to auto tune between Triton and A10 for the matmuls and it chooses the

best implementation. Um, so it also considers potential epilogue and prologue fusions with your matmul kernel. so essentially like for a matmul where you have like shapes M and K matmul K and N where K is very large compared to M and N, your standard Triton kernel has very poor utilization because there's not sufficient parallelism in your kernel. And the general Triton split K gem uses an atomic

add and that kind of makes fusion not really possible. So I kind of drew this quick diagram to illustrate this case. So in the naive case you see a matrix multiplication between two matrices for a single output element. And in the traditional with your standard Triton gem kernel you have one thread block process one output one output element. So if K is really you process the entire

row of the first matrix and the entire column of the second matrix. So you don't have very good overlapping on the GPU and utilization. The traditional um, K atomic add implementation kind of divides up the K dimension, say between two thread blocks. And so, you know, you see in this part it's like, um, each thread block has less overall work, so there can be better overlapping. Um,

and at the end you can combine kind of each thread block has partial sums of, uh, that matmul, and then you just add them together for the final result. However, if you wanted to do like an operation on the epilogue, like loading a bias and adding it to the result of your matmul, you can't really do that with the atomic add. Um, so kind of, uh, an

implementation we wanted to go with for kind of supporting, uh, better fusible, uh, split K gem implementation in Inductor was this thing called decompose K. Um, here's the source code for it. If If folks are interested in snapping a picture, it's a bit complicated, so I'll try to explain with this diagram here. Um, so we have the same kind of setup, a matrix multiplication A and B,

where K is very large. Um, what we actually do is we transform A and B into 3D matrices, and we split up the K dimension. And then we do a batch matrix multiplication. So, it's kind of like the original split K gem, but, um, after we do the batch matrix multiplication, we just call it torch.sum, and then we have our final value. So, this is like functionally

equivalent to just doing the, uh, regular gem. Um, but I think one thing I would like to point out is like the torch.sum at the end, this is actually fusible with other ops, right? So, if you had like, a, you know, like a cast or an addition or some other random point-wise op, you can fuse that onto the resulting, uh, summation at the end. Here's some benchmark

results. Um, we compared the decompose K paradigm to A10, which uses CuBLAS under the hood. We can see on uh, H100 for FP16 and FP32, decomposed K is generally around like 10 to 15% faster for most shapes. Um and we also do this benchmarking on B200 uh versus A10 um on uh FP16, we see pretty sizable speedups. Uh sometimes there's a bit of a regression. I think

this FP32 data point is very valuable. Like you know, because cuBLAS is a closed vendor library, you have no idea what kind of behaviors it can do. So sometimes it might not pick the best heuristics for like, you know, utilizing split K because it might not think it can use the tensor cores for FP32. Um and so, you know, with decomposed K you can kind of guarantee

like you always do this sort of kernel implementation and it gives you more control um over kind of what you're working with. And this was all done in with uh TF32 enabled. So, moving on, like what about fusions, right? Like we have the sum at the end. What if our model also had a like dot two afterwards that we wanted to fuse in the same kernel um

to have like more end-to-end speedups? So, uh this is something that we also tested, right? Because we can fuse arbitrary ops onto the end of decomposed K. Uh we had some benchmarking numbers here with decomposed K and ReLU. And um essentially the custom op here is a decomposed K plus ReLU. Torch compile was the kind of standard decomposed K but no fusion. And then baseline is A10.

And you can see like in a lot of cases we can get up to like 30 to 40% speedup over the like state-of-the-art cuBLAS baseline. So, now we kind of move on to like, okay, like what does it take to productionize something like this, right? Um we see a very good example in which decomposed K can be very useful in Deep Seek V3 uh in the MoE

router. So, essentially like the router gem routes tokens to different experts, and T is represents the batch size which can vary significantly. Um, so, you know, like the the for decomposed K, there's kind of a hyper parameter you have to pick is the num splits. Like how much exactly do you split the K dimension? based on like your batch size, you might want different splits to reach

peak performance. So, what we do is we actually do per range auto tuning, um, and we kind of dispatch based on a range of batch sizes. So, this is kind of the code here if folks are interested for how to do this natively, um, with torch compile to do the custom op auto tuning. Um, this is kind of like a better description of what's going on. So,

it's like, you know, if if the batch size is less than or equal to one, I dispatch to one implementation. Uh, if the batch size is less than or equal to 128, we use a certain, uh, num splits hyperparameter for doing decomposed K versus if it's like less than or equal to 512, we do a different hyperparameter, etc. Right? So, this kind of allows us to not

have to hardcode one hyperparameter for all shapes, but to be able to dispatch based on the batch size, which is dynamic. Um, here are some benchmarking results for the router gem using decomposed K, and a gemv decomposition. You can see that like for, uh, you know, certain batch sizes, we can reach up to like 4x the default, um, Uh, and this is on V100. Um, some So,

some future plans that we have. So, right now this is supported via torch compile, but it might not always pick the globally best fusion, right? What we do is we benchmark decomposed K itself unfused, and then we pick the fastest implementation, and then we fuse whatever arbitrary fusion that we want to at the end of it. However, like to get really the global optimum of what's going

on, we want to to take each fused version initially and then tune the hyper parameters in Auto Tune and then pick the fastest one. Um, so that's currently kind of one of our future plans for um, reaching peak performance. Yeah, like another example of this that was kind of an extension of this original work is like uh, around distributed collectives. So, this is kind of like an

introduction to how distributed collectives are Auto Tune. Uh, you to rank has uh, like when you do like an all gather um, matmul kernel, um, you take the average of each of them with an all reduce uh, and pick the fastest. um, in a lot of cases for BLM, uh, you know, like based on your batch size, you might want different implementations, right? Sometimes you might just

want the all gather and matmul serialized back-to-back. Other times you want to overlap them and parallelize them uh, and do pipelining. So, here just shows that which which batch size, which implementation wins. Um, another future work that we have planned is really to like instead of doing this hand manual tuning of like, you know, if it's less than or equal to this batch size, we dispatch to

a certain variant, otherwise we use a different variant. Uh, we want to build like a more auto heuristic pipeline to where like, you know, you can build a decision tree based on whatever custom op you're working with um, and like like the sizes that change. And this is some like example of like what this would look like. And cool. Yeah, so that's that's all I got. Um,

feel free to let me know if you have any questions. Oh, yes. Yeah, thanks for the great talk. Just one quick question regarding the talk. I noticed that you guys are using Yeah. So, you guys are using that last size of the one I guess the computer and then next number of those one I think is >> Right, right. So, you mean for this one, Yeah, I

think like, you know, like generally this was more like our hand tuning. We kind of realized that, you um if a batch size fell in a certain range, then a certain split was was good enough, right? Like if we really Right, so we did it manually, but I think um kind of what we aim to want to do with like the auto heuristic pipeline is like let

like a machine like a basic machine learning model learn these heuristics for you, right? To achieve like pretty optimal performance, um and you know, like condition on the right things, right?