From 19d590230f27815ff68c918df5e3e52be5732fea Mon Sep 17 00:00:00 2001
From: tomiock <ockier1@gmail.com>
Date: Tue, 9 Jun 2026 10:30:14 +0200
Subject: [PATCH 1/2] [fix] convertion script did not load state_dict properly

---
 utils/convertion_script.py | 120 ++++++++++++++++++++++++++-----------
 1 file changed, 84 insertions(+), 36 deletions(-)

diff --git a/utils/convertion_script.py b/utils/convertion_script.py
index a7785dd..148f70f 100644
--- a/utils/convertion_script.py
+++ b/utils/convertion_script.py
@@ -1,69 +1,117 @@
+"""Convert trainer DCP checkpoints -> HF safetensors snapshots.
+
+KEY MAPPING (this is where the old version was silently broken):
+  The trainer saves `self.model` (a Qwen3VL *ForConditionalGeneration*) inside a
+  state_dict container, and torch.compile wraps submodules, so DCP keys look like:
+
+      model.model.language_model.layers.0.self_attn.q_proj.weight   (double `model.`)
+      model.model.visual.merger._orig_mod.norm.weight               (compile `._orig_mod`)
+      model.lm_head._orig_mod.weight
+
+  To restore into an HF model we map each DCP key -> HF key by stripping the OUTER
+  `model.` container prefix and removing every `._orig_mod`:
+
+      model.model.language_model....  ->  model.language_model....   (HF ForConditionalGeneration)
+      model.lm_head._orig_mod.weight  ->  lm_head.weight
+
+  The OLD script used `AutoModel` (the inner Qwen3VLModel, keys `language_model.*`)
+  and `target = "model." + key` (single `model.`), and ignored `._orig_mod`. That
+  matched 0/625 weights, so `load_state_dict(strict=False)` kept the BASE weights and
+  every snapshot was just the untrained base model. We now use
+  AutoModelForImageTextToText (-> ForConditionalGeneration, keys `model.<...>` + lm_head)
+  and assert a high match rate so a silent miss can never happen again.
+
+Run:
+    python utils/convertion_script.py --base_model <hf_dir> --checkpoint_dir <dir_with_checkpoint-step-*>
+"""
+
 import torch
 import torch.distributed.checkpoint as dcp
-from transformers import AutoModel, AutoProcessor
+from transformers import AutoModelForImageTextToText, AutoProcessor
 import argparse
 import os
 import glob
 
-def convert_nested_dcp_batch(base_model_path, checkpoint_dir):
+
+def dcp_to_hf(dcp_key: str) -> str:
+    """Strip the outer trainer `model.` container and any torch.compile `._orig_mod`."""
+    key = dcp_key[len("model."):] if dcp_key.startswith("model.") else dcp_key
+    return key.replace("._orig_mod", "")
+
+
+def convert_nested_dcp_batch(base_model_path, checkpoint_dir, min_match_ratio=0.95):
     models_out_dir = os.path.join(checkpoint_dir, "models")
     os.makedirs(models_out_dir, exist_ok=True)
 
     search_pattern = os.path.join(checkpoint_dir, "checkpoint-step-*")
     checkpoint_dirs = [d for d in glob.glob(search_pattern) if os.path.isdir(d)]
-    
     if not checkpoint_dirs:
         print(f"No checkpoints found in {checkpoint_dir} matching 'checkpoint-step-*'")
         return
 
     print(f"Loading base model from {base_model_path}...")
-    model = AutoModel.from_pretrained(
-        base_model_path,
-        torch_dtype=torch.bfloat16, 
-        trust_remote_code=True, 
+    model = AutoModelForImageTextToText.from_pretrained(
+        base_model_path, torch_dtype=torch.bfloat16, trust_remote_code=True,
     )
     processor = AutoProcessor.from_pretrained(base_model_path, trust_remote_code=True)
-    
-    hf_state_dict = model.state_dict()
-    NESTING_PREFIX = "model."
+
+    ref_state = model.state_dict()          # HF ForConditionalGeneration keys
+    ref_keys = set(ref_state.keys())
 
     for ckpt_path in sorted(checkpoint_dirs):
         step_name = os.path.basename(ckpt_path)
-        step_num = step_name.split('-')[-1]
+        step_num = step_name.split("-")[-1]
         output_path = os.path.join(models_out_dir, f"step-{step_num}")
-        
         print(f"Processing {step_name} -> {output_path}")
 
         reader = dcp.FileSystemReader(ckpt_path)
-        metadata = reader.read_metadata()
-        checkpoint_keys = set(metadata.state_dict_metadata.keys())
-        
+        checkpoint_keys = set(reader.read_metadata().state_dict_metadata.keys())
+
+        # Build the load plan keyed by the DCP key, with correctly-shaped CPU buffers
+        # taken from the reference HF model. dcp_key -> hf_key.
+        dcp_to_hf_map = {}
         load_plan = {}
-        for hf_key, tensor in hf_state_dict.items():
-            target_key = NESTING_PREFIX + hf_key
-            if target_key in checkpoint_keys:
-                load_plan[target_key] = tensor
-
-        dcp.load(
-            state_dict=load_plan,
-            checkpoint_id=ckpt_path,
-        )
-
-        restored_state_dict = {}
-        for hf_key in hf_state_dict.keys():
-            target_key = NESTING_PREFIX + hf_key
-            if target_key in load_plan:
-                restored_state_dict[hf_key] = load_plan[target_key]
-                
-        model.load_state_dict(restored_state_dict, strict=False)
-        
+        for dk in checkpoint_keys:
+            hk = dcp_to_hf(dk)
+            if hk in ref_keys:
+                dcp_to_hf_map[dk] = hk
+                load_plan[dk] = torch.empty_like(ref_state[hk], device="cpu")
+
+        n_match, n_ref = len(load_plan), len(ref_keys)
+        # lm_head is tied -> absent from the checkpoint; don't count it against us.
+        tied = getattr(getattr(model, "config", None), "tie_word_embeddings", False)
+        effective_ref = n_ref - (1 if tied and "lm_head.weight" in ref_keys else 0)
+        ratio = n_match / max(1, effective_ref)
+        print(f"  matched {n_match}/{n_ref} HF tensors from checkpoint (ratio={ratio:.3f})")
+        if ratio < min_match_ratio:
+            raise RuntimeError(
+                f"Only {n_match}/{effective_ref} weights matched (ratio {ratio:.3f} < "
+                f"{min_match_ratio}). Key mapping is wrong -- refusing to write a snapshot "
+                f"that would silently be the untrained base model. "
+                f"Example checkpoint keys: {sorted(checkpoint_keys)[:3]}"
+            )
+
+        dcp.load(state_dict=load_plan, checkpoint_id=ckpt_path)
+
+        restored_state_dict = {dcp_to_hf_map[dk]: t for dk, t in load_plan.items()}
+        missing, unexpected = model.load_state_dict(restored_state_dict, strict=False)
+        missing = [m for m in missing if not (tied and m == "lm_head.weight")]
+        if missing:
+            print(f"  WARNING: {len(missing)} weights NOT restored (kept base), e.g. {missing[:5]}")
+        if unexpected:
+            print(f"  WARNING: {len(unexpected)} unexpected keys, e.g. {unexpected[:5]}")
+
         processor.save_pretrained(output_path)
         model.save_pretrained(output_path, safe_serialization=True)
         print(f"Saved HF snapshot to {output_path}\n")
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--base_model", required=True, help="Path to original HF model")
-    parser.add_argument("--checkpoint_dir", required=True, help="Path to the directory containing checkpoint folders")
-    
+    parser.add_argument("--checkpoint_dir", required=True,
+                        help="Path to the directory containing checkpoint folders")
+    parser.add_argument("--min_match_ratio", type=float, default=0.95,
+                        help="Fail if fewer than this fraction of model weights are restored")
     args = parser.parse_args()
-    convert_nested_dcp_batch(args.base_model, args.checkpoint_dir)
+    convert_nested_dcp_batch(args.base_model, args.checkpoint_dir, args.min_match_ratio)

From f812603905c5230963c880ffd449dc5d41a410cc Mon Sep 17 00:00:00 2001
From: tomiock <ockier1@gmail.com>
Date: Tue, 9 Jun 2026 11:01:16 +0200
Subject: [PATCH 2/2] [feat] Multi-Token Prediction (MTP) for Qwen3-VL and
 Qwen3.5

Depth-1 sequential MTP head (DeepSeek-V3 / Qwen3-Next style): a top-level
`mtp` module predicts token i+2 from the trunk hidden + embed(token i+1),
sharing embed_tokens and the tied lm_head. Param layout matches the released
Qwen3.5 checkpoints (mtp.fc / pre_fc_norm_{embedding,hidden} / layers / norm);
fc concat order is [embedding, hidden].

- Config-gated by text_config.mtp_num_hidden_layers (0 = disabled, fully
  backward-compatible; base Qwen3-VL checkpoints unaffected).
- Forward returns total = main_ce + mtp_loss_weight * mtp_ce; build_mtp_targets
  masks packed-segment boundaries so no supervision leaks across samples.
- Training wiring: train_mtp / random_init_mtp / lr_mtp / mtp_loss_weight,
  init_mtp, MTP optimizer group, FSDP/TP/compile, FLOPs accounting, warning_once
  on architecture-vs-config mismatch.
- eval/benchmark_mtp_decode.py: self-speculative greedy decode benchmark
  (verified ~1.61x decode speedup, ~0.77 acceptance, lossless on Qwen3.5-2B).
- models/tests/test_mtp.py: masking, forward/backward, checkpoint-load parity,
  and a draft-quality regression guard.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 eval/benchmark_mtp_decode.py | 313 +++++++++++++++++++++++++++++++++++
 models/qwen3_5/model.py      |  81 +++++++++
 models/qwen3_5/utils.py      |  26 +++
 models/qwen3_vl/model.py     | 108 ++++++++++++
 models/tests/test_mtp.py     | 193 +++++++++++++++++++++
 train/config.py              |   8 +-
 train/flops_estimation.py    |  10 ++
 train/infra.py               |  15 ++
 train/logger.py              |   8 +
 train/train_qwen.py          |  37 ++++-
 train/utils.py               |  29 +++-
 11 files changed, 818 insertions(+), 10 deletions(-)
 create mode 100644 eval/benchmark_mtp_decode.py
 create mode 100644 models/tests/test_mtp.py

diff --git a/eval/benchmark_mtp_decode.py b/eval/benchmark_mtp_decode.py
new file mode 100644
index 0000000..88aa5b5
--- /dev/null
+++ b/eval/benchmark_mtp_decode.py
@@ -0,0 +1,313 @@
+"""Verify decode-speed gain from the MTP head via self-speculative decoding.
+
+Compares plain greedy autoregressive decoding ("baseline") against depth-1 MTP
+self-speculation ("spec") on the native models, and reports the draft acceptance
+rate, tokens per trunk forward, and wall-clock tokens/sec.
+
+How depth-1 MTP speculation works here (greedy, exact):
+  - The trunk predicts the next token ``t`` (and gives hidden states).
+  - The MTP head drafts ``d`` = the token *after* ``t`` (from the trunk hidden +
+    embed(t), run over the full prefix so it attends to context).
+  - One trunk forward on ``[ids, t, d]`` verifies both positions in parallel:
+      * position of ``t`` predicts the true token after ``t`` (``r1``);
+      * position of ``d`` is a bonus predicting the token after ``d`` (``r2``).
+    If ``r1 == d`` the draft is accepted -> 2 tokens from 1 trunk forward;
+    otherwise 1 token (and ``r1`` is the corrected continuation).
+  Greedy spec output is identical to greedy baseline output, so the script also
+  asserts the two token streams match (a correctness check on the head).
+
+Caveats:
+  - No KV cache: every forward reprocesses the whole prefix (training-shaped
+    varlen forward). The trunk (many layers) dominates wall-clock, so the
+    speedup tracks "tokens per trunk forward". A KV-cache deployment would shift
+    absolute numbers but the acceptance rate is the portable metric.
+  - Run uncompiled (eager) to avoid recompilation noise across lengths.
+  - A randomly-initialized MTP head (e.g. a fresh Qwen3-VL head) drafts garbage
+    -> ~0 acceptance -> spec is slower. Use a checkpoint whose mtp.* weights are
+    trained (e.g. Qwen3.5) to see the real gain.
+
+Run:
+    CUDA_VISIBLE_DEVICES=7 python eval/benchmark_mtp_decode.py \
+        --model_dir /data/151-1/users/tockier/qwen_finetune/cache/qwen35_2b \
+        --model_type qwen3_5 --num_tokens 128
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+import time
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+PROMPTS = [
+    "The history of the Roman Empire began when",
+    "Here is a simple recipe for chocolate chip cookies. First,",
+    "In the field of machine learning, a transformer is",
+    "Once upon a time, in a small village by the sea, there lived",
+]
+
+
+def load_model(model_dir: str, model_type: str, device: str):
+    if model_type == "qwen3_5":
+        from models.qwen3_5.model import Qwen3_5ForCausalLM as M
+    elif model_type == "qwen3_vl":
+        from models.qwen3_vl.model import Qwen3VLForCausalLM as M
+    else:
+        raise ValueError(f"unknown model_type {model_type}")
+    model, cfg = M.from_pretrained(model_dir, dtype=torch.bfloat16, device=device, load_vision=False)
+    model.eval()
+    if model.mtp is None:
+        raise RuntimeError(
+            "model has no MTP head (mtp_num_hidden_layers=0 in config.json); "
+            "nothing to speculate with."
+        )
+    return model, cfg
+
+
+def _text_rope(model, length: int, device, dtype):
+    pos = torch.arange(length, device=device).view(1, 1, -1).expand(3, 1, -1)
+    cos, sin = model._compute_cos_sin(pos)
+    return cos.to(dtype), sin.to(dtype)
+
+
+def _trunk(model, seq: torch.Tensor):
+    """Return (hidden, logits) for a packed single-segment text row (1, M)."""
+    device = seq.device
+    dtype = model.lm_head.weight.dtype
+    M = seq.shape[1]
+    cos, sin = _text_rope(model, M, device, dtype)
+    cu = torch.tensor([0, M], dtype=torch.int32, device=device)
+    embeds = model.model.language_model.embed_tokens(seq)
+    h = model.model.language_model(embeds, cos, sin, cu, M)
+    return h, model.lm_head(h)
+
+
+def _mtp_draft(model, seq: torch.Tensor, h_seq: torch.Tensor, next_tok: torch.Tensor):
+    """Draft the token after ``next_tok`` using the MTP head over the full prefix.
+
+    ``h_seq`` is the trunk hidden over ``seq`` (its last position predicts
+    ``next_tok``). MTP consumes hidden_i + embed(token_{i+1}); at the last
+    position token_{i+1} is the just-predicted ``next_tok``.
+    """
+    device = seq.device
+    dtype = model.lm_head.weight.dtype
+    M = seq.shape[1]
+    cos, sin = _text_rope(model, M, device, dtype)
+    cu = torch.tensor([0, M], dtype=torch.int32, device=device)
+    shifted = torch.cat([seq[:, 1:], next_tok.view(1, 1)], dim=1)
+    next_embeds = model.model.language_model.embed_tokens(shifted)
+    mtp_h = model.mtp(h_seq, next_embeds, cos, sin, cu, M)
+    return model.lm_head(mtp_h[:, -1]).argmax(-1)[0]  # 0-dim token id
+
+
+@torch.no_grad()
+def baseline_decode(model, ids: torch.Tensor, n_tokens: int):
+    out = []
+    n_trunk = 0
+    while len(out) < n_tokens:
+        _, logits = _trunk(model, ids)
+        n_trunk += 1
+        t = logits[0, -1].argmax()
+        ids = torch.cat([ids, t.view(1, 1)], dim=1)
+        out.append(int(t))
+    return out, {"trunk_forwards": n_trunk}
+
+
+@torch.no_grad()
+def spec_decode(model, ids: torch.Tensor, n_tokens: int):
+    out = []
+    n_trunk = 0
+    n_steps = 0
+    n_accept = 0
+
+    h, logits = _trunk(model, ids)
+    n_trunk += 1
+    t = logits[0, -1].argmax()
+    d = _mtp_draft(model, ids, h, t)
+
+    while len(out) < n_tokens:
+        verify = torch.cat([ids, t.view(1, 1), d.view(1, 1)], dim=1)
+        h, logits = _trunk(model, verify)
+        n_trunk += 1
+        n_steps += 1
+        r1 = logits[0, -2].argmax()  # true token after t
+        r2 = logits[0, -1].argmax()  # bonus: token after d
+        if torch.equal(r1, d):
+            n_accept += 1
+            ids = verify
+            out.append(int(t)); out.append(int(d))
+            t = r2
+            d = _mtp_draft(model, ids, h, t)
+        else:
+            ids = torch.cat([ids, t.view(1, 1)], dim=1)
+            out.append(int(t))
+            t = r1
+            d = _mtp_draft(model, ids, h[:, : ids.shape[1]], t)
+
+    stats = {
+        "trunk_forwards": n_trunk,
+        "steps": n_steps,
+        "accept_rate": n_accept / max(n_steps, 1),
+    }
+    return out[:n_tokens], stats
+
+
+def _timed(fn, *args):
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    res = fn(*args)
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    return res, time.perf_counter() - t0
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model_dir", required=True)
+    ap.add_argument("--model_type", default="qwen3_5", choices=["qwen3_5", "qwen3_vl"])
+    ap.add_argument("--num_tokens", type=int, default=128)
+    ap.add_argument("--num_prompts", type=int, default=len(PROMPTS))
+    ap.add_argument("--device", default="cuda")
+    ap.add_argument("--out", default=None,
+                    help="write a presentation-ready Markdown report to this path")
+    ap.add_argument("--show_chars", type=int, default=320,
+                    help="max chars of generated text to show in the console")
+    args = ap.parse_args()
+
+    from transformers import AutoTokenizer
+
+    device = args.device
+    model, cfg = load_model(args.model_dir, args.model_type, device)
+    tok = AutoTokenizer.from_pretrained(args.model_dir, use_fast=False)
+    print(f"model={args.model_type}  mtp_num_hidden_layers={cfg.text.mtp_num_hidden_layers}")
+
+    # warmup (compile-free, but warms kernels / caches)
+    warm = tok(PROMPTS[0], return_tensors="pt").input_ids.to(device)
+    baseline_decode(model, warm, 8)
+    spec_decode(model, warm, 8)
+
+    results = []  # one dict per prompt
+    tot_base_tok = tot_base_t = tot_spec_tok = tot_spec_t = 0.0
+    spec_trunk = base_trunk = 0
+    for prompt in PROMPTS[: args.num_prompts]:
+        ids = tok(prompt, return_tensors="pt").input_ids.to(device)
+
+        (base_out, base_stats), base_t = _timed(baseline_decode, model, ids, args.num_tokens)
+        (spec_out, spec_stats), spec_t = _timed(spec_decode, model, ids, args.num_tokens)
+
+        match = base_out == spec_out  # greedy spec is exact -> identical text
+        text = tok.decode(spec_out, skip_special_tokens=True)
+
+        base_tps = args.num_tokens / base_t
+        spec_tps = args.num_tokens / spec_t
+        spec_trunk += spec_stats["trunk_forwards"]
+        base_trunk += base_stats["trunk_forwards"]
+        tot_base_tok += args.num_tokens; tot_base_t += base_t
+        tot_spec_tok += args.num_tokens; tot_spec_t += spec_t
+        results.append({
+            "prompt": prompt, "text": text,
+            "base_tps": base_tps, "spec_tps": spec_tps, "speedup": spec_tps / base_tps,
+            "accept": spec_stats["accept_rate"], "exact": match,
+        })
+
+    overall = {
+        "model": args.model_type,
+        "num_tokens": args.num_tokens,
+        "dtype": str(model.lm_head.weight.dtype).replace("torch.", ""),
+        "base_tps": tot_base_tok / tot_base_t,
+        "spec_tps": tot_spec_tok / tot_spec_t,
+        "speedup": (tot_spec_tok / tot_spec_t) / (tot_base_tok / tot_base_t),
+        "accept": sum(r["accept"] for r in results) / len(results),
+        "tok_per_fwd_base": tot_base_tok / base_trunk,
+        "tok_per_fwd_spec": tot_spec_tok / spec_trunk,
+        "all_exact": all(r["exact"] for r in results),
+    }
+    _print_console(results, overall, args.show_chars)
+    if args.out:
+        _write_markdown(args.out, results, overall)
+        print(f"\nMarkdown report written to {args.out}")
+
+
+def _print_console(results, overall, show_chars):
+    BOLD, DIM, GRN, CYN, RST = "\033[1m", "\033[2m", "\033[32m", "\033[36m", "\033[0m"
+    bar = "=" * 70
+    print(f"\n{bar}")
+    print(f"{BOLD}  MTP Self-Speculative Decoding — {overall['model']}{RST}")
+    print(bar)
+    print(f"  {DIM}greedy · no KV-cache (eager) · {overall['dtype']} · "
+          f"{overall['num_tokens']} tokens/prompt{RST}")
+    exact_note = "lossless (identical output)" if overall["all_exact"] else "OUTPUT DIFFERS!"
+    print(f"\n  {BOLD}{GRN}▶ {overall['speedup']:.2f}× faster decoding{RST}  "
+          f"({overall['base_tps']:.1f} → {overall['spec_tps']:.1f} tok/s)  ·  "
+          f"{exact_note}")
+    print(f"  {BOLD}▶ draft acceptance {overall['accept']:.0%}{RST}  ·  "
+          f"{overall['tok_per_fwd_spec']:.2f} tokens per model forward "
+          f"(vs {overall['tok_per_fwd_base']:.2f} baseline)")
+
+    print(f"\n  {BOLD}Sample generations{RST} {DIM}(prompt → continuation; "
+          f"baseline and MTP produce the same text){RST}")
+    print("  " + "-" * 66)
+    for i, r in enumerate(results, 1):
+        gen = r["text"].replace("\n", " ").strip()
+        if len(gen) > show_chars:
+            gen = gen[:show_chars].rstrip() + " …"
+        print(f"  {BOLD}[{i}]{RST} {CYN}{r['prompt']}{RST}")
+        print(f"      {gen}")
+        print(f"      {DIM}{r['speedup']:.2f}× · accept {r['accept']:.0%} · "
+              f"{'✓ identical' if r['exact'] else '✗ DIFFERS'}{RST}")
+
+    print(f"\n  {BOLD}Per-prompt metrics{RST}")
+    print(f"  {'prompt':<34}{'base t/s':>10}{'spec t/s':>10}{'speedup':>9}{'accept':>8}")
+    print("  " + "-" * 66)
+    for r in results:
+        print(f"  {r['prompt'][:32]:<34}{r['base_tps']:>10.1f}{r['spec_tps']:>10.1f}"
+              f"{r['speedup']:>8.2f}×{r['accept']:>8.0%}")
+    print(f"  {BOLD}{'OVERALL':<34}{overall['base_tps']:>10.1f}{overall['spec_tps']:>10.1f}"
+          f"{overall['speedup']:>8.2f}×{overall['accept']:>8.0%}{RST}")
+    print(bar + "\n")
+
+
+def _write_markdown(path, results, overall):
+    L = []
+    L.append(f"# MTP Self-Speculative Decoding — {overall['model']}\n")
+    L.append(f"*greedy · no KV-cache (eager) · {overall['dtype']} · "
+             f"{overall['num_tokens']} tokens/prompt*\n")
+    L.append(f"## Headline\n")
+    exact = "**lossless** — identical output to standard decoding" if overall["all_exact"] \
+        else "**WARNING: output differs from baseline**"
+    L.append(f"- 🚀 **{overall['speedup']:.2f}× faster decoding** "
+             f"({overall['base_tps']:.1f} → {overall['spec_tps']:.1f} tok/s)")
+    L.append(f"- ✅ {exact}")
+    L.append(f"- 🎯 **{overall['accept']:.0%} draft acceptance** — "
+             f"{overall['tok_per_fwd_spec']:.2f} tokens per model forward "
+             f"(vs {overall['tok_per_fwd_base']:.2f} baseline)\n")
+
+    L.append("## Throughput\n")
+    L.append("| Prompt | Baseline tok/s | MTP tok/s | Speedup | Accept |")
+    L.append("|---|--:|--:|--:|--:|")
+    for r in results:
+        L.append(f"| {r['prompt']} | {r['base_tps']:.1f} | {r['spec_tps']:.1f} "
+                 f"| {r['speedup']:.2f}× | {r['accept']:.0%} |")
+    L.append(f"| **Overall** | **{overall['base_tps']:.1f}** | **{overall['spec_tps']:.1f}** "
+             f"| **{overall['speedup']:.2f}×** | **{overall['accept']:.0%}** |\n")
+
+    L.append("## Sample generations\n")
+    L.append("_Baseline and MTP speculative decoding produce the **same** text "
+             "(greedy is exact); MTP just gets there faster._\n")
+    for i, r in enumerate(results, 1):
+        L.append(f"**{i}. {r['prompt']}**\n")
+        L.append(f"> {r['text'].strip()}\n")
+        L.append(f"<sub>{r['speedup']:.2f}× faster · {r['accept']:.0%} draft acceptance · "
+                 f"{'identical to baseline ✓' if r['exact'] else 'DIFFERS ✗'}</sub>\n")
+
+    with open(path, "w") as f:
+        f.write("\n".join(L))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/models/qwen3_5/model.py b/models/qwen3_5/model.py
index c180035..f9f5dea 100644
--- a/models/qwen3_5/model.py
+++ b/models/qwen3_5/model.py
@@ -20,6 +20,7 @@
     _local,
     CausalLMOutput,
     causal_lm_loss,
+    build_mtp_targets,
     apply_rope,
     mrope_cos_sin,
     apply_rope_vision,
@@ -306,6 +307,62 @@ def forward(
                 )
         return self.norm(x)
 
+class MTP(nn.Module):
+    """Multi-Token Prediction module (HF name: ``mtp``, top-level sibling of
+    ``model`` and ``lm_head``).
+
+    Matches the ``mtp.*`` parameter layout shipped in Qwen3.5/Qwen3-Next
+    checkpoints::
+
+        mtp.fc.weight                       Linear(2*hidden -> hidden, bias=False)
+        mtp.pre_fc_norm_embedding.weight    OffsetRMSNorm on the shifted embedding
+        mtp.pre_fc_norm_hidden.weight       OffsetRMSNorm on the trunk hidden
+        mtp.layers.{i}.*                    full-attention decoder layers
+        mtp.norm.weight                     final OffsetRMSNorm before shared head
+
+    ``embed_tokens`` and ``lm_head`` are shared with the main model (not stored
+    here). A single module predicts one extra token (offset +2).
+    """
+
+    def __init__(self, cfg: Qwen3_5TextConfig):
+        super().__init__()
+        self.pre_fc_norm_embedding = OffsetRMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps)
+        self.pre_fc_norm_hidden = OffsetRMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps)
+        self.fc = nn.Linear(2 * cfg.hidden_size, cfg.hidden_size, bias=False)
+        self.layers = nn.ModuleList(
+            [DecoderLayer(cfg, "full_attention") for _ in range(cfg.mtp_num_hidden_layers)]
+        )
+        self.norm = OffsetRMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        next_embeds: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+    ) -> torch.Tensor:
+        nh = self.pre_fc_norm_hidden(hidden_states)
+        ne = self.pre_fc_norm_embedding(next_embeds)
+        # fc expects [embedding ; hidden] (embedding in the first half) — this
+        # ordering matches the released Qwen3.5/Qwen3-Next mtp.fc weights.
+        x = self.fc(torch.cat([ne, nh], dim=-1))
+        for layer in self.layers:
+            x = layer(x, cos, sin, cu_seqlens, max_seqlen)
+        return self.norm(x)
+
+    @torch.no_grad()
+    def init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, OffsetRMSNorm):
+                # (1 + weight) parametrization → identity scale at weight = 0.
+                nn.init.zeros_(m.weight)
+
 class VisionPatchEmbed(nn.Module):
     def __init__(self, cfg: Qwen3_5VisionConfig):
         super().__init__()
@@ -577,6 +634,13 @@ def __init__(self, cfg: Qwen3_5Config, **kwargs):
         if cfg.tie_word_embeddings:
             self.lm_head.weight = self.model.language_model.embed_tokens.weight
 
+        # Multi-Token Prediction (depth-1). Built when the config declares MTP
+        # layers (the released Qwen3.5 checkpoints carry mtp.* weights).
+        self.mtp_loss_weight = 0.1
+        self.mtp = (
+            MTP(cfg.text) if cfg.text.mtp_num_hidden_layers > 0 else None
+        )
+
         # Text rope: store only inv_freq; cos/sin are computed per-forward via
         # MRoPE (3D position ids). For text-only inputs the 3 axes share the
         # same arange, which collapses to plain 1D rope.
@@ -793,6 +857,23 @@ def forward(
         if labels.dim() == 1:
             labels = labels.unsqueeze(0)
         loss = causal_lm_loss(logits, labels)
+
+        if self.mtp is not None and input_ids is not None:
+            # Depth-1 MTP: position i consumes embed(token i+1) + trunk hidden i,
+            # and predicts token i+2. Reuses the trunk's cos/sin/cu_seqlens.
+            shifted_ids = torch.roll(input_ids, shifts=-1, dims=-1)
+            next_embeds = self.model.language_model.embed_tokens(shifted_ids)
+            mtp_hidden = self.mtp(h, next_embeds, cos, sin, cu_seqlens, max_seqlen)
+            mtp_logits = self.lm_head(mtp_hidden)
+            mtp_targets = build_mtp_targets(labels, cu_seqlens, offset=2)
+            mtp_loss = F.cross_entropy(
+                mtp_logits.reshape(-1, mtp_logits.size(-1)).float(),
+                mtp_targets.reshape(-1),
+                ignore_index=-100,
+            )
+            loss = loss + self.mtp_loss_weight * mtp_loss
+            return CausalLMOutput(loss=loss, logits=logits, mtp_loss=mtp_loss.detach())
+
         return CausalLMOutput(loss=loss, logits=logits)
 
     @classmethod
diff --git a/models/qwen3_5/utils.py b/models/qwen3_5/utils.py
index 1523439..b23ebf5 100644
--- a/models/qwen3_5/utils.py
+++ b/models/qwen3_5/utils.py
@@ -45,6 +45,7 @@ def _local(param: torch.Tensor) -> torch.Tensor:
 class CausalLMOutput:
     loss: torch.Tensor
     logits: torch.Tensor
+    mtp_loss: torch.Tensor | None = None
 
 def causal_lm_loss(
     logits: torch.Tensor,
@@ -60,6 +61,30 @@ def causal_lm_loss(
         ignore_index=ignore_index,
     )
 
+def build_mtp_targets(
+    labels: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    offset: int = 2,
+    ignore_index: int = -100,
+) -> torch.Tensor:
+    """Targets for a depth-1 MTP head over a packed varlen row.
+
+    The MTP hidden at position ``i`` predicts token ``i+offset`` (offset=2: the
+    token *after* the main model's next-token). Targets are the main ``labels``
+    rolled left by ``offset`` so they inherit the assistant-only ``-100`` mask.
+    Positions whose target would cross a packed-sample boundary (the last
+    ``offset`` positions of each segment in ``cu_seqlens``) are masked.
+    """
+    total = labels.shape[-1]
+    device = labels.device
+    seg_id = torch.bucketize(
+        torch.arange(total, device=device), cu_seqlens[1:-1].to(device), right=True
+    )
+    valid = seg_id == torch.roll(seg_id, -offset)
+    valid[-offset:] = False
+    tgt = torch.roll(labels, shifts=-offset, dims=-1)
+    return torch.where(valid.expand_as(tgt), tgt, torch.full_like(tgt, ignore_index))
+
 def precompute_rope_cache(
     head_dim: int,
     max_seq_len: int,
@@ -203,5 +228,6 @@ def load_safetensors_into(
         missing.discard("lm_head.weight")
     if not load_vision:
         missing = {m for m in missing if not m.startswith("model.visual.")}
+    missing = {m for m in missing if not m.startswith("mtp.")}
     if missing:
         raise RuntimeError(f"Missing weights after load: {sorted(missing)[:8]} ... ({len(missing)} total)")
diff --git a/models/qwen3_vl/model.py b/models/qwen3_vl/model.py
index f328a8b..cf5def5 100644
--- a/models/qwen3_vl/model.py
+++ b/models/qwen3_vl/model.py
@@ -49,6 +49,7 @@ def dispatch_varlen_attention(
 class CausalLMOutput:
     loss: torch.Tensor
     logits: torch.Tensor
+    mtp_loss: torch.Tensor | None = None
 
 def causal_lm_loss(
     logits: torch.Tensor,
@@ -63,6 +64,31 @@ def causal_lm_loss(
         ignore_index=ignore_index,
     )
 
+def build_mtp_targets(
+    labels: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    offset: int = 2,
+    ignore_index: int = -100,
+) -> torch.Tensor:
+    """Targets for a depth-1 MTP head over a packed varlen row.
+
+    The MTP hidden at position ``i`` predicts token ``i+offset`` (offset=2: the
+    token *after* the main model's next-token). Targets are the main ``labels``
+    rolled left by ``offset`` so they inherit the assistant-only ``-100`` mask.
+    Positions whose target would cross a packed-sample boundary (the last
+    ``offset`` positions of each segment in ``cu_seqlens``) are set to
+    ``ignore_index`` so no supervision leaks across samples.
+    """
+    total = labels.shape[-1]
+    device = labels.device
+    seg_id = torch.bucketize(
+        torch.arange(total, device=device), cu_seqlens[1:-1].to(device), right=True
+    )
+    valid = seg_id == torch.roll(seg_id, -offset)
+    valid[-offset:] = False
+    tgt = torch.roll(labels, shifts=-offset, dims=-1)
+    return torch.where(valid.expand_as(tgt), tgt, torch.full_like(tgt, ignore_index))
+
 @dataclass
 class Qwen3VLTextConfig:
     vocab_size: int
@@ -78,6 +104,8 @@ class Qwen3VLTextConfig:
     tie_word_embeddings: bool
     mrope_section: list[int] | None = None
     mrope_interleaved: bool = True
+    mtp_num_hidden_layers: int = 0
+    mtp_use_dedicated_embeddings: bool = False
 
 @dataclass
 class Qwen3VLVisionConfig:
@@ -125,6 +153,8 @@ def from_json(cls, path: str | Path) -> "Qwen3VLConfig":
             tie_word_embeddings=tc.get("tie_word_embeddings", raw.get("tie_word_embeddings", False)),
             mrope_section=rs.get("mrope_section"),
             mrope_interleaved=rs.get("mrope_interleaved", True),
+            mtp_num_hidden_layers=tc.get("mtp_num_hidden_layers", 0),
+            mtp_use_dedicated_embeddings=tc.get("mtp_use_dedicated_embeddings", False),
         )
         vc = raw["vision_config"]
         vision = Qwen3VLVisionConfig(
@@ -333,6 +363,61 @@ def forward(
                 )
         return self.norm(x)
 
+class Qwen3VLMTP(nn.Module):
+    """Multi-Token Prediction module (HF name: ``mtp``, top-level sibling of
+    ``model`` and ``lm_head``).
+
+    Parameter layout mirrors Qwen3.5/Qwen3-Next ``mtp.*`` checkpoints::
+
+        mtp.fc.weight                       Linear(2*hidden -> hidden, bias=False)
+        mtp.pre_fc_norm_embedding.weight    RMSNorm on the shifted token embedding
+        mtp.pre_fc_norm_hidden.weight       RMSNorm on the trunk hidden state
+        mtp.layers.{i}.*                    standard Qwen3-VL decoder layers
+        mtp.norm.weight                     final RMSNorm before the shared head
+
+    ``embed_tokens`` and ``lm_head`` are shared with the main model and are NOT
+    stored here (no dedicated embeddings; matches the checkpoints). A single
+    module predicts one extra token (offset +2).
+    """
+
+    def __init__(self, cfg: Qwen3VLTextConfig):
+        super().__init__()
+        self.pre_fc_norm_embedding = RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps)
+        self.pre_fc_norm_hidden = RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps)
+        self.fc = nn.Linear(2 * cfg.hidden_size, cfg.hidden_size, bias=False)
+        self.layers = nn.ModuleList(
+            [Qwen3VLTextLayer(cfg) for _ in range(cfg.mtp_num_hidden_layers)]
+        )
+        self.norm = RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        next_embeds: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+    ) -> torch.Tensor:
+        nh = self.pre_fc_norm_hidden(hidden_states)
+        ne = self.pre_fc_norm_embedding(next_embeds)
+        # fc expects [embedding ; hidden] (embedding in the first half) — this
+        # ordering matches the released Qwen3.5/Qwen3-Next mtp.fc weights.
+        x = self.fc(torch.cat([ne, nh], dim=-1))
+        for layer in self.layers:
+            x = layer(x, cos, sin, cu_seqlens, max_seqlen)
+        return self.norm(x)
+
+    @torch.no_grad()
+    def init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, RMSNorm):
+                nn.init.ones_(m.weight)
+
 class Qwen3VLVisionPatchEmbed(nn.Module):
     def __init__(self, cfg: Qwen3VLVisionConfig):
         super().__init__()
@@ -629,6 +714,11 @@ def __init__(self, cfg: Qwen3VLConfig, **kwargs):
         if cfg.tie_word_embeddings:
             self.lm_head.weight = self.model.language_model.embed_tokens.weight
 
+        self.mtp_loss_weight = 0.1
+        self.mtp = (
+            Qwen3VLMTP(cfg.text) if cfg.text.mtp_num_hidden_layers > 0 else None
+        )
+
         # Text rope: store only inv_freq; cos/sin are computed per-forward via
         # MRoPE (3D position ids). For text-only inputs the 3 axes share the
         # same arange, which collapses to plain 1D rope.
@@ -838,6 +928,23 @@ def forward(
         if labels.dim() == 1:
             labels = labels.unsqueeze(0)
         loss = causal_lm_loss(logits, labels)
+
+        if self.mtp is not None and input_ids is not None:
+            # Depth-1 MTP: position i consumes embed(token i+1) + trunk hidden i,
+            # and predicts token i+2. Reuses the trunk's cos/sin/cu_seqlens.
+            shifted_ids = torch.roll(input_ids, shifts=-1, dims=-1)
+            next_embeds = self.model.language_model.embed_tokens(shifted_ids)
+            mtp_hidden = self.mtp(h, next_embeds, cos, sin, cu_seqlens, max_seqlen)
+            mtp_logits = self.lm_head(mtp_hidden)
+            mtp_targets = build_mtp_targets(labels, cu_seqlens, offset=2)
+            mtp_loss = F.cross_entropy(
+                mtp_logits.reshape(-1, mtp_logits.size(-1)).float(),
+                mtp_targets.reshape(-1),
+                ignore_index=-100,
+            )
+            loss = loss + self.mtp_loss_weight * mtp_loss
+            return CausalLMOutput(loss=loss, logits=logits, mtp_loss=mtp_loss.detach())
+
         return CausalLMOutput(loss=loss, logits=logits)
 
     @classmethod
@@ -941,6 +1048,7 @@ def load_safetensors_into(
         missing.discard("lm_head.weight")
     if not load_vision:
         missing = {m for m in missing if not m.startswith("model.visual.")}
+    missing = {m for m in missing if not m.startswith("mtp.")}
     if missing:
         raise RuntimeError(f"Missing weights after load: {sorted(missing)[:8]} ... ({len(missing)} total)")
 
diff --git a/models/tests/test_mtp.py b/models/tests/test_mtp.py
new file mode 100644
index 0000000..927129a
--- /dev/null
+++ b/models/tests/test_mtp.py
@@ -0,0 +1,193 @@
+"""Multi-Token Prediction (MTP) tests for Qwen3-VL and Qwen3.5.
+
+Covers:
+  * boundary masking of MTP targets over a packed varlen row (CPU, no kernels);
+  * the MTP module is absent and the loss path is unchanged when disabled;
+  * a depth-1 MTP forward/backward produces finite losses and grads (CUDA);
+  * the native Qwen3.5 loader maps the checkpoint's ``mtp.*`` weights exactly.
+
+Run:
+    pytest models/tests/test_mtp.py
+    python models/tests/test_mtp.py
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+
+from models.qwen3_vl.model import (
+    Qwen3VLConfig,
+    Qwen3VLForCausalLM,
+    Qwen3VLTextConfig,
+    Qwen3VLVisionConfig,
+    build_mtp_targets,
+)
+
+QWEN35_SNAPSHOT = os.environ.get(
+    "QWEN35_SNAPSHOT",
+    "/data/151-1/users/tockier/qwen_finetune/cache/qwen35_2b",
+)
+
+
+def _vl_cfg(mtp_layers: int) -> Qwen3VLConfig:
+    text = Qwen3VLTextConfig(
+        vocab_size=256, hidden_size=64, intermediate_size=128, num_hidden_layers=2,
+        num_attention_heads=4, num_key_value_heads=2, head_dim=16,
+        max_position_embeddings=512, rms_norm_eps=1e-6, rope_theta=10000.0,
+        tie_word_embeddings=True, mtp_num_hidden_layers=mtp_layers,
+    )
+    vision = Qwen3VLVisionConfig(
+        depth=1, hidden_size=32, intermediate_size=64, num_heads=2, in_channels=3,
+        patch_size=16, temporal_patch_size=2, spatial_merge_size=2,
+        num_position_embeddings=64, out_hidden_size=64, hidden_act="gelu",
+        deepstack_visual_indexes=[],
+    )
+    return Qwen3VLConfig(
+        text=text, vision=vision, image_token_id=200, video_token_id=201,
+        vision_start_token_id=202, vision_end_token_id=203, tie_word_embeddings=True,
+    )
+
+
+def test_build_mtp_targets_no_cross_segment_leak():
+    # Two packed samples: [0,6) and [6,10).
+    cu = torch.tensor([0, 6, 10], dtype=torch.int32)
+    labels = torch.arange(100, 110).view(1, 10)
+    tgt = build_mtp_targets(labels, cu, offset=2)
+    # offset=2: position i predicts token i+2; last two positions of each
+    # segment have no in-segment target and must be masked.
+    expected = [102, 103, 104, 105, -100, -100, 108, 109, -100, -100]
+    assert tgt[0].tolist() == expected
+
+
+def test_build_mtp_targets_inherits_ignore_mask():
+    cu = torch.tensor([0, 8], dtype=torch.int32)
+    labels = torch.tensor([[10, -100, 12, 13, -100, 15, 16, 17]])
+    tgt = build_mtp_targets(labels, cu, offset=2)
+    # tgt[i] = labels[i+2]; last two positions masked by the boundary rule.
+    assert tgt[0].tolist() == [12, 13, -100, 15, 16, 17, -100, -100]
+
+
+def test_mtp_disabled_is_noop():
+    model = Qwen3VLForCausalLM(_vl_cfg(0))
+    assert model.mtp is None
+    assert not any(n.startswith("mtp.") for n, _ in model.named_parameters())
+
+
+@pytest.mark.cuda_only
+def test_mtp_forward_backward():
+    if not torch.cuda.is_available():
+        pytest.skip("requires CUDA (varlen flash kernel)")
+    dev = "cuda"
+    ids = torch.randint(0, 200, (1, 10), device=dev)
+    cu = torch.tensor([0, 6, 10], dtype=torch.int32, device=dev)
+
+    base = Qwen3VLForCausalLM(_vl_cfg(0)).to(dev).to(torch.bfloat16)
+    out0 = base(input_ids=ids, attention_mask=cu, labels=ids.clone())
+    assert out0.mtp_loss is None
+
+    model = Qwen3VLForCausalLM(_vl_cfg(1)).to(dev).to(torch.bfloat16)
+    model.mtp.init_weights()
+    out = model(input_ids=ids, attention_mask=cu, labels=ids.clone())
+    assert out.mtp_loss is not None and torch.isfinite(out.mtp_loss)
+    out.loss.backward()
+    g = model.mtp.fc.weight.grad
+    assert g is not None and torch.isfinite(g).all() and g.norm() > 0
+
+    # MTP param layout matches Qwen3.5/Qwen3-Next checkpoints.
+    names = {n for n, _ in model.named_parameters() if n.startswith("mtp.")}
+    for required in (
+        "mtp.fc.weight",
+        "mtp.pre_fc_norm_embedding.weight",
+        "mtp.pre_fc_norm_hidden.weight",
+        "mtp.norm.weight",
+        "mtp.layers.0.self_attn.q_proj.weight",
+    ):
+        assert required in names, required
+
+
+@pytest.mark.cuda_only
+def test_qwen35_loads_checkpoint_mtp_weights():
+    import json
+
+    from safetensors import safe_open
+
+    if not os.path.isdir(QWEN35_SNAPSHOT):
+        pytest.skip(f"no Qwen3.5 snapshot at {QWEN35_SNAPSHOT}")
+    from models.qwen3_5.model import Qwen3_5ForCausalLM
+
+    model, cfg = Qwen3_5ForCausalLM.from_pretrained(
+        QWEN35_SNAPSHOT, dtype=torch.bfloat16, device="cpu", load_vision=True
+    )
+    assert model.mtp is not None and cfg.text.mtp_num_hidden_layers > 0
+
+    sd = dict(model.state_dict())
+    wm = json.load(
+        open(os.path.join(QWEN35_SNAPSHOT, "model.safetensors.index.json"))
+    )["weight_map"]
+    mtp_keys = [k for k in wm if k.startswith("mtp.")]
+    assert len(mtp_keys) > 0
+    for k in mtp_keys:
+        with safe_open(os.path.join(QWEN35_SNAPSHOT, wm[k]), framework="pt") as f:
+            ref = f.get_tensor(k).to(torch.bfloat16)
+        assert torch.equal(sd[k], ref), k
+
+
+@pytest.mark.cuda_only
+def test_qwen35_mtp_head_drafts_usefully():
+    """Teacher-forced guard on the trained head: it must predict token i+2 well
+    above chance. Catches mis-invocation (e.g. wrong mtp.fc concat order, which
+    previously dropped this from ~0.55 to ~0.05 yet passed every other test)."""
+    if not torch.cuda.is_available():
+        pytest.skip("requires CUDA")
+    if not os.path.isdir(QWEN35_SNAPSHOT):
+        pytest.skip(f"no Qwen3.5 snapshot at {QWEN35_SNAPSHOT}")
+    from transformers import AutoTokenizer
+
+    from models.qwen3_5.model import Qwen3_5ForCausalLM
+
+    model, _ = Qwen3_5ForCausalLM.from_pretrained(
+        QWEN35_SNAPSHOT, dtype=torch.bfloat16, device="cuda", load_vision=False
+    )
+    model.eval()
+    tok = AutoTokenizer.from_pretrained(QWEN35_SNAPSHOT, use_fast=False)
+    ids = tok(
+        "The history of the Roman Empire is a long and complex story that spans "
+        "many centuries, from the founding of the city to the fall of the west.",
+        return_tensors="pt",
+    ).input_ids.to("cuda")
+    L = ids.shape[1]
+    dtype = model.lm_head.weight.dtype
+    lm = model.model.language_model
+    pos = torch.arange(L, device="cuda").view(1, 1, -1).expand(3, 1, -1)
+    cos, sin = model._compute_cos_sin(pos)
+    cos, sin = cos.to(dtype), sin.to(dtype)
+    cu = torch.tensor([0, L], dtype=torch.int32, device="cuda")
+    with torch.no_grad():
+        h = lm(lm.embed_tokens(ids), cos, sin, cu, L)
+        next_embeds = lm.embed_tokens(torch.roll(ids, -1, dims=1))
+        mtp_h = model.mtp(h, next_embeds, cos, sin, cu, L)
+        pred = model.lm_head(mtp_h)[0].argmax(-1)
+        tgt = torch.roll(ids, -2, dims=1)[0]
+        acc = (pred[: L - 2] == tgt[: L - 2]).float().mean().item()
+    assert acc > 0.3, f"MTP target+2 top1={acc:.3f} too low — head mis-invoked"
+
+
+if __name__ == "__main__":
+    test_build_mtp_targets_no_cross_segment_leak()
+    test_build_mtp_targets_inherits_ignore_mask()
+    test_mtp_disabled_is_noop()
+    print("CPU tests passed")
+    if torch.cuda.is_available():
+        test_mtp_forward_backward()
+        test_qwen35_loads_checkpoint_mtp_weights()
+        test_qwen35_mtp_head_drafts_usefully()
+        print("CUDA tests passed")
+    else:
+        print("CUDA unavailable; skipped GPU tests")
diff --git a/train/config.py b/train/config.py
index f4e33c5..16dbf5b 100644
--- a/train/config.py
+++ b/train/config.py
@@ -25,6 +25,8 @@ class Model:
     train_llm: bool = True
     train_mlp: bool = True
     train_vit: bool = False
+    train_mtp: bool = False
+    # warn: the MTP can be disabled at the model's config
 
 @dataclass
 class Wandb:
@@ -69,9 +71,13 @@ class Training:
     lr_llm: float = 2e-6
     lr_mlp: float = 1e-5
     lr_vit: float = 1e-6
+    lr_mtp: float = 1e-4
+
+    # weight of the MTP auxiliary loss: total = main_ce + mtp_loss_weight * mtp_ce
+    mtp_loss_weight: float = 0.1
 
-    # init of the projecter and deepstack layers
     random_init: bool = False
+    random_init_mtp: bool = False
 
     # gradient accumulation
     tpi_multiplier: float = 1.0
diff --git a/train/flops_estimation.py b/train/flops_estimation.py
index b813a1f..12219ea 100644
--- a/train/flops_estimation.py
+++ b/train/flops_estimation.py
@@ -152,6 +152,11 @@ def qwen3_vl_flops(model_config: Qwen3VLConfig | Qwen3_5Config):
             + logits_term
         )
 
+        # MTP head: mtp_num_hidden_layers decoder layers + one extra logits proj.
+        mtp_layers = getattr(model_config.text, "mtp_num_hidden_layers", 0)
+        if mtp_layers > 0:
+            text_total_flops += mtp_layers * (self_attn_term + mlp_term) + logits_term
+
         return text_total_flops + vision_flops(model_config)
 
     def qwen3_5_flops(model_config: Qwen3_5Config):
@@ -197,6 +202,11 @@ def qwen3_5_flops(model_config: Qwen3_5Config):
             + logits_term
         )
 
+        # MTP head: full-attention decoder layer(s) + one extra logits proj.
+        mtp_layers = getattr(text, "mtp_num_hidden_layers", 0)
+        if mtp_layers > 0:
+            text_total_flops += mtp_layers * (full_attn_term + mlp_term) + logits_term
+
         return text_total_flops + vision_flops(model_config)
 
     if model_type is ModelType.Qwen3_vl:
diff --git a/train/infra.py b/train/infra.py
index 9483b92..c65b6e9 100644
--- a/train/infra.py
+++ b/train/infra.py
@@ -247,6 +247,10 @@ def compile_model(model: torch.nn.Module):
 
     inner.visual.merger = torch.compile(inner.visual.merger, fullgraph=False, mode='max-autotune-no-cudagraphs')
 
+    if getattr(model, "mtp", None) is not None:
+        for transformer_block in model.mtp.layers:
+            transformer_block.compile(dynamic=True, fullgraph=False, mode='default')
+
 def apply_fsdp(model_type, model, **kwargs):
     if model_type == ModelType.Qwen3_text:
         apply_fsdp_qwen3(model, **kwargs)
@@ -288,6 +292,7 @@ def apply_fsdp_qwen3_vl(model, mesh, reshard_after_forward_policy='never'):
 
     fully_shard(model.lm_head, mesh=mesh, reshard_after_forward=False)
 
+    outer = model
     model = model.model
 
     match reshard_after_forward_policy:
@@ -343,6 +348,16 @@ def apply_fsdp_qwen3_vl(model, mesh, reshard_after_forward_policy='never'):
             reshard_after_forward=reshard_after_forward_policy == "always",
     )
 
+    # MTP head (top-level sibling of `model`).
+    if getattr(outer, "mtp", None) is not None:
+        for transformer_block in outer.mtp.layers:
+            fully_shard(
+                transformer_block,
+                mesh=mesh,
+                reshard_after_forward=reshard_after_forward,
+            )
+        fully_shard(outer.mtp, mesh=mesh, reshard_after_forward=reshard_after_forward)
+
     fully_shard(model, mesh=mesh)
 
 def apply_tp(
diff --git a/train/logger.py b/train/logger.py
index d1c2cd6..60003d1 100644
--- a/train/logger.py
+++ b/train/logger.py
@@ -18,6 +18,14 @@ class Color:
 
 logger = logging.getLogger("train_logger")
 
+_warned_once: set[str] = set()
+
+def warning_once(msg: str) -> None:
+    """Emit ``msg`` at WARNING level at most once per process (dedup by text)."""
+    if msg not in _warned_once:
+        _warned_once.add(msg)
+        logger.warning(msg)
+
 def init_logger():
     # Clear existing handlers to avoid duplicates
     if logger.handlers:
diff --git a/train/train_qwen.py b/train/train_qwen.py
index 97cf011..cf91212 100644
--- a/train/train_qwen.py
+++ b/train/train_qwen.py
@@ -26,7 +26,7 @@
 # training imports
 from train.config_manager import ConfigManager
 from train.config import Config, ModelType
-from train.logger import init_logger, logger, Color
+from train.logger import init_logger, logger, warning_once, Color
 from train.infra import (
     get_mesh,
     get_tp_group,
@@ -44,6 +44,7 @@
 
     init_qwen35,
     init_qwen3vl,
+    init_mtp,
 
     dist_mean,
     dist_max,
@@ -168,6 +169,23 @@ def __init__(self, cfg: Config):
             else:
                 logger.info('model not initlized, incompatible')
 
+        # MTP head: forward computes total = main_ce + mtp_loss_weight * mtp_ce.
+        self.model.mtp_loss_weight = self.training_args.mtp_loss_weight
+        if getattr(self.model, "mtp", None) is None and (
+            self.model_args.train_mtp or self.training_args.random_init_mtp
+        ):
+            warning_once(
+                "MTP is enabled in the training config "
+                f"(train_mtp={self.model_args.train_mtp}, "
+                f"random_init_mtp={self.training_args.random_init_mtp}) but the model "
+                "has no MTP head (mtp_num_hidden_layers=0 / absent in the model's "
+                "config.json). The MTP flags are ignored. To enable it, set "
+                '"mtp_num_hidden_layers" > 0 in text_config of the model config.json.'
+            )
+        if self.training_args.random_init_mtp:
+            logger.info('randomly initializing MTP head')
+            init_mtp(self.model)
+
         # replace flash_attn
         self.model.train()
         if self.model_args.model_impl == "hf":
@@ -277,17 +295,19 @@ def create_optimizer(self):
         lr_mlp = self.training_args.lr_mlp
         lr_vit = self.training_args.lr_vit
         lr_llm = self.training_args.lr_llm
+        lr_mtp = self.training_args.lr_mtp
 
         mlp_params = []
         vision_params = []
         llm_params = []
+        mtp_params = []
 
         for n, p in self.model.named_parameters():
             if not p.requires_grad:
                 continue
-            if "visual.merger" in n:
-                mlp_params.append(p)
-            elif "visual.deepstack_merger_list":
+            if "mtp." in n:
+                mtp_params.append(p)
+            elif "visual.merger" in n or "visual.deepstack_merger_list" in n:
                 mlp_params.append(p)
             elif "visual.patch_embed" in n:
                 vision_params.append(p)
@@ -296,6 +316,7 @@ def create_optimizer(self):
             else:
                 llm_params.append(p)
 
+        # AdamW skips empty param groups, so an absent MTP head is harmless.
         optimizer_grouped_parameters = [
             {
                 "params": mlp_params,
@@ -309,6 +330,10 @@ def create_optimizer(self):
                 "params": llm_params,
                 "lr": lr_llm,
             },
+            {
+                "params": mtp_params,
+                "lr": lr_mtp,
+            },
         ]
 
         # TODO: add weight decay exclusion for bias and LayerNorm
@@ -523,6 +548,9 @@ def log(self, avg_loss, max_loss, global_tokens, global_assistant_tokens, global
         if gathered is not None:
             log_metrics.update(self._topk_metrics(gathered))
 
+        if getattr(self, "mtp_loss_val", None) is not None:
+            log_metrics["train/mtp_loss"] = self.mtp_loss_val.item()
+
         wandb.log(log_metrics, step=self.global_step)
 
     def setup_accumulation(self, tpi_multiplier=1.5):
@@ -541,6 +569,7 @@ def train_step(self, data_iterator, optimizer):
                     **batch
                 )
                 loss = outputs.loss
+                self.mtp_loss_val = getattr(outputs, "mtp_loss", None)
 
         with record_function("backward_pass"):
             scaled_loss = loss / self.current_accum_target
diff --git a/train/utils.py b/train/utils.py
index 6534e95..ac3669b 100644
--- a/train/utils.py
+++ b/train/utils.py
@@ -85,6 +85,21 @@ def init_weights(m):
     for param in model.visual.merger.parameters():
         torch.distributed.broadcast(param.data, src=0)
 
+def init_mtp(model):
+    """Initialize a freshly-added MTP head and broadcast from rank 0.
+
+    Call with the outer ``*ForCausalLM`` (the module that owns ``.mtp``). No-op
+    when the model has no MTP head, or when the head was loaded from a checkpoint
+    that already contained ``mtp.*`` weights (the caller decides).
+    """
+    mtp = getattr(model, "mtp", None)
+    if mtp is None:
+        return
+    torch.manual_seed(42)
+    mtp.init_weights()
+    for param in mtp.parameters():
+        torch.distributed.broadcast(param.data, src=0)
+
 def generate_accumulation_pattern(target_multiplier: float, pattern_length: int = 100) -> list[int]:
     if target_multiplier < 1.0:
         raise ValueError("Multiplier must be >= 1.0")
@@ -150,11 +165,10 @@ def set_model_qwen3_5(model_args: ModelArgs, model):
         p.requires_grad = model_args.train_llm
     model.lm_head.requires_grad = model_args.train_llm
 
-    # MTP Heads (Tie to LLM if computing MTP loss, otherwise force False)
-    for n, p in model.named_parameters():
-        if "mtp" in n.lower():
-            # TODO: implement MTP and unfreeze the Module
-            p.requires_grad = False
+    # MTP head (independent of the backbone freeze flags)
+    if getattr(model, "mtp", None) is not None:
+        for n, p in model.mtp.named_parameters():
+            p.requires_grad = model_args.train_mtp
 
     return model
 
@@ -174,6 +188,11 @@ def set_model_qwen3vl(model_args: ModelArgs, model):
         p.requires_grad = model_args.train_llm
     model.lm_head.requires_grad = model_args.train_llm
 
+    # MTP head (independent of the backbone freeze flags)
+    if getattr(model, "mtp", None) is not None:
+        for n, p in model.mtp.named_parameters():
+            p.requires_grad = model_args.train_mtp
+
     return model
 
 def set_model_qwen3(model_args: ModelArgs, model):