[release/2.13] [ROCm] Optimize AMD normalization backward kernel by using tiled and …

[anatoliylitv]

https://amd-hub.atlassian.net/browse/ROCM-24994

Upstream PR: pytorch#189405

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 Backward Benchmark: Tiled only vs Tiled+Two pass

Device: AMD Instinct MI350X | Warmup: 20 | Iters: 100 | Runs averaged: 3

Summary (avg µs over 3 runs)

Benchmark	Op	Shape	dtype	Branch	Tiled only	Tiled+Two pass	Δ (µs)	Speedup	Winner
tile8_small	layer_norm	(32, 512)	float16	Tile-8	7.94	8.16	-0.22	0.97×	~tie
tile64_medium	layer_norm	(96, 768)	float16	Tile-64	20.71	20.78	-0.07	1.00×	~tie
tile128_medium	layer_norm	(192, 1024)	bfloat16	Tile-128	8.20	8.51	-0.31	0.96×	~tie
tile256_large	layer_norm	(4096, 1024)	float16	Tile-256	31.07	13.12	+17.95	2.37×	Two pass
tile256_bert	layer_norm	(1024, 768)	float16	Tile-256 BERT	11.15	11.05	+0.10	1.01×	~tie
tile256_large_bf16	layer_norm	(4096, 1024)	bfloat16	Tile-256	31.28	13.10	+18.18	2.39×	Two pass
tile256_large_fp32	layer_norm	(4096, 1024)	float32	Tile-256	30.87	15.90	+14.97	1.94×	Two pass
two_pass_huge_M	layer_norm	(131072, 64)	float16	Two-pass M-parallel	26.48	27.70	-1.22	0.96×	~tie
llm_hidden_4096	layer_norm	(32, 4096, 4096)	bfloat16	Tile-256 LLM	509.65	513.38	-3.73	0.99×	~tie
gpt2_style	layer_norm	(8, 1024, 768)	float16	Tile-256 GPT2	58.22	16.56	+41.66	3.52×	Two pass
rms_tile256	rms_norm	(4096, 1024)	float16	Tile-256 rms	22.34	11.18	+11.16	2.00×	Two pass
rms_llm	rms_norm	(32, 4096, 4096)	bfloat16	Tile-256 rms LLM	476.18	476.03	+0.15	1.00×	~tie
rms_two_pass	rms_norm	(131072, 64)	float16	Two-pass rms	16.73	17.29	-0.56	0.97×	~tie

Δ = Tiled only − Tiled+Two pass (negative = Two pass faster). Speedup = Tiled only / Tiled+Two pass.

Also result of generated test for correcteness:
tests/test_correctness.py::test_layer_norm_gamma_beta_backward_matches_cpu_reference[gpt2_style] PASSED [ 4%]
tests/test_correctness.py::test_edge_case_shape_M_eq_1 PASSED [ 8%]
tests/test_correctness.py::test_rms_norm_gamma_backward_matches_cpu_reference[rms_llm] PASSED [ 12%]
tests/test_correctness.py::test_layer_norm_gamma_beta_backward_matches_cpu_reference[tile8_small] PASSED [ 16%]
tests/test_correctness.py::test_noncontiguous_input_matches_reference PASSED [ 20%]
tests/test_correctness.py::test_layer_norm_gamma_beta_backward_matches_cpu_reference[two_pass_huge_M] PASSED [ 25%]
tests/test_correctness.py::test_layer_norm_autograd_matches_cpu_reference PASSED [ 29%]
tests/test_correctness.py::test_layer_norm_gamma_beta_backward_matches_cpu_reference[tile256_large_bf16] PASSED [ 33%]
tests/test_correctness.py::test_output_mask_selects_expected_grads[dgamma_and_dbeta] PASSED [ 37%]
tests/test_correctness.py::test_configs_cover_all_tile_branches PASSED [ 41%]
tests/test_correctness.py::test_layer_norm_gamma_beta_backward_matches_cpu_reference[tile128_medium] PASSED [ 45%]
tests/test_correctness.py::test_rms_norm_autograd_matches_cpu_reference PASSED [ 50%]
tests/test_correctness.py::test_edge_case_shape_N_eq_1 PASSED [ 54%]
tests/test_correctness.py::test_layer_norm_gamma_beta_backward_matches_cpu_reference[tile256_bert] PASSED [ 58%]
tests/test_correctness.py::test_output_mask_selects_expected_grads[dbeta_only] PASSED [ 62%]
tests/test_correctness.py::test_layer_norm_gamma_beta_backward_matches_cpu_reference[tile256_large] PASSED [ 66%]
tests/test_correctness.py::test_rms_norm_gamma_backward_matches_cpu_reference[rms_tile256] PASSED [ 70%]
tests/test_correctness.py::test_layer_norm_gamma_beta_backward_matches_cpu_reference[llm_hidden_4096] PASSED [ 75%]
tests/test_correctness.py::test_layer_norm_gamma_beta_backward_matches_cpu_reference[tile256_large_fp32] PASSED [ 79%]
tests/test_correctness.py::test_rms_norm_gamma_backward_matches_cpu_reference[rms_two_pass] PASSED [ 83%]
tests/test_correctness.py::test_output_mask_selects_expected_grads[dgamma_only] PASSED [ 87%]
tests/test_correctness.py::test_layer_norm_gamma_beta_backward_matches_cpu_reference[tile64_medium] PASSED [ 91%]
tests/test_huge_m_buffer_bug.py::test_layer_norm_huge_M_multidim_normalized_shape_oob XFAIL (ROCM-24994: LaunchGammaBetaBackwardCUDAKernel's huge-M branch sizes its dgamma_blocks/dbeta_blocks scratch buffers from dgamma->size(-1) instead of ...) [ 95%]
tests/test_huge_m_buffer_bug.py::test_layer_norm_huge_M_undefined_gamma_dbeta_shape XFAIL (ROCM-24994: LaunchGammaBetaBackwardCUDAKernel's huge-M branch sizes its dgamma_blocks/dbeta_blocks scratch buffers from dgamma->size(-1) instead of N ...) [100%]

============================================================================================================== 22 passed, 2 xfailed in 14.14s ===============================================================================================================
Used AI assistance from Cursor.

Short performance reproducer:
layer_norm_gamma_beta_backward_reproducer.py
run.sh
Pull Request resolved: pytorch#189405
Approved by: https://github.com/jeffdaily

[anatoliylitv]

Test Report: LayerNorm ROCm Commits

Date: 2026-08-14
Repo: /myworkspace/pytorch
Build: 2.13.0+git26e8e73 (HIP 7.14.60850) — new build, rebased onto release/2.13
Device: AMD Instinct MI350X

Previous run (2.12.0+git1526dfd) results are superseded by this run. See "Build History" below.

Commits Under Test

Same two fixes, now present as re-committed/cherry-picked equivalents on the current HEAD (26e8e738ff1) rather than the original hashes (base branch changed from main-style history to release/2.13):

Original commit	Equivalent on current HEAD	Description
14f8247	06aff09 / c3534e4	[ROCm] Fix LayerNorm backward kernel for AMD Strix Halo GPUs (pytorch#183864)
1526dfd	26e8e73 / 063ee08	[ROCm] Optimize AMD normalization backward kernel by using tiled and two pass implementations (pytorch#189405)

Verified both fixes' test coverage (test_layer_norm_gamma_beta_backward_dispatch_bands, _boundaries, _two_pass_at_huge_M, test_layer_norm_backward_undefined_gamma, test_layer_norm_backwards_eps) is present in test/test_nn.py on this build before running.

Command

cd /myworkspace/pytorch

python test/test_nn.py \
  TestNNDeviceTypeCUDA.test_layer_norm_gamma_beta_backward_dispatch_bands_cuda \
  TestNNDeviceTypeCUDA.test_layer_norm_gamma_beta_backward_dispatch_boundaries_cuda \
  TestNNDeviceTypeCUDA.test_layer_norm_backward_undefined_gamma_cuda \
  -v

python test/test_nn.py \
  TestNNDeviceTypeCUDA.test_layer_norm_gamma_beta_backward_two_pass_at_huge_M_cuda \
  -v

python test/test_nn.py TestNN.test_layer_norm_backwards_eps -v

Results

Test	Class	Status	Time
test_layer_norm_gamma_beta_backward_dispatch_bands_cuda	TestNNDeviceTypeCUDA	PASS	included in 11.578s (3 tests)
test_layer_norm_gamma_beta_backward_dispatch_boundaries_cuda	TestNNDeviceTypeCUDA	PASS	included in 11.578s (3 tests)
test_layer_norm_backward_undefined_gamma_cuda	TestNNDeviceTypeCUDA	PASS	included in 11.578s (3 tests)
test_layer_norm_gamma_beta_backward_two_pass_at_huge_M_cuda	TestNNDeviceTypeCUDA	PASS	0.532s
test_layer_norm_backwards_eps	TestNN	PASS	5.725s

Total: 5/5 passed, 0 failed, 0 skipped

Raw Output

Run 1: bands / boundaries / undefined_gamma

test_layer_norm_gamma_beta_backward_dispatch_bands_cuda (__main__.TestNNDeviceTypeCUDA.test_layer_norm_gamma_beta_backward_dispatch_bands_cuda) ... ok
test_layer_norm_gamma_beta_backward_dispatch_boundaries_cuda (__main__.TestNNDeviceTypeCUDA.test_layer_norm_gamma_beta_backward_dispatch_boundaries_cuda) ... ok
test_layer_norm_backward_undefined_gamma_cuda (__main__.TestNNDeviceTypeCUDA.test_layer_norm_backward_undefined_gamma_cuda) ... ok

----------------------------------------------------------------------
Ran 3 tests in 11.578s

OK

Run 2: huge-M two-pass

test_layer_norm_gamma_beta_backward_two_pass_at_huge_M_cuda (__main__.TestNNDeviceTypeCUDA.test_layer_norm_gamma_beta_backward_two_pass_at_huge_M_cuda) ... ok

----------------------------------------------------------------------
Ran 1 test in 0.532s

OK

Run 3: test_layer_norm_backwards_eps

test_layer_norm_backwards_eps (__main__.TestNN.test_layer_norm_backwards_eps) ... ok

----------------------------------------------------------------------
Ran 1 test in 5.725s

OK

Build History

Run	Build	HEAD commit	Result
1 (previous)	2.12.0+git1526dfd	1526dfd	5/5 PASS
2 (this run)	2.13.0+git26e8e73	26e8e73	5/5 PASS

Notes on the rebuild:

• HEAD moved from a main-based checkout to a release/2.13-based checkout; the same two PRs ([ROCm] Fix LayerNorm backward kernel for AMD Strix Halo GPUs pytorch/pytorch#183864, [ROCm] Optimize AMD normalization backward kernel by using tiled and two pass implementations where they are most efficient. pytorch/pytorch#189405) exist as re-committed equivalents (git log --oneline --all | grep finds three copies of each PR title across branches: main, an intermediate branch, and release/2.13).
• No content drift was found in the relevant test methods between the two builds; timings differ slightly (11.578s vs 9.097s for the 3-test batch; 5.725s vs 3.220s for test_layer_norm_backwards_eps) which is normal run-to-run/GPU-contention variance, not a regression.

Notes

• No PYTORCH_TEST_WITH_ROCM or other ROCm-specific environment variable was required; TEST_WITH_ROCM is auto-detected from torch.version.hip.
• Device-type test methods require the _cuda suffix when selected by full name (e.g. TestNNDeviceTypeCUDA.test_foo_cuda), since instantiate_device_type_tests appends the device suffix to @onlyCUDA methods.
• No source changes were made in this session; this was a test-only verification run against the new build.