From 843a9394411509dd5ecc17316e0b786a04aa9382 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Mon, 29 Jun 2026 14:55:59 +0200 Subject: [PATCH] Restore original PR #2047 (#2053) * Split mode graph for GLM4MoE MTP in split_mode_tensor_parallel Analogous to what was done for Qwen35 dense MTP in PR 2027: - ctx_for_layer_split(): return buft_matrix for GLM4_MOE MTP tail layers instead of buft, so split tensor preparation uses the correct split context. - create_glm4_moe_tensors(): remove the ctx_split = ctx_layer override for MTP tail layers. ctx_for_layer_split(i) now returns buft_matrix for GLM4_MOE MTP tails, so regular layer tensors (attn, ffn) are created in the split context. NextN tensors (eh_proj, enorm, hnorm, shared_head_head, shared_head_norm) stay monolithic via ctx_for_layer(). - create_tensors(): add LLM_ARCH_GLM4_MOE to the MTP tail layer splitting exclusion so split processing visits them. - build_glm4_moe_mtp(): pass inp_out_ids to build_std_attention instead of post-processing with ggml_get_rows. Use build_output for the output projection to properly handle split mode. - build_output(): add LLM_ARCH_GLM4_MOE to the is_qwen_mtp check to ensure MTP output is properly materialized. * Skip loading shared_head_head for GLM4MoE MTP Add TENSOR_SKIP flag to shared_head_head so it is never loaded, even when present in the GGUF file. The graph code already falls back to model.output when shared_head_head is nullptr (line 375-376), which frees ~306 MiB on CUDA0 for models that include this tensor (e.g., GLM-Steam-106B). The 355B GLM-4.6 model does not have this tensor and already uses the same fallback path. * cuda-graph: GLM4_MOE - Don't load layer.nextn.embed_tokens --------- Co-authored-by: Nexesenex <124105151+Nexesenex@users.noreply.github.com> --- src/graphs/build_glm4.cpp | 14 +++++--------- src/llama-build-context.cpp | 3 ++- src/llama-load-tensors.cpp | 14 +++++++------- 3 files changed, 14 insertions(+), 17 deletions(-) diff --git a/src/graphs/build_glm4.cpp b/src/graphs/build_glm4.cpp index cb567cc4..af04f122 100644 --- a/src/graphs/build_glm4.cpp +++ b/src/graphs/build_glm4.cpp @@ -319,7 +319,7 @@ struct ggml_tensor * llm_build_context::build_glm4_moe_mtp( ggml_tensor * ffn_inp; if (rope_cache == nullptr) { cur = build_std_attention(gf, mtp_layer.attn_norm, cur, - inp_pos, nullptr, nullptr, + inp_pos, inp_out_ids, nullptr, KQ_mask, nullptr, nullptr, kq_scale, 0.0f, 0, il, true, false, true, false, false, nullptr); ffn_inp = cur; @@ -347,10 +347,9 @@ struct ggml_tensor * llm_build_context::build_glm4_moe_mtp( kq_scale, cb, il); ffn_inp = ggml_add(ctx0, cur, inpSA); cb(ffn_inp, "mtp_ffn_inp", il); - } - - if (inp_out_ids) { - ffn_inp = ggml_get_rows(ctx0, ffn_inp, inp_out_ids); + if (inp_out_ids) { + ffn_inp = ggml_get_rows(ctx0, ffn_inp, inp_out_ids); + } } // MoE FFN @@ -371,15 +370,12 @@ struct ggml_tensor * llm_build_context::build_glm4_moe_mtp( cur = lctx.cvec.apply_to(ctx0, cur, il); cb(cur, "ffn_out", il); - cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.nextn.shared_head_norm, NULL, LLM_NORM_RMS, cb, il); - cb(cur, "result_norm", -1); - // If nextn.shared_head_head is missing (GLM-4.6), use model.output (Main LM Head) ggml_tensor * mtp_head_weights = mtp_layer.nextn.shared_head_head; if (mtp_head_weights == nullptr) { mtp_head_weights = model.output; } - cur = llm_build_lora_mm(lctx, ctx0, mtp_head_weights, cur); + cur = build_output(lctx, ctx0, cur, mtp_head_weights, mtp_layer.nextn.shared_head_norm, cb); cb(cur, "result_output", -1); return cur; diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index 93d9a607..36f5f274 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -2140,7 +2140,8 @@ ggml_tensor * llm_build_context::build_output(llama_context & lctx, ggml_context int idx_out = ggml_backend_sched_get_backend_idx(lctx.sched, lctx.model.output->buffer); if (idx_out >= 0) idx = idx_out; const bool is_qwen_mtp = (lctx.model.arch == LLM_ARCH_QWEN35 || - lctx.model.arch == LLM_ARCH_QWEN35MOE) && lctx.cparams.mtp; + lctx.model.arch == LLM_ARCH_QWEN35MOE || + lctx.model.arch == LLM_ARCH_GLM4_MOE) && lctx.cparams.mtp; if (cur->op == GGML_OP_REDUCE && cur->src[idx] && !is_qwen_mtp) { // avoid copy to main GPU cur->view_src = cur->src[idx]; diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index 0e1464ee..3c045d71 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -182,7 +182,10 @@ struct create_tensors_helper : public create_tensors_helper_interface { inline ggml_context * ctx_for_layer_split(int i, bool force_split = false) const { const bool is_mtp_layer = model.hparams.nextn_predict_layers > 0 && static_cast(i) >= model.hparams.n_layer - model.hparams.nextn_predict_layers; - return is_mtp_layer && !force_split ? ctx_map.at(model.buft_layer[i].buft) : ctx_map.at(model.buft_layer[i].buft_matrix); + if (is_mtp_layer && !force_split && model.arch != LLM_ARCH_GLM4_MOE) { + return ctx_map.at(model.buft_layer[i].buft); + } + return ctx_map.at(model.buft_layer[i].buft_matrix); } std::map buft_layer_count; @@ -2870,10 +2873,6 @@ bool create_tensors_helper::create_glm4_moe_tensors(const LLM_TN & tn) { const bool is_mtp_layer = hparams.nextn_predict_layers > 0 && static_cast(i) >= n_layer - hparams.nextn_predict_layers; - if (is_mtp_layer) { - ctx_split = ctx_layer; - } - int flags = 0; // Skip loading MTP layers if the feature is disabled if (!model.mtp) { @@ -2956,7 +2955,7 @@ bool create_tensors_helper::create_glm4_moe_tensors(const LLM_TN & tn) { layer.nextn.embed_tokens = create_tensor(nextn_host_ctx, tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", final_layer), { n_embd, n_vocab }, - flags | llama_model_loader::TENSOR_NOT_REQUIRED); + flags | llama_model_loader::TENSOR_NOT_REQUIRED | llama_model_loader::TENSOR_SKIP); layer.nextn.enorm = create_tensor(nextn_ctx, tn(LLM_TENSOR_NEXTN_ENORM, "weight", final_layer), { n_embd }, @@ -2968,7 +2967,7 @@ bool create_tensors_helper::create_glm4_moe_tensors(const LLM_TN & tn) { layer.nextn.shared_head_head = create_tensor(nextn_ctx, tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", final_layer), { n_embd, n_vocab }, - flags | llama_model_loader::TENSOR_NOT_REQUIRED); + flags | llama_model_loader::TENSOR_NOT_REQUIRED | llama_model_loader::TENSOR_SKIP); layer.nextn.shared_head_norm = create_tensor(nextn_ctx, tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", final_layer), { n_embd }, @@ -4610,6 +4609,7 @@ bool create_tensors_helper::create_tensors() { // For now only run MTP into the per-layer if (model.mtp && hparams.nextn_predict_layers > 0 && model.arch != LLM_ARCH_QWEN35 && + model.arch != LLM_ARCH_GLM4_MOE && static_cast(il) >= static_cast(n_layer) - hparams.nextn_predict_layers) { LLAMA_LOG_DEBUG("%s: not splitting MTP tail layer %d (forced non-split)\n", __func__, il); continue;