Rewrite Triton normalization backward kernel_1 (#499)

[jlamypoirier]

Summary

Closes the backward-pass gap in layer_norm/rms_norm (issue #499). On H100 the Triton backward trailed apex and torch-compiled by ~1.1–1.6× at most hidden sizes — worst on tall-narrow shapes. kernel_1 was both the bottleneck and the source of an oversized partial-reduction buffer for kernel_2.

kernel_1 rewrite

1. Decouple the register tile from n_cols. A block_size_row × block_size_col tile grid-strides the columns, so occupancy no longer collapses as the hidden size grows. Rows wider than one chunk use a two-pass scheme (reduce the per-row corrections, then re-read to write grad_input and the partials); narrower rows stay single-pass with no re-read.
2. Bound the partial-reduction work, the way apex does. apex (both the general fused path and the hand-tuned fast path) does not avoid a second reduction kernel — it bounds the number of partial rows to a small constant via row grid-striding. Previously kernel_1 emitted one partial row per block_size_row input rows, so the buffer kernel_2 reduces grew with the row count (4096 rows at 32768×1024 → kernel_2 ran at ~10% of bandwidth). Single-pass now grid-strides the rows with a program count fixed at multi_processor_count × 2, folding many row tiles into one fp32-accumulated partial. The buffer is then independent of the row count. Two waves per SM is the measured knee — one wave starves grad_input latency-hiding; more only re-inflates kernel_2.

Config tuning

An offline sweep of the kernel_1 config space (single-pass threshold, block_size_col, block_size_row, num_warps, num_stages), validated per shape against the prior config, drives the launch heuristic:

• Single-pass is extended to wide rows. A wide row stays single-pass when a warp-saturated one-row tile spanning the whole row fits in registers — which avoids the two-pass column re-read entirely. It fits up to the block-size cap without bias and half of it with bias (bias roughly doubles live registers per element), and wins once there are enough rows to fill the SMs. This is the main lever for wide hidden sizes.
• The remaining two-pass path is tuned (wider column chunk, more warps), and num_stages is threaded through the launch.

Parameter-grad partials reduce in fp32 (the store casts to the buffer dtype) — reducing in bf16 measurably degrades the parameter gradients.

Results (H100, bf16)

Backward µs vs. the fastest competitor (apex_fast for LN, torch-compiled-max otherwise). Bold = match-or-beat. Triton match-or-beats the best alternative on 9/15 LN and 10/15 RMS shapes; apex's general fused path is beaten everywhere (1.5–3× slower on backward).

layer_norm	ours	best alt	ratio		rms_norm	ours	best alt	ratio
1024×8192	36.2	47.7	0.76×		4096×4096	46.5	62.5	0.74×
2048×4096	32.3	35.8	0.90×		4096×8192	86.6	117	0.74×
8192×1024	27.6	30.2	0.91×		2048×8192	48.2	62.2	0.77×
4096×8192	119	129	0.93×		8192×1024	24.0	30.7	0.78×
2048×8192	67.7	71.8	0.94×		2048×4096	26.8	31.6	0.85×
16384×1024	49.1	50.7	0.97×		1024×8192	28.1	32.5	0.86×
4096×4096	56.2	56.9	0.99×		32768×1024	78.5	86.9	0.90×
4096×2048	29.0	29.4	0.99×		16384×1024	43.1	47.4	0.91×
32768×1024	89.4	88.6	1.01×		16384×2048	82.0	83.8	0.98×
512×16384	41.5	39.2	1.06×		2048×16384	114	115	0.99×
8192×2048	51.5	47.8	1.08×		1024×16384	64.1	62.3	1.03×
1024×16384	77.6	70.3	1.10×		8192×4096	83.2	79.6	1.05×
16384×2048	89.4	81.6	1.10×		512×16384	34.7	31.9	1.09×
2048×16384	148	128	1.16×		8192×2048	45.9	40.6	1.13×
8192×4096	97.2	83.4	1.17×		4096×2048	25.9	22.3	1.16×

The wide hidden sizes improved up to 1.4× from the config tuning (e.g. rms_norm 4096×8192 121→87 µs) by removing the re-read where it can be afforded, with no regression elsewhere. kernel_2 is no longer a factor (2–5 µs on single-pass shapes, was up to 47 µs).

Remaining sub-parity shapes are a mix of (a) narrow shapes (n_cols ≤ 4096) where apex_fast/compiled are very tight, and (b) layer_norm at the widest hidden size (16384), where the bias term spills the wide single-pass tile so the kernel must fall back to the two-pass re-read. Closing (b) would need a bias-aware shared-memory single-pass; it is the natural follow-up.

Forward is at parity across implementations and is unchanged.

Default implementation

Now that the Triton backward is competitive with or faster than apex, NormalizationConfig.implementation = auto resolves to Triton when Triton is enabled (or required, for zero-centered weights) and falls back to PyTorch otherwise — apex is dropped from the auto path. The apex fast/fused implementations stay available via explicit selection.

Benchmark harness

tools/benchmark/triton_kernels:

• Isolated, cold-L2 backward timing (forward untimed, L2 flushed, then the backward timed) — training-representative. The prior fwd_bwd − fwd number had a warm-L2 confound: the forward left the saved output partly resident in L2, flattering the backward in a way real training never sees.
• Per-kernel device-time breakdown, so kernel_1 and kernel_2 can be attributed separately.

Validation

tests/layers/ and tests/tools/test_triton_benchmark.py: 733 passed, 27 skipped (H100). Parameter-grad precision is bit-equivalent to the previous kernel (grad_weight rel-rms ≈ 2.8–2.9e-3).

---

Authored by Claude Opus 4.8 (Claude Code).

🤖 Generated with Claude Code