Optimize discarded global random state updates

[MicroMilo]

Proposed changes

Fixes #2711.

MLX's global random state advances through a chain of split operations. When random values are created but discarded, their sample branches are dead, but the state branch remains live. Evaluating a later random value therefore dispatches every intermediate RandomBits kernel.

This PR:

• adds a RandomAdvance primitive for CPU, Metal, and CUDA that applies repeated split-state transitions in one dispatch;
• folds candidate discovery into eval's existing degree walk;
• rewrites only runs of at least 16 transitions whose intermediate sample-key branches are not reachable from the current eval roots;
• preserves shared chains when random outputs are used or an intermediate state is itself an eval root;
• adds exact-sequence, held-output, all-used-output, async-eval, vmap, and compile tests.

The rewrite does not reduce the number of Threefry state transitions. It reduces graph and dispatch overhead by executing those transitions inside one primitive.

Correctness

The new primitive was compared bit-for-bit with repeated mx.random.split(key)[0] across multiple seeds and step counts. The integrated rewrite also checks that:

• the next generated sample exactly matches the explicit-key reference;
• a previously held random output remains valid after evaluating a compressed tail;
• fully used random-output chains are left unchanged;
• async evaluation preserves the same sequence.

Local validation on an Apple M5:

• source build with CPU and Metal backends;
• python/tests/test_random.py: 18 tests passed;
• test_eval.py, test_vmap.py, test_compile.py, and test_export_import.py: 117 tests passed;
• direct CPU/Metal state-equivalence checks for 0-1000 transitions passed;
• pre-commit run --all-files passed.

The CUDA implementation is included but was not compiled locally. Its build and runtime behavior still needs CI or reviewer validation in a CUDA-capable environment; this limitation does not affect the CPU/Metal evidence above.

Performance

An integrated 61-repetition alternating A/B run on Apple M5 measured:

Discarded random calls	Median speedup	Bootstrap 95% CI
16	1.32x	[1.28, 1.34]
32	1.46x	[1.44, 1.49]
250	2.35x	[2.22, 2.45]
1000	3.27x	[3.22, 3.32]
5000	3.75x	[3.68, 3.90]

Each graph was constructed before timing. The timed region starts after mx.synchronize(), then runs mx.eval(...) followed by another synchronization. It therefore includes eval-time graph traversal, rewrite detection, scheduling, dispatch, and execution, but excludes Python graph construction. Baseline/candidate order alternated, and the report retained every raw paired sample.

A fresh separate-build smoke comparison of this branch against main reproduced approximately 1.25x at 16 calls, 3.29x at 1000 calls, and 3.44-3.71x at 5000 calls.

Fully used RNG chains showed no significant regression in the alternating A/B experiment. A 500-node non-random graph showed a small ~1.9% detection-overhead signal in that experiment, while smaller non-random graphs did not. A fresh separate-build smoke run did not show a consistent slowdown, but the cost of performing this detection in eval remains an explicit review point.

Design points for review

1. Rewrite placement. I placed candidate discovery in eval's existing degree walk because eligibility depends on reachability from the current eval roots. Would a separate graph-rewrite pass be a better architectural fit, despite requiring the same root-relative liveness information?
2. API surface. The current patch exposes random::advance in the C++ API and binds it as the private Python hook _advance for direct equivalence and transform tests. Would you prefer keeping the C++ helper internal and removing the Python hook once the integrated tests are sufficient?
3. Backend policy. The 16-transition threshold is based on the measured crossover on Apple M5/Metal, while the current rewrite can also run on CPU and CUDA. Should the initial optimization be restricted to Metal until backend-specific crossover data is available, or is a conservative shared threshold preferable?

Checklist

• 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 (pending the API-surface decision above; no public Python API is added)