From 7e8324f5feba1d8be3f966e1791e84107e40cdd7 Mon Sep 17 00:00:00 2001 From: Niko Maroulis Date: Wed, 7 Oct 2026 09:41:50 -0400 Subject: [PATCH 1/3] metal : fix MUL_MAT+ADD fusion when the residual is itself a MUL_MAT (#30100) * metal : fix MUL_MAT+ADD fusion when the residual is itself a MUL_MAT ggml_metal_op_mul_mat_mma picks the residual of a fused MUL_MAT+ADD as "the ADD operand whose op is not MUL_MAT". When both operands of the ADD are mat-mul outputs (x = W1 @ u + W2 @ v), that test is true for both, so the residual resolves to the fused mat-mul's own, never-written output and the kernel adds whatever that buffer holds. The fusion check (ggml_metal_mul_mat_add_operand) already selects the operand by identity; make the encoder do the same. Clef decision models hit this in their head (proj_option_context @ ctx + proj_option_lexical @ lex, 9 option rows): on Metal, /v1/systemone probabilities collapse toward uniform (billing 0.28 where the CPU backend gives 0.977, Cloudflare_clef-flash Q8_0), deterministic per memory layout, correct with GGML_METAL_FUSION_DISABLE=1. Not a quantization issue: the same file is right on CPU. Add a MUL_MAT_ADD mode to test-backend-ops where the residual is a second mat-mul; on Metal it fails 27 of 28 cases before this change (the one pass is f16 n=2, under the MMA row threshold, so nothing fuses). * Update tests/test-backend-ops.cpp Co-authored-by: Georgi Gerganov --------- Co-authored-by: Georgi Gerganov --- ggml/src/ggml-metal/ggml-metal-ops.cpp | 3 ++- tests/test-backend-ops.cpp | 10 +++++++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 2a14cd0c2..597d5300c 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2476,7 +2476,8 @@ static int ggml_metal_op_mul_mat_mma(ggml_metal_op_t ctx, int idx) { if (fuse_add) { dst = ctx->node(idx + n_fuse - 1); - res = dst->src[0]->op == GGML_OP_MUL_MAT ? dst->src[1] : dst->src[0]; + // the residual is the other operand of the ADD, by identity: it can itself be a MUL_MAT output + res = dst->src[0] == op ? dst->src[1] : dst->src[0]; } auto pipeline = ggml_metal_library_get_pipeline_mul_mv_mma_auto(lib, op, fuse_add); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index b080ab496..efc1777ef 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7233,6 +7233,7 @@ enum mul_mat_add_mode { MUL_MAT_ADD_ROW, // mm + a one-row res, which a same-shape fusion must leave alone MUL_MAT_ADD_RES_INPLACE, // res += mm MUL_MAT_ADD_B_INPLACE, // b += mm: the sum overwrites the mat-mul input (m == k) + MUL_MAT_ADD_MM_MM, // mm2 + mm: the residual is itself a mat-mul output, so both operands are MUL_MAT }; static std::string var_to_str(mul_mat_add_mode mode) { @@ -7242,6 +7243,7 @@ static std::string var_to_str(mul_mat_add_mode mode) { case MUL_MAT_ADD_ROW: return "mm+row"; case MUL_MAT_ADD_RES_INPLACE: return "res+=mm"; case MUL_MAT_ADD_B_INPLACE: return "b+=mm"; + case MUL_MAT_ADD_MM_MM: return "mm2+mm"; } return "unknown"; } @@ -7287,6 +7289,12 @@ struct test_mul_mat_add : public test_case { case MUL_MAT_ADD_RES_MM: out = ggml_add(ctx, res, mm); break; case MUL_MAT_ADD_RES_INPLACE: out = ggml_add_inplace(ctx, res, mm); break; case MUL_MAT_ADD_B_INPLACE: out = ggml_add_inplace(ctx, b, mm); break; + case MUL_MAT_ADD_MM_MM: + { + ggml_tensor * a2 = ggml_new_tensor_2d(ctx, type_a, k, m); + ggml_tensor * mm2 = ggml_mul_mat(ctx, a2, b); + out = ggml_add(ctx, mm2, mm); + } break; } ggml_set_name(out, "out"); @@ -10409,7 +10417,7 @@ static std::vector> make_test_cases_eval() { } } for (ggml_type type_a : {GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q5_K, GGML_TYPE_F16}) { - for (mul_mat_add_mode mode : {MUL_MAT_ADD_MM_RES, MUL_MAT_ADD_RES_MM, MUL_MAT_ADD_ROW, MUL_MAT_ADD_RES_INPLACE}) { + for (mul_mat_add_mode mode : {MUL_MAT_ADD_MM_RES, MUL_MAT_ADD_RES_MM, MUL_MAT_ADD_ROW, MUL_MAT_ADD_RES_INPLACE, MUL_MAT_ADD_MM_MM}) { test_cases.emplace_back(new test_mul_mat_add(type_a, 1000, n, 1024, mode)); } test_cases.emplace_back(new test_mul_mat_add(type_a, 2048, n, 2048, MUL_MAT_ADD_B_INPLACE)); From 988190680d5a89fce97de3c20df2c2813731fd61 Mon Sep 17 00:00:00 2001 From: pratiknarola-t Date: Wed, 7 Oct 2026 19:19:37 +0530 Subject: [PATCH 2/3] metal : few-row MMA mat-mul for the remaining src0 types (#30065) The generic few-row MMA kernel works for any type with a 16-weight dequantizer, so it now also takes BF16, Q1_0, Q2_0, MXFP4, Q2_K, Q3_K, TQ2_0 and the IQ types. Each type starts at the row count where it beats the current kernels on an M3 Ultra: 5 rows for TQ2_0, 4 for BF16, 3 for MXFP4, Q2_0, Q2_K and IQ4_NL, and 2 for the others. test-backend-ops perf -o MUL_MAT, m=4096, k=14336, M3 Ultra, time of this change over master (mean of two interleaved runs each): 0.23 to 0.98 from the threshold to 8 rows, 0.24 to 0.33 at 9 to 16 rows, and 0.99 to 1.01 at 1 and 512 rows. --- ggml/src/ggml-metal/ggml-metal-common.cpp | 24 ++++++++++++++++++++ ggml/src/ggml-metal/kernels/mul_mv_mma.metal | 19 ++++++++++++++++ tests/test-backend-ops.cpp | 9 +++++--- 3 files changed, 49 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-common.cpp b/ggml/src/ggml-metal/ggml-metal-common.cpp index 8b065f54a..a00f87c99 100644 --- a/ggml/src/ggml-metal/ggml-metal-common.cpp +++ b/ggml/src/ggml-metal/ggml-metal-common.cpp @@ -80,6 +80,22 @@ static bool ggml_metal_mul_mv_mma_type_supported(enum ggml_type type) { case GGML_TYPE_Q4_K: case GGML_TYPE_Q5_K: case GGML_TYPE_Q6_K: + case GGML_TYPE_BF16: + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q2_0: + case GGML_TYPE_MXFP4: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_TQ2_0: return true; default: return false; @@ -105,10 +121,18 @@ static int64_t ggml_metal_mul_mv_mma_rows_min(enum ggml_type type) { switch (type) { case GGML_TYPE_F32: return 6; + case GGML_TYPE_TQ2_0: + return 5; + case GGML_TYPE_BF16: + return 4; case GGML_TYPE_F16: + case GGML_TYPE_Q2_0: case GGML_TYPE_Q4_K: case GGML_TYPE_Q5_0: case GGML_TYPE_Q5_1: + case GGML_TYPE_Q2_K: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_MXFP4: return 3; default: return 2; diff --git a/ggml/src/ggml-metal/kernels/mul_mv_mma.metal b/ggml/src/ggml-metal/kernels/mul_mv_mma.metal index fcf981a93..5cf245d52 100644 --- a/ggml/src/ggml-metal/kernels/mul_mv_mma.metal +++ b/ggml/src/ggml-metal/kernels/mul_mv_mma.metal @@ -605,4 +605,23 @@ MUL_MV_MMA_GEN("q5_1", block_q5_1, 2, dequantize_q5_1) MUL_MV_MMA_GEN("q4_K", block_q4_K, QK_NL, dequantize_q4_K) MUL_MV_MMA_GEN("q6_K", block_q6_K, QK_NL, dequantize_q6_K) +#if defined(GGML_METAL_HAS_BF16) +MUL_MV_MMA_GEN("bf16", bfloat4x4, 1, dequantize_bf16) +#endif +MUL_MV_MMA_GEN("q1_0", block_q1_0, 8, dequantize_q1_0) +MUL_MV_MMA_GEN("q2_0", block_q2_0, 4, dequantize_q2_0) +MUL_MV_MMA_GEN("mxfp4", block_mxfp4, 2, dequantize_mxfp4) +MUL_MV_MMA_GEN("q2_K", block_q2_K, QK_NL, dequantize_q2_K) +MUL_MV_MMA_GEN("q3_K", block_q3_K, QK_NL, dequantize_q3_K) +MUL_MV_MMA_GEN("iq2_xxs", block_iq2_xxs, QK_NL, dequantize_iq2_xxs) +MUL_MV_MMA_GEN("iq2_xs", block_iq2_xs, QK_NL, dequantize_iq2_xs) +MUL_MV_MMA_GEN("iq2_s", block_iq2_s, QK_NL, dequantize_iq2_s) +MUL_MV_MMA_GEN("iq3_xxs", block_iq3_xxs, QK_NL, dequantize_iq3_xxs) +MUL_MV_MMA_GEN("iq3_s", block_iq3_s, QK_NL, dequantize_iq3_s) +MUL_MV_MMA_GEN("iq1_s", block_iq1_s, QK_NL, dequantize_iq1_s) +MUL_MV_MMA_GEN("iq1_m", block_iq1_m, QK_NL, dequantize_iq1_m) +MUL_MV_MMA_GEN("iq4_nl", block_iq4_nl, 2, dequantize_iq4_nl) +MUL_MV_MMA_GEN("iq4_xs", block_iq4_xs, QK_NL, dequantize_iq4_xs) +MUL_MV_MMA_GEN("tq2_0", block_tq2_0, QK_NL, dequantize_tq2_0) + #undef MUL_MV_MMA_GEN diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index efc1777ef..89554d4d2 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10404,8 +10404,11 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, m, n, k, {1, 1}, {1, 1})); } test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 64, n, 256, {3, 2}, {2, 1})); - for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, GGML_TYPE_Q8_0, - GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K}) { + for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, + GGML_TYPE_Q8_0, GGML_TYPE_Q1_0, GGML_TYPE_Q2_0, GGML_TYPE_MXFP4, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, + GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_TQ2_0, + GGML_TYPE_IQ2_XXS, GGML_TYPE_IQ2_XS, GGML_TYPE_IQ2_S, GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ3_S, + GGML_TYPE_IQ1_S, GGML_TYPE_IQ1_M, GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_XS}) { test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 48, n, 2560, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 1000, n, 1024, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 3000, n, 512, {1, 1}, {1, 1})); @@ -10428,7 +10431,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 1000, n, 1024, {1, 1}, {1, 1}, {0, 1, 2, 3}, 1280)); } // a src1 with a nonzero mean, for the zero points and mins of quantized src0 types - for (ggml_type type_a : {GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q8_0, GGML_TYPE_Q4_K, GGML_TYPE_Q5_K}) { + for (ggml_type type_a : {GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q8_0, GGML_TYPE_Q2_K, GGML_TYPE_Q4_K, GGML_TYPE_Q5_K}) { test_cases.emplace_back(new test_mul_mat_pos(type_a, GGML_TYPE_F32, 1000, n, 1024, {1, 1}, {1, 1})); } } From 448147d42a71d0147c6dc844f99f9ffe7419183a Mon Sep 17 00:00:00 2001 From: Pascal Date: Wed, 7 Oct 2026 16:30:00 +0200 Subject: [PATCH 3/3] llama: share the nextn tensor flags between models (#30097) * llama: share the nextn tensor flags between models Follow-up of the TODO in glm5-next: move the trunk-only and MTP-only detection that each model copied into a nextn_flags helper of llama_model_base. It probes the first trunk layer and the first NextN layer, and adds TENSOR_SKIP when MTP is not loaded. qwen4exp probes hc_attn_norm since it has no attn_norm. deepseek4, nemotron-h, qwen35, qwen35moe, qwen3next and qwen4exp now also accept a trunk-only file, like the other models. * llama: avoid capturing structured bindings in the nextn flags Lambdas that capture structured bindings need C++20, and GCC 15 rejects them under -Werror, so the models read the trunk and MTP flags into plain variables. --- src/llama-model.cpp | 21 +++++++++++++++++++++ src/llama-model.h | 7 +++++++ src/models/bailingmoe3.cpp | 12 +++--------- src/models/cohere2moe.cpp | 15 +++------------ src/models/deepseek2.cpp | 12 +++--------- src/models/deepseek32.cpp | 12 +++--------- src/models/deepseek4.cpp | 6 +++--- src/models/glm-dsa.cpp | 16 +++------------- src/models/glm4-moe.cpp | 12 +++--------- src/models/glm5-next.cpp | 13 +++---------- src/models/hy-v3.cpp | 15 +++------------ src/models/mimo2.cpp | 12 +++--------- src/models/nemotron-h.cpp | 6 +++--- src/models/qwen35.cpp | 6 +++--- src/models/qwen35moe.cpp | 6 +++--- src/models/qwen3next.cpp | 6 +++--- src/models/qwen4exp.cpp | 7 +++---- src/models/step35.cpp | 15 +++------------ 18 files changed, 76 insertions(+), 123 deletions(-) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 1f3b80b08..9f05a66e4 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -3461,6 +3461,27 @@ void llama_model_base::create_tensor_qkv(llama_layer & layer, int bid, } } +llama_model_base::nextn_flags_t llama_model_base::nextn_flags(llama_model_loader & ml, llm_tensor trunk_probe) const { + int trunk = 0; + int mtp = 0; + + // a file without the first trunk layer is MTP-only, a file without the first NextN layer is trunk-only + if (hparams.n_layer_nextn > 0) { + if (ml.get_weight(tn(trunk_probe, "weight", 0).str().c_str()) == nullptr) { + trunk = TENSOR_NOT_REQUIRED; + } + if (ml.get_weight(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", hparams.n_layer()).str().c_str()) == nullptr) { + mtp = TENSOR_NOT_REQUIRED; + } + } + + if (!ml.load_mtp) { + mtp |= TENSOR_SKIP; + } + + return { trunk, mtp }; +} + void llama_model_base::load_swa_pattern(llama_model_loader & ml, uint32_t n_pattern, bool dense_first) { if (ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, false)) { return; diff --git a/src/llama-model.h b/src/llama-model.h index a86cc5693..82124471f 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -849,6 +849,13 @@ struct llama_model_base : public llama_model { int64_t n_embd_, int64_t n_embd_q_, int64_t n_embd_k_, int64_t n_embd_v_, int flags); + // helper: tensor flags for a file that holds the trunk, the NextN layers, or both + struct nextn_flags_t { + int trunk; // TENSOR_NOT_REQUIRED when the file holds only the NextN layers + int mtp; // TENSOR_NOT_REQUIRED when the file holds only the trunk, TENSOR_SKIP when MTP is not loaded + }; + nextn_flags_t nextn_flags(llama_model_loader & ml, llm_tensor trunk_probe = LLM_TENSOR_ATTN_NORM) const; + // helper: read the SWA pattern as one flag per layer, or as a period expanded by set_swa_pattern void load_swa_pattern(llama_model_loader & ml, uint32_t n_pattern, bool dense_first = false); diff --git a/src/models/bailingmoe3.cpp b/src/models/bailingmoe3.cpp index b04b58449..2458b1c1a 100644 --- a/src/models/bailingmoe3.cpp +++ b/src/models/bailingmoe3.cpp @@ -64,15 +64,9 @@ void llama_model_bailingmoe3::load_arch_tensors(llama_model_loader & ml) { const int64_t qk_head_dim = hparams.n_embd_head_k_mla(); const int64_t v_head_dim = hparams.n_embd_head_v_mla(); - 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) { - mtp_flags |= TENSOR_SKIP; - } + const auto nf = nextn_flags(ml); + const int trunk_flags = nf.trunk; + const int mtp_flags = nf.mtp; for (int il = 0; il < n_layer; ++il) { auto & layer = layers[il]; diff --git a/src/models/cohere2moe.cpp b/src/models/cohere2moe.cpp index 3baf3950c..a379e2f60 100644 --- a/src/models/cohere2moe.cpp +++ b/src/models/cohere2moe.cpp @@ -40,18 +40,9 @@ void llama_model_cohere2moe::load_arch_hparams(llama_model_loader & ml) { void llama_model_cohere2moe::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); - // Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP - // tensors live in a separate file. Mark MTP tensors NOT_REQUIRED so the - // trunk loads cleanly. - 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) { - mtp_flags |= TENSOR_SKIP; - } + 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); diff --git a/src/models/deepseek2.cpp b/src/models/deepseek2.cpp index 47fc2f1ac..f1067a516 100644 --- a/src/models/deepseek2.cpp +++ b/src/models/deepseek2.cpp @@ -56,15 +56,9 @@ void llama_model_deepseek2::load_arch_tensors(llama_model_loader & ml) { LLAMA_LOAD_LOCALS; const int64_t n_expert_shared = hparams.n_expert_shared; - 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) { - mtp_flags |= TENSOR_SKIP; - } + const auto nf = nextn_flags(ml); + const int trunk_flags = nf.trunk; + const int mtp_flags = nf.mtp; const bool is_mla = hparams.is_mla(); diff --git a/src/models/deepseek32.cpp b/src/models/deepseek32.cpp index ade9b9ed9..76f763c76 100644 --- a/src/models/deepseek32.cpp +++ b/src/models/deepseek32.cpp @@ -46,15 +46,9 @@ void llama_model_deepseek32::load_arch_hparams(llama_model_loader & ml) { void llama_model_deepseek32::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) { - mtp_flags |= TENSOR_SKIP; - } + const auto nf = nextn_flags(ml); + const int trunk_flags = nf.trunk; + const int mtp_flags = nf.mtp; const bool is_mla = hparams.is_mla(); if (!is_mla) { diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 4edf0c2b4..0974b63b3 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -93,9 +93,9 @@ void llama_model_deepseek4::load_arch_tensors(llama_model_loader & ml) { const int64_t hc_dim = hc_mult * n_embd; const int64_t hc_mix_dim = (2 + hc_mult) * hc_mult; - const bool mtp_only = (n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr); - const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; - const int mtp_flags = ml.load_mtp ? 0 : TENSOR_SKIP; + 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); diff --git a/src/models/glm-dsa.cpp b/src/models/glm-dsa.cpp index 32581f840..3e28ef271 100644 --- a/src/models/glm-dsa.cpp +++ b/src/models/glm-dsa.cpp @@ -75,19 +75,9 @@ void llama_model_glm_dsa::load_arch_tensors(llama_model_loader & ml) { LLAMA_LOAD_LOCALS; const int64_t n_expert_shared = hparams.n_expert_shared; - // MTP-only: the GGUF carries only the NextN/MTP block(s) (user split target/draft). - const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr); - // Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP - // tensors live in a separate file (or were stripped at conversion). Mark - // MTP tensors NOT_REQUIRED so the trunk loads cleanly. - 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) { - mtp_flags |= TENSOR_SKIP; - } + const auto nf = nextn_flags(ml); + const int trunk_flags = nf.trunk; + const int mtp_flags = nf.mtp; const bool is_mla = hparams.is_mla(); if (!is_mla) { diff --git a/src/models/glm4-moe.cpp b/src/models/glm4-moe.cpp index d65cace1f..4b41b5958 100644 --- a/src/models/glm4-moe.cpp +++ b/src/models/glm4-moe.cpp @@ -29,15 +29,9 @@ void llama_model_glm4_moe::load_arch_tensors(llama_model_loader & ml) { LLAMA_LOAD_LOCALS; const int64_t n_expert_shared = hparams.n_expert_shared; - 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) { - mtp_flags |= TENSOR_SKIP; - } + const auto nf = nextn_flags(ml); + const int trunk_flags = nf.trunk; + const int mtp_flags = nf.mtp; GGML_ASSERT(hparams.n_expert > 0 && "n_expert must be > 0 for GLM4_MOE MoE layers"); GGML_ASSERT(hparams.n_expert_used() > 0 && "n_expert_used must be > 0 for GLM4_MOE MoE layers"); diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 4343e6ac6..48b07f1e2 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -64,16 +64,9 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) { const int64_t hc = hparams.dsv4_hc_mult; const int64_t hc_mix_dim = (2 + hc)*hc; - // TODO: consolidate the trunk_only/mtp_only detection and flags shared with glm4-moe and qwen35moe - 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) { - mtp_flags |= TENSOR_SKIP; - } + 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); diff --git a/src/models/hy-v3.cpp b/src/models/hy-v3.cpp index bb0e0a2b8..c4c55c3f9 100644 --- a/src/models/hy-v3.cpp +++ b/src/models/hy-v3.cpp @@ -22,18 +22,9 @@ void llama_model_hy_v3::load_arch_hparams(llama_model_loader & ml) { void llama_model_hy_v3::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); - // Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP - // tensors live in a separate file (e.g. user split target/draft). Mark - // MTP tensors NOT_REQUIRED so the trunk loads cleanly. - 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) { - mtp_flags |= TENSOR_SKIP; - } + 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); diff --git a/src/models/mimo2.cpp b/src/models/mimo2.cpp index 866959c83..a70350191 100644 --- a/src/models/mimo2.cpp +++ b/src/models/mimo2.cpp @@ -25,15 +25,9 @@ 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) { - mtp_flags |= TENSOR_SKIP; - } + 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); diff --git a/src/models/nemotron-h.cpp b/src/models/nemotron-h.cpp index cd90e24d2..cf78ebfbe 100644 --- a/src/models/nemotron-h.cpp +++ b/src/models/nemotron-h.cpp @@ -49,9 +49,9 @@ void llama_model_nemotron_h::load_arch_hparams(llama_model_loader & ml) { void llama_model_nemotron_h::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 int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; - const int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0; + const auto nf = nextn_flags(ml); + const int trunk_flags = nf.trunk; + const int mtp_flags = nf.mtp; // mamba2 Mixer SSM params // NOTE: int64_t for tensor dimensions diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index e484850da..350cea808 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -33,9 +33,9 @@ void llama_model_qwen35::load_arch_hparams(llama_model_loader & ml) { void llama_model_qwen35::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 int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; - int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0; + 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); diff --git a/src/models/qwen35moe.cpp b/src/models/qwen35moe.cpp index 6e28028a8..f9cc8a2c6 100644 --- a/src/models/qwen35moe.cpp +++ b/src/models/qwen35moe.cpp @@ -36,9 +36,9 @@ void llama_model_qwen35moe::load_arch_hparams(llama_model_loader & ml) { void llama_model_qwen35moe::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 int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; - int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0; + 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); diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index 2d56c38ca..620161ac7 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -35,9 +35,9 @@ void llama_model_qwen3next::load_arch_tensors(llama_model_loader & ml) { throw std::runtime_error(arch_name() + " model cannot have zero experts"); } - const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr); - const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; - int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0; + 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); diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index d05f72fc3..88e2e8891 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -175,10 +175,9 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { const int64_t hc_dim = hc * n_embd; const int64_t hc_lr = hparams.hc_low_rank; - // an MTP-only file carries the MTP block, the embeddings and the LM head, but no trunk - const bool mtp_only = n_layer_nextn > 0 && ml.get_weight(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", 0).str().c_str()) == nullptr; - const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; - const int mtp_flags = ml.load_mtp ? 0 : TENSOR_SKIP; + const auto nf = nextn_flags(ml, LLM_TENSOR_HC_ATTN_NORM); + 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); diff --git a/src/models/step35.cpp b/src/models/step35.cpp index f4660bf19..bfa80fab3 100644 --- a/src/models/step35.cpp +++ b/src/models/step35.cpp @@ -37,18 +37,9 @@ void llama_model_step35::load_arch_hparams(llama_model_loader & ml) { void llama_model_step35::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); - // Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP - // tensors live in a separate file (e.g. user split target/draft). Mark - // MTP tensors NOT_REQUIRED so the trunk loads cleanly. - 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) { - mtp_flags |= TENSOR_SKIP; - } + 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);