diff --git a/common/arg.cpp b/common/arg.cpp index e1e9ef36e..ac59481b1 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -3363,7 +3363,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex {"--tools"}, "TOOL1,TOOL2,...", "experimental: whether to enable built-in tools for AI agents - do not enable in untrusted environments (default: no tools)\n" "specify \"all\" to enable all tools\n" - "available tools: read_file, file_glob_search, grep_search, exec_shell_command, write_file, edit_file, get_datetime, get_info\n" + "available tools: read_file, file_glob_search, grep_search, exec_shell_command, write_file, edit_file, get_info\n" "note: for security reasons, this will limit --cors-origins to localhost by default", [](common_params & params, const std::string & value) { params.server_tools = parse_csv_row(value); diff --git a/common/chat-diff-analyzer.cpp b/common/chat-diff-analyzer.cpp index 7db1dcb0f..d6d2af2d5 100644 --- a/common/chat-diff-analyzer.cpp +++ b/common/chat-diff-analyzer.cpp @@ -193,6 +193,14 @@ static std::vector void { + if (tmpl.src.find("Bailing V3 chat template") != std::string::npos) { + analysis.tools.arguments.value_suffix = trim_whitespace(analysis.tools.arguments.value_suffix); + analysis.tools.arguments.tolerate_intertag_whitespace = true; + LOG_DBG(ANSI_ORANGE "[Patch: Bailing V3]\n" ANSI_RESET); + } + }, }); diff --git a/common/chat.cpp b/common/chat.cpp index 803b4021b..2fa9d2e07 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -479,36 +479,80 @@ std::vector common_chat_msgs_parse_oaicompat(const json & messa return msgs; } +struct messages_inp_normalizer { + const jinja::caps & caps; + + messages_inp_normalizer(const jinja::caps & c) : caps(c) {} + + // handle supports_string_content / supports_typed_content + // if string=true and array=false, convert array to string + // if string=false and array=true, convert string to array + // if both are true, do nothing + json normalize(const json & messages) { + bool only_string = caps.supports_string_content && !caps.supports_typed_content; + bool only_typed = !caps.supports_string_content && caps.supports_typed_content; + if ((!only_string && !only_typed) || !messages.is_array()) { + return messages; + } + json normalized = json::array(); + for (const auto & msg : messages) { + json copy = msg; + auto it = copy.find("content"); + if (it != copy.end()) { + if (only_typed && it->is_string()) { + *it = json::array({ + json{ + {"type", "text"}, + {"text", it->get()}, + } + }); + } else if (only_string && it->is_array()) { + *it = concat_content_parts(*it); + } + } + normalized.push_back(std::move(copy)); + } + return normalized; + } + + // join parts with newline, do not add newline before or after media markers + static std::string concat_content_parts(const json & parts) { + std::string text; + bool last_was_media_marker = false; + for (const auto & part : parts) { + std::string type = part.value("type", ""); + bool add_new_line = true; + if (type == "text") { + add_new_line = !last_was_media_marker && !text.empty(); + last_was_media_marker = false; + } else if (type == "media_marker") { + add_new_line = false; + last_was_media_marker = true; + } else { + LOG_WRN("Ignoring content part type: %s\n", type.c_str()); + continue; + } + + if (add_new_line) { + text += '\n'; + } + + text += part.value("text", ""); + } + return text; + } +}; + static json render_message_to_json(const std::vector & msgs, const jinja::caps & c) { if (!c.supports_string_content && !c.supports_typed_content) { LOG_WRN("%s: Neither string content nor typed content is supported by the template. This is unexpected and may lead to issues.\n", __func__); } - bool only_string_accepted = c.supports_string_content && !c.supports_typed_content; - bool only_typed_accepted = !c.supports_string_content && c.supports_typed_content; - json messages = json::array(); for (const auto & msg : msgs) { - if (only_string_accepted) { - json jmsg = msg.to_json_oaicompat(/* concat_typed_text= */ true); - messages.push_back(jmsg); - } else if (only_typed_accepted) { - json jmsg = msg.to_json_oaicompat(/* concat_typed_text= */ false); - if (jmsg.at("content").is_string()) { - jmsg["content"] = json::array({ - json{ - {"type", "text"}, - {"text", jmsg.at("content").get()}, - } - }); - } - messages.push_back(jmsg); - } else { - json jmsg = msg.to_json_oaicompat(/* concat_typed_text= */ false); - messages.push_back(jmsg); - } + messages.push_back(msg.to_json_oaicompat(/* concat_typed_text= */ false)); } - return messages; + return messages_inp_normalizer(c).normalize(messages); } // DEPRECATED: only used in tests @@ -906,8 +950,11 @@ static std::string common_chat_template_direct_apply_impl( const std::optional & additional_context = std::nullopt) { jinja::context ctx(tmpl.source()); + // messages_override is already built for this template, do not touch its content parts nlohmann::ordered_json inp = nlohmann::ordered_json{ - {"messages", messages_override.has_value() ? *messages_override : inputs.messages}, + {"messages", messages_override.has_value() + ? *messages_override + : messages_inp_normalizer(tmpl.original_caps()).normalize(inputs.messages)}, {"bos_token", tmpl.bos_token()}, {"eos_token", tmpl.eos_token()}, {"enable_thinking", inputs.enable_thinking}, @@ -971,14 +1018,12 @@ static std::string common_chat_template_generation_prompt_impl( const std::optional & tools_override = std::nullopt, const std::optional & additional_context = std::nullopt) { - auto adjusted_messages = messages_override ? *messages_override : inputs.messages; - autoparser::generation_params params = inputs; params.add_generation_prompt = false; params.continue_final_message = COMMON_CHAT_CONTINUATION_NONE; - std::string no_gen_prompt = common_chat_template_direct_apply_impl(tmpl, params, adjusted_messages, tools_override, additional_context); + std::string no_gen_prompt = common_chat_template_direct_apply_impl(tmpl, params, messages_override, tools_override, additional_context); params.add_generation_prompt = true; - std::string gen_prompt = common_chat_template_direct_apply_impl(tmpl, params, adjusted_messages, tools_override, additional_context); + std::string gen_prompt = common_chat_template_direct_apply_impl(tmpl, params, messages_override, tools_override, additional_context); size_t prefix_len = 0; size_t min_size = std::min(no_gen_prompt.size(), gen_prompt.size()); @@ -2339,6 +2384,179 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha return data; } +// Kimi K3 - XTML tagged format, built by open_tag/close_tag macros: +// open_tag(t, attrs) = <|open|>t k="v"...<|sep|> close_tag(t) = <|close|>t<|sep|> +// assistant := [think] [response] [tools] close_tag(message) <|end_of_msg|> +// the generation prompt already opens the think (or response) section, so the +// section opener is optional here - same as Kimi K2 Thinking +static common_chat_params common_chat_params_init_kimi_k3(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.supports_thinking = true; + + const std::string SEP = "<|sep|>"; + const std::string MSG_START = "<|open|>message role=\"assistant\"<|sep|>"; + const std::string THINK_START = "<|open|>think<|sep|>"; + const std::string THINK_END = "<|close|>think<|sep|>"; + const std::string RESP_START = "<|open|>response<|sep|>"; + const std::string RESP_END = "<|close|>response<|sep|>"; + const std::string TOOLS_START = "<|open|>tools<|sep|>"; + const std::string TOOLS_END = "<|close|>tools<|sep|>"; + const std::string CALL_START = "<|open|>call tool=\""; + const std::string CALL_END = "<|close|>call<|sep|>"; + const std::string ARG_START = "<|open|>argument key=\""; + const std::string ARG_END = "<|close|>argument<|sep|>"; + const std::string MSG_END = "<|close|>message<|sep|>"; + const std::string EOM_TOKEN = "<|end_of_msg|>"; + + // only the markers are special tokens. tag names ("think", "response", ...) are + // normal tokens and must not be preserved, or prose with those words is broken + data.preserved_tokens = { + "<|open|>", + "<|close|>", + "<|sep|>", + "<|end_of_msg|>", + }; + + data.thinking_start_tag = THINK_START; + data.thinking_end_tags = { THINK_END }; + + // per-role message-start delimiters. user/assistant messages only have the role + // attribute, so the full opener is used. system and tool messages have more + // attributes, so those delimiters stop after the closing quote of the role + data.message_delimiters = { + { COMMON_CHAT_ROLE_ASSISTANT, "<|open|>message role=\"assistant\"<|sep|>" }, + { COMMON_CHAT_ROLE_USER, "<|open|>message role=\"user\"<|sep|>" }, + { COMMON_CHAT_ROLE_TOOL, "<|open|>message role=\"tool\"" }, + { COMMON_CHAT_ROLE_SYSTEM, "<|open|>message role=\"system\"" }, + }; + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE; + + if (inputs.has_continuation()) { + const auto & msg = inputs.continue_msg; + + data.generation_prompt = MSG_START + THINK_START + msg.reasoning_content; + if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + data.generation_prompt += THINK_END + RESP_START + msg.render_content(); + } + + data.prompt += data.generation_prompt; + } + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + auto end = p.end(); + + auto start = p.optional(p.literal(MSG_START)); + + // the think section is always consumed, even with reasoning extraction off: + // the generation prompt ends with open_tag('think'), so it is always present. + // reasoning stops at its own closer, or at the response opener if the model + // skips the closer + auto think_body = extract_reasoning ? p.reasoning(p.until_one_of({ THINK_END, RESP_START })) : + p.content(p.until_one_of({ THINK_END, RESP_START })); + + auto reasoning = p.optional(p.optional(p.literal(THINK_START)) + think_body + + p.optional(p.literal(THINK_END))); + + // content runs to the response closer, or to the next section if truncated + auto response = p.optional(p.literal(RESP_START)) + + p.content(p.until_one_of({ RESP_END, TOOLS_START, MSG_END })) + + p.optional(p.literal(RESP_END)); + + // the EOG token after the message closer reaches the parser as text, + // so it must be consumed or the parse stays incomplete + auto trailer = p.optional(p.literal(MSG_END)) + p.optional(p.literal(EOM_TOKEN)); + + if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { + return start + reasoning + response + trailer + end; + } + + auto tool_choices = p.choice(); + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + std::string name = function.at("name"); + const json schema = function.contains("parameters") ? function.at("parameters") : json::object(); + + // arguments come one tag per key, with the JSON type in a type="..." + // attribute. the type is taken from the tool schema instead, as it tells + // us if the value is JSON or a literal string + auto args = p.eps(); + if (schema.contains("properties") && !schema.at("properties").empty()) { + auto arg_choices = p.choice(); + for (const auto & prop : schema.at("properties").items()) { + const std::string & key = prop.key(); + + std::string type = "string"; + if (prop.value().is_object() && prop.value().contains("type") && + prop.value().at("type").is_string()) { + type = prop.value().at("type").get(); + } + + auto value = type == "string" ? p.tool_arg_string_value(p.until(ARG_END)) : + p.tool_arg_value(p.until(ARG_END)); + + // skip the trailing type="..." attribute: anything up to <|sep|> + arg_choices |= p.rule("kimi-k3-arg-" + name + "-" + key, + p.tool_arg(p.tool_arg_open(p.literal(ARG_START)) + + p.tool_arg_name(p.literal(key)) + p.literal("\"") + + p.until(SEP) + p.literal(SEP) + value + + p.tool_arg_close(p.literal(ARG_END)))); + } + args = p.zero_or_more(arg_choices); + } + + // skip the trailing index="N" attribute the same way + auto call = p.tool(p.tool_open(p.literal(CALL_START) + p.tool_name(p.literal(name)) + p.literal("\"") + + p.until(SEP) + p.literal(SEP)) + + p.tool_args(args) + p.tool_close(p.literal(CALL_END))); + + tool_choices |= p.rule("kimi-k3-tool-" + name, call); + }); + + // all calls go inside one tools section, then the message is closed. the + // message closer is part of the trigger rule, or else the lazy grammar + // rejects it once tool calls have started + auto tools_section = + p.trigger_rule("kimi-k3-tool-call", p.literal(TOOLS_START) + p.one_or_more(tool_choices) + + p.literal(TOOLS_END) + p.optional(p.literal(MSG_END)) + + p.optional(p.literal(EOM_TOKEN))); + + auto tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? tools_section : + p.optional(tools_section); + + return start + reasoning + response + tools + trailer + end; + }); + + data.parser = parser.save(); + + if (include_grammar) { + data.grammar_lazy = inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_REQUIRED; + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + if (function.contains("parameters")) { + auto schema = function.at("parameters"); + builder.resolve_refs(schema); + } + }); + parser.build_grammar(builder, data.grammar_lazy); + }); + + data.grammar_triggers = { + { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, TOOLS_START }, + }; + } + + return data; +} + // Cohere2 MoE (a.k.a. "North Code") parser. // // The assistant turn is fully marker-wrapped: @@ -3307,6 +3525,13 @@ std::optional common_chat_try_specialized_template( return common_chat_params_init_kimi_k2(tmpl, params); } + // Kimi K3 - the <|open|>/<|close|>/<|end_of_msg|> markers are unique to it + if (src.find("<|open|>") != std::string::npos && src.find("<|close|>") != std::string::npos && + src.find("<|end_of_msg|>") != std::string::npos) { + LOG_DBG("Using specialized template: Kimi K3\n"); + return common_chat_params_init_kimi_k3(tmpl, params); + } + // Cohere2 MoE / North Code - marker-wrapped format with <|START_TEXT|> content and // <|START_ACTION|> JSON tool calls. <|START_TEXT|> is unique to this template (the older // Command-R templates use <|START_RESPONSE|>). diff --git a/common/jinja/caps.cpp b/common/jinja/caps.cpp index 00b02119a..5e6b34f76 100644 --- a/common/jinja/caps.cpp +++ b/common/jinja/caps.cpp @@ -26,7 +26,7 @@ void caps_apply_preserve_reasoning(jinja::context & ctx, bool enabled) { ctx.set_val("preserve_thinking", mk_val(enabled)); ctx.set_val("clear_thinking", mk_val(!enabled)); ctx.set_val("truncate_history_thinking", mk_val(!enabled)); - ctx.set_val("drop_thinking", mk_val(!enabled)); + ctx.set_val("drop_thinking", mk_val(!enabled)); } void caps_apply_reasoning_effort(jinja::context & ctx, const std::string & effort) { @@ -120,6 +120,8 @@ caps caps_get(jinja::program & prog) { JJ_DEBUG("%s\n", ">>> Running capability check: typed content"); + static const std::string content_marker = "STRING_MARKER"; + // case: typed content support caps_try_execute( prog, @@ -128,22 +130,26 @@ caps caps_get(jinja::program & prog) { return json::array({ { {"role", "user"}, - {"content", "content"} + {"content", content_marker} } }); }, nullptr, // ctx_fn nullptr, // tools_fn - [&](context &, bool success, value & messages, value &, const std::string &) { + [&](context &, bool success, value & messages, value &, const std::string & rendered) { auto & content = messages->at(0)->at("content"); caps_print_stats(content, "messages[0].content"); - if (has_op(content, "selectattr") || has_op(content, "array_access")) { + bool used_as_array = has_op(content, "selectattr") || has_op(content, "array_access"); + if (used_as_array) { // accessed as an array result.supports_typed_content = true; } if (!success) { // failed to execute with content as string result.supports_string_content = false; + } else if (used_as_array && rendered.find(content_marker) == std::string::npos) { + // edge case: string may be accessed for checking, but does not appear in the output + result.supports_string_content = false; } } ); diff --git a/common/speculative.cpp b/common/speculative.cpp index aec94199b..08ac810cd 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -926,6 +926,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout const bool is_dspark; + // dspark speculators + bool sample_from_anchor = true; + const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices uint32_t target_layer_ids_n = 0; @@ -960,16 +963,20 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { if (llama_model_meta_val_str(model_dft, "dflash.block_size", buf, sizeof(buf)) >= 0) { block_size = std::atoi(buf); } + if (llama_model_meta_val_str(model_dft, "dflash.sample_from_anchor", buf, sizeof(buf)) >= 0) { + sample_from_anchor = std::strcmp(buf, "true") == 0; + } } mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft)); LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str()); LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min); - LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u\n", __func__, block_size, mask_token_id, target_layer_ids_n); + LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u, sample_from_anchor=%s\n", __func__, + block_size, mask_token_id, target_layer_ids_n, sample_from_anchor ? "true" : "false"); // DFlash input is [id_last, * (block_size-1)]: in-place denoising yields at most - // block_size-1 draft tokens, DSpark yield a full block_size draft tokens - const int32_t n_draft_max = is_dspark ? block_size : block_size - 1; + // block_size-1 draft tokens, anchor-first DSpark yields a full block_size draft tokens + const int32_t n_draft_max = is_dspark && sample_from_anchor ? block_size : block_size - 1; if (this->params.n_max > n_draft_max || this->params.n_min > n_draft_max) { LOG_WRN("%s: requested draft size (n_max=%d, n_min=%d) exceeds the trained block size %d -- clamping to %d\n", __func__, this->params.n_max, this->params.n_min, block_size, n_draft_max); @@ -1175,7 +1182,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { const int32_t n_draft = params.n_max; - const int32_t n_block_tokens = n_draft + (is_dspark ? 0 : 1); + const int32_t n_block_tokens = n_draft + (is_dspark && sample_from_anchor ? 0 : 1); i_block_beg[seq_id] = batch.n_tokens; n_block [seq_id] = n_block_tokens; for (int32_t i = 0; i < n_block_tokens; ++i) { @@ -1208,11 +1215,11 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { auto & result = *dp.result; if (is_dspark) { - // DSpark predicts the next token from position 0 and optionally truncates - // at the first position below the confidence threshold. + // DSpark: read from the first draft slot, truncate below the confidence threshold const float * conf = params.p_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr; - - for (int32_t i = 0; i < n_block_tokens; ++i) { + // bonus-anchor drafts read the mask positions only, like DFlash + const int32_t i_draft_beg = sample_from_anchor ? 0 : 1; + for (int32_t i = i_draft_beg; i < n_block_tokens; ++i) { const int32_t idx = beg + i; if (conf && conf[(size_t) idx * n_embd_dec] < params.p_min) { diff --git a/conversion/__init__.py b/conversion/__init__.py index f4d475de7..3232a1050 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -27,6 +27,7 @@ TEXT_MODEL_MAP: dict[str, str] = { "BaichuanForCausalLM": "baichuan", "BailingMoeForCausalLM": "bailingmoe", "BailingMoeV2ForCausalLM": "bailingmoe", + "BailingMoeV3ForCausalLM": "bailingmoe3", "BambaForCausalLM": "granite", "BertForMaskedLM": "bert", "BertForSequenceClassification": "bert", @@ -54,6 +55,8 @@ TEXT_MODEL_MAP: dict[str, str] = { "DeepseekV32ForCausalLM": "deepseek", "DFlashDraftModel": "qwen", "Qwen3DSparkModel": "qwen", + "DSparkDraftModel": "qwen", + "DSparkSpeculator": "qwen", "DeepseekV4ForCausalLM": "deepseek", "DeepseekV4DSparkModel": "deepseek", "DistilBertForMaskedLM": "bert", @@ -125,6 +128,7 @@ TEXT_MODEL_MAP: dict[str, str] = { "JinaEmbeddingsV5Model": "bert", "KORMoForCausalLM": "qwen", "KimiK25ForConditionalGeneration": "deepseek", + "KimiK3ForConditionalGeneration": "kimi_k3", "KimiLinearForCausalLM": "kimi_linear", "KimiLinearModel": "kimi_linear", "KimiVLForConditionalGeneration": "deepseek", diff --git a/conversion/afmoe.py b/conversion/afmoe.py index 5e66a51da..844925dca 100644 --- a/conversion/afmoe.py +++ b/conversion/afmoe.py @@ -13,6 +13,7 @@ from .llama import LlamaModel @ModelBase.register("AfmoeForCausalLM") +@ModelBase.example("arcee-ai/Trinity-Large-Thinking") class AfmoeModel(LlamaModel): model_arch = gguf.MODEL_ARCH.AFMOE diff --git a/conversion/arctic.py b/conversion/arctic.py index 775cacaab..843e24a7b 100644 --- a/conversion/arctic.py +++ b/conversion/arctic.py @@ -16,6 +16,7 @@ from .llama import LlamaModel @ModelBase.register("ArcticForCausalLM") +@ModelBase.example("Snowflake/snowflake-arctic-instruct") class ArcticModel(TextModel): model_arch = gguf.MODEL_ARCH.ARCTIC diff --git a/conversion/baichuan.py b/conversion/baichuan.py index 4cf34057c..769bdd567 100644 --- a/conversion/baichuan.py +++ b/conversion/baichuan.py @@ -9,6 +9,7 @@ from .base import ModelBase, TextModel, gguf, logger @ModelBase.register("BaichuanForCausalLM", "BaiChuanForCausalLM") +@ModelBase.example("baichuan-inc/Baichuan2-7B-Chat", "baichuan-inc/Baichuan-7B") class BaichuanModel(TextModel): model_arch = gguf.MODEL_ARCH.BAICHUAN diff --git a/conversion/bailingmoe.py b/conversion/bailingmoe.py index 2c6425cb6..351be1df1 100644 --- a/conversion/bailingmoe.py +++ b/conversion/bailingmoe.py @@ -11,6 +11,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("BailingMoeForCausalLM") +@ModelBase.example("inclusionAI/Ling-lite") class BailingMoeModel(TextModel): model_arch = gguf.MODEL_ARCH.BAILINGMOE @@ -108,6 +109,7 @@ class BailingMoeModel(TextModel): @ModelBase.register("BailingMoeV2ForCausalLM") +@ModelBase.example("inclusionAI/Ling-mini-2.0") class BailingMoeV2Model(TextModel): model_arch = gguf.MODEL_ARCH.BAILINGMOE2 @@ -189,6 +191,7 @@ class BailingMoeV2Model(TextModel): @ModelBase.register("SarvamMoEForCausalLM", "modeling_sarvam_moe.SarvamMoEForCausalLM") +@ModelBase.example("sarvamai/sarvam-30b") class SarvamMoEModel(BailingMoeV2Model): model_arch = gguf.MODEL_ARCH.BAILINGMOE2 # Sarvam-MoE shares the BailingMoeV2 architecture; only differences: diff --git a/conversion/bailingmoe3.py b/conversion/bailingmoe3.py new file mode 100644 index 000000000..20bba23e5 --- /dev/null +++ b/conversion/bailingmoe3.py @@ -0,0 +1,193 @@ +from __future__ import annotations + +import re + +from typing import Callable, Iterable, TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from torch import Tensor + +from .base import ModelBase, TextModel, gguf + + +@ModelBase.register("BailingMoeV3ForCausalLM") +@ModelBase.example("inclusionAI/Ling-3.0-tiny", "inclusionAI/Ling-3.0-flash") +class BailingMoeV3Model(TextModel): + model_arch = gguf.MODEL_ARCH.BAILINGMOE3 + supports_mtp_export = True + + _experts: list[dict[str, Tensor]] | None = None + _main_layers: int | None = None + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + nextn_layers = self.hparams.get("num_nextn_predict_layers", 0) or 0 + if self.no_mtp: + nextn_layers = 0 + self.block_count = self.hparams["num_hidden_layers"] + nextn_layers + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) + + def index_tensors(self, remote_hf_model_id: str | None = None): + type(self)._main_layers = self.hparams["num_hidden_layers"] + return super().index_tensors(remote_hf_model_id=remote_hf_model_id) + + def set_vocab(self): + self._set_vocab_gpt2() + + def is_full_attention(self, bid: int) -> bool: + n_layer = self.hparams["num_hidden_layers"] + layer_group_size = self.hparams["layer_group_size"] + return bid >= n_layer or (bid + 1) % layer_group_size == 0 or bid >= n_layer // layer_group_size * layer_group_size + + def set_gguf_parameters(self): + if not self.hparams.get("no_kda_lora", False): + raise ValueError("BailingMoeV3 KDA LoRA projections are not supported") + if not self.hparams.get("kda_safe_gate", False): + raise ValueError("BailingMoeV3 non-safe KDA gates are not supported") + if self.hparams.get("gated_attention_proj_granularity_type") != "head_wise": + raise ValueError("BailingMoeV3 requires head-wise attention gates") + + self.hparams["num_key_value_heads"] = 1 + super().set_gguf_parameters() + + n_head_kv = [1 if self.is_full_attention(il) else 0 for il in range(self.block_count)] + self.gguf_writer.add_head_count_kv(n_head_kv) + + self.gguf_writer.add_vocab_size(self.hparams["vocab_size"]) + self.gguf_writer.add_ssm_conv_kernel(self.hparams["short_conv_kernel_size"]) + self.gguf_writer.add_kda_head_dim(self.hparams["head_dim"]) + self.gguf_writer.add_kda_safe_gate(self.hparams["kda_safe_gate"]) + self.gguf_writer.add_kda_gate_lower_bound(self.hparams["kda_lower_bound"]) + + kv_lora_rank = self.hparams["kv_lora_rank"] + qk_nope_head_dim = self.hparams["qk_nope_head_dim"] + qk_rope_head_dim = self.hparams["qk_rope_head_dim"] + if (q_lora_rank := self.hparams.get("q_lora_rank")) is not None: + self.gguf_writer.add_q_lora_rank(q_lora_rank) + self.gguf_writer.add_kv_lora_rank(kv_lora_rank) + self.gguf_writer.add_rope_dimension_count(qk_rope_head_dim) + self.gguf_writer.add_key_length(kv_lora_rank + qk_rope_head_dim) + self.gguf_writer.add_key_length_mla(qk_nope_head_dim + qk_rope_head_dim) + self.gguf_writer.add_value_length_mla(self.hparams["v_head_dim"]) + + self.gguf_writer.add_expert_feed_forward_length(self.hparams["moe_intermediate_size"]) + self.gguf_writer.add_expert_shared_feed_forward_length(self.hparams["moe_shared_expert_intermediate_size"]) + self.gguf_writer.add_expert_shared_count(self.hparams["num_shared_experts"]) + self.gguf_writer.add_leading_dense_block_count(self.hparams["first_k_dense_replace"]) + self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"]) + self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"]) + + def clamp_limits(key: str) -> list[float] | None: + values = self.hparams.get(key) + if values is None: + return None + values = [0.0 if value is None else float(value) for value in values[:self.block_count]] + return values + [0.0] * (self.block_count - len(values)) + + if (values := clamp_limits("expert_swiglu_limit_list")) is not None: + self.gguf_writer.add_swiglu_clamp_exp(values) + if (values := clamp_limits("share_expert_swiglu_limit_list")) is not None: + self.gguf_writer.add_swiglu_clamp_shexp(values) + + if not self.no_mtp and (nextn_layers := self.hparams.get("num_nextn_predict_layers", 0)): + self.gguf_writer.add_nextn_predict_layers(nextn_layers) + + def prepare_metadata(self, vocab_only: bool): + from_dir = self.fname_out.is_dir() + super().prepare_metadata(vocab_only=vocab_only) + + if not self.mtp_only or not from_dir: + return + + output_type: str = self.ftype.name.partition("_")[2] + fname_default: str = gguf.naming_convention( + self.metadata.name, self.metadata.basename, self.metadata.finetune, + self.metadata.version, size_label=None, output_type=output_type, model_type=None) + self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf" + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + if name.endswith(".expert_bias"): + name += ".bias" + + if cls._main_layers is None: + return super().filter_tensors((name, gen)) + + m = re.match(r"model\.layers\.(\d+)\.", name) + is_mtp = m is not None and int(m.group(1)) >= cls._main_layers + + if is_mtp and cls.no_mtp: + return None + if cls.mtp_only and not is_mtp and name not in ( + "model.word_embeddings.weight", "model.norm.weight", "lm_head.weight", + ): + return None + + return super().filter_tensors((name, gen)) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name.endswith((".q_conv1d.weight", ".k_conv1d.weight", ".v_conv1d.weight")) and data_torch.ndim in (2, 3): + d_inner = data_torch.shape[0] + d_conv = data_torch.shape[-1] + data_torch = data_torch.reshape(1, d_inner, 1, d_conv) + + if name.endswith(".A_log"): + data_torch = torch.exp(data_torch).reshape(-1, 1) + + if name.endswith(".dt_bias"): + name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias" + + if name.endswith(".attention.f_proj.weight"): + assert bid is not None + if self.is_full_attention(bid): + raise ValueError(f"unexpected f_proj on full-attention layer {bid}") + name = self.format_tensor_name(gguf.MODEL_TENSOR.SSM_F_A, bid) + + if name.endswith(".attention.g_proj.weight"): + assert bid is not None + tensor = gguf.MODEL_TENSOR.ATTN_GATE if self.is_full_attention(bid) else gguf.MODEL_TENSOR.SSM_G_A + name = self.format_tensor_name(tensor, bid) + + if ".mlp.experts." in name: + n_experts = self.hparams["num_experts"] + assert bid is not None + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + + self._experts[bid][name] = data_torch + if len(self._experts[bid]) >= n_experts * 3: + for weight_name in ("down_proj", "gate_proj", "up_proj"): + tensors = [] + for expert_id in range(n_experts): + expert_name = f"model.layers.{bid}.mlp.experts.{expert_id}.{weight_name}.weight" + tensors.append(self._experts[bid].pop(expert_name)) + merged_name = f"model.layers.{bid}.mlp.experts.{weight_name}.weight" + yield from super().modify_tensors(torch.stack(tensors, dim=0), merged_name, bid) + return + + if name.endswith(".attention.kv_b_proj.weight"): + assert bid is not None + n_head = self.hparams["num_attention_heads"] + v_head_dim = self.hparams["v_head_dim"] + qk_nope_head_dim = self.hparams["qk_nope_head_dim"] + assert data_torch.shape[0] == n_head * (v_head_dim + qk_nope_head_dim) + kv_b = data_torch.view(n_head, v_head_dim + qk_nope_head_dim, data_torch.shape[-1]) + k_b, v_b = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1) + name_k = self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K_B, bid) + name_v = self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V_B, bid) + yield from super().modify_tensors(k_b.transpose(1, 2), name_k, bid) + yield from super().modify_tensors(v_b, name_v, bid) + return + + yield from super().modify_tensors(data_torch, name, bid) + + def prepare_tensors(self): + super().prepare_tensors() + if self._experts is not None: + experts = [name for layer in self._experts for name in layer] + if experts: + raise ValueError(f"Unprocessed experts: {experts}") diff --git a/conversion/base.py b/conversion/base.py index 718d53944..56547ace0 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -658,6 +658,43 @@ class ModelBase: def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: return () + @staticmethod + def repack_mxfp4_blocks(packed: Tensor, scale: Tensor) -> np.ndarray: + """ + Repack 4-bit MX weights into ggml `block_mxfp4`. Lossless - only moves bits. + + Source (compressed-tensors "mxfp4-pack-quantized", also used by DeepSeek-V4): + packed uint8 [rows, cols/2] element 2i in the low nibble, 2i+1 in the high one + scale uint8 [rows, cols/32] one E8M0 biased exponent per 32-element group + + Destination, per group: one scale byte then 16 code bytes, where byte j holds + element j in the low nibble and element j+16 in the high one. + + The 4-bit codes need no remapping: both sides index into ggml's kvalues_mxfp4 + order. ggml doubles the kvalues and halves the scale, so the value is the same. + """ + p = packed.contiguous().view(torch.uint8) + s = scale.contiguous().view(torch.uint8) + + rows, packed_cols = p.shape + cols = packed_cols * 2 + if cols % 32 != 0: + raise ValueError(f"MXFP4 source row has {cols} values, expected a multiple of 32") + + n_blocks = cols // 32 + if tuple(s.shape) != (rows, n_blocks): + raise ValueError(f"MXFP4 scale shape {tuple(s.shape)} does not match {(rows, n_blocks)}") + + src = p.reshape(rows, n_blocks, 16) + lo = src & 0x0F # elements 0, 2, 4, ... + hi = (src >> 4) & 0x0F # elements 1, 3, 5, ... + + vals = torch.stack((lo, hi), dim=-1).reshape(rows, n_blocks, 32) + qs = vals[:, :, :16] | (vals[:, :, 16:] << 4) + + raw = torch.cat((s.unsqueeze(-1), qs.to(torch.uint8)), dim=-1) + return raw.reshape(rows, n_blocks * 17).cpu().numpy() + @staticmethod def _nvfp4_pack(weight: Tensor, scale: Tensor) -> tuple[np.ndarray, list[int]]: """Repack NVFP4 ModelOpt tensors into ggml super-block layout. @@ -1112,6 +1149,14 @@ class ModelBase: return modelcls return func + @classmethod + def example(cls, *hf_repos: str) -> Callable[[AnyModel], AnyModel]: + del hf_repos # unused + + def func(modelcls: AnyModel) -> AnyModel: + return modelcls + return func + @classmethod def print_registered_models(cls): for model_type, model_classes in cls._model_classes.items(): @@ -2661,7 +2706,10 @@ def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> st # Step3-VL keeps text config under text_config but uses a custom top-level architecture. # For text conversion we route to a dedicated text-only class. # TODO: refactor this later to avoid adding exception here - if model_type == ModelType.TEXT and arch in ("StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration"): + # Kimi-K3's text_config reports "KimiLinearForCausalLM", which is the older + # Kimi-Linear-48B architecture and cannot load K3 (no attention residuals, + # latent MoE, situ, ...). Route on the top-level architecture instead. + if model_type == ModelType.TEXT and arch in ("StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration"): return arch # if "architectures" is found in the sub-config, use that instead diff --git a/conversion/bert.py b/conversion/bert.py index 0d25d0d62..8ea6c42dc 100644 --- a/conversion/bert.py +++ b/conversion/bert.py @@ -15,6 +15,7 @@ from .base import ModelBase, SentencePieceTokenTypes, TextModel, gguf, logger @ModelBase.register("BertModel", "BertForMaskedLM", "CamembertModel", "BertForSequenceClassification") +@ModelBase.example("BAAI/bge-small-en-v1.5", "dangvantuan/sentence-camembert-base") class BertModel(TextModel): model_arch = gguf.MODEL_ARCH.BERT @@ -240,6 +241,7 @@ class BertModel(TextModel): @ModelBase.register("DistilBertModel", "DistilBertForMaskedLM", "DistilBertForSequenceClassification") +@ModelBase.example("distilbert/distilbert-base-uncased") class DistilBertModel(BertModel): model_arch = gguf.MODEL_ARCH.BERT @@ -263,6 +265,7 @@ class DistilBertModel(BertModel): @ModelBase.register("RobertaModel", "RobertaForSequenceClassification") +@ModelBase.example("sentence-transformers/stsb-roberta-base") class RobertaModel(BertModel): model_arch = gguf.MODEL_ARCH.BERT @@ -312,6 +315,7 @@ class RobertaModel(BertModel): @ModelBase.register("NomicBertModel") +@ModelBase.example("nomic-ai/nomic-embed-text-v1.5") class NomicBertModel(BertModel): model_arch = gguf.MODEL_ARCH.BERT @@ -400,6 +404,7 @@ class NomicBertModel(BertModel): @ModelBase.register("NeoBERT", "NeoBERTLMHead", "NeoBERTForSequenceClassification") +@ModelBase.example("chandar-lab/NeoBERT") class NeoBert(BertModel): model_arch = gguf.MODEL_ARCH.NEO_BERT @@ -431,6 +436,7 @@ class NeoBert(BertModel): @ModelBase.register("EuroBertModel", "JinaEmbeddingsV5Model") +@ModelBase.example("hf-tiny-v2/tiny-random-EuroBertModel", "jinaai/jina-embeddings-v5-text-nano") class EuroBertModel(TextModel): model_arch = gguf.MODEL_ARCH.EUROBERT @@ -459,6 +465,7 @@ class EuroBertModel(TextModel): @ModelBase.register("XLMRobertaModel", "XLMRobertaForSequenceClassification") +@ModelBase.example("BAAI/bge-m3") class XLMRobertaModel(BertModel): model_arch = gguf.MODEL_ARCH.BERT _lora_files = {} @@ -561,6 +568,7 @@ class XLMRobertaModel(BertModel): @ModelBase.register("JinaBertModel", "JinaBertForMaskedLM") +@ModelBase.example("jinaai/jina-embeddings-v2-base-en") class JinaBertV2Model(BertModel): model_arch = gguf.MODEL_ARCH.JINA_BERT_V2 @@ -588,6 +596,7 @@ class JinaBertV2Model(BertModel): @ModelBase.register("ModernBertModel", "ModernBertForMaskedLM", "ModernBertForSequenceClassification") +@ModelBase.example("answerdotai/ModernBERT-base") class ModernBertModel(BertModel): model_arch = gguf.MODEL_ARCH.MODERN_BERT diff --git a/conversion/bitnet.py b/conversion/bitnet.py index 0c2baee87..82bcadaf9 100644 --- a/conversion/bitnet.py +++ b/conversion/bitnet.py @@ -9,6 +9,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("BitnetForCausalLM", "BitNetForCausalLM") +@ModelBase.example("microsoft/bitnet-b1.58-2B-4T") class BitnetModel(TextModel): model_arch = gguf.MODEL_ARCH.BITNET diff --git a/conversion/bloom.py b/conversion/bloom.py index d98edf6d5..9654cd4a0 100644 --- a/conversion/bloom.py +++ b/conversion/bloom.py @@ -13,6 +13,7 @@ from .base import ModelBase, TextModel, gguf, logger @ModelBase.register("BloomForCausalLM", "BloomModel") +@ModelBase.example("bigscience/bloom-560m") class BloomModel(TextModel): model_arch = gguf.MODEL_ARCH.BLOOM diff --git a/conversion/chameleon.py b/conversion/chameleon.py index a996bfa53..8f2065df6 100644 --- a/conversion/chameleon.py +++ b/conversion/chameleon.py @@ -12,6 +12,8 @@ from .llama import LlamaModel @ModelBase.register("ChameleonForConditionalGeneration") @ModelBase.register("ChameleonForCausalLM") # obsolete +# [TAG_HF_EXAMPLE_GATED] facebook/chameleon-7b is gated +# [TAG_HF_EXAMPLE_MISSING] class ChameleonModel(TextModel): model_arch = gguf.MODEL_ARCH.CHAMELEON diff --git a/conversion/chatglm.py b/conversion/chatglm.py index d63855038..9b902dae3 100644 --- a/conversion/chatglm.py +++ b/conversion/chatglm.py @@ -9,6 +9,7 @@ from .base import ModelBase, SentencePieceTokenTypes, TextModel, gguf @ModelBase.register("GlmForCausalLM", "ChatGLMModel", "ChatGLMForConditionalGeneration") +@ModelBase.example("THUDM/chatglm3-6b", "zai-org/glm-4-9b-chat-hf") class ChatGLMModel(TextModel): model_arch = gguf.MODEL_ARCH.CHATGLM diff --git a/conversion/codeshell.py b/conversion/codeshell.py index 8bfc3178d..1c7f1129b 100644 --- a/conversion/codeshell.py +++ b/conversion/codeshell.py @@ -4,6 +4,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("CodeShellForCausalLM") +@ModelBase.example("WisdomShell/CodeShell-7B") class CodeShellModel(TextModel): model_arch = gguf.MODEL_ARCH.CODESHELL diff --git a/conversion/cogvlm.py b/conversion/cogvlm.py index d92df55d4..13c314441 100644 --- a/conversion/cogvlm.py +++ b/conversion/cogvlm.py @@ -11,6 +11,7 @@ from .llama import LlamaModel @ModelBase.register("CogVLMForCausalLM") +@ModelBase.example("THUDM/cogvlm2-llama3-chat-19B", "THUDM/cogvlm-chat-hf") class CogVLMVisionModel(MmprojModel): def set_gguf_parameters(self): @@ -29,5 +30,6 @@ class CogVLMVisionModel(MmprojModel): @ModelBase.register("CogVLMForCausalLM") +@ModelBase.example("THUDM/cogvlm2-llama3-chat-19B", "THUDM/cogvlm-chat-hf") class CogVLMModel(LlamaModel): model_arch = gguf.MODEL_ARCH.COGVLM diff --git a/conversion/command_r.py b/conversion/command_r.py index 118565c66..971f93ebd 100644 --- a/conversion/command_r.py +++ b/conversion/command_r.py @@ -12,6 +12,8 @@ from .base import ModelBase, TextModel, gguf, logger @ModelBase.register("CohereForCausalLM") +# [TAG_HF_EXAMPLE_GATED] CohereLabs/c4ai-command-r-v01 is gated +# [TAG_HF_EXAMPLE_MISSING] class CommandR2Model(TextModel): model_arch = gguf.MODEL_ARCH.COMMAND_R @@ -30,6 +32,8 @@ class CommandR2Model(TextModel): @ModelBase.register("Cohere2ForCausalLM") +# [TAG_HF_EXAMPLE_GATED] CohereLabs/c4ai-command-r7b-12-2024 is gated +@ModelBase.example("hf-tiny-v2/tiny-random-Cohere2ForCausalLM") class Cohere2Model(TextModel): model_arch = gguf.MODEL_ARCH.COHERE2 @@ -59,6 +63,7 @@ class Cohere2Model(TextModel): @ModelBase.register("Cohere2MoeForCausalLM") +@ModelBase.example("CohereLabs/North-Mini-Code-1.0") class Cohere2MoeModel(TextModel): model_arch = gguf.MODEL_ARCH.COHERE2MOE _n_main_layers: int | None = None diff --git a/conversion/dbrx.py b/conversion/dbrx.py index 207ebcb89..d37ce83e7 100644 --- a/conversion/dbrx.py +++ b/conversion/dbrx.py @@ -9,6 +9,7 @@ from .base import ModelBase, TextModel, gguf, logger @ModelBase.register("DbrxForCausalLM") +@ModelBase.example("alpindale/dbrx-instruct") class DbrxModel(TextModel): model_arch = gguf.MODEL_ARCH.DBRX diff --git a/conversion/deci.py b/conversion/deci.py index be446eefa..2ccaa92a9 100644 --- a/conversion/deci.py +++ b/conversion/deci.py @@ -13,6 +13,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("DeciLMForCausalLM") +@ModelBase.example("nvidia/Llama-3_1-Nemotron-51B-Instruct", "Deci/DeciLM-7B") class DeciModel(TextModel): model_arch = gguf.MODEL_ARCH.DECI diff --git a/conversion/deepseek.py b/conversion/deepseek.py index 1846ca401..225f8645d 100644 --- a/conversion/deepseek.py +++ b/conversion/deepseek.py @@ -18,6 +18,7 @@ from .qwen import QwenModel @ModelBase.register("DeepseekOCRForCausalLM") +@ModelBase.example("deepseek-ai/DeepSeek-OCR") class DeepseekOCRVisionModel(MmprojModel): # HF dynamic_preprocess() max_num, which differs per model preproc_max_tiles = 9 @@ -100,11 +101,13 @@ class DeepseekOCRVisionModel(MmprojModel): @ModelBase.register("UnlimitedOCRForCausalLM") +@ModelBase.example("baidu/Unlimited-OCR") class UnlimitedOCRVisionModel(DeepseekOCRVisionModel): preproc_max_tiles = 32 @ModelBase.register("DeepseekOCR2ForCausalLM") +@ModelBase.example("deepseek-ai/DeepSeek-OCR-2") class DeepseekOCR2VisionModel(DeepseekOCRVisionModel): preproc_max_tiles = 6 @@ -134,6 +137,7 @@ class DeepseekOCR2VisionModel(DeepseekOCRVisionModel): @ModelBase.register("DeepseekForCausalLM") +@ModelBase.example("deepseek-ai/deepseek-moe-16b-chat") class DeepseekModel(TextModel): model_arch = gguf.MODEL_ARCH.DEEPSEEK @@ -228,6 +232,7 @@ class DeepseekModel(TextModel): "YoutuForCausalLM", "YoutuVLForConditionalGeneration", ) +@ModelBase.example("deepseek-ai/DeepSeek-V2-Lite", "deepseek-ai/DeepSeek-V3") class DeepseekV2Model(TextModel): model_arch = gguf.MODEL_ARCH.DEEPSEEK2 @@ -457,6 +462,7 @@ class DeepseekV2Model(TextModel): @ModelBase.register("DeepseekV32ForCausalLM") +@ModelBase.example("deepseek-ai/DeepSeek-V3.2-Exp") class DeepseekV32Model(DeepseekV2Model): model_arch = gguf.MODEL_ARCH.DEEPSEEK32 skip_mtp = False @@ -517,6 +523,7 @@ class DeepseekV32Model(DeepseekV2Model): @ModelBase.register("DeepseekV4ForCausalLM") +@ModelBase.example("deepseek-ai/DeepSeek-V4-Flash-Base") class DeepseekV4Model(TextModel): model_arch = gguf.MODEL_ARCH.DEEPSEEK4 supports_mtp_export = True @@ -709,31 +716,6 @@ class DeepseekV4Model(TextModel): for name in tensors_to_remove: del self.model_tensors[name] - @staticmethod - def _pack_mxfp4_blocks(weight: Tensor, scale: Tensor) -> np.ndarray: - packed = weight.contiguous().view(torch.uint8) - scale_u8 = scale.contiguous().view(torch.uint8) - - out_features, packed_cols = packed.shape - logical_cols = packed_cols * 2 - if logical_cols % 32 != 0: - raise ValueError(f"MXFP4 source row has {logical_cols} values, expected a multiple of 32") - - n_blocks = logical_cols // 32 - if tuple(scale_u8.shape) != (out_features, n_blocks): - raise ValueError(f"MXFP4 scale shape {tuple(scale_u8.shape)} does not match {(out_features, n_blocks)}") - - src = packed.reshape(out_features, n_blocks, 16) - low = src & 0x0F - high = (src >> 4) & 0x0F - - # The safetensors bytes store adjacent values as low/high nibbles. - # ggml MXFP4 blocks store values 0..15 in low nibbles and 16..31 in high nibbles. - vals = torch.stack((low, high), dim=-1).reshape(out_features, n_blocks, 32) - qs = vals[:, :, :16] | (vals[:, :, 16:] << 4) - raw = torch.cat((scale_u8.unsqueeze(-1), qs.to(torch.uint8)), dim=-1) - return raw.reshape(out_features, n_blocks * 17).cpu().numpy() - def _write_mxfp4_expert_tensor(self, bid: int, proj: str, tensor_key: gguf.MODEL_TENSOR) -> list[str]: n_experts = self.hparams["n_routed_experts"] data: np.ndarray | None = None @@ -747,7 +729,7 @@ class DeepseekV4Model(TextModel): weight = LazyTorchTensor.to_eager(self.model_tensors[weight_name]()) scale = LazyTorchTensor.to_eager(self.model_tensors[scale_name]()) - packed = self._pack_mxfp4_blocks(weight, scale) + packed = self.repack_mxfp4_blocks(weight, scale) if data is None: data = np.empty((n_experts, *packed.shape), dtype=packed.dtype) data[eid] = packed @@ -936,6 +918,7 @@ class DeepseekV4Model(TextModel): @ModelBase.register("DeepseekV4DSparkModel") +@ModelBase.example("deepseek-ai/DeepSeek-V4-Flash-DSpark") class DeepseekV4DSparkModel(DeepseekV4Model): model_arch = gguf.MODEL_ARCH.DFLASH diff --git a/conversion/dots1.py b/conversion/dots1.py index 7ac299a6e..ffa3b6db4 100644 --- a/conversion/dots1.py +++ b/conversion/dots1.py @@ -11,6 +11,7 @@ from .qwen import Qwen2MoeModel @ModelBase.register("Dots1ForCausalLM") +@ModelBase.example("rednote-hilab/dots.llm1.inst") class Dots1Model(Qwen2MoeModel): model_arch = gguf.MODEL_ARCH.DOTS1 diff --git a/conversion/dotsocr.py b/conversion/dotsocr.py index f87f62abd..ace6aa9a1 100644 --- a/conversion/dotsocr.py +++ b/conversion/dotsocr.py @@ -9,6 +9,7 @@ from .base import MmprojModel, ModelBase, gguf @ModelBase.register("DotsOCRForCausalLM") +@ModelBase.example("rednote-hilab/dots.ocr") class DotsOCRVisionModel(MmprojModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) diff --git a/conversion/dream.py b/conversion/dream.py index 459e8d46a..14f25404d 100644 --- a/conversion/dream.py +++ b/conversion/dream.py @@ -9,6 +9,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("DreamModel") +@ModelBase.example("Dream-org/Dream-v0-Instruct-7B") class DreamModel(TextModel): model_arch = gguf.MODEL_ARCH.DREAM diff --git a/conversion/ernie.py b/conversion/ernie.py index aa8a3bc8e..3c4226a25 100644 --- a/conversion/ernie.py +++ b/conversion/ernie.py @@ -15,6 +15,7 @@ from .base import MmprojModel, ModelBase, TextModel, gguf @ModelBase.register("Ernie4_5_ForCausalLM", "Ernie4_5ForCausalLM") +@ModelBase.example("baidu/ERNIE-4.5-0.3B-PT") class Ernie4_5Model(TextModel): model_arch = gguf.MODEL_ARCH.ERNIE4_5 @@ -73,6 +74,7 @@ class Ernie4_5Model(TextModel): @ModelBase.register("Ernie4_5_MoeForCausalLM") +@ModelBase.example("baidu/ERNIE-4.5-21B-A3B-PT") class Ernie4_5MoeModel(Ernie4_5Model): model_arch = gguf.MODEL_ARCH.ERNIE4_5_MOE _experts: list[dict[str, Tensor]] | None = None @@ -156,11 +158,13 @@ class Ernie4_5MoeModel(Ernie4_5Model): @ModelBase.register("PaddleOCRVLForConditionalGeneration") +@ModelBase.example("PaddlePaddle/PaddleOCR-VL") class PaddleOCRModel(Ernie4_5Model): model_arch = gguf.MODEL_ARCH.PADDLEOCR @ModelBase.register("PaddleOCRVisionModel") +@ModelBase.example("PaddlePaddle/PaddleOCR-VL") class PaddleOCRVisionModel(MmprojModel): # PaddleOCR-VL uses a modified version of Siglip min_pixels: int = 0 diff --git a/conversion/exaone.py b/conversion/exaone.py index 1cd2244db..0919d2ffa 100644 --- a/conversion/exaone.py +++ b/conversion/exaone.py @@ -15,6 +15,7 @@ from .qwenvl import Qwen2VLVisionModel @ModelBase.register("ExaoneForCausalLM") +@ModelBase.example("LGAI-EXAONE/EXAONE-3.5-2.4B-Instruct") class ExaoneModel(TextModel): model_arch = gguf.MODEL_ARCH.EXAONE @@ -60,6 +61,7 @@ class ExaoneModel(TextModel): @ModelBase.register("Exaone4ForCausalLM") +@ModelBase.example("LGAI-EXAONE/EXAONE-4.0-32B") class Exaone4Model(TextModel): model_arch = gguf.MODEL_ARCH.EXAONE4 @@ -126,6 +128,7 @@ class Exaone4Model(TextModel): # note: transformers >= 5.1 renamed the class to "ExaoneMoeForCausalLM" (lowercase 'e'), # so accept both spellings - LG AI have updated the configs of already-released models @ModelBase.register("ExaoneMoEForCausalLM", "ExaoneMoeForCausalLM") +@ModelBase.example("LGAI-EXAONE/K-EXAONE-236B-A23B") class ExaoneMoEModel(Exaone4Model): model_arch = gguf.MODEL_ARCH.EXAONE_MOE @@ -214,6 +217,7 @@ class ExaoneMoEModel(Exaone4Model): @ModelBase.register("Exaone4_5_ForConditionalGeneration") +@ModelBase.example("LGAI-EXAONE/EXAONE-4.5-33B") class Exaone4_5_TextModel(Exaone4Model): """Text tower of EXAONE 4.5; Tensors match EXAONE4""" @@ -267,6 +271,7 @@ class Exaone4_5_TextModel(Exaone4Model): @ModelBase.register("Exaone4_5_ForConditionalGeneration") +@ModelBase.example("LGAI-EXAONE/EXAONE-4.5-33B") class Exaone4_5VisionModel(Qwen2VLVisionModel): """Vision tower for EXAONE 4.5; Qwen2-VL-style ViT (GQA) + patch merger""" diff --git a/conversion/falcon.py b/conversion/falcon.py index 085fd4cd3..2c55511a0 100644 --- a/conversion/falcon.py +++ b/conversion/falcon.py @@ -11,6 +11,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("FalconForCausalLM", "RWForCausalLM") +@ModelBase.example("tiiuae/falcon-7b") class FalconModel(TextModel): model_arch = gguf.MODEL_ARCH.FALCON diff --git a/conversion/falcon_h1.py b/conversion/falcon_h1.py index a8bc880b2..6686f7001 100644 --- a/conversion/falcon_h1.py +++ b/conversion/falcon_h1.py @@ -12,6 +12,7 @@ from .mamba import Mamba2Model @ModelBase.register("FalconH1ForCausalLM") +@ModelBase.example("tiiuae/Falcon-H1-0.5B-Base") class FalconH1Model(Mamba2Model): model_arch = gguf.MODEL_ARCH.FALCON_H1 diff --git a/conversion/gemma.py b/conversion/gemma.py index f15a10a38..6b4d7d171 100644 --- a/conversion/gemma.py +++ b/conversion/gemma.py @@ -14,6 +14,8 @@ from .base import MmprojModel, ModelBase, TextModel, gguf, logger @ModelBase.register("GemmaForCausalLM") +# [TAG_HF_EXAMPLE_GATED] google/gemma-2b is gated +@ModelBase.example("trl-internal-testing/tiny-GemmaForCausalLM") class GemmaModel(TextModel): model_arch = gguf.MODEL_ARCH.GEMMA @@ -68,6 +70,8 @@ class GemmaModel(TextModel): @ModelBase.register("Gemma2ForCausalLM") +# [TAG_HF_EXAMPLE_GATED] google/gemma-2-9b-it is gated +@ModelBase.example("trl-internal-testing/tiny-Gemma2ForCausalLM") class Gemma2Model(TextModel): model_arch = gguf.MODEL_ARCH.GEMMA2 @@ -118,6 +122,8 @@ class Gemma2Model(TextModel): @ModelBase.register("Gemma3ForCausalLM", "Gemma3ForConditionalGeneration") +# [TAG_HF_EXAMPLE_GATED] google/gemma-3-4b-it is gated +@ModelBase.example("trl-internal-testing/tiny-Gemma3ForConditionalGeneration", "hf-tiny-v2/tiny-random-Gemma3ForCausalLM") class Gemma3Model(TextModel): model_arch = gguf.MODEL_ARCH.GEMMA3 @@ -174,6 +180,8 @@ class Gemma3Model(TextModel): @ModelBase.register("Gemma3TextModel") +# [TAG_HF_EXAMPLE_GATED] google/embeddinggemma-300m is gated +@ModelBase.example("hf-tiny-v2/tiny-random-Gemma3TextModel") class EmbeddingGemma(Gemma3Model): model_arch = gguf.MODEL_ARCH.GEMMA_EMBEDDING module_paths = [] @@ -248,6 +256,8 @@ class EmbeddingGemma(Gemma3Model): @ModelBase.register("Gemma3ForConditionalGeneration") +# [TAG_HF_EXAMPLE_GATED] google/gemma-3-4b-it is gated +@ModelBase.example("trl-internal-testing/tiny-Gemma3ForConditionalGeneration") class Gemma3VisionModel(MmprojModel): def set_gguf_parameters(self): super().set_gguf_parameters() @@ -352,6 +362,8 @@ class ConformerAudioModel(MmprojModel): @ModelBase.register("Gemma3nForConditionalGeneration") +# [TAG_HF_EXAMPLE_GATED] google/gemma-3n-E2B-it is gated +@ModelBase.example("hf-tiny-v2/tiny-random-Gemma3nForConditionalGeneration") class Gemma3nVisionAudioModel(ConformerAudioModel): has_audio_encoder = True has_vision_encoder = True @@ -471,6 +483,8 @@ class Gemma3nVisionAudioModel(ConformerAudioModel): @ModelBase.register("Gemma3nForCausalLM", "Gemma3nForConditionalGeneration") +# [TAG_HF_EXAMPLE_GATED] google/gemma-3n-E2B-it is gated +@ModelBase.example("hf-tiny-v2/tiny-random-Gemma3nForConditionalGeneration") class Gemma3NModel(Gemma3Model): model_arch = gguf.MODEL_ARCH.GEMMA3N @@ -615,6 +629,7 @@ class Gemma3NModel(Gemma3Model): @ModelBase.register("Gemma4ForConditionalGeneration", "Gemma4ForCausalLM") +@ModelBase.example("google/gemma-4-31B-it", "google/gemma-4-26B-A4B-it", "google/gemma-4-E2B-it") class Gemma4Model(Gemma3Model): model_arch = gguf.MODEL_ARCH.GEMMA4 @@ -795,6 +810,7 @@ class Gemma4Model(Gemma3Model): @ModelBase.register("Gemma4UnifiedForConditionalGeneration") +@ModelBase.example("hf-tiny-v2/tiny-random-Gemma4UnifiedForConditionalGeneration") class Gemma4UnifiedModel(Gemma4Model): model_arch = gguf.MODEL_ARCH.GEMMA4 @@ -815,6 +831,7 @@ class Gemma4UnifiedModel(Gemma4Model): @ModelBase.register("Gemma4AssistantForCausalLM", "Gemma4UnifiedAssistantForCausalLM") +@ModelBase.example("google/gemma-4-31B-it-assistant", "google/gemma-4-26B-A4B-it-assistant", "google/gemma-4-E2B-it-assistant") class Gemma4AssistantModel(Gemma4Model): model_arch = gguf.MODEL_ARCH.GEMMA4_ASSISTANT @@ -835,6 +852,7 @@ class Gemma4AssistantModel(Gemma4Model): @ModelBase.register("Gemma4ForConditionalGeneration") +@ModelBase.example("google/gemma-4-31B-it", "google/gemma-4-26B-A4B-it", "google/gemma-4-E2B-it") class Gemma4VisionAudioModel(MmprojModel): has_audio_encoder = True has_vision_encoder = True @@ -913,6 +931,7 @@ class Gemma4VisionAudioModel(MmprojModel): @ModelBase.register("Gemma4UnifiedForConditionalGeneration") +@ModelBase.example("hf-tiny-v2/tiny-random-Gemma4UnifiedForConditionalGeneration") class Gemma4UnifiedVisionAudioModel(Gemma4VisionAudioModel): has_audio_encoder = True has_vision_encoder = True diff --git a/conversion/glm.py b/conversion/glm.py index e28f54574..abd7f279f 100644 --- a/conversion/glm.py +++ b/conversion/glm.py @@ -15,6 +15,7 @@ from .deepseek import DeepseekV2Model @ModelBase.register("Glm4ForCausalLM", "Glm4vForConditionalGeneration") +@ModelBase.example("zai-org/GLM-4-9B-0414") class Glm4Model(TextModel): model_arch = gguf.MODEL_ARCH.GLM4 use_mrope = False @@ -86,6 +87,7 @@ class Glm4Model(TextModel): @ModelBase.register("GlmOcrForConditionalGeneration") +@ModelBase.example("zai-org/GLM-OCR") class GlmOCRModel(Glm4Model): model_arch = gguf.MODEL_ARCH.GLM4 use_mrope = False @@ -107,6 +109,7 @@ class GlmOCRModel(Glm4Model): @ModelBase.register("Glm4MoeForCausalLM", "Glm4vMoeForConditionalGeneration") +@ModelBase.example("zai-org/GLM-4.5-Air") class Glm4MoeModel(TextModel): model_arch = gguf.MODEL_ARCH.GLM4_MOE @@ -204,6 +207,7 @@ class Glm4MoeModel(TextModel): @ModelBase.register("Glm4MoeLiteForCausalLM") +@ModelBase.example("zai-org/GLM-4.7-Flash") class Glm4MoeLiteModel(DeepseekV2Model): model_arch = gguf.MODEL_ARCH.DEEPSEEK2 skip_mtp = False @@ -272,6 +276,7 @@ class Glm4MoeLiteModel(DeepseekV2Model): @ModelBase.register("GlmMoeDsaForCausalLM") +@ModelBase.example("zai-org/GLM-5.2") class GlmMoeDsaModel(DeepseekV2Model): model_arch = gguf.MODEL_ARCH.GLM_DSA skip_mtp = False @@ -340,6 +345,7 @@ class GlmMoeDsaModel(DeepseekV2Model): @ModelBase.register("SolarOpenForCausalLM") +@ModelBase.example("upstage/Solar-Open-100B") class SolarOpenModel(Glm4MoeModel): model_arch = gguf.MODEL_ARCH.GLM4_MOE diff --git a/conversion/gpt2.py b/conversion/gpt2.py index 1cf06ae8b..06dff9e4c 100644 --- a/conversion/gpt2.py +++ b/conversion/gpt2.py @@ -11,6 +11,7 @@ from .base import ModelBase, TextModel, gguf, logger @ModelBase.register("GPT2LMHeadModel") +@ModelBase.example("openai-community/gpt2") class GPT2Model(TextModel): model_arch = gguf.MODEL_ARCH.GPT2 @@ -38,6 +39,7 @@ class GPT2Model(TextModel): @ModelBase.register("RuGPT3XLForCausalLM") +@ModelBase.example("evilfreelancer/ruGPT3XL") class RuGPT3XLModel(TextModel): model_arch = gguf.MODEL_ARCH.GPT2 diff --git a/conversion/gpt_oss.py b/conversion/gpt_oss.py index d2c70c0bb..7542ec0ea 100644 --- a/conversion/gpt_oss.py +++ b/conversion/gpt_oss.py @@ -11,6 +11,7 @@ from .base import ModelBase, TextModel, gguf, logger @ModelBase.register("GptOssForCausalLM") +@ModelBase.example("openai/gpt-oss-20b") class GptOssModel(TextModel): model_arch = gguf.MODEL_ARCH.GPT_OSS diff --git a/conversion/gptneox.py b/conversion/gptneox.py index 6a42b12b1..0b0e91c4f 100644 --- a/conversion/gptneox.py +++ b/conversion/gptneox.py @@ -13,6 +13,7 @@ from .base import ModelBase, TextModel, gguf, logger @ModelBase.register("GPTNeoXForCausalLM") +@ModelBase.example("EleutherAI/pythia-70m") class GPTNeoXModel(TextModel): model_arch = gguf.MODEL_ARCH.GPTNEOX diff --git a/conversion/granite.py b/conversion/granite.py index 956342e6d..5f1e3e847 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -15,6 +15,7 @@ from .mamba import Mamba2Model @ModelBase.register("GraniteForCausalLM") +@ModelBase.example("ibm-granite/granite-3.3-2b-instruct") class GraniteModel(LlamaModel): """Conversion for IBM's GraniteForCausalLM""" model_arch = gguf.MODEL_ARCH.GRANITE @@ -74,6 +75,7 @@ class GraniteModel(LlamaModel): @ModelBase.register("GraniteMoeForCausalLM", "GraniteMoeSharedForCausalLM") +@ModelBase.example("ibm-granite/granite-3.1-3b-a800m-instruct") class GraniteMoeModel(GraniteModel): """Conversion for IBM's GraniteMoeForCausalLM""" model_arch = gguf.MODEL_ARCH.GRANITE_MOE @@ -124,6 +126,7 @@ class GraniteMoeModel(GraniteModel): @ModelBase.register("GraniteSwitchForCausalLM") +@ModelBase.example("ibm-granite/granite-switch-4.1-3b-preview") class GraniteSwitchModel(GraniteMoeModel): """Dense, all-attention Granite with N per-token embedded LoRA adapters, stacked over the adapter dim with a zero adapter at slot 0 (N = num_adapters + 1).""" @@ -284,6 +287,7 @@ class GraniteSwitchModel(GraniteMoeModel): @ModelBase.register("GraniteMoeHybridForCausalLM", "BambaForCausalLM") +@ModelBase.example("ibm-granite/granite-4.0-h-tiny", "ibm-ai-platform/Bamba-9B-v2") class GraniteHybridModel(Mamba2Model, GraniteMoeModel): """GraniteHybrid is a hybrid SSM + Attention model that uses Mamba2 SSM layers and optionally uses MoE w/ a shared expert""" @@ -426,6 +430,7 @@ class GraniteHybridModel(Mamba2Model, GraniteMoeModel): @ModelBase.register("GraniteSpeechForConditionalGeneration") +@ModelBase.example("ibm-granite/granite-speech-3.3-2b", "ibm-granite/granite-4.0-1b-speech") class GraniteSpeechMmprojModel(MmprojModel): has_vision_encoder = False has_audio_encoder = True @@ -509,6 +514,7 @@ class GraniteSpeechMmprojModel(MmprojModel): @ModelBase.register("GraniteSpeechPlusForConditionalGeneration") +@ModelBase.example("ibm-granite/granite-speech-4.1-2b-plus") class GraniteSpeechPlusMmprojModel(GraniteSpeechMmprojModel): """Conversion for GraniteSpeechPlus - extends GraniteSpeech with feature layer concatenation""" has_vision_encoder = False @@ -537,6 +543,7 @@ class GraniteSpeechPlusMmprojModel(GraniteSpeechMmprojModel): @ModelBase.register("Granite4VisionForConditionalGeneration") +@ModelBase.example("ibm-granite/granite-4.0-3b-vision") class Granite4VisionMmprojModel(MmprojModel): has_vision_encoder = True has_audio_encoder = False diff --git a/conversion/grok.py b/conversion/grok.py index 9098e514a..b966361d2 100644 --- a/conversion/grok.py +++ b/conversion/grok.py @@ -13,6 +13,7 @@ from .base import ModelBase, TextModel, gguf, logger @ModelBase.register("GrokForCausalLM", "Grok1ForCausalLM") +@ModelBase.example("keyfan/grok-1-hf") class GrokModel(TextModel): model_arch = gguf.MODEL_ARCH.GROK diff --git a/conversion/grovemoe.py b/conversion/grovemoe.py index a8be931cb..f418f18ac 100644 --- a/conversion/grovemoe.py +++ b/conversion/grovemoe.py @@ -11,6 +11,7 @@ from .base import ModelBase, TextModel, gguf, logger @ModelBase.register("GroveMoeForCausalLM", "modeling_grove_moe.GroveMoeForCausalLM") +@ModelBase.example("inclusionAI/GroveMoE-Inst") class GroveMoeModel(TextModel): model_arch = gguf.MODEL_ARCH.GROVEMOE diff --git a/conversion/hunyuan.py b/conversion/hunyuan.py index f5ac8a4fb..ee1a10654 100644 --- a/conversion/hunyuan.py +++ b/conversion/hunyuan.py @@ -17,6 +17,7 @@ from .qwen import QwenModel @ModelBase.register("HunYuanMoEV1ForCausalLM") +@ModelBase.example("tencent/Hunyuan-A13B-Instruct") class HunYuanMoEModel(TextModel): model_arch = gguf.MODEL_ARCH.HUNYUAN_MOE @@ -154,6 +155,7 @@ class HunYuanMoEModel(TextModel): @ModelBase.register("HunYuanDenseV1ForCausalLM") +@ModelBase.example("tencent/Hunyuan-4B-Instruct") class HunYuanModel(TextModel): model_arch = gguf.MODEL_ARCH.HUNYUAN_DENSE @@ -290,6 +292,7 @@ class HunYuanModel(TextModel): @ModelBase.register("HunYuanVLForConditionalGeneration") +@ModelBase.example("tencent/HunyuanOCR") class HunyuanVLVisionModel(MmprojModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -333,6 +336,7 @@ class HunyuanVLVisionModel(MmprojModel): @ModelBase.register("HunYuanVLForConditionalGeneration") +@ModelBase.example("tencent/HunyuanOCR") class HunyuanVLTextModel(HunYuanModel): model_arch = gguf.MODEL_ARCH.HUNYUAN_VL @@ -365,6 +369,7 @@ class HunyuanVLTextModel(HunYuanModel): @ModelBase.register("HYV3ForCausalLM") +@ModelBase.example("tencent/Hy3") class HYV3Model(TextModel): model_arch = gguf.MODEL_ARCH.HY_V3 supports_mtp_export = True diff --git a/conversion/internlm.py b/conversion/internlm.py index 7e11aca3c..df2668474 100644 --- a/conversion/internlm.py +++ b/conversion/internlm.py @@ -14,6 +14,7 @@ from .llama import LlamaModel @ModelBase.register("InternLM2ForCausalLM") +@ModelBase.example("internlm/internlm2-chat-7b") class InternLM2Model(TextModel): model_arch = gguf.MODEL_ARCH.INTERNLM2 @@ -170,6 +171,7 @@ class InternLM2Model(TextModel): @ModelBase.register("InternLM3ForCausalLM") +@ModelBase.example("internlm/internlm3-8b-instruct") class InternLM3Model(TextModel): model_arch = gguf.MODEL_ARCH.LLAMA diff --git a/conversion/internvl.py b/conversion/internvl.py index 9a2a1e43d..799e23f5f 100644 --- a/conversion/internvl.py +++ b/conversion/internvl.py @@ -9,6 +9,7 @@ from .base import MmprojModel, ModelBase, gguf @ModelBase.register("InternVisionModel") +@ModelBase.example("OpenGVLab/InternVL3-2B", "OpenGVLab/InternVL2_5-1B") class InternVisionModel(MmprojModel): min_dynamic_tiles: int = 0 diff --git a/conversion/jais.py b/conversion/jais.py index 00add4c77..f3f96c3ef 100644 --- a/conversion/jais.py +++ b/conversion/jais.py @@ -11,6 +11,8 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("Jais2ForCausalLM") +# [TAG_HF_EXAMPLE_GATED] inceptionai/Jais-2-8B-Chat is gated +# [TAG_HF_EXAMPLE_MISSING] class Jais2Model(TextModel): model_arch = gguf.MODEL_ARCH.JAIS2 @@ -22,6 +24,7 @@ class Jais2Model(TextModel): @ModelBase.register("JAISLMHeadModel") +@ModelBase.example("inceptionai/jais-family-590m") class JaisModel(TextModel): model_arch = gguf.MODEL_ARCH.JAIS diff --git a/conversion/jamba.py b/conversion/jamba.py index da712ba50..a2e642cb0 100644 --- a/conversion/jamba.py +++ b/conversion/jamba.py @@ -11,6 +11,7 @@ from .base import ModelBase, TextModel, gguf, logger @ModelBase.register("JambaForCausalLM") +@ModelBase.example("ai21labs/Jamba-v0.1") class JambaModel(TextModel): model_arch = gguf.MODEL_ARCH.JAMBA diff --git a/conversion/januspro.py b/conversion/januspro.py index b49691205..0f71ab3cd 100644 --- a/conversion/januspro.py +++ b/conversion/januspro.py @@ -11,6 +11,7 @@ from .llama import LlamaModel @ModelBase.register("JanusForConditionalGeneration") +@ModelBase.example("deepseek-community/Janus-Pro-1B") class JanusProModel(LlamaModel): model_arch = gguf.MODEL_ARCH.LLAMA # reuse Llama arch @@ -34,6 +35,7 @@ class JanusProModel(LlamaModel): @ModelBase.register("JanusForConditionalGeneration") +@ModelBase.example("deepseek-community/Janus-Pro-1B") class JanusProVisionModel(MmprojModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) diff --git a/conversion/kimi_k3.py b/conversion/kimi_k3.py new file mode 100644 index 000000000..d15d1d64b --- /dev/null +++ b/conversion/kimi_k3.py @@ -0,0 +1,376 @@ +from __future__ import annotations + +import re +from pathlib import Path +from typing import Callable, Iterable, Iterator, TYPE_CHECKING + +import numpy as np +import torch + +if TYPE_CHECKING: + from torch import Tensor + +from .base import LazyTorchTensor, ModelBase, TextModel, gguf, logger + +from .kimi_linear import KimiLinearModel + + +@ModelBase.register("KimiK3ForConditionalGeneration") +@ModelBase.example("moonshotai/Kimi-K3") +class KimiK3Model(TextModel): + """ + Kimi-K3 text model (KimiLinearForCausalLM under a `language_model.` prefix). + + Shares the hybrid MLA + KDA skeleton with kimi-linear, but that converter + cannot load it: K3 adds cross-layer attention residuals, a latent MoE, the + situ activation, an MLA output gate and a full-rank KDA gate. + + The vision tower and mm_projector are skipped - text only for now. + """ + + model_arch = gguf.MODEL_ARCH.KIMI_K3 + + _experts: list[dict[str, Tensor]] | None = None + + # `_res_norm.weight` and `_res_proj.weight` are only used as their + # elementwise product, so they are fused into one [n_embd] vector here. + # they arrive apart, so buffer the first one and tag it with its kind. + _res_parts: dict[str, tuple[str, Tensor]] + + # HF suffix -> (gguf tensor, per-layer?) + _RES_FUSIONS = { + "self_attention_res": (gguf.MODEL_TENSOR.ATTN_RES_SCORE, True), + "mlp_res": (gguf.MODEL_TENSOR.FFN_RES_SCORE, True), + "output_attn_res": (gguf.MODEL_TENSOR.OUTPUT_RES_SCORE, False), + } + + # compressed-tensors MXFP4. the `language_model.` prefix is still there, as + # self.model_tensors is keyed by the raw checkpoint names + _MXFP4_FORMAT = "mxfp4-pack-quantized" + _MXFP4_EXPERT_RE = re.compile( + r"^(?:language_model\.)?model\.layers\.(\d+)" + r"\.block_sparse_moe\.experts\.(\d+)\.(w[123])\.weight_packed$" + ) + _MXFP4_PROJ = { + "w1": gguf.MODEL_TENSOR.FFN_GATE_EXP, + "w2": gguf.MODEL_TENSOR.FFN_DOWN_EXP, + "w3": gguf.MODEL_TENSOR.FFN_UP_EXP, + } + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._res_parts = {} + + def set_vocab(self): + # K3 has the same TikToken vocab as K2, so kimi-linear's vocab handling works. + # borrowed, not inherited: the method only touches TextModel members, and K3 + # shares none of kimi-linear's tensor layout. + KimiLinearModel.set_vocab(self) # ty: ignore[invalid-argument-type] + + # ...but that forces eos to the tokenizer's eos_id, which is [EOS], the + # document terminator. K3's config says <|end_of_msg|>, the turn terminator; + # with [EOS] the generation never stops at the end of a turn. + if (eos := self.hparams.get("eos_token_id")) is not None: + logger.info(f"restoring configured eos_token_id {eos} (kimi-linear forces the tokenizer's)") + self.gguf_writer.add_eos_token_id(eos) + + # K3 renders chats in python (encoding_k3.py) and ships no jinja template, + # so add the bundled one when the model has none + if gguf.SpecialVocab(self.dir_model, load_merges=False).chat_template is None: + template_path = Path(__file__).parent.parent / "models" / "templates" / "Kimi-K3.jinja" + logger.info(f"gguf: model has no chat template, using {template_path.name}") + self.gguf_writer.add_chat_template(template_path.read_text(encoding="utf-8")) + + # + # compressed-tensors MXFP4 -> ggml MXFP4 + # + + def _is_mxfp4_packed(self) -> bool: + quant_config = self.hparams.get("quantization_config") or {} + return (quant_config.get("quant_method") == "compressed-tensors" + and quant_config.get("format") == self._MXFP4_FORMAT) + + def dequant_model(self): + if not self._is_mxfp4_packed(): + return super().dequant_model() + + # skipping base.py's dequant is only safe if the experts are the only + # quantized tensors, so check it + stray = [n for n in self.model_tensors + if n.endswith(".weight_packed") and not self._MXFP4_EXPERT_RE.match(n)] + if stray: + raise NotImplementedError( + f"{len(stray)} MXFP4 tensor(s) outside the routed experts, e.g. {stray[0]!r}; " + "only the routed experts have a repack path" + ) + + def _mxfp4_expert_tensor(self, loaders: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]): + """ + One stacked [n_expert, rows, cols] MXFP4 tensor, built lazily. + + gguf_writer holds every added tensor until the final write, so building + this eagerly (like the DeepSeek-V4 path does) keeps all ~1.38 TB of + experts in memory. lazy means only the tensor being written is resident. + """ + # meta shapes, so this does not read any weights + rows, packed_cols = loaders[0][0]().shape + n_blocks = (packed_cols * 2) // 32 + byte_shape = (len(loaders), rows, n_blocks * 17) + + def load(fns: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]) -> np.ndarray: + out = np.empty(byte_shape, dtype=np.uint8) + for eid, (packed_fn, scale_fn) in enumerate(fns): + out[eid] = self.repack_mxfp4_blocks( + LazyTorchTensor.to_eager(packed_fn()), + LazyTorchTensor.to_eager(scale_fn()), + ) + return out + + # loaders goes through args, not the closure, so that `func` matches + # LazyBase's single-argument shape + return gguf.LazyNumpyTensor( + meta=gguf.LazyNumpyTensor.meta_with_dtype_and_shape(np.uint8, byte_shape), + args=(loaders,), + func=load, + ) + + def _write_mxfp4_experts(self) -> None: + n_experts = self.hparams["num_experts"] + + # (bid, wid) -> {expert id: (packed name, scale name)} + groups: dict[tuple[int, str], dict[int, tuple[str, str]]] = {} + for name in self.model_tensors: + m = self._MXFP4_EXPERT_RE.match(name) + if m is None: + continue + bid, eid, wid = int(m.group(1)), int(m.group(2)), m.group(3) + scale_name = name.removesuffix("_packed") + "_scale" + if scale_name not in self.model_tensors: + raise KeyError(f"missing {scale_name} for {name}") + groups.setdefault((bid, wid), {})[eid] = (name, scale_name) + + consumed: list[str] = [] + for (bid, wid), experts in sorted(groups.items()): + missing = [e for e in range(n_experts) if e not in experts] + if missing: + raise KeyError( + f"layer {bid} {wid}: {len(missing)} of {n_experts} experts missing, " + f"first is {missing[0]}" + ) + if len(experts) != n_experts: + raise KeyError(f"layer {bid} {wid}: {len(experts)} experts, expected {n_experts}") + + loaders = [] + for eid in range(n_experts): + packed_name, scale_name = experts[eid] + loaders.append((self.model_tensors[packed_name], self.model_tensors[scale_name])) + consumed += [packed_name, scale_name] + + data = self._mxfp4_expert_tensor(loaders) + new_name = self.format_tensor_name(self._MXFP4_PROJ[wid], bid) + shape = gguf.quant_shape_from_byte_shape(data.shape, gguf.GGMLQuantizationType.MXFP4) + logger.info( + f"{new_name}: repacked {n_experts} experts to MXFP4, " + f"shape = {{{', '.join(str(n) for n in reversed(shape))}}}" + ) + self.gguf_writer.add_tensor(new_name, data, raw_dtype=gguf.GGMLQuantizationType.MXFP4) + + for name in consumed: + del self.model_tensors[name] + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + # not a generator on purpose: base.py chains this with get_tensors(), so the + # tensors used here must be removed from model_tensors before that starts + if self._is_mxfp4_packed(): + self._write_mxfp4_experts() + return () + + def get_tensors(self) -> Iterator[tuple[str, Tensor]]: + for name, data in super().get_tensors(): + if name.startswith(("vision_tower.", "mm_projector.")): + continue # text only + if name.startswith("language_model."): + name = name[len("language_model."):] + yield name, data + + def set_gguf_parameters(self): + # MLA is served as MQA with a single large head, then decompressed + self.hparams["num_key_value_heads"] = 1 + + super().set_gguf_parameters() + self.gguf_writer.add_vocab_size(self.hparams["vocab_size"]) + + linear_attn_config = self.hparams["linear_attn_config"] + + # n_head_kv == 0 marks a KDA (recurrent) layer. the layer lists are 1-indexed, + # as KimiLinearConfig.is_kda_layer uses (layer_idx + 1) + full_attn_layers = linear_attn_config["full_attn_layers"] + n_kv_heads = [ + self.hparams["num_key_value_heads"] if (il + 1) in full_attn_layers else 0 + for il in range(self.hparams["num_hidden_layers"]) + ] + assert len(n_kv_heads) == self.hparams["num_hidden_layers"] + self.gguf_writer.add_head_count_kv(n_kv_heads) + + # --- KDA --- + self.gguf_writer.add_ssm_conv_kernel(linear_attn_config["short_conv_kernel_size"]) + self.gguf_writer.add_kda_head_dim(linear_attn_config["head_dim"]) + if (lb := linear_attn_config.get("gate_lower_bound")) is not None: + self.gguf_writer.add_kda_gate_lower_bound(lb) + + # --- MLA --- + if (q_lora_rank := self.hparams.get("q_lora_rank")) is not None: + self.gguf_writer.add_q_lora_rank(q_lora_rank) + kv_lora_rank = self.hparams["kv_lora_rank"] + self.gguf_writer.add_kv_lora_rank(kv_lora_rank) + + qk_nope_head_dim = self.hparams["qk_nope_head_dim"] + qk_rope_head_dim = self.hparams["qk_rope_head_dim"] + v_head_dim = self.hparams["v_head_dim"] + # K3 is nope-only; qk_rope_head_dim still sizes the un-absorbed part of K + assert self.hparams.get("mla_use_nope"), "K3 MLA is expected to be nope-only" + self.gguf_writer.add_rope_dimension_count(qk_rope_head_dim) + # MLA is served as MQA, so the cache holds the compressed latent + self.gguf_writer.add_key_length(kv_lora_rank + qk_rope_head_dim) + self.gguf_writer.add_value_length(kv_lora_rank) + self.gguf_writer.add_key_length_mla(qk_nope_head_dim + qk_rope_head_dim) + self.gguf_writer.add_value_length_mla(v_head_dim) + + # --- MoE --- + self.gguf_writer.add_expert_feed_forward_length(self.hparams["moe_intermediate_size"]) + self.gguf_writer.add_expert_shared_count(self.hparams["num_shared_experts"]) + self.gguf_writer.add_leading_dense_block_count(self.hparams["first_k_dense_replace"]) + self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"]) + self.gguf_writer.add_expert_weights_norm(self.hparams["moe_renormalize"]) + assert self.hparams["moe_router_activation_func"] == "sigmoid" + self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) + # latent MoE: routed experts live in a down-projected space + if (latent := self.hparams.get("routed_expert_hidden_size")) is not None: + self.gguf_writer.add_expert_latent_length(latent) + + # --- situ activation --- + assert self.hparams["hidden_act"] == "situ", \ + f"unexpected hidden_act {self.hparams['hidden_act']!r}" + self.gguf_writer.add_activation_situ_beta(self.hparams["activation_situ_beta"]) + self.gguf_writer.add_activation_situ_linear_beta(self.hparams["activation_situ_linear_beta"]) + + # --- cross-layer attention residuals --- + self.gguf_writer.add_attn_res_block_size(self.hparams["attn_res_block_size"]) + + def prepare_tensors(self): + super().prepare_tensors() + if self._experts is not None: + leftover = [k for d in self._experts for k in d.keys()] + if leftover: + raise ValueError(f"Unprocessed experts: {leftover}") + if self._res_parts: + raise ValueError(f"Unpaired attention-residual tensors: {sorted(self._res_parts)}") + if self._is_mxfp4_packed(): + # label the file for what it is; prepare_metadata runs after this + self._is_mxfp4 = True + self.ftype = gguf.LlamaFileType.MOSTLY_MXFP4_MOE + + def _try_fuse_res(self, data_torch: Tensor, name: str, bid: int | None): + """ + Pair _res_norm.weight with _res_proj.weight and emit their product. + + Returns None if this is not a res tensor, [] if buffered until its pair. + """ + for prefix, (tensor_id, per_layer) in self._RES_FUSIONS.items(): + for kind in ("norm", "proj"): + if not name.endswith(f"{prefix}_{kind}.weight"): + continue + key = f"{prefix}.{bid}" + other = self._res_parts.pop(key, None) + if other is None: + self._res_parts[key] = (kind, data_torch) + return [] + other_kind, other_data = other + assert other_kind != kind, f"duplicate {kind} for {key}" + norm = data_torch if kind == "norm" else other_data + proj = data_torch if kind == "proj" else other_data + fused = norm.float().flatten() * proj.float().flatten() + # ".weight" suffix matches the convention map_tensor_name applies + new_name = (self.format_tensor_name(tensor_id, bid) if per_layer + else gguf.TENSOR_NAMES[tensor_id] + ".weight") + logger.info(f"fused {prefix}_norm * {prefix}_proj -> {new_name}") + return [(new_name, fused)] + return None + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # --- cross-layer attention residuals: fuse norm * proj --- + fused = self._try_fuse_res(data_torch, name, bid) + if fused is not None: + yield from fused + return + + # --- KDA conv1d: HF [d_inner, 1, d_conv] -> ggml ne [d_conv, 1, d_inner, 1] --- + # GGUF reverses the numpy shape on write, so target numpy (1, d_inner, 1, d_conv). + # conv_step varies fastest in both layouts, so this is a pure reshape. + if name.endswith((".q_conv1d.weight", ".k_conv1d.weight", ".v_conv1d.weight")): + if data_torch.ndim == 3: # [d_inner, 1, d_conv] + d_inner, _, d_conv = data_torch.shape + elif data_torch.ndim == 2: # [d_inner, d_conv] + d_inner, d_conv = data_torch.shape + else: + raise ValueError(f"unexpected conv1d rank {data_torch.ndim} for {name}") + data_torch = data_torch.reshape(1, d_inner, 1, d_conv) + + # -exp(A_log) is folded here so the graph does not have to + if name.endswith(".A_log"): + n_head = self.hparams["num_attention_heads"] + data_torch = -torch.exp(data_torch.float()[:n_head]) + + # dt_bias -> the name SSM_DT's mapping expects + if name.endswith(".dt_bias"): + name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias" + + # --- g_proj is two different tensors sharing one HF name --- + # KDA layers: full-rank gate, [d_inner, n_embd] (replaces g_a/g_b) + # MLA layers: output gate, [n_head*v_head_dim, n_embd] + # Name-based mapping cannot tell them apart, so resolve by layer type. + if name.endswith(".self_attn.g_proj.weight"): + assert bid is not None + is_kda = (bid + 1) not in self.hparams["linear_attn_config"]["full_attn_layers"] + tensor_id = gguf.MODEL_TENSOR.SSM_G if is_kda else gguf.MODEL_TENSOR.ATTN_GATE + yield self.format_tensor_name(tensor_id, bid), data_torch + return + + # --- routed experts: stack per-expert 2D weights into one 3D tensor --- + if ".block_sparse_moe.experts." in name: + n_experts = self.hparams["num_experts"] + assert bid is not None + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + self._experts[bid][name] = data_torch + + if len(self._experts[bid]) < n_experts * 3: + return + + # w1: gate, w2: down, w3: up + for wid, tensor_id in (("w1", gguf.MODEL_TENSOR.FFN_GATE_EXP), + ("w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP), + ("w3", gguf.MODEL_TENSOR.FFN_UP_EXP)): + datas = [] + for xid in range(n_experts): + ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight" + datas.append(self._experts[bid].pop(ename)) + stacked = torch.stack(datas, dim=0) + yield from super().modify_tensors(stacked, self.format_tensor_name(tensor_id, bid), bid) + return + + # --- MLA absorption: split kv_b into k_b (transposed) and v_b --- + if name.endswith("kv_b_proj.weight"): + n_head_kv = self.hparams["num_key_value_heads"] + v_head_dim = self.hparams["v_head_dim"] + qk_nope_head_dim = self.hparams["qk_nope_head_dim"] + assert data_torch.shape[0] == n_head_kv * (v_head_dim + qk_nope_head_dim) + kv_b = data_torch.view(n_head_kv, v_head_dim + qk_nope_head_dim, data_torch.shape[-1]) + k_b, v_b = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1) + k_b = k_b.transpose(1, 2) + yield from super().modify_tensors(k_b, name.replace("kv_b_proj", "k_b_proj"), bid) + yield from super().modify_tensors(v_b, name.replace("kv_b_proj", "v_b_proj"), bid) + return + + yield from super().modify_tensors(data_torch, name, bid) diff --git a/conversion/kimi_linear.py b/conversion/kimi_linear.py index f2e6cda83..697ab1b4a 100644 --- a/conversion/kimi_linear.py +++ b/conversion/kimi_linear.py @@ -13,6 +13,7 @@ from .qwen import QwenModel @ModelBase.register("KimiLinearModel", "KimiLinearForCausalLM") +@ModelBase.example("moonshotai/Kimi-Linear-48B-A3B-Instruct") class KimiLinearModel(TextModel): """Kimi-Linear model with hybrid MLA+KDA architecture""" model_arch = gguf.MODEL_ARCH.KIMI_LINEAR diff --git a/conversion/kimivl.py b/conversion/kimivl.py index 5ff3c39ca..ae60abf30 100644 --- a/conversion/kimivl.py +++ b/conversion/kimivl.py @@ -11,6 +11,7 @@ from .base import MmprojModel, ModelBase, gguf @ModelBase.register("KimiVLForConditionalGeneration") +@ModelBase.example("moonshotai/Kimi-VL-A3B-Instruct") class KimiVLModel(MmprojModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -52,6 +53,7 @@ class KimiVLModel(MmprojModel): @ModelBase.register("KimiK25ForConditionalGeneration") +@ModelBase.example("moonshotai/Kimi-K2.5") class KimiK25Model(MmprojModel): """Kimi-K2.5 with MoonViT3d vision encoder""" @@ -155,6 +157,7 @@ class KimiK25Model(MmprojModel): @ModelBase.register("Glm5vForConditionalGeneration") +# [TAG_HF_EXAMPLE_MISSING] class Glm5vModel(KimiK25Model): """GLM-5.2-Vision MoonViT3d encoder and projector diff --git a/conversion/laguna.py b/conversion/laguna.py index a90f355ca..29e0b3d6b 100644 --- a/conversion/laguna.py +++ b/conversion/laguna.py @@ -13,6 +13,7 @@ from .base import ModelBase, TextModel, gguf, logger @ModelBase.register("LagunaForCausalLM") +@ModelBase.example("poolside/Laguna-XS.2", "poolside/Laguna-S-2.1") class LagunaModel(TextModel): model_arch = gguf.MODEL_ARCH.LAGUNA _experts: list[dict] | None = None diff --git a/conversion/lfm2.py b/conversion/lfm2.py index 70ce45658..984f44480 100644 --- a/conversion/lfm2.py +++ b/conversion/lfm2.py @@ -13,6 +13,7 @@ from .gemma import ConformerAudioModel @ModelBase.register("Lfm2ForCausalLM", "LFM2ForCausalLM") +@ModelBase.example("LiquidAI/LFM2-1.2B", "LiquidAI/LFM2.5-350M") class LFM2Model(TextModel): model_arch = gguf.MODEL_ARCH.LFM2 @@ -65,6 +66,7 @@ class LFM2Model(TextModel): @ModelBase.register("Lfm2Model", "Lfm2BidirectionalModel") +@ModelBase.example("LiquidAI/LFM2.5-ColBERT-350M", "LiquidAI/LFM2.5-Embedding-350M") class LFM2ColBertModel(LFM2Model): model_arch = gguf.MODEL_ARCH.LFM2 dense_tensor_name = "dense_2" @@ -93,6 +95,7 @@ class LFM2ColBertModel(LFM2Model): @ModelBase.register("Lfm2MoeForCausalLM") +@ModelBase.example("LiquidAI/LFM2-8B-A1B") class LFM2MoeModel(TextModel): model_arch = gguf.MODEL_ARCH.LFM2MOE @@ -166,6 +169,7 @@ class LFM2MoeModel(TextModel): @ModelBase.register("Lfm2VlForConditionalGeneration") +@ModelBase.example("LiquidAI/LFM2-VL-450M") class LFM2VLModel(MmprojModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -200,6 +204,7 @@ class LFM2VLModel(MmprojModel): @ModelBase.register("Lfm2AudioForConditionalGeneration") +@ModelBase.example("LiquidAI/LFM2.5-Audio-1.5B", "LiquidAI/LFM2-Audio-1.5B") class LFM2AudioModel(ConformerAudioModel): has_vision_encoder = False has_audio_encoder = True @@ -238,6 +243,7 @@ class LFM2AudioModel(ConformerAudioModel): @ModelBase.register("Lfm25AudioTokenizer") +@ModelBase.example("LiquidAI/LFM2.5-Audio-1.5B") class LFM25AudioTokenizer(LFM2Model): model_arch = gguf.MODEL_ARCH.LFM2 diff --git a/conversion/lighton_ocr.py b/conversion/lighton_ocr.py index ead3200ac..8686fe5c9 100644 --- a/conversion/lighton_ocr.py +++ b/conversion/lighton_ocr.py @@ -11,6 +11,7 @@ from .llava import LlavaVisionModel @ModelBase.register("LightOnOCRForConditionalGeneration") +@ModelBase.example("lightonai/LightOnOCR-1B-1025") class LightOnOCRVisionModel(LlavaVisionModel): is_mistral_format = False use_break_tok = False diff --git a/conversion/llada.py b/conversion/llada.py index 98dc9de95..c03607191 100644 --- a/conversion/llada.py +++ b/conversion/llada.py @@ -11,6 +11,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("LLaDAModelLM") +@ModelBase.example("GSAI-ML/LLaDA-8B-Instruct") class LLaDAModel(TextModel): model_arch = gguf.MODEL_ARCH.LLADA undo_permute = True @@ -114,6 +115,7 @@ class LLaDAModel(TextModel): @ModelBase.register("LLaDAMoEModel", "LLaDAMoEModelLM") +@ModelBase.example("inclusionAI/LLaDA-MoE-7B-A1B-Instruct") class LLaDAMoEModel(TextModel): model_arch = gguf.MODEL_ARCH.LLADA_MOE diff --git a/conversion/llama.py b/conversion/llama.py index 1aced49c5..41d8c2309 100644 --- a/conversion/llama.py +++ b/conversion/llama.py @@ -28,6 +28,8 @@ from .base import ModelBase, TextModel, gguf, logger "Eagle3DraftModel", "IQuestCoderForCausalLM", "LlamaModel") +# [TAG_HF_EXAMPLE_GATED] meta-llama/Llama-3.2-1B-Instruct is gated +@ModelBase.example("unsloth/Llama-3.2-1B-Instruct", "mistralai/Mistral-7B-Instruct-v0.3", "mistralai/Mixtral-8x7B-Instruct-v0.1") class LlamaModel(TextModel): model_arch = gguf.MODEL_ARCH.LLAMA undo_permute = True @@ -359,6 +361,7 @@ class LlamaModel(TextModel): @ModelBase.register("ArceeForCausalLM") +@ModelBase.example("arcee-ai/AFM-4.5B") class ArceeModel(LlamaModel): model_arch = gguf.MODEL_ARCH.ARCEE @@ -371,6 +374,8 @@ class ArceeModel(LlamaModel): "Llama4ForConditionalGeneration", "Llama4ForCausalLM", ) +# [TAG_HF_EXAMPLE_GATED] meta-llama/Llama-4-Scout-17B-16E-Instruct is gated +@ModelBase.example("unsloth/Llama-4-Scout-17B-16E-Instruct") class Llama4Model(LlamaModel): model_arch = gguf.MODEL_ARCH.LLAMA4 undo_permute = False @@ -412,16 +417,19 @@ class Llama4Model(LlamaModel): @ModelBase.register("LlamaBidirectionalModel") +@ModelBase.example("nvidia/llama-embed-nemotron-8b") class LlamaEmbedNemotronModel(LlamaModel): model_arch = gguf.MODEL_ARCH.LLAMA_EMBED @ModelBase.register("SmolLM3ForCausalLM") +@ModelBase.example("HuggingFaceTB/SmolLM3-3B") class SmolLM3Model(LlamaModel): model_arch = gguf.MODEL_ARCH.SMOLLM3 @ModelBase.register("ApertusForCausalLM") +@ModelBase.example("swiss-ai/Apertus-8B-Instruct-2509") class ApertusModel(LlamaModel): model_arch = gguf.MODEL_ARCH.APERTUS undo_permute = False diff --git a/conversion/llama4.py b/conversion/llama4.py index f84c76296..280e309dd 100644 --- a/conversion/llama4.py +++ b/conversion/llama4.py @@ -9,6 +9,8 @@ from .base import MmprojModel, ModelBase, gguf @ModelBase.register("Llama4ForConditionalGeneration") +# [TAG_HF_EXAMPLE_GATED] meta-llama/Llama-4-Scout-17B-16E-Instruct is gated +@ModelBase.example("unsloth/Llama-4-Scout-17B-16E-Instruct") class Llama4VisionModel(MmprojModel): def set_gguf_parameters(self): super().set_gguf_parameters() diff --git a/conversion/llava.py b/conversion/llava.py index 31d6e2ad8..98a004f98 100644 --- a/conversion/llava.py +++ b/conversion/llava.py @@ -16,6 +16,7 @@ from .llama import LlamaModel "LlavaForConditionalGeneration", # pixtral "Mistral3ForConditionalGeneration", # mistral small 3.1 ) +@ModelBase.example("mistral-community/pixtral-12b", "mistralai/Mistral-Small-3.1-24B-Instruct-2503") class LlavaVisionModel(MmprojModel): img_break_tok_id = -1 use_break_tok = True diff --git a/conversion/maincoder.py b/conversion/maincoder.py index 18b625b08..2e291b8a9 100644 --- a/conversion/maincoder.py +++ b/conversion/maincoder.py @@ -4,6 +4,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("MaincoderForCausalLM") +@ModelBase.example("Maincode/Maincoder-1B") class MaincoderModel(TextModel): model_arch = gguf.MODEL_ARCH.MAINCODER diff --git a/conversion/mamba.py b/conversion/mamba.py index 43d559ffb..8a2a46375 100644 --- a/conversion/mamba.py +++ b/conversion/mamba.py @@ -14,6 +14,7 @@ from .base import ModelBase, TextModel, gguf, logger @ModelBase.register("MambaForCausalLM", "MambaLMHeadModel", "FalconMambaForCausalLM") +@ModelBase.example("state-spaces/mamba-130m-hf", "tiiuae/falcon-mamba-7b") class MambaModel(TextModel): model_arch = gguf.MODEL_ARCH.MAMBA @@ -100,6 +101,7 @@ class MambaModel(TextModel): @ModelBase.register("Mamba2ForCausalLM") +@ModelBase.example("mistralai/Mamba-Codestral-7B-v0.1") class Mamba2Model(TextModel): model_arch = gguf.MODEL_ARCH.MAMBA2 diff --git a/conversion/mellum.py b/conversion/mellum.py index 79bc6755c..1e50f92ae 100644 --- a/conversion/mellum.py +++ b/conversion/mellum.py @@ -11,6 +11,7 @@ from .base import ModelBase, TextModel, gguf, logger @ModelBase.register("MellumForCausalLM") +@ModelBase.example("JetBrains/Mellum2-12B-A2.5B-Base") class MellumModel(TextModel): model_arch = gguf.MODEL_ARCH.MELLUM diff --git a/conversion/mimo.py b/conversion/mimo.py index ca2ed28ad..15dbeb7e7 100644 --- a/conversion/mimo.py +++ b/conversion/mimo.py @@ -14,6 +14,7 @@ from .base import MmprojModel, ModelBase, TextModel, gguf @ModelBase.register("MiMoV2FlashForCausalLM", "MiMoV2ForCausalLM") +@ModelBase.example("XiaomiMiMo/MiMo-V2.5") class MimoV2Model(TextModel): model_arch = gguf.MODEL_ARCH.MIMO2 @@ -230,6 +231,7 @@ class MimoV2Model(TextModel): @ModelBase.register("MiMoV2ForCausalLM") +@ModelBase.example("XiaomiMiMo/MiMo-V2.5") class MiMoV2VisionAudioModel(MmprojModel): has_audio_encoder = True diff --git a/conversion/minicpm.py b/conversion/minicpm.py index bf3fa8142..678d7bec1 100644 --- a/conversion/minicpm.py +++ b/conversion/minicpm.py @@ -14,6 +14,7 @@ from .qwen import Qwen3_5TextModel @ModelBase.register("MiniCPMForCausalLM") +@ModelBase.example("openbmb/MiniCPM-2B-sft-bf16") class MiniCPMModel(TextModel): model_arch = gguf.MODEL_ARCH.MINICPM @@ -61,6 +62,7 @@ class MiniCPMModel(TextModel): @ModelBase.register("MiniCPM3ForCausalLM") +@ModelBase.example("openbmb/MiniCPM3-4B") class MiniCPM3Model(TextModel): model_arch = gguf.MODEL_ARCH.MINICPM3 @@ -117,6 +119,7 @@ class MiniCPM3Model(TextModel): # the LM (text mode) and once as the mmproj (vision mode), mirroring the Qwen3-VL setup. @ModelBase.register("MiniCPMV4_6ForConditionalGeneration") +@ModelBase.example("openbmb/MiniCPM-V-4_6") class MiniCPMV4_6TextModel(Qwen3_5TextModel): model_arch = gguf.MODEL_ARCH.QWEN35 @@ -134,6 +137,7 @@ class MiniCPMV4_6TextModel(Qwen3_5TextModel): @ModelBase.register("MiniCPMV4_6ForConditionalGeneration") +@ModelBase.example("openbmb/MiniCPM-V-4_6") class MiniCPMV4_6VisionModel(MmprojModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) diff --git a/conversion/minimax.py b/conversion/minimax.py index d7a00bac9..53a9ff60f 100644 --- a/conversion/minimax.py +++ b/conversion/minimax.py @@ -12,6 +12,7 @@ from .base import ModelBase, TextModel, MmprojModel, gguf, logger @ModelBase.register("MiniMaxText01ForCausalLM") @ModelBase.register("MiniMaxM1ForCausalLM") +@ModelBase.example("MiniMaxAI/MiniMax-Text-01", "MiniMaxAI/MiniMax-M1-40k") class MiniMaxText01Model(TextModel): model_arch = gguf.MODEL_ARCH.MINIMAX01 @@ -119,6 +120,7 @@ class MiniMaxText01Model(TextModel): @ModelBase.register("MiniMaxM2ForCausalLM") +@ModelBase.example("MiniMaxAI/MiniMax-M2") class MiniMaxM2Model(TextModel): model_arch = gguf.MODEL_ARCH.MINIMAXM2 _experts_cache: dict[int, dict[str, Tensor]] = {} @@ -163,6 +165,7 @@ class MiniMaxM2Model(TextModel): @ModelBase.register("MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration") +@ModelBase.example("MiniMaxAI/MiniMax-M3") class MiniMaxM3Model(MiniMaxM2Model): model_arch = gguf.MODEL_ARCH.MINIMAXM3 @@ -203,6 +206,7 @@ class MiniMaxM3Model(MiniMaxM2Model): @ModelBase.register("MiniMaxM3SparseForConditionalGeneration", "MiniMaxM3VLForConditionalGeneration") +@ModelBase.example("MiniMaxAI/MiniMax-M3") class MiniMaxM3VisionModel(MmprojModel): @classmethod def filter_tensors(cls, item): diff --git a/conversion/mistral3.py b/conversion/mistral3.py index af9438ae7..fee039b35 100644 --- a/conversion/mistral3.py +++ b/conversion/mistral3.py @@ -15,6 +15,7 @@ from .llama import LlamaModel "Mistral3ForConditionalGeneration", "Ministral3ForCausalLM", ) +@ModelBase.example("mistralai/Mistral-Small-3.1-24B-Instruct-2503", "hf-tiny-v2/tiny-random-Ministral3ForCausalLM") class Mistral3Model(TextModel): class Ministral3Model(LlamaModel): model_arch = gguf.MODEL_ARCH.MISTRAL3 diff --git a/conversion/mpt.py b/conversion/mpt.py index 9557ab7fa..d5d849ff3 100644 --- a/conversion/mpt.py +++ b/conversion/mpt.py @@ -9,6 +9,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("MPTForCausalLM") +@ModelBase.example("anas-awadalla/mpt-7b") class MPTModel(TextModel): model_arch = gguf.MODEL_ARCH.MPT diff --git a/conversion/muse_glimmer.py b/conversion/muse_glimmer.py index cc588e832..b205f70a0 100644 --- a/conversion/muse_glimmer.py +++ b/conversion/muse_glimmer.py @@ -24,6 +24,7 @@ def _unpermute_for_rope(tensor: "Tensor", n_heads: int) -> "Tensor": @ModelBase.register("MuseGlimmerForConditionalGeneration") +@ModelBase.example("meta-models/Muse-Glimmer-30B") class MuseGlimmerModel(TextModel): model_arch = gguf.MODEL_ARCH.MUSE_GLIMMER @@ -78,6 +79,7 @@ class MuseGlimmerModel(TextModel): @ModelBase.register("MuseGlimmerForConditionalGeneration") +@ModelBase.example("meta-models/Muse-Glimmer-30B") class MuseGlimmerVisionModel(MmprojModel): def get_vision_config(self) -> dict[str, Any] | None: c = self.global_config.get("vision_config") @@ -131,6 +133,7 @@ class MuseGlimmerVisionModel(MmprojModel): @ModelBase.register("MuseGlimmerAssistantModel") +@ModelBase.example("meta-models/Muse-Glimmer-30B-assistant") class MuseGlimmerAssistantModel(TextModel): model_arch = gguf.MODEL_ARCH.DFLASH diff --git a/conversion/nanbeige.py b/conversion/nanbeige.py index f1fc425b3..a5b269a7a 100644 --- a/conversion/nanbeige.py +++ b/conversion/nanbeige.py @@ -5,6 +5,7 @@ from .llama import LlamaModel @ModelBase.register("NanbeigeForCausalLM") +@ModelBase.example("Nanbeige/Nanbeige4.2-3B") class NanbeigeModel(LlamaModel): model_arch = gguf.MODEL_ARCH.NANBEIGE undo_permute = True diff --git a/conversion/nemotron.py b/conversion/nemotron.py index c46cec143..3e37c7b46 100644 --- a/conversion/nemotron.py +++ b/conversion/nemotron.py @@ -16,6 +16,7 @@ from .granite import GraniteHybridModel "NemotronH_Nano_VL_V2", "RADIOModel", ) +@ModelBase.example("nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16") class NemotronNanoV2VLModel(MmprojModel): # ViT-Huge architecture parameters for RADIO v2.5-h _vit_hidden_size = 1280 @@ -151,6 +152,7 @@ class NemotronNanoV2VLModel(MmprojModel): @ModelBase.register("NemotronForCausalLM") +@ModelBase.example("nvidia/Minitron-4B-Base") class NemotronModel(TextModel): model_arch = gguf.MODEL_ARCH.NEMOTRON @@ -193,6 +195,7 @@ class NemotronModel(TextModel): @ModelBase.register("NemotronHForCausalLM") +@ModelBase.example("nvidia/Nemotron-H-8B-Base-8K") class NemotronHModel(GraniteHybridModel): """Hybrid mamba2/attention model from NVIDIA""" model_arch = gguf.MODEL_ARCH.NEMOTRON_H diff --git a/conversion/olmo.py b/conversion/olmo.py index 1664c30e4..e6faa1975 100644 --- a/conversion/olmo.py +++ b/conversion/olmo.py @@ -14,6 +14,7 @@ from .llama import LlamaModel @ModelBase.register("OlmoForCausalLM") @ModelBase.register("OLMoForCausalLM") +@ModelBase.example("allenai/OLMo-1.7-7B-hf") class OlmoModel(TextModel): model_arch = gguf.MODEL_ARCH.OLMO @@ -39,12 +40,14 @@ class OlmoModel(TextModel): @ModelBase.register("SeedOssForCausalLM") +@ModelBase.example("ByteDance-Seed/Seed-OSS-36B-Instruct") class SeedOssModel(TextModel): model_arch = gguf.MODEL_ARCH.SEED_OSS @ModelBase.register("Olmo2ForCausalLM") @ModelBase.register("Olmo3ForCausalLM") +@ModelBase.example("allenai/OLMo-2-1124-7B-Instruct", "allenai/Olmo-3-7B-Instruct") class Olmo2Model(TextModel): model_arch = gguf.MODEL_ARCH.OLMO2 @@ -67,6 +70,7 @@ class Olmo2Model(TextModel): @ModelBase.register("OlmoeForCausalLM") +@ModelBase.example("allenai/OLMoE-1B-7B-0924") class OlmoeModel(TextModel): model_arch = gguf.MODEL_ARCH.OLMOE diff --git a/conversion/openelm.py b/conversion/openelm.py index ecc746dc3..8863378bb 100644 --- a/conversion/openelm.py +++ b/conversion/openelm.py @@ -9,6 +9,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("OpenELMForCausalLM") +@ModelBase.example("apple/OpenELM-270M") class OpenELMModel(TextModel): model_arch = gguf.MODEL_ARCH.OPENELM diff --git a/conversion/orion.py b/conversion/orion.py index 8dfceeed1..3e4c633c1 100644 --- a/conversion/orion.py +++ b/conversion/orion.py @@ -4,6 +4,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("OrionForCausalLM") +@ModelBase.example("OrionStarAI/Orion-14B-Base") class OrionModel(TextModel): model_arch = gguf.MODEL_ARCH.ORION diff --git a/conversion/pangu.py b/conversion/pangu.py index 42016ba02..74c76532b 100644 --- a/conversion/pangu.py +++ b/conversion/pangu.py @@ -11,6 +11,7 @@ from .base import ModelBase, TextModel, gguf, logger @ModelBase.register("PanguEmbeddedForCausalLM") +@ModelBase.example("FreedomIntelligence/openPangu-Embedded-7B-V1.1") class PanguEmbeddedModel(TextModel): model_arch = gguf.MODEL_ARCH.PANGU_EMBED diff --git a/conversion/phi.py b/conversion/phi.py index df4bfe809..7d2532067 100644 --- a/conversion/phi.py +++ b/conversion/phi.py @@ -14,6 +14,7 @@ from .base import MmprojModel, ModelBase, SentencePieceTokenTypes, TextModel, gg @ModelBase.register("PhiForCausalLM") +@ModelBase.example("microsoft/phi-2") class Phi2Model(TextModel): model_arch = gguf.MODEL_ARCH.PHI2 @@ -36,6 +37,7 @@ class Phi2Model(TextModel): @ModelBase.register("Phi3ForCausalLM", "Phi4ForCausalLMV") +@ModelBase.example("microsoft/Phi-3-mini-4k-instruct") class Phi3MiniModel(TextModel): model_arch = gguf.MODEL_ARCH.PHI3 @@ -210,6 +212,7 @@ class Phi3MiniModel(TextModel): @ModelBase.register("Phi4ForCausalLMV") +# [TAG_HF_EXAMPLE_MISSING] class Phi4VisionMmprojModel(MmprojModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -336,6 +339,7 @@ class Phi4VisionMmprojModel(MmprojModel): @ModelBase.register("PhiMoEForCausalLM") +@ModelBase.example("microsoft/Phi-3.5-MoE-instruct") class PhiMoeModel(Phi3MiniModel): model_arch = gguf.MODEL_ARCH.PHIMOE diff --git a/conversion/plamo.py b/conversion/plamo.py index c4bcbdf06..31c6455aa 100644 --- a/conversion/plamo.py +++ b/conversion/plamo.py @@ -13,6 +13,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("PlamoForCausalLM") +@ModelBase.example("pfnet/plamo-13b") class PlamoModel(TextModel): model_arch = gguf.MODEL_ARCH.PLAMO @@ -58,6 +59,7 @@ class PlamoModel(TextModel): @ModelBase.register("Plamo2ForCausalLM", "PLaMo2ForCausalLM") +@ModelBase.example("pfnet/plamo-2-1b") class Plamo2Model(TextModel): model_arch = gguf.MODEL_ARCH.PLAMO2 @@ -147,6 +149,8 @@ class Plamo2Model(TextModel): @ModelBase.register("Plamo3ForCausalLM", "PLaMo3ForCausalLM") +# [TAG_HF_EXAMPLE_GATED] pfnet/plamo-3-nict-2b-base is gated +@ModelBase.example("midorin-Linux/plamo-3-12b-self-merged-base") class Plamo3Model(TextModel): model_arch = gguf.MODEL_ARCH.PLAMO3 diff --git a/conversion/plm.py b/conversion/plm.py index 3fde48708..bca0147e6 100644 --- a/conversion/plm.py +++ b/conversion/plm.py @@ -4,6 +4,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("PLMForCausalLM") +@ModelBase.example("PLM-Team/PLM-1.8B-Instruct") class PLMModel(TextModel): model_arch = gguf.MODEL_ARCH.PLM diff --git a/conversion/pockettts.py b/conversion/pockettts.py index 62ecb5acd..1c99e58cf 100644 --- a/conversion/pockettts.py +++ b/conversion/pockettts.py @@ -77,6 +77,7 @@ def _load_hparams(dir_model: Path) -> dict[str, Any]: @ModelBase.register("PocketTTSModel") +# [TAG_HF_EXAMPLE_MISSING] model is gated, and the checkpoint requires cd to subdir, not supported here class PocketTTSModel(TextModel): model_arch = gguf.MODEL_ARCH.POCKETTTS @@ -174,6 +175,7 @@ class PocketTTSModel(TextModel): @ModelBase.register("PocketTTSModel") +# [TAG_HF_EXAMPLE_MISSING] model is gated, and the checkpoint requires cd to subdir, not supported here class PocketTTSMmprojModel(MmprojModel): has_audio_encoder = True has_vision_encoder = False diff --git a/conversion/qwen.py b/conversion/qwen.py index ead435455..26b10452b 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -4,15 +4,17 @@ import json from typing import Any, Callable, Iterable, TYPE_CHECKING +import numpy as np import torch if TYPE_CHECKING: from torch import Tensor -from .base import ModelBase, TextModel, gguf, logger +from .base import LazyTorchTensor, ModelBase, TextModel, gguf, logger @ModelBase.register("QWenLMHeadModel") +@ModelBase.example("Qwen/Qwen-7B") class QwenModel(TextModel): model_arch = gguf.MODEL_ARCH.QWEN @@ -51,6 +53,7 @@ class QwenModel(TextModel): "AudioFlamingo3ForConditionalGeneration", "DotsOCRForCausalLM", ) +@ModelBase.example("Qwen/Qwen2.5-7B-Instruct") class Qwen2Model(TextModel): model_arch = gguf.MODEL_ARCH.QWEN2 @@ -71,6 +74,7 @@ class Qwen2Model(TextModel): @ModelBase.register("Qwen2MoeForCausalLM") +@ModelBase.example("Qwen/Qwen1.5-MoE-A2.7B") class Qwen2MoeModel(TextModel): model_arch = gguf.MODEL_ARCH.QWEN2MOE @@ -153,6 +157,7 @@ class Qwen2MoeModel(TextModel): @ModelBase.register("Qwen3ForCausalLM", "Qwen3Model") +@ModelBase.example("Qwen/Qwen3-8B") class Qwen3Model(Qwen2Model): model_arch = gguf.MODEL_ARCH.QWEN3 @@ -251,6 +256,7 @@ class Qwen3Model(Qwen2Model): @ModelBase.register("Qwen3MoeForCausalLM") +@ModelBase.example("Qwen/Qwen3-30B-A3B") class Qwen3MoeModel(Qwen2MoeModel): model_arch = gguf.MODEL_ARCH.QWEN3MOE @@ -362,6 +368,7 @@ class _QwenMtpMixin: @ModelBase.register("Qwen3NextForCausalLM") +@ModelBase.example("Qwen/Qwen3-Next-80B-A3B-Instruct") class Qwen3NextModel(_QwenMtpMixin, Qwen2MoeModel): model_arch = gguf.MODEL_ARCH.QWEN3NEXT @@ -421,6 +428,7 @@ class Qwen3NextModel(_QwenMtpMixin, Qwen2MoeModel): @ModelBase.register("RND1") +@ModelBase.example("radicalnumerics/RND1-Base-0910") class RND1Model(Qwen2MoeModel): model_arch = gguf.MODEL_ARCH.RND1 @@ -620,16 +628,19 @@ class _Qwen35MRopeMixin: @ModelBase.register("Qwen3_5ForConditionalGeneration", "Qwen3_5ForCausalLM") +@ModelBase.example("Qwen/Qwen3.5-9B") class Qwen3_5TextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase): model_arch = gguf.MODEL_ARCH.QWEN35 @ModelBase.register("Qwen3_5MoeForConditionalGeneration", "Qwen3_5MoeForCausalLM") +@ModelBase.example("Qwen/Qwen3.5-35B-A3B") class Qwen3_5MoeTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase): model_arch = gguf.MODEL_ARCH.QWEN35MOE @ModelBase.register("DFlashDraftModel") +@ModelBase.example("z-lab/Qwen3.5-9B-DFlash") class DFlashModel(Qwen3Model): model_arch = gguf.MODEL_ARCH.DFLASH @@ -698,21 +709,82 @@ class DFlashModel(Qwen3Model): yield from super().modify_tensors(data_torch, name, bid) -@ModelBase.register("Qwen3DSparkModel") +@ModelBase.register("Qwen3DSparkModel", "DSparkDraftModel", "DSparkSpeculator") +@ModelBase.example("satgeze/Qwen3.6-27B-DSpark") class DSparkModel(DFlashModel): - # DSpark = DFlash + a semi-autoregressive Markov head + # DSpark = DFlash + a semi-autoregressive Markov head. model_arch = gguf.MODEL_ARCH.DFLASH - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - # normalize the flat DeepSpec schema to DFlash's nested dflash_config - self.hparams.setdefault("dflash_config", { - k: self.hparams[k] for k in ("target_layer_ids", "mask_token_id") if k in self.hparams - }) + def __init__(self, dir_model, *args, **kwargs): + hparams = kwargs.pop("hparams", None) + if hparams is None: + hparams = ModelBase.load_hparams(dir_model, False) + + # EAGLE3-style exports use the 1+N bonus-anchor block, DFlash-lineage exports sample from the anchor + self._sample_from_anchor = hparams.get( + "sample_from_anchor", + "transformer_layer_config" not in hparams and "aux_hidden_state_layer_ids" not in hparams) + if "transformer_layer_config" in hparams: + hparams = {**hparams, **hparams["transformer_layer_config"]} + + super().__init__(dir_model, *args, hparams=hparams, **kwargs) + + # normalize both schemas to DFlash's nested dflash_config + if "aux_hidden_state_layer_ids" in self.hparams: + self.hparams.setdefault("dflash_config", { + "mask_token_id": self.hparams.get("mask_token_id"), + "target_layer_ids": [i - 1 for i in self.hparams["aux_hidden_state_layer_ids"]], + }) + else: + self.hparams.setdefault("dflash_config", { + k: self.hparams[k] for k in ("target_layer_ids", "mask_token_id") if k in self.hparams + }) + + if (markov_head_type := self.hparams.get("markov_head_type", "vanilla")) != "vanilla": + raise ValueError(f"unsupported markov_head_type {markov_head_type!r} (only 'vanilla' is supported)") + + n_vocab = self.hparams["vocab_size"] + self._n_vocab_draft = self.hparams.get("draft_vocab_size") or n_vocab + if self._n_vocab_draft > n_vocab: + raise ValueError(f"draft_vocab_size {self._n_vocab_draft} exceeds vocab_size {n_vocab}") + self._d2t: Tensor | None = None + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_sample_from_anchor(self._sample_from_anchor) @classmethod def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: - name, gen = item - if name.endswith(("embed_tokens.weight", "lm_head.weight")): + if item[0] == "t2d": # not used at runtime return None - return super().filter_tensors((name, gen)) + return super().filter_tensors(item) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name == "model.d2t": + self._d2t = data_torch + return + + if self._n_vocab_draft == self.hparams["vocab_size"] and name.endswith(("embed_tokens.weight", "lm_head.weight")): + return + + yield from super().modify_tensors(data_torch, name, bid) + + def prepare_tensors(self): + super().prepare_tensors() + + n_vocab = self.hparams["vocab_size"] + if self._n_vocab_draft < n_vocab and self._d2t is None: + raise ValueError(f"draft_vocab_size {self._n_vocab_draft} < vocab_size {n_vocab} but no d2t table found") + + # write d2t as absolute target token ids + if self._d2t is not None: + data = LazyTorchTensor.to_eager(self._d2t).to(torch.int64).cpu().numpy().reshape(-1) + if data.size != self._n_vocab_draft: + raise ValueError(f"d2t size {data.size} does not match draft_vocab_size {self._n_vocab_draft}") + data = data + np.arange(data.size, dtype=np.int64) + if np.any((data < 0) | (data >= n_vocab)): + raise ValueError(f"d2t target ids out of range for target vocab size {n_vocab}") + if np.unique(data).size != data.size: + raise ValueError("d2t contains duplicate target ids") + logger.info(f"{'d2t,':<30} --> I64, shape = {{{data.size}}}") + self.gguf_writer.add_tensor("d2t", data, raw_dtype=gguf.GGMLQuantizationType.I64) diff --git a/conversion/qwen3tts.py b/conversion/qwen3tts.py index d21a50595..1f6b9a1b0 100644 --- a/conversion/qwen3tts.py +++ b/conversion/qwen3tts.py @@ -37,6 +37,7 @@ _ACT2FN = { @ModelBase.register("Qwen3TTSForConditionalGeneration") +@ModelBase.example("Qwen/Qwen3-TTS-12Hz-1.7B-Base") class Qwen3TTSTalkerModel(TextModel): model_arch = gguf.MODEL_ARCH.QWEN3TTS @@ -185,6 +186,7 @@ class Qwen3TTSTalkerModel(TextModel): @ModelBase.register("Qwen3TTSForConditionalGeneration") +@ModelBase.example("Qwen/Qwen3-TTS-12Hz-1.7B-Base") class Qwen3TTSSpeakerEncoderModel(MmprojModel): has_vision_encoder = False has_audio_encoder = True diff --git a/conversion/qwen3vl.py b/conversion/qwen3vl.py index 9f1175769..4fec708c9 100644 --- a/conversion/qwen3vl.py +++ b/conversion/qwen3vl.py @@ -14,6 +14,7 @@ from .qwenvl import Qwen25AudioModel @ModelBase.register("Qwen3VLForConditionalGeneration", "Qwen3VLMoeForConditionalGeneration", "Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration") +@ModelBase.example("Qwen/Qwen3-VL-4B-Instruct", "Qwen/Qwen3-VL-30B-A3B-Instruct", "Qwen/Qwen3.5-9B", "Qwen/Qwen3.5-35B-A3B") class Qwen3VLVisionModel(MmprojModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -144,6 +145,7 @@ class Qwen3VLVisionModel(MmprojModel): @ModelBase.register("Qwen3OmniMoeForConditionalGeneration") +@ModelBase.example("Qwen/Qwen3-Omni-30B-A3B-Instruct") class Qwen3OmniMmprojModel(Qwen3VLVisionModel, Qwen25AudioModel): has_audio_encoder = True has_vision_encoder = True @@ -217,12 +219,14 @@ class Qwen3OmniMmprojModel(Qwen3VLVisionModel, Qwen25AudioModel): @ModelBase.register("Qwen3ASRForConditionalGeneration") +@ModelBase.example("Qwen/Qwen3-ASR-0.6B-hf") class Qwen3ASRMmprojModel(Qwen3OmniMmprojModel): has_audio_encoder = True has_vision_encoder = False @ModelBase.register("Glm4vForConditionalGeneration", "Glm4vMoeForConditionalGeneration", "GlmOcrForConditionalGeneration") +@ModelBase.example("zai-org/GLM-4.1V-9B-Thinking", "zai-org/GLM-4.5V") class Glm4VVisionModel(Qwen3VLVisionModel): def set_gguf_parameters(self): MmprojModel.set_gguf_parameters(self) # skip Qwen3VLVisionModel parameters @@ -246,6 +250,7 @@ class Glm4VVisionModel(Qwen3VLVisionModel): @ModelBase.register("Qwen3VLForConditionalGeneration") +@ModelBase.example("Qwen/Qwen3-VL-4B-Instruct") class Qwen3VLTextModel(Qwen3Model): model_arch = gguf.MODEL_ARCH.QWEN3VL @@ -268,6 +273,7 @@ class Qwen3VLTextModel(Qwen3Model): @ModelBase.register("Qwen3VLMoeForConditionalGeneration") +@ModelBase.example("Qwen/Qwen3-VL-30B-A3B-Instruct") class Qwen3VLMoeTextModel(Qwen3MoeModel): model_arch = gguf.MODEL_ARCH.QWEN3VLMOE @@ -317,6 +323,7 @@ class Qwen3VLMoeTextModel(Qwen3MoeModel): @ModelBase.register("Qwen3OmniMoeForConditionalGeneration") +@ModelBase.example("Qwen/Qwen3-Omni-30B-A3B-Instruct") class Qwen3OmniMoeTextModel(Qwen3VLMoeTextModel): model_arch = gguf.MODEL_ARCH.QWEN3VLMOE @@ -338,6 +345,7 @@ class Qwen3OmniMoeTextModel(Qwen3VLMoeTextModel): @ModelBase.register("Qwen3ASRForConditionalGeneration") +@ModelBase.example("Qwen/Qwen3-ASR-0.6B-hf") class Qwen3ASRTextModel(Qwen3VLTextModel): model_arch = gguf.MODEL_ARCH.QWEN3VL diff --git a/conversion/qwenvl.py b/conversion/qwenvl.py index 202a47961..579a86a99 100644 --- a/conversion/qwenvl.py +++ b/conversion/qwenvl.py @@ -17,6 +17,7 @@ from .base import MmprojModel, ModelBase, TextModel, gguf "Qwen2_5_VLForConditionalGeneration", "Qwen2_5OmniModel", ) +@ModelBase.example("Qwen/Qwen2-VL-2B-Instruct", "Qwen/Qwen2.5-VL-3B-Instruct") class Qwen2VLModel(TextModel): model_arch = gguf.MODEL_ARCH.QWEN2VL @@ -40,6 +41,7 @@ class Qwen2VLModel(TextModel): @ModelBase.register("Qwen2VLModel", "Qwen2VLForConditionalGeneration", "Qwen2_5_VLForConditionalGeneration") +@ModelBase.example("Qwen/Qwen2-VL-2B-Instruct", "Qwen/Qwen2.5-VL-3B-Instruct") class Qwen2VLVisionModel(MmprojModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -161,6 +163,7 @@ class Qwen25AudioModel(MmprojModel): @ModelBase.register("Qwen2_5OmniModel") +@ModelBase.example("Qwen/Qwen2.5-Omni-3B") class Qwen25OmniModel(Qwen2VLVisionModel, Qwen25AudioModel): has_audio_encoder = True has_vision_encoder = True diff --git a/conversion/refact.py b/conversion/refact.py index 1170cddeb..d6361512f 100644 --- a/conversion/refact.py +++ b/conversion/refact.py @@ -9,6 +9,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("GPTRefactForCausalLM") +@ModelBase.example("smallcloudai/Refact-1_6-base") class RefactModel(TextModel): model_arch = gguf.MODEL_ARCH.REFACT diff --git a/conversion/rwkv.py b/conversion/rwkv.py index 2de0aa534..e6fa84264 100644 --- a/conversion/rwkv.py +++ b/conversion/rwkv.py @@ -11,6 +11,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("Rwkv6ForCausalLM") +@ModelBase.example("RWKV/v6-Finch-1B6-HF") class Rwkv6Model(TextModel): model_arch = gguf.MODEL_ARCH.RWKV6 @@ -83,6 +84,7 @@ class Rwkv6Model(TextModel): @ModelBase.register("RWKV6Qwen2ForCausalLM") +@ModelBase.example("recursal/QRWKV6-32B-Instruct-Preview-v0.1") class RWKV6Qwen2Model(Rwkv6Model): model_arch = gguf.MODEL_ARCH.RWKV6QWEN2 @@ -136,6 +138,7 @@ class RWKV6Qwen2Model(Rwkv6Model): @ModelBase.register("Rwkv7ForCausalLM", "RWKV7ForCausalLM") +@ModelBase.example("fla-hub/rwkv7-1.5B-world") class Rwkv7Model(TextModel): model_arch = gguf.MODEL_ARCH.RWKV7 @@ -261,6 +264,7 @@ class Rwkv7Model(TextModel): @ModelBase.register("RwkvHybridForCausalLM") +@ModelBase.example("RWKV-Red-Team/ARWKV-7B-Preview-0.1") class ARwkv7Model(Rwkv7Model): model_arch = gguf.MODEL_ARCH.ARWKV7 diff --git a/conversion/sarashina2.py b/conversion/sarashina2.py index 05448db81..fdb3e78da 100644 --- a/conversion/sarashina2.py +++ b/conversion/sarashina2.py @@ -12,6 +12,7 @@ from .qwenvl import Qwen2VLVisionModel @ModelBase.register("Sarashina2VisionForCausalLM") +@ModelBase.example("sbintuitions/sarashina2.2-vision-3b") class Sarashina2VLTextModel(LlamaModel): model_arch = gguf.MODEL_ARCH.LLAMA @@ -26,6 +27,7 @@ class Sarashina2VLTextModel(LlamaModel): @ModelBase.register("Sarashina2VisionForCausalLM") +@ModelBase.example("sbintuitions/sarashina2.2-vision-3b") class Sarashina2VLVisionModel(Qwen2VLVisionModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) diff --git a/conversion/smallthinker.py b/conversion/smallthinker.py index 1b0f79aa3..73d07b51a 100644 --- a/conversion/smallthinker.py +++ b/conversion/smallthinker.py @@ -11,6 +11,7 @@ from .base import ModelBase, TextModel, gguf, logger @ModelBase.register("SmallThinkerForCausalLM") +@ModelBase.example("PowerInfer/SmallThinker-4BA0.6B-Instruct") class SmallThinkerModel(TextModel): model_arch = gguf.MODEL_ARCH.SMALLTHINKER diff --git a/conversion/smolvlm.py b/conversion/smolvlm.py index 30e9dca32..0cccb8f6f 100644 --- a/conversion/smolvlm.py +++ b/conversion/smolvlm.py @@ -9,6 +9,7 @@ from .base import MmprojModel, ModelBase, gguf @ModelBase.register("Idefics3ForConditionalGeneration", "SmolVLMForConditionalGeneration") +@ModelBase.example("HuggingFaceTB/SmolVLM-Instruct", "HuggingFaceM4/Idefics3-8B-Llama3") class SmolVLMModel(MmprojModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) diff --git a/conversion/stablelm.py b/conversion/stablelm.py index 6e16378a0..ac3a1ca9e 100644 --- a/conversion/stablelm.py +++ b/conversion/stablelm.py @@ -11,6 +11,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("StableLmForCausalLM", "StableLMEpochForCausalLM", "LlavaStableLMEpochForCausalLM") +@ModelBase.example("stabilityai/stablelm-2-1_6b") class StableLMModel(TextModel): model_arch = gguf.MODEL_ARCH.STABLELM diff --git a/conversion/starcoder.py b/conversion/starcoder.py index 0b4ffd847..4a726ac36 100644 --- a/conversion/starcoder.py +++ b/conversion/starcoder.py @@ -4,6 +4,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("GPTBigCodeForCausalLM") +@ModelBase.example("bigcode/gpt_bigcode-santacoder") class StarCoderModel(TextModel): model_arch = gguf.MODEL_ARCH.STARCODER @@ -19,5 +20,6 @@ class StarCoderModel(TextModel): @ModelBase.register("Starcoder2ForCausalLM") +@ModelBase.example("bigcode/starcoder2-3b") class StarCoder2Model(TextModel): model_arch = gguf.MODEL_ARCH.STARCODER2 diff --git a/conversion/step3.py b/conversion/step3.py index f7cdc997e..93eb3134e 100644 --- a/conversion/step3.py +++ b/conversion/step3.py @@ -16,6 +16,7 @@ from .qwen import Qwen3Model @ModelBase.register("StepVLForConditionalGeneration", "Step3p7ForConditionalGeneration") +@ModelBase.example("stepfun-ai/Step3-VL-10B", "stepfun-ai/Step-3.7-Flash") class Step3VLVisionModel(MmprojModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -91,11 +92,13 @@ class Step3VLVisionModel(MmprojModel): @ModelBase.register("StepVLForConditionalGeneration") +@ModelBase.example("stepfun-ai/Step3-VL-10B") class Step3VLTextModel(Qwen3Model): model_arch = gguf.MODEL_ARCH.QWEN3 @ModelBase.register("Step3p5ForCausalLM", "Step3p7ForConditionalGeneration") +@ModelBase.example("stepfun-ai/Step-3.7-Flash") class Step35Model(TextModel): model_arch = gguf.MODEL_ARCH.STEP35 supports_mtp_export = True diff --git a/conversion/t5.py b/conversion/t5.py index 73dcfd1a2..3466ce49d 100644 --- a/conversion/t5.py +++ b/conversion/t5.py @@ -16,6 +16,7 @@ from .base import ModelBase, SentencePieceTokenTypes, TextModel, gguf, logger @ModelBase.register("MT5ForConditionalGeneration") @ModelBase.register("UMT5ForConditionalGeneration") @ModelBase.register("UMT5Model") +@ModelBase.example("google-t5/t5-small", "google/flan-t5-small", "google/umt5-small") class T5Model(TextModel): model_arch = gguf.MODEL_ARCH.T5 @@ -153,6 +154,7 @@ class T5Model(TextModel): @ModelBase.register("T5EncoderModel") +@ModelBase.example("sentence-transformers/sentence-t5-base") class T5EncoderModel(TextModel): model_arch = gguf.MODEL_ARCH.T5ENCODER diff --git a/conversion/talkie.py b/conversion/talkie.py index a970b32d3..31445243d 100644 --- a/conversion/talkie.py +++ b/conversion/talkie.py @@ -11,6 +11,7 @@ from .base import LazyTorchTensor, ModelBase, TextModel, gguf @ModelBase.register("TalkieForCausalLM") +@ModelBase.example("lewtun/talkie-1930-13b-it-hf") class TalkieModel(TextModel): model_arch = gguf.MODEL_ARCH.TALKIE diff --git a/conversion/ultravox.py b/conversion/ultravox.py index 347188733..62819e574 100644 --- a/conversion/ultravox.py +++ b/conversion/ultravox.py @@ -9,6 +9,7 @@ from .base import MmprojModel, ModelBase, TextModel, gguf @ModelBase.register("UltravoxModel") +@ModelBase.example("fixie-ai/ultravox-v0_5-llama-3_2-1b") class UltravoxModel(TextModel): model_arch = gguf.MODEL_ARCH.LLAMA # dummy @@ -18,6 +19,7 @@ class UltravoxModel(TextModel): @ModelBase.register("GlmasrModel") +@ModelBase.example("zai-org/GLM-ASR-Nano-2512") class GlmASRWhisperEncoderModel(MmprojModel): has_vision_encoder = False has_audio_encoder = True @@ -82,6 +84,7 @@ class GlmASRWhisperEncoderModel(MmprojModel): @ModelBase.register("Qwen2AudioForConditionalGeneration") +@ModelBase.example("Qwen/Qwen2-Audio-7B-Instruct") class WhisperEncoderModel(MmprojModel): has_vision_encoder = False # no vision encoder has_audio_encoder = True @@ -123,6 +126,7 @@ class WhisperEncoderModel(MmprojModel): @ModelBase.register("UltravoxModel") +@ModelBase.example("fixie-ai/ultravox-v0_5-llama-3_2-1b") class UltravoxWhisperEncoderModel(WhisperEncoderModel): has_vision_encoder = False # no vision encoder has_audio_encoder = True @@ -134,6 +138,7 @@ class UltravoxWhisperEncoderModel(WhisperEncoderModel): @ModelBase.register("MERaLiON2ForConditionalGeneration") +@ModelBase.example("MERaLiON/MERaLiON-2-3B") class MERaLiONWhisperEncoderModel(WhisperEncoderModel): has_vision_encoder = False has_audio_encoder = True @@ -180,6 +185,7 @@ class MERaLiONWhisperEncoderModel(WhisperEncoderModel): @ModelBase.register("VoxtralForConditionalGeneration") +@ModelBase.example("mistralai/Voxtral-Mini-3B-2507") class VoxtralWhisperEncoderModel(WhisperEncoderModel): has_vision_encoder = False # no vision encoder has_audio_encoder = True @@ -191,6 +197,7 @@ class VoxtralWhisperEncoderModel(WhisperEncoderModel): @ModelBase.register("AudioFlamingo3ForConditionalGeneration") +@ModelBase.example("nvidia/audio-flamingo-3-hf") class AudioFlamingo3WhisperEncoderModel(WhisperEncoderModel): def set_gguf_parameters(self): super().set_gguf_parameters() diff --git a/conversion/wavtokenizer.py b/conversion/wavtokenizer.py index 7d25447be..c9a4b505d 100644 --- a/conversion/wavtokenizer.py +++ b/conversion/wavtokenizer.py @@ -9,6 +9,7 @@ from .base import ModelBase, TextModel, gguf, logger @ModelBase.register("WavTokenizerDec") +@ModelBase.example("novateur/WavTokenizer-large-speech-75token") class WavTokenizerDecModel(TextModel): model_arch = gguf.MODEL_ARCH.WAVTOKENIZER_DEC diff --git a/conversion/xverse.py b/conversion/xverse.py index fa8a31a13..aa3b33880 100644 --- a/conversion/xverse.py +++ b/conversion/xverse.py @@ -11,6 +11,7 @@ from .base import ModelBase, TextModel, gguf @ModelBase.register("XverseForCausalLM") +@ModelBase.example("xverse/XVERSE-7B") class XverseModel(TextModel): model_arch = gguf.MODEL_ARCH.XVERSE diff --git a/conversion/youtuvl.py b/conversion/youtuvl.py index cabc44445..e97261077 100644 --- a/conversion/youtuvl.py +++ b/conversion/youtuvl.py @@ -9,6 +9,7 @@ from .base import MmprojModel, ModelBase, gguf, logger @ModelBase.register("YoutuVLForConditionalGeneration") +@ModelBase.example("tencent/Youtu-VL-4B-Instruct") class YoutuVLVisionModel(MmprojModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) diff --git a/docs/release.md b/docs/release.md index 4335ef9d4..e0c9c486b 100644 --- a/docs/release.md +++ b/docs/release.md @@ -29,6 +29,12 @@ identify which PRs require a version bump before cutting a release._ Releases are created by running the [make-release](.github/workflows/make-release.yml) which is a manual workflow. +The workflow runs against the branch selected in the "Run workflow" dialog +(default `master`) and takes an optional `commit` SHA. When a commit is given, +the workflow validates that the commit belongs to the branch and is not older +than 3 days from the branch HEAD, then releases that commit instead of the +branch HEAD. + The workflow creates an annotated git tag (e.g. `v0.1.0`) and pushes it to the remote. No GitHub Release object is created, the tag is the release artifact. diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index a312c8aa0..d043c9b6e 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -124,6 +124,7 @@ class Keys: EXPERT_WEIGHTS_NORM = "{arch}.expert_weights_norm" EXPERT_GATING_FUNC = "{arch}.expert_gating_func" EXPERT_GROUP_SCALE = "{arch}.expert_group_scale" + EXPERT_LATENT_LENGTH = "{arch}.expert_latent_length" EXPERTS_PER_GROUP = "{arch}.experts_per_group" MOE_EVERY_N_LAYERS = "{arch}.moe_every_n_layers" MOE_LATENT_SIZE = "{arch}.moe_latent_size" @@ -161,6 +162,7 @@ class Keys: TARGET_LAYERS = "{arch}.target_layers" TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size" BLOCK_SIZE = "{arch}.block_size" + SAMPLE_FROM_ANCHOR = "{arch}.sample_from_anchor" NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" NORM_BEFORE_FC = "{arch}.norm_before_fc" @@ -238,6 +240,13 @@ class Keys: SCALING_YARN_BETA_FAST = "{arch}.rope.scaling.yarn_beta_fast" SCALING_YARN_BETA_SLOW = "{arch}.rope.scaling.yarn_beta_slow" + class Activation: + SITU_BETA = "{arch}.activation.situ_beta" + SITU_LINEAR_BETA = "{arch}.activation.situ_linear_beta" + + class AttnRes: + BLOCK_SIZE = "{arch}.attn_res.block_size" + class Split: LLM_KV_SPLIT_NO = "split.no" LLM_KV_SPLIT_COUNT = "split.count" @@ -252,7 +261,9 @@ class Keys: DT_B_C_RMS = "{arch}.ssm.dt_b_c_rms" class KDA: - HEAD_DIM = "{arch}.kda.head_dim" + HEAD_DIM = "{arch}.kda.head_dim" + SAFE_GATE = "{arch}.kda.safe_gate" + GATE_LOWER_BOUND = "{arch}.kda.gate_lower_bound" class WKV: HEAD_SIZE = "{arch}.wkv.head_size" @@ -543,6 +554,7 @@ class MODEL_ARCH(IntEnum): PLM = auto() BAILINGMOE = auto() BAILINGMOE2 = auto() + BAILINGMOE3 = auto() DOTS1 = auto() ARCEE = auto() AFMOE = auto() @@ -580,6 +592,7 @@ class MODEL_ARCH(IntEnum): LLAMA_EMBED = auto() MAINCODER = auto() KIMI_LINEAR = auto() + KIMI_K3 = auto() TALKIE = auto() MELLUM = auto() NANBEIGE = auto() @@ -698,6 +711,13 @@ class MODEL_TENSOR(IntEnum): SSM_BETA = auto() # Kimi Linear qwen3.5 SSM_G_A = auto() # Kimi Linear SSM_G_B = auto() # Kimi Linear + SSM_G = auto() # Kimi K3 (full-rank KDA gate, replaces SSM_G_A/SSM_G_B) + ATTN_RES_SCORE = auto() # Kimi K3 (fused res_norm * res_proj, pre-attention) + FFN_RES_SCORE = auto() # Kimi K3 (fused res_norm * res_proj, pre-FFN) + OUTPUT_RES_SCORE = auto() # Kimi K3 (fused res_norm * res_proj, final) + FFN_ROUTED_DOWN = auto() # Kimi K3 (latent MoE: hidden -> latent) + FFN_ROUTED_UP = auto() # Kimi K3 (latent MoE: latent -> hidden) + FFN_ROUTED_NORM = auto() # Kimi K3 (latent MoE: norm on expert output) TIME_MIX_W0 = auto() TIME_MIX_W1 = auto() TIME_MIX_W2 = auto() @@ -1250,6 +1270,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.PLM: "plm", MODEL_ARCH.BAILINGMOE: "bailingmoe", MODEL_ARCH.BAILINGMOE2: "bailingmoe2", + MODEL_ARCH.BAILINGMOE3: "bailingmoe3", MODEL_ARCH.DOTS1: "dots1", MODEL_ARCH.ARCEE: "arcee", MODEL_ARCH.AFMOE: "afmoe", @@ -1288,6 +1309,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.LLAMA_EMBED: "llama-embed", MODEL_ARCH.MAINCODER: "maincoder", MODEL_ARCH.KIMI_LINEAR: "kimi-linear", + MODEL_ARCH.KIMI_K3: "kimi-k3", MODEL_ARCH.TALKIE: "talkie", MODEL_ARCH.MELLUM: "mellum", MODEL_ARCH.NANBEIGE: "nanbeige", @@ -1404,6 +1426,13 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.SSM_BETA: "blk.{bid}.ssm_beta", # Kimi Linear qwen3.5 MODEL_TENSOR.SSM_G_A: "blk.{bid}.ssm_g_a", # Kimi Linear MODEL_TENSOR.SSM_G_B: "blk.{bid}.ssm_g_b", # Kimi Linear + MODEL_TENSOR.SSM_G: "blk.{bid}.ssm_g", # Kimi K3 + MODEL_TENSOR.ATTN_RES_SCORE: "blk.{bid}.attn_res_score", # Kimi K3 + MODEL_TENSOR.FFN_RES_SCORE: "blk.{bid}.ffn_res_score", # Kimi K3 + MODEL_TENSOR.OUTPUT_RES_SCORE: "output_res_score", # Kimi K3 + MODEL_TENSOR.FFN_ROUTED_DOWN: "blk.{bid}.ffn_routed_down", # Kimi K3 + MODEL_TENSOR.FFN_ROUTED_UP: "blk.{bid}.ffn_routed_up", # Kimi K3 + MODEL_TENSOR.FFN_ROUTED_NORM: "blk.{bid}.ffn_routed_norm", # Kimi K3 MODEL_TENSOR.TIME_MIX_W0: "blk.{bid}.time_mix_w0", MODEL_TENSOR.TIME_MIX_W1: "blk.{bid}.time_mix_w1", MODEL_TENSOR.TIME_MIX_W2: "blk.{bid}.time_mix_w2", @@ -4209,6 +4238,50 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, MODEL_TENSOR.LAYER_OUT_NORM, ], + MODEL_ARCH.BAILINGMOE3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.FFN_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.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.SSM_CONV1D_Q, + MODEL_TENSOR.SSM_CONV1D_K, + MODEL_TENSOR.SSM_CONV1D_V, + MODEL_TENSOR.SSM_F_A, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_G_A, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.LAYER_OUT_NORM, + ], MODEL_ARCH.DOTS1: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, @@ -4747,6 +4820,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { ], MODEL_ARCH.DFLASH: [ MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.ATTN_NORM, MODEL_TENSOR.ATTN_Q, @@ -4786,6 +4860,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_UP_SHEXP, MODEL_TENSOR.FC, MODEL_TENSOR.ENC_OUTPUT_NORM, + MODEL_TENSOR.D2T, # optional DSpark heads MODEL_TENSOR.DSPARK_MARKOV_W1, MODEL_TENSOR.DSPARK_MARKOV_W2, @@ -4960,6 +5035,56 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_DOWN_SHEXP, MODEL_TENSOR.FFN_UP_SHEXP, ], + MODEL_ARCH.KIMI_K3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.OUTPUT_RES_SCORE, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_RES_SCORE, + MODEL_TENSOR.FFN_RES_SCORE, + # MLA (full-attention layers) + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + # KDA (linear-attention layers) + MODEL_TENSOR.SSM_CONV1D_Q, + MODEL_TENSOR.SSM_CONV1D_K, + MODEL_TENSOR.SSM_CONV1D_V, + MODEL_TENSOR.SSM_F_A, + MODEL_TENSOR.SSM_F_B, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_G, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + # FFN + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_ROUTED_DOWN, + MODEL_TENSOR.FFN_ROUTED_UP, + MODEL_TENSOR.FFN_ROUTED_NORM, + ], MODEL_ARCH.TALKIE: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT, @@ -5412,7 +5537,9 @@ KEY_SSM_GROUP_COUNT = Keys.SSM.GROUP_COUNT KEY_SSM_DT_B_C_RMS = Keys.SSM.DT_B_C_RMS # KDA -KEY_KDA_HEAD_DIM = Keys.KDA.HEAD_DIM +KEY_KDA_HEAD_DIM = Keys.KDA.HEAD_DIM +KEY_KDA_SAFE_GATE = Keys.KDA.SAFE_GATE +KEY_KDA_GATE_LOWER_BOUND = Keys.KDA.GATE_LOWER_BOUND # tokenization KEY_TOKENIZER_MODEL = Keys.Tokenizer.MODEL diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 05f86396d..9e0914fd8 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -981,6 +981,9 @@ class GGUFWriter: def add_block_size(self, value: int) -> None: self.add_uint32(Keys.LLM.BLOCK_SIZE.format(arch=self.arch), value) + def add_sample_from_anchor(self, value: bool) -> None: + self.add_bool(Keys.LLM.SAMPLE_FROM_ANCHOR.format(arch=self.arch), value) + def add_target_layers(self, value: Sequence[int]) -> None: self.add_array(Keys.LLM.TARGET_LAYERS.format(arch=self.arch), value) @@ -1103,9 +1106,27 @@ class GGUFWriter: def add_ssm_dt_b_c_rms(self, value: bool) -> None: self.add_bool(Keys.SSM.DT_B_C_RMS.format(arch=self.arch), value) + def add_expert_latent_length(self, value: int) -> None: + self.add_uint32(Keys.LLM.EXPERT_LATENT_LENGTH.format(arch=self.arch), value) + + def add_activation_situ_beta(self, value: float) -> None: + self.add_float32(Keys.Activation.SITU_BETA.format(arch=self.arch), value) + + def add_activation_situ_linear_beta(self, value: float) -> None: + self.add_float32(Keys.Activation.SITU_LINEAR_BETA.format(arch=self.arch), value) + + def add_attn_res_block_size(self, value: int) -> None: + self.add_uint32(Keys.AttnRes.BLOCK_SIZE.format(arch=self.arch), value) + def add_kda_head_dim(self, value: int) -> None: self.add_uint32(Keys.KDA.HEAD_DIM.format(arch=self.arch), value) + def add_kda_safe_gate(self, value: bool) -> None: + self.add_bool(Keys.KDA.SAFE_GATE.format(arch=self.arch), value) + + def add_kda_gate_lower_bound(self, value: float) -> None: + self.add_float32(Keys.KDA.GATE_LOWER_BOUND.format(arch=self.arch), value) + def add_tokenizer_model(self, model: str) -> None: self.add_string(Keys.Tokenizer.MODEL, model) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index ee3a8d3aa..3292942b4 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -76,14 +76,14 @@ class TensorNameMap: # Output MODEL_TENSOR.OUTPUT: ( "embed_out", # gptneox - "lm_head", # gpt2 mpt falcon llama-hf baichuan qwen mamba dbrx jais nemotron exaone olmoe olmo2 phimoe plamo2 + "lm_head", # gpt2 mpt falcon llama-hf baichuan qwen mamba dbrx jais nemotron exaone olmoe olmo2 phimoe plamo2 llama4 "output", # llama-pth bloom internlm2 "word_embeddings_for_head", # persimmon "lm_head.linear", # phi2 "output_layer", # chatglm "head", # rwkv "head.out", # wavtokenizer - "lm_head", # llama4 + "model.lm_head", # dflash "model.transformer.ff_out", # llada "head.decoder", # modern-bert ), @@ -255,6 +255,7 @@ class TensorNameMap: # Attention query MODEL_TENSOR.ATTN_Q: ( "model.layers.{bid}.self_attn.q_proj", # llama-hf nemotron olmoe olmo2 phimoe + "model.layers.{bid}.attention.q_proj", # bailingmoe3 "layers.{bid}.self_attn.q_proj", # embeddinggemma "model.layers.{bid}.self_attn.q_proj_no_perm", # llama-custom "layers.{bid}.attention.wq", # llama-pth @@ -275,6 +276,7 @@ class TensorNameMap: # Attention key MODEL_TENSOR.ATTN_K: ( "model.layers.{bid}.self_attn.k_proj", # llama-hf nemotron olmoe olmo2 phimoe + "model.layers.{bid}.attention.k_proj", # bailingmoe3 "layers.{bid}.self_attn.k_proj", # embeddinggemma "model.layers.{bid}.self_attn.k_proj_no_perm", # llama-custom "layers.{bid}.attention.wk", # llama-pth @@ -296,6 +298,7 @@ class TensorNameMap: # Attention value MODEL_TENSOR.ATTN_V: ( "model.layers.{bid}.self_attn.v_proj", # llama-hf nemotron olmoe olmo2 phimoe + "model.layers.{bid}.attention.v_proj", # bailingmoe3 "layers.{bid}.self_attn.v_proj", # embeddinggemma "layers.{bid}.attention.wv", # llama-pth "encoder.layer.{bid}.attention.self.value", # bert @@ -321,6 +324,8 @@ class TensorNameMap: "transformer.h.{bid}.self_attention.dense", # falcon "h.{bid}.self_attention.dense", # bloom "model.layers.{bid}.self_attn.o_proj", # llama-hf nemotron olmoe olmo2 phimoe + "model.layers.{bid}.attention.o_proj", # bailingmoe3 + "model.layers.{bid}.attention.dense", # bailingmoe3 MLA "layers.{bid}.self_attn.o_proj", # embeddinggemma "model.layers.{bid}.self_attn.out_proj", # lfm2 minimax-01 "model.layers.{bid}.self_attn.linear_attn", # deci @@ -834,6 +839,7 @@ class TensorNameMap: "model.layers.{bid}.linear_attn.dt_proj", # qwen3next "backbone.layers.{bid}.mixer.dt", # nemotron-h-moe "model.layers.{bid}.self_attn.dt_proj", # kimi + "model.layers.{bid}.attention.dt_proj", # bailingmoe3 ), MODEL_TENSOR.SSM_DT_NORM: ( @@ -848,6 +854,7 @@ class TensorNameMap: "model.layers.layers.{bid}.mixer.A_log", # plamo2 "model.layers.{bid}.linear_attn.A_log", # qwen3next "model.layers.{bid}.self_attn.A_log", # kimi + "model.layers.{bid}.attention.A_log", # bailingmoe3 ), MODEL_TENSOR.SSM_B_NORM: ( @@ -874,6 +881,7 @@ class TensorNameMap: "model.layers.{bid}.linear_attn.norm", # qwen3next "backbone.layers.{bid}.mixer.norm", # mamba2 "model.layers.{bid}.self_attn.o_norm", # kimi + "model.layers.{bid}.attention.o_norm", # bailingmoe3 ), MODEL_TENSOR.SSM_OUT: ( @@ -895,12 +903,15 @@ class TensorNameMap: # Kimi Linear KDA (using SSM_ prefix for consistency) MODEL_TENSOR.SSM_CONV1D_Q: ( "model.layers.{bid}.self_attn.q_conv1d", + "model.layers.{bid}.attention.q_conv1d", ), MODEL_TENSOR.SSM_CONV1D_K: ( "model.layers.{bid}.self_attn.k_conv1d", + "model.layers.{bid}.attention.k_conv1d", ), MODEL_TENSOR.SSM_CONV1D_V: ( "model.layers.{bid}.self_attn.v_conv1d", + "model.layers.{bid}.attention.v_conv1d", ), MODEL_TENSOR.SSM_F_A: ( "model.layers.{bid}.self_attn.f_a_proj", @@ -911,7 +922,21 @@ class TensorNameMap: MODEL_TENSOR.SSM_BETA: ( "model.layers.{bid}.linear_attn.in_proj_b", # qwen3.5 "model.layers.{bid}.self_attn.b_proj", # Kimi Linear + "model.layers.{bid}.attention.b_proj", # bailingmoe3 ), + # Kimi K3 latent MoE: routed experts operate in a down-projected space + MODEL_TENSOR.FFN_ROUTED_DOWN: ( + "model.layers.{bid}.block_sparse_moe.routed_expert_down_proj", + ), + + MODEL_TENSOR.FFN_ROUTED_UP: ( + "model.layers.{bid}.block_sparse_moe.routed_expert_up_proj", + ), + + MODEL_TENSOR.FFN_ROUTED_NORM: ( + "model.layers.{bid}.block_sparse_moe.routed_expert_norm", + ), + MODEL_TENSOR.SSM_G_A: ( "model.layers.{bid}.self_attn.g_a_proj", ), @@ -1090,40 +1115,48 @@ class TensorNameMap: MODEL_TENSOR.ATTN_Q_A: ( "model.layers.{bid}.self_attn.q_a_proj", # deepseek2 + "model.layers.{bid}.attention.q_a_proj", # bailingmoe3 (Ling-3.0-tiny) "layers.{bid}.attention.wq_a", # mistral-large ), MODEL_TENSOR.ATTN_Q_B: ( "model.layers.{bid}.self_attn.q_b_proj", # deepseek2 + "model.layers.{bid}.attention.q_b_proj", # bailingmoe3 (Ling-3.0-tiny) "layers.{bid}.attention.wq_b", # mistral-large ), MODEL_TENSOR.ATTN_KV_A_MQA: ( "model.layers.{bid}.self_attn.kv_a_proj_with_mqa", # deepseek2 + "model.layers.{bid}.attention.kv_a_proj_with_mqa", # bailingmoe3 "layers.{bid}.attention.wkv_a_with_mqa", # mistral-large ), MODEL_TENSOR.ATTN_KV_B: ( "model.layers.{bid}.self_attn.kv_b_proj", # deepseek2 + "model.layers.{bid}.attention.kv_b_proj", # bailingmoe3 ), MODEL_TENSOR.ATTN_K_B: ( "model.layers.{bid}.self_attn.k_b_proj", # deepseek2 + "model.layers.{bid}.attention.k_b_proj", # bailingmoe3 "layers.{bid}.attention.k_b_proj", # mistral-large ), MODEL_TENSOR.ATTN_V_B: ( "model.layers.{bid}.self_attn.v_b_proj", # deepseek2 + "model.layers.{bid}.attention.v_b_proj", # bailingmoe3 "layers.{bid}.attention.v_b_proj", # mistral-large ), MODEL_TENSOR.ATTN_Q_A_NORM: ( "model.layers.{bid}.self_attn.q_a_layernorm", # deepseek2 + "model.layers.{bid}.attention.q_a_layernorm", # bailingmoe3 (Ling-3.0-tiny) "layers.{bid}.attention.q_a_norm", # mistral-large ), MODEL_TENSOR.ATTN_KV_A_NORM: ( "model.layers.{bid}.self_attn.kv_a_layernorm", # deepseek2 + "model.layers.{bid}.attention.kv_a_layernorm", # bailingmoe3 "layers.{bid}.attention.kv_a_norm", # mistral-large ), diff --git a/src/llama-adapter.cpp b/src/llama-adapter.cpp index 3e0fe66af..e6678a66d 100644 --- a/src/llama-adapter.cpp +++ b/src/llama-adapter.cpp @@ -396,8 +396,11 @@ static void llama_adapter_lora_init_impl(llama_model & model, const char * path_ llama_file gguf_file(path_lora, "rb"); std::vector read_buf; auto set_tensor = [&](ggml_tensor * orig, ggml_tensor * dev) { - size_t offs = gguf_get_data_offset(ctx_gguf.get()) + gguf_get_tensor_offset(ctx_gguf.get(), gguf_find_tensor(ctx_gguf.get(), orig->name)); - size_t size = ggml_nbytes(orig); + const size_t offs = gguf_get_data_offset(ctx_gguf.get()) + gguf_get_tensor_offset(ctx_gguf.get(), gguf_find_tensor(ctx_gguf.get(), orig->name)); + const size_t size = ggml_nbytes(orig); + if (offs + size < offs || offs + size > gguf_file.size()) { + throw std::runtime_error(format("LoRA tensor '%s' data is not within the file bounds, file is corrupted or incomplete", orig->name)); + } read_buf.resize(size); gguf_file.seek(offs, SEEK_SET); gguf_file.read_raw(read_buf.data(), size); diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 1cd9b758f..5b88bde14 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -107,6 +107,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_PLM, "plm" }, { LLM_ARCH_BAILINGMOE, "bailingmoe" }, { LLM_ARCH_BAILINGMOE2, "bailingmoe2" }, + { LLM_ARCH_BAILINGMOE3, "bailingmoe3" }, { LLM_ARCH_DOTS1, "dots1" }, { LLM_ARCH_ARCEE, "arcee" }, { LLM_ARCH_AFMOE, "afmoe" }, @@ -144,6 +145,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_LLAMA_EMBED, "llama-embed" }, { LLM_ARCH_MAINCODER, "maincoder" }, { LLM_ARCH_KIMI_LINEAR, "kimi-linear" }, + { LLM_ARCH_KIMI_K3, "kimi-k3" }, { LLM_ARCH_TALKIE, "talkie" }, { LLM_ARCH_MELLUM, "mellum" }, { LLM_ARCH_NANBEIGE, "nanbeige" }, @@ -187,6 +189,9 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_FEATURES_LENGTH, "%s.features_length" }, { LLM_KV_BLOCK_COUNT, "%s.block_count" }, { LLM_KV_LEADING_DENSE_BLOCK_COUNT, "%s.leading_dense_block_count" }, + { LLM_KV_ATTN_RES_BLOCK_SIZE, "%s.attn_res.block_size" }, + { LLM_KV_ACTIVATION_SITU_BETA, "%s.activation.situ_beta" }, + { LLM_KV_ACTIVATION_SITU_LINEAR_BETA, "%s.activation.situ_linear_beta" }, { LLM_KV_FEED_FORWARD_LENGTH, "%s.feed_forward_length" }, { LLM_KV_EXPERT_FEED_FORWARD_LENGTH, "%s.expert_feed_forward_length" }, { LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, "%s.expert_shared_feed_forward_length" }, @@ -202,6 +207,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_EXPERT_GROUP_USED_COUNT, "%s.expert_group_used_count" }, { LLM_KV_EXPERT_WEIGHTS_SCALE, "%s.expert_weights_scale" }, { LLM_KV_EXPERT_WEIGHTS_NORM, "%s.expert_weights_norm" }, + { LLM_KV_EXPERT_LATENT_LENGTH, "%s.expert_latent_length" }, { LLM_KV_EXPERT_GATING_FUNC, "%s.expert_gating_func" }, { LLM_KV_EXPERT_GROUP_SCALE, "%s.expert_group_scale" }, { LLM_KV_EXPERTS_PER_GROUP, "%s.experts_per_group" }, @@ -312,7 +318,9 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_SSM_GROUP_COUNT, "%s.ssm.group_count" }, { LLM_KV_SSM_DT_B_C_RMS, "%s.ssm.dt_b_c_rms" }, - { LLM_KV_KDA_HEAD_DIM, "%s.kda.head_dim" }, + { LLM_KV_KDA_HEAD_DIM, "%s.kda.head_dim" }, + { LLM_KV_KDA_SAFE_GATE, "%s.kda.safe_gate" }, + { LLM_KV_KDA_GATE_LOWER_BOUND, "%s.kda.gate_lower_bound" }, { LLM_KV_WKV_HEAD_SIZE, "%s.wkv.head_size" }, @@ -463,6 +471,13 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_SSM_F_B, "blk.%d.ssm_f_b" }, { LLM_TENSOR_SSM_BETA, "blk.%d.ssm_beta" }, { LLM_TENSOR_SSM_G_A, "blk.%d.ssm_g_a" }, + { LLM_TENSOR_SSM_G, "blk.%d.ssm_g" }, + { LLM_TENSOR_ATTN_RES_SCORE, "blk.%d.attn_res_score" }, + { LLM_TENSOR_FFN_RES_SCORE, "blk.%d.ffn_res_score" }, + { LLM_TENSOR_OUTPUT_RES_SCORE, "output_res_score" }, + { LLM_TENSOR_FFN_ROUTED_DOWN, "blk.%d.ffn_routed_down" }, + { LLM_TENSOR_FFN_ROUTED_UP, "blk.%d.ffn_routed_up" }, + { LLM_TENSOR_FFN_ROUTED_NORM, "blk.%d.ffn_routed_norm" }, { LLM_TENSOR_SSM_G_B, "blk.%d.ssm_g_b" }, { LLM_TENSOR_SSM_NORM, "blk.%d.ssm_norm" }, { LLM_TENSOR_ATTN_Q_A_NORM, "blk.%d.attn_q_a_norm" }, @@ -756,6 +771,13 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_SSM_F_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_SSM_BETA, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_SSM_G_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_SSM_G, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_RES_SCORE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_FFN_RES_SCORE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_OUTPUT_RES_SCORE, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, + {LLM_TENSOR_FFN_ROUTED_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_FFN_ROUTED_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_FFN_ROUTED_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_SSM_G_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_TIME_MIX_LERP_X, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_TIME_MIX_LN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, @@ -976,6 +998,8 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_QWEN3NEXT: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_BAILINGMOE3: + case LLM_ARCH_KIMI_K3: case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: case LLM_ARCH_DEEPSEEK4: @@ -1040,6 +1064,8 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_MISTRAL4: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_BAILINGMOE3: + case LLM_ARCH_KIMI_K3: case LLM_ARCH_QWEN3TTS: return false; default: diff --git a/src/llama-arch.h b/src/llama-arch.h index 24252b9da..8042120a2 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -112,6 +112,7 @@ enum llm_arch { LLM_ARCH_PLM, LLM_ARCH_BAILINGMOE, LLM_ARCH_BAILINGMOE2, + LLM_ARCH_BAILINGMOE3, LLM_ARCH_DOTS1, LLM_ARCH_ARCEE, LLM_ARCH_AFMOE, @@ -145,6 +146,7 @@ enum llm_arch { LLM_ARCH_LLAMA_EMBED, LLM_ARCH_MAINCODER, LLM_ARCH_KIMI_LINEAR, + LLM_ARCH_KIMI_K3, LLM_ARCH_TALKIE, LLM_ARCH_MELLUM, LLM_ARCH_EAGLE3, @@ -192,6 +194,9 @@ enum llm_kv { LLM_KV_FEATURES_LENGTH, LLM_KV_BLOCK_COUNT, LLM_KV_LEADING_DENSE_BLOCK_COUNT, + LLM_KV_ATTN_RES_BLOCK_SIZE, + LLM_KV_ACTIVATION_SITU_BETA, + LLM_KV_ACTIVATION_SITU_LINEAR_BETA, LLM_KV_FEED_FORWARD_LENGTH, LLM_KV_EXPERT_FEED_FORWARD_LENGTH, LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, @@ -207,6 +212,7 @@ enum llm_kv { LLM_KV_EXPERT_GROUP_USED_COUNT, LLM_KV_EXPERT_WEIGHTS_SCALE, LLM_KV_EXPERT_WEIGHTS_NORM, + LLM_KV_EXPERT_LATENT_LENGTH, LLM_KV_EXPERT_GATING_FUNC, LLM_KV_EXPERT_GROUP_SCALE, LLM_KV_EXPERTS_PER_GROUP, @@ -318,6 +324,8 @@ enum llm_kv { LLM_KV_SSM_DT_B_C_RMS, LLM_KV_KDA_HEAD_DIM, + LLM_KV_KDA_SAFE_GATE, + LLM_KV_KDA_GATE_LOWER_BOUND, LLM_KV_WKV_HEAD_SIZE, @@ -492,6 +500,13 @@ enum llm_tensor { LLM_TENSOR_SSM_BETA, // kimi: beta mixing coefficient and qwen3.5 LLM_TENSOR_SSM_G_A, // kimi: output gate projection A LLM_TENSOR_SSM_G_B, // kimi: output gate projection B + LLM_TENSOR_SSM_G, // kimi-k3: full-rank KDA gate + LLM_TENSOR_ATTN_RES_SCORE, // kimi-k3: fused res_norm*res_proj (pre-attn) + LLM_TENSOR_FFN_RES_SCORE, // kimi-k3: fused res_norm*res_proj (pre-ffn) + LLM_TENSOR_OUTPUT_RES_SCORE, // kimi-k3: fused res_norm*res_proj (final) + LLM_TENSOR_FFN_ROUTED_DOWN, // kimi-k3: latent MoE down + LLM_TENSOR_FFN_ROUTED_UP, // kimi-k3: latent MoE up + LLM_TENSOR_FFN_ROUTED_NORM, // kimi-k3: latent MoE norm LLM_TENSOR_TIME_MIX_W0, LLM_TENSOR_TIME_MIX_W1, LLM_TENSOR_TIME_MIX_W2, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index cebdd9058..b22623c34 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2303,8 +2303,12 @@ void llama_context::output_reorder() { uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { uint32_t res; - if (model.arch == LLM_ARCH_QWEN3NEXT || + if (model.arch == LLM_ARCH_KIMI_K3) { + // the n_tokens*40 budget below is exhausted at ubatch 3840 + res = std::max(n_tokens * 160, 64u * model.n_tensors()); + } else if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || + model.arch == LLM_ARCH_BAILINGMOE3 || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_DEEPSEEK4 || diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 55f1fff33..ba10f0aa4 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1836,6 +1836,8 @@ ggml_tensor * llm_graph_context::build_ffn( cur = ggml_reglu(ctx0, cur); cb(cur, "ffn_reglu", il); } break; + case LLM_FFN_SITU: + GGML_ABORT("not yet supported"); default: GGML_ABORT("fatal error"); } @@ -2175,6 +2177,21 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cur = ggml_silu(ctx0, cur); cb(cur, "ffn_moe_silu", il); } break; + case LLM_FFN_SITU: + { + // situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * lb*tanh(up/lb) + GGML_ASSERT(has_gate); + const float beta = hparams.situ_beta; + const float lb = hparams.situ_linear_beta; + + ggml_tensor * act = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, cur, 1.0f/beta)), beta); + act = ggml_mul(ctx0, act, ggml_sigmoid(ctx0, cur)); + if (lb > 0.0f) { + up = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, up, 1.0f/lb)), lb); + } + cur = ggml_mul(ctx0, act, up); + cb(cur, "ffn_moe_situ", il); + } break; case LLM_FFN_GELU: if (has_gate) { cur = ggml_geglu_split(ctx0, cur, up); diff --git a/src/llama-graph.h b/src/llama-graph.h index 75bc0fe80..94324c745 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -59,6 +59,7 @@ enum llm_ffn_op_type : int { LLM_FFN_GEGLU, LLM_FFN_REGLU, LLM_FFN_SWIGLU_OAI_MOE, + LLM_FFN_SITU, // kimi-k3 }; enum llm_ffn_gate_type { diff --git a/src/llama-hparams.h b/src/llama-hparams.h index d9ac112b7..e91ce1cc3 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -4,10 +4,11 @@ #include #include +#include // bump if necessary #define LLAMA_MAX_LAYERS 512 -#define LLAMA_MAX_EXPERTS 512 // Qwen3 Next +#define LLAMA_MAX_EXPERTS 1024 // Kimi K3 enum llama_expert_gating_func_type { LLAMA_EXPERT_GATING_FUNC_TYPE_NONE = 0, @@ -169,6 +170,14 @@ struct llama_hparams { // for Kimi Linear KDA uint32_t n_embd_head_kda = 0; + bool kda_safe_gate = false; + + // kimi-k3 + uint32_t n_expert_latent = 0; // routed_expert_hidden_size (0 = experts run at n_embd) + uint32_t attn_res_block_size = 0; // 0 = no cross-layer attention residuals + float kda_gate_lower_bound = -INFINITY; + float situ_beta = 1.0f; + float situ_linear_beta = 0.0f; // 0 = no linear-beta transform on the up branch bool ssm_dt_b_c_rms = false; diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index abca773a9..be9524d40 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -121,6 +121,7 @@ void llama_model_saver::add_kv(const enum llm_kv key, const Container & value, c } // instantiate for external usage: template void llama_model_saver::add_kv>(const enum llm_kv, const std::vector &, const bool); +template void llama_model_saver::add_kv>(const enum llm_kv, const std::vector &, const bool); void llama_model_saver::add_kv(const enum llm_kv key, const std::vector & value) { std::vector tmp(value.size()); @@ -213,10 +214,13 @@ void llama_model_saver::add_kv_from_model() { 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_LATENT_LENGTH, hparams.n_expert_latent); 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, hparams.swiglu_clamp_exp); - add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp); + add_kv(LLM_KV_SWIGLU_CLAMP_EXP, std::vector( + hparams.swiglu_clamp_exp.begin(), hparams.swiglu_clamp_exp.begin() + hparams.n_layer_all)); + add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, std::vector( + hparams.swiglu_clamp_shexp.begin(), hparams.swiglu_clamp_shexp.begin() + hparams.n_layer_all)); add_kv(LLM_KV_USE_PARALLEL_RESIDUAL, hparams.use_par_res); // add_kv(LLM_KV_TENSOR_DATA_LAYOUT, ???); add_kv(LLM_KV_EXPERT_COUNT, hparams.n_expert); @@ -319,6 +323,8 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_SSM_DT_B_C_RMS, hparams.ssm_dt_b_c_rms); add_kv(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda); + add_kv(LLM_KV_KDA_SAFE_GATE, hparams.kda_safe_gate); + add_kv(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound); add_kv(LLM_KV_WKV_HEAD_SIZE, hparams.wkv_head_size); @@ -376,6 +382,10 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_XIELU_BETA, hparams.xielu_beta); add_kv(LLM_KV_XIELU_EPS, hparams.xielu_eps); + add_kv(LLM_KV_ATTN_RES_BLOCK_SIZE, hparams.attn_res_block_size); + add_kv(LLM_KV_ACTIVATION_SITU_BETA, hparams.situ_beta); + add_kv(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, hparams.situ_linear_beta); + // deprecated // add_kv(LLM_KV_TOKENIZER_PREFIX_ID, ???); // add_kv(LLM_KV_TOKENIZER_SUFFIX_ID, ???); @@ -403,6 +413,7 @@ void llama_model_saver::add_tensors_from_model() { add_tensor(model->output_norm_enc); add_tensor(model->output_s); add_tensor(model->output_in_s); + add_tensor(model->output_res_score); add_tensor(model->cls); add_tensor(model->cls_b); add_tensor(model->cls_out); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 587300757..e35a683d1 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -47,6 +47,7 @@ #include "models/baichuan.cpp" #include "models/bailingmoe.cpp" #include "models/bailingmoe2.cpp" +#include "models/bailingmoe3.cpp" #include "models/bert.cpp" #include "models/bitnet.cpp" #include "models/bloom.cpp" @@ -105,6 +106,7 @@ #include "models/jamba.cpp" #include "models/jina-bert-v2.cpp" #include "models/jina-bert-v3.cpp" +#include "models/kimi-k3.cpp" #include "models/kimi-linear.cpp" #include "models/laguna.cpp" #include "models/lfm2.cpp" @@ -403,6 +405,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_bailingmoe(params); case LLM_ARCH_BAILINGMOE2: return new llama_model_bailingmoe2(params); + case LLM_ARCH_BAILINGMOE3: + return new llama_model_bailingmoe3(params); case LLM_ARCH_SEED_OSS: return new llama_model_seed_oss(params); case LLM_ARCH_DOTS1: @@ -469,6 +473,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_mimo2(params); case LLM_ARCH_KIMI_LINEAR: return new llama_model_kimi_linear(params); + case LLM_ARCH_KIMI_K3: + return new llama_model_kimi_k3(params); case LLM_ARCH_STEP35: return new llama_model_step35(params); default: @@ -966,6 +972,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_A13B: return "A13B"; case LLM_TYPE_7B_A1B: return "7B.A1B"; case LLM_TYPE_8B_A1B: return "8B.A1B"; + case LLM_TYPE_7_9B_A1_3B: return "7.9B.A1.3B"; case LLM_TYPE_12B_A2_5B: return "12B.A2.5B"; case LLM_TYPE_16B_A1B: return "16B.A1B"; case LLM_TYPE_21B_A3B: return "21B.A3B"; @@ -982,6 +989,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_118B_A8B: return "118B.A8B"; case LLM_TYPE_120B_A12B: return "120B.A12B"; case LLM_TYPE_122B_A10B: return "122B.A10B"; + case LLM_TYPE_124B_A5_1B: return "124B.A5.1B"; case LLM_TYPE_196B_A11B: return "196B.A11B"; case LLM_TYPE_230B_A10B: return "230B.A10B"; case LLM_TYPE_428B_A23B: return "428B.A23B"; @@ -992,6 +1000,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_397B_A17B: return "397B.A17B"; case LLM_TYPE_685B_A37B: return "685B.A37B"; case LLM_TYPE_744B_A40B: return "744B.A40B"; + case LLM_TYPE_2_8T_A50B: return "2.8T.A50B"; case LLM_TYPE_E2B: return "E2B"; case LLM_TYPE_E4B: return "E4B"; default: return "?B"; @@ -2104,7 +2113,7 @@ void llama_model::print_info() const { LLAMA_LOG_INFO("%s: expert_weights_norm = %d\n", __func__, hparams.expert_weights_norm); } - if (arch == LLM_ARCH_BAILINGMOE2) { + if (arch == LLM_ARCH_BAILINGMOE2 || arch == LLM_ARCH_BAILINGMOE3) { LLAMA_LOG_INFO("%s: n_layer_dense_lead = %d\n", __func__, hparams.n_layer_dense_lead); LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp); LLAMA_LOG_INFO("%s: n_ff_shexp = %d\n", __func__, hparams.n_ff_shexp); @@ -2399,11 +2408,11 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // checks default: { - // The MTP head is dense-attention only on hybrid Qwen3-Next/3.5/3.6, so use a plain - // attention KV cache for the MTP context instead of the hybrid wrapper. + // Dense MTP heads use a plain attention KV cache instead of the hybrid wrapper. const bool mtp_on_hybrid_qwen = params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && - (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE); + (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || + arch == LLM_ARCH_BAILINGMOE3); const bool mtp_on_hybrid_nemotron = params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && arch == LLM_ARCH_NEMOTRON_H_MOE; @@ -2749,6 +2758,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_NEMOTRON_H: case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_KIMI_K3: return LLAMA_ROPE_TYPE_NONE; // use what we call a normal RoPE, operating on pairs of consecutive head values @@ -2780,6 +2790,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_GRANITE_SWITCH: case LLM_ARCH_CHAMELEON: case LLM_ARCH_BAILINGMOE: + case LLM_ARCH_BAILINGMOE3: case LLM_ARCH_NEO_BERT: case LLM_ARCH_SMOLLM3: case LLM_ARCH_ARCEE: diff --git a/src/llama-model.h b/src/llama-model.h index 510ed8498..4412ef08e 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -118,6 +118,7 @@ enum llm_type { LLM_TYPE_A13B, LLM_TYPE_7B_A1B, LLM_TYPE_8B_A1B, // lfm2moe + LLM_TYPE_7_9B_A1_3B, // Ling-3.0-tiny LLM_TYPE_12B_A2_5B, LLM_TYPE_16B_A1B, LLM_TYPE_21B_A3B, // Ernie MoE small @@ -134,6 +135,7 @@ enum llm_type { LLM_TYPE_118B_A8B, // Laguna-S-2 LLM_TYPE_120B_A12B, // Nemotron 3 Super LLM_TYPE_122B_A10B, // Qwen3.5 + LLM_TYPE_124B_A5_1B, // Ling-3.0-flash LLM_TYPE_196B_A11B, // Step3.5-Flash LLM_TYPE_230B_A10B, // Minimax M2 LLM_TYPE_428B_A23B, // Minimax M3 @@ -144,6 +146,7 @@ enum llm_type { LLM_TYPE_397B_A17B, // Qwen3.5 LLM_TYPE_685B_A37B, // DeepSeek V3.2 LLM_TYPE_744B_A40B, // GLM-5 + LLM_TYPE_2_8T_A50B, // Kimi-K3 LLM_TYPE_E2B, LLM_TYPE_E4B, }; @@ -530,6 +533,14 @@ struct llama_layer { struct ggml_tensor * ssm_g_b = nullptr; struct ggml_tensor * ssm_o_norm = nullptr; + // kimi-k3 + struct ggml_tensor * ssm_g = nullptr; // full-rank KDA gate (replaces ssm_g_a/ssm_g_b) + struct ggml_tensor * attn_res_score = nullptr; // fused res_norm*res_proj, pre-attention + struct ggml_tensor * ffn_res_score = nullptr; // fused res_norm*res_proj, pre-FFN + struct ggml_tensor * ffn_routed_down = nullptr; // latent MoE: n_embd -> n_expert_latent + struct ggml_tensor * ffn_routed_up = nullptr; // latent MoE: n_expert_latent -> n_embd + struct ggml_tensor * ffn_routed_norm = nullptr; + // DSA (deepseek sparse attention) struct ggml_tensor * indexer_k_norm = nullptr; struct ggml_tensor * indexer_k_norm_b = nullptr; @@ -589,6 +600,7 @@ struct llama_model { struct ggml_tensor * tok_norm_b = nullptr; struct ggml_tensor * output_norm = nullptr; + struct ggml_tensor * output_res_score = nullptr; // kimi-k3: final cross-layer residual mix struct ggml_tensor * output_norm_b = nullptr; struct ggml_tensor * output = nullptr; struct ggml_tensor * output_b = nullptr; diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 5e766d57a..20ba3827f 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -476,7 +476,12 @@ static ggml_type llama_tensor_get_type_impl(quantize_state_impl & qs, ggml_type } else if (ftype == LLAMA_FTYPE_MOSTLY_MXFP4_MOE) { // MoE tensors -> MXFP4 // other tensors -> Q8_0 - if (tensor->ne[2] > 1) { + // MLA projection tensors are also 3D, so match expert tensor roles explicitly. + const bool is_bailingmoe3_expert = arch == LLM_ARCH_BAILINGMOE3 && + (category == tensor_category::FFN_UP || + category == tensor_category::FFN_GATE || + category == tensor_category::FFN_DOWN); + if (tensor->ne[2] > 1 && (arch != LLM_ARCH_BAILINGMOE3 || is_bailingmoe3_expert)) { new_type = GGML_TYPE_MXFP4; } else { new_type = GGML_TYPE_Q8_0; diff --git a/src/models/bailingmoe3.cpp b/src/models/bailingmoe3.cpp new file mode 100644 index 000000000..f5855696e --- /dev/null +++ b/src/models/bailingmoe3.cpp @@ -0,0 +1,532 @@ +#include "models.h" +#include "llama-memory-recurrent.h" + +void llama_model_bailingmoe3::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl); + ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv); + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q, false); + ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); + ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda); + if (!ml.get_key(LLM_KV_KDA_SAFE_GATE, hparams.kda_safe_gate, false)) { + hparams.kda_safe_gate = true; + } + ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound); + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all, false); + ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, false); + + if (hparams.n_ff_shexp == 0) { + hparams.n_ff_shexp = hparams.n_ff_exp * std::max(1u, hparams.n_expert_shared); + } + + GGML_ASSERT(hparams.kda_safe_gate); + GGML_ASSERT(hparams.kda_gate_lower_bound < 0.0f); + + for (uint32_t il = 0; il < hparams.n_layer(); ++il) { + hparams.is_recr_impl[il] = hparams.n_head_kv(il) == 0; + } + + switch (hparams.n_layer()) { + case 24: type = hparams.n_embd == 1536 && hparams.n_expert == 128 ? LLM_TYPE_7_9B_A1_3B : LLM_TYPE_UNKNOWN; break; + case 42: type = hparams.n_embd == 2560 && hparams.n_expert == 512 ? LLM_TYPE_124B_A5_1B : LLM_TYPE_UNKNOWN; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_bailingmoe3::load_arch_tensors(llama_model_loader & ml) { + LLAMA_LOAD_LOCALS; + + 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 }, TENSOR_NOT_REQUIRED); + if (output == nullptr) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED); + } + + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_inner = head_dim * n_head; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t qk_rope_head_dim = hparams.n_rot(); + const int64_t qk_head_dim = hparams.n_embd_head_k_mla(); + const int64_t v_head_dim = hparams.n_embd_head_v_mla(); + + const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr); + const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight"; + const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr); + const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; + int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0; + + if (!ml.load_mtp) { + mtp_flags |= TENSOR_SKIP; + } + + for (int il = 0; il < n_layer; ++il) { + auto & layer = layers[il]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, trunk_flags); + + if (hparams.is_recr(il)) { + layer.ssm_q_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_Q, "weight", il), { d_conv, 1, d_inner, 1 }, trunk_flags); + layer.ssm_k_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_K, "weight", il), { d_conv, 1, d_inner, 1 }, trunk_flags); + layer.ssm_v_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_V, "weight", il), { d_conv, 1, d_inner, 1 }, trunk_flags); + + create_tensor_qkv(layer, il, n_embd, d_inner, d_inner, d_inner, trunk_flags); + layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", il), { n_embd, d_inner }, trunk_flags); + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_head }, trunk_flags); + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, il), { 1, n_head }, trunk_flags); + layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { d_inner }, trunk_flags); + layer.ssm_g_a = create_tensor(tn(LLM_TENSOR_SSM_G_A, "weight", il), { n_embd, d_inner }, trunk_flags); + layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_dim }, trunk_flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { d_inner, n_embd }, trunk_flags); + } else { + if (q_lora_rank > 0) { + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", il), { n_embd, q_lora_rank }, trunk_flags); + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", il), { q_lora_rank }, trunk_flags); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", il), { q_lora_rank, n_head * qk_head_dim }, trunk_flags); + } else { + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", il), { n_embd, n_head * qk_head_dim }, trunk_flags); + } + layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", il), { n_embd, kv_lora_rank + qk_rope_head_dim }, trunk_flags); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", il), { kv_lora_rank }, trunk_flags); + layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", il), { qk_head_dim - qk_rope_head_dim, kv_lora_rank, n_head }, trunk_flags); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", il), { kv_lora_rank, v_head_dim, n_head }, trunk_flags); + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, n_head }, trunk_flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_head * v_head_dim, n_embd }, trunk_flags); + } + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", il), { n_embd }, trunk_flags); + if ((uint32_t) il < hparams.n_layer_dense_lead) { + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", il), { n_embd, n_ff }, trunk_flags); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", il), { n_embd, n_ff }, trunk_flags); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", il), { n_ff, n_embd }, trunk_flags); + } else { + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, trunk_flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", il), { n_expert }, trunk_flags); + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", il), { n_embd, hparams.n_ff_exp, n_expert }, trunk_flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", il), { n_embd, hparams.n_ff_exp, n_expert }, trunk_flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { hparams.n_ff_exp, n_embd, n_expert }, trunk_flags); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, trunk_flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, trunk_flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { hparams.n_ff_shexp, n_embd }, trunk_flags); + } + } + + for (int il = n_layer; il < n_layer_all; ++il) { + auto & layer = layers[il]; + const int flags = mtp_flags; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, flags); + if (q_lora_rank > 0) { + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", il), { n_embd, q_lora_rank }, flags); + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", il), { q_lora_rank }, flags); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", il), { q_lora_rank, n_head * qk_head_dim }, flags); + } else { + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", il), { n_embd, n_head * qk_head_dim }, flags); + } + layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", il), { n_embd, kv_lora_rank + qk_rope_head_dim }, flags); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", il), { kv_lora_rank }, flags); + layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", il), { qk_head_dim - qk_rope_head_dim, kv_lora_rank, n_head }, flags); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", il), { kv_lora_rank, v_head_dim, n_head }, flags); + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, n_head }, flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_head * v_head_dim, n_embd }, flags); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", il), { n_embd }, flags); + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", il), { n_expert }, flags); + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", il), { n_embd, hparams.n_ff_exp, n_expert }, flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", il), { n_embd, hparams.n_ff_exp, n_expert }, flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { hparams.n_ff_exp, n_embd, n_expert }, flags); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { hparams.n_ff_shexp, n_embd }, flags); + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, flags); + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, flags); + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_LAYER_OUT_NORM, "weight", il), { n_embd }, flags); + } +} + +std::unique_ptr llama_model_bailingmoe3::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } + return std::make_unique(*this, params); +} + +static ggml_tensor * bailingmoe3_causal_conv1d( + ggml_cgraph * gf, + ggml_context * ctx0, + ggml_tensor * conv_states_all, + ggml_tensor * conv_state_all, + int64_t qkv, + ggml_tensor * x, + ggml_tensor * proj_w, + ggml_tensor * conv_w, + int64_t d_conv, + int64_t head_dim, + int64_t n_head, + int64_t n_seq_tokens, + int64_t n_seqs, + int64_t n_tokens, + int64_t cache_head) { + const int64_t d_inner = head_dim * n_head; + const int64_t conv_state_size = (d_conv - 1) * d_inner; + const int64_t total_state_size = 3 * conv_state_size; + + ggml_tensor * conv_state = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_state_all), + total_state_size * ggml_element_size(conv_state_all), + qkv * conv_state_size * ggml_element_size(conv_state_all)); + + ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x); + x_proj = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); + ggml_tensor * conv_x = ggml_concat(ctx0, conv_state, ggml_transpose(ctx0, x_proj), 0); + + ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs, + conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv_x, + ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_states_all), + total_state_size * ggml_element_size(conv_states_all), + (cache_head * total_state_size + qkv * conv_state_size) * ggml_element_size(conv_states_all)))); + + ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner); + ggml_tensor * out = ggml_ssm_conv(ctx0, conv_x, conv_weight); + out = ggml_silu(ctx0, ggml_reshape_2d(ctx0, out, d_inner, n_tokens)); + return ggml_reshape_4d(ctx0, out, head_dim, n_head, n_seq_tokens, n_seqs); +} + +llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_build_delta_net_base(params), model(model) { + ggml_tensor * inpL = build_inp_embd(model.tok_embd); + cb(inpL, "model.input_embed", -1); + + auto * inp = build_inp_mem_hybrid_k(); + auto * inp_rs = inp->get_recr(); + auto * inp_attn = inp->get_attn(); + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + const int64_t n_head = hparams.n_head(); + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_inner = n_head * head_dim; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + const int64_t qk_head_dim = hparams.n_embd_head_k_mla(); + const int64_t v_head_dim = hparams.n_embd_head_v_mla(); + const int64_t qk_rope_head_dim = hparams.n_rot(); + const int64_t qk_nope_head_dim = qk_head_dim - qk_rope_head_dim; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const float kq_scale = 1.0f / sqrtf((float) qk_head_dim); + + GGML_ASSERT(n_seqs > 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + ggml_tensor * inpSA = inpL; + ggml_tensor * cur = build_norm(inpL, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + if (hparams.is_recr(il)) { + const auto * mctx_cur = inp_rs->mctx; + const auto cache_head = mctx_cur->get_head(); + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs); + + ggml_tensor * q = bailingmoe3_causal_conv1d( + gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head); + ggml_tensor * k = bailingmoe3_causal_conv1d( + gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head); + ggml_tensor * v = bailingmoe3_causal_conv1d( + gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head); + + ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_f_a, cur); + gate = ggml_add(ctx0, gate, layer.ssm_dt_b); + gate = ggml_reshape_3d(ctx0, gate, head_dim, n_head, n_tokens); + ggml_tensor * a = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head, 1); + gate = ggml_scale(ctx0, ggml_sigmoid(ctx0, ggml_mul(ctx0, gate, a)), hparams.kda_gate_lower_bound); + gate = ggml_reshape_4d(ctx0, gate, head_dim, n_head, n_seq_tokens, n_seqs); + cb(gate, "kda_gate", il); + + ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur); + beta = ggml_sigmoid(ctx0, ggml_reshape_4d(ctx0, beta, 1, n_head, n_seq_tokens, n_seqs)); + + q = ggml_l2_norm(ctx0, q, hparams.f_norm_rms_eps); + k = ggml_l2_norm(ctx0, k, hparams.f_norm_rms_eps); + + ggml_tensor * states_all = mctx_cur->get_s_l(il); + ggml_tensor * state = build_rs(inp_rs, states_all, hparams.n_embd_s(), n_seqs); + state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head, n_seqs); + + auto result = build_delta_net(q, k, v, gate, beta, state, il); + ggml_tensor * out = ggml_cont(ctx0, result.first); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, result.second, + ggml_view_1d(ctx0, states_all, hparams.n_embd_s() * n_seqs, + cache_head * hparams.n_embd_s() * ggml_element_size(states_all)))); + + ggml_tensor * out_gate = ggml_mul_mat(ctx0, layer.ssm_g_a, cur); + out_gate = ggml_reshape_3d(ctx0, out_gate, head_dim, n_head, n_tokens); + out = ggml_reshape_3d(ctx0, out, head_dim, n_head, n_tokens); + out = build_norm(out, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il); + out = ggml_mul(ctx0, out, ggml_sigmoid(ctx0, out_gate)); + cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, out, d_inner, n_tokens)); + cb(cur, "kda_out", il); + } else { + ggml_tensor * attn_input = cur; + ggml_tensor * q_all; + if (layer.wq_a) { + q_all = ggml_mul_mat(ctx0, layer.wq_a, cur); + cb(q_all, "q_a", il); + q_all = build_norm(q_all, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + cb(q_all, "q_a_norm", il); + q_all = ggml_mul_mat(ctx0, layer.wq_b, q_all); + cb(q_all, "q_b", il); + } else { + q_all = ggml_mul_mat(ctx0, layer.wq, cur); + } + ggml_tensor * q_nope = ggml_view_3d(ctx0, q_all, qk_nope_head_dim, n_head, n_tokens, + ggml_row_size(q_all->type, qk_head_dim), + ggml_row_size(q_all->type, qk_head_dim) * n_head, 0); + ggml_tensor * q_pe = ggml_view_3d(ctx0, q_all, qk_rope_head_dim, n_head, n_tokens, + ggml_row_size(q_all->type, qk_head_dim), + ggml_row_size(q_all->type, qk_head_dim) * n_head, + ggml_row_size(q_all->type, qk_nope_head_dim)); + + ggml_tensor * kv_all = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + ggml_tensor * kv = ggml_view_2d(ctx0, kv_all, kv_lora_rank, n_tokens, + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), 0); + ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_all, qk_rope_head_dim, 1, n_tokens, + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), + ggml_row_size(kv_all->type, kv_lora_rank)); + + q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + kv = build_norm(kv, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + q_nope = ggml_mul_mat(ctx0, layer.wk_b, q_nope); + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + + ggml_tensor * q = ggml_concat(ctx0, q_nope, q_pe, 0); + kv = ggml_reshape_3d(ctx0, kv, kv_lora_rank, 1, n_tokens); + ggml_tensor * k = ggml_concat(ctx0, kv, k_pe, 0); + + cur = build_attn(inp_attn, nullptr, nullptr, nullptr, + q, k, kv, nullptr, nullptr, layer.wv_b, kq_scale, il); + + ggml_tensor * attn_gate = ggml_mul_mat(ctx0, layer.wqkv_gate, attn_input); + attn_gate = ggml_sigmoid(ctx0, ggml_reshape_3d(ctx0, attn_gate, 1, n_head, n_tokens)); + cur = ggml_reshape_3d(ctx0, cur, v_head_dim, n_head, n_tokens); + cur = ggml_mul(ctx0, cur, attn_gate); + cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, cur, v_head_dim * n_head, n_tokens)); + cb(cur, "mla_out", il); + } + + if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) { + 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); + cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + + if ((uint32_t) il < hparams.n_layer_dense_lead) { + cur = build_ffn(cur, + layer.ffn_up, nullptr, nullptr, + layer.ffn_gate, nullptr, nullptr, + layer.ffn_down, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + } else { + ggml_tensor * moe = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_SILU, + hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + 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); + cur = ggml_add(ctx0, moe, shared); + } + + cur = ggml_add(ctx0, cur, ffn_inp); + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + inpL = cur; + } + + ggml_tensor * cur = build_norm(inpL, model.output_norm, nullptr, LLM_NORM_RMS, -1); + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + if (!cparams.embeddings_nextn_masked && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + ggml_build_forward_expand(gf, cur); +} + +llama_model_bailingmoe3::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) : + llm_graph_context(params) { + GGML_ASSERT(hparams.n_layer_nextn == 1 && "BailingMoE3 MTP requires one NextN layer"); + + const int il = hparams.n_layer() + cparams.nextn_layer_offset; + GGML_ASSERT(cparams.nextn_layer_offset >= 0 && + cparams.nextn_layer_offset < (int) hparams.n_layer_nextn && + "nextn_layer_offset out of range"); + const auto & layer = model.layers[il]; + + GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj"); + GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm"); + GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm"); + GGML_ASSERT(layer.nextn.shared_head_norm && "MTP block missing final norm"); + + const int64_t n_head = hparams.n_head(); + const int64_t qk_head_dim = hparams.n_embd_head_k_mla(); + const int64_t v_head_dim = hparams.n_embd_head_v_mla(); + const int64_t qk_rope_head_dim = hparams.n_rot(); + const int64_t qk_nope_head_dim = qk_head_dim - qk_rope_head_dim; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const float kq_scale = 1.0f / sqrtf((float) qk_head_dim); + + auto inp = std::make_unique(hparams.n_embd); + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens); + ggml_set_input(inp->embd); + ggml_set_name(inp->embd, "mtp_h_input"); + + ggml_tensor * tok_embd = ggml_get_rows(ctx0, model.tok_embd, inp->tokens); + ggml_tensor * h_norm = build_norm(inp->embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); + ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); + ggml_tensor * cur = ggml_mul_mat(ctx0, layer.nextn.eh_proj, ggml_concat(ctx0, e_norm, h_norm, 0)); + cb(cur, "mtp_eh_proj", il); + + res->add_input(std::move(inp)); + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + auto * inp_attn = build_attn_inp_k(); + + ggml_tensor * inpSA = cur; + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + ggml_tensor * attn_input = cur; + + ggml_tensor * q_all; + if (layer.wq_a) { + q_all = ggml_mul_mat(ctx0, layer.wq_a, cur); + cb(q_all, "q_a", il); + q_all = build_norm(q_all, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + cb(q_all, "q_a_norm", il); + q_all = ggml_mul_mat(ctx0, layer.wq_b, q_all); + cb(q_all, "q_b", il); + } else { + q_all = ggml_mul_mat(ctx0, layer.wq, cur); + } + ggml_tensor * q_nope = ggml_view_3d(ctx0, q_all, qk_nope_head_dim, n_head, n_tokens, + ggml_row_size(q_all->type, qk_head_dim), + ggml_row_size(q_all->type, qk_head_dim) * n_head, 0); + ggml_tensor * q_pe = ggml_view_3d(ctx0, q_all, qk_rope_head_dim, n_head, n_tokens, + ggml_row_size(q_all->type, qk_head_dim), + ggml_row_size(q_all->type, qk_head_dim) * n_head, + ggml_row_size(q_all->type, qk_nope_head_dim)); + + ggml_tensor * kv_all = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + ggml_tensor * kv = ggml_view_2d(ctx0, kv_all, kv_lora_rank, n_tokens, + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), 0); + ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_all, qk_rope_head_dim, 1, n_tokens, + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), + ggml_row_size(kv_all->type, kv_lora_rank)); + + q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + kv = build_norm(kv, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + q_nope = ggml_mul_mat(ctx0, layer.wk_b, q_nope); + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + + ggml_tensor * q = ggml_concat(ctx0, q_nope, q_pe, 0); + kv = ggml_reshape_3d(ctx0, kv, kv_lora_rank, 1, n_tokens); + ggml_tensor * k = ggml_concat(ctx0, kv, k_pe, 0); + + cur = build_attn(inp_attn, nullptr, nullptr, nullptr, + q, k, kv, nullptr, nullptr, layer.wv_b, kq_scale, il); + + ggml_tensor * attn_gate = ggml_mul_mat(ctx0, layer.wqkv_gate, attn_input); + attn_gate = ggml_sigmoid(ctx0, ggml_reshape_3d(ctx0, attn_gate, 1, n_head, n_tokens)); + cur = ggml_reshape_3d(ctx0, cur, v_head_dim, n_head, n_tokens); + cur = ggml_mul(ctx0, cur, attn_gate); + cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, cur, v_head_dim * n_head, n_tokens)); + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + + ggml_tensor * moe = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_SILU, + hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + 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); + cur = ggml_add(ctx0, moe, shared); + cur = ggml_add(ctx0, cur, ffn_inp); + cur = build_norm(cur, layer.nextn.shared_head_norm, nullptr, LLM_NORM_RMS, -1); + + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/deepseek32.cpp b/src/models/deepseek32.cpp index 8a07a0b71..08555a801 100644 --- a/src/models/deepseek32.cpp +++ b/src/models/deepseek32.cpp @@ -180,10 +180,11 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_ const int64_t n_indexer_head = hparams.indexer_n_head; const int64_t n_embd_indexer_head = hparams.indexer_head_size; - const int64_t n_embd_indexer_head_rope = hparams.n_rot(); - const int64_t n_embd_indexer_head_nope = n_embd_indexer_head - n_embd_indexer_head_rope; const uint32_t n_indexer_top_k = hparams.indexer_top_k; + // the indexer head layous is [rope | nope] + GGML_ASSERT(hparams.n_rot() <= n_embd_indexer_head); + const uint32_t kv_lora_rank = hparams.n_lora_kv; // We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly. @@ -233,28 +234,11 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_ ggml_tensor * indexer_q = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_q_b, qr); cb(indexer_q, "indexer_q", il); - // split into {n_embd_indexer_head_rope, n_indexer_head, n_tokens} - ggml_tensor * indexer_q_pe = - ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_rope, n_indexer_head, n_tokens, - ggml_row_size(indexer_q->type, n_embd_indexer_head), - ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head, 0); - cb(indexer_q_pe, "indexer_q_pe", il); - - // and {n_embd_indexer_head_nope, n_indexer_head, n_tokens} - ggml_tensor * indexer_q_nope = - ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_nope, n_indexer_head, n_tokens, - ggml_row_size(indexer_q->type, n_embd_indexer_head), - ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head, - ggml_row_size(indexer_q->type, n_embd_indexer_head_nope)); - cb(indexer_q_nope, "indexer_q_nope", il); - - indexer_q_pe = ggml_rope_ext(ctx0, indexer_q_pe, inp_pos, nullptr, n_rot, + // {n_embd_indexer_head, n_indexer_head, n_tokens} + indexer_q = ggml_reshape_3d(ctx0, indexer_q, n_embd_indexer_head, n_indexer_head, n_tokens); + indexer_q = ggml_rope_ext(ctx0, indexer_q, inp_pos, nullptr, n_rot, LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - cb(indexer_q_pe, "indexer_q_pe", il); - - // {n_embd_indexer_head_rope + n_embd_indexer_head_nope, n_head, n_tokens} - indexer_q = ggml_concat(ctx0, indexer_q_pe, indexer_q_nope, 0); cb(indexer_q, "indexer_q", il); ggml_tensor * indexer_k = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_k, cur); @@ -263,28 +247,11 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_ indexer_k = build_norm(indexer_k, model.layers[il].indexer_k_norm, model.layers[il].indexer_k_norm_b, LLM_NORM, il); cb(indexer_k, "indexer_k", il); - // split into {n_embd_indexer_head_rope, 1, n_tokens} - ggml_tensor * indexer_k_pe = - ggml_view_3d(ctx0, indexer_k, n_embd_indexer_head_rope, 1, n_tokens, - ggml_row_size(indexer_k->type, n_embd_indexer_head), - ggml_row_size(indexer_k->type, n_embd_indexer_head) * 1, 0); - cb(indexer_k_pe, "indexer_k_pe", il); - - // and {n_embd_indexer_head_nope, 1, n_tokens} - ggml_tensor * indexer_k_nope = - ggml_view_3d(ctx0, indexer_k, n_embd_indexer_head_nope, 1, n_tokens, - ggml_row_size(indexer_k->type, n_embd_indexer_head), - ggml_row_size(indexer_k->type, n_embd_indexer_head) * 1, - ggml_row_size(indexer_k->type, n_embd_indexer_head_nope)); - cb(indexer_k_nope, "indexer_k_nope", il); - - indexer_k_pe = ggml_rope_ext(ctx0, indexer_k_pe, inp_pos, nullptr, n_rot, + // {n_embd_indexer_head, 1, n_tokens} + indexer_k = ggml_reshape_3d(ctx0, indexer_k, n_embd_indexer_head, 1, n_tokens); + indexer_k = ggml_rope_ext(ctx0, indexer_k, inp_pos, nullptr, n_rot, LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - cb(indexer_k_pe, "indexer_k_pe", il); - - // {n_embd_indexer_head_rope + n_embd_indexer_head_nope, 1, n_tokens} - indexer_k = ggml_concat(ctx0, indexer_k_pe, indexer_k_nope, 0); cb(indexer_k, "indexer_k", il); // perform Hadamard transform on indexer q and k diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 39e4f5dd8..5b70a5179 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -85,6 +85,16 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { const int64_t n_embd_inp = hparams.n_embd_inp_enc(); tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED); + + // reduced draft vocab (optional): d2t maps draft rows to target token ids + int64_t n_vocab_draft = n_vocab; + const struct ggml_tensor * d2t_meta = ml->get_tensor_meta("d2t"); + if (d2t_meta) { + n_vocab_draft = d2t_meta->ne[0]; + d2t = create_tensor(tn(LLM_TENSOR_D2T), { n_vocab_draft }, 0); + LLAMA_LOG_INFO("%s: DFlash using d2t mapping (draft_vocab_size = %lld)\n", __func__, (long long) n_vocab_draft); + } + // DSpark = DFlash + a semi-autoregressive Markov head and Confidence head // // TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4) @@ -94,7 +104,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { const int64_t dspark_markov_rank = markov_meta->ne[0]; dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0); - dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab }, 0); + dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab_draft }, 0); dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, 0); dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED); @@ -157,6 +167,9 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { return; } + // optional: reduced-vocab drafts ship their own, full-vocab drafts share the target's via ctx_other + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab_draft }, TENSOR_NOT_REQUIRED); + for (int i = 0; i < n_layer; ++i) { auto & layer = layers[i]; @@ -242,6 +255,11 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & const int64_t block_size = std::stoi(it->second); GGML_ASSERT(block_size > 0); + // bonus anchor (SpecForge exports): slot 0 is a bonus token, not a prediction slot + const auto it_anchor = model.gguf_kv.find("dflash.sample_from_anchor"); + const bool sample_from_anchor = it_anchor == model.gguf_kv.end() || it_anchor->second == "true"; + const int64_t i_draft_beg = sample_from_anchor ? 0 : 1; + const int64_t n_blocks = g.ubatch.n_seqs_unq; GGML_ASSERT(n_blocks > 0 && n_tok % n_blocks == 0 && "DSpark markov head requires equal-size blocks"); // runtime tokens per block in this ubatch (anchor + drafted positions), bounded by training block_size @@ -263,11 +281,26 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & ggml_tensor * cat = nullptr; ggml_tensor * cat_conf = nullptr; + if (!sample_from_anchor) { + // bonus anchor slot: pass the logits through unbiased, pad the (unread) confidence column + cat = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0)); + cat_conf = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0))); + } + // TODO: the in-graph chain is greedy (argmax); sampling params affect only the final // token pick, not the Markov conditioning path - for (int64_t i = 0; i < block_drafts; ++i) { + for (int64_t i = i_draft_beg; i < block_drafts; ++i) { ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev); // [R, n_blocks] - ggml_tensor * bias = ggml_mul_mat(ctx0, w2, w1_prev); // [n_vocab, n_blocks] + ggml_tensor * bias = ggml_mul_mat(ctx0, w2, w1_prev); // [n_vocab_draft, n_blocks] + if (model.d2t) { + // reduced draft vocab: scatter the bias to the target rows (base is -inf on the others) + const int64_t n_draft_vocab = bias->ne[0]; + ggml_tensor * full = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_vocab, n_blocks), 0.0f); + bias = ggml_set_rows(ctx0, full, + ggml_reshape_3d(ctx0, bias, 1, n_draft_vocab, n_blocks), + ggml_reshape_3d(ctx0, model.d2t, n_draft_vocab, 1, 1)); + bias = ggml_reshape_2d(ctx0, bias, n_vocab, n_blocks); + } // position i of every block: strided view [n_vocab, n_blocks] ggml_tensor * base_i = ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, i*base->nb[1]); @@ -497,6 +530,22 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra } cur = build_lora_mm(output, cur, output_s); + + // reduced-draft-vocab exports: scatter the draft logits to the target vocabulary via d2t + if (model.d2t) { + const int64_t n_draft_vocab = cur->ne[0]; + const int64_t n_outputs = cur->ne[1]; + const int64_t n_vocab = (int64_t) model.vocab.n_tokens(); + + GGML_ASSERT(model.d2t->type == GGML_TYPE_I64); + GGML_ASSERT(model.d2t->ne[0] == n_draft_vocab); + + ggml_tensor * logits = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_vocab, n_outputs), -INFINITY); + cur = ggml_set_rows(ctx0, logits, + ggml_reshape_3d(ctx0, cur, 1, n_draft_vocab, n_outputs), + ggml_reshape_3d(ctx0, model.d2t, n_draft_vocab, 1, 1)); + cur = ggml_reshape_2d(ctx0, cur, n_vocab, n_outputs); + } cb(cur, "result_output", -1); res->t_logits = cur; diff --git a/src/models/glm-dsa.cpp b/src/models/glm-dsa.cpp index 360c2ee77..803ef7674 100644 --- a/src/models/glm-dsa.cpp +++ b/src/models/glm-dsa.cpp @@ -216,10 +216,11 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par const int64_t n_indexer_head = hparams.indexer_n_head; const int64_t n_embd_indexer_head = hparams.indexer_head_size; - const int64_t n_embd_indexer_head_rope = hparams.n_rot(); - const int64_t n_embd_indexer_head_nope = n_embd_indexer_head - n_embd_indexer_head_rope; const uint32_t n_indexer_top_k = hparams.indexer_top_k; + // the indexer head layout is [rope | nope] + GGML_ASSERT(hparams.n_rot() <= n_embd_indexer_head); + const uint32_t kv_lora_rank = hparams.n_lora_kv; // We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly. @@ -273,28 +274,11 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par ggml_tensor * indexer_q = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_q_b, qr); cb(indexer_q, "indexer_q", il); - // split into {n_embd_indexer_head_rope, n_indexer_head, n_tokens} - ggml_tensor * indexer_q_pe = - ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_rope, n_indexer_head, n_tokens, - ggml_row_size(indexer_q->type, n_embd_indexer_head), - ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head, 0); - cb(indexer_q_pe, "indexer_q_pe", il); - - // and {n_embd_indexer_head_nope, n_indexer_head, n_tokens} - ggml_tensor * indexer_q_nope = - ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_nope, n_indexer_head, n_tokens, - ggml_row_size(indexer_q->type, n_embd_indexer_head), - ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head, - ggml_row_size(indexer_q->type, n_embd_indexer_head_nope)); - cb(indexer_q_nope, "indexer_q_nope", il); - - indexer_q_pe = ggml_rope_ext(ctx0, indexer_q_pe, inp_pos, nullptr, n_rot, + // {n_embd_indexer_head, n_indexer_head, n_tokens} + indexer_q = ggml_reshape_3d(ctx0, indexer_q, n_embd_indexer_head, n_indexer_head, n_tokens); + indexer_q = ggml_rope_ext(ctx0, indexer_q, inp_pos, nullptr, n_rot, LLAMA_ROPE_TYPE_NORM, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - cb(indexer_q_pe, "indexer_q_pe", il); - - // {n_embd_indexer_head_rope + n_embd_indexer_head_nope, n_head, n_tokens} - indexer_q = ggml_concat(ctx0, indexer_q_pe, indexer_q_nope, 0); cb(indexer_q, "indexer_q", il); ggml_tensor * indexer_k = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_k, cur); @@ -303,28 +287,11 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par indexer_k = build_norm(indexer_k, model.layers[il].indexer_k_norm, model.layers[il].indexer_k_norm_b, LLM_NORM, il); cb(indexer_k, "indexer_k", il); - // split into {n_embd_indexer_head_rope, 1, n_tokens} - ggml_tensor * indexer_k_pe = - ggml_view_3d(ctx0, indexer_k, n_embd_indexer_head_rope, 1, n_tokens, - ggml_row_size(indexer_k->type, n_embd_indexer_head), - ggml_row_size(indexer_k->type, n_embd_indexer_head) * 1, 0); - cb(indexer_k_pe, "indexer_k_pe", il); - - // and {n_embd_indexer_head_nope, 1, n_tokens} - ggml_tensor * indexer_k_nope = - ggml_view_3d(ctx0, indexer_k, n_embd_indexer_head_nope, 1, n_tokens, - ggml_row_size(indexer_k->type, n_embd_indexer_head), - ggml_row_size(indexer_k->type, n_embd_indexer_head) * 1, - ggml_row_size(indexer_k->type, n_embd_indexer_head_nope)); - cb(indexer_k_nope, "indexer_k_nope", il); - - indexer_k_pe = ggml_rope_ext(ctx0, indexer_k_pe, inp_pos, nullptr, n_rot, + // {n_embd_indexer_head, 1, n_tokens} + indexer_k = ggml_reshape_3d(ctx0, indexer_k, n_embd_indexer_head, 1, n_tokens); + indexer_k = ggml_rope_ext(ctx0, indexer_k, inp_pos, nullptr, n_rot, LLAMA_ROPE_TYPE_NORM, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - cb(indexer_k_pe, "indexer_k_pe", il); - - // {n_embd_indexer_head_rope + n_embd_indexer_head_nope, 1, n_tokens} - indexer_k = ggml_concat(ctx0, indexer_k_pe, indexer_k_nope, 0); cb(indexer_k, "indexer_k", il); // perform Hadamard transform on indexer q and k diff --git a/src/models/kimi-k3.cpp b/src/models/kimi-k3.cpp new file mode 100644 index 000000000..d952d72cd --- /dev/null +++ b/src/models/kimi-k3.cpp @@ -0,0 +1,614 @@ +#include "models.h" +#include "llama-memory-recurrent.h" + +// +// Kimi-K3 text model: hybrid KDA (linear) + MLA (full) attention, as in kimi-linear. +// Parts that kimi-linear does not have: +// 1. cross-layer residual attention (attn_res_block_size) +// 2. latent MoE (routed experts run at n_expert_latent) +// 3. situ activation (replaces SwiGLU everywhere) +// 4. MLA output gate (sigmoid gate before o_proj) +// 5. full-rank KDA gate (single ssm_g instead of ssm_g_a/ssm_g_b) +// + +void llama_model_kimi_k3::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl); + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q, false); + ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv); + ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); + ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda); + ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound, false); + + // the MLA cache holds the compressed latent + // set it here too, as older GGUFs have no value_length key + hparams.n_embd_head_v_full = hparams.n_lora_kv; + + // n_head_kv == 0 marks a KDA (recurrent) layer, as in kimi-linear + for (uint32_t i = 0; i < hparams.n_layer(); ++i) { + hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0; + } + + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); + ml.get_key(LLM_KV_EXPERT_LATENT_LENGTH, hparams.n_expert_latent, false); + + ml.get_key(LLM_KV_ATTN_RES_BLOCK_SIZE, hparams.attn_res_block_size); + ml.get_key(LLM_KV_ACTIVATION_SITU_BETA, hparams.situ_beta); + ml.get_key(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, hparams.situ_linear_beta); + + switch (hparams.n_layer()) { + case 93: type = LLM_TYPE_2_8T_A50B; break; // Kimi-K3 + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_kimi_k3::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + const int64_t n_embd_latent = hparams.n_expert_latent > 0 ? hparams.n_expert_latent : n_embd; + + 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); + + if (hparams.attn_res_block_size > 0) { + output_res_score = create_tensor(tn(LLM_TENSOR_OUTPUT_RES_SCORE, "weight"), {n_embd}, 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.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + + if (hparams.attn_res_block_size > 0) { + layer.attn_res_score = create_tensor(tn(LLM_TENSOR_ATTN_RES_SCORE, "weight", i), {n_embd}, 0); + layer.ffn_res_score = create_tensor(tn(LLM_TENSOR_FFN_RES_SCORE, "weight", i), {n_embd}, 0); + } + + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t d_inner = head_dim * n_head; + + if (hparams.is_recr(i)) { + // conv1d may be stored 4D [d_conv, 1, d_inner, 1] or 3D (quantization drops the trailing 1) + auto conv = [&](llm_tensor tid) { + ggml_tensor * t = create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner, 1}, TENSOR_NOT_REQUIRED); + return t ? t : create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner}, 0); + }; + layer.ssm_q_conv = conv(LLM_TENSOR_SSM_CONV1D_Q); + layer.ssm_k_conv = conv(LLM_TENSOR_SSM_CONV1D_K); + layer.ssm_v_conv = conv(LLM_TENSOR_SSM_CONV1D_V); + + create_tensor_qkv(layer, i, n_embd, d_inner, d_inner, d_inner, 0); + + layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", i), {n_embd, head_dim}, 0); + layer.ssm_f_b = create_tensor(tn(LLM_TENSOR_SSM_F_B, "weight", i), {head_dim, d_inner}, 0); + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0); + + // K3's A_log is a plain 1-D [n_head] tensor (kimi-linear's is padded) + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {n_head}, 0); + layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, 0); + + // K3 uses a single full-rank gate instead of kimi-linear's g_a/g_b pair + layer.ssm_g = create_tensor(tn(LLM_TENSOR_SSM_G, "weight", i), {n_embd, d_inner}, 0); + layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_dim}, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {d_inner, n_embd}, 0); + } else { + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t n_embd_head_k = hparams.n_embd_head_k_mla(); + const int64_t n_embd_head_v = hparams.n_embd_head_v_mla(); + const int64_t qk_rope_head_dim = hparams.n_rot(); + const int64_t qk_nope_head_dim = n_embd_head_k - qk_rope_head_dim; + + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, TENSOR_NOT_REQUIRED); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, 0); + + if (layer.attn_q_a_norm) { + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k}, 0); + } else { + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_head * n_embd_head_k}, 0); + } + + layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + qk_rope_head_dim}, 0); + layer.wkv_b = create_tensor(tn(LLM_TENSOR_ATTN_KV_B, "weight", i), + {kv_lora_rank, n_head * (qk_nope_head_dim + n_embd_head_v)}, + TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL); + if (!layer.wkv_b) { + layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {qk_nope_head_dim, kv_lora_rank, n_head}, 0); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v, n_head}, 0); + } + + // K3: sigmoid output gate applied to the attention output before o_proj + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head * n_embd_head_v}, TENSOR_NOT_REQUIRED); + + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v, n_embd}, 0); + } + + if (i < (int) hparams.n_layer_dense_lead) { + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + } else { + const int64_t n_ff_exp = hparams.n_ff_exp; + + 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); + + // routed experts live in the latent space + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd_latent, 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_latent, n_expert}, 0); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd_latent, n_ff_exp, n_expert}, 0); + + if (hparams.n_expert_latent > 0) { + layer.ffn_routed_down = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_DOWN, "weight", i), {n_embd, n_embd_latent}, 0); + layer.ffn_routed_up = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_UP, "weight", i), {n_embd_latent, n_embd}, 0); + layer.ffn_routed_norm = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_NORM, "weight", i), {n_embd_latent}, TENSOR_NOT_REQUIRED); + } + + // shared experts stay at n_embd, width = moe_intermediate_size * n_expert_shared + const int64_t n_ff_shexp = n_ff_exp * (hparams.n_expert_shared > 0 ? hparams.n_expert_shared : 1); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, TENSOR_NOT_REQUIRED); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, TENSOR_NOT_REQUIRED); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, TENSOR_NOT_REQUIRED); + } + } +} + +std::unique_ptr llama_model_kimi_k3::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +// situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta) +// linear_beta <= 0 disables the transform on the up branch +static ggml_tensor * kimi_k3_situ(ggml_context * ctx0, ggml_tensor * gate, ggml_tensor * up, + float beta, float linear_beta) { + ggml_tensor * a = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, gate, 1.0f/beta)), beta); + a = ggml_mul(ctx0, a, ggml_sigmoid(ctx0, gate)); + + if (linear_beta > 0.0f) { + up = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, up, 1.0f/linear_beta)), linear_beta); + } + return ggml_mul(ctx0, a, up); +} + +// +// cross-layer residual attention +// + +// layout is [n_embd, n_ckpt, n_tokens]: rms_norm reduces over ne0, dsv4_hc_pre over ne1 +// append the new checkpoint, do not re-fold the whole chain +void llama_model_kimi_k3::graph::res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens) { + ggml_tensor * ckpt = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens); + + resi_stack = resi_stack ? ggml_concat(ctx0, resi_stack, ckpt, 1) : ckpt; +} + +ggml_tensor * llama_model_kimi_k3::graph::res_mix(ggml_tensor * cur, ggml_tensor * score_w, + int64_t n_tokens, int il) { + if (!resi_stack) { + return cur; // layer 0: nothing banked yet + } + + const int n_ckpt = (int) resi_stack->ne[1]; + const float eps = hparams.f_norm_rms_eps; + + ggml_tensor * src = resi_stack; // [n_embd, n_ckpt, n_tokens] + + // one rms_norm scores all checkpoints at once + // note: the scores use the normalized values, but the sum below uses the raw ones + ggml_tensor * sc_src = ggml_rms_norm(ctx0, src, eps); + sc_src = ggml_mul(ctx0, sc_src, score_w); + sc_src = ggml_sum_rows(ctx0, sc_src); // [1, n_ckpt, n_tokens] + sc_src = ggml_reshape_2d(ctx0, sc_src, n_ckpt, n_tokens); + + // the current residual stream is scored apart, so the stack stays append-only + ggml_tensor * sc_cur = ggml_rms_norm(ctx0, cur, eps); + sc_cur = ggml_mul(ctx0, sc_cur, score_w); + sc_cur = ggml_sum_rows(ctx0, sc_cur); // [1, n_tokens] + + ggml_tensor * scores = ggml_concat(ctx0, sc_src, sc_cur, 0); // [n_ckpt+1, n_tokens] + ggml_tensor * probs = ggml_soft_max(ctx0, scores); // over ne0 = n_ckpt+1 + cb(probs, "res_probs", il); + + // split the sum: hc_pre handles the stack, a broadcast-multiply the current stream + ggml_tensor * p_src = ggml_cont(ctx0, ggml_view_2d(ctx0, probs, n_ckpt, n_tokens, probs->nb[1], 0)); + ggml_tensor * p_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, probs, 1, n_tokens, probs->nb[1], + probs->nb[0] * n_ckpt)); + + ggml_tensor * out = ggml_dsv4_hc_pre(ctx0, src, p_src); + out = ggml_add(ctx0, out, ggml_mul(ctx0, cur, p_cur)); + + return out; +} + +llama_model_kimi_k3::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_build_delta_net_base(params), model(model) { + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + cb(inpL, "inp_embd", -1); + + // K3 MLA is nope-only, so there is no position input + + auto * inp_kv = !hparams.is_mla() ? build_inp_mem_hybrid() : nullptr; + auto * inp_k = hparams.is_mla() ? build_inp_mem_hybrid_k() : nullptr; + auto * inp_rs = hparams.is_mla() ? inp_k->get_recr() : inp_kv->get_recr(); + auto * inp_attn_kv = !hparams.is_mla() ? inp_kv->get_attn() : nullptr; + auto * inp_attn_k = hparams.is_mla() ? inp_k->get_attn() : nullptr; + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + const int64_t n_head_kda = hparams.n_head(); + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t d_inner = n_head_kda * head_dim; + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + GGML_ASSERT(n_seqs != 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); + + const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla(); + const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla(); + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t n_embd_head_qk_rope = hparams.n_rot(); + const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope; + const float kq_scale_mla = 1.0f / sqrtf((float) n_embd_head_k_mla); + + const uint32_t res_bs = hparams.attn_res_block_size; + const bool use_attn_res = res_bs > 0; + const int64_t n_embd_latent = hparams.n_expert_latent > 0 ? hparams.n_expert_latent : n_embd; + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + + // the residual stream, banked on checkpoint layers and then restarted + // from the attention output alone + ggml_tensor * prefix_sum = inpL; + + cur = use_attn_res ? res_mix(prefix_sum, layer.attn_res_score, n_tokens, il) + : prefix_sum; + + bool banked = false; + if (use_attn_res && (uint32_t) il % res_bs == 0) { + res_push(prefix_sum, n_embd, n_tokens); // banks the RAW layer input, not `cur` + banked = true; + } + + cur = build_norm(cur, layer.attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + ggml_build_forward_expand(gf, cur); + + if (hparams.is_recr(il)) { + cur = build_kda_layer(cur, layer, inp_rs, d_conv, head_dim, n_head_kda, + d_inner, n_seq_tokens, n_seqs, il); + } else { + cur = build_mla_layer(cur, layer, inp_attn_k, inp_attn_kv, + n_embd_head_k_mla, n_embd_head_v_mla, kv_lora_rank, + n_embd_head_qk_rope, n_embd_head_qk_nope, kq_scale_mla, il); + } + + prefix_sum = banked ? cur : ggml_add(ctx0, prefix_sum, cur); + cb(prefix_sum, "prefix_sum_attn", il); + + cur = use_attn_res ? res_mix(prefix_sum, layer.ffn_res_score, n_tokens, il) + : prefix_sum; + + cur = build_norm(cur, layer.ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + if ((uint32_t) il < hparams.n_layer_dense_lead) { + ggml_tensor * g = ggml_mul_mat(ctx0, layer.ffn_gate, cur); + ggml_tensor * u = ggml_mul_mat(ctx0, layer.ffn_up, cur); + cur = kimi_k3_situ(ctx0, g, u, hparams.situ_beta, hparams.situ_linear_beta); + cur = ggml_mul_mat(ctx0, layer.ffn_down, cur); + cb(cur, "ffn_out", il); + } else { + cur = build_latent_moe(cur, layer, n_embd_latent, il); + } + + prefix_sum = ggml_add(ctx0, prefix_sum, cur); + prefix_sum = build_cvec(prefix_sum, il); + cb(prefix_sum, "l_out", il); + + inpL = prefix_sum; + } + + cur = inpL; + + // final mix, then narrow to the output tokens + if (use_attn_res) { + cur = res_mix(cur, model.output_res_score, n_tokens, -1); + } + if (inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + +// +// KDA layer +// + +// causal conv1d over one of Q/K/V. `qkv` selects which third of the conv state to use +static ggml_tensor * kimi_k3_conv1d(ggml_cgraph * gf, ggml_context * ctx0, + ggml_tensor * conv_states_all, ggml_tensor * conv_state_all, + int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w, + int64_t d_conv, int64_t head_dim, int64_t n_head, + int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head) { + const int64_t d_inner = head_dim * n_head; + const int64_t conv_state_size = (d_conv - 1) * d_inner; + const int64_t n_embd_r_total = 3 * conv_state_size; + + ggml_tensor * conv_state_x = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_state_all), + n_embd_r_total * ggml_element_size(conv_state_all), + qkv * conv_state_size * ggml_element_size(conv_state_all)); + + ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x); + ggml_tensor * x_3d = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); + ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0); + + ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs, + conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]); + ggml_build_forward_expand(gf, + ggml_cpy(ctx0, last_conv_x, + ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_states_all), + n_embd_r_total * ggml_element_size(conv_states_all), + (kv_head * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all)))); + + ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner); + ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight); + Xcur = ggml_reshape_2d(ctx0, Xcur, d_inner, n_tokens); + Xcur = ggml_silu(ctx0, Xcur); + + return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head, n_seq_tokens, n_seqs); +} + +ggml_tensor * llama_model_kimi_k3::graph::build_kda_layer( + ggml_tensor * cur, const llama_layer & layer, llm_graph_input_rs * inp_rs, + int64_t d_conv, int64_t head_dim, int64_t n_head_kda, + int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il) { + + const auto * mctx_cur = inp_rs->mctx; + const auto kv_head = mctx_cur->get_head(); + + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs); + + ggml_tensor * Qcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * Kcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * Vcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head); + cb(Qcur, "kda_q_conv", il); + cb(Kcur, "kda_k_conv", il); + cb(Vcur, "kda_v_conv", il); + + // gate_lower_bound is not a clamp - when set, it swaps the decay gate activation: + // unset (kimi-linear): g = -exp(A_log) * softplus(f_b(f_a(x)) + dt_bias) + // set (K3, -5.0): g = lower_bound * sigmoid(exp(A_log) * (f_b(f_a(x)) + dt_bias)) + // ssm_a holds -exp(A_log) (folded at conversion time), so exp(A_log) == -ssm_a + ggml_tensor * f_a = ggml_mul_mat(ctx0, layer.ssm_f_a, cur); + ggml_tensor * g1 = ggml_mul_mat(ctx0, layer.ssm_f_b, f_a); + g1 = ggml_add(ctx0, g1, layer.ssm_dt_b); + + ggml_tensor * A = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head_kda, 1); + + if (hparams.kda_gate_lower_bound > -INFINITY) { + g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens); + g1 = ggml_mul(ctx0, g1, A); // -exp(A_log) * (...) + g1 = ggml_sigmoid(ctx0, ggml_scale(ctx0, g1, -1.0f)); + g1 = ggml_scale(ctx0, g1, hparams.kda_gate_lower_bound); + } else { + g1 = ggml_softplus(ctx0, g1); + g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens); + g1 = ggml_mul(ctx0, g1, A); + } + cb(g1, "kda_g1", il); + + g1 = ggml_reshape_4d(ctx0, g1, head_dim, n_head_kda, n_seq_tokens, n_seqs); + + ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur); + beta = ggml_reshape_4d(ctx0, beta, 1, n_head_kda, n_seq_tokens, n_seqs); + beta = ggml_sigmoid(ctx0, beta); + cb(beta, "kda_beta", il); + + ggml_tensor * cur_3d = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs); + + ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il); + ggml_tensor * state = build_rs(inp_rs, ssm_states_all, hparams.n_embd_s(), n_seqs); + state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head_kda, n_seqs); + + const float eps = hparams.f_norm_rms_eps; + Qcur = ggml_l2_norm(ctx0, Qcur, eps); + Kcur = ggml_l2_norm(ctx0, Kcur, eps); + + auto attn_out = build_delta_net(Qcur, Kcur, Vcur, g1, beta, state, il); + + ggml_tensor * output = ggml_cont(ctx0, attn_out.first); + cb(output, "kda_scan_out", il); + ggml_tensor * new_state = attn_out.second; + + ggml_build_forward_expand(gf, + ggml_cpy(ctx0, new_state, + ggml_view_1d(ctx0, ssm_states_all, hparams.n_embd_s() * n_seqs, + kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all)))); + + // K3: single full-rank gate (kimi-linear factors this as g_b(g_a(x))) + ggml_tensor * cur_2d = ggml_reshape_2d(ctx0, cur_3d, cur_3d->ne[0], n_seq_tokens * n_seqs); + ggml_tensor * g2 = ggml_mul_mat(ctx0, layer.ssm_g, cur_2d); + g2 = ggml_reshape_3d(ctx0, g2, head_dim, n_head_kda, n_seq_tokens * n_seqs); + + ggml_tensor * o = ggml_reshape_3d(ctx0, output, head_dim, n_head_kda, n_seq_tokens * n_seqs); + ggml_tensor * normed = build_norm(o, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il); + cb(g2, "kda_g2", il); + cb(normed, "kda_normed", il); + ggml_tensor * gated = ggml_mul(ctx0, normed, ggml_sigmoid(ctx0, g2)); + + gated = ggml_cont_2d(ctx0, gated, d_inner, n_tokens); + cur = ggml_mul_mat(ctx0, layer.wo, gated); + cb(cur, "kda_out", il); + + return cur; +} + +// +// MLA layer (nope-only, with K3's sigmoid output gate) +// + +ggml_tensor * llama_model_kimi_k3::graph::build_mla_layer( + ggml_tensor * cur, const llama_layer & layer, + llm_graph_input_attn_k * inp_attn_k, llm_graph_input_attn_kv * inp_attn_kv, + int64_t n_embd_head_k_mla, int64_t n_embd_head_v_mla, int64_t kv_lora_rank, + int64_t n_embd_head_qk_rope, int64_t n_embd_head_qk_nope, float kq_scale, int il) { + + ggml_tensor * inp_gate = cur; // the output gate reads the *normed* layer input + + ggml_tensor * Qcur; + if (layer.wq_a) { + Qcur = ggml_mul_mat(ctx0, layer.wq_a, cur); + Qcur = build_norm(Qcur, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + Qcur = ggml_mul_mat(ctx0, layer.wq_b, Qcur); + } else { + Qcur = ggml_mul_mat(ctx0, layer.wq, cur); + } + + ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + + ggml_tensor * kv_cmpr = ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens, + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0); + ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens, + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank)); + + // no RoPE: mla_use_nope is asserted at conversion time + kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + + ggml_tensor * out; + if (layer.wk_b && layer.wv_b) { + ggml_tensor * q_nope = ggml_view_3d(ctx0, Qcur, n_embd_head_qk_nope, n_head, n_tokens, + ggml_row_size(Qcur->type, n_embd_head_k_mla), + ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head, 0); + ggml_tensor * q_pe = ggml_view_3d(ctx0, Qcur, n_embd_head_qk_rope, n_head, n_tokens, + ggml_row_size(Qcur->type, n_embd_head_k_mla), + ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head, + ggml_row_size(Qcur->type, n_embd_head_qk_nope)); + + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope); + q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3); + + ggml_tensor * Q = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0); + ggml_tensor * kv_cmpr_3d = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens); + ggml_tensor * K = ggml_concat(ctx0, kv_cmpr_3d, k_pe, 0); + ggml_tensor * V = kv_cmpr_3d; + + // wo == NULL: the output projection is applied after the gate below + out = build_attn(inp_attn_k, nullptr, NULL, nullptr, Q, K, V, nullptr, nullptr, layer.wv_b, kq_scale, il); + } else { + ggml_tensor * Q = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k_mla, n_head, n_tokens); + ggml_tensor * kv = ggml_mul_mat(ctx0, layer.wkv_b, kv_cmpr); + const int64_t kv_per_head = n_embd_head_qk_nope + n_embd_head_v_mla; + + ggml_tensor * k_nope = ggml_view_3d(ctx0, kv, n_embd_head_qk_nope, n_head, n_tokens, + ggml_row_size(kv->type, kv_per_head), ggml_row_size(kv->type, kv_per_head * n_head), 0); + ggml_tensor * V = ggml_cont(ctx0, ggml_view_3d(ctx0, kv, n_embd_head_v_mla, n_head, n_tokens, + ggml_row_size(kv->type, kv_per_head), ggml_row_size(kv->type, kv_per_head * n_head), + ggml_row_size(kv->type, n_embd_head_qk_nope))); + + ggml_tensor * k_pe_t = ggml_new_tensor_3d(ctx0, k_pe->type, n_embd_head_qk_rope, n_head, n_tokens); + ggml_tensor * K = ggml_concat(ctx0, ggml_repeat(ctx0, k_pe, k_pe_t), k_nope, 0); + + out = build_attn(inp_attn_kv, nullptr, NULL, nullptr, Q, K, V, nullptr, nullptr, nullptr, kq_scale, il); + } + + // K3: attn_output *= sigmoid(g_proj(x)), then o_proj + if (layer.wqkv_gate) { + ggml_tensor * g = ggml_sigmoid(ctx0, ggml_mul_mat(ctx0, layer.wqkv_gate, inp_gate)); + out = ggml_mul(ctx0, out, g); + cb(out, "mla_gated", il); + } + + out = ggml_mul_mat(ctx0, layer.wo, out); + cb(out, "mla_out", il); + + return out; +} + +// +// latent MoE: down-project, run the routed experts in the latent space, norm, up-project; +// shared experts stay at n_embd and read the un-projected input. +// + +ggml_tensor * llama_model_kimi_k3::graph::build_latent_moe( + ggml_tensor * cur, const llama_layer & layer, int64_t n_embd_latent, int il) { + + ggml_tensor * identity = cur; + + ggml_tensor * routed_in = layer.ffn_routed_down + ? ggml_mul_mat(ctx0, layer.ffn_routed_down, cur) + : cur; + + // the router scores the full-width input while the experts take the latent one, + // so the logits are computed here and passed to build_moe_ffn + ggml_tensor * logits = ggml_mul_mat(ctx0, layer.ffn_gate_inp, identity); + cb(logits, "ffn_moe_logits", il); + + ggml_tensor * moe_out = build_moe_ffn(routed_in, + nullptr, // gate_inp unused: the logits above are passed instead + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + hparams.n_expert, + hparams.n_expert_used, + LLM_FFN_SITU, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il, + logits); + cb(moe_out, "ffn_moe_out", il); + + if (layer.ffn_routed_norm) { + moe_out = build_norm(moe_out, layer.ffn_routed_norm, NULL, LLM_NORM_RMS, il); + } + if (layer.ffn_routed_up) { + moe_out = ggml_mul_mat(ctx0, layer.ffn_routed_up, moe_out); + } + GGML_UNUSED(n_embd_latent); + + if (layer.ffn_gate_shexp) { + ggml_tensor * g = ggml_mul_mat(ctx0, layer.ffn_gate_shexp, identity); + ggml_tensor * u = ggml_mul_mat(ctx0, layer.ffn_up_shexp, identity); + ggml_tensor * sh = kimi_k3_situ(ctx0, g, u, hparams.situ_beta, hparams.situ_linear_beta); + sh = ggml_mul_mat(ctx0, layer.ffn_down_shexp, sh); + cb(sh, "ffn_shexp", il); + moe_out = ggml_add(ctx0, moe_out, sh); + } + + cb(moe_out, "ffn_out", il); + return moe_out; +} diff --git a/src/models/models.h b/src/models/models.h index cefaf56eb..180b30a46 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1784,6 +1784,25 @@ struct llama_model_bailingmoe2 : public llama_model_base { }; +struct llama_model_bailingmoe3 : public llama_model_base { + llama_model_bailingmoe3(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_build_delta_net_base { + graph(const llama_model & model, const llm_graph_params & params); + + const llama_model & model; + }; + + struct graph_mtp : public llm_graph_context { + graph_mtp(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_seed_oss : public llama_model_base { llama_model_seed_oss(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; @@ -2285,6 +2304,42 @@ struct llama_model_mimo2 : public llama_model_base { }; +struct llama_model_kimi_k3 : public llama_model_base { + llama_model_kimi_k3(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_build_delta_net_base { + graph(const llama_model & model, const llm_graph_params & params); + + const llama_model & model; + + // Cross-layer residual attention (K3's `_apply_attn_res`). + ggml_tensor * resi_stack = nullptr; + + void res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens); + ggml_tensor * res_mix(ggml_tensor * cur, ggml_tensor * score_w, + int64_t n_tokens, int il); + + ggml_tensor * build_kda_layer(ggml_tensor * cur, const llama_layer & layer, + llm_graph_input_rs * inp_rs, + int64_t d_conv, int64_t head_dim, int64_t n_head_kda, + int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il); + + ggml_tensor * build_mla_layer(ggml_tensor * cur, const llama_layer & layer, + llm_graph_input_attn_k * inp_attn_k, + llm_graph_input_attn_kv * inp_attn_kv, + int64_t n_embd_head_k_mla, int64_t n_embd_head_v_mla, + int64_t kv_lora_rank, int64_t n_embd_head_qk_rope, + int64_t n_embd_head_qk_nope, float kq_scale, int il); + + ggml_tensor * build_latent_moe(ggml_tensor * cur, const llama_layer & layer, + int64_t n_embd_latent, int il); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + struct llama_model_kimi_linear : public llama_model_base { llama_model_kimi_linear(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp index e795b47e0..d4b5ec0d9 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -1606,12 +1606,16 @@ mtmd_image_preproc_out mtmd_image_preprocessor_granite::preprocess(const clip_im const clip_image_size orig_size = img.get_size(); const int tile_size = hparams.image_size; + GGML_ASSERT(tile_size > 0); // llava-next always encodes an overview plus a grid of tiles, even for small images const clip_image_size refined_size = select_best_resolution(orig_size, hparams.image_res_candidates); const int grid_x = refined_size.width / tile_size; const int grid_y = refined_size.height / tile_size; + // the tiles are stacked on the Y axis, a big grid overflows the stacked image height + GGML_ASSERT(grid_x >= 0 && grid_x <= 1024 && grid_y >= 0 && grid_y <= 1024); + clip_image_u8 overview; img_tool::resize(img, overview, {tile_size, tile_size}, hparams.image_resize_algo_ov, hparams.image_pad_ov, hparams.image_pad_color_ov); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a7d382886..fdc8f2b80 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2947,8 +2947,10 @@ private: }); // generate the actual drafts (if any) - { - common_speculative_draft(spec.get()); + if (!drafting.empty()) { + queue_tasks.yield_to_queue([&]() { + common_speculative_draft(spec.get()); + }); } // make checkpoints if needed @@ -3578,8 +3580,8 @@ private: has_output |= batch.tokens[i].output; } - // decode on the worker thread, so we can still handle metrics tasks while waiting - // note: the sync is done here too, so that the wait also happens off the main thread + // yield to the queue, so we can still handle metrics tasks while decoding + // note: the sync is done here too, so that the wait is also covered by the yield int ret = 0; queue_tasks.yield_to_queue([&]() { ret = llama_decode(ctx_tgt, batch_view); @@ -3644,11 +3646,18 @@ private: // TODO: avoid restoring the draft context and re-evaluating the drafted tokens when not needed [TAG_SPEC_AVOID_DRAFT_REEVAL] // for now, always re-evaluate for simplicity // ref: https://github.com/ggml-org/llama.cpp/pull/22728#issuecomment-4400925384 - if (!common_speculative_process(spec.get(), batch_view)) { - SRV_ERR("%s", "failed to process speculative batch\n"); + if (spec) { + bool ok = true; + queue_tasks.yield_to_queue([&]() { + ok = common_speculative_process(spec.get(), batch_view); + }); - // TODO: handle error - throw std::runtime_error("failed to process speculative batch"); + if (!ok) { + SRV_ERR("%s", "failed to process speculative batch\n"); + + // TODO: handle error + throw std::runtime_error("failed to process speculative batch"); + } } // handle `n_cmpl > 1` tasks - when the main prompt is processed, activate all child tasks too diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index d22527ffe..2bcc9bd8f 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -150,31 +150,46 @@ bool server_queue::process_new_tasks(bool is_yielding) { void server_queue::worker_loop() { while (true) { - std::function work; { std::unique_lock lock(mutex_tasks); + // wait on busy instead of yielding - busy stays set even when the yield already ended worker.cv.wait(lock, [&]{ - return worker.stop || worker.work != nullptr; + return worker.stop || worker.busy; }); if (worker.stop) { return; } - work = std::move(worker.work); - worker.work = nullptr; } - // note: do not hold any lock here, work() may post new tasks - std::exception_ptr exception; - try { - work(); - } catch (...) { - exception = std::current_exception(); + // process tasks while the yield is active + while (true) { + bool terminated = false; + try { + // note: do not hold any lock here, the callback may post new tasks + terminated = process_new_tasks(true); + } catch (...) { + std::unique_lock lock(mutex_tasks); + worker.exception = std::current_exception(); + break; + } + + std::unique_lock lock(mutex_tasks); + if (terminated || worker.stop || !worker.yielding) { + break; + } + if (!queue_tasks.empty()) { + continue; // a new task arrived in the meantime + } + condition_tasks.wait(lock, [&]{ + return worker.stop || !running || !worker.yielding || !queue_tasks.empty(); + }); } - // signal completion to yield_to_queue() - std::unique_lock lock(mutex_tasks); - worker.exception = std::move(exception); - worker.busy = false; + // signal to yield_to_queue() that no more tasks will be processed + { + std::unique_lock lock(mutex_tasks); + worker.busy = false; + } condition_tasks.notify_all(); } } @@ -188,6 +203,7 @@ void server_queue::worker_stop() { worker.stop = true; } worker.cv.notify_one(); + condition_tasks.notify_all(); worker.thread.join(); } @@ -199,29 +215,29 @@ void server_queue::yield_to_queue(std::function && work) { { std::unique_lock lock(mutex_tasks); GGML_ASSERT(!worker.busy && "yield_to_queue() cannot be nested"); - worker.busy = true; - worker.work = std::move(work); + worker.busy = true; + worker.yielding = true; } worker.cv.notify_one(); - while (true) { - // note: on terminate this is a no-op, but we still wait for the work to finish - process_new_tasks(true); - - std::unique_lock lock(mutex_tasks); - // declined tasks are moved to queue_tasks_unhandled, so a non-empty queue always has something new - condition_tasks.wait(lock, [&]{ - return !worker.busy || (running && !queue_tasks.empty()); - }); - if (!worker.busy) { - break; - } + // run the work on the current thread, so that all ggml compute stays on the same thread + std::exception_ptr exception; + try { + work(); + } catch (...) { + exception = std::current_exception(); } - std::exception_ptr exception; { std::unique_lock lock(mutex_tasks); + // the yield is over, wait for the worker to finish its current task + worker.yielding = false; + condition_tasks.notify_all(); + condition_tasks.wait(lock, [&]{ + return !worker.busy; + }); + // put the declined tasks back, keeping their order while (!queue_tasks_unhandled.empty()) { queue_tasks.push_front(std::move(queue_tasks_unhandled.back())); @@ -231,8 +247,12 @@ void server_queue::yield_to_queue(std::function && work) { // make sure to avoid idle timeout here time_last_task = ggml_time_ms(); - // the worker is idle now, take the exception it may have left behind - std::swap(exception, worker.exception); + // an exception from work() takes precedence over the one from the worker + if (!exception) { + std::swap(exception, worker.exception); + } else { + worker.exception = nullptr; + } } QUE_DBG("%s", "done yielding to queue\n"); @@ -249,7 +269,9 @@ void server_queue::start_loop(int64_t idle_sleep_ms) { // spawn the worker thread used by yield_to_queue() GGML_ASSERT(!worker.thread.joinable() && "start_loop() is already running"); - worker.stop = false; + worker.stop = false; + worker.busy = false; + worker.yielding = false; worker.thread = std::thread([this]() { worker_loop(); }); constexpr auto max_wait_time = std::chrono::seconds(1); diff --git a/tools/server/server-queue.h b/tools/server/server-queue.h index 69a3e554f..52d30095c 100644 --- a/tools/server/server-queue.h +++ b/tools/server/server-queue.h @@ -33,11 +33,11 @@ private: // used by yield_to_queue, all fields are guarded by mutex_tasks struct worker_t { std::thread thread; - std::condition_variable cv; // the worker sleeps on this until there is work - std::function work; // pending work, picked up by the thread - std::exception_ptr exception; // exception thrown by work(), if any - bool stop = false; - bool busy = false; + std::condition_variable cv; // the worker sleeps on this until a yield starts + std::exception_ptr exception; // exception thrown while processing tasks, if any + bool stop = false; + bool busy = false; // set by yield_to_queue(), cleared by the worker once it is done processing tasks + bool yielding = false; // work() is still running on the start_loop() thread }; worker_t worker; @@ -93,7 +93,7 @@ public: */ void start_loop(int64_t idle_sleep_ms = -1); - // run work() on a separate thread, while the current thread calls process_new_tasks + // while waiting for work() to finish, run process_new_tasks on the worker thread // returns once work() is done (may throw exceptions) // must be called from start_loop() thread (ideally inside callback_update_slots) // use case: return metrics while encode/decode is running @@ -116,6 +116,7 @@ public: // the second argument tells whether the queue is currently yielding (see yield_to_queue) // only then may the callback return false to decline the task, and it must leave it // untouched, so that it can be put back in the queue later + // note: while yielding, the callback runs on worker thread, not main thread void on_new_task(std::function callback) { callback_new_task = std::move(callback); } diff --git a/tools/server/server-tools.cpp b/tools/server/server-tools.cpp index fd0ff8ddd..5e5e60efd 100644 --- a/tools/server/server-tools.cpp +++ b/tools/server/server-tools.cpp @@ -8,7 +8,6 @@ #include #include #include -#include #include #include #include @@ -1692,61 +1691,6 @@ private: } }; -// -// get_datetime: returns the current date and time -// - -struct server_tool_get_datetime : server_tool { - server_tool_get_datetime() { - name = "get_datetime"; - display_name = "Get Date & Time"; - permission_write = false; - } - - json get_definition() const override { - return { - {"type", "function"}, - {"function", { - {"name", name}, - {"description", "Returns the current date and time in UTC"}, - {"parameters", { - {"type", "object"}, - {"properties", { - {"format", { - {"type", "string"}, - {"description", - "strftime()-style format string for the output (default: \"%Y-%m-%dT%H:%M:%SZ\", " - "e.g. ISO 8601). Choose your own format if you need something else, " - "e.g. \"%A, %B %d %Y\" for a human-readable date."}, - }}, - }}, - }}, - }}, - }; - } - - json invoke(json params, server_tool::stream *) const override { - std::string format = json_value(params, "format", std::string("%Y-%m-%dT%H:%M:%SZ")); - - auto now = std::chrono::system_clock::now(); - auto time = std::chrono::system_clock::to_time_t(now); - std::tm tm_utc; -#ifdef _WIN32 - gmtime_s(&tm_utc, &time); -#else - gmtime_r(&time, &tm_utc); -#endif - - char buf[256]; - size_t len = std::strftime(buf, sizeof(buf), format.c_str(), &tm_utc); - if (len == 0) { - return {{"error", "invalid format string"}}; - } - - return {{"result", std::string(buf, len)}}; - } -}; - // // get_info: returns runtime info (OS name/version and cwd) // @@ -2005,6 +1949,10 @@ static server_tool & find_tool(std::vector> & tools // static std::vector> build_tools() { + // IMPORTANT: for contributors, please keep this array of tools as minimal as possible + // we only accept minimal i/o and shell command tools here + // for example, do not add: web search, get date time, etc. + // high-level functionality should be added either via MCP or web UI std::vector> tools; tools.push_back(std::make_unique()); tools.push_back(std::make_unique()); @@ -2012,7 +1960,6 @@ static std::vector> build_tools() { tools.push_back(std::make_unique()); tools.push_back(std::make_unique()); tools.push_back(std::make_unique()); - tools.push_back(std::make_unique()); tools.push_back(std::make_unique()); return tools; } diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockGetDatetime.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockGetDatetime.svelte index 23de7cbd6..44b3ec645 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockGetDatetime.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockGetDatetime.svelte @@ -33,7 +33,7 @@ if (typeof obj.result === 'string') return { dateString: obj.result.trim() }; } } catch { - return { dateString: toolResultString.trim() }; + // not JSON - nothing to show } return {}; diff --git a/tools/ui/src/lib/components/app/server/ServerErrorSplash.svelte b/tools/ui/src/lib/components/app/server/ServerErrorSplash.svelte index a2e23eecd..fbdbf19f2 100644 --- a/tools/ui/src/lib/components/app/server/ServerErrorSplash.svelte +++ b/tools/ui/src/lib/components/app/server/ServerErrorSplash.svelte @@ -158,6 +158,8 @@
> label: 'Search files', source: ToolSource.BUILTIN }, - [BuiltInTool.GET_DATETIME]: { icon: Clock, label: 'Current time', source: ToolSource.BUILTIN }, + [BuiltInTool.GET_DATETIME]: { icon: Clock, label: 'Current time', source: ToolSource.FRONTEND }, [BuiltInTool.GET_INFO]: { icon: Info, label: 'Runtime info', source: ToolSource.BUILTIN }, [BuiltInTool.GREP_SEARCH]: { icon: SearchCode, diff --git a/tools/ui/src/lib/constants/cli-flags.constants.ts b/tools/ui/src/lib/constants/cli-flags.constants.ts index 4fbee8a36..c4af2b6f4 100644 --- a/tools/ui/src/lib/constants/cli-flags.constants.ts +++ b/tools/ui/src/lib/constants/cli-flags.constants.ts @@ -1,4 +1,5 @@ export const CLI_FLAGS = { + AGENT: '--agent', API_KEY: '--api-key', MCP_PROXY: '--ui-mcp-proxy', SLOTS: '--slots', diff --git a/tools/ui/src/lib/constants/get-datetime.ts b/tools/ui/src/lib/constants/get-datetime.ts new file mode 100644 index 000000000..c8726d9b5 --- /dev/null +++ b/tools/ui/src/lib/constants/get-datetime.ts @@ -0,0 +1,20 @@ +import { BuiltInTool, JsonSchemaType, ToolCallType } from '$lib/enums'; +import type { OpenAIToolDefinition } from '$lib/types'; + +export const GET_DATETIME_TOOL_NAME = BuiltInTool.GET_DATETIME; + +export function buildGetDatetimeToolDefinition(): OpenAIToolDefinition { + return { + function: { + description: + 'Returns the current local date and time in ISO 8601 format, with the IANA time zone name', + name: GET_DATETIME_TOOL_NAME, + parameters: { + properties: {}, + required: [], + type: JsonSchemaType.OBJECT + } + }, + type: ToolCallType.FUNCTION + }; +} diff --git a/tools/ui/src/lib/constants/index.ts b/tools/ui/src/lib/constants/index.ts index 17e5f4d00..289239113 100644 --- a/tools/ui/src/lib/constants/index.ts +++ b/tools/ui/src/lib/constants/index.ts @@ -59,3 +59,5 @@ export * from './uri-template.constants'; export * from './url.constants'; export * from './working-directory.constants'; export * from './read-media'; +export * from './get-datetime'; +export * from './browser-info'; diff --git a/tools/ui/src/lib/constants/settings-registry.constants.ts b/tools/ui/src/lib/constants/settings-registry.constants.ts index 9de4a9720..0aaa98f2a 100644 --- a/tools/ui/src/lib/constants/settings-registry.constants.ts +++ b/tools/ui/src/lib/constants/settings-registry.constants.ts @@ -324,6 +324,7 @@ const SETTINGS_REGISTRY: Record = { { defaultValue: '', help: `Set the API Key if you are using ${CLI_FLAGS.API_KEY} option for the server.`, + isPrivate: true, key: SETTINGS_KEYS.API_KEY, label: 'API Key', section: SETTINGS_SECTION_SLUGS.GENERAL, @@ -713,6 +714,7 @@ export const SETTINGS_CHAT_SECTIONS: SettingsSection[] = [ help: s.help, isExperimental: s.isExperimental, isPositiveInteger: s.isPositiveInteger, + isPrivate: s.isPrivate, key: s.key, label: s.label, max: s.max, diff --git a/tools/ui/src/lib/services/mcp.service.ts b/tools/ui/src/lib/services/mcp.service.ts index 43f905c33..65e9e59d6 100644 --- a/tools/ui/src/lib/services/mcp.service.ts +++ b/tools/ui/src/lib/services/mcp.service.ts @@ -18,7 +18,8 @@ import { DEFAULT_CLIENT_VERSION, DEFAULT_IMAGE_MIME_TYPE, DEFAULT_MCP_CONFIG, - HEADERS + HEADERS, + NEWLINE } from '$lib/constants'; import { MCPConnectionPhase, @@ -70,6 +71,7 @@ interface ToolResultContentItem { interface ToolCallResult { content?: ToolResultContentItem[]; + structuredContent?: Record; isError?: boolean; _meta?: Record; } @@ -1012,10 +1014,20 @@ export class MCPService { if (!Array.isArray(content)) return ''; - return content + const formatted = content .map((item) => this.formatSingleContent(item)) .filter(Boolean) - .join('\n'); + .join(NEWLINE); + + if (formatted !== '') { + return formatted; + } + + if (result.structuredContent && typeof result.structuredContent === 'object') { + return JSON.stringify(result.structuredContent); + } + + return ''; } private static formatSingleContent(content: ToolResultContentItem): string { diff --git a/tools/ui/src/lib/stores/agentic.svelte.ts b/tools/ui/src/lib/stores/agentic.svelte.ts index 075c46c97..06c7661fe 100644 --- a/tools/ui/src/lib/stores/agentic.svelte.ts +++ b/tools/ui/src/lib/stores/agentic.svelte.ts @@ -56,7 +56,8 @@ import type { AgenticSession, McpServerOverride, MCPToolCall, - SettingsConfigType + SettingsConfigType, + ToolExecutionResult } from '$lib/types'; import type { AgenticFlowCallbacks, @@ -83,7 +84,12 @@ import type { DatabaseMessageExtraAudioFile, DatabaseMessageExtraImageFile } from '$lib/types/database'; -import { getAudioInputFormat, isAbortError } from '$lib/utils'; +import { + executeBrowserInfoTool, + executeGetDatetimeTool, + getAudioInputFormat, + isAbortError +} from '$lib/utils'; import { SvelteMap } from 'svelte/reactivity'; function createDefaultSession(): AgenticSession { @@ -942,18 +948,26 @@ class AgenticStore { if (executionResult.isError) toolSuccess = false; } else if (toolSource === ToolSource.FRONTEND) { const args = this.parseToolArguments(toolCall.function.arguments); - const executionResult = - toolName === BuiltInTool.READ_MEDIA - ? await ReadMediaService.executeTool( - args, - { - audio: modelsStore.modelSupportsAudio(effectiveModel), - vision: modelsStore.modelSupportsVision(effectiveModel) - }, - signal, - conversationsStore.activeConversation?.cwd - ) - : await SandboxService.executeTool(toolName, args, signal); + + let executionResult: ToolExecutionResult; + + if (toolName === BuiltInTool.GET_DATETIME) { + executionResult = executeGetDatetimeTool(); + } else if (toolName === BuiltInTool.GET_INFO) { + executionResult = executeBrowserInfoTool(); + } else if (toolName === BuiltInTool.READ_MEDIA) { + executionResult = await ReadMediaService.executeTool( + args, + { + audio: modelsStore.modelSupportsAudio(effectiveModel), + vision: modelsStore.modelSupportsVision(effectiveModel) + }, + signal, + conversationsStore.activeConversation?.cwd + ); + } else { + executionResult = await SandboxService.executeTool(toolName, args, signal); + } result = executionResult.content; diff --git a/tools/ui/src/lib/stores/tools.svelte.ts b/tools/ui/src/lib/stores/tools.svelte.ts index 642662928..4cfb9f370 100644 --- a/tools/ui/src/lib/stores/tools.svelte.ts +++ b/tools/ui/src/lib/stores/tools.svelte.ts @@ -1,5 +1,7 @@ import { browser } from '$app/environment'; import { + buildBrowserInfoToolDefinition, + buildGetDatetimeToolDefinition, buildReadMediaToolDefinition, DISABLED_TOOL_KEYS_LOCALSTORAGE_KEY, HOME_TILDE, @@ -175,7 +177,7 @@ class ToolsStore { } get frontendTools(): OpenAIToolDefinition[] { - const tools: OpenAIToolDefinition[] = []; + const tools: OpenAIToolDefinition[] = [buildGetDatetimeToolDefinition()]; if (settingsStore.config.jsSandboxEnabled) { tools.push(buildSandboxToolDefinition(!!settingsStore.config.symbolicMathEnabled)); @@ -185,9 +187,18 @@ class ToolsStore { if (readMedia) tools.push(readMedia); + // provide browser's get_info tool if server doesn't provide one + if (!this.hasBuiltinTool(BuiltInTool.GET_INFO)) { + tools.push(buildBrowserInfoToolDefinition()); + } + return tools; } + private hasBuiltinTool(name: BuiltInTool): boolean { + return this._builtinTools.some((def) => def.function.name === name); + } + /** * `read_media` runs in the frontend on top of the server's `read_file`, so it * exists only when that tool is served and the active model can perceive the @@ -195,11 +206,7 @@ class ToolsStore { * conversation uses. */ private readMediaTool(): OpenAIToolDefinition | null { - const hasReadFile = this._builtinTools.some( - (def) => def.function.name === BuiltInTool.READ_FILE - ); - - if (!hasReadFile) return null; + if (!this.hasBuiltinTool(BuiltInTool.READ_FILE)) return null; const model = modelsStore.selectedModelName ?? modelsStore.models[0]?.model ?? ''; diff --git a/tools/ui/src/lib/types/settings.d.ts b/tools/ui/src/lib/types/settings.d.ts index 55f304ea2..d04837727 100644 --- a/tools/ui/src/lib/types/settings.d.ts +++ b/tools/ui/src/lib/types/settings.d.ts @@ -31,6 +31,7 @@ export interface SettingsEntry { radioOptions?: Array<{ value: string; label: string; key: string; isExperimental?: boolean }>; isExperimental?: boolean; isPositiveInteger?: boolean; + isPrivate?: boolean; placeholder?: string; min?: number; max?: number; @@ -55,6 +56,7 @@ export interface SettingsFieldConfig { type: SettingsFieldType; isExperimental?: boolean; isPositiveInteger?: boolean; + isPrivate?: boolean; placeholder?: string; min?: number; max?: number; diff --git a/tools/ui/src/lib/utils/browser-info.ts b/tools/ui/src/lib/utils/browser-info.ts new file mode 100644 index 000000000..c96abb01e --- /dev/null +++ b/tools/ui/src/lib/utils/browser-info.ts @@ -0,0 +1,39 @@ +/** + * Browser fallback for the server's `get_info` tool, offered only when the + * server does not serve one (llama-server without --agent). It tells the model + * which OS the browser runs on and that there is no local file or shell access, + * so it does not plan around tools that are not there. + * + * @see server_tool_get_info in tools/server/server-tools.cpp - the served variant + * @see buildBrowserInfoToolDefinition in constants/browser-info.ts - tool schema sent to the LLM + */ + +import { browser } from '$app/environment'; +import { + BROWSER_INFO_NOTE, + BROWSER_INFO_OS_UA_PATTERNS, + BROWSER_INFO_OS_UNKNOWN +} from '$lib/constants'; +import type { ToolExecutionResult } from '$lib/types'; + +function detectOs(userAgent: string): string { + for (const [pattern, os] of BROWSER_INFO_OS_UA_PATTERNS) { + if (pattern.test(userAgent)) return os; + } + + return BROWSER_INFO_OS_UNKNOWN; +} + +/** + * Result shape mirrors the server tool's JSON so the `get_info` renderer reads + * `os` the same way, minus `cwd` - there is no working directory to report. + */ +export function executeBrowserInfoTool(): ToolExecutionResult { + return { + content: JSON.stringify({ + note: BROWSER_INFO_NOTE, + os: browser ? detectOs(navigator.userAgent) : BROWSER_INFO_OS_UNKNOWN + }), + isError: false + }; +} diff --git a/tools/ui/src/lib/utils/get-datetime.ts b/tools/ui/src/lib/utils/get-datetime.ts new file mode 100644 index 000000000..cd17bb1e8 --- /dev/null +++ b/tools/ui/src/lib/utils/get-datetime.ts @@ -0,0 +1,38 @@ +/** + * Frontend executor for the `get_datetime` tool. It runs in the browser, so it + * reports the user's own clock and time zone instead of the server's UTC time - + * a chat about "tomorrow" means the user's tomorrow, not the host's. + * + * @see buildGetDatetimeToolDefinition in constants/get-datetime.ts - tool schema sent to the LLM + */ + +import type { ToolExecutionResult } from '$lib/types'; + +function pad(value: number): string { + return String(value).padStart(2, '0'); +} + +/** ISO 8601 in local time, e.g. `2026-08-17T14:05:09+02:00` */ +function localIsoString(date: Date): string { + // getTimezoneOffset() counts minutes behind UTC, ISO 8601 counts them ahead + const offset = -date.getTimezoneOffset(); + const sign = offset < 0 ? '-' : '+'; + const absOffset = Math.abs(offset); + const day = `${date.getFullYear()}-${pad(date.getMonth() + 1)}-${pad(date.getDate())}`; + const time = `${pad(date.getHours())}:${pad(date.getMinutes())}:${pad(date.getSeconds())}`; + + return `${day}T${time}${sign}${pad(Math.floor(absOffset / 60))}:${pad(absOffset % 60)}`; +} + +/** The `result` field keeps the shape the `get_datetime` renderer already reads. */ +export function executeGetDatetimeTool(): ToolExecutionResult { + const now = new Date(); + + return { + content: JSON.stringify({ + result: localIsoString(now), + timezone: Intl.DateTimeFormat().resolvedOptions().timeZone + }), + isError: false + }; +} diff --git a/tools/ui/src/lib/utils/index.ts b/tools/ui/src/lib/utils/index.ts index 3fb701b32..dd997c206 100644 --- a/tools/ui/src/lib/utils/index.ts +++ b/tools/ui/src/lib/utils/index.ts @@ -331,6 +331,12 @@ export { getChatCommands } from './chat-commands'; // SANDBOX_TOOL_DEFINITION is deprecated; kept for backward compatibility. export { buildSandboxToolDefinition, SANDBOX_TOOL_DEFINITION } from './sandbox-tool'; +// Frontend `get_datetime` executor (the browser clock, not the server's) +export { executeGetDatetimeTool } from './get-datetime'; + +// Browser fallback for the server's get_info tool +export { executeBrowserInfoTool } from './browser-info'; + // Cryptography utilities export { uuid } from './uuid'; diff --git a/tools/ui/tests/unit/mcp-service.test.ts b/tools/ui/tests/unit/mcp-service.test.ts index 358881b47..97879f234 100644 --- a/tools/ui/tests/unit/mcp-service.test.ts +++ b/tools/ui/tests/unit/mcp-service.test.ts @@ -2,7 +2,7 @@ import { Client } from '@modelcontextprotocol/sdk/client'; import { CORS_PROXY } from '$lib/constants'; import { MCPConnectionPhase, MCPTransportType } from '$lib/enums'; import { MCPService } from '$lib/services/mcp.service'; -import type { MCPConnectionLog, MCPServerConfig } from '$lib/types'; +import type { MCPConnection, MCPConnectionLog, MCPServerConfig } from '$lib/types'; import { afterEach, describe, expect, it, vi } from 'vitest'; type DiagnosticFetchFactory = ( @@ -329,4 +329,21 @@ describe('MCPService', () => { ) ).toHaveLength(0); }); + + it('falls back to structuredContent when content array is empty', async () => { + const connection = { + client: { + callTool: vi.fn().mockResolvedValue({ + content: [], + structuredContent: { accounts: [{ id: 1 }], total: 1 } + }) + }, + requestTimeoutMs: 9000, + serverName: 'test-server' + } as unknown as MCPConnection; + const result = await MCPService.callTool(connection, { arguments: {}, name: 'tool' }); + + expect(result.isError).toBe(false); + expect(result.content).toBe('{"accounts":[{"id":1}],"total":1}'); + }); }); diff --git a/tools/ui/tests/unit/settings-private-fields.test.ts b/tools/ui/tests/unit/settings-private-fields.test.ts new file mode 100644 index 000000000..0716b3cba --- /dev/null +++ b/tools/ui/tests/unit/settings-private-fields.test.ts @@ -0,0 +1,12 @@ +import { SETTINGS_CHAT_SECTIONS, SETTINGS_KEYS } from '$lib/constants'; +import { describe, expect, it } from 'vitest'; + +describe('checkApiKeyField', () => { + it('should have isPrivate set to true', () => { + const fields = SETTINGS_CHAT_SECTIONS.flatMap((section) => section.fields); + const apiKeyField = fields.find((field) => field?.key === SETTINGS_KEYS.API_KEY); + + expect(apiKeyField).toBeDefined(); + expect(apiKeyField?.isPrivate).toBe(true); + }); +});