Repository navigation
test(rocm): Cover donated-cotangent LayerNorm VJP, gemv fallback, empty norms; fix item 46 #2663
Copy link
Copy link
Closed
Labels
platform:linuxLinux (CUDA / packaging) specificLinux (CUDA / packaging) specificpriority:lowLow priorityLow prioritystatus:doneCompletedCompletedtype:testTest related changesTest related changes
Description
Problem
Four small gaps in the ROCm overlay's test coverage and records, all in tests/rocm_hip_grid_clamp/ and src/lib/mlxcelverse/overlays/rocm/. Each code path below currently passes its tests only because no test reaches it, so a regression there would go unnoticed. Implement everything in one PR with one verification pass at the end and as few pushes as possible.
Items
- Donated cotangent in
LayerNormVJP::eval_gpu(mlx/backend/rocm/layer_norm.hip:415-494, fix(rocm): write the layer_norm VJP bias gradient #2635).gbis column-reduced before the VJP kernel because a donatedg(g_in_gxorg_in_gw) is overwritten by it.fast_layer_norm_vjp_for_test(src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:4272) passes a caller-held cotangent, sodonate_gis never true in today's tests. The cheapest route needs no bridge change: a non-row-contiguousg(for example a transposed view) is copied bycheck_input, which setsdonate_g |= g_copied, so with a held contiguousxthe copy lands ingx. Add a test intests/rocm_hip_grid_clamp/layer_norm_bias_grad.rsthat does this, with and without a weight, and checksgx,gwandgbagainst a host f64 reference. - Non-inline
gemv_batched(mlx/backend/rocm/gemms/gemv.hip:140, launched at:489-505only whenbatch_ndim > kMaxInlineBatchDims= 8, widened to int64 in fix(rocm): guard empty and 64-bit sizes in norm VJP, gemv, GatherQMM #2636). Add a test insize_guards.rswith at least 9 batch dims thatcollapse_contiguous_dimscannot merge (for example size-2 dims with broadcast strides alternating between the matrix and the vector), compared with a host reference. - Norm edge cases in
size_guards.rs(norms_of_zero_rows_return_empty_and_zero_gradients): add zero-rowfast_rms_norm_no_weight, and an empty input that takes thestridedbranch ofRMSNorm::eval_gpu(rms_norm.hip:309, last dim contiguous, rows not packed). If MLX flags every empty array as packed so that branch is unreachable, say so in the test comment and keep the case. - Stale grid-limit record.
LOCAL_FIXES.mditem 46 says y grids above 65535 blocks are refused withhipErrorInvalidConfigurationand that the row and tile tests "fail without the fix outright"; item 61 (fix(rocm): gather QMV large-batch launch and 64-bit output offsets #2620) measured on ROCm 7.15 (HIP 71526333) that y up to 262144 and z up to 16777216 blocks launch and run. The same claim is intests/rocm_hip_grid_clamp.rs:19-25,47-51andtests/rocm_hip_grid_clamp/gather_qmv.rs:15,39. Re-measure with a smallhipLaunchKernelprobe, then correct item 46 and those comments to what this runtime does, keeping the caps.
Acceptance Criteria
- Item 1: the new test passes, and fails when the
col_reduceblock is temporarily moved after the VJP kernel launch (record the failing output in the PR) - Item 2: the new test passes, and fails when the fallback launch is temporarily broken (for example zeroing
vec_batch_strides), proving it reachesgemv_batched - Item 3: both cases pass; the strided case states whether it reached the strided branch
- Item 4: probe output (runtime version, grid sizes, launch result) recorded in the PR; item 46 and the test comments match it
- New
LOCAL_FIXES.mditem 68 (the current highest is 67) records items 1 to 4
Verification
scripts/rocm_gpu_guard.sh --hold -- cargo test --profile test-fast --features rocm --test rocm_hip_grid_clamp -- --test-threads=1
make verify-rocm-overlay verify-fmt verify-clippy-rocmMetadata
Metadata
Assignees
Labels
platform:linuxLinux (CUDA / packaging) specificLinux (CUDA / packaging) specificpriority:lowLow priorityLow prioritystatus:doneCompletedCompletedtype:testTest related changesTest related changes
Activity
test(rocm): cover donated LayerNorm cotangent, gemv fallback, empty n…