feat: add mtp support

configs/qwen3.5-4b-mtp.json

@@ -0,0 +1,37 @@
+{
+  "architectures": [
+    "Qwen3_5MTPDraftModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "mtp.Qwen3_5MTPDraftModel"
+  },
+  "attn_output_gate": true,
+  "bos_token_id": 248043,
+  "eos_token_id": 248044,
+  "head_dim": 256,
+  "hidden_act": "silu",
+  "hidden_size": 2560,
+  "initializer_range": 0.02,
+  "intermediate_size": 9216,
+  "max_position_embeddings": 262144,
+  "model_type": "qwen3",
+  "mtp_config": {
+    "share_lm_head": true
+  },
+  "num_attention_heads": 16,
+  "num_hidden_layers": 1,
+  "num_key_value_heads": 4,
+  "pad_token_id": 248044,
+  "partial_rotary_factor": 0.25,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 10000000,
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 248320
+}

docs/basic_usage/training.md

@@ -234,6 +234,7 @@ The checked-in examples are the canonical starting points:
 | EAGLE3 offline disaggregated | [`qwen3-8b-eagle3-offline-disaggregated.yaml`](../../examples/configs/qwen3-8b-eagle3-offline-disaggregated.yaml) |
 | Ascend NPU DFlash online | [`qwen3.5-4b-dflash-online-npu.yaml`](../../examples/configs/qwen3.5-4b-dflash-online-npu.yaml) |
 | Ascend NPU Domino online | [`qwen3.5-4b-domino-online-npu.yaml`](../../examples/configs/qwen3.5-4b-domino-online-npu.yaml) |
+| Ascend NPU MTP disaggregated | [`qwen3.5-4b-mtp-disaggregated-npu.yaml`](../../examples/configs/qwen3.5-4b-mtp-disaggregated-npu.yaml) |
 
 ## Online and offline data
 
@@ -262,6 +263,7 @@ The unified runtime supports text training in these combinations:
 | DFlash | Yes, consumer DP | Yes, DP | Yes, consumer DP |
 | Domino | Yes, consumer DP | Yes, DP | Yes, consumer DP |
 | DSpark | Yes, consumer DP | Yes, DP | Yes, consumer DP |
+| MTP | Yes, consumer DP | Yes, DP | Yes, consumer DP |
 | P-EAGLE | Yes, consumer DP, batch size 1 | No | No |
 
 Unsupported combinations fail explicitly during config validation or run
@@ -274,10 +276,10 @@ assembly. In particular:
 - attention backends are strategy-specific: EAGLE3 accepts `sdpa`,
   `flex_attention`, `fa`, or offline `usp`; P-EAGLE requires
   `flex_attention`; DFlash, Domino, and DSpark accept `eager`, `sdpa`, or
-  `flex_attention`;
+  `flex_attention`; MTP accepts `eager` or `sdpa`;
 - P-EAGLE requires `training.batch_size=1` and reuses EAGLE3's server capture
   schema;
-- offline feature training supports EAGLE3, DFlash, Domino, and DSpark;
+- offline feature training supports EAGLE3, DFlash, Domino, DSpark, and MTP;
 - every online run is disaggregated and uses `model.target_backend=sglang`;
   finite runs may omit both step fields so the producer can publish the exact
   optimizer horizon derived from the prepared prompt plan;
@@ -509,6 +511,23 @@ specforge export --to hf \
 Pass `--vocab-mapping /path/to/mapping.pt` when the checkpoint predates the
 mapping buffers or when you intentionally need to refresh them.
 
+MTP is deployed by merging its trained native head back into the target model.
+The merge command accepts the same runtime checkpoint shapes as the generic
+exporter (`training_state.pt`, a step/latest directory, or the run output
+directory):
+
+```bash
+python scripts/merge_mtp_to_base.py \
+  --base-model-path Qwen/Qwen3.5-4B \
+  --mtp-checkpoint-path ./outputs/qwen3.5-4b-mtp/qwen3.5-4b-mtp-latest \
+  --draft-config configs/qwen3.5-4b-mtp.json \
+  --output-path ./exports/Qwen3.5-4B-MTP
+```
+
+An already-exported HF MTP draft directory can also be passed as
+`--mtp-checkpoint-path`; in that case its own `config.json` is used and
+`--draft-config` may be omitted.
+
 ## Troubleshooting
 
 ### Late OOM or non-finite hidden states on online runs

examples/README.md

@@ -38,6 +38,7 @@ NPU, offline, and managed/external-service variants, is in
 | `examples/configs/qwen3.5-4b-dflash-online-npu.yaml` | Disaggregated NPU SGLang capture | DFlash |
 | `examples/configs/qwen3.5-4b-dflash-disaggregated-npu.yaml` | Managed local NPU full stack | DFlash |
 | `examples/configs/qwen3.5-4b-domino-online-npu.yaml` | Disaggregated NPU SGLang capture | Domino |
+| `examples/configs/qwen3.5-4b-mtp-disaggregated-npu.yaml` | Managed local NPU full stack | MTP |
 
 Online configs point `data.train_data_path` at raw conversation data. Offline
 configs expect strategy-specific feature checkpoints in

examples/configs/qwen3.5-4b-mtp-disaggregated-npu.yaml

@@ -0,0 +1,77 @@
+# Single-node managed full stack for Ascend NPU, launched with one
+# `specforge train` command. Device ordinals are injected via
+# ASCEND_RT_VISIBLE_DEVICES.
+model:
+  target_model_path: "Qwen/Qwen3.5-4B"
+  draft_model_config: "configs/qwen3.5-4b-mtp.json"
+  target_backend: sglang
+  trust_remote_code: true
+  # Qwen3.5-4B nests its text decoder under model.language_model and ties
+  # lm_head to the embedding, so both keys point at the same tensor.
+  embedding_key: "model.language_model.embed_tokens.weight"
+  lm_head_key: "model.language_model.embed_tokens.weight"
+  torch_dtype: "bfloat16"
+  sglang_attention_backend: ascend
+data:
+  train_data_path: "./cache/dataset/train_regen.jsonl"
+  max_length: 32768
+  chat_template: "qwen3.5"
+  build_dataset_num_proc: 32
+  cache_dir: "./cache"
+training:
+  strategy: "mtp"
+  num_epochs: 10
+  max_steps: 10000
+  batch_size: 2
+  accumulation_steps: 4
+  learning_rate: 0.0006
+  warmup_ratio: 0.04
+  max_grad_norm: 1
+  attention_backend: "sdpa"
+  save_interval: 10000
+  log_interval: 50
+  dist_timeout: 30
+  seed: 42
+tracking:
+  report_to: "tensorboard"
+run_id: "qwen3.5-4b-mtp-npu-managed"
+output_dir: "./outputs/qwen3.5-4b-mtp-npu-managed"
+
+deployment:
+  mode: disaggregated
+  trainer:
+    nnodes: 1
+    nproc_per_node: 10
+  disaggregated:
+    control_dir: outputs/qwen3.5-4b-mtp-npu-managed/control
+    consumer_state_dir: outputs/qwen3.5-4b-mtp-npu-managed/consumer-state
+    backend: mooncake
+    managed_local:
+      # Validated 16-card A3 layout: 6 TP=1 capture servers own devices 0-5,
+      # the 10-rank trainer owns devices 6-15 (global batch = 2*4*10 = 80).
+      # The pipeline is capture-bound, so scale capture servers first.
+      trainer_cuda_visible_devices:
+        ["6", "7", "8", "9", "10", "11", "12", "13", "14", "15"]
+      mooncake:
+        # Host segment mounted by the capture-server sink for feature storage.
+        # 16 GiB fits a 64GB-card host next to the target model and KV cache.
+        global_segment_size_bytes: 17179869184
+      capture_servers:
+        - port: 30000
+          cuda_visible_devices: ["0"]
+          tp_size: 1
+        - port: 30001
+          cuda_visible_devices: ["1"]
+          tp_size: 1
+        - port: 30002
+          cuda_visible_devices: ["2"]
+          tp_size: 1
+        - port: 30003
+          cuda_visible_devices: ["3"]
+          tp_size: 1
+        - port: 30004
+          cuda_visible_devices: ["4"]
+          tp_size: 1
+        - port: 30005
+          cuda_visible_devices: ["5"]
+          tp_size: 1

scripts/merge_mtp_to_base.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python3
+# coding=utf-8
+"""Merge a trained MTP draft checkpoint back into the base target model.
+
+Thin CLI wrapper around ``specforge.export.mtp.merge_mtp_into_base``; the merge
+logic and the per-family key mapping live in the package.
+
+Example:
+    python scripts/merge_mtp_to_base.py \
+        --base-model-path PATH/TO/Qwen3.5-4B \
+        --mtp-checkpoint-path PATH/TO/outputs/qwen3.5-4b-mtp/RUN-latest \
+        --draft-config configs/qwen3.5-4b-mtp.json \
+        --output-path PATH/TO/Qwen3.5-4B-MTP \
+        --key-format sglang
+"""
+
+import argparse
+
+from specforge.export.mtp import merge_mtp_into_base
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Merge trained MTP weights back into the base Qwen3.5 model."
+    )
+    parser.add_argument(
+        "--base-model-path",
+        type=str,
+        required=True,
+        help="Path to the original Qwen3.5 base model checkpoint.",
+    )
+    parser.add_argument(
+        "--mtp-checkpoint-path",
+        type=str,
+        required=True,
+        help=(
+            "SpecForge runtime checkpoint/output path, or an already-exported "
+            "HF MTP draft directory."
+        ),
+    )
+    parser.add_argument(
+        "--draft-config",
+        type=str,
+        default=None,
+        help=(
+            "Draft config JSON (required for a SpecForge runtime checkpoint; "
+            "the exported HF directory already contains config.json)."
+        ),
+    )
+    parser.add_argument(
+        "--output-path",
+        type=str,
+        required=True,
+        help="Directory to write the merged checkpoint.",
+    )
+    parser.add_argument(
+        "--key-format",
+        type=str,
+        default="sglang",
+        choices=["sglang", "hf"],
+        help=(
+            "MTP key layout. Both 'sglang' and 'hf' produce the flat native "
+            "layout (mtp.layers.0.* / mtp.norm.weight) that SGLang's flat "
+            "Qwen3_5ForCausalLMMTP and HF/vLLM MTP modules expect; the argument "
+            "is kept for backward compatibility."
+        ),
+    )
+    args = parser.parse_args()
+
+    merge_mtp_into_base(
+        args.base_model_path,
+        args.mtp_checkpoint_path,
+        args.output_path,
+        args.key_format,
+        draft_config_path=args.draft_config,
+    )
+
+
+if __name__ == "__main__":
+    main()

specforge/algorithms/builtin.py

@@ -6,14 +6,15 @@
 from specforge.algorithms.domino.providers import create_registration as domino
 from specforge.algorithms.dspark.providers import create_registration as dspark
 from specforge.algorithms.eagle3.providers import create_registration as eagle3
+from specforge.algorithms.mtp.providers import create_registration as mtp
 from specforge.algorithms.peagle.providers import create_registration as peagle
 from specforge.algorithms.registry import AlgorithmRegistry
 
 
 def builtin_algorithm_registry() -> AlgorithmRegistry:
     """Return a fresh immutable catalog without module-level mutation."""
 
-    return AlgorithmRegistry((eagle3(), peagle(), dflash(), domino(), dspark()))
+    return AlgorithmRegistry((eagle3(), peagle(), dflash(), domino(), dspark(), mtp()))
 
 
 __all__ = ["builtin_algorithm_registry"]

specforge/algorithms/common/dflash_family_data.py → specforge/algorithms/common/hidden_states_data.py

@@ -1,4 +1,7 @@
-"""Shared DFlash-family normalization and padding adapters."""
+"""Shared hidden-states normalization and padding adapters.
+
+Used by the DFlash-family (DFlash/Domino/DSpark) and MTP algorithms.
+"""
 
 from __future__ import annotations
 
@@ -9,6 +12,7 @@
 
 NORMALIZER_ID = "dflash_family_offline_v1"
 DSPARK_NORMALIZER_ID = "dspark_offline_v1"
+MTP_NORMALIZER_ID = "mtp_offline_v1"
 
 
 def _normalize_hidden_states(
@@ -185,15 +189,99 @@ def collate(features):
     return collate
 
 
+def normalize_mtp_offline_sample(raw, max_len: int):
+    """Normalize MTP capture tensors (no aux-layer concat, final hidden only)."""
+
+    input_ids = raw["input_ids"][:max_len].unsqueeze(0)
+    loss_mask = raw["loss_mask"][:max_len].unsqueeze(0)
+    target_last_hidden_states = _normalize_hidden_states(
+        raw,
+        "target_last_hidden_states",
+        max_len,
+        description="MTP target_last_hidden_states",
+    )
+    lengths = {
+        input_ids.shape[1],
+        loss_mask.shape[1],
+        target_last_hidden_states.shape[1],
+    }
+    if len(lengths) != 1:
+        raise ValueError(
+            "offline MTP features have mismatched sequence lengths after "
+            f"truncation: input_ids={input_ids.shape[1]}, "
+            f"loss_mask={loss_mask.shape[1]}, "
+            f"target_last_hidden_states={target_last_hidden_states.shape[1]}"
+        )
+    return {
+        "input_ids": input_ids,
+        "loss_mask": loss_mask,
+        "target_last_hidden_states": target_last_hidden_states,
+    }
+
+
+def build_mtp_offline_reader(
+    strategy,
+    hidden_states_path,
+    *,
+    run_id,
+    ttt_length,
+    max_len,
+):
+    # Transitional runtime import; the composition root will inject this port.
+    from specforge.runtime.data_plane.offline_reader import OfflineManifestReader
+
+    return OfflineManifestReader(
+        hidden_states_path,
+        run_id=run_id,
+        strategy=strategy,
+        feature_keys=(
+            "input_ids",
+            "loss_mask",
+            "target_last_hidden_states",
+        ),
+        target_repr="hidden_state",
+        ttt_length=ttt_length,
+        max_len=max_len,
+    )
+
+
+def build_mtp_offline_normalizer(max_len, **_topology):
+    return partial(normalize_mtp_offline_sample, max_len=max_len)
+
+
+def build_mtp_collator():
+    def collate(features):
+        return pad_and_concatenate_features(
+            features,
+            sequence_axes={
+                "input_ids": 1,
+                "loss_mask": 1,
+                "target_last_hidden_states": 1,
+            },
+            required_keys=(
+                "input_ids",
+                "loss_mask",
+                "target_last_hidden_states",
+            ),
+        )
+
+    return collate
+
+
 __all__ = [
     "DSPARK_NORMALIZER_ID",
+    "MTP_NORMALIZER_ID",
     "NORMALIZER_ID",
     "build_collator",
     "build_dspark_collator",
     "build_dspark_offline_normalizer",
     "build_dspark_offline_reader",
+    "build_mtp_collator",
+    "build_mtp_offline_normalizer",
+    "build_mtp_offline_reader",
     "build_offline_normalizer",
     "build_offline_reader",
     "normalize_dspark_offline_sample",
+    "normalize_mtp_offline_sample",
     "normalize_offline_sample",
 ]