model : Port Minimax M2 from mainline (#907)

Co-authored-by: firecoperana <firecoperana>

src/llama-build-context.cpp

@@ -7847,7 +7847,6 @@ ggml_cgraph * llm_build_context::build_ernie4_5_moe() {
         cur = llm_build_norm(ctx0, inpL, hparams, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, cb, il);
         cb(cur, "attn_norm", il);
 
-        // self-attention
         // self-attention
         {
             // Q, K, V projections
@@ -7899,7 +7898,6 @@ ggml_cgraph * llm_build_context::build_ernie4_5_moe() {
         }
 
         if (il == n_layer - 1 && inp_out_ids) {
-            ggml_tensor * inp_out_ids = build_inp_out_ids();
             cur = ggml_get_rows(ctx0, cur, inp_out_ids);
             inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
         }
@@ -8366,6 +8364,129 @@ ggml_cgraph * llm_build_context::build_bailingmoe2() {
     return gf;
 }
 
+ggml_cgraph* llm_build_context::build_minimaxm2() {
+    ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model), false);
+    const int64_t n_embd_head = hparams.n_embd_head_v;
+    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
+    // GGML_ASSERT(n_embd_head == hparams.n_rot); this is wrong in case of minimax, head_dim = 128, n_rot = 64
+
+    ggml_tensor * cur;
+    ggml_tensor * inpL;
+
+    inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb);
+
+    ggml_tensor * inp_pos = build_inp_pos();
+
+
+    //auto * inp_attn = build_attn_inp_kv();
+    ggml_tensor * inp_out_ids = build_inp_out_ids();
+    ggml_tensor * KQ_mask = build_inp_KQ_mask();
+
+    for (int il = 0; il < n_layer; ++il) {
+        ggml_tensor* inpSA = inpL;
+
+        cur = inpL;
+
+        // self_attention
+        {
+            cur = llm_build_norm(ctx0, inpL, hparams, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, cb, il);
+            cb(cur, "attn_norm", il);
+
+            // Q, K, V projections
+            ggml_tensor * Qcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wq, cur);
+            cb(Qcur, "Qcur", il);
+
+            ggml_tensor * Kcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wk, cur);
+            cb(Kcur, "Kcur", il);
+
+            ggml_tensor * Vcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wv, cur);
+            cb(Vcur, "Vcur", il);
+
+            Qcur = llm_build_norm(ctx0, Qcur, hparams, model.layers[il].attn_q_norm, NULL,
+                LLM_NORM_RMS, cb, il);
+            cb(Qcur, "Qcur_normed", il);
+
+            Kcur = llm_build_norm(ctx0, Kcur, hparams, model.layers[il].attn_k_norm, NULL,
+                LLM_NORM_RMS, cb, il);
+            cb(Kcur, "Kcur_normed", il);
+
+            // reshape for multi-head
+            Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
+            Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
+            // Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
+
+
+            // apply RoPE
+            Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr,
+                n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+                ext_factor, attn_factor, beta_fast, beta_slow);
+            Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr,
+                n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+                ext_factor, attn_factor, beta_fast, beta_slow);
+            cb(Qcur, "Qcur", il);
+            cb(Kcur, "Kcur", il);
+            cb(Vcur, "Vcur", il);
+
+            cur = llm_build_kv(ctx0, lctx, kv_self, gf,
+                model.layers[il].wo, NULL,
+                Kcur, Vcur, Qcur, KQ_mask,
+                n_tokens, kv_head, n_kv,
+                1.0f / sqrtf(float(n_embd_head)), cb, il);
+        }
+
+        if (il == n_layer - 1 && inp_out_ids) {
+            cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
+        }
+
+        ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+        cb(ffn_inp, "ffn_inp", il);
+
+        // MoE branch
+        cur = llm_build_norm(ctx0, ffn_inp, hparams,
+            model.layers[il].ffn_norm, NULL,
+            LLM_NORM_RMS,cb, il);
+        cb(cur, "ffn_norm", il);
+
+        cur = llm_build_moe_ffn(ctx0, lctx, cur,
+                model.layers[il].ffn_gate_inp,
+                model.layers[il].ffn_up_exps,
+                model.layers[il].ffn_gate_exps,
+                model.layers[il].ffn_down_exps,
+                model.layers[il].ffn_exp_probs_b,
+                n_expert, n_expert_used,
+                LLM_FFN_SILU, true,
+                false, 0,
+                (llm_expert_gating_func_type)hparams.expert_gating_func,
+                cb, il, gf);
+        cb(cur, "ffn_moe_out", il);
+
+        cur = ggml_add(ctx0, cur, ffn_inp);
+
+        cur = lctx.cvec.apply_to(ctx0, cur, il);
+        cb(cur, "l_out", il);
+
+        // input for next layer
+        inpL = cur;
+    }
+
+    cur = inpL;
+
+    cur = llm_build_norm(ctx0, cur,
+        hparams, model.output_norm, NULL,
+        LLM_NORM_RMS, cb, -1);
+
+    cb(cur, "result_norm", -1);
+
+    // lm_head
+    cur = llm_build_lora_mm(lctx, ctx0, model.output, cur);
+
+    cb(cur, "result_output", -1);
+
+    ggml_build_forward_expand(gf, cur);
+    return gf;
+}
+
 ggml_cgraph * llm_build_context::llama_build_graph_defrag(llama_context & lctx, const std::vector<uint32_t> & ids) {
     llama_batch dummy;
     dummy.n_tokens = 0;
@@ -8712,6 +8833,10 @@ ggml_cgraph * llm_build_context::llama_build_graph(
             {
                 result = llm.build_bailingmoe2();
             } break;
+        case LLM_ARCH_MINIMAX_M2:
+            {
+                result = llm.build_minimaxm2();
+            } break;
         default:
             GGML_ABORT("fatal error");
     }