From f805c57a2d0b7cc171e599303ce2040f6e1bfe15 Mon Sep 17 00:00:00 2001 From: Aman Gupta Date: Fri, 25 Sep 2026 11:15:33 +0800 Subject: [PATCH] llama : fix tensor split for fused qkv with uneven K/V head sizes (#29294) * llama : fix tensor split for fused qkv with uneven K/V head sizes Assisted-by: Qwen3.8-27B * fix v granularity * convert: fix mtp conversion * convert: add support for mtp flags * fix loader --- conversion/mimo.py | 34 +++++++++++++++++++++++++++++++--- src/llama-model.cpp | 30 ++++++++++++++++++++++-------- src/models/mimo2.cpp | 4 +++- 3 files changed, 56 insertions(+), 12 deletions(-) diff --git a/conversion/mimo.py b/conversion/mimo.py index 8a2689b965..982594e858 100644 --- a/conversion/mimo.py +++ b/conversion/mimo.py @@ -17,6 +17,7 @@ from .base import MmprojModel, ModelBase, TextModel, gguf, logger @ModelBase.example("XiaomiMiMo/MiMo-V2.5") class MimoV2Model(TextModel): model_arch = gguf.MODEL_ARCH.MIMO2 + supports_mtp_export = True # 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. @@ -25,6 +26,8 @@ class MimoV2Model(TextModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + if self.no_mtp: + self._n_nextn = 0 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) @@ -101,7 +104,7 @@ class MimoV2Model(TextModel): 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$") + pat = re.compile(r"^model\.(mtp\.)?layers\.(\d+)\.self_attn\.qkv_proj\.weight_scale_inv$") for name in list(self.model_tensors.keys()): m = pat.match(name) if not m: @@ -109,10 +112,13 @@ class MimoV2Model(TextModel): weight_name = name.removesuffix("_scale_inv") if weight_name not in self.model_tensors: continue + bid = int(m.group(2)) + if m.group(1) is not None: + bid += self.hparams["num_hidden_layers"] qkv_overrides[weight_name] = ( self.model_tensors[weight_name], self.model_tensors[name], - int(m.group(1)), + bid, ) super().dequant_model() @@ -165,7 +171,8 @@ class MimoV2Model(TextModel): 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) + if self._n_nextn > 0: + self.gguf_writer.add_nextn_predict_layers(self._n_nextn) _MXFP4_EXPERT_RE = re.compile( r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.weight$" @@ -251,11 +258,32 @@ class MimoV2Model(TextModel): def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: name, gen = item + is_mtp = name.startswith("model.mtp.layers.") + if is_mtp and cls.no_mtp: + return None + if cls.mtp_only and not is_mtp and name not in ( + "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", + ): + return None + if "attention_sink" in name and not name.endswith(".weight"): name += ".weight" return super().filter_tensors((name, gen)) + 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, 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 diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 3801e5cbe7..6fdde70ce0 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -666,11 +666,16 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str } if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) { - const int64_t n_embd = hparams.n_head(il) * hparams.n_embd_head_k(il); - const int64_t n_embd_gqa = hparams.n_embd_v_gqa(il); - GGML_ASSERT(hparams.n_embd_k_gqa(il) == n_embd_gqa); - GGML_ASSERT(tensor->ne[axis] == n_embd + 2*n_embd_gqa); - return {{n_embd, 1}, {n_embd_gqa, 2}}; + const int64_t n_embd_q = hparams.n_head(il) * hparams.n_embd_head_k(il); + const int64_t n_embd_k = hparams.n_embd_k_gqa(il); + const int64_t n_embd_v = hparams.n_embd_v_gqa(il); + GGML_ASSERT(tensor->ne[axis] == n_embd_q + n_embd_k + n_embd_v); + if (n_embd_k == n_embd_v) { + return {{n_embd_q, 1}, {n_embd_k, 2}}; + } + // uneven K/V head sizes (e.g. MiMo d_k=192 d_v=128): split K and V as separate + // segments so each device gets whole heads of both + return {{n_embd_q, 1}, {n_embd_k, 1}, {n_embd_v, 1}}; } if (std::regex_match(tensor_name, pattern_ffn_up_weight) || std::regex_match(tensor_name, pattern_ffn_up_bias)) { const int64_t n_ff = hparams.n_ff(il); @@ -763,7 +768,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str } if (std::regex_match(tensor_name, pattern_attn_out_weight)) { GGML_ASSERT(segments.size() == 1); - return {granularity_q}; + return {granularity_head * hparams.n_embd_head_v(il)}; } if (std::regex_match(tensor_name, pattern_attn_gate_weight)) { GGML_ASSERT(segments.size() == 1); @@ -774,20 +779,29 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str } const int64_t granularity_kv = granularity_q / n_gqa; + // the V head size can differ from the K head size (e.g. MiMo d_k=192 d_v=128): + // align V tensors to whole V heads at the same head-index scale as Q and K so all + // three stay in lockstep per device + const int64_t granularity_v = (granularity_kv / hparams.n_embd_head_k(il)) * hparams.n_embd_head_v(il); if (std::regex_match(tensor_name, pattern_kv_weight) || std::regex_match(tensor_name, pattern_kv_bias) || std::regex_match(tensor_name, pattern_kv_cache)) { GGML_ASSERT(segments.size() == 1); - return {granularity_kv}; + const bool is_v = tensor_name.find("attn_v") != std::string::npos || tensor_name.find("cache_v") != std::string::npos; + return {is_v ? granularity_v : granularity_kv}; } if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) { - GGML_ASSERT(segments.size() == 2); // fused full attention layers need Q gate tensors handled like above: // TODO: deduplicate condition [TAG_SPLIT_QGATE_QWEN] if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE || ud->model->arch == LLM_ARCH_QWEN4EXP) { return {std::lcm(2*n_embd_q, blck_size_perf), granularity_kv}; } + if (segments.size() == 3) { + // uneven K/V head sizes: per-segment granularity, V aligned to whole V heads + return {granularity_q, granularity_kv, granularity_v}; + } + GGML_ASSERT(segments.size() == 2); return {granularity_q, granularity_kv}; } } diff --git a/src/models/mimo2.cpp b/src/models/mimo2.cpp index a466984af0..b6d7acedac 100644 --- a/src/models/mimo2.cpp +++ b/src/models/mimo2.cpp @@ -25,8 +25,10 @@ void llama_model_mimo2::load_arch_hparams(llama_model_loader & ml) { void llama_model_mimo2::load_arch_tensors(llama_model_loader & ml) { LLAMA_LOAD_LOCALS; + const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr); const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight"; const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr); + const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0; if (!ml.load_mtp) { @@ -46,7 +48,7 @@ void llama_model_mimo2::load_arch_tensors(llama_model_loader & ml) { uint32_t n_head = hparams.n_head(i); const bool is_nextn = i >= n_layer; - const int flags = is_nextn ? mtp_flags : 0; + const int flags = is_nextn ? mtp_flags : trunk_flags; create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, flags); layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_v * n_head, n_embd }, flags);