Lightning Talk: Combo Kernels: Horizontal Fusion Optimizat... Karthick Panner Selvam & Elias Ellison

This talk discusses combo kernels and horizontal fusion optimization in the context of torch.compile, presented by Karthik from Super Assistant Lab. The speaker explains how multiple small independent kernels can be grouped into a single Triton kernel, allowing them to run concurrently on a GPU. This optimization process overcomes limitations such as the overhead of launching separate kernels and underutilization of streaming multiprocessors (SM) in GPUs. By merging independent operations into one kernel, the efficiency of GPU utilization increases, which is demonstrated through a visual representation of performance metrics. The talk delves into the architecture of the combo kernel pipeline, detailing stages such as scheduling, partitioning, code generation, and auto-tuning, ultimately showcasing the performance benefits achieved through optimized kernel launches. The speaker concludes by mentioning ongoing work to enhance memory awareness and tuning processes related to combo kernels.

Good afternoon everyone. I'm Karthik from Super assistant lab and Elias and I going to present a combo kernels horizontal fusion optimization in torch.compile. So this work has been a shaped by many contributors like across by torch compiler team. So what are combo kernels? When you have a multiple like small independent kernels instead of launching one by one. So we group together into a single Triton kernel so

we can launch and it will run concurrently in a GPU. So let's start with a like something familiar. So this is the function that takes a four tensors and applies the independent reduction mean max sum and another sum. So we decorate with torch.compile and run it on a CUDA tensor. So you see this pattern all the time in the real models the multiple independent operations on tensors

with no dependency between them. They could run in parallel on GPU but today they just launch one by one. So now what it does inductor actually do with this? So it compiles each reductions into its own Triton kernel. The four operations the four kernel launches. You can see that four separate kernel names here. The problem is that it isn't the quality of the like each Triton kernel

is a quantity of the And this is what it look likes in the trace the four distinct kernel launches with its own overhead. Each kernel launch has the overhead like microsecond through the CUDA driver on the CPU side plus the GPU side scheduling cost. For bigger kernels this is a noise but small reductions and point wise this overhead can be comparable to the actual compute. And there's

a second problem like SM utilization a small kernel doesn't generate like enough thread blocks to keep all the SM busy. Most of this GPUs just idle and waiting for the next launch. After the inductor done this vertical fusion, you are left with many independent operation like this. They're like the independent kernels doesn't have no producer consumer relationship. Like can't fuse them vertically. Like each one like pays

the like overhead cost as well as underutilized the GPU. The solution is that is on the fusion. So we take the end independent like kernels and merge them into a one Triton kernel that internally rooms thread block to the right sub kernel using the like TL program ID. Think of which is a dispatch cable inside the kernel itself. The first K thread blocks execute the sub kernel

zero, the next M thread blocks execute the sub kernel one and so on. So one launch instead of N. So this eliminating the dispatch overhead and packing all sub kernels thread blocks into a single grid. So SMs will be more active and simultaneously. Currently this combo kernel is an experimental future. Like it's off by default. You enable them with this combo kernel flags. The key config knob

is per sub kernel flag. It gives each sub kernels its own independent block sizes. Like X block zero, X block one and so on. Instead of forcing them using the shared configuration. This matters because the two reduction kernels with different shapes can have a different optimal block sizes. With combo kernels enabled, this four separate reductions horizontally fused into a single kernel. Like one kernel and one launch.

The operations themselves are unchanged. It's only purely like fusion level optimization. And here's the trace with combo kernels. Like where you visually saw the four bars for like four different Triton kernels. Like the launch overhead and idle SMs. Now you see like one wider ball, the same compute but GPUs busy. I will just go through this combo kernel architecture. So the combo kernel pipeline has four stages,

scheduler, horizontal partitioning, and code gen and Triton heuristics. So it starts in the scheduler that after inductor build this dependency graph and does the topological sorting, we find the node that are same topological level. These nodes with no data dependency between them. So they can safely run in parallel. We filter them for eligibility like compatible nodes types and same devices and no order violations. There is also

an optional flag like benchmarking gate. So we compile both fused and unfused version and benchmark them and only keep the combo kernel if actually wins. Like the ideal goal we plan to is make this combo kernel always profitable like without needing the benchmarking this pass. So and then we also don't blind and like merge everything. So the partition happening like multiple pieces. First we isolate the large

point wise kernels that already saturate the GPU. Fusing them just add the overhead. So no benefit for us. And likewise the separating the long reduction with the short reduction and we also enforce the odd limits. And when you are fusing the several kernels, each of this own like input and outputs, the odd limits will be like you know quickly get the limits. And if the groups ends

with just one kernel like it will stays as the regular Triton kernel instead of the So now we get the code generation part. So this is where the actual Triton code gen is produced. The first we select these dispatch strategy. There are few options like we prefer the flatten grid approach. So, this strategy determines how thread blocks get mapped into the sub-kernel. Then we generate the PID

routing logic, a chain of if and else branches. Each thread blocks gets like the TL program ID. Like compare against the offset to figure out which sub-kernel it's belongs to. Then inside this each branch, the sub-kernel body is fully inlined with like like kernels read, compute, and writes. And the code gen builds with the combo grid meta, which will be at the dictionary that tells the runtime

how to compute this total like, you know, launch grids. So, last stage is in auto tuning. So, this is where like it's get interesting when like per sub-kernel config has been enabled. So, each sub-kernel calls its heuristics like pointwise reduction and persistent reductions. So, based on its operation type, so the block size parameter gets suffixed with like X block zero, X block one, and so on. So,

each sub-kernel has its own independent tuning space. So, this grid computation sums up the thread block across all sub-kernel to get the total launch grid. Then we auto-tune with chained sequential approach. Tune the largest sub-kernel first with holding the rest fixed, then next one and so on. At the end we do the final pass that tune the shared parameter like number of warps across the old combo

kernel. So, here you can see the actual generated Triton code. So, the few things to notice like per sub-kernel config at the top, each sub-kernel has its X block zero or block zero. Like they are not sharing the uh like single configuration. And PID routing law branches if uh PID is lesser than offset zero, and elif PID is lesser than offset one, that's our the dispatch table.

So, each thread blocks checks with this range it falls into and jumps to the right subkernel. And inside the each branch that full inside like inline subkernel uh body load, compute, and store exactly it would appear in the standalone Triton So, now now with the numbers. So, on our CA perf suite on H100 GPU, the combo kernels delivered geo mean speedup of 7.3% on again phase and

a 5.97% on torch bench. Some models see like much larger gains like uh particularly those with small independent operations. The model dominated by like large matmuls see less benefit that is expected. So, those kernels already saturated the GPU. So, what's compelling here is that it's a pure compiler optimization, no model change, no user code change. So, the gains come from two sources, one from like fewer launchers

and better SM utilization. So, you flip the config and compiler does the rest. So, let me close with like what we are actively working on this combo kernel. So, first is memory aware fusion. So, when you fuse the kernels, the live tensors overlap, which can increases the peak memory. So, we are adding the analysis, so we fuse only if doesn't increase the peak memory by given threshold.

And uh improving the coordinate descent tuning. So, the chained auto tuning I described earlier is a good starting config through a grid search. So, we are extending the coordinate descent tuning as a hill climber to fine-tune each subkernel block sizes independently after that. And using the pointwise with reduction. So, currently combo kernel fuse like reduction with reduction and point-wise point-wise and fusing across these types requires careful

handling and but the SM utilization like gains make it worth pursuing and we also plan to enable by default in OSS. Yeah, that's it. Thank you.