RMSNormVJP backward writes a full {n_rows, D} gw_temp intermediate

[JasonHonKL]

Proposed changes

The GPU backward pass of mx.fast.rms_norm computes the weight gradient gw
in two full-tensor passes:

1. vjp_rms_single_row / vjp_rms_looped (one threadgroup per row) write each
   row's gw contribution into a temporary of shape {n_rows, D} — the same
   size as the whole input (mlx/backend/metal/normalization.cpp:142-152).
2. A separate strided reduction re-reads all of gw_temp and sums the rows
   into the final {D} output (normalization.cpp:200-205).

For each element of gw this moves 2 · n_rows · D · itemsize bytes through
DRAM (write + re-read) that contain no information beyond partial sums waiting
to be added. On a GPU this is pure bandwidth waste — the arithmetic to combine
rows inside the kernel is free by comparison. It also forces a transient
allocation the size of the input.

Result

variant	time / 32 iters	per grad call
eager autograd	242–498 ms	~8–15 ms
fast (before)	~77 ms	~2.4 ms
fast (after)	52.8 ms	~1.65 ms

POC

(AI generated)

import time, mlx.core as mx

def t(fn, n=50):                      # timed loop with warmup + eval
    for _ in range(5): mx.eval(fn())
    tic = time.perf_counter()
    for _ in range(n): mx.eval(fn())
    return 1e3 * (time.perf_counter() - tic) / n

x = mx.random.uniform(shape=(8192, 4096)).astype(mx.float16)
w = mx.random.uniform(shape=(4096,)).astype(mx.float16)
y = mx.random.uniform(shape=(8192, 4096)).astype(mx.float16)
mx.eval(x, w, y)

# correctness guard: fast grad must match eager autograd
gf  = mx.grad(lambda x, w: (mx.fast.rms_norm(x, w, 1e-5) * y).sum(), argnums=(0, 1))
ref = mx.grad(lambda x, w: ((mx.rsqrt(x.astype(mx.float32).square().mean(-1, True) + 1e-5)
       * x).astype(mx.float32) * w * y).sum(), argnums=(0, 1))
gx, gw = gf(x, w); rgx, rgw = ref(x, w)
assert mx.abs(gx - rgx).max().item() < 1e-4 and mx.abs(gw - rgw).max().item() < 1e-3, "WRONG"

fwd = t(lambda: mx.fast.rms_norm(x, w, 1e-5))   # unchanged by the fix
bwd = t(lambda: gf(x, w))                        # forward recompute + VJP
print(f"correct  | fwd {fwd:.3f} ms | grad {bwd:.3f} ms | VJP-only {bwd-fwd:.3f} ms")

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[JasonHonKL]

@zcbenz await verification I guess haha. Quite some changes mostly related to calculate before merge like divide and conquer.