From b416a403cb21551f03aad8b435cfab26aacd8696 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sat, 22 Aug 2026 15:10:56 +0900 Subject: [PATCH] perf(speculative): close the early-exit verify walk question as a measured no Issue #1179 asked three things: correct the stale comment claiming the deferred verify path projects one position at a time (it has projected the whole block in one graph since fcf9e205), decide the fate of MLXCEL_ENABLE_MTP_DEFERRED, and measure whether a device-side early-exit walk is worth building. Measured, and the walk is not built. The LM-head projection cost is not flat in the projected width, which cuts both ways: the early-exit ceiling t(K) - t(1) at production widths is 2 to 4% of the verify forward under the production narrow kernel selection, and a sequential walk cannot collect even that, because every walked position re-reads the full 566 to 715 MB weight matrix. At the measured operating points (K = 3 to 4, acceptance 0.80 to 0.88) the walk's expected cost is 1.2 to 1.6x the batched projection it replaces. The one width where it gets ahead (K = 8, ~1.9 ms) is not an operating point on either pairing and is collected cheaper by padding the projection past the qmv batch limit into the matrix-matrix kernel (2.56 ms against 4.51 at width 8), noted in the record for whoever revisits wide blocks. Full numbers in docs/benchmark_results/mtp-verify-early-exit-decision-m5max-2026-08-22.md, taken with examples/mtp_projection_width_bench (guarded per row by the projected logits shape) plus a production K sweep on both pairings. MLXCEL_ENABLE_MTP_DEFERRED is removed rather than renamed: it gated a path that split the verify forward from the LM-head projection without deferring any work (identical compute, one extra bridge crossing, one extra intermediate). Its orphaned helper chain goes with it (argmax_from_hidden_positions, forward_hidden_with_speculative_sinks, speculative_logits_from_hidden at both levels), the stale gate comment is replaced by the measured reason, and the two doc references to the flag are reworded. Byte identity re-verified after the removal: Gemma 4 MTP output equals classic decode under the same kernel selection. Closes #1179 --- ...fy-early-exit-decision-m5max-2026-08-22.md | 146 ++++++++++++ docs/environment-variables.md | 1 - examples/mtp_projection_width_bench.rs | 215 ++++++++++++++++++ .../src/drafter/dflash/round_loop.rs | 4 +- src/models/gated_delta.rs | 3 +- src/models/gemma4.rs | 87 +------ src/models/gemma4_mtp_target.rs | 139 +++-------- 7 files changed, 398 insertions(+), 197 deletions(-) create mode 100644 docs/benchmark_results/mtp-verify-early-exit-decision-m5max-2026-08-22.md create mode 100644 examples/mtp_projection_width_bench.rs diff --git a/docs/benchmark_results/mtp-verify-early-exit-decision-m5max-2026-08-22.md b/docs/benchmark_results/mtp-verify-early-exit-decision-m5max-2026-08-22.md new file mode 100644 index 000000000..e385aa82d --- /dev/null +++ b/docs/benchmark_results/mtp-verify-early-exit-decision-m5max-2026-08-22.md @@ -0,0 +1,146 @@ +# MTP verify early-exit walk: decided by measurement, not built (issue #1179) + +Issue #1179 asked whether a device-side early-exit verify walk beats the +batched full-logits verifier, with a stated prior: the LM head is +weight-read-bound, so skipping tail positions should save almost nothing. +The measurement below answers no, the walk should not be built, but for a +sharper reason than the prior: the projection cost is *not* flat in the +projected width, and the very effect that makes tail positions expensive +also makes a sequential walk pay the full weight read once per walked +position, which loses to the batched projection outright at every +production block size. + +## Environment + +| Field | Value | +|---|---| +| Host | Apple M5 Max, 128 GB unified memory, macOS 26.6.1 | +| Build | `cargo build --release --features metal,accelerate`, branch `perf/issue-1179-verify-early-exit-decision` (main at `77c71402` plus this branch's cleanup) | +| Harness (a) | `examples/mtp_projection_width_bench` 16 5: T=16 projections per batch folded into one eval, best of 5 rounds, weights round-robined over 2 copies (566 to 715 MB each) so every read hits DRAM | +| Harness (b) | offline CLI, 300 tokens, temperature 0, `MLXCEL_MTP_BLOCK_CONTROLLER=requested` (block width pinned, #1207), `MLXCEL_MTP_ALLOW_INEXACT=1` (widths above the narrow batch limit cannot pass the exactness probe; this sweep measures cost, not output), 15 s cooldowns, Time Machine stopped and verified stopped at both ends, indexers paused | +| Guard | each width row asserts and prints its logits shape `[1, W, vocab]`; the two arms of the question differ by that shape and nothing else | + +## (a) LM-head projection cost against projected width + +Projection + argmax at the real head shapes, affine 4-bit group 64. + +Default kernel selection (`qmv_wide` on, what a non-MTP process runs): + +| W | gemma4-12b tied head (262144 x 3840) | GB/s | qwen3.8-27b lm_head (248320 x 5120) | GB/s | +|---:|---:|---:|---:|---:| +| 1 | 0.985 ms | 575 | 1.234 ms | 580 | +| 2 | 1.020 ms | 555 | 1.301 ms | 550 | +| 3 | 1.149 ms | 493 | 1.419 ms | 504 | +| 4 | 1.352 ms | 419 | 1.903 ms | 376 | +| 5 | 1.846 ms | 307 | 2.246 ms | 318 | +| 8 | 3.406 ms | 166 | 4.299 ms | 166 | +| 16 | 2.561 ms | 221 | 3.266 ms | 219 | +| 32 | 2.452 ms | 231 | 3.111 ms | 230 | + +Narrow kernel selection (`MLXCEL_QMV_WIDE=0`, what the MTP verify path +actually runs on generation 15+ after the #1199 exactness retry): + +| W | gemma4-12b tied head | GB/s | qwen3.8-27b lm_head | GB/s | +|---:|---:|---:|---:|---:| +| 1 | 0.985 ms | 575 | 1.225 ms | 584 | +| 2 | 1.243 ms | 456 | 1.453 ms | 492 | +| 3 | 1.811 ms | 313 | 1.960 ms | 365 | +| 4 | 2.331 ms | 243 | 2.849 ms | 251 | +| 5 | 3.141 ms | 180 | 3.353 ms | 213 | +| 8 | 4.510 ms | 126 | 5.030 ms | 142 | +| 16 | 2.559 ms | 221 | 3.256 ms | 220 | +| 32 | 2.450 ms | 231 | 3.110 ms | 230 | + +W = 1 is identical in both selections (`M = 1` takes `qmv` regardless), +and so are W = 16 and 32 (the matrix-matrix kernel ignores the flag). The +`M` in between is where the narrow pin costs, which is the #1261/#1278 +collateral measured at the head shape. + +Two structural facts, visible in both selections: + +- The curve is not flat. W = 1 runs at the memory-bandwidth roof + (575 to 580 GB/s); the `M >= 2` kernels degrade steadily to about + 166 GB/s at W = 8; the matrix-matrix kernel takes over at the batch + limit and W = 16 and 32 cost *less* than W = 8. The issue's prior + (projection cost invariant in W, ceiling near zero) was wrong in the + letter but right in the verdict. +- The early-exit ceiling `t(K) - t(1)` at the widths anyone runs is + 0.74 to 0.83 ms per round at K = 3 and 1.35 to 1.62 ms at K = 4 under + the production (narrow) selection, against verify forwards of 31 to + 40 ms per round: about 2 to 4% of the verify forward, before paying + anything for the walk itself. And the walk cannot collect it, per the + next section. + +## Why the walk loses even that ceiling + +A device-side early-exit walk projects position by position and stops at +the first mismatch, so it pays `E[A] x t(1)` where `E[A]` is the expected +number of positions walked, and every one of those positions re-reads the +full weight matrix. At the measured production operating points: + +`E[A]` below is `E[accepted] + 1` under a per-position independence +approximation of the measured acceptance rate, capped at K. `t(K)` is the +narrow-selection column, because that is what the MTP verify path runs. + +| pairing | K | acceptance | E[A] | walk cost `E[A] x t(1)` | batched `t(K)` | walk wins? | +|---|---:|---:|---:|---:|---:|---| +| gemma4-12b | 3 | 0.876 | 2.64 | 2.60 ms | 1.81 ms | no, loses 1.4x | +| gemma4-12b | 4 | 0.799 | 2.95 | 2.91 ms | 2.33 ms | no, loses 1.2x | +| gemma4-12b | 8 | 0.636 | 2.67 | 2.63 ms | 4.51 ms | ~1.9 ms, see below | +| qwen3.8-27b | 3 | 0.827 | 2.51 | 3.10 ms | 1.96 ms | no, loses 1.6x | + +The one width where the walk gets ahead, K = 8, is dominated by a far +simpler lever the width curve exposes: padding the projection to the +matrix-matrix width costs 2.56 ms there, which collects the same ~1.9 ms +with a reshape instead of a sequential dispatch-compare-branch walk, and +K = 8 is not an operating point on either pairing anyway (see below). The +walk estimate also charges nothing for its own K sequential +dispatch/compare steps. + +The one cell where the walk is marginally ahead (K = 8, ~0.7 ms against a +57 ms verify forward, ~1.2%) is not an operating point: the production +sweep below puts Gemma 4's throughput optimum at K = 4 and the #1207 +controller now finds that optimum by measurement, and the estimate charges +the walk nothing for its K sequential dispatch/compare steps, which a real +implementation would pay. + +## (b) Production verify cost against block size + +Same protocol as (a)'s harness row. `verify_forward_ms / rounds`: + +| K | gemma verify ms/round | gemma tok/s | qwen verify ms/round (relative) | qwen tok/s | +|---:|---:|---:|---:|---:| +| 3 | 31.2 | 83.1 | 107.9 | 22.6 | +| 4 | 35.5 | 89.1 | 136.6 | 22.6 | +| 8 | 57.0 | 85.9 | 255.5 | 15.7 | +| 16 | 61.1 | 75.2 | 301.4 | 12.5 | +| 32 | 155.4 | 28.6 | 377.5 | 8.9 | + +The Gemma column is consistent with this host's clean baselines (89.15 +tok/s pinned at K = 4 in the #1207 record). The Qwen block ran late in the +sweep and its absolute values carry sustained-load throttling (this host's +known failure mode; its clean K = 3 verify round is ~40 ms), so read only +its ordering: K = 3 and 4 are equivalent, everything wider loses. On both +pairings the operating range stays at K <= 8, where (a) says early exit +has nothing to win. + +An observation for whoever revisits wide blocks: the W = 8 anomaly (more +expensive than W = 16) means a K = 8 verify would project cheaper padded +to the matrix-matrix width than early-exited. If wide blocks ever become +an operating point, pad-to-qmm is the lever to measure first, not the +walk. + +## Decision + +Not built, closed as a measured no (the outcome #1179 names as success): + +- At the production widths (K = 3 to 5) the batched projection beats a + sequential early-exit walk by 1.2 to 1.6x on the walk's own best case, + because the walk re-reads the weight per walked position. +- The ceiling the walk chases there is 2 to 4% of the verify forward, + and where a wide block would make it larger, pad-to-qmm collects it + cheaper. +- The stale gate comment and the `MLXCEL_ENABLE_MTP_DEFERRED` flag, which + gated a path that split the verify forward from the projection without + deferring any work (identical compute, one extra bridge crossing), were + removed with their orphaned helper chain. diff --git a/docs/environment-variables.md b/docs/environment-variables.md index e07cea435..f3fca7421 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -183,7 +183,6 @@ The OpenAI audio endpoints (`/v1/audio/speech`, `/v1/audio/transcriptions`, `/v1 | `MLXCEL_MTP_ADAPTIVE` | `0`/`false`/`no`/`off` to disable, any other value (or unset) to enable | on | Adaptive B=1 MTP policy (issue #333). When on, the server profiles the first few B=1 MTP bursts of each (target, drafter, hardware, block_size) pairing (acceptance length, verify latency, drafter latency, batch size, prompt shape) and settles to a data-driven enable/decline verdict that overrides the static per-hardware gate when the measured profile is clearly favorable or unfavorable, falling back to the static default otherwise. The verdict (enable/decline plus the coarse acceptance rate, no prompt data) is persisted at `${MLXCEL_CACHE_DIR:-$HOME/.cache/mlxcel}/mtp-policy/.json` (hint format v3; v2 hints were settled against the pre-#725 verify kernel and the pre-#736 estimator, so they are ignored and the pairing re-profiles once), so profiling runs once per pairing and a restart reuses the verdict. Changing `MLXCEL_DRAFT_BLOCK_SIZE` changes the block_size dimension of the key, so the old hint is discarded and profiling restarts for the new K. Set to an off value to disable profiling and use the pre-#333 static per-hardware gates. `MLXCEL_ENABLE_MTP_B1` still pins the decision and, when set, suppresses profiling. The experimental batched (B>1) path is unaffected and stays behind `MLXCEL_ENABLE_MTP_BATCH`. The speedup estimate is measured, not modeled (issue #736): while profiling, each burst runs a couple of classic-step probe rounds (drafterless rounds whose `[1, 1]` verify forward is shape-identical to a classic decode step; each emits one real greedy token, so nothing is wasted and temperature-0 output stays byte-identical), and the estimator compares the measured speculative round cost (verify + drafter + walk/finalize overhead) against the measured classic step time, which is taken as the median (not the mean) of the per-burst probe means so the first burst's one-time CUDA kernel/graph compilation for the `[1, 1]` verify shape cannot skew the estimate toward a falsely slow classic step. This makes the verdict correct across backends and kernel eras without hardware heuristics: on GB10 with the multirow qmv verify (`MLXCEL_QMV_MULTIROW`, issue #725) the Gemma 4 12B pairing profiles to about 1.5× and enables, while the same pairing on the pre-#725 per-row verify profiles to about 0.5× and declines with margin. When a window collects no probe signal the pre-#736 shape heuristic (issue #638: 1.0 on Apple Silicon, `sqrt(K)` elsewhere) remains as the fallback. The resulting state is readable at `GET /v1/internal/mtp-policy`; see [Adaptive MTP policy API](mtp-policy-api.md), and read it there rather than parsing the hint files, whose format is private. | | `MLXCEL_ENABLE_MTP_B1` | `0`/`false`/`no`/`off` to disable, any other value to force on | adaptive (per hardware) | Manual override for the singleton (B=1) MTP burst, in both directions. When set it pins the decision and disables adaptive profiling (issue #333). When unset, the adaptive policy decides (see `MLXCEL_MTP_ADAPTIVE`); with `MLXCEL_MTP_ADAPTIVE=0` the decision is the static per-hardware default (issue #165): non-batchable targets (`gemma4_unified` 12B pairs, whose only decode path is B=1) default **on** everywhere (measured across three prompts: 1.90x to 3.14x on M5 Max, 1.74x to 2.61x on M3 Ultra, 0.95x to 1.48x on M1 Ultra, where the prose end is a 5% loss rather than a gain); batch-capable targets (the 31B + bf16 assistant) default **on from Apple GPU generation 15** (M3, M4, M5), and fall back to classic decode on generation 13 (M1, M2). This was `has_neural_accelerator` (M5 only) until issue #1217, on the strength of ~1.2 to 1.4× on M5 Max against a ~0.75 to 0.96× regression on M1 Ultra, both measured before #1194/#1199/#1203/#1208/#1215 and neither re-measured since. M3 Ultra, which the old predicate lumped with M1 Ultra, was never measured on this pairing at all until #1217 did it: 1.95× (prose), 2.41× (source code) and 2.65× (enumeration) on 2026-08-20 under the #1215 protocol, at a verify-round cost of 1.51 classic steps against 2.71 on M1 Ultra. The discriminator is the `use_qmv_wide` split (an affine-quantized projection at `M >= 2` runs as one wide pass from generation 15 and as `K` narrow passes below it), not the Neural Accelerator. M4 is grouped by that shared dispatch rather than measured. Generation 13 keeps declining and the width sweep supports that: round cost fits `0.83 + 0.170 K` classic steps for this pairing on M3 Ultra against `1.14 + 0.090 K` for the 12B pairing there, so the bf16 drafter costs about 1.9x as much per block position, and carrying that ratio onto generation 13's `1.35 + 0.346 K` puts a block-4 round near 3.6 classic steps, which 2.96 to 3.99 emitted tokens would only just cover. See `docs/benchmark_results/mtp-b1-gate-m3ultra-2026-08-20.md`. | | `MLXCEL_ENABLE_MTP_BATCH` | truthy value | off | **Advanced.** Forces the batched Gemma 4 MTP burst path for parity/debug testing. Not governed by the adaptive policy (issue #333), which scopes to the validated, byte-identical B=1 path. | -| `MLXCEL_ENABLE_MTP_DEFERRED` | `1` | off | **Advanced.** Enables the deferred greedy verifier path for Gemma 4 MTP when sampling settings allow it. | | `MLXCEL_METAL4_ATTENTION` | `0`/`false`/`no`/`off` to disable; unset or any other value to enable | on (where the hardware has it) | **Advanced, diagnostic kill switch.** Forces the M5 neural-accelerator fused attention route off on hardware that has it, so `layers::metal4_causal_attention` is skipped and the ordinary SDPA path runs instead. Off-switch only: it cannot turn the route on where `has_neural_accelerator && macos_supports_na` is false, so setting it on an M1 is inert. This route is the first suspect whenever an M5 disagrees numerically with an earlier Apple GPU generation (issue #1065), and before this switch the only way to A/B the hypothesis was to patch `should_use_metal4_attention` and rebuild, which is what the #1182 M5 investigation had to do. Read once per process, so set it before starting `mlxcel` or `mlxcel-server`. Inert on non-Metal builds. | | `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. | diff --git a/examples/mtp_projection_width_bench.rs b/examples/mtp_projection_width_bench.rs new file mode 100644 index 000000000..388e5cd11 --- /dev/null +++ b/examples/mtp_projection_width_bench.rs @@ -0,0 +1,215 @@ +// MTP verify LM-head projection cost as a function of projected width. +// +// Issue #1179 asks whether a device-side early-exit verify walk is worth +// building. Early exit cannot skip the transformer forward (the whole block +// is one forward); all it can skip is the tail of the hidden-to-logits +// projection plus its argmax. This bench measures the ceiling of that +// saving: the cost of projecting W positions through the LM head at the +// real head shapes, for W from 1 to the widest block anyone runs. If the +// curve is flat in W, the projection is weight-read-bound and early exit +// has nothing to win at any accept rate. +// +// Method mirrors `qmm_gemv_microbench.rs`: quantize a bf16 weight at the +// real (out, in, bits, group) shape into enough independent copies that the +// working set defeats caches, then time T back-to-back projection+argmax +// graphs folded into one root so one eval submits the batch. The argmax is +// part of the measured graph because the device walk would fuse it. +// +// Guard (issue #1179 "the two arms must be provably different"): each row +// re-derives the logits and argmax shapes from a check call and prints +// them; a row only counts if the logits shape is [1, W, out]. +// +// Usage: cargo run --release --features metal,accelerate \ +// --example mtp_projection_width_bench [ T_PER_BATCH ] [ ROUNDS ] +// Defaults: T_PER_BATCH=16, ROUNDS=5 (best round reported). + +use mlxcel_core::{ + MlxArray, UniquePtr, add, argmax_last_axis, astype, eval, from_slice_f32, + quantize_weights_biases, quantize_weights_scales, quantize_weights_w, quantized_matmul, + synchronize_default, +}; +use std::time::Instant; + +const GROUP_SIZE: i32 = 64; +const BF16: i32 = 12; + +struct Head { + name: &'static str, + out: i32, + inp: i32, + bits: i32, +} + +// The two MTP-target LM heads in production today, from the local +// checkpoints' config.json: Gemma 4 12B ties the 262144-token embedding at +// hidden 3840; Qwen 3.8 27B has an untied 248320 x 5120 head. Both affine +// 4-bit group 64 in the mlx-community checkpoints. +const HEADS: &[Head] = &[ + Head { + name: "gemma4-12b tied head", + out: 262144, + inp: 3840, + bits: 4, + }, + Head { + name: "qwen3.8-27b lm_head", + out: 248320, + inp: 5120, + bits: 4, + }, +]; + +const WIDTHS: &[i32] = &[1, 2, 3, 4, 5, 8, 16, 32]; + +fn weight_bytes(out: i64, inp: i64, bits: i64) -> i64 { + let packed = out * inp * bits / 8; + let groups = out * (inp / GROUP_SIZE as i64); + packed + groups * 4 +} + +fn make_bf16(shape: &[i32], seed: usize) -> UniquePtr { + let total: usize = shape.iter().map(|&d| d as usize).product(); + let data: Vec = (0..total) + .map(|i| (((i + seed * 7919) as f32) * 0.000271).sin() * 0.05) + .collect(); + let f32arr = from_slice_f32(&data, shape); + astype(&f32arr, BF16) +} + +struct QuantizedWeight { + w: UniquePtr, + scales: UniquePtr, + biases: UniquePtr, +} + +fn quantize(head: &Head, seed: usize) -> QuantizedWeight { + let wf = make_bf16(&[head.out, head.inp], seed); + let w = quantize_weights_w(&wf, GROUP_SIZE, head.bits); + let scales = quantize_weights_scales(&wf, GROUP_SIZE, head.bits); + let biases = quantize_weights_biases(&wf, GROUP_SIZE, head.bits); + eval(&w); + eval(&scales); + eval(&biases); + QuantizedWeight { w, scales, biases } +} + +/// One projection+argmax graph: `[1, W, in] -> [1, W, out] -> [1, W]`. +fn project_argmax(x: &MlxArray, qw: &QuantizedWeight, bits: i32) -> UniquePtr { + // SAFETY: all arrays outlive the call; the biases pointer is valid. + let logits = unsafe { + quantized_matmul( + x, + &qw.w, + &qw.scales, + qw.biases.as_ref().unwrap() as *const MlxArray, + true, + GROUP_SIZE, + bits, + "affine", + ) + }; + argmax_last_axis(&logits) +} + +fn run_batch(x: &MlxArray, weights: &[QuantizedWeight], t: usize, bits: i32) -> f64 { + let start = Instant::now(); + let mut acc: Option> = None; + for i in 0..t { + let ids = project_argmax(x, &weights[i % weights.len()], bits); + acc = Some(match acc { + None => ids, + Some(a) => add(&a, &ids), + }); + } + let root = acc.unwrap(); + eval(&root); + synchronize_default(); + start.elapsed().as_secs_f64() +} + +fn main() { + let args: Vec = std::env::args().collect(); + let t_per_batch: usize = args.get(1).and_then(|s| s.parse().ok()).unwrap_or(16); + let rounds: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(5); + + println!( + "=== MTP LM-head projection width bench (issue #1179) === T={} ROUNDS={} group_size={}", + t_per_batch, rounds, GROUP_SIZE + ); + + for head in HEADS { + let bytes = weight_bytes(head.out as i64, head.inp as i64, head.bits as i64); + let copies = ((256 * 1024 * 1024) / bytes).clamp(2, 8) as usize; + let weights: Vec = (0..copies).map(|s| quantize(head, s + 1)).collect(); + + println!( + "\n{} ({} x {}, {}-bit, {:.0} MB/read, {} copies)", + head.name, + head.out, + head.inp, + head.bits, + bytes as f64 / 1e6, + copies + ); + println!( + "{:>5} {:>22} {:>10} {:>9} {:>12}", + "W", "logits shape (guard)", "ms/call", "GB/s", "vs W=1" + ); + + let mut per_w_ms: Vec<(i32, f64)> = Vec::new(); + for &w in WIDTHS { + let x = make_bf16(&[1, w, head.inp], 0); + eval(&x); + synchronize_default(); + + // Guard: derive the shapes this arm actually projects. + let check = unsafe { + quantized_matmul( + &x, + &weights[0].w, + &weights[0].scales, + weights[0].biases.as_ref().unwrap() as *const MlxArray, + true, + GROUP_SIZE, + head.bits, + "affine", + ) + }; + let logits_shape = mlxcel_core::array_shape(&check); + assert_eq!( + logits_shape, + vec![1, w, head.out], + "guard: this arm is not projecting W={w} positions" + ); + + let _ = run_batch(&x, &weights, t_per_batch, head.bits); // warmup + let mut best = f64::MAX; + for _ in 0..rounds { + best = best.min(run_batch(&x, &weights, t_per_batch, head.bits)); + } + let ms_per_call = best * 1e3 / t_per_batch as f64; + let gbs = (bytes as f64 * t_per_batch as f64) / best / 1e9; + let vs_w1 = per_w_ms + .first() + .map(|&(_, base)| format!("{:+.1}%", (ms_per_call / base - 1.0) * 100.0)) + .unwrap_or_default(); + println!( + "{:>5} {:>22} {:>10.3} {:>9.0} {:>12}", + w, + format!("{:?}", logits_shape), + ms_per_call, + gbs, + vs_w1 + ); + per_w_ms.push((w, ms_per_call)); + } + + // The early-exit ceiling at block K with A positions kept is + // t(K) - t(A); print the widest useful contrast per K. + println!("early-exit ceiling (t(K) - t(1), the most a device walk could save):"); + for &(w, ms) in per_w_ms.iter().skip(1) { + let base = per_w_ms[0].1; + println!(" K={:<3} ceiling = {:.3} ms/round", w, ms - base); + } + } +} diff --git a/src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs b/src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs index 23829b516..86e60648c 100644 --- a/src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs +++ b/src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs @@ -919,8 +919,8 @@ impl DFlashGenerator { /// path (not DFlash). If a model with a non-empty suppressed-id set ever /// adopts the DFlash verify path, this call must be preceded by /// `mlxcel_core::sampling::apply_token_bias(logits, token_bias)` — exactly as -/// `Gemma4MtpTargetAdapter::argmax_from_hidden_positions` does — so that -/// suppressed placeholder token ids cannot win a verify position and +/// the Gemma 4 MTP verify forward does before its per-position argmax — so +/// that suppressed placeholder token ids cannot win a verify position and /// reintroduce the #350 leak. fn argmax_logits_to_array(logits: &MlxArray, seq_len: i32) -> UniquePtr { let shape = ffi::array_shape(logits); diff --git a/src/models/gated_delta.rs b/src/models/gated_delta.rs index 67f4b0911..7dcff88f0 100644 --- a/src/models/gated_delta.rs +++ b/src/models/gated_delta.rs @@ -742,8 +742,7 @@ pub fn gated_delta_update_chain_parity( gated_delta_ops_with_parity(q, k, v, &g, &beta, state, mask, chain_parity_enabled()) } -/// Diagnostic escape hatch for the chain-parity kernel, mirroring the -/// `MLXCEL_ENABLE_MTP_DEFERRED` parity-experiment precedent: +/// Diagnostic escape hatch for the chain-parity kernel: /// `MLXCEL_GDN_CHAIN_PARITY=0` restores the pre-#1165 block numerics /// (float32 in-block state carry) on the speculative verify / rollback /// paths for A/B attribution of the parity kernel's cost and acceptance diff --git a/src/models/gemma4.rs b/src/models/gemma4.rs index 0f07d6ade..c202f9bba 100644 --- a/src/models/gemma4.rs +++ b/src/models/gemma4.rs @@ -4494,60 +4494,13 @@ impl Gemma4Model { logits } - /// Run the transformer with speculative sinks but skip the tied LM head. - /// - /// Used by: Gemma 4 MTP deferred greedy verification, which needs the - /// pre-norm hidden states and shared K/V slabs but can project only the - /// positions required by the speculative walk. - #[allow(clippy::too_many_arguments)] - fn forward_hidden_with_caches_and_speculative_sinks( - &self, - input_ids: &MlxArray, - input_embeddings: Option<&MlxArray>, - caches: &mut [Cache], - mask: Option<&MlxArray>, - per_layer_inputs: Option<&MlxArray>, - capture_layer_ids: Option<&[usize]>, - sinks: Option<&mut Gemma4SpeculativeSinks>, - skip_final_norm: bool, - left_padding: Option<&[i32]>, - ) -> UniquePtr { - self.text_model.forward_with_speculative_sinks( - input_ids, - input_embeddings, - caches, - mask, - per_layer_inputs, - capture_layer_ids, - sinks, - skip_final_norm, - None, - left_padding, - None, - ) - } - /// Apply the Gemma 4 final norm to a pre-norm decoder hidden state. /// - /// Used by: Gemma 4 MTP drafter hidden preparation and deferred - /// hidden-to-logits verification. + /// Used by: Gemma 4 MTP drafter hidden preparation. fn speculative_draft_hidden(&self, hidden: &MlxArray) -> UniquePtr { self.text_model.norm.forward(hidden) } - /// Project a pre-norm decoder hidden state to logits using the tied LM - /// head and optional final-logit softcap. - /// - /// Used by: Gemma 4 MTP deferred greedy verification. - fn speculative_logits_from_hidden(&self, hidden: &MlxArray) -> UniquePtr { - let hidden = self.speculative_draft_hidden(hidden); - let mut logits = self.text_model.embed_tokens.as_linear(&hidden); - if let Some(cap) = self.config.final_logit_softcapping { - logits = mlxcel_core::compiled_softcap(&logits, cap); - } - logits - } - pub(crate) fn make_caches(&self) -> Vec { self.text_model.make_caches() } @@ -5547,36 +5500,6 @@ impl Gemma4Wrapper { ) } - pub(crate) fn forward_hidden_with_speculative_sinks( - &self, - input_ids: &MlxArray, - input_embeddings: Option<&MlxArray>, - per_layer_inputs: Option<&MlxArray>, - mask: Option<&MlxArray>, - seq_id: Option, - capture_layer_ids: Option<&[usize]>, - sinks: Option<&mut Gemma4SpeculativeSinks>, - skip_final_norm: bool, - ) -> UniquePtr { - self.sequence_state.with_or_create_sequence_state( - seq_id, - || self.model.make_caches(), - |sequence_caches| { - self.model.forward_hidden_with_caches_and_speculative_sinks( - input_ids, - input_embeddings, - sequence_caches, - mask, - per_layer_inputs, - capture_layer_ids, - sinks, - skip_final_norm, - None, - ) - }, - ) - } - /// Normalize a pre-norm hidden state before handing it to the MTP /// assistant drafter. /// @@ -5585,14 +5508,6 @@ impl Gemma4Wrapper { self.model.speculative_draft_hidden(hidden) } - /// Project a pre-norm hidden state to logits without rerunning the - /// transformer. - /// - /// Used by: Gemma 4 MTP deferred greedy verification. - pub(crate) fn speculative_logits_from_hidden(&self, hidden: &MlxArray) -> UniquePtr { - self.model.speculative_logits_from_hidden(hidden) - } - /// Allocate a fresh per-layer cache vector for a batched MTP burst /// Every cache starts empty; the batched verify pass /// grows them with a leading batch dim `B` once the first `[B, L]` diff --git a/src/models/gemma4_mtp_target.rs b/src/models/gemma4_mtp_target.rs index 09bdc8c8b..4fe469c4b 100644 --- a/src/models/gemma4_mtp_target.rs +++ b/src/models/gemma4_mtp_target.rs @@ -229,61 +229,36 @@ impl<'a> Gemma4MtpTargetAdapter<'a> { let verify_arr = mlxcel_core::from_slice_i32(verify_input, &[1, verify_input.len() as i32]); let mut sinks = Gemma4SpeculativeSinks::with_hidden_and_shared_kv(); sinks.tree_positions = tree_positions; - // Greedy/no-logprobs can use the latest upstream deferred path: run - // the transformer once with `skip_final_norm=True`, capture pre-norm - // hidden/shared K/V, then project hidden positions to logits only as - // needed. Keep the full-logits path for non-greedy or logprob - // requests so existing sampler/logprob semantics stay unchanged. - // The upstream Python reference uses deferred greedy hidden→logits - // projection by default. In Rust/MLX today that path projects one - // position at a time across the cxx bridge and is slower than the - // batched `[K, vocab]` LM-head projection for Gemma 4 31B on local - // Apple Silicon runs. Keep it available for parity experiments, but - // leave the faster full-logits verifier as the default until we have - // a fused/graph-side deferred walk. - let use_deferred_greedy = std::env::var("MLXCEL_ENABLE_MTP_DEFERRED").ok().as_deref() - == Some("1") - && sampler.temperature == 0.0 - && !logprobs_config.enabled; - let logits = if use_deferred_greedy { - let _ = self.wrapper.forward_hidden_with_speculative_sinks( - &verify_arr, - None, - None, - mask, - self.seq_id, - None, - Some(&mut sinks), - true, - ); - None + // One forward returns the full `[1, K, vocab]` verify logits: the LM + // head projects all K positions in one batched call, and the argmax + // walk reads them after one eval. The projection deliberately does + // not early-stop at the first draft mismatch: the head is + // weight-read-bound, so projecting fewer positions saves almost + // nothing (measured on #1179, where a hidden-first "deferred" split + // of this call was also retired for adding a bridge crossing without + // deferring any work). + let raw_logits = self.wrapper.forward_with_speculative_sinks( + &verify_arr, + None, + None, + mask, + self.seq_id, + None, + Some(&mut sinks), + None, + ); + // issue #350: mask the model's output-illegal placeholder ids on the + // verify logits BEFORE the per-position argmax AND the logprob + // extraction, so a suppressed id can never win a near-tie verify + // position (the leak) and the emitted token plus its logprob stay + // consistent. The bias is the model's suppressed-id set, forwarded + // from the server `enqueue_request` through `MtpGenerator` into + // `sampler`. An empty map (every non-multimodal model) short-circuits + // here to the raw logits, so the baseline stays bit-exact. + let logits = if sampler.token_bias.is_empty() { + raw_logits } else { - let raw_logits = self.wrapper.forward_with_speculative_sinks( - &verify_arr, - None, - None, - mask, - self.seq_id, - None, - Some(&mut sinks), - None, - ); - // issue #350: mask the model's output-illegal placeholder ids on the - // verify logits BEFORE the per-position argmax AND the logprob - // extraction, so a suppressed id can never win a near-tie verify - // position (the leak) and the emitted token plus its logprob stay - // consistent. The bias is the model's suppressed-id set, forwarded - // from the server `enqueue_request` through `MtpGenerator` into - // `sampler`. An empty map (every non-multimodal model) short-circuits - // here to the raw logits, so the baseline stays bit-exact. - if sampler.token_bias.is_empty() { - Some(raw_logits) - } else { - Some(mlxcel_core::sampling::apply_token_bias( - &raw_logits, - &sampler.token_bias, - )) - } + mlxcel_core::sampling::apply_token_bias(&raw_logits, &sampler.token_bias) }; // Greedy-parity gate: pull the per-position argmax tokens from @@ -293,24 +268,17 @@ impl<'a> Gemma4MtpTargetAdapter<'a> { // At temperature > 0 a future enhancement plumbs the sampler // through per-position; the round-loop driver's perf-sensitive // path is greedy, so we keep argmax-only for now. - let target_tokens = if let Some(logits) = logits.as_ref() { - let ids = Self::argmax_per_position(logits); - Self::log_top_two_gaps(logits, &ids); + let target_tokens = { + let ids = Self::argmax_per_position(&logits); + Self::log_top_two_gaps(&logits, &ids); ids - } else { - // The hidden sink is still owned by `sinks`; pull it below but - // compute after extraction so the hidden handle is available for - // both token projection and `VerifyCaptured`. - Vec::new() }; // Per-position log-probability data, aligned 1:1 with // `target_tokens`. `None` (zero-overhead) when logprobs are // disabled; the round loop forwards the entries for accepted // positions on to `finalize_burst_success`. - let target_logprobs = logits.as_ref().and_then(|logits| { - Self::per_position_logprobs(logits, &target_tokens, logprobs_config) - }); + let target_logprobs = Self::per_position_logprobs(&logits, &target_tokens, logprobs_config); // Capture the hidden + pre-slice shared K/V for the finalize step. let hidden_sink = sinks @@ -320,11 +288,6 @@ impl<'a> Gemma4MtpTargetAdapter<'a> { .into_iter() .next_back() .expect("hidden sink must carry at least one entry"); - let target_tokens = if target_tokens.is_empty() && use_deferred_greedy { - self.argmax_from_hidden_positions(hidden_full.as_ref().unwrap(), &sampler.token_bias) - } else { - target_tokens - }; let shared_kv_map = sinks .shared_kv_sink @@ -444,42 +407,6 @@ impl<'a> Gemma4MtpTargetAdapter<'a> { .speculative_draft_hidden(hidden.as_ref().unwrap()) } - /// Greedy target-token extraction from pre-norm hidden states without a - /// Rust-side per-position FFI loop. - /// - /// This keeps the upstream-style `skip_final_norm=True` verify path but - /// projects the whole `[B=1, K, H]` hidden block through - /// `speculative_logits_from_hidden()` in one MLX graph, then materializes - /// the `[K]` argmax tensor with one host copy. It deliberately does not - /// early-stop on the first mismatch: for the small Gemma 4 MTP block sizes - /// we use today, avoiding `K` separate cxx/MLX calls is more important - /// than skipping the tail projection on low-accept rounds. - /// - /// `token_bias` is the model's output-suppression set (issue #350). When it - /// is non-empty the projected logits are biased before the argmax so a - /// suppressed placeholder id can never win a deferred-greedy verify - /// position. An empty map (every non-multimodal model) takes the raw argmax, - /// preserving the bit-exact baseline. - fn argmax_from_hidden_positions( - &self, - hidden_full: &MlxArray, - token_bias: &mlxcel_core::sampling::TokenBiasMap, - ) -> Vec { - let shape = mlxcel_core::array_shape(hidden_full); - debug_assert_eq!(shape.len(), 3, "hidden must be 3-D [B, T, H]"); - let expected_len = shape[1].max(0) as usize; - let logits = self.wrapper.speculative_logits_from_hidden(hidden_full); - let argmax = if token_bias.is_empty() { - mlxcel_core::argmax_last_axis(logits.as_ref().unwrap()) - } else { - let biased = - mlxcel_core::sampling::apply_token_bias(logits.as_ref().unwrap(), token_bias); - mlxcel_core::argmax_last_axis(&biased) - }; - mlxcel_core::eval(&argmax); - materialize_argmax_i32_vec(&argmax, expected_len) - } - /// Slice the captured shared K/V tensors along the seq-len axis by /// `rejected = block_size - accepted - 1` so the post-rollback K/V /// matches the trimmed cache.