Skip to content

feat(mtp): add native MTP conditioning for DSpark and DFlash #23

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhouaihui merged 2 commits into from
Nov 12, 2026
Merged

feat(mtp): add native MTP conditioning for DSpark and DFlash #23

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0551322
feat(mtp): condition DSpark and DFlash on native MTP features
zhouaihui Nov 11, 2026
6ce63d1
fix(mtp): clarify suffix metrics and validate fusion
zhouaihui Nov 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .github/workflows/ci.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,9 @@ jobs:
tests/test_eagle4_validation.py \
tests/test_fa4_import_compat.py \
tests/test_flex_attention_backend.py \
tests/test_dmtp.py \
tests/test_training_resume_rng.py \
tests/test_evaluation_contract.py \
tests/test_open_source_entrypoints.py \
tests/test_release_audit.py \
tests/test_shipped_config_defaults.py \
Expand Down
17 changes: 17 additions & 0 deletions aspec/commands/_recipe_smoke.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,19 @@
def audit_state(trainer):
result = _base_audit_state(trainer)
draft = getattr(trainer, "draft_model", None)
if getattr(getattr(draft, "config", None), "mtp_conditioning", None):
from aspec.commands._checkpoint_audit import tensor_hash

result["mtp_conditioning"] = {
"policy": draft.config.mtp_conditioning,
"frozen": tensor_hash((n, p) for n, p in draft.named_parameters()
if n.startswith(("mtp_source.", "embed_tokens."))),
"fusion": tensor_hash(draft.mtp_fusion.named_parameters())
if draft.mtp_fusion is not None else None,
}
expected = [p for p in draft.parameters() if p.requires_grad]
if [id(p) for p in trainer.optimizer.model_params] != [id(p) for p in expected]:
raise AssertionError("MTP conditioning optimizer ownership differs from the Draft")
if getattr(getattr(draft, "config", None), "ffn_type", "dense") == "ling_moe":
result["ling_expert_bias"] = [layer.mlp.gate.expert_bias.detach().cpu().tolist()
for layer in draft.layers]
Expand Down Expand Up @@ -117,6 +130,10 @@ def train_from_queue(self, step, num_batches):
def save_model(self, step, force_sync=False):
super().save_model(step, force_sync)
current = audit_state(self._trainer)
if current.get("mtp_conditioning"):
initial = json.loads((self._audit_root / f"recipe-audit-initial{self._audit_suffix}.json").read_text())
if current["mtp_conditioning"]["frozen"] != initial["mtp_conditioning"]["frozen"]:
raise AssertionError("MTP conditioning changed its frozen source/embedding")
if current.get("eagle4", {}).get("policy", {}).get("stage") == "candidate_head":
initial = json.loads((self._audit_root / f"recipe-audit-initial{self._audit_suffix}.json").read_text())
if initial["eagle4"]["body"] != current["eagle4"]["body"]:
Expand Down
11 changes: 9 additions & 2 deletions aspec/commands/convert_to_hf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -163,6 +163,9 @@ def _serving_config(raw_config: dict) -> dict:
from aspec.config.draft_attention import resolve_draft_attention
from aspec.config.draft_structure import resolve_draft_structure

if raw_config.get("mtp_conditioning"):
raise NotImplementedError("MTP-conditioned Drafts require native --format aspec export")

if raw_config.get("architectures") == ["NativeMTPDraftModel"]:
native = json.loads(json.dumps(raw_config["native_config"]))
native["architectures"] = ["Qwen3_5ForCausalLMMTP"]
Expand Down Expand Up @@ -461,6 +464,10 @@ def convert_checkpoint(
}
if eagle4_asset is not None:
manifest["eagle4"] = eagle4_asset
if raw_config.get("mtp_conditioning"):
from aspec.config.mtp_conditioning import mtp_conditioning_semantics

manifest["mtp_conditioning"] = mtp_conditioning_semantics(raw_config)
if artifact_format == "aspec":
manifest["compute_dtype"] = str(model.embed_tokens.weight.dtype)
manifest["tensor_dtypes"] = dict(Counter(str(t.dtype) for t in tensors.values()))
Expand All @@ -473,7 +480,7 @@ def convert_checkpoint(
if policy.get("dsa") is not None:
manifest["draft_structure"]["attention"]["dsa"] = policy["dsa"]
algorithm = training_contract["algorithm"]
block_proposal = algorithm in {"dspark", "dflash"}
block_proposal = algorithm in {"dspark", "dflash"} and not raw_config.get("mtp_conditioning")
mla_proposal = raw_config.get("architectures") == ["Eagle3DeepseekV2ForCausalLM"]
proposal = ("aspec.draft_artifact.MLAEagleConsumer" if mla_proposal else
"aspec.draft_artifact.BlockDraftConsumer" if block_proposal else None)
Expand All @@ -483,7 +490,7 @@ def convert_checkpoint(
"validation": "required_per_artifact",
"scope": ("asset and proposal; caller owns verifier features/cache/verification"
if proposal else "asset loading and weights-only warmstart; no native proposal component"),
"borrowed_tensors": ["Target output head"] if block_proposal else [],
"borrowed_tensors": ["Target output head"] if algorithm in {"dspark", "dflash"} else [],
}
(output_dir / "export_manifest.json").write_text(
json.dumps(manifest, indent=2, sort_keys=True) + "\n",
Expand Down
106 changes: 106 additions & 0 deletions aspec/config/dflash_qwen3_5_9b_dmtp_config.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
{
"architectures": [
"DFlashDraftModel"
],
"model_type": "dflash",
"hidden_size": 4096,
"intermediate_size": 12288,
"head_dim": 128,
"num_hidden_layers": 5,
"num_attention_heads": 32,
"num_key_value_heads": 8,
"vocab_size": 248320,
"rms_norm_eps": 1e-06,
"max_position_embeddings": 262144,
"rope_theta": 10000000.0,
"num_target_layers": 5,
"target_hidden_size": 4096,
"target_num_hidden_layers": 32,
"target_layer_ids": [
1,
8,
15,
22,
29
],
"mask_token_id": 248070,
"tie_word_embeddings": false,
"mtp_conditioning": {
"mode": "hidden",
"native_config": {
"attention_bias": false,
"attention_dropout": 0,
"attn_output_gate": true,
"dtype": "bfloat16",
"eos_token_id": 248044,
"full_attention_interval": 4,
"head_dim": 256,
"hidden_act": "silu",
"hidden_size": 4096,
"initializer_range": 0.02,
"intermediate_size": 12288,
"layer_types": [
"linear_attention",
"linear_attention",
"linear_attention",
"full_attention",
"linear_attention",
"linear_attention",
"linear_attention",
"full_attention",
"linear_attention",
"linear_attention",
"linear_attention",
"full_attention",
"linear_attention",
"linear_attention",
"linear_attention",
"full_attention",
"linear_attention",
"linear_attention",
"linear_attention",
"full_attention",
"linear_attention",
"linear_attention",
"linear_attention",
"full_attention",
"linear_attention",
"linear_attention",
"linear_attention",
"full_attention",
"linear_attention",
"linear_attention",
"linear_attention",
"full_attention"
],
"linear_conv_kernel_dim": 4,
"linear_key_head_dim": 128,
"linear_num_key_heads": 16,
"linear_num_value_heads": 32,
"linear_value_head_dim": 128,
"max_position_embeddings": 262144,
"mlp_only_layers": [],
"model_type": "qwen3_5_text",
"mtp_num_hidden_layers": 1,
"mtp_use_dedicated_embeddings": false,
"num_attention_heads": 16,
"num_hidden_layers": 32,
"num_key_value_heads": 4,
"rms_norm_eps": 1e-06,
"use_cache": true,
"vocab_size": 248320,
"mamba_ssm_dtype": "float32",
"rope_parameters": {
"mrope_interleaved": true,
"mrope_section": [
11,
11,
10
],
"rope_type": "default",
"rope_theta": 10000000,
"partial_rotary_factor": 0.25
}
}
}
}
7 changes: 6 additions & 1 deletion aspec/config/draft_algorithm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,4 +73,9 @@ def draft_model_semantics(config) -> dict:
"enable_post_mlp_norm",
"input_embedding_scale", "output_multiplier", "final_logit_softcapping", "is_causal",
)
return {key: raw[key] for key in fields if key in raw}
semantics = {key: raw[key] for key in fields if key in raw}
if raw.get("mtp_conditioning") is not None:
from aspec.config.mtp_conditioning import resolve_mtp_conditioning

semantics["mtp_conditioning"] = resolve_mtp_conditioning(raw)
return semantics
110 changes: 110 additions & 0 deletions aspec/config/dspark_qwen3_5_9b_dmtp_config.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,110 @@
{
"architectures": [
"DSparkDraftModel"
],
"model_type": "dspark",
"hidden_size": 4096,
"intermediate_size": 12288,
"head_dim": 128,
"num_hidden_layers": 5,
"num_attention_heads": 32,
"num_key_value_heads": 8,
"vocab_size": 248320,
"rms_norm_eps": 1e-06,
"max_position_embeddings": 262144,
"rope_theta": 10000000.0,
"num_target_layers": 5,
"target_hidden_size": 4096,
"target_num_hidden_layers": 32,
"target_layer_ids": [
1,
8,
15,
22,
29
],
"mask_token_id": 248070,
"markov_rank": 256,
"markov_head_type": "vanilla",
"enable_confidence_head": true,
"confidence_head_with_markov": true,
"tie_word_embeddings": false,
"mtp_conditioning": {
"mode": "hidden",
"native_config": {
"attention_bias": false,
"attention_dropout": 0,
"attn_output_gate": true,
"dtype": "bfloat16",
"eos_token_id": 248044,
"full_attention_interval": 4,
"head_dim": 256,
"hidden_act": "silu",
"hidden_size": 4096,
"initializer_range": 0.02,
"intermediate_size": 12288,
"layer_types": [
"linear_attention",
"linear_attention",
"linear_attention",
"full_attention",
"linear_attention",
"linear_attention",
"linear_attention",
"full_attention",
"linear_attention",
"linear_attention",
"linear_attention",
"full_attention",
"linear_attention",
"linear_attention",
"linear_attention",
"full_attention",
"linear_attention",
"linear_attention",
"linear_attention",
"full_attention",
"linear_attention",
"linear_attention",
"linear_attention",
"full_attention",
"linear_attention",
"linear_attention",
"linear_attention",
"full_attention",
"linear_attention",
"linear_attention",
"linear_attention",
"full_attention"
],
"linear_conv_kernel_dim": 4,
"linear_key_head_dim": 128,
"linear_num_key_heads": 16,
"linear_num_value_heads": 32,
"linear_value_head_dim": 128,
"max_position_embeddings": 262144,
"mlp_only_layers": [],
"model_type": "qwen3_5_text",
"mtp_num_hidden_layers": 1,
"mtp_use_dedicated_embeddings": false,
"num_attention_heads": 16,
"num_hidden_layers": 32,
"num_key_value_heads": 4,
"rms_norm_eps": 1e-06,
"use_cache": true,
"vocab_size": 248320,
"mamba_ssm_dtype": "float32",
"rope_parameters": {
"mrope_interleaved": true,
"mrope_section": [
11,
11,
10
],
"rope_type": "default",
"rope_theta": 10000000,
"partial_rotary_factor": 0.25
}
}
}
}
Loading
Loading