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

@@ -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,11 +347,10 @@ 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);
         }
+    }
 
     // MoE FFN
     cur = llm_build_std_moe_ffn(ctx0, lctx, mtp_layer.ffn_norm, ffn_inp,
@@ -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;

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];

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<uint32_t>(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<ggml_backend_buffer_type_t, int> 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<uint32_t>(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<uint32_t>(il) >= static_cast<uint32_t>(n_layer) - hparams.nextn_predict_layers) {
                 LLAMA_LOG_DEBUG("%s: not splitting MTP tail layer %d (forced non-split)\n", __func__, il);
                 continue;