ciflow/trunk/183864: Fix LayerNorm backward kernel for AMD Strix Halo GPUs

LayerNorm/RMSNorm's gamma beta backward kernel produced wrong results for

AMD Strix Halo and other RDNA GPUs.

It was caused by a hardcoded warp size (incorrect for RDNA/Strix Halo),

and mixed usage of thread block and warp coordinates. This caused half of the

values (along the x-axis/features) being dropped when saving the gradients

after partial sum reduction.

The fix is to query and use correct warp size instead of the hardcoded values.

• No change to the working AMD and CUDA code path.


• This will bring Strix Halo/RNDA to the proven/performant CUDA code path.

Issue reproduced and tested on local Strix Halo hardware using:

$python test/test_nn.py TestNN.test_layer_norm_backwards_eps -v

See issue #183861