chore(deps): report the CUDA RMSNorm small-axis regression upstream and drop the local overlay once fixed

[inureyes]

Summary

Report the small-axis CUDA RMSNorm regression to upstream ml-explore/mlx and track dropping the local rms_norm.cu overlay once upstream fixes it. This is the follow-up noted in #829 (which fixed #824 by overlaying the pre-regression kernel while keeping the MLX pin at the latest commit).

The upstream problem

MLX #3792 ("Fix CUDA RMSNorm small-row dispatch", commit a5a684d, first shipped in the 0.32.1 pin) regressed the small-axis CUDA RMSNorm. For the dispatch config dispatch_group_dim selects when axis_size falls in the (n_per_thread*32, n_per_thread*32*2] band (with N_READS == 8 for half-precision, that band is (256, 512]), #3792 changed groups_per_block from 2 to 1 with block_dim == n_groups*group_dim == 64 (a multi-warp block). The result is a numerically-wrong normalizer for that config.

DeepSeek-V2's MLA kv_a_layernorm normalizes over kv_lora_rank == 512, which lands exactly on that band, so DeepSeek-V2 text generation degenerated into repeated tokens on CUDA. Models whose RMSNorm axis is larger (qwen3 hidden_size 1024+, etc.) take a different dispatch path and are unaffected. compute-sanitizer initcheck and memcheck are both clean, so it is a numerical kernel error, not a memory-safety bug. The later #3850 rewrite ("RMSNorm forward speed up") kept the same axis broken (with a different, alignment-dependent symptom).

Pinned by a deterministic bisect on our side: 1700b39 (the commit before #3792) is coherent, a5a684d (#3792) is garbage, using an MLX_DISABLE_COMPILE=1 MLX_USE_CUDA_GRAPHS=0 probe to remove the graph-capture nondeterminism that masked the determinism.

What to file upstream

An issue (and ideally a minimal repro + fix) on ml-explore/mlx: CUDA RMSNorm produces incorrect output for axis_size in the small multi-warp dispatch band (e.g. 512 with fp16, block_dim == 64), regressed by #3792 and still wrong after the #3850 rewrite. Include the dispatch config, the deepseek-v2-lite repro, and the bisect boundary.

Definition of done

• Upstream issue/PR filed on ml-explore/mlx and linked here.
• When an upstream fix lands and mlxcel bumps its MLX pin past it, delete the src/lib/mlx-cpp/patches/mlx/backend/cuda/rms_norm.cu overlay (added in fix(cuda): restore DeepSeek-V2 output by reverting the regressed RMSNorm kernel #829) and re-verify deepseek-v2-lite stays coherent with the upstream fast kernel, restoring the #3850 speedup. The overlay header comment already flags this.

Refs

Follow-up to #824 / #829. Sibling to #821 (the analogous upstream tracking for the CUDA graph-cache abort).