rmsnorm bwd: make blocking math more robust

[v0i0]

Stack from ghstack (oldest at bottom):

• -> rmsnorm bwd: make blocking math more robust #532

[xuzhao9]

Just curious, can you give more details on why it makes block math more robust?

[v0i0]

Just curious, can you give more details on why it makes block math more robust?

for M <= 8 * num_sm, e.g. M=512, the current code will yield BLOCK_SIZE_M=0, and then PARTIAL_SIZE's computation will cause an exception