Repository navigation
fix(rocm): guard empty and 64-bit sizes in norm VJP, gemv, GatherQMM - #2636
Merged
Merged
Conversation
Three ROCm overlay launch paths narrowed 64-bit sizes to int or divided before an empty-output check. - rms_norm and layer_norm: zero rows launched row_grid(0), an invalid configuration. Forward and VJP now return early; the VJPs fill dw and db with zeros. The VJPs size gw_temp and the reduction plans with 32-bit dimensions, so past INT_MAX rows they now throw a named error before copying any input instead of wrapping and faulting in col_reduce_looped. - gemv: the row offset in gemv_impl and the batch output offset in both batched kernels are 64-bit, as #2624 did for the gather kernels. - GatherQMM: the empty-output return now precedes B = out.size() / M / N, which raised SIGFPE for M == 0 or N == 0 in every mode. tests/rocm_hip_grid_clamp/size_guards.rs covers each case; all six tests fail or crash on main and pass with the fix on gfx1151 (the gemv and row-count tests are ignored for memory). LOCAL_FIXES.md item 65. Closes #2633
inureyes
added
status:review
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.
Three ROCm overlay launch paths narrowed 64-bit sizes to
intor divided before an empty-output check. All three are fixed here with one test module,tests/rocm_hip_grid_clamp/size_guards.rs, andLOCAL_FIXES.mditem 65. No core file changes.Fixes
layer_norm.hip,rms_norm.hip): forward returns after allocating an empty output, and the VJPs return after allocatinggx, withgwandgbfilled with zeros throughfill_gpu(the sum over no rows). The VJPs sizegw_temp, thegbview from fix(rocm): write the layer_norm VJP bias gradient #2635 and the column-reduce plans with 32-bit dimensions (Shapeis int32), so the wider type does not fit; pastINT_MAXrows with a weight or bias gradient they now throw "rows exceed INT_MAX" before copying any input.gemms/gemv.hip):mat + row * colsingemv_implandout + batch_idx * rowsingemv_batchedandgemv_batched_inlineare 64-bit.quantized/qmm.hip):if (out.size() == 0) return;moved aboveB = out.size() / M / N, after theset_input_array/set_output_arraycalls, for every mode.Tests (gfx1151, red on main, green with the fix)
norms_of_zero_rows_return_empty_and_zero_gradients(forward, VJP everywrt, with and without weight)hipLaunchKernel ... invalid configuration argumentnorm_vjp_past_int_max_rows_throws(ignored; 2^31 + 2 rows)col_reduce_loopedgemv_past_i32_matrix_elements_matches_host(ignored; 65537 x 32768 f16)gemv_singlebatched_gemv_past_i32_output_elements_matches_host(ignored; 65535 x 32800 output)gemv_batched_inlinegather_qmm_affine_with_empty_output_returns_empty(M = 0, N = 0)gather_qmm_nvfp4_with_empty_output_returns_empty(M = 0, N = 0)The ignored tests pass with
--ignored(2.9 s together; the row-count test allocates nothing with the guard). With the fix, the fullrocm_hip_grid_clampbinary (25 passed, 10 ignored),rocm_nvfp4_quantandrocm_gather_qmm_expert_batchedpass. Clippy onmlxcelandmlxcel-corelib and tests is clean, and the contract checks (verify-llama-compat,verify-env-var-docs,verify-kernel-port-dispatch,verify-rocm-overlay,verify-kernel-dtype-keys) pass.Closes #2633