[apullin]

Summary

Major speedup for LayerNorm on small feature shapes via tiling compute. Origin is in my own application of liger_kernel training very small (edge) networks. Working well for me on H100, and appears to be a general win, though open to redesign or scope limitation if needed.

• One program per row starves the GPU on narrow widths (16-256 channels with many rows). Widths up to 256 columns (BLOCK_SIZE <= 256, n_rows >= 4096) now process up to 64 rows per 4,096-element tile (_block_layer_norm_forward_kernel / _block_layer_norm_backward_kernel), and the backward kernel runs up to four persistent programs per SM on that path.
• Wider rows (> 256 columns) and small row counts (< 4,096 rows) keep the existing single-row path unchanged, and row_mode=True forces the single-row path (matching ops/rms_norm.py).

Benchmark (NVIDIA H100 80GB HBM3, 7,864,320 rows, bfloat16, triton.testing.do_bench median ms)

hidden	fwd single-row	fwd block-row	fwd speedup	bwd single-row	bwd block-row	bwd speedup	total single-row	total block-row	total speedup
16	4.728	0.192	24.58x	25.824	0.682	37.84x	30.552	0.875	34.93x
32	4.732	0.358	13.22x	33.009	1.136	29.07x	37.741	1.494	25.27x
64	4.739	0.689	6.88x	43.803	2.572	17.03x	48.542	3.260	14.89x
128	4.900	1.351	3.63x	45.531	5.558	8.19x	50.431	6.908	7.30x
256	5.287	2.677	1.98x	48.364	16.700	2.90x	53.652	19.377	2.77x
512 (ctrl)	6.272	6.295	1.00x	55.784	55.655	1.00x	62.055	61.949	1.00x

Testing Done

• Hardware Type: NVIDIA H100 80GB HBM3
• run make test to ensure correctness (5325 passed, 1146 skipped, 21 xfailed, including test_liger_layer_norm_block_row across ragged row counts and odd hidden sizes)
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[apullin]

There is also a known bug in LigerLayerNorm bias handling, I can follow up with another PR after this one.