trunk/180b3630b44c3f9f40a29d5f9e636fb1d190dffb: [MPS] Migrate randperm off MPSGraph to Metal (#182528)

Drops the MPSGraph randperm path (mps::random_mps_impl<int64_t> plus the RandomCachedGraph / RandomOpBlock machinery, which had no other callers after the dispatch unification) and reimplements randperm_out_mps with three Metal-native paths selected by n:

• Small n (n <= 384, int32/int64 contiguous result): a single-threadgroup Fisher-Yates kernel writes the permutation directly into result. Thread 0 walks i = n - 1 .. 1 swapping buf[i] with buf[Philox % (i + 1)] in threadgroup memory; the rest of the threadgroup cooperates on the identity init and the writeout. No keys buffer, no argsort, no cast. The 384 cutoff is the measured crossover where the serial swap loop starts losing to launching a sort.


• Mid n (384 < n < 32768, or non-int dtypes): uniform float keys + argsort, routing int64 indices straight into result via at::sort_out when the result is already int64 and contiguous.


• Large n (n >= 32768): random integer keys + a partial radix sort, then an island-dedup pass. A permutation of [0, n) does not need a full sort - sorting random keys by only their low bits already orders the distinct keys uniformly. We generate 24-bit keys and sort just those 24 bits (3 radix passes instead of the 4 a full 32-bit-key sort needs), then randperm_dedup_islands runs an exact Fisher-Yates shuffle within each equal-key island to scrub the residual tie-bias. The result is uniform over all n!, matching CPU (Fisher-Yates) and CUDA (radix_sort_pairs + randperm_handle_duplicate_keys). The limited-bit sort reuses the existing Metal radix sorter in Sort.mm via a new max_passes cap (default -1, so the general sort/argsort path is unchanged).

M4 Max, int64 result, per-call latency, baseline is the 2.12 release (still MPSGraph):

We beat MPSGraph at every size except n=1M, where its fused randperm still edges out the 3-pass radix (0.57 vs 0.74 ms). The island-dedup costs ~7% at 10M and is essentially free at 1M. Fully beating MPSGraph at 1M would need a non-sort algorithm (e.g. a fused random-slot scatter), left as follow-up.

Validity (sort(randperm(n)) == arange(n)) checked for n in {32768, 1e5, 1e6, 2e6} across int32/int64. Uniformity verified with a full-distribution chi-square (n=7, all 5040 permutations) and an adjacent-ascending-pair test at n=1M: this path matches CPU and the 2.12 MPSGraph baseline (z within +-1 of uniform), while a deliberately tie-biased sort-without-dedup control is flagged at z=+62.

Authored with Claude.

Pull Request resolved: