diff --git a/conversion/__init__.py b/conversion/__init__.py index 01c8db296..e3c10d5fc 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -146,6 +146,7 @@ TEXT_MODEL_MAP: dict[str, str] = { "JinaBertModel": "bert", "JinaEmbeddingsV5Model": "bert", "K2HorizonForCausalLM": "k2_horizon", + "Kolibri1ForCausalLM": "kolibri1", "KORMoForCausalLM": "qwen", "KimiK25ForConditionalGeneration": "deepseek", "KimiK3ForConditionalGeneration": "kimi_k3", diff --git a/conversion/base.py b/conversion/base.py index e3d265157..72c1256ba 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -1872,6 +1872,9 @@ class TextModel(ModelBase): if chkhsh == "f5c965e4198da6e7525072d87d1ecf4c51d981fd486c7c4ec728ef9cec7e03fc": # ref: https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B res = "gpt-4o" + if chkhsh == "6e040dfe72e4b85855588c53acf4909ac4e98e55bc3a33cd7db499180dc42a78": + # ref: https://huggingface.co/Aleph-Alpha/Kolibri-1 + res = "qwen2" if chkhsh == "b3d1dd861f1d4c5c0d2569ce36baf3f90fe8a102db3de50dd71ff860d91be3df": # ref: https://huggingface.co/aari1995/German_Semantic_V3 res = "jina-v2-de" diff --git a/conversion/kolibri1.py b/conversion/kolibri1.py new file mode 100644 index 000000000..845501f4e --- /dev/null +++ b/conversion/kolibri1.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from .afmoe import AfmoeModel +from .base import ModelBase, gguf + + +@ModelBase.register("Kolibri1ForCausalLM") +@ModelBase.example("Aleph-Alpha/Kolibri-1") +class Kolibri1Model(AfmoeModel): + model_arch = gguf.MODEL_ARCH.KOLIBRI1 + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # AfmoeModel reads these keys, this config stores the same values under other names + self.hparams["num_shared_experts"] = 1 + self.hparams["route_norm"] = self.hparams["norm_topk_prob"] + self.hparams["intermediate_size"] = self.hparams["moe_intermediate_size"] + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_expert_shared_feed_forward_length(self.hparams["shared_expert_intermediate_size"]) + self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in self.hparams["layer_types"]]) + self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) diff --git a/convert_hf_to_gguf_update.py b/convert_hf_to_gguf_update.py index d49fb8f3f..d3aa3c380 100755 --- a/convert_hf_to_gguf_update.py +++ b/convert_hf_to_gguf_update.py @@ -191,6 +191,8 @@ pre_computed_hashes = [ {"name": "grok-2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/alvarobartt/grok-2-tokenizer", "chkhsh": "66b8d4e19ab16c3bfd89bce5d785fb7e0155e8648708a1f42077cb9fe002c273"}, # muse-spark uses the l4_202k tokenizer, whose pre-tokenizer regex is gpt-4o's {"name": "gpt-4o", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B", "chkhsh": "f5c965e4198da6e7525072d87d1ecf4c51d981fd486c7c4ec728ef9cec7e03fc"}, + # kolibri1 uses a 128k vocabulary with the qwen2 pre-tokenizer regex + {"name": "qwen2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/Aleph-Alpha/Kolibri-1", "chkhsh": "6e040dfe72e4b85855588c53acf4909ac4e98e55bc3a33cd7db499180dc42a78"}, # jina-v2-de variants {"name": "jina-v2-de", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/aari1995/German_Semantic_V3", "chkhsh": "b3d1dd861f1d4c5c0d2569ce36baf3f90fe8a102db3de50dd71ff860d91be3df"}, {"name": "gpt-2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/evilfreelancer/ruGPT3XL", "chkhsh": "0fe1cf6eda062318a1af7270f3331a85c539a01778ff948e24388e949c5282f4"}, diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index bfeab6adf..32e088d88 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -701,6 +701,7 @@ class MODEL_ARCH(IntEnum): QWEN3TTS = auto() POCKETTTS = auto() K2HORIZON = auto() + KOLIBRI1 = auto() class VISION_PROJECTOR_TYPE(IntEnum): @@ -1518,6 +1519,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.QWEN3TTS: "qwen3tts", MODEL_ARCH.POCKETTTS: "pockettts", MODEL_ARCH.K2HORIZON: "k2-horizon", + MODEL_ARCH.KOLIBRI1: "kolibri1", } VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = { @@ -6053,6 +6055,29 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_UP_SHEXP, MODEL_TENSOR.FFN_DOWN_SHEXP, ], + MODEL_ARCH.KOLIBRI1: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_POST_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + ], } # tensors that will not be serialized diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index ab6abf9cb..3586d6528 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -416,6 +416,7 @@ class TensorNameMap: "model.layers.{bid}.post_attention_layernorm", # gemma2 olmo2 # ge "layers.{bid}.post_attention_layernorm", # embeddinggemma "model.layers.{bid}.post_self_attn_layernorm", # glm-4-0414 + "model.layers.{bid}.post_attn_norm", # kolibri1 "model.layers.layers.{bid}.post_mixer_norm.weight", # plamo2 ), @@ -498,6 +499,7 @@ class TensorNameMap: "model.layers.{bid}.post_feedforward_layernorm", # gemma2 olmo2 "layers.{bid}.post_feedforward_layernorm", # embeddinggemma "model.layers.{bid}.post_mlp_layernorm", # glm-4-0414 + "model.layers.{bid}.post_ffn_norm", # kolibri1 "model.layers.layers.{bid}.post_mlp_norm.weight", # plamo2 "model.layers.{bid}.feed_forward.up_proj", "model.layers.{bid}.post_moe_norm", # grok-2 @@ -542,6 +544,7 @@ class TensorNameMap: "model.layers.{bid}.mlp.moe_statics.e_score_correction", # ernie4.5-moe "model.layers.{bid}.mlp.gate.expert_bias", # bailingmoe2 "model.layers.{bid}.mlp.expert_bias", # afmoe + "model.layers.{bid}.moe.router.expert_bias", # kolibri1 "model.layers.{bid}.feed_forward.expert_bias", # lfm2moe "model.layers.{bid}.block_sparse_moe.e_score_correction", # minimax-m2 "backbone.layers.{bid}.mixer.gate.e_score_correction", # nemotron-h-moe diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index bef2d45c0..ffa35bdea 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -166,6 +166,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_POCKETTTS, "pockettts" }, { LLM_ARCH_K2_HORIZON, "k2-horizon" }, { LLM_ARCH_BEAM, "beam" }, + { LLM_ARCH_KOLIBRI1, "kolibri1" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; diff --git a/src/llama-arch.h b/src/llama-arch.h index c305c1766..14ce29ab3 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -171,6 +171,7 @@ enum llm_arch { LLM_ARCH_HRM_TEXT, LLM_ARCH_K2_HORIZON, LLM_ARCH_BEAM, + LLM_ARCH_KOLIBRI1, LLM_ARCH_UNKNOWN, }; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 099c7a052..803a8329e 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -360,6 +360,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_spark2_5(params); case LLM_ARCH_K2_HORIZON: return new llama_model_k2_horizon(params); + case LLM_ARCH_KOLIBRI1: + return new llama_model_kolibri1(params); default: throw std::runtime_error(std::string("unsupported model architecture: '") + llm_arch_name(arch) + "'"); } @@ -1002,6 +1004,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_10B_128x3_66B: return "10B+128x3.66B"; case LLM_TYPE_57B_A14B: return "57B.A14B"; case LLM_TYPE_816B_A42B: return "816B.A42B"; + case LLM_TYPE_78B_A3_5B: return "78B.A3.5B"; case LLM_TYPE_17B_16E: return "17Bx16E (Scout)"; case LLM_TYPE_17B_128E: return "17Bx128E (Maverick)"; case LLM_TYPE_A13B: return "A13B"; @@ -3615,6 +3618,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_MAPLE: case LLM_ARCH_HRM_TEXT: case LLM_ARCH_BEAM: + case LLM_ARCH_KOLIBRI1: return LLAMA_ROPE_TYPE_NEOX; case LLM_ARCH_DFLASH: diff --git a/src/llama-model.h b/src/llama-model.h index 6487b0e9c..a187c2749 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -115,6 +115,7 @@ enum llm_type { LLM_TYPE_10B_128x3_66B, LLM_TYPE_57B_A14B, LLM_TYPE_816B_A42B, + LLM_TYPE_78B_A3_5B, // Kolibri1 LLM_TYPE_17B_16E, // llama4 Scout LLM_TYPE_17B_128E, // llama4 Maverick LLM_TYPE_A13B, diff --git a/src/models/kolibri1.cpp b/src/models/kolibri1.cpp new file mode 100644 index 000000000..acc1dd159 --- /dev/null +++ b/src/models/kolibri1.cpp @@ -0,0 +1,187 @@ +#include "models.h" + +void llama_model_kolibri1::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_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + + if (hparams.n_swa == 0 || hparams.n_expert_shared != 1 || hparams.n_ff_shexp == 0) { + throw std::runtime_error("Kolibri-1: invalid sliding window or shared expert configuration"); + } + + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl); + + hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train; + hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train; + + switch (hparams.n_layer()) { + case 50: type = LLM_TYPE_78B_A3_5B; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_kolibri1::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + const int64_t n_ff_exp = hparams.n_ff_exp(); + const int64_t n_ff_shexp = hparams.n_ff_shexp; + + 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) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_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.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); + + // Routing is sigmoid, the selection bias is added to the logits before top-k. + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {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, n_embd, n_expert}, 0); + + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, 0); + } +} + +llama_model_kolibri1::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_v(); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + + ggml_tensor * cur; + 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) { + const auto & layer = model.layers[il]; + + ggml_tensor * residual = inpL; + + cur = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + // QK-norm, then RoPE on the sliding layers only: the global layers are NoPE. + auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, n_head, n_head_kv, il); + + Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); + Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); + cb(Qcur, "Qcur_normed", il); + cb(Kcur, "Kcur_normed", il); + + 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); + cb(Qcur, "Qcur_rope", il); + + 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); + cb(Kcur, "Kcur_rope", il); + } + + cur = build_attn(inp_attn, + layer.wo, NULL, layer.wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + cb(cur, "attn_out", il); + + cur = build_norm(cur, layer.attn_post_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_post_norm", il); + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + residual = ggml_get_rows(ctx0, residual, inp_out_ids); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, residual); + cb(ffn_inp, "ffn_inp", il); + + cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + // The router runs in f32: the selection adds the bias to the logits, the weights + // are the sigmoid of the unbiased logits, so the same logits feed both. + ggml_tensor * logits = build_lora_mm(layer.ffn_gate_inp, cur); + ggml_prec_set_acc(logits, GGML_PREC_F32); + ggml_prec_set_src(logits, GGML_PREC_F32, 1); + cb(logits, "ffn_moe_logits", il); + + ggml_tensor * selection = ggml_add(ctx0, logits, layer.ffn_exp_probs_b); + cb(selection, "ffn_moe_logits_biased", il); + + ggml_tensor * selected = ggml_argsort_top_k(ctx0, selection, n_expert_used); + cb(selected, "ffn_moe_topk", il); + + ggml_tensor * moe_out = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + nullptr, + n_expert, n_expert_used, + LLM_FFN_SILU, hparams.expert_weights_norm, 1.0f, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il, + logits, + nullptr, nullptr, nullptr, nullptr, selected); + cb(moe_out, "ffn_moe_out", il); + + ggml_tensor * shared = build_ffn(cur, + layer.ffn_up_shexp, nullptr, nullptr, + layer.ffn_gate_shexp, nullptr, nullptr, + layer.ffn_down_shexp, nullptr, nullptr, nullptr, + LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(shared, "ffn_shexp_out", il); + + cur = ggml_add(ctx0, moe_out, shared); + cur = build_norm(cur, layer.ffn_post_norm, NULL, LLM_NORM_RMS, il); + cur = ggml_add(ctx0, cur, ffn_inp); + + inpL = build_cvec(cur, il); + cb(inpL, "l_out", il); + } + + 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; +} + +std::unique_ptr llama_model_kolibri1::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} diff --git a/src/models/models.h b/src/models/models.h index 98f6b6e35..e6cf9d213 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2974,3 +2974,16 @@ struct llama_model_k2_horizon : public llama_model_base { std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; + + +struct llama_model_kolibri1 : public llama_model_base { + llama_model_kolibri1(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 build_arch_graph(const llm_graph_params & params) const override; +}; diff --git a/tests/fusion/MTL.csv b/tests/fusion/MTL.csv index d31a9f23e..9ab1f1f3f 100644 --- a/tests/fusion/MTL.csv +++ b/tests/fusion/MTL.csv @@ -139,6 +139,9 @@ kimi-linear ,0 ,any ,GATED_DELTA_NET+CPY , 1 kimi-linear ,0 ,any ,MUL+ADD , 2 kimi-linear ,0 ,any ,RMS_NORM+MUL , 7 kimi-linear ,0 ,any ,RMS_NORM+SCALE , 2 +kolibri1 ,0 ,any ,MUL+ADD , 4 +kolibri1 ,0 ,any ,RMS_NORM+MUL , 10 +kolibri1 ,0 ,any ,RMS_NORM+MUL+ADD , 3 laguna ,0 ,any ,ADD+ADD , 1 laguna ,0 ,any ,MUL+ADD , 2 laguna ,0 ,any ,RMS_NORM+MUL , 9 diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 14ca203be..143c1335e 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -185,7 +185,8 @@ 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_MUSE_SPARK || arch == LLM_ARCH_AFMOE || arch == LLM_ARCH_MAION_CODER) { + } else if (arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_MUSE_SPARK || arch == LLM_ARCH_AFMOE || arch == LLM_ARCH_MAION_CODER || + arch == LLM_ARCH_KOLIBRI1) { n_head_kv = 2; // GQA coverage } if (arch == LLM_ARCH_BEAM) { @@ -316,7 +317,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(5)); } else if (arch == LLM_ARCH_COHERE2MOE || arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_SPARK2_5 || arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_MUSE_SPARK || arch == LLM_ARCH_GRANITE_SWA || arch == LLM_ARCH_DOTS3NOTE || - arch == LLM_ARCH_MAPLE) { + arch == LLM_ARCH_KOLIBRI1 || arch == LLM_ARCH_MAPLE) { std::vector pattern; pattern.reserve(n_layer); for (uint32_t il = 0; il < n_layer; il++) { @@ -336,6 +337,10 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); } + if (arch == LLM_ARCH_KOLIBRI1) { + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, false); + } + // MSA requires one indexer head per GQA (KV) head, unlike the DSA archs where the // indexer head count is independent of the main attention head count. if (arch == LLM_ARCH_QWEN4EXP || arch == LLM_ARCH_GLM5_NEXT) { @@ -758,6 +763,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_LAGUNA: case LLM_ARCH_MAPLE: case LLM_ARCH_MUSE_SPARK: + case LLM_ARCH_KOLIBRI1: return true; default: return false;