Anatoliylitv/tiled and 2 pass layer norm

[anatoliylitv]

Hugging Face model dropped 5-10% after switching to tiled kernel (10+ model and tests).

Implemented mixed approach, using combination of Tiled and Two Pass:

Hugging Face (huggingface_bart) performance is back:
Legacy Two pass performance=1652, 1644
Tiled performance =1574, 1568
Mixed performance =1649, 1652

Synthetic reproducer where performance from chess board like became uniform, keeping benefits from both implementation:

Layer Norm Gamma/Beta Backward — Benchmark Summary

Device: AMD Instinct MI350X (ROCm, warp_size=64)
Warmup: 20 · Iters: 100 · all times in µs, averaged over 3 runs
Legacy: cuComputePartGradGammaBeta + cuComputeGradGammaBeta (original two-pass, M ≥ 128)
Tiled: LaunchGammaBetaBackwardCUDAKernel (single-pass tiled, M ≥ 128)
Mixed: new dispatcher — picks legacy for 128 ≤ M ≤ 65536, tiled for M > 65536 && N/64 < sm_count/2
(M < 128 always uses GammaBetaBackwardSimpleCUDAKernel in all three runs)

Name	Op	Shape (M×N)	dtype	(Legacy) Two pass (µs)	Tiled (µs)	Mixed (µs) Tiled & Two pass	Mixed	vs Tiled
tile8_small	layer_norm	32×512	fp16	7.89	7.93	8.00	simple (M<128)	≈0%
tile64_medium	layer_norm	96×768	fp16	20.69	20.73	20.78	simple (M<128)	≈0%
tile128_medium	layer_norm	192×1024	bf16	7.86	8.14	12.53	legacy ⚠️	−54%
tile256_large	layer_norm	4096×1024	fp16	12.73	31.06	12.72	legacy	+59%
tile256_bert	layer_norm	1024×768	fp16	12.70	11.04	12.54	legacy ⚠️	−14%
tile256_large_bf16	layer_norm	4096×1024	bf16	12.65	31.25	12.91	legacy	+59%
tile256_large_fp32	layer_norm	4096×1024	fp32	15.81	30.88	15.93	legacy	+48%
two_pass_huge_M	layer_norm	131072×64	fp16	114.36	26.66	26.79	tiled	≈0%
llm_hidden_4096	layer_norm	131072×4096	bf16	950.26	513.72	513.33	tiled	≈0%
gpt2_style	layer_norm	8192×768	fp16	16.02	58.02	16.79	legacy	+71%
rms_tile256	rms_norm	4096×1024	fp16	13.03	22.34	11.72	legacy	+48%
rms_llm	rms_norm	131072×4096	bf16	927.19	477.92	477.17	tiled	≈0%
rms_two_pass	rms_norm	131072×64	fp16	109.44	16.60	16.84	tiled	≈0%

vs Tiled = positive means mixed is faster than pure tiled.

Observations

Wins (mixed picks the right kernel):

• tile256_large / _bf16 / _fp32 (M=4096): ~2.4× faster than tiled
• gpt2_style (M=8192): ~3.5× faster than tiled
• rms_tile256 (M=4096): ~1.9× faster than tiled
• two_pass_huge_M, llm_hidden_4096, rms_llm, rms_two_pass (M=131072): no regression vs tiled, massive win vs legacy (4–7×)

cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @jataylo @hongxiayang @naromero77amd @pragupta @jerrymannil @xinyazhang

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[rocm-repo-management-api]

Jenkins build for 20a3ce6fac4bcf46d9846f98270b85ce3346efc5 commit finished as FAILURE
Links: Pipeline Overview / Build artifacts / Test Results