Fix silent Metal grid truncation in softmax, logsumexp, rms_norm and layer_norm past 2^27 rows

[ArupKantiDas]

Issue

For running operations like softmax, logsumexp, rms_norm and layer_norm on the Mac GPU, each row used to get 1 threadgroup launched in one long grid, set up one after the other.
The grid passed 2^32 threads at somewhere around 2^27 rows (of 128 numbers or fewer).
Metal cuts the grid off silently, because of which most of the rows were never even computed and came back as 0.
The layer_norm gradient and the long-row rms_norm gradient have the exact same problem.

Fix

The new helper in utils.h named dispatch_threadgroups_in_slices splits the oversized launches into slices that actually fit, and each slice then binds the input and output at its first row, so the GPU kernels stay unchanged.
The inputs that already fit still use exactly one launch, which is the same idea as the reductions fix done in #4437.
The row counts changed from int to int64_t, and the rms_norm gradient for short rows is also untouched, since it already uses at most about 512 teams.

Test and Reproduction

First I reproduced the issue on my M5 Pro, 0.32.4, at the following boundaries: [2^27 - 1] rows correct, [2^27] rows all wrong, [2^27 + 1000] rows with only first 1000 correct.
Then 134M and 402M random rows compared against MLX's CPU result, of which there were 0 mismatches after the fix.
Added 3 new tests (softmax/logsumexp, rms_norm/layer_norm, layer_norm gradient). They fail on 0.32.4 and pass with the fix applied.
The gradient test is skipped in CI because it uses about 6.5 GB of memory.
The full suite ran with 996 of them, out of which 56 were skipped.

Fixes #4624

Benchmarks

Milliseconds per call, float16 inputs, both main and the fix built with Xcode, measured on my M5 Pro.

shape	op	main (ms)	fix (ms)
(16384, 4096)	softmax	1.216	1.207
(16384, 4096)	layer_norm	1.218	1.200
(262144, 1024)	rms_norm	4.184	4.175
(4194304, 64)	logsumexp	4.434	4.414

The cross_entropy and the attention fallback use the same launch pattern and are not fixed in this PR. Can be taken up in a later PR.

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: Claude Code helped do the repro, write the fix and the tests on my guidance. I directed the work and reviewed the result.

[zcbenz]

Would it worth replacing the dispatch_col_reduce introduced in #4437 with dispatch_threadgroups_in_slices?