diff --git a/CHANGELOG.md b/CHANGELOG.md index 0f40cf578..de6b2941a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ### Added - **`mlx-community/Qwen3.8-27B-4bit` is qualified on the existing `qwen3_5` path** (#1163). Qwen3.8 declares `model_type: "qwen3_5"` / `Qwen3_5ForConditionalGeneration` and is architecturally identical to Qwen3.5-27B down to a byte-identical weight-map key set, so it already loaded and ran on `main` with zero code changes; that was accidental rather than guaranteed, since `Qwen35Config` carries no `deny_unknown_fields` and silently dropped three keys the generation added that are load-bearing upstream (`output_gate_type`, `rope_parameters.mrope_interleaved`, and the top-level `language_model_only`). Those three are now read at load time and rejected with a named error when they ask for behavior mlxcel does not implement (`output_gate_type: "sigmoid"`, `mrope_interleaved: false`, `language_model_only: true`), across every Qwen3.5-family parse site including the MiniCPM-V 4.6 text backbone, instead of being dropped silently and producing wrong output. **Upgrade note:** `vision_start_token_id` is now mandatory for the Qwen3.5-family VLM path; the stale `248045` default this replaces was wrong for every shipped checkpoint in the family (all use `248053`), and loading with the wrong id used to mis-segment MRoPE vision spans without failing. Every known checkpoint in the family supplies the key, and the new failure is fail-closed at startup with a message naming the missing key, but a deployment running a checkpoint whose `config.json` happens to omit it will see a new startup failure on upgrade rather than the previous silent misbehavior. +- **Qwen 3.5 MTP speculative decoding, driven by the `qwen3_5_mtp` drafter family (`mlx-community/Qwen3.8-27B-MTP-bf16` and siblings)** (#1165). The Qwen 3.5 (dense and MoE, text and VLM) targets now implement `MtpTarget` alongside Gemma 4, so the same `MtpGenerator` round loop, the tick-cooperative slice path, and the run-to-completion burst all drive this pairing on both the offline CLI and the server. Temperature-0 output is byte-identical to classic decode on this pairing only when the Metal chain-parity gated-delta kernel (`gated_delta_step_seqpar`, `MLXCEL_GDN_CHAIN_PARITY`) is taken: the standard kernel carries float32 recurrent state across a verify block and rounds it to the storage dtype once at the end, while the classic single-token decode chain rounds after every token, so without the parity kernel a near-tie argmax can flip about once every 100-250 generated tokens. `mtp_capable_target` gates this pairing on both Metal availability and the checkpoint's GDN shape satisfying the parity kernel's contract, declining to classic decode rather than silently forfeiting exactness when either does not hold. **Measured verdict (M1 Ultra, Mac Studio):** MTP decode is 0.70x classic at the drafter's own configured block size (3) and 0.59x at the historical block-4 default, both below the 1.0x break-even, because the target's multi-token verify forward does not amortize across the block on this GPU generation (about 29 ms per extra block position, roughly two thirds of a full decode step); the adaptive MTP policy (issue #333) profiles and declines this pairing by default on pre-M5 hardware (`mtp_b1_default`), matching the precedent already recorded for Gemma 4 MTP. M5-class hardware is not measured; that cell is open. See `docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md` for the full record, including the DFlash blast-radius A/B (`forward_speculative` is shared verify machinery) and the explicitly unmeasured acceptance-under-concurrent-slice-rotation cell. The drafter-configured `block_size` (3) is now the default when `--draft-block-size` is not passed, instead of the flat Gemma-4-derived constant (4); the PR's own measurement found the configured value faster (16.48 vs 13.81 tok/s, 0.591 vs 0.465 acceptance). ### Fixed diff --git a/TECHNICAL_REPORTS/1182-qwen35-mtp-drafter-20260816.en.md b/TECHNICAL_REPORTS/1182-qwen35-mtp-drafter-20260816.en.md new file mode 100644 index 000000000..d7eea37b8 --- /dev/null +++ b/TECHNICAL_REPORTS/1182-qwen35-mtp-drafter-20260816.en.md @@ -0,0 +1,315 @@ +# Technical Report: PR #1182 - feat(inference): support the qwen3_5_mtp drafter for Qwen 3.5 MTP speculative decoding + +**Date**: 2026-08-16 +**Author**: AI Code Reviewer +**Status**: Completed +**Languages**: Rust, C++ (Metal kernel source) +**Risk Level**: Low (the pairing is gated off by default on every hardware generation measured so far; every dispatch site declines to classic decode rather than silently violating the exactness contract on an unsupported pairing) + +--- + +## Executive Summary + +PR #1182 (issue #1165) adds the split-out `qwen3_5_mtp` drafter (e.g. `mlx-community/Qwen3.8-27B-MTP-bf16`) as a first-class MTP drafter and implements mlxcel's `MtpTarget` trait for the Qwen 3.5 family (dense and MoE, text and VLM), so the existing `MtpGenerator` round loop, the tick-cooperative slice path, and the run-to-completion burst all drive this pairing exactly as they already drive Gemma 4. The result only ships correctly because of one specific numerical fix: the temperature-0 exactness gate initially failed with a near-tie argmax flip roughly once every 100-250 generated tokens, traced to the shared gated-delta Metal kernel carrying its recurrent state in float32 across a `T = K` verify block and quantizing to the storage dtype only at the final store, while the classic single-token decode chain quantizes after every token. A `T = K` block was therefore not bit-identical to `K` consecutive `T = 1` decode steps, and a last-ulp state difference flips near-tie logits. The fix is a chain-parity kernel variant (`gated_delta_step_seqpar`) that rounds the state through the storage dtype at the end of every in-block step, making the two shapes provably bit-identical (pinned in `parity_kernel_block_is_bitwise_equal_to_single_token_chain` and confirmed over full 240-token decode chains); it is used only by the speculative verify capture path and the rollback replay, so classic prefill and decode are untouched and DFlash inherits the same numerics for free through the shared `forward_speculative` path. + +The measured hardware verdict is negative on the machine available for this PR: MTP decode measures **0.70x** classic at the drafter's own configured block size (3) and 0.59x at the historical block-4 default, on an Apple M1 Ultra. The cause is not a defect in the implementation, every diagnostic in the round loop is healthy, but a hardware ceiling: `verify_forward_cost_scaling` shows the target's own plain multi-token forward gains almost no amortization at this GPU generation's block sizes (about 29 ms per extra block position, roughly two-thirds of a full decode step, on both the verify path and the plain batched-prefill path equally). This is the same pre-M5 regression already encoded for Gemma 4 MTP in `mtp_b1_default` (issue #165: M1 Ultra 0.75-0.96x, M5 Max 1.2-1.4x), and the adaptive MTP policy (issue #333) gates Qwen 3.5 MTP through the identical static-default-plus-profiling mechanism, so a deployment on this hardware class declines to classic decode by default rather than silently regressing. The M5 cell (and M3/M4) for this specific pairing is open; the PR includes a reproduction guide in the issue thread for anyone with different hardware to fill it in. + +The third load-bearing fact is a premise correction made before implementation started. The issue body originally scoped "the hard part" as inventing partial-acceptance rollback for a target that mixes 16 full-attention layers with 48 gated-delta linear layers, on the theory that a recurrent state cannot be trimmed and must be snapshotted and restored. That machinery already existed: `Qwen35Model::rollback_speculative_cache` implements exactly this for the DFlash path, trimming attention KV by the unaccepted tail and replaying gated-delta snapshots over the accepted prefix, and the verify pass already returns the hidden states an MTP drafter needs. The genuinely new work was therefore wiring, not cache surgery: a new drafter model, the `MtpTarget` trait implementation wrapping the existing verify/rollback surface, and lifting roughly a dozen Gemma-4-only gates across the server and offline CLI dispatch paths. + +Review and security passes on the initial PR fixed two HIGH findings (commit `60ef14b2`) and one HIGH finding (commit `ef3aa95e`) before this finalization pass. This pass closed ten additional findings the review flagged: an unattributed benchmark that could report self-comparison noise as a result, a capability gate that checked Metal availability but not the chain-parity kernel's actual shape contract, a shipped default block size that the PR's own measurement showed was the slower of the two configurations, a C++ dispatch path that silently drops the chain-parity request on an unreachable-today shape, an unmeasured acceptance-under-concurrent-load cell reported as if it were the measured single-request number, a missing environment-variable doc entry, roughly ten stale Gemma-4-only doc claims, three now-false claims that the Qwen drafter's `reset` is a no-op, a 335 MB seed hidden state held a round longer than necessary, and dead `if (true)` scaffolding copied into the new Metal kernel source. + +--- + +## 1. Problem Statement + +### 1.1 Background + +Issue #1165 established that Qwen3.8-27B's MTP head is not part of the target checkpoint in any MLX conversion: the upstream `Qwen/Qwen3.8-27B` base ships 15 `mtp.*` tensors alongside `model.*`, but every `mlx-community/Qwen3.8-27B-*` conversion drops them, and mlx-community republishes them as a standalone drafter (`mlx-community/Qwen3.8-27B-MTP-bf16`, `model_type: "qwen3_5_mtp"`, `block_size: 3`). Before this PR, `mlxcel generate --draft-model ` failed with `Unsupported model type: qwen3_5_mtp` regardless of `--draft-kind`, both messages accurate and fail-closed, so this was a missing capability rather than a latent bug. + +### 1.2 Existing Issues (from the review/security pass, addressed in finalization) + +- **F1** (High, process/measurement integrity): `parity_kernel_cost_vs_standard` had no arm-attribution guard. With `MLXCEL_GDN_CHAIN_PARITY=0` in the environment, `gated_delta_update_chain_parity` silently falls back to the standard kernel (`chain_parity_enabled()` reads the env var), so both arms of the benchmark resolve to the same code path and the test would compare it against itself. The recorded ratios (1.089 / 1.117 / 0.834 / 1.008) were indistinguishable from that failure mode, precisely the pattern CLAUDE.md's "prove the two arms differ before reporting any number" rule targets (epic #909 named three prior benchmarks with this exact defect). +- **F2** (High, correctness): `mtp_capable_target` gated the Qwen 3.5 family on `metal_is_available()` alone. Byte-identity additionally requires `supports_metal_gated_delta_kernel(hk, hv, dk, dv)` (`dk >= 32 && dk % 32 == 0 && hv % hk == 0`); a Metal host whose checkpoint has an incompatible `linear_key_head_dim` would silently take the ops-based fallback and forfeit the temperature-0 exactness contract with no signal. The tree's own test fixtures already use `linear_key_head_dim: 4` (not a hypothetical shape). +- **F3** (Medium, performance/correctness of defaults): `prefer_requested_block_size() = true` for the Qwen drafter means the drafter's own `configured_block_size()` (3) was never consulted; the per-kind CLI default (4, a Gemma-4-derived constant) always won. The PR's own measurement shows block 3 at 16.48 tok/s / 0.591 acceptance versus block 4 at 13.81 / 0.465, so the shipped default was the slower measured configuration. +- **F4** (Medium, defense in depth): the C++ dispatch in `metal_gated_delta_forward_impl` silently takes the standard kernel branch whenever the gate is vectorized or a mask is present, even when the caller requested `chain_parity = true`. Not reachable from the Qwen 3.5 verify path today (`mask = None`, `g` is always 3-dim there), but nothing enforced that, so a future caller could forfeit byte-identity with no signal. +- **F5** (Medium, documentation honesty): the benchmark record's 0.591 acceptance figure is single-request only; it did not state that `Qwen35MtpDraftModel::reset` destroys the drafter's accumulated KV history (unlike Gemma 4's no-op reset) and fires on every park during concurrent slice-grant rotation, so acceptance under load is structurally worse and was not characterized. +- **F6-F10** (Low, documentation and cleanup): `MLXCEL_GDN_CHAIN_PARITY` was undocumented in `docs/environment-variables.md`; roughly ten module/API doc comments across `speculative_slice.rs`, `speculative_burst.rs`, `scheduler.rs`, `drafter/mod.rs`, and `speculative/mtp/{mod,generator}.rs` still asserted MTP was Gemma-4-only; three of those also claimed the MTP `Drafter::reset` is universally the trait-default no-op, which is false for the Qwen drafter; the seed-round `verify_hidden_full` (`[1, P, H]`, about 335 MB at P=32k/H=5120 in f16) was held alive through an entire round after its one-shot consumer instead of being cleared immediately; and the new chain-parity Metal kernel source carried a dead `if (true) { ... }` wrapper copied from the masked kernel variant. + +### 1.3 Risk Assessment + +| Risk | Impact | Likelihood | +|------|--------|------------| +| Benchmark numbers indistinguishable from self-comparison noise reported as a real ratio (F1) | Medium (misleads future capacity/perf decisions, no user-facing effect) | Confirmed reproducible: setting `MLXCEL_GDN_CHAIN_PARITY=0` and re-running the original test made both arms identical | +| A Metal host with an incompatible GDN shape silently loses temperature-0 exactness (F2) | High if it occurs (silently wrong generation output under a documented guarantee) | Low today (every published `mlx-community` Qwen 3.5 checkpoint has `dk = 128`), but the gate did not check it | +| Shipping the slower of two measured block-size configurations as the silent default (F3) | Medium (throughput regression an operator would not discover without manually overriding the flag) | Certain before the fix: every unconfigured Qwen 3.5 MTP request used block 4 | +| C++ silently dropping `chain_parity` on vectorized/masked shapes (F4) | High if it occurs (same exactness forfeiture as F2) | Very low today (unreachable from any current caller), real for future callers | +| Acceptance under concurrent load reported implicitly as the single-request number (F5) | Low-Medium (an operator sizing concurrent MTP capacity from the doc gets an optimistic number) | Certain: no concurrent-rotation measurement exists anywhere in the tree | +| Stale docs (F6-F10) | Low (reader/maintainer confusion, no functional impact) | Certain prior to this pass; corrected here | + +--- + +## 2. Technical Review + +### 2.1 Security + +Review and security passes were completed by the orchestrator, `pr-reviewer`, and `pr-security-checker` before this finalization pass: two HIGH findings fixed in `60ef14b2` (MTP dispatch-site coverage test, chain-parity test lint) and one HIGH finding fixed in `ef3aa95e` (bounding the `qwen3_5_mtp` drafter's config scalars against a divide-by-zero and an unbounded-allocation abort at load, since the drafter loads lazily inside a core worker thread that aborts the process on an uncaught panic). This pass did not revisit either; it closed the ten findings in 1.2 above. + +**Issues Found (this pass):** + +| Issue | Severity | Status | +|-------|----------|--------| +| F1: no arm-attribution guard in `parity_kernel_cost_vs_standard` | High | Fixed | +| F2: `mtp_capable_target` checks Metal availability but not the parity kernel's shape contract | High | Fixed | +| F3: shipped default block size is the slower measured configuration | Medium | Fixed | +| F4: C++ dispatch silently ignores `chain_parity` on vectorized/masked shapes | Medium | Fixed | +| F5: acceptance-under-load left uncharacterized and unlabeled in the benchmark record | Medium | Fixed (documented as an explicit open cell) | +| F6-F10: env var doc, ~10 stale Gemma-4-only claims, 3 false `reset`-is-no-op claims, held-too-long seed hidden, dead kernel scaffolding | Low | Fixed | + +### 2.2 Performance + +The measured verdict (see Executive Summary and `docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md`) is negative on the one hardware generation available for this PR, and F3's fix changes which negative number ships by default (0.70x at block 3, versus 0.59x at the pre-fix block-4 default). Both remain below the 1.0x break-even the adaptive MTP policy uses to decide engagement, so the practical effect of F3 is a better-informed default for the case where the policy is bypassed (`MLXCEL_ENABLE_MTP_B1` pinned on, or profiling disabled), not a change in whether the pairing engages by default. + +F1's re-measurement after adding the arm-attribution guard produced five fresh runs, medians 0.979 / 1.007 / 0.994 / 1.022 at T=1/3/4/64 (min-max 0.877-1.043 across runs), still inside the dispatch-noise band the original (unattributed) numbers claimed, but now provably attributable: the guard dispatches both kernels once at `T=3`, asserts their final GDN state differs (which the two kernels are defined to do whenever `T > 1`), and panics before printing anything if that assertion fails. The corrected numbers differ from the originally recorded ones at the individual-ratio level (this is expected: the original numbers were not attributable to begin with, and dispatch-level noise on a loaded box does not reproduce ratio-for-ratio across runs), but the qualitative conclusion, that the kernel's added per-step register round costs nothing measurable outside noise, is unchanged. + +### 2.3 Compatibility & Dependencies + +- **Breaking Changes**: none. `resolve_draft_block_size`'s signature gained a `model_path: &Path` parameter (F3's fix); both of its call sites (`src/server/speculative_dispatch.rs`, `src/commands/generate.rs`) already have the drafter path in scope at that point. `supports_metal_gated_delta_kernel` changed from private to `pub(crate)` (F2) with no external surface change. +- **New Dependencies**: none. +- **Behavior change**: an operator who runs `mlxcel generate --draft-model --draft-kind mtp` (or the server equivalent) without an explicit `--draft-block-size` now gets block size 3 instead of 4, per F3. An explicit `--draft-block-size` override is unaffected. A Metal host whose target's GDN shape does not satisfy the parity kernel's contract now declines to classic decode (F2) instead of silently running without byte-identity; no published checkpoint is affected by this today. + +### 2.4 Code Quality + +- **Test Coverage**: this pass added `assert_parity_and_standard_arms_are_distinguishable` (the F1 guard, `tests/qwen38_mtp_chain_parity.rs`), three tests for `Qwen35Model::supports_chain_parity_kernel` / `qwen35_config_supports_chain_parity_kernel` (F2, `src/models/qwen3_5.rs`), five tests for `resolve_draft_block_size` and `peek_qwen35_mtp_configured_block_size` across the `mlxcel` and `mlxcel-core` crates (F3), and `chain_parity_forfeited_by_shape_matches_the_cpp_dispatch_branches` plus the underlying pure predicate (F4, `src/models/gated_delta.rs` / `_tests.rs`). +- **Code Complexity**: each fix is localized. F2 factors the GDN-shape check into `Qwen35Model::supports_chain_parity_kernel`, itself a thin wrapper over a free function (`qwen35_config_supports_chain_parity_kernel`) taking a bare `&Qwen35Config`, kept separate specifically so it is unit-testable without constructing a full model with real weights. F4 similarly factors the C++ shape predicate into a pure `chain_parity_forfeited_by_shape(bool, bool) -> bool` so the warning condition is testable without a Metal device. +- **Technical Debt**: decreased. The benchmark record now states its own attribution mechanism rather than asserting a number; the capability gate and the CLI default both now reflect the checkpoint's own declared shape/config instead of a Gemma-4-derived constant applied uniformly across an unrelated drafter family; and roughly a dozen doc sites that would have misled the next contributor into thinking Qwen 3.5 MTP support does not exist, or that its drafter reset is a no-op, are corrected. + +--- + +## 3. Technical Decisions + +### 3.1 Arm-attribution guard: assert-and-panic before the table, not a printed caveat + +**Alternatives Considered:** + +| Option | Pros | Cons | +|--------|------|------| +| Print a warning alongside the numbers when attribution cannot be confirmed | Keeps the test running, numbers still visible | Exactly the failure mode CLAUDE.md's rule targets: numbers get quoted regardless of the caveat | +| **Chosen: panic before printing any ratio if the arms are provably identical at `T=3`** | No route to a misleading number reaching a report; failure message names the likely cause (`MLXCEL_GDN_CHAIN_PARITY=0`) | The whole `--ignored --nocapture` invocation aborts on the check failing, rather than degrading gracefully | + +**Rationale:** the review explicitly named "make the test refuse to report when attribution fails" as the requirement; a warning-and-continue design still lets a copy-pasted number end up in a document. `T = 1` cannot serve as the discriminator (the two kernels are defined to be byte-identical there), so the check runs once at `T = 3`, matching the shape already pinned bitwise-different by the sibling `parity_kernel_block_is_bitwise_equal_to_single_token_chain` test. + +### 3.2 Capability gate: fold the shape predicate into `mtp_capable_target`, not a warning + +**Alternatives Considered:** + +| Option | Pros | Cons | +|--------|------|------| +| Emit a one-time `tracing::warn!` when parity is requested but the shape is incompatible, keep serving | Minimal code change | The exactness contract is documented as absolute ("temperature-0 output byte-identical to classic decode"); a warning does not stop a wrong answer from reaching a client | +| **Chosen: fold `supports_metal_gated_delta_kernel` into `mtp_capable_target` itself** | A checkpoint that cannot honor the contract declines to classic decode automatically, the same fail-closed posture every other MTP gate in this codebase already uses | Requires exposing the shape predicate from `mlxcel-core` and adding a config-reading accessor on `Qwen35Model` | + +**Rationale:** the review offered both options; "better" was explicitly named as folding the predicate into the gate. Given every other MTP capability check in this tree (Metal availability, target/drafter compatibility, hidden-size and vocab match) already declines rather than warns, a warn-only fix here would have been the one exception to an otherwise-consistent fail-closed design. + +### 3.3 Block-size default: peek the drafter's own config, don't change the shared per-kind constant + +**Alternatives Considered:** + +| Option | Pros | Cons | +|--------|------|------| +| Change `DEFAULT_MTP_BLOCK_SIZE` from 4 to 3 | One-line fix | `DrafterKind::Mtp` covers both Gemma 4 and Qwen 3.5; this would change Gemma 4's default too, on evidence measured only for Qwen 3.5 | +| Flip `prefer_requested_block_size()` to `false` for the Qwen drafter | Reuses the existing adaptive `effective_mtp_block_size` machinery | Upstream mlx-vlm sets `prefer_requested_block_size = True` specifically for this drafter; flipping it breaks that parity for a reason unrelated to the actual bug (the bug is what value gets treated as "requested", not whether the requested value should be preferred) | +| **Chosen: peek the drafter checkpoint's own `block_size` and use it as the default "requested" value when the operator passes no explicit override** | Fixes the actual defect (a Gemma-4 constant standing in for "no request" on an unrelated family) without touching Gemma 4's default or upstream parity; scoped narrowly via a `model_type == "qwen3_5_mtp"` check so a same-named field on any other drafter shape is never misread | Touches the shared `resolve_draft_block_size` signature (adds a `model_path` parameter), both of its call sites, and adds a new narrow peek function in `mlxcel-core` | + +**Rationale:** the actual defect is upstream of `prefer_requested_block_size` entirely: the "requested" value itself was never actually requested by the operator, it was a flat constant standing in for "unset." Reading the specific checkpoint's own declared value when nothing was explicitly requested addresses the defect at its source and generalizes to any future Qwen 3.5 MTP checkpoint that ships a different `block_size`, rather than hardcoding today's measured value. + +### 3.4 C++ dispatch warning: `eprintln!`, not `tracing::warn!` + +**Rationale:** matches existing precedent in this tree (`report_fused_rope_bypass_once` in `src/models/llama3.rs`): only `mlxcel-server` installs a tracing subscriber unconditionally, the offline `mlxcel` CLI only does so when `RUST_LOG` is set, and the chain-parity path is reachable from both binaries. A `tracing::warn!` here would be silently swallowed on the CLI in the common case. + +--- + +## 4. Implementation Details + +### 4.1 F1: arm-attribution guard (`tests/qwen38_mtp_chain_parity.rs`) + +```rust +fn assert_parity_and_standard_arms_are_distinguishable() { + let t = 3; + let (q, k, v, a, b) = gdn_inputs(t, 0.777); + // ... a_log, dt_bias setup ... + let (_, state_std) = gated_delta_update(&q, &k, &v, &a, &b, &a_log, &dt_bias, None, None); + let (_, state_par) = + gated_delta_update_chain_parity(&q, &k, &v, &a, &b, &a_log, &dt_bias, None, None); + assert_ne!( + raw_bytes(&state_std), + raw_bytes(&state_par), + "arm-attribution guard failed: ... MLXCEL_GDN_CHAIN_PARITY=0 is set ..." + ); +} +``` + +Called at the top of `parity_kernel_cost_vs_standard`, before any timing loop runs. Verified directly: running with `MLXCEL_GDN_CHAIN_PARITY=0` set makes the assertion fail with the expected message; running with the variable unset (the default) passes and the timing loop proceeds. + +### 4.2 F2: capability gate shape check (`src/models/qwen3_5.rs`, `src/server/batch/speculative_burst.rs`) + +```rust +// qwen3_5.rs +pub(crate) fn supports_chain_parity_kernel(&self) -> bool { + qwen35_config_supports_chain_parity_kernel(&self.config) +} + +// speculative_burst.rs +LoadedModel::Qwen35(m) | LoadedModel::Qwen35Moe(m) => { + mlxcel_core::metal_is_available() && m.supports_chain_parity_kernel() +} +LoadedModel::Qwen35VLM(vlm) | LoadedModel::Qwen35MoeVLM(vlm) => { + mlxcel_core::metal_is_available() && vlm.text_model.supports_chain_parity_kernel() +} +``` + +`qwen35_config_supports_chain_parity_kernel` is a free function taking `&Qwen35Config` (not `&self`) specifically so `chain_parity_gate_tests` can exercise it against a bare config, including the tree's own `linear_key_head_dim: 4` fixture shape, without constructing a full model. + +### 4.3 F3: drafter-configured block-size default (`src/lib/mlxcel-core/src/drafter/mod.rs`, `src/cli/speculative_args.rs`) + +```rust +// drafter/mod.rs +pub fn peek_qwen35_mtp_configured_block_size(model_path: &Path) -> Option { + let peek = read_drafter_config_peek(model_path)?; + if peek.model_type.as_deref() != Some(QWEN35_MTP_MODEL_TYPE) { + return None; + } + peek.block_size +} + +// speculative_args.rs +pub fn resolve_draft_block_size( + override_value: Option, + kind: DrafterKind, + model_path: &std::path::Path, +) -> u32 { + if let Some(n) = override_value { + return n; + } + if kind == DrafterKind::Mtp + && let Some(configured) = mlxcel_core::drafter::peek_qwen35_mtp_configured_block_size(model_path) + && let Ok(n) = u32::try_from(configured) + { + return n; + } + default_block_size_for_kind(kind) +} +``` + +### 4.4 F4: C++ silent-forfeiture warning (`src/models/gated_delta.rs`) + +```rust +fn chain_parity_forfeited_by_shape(vectorized_gate: bool, has_mask: bool) -> bool { + vectorized_gate || has_mask +} + +fn warn_if_chain_parity_forfeited_by_shape(g: &MlxArray, mask: Option<&MlxArray>) { + let vectorized_gate = mlxcel_core::array_ndim(g) == 4; + if !chain_parity_forfeited_by_shape(vectorized_gate, mask.is_some()) { + return; + } + static WARNED: std::sync::Once = std::sync::Once::new(); + WARNED.call_once(|| { eprintln!("mlxcel: chain-parity gated-delta kernel was requested but ..."); }); +} +``` + +Called immediately before the `chain_parity` dispatch branch in `gated_delta_ops_with_parity`, mirroring the C++ `if (vectorized && has_mask) ... else if (vectorized) ... else if (has_mask) ... else if (chain_parity)` branch order exactly. + +### 4.5 F5: explicit unmeasured cell (`docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md`) + +Added directly under the 0.591 single-request acceptance figure: + +> **Not measured: acceptance under concurrent slice-grant rotation.** ... `Qwen35MtpDraftModel::reset` ... DESTROYS the drafter's accumulated KV history ... and it fires on every park (tick-slice grant rotation) ... Acceptance under that load regime is therefore structurally worse than the 0.591 measured here, and by how much is uncharacterized. + +### 4.6 F6-F10: documentation and cleanup + +F6 added a `MLXCEL_GDN_CHAIN_PARITY` row to `docs/environment-variables.md`'s speculative-decoding table, matching every sibling MTP variable's documentation depth. F7/F8 corrected the Gemma-4-only and reset-is-no-op claims in `speculative_slice.rs`, `speculative_burst.rs`, `scheduler.rs`, `drafter/mod.rs`, and `speculative/mtp/{mod,generator}.rs`, in each case naming what actually differs between the Gemma 4 and Qwen 3.5 MTP families rather than deleting the distinction. F9 changed `state.verify_out.verify_hidden_full.as_ref()` to `.take()` in the one-shot prompt-prefill hook in `generator.rs`, releasing the ~335 MB seed hidden immediately after its only consumer instead of at the end of the round. F10 removed the dead `if (true) { ... }` wrapper from the new `GATED_DELTA_METAL_SOURCE_SEQPAR` Metal kernel source (pure de-indent, verified bitwise-identical output via `parity_kernel_block_is_bitwise_equal_to_single_token_chain` after a full rebuild). + +--- + +## 5. Premise Correction: the recurrent rollback already existed + +The issue body's implementation plan stated: "The rollback is the hard part: partial acceptance must trim the KV of the 16 full-attention layers and restore the gated-delta recurrent state of the 48 linear layers to the accepted prefix. A recurrent state cannot be trimmed; it must be snapshotted before the draft block and restored on partial accept, or recomputed." A comment on the issue, filed before implementation began, corrected this: `Qwen35Model::rollback_speculative_cache` already implements exactly this for the DFlash path (attention layers trimmed by `block_size - (max(accepted) + 1)` with per-row tail zeroing for `B > 1`; gated-delta layers replayed from `GdnRollbackSnapshot` over the accepted prefix), and the verify pass already returns the hidden states an MTP drafter needs (`VerifyOutput { logits, hidden_states, gdn_states }`). + +The corrected scope, wiring rather than cache surgery, is what actually shipped: `5bb9bd1d` adds the drafter and the `MtpTarget` adapter wrapping the existing verify/rollback surface without adding new recurrent-state machinery; `89bcbae6` lifts the Gemma-4-only gates. Both commits are consistent with the corrected premise, not the original one. Recording this here follows the project's own convention (`docs/code-guidelines.md` / CLAUDE.md performance-issue rules): a corrected premise, caught before implementation, is the deliverable for that half of the issue, not a footnote. + +--- + +## 6. Change Summary + +### Statistics (this finalization pass) + +| Item | Value | +|------|-------| +| Files changed | 18 | +| Lines added / removed | +642 / -126 (approximate, pre-final-commit) | +| New tests | 5 (`resolve_draft_block_size` / `peek_qwen35_mtp_configured_block_size`), 3 (`chain_parity_gate_tests`), 2 (`chain_parity_forfeited_by_shape`), 1 (arm-attribution guard) = 11 new test functions | +| New env var documented | 1 (`MLXCEL_GDN_CHAIN_PARITY`) | +| Benchmark record numbers re-measured | 4 (isolated kernel-cost ratios at T=1/3/4/64) | + +### Changes by Category + +| Category | Count | Summary | +|----------|-------|---------| +| Test/measurement integrity | 1 | F1: arm-attribution guard, benchmark re-measured | +| Correctness hardening | 2 | F2: capability gate shape check; F4: C++ silent-forfeiture warning | +| Default-value correctness | 1 | F3: drafter-configured block size | +| Documentation honesty | 1 | F5: explicit unmeasured-cell note | +| Documentation accuracy | 3 | F6: env var doc; F7/F8: ~10 stale Gemma-4-only / reset-no-op claims | +| Memory / cleanup | 2 | F9: early-clear the seed hidden; F10: dead kernel scaffolding | + +### Related Commits + +| Hash | Type | Message | +|------|------|---------| +| `5bb9bd1d` | feat | add the qwen3_5_mtp drafter and Qwen 3.5 MTP target adapter (parent PR) | +| `89bcbae6` | feat | wire Qwen 3.5 targets into the MTP burst, slice, and offline paths (parent PR) | +| `9575bad4` | test | add Qwen 3.5 MTP chain-parity diagnostic against the single-token chain (parent PR) | +| `6fce98ba` | fix | make the Qwen 3.5 verify block bit-identical to the decode chain (parent PR) | +| `f6d075be` | perf | gate Qwen MTP on Metal, measure the pairing, record the hardware verdict (parent PR) | +| `60ef14b2` | test | pin MTP dispatch-site coverage and fix the chain-parity test lint (review pass, 2 HIGH) | +| `ef3aa95e` | fix | bound the qwen3_5_mtp drafter config scalars at load (security pass, 1 HIGH) | +| `a88af1ff` | fix | harden Qwen 3.5 MTP exactness gate and defaults (this pass, F1-F4) | +| `775afe95` | fix | correct stale MTP docs and trim a held-too-long buffer (this pass, F6-F10) | + +--- + +## 7. Follow-up Actions + +### Required + +- [ ] None from this pass; all ten findings the review raised are addressed. + +### Future Improvements (recorded as known limitations, not fixed here) + +- **M5-class hardware is unmeasured for this pairing.** The Gemma 4 precedent (verify amortization holds, 1.2-1.4x on M5 Max) projects a positive sign, but no M5 machine was available. The issue thread includes a reproduction guide for filling this cell. +- **Acceptance under concurrent slice-grant rotation is uncharacterized.** F5 documents this as an explicit open cell rather than fixing it; filling it requires a multi-client load harness exercising the tick-slice grant-rotation path specifically, which is out of scope for this pass. +- **B > 1 batched MTP remains Gemma-4-only.** Qwen 3.5 MTP is B = 1 only by design (issue #1165's own scope); the batched adapter methods return `DrafterError::DraftFailed` per the trait default for this family, matching the documented limitation. +- Splitting an MTP head out of a raw HF checkpoint (so users are not solely dependent on `mlx-community/Qwen3.8-27B-MTP-bf16` continuing to exist) remains a deliberately deferred follow-up, tracked in the issue thread. + +--- + +## Appendix + +### A. Test Results + +- `cargo test --release --features metal,accelerate --test qwen38_mtp_chain_parity` (non-ignored): 1 passed (`parity_kernel_block_is_bitwise_equal_to_single_token_chain`), 5 ignored (real-checkpoint gated), 0 failed. Verified again after the F10 kernel-source cleanup (full LTO rebuild). +- `cargo test --release --features metal,accelerate --test qwen38_mtp_chain_parity -- --ignored --nocapture parity_kernel_cost`: 1 passed, five repeated runs recorded in the updated benchmark record. +- `cargo test --release --features metal,accelerate -p mlxcel-core drafter::qwen3_5_mtp`: 20 passed, 0 failed. +- `cargo test --release --features metal,accelerate -p mlxcel-core drafter::` (gated_delta shape-gate and chain-parity-forfeiture tests included via the `mlxcel` crate's own `models::gated_delta` module): 6 passed, 11 ignored (serial-MLX gated), 0 failed. +- `cargo fmt --check`: clean. +- Manual sanity check of the arm-attribution guard: `MLXCEL_GDN_CHAIN_PARITY=0 cargo test ... parity_kernel_cost -- --ignored --nocapture` panics with the expected message; unset (default) passes. + +### B. References + +- Issue #1165 (specification, premise-correction comment, hardware reproduction guide) +- `docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md` (measurement record, updated by this pass) +- `docs/environment-variables.md` (`MLXCEL_GDN_CHAIN_PARITY`, new) +- `src/models/gated_delta.rs` (chain-parity kernel dispatch, arm-attribution-adjacent warning) +- `src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp` (`GATED_DELTA_METAL_SOURCE_SEQPAR`) +- `src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/` (drafter model, config, tests) +- `src/models/qwen3_5_mtp_target.rs` (`MtpTarget` adapter) +- CLAUDE.md, "Performance issue completion criteria" (arm-attribution rule, epic #909 precedent) diff --git a/TECHNICAL_REPORTS/1182-qwen35-mtp-drafter-20260816.ko.md b/TECHNICAL_REPORTS/1182-qwen35-mtp-drafter-20260816.ko.md new file mode 100644 index 000000000..b6d1d7a14 --- /dev/null +++ b/TECHNICAL_REPORTS/1182-qwen35-mtp-drafter-20260816.ko.md @@ -0,0 +1,315 @@ +# 기술 보고서: PR #1182 - feat(inference): support the qwen3_5_mtp drafter for Qwen 3.5 MTP speculative decoding + +**작성일**: 2026-08-16 +**작성자**: AI Code Reviewer +**상태**: 완료 +**언어**: Rust, C++ (Metal 커널 소스) +**위험도**: Low (지금까지 측정한 모든 하드웨어 세대에서 이 페어링은 기본값으로 꺼져 있고, 지원하지 않는 조합을 만나면 모든 디스패치 지점이 정확성 계약을 조용히 어기는 대신 classic decode로 물러난다) + +--- + +## 요약 + +PR #1182(이슈 #1165)는 별도로 분리된 `qwen3_5_mtp` drafter(예: `mlx-community/Qwen3.8-27B-MTP-bf16`)를 정식 MTP drafter로 추가하고, Qwen 3.5 계열(dense와 MoE, text와 VLM 모두)에 mlxcel의 `MtpTarget` 트레이트를 구현한다. 그 결과 기존 `MtpGenerator` 라운드 루프와 tick-cooperative slice 경로, run-to-completion burst가 모두 Gemma 4를 몰던 것과 똑같은 방식으로 이 페어링을 몬다. 이 결과물이 제대로 동작하는 이유는 구체적인 수치 수정 하나 덕분이다. temperature-0 정확성 게이트가 처음에는 100~250 토큰마다 한 번꼴로 near-tie argmax가 뒤집히며 실패했는데, 원인을 추적해보니 공유 gated-delta Metal 커널이 `T = K` verify 블록 전체에 걸쳐 recurrent state를 float32로 들고 있다가 마지막 저장 시점에만 storage dtype으로 양자화하는 반면, classic한 single-token decode chain은 토큰마다 매번 양자화한다는 데 있었다. 그래서 `T = K` 블록은 그것이 재현해야 할 `K`번의 연속된 `T = 1` decode step과 비트 단위로 같지 않았고, 마지막 ulp 하나의 state 차이가 near-tie logit을 뒤집었다. 수정은 in-block step이 끝날 때마다 state를 storage dtype으로 round하는 chain-parity 커널 변형(`gated_delta_step_seqpar`)이고, 이렇게 하면 두 형태가 증명 가능하게 비트 단위로 같아진다(`parity_kernel_block_is_bitwise_equal_to_single_token_chain`으로 고정했고 240 토큰 전체 decode chain에서도 확인했다). 이 커널은 speculative verify capture 경로와 rollback replay에서만 쓰이므로 classic prefill과 decode는 손대지 않고, DFlash도 같은 `forward_speculative` 경로를 공유하는 덕분에 같은 수치를 별도 작업 없이 물려받는다. + +이 PR을 검증할 수 있었던 하드웨어 한 대에서는 측정된 하드웨어 판정이 부정적이다. Apple M1 Ultra에서 MTP decode는 drafter 자신이 설정한 block size(3)에서 classic 대비 **0.70배**, 역사적으로 쓰이던 block-4 기본값에서는 0.59배를 기록한다. 이건 구현의 결함이 아니다. 라운드 루프의 모든 진단 지표는 건강하다. 원인은 하드웨어 한계다. `verify_forward_cost_scaling`을 보면 target 자신의 평범한 multi-token forward가 이 GPU 세대의 블록 크기에서는 거의 amortize되지 않는다(블록 위치 하나가 늘 때마다 약 29ms, 이는 verify 경로와 평범한 batched-prefill 경로 양쪽에서 똑같이, 완전한 decode step 하나의 3분의 2 정도를 차지한다). 이건 `mtp_b1_default`에 이미 기록되어 있는 Gemma 4 MTP의 pre-M5 회귀와 같은 현상이다(이슈 #165: M1 Ultra 0.75~0.96배, M5 Max 1.2~1.4배). adaptive MTP 정책(이슈 #333)이 Qwen 3.5 MTP를 정확히 같은 정적 기본값과 프로파일링 메커니즘으로 게이트하므로, 이 하드웨어군에 배포된 서버는 조용히 성능이 떨어지는 대신 기본값으로 classic decode를 선택한다. 이 페어링의 M5 셀(그리고 M3/M4)은 아직 비어 있고, 이슈 스레드에는 다른 하드웨어를 가진 사람이 이 셀을 채울 수 있도록 재현 안내가 함께 실려 있다. + +세 번째로 이 PR을 떠받치는 사실은 구현이 시작되기 전에 이뤄진 전제 교정이다. 이슈 본문은 애초에 "어려운 부분"을, full-attention 레이어 16개와 gated-delta linear 레이어 48개가 섞인 target을 위한 partial-acceptance rollback을 새로 만드는 일로 잡아뒀다. recurrent state는 자를 수 없으니 draft 블록 이전에 스냅샷을 찍고 partial accept 시점에 복원하거나 재계산해야 한다는 논리였다. 그런데 이 장치는 이미 존재했다. `Qwen35Model::rollback_speculative_cache`가 DFlash 경로를 위해 정확히 이 일을 이미 하고 있었다. attention KV는 미승인 꼬리만큼 자르고, gated-delta 레이어는 승인된 접두사 위에서 `GdnRollbackSnapshot`을 replay하는 식이다. verify pass도 MTP drafter가 필요로 하는 hidden state를 이미 돌려주고 있었다. 그래서 이 이슈에서 진짜 새로 필요했던 일은 캐시 수술이 아니라 배선이었다. 새 drafter 모델 하나, 기존 verify/rollback 표면을 감싸는 `MtpTarget` 트레이트 구현, 그리고 서버와 offline CLI 디스패치 경로 곳곳에 흩어진 Gemma-4-전용 게이트 열 곳 남짓을 걷어내는 작업이었다. + +최초 PR에 대한 리뷰와 보안 검토는 이번 finalization 이전에 HIGH 두 건(커밋 `60ef14b2`)과 HIGH 한 건(커밋 `ef3aa95e`)을 이미 고쳤다. 이번 finalization은 리뷰가 지적한 열 건을 추가로 닫았다. 자기 자신과 비교하는 노이즈를 결과로 보고할 수 있었던 벤치마크, Metal 가용성만 확인하고 chain-parity 커널의 실제 shape 계약은 확인하지 않던 capability gate, PR 자신의 측정으로 이미 더 느리다고 드러난 설정을 기본값으로 출하하고 있던 block size, 오늘은 도달할 수 없는 shape에서 chain-parity 요청을 조용히 무시하는 C++ 디스패치 경로, 단일 요청 수치를 측정치인 것처럼 보고하면서 동시 부하 아래 acceptance는 측정하지 않은 채 남겨둔 벤치마크 기록, 빠져 있던 환경변수 문서 항목, Gemma-4-전용이라고 낡게 서술된 문서 열 곳 남짓, Qwen drafter의 `reset`이 no-op라는 이제는 틀린 서술 세 곳, 필요 이상으로 한 라운드 더 붙들려 있던 335MB짜리 seed hidden state, 새 Metal 커널 소스에 복사되어 들어온 죽은 `if (true)` 스캐폴딩이 그 대상이다. + +--- + +## 1. 문제 정의 + +### 1.1 배경 + +이슈 #1165는 Qwen3.8-27B의 MTP head가 어떤 MLX 변환본에서도 target 체크포인트 안에 들어 있지 않다는 사실을 먼저 확인했다. 업스트림 `Qwen/Qwen3.8-27B` base는 `model.*`와 나란히 `mtp.*` 텐서 15개를 담고 있지만, `mlx-community/Qwen3.8-27B-*` 변환본은 모두 이를 빼버리고, mlx-community는 이를 별도의 drafter(`mlx-community/Qwen3.8-27B-MTP-bf16`, `model_type: "qwen3_5_mtp"`, `block_size: 3`)로 다시 배포한다. 이 PR 이전에는 `mlxcel generate --draft-model <해당 체크포인트>`가 `--draft-kind`와 무관하게 `Unsupported model type: qwen3_5_mtp`로 실패했고, 두 메시지 모두 정확했고 fail-closed였다. 그러니 이건 잠재적 버그가 아니라 빠져 있던 기능이었다. + +### 1.2 기존 문제점 (리뷰·보안 검토에서 나온 항목, 이번 finalization에서 해결) + +- **F1** (High, 측정 신뢰성): `parity_kernel_cost_vs_standard`에는 arm-attribution 가드가 없었다. 환경에 `MLXCEL_GDN_CHAIN_PARITY=0`이 설정되어 있으면 `gated_delta_update_chain_parity`는 조용히 standard 커널로 물러나는데(`chain_parity_enabled()`가 그 환경변수를 읽는다), 그러면 벤치마크의 두 arm이 같은 코드 경로로 귀결되어 자기 자신과 비교하는 결과가 나온다. 기록되어 있던 비율(1.089 / 1.117 / 0.834 / 1.008)은 이 실패 모드와 구분할 수 없었다. CLAUDE.md의 "숫자를 보고하기 전에 두 arm이 실제로 다르다는 것을 증명하라"는 규칙이 정확히 겨냥하는 패턴이고, 에픽 #909는 같은 결함을 가진 벤치마크 세 건을 이미 이름까지 붙여 남겨뒀다. +- **F2** (High, 정확성): `mtp_capable_target`은 Qwen 3.5 계열을 `metal_is_available()` 하나만으로 게이트했다. 비트 단위 동일성은 여기에 더해 `supports_metal_gated_delta_kernel(hk, hv, dk, dv)`(`dk >= 32 && dk % 32 == 0 && hv % hk == 0`)를 요구한다. `linear_key_head_dim`이 이 계약과 맞지 않는 체크포인트를 올린 Metal 호스트는 아무 신호 없이 ops 기반 fallback으로 조용히 빠지면서 temperature-0 정확성 계약을 잃는다. 이 트리 자신의 테스트 fixture가 이미 `linear_key_head_dim: 4`를 쓰고 있으니, 이건 가상의 형태가 아니다. +- **F3** (Medium, 기본값의 정확성): Qwen drafter는 `prefer_requested_block_size() = true`라서 drafter 자신의 `configured_block_size()`(3)는 한 번도 참조되지 않았다. Gemma 4에서 물려받은 상수인 per-kind CLI 기본값(4)이 항상 이겼다. PR 자신의 측정은 block 3가 16.48 tok/s / 0.591 acceptance, block 4가 13.81 / 0.465임을 보여주므로, 출하되던 기본값은 측정으로 확인된 두 설정 중 더 느린 쪽이었다. +- **F4** (Medium, 심층 방어): `metal_gated_delta_forward_impl`의 C++ 디스패치는 gate가 vectorized이거나 mask가 있으면, 호출자가 `chain_parity = true`를 요청했더라도 조용히 standard 커널 분기를 탄다. 오늘 Qwen 3.5 verify 경로에서는 도달할 수 없는 형태이지만(`mask = None`이고 `g`는 항상 3차원이다), 이를 강제하는 장치는 아무것도 없었다. 미래의 호출자가 아무 신호 없이 비트 단위 동일성을 잃을 수 있는 셈이다. +- **F5** (Medium, 문서의 정직성): 벤치마크 기록의 0.591 acceptance 수치는 단일 요청만 측정한 값이다. `Qwen35MtpDraftModel::reset`이 Gemma 4의 no-op reset과 달리 drafter가 쌓아둔 KV history를 파괴하며, 동시 slice-grant rotation 아래에서 park될 때마다 이게 일어난다는 사실은 서술되어 있지 않았다. 그래서 부하 아래 acceptance는 구조적으로 더 나쁠 수밖에 없는데, 그 정도는 특성화되지 않은 채 남아 있었다. +- **F6~F10** (Low, 문서와 정리): `docs/environment-variables.md`에 `MLXCEL_GDN_CHAIN_PARITY` 항목이 빠져 있었다. `speculative_slice.rs`, `speculative_burst.rs`, `scheduler.rs`, `drafter/mod.rs`, `speculative/mtp/{mod,generator}.rs`에 걸쳐 모듈·API 문서 열 곳 남짓이 여전히 MTP를 Gemma-4-전용이라고 서술하고 있었다. 그중 세 곳은 MTP `Drafter::reset`이 보편적으로 트레이트 기본값인 no-op라고도 주장했는데, 이는 Qwen drafter에는 해당하지 않는다. 시드 라운드의 `verify_hidden_full`(`[1, P, H]`, P=32k/H=5120 f16 기준 약 335MB)이 유일한 소비처가 쓰고 난 뒤에도 즉시 해제되지 않고 라운드 하나 전체 동안 살아 있었다. 그리고 새 chain-parity Metal 커널 소스에는 masked 변형에서 그대로 복사되어 들어온 죽은 `if (true) { ... }` 래퍼가 남아 있었다. + +### 1.3 위험성 + +| 위험 | 영향도 | 발생 가능성 | +|-----|-------|-----------| +| 벤치마크 수치가 자기 자신과 비교한 노이즈와 구분되지 않는데 실제 비율처럼 보고됨 (F1) | Medium (향후 용량·성능 판단을 오도, 사용자에게 직접 영향은 없음) | 재현 확인됨: `MLXCEL_GDN_CHAIN_PARITY=0`을 설정하고 기존 테스트를 재실행하면 두 arm이 그대로 같아진다 | +| GDN shape가 맞지 않는 Metal 호스트가 조용히 temperature-0 정확성을 잃음 (F2) | High, 실제로 일어나면 (문서화된 보증 아래에서 조용히 틀린 생성 결과) | 오늘은 낮음(공개된 `mlx-community` Qwen 3.5 체크포인트는 모두 `dk = 128`이다), 다만 게이트가 이를 확인하지 않고 있었다 | +| 측정으로 확인된 두 block size 설정 중 더 느린 쪽을 조용한 기본값으로 출하 (F3) | Medium (플래그를 수동으로 바꾸지 않는 한 운영자가 알아채지 못할 처리량 저하) | 수정 전에는 확실했다. 설정하지 않은 모든 Qwen 3.5 MTP 요청이 block 4를 썼다 | +| C++가 vectorized/masked shape에서 `chain_parity`를 조용히 무시함 (F4) | High, 실제로 일어나면 (F2와 같은 정확성 상실) | 오늘은 매우 낮음(현재 어떤 호출자도 이 경로에 도달하지 않는다), 미래의 호출자에게는 실재 | +| 동시 부하 아래 acceptance가 단일 요청 수치로 암묵적으로 대체되어 보고됨 (F5) | Low-Medium (동시 MTP 용량을 산정하려는 운영자가 낙관적인 수치를 받는다) | 확실함: 트리 어디에도 동시 rotation 측정이 없다 | +| 낡은 문서 (F6~F10) | Low (독자·관리자 혼란, 기능 영향 없음) | 이번 수정 전에는 확실했고, 이번에 바로잡았다 | + +--- + +## 2. 기술적 검토 사항 + +### 2.1 보안 관점 + +리뷰와 보안 검토는 오케스트레이터와 `pr-reviewer`, `pr-security-checker`가 이번 finalization 이전에 이미 완료했다. HIGH 항목 두 건은 `60ef14b2`(MTP 디스패치 지점 커버리지 테스트, chain-parity 테스트 lint)에서, HIGH 항목 한 건은 `ef3aa95e`(`qwen3_5_mtp` drafter의 config 스칼라 값을 로드 시점에 0으로 나누기와 무제한 할당 abort로부터 지켜냄. drafter는 batch worker 안에서 지연 로드되는데, 이 코어 워커 스레드는 잡히지 않은 panic을 프로세스 abort로 바꿔버리기 때문이다)에서 이미 고쳤다. 이번 단계는 그 둘을 다시 다루지 않았고, 위 1.2절의 열 건을 닫았다. + +**이번 단계에서 발견된 이슈:** + +| 이슈 | 심각도 | 상태 | +|-----|-------|-----| +| F1: `parity_kernel_cost_vs_standard`에 arm-attribution 가드 없음 | High | Fixed | +| F2: `mtp_capable_target`이 Metal 가용성만 보고 parity 커널의 shape 계약은 확인하지 않음 | High | Fixed | +| F3: 출하된 기본 block size가 측정상 더 느린 설정 | Medium | Fixed | +| F4: C++ 디스패치가 vectorized/masked shape에서 `chain_parity`를 조용히 무시 | Medium | Fixed | +| F5: 부하 아래 acceptance가 벤치마크 기록에 특성화도, 표기도 없이 빠져 있음 | Medium | Fixed (명시적 미측정 셀로 기록) | +| F6~F10: 환경변수 문서, Gemma-4-전용 서술 약 10곳, `reset`이 no-op라는 틀린 서술 3곳, 너무 오래 붙들린 seed hidden, 죽은 커널 스캐폴딩 | Low | Fixed | + +### 2.2 성능 관점 + +측정된 판정(요약 절과 `docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md` 참고)은 이 PR을 검증할 수 있었던 유일한 하드웨어 세대에서 부정적이고, F3의 수정은 그중 어떤 부정적 숫자가 기본값으로 나가는지를 바꾼다(수정 전 block-4 기본값에서 0.59배였던 것이 block 3에서 0.70배가 된다). 둘 다 adaptive MTP 정책이 참여 여부를 판단하는 데 쓰는 1.0배 손익분기점 아래이므로, F3의 실질 효과는 정책이 우회되는 경우(`MLXCEL_ENABLE_MTP_B1`을 켜서 고정했거나 프로파일링을 끈 경우)에 더 나은 정보로 채워진 기본값이지, 기본 참여 여부 자체가 바뀌는 것은 아니다. + +F1이 arm-attribution 가드를 추가한 뒤 다시 측정한 결과는 다섯 번의 실행에서 T=1/3/4/64 기준 중앙값 0.979 / 1.007 / 0.994 / 1.022(실행 간 min-max 0.877~1.043)로 나왔다. 원래(귀속되지 않은) 수치가 주장하던 것과 마찬가지로 dispatch 노이즈 대역 안에 있지만, 이번에는 증명 가능하게 귀속된 수치다. 가드는 `T=3`에서 두 커널을 한 번씩 실행해서 최종 GDN state가 서로 다른지 단언하고(두 커널은 `T > 1`이면 항상 다르도록 설계되어 있다), 이 단언이 실패하면 아무 것도 출력하기 전에 패닉한다. 교정된 수치는 원래 기록된 값과 개별 비율 수준에서는 다르다(이는 예상된 일이다. 원래 수치는 애초에 귀속 가능하지 않았고, 부하가 걸린 박스에서의 dispatch 수준 노이즈는 실행마다 비율 단위로 재현되지 않는다). 다만 정성적 결론, 즉 커널이 추가한 per-step register round 하나가 노이즈 바깥에서 측정 가능한 비용을 만들지 않는다는 결론은 바뀌지 않았다. + +### 2.3 호환성/의존성 관점 + +- **Breaking Changes**: 없다. `resolve_draft_block_size`의 시그니처에 `model_path: &Path` 매개변수가 추가됐다(F3의 수정). 이 함수의 호출 지점 두 곳(`src/server/speculative_dispatch.rs`, `src/commands/generate.rs`) 모두 그 시점에 이미 drafter 경로를 갖고 있다. `supports_metal_gated_delta_kernel`은 private에서 `pub(crate)`로 바뀌었지만(F2), 외부에 노출되는 표면은 그대로다. +- **새로운 의존성**: 없다. +- **동작 변화**: `--draft-block-size`를 명시하지 않고 `mlxcel generate --draft-model --draft-kind mtp`(또는 서버에서 같은 조합)를 실행하는 운영자는 이제 F3에 따라 4가 아니라 block size 3를 받는다. 명시적으로 `--draft-block-size`를 준 경우는 영향이 없다. GDN shape가 parity 커널의 계약을 만족하지 않는 Metal 호스트는 이제 비트 단위 동일성 없이 조용히 도는 대신 classic decode로 물러난다(F2). 공개된 체크포인트 중에는 오늘 이 영향을 받는 것이 없다. + +### 2.4 코드 품질 관점 + +- **테스트 커버리지**: 이번 단계에서 `assert_parity_and_standard_arms_are_distinguishable`(F1 가드, `tests/qwen38_mtp_chain_parity.rs`), `Qwen35Model::supports_chain_parity_kernel` / `qwen35_config_supports_chain_parity_kernel`을 위한 테스트 세 건(F2, `src/models/qwen3_5.rs`), `mlxcel`과 `mlxcel-core` 두 크레이트에 걸친 `resolve_draft_block_size` / `peek_qwen35_mtp_configured_block_size` 테스트 다섯 건(F3), `chain_parity_forfeited_by_shape_matches_the_cpp_dispatch_branches`와 그 아래 순수 predicate(F4, `src/models/gated_delta.rs` / `_tests.rs`)를 추가했다. +- **코드 복잡도**: 각 수정은 국지적이다. F2는 GDN-shape 검사를 `Qwen35Model::supports_chain_parity_kernel`로 뽑아냈는데, 이 함수 자체는 bare `&Qwen35Config`를 받는 자유 함수(`qwen35_config_supports_chain_parity_kernel`)를 감싸는 얇은 wrapper다. 굳이 분리한 이유는 실제 가중치로 전체 모델을 만들지 않고도 단위 테스트로 검증할 수 있게 하기 위해서다. F4도 마찬가지로 C++ shape predicate를 순수 함수 `chain_parity_forfeited_by_shape(bool, bool) -> bool`로 뽑아내서, Metal 장치 없이도 경고 조건을 테스트할 수 있게 했다. +- **기술 부채**: 감소했다. 벤치마크 기록은 이제 숫자를 그냥 단언하는 대신 자신의 귀속 메커니즘을 스스로 서술한다. capability gate와 CLI 기본값 둘 다 이제는 관련 없는 drafter family에 일괄 적용되던 Gemma-4 유래 상수 대신 체크포인트 자신이 선언한 shape와 config를 반영한다. 그리고 다음 기여자를 Qwen 3.5 MTP 지원이 아예 없다거나 그 drafter의 reset이 no-op라고 오도할 뻔했던 문서 지점 열 곳 남짓이 바로잡혔다. + +--- + +## 3. 기술적 선택과 그 이유 + +### 3.1 arm-attribution 가드: 표로 출력하기 전에 단언하고 실패하기, 출력에 주석을 다는 방식이 아니라 + +**고려한 대안:** + +| 옵션 | 장점 | 단점 | +|-----|-----|-----| +| 귀속을 확신할 수 없을 때 숫자 옆에 경고를 출력 | 테스트가 계속 돌고 숫자도 그대로 보임 | CLAUDE.md 규칙이 정확히 겨냥하는 실패 모드다. 주석이 붙어 있어도 숫자는 그대로 인용될 수 있다 | +| **선택: `T=3`에서 두 arm이 증명 가능하게 같으면 어떤 비율도 출력하기 전에 패닉** | 오도할 수 있는 숫자가 보고서에 도달할 경로 자체가 없음. 실패 메시지가 유력한 원인(`MLXCEL_GDN_CHAIN_PARITY=0`)을 지목함 | `--ignored --nocapture` 호출 전체가 이 검사 실패 시 우아하게 저하되는 대신 중단됨 | + +**선택 이유**: 리뷰는 "귀속에 실패하면 테스트가 보고를 거부하게 만들라"고 명시적으로 요구했다. 경고하고 계속 진행하는 설계는 여전히 복사-붙여넣기된 숫자가 문서에 들어갈 여지를 남긴다. `T = 1`은 판별자로 쓸 수 없다(두 커널은 그 지점에서 비트 단위로 같도록 설계되어 있다). 그래서 이 검사는 `T = 3`에서 한 번만 도는데, 이는 자매 테스트인 `parity_kernel_block_is_bitwise_equal_to_single_token_chain`이 이미 비트 단위로 다르다고 고정해둔 형태와 일치한다. + +### 3.2 capability gate: shape predicate를 경고가 아니라 `mtp_capable_target`에 접어넣기 + +**고려한 대안:** + +| 옵션 | 장점 | 단점 | +|-----|-----|-----| +| parity가 요청됐는데 shape가 맞지 않으면 일회성 `tracing::warn!`을 내고 계속 서빙 | 코드 변경이 최소 | 정확성 계약은 "temperature-0 출력이 classic decode와 비트 단위로 같다"는 절대적인 문구로 문서화되어 있다. 경고 하나로는 틀린 답이 클라이언트에 닿는 것을 막지 못한다 | +| **선택: `supports_metal_gated_delta_kernel`을 `mtp_capable_target` 자체에 접어넣기** | 계약을 지킬 수 없는 체크포인트가 자동으로 classic decode로 물러난다. 이 코드베이스의 다른 모든 MTP 게이트가 이미 쓰고 있는 fail-closed 자세와 같다 | `mlxcel-core`에서 shape predicate를 노출하고, `Qwen35Model`에 config를 읽는 접근자를 새로 붙여야 한다 | + +**선택 이유**: 리뷰는 두 옵션을 모두 제시했고, "더 나은 쪽"으로 predicate를 게이트에 접어넣는 방식을 명시적으로 이름 붙였다. 이 트리의 다른 모든 MTP capability 검사(Metal 가용성, target/drafter 호환성, hidden-size와 vocab 일치)가 이미 경고가 아니라 거부로 동작하는 만큼, 여기서만 경고로 처리하면 그 자체로 일관성 있던 fail-closed 설계의 유일한 예외가 될 뻔했다. + +### 3.3 block size 기본값: 공유 상수를 바꾸는 대신 drafter 자신의 config를 들여다보기 + +**고려한 대안:** + +| 옵션 | 장점 | 단점 | +|-----|-----|-----| +| `DEFAULT_MTP_BLOCK_SIZE`를 4에서 3으로 변경 | 한 줄짜리 수정 | `DrafterKind::Mtp`는 Gemma 4와 Qwen 3.5를 모두 포괄한다. 이렇게 하면 Qwen 3.5에서만 측정한 근거로 Gemma 4의 기본값까지 바뀐다 | +| Qwen drafter의 `prefer_requested_block_size()`를 `false`로 뒤집기 | 기존의 adaptive `effective_mtp_block_size` 장치를 재사용 | 업스트림 mlx-vlm은 이 drafter에 한해 구체적으로 `prefer_requested_block_size = True`를 설정한다. 이를 뒤집으면 실제 버그와는 무관한 이유로 그 parity가 깨진다(버그는 "요청된" 값을 우선할지 여부가 아니라, 무엇이 "요청된" 값으로 취급되는지에 있다) | +| **선택: drafter 체크포인트 자신의 `block_size`를 확인해서, 운영자가 명시적 override를 주지 않았을 때 그 값을 기본 "요청" 값으로 쓰기** | Gemma 4의 기본값이나 업스트림 parity를 건드리지 않고 실제 결함(관련 없는 family에 Gemma-4 상수가 "요청 없음"을 대신하고 있던 문제)을 고침. `model_type == "qwen3_5_mtp"`일 때만 좁게 확인하므로 다른 drafter shape에 같은 이름의 필드가 있어도 잘못 읽히지 않음 | 공유 함수 `resolve_draft_block_size`의 시그니처(경로 매개변수 추가)와 그 호출 지점 두 곳을 건드리고, `mlxcel-core`에 새 peek 함수를 하나 추가해야 함 | + +**선택 이유**: 실제 결함은 `prefer_requested_block_size`보다 한 단계 앞에 있다. "요청된" 값 자체가 운영자가 실제로 요청한 적이 없는 값이었다. 그건 그냥 "설정 안 함"을 대신하던 고정 상수였다. 아무 것도 명시적으로 요청되지 않았을 때 그 특정 체크포인트가 스스로 선언한 값을 읽으면 결함을 근원에서 고칠 수 있고, 오늘 측정된 값을 하드코딩하는 대신 앞으로 다른 `block_size`를 선언하는 어떤 Qwen 3.5 MTP 체크포인트에도 자연스럽게 일반화된다. + +### 3.4 C++ 디스패치 경고: `tracing::warn!`이 아니라 `eprintln!` + +**선택 이유**: 이 트리의 기존 선례(`src/models/llama3.rs`의 `report_fused_rope_bypass_once`)를 따른다. tracing subscriber를 무조건 설치하는 건 `mlxcel-server`뿐이고, offline `mlxcel` CLI는 `RUST_LOG`가 설정되어 있을 때만 설치한다. chain-parity 경로는 두 바이너리 모두에서 도달 가능하다. 여기서 `tracing::warn!`을 쓰면 흔한 경우인 CLI에서 조용히 삼켜진다. + +--- + +## 4. 구현 상세 + +### 4.1 F1: arm-attribution 가드 (`tests/qwen38_mtp_chain_parity.rs`) + +```rust +fn assert_parity_and_standard_arms_are_distinguishable() { + let t = 3; + let (q, k, v, a, b) = gdn_inputs(t, 0.777); + // ... a_log, dt_bias 준비 ... + let (_, state_std) = gated_delta_update(&q, &k, &v, &a, &b, &a_log, &dt_bias, None, None); + let (_, state_par) = + gated_delta_update_chain_parity(&q, &k, &v, &a, &b, &a_log, &dt_bias, None, None); + assert_ne!( + raw_bytes(&state_std), + raw_bytes(&state_par), + "arm-attribution guard failed: ... MLXCEL_GDN_CHAIN_PARITY=0 is set ..." + ); +} +``` + +`parity_kernel_cost_vs_standard`의 맨 앞, 어떤 타이밍 루프도 돌기 전에 호출한다. 직접 확인한 바로는 `MLXCEL_GDN_CHAIN_PARITY=0`을 설정하면 예상한 메시지와 함께 이 단언이 실패하고, 설정하지 않으면(기본값) 통과해서 타이밍 루프가 이어진다. + +### 4.2 F2: capability gate의 shape 검사 (`src/models/qwen3_5.rs`, `src/server/batch/speculative_burst.rs`) + +```rust +// qwen3_5.rs +pub(crate) fn supports_chain_parity_kernel(&self) -> bool { + qwen35_config_supports_chain_parity_kernel(&self.config) +} + +// speculative_burst.rs +LoadedModel::Qwen35(m) | LoadedModel::Qwen35Moe(m) => { + mlxcel_core::metal_is_available() && m.supports_chain_parity_kernel() +} +LoadedModel::Qwen35VLM(vlm) | LoadedModel::Qwen35MoeVLM(vlm) => { + mlxcel_core::metal_is_available() && vlm.text_model.supports_chain_parity_kernel() +} +``` + +`qwen35_config_supports_chain_parity_kernel`은 `&self`가 아니라 `&Qwen35Config`를 받는 자유 함수다. `chain_parity_gate_tests`가 실제 모델을 만들지 않고도 이 트리 자신의 `linear_key_head_dim: 4` fixture 형태를 포함해 bare config만으로 이 함수를 검증할 수 있게 하려는 목적이다. + +### 4.3 F3: drafter가 설정한 block-size 기본값 (`src/lib/mlxcel-core/src/drafter/mod.rs`, `src/cli/speculative_args.rs`) + +```rust +// drafter/mod.rs +pub fn peek_qwen35_mtp_configured_block_size(model_path: &Path) -> Option { + let peek = read_drafter_config_peek(model_path)?; + if peek.model_type.as_deref() != Some(QWEN35_MTP_MODEL_TYPE) { + return None; + } + peek.block_size +} + +// speculative_args.rs +pub fn resolve_draft_block_size( + override_value: Option, + kind: DrafterKind, + model_path: &std::path::Path, +) -> u32 { + if let Some(n) = override_value { + return n; + } + if kind == DrafterKind::Mtp + && let Some(configured) = mlxcel_core::drafter::peek_qwen35_mtp_configured_block_size(model_path) + && let Ok(n) = u32::try_from(configured) + { + return n; + } + default_block_size_for_kind(kind) +} +``` + +### 4.4 F4: C++의 조용한 계약 상실에 대한 경고 (`src/models/gated_delta.rs`) + +```rust +fn chain_parity_forfeited_by_shape(vectorized_gate: bool, has_mask: bool) -> bool { + vectorized_gate || has_mask +} + +fn warn_if_chain_parity_forfeited_by_shape(g: &MlxArray, mask: Option<&MlxArray>) { + let vectorized_gate = mlxcel_core::array_ndim(g) == 4; + if !chain_parity_forfeited_by_shape(vectorized_gate, mask.is_some()) { + return; + } + static WARNED: std::sync::Once = std::sync::Once::new(); + WARNED.call_once(|| { eprintln!("mlxcel: chain-parity gated-delta kernel was requested but ..."); }); +} +``` + +`gated_delta_ops_with_parity`의 `chain_parity` 디스패치 분기 바로 앞에서 호출한다. C++의 `if (vectorized && has_mask) ... else if (vectorized) ... else if (has_mask) ... else if (chain_parity)` 분기 순서를 그대로 거울처럼 따라간다. + +### 4.5 F5: 명시적 미측정 셀 (`docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md`) + +0.591 단일 요청 acceptance 수치 바로 아래에 다음을 추가했다. + +> **Not measured: acceptance under concurrent slice-grant rotation.** ... `Qwen35MtpDraftModel::reset` ... DESTROYS the drafter's accumulated KV history ... and it fires on every park (tick-slice grant rotation) ... Acceptance under that load regime is therefore structurally worse than the 0.591 measured here, and by how much is uncharacterized. + +### 4.6 F6~F10: 문서와 정리 + +F6은 `docs/environment-variables.md`의 speculative-decoding 표에 `MLXCEL_GDN_CHAIN_PARITY` 행을 추가했다. 다른 MTP 관련 환경변수와 같은 깊이로 서술했다. F7/F8은 `speculative_slice.rs`, `speculative_burst.rs`, `scheduler.rs`, `drafter/mod.rs`, `speculative/mtp/{mod,generator}.rs`에서 Gemma-4-전용이라는 서술과 reset이 no-op라는 서술을 바로잡았는데, 각 지점마다 그 구분을 지우는 대신 Gemma 4와 Qwen 3.5 MTP 두 family가 실제로 무엇이 다른지 구체적으로 서술했다. F9는 `generator.rs`의 일회성 prompt-prefill hook에서 `state.verify_out.verify_hidden_full.as_ref()`를 `.take()`로 바꿔서, 약 335MB짜리 seed hidden을 라운드가 끝날 때가 아니라 유일한 소비처가 쓰고 난 직후 바로 해제하게 했다. F10은 새로 추가된 `GATED_DELTA_METAL_SOURCE_SEQPAR` Metal 커널 소스에서 죽은 `if (true) { ... }` 래퍼를 제거했다(순수하게 들여쓰기만 바꾼 변경이고, 전체 재빌드 후 `parity_kernel_block_is_bitwise_equal_to_single_token_chain`으로 출력이 비트 단위로 그대로임을 확인했다). + +--- + +## 5. 전제 교정: recurrent rollback은 이미 존재하고 있었다 + +이슈 본문의 구현 계획은 이렇게 서술했다. "rollback이 어려운 부분이다: partial acceptance는 full-attention 레이어 16개의 KV를 자르는 동시에 linear 레이어 48개의 gated-delta recurrent state를 승인된 접두사까지 복원해야 한다. recurrent state는 자를 수 없다. draft 블록 이전에 스냅샷을 찍고 partial accept 시점에 복원하거나 재계산해야 한다." 구현이 시작되기 전에 이슈에 달린 코멘트 하나가 이를 바로잡았다. `Qwen35Model::rollback_speculative_cache`가 DFlash 경로를 위해 이미 정확히 이 일을 하고 있었다(attention 레이어는 `block_size - (max(accepted) + 1)`만큼 자르고 `B > 1`일 때는 행 단위로 꼬리를 0으로 채운다. gated-delta 레이어는 승인된 접두사 위에서 `GdnRollbackSnapshot`으로부터 replay한다). verify pass도 MTP drafter가 필요로 하는 hidden state를 이미 돌려주고 있었다(`VerifyOutput { logits, hidden_states, gdn_states }`). + +실제로 출하된 것은 이 교정된 범위, 즉 캐시 수술이 아니라 배선이다. `5bb9bd1d`는 새로운 recurrent-state 장치를 추가하지 않고 기존 verify/rollback 표면을 감싸는 drafter와 `MtpTarget` 어댑터를 추가했다. `89bcbae6`은 Gemma-4-전용 게이트를 걷어냈다. 두 커밋 모두 원래 전제가 아니라 교정된 전제와 일치한다. 이걸 여기 남겨두는 것은 프로젝트 자신의 관례(`docs/code-guidelines.md`, CLAUDE.md의 성능 이슈 규칙)를 따르는 일이다. 구현 전에 잡아낸 전제 교정은 이슈 절반의 각주가 아니라 그 절반의 산출물 자체다. + +--- + +## 6. 변경 요약 + +### 통계 (이번 finalization 단계) + +| 항목 | 값 | +|-----|---| +| 변경된 파일 수 | 18 | +| 추가/삭제 라인 (근사치, 최종 커밋 전) | +642 / -126 | +| 새 테스트 | `resolve_draft_block_size` / `peek_qwen35_mtp_configured_block_size` 5건, `chain_parity_gate_tests` 3건, `chain_parity_forfeited_by_shape` 2건, arm-attribution 가드 1건, 총 11개 테스트 함수 | +| 새로 문서화한 환경변수 | 1개 (`MLXCEL_GDN_CHAIN_PARITY`) | +| 다시 측정한 벤치마크 수치 | 4개 (T=1/3/4/64에서의 isolated kernel-cost 비율) | + +### 카테고리별 변경 + +| 카테고리 | 변경 수 | 주요 내용 | +|---------|--------|----------| +| 테스트/측정 신뢰성 | 1 | F1: arm-attribution 가드, 벤치마크 재측정 | +| 정확성 강화 | 2 | F2: capability gate의 shape 검사, F4: C++의 조용한 계약 상실 경고 | +| 기본값의 정확성 | 1 | F3: drafter가 설정한 block size | +| 문서의 정직성 | 1 | F5: 명시적 미측정 셀 표기 | +| 문서의 정확성 | 3 | F6: 환경변수 문서, F7/F8: Gemma-4-전용/reset-no-op 서술 약 10곳 | +| 메모리/정리 | 2 | F9: seed hidden 조기 해제, F10: 죽은 커널 스캐폴딩 제거 | + +### 관련 커밋 + +| Hash | Type | Message | +|------|------|---------| +| `5bb9bd1d` | feat | add the qwen3_5_mtp drafter and Qwen 3.5 MTP target adapter (상위 PR) | +| `89bcbae6` | feat | wire Qwen 3.5 targets into the MTP burst, slice, and offline paths (상위 PR) | +| `9575bad4` | test | add Qwen 3.5 MTP chain-parity diagnostic against the single-token chain (상위 PR) | +| `6fce98ba` | fix | make the Qwen 3.5 verify block bit-identical to the decode chain (상위 PR) | +| `f6d075be` | perf | gate Qwen MTP on Metal, measure the pairing, record the hardware verdict (상위 PR) | +| `60ef14b2` | test | pin MTP dispatch-site coverage and fix the chain-parity test lint (리뷰 단계, HIGH 2건) | +| `ef3aa95e` | fix | bound the qwen3_5_mtp drafter config scalars at load (보안 단계, HIGH 1건) | +| `a88af1ff` | fix | harden Qwen 3.5 MTP exactness gate and defaults (이번 단계, F1~F4) | +| `775afe95` | fix | correct stale MTP docs and trim a held-too-long buffer (이번 단계, F6~F10) | + +--- + +## 7. 후속 조치 + +### 완료 필요 + +- [ ] 없음. 이번 단계에서 리뷰가 지적한 열 건 모두 해결했다. + +### 향후 개선 사항 (알려진 제약으로 기록, 이번 PR에서는 수정하지 않음) + +- **이 페어링에는 M5급 하드웨어 측정이 없다.** Gemma 4 선례(verify amortization이 성립하고 M5 Max에서 1.2~1.4배)는 긍정적인 방향을 시사하지만, 이번 작업에는 M5 머신을 쓸 수 없었다. 이슈 스레드에 이 셀을 채우기 위한 재현 안내를 실어뒀다. +- **동시 slice-grant rotation 아래 acceptance는 특성화되지 않았다.** F5는 이를 고치는 대신 명시적인 미측정 셀로 문서화했다. 이를 채우려면 tick-slice grant-rotation 경로를 구체적으로 구동하는 다중 클라이언트 부하 하네스가 필요한데, 이번 단계의 범위 밖이다. +- **B > 1 batched MTP는 여전히 Gemma-4-전용이다.** Qwen 3.5 MTP는 설계상 B = 1만 지원한다(이슈 #1165 자신의 범위다). batched 어댑터 메서드는 이 family에 대해 트레이트 기본값대로 `DrafterError::DraftFailed`를 반환하며, 이는 문서화된 제약과 일치한다. +- raw HF 체크포인트에서 MTP head를 분리해내는 기능(사용자가 `mlx-community/Qwen3.8-27B-MTP-bf16`이 계속 존재하는 데만 의존하지 않도록)은 이슈 스레드에서 계속 추적되는, 의도적으로 미룬 후속 작업이다. + +--- + +## 부록 + +### A. 테스트 결과 + +- `cargo test --release --features metal,accelerate --test qwen38_mtp_chain_parity` (non-ignored): 1건 통과(`parity_kernel_block_is_bitwise_equal_to_single_token_chain`), 5건 무시(실제 체크포인트 필요), 0건 실패. F10의 커널 소스 정리 이후 전체 LTO 재빌드를 거쳐 다시 확인했다. +- `cargo test --release --features metal,accelerate --test qwen38_mtp_chain_parity -- --ignored --nocapture parity_kernel_cost`: 1건 통과, 다섯 번 반복 실행한 결과를 갱신된 벤치마크 기록에 남겼다. +- `cargo test --release --features metal,accelerate -p mlxcel-core drafter::qwen3_5_mtp`: 20건 통과, 0건 실패. +- `cargo test --release --features metal,accelerate -p mlxcel-core drafter::` (`mlxcel` 크레이트 자체의 `models::gated_delta` 모듈을 통한 shape-gate·chain-parity-forfeiture 테스트 포함): 6건 통과, 11건 무시(serial-MLX 게이트), 0건 실패. +- `cargo fmt --check`: 클린. +- arm-attribution 가드 수동 확인: `MLXCEL_GDN_CHAIN_PARITY=0 cargo test ... parity_kernel_cost -- --ignored --nocapture`는 예상한 메시지와 함께 패닉한다. 설정하지 않으면(기본값) 통과한다. + +### B. 참고 자료 + +- 이슈 #1165 (사양, 전제 교정 코멘트, 하드웨어 재현 안내) +- `docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md` (측정 기록, 이번 단계에서 갱신) +- `docs/environment-variables.md` (`MLXCEL_GDN_CHAIN_PARITY`, 신규) +- `src/models/gated_delta.rs` (chain-parity 커널 디스패치, arm-attribution과 맞닿은 경고) +- `src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp` (`GATED_DELTA_METAL_SOURCE_SEQPAR`) +- `src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/` (drafter 모델, config, 테스트) +- `src/models/qwen3_5_mtp_target.rs` (`MtpTarget` 어댑터) +- CLAUDE.md, "Performance issue completion criteria" (arm-attribution 규칙, 에픽 #909 선례) diff --git a/docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md b/docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md new file mode 100644 index 000000000..5fc930910 --- /dev/null +++ b/docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md @@ -0,0 +1,171 @@ +# Qwen3.8-27B MTP speculative decoding — M1 Ultra, 2026-08-16 + +Measurement record for issue #1165 (qwen3_5_mtp drafter support). All numbers +measured through the production call paths on the branch that adds the +feature, at commit range `5bb9bd1d..` (drafter + adapter + chain-parity +kernel). + +## Environment + +- Mac Studio, Apple M1 Ultra, 128 GB unified memory, macOS 26.6.1. +- Target: `models/qwen3.8-27b-4bit` (mlx-community conversion, 4-bit affine, + untied lm_head, vocab 248320). Drafter: `models/qwen3.8-27b-mtp-bf16` + (`mlx-community/Qwen3.8-27B-MTP-bf16`, 15 tensors, bf16, loaded as f16). +- Machine load (1-min avg) recorded per section below. The box carries + background load; every comparison below is a within-run ratio between arms + measured back to back, which survives it. Within-arm dispersion is shown by + min-max over the repetitions. + +## End-to-end: MTP vs classic decode (offline CLI, production path) + +`mlxcel generate -m -p -n 300 -t 0.0 [--draft-model ]`, +long-form prompt ("Explain in detail how a transformer neural network +works..."), 300 generated tokens, 5 repetitions per arm, decode tok/s as +printed by the CLI. Load 1.9-5.0 across the section. + +| arm | decode tok/s (median) | min-max | acceptance rate | emitted/verify | vs classic | +|---|---|---|---|---|---| +| classic | 23.45 | 23.44-23.46 | n/a | n/a | 1.00x | +| MTP, block 4 (CLI default) | 13.81 | 13.79-13.85 | 0.465 | 2.39 | 0.59x | +| MTP, block 3 (drafter-configured) | 16.48 | 16.39-16.56 | 0.591 | 2.18 | 0.70x | + +Acceptance is deterministic across repetitions (temperature 0). Round mix at +block 3: 55 full / 52 partial / 30 zero accepts over 137 rounds, i.e. +**first-draft acceptance 78%** — the jointly-trained head drafts well. +Exactness: temperature-0 output is byte-identical to classic decode on a +5-prompt gate (offline CLI) and token-identical on the server path (see the +detokenizer note below). + +**Not measured: acceptance under concurrent slice-grant rotation.** The +0.591 acceptance rate above is single-request only (offline CLI, one +sequence at a time, no grant rotation). `Qwen35MtpDraftModel::reset` +(`src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/model.rs`) DESTROYS the +drafter's accumulated KV history, unlike the Gemma 4 assistant drafter's +`reset` (the trait-default no-op), and it fires on every park (tick-slice +grant rotation, `super::speculative_slice`) and every `return_drafter`. Under +the server's continuous-batching scheduler with more than one concurrent MTP +request, a parked drafter resumes via `set_shared_kv` re-anchoring into the +empty-cache mode rather than resuming its history, so the drafter repeatedly +restarts from reduced context between rounds. Acceptance under that load +regime is therefore structurally worse than the 0.591 measured here, and by +how much is uncharacterized: this run did not exercise concurrent grant +rotation, and no number for it is reported anywhere in this document. + +## Where the time goes (per-round split, from the round-loop diagnostics) + +Block 4 run: 125 rounds, decode 21.67 s -> 173 ms/round emitting 2.39 tokens. + +| phase | ms/round | share | +|---|---|---| +| verify forward (K=4) | 137.6 | 68% | +| draft (3 drafter steps) | 19.9 | 10% | +| accept-hook extension (est.) | ~13 | 6% | +| verify finalize (rollback) | 2.7 | 1% | +| walk + shared-kv re-arm | <0.3 | ~0% | + +A classic decode step is 42.6 ms. The verify forward at K=4 costs 3.2 classic +steps, which is the entire story: at 137 ms the round emits 2.39 tokens +(72 ms/token) against classic's 42.6 ms/token. + +## Root cause of the verify cost: the target's own multi-token forward + +Isolation (`tests/qwen38_mtp_chain_parity.rs::verify_forward_cost_scaling`, +medians over 20 timed calls per shape): + +| T | `forward_speculative` (per-position verify attention + GDN snapshot capture) | plain batched prefill path (same block, no verify machinery) | +|---|---|---| +| 1 | 43.6 ms | 43.6 ms | +| 2 | 75.7 ms | 75.0 ms | +| 3 | 106.2 ms | 104.4 ms | +| 4 | 135.5 ms | 133.2 ms | +| 8 | 244.0 ms | 238.7 ms | + +The verify machinery itself (per-position attention dispatch loop, GDN +rollback snapshots) costs under 2% at every T. The cost lives in the target's +**plain multi-token forward**: every extra token in a block adds ~29 ms, about +two thirds of a full single-token step, on both paths equally. The +K-token-block-for-the-price-of-one premise of speculative decoding does not +hold on this GPU generation for this model. This is precisely the pre-M5 +regression already encoded for Gemma 4 in `mtp_b1_default` +(`src/server/batch/speculative_burst.rs`): M1 Ultra measured 0.75-0.96x for +Gemma 4 MTP while M5 Max measured 1.2-1.4x, and the static per-hardware +default plus the adaptive MTP policy (issue #333) exist exactly to gate this. +Qwen fares somewhat worse than Gemma here because each drafter step and +accept-hook extension pays a full-vocab (248320) LM-head projection. + +**Not measured: M5-class hardware.** The Gemma precedent (verify amortization +holds, 1.2-1.4x) projects a positive sign for this pairing there, but no M5 +machine was available to this run; the cell is open. + +## Chain-parity kernel cost (correctness requirement, not an optimization) + +The temperature-0 exactness gate initially failed: one near-tie argmax flip +per ~100-250 generated tokens. Cause: the gated-delta Metal kernel emits its +recurrent state in the storage dtype (f16), so a T=1 decode chain rounds the +state after every token while a T=K verify block carried float32 state across +the block and rounded once. The chain-parity kernel variant +(`gated_delta_step_seqpar`) rounds at every in-block step, making T=K +bit-identical to K single-token steps (pinned bitwise in +`parity_kernel_block_is_bitwise_equal_to_single_token_chain`). It is used +ONLY by the speculative verify capture path and the rollback replay; classic +prefill and decode are untouched. + +- Isolated kernel cost (27B GDN geometry, medians over 100 dispatches per + run). **Corrected numbers**: the test (`parity_kernel_cost_vs_standard`, + `tests/qwen38_mtp_chain_parity.rs`) originally had no arm-attribution + guard. With `MLXCEL_GDN_CHAIN_PARITY=0` in the environment both the + "standard" and "parity" arms silently resolve to the same standard kernel, + and the test would compare that path against itself without noticing. It + now runs `assert_parity_and_standard_arms_are_distinguishable` first, + which dispatches both kernels once at `T=3` and panics if their final GDN + state is bitwise equal (the two kernels are defined to diverge whenever + `T > 1`); the run this table is drawn from confirmed the guard passes + (arms are distinguishable) before printing any ratio. Five repeated runs, + medians of each run's 100-dispatch-median shown, min-max across the five + runs in parens: ratios parity/standard 0.979 (0.877-1.013, T=1), 1.007 + (0.987-1.024, T=3), 0.994 (0.922-1.003, T=4), 1.022 (0.995-1.043, T=64), + still inside the dispatch-noise band. The qualitative conclusion is + unchanged from the originally recorded (unattributed) 1.089 / 1.117 / + 0.834 / 1.008: the kernel's time loop was already sequential in T, the + variant only adds a per-step register round. The absolute ratios moved + run-to-run on this box (load 4.8-5.2 during the corrected run) by more + than the originally reported figures suggested, which is itself evidence + the earlier single-sample numbers were not attributable to anything more + precise than noise. +- End-to-end attribution through the production MTP path + (`MLXCEL_GDN_CHAIN_PARITY=0` diagnostic toggle, 3 reps): 14.29 tok/s + without parity vs 13.81 with, ~3.4% — the price of byte-exactness, and not + the reason MTP loses on this box (see above). + +## DFlash blast radius (shared verify machinery) + +`forward_speculative` is also DFlash's verify path, so the parity kernel +changes a shipped feature. Same-binary A/B on identical prompts +(`qwen3.5-27b-4bit` + `qwen3.5-27b-dflash`, server path, temp 0): + +| request | arm | wall | rounds | acceptance | emitted/verify | +|---|---|---|---|---|---| +| 120 tok | parity OFF (pre-change numerics) | 17.09 s | 37 | 0.152 | 3.22 | +| 120 tok | parity ON | 16.79 s | 38 | 0.146 | 3.13 | +| 80 tok | parity OFF | 16.69 s | 41 | 0.0705 | 1.93 | +| 80 tok | parity ON | 16.90 s | 41 | 0.068 | 1.93 | + +No material change: walls within run noise, acceptance within near-tie +reshuffling. (A separately quoted main-branch baseline with acceptance 0.274 +at 120 tokens used different prompts and is not attributable; the controlled +same-prompt A/B above is the before/after.) + +## Server E2E and the boundary detokenizer note + +`mlxcel-server` with the MTP pair engages the tick-slice path (55-57 slices +for a 150-token completion) and produces **token-identical** output to the +classic server (150/150 token ids equal, `logprobs` comparison); a 170-token +run is byte-identical in text as well. A 150-token run's *text* differed in +the final characters only because the token budget cut mid-BPE-merge and the +classic streaming detokenizer flushes a trailing partial merge differently +from the burst finalizer — a pre-existing boundary artifact independent of +MTP (token ids agree). + +The adaptive MTP policy (issue #333) governs engagement: on this box its +profiling phase attempts MTP and its estimator settles per pairing; the +static pre-M5 default declines. `MLXCEL_ENABLE_MTP_B1` pins either way. diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 72f5019ff..abf9e6d00 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -184,6 +184,7 @@ The OpenAI audio endpoints (`/v1/audio/speech`, `/v1/audio/transcriptions`, `/v1 | `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 ~1.87× on M5 Max and ~1.1 to 1.4× on M1 Ultra); batch-capable targets (the 31B + bf16 assistant) default **on only on M5+** (Neural Accelerator generation): M5 Max measured ~1.2 to 1.4×, while M1 Ultra measured a consistent ~0.75 to 0.96× regression, so pre-M5 chips fall back to classic decode. | | `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_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). See `docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md` for the measured kernel cost (inside the dispatch-noise band). | | `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/docs/supported-models.md b/docs/supported-models.md index d935ab328..5f07e4376 100644 --- a/docs/supported-models.md +++ b/docs/supported-models.md @@ -533,7 +533,8 @@ until each path has explicit mixed-cache and multimodal validation. | Drafter | Target families | Notes | |---------|-----------------|-------| | MTP | Gemma 4 target paths | Available through shared speculative decoding flags, on `mlxcel-server` and on offline `mlxcel generate --draft-kind mtp`. | -| DFlash | Qwen 3.5 text/VLM paths | `mlxcel-server` only. Available through shared speculative decoding flags there; offline `mlxcel generate` rejects a DFlash drafter with a named error because it does not construct the DFlash round loop. Qwen3.8 rides the same `qwen3_5` path but has no published DFlash drafter, and its MTP tensors are dropped by the mlx-community conversion, so it runs unspeculated (#1165). | +| MTP | Qwen 3.5 / 3.6 / 3.8 text and VLM paths (B = 1) | Requires the split-out `qwen3_5_mtp` drafter checkpoint (for Qwen3.8-27B: [`mlx-community/Qwen3.8-27B-MTP-bf16`](https://huggingface.co/mlx-community/Qwen3.8-27B-MTP-bf16); the mlx-community target conversions drop the in-model `mtp.*` tensors, so the drafter is always a separate directory). Auto-detected from the drafter's `model_type` with no `--draft-kind`, on `mlxcel-server` and offline `mlxcel generate`. Single-request only: batched (B > 1) MTP windows decline to classic decode. Image/video prefill runs on the target; MTP accelerates only the text decode tail. Metal-only (the temperature-0 exactness contract rests on the Metal chain-parity gated-delta kernel). Whether MTP pays off is hardware-generation-dependent, exactly as for Gemma 4: on M1 Ultra the target's own multi-token verify forward barely amortizes and MTP measures 0.59-0.70x classic; the server's per-hardware default plus the adaptive MTP policy gate this automatically. See `benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md`. | +| DFlash | Qwen 3.5 text/VLM paths | `mlxcel-server` only. Available through shared speculative decoding flags there; offline `mlxcel generate` rejects a DFlash drafter with a named error because it does not construct the DFlash round loop. Qwen3.8 rides the same `qwen3_5` path but has no published DFlash drafter; pair it with the MTP drafter row above instead (#1165). | Use auto-detection by default. Override only when you know the target and drafter checkpoint pair are compatible. diff --git a/src/cli/speculative_args.rs b/src/cli/speculative_args.rs index 38e280818..b97afad96 100644 --- a/src/cli/speculative_args.rs +++ b/src/cli/speculative_args.rs @@ -54,6 +54,12 @@ use mlxcel_core::drafter::{DrafterKind, KNOWN_DRAFTER_KINDS}; /// /// - **MTP** → `4`: the Gemma 4 MTP "assistant" draft block length used /// by https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/speculative/drafters/gemma4_assistant/config.py. +/// `DrafterKind::Mtp` also covers the unrelated Qwen 3.5 MTP family; for +/// that family [`resolve_draft_block_size`] prefers the drafter +/// checkpoint's own configured `block_size` +/// (`mlxcel_core::drafter::peek_qwen35_mtp_configured_block_size`) over +/// this constant, so this Gemma-4-derived value is only the fallback +/// when no Qwen 3.5 MTP hint is found. /// - **DFlash** → `16`: the Qwen 3.5 DFlash drafter's `block_size` /// declared in https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/speculative/drafters/qwen3_dflash/config.py#L31. /// Mirrors [`mlxcel_core::drafter::dflash::DEFAULT_BLOCK_SIZE`]. @@ -83,14 +89,16 @@ pub struct SpeculativeArgs { /// `mlxcel-server`), the kind is auto-detected from the drafter's /// `config.json::model_type` via /// `mlxcel_core::drafter::resolve_drafter_kind`. Auto-detect maps - /// `gemma4_assistant -> mtp`; everything else falls back to - /// `dflash` (matching the upstream - /// `DEFAULT_DRAFTER_KIND = "dflash"` convention). + /// `gemma4_assistant`, `gemma4_unified_assistant`, and `qwen3_5_mtp` + /// to `mtp`; everything else falls back to `dflash` (matching the + /// upstream `DEFAULT_DRAFTER_KIND = "dflash"` convention). /// /// When unset AND no drafter path is supplied, mlxcel runs without a /// speculative drafter. The offline `mlxcel generate` command keeps the - /// classic non-MTP / non-DFlash `SpeculativeGenerator` path when a - /// drafter is supplied without an explicit kind. + /// classic `SpeculativeGenerator` path when a drafter without an MTP + /// `model_type` is supplied without an explicit kind; an auto-detected + /// MTP drafter routes to the MTP round loop (none of the MTP drafter + /// checkpoints can load as a standalone classic draft model). /// /// Parity note (offline `mlxcel generate` with `--draft-kind mtp`): /// at temperature 0 with no sampling penalties the output matches the @@ -187,15 +195,41 @@ pub fn default_block_size_for_kind(kind: DrafterKind) -> u32 { } } -/// Resolve the effective draft block size given an explicit CLI override -/// and the resolved drafter kind. +/// Resolve the effective draft block size given an explicit CLI override, +/// the resolved drafter kind, and the drafter checkpoint path. /// /// When `override_value` is `Some(n)`, returns `n` (with no further /// validation here; concrete generators enforce their own minimums). -/// When `override_value` is `None`, returns -/// [`default_block_size_for_kind`] for `kind`. -pub fn resolve_draft_block_size(override_value: Option, kind: DrafterKind) -> u32 { - override_value.unwrap_or_else(|| default_block_size_for_kind(kind)) +/// +/// When `override_value` is `None` AND `kind == DrafterKind::Mtp`, this +/// first tries [`mlxcel_core::drafter::peek_qwen35_mtp_configured_block_size`]: +/// `DrafterKind::Mtp` covers two unrelated drafter families (Gemma 4 +/// assistant and Qwen 3.5 MTP), and [`DEFAULT_MTP_BLOCK_SIZE`] is the +/// Gemma-4-derived constant. Applying it unconditionally to a Qwen 3.5 MTP +/// drafter ships whatever that constant happens to be rather than the +/// checkpoint's own declared block size. Issue #1165's own measurement +/// found the published checkpoint's configured block size (3) faster than +/// the flat constant (4): 16.48 vs 13.81 tok/s, 0.591 vs 0.465 acceptance +/// (`docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md`). When the +/// peek finds no Qwen 3.5 MTP `block_size` (wrong family, missing config, +/// or the checkpoint omits an explicit value), falls back to +/// [`default_block_size_for_kind`] as before. +pub fn resolve_draft_block_size( + override_value: Option, + kind: DrafterKind, + model_path: &std::path::Path, +) -> u32 { + if let Some(n) = override_value { + return n; + } + if kind == DrafterKind::Mtp + && let Some(configured) = + mlxcel_core::drafter::peek_qwen35_mtp_configured_block_size(model_path) + && let Ok(n) = u32::try_from(configured) + { + return n; + } + default_block_size_for_kind(kind) } /// Apply the `MLXCEL_DRAFT_KIND` env-var fallback to the raw diff --git a/src/cli/speculative_args_tests.rs b/src/cli/speculative_args_tests.rs index 0d31d8701..23f1a7e3b 100644 --- a/src/cli/speculative_args_tests.rs +++ b/src/cli/speculative_args_tests.rs @@ -117,14 +117,70 @@ fn default_block_size_for_internal_mtp_shares_dflash_default() { #[test] fn resolve_draft_block_size_uses_override_when_provided() { - assert_eq!(resolve_draft_block_size(Some(8), DrafterKind::Mtp), 8); - assert_eq!(resolve_draft_block_size(Some(32), DrafterKind::Dflash), 32); + // An explicit override always wins, even against a drafter path whose + // own config declares a different value. + let dir = tempfile::tempdir().expect("tempdir"); + std::fs::write( + dir.path().join("config.json"), + r#"{"model_type": "qwen3_5_mtp", "block_size": 3}"#, + ) + .expect("write config.json"); + assert_eq!( + resolve_draft_block_size(Some(8), DrafterKind::Mtp, dir.path()), + 8 + ); + assert_eq!( + resolve_draft_block_size(Some(32), DrafterKind::Dflash, dir.path()), + 32 + ); +} + +#[test] +fn resolve_draft_block_size_falls_back_to_per_kind_default_without_a_drafter_hint() { + // No config.json at all (nonexistent path). DFlash never peeks, and + // an MTP drafter with no discoverable config falls back to the flat + // Gemma-4-derived constant. + let missing = std::path::Path::new("/nonexistent/mlxcel-drafter-test-path"); + assert_eq!(resolve_draft_block_size(None, DrafterKind::Mtp, missing), 4); + assert_eq!( + resolve_draft_block_size(None, DrafterKind::Dflash, missing), + 16 + ); +} + +#[test] +fn resolve_draft_block_size_honors_the_qwen35_mtp_drafters_own_configured_value() { + // Issue #1165 hardening: a Qwen 3.5 MTP drafter's own `block_size` + // (the published checkpoint declares 3) must win over the flat + // Gemma-4-derived MTP constant (4) when the operator passes no + // explicit `--draft-block-size` override. + let dir = tempfile::tempdir().expect("tempdir"); + std::fs::write( + dir.path().join("config.json"), + r#"{"model_type": "qwen3_5_mtp", "block_size": 3}"#, + ) + .expect("write config.json"); + assert_eq!( + resolve_draft_block_size(None, DrafterKind::Mtp, dir.path()), + 3 + ); } #[test] -fn resolve_draft_block_size_falls_back_to_per_kind_default() { - assert_eq!(resolve_draft_block_size(None, DrafterKind::Mtp), 4); - assert_eq!(resolve_draft_block_size(None, DrafterKind::Dflash), 16); +fn resolve_draft_block_size_ignores_block_size_from_a_non_qwen35_mtp_drafter() { + // A Gemma 4 assistant drafter (or any non-`qwen3_5_mtp` model_type) + // must not have a same-named `block_size` field misread as the hint; + // the flat per-kind default still applies. + let dir = tempfile::tempdir().expect("tempdir"); + std::fs::write( + dir.path().join("config.json"), + r#"{"model_type": "gemma4_assistant", "block_size": 7}"#, + ) + .expect("write config.json"); + assert_eq!( + resolve_draft_block_size(None, DrafterKind::Mtp, dir.path()), + 4 + ); } #[test] diff --git a/src/commands/generate.rs b/src/commands/generate.rs index 8ae7eb284..6c55444d4 100644 --- a/src/commands/generate.rs +++ b/src/commands/generate.rs @@ -1406,17 +1406,22 @@ pub(super) fn run_generation_mode( .map_err(|e| anyhow::anyhow!("{e}"))?; let resolved_kind = resolve_drafter_kind(draft_model_path, explicit_kind) .map_err(|e| anyhow::anyhow!("--draft-kind / drafter config: {e}"))?; - let block_size = resolve_draft_block_size(args.speculative.draft_block_size, resolved_kind); + let block_size = resolve_draft_block_size( + args.speculative.draft_block_size, + resolved_kind, + draft_model_path, + ); let user_requested_explicit_kind = explicit_kind.is_some(); - // issue #166: when the operator explicitly passes `--draft-kind mtp`, + // issue #166 / #1165: when the resolved kind is MTP (explicit + // `--draft-kind mtp`, or auto-detected from an MTP `model_type`), // drive the kind-specific `MtpGenerator` round loop here, reusing the // SAME per-target `MtpTarget` adapters the server burst path uses - // (`src/models/gemma4_mtp_target.rs`). This runs BEFORE - // `load_model(draft_model_path)` because an MTP assistant is loaded as a - // `Drafter` (via `load_drafter`), not as a full `LoadedModel`. DFlash / - // InternalMtp explicit kinds and the auto-detect classic path are - // untouched and fall through below. + // (`src/models/gemma4_mtp_target.rs`, `src/models/qwen3_5_mtp_target.rs`). + // This runs BEFORE `load_model(draft_model_path)` because an MTP + // drafter is loaded as a `Drafter` (via `load_drafter`), not as a full + // `LoadedModel`. DFlash / InternalMtp explicit kinds and the + // auto-detect classic path are untouched and fall through below. if should_route_offline_mtp(user_requested_explicit_kind, resolved_kind) { if vlm_embeddings.is_some() { return Err(anyhow!( @@ -1583,10 +1588,19 @@ pub(super) fn run_generation_mode( /// through to the deferred-error branch. Extracted as a pure function so the /// loop-construction decision is unit-testable without loading a model. fn should_route_offline_mtp( - user_requested_explicit_kind: bool, + _user_requested_explicit_kind: bool, resolved_kind: DrafterKind, ) -> bool { - user_requested_explicit_kind && resolved_kind == DrafterKind::Mtp + // Route on the RESOLVED kind, explicit or auto-detected (#1165 + // acceptance: a `qwen3_5_mtp` drafter auto-resolves to MTP with no + // `--draft-kind`). Auto-detect only resolves to `Mtp` for a genuinely + // MTP `model_type` (`gemma4_assistant`, `gemma4_unified_assistant`, + // `qwen3_5_mtp`), none of which can load as a standalone model, so the + // classic fall-through those drafters used to hit was always a + // downstream load error, never a working path. Unknown model_types + // still auto-resolve to the DFlash default and keep the classic + // SpeculativeGenerator (the small-full-model-as-drafter workflow). + resolved_kind == DrafterKind::Mtp } /// Reject a DFlash speculative drafter before the offline path tries to load @@ -1726,12 +1740,31 @@ fn run_offline_mtp( LoadedModel::Gemma4(wrapper) => wrapper as &dyn LanguageModel, LoadedModel::Gemma4VLM(vlm) => vlm as &dyn LanguageModel, LoadedModel::Gemma4Unified(unified) => unified as &dyn LanguageModel, + // Metal-only for the Qwen 3.5 family: the temperature-0 exactness + // contract rests on the Metal chain-parity gated-delta kernel + // (issue #1165); on other backends the verify block is not + // bit-identical to classic decode, so fail closed with the reason. + LoadedModel::Qwen35(_) + | LoadedModel::Qwen35Moe(_) + | LoadedModel::Qwen35VLM(_) + | LoadedModel::Qwen35MoeVLM(_) + if !mlxcel_core::metal_is_available() => + { + return Err(anyhow!( + "Qwen 3.5 MTP speculative decoding requires the Metal backend: its \ + temperature-0 exactness guarantee rests on the Metal chain-parity \ + gated-delta kernel, which has no CUDA/CPU port yet. Omit \ + --draft-model to run classic decode." + )); + } + LoadedModel::Qwen35(qwen) | LoadedModel::Qwen35Moe(qwen) => qwen as &dyn LanguageModel, + LoadedModel::Qwen35VLM(vlm) | LoadedModel::Qwen35MoeVLM(vlm) => vlm as &dyn LanguageModel, _ => { return Err(anyhow!( "--draft-kind mtp is only supported for Gemma 4 (text, VLM, or \ - Unified) targets; the loaded target is not MTP-capable. Omit \ - --draft-kind to use the classic SpeculativeGenerator with your \ - --draft-model drafter." + Unified) and Qwen 3.5 (text or VLM) targets; the loaded target \ + is not MTP-capable. Omit --draft-kind to use the classic \ + SpeculativeGenerator with your --draft-model drafter." )); } }; @@ -1801,6 +1834,27 @@ fn run_offline_mtp( &token_history, block_size, ), + // Qwen 3.5 family (#1165): `seq_id = None` selects the model's + // internal fallback cache slot, the documented offline / single-row + // CLI shape. + LoadedModel::Qwen35(qwen) | LoadedModel::Qwen35Moe(qwen) => drive_offline_mtp( + mlxcel::models::qwen3_5_mtp_target::Qwen35MtpTargetAdapter::new(qwen, None), + drafter, + prompt_tokens, + max_tokens, + &sampling, + &token_history, + block_size, + ), + LoadedModel::Qwen35VLM(vlm) | LoadedModel::Qwen35MoeVLM(vlm) => drive_offline_mtp( + mlxcel::models::qwen3_5_mtp_target::Qwen35VLMtpTargetAdapter::new(vlm, None), + drafter, + prompt_tokens, + max_tokens, + &sampling, + &token_history, + block_size, + ), // Unreachable: the variant gate above already returned an error for any // non-MTP-capable target. Kept for match exhaustiveness. _ => unreachable!("non-MTP-capable target rejected by the variant gate above"), diff --git a/src/commands/generate_tests.rs b/src/commands/generate_tests.rs index 05d25d7a6..331f8bcc3 100644 --- a/src/commands/generate_tests.rs +++ b/src/commands/generate_tests.rs @@ -67,21 +67,23 @@ fn xla_cli_accepts_empty_image_list_for_audio_only_requests() { assert!(decode_xla_cli_images(&[]).unwrap().is_empty()); } -// issue #166: the offline MTP loop is constructed only when the operator -// explicitly passes `--draft-kind mtp`. An auto-detected MTP shape (no explicit -// flag) keeps the classic SpeculativeGenerator path, and DFlash / InternalMtp -// explicit kinds fall through to the deferred-error branch. These pins guard the -// routing decision without loading a model. -#[test] -fn should_route_offline_mtp_only_for_explicit_mtp() { +// issue #166 / #1165: the offline MTP loop is constructed whenever the +// RESOLVED kind is MTP — an explicit `--draft-kind mtp`, or an auto-detected +// MTP `model_type` (`gemma4_assistant`, `gemma4_unified_assistant`, +// `qwen3_5_mtp`), none of which can load as a standalone model, so the +// classic fall-through was always a downstream load error. DFlash / +// InternalMtp explicit kinds fall through to the deferred-error branch. These +// pins guard the routing decision without loading a model. +#[test] +fn should_route_offline_mtp_for_explicit_mtp() { assert!(should_route_offline_mtp(true, DrafterKind::Mtp)); } #[test] -fn should_route_offline_mtp_rejects_auto_detected_mtp() { - // Auto-detect resolved to MTP but no explicit `--draft-kind`: stay on the - // classic path for backward compatibility. - assert!(!should_route_offline_mtp(false, DrafterKind::Mtp)); +fn should_route_offline_mtp_for_auto_detected_mtp() { + // #1165 acceptance: a qwen3_5_mtp drafter with no --draft-kind + // auto-resolves to MTP and takes the MTP round loop. + assert!(should_route_offline_mtp(false, DrafterKind::Mtp)); } #[test] diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index 22f2a25dd..a092b4924 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -1798,6 +1798,25 @@ void metal_gated_delta_forward( std::unique_ptr& new_state // [B, Hv, Dv, Dk] ); +// Chain-parity variant for the speculative verify / rollback replay paths +// (issue #1165): rounds the recurrent state through the storage dtype after +// every in-block step so a T=K block is bit-identical to K consecutive T=1 +// decode calls. Identical to metal_gated_delta_forward at T=1. Only the +// scalar-gate / no-mask shape gets the parity kernel; other shapes fall back +// to the standard kernels. +// Used by: Qwen3.5 forward_speculative + rollback replay. +void metal_gated_delta_forward_chain_parity( + const MlxArray& q, + const MlxArray& k, + const MlxArray& v, + const MlxArray& g, + const MlxArray& beta, + const MlxArray& state, + const MlxArray* mask, // nullable: [B, T] + std::unique_ptr& output, + std::unique_ptr& new_state +); + // Quantization mode is "affine" (standard mlx-community models). void fused_mamba2_forward( // Input diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp index 6b1c6b19b..e77d2bd23 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp @@ -605,6 +605,93 @@ namespace { } )"; + // Variant 1b: scalar gate, no mask, CHAIN-PARITY state rounding. + // + // Identical arithmetic to variant 1 except the recurrent state is + // rounded through the storage dtype (InT, f16/bf16 in practice) at the + // END of every in-block step. The standard kernel carries float32 state + // across the whole T-block and quantizes only at the final store, so a + // T=K block is NOT bit-identical to K consecutive T=1 calls (each of + // which round-trips the state through InT between steps). The MTP / + // DFlash speculative verify pass and the rollback replay need exactly + // that bit-identity: at temperature 0 the verify block's argmax must + // match the classic single-token decode chain, and a last-ulp state + // difference flips near-tie logits (issue #1165's exactness gate). + // Rounding at each step end makes T=K bit-equal to the T=1 chain; for + // T=1 the round is idempotent with the store cast, so the kernel is + // byte-identical to variant 1 in the decode shape. + static const char* GATED_DELTA_METAL_SOURCE_SEQPAR = R"( + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + constexpr int n_per_t = Dk / 32; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + y += b_idx * T * Hv * Dv + hv_idx * Dv; + + auto dk_idx = thread_position_in_threadgroup.x; + auto dv_idx = thread_position_in_grid.y; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + float state[n_per_t]; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = static_cast(i_state[s_idx]); + } + + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + auto beta_ = beta + b_idx * T * Hv; + + for (int t = 0; t < T; ++t) { + float kv_mem = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = state[i] * g_[hv_idx]; + kv_mem += state[i] * k_[s_idx]; + } + kv_mem = simd_sum(kv_mem); + + auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx]; + + float out = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = state[i] + k_[s_idx] * delta; + out += state[i] * q_[s_idx]; + } + out = simd_sum(out); + if (thread_index_in_simdgroup == 0) { + y[dv_idx] = static_cast(out); + } + // Chain parity: round the state through the storage dtype at + // the end of every step, exactly where the T=1 chain's store + // + reload would round it. + for (int i = 0; i < n_per_t; ++i) { + state[i] = static_cast(static_cast(state[i])); + } + q_ += Hk * Dk; + k_ += Hk * Dk; + v_ += Hv * Dv; + y += Hv * Dv; + g_ += Hv; + beta_ += Hv; + } + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + o_state[s_idx] = static_cast(state[i]); + } + )"; + // Variant 2: scalar gate, with mask static const char* GATED_DELTA_METAL_SOURCE_MASK = R"( auto n = thread_position_in_grid.z; @@ -831,6 +918,10 @@ namespace { static GatedDeltaKernelHolder holder; return holder; } + static GatedDeltaKernelHolder& get_gd_kernel_seqpar() { + static GatedDeltaKernelHolder holder; + return holder; + } static GatedDeltaKernelHolder& get_gd_kernel_mask() { static GatedDeltaKernelHolder holder; return holder; @@ -889,7 +980,7 @@ void metal_stop_capture() { mlx::core::metal::stop_capture(); } -void metal_gated_delta_forward( +static void metal_gated_delta_forward_impl( const MlxArray& q, const MlxArray& k, const MlxArray& v, @@ -897,6 +988,7 @@ void metal_gated_delta_forward( const MlxArray& beta, const MlxArray& state, const MlxArray* mask, + bool chain_parity, std::unique_ptr& output, std::unique_ptr& new_state ) { @@ -954,6 +1046,17 @@ void metal_gated_delta_forward( holder = &get_gd_kernel_mask(); kernel_name = "gated_delta_step_mask"; kernel_source = GATED_DELTA_METAL_SOURCE_MASK; + } else if (chain_parity) { + // Chain-parity variant: per-step state rounding so a T=K verify + // block is bit-identical to K consecutive T=1 decode calls. Only + // defined for the scalar-gate/no-mask shape the speculative verify + // and rollback-replay paths use; the vectorized/masked shapes above + // take the standard kernels (their callers never request parity). + input_names = {"q", "k", "v", "g", "beta", "state_in", "T"}; + inputs = {q.inner, k.inner, v.inner, g_cast, beta_cast, state_cast, T_arr}; + holder = &get_gd_kernel_seqpar(); + kernel_name = "gated_delta_step_seqpar"; + kernel_source = GATED_DELTA_METAL_SOURCE_SEQPAR; } else { input_names = {"q", "k", "v", "g", "beta", "state_in", "T"}; inputs = {q.inner, k.inner, v.inner, g_cast, beta_cast, state_cast, T_arr}; @@ -997,6 +1100,41 @@ void metal_gated_delta_forward( new_state = std::make_unique(std::move(results[1])); } +void metal_gated_delta_forward( + const MlxArray& q, + const MlxArray& k, + const MlxArray& v, + const MlxArray& g, + const MlxArray& beta, + const MlxArray& state, + const MlxArray* mask, + std::unique_ptr& output, + std::unique_ptr& new_state +) { + metal_gated_delta_forward_impl(q, k, v, g, beta, state, mask, false, output, new_state); +} + +// Chain-parity gated-delta forward for the speculative verify / rollback +// replay paths (issue #1165): a T=K block is bit-identical to K consecutive +// T=1 decode calls because the recurrent state is rounded through the +// storage dtype after every in-block step. Behaviour is identical to +// metal_gated_delta_forward for T=1. Parity is only defined for the +// scalar-gate / no-mask shape; vectorized or masked inputs fall back to the +// standard kernels. +void metal_gated_delta_forward_chain_parity( + const MlxArray& q, + const MlxArray& k, + const MlxArray& v, + const MlxArray& g, + const MlxArray& beta, + const MlxArray& state, + const MlxArray* mask, + std::unique_ptr& output, + std::unique_ptr& new_state +) { + metal_gated_delta_forward_impl(q, k, v, g, beta, state, mask, true, output, new_state); +} + // Compiled MoE gate: sigmoid + bias + topk + normalize + scale // Uses mx::core::compile for kernel fusion matching Python's @mx.compile group_expert_select namespace { diff --git a/src/lib/mlxcel-core/src/drafter/dflash/drafter.rs b/src/lib/mlxcel-core/src/drafter/dflash/drafter.rs index 9f5fd7c85..4461c53ef 100644 --- a/src/lib/mlxcel-core/src/drafter/dflash/drafter.rs +++ b/src/lib/mlxcel-core/src/drafter/dflash/drafter.rs @@ -165,7 +165,10 @@ impl DFlashDrafter { /// /// `weights` is mutated in place; non-bf16 tensors and quantization /// auxiliaries (scales, biases) are untouched. -fn convert_bf16_to_f16_non_quantized(weights: &mut WeightMap) { +/// +/// Used by: DFlash drafter load, Qwen 3.5 MTP drafter load +/// (`crate::drafter::qwen3_5_mtp::model::Qwen35MtpDraftModel::from_path`). +pub(crate) fn convert_bf16_to_f16_non_quantized(weights: &mut WeightMap) { let bf16_keys: Vec = weights .iter() .filter(|(k, v)| { diff --git a/src/lib/mlxcel-core/src/drafter/dflash/mlp.rs b/src/lib/mlxcel-core/src/drafter/dflash/mlp.rs index 05eec4778..b8d28d82a 100644 --- a/src/lib/mlxcel-core/src/drafter/dflash/mlp.rs +++ b/src/lib/mlxcel-core/src/drafter/dflash/mlp.rs @@ -21,6 +21,9 @@ use cxx::UniquePtr; /// SwiGLU MLP block reused by every drafter layer. /// +/// Used by: DFlash drafter layers, Qwen 3.5 MTP drafter layer +/// (`crate::drafter::qwen3_5_mtp::layer::Qwen35MtpDecoderLayer`). +/// /// Mirrors upstream `Qwen3MLP(config.hidden_size, config.intermediate_size)` /// (`gate_proj`, `up_proj`, `down_proj`, all bias-free). The non-quantized /// path uses the fused compiled FP MLP (`compiled_swiglu_mlp_fp16`), so diff --git a/src/lib/mlxcel-core/src/drafter/mod.rs b/src/lib/mlxcel-core/src/drafter/mod.rs index 22dc08222..36c4d9099 100644 --- a/src/lib/mlxcel-core/src/drafter/mod.rs +++ b/src/lib/mlxcel-core/src/drafter/mod.rs @@ -68,7 +68,7 @@ //! //! | Variant | Concrete impl | Wired by | //! |---------|---------------|----------| -//! | [`DrafterKind::Mtp`] | `Gemma4AssistantDraftModel` | | +//! | [`DrafterKind::Mtp`] | `Gemma4AssistantDraftModel`, and (since issue #1165) `Qwen35MtpDraftModel` (`qwen3_5_mtp` model_type) | | //! | [`DrafterKind::Dflash`] | `DFlashDraftModel` | | //! | [`DrafterKind::InternalMtp`] | `InternalMtpDrafter` | | //! @@ -99,6 +99,9 @@ pub mod gemma4_assistant; /// — landed here independently so the layer can /// be unit-tested in isolation before integration. pub mod masked_embedder; +/// Concrete Qwen 3.5 / 3.6 / 3.8 MTP drafter implementation (the split-out +/// `qwen3_5_mtp` head). Wired into [`load_drafter`]'s `Mtp` arm. +pub mod qwen3_5_mtp; /// Test-only capture of this module's `tracing` output, used by the /// assertions on [`resolve_drafter_kind`]'s operator-facing diagnostics. #[cfg(test)] @@ -215,6 +218,13 @@ pub fn drafter_kind_by_model_type() -> &'static HashMap<&'static str, DrafterKin // `Gemma4AssistantDraftModel` class (identical centroid head + 4-layer // shared-K/V stack) and therefore resolves to the same MTP round loop. m.insert("gemma4_unified_assistant", DrafterKind::Mtp); + // The Qwen 3.5 / 3.6 / 3.8 split-out MTP head (e.g. + // mlx-community/Qwen3.8-27B-MTP-bf16). Without this entry a + // `qwen3_5_mtp` drafter auto-resolves to the DFlash default and dies + // in the generic loader with "Unsupported model type". Loads with + // `Qwen35MtpDraftModel` via [`load_drafter`]'s per-model_type Mtp + // dispatch. + m.insert(QWEN35_MTP_MODEL_TYPE, DrafterKind::Mtp); m }) } @@ -352,12 +362,27 @@ pub enum DrafterError { DraftBlockMissingHidden, } -/// Subset of the drafter's `config.json` that [`resolve_drafter_kind`] -/// needs. Unknown / extra fields are ignored. +/// Subset of the drafter's `config.json` that [`resolve_drafter_kind`] and +/// [`peek_qwen35_mtp_configured_block_size`] need. Unknown / extra fields +/// are ignored. #[derive(Debug, Deserialize)] struct DrafterConfigPeek { #[serde(default)] model_type: Option, + #[serde(default)] + block_size: Option, +} + +/// `model_type` value that identifies a Qwen 3.5 MTP drafter checkpoint +/// (see [`qwen3_5_mtp::Qwen35MtpConfig`]). Named so the block-size peek +/// below and the kind map in [`drafter_kind_by_model_type`] cannot drift +/// independently. +const QWEN35_MTP_MODEL_TYPE: &str = "qwen3_5_mtp"; + +fn read_drafter_config_peek(model_path: &Path) -> Option { + let cfg_path = model_path.join("config.json"); + let bytes = fs::read(&cfg_path).ok()?; + serde_json::from_slice::(&bytes).ok() } /// Read `model_path/config.json` and return the `model_type` field if it @@ -366,18 +391,41 @@ struct DrafterConfigPeek { /// OSError) -> None` behaviour, which is load-bearing for the DFlash /// fallback path (DFlash configs intentionally omit `model_type`). fn peek_drafter_model_type(model_path: &Path) -> Result, DrafterError> { - let cfg_path = model_path.join("config.json"); - let bytes = match fs::read(&cfg_path) { - Ok(b) => b, - // Treat any I/O error as "model_type unknown" — this matches the - // upstream exception-swallowing semantics in `_peek_drafter_model_type`. - Err(_) => return Ok(None), - }; - match serde_json::from_slice::(&bytes) { - Ok(peek) => Ok(peek.model_type), - // Malformed JSON: also treated as unknown, same as upstream. - Err(_) => Ok(None), + Ok(read_drafter_config_peek(model_path).and_then(|peek| peek.model_type)) +} + +/// Read `model_path/config.json` and return its top-level `block_size` +/// field, but ONLY when `model_type == "qwen3_5_mtp"`. This is a +/// deliberately narrow peek, not a general block-size reader, so a +/// same-named field in an unrelated drafter shape can never be +/// misinterpreted as a block-size hint. +/// +/// Used by `resolve_draft_block_size` in the `mlxcel` binary crate's +/// `cli::speculative_args` module (not reachable from this doc comment, +/// `mlxcel-core` sits below `mlxcel` in the dependency graph) so the +/// operator-visible `--draft-block-size` default for an MTP drafter +/// reflects the specific checkpoint's own configured value (the published +/// `mlx-community/Qwen3.8-27B-MTP-bf16` declares `3`) rather than the flat +/// Gemma-4-derived constant (`4`) that was previously applied to every +/// `DrafterKind::Mtp` drafter regardless of family. Issue #1165 hardening: +/// the PR's own measurement showed block 3 (drafter-configured) +/// outperforming block 4 (the flat constant) on the published pairing +/// (16.48 vs 13.81 tok/s, 0.591 vs 0.465 acceptance, see +/// `docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md`), so shipping +/// the flat constant as the silent default was shipping the slower +/// measured configuration. +/// +/// Returns `None` (falls back to the flat per-kind default) when the +/// config is missing, unparseable, not a Qwen 3.5 MTP drafter, or declares +/// no explicit `block_size` (the checkpoint then derives one from +/// `mtp_num_hidden_layers` at full load time, which this lightweight peek +/// does not replicate). +pub fn peek_qwen35_mtp_configured_block_size(model_path: &Path) -> Option { + let peek = read_drafter_config_peek(model_path)?; + if peek.model_type.as_deref() != Some(QWEN35_MTP_MODEL_TYPE) { + return None; } + peek.block_size } /// Reconcile the caller's `kind` choice with the drafter's actual @@ -627,9 +675,63 @@ pub trait Drafter { Vec::new() } + /// Prefill a **stateful** MTP drafter's own cache from the target's + /// prompt hidden states, and compute its first seed proposal. + /// + /// **Qwen 3.5 MTP only.** The Qwen MTP head accumulates one cache entry + /// per emitted target position; its history begins with the shifted + /// prompt (`prompt_tokens[1..] ++ first_bonus`) paired with the target's + /// post-final-norm hidden for the whole prompt (`[1, P, H]`). The MTP + /// round loop calls this once, on the first round of a session, and only + /// when the target advertised a full prompt hidden + /// ([`crate::speculative::mtp::target::MtpVerifyOutput::verify_hidden_full`]). + /// Stateless drafters (Gemma 4 assistant, DFlash) keep the no-op + /// default; a hook failure degrades draft quality only — the round loop + /// logs and continues with an empty drafter history. + #[allow(unused_variables)] + fn prefill_from_target_hidden( + &mut self, + prompt_tokens: &[i32], + hidden: &MlxArray, + first_bonus: i32, + sampler: &crate::generate::SamplingConfig, + ) -> Result<(), DrafterError> { + Ok(()) + } + + /// Extend a **stateful** MTP drafter's cache after a verify round. + /// + /// **Qwen 3.5 MTP only.** Called by the MTP round loop after the + /// speculative walk and the target-side rollback, with the target's full + /// post-final-norm verify hidden (`[1, block, H]`, one slot per verify + /// position), the round's draft proposals, the walk's `accepted` count, + /// and the emitted `new_tokens`. The drafter trims the rejected in-round + /// cache tail, appends the accepted tokens (paired with the target's + /// true hidden), and precomputes the next round's seed proposal. + /// Stateless drafters keep the no-op default; failures degrade draft + /// quality only. + #[allow(unused_variables)] + fn accept_verified_tokens( + &mut self, + verify_hidden: &MlxArray, + draft_tokens: &[i32], + accepted: usize, + new_tokens: &[i32], + sampler: &crate::generate::SamplingConfig, + ) -> Result<(), DrafterError> { + Ok(()) + } + /// Reset the drafter's own KV cache between full generation calls. /// - /// **DFlash- and InternalMtp-only.** Default no-op for MTP. + /// **DFlash- and InternalMtp-only** by original design; this trait + /// default (`Ok(())`, a no-op) is what the Gemma 4 assistant MTP + /// drafter uses, since it has no own KV cache to reset (shares K/V from + /// the target instead). The Qwen 3.5 MTP drafter (issue #1165) is the + /// exception: it DOES own accumulated KV history and overrides this to + /// destroy it on every reset (see `Qwen35MtpDraftModel::reset`), so + /// "default no-op for MTP" no longer holds for every `DrafterKind::Mtp` + /// implementation, only for the ones that do not override it. #[allow(unused_variables)] fn reset(&mut self, target: &dyn LanguageModel) -> Result<(), DrafterError> { Ok(()) @@ -817,9 +919,23 @@ pub fn load_drafter(path: &Path, kind: Option) -> Result { - // Wired in — Gemma 4 MTP assistant drafter. - let model = gemma4_assistant::Gemma4AssistantDraftModel::from_path(path)?; - Ok((Box::new(model), resolved)) + // Two concrete MTP drafter families share the kind; dispatch on + // the drafter's own `model_type`. `qwen3_5_mtp` loads the Qwen + // MTP head; the Gemma 4 assistant loader remains the fallback for + // the two `gemma4*_assistant` spellings AND for unknown/missing + // model_types under an explicit `--draft-kind mtp` (preserving + // the pre-existing behavior for hand-rolled Gemma fixtures). + let model_type = peek_drafter_model_type(path)?; + match model_type.as_deref() { + Some("qwen3_5_mtp") => { + let model = qwen3_5_mtp::Qwen35MtpDraftModel::from_path(path)?; + Ok((Box::new(model), resolved)) + } + _ => { + let model = gemma4_assistant::Gemma4AssistantDraftModel::from_path(path)?; + Ok((Box::new(model), resolved)) + } + } } // Remaining variant lands in peer — returning a typed // error rather than `unimplemented!()` here gives users an @@ -935,10 +1051,99 @@ mod tests { assert_eq!(map.get("gemma4_assistant"), Some(&DrafterKind::Mtp)); // The Gemma 4 Unified assistant resolves to the same MTP round loop. assert_eq!(map.get("gemma4_unified_assistant"), Some(&DrafterKind::Mtp)); - // Both assistant spellings, nothing else: parity with upstream - // `DRAFTER_KIND_BY_MODEL_TYPE = {"gemma4_assistant": "mtp", - // "gemma4_unified_assistant": "mtp"}`. - assert_eq!(map.len(), 2); + // The Qwen 3.5 / 3.6 / 3.8 split-out MTP head also resolves to the + // MTP round loop (upstream `"qwen3_5_mtp": "mtp"`). + assert_eq!(map.get("qwen3_5_mtp"), Some(&DrafterKind::Mtp)); + // The three MTP spellings, nothing else: parity with upstream + // `DRAFTER_KIND_BY_MODEL_TYPE`. + assert_eq!(map.len(), 3); + } + + #[test] + fn auto_detect_qwen3_5_mtp_resolves_to_mtp() { + let dir = tempdir().unwrap(); + write_drafter_config(&dir, Some("qwen3_5_mtp")); + let resolved = resolve(dir.path(), None).unwrap(); + assert_eq!(resolved, DrafterKind::Mtp); + } + + // ----- peek_qwen35_mtp_configured_block_size --------------------------- + + #[test] + fn peek_configured_block_size_reads_the_qwen35_mtp_value() { + let dir = tempdir().unwrap(); + fs::write( + dir.path().join("config.json"), + r#"{"model_type": "qwen3_5_mtp", "block_size": 3}"#, + ) + .expect("write config.json"); + assert_eq!( + super::peek_qwen35_mtp_configured_block_size(dir.path()), + Some(3) + ); + } + + #[test] + fn peek_configured_block_size_ignores_other_model_types() { + // A same-named `block_size` field on an unrelated (or Gemma 4) + // drafter must never be misread as a Qwen 3.5 MTP hint. + let dir = tempdir().unwrap(); + fs::write( + dir.path().join("config.json"), + r#"{"model_type": "gemma4_assistant", "block_size": 7}"#, + ) + .expect("write config.json"); + assert_eq!( + super::peek_qwen35_mtp_configured_block_size(dir.path()), + None + ); + } + + #[test] + fn peek_configured_block_size_is_none_when_absent_missing_or_unparseable() { + // No block_size field at all. + let dir = tempdir().unwrap(); + write_drafter_config(&dir, Some("qwen3_5_mtp")); + assert_eq!( + super::peek_qwen35_mtp_configured_block_size(dir.path()), + None + ); + + // Missing config.json entirely. + let empty_dir = tempdir().unwrap(); + assert_eq!( + super::peek_qwen35_mtp_configured_block_size(empty_dir.path()), + None + ); + + // Malformed JSON. + let bad_dir = tempdir().unwrap(); + fs::write(bad_dir.path().join("config.json"), "not json").expect("write config.json"); + assert_eq!( + super::peek_qwen35_mtp_configured_block_size(bad_dir.path()), + None + ); + } + + #[test] + fn load_drafter_routes_qwen3_5_mtp_to_qwen_loader() { + // A `qwen3_5_mtp` model_type must dispatch into + // `Qwen35MtpDraftModel::from_path`, not the Gemma 4 assistant loader. + // The bare fixture (no text_config) fails inside the Qwen config + // normalize with a message naming `text_config`, which the Gemma + // loader's error does not produce for this model_type. + let dir = tempdir().unwrap(); + write_drafter_config(&dir, Some("qwen3_5_mtp")); + let err = load(dir.path(), None).expect_err("bare fixture has no text_config"); + match err { + DrafterError::Config(msg) => { + assert!( + msg.contains("Qwen35MtpConfig"), + "expected the Qwen MTP loader's config error, got: {msg}" + ); + } + other => panic!("expected Config error from the Qwen MTP loader, got {other:?}"), + } } // ----- resolve_drafter_kind ------------------------------------------- diff --git a/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/config.rs b/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/config.rs new file mode 100644 index 000000000..e4c095955 --- /dev/null +++ b/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/config.rs @@ -0,0 +1,557 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Configuration for the Qwen 3.5 MTP drafter (`model_type: "qwen3_5_mtp"`). +//! +//! Mirrors the upstream `Qwen3_5MTPConfig` shape: a small top level +//! (`block_size`, `tie_word_embeddings`) plus a nested `text_config` that is a +//! full copy of the target's text config. See +//! https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/speculative/drafters/qwen3_5_mtp/config.py. +//! +//! The drafter exercises only a subset of the nested text config (the MTP +//! head is a single full-attention decoder layer sized on the target's +//! `hidden_size`), so [`Qwen35MtpTextConfig`] deserializes just that subset +//! and ignores the rest. The full target-side `Qwen35Config` lives in the +//! `mlxcel` binary crate, above `mlxcel-core` in the dependency graph, so it +//! cannot be reused here; field names, defaults, and the +//! `rope_parameters`-derived accessors mirror it exactly. + +use serde::Deserialize; + +/// Quantization arguments for a (hypothetical) quantized drafter checkpoint. +/// The published `mlx-community/Qwen3.8-27B-MTP-bf16` drafter is unquantized; +/// the field exists so a future quantized split loads without a config edit. +#[derive(Debug, Clone, Deserialize)] +pub struct Qwen35MtpQuantization { + pub group_size: i32, + pub bits: i32, +} + +/// Subset of the target-mirroring `text_config` the MTP drafter consumes. +/// +/// Unknown fields (`layer_types`, `linear_*`, MoE knobs, …) are ignored: the +/// MTP layer is always full attention (upstream builds it with +/// `full_attention_interval=1`), so none of the linear-attention or layout +/// fields participate. +#[derive(Debug, Clone, Deserialize)] +pub struct Qwen35MtpTextConfig { + #[serde(default)] + pub model_type: String, + pub hidden_size: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + #[serde(default)] + pub head_dim: Option, + #[serde(default = "default_rms_norm_eps")] + pub rms_norm_eps: f32, + pub intermediate_size: usize, + pub vocab_size: usize, + /// RoPE parameters in the same free-form dict shape the target's + /// `Qwen35Config` parses (`rope_theta`, `partial_rotary_factor`, plus + /// MRoPE keys the drafter ignores). + #[serde(default)] + pub rope_parameters: Option, + /// Number of MTP decoder layers. `1` on every published Qwen 3.5 / 3.6 / + /// 3.8 checkpoint. + #[serde(default = "default_mtp_num_hidden_layers")] + pub mtp_num_hidden_layers: usize, + #[serde(default)] + pub tie_word_embeddings: bool, + #[serde(default)] + pub quantization: Option, +} + +fn default_rms_norm_eps() -> f32 { + 1e-6 +} + +fn default_mtp_num_hidden_layers() -> usize { + 1 +} + +impl Qwen35MtpTextConfig { + /// Effective quantization group size (matches `Qwen35Config::group_size`). + pub fn group_size(&self) -> i32 { + self.quantization + .as_ref() + .map(|q| q.group_size) + .unwrap_or(64) + } + + /// Effective quantization bit width (matches `Qwen35Config::bits`). + pub fn bits(&self) -> i32 { + self.quantization.as_ref().map(|q| q.bits).unwrap_or(4) + } + + /// RoPE base frequency. Default mirrors the target-side + /// `Qwen35Config::rope_theta` fallback. + pub fn rope_theta(&self) -> f32 { + self.rope_parameters + .as_ref() + .and_then(|rp| rp.get("rope_theta")) + .and_then(|v| v.as_f64()) + .map(|v| v as f32) + .unwrap_or(100_000.0) + } + + /// Partial rotary factor. Default mirrors the target-side + /// `Qwen35Config::partial_rotary_factor` fallback. + pub fn partial_rotary_factor(&self) -> f32 { + self.rope_parameters + .as_ref() + .and_then(|rp| rp.get("partial_rotary_factor")) + .and_then(|v| v.as_f64()) + .map(|v| v as f32) + .unwrap_or(0.25) + } + + /// Per-head dimension, falling back to `hidden_size / num_attention_heads`. + /// + /// The fallback is lazy and division-safe. It used to be + /// `unwrap_or(self.hidden_size / self.num_attention_heads)`, and + /// `Option::unwrap_or` takes a *value*: the division was evaluated on every + /// call, including when `head_dim` was `Some`. A `num_attention_heads: 0` in + /// a drafter `config.json` therefore panicked here with "attempt to divide + /// by zero" even for a config that otherwise specified the head geometry in + /// full. That mattered because the drafter is loaded from disk inside the + /// batch worker at the first speculative request rather than at startup, and + /// every core inference worker runs under `run_core_thread_or_abort`, which + /// converts an uncaught panic into `std::process::abort()`: a config scalar + /// took down the whole server process. + /// + /// [`Qwen35MtpConfig::normalize`] is the real gate (it rejects a zero head + /// count with a named error before any weight is touched). The `checked_div` + /// here is defense in depth for any path that resolves head geometry without + /// having normalized first: 0 is a benign sentinel because every downstream + /// consumer of a head dim is a reshape or a projection load that rejects it + /// with a shape error, which is recoverable where a panic is not. + pub fn head_dim_resolved(&self) -> usize { + self.head_dim.unwrap_or_else(|| { + self.hidden_size + .checked_div(self.num_attention_heads) + .unwrap_or(0) + }) + } + + /// Rotary dimensions per head (`head_dim * partial_rotary_factor`), + /// matching the target-side `Qwen35Config::rope_dims`. + pub fn rope_dims(&self) -> i32 { + (self.head_dim_resolved() as f32 * self.partial_rotary_factor()) as i32 + } + + /// Whether the mirrored target family is a MoE variant. The dense MTP + /// drafter does not implement the MoE decoder layer, so loaders reject + /// this early with a named error instead of failing on a missing + /// `switch_mlp` weight. + pub fn is_moe(&self) -> bool { + self.model_type.contains("moe") + } + + /// Bounds-check the numeric fields the drafter actually consumes. + /// + /// Called from [`Qwen35MtpConfig::normalize`], which is the single gate + /// every load path goes through, so a rejected value never reaches a + /// division, a reshape, or an allocation. Checks are ordered so the most + /// specific message wins: the zero checks on the two head counts precede the + /// GQA divisibility check, whose `%` would itself divide by zero on + /// `num_key_value_heads: 0`. + fn validate_bounds(&self) -> Result<(), String> { + if self.hidden_size == 0 { + return Err(format!( + "Qwen35MtpConfig.text_config.hidden_size ({}) must be > 0: it sizes every \ + projection and the residual stream", + self.hidden_size + )); + } + if self.num_attention_heads == 0 { + return Err(format!( + "Qwen35MtpConfig.text_config.num_attention_heads ({}) must be > 0: it is the \ + divisor in the head_dim fallback and a reshape extent in the MTP attention \ + layer, so zero is an uncatchable abort at first draft rather than a load error", + self.num_attention_heads + )); + } + if self.num_key_value_heads == 0 { + return Err(format!( + "Qwen35MtpConfig.text_config.num_key_value_heads ({}) must be > 0: it is a \ + reshape extent for the K/V projections in the MTP attention layer", + self.num_key_value_heads + )); + } + if !self + .num_attention_heads + .is_multiple_of(self.num_key_value_heads) + { + return Err(format!( + "Qwen35MtpConfig.text_config.num_attention_heads ({}) must be a multiple of \ + num_key_value_heads ({}): the MTP attention layer reshapes on both counts and \ + repeats K/V across each query group, which has no meaning for a ragged layout", + self.num_attention_heads, self.num_key_value_heads + )); + } + if let Some(head_dim) = self.head_dim + && head_dim == 0 + { + return Err( + "Qwen35MtpConfig.text_config.head_dim (0) must be > 0 when present: it is a \ + reshape extent and the denominator of the attention scale" + .to_string(), + ); + } + if self.vocab_size == 0 { + return Err(format!( + "Qwen35MtpConfig.text_config.vocab_size ({}) must be > 0: it sizes the drafter's \ + embedding table and LM head", + self.vocab_size + )); + } + if self.intermediate_size == 0 { + return Err(format!( + "Qwen35MtpConfig.text_config.intermediate_size ({}) must be > 0: it sizes the \ + SwiGLU MLP in the MTP decoder layer", + self.intermediate_size + )); + } + if !(1..=MAX_MTP_NUM_HIDDEN_LAYERS).contains(&self.mtp_num_hidden_layers) { + return Err(format!( + "Qwen35MtpConfig.text_config.mtp_num_hidden_layers ({}) must be between 1 and \ + {MAX_MTP_NUM_HIDDEN_LAYERS}: the count is used as an allocation size before any \ + weight is read, so an unbounded value is an allocation-failure abort rather than \ + a load error, and 0 layers builds a drafter that cannot draft", + self.mtp_num_hidden_layers + )); + } + Ok(()) + } +} + +/// Drafter config for the Qwen 3.5 MTP head. +/// +/// Mirrors upstream `Qwen3_5MTPConfig`: `block_size` defaults to +/// `mtp_num_hidden_layers + 2` when absent (the **total** verify-round budget +/// including the bonus token, so the published `block_size: 3` checkpoints +/// draft 2 tokens per round). +#[derive(Debug, Clone, Deserialize)] +pub struct Qwen35MtpConfig { + #[serde(default = "default_model_type")] + pub model_type: String, + /// Total round budget including the bonus token. `None` in JSON resolves + /// to `mtp_num_hidden_layers + 2` in [`Self::normalize`], mirroring + /// upstream `from_dict`'s `flat.setdefault("block_size", mtp_depth + 2)`. + #[serde(default)] + pub block_size: Option, + #[serde(default = "default_tie_word_embeddings")] + pub tie_word_embeddings: bool, + pub text_config: Option, +} + +fn default_model_type() -> String { + "qwen3_5_mtp".to_string() +} + +fn default_tie_word_embeddings() -> bool { + true +} + +/// Upper bound on `text_config.mtp_num_hidden_layers`. +/// +/// This is an allocation bound, not a claim about how deep an MTP head can +/// usefully be. Every published Qwen 3.5 / 3.6 / 3.8 MTP checkpoint ships +/// `mtp_num_hidden_layers: 1`, and upstream derives `block_size` from it as +/// `mtp_num_hidden_layers + 2`, so the field is a draft depth rather than a +/// model dimension: it does not grow with parameter count the way +/// `num_hidden_layers` does. +/// +/// The bound is needed because the count is consumed as an allocation size +/// *before* any weight is consulted. `Qwen35MtpDraftModel::from_weights` builds +/// an expected-weight inventory of eleven formatted `String`s per layer and a +/// `Vec::with_capacity` of this size, so a hostile `mtp_num_hidden_layers: +/// 100000000` asks for on the order of 1.1e9 allocations up front. Allocation +/// failure in Rust is an abort, not a catchable panic, and the drafter loads +/// inside the batch worker at first request, so that is a whole-server abort +/// driven by a config scalar. +/// +/// 8 leaves eight times the headroom over every checkpoint that exists while +/// keeping the worst case at a few hundred small allocations. If a real +/// checkpoint ever ships a deeper head, raise this constant: that makes the +/// change one reviewed line rather than a silently unbounded capacity, and the +/// rejection names the field and the cap so the reason is obvious from the +/// error alone. +const MAX_MTP_NUM_HIDDEN_LAYERS: usize = 8; + +impl Qwen35MtpConfig { + /// Validate and apply the upstream post-init fixups: + /// + /// - `text_config` must be present (upstream raises `ValueError`). + /// - `block_size` defaults to `mtp_num_hidden_layers + 2`. + /// - `tie_word_embeddings` follows `text_config.tie_word_embeddings` + /// (upstream `__post_init__`). + /// - MoE text configs are rejected: the dense drafter has no + /// `Qwen3_5MoeDecoderLayer` port. + /// - Numeric fields are bounds-checked (see below). + /// + /// The numeric checks exist because two `text_config` scalars reach code + /// that aborts the process rather than returning an error. A + /// `num_attention_heads: 0` divides by zero in + /// [`Qwen35MtpTextConfig::head_dim_resolved`], and an unbounded + /// `mtp_num_hidden_layers` is used as an allocation size in + /// `Qwen35MtpDraftModel::from_weights` before any weight is read (see + /// [`MAX_MTP_NUM_HIDDEN_LAYERS`]). The drafter is loaded lazily inside the + /// batch worker at the first speculative request, and core worker threads + /// run under `run_core_thread_or_abort`, which turns an uncaught panic into + /// `std::process::abort()`; an allocation failure aborts outright. So a + /// malformed drafter `config.json` used to take down the whole server on + /// first request instead of declining to classic decode with a named error. + /// + /// Rejecting at normalize time is the same policy the tree already applies + /// to quantization scalars in `crate::layers::validate_quantization_params` + /// (issue #929): refuse a config-driven value that can only fail + /// uncatchably later, at parse time, with the field and the value in the + /// message. + pub fn normalize(mut self) -> Result { + let text_cfg = self + .text_config + .as_ref() + .ok_or_else(|| "Qwen35MtpConfig.text_config must be set".to_string())?; + if text_cfg.is_moe() { + return Err(format!( + "qwen3_5_mtp drafter: MoE text_config (model_type {:?}) is not supported; \ + only the dense MTP decoder layer is implemented", + text_cfg.model_type + )); + } + text_cfg.validate_bounds()?; + if self.block_size.is_none() { + self.block_size = Some(text_cfg.mtp_num_hidden_layers + 2); + } + self.tie_word_embeddings = text_cfg.tie_word_embeddings; + Ok(self) + } + + /// Nested text config accessor. Call after [`Self::normalize`]. + pub fn text_config(&self) -> &Qwen35MtpTextConfig { + self.text_config + .as_ref() + .expect("Qwen35MtpConfig.text_config must be set (call normalize first)") + } + + /// Resolved total round budget (bonus token included). Call after + /// [`Self::normalize`]. + pub fn block_size(&self) -> usize { + self.block_size + .expect("Qwen35MtpConfig.block_size resolved by normalize") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The published `mlx-community/Qwen3.8-27B-MTP-bf16` config shape: + /// `block_size: 3` at top level, `text_config` mirroring the target + /// (with keys the drafter ignores), `mtp_num_hidden_layers: 1`. + #[test] + fn parses_published_qwen38_drafter_config_shape() { + let json = r#"{ + "block_size": 3, + "model_type": "qwen3_5_mtp", + "tie_word_embeddings": false, + "text_config": { + "model_type": "qwen3_5_text", + "attn_output_gate": true, + "full_attention_interval": 4, + "head_dim": 256, + "hidden_size": 5120, + "intermediate_size": 17408, + "layer_types": ["linear_attention", "full_attention"], + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": false, + "num_attention_heads": 24, + "num_key_value_heads": 4, + "num_hidden_layers": 64, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "mrope_interleaved": true, + "mrope_section": [11, 11, 10], + "partial_rotary_factor": 0.25, + "rope_theta": 10000000, + "rope_type": "default" + }, + "tie_word_embeddings": false, + "vocab_size": 248320 + } + }"#; + let cfg: Qwen35MtpConfig = serde_json::from_str(json).expect("parse config"); + let cfg = cfg.normalize().expect("normalize"); + assert_eq!(cfg.model_type, "qwen3_5_mtp"); + assert_eq!(cfg.block_size(), 3); + assert!(!cfg.tie_word_embeddings); + let tc = cfg.text_config(); + assert_eq!(tc.hidden_size, 5120); + assert_eq!(tc.num_attention_heads, 24); + assert_eq!(tc.num_key_value_heads, 4); + assert_eq!(tc.head_dim_resolved(), 256); + assert_eq!(tc.mtp_num_hidden_layers, 1); + assert_eq!(tc.vocab_size, 248320); + assert_eq!(tc.rope_theta(), 10_000_000.0); + // 256 * 0.25 = 64 rotary dims — the partial-rotary geometry the + // target's full-attention layers use. + assert_eq!(tc.rope_dims(), 64); + } + + /// `block_size` omitted resolves to `mtp_num_hidden_layers + 2` + /// (upstream `from_dict` setdefault). + #[test] + fn block_size_defaults_to_mtp_depth_plus_two() { + let json = r#"{ + "model_type": "qwen3_5_mtp", + "text_config": { + "model_type": "qwen3_5_text", + "hidden_size": 64, + "intermediate_size": 128, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "vocab_size": 512, + "mtp_num_hidden_layers": 1 + } + }"#; + let cfg: Qwen35MtpConfig = serde_json::from_str(json).expect("parse"); + let cfg = cfg.normalize().expect("normalize"); + assert_eq!(cfg.block_size(), 3); + } + + #[test] + fn normalize_rejects_missing_text_config() { + let cfg: Qwen35MtpConfig = + serde_json::from_str(r#"{"model_type": "qwen3_5_mtp"}"#).expect("parse"); + let err = cfg.normalize().expect_err("must reject"); + assert!(err.contains("text_config")); + } + + #[test] + fn normalize_rejects_moe_text_config() { + let json = r#"{ + "model_type": "qwen3_5_mtp", + "text_config": { + "model_type": "qwen3_5_moe_text", + "hidden_size": 64, + "intermediate_size": 128, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "vocab_size": 512 + } + }"#; + let cfg: Qwen35MtpConfig = serde_json::from_str(json).expect("parse"); + let err = cfg.normalize().expect_err("must reject MoE"); + assert!(err.contains("MoE"), "got: {err}"); + } + + /// `num_attention_heads: 0` must be rejected at normalize time, not divided + /// by at first draft. + /// + /// `head_dim` is explicitly `256` here, which is the point of the fixture: + /// the config looks fully specified, so nothing should ever need the + /// `hidden_size / num_attention_heads` fallback. The old + /// `unwrap_or(hidden_size / num_attention_heads)` evaluated that division + /// eagerly anyway, and `head_dim_resolved()` panicked with "attempt to + /// divide by zero" on a config that never asked for the fallback. + #[test] + fn normalize_rejects_zero_attention_heads() { + let json = r#"{ + "model_type": "qwen3_5_mtp", + "text_config": { + "model_type": "qwen3_5_text", + "head_dim": 256, + "hidden_size": 64, + "intermediate_size": 128, + "num_attention_heads": 0, + "num_key_value_heads": 2, + "vocab_size": 512 + } + }"#; + let cfg: Qwen35MtpConfig = serde_json::from_str(json).expect("parse"); + let err = cfg + .normalize() + .expect_err("must reject zero attention heads"); + assert!(err.contains("num_attention_heads"), "got: {err}"); + } + + /// [`Qwen35MtpTextConfig::head_dim_resolved`] must not panic even when + /// called without a preceding `normalize()`. Reaching this assertion at all + /// is the test: the pre-fix implementation aborted the process here. + #[test] + fn head_dim_resolved_does_not_divide_by_zero() { + let tc = Qwen35MtpTextConfig { + model_type: "qwen3_5_text".to_string(), + hidden_size: 64, + num_attention_heads: 0, + num_key_value_heads: 2, + head_dim: None, + rms_norm_eps: 1e-6, + intermediate_size: 128, + vocab_size: 512, + rope_parameters: None, + mtp_num_hidden_layers: 1, + tie_word_embeddings: false, + quantization: None, + }; + assert_eq!(tc.head_dim_resolved(), 0); + assert_eq!(tc.rope_dims(), 0); + } + + /// An unbounded `mtp_num_hidden_layers` is an allocation size the model + /// loader consumes before reading any weight, so it is bounded here. + #[test] + fn normalize_rejects_absurd_mtp_layer_count() { + let json = r#"{ + "model_type": "qwen3_5_mtp", + "text_config": { + "model_type": "qwen3_5_text", + "hidden_size": 64, + "intermediate_size": 128, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "vocab_size": 512, + "mtp_num_hidden_layers": 100000000 + } + }"#; + let cfg: Qwen35MtpConfig = serde_json::from_str(json).expect("parse"); + let err = cfg + .normalize() + .expect_err("must reject absurd mtp layer count"); + assert!(err.contains("mtp_num_hidden_layers"), "got: {err}"); + assert!(err.contains("100000000"), "got: {err}"); + } + + /// GQA layout: `Qwen35MtpAttention` reshapes on both head counts and + /// repeats K/V across each query group, so a ragged layout is rejected. + #[test] + fn normalize_rejects_non_divisible_gqa_head_counts() { + let json = r#"{ + "model_type": "qwen3_5_mtp", + "text_config": { + "model_type": "qwen3_5_text", + "hidden_size": 64, + "intermediate_size": 128, + "num_attention_heads": 6, + "num_key_value_heads": 4, + "vocab_size": 512 + } + }"#; + let cfg: Qwen35MtpConfig = serde_json::from_str(json).expect("parse"); + let err = cfg.normalize().expect_err("must reject ragged GQA layout"); + assert!(err.contains("num_attention_heads"), "got: {err}"); + assert!(err.contains("num_key_value_heads"), "got: {err}"); + } +} diff --git a/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/layer.rs b/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/layer.rs new file mode 100644 index 000000000..c13e5d2be --- /dev/null +++ b/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/layer.rs @@ -0,0 +1,263 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Decoder layer for the Qwen 3.5 MTP drafter. +//! +//! Upstream builds each MTP layer as `Qwen3_5DecoderLayer(replace(text_config, +//! num_hidden_layers=mtp_layers, full_attention_interval=1))` — the +//! `full_attention_interval=1` forces every MTP layer onto the full-attention +//! path, never gated-delta. The published Qwen3.8-27B MTP tensors confirm it +//! (`self_attn.*` only, no `linear_attn.*`). +//! +//! This is an in-crate port of the target-side full-attention layer +//! (`Qwen3NextAttention` + dense `MLP` in `src/models/qwen3_next.rs` of the +//! `mlxcel` binary crate, which sits above `mlxcel-core` and is therefore not +//! reachable from here — the same layering constraint that gave DFlash its own +//! in-crate attention/MLP). Differences from the target-side layer: +//! +//! - No MRoPE branch: the drafter never sees multimodal position ids +//! (upstream's drafter passes plain integer `position_ids`, and the +//! target-side DFlash speculative forward takes the same standard-RoPE path). +//! - The RoPE offset is an explicit `rope_offset` argument instead of +//! `cache.offset`: the drafter tracks its own logical `_next_position`, +//! which can run ahead of its cache length when the prompt-side drafter +//! prefill was skipped (adopted prompt prefix, or a reset between slice +//! grants). Cache storage geometry still comes from the cache itself. +//! - No `target_verify` per-position path: drafter outputs never need +//! bit-parity with single-token decode — greedy parity is enforced on the +//! TARGET's verify pass, and drafter numerics only influence acceptance. +//! +//! The Q projection carries the Qwen 3.5 `attn_output_gate` layout: `q_proj` +//! emits queries and a sigmoid gate concatenated per head +//! (`[B, L, num_heads, 2 * head_dim]`), and the attention output is +//! `o_proj(attn * sigmoid(gate))`. + +use crate::cache::KVCache; +use crate::drafter::dflash::mlp::DFlashMlp; +use crate::ffi::{self, MlxArray}; +use crate::layers::{RMSNorm, UnifiedLinear}; +use crate::weights::WeightMap; +use cxx::UniquePtr; + +use super::config::Qwen35MtpTextConfig; + +/// Full-attention block of one MTP decoder layer (gated-Q Qwen 3.5 layout). +pub struct Qwen35MtpAttention { + q_proj: UnifiedLinear, + k_proj: UnifiedLinear, + v_proj: UnifiedLinear, + o_proj: UnifiedLinear, + q_norm: RMSNorm, + k_norm: RMSNorm, + num_heads: i32, + num_kv_heads: i32, + head_dim: i32, + scale: f32, + rope_dims: i32, + rope_base: f32, +} + +impl Qwen35MtpAttention { + /// Attention forward over the drafter's own KV cache. + /// + /// - `x`: `[B, L, hidden_size]` input (post `input_layernorm`). + /// - `mask`: `Some(causal mask)` for multi-token forwards (drafter + /// prompt prefill and accepted-token extension), `None` for the + /// single-token draft steps. + /// - `cache`: this layer's own KV cache; keys/values are appended. + /// - `rope_offset`: absolute position of `x`'s first token in the + /// target sequence (the drafter's `next_position`). + pub fn forward( + &self, + x: &MlxArray, + mask: Option<&MlxArray>, + cache: &mut KVCache, + rope_offset: i32, + ) -> UniquePtr { + let shape = ffi::array_shape(x); + let b = shape[0]; + let l = shape[1]; + + // Gated Q projection: [B, L, 2 * num_heads * head_dim], reshaped to + // [B, L, num_heads, 2 * head_dim] and split into queries + gate. + let q_proj_output = self.q_proj.forward(x); + let q_proj_reshaped = ffi::reshape(&q_proj_output, &[b, l, self.num_heads, -1]); + let queries = ffi::slice( + &q_proj_reshaped, + &[0, 0, 0, 0], + &[b, l, self.num_heads, self.head_dim], + ); + let q_last_dim = ffi::array_shape(&q_proj_reshaped)[3]; + let gate = ffi::slice( + &q_proj_reshaped, + &[0, 0, 0, self.head_dim], + &[b, l, self.num_heads, q_last_dim], + ); + let gate = ffi::reshape(&gate, &[b, l, -1]); + + let keys = self.k_proj.forward(x); + let values = self.v_proj.forward(x); + + let queries = ffi::reshape(&queries, &[b, l, self.num_heads, self.head_dim]); + let keys = ffi::reshape(&keys, &[b, l, self.num_kv_heads, self.head_dim]); + let values = ffi::reshape(&values, &[b, l, self.num_kv_heads, self.head_dim]); + + // Per-head RMSNorm over head_dim, applied before transpose — same + // order as the target-side layer. + let queries = self.q_norm.forward(&queries); + let keys = self.k_norm.forward(&keys); + + let queries = ffi::transpose_axes(&queries, &[0, 2, 1, 3]); + let keys = ffi::transpose_axes(&keys, &[0, 2, 1, 3]); + let values = ffi::transpose_axes(&values, &[0, 2, 1, 3]); + + // Partial rotary embedding at the drafter's logical position. + let queries = ffi::fast_rope( + &queries, + self.rope_dims, + false, + self.rope_base, + 1.0, + rope_offset, + ); + let keys = ffi::fast_rope( + &keys, + self.rope_dims, + false, + self.rope_base, + 1.0, + rope_offset, + ); + + let (cache_k, cache_v) = cache.update_and_fetch(keys, values); + + let attn_out = + crate::layers::attention(&queries, &cache_k, &cache_v, self.scale, mask, 0.0, 0); + + let output = ffi::transpose_axes(&attn_out, &[0, 2, 1, 3]); + let output = ffi::reshape(&output, &[b, l, -1]); + + // Sigmoid output gate, then o_proj — the `attn_output_gate` path. + let gate_sigmoid = ffi::sigmoid(&gate); + let gated = ffi::multiply(&output, &gate_sigmoid); + self.o_proj.forward(&gated) + } + + /// Load one layer's attention weights from `{prefix}.{q,k,v,o}_proj` and + /// `{prefix}.{q,k}_norm`. + pub fn from_weights( + weights: &WeightMap, + prefix: &str, + config: &Qwen35MtpTextConfig, + ) -> Result { + let group_size = config.group_size(); + let bits = config.bits(); + + let q_proj = + UnifiedLinear::from_weights(weights, &format!("{prefix}.q_proj"), group_size, bits)?; + let k_proj = + UnifiedLinear::from_weights(weights, &format!("{prefix}.k_proj"), group_size, bits)?; + let v_proj = + UnifiedLinear::from_weights(weights, &format!("{prefix}.v_proj"), group_size, bits)?; + let o_proj = + UnifiedLinear::from_weights(weights, &format!("{prefix}.o_proj"), group_size, bits)?; + + let q_norm_w = weights + .get(&format!("{prefix}.q_norm.weight")) + .map(|w| ffi::copy(w)) + .ok_or_else(|| format!("Weight not found: {prefix}.q_norm.weight"))?; + let k_norm_w = weights + .get(&format!("{prefix}.k_norm.weight")) + .map(|w| ffi::copy(w)) + .ok_or_else(|| format!("Weight not found: {prefix}.k_norm.weight"))?; + + let head_dim = config.head_dim_resolved() as i32; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + q_norm: RMSNorm::new(q_norm_w, config.rms_norm_eps), + k_norm: RMSNorm::new(k_norm_w, config.rms_norm_eps), + num_heads: config.num_attention_heads as i32, + num_kv_heads: config.num_key_value_heads as i32, + head_dim, + scale: 1.0 / (head_dim as f32).sqrt(), + rope_dims: config.rope_dims(), + rope_base: config.rope_theta(), + }) + } +} + +/// One MTP decoder layer: pre-norm attention + pre-norm SwiGLU MLP with +/// residual connections, mirroring the target-side `Qwen35DecoderLayer` +/// full-attention arm. +pub struct Qwen35MtpDecoderLayer { + attention: Qwen35MtpAttention, + mlp: DFlashMlp, + input_layernorm: RMSNorm, + post_attention_layernorm: RMSNorm, +} + +impl Qwen35MtpDecoderLayer { + /// Layer forward. Argument semantics match + /// [`Qwen35MtpAttention::forward`]. + pub fn forward( + &self, + x: &MlxArray, + mask: Option<&MlxArray>, + cache: &mut KVCache, + rope_offset: i32, + ) -> UniquePtr { + let normed = self.input_layernorm.forward(x); + let r = self.attention.forward(&normed, mask, cache, rope_offset); + let h = ffi::add(x, &r); + let mlp_out = self.mlp.forward(&self.post_attention_layernorm.forward(&h)); + ffi::add(&h, &mlp_out) + } + + /// Load one layer from `{prefix}.self_attn.*`, `{prefix}.mlp.*`, + /// `{prefix}.input_layernorm.weight`, + /// `{prefix}.post_attention_layernorm.weight`. + pub fn from_weights( + weights: &WeightMap, + prefix: &str, + config: &Qwen35MtpTextConfig, + ) -> Result { + let attention = + Qwen35MtpAttention::from_weights(weights, &format!("{prefix}.self_attn"), config)?; + let mlp = DFlashMlp::from_weights( + weights, + &format!("{prefix}.mlp"), + config.group_size(), + config.bits(), + )?; + + let input_norm_w = weights + .get(&format!("{prefix}.input_layernorm.weight")) + .map(|w| ffi::copy(w)) + .ok_or_else(|| format!("Weight not found: {prefix}.input_layernorm.weight"))?; + let post_norm_w = weights + .get(&format!("{prefix}.post_attention_layernorm.weight")) + .map(|w| ffi::copy(w)) + .ok_or_else(|| format!("Weight not found: {prefix}.post_attention_layernorm.weight"))?; + + Ok(Self { + attention, + mlp, + input_layernorm: RMSNorm::new(input_norm_w, config.rms_norm_eps), + post_attention_layernorm: RMSNorm::new(post_norm_w, config.rms_norm_eps), + }) + } +} diff --git a/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/mod.rs b/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/mod.rs new file mode 100644 index 000000000..86791eabf --- /dev/null +++ b/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/mod.rs @@ -0,0 +1,40 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Qwen 3.5 / 3.6 / 3.8 Multi-Token Prediction (MTP) drafter — the split-out +//! MTP head (`fc` fusion + one full-attention decoder layer) that pairs with +//! `qwen3_5`-family targets for MTP speculative decoding. +//! +//! Top-level overview: +//! +//! - [`config`] — `Qwen35MtpConfig` (top-level `block_size`) and the +//! target-mirroring `Qwen35MtpTextConfig` subset. +//! - [`layer`] — `Qwen35MtpDecoderLayer` (gated-Q full attention + SwiGLU +//! MLP, drafter-owned KV cache). +//! - [`model`] — `Qwen35MtpDraftModel` implementing +//! [`crate::drafter::Drafter`], including the stateful +//! prompt-prefill / accept-verified lifecycle hooks. +//! +//! Upstream reference: +//! https://github.com/Blaizzy/mlx-vlm/tree/main/mlx_vlm/speculative/drafters/qwen3_5_mtp. + +pub mod config; +pub mod layer; +pub mod model; + +#[cfg(test)] +mod tests; + +pub use config::{Qwen35MtpConfig, Qwen35MtpTextConfig}; +pub use model::Qwen35MtpDraftModel; diff --git a/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/model.rs b/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/model.rs new file mode 100644 index 000000000..4a056a7c3 --- /dev/null +++ b/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/model.rs @@ -0,0 +1,739 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! `Qwen35MtpDraftModel` — Rust port of the Qwen 3.5 / 3.6 / 3.8 MTP drafter. +//! +//! Mirrors +//! https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/speculative/drafters/qwen3_5_mtp/qwen3_5_mtp.py +//! (B = 1 surface; the batched drafter path is out of scope and declines via +//! the trait's erroring `draft_block_batched` default). +//! +//! ## Lifecycle +//! +//! Unlike the Gemma 4 assistant (stateless per round, cross-attends into the +//! target's shared K/V), the Qwen MTP drafter is **stateful**: it owns one KV +//! cache per MTP layer and accumulates one entry per emitted target position, +//! exactly like the trained MTP head does. Position `i` of the drafter's +//! sequence is the pair `(token_{i+1} embedding, target hidden at i)` fused +//! through `fc`. +//! +//! 1. **Load.** [`Qwen35MtpDraftModel::from_path`] parses `config.json`, +//! loads and sanitizes the 15-tensor head, converts bf16 → f16 (Apple +//! Silicon precision rules), and constructs the model. +//! 2. **Bind.** [`Drafter::bind`] borrows the target's `embed_tokens` and +//! `lm_head` (falling back to the target embedding's `as_linear` for tied +//! checkpoints). The drafter owns neither. +//! 3. **Prompt prefill.** [`Drafter::prefill_from_target_hidden`] runs the +//! shifted prompt (`prompt[1..] ++ bonus`) paired with the target's +//! post-final-norm prompt hidden through the MTP stack, populating the +//! drafter cache and computing the first seed token. +//! 4. **Rounds.** [`Drafter::set_shared_kv`] discards the tensors (Qwen has +//! no shared-K/V concept) and only aligns positions; +//! [`Drafter::draft_block`] consumes the seed and runs small +//! autoregressive forwards; after the target verify, +//! [`Drafter::accept_verified_tokens`] trims the rejected in-round cache +//! tail and extends the cache with the accepted tokens paired with the +//! target's verify hidden, computing the next seed. +//! +//! A drafter whose cache was cleared mid-session (slice-grant rotation resets +//! the shared worker handle, or a hook failure poisoned the state) degrades +//! gracefully: `set_shared_kv` re-anchors `next_position` to the target's +//! cache offset and drafts continue with an empty history. That costs +//! acceptance, never correctness — greedy parity is enforced by the target's +//! verify pass alone. + +use crate::cache::KVCache; +use crate::drafter::{Drafter, DrafterError, DrafterKind, SharedKv}; +use crate::ffi::{self, MlxArray}; +use crate::generate::{LanguageModel, SamplingConfig}; +use crate::layers::{RMSNorm, UnifiedEmbedding, UnifiedLinear}; +use crate::weights::WeightMap; +use cxx::UniquePtr; +use std::path::Path; + +use super::config::Qwen35MtpConfig; +use super::layer::Qwen35MtpDecoderLayer; + +/// LM head borrowed from the target at `bind()` time. +enum MtpLmHead { + /// Target has an untied `lm_head` (Qwen3.8-27B). + Linear(UnifiedLinear), + /// Tied checkpoints project through the target embedding table. + TiedEmbed(UnifiedEmbedding), +} + +/// Qwen 3.5 MTP drafter — `fc`-fused single-decoder-layer head with its own +/// per-layer KV cache, borrowing embeddings and the LM head from the target. +/// +/// Implements [`Drafter`] and is wired into the `Mtp` arm of +/// [`crate::drafter::load_drafter`] for `model_type == "qwen3_5_mtp"`. +pub struct Qwen35MtpDraftModel { + config: Qwen35MtpConfig, + fc: UnifiedLinear, + pre_fc_norm_embedding: RMSNorm, + pre_fc_norm_hidden: RMSNorm, + norm: RMSNorm, + layers: Vec, + + /// Per-layer drafter-owned KV cache. One entry per MTP layer. + cache: Vec, + + /// Target's embedding table, captured by `bind()`. + target_embed: Option, + /// Upstream reads `inner.embed_scale` defaulting to 1.0; Qwen 3.5 has no + /// embedding scale, so this stays 1.0 and the multiply is skipped. + target_embed_scale: f32, + /// Target's LM head, captured by `bind()`. + lm_head: Option, + + /// Precomputed seed for the next `draft_block`: the MTP head's prediction + /// one past the last accepted token, computed by the prompt prefill / + /// accept hooks from TARGET hidden (so its cache entry uses the true + /// pair, unlike in-round entries that substitute the drafter's own + /// hidden). + seed_token: Option, + seed_hidden: Option>, + + /// Absolute target-sequence position of the next drafter cache append. + /// Tracks the target cache offset exactly while state is intact; also the + /// RoPE offset for the next forward. + next_position: i32, + /// Number of in-round cache appends since the last accept, i.e. how many + /// entries the current draft block added. The accept hook keeps + /// `min(accepted, round_appended)` of them and trims the rest. + round_appended: i32, + /// Diagnostics only — the trait's `kv_offset` / `position` arguments. + kv_valid_len: i32, + position: i32, +} + +impl std::fmt::Debug for Qwen35MtpDraftModel { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Qwen35MtpDraftModel") + .field("model_type", &self.config.model_type) + .field("block_size", &self.config.block_size()) + .field("num_layers", &self.layers.len()) + .field("bound", &self.lm_head.is_some()) + .field("cache_offset", &self.cache.first().map(|c| c.offset)) + .field("next_position", &self.next_position) + .field("round_appended", &self.round_appended) + .field("has_seed", &self.seed_token.is_some()) + .finish() + } +} + +impl Qwen35MtpDraftModel { + /// Construct from a checkpoint directory containing `config.json` and + /// safetensors shards. Used by [`crate::drafter::load_drafter`]'s `Mtp` + /// arm when the drafter's `model_type` is `qwen3_5_mtp`. + pub fn from_path(path: &Path) -> Result { + let cfg_path = path.join("config.json"); + let bytes = std::fs::read(&cfg_path).map_err(|e| DrafterError::ConfigIo { + path: cfg_path.display().to_string(), + source: e, + })?; + let config: Qwen35MtpConfig = + serde_json::from_slice(&bytes).map_err(|e| DrafterError::ConfigParse { + path: cfg_path.display().to_string(), + source: e, + })?; + let config = config.normalize().map_err(DrafterError::Config)?; + + let mut weights = crate::weights::load_weights_from_dir(path) + .map_err(|reason| DrafterError::WeightLoad { reason })?; + Self::sanitize_weights(&mut weights); + // Apple Silicon precision: bf16 → f16 on non-quantized tensors at + // the weight-loading boundary, matching the target model loaders and + // the DFlash drafter loader. The published drafter is bf16; the + // paired 4-bit target's activations are f16, and a dtype-mixed + // concat/matmul would silently promote to f32. + crate::drafter::dflash::drafter::convert_bf16_to_f16_non_quantized(&mut weights); + Self::from_weights(&weights, config) + } + + /// Construct from an in-memory weight map (already sanitized). Used by + /// `from_path` and unit tests with synthetic fixtures. + pub fn from_weights( + weights: &WeightMap, + config: Qwen35MtpConfig, + ) -> Result { + let text_cfg = config.text_config().clone(); + + Self::check_weight_inventory(weights, text_cfg.mtp_num_hidden_layers)?; + + let fc = UnifiedLinear::from_weights(weights, "fc", text_cfg.group_size(), text_cfg.bits()) + .map_err(|reason| DrafterError::WeightLoad { reason })?; + let norm_w = |key: &str| -> Result, DrafterError> { + weights + .get(key) + .map(|w| ffi::copy(w)) + .ok_or_else(|| DrafterError::WeightLoad { + reason: format!("Weight not found: {key}"), + }) + }; + let pre_fc_norm_embedding = RMSNorm::new( + norm_w("pre_fc_norm_embedding.weight")?, + text_cfg.rms_norm_eps, + ); + let pre_fc_norm_hidden = + RMSNorm::new(norm_w("pre_fc_norm_hidden.weight")?, text_cfg.rms_norm_eps); + let norm = RMSNorm::new(norm_w("norm.weight")?, text_cfg.rms_norm_eps); + + let mut layers = Vec::with_capacity(text_cfg.mtp_num_hidden_layers); + for i in 0..text_cfg.mtp_num_hidden_layers { + layers.push( + Qwen35MtpDecoderLayer::from_weights(weights, &format!("layers.{i}"), &text_cfg) + .map_err(|reason| DrafterError::WeightLoad { reason })?, + ); + } + + let cache = (0..layers.len()).map(|_| KVCache::new()).collect(); + + Ok(Self { + config, + fc, + pre_fc_norm_embedding, + pre_fc_norm_hidden, + norm, + layers, + cache, + target_embed: None, + target_embed_scale: 1.0, + lm_head: None, + seed_token: None, + seed_hidden: None, + next_position: 0, + round_appended: 0, + kv_valid_len: 0, + position: 0, + }) + } + + /// Apply the upstream `sanitize` rules to a freshly loaded weight map: + /// + /// - Strip the `mtp.` prefix from raw-HF-layout keys. + /// - Apply the `+1.0` RMSNorm offset ONLY to `mtp.`-prefixed norm keys + /// (the seven suffixes upstream lists). The published + /// `mlx-community/Qwen3.8-27B-MTP-bf16` layout is already stripped — + /// i.e. the conversion already applied the shift — so already-stripped + /// keys MUST NOT be shifted again; a double shift corrupts every norm. + /// + /// The MoE `experts.gate_up_proj` split from upstream `sanitize` is + /// intentionally not ported: MoE text configs are rejected at + /// [`Qwen35MtpConfig::normalize`]. + pub fn sanitize_weights(weights: &mut WeightMap) { + const NORM_SUFFIXES: [&str; 7] = [ + ".input_layernorm.weight", + ".post_attention_layernorm.weight", + ".q_norm.weight", + ".k_norm.weight", + "norm.weight", + "pre_fc_norm_embedding.weight", + "pre_fc_norm_hidden.weight", + ]; + let keys: Vec = weights.keys().cloned().collect(); + for key in keys { + let Some(stripped) = key.strip_prefix("mtp.") else { + continue; + }; + let stripped = stripped.to_string(); + let mut value = weights.remove(&key).expect("key just enumerated"); + if NORM_SUFFIXES.iter().any(|s| stripped.ends_with(s)) { + let shape = ffi::array_shape(&value); + if shape.len() == 1 { + let one = ffi::ones(&shape, ffi::array_dtype(&value)); + value = ffi::add(&value, &one); + } + } + weights.insert(stripped, value); + } + } + + /// Strict inventory gate: every key in the (sanitized) weight map must be + /// one this drafter consumes. A leftover unknown key means the directory + /// is not a standalone `qwen3_5_mtp` head (e.g. a full target checkpoint, + /// or a layout this port does not understand) — fail closed with the + /// offending names instead of silently ignoring tensors. + fn check_weight_inventory(weights: &WeightMap, num_layers: usize) -> Result<(), DrafterError> { + let mut expected: Vec = vec![ + "fc.weight".into(), + "pre_fc_norm_embedding.weight".into(), + "pre_fc_norm_hidden.weight".into(), + "norm.weight".into(), + ]; + for i in 0..num_layers { + for name in [ + "input_layernorm.weight", + "post_attention_layernorm.weight", + "self_attn.q_proj.weight", + "self_attn.k_proj.weight", + "self_attn.v_proj.weight", + "self_attn.o_proj.weight", + "self_attn.q_norm.weight", + "self_attn.k_norm.weight", + "mlp.gate_proj.weight", + "mlp.up_proj.weight", + "mlp.down_proj.weight", + ] { + expected.push(format!("layers.{i}.{name}")); + } + } + let is_expected = |key: &str| -> bool { + expected.iter().any(|e| { + key == e + || key.strip_suffix(".scales") == Some(strip_weight(e)) + || key.strip_suffix(".biases") == Some(strip_weight(e)) + }) + }; + let mut unknown: Vec<&String> = weights.keys().filter(|k| !is_expected(k)).collect(); + if !unknown.is_empty() { + unknown.sort(); + let shown: Vec<&str> = unknown.iter().take(5).map(|s| s.as_str()).collect(); + return Err(DrafterError::WeightLoad { + reason: format!( + "qwen3_5_mtp drafter: {} unexpected tensor(s) in checkpoint (first: {}); \ + this directory does not look like a standalone qwen3_5_mtp head", + unknown.len(), + shown.join(", ") + ), + }); + } + Ok(()) + } + + /// Clear all per-run state (cache, seed, counters) while keeping the + /// bind. Leaves the drafter in the "empty cache" mode `set_shared_kv` + /// re-anchors from. + fn clear_runtime_state(&mut self) { + self.cache = (0..self.layers.len()).map(|_| KVCache::new()).collect(); + self.seed_token = None; + self.seed_hidden = None; + self.next_position = 0; + self.round_appended = 0; + } + + fn require_bound(&self) -> Result<(), DrafterError> { + if self.target_embed.is_none() || self.lm_head.is_none() { + return Err(DrafterError::BindNotCalled); + } + Ok(()) + } + + /// One MTP-stack forward over `token_ids` (length `S`) paired with the + /// matching target hidden `[1, S, H]`: + /// `norm(layers(fc(concat(pre_fc_norm_embedding(embed(tokens)), + /// pre_fc_norm_hidden(hidden)))))`. Appends `S` entries to every layer + /// cache and advances `next_position` by `S`. Returns the post-`norm` + /// hidden `[1, S, H]`. + fn forward_hidden_stack( + &mut self, + token_ids: &[i32], + hidden: &MlxArray, + ) -> Result, DrafterError> { + let target_embed = self + .target_embed + .as_ref() + .ok_or(DrafterError::BindNotCalled)?; + let s = token_ids.len() as i32; + debug_assert!(s >= 1, "forward_hidden_stack requires at least one token"); + debug_assert_eq!( + ffi::array_shape(hidden)[1], + s, + "hidden seq len must match token count" + ); + + let ids = ffi::from_slice_i32(token_ids, &[1, s]); + let mut tok_embed = target_embed.forward(&ids); + if self.target_embed_scale != 1.0 { + tok_embed = crate::ops::multiply_scalar(&tok_embed, self.target_embed_scale); + } + + let a = self.pre_fc_norm_embedding.forward(&tok_embed); + let b = self.pre_fc_norm_hidden.forward(hidden); + let fused = crate::ops::concatenate(&a, &b, -1); + let mut h = self.fc.forward(&fused); + + // Multi-token forwards use a causal mask offset by the drafter + // cache's key length (mirrors upstream + // `create_attention_mask(h, layer_cache)`); single-token draft steps + // need none. The RoPE offset is the drafter's logical position, which + // equals the cache offset while state is intact and runs ahead of it + // in the empty-cache degraded mode. + let mask = if s > 1 { + let cache_offset = self.cache.first().map(|c| c.offset).unwrap_or(0); + Some(crate::utils::create_causal_mask(s, cache_offset)) + } else { + None + }; + let rope_offset = self.next_position; + for (layer, cache) in self.layers.iter().zip(self.cache.iter_mut()) { + h = layer.forward(&h, mask.as_deref(), cache, rope_offset); + } + let h = self.norm.forward(&h); + self.next_position += s; + Ok(h) + } + + /// Project drafter hidden through the borrowed target LM head. + fn project_logits(&self, h: &MlxArray) -> Result, DrafterError> { + match self.lm_head.as_ref().ok_or(DrafterError::BindNotCalled)? { + MtpLmHead::Linear(lm) => Ok(lm.forward(h)), + MtpLmHead::TiedEmbed(embed) => Ok(embed.as_linear(h)), + } + } + + /// Single-token sample from the last position of `logits`. Temperature 0 + /// is argmax; non-greedy configs sample through the fused kernel (the + /// round loop owns the quality consequences, exactly as with the Gemma 4 + /// assistant drafter). + fn sample_one(logits: &MlxArray, sampler: &SamplingConfig) -> i32 { + let last = ffi::slice_last_logits(logits); + let tok = ffi::fused_sample( + &last, + sampler.temperature, + sampler.top_k, + sampler.top_p, + sampler.min_p, + ); + ffi::eval(&tok); + ffi::item_i32(&tok) + } + + /// Compute and stash the next-round seed from the last position of a + /// post-`norm` drafter hidden block. + fn set_seed_from_hidden( + &mut self, + h: &MlxArray, + sampler: &SamplingConfig, + ) -> Result<(), DrafterError> { + let shape = ffi::array_shape(h); + let last = shape[1] - 1; + let h_last = ffi::slice(h, &[0, last, 0], &[shape[0], last + 1, shape[2]]); + let logits = self.project_logits(&h_last)?; + let tok = Self::sample_one(&logits, sampler); + self.seed_token = Some(tok); + self.seed_hidden = Some(h_last); + Ok(()) + } + + /// Test/diagnostic accessor: `(cache_offset, next_position, + /// round_appended, has_seed)`. + pub fn state_probe(&self) -> (i32, i32, i32, bool) { + ( + self.cache.first().map(|c| c.offset).unwrap_or(0), + self.next_position, + self.round_appended, + self.seed_token.is_some(), + ) + } +} + +fn strip_weight(key: &str) -> &str { + key.strip_suffix(".weight").unwrap_or(key) +} + +impl Drafter for Qwen35MtpDraftModel { + fn bind(&mut self, target: &dyn LanguageModel) -> Result<(), DrafterError> { + let embed = target + .embed_tokens_module() + .ok_or(DrafterError::TargetMissingFeature { + feature: "embed_tokens_module", + })?; + // Upstream: `lm_head = target.lm_head or embed_tokens.as_linear`. + self.lm_head = Some(match target.lm_head_module() { + Some(lm) => MtpLmHead::Linear(lm), + None => MtpLmHead::TiedEmbed(embed.clone_shared()), + }); + // Upstream reads `inner.embed_scale` with a 1.0 default; the Qwen 3.5 + // family carries no embedding scale. + self.target_embed_scale = 1.0; + self.target_embed = Some(embed); + Ok(()) + } + + /// Reject a target whose text hidden size or vocabulary does not match + /// this MTP head. The head consumes the target's post-final-norm hidden + /// at `hidden_size` width and emits through the target's own LM head, so + /// both must line up. + fn validate_target_compat(&self, target: &dyn LanguageModel) -> Result<(), DrafterError> { + let expected_hidden = self.config.text_config().hidden_size as i32; + let sentinel = ffi::from_slice_i32(&[0_i32], &[1, 1]); + let embedded = + target + .embed_tokens(&sentinel) + .ok_or(DrafterError::TargetMissingFeature { + feature: "embed_tokens", + })?; + let target_hidden = ffi::array_shape(&embedded).last().copied().unwrap_or(0); + if target_hidden != expected_hidden { + return Err(DrafterError::BindFailed { + reason: format!( + "qwen3_5_mtp drafter is incompatible with this target: drafter \ + text_config.hidden_size = {expected_hidden} but the target's text hidden \ + size = {target_hidden}. The MTP head consumes the target's last hidden \ + state, so these must be equal; pair the drafter split from the same \ + checkpoint family (e.g. mlx-community/Qwen3.8-27B-MTP-bf16 with a \ + Qwen3.8-27B target)." + ), + }); + } + + // Vocabulary probe via the head the drafter would actually borrow. + let drafter_vocab = self.config.text_config().vocab_size as i32; + let zero_hidden = ffi::zeros(&[1, 1, target_hidden], crate::dtype::FLOAT32); + let logits = match target.lm_head_module() { + Some(lm) => lm.forward(&zero_hidden), + None => match target.embed_tokens_module() { + Some(embed) => embed.as_linear(&zero_hidden), + // Mock targets that only expose `embed_tokens` pass on the + // hidden-size gate alone, mirroring the Gemma 4 adapter. + None => return Ok(()), + }, + }; + let target_vocab = ffi::array_shape(&logits).last().copied().unwrap_or(0); + if target_vocab != drafter_vocab { + return Err(DrafterError::BindFailed { + reason: format!( + "qwen3_5_mtp drafter vocabulary is incompatible with this target: drafter \ + text_config.vocab_size = {drafter_vocab} but the target vocabulary = \ + {target_vocab}. Draft token ids must index the same vocabulary the \ + target verifies against." + ), + }); + } + Ok(()) + } + + /// Qwen MTP has no shared-K/V concept: the tensors are ignored whatever + /// their count (the target adapter passes an empty slice) and only the + /// position metadata is consumed, mirroring upstream `set_shared_kv`'s + /// `del shared_kv_states`. + /// + /// With an empty cache the drafter re-anchors `next_position` to the + /// target's post-rollback cache offset; with an intact cache the two are + /// already equal by construction. A non-empty cache whose position + /// disagrees is stale (a hook failure or an out-of-band reset) and is + /// cleared so the round continues in the empty-cache degraded mode + /// instead of drafting from corrupt state. + fn set_shared_kv( + &mut self, + _shared_kv: SharedKv<'_>, + kv_offset: usize, + position: usize, + _left_padding: usize, + ) -> Result<(), DrafterError> { + self.kv_valid_len = kv_offset as i32; + self.position = position as i32; + let cache_empty = self.cache.iter().all(|c| c.offset == 0); + if cache_empty { + self.next_position = kv_offset as i32; + } else if self.next_position != kv_offset as i32 { + tracing::debug!( + next_position = self.next_position, + kv_offset, + "qwen3_5_mtp drafter position drifted from target cache; clearing drafter \ + state and re-anchoring (draft context lost, correctness unaffected)" + ); + self.clear_runtime_state(); + self.next_position = kv_offset as i32; + } + Ok(()) + } + + /// The drafter owns its cache internally (`self.cache`), so the trait's + /// external-cache factory stays empty, like the Gemma 4 assistant. + fn make_cache(&self) -> Vec { + Vec::new() + } + + /// Clear all per-run state and re-bind against `target`. + /// + /// NOTE for the tick-slice grant rotation: unlike the Gemma 4 assistant + /// (whose reset is the trait default no-op), this reset DESTROYS the + /// drafter's accumulated history. A parked slice session whose drafter + /// was reset resumes correctly — `set_shared_kv` re-anchors into the + /// empty-cache mode — but with reduced draft context until the history + /// rebuilds from accepted tokens. + fn reset(&mut self, target: &dyn LanguageModel) -> Result<(), DrafterError> { + self.clear_runtime_state(); + self.bind(target) + } + + fn configured_block_size(&self) -> Option { + Some(self.config.block_size()) + } + + /// Upstream sets `prefer_requested_block_size = True`: the user-requested + /// block size is honored directly instead of the Gemma-style adaptive + /// configured-depth warm-up. + fn prefer_requested_block_size(&self) -> bool { + true + } + + fn prefill_from_target_hidden( + &mut self, + prompt_tokens: &[i32], + hidden: &MlxArray, + first_bonus: i32, + sampler: &SamplingConfig, + ) -> Result<(), DrafterError> { + self.require_bound()?; + let p = prompt_tokens.len(); + if p == 0 { + return Ok(()); + } + let hshape = ffi::array_shape(hidden); + if hshape.len() != 3 || hshape[1] != p as i32 { + return Err(DrafterError::DraftFailed { + reason: format!( + "qwen3_5_mtp drafter prefill: hidden shape {hshape:?} does not cover the \ + {p}-token prompt" + ), + }); + } + // Start clean: the prompt prefill defines position 0 of the drafter + // sequence, so any prior state is stale by construction. + self.clear_runtime_state(); + + // Position i pairs token_{i+1} with hidden_i: shift the prompt left + // by one and append the just-sampled first bonus. + let mut shifted: Vec = prompt_tokens[1..].to_vec(); + shifted.push(first_bonus); + let h = self.forward_hidden_stack(&shifted, hidden)?; + self.set_seed_from_hidden(&h, sampler) + } + + fn accept_verified_tokens( + &mut self, + verify_hidden: &MlxArray, + draft_tokens: &[i32], + accepted: usize, + new_tokens: &[i32], + sampler: &SamplingConfig, + ) -> Result<(), DrafterError> { + self.require_bound()?; + let hshape = ffi::array_shape(verify_hidden); + let block = draft_tokens.len() + 1; + if hshape.len() != 3 || (hshape[1] as usize) < block || accepted > draft_tokens.len() { + // Validate BEFORE mutating: a malformed accept would otherwise + // leave a half-trimmed cache. Poisoned state is cleared so the + // next `set_shared_kv` re-anchors cleanly. + self.clear_runtime_state(); + return Err(DrafterError::DraftFailed { + reason: format!( + "qwen3_5_mtp drafter accept: verify_hidden shape {hshape:?} does not cover \ + block={block} / accepted={accepted}" + ), + }); + } + + // Trim the in-round cache entries beyond the accepted prefix. Entries + // appended during `draft_block` used the drafter's own hidden as the + // pair input; accepted ones are kept as-is (upstream does the same), + // rejected ones are removed. + let keep = accepted.min(self.round_appended.max(0) as usize); + let trim = self.round_appended - keep as i32; + if trim > 0 { + for cache in &mut self.cache { + cache.trim(trim); + } + self.next_position -= trim; + } + + // Extend with the accepted tokens not yet in the cache, paired with + // the target's true verify hidden, plus the newly emitted bonus. + let h_dim = hshape[2]; + let mut tokens: Vec = Vec::new(); + let mut hidden_cat: Option> = None; + let push_slice = |pos: usize, hidden_cat: &mut Option>| { + let pos = pos as i32; + let s = ffi::slice(verify_hidden, &[0, pos, 0], &[hshape[0], pos + 1, h_dim]); + *hidden_cat = Some(match hidden_cat.take() { + None => s, + Some(prev) => crate::ops::concatenate(&prev, &s, 1), + }); + }; + for (draft_idx, &draft_tok) in draft_tokens.iter().enumerate().take(accepted).skip(keep) { + tokens.push(draft_tok); + push_slice(draft_idx, &mut hidden_cat); + } + if let Some(&last) = new_tokens.last() { + tokens.push(last); + push_slice(accepted, &mut hidden_cat); + } + + if let (false, Some(hiddens)) = (tokens.is_empty(), hidden_cat) { + let h = self.forward_hidden_stack(&tokens, &hiddens)?; + self.set_seed_from_hidden(&h, sampler)?; + } + self.round_appended = 0; + Ok(()) + } + + fn draft_block( + &mut self, + last_bonus: i32, + hidden: Option<&MlxArray>, + block_size: usize, + sampler: &SamplingConfig, + ) -> Result, DrafterError> { + self.require_bound()?; + if block_size <= 1 { + // Degenerate budget: drop any pending seed (its position + // assumption still holds, but the round loop only reaches here + // on terminal budgets). + return Ok(Vec::new()); + } + + let mut tokens: Vec = Vec::with_capacity(block_size - 1); + self.round_appended = 0; + + // Seed fast path: the prompt prefill / accept hook already computed + // this round's first proposal from TARGET hidden (and its cache entry + // is already appended), so consume it instead of re-deriving from the + // caller's (bonus, hidden) pair — the two computations are identical + // by construction. + let (mut tok, mut h_prev): (i32, UniquePtr) = + match (self.seed_token.take(), self.seed_hidden.take()) { + (Some(seed_tok), Some(seed_hidden)) => { + tokens.push(seed_tok); + (seed_tok, seed_hidden) + } + _ => { + let hidden = hidden.ok_or(DrafterError::DraftBlockMissingHidden)?; + (last_bonus, ffi::copy(hidden)) + } + }; + + while tokens.len() < block_size - 1 { + let h = self.forward_hidden_stack(&[tok], &h_prev)?; + self.round_appended += 1; + let logits = self.project_logits(&h)?; + tok = Self::sample_one(&logits, sampler); + tokens.push(tok); + h_prev = h; + } + tokens.truncate(block_size - 1); + Ok(tokens) + } + + fn sanitize(&mut self, weights: &mut WeightMap) -> Result<(), DrafterError> { + Self::sanitize_weights(weights); + Ok(()) + } + + fn kind(&self) -> DrafterKind { + DrafterKind::Mtp + } +} diff --git a/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/tests.rs b/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/tests.rs new file mode 100644 index 000000000..e2c5e2ae4 --- /dev/null +++ b/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/tests.rs @@ -0,0 +1,403 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Unit tests for the Qwen 3.5 MTP drafter: sanitize semantics, the strict +//! weight inventory, and the stateful cache/position lifecycle +//! (prompt prefill → draft → accept) that upstream drives implicitly. + +use super::config::Qwen35MtpConfig; +use super::model::Qwen35MtpDraftModel; +use crate::drafter::{Drafter, DrafterError, SharedKv}; +use crate::ffi::{self, MlxArray}; +use crate::generate::{LanguageModel, SamplingConfig}; +use crate::layers::{UnifiedEmbedding, UnifiedLinear}; +use crate::weights::WeightMap; +use cxx::UniquePtr; + +const HIDDEN: i32 = 8; +const VOCAB: i32 = 16; + +fn tiny_config() -> Qwen35MtpConfig { + let cfg: Qwen35MtpConfig = serde_json::from_str( + r#"{ + "model_type": "qwen3_5_mtp", + "block_size": 3, + "text_config": { + "model_type": "qwen3_5_text", + "hidden_size": 8, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 4, + "intermediate_size": 16, + "vocab_size": 16, + "rms_norm_eps": 1e-6, + "rope_parameters": {"rope_theta": 10000.0, "partial_rotary_factor": 0.5}, + "mtp_num_hidden_layers": 1, + "tie_word_embeddings": false + } + }"#, + ) + .expect("tiny config parses"); + cfg.normalize().expect("tiny config normalizes") +} + +fn insert_val(w: &mut WeightMap, key: &str, shape: &[i32], val: f32) { + let n: i32 = shape.iter().product(); + let data = vec![val; n as usize]; + w.insert(key.to_string(), ffi::from_slice_f32(&data, shape)); +} + +/// The published 15-tensor inventory (already `mtp.`-prefix stripped), with +/// deterministic small values so real MLX forwards run. +fn tiny_weights() -> WeightMap { + let mut w: WeightMap = WeightMap::new(); + insert_val(&mut w, "fc.weight", &[8, 16], 0.05); + insert_val(&mut w, "pre_fc_norm_embedding.weight", &[8], 1.0); + insert_val(&mut w, "pre_fc_norm_hidden.weight", &[8], 1.0); + insert_val(&mut w, "norm.weight", &[8], 1.0); + insert_val(&mut w, "layers.0.input_layernorm.weight", &[8], 1.0); + insert_val( + &mut w, + "layers.0.post_attention_layernorm.weight", + &[8], + 1.0, + ); + insert_val(&mut w, "layers.0.self_attn.q_proj.weight", &[16, 8], 0.05); + insert_val(&mut w, "layers.0.self_attn.k_proj.weight", &[4, 8], 0.05); + insert_val(&mut w, "layers.0.self_attn.v_proj.weight", &[4, 8], 0.05); + insert_val(&mut w, "layers.0.self_attn.o_proj.weight", &[8, 8], 0.05); + insert_val(&mut w, "layers.0.self_attn.q_norm.weight", &[4], 1.0); + insert_val(&mut w, "layers.0.self_attn.k_norm.weight", &[4], 1.0); + insert_val(&mut w, "layers.0.mlp.gate_proj.weight", &[16, 8], 0.05); + insert_val(&mut w, "layers.0.mlp.up_proj.weight", &[16, 8], 0.05); + insert_val(&mut w, "layers.0.mlp.down_proj.weight", &[8, 16], 0.05); + w +} + +fn build_drafter() -> Qwen35MtpDraftModel { + Qwen35MtpDraftModel::from_weights(&tiny_weights(), tiny_config()).expect("tiny drafter builds") +} + +/// Mock Qwen-style target: real embedding table + untied LM head so the +/// drafter's bind captures working shared-buffer modules. +struct MockQwenTarget { + embed: UnifiedEmbedding, + lm_head: UnifiedLinear, +} + +impl MockQwenTarget { + fn new() -> Self { + let mut w: WeightMap = WeightMap::new(); + insert_val(&mut w, "embed_tokens.weight", &[VOCAB, HIDDEN], 0.1); + insert_val(&mut w, "lm_head.weight", &[VOCAB, HIDDEN], 0.1); + let embed = + UnifiedEmbedding::from_weights(&w, "embed_tokens", 64, 4).expect("mock embed builds"); + let lm_head = UnifiedLinear::from_weights(&w, "lm_head", 64, 4).expect("mock head builds"); + Self { embed, lm_head } + } +} + +impl LanguageModel for MockQwenTarget { + fn forward( + &self, + _input_ids: &MlxArray, + _caches: &mut [crate::layers::KVCache], + _mask: Option<&MlxArray>, + ) -> UniquePtr { + unreachable!("drafter tests do not invoke the target forward") + } + + fn make_caches(&self) -> Vec { + Vec::new() + } + + fn num_layers(&self) -> usize { + 0 + } + + fn eos_token_ids(&self) -> Vec { + Vec::new() + } + + fn embed_tokens(&self, input_ids: &MlxArray) -> Option> { + Some(self.embed.forward(input_ids)) + } + + fn embed_tokens_module(&self) -> Option { + Some(self.embed.clone_shared()) + } + + fn lm_head_module(&self) -> Option { + Some(self.lm_head.clone_shared()) + } +} + +fn greedy() -> SamplingConfig { + SamplingConfig { + temperature: 0.0, + ..SamplingConfig::default() + } +} + +fn hidden_block(seq: i32) -> UniquePtr { + let n = (seq * HIDDEN) as usize; + let data: Vec = (0..n).map(|i| 0.01 * (i % 7) as f32).collect(); + ffi::from_slice_f32(&data, &[1, seq, HIDDEN]) +} + +fn first_f32(arr: &MlxArray) -> f32 { + ffi::eval(arr); + let bytes = ffi::array_to_raw_bytes(arr); + f32::from_ne_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) +} + +// ── Sanitize ────────────────────────────────────────────────────────────── + +/// Raw-HF layout (`mtp.` prefix) gets stripped and its 1-D norm weights get +/// the `+1.0` shift on exactly the seven upstream suffixes. +#[test] +fn sanitize_strips_mtp_prefix_and_shifts_norms() { + let mut w: WeightMap = WeightMap::new(); + w.insert("mtp.norm.weight".into(), ffi::from_slice_f32(&[0.5], &[1])); + w.insert( + "mtp.layers.0.self_attn.q_norm.weight".into(), + ffi::from_slice_f32(&[0.25], &[1]), + ); + // Non-norm key: stripped but NOT shifted. + w.insert( + "mtp.fc.weight".into(), + ffi::from_slice_f32(&[0.5, 0.5], &[1, 2]), + ); + Qwen35MtpDraftModel::sanitize_weights(&mut w); + + assert!(!w.keys().any(|k| k.starts_with("mtp.")), "prefix stripped"); + assert_eq!(first_f32(w.get("norm.weight").unwrap()), 1.5); + assert_eq!( + first_f32(w.get("layers.0.self_attn.q_norm.weight").unwrap()), + 1.25 + ); + assert_eq!(first_f32(w.get("fc.weight").unwrap()), 0.5); +} + +/// The published `mlx-community/Qwen3.8-27B-MTP-bf16` layout is already +/// stripped, i.e. the conversion already applied the shift: sanitize MUST +/// leave it untouched (a second `+1.0` corrupts every norm). +#[test] +fn sanitize_leaves_already_stripped_layout_unshifted() { + let mut w: WeightMap = WeightMap::new(); + w.insert("norm.weight".into(), ffi::from_slice_f32(&[1.5], &[1])); + w.insert( + "pre_fc_norm_hidden.weight".into(), + ffi::from_slice_f32(&[0.75], &[1]), + ); + Qwen35MtpDraftModel::sanitize_weights(&mut w); + assert_eq!(first_f32(w.get("norm.weight").unwrap()), 1.5); + assert_eq!(first_f32(w.get("pre_fc_norm_hidden.weight").unwrap()), 0.75); +} + +// ── Weight inventory ────────────────────────────────────────────────────── + +/// The exact published 15-tensor inventory constructs. +#[test] +fn from_weights_accepts_the_published_inventory() { + let _ = build_drafter(); +} + +/// A missing required tensor fails with a message naming it. +#[test] +fn from_weights_rejects_missing_tensor() { + let mut w = tiny_weights(); + w.remove("layers.0.self_attn.k_norm.weight"); + let err = Qwen35MtpDraftModel::from_weights(&w, tiny_config()).expect_err("must fail"); + let msg = format!("{err}"); + assert!(msg.contains("k_norm"), "got: {msg}"); +} + +/// Unknown extra tensors mean the directory is not a standalone MTP head; +/// fail closed naming them instead of silently ignoring. +#[test] +fn from_weights_rejects_unknown_extra_tensor() { + let mut w = tiny_weights(); + insert_val( + &mut w, + "model.layers.0.self_attn.q_proj.weight", + &[4, 4], + 0.0, + ); + let err = Qwen35MtpDraftModel::from_weights(&w, tiny_config()).expect_err("must fail"); + let msg = format!("{err}"); + assert!(msg.contains("unexpected tensor"), "got: {msg}"); + assert!(msg.contains("model.layers.0"), "got: {msg}"); +} + +// ── Lifecycle errors ────────────────────────────────────────────────────── + +#[test] +fn draft_block_requires_bind_then_hidden() { + let mut drafter = build_drafter(); + let err = drafter + .draft_block(1, None, 3, &greedy()) + .expect_err("unbound must fail"); + assert!(matches!(err, DrafterError::BindNotCalled)); + + let target = MockQwenTarget::new(); + drafter.bind(&target).expect("bind"); + // No seed, no hidden: the MTP path needs one of the two. + let err = drafter + .draft_block(1, None, 3, &greedy()) + .expect_err("hidden required without a seed"); + assert!(matches!(err, DrafterError::DraftBlockMissingHidden)); +} + +#[test] +fn validate_target_compat_accepts_matching_mock() { + let drafter = build_drafter(); + let target = MockQwenTarget::new(); + drafter + .validate_target_compat(&target) + .expect("mock geometry matches the tiny config"); +} + +// ── Stateful lifecycle ──────────────────────────────────────────────────── + +/// The load-bearing state machine: prompt prefill, seeded draft, full and +/// partial accepts. After every phase the drafter's cache offset and logical +/// `next_position` must track the target cache offset exactly (that +/// invariant is what `set_shared_kv` re-anchoring checks against). +#[test] +fn prefill_draft_accept_keeps_positions_in_sync() { + let mut drafter = build_drafter(); + let target = MockQwenTarget::new(); + drafter.bind(&target).expect("bind"); + let sampler = greedy(); + + // Prompt prefill: 4 tokens + first bonus -> 4 cache entries, seed ready. + let prompt = [3_i32, 5, 7, 2]; + drafter + .prefill_from_target_hidden(&prompt, &hidden_block(4), 9, &sampler) + .expect("prefill"); + assert_eq!(drafter.state_probe(), (4, 4, 0, true)); + + // Round arm: target cache offset after target prefill = 4. Cache is + // non-empty and in sync, so nothing is cleared. + drafter + .set_shared_kv(SharedKv::new(&[]), 4, 3, 0) + .expect("arm"); + assert_eq!(drafter.state_probe(), (4, 4, 0, true)); + + // Draft block (block_size 3): the seed supplies the first proposal, one + // forward produces the second. One in-round cache append. + let draft = drafter + .draft_block(9, None, 3, &sampler) + .expect("draft with seed"); + assert_eq!(draft.len(), 2); + assert_eq!(drafter.state_probe(), (5, 5, 1, false)); + + // Full accept: accepted = 2, new_tokens = both drafts + target bonus. + // keep = min(2, 1) = 1 in-round entry, no trim; appends draft[1] and the + // bonus paired with verify hidden -> cache 4 + 1 + 2 = 7. The target + // cache after a full-accept verify is 4 + 3 = 7: in sync. + let new_tokens = vec![draft[0], draft[1], 11]; + drafter + .accept_verified_tokens(&hidden_block(3), &draft, 2, &new_tokens, &sampler) + .expect("full accept"); + assert_eq!(drafter.state_probe(), (7, 7, 0, true)); + drafter + .set_shared_kv(SharedKv::new(&[]), 7, 6, 0) + .expect("re-arm stays in sync"); + assert_eq!(drafter.state_probe(), (7, 7, 0, true)); + + // Next round: draft again (seed + 1 forward -> 8 entries), then a + // zero-accept round: trim the 1 in-round entry back to 7, append only + // the target bonus -> 8. Target: 7 + (0 + 1) = 8. In sync. + let draft2 = drafter.draft_block(11, None, 3, &sampler).expect("draft 2"); + assert_eq!(draft2.len(), 2); + assert_eq!(drafter.state_probe(), (8, 8, 1, false)); + drafter + .accept_verified_tokens(&hidden_block(3), &draft2, 0, &[13], &sampler) + .expect("zero accept"); + assert_eq!(drafter.state_probe(), (8, 8, 0, true)); +} + +/// Empty-cache mode: `set_shared_kv` anchors `next_position` to the target +/// offset, and a draft block runs from the caller's (bonus, hidden) pair +/// without any seed or history. +#[test] +fn set_shared_kv_anchors_empty_cache_and_drafts_without_history() { + let mut drafter = build_drafter(); + let target = MockQwenTarget::new(); + drafter.bind(&target).expect("bind"); + drafter + .set_shared_kv(SharedKv::new(&[]), 10, 9, 0) + .expect("arm"); + assert_eq!(drafter.state_probe(), (0, 10, 0, false)); + + let draft = drafter + .draft_block(4, Some(hidden_block(1).as_ref().unwrap()), 3, &greedy()) + .expect("seedless draft"); + assert_eq!(draft.len(), 2); + // Two forwards appended (no seed): cache 2, position 12. + assert_eq!(drafter.state_probe(), (2, 12, 2, false)); +} + +/// A non-empty cache whose position disagrees with the target offset is +/// stale (out-of-band reset, hook failure): `set_shared_kv` clears it and +/// re-anchors instead of drafting from corrupt state. +#[test] +fn set_shared_kv_clears_stale_state_on_position_mismatch() { + let mut drafter = build_drafter(); + let target = MockQwenTarget::new(); + drafter.bind(&target).expect("bind"); + drafter + .prefill_from_target_hidden(&[3, 5, 7, 2], &hidden_block(4), 9, &greedy()) + .expect("prefill"); + assert_eq!(drafter.state_probe(), (4, 4, 0, true)); + + // Target says the cache sits at 9; the drafter thinks 4 -> stale. + drafter + .set_shared_kv(SharedKv::new(&[]), 9, 8, 0) + .expect("arm"); + assert_eq!(drafter.state_probe(), (0, 9, 0, false)); +} + +/// `reset` clears everything and re-binds, leaving the drafter ready for the +/// next session's empty-cache re-anchor (the slice-grant rotation path). +#[test] +fn reset_clears_state_and_rebinds() { + let mut drafter = build_drafter(); + let target = MockQwenTarget::new(); + drafter.bind(&target).expect("bind"); + drafter + .prefill_from_target_hidden(&[3, 5, 7, 2], &hidden_block(4), 9, &greedy()) + .expect("prefill"); + drafter.reset(&target).expect("reset"); + assert_eq!(drafter.state_probe(), (0, 0, 0, false)); + // Still bound: a seedless draft with an explicit hidden works. + drafter + .set_shared_kv(SharedKv::new(&[]), 4, 3, 0) + .expect("arm"); + let draft = drafter + .draft_block(4, Some(hidden_block(1).as_ref().unwrap()), 3, &greedy()) + .expect("post-reset draft"); + assert_eq!(draft.len(), 2); +} + +/// `prefer_requested_block_size` and the configured block size mirror the +/// upstream drafter flags. +#[test] +fn block_size_flags_match_upstream() { + let drafter = build_drafter(); + assert!(drafter.prefer_requested_block_size()); + assert_eq!(drafter.configured_block_size(), Some(3)); +} diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index a59ecdcd8..3e2fb95e3 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -1809,6 +1809,26 @@ mod ffi { new_state: &mut UniquePtr, ); + /// Chain-parity GatedDeltaNet forward for the speculative verify / + /// rollback replay paths (issue #1165): rounds the recurrent state + /// through the storage dtype after every in-block step so a T=K + /// block is bit-identical to K consecutive T=1 decode calls (the + /// classic greedy chain the exactness gate compares against). + /// Identical to `metal_gated_delta_forward` at T=1. + /// Used by: Qwen3.5 `forward_speculative` and rollback replay. + #[allow(clippy::too_many_arguments)] + unsafe fn metal_gated_delta_forward_chain_parity( + q: &MlxArray, + k: &MlxArray, + v: &MlxArray, + g: &MlxArray, + beta: &MlxArray, + state: &MlxArray, + mask: *const MlxArray, // nullable + output: &mut UniquePtr, + new_state: &mut UniquePtr, + ); + // Fused Mamba2 mixer forward for single-token decode. /// Combines in_proj + conv1d + SSM kernel + MambaRMSNormGated + out_proj into one C++ call. /// Replaces ~23 FFI round-trips for the hot decode path. diff --git a/src/lib/mlxcel-core/src/speculative/mtp/generator.rs b/src/lib/mlxcel-core/src/speculative/mtp/generator.rs index 58ac9d2d3..b1d8173de 100644 --- a/src/lib/mlxcel-core/src/speculative/mtp/generator.rs +++ b/src/lib/mlxcel-core/src/speculative/mtp/generator.rs @@ -12,7 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! [`MtpGenerator`] — round-loop driver for Gemma 4 MTP speculative decoding +//! [`MtpGenerator`]: round-loop driver for MTP speculative decoding. +//! Generic over any [`MtpTarget`] implementation: originally the Gemma 4 +//! assistant family, and since issue #1165 also the Qwen 3.5 MTP family. //! //! Top-level lifecycle (single-batch path): //! @@ -257,6 +259,13 @@ pub struct MtpSessionState { verify_out: MtpVerifyOutput, /// Current bonus token (the last emitted token). bonus: i32, + /// Prompt tokens pending a one-time stateful-drafter prompt prefill + /// ([`crate::drafter::Drafter::prefill_from_target_hidden`]). `Some` only + /// when the seed capture carried a full prompt hidden + /// ([`MtpVerifyOutput::verify_hidden_full`]); consumed by the first + /// round (probe or speculative). `None` for stateless-drafter targets + /// (Gemma 4 family), keeping their sessions byte-identical. + drafter_prefill_tokens: Option>, /// Tokens emitted so far, including the first bonus. The session does /// not keep the token stream itself: each step hands its new tokens to /// the caller, who owns accumulation/streaming. @@ -332,14 +341,15 @@ impl MtpStepOutput { } } -/// Round-loop driver for Gemma 4 MTP speculative decoding (B=1). +/// Round-loop driver for MTP speculative decoding (B=1): Gemma 4 assistant +/// and, since issue #1165, Qwen 3.5 MTP. /// /// Generic over `T: MtpTarget` so we get static dispatch into the target's /// sink-aware forward / rollback hooks (one v-table hop per call would /// otherwise show up as measurable noise at K=4 and decode-bound rounds). /// The drafter is held as `Box` so the call site can swap in -/// Gemma 4 assistant / DFlash / future MTP shapes without touching the -/// generator type. +/// the Gemma 4 assistant drafter, the Qwen 3.5 MTP drafter, or a future MTP +/// shape without touching the generator type. /// /// ## Construction /// @@ -646,9 +656,18 @@ impl MtpGenerator { new_logprobs.push(first_bonus_lp); } let finished = eos_tokens.contains(&first_bonus) || max_tokens == 1; + // Stash the prompt for the one-time stateful-drafter prefill hook + // only when the target advertised a full prompt hidden. Stateless + // targets (Gemma 4) leave `verify_hidden_full` as `None`, so this is + // `None` and the first round is unchanged for them. + let drafter_prefill_tokens = verify_out + .verify_hidden_full + .is_some() + .then(|| prompt_tokens.to_vec()); let state = MtpSessionState { verify_out, bonus: first_bonus, + drafter_prefill_tokens, emitted_count: 1, accept_lens: Vec::new(), probes_remaining: self.profile_probe_rounds, @@ -775,6 +794,41 @@ impl MtpGenerator { state.finished = true; return MtpStepOutput::finished_empty(); } + // One-time stateful-drafter prompt prefill (Qwen 3.5 MTP). Runs on + // the FIRST round of the session — before the probe branch, because + // probe rounds also extend the drafter history through the accept + // hook below and need the prompt prefix in place. Consuming the + // stashed prompt makes this idempotent across resumed sessions. A + // hook failure degrades draft quality only (the drafter continues + // with an empty history), never correctness — greedy parity rests on + // the target's verify pass alone — so it is logged, not fatal. + // + // `verify_hidden_full` is taken (not borrowed) here for the same + // one-shot reason: this hook is its only consumer, and at + // P = 32k prompt tokens / H = 5120 the `[1, P, H]` f16 seed hidden + // is about 335 MB. Taking it clears `state.verify_out` immediately + // instead of letting it ride, still allocated, through the rest of + // this round until `verify_finalize` overwrites `state.verify_out` + // wholesale at the end. This is unrelated to the PER-ROUND + // `verify_hidden_full` that `verify_finalize` sets below (a much + // smaller `[1, block, H]` slab consumed every round by + // `accept_verified_tokens`); that one is untouched by this take, + // because on any round after the first `drafter_prefill_tokens` is + // already `None` and this whole block does not run. + if let Some(prompt_tokens) = state.drafter_prefill_tokens.take() + && let Some(hidden_full) = state.verify_out.verify_hidden_full.take() + && let Err(e) = self.drafter.prefill_from_target_hidden( + &prompt_tokens, + &hidden_full, + state.bonus, + sampling, + ) + { + tracing::warn!( + "MTP drafter prompt prefill failed: {e}; continuing with an empty \ + drafter history" + ); + } // [#736] Classic-step probe rounds. While the adaptive policy is // profiling, the first `profile_probe_rounds` rounds skip the // drafter and verify only the bonus token: the resulting `[1, 1]` @@ -937,6 +991,30 @@ impl MtpGenerator { } } + // Stateful-drafter history extension (Qwen 3.5 MTP): after the + // rollback, hand the drafter the full verify hidden so it trims its + // rejected in-round tail, appends the accepted tokens paired with the + // target's true hidden, and precomputes the next round's seed. + // `state.verify_out` was just replaced by `verify_finalize`, so its + // `verify_hidden_full` is THIS round's block hidden. Probe rounds + // take the same path (draft_tokens is empty, accepted == 0), which + // keeps the drafter history position-synchronized with the target + // cache across probes. Stateless targets carry `None` and skip. + // Early-finish rounds returned above without this call — the drafter + // is reset before its next session, so nothing is lost. Failures + // degrade draft quality only and are logged, not fatal. + if let Some(hidden_full) = state.verify_out.verify_hidden_full.as_ref() + && let Err(e) = self.drafter.accept_verified_tokens( + hidden_full, + &draft_tokens, + walk.accepted, + &new_tokens, + sampling, + ) + { + tracing::warn!("MTP drafter accept_verified_tokens failed: {e}; drafter history reset"); + } + // 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/mod.rs b/src/lib/mlxcel-core/src/speculative/mtp/mod.rs index c1bf1c451..04015e5ae 100644 --- a/src/lib/mlxcel-core/src/speculative/mtp/mod.rs +++ b/src/lib/mlxcel-core/src/speculative/mtp/mod.rs @@ -12,8 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Multi-Token Prediction (MTP) round-loop generator for the Gemma 4 -//! assistant drafter family. +//! Multi-Token Prediction (MTP) round-loop generator. Originally built for +//! the Gemma 4 assistant drafter family; since issue #1165 also drives the +//! Qwen 3.5 MTP drafter family through the same generic [`MtpGenerator`], +//! parameterized over any target implementing [`target::MtpTarget`]. //! //! This module is a **peer** to [`crate::speculative::SpeculativeGenerator`], //! not an extension. The classic speculative path and the MTP path differ in @@ -21,11 +23,11 @@ //! //! | Aspect | Classic speculative | MTP | //! |---------------------------|-------------------------------|-------------------------------------------| -//! | Drafter KV cache | Owns its own KV cache | Has no own KV cache | +//! | Drafter KV cache | Owns its own KV cache | Gemma 4: no own KV cache (shares the target's). Qwen 3.5 MTP: owns accumulated KV history, destroyed on `reset` | //! | Drafter input | Last accepted token | Bonus token + target's last hidden | //! | Verify | Single batched forward | Single batched forward (same) | //! | Rollback | `trim_caches` | `rollback_speculative_cache` (per-row) | -//! | Cross-attention | Each model attends own cache | Drafter attends target's shared K/V | +//! | Cross-attention | Each model attends own cache | Gemma 4: drafter attends target's shared K/V. Qwen 3.5 MTP: no cross-attention, consumes the target's post-norm hidden state | //! //! Both paths satisfy the same external contract — the //! [`crate::generate::LanguageModel`] target trait + a [`crate::drafter::Drafter`] @@ -41,10 +43,12 @@ //! //! - [`walk`] — `_speculative_walk` accept logic (standalone, well-tested). //! - [`generator`] — [`MtpGenerator`] round-loop driver. -//! - [`target`] — [`MtpTarget`] trait that the Gemma 4 wrapper implements -//! to expose its speculative hooks (`forward_with_speculative_sinks`, -//! `rollback_speculative_cache`) to the round-loop driver without -//! forcing this crate to depend on the outer `mlxcel` crate. +//! - [`target`]: [`MtpTarget`] trait that a target model implements to +//! expose its speculative hooks (`forward_with_speculative_sinks`, +//! `rollback_speculative_cache`) to the round-loop driver without forcing +//! this crate to depend on the outer `mlxcel` crate. Implemented by the +//! Gemma 4 wrapper and, since issue #1165, the Qwen 3.5 (dense and MoE, +//! text and VLM) targets. pub(crate) mod adaptive; pub mod generator; 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 1e913ca1a..7233ff6f9 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 @@ -1090,6 +1090,7 @@ mod tests { next_shared_kv, kv_offset, bonus_position: kv_offset, + verify_hidden_full: None, } } } diff --git a/src/lib/mlxcel-core/src/speculative/mtp/target.rs b/src/lib/mlxcel-core/src/speculative/mtp/target.rs index 42fa002f8..cade7f142 100644 --- a/src/lib/mlxcel-core/src/speculative/mtp/target.rs +++ b/src/lib/mlxcel-core/src/speculative/mtp/target.rs @@ -156,6 +156,16 @@ pub struct MtpVerifyOutput { pub kv_offset: usize, /// Position id of the bonus token. pub bonus_position: usize, + /// Full multi-position hidden for **stateful** drafters (Qwen 3.5 MTP): + /// the target's post-final-norm hidden at every forwarded position of + /// the pass that produced this output — `[1, P, H]` over the prompt for + /// the `prefill_and_seed` seed, `[1, block, H]` over the verify block + /// afterwards. `None` for targets whose drafters are stateless (Gemma 4 + /// family) and for seeds whose prefill only covered an adopted-prefix + /// suffix; the round loop then skips the corresponding drafter-history + /// hook ([`crate::drafter::Drafter::prefill_from_target_hidden`] / + /// [`crate::drafter::Drafter::accept_verified_tokens`]). + pub verify_hidden_full: Option>, } impl std::fmt::Debug for MtpVerifyOutput { @@ -164,6 +174,7 @@ impl std::fmt::Debug for MtpVerifyOutput { .field("next_shared_kv_len", &self.next_shared_kv.len()) .field("kv_offset", &self.kv_offset) .field("bonus_position", &self.bonus_position) + .field("has_verify_hidden_full", &self.verify_hidden_full.is_some()) .finish() } } diff --git a/src/lib/mlxcel-core/src/speculative/mtp/tests.rs b/src/lib/mlxcel-core/src/speculative/mtp/tests.rs index ad6b89982..ea26e8335 100644 --- a/src/lib/mlxcel-core/src/speculative/mtp/tests.rs +++ b/src/lib/mlxcel-core/src/speculative/mtp/tests.rs @@ -100,6 +100,7 @@ impl MockMtpTarget { next_shared_kv, kv_offset, bonus_position: kv_offset, + verify_hidden_full: None, } } } diff --git a/src/models/gated_delta.rs b/src/models/gated_delta.rs index 90ada83ed..67f4b0911 100644 --- a/src/models/gated_delta.rs +++ b/src/models/gated_delta.rs @@ -223,6 +223,30 @@ pub fn gated_delta_ops( beta: &MlxArray, state: Option<&MlxArray>, mask: Option<&MlxArray>, +) -> (UniquePtr, UniquePtr) { + gated_delta_ops_with_parity(q, k, v, g, beta, state, mask, false) +} + +/// Body of [`gated_delta_ops`] with the chain-parity switch. +/// +/// `chain_parity == true` routes the Metal path through the chain-parity +/// kernel (per-step storage-dtype state rounding, issue #1165) so a `T = K` +/// block is bit-identical to `K` consecutive `T = 1` calls. The non-Metal +/// fallback ignores the flag: the parity guarantee is Metal-only today, and +/// the speculative verify paths that need it are Metal-gated in practice. +/// +/// Used by: [`gated_delta_ops`], [`gated_delta_update`], +/// [`gated_delta_update_chain_parity`] +#[allow(clippy::too_many_arguments)] +fn gated_delta_ops_with_parity( + q: &MlxArray, + k: &MlxArray, + v: &MlxArray, + g: &MlxArray, + beta: &MlxArray, + state: Option<&MlxArray>, + mask: Option<&MlxArray>, + chain_parity: bool, ) -> (UniquePtr, UniquePtr) { let shape = mlxcel_core::array_shape(q); let b = shape[0]; @@ -256,17 +280,32 @@ pub fn gated_delta_ops( // SAFETY: Metal kernel reads from valid input arrays, writes to output/new_state. // mask_ptr is null when mask is None, which the C++ side handles correctly. unsafe { - mlxcel_core::metal_gated_delta_forward( - q, - k, - v, - g, - beta, - ¤t_state, - mask_ptr, - &mut output, - &mut new_state, - ); + if chain_parity { + warn_if_chain_parity_forfeited_by_shape(g, mask); + mlxcel_core::metal_gated_delta_forward_chain_parity( + q, + k, + v, + g, + beta, + ¤t_state, + mask_ptr, + &mut output, + &mut new_state, + ); + } else { + mlxcel_core::metal_gated_delta_forward( + q, + k, + v, + g, + beta, + ¤t_state, + mask_ptr, + &mut output, + &mut new_state, + ); + } } return (output, new_state); } @@ -638,8 +677,9 @@ fn gated_delta_chunked( /// lane, so `Dk` must cover at least one full SIMD group and be exactly /// divisible by 32. Its GQA mapping also requires an integral `Hv / Hk`. /// -/// Used by: Qwen3Next, Qwen3.5, KimiLinear -fn supports_metal_gated_delta_kernel(hk: i32, hv: i32, dk: i32, dv: i32) -> bool { +/// Used by: Qwen3Next, Qwen3.5, KimiLinear, `speculative_burst::mtp_capable_target` +/// (the Qwen 3.5 MTP capability gate, issue #1165 hardening) +pub(crate) fn supports_metal_gated_delta_kernel(hk: i32, hv: i32, dk: i32, dv: i32) -> bool { hk > 0 && hv > 0 && dv > 0 && dk >= 32 && dk % 32 == 0 && hv >= hk && hv % hk == 0 } @@ -670,6 +710,105 @@ pub fn gated_delta_update( gated_delta_ops(q, k, v, &g, &beta, state, mask) } +/// Chain-parity variant of [`gated_delta_update`] for the speculative verify +/// and rollback-replay paths (issue #1165). +/// +/// The standard Metal kernel carries float32 recurrent state across a `T = K` +/// block and quantizes it to the storage dtype only at the final store, while +/// the classic decode chain quantizes after every token — so a verify block +/// is NOT bit-identical to the single-token chain it must match at +/// temperature 0, and last-ulp state differences flip near-tie argmaxes. +/// This variant rounds the state through the storage dtype at the end of +/// every in-block step, making `T = K` bit-equal to `K` consecutive `T = 1` +/// calls. For `T = 1` it is byte-identical to [`gated_delta_update`]. +/// +/// Used by: `Qwen35GatedDeltaNet::forward_hidden_internal_with_capture` +/// (the speculative verify pass), `Qwen35Model::rollback_speculative_cache` +/// (the GDN replay). +#[allow(clippy::too_many_arguments)] +pub fn gated_delta_update_chain_parity( + q: &MlxArray, + k: &MlxArray, + v: &MlxArray, + a: &MlxArray, + b: &MlxArray, + a_log: &MlxArray, + dt_bias: &MlxArray, + state: Option<&MlxArray>, + mask: Option<&MlxArray>, +) -> (UniquePtr, UniquePtr) { + let beta = mlxcel_core::sigmoid(b); + let g = compute_g(a_log, a, dt_bias); + 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: +/// `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 +/// effect. Default ON — turning it off forfeits the temperature-0 +/// byte-identity of speculative output against classic decode. +fn chain_parity_enabled() -> bool { + static FLAG: std::sync::OnceLock = std::sync::OnceLock::new(); + *FLAG.get_or_init(|| { + std::env::var("MLXCEL_GDN_CHAIN_PARITY") + .map(|v| v != "0") + .unwrap_or(true) + }) +} + +/// Warn once when a chain-parity request cannot be honored by shape. +/// +/// The C++ dispatch (`metal_gated_delta_forward_impl` in +/// `src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp`) only defines the +/// chain-parity kernel variant for the scalar-gate, no-mask shape; a +/// vectorized gate (`g` is 4-dim, `[B, T, Hv, Dk]`) or an explicit mask +/// silently falls through to the standard kernel regardless of the +/// `chain_parity` argument, forfeiting the T=K-block-equals-K-single-steps +/// byte-identity contract with no signal at the call site. Neither shape is +/// reachable from the Qwen 3.5 MTP verify/rollback paths today (`mask = +/// None`, `g` is always 3-dim there), so this has never fired in +/// production; the warning exists so a future caller that adds a mask or a +/// vectorized gate to a chain-parity-requesting path fails loudly instead +/// of silently losing exactness (issue #1165 hardening). +/// +/// `eprintln!` rather than `tracing::warn!`: only `mlxcel-server` installs a +/// tracing subscriber unconditionally, the offline `mlxcel` CLI only does so +/// when `RUST_LOG` is set (see `init_cli_tracing` in `src/main.rs`), and this +/// path is reachable from both binaries. +/// +/// Mirrors the C++ shape dispatch in `metal_gated_delta_forward_impl` +/// exactly: `vectorized || has_mask` takes the standard kernel branches +/// (`gated_delta_step_vec[_mask]` / `gated_delta_step_mask`) even when the +/// caller passed `chain_parity = true`; only `!vectorized && !has_mask` +/// reaches the `chain_parity` check and dispatches +/// `gated_delta_step_seqpar`. Pulled out as a pure function (rather than +/// inlined into the warn helper) so the shape predicate is unit-testable +/// without a Metal device or a live `MlxArray`. +fn chain_parity_forfeited_by_shape(vectorized_gate: bool, has_mask: bool) -> bool { + vectorized_gate || has_mask +} + +fn warn_if_chain_parity_forfeited_by_shape(g: &MlxArray, mask: Option<&MlxArray>) { + let vectorized_gate = mlxcel_core::array_ndim(g) == 4; + if !chain_parity_forfeited_by_shape(vectorized_gate, mask.is_some()) { + return; + } + static WARNED: std::sync::Once = std::sync::Once::new(); + WARNED.call_once(|| { + eprintln!( + "mlxcel: chain-parity gated-delta kernel was requested but the input shape \ + (vectorized_gate={vectorized_gate}, masked={}) is outside the parity kernel's \ + contract; the C++ dispatch silently falls back to the standard kernel, which \ + does NOT guarantee T=K-block byte-identity with K single-token decode steps. \ + The MTP temperature-0 exactness contract is forfeited for this call. See \ + `metal_gated_delta_forward_impl` in src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp.", + mask.is_some() + ); + }); +} + /// Fast RMS normalization without a learned scale, followed by scalar scaling. /// /// Mirrors mlx-lm's `mx.fast.rms_norm(x, None, eps) * scale` path for linear diff --git a/src/models/gated_delta_tests.rs b/src/models/gated_delta_tests.rs index f10aeffca..a476d1361 100644 --- a/src/models/gated_delta_tests.rs +++ b/src/models/gated_delta_tests.rs @@ -13,8 +13,9 @@ // limitations under the License. use super::{ - GatedDeltaCache, RMSNormGated, compute_g, gated_delta_chunked, gated_delta_ops, - gated_delta_step, precise_swiglu_gate, restore_dtype, supports_metal_gated_delta_kernel, + GatedDeltaCache, RMSNormGated, chain_parity_forfeited_by_shape, compute_g, gated_delta_chunked, + gated_delta_ops, gated_delta_step, precise_swiglu_gate, restore_dtype, + supports_metal_gated_delta_kernel, }; use mlxcel_core::{dtype, generate::ModelStateSnapshot}; @@ -203,6 +204,27 @@ fn metal_gated_delta_shape_gate_rejects_unsupported_contracts() { assert!(!supports_metal_gated_delta_kernel(3, 4, 64, 2)); } +// Issue #1165 hardening: the C++ dispatch (`metal_gated_delta_forward_impl`) +// silently ignores `chain_parity=true` whenever the gate is vectorized or a +// mask is present, taking the standard kernel instead of +// `gated_delta_step_seqpar`. `chain_parity_forfeited_by_shape` is the pure +// predicate behind the one-time warning that makes that forfeiture loud; +// pin it directly against every combination the C++ `if`/`else if` chain in +// `metal_gated_delta_forward_impl` distinguishes. +#[test] +fn chain_parity_forfeited_by_shape_matches_the_cpp_dispatch_branches() { + // Scalar gate, no mask: the ONLY shape that reaches + // `gated_delta_step_seqpar`. Parity is honored, nothing to warn about. + assert!(!chain_parity_forfeited_by_shape(false, false)); + // Vectorized gate (with or without a mask) takes `gated_delta_step_vec` + // / `gated_delta_step_vec_mask` before the `chain_parity` check is ever + // reached. + assert!(chain_parity_forfeited_by_shape(true, false)); + assert!(chain_parity_forfeited_by_shape(true, true)); + // Scalar gate with a mask takes `gated_delta_step_mask`, same forfeiture. + assert!(chain_parity_forfeited_by_shape(false, true)); +} + #[test] #[ignore = "requires serial MLX execution"] fn gated_delta_ops_single_token_output_shape() { diff --git a/src/models/gemma4_mtp_target.rs b/src/models/gemma4_mtp_target.rs index 86c5523a6..6696ff8cd 100644 --- a/src/models/gemma4_mtp_target.rs +++ b/src/models/gemma4_mtp_target.rs @@ -415,7 +415,7 @@ impl<'a> Gemma4MtpTargetAdapter<'a> { /// /// Returns `None` when `logprobs_config.enabled` is false (the /// zero-overhead path — no slicing, no log-softmax). - fn per_position_logprobs( + pub(crate) fn per_position_logprobs( logits: &MlxArray, target_tokens: &[i32], logprobs_config: &mlxcel_core::sampling::LogprobsConfig, @@ -558,6 +558,10 @@ impl<'a> MtpTarget for Gemma4MtpTargetAdapter<'a> { next_shared_kv, kv_offset, bonus_position, + // The Gemma 4 assistant drafter is stateless per round; no + // full-block hidden is advertised, so the round loop skips the + // stateful-drafter history hooks. + verify_hidden_full: None, }; (first_bonus, seed, first_bonus_lp) } @@ -763,6 +767,7 @@ impl<'a> MtpTarget for Gemma4MtpTargetAdapter<'a> { next_shared_kv, kv_offset, bonus_position, + verify_hidden_full: None, } } diff --git a/src/models/mod.rs b/src/models/mod.rs index 51b4b295a..f750946af 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -139,6 +139,7 @@ pub mod qwen2_moe; pub mod qwen2_vl; pub mod qwen3; pub mod qwen3_5; +pub mod qwen3_5_mtp_target; pub mod qwen3_moe; pub mod qwen3_next; pub mod qwen3_vl; diff --git a/src/models/qwen3_5.rs b/src/models/qwen3_5.rs index 892714913..e4ee3bd04 100644 --- a/src/models/qwen3_5.rs +++ b/src/models/qwen3_5.rs @@ -29,7 +29,8 @@ use crate::distributed::pipeline::LayerFilter; use crate::distributed::pipeline::StageExecutionOutput; use crate::distributed::pipeline::partial_loading::filter_weight_map; use crate::models::gated_delta::{ - GatedDeltaCache, RMSNormGated, gated_delta_update, scaled_fast_rms_norm_no_weight, + GatedDeltaCache, RMSNormGated, gated_delta_update, gated_delta_update_chain_parity, + scaled_fast_rms_norm_no_weight, }; use crate::models::model_owned::ModelOwnedSequenceState; use crate::models::qwen_mrope_state::MRopeState; @@ -454,6 +455,24 @@ impl Qwen35Config { } } +/// Whether a Qwen 3.5-family config's GDN geometry satisfies the Metal +/// chain-parity kernel's shape contract +/// ([`crate::models::gated_delta::supports_metal_gated_delta_kernel`]). +/// +/// Free function (rather than a method taking `&self`) so it can be +/// exercised against a bare [`Qwen35Config`] in unit tests without +/// constructing a full [`Qwen35Model`] with real weights. +/// +/// Used by: [`Qwen35Model::supports_chain_parity_kernel`] +fn qwen35_config_supports_chain_parity_kernel(config: &Qwen35Config) -> bool { + crate::models::gated_delta::supports_metal_gated_delta_kernel( + config.linear_num_key_heads as i32, + config.linear_num_value_heads as i32, + config.linear_key_head_dim as i32, + config.linear_value_head_dim as i32, + ) +} + // GatedDeltaNet - Qwen3.5 variant with separate projections. /// GatedDeltaNet for Qwen3.5 with separate in_proj_qkv, in_proj_z, in_proj_b, in_proj_a #[allow(dead_code)] @@ -864,8 +883,16 @@ impl Qwen35GatedDeltaNet { conv_input: mlxcel_core::copy(&conv_input), }); - // Run gated_delta_update for the full verify block. - let (out, new_state) = gated_delta_update( + // Run the gated-delta update for the full verify block through the + // CHAIN-PARITY kernel (issue #1165): the standard kernel carries + // float32 state across the block and quantizes once at the end, + // which is not bit-identical to the per-token quantization of the + // classic decode chain the verify argmax must match at temperature + // 0. The parity kernel rounds the state through the storage dtype + // after every in-block step, so a T=K verify block equals K + // single-token decode steps bit-for-bit (and equals the standard + // kernel exactly at T=1). + let (out, new_state) = gated_delta_update_chain_parity( &q, &k, &v, @@ -1424,6 +1451,25 @@ impl Qwen35Model { self.make_internal_caches() } + /// Whether this checkpoint's GDN geometry satisfies the Metal chain-parity + /// kernel's shape contract (`dk >= 32 && dk % 32 == 0 && hv % hk == 0`, + /// [`crate::models::gated_delta::supports_metal_gated_delta_kernel`]). + /// + /// MTP speculative decoding's temperature-0 byte-identity guarantee rests + /// entirely on `gated_delta_update_chain_parity` taking the Metal kernel + /// path (issue #1165). `metal_is_available()` alone does not imply that: + /// a shape outside the kernel's contract silently falls back to the + /// ops-based path, which carries float32 state across a verify block + /// instead of rounding per step, forfeiting byte-identity with no signal. + /// Callers that gate MTP eligibility on Metal availability must also + /// check this so an incompatible checkpoint declines to classic decode + /// instead of silently violating the exactness contract. + /// + /// Used by: `speculative_burst::mtp_capable_target` + pub(crate) fn supports_chain_parity_kernel(&self) -> bool { + qwen35_config_supports_chain_parity_kernel(&self.config) + } + fn visible_len(cache: &Qwen3NextCache) -> usize { match cache { Qwen3NextCache::Attention(kv) => kv.seq_len().max(0) as usize, @@ -2040,7 +2086,11 @@ impl Qwen35Model { ], ); - let (_y, replayed_state) = gated_delta_update( + // Chain-parity replay (issue #1165): the replayed prefix state + // must equal the state a classic decode chain would hold after + // the accepted tokens, so the replay uses the same per-step + // rounding kernel as the verify pass that captured the block. + let (_y, replayed_state) = gated_delta_update_chain_parity( &q_n, &k_n, &v_n, @@ -2097,6 +2147,161 @@ impl Qwen35Model { max_a } + + // === Sequence-routed speculative hooks (MTP target adapter) ============ + // + // The MtpTarget trait is `&self`-only with no cache parameter: the target + // resolves its own per-sequence cache. These wrappers route the existing + // caller-owned-slice speculative hooks above through + // `ModelOwnedSequenceState`, exactly like `forward_with_sequence_caches` + // does for the classic path, so the MTP adapter + // (`crate::models::qwen3_5_mtp_target`) stays a stateless view over the + // model — the shape the tick-cooperative slice path requires (it + // reconstructs the adapter every scheduler tick). + + /// Prefill `input_ids` through the standard batched-attention forward + /// while also returning the last decoder layer's **pre-final-norm** + /// hidden for every position. + /// + /// This is deliberately NOT [`Self::forward_speculative`]: that path runs + /// the per-position target-verify attention (`attend_per_position`), + /// which is byte-aligned with single-token *decode* but is O(seq_len) + /// attention dispatches and does not match the classic *prefill* + /// numerics. The classic path's first token is sampled from the batched + /// causal prefill logits, so MTP's first bonus must come from the same + /// computation to stay byte-identical — upstream's MTP prefill likewise + /// runs the plain forward (`capture_layer_ids=None`, so + /// `target_verify=False`). GDN rollback snapshots are not captured: a + /// prefill needs no rollback, and holding every linear layer's + /// full-prompt q/k/v tensors alive simultaneously would balloon memory. + /// + /// Keep the loop in lockstep with [`Self::forward_internal`] (mask + /// construction, layer dispatch, final norm, LM head). + /// + /// Used by: `Qwen35MtpTargetAdapter::prefill_and_seed`. + /// Public test seam: out-of-crate diagnostics (the chain-parity / + /// verify-cost tests in `tests/`) time this batched-attention path + /// against `forward_speculative`'s per-position verify path. Not for + /// production callers outside this module — use the sequence-routed + /// wrappers. + #[doc(hidden)] + pub fn forward_prefill_with_last_hidden( + &self, + input_ids: &MlxArray, + caches: &mut [Qwen3NextCache], + ) -> (UniquePtr, UniquePtr) { + let mut h = self.embed_tokens.forward(input_ids); + + let shape = mlxcel_core::array_shape(&h); + let seq_len = shape[1]; + + let fa_idx = self.config.full_attention_interval - 1; + let fa_mask = if seq_len > 1 { + let offset = if fa_idx < caches.len() { + caches[fa_idx].offset() + } else { + 0 + }; + Some(create_causal_mask(seq_len, offset)) + } else { + None + }; + + for (layer, cache) in self.layers.iter().zip(caches.iter_mut()) { + let mask = if layer.is_linear { + None + } else { + fa_mask.as_deref() + }; + h = layer.forward(&h, mask, cache, None); + } + + let normed = self.norm.forward(&h); + let logits = if let Some(ref lm_head) = self.lm_head { + lm_head.forward(&normed) + } else { + self.embed_tokens.as_linear(&normed) + }; + (logits, h) + } + + /// Apply the final RMSNorm to a captured pre-norm hidden block. + /// + /// The MTP drafter consumes the target's **post-final-norm** hidden: + /// upstream's `return_hidden` appends the model output (`self.model(...)` + /// ends in `self.norm(h)`), and the Qwen 3.5 family has no + /// `speculative_draft_hidden` hook to do this in the round loop. The + /// hidden captured by [`Self::forward_speculative`] / + /// [`Self::forward_prefill_with_last_hidden`] is pre-norm, so the MTP + /// adapter normalizes it here before handing it to the drafter. + /// + /// Used by: `Qwen35MtpTargetAdapter`. + pub(crate) fn apply_final_norm(&self, hidden: &MlxArray) -> UniquePtr { + self.norm.forward(hidden) + } + + /// Sequence-routed [`Self::forward_prefill_with_last_hidden`]. + /// + /// Used by: `Qwen35MtpTargetAdapter::prefill_and_seed`. + pub(crate) fn forward_prefill_with_last_hidden_for_sequence( + &self, + input_ids: &MlxArray, + seq_id: Option, + ) -> (UniquePtr, UniquePtr) { + self.sequence_state.with_or_create_sequence_state( + seq_id, + || self.make_internal_caches(), + |caches| self.forward_prefill_with_last_hidden(input_ids, caches), + ) + } + + /// Sequence-routed [`Self::forward_speculative`]. + /// + /// Used by: `Qwen35MtpTargetAdapter::verify_forward`. + pub(crate) fn forward_speculative_for_sequence( + &self, + input_ids: &MlxArray, + seq_id: Option, + capture_layer_ids: &[usize], + ) -> VerifyOutput { + self.sequence_state.with_or_create_sequence_state( + seq_id, + || self.make_internal_caches(), + |caches| self.forward_speculative(input_ids, caches, capture_layer_ids), + ) + } + + /// Sequence-routed [`Self::rollback_speculative_cache`]. + /// + /// Used by: `Qwen35MtpTargetAdapter::verify_finalize`. + pub(crate) fn rollback_speculative_cache_for_sequence( + &self, + seq_id: Option, + gdn_states: &[GdnRollbackSnapshot], + accepted: &[i32], + block_size: i32, + ) -> i32 { + self.sequence_state.with_sequence_state(seq_id, |caches| { + self.rollback_speculative_cache(caches, gdn_states, accepted, block_size) + }) + } + + /// Absolute offset of the sequence's first attention-layer KV cache + /// (the value upstream reads as `prompt_cache[0].offset` when rebinding + /// the drafter after a rollback). + /// + /// Used by: `Qwen35MtpTargetAdapter::{prefill_and_seed, verify_finalize}`. + pub(crate) fn speculative_cache_offset_for_sequence(&self, seq_id: Option) -> i32 { + self.sequence_state.with_sequence_state(seq_id, |caches| { + caches + .iter() + .find_map(|c| match c { + Qwen3NextCache::Attention(kv) => Some(kv.offset), + _ => None, + }) + .unwrap_or(0) + }) + } } /// DFlash speculative-decoding target adapter for the Qwen 3.5 text @@ -3449,3 +3654,57 @@ mod sanitize_tests { } } } + +/// Regression coverage for the `mtp_capable_target` hardening (issue #1165 +/// review): `metal_is_available()` is necessary but not sufficient for MTP +/// byte-identity, the checkpoint's GDN geometry must also satisfy the +/// chain-parity kernel's shape contract. This module pins the config-level +/// predicate directly, without needing a Metal device or a full model. +#[cfg(test)] +mod chain_parity_gate_tests { + use super::{Qwen35Config, qwen35_config_supports_chain_parity_kernel}; + + /// `linear_key_head_dim` matches the published `mlx-community` Qwen 3.5 / + /// 3.8 checkpoints (Dk=128); the shape the parity kernel was built for. + fn published_shape_config() -> Qwen35Config { + serde_json::from_value(serde_json::json!({ + "model_type": "qwen3_5", + "hidden_size": 16, + "num_hidden_layers": 1, + "intermediate_size": 32, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "linear_num_value_heads": 48, + "linear_num_key_heads": 16, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "full_attention_interval": 4, + "vocab_size": 100 + })) + .expect("minimal qwen3.5 test config") + } + + #[test] + fn accepts_the_published_checkpoint_geometry() { + assert!(qwen35_config_supports_chain_parity_kernel( + &published_shape_config() + )); + } + + #[test] + fn declines_a_key_head_dim_not_a_multiple_of_32() { + // Mirrors the `linear_key_head_dim: 4` shape already exercised by + // `sanitize_output_loads_through_switch_glu_for_all_naming_conventions`. + // Not hypothetical: a shape this small tree already tests with. + let mut config = published_shape_config(); + config.linear_key_head_dim = 4; + assert!(!qwen35_config_supports_chain_parity_kernel(&config)); + } + + #[test] + fn declines_a_non_integral_value_to_key_head_ratio() { + let mut config = published_shape_config(); + config.linear_num_value_heads = 15; // 15 % 16 != 0 + assert!(!qwen35_config_supports_chain_parity_kernel(&config)); + } +} diff --git a/src/models/qwen3_5_mtp_target.rs b/src/models/qwen3_5_mtp_target.rs new file mode 100644 index 000000000..25b0709a0 --- /dev/null +++ b/src/models/qwen3_5_mtp_target.rs @@ -0,0 +1,521 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! MTP target adapter for the Qwen 3.5 family. +//! +//! Glue layer that wires [`crate::models::qwen3_5::Qwen35Model`] (and its VLM +//! wrapper [`crate::vision::Qwen35VLModel`]) to the +//! [`mlxcel_core::speculative::mtp::target::MtpTarget`] trait, the same role +//! [`crate::models::gemma4_mtp_target`] plays for Gemma 4. +//! +//! ## Cache ownership +//! +//! `MtpTarget`'s methods are `&self` with no cache parameter, while the +//! Qwen 3.5 speculative hooks (`forward_speculative`, +//! `rollback_speculative_cache`) take a caller-owned `&mut [Qwen3NextCache]`. +//! The adapter bridges the two through the model's **sequence-routed** +//! wrappers (`*_for_sequence`, added alongside this adapter), which resolve +//! `seq_id` into `ModelOwnedSequenceState` internally. That keeps the adapter +//! a stateless, cheaply reconstructible view over the model — the shape the +//! tick-cooperative slice path depends on (it rebuilds the adapter every +//! scheduler tick) and the reason a `RefCell>`-inside-the- +//! adapter design was rejected. +//! +//! ## Hidden-state semantics +//! +//! The Qwen MTP drafter consumes the target's **post-final-norm** hidden: +//! upstream's `return_hidden` appends the model output (which ends in +//! `self.norm(h)`), and Qwen has no `speculative_draft_hidden` hook. The +//! captured per-layer hidden here is pre-norm, so every hidden handed out +//! (seed, per-round `next_hidden`, and the full-block +//! `verify_hidden_full` for the stateful drafter's history hooks) goes +//! through [`Qwen35Model::apply_final_norm`] first. +//! +//! ## Shared K/V +//! +//! Qwen 3.5 has no shared-K/V concept (the drafter owns its own per-layer KV +//! cache), so `next_shared_kv` is always empty and the drafter's +//! `set_shared_kv` consumes only the position metadata, mirroring upstream. +//! +//! ## Scope +//! +//! B = 1 only. The batched `MtpTarget` methods keep their erroring trait +//! defaults, so a B > 1 dispatch declines with the existing named error and +//! the burst driver falls back to classic decode. + +use mlxcel_core::cache::SequenceId; +use mlxcel_core::generate::SamplingConfig; +use mlxcel_core::sampling::{LogprobsConfig, TokenLogprobData}; +use mlxcel_core::speculative::mtp::target::{ + MtpTarget, MtpVerifyOutput, VerifyCaptured, VerifyForwardOutput, +}; +use mlxcel_core::{MlxArray, UniquePtr}; + +use crate::models::gemma4_mtp_target::Gemma4MtpTargetAdapter; +use crate::models::qwen3_5::{GdnRollbackSnapshot, Qwen35Model}; + +/// Flatten the post-norm verify hidden plus the GDN rollback snapshots into +/// the opaque [`VerifyCaptured`] slots. +/// +/// Encoding (pinned by `captured_round_trip_preserves_gdn_snapshots`): +/// +/// - `scalars[0]` = number of snapshots `N`; then per snapshot `i`: +/// `scalars[1 + 2i]` = `layer_idx`, `scalars[2 + 2i]` = `has_init_state` +/// (0/1). +/// - `tensors[0]` = post-norm hidden `[1, block, H]`; then per snapshot, in +/// snapshot order: `q, k, v, a, b, conv_input`, plus `init_state` when the +/// flag is set. A mis-ordered flatten produces subtly wrong recurrent +/// state rather than a crash, so any change here must update the +/// round-trip test in lockstep. +fn flatten_captured( + hidden_post_norm: UniquePtr, + gdn_states: Vec, +) -> VerifyCaptured { + let mut scalars: Vec = Vec::with_capacity(1 + 2 * gdn_states.len()); + scalars.push(gdn_states.len() as i32); + let mut tensors: Vec> = Vec::with_capacity(1 + 7 * gdn_states.len()); + tensors.push(hidden_post_norm); + for snap in gdn_states { + scalars.push(snap.layer_idx as i32); + scalars.push(snap.init_state.is_some() as i32); + tensors.push(snap.q); + tensors.push(snap.k); + tensors.push(snap.v); + tensors.push(snap.a); + tensors.push(snap.b); + tensors.push(snap.conv_input); + if let Some(init_state) = snap.init_state { + tensors.push(init_state); + } + } + VerifyCaptured { tensors, scalars } +} + +/// Inverse of [`flatten_captured`]. Panics on a malformed encoding — the +/// captured state never crosses a serialization boundary, so a mismatch is a +/// programming error in this module, not an input condition. +fn unflatten_captured(captured: VerifyCaptured) -> (UniquePtr, Vec) { + let VerifyCaptured { tensors, scalars } = captured; + let mut tensors = tensors.into_iter(); + let hidden = tensors + .next() + .expect("VerifyCaptured must carry the hidden tensor at index 0"); + let n = *scalars + .first() + .expect("VerifyCaptured scalars must carry the snapshot count") as usize; + assert_eq!( + scalars.len(), + 1 + 2 * n, + "VerifyCaptured scalars must carry (layer_idx, has_init) per snapshot" + ); + let mut gdn_states = Vec::with_capacity(n); + for i in 0..n { + let layer_idx = scalars[1 + 2 * i] as usize; + let has_init = scalars[2 + 2 * i] != 0; + let q = tensors.next().expect("snapshot q"); + let k = tensors.next().expect("snapshot k"); + let v = tensors.next().expect("snapshot v"); + let a = tensors.next().expect("snapshot a"); + let b = tensors.next().expect("snapshot b"); + let conv_input = tensors.next().expect("snapshot conv_input"); + let init_state = has_init.then(|| tensors.next().expect("snapshot init_state")); + gdn_states.push(GdnRollbackSnapshot { + layer_idx, + q, + k, + v, + a, + b, + init_state, + conv_input, + }); + } + assert!( + tensors.next().is_none(), + "VerifyCaptured tensor count disagrees with the scalar encoding" + ); + (hidden, gdn_states) +} + +/// MTP target adapter binding a [`Qwen35Model`] to a per-sequence cache slot. +/// +/// Constructed by the server dispatch (or the offline CLI with +/// `seq_id = None`, which selects the model's internal fallback slot) and +/// consumed by [`mlxcel_core::speculative::mtp::MtpGenerator`] over the round +/// loop. The adapter does NOT own the model and holds no tensor state. +pub struct Qwen35MtpTargetAdapter<'a> { + model: &'a Qwen35Model, + seq_id: Option, + /// Adopted prompt-cache prefix length (issue #518). When `> 0`, + /// `prefill_and_seed` forwards only `prompt_tokens[prefill_start_offset..]` + /// over the KV already restored into the sequence slot. The seed then + /// advertises no full prompt hidden (`verify_hidden_full = None`), so the + /// stateful drafter skips its prompt prefill and builds history from + /// accepted tokens only — reduced draft context, identical correctness. + prefill_start_offset: usize, +} + +impl<'a> Qwen35MtpTargetAdapter<'a> { + /// Construct an adapter routing every call through the per-sequence + /// cache slot at `seq_id` (`None` = the model's internal fallback slot, + /// the offline CLI shape). + pub fn new(model: &'a Qwen35Model, seq_id: Option) -> Self { + Self { + model, + seq_id, + prefill_start_offset: 0, + } + } + + /// Set the adopted prompt-cache prefix length (issue #518). See the + /// field doc for the drafter-history consequence. + #[must_use] + pub fn with_prefill_start_offset(mut self, prefill_start_offset: usize) -> Self { + self.prefill_start_offset = prefill_start_offset; + self + } + + /// Capture-layer list for the verify pass: the last decoder layer only. + fn capture_last_layer(&self) -> [usize; 1] { + [self.model.num_layers().saturating_sub(1)] + } +} + +impl<'a> MtpTarget for Qwen35MtpTargetAdapter<'a> { + fn prefill_and_seed( + &self, + prompt_tokens: &[i32], + sampler: &SamplingConfig, + token_history: &[i32], + logprobs_config: &LogprobsConfig, + ) -> (i32, MtpVerifyOutput, Option) { + // Suffix-only prefill for an adopted prompt-cache prefix (issue + // #518): the sequence slot already holds KV for `[..offset]` + // (restored by the scheduler's APC snapshot adoption), so forward + // only the suffix; RoPE continues from the cache's restored offset. + let offset = self.prefill_start_offset.min(prompt_tokens.len()); + let forward_tokens = &prompt_tokens[offset..]; + let prompt_arr = + mlxcel_core::from_slice_i32(forward_tokens, &[1, forward_tokens.len() as i32]); + + // Batched-causal prefill, byte-identical to the classic path (NOT + // the per-position verify attention — see + // `forward_prefill_with_last_hidden`'s doc for why that matters for + // first-token parity and prefill cost). + let (logits, hidden_pre) = self + .model + .forward_prefill_with_last_hidden_for_sequence(&prompt_arr, self.seq_id); + + // First bonus from the last-position logits, with the + // history-dependent-penalty context and the sampler's token bias — + // the same computation the classic decode path's first token uses. + let logits_shape = mlxcel_core::array_shape(&logits); + let last_pos = logits_shape[1] - 1; + let vocab = logits_shape[2]; + let last_logits = mlxcel_core::slice( + &logits, + &[0, last_pos, 0], + &[logits_shape[0], last_pos + 1, vocab], + ); + let (token_arr, adjusted_logits) = + mlxcel_core::sampling::sample_token_optimized(&last_logits, sampler, token_history); + mlxcel_core::eval(&token_arr); + let first_bonus = mlxcel_core::item_i32(&token_arr); + let first_bonus_lp = + mlxcel_core::sampling::compute_logprobs(&adjusted_logits, first_bonus, logprobs_config); + + // Post-final-norm hidden for the drafter: the seed `next_hidden` is + // the last forwarded position; the full block feeds the stateful + // drafter's prompt prefill, but ONLY when it covers the whole prompt + // (an adopted-prefix suffix does not). + let hidden_post = self.model.apply_final_norm(&hidden_pre); + let next_hidden = Gemma4MtpTargetAdapter::last_position_hidden(&hidden_post); + let verify_hidden_full = (offset == 0).then_some(hidden_post); + + // Absolute post-prefill cache offset (= prompt length when cold, or + // restored offset + suffix length under adoption). + let kv_offset = self + .model + .speculative_cache_offset_for_sequence(self.seq_id) + .max(0) as usize; + let bonus_position = kv_offset.saturating_sub(1); + + let seed = MtpVerifyOutput { + next_hidden, + // Qwen has no shared-K/V concept; the drafter ignores tensors. + next_shared_kv: Vec::new(), + kv_offset, + bonus_position, + verify_hidden_full, + }; + (first_bonus, seed, first_bonus_lp) + } + + fn embed_token(&self, token_id: i32) -> UniquePtr { + let input_ids = mlxcel_core::from_slice_i32(&[token_id], &[1, 1]); + ::embed_tokens(self.model, &input_ids) + .expect("Qwen35Model exposes its embed_tokens table") + } + + fn verify_forward( + &self, + verify_input: &[i32], + sampler: &SamplingConfig, + logprobs_config: &LogprobsConfig, + ) -> VerifyForwardOutput { + let verify_arr = mlxcel_core::from_slice_i32(verify_input, &[1, verify_input.len() as i32]); + let capture = self.capture_last_layer(); + let out = self + .model + .forward_speculative_for_sequence(&verify_arr, self.seq_id, &capture); + + // issue #350: apply the model's output-suppression bias BEFORE the + // per-position argmax and the logprob extraction, exactly as the + // Gemma 4 adapter does. An empty map short-circuits to the raw + // logits, preserving the bit-exact baseline. + let logits = if sampler.token_bias.is_empty() { + out.logits + } else { + mlxcel_core::sampling::apply_token_bias(&out.logits, &sampler.token_bias) + }; + + // Greedy-parity gate: per-position argmax. At temperature 0 this is + // byte-identical to the drafter-less target's own argmax extension + // (the verify pass runs the per-position target-verify attention). + let target_tokens = Gemma4MtpTargetAdapter::argmax_per_position(&logits); + let target_logprobs = + Gemma4MtpTargetAdapter::per_position_logprobs(&logits, &target_tokens, logprobs_config); + + // Post-norm full-block hidden: consumed by finalize for + // `next_hidden` / the drafter's accept hook. + let mut hidden_states = out.hidden_states; + let hidden_pre = hidden_states + .pop() + .expect("forward_speculative returns the requested last-layer capture"); + let hidden_post = self.model.apply_final_norm(&hidden_pre); + + VerifyForwardOutput { + target_tokens, + target_logprobs, + captured: flatten_captured(hidden_post, out.gdn_states), + } + } + + fn verify_finalize( + &self, + accepted: usize, + block_size: usize, + captured: VerifyCaptured, + ) -> MtpVerifyOutput { + let (hidden_post, gdn_states) = unflatten_captured(captured); + + // Hidden at the accepted position seeds the drafter's next round. + let next_hidden = Gemma4MtpTargetAdapter::hidden_at_position(&hidden_post, accepted); + + // Roll the KV + GDN caches back to the accepted prefix. Upstream + // skips the call entirely on full accept (`accepted < bs - 1` + // guard): the trim would be zero and the GDN replay would recompute + // the state it already has. + if accepted + 1 < block_size { + let _ = self.model.rollback_speculative_cache_for_sequence( + self.seq_id, + &gdn_states, + &[accepted as i32], + block_size as i32, + ); + } + + // Absolute post-rollback cache offset, read from the model-owned + // cache (mirrors upstream's `prompt_cache[0].offset` rebind). + let kv_offset = self + .model + .speculative_cache_offset_for_sequence(self.seq_id) + .max(0) as usize; + + MtpVerifyOutput { + next_hidden, + next_shared_kv: Vec::new(), + kv_offset, + bonus_position: kv_offset.saturating_sub(1), + // Full-block post-norm hidden for the stateful drafter's accept + // hook (trim + extend + next-seed precompute). + verify_hidden_full: Some(hidden_post), + } + } + + fn num_layers(&self) -> usize { + self.model.num_layers() + } + + fn eos_token_ids(&self) -> Vec { + ::eos_token_ids(self.model) + } +} + +/// MTP target adapter for the Qwen 3.5 VLM wrapper. +/// +/// Pure delegation to the inner text model, mirroring the DFlash VL adapter +/// (`src/vision/qwen3_5_vl.rs`) and the Gemma 4 VL MTP shell: vision +/// features are fully prefilled before the MTP round loop begins, so the +/// round loop only touches the text backbone. Multimodal requests never +/// reach here — the burst gate rejects them — and the speculative forward +/// follows the DFlash convention of standard (non-MRoPE) positions for +/// text-only requests on a VL target. +pub struct Qwen35VLMtpTargetAdapter<'a> { + inner: Qwen35MtpTargetAdapter<'a>, +} + +impl<'a> Qwen35VLMtpTargetAdapter<'a> { + /// Construct an adapter routing every call through the inner text + /// model's per-sequence cache slot at `seq_id`. + pub fn new(vlm: &'a crate::vision::Qwen35VLModel, seq_id: Option) -> Self { + Self { + inner: Qwen35MtpTargetAdapter::new(&vlm.text_model, seq_id), + } + } + + /// Set the adopted prompt-cache prefix length (issue #518), delegating + /// to the inner text-model adapter. + #[must_use] + pub fn with_prefill_start_offset(mut self, prefill_start_offset: usize) -> Self { + self.inner = self.inner.with_prefill_start_offset(prefill_start_offset); + self + } +} + +impl<'a> MtpTarget for Qwen35VLMtpTargetAdapter<'a> { + fn prefill_and_seed( + &self, + prompt_tokens: &[i32], + sampler: &SamplingConfig, + token_history: &[i32], + logprobs_config: &LogprobsConfig, + ) -> (i32, MtpVerifyOutput, Option) { + self.inner + .prefill_and_seed(prompt_tokens, sampler, token_history, logprobs_config) + } + + fn embed_token(&self, token_id: i32) -> UniquePtr { + self.inner.embed_token(token_id) + } + + fn verify_forward( + &self, + verify_input: &[i32], + sampler: &SamplingConfig, + logprobs_config: &LogprobsConfig, + ) -> VerifyForwardOutput { + self.inner + .verify_forward(verify_input, sampler, logprobs_config) + } + + fn verify_finalize( + &self, + accepted: usize, + block_size: usize, + captured: VerifyCaptured, + ) -> MtpVerifyOutput { + self.inner.verify_finalize(accepted, block_size, captured) + } + + fn num_layers(&self) -> usize { + self.inner.num_layers() + } + + fn eos_token_ids(&self) -> Vec { + self.inner.eos_token_ids() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn tensor(vals: &[f32], shape: &[i32]) -> UniquePtr { + mlxcel_core::from_slice_f32(vals, shape) + } + + fn first_val(arr: &MlxArray) -> f32 { + mlxcel_core::eval(arr); + let bytes = mlxcel_core::array_to_raw_bytes(arr); + f32::from_ne_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) + } + + /// The flatten/unflatten pair must preserve snapshot order, per-snapshot + /// tensor identity, `layer_idx`, and the optional `init_state` — a + /// mis-ordered flatten produces subtly wrong recurrent state after + /// rollback rather than a crash, which is exactly the failure mode this + /// pin exists to catch. + #[test] + fn captured_round_trip_preserves_gdn_snapshots() { + let hidden = tensor(&[42.0, 43.0], &[1, 2, 1]); + // Two snapshots with distinguishable leading values; the second has + // an init_state, the first does not. + let snap_a = GdnRollbackSnapshot { + layer_idx: 3, + q: tensor(&[1.0], &[1, 1, 1, 1]), + k: tensor(&[2.0], &[1, 1, 1, 1]), + v: tensor(&[3.0], &[1, 1, 1, 1]), + a: tensor(&[4.0], &[1, 1, 1]), + b: tensor(&[5.0], &[1, 1, 1]), + init_state: None, + conv_input: tensor(&[6.0], &[1, 1, 1]), + }; + let snap_b = GdnRollbackSnapshot { + layer_idx: 7, + q: tensor(&[10.0], &[1, 1, 1, 1]), + k: tensor(&[20.0], &[1, 1, 1, 1]), + v: tensor(&[30.0], &[1, 1, 1, 1]), + a: tensor(&[40.0], &[1, 1, 1]), + b: tensor(&[50.0], &[1, 1, 1]), + init_state: Some(tensor(&[70.0], &[1, 1, 1, 1])), + conv_input: tensor(&[60.0], &[1, 1, 1]), + }; + + let captured = flatten_captured(hidden, vec![snap_a, snap_b]); + assert_eq!(captured.scalars, vec![2, 3, 0, 7, 1]); + // 1 hidden + 6 (snapshot without init) + 7 (snapshot with init). + assert_eq!(captured.tensors.len(), 14); + + let (hidden, gdn) = unflatten_captured(captured); + assert_eq!(first_val(hidden.as_ref().unwrap()), 42.0); + assert_eq!(gdn.len(), 2); + + assert_eq!(gdn[0].layer_idx, 3); + assert!(gdn[0].init_state.is_none()); + assert_eq!(first_val(gdn[0].q.as_ref().unwrap()), 1.0); + assert_eq!(first_val(gdn[0].b.as_ref().unwrap()), 5.0); + assert_eq!(first_val(gdn[0].conv_input.as_ref().unwrap()), 6.0); + + assert_eq!(gdn[1].layer_idx, 7); + assert_eq!(first_val(gdn[1].q.as_ref().unwrap()), 10.0); + assert_eq!(first_val(gdn[1].conv_input.as_ref().unwrap()), 60.0); + assert_eq!( + first_val(gdn[1].init_state.as_ref().unwrap().as_ref().unwrap()), + 70.0 + ); + } + + /// Zero snapshots (a hypothetical all-attention layout) round-trips to + /// an empty snapshot list without touching the hidden slot. + #[test] + fn captured_round_trip_handles_zero_snapshots() { + let captured = flatten_captured(tensor(&[9.0], &[1, 1, 1]), Vec::new()); + assert_eq!(captured.scalars, vec![0]); + assert_eq!(captured.tensors.len(), 1); + let (hidden, gdn) = unflatten_captured(captured); + assert_eq!(first_val(hidden.as_ref().unwrap()), 9.0); + assert!(gdn.is_empty()); + } +} diff --git a/src/server/batch/scheduler.rs b/src/server/batch/scheduler.rs index c65e96793..752867002 100644 --- a/src/server/batch/scheduler.rs +++ b/src/server/batch/scheduler.rs @@ -1607,9 +1607,11 @@ impl BatchScheduler { /// MTP) or /// [`mlxcel_core::drafter::dflash::SpeculativeTarget`] (for /// DFlash). Today that means: - /// - **MTP**: `Gemma4Wrapper` / `Gemma4VLModel` — see the - /// `MtpTarget` impls in - /// [`crate::models::gemma4_mtp_target`]. + /// - **MTP**: `Gemma4Wrapper` / `Gemma4VLModel` (Gemma 4 assistant + /// drafter, see the `MtpTarget` impls in + /// [`crate::models::gemma4_mtp_target`]) and, since issue #1165, + /// `Qwen35Model` / `Qwen35VLModel` (`qwen3_5_mtp` drafter, Metal + /// only, see [`crate::models::qwen3_5_mtp_target`]). /// - **DFlash**: `Qwen35Model` / `Qwen35VLModel` — see the /// `SpeculativeTarget` impl in `crate::models::qwen3_5`. /// 3. The drafter weights are loadable at the recorded @@ -4256,10 +4258,7 @@ impl BatchScheduler { }; block_size >= 2 && super::speculative_slice::mtp_tick_slice_enabled() - && matches!( - self.model, - LoadedModel::Gemma4(_) | LoadedModel::Gemma4VLM(_) | LoadedModel::Gemma4Unified(_) - ) + && super::speculative_burst::mtp_capable_target(&self.model) && !seq.prompt_tokens.is_empty() && super::speculative_burst::mtp_prefill_suffix_start( seq.prefill_start_offset, @@ -4539,7 +4538,8 @@ impl BatchScheduler { let seq = window.into_iter().next().expect("window has the head"); // Tick-cooperative MTP slice (issue #734): a B=1 MTP request on - // the Gemma 4 family is served one speculative round per + // any `mtp_capable_target` family (Gemma 4 assistant, and since + // issue #1165 Qwen 3.5 MTP) is served one speculative round per // scheduler tick instead of a run-to-completion burst, so // concurrent classic rows advance between rounds. The DFlash // arm keeps the legacy burst (its generator has no resumable @@ -4689,6 +4689,12 @@ impl BatchScheduler { LoadedModel::Gemma4(wrapper) => compat_and_bind(drafter, wrapper), LoadedModel::Gemma4VLM(vlm) => compat_and_bind(drafter, vlm), LoadedModel::Gemma4Unified(unified) => compat_and_bind(drafter, unified), + LoadedModel::Qwen35(qwen) | LoadedModel::Qwen35Moe(qwen) => { + compat_and_bind(drafter, qwen) + } + LoadedModel::Qwen35VLM(vlm) | LoadedModel::Qwen35MoeVLM(vlm) => { + compat_and_bind(drafter, vlm) + } // Unreachable per the callers' variant gates; produce a clean // per-request error rather than panicking. _ => Err( @@ -4732,13 +4738,11 @@ impl BatchScheduler { // Variant gate BEFORE any drafter IO, same rationale and message // as `run_mtp_burst`: an unsupported pairing declines to classic // without surfacing a confusing drafter-load error. - if !matches!( - self.model, - LoadedModel::Gemma4(_) | LoadedModel::Gemma4VLM(_) | LoadedModel::Gemma4Unified(_) - ) { + if !super::speculative_burst::mtp_capable_target(&self.model) { tracing::warn!( "MTP speculative dispatch declined: target is not \ - Gemma 4 (text, VLM, or Unified); falling back to classic decode", + Gemma 4 (text, VLM, or Unified) or Qwen 3.5 (text or VLM); \ + falling back to classic decode", ); return Some(seq); } @@ -4877,6 +4881,46 @@ impl BatchScheduler { &token_history, )) } + // Qwen 3.5 family (#1165): same slice-0 shape; the adapter is a + // stateless per-tick view over the model's sequence slot, which + // is exactly why the cache-ownership design routes through + // `*_for_sequence` wrappers instead of an adapter-owned cache. + LoadedModel::Qwen35(qwen) | LoadedModel::Qwen35Moe(qwen) => { + let adapter = crate::models::qwen3_5_mtp_target::Qwen35MtpTargetAdapter::new( + qwen, + Some(seq.seq_id), + ) + .with_prefill_start_offset(prefill_start_offset); + Ok(super::speculative_slice::begin_slice_session( + adapter, + drafter, + seq, + &self.tokenizer, + model_eos, + block_size, + probe_rounds, + prefill_start_offset, + &token_history, + )) + } + LoadedModel::Qwen35VLM(vlm) | LoadedModel::Qwen35MoeVLM(vlm) => { + let adapter = crate::models::qwen3_5_mtp_target::Qwen35VLMtpTargetAdapter::new( + vlm, + Some(seq.seq_id), + ) + .with_prefill_start_offset(prefill_start_offset); + Ok(super::speculative_slice::begin_slice_session( + adapter, + drafter, + seq, + &self.tokenizer, + model_eos, + block_size, + probe_rounds, + prefill_start_offset, + &token_history, + )) + } // Defensive arm rather than `unreachable!()` so a future // LoadedModel variant admitted by the gate above surfaces as a // clean per-request error instead of a worker panic. @@ -4992,11 +5036,30 @@ impl BatchScheduler { super::speculative_slice::step_slice_session(adapter, &mut job, &self.tokenizer); true } + LoadedModel::Qwen35(qwen) | LoadedModel::Qwen35Moe(qwen) => { + let adapter = crate::models::qwen3_5_mtp_target::Qwen35MtpTargetAdapter::new( + qwen, + Some(job.seq.seq_id), + ) + .with_prefill_start_offset(job.prefill_start_offset); + super::speculative_slice::step_slice_session(adapter, &mut job, &self.tokenizer); + true + } + LoadedModel::Qwen35VLM(vlm) | LoadedModel::Qwen35MoeVLM(vlm) => { + let adapter = crate::models::qwen3_5_mtp_target::Qwen35VLMtpTargetAdapter::new( + vlm, + Some(job.seq.seq_id), + ) + .with_prefill_start_offset(job.prefill_start_offset); + super::speculative_slice::step_slice_session(adapter, &mut job, &self.tokenizer); + true + } _ => false, }; if !stepped { // Defensive: the model cannot change mid-flight (slice 0 only - // starts on the Gemma 4 family). Fail the request cleanly; the + // starts on an `mtp_capable_target` family: Gemma 4 assistant or + // Qwen 3.5 MTP). Fail the request cleanly; the // drafter is dropped with the job so the next speculative // request lazily reloads it. self.fail_inflight_slice_job( @@ -5035,12 +5098,21 @@ impl BatchScheduler { /// `Drafter::reset` is the trait default no-op for the MTP assistant /// drafter; see `MtpSliceJob::attach_drafter` for the correctness /// argument) and push the job onto the grant backlog ring. + /// + /// The Qwen 3.5 MTP drafter's reset is NOT a no-op: it clears the + /// drafter-owned KV history. That is still correct at a park boundary — + /// the shared worker handle may serve other grantees before this job is + /// promoted, so per-session drafter state cannot survive rotation, and + /// the resumed session's `set_shared_kv` re-anchors into the documented + /// empty-cache mode (reduced draft context, identical output). fn park_speculative_slice(&mut self, mut job: Box) { if let Some(drafter) = job.take_drafter() { let target_lm: Option<&dyn LanguageModel> = match &self.model { LoadedModel::Gemma4(wrapper) => Some(wrapper), LoadedModel::Gemma4VLM(vlm) => Some(vlm), LoadedModel::Gemma4Unified(unified) => Some(unified), + LoadedModel::Qwen35(qwen) | LoadedModel::Qwen35Moe(qwen) => Some(qwen), + LoadedModel::Qwen35VLM(vlm) | LoadedModel::Qwen35MoeVLM(vlm) => Some(vlm), _ => None, }; match target_lm { @@ -5127,9 +5199,13 @@ impl BatchScheduler { // Re-acquire the worker drafter returned at park time. // `ensure_loaded` is a no-op while the handle sits in // the slot; it reloads from disk only if the park-time - // reset failed and dropped the handle (cannot happen - // for the MTP arm, whose reset is the trait default - // no-op, but handled for symmetry with the burst + // reset failed and dropped the handle. For the Gemma 4 + // assistant drafter (trait-default no-op reset) this + // cannot happen; for the Qwen 3.5 MTP drafter, whose + // reset re-binds against the target and is therefore + // fallible in principle, this branch is the defensive + // path (not expected to fire against an already-bound + // live target, but handled for symmetry with the burst // paths). A cancelled parked job is promoted normally: // its next `step_session` observes the cancel at the // round top and finishes with the tokens emitted so @@ -5258,14 +5334,19 @@ impl BatchScheduler { // holds the slot, the BETWEEN-slice holds skip `return_drafter` // (there is nothing to hand off); the END-of-session return here // resets, exactly like the legacy burst, and the park boundary - // (issue #746) routes through the same `return_drafter` plumbing - // (safe for the MTP arm because its reset is the trait default - // no-op; see `MtpSliceJob::attach_drafter`). + // (issue #746) routes through the same `return_drafter` plumbing. + // Correctness does not depend on what `reset` does to drafter-side + // state either way (the Gemma 4 assistant drafter's reset is the + // trait default no-op; the Qwen 3.5 MTP drafter's reset destroys + // its accumulated KV history instead): see + // `MtpSliceJob::attach_drafter` for the full argument. if let Some(drafter) = job.take_drafter() { let target_lm: Option<&dyn LanguageModel> = match &self.model { LoadedModel::Gemma4(wrapper) => Some(wrapper), LoadedModel::Gemma4VLM(vlm) => Some(vlm), LoadedModel::Gemma4Unified(unified) => Some(unified), + LoadedModel::Qwen35(qwen) | LoadedModel::Qwen35Moe(qwen) => Some(qwen), + LoadedModel::Qwen35VLM(vlm) | LoadedModel::Qwen35MoeVLM(vlm) => Some(vlm), _ => None, }; match target_lm { diff --git a/src/server/batch/speculative_burst.rs b/src/server/batch/speculative_burst.rs index ecd102199..e82b991c2 100644 --- a/src/server/batch/speculative_burst.rs +++ b/src/server/batch/speculative_burst.rs @@ -65,6 +65,15 @@ //! [`crate::models::gemma4_mtp_target::Gemma4MtpTargetAdapter`] / //! [`crate::models::gemma4_mtp_target::Gemma4VLMtpTargetAdapter`] / //! [`crate::models::gemma4_mtp_target::Gemma4UnifiedMtpTargetAdapter`]. +//! - **MTP / Qwen 3.5** (issue #1165): [`crate::LoadedModel::Qwen35`], +//! [`crate::LoadedModel::Qwen35Moe`], and their Qwen 3.5 VLM-wrapped +//! variants, gated additionally on the Metal backend (see +//! [`mtp_capable_target`]). Drives +//! [`mlxcel_core::speculative::mtp::MtpGenerator`] through +//! [`crate::models::qwen3_5_mtp_target::Qwen35MtpTargetAdapter`] / +//! [`crate::models::qwen3_5_mtp_target::Qwen35VLMtpTargetAdapter`], paired +//! with the `qwen3_5_mtp` drafter (a DIFFERENT drafter family from the +//! DFlash bullet below despite targeting the same model family). //! - **DFlash / Qwen 3.5** — [`crate::LoadedModel::Qwen35`], //! [`crate::LoadedModel::Qwen35Moe`], and their Qwen 3.5 VLM-wrapped //! variants for text-only requests. True multimodal requests still @@ -624,6 +633,50 @@ fn ragged_target_sliding_window(model: &LoadedModel) -> Option { } } +/// Whether `model` implements the B = 1 [`MtpTarget`] surface, i.e. can be +/// driven by `run_mtp_burst` and the tick-cooperative slice path. +/// +/// This is the single source of truth for the pure-boolean MTP capability +/// gates (`can_wait_for_slice_grant`, `start_mtp_slice_b1`'s variant gate, +/// the offline CLI gate). The arm-bearing matches that construct the +/// per-family adapters (`run_mtp_burst`, the slice start/step dispatch, the +/// offline adapter selection) necessarily repeat the variant list — when +/// adding a new MTP-capable family, extend this helper AND every +/// adapter-constructing match, or the family silently declines to classic on +/// the paths you missed. +/// +/// Batched (B > 1) MTP capability is narrower: only the Gemma 4 family +/// implements the batched adapter today (`run_mtp_burst_batched`'s own gate); +/// Qwen 3.5 MTP is B = 1 only (#1165). +/// +/// The Qwen 3.5 family additionally requires the Metal backend: its +/// temperature-0 exactness contract rests on the chain-parity gated-delta +/// kernel (`gated_delta_update_chain_parity`), which only exists as a Metal +/// kernel today. On CUDA/CPU builds the multi-token verify block would not +/// be bit-identical to the classic decode chain, so the family declines to +/// classic decode rather than silently violating the contract. +/// +/// Metal availability alone is not sufficient, either: the chain-parity +/// kernel additionally requires the checkpoint's GDN geometry to satisfy +/// `supports_metal_gated_delta_kernel` (`dk >= 32 && dk % 32 == 0 && +/// hv % hk == 0`). A Metal host whose `linear_key_head_dim` falls outside +/// that contract would otherwise silently take the ops-based fallback and +/// forfeit byte-identity with no signal (issue #1165 hardening). Every +/// published `mlx-community` Qwen 3.5 checkpoint satisfies the shape today +/// (`dk = 128`), but the gate checks it rather than assuming it. +pub(crate) fn mtp_capable_target(model: &LoadedModel) -> bool { + match model { + LoadedModel::Gemma4(_) | LoadedModel::Gemma4VLM(_) | LoadedModel::Gemma4Unified(_) => true, + LoadedModel::Qwen35(m) | LoadedModel::Qwen35Moe(m) => { + mlxcel_core::metal_is_available() && m.supports_chain_parity_kernel() + } + LoadedModel::Qwen35VLM(vlm) | LoadedModel::Qwen35MoeVLM(vlm) => { + mlxcel_core::metal_is_available() && vlm.text_model.supports_chain_parity_kernel() + } + _ => false, + } +} + /// Successful burst outcome returned to the scheduler. /// /// The scheduler uses `tokens_generated` to update the per-request @@ -942,7 +995,8 @@ impl<'a> BurstContext<'a> { } } -/// MTP B=1 burst — Gemma 4 / Gemma 4 VLM target. +/// MTP B=1 burst: Gemma 4 (assistant drafter) or Qwen 3.5 (`qwen3_5_mtp` +/// drafter, issue #1165) target. /// /// **Variant gate runs before drafter load.** Surfacing /// "unsupported target" is cheap (no IO) and the operator should see @@ -951,9 +1005,9 @@ impl<'a> BurstContext<'a> { /// wrong (e.g. `--model qwen3.5 --draft-kind mtp`). This ordering also /// matches the DFlash burst's intent below. /// -/// **Drafter bind happens before MtpGenerator construction.** The -/// underlying [`Gemma4AssistantDraftModel`] (the only currently -/// supported MTP drafter shape) requires +/// **Drafter bind happens before MtpGenerator construction.** Every +/// MTP-style drafter shape ([`Gemma4AssistantDraftModel`] and, since issue +/// #1165, `Qwen35MtpDraftModel`) requires /// [`mlxcel_core::drafter::Drafter::bind`] to be called before its /// first [`mlxcel_core::drafter::Drafter::draft_block`] call — bind /// captures the target's `embed_tokens` and resolves the drafter's @@ -992,10 +1046,32 @@ fn run_mtp_burst( LoadedModel::Gemma4(wrapper) => wrapper as &dyn LanguageModel, LoadedModel::Gemma4VLM(vlm) => vlm as &dyn LanguageModel, LoadedModel::Gemma4Unified(unified) => unified as &dyn LanguageModel, + // Qwen 3.5 family targets (text + VLM wrappers, dense + MoE) pair + // with the qwen3_5_mtp drafter (#1165). The drafter binds against + // the text backbone's embed_tokens / lm_head, which the VLM wrapper + // exposes through its LanguageModel impl exactly like the text + // model. Metal-only: the exactness contract needs the chain-parity + // gated-delta kernel (see `mtp_capable_target`). + LoadedModel::Qwen35(_) + | LoadedModel::Qwen35Moe(_) + | LoadedModel::Qwen35VLM(_) + | LoadedModel::Qwen35MoeVLM(_) + if !mlxcel_core::metal_is_available() => + { + tracing::warn!( + "MTP speculative dispatch declined: Qwen 3.5 MTP requires the Metal \ + backend (its temperature-0 exactness rests on the Metal chain-parity \ + gated-delta kernel); falling back to classic decode" + ); + return Err(BurstOutcome::DeclineToClassic); + } + LoadedModel::Qwen35(qwen) | LoadedModel::Qwen35Moe(qwen) => qwen as &dyn LanguageModel, + LoadedModel::Qwen35VLM(vlm) | LoadedModel::Qwen35MoeVLM(vlm) => vlm as &dyn LanguageModel, _ => { tracing::warn!( "MTP speculative dispatch declined: target is {:?}, expected \ - Gemma 4 (text, VLM, or Unified); falling back to classic decode", + Gemma 4 (text, VLM, or Unified) or Qwen 3.5 (text or VLM); \ + falling back to classic decode", model_variant_label(ctx.model), ); return Err(BurstOutcome::DeclineToClassic); @@ -1161,9 +1237,51 @@ fn run_mtp_burst( profile_probe_rounds, ) } + // Qwen 3.5 family (#1165): the adapter routes through the model's + // per-sequence cache slot; there is no rotating-cache buffer to arm + // (Qwen attention caches are plain growing KVCaches), so no + // block-size-derived constructor parameter. + LoadedModel::Qwen35(qwen) | LoadedModel::Qwen35Moe(qwen) => { + let adapter = crate::models::qwen3_5_mtp_target::Qwen35MtpTargetAdapter::new( + qwen, + Some(seq.seq_id), + ) + .with_prefill_start_offset(prefill_start_offset); + drive_mtp_generator( + adapter, + owned_drafter, + &prompt, + max_tokens, + &sampling, + &token_history, + block_size, + cancel, + &logprobs_config, + profile_probe_rounds, + ) + } + LoadedModel::Qwen35VLM(vlm) | LoadedModel::Qwen35MoeVLM(vlm) => { + let adapter = crate::models::qwen3_5_mtp_target::Qwen35VLMtpTargetAdapter::new( + vlm, + Some(seq.seq_id), + ) + .with_prefill_start_offset(prefill_start_offset); + drive_mtp_generator( + adapter, + owned_drafter, + &prompt, + max_tokens, + &sampling, + &token_history, + block_size, + cancel, + &logprobs_config, + profile_probe_rounds, + ) + } // Unreachable: the hoisted variant check above already - // returned for any model that is neither `Gemma4` nor - // `Gemma4VLM`. Keeping a defensive arm rather than + // returned for any model outside the Gemma 4 / Qwen 3.5 MTP + // families. Keeping a defensive arm rather than // `unreachable!()` so a future LoadedModel variant added to // the enum surfaces as a clean burst error instead of a // panic at request time. @@ -2285,10 +2403,15 @@ fn run_mtp_burst_batched( LoadedModel::Gemma4(wrapper) => wrapper as &dyn LanguageModel, LoadedModel::Gemma4VLM(vlm) => vlm as &dyn LanguageModel, LoadedModel::Gemma4Unified(unified) => unified as &dyn LanguageModel, + // Qwen 3.5 MTP is B = 1 only (#1165): its batched MtpTarget methods + // keep the erroring trait defaults, so the batched arm declines here + // rather than failing mid-burst. Rows fall back to classic decode; + // singleton Qwen requests still take the B = 1 MTP arm. _ => { tracing::warn!( "MTP batched speculative dispatch declined: target is {:?}, expected \ - Gemma 4 (text, VLM, or Unified); falling back to classic decode", + Gemma 4 (text, VLM, or Unified; Qwen 3.5 MTP is B = 1 only); \ + falling back to classic decode", model_variant_label(ctx.model), ); return Err(BurstOutcome::DeclineToClassic); diff --git a/src/server/batch/speculative_burst_tests.rs b/src/server/batch/speculative_burst_tests.rs index 86abeca77..5cd0e31b4 100644 --- a/src/server/batch/speculative_burst_tests.rs +++ b/src/server/batch/speculative_burst_tests.rs @@ -535,6 +535,7 @@ impl MockMtpTarget { next_shared_kv, kv_offset, bonus_position: kv_offset, + verify_hidden_full: None, } } } @@ -1779,3 +1780,200 @@ fn sampling_config_eq_excludes_history_dependent_penalty_configs() { ); } } + +// =========================================================================== +// MTP dispatch-site coverage (issue #1165) +// =========================================================================== +// +// Every MTP dispatch site is a `match` over `LoadedModel` with a `_ =>` +// decline/defensive arm, so omitting a family from one of them does NOT fail +// to compile: the family silently declines to classic decode on the path that +// was missed. `mtp_capable_target` is the single source of truth for the +// pure-boolean gates, but the arm-bearing matches that construct the +// per-family adapters necessarily repeat the variant list, and nothing in the +// type system keeps the two in agreement. +// +// This test closes that gap mechanically: it reads the variant list out of +// `mtp_capable_target` and asserts every one of those variants is named in +// every adapter-constructing / drafter-routing site. Adding a family to the +// helper without wiring it into a site now fails at `cargo test` time instead +// of producing a silent decline in production. + +/// Source of the three files that carry an MTP dispatch site. +const SRC_SPECULATIVE_BURST: &str = include_str!("speculative_burst.rs"); +const SRC_SCHEDULER: &str = include_str!("scheduler.rs"); +const SRC_OFFLINE_GENERATE: &str = include_str!("../../commands/generate.rs"); + +/// Every function that must handle every MTP-capable `LoadedModel` variant, +/// as `(human label, source, `fn` anchor)`. +/// +/// A site belongs here when omitting a variant from it makes that family +/// silently take a non-MTP path. Sites that are deliberately Gemma-only +/// (the B > 1 batched burst, `ragged_target_sliding_window`) are excluded: +/// Qwen 3.5 MTP is B = 1 only (#1165). +const MTP_DISPATCH_SITES: &[(&str, &str, &str)] = &[ + ( + "run_mtp_burst (variant gate + adapter selection)", + SRC_SPECULATIVE_BURST, + "fn run_mtp_burst(", + ), + ( + "BatchScheduler::bind_drafter_to_target (compat_and_bind)", + SRC_SCHEDULER, + "fn bind_drafter_to_target(", + ), + ( + "BatchScheduler::start_mtp_slice_b1 (slice 0 adapter)", + SRC_SCHEDULER, + "fn start_mtp_slice_b1(", + ), + ( + "BatchScheduler::execute_speculative_slice_round (per-tick adapter)", + SRC_SCHEDULER, + "fn execute_speculative_slice_round(", + ), + ( + "BatchScheduler::park_speculative_slice (drafter return)", + SRC_SCHEDULER, + "fn park_speculative_slice(", + ), + ( + "BatchScheduler::finalize_speculative_slice (drafter return)", + SRC_SCHEDULER, + "fn finalize_speculative_slice(", + ), + ( + "run_offline_mtp (offline CLI gate + adapter selection)", + SRC_OFFLINE_GENERATE, + "fn run_offline_mtp(", + ), +]; + +/// Extract the brace-delimited body that follows `anchor` in `src`. +/// +/// Brace matching skips `//` line comments, `/* */` block comments, and +/// double-quoted string literals so a `format!("{x}")` or a brace in prose +/// cannot unbalance the scan. +fn mtp_fn_body<'a>(src: &'a str, anchor: &str) -> &'a str { + let start = src + .find(anchor) + .unwrap_or_else(|| panic!("MTP dispatch-site anchor {anchor:?} not found; if the function was renamed, update MTP_DISPATCH_SITES")); + let bytes = src.as_bytes(); + let open = start + + src[start..] + .find('{') + .unwrap_or_else(|| panic!("no opening brace after {anchor:?}")); + let mut depth = 0usize; + let mut i = open; + while i < bytes.len() { + match bytes[i] { + b'/' if bytes.get(i + 1) == Some(&b'/') => { + while i < bytes.len() && bytes[i] != b'\n' { + i += 1; + } + } + b'/' if bytes.get(i + 1) == Some(&b'*') => { + i += 2; + while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') { + i += 1; + } + i += 2; + } + b'"' => { + i += 1; + while i < bytes.len() && bytes[i] != b'"' { + i += if bytes[i] == b'\\' { 2 } else { 1 }; + } + i += 1; + } + b'{' => { + depth += 1; + i += 1; + } + b'}' => { + depth -= 1; + if depth == 0 { + return &src[open..=i]; + } + i += 1; + } + _ => i += 1, + } + } + panic!("unbalanced braces in the body following {anchor:?}"); +} + +/// Collect every `LoadedModel::` identifier named in `body`. +fn loaded_model_variants(body: &str) -> Vec { + let mut out: Vec = Vec::new(); + let mut rest = body; + while let Some(idx) = rest.find("LoadedModel::") { + let tail = &rest[idx + "LoadedModel::".len()..]; + let end = tail + .find(|c: char| !c.is_ascii_alphanumeric() && c != '_') + .unwrap_or(tail.len()); + let name = &tail[..end]; + if !name.is_empty() && !out.iter().any(|v| v == name) { + out.push(name.to_string()); + } + rest = &tail[end..]; + } + out +} + +/// The variant list `mtp_capable_target` accepts. The helper must name only +/// MTP-capable variants and decline everything else through its `_ =>` arm; +/// an explicit `LoadedModel::X => false` arm would be read here as capable. +fn mtp_capable_variants() -> Vec { + let body = mtp_fn_body(SRC_SPECULATIVE_BURST, "fn mtp_capable_target("); + let variants = loaded_model_variants(body); + assert!( + variants.len() >= 3, + "mtp_capable_target should name at least the three Gemma 4 variants, found {variants:?}" + ); + variants +} + +#[test] +fn mtp_capable_target_names_the_qwen35_family() { + // Guards the #1165 wiring itself: the four Qwen 3.5 variants (dense and + // MoE, text and VLM) must be MTP-capable alongside the Gemma 4 three. + let variants = mtp_capable_variants(); + for expected in [ + "Gemma4", + "Gemma4VLM", + "Gemma4Unified", + "Qwen35", + "Qwen35Moe", + "Qwen35VLM", + "Qwen35MoeVLM", + ] { + assert!( + variants.iter().any(|v| v == expected), + "mtp_capable_target must accept LoadedModel::{expected}; found {variants:?}" + ); + } +} + +#[test] +fn every_mtp_dispatch_site_covers_every_capable_variant() { + let capable = mtp_capable_variants(); + let mut missing: Vec = Vec::new(); + for (label, src, anchor) in MTP_DISPATCH_SITES { + let body = mtp_fn_body(src, anchor); + let named = loaded_model_variants(body); + for variant in &capable { + if !named.iter().any(|v| v == variant) { + missing.push(format!("{label}: LoadedModel::{variant}")); + } + } + } + assert!( + missing.is_empty(), + "these MTP-capable variants are accepted by `mtp_capable_target` but are not \ + handled at every dispatch site, so they would SILENTLY decline to classic \ + decode on the paths below:\n {}\n\ + Add the missing arm(s), or narrow `mtp_capable_target`.", + missing.join("\n "), + ); +} diff --git a/src/server/batch/speculative_slice.rs b/src/server/batch/speculative_slice.rs index b5948e83c..a94a44e21 100644 --- a/src/server/batch/speculative_slice.rs +++ b/src/server/batch/speculative_slice.rs @@ -19,7 +19,9 @@ //! scheduler tick, so every concurrent classic-decode row stalls for the //! whole burst (the head-of-line block measured by //! `BurstFinalized::burst_wall_ms`, issue #638 / PR #733). This module -//! removes that block for the B=1 MTP arm on the Gemma 4 family: the +//! removes that block for the B=1 MTP arm on every family +//! [`super::speculative_burst::mtp_capable_target`] accepts (the Gemma 4 +//! assistant drafter and, since issue #1165, the Qwen 3.5 MTP drafter): the //! request is served as a sequence of SLICES, one per scheduler tick: //! slice 0 is the prefill + seed //! ([`mlxcel_core::speculative::mtp::MtpGenerator::begin_session`]), and @@ -31,10 +33,11 @@ //! //! ## Why the generator is reconstructed every tick //! -//! `MtpGenerator` owns its target adapter, and the Gemma 4 adapters -//! (`Gemma4MtpTargetAdapter` and friends) BORROW the model, so the -//! generator is self-referential with respect to the scheduler and cannot -//! be stashed across ticks (the blocker PR #733 documented). What CAN be +//! `MtpGenerator` owns its target adapter, and every B=1 adapter +//! (`Gemma4MtpTargetAdapter` and friends, and `Qwen35MtpTargetAdapter` / +//! `Qwen35VLMtpTargetAdapter`) BORROWS the model, so the generator is +//! self-referential with respect to the scheduler and cannot be stashed +//! across ticks (the blocker PR #733 documented). What CAN be //! stashed is everything the round loop actually carries between rounds: //! //! - [`MtpSliceJob`] owns the request's `SequenceInfo`, the generator's @@ -53,9 +56,17 @@ //! the resetting //! [`super::speculative_burst::WorkerDrafterSlot::return_drafter`] runs //! when the whole session finishes, and (since issue #746) at every -//! park boundary when the slot rotates to another grantee, which is safe -//! because the MTP assistant drafter's reset is the trait default no-op -//! (see [`MtpSliceJob::attach_drafter`]). +//! park boundary when the slot rotates to another grantee. Correctness +//! does not depend on what `reset` does to drafter-side state: the next +//! slice always re-anchors via `set_shared_kv` regardless. The two families +//! differ in what a park boundary costs, though. The Gemma 4 assistant +//! drafter's `reset` is the trait-default no-op, so a park loses nothing. +//! The Qwen 3.5 MTP drafter's `reset` DESTROYS its accumulated KV history +//! (see `Qwen35MtpDraftModel::reset`), so a park under grant rotation +//! restarts that drafter's draft context from empty; acceptance under +//! concurrent rotation for that family is correspondingly worse than a +//! single-request measurement shows (see +//! `docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md`, issue #1165). //! //! ## Streaming //! @@ -68,10 +79,12 @@ //! //! ## Scope //! -//! B=1 MTP on the Gemma 4 family only. The DFlash arm and the B>1 batched -//! arms keep the legacy run-to-completion burst (their generators do not -//! expose a resumable step API yet), and `MLXCEL_MTP_TICK_SLICE=0` forces -//! the MTP arm back onto the legacy burst as an operator escape hatch. +//! B=1 MTP for every [`super::speculative_burst::mtp_capable_target`] +//! family: the Gemma 4 assistant drafter and, since issue #1165, the +//! Qwen 3.5 MTP drafter. The DFlash arm and the B>1 batched arms keep the +//! legacy run-to-completion burst (their generators do not expose a +//! resumable step API yet), and `MLXCEL_MTP_TICK_SLICE=0` forces the MTP +//! arm back onto the legacy burst as an operator escape hatch. //! //! ## Slot rotation (issue #746) //! @@ -390,18 +403,28 @@ impl MtpSliceJob { /// Correctness of the round-trip through /// [`super::speculative_burst::WorkerDrafterSlot::return_drafter`] /// (which calls [`Drafter::reset`]) and back: `step_session`'s - /// resumability contract requires that nothing the shared-KV re-arm - /// does not rebuild is destroyed between rounds. For the MTP arm this - /// holds because (a) the Gemma 4 assistant drafter does not override - /// `reset`, so its reset is the trait default no-op (verified: it - /// unloads no weights and unbinds nothing), (b) its bind state - /// (captured target embedding + resolved LM head) is derived from the - /// worker's one target model and is recomputed identically by the - /// promotion-time re-bind, and (c) every per-round value the drafter - /// consumes (`shared_kv`, offsets, RoPE position) is a stateless - /// overwrite performed by `set_shared_kv` at the top of EVERY round - /// from the session's own stored `MtpVerifyOutput`, so rounds of - /// another session in between cannot leak state into this one. + /// resumability contract requires only that nothing `set_shared_kv`'s + /// per-round overwrite does not already rebuild is load-bearing across + /// the round-trip. For the MTP arm this holds because (a) what `reset` + /// does to drafter-side history varies by family and correctness does + /// NOT depend on it either way: the Gemma 4 assistant drafter does not + /// override `reset` (trait default no-op, verified: it unloads no + /// weights and unbinds nothing), while the Qwen 3.5 MTP drafter's + /// `reset` DESTROYS its accumulated KV history (see + /// `Qwen35MtpDraftModel::reset`) rather than preserving it, (b) its + /// bind state (captured target embedding + resolved LM head) is derived + /// from the worker's one target model and is recomputed identically by + /// the promotion-time re-bind, and (c) every per-round value the + /// drafter consumes (`shared_kv`, offsets, RoPE position) is a + /// stateless overwrite performed by `set_shared_kv` at the top of EVERY + /// round from the session's own stored `MtpVerifyOutput`, so rounds of + /// another session in between cannot leak state into this one. What + /// destroyed history costs the Qwen 3.5 MTP family is not correctness + /// but draft quality: a park boundary restarts that drafter's draft + /// context from empty, so acceptance under grant rotation is worse than + /// a single-request measurement shows (see + /// `docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md`, issue + /// #1165, uncharacterized). pub(crate) fn attach_drafter(&mut self, drafter: Box) { debug_assert!( self.drafter.is_none(), diff --git a/src/server/batch/speculative_slice_tests.rs b/src/server/batch/speculative_slice_tests.rs index 77a6ee34a..cddc8662c 100644 --- a/src/server/batch/speculative_slice_tests.rs +++ b/src/server/batch/speculative_slice_tests.rs @@ -115,6 +115,7 @@ impl MockMtpTarget { ], kv_offset, bonus_position: kv_offset, + verify_hidden_full: None, } } } diff --git a/src/server/speculative_dispatch.rs b/src/server/speculative_dispatch.rs index 5d22031c9..87ed0f543 100644 --- a/src/server/speculative_dispatch.rs +++ b/src/server/speculative_dispatch.rs @@ -245,7 +245,8 @@ impl SpeculativeDispatch { })?; let user_requested_explicit_kind = explicit_kind.is_some(); - let block_size = resolve_draft_block_size(config.draft_block_size, resolved_kind); + let block_size = + resolve_draft_block_size(config.draft_block_size, resolved_kind, &draft_model_path); // Dispatch matrix: // - explicit MTP / DFlash → kind-specific generator diff --git a/tests/qwen38_mtp_chain_parity.rs b/tests/qwen38_mtp_chain_parity.rs new file mode 100644 index 000000000..adf704b4d --- /dev/null +++ b/tests/qwen38_mtp_chain_parity.rs @@ -0,0 +1,628 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Qwen 3.5-family MTP **chain** parity: multi-token verify blocks vs the +//! single-token decode chain. +//! +//! ## What this pins (and what the sibling test cannot see) +//! +//! `tests/qwen35_mtp_verify_parity.rs` compares a verify block's per-position +//! argmax against a fresh full-prefix re-prefill per position — both arms are +//! multi-token block forwards, so their gated-delta recurrent state evolves +//! through the same in-block float32 carry. The MTP/DFlash **round loop**, +//! however, interleaves multi-token verify blocks into a chain whose +//! byte-parity reference is the classic **single-token** decode chain, where +//! the recurrent state is materialized (and dtype-rounded) after every +//! token. This test drives both chain shapes over the same token sequence +//! from the same prefill and reports the first argmax divergence, which is +//! the exact mechanism behind a temperature-0 MTP output drifting from +//! classic decode at a near-tie logit. +//! +//! ## Method +//! +//! 1. **Reference chain** — prefill a fixed prompt, then decode N tokens one +//! at a time through `forward_speculative` (`T = 1` per call; the capture +//! path's single-token branch is op-identical to classic decode). +//! 2. **Block chain** — fresh caches, same prefill, then replay the +//! reference chain's tokens in `T = K` verify blocks (fully-accepted +//! rounds: no rollback involved) and compare every position's argmax +//! against the reference chain. +//! +//! A mismatch here means the multi-token forward itself (gated-delta block +//! scan and/or attention) is not byte-equal to the single-token chain — no +//! drafter, no rollback, no sampling in sight. +//! +//! ## Invocation +//! +//! ```bash +//! cargo test --test qwen38_mtp_chain_parity --release --features metal,accelerate -- --ignored --nocapture +//! ``` +//! +//! `#[ignore]`-gated as real-model heavy. Uses the first present candidate +//! checkpoint (Qwen3.8-27B preferred, smaller Qwen 3.5 sizes accepted). + +mod common; + +use common::repo_model_dir; + +use mlxcel::models::Qwen35Model; +use mlxcel::models::qwen3_next::Qwen3NextCache; +use mlxcel::{LoadedModel, initialize_runtime, load_model}; + +const CANDIDATE_TARGETS: &[&str] = &[ + "qwen3.8-27b-4bit", + "qwen3.5-0.8b-4bit", + "qwen3.5-2b-4bit", + "qwen3.5-4b-4bit", +]; + +/// Fixed in-vocab prompt ids (arbitrary; only chain agreement matters). +const PROMPT_TOKENS: &[i32] = &[785, 3974, 13876, 8533, 4290, 374, 264, 1273]; + +/// Chain length. Long enough to catch the observed divergence density +/// (flips showed up within 40-200 generated tokens at 27B). +const CHAIN_LEN: usize = 240; + +/// Verify-block width for the block chain (draft 2 + bonus, the published +/// `block_size: 3` shape). +const BLOCK: usize = 3; + +fn argmax_positions(logits: &mlxcel_core::MlxArray) -> Vec { + let shape = mlxcel_core::array_shape(logits); + let positions = shape[1]; + let argmax = mlxcel_core::argmax_last_axis(logits); + mlxcel_core::eval(&argmax); + let mut out = Vec::with_capacity(positions as usize); + for i in 0..positions { + let cell = mlxcel_core::slice(&argmax, &[0, i], &[1, i + 1]); + out.push(mlxcel_core::item_i32(&cell)); + } + out +} + +/// Resolve the inner text model, reaching through the VLM wrapper (most +/// mlx-community Qwen 3.5-family checkpoints load as the VLM variant). +fn as_qwen35(loaded: &LoadedModel) -> Option<&Qwen35Model> { + match loaded { + LoadedModel::Qwen35(m) | LoadedModel::Qwen35Moe(m) => Some(m), + LoadedModel::Qwen35VLM(vlm) | LoadedModel::Qwen35MoeVLM(vlm) => Some(&vlm.text_model), + _ => None, + } +} + +fn forward_tokens(model: &Qwen35Model, tokens: &[i32], caches: &mut [Qwen3NextCache]) -> Vec { + let arr = mlxcel_core::from_slice_i32(tokens, &[1, tokens.len() as i32]); + let out = model.forward_speculative(&arr, caches, &[]); + argmax_positions(&out.logits) +} + +#[test] +#[ignore] +fn block_verify_chain_matches_single_token_chain() { + let target_dir = CANDIDATE_TARGETS + .iter() + .map(|name| repo_model_dir(name)) + .find(|dir| dir.join("config.json").exists()); + let Some(target_dir) = target_dir else { + eprintln!("skipping: no candidate Qwen 3.5-family checkpoint under models/"); + return; + }; + + let _runtime = initialize_runtime(); + let (loaded, _tokenizer) = load_model(&target_dir).expect("load target"); + let model = as_qwen35(&loaded).expect("expected a Qwen 3.5-family target"); + + // Reference chain: prefill + T=1 decode steps. + let mut ref_caches = model.make_speculative_caches_for_test(); + let prefill_argmax = forward_tokens(model, PROMPT_TOKENS, &mut ref_caches); + let mut chain: Vec = vec![*prefill_argmax.last().expect("prefill argmax")]; + for _ in 0..CHAIN_LEN { + let last = *chain.last().expect("non-empty chain"); + let next = forward_tokens(model, &[last], &mut ref_caches); + chain.push(next[0]); + } + + // Block chain: fresh caches, same prefill, replay the reference tokens + // in fully-accepted T=BLOCK verify rounds. + let mut blk_caches = model.make_speculative_caches_for_test(); + let _ = forward_tokens(model, PROMPT_TOKENS, &mut blk_caches); + let mut mismatches: Vec<(usize, i32, i32)> = Vec::new(); + let mut pos = 0usize; // index into `chain`: the block starts at chain[pos] + while pos + BLOCK < chain.len() { + let block: Vec = chain[pos..pos + BLOCK].to_vec(); + let argmax = forward_tokens(model, &block, &mut blk_caches); + for (offset, &got) in argmax.iter().enumerate() { + let expect = chain[pos + offset + 1]; + if got != expect { + mismatches.push((pos + offset + 1, expect, got)); + } + } + pos += BLOCK; + } + + if mismatches.is_empty() { + eprintln!( + "block chain matches the single-token chain over {} tokens (block={})", + chain.len(), + BLOCK + ); + } else { + let first = mismatches[0]; + panic!( + "block-vs-single-token chain diverged at {} position(s) over {} tokens \ + (block={}); first at chain index {} (single-token argmax {}, block argmax {})", + mismatches.len(), + chain.len(), + BLOCK, + first.0, + first.1, + first.2, + ); + } +} + +// =========================================================================== +// Gated-delta chain-parity kernel: bitwise pin + isolated cost measurement. +// =========================================================================== + +use mlxcel::models::gated_delta::{gated_delta_update, gated_delta_update_chain_parity}; + +/// Qwen3.8-27B GDN geometry: Hk=16, Hv=48, Dk=Dv=128 (f16 activations). +const GDN_HK: i32 = 16; +const GDN_HV: i32 = 48; +const GDN_D: i32 = 128; + +fn gdn_inputs( + t: i32, + seed: f32, +) -> ( + mlxcel_core::UniquePtr, + mlxcel_core::UniquePtr, + mlxcel_core::UniquePtr, + mlxcel_core::UniquePtr, + mlxcel_core::UniquePtr, +) { + let fill = |n: usize, scale: f32| -> Vec { + (0..n) + .map(|i| ((i as f32 * 0.37 + seed).sin()) * scale) + .collect() + }; + let qk_n = (t * GDN_HK * GDN_D) as usize; + let v_n = (t * GDN_HV * GDN_D) as usize; + let ab_n = (t * GDN_HV) as usize; + let f16 = |data: &[f32], shape: &[i32]| { + let arr = mlxcel_core::from_slice_f32(data, shape); + mlxcel_core::astype(&arr, mlxcel_core::dtype::FLOAT16) + }; + ( + f16(&fill(qk_n, 0.5), &[1, t, GDN_HK, GDN_D]), + f16(&fill(qk_n, 0.4), &[1, t, GDN_HK, GDN_D]), + f16(&fill(v_n, 0.3), &[1, t, GDN_HV, GDN_D]), + f16(&fill(ab_n, 0.2), &[1, t, GDN_HV]), + f16(&fill(ab_n, 0.1), &[1, t, GDN_HV]), + ) +} + +fn gdn_slice_t( + arr: &mlxcel_core::MlxArray, + t: i32, +) -> mlxcel_core::UniquePtr { + let shape = mlxcel_core::array_shape(arr); + let mut starts = vec![0; shape.len()]; + let mut stops = shape.clone(); + starts[1] = t; + stops[1] = t + 1; + mlxcel_core::slice(arr, &starts, &stops) +} + +fn raw_bytes(arr: &mlxcel_core::MlxArray) -> Vec { + mlxcel_core::eval(arr); + mlxcel_core::array_to_raw_bytes(arr) +} + +/// The chain-parity contract itself, pinned bitwise on synthetic 27B-geometry +/// inputs: a `T = 3` parity-kernel block must produce the SAME per-position +/// outputs and the SAME final state as three consecutive `T = 1` calls of the +/// standard kernel (the classic decode chain). This is the invariant the MTP +/// temperature-0 exactness gate rests on; the standard kernel fails it by +/// design (float32 in-block state carry). +#[test] +fn parity_kernel_block_is_bitwise_equal_to_single_token_chain() { + if !mlxcel_core::metal_is_available() { + eprintln!("skipping: chain-parity kernel is Metal-only"); + return; + } + let _runtime = initialize_runtime(); + let t = 3; + let (q, k, v, a, b) = gdn_inputs(t, 0.123); + let a_log_v: Vec = (0..GDN_HV).map(|i| (i as f32 * 0.05).cos() * 0.5).collect(); + let dt_v: Vec = (0..GDN_HV).map(|i| (i as f32 * 0.03).sin() * 0.5).collect(); + let a_log = mlxcel_core::from_slice_f32(&a_log_v, &[GDN_HV]); + let dt_bias = mlxcel_core::from_slice_f32(&dt_v, &[GDN_HV]); + + // Block: parity kernel over T=3, state=None. + let (y_block, state_block) = + gated_delta_update_chain_parity(&q, &k, &v, &a, &b, &a_log, &dt_bias, None, None); + + // Chain: three standard T=1 calls, state carried between calls exactly + // as the classic decode chain carries it (storage-dtype round trip). + let mut state: Option> = None; + let mut y_chain: Vec> = Vec::new(); + for step in 0..t { + let (y_t, next_state) = gated_delta_update( + &gdn_slice_t(&q, step), + &gdn_slice_t(&k, step), + &gdn_slice_t(&v, step), + &gdn_slice_t(&a, step), + &gdn_slice_t(&b, step), + &a_log, + &dt_bias, + state.as_deref(), + None, + ); + y_chain.push(y_t); + state = Some(next_state); + } + + for step in 0..t { + assert_eq!( + raw_bytes(&gdn_slice_t(&y_block, step)), + raw_bytes(y_chain[step as usize].as_ref().unwrap()), + "parity block output at step {step} must be bitwise equal to the T=1 chain" + ); + } + assert_eq!( + raw_bytes(&state_block), + raw_bytes(state.unwrap().as_ref().unwrap()), + "parity block final state must be bitwise equal to the T=1 chain state" + ); +} + +/// Arm-attribution guard for [`parity_kernel_cost_vs_standard`]. +/// +/// `gated_delta_update_chain_parity` reads the `MLXCEL_GDN_CHAIN_PARITY` env +/// var (default on, see `chain_parity_enabled` in `src/models/gated_delta.rs`) +/// and silently degrades to the standard kernel when it is set to `0`. If +/// that happens, the "parity" arm below is bit-for-bit the same dispatch as +/// the "standard" arm and any cost ratio printed would be a path compared +/// against itself, indistinguishable from dispatch noise. This is precisely +/// the failure mode CLAUDE.md's "prove the two arms differ before reporting +/// any number" rule exists to prevent (see epic #909, three prior +/// benchmarks that reported self-comparison noise as a result). +/// +/// The discriminator: the two kernels are defined to diverge in their final +/// recurrent state whenever `T > 1` (the parity kernel rounds state through +/// the storage dtype after every in-block step; the standard kernel carries +/// float32 state across the whole block and only rounds once at the end). +/// `T = 1` is documented to be byte-identical BY DESIGN, with no in-block +/// state carry to round differently, so it cannot serve as the +/// discriminator; this check therefore always runs at `T = 3`, matching the +/// shape already pinned bitwise-equal-vs-different by +/// [`parity_kernel_block_is_bitwise_equal_to_single_token_chain`] and +/// [`block_verify_chain_matches_single_token_chain`] respectively. A failure +/// here means the whole table below would be unattributable, so the test +/// panics instead of printing it. +fn assert_parity_and_standard_arms_are_distinguishable() { + let t = 3; + let (q, k, v, a, b) = gdn_inputs(t, 0.777); + let a_log_v: Vec = (0..GDN_HV).map(|i| (i as f32 * 0.05).cos() * 0.5).collect(); + let dt_v: Vec = (0..GDN_HV).map(|i| (i as f32 * 0.03).sin() * 0.5).collect(); + let a_log = mlxcel_core::from_slice_f32(&a_log_v, &[GDN_HV]); + let dt_bias = mlxcel_core::from_slice_f32(&dt_v, &[GDN_HV]); + + let (_, state_std) = gated_delta_update(&q, &k, &v, &a, &b, &a_log, &dt_bias, None, None); + let (_, state_par) = + gated_delta_update_chain_parity(&q, &k, &v, &a, &b, &a_log, &dt_bias, None, None); + assert_ne!( + raw_bytes(&state_std), + raw_bytes(&state_par), + "arm-attribution guard failed: the parity kernel produced the SAME final state as the \ + standard kernel at T={t}, so the two arms are indistinguishable in this run and no \ + cost ratio below would be attributable to the parity kernel. This almost always means \ + MLXCEL_GDN_CHAIN_PARITY=0 is set in the environment, which makes \ + `gated_delta_update_chain_parity` silently fall back to the standard kernel (see \ + `chain_parity_enabled` in src/models/gated_delta.rs). Unset it (or set it to a value \ + other than \"0\") and re-run." + ); +} + +/// Isolated cost of the parity kernel vs the standard kernel on the real +/// verify shapes (medians over repeated dispatches). Run manually: +/// +/// ```bash +/// cargo test --test qwen38_mtp_chain_parity --release --features metal,accelerate -- --ignored --nocapture parity_kernel_cost +/// ``` +#[test] +#[ignore] +fn parity_kernel_cost_vs_standard() { + if !mlxcel_core::metal_is_available() { + eprintln!("skipping: chain-parity kernel is Metal-only"); + return; + } + let _runtime = initialize_runtime(); + assert_parity_and_standard_arms_are_distinguishable(); + let a_log_v: Vec = (0..GDN_HV).map(|i| (i as f32 * 0.05).cos() * 0.5).collect(); + let dt_v: Vec = (0..GDN_HV).map(|i| (i as f32 * 0.03).sin() * 0.5).collect(); + let a_log = mlxcel_core::from_slice_f32(&a_log_v, &[GDN_HV]); + let dt_bias = mlxcel_core::from_slice_f32(&dt_v, &[GDN_HV]); + + for t in [1_i32, 3, 4, 64] { + let (q, k, v, a, b) = gdn_inputs(t, 0.5); + let time = |parity: bool| -> f64 { + // Warmup (includes the one-time kernel JIT). + for _ in 0..20 { + let (y, s) = if parity { + gated_delta_update_chain_parity( + &q, &k, &v, &a, &b, &a_log, &dt_bias, None, None, + ) + } else { + gated_delta_update(&q, &k, &v, &a, &b, &a_log, &dt_bias, None, None) + }; + mlxcel_core::eval(&y); + mlxcel_core::eval(&s); + } + let mut samples: Vec = Vec::with_capacity(100); + for _ in 0..100 { + let start = std::time::Instant::now(); + let (y, s) = if parity { + gated_delta_update_chain_parity( + &q, &k, &v, &a, &b, &a_log, &dt_bias, None, None, + ) + } else { + gated_delta_update(&q, &k, &v, &a, &b, &a_log, &dt_bias, None, None) + }; + mlxcel_core::eval(&y); + mlxcel_core::eval(&s); + samples.push(start.elapsed().as_secs_f64() * 1e6); + } + samples.sort_by(|x, y| x.partial_cmp(y).unwrap()); + samples[samples.len() / 2] + }; + let std_us = time(false); + let par_us = time(true); + eprintln!( + "T={t}: standard {std_us:.1} us, parity {par_us:.1} us, ratio {:.3}", + par_us / std_us + ); + } +} + +// =========================================================================== +// Verify-forward cost scaling: where does the T>1 block spend its time? +// =========================================================================== + +/// Times three shapes of the target forward per T, medians over repeats: +/// +/// - `capture`: `forward_speculative` (per-position verify attention + GDN +/// snapshot capture) — the MTP/DFlash verify pass. +/// - `batched`: `forward_prefill_with_last_hidden` (batched causal +/// attention, normal GDN, no capture) — the classic prefill shape over the +/// same block. +/// - The `T = 1` `capture` row is the classic decode step cost reference. +/// +/// The gap between `capture` and `batched` at the same T attributes the +/// verify machinery's own cost (per-position attention dispatches + snapshot +/// copies); the gap between `batched` at T and `T = 1` attributes how well +/// the block amortizes the weight read. Run manually: +/// +/// ```bash +/// cargo test --test qwen38_mtp_chain_parity --release --features metal,accelerate -- --ignored --nocapture verify_forward_cost +/// ``` +#[test] +#[ignore] +fn verify_forward_cost_scaling() { + let Some(target_dir) = CANDIDATE_TARGETS + .iter() + .map(|name| repo_model_dir(name)) + .find(|dir| dir.join("config.json").exists()) + else { + eprintln!("skipping: no candidate Qwen 3.5-family checkpoint under models/"); + return; + }; + let _runtime = initialize_runtime(); + let (loaded, _tokenizer) = load_model(&target_dir).expect("load target"); + let model = as_qwen35(&loaded).expect("expected a Qwen 3.5-family target"); + let last_layer = [model.num_layers() - 1]; + + // Warm caches with a small prompt so every timed call runs in the + // decode-tail regime (non-empty prefix), then time repeated blocks. + for t in [1_i32, 2, 3, 4, 8] { + let tokens: Vec = (0..t).map(|i| 700 + i).collect(); + let arr = mlxcel_core::from_slice_i32(&tokens, &[1, t]); + + let time_shape = |capture: bool| -> f64 { + let mut caches = model.make_speculative_caches_for_test(); + let prompt = + mlxcel_core::from_slice_i32(PROMPT_TOKENS, &[1, PROMPT_TOKENS.len() as i32]); + let (logits, _hidden) = model.forward_prefill_with_last_hidden(&prompt, &mut caches); + mlxcel_core::eval(&logits); + let mut samples: Vec = Vec::new(); + for i in 0..24 { + let start = std::time::Instant::now(); + if capture { + let out = model.forward_speculative(&arr, &mut caches, &last_layer); + mlxcel_core::eval(&out.logits); + } else { + let (logits, _h) = model.forward_prefill_with_last_hidden(&arr, &mut caches); + mlxcel_core::eval(&logits); + } + if i >= 4 { + samples.push(start.elapsed().as_secs_f64() * 1000.0); + } + } + samples.sort_by(|x, y| x.partial_cmp(y).unwrap()); + samples[samples.len() / 2] + }; + + let capture_ms = time_shape(true); + let batched_ms = time_shape(false); + eprintln!("T={t}: capture {capture_ms:.1} ms, batched {batched_ms:.1} ms"); + } +} + +// =========================================================================== +// Adapter-level chain parity: the REAL MtpTarget adapter vs classic decode. +// =========================================================================== + +use mlxcel::models::qwen3_5_mtp_target::Qwen35MtpTargetAdapter; +use mlxcel_core::generate::{LanguageModel, SamplingConfig}; +use mlxcel_core::sampling::LogprobsConfig; +use mlxcel_core::speculative::mtp::target::MtpTarget; + +/// Classic-path chain: prefill via `LanguageModel::forward` (the exact +/// numeric path the drafter-less CLI decode takes for a text-only request on +/// this family) then `T = 1` decode steps, greedy argmax throughout. +fn classic_chain(model: &Qwen35Model, prompt: &[i32], len: usize) -> Vec { + ::reset_runtime_state(model); + let prompt_arr = mlxcel_core::from_slice_i32(prompt, &[1, prompt.len() as i32]); + let logits = ::forward(model, &prompt_arr, &mut [], None); + let mut chain = vec![*argmax_positions(&logits).last().expect("prefill argmax")]; + for _ in 0..len { + let last = *chain.last().expect("non-empty"); + let arr = mlxcel_core::from_slice_i32(&[last], &[1, 1]); + let logits = ::forward(model, &arr, &mut [], None); + chain.push(argmax_positions(&logits)[0]); + } + chain +} + +fn greedy() -> SamplingConfig { + SamplingConfig { + temperature: 0.0, + ..SamplingConfig::default() + } +} + +/// Drive the real [`Qwen35MtpTargetAdapter`] with **perfect drafts** (the +/// classic chain's own tokens): every round fully accepts, no rollback ever +/// runs. Any divergence from the classic chain is therefore in the adapter's +/// prefill or the multi-token verify forward itself. +#[test] +#[ignore] +fn adapter_chain_with_perfect_drafts_matches_classic_chain() { + let Some(target_dir) = CANDIDATE_TARGETS + .iter() + .map(|name| repo_model_dir(name)) + .find(|dir| dir.join("config.json").exists()) + else { + eprintln!("skipping: no candidate Qwen 3.5-family checkpoint under models/"); + return; + }; + let _runtime = initialize_runtime(); + let (loaded, _tokenizer) = load_model(&target_dir).expect("load target"); + let model = as_qwen35(&loaded).expect("expected a Qwen 3.5-family target"); + + let chain = classic_chain(model, PROMPT_TOKENS, 240); + + ::reset_runtime_state(model); + let adapter = Qwen35MtpTargetAdapter::new(model, None); + let sampler = greedy(); + let logprobs = LogprobsConfig::default(); + let (bonus, _seed, _lp) = adapter.prefill_and_seed(PROMPT_TOKENS, &sampler, &[], &logprobs); + assert_eq!( + bonus, chain[0], + "first bonus from the adapter prefill differs from the classic prefill argmax" + ); + + let mut mismatches: Vec<(usize, i32, i32)> = Vec::new(); + let mut pos = 0usize; + while pos + BLOCK < chain.len() { + let verify_input: Vec = chain[pos..pos + BLOCK].to_vec(); + let out = adapter.verify_forward(&verify_input, &sampler, &logprobs); + for (offset, &got) in out.target_tokens.iter().enumerate() { + let expect = chain[pos + offset + 1]; + if got != expect { + mismatches.push((pos + offset + 1, expect, got)); + } + } + // Fully accepted: BLOCK - 1 drafts all matched (by construction the + // drafts ARE the classic tokens; a target mismatch is recorded above + // but the chain replay continues on the classic tokens). + let _ = adapter.verify_finalize(BLOCK - 1, BLOCK, out.captured); + pos += BLOCK; + } + + assert!( + mismatches.is_empty(), + "adapter verify chain (perfect drafts, no rollback) diverged from classic decode at \ + {} position(s); first at chain index {} (classic {}, verify {})", + mismatches.len(), + mismatches[0].0, + mismatches[0].1, + mismatches[0].2, + ); + eprintln!( + "adapter perfect-draft chain matches classic decode over {} tokens", + chain.len() + ); +} + +/// Drive the adapter with **always-wrong drafts**: every round rejects both +/// drafts (`accepted = 0`), so `verify_finalize` runs the KV trim + GDN +/// rollback replay on every round. Position 0 of each verify block only +/// conditions on the accepted prefix, so its argmax must still follow the +/// classic chain; any divergence isolates the rollback path. +#[test] +#[ignore] +fn adapter_chain_with_rejected_drafts_matches_classic_chain() { + let Some(target_dir) = CANDIDATE_TARGETS + .iter() + .map(|name| repo_model_dir(name)) + .find(|dir| dir.join("config.json").exists()) + else { + eprintln!("skipping: no candidate Qwen 3.5-family checkpoint under models/"); + return; + }; + let _runtime = initialize_runtime(); + let (loaded, _tokenizer) = load_model(&target_dir).expect("load target"); + let model = as_qwen35(&loaded).expect("expected a Qwen 3.5-family target"); + + let chain = classic_chain(model, PROMPT_TOKENS, 160); + + ::reset_runtime_state(model); + let adapter = Qwen35MtpTargetAdapter::new(model, None); + let sampler = greedy(); + let logprobs = LogprobsConfig::default(); + let (bonus, _seed, _lp) = adapter.prefill_and_seed(PROMPT_TOKENS, &sampler, &[], &logprobs); + assert_eq!(bonus, chain[0], "first bonus differs from classic prefill"); + + let mut mismatches: Vec<(usize, i32, i32)> = Vec::new(); + for pos in 0..chain.len() - 1 { + // Drafts guaranteed wrong: shift the classic continuation by one id. + let wrong = |tok: i32| (tok + 1) % 1000 + 10; + let verify_input = vec![chain[pos], wrong(chain[pos + 1]), wrong(chain[pos + 1])]; + let out = adapter.verify_forward(&verify_input, &sampler, &logprobs); + let got = out.target_tokens[0]; + let expect = chain[pos + 1]; + if got != expect { + mismatches.push((pos + 1, expect, got)); + } + // Zero accepted: rollback replay runs every round. + let _ = adapter.verify_finalize(0, BLOCK, out.captured); + } + + assert!( + mismatches.is_empty(), + "adapter verify chain (rejected drafts, rollback every round) diverged from classic \ + decode at {} position(s); first at chain index {} (classic {}, verify {})", + mismatches.len(), + mismatches[0].0, + mismatches[0].1, + mismatches[0].2, + ); + eprintln!( + "adapter rejected-draft chain matches classic decode over {} tokens", + chain.len() + ); +}