super().set_vocab()
+class _QwenMtpMixin:
+ """Shared MTP wiring for Qwen3-Next and Qwen3.5/3.6 text variants. The HF
+ config carries the MTP block under `mtp_num_hidden_layers` (computed from
+ the checkpoint when absent, e.g. Qwen3-Next) and the tensors under
+ `mtp.*`; we extend block_count, emit the nextn metadata key, and remap
+ `mtp.*` to the standard layer-indexed nextn naming so the existing
+ tensor_map handles them."""
+
+ supports_mtp_export = True
+ hparams: dict[str, Any]
+ model_arch: gguf.MODEL_ARCH
+ gguf_writer: gguf.GGUFWriter
+ block_count: int
+ tensor_map: gguf.TensorNameMap
+ no_mtp: bool
+ mtp_only: bool
+ _original_block_count: int | None = None
+ opt_num_mtp_layers: int = 0
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.block_count = self.hparams["num_hidden_layers"]
+ if not self.no_mtp:
+ n_mtp = self.hparams.get("mtp_num_hidden_layers", 0)
+ # Qwen-3-Next doesn't include `mtp_num_hidden_layers` in config.
+ if n_mtp == 0:
+ assert self.opt_num_mtp_layers != 0
+ n_mtp = self.opt_num_mtp_layers
+ self.block_count += n_mtp
+ self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
+
+ def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
+ hparams = {**self.hparams, **self.hparams.get("text_config", {})}
+ key = next((k for k in ["n_layers", "num_hidden_layers", "n_layer", "num_layers"] if k in hparams), None)
+ type(self)._original_block_count = hparams.get(key)
+ type(self).opt_num_mtp_layers = 0
+ return super().index_tensors(remote_hf_model_id=remote_hf_model_id) # ty: ignore[unresolved-attribute]
+
+ @classmethod
+ def filter_tensors(cls, item):
+ assert cls._original_block_count is not None
+ # TODO: change TextModel to super()
+ if (titem := TextModel.filter_tensors(item)) is None:
+ return None
+ name, gen = titem
+ if name.startswith("model.mtp."):
+ name = name.replace("model.", "", 1)
+ if name.startswith("mtp."):
+ if cls.no_mtp:
+ return None
+ remapper = {
+ "fc": "eh_proj",
+ "pre_fc_norm_embedding": "enorm",
+ "pre_fc_norm_hidden": "hnorm",
+ "norm": "shared_head.norm",
+ }
+ parts = name.split(".", 3)
+ if len(parts) == 4 and parts[1] == "layers" and parts[2].isdecimal():
+ mtp_idx = int(parts[2])
+ name = f"model.layers.{cls._original_block_count + mtp_idx}.{parts[3]}"
+ cls.opt_num_mtp_layers = max(cls.opt_num_mtp_layers, mtp_idx + 1)
+ 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 (
+ "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
+ "embed_tokens.weight", "norm.weight",
+ )
+ if not keep:
+ return None
+ return name, gen
+
+ def set_gguf_parameters(self):
+ super().set_gguf_parameters() # ty: ignore[unresolved-attribute]
+ if self.no_mtp:
+ return
+ if (n := self.block_count - self.hparams["num_hidden_layers"]) > 0:
+ self.gguf_writer.add_nextn_predict_layers(n)
+
+ def prepare_metadata(self, vocab_only: bool):
+ from_dir = self.fname_out.is_dir()
+ super().prepare_metadata(vocab_only=vocab_only) # ty: ignore[unresolved-attribute]
+
+ if not self.mtp_only or not from_dir:
+ return
+
+ output_type: str = self.ftype.name.partition("_")[2] # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
+ fname_default: str = gguf.naming_convention(
+ self.metadata.name, self.metadata.basename, self.metadata.finetune, # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
+ self.metadata.version, size_label=None, output_type=output_type, model_type=None) # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
+ self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
+
+
@ModelBase.register("Qwen3NextForCausalLM")
-class Qwen3NextModel(Qwen2MoeModel):
+class Qwen3NextModel(_QwenMtpMixin, Qwen2MoeModel):
model_arch = gguf.MODEL_ARCH.QWEN3NEXT
def set_gguf_parameters(self):
rope_dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]
self.gguf_writer.add_rope_dimension_count(int(rope_dim * self.rope_parameters.get("partial_rotary_factor", 0.25)))
- @classmethod
- def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
- name, gen = item
-
- if name.startswith("mtp"):
- # ignore MTP layers for now
- return None
-
- return super().filter_tensors(item)
-
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
if name.endswith(".A_log"):
data_torch = -torch.exp(data_torch)
self.gguf_writer.add_rope_dimension_sections(self._QWEN35_DEFAULT_MROPE_SECTION)
-class _Qwen35MtpMixin:
- """Shared MTP wiring for Qwen3.5/3.6 text variants. The HF config carries
- the MTP block under `mtp_num_hidden_layers` and the tensors under
- `mtp.*`; we extend block_count, emit the nextn metadata key, and remap
- `mtp.*` to the standard layer-indexed nextn naming so the existing
- tensor_map handles them."""
-
- supports_mtp_export = True
- hparams: dict[str, Any]
- model_arch: gguf.MODEL_ARCH
- gguf_writer: gguf.GGUFWriter
- block_count: int
- tensor_map: gguf.TensorNameMap
- no_mtp: bool
- mtp_only: bool
- _original_block_count: int | None = None
-
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
- self.block_count = self.hparams["num_hidden_layers"]
- if not self.no_mtp:
- self.block_count += self.hparams.get("mtp_num_hidden_layers", 0)
- self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
-
- def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
- hparams = {**self.hparams, **self.hparams.get("text_config", {})}
- key = next((k for k in ["n_layers", "num_hidden_layers", "n_layer", "num_layers"] if k in hparams), None)
- type(self)._original_block_count = hparams.get(key)
- return super().index_tensors(remote_hf_model_id=remote_hf_model_id) # ty: ignore[unresolved-attribute]
-
- @classmethod
- def filter_tensors(cls, item):
- assert cls._original_block_count is not None
- # TODO: change TextModel to super()
- if (titem := TextModel.filter_tensors(item)) is None:
- return None
- name, gen = titem
- if name.startswith("model.mtp."):
- name = name.replace("model.", "", 1)
- if name.startswith("mtp."):
- if cls.no_mtp:
- return None
- remapper = {
- "fc": "eh_proj",
- "pre_fc_norm_embedding": "enorm",
- "pre_fc_norm_hidden": "hnorm",
- "norm": "shared_head.norm",
- }
- parts = name.split(".", 3)
- if len(parts) == 4 and parts[1] == "layers" and parts[2].isdecimal():
- mtp_idx = int(parts[2])
- name = f"model.layers.{cls._original_block_count + mtp_idx}.{parts[3]}"
- 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 (
- "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
- "embed_tokens.weight", "norm.weight",
- )
- if not keep:
- return None
- return name, gen
-
- def set_gguf_parameters(self):
- super().set_gguf_parameters() # ty: ignore[unresolved-attribute]
- if self.no_mtp:
- return
- if (n := self.hparams.get("mtp_num_hidden_layers", 0)) > 0:
- self.gguf_writer.add_nextn_predict_layers(n)
-
- def prepare_metadata(self, vocab_only: bool):
- from_dir = self.fname_out.is_dir()
- super().prepare_metadata(vocab_only=vocab_only) # ty: ignore[unresolved-attribute]
-
- if not self.mtp_only or not from_dir:
- return
-
- output_type: str = self.ftype.name.partition("_")[2] # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
- fname_default: str = gguf.naming_convention(
- self.metadata.name, self.metadata.basename, self.metadata.finetune, # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
- self.metadata.version, size_label=None, output_type=output_type, model_type=None) # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
- self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
-
-
@ModelBase.register("Qwen3_5ForConditionalGeneration", "Qwen3_5ForCausalLM")
-class Qwen3_5TextModel(_Qwen35MtpMixin, _Qwen35MRopeMixin, _LinearAttentionVReorderBase):
+class Qwen3_5TextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
model_arch = gguf.MODEL_ARCH.QWEN35
@ModelBase.register("Qwen3_5MoeForConditionalGeneration", "Qwen3_5MoeForCausalLM")
-class Qwen3_5MoeTextModel(_Qwen35MtpMixin, _Qwen35MRopeMixin, _LinearAttentionVReorderBase):
+class Qwen3_5MoeTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
model_arch = gguf.MODEL_ARCH.QWEN35MOE
ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);
ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group);
- // Mark recurrent layers (linear attention layers)
+ // NextN/MTP: extra decoder block appended beyond the main stack
+ 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 < n_layer_all");
+
+ // Mark recurrent layers (linear attention layers).
if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) {
uint32_t full_attn_interval = 4;
ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false);
}
}
-void llama_model_qwen3next::load_arch_tensors(llama_model_loader &) {
+void llama_model_qwen3next::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
if (n_expert == 0) {
throw std::runtime_error(arch_name() + " model cannot have zero experts");
}
+ const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
+ const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
+ int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
+
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
// output
const int64_t qkvz_dim = key_dim * 2 + value_dim * 2;
const int64_t ba_dim = n_v_heads * 2;
- for (int i = 0; i < n_layer; ++i) {
- auto & layer = layers[i];
- const uint32_t n_ff_shexp = hparams.n_ff_shexp > 0 ? hparams.n_ff_shexp : hparams.n_ff(i);
+ auto load_block_trunk = [&](int il, int flags) {
+ auto & layer = layers[il];
+ const uint32_t n_ff_shexp = hparams.n_ff_shexp > 0 ? hparams.n_ff_shexp : hparams.n_ff(il);
- layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0);
- layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), { n_embd }, 0);
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, flags);
+ layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, flags);
- if (!hparams.is_recr(i)) {
+ if (!hparams.is_recr(il)) {
// Attention layers
- create_tensor_qkv(layer, i, 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", i), { 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);
// Q/K normalization for attention layers
- layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0);
- layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { 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);
} else {
// Linear attention (gated delta net) specific tensors
// Create tensors with calculated dimensions
// note: ssm_in is used by legacy GGUF
- layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), { n_embd, qkvz_dim }, TENSOR_NOT_REQUIRED);
- layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), { n_embd, key_dim * 2 + value_dim }, TENSOR_NOT_REQUIRED);
- layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), { n_embd, value_dim }, TENSOR_NOT_REQUIRED);
- layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), { hparams.ssm_d_conv, conv_dim }, 0);
- layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), { hparams.ssm_dt_rank }, 0);
- layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), { hparams.ssm_dt_rank }, 0);
- layer.ssm_beta_alpha = create_tensor(tn(LLM_TENSOR_SSM_BETA_ALPHA, "weight", i), { n_embd, ba_dim }, 0);
- layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), { head_v_dim }, 0);
- layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), { value_dim, n_embd }, 0);
+ layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", il), { n_embd, qkvz_dim }, TENSOR_NOT_REQUIRED | flags);
+ layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, key_dim * 2 + value_dim }, TENSOR_NOT_REQUIRED | flags);
+ layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, TENSOR_NOT_REQUIRED | 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_alpha = create_tensor(tn(LLM_TENSOR_SSM_BETA_ALPHA, "weight", il), { n_embd, ba_dim }, 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);
}
- layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert }, 0);
- layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), { n_ff_exp, n_embd, n_expert }, 0);
- create_tensor_gate_up_exps(layer, i, 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);
// Shared experts
- layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", i), { n_embd }, 0);
- layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), { n_embd, n_ff_shexp }, 0);
- layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), { n_embd, n_ff_shexp }, 0);
- layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_shexp, n_embd }, 0);
+ 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);
+ };
+
+ auto load_block_mtp = [&](int il) {
+ // MTP head is identical to the trunk block (full attention + FFN)
+ load_block_trunk(il, mtp_flags);
+
+ auto & layer = layers[il];
+
+ // NextN-specific tensors that define the MTP block.
+ layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, mtp_flags);
+ layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, mtp_flags);
+ layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, mtp_flags);
+ layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, mtp_flags | TENSOR_NOT_REQUIRED);
+ layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab }, mtp_flags | TENSOR_NOT_REQUIRED);
+ layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, mtp_flags | TENSOR_NOT_REQUIRED);
+ };
+
+ for (int i = 0; i < n_layer; i++) {
+ load_block_trunk(i, trunk_flags);
+ }
+ for (int i = n_layer; i < n_layer_all; i++) {
+ load_block_mtp(i);
}
}
std::unique_ptr<llm_graph_context> llama_model_qwen3next::build_arch_graph(const llm_graph_params & params) const {
+ if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
+ return std::make_unique<graph_mtp>(*this, params);
+ }
return std::make_unique<graph>(*this, params);
}
ggml_tensor * inp_pos = build_inp_pos();
ggml_tensor * inp_out_ids = build_inp_out_ids();
+ // MTP/NextN layers are loaded as extra decoder blocks but not executed in the main pass.
for (int il = 0; il < n_layer; ++il) {
res->t_layer_inp[il] = inpL;
cur = build_layer_attn(inp->get_attn(), cur, inp_pos, il);
}
- if (il == n_layer - 1 && inp_out_ids) {
+ if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
}
cur = inpL;
- // Final norm
+ // post-norm hidden state is input to both the LM head and the MTP head
cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
+ cb(cur, "h_nextn", -1);
+ res->t_h_nextn = cur;
+
+ if (!cparams.embeddings_nextn_masked && inp_out_ids) {
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+ }
+
cb(cur, "result_norm", -1);
res->t_embd = cur;
ggml_build_forward_expand(gf, cur);
}
-// utility to get one slice from the third dimension
-// input dim: [x, y, c, b]
-// output dim: [x, y, 1, b]
-static ggml_tensor * get_slice_2d(ggml_context * ctx0, ggml_tensor * t, int64_t c) {
- return ggml_view_4d(ctx0, t, t->ne[0], t->ne[1], 1, t->ne[3],
- t->nb[1], t->nb[2], t->nb[3], t->nb[2] * c);
-}
-
ggml_tensor * llama_model_qwen3next::graph::build_norm_gated(
ggml_tensor * input,
ggml_tensor * weights,
// Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention
// Qwen3Next uses a single Q projection that outputs query + gate
- ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur);
+ ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s);
cb(Qcur_full, "Qcur_full", il);
Qcur_full = ggml_reshape_4d(ctx0, Qcur_full, n_embd_head * 2, n_head, n_tokens, 1);
Qcur_full->nb[1], Qcur_full->nb[2], Qcur_full->nb[3], n_embd_head * ggml_element_size(Qcur_full));
cb(gate, "gate", il);
- ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur);
+ ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s);
cb(Kcur, "Kcur", il);
- ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur);
+ ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s);
cb(Vcur, "Vcur", il);
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
gate = ggml_sigmoid(ctx0, gate);
cb(gate, "gate_sigmoid", il);
- gate = ggml_reshape_2d(ctx0, gate, n_embd_head * n_head, n_tokens);
-
cur = ggml_mul(ctx0, cur, gate);
cb(cur, "attn_gated", il);
LLM_FFN_SILU, true,
hparams.expert_weights_scale,
LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, il,
- nullptr, model.layers[il].ffn_gate_up_exps);
+ nullptr, model.layers[il].ffn_gate_up_exps,
+ model.layers[il].ffn_up_exps_s,
+ model.layers[il].ffn_gate_exps_s,
+ model.layers[il].ffn_down_exps_s);
cb(moe_out, "ffn_moe_out", il);
// Add shared experts if present - following Qwen3Next reference implementation
if (model.layers[il].ffn_up_shexp != nullptr) {
ggml_tensor * ffn_shexp =
build_ffn(cur,
- model.layers[il].ffn_up_shexp, NULL, NULL,
- model.layers[il].ffn_gate_shexp, NULL, NULL,
- model.layers[il].ffn_down_shexp, NULL, NULL,
+ model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s,
+ model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s,
+ model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s,
NULL,
LLM_FFN_SILU, LLM_FFN_PAR, il);
cb(ffn_shexp, "ffn_shexp", il);
}
return cur;
}
+
+// LLM_GRAPH_TYPE_DECODER_MTP draft head for Qwen3-Next
+llama_model_qwen3next::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
+ : llm_graph_context(params) {
+ GGML_ASSERT(hparams.n_layer_nextn > 0 && "QWEN3NEXT MTP requires n_layer_nextn > 0");
+ GGML_ASSERT(hparams.n_layer_nextn == 1 && "QWEN3NEXT MTP currently only supports a single MTP block");
+
+ const int64_t n_embd_head = hparams.n_embd_head_v();
+ GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
+
+ 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.ffn_gate_inp && "MTP block missing ffn_gate_inp");
+
+ // TODO: extract in a common llm_graph_context::build_inp_embd_h()
+ auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);
+
+ 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, hparams.n_embd_inp(), n_tokens);
+ ggml_set_input(inp->embd);
+
+ // TODO: make static using `ggml_build_forward_select()`
+ // see llm_graph_context::build_inp_embd() for reference
+ ggml_tensor * tok_embd;
+ if (ubatch.token) {
+ ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
+
+ tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
+ } else {
+ tok_embd = inp->embd;
+ }
+ cb(tok_embd, "mtp_tok_embd", il);
+
+ inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
+ ggml_set_input(inp->h);
+ ggml_set_name(inp->h, "mtp_h_input");
+
+ ggml_tensor * h_embd = inp->h;
+
+ 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 = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
+ cb(h_norm, "mtp_hnorm", il);
+
+ ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
+ cb(e_norm, "mtp_enorm", il);
+
+ ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);
+ cb(concat, "mtp_concat", il);
+
+ ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
+ cb(cur, "mtp_eh_proj", il);
+
+ ggml_tensor * inpSA = cur;
+
+ cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_attn_norm", il);
+
+ 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 * 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);
+
+ Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr,
+ n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+ Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr,
+ n_rot, 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);
+ cb(Vcur, "mtp_Vcur", 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);
+
+ 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);
+
+ // TODO: CUDA is missing non-contiguous unary ops. when implemented: remove this cont
+ gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens);
+ cb(gate, "mtp_gate", il);
+
+ cur = ggml_mul(ctx0, cur, ggml_sigmoid(ctx0, gate));
+ 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);
+ inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
+ }
+
+ cur = ggml_add(ctx0, cur, inpSA);
+ cb(cur, "mtp_attn_residual", il);
+
+ ggml_tensor * ffn_residual = cur;
+ cur = build_norm(cur, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_attn_post_norm", il);
+
+ // MoE FFN — routed experts plus gated shared expert (mirrors the trunk).
+ ggml_tensor * moe_out =
+ build_moe_ffn(cur,
+ layer.ffn_gate_inp,
+ layer.ffn_up_exps,
+ layer.ffn_gate_exps,
+ layer.ffn_down_exps,
+ nullptr,
+ n_expert, n_expert_used,
+ LLM_FFN_SILU, true,
+ hparams.expert_weights_scale,
+ LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, il,
+ nullptr, layer.ffn_gate_up_exps,
+ layer.ffn_up_exps_s,
+ layer.ffn_gate_exps_s,
+ layer.ffn_down_exps_s);
+ cb(moe_out, "mtp_ffn_moe_out", il);
+
+ if (layer.ffn_up_shexp != nullptr) {
+ ggml_tensor * ffn_shexp =
+ build_ffn(cur,
+ layer.ffn_up_shexp, nullptr, layer.ffn_up_shexp_s,
+ layer.ffn_gate_shexp, nullptr, layer.ffn_gate_shexp_s,
+ layer.ffn_down_shexp, nullptr, layer.ffn_down_shexp_s,
+ nullptr,
+ LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(ffn_shexp, "mtp_ffn_shexp", il);
+
+ ggml_tensor * shared_gate = build_lora_mm(layer.ffn_gate_inp_shexp, cur);
+ shared_gate = ggml_sigmoid(ctx0, shared_gate);
+ cb(shared_gate, "mtp_shared_expert_gate_sigmoid", il);
+
+ ffn_shexp = ggml_mul(ctx0, ffn_shexp, shared_gate);
+ cb(ffn_shexp, "mtp_ffn_shexp_gated", il);
+
+ cur = ggml_add(ctx0, moe_out, ffn_shexp);
+ } else {
+ cur = moe_out;
+ }
+ cb(cur, "mtp_ffn_out", il);
+
+ cur = ggml_add(ctx0, cur, ffn_residual);
+ cb(cur, "mtp_post_ffn", il);
+
+ ggml_tensor * head_norm_w = layer.nextn.shared_head_norm
+ ? layer.nextn.shared_head_norm
+ : model.output_norm;
+ GGML_ASSERT(head_norm_w && "QWEN3NEXT MTP: missing both nextn.shared_head_norm and output_norm");
+ cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);
+
+ cb(cur, "h_nextn", -1);
+ res->t_h_nextn = cur;
+
+ 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 && "QWEN3NEXT 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);
+}
overview / pkg / ggml / sources / llama.cpp / commitdiff