[luciusacarr]

Fixes #4759

In the single-row forward and vjp kernels, n = axis_size - lid * N_READS is negative for threads after the end of the row. The fill loop wrote out of bounds because of this, up to 240 iterations. This keeps the whole SIMD group waiting. Letting n = max(int(axis_size) - int(lid * N_READS), 0) in both kernels keeps behavior as intended without undefined behavior, without slowdowns.

On an M2 Pro, fp16, 2^26 values per call, in milliseconds:
Forward (mx.fast.layer_norm)

row	main	this PR	speedup	rms_norm
16	236.82	20.45	11.6×	14.95
64	48.01	5.20	9.2×	3.61
256	1.64	1.66	1.0×	1.67
264	15.95	2.22	7.2×	1.69

Backward (grad w.r.t. x, w, b)

row	main	this PR	speedup
16	251.70	33.01	7.6×
64	53.82	10.55	5.1×
256	6.18	6.27	1.0×
264	21.26	6.91	3.1×

Outputs are identical on main and this PR across 521 row sizes, three dtypes, forward and backward.

Thanks to ArupKantiDas for flagging the issue and creating a fix independently.

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: Used for aid in profiling and initial diagnosis.

[ArupKantiDas]

Thanks for fixing this. Your change is correct and gives the same behaviour as the fix on my branch.

I pushed it yesterday but couldn't open a PR because GitHub said “open pull request limit has been reached for this user”. #4717 was still open then and merged today.

My version clamps n at 0 in both kernels, while this PR starts the fill loop at max(n, 0). Both keep positive values unchanged. Mine also uses int casts to make the subtraction signed on purpose. Both fix the issue, so I'm happy to leave the choice to the maintainers and not open another PR if this one goes ahead.

The results may be useful here too. On an M5 Pro, using real builds of main at 56c2bf9 with and without the fix, the comparison gave bit-for-bit identical outputs across 900 arrays covering 60 row lengths, 3 dtypes, forward and all gradients. The full suite ran 1,005 tests and reported OK, with 56 skipped.

For float16 rows of 16, forward went from 91.52 to 8.55 ms (10.7x) and backward from 94.84 to 14.63 ms (6.5x), which is close to your M2 Pro results.

[luciusacarr]

Thanks for the heads up, apologies for that. I went ahead and added you as a co-author, and adopted your method for clamping n, since I believe it serves as a more robust solution. This should keep your PR slot open, so you have the ability to make any other PRs if you wish, as an added bonus.

I am also fine with closing this one, if you wish to make a PR off your own branch and have it be your name alone.

[ArupKantiDas]

Thanks for adding me as a co-author and adopting the clamp. I'm happy to be a co-author on the commit and let #4772 go ahead, so there's no need to close it or for me to open a separate PR.