Sorting-free top-k/top-p/min-p sampling with dual-pivot rejection kernels #901

Blocked by #900 (Gumbel-max sampling kernel: establishes the sampling-kernel plumbing, RNG semantics, and parity harness this issue extends).

Background

When top-k/top-p/min-p are active, mlxcel's fused_sample (src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp, ~4525-4597) runs argpartition for top-k, a top-p filter, a min-p compiled kernel, and random::categorical. argpartition is a partition-sort-class operation over the full vocabulary executed every token, and the chain materializes intermediate filtered distributions.

Dual-pivot rejection sampling eliminates sorting entirely. The algorithms, per filter:

• Top-k: loop over a shrinking probability interval [low, high]. Each round: (a) draw one candidate with a single-pass CDF scan restricted to prob > low (block-wide inclusive prefix sum plus an adjacent-difference step to flag the index where the running CDF crosses the drawn uniform); (b) set pivot_0 = prob[candidate], pivot_1 = (pivot_0 + high)/2; (c) block-reduce the (count, mass) of entries greater than each pivot; (d) if fewer than k entries exceed pivot_0, accept the candidate; otherwise shrink the interval to whichever pivot still contains the top-k boundary and renormalize by the counted mass. Converges in a few rounds; each round is one vocabulary scan with no sort.
• Top-p: same structure; the accept test compares the tail mass above the pivot against top_p.
• Min-p: single pass, no iteration; compute the row max, set pivot = max * p, then one filtered renormalized sampling pass.
• Joint top-k + top-p: fuse both accept tests into one rejection loop.
• Numerical hardening: if the backend JIT compiles with fast-math, flush-to-zero can corrupt subnormal pivot arithmetic and break the shrinking-interval invariant. Guard the pivot math with bit-exact IEEE multiply/add helpers (or disable fast-math for this kernel); add a test with subnormal-scale probabilities.
• Determinism option: float addition is non-associative, so a scheduling-dependent reduction order can change which index the CDF scan flags. A fixed-order scan (deterministic prefix sum) removes that; implement if cheap on our backends, otherwise document that per-run determinism relies on fixed launch geometry.

Required primitives (block reduce, inclusive scan, adjacent difference) map to CUB on CUDA and simdgroup/threadgroup intrinsics on Metal.

Task

1. Kernels (Metal + CUDA JIT strings, extending the sampling kernel file from Softmax-free GPU sampling: Gumbel-max categorical sampling kernel #900): implement min-p (single pass), top-k, top-p, and joint top-k+top-p dual-pivot rejection sampling operating on probabilities. Add a fused softmax pass (or reuse a single MLX softmax) to produce probs from the post-penalty logits; do not re-implement the logit-level pre-steps (bias/penalties/XTC stay upstream in sample_token_optimized_core).
2. Convergence guard: cap rejection rounds (e.g. 32); on cap overflow fall back to the existing argpartition path for that row and count the event in a metric. Add an adversarial test (near-uniform distribution over 152K entries) exercising the cap.
3. Wiring: fused_sample routes to the new kernels whenever top-k/top-p/min-p are active; batch-wide single launch covering all rows with per-row parameters (rows with different k/p values in one launch). Env kill switch MLXCEL_SAMPLING_REJECTION=0 restores the current chain.
4. XTC interplay: XTC masks logits before sampling; verify masked entries (zero prob) are never sampled and that renormalization inside the rejection loop is consistent with the masked distribution.

Performance validation (mandatory)

• Extend the sampling microbenchmark from Softmax-free GPU sampling: Gumbel-max categorical sampling kernel #900: vocab {32K, 64K, 152K} x batch {1, 4, 8} x config {top-k=40, top-p=0.9, min-p=0.05, top-k+top-p}, baseline vs rejection kernels, Metal and CUDA.
• Required outcomes: geometric-mean sampling-step improvement across the matrix, and no single configuration slower than baseline beyond 5%; end-to-end decode with default llama-server-style sampling (top-k 40, top-p 0.9) at batch 4 shows measurable tok/s improvement, batch 1 within noise.
• Record in docs/benchmark_results/rejection-sampling-<hw>-<date>.md.

Regression guard (mandatory)

• Distribution equivalence: for top-k, the sampled support must exactly equal the argpartition-derived top-k set (property test over random distributions with ties), and empirical frequencies within the support must pass a chi-square test against the renormalized truncated distribution (>= 1e6 draws per case). Same for top-p (support = minimal prefix set reaching p, matching the current filter's tie-breaking semantics; document any intentional difference) and min-p.
• Determinism: same seed => same stream on the same backend; cap-overflow fallback produces valid samples.
• Greedy and pure-temperature paths (from Softmax-free GPU sampling: Gumbel-max categorical sampling kernel #900) unaffected; kill switch restores baseline; full sampling tests and cargo test -p mlxcel-core pass.
• End-to-end: pinned-prompt greedy decode across 3 model families byte-identical to baseline (sampling change must not touch the greedy path).

References

• src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp (fused_sample, top_p_filter)
• src/lib/mlxcel-core/src/sampling.rs

Line numbers are indicative as of current main; search for the named symbols.

Part of epic #909.