diff --git a/common/speculative.cpp b/common/speculative.cpp index 851a47b9a58..c9709961df1 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2542,6 +2542,8 @@ common_speculative_init_result::common_speculative_init_result( model_path = params.speculative.draft.mparams.path; LOG_INF("%s: loading draft model '%s'\n", __func__, model_path.c_str()); + mparams.model_shared = model_tgt; + llama_model * model_dft = llama_model_load_from_file(params.model.path.c_str(), mparams); if (model_dft == NULL) { LOG_ERR("%s: failed to load draft model, '%s'\n", __func__, model_path.c_str()); diff --git a/conversion/bailingmoe3.py b/conversion/bailingmoe3.py index 20bba23e51c..9ba3112ebc5 100644 --- a/conversion/bailingmoe3.py +++ b/conversion/bailingmoe3.py @@ -121,9 +121,9 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca if is_mtp and cls.no_mtp: return None - if cls.mtp_only and not is_mtp and name not in ( + if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in ( "model.word_embeddings.weight", "model.norm.weight", "lm_head.weight", - ): + )): return None return super().filter_tensors((name, gen)) diff --git a/conversion/base.py b/conversion/base.py index daae28e92ad..c0dd413b596 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -120,6 +120,7 @@ class ModelBase: supports_mtp_export: bool = False mtp_only: bool = False no_mtp: bool = False + mtp_shared_embd: bool = False def __init__(self, dir_model: Path, ftype: gguf.LlamaFileType, fname_out: Path, *, is_big_endian: bool = False, use_temp_file: bool = False, eager: bool = False, @@ -1032,6 +1033,10 @@ def set_type(self): def prepare_metadata(self, vocab_only: bool): + # tells the loader they are missing on purpose + if self.mtp_only and self.mtp_shared_embd: + self.gguf_writer.add_nextn_shared_target_tensors(True) + total_params, shared_params, expert_params, expert_count = self.gguf_writer.get_total_parameter_count() self.metadata = gguf.Metadata.load(self.metadata_override, self.dir_model_card, self.model_name, total_params) diff --git a/conversion/command_r.py b/conversion/command_r.py index 971f93ebdf1..2b513509d55 100644 --- a/conversion/command_r.py +++ b/conversion/command_r.py @@ -131,9 +131,9 @@ def filter_tensors(cls, item): is_mtp = (m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._n_main_layers if is_mtp and cls.no_mtp: return None - if cls.mtp_only and not is_mtp and name not in ( + if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in ( "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", - ): + )): return None return name, gen diff --git a/conversion/dots3.py b/conversion/dots3.py index c7ac2319e24..e8d3f350c74 100644 --- a/conversion/dots3.py +++ b/conversion/dots3.py @@ -99,9 +99,9 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca # --no-mtp: drop the NextN/MTP block; --mtp: keep only that block plus the shared embeddings/norm/lm_head if is_mtp and cls.no_mtp: return None - if cls.mtp_only and not is_mtp and name not in ( + if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in ( "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", - ): + )): return None return name, gen diff --git a/conversion/glm.py b/conversion/glm.py index 7544f850cb2..245f01f84bf 100644 --- a/conversion/glm.py +++ b/conversion/glm.py @@ -138,9 +138,9 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca if is_mtp and cls.no_mtp: return None - if cls.mtp_only and not is_mtp and name not in ( + if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in ( "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", - ): + )): return None return name, gen @@ -292,9 +292,9 @@ def filter_tensors(cls, item): is_mtp = match is not None and int(match.group(1)) >= cls._n_main_layers if is_mtp and cls.no_mtp: return None - if cls.mtp_only and not is_mtp and name not in ( + if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in ( "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", - ): + )): return None return name, gen @@ -352,9 +352,9 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca return None # --mtp: keep ONLY NextN-block tensors plus the shared embeddings/ # norm/lm_head (so the resulting GGUF carries just the draft head). - if cls.mtp_only and not is_mtp and name not in ( + if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in ( "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", - ): + )): return None return name, gen diff --git a/conversion/qwen.py b/conversion/qwen.py index 419611896fc..ca89d27ba4a 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -338,7 +338,7 @@ def filter_tensors(cls, item): elif len(parts) == 3 and parts[1] in remapper: name = f"model.layers.{cls._original_block_count}.{remapper[parts[1]]}.{parts[2]}" elif cls.mtp_only: - keep = name in ( + keep = not cls.mtp_shared_embd and name in ( "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", "embed_tokens.weight", "norm.weight", ) diff --git a/conversion/qwen4exp.py b/conversion/qwen4exp.py index 168796d616b..5b2b0495e92 100644 --- a/conversion/qwen4exp.py +++ b/conversion/qwen4exp.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Iterable, cast +from typing import Callable, Iterable, cast import torch from torch import Tensor @@ -21,20 +21,56 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase): Shares the Qwen3.5 gated delta net and interleaved mrope, and adds three things: hyper-connections in place of every layer norm, QSA sparse attention on the full attention layers, and PLE n-gram hash embeddings on a single layer. + + The checkpoint also carries a NextN/MTP draft head under `mtp.*`, exported as a + trailing block; pass --no-nextn to leave it out. """ model_arch = gguf.MODEL_ARCH.QWEN4EXP - # the MTP block is a separate draft head; vLLM drops it too - supports_mtp_export = False - no_mtp = True - def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # only the shard names, so the table itself is never held self._ple_shards: dict[int, str] = {} self._ple_row_dim: int | None = None + # _QwenMtpMixin renames mtp.layers.0.* to the trailing block index, so the head reuses the + # existing qwen4exp mappings; only the two pieces below differ. + _MTP_MIXER_PREFIX = "mtp.hyper_connection_mixer." + + @classmethod + def filter_tensors(cls, item): + # unindexed in the checkpoint, per-block in the GGUF + name, gen = item + if name.startswith("model." + cls._MTP_MIXER_PREFIX): + name = name.replace("model.", "", 1) + if name.startswith(cls._MTP_MIXER_PREFIX): + if cls.no_mtp: + return None + assert cls._original_block_count is not None + return f"model.layers.{cls._original_block_count}.{name[len('mtp.'):]}", gen + return super().filter_tensors((name, gen)) + + def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]: + # W_e@e + W_h@h == [W_e|W_h] @ concat(e, h), so fc_embedding and fc_hidden fuse into eh_proj + tensors = super().index_tensors(remote_hf_model_id=remote_hf_model_id) + + emb = tensors.pop("mtp.fc_embedding.weight", None) + hid = tensors.pop("mtp.fc_hidden.weight", None) + if emb is None and hid is None: + return tensors + if emb is None or hid is None: + raise ValueError( + "the qwen4exp MTP combiner needs both mtp.fc_embedding.weight and " + "mtp.fc_hidden.weight; pass --no-nextn to convert without the draft head" + ) + + assert self._original_block_count is not None + # fc_embedding first: the graph concatenates the embedding ahead of the hidden state + name = f"model.layers.{self._original_block_count}.eh_proj.weight" + tensors[name] = lambda: torch.cat([emb(), hid()], dim=1) + return tensors + def _read_hash_constants(self, suffix: str) -> list[int]: """Read an int64 PLE constant straight from the checkpoint. @@ -63,14 +99,15 @@ def set_gguf_parameters(self): self.gguf_writer.add_indexer_top_k(hp["indexer_budget"]) ratio = hp["indexer_compress_ratio"] layer_types = hp["layer_types"] - self.gguf_writer.add_attention_compress_ratios( - [ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)] - ) + ratios = [ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)] + # read with length block_count; 0 selects dense, which is how the MTP blocks attend + ratios += [0] * (self.block_count - n_layer) + self.gguf_writer.add_attention_compress_ratios(ratios) # ple_layer_ids is 1-based in the HF config; empty means no n-gram table, - # so emit no PLE keys rather than optional ones + # so emit no PLE keys rather than optional ones. a draft-only export has no PLE table either. ple_layers = [i - 1 for i in hp["ple_layer_ids"]] - if not ple_layers: + if not ple_layers or self.mtp_only: return self.gguf_writer.add_ple_layers(ple_layers) self.gguf_writer.add_ple_ngram_size(hp["ngram_size"]) diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index 78ad26c6563..6e7dddfa661 100755 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -125,6 +125,10 @@ def parse_args() -> argparse.Namespace: "--no-nextn", "--no-mtp", dest="no_mtp", action="store_true", help="Exclude NextN speculative draft tensors from the converted GGUF. Pair with --mtp or --dspark on a second run to publish target and draft as two files.", ) + parser.add_argument( + "--mtp-shared-embd", action="store_true", + help="With --mtp, leave the token embeddings, output norm and LM head out of the draft and take them from the target model at load time. Much smaller draft, but it needs a llama.cpp new enough to read it.", + ) parser.add_argument( "--dspark", action="store_true", help="Export only the DeepSeek-V4 DSpark draft tensors as a separate GGUF.", @@ -278,6 +282,12 @@ def main() -> None: if args.mtp: model_class.mtp_only = True + if args.mtp_shared_embd: + if not args.mtp: + logger.error("--mtp-shared-embd only applies together with --mtp") + sys.exit(1) + model_class.mtp_shared_embd = True + model_instance = model_class(dir_model, output_type, fname_out, is_big_endian=args.bigendian, use_temp_file=args.use_temp_file, eager=args.no_lazy, diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index e5ccd1feab1..3ea72dbce9f 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -1425,13 +1425,13 @@ struct ggml_backend_cuda_context { int curr_stream_no = 0; #ifdef USE_CUDA_GRAPH - // Map from first_node_ptr to cuda_graph - allows multiple graphs per context - // when the computation is split across CPU/GPU (e.g., with --n-cpu-moe) - std::unordered_map> cuda_graphs; + std::unordered_map> cuda_graphs; + + static const size_t max_cuda_graphs = 64; int64_t last_graph_eviction_sweep = 0; - ggml_cuda_graph * cuda_graph(const void * first_node_ptr) { + ggml_cuda_graph * cuda_graph(uint64_t graph_key) { const int64_t time_now = ggml_time_us(); // sweep every 5s, evicting cuda graphs unused for >=10s @@ -1446,9 +1446,18 @@ struct ggml_backend_cuda_context { } } - auto it = cuda_graphs.find(first_node_ptr); + auto it = cuda_graphs.find(graph_key); if (it == cuda_graphs.end()) { - it = cuda_graphs.emplace(first_node_ptr, std::make_unique()).first; + while (cuda_graphs.size() >= max_cuda_graphs) { + auto lru = cuda_graphs.begin(); + for (auto c = cuda_graphs.begin(); c != cuda_graphs.end(); ++c) { + if (c->second->last_used_time < lru->second->last_used_time) { + lru = c; + } + } + cuda_graphs.erase(lru); + } + it = cuda_graphs.emplace(graph_key, std::make_unique()).first; } it->second->last_used_time = time_now; return it->second.get(); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index a6fc655c41c..387339c160c 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2582,14 +2582,30 @@ static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { return use_cuda_graph; } -static const void * ggml_cuda_graph_get_key(ggml_cgraph * cgraph) { - return cgraph->nodes[0]; +// a captured graph hard-codes its shapes, so with one key per split an alternating shape +// (a speculative verify batch) resets warmup forever. O(1) on purpose: walking nodes undoes the +// point of a cuda graph. A shape this fails to separate re-captures as before, so it cannot regress. +static uint64_t ggml_cuda_graph_get_key(ggml_cgraph * cgraph) { + uint64_t key = (uint64_t) (uintptr_t) cgraph->nodes[0]; + + auto mix = [&key](uint64_t v) { + key = (key ^ v) * 0x100000001b3ull; + }; + + mix(cgraph->n_nodes); + + for (int d = 0; d < GGML_MAX_DIMS; d++) { + mix(cgraph->nodes[0]->ne[d]); + mix(cgraph->nodes[cgraph->n_nodes - 1]->ne[d]); + } + + return key; } static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph) { bool res = false; - const void * graph_key = ggml_cuda_graph_get_key(cgraph); + const uint64_t graph_key = ggml_cuda_graph_get_key(cgraph); ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); if (cgraph->uid != 0 && @@ -2628,7 +2644,7 @@ static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx return res; } -static void ggml_cuda_graph_update_executable(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { +static void ggml_cuda_graph_update_executable(ggml_backend_cuda_context * cuda_ctx, uint64_t graph_key) { ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); #if CUDART_VERSION >= 12000 @@ -4019,7 +4035,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return 0; } -static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, const bool use_cuda_graph, const bool cuda_graph_update_required, const void * graph_key) { +static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, const bool use_cuda_graph, const bool cuda_graph_update_required, uint64_t graph_key) { bool graph_evaluated_or_captured = false; // flag used to determine whether it is an integrated_gpu @@ -4238,7 +4254,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } #ifdef USE_CUDA_GRAPH -static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { +static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, uint64_t graph_key) { ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); if (graph->graph == nullptr) { @@ -4261,7 +4277,7 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, bool use_cuda_graph = false; bool cuda_graph_update_required = false; - const void * graph_key = nullptr; + uint64_t graph_key = 0; #ifdef USE_CUDA_GRAPH graph_key = ggml_cuda_graph_get_key(cgraph); @@ -4344,7 +4360,7 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; #ifdef USE_CUDA_GRAPH - const void * graph_key = ggml_cuda_graph_get_key(cgraph); + const uint64_t graph_key = ggml_cuda_graph_get_key(cgraph); const bool use_cuda_graph = ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); #else const bool use_cuda_graph = false; diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index c99feb3c795..4e7f2d42cc5 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -129,6 +129,7 @@ class LLM: MOE_EVERY_N_LAYERS = "{arch}.moe_every_n_layers" MOE_LATENT_SIZE = "{arch}.moe_latent_size" NEXTN_PREDICT_LAYERS = "{arch}.nextn_predict_layers" + NEXTN_SHARED_TARGET_TENSORS = "{arch}.nextn_shared_target_tensors" NUM_DEEPSTACK_LAYERS = "{arch}.n_deepstack_layers" DEEPSTACK_MAPPING = "{arch}.deepstack_mapping" POOLING_TYPE = "{arch}.pooling_type" @@ -1175,6 +1176,10 @@ class MODEL_TENSOR(IntEnum): NEXTN_HNORM = auto() NEXTN_SHARED_HEAD_HEAD = auto() NEXTN_SHARED_HEAD_NORM = auto() + # qwen4exp: the MTP head's own hyper-connection mixer, in place of an output norm + NEXTN_HC_HEAD_NORM = auto() + NEXTN_HC_HEAD_DOWN = auto() + NEXTN_HC_HEAD_UP = auto() # eagle3 FC = auto() # feature fusion layer D2T = auto() # draft to target vocabulary mapping @@ -1952,6 +1957,9 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_HNORM: "blk.{bid}.nextn.hnorm", MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD: "blk.{bid}.nextn.shared_head_head", MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: "blk.{bid}.nextn.shared_head_norm", + MODEL_TENSOR.NEXTN_HC_HEAD_NORM: "blk.{bid}.nextn.hc_head_norm", + MODEL_TENSOR.NEXTN_HC_HEAD_DOWN: "blk.{bid}.nextn.hc_head_down", + MODEL_TENSOR.NEXTN_HC_HEAD_UP: "blk.{bid}.nextn.hc_head_up", MODEL_TENSOR.FC: "fc", MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1", MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2", @@ -2914,6 +2922,14 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.PLE_NORM_QUERY, MODEL_TENSOR.PLE_NORM_CONV, MODEL_TENSOR.PLE_CONV1D, + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_HC_HEAD_NORM, + MODEL_TENSOR.NEXTN_HC_HEAD_DOWN, + MODEL_TENSOR.NEXTN_HC_HEAD_UP, ], MODEL_ARCH.PLAMO: [ MODEL_TENSOR.TOKEN_EMBD, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index d95fe9b1ac3..87f32c32962 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -905,6 +905,9 @@ def add_moe_latent_size(self, value: int) -> None: def add_nextn_predict_layers(self, count: int) -> None: self.add_uint32(Keys.LLM.NEXTN_PREDICT_LAYERS.format(arch=self.arch), count) + def add_nextn_shared_target_tensors(self, value: bool) -> None: + self.add_bool(Keys.LLM.NEXTN_SHARED_TARGET_TENSORS.format(arch=self.arch), value) + def add_swin_norm(self, value: bool) -> None: self.add_bool(Keys.LLM.SWIN_NORM.format(arch=self.arch), value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 861acfe181f..c9a6574d9b2 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -2742,6 +2742,15 @@ class TensorNameMap: MODEL_TENSOR.HC_HEAD_UP: ( "model.hyper_connection_mixer.input_mix_weight_up", ), + MODEL_TENSOR.NEXTN_HC_HEAD_NORM: ( + "model.layers.{bid}.hyper_connection_mixer.hc_norm", + ), + MODEL_TENSOR.NEXTN_HC_HEAD_DOWN: ( + "model.layers.{bid}.hyper_connection_mixer.input_mix_weight_down", + ), + MODEL_TENSOR.NEXTN_HC_HEAD_UP: ( + "model.layers.{bid}.hyper_connection_mixer.input_mix_weight_up", + ), MODEL_TENSOR.INDEXER_Q_NORM: ( "model.layers.{bid}.self_attn.indexer.q_layernorm", ), diff --git a/include/llama.h b/include/llama.h index ef7a012c43a..41b9123042b 100644 --- a/include/llama.h +++ b/include/llama.h @@ -340,6 +340,9 @@ extern "C" { // override key-value pairs of the model meta data const struct llama_model_kv_override * kv_overrides; + // target for a draft head that declares nextn_shared_target_tensors; must outlive this model + const struct llama_model * model_shared; + // Keep the booleans together to avoid misalignment during copy-by-value. bool vocab_only; // only load the vocabulary, no weights bool check_tensors; // validate model tensor data diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 5e61f61f7f0..c9f10334eba 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -217,6 +217,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_MOE_EVERY_N_LAYERS, "%s.moe_every_n_layers" }, { LLM_KV_MOE_LATENT_SIZE, "%s.moe_latent_size" }, { LLM_KV_NEXTN_PREDICT_LAYERS, "%s.nextn_predict_layers" }, + { LLM_KV_NEXTN_SHARED_TARGET_TENSORS, "%s.nextn_shared_target_tensors" }, { LLM_KV_NUM_DEEPSTACK_LAYERS, "%s.n_deepstack_layers" }, { LLM_KV_DEEPSTACK_MAPPING, "%s.deepstack_mapping" }, { LLM_KV_HIDDEN_ACT, "%s.hidden_activation" }, @@ -574,6 +575,9 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_NEXTN_HNORM, "blk.%d.nextn.hnorm" }, { LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "blk.%d.nextn.shared_head_head" }, { LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "blk.%d.nextn.shared_head_norm" }, + { LLM_TENSOR_NEXTN_HC_HEAD_NORM, "blk.%d.nextn.hc_head_norm" }, + { LLM_TENSOR_NEXTN_HC_HEAD_DOWN, "blk.%d.nextn.hc_head_down" }, + { LLM_TENSOR_NEXTN_HC_HEAD_UP, "blk.%d.nextn.hc_head_up" }, { LLM_TENSOR_ATTN_SUB_NORM, "blk.%d.attn_sub_norm" }, { LLM_TENSOR_FFN_SUB_NORM, "blk.%d.ffn_sub_norm" }, { LLM_TENSOR_DEC_OUTPUT_NORM, "dec.output_norm" }, @@ -962,6 +966,9 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_NEXTN_HNORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_NEXTN_HC_HEAD_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_NEXTN_HC_HEAD_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_NEXTN_HC_HEAD_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, // Nemotron 3 Super // latent projections feed ggml_mul_mat, the buft probe must use MUL_MAT to keep them on GPU {LLM_TENSOR_FFN_LATENT_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, diff --git a/src/llama-arch.h b/src/llama-arch.h index ca7d55a5fd7..835398f9782 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -222,6 +222,7 @@ enum llm_kv { LLM_KV_MOE_EVERY_N_LAYERS, LLM_KV_MOE_LATENT_SIZE, LLM_KV_NEXTN_PREDICT_LAYERS, + LLM_KV_NEXTN_SHARED_TARGET_TENSORS, LLM_KV_NUM_DEEPSTACK_LAYERS, LLM_KV_DEEPSTACK_MAPPING, LLM_KV_HIDDEN_ACT, @@ -686,6 +687,9 @@ enum llm_tensor { LLM_TENSOR_NEXTN_HNORM, LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, + LLM_TENSOR_NEXTN_HC_HEAD_NORM, + LLM_TENSOR_NEXTN_HC_HEAD_DOWN, + LLM_TENSOR_NEXTN_HC_HEAD_UP, LLM_TENSOR_MASKED_EMBD_CENTROIDS, LLM_TENSOR_MASKED_EMBD_ORDERING, LLM_TENSOR_FC, diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 7663797ba00..13e77ab218b 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1106,6 +1106,68 @@ bool llama_model_loader::lazy_read::add(const std::string & name, const ggml_ten return true; } +// declared in llama-model.h, which this file does not include +const std::vector> & llama_internal_get_tensor_map(const llama_model * model); + +struct ggml_tensor * llama_model_loader::borrow_shared_tensor(const LLM_TN_IMPL & tn, const std::initializer_list & ne) { + // checked first so no other tensor in any model pays a metadata lookup + if (tn.tensor != LLM_TENSOR_TOKEN_EMBD && tn.tensor != LLM_TENSOR_OUTPUT && tn.tensor != LLM_TENSOR_OUTPUT_NORM) { + return nullptr; + } + + if (shared_target_tensors < 0) { + bool shared = false; + get_key(LLM_KV_NEXTN_SHARED_TARGET_TENSORS, shared, false); + shared_target_tensors = shared ? 1 : 0; + } + if (shared_target_tensors == 0) { + return nullptr; + } + + const std::string name = tn.str(); + if (get_weight(name.c_str()) != nullptr) { + return nullptr; + } + + if (model_shared == nullptr) { + throw std::runtime_error(format("%s: this model is a draft head without its own '%s'; " + "load it as a draft of its target model, not on its own", __func__, name.c_str())); + } + + ggml_tensor * src = nullptr; + for (const auto & [n, t] : llama_internal_get_tensor_map(model_shared)) { + if (n == name) { + src = t; + break; + } + } + if (src == nullptr) { + throw std::runtime_error(format("%s: draft needs tensor '%s' from the target, which does not have it", + __func__, name.c_str())); + } + + // used directly, so the shapes must agree exactly + size_t dim = 0; + for (const int64_t n : ne) { + if (dim >= GGML_MAX_DIMS || src->ne[dim] != n) { + throw std::runtime_error(format("%s: draft and target disagree on '%s': target has %s, draft wants %s", + __func__, name.c_str(), llama_format_tensor_shape(src).c_str(), llama_format_tensor_shape(ne).c_str())); + } + dim++; + } + for (; dim < GGML_MAX_DIMS; dim++) { + if (src->ne[dim] != 1) { + throw std::runtime_error(format("%s: draft and target disagree on '%s': target has %s, draft wants %s", + __func__, name.c_str(), llama_format_tensor_shape(src).c_str(), llama_format_tensor_shape(ne).c_str())); + } + } + + LLAMA_LOG_INFO("%s: tensor %s taken from the target model\n", __func__, name.c_str()); + + // not counted in n_created/size_data: not in this file, neither allocated nor freed here + return src; +} + struct ggml_tensor * llama_model_loader::create_tensor( const llama_hparams & hparams, const buft_list_t * buft_list_cpu, const buft_list_t * buft_list_input, const buft_list_t * buft_list_output, const buft_list_t * buft_list_layer, const LLM_TN_IMPL & tn, const std::initializer_list & ne, int flags) { @@ -1326,6 +1388,11 @@ struct ggml_tensor * llama_model_loader::create_tensor( return ret; } + // must precede check_tensor_dims, and must win over the arch fallback that ties output to token_embd + if (ggml_tensor * shared = borrow_shared_tensor(tn, ne)) { + return shared; + } + LLAMA_LOG_DEBUG("%s: loading tensor %s\n", __func__, tn.str().c_str()); const struct ggml_tensor * cur = check_tensor_dims(tn.str(), ne, !(flags & TENSOR_NOT_REQUIRED), flags & TENSOR_ALLOW_RESHAPE); if (cur == NULL) { diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index 9e51d0ce750..7cf1d823cde 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -117,6 +117,11 @@ struct llama_model_loader { std::set tensors; } lazy; + const struct llama_model * model_shared = nullptr; + + // cached nextn_shared_target_tensors, -1 until first read + int shared_target_tensors = -1; + llama_files files; llama_ftype ftype; llama_fver fver; @@ -238,6 +243,9 @@ struct llama_model_loader { const llama_hparams & hparams, const buft_list_t * buft_list_cpu, const buft_list_t * buft_list_input, const buft_list_t * buft_list_output, const buft_list_t * buft_list_layer, const LLM_TN_IMPL & tn, const std::initializer_list & ne, int flags); + // token_embd/output/output_norm from the target. null unless the file declares the flag. + struct ggml_tensor * borrow_shared_tensor(const LLM_TN_IMPL & tn, const std::initializer_list & ne); + void done_getting_tensors(bool partial = false) const; void init_mappings(bool prefetch = true, llama_mlocks * mlock_mmaps = nullptr); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index e679b24e87f..30efa02a338 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2434,7 +2434,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, const bool mtp_on_hybrid_qwen = params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || - arch == LLM_ARCH_BAILINGMOE3); + arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_QWEN4EXP); const bool mtp_on_hybrid_nemotron = params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && arch == LLM_ARCH_NEMOTRON_H_MOE; @@ -2692,6 +2692,7 @@ llama_model_params llama_model_default_params() { /*.progress_callback =*/ nullptr, /*.progress_callback_user_data =*/ nullptr, /*.kv_overrides =*/ nullptr, + /*.model_shared =*/ nullptr, /*.vocab_only =*/ false, /*.check_tensors =*/ false, /*.use_extra_bufts =*/ true, diff --git a/src/llama-model.h b/src/llama-model.h index 38066538ed1..f58352f0835 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -227,6 +227,11 @@ struct llama_layer_nextn { struct ggml_tensor * shared_head_head_s = nullptr; struct ggml_tensor * shared_head_head_in_s = nullptr; struct ggml_tensor * shared_head_norm = nullptr; + + // qwen4exp: the MTP head's mixer; collapses the streams and stands in for the output norm + struct ggml_tensor * hc_head_norm = nullptr; + struct ggml_tensor * hc_head_down = nullptr; + struct ggml_tensor * hc_head_up = nullptr; }; struct llama_layer_switch_lora { diff --git a/src/llama.cpp b/src/llama.cpp index 633db658c95..7c49b3a2462 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -318,7 +318,8 @@ static std::pair llama_model_load(struct gguf_context * meta llama_model_loader ml(metadata, set_tensor_data, set_tensor_data_ud, fname, splits, file, params.load_mode, params.check_tensors, params.no_alloc, params.load_mtp, params.kv_overrides, params.tensor_buft_overrides); - ml.lazy.mode = params.lazy_mode; + ml.lazy.mode = params.lazy_mode; + ml.model_shared = params.model_shared; ml.print_info(); std::unique_ptr model_ptr(llama_model_create(ml, params)); diff --git a/src/models/models.h b/src/models/models.h index 9b87a40d5af..51e09a699e4 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2285,7 +2285,12 @@ struct llama_model_qwen4exp : public llama_model_base { struct graph : public llm_build_delta_net_base { graph(const llama_model & model, const llm_graph_params & params); - private: + protected: + // graph_mtp ctor: binds the members without building the trunk + struct no_build_t {}; + graph(const llama_model & model, const llm_graph_params & params, no_build_t) : + llm_build_delta_net_base(params), model(model) {} + // HC replaces every layer norm: residual is [n_embd, hc, n_tokens] ggml_tensor * build_hc_mix( ggml_tensor * x, @@ -2377,6 +2382,10 @@ struct llama_model_qwen4exp : public llama_model_base { const llama_model & model; }; + struct graph_mtp : public graph { + graph_mtp(const llama_model & model, const llm_graph_params & params); + }; + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index abf6a0502fb..1ab0cef048d 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -7,6 +7,10 @@ #include void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { + // must precede the per-layer arrays: n_layer() == n_layer_all - n_layer_nextn. + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < block_count"); + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false); ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false); ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); @@ -112,12 +116,16 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { const int64_t hc_dim = hc * n_embd; const int64_t hc_lr = hparams.hc_low_rank; + // a draft-only export declares the full block count but ships the MTP block alone. + const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.hc_attn_norm.weight") == nullptr); + const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0); - // there is no output_norm: the final hyper-connection mixer carries it - hc_head_norm = create_tensor(tn(LLM_TENSOR_HC_HEAD_NORM, "weight"), { hc_dim }, 0); - hc_head_down = create_tensor(tn(LLM_TENSOR_HC_HEAD_DOWN, "weight"), { hc_dim, hc_lr }, 0); - hc_head_up = create_tensor(tn(LLM_TENSOR_HC_HEAD_UP, "weight"), { hc_lr, hc_dim }, 0); + // no output_norm: this mixer carries it. the MTP head has its own in nextn.hc_head_*. + hc_head_norm = create_tensor(tn(LLM_TENSOR_HC_HEAD_NORM, "weight"), { hc_dim }, trunk_flags); + hc_head_down = create_tensor(tn(LLM_TENSOR_HC_HEAD_DOWN, "weight"), { hc_dim, hc_lr }, trunk_flags); + hc_head_up = create_tensor(tn(LLM_TENSOR_HC_HEAD_UP, "weight"), { hc_lr, hc_dim }, trunk_flags); output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED); if (output == NULL) { @@ -140,9 +148,13 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { { hparams.ple_head_dim, ple_rows }, TENSOR_READ_LAZY); } - for (int il = 0; il < n_layer; ++il) { + const int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0; + + for (int il = 0; il < (int) hparams.n_layer_all; ++il) { auto & layer = layers[il]; + const int flags = il < n_layer ? trunk_flags : mtp_flags; + const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used; const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff; @@ -155,61 +167,80 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { const int64_t conv_dim = key_dim * 2 + value_dim; // two HC modules per layer: before the token mixer, before the MoE - layer.hc_attn_norm = create_tensor(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", il), { hc_dim }, 0); - layer.hc_attn_down = create_tensor(tn(LLM_TENSOR_HC_ATTN_DOWN, "weight", il), { hc_dim, hc_lr }, 0); - layer.hc_attn_up = create_tensor(tn(LLM_TENSOR_HC_ATTN_UP, "weight", il), { hc_lr, hc_dim }, 0); - layer.hc_attn_inject = create_tensor(tn(LLM_TENSOR_HC_ATTN_INJECT, "weight", il), { hc_dim, hc }, 0); - layer.hc_ffn_norm = create_tensor(tn(LLM_TENSOR_HC_FFN_NORM, "weight", il), { hc_dim }, 0); - layer.hc_ffn_down = create_tensor(tn(LLM_TENSOR_HC_FFN_DOWN, "weight", il), { hc_dim, hc_lr }, 0); - layer.hc_ffn_up = create_tensor(tn(LLM_TENSOR_HC_FFN_UP, "weight", il), { hc_lr, hc_dim }, 0); - layer.hc_ffn_inject = create_tensor(tn(LLM_TENSOR_HC_FFN_INJECT, "weight", il), { hc_dim, hc }, 0); + layer.hc_attn_norm = create_tensor(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", il), { hc_dim }, flags); + layer.hc_attn_down = create_tensor(tn(LLM_TENSOR_HC_ATTN_DOWN, "weight", il), { hc_dim, hc_lr }, flags); + layer.hc_attn_up = create_tensor(tn(LLM_TENSOR_HC_ATTN_UP, "weight", il), { hc_lr, hc_dim }, flags); + layer.hc_attn_inject = create_tensor(tn(LLM_TENSOR_HC_ATTN_INJECT, "weight", il), { hc_dim, hc }, flags); + layer.hc_ffn_norm = create_tensor(tn(LLM_TENSOR_HC_FFN_NORM, "weight", il), { hc_dim }, flags); + layer.hc_ffn_down = create_tensor(tn(LLM_TENSOR_HC_FFN_DOWN, "weight", il), { hc_dim, hc_lr }, flags); + layer.hc_ffn_up = create_tensor(tn(LLM_TENSOR_HC_FFN_UP, "weight", il), { hc_lr, hc_dim }, flags); + layer.hc_ffn_inject = create_tensor(tn(LLM_TENSOR_HC_FFN_INJECT, "weight", il), { hc_dim, hc }, flags); if (!hparams.is_recr(il)) { // full attention: wq holds [q|gate] interleaved per head - create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, 0); - layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, 0); + create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, flags); - layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, 0); - layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, 0); + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, flags); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, flags); const int64_t idx_dim = hparams.indexer_head_size; - layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", il), { n_embd, hparams.indexer_n_head * idx_dim }, 0); - layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", il), { n_embd, idx_dim }, 0); - layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", il), { idx_dim }, 0); - layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", il), { idx_dim }, 0); + layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", il), { n_embd, hparams.indexer_n_head * idx_dim }, flags); + layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", il), { n_embd, idx_dim }, flags); + layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", il), { idx_dim }, flags); + layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", il), { idx_dim }, flags); } else { - layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, key_dim * 2 + value_dim }, 0); - layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, 0); - layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, 0); - layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, 0); - layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { hparams.ssm_dt_rank }, 0); - layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_v_heads }, 0); - layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", il), { n_embd, n_v_heads }, 0); - layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_v_dim }, 0); - layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", il), { value_dim, n_embd }, 0); + layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, key_dim * 2 + value_dim }, flags); + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, flags); + layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, flags); + layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, flags); + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { hparams.ssm_dt_rank }, flags); + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_v_heads }, flags); + layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", il), { n_embd, n_v_heads }, flags); + layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_v_dim }, flags); + layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", il), { value_dim, n_embd }, flags); } if (hparams.is_ple(il)) { - layer.ple_key = create_tensor(tn(LLM_TENSOR_PLE_KEY, "weight", il), { n_embd, hc_dim }, 0); - layer.ple_value = create_tensor(tn(LLM_TENSOR_PLE_VALUE, "weight", il), { n_embd, n_embd }, 0); - layer.ple_norm_key = create_tensor(tn(LLM_TENSOR_PLE_NORM_KEY, "weight", il), { hc_dim }, 0); - layer.ple_norm_query = create_tensor(tn(LLM_TENSOR_PLE_NORM_QUERY, "weight", il), { hc_dim }, 0); - layer.ple_norm_conv = create_tensor(tn(LLM_TENSOR_PLE_NORM_CONV, "weight", il), { hc_dim }, 0); - layer.ple_conv1d = create_tensor(tn(LLM_TENSOR_PLE_CONV1D, "weight", il), { hparams.ple_conv_kernel, hc_dim }, 0); + layer.ple_key = create_tensor(tn(LLM_TENSOR_PLE_KEY, "weight", il), { n_embd, hc_dim }, flags); + layer.ple_value = create_tensor(tn(LLM_TENSOR_PLE_VALUE, "weight", il), { n_embd, n_embd }, flags); + layer.ple_norm_key = create_tensor(tn(LLM_TENSOR_PLE_NORM_KEY, "weight", il), { hc_dim }, flags); + layer.ple_norm_query = create_tensor(tn(LLM_TENSOR_PLE_NORM_QUERY, "weight", il), { hc_dim }, flags); + layer.ple_norm_conv = create_tensor(tn(LLM_TENSOR_PLE_NORM_CONV, "weight", il), { hc_dim }, flags); + layer.ple_conv1d = create_tensor(tn(LLM_TENSOR_PLE_CONV1D, "weight", il), { hparams.ple_conv_kernel, hc_dim }, flags); + } + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, flags); + create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, flags); + + layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, flags); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, flags); + + if (il < n_layer) { + continue; } - layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, 0); - layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, 0); - create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, 0); + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { hc_dim }, flags); + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, flags); - layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, 0); - layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0); - layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0); - layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, 0); + layer.nextn.hc_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_HC_HEAD_NORM, "weight", il), { hc_dim }, flags); + layer.nextn.hc_head_down = create_tensor(tn(LLM_TENSOR_NEXTN_HC_HEAD_DOWN, "weight", il), { hc_dim, hc_lr }, flags); + layer.nextn.hc_head_up = create_tensor(tn(LLM_TENSOR_NEXTN_HC_HEAD_UP, "weight", il), { hc_lr, hc_dim }, flags); + + // absent when mtp_use_dedicated_embeddings=false (qwen4exp); the head falls back to the trunk's. + layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, flags | TENSOR_NOT_REQUIRED); + layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab }, flags | TENSOR_NOT_REQUIRED); } } std::unique_ptr llama_model_qwen4exp::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } return std::make_unique(*this, params); } @@ -349,7 +380,10 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa cur = build_layer_attn(inp->get_attn(), mctx_hyb, cur, inp_pos, sections, il); } - if (il == n_layer - 1 && inp_out_ids) { + // an unmasked MTP export needs every token's row, so it defers the gather until after t_h_nextn. + const bool gather_now = !cparams.embeddings_nextn || cparams.embeddings_nextn_masked; + + if (il == n_layer - 1 && inp_out_ids && gather_now) { // everything below is per token, so drop the rows that produce no output cur = ggml_get_rows(ctx0, cur, inp_out_ids); inject = ggml_get_rows(ctx0, inject, inp_out_ids); @@ -377,6 +411,18 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa cb(res_hc, "l_last", il); } + // export res_hc itself, never a reshape view: a pure view gets no backend assignment to read back. + if (cparams.embeddings_nextn) { + cb(res_hc, "h_nextn", -1); + res->t_h_nextn = res_hc; + + if (!cparams.embeddings_nextn_masked && inp_out_ids) { + res_hc = ggml_reshape_2d(ctx0, res_hc, n_embd*hc, res_hc->ne[2]); + res_hc = ggml_get_rows(ctx0, res_hc, inp_out_ids); + res_hc = ggml_reshape_3d(ctx0, res_hc, n_embd, hc, res_hc->ne[1]); + } + } + // the final mixer is the output norm: there is no separate one ggml_tensor * cur = build_hc_mix(res_hc, model.hc_head_norm, model.hc_head_down, model.hc_head_up, @@ -392,6 +438,177 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa ggml_build_forward_expand(gf, cur); } +// LLM_GRAPH_TYPE_DECODER_MTP draft head for qwen4exp. Attends densely: QSA only prunes context +// past a 2048-token budget, so dense is a numerical superset and drafts are verified regardless. +// TODO: wire up QSA here for long-context draft fidelity. +llama_model_qwen4exp::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) : + graph(model, params, no_build_t{}) { + GGML_ASSERT(hparams.n_layer_nextn > 0 && "QWEN4EXP MTP requires n_layer_nextn > 0"); + GGML_ASSERT(hparams.n_layer_nextn == 1 && "QWEN4EXP MTP currently only supports a single MTP block"); + GGML_ASSERT(ubatch.token && "QWEN4EXP MTP requires token input"); + + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc * n_embd; + GGML_ASSERT(hparams.n_embd_out() == (uint32_t) hc_dim && "QWEN4EXP MTP hidden width mismatch"); + + const int il = hparams.n_layer(); + const auto & layer = model.layers[il]; + + GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj"); + GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm"); + GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm"); + GGML_ASSERT(layer.nextn.hc_head_norm && "MTP block missing nextn.hc_head_norm"); + + int sections[4]; + std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections); + + auto inp = std::make_unique(hc_dim); + + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hc_dim, n_tokens); + ggml_set_input(inp->embd); + + inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hc_dim, n_tokens); + ggml_set_input(inp->h); + ggml_set_name(inp->h, "mtp_h_input"); + + ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd; + ggml_tensor * tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens); + cb(tok_embd, "mtp_tok_embd", il); + + ggml_tensor * h_state = ggml_reshape_3d(ctx0, inp->h, n_embd, hc, n_tokens); + cb(h_state, "mtp_h_state", il); + + res->add_input(std::move(inp)); + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + auto * inp_attn = build_attn_inp_kv(); + + ggml_tensor * h_norm = ggml_rms_norm(ctx0, h_state, hparams.f_norm_rms_eps); + h_norm = ggml_reshape_2d(ctx0, h_norm, hc_dim, n_tokens); + h_norm = ggml_mul(ctx0, h_norm, layer.nextn.hnorm); + h_norm = ggml_reshape_3d(ctx0, h_norm, n_embd, hc, n_tokens); + cb(h_norm, "mtp_hnorm", il); + + ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); + e_norm = ggml_repeat_4d(ctx0, + ggml_reshape_3d(ctx0, e_norm, n_embd, 1, n_tokens), + n_embd, hc, n_tokens, 1); + cb(e_norm, "mtp_enorm", il); + + // per stream, not pooled: pooling before the projection discards the hyper-connection residual. + ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0); + cb(concat, "mtp_concat", il); + + ggml_tensor * res_hc = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s); + cb(res_hc, "mtp_eh_proj", il); + + ggml_tensor * inject = nullptr; + ggml_tensor * cur = build_hc_mix(res_hc, + layer.hc_attn_norm, layer.hc_attn_down, layer.hc_attn_up, layer.hc_attn_inject, + &inject, il); + cb(cur, "mtp_hc_attn_pre", il); + + const int64_t n_embd_head = hparams.n_embd_head_v(); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + + ggml_tensor * Qcur_full = build_lora_mm(layer.wq, cur, layer.wq_s); + cb(Qcur_full, "mtp_Qcur_full", il); + + ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, + ggml_element_size(Qcur_full) * n_embd_head * 2, + ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head, 0); + Qcur = build_norm(Qcur, layer.attn_q_norm, nullptr, LLM_NORM_RMS, il); + cb(Qcur, "mtp_Qcur_normed", il); + + ggml_tensor * gate = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, + ggml_element_size(Qcur_full) * n_embd_head * 2, + ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head, + ggml_element_size(Qcur_full) * n_embd_head); + gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens); + cb(gate, "mtp_gate", il); + + ggml_tensor * Kcur = build_lora_mm(layer.wk, cur, layer.wk_s); + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); + Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il); + cb(Kcur, "mtp_Kcur_normed", il); + + ggml_tensor * Vcur = build_lora_mm(layer.wv, cur, layer.wv_s); + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); + cb(Vcur, "mtp_Vcur", il); + + Qcur = ggml_rope_multi(ctx0, Qcur, inp_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + Kcur = ggml_rope_multi(ctx0, Kcur, inp_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(Qcur, "mtp_Qcur", il); + cb(Kcur, "mtp_Kcur", il); + + const float kq_scale = hparams.f_attention_scale == 0.0f + ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale; + + cur = build_attn(inp_attn, + nullptr, nullptr, nullptr, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + cb(cur, "mtp_attn_pregate", il); + + cur = ggml_mul(ctx0, cur, ggml_sigmoid(ctx0, gate)); + cb(cur, "mtp_attn_gated", il); + + cur = build_lora_mm(layer.wo, cur, layer.wo_s); + cb(cur, "mtp_attn_out", il); + + if (inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inject = ggml_get_rows(ctx0, inject, inp_out_ids); + + res_hc = ggml_reshape_2d(ctx0, res_hc, hc_dim, res_hc->ne[2]); + res_hc = ggml_get_rows(ctx0, res_hc, inp_out_ids); + res_hc = ggml_reshape_3d(ctx0, res_hc, n_embd, hc, res_hc->ne[1]); + } + + res_hc = build_hc_combine(res_hc, cur, inject, il); + cb(res_hc, "mtp_hc_attn_post", il); + + cur = build_hc_mix(res_hc, + layer.hc_ffn_norm, layer.hc_ffn_down, layer.hc_ffn_up, layer.hc_ffn_inject, + &inject, il); + cb(cur, "mtp_hc_ffn_pre", il); + + cur = build_layer_ffn(cur, il); + cb(cur, "mtp_ffn_out", il); + + res_hc = build_hc_combine(res_hc, cur, inject, il); + cb(res_hc, "mtp_hc_ffn_post", il); + + // the next draft step re-enters here, so export the wide stream before it is collapsed. + cb(res_hc, "h_nextn", -1); + res->t_h_nextn = res_hc; + + cur = build_hc_mix(res_hc, + layer.nextn.hc_head_norm, layer.nextn.hc_head_down, layer.nextn.hc_head_up, + nullptr, nullptr, -1); + cb(cur, "mtp_hc_head", -1); + + // no res->t_embd: it is n_embd wide, but the context sizes that buffer by n_embd_out. + + ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output; + ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s; + GGML_ASSERT(head_w && "QWEN4EXP MTP: missing LM head (nextn.shared_head_head or model.output)"); + + cur = build_lora_mm(head_w, cur, head_s); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + std::pair llama_model_qwen4exp::graph::build_qkvz( ggml_tensor * input, int il) {