[ArupKantiDas]

Issue

fast.cross_entropy has the possibility of silently skipping rows when the Metal grid gets too large. This is the same kind of issue as #4624, fixed in #4717.

The launch uses 1 threadgroup per row and dispatches n_rows * threadgroup_size threads in a single grid. Metal truncates the grid dimension past UINT32_MAX without reporting an error. With up to 128 classes, the threadgroup has 32 threads, so this happens at 2^27 rows. Rows beyond the wrapped grid are never written, leaving those outputs with their previous buffer contents.

Fix

This uses dispatch_threadgroups_in_slices for forward and VJP, the same helper used by softmax since #4717. Each slice binds the buffers at its starting row, so the kernels don't need any changes.

n_rows is now int64_t and cross_entropy_dims becomes cross_entropy_threadgroup_size since it only needs to return the threadgroup size. Below the limit there is still 1 slice and 1 dispatch.

Test and reproduction

AI Assistant ran the comparison on an M5 Pro with 48 GB, using distinct random rows with seed 0, random targets and 8 classes. The reference was calculated in NumPy float64 in chunks of 2^22 rows. A row was counted as wrong if the error was greater than 1e-3.

At 2^27 - 16 rows, float32 had 0 wrong rows in forward and VJP on PyPI 0.32.4, a clean build of main at 8329918 and the fix.

At 2^27 + 2^20 rows:

• PyPI 0.32.4 and clean main both had 134,217,728 wrong rows in forward and VJP for float32. The first wrong row was 1,048,576, which matches the wrapped grid computing only the first 2^20 rows.
• The fix had 0 wrong rows. Maximum error was 7.7e-7 for loss and 4.6e-7 for the gradient.
• For float16, main had 134,217,728 wrong rows in both forward and VJP. The fix had 0 wrong rows.

The full suite on the cross_entropy-only commit ran 1,010 tests and reported OK, with 56 skipped. pre-commit also passes.

No new test is included because the 2^27-row tests in #4717 were removed as too slow. I can add one if wanted.

Benchmarks

These compare a clean build of main at 8329918 against the fix on the same base. Timings are milliseconds, using the median of 3 alternating runs per build. Each run takes the median of 7 batches of 20 iterations.

shape	dtype	forward main	forward fix	forward + backward main	forward + backward fix
1048576 x 8	float32	1.472	1.469	2.403	2.394
1048576 x 128	float32	2.162	2.145	6.359	6.364
65536 x 1024	float16	0.811	0.807	1.851	1.854
8192 x 32000	bfloat16	2.042	2.061	6.123	6.073
2048 x 151936	bfloat16	2.443	2.430	7.341	7.263

The largest difference is 1.06%, with the fix faster in that case. The 3 runs of a single build varied by up to 4.7%.

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: AI Assistant wrote the code, reproduction script and benchmarks and ran the checks above. I directed the work, chose the scope and reviewed it.