Repository navigation
fix(rocm): split per-row softmax and norm grids past the HIP limit - #2627
Merged
Merged
Conversation
The ROCm softmax, logsumexp, rms_norm (plain, strided, VJP) and layer_norm (forward, VJP) launchers 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 on gfx1151 with "invalid configuration argument". Each launcher now caps grid x at UINT32_MAX / BLOCK_DIM and spills into y, as LOCAL_FIXES item 11 does for arg_reduce; kernels take n_rows, compute a 64-bit row from blockIdx.x and blockIdx.y and return past it. Below the cap the grid is dim3(n_rows) as before. This is the ROCm side of ml-explore/mlx#4717. Recorded as LOCAL_FIXES item 62. Tests: a row_grid module in tests/rocm_hip_grid_clamp.rs runs all seven launches on 2^24 + 1 x 4 f32 against an f64 host reference (7 failed before, 7 pass after on gfx1151), plus ignored 512 and 1024-thread softmax arm tests. Two test-only FFI entries expose the rms_norm and layer_norm VJPs through mlx::core::vjp. Closes #2619
inureyes
added
status:review
6 tasks done
inureyes
marked this pull request as ready for review
October 26, 2026 10:34
5 of 6 tasks
inureyes
added a commit
that referenced
this pull request
Oct 26, 2026
## Summary ROCm had no NVFP4 kernel, so ModelOpt NVFP4 checkpoints were converted to affine 4-bit at load and MLX-native NVFP4 planes were refused (#1806). This adds native NVFP4 (E2M1 values, E4M3 scales per 16, optional float32 global scale per expert) to `QuantizedMatmul` and `GatherQMM` and marks `(Rocm, Nvfp4)` `Native`. - New overlay file `quantized/nvfp4_qmv.hip` (included by `qmm.hip`): `nvfp4_qmv_kernel` computes eight rows per block (one for M = 1 and unsorted gather decode) and reloads a weight group only when the expert changes, so dense multi-row calls (speculative-decode verify widths) and sorted gather prefill read each group once per eight rows. It serves `QuantizedMatmul` below the dequantize + GEMM crossover and every NVFP4 `GatherQMM` call, applying `global_scale[expert] / (448 * 6)` once per output. `nvfp4_dequantize_kernel` expands a plane exactly; `select_qmm_route` admits NVFP4 to the dequantize + hipBLASLt/rocBLAS route, which covers prefill-size and non-transposed GEMMs. - `qdequant.hpp`: branch-free E4M3 decode beside the moved fp4 decode (both exact, power-of-two offsets folded into the final multiply). - Wave64: the kernels use 16-lane column groups with width-16 shuffles and never read `WARP_SIZE`, so they are lane-free and need no wavefront gate (not testable on this host). - The NVFP4 gather path returns before the mxfp4/mxfp8 group-32 check. `QuantizedMatmul` now refuses an unexpected fourth input in a non-affine mode instead of ignoring it. - `validate_native_global_scale` (switch_layers) admitted only Metal and CUDA; it now admits ROCm. Sanitize route tests, layers route-line tests and docs (installation ROCm table, env vars, supported models) follow the table change. `LOCAL_FIXES.md` item 63 (rebased onto #2627, which took 62). ## Numerical results (gfx1151, vs f32 `x @ dequantize(w).T` on CPU) Max abs error / error relative to max |reference|, worst case per path: | Path | f32 | f16 | bf16 | |---|---|---|---| | qmm qmv M=1 | 7.6e-5 / 5.2e-7 | 5.9e-2 / 4.0e-4 | 4.7e-1 / 3.1e-3 | | qmm multi-row qmv M=6,16,33 | 9.5e-5 / 6.1e-7 | 6.2e-2 / 3.6e-4 | 5.0e-1 / 2.9e-3 | | qmm dequant+GEMM M=64,300 | 4.7e-4 / 2.3e-6 | 6.3e-2 / 4.1e-4 | 5.0e-1 / 3.3e-3 | | qmm batched w (qmv M=3 / GEMM M=64) | 1.1e-5 / 1.4e-7, 9.9e-5 / 1.0e-6 | n/a | 2.5e-1 / 2.9e-3, 2.5e-1 / 2.4e-3 | | qmm non-transposed (GEMM) | 5.7e-5 / 6.1e-7 | n/a | 2.4e-1 / 2.7e-3 | | qmm + global scale after matmul | 1.5e-4 / 1.2e-6 | 6.8e-2 / 4.9e-4 | 5.3e-1 / 5.4e-3 | | gather unsorted decode, no gs | 2.5e-5 / 2.2e-7 | 3.1e-2 / 3.0e-4 | 4.0e-1 / 2.7e-3 | | gather unsorted decode, per-expert gs | 2.3e-5 / 2.2e-7 | 3.1e-2 / 2.5e-4 | 2.5e-1 / 2.2e-3 | | gather sorted prefill (B=128), no gs | 3.1e-5 / 2.3e-7 | 5.9e-2 / 4.4e-4 | 3.0e-1 / 2.2e-3 | | gather sorted prefill, per-expert gs | 3.1e-5 / 2.3e-7 | 3.1e-2 / 2.0e-4 | 3.0e-1 / 2.0e-3 | | gather M=3 rows per expert, no gs / gs | 3.1e-5 / 2.2e-7, 2.3e-5 / 1.8e-7 | 3.1e-2 / 2.4e-4, 3.1e-2 / 2.5e-4 | 2.5e-1 / 2.1e-3, 3.9e-1 / 2.8e-3 | f16/bf16 errors are output rounding; f32 matches to accumulation order. ## Test plan - [x] `tests/rocm_nvfp4_quant.rs` and `tests/rocm_nvfp4_quant/gather.rs` (new): 7 tests pass on gfx1151 under the GPU guard after the post-rebase overlay rebuild (worst relative error f32 2.1e-6, f16 4.9e-4, bf16 5.4e-3) - [x] The six formerly gated tests take the native branch and pass: the three `switch_layers::global_scale_tests` cases, `gemma_unified::...unified_nvfp4_fused_experts_keep_their_native_global_scale`, `mistral4::modelopt_stack_tests::...`, `phimoe::sanitize_tests::...`; also the ROCm-skipped `sanitize_tests` NVFP4 loads and `ffi_tests` native-NVFP4 global-scale tests (post-rebase rerun: mlxcel lib filters 56 passed, mlxcel-core 11 passed) - [x] Regression: `rocm_mxfp4_quant`, `rocm_hip_grid_clamp gather_qmv`, `rocm_gather_qmm_expert_batched` - [x] `make verify-llama-compat verify-env-var-docs verify-kernel-port-dispatch verify-kernel-dtype-keys verify-rocm-overlay` - [x] `cargo clippy --profile test-fast --features rocm -p mlxcel -p mlxcel-core --lib --tests -- -D warnings` - [ ] Real-checkpoint smoke: not run. No NVFP4 checkpoint is in the local HF cache and large downloads were out of scope on this host. ## Changes during review - The K-multiple-of-16 check now applies to the quantized axis (`transpose_ ? K : N`), so a non-transposed weight with K = 17 reaches the dequantize + GEMM route instead of being refused. - `nvfp4_qmv` caps grid x at UINT32_MAX / 16 (blockDim.x is 16) with a clear error. - `check_nvfp4_planes` requires `scales.shape == w.shape[:-1] + [w.shape[-1] / 2]` with one-byte scales in both the `QuantizedMatmul` and `GatherQMM` arms, and the dequantize sizes its grid from the weight. Before this, a malformed scale plane (which MLX's `validate_quantized_input` lets through) could make the dequantize kernel write past its output and the kernels read past the scales. - New tests: a partial 16-column block (N = 272), out-of-range expert indices (E, E+1, i32::MAX) that must produce exact zeros, an untransposed K = 17 case, a mismatched scale plane that must be refused, and `assert_reconstructs`, which ties the global-scale reference to the source weights rather than to the kernel's constant. - The global-scale factor `global_scale / (448 * 6)` was checked against MLX's `fp_quantize`/`fp_dequantize` in `mlx/ops.cpp` and the CUDA `fp_gather_qmv`. Follow-ups: an expert-batched or grouped dequantize + GEMM path for NVFP4 gather prefill (it currently reads each expert's weights once per eight sorted rows); wave64 not measured; `nvfp4_qmv_kernel` does not bound `lhs_indices` against the x batch (same as most existing gather kernels; mlxcel generates them); sorted NVFP4 gather prefill may be slower than the former affine conversion with its expert-batched routes (not measured, no checkpoint); `GatherQMM::eval_gpu` divides by M*N before an empty-output check (pre-existing). Closes #2617
This was referenced Oct 26, 2026
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.
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
intn_rows. HIP refusesgridDim.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.UINT32_MAX / BLOCK_DIMand spills into y (thearg_reduce.hippattern of LOCAL_FIXES item 11). Kernels taken_rows, compute a 64-bit row fromblockIdx.x + blockIdx.y * gridDim.x, and return past it before any barrier. Below the cap the grid isdim3(n_rows)as before.fast_rms_norm_vjp_for_testandfast_layer_norm_vjp_for_test, run the VJP primitives throughmlx::core::vjp; the crate had no VJP entry point.Tests (gfx1151)
New
row_gridmodule intests/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:After:
7 passed; 0 failed; 2 ignored; finished in 7.12s. The wholerocm_hip_grid_clampfile 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
dxchecks that every row ran.Existing tests, post-fix, under
rocm_gpu_guard.sh --hold: mlxcel-core lib filterssoftmax(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-overlayand 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 themlxcellib test build).Closes #2619