feat(gguf): auto-detect + export Qwen3.5/3.8 MTP self-speculative head from GGUF - #529
Open
justinchuby wants to merge 5 commits into
Open
feat(gguf): auto-detect + export Qwen3.5/3.8 MTP self-speculative head from GGUF#529justinchuby wants to merge 5 commits into
justinchuby wants to merge 5 commits into
Conversation
The dense Qwen3.5/3.8 GGUF (e.g. unsloth/Qwen3.8-27B) ships a trained
multi-token-prediction ("nextn") head as the trailing blk.<N> block,
tagged by <arch>.nextn_predict_layers. gguf_to_config read that key only
to subtract the block from num_hidden_layers, so the text-only backbone
build silently dropped the head and the exported ONNX had no
self-speculative drafter.
Recover the head as a mtp/ sidecar:
- _config_mapping: surface _gguf_nextn_predict_layers / _gguf_mtp_block_indices.
- _mtp.py (new): map blk.<N>.nextn.* + attn/ffn to Qwen35MtpModel stems,
derive Qwen35MtpConfig (all dims inherited from the backbone; no hardcoding),
and build the head via the existing quantized/dequantized loaders through an
injectable name_mapper. The two pre_fc_norm_* OffsetRMSNorm +1 offsets missed
by the generic pass are stripped explicitly.
- _builder: expose hidden_states.<N-1> as the head seed and attach pkg.mtp_head.
- __main__: save pkg.mtp_head to output_dir/mtp/ and write speculator metadata.
- inference_metadata: write_mtp_speculator_metadata (proposal_type: mtp).
- _mtp_test.py (new): mapping, config derivation, and tiny-GGUF head emission.
Verified the real Qwen3.8-27B GGUF blk.64 head tensor shapes match the
metadata-derived head (doubled-Q gating: attn_q = 2*NH*HD rows; eh_proj = H x 2H).
Unblocks MTP self-speculative decode for Qwen3.8-27B (engine side: onnx-genai #1598).
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| @@ -0,0 +1,217 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
Performance Comparison
|
The Qwen3.5/3.8 MTP self-speculative head is now gated behind an explicit opt-in flag instead of being emitted whenever the GGUF ships a nextn block. Default behavior is unchanged: a text-only export with no head sidecar, so existing exports are byte-for-byte identical. - build_from_gguf gains include_mtp: bool = False. Both the backbone hidden-state seed (output_layer_indices=[N-1], step 4b) and the head sidecar build (step 10) are gated on include_mtp AND has_mtp_head. When include_mtp=True but the GGUF has no nextn head, a warning is logged and the build proceeds text-only; when a head is present but the flag is off, an info log points at --include-mtp. Rejects include_mtp with static_cache or mmproj (ValueError). - CLI: add `--include-mtp` (store_true, off by default), threaded through _cmd_build_gguf; the mtp/ sidecar save + SpeculatorConfig metadata write already key off pkg.mtp_head, so they only fire when the flag is set. CLI-level SystemExit guards mirror the build_from_gguf conflicts. - Tests: gating coverage (default omits head; --include-mtp attaches it with mtp_hidden output; static_cache conflict raises; no-head GGUF warns). MTP dims are still derived entirely from GGUF/source metadata; the flag only gates whether the subgraph is emitted. One driver now serves both the text-only and MTP paths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
justinchuby
changed the title
Supersedes the opt-in flag: the MTP self-speculative head is a purely additive sidecar (mtp/model.onnx + shared embedding/lm_head + SpeculatorConfig metadata) that existing consumers ignore, so there is no reason to gate it behind a user flag. Emit it iff the source GGUF actually ships the head. - Presence check (has_mtp_head): the config carries _gguf_mtp_block_indices, derived in gguf_to_config from <arch>.nextn_predict_layers > 0 (the trailing blk.<N>.nextn.* head block). When present, always emit the sidecar + the backbone hidden_states.<N-1> seed + SpeculatorConfig. When absent, skip — text-only sources produce byte-identical output to before. - Remove build_from_gguf's include_mtp parameter and the CLI --include-mtp flag plus its conflict guards; steps 4b/10 gate on has_mtp_head alone. static_cache still skips the head (its dynamic-cache requirement) via a logged info, keeping static exports unchanged. - Tests: replace flag-gating with auto-detect coverage (nextn GGUF auto-emits the head with mtp_hidden output; plain GGUF omits it). Dims remain fully derived from GGUF/source metadata (no hardcoding). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
changed the title
…etadata across dtype replace
Two defects prevented the real int4/bf16 Qwen3.8-27B export from emitting a
usable MTP self-speculative sidecar:
1. Qwen35TextModel.forward ignored output_layer_indices, so the backbone never
produced the hidden_states.{N-1} seed the head consumes (it always returned a
2-tuple). Add the same capture path as base.py TextModel: return the selected
per-layer post-residual hidden states so HybridCausalLMTask emits
hidden_states.{idx}. Without this the speculator.target_hidden_output
referenced a non-existent output.
2. build_from_gguf lost the private _gguf_* MTP metadata whenever an explicit
--dtype (or quantization) triggered dataclasses.replace, which returns a
fresh config without the plain attributes gguf_to_config set. Auto-detection
(has_mtp_head) then saw nothing and silently skipped the head. Capture the
nextn metadata into locals right after gguf_to_config (like _gguf_model_type)
and re-attach it onto the final config before step 4b.
Also extend write_mtp_speculator_metadata with embedding_weights,
lm_head_weights, and vocab_size so the SpeculatorConfig block names the shared
tables and vocab the runtime needs.
Tests: regression for the dtype-replace path and an assertion that the backbone
exposes hidden_states.{N-1} when the source ships the nextn head.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…latorConfig schema
The emitted `speculator:` block used field names the onnx-genai runtime cannot
parse, so InferenceMetadata.speculative deserialized to None and the MTP sidecar
was silently ignored (text-only decode). Conform exactly to the authoritative
schema in onnx-genai crates/onnx-genai-metadata (schema/generation.rs
SpeculatorConfig, parser.rs resolve_mtp, config.rs validate_resolved_mtp_config):
- Publish under the top-level `speculative` key (was `speculator`, an unknown
key that is dropped).
- `model_path` -> `model`; `hidden_size` -> `target_hidden_size`.
- `kv_mode: hidden_threaded` (engine-internal enum) -> `proposal_local`
(the only valid MtpKvMode for this k=1 head).
- `embedding_weights`/`lm_head_weights` flat strings -> `embedding`/`lm_head`
MtpTargetInitializer objects {source: target_initializer, name: ...}.
- Add `target_hidden_layout: BSH` (rank-3 [batch,seq,hidden] seed) and
`hc_mult: 1` (resolve_mtp requires hc_mult>0; validate pins it to 1 for BSH).
- Omit `mtp_state_output` (hc_mult==1 head has no recurrent HC state; the
sidecar emits only mtp_hidden + present.0.{key,value}).
Test: assert the emitted YAML has the exact schema keys/values and validates
against onnx-genai's published inference_metadata.schema.json.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Recovers the Multi-Token-Prediction (MTP / "nextn") self-speculative head that the Qwen3.8-27B ONNX export was silently dropping, and emits it automatically whenever the source GGUF ships the head — no CLI flag, no config toggle. The MTP sidecar is a purely-additive artifact (
mtp/model.onnx+ shared embedding/lm_head +SpeculatorConfigmetadata) that existing consumers ignore, so it is driven entirely by source-tensor presence: emit MTP iff available.Auto-detect (presence-driven, no flag)
The decision is made by
has_mtp_head(config)— true when the config carries_gguf_mtp_block_indices, whichgguf_to_configpopulates from the source GGUF. Exact presence check:<arch>.nextn_predict_layers > 0(for Qwen3.8-27B:qwen35.nextn_predict_layers = 1,block_count = 65), which yields_gguf_mtp_block_indices = [decoder_layers .. decoder_layers + nextn - 1](here[64]), andblk.<N>.nextn.*tensors —nextn.eh_proj(thefc),nextn.enorm,nextn.hnorm,nextn.shared_head_norm— plus that block's own attention (attn_q/k/v/output,attn_q_norm/k_norm,attn_norm,post_attention_norm) and FFN (ffn_gate/up/down) sublayer.Behavior:
mtp/model.onnxsidecar + the backbonehidden_states.<N-1>seed output + aSpeculatorConfig(proposal_type: mtp) block ininference_metadata.yaml.--static-cacheskips the head (it needs the dynamic concat-grow cache), logged as info; those exports are unchanged.MTP dims are derived entirely from GGUF/source metadata (hidden size, heads, head_dim, KV heads, rope, quant) — no hardcoding.
Root cause
The dense Qwen3.5/3.8 GGUF (e.g.
unsloth/Qwen3.8-27B) ships a trained MTP head as the trailingblk.<N>block (for 27B:blk.64), tagged by<arch>.nextn_predict_layers.gguf_to_configread that key only to subtract the block fromnum_hidden_layers, and the text-only backbone build never made a second pass — so the head tensors were dropped.Changes
integrations/gguf/_config_mapping.py— surface_gguf_nextn_predict_layers/_gguf_mtp_block_indices(the trailing block previously only subtracted).integrations/gguf/_mtp.py(new) —has_mtp_head(presence check),map_gguf_mtp_to_hf_names(GGUFblk.<N>.nextn.*+ attn/ffn ->Qwen35MtpModelstems),derive_mtp_config(all dims inherited from the backbone config — no hardcoding),build_mtp_head_from_gguf(reusesQwen35MtpModel/Qwen35MtpTaskand the existing quantized/dequantized loaders via an injectablename_mapper;fcis quantized on the same main-layer path). The twopre_fc_norm_*OffsetRMSNorm+1offsets missed by the generic norm strip are handled explicitly.integrations/gguf/_builder.py— steps 4b (seed output) and 10 (attachpkg.mtp_head) gate onhas_mtp_head(config)alone (auto-detect);static_cacheskips the head via a logged info. The quantized/dequantized loaders gained optionalname_mapper/warn_unmappedparams (backward compatible).__main__.py— whenpkg.mtp_headis present, save it intooutput_dir/mtp/and write speculator metadata. (No new CLI flag.)integrations/onnx_genai/inference_metadata.py—write_mtp_speculator_metadata(proposal_type: mtp,target_hidden_output,mtp_hidden, single-layer KV).integrations/gguf/_mtp_test.py(new) — mapping, config derivation, tiny-GGUF head emission, and auto-detect (nextn GGUF auto-emits the head withmtp_hiddenoutput; plain GGUF omits it).Validation (partial — stated explicitly)
build_from_gguf(source with nextn ->pkg.mtp_headattached; source without -> absent). Emitted head initializers includefc.weight_t,pre_fc_norm_embedding.weight,pre_fc_norm_hidden.weight,norm.weight,layers.0.self_attn.q_proj.weight_t; graph outputmtp_hidden; every initializer has backing weights._config_mapping/_tensor_mappingsuites still pass (74 total).blk.64shape contract VERIFIED against the metadata-derived head — all 15 head tensors map, backbone excluded:attn_q= (12288, 5120) = (2NHHD, H) -> doubled-Q output gatingattn_k/attn_v= (1024, 5120) = (NKVHD, H);attn_output= (5120, 6144) = (H, NHHD)nextn.eh_proj= (5120, 10240) = (H, 2H) — thefcGEMM (Q8_0 while projections are Q4_0; mixed quant handled per-tensor)onnx_irlacksmax_shard_size_bytes, which also breaks the pre-existingtest_default_quantized_package_save_reload— unrelated to this change).Real artifact produced + two export-path fixes
Producing the real int4/bf16 27B artifact surfaced two defects that made the MTP sidecar silently degenerate. Both are now fixed (commit
987c8d04):Qwen35TextModel.forwardignoredoutput_layer_indices— it always returned a 2-tuple, so the backbone never emitted thehidden_states.{N-1}seed the head consumes;speculator.target_hidden_outputpointed at a non-existent output. Added the same capture path asbase.pyTextModel(returns the selected per-layer post-residual hidden states soHybridCausalLMTaskemitshidden_states.{idx}).build_from_ggufdropped the_gguf_*nextn metadata under--dtype— an explicit dtype (or quantization) triggersdataclasses.replace, which returns a fresh config without the plain_gguf_*attributes, sohas_mtp_headsaw nothing and skipped the head. Now captured into locals right aftergguf_to_config(like_gguf_model_type) and re-attached before step 4b. Regression test added (test_mtp_survives_dtype_replace).write_mtp_speculator_metadataalso now emitsembedding_weights,lm_head_weights, andvocab_size.Full export — VALIDATED end-to-end
Command:
model.onnx(2.2 MB) +model.onnx.data(17.4 GB),mtp/model.onnx(38 KB) +mtp/model.onnx.data(354 MB),inference_metadata.yaml,tokenizer.json. Total 17 GB.hidden_states.63=[batch, sequence_len, 5120]— exactlyspeculator.target_hidden_output.inputs_embeds,hidden_states,attention_mask,position_ids,past_key_values.0.{key,value}; outputsmtp_hidden+present.0.{key,value}(single full-attention layer KV).fc.weight_t= (10240, 5120); q/k/v/o + gate/up/down projections are int4MatMulNBits(UINT8 packed, block-32) following the main-layer path;eh_projstays bf16 (Q8_0 in source). Shared embedding/lm_head are reused from the main model (named in the speculator block), not duplicated.proposal_type: mtp,num_speculative_tokens: 1,model_path: mtp/model.onnx,mtp_hidden_output: mtp_hidden,kv_mode: hidden_threaded,embedding_weights: model.embed_tokens.weight,lm_head_weights: lm_head.weight,target_hidden_output: hidden_states.63,hidden_size: 5120,vocab_size: 248320.Runtime schema-conformance fix (metadata block)
The first emitted metadata used a
speculator:block whose field names the onnx-genai runtime cannot parse, soInferenceMetadata.speculativedeserialized toNoneand the sidecar was silently ignored (text-only decode). Now conformed exactly to the authoritative schema (onnx-genaicrates/onnx-genai-metadata:schema/generation.rsSpeculatorConfig,parser.rsresolve_mtp,config.rsvalidate_resolved_mtp_config) — commitb6747e63:speculator->speculative(the runtime key; aliasspeculator_config).speculatorwas unknown -> dropped.model_path->model;hidden_size->target_hidden_size.kv_mode: hidden_threaded(engine-internal enum) ->proposal_local(only validMtpKvModefor this k=1 head).embedding_weights/lm_head_weightsflat strings -> nestedembedding/lm_headMtpTargetInitializerobjects{source: target_initializer, name: ...}.target_hidden_layout: BSH(rank-3[batch,seq,hidden]seed) +hc_mult: 1(resolve_mtprequireshc_mult>0;validate_resolved_mtp_configpins it to 1 for BSH).mtp_state_output(anhc_mult==1head has no recurrent Hyper-Connection state; the sidecar emits onlymtp_hidden+present.0.{key,value}).Emitted block:
Test
TestMtpSpeculatorMetadataasserts the exact keys/values and validates the emitted YAML against onnx-genai's publishedinference_metadata.schema.json. (The existing 17 GB artifact's metadata was hand-patched for the benchmark; not re-exported.)Why
Unblocks MTP self-speculative decode for Qwen3.8-27B batch=1 decode on H200. The engine side already landed in onnx-genai #1598 (recurrent-state commit-by-accepted-prefix);
mtp.rs/MtpDecodeSession+SpeculatorConfigconsume exactly this sidecar. Any Qwen3.5/3.8 GGUF that ships the head now exports it end-to-end automatically.Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com