From 5eb1e5d4ef5a51c664d5d31c6fe45272a3a28b3c Mon Sep 17 00:00:00 2001 From: "mikhail.podvitskii" Date: Mon, 8 Jun 2026 10:29:28 +0200 Subject: [PATCH 1/2] model : MTP support for Mellum --- conversion/mellum.py | 74 +++++++++++- gguf-py/gguf/constants.py | 5 + src/models/mellum.cpp | 243 ++++++++++++++++++++++++++++++++------ src/models/models.h | 5 + 4 files changed, 289 insertions(+), 38 deletions(-) diff --git a/conversion/mellum.py b/conversion/mellum.py index 6bcb81df8cfe..9dc2192f3778 100644 --- a/conversion/mellum.py +++ b/conversion/mellum.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Iterable, TYPE_CHECKING +from typing import Callable, Iterable, TYPE_CHECKING import torch @@ -14,6 +14,54 @@ @ModelBase.example("JetBrains/Mellum2-12B-A2.5B-Base") class MellumModel(TextModel): model_arch = gguf.MODEL_ARCH.MELLUM + supports_mtp_export = True + + _n_main_layers: int | None = None + + _experts: list[dict[str, Tensor]] | None = None + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # number of exported MTP blocks, appended after the trunk layers + self.n_nextn = 0 if self.no_mtp else (self.find_hparam(["num_nextn_predict_layers", "num_mtp_layers"], optional=True) or 0) + if self.mtp_only and self.n_nextn == 0: + raise ValueError("--mtp requested, but the checkpoint has no MTP layers (num_nextn_predict_layers is 0)") + if self.n_nextn > 0: + if any(t != "sparse" for t in self.hparams.get("mtp_mlp_layer_types") or []): + raise ValueError("Mellum MTP layers with dense MLP are not supported") + self.block_count += self.n_nextn + 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): + type(self)._n_main_layers = self.hparams["num_hidden_layers"] + return super().index_tensors(remote_hf_model_id=remote_hf_model_id) + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + if (titem := super().filter_tensors(item)) is None: + return None + name, gen = titem + + if name.startswith("mtp."): + if cls.no_mtp: + return None + assert cls._n_main_layers is not None + # mtp.layers.{i}.* -> model.layers.{n_main + i}.*, mtp.{fc,norm,...} -> nextn names of the first MTP block + 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(): + name = f"model.layers.{cls._n_main_layers + int(parts[2])}.{parts[3]}" + elif len(parts) == 3 and parts[1] in remapper: + name = f"model.layers.{cls._n_main_layers}.{remapper[parts[1]]}.{parts[2]}" + elif cls.mtp_only and name not in ("model.embed_tokens.weight", "model.norm.weight", "lm_head.weight"): + return None + + return name, gen def set_gguf_parameters(self): super().set_gguf_parameters() @@ -26,13 +74,31 @@ def set_gguf_parameters(self): if (use_sliding_window is True or use_sliding_window is None) and sliding_window is not None: self.gguf_writer.add_sliding_window(sliding_window) logger.info(f"gguf: sliding window = {sliding_window}") - self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in self.hparams["layer_types"]]) - logger.info(f"gguf: sliding window pattern length = {len(self.hparams['layer_types'])}") + layer_types = list(self.hparams["layer_types"]) + # MTP blocks default to sliding window attention, as in transformers + if self.n_nextn > 0: + layer_types += self.hparams.get("mtp_layer_types") or ["sliding_attention"] * self.n_nextn + self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in layer_types]) + logger.info(f"gguf: sliding window pattern length = {len(layer_types)}") if (norm_topk_prob := self.hparams.get("norm_topk_prob")) is not None: self.gguf_writer.add_expert_weights_norm(norm_topk_prob) - _experts: list[dict[str, Tensor]] | None = None + if self.n_nextn > 0: + self.gguf_writer.add_nextn_predict_layers(self.n_nextn) + + def prepare_metadata(self, vocab_only: bool): + from_dir = self.fname_out.is_dir() + super().prepare_metadata(vocab_only=vocab_only) + + if not self.mtp_only or not from_dir: + return + + output_type: str = self.ftype.name.partition("_")[2] + fname_default: str = gguf.naming_convention( + self.metadata.name, self.metadata.basename, self.metadata.finetune, + self.metadata.version, size_label=None, output_type=output_type, model_type=None) + self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf" def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: if name.find("experts") != -1: diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index a39d046fe174..540b8f100a77 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -5890,6 +5890,11 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_GATE_EXP, MODEL_TENSOR.FFN_DOWN_EXP, MODEL_TENSOR.FFN_UP_EXP, + # NextN/MTP tensors + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], MODEL_ARCH.NANBEIGE: [ MODEL_TENSOR.TOKEN_EMBD, diff --git a/src/models/mellum.cpp b/src/models/mellum.cpp index bc36dfa09cdb..b0008843b83e 100644 --- a/src/models/mellum.cpp +++ b/src/models/mellum.cpp @@ -29,9 +29,13 @@ void llama_model_mellum::load_arch_hparams(llama_model_loader & ml) { } } -void llama_model_mellum::load_arch_tensors(llama_model_loader &) { +void llama_model_mellum::load_arch_tensors(llama_model_loader & ml) { LLAMA_LOAD_LOCALS; + const auto nf = nextn_flags(ml); + const int trunk_flags = nf.trunk; + const int mtp_flags = nf.mtp; + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); // output @@ -45,36 +49,78 @@ void llama_model_mellum::load_arch_tensors(llama_model_loader &) { throw std::runtime_error("n_expert_used must be > 0 for Mellum"); } - for (int i = 0; i < n_layer; ++i) { + auto load_block = [&](int i, int flags) { auto & layer = layers[i]; - layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags); - create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_gqa, n_embd_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, i, n_embd, n_embd_head_k * n_head, n_embd_gqa, n_embd_gqa, flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, flags); - 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", i), {n_embd_head_k}, 0); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, flags); + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, flags); - layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags); - layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags); const int64_t n_ff_exp = hparams.n_ff_exp() ? hparams.n_ff_exp() : n_ff / n_expert_used; - layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { n_embd, n_ff_exp, 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); - layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags); + }; + + for (int i = 0; i < n_layer; ++i) { + load_block(i, trunk_flags); + } + + // MTP block is a regular Mellum decoder block plus the NextN projection and norms + for (int i = n_layer; i < n_layer_all; ++i) { + auto & layer = layers[i]; + + load_block(i, mtp_flags); + + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, mtp_flags); + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, mtp_flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, mtp_flags); + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, mtp_flags); } } std::unique_ptr llama_model_mellum::build_arch_graph(const llm_graph_params & params) const { - if (hparams.swa_type == LLAMA_SWA_TYPE_STANDARD) { + const bool iswa = hparams.swa_type == LLAMA_SWA_TYPE_STANDARD; + + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + if (iswa) { + return std::make_unique>(*this, params); + } + return std::make_unique>(*this, params); + } + + if (iswa) { return std::make_unique>(*this, params); } return std::make_unique>(*this, params); } +static ggml_tensor * mellum_build_rope(const llm_graph_context & g, const llama_model & model, ggml_tensor * cur, ggml_tensor * inp_pos, int il) { + const bool is_swa = g.hparams.is_swa(il); + + const float freq_base_l = model.get_rope_freq_base(g.cparams, il); + const float freq_scale_l = model.get_rope_freq_scale(g.cparams, il); + const float ext_factor_l = is_swa ? 0.0f : g.ext_factor; + const float attn_factor_l = is_swa ? 1.0f : g.attn_factor; + const float beta_fast_l = is_swa ? 32.0f : g.beta_fast; + const float beta_slow_l = is_swa ? 1.0f : g.beta_slow; + + return ggml_rope_ext( + g.ctx0, cur, inp_pos, nullptr, + g.n_rot, g.rope_type, g.n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l + ); +} + template llama_model_mellum::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { const int64_t n_embd_head = hparams.n_embd_head_v(); @@ -122,26 +168,8 @@ llama_model_mellum::graph::graph(const llama_model & model, const llm_grap Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il); cb(Kcur, "Kcur_normed", il); - const bool is_swa = hparams.is_swa(il); - - const float freq_base_l = model.get_rope_freq_base(cparams, il); - const float freq_scale_l = model.get_rope_freq_scale(cparams, il); - const float ext_factor_l = is_swa ? 0.0f : ext_factor; - const float attn_factor_l = is_swa ? 1.0f : attn_factor; - const float beta_fast_l = is_swa ? 32.0f : beta_fast; - const float beta_slow_l = is_swa ? 1.0f : beta_slow; - - Qcur = ggml_rope_ext( - ctx0, Qcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, - ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l - ); - - Kcur = ggml_rope_ext( - ctx0, Kcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, - ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l - ); + Qcur = mellum_build_rope(*this, model, Qcur, inp_pos, il); + Kcur = mellum_build_rope(*this, model, Kcur, inp_pos, il); cb(Qcur, "Qcur", il); cb(Kcur, "Kcur", il); @@ -151,7 +179,7 @@ llama_model_mellum::graph::graph(const llama_model & model, const llm_grap model.layers[il].wo, model.layers[il].wo_b, model.layers[il].wo_s, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il); } - if (il == n_layer - 1 && inp_out_ids) { + if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) { cur = ggml_get_rows(ctx0, cur, inp_out_ids); inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); } @@ -198,6 +226,14 @@ llama_model_mellum::graph::graph(const llama_model & model, const llm_grap model.output_norm, nullptr, LLM_NORM_RMS, -1); + // post-norm hidden state feeds both the LM head and the MTP head + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + if (crop_after_nextn(inp_out_ids)) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + cb(cur, "result_norm", -1); res->t_embd = cur; @@ -208,5 +244,144 @@ llama_model_mellum::graph::graph(const llama_model & model, const llm_grap res->t_logits = cur; } +template +llama_model_mellum::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + GGML_ASSERT(hparams.n_layer_nextn == 1 && "MELLUM 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()); + GGML_ASSERT(n_embd_head == n_rot); + + 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.shared_head_norm && "MTP block missing nextn.shared_head_norm"); + + // TODO: extract in a common llm_graph_context::build_inp_embd_h() + auto inp = std::make_unique(hparams.n_embd_inp(), 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); + + 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"); + + // TODO: make static using `ggml_build_forward_select()` + // see llm_graph_context::build_inp_embd() for reference + ggml_tensor * tok_embd; + ASSERT_EMBD_OR_TOKEN(ubatch); + if (ubatch.token) { + tok_embd = ggml_get_rows(ctx0, model.tok_embd, inp->tokens); + } else { + tok_embd = inp->embd; + } + cb(tok_embd, "mtp_tok_embd", il); + + 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(); + + using inp_attn_type = std::conditional_t; + inp_attn_type * inp_attn = nullptr; + + if constexpr (iswa) { + inp_attn = build_attn_inp_kv_iswa(); + } else { + inp_attn = build_attn_inp_kv(); + } + + ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); + cb(e_norm, "mtp_enorm", il); + + ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); + cb(h_norm, "mtp_hnorm", il); + + ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, 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); + + // self_attention + { + auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, n_head, n_head_kv, il); + + Qcur = build_norm(Qcur, layer.attn_q_norm, nullptr, LLM_NORM_RMS, il); + cb(Qcur, "mtp_Qcur_normed", il); + + Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il); + cb(Kcur, "mtp_Kcur_normed", il); + + Qcur = mellum_build_rope(*this, model, Qcur, inp_pos, il); + Kcur = mellum_build_rope(*this, model, Kcur, inp_pos, il); + + cb(Qcur, "mtp_Qcur", il); + cb(Kcur, "mtp_Kcur", il); + cb(Vcur, "mtp_Vcur", il); + + cur = build_attn(inp_attn, + layer.wo, layer.wo_b, layer.wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il); + cb(cur, "mtp_attn_out", il); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "mtp_ffn_inp", il); + + cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_ffn_norm", il); + + cur = 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, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il, + nullptr, nullptr, + layer.ffn_up_exps_s, + layer.ffn_gate_exps_s, + layer.ffn_down_exps_s); + cb(cur, "mtp_ffn_moe_out", il); + + cur = ggml_add(ctx0, cur, ffn_inp); + cb(cur, "mtp_ffn_out", il); + + cur = build_norm(cur, layer.nextn.shared_head_norm, nullptr, LLM_NORM_RMS, -1); + + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + cb(cur, "mtp_shared_head_norm", -1); + + cur = build_lora_mm(model.output, cur, model.output_s); + + cb(cur, "result_output", -1); + res->t_logits = cur; +} + template struct llama_model_mellum::graph; template struct llama_model_mellum::graph; +template struct llama_model_mellum::graph_mtp; +template struct llama_model_mellum::graph_mtp; diff --git a/src/models/models.h b/src/models/models.h index 88fd07d973cd..99b92ac9d83d 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -466,6 +466,11 @@ struct llama_model_mellum : public llama_model_base { graph(const llama_model & model, const llm_graph_params & params); }; + template + struct graph_mtp : public llm_graph_context { + graph_mtp(const llama_model & model, const llm_graph_params & params); + }; + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; From ab21a6288839d3caa5097faa35587085df767eb3 Mon Sep 17 00:00:00 2001 From: "mikhail.podvitskii" Date: Thu, 15 Oct 2026 20:26:42 +0200 Subject: [PATCH 2/2] model : inline rope construction in Mellum --- src/models/mellum.cpp | 61 ++++++++++++++++++++++++++++--------------- 1 file changed, 40 insertions(+), 21 deletions(-) diff --git a/src/models/mellum.cpp b/src/models/mellum.cpp index b0008843b83e..eaa540310c7f 100644 --- a/src/models/mellum.cpp +++ b/src/models/mellum.cpp @@ -104,23 +104,6 @@ std::unique_ptr llama_model_mellum::build_arch_graph(const ll return std::make_unique>(*this, params); } -static ggml_tensor * mellum_build_rope(const llm_graph_context & g, const llama_model & model, ggml_tensor * cur, ggml_tensor * inp_pos, int il) { - const bool is_swa = g.hparams.is_swa(il); - - const float freq_base_l = model.get_rope_freq_base(g.cparams, il); - const float freq_scale_l = model.get_rope_freq_scale(g.cparams, il); - const float ext_factor_l = is_swa ? 0.0f : g.ext_factor; - const float attn_factor_l = is_swa ? 1.0f : g.attn_factor; - const float beta_fast_l = is_swa ? 32.0f : g.beta_fast; - const float beta_slow_l = is_swa ? 1.0f : g.beta_slow; - - return ggml_rope_ext( - g.ctx0, cur, inp_pos, nullptr, - g.n_rot, g.rope_type, g.n_ctx_orig, freq_base_l, freq_scale_l, - ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l - ); -} - template llama_model_mellum::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { const int64_t n_embd_head = hparams.n_embd_head_v(); @@ -168,8 +151,26 @@ llama_model_mellum::graph::graph(const llama_model & model, const llm_grap Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il); cb(Kcur, "Kcur_normed", il); - Qcur = mellum_build_rope(*this, model, Qcur, inp_pos, il); - Kcur = mellum_build_rope(*this, model, Kcur, inp_pos, il); + const bool is_swa = hparams.is_swa(il); + + const float freq_base_l = model.get_rope_freq_base(cparams, il); + const float freq_scale_l = model.get_rope_freq_scale(cparams, il); + const float ext_factor_l = is_swa ? 0.0f : ext_factor; + const float attn_factor_l = is_swa ? 1.0f : attn_factor; + const float beta_fast_l = is_swa ? 32.0f : beta_fast; + const float beta_slow_l = is_swa ? 1.0f : beta_slow; + + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l + ); + + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l + ); cb(Qcur, "Qcur", il); cb(Kcur, "Kcur", il); @@ -328,8 +329,26 @@ llama_model_mellum::graph_mtp::graph_mtp(const llama_model & model, const Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il); cb(Kcur, "mtp_Kcur_normed", il); - Qcur = mellum_build_rope(*this, model, Qcur, inp_pos, il); - Kcur = mellum_build_rope(*this, model, Kcur, inp_pos, il); + const bool is_swa = hparams.is_swa(il); + + const float freq_base_l = model.get_rope_freq_base(cparams, il); + const float freq_scale_l = model.get_rope_freq_scale(cparams, il); + const float ext_factor_l = is_swa ? 0.0f : ext_factor; + const float attn_factor_l = is_swa ? 1.0f : attn_factor; + const float beta_fast_l = is_swa ? 32.0f : beta_fast; + const float beta_slow_l = is_swa ? 1.0f : beta_slow; + + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l + ); + + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l + ); cb(Qcur, "mtp_Qcur", il); cb(Kcur, "mtp_Kcur", il);