[inureyes]

Summary

Adds modified rejection sampling as an opt-in acceptance rule for the classic speculative path (MLXCEL_SPECULATIVE_STOCHASTIC_ACCEPT=1, default off), together with the proof, the regression guards, the closed-form diagnostic, and a correction to the issue's premise that changes what this issue should be understood to be asking for.

The premise correction, and why it changes the ask

Issue #902 states that the verify loop "accepts a draft token iff it equals the target model's argmax", and that this biases the output at temperature > 0. For the path the issue's line numbers point at, that is not true.

SpeculativeGenerator selects the target token with sample_token_optimized(pos_logits, target_sampling, history), a fresh draw from the target sampler, which is argmax only at temperature == 0. It emits that draw on both branches: the drafted token on acceptance (equal to the draw by the accept test) and the draw itself on rejection. Every emitted token is therefore a fresh conditional draw from the target, and the stream was already exactly a target-only sample. Call this rule sampler-match. It was never biased; what it loses is acceptance rate, sum_x p(x) q(x).

This is measured, not argued: the_default_sampler_match_rule_also_matches_the_target_distribution passes the same chi-square that rejects argmax-against-argmax at ~1825.

The biased rule the issue describes is real, but it lives in Gemma4MtpTargetAdapter::verify_forward (argmax_per_position, src/models/gemma4_mtp_target.rs:655) and the DFlash round loop (argmax_logits_to_array, drafter/dflash/round_loop.rs:728), both regardless of sampler.temperature. Neither is changed here.

The acceptance ceiling: why "strictly improves vs argmax" is unachievable

For a q-distributed proposal and a p-distributed emission, the largest probability the two can coincide is sum_x min(p(x), q(x)) — the maximal-coupling bound, equivalently 1 - TV(p, q). Every distribution-preserving rule is capped there.

Rule	Lossless	Acceptance probability
Sampler-match (the default, pre-#902)	yes	sum_x p(x) q(x)
Modified rejection sampling (this PR, opt-in)	yes	sum_x min(p(x), q(x)) — the ceiling
Argmax-against-argmax (MTP, DFlash)	no	q(argmax p) — not bounded by the ceiling

Since min(a, b) >= a * b on [0, 1], rejection sampling always accepts at least as often as sampler-match. But q(argmax p) is under no such constraint: with p(a) = 0.3 and q(a) = 0.95 the argmax rule accepts at 0.95 where the ceiling for any correct rule is about 0.30.

So the reviewer's reading is correct and I agree with it: no implementation can make a distribution-preserving rule strictly improve mean accepted length against argmax acceptance. That is not a defect to fix, it is the price of unbiasedness, and it means the issue's required outcome is unachievable as literally written. Read against sampler-match — the rule this generator actually ran — the improvement is real and provable, and that is what the tests assert.

Recommendation: opt-in, not default

Modified rejection sampling is off by default in this PR. MLXCEL_SPECULATIVE_STOCHASTIC_ACCEPT=1 enables it; SpeculativeGenerator::with_stochastic_acceptance(bool) overrides programmatically.

The reasoning: on the only path this rule reaches, the rule it replaces was already distribution-preserving, so it buys acceptance rate and nothing else. The available gain is exactly sum_x min(p, q) / sum_x p(x) q(x), and that ratio collapses toward 1 whenever the drafter is confident, because q(t*) ~ 1 makes min(p, q) and p * q coincide. Measured at ~1.02 on a Llama-3.1-8B / Llama-3.2-1B pair at temperature 0.7. Two extra full-vocabulary passes and an extra host sync per verified position do not pay for two percent of theoretical acceptance.

This follows MLXCEL_FUSED_QK_NORM (#326) and the #905 fusions in this same epic: the machinery, the proof and the tests land and are available, and the default does not move until a measurement justifies it. The place this rule earns its default is a verify path whose current rule is genuinely biased — MTP and DFlash — where it would buy correctness rather than only acceptance rate, and where trading acceptance for unbiasedness is a real decision.

With the env default off, every distributional test would have silently stopped exercising this module, which is exactly how #899 nearly shipped a dead path. The programmatic override exists so the tests force the rule explicitly, and stochastic_runs_report_the_stochastic_rule_and_a_longer_accepted_run fails if a run reports anything but rule=stochastic.

A separate bug this work uncovered: #994

The draft KV cache rewind trims min(rejected + 1, n) where the correct trim is rejected - 1. At accepted == 0 the draft cache does not advance at all, so the drafter proposes from a frozen prefix while the target advances, in a self-reinforcing loop. It predates #902, affects every acceptance rule identically, and is the most likely explanation for measured acceptance of 0.12-0.46 on an 8B/1B pair against the 0.6-0.8 reported upstream. It also makes any speculative A/B on this path noisy: two 200-token runs produced closed_form_sum_min of 0.221 and 0.396 on the same model pair.

Filed as #994 with the derivation. draft_cache_falls_behind_the_main_cache_every_round_pre_existing_defect pins the current wrong behavior so a fix produces a failing test rather than silence. Not fixed here: it changes token streams at every temperature and needs its own parity story.

What changed

src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp — the filter chain of fused_sample_impl (temperature scale, top-k, top-p, min-p) is extracted into fused_sample_filter_logits, and a new fused_sample_probs runs that same chain and returns softmax(...) in float32 instead of drawing. The effective distribution is therefore derived from the sampler's own code rather than reimplemented next to it, so it cannot drift. A greedy configuration (temperature == 0 or top_k == 1) returns the one-hot argmax indicator, which is the degenerate q a greedily-proposing drafter needs. Float32 is deliberate: the caller reads single entries back to the host and compares them against a uniform draw, and an f16 probability underflows below ~6e-8, turning a legitimately small acceptance ratio into an unconditional accept.

src/lib/mlxcel-core/src/sampling.rs — the pre-sampler steps (last-position slice, token bias, repetition / DRY / frequency-presence penalties, XTC) are split out of sample_token_optimized_core into preprocess_logits_for_sampling, with no behavior change. Two new public functions build on it: effective_token_distribution returns p or q for a given logits/config/history triple and consumes no randomness, and sample_token_with_distribution returns a token together with the distribution it was drawn from from a single pre-step, which is what guarantees q is the actual proposal distribution rather than a later reconstruction.

src/lib/mlxcel-core/src/speculative/stochastic_accept.rs (new) — acceptance_rule, accept_draft_token, residual_resample, verify_draft_token, the kill switch and the one-shot info logging. The module header carries the losslessness proof and the reason it holds for any proposal q. The accept test is p(t) > 0 && u * q(t) <= p(t); the first conjunct is load-bearing, not cosmetic, because a q(t) that has underflowed float32 to zero would otherwise make the product test read 0 <= 0 and accept a token the target's own filter excluded. The residual draw goes through ffi::fused_sample(log(relu(p - q)), 1.0, 0, 1.0, 0.0), so softmax(log r) == r / sum(r) is the normalized residual by construction and the draw reuses the #900 Gumbel-max kernel where the backend supports it.

src/lib/mlxcel-core/src/speculative/mod.rs — the verify loop dispatches on the rule. On the stochastic arm the draft loop captures q per position (left unevaluated, so a chain rejected early never materializes the tail) and the verify loop builds p at each position from the history including the tokens accepted so far in the round. Both arms produce the same (accept, replacement) pair, so all the downstream bookkeeping, EOS handling and cache rewind is shared and unchanged. New SpeculativeAcceptanceStats counts rounds, proposals and acceptances.

Docs: docs/speculative-acceptance.md (new), MLXCEL_SPECULATIVE_STOCHASTIC_ACCEPT in docs/environment-variables.md, links from README.md, docs/README.md, docs/architecture.md.

The distributional argument, and how the test verifies it

Null hypothesis. The token stream emitted by SpeculativeGenerator::generate at temperature > 0 is distributed identically to sampling the target model alone at the same temperature.

Why the test can check it directly. The test target model's logits do not depend on its input, so target-only sampling emits i.i.d. draws from one known categorical p = softmax(target_logits / 0.7). Under H0 every token the speculative generator emits, whatever mixture of accepted drafts, residual resamples and bonus tokens produced it, is also an i.i.d. draw from that same p, so a Pearson goodness-of-fit test against p tests H0 directly rather than by proxy.

Statistic, samples, threshold. Pearson chi-square over a 6-entry vocabulary, dof = 5. The committed default draws 12000 emitted tokens (SAMPLE_TOKENS, overridable with MLXCEL_SPEC_DIST_SAMPLES); the num_draft = 1 companion test draws half that. Rejected above 25.7448, the upper-tail quantile at alpha = 1e-4, so a correct implementation fails one run in ten thousand.

Demonstrated power, not assumed power. chi_square_rejects_known_wrong_acceptance_rules replays the same statistic against two deliberately wrong rules and requires rejection. Measured on this branch at n = 4000 verify positions: correct rule 4.88, "resample from p instead of the residual" 136.48, argmax-against-argmax 1825.13, threshold 25.7448. The middle arm is the specific mistake a naive implementation makes, and it is detected with a five-fold margin at a quarter of the end-to-end sample count.

Where chi-square is blind, zero-tolerance assertions cover. A goodness-of-fit test cannot see a single wrong token, so: filtered_target_support_is_never_violated runs a top_k = 2 target against an unfiltered drafter over 1500 emitted tokens and fails on the first token outside the target's two-token support (exactly what a missing p(t) > 0 conjunct produces); tokens_outside_the_target_support_are_never_accepted and underflowed_proposal_mass_does_not_bypass_the_target_support assert the same property on the primitive over 400 and 200 trials with no tolerance; residual_draw_follows_the_normalized_relu_difference fails if the residual ever emits a token with p <= q; and the temperature-0 test fails on one differing token.

Temperature 0

acceptance_rule returns Argmax whenever the target sampler is greedy (temperature == 0.0 or top_k == 1, matching the C++ sampler's own short-circuit, pinned by sampler_is_greedy_matches_fused_sampler_short_circuit). On that arm no distribution tensor is built, no randomness is drawn, and the pre-existing comparison runs verbatim. temperature_zero_stream_is_byte_identical_to_greedy_target_only walks a deterministic 6-cycle target against two drafters (one agreeing on one token in six, one never agreeing) at num_draft 1, 2 and 4, and asserts the emitted vector equals the target's own greedy continuation exactly.

Proving from a log which rule ran

Both lines are at info. The first time a process runs each distinct rule it logs speculative acceptance rule active rule="stochastic (modified rejection sampling)" distribution_preserving=true; the latch is per rule kind, not one global flag, so a process that serves a greedy request and later a temperature > 0 one logs both lines. The other three labels are argmax (greedy target sampler), argmax (stochastic acceptance disabled by env) and argmax (drafter reported no proposal distribution). The first occurrence of each per-token outcome is logged the same way (accept, residual resample, and the degenerate residual fallback). Every generate call then logs speculative decode finished rule=Stochastic rounds=.. proposed_draft_tokens=.. accepted_draft_tokens=.. acceptance_rate=.. mean_accepted_len=.. generated_tokens=.., which is both the A/B measurement and the proof of which arm produced it. stochastic_runs_report_the_stochastic_rule_and_a_longer_accepted_run fails if a temperature > 0 run reports anything other than Stochastic.

Kill switch

MLXCEL_SPECULATIVE_STOCHASTIC_ACCEPT=0 (also false, no, off, case- and whitespace-insensitive) restores the previous rule at every temperature, read once per process.

RNG

The accept test draws one U[0,1) per verified position and the residual draws one categorical per rejection, both through MLX's default key sequence, so --seed still reproduces a run exactly. But the stream is consumed differently than before, so at an equal seed the emitted tokens differ from a pre-change run even though the distribution is identical (the same relationship MLXCEL_SAMPLING_GUMBEL has). Set the kill switch to reproduce a stream recorded before this landed. Documented in docs/speculative-acceptance.md and the env-var table.

Observability: how to prove which rule ran

The mlxcel CLI installs no tracing subscriber, so info-level instrumentation is discarded on the only binary that runs SpeculativeGenerator. Every speculative run therefore prints an unconditional stdout summary:

[Speculative acceptance] rule=stochastic rounds=153 proposed=612 positions_tested=195 accepted=46 per_position_acceptance=0.2359 acceptance_rate=0.0752 mean_accepted_len=0.3007 (stochastic (modified rejection sampling))

rule= is a stable, whitespace-free, greppable id: greedy-argmax, sampler-match, stochastic, sampler-match-no-proposal-distribution. Its shape is pinned by cli_summary_line_carries_the_rule_id_and_mean_accepted_len. Setting RUST_LOG additionally installs a subscriber for the offline CLI commands (stderr, so piped stdout stays clean); default output is unchanged when it is unset, and Serve is skipped because the server installs its own.

Two metric corrections found by making it visible

acceptance_rate is the wrong metric for an A/B. It is accepted / proposed, but a chain stops at its first rejection and never tests the block positions behind it, so proposed is not the denominator theory predicts against. On a real Llama 8B/1B pair this made a correct implementation read as a 2x regression (0.059 stochastic vs 0.119 kill-switch). positions_tested is now counted identically on both rules and per_position_acceptance is reported beside it. On the same runs that reads 0.2359 vs its own closed form 0.2213: correct.

The drafter-kind line was misleading. With no explicit --draft-kind, the CLI printed the kind auto-detected from the drafter's config.json (usually dflash) and then ignored it and ran the classic SpeculativeGenerator anyway, so a benchmark run was indistinguishable from a DFlash round-loop run. The path actually taken is now printed:

Resolved drafter kind: dflash (block_size = 16, default)
Drafter kind was auto-detected only; running the classic SpeculativeGenerator path (pass --draft-kind explicitly to select a kind-specific round loop).

Passing --draft-kind explicitly does not give you a DFlash round loop here; it returns an error saying the offline path does not construct one.

Verifying the gain before trusting a throughput number

MLXCEL_SPECULATIVE_ACCEPT_DIAG=1 adds closed_form_sum_min (sum_x min(p, q), what modified rejection sampling accepts) and closed_form_sum_prod (sum_x p(x) q(x), what the old rule accepts), averaged over the positions the run actually tested, and makes the argmax arm capture q so both halves of an A/B report the same quantity. Because min(a, b) >= a * b on [0, 1], sum_min >= sum_prod holds for every pair of distributions, so the two numbers convert an acceptance anomaly into an arithmetic statement: the measured per-position rate must sit at sum_min, and sum_min can never fall below sum_prod. Off by default (two full-vocabulary reductions per tested position).

Measured on Llama-3.1-8B-Instruct-4bit with a Llama-3.2-1B-Instruct drafter, 200 tokens, --num-draft-tokens 4, one trajectory per row:

Temperature	closed_form_sum_min	closed_form_sum_prod	Ratio	measured per_position_acceptance
0.7	0.2808	0.2740	1.02x	0.2902
1.0	0.1738	0.1613	1.08x	0.1735
1.5	0.6809	0.0494	13.8x	0.6966

The measured value sits on sum_min at every temperature: that is the correctness result, and it is the strongest evidence in this PR that the accept test is right on real models rather than only on synthetic ones. The ratio column is the throughput headroom, and on this pair at 0.7 there is essentially none, because sum_x min(p, q) collapses onto sum_x p(x) q(x) whenever the drafter is near-deterministic (q(t*) ~ 1 makes min(p(t*), q(t*)) = p(t*) and p(t*) q(t*) ~ p(t*)). The gain appears when q is diffuse. Between-trajectory variance is also large: two 200-token runs at 0.7 differed by nearly 2x in sum_min, so a single short run there measures noise, not the change. Read sum_min / sum_prod first; it bounds what any throughput measurement can show. The 1.5 row is a real measurement of the mechanism but comes from a degenerate high-temperature trajectory where both models agree on repetitive text, so it is not a production number.

How to run the throughput benchmark

Throughput timings were not run here, by instruction. The classic SpeculativeGenerator is reachable only from the offline CLI (src/server/speculative_dispatch.rs:41 records that the server does not use it), so the A/B is mlxcel generate --draft-model, with the kill switch selecting the arm:

MLXCEL_SPECULATIVE_ACCEPT_DIAG=1 ./target/release/mlxcel generate -m models/<target> --draft-model models/<draft> --num-draft-tokens 4 --temp 0.7 --seed 1234 -n 256 -p "<prompt>" for the default (sampler-match) arm, and the same command with MLXCEL_SPECULATIVE_STOCHASTIC_ACCEPT=1 for the opt-in arm. Confirm rule=argmax-kill-switch vs rule=stochastic, then compare per_position_acceptance and mean_accepted_len. At --temp 0.0 both arms must produce identical output and identical rule=greedy-argmax.

Do not measure this change with speculative_bench. Every pairing it benches is MTP or DFlash, which this PR does not touch, so a before/after would compare the unchanged fallback against itself and produce a clean-looking null. docs/speculative-acceptance.md says this explicitly so the next person does not repeat it.

What this PR does not do

Task 1 of the issue (expose draft probabilities on Drafter, speculative/mtp/, drafter/dflash, drafter/gemma4_assistant) and the MTP/DFlash half of task 3 are not implemented. Drafter::draft_block returns Vec<i32> and MtpTarget::verify_forward returns VerifyForwardOutput { target_tokens: Vec<i32>, .. }; the target logits are a local inside Gemma4MtpTargetAdapter::verify_forward that is dropped at the end of the function, and MtpBatchedVerifyForwardOutput has no logprobs field at all. Wiring stochastic acceptance there means widening both traits and every implementor, on the perf-critical deferred-greedy path, with no Gemma 4 or DFlash checkpoint available in this environment to validate against. Adding that surface without a consumer that can be exercised is how #899 shipped a path that never activated, so it is deliberately left out and recorded in docs/speculative-acceptance.md's per-path table instead. The ArgmaxNoProposalDistribution rule variant exists so those paths have a correct, logged answer when they are wired.

The mtp_policy half of task 3 needs no change: MtpBurstProfile::from_summary consumes MtpAcceptanceSummary, which MtpRoundDiagnostics::record_round fills from the MTP walk's own accepted counts. That plumbing is untouched and keeps working; had the MTP path been converted, its higher acceptance profile would flow through the same counters with no code change.

The issue's docs/benchmark_results/chain-speculative-acceptance-<hw>-<date>.md is not written, because writing a results file without having run the benchmark would be fabrication.

CUDA is unvalidated. No CUDA hardware and no nvcc in this environment. fused_sample_probs uses only backend-agnostic MLX ops and inherits fused_sample's existing single_dtype_scalars (non-Metal) branch through the shared filter chain, but nothing on the CUDA path was compiled or run.

Test plan

• cargo test --release -p mlxcel-core --lib --features metal,accelerate speculative:: — 95 passed, 0 failed.
• cargo test --release --test speculative_parity --features metal,accelerate — passed.
• cargo check --release --lib --tests --features metal,accelerate — clean.
• cargo test --release -p mlxcel-core --lib --features metal,accelerate — 1285 passed, 1 failed (pre-existing fix(tests): fused MoE GeGLU parity threshold is tighter than the gather_qmm reference's own jitter #964; baseline was 1250+1).
• cargo clippy --release -p mlxcel-core --lib --tests --features metal,accelerate -- -D warnings — clean.
• cargo clippy --release --bin mlxcel --features metal,accelerate -- -D warnings — clean.
• End-to-end CLI check on a real Llama 8B/1B pair: all three arms print distinct rule= ids (stochastic, sampler-match, argmax-greedy) and the diagnostic confirms measured per-position acceptance sits at sum_min at temperatures 0.7, 1.0 and 1.5.
• cargo fmt --all -- --check — clean.
• python3 scripts/ci/check_cross_repo_refs.py — exit 0.
• Throughput A/B at temperatures 0.0 / 0.7 / 1.0 — deferred to the maintainer per instruction; commands above.

Closes #902