From 82bf64d1efa022607d0bcc1225c5775125cdb466 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sat, 22 Aug 2026 13:42:36 +0900 Subject: [PATCH 1/2] perf(speculative): decide the MTP verify width by measured throughput, not an acceptance proxy The adaptive block-size controller held a drafter at its configured depth until the configured prefix was usually fully accepted. That proxy cannot pass on the Gemma 4 12B pairing (acceptance 0.585 over a 3-proposal prefix), so the pairing ran about 5% below its own measured optimum: 93.54 tok/s at the held width 4 against 98.16 at the requested width 5 on an M5 Max (issue #1207). The B = 1 round loop now decides with the signal widening exists to improve: emitted tokens per millisecond of round time, measured on the rounds it already runs. BlockThroughputController alternates 32-round measurement windows between the configured depth and the requested ceiling, holds whichever measures faster (a challenger needs a 2% lead; ties go to the configured depth), and re-challenges the loser on a 4/16/64-window backoff. A collapsing challenger is aborted after 4 rounds once it trails by more than 35%, so a harmful ceiling (Qwen 3.8 at width 12 measures 5.80 against 21.30 tok/s) costs rounds, not windows. The search space stays the proxy's two widths: --draft-block-size is a ceiling the user set, not a hint to wander from, and both widths are bounded by the one the exactness probe ran at, so no new verify shape reaches the gate. Evidence lives in the generator, so a server process keeps it across requests. Drafters that set prefer_requested_block_size (Qwen 3.5 MTP) bypass both controllers, exactly as before. MLXCEL_MTP_BLOCK_CONTROLLER=proxy restores the upstream gate; the batched (B > 1) loop stays on it regardless, because the row-averaged accept length is the only per-round signal that loop measures today, and the #1207 measurement is B = 1. Refs #1207 --- docs/environment-variables.md | 1 + .../src/speculative/mtp/adaptive.rs | 341 +++++++++++++++++- .../src/speculative/mtp/generator.rs | 104 +++++- .../src/speculative/mtp/round_loop_batched.rs | 6 + 4 files changed, 426 insertions(+), 26 deletions(-) diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 1ba409e1e..9a172ade1 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -188,6 +188,7 @@ The OpenAI audio endpoints (`/v1/audio/speech`, `/v1/audio/transcriptions`, `/v1 | `MLXCEL_GDN_CHAIN_PARITY` | `0` to disable, any other value (or unset) to enable | on | **Advanced, diagnostic escape hatch.** Gates the chain-parity gated-delta Metal kernel used by Qwen 3.5 MTP's speculative verify and rollback-replay paths (issue #1165). The standard gated-delta kernel carries float32 recurrent state across a `T = K` verify block and rounds it to the storage dtype only once at the end, while the classic single-token decode chain rounds after every token; a `T = K` verify block is therefore NOT bit-identical to `K` consecutive single-token decode steps unless the state is rounded per in-block step. The chain-parity kernel (`gated_delta_step_seqpar`) does that rounding, which is what makes Qwen 3.5 MTP's temperature-0 output byte-identical to classic decode. **Setting this to `0` forfeits that exactness contract**, restoring the pre-#1165 block numerics for A/B attribution of the parity kernel's own cost and acceptance effect; do not set it to `0` in a deployment that needs byte-identical speculative output. Metal-only: the non-Metal ops fallback ignores the flag (the parity guarantee does not exist off Metal today). **The kernel is necessary but not sufficient**: byte-identity also requires every quantized projection to dispatch to the same MLX kernel at `M = block_size` as at `M = 1`, which is not true on every GPU generation or at every block width, so the runtime probe behind `MLXCEL_MTP_ALLOW_INEXACT` is what actually decides whether MTP engages. See `docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md` for the measured kernel cost (inside the dispatch-noise band). | | `MLXCEL_MTP_ALLOW_INEXACT` | `1`/`true`/`yes`/`on` to enable; unset or anything else to disable | **off** | Engage Qwen 3.5 MTP speculative decoding even when the startup exactness probe reports that the multi-token verify block is **not** byte-identical to the single-token decode chain. Before enabling MTP the runtime now measures the property instead of predicting it: one synthetic verify block and the equivalent single-token chain are run from the same state on the loaded checkpoint at the configured `--draft-block-size`, and their logits are compared byte for byte (three independent synthetic inputs, each two short prefills plus `K + 1` forwards; measured 4.9 s for the first call and 1.3 s for a later one per input on a Qwen3.8-27B 4-bit target on an M1 Ultra, the difference being MLX's one-time kernel compilation; more than one input because a kernel pair can disagree by only a byte or two out of ten thousand, at which amplitude a single draw can read as equal; memoized per (model, block width) and warmed at worker startup so it never lands on the request path). A divergence means temperature-0 speculative output would silently differ from `mlxcel generate` without `--draft-model`, so the default is to decline and run classic decode. The static conditions (Metal backend, `supports_metal_gated_delta_kernel` geometry) still apply and are checked first; this probe covers what they cannot, namely which MLX kernel each **quantized projection** dispatches to at `M = K` versus `M = 1`. That choice depends on the GPU generation, the quantization mode, the operand sizes and the block width: `use_qmv_wide` in [`mlx/backend/metal/quantized.cpp`](https://github.com/ml-explore/mlx/blob/main/mlx/backend/metal/quantized.cpp) sends `M >= 2` to a different reduction whenever `mode != "affine" || arch_gen >= 15`, and `get_qmv_batch_limit` sends `M` above 10, 12, 18 or 32 (by architecture size and generation) to the matrix-matrix kernel. Measured: an affine 4-bit Qwen3.8-27B target on an M1 Ultra is byte-identical at block widths 1 through 11 and diverges at 12 (the `arch_size == 'd'` branch); the same checkpoint on an M5 Max diverges from block width 2, where the `M >= 2` split fires before any batch limit can be observed. Within one checkpoint the limit is per projection, not per model: Gemma 4 12B's attention shapes hold to 17 on an M1 Ultra while its MLP shapes break at 12, so a model's own cliff is the minimum over its shapes, which is why this is measured rather than tabulated. Note the ordering that #1199 introduced: on a failing probe the gate first retries with `qmv_wide` disabled and keeps the narrow kernel when that restores exactness, and only a probe that fails **both** ways consults this flag. On Apple GPU generation 15+ the narrow retry passes, so this flag alone is inert there: the process is pinned narrow, output stays byte-identical, and the log shows the retry's INFO line rather than the override warning (verified live on M3 Ultra, 2026-08-22, byte-identical output with and without the flag; see `benchmark_results/qmv-wide-pin-tax-m3ultra-2026-08-22.md`). To research the fast kernel there, set `MLXCEL_QMV_WIDE=1` together with this flag: the pin skips the retry and this flag then engages MTP on the wide kernel, forfeiting byte-identity with the loud WARN. This flag alone is load-bearing only where no exact kernel selection exists at the configured block width. Read once per process. | | `MLXCEL_QMV_WIDE` | `0`/`false`/`no`/`off` to disable; `1` (or any other value) to pin wide | unset (wide, until the MTP gate's retry turns it off) | Operator pin for MLX's `qmv_wide` kernel, the faster reduction for `M >= 2` quantized matmuls on Apple GPU generation 15+ (overlay in `src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp`, added by #1199). **Setting the variable at all, to any value, counts as an operator pin**: the MTP exactness gate's retry (`retry_without_qmv_wide`) is skipped in both directions, so `MLXCEL_QMV_WIDE=1` keeps the wide kernel and makes a failing probe decline MTP instead of buying exactness back, and `MLXCEL_QMV_WIDE=0` runs the whole process narrow from the start. Unset, the kernel is wide until a failing MTP probe's retry finds the narrow kernel exact and pins the process narrow for good. The pin is process-wide and sits on the dispatch path of every quantized matmul; what non-MTP work pays for the narrow state is measured in `benchmark_results/qmv-wide-pin-tax-m3ultra-2026-08-22.md` (nothing measurable on batched decode, about 15 ms per prompt-cache-hit request's suffix prefill). Read once per process at first dispatch; `mlxcel_core::set_qmv_wide` can move it at runtime and the gate is its only caller. | +| `MLXCEL_MTP_BLOCK_CONTROLLER` | `proxy` to pin the acceptance-proxy gate; any other value (or unset) selects the throughput comparator | throughput | Which controller decides the B=1 MTP verify width when `--draft-block-size` exceeds the drafter's configured depth (issue #1207). The default measures the decision: after a short warm-up the round loop alternates measurement windows (32 rounds) between the configured depth and the requested ceiling, compares emitted tokens per millisecond of round time, holds whichever measures faster (a challenger needs a 2% lead; ties go to the configured depth), and re-challenges the loser on a growing backoff (4, 16, then every 64 windows), with a collapsing challenger aborted after 4 rounds once it trails by more than 35% so a harmful ceiling (the Qwen 3.8 pairing measures 5.80 against 21.30 tok/s at width 12) costs rounds rather than windows. Evidence lives in the generator, so a server process keeps it across requests; drafters that set `prefer_requested_block_size` (Qwen 3.5 MTP) bypass both controllers and always honour the request, exactly as before. Set to `proxy` to restore upstream's fully-accepted-prefix gate, which issue #1207 measured holding the Gemma 4 12B pairing about 5% below its optimum (93.54 against 98.16 tok/s at requested width 5 on an M5 Max) because at 0.585 acceptance the configured prefix is rarely fully accepted no matter how profitable widening is. The batched (B>1) loop stays on the proxy gate regardless: the row-averaged accept length is the only per-round signal it measures today. Read once per generator. | | `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). | diff --git a/src/lib/mlxcel-core/src/speculative/mtp/adaptive.rs b/src/lib/mlxcel-core/src/speculative/mtp/adaptive.rs index c12e6d6e3..a0dba0c18 100644 --- a/src/lib/mlxcel-core/src/speculative/mtp/adaptive.rs +++ b/src/lib/mlxcel-core/src/speculative/mtp/adaptive.rs @@ -12,14 +12,33 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Adaptive MTP block-size controller. +//! Adaptive MTP block-size controllers. //! -//! Ports upstream `mlx_vlm.speculative.mtp._effective_mtp_block_size`. -//! Gemma 4 assistants are configured for a 4-token verify block, but users -//! may request a larger `--draft-block-size`. The reference treats that -//! larger value as a ceiling: stay at the configured depth until the recent -//! acceptance history shows the configured prefix is usually fully accepted, -//! then expand to the requested ceiling. +//! Two of them, one signal each: +//! +//! [`effective_mtp_block_size`] ports upstream +//! `mlx_vlm.speculative.mtp._effective_mtp_block_size`. Gemma 4 assistants +//! are configured for a 4-token verify block, but users may request a larger +//! `--draft-block-size`. The reference treats that larger value as a +//! ceiling: stay at the configured depth until the recent acceptance history +//! shows the configured prefix is usually fully accepted, then expand to the +//! requested ceiling. Issue #1207 measured the flaw in that proxy: on the +//! Gemma 4 12B pairing acceptance never clears the bar, the controller never +//! expands, and the pairing runs about 5% below its own measured optimum +//! (93.54 against 98.16 tok/s at requested width 5 on an M5 Max). The proxy +//! remains in use on the batched (B > 1) round loop, whose per-row averaged +//! acceptance is the only per-round signal that loop currently measures, and +//! behind `MLXCEL_MTP_BLOCK_CONTROLLER=proxy` as the escape hatch. +//! +//! [`BlockThroughputController`] replaces the proxy on the B = 1 round loop +//! with the direct signal: emitted tokens per millisecond of round time, +//! which is the quantity widening exists to improve and which the round +//! loop already pays to know. It alternates measurement windows between the +//! configured depth and the requested ceiling, adopts whichever measures +//! faster, and re-challenges the loser on a backoff schedule. The search +//! space is deliberately the same two widths the proxy chose between: +//! `--draft-block-size` stays a ceiling the user set, not a hint a walk +//! wanders away from. /// Minimum number of completed MTP rounds before expanding above the /// drafter's configured block size. @@ -68,9 +87,317 @@ pub(crate) fn effective_mtp_block_size( } } +/// Rounds run at the configured depth before the first measurement window +/// opens. Mirrors [`MIN_HISTORY_FOR_EXPANSION`]: the first rounds of a +/// session carry prefill warm-up and kernel compilation, which would bias +/// whichever arm measured first. +const WARMUP_ROUNDS: usize = 8; +/// Rounds per measurement window. At the ~30 ms rounds of the measured +/// Gemma 4 pairing a window is about one second; acceptance variance makes +/// a single window a noisy estimate of a 5% difference, which the margin, +/// the re-challenge schedule, and the bounded two-arm search space are +/// sized to tolerate (a wrong adoption costs the measured 5%, is bounded by +/// the user's own ceiling, and is revisited). +const WINDOW_ROUNDS: usize = 32; +/// Relative lead the challenging arm must measure over the held arm's most +/// recent window to be adopted. Below this the tie goes to the held arm, +/// so measurement noise does not flap the width round-to-round. +const ADOPT_MARGIN: f64 = 0.02; +/// A challenge window aborts once it has this many rounds and trails the +/// held arm by more than [`EARLY_ABORT_DEFICIT`]. This is the Qwen 3.8 +/// guard: at a requested width of 12 that pairing measures 5.80 against +/// 21.30 tok/s (issue #1207), and the deficit is visible within a few +/// rounds, so the collapsed arm is charged a handful of rounds rather +/// than a full window. +const EARLY_ABORT_MIN_ROUNDS: usize = 4; +const EARLY_ABORT_DEFICIT: f64 = 0.35; +/// Re-challenge schedule for the losing arm, in windows of the held arm: +/// starts at the base, multiplies by 4 per consecutive loss, and caps. A +/// consistently losing arm ends up probed for at most a few rounds per +/// couple of thousand, which prices the Qwen-style collapsed ceiling at +/// well under a percent of steady-state throughput. +const RECHALLENGE_BASE_WINDOWS: usize = 4; +const RECHALLENGE_GROWTH: usize = 4; +const RECHALLENGE_CAP_WINDOWS: usize = 64; + +/// Which of the controller's two arms a round is charged to. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Arm { + Configured, + Requested, +} + +/// Two-arm throughput comparator for the B = 1 MTP verify width +/// (issue #1207). +/// +/// Owned by the generator rather than the session, so a server process +/// keeps its evidence across requests; a CLI run is one session either +/// way. All state advances in [`Self::record_round`]; [`Self::decide`] +/// only reads the current arm and applies the remaining-budget cap. +/// +/// Exactness scope: the widths this controller can pick are the two the +/// proxy controller already picked between, both bounded by the requested +/// width the exactness probe ran at, so it introduces no verify shape the +/// gate has not covered. +#[derive(Debug)] +pub(crate) struct BlockThroughputController { + configured: usize, + requested: usize, + /// Arm the controller has settled on outside challenge windows. + held: Arm, + /// Arm the in-progress window is charged to. Differs from `held` + /// exactly during a challenge. + measuring: Arm, + /// Most recent completed-window rate (emitted tokens per ms) per arm. + rate_configured: Option, + rate_requested: Option, + window_rounds: usize, + window_emitted: usize, + window_ms: f64, + warmup_rounds_left: usize, + /// Held-arm windows to complete before the losing arm is re-tried. + windows_until_challenge: usize, + /// Current re-challenge delay, in windows. Grows per consecutive loss. + challenge_backoff: usize, +} + +impl BlockThroughputController { + pub(crate) fn new(requested: usize, configured: usize) -> Self { + Self { + configured, + requested, + held: Arm::Configured, + measuring: Arm::Configured, + rate_configured: None, + rate_requested: None, + window_rounds: 0, + window_emitted: 0, + window_ms: 0.0, + warmup_rounds_left: WARMUP_ROUNDS, + // The first challenge follows the first completed configured + // window, so the requested width is measured once per session + // (or process) even when it never wins. + windows_until_challenge: 1, + challenge_backoff: RECHALLENGE_BASE_WINDOWS, + } + } + + /// Whether the two arms actually differ. When they do not (the request + /// is at or below the configured depth), the controller is inert and + /// [`Self::decide`] reproduces the proxy controller's early return. + fn active(&self) -> bool { + self.requested > self.configured && self.configured > 1 + } + + fn arm_width(&self, arm: Arm) -> usize { + match arm { + Arm::Configured => self.configured, + Arm::Requested => self.requested, + } + } + + fn rate_of(&self, arm: Arm) -> Option { + match arm { + Arm::Configured => self.rate_configured, + Arm::Requested => self.rate_requested, + } + } + + fn set_rate(&mut self, arm: Arm, rate: f64) { + match arm { + Arm::Configured => self.rate_configured = Some(rate), + Arm::Requested => self.rate_requested = Some(rate), + } + } + + fn other(arm: Arm) -> Arm { + match arm { + Arm::Configured => Arm::Requested, + Arm::Requested => Arm::Configured, + } + } + + /// The verify width the next round should use, bounded by the + /// remaining emission budget (`remaining_budget` includes the prefix + /// bonus position, exactly as [`effective_mtp_block_size`] takes it). + pub(crate) fn decide(&self, remaining_budget: usize) -> usize { + let ceiling = self.requested.min(remaining_budget); + if !self.active() { + return ceiling; + } + if self.warmup_rounds_left > 0 { + return self.configured.min(ceiling); + } + self.arm_width(self.measuring).min(ceiling) + } + + /// Feed one completed speculative round. `width` is the block the round + /// actually verified: rounds the budget forced below the measuring + /// arm's width are ignored rather than mis-charged. + pub(crate) fn record_round(&mut self, width: usize, emitted: usize, round_ms: f64) { + if !self.active() || round_ms <= 0.0 { + return; + } + if self.warmup_rounds_left > 0 { + self.warmup_rounds_left -= 1; + return; + } + if width != self.arm_width(self.measuring) { + return; + } + self.window_rounds += 1; + self.window_emitted += emitted; + self.window_ms += round_ms; + + let in_challenge = self.measuring != self.held; + if in_challenge + && self.window_rounds >= EARLY_ABORT_MIN_ROUNDS + && let Some(held_rate) = self.rate_of(self.held) + { + let partial = self.window_emitted as f64 / self.window_ms; + if partial < held_rate * (1.0 - EARLY_ABORT_DEFICIT) { + // The challenger is collapsing; close its window now so a + // Qwen-12-style arm costs rounds, not a full window. + self.finish_window(); + return; + } + } + if self.window_rounds >= WINDOW_ROUNDS { + self.finish_window(); + } + } + + fn finish_window(&mut self) { + let rate = self.window_emitted as f64 / self.window_ms; + let arm = self.measuring; + self.set_rate(arm, rate); + self.window_rounds = 0; + self.window_emitted = 0; + self.window_ms = 0.0; + + if arm != self.held { + // A challenge window closed: adopt on a clear lead, otherwise + // return to the held arm and back the loser off. + let held_rate = self.rate_of(self.held); + let adopted = match held_rate { + Some(h) => rate > h * (1.0 + ADOPT_MARGIN), + None => true, + }; + if adopted { + self.held = arm; + self.challenge_backoff = RECHALLENGE_BASE_WINDOWS; + } else { + self.challenge_backoff = + (self.challenge_backoff * RECHALLENGE_GROWTH).min(RECHALLENGE_CAP_WINDOWS); + } + self.measuring = self.held; + self.windows_until_challenge = self.challenge_backoff; + } else { + // A held-arm window closed: refresh its rate and count down to + // the next challenge. + self.windows_until_challenge = self.windows_until_challenge.saturating_sub(1); + if self.windows_until_challenge == 0 { + self.measuring = Self::other(self.held); + } + } + } +} + #[cfg(test)] mod tests { use super::effective_mtp_block_size; + use super::{Arm, BlockThroughputController, WARMUP_ROUNDS, WINDOW_ROUNDS}; + + /// Drive `rounds` rounds through the controller, asking `decide` first + /// (as the round loop does) and answering with the arm's synthetic + /// (emitted, ms) profile. Returns how many rounds ran at each width. + fn drive( + c: &mut BlockThroughputController, + rounds: usize, + budget: usize, + profile: impl Fn(usize) -> (usize, f64), + ) -> std::collections::HashMap { + let mut widths = std::collections::HashMap::new(); + for _ in 0..rounds { + let w = c.decide(budget); + *widths.entry(w).or_insert(0) += 1; + let (emitted, ms) = profile(w); + c.record_round(w, emitted, ms); + } + widths + } + + /// The measured Gemma 4 12B pairing (issue #1207): configured 4 at + /// 93.5 tok/s, requested 5 at 98.2. The controller must adopt 5. + #[test] + fn adopts_the_requested_width_when_it_measures_faster() { + let mut c = BlockThroughputController::new(5, 4); + let profile = |w: usize| match w { + 4 => (2705, 29160.0), // 2.705 emitted / 29.16 ms, scaled x1000 + 5 => (3132, 32170.0), + other => panic!("unexpected width {other}"), + }; + drive(&mut c, WARMUP_ROUNDS + 3 * WINDOW_ROUNDS, 1 << 20, profile); + assert_eq!(c.held, Arm::Requested); + assert_eq!(c.decide(1 << 20), 5); + } + + /// The measured Qwen 3.8 pairing: configured 3 at 21.30 tok/s, + /// requested 12 at 5.80. The challenge must abort early and the + /// controller must keep refusing to widen. + #[test] + fn refuses_a_requested_width_that_collapses() { + let mut c = BlockThroughputController::new(12, 3); + let profile = |w: usize| match w { + 3 => (2, 94.0), // ~21.3 tok/s at ~2 emitted per round + 12 => (2, 345.0), // ~5.8 tok/s + other => panic!("unexpected width {other}"), + }; + let widths = drive(&mut c, WARMUP_ROUNDS + 20 * WINDOW_ROUNDS, 1 << 20, profile); + assert_eq!(c.held, Arm::Configured); + assert_eq!(c.decide(1 << 20), 3); + // The collapsed arm was charged rounds, but only a few per + // backoff cycle: bound it well below one full window per cycle. + let wide_rounds = widths.get(&12).copied().unwrap_or(0); + assert!( + wide_rounds > 0 && wide_rounds < 3 * super::EARLY_ABORT_MIN_ROUNDS + 2, + "collapsed arm charged {wide_rounds} rounds" + ); + } + + /// A small lead inside the adoption margin does not flap the width. + #[test] + fn a_lead_inside_the_margin_stays_configured() { + let mut c = BlockThroughputController::new(5, 4); + let profile = |w: usize| match w { + 4 => (1000, 10000.0), + 5 => (1010, 10000.0), // +1.0%, inside the 2% margin + other => panic!("unexpected width {other}"), + }; + drive(&mut c, WARMUP_ROUNDS + 6 * WINDOW_ROUNDS, 1 << 20, profile); + assert_eq!(c.held, Arm::Configured); + } + + /// Budget caps every width, and capped rounds are not charged to the + /// measuring arm's window. + #[test] + fn budget_caps_the_width_and_capped_rounds_are_ignored() { + let mut c = BlockThroughputController::new(5, 4); + assert_eq!(c.decide(3), 3); + let before = c.window_rounds; + c.record_round(3, 3, 30.0); + assert_eq!(c.window_rounds, before); + } + + /// Inert when the request is not above the configured depth: behaves + /// like the proxy controller's early return. + #[test] + fn inert_at_or_below_the_configured_depth() { + let c = BlockThroughputController::new(4, 4); + assert_eq!(c.decide(1 << 20), 4); + let c = BlockThroughputController::new(3, 4); + assert_eq!(c.decide(1 << 20), 3); + } #[test] fn stays_at_requested_when_not_above_configured() { diff --git a/src/lib/mlxcel-core/src/speculative/mtp/generator.rs b/src/lib/mlxcel-core/src/speculative/mtp/generator.rs index 7306a74c7..de6cdaf3d 100644 --- a/src/lib/mlxcel-core/src/speculative/mtp/generator.rs +++ b/src/lib/mlxcel-core/src/speculative/mtp/generator.rs @@ -46,7 +46,7 @@ use crate::sampling::{LogprobsConfig, TokenLogprobData}; use std::sync::atomic::{AtomicBool, Ordering}; use std::time::{Duration, Instant}; -use super::adaptive::effective_mtp_block_size; +use super::adaptive::{BlockThroughputController, effective_mtp_block_size}; use super::target::{MtpTarget, MtpVerifyOutput}; use super::tree::DraftTree; use super::walk::{WalkResult, speculative_walk}; @@ -354,8 +354,10 @@ pub struct MtpSessionState { /// not keep the token stream itself: each step hands its new tokens to /// the caller, who owns accumulation/streaming. emitted_count: usize, - /// Per-round accepted counts feeding the adaptive block-size controller - /// ([`effective_mtp_block_size`]). Probe rounds are excluded. + /// Per-round accepted counts feeding the acceptance-proxy block + /// controller ([`effective_mtp_block_size`]), which only decides when + /// `MLXCEL_MTP_BLOCK_CONTROLLER=proxy` pins it (issue #1207). Probe + /// rounds are excluded. accept_lens: Vec, /// Whether the narrower-than-requested block warning has already fired. /// Once per session: the condition holds every round it holds at all, and @@ -497,6 +499,24 @@ pub struct MtpGenerator { /// and [`MtpTarget::tree_round_is_available`] is the per-round capability /// question asked on top of it (issue #1204). tree_drafting: bool, + /// Measured-throughput block-width comparator (issue #1207). Owned by + /// the generator rather than the session so a server process keeps its + /// evidence across requests. Consulted only when + /// `prefer_requested_block_size` is false and the proxy controller is + /// not pinned via `MLXCEL_MTP_BLOCK_CONTROLLER=proxy`. + block_controller: BlockThroughputController, + /// `MLXCEL_MTP_BLOCK_CONTROLLER=proxy`: keep the upstream + /// fully-accepted-prefix gate instead of the throughput comparator. + proxy_block_controller: bool, +} + +/// Whether `MLXCEL_MTP_BLOCK_CONTROLLER` pins the upstream acceptance-proxy +/// block controller. Any other value (or unset) selects the measured +/// throughput comparator, the default since issue #1207. +fn proxy_block_controller_pinned() -> bool { + std::env::var("MLXCEL_MTP_BLOCK_CONTROLLER") + .map(|value| value.trim().eq_ignore_ascii_case("proxy")) + .unwrap_or(false) } /// Whether tree drafting is switched on for this process. @@ -539,6 +559,8 @@ impl MtpGenerator { last_acceptance: None, profile_probe_rounds: 0, tree_drafting: tree_drafting_enabled(), + block_controller: BlockThroughputController::new(block_size, configured_block_size), + proxy_block_controller: proxy_block_controller_pinned(), } } @@ -1002,28 +1024,39 @@ impl MtpGenerator { ); } } + // (width, start) of a round the throughput comparator will be fed at + // the end of this step. `None` for probe rounds, for the proxy / + // prefer-requested paths, and for rounds that bail out early (a + // drafter failure or rollback refusal spends time this round, but + // charging a terminal round to either arm would bias the window it + // lands in). + let mut throughput_round: Option<(usize, Instant)> = None; let (draft_tokens, draft_tree): (Vec, Option) = if is_probe { state.probes_remaining -= 1; (Vec::new(), None) } else { // Bound the block size by the remaining budget. When the // operator requested a block larger than the drafter's - // configured depth, mirror upstream's adaptive controller: - // stay at configured depth until recent acceptance proves the - // configured prefix is usually fully accepted, then expand to - // the requested ceiling. The `+1` is because the verify input - // is `[bonus, draft_0, …, draft_{K-2}]`: one prefix bonus - // position that the round-loop already counts as emitted. + // configured depth, the throughput comparator alternates + // measurement windows between the two widths and holds + // whichever measures more emitted tokens per millisecond + // (issue #1207); `MLXCEL_MTP_BLOCK_CONTROLLER=proxy` restores + // upstream's fully-accepted-prefix gate. The `+1` is because + // the verify input is `[bonus, draft_0, …, draft_{K-2}]`: one + // prefix bonus position that the round-loop already counts as + // emitted. let remaining = state.max_tokens - state.emitted_count + 1; let bs = if self.prefer_requested_block_size { self.block_size.min(remaining) - } else { + } else if self.proxy_block_controller { effective_mtp_block_size( self.block_size, self.configured_block_size, &state.accept_lens, remaining, ) + } else { + self.block_controller.decide(remaining) }; if bs <= 1 { state.finished = true; @@ -1035,18 +1068,37 @@ impl MtpGenerator { // A request silently reduced to the configured depth is the case // that makes `--draft-block-size` look tunable when it is not // (issue #1206); a reduction the remaining budget forced is - // ordinary end-of-generation behavior and stays quiet. + // ordinary end-of-generation behavior and stays quiet. Under the + // throughput comparator the narrower rounds are measurement, not + // a verdict, so the message names the mechanism that would + // expand it. if bs < self.block_size && self.block_size <= remaining && !state.warned_block_override { state.warned_block_override = true; - tracing::warn!( - requested = self.block_size, - effective = bs, - configured = self.configured_block_size, - "MTP drafted at a narrower block than requested: the adaptive controller \ - holds this drafter at its configured depth until recent acceptance shows \ - that prefix is usually fully accepted" - ); + if self.proxy_block_controller { + tracing::warn!( + requested = self.block_size, + effective = bs, + configured = self.configured_block_size, + "MTP drafted at a narrower block than requested: the proxy controller \ + holds this drafter at its configured depth until recent acceptance shows \ + that prefix is usually fully accepted (MLXCEL_MTP_BLOCK_CONTROLLER=proxy)" + ); + } else { + tracing::warn!( + requested = self.block_size, + effective = bs, + configured = self.configured_block_size, + "MTP drafted at a narrower block than requested: the throughput \ + comparator holds the configured depth until the requested width \ + measures more emitted tokens per millisecond over a recent window \ + (issue #1207)" + ); + } + } + + if !self.prefer_requested_block_size && !self.proxy_block_controller { + throughput_round = Some((bs, Instant::now())); } // Arm the drafter's shared K/V from the stored verify output @@ -1322,6 +1374,20 @@ impl MtpGenerator { } } + // Feed the throughput comparator with what this round actually + // delivered: everything from the shared-K/V arm through the accept + // hook, against everything it emitted. Per-round attribution of + // lazily-evaluated GPU work is approximate (a round can execute a + // predecessor's deferred tail), which is why the comparator only + // ever reads window sums. + if let Some((bs, started)) = throughput_round { + self.block_controller.record_round( + bs, + new_tokens.len(), + duration_ms(started.elapsed()), + ); + } + // Next round's bonus is the last emitted token. `walk.new_tokens` // always carries at least the target's own token at position 0 // (the budget is >= 1 at the top of the round), so the fallback diff --git a/src/lib/mlxcel-core/src/speculative/mtp/round_loop_batched.rs b/src/lib/mlxcel-core/src/speculative/mtp/round_loop_batched.rs index 7233ff6f9..7ce5c23e9 100644 --- a/src/lib/mlxcel-core/src/speculative/mtp/round_loop_batched.rs +++ b/src/lib/mlxcel-core/src/speculative/mtp/round_loop_batched.rs @@ -311,6 +311,12 @@ impl MtpBatchedGenerator { let bs = if self.prefer_requested_block_size { self.block_size.min(remaining_min) } else { + // The batched loop stays on the acceptance-proxy controller + // while the B = 1 loop moved to the measured-throughput + // comparator (issue #1207): the row-averaged accept length is + // the only per-round signal this loop measures today, its + // per-round wall time mixes rows at divergent depths, and the + // #1207 measurement that justified the switch is B = 1. effective_mtp_block_size( self.block_size, self.configured_block_size, From 2c1da8a4070d5297f1a090d0e1abeaa12d823f3d Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sat, 22 Aug 2026 14:29:52 +0900 Subject: [PATCH 2/2] feat(speculative): add a width-sweep measurement mode and a challenge verdict log MLXCEL_MTP_BLOCK_CONTROLLER grows a third value: 'requested' honours the requested width from the first round, controller-free. #1207's own width sweep needed a temporary prefer_requested_block_size code patch to hold the widths it compared; this value is that patch as a switch, and it is what the ground-truth measurements in the PR were taken with. Not a deployment setting: the harmful-ceiling protection is off under it. The throughput comparator also logs one INFO line per closed challenge window (challenger and held widths and their measured tokens-per-ms rates, and whether the challenger was adopted), so a sweep can see which width a run settled on and why without an instrumented build. Verdicts are rare, so the per-round path is untouched. Refs #1207 --- docs/environment-variables.md | 2 +- .../src/speculative/mtp/adaptive.rs | 12 ++++ .../src/speculative/mtp/generator.rs | 62 ++++++++++++------- 3 files changed, 54 insertions(+), 22 deletions(-) diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 9a172ade1..e07cea435 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -188,7 +188,7 @@ The OpenAI audio endpoints (`/v1/audio/speech`, `/v1/audio/transcriptions`, `/v1 | `MLXCEL_GDN_CHAIN_PARITY` | `0` to disable, any other value (or unset) to enable | on | **Advanced, diagnostic escape hatch.** Gates the chain-parity gated-delta Metal kernel used by Qwen 3.5 MTP's speculative verify and rollback-replay paths (issue #1165). The standard gated-delta kernel carries float32 recurrent state across a `T = K` verify block and rounds it to the storage dtype only once at the end, while the classic single-token decode chain rounds after every token; a `T = K` verify block is therefore NOT bit-identical to `K` consecutive single-token decode steps unless the state is rounded per in-block step. The chain-parity kernel (`gated_delta_step_seqpar`) does that rounding, which is what makes Qwen 3.5 MTP's temperature-0 output byte-identical to classic decode. **Setting this to `0` forfeits that exactness contract**, restoring the pre-#1165 block numerics for A/B attribution of the parity kernel's own cost and acceptance effect; do not set it to `0` in a deployment that needs byte-identical speculative output. Metal-only: the non-Metal ops fallback ignores the flag (the parity guarantee does not exist off Metal today). **The kernel is necessary but not sufficient**: byte-identity also requires every quantized projection to dispatch to the same MLX kernel at `M = block_size` as at `M = 1`, which is not true on every GPU generation or at every block width, so the runtime probe behind `MLXCEL_MTP_ALLOW_INEXACT` is what actually decides whether MTP engages. See `docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md` for the measured kernel cost (inside the dispatch-noise band). | | `MLXCEL_MTP_ALLOW_INEXACT` | `1`/`true`/`yes`/`on` to enable; unset or anything else to disable | **off** | Engage Qwen 3.5 MTP speculative decoding even when the startup exactness probe reports that the multi-token verify block is **not** byte-identical to the single-token decode chain. Before enabling MTP the runtime now measures the property instead of predicting it: one synthetic verify block and the equivalent single-token chain are run from the same state on the loaded checkpoint at the configured `--draft-block-size`, and their logits are compared byte for byte (three independent synthetic inputs, each two short prefills plus `K + 1` forwards; measured 4.9 s for the first call and 1.3 s for a later one per input on a Qwen3.8-27B 4-bit target on an M1 Ultra, the difference being MLX's one-time kernel compilation; more than one input because a kernel pair can disagree by only a byte or two out of ten thousand, at which amplitude a single draw can read as equal; memoized per (model, block width) and warmed at worker startup so it never lands on the request path). A divergence means temperature-0 speculative output would silently differ from `mlxcel generate` without `--draft-model`, so the default is to decline and run classic decode. The static conditions (Metal backend, `supports_metal_gated_delta_kernel` geometry) still apply and are checked first; this probe covers what they cannot, namely which MLX kernel each **quantized projection** dispatches to at `M = K` versus `M = 1`. That choice depends on the GPU generation, the quantization mode, the operand sizes and the block width: `use_qmv_wide` in [`mlx/backend/metal/quantized.cpp`](https://github.com/ml-explore/mlx/blob/main/mlx/backend/metal/quantized.cpp) sends `M >= 2` to a different reduction whenever `mode != "affine" || arch_gen >= 15`, and `get_qmv_batch_limit` sends `M` above 10, 12, 18 or 32 (by architecture size and generation) to the matrix-matrix kernel. Measured: an affine 4-bit Qwen3.8-27B target on an M1 Ultra is byte-identical at block widths 1 through 11 and diverges at 12 (the `arch_size == 'd'` branch); the same checkpoint on an M5 Max diverges from block width 2, where the `M >= 2` split fires before any batch limit can be observed. Within one checkpoint the limit is per projection, not per model: Gemma 4 12B's attention shapes hold to 17 on an M1 Ultra while its MLP shapes break at 12, so a model's own cliff is the minimum over its shapes, which is why this is measured rather than tabulated. Note the ordering that #1199 introduced: on a failing probe the gate first retries with `qmv_wide` disabled and keeps the narrow kernel when that restores exactness, and only a probe that fails **both** ways consults this flag. On Apple GPU generation 15+ the narrow retry passes, so this flag alone is inert there: the process is pinned narrow, output stays byte-identical, and the log shows the retry's INFO line rather than the override warning (verified live on M3 Ultra, 2026-08-22, byte-identical output with and without the flag; see `benchmark_results/qmv-wide-pin-tax-m3ultra-2026-08-22.md`). To research the fast kernel there, set `MLXCEL_QMV_WIDE=1` together with this flag: the pin skips the retry and this flag then engages MTP on the wide kernel, forfeiting byte-identity with the loud WARN. This flag alone is load-bearing only where no exact kernel selection exists at the configured block width. Read once per process. | | `MLXCEL_QMV_WIDE` | `0`/`false`/`no`/`off` to disable; `1` (or any other value) to pin wide | unset (wide, until the MTP gate's retry turns it off) | Operator pin for MLX's `qmv_wide` kernel, the faster reduction for `M >= 2` quantized matmuls on Apple GPU generation 15+ (overlay in `src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp`, added by #1199). **Setting the variable at all, to any value, counts as an operator pin**: the MTP exactness gate's retry (`retry_without_qmv_wide`) is skipped in both directions, so `MLXCEL_QMV_WIDE=1` keeps the wide kernel and makes a failing probe decline MTP instead of buying exactness back, and `MLXCEL_QMV_WIDE=0` runs the whole process narrow from the start. Unset, the kernel is wide until a failing MTP probe's retry finds the narrow kernel exact and pins the process narrow for good. The pin is process-wide and sits on the dispatch path of every quantized matmul; what non-MTP work pays for the narrow state is measured in `benchmark_results/qmv-wide-pin-tax-m3ultra-2026-08-22.md` (nothing measurable on batched decode, about 15 ms per prompt-cache-hit request's suffix prefill). Read once per process at first dispatch; `mlxcel_core::set_qmv_wide` can move it at runtime and the gate is its only caller. | -| `MLXCEL_MTP_BLOCK_CONTROLLER` | `proxy` to pin the acceptance-proxy gate; any other value (or unset) selects the throughput comparator | throughput | Which controller decides the B=1 MTP verify width when `--draft-block-size` exceeds the drafter's configured depth (issue #1207). The default measures the decision: after a short warm-up the round loop alternates measurement windows (32 rounds) between the configured depth and the requested ceiling, compares emitted tokens per millisecond of round time, holds whichever measures faster (a challenger needs a 2% lead; ties go to the configured depth), and re-challenges the loser on a growing backoff (4, 16, then every 64 windows), with a collapsing challenger aborted after 4 rounds once it trails by more than 35% so a harmful ceiling (the Qwen 3.8 pairing measures 5.80 against 21.30 tok/s at width 12) costs rounds rather than windows. Evidence lives in the generator, so a server process keeps it across requests; drafters that set `prefer_requested_block_size` (Qwen 3.5 MTP) bypass both controllers and always honour the request, exactly as before. Set to `proxy` to restore upstream's fully-accepted-prefix gate, which issue #1207 measured holding the Gemma 4 12B pairing about 5% below its optimum (93.54 against 98.16 tok/s at requested width 5 on an M5 Max) because at 0.585 acceptance the configured prefix is rarely fully accepted no matter how profitable widening is. The batched (B>1) loop stays on the proxy gate regardless: the row-averaged accept length is the only per-round signal it measures today. Read once per generator. | +| `MLXCEL_MTP_BLOCK_CONTROLLER` | `proxy` to pin the acceptance-proxy gate; `requested` to honour the requested width immediately (measurement mode); any other value (or unset) selects the throughput comparator | throughput | Which controller decides the B=1 MTP verify width when `--draft-block-size` exceeds the drafter's configured depth (issue #1207). The default measures the decision: after a short warm-up the round loop alternates measurement windows (32 rounds) between the configured depth and the requested ceiling, compares emitted tokens per millisecond of round time, holds whichever measures faster (a challenger needs a 2% lead; ties go to the configured depth), and re-challenges the loser on a growing backoff (4, 16, then every 64 windows), with a collapsing challenger aborted after 4 rounds once it trails by more than 35% so a harmful ceiling (the Qwen 3.8 pairing measures 5.80 against 21.30 tok/s at width 12) costs rounds rather than windows. Evidence lives in the generator, so a server process keeps it across requests; drafters that set `prefer_requested_block_size` (Qwen 3.5 MTP) bypass both controllers and always honour the request, exactly as before. Set to `proxy` to restore upstream's fully-accepted-prefix gate, which issue #1207 measured holding the Gemma 4 12B pairing about 5% below its optimum (93.54 against 98.16 tok/s at requested width 5 on an M5 Max) because at 0.585 acceptance the configured prefix is rarely fully accepted no matter how profitable widening is. The batched (B>1) loop stays on the proxy gate regardless: the row-averaged accept length is the only per-round signal it measures today. Set to `requested` to bypass both controllers and draft at the requested width from the first round: this is the width-sweep measurement mode (#1207's own sweep needed a temporary code patch to hold widths; this value is that patch as a switch), not a deployment setting, and the harmful-ceiling protection is off under it. Read once per generator. | | `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). | diff --git a/src/lib/mlxcel-core/src/speculative/mtp/adaptive.rs b/src/lib/mlxcel-core/src/speculative/mtp/adaptive.rs index a0dba0c18..a8d96ed30 100644 --- a/src/lib/mlxcel-core/src/speculative/mtp/adaptive.rs +++ b/src/lib/mlxcel-core/src/speculative/mtp/adaptive.rs @@ -283,6 +283,18 @@ impl BlockThroughputController { Some(h) => rate > h * (1.0 + ADOPT_MARGIN), None => true, }; + // One line per verdict so a sweep can see which width a run + // settled on and why, without instrumented builds. Verdicts are + // rare (once per challenge window at the least), so this stays + // out of the per-round path. + tracing::info!( + challenger_width = self.arm_width(arm), + held_width = self.arm_width(self.held), + challenger_rate = format!("{rate:.5}"), + held_rate = held_rate.map(|h| format!("{h:.5}")).unwrap_or_default(), + adopted, + "MTP block-width challenge closed (tokens per ms, issue #1207)" + ); if adopted { self.held = arm; self.challenge_backoff = RECHALLENGE_BASE_WINDOWS; diff --git a/src/lib/mlxcel-core/src/speculative/mtp/generator.rs b/src/lib/mlxcel-core/src/speculative/mtp/generator.rs index de6cdaf3d..a798c1163 100644 --- a/src/lib/mlxcel-core/src/speculative/mtp/generator.rs +++ b/src/lib/mlxcel-core/src/speculative/mtp/generator.rs @@ -505,18 +505,34 @@ pub struct MtpGenerator { /// `prefer_requested_block_size` is false and the proxy controller is /// not pinned via `MLXCEL_MTP_BLOCK_CONTROLLER=proxy`. block_controller: BlockThroughputController, - /// `MLXCEL_MTP_BLOCK_CONTROLLER=proxy`: keep the upstream - /// fully-accepted-prefix gate instead of the throughput comparator. - proxy_block_controller: bool, + /// Which block-width controller `MLXCEL_MTP_BLOCK_CONTROLLER` selected. + block_controller_mode: BlockControllerMode, } -/// Whether `MLXCEL_MTP_BLOCK_CONTROLLER` pins the upstream acceptance-proxy -/// block controller. Any other value (or unset) selects the measured -/// throughput comparator, the default since issue #1207. -fn proxy_block_controller_pinned() -> bool { - std::env::var("MLXCEL_MTP_BLOCK_CONTROLLER") - .map(|value| value.trim().eq_ignore_ascii_case("proxy")) - .unwrap_or(false) +/// `MLXCEL_MTP_BLOCK_CONTROLLER` values. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BlockControllerMode { + /// Default: the measured-throughput comparator (issue #1207). + Throughput, + /// The upstream fully-accepted-prefix gate. + Proxy, + /// Honour the requested width immediately, controller-free. The + /// measurement mode: #1207's own width sweep needed a temporary code + /// patch to hold the widths it compared, and this value is that patch + /// as a switch. Not a deployment setting. + Requested, +} + +/// Parse `MLXCEL_MTP_BLOCK_CONTROLLER`. Unknown values fall back to the +/// default rather than erroring, matching the other MLXCEL_* switches. +fn block_controller_mode() -> BlockControllerMode { + match std::env::var("MLXCEL_MTP_BLOCK_CONTROLLER") { + Ok(value) if value.trim().eq_ignore_ascii_case("proxy") => BlockControllerMode::Proxy, + Ok(value) if value.trim().eq_ignore_ascii_case("requested") => { + BlockControllerMode::Requested + } + _ => BlockControllerMode::Throughput, + } } /// Whether tree drafting is switched on for this process. @@ -560,7 +576,7 @@ impl MtpGenerator { profile_probe_rounds: 0, tree_drafting: tree_drafting_enabled(), block_controller: BlockThroughputController::new(block_size, configured_block_size), - proxy_block_controller: proxy_block_controller_pinned(), + block_controller_mode: block_controller_mode(), } } @@ -1048,15 +1064,17 @@ impl MtpGenerator { let remaining = state.max_tokens - state.emitted_count + 1; let bs = if self.prefer_requested_block_size { self.block_size.min(remaining) - } else if self.proxy_block_controller { - effective_mtp_block_size( - self.block_size, - self.configured_block_size, - &state.accept_lens, - remaining, - ) } else { - self.block_controller.decide(remaining) + match self.block_controller_mode { + BlockControllerMode::Requested => self.block_size.min(remaining), + BlockControllerMode::Proxy => effective_mtp_block_size( + self.block_size, + self.configured_block_size, + &state.accept_lens, + remaining, + ), + BlockControllerMode::Throughput => self.block_controller.decide(remaining), + } }; if bs <= 1 { state.finished = true; @@ -1075,7 +1093,7 @@ impl MtpGenerator { if bs < self.block_size && self.block_size <= remaining && !state.warned_block_override { state.warned_block_override = true; - if self.proxy_block_controller { + if self.block_controller_mode == BlockControllerMode::Proxy { tracing::warn!( requested = self.block_size, effective = bs, @@ -1097,7 +1115,9 @@ impl MtpGenerator { } } - if !self.prefer_requested_block_size && !self.proxy_block_controller { + if !self.prefer_requested_block_size + && self.block_controller_mode == BlockControllerMode::Throughput + { throughput_round = Some((bs, Instant::now())); }