Reduce peak memory of large half-precision random.uniform/normal

[dogukanveziroglu]

Hi guys, so I've been messing around with mlx.random for the last few days initially just trying to figure out why mx.random.normal((46341, 46341), dtype=mx.bfloat16) was crashing on my M4 16GB. I realised it uses too much memory then it supposed to do. I tried to make some changes about the calculation. Opening as a draft because I want your gut check before I polish anything further. I might try to climb a vertical flat wall but I am not sure :D. If what I made is dumb pls tell and guide me I would love to get some critisim.

Heads up: I didn't tried on nvidia GPU yet I will try it soon...

What

mx.random.normal((46341, 46341), dtype=mx.bfloat16) aborts with
"Insufficient Memory" on a 16 GB Apple-Silicon device. The standard
bits → divide → cast → minimum → mul → add chain holds three
fp32-sized buffers at the same time — about 12.88 GB of peak for a
4.3 GB output.

This PR adds two narrow GPU paths to fix it. The headline canary
goes from abort (12.88 GB) to success at 4.69 GB peak on
M4-16GB.

How

Two new dispatch paths in mlx/random.cpp:

1. Fused RandomUniform Metal kernel for half-precision uniform
when bounds are scalar, total size is even, and a single key is in
flight. Computes bits → divide → clip → cast → affine per thread
in registers, so no fp32 intermediate ever lands in global memory.
The output is bit-identical to vanilla, peak drops from ~3x to
1x, and it's 5–10× faster on large shapes.

2. Chunked path that splits the existing fp32 pipeline along
axis 0 into K independent sub-keys (K = ⌈bits_bytes / 256MB⌉,
clamped to [4, 256]). Triggers when the fp32-equivalent size is
≥ 512 MB. Peak drops to about (1 + 2/K) × output. Bytes differ
from vanilla because sub-keys differ, each chunk still does
fp32-then-cast.

Both paths are skipped when detail::in_tracing() is true, so
anything inside vmap / compile / vjp / jvp falls back to
the vanilla pipeline. fp32 is excluded from the chunked path on
purpose — vanilla fp32 already runs at ~1x output peak (the
intermediate IS the target dtype), so chunking only adds overhead.

Dispatch summary

When	What runs
GPU + bf16/fp16 + scalar bounds + even N + single key + outside transforms	fused kernel
(bit-equal vanilla)
GPU + bf16/fp16 + ≥ 512 MB + scalar bounds + single key + outside transforms	chunked path
Everything else	vanilla (unchanged)

Numbers (M4-16GB, 30 reps, fresh subprocess)

Half-precision uniform (fused kernel):

Shape	Vanilla ms / GB	This PR ms / GB
4096² bf16	7.83 / 0.10	1.67 / 0.034
16384² bf16	113 / 1.61	12.2 / 0.54
32768² bf16	448 / 6.44	44.7 / 2.15

Half-precision normal (chunked):

Shape	Vanilla ms / GB	This PR ms / GB	Δ ms	Δ peak
16384² bf16	164 / 1.61	182 / 0.94	+11%	−42%
32768² bf16	649 / 6.44	718 / 2.55	+11%	−60%
46341² bf16	abort 12.88 GB	1502 ms / 4.69 GB	n/a	fits

fp32 (deliberately unchanged):

Shape	Vanilla ms / GB	This PR ms / GB
16384² fp32 normal	164.61 / 1.07	164.57 / 1.07

Quality

• bf16 uniform fused output is bit-equal to vanilla
  (mean/var/min/max identical for 10 tested seeds at 16384²).
• bf16 normal chunked output: KS one-sample distribution matches
  vanilla; inter-chunk Pearson correlation < 0.011 at all chunk
  boundaries; unique bf16 value count 2315 at 100K samples
• mlx-lm Qwen2.5-0.5B inference: identical sampled tokens at
  identical seeds, same peak memory, ~2% better steady-state tok/s.

Trade-offs

Real but small:

• Half-precision uniform at very small shapes (≤256²) is +35–43%
  slower because the fused kernel's launch overhead dominates.
  Crossover above ~1K² gives the 5–90× speedups. Easy to add a
  min-size guard if you'd prefer.
• Half-precision normal at chunked-trigger shapes (≥11K²) is
  +9–11% slower in exchange for the 42–60% peak cut.
• Shapes with odd total element count fall back to vanilla on the
  uniform path (the kernel emits two outputs per thread).

Tests

• 22/22 in python/tests/test_random.py (8 new in TestRandomChunked
  with mathematically derived 5σ/√N tolerances)
• 708/708 full pytest

Files

• mlx/random.cpp — dispatch + chunking helper
• mlx/primitives.{h,cpp} — RandomUniform class
• mlx/backend/metal/kernels/random.metal — runiformc<T> kernel
• mlx/backend/metal/primitives.cpp — RandomUniform::eval_gpu
• mlx/backend/cpu/primitives.cpp — RandomUniform::eval_cpu stub
• mlx/backend/cuda/random.cu — CUDA mirror
• python/tests/test_random.py — TestRandomChunked

Repro the canary

python -c "                                          
import mlx.core as mx                        
mx.clear_cache(); mx.reset_peak_memory()
mx.random.seed(0)                                                                                  
a = mx.random.normal(shape=(46341, 46341), dtype=mx.bfloat16); mx.eval(a)
print(f'CANARY peak={mx.get_peak_memory()/1e9:.2f} GB')                                            
"                                                    
# Vanilla: abort. This PR: 4.69 GB.                                                                

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)

[dogukanveziroglu]

@angeloskath could you give some thoughts please, thank you.

[zcbenz]

Wouldn't mx.compile fuse the ops?

[dogukanveziroglu]

Wouldn't mx.compile fuse the ops?

That was actually the first thing that I tried but it gave me the same memory peak with 3x more latency. The problem is the two compile fusion regression for random.uniform and random.normal. First I tried to fix the donation guard in mlx/backend/common/compiled.cpp but as far as I understand @awni intentionally made bits fp32 sized at #2361 for quality so I decided to try this way.

[zcbenz]

I think we should try to make mx.compile fuse the ops, and only add a new primitive as the last resort when there is a very good reason that we couldn't do it. @angeloskath

[dogukanveziroglu]

Ok, I will try to make mx.compile fuse this without dropping the randomization quality.