DeepSeek 4 MTP implementation (#2216)

* add standalone DeepSeek V4 MTP

* fix openPangu indexer tensor identities

* spec-bench: checkpoint DeepSeek V4 before draft

* minor changes in comments

common/speculative.cpp

@@ -1906,7 +1906,7 @@ bool common_speculative_load_draft_model(
     }
 
     gpt_params params_dft = params_base;
-    params_dft.devices          = params.devices;
+    params_dft.devices          = params.devices.empty() ? params_base.devices : params.devices;
     params_dft.model            = params.model;
     params_dft.n_gpu_layers     = params.n_gpu_layers;
     params_dft.cache_type_k     = params.cache_type_k.empty() ? params_base.cache_type_k : params.cache_type_k;
@@ -1923,6 +1923,15 @@ bool common_speculative_load_draft_model(
         free_command_line(argc, argv);
     }
 
+    // We likely dont want to inehit offload policy for MTP
+    if (params.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP)) {
+        params_dft.ncmoe = 0;
+        params_dft.tensor_buft_overrides.clear();
+        params_dft.offload_policy.clear();
+        LOG_INF("%s: MTP draft ignores target CPU-MoE/tensor placement overrides\n",
+                __func__);
+    }
+
     LOG_INF("%s: loading draft model '%s'\n", __func__, params_dft.model.c_str());
 
     if (params_dft.n_ctx == 0) {
@@ -2033,6 +2042,15 @@ bool common_speculative_finalize_startup(
         const llama_model * model) {
     auto & params = params_base.speculative;
 
+    if (params.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP) &&
+        model != nullptr &&
+        llama_model_is_deepseek4(model) &&
+        llama_model_n_nextn_layer(model) > 1) {
+        LOG_ERR("%s: DeepSeek-V4 MTP supports exactly one NextN predictor layer, got %d.\n",
+                __func__, llama_model_n_nextn_layer(model));
+        return false;
+    }
+
     if (!params.needs_dft_model()) {
         params.clear_dft();
     }
@@ -2042,11 +2060,22 @@ bool common_speculative_finalize_startup(
         if (!common_speculative_load_draft_model(params, params_base)) {
             return false;
         }
+
+        if (params.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP) &&
+            params.model_dft != nullptr &&
+            llama_model_is_deepseek4(params.model_dft) &&
+            llama_model_n_nextn_layer(params.model_dft) != 1) {
+            LOG_ERR("%s: DeepSeek-V4 MTP draft requires exactly one NextN predictor layer, got %d.\n",
+                    __func__, llama_model_n_nextn_layer(params.model_dft));
+            return false;
+        }
     }
 
     params_base.has_mtp = params.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP);
-    const bool has_external_mtp = params_base.has_mtp &&
-        llama_model_is_gemma4_mtp_assistant(params.model_dft);
+    const bool has_external_mtp = params_base.has_mtp && params.model_dft &&
+        (llama_model_is_gemma4_mtp_assistant(params.model_dft) ||
+         (llama_model_is_deepseek4(params.model_dft) &&
+          llama_model_n_nextn_layer(params.model_dft) == 1));
 
     params_base.has_mtp = common_speculative_prepare_mtp_runtime(
         params,
@@ -2385,6 +2414,7 @@ bool common_speculative_checkpoint_restore(
 
         if (!ids.empty()) {
             const int n_re = (int) ids.size();
+            std::vector<float> redecoded_hidden;
             llama_batch re_batch = llama_batch_init(n_re, 0, 1);
             common_batch_add(re_batch, ckpt.sampled, ckpt.n_past, { seq_id }, n_re == 1);
             for (int j = 0; j < n_re - 1; ++j) {
@@ -3236,7 +3266,7 @@ common_speculative_round_result common_speculative_run_round(
         return result;
     }
 
-    if (llama_model_has_recurrent(model) || llama_model_is_openpangu(model)) {
+    if (common_speculative_needs_checkpoint(model)) {
         if (!common_speculative_before_draft(
                 spec,
                 model,

convert_hf_to_gguf.py

@@ -46,6 +46,7 @@ AnyModel = TypeVar("AnyModel", bound="type[Model]")
 
 class Model:
     _model_classes: dict[str, type[Model]] = {}
+    mtp_only = False
 
     dir_model: Path
     ftype: gguf.LlamaFileType
@@ -279,7 +280,7 @@ class Model:
             old_dtype = data_torch.dtype
 
             # convert any unsupported data types to float32
-            if data_torch.dtype not in (torch.float16, torch.float32):
+            if data_torch.dtype not in (torch.float16, torch.float32) and not self.mtp_only:
                 data_torch = data_torch.to(torch.float32)
 
             # use the first number-like part of the tensor name as the block id
@@ -289,7 +290,10 @@ class Model:
                     bid = int(part)
                     break
 
-            for new_name, data in ((n, d.squeeze().numpy()) for n, d in self.modify_tensors(data_torch, name, bid)):
+            for new_name, data in ((
+                n,
+                (d if self.mtp_only and self.model_arch == gguf.MODEL_ARCH.DEEPSEEK4 and n == "output_hc_scale.weight" else d.squeeze()).numpy(),
+            ) for n, d in self.modify_tensors(data_torch, name, bid)):
                 data: np.ndarray  # type hint
                 n_dims = len(data.shape)
                 data_qtype: gguf.GGMLQuantizationType | bool = self.tensor_force_quant(name, new_name, bid, n_dims)
@@ -703,6 +707,9 @@ class Model:
         if chkhsh == "877081d19cf6996e2c4ff0e1236341e9b7bde288f5311a56a937f0afbbb3aeb5":
             # ref: https://huggingface.co/deepseek-ai/DeepSeek-V3
             res = "deepseek-v3"
+        if chkhsh == "b4b8ca1f9769494fbd956ebc4c249de6131fb277a4a3345a7a92c7dd7a55808d":
+            # ref: https://huggingface.co/ddh0/DeepSeek-V4-Flash-GGUF
+            res = "joyai-llm"
         if chkhsh == "d5f1dd6f980fec569fb218a81a7658ac45fc56b38c5a0adeb1c232fbe04ef5ec":
             # ref: https://huggingface.co/ByteDance-Seed/Seed-Coder-8B-Base
             res = "seed-coder"
@@ -4613,13 +4620,17 @@ class DeepseekV2Model(Model):
         super().set_gguf_parameters()
         hparams = self.hparams
 
-        self.gguf_writer.add_leading_dense_block_count(hparams["first_k_dense_replace"])
+        self.gguf_writer.add_leading_dense_block_count(hparams.get("first_k_dense_replace", 0))
         self.gguf_writer.add_vocab_size(hparams["vocab_size"])
         if "q_lora_rank" in hparams and hparams["q_lora_rank"] is not None:
             self.gguf_writer.add_q_lora_rank(hparams["q_lora_rank"])
-        self.gguf_writer.add_kv_lora_rank(hparams["kv_lora_rank"])
-        self.gguf_writer.add_key_length(hparams["qk_nope_head_dim"] + hparams["qk_rope_head_dim"])
-        self.gguf_writer.add_value_length(hparams["v_head_dim"])
+        if hparams.get("kv_lora_rank") is not None:
+            self.gguf_writer.add_kv_lora_rank(hparams["kv_lora_rank"])
+        key_length = hparams.get("head_dim")
+        if key_length is None:
+            key_length = hparams["qk_nope_head_dim"] + hparams["qk_rope_head_dim"]
+        self.gguf_writer.add_key_length(key_length)
+        self.gguf_writer.add_value_length(hparams.get("v_head_dim", key_length))
         self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"])
         self.gguf_writer.add_expert_count(hparams["n_routed_experts"])
         self.gguf_writer.add_expert_shared_count(hparams["n_shared_experts"])
@@ -4630,17 +4641,20 @@ class DeepseekV2Model(Model):
             self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)
         elif hparams["scoring_func"] == "softmax":
             self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SOFTMAX)
+        elif hparams["scoring_func"] == "sqrtsoftplus":
+            self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SQRTSOFTPLUS)
         else:
             raise ValueError(f"Unsupported scoring_func value: {hparams['scoring_func']}")
 
         self.gguf_writer.add_rope_dimension_count(hparams["qk_rope_head_dim"])
 
-        if self.hparams.get("rope_scaling") is not None and "factor" in self.hparams["rope_scaling"]:
-            if self.hparams["rope_scaling"].get("type") == "yarn":
+        rope_scaling = self.hparams.get("rope_scaling")
+        if rope_scaling is not None and "factor" in rope_scaling:
+            if rope_scaling.get("type") == "yarn" and "mscale_all_dim" in rope_scaling:
                 self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.YARN)
-                self.gguf_writer.add_rope_scaling_factor(self.hparams["rope_scaling"]["factor"])
-                self.gguf_writer.add_rope_scaling_orig_ctx_len(self.hparams["rope_scaling"]["original_max_position_embeddings"])
-                self.gguf_writer.add_rope_scaling_yarn_log_mul(0.1 * hparams["rope_scaling"]["mscale_all_dim"])
+                self.gguf_writer.add_rope_scaling_factor(rope_scaling["factor"])
+                self.gguf_writer.add_rope_scaling_orig_ctx_len(rope_scaling["original_max_position_embeddings"])
+                self.gguf_writer.add_rope_scaling_yarn_log_mul(0.1 * rope_scaling["mscale_all_dim"])
 
     _experts: list[dict[str, Tensor]] | None = None
 
@@ -4726,20 +4740,387 @@ class DeepseekV2Model(Model):
 @Model.register("DeepseekV4ProForCausalLM")
 class DeepseekV4Model(DeepseekV2Model):
     model_arch = gguf.MODEL_ARCH.DEEPSEEK4
+    supports_mtp_export = True
+    mtp_only = False
+
+    _mtp_projection_parts: dict[str, Tensor]
+    _mtp_experts: dict[tuple[int, str], dict[int, Tensor]]
+    _mtp_scales: dict[str, Tensor]
+    _mtp_pending_weights: dict[str, Tensor]
+    _mtp_expected_scales: set[str]
+    _mtp_raw_tensors: dict[str, Tensor]
 
     def __init__(self, *args, **kwargs):
+        type(self).mtp_only = bool(type(self).mtp_only)
         super().__init__(*args, **kwargs)
-        self.block_count = self.hparams["num_hidden_layers"] + self.hparams.get("num_nextn_predict_layers", 0)
+
+        main_layers = int(self.hparams["num_hidden_layers"])
+        nextn_layers = int(self.hparams.get("num_nextn_predict_layers", 0) or 0)
+        if self.mtp_only:
+            if nextn_layers != 1:
+                raise ValueError(f"DeepSeek-V4 MTP export requires one predictor layer, got {nextn_layers}")
+        self.block_count = main_layers + nextn_layers
         self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
 
+        self._mtp_projection_parts = {}
+        self._mtp_experts = {}
+        self._mtp_scales = {}
+        self._mtp_pending_weights = {}
+        self._mtp_expected_scales = set()
+        self._mtp_raw_tensors = {}
+
+    def _mtp_source_selection(self) -> tuple[set[str], list[str]]:
+        index_name = "model.safetensors.index.json" if self.is_safetensors else "pytorch_model.bin.index.json"
+        index_path = self.dir_model / index_name
+        def is_mtp_tensor(name: str) -> bool:
+            return (
+                name in {"embed.weight", "norm.weight", "head.weight", "head.scale"}
+                or name.startswith("mtp.0.")
+            )
+
+        if index_path.is_file():
+            with open(index_path, "r", encoding="utf-8") as f:
+                index: dict[str, Any] = json.load(f)
+            weight_map = index.get("weight_map")
+            if not isinstance(weight_map, dict):
+                raise ValueError(f"Can't load 'weight_map' from {index_name!r}")
+
+            selected = {name for name in weight_map if is_mtp_tensor(name)}
+            parts = sorted({str(weight_map[name]) for name in selected})
+            missing_parts = [name for name in parts if not (self.dir_model / name).is_file()]
+            if missing_parts:
+                raise FileNotFoundError(
+                    "DeepSeek-V4 MTP conversion requires missing index-derived shard(s): "
+                    + ", ".join(missing_parts)
+                )
+        else:
+            selected = set()
+            parts = []
+            for part_name in self.part_names:
+                if self.is_safetensors:
+                    from safetensors import safe_open
+                    with safe_open(self.dir_model / part_name, framework="pt", device="cpu") as model_part:
+                        part_selected = {name for name in model_part.keys() if is_mtp_tensor(name)}
+                else:
+                    model_part = torch.load(
+                        str(self.dir_model / part_name),
+                        map_location="cpu",
+                        mmap=True,
+                        weights_only=True,
+                    )
+                    part_selected = {name for name in model_part if is_mtp_tensor(name)}
+
+                if part_selected:
+                    selected.update(part_selected)
+                    parts.append(part_name)
+
+        required_roots = {"embed.weight", "norm.weight", "head.weight"}
+        missing_roots = required_roots - selected
+        if missing_roots:
+            source = index_name if index_path.is_file() else "model shards"
+            raise ValueError(
+                f"DeepSeek-V4 MTP conversion is missing root tensor(s) in {source}: {sorted(missing_roots)}"
+            )
+        if not selected:
+            raise ValueError("DeepSeek-V4 MTP conversion found no standalone MTP tensors in the model shards")
+
+        self._mtp_expected_scales = {
+            self._map_mtp_source_name(name).removesuffix(".scale") + ".weight"
+            for name in selected if name.endswith(".scale")
+        }
+        return selected, parts
+
+    def _map_mtp_source_name(self, name: str) -> str:
+        if not name.startswith("mtp."):
+            return name
+        parts = name.split(".", 2)
+        if len(parts) != 3 or not parts[1].isdecimal():
+            raise ValueError(f"Unexpected DeepSeek-V4 MTP tensor {name!r}")
+        if int(parts[1]) != 0:
+            raise ValueError(f"DeepSeek-V4 MTP export supports predictor 0 only, got {parts[1]}")
+
+        bid = int(self.hparams["num_hidden_layers"])
+        suffix = parts[2]
+        if suffix in {"hc_head_fn", "hc_head_base", "hc_head_scale"}:
+            return suffix
+        if suffix in {"e_proj.weight", "e_proj.scale", "h_proj.weight", "h_proj.scale"}:
+            return f"layers.{bid}.nextn.{suffix}"
+        if suffix == "enorm.weight":
+            return f"layers.{bid}.nextn.enorm.weight"
+        if suffix == "hnorm.weight":
+            return f"layers.{bid}.nextn.hnorm.weight"
+        if suffix == "norm.weight":
+            return f"layers.{bid}.nextn.shared_head_norm.weight"
+        return f"layers.{bid}.{suffix}"
+
+    def get_tensors(self) -> Iterator[tuple[str, Tensor]]:
+        if not self.mtp_only:
+            yield from super().get_tensors()
+            return
+
+        selected, parts = self._mtp_source_selection()
+        seen: set[str] = set()
+        for part_name in parts:
+            logger.info(f"gguf: loading selected MTP model part '{part_name}'")
+            ctx: ContextManager[Any]
+            if self.is_safetensors:
+                from safetensors import safe_open
+                ctx = cast(ContextManager[Any], safe_open(self.dir_model / part_name, framework="pt", device="cpu"))
+            else:
+                ctx = contextlib.nullcontext(torch.load(str(self.dir_model / part_name), map_location="cpu", mmap=True, weights_only=True))
+
+            with ctx as model_part:
+                for name in model_part.keys():
+                    if name not in selected:
+                        continue
+                    seen.add(name)
+                    if self.is_safetensors:
+                        if self.lazy:
+                            data = LazyTorchTensor.from_safetensors_slice(model_part.get_slice(name))
+                        else:
+                            data = model_part.get_tensor(name)
+                    else:
+                        data = model_part[name]
+                        if self.lazy:
+                            data = LazyTorchTensor.from_eager(data)
+                    yield self._map_mtp_source_name(name), data
+
+        if missing := selected - seen:
+            raise ValueError(f"DeepSeek-V4 MTP index names missing from selected model parts: {sorted(missing)}")
+
+    def _format_dsv4_tensor_name(self, key: gguf.MODEL_TENSOR, bid: int | None, suffix: str = ".weight") -> str:
+        return self.format_tensor_name(key, bid, suffix)
+
+    def _map_dsv4_tensor_name(self, name: str, bid: int | None) -> tuple[gguf.MODEL_TENSOR, str]:
+        root_map: dict[str, tuple[gguf.MODEL_TENSOR, str]] = {
+            "embed.weight": (gguf.MODEL_TENSOR.TOKEN_EMBD, ".weight"),
+            "norm.weight": (gguf.MODEL_TENSOR.OUTPUT_NORM, ".weight"),
+            "head.weight": (gguf.MODEL_TENSOR.OUTPUT, ".weight"),
+            "hc_head_fn": (gguf.MODEL_TENSOR.HC_HEAD_FN, ".weight"),
+            "hc_head_base": (gguf.MODEL_TENSOR.HC_HEAD_BASE, ".weight"),
+            "hc_head_scale": (gguf.MODEL_TENSOR.HC_HEAD_SCALE, ".weight"),
+        }
+        if name in root_map:
+            return root_map[name]
+
+        match = re.match(r"layers\.(\d+)\.(.+)$", name)
+        if match is None:
+            raise ValueError(f"Unsupported DeepSeek-V4 tensor {name!r}")
+        layer = int(match.group(1))
+        if bid != layer:
+            raise ValueError(f"Tensor {name!r} parsed bid {bid} but layer name has {layer}")
+        layer_name = match.group(2)
+
+        layer_map: dict[str, tuple[gguf.MODEL_TENSOR, str]] = {
+            "hc_attn_fn": (gguf.MODEL_TENSOR.HC_ATTN_FN, ".weight"),
+            "hc_attn_base": (gguf.MODEL_TENSOR.HC_ATTN_BASE, ".weight"),
+            "hc_attn_scale": (gguf.MODEL_TENSOR.HC_ATTN_SCALE, ".weight"),
+            "hc_ffn_fn": (gguf.MODEL_TENSOR.HC_FFN_FN, ".weight"),
+            "hc_ffn_base": (gguf.MODEL_TENSOR.HC_FFN_BASE, ".weight"),
+            "hc_ffn_scale": (gguf.MODEL_TENSOR.HC_FFN_SCALE, ".weight"),
+            "attn.attn_sink": (gguf.MODEL_TENSOR.ATTN_SINKS, ".weight"),
+            "attn.wq_a.weight": (gguf.MODEL_TENSOR.ATTN_Q_A, ".weight"),
+            "attn.wq_b.weight": (gguf.MODEL_TENSOR.ATTN_Q_B, ".weight"),
+            "attn.q_norm.weight": (gguf.MODEL_TENSOR.ATTN_Q_A_NORM, ".weight"),
+            "attn.wkv.weight": (gguf.MODEL_TENSOR.ATTN_KV, ".weight"),
+            "attn.kv_norm.weight": (gguf.MODEL_TENSOR.ATTN_KV_NORM, ".weight"),
+            "attn.wo_a.weight": (gguf.MODEL_TENSOR.ATTN_OUT_A, ".weight"),
+            "attn.wo_b.weight": (gguf.MODEL_TENSOR.ATTN_OUT_B, ".weight"),
+            "attn_norm.weight": (gguf.MODEL_TENSOR.ATTN_NORM, ".weight"),
+            "ffn_norm.weight": (gguf.MODEL_TENSOR.FFN_NORM, ".weight"),
+            "ffn.gate.weight": (gguf.MODEL_TENSOR.FFN_GATE_INP, ".weight"),
+            "ffn.gate.bias": (gguf.MODEL_TENSOR.FFN_EXP_PROBS_B, ".bias"),
+            "ffn.shared_experts.w1.weight": (gguf.MODEL_TENSOR.FFN_GATE_SHEXP, ".weight"),
+            "ffn.shared_experts.w2.weight": (gguf.MODEL_TENSOR.FFN_DOWN_SHEXP, ".weight"),
+            "ffn.shared_experts.w3.weight": (gguf.MODEL_TENSOR.FFN_UP_SHEXP, ".weight"),
+            "nextn.eh_proj.weight": (gguf.MODEL_TENSOR.NEXTN_EH_PROJ, ".weight"),
+            "nextn.enorm.weight": (gguf.MODEL_TENSOR.NEXTN_ENORM, ".weight"),
+            "nextn.hnorm.weight": (gguf.MODEL_TENSOR.NEXTN_HNORM, ".weight"),
+            "nextn.shared_head_norm.weight": (gguf.MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ".weight"),
+        }
+        if layer_name in layer_map:
+            return layer_map[layer_name]
+        raise ValueError(f"Unsupported DeepSeek-V4 tensor {name!r}")
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if not self.mtp_only:
+            return super().modify_tensors(data_torch, name, bid)
+
+        if name == "head.scale":
+            return []
+
+        expert_match = re.match(r"layers\.(\d+)\.ffn\.experts\.(\d+)\.w([123])\.(weight|scale)$", name)
+
+        if name.endswith(".scale"):
+            weight_name = name.removesuffix(".scale") + ".weight"
+            self._mtp_scales[weight_name] = data_torch
+            if weight_name not in self._mtp_pending_weights:
+                return []
+            weight = self._mtp_pending_weights.pop(weight_name)
+            scale = self._mtp_scales.pop(weight_name)
+            if expert_match is not None:
+                return self._record_mtp_expert(weight_name, self._pack_mxfp4_blocks(weight, scale))
+            data_torch = self._dequantize_mtp_weight(weight, scale)
+            name = weight_name
+        elif name.endswith(".weight") and name in self._mtp_expected_scales:
+            if name not in self._mtp_scales:
+                self._mtp_pending_weights[name] = data_torch
+                return []
+            scale = self._mtp_scales.pop(name)
+            if expert_match is not None:
+                return self._record_mtp_expert(name, self._pack_mxfp4_blocks(data_torch, scale))
+            data_torch = self._dequantize_mtp_weight(data_torch, scale)
+
+        if data_torch.dtype not in (torch.float16, torch.float32):
+            data_torch = data_torch.to(torch.float32)
+
+        # A source scale is stored as one scale per 128x128 FP8 tile, or as
+        # one scale per 16-column I8 expert block. Normalize both forms here;
+        # the GGUF writer then applies the requested output quantization.
+
+        projection_match = re.match(r"layers\.(\d+)\.nextn\.(e_proj|h_proj)\.weight$", name)
+        if projection_match:
+            self._mtp_projection_parts[name] = data_torch
+            layer = int(projection_match.group(1))
+            e_name = f"layers.{layer}.nextn.e_proj.weight"
+            h_name = f"layers.{layer}.nextn.h_proj.weight"
+            if e_name not in self._mtp_projection_parts or h_name not in self._mtp_projection_parts:
+                return []
+            e_proj = self._mtp_projection_parts.pop(e_name)
+            h_proj = self._mtp_projection_parts.pop(h_name)
+            out_name = self._format_dsv4_tensor_name(gguf.MODEL_TENSOR.NEXTN_EH_PROJ, layer)
+            return [(out_name, torch.cat((e_proj, h_proj), dim=1).contiguous())]
+
+        expert_match = re.match(r"layers\.(\d+)\.ffn\.experts\.(\d+)\.w([123])\.weight$", name)
+        if expert_match:
+            layer = int(expert_match.group(1))
+            expert = int(expert_match.group(2))
+            proj = expert_match.group(3)
+            key = (layer, proj)
+            self._mtp_experts.setdefault(key, {})[expert] = data_torch
+            n_experts = int(self.hparams.get("n_routed_experts", self.hparams.get("num_experts", 0)) or 0)
+            if n_experts <= 0 or len(self._mtp_experts[key]) < n_experts:
+                return []
+            experts = self._mtp_experts.pop(key)
+            if set(experts) != set(range(n_experts)):
+                raise ValueError(f"Incomplete DeepSeek-V4 MTP expert set for layer {layer}, projection w{proj}")
+            stacked = torch.stack([experts[eid] for eid in range(n_experts)], dim=0)
+            expert_key = {
+                "1": gguf.MODEL_TENSOR.FFN_GATE_EXP,
+                "2": gguf.MODEL_TENSOR.FFN_DOWN_EXP,
+                "3": gguf.MODEL_TENSOR.FFN_UP_EXP,
+            }[proj]
+            return [(self._format_dsv4_tensor_name(expert_key, layer), stacked)]
+
+        tensor_key, suffix = self._map_dsv4_tensor_name(name, bid)
+        return [(self._format_dsv4_tensor_name(tensor_key, bid, suffix), data_torch)]
+
+    def _record_mtp_expert(self, weight_name: str, packed: Tensor) -> Iterable[tuple[str, Tensor]]:
+        match = re.match(r"layers\.(\d+)\.ffn\.experts\.(\d+)\.w([123])\.weight$", weight_name)
+        if match is None:
+            raise ValueError(f"Unexpected packed DeepSeek-V4 expert tensor {weight_name!r}")
+
+        layer = int(match.group(1))
+        expert = int(match.group(2))
+        proj = match.group(3)
+        key = (layer, proj)
+        self._mtp_experts.setdefault(key, {})[expert] = packed
+        n_experts = int(self.hparams.get("n_routed_experts", self.hparams.get("num_experts", 0)) or 0)
+        if n_experts <= 0 or len(self._mtp_experts[key]) < n_experts:
+            return []
+        experts = self._mtp_experts.pop(key)
+        if set(experts) != set(range(n_experts)):
+            raise ValueError(f"Incomplete DeepSeek-V4 MTP expert set for layer {layer}, projection w{proj}")
+
+        expert_key = {
+            "1": gguf.MODEL_TENSOR.FFN_GATE_EXP,
+            "2": gguf.MODEL_TENSOR.FFN_DOWN_EXP,
+            "3": gguf.MODEL_TENSOR.FFN_UP_EXP,
+        }[proj]
+        out_name = self._format_dsv4_tensor_name(expert_key, layer)
+        self._mtp_raw_tensors[out_name] = torch.stack([experts[eid] for eid in range(n_experts)], dim=0).contiguous()
+        return []
+
+    @staticmethod
+    def _pack_mxfp4_blocks(weight: Tensor, scale: Tensor) -> Tensor:
+        weight = LazyTorchTensor.to_eager(weight)
+        scale = LazyTorchTensor.to_eager(scale)
+        packed = weight.contiguous().view(torch.uint8)
+        scale_u8 = scale.contiguous().view(torch.uint8)
+
+        out_features, packed_cols = packed.shape
+        logical_cols = packed_cols * 2
+        if logical_cols % 32 != 0:
+            raise ValueError(f"MXFP4 source row has {logical_cols} values, expected a multiple of 32")
+
+        n_blocks = logical_cols // 32
+        if tuple(scale_u8.shape) != (out_features, n_blocks):
+            raise ValueError(f"MXFP4 scale shape {tuple(scale_u8.shape)} does not match {(out_features, n_blocks)}")
+
+        src = packed.reshape(out_features, n_blocks, 16)
+        low = src & 0x0F
+        high = (src >> 4) & 0x0F
+        vals = torch.stack((low, high), dim=-1).reshape(out_features, n_blocks, 32)
+        qs = vals[:, :, :16] | (vals[:, :, 16:] << 4)
+        raw = torch.cat((scale_u8.unsqueeze(-1), qs.to(torch.uint8)), dim=-1)
+        return raw.reshape(out_features, n_blocks * 17)
+
+    @staticmethod
+    def _dequantize_mtp_weight(weight: Tensor, scale: Tensor) -> Tensor:
+        if len(weight.shape) != 2 or len(scale.shape) != 2:
+            return weight
+
+        scale_f = scale.view(torch.uint8).float()
+        scale_f = torch.exp2(scale_f - 127.0)
+        row_repeat = max(1, (weight.shape[0] + scale.shape[0] - 1) // scale.shape[0])
+        col_repeat = max(1, (weight.shape[1] + scale.shape[1] - 1) // scale.shape[1])
+        scale_f = scale_f.repeat_interleave(row_repeat, 0)[:weight.shape[0]]
+        scale_f = scale_f.repeat_interleave(col_repeat, 1)[:, :weight.shape[1]]
+        return weight.float() * scale_f
+
+    def prepare_tensors(self):
+        super().prepare_tensors()
+        if self.mtp_only:
+            for name, data in self._mtp_raw_tensors.items():
+                self.gguf_writer.add_tensor(
+                    name,
+                    data.cpu().numpy(),
+                    raw_dtype=gguf.GGMLQuantizationType.MXFP4,
+                )
+            if self._mtp_projection_parts:
+                raise ValueError(f"Unpaired DeepSeek-V4 MTP projection tensors: {sorted(self._mtp_projection_parts)}")
+            if self._mtp_experts:
+                raise ValueError(f"Unprocessed DeepSeek-V4 MTP expert tensors: {sorted(self._mtp_experts)}")
+            if self._mtp_pending_weights:
+                raise ValueError(f"Unpaired DeepSeek-V4 MTP weight scales: {sorted(self._mtp_pending_weights)}")
+
     def set_gguf_parameters(self):
         super().set_gguf_parameters()
+        hparams = self.hparams
+        arch = gguf.MODEL_ARCH_NAMES[self.model_arch]
 
-        if (indexer_heads := self.hparams.get("num_indexer_heads")) is not None:
+        # DeepSeek-V4 structural metadata is required even before its graph exists.
+        self.gguf_writer.add_array(f"{arch}.swiglu_clamp_exp", [float(hparams["swiglu_limit"])] * self.block_count)
+        self.gguf_writer.add_array(f"{arch}.swiglu_clamp_shexp", [float(hparams["swiglu_limit"])] * self.block_count)
+        self.gguf_writer.add_sliding_window(int(hparams["sliding_window"]))
+        self.gguf_writer.add_uint32(f"{arch}.attention.output_group_count", int(hparams["o_groups"]))
+        self.gguf_writer.add_uint32(f"{arch}.attention.output_lora_rank", int(hparams["o_lora_rank"]))
+        self.gguf_writer.add_array(f"{arch}.attention.compress_ratios", hparams["compress_ratios"])
+        self.gguf_writer.add_float32(f"{arch}.attention.compress_rope_freq_base", float(hparams["compress_rope_theta"]))
+        self.gguf_writer.add_uint32(f"{arch}.hyper_connection.count", int(hparams["hc_mult"]))
+        self.gguf_writer.add_uint32(f"{arch}.hyper_connection.sinkhorn_iterations", int(hparams["hc_sinkhorn_iters"]))
+        self.gguf_writer.add_float32(f"{arch}.hyper_connection.epsilon", float(hparams["hc_eps"]))
+        if (hash_layers := hparams.get("num_hash_layers")) is not None:
+            self.gguf_writer.add_uint32(f"{arch}.hash_layer_count", int(hash_layers))
+
+        if self.mtp_only:
+            self.gguf_writer.add_embedding_length_out(int(hparams["hidden_size"]) * int(hparams["hc_mult"]))
+
+        if (indexer_heads := self.hparams.get("num_indexer_heads", self.hparams.get("index_n_heads"))) is not None:
             self.gguf_writer.add_attention_indexer_head_count(indexer_heads)
-        if (indexer_dim := self.hparams.get("indexer_head_dim")) is not None:
+        if (indexer_dim := self.hparams.get("indexer_head_dim", self.hparams.get("index_head_dim"))) is not None:
             self.gguf_writer.add_attention_indexer_key_length(indexer_dim)
-        if (indexer_top_k := self.hparams.get("indexer_topk")) is not None:
+        if (indexer_top_k := self.hparams.get("indexer_topk", self.hparams.get("index_topk"))) is not None:
             self.gguf_writer.add_attention_indexer_top_k(indexer_top_k)
         if (nextn_layers := self.hparams.get("num_nextn_predict_layers")) is not None:
             self.gguf_writer.add_nextn_predict_layers(nextn_layers)
@@ -5954,6 +6335,7 @@ class LazyTorchTensor(gguf.LazyBase):
         "BOOL": torch.bool,
         "F8_E4M3": torch.float8_e4m3fn,
         "F8_E5M2": torch.float8_e5m2,
+        "F8_E8M0": getattr(torch, "float8_e8m0fnu", torch.uint8),
     }
 
     def numpy(self) -> gguf.LazyNumpyTensor:
@@ -6051,6 +6433,10 @@ def parse_args() -> argparse.Namespace:
         "--target-model-dir", type=Path,
         help="matching target model directory; required for DFlash conversion to reuse tokenizer and infer target feature width",
     )
+    parser.add_argument(
+        "--mtp", action="store_true",
+        help="export a standalone DeepSeek-V4 MTP predictor companion",
+    )
 
     return parser.parse_args()
 
@@ -6124,6 +6510,11 @@ def main() -> None:
             logger.error(f"Model {model_architecture} is not supported")
             sys.exit(1)
 
+        if args.mtp:
+            if not getattr(model_class, "supports_mtp_export", False):
+                raise ValueError(f"Architecture {model_architecture!r} does not support standalone MTP export")
+            model_class.mtp_only = True
+
         model_instance = model_class(dir_model=dir_model, ftype=output_type, fname_out=fname_out,
                                      is_big_endian=args.bigendian, use_temp_file=args.use_temp_file,
                                      eager=args.no_lazy,

gguf-py/gguf/constants.py

@@ -78,6 +78,7 @@ class Keys:
         VOCAB_SIZE                        = "{arch}.vocab_size"
         CONTEXT_LENGTH                    = "{arch}.context_length"
         EMBEDDING_LENGTH                  = "{arch}.embedding_length"
+        EMBEDDING_LENGTH_OUT              = "{arch}.embedding_length_out"
         BLOCK_COUNT                       = "{arch}.block_count"
         LEADING_DENSE_BLOCK_COUNT         = "{arch}.leading_dense_block_count"
         FEED_FORWARD_LENGTH               = "{arch}.feed_forward_length"
@@ -291,6 +292,9 @@ class MODEL_TENSOR(IntEnum):
     POS_EMBD             = auto()
     OUTPUT               = auto()
     OUTPUT_NORM          = auto()
+    HC_HEAD_FN           = auto()
+    HC_HEAD_BASE         = auto()
+    HC_HEAD_SCALE        = auto()
     ROPE_FREQS           = auto()
     ROPE_FACTORS_LONG    = auto()
     ROPE_FACTORS_SHORT   = auto()
@@ -399,6 +403,16 @@ class MODEL_TENSOR(IntEnum):
     DFLASH_FC            = auto()
     DFLASH_HIDDEN_NORM   = auto()
     DFLASH_AUX_HIDDEN_NORM = auto()
+    ATTN_KV              = auto()
+    ATTN_KV_NORM         = auto()
+    ATTN_OUT_A           = auto()
+    ATTN_OUT_B           = auto()
+    HC_ATTN_FN           = auto()
+    HC_ATTN_BASE         = auto()
+    HC_ATTN_SCALE        = auto()
+    HC_FFN_FN            = auto()
+    HC_FFN_BASE          = auto()
+    HC_FFN_SCALE         = auto()
     # openPangu-2.0 (MoME causal-conv on MLA latents)
     ATTN_QA_CONV         = auto()
     ATTN_KV_CONV         = auto()   # compresskv_conv
@@ -496,6 +510,9 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
     MODEL_TENSOR.POS_EMBD:             "position_embd",
     MODEL_TENSOR.OUTPUT_NORM:          "output_norm",
     MODEL_TENSOR.OUTPUT:               "output",
+    MODEL_TENSOR.HC_HEAD_FN:           "output_hc_fn",
+    MODEL_TENSOR.HC_HEAD_BASE:         "output_hc_base",
+    MODEL_TENSOR.HC_HEAD_SCALE:        "output_hc_scale",
     MODEL_TENSOR.ROPE_FREQS:           "rope_freqs",
     MODEL_TENSOR.ROPE_FACTORS_LONG:    "rope_factors_long",
     MODEL_TENSOR.ROPE_FACTORS_SHORT:   "rope_factors_short",
@@ -557,6 +574,16 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
     MODEL_TENSOR.ATTN_V_B:             "blk.{bid}.attn_v_b",
     MODEL_TENSOR.ATTN_Q_A_NORM:        "blk.{bid}.attn_q_a_norm",
     MODEL_TENSOR.ATTN_KV_A_NORM:       "blk.{bid}.attn_kv_a_norm",
+    MODEL_TENSOR.ATTN_KV:              "blk.{bid}.attn_kv",
+    MODEL_TENSOR.ATTN_KV_NORM:         "blk.{bid}.attn_kv_a_norm",
+    MODEL_TENSOR.ATTN_OUT_A:           "blk.{bid}.attn_output_a",
+    MODEL_TENSOR.ATTN_OUT_B:           "blk.{bid}.attn_output_b",
+    MODEL_TENSOR.HC_ATTN_FN:           "blk.{bid}.hc_attn_fn",
+    MODEL_TENSOR.HC_ATTN_BASE:         "blk.{bid}.hc_attn_base",
+    MODEL_TENSOR.HC_ATTN_SCALE:        "blk.{bid}.hc_attn_scale",
+    MODEL_TENSOR.HC_FFN_FN:            "blk.{bid}.hc_ffn_fn",
+    MODEL_TENSOR.HC_FFN_BASE:          "blk.{bid}.hc_ffn_base",
+    MODEL_TENSOR.HC_FFN_SCALE:         "blk.{bid}.hc_ffn_scale",
     MODEL_TENSOR.ATTN_SUB_NORM:        "blk.{bid}.attn_sub_norm",
     MODEL_TENSOR.FFN_SUB_NORM:         "blk.{bid}.ffn_sub_norm",
     MODEL_TENSOR.DEC_ATTN_NORM:        "dec.blk.{bid}.attn_norm",
@@ -1324,6 +1351,34 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.TOKEN_EMBD,
         MODEL_TENSOR.OUTPUT_NORM,
         MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.HC_HEAD_FN,
+        MODEL_TENSOR.HC_HEAD_BASE,
+        MODEL_TENSOR.HC_HEAD_SCALE,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_SINKS,
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_KV,
+        MODEL_TENSOR.ATTN_KV_NORM,
+        MODEL_TENSOR.ATTN_OUT_A,
+        MODEL_TENSOR.ATTN_OUT_B,
+        MODEL_TENSOR.HC_ATTN_FN,
+        MODEL_TENSOR.HC_ATTN_BASE,
+        MODEL_TENSOR.HC_ATTN_SCALE,
+        MODEL_TENSOR.HC_FFN_FN,
+        MODEL_TENSOR.HC_FFN_BASE,
+        MODEL_TENSOR.HC_FFN_SCALE,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
         MODEL_TENSOR.ROPE_FREQS,
         MODEL_TENSOR.ATTN_NORM,
         MODEL_TENSOR.ATTN_Q,
@@ -1934,8 +1989,9 @@ class GGMLQuantizationType(IntEnum):
 
 
 class ExpertGatingFuncType(IntEnum):
-    SOFTMAX  = 1
-    SIGMOID  = 2
+    SOFTMAX       = 1
+    SIGMOID       = 2
+    SQRTSOFTPLUS  = 4
 
 
 # TODO: add GGMLFileType from ggml_ftype in ggml.h

gguf-py/gguf/gguf_writer.py

@@ -634,6 +634,9 @@ class GGUFWriter:
     def add_embedding_length(self, length: int) -> None:
         self.add_uint32(Keys.LLM.EMBEDDING_LENGTH.format(arch=self.arch), length)
 
+    def add_embedding_length_out(self, length: int) -> None:
+        self.add_uint32(Keys.LLM.EMBEDDING_LENGTH_OUT.format(arch=self.arch), length)
+
     def add_features_length(self, length: int) -> None:
         self.add_uint32(Keys.LLM.FEATURES_LENGTH.format(arch=self.arch), length)
 

src/graphs/build_deepseek4.cpp

@@ -1239,9 +1239,7 @@ static ggml_tensor * ds4_attention(ggml_cgraph * gf, ggml_context * ctx0, llm_bu
 ggml_cgraph * llm_build_context::build_deepseek4() {
     ggml_cgraph * gf = new_graph_custom();
 
-    if (lctx.cparams.mtp_op_type != MTP_OP_NONE) {
-        GGML_ABORT("DeepSeek4 MTP execution is not implemented");
-    }
+    const bool is_mtp = lctx.cparams.mtp_op_type != MTP_OP_NONE;
 
     const int64_t n_embd_head = hparams.n_embd_head_k(0);
     const int64_t n_embd_head_rope = hparams.n_rot;
@@ -1258,19 +1256,59 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
     dsv4_build_plan_inputs(ctx0, lctx.dsv4.inputs.hca, lctx.dsv4.hca_plan, "dsv4_hca", n_tokens, true, lctx.cparams.flash_attn);
     dsv4_build_plan_inputs(ctx0, lctx.dsv4.inputs.lid, lctx.dsv4.lid_plan, "dsv4_lid", n_tokens, false, lctx.cparams.flash_attn);
 
-    ggml_tensor * inp = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb);
     ggml_tensor * inp_pos = build_inp_pos();
     ggml_tensor * KQ_mask = hparams.n_swa > 0 ? build_inp_KQ_mask_swa() : build_inp_KQ_mask();
-    ggml_tensor * inpL = ggml_reshape_3d(ctx0, inp, n_embd, 1, n_tokens);
-    inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1);
-    cb(inpL, "hc_init", -1);
+    ggml_tensor * inpL = nullptr;
 
     ggml_tensor * append_csa_state = nullptr;
     ggml_tensor * append_csa_score = nullptr;
     ggml_tensor * append_lid_state = nullptr;
     ggml_tensor * append_lid_score = nullptr;
 
-    for (int il = 0; il < n_layer; ++il) {
+    if (is_mtp) {
+        GGML_ASSERT(model.mtp && hparams.nextn_predict_layers == 1);
+        GGML_ASSERT(n_layer > hparams.nextn_predict_layers);
+
+        const int64_t n_hidden = n_embd * hc;
+        ggml_tensor * hidden_state = nullptr;
+        if (lctx.cparams.mtp_op_type == MTP_OP_WARMUP || lctx.cparams.mtp_op_type == MTP_OP_UPDATE_ACCEPTED) {
+            hidden_state = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_hidden, n_tokens);
+        } else {
+            hidden_state = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_hidden);
+        }
+        ggml_set_name(hidden_state, "inp_mtp_states");
+        ggml_set_input(hidden_state);
+        lctx.inp_mtp_states = hidden_state;
+
+        ggml_tensor * tok_embd = build_inp_embd_mtp(model.tok_embd);
+        const int il_mtp = n_layer - hparams.nextn_predict_layers;
+        const auto & mtp_layer = model.layers[il_mtp];
+
+        ggml_tensor * h_state = ggml_reshape_3d(ctx0, hidden_state, n_embd, hc, n_tokens);
+        cb(h_state, "mtp_h_state", il_mtp);
+        ggml_tensor * h_norm = llm_build_norm(ctx0, h_state, hparams, mtp_layer.nextn.hnorm,
+                nullptr, LLM_NORM_RMS, cb, il_mtp);
+        cb(h_norm, "mtp_hnorm", il_mtp);
+
+        ggml_tensor * e_norm = llm_build_norm(ctx0, tok_embd, hparams, mtp_layer.nextn.enorm,
+                nullptr, LLM_NORM_RMS, cb, il_mtp);
+        e_norm = ggml_reshape_3d(ctx0, e_norm, n_embd, 1, n_tokens);
+        e_norm = ggml_repeat_4d(ctx0, e_norm, n_embd, hc, n_tokens, 1);
+        cb(e_norm, "mtp_enorm", il_mtp);
+
+        ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, 0);
+        cb(concat, "mtp_concat", il_mtp);
+        inpL = llm_build_lora_mm(lctx, ctx0, mtp_layer.nextn.eh_proj, concat);
+        cb(inpL, "mtp_eh_proj", il_mtp);
+    } else {
+        ggml_tensor * inp = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb);
+        inpL = ggml_reshape_3d(ctx0, inp, n_embd, 1, n_tokens);
+        inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1);
+        cb(inpL, "hc_init", -1);
+    }
+
+    const int n_layer_begin = is_mtp ? n_layer - hparams.nextn_predict_layers : 0;
+    for (int il = n_layer_begin; il < n_layer; ++il) {
 
         auto cur = ds4_attention(gf, ctx0, *this, inpL,
                              &append_csa_state, &append_csa_score,
@@ -1386,6 +1424,41 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
         cb(inpL, "l_out", il);
     }
 
+    if (is_mtp) {
+        const int il_mtp = n_layer - hparams.nextn_predict_layers;
+        ggml_tensor * inp_out_ids = build_inp_out_ids();
+        ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
+        ggml_tensor * h_nextn = ggml_get_rows(ctx0, flat, inp_out_ids);
+        cb(h_nextn, "result_mtp_embd", -1);
+        ggml_set_output(h_nextn);
+        ggml_build_forward_expand(gf, h_nextn);
+
+        inpL = ggml_reshape_3d(ctx0, h_nextn, n_embd, hc, n_outputs);
+        ggml_tensor * out = build_hc_head(ctx0, *this, hparams, n_embd, hparams.f_norm_rms_eps,
+                inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base);
+        cb(out, "mtp_hc_head", -1);
+
+        ggml_tensor * head_norm = model.layers[il_mtp].nextn.shared_head_norm
+                ? model.layers[il_mtp].nextn.shared_head_norm : model.output_norm;
+        GGML_ASSERT(head_norm != nullptr);
+        out = llm_build_norm(ctx0, out, hparams, head_norm, nullptr, LLM_NORM_RMS, cb, -1);
+        cb(out, "mtp_shared_head_norm", -1);
+
+        out = build_output(lctx, ctx0, out, model.output, nullptr, cb);
+        cb(out, "result_output", -1);
+        ggml_build_forward_expand(gf, out);
+        return gf;
+    }
+
+    if (lctx.cparams.mtp && (hparams.nextn_predict_layers > 0 || model.arch == LLM_ARCH_DEEPSEEK4)) {
+        ggml_tensor * inp_out_ids = build_inp_out_ids();
+        ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
+        ggml_tensor * h_nextn = ggml_get_rows(ctx0, flat, inp_out_ids);
+        cb(h_nextn, "result_mtp_embd", -1);
+        ggml_set_output(h_nextn);
+        ggml_build_forward_expand(gf, h_nextn);
+    }
+
     if (n_outputs != n_tokens) {
         ggml_tensor * inp_out_ids = build_inp_out_ids();
         ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);

src/llama-arch.cpp

@@ -127,6 +127,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
     { LLM_KV_VOCAB_SIZE,                        "%s.vocab_size"                        },
     { LLM_KV_CONTEXT_LENGTH,                    "%s.context_length"                    },
     { LLM_KV_EMBEDDING_LENGTH,                  "%s.embedding_length"                  },
+    { LLM_KV_EMBEDDING_LENGTH_OUT,              "%s.embedding_length_out"              },
     { LLM_KV_EMBEDDING_LENGTH_PER_LAYER,        "%s.embedding_length_per_layer_input"  },
     { LLM_KV_BLOCK_COUNT,                       "%s.block_count"                       },
     { LLM_KV_LEADING_DENSE_BLOCK_COUNT,         "%s.leading_dense_block_count"         },

src/llama-arch.h

@@ -106,6 +106,7 @@ enum llm_kv {
     LLM_KV_VOCAB_SIZE,
     LLM_KV_CONTEXT_LENGTH,
     LLM_KV_EMBEDDING_LENGTH,
+    LLM_KV_EMBEDDING_LENGTH_OUT,
     LLM_KV_BLOCK_COUNT,
     LLM_KV_LEADING_DENSE_BLOCK_COUNT,
     LLM_KV_FEED_FORWARD_LENGTH,

src/llama-dsv4.cpp

@@ -1664,6 +1664,31 @@ bool llama_prepare_dsv4_graph_inputs(llama_context & lctx, const llama_batch & b
         return false;
     }
 
+    // Standalone companions contain only the predictor block, skip target state planning.
+    const bool is_dsv4_mtp = lctx.model.mtp &&
+        lctx.cparams.mtp_op_type != MTP_OP_NONE &&
+        lctx.model.hparams.nextn_predict_layers > 0 &&
+        lctx.model.hparams.dsv4_compress_ratios[(size_t) (lctx.model.hparams.n_layer - lctx.model.hparams.nextn_predict_layers)] == 0;
+    if (is_dsv4_mtp) {
+        lctx.dsv4.raw = {};
+        if (!reserve_plan && !dsv4_build_raw_context(lctx, batch, lctx.dsv4.raw)) {
+            return false;
+        }
+        lctx.dsv4.csa_plan = {};
+        lctx.dsv4.hca_plan = {};
+        lctx.dsv4.lid_plan = {};
+        lctx.dsv4.csa_ctx = {};
+        lctx.dsv4.hca_ctx = {};
+        lctx.dsv4.lid_ctx = {};
+
+        if (set_tensors) {
+            dsv4_set_input_tensor(lctx.dsv4.inputs.raw_k_write_src_idxs, lctx.dsv4.raw.write_src_idxs);
+            dsv4_set_input_tensor(lctx.dsv4.inputs.raw_k_write_idxs, lctx.dsv4.raw.write_dst_idxs);
+            dsv4_set_input_tensor(lctx.dsv4.inputs.raw_k_read_idxs, lctx.dsv4.raw.read_dst_idxs);
+        }
+        return true;
+    }
+
     if (!lctx.ensure_dsv4_cache_tensors()) {
         return false;
     }

src/llama-hparams.cpp

@@ -161,6 +161,10 @@ void llm_load_hparams(
 
     ml.get_key(LLM_KV_CONTEXT_LENGTH,    hparams.n_ctx_train);
     ml.get_key(LLM_KV_EMBEDDING_LENGTH,  hparams.n_embd);
+    ml.get_key(LLM_KV_EMBEDDING_LENGTH_OUT, hparams.n_embd_out, false);
+    if (hparams.n_embd_out == 0) {
+        hparams.n_embd_out = hparams.n_embd;
+    }
     ml.get_key(LLM_KV_EXPERT_COUNT,      hparams.n_expert,      false);
     ml.get_key(LLM_KV_EXPERT_USED_COUNT, hparams.n_expert_used, false);
 
@@ -1659,6 +1663,12 @@ void llm_load_hparams(
         case LLM_ARCH_DEEPSEEK4:
         case LLM_ARCH_GLM_DSA:
             {
+                ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false);
+                // Probe the first appended predictor block, or n_layer - 1 for base GGUFs.
+                const uint32_t dsv4_probe_offset = std::max<uint32_t>(1, hparams.nextn_predict_layers);
+                const uint32_t dsv4_probe_layer = hparams.n_layer > dsv4_probe_offset
+                    ? hparams.n_layer - dsv4_probe_offset
+                    : 0;
                 ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH,     hparams.n_ff_exp);
                 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS,    hparams.f_norm_rms_eps);
                 // GLM-DSA lightning-indexer k_norm is a (non-RMS) LayerNorm built via LLM_NORM,
@@ -1692,16 +1702,17 @@ void llm_load_hparams(
                 ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK,      hparams.n_lora_q);
                 ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK,     hparams.n_lora_kv, false);
                 if (model.arch == LLM_ARCH_DEEPSEEK4 && hparams.n_lora_kv == 0) {
-                    if (auto * kv_norm = ml.get_tensor_meta("blk.0.attn_kv_a_norm.weight")) {
+                    const uint32_t probe_layer = dsv4_probe_layer;
+                    if (auto * kv_norm = ml.get_tensor_meta(format("blk.%u.attn_kv_a_norm.weight", probe_layer).c_str())) {
                         hparams.n_lora_kv = (uint32_t) kv_norm->ne[0];
-                    } else if (auto * kv = ml.get_tensor_meta("blk.0.attn_kv.weight")) {
+                    } else if (auto * kv = ml.get_tensor_meta(format("blk.%u.attn_kv.weight", probe_layer).c_str())) {
                         const int64_t kv_inner = kv->ne[0] == hparams.n_embd ? kv->ne[1] : kv->ne[0];
                         hparams.n_lora_kv = (uint32_t) kv_inner;
                     } else {
-                        auto * kv_a = ml.get_tensor_meta("blk.0.attn_kv_latent.weight");
+                        auto * kv_a = ml.get_tensor_meta(format("blk.%u.attn_kv_latent.weight", probe_layer).c_str());
                         bool subtract_rope = false;
                         if (kv_a == nullptr) {
-                            kv_a = ml.require_tensor_meta("blk.0.attn_kv_a_mqa.weight");
+                            kv_a = ml.require_tensor_meta(format("blk.%u.attn_kv_a_mqa.weight", probe_layer).c_str());
                             subtract_rope = true;
                         }
 
@@ -1740,8 +1751,9 @@ void llm_load_hparams(
                     }
 
                     const auto * hc_head_base = ml.get_tensor_meta("hc_head_base");
-                    const auto * wo_a_0 = ml.get_tensor_meta("blk.0.attn_output_a.weight");
-                    const auto * wo_b_0 = ml.get_tensor_meta("blk.0.attn_output_b.weight");
+                    const uint32_t probe_layer = dsv4_probe_layer;
+                    const auto * wo_a_0 = ml.get_tensor_meta(format("blk.%u.attn_output_a.weight", probe_layer).c_str());
+                    const auto * wo_b_0 = ml.get_tensor_meta(format("blk.%u.attn_output_b.weight", probe_layer).c_str());
 
                     if (!ml.get_key(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, hparams.dsv4_o_group_count, false) && wo_a_0 != nullptr) {
                         GGML_ASSERT(wo_a_0->ne[0] > 0);
@@ -1763,6 +1775,10 @@ void llm_load_hparams(
                             hparams.dsv4_hc_mult = (uint32_t) (wo_a_0->ne[1] / hparams.n_embd);
                         }
                     }
+                    // Base GGUFs lack companion output width, derive it from target HC width.
+                    if (hparams.n_embd_out == hparams.n_embd && hparams.dsv4_hc_mult > 1) {
+                        hparams.n_embd_out = hparams.n_embd * hparams.dsv4_hc_mult;
+                    }
                     if (!ml.get_key(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters, false)) {
                         hparams.dsv4_hc_sinkhorn_iters = 3;
                     }

src/llama-hparams.h

@@ -24,6 +24,7 @@ struct llama_hparams {
     uint32_t n_vocab;
     uint32_t n_ctx_train; // context size the model was trained on
     uint32_t n_embd;
+    uint32_t n_embd_out = 0;
     uint32_t n_layer;
     int32_t n_layer_kv_from_start = -1; // if non-negative, the first n_layer_kv_from_start layers have KV cache
     uint32_t n_rot;

src/llama-load-tensors.cpp

@@ -2785,13 +2785,63 @@ bool create_tensors_helper::create_deepseek4_tensors(const LLM_TN &) {
         return format("blk.%d.%s.weight", i, stem);
     };
 
+    const int mtp_layer = n_layer - (int) hparams.nextn_predict_layers;
+    const bool is_standalone_mtp = hparams.nextn_predict_layers == 1 &&
+        ml.get_tensor_meta("blk.0.attn_norm.weight") == nullptr &&
+        ml.get_tensor_meta(format("blk.%d.nextn.eh_proj.weight", mtp_layer).c_str()) != nullptr;
+
     model.tok_embd    = create_tensor_from_meta(ctx_input,  "token_embd.weight");
     model.output_norm = create_tensor_from_meta(ctx_output, "output_norm.weight");
     model.output      = create_tensor_from_meta(ctx_output, "output.weight");
 
-    model.hc_head_base  = create_tensor_from_meta(ctx_output, pick_tensor_name({"hc_head_base.weight", "output_hc_base.weight"}), llama_model_loader::TENSOR_NOT_REQUIRED);
-    model.hc_head_fn    = create_tensor_from_meta(ctx_output, pick_tensor_name({"hc_head_fn.weight", "output_hc_fn.weight"}), llama_model_loader::TENSOR_NOT_REQUIRED);
-    model.hc_head_scale = create_tensor_from_meta(ctx_output, pick_tensor_name({"hc_head_scale.weight", "output_hc_scale.weight"}), llama_model_loader::TENSOR_NOT_REQUIRED);
+    const int hc_head_flags = is_standalone_mtp ? 0 : llama_model_loader::TENSOR_NOT_REQUIRED;
+    model.hc_head_base  = create_tensor_from_meta(ctx_output, pick_tensor_name({"hc_head_base.weight", "output_hc_base.weight"}), hc_head_flags);
+    model.hc_head_fn    = create_tensor_from_meta(ctx_output, pick_tensor_name({"hc_head_fn.weight", "output_hc_fn.weight"}), hc_head_flags);
+    model.hc_head_scale = create_tensor_from_meta(ctx_output, pick_tensor_name({"hc_head_scale.weight", "output_hc_scale.weight"}), hc_head_flags);
+
+    // Standalone companions declare the full block count but contain only predictor tensors.
+    if (is_standalone_mtp) {
+        const int i = mtp_layer;
+        ggml_context * ctx_split = ctx_for_layer_split(i);
+        auto & layer = model.layers[i];
+
+        layer.attn_norm      = create_tensor_from_meta(ctx_split, format("blk.%d.attn_norm.weight", i));
+        layer.attn_sinks     = create_tensor_from_meta(ctx_split, format("blk.%d.attn_sinks.weight", i));
+        layer.wq_a           = create_tensor_from_meta(ctx_split, format("blk.%d.attn_q_a.weight", i));
+        layer.attn_q_a_norm  = create_tensor_from_meta(ctx_split, format("blk.%d.attn_q_a_norm.weight", i));
+        layer.wq_b           = create_tensor_from_meta(ctx_split, format("blk.%d.attn_q_b.weight", i));
+        layer.wkv_latent     = create_tensor_from_meta(ctx_split, format("blk.%d.attn_kv.weight", i));
+        layer.wkv_b          = layer.wkv_latent;
+        layer.wkv_a_mqa      = layer.wkv_latent;
+        layer.attn_kv_a_norm = create_tensor_from_meta(ctx_split, format("blk.%d.attn_kv_a_norm.weight", i));
+        layer.attn_kv_norm   = layer.attn_kv_a_norm;
+        layer.wo_a           = create_tensor_from_meta(ctx_split, format("blk.%d.attn_output_a.weight", i));
+        layer.wo_b           = create_tensor_from_meta(ctx_split, format("blk.%d.attn_output_b.weight", i));
+        layer.wo             = layer.wo_b;
+
+        layer.hc_attn_base  = create_tensor_from_meta(ctx_split, format("blk.%d.hc_attn_base.weight", i));
+        layer.hc_attn_fn    = create_tensor_from_meta(ctx_split, format("blk.%d.hc_attn_fn.weight", i));
+        layer.hc_attn_scale = create_tensor_from_meta(ctx_split, format("blk.%d.hc_attn_scale.weight", i));
+        layer.hc_ffn_base   = create_tensor_from_meta(ctx_split, format("blk.%d.hc_ffn_base.weight", i));
+        layer.hc_ffn_fn     = create_tensor_from_meta(ctx_split, format("blk.%d.hc_ffn_fn.weight", i));
+        layer.hc_ffn_scale  = create_tensor_from_meta(ctx_split, format("blk.%d.hc_ffn_scale.weight", i));
+
+        layer.ffn_norm       = create_tensor_from_meta(ctx_split, format("blk.%d.ffn_norm.weight", i));
+        layer.ffn_gate_inp   = create_tensor_from_meta(ctx_split, format("blk.%d.ffn_gate_inp.weight", i));
+        layer.ffn_exp_probs_b = create_tensor_from_meta(ctx_split, format("blk.%d.exp_probs_b.bias", i));
+        layer.ffn_gate_exps  = create_tensor_from_meta(ctx_split, format("blk.%d.ffn_gate_exps.weight", i));
+        layer.ffn_down_exps  = create_tensor_from_meta(ctx_split, format("blk.%d.ffn_down_exps.weight", i));
+        layer.ffn_up_exps    = create_tensor_from_meta(ctx_split, format("blk.%d.ffn_up_exps.weight", i));
+        layer.ffn_gate_shexp = create_tensor_from_meta(ctx_split, format("blk.%d.ffn_gate_shexp.weight", i));
+        layer.ffn_down_shexp = create_tensor_from_meta(ctx_split, format("blk.%d.ffn_down_shexp.weight", i));
+        layer.ffn_up_shexp   = create_tensor_from_meta(ctx_split, format("blk.%d.ffn_up_shexp.weight", i));
+
+        layer.nextn.eh_proj          = create_tensor_from_meta(ctx_split, format("blk.%d.nextn.eh_proj.weight", i));
+        layer.nextn.enorm            = create_tensor_from_meta(ctx_split, format("blk.%d.nextn.enorm.weight", i));
+        layer.nextn.hnorm            = create_tensor_from_meta(ctx_split, format("blk.%d.nextn.hnorm.weight", i));
+        layer.nextn.shared_head_norm = create_tensor_from_meta(ctx_split, format("blk.%d.nextn.shared_head_norm.weight", i));
+        return use_mmap_buffer;
+    }
 
     for (int i = 0; i < n_layer; ++i) {
         ggml_context * ctx_split = ctx_for_layer_split(i);

src/llama-spec-features.cpp

@@ -15,6 +15,10 @@ uint32_t llama_mtp_state_n_embd(const struct llama_context * ctx) {
         return hparams.mtp_backbone_n_embd;
     }
 
+    if (ctx->cparams.mtp && ctx->model.arch == LLM_ARCH_DEEPSEEK4 && hparams.n_embd_out > hparams.n_embd) {
+        return hparams.n_embd_out;
+    }
+
     return hparams.n_embd;
 }
 

src/llama.cpp

@@ -3714,7 +3714,7 @@ static std::pair<std::vector<double>, double> get_layer_sizes(const llama_model_
         if (name == "output_norm.weight") {
             continue;
         }
-        if (auto pos = name.find("output_hc_"); pos == 0) {
+        if (name.find("output_hc_") == 0 || name.find("hc_head_") == 0) {
             ow_size += size;
             continue;
         }
@@ -5653,7 +5653,8 @@ static bool llama_context_has_mtp_outputs(const llama_context & lctx) {
         lctx.model.hparams.nextn_predict_layers > 0 ||
         lctx.model.arch == LLM_ARCH_GEMMA4 ||
         lctx.model.arch == LLM_ARCH_GEMMA4_MTP ||
-        lctx.model.arch == LLM_ARCH_GEMMA4_ASSISTANT);
+        lctx.model.arch == LLM_ARCH_GEMMA4_ASSISTANT ||
+        lctx.model.arch == LLM_ARCH_DEEPSEEK4);
 }
 
 static size_t llama_output_reserve(llama_context & lctx, size_t n_outputs) {
@@ -6202,24 +6203,28 @@ static int llama_decode_internal(
             const bool has_mtp = llama_context_has_mtp_outputs(lctx);
             const bool use_raw_mtp_embd = has_mtp && (lctx.model.arch == LLM_ARCH_GEMMA4    ||
                                                       lctx.model.arch == LLM_ARCH_GEMMA4_MTP||
-                                                      lctx.model.arch == LLM_ARCH_GEMMA4_ASSISTANT);
+                                                      lctx.model.arch == LLM_ARCH_GEMMA4_ASSISTANT ||
+                                                      lctx.model.arch == LLM_ARCH_DEEPSEEK4);
+            // For DSV4 we want to extract the 16,384-dim embedding first
             if (cparams.embeddings || has_mtp) {
-                for (int i = gf->n_nodes - 1; i >= 0; --i) {
-                    if (use_raw_mtp_embd && strcmp(gf->nodes[i]->name, "result_mtp_embd") == 0) {
-                        // MTP recurrent state can be wider/different than the logits head hidden state.
-                        embd = gf->nodes[i];
-                        break;
+                if (use_raw_mtp_embd) {
+                    for (int i = gf->n_nodes - 1; i >= 0; --i) {
+                        if (strcmp(gf->nodes[i]->name, "result_mtp_embd") == 0) {
+                            embd = gf->nodes[i];
+                            break;
+                        }
                     }
-                    if (strcmp(gf->nodes[i]->name, "result_embd_pooled") == 0) {
-                        embd = gf->nodes[i];
-                        break;
-                    }
-                    // Strictly speaking we should use if (!use_raw_mtp_embd && strcmp(gf->nodes[i]->name, "result_norm") == 0)
-                    // as Gemma4 MTP is supposed to be using embeddings before rms_norm.
-                    // I don't see any significant difference between this and what we had before, so not making the change (yet).
-                    if (strcmp(gf->nodes[i]->name, "result_norm") == 0) {
-                        embd = gf->nodes[i];
-                        break;
+                }
+                if (!embd) {
+                    for (int i = gf->n_nodes - 1; i >= 0; --i) {
+                        if (strcmp(gf->nodes[i]->name, "result_embd_pooled") == 0) {
+                            embd = gf->nodes[i];
+                            break;
+                        }
+                        if (strcmp(gf->nodes[i]->name, "result_norm") == 0) {
+                            embd = gf->nodes[i];
+                            break;
+                        }
                     }
                 }
             }
@@ -6248,6 +6253,7 @@ static int llama_decode_internal(
         llama_graph_compute(lctx, gf, n_threads);
 
         if (lctx.model.arch == LLM_ARCH_DEEPSEEK4 &&
+            lctx.cparams.mtp_op_type == MTP_OP_NONE &&
             lctx.kv_self.ckpt.selected_spec_mode == LLAMA_SPEC_CKPT_PER_STEP &&
             !llama_dsv4_spec_ckpt_capture_rows(&lctx)) {
             return GGML_STATUS_FAILED;
@@ -8337,10 +8343,10 @@ struct llama_context * llama_init_from_model(
         }
     }
 
-    if (cparams.mtp && hparams.nextn_predict_layers > 0) {
+    if (cparams.mtp && (hparams.nextn_predict_layers > 0 || model->arch == LLM_ARCH_DEEPSEEK4)) {
         const auto n_batch = cparams.n_batch;
         const auto n_vocab = hparams.n_vocab;
-        const auto n_embd  = hparams.n_embd;
+        const auto n_embd  = llama_output_embd_width(*ctx);
 
         const size_t logits_size = n_vocab*n_batch;
         const size_t embd_size   = n_embd*n_batch;
@@ -12328,7 +12334,7 @@ void llama_set_draft_input_hidden_state(struct llama_context * ctx, const float
     ctx->draft_input_hidden_state = hidden_state;
     ctx->draft_input_hidden_state_n_floats = ctx->inp_mtp_states
         ? ggml_nbytes(ctx->inp_mtp_states) / sizeof(float)
-        : 0;
+        : llama_mtp_state_n_embd(ctx);
 }
 
 void llama_set_mtp_target_context(struct llama_context * ctx, struct llama_context * target_ctx) {