feat(Part2): Fused Add + RMSNorm Triton kernel with MASE transform pass

[Dorijan9]

Fused Add + RMSNorm Triton Kernel with MASE Transform Pass

Part 2 of the ADLS kernel-fusion-aware optimisation pipeline.

Problem

Every transformer decoder layer runs residual addition and RMSNorm as two separate CUDA kernels, writing the intermediate tensor to HBM and reading it back. For a 32-layer Llama model, that's 64 wasted memory round-trips per forward pass.

Solution

A hand-written Triton kernel fuses both operations into a single GPU kernel launch, plus an FX graph transform pass that automatically pattern-matches add → RMSNorm and swaps in the fused module.

Files

File	What
src/chop/passes/graph/transforms/fused_rmsnorm/triton_fused_add_rmsnorm.py	Triton fwd/bwd kernels, autograd function, nn.Module wrapper
src/chop/passes/graph/transforms/fused_rmsnorm/fused_rmsnorm_transform.py	FX graph transform pass
src/chop/passes/graph/transforms/fused_rmsnorm/init.py	Package exports
test/passes/graph/transforms/test_fused_add_rmsnorm.py	Correctness tests + benchmarks

Tests

• Forward correctness: 144/144 (8 shapes × 3 dtypes × 3 casting modes × 2 offsets)
• Backward correctness: 36/36
• nn.Module wrapper: ✅

Supports llama, gemma, and none casting modes. Recognises all major HuggingFace RMSNorm variants (Llama, Mistral, Gemma, Qwen2).