Repository navigation
perf(sampling): sorting-free top-p sampling with dual-pivot rejection kernels - #993
Merged
Merged
Conversation
inureyes
added
status:review
inureyes
changed the title
The filtered sampling path runs `argpartition` over the whole vocabulary for top-k, an `argsort` + `cumsum` + two `take_along_axis` chain for top-p, a compiled softmax/max/mask for min-p, and then `random::categorical`, materialising a full `[B, V]` intermediate at every stage, once per token. `argsort` alone is O(V log V) over 32K-152K entries.
All three filters are threshold filters on the probability value, so the filtered support is `{p > low}` for one scalar per row and the whole job is to find that scalar without sorting. The new kernel does it by rejection sampling on a shrinking interval, one threadgroup per row: draw a candidate from the current proposal with a fixed-order block scan, take the candidate's probability as one pivot and the midpoint of the bracket as the other, reduce `(count, mass)` above both in one vocabulary sweep, accept when the candidate passes every active filter test, otherwise raise the interval floor and repeat. The accepted draw is distributed exactly as the truncated renormalised distribution because the proposal always contains the target support.
The bisection pivot is the midpoint of the IEEE-754 BIT PATTERNS, not the arithmetic mean. Pure integer arithmetic, so a fast-math backend with flush-to-zero cannot corrupt the shrink step for subnormal probabilities, and it isolates a single float in at most 32 rounds at any magnitude where arithmetic bisection of [0, 1] never resolves a boundary sitting at 1e-6. min-p is folded into the initial interval rather than tested per round, so it costs no rejection round at all.
Wiring: `fused_sample` routes the filtered stochastic path here whenever the backend supports it and `MLXCEL_SAMPLING_REJECTION` is not falsy. Greedy stays on `argmax` and the no-filter path stays on the #900 Gumbel-max kernel, both untouched. One launch covers the whole batch, and the kernel reads `{top_k, top_p, min_p}` per row so a batch whose rows disagree needs no second launch.
Observability first, not last. #899 shipped a fused decode path that never activated and whose benchmark compared the fallback against itself, because the declines were reported at `debug` and a real server emits none. Every `fused_sample` branch here records its outcome, and the Rust caller announces the first occurrence of each distinct kind at INFO, per kind rather than one global one-shot. Cap overflow is both an INFO line and a reachable counter (`rejection_cap_overflow_rows` / `_launches`).
The CUDA JIT string is written for structural parity with the Metal one and is UNVALIDATED: no CUDA hardware and no nvcc were available, so it has never been compiled or run.
Refs #901.
Support equality and goodness of fit are the whole correctness claim for a rejection sampler, so they are tested directly rather than through a bitwise comparison that no longer exists. `sampling_rejection_tests.rs` (16 tests): top-k support equals the `argpartition` mask including a deliberate 12-way tie at the k-th rank; top-p support equals the stock nucleus prefix; min-p support equals `p >= min_p * max`; all three compose into one threshold; frequencies inside each support pass chi-square against the renormalised truncated distribution at upper-tail p = 1e-6 over 1e6 draws for the three single-filter cases; `-inf`-masked entries (token bias, XTC) are never drawn and do not shift the support; a 220-nat logit spread that reaches subnormal probabilities converges and stays exact, which is what the bit-space bisection exists for; a fixed seed reproduces a stream; a near-uniform 152064-entry row converges inside the 32-round cap; a one-round cap forces the fallback, produces valid tokens, moves both cap-overflow counters and announces itself; and every filtered config routes to the kernel while greedy, the no-filter path and the explicit reference arm report that they do not. The joint top-k + top-p case is a documented superset rather than an equality, and the test says so and pins the containment. The stock chain masks to the top-k set and renormalises before applying top-p, so its mass target is `top_p * Z_k`; the kernel applies both tests to the untruncated distribution, so its target is `top_p * total`. Resolving `Z_k` exactly would mean pinning the top-k threshold to a single float before sampling, which costs the bisection rounds the algorithm exists to avoid. Every other configuration is exact. `tests/sampling_rejection_kill_switch.rs` owns its process, the way the #900 switch test does, because the C++ gate resolves `MLXCEL_SAMPLING_REJECTION` into a `static` on first use. It pins that all four filtered configurations return to the `argpartition` chain bit-for-bit, that the decline is announced, and that the kernel entry point stays callable. `examples/rejection_sampling_microbench.rs` is the #900 harness over the matrix the issue specifies: 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}. It prints, before the table, the dispatch outcome each arm recorded, and it times the rejection arm WITH the host readback of the converged flag, because that readback is a sync a real decode step pays. Each row also reports the worst per-row round count, which is the number that says how close a shape runs to the fallback. Two things the tests found and that changed as a result. `sampling_gumbel_tests::filtered_configs_stay_bit_identical_to_the_categorical_path` asserted an invariant this issue deliberately breaks; it is now `filtered_configs_never_reach_the_gumbel_kernel`, which pins what is still true (a filtered config must never reach a kernel that draws from the full softmax). And the dispatch record needed a non-destructive read: the INFO logger's drain pops what it reports, which makes it useless to a test sharing a binary with other samplers, so `sampling_dispatch_recorded_report` was added for tests and the benchmark. No benchmark numbers here. The sweep is the maintainer's to run and record in `docs/benchmark_results/`. Refs #901.
The Rust `top_p_filter` reference copy said production routes through the C++ `top_p_filter`; since #901 that is only true under `MLXCEL_SAMPLING_REJECTION=0` or after a convergence-cap fallback. The batched seed-determinism helper said `top_p < 1` forces the categorical path; it now forces the rejection kernel, and the load-bearing part (RNG-consuming, not argmax) is what the comment should have said. Refs #901.
…rnel `prime_lookahead_with_input` schedules its forward and its sample without reading anything to host, so the GPU runs ahead on step n while the scheduler reads step n-1. The rejection path's converged-flag readback would have blocked exactly there and collapsed that pipeline for every filtered sampling config, which is the llama-server default. `fused_sample_lazy` is `fused_sample` minus the readback: same routing, same kernel, same draw, pinned by a test that compares the two streams at one seed. It is what the lookahead prime calls now. Skipping the check is sound rather than merely convenient, and the header now carries the argument. The bracket between `low` and `high` is bisected in IEEE-754 bit space, so its width halves every round from a starting distance below 2^30; after thirty rounds `low` and `high` are adjacent floats, at which point the proposal set equals the filtered support and the next draw is accepted unconditionally. Thirty-one rounds is therefore an upper bound and the 32-round cap cannot be reached. The `ok` output and the `argpartition` fallback stay on the verified path as the guard that would catch a future change to the pivot arithmetic; an arithmetic midpoint, for one, does not halve the bit distance and stalls outright on small probabilities. If the cap were ever reached on the lazy path, the row returns its own argmax, which the kernel guarantees is inside the filtered support: one greedy draw on a speculative step, not an out-of-support token. The lazy launch identifies itself as "rejection kernel, unverified" in the dispatch report, so the two are distinguishable in a log. Refs #901.
The measured matrix does not support routing the whole filtered path, so it no longer does. M1 Ultra, three repetitions of vocab {32K, 64K, 152K} x batch {1, 4, 8}, both arms timed inside one run: top-p alone 1.28x to 2.35x at every cell, top-k alone 0.31x to 0.97x at every cell, min-p alone 0.47x to 0.88x at 152K, and top-k with top-p 1.27x to 1.64x at 32K but 0.71x to 0.83x at 152K. The direction never flipped across repetitions, so this is structural rather than noise, and the issue's "no configuration slower than baseline beyond 5%" is not met by top-k.
The kernel's own round counter, which the microbenchmark reports, says why. top-p accepts in one or two rounds because a probability-weighted draw lands in the high-mass head and `mass(> p) <= top_p` is satisfied immediately. top-k needs the draw to land inside the top forty of the whole vocabulary, so it takes two to seven rounds and the count grows with the vocabulary. Every extra round is another full-row sweep, and the loop runs on ONE threadgroup per row because the shrinking interval is carried across sweeps and MLX custom kernels have no grid-wide barrier. So the kernel pays single-core bandwidth per round while `argpartition` and the compiled min-p filter run across the whole GPU. The kernel replaces a SORT; where the stock chain does not sort, it cannot win.
Routing is therefore: top-p active, and when top-k is active as well, vocabulary at most 32768. `sampling_rejection_routes` is a pure function of `(vocab, top_k, top_p, min_p)` exported to Rust, so `the_routing_policy_matches_the_measured_matrix` pins every measured cell without a GPU. 65536 is excluded for the joint case because it has not been measured for that combination, not because it lost. min-p rides along with top-p because it is folded into the initial interval and costs no round; that specific pairing is an extrapolation from the two filters measured separately and is labelled as one.
This means the llama-server default (top-k 40 with top-p 0.9) keeps the stock chain on Qwen, Llama-3 and Gemma vocabularies and only takes the kernel on 32K-vocabulary models. Declines are announced once at INFO carrying the measured numbers and the vocabulary that caused them. Correctness is still tested for every filter combination through the forced entry point, so widening the policy later is a one-constant change against tests that already prove the kernel right.
The routing tests now decide which path ran by comparing streams at a fixed seed rather than by reading the dispatch record. The record is process-global and one-shot per kind, so a concurrent test in the same binary can claim the "rejection kernel" slot and make a negative assertion a coin flip; the group run caught exactly that.
Also fixes the segfault at exit, which was not the new kernel. `examples/rejection_sampling_microbench` exited 139 after printing a complete table; the fault is in `__cxa_finalize_ranges` and reproduces with nothing in the process but one `fused_sample_categorical` call carrying a min-p. The compiled min-p filter is a function-local `static std::function` holding handles into MLX's process-global compile cache, and `exit` destroys the two in an unspecified order. It is now leaked deliberately, which is the standard fix for a process-lifetime singleton with a non-trivial destructor, and the two JIT kernel holders in `sampling.cpp` and `sampling_rejection.cpp` get the same treatment because they have the same shape. Both microbenchmarks now exit 0.
Refs #901.
… loop The routed configuration measured 1.7x SLOWER in end-to-end decode (Qwen3-0.6B, `--top-p 0.9 --temp 0.8`, 144.7 and 154.6 tok/s against 258.3 and 261.3 with the kill switch) while the op-level matrix scored the same cell 1.14x to 1.17x faster. The readback was the cause, and it is now measured rather than argued. `the_production_sampling_call_never_synchronizes` enqueues 24 chained 1024x1024 matmuls, calls the sampler, and times both the call and the drain that follows. Against the previous code the routed top-p path took **86.3ms to "build" while 63us of GPU work remained afterwards**: it had absorbed the entire queue. That is a structural witness, not a timing threshold, and it fails loudly on any future sampler that evaluates something. Both decode drivers are software pipelines. `generate.rs` builds step n+1 and `async_eval`s it, then reads step n; the batch scheduler's lookahead prime does the same. Evaluating the rejection kernel's per-row converged flags inside that build drains the queue and collapses the overlap. I had already seen this on the scheduler call site and papered over it with a second entry point, `fused_sample_lazy`, instead of recognising that the CLI loop has the identical shape. That entry point is gone: the production path is now the non-blocking one, everywhere, and there is one entry point again. Not waiting is sound, not merely cheap. The bit-space bracket bounds the loop at 31 rounds, so the 32-round cap is unreachable, and the flags are still checked, just never waited on: they are stashed and inspected on a later call once `array::is_available()` reports the launch has landed, which in a decode loop is by the next token. An unconverged row has already returned its own argmax by then, which the kernel guarantees is inside the filtered support, so the degradation would be one greedy draw rather than an invalid token, and it is counted and announced instead of silent. `cap_overflow_is_detected_without_waiting_on_the_production_path` drives that mechanism through a test-only entry point that takes an explicit cap. The forced entry point keeps the synchronous readback and the `argpartition` fallback, and `the_forced_entry_point_is_the_one_that_synchronizes` asserts it still blocks, so the two cannot quietly converge and leave the fallback dead. The harness could not see any of this, which is the third time in this epic an op-level number has failed to survive contact with a real call site, so it now has two modes. `iso` is the old one. `pipe` reproduces the decode loop: enqueue a calibrated synthetic forward, build the sample, `async_eval` both, read the previous step's tokens, one synchronize at the end. A `routed=no` row reads about 1.00x in `pipe` because both arms are then the same code, which doubles as a check that the routing policy is live. `docs/benchmarks.md` states the general rule: an op-level harness measures an operation with the caller's pipeline already drained, so anything that touches host memory or branches on device state needs a pipelined arm before its number means anything. Refs #901.
Rebasing onto #902 put two deliberate decompositions of the same region side by side. #902 extracted `fused_sample_filter_logits` so `fused_sample_probs` reproduces the sampler's support and normalization by construction; #901 extracted the chain plus its categorical draw as the reference arm, the kill-switch target and the cap-overflow fallback. Both survive: `fused_sample_argpartition_chain` is now a two-line wrapper over #902's function, so there is still exactly one filter chain in the file and a fallback draw cannot drift from the distribution that function reports. The semantic half is the part that would have broken quietly. #902's contract is that `fused_sample_probs` describes what `fused_sample` draws from, and #901 gave the sampler a second support: the rejection kernel evaluates top-k and top-p against the untruncated row while the stock chain renormalises between them, so for that one combination the kernel's support is a superset. A `fused_sample_probs` that kept reporting the stock support would have handed the speculative accept test a `p` missing part of the target's support, rejecting tokens the target could genuinely have produced, and no existing test would have caught it: #902's own cases all use configurations #901 does not route. Rather than reimplement the kernel's support, `fused_sample_filter_logits` takes a `rejection_semantics` flag that REORDERS the same stages: run the top-k and top-p masks independently against the untruncated row and intersect them with `minimum`, which is exact because the masked fill is `lowest()` and therefore below any real logit. min-p stays last either way, since its threshold is `min_p * p_max` and the row maximum survives every filter. With either filter alone the two orders coincide, so the flag changes nothing outside the one combination. `fused_sample_probs` passes the same predicate the sampler used, and the cap-overflow fallback passes `true` because its configuration is routed. Two new guards. `the_reported_distribution_matches_what_the_sampler_draws` draws through `fused_sample` for six configurations, routed and declined, and requires every sampled token to carry mass in the reported distribution and every well-supported token in that distribution to be reachable. `the_two_filter_orderings_really_do_differ_for_top_k_plus_top_p` proves the first is not passing vacuously by showing the two orderings produce different supports on a distribution that straddles the boundary. Three problems the rebase surfaced in my own deferred-verification code, all fixed. The row count now comes from the flags array itself instead of a separately stored field, so a mismatch cannot read past the buffer and report garbage as cap overflow. A row counts as overflow only if it BOTH reports failure and reports consuming the whole round budget, so a buffer that was never written (which reads as zeros) cannot be reported as a kernel regression; under-reporting a diagnostic is recoverable, sending someone to hunt a nonexistent bug in the pivot arithmetic is not. And the single stash slot became a four-slot ring, because a launch that had not landed yet was being overwritten by the next one, silently dropping the only signal that the kernel regressed. The cap-overflow test now pins the counting rule through the same `count_and_report_overflow` the drain uses, instead of depending on the ring surviving concurrent sampling. The ring is best-effort by design and that is the right trade for a hot-path diagnostic, but the wrong basis for an assertion. The test also needed top-k active to overflow at all: top-p alone accepts in the first round, so a one-round cap converges and there is nothing to count. Verified: `sampling` group 100 passed across six consecutive runs, #902's `stochastic_accept` 18 passed and `distribution_tests` 14 passed, clippy and fmt clean. Refs #901.
inureyes
force-pushed
the
feature/issue-901-rejection-sampling
branch
from
August 2, 2026 04:46
b86bcbc to
2c2f437
Compare
…M1 Ultra Adds the performance-validation report issue #901 requires, covering the full vocab by batch by filter matrix in two measurement modes plus end-to-end decode, with the routing decision derived from the numbers rather than asserted. The routing policy follows from one fact: the kernel replaces a sort, so where the stock chain does not sort it cannot win. top-p converges in one or two rounds at every vocabulary and gains 1.17x to 1.82x pipelined; top-k alone needs two to seven rounds with the count growing in vocabulary and loses; min-p alone loses because there is no sort to remove. The joint case routes only at vocab 32768 or below, where batch 4 at 65536 lost reproducibly across three repetitions. States plainly that the llama-server default of top-k 40 with top-p 0.9 therefore keeps the argpartition chain on Qwen, Llama-3 and Gemma vocabularies. Records why the harness now reports two speedups per row. The first version of this kernel measured 1.17x isolated and 0.58x end to end, because the eager entry point read convergence flags back to the host every token and absorbed the outstanding GPU queue. An isolated harness cannot see that cost, since it already synchronizes. The production path now evaluates nothing, the bit-space bracket bounds the loop below the round cap, and a structural test pins the property by timing the drain after a sampler call: 86.293 ms of build time against 63.3 us of outstanding work is a hardware-independent witness of a blocking call. Also records the semantic reconciliation with #902. The kernel evaluates top-k and top-p against the untruncated row while the stock chain renormalizes between them, so `fused_sample_probs` would have reported a support missing part of the target's, and the speculative accept test would have rejected tokens the target could genuinely produce. No existing test covered it because #902's cases are all configurations #901 does not route.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Replaces the
argsort+cumsumnucleus filter and the trailingrandom::categoricalwith a single custom kernel that resolves the truncated sampling support by dual-pivot rejection sampling on a shrinking probability interval, with no sort anywhere. One launch covers the whole batch and the kernel takes per-row{top_k, top_p, min_p}.MLXCEL_SAMPLING_REJECTION=0restores the previous chain exactly.Rebased onto
origin/main(a93d365e) with #902 and #907 merged. The reconciliation with #902 is semantic, not textual, and is described below.Measured
End-to-end decode, Qwen3-0.6B-4bit,
--top-p 0.9 --temp 0.8 --seed 99 -n 200, three reps per arm: 274.55 tok/s with the kernel against 264.95 tok/s with the kill switch, so 1.04x. The gain is small because sampling is a minor share of decode on a 0.6B model, which is the same reason #900's end-to-end came in at +2.3%; it is the honest number and it is published as such.Op level, M1 Ultra, three repetitions of vocab {32K, 64K, 152K} x batch {1, 4, 8}.
isois the classic per-iteration measurement,pipereproduces the decode loop.Every
routed=norow reads about 1.00x pipelined, which is the check that the routing policy is in force rather than merely declared. The 64K joint case is excluded by measurement rather than omission.The three things this PR got wrong and then fixed
Routing. The first cut routed the whole filtered path. The kernel replaces a sort, so it wins exactly where the stock chain sorts, which is when top-p is active; the round counter shows why (top-p accepts in one or two rounds, top-k needs two to seven and the count grows with the vocabulary, and every round is another full-row sweep on a single threadgroup).
sampling_rejection_routesis a pure function pinned cell by cell against the table above. The llama-server default (top-k 40 with top-p 0.9) keeps the stock chain on Qwen, Llama-3 and Gemma vocabularies and routes only on 32K-vocabulary models.Synchronization. The routed path read the per-row converged flags back to host per token. Both decode drivers are software pipelines that build step n+1 and
async_evalit before reading step n, so that drained the queue inside the caller's build phase: 1.7x slower end to end whileisocalled it 1.14x faster.the_production_sampling_call_never_synchronizesis the permanent guard and was the evidence: against the pre-fix code the routed path took 86.293 ms to build while 63.3 us of GPU work remained afterwards. The production path now evaluates nothing. That is sound rather than convenient, because the bit-space bracket bounds the loop at 31 rounds so the 32-round cap is unreachable; the flags are still checked, just never waited on, through MLX's non-blockingarray::is_available().Agreement with #902. See below.
Reconciling with #902
#902 extracted
fused_sample_filter_logitssofused_sample_probsreproduces the sampler's support and normalization by construction; #901 extracted the chain plus its categorical draw as the reference arm, the kill-switch target and the cap-overflow fallback. Both survive:fused_sample_argpartition_chainis now a two-line wrapper over #902's function, so the file still has exactly one filter chain and a fallback draw cannot drift from the distribution that function reports.The semantic half is what would have broken quietly. #902's contract is that
fused_sample_probsdescribes whatfused_sampledraws from, and #901 gives the sampler a second support: the rejection kernel evaluates top-k and top-p against the untruncated row while the stock chain renormalises between them, so for that one combination the kernel's support is a superset. Afused_sample_probsthat kept reporting the stock support would have handed the speculative accept test apmissing part of the target's support, rejecting tokens the target could genuinely have produced, and no existing test would have caught it: #902's own cases all use configurations #901 does not route.Rather than reimplement the kernel's support alongside it,
fused_sample_filter_logitstakes arejection_semanticsflag that reorders the same stages: run the top-k and top-p masks independently against the untruncated row and intersect them withminimum, which is exact because the masked fill islowest()and therefore below any real logit. min-p stays last either way, since its threshold ismin_p * p_maxand the row maximum survives every filter. With either filter alone the two orders coincide, so the flag changes nothing outside the one combination.fused_sample_probspasses the same predicate the sampler used, and the cap-overflow fallback passestruebecause its configuration is routed.Two new guards:
the_reported_distribution_matches_what_the_sampler_drawsdraws throughfused_samplefor six configurations, routed and declined, and requires every sampled token to carry mass in the reported distribution and every well-supported token in that distribution to be reachable;the_two_filter_orderings_really_do_differ_for_top_k_plus_top_pproves the first is not passing vacuously.The rebase also surfaced three problems in my own deferred-verification code. The row count now comes from the flags array itself rather than a separately stored field, so a mismatch cannot read past the buffer. A row counts as overflow only if it both reports failure and reports consuming the whole round budget, so a buffer that was never written (which reads as zeros) cannot be reported as a kernel regression; under-reporting a diagnostic is recoverable, sending someone to hunt a nonexistent bug in the pivot arithmetic is not. And the single stash slot became a four-slot ring, because a launch that had not landed was being overwritten by the next one, dropping the only signal that the kernel regressed.
The algorithm
All three filters are threshold filters on the probability value, so the filtered support is
{p > low}for one scalar per row and the job is finding that scalar without sorting. Each round draws a candidate from the current proposal with a fixed-order block scan, setspivot_0to the candidate's probability andpivot_1to the midpoint of the bracket, reduces(count, mass)above both in one vocabulary sweep, and accepts as soon as the candidate passes every active test. The accepted draw is distributed exactly as the truncated renormalised distribution, because the proposal always contains the target support.pivot_1is the midpoint of the IEEE-754 bit patterns, not the arithmetic mean. Pure integer arithmetic, so a fast-math backend with flush-to-zero cannot corrupt the shrink step for subnormal probabilities. It also turns the round cap into a proof: the bit distance at least halves every round from a starting value below 2^30, so after thirty roundslowandhighare adjacent floats, the proposal equals the support, and the next draw is accepted unconditionally. min-p costs no round at all: it is folded into the initial interval.Harness
Two speedups per row.
isois the old measurement;pipereproduces the decode loop (calibrated synthetic forward, build the sample,async_evalboth, read the previous step's tokens, one synchronize at the end). The pipelined arm samples throughfused_sample, so arouted=norow reads about 1.00x because both arms are the same code.docs/benchmarks.mdrecords the general rule: an op-level harness measures an operation with the caller's pipeline already drained, so anything that touches host memory or branches on device state needs a pipelined arm before its number means anything.Observability
Every
fused_samplebranch records its outcome and the Rust caller announces the first occurrence of each distinct kind at INFO (never debug), one-shot per kind. Ten kinds, includingnot routedcarrying the measured numbers and the vocabulary that caused the decline, and separate kinds for the production launch ("converged flags checked without waiting") and the forced one ("reads the converged flags back and therefore synchronizes"). Cap overflow is both an INFO line and a reachable counter.Statistical guards and their actual sample counts
sampling_rejection_tests.rs, 26 tests. Support equality against the argpartition-derived set including ties (200000 draws, 12-way tie at the k-th rank). Chi-square inside the support against the renormalised truncated distribution at upper-tail p = 1e-6 at the full 1e6 draws for top-k=40, top-p=0.95 and min-p=0.05. Composite cases draw 200000, not 1e6, stated plainly. Subnormal hardening over a 220-nat spread reaching below the f32 subnormal floor. Determinism over an eight-launch stream. The adversarial near-uniform 152064 row at both the production cap and a cap of 1. Routing decided by comparing streams at a fixed seed rather than by reading the process-global one-shot dispatch record, which a concurrent test can claim.Also fixed: a pre-existing segfault at exit
examples/rejection_sampling_microbenchexited 139 after printing a complete table. The fault is in__cxa_finalize_rangesand reproduces with nothing in the process but onefused_sample_categoricalcall carrying a min-p, so it predates this issue: the compiled min-p filter is a function-localstatic std::functionholding handles into MLX's process-global compile cache, and C++ leaves their destruction order unspecified. Now leaked deliberately, as are the JIT kernel holders. Both microbenchmarks exit 0.What is NOT validated
CUDA is unvalidated. No CUDA hardware and no
nvcc, so the CUDA JIT string has never been compiled or run. Marked UNVALIDATED in the file, the same status as the CUDA strings insampling.cpp,fused_norm.cppandfused_rope_append.cpp.top-p + min-pis an extrapolation, not a measured cell.The end-to-end and matrix numbers above predate the rebase. Nothing in the reconciliation changes the routed path's arithmetic, but they are worth re-running once on the rebased branch before merge.
Test plan
cargo test --release -p mlxcel-core --lib --features metal,accelerate sampling-> 100 passed, 0 failed, six consecutive runs (the module shares process-global sampler state, so the repeats are the real check)cargo test --release -p mlxcel-core --lib --features metal,accelerate stochastic_accept-> 18 passed (Distribution-preserving speculative-decoding acceptance (chain speculative sampling) #902)cargo test --release -p mlxcel-core --lib --features metal,accelerate speculative::distribution_tests-> 14 passed (Distribution-preserving speculative-decoding acceptance (chain speculative sampling) #902's distributional regression guard)cargo test --release --test sampling_rejection_kill_switch --features metal,accelerate-> 1 passedcargo clippy -p mlxcel-core --lib --tests --features metal,accelerate -- -D warnings-> cleancargo clippy --lib --tests --example rejection_sampling_microbench --features metal,accelerate -- -D warnings-> cleancargo fmt --all -- --check-> cleanCloses #901