Skip to content

test(rocm): Cover donated-cotangent LayerNorm VJP, gemv fallback, empty norms; fix item 46 #2663

Closed

Description

@inureyes

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

  1. Donated cotangent in LayerNormVJP::eval_gpu (mlx/backend/rocm/layer_norm.hip:415-494, fix(rocm): write the layer_norm VJP bias gradient #2635). gb is column-reduced before the VJP kernel because a donated g (g_in_gx or g_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, so donate_g is never true in today's tests. The cheapest route needs no bridge change: a non-row-contiguous g (for example a transposed view) is copied by check_input, which sets donate_g |= g_copied, so with a held contiguous x the copy lands in gx. Add a test in tests/rocm_hip_grid_clamp/layer_norm_bias_grad.rs that does this, with and without a weight, and checks gx, gw and gb against a host f64 reference.
  2. Non-inline gemv_batched (mlx/backend/rocm/gemms/gemv.hip:140, launched at :489-505 only when batch_ndim > kMaxInlineBatchDims = 8, widened to int64 in fix(rocm): guard empty and 64-bit sizes in norm VJP, gemv, GatherQMM #2636). Add a test in size_guards.rs with at least 9 batch dims that collapse_contiguous_dims cannot merge (for example size-2 dims with broadcast strides alternating between the matrix and the vector), compared with a host reference.
  3. Norm edge cases in size_guards.rs (norms_of_zero_rows_return_empty_and_zero_gradients): add zero-row fast_rms_norm_no_weight, and an empty input that takes the strided branch of RMSNorm::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.
  4. Stale grid-limit record. LOCAL_FIXES.md item 46 says y grids above 65535 blocks are refused with hipErrorInvalidConfiguration and 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 in tests/rocm_hip_grid_clamp.rs:19-25,47-51 and tests/rocm_hip_grid_clamp/gather_qmv.rs:15,39. Re-measure with a small hipLaunchKernel probe, 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_reduce block 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 reaches gemv_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.md item 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-rocm

Activity

  1. added a commit that references this issue on Oct 27, 2026
    2df0f88
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions