diff --git a/common/speculative.cpp b/common/speculative.cpp index d34d1c9c595..988b8c9bb75 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2582,7 +2582,7 @@ 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()); - llama_model * model_dft = llama_model_load_from_file(params.model.path.c_str(), mparams); + llama_model * model_dft = llama_model_load_from_file(model_path.c_str(), mparams); if (model_dft == NULL) { LOG_ERR("%s: failed to load draft model, '%s'\n", __func__, model_path.c_str()); return; diff --git a/conversion/qwen4exp.py b/conversion/qwen4exp.py index 168796d616b..5de29934605 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,64 @@ 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 + # The MTP head is one trunk-shaped block (dense attention + MoE, wrapped in + # hyper-connections) plus a combiner, so once _QwenMtpMixin renames + # `mtp.layers.0.*` to the trailing block index its tensors ride the existing + # qwen4exp mappings unchanged. Only the two head-level pieces below differ. + + _MTP_MIXER_PREFIX = "mtp.hyper_connection_mixer." + + @classmethod + def filter_tensors(cls, item): + # the head carries its own copy of the trunk's hc_head_* output mixer, + # which qwen4exp has in place of a final norm; it is unindexed in the + # checkpoint and 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]]: + # qwen4exp splits the combiner the shared NextN code calls eh_proj into + # fc_embedding and fc_hidden; W_e@e + W_h@h == [W_e|W_h] @ concat(e, h), + # so the two fuse back into the single expected matmul + 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 token embedding ahead of + # the hidden state, so the fused weight has to be ordered to match + 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 +107,18 @@ 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)] + # llama.cpp reads this array with length block_count, and the MTP blocks + # trailing the trunk attend densely, which is what a ratio of 0 selects + 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 carries no trunk tensors, so it carries no PLE table + # to describe 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/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index c99feb3c795..1b0d439c7de 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -1175,6 +1175,11 @@ 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, which stands in for the + # output norm the trunk does not have + 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,15 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.PLE_NORM_QUERY, MODEL_TENSOR.PLE_NORM_CONV, MODEL_TENSOR.PLE_CONV1D, + # NextN/MTP draft head + 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/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 861acfe181f..7bd73c38a66 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -2742,6 +2742,16 @@ class TensorNameMap: MODEL_TENSOR.HC_HEAD_UP: ( "model.hyper_connection_mixer.input_mix_weight_up", ), + # the MTP head carries its own copy of the head mixer above + 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/src/llama-arch.cpp b/src/llama-arch.cpp index 5e61f61f7f0..7daee813809 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -574,6 +574,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 +965,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..f013137665a 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -686,6 +686,11 @@ enum llm_tensor { LLM_TENSOR_NEXTN_HNORM, LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, + // qwen4exp: the MTP head ends in its own hyper-connection mixer rather than a + // plain RMSNorm, mirroring the trunk's hc_head_* (which is its output 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.cpp b/src/llama-model.cpp index e679b24e87f..e06a78f1d43 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; diff --git a/src/llama-model.h b/src/llama-model.h index 38066538ed1..c67d86f823d 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -227,6 +227,12 @@ 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 own final hyper-connection mixer, which stands in for both + // the stream collapse and the output norm (the trunk has no separate output_norm either) + 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/models/models.h b/src/models/models.h index 9b87a40d5af..f93518943ca 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: + // tag-dispatched ctor for graph_mtp: 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,11 @@ struct llama_model_qwen4exp : public llama_model_base { const llama_model & model; }; + // LLM_GRAPH_TYPE_DECODER_MTP draft head: one HC-wrapped dense-attention + MoE block + 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..18a6bddca00 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -7,6 +7,11 @@ #include void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { + // NextN/MTP: an extra decoder block appended past the trunk. Read this first, since + // n_layer() == n_layer_all - n_layer_nextn feeds every per-layer array below. + 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); @@ -114,10 +119,16 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0); + // draft-head-only files (no trunk tensors) ship without the hc head mixer; + // the MTP graph falls back to nextn.shared_head_norm + the block's own ffn + const bool mtp_only = (hparams.n_layer_nextn > 0) && + (ml.get_weight(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", 0).str().c_str()) == nullptr); + const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 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); + 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) { @@ -125,7 +136,7 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { } // flat [ple_head_dim, n_rows] gather target; n_rows is padded, so read it back - if (hparams.ple_n_heads > 0) { + if (!mtp_only && hparams.ple_n_heads > 0) { const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").str(); const auto & ple_w = ml.require_weight(ple_name.c_str()); const int64_t ple_rows = ple_w.tensor->ne[1]; @@ -140,9 +151,20 @@ 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) { + // MTP tensors sit in the trailing blocks; skip them entirely unless a draft head was asked for + 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]; + if (mtp_only && il < n_layer) { + continue; + } + + // the MTP block is structurally a trunk block: is_recr()/is_ple() are both false past + // the trunk, so it takes the full-attention + MoE path below with no special casing + const int flags = il < n_layer ? 0 : 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 +177,87 @@ 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 }, 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.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); - 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); + if (il < n_layer) { + continue; + } + + // NextN/MTP head. enorm/hnorm gate the two inputs; eh_proj is the checkpoint's + // fc_embedding and fc_hidden fused side by side, so one matmul over + // concat(e, h) computes fc_embedding@e + fc_hidden@h. + 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); + + // the head's own output mixer, mirroring the trunk's hc_head_*: it collapses the + // hc streams and stands in for the output norm, of which qwen4exp has none + layer.nextn.hc_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_HC_HEAD_NORM, "weight", il), { hc_dim }, TENSOR_NOT_REQUIRED | flags); + layer.nextn.hc_head_down = create_tensor(tn(LLM_TENSOR_NEXTN_HC_HEAD_DOWN, "weight", il), { hc_dim, hc_lr }, TENSOR_NOT_REQUIRED | flags); + layer.nextn.hc_head_up = create_tensor(tn(LLM_TENSOR_NEXTN_HC_HEAD_UP, "weight", il), { hc_lr, hc_dim }, TENSOR_NOT_REQUIRED | flags); + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { hc_dim }, TENSOR_NOT_REQUIRED | flags); + + // qwen4exp sets mtp_use_dedicated_embeddings=false, so these are absent and the + // head falls back to the trunk's embedding table and LM head + 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 +397,11 @@ 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 a hidden row for every token, so in that case the + // gather is deferred until after t_h_nextn is taken below + 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 +429,23 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa cb(res_hc, "l_last", il); } + // The MTP head consumes the wide residual, before the head mixer collapses it. Export the + // combine result itself rather than a reshape of it: a pure view gets no backend assignment + // from the scheduler, and the readback in llama_context looks one up. It is contiguous, so + // [n_embd, hc, rows] already has the [n_embd_out, rows] layout the reader expects, and it + // carries exactly the right rows either way -- gathered above when masked, ungathered when not. + if (cparams.embeddings_nextn) { + cb(res_hc, "h_nextn", -1); + res->t_h_nextn = res_hc; + + // deferred from the last layer: collapse to the output rows now that the export is taken + 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 +461,203 @@ 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. +// +// The head folds the next token's embedding into the trunk's wide hyper-connection residual, +// runs one trunk-style block over it, and collapses the result with its own mixer before +// reusing the trunk's LM head. The wide post-block residual is exported as t_h_nextn so the +// speculative driver can feed it straight back in for the next draft step. +// +// v1 simplification: the block attends densely. The trunk's QSA only prunes context past a +// 2048-token budget, so dense is a numerical superset; drafts are verified by the target +// either way. The indexer tensors are still loaded so the GGUF stays complete. +// 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"); + + // files shipped without the head's own mixer use the nextn norm and the + // block's ffn projections instead + ggml_tensor * head_norm = layer.nextn.hc_head_norm ? layer.nextn.hc_head_norm : layer.nextn.shared_head_norm; + ggml_tensor * head_down = layer.nextn.hc_head_down ? layer.nextn.hc_head_down : layer.hc_ffn_down; + ggml_tensor * head_up = layer.nextn.hc_head_up ? layer.nextn.hc_head_up : layer.hc_ffn_up; + GGML_ASSERT(head_norm && head_down && head_up && "MTP block missing head mixer tensors"); + + 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(); + + // grouped RMSNorm over the wide stream: normalise each hc stream, then scale the flattened + // [hc_dim] vector with the head's gamma, exactly as build_hc_mix does + 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); + + // the token embedding is shared across the streams, so broadcast it to hc copies + 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); + + // eh_proj holds fc_embedding and fc_hidden side by side, so this one matmul is + // fc_embedding @ e_norm + fc_hidden @ h_norm, applied to each stream independently. + // Keeping the streams distinct here is the point of the hyper-connection residual: + // pooling them before the projection would throw that away. + 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); + + // ---- dense attention, mirroring the trunk's full-attention branch ---- + 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); + + // IMRoPE, same convention and freq_base as the trunk + 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); + + // ---- MoE, identical to the trunk's build_layer_ffn ---- + 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. + // As in the trunk, export the combine result rather than a reshape view of it. + cb(res_hc, "h_nextn", -1); + res->t_h_nextn = res_hc; + + // the head's own mixer collapses the streams and doubles as the output norm + cur = build_hc_mix(res_hc, + head_norm, head_down, head_up, + nullptr, nullptr, -1); + cb(cur, "mtp_hc_head", -1); + + // deliberately no res->t_embd: it would be n_embd wide while the context sizes its + // embedding buffer by n_embd_out (the wide stream). The driver reads t_h_nextn instead. + + 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) {