MTP and DFlash¶

Multi-Token Prediction (MTP) and DFlash are two ways to build a target-specific drafter. Both try to make a verification round yield several target-approved tokens, but their proposal mechanisms are different: MTP learns explicit future-token predictions, while DFlash uses a lightweight block-diffusion model to refine several proposal positions in parallel.

Neither term names a universal checkpoint format. The target, drafter, and runtime protocol must match. This page first explains the ideas, then gives the exact mlxcel support matrix and commands.

MTP: Learning More Than the Immediate Next Token¶

Ordinary causal-language-model training predicts the next token at each position. An MTP objective adds losses for tokens farther into the future. In a simplified form, a representation at position t is used to predict x_(t+1), x_(t+2), ... x_(t+K):

L_MTP = sum from k=1 to K of lambda_k * CE(p_k(. | h_t), x_(t+k))

The concrete architecture varies. A model can carry several prediction heads, a small recurrent future-prediction stack, or a separate assistant that consumes target hidden states. During inference, those learned components propose a chain or block; the full target still verifies it.

Why MTP Helps Decoding¶

MTP changes the drafter economics in two ways:

1. Target alignment: the drafter was trained to predict futures for a specific target family, so its accepted prefix can be longer than that of an unrelated small LM.
2. Reduced duplicated work: an assistant may reuse target embeddings, target hidden states, or prefix/KV information instead of running another complete language model from scratch.

It does not make the target intrinsically non-autoregressive. If the target rejects early, later proposals are discarded. Speed comes from a cheap proposal path plus efficient target block verification, not from trusting unverified MTP output.

MTP Is Also a Training Objective¶

Multi-token prediction can improve learned representations even when it is never used for speculative inference. Conversely, weights trained with an MTP objective are not automatically usable by every serving runtime. Inference needs the exact head/assistant architecture, masks, cache protocol, positional convention, and verification algorithm.

Gemma 4 Assistant Models¶

Gemma 4 publishes dedicated, small text-only assistant models for the Gemma 4 family. The assistant consumes the target model's last-layer activations, shares the target input embedding, and can reuse the target prefix state rather than performing an independent full prefill. It predicts several future tokens, which the Gemma 4 target validates in a dense verification pass.

This design is different from selecting an arbitrary small Gemma checkpoint. Use an assistant model card that explicitly names the target family and size. A text-only assistant can accelerate text generation from a multimodal target checkpoint, but that does not imply that the speculative tail understands image, audio, or video embeddings.

Gemma 4 MTP engages on Metal and CUDA. Before the first burst, the worker runs the same block-versus-chain exactness probe the Qwen 3.5 family uses (see MTP Exactness Probe), so a Gemma 4 pair declines to plain decode when its verify block is not byte-identical to single-token decoding on the running hardware.

Qwen 3.5-Family MTP¶

The Qwen ecosystem publishes models and configurations with multi-token prediction support; other runtimes can expose in-checkpoint NEXTN/MTP heads. mlxcel has a narrower, explicit contract:

• Pair a Qwen 3.5/Qwen3.8-family target with a separate drafter whose model_type is qwen3_5_mtp.
• The drafter's declared block_size is used when --draft-block-size is omitted; published checkpoints commonly use 3.
• The path is text-only, including text-only requests against a supported VLM-wrapped target.
• Built-in mtp.layers.* weights inside a target checkpoint are not executed. The loader excludes those weights, internal-mtp is not user-selectable, and its dedicated round loop remains unimplemented. Do not advertise a target checkpoint alone as sufficient for mlxcel MTP. The one shape that runs its checkpoint's own head is qwen4_exp, described below.

The distinction matters operationally: an upstream model card can correctly say “MTP included,” while mlxcel requires the external qwen3_5_mtp companion checkpoint.

GLM-4.7-Flash MTP¶

GLM-4.7-Flash (glm4_moe_lite) keeps one DeepSeek-V3-style next-token-prediction layer as model.layers.47.*. The community 4-bit conversions drop those tensors, so mlxcel runs the layer as a separate drafter directory that mlxcel split-mtp writes from the raw zai-org/GLM-4.7-Flash checkpoint:

mlxcel split-mtp -m ./models/glm-4.7-flash-bf16 -o ./models/glm-4.7-flash-mtp-4bit --q-bits 4

The result has model_type: glm4_moe_lite_mtp and pairs with any glm4_moe_lite target of the same geometry, offline and on the server. The drafter is fed the target's post-final-norm hidden state (MLXCEL_GLM_MTP_HIDDEN_TAP=pre selects the pre-norm residual for acceptance measurements; it never changes emitted tokens). The default verify block is num_nextn_predict_layers + 1, which is 2 on the published checkpoint, and split-mtp --block-size records a different width up to 16. Startup runs the same block-versus-chain exactness probe as Qwen MTP. Two limits apply: a request that adopted a prompt-cache prefix is served by plain decode, and an MTP-served request donates nothing to the prefix cache because its target K/V lives in a model-owned slot. See mlxcel split-mtp for the command's flags.

Qwen3.8-Flash-Next (qwen4_exp) MTP¶

qwen4_exp owns its per-sequence state, so speculative decoding needs a rollback that rebuilds that state from the accepted prefix. mlxcel implements one (src/models/qwen4_exp_rollback.rs): the KV cache and QSA index are trimmed, summaries of blocks that no longer complete are dropped and recomputed lazily, the GatedDeltaNet recurrence and both convolution windows are replayed over the accepted tokens, and the PLE n-gram context is recomputed from them so a rejected EOS never resets it. The same rollback backs --prompt-lookup on this family.

Two drafter arrangements work. A split directory whose model_type is qwen4_exp_mtp is the plain form. A target conversion that kept its mtp.* tensors can also serve as its own drafter: pass the target directory again as --draft-model (offline) or --model-draft (server). No flag selects it; the loader resolves mtp when the index or headers hold mtp.*, model.mtp.*, or language_model.mtp.* tensors. --draft-kind mtp on a checkpoint without the head fails with NoEmbeddedMtpHead, and a directory without those tensors resolves as before.

# Embedded head kept by the conversion
mlxcel generate -m models/Qwen3.8-Flash-Next-4bit-mtp \
  --draft-model models/Qwen3.8-Flash-Next-4bit-mtp -p "Hello" -n 256 --temp 0

# Split drafter directory
mlxcel-server -m models/Qwen3.8-Flash-Next-OptiQ-2bit \
  --model-draft models/Qwen3.8-Flash-Next-MTP --port 8080

--draft-block-size defaults to 4; the split drafter's declared block_size is the adaptive start depth. The MTP target prefill follows the classic plan (2048-token chunks, or the server's --prefill-chunk-size), because a single-call prefill diverged from classic decode on long prompts. The verify block is decode-equivalent (one query per QSA attention call, convolutions one row at a time, the fused GatedDeltaNet projection, dense gate projections row by row, unsorted expert gathers); without it the exactness probe declines the real checkpoint at block 4. On GPU generation 15 and later the remaining divergence source is qmv_wide, which the gate switches off, and prompt lookup on this family runs that gate too at max_draft + 1. A drafter with a different hc_count is refused before bind.

The path is Metal only, at B = 1. Requests with images or video, and B > 1 windows, decline to plain decode. The embedded head stays resident, about 1.61 GiB for the 4-bit checkpoint, and mlxcel inspect counts it once because the index total_size already includes it.

DFlash: Diffusion Used as a Drafter¶

DFlash is a lightweight block-diffusion drafter. Instead of generating draft tokens one by one, it starts with a block of unknown/noisy positions and iteratively refines them using target-conditioned features. The resulting block is only a proposal. The autoregressive target verifies the candidates and preserves authority over output.

This separates DFlash from a native diffusion language model:

• In DFlash speculation, diffusion is an auxiliary proposal engine and the target is still autoregressive.
• In DiffusionGemma or LLaDA-2 generation, block denoising/unmasking is the target model's native decoding algorithm; there is no separate autoregressive target validating every candidate. See Block Diffusion.

DFlash can reduce the serial work inside the drafter and often uses a wider block than MTP. Its costs are target-specific training, multiple draft refinement operations, specialized attention masks/cache binding, and wasted block work after early target disagreement.

DFlash-Family Drafters in mlxcel¶

All of these drafters borrow embed_tokens and lm_head from the target, so none of them loads as a standalone model; passing one to -m is rejected with that explanation. Each is paired only with its own target family, and a mismatched pair is refused by name.

The LFM2, Muse Glimmer, and Laguna pairings run behind a measured block-versus-chain exactness probe on the server, and the Laguna pairing also offline. When the probe declines, the server serves plain decode and offline generate exits with the verdict. MLXCEL_MTP_ALLOW_INEXACT=1 engages anyway and forfeits byte-identity with plain decode.

In the batched Qwen 3.5 DFlash round loop each row commits its own accept count instead of clamping the round to the shortest live row. The attention caches compact back to a per-row contiguous prefix, queries and keys rotate at each row's own positions with its stale gap masked, and each linear-attention layer replays its convolution and recurrent state over the accepted prefix. The shared drafter context is still a rectangular prefix, so the accept walk stays below what independent drafting would give. MLXCEL_DFLASH_PER_ROW_ACCEPTS=0 restores the clamp for A/B work, and the clamp is taken automatically when the target's attention caches are quantized, paged, or head-trimmed because those cannot be compacted row by row. At B = 4 on the validated M3 Ultra host the per-row path measured 116.05 tok/s against 83.06 for the clamp.

MTP and DFlash Compared¶

Support Matrix¶

Quick Commands¶

Gemma 4 MTP Server¶

mlxcel-server \
  --model ./models/gemma-4-31b-it-4bit \
  --model-draft ./models/gemma-4-31B-it-assistant-bf16 \
  --draft-kind mtp \
  --draft-block-size 4 \
  --parallel 1 \
  --max-batch-size 1 \
  --port 8080

Gemma 4 Offline MTP¶

mlxcel generate \
  --model ./models/gemma-4-31b-it-4bit \
  --draft-model ./models/gemma-4-31B-it-assistant-bf16 \
  --draft-kind mtp \
  --prompt "Summarize the key ideas in the Rust ownership model." \
  --max-tokens 400 \
  --profile

The --draft-model spelling in this example is intentional: offline mlxcel generate registers only that name. The mlxcel serve and standalone mlxcel-server entry points accept both --draft-model and --model-draft; the standalone server help presents the llama.cpp-compatible --model-draft spelling as primary.

Qwen MTP Server¶

mlxcel-server \
  --model ./models/qwen3.8-27b-4bit \
  --model-draft ./models/qwen3.8-27b-mtp \
  --draft-kind mtp \
  --draft-block-size 3 \
  --max-batch-size 1 \
  --port 8080

Omit --draft-block-size to honor a qwen3_5_mtp checkpoint's own block_size when present.

Qwen DFlash Server¶

mlxcel-server \
  --model ./models/qwen3.5-27b-4bit \
  --model-draft ./models/Qwen3.5-27B-DFlash \
  --draft-kind dflash \
  --draft-block-size 16 \
  --max-batch-size 1 \
  --port 8080

mlxcel serve accepts --draft-model; standalone mlxcel-server presents llama.cpp-compatible --model-draft. Both spellings are aliases in both server paths. For reproducible operations, pass --draft-kind explicitly even though known assistant model_type values can be auto-detected.

LFM2 DSpark Server¶

mlxcel-server \
  --model ./models/lfm2.5-2.6b-bf16 \
  --model-draft ./models/lfm2.5-2.6b-dspark-bf16 \
  --port 8080

The DSpark kind is auto-detected from the drafter's architecture and dflash_config block. Send greedy requests (temperature: 0); sampling requests are served by plain decode.

GLM-4.7-Flash MTP¶

mlxcel split-mtp -m ./models/glm-4.7-flash-bf16 -o ./models/glm-4.7-flash-mtp-4bit --q-bits 4

mlxcel generate \
  --model ./models/glm-4.7-flash-4bit \
  --draft-model ./models/glm-4.7-flash-mtp-4bit \
  --prompt "Explain in two sentences why the sky is blue." \
  --temp 0

The same drafter directory works as --model-draft on mlxcel-server.

Laguna DFlash Offline¶

mlxcel generate \
  --model ./models/laguna-xs-2.1-nvfp4 \
  --draft-model ./models/laguna-xs-2.1-dflash \
  --draft-kind dflash \
  --prompt "def retry_with_backoff(" \
  --no-chat-template \
  --temp 0

--draft-kind dflash on mlxcel generate is accepted only for a Laguna target with its Laguna DFlash drafter. Any other target reports that its DFlash pairing runs through the server.

CLI and Environment Reference¶

Draft Block Width¶

When --draft-block-size, MLXCEL_DRAFT_BLOCK_SIZE, and LLAMA_ARG_DRAFT_BLOCK_SIZE are all unset, the width is resolved once at startup in this order:

1. A width the drafter checkpoint declares for itself: the Qwen 3.5 MTP block_size, the LFM2 DSpark runtime_block_size, the Muse Glimmer assistant depth, the Inkling MTP depth, or the GLM-4.7-Flash num_nextn_predict_layers + 1.
2. A measured default for the running device and the target's quantization. The one entry is a DFlash drafter against an affine-quantized target on CUDA compute capability 12.1 (GB10), which resolves to 4; the flat 16 measured slower than plain decode there.
3. The flat per-kind default: 4 for MTP, 16 for DFlash, and 4 for a full model.

The startup speculative= line logs the resolved width together with the rule that produced it (for example block_size=4, block_size_from=measured default for sm_121 with an affine-quantized target), and the per-request DFlash diagnostics line repeats the width. The width is worker-owned, so changing it needs a restart.

B=1 Adaptive MTP Policy¶

Server-side B=1 MTP is adaptive by default. For each target/drafter/hardware/block-size combination, mlxcel profiles the first four qualifying bursts, estimates benefit from acceptance and latency, and persists a coarse enable/decline verdict under ${MLXCEL_CACHE_DIR:-$HOME/.cache/mlxcel}/mtp-policy/. It stores no prompt text or token IDs. An ambiguous measurement falls back to the static hardware policy.

Do not infer production B>1 benefit from a B=1 microbenchmark. Queueing, scheduler cooperation, target batchability, and verify shape all change.

MTP Exactness Probe¶

For unpenalized temperature-0 requests, mlxcel tries to make Qwen 3.5-family and Gemma 4 MTP output byte-identical to plain decode. A block forward of width K can nevertheless differ from K single-token forwards because MLX may select a different quantized projection kernel or reduction order.

At worker startup, the runtime compares a synthetic verify block with the equivalent single-token chain from the same state for three independent inputs. The result is memoized per (model, block width) and warmed off the request path. If it diverges, mlxcel retries with MLX's qmv_wide kernel disabled. If the retry is exact, that selection stays disabled and MTP engages; if it is still different, the runtime declines to plain decoding. The qmv_wide retry runs only where the Metal backend is available; other backends go straight to the verdict. A decline names the probe verdict and is logged once per process.

• MLXCEL_QMV_WIDE explicitly pins kernel selection and takes precedence over the gate's retry.
• MLXCEL_MTP_ALLOW_INEXACT=1 engages Qwen 3.5 or Gemma 4 MTP despite a failed probe. This is a throughput-research override, not an exactness setting.
• MLXCEL_MTP_PROBE_PROMPT_LEN sets how many prompt tokens the probe prefills before comparing (default 8). It exists to repeat the length sweep on another checkpoint; longer values cost one extra prefill per probe draw at worker startup.
• Penalties weaken the byte-identity contract: only the first bonus token receives the same penalty processing before later verification tokens continue greedily. f16 reduction jitter can also cause a pair to decline.

This measured gate is preferable to a hard-coded hardware allowlist because checkpoint dimensions, quantization, block width, GPU generation, MLX version, and kernel selection all affect equality.

Drafter Cost and Diagnostics¶

Qwen MTP drafters are quantized at load by default to reduce weight traffic. MLXCEL_MTP_QUANTIZE_DRAFTER=0 opts out. Quantization does not bypass verification; it changes proposal cost, not target authority.

These diagnostics help explain a round:

• MLXCEL_MTP_DRAFT_PROFILE=1 synchronizes individual drafter components to attribute GPU time. The extra synchronization makes it unsuitable as the number to quote for ordinary step latency; compare with an unprofiled run.
• MLXCEL_MTP_GAP_LOG=1 logs the candidate token and top-two logit gap at each verify position for exactness and branching research.
• MLXCEL_PRINT_TOKEN_IDS=1 makes offline mlxcel generate print the prompt and generated token ids to stderr, so a plain run and a --draft-model run can be compared on ids rather than decoded text.

The round log reports the effective block width, which can be smaller than the configured width near the output-token budget or cache limit.

Pairing Community Checkpoints Safely¶

Before loading a community assistant, MTP, DFlash, Medusa, or EAGLE repository, answer all of these:

1. Which exact target repository and revision was used for training?
2. Does the runtime implement this drafter architecture and model_type, rather than only the general algorithm name?
3. Are target tokenizer IDs, embedding dimensions, hidden-state tap, number of layers, and positional convention compatible?
4. Does the model card state a block width and verification/sampling contract?
5. Are draft weights already quantized, and does that quantization have runtime support?
6. Is the published speed measured end to end on a comparable target, batch, output length, and hardware?

An incompatible specialized drafter should be treated like an incompatible adapter, not like a lower-quality fallback model. Prefer a startup error or an explicit baseline fallback over silently emitting unverified proposals.

Further Reading¶

• Gloeckle et al., Better & Faster Large Language Models via Multi-token Prediction
• Google AI for Developers, Gemma Multi-Token Prediction overview
• Hugging Face Transformers, Gemma 4 assistant model
• Qwen, Qwen3.5 model card and MTP deployment examples
• Z Lab, DFlash project and DFlash paper
• DFlash reference implementation