Commit f805c57a2 for llama.cpp
commit f805c57a2d0b7cc171e599303ce2040f6e1bfe15
Author: Aman Gupta <amangupta052@gmail.com>
Date: Fri Sep 25 11:15:33 2026 +0800
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
diff --git a/conversion/mimo.py b/conversion/mimo.py
index 8a2689b96..982594e85 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 3801e5cbe..6fdde70ce 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 a466984af..b6d7aceda 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);