[inureyes]

Summary

Issue #830 asked for an upstream bug report against ml-explore/mlx: MLX #3792 was said to have regressed the small-axis CUDA RMSNorm, with #3850 keeping the same axis broken. Verifying that against real upstream source inverted it, so no upstream report is filed and none should be. The kernel at mlxcel's MLX pin is correct on the affected axis. The pre-#3792 kernel that mlxcel carries as src/lib/mlx-cpp/patches/mlx/backend/cuda/rms_norm.cu is the one that computes a wrong normalizer, and it reads out of bounds while doing it.

What was verified, and against what

The pin is not 0.32.1 as #830 states. src/lib/mlx-cpp/CMakeLists.txt:95 pins b7c3dd6d27f45b5365b08a840310187dc503f1db, which is upstream PR #3850 itself. The on-disk MLX checkout under target/ already has the overlay applied, so all three kernels were fetched from GitHub at their own commits instead: 1700b39a (the direct parent of #3792, confirmed via the commit's parents), a5a684d (#3792), and the pin.

Measured on GB10 (sm_121, CUDA 13) with the three kernel bodies copied verbatim, each launched with its own commit's dispatch configuration and diffed against a per-row float64 reference:

Kernel	Failing axis sizes, 8 to 8192, row counts 1 to 17
1700b39 (the overlay mlxcel ships)	all 34 swept sizes in [257, 512]
a5a684d (#3792)	none
b7c3dd6d (#3850, the pin)	none

compute-sanitizer --tool memcheck reports Invalid __shared__ read of size 4 bytes for the overlay's launch geometry only; the other two are clean.

Root cause

At axis_size in (256, 512] for 16-bit dtypes, the pre-#3792 dispatch pairs block_dim == 64 with groups_per_block == 2. That launches a 128-thread block, but the reduction is BlockBroadcastReduce<float, 64, 32>, whose TempStorage is sized BLOCK_DIM / WARP_SIZE == 2 while cg::tiled_partition<32> of the launched block reports meta_group_size() == 4. The second reduction stage indexes two floats past the end of the scratch array and folds in the partials belonging to the other row that shares the block.

Two failure modes follow. Above one row the sum covers both rows, so the normalizer is deterministically wrong (0.55 to 0.94 relative error). At one row the threadIdx.y == 1 threads take the row-bounds early return and never write their partials, so the normalizer picks up residual shared memory. Seeding that window with 0, 1, 100 and 10000 moves the applied normalizer to 99.99%, 92.79%, 24.19% and 2.49% of the float64 reference.

#3792 changed that one value to 1 and added static_assert(block_dim <= 32 || groups_per_block() == 1) to pin the invariant. It is a fix, and its title says so. #3850 rewrote the forward path to one row per block, where block size and template parameter agree by construction.

#830 got the band, the config, the changed value, and DeepSeek-V2's kv_lora_rank == 512 landing on it all correct. Only the direction is backwards.

Not reconciled

#829 measured deepseek-v2-lite going 0/4 to 4/4 coherent with the overlay as the only change. Against these measurements the overlay should have made that axis worse. The document says so plainly rather than papering over it, and notes two unverified leads: #817's original "nondeterministic across fresh processes" symptom matches the uninitialized read exactly, and graph capture changes which kernels precede the norm, which is what determines the overlay's answer at one row. Settling it needs a build with the overlay removed, measured the way #829 measured. That is a separate change with real risk and is deliberately not attempted here, so no runtime behavior changes in this PR.

What changed

• docs/upstream/mlx-cuda-rmsnorm-small-axis-regression.md (new): the report, following the two existing docs/upstream/ examples. Summary, evidence, root cause with quoted source at every cited line, minimal fix, repro, affected surface. The standalone CUDA repro was run; the mlx.core sketch and the model-level repro are labelled not run.
• src/lib/mlx-cpp/patches/mlx/backend/cuda/rms_norm.cu: the tracking comment no longer dangles on "report filed as a follow-up". It points at the document, marks the superseded paragraph, and carries the corrected removal condition: nothing needs to land upstream first, so the overlay can be deleted once deepseek-v2-lite is re-verified on the upstream fast kernel, recovering the #3850 speedup. Comment lines above the ---- begin upstream ---- marker only; the vendored kernel body is untouched.

The patch-file sanity check at CMakeLists.txt:115-126 globs patches/mlx/backend/metal/*.cpp and does not match this file.

Follow-up worth filing

The overlay is wrong on an axis nine model files reach through kv_a_layernorm. Dropping it and re-verifying deepseek-v2-lite is the obvious next step, but it needs a real model run and is out of scope for a documentation change.

Test plan

• cargo fmt --all -- --check
• cargo clippy -p mlxcel-core -- -D warnings (clean, 16.23s)
• cargo check -p mlxcel-core --features cuda (clean, 29.20s; re-applied the overlay and recompiled rms_norm.cu.o through nvcc, so the comment edit is confirmed to build)
• Overlay diff is comment-only: zero non-// lines changed
• No em dashes and no AI attribution in the diff or the commit message
• cargo clippy -p mlxcel-core --features cuda skipped: it hashes to a fresh build directory and starts a cold full MLX CUDA rebuild. Aborted after two minutes and the partial directory removed. No Rust changed in this PR, so it carried no signal.

Closes #830