From 8f0ff728caa1313dce1809b4258e68246a0845b5 Mon Sep 17 00:00:00 2001
From: Jesse Posner <jesse.posner@gmail.com>
Date: Fri, 13 Mar 2026 16:55:25 -0700
Subject: [PATCH] gguf: include MTP tensors for Qwen3-Next and Qwen3.5 models

Qwen3-Next and Qwen3.5 models ship with multi-token prediction (MTP)
heads, but the GGUF converter currently ignores them. This prevents
any MTP-based speculative decoding from being tested on these models.

Changes:
- convert_hf_to_gguf.py: Map MTP layer tensors to block indices above
  num_hidden_layers, and map non-layer MTP tensors (fc, norms) to NEXTN
  tensor names. Set nextn_predict_layers metadata.
- llama-arch.cpp: Add FFN_NORM and NEXTN_* tensors to QWEN3NEXT and
  QWEN35MOE architecture definitions so the loader accepts them.

The MTP layer in Qwen3.5-122B-A10B is a full attention + MoE block
(not DeltaNet), matching the model's full_attention_interval pattern.
At layer index 48 (= num_hidden_layers), 48 % 4 == 0, so it naturally
falls on an attention layer slot.

This is conversion-only. No graph builder or inference changes.
The tensors are loaded but unused until a MTP graph builder is added
(see PR #18886 for the MTP API effort).
---
 convert_hf_to_gguf.py | 44 +++++++++++++++++++++++++++++++++++++++++--
 src/llama-arch.cpp    | 14 ++++++++++++++
 2 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py
index eec0ea14e3a6..7833ef4f4c11 100755
--- a/convert_hf_to_gguf.py
+++ b/convert_hf_to_gguf.py
@@ -4662,9 +4662,49 @@ def set_gguf_parameters(self):
             rope_dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]
         self.gguf_writer.add_rope_dimension_count(int(rope_dim * self.hparams.get("partial_rotary_factor", 0.25)))
 
+        # MTP (multi-token prediction) support
+        n_mtp = self.hparams.get("mtp_num_hidden_layers", 0)
+        if n_mtp > 0:
+            self.block_count = self.hparams["num_hidden_layers"] + n_mtp
+            self.gguf_writer.add_block_count(self.block_count)
+            self.gguf_writer.add_nextn_predict_layers(n_mtp)
+
+    _mtp_remapper = {
+        "mtp.fc":                  "model.layers.{bid}.eh_proj",
+        "mtp.pre_fc_norm_embedding": "model.layers.{bid}.enorm",
+        "mtp.pre_fc_norm_hidden":  "model.layers.{bid}.hnorm",
+        "mtp.norm":                "model.layers.{bid}.shared_head.norm",
+    }
+
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        if name.startswith("mtp"):
-            return  # ignore MTP layers for now
+        if name.startswith("mtp."):
+            n_main = self.hparams["num_hidden_layers"]
+            n_mtp = self.hparams.get("mtp_num_hidden_layers", 0)
+            if n_mtp == 0:
+                return  # MTP not configured
+
+            if "layers." in name:
+                # mtp.layers.{i}.* -> model.layers.{n_main + i}.*
+                import re
+                m = re.match(r"mtp\.layers\.(\d+)\.(.*)", name)
+                if m:
+                    mtp_idx = int(m.group(1))
+                    rest = m.group(2)
+                    name = f"model.layers.{n_main + mtp_idx}.{rest}"
+                    bid = n_main + mtp_idx
+            else:
+                # Non-layer MTP tensors (fc, norms): map to NEXTN names
+                stem = name.rsplit(".", 1)[0]  # e.g. "mtp.fc"
+                suffix = "." + name.rsplit(".", 1)[1]  # e.g. ".weight"
+                if stem in self._mtp_remapper:
+                    template = self._mtp_remapper[stem]
+                    for mtp_bid in range(n_main, n_main + n_mtp):
+                        new_name = template.format(bid=mtp_bid) + suffix
+                        yield from super().modify_tensors(data_torch, new_name, mtp_bid)
+                    return
+                else:
+                    logger.warning(f"Unknown MTP tensor: {name}")
+                    return
         if name.endswith(".A_log"):
             data_torch = -torch.exp(data_torch)
         elif name.endswith(".dt_bias"):
diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
index 799d16167ba7..0474d8c8257a 100644
--- a/src/llama-arch.cpp
+++ b/src/llama-arch.cpp
@@ -1024,6 +1024,13 @@ static std::set<llm_tensor> llm_get_tensor_names(llm_arch arch) {
                 LLM_TENSOR_SSM_IN,
                 LLM_TENSOR_SSM_NORM,
                 LLM_TENSOR_SSM_OUT,
+                LLM_TENSOR_FFN_NORM,
+                LLM_TENSOR_NEXTN_EH_PROJ,
+                LLM_TENSOR_NEXTN_EMBED_TOKENS,
+                LLM_TENSOR_NEXTN_ENORM,
+                LLM_TENSOR_NEXTN_HNORM,
+                LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD,
+                LLM_TENSOR_NEXTN_SHARED_HEAD_NORM,
             };
         case LLM_ARCH_QWEN35:
             return {
@@ -1082,6 +1089,13 @@ static std::set<llm_tensor> llm_get_tensor_names(llm_arch arch) {
                 LLM_TENSOR_SSM_ALPHA,
                 LLM_TENSOR_SSM_NORM,
                 LLM_TENSOR_SSM_OUT,
+                LLM_TENSOR_FFN_NORM,
+                LLM_TENSOR_NEXTN_EH_PROJ,
+                LLM_TENSOR_NEXTN_EMBED_TOKENS,
+                LLM_TENSOR_NEXTN_ENORM,
+                LLM_TENSOR_NEXTN_HNORM,
+                LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD,
+                LLM_TENSOR_NEXTN_SHARED_HEAD_NORM,
             };
         case LLM_ARCH_QWEN3VL:
         case LLM_ARCH_CHAMELEON:
