diff --git a/common/speculative.cpp b/common/speculative.cpp index 53c75e647..8de714dc0 100644 --- a/common/speculative.cpp +++ b/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 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, diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index 3b581cdbc..688193a36 100644 --- a/convert_hf_to_gguf.py +++ b/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, diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index f1a484bab..630605cc5 100644 --- a/gguf-py/gguf/constants.py +++ b/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 diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 4f606c640..84a6a9656 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/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) diff --git a/src/graphs/build_deepseek4.cpp b/src/graphs/build_deepseek4.cpp index d7a8c8cec..37d1e68a5 100644 --- a/src/graphs/build_deepseek4.cpp +++ b/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); diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index fcde0a95b..9c05d3b8b 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -127,6 +127,7 @@ static const std::map 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" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index bdd9e2a38..2f4224a34 100644 --- a/src/llama-arch.h +++ b/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, diff --git a/src/llama-dsv4.cpp b/src/llama-dsv4.cpp index fee18f3f1..6edd57282 100644 --- a/src/llama-dsv4.cpp +++ b/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; } diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 6ceb548a6..2af8704d1 100644 --- a/src/llama-hparams.cpp +++ b/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(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; } diff --git a/src/llama-hparams.h b/src/llama-hparams.h index ec8daf2d2..5fe94bde2 100644 --- a/src/llama-hparams.h +++ b/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; diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index a75642b0a..e21aca11f 100644 --- a/src/llama-load-tensors.cpp +++ b/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); diff --git a/src/llama-spec-features.cpp b/src/llama-spec-features.cpp index 2e93b5fa9..5e6a6ff61 100644 --- a/src/llama-spec-features.cpp +++ b/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; } diff --git a/src/llama.cpp b/src/llama.cpp index 79e66adf4..5b6515aee 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -3714,7 +3714,7 @@ static std::pair, 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) {