Repository navigation
Fix Metal grid truncation in softmax/logsumexp/rms_norm/layer_norm past 2^27 rows - #4717
Merged
zcbenz merged 3 commits intoOct 26, 2026
Merged
Conversation
zcbenz
added
the
await response
Contributor
Author
|
Yes, it is, but only after one small extension: optional y/z grid sizes, default to 1. The reduction's offset, which is set from the callback, stays as a kernel argument because it needs the absolute column for its bounds check. The dispatch_col_reduce goes away, and its 4 callers in reduce.cpp use the shared helper. Should I add it to this PR, or as a follow-up if you'd rather keep this one small. |
Member
|
Please add it to this PR, thanks for looking into it! |
ArupKantiDas
force-pushed
the
fix-row-kernel-grid-overflow
branch
from
October 24, 2026 05:24
8c08681 to
6ca0320
Compare
Contributor
Author
|
Thanks for the confirmation. It's done, I've added the change to this PR as a separate commit and also rebased onto the latest Reduction benchmarks (float16, ms per call, M5 Pro):
|
zcbenz
removed
the
await response
zcbenz
force-pushed
the
fix-row-kernel-grid-overflow
branch
from
October 25, 2026 23:46
6ca0320 to
6bbebd9
Compare
softmax, logsumexp, rms_norm and layer_norm (and the layer_norm VJP and the looped rms_norm VJP) launch one threadgroup per row in a 1D grid. Once rows * threadgroup size passes 2^32, Metal truncates the grid width, so most rows are never computed and silently come back as zeros. Launch the rows in slices that each fit in one grid, and bind the input and output buffers at each slice's first row. Launches that already fit are unchanged and still use a single dispatch. Fixes ml-explore#4624
dispatch_col_reduce from ml-explore#4437 does the same slicing as the new helper. Give the helper optional grid_y and grid_z sizes, use it for the four column reduction dispatches, and remove dispatch_col_reduce. The reductions still set each slice's start as the column_offset kernel argument (from the bind callback), because the kernel needs the absolute column for its bounds check.
zcbenz
force-pushed
the
fix-row-kernel-grid-overflow
branch
from
October 26, 2026 00:32
6bbebd9 to
f10f293
Compare
zcbenz
approved these changes
Oct 26, 2026
zcbenz
changed the title
This was referenced Oct 26, 2026
Closed
inureyes
added a commit
to lablup/mlxcel
that referenced
this pull request
Oct 26, 2026
…2627) ## Summary The ROCm per-row launchers (softmax, logsumexp, rms_norm plain/strided/VJP, layer_norm forward/VJP) put one block per row in grid x with an `int` `n_rows`. HIP refuses `gridDim.x * blockDim.x >= 2^32`, so 2^24 + 1 rows at 256 threads failed to launch. This is the ROCm side of ml-explore/mlx#4717. - Each launcher caps grid x at `UINT32_MAX / BLOCK_DIM` and spills into y (the `arg_reduce.hip` pattern of LOCAL_FIXES item 11). Kernels take `n_rows`, compute a 64-bit row from `blockIdx.x + blockIdx.y * gridDim.x`, and return past it before any barrier. Below the cap the grid is `dim3(n_rows)` as before. - The strided rms_norm kernel decomposes its row in 64 bits. - LOCAL_FIXES item 62 (61 was taken by #2624 while this was in flight). - Two test-only FFI entries, `fast_rms_norm_vjp_for_test` and `fast_layer_norm_vjp_for_test`, run the VJP primitives through `mlx::core::vjp`; the crate had no VJP entry point. ## Tests (gfx1151) New `row_grid` module in `tests/rocm_hip_grid_clamp.rs`: seven launches on 2^24 + 1 x 4 f32, every element compared on the GPU with an f64 host reference, plus `#[ignore]`d tests for the 512 and 1024-thread softmax arms (about 17 GiB each). Before the fix (overlay reverted, same tests): `0 passed; 7 failed; 2 ignored`, each with: ``` GPU launch failed: hipLaunchKernel( func, grid_dim, block_dim, params, smem_bytes, stream_) failed: invalid configuration argument ``` After: `7 passed; 0 failed; 2 ignored; finished in 7.12s`. The whole `rocm_hip_grid_clamp` file after rebasing onto #2624: `20 passed; 0 failed; 7 ignored`. A CPU-stream reference was tried first and dropped: about 5 minutes a test from host reads of device memory, and its f32 sum over 2^24 rows for the layer_norm weight gradient was 2.7e5 off the GPU value, which is within 1.7e3 of the f64 sum. The weight gradients use rtol 1e-3 for that reason; the per-row `dx` checks that every row ran. Existing tests, post-fix, under `rocm_gpu_guard.sh --hold`: mlxcel-core lib filters `softmax` (10), `rms_norm` (15), `layer_norm` (2), `fused_norm_parity` (10), all pass. `make verify-llama-compat verify-env-var-docs verify-kernel-port-dispatch verify-kernel-dtype-keys verify-rocm-overlay` and clippy (`-p mlxcel-core --lib --tests`, `-p mlxcel --test rocm_hip_grid_clamp`) are clean. Not run: the two ignored softmax arm tests (17 GiB on a 31 GiB shared host), `layer_norm_multiwave_mean_snapshot_is_stable` (needs the `mlxcel` lib test build). Closes #2619
This was referenced Oct 26, 2026
zcbenz
pushed a commit
to ArupKantiDas/mlx
that referenced
this pull request
Oct 31, 2026
The forward and VJP kernels run one threadgroup per row as a single grid of n_rows * threadgroup_size threads. Metal truncates a grid dimension past UINT32_MAX without an error, so with a small vocabulary (32-thread threadgroups) rows past 2^27 were never written. Use dispatch_threadgroups_in_slices like softmax and the norm kernels (ml-explore#4717), binding each buffer at its row offset. n_rows is now int64_t.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.hnameddispatch_threadgroups_in_slicessplits 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
inttoint64_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.
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.