feat(speculative): opt-in acceptance-optimal rejection sampling, and a premise correction

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)

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
 

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

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)