Superminaren/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-27 21:46:57 +02:00
Code Issues Packages Projects Releases Wiki Activity

llama : fix tensor split for fused qkv with uneven K/V head sizes (#29294)

Browse Source 
* 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
This commit is contained in:
Aman Gupta
2026-09-25 11:15:33 +08:00
committed by GitHub
parent 4de0926596
commit f805c57a2d
3 changed files with 56 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
conversion/mimo.py
+31 -3
View File
@@ -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
src/llama-model.cpp
+22 -8
View File
@@ -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};
}
}
src/models/mimo2.cpp
+3 -1
View File
@@ -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);