[Test] Expand RMSNorm backward configs, cache reuse, and Torch benchmarks

[jhinpan]

Summary

Follow-up to #800's final review and the post-merge note:

• centralize the plain and fused-add backward/autograd config matrices so their coverage cannot drift;
• add the missing unaligned N > 2048 generic-path coverage to fused-add backward and autograd;
• use the existing _run_compiled launcher cache in the direct backward tests, including an assertion that the second fused residual-gradient case reuses the compiled callable;
• add an opt-in, correctness-checked RMSNorm and fused-add RMSNorm backward comparison against PyTorch eager.

This deliberately does not add kernel tuning knobs or another production cache. Production already has device-aware caches, and structural RMSNorm autotuning remains in #785.

Benchmark contract

Run with:

ROCDSL_COMPARE_TORCH=1 \
  python3 -m pytest \
  tests/kernels/test_rmsnorm.py::test_rmsnorm_backward_torch_benchmark \
  -q -s

Optional shapes use ROCDSL_RMSNORM_BWD_BENCH_SHAPES="M,N,dtype;...". This is intentionally separate from the existing forward benchmark override.

The comparison is public-autograd backward vs public-autograd backward. Forward graph construction and FlyDSL compilation happen before timing; backward allocations, the mandatory fp32 dweight zero, and the fp32-to-parameter-dtype cast are included. Each shape compares gradients before recording events. Timings below are hot-cache GPU-event averages with 10 warmups and 100 iterations on one AMD Instinct MI355X, PyTorch 2.9.1+rocm7.2.0.git7e1940d4.

Torch/FlyDSL < 1 means FlyDSL is slower.

op	shape	dtype	FlyDSL (us)	Torch (us)	Torch/FlyDSL
rmsnorm_bwd	1x4096	bf16	85.4	42.8	0.50x
fused_add_bwd	1x4096	bf16	98.6	83.2	0.84x
rmsnorm_bwd	128x4096	bf16	92.2	57.5	0.62x
fused_add_bwd	128x4096	bf16	114.5	76.7	0.67x
rmsnorm_bwd	4096x4096	bf16	157.2	67.7	0.43x
fused_add_bwd	4096x4096	bf16	160.2	82.4	0.51x
rmsnorm_bwd	512x2048	bf16	81.2	48.5	0.60x
fused_add_bwd	512x2048	bf16	100.3	64.0	0.64x
rmsnorm_bwd	512x3000	f32	77.6	48.6	0.63x
fused_add_bwd	512x3000	f32	94.8	63.4	0.67x
rmsnorm_bwd	1024x8192	bf16	81.1	47.5	0.59x
fused_add_bwd	1024x8192	bf16	102.0	74.7	0.73x

These data establish a reproducible baseline and show that the current public backward paths still have optimization headroom; they do not by themselves isolate the bottleneck.

Validation

• black --check tests/kernels/test_rmsnorm.py
• ruff check tests/kernels/test_rmsnorm.py
• python3 -m py_compile tests/kernels/test_rmsnorm.py
• RMSNorm non-large, non-multi-GPU, non-benchmark suite: 12 passed
• multi-GPU/device-cache regressions: 2 passed
• opt-in Torch benchmark, six shapes x plain/fused: 1 passed / 12 rows
• verified the forward ROCDSL_RMSNORM_SHAPES override does not alter backward/autograd configs

Related review threads: compiled launcher reuse, more config coverage.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.