diff --git a/README.md b/README.md index 09f9d2f49..cb0eb2dbc 100644 --- a/README.md +++ b/README.md @@ -314,6 +314,7 @@ Install with `pip install ./python`. See [Python client](docs/python-client.md) - [Architecture overview](docs/architecture.md) - [Tensor and pipeline parallelism](docs/distributed.md) - [TurboQuant KV cache](docs/turbo-kv-cache.md) +- [Speculative-decoding acceptance](docs/speculative-acceptance.md) - [OpenAI Responses API](docs/responses-api.md) - [Audio input preprocessing](docs/audio-preprocessing.md) - [Python client](docs/python-client.md) diff --git a/docs/README.md b/docs/README.md index 162086cf2..bc4fc68b6 100644 --- a/docs/README.md +++ b/docs/README.md @@ -27,6 +27,7 @@ Current GitHub-facing docs: 12. `block-diffusion.md` — DiffusionGemma block-diffusion generation: canvas denoising vs autoregressive, CLI flags, throughput, and phase 1 limitations. 13. `python-client.md`: the `mlxcel` Python client over the OpenAI-compatible server, covering managed and connect modes, streaming, chat, structured output, the `openai_client` escape hatch, async usage, and troubleshooting. 14. `audio-preprocessing.md` — shared model-input WAV normalization, loaded family policy, resource limits, cancellation, metrics, and current XLA capability boundary. +15. `speculative-acceptance.md` — speculative-decoding acceptance rules, which rule each code path runs, the distribution-preservation guarantee and its RNG dependency, the kill switch, and how to read the active rule off a log. ## Architecture Decision Records diff --git a/docs/architecture.md b/docs/architecture.md index 4787b8124..e6e626f91 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -41,6 +41,8 @@ src/ - `src/lib/mlxcel-core/src/sampling.rs` — penalties and token sampling shared by CLI/server paths. - `src/lib/mlxcel-core/src/generate.rs` — `LanguageModel` trait and generation loops. - `src/lib/mlxcel-core/src/drafter/` and `src/lib/mlxcel-core/src/speculative/` — speculative decoding support. + `speculative/stochastic_accept.rs` holds the acceptance rules and the distribution-preservation + guarantee; see [`speculative-acceptance.md`](speculative-acceptance.md) for which rule each path runs. - `src/lib/mlxcel-core/src/layers.rs`, `src/lib/mlxcel-core/src/weights.rs`, `src/lib/mlxcel-core/src/utils.rs` — model building blocks, SafeTensors loading, masks, and helper operations. - `src/lib/mlxcel-core/src/autotune/`: shape-bucketed kernel autotuner (issue #906), covering the diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 8b102c2b7..7cabfd446 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -165,6 +165,8 @@ The OpenAI audio endpoints (`/v1/audio/speech`, `/v1/audio/transcriptions`, `/v1 | `MLXCEL_ENABLE_MTP_DEFERRED` | `1` | off | **Advanced.** Enables the deferred greedy verifier path for Gemma 4 MTP when sampling settings allow it. | | `MLXCEL_MTP_TICK_SLICE` | `0`/`false`/`no`/`off` to disable, any other value (or unset) to enable | on | Tick-cooperative B=1 MTP serving (issue #734). When on (the default), a B=1 MTP request on the Gemma 4 family is served one speculative round per scheduler tick, alternating with the classic decode/prefill actions, so concurrent classic-decode rows advance between rounds and the head-of-line stall a speculative request imposes drops from the whole burst to about one round (`burst_wall_ms` in the finalize log reports the max single-tick wall). Tokens stream per round instead of in one end-of-burst lump. Set to an off value to restore the legacy run-to-completion burst (the whole request served inside one tick). The interleaving trades roughly 27% of the speculative request's own aggregate decode throughput (cross-tick round gaps) for that bounded stall, so a deployment serving speculative requests without concurrent classic traffic can turn it off to keep the full-throughput burst. Greedy output, acceptance accounting, and every other env gate are unchanged in both modes; DFlash and the batched B>1 paths always run to completion regardless of this flag. | | `MLXCEL_MTP_SLICE_GRANT_ROUNDS` | non-negative integer | `8` | Grant budget for one hold of the tick-slice speculative slot (issue #746), counted in executed slices (slice 0, the prefill + seed, counts as the first slice of a grant). While a slice is in flight, up to 2 further tick-slice-eligible requests park in a grant backlog instead of permanently falling back to classic decode; once the active request has run this many slices with the backlog non-empty, it parks at the next round boundary and the slot is granted to the next request (priority lane first, FIFO within a lane, with an anti-starvation floor: an entry passed over by 2 grant decisions is granted next regardless of lane), so concurrent long streams share speculative acceleration in bounded turns. The budget is read once per grant and per admission decision (cached for the per-round expiry check), so changes apply from the next grant. The budget binds only under contention: a single speculative request never rotates and behaves exactly as under #734. Rotation preserves per-request token streams byte-identically (the drafter is re-armed from the session's own stored verify output at every round). `0` disables rotation and restores the pre-#746 behavior: the active request holds the slot for its whole generation and every concurrent speculative request falls back to classic decode. Unparseable values fall back to the default. | +| `MLXCEL_SPECULATIVE_STOCHASTIC_ACCEPT` | `1`/`true`/`yes`/`on` to enable; unset or anything else to disable | **off** | Acceptance-optimal speculative acceptance for the classic `SpeculativeGenerator` path (offline `mlxcel generate --draft-model`), issue #902. When on, `temperature > 0` verification uses modified rejection sampling (accept the drafted token `t` iff `u * q(t) <= p(t)` for a fresh `u ~ U[0,1)`, and on the first rejection emit a draw from the normalized residual `relu(p - q)`) instead of the default sampler-match rule (accept iff the draft equals an independent draw from the target sampler). **Both rules are distribution-preserving**: the emitted stream is a target-only sample either way, which is the central correction to the issue's premise. What changes is the acceptance probability, which rises from `sum_x p(x) q(x)` to `sum_x min(p(x), q(x))`, the maximal-coupling ceiling for any correct rule. **Opt-in rather than default** because the gain is the ratio between those two quantities and it collapses toward 1 whenever the drafter is confident (`q(t*) ~ 1` makes `min(p, q)` and `p * q` coincide); measured at about 1.02 on a Llama-3.1-8B / Llama-3.2-1B pair at temperature 0.7, which does not pay for two extra full-vocabulary passes and a host sync per verified position. Check the available gain with `MLXCEL_SPECULATIVE_ACCEPT_DIAG=1` before enabling. Enabling changes the RNG stream, so at an equal seed the emitted tokens differ from a default run even though the distribution is identical. Greedy (`temperature == 0` or `top_k == 1`) never reaches either rule and is byte-identical. The Gemma 4 MTP and DFlash round loops are unaffected: they select the target token by argmax regardless of temperature, so this switch is inert there. `SpeculativeGenerator::with_stochastic_acceptance(bool)` overrides it programmatically. Read once per process. See [`speculative-acceptance.md`](speculative-acceptance.md). | +| `MLXCEL_SPECULATIVE_ACCEPT_DIAG` | `1`/`true`/`on`/`yes` to enable; unset or falsy to disable | off | Per-position acceptance diagnostic for the classic speculative path (issue #902). Adds a second stdout line after a speculative run reporting `closed_form_sum_min` (`sum_x min(p(x), q(x))`, the probability modified rejection sampling accepts) and `closed_form_sum_prod` (`sum_x p(x) q(x)`, the probability the pre-#902 rule accepts), both averaged over the verify positions the run actually tested. Because `min(a, b) >= a * b` for `a, b` in `[0, 1]`, `sum_min >= sum_prod` holds for every pair of distributions, so the two numbers turn an acceptance-rate anomaly into an arithmetic statement: the measured `per_position_acceptance` must sit at `sum_min`, and their ratio is the upper bound on any throughput gain the change can produce for that model pair. Enabling it also makes the argmax arm capture the drafter's proposal distribution so both halves of an A/B report the same quantity. Costs two full-vocabulary reductions and a host readback per tested position, which is why it is off by default. Read once per process. See [`speculative-acceptance.md`](speculative-acceptance.md). | ## Generation loop detection (issue #432) diff --git a/docs/speculative-acceptance.md b/docs/speculative-acceptance.md new file mode 100644 index 000000000..598376980 --- /dev/null +++ b/docs/speculative-acceptance.md @@ -0,0 +1,289 @@ +# Speculative-decoding acceptance rules + +Speculative decoding proposes tokens with a cheap draft model and verifies them +with the expensive target model. The *acceptance rule* decides which proposals +survive. It determines two things at once: how many tokens a round emits per +target forward (throughput) and what distribution the emitted stream follows +(correctness). This document states which rule each code path runs, what each +one guarantees, and how to tell from a log which one a given run used. + +## The rules + +Write `p` for the target model's effective next-token distribution at a verify +position (after temperature, token bias, penalties, XTC, and top-k / top-p / +min-p filtering) and `q` for the distribution the drafter actually drew its +proposal from. + +### Sampler-match (the default at `temperature > 0`) + +Accept the drafted token iff it equals a fresh draw from `p`; emit that draw on +both branches. **Lossless**: every emitted token is a fresh conditional draw +from the target, so the stream is exactly a target-only sample. Its acceptance +probability is `sum_x p(x) q(x)`. + +This is what `SpeculativeGenerator` has always done. Issue #902 describes the +verify loop as comparing against the target's *argmax*; for this generator that +is inaccurate, and the difference is the whole point: the code calls +`sample_token_optimized(pos_logits, target_sampling, history)`, which is `argmax` +only when `temperature == 0`. + +### Greedy argmax (`temperature == 0`) + +Sampler-match where the sampler happens to be `argmax`. Lossless by the same +argument; it is exactly greedy target decoding. + +### Argmax-against-argmax (MTP and DFlash only) + +Accept iff the drafted token equals `argmax(target_logits)`, regardless of +temperature; emit the argmax on mismatch. **Not lossless** at `temperature > 0`: +the served stream is the target's greedy stream, not a sample from `p`. Its +acceptance probability is `q(argmax p)`. + +### Modified rejection sampling (opt-in) + +From Leviathan et al. 2023, "Fast Inference from Transformers via Speculative +Decoding" (Algorithm 1) and Chen et al. 2023, "Accelerating Large Language +Model Decoding with Speculative Sampling". + +1. Draw `u ~ U[0, 1)`. Accept the drafted token `t` iff `u * q(t) <= p(t)`. +2. On the first rejection, emit one token drawn from `normalize(relu(p - q))` + and end the chain. +3. If every drafted token was accepted, emit a bonus token drawn from `p`. + +Lossless, because `min(q, p) + relu(p - q) == p` pointwise (the residual's +normalizer cancels the rejection probability exactly). Acceptance probability +`sum_x min(p(x), q(x))`. + +## The acceptance ceiling + +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 | yes | `sum_x p(x) q(x)` | +| Modified rejection sampling | yes | `sum_x min(p(x), q(x))` (the ceiling) | +| Argmax-against-argmax | **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. + +**Consequence, stated plainly:** 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. Issue #902's required outcome ("mean accepted draft length at +temp 0.7 strictly improves vs argmax acceptance") is unachievable in general +and should be read as applying against sampler-match, where the improvement is +real and provable. + +## Which rule each path runs today + +| Path | `temperature == 0` | `temperature > 0` default | `temperature > 0` opt-in | +|------|--------------------|---------------------------|--------------------------| +| `SpeculativeGenerator` (classic; `mlxcel generate --draft-model`) | greedy argmax (lossless) | **sampler-match** (lossless) | modified rejection sampling (lossless, acceptance-optimal) | +| Gemma 4 MTP round loop | argmax (lossless here) | argmax-against-argmax (**biased**) | not wired | +| DFlash round loop | argmax (lossless here) | argmax-against-argmax (**biased**) | not wired | + +The MTP and DFlash verify paths select the target token with +`argmax_per_position` / `argmax_logits_to_array` regardless of the sampler, and +their drafters return token ids without the distribution they were drawn from +(`Drafter::draft_block` returns `Vec`). Making them lossless requires +widening both the drafter interface and `MtpTarget::verify_forward`'s output. +Those are the paths where this feature would buy correctness rather than only +acceptance rate, and where the acceptance trade against argmax is a real +decision rather than a free win. + +## RNG dependency + +Modified rejection sampling draws randomness the previous rules did not: + +* one `U[0, 1)` per verified draft position, for the accept test; +* one categorical draw per rejection, for the residual. + +Both go through MLX's default key sequence, the same stream the fused sampler +consumes, so `mlxcel`'s `--seed` (MLX `random::seed`) still reproduces a run +exactly. But the stream is consumed differently than before: **at an equal seed +the emitted tokens differ from a pre-change run**, in the same sense that +`MLXCEL_SAMPLING_GUMBEL` changes the stream without changing the distribution. +To reproduce a token stream recorded before this landed, set +`MLXCEL_SPECULATIVE_STOCHASTIC_ACCEPT=0`. + +Greedy (`temperature == 0`, or `top_k == 1`) draws nothing new: the rule +selector returns the argmax rule before any distribution tensor is built, so +temperature-0 output is byte-identical. + +## Opt-in, and why it is not the default + +`MLXCEL_SPECULATIVE_STOCHASTIC_ACCEPT=1` (also `true`, `yes`, `on`) enables +modified rejection sampling. It is **off by default**, and +`SpeculativeGenerator::with_stochastic_acceptance(bool)` overrides the +environment programmatically. + +The default is off because on the only path it currently reaches, the rule it +replaces was already lossless. It therefore buys acceptance rate and nothing +else, and the available gain is exactly +`sum_x min(p, q) / sum_x p(x) q(x)`, which collapses toward 1 whenever the +drafter is confident: `q(t*) ~ 1` makes `min(p, q)` and `p * q` coincide. +Measured at about 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 is not worth two percent of theoretical acceptance. + +This follows the same pattern as `MLXCEL_FUSED_QK_NORM` (#326) and the #905 +fusions: the machinery, the proof and the tests land and are available, but the +default does not move until a measurement justifies it. + +## Reading the rule off a run + +**The `mlxcel` CLI installs no `tracing` subscriber by default**, so the +info-level instrumentation below is invisible on the only binary that runs +`SpeculativeGenerator`. The acceptance summary is therefore printed on +**stdout**, unconditionally, after every speculative run: + +``` +[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 the stable identifier, one of `stochastic`, `greedy-argmax`, +`sampler-match`, `sampler-match-no-proposal-distribution`. It is what proves an +A/B arm is not silently the fallback. + +### The two acceptance rates are not interchangeable + +| Field | Definition | Use | +|-------|------------|-----| +| `per_position_acceptance` | `accepted / positions_tested` | **The one theory predicts.** Compare across arms and against `sum_x min(p, q)`. | +| `acceptance_rate` | `accepted / proposed` | How much of each drafted block survived. Moves with block length. | +| `mean_accepted_len` | `accepted / rounds` | Tokens gained per target forward. Drives throughput. | + +`proposed` counts every position of every drafted block, but a chain stops at +its first rejection and never tests the positions behind it. So +`acceptance_rate` is depressed by exactly the amount the chain truncated, and +comparing it across two arms with different chain-length profiles is +meaningless. Confusing these two is enough to make a correct implementation +look like a 2x regression. + +### Optional tracing + +Setting `RUST_LOG` now installs a subscriber for the offline CLI commands +(writing to stderr, so piped stdout stays clean), which surfaces the one-shot +`speculative acceptance rule active` line, the per-call `speculative decode +finished` line, and the drafter auto-detection line. Default output is +unchanged when `RUST_LOG` is unset. `mlxcel serve` is unaffected: it installs +its own subscriber. + +## Verifying the gain before trusting a throughput number + +`MLXCEL_SPECULATIVE_ACCEPT_DIAG=1` adds a second line: + +``` +[Speculative acceptance diagnostic] closed_form_sum_min=0.2213 closed_form_sum_prod=0.2087 (measured per-position acceptance must sit at sum_min, and sum_min >= sum_prod always) +``` + +`closed_form_sum_min` is `sum_x min(p(x), q(x))` averaged over the verify +positions the run actually tested, which is the probability modified rejection +sampling accepts. `closed_form_sum_prod` is `sum_x p(x) q(x)`, the probability +the pre-#902 rule accepts. Because `min(a, b) >= a * b` for `a, b` in `[0, 1]`, +`sum_min >= sum_prod` holds for every pair of distributions, always. The +diagnostic therefore turns "the acceptance rate looks wrong" into an arithmetic +statement: + +* measured `per_position_acceptance` should sit at `closed_form_sum_min`. If it + does not, the accept test is wrong. +* `closed_form_sum_min` should be at or above `closed_form_sum_prod`. If it is + not, `p` or `q` is wrong. + +It costs two full-vocabulary reductions and a host readback per tested +position, so it is off by default. + +### The size of the gain depends on the drafter's entropy, not on the code + +`sum_x min(p, q)` collapses onto `sum_x p(x) q(x)` whenever the drafter is +nearly deterministic: if `q(t*) ~ 1` then `min(p(t*), q(t*)) = p(t*)` and +`p(t*) q(t*) ~ p(t*)`, so the two rules accept at the same rate and this change +buys nothing. The gain appears when `q` is diffuse. Measured on +Llama-3.1-8B-Instruct-4bit with a Llama-3.2-1B-Instruct drafter, 200 tokens, +`--num-draft-tokens 4`, one trajectory each: + +| 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, which is the +correctness result. The *ratio* column is the throughput headroom, and on this +pair at 0.7 there is essentially none: the 1B drafter is close to deterministic +there. Between-trajectory variance in the agreement profile is large (two +200-token runs at the same temperature differed by nearly 2x in `sum_min`), so +a single short throughput run at 0.7 on this pair measures noise, not the +change. Read `closed_form_sum_min / closed_form_sum_prod` first: it is the +upper bound on what any throughput measurement can show. + +The temperature 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; do not read it as a typical production number. + +## Measuring the change + +The classic `SpeculativeGenerator` is the path this rule changed, and it is +reachable only from the offline CLI (the server routes speculative work through +the MTP and DFlash round loops). Run the A/B against `mlxcel generate` with a +draft model, not against `speculative_bench`, whose pairings are all MTP or +DFlash and would compare the unchanged fallback against itself: + +```bash +# Baseline arm: the default sampler-match rule. +MLXCEL_SPECULATIVE_ACCEPT_DIAG=1 \ + ./target/release/mlxcel generate \ + -m models/ --draft-model models/ --num-draft-tokens 4 \ + --temp 0.7 --seed 1234 -n 256 -p "" + +# Opt-in arm: modified rejection sampling. +MLXCEL_SPECULATIVE_STOCHASTIC_ACCEPT=1 MLXCEL_SPECULATIVE_ACCEPT_DIAG=1 \ + ./target/release/mlxcel generate \ + -m models/ --draft-model models/ --num-draft-tokens 4 \ + --temp 0.7 --seed 1234 -n 256 -p "" +``` + +Confirm the `rule=` field differs between arms (`sampler-match` vs `stochastic`), then compare `per_position_acceptance` and `mean_accepted_len`. +Repeat at `--temp 0.0`, where both arms must produce identical output and +identical `rule=greedy-argmax`. + +**A note on the drafter-kind line.** With no explicit `--draft-kind`, the CLI +prints the drafter kind it *auto-detected* from the drafter's `config.json`, +which is frequently `dflash`, and then ignores it and runs the classic +`SpeculativeGenerator` anyway. A second line now states which path actually +runs: + +``` +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** get you a DFlash round loop in +`mlxcel generate`; it returns an error saying the offline path does not +construct one. + +## Regression guards + +In `src/lib/mlxcel-core/src/speculative/`: + +* `distribution_tests.rs::speculative_stream_matches_the_target_distribution` — + chi-square goodness of fit of the emitted stream against the target's exact + categorical, over 12000 tokens (`MLXCEL_SPEC_DIST_SAMPLES` overrides), dof 5, + rejected above 25.7448 (alpha = 1e-4). +* `distribution_tests.rs::chi_square_rejects_known_wrong_acceptance_rules` — + power calibration. The same statistic must reject a rule that resamples from + `p` instead of the residual, and an argmax-against-argmax rule. +* `distribution_tests.rs::filtered_target_support_is_never_violated` — zero + tolerance: one emitted token outside a `top_k = 2` target's support fails. +* `distribution_tests.rs::temperature_zero_stream_is_byte_identical_to_greedy_target_only` + — zero tolerance on greedy output. +* `distribution_tests.rs::kv_cache_length_is_exact_in_both_termination_regimes` + — cache rewind arithmetic, pinned in the all-reject and all-accept regimes. +* `stochastic_accept_tests.rs` — the accept rate against `min(1, p/q)`, the + residual against `normalize(relu(p - q))`, and the `sum min(p, q)` optimum. diff --git a/src/commands/generate.rs b/src/commands/generate.rs index e4fe4dd1e..8890c274a 100644 --- a/src/commands/generate.rs +++ b/src/commands/generate.rs @@ -1300,6 +1300,12 @@ fn run_generation_mode( println!("Loading draft model from {:?}...", draft_model_path); let (draft_model, _draft_tokenizer) = select_backend().load_model(draft_model_path)?; println!("Draft model loaded."); + // This line reports what the drafter's `config.json` *auto-detected*, + // which is not the same thing as the path this command will run. With + // no explicit `--draft-kind`, an auto-detected `dflash` still falls + // through to the classic `SpeculativeGenerator` below. Printing only + // the detected kind made a benchmark run indistinguishable from a + // DFlash round-loop run, so the path actually taken is now printed too. println!( "Resolved drafter kind: {} (block_size = {block_size}{})", resolved_kind, @@ -1309,6 +1315,13 @@ fn run_generation_mode( ", default" }, ); + if !user_requested_explicit_kind { + println!( + "Drafter kind was auto-detected only; running the classic \ + SpeculativeGenerator path (pass --draft-kind explicitly to select a \ + kind-specific round loop)." + ); + } // MTP is handled above (issue #166). The remaining explicit kinds // (DFlash, InternalMtp) still need their kind-specific round loops and @@ -1369,14 +1382,24 @@ fn run_generation_mode( let mut spec_generator = SpeculativeGenerator::new(main_num_layers, draft_num_layers) .with_token_bias(token_bias); - spec_generator.generate( + let result = spec_generator.generate( model, &draft_model, prompt_tokens, args.generation.max_tokens, args.model.num_draft_tokens, sampling_config, - ) + ); + + // The CLI installs no tracing subscriber, so the info-level acceptance + // instrumentation inside `SpeculativeGenerator` never reaches a + // terminal. Print the summary on stdout unconditionally: it names the + // acceptance rule that actually ran and the mean accepted draft length, + // which are the two facts an acceptance-rate A/B has to be able to + // state about itself. + println!("{}", spec_generator.acceptance_stats().summary_line()); + + result } else if let Some(embeddings) = vlm_embeddings { generate_with_embeddings( model, diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp index eca9d085e..84902b874 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp @@ -4583,30 +4583,18 @@ static bool gumbel_sample_applies( // - Categorical sampling: temp scaling + random::categorical in one kernel // - Min-p filtering: softmax + max + mask in one kernel // - Top-p filtering: uncompiled (cumsum lacks output_shapes in MLX v0.31.x) -static std::unique_ptr fused_sample_impl( - const MlxArray& logits, +// +// The filter chain itself lives in `fused_sample_filter_logits` so that +// `fused_sample_probs` (issue #902) can reproduce the sampler's support and +// normalization *by construction* rather than by a parallel reimplementation +// that could drift. +static mlx::core::array fused_sample_filter_logits( + mlx::core::array x, float temperature, int32_t top_k, float top_p, - float min_p, - bool allow_gumbel + float min_p ) { - auto x = logits.inner; - - // Greedy path: argmax - if (temperature == 0.0f || top_k == 1) { - return std::make_unique(mlx::core::argmax(x, -1, false)); - } - - // Softmax-free stochastic sampling (#900). Batch-wide: one launch covers - // every row of `[B, vocab]`, so the batched fused sampler needs no per-row - // loop here. Sampled ids differ from the `categorical` path at equal seeds - // (a different RNG consumer), but the distribution is identical. - if (allow_gumbel && gumbel_sample_applies(x, temperature, top_k, top_p, min_p)) { - return std::make_unique( - mlxcel::turbo::gumbel_max_sample(x, temperature)); - } - // Build sampler scalars in the logit dtype on non-Metal backends only // (issue #636). On CUDA the bf16 promotion patch keeps bf16 op boundaries // in bf16, but a bare `array(f32)` scalar against a bf16/f16 logit tensor @@ -4663,6 +4651,36 @@ static std::unique_ptr fused_sample_impl( x = std::move(result[0]); } + return x; +} + +static std::unique_ptr fused_sample_impl( + const MlxArray& logits, + float temperature, + int32_t top_k, + float top_p, + float min_p, + bool allow_gumbel +) { + auto x = logits.inner; + + // Greedy path: argmax + if (temperature == 0.0f || top_k == 1) { + return std::make_unique(mlx::core::argmax(x, -1, false)); + } + + // Softmax-free stochastic sampling (#900). Batch-wide: one launch covers + // every row of `[B, vocab]`, so the batched fused sampler needs no per-row + // loop here. Sampled ids differ from the `categorical` path at equal seeds + // (a different RNG consumer), but the distribution is identical. + if (allow_gumbel && gumbel_sample_applies(x, temperature, top_k, top_p, min_p)) { + return std::make_unique( + mlxcel::turbo::gumbel_max_sample(x, temperature)); + } + + x = fused_sample_filter_logits( + std::move(x), temperature, top_k, top_p, min_p); + // Categorical sampling (not compiled — random ops need known shapes at trace time) return std::make_unique(mlx::core::random::categorical(x, -1)); } @@ -4688,6 +4706,49 @@ std::unique_ptr fused_sample_categorical( return fused_sample_impl(logits, temperature, top_k, top_p, min_p, false); } +// The exact categorical distribution `fused_sample` draws from, as a float32 +// `[batch, vocab]` row-normalized probability tensor (issue #902). +// +// Correctness contract: the filter chain is *the same code* `fused_sample` +// runs (`fused_sample_filter_logits`), so the support and the relative masses +// cannot drift from the sampler. Both stochastic sampler arms agree with it: +// `random::categorical(x, -1)` draws from `softmax(x, -1)` by definition, and +// the Gumbel-max kernel draws from `softmax(logits / temperature)`, which is +// what the filter chain leaves when no top-k / top-p / min-p narrows the row. +// +// A greedy sampler configuration (`temperature == 0` or `top_k == 1`) draws a +// point mass, so this returns the one-hot indicator at the argmax. That is the +// degenerate `p_draft` a greedily-proposing drafter needs for the modified +// rejection-sampling accept test. +// +// The softmax runs in float32 regardless of the logit dtype: the caller reads +// individual entries back to the host and compares them against a uniform +// draw, where an f16 probability (which underflows below ~6e-8) would turn a +// legitimately small acceptance ratio into an unconditional accept. +std::unique_ptr fused_sample_probs( + const MlxArray& logits, + float temperature, + int32_t top_k, + float top_p, + float min_p +) { + auto x = logits.inner; + + if (temperature == 0.0f || top_k == 1) { + auto idx = mlx::core::argmax(x, -1, /*keepdims=*/true); + const auto vocab = static_cast(x.shape(-1)); + auto ids = mlx::core::arange(0.0, vocab, mlx::core::int32); + auto one_hot = mlx::core::equal(ids, mlx::core::astype(idx, mlx::core::int32)); + return std::make_unique( + mlx::core::astype(one_hot, mlx::core::float32)); + } + + x = fused_sample_filter_logits( + std::move(x), temperature, top_k, top_p, min_p); + return std::make_unique( + mlx::core::softmax(mlx::core::astype(x, mlx::core::float32), -1, /*precise=*/true)); +} + std::unique_ptr gumbel_max_sample( const MlxArray& logits, float temperature diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index c92a49089..86974d871 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -1294,6 +1294,19 @@ std::unique_ptr fused_sample_categorical( float min_p ); +// The exact categorical distribution `fused_sample` draws from (#902), as a +// float32 [batch, vocab] row-normalized probability tensor. Shares the filter +// chain with `fused_sample` so support and masses cannot drift. A greedy +// configuration (temperature == 0 or top_k == 1) returns the one-hot argmax +// indicator. +std::unique_ptr fused_sample_probs( + const MlxArray& logits, + float temperature, + int32_t top_k, + float top_p, + float min_p +); + // Softmax-free Gumbel-max categorical sampling (#900), called directly. // `logits` is 2D [batch, vocab]; returns [batch] uint32 token ids. std::unique_ptr gumbel_max_sample( diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index fed1841fc..68988d414 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -2053,6 +2053,27 @@ mod ffi { min_p: f32, ) -> UniquePtr; + /// The exact categorical distribution [`fused_sample`] draws from + /// (#902), as a float32 `[batch, vocab]` row-normalized probability + /// tensor. Shares the filter chain with [`fused_sample`] in C++, so + /// its support and relative masses cannot drift from the sampler. + /// + /// A greedy configuration (`temperature == 0.0` or `top_k == 1`) + /// returns the one-hot indicator at the argmax, which is the + /// degenerate proposal distribution of a greedily-proposing drafter. + /// + /// The softmax runs in float32 regardless of the logit dtype: entries + /// are read back to the host and compared against a uniform draw, + /// where an f16 probability would underflow below ~6e-8 and turn a + /// legitimately small acceptance ratio into an unconditional accept. + fn fused_sample_probs( + logits: &MlxArray, + temperature: f32, + top_k: i32, + top_p: f32, + min_p: f32, + ) -> UniquePtr; + /// Softmax-free Gumbel-max categorical sampling (#900), called /// directly. `logits` is 2-D `[batch, vocab]`; returns a `[batch]` /// uint32 token-id array drawn from `softmax(logits / temperature)`. diff --git a/src/lib/mlxcel-core/src/sampling.rs b/src/lib/mlxcel-core/src/sampling.rs index d47e28711..80b0d1d81 100644 --- a/src/lib/mlxcel-core/src/sampling.rs +++ b/src/lib/mlxcel-core/src/sampling.rs @@ -298,8 +298,38 @@ fn sample_token_optimized_core( logits: &MlxArray, config: &SamplingConfig, token_history: &[i32], - mut state: Option<&mut SamplerState>, + state: Option<&mut SamplerState>, ) -> (UniquePtr, UniquePtr) { + let last_logits = preprocess_logits_for_sampling(logits, config, token_history, state); + + let token = ffi::fused_sample( + &last_logits, + config.temperature, + config.top_k, + config.top_p, + config.min_p, + ); + (token, last_logits) +} + +/// Everything [`sample_token_optimized`] does to the logits *before* the fused +/// temperature / top-k / top-p / min-p sampler: last-position slice, token +/// bias, repetition penalty, DRY, frequency/presence penalty, XTC. +/// +/// Split out so the speculative acceptance path (issue #902) can obtain the +/// exact same pre-sampler logits and hand them to [`ffi::fused_sample_probs`], +/// which then applies the identical filter chain the sampler itself runs. The +/// effective categorical distribution is therefore derived from the sampler's +/// own code rather than reconstructed alongside it. +/// +/// Used by: [`sample_token_optimized_core`], [`effective_token_distribution`], +/// [`sample_token_with_distribution`] +fn preprocess_logits_for_sampling( + logits: &MlxArray, + config: &SamplingConfig, + token_history: &[i32], + mut state: Option<&mut SamplerState>, +) -> UniquePtr { // Use optimized slice_last_logits: [batch, seq, vocab] -> [batch, vocab]. let last_logits = ffi::slice_last_logits(logits); @@ -385,20 +415,78 @@ fn sample_token_optimized_core( // default disabled state and skips this entirely — no array ops, and no // draw from the per-request RNG stream, which keeps every existing // request's token stream byte-identical to before this feature existed. - let last_logits = if config.xtc_probability > 0.0 { + if config.xtc_probability > 0.0 { apply_xtc_step(&last_logits, config) } else { last_logits - }; + } +} + +/// The exact categorical distribution [`sample_token_optimized`] would draw +/// from for `logits` under `config`, as a float32 `[batch, vocab]` +/// row-normalized probability tensor. +/// +/// This is the `p` and `q` of modified rejection sampling (issue #902). It is +/// built from the sampler's own pre-steps ([`preprocess_logits_for_sampling`]) +/// and the sampler's own filter chain ([`ffi::fused_sample_probs`]), so a +/// change to either automatically moves the distribution with it. +/// +/// A greedy config (`temperature == 0.0` or `top_k == 1`) yields the one-hot +/// indicator at the argmax, which is the correct degenerate proposal +/// distribution for a greedily-proposing drafter. +/// +/// Consumes no randomness: safe to call without perturbing the token stream of +/// any other sampler on the same RNG key sequence. +/// +/// Used by: `speculative::stochastic_accept` +pub fn effective_token_distribution( + logits: &MlxArray, + config: &SamplingConfig, + token_history: &[i32], +) -> UniquePtr { + let processed = preprocess_logits_for_sampling(logits, config, token_history, None); + ffi::fused_sample_probs( + &processed, + config.temperature, + config.top_k, + config.top_p, + config.min_p, + ) +} +/// Sample one token *and* return the distribution it was drawn from. +/// +/// Equivalent to [`sample_token_optimized`] followed by +/// [`effective_token_distribution`], but the (potentially expensive) bias and +/// penalty pre-steps run once instead of twice. The returned token is the same +/// array [`sample_token_optimized`] returns for the same inputs and RNG state: +/// the extra distribution tensor is a pure function of the pre-sampler logits +/// and draws nothing from the RNG stream. +/// +/// Returns `(token, probs)` where `probs` is float32 `[batch, vocab]`. +/// +/// Used by: `SpeculativeGenerator` draft loop (issue #902) +pub fn sample_token_with_distribution( + logits: &MlxArray, + config: &SamplingConfig, + token_history: &[i32], +) -> (UniquePtr, UniquePtr) { + let processed = preprocess_logits_for_sampling(logits, config, token_history, None); let token = ffi::fused_sample( - &last_logits, + &processed, config.temperature, config.top_k, config.top_p, config.min_p, ); - (token, last_logits) + let probs = ffi::fused_sample_probs( + &processed, + config.temperature, + config.top_k, + config.top_p, + config.min_p, + ); + (token, probs) } /// Batch-parallel sampling: sample one token per sequence from batched logits. diff --git a/src/lib/mlxcel-core/src/speculative/distribution_tests.rs b/src/lib/mlxcel-core/src/speculative/distribution_tests.rs new file mode 100644 index 000000000..18461a88f --- /dev/null +++ b/src/lib/mlxcel-core/src/speculative/distribution_tests.rs @@ -0,0 +1,1062 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! End-to-end distributional regression guard for speculative decoding +//! (issue #902). +//! +//! # What is being tested +//! +//! **Null hypothesis (H0):** the token stream emitted by +//! [`SpeculativeGenerator::generate`] at `temperature > 0` is distributed +//! identically to sampling the target model alone at the same temperature. +//! +//! The test models make H0 checkable in closed form. The target's logits do +//! not depend on its input, so target-only sampling emits i.i.d. draws from a +//! single known categorical `p = softmax(target_logits / temperature)`. 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`. A Pearson chi-square goodness-of-fit +//! test against `p` therefore tests H0 directly. +//! +//! **Statistic:** `sum_x (O_x - N p_x)^2 / (N p_x)` over the `V = 6` vocabulary +//! entries, `dof = V - 1 = 5`. +//! +//! **Sample count:** `SAMPLE_TOKENS` tokens, which the committed default sets +//! to 12000 (override with `MLXCEL_SPEC_DIST_SAMPLES`). Every emitted token is +//! one sample; a run collects them across as many `generate` calls as needed. +//! +//! **Threshold:** `CHI2_DOF5_ALPHA_1E4 = 25.7448`, the upper-tail chi-square +//! quantile at `alpha = 1e-4` for 5 degrees of freedom. A correct +//! implementation fails once in ten thousand runs. +//! +//! # Why this test can fail +//! +//! A goodness-of-fit test is worthless without demonstrated power, so the +//! power is measured in-tree rather than asserted. `chi_square_rejects_*` +//! replays the identical harness with two deliberately wrong acceptance rules +//! and requires the statistic to exceed the same threshold: +//! +//! * `residual_replaced_by_target_resample` — the single most common way to +//! get this algorithm wrong: on rejection, draw from `p` instead of from +//! `normalize(relu(p - q))`. Expected statistic at N = 12000 is roughly 360, +//! fourteen times the threshold. +//! * `argmax_target` — the rule the MTP and DFlash verify paths still run +//! (`Gemma4MtpTargetAdapter::verify_forward` calls `argmax_per_position` +//! regardless of temperature). With a context-independent target it collapses +//! the stream onto the argmax; the statistic runs into the thousands. This is +//! the bias the issue describes, recorded here rather than gated. +//! +//! Chi-square is blind to a single wrong token, so the zero-tolerance +//! assertions live beside it: `filtered_target_support_is_never_violated` +//! fails on one emitted token outside the target's top-k support, and the +//! temperature-0 test fails on one differing token. + +use super::*; +use crate::dtype; +use crate::sampling::effective_token_distribution; + +/// Upper-tail chi-square quantile at alpha = 1e-4, 5 degrees of freedom. +const CHI2_DOF5_ALPHA_1E4: f64 = 25.7448; + +/// Tokens collected by the committed default run of the distributional test. +/// +/// Sized from the measured power of the two wrong-rule arms: at 12000 the +/// residual bug lands near 360 against a threshold of 25.74. Raise it with +/// `MLXCEL_SPEC_DIST_SAMPLES` when investigating. +const SAMPLE_TOKENS: usize = 12_000; + +const VOCAB: usize = 6; +const TEMPERATURE: f32 = 0.7; + +/// Target logits. Deliberately not flat, so the chi-square has cells with very +/// different expected counts and a bias in either direction shows up. +const TARGET_LOGITS: [f32; VOCAB] = [2.0, 1.0, 0.5, 0.0, -0.5, -1.0]; +/// Draft logits: a different model that agrees with the target often enough to +/// be a realistic drafter (`sum min(p, q)` is about 0.71) but disagrees on the +/// ordering of the top two tokens. +const DRAFT_LOGITS: [f32; VOCAB] = [1.4, 1.3, 0.2, 0.6, -0.3, -0.8]; + +fn sample_budget() -> usize { + std::env::var("MLXCEL_SPEC_DIST_SAMPLES") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(SAMPLE_TOKENS) +} + +fn softmax(logits: &[f32], temperature: f32) -> Vec { + let scaled: Vec = logits + .iter() + .map(|&x| f64::from(x) / f64::from(temperature)) + .collect(); + let max = scaled.iter().cloned().fold(f64::NEG_INFINITY, f64::max); + let exp: Vec = scaled.iter().map(|x| (x - max).exp()).collect(); + let sum: f64 = exp.iter().sum(); + exp.iter().map(|x| x / sum).collect() +} + +fn chi_square(observed: &[u32], expected: &[f64]) -> f64 { + let n = f64::from(observed.iter().sum::()); + observed + .iter() + .zip(expected) + .map(|(&o, &e)| { + let e = n * e; + let d = f64::from(o) - e; + d * d / e + }) + .sum() +} + +fn stochastic_config() -> SamplingConfig { + SamplingConfig { + temperature: TEMPERATURE, + ..SamplingConfig::default() + } +} + +// --------------------------------------------------------------------------- +// Test models +// --------------------------------------------------------------------------- + +/// A model whose logits are the same at every position and independent of its +/// input, so its next-token distribution is one fixed categorical. +/// +/// Keeping the target context-independent is what makes H0 a plain i.i.d. +/// goodness-of-fit statement instead of a statement about a sequence model. +struct ContextFreeModel { + logits: Vec, + eos_tokens: Vec, +} + +impl ContextFreeModel { + fn new(logits: &[f32]) -> Self { + Self { + logits: logits.to_vec(), + eos_tokens: Vec::new(), + } + } +} + +impl LanguageModel for ContextFreeModel { + fn forward( + &self, + input_ids: &MlxArray, + caches: &mut [KVCache], + _mask: Option<&MlxArray>, + ) -> UniquePtr { + let shape = ffi::array_shape(input_ids); + let (batch, seq_len) = (shape[0] as usize, shape[1] as usize); + // Advance the KV caches so cache-offset bookkeeping and the rewind on + // rejection are exercised exactly as they are with a real model. + for cache in caches.iter_mut() { + let k = ffi::ones(&[shape[0], 2, shape[1], 4], dtype::FLOAT32); + let v = ffi::ones(&[shape[0], 2, shape[1], 4], dtype::FLOAT32); + cache.update(k, v); + } + let mut out = Vec::with_capacity(batch * seq_len * self.logits.len()); + for _ in 0..(batch * seq_len) { + out.extend_from_slice(&self.logits); + } + ffi::from_slice_f32(&out, &[shape[0], shape[1], self.logits.len() as i32]) + } + + fn make_caches(&self) -> Vec { + vec![KVCache::new()] + } + + fn num_layers(&self) -> usize { + 1 + } + + fn eos_token_ids(&self) -> Vec { + self.eos_tokens.clone() + } +} + +/// A deterministic sequence model: the logit peak at each position is a +/// function of the token *at* that position, so greedy decoding walks a fixed +/// orbit. Used for the temperature-0 byte-identity test, where a +/// context-independent model would make any acceptance rule trivially right. +struct OrbitModel { + multiplier: i32, + offset: i32, + peak: f32, +} + +impl OrbitModel { + fn next(&self, token: i32) -> i32 { + (token * self.multiplier + self.offset).rem_euclid(VOCAB as i32) + } +} + +impl LanguageModel for OrbitModel { + fn forward( + &self, + input_ids: &MlxArray, + caches: &mut [KVCache], + _mask: Option<&MlxArray>, + ) -> UniquePtr { + let shape = ffi::array_shape(input_ids); + ffi::eval(input_ids); + let bytes = ffi::array_to_raw_bytes(input_ids); + let tokens: Vec = bytes + .chunks_exact(4) + .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + for cache in caches.iter_mut() { + let k = ffi::ones(&[shape[0], 2, shape[1], 4], dtype::FLOAT32); + let v = ffi::ones(&[shape[0], 2, shape[1], 4], dtype::FLOAT32); + cache.update(k, v); + } + let mut out = vec![0.0f32; tokens.len() * VOCAB]; + for (i, &tok) in tokens.iter().enumerate() { + out[i * VOCAB + self.next(tok) as usize] = self.peak; + } + ffi::from_slice_f32(&out, &[shape[0], shape[1], VOCAB as i32]) + } + + fn make_caches(&self) -> Vec { + vec![KVCache::new()] + } + + fn num_layers(&self) -> usize { + 1 + } + + fn eos_token_ids(&self) -> Vec { + Vec::new() + } +} + +// --------------------------------------------------------------------------- +// The main distributional test +// --------------------------------------------------------------------------- + +/// Collect emitted tokens from repeated speculative generations until the +/// sample budget is met. The first token of each call comes from the prefill +/// sample, which is a plain draw from `p` and therefore also a valid sample. +fn collect_speculative_tokens(budget: usize, num_draft: usize) -> Vec { + let target = ContextFreeModel::new(&TARGET_LOGITS); + let draft = ContextFreeModel::new(&DRAFT_LOGITS); + let config = stochastic_config(); + let mut counts = vec![0u32; VOCAB]; + let mut collected = 0usize; + + while collected < budget { + let mut generator = SpeculativeGenerator::new(target.num_layers(), draft.num_layers()) + .with_stochastic_acceptance(true); + let per_call = 64.min(budget - collected); + let (tokens, _stats) = + generator.generate(&target, &draft, &[0], per_call, num_draft, &config); + assert!( + !tokens.is_empty(), + "the generator must make progress on every call" + ); + for t in &tokens { + assert!( + (0..VOCAB as i32).contains(t), + "emitted token {t} is outside the vocabulary" + ); + counts[*t as usize] += 1; + } + collected += tokens.len(); + } + counts +} + +/// H0: the speculative stream at `temperature > 0` is distributed as +/// target-only sampling. Rejected above the alpha = 1e-4 chi-square threshold. +/// +/// See the module docs for the full statement, the sample count and the power +/// calibration that makes this assertion meaningful. +#[test] +fn speculative_stream_matches_the_target_distribution() { + let expected = softmax(&TARGET_LOGITS, TEMPERATURE); + let budget = sample_budget(); + let counts = collect_speculative_tokens(budget, 4); + let n: u32 = counts.iter().sum(); + let chi2 = chi_square(&counts, &expected); + + let empirical: Vec = counts + .iter() + .map(|&c| f64::from(c) / f64::from(n)) + .collect(); + assert!( + chi2 < CHI2_DOF5_ALPHA_1E4, + "speculative stream is distinguishable from target-only sampling: \ + chi-square {chi2:.3} exceeds the alpha=1e-4 dof=5 threshold {CHI2_DOF5_ALPHA_1E4}. \ + n={n}, empirical={empirical:?}, expected={expected:?}" + ); +} + +/// The guarantee must not depend on the draft block length: a longer chain +/// runs the accept test more times per verify forward and reaches deeper +/// positions, where a context-bookkeeping error would surface. +#[test] +fn speculative_stream_matches_the_target_distribution_at_block_size_one() { + let expected = softmax(&TARGET_LOGITS, TEMPERATURE); + let counts = collect_speculative_tokens(sample_budget() / 2, 1); + let chi2 = chi_square(&counts, &expected); + assert!( + chi2 < CHI2_DOF5_ALPHA_1E4, + "num_draft=1 stream is distinguishable from the target: chi-square {chi2:.3}, \ + counts {counts:?}" + ); +} + +// --------------------------------------------------------------------------- +// Power calibration: the same statistic must reject known-wrong rules +// --------------------------------------------------------------------------- + +/// Replay one verify position under a chosen acceptance rule, returning the +/// emitted token. `p` and `q` are the fixed context-free distributions, so a +/// full chain is not needed to reproduce the emitted-token distribution: every +/// position of every chain sees the same pair. +fn emit_one_token( + rule: &str, + target_probs: &MlxArray, + draft_probs: &MlxArray, + draft_token: i32, + argmax_token: i32, +) -> i32 { + match rule { + "stochastic" => { + match stochastic_accept::verify_draft_token(target_probs, draft_probs, draft_token) { + DraftVerdict::Accept => draft_token, + DraftVerdict::Reject { replacement } => replacement, + } + } + // The common implementation mistake: reject correctly, then draw the + // replacement from `p` instead of from the residual. + "residual_replaced_by_target_resample" => { + if stochastic_accept::accept_draft_token(target_probs, draft_probs, draft_token) { + draft_token + } else { + let tok = ffi::fused_sample(&ffi::log(target_probs), 1.0, 0, 1.0, 0.0); + ffi::eval(&tok); + ffi::item_i32(&tok) + } + } + // The rule the MTP and DFlash verify paths still run: compare against + // the target's argmax and emit the argmax on mismatch. + "argmax_target" => argmax_token, + other => panic!("unknown rule {other}"), + } +} + +/// Draw one token from a `[1, vocab]` probability row. +fn draw_from(probs: &MlxArray) -> i32 { + let tok = ffi::fused_sample(&ffi::log(probs), 1.0, 0, 1.0, 0.0); + ffi::eval(&tok); + ffi::item_i32(&tok) +} + +fn power_arm_chi_square(rule: &str, trials: usize) -> f64 { + let target_logits = ffi::from_slice_f32(&TARGET_LOGITS, &[1, 1, VOCAB as i32]); + let draft_logits = ffi::from_slice_f32(&DRAFT_LOGITS, &[1, 1, VOCAB as i32]); + let config = stochastic_config(); + let target_probs = effective_token_distribution(&target_logits, &config, &[]); + let draft_probs = effective_token_distribution(&draft_logits, &config, &[]); + ffi::eval(&target_probs); + ffi::eval(&draft_probs); + + let argmax_token = TARGET_LOGITS + .iter() + .enumerate() + .max_by(|a, b| a.1.partial_cmp(b.1).unwrap()) + .map(|(i, _)| i as i32) + .unwrap(); + + let mut counts = vec![0u32; VOCAB]; + for _ in 0..trials { + let draft_token = draw_from(&draft_probs); + let emitted = emit_one_token(rule, &target_probs, &draft_probs, draft_token, argmax_token); + counts[emitted as usize] += 1; + } + chi_square(&counts, &softmax(&TARGET_LOGITS, TEMPERATURE)) +} + +/// Power calibration. The committed acceptance rule must pass the same +/// statistic that rejects both wrong rules; otherwise the goodness-of-fit test +/// above proves nothing. +/// +/// Uses a smaller trial count than the end-to-end test because each trial here +/// is one verify position rather than a whole `generate` call. The margins are +/// large enough that the conclusion is not sensitive to the exact count. +#[test] +fn chi_square_rejects_known_wrong_acceptance_rules() { + const TRIALS: usize = 4000; + + let correct = power_arm_chi_square("stochastic", TRIALS); + assert!( + correct < CHI2_DOF5_ALPHA_1E4, + "the committed rule must pass its own statistic: chi-square {correct:.3}" + ); + + let residual_bug = power_arm_chi_square("residual_replaced_by_target_resample", TRIALS); + assert!( + residual_bug > CHI2_DOF5_ALPHA_1E4, + "resampling from p instead of the residual must be detected, but the \ + chi-square was only {residual_bug:.3} against a {CHI2_DOF5_ALPHA_1E4} threshold: \ + the goodness-of-fit test above has no power and cannot be trusted" + ); + + let argmax_bug = power_arm_chi_square("argmax_target", TRIALS); + assert!( + argmax_bug > CHI2_DOF5_ALPHA_1E4, + "argmax-target acceptance must be detected, but the chi-square was only \ + {argmax_bug:.3}" + ); + + // Recorded, not gated: the margin by which each wrong rule is rejected. + println!( + "power calibration at n={TRIALS}: correct={correct:.2}, \ + residual-bug={residual_bug:.2}, argmax-target={argmax_bug:.2}, \ + threshold={CHI2_DOF5_ALPHA_1E4}" + ); +} + +// --------------------------------------------------------------------------- +// Zero-tolerance guards (chi-square cannot see a single wrong token) +// --------------------------------------------------------------------------- + +/// One emitted token outside the target's filtered support fails this test. +/// +/// With `top_k = 2` the target assigns exactly zero mass to four of the six +/// tokens, while the unfiltered drafter proposes them roughly half the time. +/// A missing `p(t) > 0` conjunct in the accept test, or a residual that leaked +/// mass outside `p`'s support, shows up here on the first offending token +/// rather than as a shifted histogram. +#[test] +fn filtered_target_support_is_never_violated() { + let target = ContextFreeModel::new(&TARGET_LOGITS); + let draft = ContextFreeModel::new(&DRAFT_LOGITS); + let config = SamplingConfig { + temperature: TEMPERATURE, + top_k: 2, + ..SamplingConfig::default() + }; + // Top-2 of TARGET_LOGITS. + let allowed = [0i32, 1]; + + let mut emitted = 0usize; + while emitted < 1500 { + let mut generator = SpeculativeGenerator::new(target.num_layers(), draft.num_layers()) + .with_stochastic_acceptance(true); + let (tokens, _) = generator.generate(&target, &draft, &[0], 64, 4, &config); + for t in &tokens { + assert!( + allowed.contains(t), + "emitted token {t} lies outside the target's top-2 support {allowed:?}; \ + the target sampler could never have produced it" + ); + } + emitted += tokens.len(); + } +} + +// --------------------------------------------------------------------------- +// Temperature 0 and cache rewind +// --------------------------------------------------------------------------- + +/// The target model used by the temperature-0 and cache-rewind tests: +/// `next(t) = t + 1 (mod 6)`, a full 6-cycle rather than a short orbit. +fn greedy_target() -> OrbitModel { + OrbitModel { + multiplier: 1, + offset: 1, + peak: 8.0, + } +} + +/// A drafter that agrees with [`greedy_target`] exactly when the current token +/// is 0 (`2t + 1 == t + 1 (mod 6)` iff `t == 0`), so a run over the target's +/// 6-cycle produces a fixed 1-in-6 mix of accepted and rejected positions. +fn mixed_draft() -> OrbitModel { + OrbitModel { + multiplier: 2, + offset: 1, + peak: 8.0, + } +} + +/// A drafter that never agrees with [`greedy_target`]: `t + 4` and `t + 1` +/// differ for every `t`. +fn always_rejecting_draft() -> OrbitModel { + OrbitModel { + multiplier: 1, + offset: 4, + peak: 8.0, + } +} + +/// At temperature 0 the speculative stream must equal the target model's own +/// greedy continuation, token for token. One differing token fails. +/// +/// The target is a deterministic 6-cycle, so the reference stream is computable +/// without a second generator. The drafter agrees on exactly one token in six, +/// which forces both the accept and the reject-plus-rewind branches to run +/// many times inside a single generation. +#[test] +fn temperature_zero_stream_is_byte_identical_to_greedy_target_only() { + let target = greedy_target(); + const PROMPT: i32 = 2; + const N: usize = 48; + + let mut expected = Vec::with_capacity(N); + let mut tok = PROMPT; + for _ in 0..N { + tok = target.next(tok); + expected.push(tok); + } + + for draft in [mixed_draft(), always_rejecting_draft()] { + for num_draft in [1usize, 2, 4] { + let mut generator = SpeculativeGenerator::new(target.num_layers(), draft.num_layers()); + let (tokens, _) = generator.generate( + &target, + &draft, + &[PROMPT], + N, + num_draft, + &SamplingConfig::greedy(), + ); + assert_eq!( + tokens, expected, + "num_draft={num_draft}, draft offset {}: temperature-0 speculative output \ + diverged from the target's greedy continuation", + draft.offset + ); + } + } + + assert_eq!( + stochastic_accept::acceptance_rule(&SamplingConfig::greedy(), true), + AcceptanceRule::GreedyArgmax, + "temperature 0 must stay on the argmax rule" + ); +} + +/// The mixed drafter really does produce both outcomes, so the byte-identity +/// test above is not silently exercising a single branch. Pins the arithmetic +/// of the agreement condition rather than trusting the comment. +#[test] +fn the_mixed_draft_model_agrees_on_exactly_one_token_in_six() { + let target = greedy_target(); + let mixed = mixed_draft(); + let rejecting = always_rejecting_draft(); + let agreements = (0..VOCAB as i32) + .filter(|&t| target.next(t) == mixed.next(t)) + .count(); + assert_eq!(agreements, 1, "mixed_draft must agree on exactly one token"); + for t in 0..VOCAB as i32 { + assert_ne!( + target.next(t), + rejecting.next(t), + "always_rejecting_draft must never agree, but it does at token {t}" + ); + } +} + +/// The main KV cache length after a run, in the two deterministic regimes. +/// +/// The invariant is that the main cache holds the prompt plus every emitted +/// token *except* the pending `current_token`, which the next round will +/// forward. The one exception is a run that stops inside the accept branch on +/// `max_tokens`: there the final emitted token was a draft token that had +/// already been forwarded as part of the verify input, so it is legitimately +/// still in the cache. Both regimes are pinned exactly here by choosing model +/// pairs whose acceptance is deterministic at temperature 0. +#[test] +fn kv_cache_length_is_exact_in_both_termination_regimes() { + const PROMPT: i32 = 2; + const N: usize = 24; + + // Regime 1: total disagreement. The draft walks a different orbit, so + // every round rejects at position 0, emits one replacement token and ends + // the round. The loop always exits through the `while` condition, so the + // last emitted token is always pending. + let target = greedy_target(); + let disagreeing_draft = always_rejecting_draft(); + for num_draft in [1usize, 3, 5] { + let mut generator = SpeculativeGenerator::new(1, 1); + let (tokens, _) = generator.generate( + &target, + &disagreeing_draft, + &[PROMPT], + N, + num_draft, + &SamplingConfig::greedy(), + ); + assert_eq!(tokens.len(), N); + let expected = 1 + tokens.len() - 1; + for (layer, cache) in generator.main_caches.iter().enumerate() { + assert_eq!( + cache.offset as usize, + expected, + "all-reject num_draft={num_draft} layer {layer}: main cache holds {} \ + entries, expected {expected} (prompt 1 + {} emitted - 1 pending). \ + A wrong trim on rejection lands here.", + cache.offset, + tokens.len() + ); + } + } + + // Regime 2: total agreement. The draft walks the target's own orbit, so + // every draft position is accepted and the run stops inside the accept + // branch on `max_tokens`, leaving the final token already forwarded. + let agreeing_draft = greedy_target(); + for num_draft in [2usize, 4] { + let mut generator = SpeculativeGenerator::new(1, 1); + let (tokens, _) = generator.generate( + &target, + &agreeing_draft, + &[PROMPT], + N, + num_draft, + &SamplingConfig::greedy(), + ); + assert_eq!(tokens.len(), N); + let expected = 1 + tokens.len(); + for (layer, cache) in generator.main_caches.iter().enumerate() { + assert_eq!( + cache.offset as usize, + expected, + "all-accept num_draft={num_draft} layer {layer}: main cache holds {} \ + entries, expected {expected} (prompt 1 + {} emitted, none pending). \ + A spurious trim on full acceptance lands here.", + cache.offset, + tokens.len() + ); + } + } +} + +/// Acceptance accounting must be exact, not approximate: it is the number the +/// performance A/B is read from, and a miscount would make the benchmark +/// report a win that did not happen. +/// +/// Both deterministic regimes are pinned. The never-agreeing drafter accepts +/// nothing, so `accepted == 0` and `mean_accepted_len == 0`. The drafter that +/// walks the target's own orbit accepts every proposal, so `accepted == +/// proposed` and the rate is exactly 1. +#[test] +fn acceptance_accounting_is_exact_in_both_deterministic_regimes() { + let target = greedy_target(); + const PROMPT: i32 = 2; + const N: usize = 24; + + for num_draft in [1usize, 3] { + let mut generator = SpeculativeGenerator::new(1, 1); + let _ = generator.generate( + &target, + &always_rejecting_draft(), + &[PROMPT], + N, + num_draft, + &SamplingConfig::greedy(), + ); + let stats = generator.acceptance_stats(); + assert_eq!(stats.rule, Some(AcceptanceRule::GreedyArgmax)); + assert_eq!(stats.accepted_draft_tokens, 0); + assert_eq!(stats.proposed_draft_tokens, stats.rounds * num_draft); + // Every round rejects at position 0, so exactly one position per round + // is ever tested no matter how long the block is. This is precisely the + // gap between `proposed` and `positions_tested`. + assert_eq!(stats.positions_tested, stats.rounds); + assert_eq!(stats.acceptance_rate(), 0.0); + assert_eq!(stats.per_position_acceptance(), 0.0); + assert_eq!(stats.mean_accepted_len(), 0.0); + + let mut generator = SpeculativeGenerator::new(1, 1); + let _ = generator.generate( + &target, + &greedy_target(), + &[PROMPT], + N, + num_draft, + &SamplingConfig::greedy(), + ); + let stats = generator.acceptance_stats(); + assert_eq!( + stats.accepted_draft_tokens, stats.proposed_draft_tokens, + "a drafter walking the target's own orbit must be accepted in full" + ); + // Nothing is rejected, so every proposed position is also tested and + // the two rates coincide. They can only coincide in this regime. + assert_eq!(stats.positions_tested, stats.proposed_draft_tokens); + assert_eq!(stats.acceptance_rate(), 1.0); + assert_eq!(stats.per_position_acceptance(), 1.0); + assert_eq!(stats.mean_accepted_len(), num_draft as f64); + } +} + +/// Expected accepted draft tokens per round when each position is accepted +/// independently with probability `a` and the chain stops at the first +/// rejection: `sum_{i=1..k} a^i`. +fn expected_accepted_len(a: f64, k: usize) -> f64 { + (1..=k).map(|i| a.powi(i as i32)).sum() +} + +/// The stochastic rule must be selected, recorded, and must actually raise the +/// mean accepted draft length. A run that silently fell back to argmax would +/// otherwise be indistinguishable from a real measurement, which is exactly +/// how an acceptance benchmark ends up comparing the fallback against itself. +/// +/// The per-position accept probability under modified rejection sampling is +/// `sum_x min(p(x), q(x))`; under the pre-#902 rule (an independent target +/// draw landing on the same token) it is `sum_x p(x) q(x)`. For this model +/// pair those are 0.705 and 0.344, which at a 4-token block predict mean +/// accepted lengths of about 1.80 and 0.52 tokens per round. The measurement +/// is required to land on the first and to clear twice the second. +#[test] +fn stochastic_runs_report_the_stochastic_rule_and_a_longer_accepted_run() { + const BLOCK: usize = 4; + let target = ContextFreeModel::new(&TARGET_LOGITS); + let draft = ContextFreeModel::new(&DRAFT_LOGITS); + + let p = softmax(&TARGET_LOGITS, TEMPERATURE); + let q = softmax(&DRAFT_LOGITS, TEMPERATURE); + let sum_min: f64 = p.iter().zip(&q).map(|(a, b)| a.min(*b)).sum(); + let sum_prod: f64 = p.iter().zip(&q).map(|(a, b)| a * b).sum(); + let predicted_new = expected_accepted_len(sum_min, BLOCK); + let predicted_old = expected_accepted_len(sum_prod, BLOCK); + + let mut rounds = 0usize; + let mut accepted = 0usize; + let mut positions = 0usize; + for _ in 0..24 { + let mut generator = SpeculativeGenerator::new(1, 1).with_stochastic_acceptance(true); + let _ = generator.generate(&target, &draft, &[0], 240, BLOCK, &stochastic_config()); + let stats = generator.acceptance_stats(); + assert_eq!( + stats.rule, + Some(AcceptanceRule::Stochastic), + "a temperature-{TEMPERATURE} run must take the stochastic rule; if this \ + reports an argmax variant the acceptance measurement is comparing the \ + fallback against itself" + ); + assert_eq!(stats.proposed_draft_tokens, stats.rounds * BLOCK); + assert!( + stats.positions_tested <= stats.proposed_draft_tokens, + "a chain cannot test more positions than it proposed" + ); + rounds += stats.rounds; + accepted += stats.accepted_draft_tokens; + positions += stats.positions_tested; + } + + // The primary assertion, stated in the quantity the theory predicts: the + // per-position acceptance probability must be `sum_x min(p(x), q(x))`. + // Binomial standard error over the pooled positions, 5 sigma. + let measured_per_position = accepted as f64 / positions as f64; + let se = (sum_min * (1.0 - sum_min) / positions as f64).sqrt(); + assert!( + (measured_per_position - sum_min).abs() <= 5.0 * se, + "per-position acceptance {measured_per_position:.4} over {positions} tested \ + positions is not the sum-min optimum {sum_min:.4} (+-{:.4})", + 5.0 * se + ); + assert!( + measured_per_position > sum_prod, + "per-position acceptance {measured_per_position:.4} must exceed the pre-#902 \ + sum-product rate {sum_prod:.4}" + ); + + // The derived per-round figure must then follow the geometric prediction. + // Looser band: `max_tokens` can truncate the final block of each call. + let measured_len = accepted as f64 / rounds as f64; + assert!( + (measured_len - predicted_new).abs() <= 0.12 * predicted_new, + "mean accepted draft length {measured_len:.4} over {rounds} rounds is not the \ + predicted {predicted_new:.4} for a per-position accept probability of {sum_min:.4}" + ); + assert!( + measured_len > 2.0 * predicted_old, + "mean accepted draft length {measured_len:.4} must clear twice the pre-#902 \ + prediction {predicted_old:.4}" + ); +} + +/// The closed-form diagnostic must agree with the measured acceptance, and must +/// respect `sum_x min(p, q) >= sum_x p(x) q(x)`. +/// +/// This is the check that turns "the acceptance rate looks wrong" into an +/// arithmetic statement. Run here on synthetic models where `p` and `q` are +/// known in closed form, so a disagreement localizes to the accept test rather +/// than to the models. +#[test] +fn closed_form_acceptance_matches_the_hand_computed_values() { + let target_logits = ffi::from_slice_f32(&TARGET_LOGITS, &[1, 1, VOCAB as i32]); + let draft_logits = ffi::from_slice_f32(&DRAFT_LOGITS, &[1, 1, VOCAB as i32]); + let config = stochastic_config(); + let p_arr = effective_token_distribution(&target_logits, &config, &[]); + let q_arr = effective_token_distribution(&draft_logits, &config, &[]); + + let (sum_min, sum_prod) = stochastic_accept::closed_form_acceptance(&p_arr, &q_arr); + + let p = softmax(&TARGET_LOGITS, TEMPERATURE); + let q = softmax(&DRAFT_LOGITS, TEMPERATURE); + let want_min: f64 = p.iter().zip(&q).map(|(a, b)| a.min(*b)).sum(); + let want_prod: f64 = p.iter().zip(&q).map(|(a, b)| a * b).sum(); + + assert!( + (sum_min - want_min).abs() < 1e-5, + "sum min {sum_min} != {want_min}" + ); + assert!( + (sum_prod - want_prod).abs() < 1e-5, + "sum prod {sum_prod} != {want_prod}" + ); + assert!( + sum_min >= sum_prod, + "min(a,b) >= a*b pointwise, so sum_min {sum_min} can never fall below \ + sum_prod {sum_prod}" + ); +} + +/// Documents a **pre-existing** draft-cache defect that predates issue #902 and +/// affects every acceptance rule identically. +/// +/// The draft loop forwards `num_draft` tokens (`current_token`, `d_0`, ..., +/// `d_{n-2}`), so the draft cache gains `n` entries per round. It should end +/// the round holding `current_token, d_0..d_{a-1}`, that is `a + 1` new +/// entries, so the correct trim is `n - a - 1 == rejected - 1`. The code trims +/// `min(rejected + 1, n)`, two entries too many, and the shortfall is never +/// recovered: for `num_draft = 4` the cache falls behind by 1, 2, 2, 2, 1 +/// tokens at `accepted` 0..4 respectively, every round, forever. +/// +/// At `accepted == 0` the draft cache does not advance at all. The drafter then +/// proposes from a frozen prefix while the target advances, its proposals +/// degrade, acceptance falls further, and the loop reinforces itself. This is +/// the most likely explanation for observed acceptance around 0.23 on an +/// 8B/1B pair against the 0.6-0.8 reported upstream. +/// +/// This test pins the current, wrong behavior so the defect is visible and so +/// whoever fixes it gets a failing test rather than silence. Fixing it is out +/// of scope for #902: it changes token streams and belongs in its own issue. +#[test] +fn draft_cache_falls_behind_the_main_cache_every_round_pre_existing_defect() { + let target = greedy_target(); + let draft = always_rejecting_draft(); + const PROMPT: i32 = 2; + const N: usize = 24; + + for num_draft in [1usize, 4] { + let mut generator = SpeculativeGenerator::new(1, 1); + let (tokens, _) = generator.generate( + &target, + &draft, + &[PROMPT], + N, + num_draft, + &SamplingConfig::greedy(), + ); + assert_eq!(tokens.len(), N); + + let main_offset = generator.main_caches[0].offset; + let draft_offset = generator.draft_caches[0].offset; + + // Every round rejects at position 0, so `accepted == 0` throughout and + // the draft cache never advances past the single prompt token. + assert_eq!( + draft_offset, 1, + "num_draft={num_draft}: draft cache should hold the prompt plus every \ + emitted token except the pending one ({}), but it is frozen at {} \ + because the trim is `min(rejected + 1, n)` instead of `rejected - 1`", + main_offset, draft_offset + ); + assert_eq!( + main_offset as usize, + 1 + tokens.len() - 1, + "num_draft={num_draft}: the main cache does track the sequence, which is \ + what makes the draft cache's drift a defect rather than a convention" + ); + assert!( + main_offset > draft_offset, + "the drafter is conditioning on {} tokens while the target conditions on {}", + draft_offset, + main_offset + ); + } +} + +/// The **default** rule (sampler-match, the rule that shipped before #902) must +/// also pass the distributional test, because it too is distribution-preserving. +/// This is the empirical half of the premise correction: if the default were +/// biased as the issue claims, this test would reject it with the same +/// statistic that rejects argmax-against-argmax at ~1825. +#[test] +fn the_default_sampler_match_rule_also_matches_the_target_distribution() { + let target = ContextFreeModel::new(&TARGET_LOGITS); + let draft = ContextFreeModel::new(&DRAFT_LOGITS); + let config = stochastic_config(); + let expected = softmax(&TARGET_LOGITS, TEMPERATURE); + let budget = sample_budget() / 2; + + let mut counts = vec![0u32; VOCAB]; + let mut collected = 0usize; + while collected < budget { + let mut generator = SpeculativeGenerator::new(1, 1).with_stochastic_acceptance(false); + let (tokens, _) = generator.generate(&target, &draft, &[0], 64, 4, &config); + assert_eq!( + generator.acceptance_stats().rule, + Some(AcceptanceRule::SamplerMatch), + "this test must exercise the default rule, not the opt-in one" + ); + for t in &tokens { + counts[*t as usize] += 1; + } + collected += tokens.len(); + } + + let chi2 = chi_square(&counts, &expected); + assert!( + chi2 < CHI2_DOF5_ALPHA_1E4, + "the pre-#902 sampler-match rule is distribution-preserving and must pass \ + the same statistic: chi-square {chi2:.3}, counts {counts:?}" + ); +} + +/// The CLI summary line is the only place the acceptance rule is observable on +/// `mlxcel generate`, because that binary installs no tracing subscriber. Its +/// shape is therefore an observable contract, not a debug convenience: pin the +/// stable `rule=` token and the presence of `mean_accepted_len`, so a rename +/// cannot silently break whoever is grepping a benchmark log to prove which arm +/// they measured. +#[test] +fn cli_summary_line_carries_the_rule_id_and_mean_accepted_len() { + let stats = SpeculativeAcceptanceStats { + rounds: 10, + proposed_draft_tokens: 40, + positions_tested: 24, + accepted_draft_tokens: 18, + rule: Some(AcceptanceRule::Stochastic), + ..SpeculativeAcceptanceStats::default() + }; + let line = stats.summary_line(); + assert!(line.starts_with("[Speculative acceptance] "), "{line}"); + assert!(line.contains("rule=stochastic "), "{line}"); + assert!(line.contains("mean_accepted_len=1.8000"), "{line}"); + assert!(line.contains("acceptance_rate=0.4500"), "{line}"); + assert!(line.contains("positions_tested=24"), "{line}"); + // 18/24, the quantity comparable to `sum min(p, q)`. Distinct from + // `acceptance_rate` (18/40), which divides by block positions the chain + // never reached. Confusing the two is what made a correct implementation + // look like a 2x regression. + assert!(line.contains("per_position_acceptance=0.7500"), "{line}"); + assert!(line.contains("rounds=10"), "{line}"); + assert!(line.contains("acceptance-optimal"), "{line}"); + + // Every rule id must be distinct and free of whitespace, or `rule=` cannot + // be grepped unambiguously. + let ids: Vec<&str> = [ + AcceptanceRule::GreedyArgmax, + AcceptanceRule::SamplerMatch, + AcceptanceRule::Stochastic, + AcceptanceRule::SamplerMatchNoProposalDistribution, + ] + .iter() + .map(|r| r.id()) + .collect(); + for id in &ids { + assert!( + !id.contains(char::is_whitespace), + "rule id {id:?} has spaces" + ); + } + let mut sorted = ids.clone(); + sorted.sort_unstable(); + sorted.dedup(); + assert_eq!( + sorted.len(), + ids.len(), + "rule ids must be distinct: {ids:?}" + ); + + // `acceptance_rate` and `per_position_acceptance` must never be confused + // for each other: they differ whenever any chain stopped early. + assert!( + (stats.per_position_acceptance() - 0.75).abs() < 1e-9, + "{}", + stats.per_position_acceptance() + ); + assert!( + (stats.acceptance_rate() - 0.45).abs() < 1e-9, + "{}", + stats.acceptance_rate() + ); + + // The default arm must be nameable too: it is one half of the A/B. + let baseline = SpeculativeAcceptanceStats { + rule: Some(AcceptanceRule::SamplerMatch), + ..stats + }; + assert!( + baseline.summary_line().contains("rule=sampler-match "), + "{}", + baseline.summary_line() + ); + + // A run that never entered the decode loop must say so rather than + // implying a rule it did not run. + let empty = SpeculativeAcceptanceStats::default(); + assert!(empty.summary_line().contains("rule=unknown ")); +} + +/// The stochastic acceptance path must leave the caches in one of the same two +/// regimes: mixed accept/reject rounds change *which* regime the run ends in, +/// never the arithmetic. Every layer must also agree with every other layer; a +/// per-layer trim mismatch corrupts attention silently and produces no other +/// visible symptom. +#[test] +fn kv_cache_stays_consistent_across_mixed_accept_reject_rounds() { + let target = ContextFreeModel::new(&TARGET_LOGITS); + let draft = ContextFreeModel::new(&DRAFT_LOGITS); + let config = stochastic_config(); + + for num_draft in [1usize, 3, 5] { + for trial in 0..8 { + let mut generator = SpeculativeGenerator::new(target.num_layers(), draft.num_layers()) + .with_stochastic_acceptance(true); + let (tokens, _) = generator.generate(&target, &draft, &[0], 40, num_draft, &config); + + let pending = 1 + tokens.len() - 1; + let forwarded = 1 + tokens.len(); + let offset = generator.main_caches[0].offset as usize; + assert!( + offset == pending || offset == forwarded, + "num_draft={num_draft} trial={trial}: main cache offset {offset} is \ + neither {pending} (last token pending) nor {forwarded} (run stopped \ + inside the accept branch); the rewind on rejection is wrong" + ); + for (layer, cache) in generator.main_caches.iter().enumerate() { + assert_eq!( + cache.offset as usize, offset, + "num_draft={num_draft} trial={trial}: main cache layer {layer} \ + disagrees with layer 0" + ); + } + let draft_offset = generator.draft_caches[0].offset; + for (layer, cache) in generator.draft_caches.iter().enumerate() { + assert_eq!( + cache.offset, draft_offset, + "num_draft={num_draft} trial={trial}: draft cache layer {layer} \ + disagrees with layer 0" + ); + } + } + } +} diff --git a/src/lib/mlxcel-core/src/speculative/mod.rs b/src/lib/mlxcel-core/src/speculative/mod.rs index 9cd58d3b8..52af43403 100644 --- a/src/lib/mlxcel-core/src/speculative/mod.rs +++ b/src/lib/mlxcel-core/src/speculative/mod.rs @@ -28,6 +28,11 @@ //! c. Accept matching prefix, rewind caches for rejected tokens //! d. Continue from the divergence point //! +//! Step 3c is the acceptance rule. At `temperature == 0` it is the argmax +//! comparison described above. At `temperature > 0` it is modified rejection +//! sampling, which preserves the target model's distribution exactly; see +//! [`stochastic_accept`]. +//! //! ## Sibling modules //! //! - [`mtp`] — Multi-Token Prediction (MTP) round-loop generator for the @@ -36,22 +41,30 @@ //! (drafter has no own KV cache; verify is a single forward over the //! whole draft block; rollback uses per-row tail-zero rather than //! `trim_caches`). See [`mtp::MtpGenerator`]. +//! - [`stochastic_accept`] — the distribution-preserving acceptance rule and +//! its residual resample, shared by every speculative verify path. pub mod mtp; +pub mod stochastic_accept; use crate::cache::can_trim_prompt_cache; use crate::ffi; +use crate::ffi::MlxArray; use crate::ffi::MlxThreadLocalStream; use crate::generate::{GenerationStats, LanguageModel, SamplingConfig}; use crate::generation_policy::{initial_token_history, merged_eos_token_ids}; use crate::hardware; use crate::layers::KVCache; -use crate::sampling::{TokenBiasMap, sample_token_optimized}; +use crate::sampling::{ + TokenBiasMap, effective_token_distribution, sample_token_optimized, + sample_token_with_distribution, +}; use crate::streams::{install_thread_local_default_stream, new_thread_local_generation_stream}; use crate::utils::{align_to_na_tile, create_padded_prefill_mask}; use cxx::UniquePtr; use std::borrow::Cow; use std::time::Instant; +use stochastic_accept::{AcceptanceRule, DraftVerdict}; /// Default chunk size for chunked prefill in speculative decoding. /// @@ -69,6 +82,122 @@ fn should_align_verification() -> bool { hw.has_neural_accelerator && hw.macos_supports_na } +/// Acceptance accounting for one [`SpeculativeGenerator::generate`] call. +/// +/// The two ratios are the quantities the acceptance change is supposed to +/// move: [`Self::acceptance_rate`] should rise from about `sum_x p(x) q(x)` +/// under the old rule to about `sum_x min(p(x), q(x))` under modified +/// rejection sampling, and [`Self::mean_accepted_len`] is the per-round +/// version that drives end-to-end throughput. +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub struct SpeculativeAcceptanceStats { + /// Verify rounds executed (one batched target forward each). + pub rounds: usize, + /// Draft tokens proposed across all rounds. + pub proposed_draft_tokens: usize, + /// Draft tokens the verify rule accepted. + pub accepted_draft_tokens: usize, + /// The acceptance rule the run actually used. + pub rule: Option, + /// Verify positions the accept test actually ran on. Differs from + /// `proposed_draft_tokens` because a chain stops at its first rejection, so + /// positions after it are proposed but never tested. + pub positions_tested: usize, + /// Running sum of `sum_x min(p(x), q(x))` over tested positions, the + /// closed-form probability modified rejection sampling accepts. Zero unless + /// [`stochastic_accept::accept_diagnostics_enabled`]. + pub closed_form_sum_min: f64, + /// Running sum of `sum_x p(x) q(x)` over tested positions, the closed-form + /// probability the pre-#902 rule accepts. Zero unless diagnostics are on. + pub closed_form_sum_prod: f64, +} + +impl SpeculativeAcceptanceStats { + /// Accepted draft tokens per *proposed* draft token. + /// + /// Not the per-position acceptance probability, and not comparable to + /// `sum_x min(p(x), q(x))`: a chain stops at its first rejection, so the + /// block positions behind it are counted as proposed but never tested. Use + /// [`Self::per_position_acceptance`] to compare against theory or across an + /// A/B; this one is the "how much of each drafted block survived" figure, + /// which also moves with the block length. + pub fn acceptance_rate(&self) -> f64 { + if self.proposed_draft_tokens == 0 { + 0.0 + } else { + self.accepted_draft_tokens as f64 / self.proposed_draft_tokens as f64 + } + } + + /// Accepted draft tokens per verify position the accept test actually ran + /// on. + /// + /// This is the quantity the theory predicts: modified rejection sampling + /// converges to `sum_x min(p(x), q(x))` and the pre-#902 rule to + /// `sum_x p(x) q(x)`, both per position. Counted identically on both + /// acceptance rules. + pub fn per_position_acceptance(&self) -> f64 { + if self.positions_tested == 0 { + 0.0 + } else { + self.accepted_draft_tokens as f64 / self.positions_tested as f64 + } + } + + /// Mean accepted draft tokens per verify round. This is the figure the + /// issue's performance gate names ("mean accepted draft length"). + pub fn mean_accepted_len(&self) -> f64 { + if self.rounds == 0 { + 0.0 + } else { + self.accepted_draft_tokens as f64 / self.rounds as f64 + } + } + + /// One-line summary for the CLI, printed on **stdout** after a speculative + /// run. + /// + /// This exists because the `mlxcel` CLI installs no `tracing` subscriber, + /// so the info-level instrumentation in this module is unreachable on the + /// only binary that runs [`SpeculativeGenerator`]. An acceptance-rate A/B + /// that cannot name the rule it measured is not a measurement, so the two + /// facts a reviewer needs (which rule ran, and the mean accepted draft + /// length) are printed unconditionally rather than behind a log level. + /// + /// The `rule=` token is [`AcceptanceRule::id`], which is stable and + /// greppable. `rule=unknown` means [`SpeculativeGenerator::generate`] never + /// reached its decode loop. + /// + /// Used by: `commands::generate` (offline `mlxcel generate --draft-model`) + pub fn summary_line(&self) -> String { + let rule = self.rule.map_or("unknown", |r| r.id()); + let label = self.rule.map_or("no speculative round ran", |r| r.label()); + let mut line = format!( + "[Speculative acceptance] rule={rule} rounds={} proposed={} positions_tested={} \ + accepted={} per_position_acceptance={:.4} acceptance_rate={:.4} \ + mean_accepted_len={:.4} ({label})", + self.rounds, + self.proposed_draft_tokens, + self.positions_tested, + self.accepted_draft_tokens, + self.per_position_acceptance(), + self.acceptance_rate(), + self.mean_accepted_len(), + ); + if self.positions_tested > 0 && self.closed_form_sum_min > 0.0 { + let n = self.positions_tested as f64; + line.push_str(&format!( + "\n[Speculative acceptance diagnostic] closed_form_sum_min={:.4} \ + closed_form_sum_prod={:.4} (measured per-position acceptance must sit at \ + sum_min, and sum_min >= sum_prod always)", + self.closed_form_sum_min / n, + self.closed_form_sum_prod / n, + )); + } + line + } +} + /// Speculative decoding generator /// /// Uses a draft model to propose candidate tokens and a main model to verify them. @@ -84,6 +213,16 @@ pub struct SpeculativeGenerator { /// synchronization stay paired even when the generator is moved /// across threads after construction. generation_stream: Option>, + /// Acceptance accounting for the most recent [`Self::generate`] call. + /// + /// Reset at the top of every call by [`Self::reset`]. Read through + /// [`Self::acceptance_stats`] and summarized once per call at info level, + /// which is what makes an acceptance-rate A/B measurable from a log rather + /// than only from a stopwatch. + acceptance: SpeculativeAcceptanceStats, + /// Per-generator override of the `MLXCEL_SPECULATIVE_STOCHASTIC_ACCEPT` + /// default. `None` follows the env; `Some(v)` forces it. + stochastic_acceptance: Option, /// Cached per-generator `TokenBiasMap` resolved from a `LangBiasConfig`. /// /// **Axis B invariant**: the bias is applied **only** to the target @@ -104,10 +243,33 @@ impl SpeculativeGenerator { draft_caches: (0..draft_num_layers).map(|_| KVCache::new()).collect(), generated_tokens: Vec::new(), generation_stream: new_thread_local_generation_stream(), + acceptance: SpeculativeAcceptanceStats::default(), + stochastic_acceptance: None, token_bias: TokenBiasMap::default(), } } + /// Acceptance accounting for the most recent [`Self::generate`] call. + pub fn acceptance_stats(&self) -> SpeculativeAcceptanceStats { + self.acceptance + } + + /// Force the acceptance-optimal (modified rejection sampling) rule on or + /// off for this generator, overriding the + /// `MLXCEL_SPECULATIVE_STOCHASTIC_ACCEPT` process default. + /// + /// The feature is opt-in because the rule it replaces was already + /// distribution-preserving on this path (see + /// [`stochastic_accept`]), so it trades measurable per-position cost for an + /// acceptance gain that is near zero against a confident drafter. This + /// setter is how a caller that has established the gain is worth taking, + /// or a test that must exercise the rule regardless of the process + /// default, turns it on. + pub fn with_stochastic_acceptance(mut self, enabled: bool) -> Self { + self.stochastic_acceptance = Some(enabled); + self + } + /// Attach a pre-resolved `TokenBiasMap` to this speculative generator. /// /// The bias is cached for the generator's lifetime and applied **only** to @@ -164,6 +326,7 @@ impl SpeculativeGenerator { *cache = KVCache::new(); } self.generated_tokens.clear(); + self.acceptance = SpeculativeAcceptanceStats::default(); } /// Get the generated tokens @@ -333,6 +496,30 @@ impl SpeculativeGenerator { } // DECODE PHASE. + // + // Acceptance policy for this call (issue #902). This generator owns + // the draft sampler, so it can always report the distribution each + // proposal was drawn from; `has_proposal_distribution` is + // unconditionally true here. The remaining gates are the target + // sampler being stochastic and the env kill switch. The chosen rule is + // logged once per process per kind at info level, so a server log + // proves which rule a run used. + let rule = stochastic_accept::acceptance_rule_with_override( + target_sampling, + true, + self.stochastic_acceptance, + ); + stochastic_accept::note_rule(rule); + let stochastic = rule == AcceptanceRule::Stochastic; + self.acceptance.rule = Some(rule); + + // The closed-form diagnostic needs `q` on *both* arms, otherwise the + // two halves of an A/B report different quantities and cannot be + // compared. Off by default, so the argmax arm captures nothing and + // stays exactly as it was. + let diagnostics = stochastic_accept::accept_diagnostics_enabled(); + let capture_proposal_probs = stochastic || diagnostics; + let decode_start = Instant::now(); let mut current_token = first_token; let mut done = false; @@ -340,6 +527,10 @@ impl SpeculativeGenerator { while self.generated_tokens.len() < max_tokens && !done { // Step 1: Generate draft tokens let mut draft_tokens = Vec::with_capacity(num_draft); + // `q` per drafted position, in proposal order. Populated only on + // the stochastic path; the argmax path allocates nothing and runs + // exactly the pre-#902 code. + let mut draft_probs: Vec> = Vec::new(); let mut draft_token = current_token; for _ in 0..num_draft { @@ -347,8 +538,22 @@ impl SpeculativeGenerator { let draft_logits = draft_model.forward(&draft_input, &mut self.draft_caches, None); // Axis B: draft sampler MUST NOT see the bias. See // `draft_sampling` for the rationale. - let (tok_arr, _) = - sample_token_optimized(&draft_logits, draft_sampling, &token_history); + let tok_arr = if capture_proposal_probs { + // The token and `q` come out of one pre-step, which is what + // guarantees `q` is the distribution this very token was + // drawn from rather than a later reconstruction. `q` is + // left unevaluated: a chain that gets rejected at an + // earlier position never materializes the tail. + let (tok_arr, probs) = sample_token_with_distribution( + &draft_logits, + draft_sampling, + &token_history, + ); + draft_probs.push(probs); + tok_arr + } else { + sample_token_optimized(&draft_logits, draft_sampling, &token_history).0 + }; ffi::eval(&tok_arr); draft_token = ffi::item_i32(&tok_arr); draft_tokens.push(draft_token); @@ -454,13 +659,51 @@ impl SpeculativeGenerator { // Reshape to [1, 1, vocab] for sample_token_optimized let pos_logits = ffi::reshape(&pos_logits, &[1, 1, main_shape[2]]); - // Axis B: target verification uses the bias-augmented sampling. - let (main_tok_arr, _) = - sample_token_optimized(&pos_logits, target_sampling, &token_history); - ffi::eval(&main_tok_arr); - let main_token = ffi::item_i32(&main_tok_arr); + // Axis B: target verification uses the bias-augmented sampling + // on both arms. + // + // `accept_token` is the token kept when the draft is accepted; + // `replacement` is the token emitted when it is not. Both arms + // produce the same pair so the bookkeeping below is shared. + // Counted on both arms: this is the denominator the per-position + // acceptance probability is measured against, and the two arms + // must report the same quantity for an A/B to mean anything. + // `proposed_draft_tokens` is *not* that denominator, because a + // chain stops at its first rejection and never tests the block + // positions behind it. + self.acceptance.positions_tested += 1; + if diagnostics { + let target_probs = + effective_token_distribution(&pos_logits, target_sampling, &token_history); + let (sum_min, sum_prod) = + stochastic_accept::closed_form_acceptance(&target_probs, &draft_probs[i]); + self.acceptance.closed_form_sum_min += sum_min; + self.acceptance.closed_form_sum_prod += sum_prod; + } - if main_token == draft_token { + let (accept_draft, replacement) = if stochastic { + // `p` at this position, conditioned on the tokens accepted + // so far in this round (`token_history` grows as we + // accept), which is the correct target conditional. + let target_probs = + effective_token_distribution(&pos_logits, target_sampling, &token_history); + match stochastic_accept::verify_draft_token( + &target_probs, + &draft_probs[i], + draft_token, + ) { + DraftVerdict::Accept => (true, draft_token), + DraftVerdict::Reject { replacement } => (false, replacement), + } + } else { + let (main_tok_arr, _) = + sample_token_optimized(&pos_logits, target_sampling, &token_history); + ffi::eval(&main_tok_arr); + let main_token = ffi::item_i32(&main_tok_arr); + (main_token == draft_token, main_token) + }; + + if accept_draft { // Accept draft token accepted += 1; @@ -479,13 +722,17 @@ impl SpeculativeGenerator { break; } } else { - // Reject: use main model's token instead - if eos_tokens.contains(&main_token) { + // Reject: emit the target-side replacement instead. On the + // argmax rule that is the target sampler's own draw; on + // the stochastic rule it is a draw from the normalized + // residual `relu(p - q)`, which is what makes the emitted + // stream distributed as `p`. + if eos_tokens.contains(&replacement) { done = true; } else { - self.generated_tokens.push(main_token); + self.generated_tokens.push(replacement); if needs_history { - token_history.push(main_token); + token_history.push(replacement); } } break; @@ -522,6 +769,10 @@ impl SpeculativeGenerator { current_token = *self.generated_tokens.last().unwrap(); } + self.acceptance.rounds += 1; + self.acceptance.proposed_draft_tokens += draft_tokens.len(); + self.acceptance.accepted_draft_tokens += accepted; + // Step 4: Rewind caches for rejected tokens let rejected = draft_tokens.len() - accepted; if rejected > 0 { @@ -555,6 +806,23 @@ impl SpeculativeGenerator { let decode_time = decode_start.elapsed(); + // One summary line per call at info level. Together with the one-shot + // rule line from `note_rule`, this is what lets an acceptance-rate A/B + // be read off a log instead of inferred from wall-clock alone, and it + // names the rule so a run can never be attributed to the wrong arm. + tracing::info!( + rule = rule.id(), + rounds = self.acceptance.rounds, + proposed_draft_tokens = self.acceptance.proposed_draft_tokens, + positions_tested = self.acceptance.positions_tested, + accepted_draft_tokens = self.acceptance.accepted_draft_tokens, + per_position_acceptance = self.acceptance.per_position_acceptance(), + acceptance_rate = self.acceptance.acceptance_rate(), + mean_accepted_len = self.acceptance.mean_accepted_len(), + generated_tokens = self.generated_tokens.len(), + "speculative decode finished" + ); + let stats = Self::build_stats( prompt_tokens.len(), self.generated_tokens.len(), @@ -606,6 +874,10 @@ fn trim_caches(caches: &mut [KVCache], n: i32) -> i32 { trimmed } +#[cfg(test)] +#[path = "distribution_tests.rs"] +mod distribution_tests; + #[cfg(test)] mod tests { use super::*; diff --git a/src/lib/mlxcel-core/src/speculative/stochastic_accept.rs b/src/lib/mlxcel-core/src/speculative/stochastic_accept.rs new file mode 100644 index 000000000..23a0bc7fe --- /dev/null +++ b/src/lib/mlxcel-core/src/speculative/stochastic_accept.rs @@ -0,0 +1,601 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Acceptance-optimal speculative acceptance (chain speculative sampling). +//! +//! # What the rule this replaces actually did +//! +//! Issue #902 opens by stating that the verify loop "accepts a draft token iff +//! it equals the target model's argmax". **For this generator that is not +//! true**, and the correction matters enough to lead with. +//! +//! [`crate::speculative::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 when `temperature == 0`. +//! It then emits that draw on **both** branches: on acceptance the drafted +//! token, which by the accept test equals the draw, and on rejection the draw +//! itself. Every emitted token is therefore a fresh conditional draw from the +//! target distribution, and the stream is already exactly a target-only sample. +//! Call this rule *sampler-match* ([`AcceptanceRule::SamplerMatch`]). +//! +//! So sampler-match was never biased. What it loses is acceptance rate: two +//! independent draws coincide with probability `sum_x p(x) q(x)`, which is +//! strictly below the best any correct rule can do. +//! +//! The biased argmax-against-argmax rule the issue describes is real, but it +//! lives elsewhere: `Gemma4MtpTargetAdapter::verify_forward` uses +//! `argmax_per_position` and the DFlash round loop uses `argmax_logits_to_array`, +//! both regardless of `sampler.temperature`. Neither is reached from here. +//! +//! # The acceptance ceiling, and why "beat argmax" is not achievable +//! +//! For a `q`-distributed proposal and a `p`-distributed emission, the largest +//! probability the two can be made to coincide is `sum_x min(p(x), q(x))`: the +//! maximal-coupling bound, equivalently `1 - TV(p, q)`. **Every** +//! distribution-preserving acceptance rule is capped there. Sampler-match +//! reaches `sum_x p(x) q(x)`; modified rejection sampling attains the bound +//! exactly and is therefore optimal among correct rules. +//! +//! Argmax-against-argmax is not bound by this, because it is not correct. It +//! accepts at `q(argmax p)`, which for a confident drafter that agrees with the +//! target's mode can be far above `sum_x min(p, q)`: with `p(a) = 0.3` and +//! `q(a) = 0.95` it accepts at 0.95 where the ceiling for any correct rule is +//! about 0.30. That extra acceptance is paid for in bias, one token at a time. +//! +//! The consequence is worth stating plainly, because it reframes the issue +//! rather than failing it: **no implementation can make a +//! distribution-preserving rule "strictly improve" mean accepted length against +//! argmax acceptance.** Asking for that is asking a correct rule to beat an +//! incorrect one at the metric the incorrect one optimizes. Against +//! sampler-match, the rule this generator actually ran, the improvement is +//! real and provable, since `sum_x min(p, q) >= sum_x p(x) q(x)` always. +//! +//! # The rule implemented here +//! +//! Modified rejection sampling, from Leviathan et al. 2023 ("Fast Inference +//! from Transformers via Speculative Decoding", Algorithm 1) and Chen et al. +//! 2023 ("Accelerating Large Language Model Decoding with Speculative +//! Sampling"). Write `p` for the target's effective distribution at the +//! position and `q` for the distribution the drafter actually proposed from: +//! +//! 1. Draw `u ~ U[0, 1)`. **Accept `t` iff `u * q(t) <= p(t)`**, i.e. accept +//! with probability `min(1, p(t) / q(t))`. +//! 2. On the first rejection, emit one token drawn from the **residual** +//! `normalize(relu(p - q))` and end the chain there. +//! 3. If every drafted token was accepted, emit a bonus token drawn from `p` +//! at the final verify position (unchanged from before). +//! +//! ## Why this preserves the target distribution +//! +//! Fix a position and let `A` be the accept event. For any token `x`: +//! +//! ```text +//! P(emit x) = P(propose x, accept x) + P(reject) * P(residual draws x) +//! = q(x) * min(1, p(x)/q(x)) + beta * relu(p(x) - q(x)) / beta +//! = min(q(x), p(x)) + relu(p(x) - q(x)) +//! = p(x) +//! ``` +//! +//! where `beta = sum_y relu(p(y) - q(y)) = 1 - sum_y min(p(y), q(y)) = P(reject)`, +//! so the normalizing constant of the residual cancels the rejection +//! probability exactly. The final line is the identity +//! `min(a, b) + max(a - b, 0) = a`. The argument holds for **any** proposal +//! `q` whose support is not required to relate to `p` in any way, which is +//! what makes the rest of the pipeline safe: a drafter that proposes greedily +//! (`q` one-hot), a drafter whose penalties were computed against a stale +//! token history, and a drafter that is a different model family all remain +//! lossless as long as `q` is *the distribution the token was actually drawn +//! from*. That is the one invariant every caller must maintain, and it is why +//! [`crate::sampling::sample_token_with_distribution`] returns the token and +//! its distribution from a single pre-step rather than recomputing `q` later. +//! +//! Chaining over a draft block is the same argument applied inductively: the +//! chain only reaches position `i` when positions `0..i` were accepted, in +//! which case the target's context at `i` is exactly the context the drafter +//! conditioned on, so `p_i` and `q_i` are distributions over the same +//! conditional. +//! +//! ## Temperature 0 is untouched +//! +//! [`acceptance_rule`] returns [`AcceptanceRule::GreedyArgmax`] whenever the +//! target sampler is greedy (`temperature == 0.0` or `top_k == 1`), and the +//! caller keeps running the pre-#902 comparison verbatim: no distribution +//! tensors are built, no randomness is drawn, and the stream is byte-identical. +//! +//! # Opt-in, not default +//! +//! `MLXCEL_SPECULATIVE_STOCHASTIC_ACCEPT=1` enables this rule; it is **off by +//! default**. Since sampler-match was already distribution-preserving on this +//! path, the change buys acceptance rate and nothing else, and the available +//! gain is the ratio `sum_x min(p, q) / sum_x p(x) q(x)`. That ratio collapses +//! toward 1 whenever the drafter is confident, because `q(t*) ~ 1` makes +//! `min(p, q)` and `p * q` coincide. Measured on a Llama-3.1-8B / +//! Llama-3.2-1B pair at temperature 0.7 it is about 1.02, which does not +//! justify two extra full-vocabulary passes and an extra host sync per +//! verified position. +//! +//! Enable it where the gain is real: a high-entropy drafter, a higher +//! temperature, or a verify path whose current rule is genuinely biased. +//! `MLXCEL_SPECULATIVE_ACCEPT_DIAG=1` reports both closed forms so the gain can +//! be checked before a throughput run rather than inferred from one. +//! +//! # Observability +//! +//! Every distinct outcome kind is logged once per process at **info** level +//! (see [`note_rule`] and [`note_outcome`]), and +//! [`crate::speculative::SpeculativeAcceptanceStats::summary_line`] prints the +//! rule and the rates on stdout after every run, because the CLI installs no +//! tracing subscriber. The one-shot latch is per kind, not global, so a +//! later-appearing kind is not swallowed by an earlier one. + +use crate::dtype; +use crate::ffi; +use crate::ffi::MlxArray; +use crate::generate::SamplingConfig; +use cxx::UniquePtr; +use std::sync::OnceLock; +use std::sync::atomic::{AtomicBool, Ordering}; + +/// Environment switch enabling modified rejection sampling. Opt-in. +pub const STOCHASTIC_ACCEPT_ENV: &str = "MLXCEL_SPECULATIVE_STOCHASTIC_ACCEPT"; + +/// Which acceptance rule a speculative verify round runs. +/// +/// **Every variant here is distribution-preserving.** That is not an accident +/// of naming, it is the central finding of issue #902: the rule this generator +/// shipped with compares the drafted token against a fresh draw from the target +/// *sampler* and emits that draw on both branches, so the emitted stream was +/// already a target-only sample. The biased argmax-against-argmax rule the +/// issue describes lives in the Gemma 4 MTP and DFlash verify paths, is not +/// selected here, and is deliberately not represented in this enum: an +/// unreachable variant would be a claim this module cannot back up. +/// +/// What the variants differ in is [`Self::is_acceptance_optimal`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +#[non_exhaustive] +pub enum AcceptanceRule { + /// The target sampler is greedy, so its draw is the argmax and the rule is + /// "the drafted token is the argmax". Kept on the original integer + /// comparison: identical outcome, none of the tensor work. + GreedyArgmax, + /// The pre-#902 default at `temperature > 0`, and still the default: + /// accept iff the drafted token equals an independent draw from the target + /// sampler, and emit that draw either way. Lossless, but its acceptance + /// probability is `sum_x p(x) q(x)`, below the achievable optimum. + SamplerMatch, + /// Modified rejection sampling with residual resample, enabled by + /// [`STOCHASTIC_ACCEPT_ENV`]. Lossless *and* acceptance-optimal. + Stochastic, + /// [`STOCHASTIC_ACCEPT_ENV`] asked for the stochastic rule but the drafter + /// could not report the distribution it proposed from, so there is no `q` + /// for the accept test and the round stays on [`Self::SamplerMatch`]. + /// Fabricating a `q` would void the losslessness proof outright. + SamplerMatchNoProposalDistribution, +} + +impl AcceptanceRule { + /// True when this rule preserves the target model's token distribution. + /// + /// True for every variant. See the type-level docs: the biased rule is in + /// MTP/DFlash, not here. Kept as an explicit, tested statement rather than + /// deleted, because "which of these is safe to serve" is the first question + /// a reader of this enum asks. + #[inline] + pub fn is_distribution_preserving(self) -> bool { + true + } + + /// True when this rule attains the maximum acceptance probability any + /// distribution-preserving rule can reach. + /// + /// That maximum is `sum_x min(p(x), q(x))`, the maximal-coupling bound: it + /// is the largest probability with which a `q`-distributed proposal and a + /// `p`-distributed emission can be made to coincide. Modified rejection + /// sampling attains it; [`Self::SamplerMatch`] reaches only + /// `sum_x p(x) q(x)`. + /// + /// A rule that is *not* distribution-preserving is not bound by this and + /// can accept more. Argmax-against-argmax accepts at `q(argmax p)`, which + /// for a confident drafter agreeing with the target's mode far exceeds + /// `sum_x min(p, q)`. That is bought by emitting a biased stream, so it is + /// not a fair comparison point for any correct rule. + #[inline] + pub fn is_acceptance_optimal(self) -> bool { + matches!(self, Self::Stochastic | Self::GreedyArgmax) + } + + /// Stable, greppable identifier for this rule. + /// + /// Printed by the CLI summary line and matched by tooling, so it is part of + /// the observable contract: renaming a token here breaks whoever is + /// grepping a benchmark log to prove which arm they measured. + pub fn id(self) -> &'static str { + match self { + Self::GreedyArgmax => "greedy-argmax", + Self::SamplerMatch => "sampler-match", + Self::Stochastic => "stochastic", + Self::SamplerMatchNoProposalDistribution => "sampler-match-no-proposal-distribution", + } + } + + /// Human-readable description shown next to [`Self::id`]. + pub fn label(self) -> &'static str { + match self { + Self::GreedyArgmax => "greedy target sampler: accept iff the draft is the argmax", + Self::SamplerMatch => { + "sampler-match: accept iff the draft equals a fresh target draw (default; \ + set MLXCEL_SPECULATIVE_STOCHASTIC_ACCEPT=1 for the acceptance-optimal rule)" + } + Self::Stochastic => "stochastic: modified rejection sampling, acceptance-optimal", + Self::SamplerMatchNoProposalDistribution => { + "sampler-match: stochastic acceptance requested but the drafter reported no \ + proposal distribution" + } + } + } + + fn latch(self) -> &'static AtomicBool { + static GREEDY: AtomicBool = AtomicBool::new(false); + static SAMPLER_MATCH: AtomicBool = AtomicBool::new(false); + static STOCHASTIC: AtomicBool = AtomicBool::new(false); + static NO_Q: AtomicBool = AtomicBool::new(false); + match self { + Self::GreedyArgmax => &GREEDY, + Self::SamplerMatch => &SAMPLER_MATCH, + Self::Stochastic => &STOCHASTIC, + Self::SamplerMatchNoProposalDistribution => &NO_Q, + } + } +} + +/// What happened to one drafted token under [`AcceptanceRule::Stochastic`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +#[non_exhaustive] +pub enum AcceptanceOutcome { + /// `u * q(t) <= p(t)` held: the drafted token was kept. + Accepted, + /// The drafted token was rejected and replaced by a draw from the + /// normalized residual `relu(p - q)`. + ResidualResample, + /// The residual carried no mass at all, so the replacement was drawn from + /// the target distribution `p` instead. Unreachable in exact arithmetic + /// (a rejection implies positive residual mass); reachable only when every + /// positive entry of `p - q` underflows float32. + ResidualDegenerateFallback, +} + +impl AcceptanceOutcome { + fn label(self) -> &'static str { + match self { + Self::Accepted => "accept", + Self::ResidualResample => "reject, resampled from normalized relu(p - q)", + Self::ResidualDegenerateFallback => { + "reject, residual mass underflowed to zero, resampled from target p" + } + } + } + + fn latch(self) -> &'static AtomicBool { + static ACCEPTED: AtomicBool = AtomicBool::new(false); + static RESIDUAL: AtomicBool = AtomicBool::new(false); + static DEGENERATE: AtomicBool = AtomicBool::new(false); + match self { + Self::Accepted => &ACCEPTED, + Self::ResidualResample => &RESIDUAL, + Self::ResidualDegenerateFallback => &DEGENERATE, + } + } +} + +/// Log the first occurrence of each distinct acceptance rule at info level. +/// +/// One latch per variant, not one global flag: a run that starts greedy and +/// later serves a `temperature > 0` request logs both lines, so the log is a +/// complete record of which rules the process ever ran. +pub fn note_rule(rule: AcceptanceRule) { + if !rule.latch().swap(true, Ordering::Relaxed) { + tracing::info!( + rule = rule.label(), + distribution_preserving = rule.is_distribution_preserving(), + "speculative acceptance rule active" + ); + } +} + +/// Log the first occurrence of each distinct per-token outcome at info level. +pub fn note_outcome(outcome: AcceptanceOutcome) { + if !outcome.latch().swap(true, Ordering::Relaxed) { + tracing::info!( + outcome = outcome.label(), + "speculative stochastic acceptance outcome first seen" + ); + } +} + +/// Reset every one-shot log latch. Test-only. +#[cfg(test)] +pub(crate) fn reset_log_latches() { + for rule in [ + AcceptanceRule::GreedyArgmax, + AcceptanceRule::SamplerMatch, + AcceptanceRule::Stochastic, + AcceptanceRule::SamplerMatchNoProposalDistribution, + ] { + rule.latch().store(false, Ordering::Relaxed); + } + for outcome in [ + AcceptanceOutcome::Accepted, + AcceptanceOutcome::ResidualResample, + AcceptanceOutcome::ResidualDegenerateFallback, + ] { + outcome.latch().store(false, Ordering::Relaxed); + } +} + +/// True when a [`SamplingConfig`] makes the fused sampler deterministic. +/// +/// Mirrors the greedy short-circuit in `fused_sample_impl` +/// (`temperature == 0.0f || top_k == 1` selects `argmax`). Keeping the two in +/// step matters: if this predicate said "stochastic" for a config the sampler +/// treats as greedy, the accept test would compare against a one-hot `p` and +/// reject every non-argmax draft for no reason. +#[inline] +pub fn sampler_is_greedy(config: &SamplingConfig) -> bool { + config.temperature == 0.0 || config.top_k == 1 +} + +/// Whether modified rejection sampling is enabled for this process. +/// +/// **Opt-in, default off.** Reads [`STOCHASTIC_ACCEPT_ENV`] once; only an +/// explicitly truthy value enables it. +/// +/// The default is off because on the one path this rule reaches, the classic +/// [`crate::speculative::SpeculativeGenerator`], the rule it replaces was +/// already distribution-preserving. So the change buys no correctness there, +/// only acceptance rate, and the acceptance gain is bounded by +/// `sum_x min(p, q) / sum_x p(x) q(x)`, which collapses toward 1 whenever the +/// drafter is confident. Measured on a Llama-3.1-8B / Llama-3.2-1B pair at +/// temperature 0.7 that ratio is about 1.02. Paying two extra +/// full-vocabulary passes and an extra host sync per verified position for a +/// two-percent theoretical acceptance gain is not a good default. +/// +/// It is worth enabling where the gain is real: a high-entropy drafter, a +/// higher temperature, or a verify path whose current rule is *not* +/// distribution-preserving (the Gemma 4 MTP and DFlash round loops, which pick +/// the target token by argmax regardless of temperature). Check the expected +/// gain first with [`ACCEPT_DIAG_ENV`], which reports both closed forms. +pub fn stochastic_accept_enabled() -> bool { + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| match std::env::var(STOCHASTIC_ACCEPT_ENV) { + Ok(raw) => matches!( + raw.trim().to_ascii_lowercase().as_str(), + "1" | "true" | "yes" | "on" + ), + Err(_) => false, + }) +} + +/// Pick the acceptance rule for a verify round. +/// +/// `has_proposal_distribution` is the caller's assertion that it can supply +/// `q` for every drafted position. A drafter that cannot report what it +/// proposed from leaves the round on argmax rather than silently inventing a +/// `q`, because an invented `q` breaks the losslessness proof outright. +pub fn acceptance_rule( + target_config: &SamplingConfig, + has_proposal_distribution: bool, +) -> AcceptanceRule { + acceptance_rule_with_override(target_config, has_proposal_distribution, None) +} + +/// [`acceptance_rule`] with an explicit per-caller override of the env default. +/// +/// `stochastic_override` of `None` consults [`stochastic_accept_enabled`]; +/// `Some(v)` uses `v` directly. This is what makes the opt-in reachable +/// programmatically rather than only through a process-wide environment +/// variable, which matters twice over: a future per-request server flag needs +/// it, and the distributional tests need to exercise the acceptance-optimal +/// rule regardless of how the process default happens to be set. A test that +/// silently followed the default would stop covering this module the moment the +/// default flipped. +pub fn acceptance_rule_with_override( + target_config: &SamplingConfig, + has_proposal_distribution: bool, + stochastic_override: Option, +) -> AcceptanceRule { + let want_stochastic = stochastic_override.unwrap_or_else(stochastic_accept_enabled); + if sampler_is_greedy(target_config) { + return AcceptanceRule::GreedyArgmax; + } + if !want_stochastic { + return AcceptanceRule::SamplerMatch; + } + if !has_proposal_distribution { + return AcceptanceRule::SamplerMatchNoProposalDistribution; + } + AcceptanceRule::Stochastic +} + +/// Environment switch for the per-position acceptance diagnostic. +pub const ACCEPT_DIAG_ENV: &str = "MLXCEL_SPECULATIVE_ACCEPT_DIAG"; + +/// Whether to accumulate the closed-form acceptance probabilities per verify +/// position. +/// +/// Off by default: it costs two full-vocabulary reductions and one host +/// readback per verified position, which is real work on a 128K vocabulary. +/// On, it turns "the acceptance rate looks wrong" into an arithmetic +/// statement, because `sum_x min(p, q) >= sum_x p(x) q(x)` holds for every +/// pair of distributions. If the measured rate does not sit at `sum min` while +/// `sum min` sits at or above `sum prod`, the defect is localized immediately +/// to either the accept test or to `p` and `q` themselves. +pub fn accept_diagnostics_enabled() -> bool { + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| match std::env::var(ACCEPT_DIAG_ENV) { + Ok(raw) => !matches!( + raw.trim().to_ascii_lowercase().as_str(), + "" | "0" | "false" | "no" | "off" + ), + Err(_) => false, + }) +} + +/// Closed-form acceptance probabilities at one verify position. +/// +/// `sum_min` is the probability modified rejection sampling accepts, and also +/// the ceiling for *any* distribution-preserving rule (maximal coupling). +/// `sum_prod` is the probability the sampler-match rule accepts. +/// `sum_min >= sum_prod` always, because `min(a, b) >= a * b` for `a, b` in +/// `[0, 1]`. Their ratio is the entire acceptance headroom this feature can +/// deliver for the given `(p, q)`. +pub fn closed_form_acceptance(target_probs: &MlxArray, proposal_probs: &MlxArray) -> (f64, f64) { + let sum_min = ffi::sum_all(&ffi::minimum(target_probs, proposal_probs)); + let sum_prod = ffi::sum_all(&ffi::multiply(target_probs, proposal_probs)); + ffi::eval(&sum_min); + ffi::eval(&sum_prod); + ( + f64::from(ffi::item_f32(&sum_min)), + f64::from(ffi::item_f32(&sum_prod)), + ) +} + +/// The verdict for one drafted token. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DraftVerdict { + /// Keep the drafted token and continue the chain. + Accept, + /// Drop the drafted token, emit `replacement`, and end the chain. + Reject { + /// Token drawn from the normalized residual (or from `p` in the + /// degenerate case recorded by [`AcceptanceOutcome`]). + replacement: i32, + }, +} + +/// Read `probs[0, token]` back to the host as an `f32`. +/// +/// `probs` is `[batch, vocab]`; only row 0 is read, which is all the B=1 +/// speculative paths need. +fn probability_of(probs: &MlxArray, token: i32) -> UniquePtr { + let index = ffi::from_slice_i32(&[token], &[1, 1]); + ffi::take_along_axis(probs, &index, -1) +} + +/// The modified rejection-sampling accept test for one drafted token. +/// +/// * `target_probs` — `p`, float32 `[1, vocab]`, from +/// [`crate::sampling::effective_token_distribution`] on the target's logits +/// at this verify position. +/// * `proposal_probs` — `q`, float32 `[1, vocab]`, the distribution the +/// drafter drew `draft_token` from. Must be the *actual* proposal +/// distribution; see the module docs. +/// +/// Accepts iff `p(t) > 0` **and** `u * q(t) <= p(t)` for a fresh `u ~ U[0, 1)`. +/// +/// The `p(t) > 0` conjunct is not cosmetic. When the target's top-k / top-p / +/// min-p filter excludes the drafted token, `p(t)` is exactly zero and the +/// token must always be rejected. Without the conjunct a `q(t)` that has +/// underflowed float32 to zero would make the product `0 <= 0` and accept a +/// token the target assigns no mass at all, which is precisely the failure the +/// whole change exists to prevent. +/// +/// One `eval` and one host readback per call, the same synchronization cost as +/// the argmax comparison it replaces. +pub fn accept_draft_token( + target_probs: &MlxArray, + proposal_probs: &MlxArray, + draft_token: i32, +) -> bool { + let p_t = probability_of(target_probs, draft_token); + let q_t = probability_of(proposal_probs, draft_token); + + // SAFETY: a null key means "draw from the thread-local default RNG state", + // the same convention `apply_xtc_step` and `layers.rs` use. Drawing through + // MLX's key sequence keeps the accept decisions reproducible under + // `ffi::random_seed`. + let u = unsafe { ffi::random_uniform(0.0, 1.0, &[1, 1], dtype::FLOAT32, std::ptr::null()) }; + + let zero = ffi::zeros(&[1, 1], dtype::FLOAT32); + let target_has_mass = ffi::greater(&p_t, &zero); + let within_ratio = ffi::less_equal(&ffi::multiply(&u, &q_t), &p_t); + let accept = ffi::logical_and(&target_has_mass, &within_ratio); + + ffi::eval(&accept); + ffi::item_bool(&accept) +} + +/// Draw the replacement token from the normalized residual `relu(p - q)`. +/// +/// Returns the token and which outcome kind produced it, so the caller can +/// feed [`note_outcome`]. +/// +/// `relu(p - q)` is non-negative by construction and sums to the rejection +/// probability, so normalizing it is exactly the residual distribution the +/// losslessness proof requires. Sampling runs through +/// [`ffi::fused_sample`] on `log(relu(p - q))` with `temperature = 1.0` and +/// every filter disabled: `softmax(log r) == r / sum(r)`, so this is the +/// normalized residual and nothing else. Routing through `fused_sample` also +/// means the residual draw uses the same #900 Gumbel-max kernel the rest of +/// the sampler uses where the backend supports it, and the `categorical` graph +/// path where it does not. Entries with zero residual become `-inf` and can +/// never win. +pub fn residual_resample( + target_probs: &MlxArray, + proposal_probs: &MlxArray, +) -> (i32, AcceptanceOutcome) { + let residual = ffi::relu(&ffi::subtract(target_probs, proposal_probs)); + let mass = ffi::sum_all(&residual); + ffi::eval(&mass); + let mass = ffi::item_f32(&mass); + + let (source, outcome) = if mass > 0.0 && mass.is_finite() { + (residual, AcceptanceOutcome::ResidualResample) + } else { + // Unreachable in exact arithmetic: reaching this function means the + // token was rejected, which implies positive residual mass. Only + // float32 underflow of every positive entry gets here. Falling back to + // `p` keeps the emitted token inside the target's support, which is + // the closest available approximation of the residual. + ( + ffi::copy(target_probs), + AcceptanceOutcome::ResidualDegenerateFallback, + ) + }; + + let token_arr = ffi::fused_sample(&ffi::log(&source), 1.0, 0, 1.0, 0.0); + ffi::eval(&token_arr); + (ffi::item_i32(&token_arr), outcome) +} + +/// Run the full per-token decision: accept test, and residual resample on +/// rejection. Records both the accept/reject outcome and, on rejection, +/// whether the residual was degenerate. +pub fn verify_draft_token( + target_probs: &MlxArray, + proposal_probs: &MlxArray, + draft_token: i32, +) -> DraftVerdict { + if accept_draft_token(target_probs, proposal_probs, draft_token) { + note_outcome(AcceptanceOutcome::Accepted); + DraftVerdict::Accept + } else { + let (replacement, outcome) = residual_resample(target_probs, proposal_probs); + note_outcome(outcome); + DraftVerdict::Reject { replacement } + } +} + +#[cfg(test)] +#[path = "stochastic_accept_tests.rs"] +mod stochastic_accept_tests; diff --git a/src/lib/mlxcel-core/src/speculative/stochastic_accept_tests.rs b/src/lib/mlxcel-core/src/speculative/stochastic_accept_tests.rs new file mode 100644 index 000000000..a1f1dfbd0 --- /dev/null +++ b/src/lib/mlxcel-core/src/speculative/stochastic_accept_tests.rs @@ -0,0 +1,571 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Unit tests for the modified rejection-sampling primitives. +//! +//! These test the accept rule and the residual draw in isolation, against +//! hand-built `p` and `q` vectors where the exact answer is known in closed +//! form. The end-to-end distributional argument for the generator lives in +//! `distribution_tests.rs`. + +use super::*; +use crate::sampling::effective_token_distribution; + +/// Upper-tail chi-square critical values at alpha = 1e-4. +/// +/// Sourced from the chi-square quantile function; indexed by degrees of +/// freedom. A false-failure rate of 1e-4 per assertion is the tolerance these +/// tests accept in exchange for keeping the power to reject a real defect (see +/// the power calibration in `distribution_tests.rs`). +const CHI2_CRIT_1E4: [f64; 8] = [ + f64::NAN, // dof 0, unused + 15.1367, // dof 1 + 18.4207, // dof 2 + 21.1075, // dof 3 + 23.5127, // dof 4 + 25.7448, // dof 5 + 27.8563, // dof 6 + 29.8775, // dof 7 +]; + +/// Pearson chi-square of `observed` counts against `expected` probabilities. +fn chi_square(observed: &[u32], expected: &[f64]) -> f64 { + let n: u32 = observed.iter().sum(); + let n = f64::from(n); + observed + .iter() + .zip(expected) + .map(|(&o, &e)| { + let e = n * e; + let d = f64::from(o) - e; + d * d / e + }) + .sum() +} + +/// Build a `[1, vocab]` float32 probability row from an explicit vector. +fn probs_row(values: &[f32]) -> UniquePtr { + ffi::from_slice_f32(values, &[1, values.len() as i32]) +} + +/// Build a `[1, 1, vocab]` logits tensor, the shape the samplers consume. +fn logits_tensor(values: &[f32]) -> UniquePtr { + ffi::from_slice_f32(values, &[1, 1, values.len() as i32]) +} + +/// Read a `[1, vocab]` float32 tensor back to the host. +fn to_host(arr: &MlxArray) -> Vec { + ffi::eval(arr); + let bytes = ffi::array_to_raw_bytes(arr); + bytes + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect() +} + +fn softmax(logits: &[f32], temperature: f32) -> Vec { + let scaled: Vec = logits + .iter() + .map(|&x| f64::from(x) / f64::from(temperature)) + .collect(); + let max = scaled.iter().cloned().fold(f64::NEG_INFINITY, f64::max); + let exp: Vec = scaled.iter().map(|x| (x - max).exp()).collect(); + let sum: f64 = exp.iter().sum(); + exp.iter().map(|x| x / sum).collect() +} + +fn stochastic_config(temperature: f32) -> SamplingConfig { + SamplingConfig { + temperature, + ..SamplingConfig::default() + } +} + +// --------------------------------------------------------------------------- +// Rule selection, greedy predicate, kill switch +// --------------------------------------------------------------------------- + +/// The greedy predicate must agree with the C++ sampler's own short-circuit +/// (`temperature == 0.0f || top_k == 1` selects `argmax`). If it drifted, the +/// accept test would compare a stochastic draft against a one-hot `p` and +/// reject nearly everything. +#[test] +fn sampler_is_greedy_matches_fused_sampler_short_circuit() { + assert!(sampler_is_greedy(&SamplingConfig::greedy())); + assert!(sampler_is_greedy(&stochastic_config(0.0))); + assert!(sampler_is_greedy(&SamplingConfig { + temperature: 0.9, + top_k: 1, + ..SamplingConfig::default() + })); + assert!(!sampler_is_greedy(&stochastic_config(0.7))); + assert!(!sampler_is_greedy(&SamplingConfig { + temperature: 1.0, + top_k: 40, + ..SamplingConfig::default() + })); +} + +/// A greedy target keeps the pre-#902 argmax rule regardless of everything +/// else, which is what makes temperature 0 byte-identical. +#[test] +fn greedy_target_always_selects_argmax_rule() { + assert_eq!( + acceptance_rule(&SamplingConfig::greedy(), true), + AcceptanceRule::GreedyArgmax + ); + assert_eq!( + acceptance_rule(&SamplingConfig::greedy(), false), + AcceptanceRule::GreedyArgmax + ); +} + +/// A drafter that cannot report `q` must not be given a fabricated one. +#[test] +fn missing_proposal_distribution_falls_back_to_argmax() { + let rule = acceptance_rule(&stochastic_config(0.7), false); + assert!(matches!( + rule, + AcceptanceRule::SamplerMatchNoProposalDistribution | AcceptanceRule::SamplerMatch + )); + // Falling back is safe: sampler-match is distribution-preserving. What is + // lost is acceptance optimality, not correctness. + assert!(rule.is_distribution_preserving()); + assert!(!rule.is_acceptance_optimal()); +} + +/// **Every** rule this module can select is distribution-preserving, including +/// the default. That is the central correction to issue #902's premise: the +/// rule the classic generator shipped with compares against a fresh draw from +/// the target *sampler*, not its argmax, and emits that draw on both branches. +/// +/// What separates the rules is acceptance optimality. `sum_x min(p, q)` is the +/// maximal-coupling ceiling for any correct rule; modified rejection sampling +/// attains it, sampler-match reaches only `sum_x p(x) q(x)`. +#[test] +fn every_selectable_rule_is_distribution_preserving() { + for rule in [ + AcceptanceRule::GreedyArgmax, + AcceptanceRule::SamplerMatch, + AcceptanceRule::Stochastic, + AcceptanceRule::SamplerMatchNoProposalDistribution, + ] { + assert!( + rule.is_distribution_preserving(), + "{} must be safe to serve: this module never selects the biased \ + argmax-against-argmax rule, which lives in the MTP and DFlash paths", + rule.id() + ); + } + + assert!(AcceptanceRule::Stochastic.is_acceptance_optimal()); + assert!(AcceptanceRule::GreedyArgmax.is_acceptance_optimal()); + assert!( + !AcceptanceRule::SamplerMatch.is_acceptance_optimal(), + "sampler-match reaches sum_x p(x) q(x), strictly below the \ + sum_x min(p, q) ceiling" + ); + assert!(!AcceptanceRule::SamplerMatchNoProposalDistribution.is_acceptance_optimal()); +} + +/// The env switch is **opt-in**: absence leaves the acceptance-optimal rule +/// off. Tested through the same matcher the runtime uses, since the +/// process-wide `OnceLock` cannot be re-read in-process. +#[test] +fn stochastic_acceptance_is_opt_in_not_opt_out() { + let truthy = |raw: &str| { + matches!( + raw.trim().to_ascii_lowercase().as_str(), + "1" | "true" | "yes" | "on" + ) + }; + for v in ["1", "true", "TRUE", " on ", "Yes"] { + assert!(truthy(v), "{v} must enable stochastic acceptance"); + } + for v in ["0", "false", "off", "no", "", "maybe"] { + assert!( + !truthy(v), + "{v} must leave the default sampler-match rule in place" + ); + } +} + +/// The default, with the variable unset, must be the pre-#902 rule. A silent +/// flip of this default would change every speculative token stream at +/// `temperature > 0` and cost throughput on pairs where the gain is nil. +#[test] +fn default_rule_at_positive_temperature_is_sampler_match() { + // Assumes the test process does not set the variable, which is how CI and + // a plain `cargo test` run. Deliberately does not mutate the environment: + // `set_var` is unsound with parallel tests and the gate is a `OnceLock` + // that a sibling test may already have resolved. + assert!( + std::env::var(STOCHASTIC_ACCEPT_ENV).is_err(), + "this test requires {STOCHASTIC_ACCEPT_ENV} to be unset" + ); + assert!( + !stochastic_accept_enabled(), + "stochastic acceptance must be opt-in" + ); + assert_eq!( + acceptance_rule(&stochastic_config(0.7), true), + AcceptanceRule::SamplerMatch + ); +} + +/// The one-shot latches fire once per kind, not once globally: a process that +/// serves a greedy request and then a stochastic one must log both. +#[test] +fn log_latches_are_per_kind_and_one_shot() { + reset_log_latches(); + assert!(!AcceptanceRule::GreedyArgmax.latch().load(Ordering::Relaxed)); + note_rule(AcceptanceRule::GreedyArgmax); + assert!(AcceptanceRule::GreedyArgmax.latch().load(Ordering::Relaxed)); + assert!( + !AcceptanceRule::Stochastic.latch().load(Ordering::Relaxed), + "logging one rule must not swallow a different rule's first occurrence" + ); + note_rule(AcceptanceRule::Stochastic); + assert!(AcceptanceRule::Stochastic.latch().load(Ordering::Relaxed)); + + note_outcome(AcceptanceOutcome::ResidualResample); + assert!( + AcceptanceOutcome::ResidualResample + .latch() + .load(Ordering::Relaxed) + ); + assert!( + !AcceptanceOutcome::ResidualDegenerateFallback + .latch() + .load(Ordering::Relaxed) + ); + reset_log_latches(); +} + +// --------------------------------------------------------------------------- +// effective_token_distribution +// --------------------------------------------------------------------------- + +/// The stochastic distribution must be the temperature-scaled softmax, to +/// float32 precision. This is the `p` of the accept test, so an error here is +/// an error in every downstream claim. +#[test] +fn effective_distribution_is_the_temperature_scaled_softmax() { + let logits = [2.0f32, 1.0, 0.5, 0.0, -0.5, -1.0]; + let config = stochastic_config(0.7); + let probs = effective_token_distribution(&logits_tensor(&logits), &config, &[]); + let got = to_host(&probs); + let want = softmax(&logits, 0.7); + + let sum: f64 = got.iter().map(|&x| f64::from(x)).sum(); + assert!( + (sum - 1.0).abs() < 1e-5, + "distribution must normalize: {sum}" + ); + for (i, (g, w)) in got.iter().zip(&want).enumerate() { + assert!( + (f64::from(*g) - w).abs() < 1e-5, + "entry {i}: got {g}, want {w}" + ); + } +} + +/// A greedy config yields the one-hot indicator at the argmax. This is the +/// degenerate `q` the issue calls for when a drafter proposes greedily; the +/// accept test then reduces to `u <= p(t)`. +#[test] +fn effective_distribution_is_one_hot_for_a_greedy_config() { + let logits = [0.1f32, 3.0, 0.2, -1.0]; + let probs = + effective_token_distribution(&logits_tensor(&logits), &SamplingConfig::greedy(), &[]); + let got = to_host(&probs); + assert_eq!(got, vec![0.0, 1.0, 0.0, 0.0]); + + // `top_k == 1` is the sampler's other greedy short-circuit. + let top1 = SamplingConfig { + temperature: 0.8, + top_k: 1, + ..SamplingConfig::default() + }; + let probs = effective_token_distribution(&logits_tensor(&logits), &top1, &[]); + assert_eq!(to_host(&probs), vec![0.0, 1.0, 0.0, 0.0]); +} + +/// A filtered target has *zero* mass outside its support. The accept test +/// depends on this being exactly zero, not merely small: it is what forces a +/// drafted token the target filtered out to be rejected. +#[test] +fn effective_distribution_zeroes_the_filtered_tail() { + let logits = [3.0f32, 2.5, 1.0, 0.5, 0.0, -2.0]; + let config = SamplingConfig { + temperature: 1.0, + top_k: 2, + ..SamplingConfig::default() + }; + let got = to_host(&effective_token_distribution( + &logits_tensor(&logits), + &config, + &[], + )); + assert!(got[0] > 0.0 && got[1] > 0.0); + for (i, &v) in got.iter().enumerate().skip(2) { + assert_eq!(v, 0.0, "token {i} is outside top-2 and must carry no mass"); + } + let sum: f32 = got.iter().sum(); + assert!((sum - 1.0).abs() < 1e-5); +} + +// --------------------------------------------------------------------------- +// accept_draft_token +// --------------------------------------------------------------------------- + +/// `p(t) >= q(t)` implies `min(1, p/q) == 1`: the token must be accepted every +/// single time, with no tolerance for a stray rejection. +#[test] +fn tokens_the_target_likes_at_least_as_much_are_always_accepted() { + let p = probs_row(&[0.6, 0.3, 0.1, 0.0]); + let q = probs_row(&[0.2, 0.2, 0.1, 0.5]); + for trial in 0..400 { + for token in [0, 1] { + assert!( + accept_draft_token(&p, &q, token), + "trial {trial}: p({token}) >= q({token}) must accept unconditionally" + ); + } + } +} + +/// A token the target assigns zero mass must be rejected every single time. +/// This is the zero-tolerance guard: one accepted token here is one token in +/// the served stream that the target's own sampler could never have produced. +#[test] +fn tokens_outside_the_target_support_are_never_accepted() { + let p = probs_row(&[0.7, 0.3, 0.0, 0.0]); + let q = probs_row(&[0.1, 0.1, 0.4, 0.4]); + for trial in 0..400 { + for token in [2, 3] { + assert!( + !accept_draft_token(&p, &q, token), + "trial {trial}: p({token}) == 0 must reject unconditionally" + ); + } + } +} + +/// A `q(t)` that has underflowed float32 to exactly zero must not turn into an +/// unconditional accept for a token the target also gives zero mass. Without +/// the `p(t) > 0` conjunct the product test reads `0 <= 0` and accepts. +#[test] +fn underflowed_proposal_mass_does_not_bypass_the_target_support() { + let p = probs_row(&[1.0, 0.0, 0.0]); + let q = probs_row(&[0.0, 0.0, 1.0]); + for _ in 0..200 { + assert!( + !accept_draft_token(&p, &q, 1), + "q(1) == 0 and p(1) == 0 must reject, not accept" + ); + } + for _ in 0..200 { + assert!( + accept_draft_token(&p, &q, 0), + "q(0) == 0 with p(0) > 0 accepts: the token was genuinely proposed" + ); + } +} + +/// The empirical acceptance rate must match `min(1, p(t)/q(t))`. +/// +/// Null hypothesis: `accept_draft_token` accepts with probability +/// `min(1, p(t)/q(t))`. Binomial with n = 4000; the assertion band is +-4 +/// standard errors, so a correct implementation fails about 6e-5 of the time +/// per token tested. +#[test] +fn acceptance_probability_matches_the_min_one_ratio() { + let p_vals = [0.50f32, 0.20, 0.20, 0.10]; + let q_vals = [0.10f32, 0.40, 0.20, 0.30]; + let p = probs_row(&p_vals); + let q = probs_row(&q_vals); + + const TRIALS: usize = 4000; + for token in 0..4usize { + let expected = f64::from(p_vals[token] / q_vals[token]).min(1.0); + let accepts = (0..TRIALS) + .filter(|_| accept_draft_token(&p, &q, token as i32)) + .count(); + let rate = accepts as f64 / TRIALS as f64; + let se = (expected * (1.0 - expected) / TRIALS as f64).sqrt(); + let band = 4.0 * se + 1e-12; + assert!( + (rate - expected).abs() <= band, + "token {token}: acceptance rate {rate:.4} outside {expected:.4} +- {band:.4}" + ); + } +} + +// --------------------------------------------------------------------------- +// residual_resample +// --------------------------------------------------------------------------- + +/// The residual draw must follow `normalize(relu(p - q))`. +/// +/// Null hypothesis: the residual sample is distributed as +/// `relu(p - q) / sum(relu(p - q))`. Pearson chi-square over the 3 tokens with +/// positive residual (dof = 2), n = 6000, rejected above 18.42 (alpha = 1e-4). +#[test] +fn residual_draw_follows_the_normalized_relu_difference() { + let p_vals = [0.50f32, 0.25, 0.15, 0.07, 0.03]; + let q_vals = [0.20f32, 0.40, 0.05, 0.30, 0.05]; + let p = probs_row(&p_vals); + let q = probs_row(&q_vals); + + let residual: Vec = p_vals + .iter() + .zip(&q_vals) + .map(|(&a, &b)| f64::from(a - b).max(0.0)) + .collect(); + let mass: f64 = residual.iter().sum(); + let expected: Vec = residual.iter().map(|r| r / mass).collect(); + + const TRIALS: usize = 6000; + let mut counts = vec![0u32; p_vals.len()]; + for _ in 0..TRIALS { + let (token, outcome) = residual_resample(&p, &q); + assert_eq!(outcome, AcceptanceOutcome::ResidualResample); + assert!( + expected[token as usize] > 0.0, + "residual drew token {token}, which has p <= q and therefore zero residual mass" + ); + counts[token as usize] += 1; + } + + // Only the support of the residual carries expected mass; restrict the + // statistic to it so no zero-expectation cell appears in the denominator. + let support: Vec = (0..p_vals.len()).filter(|&i| expected[i] > 0.0).collect(); + let obs: Vec = support.iter().map(|&i| counts[i]).collect(); + let exp: Vec = support.iter().map(|&i| expected[i]).collect(); + let dof = support.len() - 1; + let chi2 = chi_square(&obs, &exp); + assert!( + chi2 < CHI2_CRIT_1E4[dof], + "residual distribution chi-square {chi2:.3} exceeds the alpha=1e-4 \ + critical value {:.3} at dof={dof}; observed {counts:?}, expected {expected:?}", + CHI2_CRIT_1E4[dof] + ); +} + +/// `p == q` leaves no residual mass at all. The documented degenerate fallback +/// must fire, be reported as such, and still emit a token inside `p`'s support. +#[test] +fn identical_distributions_take_the_degenerate_residual_fallback() { + let vals = [0.6f32, 0.3, 0.1, 0.0]; + let p = probs_row(&vals); + let q = probs_row(&vals); + for _ in 0..64 { + let (token, outcome) = residual_resample(&p, &q); + assert_eq!(outcome, AcceptanceOutcome::ResidualDegenerateFallback); + assert!( + vals[token as usize] > 0.0, + "the fallback must stay inside the target's support, got token {token}" + ); + } +} + +/// `verify_draft_token` composes the two primitives without changing either +/// one's contract: an always-accepted token never resamples, and a token +/// outside the target support always comes back as a rejection whose +/// replacement carries target mass. +#[test] +fn verify_draft_token_composes_accept_and_residual() { + let p_vals = [0.7f32, 0.3, 0.0, 0.0]; + let p = probs_row(&p_vals); + let q = probs_row(&[0.1, 0.1, 0.4, 0.4]); + + for _ in 0..200 { + assert_eq!(verify_draft_token(&p, &q, 0), DraftVerdict::Accept); + } + for _ in 0..200 { + match verify_draft_token(&p, &q, 2) { + DraftVerdict::Accept => panic!("token 2 has zero target mass and must be rejected"), + DraftVerdict::Reject { replacement } => assert!( + p_vals[replacement as usize] > 0.0, + "replacement {replacement} must lie in the target's support" + ), + } + } +} + +// --------------------------------------------------------------------------- +// Acceptance-rate improvement +// --------------------------------------------------------------------------- + +/// The measured acceptance rate over proposals drawn from `q` must converge to +/// `sum_x min(p(x), q(x))`, the information-theoretic optimum for this pair, +/// and must beat the `sum_x p(x) q(x)` rate of the pre-#902 rule (which +/// accepts only when an independent target draw happens to land on the same +/// token). +/// +/// Null hypothesis: the accept rate equals `sum_x min(p(x), q(x))`. n = 8000 +/// proposals, band of +-4 binomial standard errors. +#[test] +fn measured_acceptance_rate_reaches_the_sum_min_optimum() { + let p_vals = [0.55f32, 0.20, 0.12, 0.08, 0.05]; + let q_vals = [0.30f32, 0.35, 0.15, 0.15, 0.05]; + let p = probs_row(&p_vals); + let q = probs_row(&q_vals); + + let sum_min: f64 = p_vals + .iter() + .zip(&q_vals) + .map(|(&a, &b)| f64::from(a.min(b))) + .sum(); + let sum_prod: f64 = p_vals + .iter() + .zip(&q_vals) + .map(|(&a, &b)| f64::from(a) * f64::from(b)) + .sum(); + assert!( + sum_min > sum_prod, + "the pair must actually favor the new rule: sum min {sum_min} vs sum prod {sum_prod}" + ); + + // Draw proposals from `q` by inverse-CDF over a deterministic sweep of the + // unit interval, so the measured rate isolates the accept test's own + // randomness instead of compounding two sampling errors. + const TRIALS: usize = 8000; + let mut cdf = Vec::with_capacity(q_vals.len()); + let mut acc = 0.0f64; + for &v in &q_vals { + acc += f64::from(v); + cdf.push(acc); + } + let mut accepts = 0usize; + for i in 0..TRIALS { + let u = (i as f64 + 0.5) / TRIALS as f64; + let token = cdf.iter().position(|&c| u < c).unwrap_or(q_vals.len() - 1); + if accept_draft_token(&p, &q, token as i32) { + accepts += 1; + } + } + let rate = accepts as f64 / TRIALS as f64; + let se = (sum_min * (1.0 - sum_min) / TRIALS as f64).sqrt(); + let band = 4.0 * se; + assert!( + (rate - sum_min).abs() <= band, + "acceptance rate {rate:.4} outside the sum-min optimum {sum_min:.4} +- {band:.4}" + ); + assert!( + rate > sum_prod, + "acceptance rate {rate:.4} must exceed the pre-#902 sum-product rate {sum_prod:.4}" + ); +} diff --git a/src/main.rs b/src/main.rs index 1cf70e79f..67cdac5d0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1780,9 +1780,43 @@ impl Default for DiffusionServeOptions { } } +/// Install a `tracing` subscriber for the offline CLI commands when `RUST_LOG` +/// is set. +/// +/// The CLI shipped without any global subscriber, so every `tracing::*` call +/// reached by `mlxcel generate`, `inspect`, `tune` and friends was silently +/// discarded no matter what `RUST_LOG` said. That is exactly backwards: the +/// instrumentation is unreachable precisely when someone has set `RUST_LOG` +/// because they are trying to diagnose something. Notably it hid the drafter +/// auto-detection line, which is the difference between a benchmark that ran +/// the classic `SpeculativeGenerator` and one that ran a kind-specific round +/// loop. +/// +/// Two deliberate constraints: +/// +/// * **Opt-in.** Nothing is installed unless `RUST_LOG` is set, so default +/// command output is unchanged byte for byte. +/// * **stderr.** Diagnostics must not contaminate the generated text on stdout, +/// which callers pipe. +/// +/// `Serve` is skipped because `server::startup` installs its own subscriber and +/// would panic on the double `init()`. `try_init` is used anyway so a future +/// second initializer degrades to a no-op rather than aborting the process. +fn init_cli_tracing(command: &Commands) { + if matches!(command, Commands::Serve(_)) || std::env::var_os("RUST_LOG").is_none() { + return; + } + let _ = tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .with_writer(std::io::stderr) + .try_init(); +} + fn main() -> anyhow::Result<()> { let cli = Cli::parse(); + init_cli_tracing(&cli.command); + // Default the CUDA kernel JIT cache to a persistent, MLX-pin-scoped dir so // the first-run kernel compilation is paid once per machine, not every boot. mlxcel_core::ensure_persistent_ptx_cache();