perf(norm): skip dW/dB calculation when parameters are frozen (#1067)

[piyushumate]

Description

Addresses #1067 (previously explored in #1068)

This PR brings the LoRA normalization optimization across the finish line:

• Rebased on top of modern main (cleanly resolving the [Test]: Refactor benchmark_geglu with standardized model configs #1116 benchmark conflicts that blocked Skip norm weight/bias gradients when frozen (LoRA/PEFT optimization) #1068).
• Incorporates maintainer feedback from @Tcc0403 on Skip norm weight/bias gradients when frozen (LoRA/PEFT optimization) #1068 (passes None directly for unused gradient pointers).
• Comprehensive unit tests: 24/24 passing on Float32 and BFloat16.
• Hardware-validated on NVIDIA A10G and NVIDIA A100 GPUs (up to 1.38x backward speedup).

Full co-author credit is preserved for @yukiu00 in the commit.

---

Motivation

In parameter-efficient fine-tuning (LoRA / PEFT), normalization parameters (weights and biases) are frozen (requires_grad=False). Previously, Liger-Kernel unconditionally:

1. Allocated intermediate reduction buffers (_dW, _DB of size sm_count * n_cols) in HBM.
2. Allocated registers and ran Triton accumulation loops across tokens.
3. Wrote intermediate gradients to global HBM.
4. Launched PyTorch .sum(dim=0) global reductions to produce tensors that PyTorch autograd immediately discarded.

---

Key Changes

• rms_norm.py: Added COMPUTE_DW: tl.constexpr to _rms_norm_backward_kernel and _block_rms_norm_backward_kernel. Prunes accumulation and memory writes when compute_dw=False; skips _dW buffer allocation and reduction; queries ctx.needs_input_grad[1].
• layer_norm.py: Added COMPUTE_DW and COMPUTE_DB flags to prune weight/bias gradients independently based on ctx.needs_input_grad[1] and [2].
• fused_add_rms_norm.py: Added COMPUTE_DW to skip _dW allocation and reduction based on ctx.needs_input_grad[2].
• test/ops/test_norm_lora_parity.py: Added 24 unit tests verifying numerical parity (FP32 & BF16), checking that w.grad is None / b.grad is None when frozen, and confirming exact x.grad equality.
• benchmark/scripts/benchmark_norm_lora.py: Added dedicated standalone LoRA backward microbenchmark.

---

Benchmark Results

NVIDIA A100-SXM4-40GB (triton.testing.do_bench, bfloat16)

Operation	Shape (B, D)	Full FT Bwd (ms)	LoRA Bwd (ms)	Speedup
LayerNorm Bwd	(2048, 4096)	0.0806	0.0585	1.38x (+38%)
RMSNorm Bwd	(2048, 4096)	0.0871	0.0707	1.23x (+23%)
LayerNorm Bwd	(4096, 4096)	0.1217	0.0991	1.23x (+23%)
RMSNorm Bwd	(4096, 4096)	0.1305	0.1177	1.11x (+11%)
LayerNorm Bwd	(4096, 8192)	0.1966	0.1729	1.14x (+14%)
RMSNorm Bwd	(4096, 8192)	0.2275	0.2135	1.07x (+7%)
LayerNorm Bwd	(8192, 4096)	0.2070	0.1843	1.12x (+12%)
RMSNorm Bwd	(8192, 4096)	0.2384	0.2243	1.06x (+6%)
LayerNorm Bwd	(8192, 8192)	0.3514	0.3255	1.08x (+8%)
RMSNorm Bwd	(8192, 8192)	0.4235	0.4077	1.04x (+4%)

NVIDIA A10G (triton.testing.do_bench, bfloat16)

• LayerNorm Bwd: 1.05x – 1.18x speedup across shapes.
• RMSNorm Bwd: 1.02x – 1.07x speedup across shapes.

---

Test Validation

• test/ops/test_norm_lora_parity.py: 24 / 24 PASSED (100%)
• Existing ops correctness: 108 / 108 PASSED (100%)
• make checkstyle (ruff check & ruff format): PASSED

[piyushumate]

@vaibhavjindal lemme know your thoughts and if any additional work is needed on this one. Thank you!