MTP (mtp_layers.0) forward semantics for self-speculative decoding?
Hi Motif team — following up after the Beta-era PolyNorm exchange (thanks again for the fixed
modeling push). I am porting Motif-3 (final) to MLX (avlp12/Motif-3-Alis-MLX-8bit,
single M3 Ultra, 20 tok/s) and are now trying to use the shipped MTP block
(model.mtp_layers.0.*) for self-speculative decoding.
Since modeling_motif.py never instantiates the MTP block, I reconstructed the forward from
tensor names/shapes (DeepSeek-V3-style): x = input_proj([h_backbone ; embed_norm(embed(t))])
→ plain residual block (GDLA + PolyNorm MLP, input_layernorm/post_attention_layernorm as
pre-norms) → final_layernorm → shared lm_head. The best draft acceptance I reach is ~33%
(k=1, greedy), which suggests one structural detail is still off (a correctly-wired trained
head usually lands 50%+).
Could you share the intended MTP forward semantics? Specifically:
- concat order and which side
embed_norm/input_layernormnormalize (I measured[h ; embed_norm(e)]≫[embed_norm(e) ; h], andinput_layernorm-as-hnorm hurts); - which backbone hidden feeds it — mean over the mHC expansion streams, pre- or
post-model.norm, at position t paired with token t+1? - the MTP attention config (full vs sliding, YaRN vs plain theta) and whether the chained
multi-step drafting feeds backfinal_layernorm(h)or the raw hidden.
A short pseudo-code snippet would be enough. I would be happy to upstream the working integration and
publish acceptance/speed numbers on the card, as with the Beta fixes.
Answering my own question for anyone following: I found the full wiring in yourmotif3-training-example repository (compute_mtp_loss / _run_mtp_block). Key details:
the anchor hidden is the post-final-norm backbone hidden, the concat is[h ; mtp_embed_norm(embed(t+1))], the MTP block is a plain (non-mHC) transformer block that
runs as a sliding-window layer (index n_layers → SWA), and the chained/hidden norm is
applied after the block before the shared lm_head. With this wiring my MLX port reaches
38–41% draft acceptance (greedy, k=1), i.e. +15–20% end-to-end decode. One remaining
question: what acceptance rate do you observe internally? I'd like to know whether ~40% is
the expected ceiling for the 1-layer head or whether I'm still leaving something on the table.
@avlp12
Thank you very much for your continued interest and contributions. Your support has been incredibly helpful to us.
Since we have made our vLLM code publicly available, it would be more helpful to explain this along with the relevant code path.
Concat order & norms
Hidden first, normed embedding second. There is no hnorm on the backbone hidden . (https://github.com/MotifTechnologies/vllm/blob/4cd9eb4129883565e69d508038d783d59ee01867/vllm/model_executor/models/motif_mtp.py#L160)Which backbone hidden feeds it
Mean over streams, post-model.norm (https://github.com/MotifTechnologies/vllm/blob/4cd9eb4129883565e69d508038d783d59ee01867/vllm/model_executor/models/motif.py#L1563)
for the pairing : position t pairs with token t+1, and the MTP block's RoPE position is tMTP attention config
All MTP layers are SWA, plain RoPE at swa_rope_theta
(https://github.com/MotifTechnologies/vllm/blob/4cd9eb4129883565e69d508038d783d59ee01867/vllm/model_executor/models/motif_mtp.py#L79)chained multi-step
chained drafting feeds back final_layernorm(h)
(https://github.com/MotifTechnologies/vllm/blob/4cd9eb4129883565e69d508038d783d59ee01867/vllm/model_executor/models/motif_mtp.py#L164)
We hope this answers your question. Please feel free to reach out if you have any further questions. Thank you for your interest and support again
@avlp12
One remaining question:what acceptance rate do you observe internally?
To answer your question, we typically observe an acceptance rate of around 70–80%. Even with the NVFP4 checkpoint(https://huggingface.co/Motif-Technologies/Motif-3-NVFP4), we haven't observed an acceptance rate below 70%.
Thank you @TaehyunKimMotif — that answers it completely, and the vLLM fork links were exactly what I needed.
Closing the loop for anyone following: the gap turned out to be two things, neither of them wiring.
Metric definition on my side. My "38–41%" was the fraction of emitted tokens that came from the draft (a/(1+a) for k=1), not per-draft acceptance. Converting: my true greedy acceptance was ~65–70% all along — consistent with your 70–80% band. (Sanity check: my measured 1.21× end-to-end is arithmetically impossible at 41% true acceptance.)
Acceptance rule at serving temperature. I was using strict sampler-equality; your stack uses rejection sampling. I've now implemented Leviathan-style rejection acceptance (accept draft x with prob min(1, p(x)/q(x)), resample from (p−q)+ on reject) in the MLX port: at T=0.8 acceptance goes 52% (equality) → 85% (rejection), +21% end-to-end, distribution-lossless.
Current numbers on a single M3 Ultra 512GB: 8-bit build 24 tok/s, 4.5-bit build 27–31 tok/s with the MTP head (k=1, matching your "1 speculative token is optimal" guidance). I'll update my model cards accordingly. Thanks again for the responsive support — the training-example and vLLM repos made this port possible.
@avlp12
This is truly an impressive result. Once again, we sincerely appreciate your continued interest and contributions.