Speed regression in rmsnorm_fwd when used with torch compile

[Maharshi-Pandya]

On a B200, when I run the benchmark_rmsnorm file WITHOUT torch compiling the rmsnorm_fwd the results look like:

=== RMSNorm Forward Pass Benchmark ===
Tensor dimensions: [9450, 5120]
Input and Output Data type: torch.bfloat16
Input tensor shapes:
x: torch.Size([9450, 5120]), dtype: torch.bfloat16
w: torch.Size([5120]), dtype: torch.float32
Executing kernel...
Kernel execution time: 0.0413 ms
Mem throughput: 4685.00 GB/s
Ref kernel execution time: 0.0649 ms
Ref mem throughput: 2981.00 GB/s

However, after doing compiled_func_cute = torch.compile(rmsnorm_fwd) the result I get is:

=== RMSNorm Forward Pass Benchmark ===
Tensor dimensions: [9450, 5120]
Input and Output Data type: torch.bfloat16
Input tensor shapes:
x: torch.Size([9450, 5120]), dtype: torch.bfloat16
w: torch.Size([5120]), dtype: torch.float32
Executing kernel...
Kernel execution time: 0.1067 ms
Mem throughput: 1815.00 GB/s
Ref kernel execution time: 0.0624 ms
Ref mem throughput: 3103.00 GB/s

I wonder where the speed regression is coming from after torch compiling because if we use this within a larger block of operations alongside torch compile, the overhead will eat all the speedup even if the CuTe kernel is compatible with torch compile. Any help is appreciated :)

[tridao]

Can you check if only 1 kernel runs when you torch compile it (e.g. with torch profiler)? This is to make sure torch compile doesn't add extra copies etc.
The only other thing I can think of is that we're not wrapping it correctly.
Tbh idk torch compile very well

[tyler-romero]

do_bench has warmup set to 5 here:

quack/benchmarks/benchmark_rmsnorm.py

Line 55 in 883eaee

avg_time = do_bench(fn, warmup=warmup_iterations, rep=iterations)

This is 5ms, not 5 iterations as the code clains, so I wonder if there isnt enough warmup for the torch.compile version.

[Maharshi-Pandya]

This is with warmup set to 1000 (ms) and rep set to 2000 (ms), it doesn't change. I think there is some overhead added by torch compile.

=== RMSNorm Forward Pass Benchmark ===
Tensor dimensions: [9450, 5120]
Input and Output Data type: torch.bfloat16
Input tensor shapes:
x: torch.Size([9450, 5120]), dtype: torch.bfloat16
w: torch.Size([5120]), dtype: torch.float32
Executing kernel...
Kernel execution time: 0.1080 ms
Mem throughput: 1791.00 GB/s
Ref kernel execution time: 0.0653 ms
Ref mem throughput: 2964.00 GB/s

[Maharshi-Pandya]

With profiling I think we can see the overhead from torch compile.

According to NSYS, the kernel execution time is similar to what we get normally without torch compile.

[tridao]

i'm surprised torch compile cpu overhead is that high (but then again i don't usually use torch compile). What about torch compiling other small kernels? Do you run into the same issue?
(e.g. can you try torch compiling the softmax kernel?)

[mattteochen]

For mem bound kernels, if you don’t use cuda graphs (e.g. reduce-overhead mode), torch compile host overhead will impact perf. For a ~30us kernel, it won’t beat eager due to its compilation machinery.

[Maharshi-Pandya]

This is what I get with "reduce-overhead" mode i.e. cudagraphs. Seems like there's still lot of overhead coming from torch compile.

=== RMSNorm Forward Pass Benchmark ===
Tensor dimensions: [9450, 5120]
Input and Output Data type: torch.bfloat16
Input tensor shapes:
x: torch.Size([9450, 5120]), dtype: torch.bfloat16
w: torch.Size([5120]), dtype: torch.float32
Executing kernel...
Kernel execution time: 0.1172 ms
Mem throughput: 1652.00 GB/s
Ref kernel execution time: 0.0658 ms
Ref mem throughput: 2943.00 GB/s

[mattteochen]

I'm happy to take a look. Could you share a repro/command @Maharshi-Pandya ?

[Maharshi-Pandya]

I run the benchmarks/benchmark_rmsnorm.py file with the following command:

python benchmark_rmsnorm.py --M 9450 --N 5120

Also, I use compiled_func_cute = torch.compile(rmsnorm_fwd) with and without "reduce-overhead" mode for benchmarking. The warmup and reps both are set to 5000 (ms) for do_bench @mattteochen

[debashishc]

Interesting, I am not getting the regression with torch.compile at least on a modal run. My script for reproducing this is here.

env settings

GPU: NVIDIA B200
torch: 2.10.0+cu130
triton: 3.6.0
CUDA: 13.0

=== eager quack + compiled ref ===
Tensor dimensions: [9450, 5120]
Input and Output Data type: torch.bfloat16
Input tensor shapes:
x: torch.Size([9450, 5120]), dtype: torch.bfloat16
w: torch.Size([5120]), dtype: torch.float32
Executing kernel...
Kernel execution time: 0.0491 ms
Mem throughput: 3945.00 GB/s
Ref kernel execution time: 0.0503 ms
Ref mem throughput: 3850.00 GB/s

reduce-overhead still shows regression though.

=== torch.compile(rmsnorm_fwd) ===
[torch.compile(rmsnorm_fwd)]
Kernel execution time: 0.0473 ms
Mem throughput: 4091.00 GB/s
[torch.compile(rmsnorm_fwd, mode='reduce-overhead')]
Kernel execution time: 0.0845 ms
Mem throughput: 2291.00 GB/s

@Maharshi-Pandya what is the torch version you are compiling with? also please let me know if any issue with the script as well.

[Maharshi-Pandya]

My env settings @debashishc

GPU: NVIDIA B200
torch: 2.9.1+cu130
triton: 3.5.1
CUDA: 13.0

[zou3519]

@Maharshi-Pandya I was unable to reproduce this on PyTorch 2.10, but I was able to see a small delta on PyTorch 2.9.1. Could you try to upgrade to PyTorch 2.10 and see if the problem goes away for you please?

• benchmark_rmsnorm2.py that does compiled_func_cute = torch.compile(rmsnorm_fwd)

On PyTorch 2.10:

$ python benchmark_rmsnorm.py --M 9120 --N 5120
=== RMSNorm Forward Pass Benchmark ===
Tensor dimensions: [9120, 5120]
Input and Output Data type: torch.bfloat16
Input tensor shapes:
x: torch.Size([9120, 5120]), dtype: torch.bfloat16
w: torch.Size([5120]), dtype: torch.float32
Executing kernel...
Kernel execution time: 0.0473 ms
Mem throughput: 3951.00 GB/s
Ref kernel execution time: 0.0491 ms
Ref mem throughput: 3802.00 GB/s

$ python benchmark_rmsnorm2.py --M 9120 --N 5120
=== RMSNorm Forward Pass Benchmark ===
Tensor dimensions: [9120, 5120]
Input and Output Data type: torch.bfloat16
Input tensor shapes:
x: torch.Size([9120, 5120]), dtype: torch.bfloat16
w: torch.Size([5120]), dtype: torch.float32
Executing kernel...
Kernel execution time: 0.0473 ms
Mem throughput: 3952.00 GB/s
Ref kernel execution time: 0.0491 ms
Ref mem throughput: 3801.00 GB/s

There is a bit of a difference on PyTorch 2.9.1 for me, but not as dramatic as yours:

$  python benchmark_rmsnorm.py --M 9120 --N 5120
=== RMSNorm Forward Pass Benchmark ===
Tensor dimensions: [9120, 5120]
Input and Output Data type: torch.bfloat16
Input tensor shapes:
x: torch.Size([9120, 5120]), dtype: torch.bfloat16
w: torch.Size([5120]), dtype: torch.float32
Executing kernel...
Kernel execution time: 0.0472 ms
Mem throughput: 3954.00 GB/s
Ref kernel execution time: 0.0473 ms
Ref mem throughput: 3950.00 GB/s

$  python benchmark_rmsnorm2.py --M 9120 --N 5120
=== RMSNorm Forward Pass Benchmark ===
Tensor dimensions: [9120, 5120]
Input and Output Data type: torch.bfloat16
Input tensor shapes:
x: torch.Size([9120, 5120]), dtype: torch.bfloat16
w: torch.Size([5120]), dtype: torch.float32
Executing kernel...
Kernel execution time: 0.0639 ms
Mem throughput: 2924.00 GB/s
Ref kernel execution time: 0.0473 ms
Ref mem throughput: 3952.00 GB/s

[Maharshi-Pandya]

Thanks for this, I’ll try with torch 2.10 and report back if I find the regression there as well @zou3519

[debashishc]

ya I think with @zou3519 analysis backing what I saw too, this could well be torch version mismatch. As an aside I'd like to understand what torch.compile fixed on their end.

My env settings @debashishc

GPU: NVIDIA B200
torch: 2.9.1+cu130
triton: 3.5.1
CUDA: 13.0

[mattteochen]

On sm120 I don't see slowdown with torch==2.9.1

=== RMSNorm Forward Pass Benchmark ===
Tensor dimensions: [9450, 5120]
Input and Output Data type: torch.bfloat16
Input tensor shapes:
x: torch.Size([9450, 5120]), dtype: torch.bfloat16
w: torch.Size([5120]), dtype: torch.float32
Executing kernel...
Kernel execution time: 0.1340 ms
Mem throughput: 1444.00 GB/s
Torch compile Cute kernel execution time: 0.1340 ms
Torch compile Cute mem throughput: 1445.00 GB/s
Ref kernel execution time: 0.1333 ms
Ref mem throughput: 1452.00 GB/s