Feat/strict qk norm reduction

[leonardHONG]

Depends on #339. Please merge #339 before this PR.

Summary

• Replace the per-row Python reduction in RMSNorm dweight with a single CUDA kernel.
• Preserve the exact ascending FP32 left-fold order with __fadd_rn, without atomics or reassociation.
• Share the CUDA fast path across the CUDA and Triton RMSNorm backward implementations while keeping the existing CPU and ROCm fallbacks.

The reduction launch count is reduced from rows to 1.

Performance

The following measurements time one complete FP32 row-reduction invocation. RMSNorm contribution generation and input preparation are excluded. Rows is the flattened product of all leading dimensions; for QK-Norm, these are typically batch, token, and head dimensions.

H100 80GB, FP32, PyTorch 2.9.1+cu128, CUDA 12.8. Each value is the median synchronized wall time across seven repeats.

QK-Norm head dimension

columns=128:

Rows	Previous reduction	New reduction	Speedup
1	0.0180 ms	0.0087 ms	2.1×
2	0.0250 ms	0.0087 ms	2.9×
4	0.0388 ms	0.0088 ms	4.4×
8	0.0658 ms	0.0087 ms	7.6×
16	0.1192 ms	0.0087 ms	13.8×
32	0.2258 ms	0.0088 ms	25.6×
64	0.4471 ms	0.0116 ms	38.5×
128	0.8653 ms	0.0214 ms	40.4×
256	1.7339 ms	0.0412 ms	42.1×
512	3.4341 ms	0.0805 ms	42.7×
1,024	6.8643 ms	0.1589 ms	43.2×
2,048	13.8142 ms	0.3167 ms	43.6×
4,096	27.5826 ms	0.6308 ms	43.7×
8,192	54.9083 ms	1.2598 ms	43.6×
16,384	109.7944 ms	2.5178 ms	43.6×

Model hidden dimension

columns=4096:

Rows	Previous reduction	New reduction	Speedup
1	0.0181 ms	0.0088 ms	2.1×
8	0.0658 ms	0.0088 ms	7.5×
32	0.2279 ms	0.0088 ms	25.8×
128	0.8682 ms	0.0224 ms	38.8×
512	3.4784 ms	0.0837 ms	41.6×
1,024	6.9332 ms	0.1657 ms	41.8×
2,048	13.9128 ms	0.4558 ms	30.5×
4,096	27.6102 ms	1.5472 ms	17.9×

Column coverage

rows=1024:

Columns	Previous reduction	New reduction	Speedup
64	6.9214 ms	0.1592 ms	43.5×
128	6.8643 ms	0.1589 ms	43.2×
256	6.9481 ms	0.1601 ms	43.4×
1,024	6.9287 ms	0.1625 ms	42.6×
2,048	6.9533 ms	0.1631 ms	42.6×
4,096	6.9332 ms	0.1657 ms	41.8×
8,192	6.8612 ms	0.2178 ms	31.5×

Validation

• Bitwise equal to the original ascending FP32 acc = acc + row loop in all 28 benchmark configurations, with max_abs_error = 0.
• Preserves the non-associative left-fold order without atomics, tree reductions, or CUB reductions.
• Covers FP16, BF16, and FP32 CUDA inputs, empty boundaries, non-contiguous inputs, multidimensional rows, and the CPU fallback.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 873bd947-820b-4039-9fb7-dfaa022ec740

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}