From 8e52631d55636029ad437985ed8c14ff80ab8714 Mon Sep 17 00:00:00 2001 From: AesSedai <7980540+AesSedai@users.noreply.github.com> Date: Thu, 7 May 2026 04:21:58 -0700 Subject: [PATCH] model: Add Mimo v2.5 model support (#22493) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * add mimo-v2.5 support * mimo-v2.5: fix modify_tensors row split * mimi-v2.5: forgot `add_attn_value_scale` plumbing * mimi-v2.5: fix tp dequant to detect tp rows * mimo-v2.5: fix TP iteration to be descending * mimo-v2.5: fix comment * mimo-v2.5: retain fused qkv * mimo-v2.5: missed the attn_value scale during merge * mimo-v2.5: fused QKV needs contiguous for scaling attention value * mimo-v2.5: move `speech_embeddings.` to TextModel filter_tensors * Update src/llama-hparams.h Co-authored-by: Sigbjørn Skjæret * Update src/models/mimo2.cpp Co-authored-by: Sigbjørn Skjæret * Update src/models/mimo2.cpp Co-authored-by: Sigbjørn Skjæret * Update convert_hf_to_gguf.py Co-authored-by: Sigbjørn Skjæret * Update convert_hf_to_gguf.py Co-authored-by: Sigbjørn Skjæret * Update src/models/mimo2.cpp Co-authored-by: Sigbjørn Skjæret * mimo-v2.5: include MTP weights in gguf --------- Co-authored-by: Sigbjørn Skjæret --- convert_hf_to_gguf.py | 149 ++++++++++++++++++++++++++++++++++-- gguf-py/gguf/constants.py | 6 ++ gguf-py/gguf/gguf_writer.py | 3 + src/llama-arch.cpp | 1 + src/llama-arch.h | 1 + src/llama-hparams.h | 2 + src/llama-model-saver.cpp | 1 + src/llama-model.cpp | 1 + src/models/mimo2.cpp | 105 ++++++++++++++++++------- 9 files changed, 235 insertions(+), 34 deletions(-) diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index 9a86aed84..fb1f5dd44 100755 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -1063,7 +1063,7 @@ class TextModel(ModelBase): name, gen = item # Skip multimodal tensors - if name.startswith(("mlp", "vit.", "vpm.", "siglip2.", "conformer.", "merger.", "resampler.", "sound_encoder.", "sound_projection.")) \ + if name.startswith(("mlp", "vit.", "vpm.", "siglip2.", "conformer.", "merger.", "resampler.", "sound_encoder.", "sound_projection.", "speech_embeddings.")) \ or "visual." in name or "vision." in name or "audio." in name or "talker." in name \ or "vision_" in name or "audio_" in name or "sam_model" in name \ or "token2wav." in name or "code2wav." in name \ @@ -9493,10 +9493,126 @@ class MiniMaxM2Model(TextModel): yield from super().modify_tensors(data_torch, name, bid) -@ModelBase.register("MiMoV2FlashForCausalLM") +@ModelBase.register("MiMoV2FlashForCausalLM", "MiMoV2ForCausalLM") class MimoV2Model(TextModel): model_arch = gguf.MODEL_ARCH.MIMO2 + # MiMo V2-Flash, V2.5 and V2.5-Pro all ship 3 trained MTP layers under model.mtp.layers.{0,1,2}. + # The HF config does not expose the count, so it's hardcoded to match the count found in the safetensors. + _n_nextn = 3 + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.block_count = self.hparams["num_hidden_layers"] + self._n_nextn + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) + + @staticmethod + def _tp_aware_qkv_dequant(weight: Tensor, scale_inv: Tensor, + n_q: int, n_kv: int, hd: int, vhd: int, + bs: int = 128) -> Tensor: + # MiMo-V2.5 (TP=4) and V2.5-Pro (TP=8) ship qkv_proj sharded across TP + # ranks; per rank, rows are stacked as [Q_per | K_per | V_per]. + # weight_scale_inv has ceil(rows_per_rank/bs) block-rows per rank (last + # may extend past rows_per_rank with phantom rows not in the weight). + # Naive repeat_interleave aligns rank 0 only and mis-applies scales to + # later ranks once rows_per_rank isn't a multiple of bs. + # Re-group the per-rank [Q_per|K_per|V_per] rows into a single fused + # [Q | K | V] tensor matching the un-sharded original layout. + q_size = n_q * hd + k_size = n_kv * hd + v_size = n_kv * vhd + total_rows = q_size + k_size + v_size + if weight.shape[0] != total_rows: + raise ValueError(f"qkv_proj weight rows {weight.shape[0]} != q+k+v {total_rows}") + + # detect TP from scale_inv block count, descending order so larger matches first + tp = None + for cand in (8, 4): + if total_rows % cand != 0: + continue + rpr = total_rows // cand + bpr = (rpr + bs - 1) // bs + if scale_inv.shape[0] == cand * bpr: + tp = cand + break + if tp is None: + raise ValueError( + f"qkv_proj: cannot detect TP - scale_inv rows {scale_inv.shape[0]}, " + f"q+k+v {total_rows}") + + q_per = q_size // tp + k_per = k_size // tp + v_per = v_size // tp + rows_per_rank = q_per + k_per + v_per + blocks_per_rank = (rows_per_rank + bs - 1) // bs + + scale_inv = scale_inv.float() + # per-row scale-row index: rank * blocks_per_rank + (rr_in_rank // bs) + row_idx = torch.arange(total_rows) + rr = row_idx % rows_per_rank + rank = row_idx // rows_per_rank + scale_row_idx = rank * blocks_per_rank + (rr // bs) + # gather: (total_rows, n_col_blocks) + scale_per_row_block = scale_inv[scale_row_idx] + # expand col-blocks -> cols: each block-col covers `bs` weight cols + scale_full = scale_per_row_block.repeat_interleave(bs, dim=1) + # crop to weight col count (in case last col-block isn't full) + scale_full = scale_full[:, : weight.shape[1]] + dequant = weight.float() * scale_full + + if tp == 1: + return dequant + + # Re-group per-rank [Q_per|K_per|V_per] rows into unified [Q | K | V] + qs, ks, vs = [], [], [] + for r in range(tp): + base = r * rows_per_rank + qs.append(dequant[base : base + q_per]) + ks.append(dequant[base + q_per : base + q_per + k_per]) + vs.append(dequant[base + q_per + k_per : base + rows_per_rank]) + return torch.cat(qs + ks + vs, dim=0) + + def dequant_model(self): + # Capture raw FP8 (weight, scale_inv) lambdas for qkv_proj BEFORE super + # rewrites them with the existing dequant. Replace super's lambda after + # it runs so scale_inv removal still happens via the standard path. + qkv_overrides: dict[str, tuple[Callable, Callable, int]] = {} + qc = self.hparams.get("quantization_config") + if isinstance(qc, dict) and qc.get("quant_method") == "fp8": + pat = re.compile(r"^model\.layers\.(\d+)\.self_attn\.qkv_proj\.weight_scale_inv$") + for name in list(self.model_tensors.keys()): + m = pat.match(name) + if not m: + continue + weight_name = name.removesuffix("_scale_inv") + if weight_name not in self.model_tensors: + continue + qkv_overrides[weight_name] = ( + self.model_tensors[weight_name], + self.model_tensors[name], + int(m.group(1)), + ) + + super().dequant_model() + + if not qkv_overrides: + return + + n_q = self.hparams["num_attention_heads"] + hd = self.hparams["head_dim"] + vhd = self.hparams["v_head_dim"] + hybrid = self.hparams["hybrid_layer_pattern"] + n_layer_text = self.hparams["num_hidden_layers"] + for weight_name, (w_fn, s_fn, bid) in qkv_overrides.items(): + # MTP layers (bid >= n_layer_text) use SWA-style attention dims + is_swa = True if bid >= n_layer_text else hybrid[bid] == 1 + n_kv = self.hparams["swa_num_key_value_heads" if is_swa else "num_key_value_heads"] + self.model_tensors[weight_name] = ( + lambda w_fn=w_fn, s_fn=s_fn, n_q=n_q, n_kv=n_kv, hd=hd, vhd=vhd: + MimoV2Model._tp_aware_qkv_dequant(w_fn(), s_fn(), n_q, n_kv, hd, vhd) + ) + def set_gguf_parameters(self): super().set_gguf_parameters() @@ -9507,11 +9623,14 @@ class MimoV2Model(TextModel): n_head_kv = self.hparams["num_key_value_heads"] n_head_kv_swa = self.hparams["swa_num_key_value_heads"] - n_head_kv_arr = [n_head_kv_swa if use_swa == 1 else n_head_kv for use_swa in self.hparams["hybrid_layer_pattern"]] + # Extend the per-layer pattern with SWA entries for the MTP blocks so the + # runtime arrays (sized to extended block_count) are fully populated. + hybrid = list(self.hparams["hybrid_layer_pattern"]) + [1] * self._n_nextn + n_head_kv_arr = [n_head_kv_swa if use_swa == 1 else n_head_kv for use_swa in hybrid] self.gguf_writer.add_head_count_kv(n_head_kv_arr) self.gguf_writer.add_sliding_window(self.hparams["sliding_window"]) - self.gguf_writer.add_sliding_window_pattern(self.hparams["hybrid_layer_pattern"]) + self.gguf_writer.add_sliding_window_pattern(hybrid) self.gguf_writer.add_value_length(self.hparams["v_head_dim"]) self.gguf_writer.add_expert_count(self.hparams["n_routed_experts"]) self.gguf_writer.add_expert_feed_forward_length(self.hparams["moe_intermediate_size"]) @@ -9521,6 +9640,12 @@ class MimoV2Model(TextModel): self.gguf_writer.add_layer_norm_rms_eps(self.hparams.get("layernorm_epsilon", 1e-5)) + v_scale = self.hparams.get("attention_value_scale") + if v_scale is not None: + self.gguf_writer.add_attn_value_scale(float(v_scale)) + + self.gguf_writer.add_nextn_predict_layers(self._n_nextn) + _experts: list[dict[str, Tensor]] | None = None @classmethod @@ -9530,13 +9655,21 @@ class MimoV2Model(TextModel): if "attention_sink" in name and not name.endswith(".weight"): name += ".weight" - # TODO: mimo v2 does not indicate the number of next-token-prediction layers, therefore we cannot do the same way as GLM4_MOE - if "model.mtp." in name: - return None - return super().filter_tensors((name, gen)) def modify_tensors(self, data_torch, name, bid): + # Remap MTP/NextN tensors to additional layer slots so the standard tensor map handles them. + # HF: model.mtp.layers.{i}.foo -> model.layers.{n_layer_text + i}.foo + m = re.match(r"^model\.mtp\.layers\.(\d+)\.(.*)$", name) + if m is not None: + mtp_idx = int(m.group(1)) + assert mtp_idx < self._n_nextn, f"MTP layer index {mtp_idx} >= _n_nextn ({self._n_nextn})" + rest = m.group(2) + n_layer_text = self.hparams["num_hidden_layers"] + new_bid = n_layer_text + mtp_idx + name = f"model.layers.{new_bid}.{rest}" + bid = new_bid + # process the experts separately if name.find("mlp.experts") != -1: n_experts = self.hparams["n_routed_experts"] diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index c68ec548f..308ebe1f4 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -175,6 +175,7 @@ class Keys: SLIDING_WINDOW = "{arch}.attention.sliding_window" SCALE = "{arch}.attention.scale" OUTPUT_SCALE = "{arch}.attention.output_scale" + VALUE_SCALE = "{arch}.attention.value_scale" TEMPERATURE_LENGTH = "{arch}.attention.temperature_length" KEY_LENGTH_MLA = "{arch}.attention.key_length_mla" VALUE_LENGTH_MLA = "{arch}.attention.value_length_mla" @@ -3868,6 +3869,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3882,6 +3884,10 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_DOWN_EXP, MODEL_TENSOR.FFN_UP_EXP, MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.LAYER_OUT_NORM, + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, ], MODEL_ARCH.STEP35: [ MODEL_TENSOR.TOKEN_EMBD, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 37f4a0454..35fb01470 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -943,6 +943,9 @@ class GGUFWriter: def add_attn_output_scale(self, value: float) -> None: self.add_float32(Keys.Attention.OUTPUT_SCALE.format(arch=self.arch), value) + def add_attn_value_scale(self, value: float) -> None: + self.add_float32(Keys.Attention.VALUE_SCALE.format(arch=self.arch), value) + def add_attn_temperature_length(self, value: int) -> None: self.add_uint32(Keys.Attention.TEMPERATURE_LENGTH.format(arch=self.arch), value) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 633a66fc6..59dde99e3 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -232,6 +232,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, "%s.attention.sliding_window_pattern" }, { LLM_KV_ATTENTION_SCALE, "%s.attention.scale" }, { LLM_KV_ATTENTION_OUTPUT_SCALE, "%s.attention.output_scale" }, + { LLM_KV_ATTENTION_VALUE_SCALE, "%s.attention.value_scale" }, { LLM_KV_ATTENTION_TEMPERATURE_LENGTH, "%s.attention.temperature_length" }, { LLM_KV_ATTENTION_TEMPERATURE_SCALE, "%s.attention.temperature_scale" }, { LLM_KV_ATTENTION_KEY_LENGTH_MLA, "%s.attention.key_length_mla" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index 8f335f5c7..e37d548c9 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -236,6 +236,7 @@ enum llm_kv { LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, LLM_KV_ATTENTION_SCALE, LLM_KV_ATTENTION_OUTPUT_SCALE, + LLM_KV_ATTENTION_VALUE_SCALE, LLM_KV_ATTENTION_TEMPERATURE_LENGTH, LLM_KV_ATTENTION_TEMPERATURE_SCALE, LLM_KV_ATTENTION_KEY_LENGTH_MLA, diff --git a/src/llama-hparams.h b/src/llama-hparams.h index ac7f9ee86..0160a89ca 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -166,6 +166,8 @@ struct llama_hparams { float f_attn_out_scale = 0.0f; uint32_t attn_temp_length = 0; + float f_attn_value_scale = 0.0f; + bool causal_attn = true; bool use_alibi = false; bool attn_soft_cap = false; diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 26864c18e..e83056557 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -268,6 +268,7 @@ void llama_model_saver::add_kv_from_model() { // add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, ???); add_kv(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale); add_kv(LLM_KV_ATTENTION_OUTPUT_SCALE, hparams.f_attn_out_scale); + add_kv(LLM_KV_ATTENTION_VALUE_SCALE, hparams.f_attn_value_scale); add_kv(LLM_KV_ATTENTION_TEMPERATURE_LENGTH, hparams.attn_temp_length); add_kv(LLM_KV_ATTENTION_TEMPERATURE_SCALE, hparams.f_attn_temp_scale); add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 67d648d3b..9d011ff34 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1671,6 +1671,7 @@ void llama_model::print_info() const { LLAMA_LOG_INFO("%s: f_max_alibi_bias = %.1e\n", __func__, hparams.f_max_alibi_bias); LLAMA_LOG_INFO("%s: f_logit_scale = %.1e\n", __func__, hparams.f_logit_scale); LLAMA_LOG_INFO("%s: f_attn_scale = %.1e\n", __func__, hparams.f_attention_scale); + LLAMA_LOG_INFO("%s: f_attn_value_scale = %.4f\n", __func__, hparams.f_attn_value_scale); LLAMA_LOG_INFO("%s: n_ff = %s\n", __func__, print_f([&](uint32_t il) { return hparams.n_ff(il); }, hparams.n_layer).c_str()); LLAMA_LOG_INFO("%s: n_expert = %u\n", __func__, hparams.n_expert); LLAMA_LOG_INFO("%s: n_expert_used = %u\n", __func__, hparams.n_expert_used); diff --git a/src/models/mimo2.cpp b/src/models/mimo2.cpp index eb10d5807..719966166 100644 --- a/src/models/mimo2.cpp +++ b/src/models/mimo2.cpp @@ -10,7 +10,16 @@ void llama_model_mimo2::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false); ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.swa_layers, hparams.n_layer); - switch (hparams.n_layer) { + float value_scale = 0.0f; + if (ml.get_key(LLM_KV_ATTENTION_VALUE_SCALE, value_scale, false) && value_scale != 1.0f) { + hparams.f_attn_value_scale = value_scale; + } + + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false); + GGML_ASSERT(hparams.nextn_predict_layers < hparams.n_layer && "nextn_predict_layers must be < n_layer"); + hparams.n_layer_kv_from_start = hparams.n_layer - hparams.nextn_predict_layers; + + switch (hparams.n_layer - hparams.nextn_predict_layers) { case 48: type = LLM_TYPE_310B_A15B; break; default: type = LLM_TYPE_UNKNOWN; } @@ -25,32 +34,45 @@ void llama_model_mimo2::load_arch_tensors(llama_model_loader &) { output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); + const uint32_t n_nextn = hparams.nextn_predict_layers; + for (int i = 0; i < n_layer; ++i) { auto & layer = layers[i]; uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(i); uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(i); uint32_t n_head = hparams.n_head(i); - create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0); - layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_v * n_head, n_embd }, 0); + // NextN/MTP layers (the last n_nextn blocks) are preserved but disabled pending support + const bool is_nextn = (n_nextn > 0) && (static_cast(i) >= n_layer - n_nextn); + const int skip = is_nextn ? TENSOR_SKIP : 0; - layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); - layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, TENSOR_NOT_REQUIRED); + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, skip); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_v * n_head, n_embd }, skip); - layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, skip); + layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, TENSOR_NOT_REQUIRED | skip); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, skip); // non-MoE branch - layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED); - layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, TENSOR_NOT_REQUIRED); - layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED); + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED | skip); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, TENSOR_NOT_REQUIRED | skip); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED | skip); // MoE branch int64_t n_ff_exp = hparams.n_ff_exp; - layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, TENSOR_NOT_REQUIRED); - layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED); - layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, TENSOR_NOT_REQUIRED); - layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED); + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, TENSOR_NOT_REQUIRED | skip); + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED | skip); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, TENSOR_NOT_REQUIRED | skip); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED | skip); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED | skip); + + if (is_nextn) { + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, skip); + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, skip); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, skip); + layer.layer_out_norm = create_tensor(tn(LLM_TENSOR_LAYER_OUT_NORM, "weight", i), {n_embd}, skip); + } } } @@ -68,7 +90,12 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param auto * inp_attn = build_attn_inp_kv_iswa(); ggml_tensor * inp_out_ids = build_inp_out_ids(); - for (int il = 0; il < n_layer; ++il) { + const float v_scale = hparams.f_attn_value_scale; + + // The last hparams.nextn_predict_layers blocks are MTP heads, currently inactive + const int n_transformer_layers = n_layer - hparams.nextn_predict_layers; + + for (int il = 0; il < n_transformer_layers; ++il) { ggml_tensor * inpSA = inpL; uint32_t n_head_l = hparams.n_head(il); @@ -83,19 +110,39 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); cb(cur, "attn_norm", il); - // compute Q and K and RoPE them - ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur); - cb(Qcur, "Qcur", il); + ggml_tensor * Qcur; + ggml_tensor * Kcur; + ggml_tensor * Vcur; - ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur); - cb(Kcur, "Kcur", il); + if (model.layers[il].wqkv) { + // Fused qkv_proj - Q/K share head_dim_k, V uses head_dim_v + ggml_tensor * qkv = build_lora_mm(model.layers[il].wqkv, cur); + cb(qkv, "wqkv", il); - ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur); - cb(Vcur, "Vcur", il); + const size_t row_k = ggml_row_size(qkv->type, n_embd_head_k); + const size_t row_v = ggml_row_size(qkv->type, n_embd_head_v); + const size_t row_full = qkv->nb[1]; + const size_t k_off = row_k * n_head_l; + const size_t v_off = k_off + row_k * n_head_kv_l; - Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k, n_head_l, n_tokens); - Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head_k, n_head_kv_l, n_tokens); - Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head_v, n_head_kv_l, n_tokens); + Qcur = ggml_view_3d(ctx0, qkv, n_embd_head_k, n_head_l, n_tokens, row_k, row_full, 0); + Kcur = ggml_view_3d(ctx0, qkv, n_embd_head_k, n_head_kv_l, n_tokens, row_k, row_full, k_off); + Vcur = ggml_view_3d(ctx0, qkv, n_embd_head_v, n_head_kv_l, n_tokens, row_v, row_full, v_off); + } else { + // Split path + Qcur = build_lora_mm(model.layers[il].wq, cur); + cb(Qcur, "Qcur", il); + + Kcur = build_lora_mm(model.layers[il].wk, cur); + cb(Kcur, "Kcur", il); + + Vcur = build_lora_mm(model.layers[il].wv, cur); + cb(Vcur, "Vcur", il); + + Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k, n_head_l, n_tokens); + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head_k, n_head_kv_l, n_tokens); + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head_v, n_head_kv_l, n_tokens); + } Qcur = ggml_rope_ext( ctx0, Qcur, inp_pos, nullptr, @@ -118,9 +165,15 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param cur = build_attn(inp_attn, model.layers[il].wo, NULL, model.layers[il].wo_s, Qcur, Kcur, Vcur, nullptr, sinks, nullptr, 1.0f/sqrtf(float(n_embd_head_k)), il); + cb(cur, "attn_out", il); + + if (v_scale) { + cur = ggml_scale(ctx0, cur, v_scale); + cb(cur, "attn_out_scaled", il); + } } - if (il == n_layer - 1 && inp_out_ids) { + if (il == n_transformer_layers - 1 && inp_out_ids) { cur = ggml_get_rows(ctx0, cur, inp_out_ids); inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); }