feat: RMSNorm backward pass kernels

[aghilann]

Description

Adds RMSNorm backward pass to TileGym - the first backward kernel implementation

Implementation:

• Single CuTile kernel (rms_norm_backward_kernel_dx) computes dx row-parallel and stores intermediate values (dy * x * rstd) into a float32 temp_buffer
• dw is computed via temp_buffer.sum(dim=0) using PyTorch's optimized reduction (avoids a second kernel with different access pattern)
• All accumulations are done in FP32 regardless of input dtype for numerical stability
• Added PyTorch reference impl (rms_norm_backward_torch) for testing/benchmarking
• Added test_rmsnorm_backward.py and bench_rmsnorm_backward.py

Performance

2-5x faster than PyTorch across all dtypes.
Both Torch Reference and cuTILE still do accumulates in FP32

bfloat16

N	CuTile (GB/s)	PyTorch (GB/s)
1024	1537.03	472.44
2048	2280.10	533.60
4096	2813.37	515.98
8192	3640.76	526.04
16384	3916.75	542.96

float16

N	CuTile (GB/s)	PyTorch (GB/s)
1024	1563.98	484.67
2048	2361.91	535.59
4096	2927.92	520.76
8192	3566.86	532.16
16384	4023.40	549.59

float32

N	CuTile (GB/s)	PyTorch (GB/s)
1024	2255.25	853.33
2048	2922.93	931.21
4096	3892.06	901.09
8192	3645.70	922.79
16384	2638.26	955.57

CI Configuration

config:
  build: true
  test: ["ops", "benchmark"]

Checklist

• Code formatted (./format.sh)
• Documentation updated (if needed)
• CI configuration reviewed

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[vgoklani]

why is the bfloat16 implementation so much slower?

[aghilann]

why is the bfloat16 implementation so much slower?

I changed my the design of the second kernel and removed it altogether since it was a huge bottleneck and there was an easier way. I honestly have no clue why bf16 was so much slower then fp16 since I deleted it before investigating it.

[vgoklani]

We use all bfloat16 so this would be a nonstarter.

[aghilann]

@vgoklani Yes I'm aware, that's why I changed the algorithm (it was slow regardless of BF16 vs FP16, though not sure why BF16 was so much worse). Anyway with my new implementation, BF16 vs FP16 performance are now pretty similar.

[aghilann]

@hannahli-nv could I get a review + CI on this PR?

[hannahli-nv]

/ok to test 24e65b8

[hannahli-nv]

/ok to test 077e8fd

[aghilann]

Oops don't know how that # made it in there. Thanks for fixing it

[hannahli-nv]

Overall LGTM, thanks for the contribution!