[inureyes]

--draft-model <dspark-dir> on an lfm2 / lfm2_moe target now runs the DFlash round loop with the greedy exact-match verify rule. A DSpark drafter is a 5-layer DFlash backbone fed five of the target's residual streams plus a rank-256 Markov token-transition head, so this adds the head, the DSpark config fields (markov_rank, rope_is_neox_style, runtime_block_size, enable_confidence_head, and num_target_layers lifted out of the nested dflash_config), the checkpoint's traditional RoPE pairing, the LFM2 verify forward with residual-stream capture, and the partial-accept rollback. Attention layers trim their KV cache; a short-conv state cannot be trimmed, so it is recomputed from the block's gated input at the committed prefix. Qwen35DFlashTarget generalizes into DFlashTargetModel, carrying make_dflash_caches, the first-draft prompt-row policy (Qwen keeps the last position, DSpark wants every prompt row) and the exactness gate, so one driver serves both families and #1343 can add a third with one impl.

Two LFM2 loader gaps that the bf16 originals expose are fixed alongside, because the pairing cannot load without them: rope_theta nested under rope_parameters was read as the 1e6 default, and unstacked per-expert feed_forward.experts.{e}.w1/w2/w3 tensors were never renamed, so MoE expert stacking never fired.

DSpark is greedy-only: a request with temperature > 0 and top_k != 1 is served by classic decode with one named warning, because the Markov head proposes a chain trained against greedy targets and there is no stochastic acceptance rule.

Real-checkpoint validation

Apple M5 Max (GPU generation 17), port 19339, against the release binaries built from this tree. The machine served other work during the session, so each throughput figure is one sample; the acceptance counters are deterministic and reproduced byte for byte across repeat runs and across a rebuild.

./target/release/mlxcel-server -m models/lfm2.5-2.6b-bf16 --model-draft models/lfm2.5-2.6b-dspark-bf16 --port 19339
curl -s localhost:19339/v1/chat/completions -H 'Content-Type: application/json' -d '{"model":"m","messages":[{"role":"user","content":"Write a detailed technical explanation of how speculative decoding works in large language model inference, covering the drafter, the verification step, and the acceptance rule."}],"max_tokens":256,"temperature":0,"logprobs":true}'

The baseline is the same request against a server started without --model-draft; identity is compared over the 256 logprobs.content[].token entries, both arms finishing on length at exactly 256 tokens.

The block-versus-chain probe declines on this host, for both pairs. That is the documented generation-15-and-newer condition and the shipped default: the burst falls back to classic decode, emits no timings, and serves the baseline's tokens (verified on the MoE pair with the drafter configured and the gate honored, 256 of 256 identical). The runs below therefore used MLXCEL_MTP_ALLOW_INEXACT=1, the existing research override.

• Dense pair (models/lfm2.5-2.6b-bf16 + models/lfm2.5-2.6b-dspark-bf16): all 256 token ids identical to the baseline at the default verify width 8 and again at --draft-block-size 10. Mean accepted length 3.59 at width 8 (draft_n 489, draft_n_accepted 184, draft_rounds 71) and 3.75 at width 10 (599 / 187 / 68). 203.3 and 212.1 tok/s against 72.3 tok/s classic (256 tokens in 3594 ms of request wall clock including a 66 ms prefill), so about 2.8x.
• MoE pair (models/lfm2.5-8b-a1b-bf16 + models/lfm2.5-8b-a1b-dspark-bf16, MLXCEL_FUSED_MOE=0 for the reference diff): mean accepted length 2.61 at width 8 (674 / 158 / 98) and 2.88 at width 10; 115.5 tok/s against 104.3 tok/s classic (2525 ms wall, 81 ms prefill), so about 1.11x, and the fused default measured the same classic figure (2532 ms). Token ids match for 153 of 256 and then diverge on a tie the two arms resolve differently: the baseline takes ':' at logprob -1.0312 and the block path ' formula' at -1.0000, and on the issue's shorter prompt both candidates read -1.4375. That is the byte-identity the probe declined, not a semantic difference; it reproduces at the same position across both widths and across runs, while the fused and unfused classic arms agree with each other.
• Temperature. "temperature": 0.7 on the dense pair is served by classic decode, with no timings and exactly one log line: DFlash speculative dispatch declined for seq seq-2: the drafter is greedy-only (DSpark) and the request samples with temperature 0.7 / top_k 50; serving it with classic decode.

Re-validated after the review and security fixes, against release binaries rebuilt from the final tree (no source file newer than the binary). The dense pair is 128 of 128 token ids identical to the same server without --model-draft, and the acceptance counters reproduce the pre-fix run digit for digit: 34 rounds, 234 proposed, 94 accepted, acceptance 0.402, 3.735 emitted per verify. That is the point of the re-run: the two allocation fixes below change what is kept, never what is computed, and identical counters are what says so. The temperature-0.7 decline reproduces with its one log line, and the mismatched pairing (--model-draft models/qwen3.5-4b-dflash on an LFM2 target, the case the new gate exists for) is refused by name at warmup and on the request, with the server still alive afterwards rather than aborted.

Review fixes

The implementation review found no CRITICAL and no HIGH. One MEDIUM was a regression this branch introduced and is fixed here, along with two unbounded checkpoint-supplied numbers and two unchecked weight shapes found while fixing it, and three doc statements the review showed to be false.

An LFM2 target paired with a non-DSpark DFlash drafter had no gate. Drafter::validate_target_compat cannot make that call: it reads the target as a LanguageModel, which carries no architecture string, so a plain DFlash drafter has no way to tell an LFM2 target from the Qwen 3.5 one it was published for, and returns Ok. DFlashDraftModel::forward does not check its fc input width either, so the mismatch reached MLX as a shape throw, and an MLX C++ exception crossing the cxx bridge aborts the process rather than failing the request. Before this branch the pairing declined at the burst's variant gate, which no longer rejects LFM2. The target is the side that knows its own family, so the policy is now DFlashTargetModel::requires_dspark_drafter(), true for Lfm2Model and Lfm2VlModel and the permissive default everywhere else, read by both run arms before any forward and answered with one operator-facing message naming --model-draft and the published pairings. Qwen 3.5 keeps the default; widening it there would change which DFlash pairings that family accepts, which is a separate call.

Two checkpoint-supplied numbers reached a buffer index or a round size unbounded, and are now in the DSpark pairing gate beside the fc-width and target_layer_ids checks. mask_token_id indexes the target's embedding table, because a DSpark drafter ships none of its own, and MLX range-checks no positive gather index, so an id past the last row read whatever followed the table in the buffer and fed it to the logits; the target half of the same gate already pins the target's vocabulary, so bounding the id against vocab_size bounds the gather. A runtime_verify_width() below two rows proposes nothing and would emit one token per burst. verify_width() saturates rather than wrapping, since usize::MAX there panics in a debug build and wraps to a zero-row width in a release one.

The Markov head's two factors were never measured against the config that sizes them. markov_w1 is gathered at a token id once per chain step, so a table with fewer rows than the vocabulary the chain draws from read past its own buffer for the same reason; a markov_w2 of the wrong width threw inside MLX, which crosses the bridge as a process abort rather than a load error. VanillaMarkovHead::from_weights now takes the vocabulary as well as the rank and checks both factors before either becomes a layer, skipping the rank check for a quantized factor (which bit-packs rank along its last axis only, so its stored width is a function of the bit depth) and not the row check.

Doc corrections, no behavior change. ShortConv::forward_with_capture claimed the snapshot costs no copy while the code takes an explicit one, which it has to, since conv_state is reassigned at the end of the call. configured_block_size implied a DFlash round-loop block-size policy that does not exist: the only readers are the MTP generator and its batched round loop, and a DSpark run gets its width from resolve_draft_block_size before the drafter is loaded. The dflash_target module doc claimed a new family costs one impl block plus a match arm, when it is three arms in speculative_burst and one of them still carries the LFM2 B = 1 decline as per-family policy outside the trait. The Drafter::is_dspark impl also delegates to the inherent method by qualified name rather than repeating its body, so the two cannot drift.

The review's second MEDIUM, a supports_batched() trait method to move that last hardcoded arm into the trait, is left for #1343, which is the family that needs it. The LOW findings (the probe measures one width, ProbeKey has no family discriminator, the two block-size hooks are declarative, the benchmark numbers stayed in this body) are recorded in the technical report with the reasons they stay open.

Security fixes

A second review pass looked specifically at what a downloaded checkpoint controls. Pointing --model-draft at a repository makes every field of its config.json and every weight shape in it a runtime input; MLX range-checks no positive gather index, and an MLX C++ exception crossing the cxx bridge aborts the process rather than failing the request. One HIGH, four MEDIUM and three LOW; six fixed, two left with the reason recorded in the technical report.

HIGH: the verify width was bounded below and not above, and this branch is what made it a control input. Before it the DFlash runtime block size came only from --draft-block-size or the flat constant 16, both operator-supplied. resolve_draft_block_size now peeks the drafter checkpoint and hands runtime_verify_width() to the scheduler as the server-wide block size, and a config setting block_size and runtime_block_size together escapes the eight-row default entirely. What that reaches first is the block-versus-chain exactness probe, which runs one single-token target forward per row, three draws over, on the scheduler thread, memoizes its verdict, and runs before the pairing gate that would otherwise refuse the config: at an absurd width that is a hang, not a slow request. runtime_verify_width() now clamps at DSPARK_MAX_VERIFY_WIDTH (32, against 8 to 10 on the published checkpoints), which is the bound that holds ahead of the probe because every caller goes through it; the gate separately refuses a config asking for more, so the operator is told rather than quietly served at 32.

MEDIUM: the mask_token_id bound was vacuous for a self-contained drafter. Its soundness argument was that the target half of the pairing gate pins the target's vocabulary, which holds only on the lazy-bind path every published DSpark checkpoint uses. A checkpoint shipping its own embed_tokens takes the other arm of bind and its table was never compared to anything. LmHead::Own had the same gap one step later, against the Markov head's vocab_size-shaped transition logits. Both are now measured at load for a DSpark config; a plain DFlash checkpoint is deliberately untouched.

MEDIUM: the Markov rank check was skipped, not narrowed, for a quantized factor, and every mlx-community conversion of these drafters is quantized, so the skipped case is the common one. The packed width is rank * bits / 32, so it is checkable: the check now accepts a width some supported bit depth explains, rather than trusting a declared bits a per-tensor override can contradict.

MEDIUM: two prompt-sized allocations that nothing read. first_hidden_for deep-copied the whole prompt hidden under the every-row policy, because mlxcel_core::copy is a real MLX Copy primitive and the signature took a borrow (168 MB at an 8k prompt on LFM2.5-2.6B, 671 MB at 32k); it now takes the array by value, which also ends the caller's retention of the original across the round loop. Separately, the short-conv rollback snapshots were captured on the prompt prefill, which is never rolled back, pinning roughly one prompt-sized buffer per conv layer alive through the eval that materialises the prefill (about 670 MB at 8k, about 2.7 GB at 32k on a 30-layer 2.6B checkpoint). SpeculativeTarget gains prefill_forward_with_capture_layers, defaulting to the verify hook so no other family changes, and LFM2 overrides it to skip the capture. The flag changes what is kept and never what is computed, which the new test pins on the logits, the captured hidden and the cache offsets at once.

LOW: an empty prompt or an empty window reached a negative slice start (a process abort) and an index panic respectively. Both were inherited unchanged from the Qwen arms and neither is reachable from the scheduler today; both are now request errors.

Test plan

• cargo test --profile test-fast --features metal,accelerate: -p mlxcel-core drafter::dflash 74 passed, and one --lib run over models::lfm2, server::batch::speculative_burst, server::batch::dflash_target, cli::speculative_args and models::detection 184 passed. Seven of those are new here: the verify-width ceiling including the usize::MAX saturating case, the self-contained DSpark vocabulary tables, the packed Markov rank check, the conv-capture flag, the mask-id and low-side width bounds, and the LFM2 non-DSpark pairing decline.
• cargo check and cargo clippy -- -D warnings clean for the root package AND for -p mlxcel-core separately, plus cargo fmt --all -- --check. Checking --lib --tests at the workspace root resolves to -p mlxcel and never compiles mlxcel-core's test target, which is how two test-only compile errors reached a suite run; both packages are checked from here on.
• --bin mlxcel offline_draft_model 4 passed.

Closes #1339

[inureyes]

Implementation Review Summary

Intent

--draft-model <dspark-dir> on an lfm2 / lfm2_moe target runs the existing DFlash round loop with the greedy exact-match verify rule at temperature 0 only, emitting output token-identical to classic decode.

Verified against the issue, line by line

• block_size 9 is read as gamma and verify_width() as block_size + 1; runtime_verify_width() caps at the eight-row default and --draft-block-size 10 restores the trained width. resolve_draft_block_size peeks the drafter config so the flat DFlash 16 never reaches a 9-proposal head, and returns None for a plain DFlash checkpoint so the Qwen default is untouched.
• The DSpark draft input is [anchor] ++ mask * (gamma - 1) at gamma rows and every one of the gamma logit positions is consumed, position 0 included. dspark_draft_block_uses_all_positions asserts that directly rather than by shape.
• The Markov chain is seeded with the anchor and feeds each step's argmax into the next step's gather, built as one lazy graph with a single host sync per round. markov_chain_feeds_previous_token_into_next_step is the test that a chain ignoring its own output would fail.
• rope_is_neox_style: false reaches all three fast_rope call sites as traditional = true, and the test asserts the attention output changes rather than just the struct field.
• Context K/V only in the drafter cache with proposal K/V concatenated after, mask = None, append-only and never trimmed: unchanged pre-existing DFlashAttention / DFlashKVCache behavior, correctly left alone.
• prefer_requested_block_size() == true and the target's tied embed_tokens borrowed as both input embedding and LM head at bind.

Short-conv rollback

n = accepted + 1, attention KVCache::trim(block_size - n), short-conv recomputed as the last L_cache - 1 rows of concat(prev_state_or_zeros, bx_block[:, :n]). Two things I checked specifically: prev_state is mlx::core::copy (a real copy node, confirmed at mlx_cxx_bridge.cpp:1078) taken after bx is computed and before *conv_state is assigned, so it cannot be the post-block state; and the zeros fallback is [bx.shape[0], L_cache - 1, bx.shape[2]] in bx's dtype, matching conv_padding()'s causal (l_cache - 1, 0). conv_rollback_matches_committed_prefix compares offset, state and the next decode step's logits against a genuine one-pass reference, which is the right shape of test for this.

DFlashTargetModel readiness for #1343

Ready. Every added hook defaults to the Qwen 3.5 behavior (first_hidden_rows to LastPromptPosition, enable_speculative_buffers to a no-op, exactness_allows to true), concat_captured_hidden and first_hidden_for(LastPromptPosition, ..) are equivalent to the code they replaced rather than merely similar, and no inherent method shadows a new trait method (Qwen35Model has mtp_exactness_allows, not exactness_allows, so the Qwen DFlash path does not begin gating on a probe it never gated on). enable_speculative_buffers is present and called by both drivers before the prefill, which is what the Muse Glimmer rotating cache needs; rotating-cache rollback is already per-family through SpeculativeTarget::rollback_partial.

Findings, all left open

No CRITICAL and no HIGH. Nothing was changed on the branch except the technical report.

• MEDIUM. A mismatched non-DSpark pairing on an LFM2 target has no gate: validate_target_compat returns Ok immediately when the drafter is not DSpark and DFlashDraftModel::forward does not check its fc input width, so an LFM2 target with a genuine Qwen 3.5 DFlash drafter reaches an MLX shape throw through a non-Result cxx shim. Pre-existing in kind (a 4B Qwen target with a 27B DFlash drafter fails identically), but LFM2 targets used to decline at the variant gate. Narrow fix: decline in the burst gate when the target is LFM2 and the drafter is not DSpark. Broad fix: an fc-width check for every DFlash drafter, which changes which Qwen pairings are accepted and so is a maintainer call.
• MEDIUM. A third family touches four match arms, not the "one impl block plus a match arm" the module doc claims: the exactness gate, the drive! dispatch, the batched gate (which hardcodes the LFM2 B = 1 decline), and model_variant_label. A supports_batched() trait method would move the third into the trait; worth doing when feat(speculative): Muse Glimmer DFlash assistant drafter (muse_glimmer_assistant, 5 sliding-attention layers, bidirectional sliding mask) on the existing DFlash round loop #1343 lands.
• LOW. The probe measures one width. dflash_exactness_allows(block_size) probes at the configured verify width, but the round loop narrows bs near the token budget, and this project's benchmark guidance records that block-versus-chain disagreement varies strongly with width. Same shape as the MTP arms, so not a regression.
• LOW. ProbeKey is {block_size, hidden_size, num_hidden_layers} with no family discriminator and a docstring that assumes one target model per process. Router mode does not. Pre-existing, widened by a third family.
• LOW. configured_block_size() and prefer_requested_block_size() are implemented as the issue asked but only the MTP generator reads them; the DFlash round loop uses its constructed block_size directly. The no-backoff requirement holds vacuously, and the comment at drafter.rs:459 implies a DFlash block-size policy that does not exist.
• LOW. ShortConv::forward_with_capture's doc says "the state is replaced, never mutated in place, so holding the previous handle costs no copy" while the code takes a copy. The copy is correct and cheap; the comment contradicts it.
• LOW. The issue's Validation section asked for tok/s at both widths in docs/benchmark_results/; they are in the PR body instead. Given the host served other work and each throughput figure is one sample, keeping single-sample numbers out of the benchmark corpus is the right call, but the deviation is worth stating.

Verification

• All stated requirements implemented
• No placeholder/mock code remaining (no todo!, unimplemented!, FIXME or stub in the diff)
• Integrated into project code flow (burst gate, both generic drivers, CLI block-size resolution, offline rejection, detection)
• Project conventions followed (no shared hotspot in layers.rs / utils.rs / switch_layers.rs / gated_delta.rs touched, no local references/... path, docs/supported-models.md and docs/speculative-acceptance.md updated)
• Existing modules reused (UnifiedEmbedding / UnifiedLinear for the Markov head, mtp_exactness_gate for the probe, materialize_argmax_i32_vec, create_causal_mask, slice_axis, concatenate_many)
• No unintended structural changes (the Qwen35DFlashTarget generalization is behavior-preserving)
• Tests pass

Reproduced locally at review time: cargo check --lib --tests, cargo clippy --lib --tests -- -D warnings and cargo fmt --all -- --check clean; -p mlxcel-core drafter::dflash 69 passed, --lib models::lfm2 29, --lib server::batch::speculative_burst 67, --lib server::batch::dflash_target 2, --lib cli::speculative_args 22, --lib models::detection 60 (the body says 58). Every count matches the recorded test plan.

A bilingual technical report is committed on the branch at TECHNICAL_REPORTS/1751-lfm2-dspark-drafter-20260910.{en,ko}.md.