Lightning Talk: Achieving SOTA GEMM Performance: A CuTeDSL Backend for PyTorch Induc... Nikhil Patel

This talk covers recent advancements in integrating QTSL gems into PyTorch Inductor to enhance performance in transformer-based models. The speaker explains that GEM (General Matrix Multiplication) operations are critical for optimizing throughput, as they account for almost 90% of runtime during inference. To address performance challenges, especially with evolving GPU architectures, the speaker discusses the implementation of a QTSL-based backend. This backend leverages NVIDIA’s Cutless API for kernel management, allowing for effective autotuning and benchmarking of various kernel configurations. The speaker highlights the benefits of achieving speed improvements in specific operational shapes and emphasizes that minimal user intervention is required to utilize this backend. Future developments include more operations support and enhanced performance optimizations.

I'm Nikil. I'm going to be talking a little today about some of our recent work to bring some more QTSL gems into PyTorch Inductor. So, just as like a little bit of background, um gems do remain one of like the primary like throughput bottlenecks in transformer based models. Um in this little diagram I've made, um I like break down the runtime spent in various operations uh doing

some inference for llama 3.370B. Um and you can see that almost uh 90% of the time is spent doing gems. Uh so they're a very important part of the model uh to optimize if we want to bring down uh latency and increase throughput. Um and users expect like torch compile to deliver high performance gems and they don't want to have to write custom kernels or do a

lot of manual tuning to find the best configs. But especially more recently, you know, as like new GPUs come out and they ship with like more complex features, um it's getting harder and harder for Inductor to maintain uh these like state-of-the-art gems. So our idea is to sort of uh explore alternate ways to increase uh torch compile gem performance and we're interested in doing that uh through

QTSL. So just as like a brief overview of what QTSL is, it's essentially like a Python API for Cutless. Um, and we're interested in it for a couple of reasons. Um, the first is it's a relatively low-level language. So, uh, we get like full access to like threads and memory hierarchy. So the performance ceiling is pretty high. Um, we also like it because of its relatively low

compile time. Uh, since we can just in time compile it. Uh this allows us to do epilog fusions and benchmark those at compile time. Um as well as do some autotuning to discover like optimal kernel configs. And we like it for the active NVIDIA investment since they are the ones maintaining this and obviously they have early access to their own hardware. Um we expect QTSL to support

new hardware features relatively fast. So we're interested in adding a QTSL gem backend to Inductor. And at a high level how that works is you have your problem shape uh along with its like dtype layout any other info and you call torch.mm on that. Uh then inductor compiles like a list of candidate kernels. Uh so A10 it's kind of like the default. Uh it'll submit one choice.

Then Triton and Cutless uh each submit like multiple candidate kernels usually with different configs. Um and now we've added QTSL will also be submitting multiple candidate configs. uh candidate kernels with different configs. And then for Triton, cutless and QTSL, all those kernels are chit compiled um in parallel. Although small note, the parallel compilation for QTSL is coming soon. Um all of those are benchmarked and then the

fastest kernel is selected as the one that will actually be used like during the runtime of your model. So zooming in a little bit into how this works specifically. Um this backend uh operates like a little differently than say like the existing Triton backend. Uh the existing Triton backend is made up of like handwritten Triton kernel templates in inductor and that can get a little hard to

maintain. Um you know as new hardware comes out sometimes these templates have to be uh rewritten. Uh if we want to support like new scaling modes we'll have to manually add new uh kernel templates to inductor. So for the um QTSL back end, we kind of went a different route and we're leveraging something from Nvidia called the cutless API. So this is a Nvidia maintained kernel library

uh where we query it with uh the aspects of our problem like the shape, dype, layout, other stuff like that and it will return all compatible kernels uh that could support that operation. Um what's nice about Nvidia owning this is that you know as new hardware comes out they can just update this library with new kernels support new scaling modes and then since um inductor is just

calling this API we'll get access to all that for free without much uh maintenance burden on our side um but since this API returns like all the compatible kernels that would be too much at compile time to compile them all and benchmark them all. So we leverage a second library from NVIDIA called NVMAL heristics. Um and this is essentially an analytical model that will estimate the performance

of different configs for your problem. And so through that we can like prune the amount of kernels that we need to autotune to around five. Um and that number can be tuned with like an environment variable that you can set. Uh but we found that five is a pretty nice sweet spot. And then those kernels are just profiled along A10 and Triton and then the fastest one

wins uh for runtime. So getting into some performance numbers uh we can see that at the um larger problems you know uh Kublos is pretty well tuned and it's hard to get performance there. But especially at some of these smaller shapes and the more like decode type shapes with like small M uh large N and K um we find that our like QTSL backend does perform uh

pretty nicely. Um it's a similar story for MXFP8. We can see especially on like the decode shapes um we get some nice speed ups over A10. And I want to mention that you know uh with this back end we are able to do epilog fusions although that's technically launching soon. It's not available right now. So um in actual models you know eventually we'll be able to fuse

other ops into these and potentially see more speed ups over these A10 Kublas kernels. And similar story for uh NVFP4. And so uh right now we can also just uh plug this into VLM and run some uh models in inference. Uh so we benchmarked uh on Llama 3.18B uh Quen 33 32B and Llama 3.370B. Um, and we can see that we do uh observe some pretty nice

speed ups, you know, especially at batch size uh 16, we see almost uh 7% speed up uh on llama 3.370B, which is a pretty nice speed up uh for something that doesn't really require um any work on the part of the user. All they really have to do is just add this new backend. So the current supported ops um we have MM BMM and then for scaled

MM we currently support MXFP8 uh MXFP4 and NVFP4 uh and we also support grouped MM. Uh coming soon uh we expect uh to have more support for hopper. Um we're also working on implementing epilog fusions uh pointwise fusions first and then eventually uh broadcasts and some reductions. Uh we plan to add a blockcaled group gem shortly. Um as well as NVMAL heristic support for group gem and

then AOT compilation. So you could use this today. Um first you need to install the dependencies. Uh right now installing the actual cutless API is a little bit weird. Um but Nvidia is hopefully planning to merge that into just the QTSL uh package and make it pip installable. Um, and then when you want to actually use this, all you need to do is add NV gem, uh,

we named it that for NVIDIA universal gem, uh, to your list of, uh, max autotune uh, backends and then just compile your model with maxtune mode turned on. And then you should start seeing these QTSL kernels, um, appear when autotuning. Um, and what's nice is that if this backend doesn't support like your op, uh, it'll automatically fall back to Triton or A10. So you don't need to

worry about uh having operations in your model where this might not support it. And if you're interested in learning more, uh we just released a PyTorch blog post about this yesterday. So you can scan that QR code. Um and yeah, I think that's pretty much all I have. So if we have any questions, uh I'd love to answer maybe one or two.