ggml-org/llama.cpp · #29778

Add Support to maion-coder arch

mailkv23 · merged Oct 15, 20269 files · 211 + / 4 −
gguf-py/gguf/constants.py27 + / 0 −
@@ -674,6 +674,7 @@ class MODEL_ARCH(IntEnum):     SPARK2_5           = auto()     LLAMA_EMBED      = auto()     MAINCODER        = auto()+    MAION_CODER      = auto()     KIMI_LINEAR      = auto()     KIMI_K3          = auto()     TALKIE           = auto()@@ -1465,6 +1466,7 @@ class MODEL_TENSOR(IntEnum):     MODEL_ARCH.SPARK2_5:         "spark2_5",     MODEL_ARCH.LLAMA_EMBED:      "llama-embed",     MODEL_ARCH.MAINCODER:        "maincoder",+    MODEL_ARCH.MAION_CODER:      "maion-coder",     MODEL_ARCH.KIMI_LINEAR:      "kimi-linear",     MODEL_ARCH.KIMI_K3:          "kimi-k3",     MODEL_ARCH.TALKIE:           "talkie",@@ -5693,6 +5695,31 @@ class MODEL_TENSOR(IntEnum):         MODEL_TENSOR.FFN_DOWN,         MODEL_TENSOR.FFN_UP,     ],+    MODEL_ARCH.MAION_CODER: [+        MODEL_TENSOR.TOKEN_EMBD,+        MODEL_TENSOR.OUTPUT_NORM,+        MODEL_TENSOR.OUTPUT,+        MODEL_TENSOR.ATTN_NORM,+        MODEL_TENSOR.ATTN_QKV,+        MODEL_TENSOR.ATTN_Q,+        MODEL_TENSOR.ATTN_Q_NORM,+        MODEL_TENSOR.ATTN_K,+        MODEL_TENSOR.ATTN_K_NORM,+        MODEL_TENSOR.ATTN_V,+        MODEL_TENSOR.ATTN_OUT,+        MODEL_TENSOR.ATTN_POST_NORM,+        MODEL_TENSOR.FFN_NORM,+        MODEL_TENSOR.FFN_POST_NORM,+        MODEL_TENSOR.FFN_GATE,+        MODEL_TENSOR.FFN_DOWN,+        MODEL_TENSOR.FFN_UP,+        MODEL_TENSOR.FFN_GATE_INP,+        MODEL_TENSOR.FFN_GATE_EXP,+        MODEL_TENSOR.FFN_DOWN_EXP,+        MODEL_TENSOR.FFN_UP_EXP,+        MODEL_TENSOR.MOE_LATENT_DOWN,+        MODEL_TENSOR.MOE_LATENT_UP,+    ],     MODEL_ARCH.KIMI_LINEAR: [         MODEL_TENSOR.TOKEN_EMBD,         MODEL_TENSOR.OUTPUT_NORM,
src/llama-arch.cpp1 + / 0 −
@@ -154,6 +154,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {     { LLM_ARCH_SPARK2_5,         "spark2_5"         },     { LLM_ARCH_LLAMA_EMBED,      "llama-embed"      },     { LLM_ARCH_MAINCODER,        "maincoder"        },+    { LLM_ARCH_MAION_CODER,      "maion-coder"      },     { LLM_ARCH_KIMI_LINEAR,      "kimi-linear"      },     { LLM_ARCH_KIMI_K3,          "kimi-k3"          },     { LLM_ARCH_GLM5_NEXT,        "glm5-next"        },
src/llama-arch.h1 + / 0 −
@@ -154,6 +154,7 @@ enum llm_arch {     LLM_ARCH_SPARK2_5,     LLM_ARCH_LLAMA_EMBED,     LLM_ARCH_MAINCODER,+    LLM_ARCH_MAION_CODER,     LLM_ARCH_KIMI_LINEAR,     LLM_ARCH_KIMI_K3,     LLM_ARCH_GLM5_NEXT,
src/llama-kv-cache.cpp1 + / 1 −
@@ -2009,7 +2009,7 @@ class llm_graph_input_k_shift : public llm_graph_input_i { void llm_graph_input_k_shift::set_input(const llama_ubatch * ubatch) {     GGML_UNUSED(ubatch); -    if (k_shift) {+    if (k_shift && k_shift->buffer) {         kv_self->set_input_k_shift(k_shift);     } 
src/llama-model-saver.cpp2 + / 1 −
@@ -225,8 +225,9 @@ void llama_model_saver::add_kv_from_model() {     add_kv(LLM_KV_BLOCK_COUNT,                       hparams.n_layer_all);     add_kv(LLM_KV_LEADING_DENSE_BLOCK_COUNT,         hparams.n_layer_dense_lead);     add_kv(LLM_KV_FEED_FORWARD_LENGTH,               hparams.n_ff_arr, true);-    add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH,        hparams.n_ff_exp());+    add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH,        hparams.n_ff_exp_arr, true);     add_kv(LLM_KV_EXPERT_LATENT_LENGTH,              hparams.n_expert_latent);+    add_kv(LLM_KV_MOE_LATENT_SIZE,                   hparams.moe_latent_size);     add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp);     add_kv(LLM_KV_EXPERT_CHUNK_FEED_FORWARD_LENGTH,  hparams.n_ff_chexp);     add_kv(LLM_KV_SWIGLU_CLAMP_EXP, std::vector<float>(
src/llama-model.cpp3 + / 0 −
@@ -52,6 +52,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params             return new llama_model_llama_embed(params);         case LLM_ARCH_MAINCODER:             return new llama_model_maincoder(params);+        case LLM_ARCH_MAION_CODER:+            return new llama_model_maion_coder(params);         case LLM_ARCH_TALKIE:             return new llama_model_talkie(params);         case LLM_ARCH_DECI:@@ -3556,6 +3558,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {         case LLM_ARCH_GEMMA3:         case LLM_ARCH_GEMMA3N:         case LLM_ARCH_GEMMA4:+        case LLM_ARCH_MAION_CODER:         case LLM_ARCH_GEMMA4_ASSISTANT:         case LLM_ARCH_GEMMA_EMBEDDING:         case LLM_ARCH_GEMMA_EMBEDDING2:
src/models/maion-coder.cppadded149 + / 0 −
@@ -0,0 +1,149 @@+#include "models.h"++void llama_model_maion_coder::load_arch_hparams(llama_model_loader & ml) {+    ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS,       hparams.f_norm_rms_eps);+    ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);+    ml.get_key(LLM_KV_INTERLEAVE_MOE_LAYER_STEP,         hparams.n_moe_layer_step);+    ml.get_key(LLM_KV_MOE_LATENT_SIZE,                   hparams.moe_latent_size);++    if (hparams.n_moe_layer_step == 0 || hparams.moe_latent_size == 0 || hparams.n_expert == 0) {+        throw std::runtime_error("maion-coder requires a positive MoE layer step, latent size and expert count");+    }++    ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);+    uint32_t n_swa_layers = 0;+    ml.get_arr_n(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, n_swa_layers);+    if (hparams.n_swa == 0 || n_swa_layers != hparams.n_layer_all) {+        throw std::runtime_error("maion-coder requires a positive sliding window and a sliding window pattern for each layer");+    }+    ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl);+    hparams.rope_pattern = hparams.is_swa_impl;+    hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;+    hparams.rope_freq_base_train_swa  = hparams.rope_freq_base_train;+    hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;++    type = LLM_TYPE_UNKNOWN;+}++void llama_model_maion_coder::load_arch_tensors(llama_model_loader &) {+    LLAMA_LOAD_LOCALS;++    const int64_t moe_n_embd = hparams.moe_latent_size;++    tok_embd    = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD,  "weight"), {n_embd, n_vocab}, 0);+    output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);+    output      = create_tensor(tn(LLM_TENSOR_OUTPUT,      "weight"), {n_embd, n_vocab}, 0);++    for (int i = 0; i < n_layer; ++i) {+        const bool is_moe_layer = (i + 1) % hparams.n_moe_layer_step == 0;+        auto & layer = layers[i];++        layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);+        create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);+        layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);+        layer.attn_q_norm    = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM,    "weight", i), {n_embd_head_k}, 0);+        layer.attn_k_norm    = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM,    "weight", i), {n_embd_head_k}, 0);+        layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 0);++        layer.ffn_norm      = create_tensor(tn(LLM_TENSOR_FFN_NORM,      "weight", i), {n_embd}, 0);+        layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0);++        if (is_moe_layer) {+            const int64_t n_ff_exp = hparams.n_ff_exp(i);+            layer.ffn_gate_inp    = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,    "weight", i), {n_embd, n_expert}, 0);+            layer.ffn_latent_down = create_tensor(tn(LLM_TENSOR_FFN_LATENT_DOWN, "weight", i), {n_embd, moe_n_embd}, 0);+            layer.ffn_gate_exps   = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS,   "weight", i), {moe_n_embd, n_ff_exp, n_expert}, 0);+            layer.ffn_up_exps     = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,     "weight", i), {moe_n_embd, n_ff_exp, n_expert}, 0);+            layer.ffn_down_exps   = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS,   "weight", i), {n_ff_exp, moe_n_embd, n_expert}, 0);+            layer.ffn_latent_up   = create_tensor(tn(LLM_TENSOR_FFN_LATENT_UP,   "weight", i), {moe_n_embd, n_embd}, 0);+        } else {+            layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);+            layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd, n_ff}, 0);+            layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);+        }+    }+}++std::unique_ptr<llm_graph_context> llama_model_maion_coder::build_arch_graph(const llm_graph_params & params) const {+    return std::make_unique<graph>(*this, params);+}++llama_model_maion_coder::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {+    const int64_t n_embd_head = hparams.n_embd_head_k();++    GGML_ASSERT(n_embd_head == hparams.n_embd_head_v());+    GGML_ASSERT(n_embd_head == n_rot);++    ggml_tensor * inpL = build_inp_embd(model.tok_embd);+    ggml_tensor * inp_pos = build_inp_pos();+    auto * inp_attn = build_attn_inp_kv_iswa();+    ggml_tensor * inp_out_ids = build_inp_out_ids();++    const float kq_scale = 1.0f / sqrtf(float(n_embd_head));++    for (int il = 0; il < n_layer; ++il) {+        // export per-layer input hidden state so speculative drafts (e.g. dflash/eagle3) can read it+        res->t_layer_inp[il] = inpL;++        ggml_tensor * inpSA = inpL;+        ggml_tensor * cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);+        cb(cur, "attn_norm", il);++        auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, n_embd_head, n_head, n_head_kv, il);+        Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);+        Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);++        // Full-attention layers do not use RoPE.+        if (hparams.is_swa(il)) {+            const float freq_base_l  = model.get_rope_freq_base(cparams, il);+            const float freq_scale_l = model.get_rope_freq_scale(cparams, il);+            Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, 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_l, freq_scale_l, ext_factor, attn_factor, beta_fast, beta_slow);+        }++        cur = build_attn(inp_attn, model.layers[il].wo, NULL, model.layers[il].wo_s, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, 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);+        }++        cur = build_norm(cur, model.layers[il].attn_post_norm, NULL, LLM_NORM_RMS, il);+        ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);+        cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il);++        if (model.layers[il].ffn_gate_inp == nullptr) {+            cur = build_ffn(cur,+                    model.layers[il].ffn_up,   NULL, NULL,+                    model.layers[il].ffn_gate, NULL, NULL,+                    model.layers[il].ffn_down, NULL, NULL,+                    NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);+        } else {+            ggml_tensor * router_logits = build_lora_mm(model.layers[il].ffn_gate_inp, cur);+            ggml_tensor * inp_latent = build_lora_mm(model.layers[il].ffn_latent_down, cur);+            cur = build_moe_ffn(inp_latent,+                    model.layers[il].ffn_gate_inp,+                    model.layers[il].ffn_up_exps,+                    model.layers[il].ffn_gate_exps,+                    model.layers[il].ffn_down_exps,+                    nullptr,+                    n_expert, hparams.n_expert_used(il),+                    LLM_FFN_SILU, true, 0.0f,+                    LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX,+                    il, router_logits);+            cur = build_lora_mm(model.layers[il].ffn_latent_up, cur);+        }++        cur = build_norm(cur, model.layers[il].ffn_post_norm, NULL, LLM_NORM_RMS, il);+        cur = ggml_add(ctx0, cur, ffn_inp);+        cur = build_cvec(cur, il);+        inpL = cur;+    }++    ggml_tensor * cur = build_norm(inpL, model.output_norm, NULL, LLM_NORM_RMS, -1);+    cb(cur, "result_norm", -1);+    res->t_embd = cur;+    cur = build_lora_mm(model.output, cur, model.output_s);+    cb(cur, "result_output", -1);+    res->t_logits = cur;+}
src/models/models.h13 + / 0 −
@@ -196,6 +196,19 @@ struct llama_model_maincoder : public llama_model_base { };  +struct llama_model_maion_coder : public llama_model_base {+    llama_model_maion_coder(const struct llama_model_params & params) : llama_model_base(params) {}+    void load_arch_hparams(llama_model_loader & ml) override;+    void load_arch_tensors(llama_model_loader & ml) override;++    struct graph : public llm_graph_context {+        graph(const llama_model & model, const llm_graph_params & params);+    };++    std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;+};++ struct llama_model_talkie : public llama_model_base {     llama_model_talkie(const struct llama_model_params & params) : llama_model_base(params) {}     void load_arch_hparams(llama_model_loader & ml) override;
tests/test-llama-archs.cpp14 + / 2 −
@@ -135,6 +135,9 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {     uint32_t n_layer = 2;     if (arch == LLM_ARCH_LLAMA4) {         n_layer = 4; // hparams.n_no_rope_layer_step is hard-coded to 4+    } else if (arch == LLM_ARCH_MAION_CODER) {+        n_head = 4;+        n_layer = 4; // dense and MoE layers, each with local and full attention     } else if (arch == LLM_ARCH_GEMMA4) {         n_embd = 128;         n_head = 2;@@ -182,7 +185,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {     uint32_t n_head_kv = n_head;     if (arch == LLM_ARCH_QWEN3) {         n_head_kv = 1; // MQA coverage-    } else if (arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE) {+    } else if (arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE || arch == LLM_ARCH_MAION_CODER) {         n_head_kv = 2; // GQA coverage     }     const uint32_t n_embd_head = n_embd / n_head;@@ -289,7 +292,9 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {     ms.add_kv(LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT, uint32_t(8));     ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW,         n_ctx/8); -    if (arch == LLM_ARCH_GEMMA4) {+    if (arch == LLM_ARCH_MAION_CODER) {+        ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, std::vector<uint32_t>{1, 1, 0, 0});+    } else if (arch == LLM_ARCH_GEMMA4) {         ms.add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER,      n_embd/2);         ms.add_kv(LLM_KV_ATTENTION_SHARED_KV_LAYERS,      uint32_t(0));         ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_SWA,        n_embd_head);@@ -443,6 +448,12 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {         }     } +    if (arch == LLM_ARCH_MAION_CODER) {+        ms.add_kv(LLM_KV_MOE_LATENT_SIZE, n_embd / 2);+        ms.add_kv(LLM_KV_EXPERT_COUNT, uint32_t(4));+        ms.add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, std::vector<uint32_t>{n_ff, n_ff / 2, n_ff, n_ff});+    }+     ms.add_kv(LLM_KV_POSNET_EMBEDDING_LENGTH,   n_embd);     ms.add_kv(LLM_KV_POSNET_BLOCK_COUNT,        n_layer);     ms.add_kv(LLM_KV_CONVNEXT_EMBEDDING_LENGTH, n_embd);@@ -654,6 +665,7 @@ static bool check_causal_attn_toggle( static bool moe_mandatory(const llm_arch arch) {     switch (arch) {         case LLM_ARCH_LLAMA4:+        case LLM_ARCH_MAION_CODER:         case LLM_ARCH_COHERE2MOE:         case LLM_ARCH_GROK:         case LLM_ARCH_QWEN2MOE: