diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cu b/ggml/src/ggml-cuda/fattn-mma-f16.cu index 50003e34..91f0a243 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cu +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cu @@ -223,18 +223,18 @@ void ggml_cuda_flash_attn_ext_mma_f16(ggml_backend_cuda_context & ctx, ggml_tens return; } - if (use_gqa_opt && gqa_ratio % 16 == 0 && Q->ne[0] == 128 && Q->ne[0] <= 8) { - if (Q->ne[1] <= 1) { - ggml_cuda_flash_attn_ext_mma_f16_case<128, 1, 16>(ctx, dst); - } - else if (Q->ne[1] <= 2) { - ggml_cuda_flash_attn_ext_mma_f16_case<128, 2, 16>(ctx, dst); - } - else { - ggml_cuda_flash_attn_ext_mma_f16_case<128, 4, 16>(ctx, dst); - } - return; - } + //if (use_gqa_opt && gqa_ratio % 16 == 0 && Q->ne[0] == 128 && Q->ne[0] <= 8) { + // if (Q->ne[1] <= 1) { + // ggml_cuda_flash_attn_ext_mma_f16_case<128, 1, 16>(ctx, dst); + // } + // else if (Q->ne[1] <= 2) { + // ggml_cuda_flash_attn_ext_mma_f16_case<128, 2, 16>(ctx, dst); + // } + // else { + // ggml_cuda_flash_attn_ext_mma_f16_case<128, 4, 16>(ctx, dst); + // } + // return; + //} if (use_gqa_opt && gqa_ratio % 8 == 0) { ggml_cuda_flash_attn_ext_mma_f16_switch_hs<8>(ctx, dst); diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 4d8e7c31..388ccb2c 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -69,6 +69,7 @@ add_library(llama unicode-script-data.cpp # graphs graphs/build_llama.cpp + graphs/build_muse_glimmer.cpp graphs/build_mistral3.cpp graphs/build_deci.cpp graphs/build_baichuan.cpp diff --git a/src/graphs/build_muse_glimmer.cpp b/src/graphs/build_muse_glimmer.cpp new file mode 100644 index 00000000..c563eb11 --- /dev/null +++ b/src/graphs/build_muse_glimmer.cpp @@ -0,0 +1,135 @@ +#include "../llama-build-context.h" +#include "../llama-model.h" +#include "../llama-context.h" + +ggml_cgraph * llm_build_context::build_muse_glimmer() { + ggml_cgraph * gf = new_graph_custom(); + + // TODO: propagate this for the post norm ops + const float post_norm_eps = 1e-8f; + + // mutable variable, needed during the last layer of the computation to skip unused tokens + int32_t n_tokens = this->n_tokens; + + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); + GGML_ASSERT(n_embd_head == hparams.n_rot); + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb); + inpL = ggml_rms_norm(ctx0, inpL, hparams.f_norm_rms_eps); + cb(inpL, "inpL_normed", -1); + + struct ggml_tensor * inp_pos = build_inp_pos(); + + ggml_tensor * KQ_mask = build_inp_KQ_mask(); + ggml_tensor * KQ_mask_swa = build_inp_KQ_mask_swa(); + + auto inp_out_ids = n_tokens > 1 ? build_inp_out_ids() : nullptr; + + const float kq_scale = 1.0f/sqrtf(float(n_embd_head)); + + ggml_tensor * ffn_inp = nullptr; + + post_norm_data pnd; + pnd.f_rms_eps = post_norm_eps; + post_norm_data * pnd_ptr = nullptr; + + bool add_input = model.split_mode == LLAMA_SPLIT_MODE_GRAPH ? false : true; + + int n_active_layer = hparams.n_layer - hparams.nextn_predict_layers; + + for (int il = 0; il < n_active_layer; ++il) { + + bool use_rope = hparams.swa_layers[il]; + auto this_KQ_mask = use_rope ? KQ_mask_swa : KQ_mask; + int this_n_swa = use_rope ? hparams.n_swa : 0; + + if (model.split_mode == LLAMA_SPLIT_MODE_GRAPH && il > 0) { + GGML_ASSERT(pnd.next_input.size() == model.devices.size()); + pnd.norm = model.layers[il-1].ffn_post_norm; + pnd_ptr = &pnd; + } + + cur = build_std_attention(gf, model.layers[il].attn_norm, inpL, + inp_pos, il == n_active_layer - 1 ? inp_out_ids : nullptr, nullptr, + this_KQ_mask, nullptr, nullptr, kq_scale, 0.0f, this_n_swa, il, use_rope, false, add_input, false, false, + model.layers[il].attn_post_norm, -1, post_norm_eps, pnd_ptr); + + if (model.split_mode == LLAMA_SPLIT_MODE_GRAPH) { + pnd_ptr = &pnd; + if (il == 0) { + pnd.next_input.resize(model.devices.size(), inpL); + } else { + GGML_ASSERT(pnd.next_input.size() == model.devices.size()); + } + pnd.norm = model.layers[il].attn_post_norm; + } + + ffn_inp = cur; + + cur = llm_build_ffn(ctx0, lctx, model.layers[il].ffn_norm, ffn_inp, + model.layers[il].ffn_up, nullptr, nullptr, + model.layers[il].ffn_gate, nullptr, nullptr, + model.layers[il].ffn_down, nullptr, nullptr, + nullptr, + LLM_FFN_SILU, LLM_FFN_PAR, cb, il, gf, add_input, false, nullptr, + model.layers[il].ffn_post_norm, post_norm_eps, pnd_ptr); + cb(cur, "ffn_out", il); + + cur = lctx.cvec.apply_to(ctx0, cur, il); + cb(cur, "l_out", il); + + // input for next layer + inpL = cur; + } + cur = inpL; + + if (model.split_mode == LLAMA_SPLIT_MODE_GRAPH) { + GGML_ASSERT(inpL->op == GGML_OP_REDUCE); + int idx = model.default_layer_device[n_active_layer]; + cur = inpL->src[idx]; + if (!cur) { + for (idx = 0; idx < int(model.devices.size()); ++idx) { + if (inpL->src[idx]) { + cur = inpL->src[idx]; break; + } + } + GGML_ASSERT(cur); + } + auto pn_extra = (ggml_split_tensor_t *)model.layers[n_active_layer-1].ffn_post_norm->extra; + GGML_ASSERT(pn_extra && pn_extra->splits[idx]); + cur = ggml_fused_rms_norm(ctx0, cur, pn_extra->splits[idx], pnd.f_rms_eps); + cb(cur, "ffn_post_norm", n_active_layer-1); + GGML_ASSERT(idx < (int)pnd.next_input.size()); + auto add = pnd.next_input[idx]; + if (!add) { + for (int j = 0; j < int(pnd.next_input.size()); ++j) { + if (pnd.next_input[j]) { + add = pnd.next_input[j]; break; + } + } + GGML_ASSERT(add); + } + cur = ggml_add(ctx0, cur, add); + cb(cur, "ffn_final", -1); + } + + // lm_head + cur = build_output(lctx, ctx0, cur, model.output, model.output_norm, cb); + cur = ggml_scale(ctx0, cur, hparams.f_logit_scale); + cb(cur, "output_scaled", -1); + + if (hparams.f_final_logit_softcapping) { + cur = ggml_softcap(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping, hparams.f_final_logit_softcapping); + } + + cb(cur, "result_output", -1); + + ggml_build_forward_expand(gf, cur); + + return gf; +} + diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index dfb83aaf..c5078c7f 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -89,6 +89,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_DFLASH_DRAFT, "dflash-draft" }, { LLM_ARCH_GEMMA4_ASSISTANT,"gemma4-assistant" }, { LLM_ARCH_OPENPANGU, "openpangu" }, + { LLM_ARCH_MUSE_GLIMMER, "muse-glimmer" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; diff --git a/src/llama-arch.h b/src/llama-arch.h index b8687d51..5fa1cb9c 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -87,6 +87,7 @@ enum llm_arch { LLM_ARCH_DFLASH_DRAFT, LLM_ARCH_GEMMA4_ASSISTANT, LLM_ARCH_OPENPANGU, + LLM_ARCH_MUSE_GLIMMER, LLM_ARCH_UNKNOWN, }; diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index a7f6067b..bd52b863 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -1093,6 +1093,28 @@ ggml_tensor * llm_build_context::do_split_norm(ggml_context * ctx, ggml_tensor * return cur; } +static ggml_tensor * llm_do_split_post_norm(ggml_context * ctx, ggml_tensor * cur, post_norm_data * pnd, int id, int n_device, const char * tag, int il_cb, const llm_build_cb & cb) { + auto pn_extra = (ggml_split_tensor_t *)pnd->norm->extra; + GGML_ASSERT(pn_extra && pn_extra->splits[id]); + GGML_ASSERT((int)pnd->next_input.size() == n_device); + cur = ggml_fused_rms_norm(ctx, cur, pn_extra->splits[id], pnd->f_rms_eps); + cb(cur, tag, il_cb); + auto add = pnd->next_input[id]; + if (!add) { + for (int j = 0; j < n_device; ++j) { + if (pnd->next_input[j]) { + add = pnd->next_input[j]; + break; + } + } + GGML_ASSERT(add); + } + cur = ggml_add(ctx, cur, add); + cb(cur, "inp_added", il_cb); + pnd->next_input[id] = cur; + return cur; +} + ggml_tensor * llm_build_context::llm_build_ffn( ggml_context * ctx, llama_context & lctx, @@ -1111,7 +1133,9 @@ ggml_tensor * llm_build_context::llm_build_ffn( llm_ffn_op_type type_op, llm_ffn_gate_type type_gate, const llm_build_cb & cb, int il, ggml_cgraph * graph, bool add_input, - bool is_norm, ggml_tensor * add_extra, ggml_tensor * post_norm) { + bool is_norm, ggml_tensor * add_extra, + ggml_tensor * post_norm, float post_norm_eps, + post_norm_data * pnd) { if (!up_b && !up_s && !gate_b && !gate_s && !down_b && !down_s && up->extra && gate->extra && down->extra && type_gate == LLM_FFN_PAR && @@ -1133,6 +1157,9 @@ ggml_tensor * llm_build_context::llm_build_ffn( GGML_ASSERT((!split_u && !split_g && !split_d) || (split_u && split_g && split_d)); if (!split_u) continue; auto cur = get_input_tensor_sm_graph(ctx, input, id); + if (pnd) { + cur = llm_do_split_post_norm(ctx, cur, pnd, id, u->n_device, "attn_post_norm", il_cb, cb); + } cur = do_split_norm(ctx, cur, ffn_norm, lctx.model.hparams, cb, id, il_cb, is_norm); if (input->op != GGML_OP_REDUCE) { cur->op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t) - 1] = 0xff; @@ -1220,7 +1247,11 @@ ggml_tensor * llm_build_context::llm_build_ffn( cb(cur, "ffn_down_s", il); } if (post_norm) { - cur = llm_build_norm(ctx, cur, lctx.model.hparams, post_norm, NULL, LLM_NORM_RMS, cb, il); + if (post_norm_eps > 0.0f) { + cur = ggml_fused_rms_norm(ctx, cur, post_norm, post_norm_eps); + } else { + cur = llm_build_norm(ctx, cur, lctx.model.hparams, post_norm, NULL, LLM_NORM_RMS, cb, il); + } cb(cur, "ffn_post_normed", il); } if (add_input) { @@ -1364,7 +1395,11 @@ ggml_tensor * llm_build_context::llm_build_ffn( } if (post_norm) { - cur = llm_build_norm(ctx, cur, lctx.model.hparams, post_norm, NULL, LLM_NORM_RMS, cb, il); + if (post_norm_eps > 0.0f) { + cur = ggml_fused_rms_norm(ctx, cur, post_norm, post_norm_eps); + } else { + cur = llm_build_norm(ctx, cur, lctx.model.hparams, post_norm, NULL, LLM_NORM_RMS, cb, il); + } cb(cur, "ffn_post_normed", il); } @@ -2680,6 +2715,10 @@ ggml_cgraph * llm_build_context::llama_build_graph( { result = llm.build_llama(); } break; + case LLM_ARCH_MUSE_GLIMMER: + { + result = llm.build_muse_glimmer(); + } break; case LLM_ARCH_DECI: { result = llm.build_deci(); @@ -3005,7 +3044,7 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens ggml_tensor * input, ggml_tensor * inp_pos, ggml_tensor * inp_out_ids, ggml_tensor * rope_factors_in, ggml_tensor * KQ_mask, ggml_tensor * sinks, ggml_tensor * inp_attn_scale, float KQ_scale, float f_attn_scale, int n_swa, int il, bool do_rope, bool add_graph_split, bool add_input, bool is_norm, bool is_multi, - ggml_tensor * post_norm, int kv_il) { + ggml_tensor * post_norm, int kv_il, float post_norm_eps, post_norm_data * pnd) { float freq_base_l = n_swa > 0 ? hparams.rope_freq_base_train_swa : cparams.rope_freq_base; float freq_scale_l = n_swa > 0 ? hparams.rope_freq_scale_train_swa : hparams.rope_freq_scale_train; @@ -3085,6 +3124,9 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens (split_wq && split_wk && split_wv && split_wo && split_kl && split_vl)); if (!split_wq) continue; auto cur = get_input_tensor_sm_graph(ctx0, input, id); + if (pnd) { + cur = llm_do_split_post_norm(ctx0, cur, pnd, id, wq->n_device, "ffn_post_norm", il_cb, cb); + } cur = do_split_norm(ctx0, cur, the_attn_norm, lctx.model.hparams, cb, id, il_cb, is_norm); auto input_normed = cur; auto the_q_norm = model.layers[il].attn_q_norm ? model.layers[il].attn_q_norm->extra ? @@ -3257,8 +3299,16 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens cur = ggml_mul(ctx0, cur, gate); } } else { - auto gate_3d = ggml_reshape_3d(ctx0, gate, 1, nh, n_tokens); - cur = ggml_fused_mul_unary(ctx0, gate_3d, attn_3d, GGML_UNARY_OP_SIGMOID); + if (gate->ne[0] == n_embd_head_v * nh) { + gate = ggml_sigmoid(ctx0, gate); + cb(gate, "gate", il_cb); + cur = ggml_reshape_2d(ctx0, cur, gate->ne[0], gate->ne[1]); + //gate = ggml_reshape_3d(ctx0, gate, cur->ne[0], cur->ne[1], cur->ne[2]); + cur = ggml_mul(ctx0, cur, gate); + } else { + auto gate_3d = ggml_reshape_3d(ctx0, gate, 1, nh, n_tokens); + cur = ggml_fused_mul_unary(ctx0, gate_3d, attn_3d, GGML_UNARY_OP_SIGMOID); + } } cb(attn_3d, "attn_gated_3d", il_cb); } @@ -3279,6 +3329,9 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens if (inp_out_ids) { cur = ggml_get_rows(ctx0, cur, inp_out_ids); cb(cur, "fa_get_rows", il_cb); + if (pnd) { + pnd->next_input[id] = ggml_get_rows(ctx0, pnd->next_input[id], inp_out_ids); + } } cur = llm_build_lora_mm(lctx, ctx0, split_wo, cur); @@ -3292,6 +3345,7 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens cb(cur, "kqv_wo_biased", il_cb); output_bias_added = true; } + if (cur->ne[1] > 32 && lctx.cparams.reduce_type != GGML_TYPE_F32) { cur = ggml_cast(ctx0, cur, lctx.cparams.reduce_type); } @@ -3393,8 +3447,15 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens cur = ggml_mul(ctx0, cur, gate); } } else { - auto gate_3d = ggml_reshape_3d(ctx0, gate, 1, n_head_l, n_tokens); - cur = ggml_fused_mul_unary(ctx0, gate_3d, attn_3d, GGML_UNARY_OP_SIGMOID); + if (gate->ne[0] == n_head_l) { + auto gate_3d = ggml_reshape_3d(ctx0, gate, 1, n_head_l, n_tokens); + cur = ggml_fused_mul_unary(ctx0, gate_3d, attn_3d, GGML_UNARY_OP_SIGMOID); + } else { + GGML_ASSERT(gate->ne[0] == n_embd_head_v * n_head_l); + gate = ggml_sigmoid(ctx0, gate); + cur = ggml_mul(ctx0, cur, gate); + //cur = ggml_fused_mul_unary(ctx0, gate, cur, GGML_UNARY_OP_SIGMOID); + } } cb(cur, "attn_gated_3d", il); cur = ggml_reshape_2d(ctx0, cur, n_embd_head_v * n_head_l, n_tokens); @@ -3438,7 +3499,11 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens } if (post_norm) { - cur = llm_build_norm(ctx0, cur, hparams, post_norm, NULL, LLM_NORM_RMS, cb, il); + if (post_norm_eps > 0) { + cur = ggml_fused_rms_norm(ctx0, cur, post_norm, post_norm_eps); + } else { + cur = llm_build_norm(ctx0, cur, hparams, post_norm, NULL, LLM_NORM_RMS, cb, il); + } cb(cur, "sa_normed", il); } diff --git a/src/llama-build-context.h b/src/llama-build-context.h index ba489eea..556f7685 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -37,6 +37,12 @@ enum llm_norm_type { LLM_NORM_RMS, }; +struct post_norm_data { + std::vector next_input; + ggml_tensor * norm; + float f_rms_eps; +}; + struct llm_build_context { const llama_model & model; llama_context & lctx; @@ -194,6 +200,8 @@ struct llm_build_context { ggml_cgraph * build_llama(); + ggml_cgraph * build_muse_glimmer(); + ggml_cgraph * build_mistral3(); ggml_cgraph * build_deci(); @@ -505,7 +513,9 @@ struct llm_build_context { llm_ffn_op_type type_op, llm_ffn_gate_type type_gate, const llm_build_cb & cb, int il, ggml_cgraph * graph = nullptr, bool add_input = false, - bool is_norm = false, ggml_tensor * add_extra = nullptr, ggml_tensor * post_norm = nullptr); + bool is_norm = false, ggml_tensor * add_extra = nullptr, + ggml_tensor * post_norm = nullptr, float post_norm_eps = 0.0f, + post_norm_data * pnd = nullptr); static ggml_tensor * build_dspark_logits(llm_build_context & llm, ggml_tensor * base_logits, ggml_tensor * input_tokens, @@ -596,7 +606,8 @@ llm_expert_gating_func_type gating_op, ggml_tensor * inp_pos, ggml_tensor * inp_out_ids, ggml_tensor * rope_factors, ggml_tensor * KQ_mask, ggml_tensor * sinks, ggml_tensor * inp_attn_scale, float KQ_scale, float f_attn_scale, int n_swa, int il, bool do_rope = true, bool add_graph_split = false, bool add_input = false, bool is_norm = false, - bool is_multi = false, ggml_tensor * post_norm = nullptr, int kv_il = -1); + bool is_multi = false, ggml_tensor * post_norm = nullptr, int kv_il = -1, float post_norm_eps = 0.0f, + post_norm_data * pnd = nullptr); static ggml_tensor * build_output(llama_context & lctx, ggml_context * ctx, ggml_tensor * cur, ggml_tensor * output, const llm_build_cb & cb); diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 81eac9e9..68ac76d9 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -1928,6 +1928,30 @@ void llm_load_hparams( validate_dflash_hparams(hparams, model.arch); } } break; + case LLM_ARCH_MUSE_GLIMMER: + { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false); + ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale); + + hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train; + ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false); + + if (!ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.swa_layers, hparams.n_layer, false)) { + uint32_t swa_period = 4; + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, swa_period); + for (int il = 0; il < hparams.n_layer; ++il) { + hparams.swa_layers[il] = (swa_period == 0 || (il % swa_period < swa_period - 1)); + } + } + + switch (hparams.n_layer - hparams.nextn_predict_layers) { + case 52: model.type = e_model::MODEL_30B; break; + default: model.type = e_model::MODEL_UNKNOWN; + } + + } break; default: (void)0; } diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index 790927f5..9d0d8eda 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -41,6 +41,8 @@ struct create_tensors_helper : public create_tensors_helper_interface { bool create_llama_tensors(const LLM_TN & tn); + bool create_muse_glimmer_tensors(const LLM_TN & tn); + bool create_deci_tensors(const LLM_TN & tn); bool create_llama4_tensors(const LLM_TN & tn); @@ -613,6 +615,40 @@ bool create_tensors_helper::create_llama_tensors(const LLM_TN & tn) { return use_mmap_buffer; } +bool create_tensors_helper::create_muse_glimmer_tensors(const LLM_TN & tn) { + LOADING_PRELUDE + create_embd_output(tn, n_embd, n_vocab, true); + + for (int i = 0; i < n_layer; ++i) { + //ggml_context * ctx_layer = ctx_for_layer(i); + ggml_context * ctx_split = ctx_for_layer_split(i); + + auto & layer = model.layers[i]; + + layer.attn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}); + layer.attn_post_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}); + + use_mmap_buffer &= !merge_qkv(tn, i, 1); + + layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}); + + layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0); + layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0); + + layer.wqkv_gate = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_embd_head_k * n_head}, 0); + + layer.ffn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}); + layer.ffn_post_norm = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}); + + // Do we need this? Don't see rope freqa in the official Meta GGUF + //layer.rope_freqs = create_tensor(ctx_split, tn(LLM_TENSOR_ROPE_FREQS, "weight"), {n_embd/n_head/2}, llama_model_loader::TENSOR_NOT_REQUIRED | (i != 0 ? llama_model_loader::TENSOR_DUPLICATED : 0)); + + create_std_ffn(i, tn, layer, n_ff, n_embd, ctx_split); + + } + return use_mmap_buffer; +} + bool create_tensors_helper::create_deci_tensors(const LLM_TN & tn) { LOADING_PRELUDE @@ -4977,6 +5013,8 @@ bool create_tensors_helper::create_tensors() { case LLM_ARCH_GRANITE_MOE: case LLM_ARCH_MISTRAL3: use_mmap_buffer = create_llama_tensors(tn); break; + case LLM_ARCH_MUSE_GLIMMER: + use_mmap_buffer = create_muse_glimmer_tensors(tn); break; case LLM_ARCH_DECI: use_mmap_buffer = create_deci_tensors(tn); break; case LLM_ARCH_LLAMA4: @@ -5221,7 +5259,7 @@ bool create_tensors_helper::create_tensors() { if (layer.attn_norm) { prepare_split_tensors(-1, ctx_split, layer.attn_norm, layer.split_attn_norm, mirror, mem_used); } - if (is_gemma4_model() && layer.attn_post_norm) { + if ((is_gemma4_model() || model.arch == LLM_ARCH_MUSE_GLIMMER) && layer.attn_post_norm) { prepare_split_tensors(-1, ctx_split, layer.attn_post_norm, layer.split_attn_post_norm, mirror, mem_used); } if (layer.rope_freqs) { @@ -5338,9 +5376,7 @@ bool create_tensors_helper::create_tensors() { } if (layer.wqkv_gate) { auto wqkv_gate_split = split_kq; - if (model.arch == LLM_ARCH_LAGUNA && layer.wqkv_gate->ne[1] == layer.wo->ne[0]) { - // Full-width Laguna M.1 gates follow the value/output partition. - // Head-wise gates still follow the K/Q partition collapsed by head size. + if (layer.wqkv_gate->ne[1] == layer.wo->ne[0]) { wqkv_gate_split = split_vo; } else { for (auto & s : wqkv_gate_split) s /= hparams.n_embd_head_k(il); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index f347c744..aaf088db 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -32,6 +32,30 @@ static const std::map> LLM_TENSOR_NA { LLM_TENSOR_FFN_GATE_UP_EXPS,"blk.%d.ffn_gate_up_exps" }, }, }, + { + LLM_ARCH_MUSE_GLIMMER, + { + { LLM_TENSOR_TOKEN_EMBD, "token_embd" }, + { LLM_TENSOR_OUTPUT_NORM, "output_norm" }, + { LLM_TENSOR_OUTPUT, "output" }, + //{ LLM_TENSOR_ROPE_FREQS, "rope_freqs" }, + { LLM_TENSOR_ATTN_NORM, "blk.%d.attn_norm" }, + { LLM_TENSOR_ATTN_POST_NORM, "blk.%d.post_attention_norm" }, + { LLM_TENSOR_ATTN_Q, "blk.%d.attn_q" }, + { LLM_TENSOR_ATTN_K, "blk.%d.attn_k" }, + { LLM_TENSOR_ATTN_V, "blk.%d.attn_v" }, + { LLM_TENSOR_ATTN_OUT, "blk.%d.attn_output" }, + { LLM_TENSOR_ATTN_Q_NORM, "blk.%d.attn_q_norm" }, + { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, + { LLM_TENSOR_ATTN_GATE, "blk.%d.attn_gate" }, + //{ LLM_TENSOR_ATTN_ROT_EMBD, "blk.%d.attn_rot_embd" }, + { LLM_TENSOR_FFN_NORM, "blk.%d.ffn_norm" }, + { LLM_TENSOR_FFN_POST_NORM, "blk.%d.post_ffw_norm" }, + { LLM_TENSOR_FFN_GATE, "blk.%d.ffn_gate" }, + { LLM_TENSOR_FFN_DOWN, "blk.%d.ffn_down" }, + { LLM_TENSOR_FFN_UP, "blk.%d.ffn_up" }, + }, + }, { LLM_ARCH_DECI, { diff --git a/src/llama.cpp b/src/llama.cpp index 40a27c13..9aa5f952 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -3884,6 +3884,7 @@ static bool is_model_split_supported(const llama_model & model) { LLM_ARCH_MISTRAL4, LLM_ARCH_MELLUM, LLM_ARCH_LAGUNA, + LLM_ARCH_MUSE_GLIMMER, }; auto it = k_supported.find(model.arch); return it != k_supported.end(); @@ -8783,6 +8784,7 @@ enum llama_rope_type llama_rope_type(const struct llama_model * model) { case LLM_ARCH_GLM_DSA: case LLM_ARCH_MISTRAL4: case LLM_ARCH_DFLASH: + case LLM_ARCH_MUSE_GLIMMER: return LLAMA_ROPE_TYPE_NORM; // the pairs of head values are offset by n_rot/2