diff --git a/common/chat-peg-parser.cpp b/common/chat-peg-parser.cpp index a309f02765..f786f5ff23 100644 --- a/common/chat-peg-parser.cpp +++ b/common/chat-peg-parser.cpp @@ -1056,3 +1056,141 @@ void common_chat_peg_gemma4_mapper::visit(const common_peg_ast_arena & arena, co visit(arena, child_id); } } + +static void minimax_m3_collect(const common_peg_ast_arena & arena, + const common_peg_ast_node & node, + const std::string & tag, + std::vector & out) { + for (auto child_id : node.children) { + const auto & child = arena.get(child_id); + if (child.tag == tag) { + out.push_back(child_id); + } else { + minimax_m3_collect(arena, child, tag, out); + } + } +} + +static common_peg_ast_id minimax_m3_value_of(const common_peg_ast_arena & arena, const common_peg_ast_node & node) { + for (auto child_id : node.children) { + const auto & tag = arena.get(child_id).tag; + if (tag == common_chat_peg_builder::TOOL_ARG_VALUE || + tag == common_chat_peg_builder::TOOL_ARG_STRING_VALUE || + tag == common_chat_peg_minimax_m3_mapper::TOOL_ARG_OBJECT || + tag == common_chat_peg_minimax_m3_mapper::TOOL_ARG_ARRAY) { + return child_id; + } + } + return COMMON_PEG_INVALID_AST_ID; +} + +static std::string minimax_m3_value_to_json(const common_peg_ast_arena & arena, common_peg_ast_id id, bool closed); + +static std::string minimax_m3_member_to_json(const common_peg_ast_arena & arena, const common_peg_ast_node & node) { + auto name_id = arena.find_by_tag(node, common_chat_peg_builder::TOOL_ARG_NAME); + if (name_id == COMMON_PEG_INVALID_AST_ID) { + return ""; + } + + return ordered_json(arena.get(name_id).text).dump() + ":" + + minimax_m3_value_to_json(arena, minimax_m3_value_of(arena, node), !node.is_partial); +} + +static std::string minimax_m3_container_to_json(const common_peg_ast_arena & arena, + const common_peg_ast_node & node, + bool is_object, + bool closed) { + const std::string tag = is_object ? common_chat_peg_builder::TOOL_ARG + : common_chat_peg_minimax_m3_mapper::TOOL_ARG_ITEM; + + std::vector entries; + minimax_m3_collect(arena, node, tag, entries); + + std::string result = is_object ? "{" : "["; + + bool add_comma = false; + for (auto entry_id : entries) { + const auto & entry = arena.get(entry_id); + + std::string text; + if (is_object) { + text = minimax_m3_member_to_json(arena, entry); + } else { + text = minimax_m3_value_to_json(arena, minimax_m3_value_of(arena, entry), !entry.is_partial); + } + + if (text.empty()) { + continue; + } + + if (add_comma) { + result += ","; + } + add_comma = true; + result += text; + } + + if (closed) { + result += is_object ? "}" : "]"; + } + return result; +} + +static std::string minimax_m3_value_to_json(const common_peg_ast_arena & arena, common_peg_ast_id id, bool closed) { + if (id == COMMON_PEG_INVALID_AST_ID) { + return ""; + } + + const auto & node = arena.get(id); + + if (node.tag == common_chat_peg_minimax_m3_mapper::TOOL_ARG_OBJECT) { + return minimax_m3_container_to_json(arena, node, /* is_object = */ true, closed); + } + + if (node.tag == common_chat_peg_minimax_m3_mapper::TOOL_ARG_ARRAY) { + return minimax_m3_container_to_json(arena, node, /* is_object = */ false, closed); + } + + if (node.tag == common_chat_peg_builder::TOOL_ARG_STRING_VALUE) { + return "\"" + escape_json_string_inner(std::string(node.text)) + (closed ? "\"" : ""); + } + + // Numbers and booleans are written verbatim by the template + return std::string(node.text); +} + +void common_chat_peg_minimax_m3_mapper::from_ast(const common_peg_ast_arena & arena, + const common_peg_parse_result & result) { + for (const auto & node : result.nodes) { + visit(arena, node); + } +} + +void common_chat_peg_minimax_m3_mapper::visit(const common_peg_ast_arena & arena, common_peg_ast_id id) { + const auto & node = arena.get(id); + + if (node.tag == common_chat_peg_builder::REASONING) { + result.reasoning_content += std::string(node.text); + return; + } + + if (node.tag == common_chat_peg_builder::CONTENT) { + result.content += std::string(node.text); + return; + } + + if (node.tag == common_chat_peg_builder::TOOL) { + auto name_id = arena.find_by_tag(node, common_chat_peg_builder::TOOL_NAME); + if (name_id != COMMON_PEG_INVALID_AST_ID) { + common_chat_tool_call call; + call.name = std::string(arena.get(name_id).text); + call.arguments = minimax_m3_container_to_json(arena, node, /* is_object = */ true, !node.is_partial); + result.tool_calls.push_back(call); + } + return; + } + + for (auto child_id : node.children) { + visit(arena, child_id); + } +} diff --git a/common/chat-peg-parser.h b/common/chat-peg-parser.h index b3ffd7de2d..cd14f2c117 100644 --- a/common/chat-peg-parser.h +++ b/common/chat-peg-parser.h @@ -40,6 +40,18 @@ class common_chat_peg_gemma4_mapper : public common_chat_peg_mapper { void visit(const common_peg_ast_arena & arena, common_peg_ast_id id); }; +class common_chat_peg_minimax_m3_mapper : public common_chat_peg_mapper { + public: + static constexpr const char * TOOL_ARG_OBJECT = "tool-arg-object"; + static constexpr const char * TOOL_ARG_ARRAY = "tool-arg-array"; + static constexpr const char * TOOL_ARG_ITEM = "tool-arg-item"; + + common_chat_peg_minimax_m3_mapper(common_chat_msg & msg) : common_chat_peg_mapper(msg) {} + virtual void from_ast(const common_peg_ast_arena & arena, const common_peg_parse_result & result); + private: + void visit(const common_peg_ast_arena & arena, common_peg_ast_id id); +}; + struct content_structure; struct tool_call_structure; diff --git a/common/chat.cpp b/common/chat.cpp index 7a6e7238cf..7740f35c0e 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -816,6 +816,8 @@ const char * common_chat_format_name(common_chat_format format) { return "peg-native"; case COMMON_CHAT_FORMAT_PEG_GEMMA4: return "peg-gemma4"; + case COMMON_CHAT_FORMAT_PEG_MINIMAX_M3: + return "peg-minimax-m3"; default: throw std::runtime_error("Unknown chat format"); } @@ -2270,6 +2272,264 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t return data; } +static common_chat_params common_chat_params_init_minimax_m3(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_MINIMAX_M3; + data.supports_thinking = true; + data.thinking_start_tag = ""; + data.thinking_end_tags = {""}; + + // M3 prefixes every tool tag with the namespace token "]<]minimax[>["; + // params use the parameter name as the tag (...). + const std::string NS = "]<]minimax[>["; + const std::string THINK_START = ""; + const std::string THINK_END = ""; + const std::string FC_START = NS + ""; + const std::string FC_END = NS + ""; + const std::string INVOKE_END = NS + ""; + + data.preserved_tokens = { + NS, + "", + "", + THINK_START, + THINK_END, + }; + + data.message_delimiters = { + { COMMON_CHAT_ROLE_ASSISTANT, "]~b]ai" }, + { COMMON_CHAT_ROLE_USER, "]~b]user" }, + { COMMON_CHAT_ROLE_TOOL, "]~b]tool" }, + { COMMON_CHAT_ROLE_SYSTEM, "]~b]developer" }, + { COMMON_CHAT_ROLE_SYSTEM, "]~b]system" }, + }; + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object(); + auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + auto include_grammar = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE); + + const std::string GEN_PROMPT = data.generation_prompt; + + using mm3 = common_chat_peg_minimax_m3_mapper; + + if (inputs.has_continuation()) { + const auto & msg = inputs.continue_msg; + + data.generation_prompt = GEN_PROMPT + THINK_START + msg.reasoning_content; + if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + data.generation_prompt += THINK_END + msg.render_content(); + } + + data.prompt += data.generation_prompt; + } + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + auto generation_prompt = p.prefix(GEN_PROMPT, THINK_START); + auto end = p.end(); + + auto reasoning = p.eps(); + if (extract_reasoning) { + auto block = inputs.enable_thinking + ? p.literal(THINK_START) + p.space() + + p.ac(p.reasoning(p.until(THINK_END)) + p.literal(THINK_END), THINK_END) + : p.literal(THINK_START) + p.ac(p.until(THINK_END) + p.literal(THINK_END), THINK_END); + + // A turn without reasoning is prefixed with a bare , written either by the + // generation prompt (thinking_mode = "disabled") or by the model itself. + reasoning = p.optional(p.choice({ block, p.literal(THINK_END) })); + } + + if (has_response_format) { + auto response_format = p.rule("response-format", + p.literal("```json") + p.space() + + p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema)) + + p.space() + p.literal("```")); + return generation_prompt + reasoning + response_format + end; + } + + if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { + return generation_prompt + reasoning + p.content(p.rest()) + end; + } + + auto alternatives_of = [](const json & schema) -> std::optional { + for (const auto * keyword : { "oneOf", "anyOf" }) { + if (schema.contains(keyword) && schema.at(keyword).is_array() && !schema.at(keyword).empty()) { + return schema.at(keyword); + } + } + return std::nullopt; + }; + + auto tool_choice = p.choice(); + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + std::string name = function.at("name"); + auto params = function.contains("parameters") ? function.at("parameters") : json::object(); + + auto schema_info = common_schema_info(); + schema_info.resolve_refs(params); + + // The template expands argument values recursively in XML (see the to_xml() macro) + std::function value_of; + std::function members_of; + + auto element_of = [&](const std::string & tag, const json & schema, const std::string & rule_name) { + const std::string close = NS + ""; + return p.rule(rule_name, + p.tool_arg( + p.tool_arg_open( + p.literal(NS + "<") + + p.tool_arg_name(p.literal(tag)) + + p.literal(">")) + + value_of(schema, rule_name, close))); + }; + + value_of = [&](const json & schema, + const std::string & rule_name, + const std::string & close) -> common_peg_parser { + auto close_tag = p.tool_arg_close(p.literal(close)); + + // A string accepts anything, so a union with a string alternative is a string + if (schema_info.resolves_to_string(schema)) { + return p.ac(p.tool_arg_string_value(p.until(close)) + close_tag, close); + } + + if (auto alternatives = alternatives_of(schema)) { + std::vector choices; + + size_t index = 0; + for (const auto & alternative : *alternatives) { + const std::string alt_name = rule_name + "-" + std::to_string(index++); + + // There is a risk that this breaks streaming deltas, but that's a risk we + // assume to provide tool arg streaming. + choices.push_back(value_of(alternative, alt_name, close)); + } + + return p.choice(choices); + } + + const std::string type = schema.contains("type") && schema.at("type").is_string() + ? schema.at("type").get() + : ""; + + if (type == "object" && schema.contains("properties")) { + return p.tag(mm3::TOOL_ARG_OBJECT, members_of(schema, rule_name)) + p.space() + close_tag; + } + + if (type == "array" && schema.contains("items")) { + const std::string item_close = NS + ""; + auto item = p.rule(rule_name + "-item", + p.tag(mm3::TOOL_ARG_ITEM, + p.literal(NS + "") + + value_of(schema.at("items"), rule_name + "-item", item_close))); + return p.tag(mm3::TOOL_ARG_ARRAY, p.repeat(p.space() + item, 0, -1)) + p.space() + close_tag; + } + + return p.tool_arg_json_value(p.schema(p.json(), rule_name + "-schema", schema, false)) + close_tag; + }; + + // Required properties in schema order, then any number of optional ones in any order. + members_of = [&](const json & schema, const std::string & rule_prefix) -> common_peg_parser { + const auto & props = schema.at("properties"); + + std::set required; + if (schema.contains("required")) { + schema.at("required").get_to(required); + } + + std::vector required_elements; + std::vector optional_elements; + for (const auto & [key, key_schema] : props.items()) { + auto element = element_of(key, key_schema, rule_prefix + "-" + key); + if (required.find(key) != required.end()) { + required_elements.push_back(element); + } else { + optional_elements.push_back(element); + } + } + + common_peg_parser members = p.eps(); + for (size_t i = 0; i < required_elements.size(); i++) { + if (i > 0) { + members = members + p.space(); + } + members = members + required_elements[i]; + } + + if (!optional_elements.empty()) { + common_peg_parser any_optional = p.choice(); + for (const auto & element : optional_elements) { + any_optional |= element; + } + members = members + p.repeat(p.space() + any_optional, 0, -1); + } + + return members; + }; + + common_peg_parser invoke_body = + params.contains("properties") ? members_of(params, "tool-" + name + "-arg") : p.eps(); + + auto func_parser = p.tool( + p.tool_open(p.literal(NS + "")) + + p.space() + invoke_body + p.space() + + p.tool_close(p.literal(INVOKE_END))); + + tool_choice |= p.rule("tool-" + name, func_parser); + }); + + auto require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; + + common_peg_parser tool_calls = p.eps(); + if (inputs.parallel_tool_calls) { + tool_calls = p.trigger_rule("tool-call", + p.literal(FC_START) + p.space() + tool_choice + + p.zero_or_more(p.space() + tool_choice) + p.space() + p.literal(FC_END)); + } else { + tool_calls = p.trigger_rule("tool-call", + p.literal(FC_START) + p.space() + tool_choice + p.space() + p.literal(FC_END)); + } + + if (!require_tools) { + tool_calls = p.optional(tool_calls); + } + + auto content_before_tools = p.content(p.until(FC_START)); + return generation_prompt + reasoning + content_before_tools + tool_calls + end; + }); + + data.parser = parser.save(); + + if (include_grammar) { + data.grammar_lazy = !(has_response_format || (has_tools && 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"); + auto schema = function.contains("parameters") ? function.at("parameters") : json::object(); + builder.resolve_refs(schema); + }); + if (has_response_format) { + auto schema = inputs.json_schema; + builder.resolve_refs(schema); + } + parser.build_grammar(builder, data.grammar_lazy); + }); + + data.grammar_triggers = { + { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, FC_START }, + }; + } + + return data; +} + namespace workaround { static void map_developer_role_to_system(json & messages) { @@ -2707,6 +2967,15 @@ std::optional common_chat_try_specialized_template( return common_chat_params_init_gigachat_v3(tmpl, params); } + // MiniMax-M3: the namespace token "]<]minimax[>[" collides with the autoparser's + // markup delimiters, so detect the template and use a dedicated parser. + if (src.find("]<]minimax[>[") != std::string::npos && + src.find("") != std::string::npos && + src.find(" mapper; if (params.format == COMMON_CHAT_FORMAT_PEG_GEMMA4) { mapper = std::make_unique(msg); + } else if (params.format == COMMON_CHAT_FORMAT_PEG_MINIMAX_M3) { + mapper = std::make_unique(msg); } else { mapper = std::make_unique(msg); } @@ -3020,6 +3291,8 @@ common_chat_msg common_chat_peg_parse(const common_peg_arena & src_pars std::unique_ptr mapper; if (params.format == COMMON_CHAT_FORMAT_PEG_GEMMA4) { mapper = std::make_unique(msg); + } else if (params.format == COMMON_CHAT_FORMAT_PEG_MINIMAX_M3) { + mapper = std::make_unique(msg); } else { mapper = std::make_unique(msg); } diff --git a/common/chat.h b/common/chat.h index d79f4ecd77..6d5b220aeb 100644 --- a/common/chat.h +++ b/common/chat.h @@ -233,6 +233,7 @@ enum common_chat_format { COMMON_CHAT_FORMAT_PEG_SIMPLE, COMMON_CHAT_FORMAT_PEG_NATIVE, COMMON_CHAT_FORMAT_PEG_GEMMA4, + COMMON_CHAT_FORMAT_PEG_MINIMAX_M3, COMMON_CHAT_FORMAT_COUNT, // Not a format, just the # formats }; diff --git a/common/common.cpp b/common/common.cpp index 82dd780fd8..9d5d2834ac 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1518,23 +1518,49 @@ done: return res; } -void common_context_seq_rm(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) { +static void common_context_seq_rm(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) { auto * mem = llama_get_memory(ctx); if (!llama_memory_seq_rm(mem, seq_id, p0, p1)) { GGML_ABORT("%s", string_format("failed to remove sequence %d with p0=%d, p1=%d\n", seq_id, p0, p1).c_str()); } } -void common_context_seq_cp(llama_context * ctx, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { +static void common_context_seq_cp(llama_context * ctx, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { auto * mem = llama_get_memory(ctx); llama_memory_seq_cp(mem, seq_id_src, seq_id_dst, p0, p1); } -void common_context_seq_add(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) { +static void common_context_seq_add(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) { auto * mem = llama_get_memory(ctx); llama_memory_seq_add(mem, seq_id, p0, p1, delta); } +void common_memory::init(llama_context * ctx_tgt, llama_context * ctx_dft) { + this->ctx_tgt = ctx_tgt; + this->ctx_dft = ctx_dft; +} + +void common_memory::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) const { + common_context_seq_rm(ctx_tgt, seq_id, p0, p1); + if (ctx_dft) { + common_context_seq_rm(ctx_dft, seq_id, p0, p1); + } +} + +void common_memory::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) const { + common_context_seq_cp(ctx_tgt, seq_id_src, seq_id_dst, p0, p1); + if (ctx_dft) { + common_context_seq_cp(ctx_dft, seq_id_src, seq_id_dst, p0, p1); + } +} + +void common_memory::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) const { + common_context_seq_add(ctx_tgt, seq_id, p0, p1, delta); + if (ctx_dft) { + common_context_seq_add(ctx_dft, seq_id, p0, p1, delta); + } +} + void common_set_adapter_lora(struct llama_context * ctx, std::vector & lora) { std::vector loras; std::vector scales; diff --git a/common/common.h b/common/common.h index 2792521836..e7c55ae925 100644 --- a/common/common.h +++ b/common/common.h @@ -173,6 +173,7 @@ enum common_speculative_type { COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, // Eagle3 speculative decoding COMMON_SPECULATIVE_TYPE_DRAFT_MTP, // Multi-token prediction COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, // DFlash speculative decoding + COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK, // DSpark speculative decoding (DFlash + Markov head) COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, // simple self-speculative decoding based on n-grams COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, // self-speculative decoding with n-gram keys only COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, // self-speculative decoding with n-gram keys and 4 m-gram values @@ -388,7 +389,7 @@ struct common_params_speculative { uint32_t need_n_rs_seq() const { bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) { - return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH; + return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK; }); return needs_rs_seq ? draft.n_max : 0u; @@ -948,10 +949,17 @@ enum common_context_seq_rm_type { // note: clears the memory of the context common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx); -// aborts execution on failure -void common_context_seq_rm (llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1); -void common_context_seq_add(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta); -void common_context_seq_cp (llama_context * ctx, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1); +struct common_memory { + llama_context * ctx_tgt = nullptr; + llama_context * ctx_dft = nullptr; + + void init(llama_context * ctx_tgt, llama_context * ctx_dft = nullptr); + + // aborts execution on failure + void seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) const; + void seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) const; + void seq_cp (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) const; +}; // // Batch utils diff --git a/common/speculative.cpp b/common/speculative.cpp index 3a6c073683..5653a90b88 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -34,6 +34,7 @@ const std::map common_speculative_type_fro {"draft-eagle3", COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3}, {"draft-mtp", COMMON_SPECULATIVE_TYPE_DRAFT_MTP}, {"draft-dflash", COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH}, + {"draft-dspark", COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK}, {"ngram-simple", COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE}, {"ngram-map-k", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K}, {"ngram-map-k4v", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V}, @@ -928,15 +929,20 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { int32_t block_size = 0; llama_token mask_token_id = 0; + // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout + const bool is_dspark; + const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices uint32_t target_layer_ids_n = 0; // scratch buffer for concatenated target features [n_tokens, n_embd_enc] std::vector features_buf; - common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, n_seq) + common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq, + common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) + : common_speculative_impl(type, n_seq) , params(params.draft) + , is_dspark(type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) { auto * ctx_tgt = this->params.ctx_tgt; auto * ctx_dft = this->params.ctx_dft; @@ -963,16 +969,18 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft)); - LOG_INF("%s: adding speculative implementation 'draft-dflash'\n", __func__); + 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); - // DFlash input is [id_last, * (block_size-1)], so it can draft at most block_size-1 tokens per step - if (this->params.n_max > block_size - 1 || this->params.n_min > block_size - 1) { - LOG_WRN("%s: requested draft size (n_max=%d, n_min=%d) exceeds the trained DFlash block size %d -- clamping to %d\n", - __func__, this->params.n_max, this->params.n_min, block_size, block_size - 1); - this->params.n_max = std::min(this->params.n_max, block_size - 1); - this->params.n_min = std::min(this->params.n_min, block_size - 1); + // 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; + 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); + this->params.n_max = std::min(this->params.n_max, n_draft_max); + this->params.n_min = std::min(this->params.n_min, n_draft_max); } batch = llama_batch_init(llama_n_batch(ctx_dft), 0, n_seq); @@ -1136,12 +1144,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { const int32_t n = (int32_t) dp.n_past; - int32_t n_draft = params.n_max; - if (dp.n_max > 0) { - n_draft = std::min(n_draft, dp.n_max); - } + const int32_t n_draft = params.n_max; - const int32_t n_block_tokens = n_draft + 1; // id_last + n_draft * + const int32_t n_block_tokens = n_draft + (is_dspark ? 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) { @@ -1173,27 +1178,57 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { auto & result = *dp.result; - // greedily read the predicted block at this sequence's noise positions 1..n_block_tokens-1 - for (int32_t i = 1; i < n_block_tokens; ++i) { - common_sampler_sample(smpl, ctx_dft, beg + i, true); + if (is_dspark) { + // DSpark predicts the next token from position 0 and optionally truncates + // at the first position below the confidence threshold. + const float * conf = params.p_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr; - const auto * cur_p = common_sampler_get_candidates(smpl, true); + for (int32_t i = 0; i < n_block_tokens; ++i) { + const int32_t idx = beg + i; - for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { - LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", - seq_id, k, i - 1, cur_p->data[k].id, cur_p->data[k].p, - common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); + if (conf && conf[(size_t) idx * n_embd_dec] < params.p_min) { + break; + } + + common_sampler_sample(smpl, ctx_dft, idx, true); + + const auto * cur_p = common_sampler_get_candidates(smpl, true); + + for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { + LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", + seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p, + common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); + } + + const llama_token id = cur_p->data[0].id; + + common_sampler_accept(smpl, id, true); + + result.push_back(id); } + } else { + // greedily read the predicted block at this sequence's noise positions 1..n_block_tokens-1 + for (int32_t i = 1; i < n_block_tokens; ++i) { + common_sampler_sample(smpl, ctx_dft, beg + i, true); - const llama_token id = cur_p->data[0].id; + const auto * cur_p = common_sampler_get_candidates(smpl, true); - if (cur_p->data[0].p < params.p_min) { - break; + for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { + LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", + seq_id, k, i - 1, cur_p->data[k].id, cur_p->data[k].p, + common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); + } + + const llama_token id = cur_p->data[0].id; + + if (cur_p->data[0].p < params.p_min) { + break; + } + + common_sampler_accept(smpl, id, true); + + result.push_back(id); } - - common_sampler_accept(smpl, id, true); - - result.push_back(id); } if (result.size() < (size_t) params.n_min) { @@ -2155,6 +2190,7 @@ std::string common_speculative_type_to_str(common_speculative_type type) { case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3: return "draft-eagle3"; case COMMON_SPECULATIVE_TYPE_DRAFT_MTP: return "draft-mtp"; case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH: return "draft-dflash"; + case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK: return "draft-dspark"; case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: return "ngram-simple"; case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: return "ngram-map-k"; case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: return "ngram-map-k4v"; @@ -2208,6 +2244,7 @@ int32_t common_speculative_n_max(const common_params_speculative * spec) { case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3: case COMMON_SPECULATIVE_TYPE_DRAFT_MTP: case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH: + case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK: n_max = std::max(n_max, std::max(0, spec->draft.n_max)); break; case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: @@ -2352,6 +2389,7 @@ common_speculative * common_speculative_init(common_params_speculative & params, bool has_draft_eagle3 = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3)) && params.draft.ctx_dft != nullptr; bool has_draft_mtp = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_MTP)) && params.draft.ctx_dft != nullptr; bool has_draft_dflash = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)) && params.draft.ctx_dft != nullptr; + bool has_draft_dspark = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK)) && params.draft.ctx_dft != nullptr; @@ -2362,7 +2400,7 @@ common_speculative * common_speculative_init(common_params_speculative & params, bool has_ngram_mod = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_NGRAM_MOD)); // when adding a new type - update here the logic above - static_assert(COMMON_SPECULATIVE_TYPE_COUNT == 10); + static_assert(COMMON_SPECULATIVE_TYPE_COUNT == 11); // this list here defines the priority of the speculators // the one with highest priority are listed first @@ -2395,6 +2433,9 @@ common_speculative * common_speculative_init(common_params_speculative & params, if (has_draft_dflash) { configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, params)); } + if (has_draft_dspark) { + configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK, params)); + } } std::vector> impls = {}; @@ -2419,6 +2460,11 @@ common_speculative * common_speculative_init(common_params_speculative & params, impls.push_back(std::make_unique(config.params, n_seq)); break; } + case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK: { + impls.push_back(std::make_unique( + config.params, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK)); + break; + } case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: { common_ngram_map ngram_map = get_common_ngram_map(config.type, config.params.ngram_simple); diff --git a/conversion/__init__.py b/conversion/__init__.py index c89d734902..31da4963cf 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -53,6 +53,7 @@ TEXT_MODEL_MAP: dict[str, str] = { "DeepseekV3ForCausalLM": "deepseek", "DeepseekV32ForCausalLM": "deepseek", "DFlashDraftModel": "qwen", + "Qwen3DSparkModel": "qwen", "DeepseekV4ForCausalLM": "deepseek", "DistilBertForMaskedLM": "bert", "DistilBertForSequenceClassification": "bert", diff --git a/conversion/nemotron.py b/conversion/nemotron.py index e44688a788..0572b42ca2 100644 --- a/conversion/nemotron.py +++ b/conversion/nemotron.py @@ -39,28 +39,48 @@ class NemotronNanoV2VLModel(MmprojModel): } return vision_config + def get_audio_config(self) -> dict[str, Any] | None: + return self.global_config.get("sound_config") + def set_gguf_parameters(self): if "image_mean" not in self.preprocessor_config: self.preprocessor_config["image_mean"] = [0.485, 0.456, 0.406] if "image_std" not in self.preprocessor_config: self.preprocessor_config["image_std"] = [0.229, 0.224, 0.225] + if self.hparams_audio is not None: + self.has_vision_encoder = True + self.has_audio_encoder = True + self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["num_mel_bins"]) + self.gguf_writer.add_audio_attention_layernorm_eps(1e-5) + self.gguf_writer.add_audio_subsampling_factor(self.hparams_audio["subsampling_factor"]) + self.gguf_writer.add_audio_conv_kernel_size(self.hparams_audio["conv_kernel_size"]) + self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.PARAKEET) + self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL) + else: + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL) + super().set_gguf_parameters() hparams = self.global_config - self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL) self.gguf_writer.add_vision_attention_layernorm_eps(1e-6) self.gguf_writer.add_vision_use_gelu(True) downsample_ratio = hparams.get("downsample_ratio", 0.5) self.gguf_writer.add_vision_projector_scale_factor(int(1.0 / downsample_ratio)) def tensor_force_quant(self, name, new_name, bid, n_dims): - if ".position_embd." in new_name or "pos_embed" in new_name: - return gguf.GGMLQuantizationType.F32 + if "sound_encoder" in name or new_name.startswith("mm.a."): + if "bias" in new_name or "norm" in new_name: + return gguf.GGMLQuantizationType.F32 + if "conv" in new_name and "weight" in new_name: + return gguf.GGMLQuantizationType.F32 + return super().tensor_force_quant(name, new_name, bid, n_dims) @classmethod def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: - name, gen = item + if (titem := super().filter_tensors(item)) is None: + return None + name, gen = titem if "input_conditioner" in name: return None @@ -69,14 +89,18 @@ class NemotronNanoV2VLModel(MmprojModel): if "radio_model.model.patch_generator.video_embedder" in name: return None - if not name.startswith("vision_model.radio_model.model.") and not name.startswith("mlp1."): + if not name.startswith(("vision_model.radio_model.model.", "mlp1.", "sound_encoder.", "sound_projection.")): return None if "patch_generator.pos_embed" in name: if not name.endswith(".weight"): name += ".weight" - return super().filter_tensors((name, gen)) + # num_batches is only used for training not inference. + if "conv.norm" in name and "num_batches" in name: + return None + + return name, gen def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: # RADIO's pos_embed doesn't have .weight suffix, but clip.cpp expects it @@ -104,7 +128,26 @@ class NemotronNanoV2VLModel(MmprojModel): n_embd = self.hparams["hidden_size"] data_torch = data_torch.reshape(n_embd, 3, patch_size, patch_size) - yield from super().modify_tensors(data_torch, name, bid) + if "depthwise_conv.weight" in name: + data_torch = data_torch.unsqueeze(-1) + data_torch = data_torch.permute(3, 1, 0, 2).contiguous() + + if "pointwise_conv" in name and name.endswith(".weight"): + if len(data_torch.shape) == 3 and data_torch.shape[2] == 1: + data_torch = data_torch.reshape(data_torch.shape[0], data_torch.shape[1]) + + if "subsampling.layers" in name and name.endswith(".bias"): + if len(data_torch.shape) == 1: + data_torch = data_torch.reshape(1, -1, 1, 1) + + if "pointwise_conv" in name and name.endswith(".bias"): + if len(data_torch.shape) == 1: + data_torch = data_torch.reshape(1, -1, 1, 1) + + for mapped_name, tensor in super().modify_tensors(data_torch, name, bid): + if name.startswith("sound_projection.") and mapped_name.startswith("mm.model.mlp."): + mapped_name = mapped_name.replace("mm.model.mlp.", "mm.a.mlp.") + yield mapped_name, tensor @ModelBase.register("NemotronForCausalLM") diff --git a/conversion/qwen.py b/conversion/qwen.py index 9bc2b99fde..d1127f7431 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -688,3 +688,23 @@ class DFlashModel(Qwen3Model): if not name.startswith("model."): name = "model." + name return super().filter_tensors((name, gen)) + + +@ModelBase.register("Qwen3DSparkModel") +class DSparkModel(DFlashModel): + # 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 + }) + + @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")): + return None + return super().filter_tensors((name, gen)) diff --git a/docs/install.md b/docs/install.md index 7198e61bf3..b36b0be267 100644 --- a/docs/install.md +++ b/docs/install.md @@ -16,22 +16,22 @@ conda-forge provides builds for: - Apple Metal (macOS) ```sh -conda install -c conda-forge llama-cpp +conda install -c conda-forge llama.cpp ``` ```sh -mamba install -c conda-forge llama-cpp +mamba install -c conda-forge llama.cpp ``` ```sh # Project-local installation -pixi add llama-cpp +pixi add llama.cpp # Global installation -pixi global install llama-cpp +pixi global install llama.cpp ``` -This distribution is managed on [`conda-forge/llama-cpp-feedstock`](https://github.com/conda-forge/llama.cpp-feedstock/). +This distribution is managed on [`conda-forge/llama.cpp-feedstock`](https://github.com/conda-forge/llama.cpp-feedstock/). Shall you have any problems, please open an issue on [its issue tracker](https://github.com/conda-forge/llama.cpp-feedstock/issues). diff --git a/docs/speculative.md b/docs/speculative.md index 4100b92f8f..3957db85c9 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -78,6 +78,38 @@ See: - #22105 +### DSpark (`draft-dspark`) + +DSpark extends DFlash with a semi-autoregressive _Markov head_: the draft still emits a whole +block per forward pass, but each block position's logits are biased by a low-rank term keyed on +the previous token, chained in-graph across the block. This keeps drafting at one decode per +block while recovering some of the left-to-right signal that pure block diffusion loses. + +The draft is a small DeepSpec checkpoint trained for a specific target (for example +[`deepseek-ai/dspark_qwen3_4b_block7`](https://huggingface.co/deepseek-ai/dspark_qwen3_4b_block7) +for `Qwen/Qwen3-4B`). Convert it with `--target-model-dir` so it inherits the target's tokenizer +and token embeddings: + +```bash +python convert_hf_to_gguf.py deepseek-ai/dspark_qwen3_4b_block7 \ + --target-model-dir Qwen/Qwen3-4B --outtype bf16 --outfile Qwen3-4B-DSpark.gguf + +llama-server -m Qwen3-4B.gguf -md Qwen3-4B-DSpark.gguf \ + --spec-type draft-dspark --spec-draft-n-max 7 -fa on --jinja +``` + +`--spec-draft-n-max` is clamped to the draft model's trained block size. + +`--spec-draft-conf-min P` truncates each drafted block at the first position whose predicted +acceptance (from the draft's confidence head, if present) falls below `P` (default 0 = disabled). + +Currently only drafts with a Qwen3 backbone are supported; support for other backbones +(e.g. Gemma4) is planned. + +See: + +- #25173 + ### n-gram Cache (`ngram-cache`) An n-gram is a sequence of n tokens. The n-gram cache implementation maintains statistics about short n-gram sequences. @@ -173,7 +205,7 @@ If a draft model is combined with a draftless decoding the draftless decoding ha ### General Speculative Parameters ``` ---spec-type [none|draft-simple|draft-eagle3|draft-dflash|draft-mtp|ngram-cache|ngram-simple|ngram-map-k|ngram-map-k4v|ngram-mod] +--spec-type [none|draft-simple|draft-eagle3|draft-dflash|draft-dspark|draft-mtp|ngram-cache|ngram-simple|ngram-map-k|ngram-map-k4v|ngram-mod] comma-separated list of types of speculative decoding to use (default: none) (env: LLAMA_ARG_SPEC_TYPE) @@ -314,6 +346,7 @@ Specifies a comma-separated list of speculative decoding types to use. | `draft-simple` | Use a simple draft model for speculation | | `draft-eagle3` | Use an EAGLE-3 draft model that reads the target's hidden states | | `draft-dflash` | Use a DFlash block-diffusion draft model that emits a block per step | +| `draft-dspark` | Use a DSpark draft model (DFlash backbone + semi-autoregressive Markov head) | | `draft-mtp` | Use Multi Token Prediction (MTP) heads from the main model | | `ngram-cache` | Use n-gram cache lookup | | `ngram-simple` | Use simple n-gram pattern matching | diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index 3022249c77..f3418c2af8 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -9,6 +9,21 @@ using namespace cub; #include "ssm-scan.cuh" + +// Minimum number of tokens to use SSD (State Space Duality) matmul path instead of scan path. +// For n_tok <= this threshold, the scan kernel is used (lower overhead for short sequences). +#define SSM_SSD_MIN_TOKENS 128 + +// prepare_dt kernel dimensions: one block per (head, seq), each block handles DT_MAX_ITEMS items. +#define SSM_SSD_DT_BLOCK 256 +#define SSM_SSD_DT_MAX_ITEMS 32 + +// Maximum tokens the SSD path supports, derived from the prepare_dt kernel block capacity. +#define SSM_SSD_MAX_TOKENS (SSM_SSD_DT_BLOCK * SSM_SSD_DT_MAX_ITEMS) + +// Chunk size for chunked SSD. Caps matmul cost at O(chunk^2) per chunk. +#define SSM_SSD_CHUNK_SIZE 256 + // We would like to keep pragma unroll for cases where L_template is not 0, // so we suppress the clang transformation warning. #ifdef __clang__ @@ -316,6 +331,429 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa } } +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +// ============================================================================ +// SSD (State Space Duality) kernels for Mamba-2 prefill (n_tok > SSM_SSD_MIN_TOKENS) +// +// Instead of a sequential scan, SSD reformulates the output as: +// Y = (L (.) (C @ B^T)) @ (X * dt) + decay * C @ s_init +// where L is a causal decay mask derived from A and dt. +// +// This converts the O(T*N) sequential scan into parallel matmuls. +// ============================================================================ +// Softplus(dt) and inclusive prefix sum per head using CUB BlockScan. +// Grid: (n_head, n_seqs) +template +__global__ void ssm_ssd_prepare_dt_kernel( + const float * __restrict__ dt_raw, + float * __restrict__ dt_sp_out, + float * __restrict__ cs_out, + const int n_head, const int n_tok, + const int dt_stride_tok, // elements between tokens in dt + const int dt_stride_seq) { // elements between sequences in dt + + const int h = blockIdx.x; + const int s = blockIdx.y; + + const float * dt_seq = dt_raw + s * dt_stride_seq; + + float * dt_sp_seq = dt_sp_out + s * n_tok * n_head; + float * cs_seq = cs_out + s * n_tok * n_head; + + const int items_per_thread = (n_tok + BLOCK_SIZE - 1) / BLOCK_SIZE; + + // Phase 1: softplus with interleaved distribution (t = i*BLOCK_SIZE + threadIdx.x). + // Each warp reads BLOCK_SIZE consecutive tokens, giving coalesced dt_raw loads + // (stride n_head between threads vs. items_per_thread*n_head in blocked layout). + float local_vals[MAX_ITEMS]; + for (int i = 0; i < items_per_thread; i++) { + const int t = i * BLOCK_SIZE + threadIdx.x; + if (t < n_tok) { + float val = dt_seq[h + t * dt_stride_tok]; + float sp = (val <= 20.0f) ? log1pf(expf(val)) : val; + local_vals[i] = sp; + dt_sp_seq[t * n_head + h] = sp; + } else { + local_vals[i] = 0.0f; + } + } + + // Phase 2+3: per-step inclusive scan to build cs[] in token order. + // With interleaved distribution the per-thread total scan would not give token-order + // prefix sums, so we scan one BLOCK_SIZE slab at a time and carry a running total. +#ifdef USE_CUB + using BlockScan = cub::BlockScan; + __shared__ typename BlockScan::TempStorage scan_temp; + __shared__ float step_total; + + float running = 0.0f; + for (int i = 0; i < items_per_thread; i++) { + float inclusive; + BlockScan(scan_temp).InclusiveSum(local_vals[i], inclusive); + const int t = i * BLOCK_SIZE + threadIdx.x; + if (t < n_tok) { + cs_seq[t * n_head + h] = running + inclusive; + } + if (threadIdx.x == BLOCK_SIZE - 1) { + step_total = inclusive; + } + __syncthreads(); + running += step_total; + } +#else + // Fallback: sequential prefix scan in shared memory, one slab at a time. + __shared__ float sdata[BLOCK_SIZE]; + float running = 0.0f; + for (int i = 0; i < items_per_thread; i++) { + const int t = i * BLOCK_SIZE + threadIdx.x; + sdata[threadIdx.x] = local_vals[i]; + __syncthreads(); + if (threadIdx.x == 0) { + for (int j = 1; j < BLOCK_SIZE; j++) { + sdata[j] += sdata[j - 1]; + } + } + __syncthreads(); + if (t < n_tok) { + cs_seq[t * n_head + h] = running + sdata[threadIdx.x]; + } + running += sdata[BLOCK_SIZE - 1]; + __syncthreads(); + } +#endif +} + +// Prepare SSD matmul inputs for one chunk: X_dt, B_weighted, C_scaled. +// T_matmul controls precision for X_dt, B_weighted (float or half). +// C_scaled is always float (pairs with float s_cur in step 3c). +// Computation is always FP32; only the final store converts to T_matmul. +// Also materializes the causal M matrix = exp(A*(cs_out - cs_in)) * CB (fused with prep to save a launch). +// Grid: (ceil(max(C*head_dim, d_state*C, chunk_len^2) / BLOCK), n_head, n_seqs) +template +__global__ void ssm_ssd_pre_matmul_kernel( + const float * __restrict__ cs, // {n_tok, n_head} cumulative dt sums + const float * __restrict__ dt_sp, // {n_tok, n_head} softplus(dt) + const float * __restrict__ A, // {1, n_head} + const float * __restrict__ x, // {head_dim, n_head, n_tok, n_seqs} + const float * __restrict__ B, // {d_state, n_group, n_tok, n_seqs} + const float * __restrict__ C_src, // {d_state, n_group, n_tok, n_seqs} + T_matmul * __restrict__ X_dt, // {head_dim, C, n_head} x * dt, d-fastest + T_matmul * __restrict__ B_weighted, // {d_state, C, n_head} B * decay_from_end + float * __restrict__ C_scaled, // {d_state, C, n_head} C * decay_to_pos (always float) + const float * __restrict__ CB, // {chunk_len, chunk_len, n_group, n_seqs} + half * __restrict__ M_out, // {chunk_len, chunk_len, n_head, n_seqs} + const int chunk_len, const int head_dim, const int n_head, const int n_group, + const int d_state, const int A_stride, + const int x_stride_tok, const int x_stride_seq, + const int B_stride_tok, const int B_stride_seq, + const int C_stride_tok, const int C_stride_seq, + const int chunk_offset, + const int n_tok_total) { + + const int h = blockIdx.y; + const int s = blockIdx.z; + const int g = h / (n_head / n_group); + + const float A_h = A[h * A_stride]; + const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x; + + const int cs_seq_off = s * n_tok_total * n_head; + const float cs_base = (chunk_offset > 0) ? cs[cs_seq_off + (chunk_offset - 1) * n_head + h] : 0.0f; + const float cs_last = cs[cs_seq_off + (chunk_offset + chunk_len - 1) * n_head + h] - cs_base; + + // Prepare X_dt = x * dt, stored d-fastest for coalesced reads and writes. + const int n_xdt = chunk_len * head_dim; + if (idx < n_xdt) { + const int d = idx % head_dim; + const int t = idx / head_dim; + + const float x_val = x[s * x_stride_seq + (chunk_offset + t) * x_stride_tok + d + h * head_dim]; + const float dt_val = dt_sp[cs_seq_off + (chunk_offset + t) * n_head + h]; + + X_dt[d + t * head_dim + h * n_xdt + s * n_xdt * n_head] = (T_matmul)(x_val * dt_val); + } + + // Prepare B_weighted and C_scaled together: both share the same index space (d_state * chunk_len) + // and the same cs_t load, so merging halves the cs[] global memory traffic. + const int n_bw = d_state * chunk_len; + if (idx < n_bw) { + const int n = idx % d_state; + const int t = idx / d_state; + + const float cs_t = cs[cs_seq_off + (chunk_offset + t) * n_head + h] - cs_base; + + const float B_val = B[s * B_stride_seq + (chunk_offset + t) * B_stride_tok + g * d_state + n]; + B_weighted[n + t * d_state + h * n_bw + s * n_bw * n_head] = (T_matmul)(B_val * __expf(A_h * (cs_last - cs_t))); + + const float C_val = C_src[s * C_stride_seq + (chunk_offset + t) * C_stride_tok + g * d_state + n]; + C_scaled[n + t * d_state + h * n_bw + s * n_bw * n_head] = C_val * __expf(A_h * cs_t); + } + + // Materialize M = exp(A*(cs_out - cs_in)) * CB with causal mask. + const int n_M = chunk_len * chunk_len; + if (idx < n_M) { + const int t_out = idx % chunk_len; + const int t_in = idx / chunk_len; + + half val; + if (t_in <= t_out) { + const float cs_out = cs[cs_seq_off + (chunk_offset + t_out) * n_head + h] - cs_base; + const float cs_in = cs[cs_seq_off + (chunk_offset + t_in) * n_head + h] - cs_base; + const float decay = __expf(A_h * (cs_out - cs_in)); + const float * CB_g = CB + (int64_t)s * chunk_len * chunk_len * n_group + + (int64_t)g * chunk_len * chunk_len; + const float cb_val = CB_g[t_out + t_in * chunk_len]; + val = __float2half(decay * cb_val); + } else { + val = __float2half(0.0f); + } + + M_out[(int64_t)s * n_M * n_head + (int64_t)h * n_M + t_in * chunk_len + t_out] = val; + } +} + +// Scale running state in-place: s_cur *= decay_total(chunk). +// Called BEFORE cuBLAS state update (beta=1) to fuse inter-chunk decay. +// Eliminates the s_old buffer and D2D memcpy vs the old approach of: +// memcpy(s_old, s_cur) -> cuBLAS(beta=0) -> s_cur += decay * s_old +// Grid: (ceil(d_state * head_dim / BLOCK), n_head, n_seqs) +template +__global__ void ssm_ssd_scale_state_kernel( + float * __restrict__ s_cur, // {d_state, head_dim, n_head, n_seqs} + const float * __restrict__ cs, // {n_tok, n_head} cumulative dt sums + const float * __restrict__ A, // {1, n_head} + const int d_state, const int head_dim, const int n_head, + const int chunk_offset, const int chunk_len, + const int n_tok_total, const int A_stride) { + + const int h = blockIdx.y; + const int s = blockIdx.z; + const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x; + const int state_per_head = d_state * head_dim; + if (idx >= state_per_head) return; + + const float A_h = A[h * A_stride]; + const int cs_seq_off = s * n_tok_total * n_head; + const float cs_base = (chunk_offset > 0) ? cs[cs_seq_off + (chunk_offset - 1) * n_head + h] : 0.0f; + const float cs_last = cs[cs_seq_off + (chunk_offset + chunk_len - 1) * n_head + h] - cs_base; + const float decay_total = __expf(A_h * cs_last); + + const int off = s * state_per_head * n_head + h * state_per_head + idx; + s_cur[off] *= decay_total; +} + +// Copy initial state from src0[ids[s]] into s_cur for each sequence. +// Grid: (ceil(d_state * head_dim * n_head / BLOCK), n_seqs) +template +__global__ void ssm_ssd_init_state_kernel( + const float * __restrict__ src0, // {d_state, head_dim, n_head, n_rs} + const int32_t * __restrict__ ids, // {n_seqs} + float * __restrict__ s_cur, // {d_state, head_dim, n_head, n_seqs} + const int state_size, // d_state * head_dim * n_head + const int64_t s0_stride_seq) { // elements between state rows + const int s = blockIdx.y; + const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x; + if (idx >= state_size) return; + + const float * s_src = src0 + (int64_t)ids[s] * s0_stride_seq; + s_cur[s * state_size + idx] = s_src[idx]; +} + +// SSD (State Space Duality) dispatch for Mamba-2 prefill. +// Chunked matmuls: CB, materialize M + cuBLAS Y, S@C, B@X_dt. +// All strides are in elements (floats), not bytes. +static void ssm_scan_ssd_f32_cuda( + ggml_backend_cuda_context & ctx, + const float * src0_d, const float * src1_d, const float * src2_d, const float * src3_d, + const float * src4_d, const float * src5_d, const int32_t * src6_d, float * dst_d, + const int64_t s0_stride_seq, // state (src0) stride between seqs + const int x_stride_tok, const int x_stride_seq, // x (src1) strides + const int dt_stride_tok, const int dt_stride_seq, // dt (src2) strides + const int A_stride, // A (src3) stride between heads + const int B_stride_tok, const int B_stride_seq, // B (src4) strides + const int C_stride_tok, const int C_stride_seq, // C (src5) strides + const int64_t s_off, const int64_t d_state, const int64_t head_dim, + const int64_t n_head, const int64_t n_group, const int64_t n_tok, const int64_t n_seq) { + + cudaStream_t stream = ctx.stream(); + const int64_t d_inner = head_dim * n_head; + + const int64_t chunk_size = SSM_SSD_CHUNK_SIZE; + const int64_t n_chunks = (n_tok + chunk_size - 1) / chunk_size; + + const int64_t state_per_head = d_state * head_dim; + + using matmul_t = half; + static constexpr cudaDataType_t matmul_dtype = CUDA_R_16F; + + ggml_cuda_pool_alloc dt_sp_buf(ctx.pool(), n_tok * n_head * n_seq); + ggml_cuda_pool_alloc cs_buf(ctx.pool(), n_tok * n_head * n_seq); + ggml_cuda_pool_alloc CB_buf(ctx.pool(), chunk_size * chunk_size * n_group * n_seq); + ggml_cuda_pool_alloc X_dt_buf(ctx.pool(), chunk_size * head_dim * n_head * n_seq); + ggml_cuda_pool_alloc B_w_buf(ctx.pool(), d_state * chunk_size * n_head * n_seq); + ggml_cuda_pool_alloc C_s_buf(ctx.pool(), d_state * chunk_size * n_head * n_seq); + float * dt_sp = dt_sp_buf.get(); + float * cs = cs_buf.get(); + float * CB = CB_buf.get(); + matmul_t * X_dt = X_dt_buf.get(); + matmul_t * B_weighted = B_w_buf.get(); + float * C_scaled = C_s_buf.get(); + float * s_cur = (float *)((char *)dst_d + s_off); // write state directly to dst + + // Step 1: softplus(dt) and parallel prefix sum over full sequence + { + dim3 grid(n_head, n_seq); + ssm_ssd_prepare_dt_kernel<<>>( + src2_d, dt_sp, cs, n_head, n_tok, dt_stride_tok, dt_stride_seq); + CUDA_CHECK(cudaGetLastError()); + } + + // Step 2: initialize running state from src0[ids[s]] + { + constexpr int BLOCK = 256; + const int64_t state_size = d_state * head_dim * n_head; + dim3 grid((state_size + BLOCK - 1) / BLOCK, n_seq); + ssm_ssd_init_state_kernel<<>>( + src0_d, src6_d, s_cur, state_size, s0_stride_seq); + CUDA_CHECK(cudaGetLastError()); + } + + // Step 3: chunked SSD loop + // Per chunk: pre_matmul (incl. M) + 4 cuBLAS (CB, Y, S@C, state update) + scale_state + cublasHandle_t handle = ctx.cublas_handle(); + CUBLAS_CHECK(cublasSetStream(handle, stream)); + const float alpha_one = 1.0f; + const float beta_zero = 0.0f; + const float beta_one = 1.0f; + const int lda_C_src = C_stride_tok; // leading dim for C in CB = C^T @ B + const int ldb_B_src = B_stride_tok; // leading dim for B in CB = C^T @ B + + // Scratch buffer for causal M matrix, reused across chunks (max size at chunk_size) + const int64_t n_M_max = chunk_size * chunk_size; + ggml_cuda_pool_alloc M_buf(ctx.pool(), n_M_max * n_head * n_seq); + half * M_mat = M_buf.get(); + + for (int64_t k = 0; k < n_chunks; k++) { + const int64_t chunk_offset = k * chunk_size; + const int64_t chunk_len = (chunk_offset + chunk_size <= n_tok) ? chunk_size : (n_tok - chunk_offset); + + // 3a: CB = C^T @ B per group + for (int64_t s = 0; s < n_seq; s++) { + const float * C_s = src5_d + s * C_stride_seq + chunk_offset * C_stride_tok; + const float * B_s = src4_d + s * B_stride_seq + chunk_offset * B_stride_tok; + float * CB_s = CB + s * chunk_len * chunk_len * n_group; + + if (n_group == 1) { + CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, + chunk_len, chunk_len, d_state, + &alpha_one, C_s, lda_C_src, B_s, ldb_B_src, + &beta_zero, CB_s, (int)chunk_len)); + } else { + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N, + chunk_len, chunk_len, d_state, + &alpha_one, + C_s, CUDA_R_32F, lda_C_src, d_state, + B_s, CUDA_R_32F, ldb_B_src, d_state, + &beta_zero, + CB_s, CUDA_R_32F, (int)chunk_len, (long long)(chunk_len * chunk_len), + n_group, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); + } + } + + // 3b: prepare X_dt, B_weighted, C_scaled + materialize causal M matrix + const int64_t n_M = chunk_len * chunk_len; + { + constexpr int BLOCK = 256; + const int64_t n_xdt = chunk_len * head_dim; + const int64_t n_bw = d_state * chunk_len; + int64_t max_work = n_xdt; + if (n_bw > max_work) max_work = n_bw; + if (n_M > max_work) max_work = n_M; + dim3 grid((max_work + BLOCK - 1) / BLOCK, n_head, n_seq); + ssm_ssd_pre_matmul_kernel<<>>( + cs, dt_sp, src3_d, src1_d, src4_d, src5_d, + X_dt, B_weighted, C_scaled, + CB, M_mat, + chunk_len, head_dim, n_head, n_group, d_state, A_stride, + x_stride_tok, x_stride_seq, B_stride_tok, B_stride_seq, C_stride_tok, C_stride_seq, + chunk_offset, n_tok); + CUDA_CHECK(cudaGetLastError()); + } + + // 3c: dst = S_cur^T @ C_scaled (state contribution) + { + const int64_t stride_S = state_per_head; + const int64_t stride_Cs = d_state * chunk_len; + + for (int64_t s = 0; s < n_seq; s++) { + float * dst_chunk = dst_d + s * d_inner * n_tok + chunk_offset * d_inner; + + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N, + head_dim, chunk_len, d_state, + &alpha_one, + s_cur + s * stride_S * n_head, CUDA_R_32F, d_state, stride_S, + C_scaled + s * stride_Cs * n_head, CUDA_R_32F, d_state, stride_Cs, + &beta_zero, + dst_chunk, CUDA_R_32F, d_inner, head_dim, + n_head, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); + } + } + + // 3d: dst += X_dt @ M^T (intra-chunk contribution, adds to 3c result) + // M is stored as M[t_out, t_in] (lower-triangular), transpose needed for Y = X @ M^T. + { + const int64_t stride_M = n_M; + const int64_t stride_X_h = (int64_t)chunk_len * head_dim; + + for (int64_t s = 0; s < n_seq; s++) { + float * dst_chunk = dst_d + s * d_inner * n_tok + chunk_offset * d_inner; + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, + head_dim, chunk_len, chunk_len, + &alpha_one, + X_dt + s * stride_X_h * n_head, matmul_dtype, head_dim, stride_X_h, + M_mat + s * stride_M * n_head, matmul_dtype, chunk_len, stride_M, + &beta_one, + dst_chunk, CUDA_R_32F, d_inner, head_dim, + n_head, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); + } + } + + // 3e: s_cur = B_weighted @ X_dt^T + decay_total * s_cur_old (state update) + { + // Scale s_cur in-place by per-head decay_total BEFORE cuBLAS overwrites it + constexpr int BLOCK = 256; + dim3 grid((state_per_head + BLOCK - 1) / BLOCK, n_head, n_seq); + ssm_ssd_scale_state_kernel<<>>( + s_cur, cs, src3_d, + d_state, head_dim, n_head, + chunk_offset, chunk_len, n_tok, A_stride); + CUDA_CHECK(cudaGetLastError()); + + // cuBLAS with beta=1: s_cur = B_weighted @ X_dt^T + 1.0 * s_cur (already scaled) + const int64_t stride_Bw = d_state * chunk_len; + const int64_t stride_X = chunk_len * head_dim; + const int64_t stride_S = state_per_head; + + for (int64_t s = 0; s < n_seq; s++) { + // X_dt is d-fastest {hd, C}, read as OP_T to get {C, hd} + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, + d_state, head_dim, chunk_len, + &alpha_one, + B_weighted + s * stride_Bw * n_head, matmul_dtype, d_state, stride_Bw, + X_dt + s * stride_X * n_head, matmul_dtype, head_dim, stride_X, + &beta_one, + s_cur + s * stride_S * n_head, CUDA_R_32F, d_state, stride_S, + n_head, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); + } + } + } +} +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const struct ggml_tensor * src0 = dst->src[0]; // s const struct ggml_tensor * src1 = dst->src[1]; // x @@ -357,6 +795,49 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { GGML_ASSERT(src6->type == GGML_TYPE_I32); GGML_ASSERT(dst->type == GGML_TYPE_F32); + // Byte strides are narrowed to int for both scan and SSD paths. + GGML_ASSERT(src0->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src0->nb[3] <= (size_t)INT_MAX); + GGML_ASSERT(src1->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src1->nb[3] <= (size_t)INT_MAX); + GGML_ASSERT(src2->nb[1] <= (size_t)INT_MAX); + GGML_ASSERT(src2->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src3->nb[1] <= (size_t)INT_MAX); + GGML_ASSERT(src4->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src4->nb[3] <= (size_t)INT_MAX); + GGML_ASSERT(src5->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src5->nb[3] <= (size_t)INT_MAX); + +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + // Mamba-2 with scalar A per head: use SSD matmul path for long sequences. + // Requires NVIDIA Turing+ otherwise fallback to scan. + const bool is_mamba2 = (src3->nb[1] == sizeof(float)); + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const bool use_ssd = is_mamba2 && n_t > SSM_SSD_MIN_TOKENS + && n_t <= SSM_SSD_MAX_TOKENS + && GGML_CUDA_CC_IS_NVIDIA(cc) + && cc >= GGML_CUDA_CC_TURING + && nr % 8 == 0; // cuBLAS requires 8-element (16-byte) alignment + + if (use_ssd) { + // ssm_ssd_init_state_kernel uses flat linear indexing within each sequence, + // so src0 must be fully contiguous across all inner dimensions. + // The scan path handles non-contiguous nb[2] via src0_nb2 but does not handle nb[1]. + GGML_ASSERT(src0->nb[1] == nc * sizeof(float)); + GGML_ASSERT(src0->nb[2] == nc * nr * sizeof(float)); + + ssm_scan_ssd_f32_cuda(ctx, + src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d, + (int64_t)(src0->nb[3] / sizeof(float)), + (int)(src1->nb[2] / sizeof(float)), (int)(src1->nb[3] / sizeof(float)), + (int)(src2->nb[1] / sizeof(float)), (int)(src2->nb[2] / sizeof(float)), + (int)(src3->nb[1] / sizeof(float)), + (int)(src4->nb[2] / sizeof(float)), (int)(src4->nb[3] / sizeof(float)), + (int)(src5->nb[2] / sizeof(float)), (int)(src5->nb[3] / sizeof(float)), + s_off, nc, nr, nh, ng, n_t, n_s); + return; + } +#endif ssm_scan_f32_cuda(src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d, src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2], src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3], diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index a05d18ee30..d07b8fe41a 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -15675,7 +15675,7 @@ static void ggml_cl_mul_mat_kq_kqv_adreno(ggml_backend_t backend, const ggml_ten // <--------------------------------------------> // extra0 = src0->view_src ? (ggml_tensor_extra_cl *)src0->view_src->extra : (ggml_tensor_extra_cl *)src0->extra; - region.origin = (extra0->offset); + region.origin = (extra0->offset + src0->view_offs); if (nb01 > nb02) { // KQ region.size = nb01 * ne01; @@ -15691,7 +15691,7 @@ static void ggml_cl_mul_mat_kq_kqv_adreno(ggml_backend_t backend, const ggml_ten // create sub-buffer for B // <--------------------------------------------> // - region.origin = (extra1->offset); + region.origin = (extra1->offset + src1->view_offs); region.size = nb10 * ne10 * ne11 * ne12; B_sub_buffer = clCreateSubBuffer((extra1->data_device), 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status); CL_CHECK(status); @@ -15712,7 +15712,7 @@ static void ggml_cl_mul_mat_kq_kqv_adreno(ggml_backend_t backend, const ggml_ten // create sub-buffer for output C // <--------------------------------------------> // - region.origin = (extrad->offset); + region.origin = (extrad->offset + dst->view_offs); region.size = ne0 * ne1 * dst->ne[2] * dst->nb[0]; // size of C in bytes D_sub_buffer = clCreateSubBuffer((extrad->data_device), 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status); CL_CHECK(status); @@ -18591,6 +18591,8 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co #ifdef GGML_OPENCL_USE_ADRENO_KERNELS if(src0t == GGML_TYPE_F16 && src1t == GGML_TYPE_F32){ if (ne01 >= 64 && ne1 >= 32 && ne00 >= 16 && (ne12 % ne02) == 0 && + // the KQ/KQV image kernels do not handle dim 3 (multi-stream batches) + ne03 == 1 && ne13 == 1 && // dst is wrapped with image1d_buffer, the size limit applies, also src0 (ne0 * ne1 * dst->ne[2] * dst->nb[0] / 4 <= backend_ctx->image_max_buffer_size)) { // For KQ diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 5dcf4503bb..49fd6e0b42 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -3490,7 +3490,7 @@ struct vk_fa_tuning_params { }; static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type, ggml_type v_type); -static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type = GGML_TYPE_F16); +static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type = GGML_TYPE_F16, ggml_type v_type = GGML_TYPE_F16); static vk_fa_tuning_params get_fa_tuning_params_scalar(const vk_device& device, uint32_t hsk, uint32_t hsv, uint32_t n_rows, uint32_t n_kv, ggml_type k_type, ggml_type v_type, bool f32acc) { @@ -3646,7 +3646,7 @@ static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_ bool shape_ok = (f32acc && device->coopmat_support_16x16x16_f32acc) || (!f32acc && device->coopmat_support_16x16x16_f16acc); const vk_fa_tuning_params params = get_fa_tuning_params_coopmat1(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc); - bool shmem_ok = ggml_vk_flash_attn_coopmat_shmem_support(device, params, hsk, hsv, f32acc, k_type); + bool shmem_ok = ggml_vk_flash_attn_coopmat_shmem_support(device, params, hsk, hsv, f32acc, k_type, v_type); if (!shape_ok || !shmem_ok) { path = FA_SCALAR; @@ -3658,11 +3658,6 @@ static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_ path = FA_SCALAR; } - // Q1_0 K/V is only implemented on coopmat2 (flash_attn_cm2); there is no scalar FA shader for it. - if ((k_type == GGML_TYPE_Q1_0 || v_type == GGML_TYPE_Q1_0) && device->coopmat2) { - path = FA_COOPMAT2; - } - switch (path) { case FA_SCALAR: return get_fa_tuning_params_scalar(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc); @@ -3904,16 +3899,27 @@ static uint32_t get_subgroup_size(const std::string &pipeline_name, const vk_dev return 0; // If no matching configuration is found } -// Whether scalar flash attention will use the MMQ path for the given k_type. -static bool ggml_vk_fa_scalar_uses_mmq(const vk_device& device, ggml_type k_type) { +// Whether scalar flash attention will use the MMQ path for the given K/V types. +static bool ggml_vk_fa_type_needs_shmem(ggml_type type) { + switch (type) { + case GGML_TYPE_IQ4_NL: + return true; + default: + return false; + } +} + +static bool ggml_vk_fa_scalar_uses_mmq(const vk_device& device, ggml_type k_type, ggml_type v_type) { #if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT) return device->integer_dot_product && device->subgroup_clustered && + !ggml_vk_fa_type_needs_shmem(v_type) && (k_type == GGML_TYPE_Q4_0 || k_type == GGML_TYPE_Q4_1 || k_type == GGML_TYPE_Q5_0 || k_type == GGML_TYPE_Q5_1 || k_type == GGML_TYPE_Q8_0); #else GGML_UNUSED(device); GGML_UNUSED(k_type); + GGML_UNUSED(v_type); return false; #endif } @@ -4246,7 +4252,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { const bool fa_ds = fa.first.subgroup_size == 0; const bool bf16_kv = fa.first.k_type == GGML_TYPE_BF16; - const bool use_mmq = ggml_vk_fa_scalar_uses_mmq(device, fa.first.k_type); + const bool use_mmq = ggml_vk_fa_scalar_uses_mmq(device, fa.first.k_type, fa.first.v_type); const void * spv_data = nullptr; size_t spv_size = 0; const char *name = nullptr; @@ -10380,7 +10386,6 @@ static void ggml_vk_mul_mat_id(ggml_backend_vk_context * ctx, vk_context& subctx static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type, ggml_type v_type) { GGML_UNUSED(f32acc); - GGML_UNUSED(v_type); // Needs to be kept up to date on shader changes const uint32_t wg_size = params.workgroup_size; const uint32_t Br = params.block_rows; @@ -10389,13 +10394,15 @@ static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, con // BF16 uses the fp32 shader (FLOAT_TYPE=float) const uint32_t float_type_size = (device->fp16 && k_type != GGML_TYPE_BF16) ? sizeof(ggml_fp16_t) : sizeof(float); - const bool mmq = ggml_vk_fa_scalar_uses_mmq(device, k_type); + const bool mmq = ggml_vk_fa_scalar_uses_mmq(device, k_type, v_type); // tmpsh is overestimated slightly const uint32_t tmpsh = wg_size * sizeof(float); const uint32_t tmpshv4 = wg_size * 4 * float_type_size; const uint32_t masksh = Bc * (Br + 1) * float_type_size; + // DATA_A_IQ4_NL is compiled into the FA shaders unconditionally, so its shared table is always allocated. + const uint32_t iq_shmem = 16 * float_type_size; uint32_t Qf, kvsh, kblocksh_size; if (mmq) { @@ -10420,7 +10427,7 @@ static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, con kblocksh_size = 0; } - const uint32_t total_size = tmpsh + tmpshv4 + masksh + Qf + kvsh + kblocksh_size; + const uint32_t total_size = tmpsh + tmpshv4 + masksh + iq_shmem + Qf + kvsh + kblocksh_size; const bool supported = total_size <= device->properties.limits.maxComputeSharedMemorySize; VK_LOG_DEBUG("ggml_vk_flash_attn_scalar_shmem_support(HSK=" << hsk << ", HSV=" << hsv << ", mmq=" << mmq << ", total_size=" << total_size << ", supported=" << supported); @@ -10428,7 +10435,8 @@ static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, con return supported; } -static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type) { +static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type, ggml_type v_type) { + GGML_UNUSED(v_type); // Needs to be kept up to date on shader changes const uint32_t Br = params.block_rows; const uint32_t Bc = params.block_cols; @@ -10444,6 +10452,8 @@ static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, co const uint32_t f16vec4 = 8; const uint32_t tmpsh = (Bc / MatBc) * sizeof(float); + // DATA_A_IQ4_NL is compiled into the FA shaders unconditionally, so its shared table is always allocated. + const uint32_t iq_shmem = 16 * sizeof(ggml_fp16_t); const uint32_t qstride = hsk_pad / 4 + 2; const uint32_t Qf = Br * qstride * f16vec4; @@ -10465,7 +10475,7 @@ static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, co const uint32_t slope = Br * acctype; - const uint32_t total_size = tmpsh + Qf + Psh + sfsh + ksh + pvsh + slope; + const uint32_t total_size = tmpsh + iq_shmem + Qf + Psh + sfsh + ksh + pvsh + slope; const bool supported = total_size <= device->properties.limits.maxComputeSharedMemorySize; VK_LOG_DEBUG("ggml_vk_flash_attn_coopmat_shmem_support(HSK=" << hsk << ", HSV=" << hsv << ", f32acc=" << f32acc << ", total_size=" << total_size << ", supported=" << supported); @@ -17617,7 +17627,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm if (op->src[3] && op->src[3]->type != GGML_TYPE_F16) { return false; } - auto fa_kv_ok = [coopmat2](ggml_type t) { + auto fa_kv_ok = [](ggml_type t) { switch (t) { case GGML_TYPE_F32: case GGML_TYPE_F16: @@ -17627,9 +17637,8 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_TYPE_Q5_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q4_0: + case GGML_TYPE_IQ4_NL: return true; - case GGML_TYPE_Q1_0: - return coopmat2; default: return false; } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp index 3192130ccf..6c264c7861 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp @@ -80,7 +80,9 @@ shared vec4 occupancy_limiter[LIMIT_OCCUPANCY_SHMEM > 0 ? LIMIT_OCCUPANCY_SHMEM void main() { #ifdef NEEDS_INIT_IQ_SHMEM - init_iq_shmem(gl_WorkGroupSize); + if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) { + init_iq_shmem(gl_WorkGroupSize); + } #endif init_indices(); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl index 66dcf61021..3c64f91dad 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl @@ -97,8 +97,8 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];}; #define FA_TYPE_Q5_0 6u #define FA_TYPE_Q5_1 7u #define FA_TYPE_Q8_0 8u +#define FA_TYPE_IQ4_NL 20u #define FA_TYPE_BF16 30u -#define FA_TYPE_Q1_0 41u #if defined(BFLOAT16) #define O_TYPE float @@ -120,8 +120,8 @@ uint fa_block_elems(uint ty) { case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0); case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1); case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0); + case FA_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL); case FA_TYPE_BF16: return 1u; - case FA_TYPE_Q1_0: return uint(QUANT_K_Q1_0); // cm2-only, harmless elsewhere default: return 1u; } } @@ -140,6 +140,13 @@ uint fa_quant_r_mmq(uint ty) { } } +bool fa_type_needs_shmem(uint ty) { + switch (ty) { + case FA_TYPE_IQ4_NL: return true; + default: return false; + } +} + // These can't be `const` globals because GLSL forbids function calls in global // const initializers, even when the spec constants would let the driver fold // them. Macros expand at the use site and fold after specialization. diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 16178e5770..057ed739aa 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -64,7 +64,9 @@ shared ACC_TYPE slope[Br]; void main() { #ifdef NEEDS_INIT_IQ_SHMEM - init_iq_shmem(gl_WorkGroupSize); + if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) { + init_iq_shmem(gl_WorkGroupSize); + } #endif init_indices(); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp index b9c03fe499..3174111530 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp @@ -46,7 +46,7 @@ float16_t faDecodeK(const decodeBufFA_K bl_in, const uint blockCoords[2], const case FA_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); case FA_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); case FA_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); - case FA_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock); default: return float16_t(0); } } @@ -59,7 +59,7 @@ float16_t faDecodeV(const decodeBufFA_V bl_in, const uint blockCoords[2], const case FA_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); case FA_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); case FA_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); - case FA_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock); default: return float16_t(0); } } @@ -67,26 +67,26 @@ float16_t faDecodeV(const decodeBufFA_V bl_in, const uint blockCoords[2], const // V=4 vector decode for K/V; dispatches to per-format _v decoders. f16vec4 faDecodeKVector(const decodeBufFA_K bl_in, const uint blockCoords[2], const uint coordInBlock[2]) { switch (FaTypeK) { - case 0u: return f16vec4(decodeBufF32(bl_in).block); - case 2u: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock); - case 3u: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock); - case 6u: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); - case 7u: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); - case 8u: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); - case 41u: return dequantFuncQ1_0_v(decodeBufQ1_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block); + case FA_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock); default: return f16vec4(0); } } f16vec4 faDecodeVVector(const decodeBufFA_V bl_in, const uint blockCoords[2], const uint coordInBlock[2]) { switch (FaTypeV) { - case 0u: return f16vec4(decodeBufF32(bl_in).block); - case 2u: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock); - case 3u: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock); - case 6u: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); - case 7u: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); - case 8u: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); - case 41u: return dequantFuncQ1_0_v(decodeBufQ1_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block); + case FA_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock); default: return f16vec4(0); } } @@ -169,6 +169,12 @@ ACC_TYPE perElemOpNonGqaSplitKStoreCol0(const in uint32_t r, const in uint32_t c } void main() { +#ifdef NEEDS_INIT_IQ_SHMEM + if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) { + init_iq_shmem(gl_WorkGroupSize); + } +#endif + init_indices(); tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutQ = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); @@ -302,7 +308,7 @@ void main() { coopmat K_T; uint32_t k_offset = ik2*p.nb12 + ik3*p.nb13; - // F16: bs_k==1 (direct load). F32: bs_k==4 (vec4 / dequantFuncF32). Q4/Q8 family: bs_k==32. Q1_0: bs_k==128. + // F16: bs_k==1 (direct load). F32: bs_k==4 (vec4 / dequantFuncF32). Quantized types: bs_k==32. #if defined(BFLOAT16) coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose); #else diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl index 8704479d96..8ba4725f33 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl @@ -27,6 +27,8 @@ layout (binding = 1) readonly buffer K_PACKED_Q5_1 { block_q5_1_packed16 data[]; layout (binding = 2) readonly buffer V_PACKED_Q5_1 { block_q5_1_packed16 data[]; } v_packed_q5_1; layout (binding = 1) readonly buffer K_PACKED_Q8_0 { block_q8_0_packed16 data[]; } k_packed_q8_0; layout (binding = 2) readonly buffer V_PACKED_Q8_0 { block_q8_0_packed16 data[]; } v_packed_q8_0; +layout (binding = 1) readonly buffer K_PACKED_IQ4_NL { block_iq4_nl_packed16 data[]; } k_packed_iq4_nl; +layout (binding = 2) readonly buffer V_PACKED_IQ4_NL { block_iq4_nl_packed16 data[]; } v_packed_iq4_nl; layout (binding = 1) readonly buffer K_PACKED_BF16 { u16vec4 data[]; } k_packed_bf16; layout (binding = 2) readonly buffer V_PACKED_BF16 { u16vec4 data[]; } v_packed_bf16; @@ -102,6 +104,17 @@ layout (binding = 1) readonly buffer K_PACKED_Q5_1_P32 { block_q5_1_packed32 dat return FLOAT_TYPE(BUF.data[a_offset + ib].d) * FLOAT_TYPEV4(v0.x, v0.y, v1.x, v1.y); \ } +#define FA_DEQUANT4_IQ4_NL(BUF) { \ + const uint shift = (iqs & 0x10) >> 2; \ + const uint qs_i = (iqs & 0xC) >> 1; \ + const uint qsw = uint(BUF.data[a_offset + ib].qs[qs_i]) \ + | (uint(BUF.data[a_offset + ib].qs[qs_i + 1u]) << 16); \ + const FLOAT_TYPE d = FLOAT_TYPE(BUF.data[a_offset + ib].d); \ + const u8vec4 q = unpack8((qsw >> shift) & 0x0F0F0F0Fu); \ + return d * FLOAT_TYPEV4(kvalues_iq4nl[q.x], kvalues_iq4nl[q.y], \ + kvalues_iq4nl[q.z], kvalues_iq4nl[q.w]); \ +} + #define FA_DEQUANT4_BF16(BUF) \ return FLOAT_TYPEV4(bf16_to_fp32(uvec4(BUF.data[(a_offset + ib) / 4]))); @@ -114,6 +127,7 @@ FLOAT_TYPEV4 dequantize4(uint ib, uint iqs, uint a_offset, uint binding_idx) { case FA_TYPE_Q5_0: FA_DEQUANT4_Q5_0(k_packed_q5_0) case FA_TYPE_Q5_1: FA_DEQUANT4_Q5_1(k_packed_q5_1) case FA_TYPE_Q8_0: FA_DEQUANT4_Q8_0(k_packed_q8_0) + case FA_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(k_packed_iq4_nl) case FA_TYPE_BF16: FA_DEQUANT4_BF16(k_packed_bf16) } } else { @@ -124,6 +138,7 @@ FLOAT_TYPEV4 dequantize4(uint ib, uint iqs, uint a_offset, uint binding_idx) { case FA_TYPE_Q5_0: FA_DEQUANT4_Q5_0(v_packed_q5_0) case FA_TYPE_Q5_1: FA_DEQUANT4_Q5_1(v_packed_q5_1) case FA_TYPE_Q8_0: FA_DEQUANT4_Q8_0(v_packed_q8_0) + case FA_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(v_packed_iq4_nl) case FA_TYPE_BF16: FA_DEQUANT4_BF16(v_packed_bf16) } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index 58d347bc54..3e6b39576e 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -673,6 +673,8 @@ void process_shaders() { fa_base_dict["ACC_TYPE"] = fp16 && f16acc ? "float16_t" : "float"; fa_base_dict["ACC_TYPEV2"] = fp16 && f16acc ? "f16vec2" : "vec2"; fa_base_dict["ACC_TYPEV4"] = fp16 && f16acc ? "f16vec4" : "vec4"; + // Compile IQ4_NL support into all FA variants so its shared LUT is available when K or V uses it. + fa_base_dict["DATA_A_IQ4_NL"] = "1"; if (fp16 && f16acc) { fa_base_dict["ACC_TYPE_MAX"] = "float16_t(65504.0)"; } diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp index bed9265b8a..99d775c576 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp @@ -73,11 +73,6 @@ inline bool ggml_webgpu_tensor_equal(const ggml_tensor * a, const ggml_tensor * return a->buffer == b->buffer && ggml_webgpu_tensor_addr(a) == ggml_webgpu_tensor_addr(b); } -inline bool ggml_webgpu_tensor_overlap(const ggml_tensor * a, const ggml_tensor * b) { - return a->buffer == b->buffer && ggml_webgpu_tensor_addr(a) < ggml_webgpu_tensor_addr(b) + ggml_nbytes(b) && - ggml_webgpu_tensor_addr(b) < ggml_webgpu_tensor_addr(a) + ggml_nbytes(a); -} - struct ggml_webgpu_shader_lib_context { ggml_tensor * src0; ggml_tensor * src1; @@ -118,6 +113,11 @@ struct ggml_webgpu_binary_shader_decisions { bool src_overlap = false; }; +struct ggml_webgpu_glu_shader_decisions { + uint32_t wg_size = 0; + bool src_overlap = false; +}; + struct ggml_webgpu_processed_shader { std::string wgsl; std::string variant; @@ -133,9 +133,12 @@ struct ggml_webgpu_ssm_scan_pipeline_key { int type; int d_state; bool xbc_overlap; + bool a_overlap; + bool ids_overlap; bool operator==(const ggml_webgpu_ssm_scan_pipeline_key & other) const { - return type == other.type && d_state == other.d_state && xbc_overlap == other.xbc_overlap; + return type == other.type && d_state == other.d_state && xbc_overlap == other.xbc_overlap && + a_overlap == other.a_overlap && ids_overlap == other.ids_overlap; } }; @@ -145,6 +148,8 @@ struct ggml_webgpu_ssm_scan_pipeline_key_hash { ggml_webgpu_hash_combine(seed, key.type); ggml_webgpu_hash_combine(seed, key.d_state); ggml_webgpu_hash_combine(seed, key.xbc_overlap); + ggml_webgpu_hash_combine(seed, key.a_overlap); + ggml_webgpu_hash_combine(seed, key.ids_overlap); return seed; } }; @@ -153,6 +158,8 @@ struct ggml_webgpu_ssm_scan_shader_decisions { uint32_t wg_size; uint32_t tokens_per_tile; bool xbc_overlap = false; + bool a_overlap = false; + bool ids_overlap = false; }; /** Argsort **/ @@ -264,7 +271,7 @@ struct ggml_webgpu_row_norm_pipeline_key_hash { struct ggml_webgpu_rms_norm_mul_pipeline_key { bool inplace; // rn_src == dst bool overlap; // mul_src == dst - bool src_overlap; // rn_src == mul_src + bool src_overlap; // rn_src binding overlaps mul_src binding bool operator==(const ggml_webgpu_rms_norm_mul_pipeline_key & other) const { return inplace == other.inplace && overlap == other.overlap && src_overlap == other.src_overlap; @@ -690,7 +697,8 @@ inline bool ggml_webgpu_flash_attn_kv_direct(const ggml_tensor * Q, inline ggml_webgpu_flash_attn_common_pipeline_key ggml_webgpu_flash_attn_make_common_pipeline_key( const ggml_webgpu_shader_lib_context & context, - uint32_t kv_direct_align) { + uint32_t kv_direct_align, + bool kv_overlap) { ggml_webgpu_flash_attn_common_pipeline_key key = {}; key.q_type = context.src0->type; key.k_type = context.src1->type; @@ -699,7 +707,7 @@ inline ggml_webgpu_flash_attn_common_pipeline_key ggml_webgpu_flash_attn_make_co key.head_dim_qk = (uint32_t) context.src0->ne[0]; key.head_dim_v = (uint32_t) context.src2->ne[0]; key.kv_direct = ggml_webgpu_flash_attn_kv_direct(context.src0, context.src1, context.src2, kv_direct_align); - key.kv_overlap = ggml_webgpu_tensor_overlap(context.src1, context.src2); + key.kv_overlap = kv_overlap; key.has_mask = context.src3 != nullptr; key.has_sinks = context.src4 != nullptr; key.uses_logit_softcap = ggml_get_op_params_f32(context.dst, 2) != 0.0f; @@ -1066,9 +1074,10 @@ struct ggml_webgpu_glu_pipeline_key { ggml_glu_op glu_op; ggml_type type; bool split; + bool src_overlap; bool operator==(const ggml_webgpu_glu_pipeline_key & other) const { - return glu_op == other.glu_op && type == other.type && split == other.split; + return glu_op == other.glu_op && type == other.type && split == other.split && src_overlap == other.src_overlap; } }; @@ -1078,6 +1087,7 @@ struct ggml_webgpu_glu_pipeline_key_hash { ggml_webgpu_hash_combine(seed, key.glu_op); ggml_webgpu_hash_combine(seed, key.type); ggml_webgpu_hash_combine(seed, key.split); + ggml_webgpu_hash_combine(seed, key.src_overlap); return seed; } }; @@ -1758,12 +1768,16 @@ class ggml_webgpu_shader_lib { return ssm_conv_pipelines[key]; } - webgpu_pipeline get_ssm_scan_pipeline(const ggml_webgpu_shader_lib_context & context) { + webgpu_pipeline get_ssm_scan_pipeline(const ggml_webgpu_shader_lib_context & context, + bool xbc_overlap, + bool a_overlap, + bool ids_overlap) { ggml_webgpu_ssm_scan_pipeline_key key = {}; key.type = context.dst->type; key.d_state = (int) context.src0->ne[0]; - key.xbc_overlap = ggml_webgpu_tensor_overlap(context.src1, context.src4) && - ggml_webgpu_tensor_overlap(context.src1, context.src5); + key.xbc_overlap = xbc_overlap; + key.a_overlap = a_overlap; + key.ids_overlap = ids_overlap; auto it = ssm_scan_pipelines.find(key); if (it != ssm_scan_pipelines.end()) { @@ -1798,7 +1812,12 @@ class ggml_webgpu_shader_lib { if (key.xbc_overlap) { defines.push_back("XBC_OVERLAP"); } - + if (key.a_overlap) { + defines.push_back("A_OVERLAP"); + } + if (key.ids_overlap) { + defines.push_back("IDS_OVERLAP"); + } variant += "_d" + std::to_string(key.d_state); auto processed = preprocessor.preprocess(wgsl_ssm_scan, defines); @@ -1806,6 +1825,8 @@ class ggml_webgpu_shader_lib { decisions->wg_size = wg_size; decisions->tokens_per_tile = tokens_per_tile; decisions->xbc_overlap = key.xbc_overlap; + decisions->a_overlap = key.a_overlap; + decisions->ids_overlap = key.ids_overlap; webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); pipeline.context = decisions; ssm_scan_pipelines[key] = pipeline; @@ -2549,11 +2570,11 @@ class ggml_webgpu_shader_lib { return unary_pipelines[key]; } - webgpu_pipeline get_rms_norm_mul_pipeline(const ggml_webgpu_shader_lib_context & context) { + webgpu_pipeline get_rms_norm_mul_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) { ggml_webgpu_rms_norm_mul_pipeline_key key = {}; key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); key.overlap = ggml_webgpu_tensor_equal(context.src1, context.dst); - key.src_overlap = ggml_webgpu_tensor_overlap(context.src0, context.src1); + key.src_overlap = src_overlap; auto it = rms_norm_mul_pipelines.find(key); if (it != rms_norm_mul_pipelines.end()) { @@ -2589,13 +2610,13 @@ class ggml_webgpu_shader_lib { return rms_norm_mul_pipelines[key]; } - webgpu_pipeline get_binary_pipeline(const ggml_webgpu_shader_lib_context & context) { + webgpu_pipeline get_binary_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) { ggml_webgpu_binary_pipeline_key key = {}; key.type = context.dst->type; key.op = context.dst->op; key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); key.overlap = ggml_webgpu_tensor_equal(context.src1, context.dst); - key.src_overlap = ggml_webgpu_tensor_overlap(context.src0, context.src1); + key.src_overlap = src_overlap; auto it = binary_pipelines.find(key); if (it != binary_pipelines.end()) { @@ -2678,10 +2699,10 @@ class ggml_webgpu_shader_lib { return pipeline; } - webgpu_pipeline get_concat_pipeline(const ggml_webgpu_shader_lib_context & context) { + webgpu_pipeline get_concat_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) { ggml_webgpu_concat_pipeline_key key = {}; key.type = context.dst->type; - key.src_overlap = ggml_webgpu_tensor_overlap(context.src0, context.src1); + key.src_overlap = src_overlap; auto it = concat_pipelines.find(key); if (it != concat_pipelines.end()) { @@ -2761,7 +2782,7 @@ class ggml_webgpu_shader_lib { return repeat_pipelines[key]; } - webgpu_pipeline get_flash_attn_pipeline(const ggml_webgpu_shader_lib_context & context) { + webgpu_pipeline get_flash_attn_pipeline(const ggml_webgpu_shader_lib_context & context, bool kv_overlap) { const bool can_use_subgroup_matrix = ggml_webgpu_flash_attn_can_use_subgroup_matrix_path( context.supports_subgroup_matrix, context.sg_mat_k, context.sg_mat_n, context.src0, context.src2); ggml_webgpu_flash_attn_decisions decisions = {}; @@ -2769,8 +2790,8 @@ class ggml_webgpu_shader_lib { decisions.q_tile = decisions.use_sg_matrix ? context.sg_mat_m : GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE; ggml_webgpu_flash_attn_pipeline_key key = {}; - key.common = - ggml_webgpu_flash_attn_make_common_pipeline_key(context, decisions.use_sg_matrix ? context.sg_mat_k : 1u); + key.common = ggml_webgpu_flash_attn_make_common_pipeline_key( + context, decisions.use_sg_matrix ? context.sg_mat_k : 1u, kv_overlap); key.common.kv_direct = decisions.use_sg_matrix && key.common.kv_direct; key.use_sg_matrix = decisions.use_sg_matrix; @@ -2824,9 +2845,10 @@ class ggml_webgpu_shader_lib { return flash_attn_pipelines[key]; } - webgpu_pipeline get_flash_attn_vec_pipeline(const ggml_webgpu_shader_lib_context & context) { + webgpu_pipeline get_flash_attn_vec_pipeline(const ggml_webgpu_shader_lib_context & context, bool kv_overlap) { ggml_webgpu_flash_attn_vec_pipeline_key key = {}; - key.common = ggml_webgpu_flash_attn_make_common_pipeline_key(context, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH); + key.common = ggml_webgpu_flash_attn_make_common_pipeline_key(context, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH, + kv_overlap); auto it = flash_attn_vec_pipelines.find(key); if (it != flash_attn_vec_pipelines.end()) { @@ -2984,11 +3006,12 @@ class ggml_webgpu_shader_lib { return cpy_pipelines[key]; } - webgpu_pipeline get_glu_pipeline(const ggml_webgpu_shader_lib_context & context) { + webgpu_pipeline get_glu_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) { ggml_webgpu_glu_pipeline_key key = {}; key.glu_op = ggml_get_glu_op(context.dst); key.type = context.dst->type; key.split = (context.src1 != nullptr); + key.src_overlap = src_overlap; auto it = glu_pipelines.find(key); if (it != glu_pipelines.end()) { @@ -3039,7 +3062,10 @@ class ggml_webgpu_shader_lib { GGML_ABORT("Unsupported type for GLU shader"); } - if (key.split) { + if (key.src_overlap) { + defines.push_back("SRC_OVERLAP"); + variant += "_src_overlap"; + } else if (key.split) { variant += "_split"; } else { defines.push_back("NO_SPLIT"); @@ -3048,8 +3074,9 @@ class ggml_webgpu_shader_lib { defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); auto processed = preprocessor.preprocess(wgsl_glu, defines); - auto decisions = std::make_shared(); + auto decisions = std::make_shared(); decisions->wg_size = context.max_wg_size; + decisions->src_overlap = key.src_overlap; webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); pipeline.context = decisions; glu_pipelines[key] = pipeline; diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 75286ec731..2add5da0b4 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -374,18 +374,59 @@ static wgpu::Buffer ggml_webgpu_tensor_buf(const ggml_tensor * tensor) { return ctx->buffer; } -static size_t ggml_webgpu_tensor_misalignment(webgpu_context & ctx, const ggml_tensor * t) { +static size_t ggml_webgpu_tensor_misalignment(const ggml_tensor * t, size_t alignment) { size_t offset = ggml_webgpu_tensor_offset(t); - return offset & (ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment - 1); + return offset & (alignment - 1); +} + +static size_t ggml_webgpu_tensor_misalignment(webgpu_context & ctx, const ggml_tensor * t) { + return ggml_webgpu_tensor_misalignment(t, ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment); +} + +static size_t ggml_webgpu_tensor_align_offset(const ggml_tensor * t, size_t alignment) { + size_t offset = ggml_webgpu_tensor_offset(t); + return offset & ~(alignment - 1); } static size_t ggml_webgpu_tensor_align_offset(webgpu_context & ctx, const ggml_tensor * t) { - size_t offset = ggml_webgpu_tensor_offset(t); - return offset & ~(ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment - 1); + return ggml_webgpu_tensor_align_offset(t, ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment); } -static size_t ggml_webgpu_tensor_binding_size(webgpu_context & ctx, ggml_tensor * t) { - return ROUNDUP_POW2(ggml_nbytes(t) + ggml_webgpu_tensor_misalignment(ctx, t), WEBGPU_STORAGE_BUF_BINDING_MULT); +static size_t ggml_webgpu_tensor_binding_size(const ggml_tensor * t, size_t alignment) { + return ROUNDUP_POW2(ggml_nbytes(t) + ggml_webgpu_tensor_misalignment(t, alignment), + WEBGPU_STORAGE_BUF_BINDING_MULT); +} + +static size_t ggml_webgpu_tensor_binding_size(webgpu_context & ctx, const ggml_tensor * t) { + return ggml_webgpu_tensor_binding_size(t, ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment); +} + +static bool ggml_webgpu_tensor_binding_overlap(const webgpu_global_context & global_ctx, + const ggml_tensor * a, + const ggml_tensor * b) { + if (a->buffer != b->buffer) { + return false; + } + + const size_t alignment = global_ctx->capabilities.limits.minStorageBufferOffsetAlignment; + const size_t a_offset = ggml_webgpu_tensor_align_offset(a, alignment); + const size_t b_offset = ggml_webgpu_tensor_align_offset(b, alignment); + return a_offset < b_offset + ggml_webgpu_tensor_binding_size(b, alignment) && + b_offset < a_offset + ggml_webgpu_tensor_binding_size(a, alignment); +} + +static bool ggml_webgpu_tensor_binding_overlap_range(const webgpu_global_context & global_ctx, + ggml_tensor * tensor, + ggml_backend_buffer_t buffer, + size_t offset, + size_t size) { + if (tensor->buffer != buffer) { + return false; + } + + const size_t alignment = global_ctx->capabilities.limits.minStorageBufferOffsetAlignment; + const size_t tensor_offset = ggml_webgpu_tensor_align_offset(tensor, alignment); + return tensor_offset < offset + size && offset < tensor_offset + ggml_webgpu_tensor_binding_size(tensor, alignment); } struct ggml_webgpu_merged_binding_range { @@ -1188,39 +1229,76 @@ static webgpu_encoded_op ggml_webgpu_ssm_scan(webgpu_context & ctx, ggml_webgpu_shader_lib_context shader_lib_ctx = {}; shader_lib_ctx.src0 = src0; shader_lib_ctx.src1 = src1; + shader_lib_ctx.src2 = src2; + shader_lib_ctx.src3 = src3; shader_lib_ctx.src4 = src4; shader_lib_ctx.src5 = src5; shader_lib_ctx.dst = dst; shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup; shader_lib_ctx.supports_subgroups = ctx->global_ctx->capabilities.supports_subgroups; + bool xbc_overlap = ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src1, src2) || + ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src1, src4) || + ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src1, src5) || + ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src2, src4) || + ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src2, src5) || + ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src4, src5); + bool a_overlap = false; + bool ids_overlap = false; + ggml_webgpu_merged_binding_range xbc_merged_range = {}; + if (xbc_overlap) { + xbc_merged_range = ggml_webgpu_tensor_merged_binding_range(ctx, { src1, src2, src4, src5 }); + a_overlap = ggml_webgpu_tensor_binding_overlap_range(ctx->global_ctx, src3, src1->buffer, + xbc_merged_range.offset, xbc_merged_range.size); + if (a_overlap) { + xbc_merged_range = ggml_webgpu_tensor_merged_binding_range(ctx, { src1, src2, src3, src4, src5 }); + } + ids_overlap = ggml_webgpu_tensor_binding_overlap_range(ctx->global_ctx, src6, src1->buffer, + xbc_merged_range.offset, xbc_merged_range.size); + if (ids_overlap) { + xbc_merged_range = + a_overlap ? ggml_webgpu_tensor_merged_binding_range(ctx, { src1, src2, src3, src4, src5, src6 }) : + ggml_webgpu_tensor_merged_binding_range(ctx, { src1, src2, src4, src5, src6 }); + } + } - webgpu_pipeline pipeline = ctx->shader_lib->get_ssm_scan_pipeline(shader_lib_ctx); - auto * decisions = static_cast(pipeline.context.get()); - const bool xbc_overlap = decisions->xbc_overlap; + webgpu_pipeline pipeline = + ctx->shader_lib->get_ssm_scan_pipeline(shader_lib_ctx, xbc_overlap, a_overlap, ids_overlap); + auto * decisions = static_cast(pipeline.context.get()); + xbc_overlap = decisions->xbc_overlap; + a_overlap = decisions->a_overlap; + ids_overlap = decisions->ids_overlap; uint32_t offset_x = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)); + uint32_t offset_dt = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src2) / ggml_type_size(src2->type)); + uint32_t offset_A = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src3) / ggml_type_size(src3->type)); uint32_t offset_B = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src4) / ggml_type_size(src4->type)); uint32_t offset_C = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src5) / ggml_type_size(src5->type)); + uint32_t offset_ids = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src6) / ggml_type_size(src6->type)); size_t xbc_bind_offset = 0; size_t xbc_bind_size = 0; if (xbc_overlap) { - const ggml_webgpu_merged_binding_range merged_range = - ggml_webgpu_tensor_merged_binding_range(ctx, { src1, src4, src5 }); - xbc_bind_offset = merged_range.offset; - xbc_bind_size = merged_range.size; - offset_x = ggml_webgpu_tensor_merged_element_offset(src1, merged_range); - offset_B = ggml_webgpu_tensor_merged_element_offset(src4, merged_range); - offset_C = ggml_webgpu_tensor_merged_element_offset(src5, merged_range); + xbc_bind_offset = xbc_merged_range.offset; + xbc_bind_size = xbc_merged_range.size; + offset_x = ggml_webgpu_tensor_merged_element_offset(src1, xbc_merged_range); + offset_dt = ggml_webgpu_tensor_merged_element_offset(src2, xbc_merged_range); + if (a_overlap) { + offset_A = ggml_webgpu_tensor_merged_element_offset(src3, xbc_merged_range); + } + offset_B = ggml_webgpu_tensor_merged_element_offset(src4, xbc_merged_range); + offset_C = ggml_webgpu_tensor_merged_element_offset(src5, xbc_merged_range); + if (ids_overlap) { + offset_ids = ggml_webgpu_tensor_merged_element_offset(src6, xbc_merged_range); + } } std::vector params = { (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)), offset_x, - (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src2) / ggml_type_size(src2->type)), - (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src3) / ggml_type_size(src3->type)), + offset_dt, + offset_A, offset_B, offset_C, - (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src6) / ggml_type_size(src6->type)), + offset_ids, (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)), (uint32_t) (src0->nb[1] / ggml_type_size(src0->type)), @@ -1260,10 +1338,19 @@ static webgpu_encoded_op ggml_webgpu_ssm_scan(webgpu_context & ctx, if (xbc_overlap) { entries.push_back( ggml_webgpu_make_bind_group_entry(1, ggml_webgpu_tensor_buf(src1), xbc_bind_offset, xbc_bind_size)); - entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, src2)); - entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 3, src3)); - entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 4, src6)); - entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 5, dst)); + if (ids_overlap) { + if (!a_overlap) { + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, src3)); + } + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, a_overlap ? 2 : 3, dst)); + } else if (a_overlap) { + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, src6)); + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 3, dst)); + } else { + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, src3)); + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 3, src6)); + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 4, dst)); + } } else { entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, src1)); entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, src2)); @@ -1381,11 +1468,10 @@ static std::optional ggml_webgpu_set_rows(webgpu_context & ct (uint32_t) (idx->ne[1]), (uint32_t) (idx->ne[2]) }; - std::vector entries = { - ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src), - ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, idx), - ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, dst), - }; + std::vector entries; + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src)); + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, idx)); + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, dst)); if (decisions->i64_idx) { entries.push_back(ggml_webgpu_make_bind_group_entry(3, ctx->set_rows_dev_error_buf, 0, @@ -1892,7 +1978,7 @@ static ggml_webgpu_flash_attn_op ggml_webgpu_flash_attn_prepare(webgpu_context & op.has_mask = mask != nullptr; op.has_sinks = sinks != nullptr; - op.kv_overlap = ggml_webgpu_tensor_overlap(K, V); + op.kv_overlap = ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, K, V); uint32_t offset_k = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, K) / ggml_type_size(K->type)); uint32_t offset_v = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, V) / ggml_type_size(V->type)); @@ -1964,7 +2050,7 @@ static uint32_t ggml_webgpu_flash_attn_vec_nwg(uint32_t vec_nwg_cap, uint32_t kv } static webgpu_encoded_op ggml_webgpu_flash_attn_direct(webgpu_context & ctx, const ggml_webgpu_flash_attn_op & op) { - webgpu_pipeline pipeline = ctx->shader_lib->get_flash_attn_pipeline(op.shader_lib_ctx); + webgpu_pipeline pipeline = ctx->shader_lib->get_flash_attn_pipeline(op.shader_lib_ctx, op.kv_overlap); auto * decisions = static_cast(pipeline.context.get()); uint32_t wg_per_head = CEIL_DIV(op.shader_lib_ctx.src0->ne[1], decisions->q_tile); uint32_t wg_x = wg_per_head * op.shader_lib_ctx.src0->ne[2] * op.shader_lib_ctx.src0->ne[3]; @@ -1979,7 +2065,7 @@ static webgpu_encoded_op ggml_webgpu_flash_attn_vec(webgpu_context & ct ggml_tensor * sinks, ggml_tensor * dst, ggml_webgpu_flash_attn_op op) { - webgpu_pipeline pipeline = ctx->shader_lib->get_flash_attn_vec_pipeline(op.shader_lib_ctx); + webgpu_pipeline pipeline = ctx->shader_lib->get_flash_attn_vec_pipeline(op.shader_lib_ctx, op.kv_overlap); auto * decisions = static_cast(pipeline.context.get()); wgpu::Buffer blk_buf = {}; @@ -2249,8 +2335,9 @@ static webgpu_encoded_op ggml_webgpu_binary_op(webgpu_context & ctx, shader_lib_ctx.dst = dst; shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup; - webgpu_pipeline pipeline = ctx->shader_lib->get_binary_pipeline(shader_lib_ctx); - auto * decisions = static_cast(pipeline.context.get()); + const bool src_overlap = ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src0, src1); + webgpu_pipeline pipeline = ctx->shader_lib->get_binary_pipeline(shader_lib_ctx, src_overlap); + auto * decisions = static_cast(pipeline.context.get()); uint32_t ne = (uint32_t) ggml_nelements(dst); @@ -2372,6 +2459,9 @@ static webgpu_encoded_op ggml_webgpu_concat(webgpu_context & ctx, ggml_tensor * dst) { uint32_t ne = (uint32_t) ggml_nelements(dst); uint32_t dim = (uint32_t) dst->op_params[0]; + if (ggml_nbytes(src0) == 0 && ggml_nbytes(src1) == 0) { + return {}; + } ggml_webgpu_shader_lib_context shader_lib_ctx = {}; shader_lib_ctx.src0 = src0; @@ -2379,20 +2469,34 @@ static webgpu_encoded_op ggml_webgpu_concat(webgpu_context & ctx, shader_lib_ctx.dst = dst; shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup; - webgpu_pipeline pipeline = ctx->shader_lib->get_concat_pipeline(shader_lib_ctx); - auto * decisions = static_cast(pipeline.context.get()); + const bool src_overlap = ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src0, src1) || + ggml_nbytes(src0) == 0 || ggml_nbytes(src1) == 0; + webgpu_pipeline pipeline = ctx->shader_lib->get_concat_pipeline(shader_lib_ctx, src_overlap); + auto * decisions = static_cast(pipeline.context.get()); uint32_t offset_src0 = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)); uint32_t offset_src1 = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)); size_t merged_offset = 0; size_t merged_size = 0; if (decisions->src_overlap) { - const ggml_webgpu_merged_binding_range merged_range = - ggml_webgpu_tensor_merged_binding_range(ctx, { src0, src1 }); - merged_offset = merged_range.offset; - merged_size = merged_range.size; - offset_src0 = ggml_webgpu_tensor_merged_element_offset(src0, merged_range); - offset_src1 = ggml_webgpu_tensor_merged_element_offset(src1, merged_range); + if (ggml_nbytes(src0) == 0) { + merged_offset = ggml_webgpu_tensor_align_offset(ctx, src1); + merged_size = ggml_webgpu_tensor_binding_size(ctx, src1); + offset_src0 = 0; + offset_src1 = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)); + } else if (ggml_nbytes(src1) == 0) { + merged_offset = ggml_webgpu_tensor_align_offset(ctx, src0); + merged_size = ggml_webgpu_tensor_binding_size(ctx, src0); + offset_src0 = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)); + offset_src1 = 0; + } else { + const ggml_webgpu_merged_binding_range merged_range = + ggml_webgpu_tensor_merged_binding_range(ctx, { src0, src1 }); + merged_offset = merged_range.offset; + merged_size = merged_range.size; + offset_src0 = ggml_webgpu_tensor_merged_element_offset(src0, merged_range); + offset_src1 = ggml_webgpu_tensor_merged_element_offset(src1, merged_range); + } } std::vector params = { ne, @@ -2518,8 +2622,9 @@ static std::optional ggml_webgpu_rms_norm_mul(webgpu_context shader_lib_ctx.dst = dst; shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup; - webgpu_pipeline pipeline = ctx->shader_lib->get_rms_norm_mul_pipeline(shader_lib_ctx); - auto * decisions = static_cast(pipeline.context.get()); + const bool src_overlap = ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, rn_src, mul_src); + webgpu_pipeline pipeline = ctx->shader_lib->get_rms_norm_mul_pipeline(shader_lib_ctx, src_overlap); + auto * decisions = static_cast(pipeline.context.get()); if (decisions->src_overlap) { const ggml_webgpu_merged_binding_range merged_range = @@ -2678,15 +2783,30 @@ static webgpu_encoded_op ggml_webgpu_glu(webgpu_context & ctx, shader_lib_ctx.dst = dst; shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup; - webgpu_pipeline pipeline = ctx->shader_lib->get_glu_pipeline(shader_lib_ctx); + const bool src_overlap = src1 != nullptr && ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src0, src1); + webgpu_pipeline pipeline = ctx->shader_lib->get_glu_pipeline(shader_lib_ctx, src_overlap); - auto * decisions = static_cast(pipeline.context.get()); + auto * decisions = static_cast(pipeline.context.get()); const int split = (src1 != nullptr); + uint32_t offset_src0 = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)); + uint32_t offset_src1 = + src1 != nullptr ? (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)) : 0; + size_t merged_offset = 0; + size_t merged_size = 0; + if (decisions->src_overlap) { + const ggml_webgpu_merged_binding_range merged_range = + ggml_webgpu_tensor_merged_binding_range(ctx, { src0, src1 }); + merged_offset = merged_range.offset; + merged_size = merged_range.size; + offset_src0 = ggml_webgpu_tensor_merged_element_offset(src0, merged_range); + offset_src1 = ggml_webgpu_tensor_merged_element_offset(src1, merged_range); + } + std::vector params = { - (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)), - src1 != nullptr ? (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)) : 0, + offset_src0, + offset_src1, (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)), (uint32_t) (src0->nb[1] / ggml_type_size(src0->type)), (uint32_t) (src0->nb[2] / ggml_type_size(src0->type)), @@ -2709,11 +2829,15 @@ static webgpu_encoded_op ggml_webgpu_glu(webgpu_context & ctx, ggml_webgpu_u32_from_f32(ggml_get_op_params_f32(dst, 3)), // limit, for swiglu_oai }; - std::vector entries = { - ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src0), - }; - uint32_t dst_binding = 1; - if (split) { + std::vector entries; + uint32_t dst_binding = 1; + if (decisions->src_overlap) { + entries.push_back( + ggml_webgpu_make_bind_group_entry(0, ggml_webgpu_tensor_buf(src0), merged_offset, merged_size)); + } else { + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src0)); + } + if (split && !decisions->src_overlap) { dst_binding = 2; entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, src1)); } @@ -4285,8 +4409,8 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const if (!supports_op) { break; } - if (ggml_webgpu_tensor_overlap(src1, src2) && src1->type != src2->type && - !ggml_is_quantized(src1->type) && !ggml_is_quantized(src2->type)) { + if (ggml_webgpu_tensor_binding_overlap(ctx->webgpu_global_ctx, src1, src2) && + src1->type != src2->type && !ggml_is_quantized(src1->type) && !ggml_is_quantized(src2->type)) { supports_op = false; break; } diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/glu.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/glu.wgsl index e6d7608cec..d03f1c207d 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/glu.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/glu.wgsl @@ -96,7 +96,22 @@ struct Params { @group(0) @binding(0) var src0: array; -#ifdef NO_SPLIT +#ifdef SRC_OVERLAP +@group(0) @binding(1) +var dst: array; + +@group(0) @binding(2) +var params: Params; + +fn a_value(base: u32) -> DataType { + return src0[base]; +} + +fn b_value(base: u32) -> DataType { + return src0[base]; +} + +#elif defined(NO_SPLIT) @group(0) @binding(1) var dst: array; diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/ssm_scan.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/ssm_scan.wgsl index 05761dec35..66bfdd6401 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/ssm_scan.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/ssm_scan.wgsl @@ -46,12 +46,29 @@ struct Params { @group(0) @binding(0) var s_in: array; #ifdef XBC_OVERLAP -@group(0) @binding(1) var x_B_C_merged: array; -@group(0) @binding(2) var dt: array; -@group(0) @binding(3) var A: array; -@group(0) @binding(4) var ids: array; -@group(0) @binding(5) var dst: array; -@group(0) @binding(6) var params: Params; +#ifdef IDS_OVERLAP +@group(0) @binding(1) var x_dt_B_C_ids_merged: array; +#ifdef A_OVERLAP +@group(0) @binding(2) var dst: array; +@group(0) @binding(3) var params: Params; +#else +@group(0) @binding(2) var A: array; +@group(0) @binding(3) var dst: array; +@group(0) @binding(4) var params: Params; +#endif +#else +@group(0) @binding(1) var x_dt_B_C_merged: array; +#ifdef A_OVERLAP +@group(0) @binding(2) var ids: array; +@group(0) @binding(3) var dst: array; +@group(0) @binding(4) var params: Params; +#else +@group(0) @binding(2) var A: array; +@group(0) @binding(3) var ids: array; +@group(0) @binding(4) var dst: array; +@group(0) @binding(5) var params: Params; +#endif +#endif #else @group(0) @binding(1) var x: array; @group(0) @binding(2) var dt: array; @@ -71,6 +88,24 @@ fn reduce_base(token_in_tile: u32) -> u32 { return token_in_tile * WG_SIZE; } +#ifdef XBC_OVERLAP +fn read_merged_f32(idx: u32) -> f32 { +#ifdef IDS_OVERLAP + return bitcast(x_dt_B_C_ids_merged[idx]); +#else + return x_dt_B_C_merged[idx]; +#endif +} +#endif + +fn read_state_slot(i3: u32) -> u32 { +#ifdef IDS_OVERLAP + return x_dt_B_C_ids_merged[params.offset_ids + i3]; +#else + return u32(ids[params.offset_ids + i3]); +#endif +} + @compute @workgroup_size(WG_SIZE) fn main( @builtin(local_invocation_id) local_id: vec3, @@ -90,13 +125,18 @@ fn main( let ir = head_seq % params.n_head; let i3 = head_seq / params.n_head; - let state_slot = u32(ids[params.offset_ids + i3]); + let state_slot = read_state_slot(i3); let g = ir / (params.n_head / params.n_group); let s_idx = params.offset_s + tid + i1 * params.stride_s1 + ir * params.stride_s2 + state_slot * params.stride_s3; var s_prev = s_in[s_idx]; - let A0 = A[params.offset_A + (tid % params.a_ne0) + ir * params.stride_A1]; + let a_idx = params.offset_A + (tid % params.a_ne0) + ir * params.stride_A1; +#ifdef A_OVERLAP + let A0 = read_merged_f32(a_idx); +#else + let A0 = A[a_idx]; +#endif for (var token_base = 0u; token_base < params.n_seq_tokens; token_base += TOKENS_PER_TILE) { if (tid < TOKENS_PER_TILE) { @@ -104,11 +144,15 @@ fn main( if (token < params.n_seq_tokens) { let x_idx = params.offset_x + i1 + ir * params.stride_x1 + token * params.stride_x2 + i3 * params.stride_x3; let dt_idx = params.offset_dt + ir + token * params.stride_dt1 + i3 * params.stride_dt2; +#ifdef XBC_OVERLAP + let dt0 = read_merged_f32(dt_idx); +#else let dt0 = dt[dt_idx]; +#endif let dtsp = select(log(1.0 + exp(dt0)), dt0, dt0 > 20.0); shared_dtsp[tid] = dtsp; #ifdef XBC_OVERLAP - shared_x_dt[tid] = x_B_C_merged[x_idx] * dtsp; + shared_x_dt[tid] = read_merged_f32(x_idx) * dtsp; #else shared_x_dt[tid] = x[x_idx] * dtsp; #endif @@ -130,7 +174,7 @@ fn main( let b_idx = params.offset_B + tid + g * params.stride_B1 + token * params.stride_B2 + i3 * params.stride_B3; let c_idx = params.offset_C + tid + g * params.stride_C1 + token * params.stride_C2 + i3 * params.stride_C3; #ifdef XBC_OVERLAP - let s = s_prev * dA + x_B_C_merged[b_idx] * x_dt; + let s = s_prev * dA + read_merged_f32(b_idx) * x_dt; #else let s = s_prev * dA + B[b_idx] * x_dt; #endif @@ -138,7 +182,7 @@ fn main( #ifdef USE_SUBGROUP_REDUCTION #ifdef XBC_OVERLAP - let subgroup_partial = subgroupAdd(s * x_B_C_merged[c_idx]); + let subgroup_partial = subgroupAdd(s * read_merged_f32(c_idx)); #else let subgroup_partial = subgroupAdd(s * C[c_idx]); #endif @@ -147,7 +191,7 @@ fn main( } #else #ifdef XBC_OVERLAP - shared_reduce[reduce_idx] = s * x_B_C_merged[c_idx]; + shared_reduce[reduce_idx] = s * read_merged_f32(c_idx); #else shared_reduce[reduce_idx] = s * C[c_idx]; #endif diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index a7d1fe7d94..59191c663e 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -7854,7 +7854,9 @@ void ggml_set_input(struct ggml_tensor * tensor) { } void ggml_set_output(struct ggml_tensor * tensor) { - tensor->flags |= GGML_TENSOR_FLAG_OUTPUT; + for (struct ggml_tensor * cur = tensor; cur != NULL; cur = cur->view_src) { + cur->flags |= GGML_TENSOR_FLAG_OUTPUT; + } } void ggml_set_param(struct ggml_tensor * tensor) { diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 0c7d7e7991..5fbce9d04d 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -373,6 +373,7 @@ class Keys: FEED_FORWARD_LENGTH = "clip.audio.feed_forward_length" PROJECTION_DIM = "clip.audio.projection_dim" BLOCK_COUNT = "clip.audio.block_count" + SUBSAMPLING_FACTOR = "clip.audio.subsampling_factor" CHUNK_SIZE = "clip.audio.chunk_size" CONV_KERNEL_SIZE = "clip.audio.conv_kernel_size" MAX_POS_EMB = "clip.audio.max_pos_emb" @@ -994,6 +995,10 @@ class MODEL_TENSOR(IntEnum): # eagle3 FC = auto() # feature fusion layer D2T = auto() # draft to target vocabulary mapping + # dspark + DSPARK_MARKOV_W1 = auto() # markov head: prev-token embed + DSPARK_MARKOV_W2 = auto() # markov head: bias projection + DSPARK_CONF_PROJ = auto() # confidence head # lfm2 audio A_ENC_NORM_CONV = auto() A_ENC_LINEAR_POS = auto() @@ -1004,6 +1009,10 @@ class MODEL_TENSOR(IntEnum): A_ENC_CONV_NORM = auto() # SSM conv A_ENC_CONV_PW1 = auto() A_ENC_CONV_PW2 = auto() + A_ENC_CONV_NORM_MEAN = auto() # parakeet + A_ENC_CONV_NORM_VAR = auto() # parakeet + A_ENC_MEL_FILTERS = auto() # parakeet + A_ENC_WINDOW = auto() # parakeet A_CTC_OUT = auto() A_CTC_OUT_MID = auto() A_ENC_ATTN_REL_POS_EMB = auto() @@ -1599,6 +1608,10 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.A_ENC_CONV_NORM: "a.blk.{bid}.conv_norm", MODEL_TENSOR.A_ENC_CONV_PW1: "a.blk.{bid}.conv_pw1", MODEL_TENSOR.A_ENC_CONV_PW2: "a.blk.{bid}.conv_pw2", + MODEL_TENSOR.A_ENC_CONV_NORM_MEAN: "a.blk.{bid}.conv_norm_mean", + MODEL_TENSOR.A_ENC_CONV_NORM_VAR: "a.blk.{bid}.conv_norm_var", + MODEL_TENSOR.A_ENC_MEL_FILTERS: "a.mel_filters", + MODEL_TENSOR.A_ENC_WINDOW: "a.window", MODEL_TENSOR.A_CTC_OUT: "a.enc_ctc_out", MODEL_TENSOR.A_CTC_OUT_MID: "a.enc_ctc_out_mid", MODEL_TENSOR.A_ENC_ATTN_REL_POS_EMB: "a.blk.{bid}.attn_rel_pos_emb", @@ -1629,6 +1642,9 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD: "blk.{bid}.nextn.shared_head_head", MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: "blk.{bid}.nextn.shared_head_norm", MODEL_TENSOR.FC: "fc", + MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1", + MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2", + MODEL_TENSOR.DSPARK_CONF_PROJ: "conf_proj", MODEL_TENSOR.D2T: "d2t", } @@ -1820,6 +1836,10 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.A_ENC_SE_CONV2, MODEL_TENSOR.A_ENC_ASP_ATTN, MODEL_TENSOR.A_ENC_ASP_TDNN, + MODEL_TENSOR.A_ENC_CONV_NORM_MEAN, + MODEL_TENSOR.A_ENC_CONV_NORM_VAR, + MODEL_TENSOR.A_ENC_MEL_FILTERS, + MODEL_TENSOR.A_ENC_WINDOW, MODEL_TENSOR.A_MM_INP_PROJ, MODEL_TENSOR.A_MM_SOFT_EMB_NORM, MODEL_TENSOR.A_MM_EMBEDDING, @@ -4379,6 +4399,10 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_UP, MODEL_TENSOR.FC, MODEL_TENSOR.ENC_OUTPUT_NORM, + # optional DSpark heads + MODEL_TENSOR.DSPARK_MARKOV_W1, + MODEL_TENSOR.DSPARK_MARKOV_W2, + MODEL_TENSOR.DSPARK_CONF_PROJ, ], MODEL_ARCH.MISTRAL4: [ MODEL_TENSOR.TOKEN_EMBD, @@ -4884,6 +4908,7 @@ class VisionProjectorType: NEMOTRON_V2_VL = "nemotron_v2_vl" QWEN3TTS_SPKENC = "qwen3tts_spkenc" # audio: ECAPA-TDNN speaker encoder HUNYUANVL = "hunyuanvl" + PARAKEET = "parakeet" # audio MINIMAXM3 = "minimax_m3" MINICPMV4_6 = "minicpmv4_6" GRANITE_SPEECH = "granite_speech" # audio diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 657ed69b68..3aa4f049f2 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1374,6 +1374,9 @@ class GGUFWriter: def add_audio_stack_factor(self, value: int) -> None: self.add_uint32(Keys.ClipAudio.Projector.STACK_FACTOR, value) + def add_audio_subsampling_factor(self, value: int) -> None: + self.add_uint32(Keys.ClipAudio.SUBSAMPLING_FACTOR, value) + def add_audio_chunk_size(self, value: int) -> None: self.add_uint32(Keys.ClipAudio.CHUNK_SIZE, value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 5bc624ac4c..7892342e47 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -1304,6 +1304,18 @@ class TensorNameMap: "model.fc", # dflash ), + MODEL_TENSOR.DSPARK_MARKOV_W1: ( + "model.markov_head.markov_w1", # dspark + ), + + MODEL_TENSOR.DSPARK_MARKOV_W2: ( + "model.markov_head.markov_w2", # dspark + ), + + MODEL_TENSOR.DSPARK_CONF_PROJ: ( + "model.confidence_head.proj", # dspark + ), + MODEL_TENSOR.CLS: ( "classifier", # jina "classifier.dense", # roberta @@ -2095,6 +2107,7 @@ class TensorNameMap: "conformer.pre_encode.conv.{bid}", # lfm2 "model.audio_tower.subsample_conv_projection.conv_{bid}.conv", # gemma3n "conformer.subsample_conv_projection.layer{bid}.conv", # gemma4 + "sound_encoder.encoder.subsampling.layers.{bid}", # parakeet "encoder.conv{bid}", # mimo-audio-tokenizer "speaker_encoder.blocks.{bid}.conv", # qwen3tts speaker encoder (only bid=0, the stem TDNN) ), @@ -2130,6 +2143,7 @@ class TensorNameMap: "conformer.layers.{bid}.self_attn.linear_q", # lfm2 "conformer.layers.{bid}.attention.attn.q_proj", # gemma3n "conformer.layers.{bid}.self_attn.q_proj", # gemma4 + "sound_encoder.encoder.layers.{bid}.self_attn.q_proj", # parakeet "encoder.layers.{bid}.attn.to_q", # granite_speech "encoder.layers.{bid}.self_attn.q_proj", # mimo-audio-tokenizer ), @@ -2139,6 +2153,7 @@ class TensorNameMap: "conformer.layers.{bid}.self_attn.linear_k", # lfm2 "conformer.layers.{bid}.attention.attn.k_proj", # gemma3n "conformer.layers.{bid}.self_attn.k_proj", # gemma4 + "sound_encoder.encoder.layers.{bid}.self_attn.k_proj", # parakeet "encoder.layers.{bid}.attn.to_k", # granite_speech (split from to_kv) "encoder.layers.{bid}.self_attn.k_proj", # mimo-audio-tokenizer ), @@ -2148,6 +2163,7 @@ class TensorNameMap: "conformer.layers.{bid}.self_attn.linear_v", # lfm2 "conformer.layers.{bid}.attention.attn.v_proj", # gemma3n "conformer.layers.{bid}.self_attn.v_proj", # gemma4 + "sound_encoder.encoder.layers.{bid}.self_attn.v_proj", # parakeet "encoder.layers.{bid}.attn.to_v", # granite_speech (split from to_kv) "encoder.layers.{bid}.self_attn.v_proj", # mimo-audio-tokenizer ), @@ -2177,6 +2193,7 @@ class TensorNameMap: "audio_tower.layers.{bid}.self_attn_layer_norm", # ultravox "conformer.layers.{bid}.norm_self_att", # lfm2 "conformer.layers.{bid}.attention.pre_attn_norm", # gemma3n + "sound_encoder.encoder.layers.{bid}.norm_self_att", # parakeet "encoder.layers.{bid}.attn.pre_norm", # granite_speech "encoder.layers.{bid}.self_attn_layer_norm", # mimo-audio-tokenizer ), @@ -2186,6 +2203,7 @@ class TensorNameMap: "conformer.layers.{bid}.self_attn.linear_out", # lfm2 "conformer.layers.{bid}.attention.post", # gemma3n "conformer.layers.{bid}.self_attn.post", # gemma4 + "sound_encoder.encoder.layers.{bid}.self_attn.o_proj", # parakeet "encoder.layers.{bid}.attn.to_out", # granite_speech "encoder.layers.{bid}.self_attn.out_proj", # mimo-audio-tokenizer ), @@ -2194,6 +2212,7 @@ class TensorNameMap: "audio_tower.layers.{bid}.final_layer_norm", # ultravox "conformer.layers.{bid}.norm_out", # lfm2 "conformer.layers.{bid}.attention.post_norm", # gemma3n + "sound_encoder.encoder.layers.{bid}.norm_out", # parakeet "encoder.layers.{bid}.post_norm", # granite_speech "encoder.layers.{bid}.final_layer_norm", # mimo-audio-tokenizer ), @@ -2202,6 +2221,7 @@ class TensorNameMap: "conformer.layers.{bid}.norm_feed_forward1", # lfm2 "conformer.layers.{bid}.ffw_layer_start.pre_layer_norm", # gemma3n "conformer.layers.{bid}.feed_forward1.pre_layer_norm", # gemma4 + "sound_encoder.encoder.layers.{bid}.norm_feed_forward1", # parakeet "encoder.layers.{bid}.ff1.pre_norm", # granite_speech ), @@ -2219,6 +2239,7 @@ class TensorNameMap: "conformer.layers.{bid}.feed_forward1.linear1", # lfm2 "conformer.layers.{bid}.ffw_layer_start.ffw_layer_1", # gemma3n "conformer.layers.{bid}.feed_forward1.ffw_layer_1", # gemma4 + "sound_encoder.encoder.layers.{bid}.feed_forward1.linear1", # parakeet "encoder.layers.{bid}.ff1.up_proj", # granite_speech "encoder.layers.{bid}.fc1", # mimo-audio-tokenizer ), @@ -2230,6 +2251,7 @@ class TensorNameMap: "conformer.layers.{bid}.feed_forward1.linear2", # lfm2 "conformer.layers.{bid}.ffw_layer_start.ffw_layer_2", # gemma3n "conformer.layers.{bid}.feed_forward1.ffw_layer_2", # gemma4 + "sound_encoder.encoder.layers.{bid}.feed_forward1.linear2", # parakeet "encoder.layers.{bid}.ff1.down_proj", # granite_speech "encoder.layers.{bid}.fc2", # mimo-audio-tokenizer ), @@ -2238,6 +2260,7 @@ class TensorNameMap: "conformer.layers.{bid}.feed_forward2.linear1", # lfm2 "conformer.layers.{bid}.ffw_layer_end.ffw_layer_1", # gemma3n "conformer.layers.{bid}.feed_forward2.ffw_layer_1", # gemma4 + "sound_encoder.encoder.layers.{bid}.feed_forward2.linear1", # parakeet "encoder.layers.{bid}.ff2.up_proj", # granite_speech ), @@ -2245,6 +2268,7 @@ class TensorNameMap: "conformer.layers.{bid}.feed_forward2.linear2", # lfm2 "conformer.layers.{bid}.ffw_layer_end.ffw_layer_2", # gemma3n "conformer.layers.{bid}.feed_forward2.ffw_layer_2", # gemma4 + "sound_encoder.encoder.layers.{bid}.feed_forward2.linear2", # parakeet "encoder.layers.{bid}.ff2.down_proj", # granite_speech ), @@ -2252,6 +2276,7 @@ class TensorNameMap: "conformer.layers.{bid}.norm_feed_forward2", # lfm2 "conformer.layers.{bid}.ffw_layer_end.pre_layer_norm", # gemma3n "conformer.layers.{bid}.feed_forward2.pre_layer_norm", # gemma4 + "sound_encoder.encoder.layers.{bid}.norm_feed_forward2", # parakeet "encoder.layers.{bid}.ff2.pre_norm", # granite_speech ), @@ -2280,20 +2305,24 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_LINEAR_POS: ( "conformer.layers.{bid}.self_attn.linear_pos", # lfm2 "conformer.layers.{bid}.attention.attn.relative_position_embedding.pos_proj", # gemma3n + "sound_encoder.encoder.layers.{bid}.self_attn.relative_k_proj", # parakeet ), MODEL_TENSOR.A_ENC_POS_BIAS_U: ( "conformer.layers.{bid}.self_attn.pos_bias_u", # lfm2 + "sound_encoder.encoder.layers.{bid}.self_attn.bias_u", # parakeet ), MODEL_TENSOR.A_ENC_POS_BIAS_V: ( "conformer.layers.{bid}.self_attn.pos_bias_v", # lfm2 + "sound_encoder.encoder.layers.{bid}.self_attn.bias_v", # parakeet ), MODEL_TENSOR.A_ENC_OUT: ( "conformer.pre_encode.out", # lfm2 "model.audio_tower.subsample_conv_projection.input_proj_linear", # gemma3n (note: it should be A_ENC_INP_PROJ, this is a mistake; it should be corrected in C++ code when it's supported) "conformer.output_proj", # gemma4 + "sound_encoder.encoder.subsampling.linear", # parakeet ), # note: some tensors below has "audio." pseudo-prefix, to prevent conflicts with vision tensors @@ -2303,6 +2332,7 @@ class TensorNameMap: "audio.multi_modal_projector.linear_{bid}", # ultravox, meralion "audio_adapter.model.{bid}", # lfm2 "audio_tower.proj{bid}", # qwen3omni + "sound_projection.linear{bid}", # parakeet (linear1, linear2) ), MODEL_TENSOR.A_MMPROJ_FC: ( @@ -2314,6 +2344,7 @@ class TensorNameMap: MODEL_TENSOR.A_MM_NORM_PRE: ( "audio.multi_modal_projector.ln_pre", # ultravox + "sound_projection.norm", # parakeet ), MODEL_TENSOR.A_MM_NORM_MID: ( @@ -2359,18 +2390,29 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_CONV_DW: ( "conformer.layers.{bid}.conv.depthwise_conv", # lfm2 "conformer.layers.{bid}.lconv1d.depthwise_conv1d", # gemma3n + "sound_encoder.encoder.layers.{bid}.conv.depthwise_conv", # parakeet "encoder.layers.{bid}.conv.depth_conv.conv", # granite_speech ), MODEL_TENSOR.A_ENC_CONV_NORM: ( "conformer.layers.{bid}.conv.batch_norm", # lfm2 "conformer.layers.{bid}.lconv1d.pre_layer_norm", # gemma3n + "sound_encoder.encoder.layers.{bid}.conv.norm", # parakeet + ), + + MODEL_TENSOR.A_ENC_CONV_NORM_MEAN: ( + "sound_encoder.encoder.layers.{bid}.conv.norm.running_mean", # parakeet + ), + + MODEL_TENSOR.A_ENC_CONV_NORM_VAR: ( + "sound_encoder.encoder.layers.{bid}.conv.norm.running_var", # parakeet "encoder.layers.{bid}.conv.batch_norm", # granite_speech ), MODEL_TENSOR.A_ENC_CONV_PW1: ( "conformer.layers.{bid}.conv.pointwise_conv1", # lfm2 "conformer.layers.{bid}.lconv1d.linear_start", # gemma3n + "sound_encoder.encoder.layers.{bid}.conv.pointwise_conv1", # parakeet "encoder.layers.{bid}.conv.up_conv", # granite_speech "speaker_encoder.blocks.{bid}.tdnn1.conv", # qwen3tts speaker encoder ), @@ -2378,6 +2420,7 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_CONV_PW2: ( "conformer.layers.{bid}.conv.pointwise_conv2", # lfm2 "conformer.layers.{bid}.lconv1d.linear_end", # gemma3n + "sound_encoder.encoder.layers.{bid}.conv.pointwise_conv2", # parakeet "encoder.layers.{bid}.conv.down_conv", # granite_speech "speaker_encoder.blocks.{bid}.tdnn2.conv", # qwen3tts speaker encoder ), @@ -2401,6 +2444,7 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_NORM_CONV: ( "conformer.layers.{bid}.norm_conv", # lfm2 "conformer.layers.{bid}.lconv1d.conv_norm", # gemma3n + "sound_encoder.encoder.layers.{bid}.norm_conv", # parakeet "encoder.layers.{bid}.conv.norm", # granite_speech ), @@ -2412,6 +2456,14 @@ class TensorNameMap: "conformer.layers.{bid}.attention.attn.per_dim_scale", # gemma4 ), + MODEL_TENSOR.A_ENC_MEL_FILTERS: ( + "sound_encoder.encoder.feature_extractor.featurizer.fb", # parakeet + ), + + MODEL_TENSOR.A_ENC_WINDOW: ( + "sound_encoder.encoder.feature_extractor.featurizer.window", # parakeet + ), + MODEL_TENSOR.A_MM_EMBEDDING: ( "model.embed_audio.embedding", # gemma3n ), diff --git a/models/templates/MiniMax-M3.jinja b/models/templates/MiniMax-M3.jinja new file mode 100644 index 0000000000..93022eb9ce --- /dev/null +++ b/models/templates/MiniMax-M3.jinja @@ -0,0 +1,247 @@ +{# ---------- special token variables ---------- #} +{%- set ns_token = ']<]minimax[>[' -%} +{%- set bod_token = ']~!b[' -%} +{%- set bos_token = ']~b]' -%} +{%- set eos_token = '[e~[' -%} +{%- set toolcall_begin_token = ns_token ~ '' -%} +{%- set toolcall_end_token = ns_token ~ '' -%} +{%- set think_begin_token = '' -%} +{%- set think_end_token = '' -%} +{%- set image_token = ']<]image[>[' -%} +{%- set video_token = ']<]video[>[' -%} +{#- Thinking mode: "enabled" / "disabled" / "adaptive" / not defined -#} +{#- Recursive XML renderer for tool_call arguments ======================== -#} +{#- None values are intentionally skipped in mapping iteration so that + `null` (which would round-trip to the literal string "null") + never appears in the rendered tool_call. The convention is: omit the + field entirely. The top-level `_args` loop applies the same rule. + The `val is none` branch below is a safety net only — upstream cleaning + (drop_none_in_tool_arguments) should ensure no None ever reaches here. -#} +{%- macro to_xml(val, ns) -%} +{%- if val is mapping -%} +{%- for k, v in val.items() if v is not none -%} +{{ ns }}<{{ k }}>{{ to_xml(v, ns) }}{{ ns }} +{%- endfor -%} +{%- elif val is iterable and val is not string -%} +{%- for item in val -%} +{{ ns }}{{ to_xml(item, ns) }}{{ ns }} +{%- endfor -%} +{%- elif val is none -%} +{#- Should be unreachable when upstream cleaning is applied. -#} +{%- elif val is boolean -%} +{{ val | tojson }} +{%- else -%} +{{ val }} +{%- endif -%} +{%- endmacro -%} +{#- Tool Rendering Functions ============================================== -#} +{%- macro render_tool_namespace(namespace_name, tool_list) -%} +{%- for tool in tool_list -%} +{{ tool.function | tojson(ensure_ascii=False) }} +{% endfor -%} +{%- endmacro -%} +{%- macro visible_text(content) -%} + {%- if content is string -%} + {{ content }} + {%- elif content is iterable and content is not mapping -%} + {%- for item in content -%} + {%- if item is mapping and item.type == 'text' -%} + {{- item.text }} + {%- elif item is mapping and item.type == 'image' -%} + {{- image_token }} + {%- elif item is mapping and item.type == 'video' -%} + {{- video_token}} + {%- elif item is string -%} + {{- item }} + {%- endif -%} + {%- endfor -%} + {%- elif content is none -%} + {{- '' }} + {%- else -%} + {{- content }} + {%- endif -%} +{%- endmacro -%} +{#- System Message Construction ============================================ -#} +{%- macro build_system_message(system_message) -%} + {%- if system_message and system_message.content -%} + {{- visible_text(system_message.content) }} + {%- else -%} + {{- 'Your model version is MiniMax-M3, developed by MiniMax. Knowledge cutoff: January 2026. Founded in early 2022, MiniMax is a global AI foundation model company committed to advancing the frontiers of AI towards AGI.' }} + {%- endif -%} + + {#- Thinking mode instructions -#} + {{- '\n\n\n' }} + {{- 'You have a thinking capability that allows you to reason step by step before responding. When thinking is enabled, wrap your reasoning in ' ~ think_begin_token ~ think_end_token ~ ' tags before your response. When thinking is disabled, begin your response directly after the ' ~ think_end_token ~ ' prefix. When thinking is adaptive, decide on your own whether to think for the current turn.\n' }} + {%- if thinking_mode is defined -%} + {%- if thinking_mode == "enabled" -%} + {{- 'Current thinking mode: enabled. You MUST think step by step before every response, including after receiving function/tool results.\n' }} + {%- elif thinking_mode == "disabled" -%} + {{- 'Current thinking mode: disabled. Do not output any thinking process.\n' }} + {%- elif thinking_mode == "adaptive" -%} + {{- 'Current thinking mode: adaptive. You are encouraged to think for complex decision-making, multi-step reasoning, or when analyzing function/tool results.\n' }} + {%- endif -%} + {%- else -%} + {{- 'Current thinking mode: adaptive. You are encouraged to think for complex decision-making, multi-step reasoning, or when analyzing function/tool results.\n' }} + {%- endif -%} + {{- '' }} +{%- endmacro -%} +{%- macro build_developer_message(developer_message) -%} + {%- if developer_message and developer_message.content -%} + {{- visible_text(developer_message.content) }} + {%- else -%} + {%- if model_identity is not defined -%} + {%- set model_identity = "You are a helpful assistant." -%} + {%- endif -%} + {{- model_identity }} + {%- endif -%} +{%- endmacro -%} +{#- Main Template Logic ================================================= -#} +{#- Role mapping: root -> system sp (high priority), system/developer -> developer sp (low priority) -#} +{%- set system_message = none -%} +{%- set developer_message = none -%} +{%- set conversation_messages = messages -%} +{%- if messages and messages[0].role == "root" -%} + {%- set system_message = messages[0] -%} + {%- set conversation_messages = messages[1:] -%} + {%- if conversation_messages and conversation_messages[0].role in ["system", "developer"] -%} + {%- set developer_message = conversation_messages[0] -%} + {%- set conversation_messages = conversation_messages[1:] -%} + {%- endif -%} +{%- elif messages and messages[0].role in ["system", "developer"] -%} + {%- set developer_message = messages[0] -%} + {%- set conversation_messages = messages[1:] -%} +{%- endif -%} +{#- Render system sp (higher priority, root role only) -#} +{{- bod_token ~ bos_token ~ 'system' ~ '\n' }} +{{- build_system_message(system_message) }} +{{- eos_token ~ '\n' }} + +{#- Render developer sp (lower priority: system/developer role + tools) -#} +{{- bos_token ~ 'developer' ~ '\n' }} +{{- build_developer_message(developer_message) }} +{%- if tools -%} + {{- '\n\n' ~ '# Tools' ~ '\n' ~ 'You may call one or more tools to assist with the user query.\nHere are the tools available in JSONSchema format:' ~ '\n' }} + {{- '\n' ~ '' ~ '\n' }} + {{- render_tool_namespace("functions", tools) }} + {{- '' ~ '\n\n' }} + {{- 'To call tools, wrap all invocations in a single ' ~ toolcall_begin_token ~ toolcall_end_token ~ ' block. Parameter values containing nested objects or arrays are recursively expanded into XML elements. Example:\n' }} + {{- '\n' ~ toolcall_begin_token ~ '\n' }} + {{- ns_token + '' }} + {{- ns_token + 'value-1' + ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + 'val-a' + ns_token + '' }} + {{- ns_token + 'val-b' + ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + '\n' }} + {{- ns_token + '' }} + {{- ns_token + 'value-1' + ns_token + '' }} + {{- ns_token + '\n' }} + {{- toolcall_end_token }} +{%- endif -%} +{{- eos_token ~ '\n' }} + +{#- Render messages -#} +{%- set last_tool_call = namespace(name=none) -%} +{%- for message in conversation_messages -%} + {%- if message.role == 'assistant' -%} + {{- bos_token ~ 'ai' ~ '\n' }} + + {%- set reasoning_content = '' %} + {%- set content = visible_text(message.content) %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if think_end_token in content %} + {%- set reasoning_content = content.split(think_end_token)[0].strip('\n').split(think_begin_token)[-1].strip('\n') %} + {%- set content = content.split(think_end_token)[-1].strip('\n') %} + {%- endif %} + {%- endif %} + + {%- if reasoning_content -%} + {#- Render thinking for every assistant turn (all-turn visible) -#} + {{- think_begin_token ~ reasoning_content ~ think_end_token }} + {%- else -%} + {#- No thinking rendered → prefix with think_end_token -#} + {{- think_end_token }} + {%- endif -%} + + {%- if content -%} + {{- content }} + {%- endif -%} + {%- if message.tool_calls -%} + {{- toolcall_begin_token ~ '\n' }} + + {%- for tool_call in message.tool_calls -%} + {%- if tool_call.function -%} + {%- set tool_call = tool_call.function -%} + {%- endif -%} +{{- ns_token + '' }} +{%- set _args = tool_call.arguments -%} +{%- for k, v in _args.items() if v is not none %} +{{- ns_token + '<' + k + '>' -}} +{{- to_xml(v, ns_token) -}} +{{- ns_token + '' }} +{%- endfor -%} +{{- ns_token + '' ~ '\n' }} + {%- endfor -%} + + {{- toolcall_end_token }} + {%- if message.tool_calls[-1].function -%} + {%- set last_tool_call.name = message.tool_calls[-1].function.name -%} + {%- else -%} + {%- set last_tool_call.name = message.tool_calls[-1].name -%} + {%- endif -%} + {%- else -%} + {%- set last_tool_call.name = none -%} + {%- endif -%} + {{- eos_token ~ '\n' }} + + {%- elif message.role == 'tool' -%} + {%- if last_tool_call.name is none -%} + {{- raise_exception("Message has tool role, but there was no previous assistant message with a tool call!") }} + {%- endif -%} + {%- if loop.first or (conversation_messages[loop.index0 - 1].role != 'tool') -%} + {{- bos_token ~ 'tool' }} + {%- endif -%} + {{- '\n' }} + {%- if message.content is string -%} + {{- message.content }} + {%- else -%} + {%- for tr in message.content -%} + {%- if tr is mapping and tr.type is defined and tr.type == 'image' -%} + {{- image_token }} + {%- elif tr is mapping and tr.type is defined and tr.type == 'video' -%} + {{- video_token }} + {%- else -%} + {{- tr.output if tr.output is defined else (tr.text if tr.type == 'text' and tr.text is defined else tr) }} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {{- '' }} + {%- if loop.last or (conversation_messages[loop.index0 + 1].role != 'tool') -%} + {{- eos_token ~ '\n' -}} + {%- endif -%} + + {%- elif message.role == 'user' -%} + {{- bos_token ~ 'user' ~ '\n' }} + {{- visible_text(message.content) }} + {{- eos_token ~ '\n' }} + {%- endif -%} +{%- endfor -%} + +{#- Generation prompt -#} +{%- if add_generation_prompt -%} +{{- bos_token ~ 'ai' ~ '\n' }} +{%- if thinking_mode is defined and thinking_mode == "disabled" -%} + {{- think_end_token }} +{%- elif thinking_mode is defined and thinking_mode == "adaptive" -%} + {#- adaptive: no prefix, let model decide -#} +{%- elif thinking_mode is defined and thinking_mode == "enabled" -%} + {#- enabled or not defined: default to think -#} + {{- think_begin_token }} +{%- else -%} + {#- adaptive: no prefix, let model decide -#} +{%- endif -%} +{%- endif -%} diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 650f5b5887..fe179461b8 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -617,6 +617,9 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_MASKED_EMBD_ORDERING, "masked_embd_ordering" }, { LLM_TENSOR_FC, "fc" }, { LLM_TENSOR_D2T, "d2t" }, + { LLM_TENSOR_DSPARK_MARKOV_W1, "markov_w1" }, + { LLM_TENSOR_DSPARK_MARKOV_W2, "markov_w2" }, + { LLM_TENSOR_DSPARK_CONF_PROJ, "conf_proj" }, }; // declare information about the model weight tensors: @@ -871,6 +874,10 @@ static const std::map LLM_TENSOR_INFOS = { // eagle3 {LLM_TENSOR_FC, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {LLM_TENSOR_D2T, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, + // dspark + {LLM_TENSOR_DSPARK_MARKOV_W1, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, + {LLM_TENSOR_DSPARK_MARKOV_W2, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DSPARK_CONF_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, }; LLM_KV::LLM_KV(llm_arch arch, const char * suffix) : arch(arch), suffix(suffix) {} diff --git a/src/llama-arch.h b/src/llama-arch.h index 0c6bde3e6f..49c2a6ac39 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -625,6 +625,9 @@ enum llm_tensor { LLM_TENSOR_MASKED_EMBD_ORDERING, LLM_TENSOR_FC, LLM_TENSOR_D2T, + LLM_TENSOR_DSPARK_MARKOV_W1, + LLM_TENSOR_DSPARK_MARKOV_W2, + LLM_TENSOR_DSPARK_CONF_PROJ, }; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 8a138e267f..38fc8e0213 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -822,6 +822,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_100B_A6B: return "100B.A6B"; case LLM_TYPE_102B_A12B: return "102B.A12B"; case LLM_TYPE_106B_A12B: return "106B.A12B"; + 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_196B_A11B: return "196B.A11B"; diff --git a/src/llama-model.h b/src/llama-model.h index 36d0480e5e..056a6efa59 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -130,6 +130,7 @@ enum llm_type { LLM_TYPE_100B_A6B, LLM_TYPE_102B_A12B, // Solar-Open LLM_TYPE_106B_A12B, // GLM-4.5-Air + LLM_TYPE_118B_A8B, // Laguna-S-2 LLM_TYPE_120B_A12B, // Nemotron 3 Super LLM_TYPE_122B_A10B, // Qwen3.5 LLM_TYPE_196B_A11B, // Step3.5-Flash @@ -606,6 +607,12 @@ struct llama_model { struct ggml_tensor * fc = nullptr; // feature fusion layer struct ggml_tensor * d2t = nullptr; // draft to target vocabulary mapping + // dspark + struct ggml_tensor * dspark_markov_w1 = nullptr; + struct ggml_tensor * dspark_markov_w2 = nullptr; + struct ggml_tensor * dspark_conf_proj = nullptr; + struct ggml_tensor * dspark_conf_proj_b = nullptr; + // unified vector to store target-model extracted layer ids in eagle3, dflash, etc. std::vector target_layer_ids; diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp index 6520e4181e..a9cb6bee5f 100644 --- a/src/llama-sampler.cpp +++ b/src/llama-sampler.cpp @@ -993,7 +993,9 @@ static void llama_sampler_greedy_backend_apply( GGML_UNUSED(gf); GGML_UNUSED(smpl); - struct ggml_tensor * curl = ggml_argmax(ctx, data->logits); + struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits)); + + struct ggml_tensor * curl = ggml_argmax(ctx, logits); ggml_set_name(curl, "greedy_argmax"); data->sampled = curl; @@ -1158,7 +1160,10 @@ static void llama_sampler_dist_backend_apply( ggml_set_name (sctx->inp_uniform, "uniform"); ggml_set_input(sctx->inp_uniform); - struct ggml_tensor * probs = ggml_soft_max(ctx, data->logits); + // flatten + struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits)); + + struct ggml_tensor * probs = ggml_soft_max(ctx, logits); ggml_set_name(probs, "dist_probs"); struct ggml_tensor * cumsum = ggml_cumsum(ctx, probs); @@ -1289,22 +1294,22 @@ static void llama_sampler_top_k_backend_apply( struct llama_sampler_data * data) { auto * sctx = (llama_sampler_top_k *) smpl->ctx; - struct ggml_tensor * top_k = ggml_top_k(ctx, data->logits, sctx->k); + struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits)); + + struct ggml_tensor * top_k = ggml_top_k(ctx, logits, sctx->k); ggml_set_name(top_k, "top_k"); if (data->candidates) { struct ggml_tensor * candidates_rows = ggml_reshape_2d(ctx, data->candidates, 1, data->candidates->ne[0]); data->candidates = ggml_get_rows(ctx, candidates_rows, top_k); - data->candidates = ggml_reshape_1d(ctx, data->candidates, sctx->k); ggml_set_name(data->candidates, "top_k_candidates"); } else { data->candidates = top_k; } - struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, data->logits, 1, data->logits->ne[0]); - struct ggml_tensor * top_k_rows = ggml_get_rows(ctx, logits_rows, top_k); - data->logits = ggml_reshape_1d(ctx, top_k_rows, sctx->k); - ggml_set_name(top_k_rows, "top_k_rows"); + struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, logits, 1, logits->ne[0]); + data->logits = ggml_get_rows(ctx, logits_rows, top_k); + ggml_set_name(data->logits, "top_k_rows"); GGML_UNUSED(gf); } @@ -1435,21 +1440,25 @@ static void llama_sampler_top_p_backend_apply( struct llama_sampler_data * data) { auto * sctx = (llama_sampler_top_p *) smpl->ctx; + // flatten + struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits)); + auto ggml_sort = [ctx](struct ggml_tensor * a, struct ggml_tensor * b) { GGML_ASSERT(ggml_nrows(a) == 1); struct ggml_tensor * a_reshaped = ggml_reshape_2d(ctx, a, 1, a->ne[0]); struct ggml_tensor * a_sorted = ggml_get_rows(ctx, a_reshaped, b); - return ggml_reshape_1d(ctx, a_sorted, a->ne[0]); + return a_sorted; }; // Get the sorted logits in descending order. - struct ggml_tensor * sorted_idx = ggml_argsort(ctx, data->logits, GGML_SORT_ORDER_DESC); + struct ggml_tensor * sorted_idx = ggml_argsort(ctx, logits, GGML_SORT_ORDER_DESC); ggml_set_name(sorted_idx, "top_p_sorted_idx"); // Do the sorting via reshape + get_rows - struct ggml_tensor * sorted_logits = ggml_sort(data->logits, sorted_idx); + struct ggml_tensor * sorted_logits = ggml_sort(logits, sorted_idx); ggml_set_name(sorted_logits, "top_p_sorted_logits"); + sorted_logits = ggml_reshape_1d(ctx, sorted_logits, ggml_nelements(sorted_logits)); struct ggml_tensor * softmax = ggml_soft_max(ctx, sorted_logits); ggml_set_name(softmax, "top_p_softmax"); @@ -1626,10 +1635,12 @@ static void llama_sampler_min_p_backend_apply( struct llama_sampler_data * data) { auto * sctx = (llama_sampler_min_p *) smpl->ctx; - struct ggml_tensor * max_idx = ggml_argmax(ctx, data->logits); + struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits)); + + struct ggml_tensor * max_idx = ggml_argmax(ctx, logits); ggml_set_name(max_idx, "max_idx"); - struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, data->logits, 1, data->logits->ne[0]); + struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, logits, 1, logits->ne[0]); ggml_set_name(logits_rows, "logits_rows"); struct ggml_tensor * max_logit = ggml_get_rows(ctx, logits_rows, max_idx); @@ -1640,7 +1651,7 @@ static void llama_sampler_min_p_backend_apply( ggml_set_name(threshold, "min_p_threshold"); // Subtract the threshold from logits. - struct ggml_tensor * sub = ggml_sub(ctx, data->logits, threshold); + struct ggml_tensor * sub = ggml_sub(ctx, logits, threshold); // Create a mask where logits below the threshold are 0 (discard), // and others are 1 (keep). @@ -1652,7 +1663,7 @@ static void llama_sampler_min_p_backend_apply( struct ggml_tensor * min_p_bias = ggml_log(ctx, mask); ggml_set_name(min_p_bias, "min_p_bias"); - data->logits = ggml_add(ctx, data->logits, min_p_bias); + data->logits = ggml_add(ctx, logits, min_p_bias); ggml_set_name(data->logits, "min_p_logits"); GGML_UNUSED(gf); @@ -1829,18 +1840,20 @@ static void llama_sampler_backend_temp_sampling( struct llama_sampler_data * data, float temp) { if (temp <= 0.0f) { + struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits)); + // Find the most probable token index. - struct ggml_tensor * max_idx = ggml_argmax(ctx, data->logits); + struct ggml_tensor * max_idx = ggml_argmax(ctx, logits); ggml_set_name(max_idx, "temp_max_idx"); if (data->candidates) { - struct ggml_tensor * candidates_rows = ggml_reshape_2d(ctx, data->candidates, 1, data->candidates->ne[0]); + struct ggml_tensor * candidates_rows = ggml_reshape_2d(ctx, data->candidates, 1, ggml_nelements(data->candidates)); data->candidates = ggml_get_rows(ctx, candidates_rows, max_idx); } else { data->candidates = max_idx; } - struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, data->logits, 1, data->logits->ne[0]); + struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, logits, 1, ggml_nelements(logits)); data->logits = ggml_get_rows(ctx, logits_rows, max_idx); return; @@ -2019,13 +2032,15 @@ static void llama_sampler_temp_ext_backend_apply( return; } + struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits)); + // Calculate min_temp, max_temp, and max_entropy. const float min_temp = std::max(0.0f, sctx->temp - sctx->delta); const float max_temp = sctx->temp + sctx->delta; - const float max_entropy = logf(data->logits->ne[0]); + const float max_entropy = logf(logits->ne[0]); // Calculate the probabilities. - struct ggml_tensor * probs = ggml_soft_max(ctx, data->logits); + struct ggml_tensor * probs = ggml_soft_max(ctx, logits); ggml_set_name(probs, "temp_ext_softmax_probs"); // Clamp probabilities to avoid log(0) which would give -inf @@ -2063,7 +2078,7 @@ static void llama_sampler_temp_ext_backend_apply( ggml_set_name(dyn_temp, "temp_ext_dyn_temp"); // Scale the logits by the dynamic temperature - struct ggml_tensor * scaled_logits = ggml_div(ctx, data->logits, dyn_temp); + struct ggml_tensor * scaled_logits = ggml_div(ctx, logits, dyn_temp); ggml_set_name(scaled_logits, "temp_ext_scaled_logits"); data->logits = scaled_logits; diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 427eed4594..dcff3aec9f 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -37,6 +37,23 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { const int64_t n_embd_inp = hparams.n_embd_inp_enc(); + // DSpark = DFlash + a semi-autoregressive Markov head and Confidence head + // + // TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4) + // need their own conversion path and graph tweaks + const struct ggml_tensor * markov_meta = ml->get_tensor_meta("markov_w1.weight"); + if (markov_meta) { + 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_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); + + LLAMA_LOG_INFO("%s: DFlash with DSpark markov head (rank = %lld)\n", __func__, (long long) dspark_markov_rank); + } + fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0); output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc) output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm @@ -105,6 +122,94 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_grap ggml_build_forward_expand(gf, cur); } +// DSpark (DFlash + Markov & Confidence head): Markov bias on the draft logits, chained per block position +static void build_dspark_markov_head(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens) { + ggml_context * ctx0 = g.ctx0; + auto & res = g.res; + + ggml_tensor * w1 = model.dspark_markov_w1; + ggml_tensor * w2 = model.dspark_markov_w2; + GGML_ASSERT(w1 && w2 && model.dspark_conf_proj && "DSpark markov/confidence weights not loaded"); + + ggml_tensor * base = res->t_logits; // [n_vocab, n_tokens] + const int64_t n_vocab = base->ne[0]; + const int64_t n_tok = base->ne[1]; + + const auto it = model.gguf_kv.find("dflash.block_size"); + GGML_ASSERT(it != model.gguf_kv.end() && "DSpark draft requires 'dflash.block_size' in GGUF metadata"); + const int64_t block_size = std::stoi(it->second); + GGML_ASSERT(block_size > 0); + + 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 + const int64_t block_drafts = n_tok / n_blocks; + if (block_drafts > block_size) { + return; + } + + // anchor (committed last) token of every block: token 0 of each block, i.e. a strided view + const size_t token_stride = (size_t) block_drafts * tokens->nb[0]; + const size_t base_stride = (size_t) block_drafts * base->nb[1]; + + ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, token_stride, 0); + prev = ggml_cont_1d(ctx0, prev, n_blocks); + + // confidence head input: predicts per-position acceptance + ggml_tensor * conf_inp = res->t_embd; // [n_embd, n_tok] + + ggml_tensor * cat = nullptr; + ggml_tensor * cat_conf = nullptr; + + // 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) { + 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] + + // 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]); + ggml_tensor * col = ggml_add(ctx0, base_i, bias); + + cat = cat ? ggml_concat(ctx0, cat, col, 1) : col; + + // conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b) -- [1, n_blocks] + ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks, + (size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]); + ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), w1_prev, 0); + ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat); + if (model.dspark_conf_proj_b) { + conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b); + } + conf = ggml_sigmoid(ctx0, conf); + + cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf; + + if (i + 1 < block_drafts) { + prev = ggml_argmax(ctx0, col); + } + } + + // cat is position-major; restore ubatch block-major order + ggml_tensor * out = ggml_reshape_3d(ctx0, cat, n_vocab, n_blocks, block_drafts); + out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, block_drafts, n_blocks] + out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok); + + { + ggml_tensor * conf = ggml_reshape_3d(ctx0, cat_conf, 1, n_blocks, block_drafts); + conf = ggml_cont(ctx0, ggml_permute(ctx0, conf, 0, 2, 1, 3)); + conf = ggml_reshape_2d(ctx0, conf, 1, n_tok); + + // note: broadcast the [1, n_tok] confidences to n_embd-wide rows to be able to reuse `llama_get_embeddings_nextn` + conf = ggml_repeat(ctx0, conf, res->t_embd); + res->t_h_nextn = conf; + ggml_build_forward_expand(g.gf, conf); + } + + res->t_logits = out; + ggml_build_forward_expand(g.gf, out); +} + // DFlash decoder, dual-mode by batch type: // * embd batch -> fused target features: project + inject K/V into the cache. // * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens @@ -210,6 +315,8 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); ggml_set_input(inp->tokens); + ggml_tensor * inp_tokens = inp->tokens; + ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens); cb(inpL, "inp_noise_embd", -1); @@ -290,4 +397,9 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra res->t_logits = cur; ggml_build_forward_expand(gf, cur); + + // DSpark: bias the draft logits with the Markov head + if (model.dspark_markov_w1) { + build_dspark_markov_head(*this, model, inp_tokens); + } } diff --git a/src/models/laguna.cpp b/src/models/laguna.cpp index fb55ec12f9..82c9a9538c 100644 --- a/src/models/laguna.cpp +++ b/src/models/laguna.cpp @@ -58,6 +58,7 @@ void llama_model_laguna::load_arch_hparams(llama_model_loader & ml) { switch (hparams.n_layer()) { case 40: type = LLM_TYPE_30B_A3B; break; // Laguna-XS.2 + case 48: type = LLM_TYPE_118B_A8B; break; // Laguna-S.2 case 70: type = LLM_TYPE_230B_A10B; break; // Laguna-M.1 default: type = LLM_TYPE_UNKNOWN; } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index f6b60e52a4..a5b660f47a 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4000,7 +4000,7 @@ struct test_ssm_scan : public test_case { test_ssm_scan(ggml_type type = GGML_TYPE_F32, int64_t d_state = 32, - int64_t head_dim = 1, // non-zero for Mamba-2 + int64_t head_dim = 1, // 1 = Mamba-1; > 1 = Mamba-2 (scalar A per head) int64_t n_head = 32, int64_t n_group = 1, int64_t n_seq_tokens = 32, @@ -4008,6 +4008,11 @@ struct test_ssm_scan : public test_case { bool xbc_overlap = false) : type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), xbc_overlap(xbc_overlap) {} + double max_nmse_err() override { + // SSD path (head_dim > 1) uses FP16 intermediates (M matrix, X_dt); Mamba-1 is pure FP32. + return (head_dim > 1) ? 2e-7 : 1e-7; + } + ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * s = ggml_new_tensor_4d(ctx, type, d_state, head_dim, n_head, n_seqs); ggml_tensor * dt = ggml_new_tensor_3d(ctx, type, n_head, n_seq_tokens, n_seqs); @@ -4034,14 +4039,14 @@ struct test_ssm_scan : public test_case { return out; } - // similar to test_mul_mat_id + void initialize_tensors(ggml_context * ctx) override { std::random_device rd; std::default_random_engine rng(rd()); for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { if (t->type == GGML_TYPE_I32) { if (ggml_is_view_op(t->op)) { continue; } - // ids + // ids: permutation of [0..n_seqs) for (int64_t r = 0; r < ggml_nrows(t); r++) { std::vector data(t->ne[0]); for (int i = 0; i < t->ne[0]; i++) { @@ -4050,6 +4055,11 @@ struct test_ssm_scan : public test_case { std::shuffle(data.begin(), data.end(), rng); ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t)); } + } else if (ggml_is_view_op(t->op)) { + continue; + } else if (t->ne[1] == n_head && t->ne[2] == 1) { + // A {1 or d_state, n_head}: negative decay (2-D tensor, ne[2]==1 distinguishes from 3-D/4-D tensors) + init_tensor_uniform(t, -1.0f, -0.5f); } else { init_tensor_uniform(t); } @@ -8770,6 +8780,9 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 32, 4)); // Mamba-2 test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 256, 64, 8, 2, 32, 4)); // Falcon-H1 test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 128, 4, 4, 16, 2, true)); // x/B/C overlap + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 256, 1)); // Nemotron-9B SSD path + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 512, 1)); // Nemotron-9B SSD multi-chunk (2 aligned chunks) + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 80, 8, 300, 2)); // Mamba-2 SSD multi-chunk (partial 2nd chunk, 2 seqs) test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 1, 1)); test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 32, 1)); @@ -9990,6 +10003,8 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_ssm_conv_bias_silu(GGML_TYPE_F32, {4, 3328, 1, 1}, {4, 3328, 1, 1}, true)); // generate test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 48, 1, 512, 1)); // prefill test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 48, 1, 1, 1)); // generate + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 512, 1)); // Nemotron-9B prefill + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 1, 1)); // Nemotron-9B generate // acc test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 1, 1}, {256, 16, 1, 1}, -1)); diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index 4dd00efddf..01b07953a6 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -730,6 +730,71 @@ static common_chat_tool imaginary_number_tool{ })", }; +static common_chat_tool nested_args_tool{ + /* .name = */ "nested_args", + /* .description = */ "Tool with nested array arguments", + /* .parameters = */ R"({ + "type": "object", + "properties": { + "tags": { + "type": "array", + "items": { "type": "string" } + }, + "entries": { + "type": "array", + "items": { + "type": "object", + "properties": { + "id": { "type": "integer" }, + "label": { "type": "string" } + }, + "required": ["id", "label"] + } + } + }, + "required": ["tags", "entries"] + })", +}; + +static common_chat_tool union_args_tool{ + /* .name = */ "union_args", + /* .description = */ "Tool with union arguments", + /* .parameters = */ R"({ + "type": "object", + "properties": { + "filter": { + "anyOf": [ + { "type": "array", "items": { "type": "string" } }, + { + "type": "object", + "properties": { + "field": { "type": "string" }, + "op": { "type": "string" } + }, + "required": ["field", "op"] + } + ] + }, + "label": { + "oneOf": [ + { "type": "string" }, + { "type": "object", "properties": { "text": { "type": "string" } } } + ] + }, + "limit": { + "oneOf": [ + { "type": "integer" }, + { + "type": "object", + "properties": { "max": { "type": "integer" } }, + "required": ["max"] + } + ] + } + } + })", +}; + static common_chat_tool nullable_string_tool{ /* .name = */ "set_nullable_str", /* .description = */ "Set a nullable string value", @@ -4850,6 +4915,370 @@ static void test_template_output_peg_parsers(bool detailed_debug) { .run(); } + // MiniMax-M3 tests - namespaced XML invoke format, the parameter name is the tag + // Format: + // ]<]minimax[>[ + // ]<]minimax[>[]<]minimax[>[Tokyo]<]minimax[>[]<]minimax[>[ + // ]<]minimax[>[ + // Reasoning uses .... The generation prompt is only "]~b]ai\n", so the model + // opens the thinking block itself; a turn without reasoning is prefixed with a bare . + { + auto tst = peg_tester("models/templates/MiniMax-M3.jinja", detailed_debug); + + // Content only (bare prefix) + tst.test("Hello, world!\nWhat's up?") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .expect(message_assist) + .expect_reconstruction() + .run(); + + // Thinking + content + tst.test("I'm\nthinkingHello, world!\nWhat's up?") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .expect(message_assist_thoughts) + .expect_reconstruction() + .run(); + + // Thinking + tool call (single, string param) + tst.test( + "Let me check the time" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[Tokyo]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ get_time_tool }) + .expect(message_with_tool_calls_and_reasoning("get_time", R"({"city": "Tokyo"})", "Let me check the time")) + .expect_reconstruction() + .run(); + + // Tool call without reasoning, integer param + tst.test( + "" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[1]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ special_function_tool }) + .expect(message_assist_call) + .expect_reconstruction() + .run(); + + // Tool call with no parameters + tst.test( + "Let's call a tool:" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ empty_args_tool }) + .expect(message_with_reasoning_and_tool_call("Let's call a tool:", "empty_args", "{}")) + .expect_reconstruction() + .run(); + + // Multiple parallel tool calls in one block + tst.test( + "Calling both" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[Paris]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[Paris]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .parallel_tool_calls(true) + .tools({ get_time_tool, get_weather_tool }) + .expect(message_with_reasoning_content_and_multiple_tool_calls( + "Calling both", "", + { { "get_time", R"({"city": "Paris"})" }, { "get_weather", R"({"city": "Paris"})" } })) + .expect_reconstruction() + .run(); + + // Content before the tool call block + tst.test( + "Thinking about it" + "Let me call the function." + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[1]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ special_function_tool }) + .expect_reasoning("Thinking about it") + .expect_content("Let me call the function.") + .expect_tool_calls({ + { "special_function", R"({"arg1": 1})", {} }, + }) + .expect_reconstruction() + .run(); + + // Negative number + tst.test( + "Test negative" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[-14]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ magic_int_tool }) + .expect_reasoning("Test negative") + .expect_tool_calls({ + { "magic_int", R"({"ref": -14})", {} }, + }) + .expect_reconstruction() + .run(); + + // Decimal number + tst.test( + "Test decimal" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[3.14]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ amount_tool }) + .expect_reasoning("Test decimal") + .expect_tool_calls({ + { "amount", R"({"orig": 3.14})", {} }, + }) + .expect_reconstruction() + .run(); + + // Boolean + tst.test( + "Test boolean" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[true]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ toggle_tool }) + .expect_reasoning("Test boolean") + .expect_tool_calls({ + { "toggle", R"({"enabled": true})", {} }, + }) + .expect_reconstruction() + .run(); + + // Multiple params of mixed types (required int first, then optional string) + tst.test( + "Multi-arg call" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[42]<]minimax[>[" + "]<]minimax[>[foo bar]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ magic_int_tool }) + .expect_reasoning("Multi-arg call") + .expect_tool_calls({ + { "magic_int", R"({"ref": 42, "name": "foo bar"})", {} }, + }) + .expect_reconstruction() + .run(); + + // Nested object param, expanded into one element per key + tst.test( + "Nested object" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[1.5]<]minimax[>[" + "]<]minimax[>[-2.5]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ imaginary_number_tool }) + .expect_reasoning("Nested object") + .expect_tool_calls({ + { "imaginary_number", R"({"number": {"real": 1.5, "imaginary": -2.5}})", {} }, + }) + .expect_reconstruction() + .run(); + + // Array params, expanded into elements (of scalars and of objects) + tst.test( + "Nested arrays" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[alpha]<]minimax[>[" + "]<]minimax[>[beta]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[1]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ nested_args_tool }) + .expect_reasoning("Nested arrays") + .expect_tool_calls({ + { "nested_args", R"({"tags": ["alpha", "beta"], "entries": [{"id": 1, "label": "one"}]})", {} }, + }) + .expect_reconstruction() + .run(); + + // Union params (anyOf/oneOf), expanded as a choice of the alternatives + tst.test( + "Union array" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[alpha]<]minimax[>[" + "]<]minimax[>[beta]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ union_args_tool }) + .expect_reasoning("Union array") + .expect_tool_calls({ + { "union_args", R"({"filter": ["alpha", "beta"]})", {} }, + }) + .expect_reconstruction() + .run(); + + // oneOf between a scalar and an object + tst.test( + "Union scalar" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[5]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ union_args_tool }) + .expect_reasoning("Union scalar") + .expect_tool_calls({ + { "union_args", R"({"limit": 5})", {} }, + }) + .expect_reconstruction() + .run(); + + tst.test( + "Union nested" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[10]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ union_args_tool }) + .expect_reasoning("Union nested") + .expect_tool_calls({ + { "union_args", R"({"limit": {"max": 10}})", {} }, + }) + .expect_reconstruction() + .run(); + + // A union with a string alternative is a string + tst.test( + "Union string" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ union_args_tool }) + .expect_reasoning("Union string") + .expect_tool_calls({ + { "union_args", R"({"label": "hi"})", {} }, + }) + .expect_reconstruction() + .run(); + + // ... even when the value looks structured + tst.test( + "Union string" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ union_args_tool }) + .expect_reasoning("Union string") + .expect_tool_calls({ + { "union_args", R"({"label": "]<]minimax[>[hi]<]minimax[>["})", {} }, + }) + .expect_reconstruction() + .run(); + + // Edge case: empty reasoning followed by a tool call + tst.test( + "" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[XYZCITY]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ get_time_tool }) + .expect(message_with_tool_calls("get_time", R"({"city": "XYZCITY"})")) + .run(); + + // Continuation tests + tst.test("world!\nWhat's up?") + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .enable_thinking(true) + .messages({ message_user, message_assist_prefill_content }) + .add_generation_prompt(false) + .continue_final_message(COMMON_CHAT_CONTINUATION_CONTENT) + .expect_reasoning("I'm thinking") + .expect_content("Hello, world!\nWhat's up?") + .run(); + + tst.test(" thinkingHello, world!\nWhat's up?") + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .enable_thinking(true) + .messages({ message_user, message_assist_prefill_reasoning }) + .add_generation_prompt(false) + .continue_final_message(COMMON_CHAT_CONTINUATION_REASONING) + .expect_reasoning("I'm thinking") + .expect_content("Hello, world!\nWhat's up?") + .run(); + } + // NVIDIA-Nemotron-Nano-v2 tests - ... format // Format: [{"name": "func", "arguments": {...}}] { diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index d02e65c9ea..a1ed2a76f8 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -428,9 +428,9 @@ static bool arch_supported(const llm_arch arch) { return false; } - // FIXME some models are segfaulting with WebGPU: + // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI. #ifdef GGML_USE_WEBGPU - if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_KIMI_LINEAR) { + if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA) { return false; } #endif // GGML_USE_WEBGPU @@ -600,9 +600,6 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg std::string status_roundtrip = "\033[1;33mSKIP\033[0m"; char nmse_str[12] = {0}; bool skip = !arch_supported(arch) || (dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && dc.devs.empty()); -#if defined(GGML_USE_WEBGPU) - skip = true; // FIXME -#endif // GGML_USE_WEBGPU if (!skip) { if (logits_cpu.empty()) { model_and_ctx_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, encode); diff --git a/tools/cli/README.md b/tools/cli/README.md index 972ea04dc7..bcddd05702 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -203,7 +203,7 @@ | `--spec-draft-device, -devd, --device-draft ` | comma-separated list of devices to use for offloading the draft model (none = don't offload)
use --list-devices to see a list of available devices | | `--spec-draft-ngl, -ngld, --gpu-layers-draft, --n-gpu-layers-draft N` | max. number of draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)
(env: LLAMA_ARG_N_GPU_LAYERS_DRAFT) | | `--spec-draft-model, -md, --model-draft FNAME` | draft model for speculative decoding (default: unused)
(env: LLAMA_ARG_SPEC_DRAFT_MODEL) | -| `--spec-type none,draft-simple,draft-eagle3,draft-mtp,draft-dflash,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)

(env: LLAMA_ARG_SPEC_TYPE) | +| `--spec-type none,draft-simple,draft-eagle3,draft-mtp,draft-dflash,draft-dspark,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)

(env: LLAMA_ARG_SPEC_TYPE) | | `--spec-ngram-mod-n-min N` | minimum number of ngram tokens to use for ngram-based speculative decoding (default: 48) | | `--spec-ngram-mod-n-max N` | maximum number of ngram tokens to use for ngram-based speculative decoding (default: 64) | | `--spec-ngram-mod-n-match N` | ngram-mod lookup length (default: 24) | diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index d7b7a812fc..67ab22c3fc 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -61,6 +61,7 @@ add_library(mtmd models/mobilenetv5.cpp models/youtuvl.cpp models/yasa2.cpp + models/parakeet.cpp ) set_target_properties(mtmd PROPERTIES diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 3981ff5bb6..f73e80a8aa 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -88,6 +88,7 @@ #define KEY_A_ATTN_WINDOW_SIZE "clip.audio.window_size" // mimo-audio-tokenizer: sliding-window radius #define KEY_A_LOCAL_BLOCK_COUNT "clip.audio.local_block_count" // mimo-v2.5: input_local_transformer layer count #define KEY_A_LOCAL_GROUP_SIZE "clip.audio.local_group_size" // mimo-v2.5: input_local_transformer grouping size +#define KEY_AUDIO_SUBSAMPLING_FACTOR "clip.audio.subsampling_factor" // // tensor name constants @@ -345,6 +346,12 @@ #define TN_YASA_STAGE_DOWN_CONV "v.stage.%d.down.conv.%s" #define TN_YASA_STAGE_BLK "v.stage.%d.blk.%d.%s.%s" +// parakeet +#define TN_MEL_FILTERS "a.mel_filters" +#define TN_WINDOW "a.window" +#define TN_CONV_NORM_MEAN "%s.blk.%d.conv_norm_mean" +#define TN_CONV_NORM_VAR "%s.blk.%d.conv_norm_var" + // align x to upper multiple of n #define CLIP_ALIGN(x, n) ((((x) + (n) - 1) / (n)) * (n)) @@ -399,6 +406,7 @@ enum projector_type { PROJECTOR_TYPE_KIMIK25, PROJECTOR_TYPE_NEMOTRON_V2_VL, PROJECTOR_TYPE_HUNYUANVL, + PROJECTOR_TYPE_PARAKEET, PROJECTOR_TYPE_EXAONE4_5, PROJECTOR_TYPE_MINICPMV4_6, PROJECTOR_TYPE_GRANITE_SPEECH, @@ -464,6 +472,7 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_GRANITE4_VISION, "granite4_vision"}, { PROJECTOR_TYPE_MIMO_AUDIO, "mimo_audio"}, { PROJECTOR_TYPE_QWEN3TTS_SPKENC, "qwen3tts_spkenc"}, + { PROJECTOR_TYPE_PARAKEET, "parakeet"}, }; static projector_type clip_projector_type_from_string(const std::string & str) { diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 472487351c..07d36b6b9c 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -110,6 +110,8 @@ struct clip_hparams { // audio int32_t n_mel_bins = 0; // whisper preprocessor int32_t proj_stack_factor = 0; // ultravox + int32_t subsampling_factor = 0; // parakeet + int32_t audio_chunk_size = 0; int32_t audio_conv_kernel_size = 0; int32_t audio_max_pos_emb = 0; @@ -124,6 +126,10 @@ struct clip_hparams { int32_t audio_window_len = -1; int32_t audio_hop_len = -1; + // parakeet + std::vector mel_filters; + std::vector window; + // mimo-audio-tokenizer: residual vector quantizer int32_t rvq_num_quantizers = 0; std::vector rvq_codebook_size; // per-quantizer bin count (ragged, e.g. 1024/1024/256/128x17) @@ -245,14 +251,16 @@ struct clip_layer { ggml_tensor * norm_conv_b = nullptr; ggml_tensor * linear_pos_w = nullptr; - ggml_tensor * conv_norm_w = nullptr; - ggml_tensor * conv_norm_b = nullptr; - ggml_tensor * conv_dw_w = nullptr; - ggml_tensor * conv_dw_b = nullptr; - ggml_tensor * conv_pw1_w = nullptr; - ggml_tensor * conv_pw1_b = nullptr; - ggml_tensor * conv_pw2_w = nullptr; - ggml_tensor * conv_pw2_b = nullptr; + ggml_tensor * conv_norm_w = nullptr; + ggml_tensor * conv_norm_b = nullptr; + ggml_tensor * conv_norm_mean = nullptr; // parakeet + ggml_tensor * conv_norm_var = nullptr; // parakeet + ggml_tensor * conv_dw_w = nullptr; + ggml_tensor * conv_dw_b = nullptr; + ggml_tensor * conv_pw1_w = nullptr; + ggml_tensor * conv_pw1_b = nullptr; + ggml_tensor * conv_pw2_w = nullptr; + ggml_tensor * conv_pw2_b = nullptr; // gemma4 audio conformer per-layer ggml_tensor * attn_pre_norm_w = nullptr; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 55f69b67c4..73e9947824 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1037,6 +1037,10 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_PARAKEET: + { + builder = std::make_unique(ctx, img); + } break; case PROJECTOR_TYPE_GRANITE4_VISION: { builder = std::make_unique(ctx, img); @@ -1360,6 +1364,20 @@ struct clip_model_loader { { get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge, false); } break; + case PROJECTOR_TYPE_PARAKEET: + { + get_u32(KEY_AUDIO_SUBSAMPLING_FACTOR, hparams.subsampling_factor); + GGML_ASSERT(hparams.subsampling_factor == 8 && + "subsampling_factor must match the conv strides in clip_graph_parakeet::build()"); + get_u32(KEY_A_CONV_KERNEL_SIZE, hparams.audio_conv_kernel_size); + GGML_ASSERT(hparams.audio_conv_kernel_size > 0 && hparams.audio_conv_kernel_size % 2 == 1 && + "audio_conv_kernel_size must be a positive odd integer"); + hparams.audio_chunk_len = 0; + hparams.audio_sample_rate = 16000; + hparams.audio_n_fft = 512; + hparams.audio_window_len = 400; + hparams.audio_hop_len = 160; + } break; case PROJECTOR_TYPE_IDEFICS3: { // use default llava-uhd preprocessing params @@ -1906,16 +1924,46 @@ struct clip_model_loader { return cur; }; - auto get_scalar = [&](const std::string & name, float default_val) { + auto get_vector = [&](const std::string & name) { + std::vector result; auto it = tensor_offset.find(name); if (it == tensor_offset.end()) { + return result; + } + + const int64_t idx = gguf_find_tensor(ctx_gguf.get(), name.c_str()); + if (idx < 0) { + throw std::runtime_error(string_format("%s: failed to find tensor %s\n", __func__, name.c_str())); + } + + if (const auto type = gguf_get_tensor_type(ctx_gguf.get(), idx); type != GGML_TYPE_F32) { + throw std::runtime_error(string_format("%s: %s must be %s, was %s\n", __func__, + name.c_str(), ggml_type_name(GGML_TYPE_F32), ggml_type_name(type))); + } + + const size_t n_bytes = gguf_get_tensor_size(ctx_gguf.get(), idx); + if (n_bytes == 0) { + throw std::runtime_error(string_format("%s: tensor %s is empty\n", __func__, name.c_str())); + } + + const size_t n_elems = n_bytes / sizeof(float); + result.resize(n_elems); + fin.seekg(it->second, std::ios::beg); + fin.read(reinterpret_cast(result.data()), n_bytes); + return result; + }; + + auto get_scalar = [&](const std::string & name, float default_val) { + auto v = get_vector(name); + if (v.empty()) { return default_val; } - size_t offset = it->second; - fin.seekg(offset, std::ios::beg); - float value; - fin.read(reinterpret_cast(&value), sizeof(float)); - return value; + if (v.size() != 1) { + throw std::runtime_error(string_format("%s: expected scalar tensor '%s' but got %d elements\n", + __func__, name.c_str(), (int) v.size())); + } + + return v[0]; }; model.class_embedding = get_tensor(TN_CLASS_EMBD, false); @@ -2855,6 +2903,68 @@ struct clip_model_loader { layer.conv_pw2_b = get_tensor(string_format(TN_CONV_PW2, prefix, il, "bias")); } } break; + case PROJECTOR_TYPE_PARAKEET: + { + + hparams.mel_filters = get_vector(TN_MEL_FILTERS); + hparams.window = get_vector(TN_WINDOW); + + // Subsampling layers (conv1d) + for (int i : {0, 2, 3, 5, 6}) { + model.pre_encode_conv_X_w[i] = get_tensor(string_format(TN_CONV1D, i, "weight")); + model.pre_encode_conv_X_b[i] = get_tensor(string_format(TN_CONV1D, i, "bias")); + } + model.pre_encode_out_w = get_tensor(string_format(TN_PRE_ENCODE_OUT, "weight")); + model.pre_encode_out_b = get_tensor(string_format(TN_PRE_ENCODE_OUT, "bias")); + + // Projection layers + model.mm_norm_pre_w = get_tensor(string_format(TN_MM_NORM_PRE, "weight"), false); + model.mm_0_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 1, "weight"), false); + model.mm_1_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 2, "weight"), false); + + // Encoder layers + for (int il = 0; il < hparams.n_layer; ++il) { + auto & layer = model.layers[il]; + + // Attention (from shared above) + + // Relative position encoding + layer.linear_pos_w = get_tensor(string_format(TN_LINEAR_POS, prefix, il, "weight")); + layer.pos_bias_u = get_tensor(string_format(TN_POS_BIAS_U, prefix, il)); + layer.pos_bias_v = get_tensor(string_format(TN_POS_BIAS_V, prefix, il)); + + // Convolution module + layer.conv_pw1_w = get_tensor(string_format(TN_CONV_PW1, prefix, il, "weight")); + layer.conv_pw1_b = get_tensor(string_format(TN_CONV_PW1, prefix, il, "bias"), false); + layer.conv_dw_w = get_tensor(string_format(TN_CONV_DW, prefix, il, "weight")); + layer.conv_dw_b = get_tensor(string_format(TN_CONV_DW, prefix, il, "bias"), false); + layer.conv_norm_w = get_tensor(string_format(TN_CONV_NORM, prefix, il, "weight")); + layer.conv_norm_b = get_tensor(string_format(TN_CONV_NORM, prefix, il, "bias")); + layer.conv_norm_mean = get_tensor(string_format(TN_CONV_NORM_MEAN, prefix, il)); + layer.conv_norm_var = get_tensor(string_format(TN_CONV_NORM_VAR, prefix, il)); + layer.conv_pw2_w = get_tensor(string_format(TN_CONV_PW2, prefix, il, "weight")); + layer.conv_pw2_b = get_tensor(string_format(TN_CONV_PW2, prefix, il, "bias"), false); + + // Feed-forward networks + layer.ff_norm_w = get_tensor(string_format(TN_FFN_NORM, prefix, il, "weight")); + layer.ff_norm_b = get_tensor(string_format(TN_FFN_NORM, prefix, il, "bias")); + + layer.ff_norm_1_w = get_tensor(string_format(TN_FFN_NORM_1, prefix, il, "weight")); + layer.ff_norm_1_b = get_tensor(string_format(TN_FFN_NORM_1, prefix, il, "bias")); + layer.ff_up_1_w = get_tensor(string_format(TN_FFN_UP_1, prefix, il, "weight")); + layer.ff_up_1_b = get_tensor(string_format(TN_FFN_UP_1, prefix, il, "bias"), false); + layer.ff_down_1_w = get_tensor(string_format(TN_FFN_DOWN_1, prefix, il, "weight")); + layer.ff_down_1_b = get_tensor(string_format(TN_FFN_DOWN_1, prefix, il, "bias"), false); + + // Layer norms + layer.norm_conv_w = get_tensor(string_format(TN_NORM_CONV, prefix, il, "weight")); + layer.norm_conv_b = get_tensor(string_format(TN_NORM_CONV, prefix, il, "bias")); + } + + model.mm_model_mlp_1_w = get_tensor(string_format(TN_MVLM_PROJ_MLP, 0, "weight")); + model.mm_model_mlp_2_w = get_tensor(string_format(TN_MVLM_PROJ_MLP, 1, "weight")); + model.mm_model_mlp_3_w = get_tensor(string_format(TN_MVLM_PROJ_MLP, 3, "weight")); + } break; case PROJECTOR_TYPE_GRANITE_SPEECH: { model.inp_proj_w = get_tensor(string_format(TN_INP_PROJ, "weight")); @@ -3700,6 +3810,10 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { } n_patches = n; } break; + case PROJECTOR_TYPE_PARAKEET: + { + n_patches = (img->nx() + (params.subsampling_factor - 1)) / params.subsampling_factor; + } break; case PROJECTOR_TYPE_GEMMA4UA: { n_patches = img->nx(); // no downsampling: one token per raw waveform frame @@ -4620,6 +4734,88 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 } set_input_f32("pos_emb", pos_emb); } break; + case PROJECTOR_TYPE_PARAKEET: + { + GGML_ASSERT(imgs.entries.size() == 1); + struct ggml_tensor * attn_mask = ggml_graph_get_tensor(gf, "attn_mask"); + const int n_q = attn_mask->ne[1]; + const int n_k = attn_mask->ne[0]; + const int n_frames = imgs.entries.front().nx(); + const int n_tokens_real = (n_frames + hparams.subsampling_factor-1) / hparams.subsampling_factor; + const float mask_value = -1e30f; + + std::vector mask_data(n_q * n_k); + if (n_k == n_q) { + // full attention: mask keys that are padding + for (int q = 0; q < n_q; ++q) { + for (int k = 0; k < n_k; ++k) { + mask_data[q * n_k + k] = (k >= n_tokens_real) ? mask_value : 0.0f; + } + } + } else { + // local attention: mask keys outside the valid window + const int att_left = n_k / 2; + for (int q = 0; q < n_q; ++q) { + for (int k = 0; k < n_k; ++k) { + const int key = q - att_left + k; + mask_data[q * n_k + k] = (key >= 0 && key < n_tokens_real) ? 0.0f : mask_value; + } + } + } + set_input_f32(attn_mask->name, mask_data); + + // local attention skew mask: zeroes out the probs that were + // computed for keys outside the valid sliding window. + if (struct ggml_tensor * local_mask = ggml_graph_get_tensor(gf, "local_mask")) { + const int lm_k = local_mask->ne[0]; + const int lm_q = local_mask->ne[1]; + const int window_size = lm_k - lm_q + 1; + std::vector lm_data(lm_q * lm_k); + for (int q = 0; q < lm_q; ++q) { + for (int k = 0; k < lm_k; ++k) { + const int rel = k - q; + lm_data[q * lm_k + k] = (rel >= 0 && rel < window_size) ? 1.0f : 0.0f; + } + } + set_input_f32(local_mask->name, lm_data); + } + + // Generate rotation frequencies for relative positional encoding. + { + const int n_state = hparams.n_embd; + const int d_half = n_state / 2; + const float log_10000 = logf(10000.0f); + std::vector freqs(d_half); + for (int k = 0; k < d_half; ++k) { + freqs[k] = expf(-(float(k * 2) * log_10000 / float(n_state))); + } + set_input_f32("pos_freqs", freqs); + } + + // Generate relative positional distance values which scaled by + // the frequency to produce the angles for sin/cos. + { + // window_size is only known after graph construction since it depends on + // n_time from the conv output, so we read it back from the graph tensor. + struct ggml_tensor * rel_pos = ggml_graph_get_tensor(gf, "rel_positions"); + const int window_size = rel_pos->ne[1]; + std::vector pos(window_size); + // local attention: window is fixed at [att_left, att_right] + // full attention: window covers the full sequence, centered + if (ggml_graph_get_tensor(gf, "local_mask")) { + const int att_left = window_size / 2; + for (int t = 0; t < window_size; ++t) { + pos[t] = float(att_left - t); + } + } else { + const int n_time = (window_size + 1) / 2; + for (int t = 0; t < window_size; ++t) { + pos[t] = float(n_time - 1 - t); + } + } + set_input_f32(rel_pos->name, pos); + } + } break; case PROJECTOR_TYPE_GRANITE_SPEECH: { const int context_size = ctx->model.hparams.audio_chunk_size; @@ -4905,6 +5101,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { return ctx->model.mm_2_w->ne[1]; case PROJECTOR_TYPE_QWEN3TTS_SPKENC: return ctx->model.mm_fc_w->ne[2]; + case PROJECTOR_TYPE_PARAKEET: + return ctx->model.mm_1_w->ne[1]; default: GGML_ABORT("Unknown projector type"); } diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 83db37ef3f..ca57a5f677 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -233,6 +233,11 @@ struct clip_graph_kimik25 : clip_graph { ggml_tensor * resize_position_embeddings_3d(uint32_t interpolation_mode); }; +struct clip_graph_parakeet : clip_graph { + clip_graph_parakeet(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; +}; + struct clip_graph_exaone4_5 : clip_graph { clip_graph_exaone4_5(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/models/parakeet.cpp b/tools/mtmd/models/parakeet.cpp new file mode 100644 index 0000000000..8be141d93b --- /dev/null +++ b/tools/mtmd/models/parakeet.cpp @@ -0,0 +1,421 @@ +#include "models.h" + +static constexpr int PARAKEET_LOCAL_ATTN_THRESHOLD = 8192; +static constexpr int PARAKEET_LOCAL_ATTN_WINDOW = 128; + +// conv subsampling + conformer encoder +ggml_cgraph * clip_graph_parakeet::build() { + + // Conv subsampling + ggml_tensor * inp = build_inp_raw(1); + inp = ggml_cont(ctx0, ggml_transpose(ctx0, inp)); + + // [freq, time, channels, batch] + ggml_tensor * cur = ggml_conv_2d(ctx0, model.pre_encode_conv_X_w[0], inp, 2, 2, 1, 1, 1, 1); + cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[0]); + cb(cur, "pre_conv_0", -1); + + cur = ggml_relu(ctx0, cur); + cb(cur, "pre_conv_0_relu", -1); + + // [freq, time, channels, batch] + cur = ggml_conv_2d_dw_direct(ctx0, model.pre_encode_conv_X_w[2], cur, 2, 2, 1, 1, 1, 1); + cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[2]); + cb(cur, "pre_conv_2", -1); + + // [freq, time, channels, batch] + cur = ggml_conv_2d(ctx0, model.pre_encode_conv_X_w[3], cur, 1, 1, 0, 0, 1, 1); + cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[3]); + cb(cur, "pre_conv_3", -1); + + cur = ggml_relu(ctx0, cur); + cb(cur, "pre_conv_3_relu", -1); + + // [freq, time, channels, batch] + cur = ggml_conv_2d_dw_direct(ctx0, model.pre_encode_conv_X_w[5], cur, 2, 2, 1, 1, 1, 1); + cb(cur, "pre_conv_5_direct", -1); + cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[5]); + cb(cur, "pre_conv_5", -1); + + // [freq, time, channels, batch] + cur = ggml_conv_2d(ctx0, model.pre_encode_conv_X_w[6], cur, 1, 1, 0, 0, 1, 1); + cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[6]); + cb(cur, "pre_conv_6", -1); + + cur = ggml_relu(ctx0, cur); + cb(cur, "pre_conv_6_relu", -1); + + // [freq, time, chan] + cur = ggml_permute(ctx0, cur, 0, 2, 1, 3); + // [freq, chan, time] + cur = ggml_cont(ctx0, cur); + + const int n_freq = cur->ne[0]; + const int n_chan = cur->ne[1]; + const int n_frames = cur->ne[2]; + + // [freq, time, chan, batch] -> [(freq * chan), time] + cur = ggml_reshape_2d(ctx0, cur, n_freq * n_chan, n_frames); + + cur = build_mm(model.pre_encode_out_w, cur); + cur = ggml_add(ctx0, cur, model.pre_encode_out_b); + + ggml_set_name(cur, "pre_enc_out"); + + // Encoder + + const auto & hparams = model.hparams; + const int n_layer = hparams.n_layer; + const int n_state = hparams.n_embd; + const float fc_factor = 0.5f; + + const int n_time = cur->ne[1]; + const bool local_attn = n_time > PARAKEET_LOCAL_ATTN_THRESHOLD; + const int att_left = local_attn ? PARAKEET_LOCAL_ATTN_WINDOW : n_time - 1; + const int att_right = local_attn ? PARAKEET_LOCAL_ATTN_WINDOW : n_time - 1; + const int window_size = local_attn ? att_left + att_right + 1 : 2 * n_time - 1; + const int d_half = n_state / 2; + const int mask_dim = local_attn ? window_size : n_time; + + // mask [key, n_time] + struct ggml_tensor * attn_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, mask_dim, n_time); + ggml_set_name(attn_mask, "attn_mask"); + ggml_set_input(attn_mask); + + struct ggml_tensor * local_mask = nullptr; + if (local_attn) { + const int chunk = att_left + att_right; + local_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, chunk + window_size - 1, chunk); + ggml_set_name(local_mask, "local_mask"); + ggml_set_input(local_mask); + } + + struct ggml_tensor * pos_freqs = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, d_half); + ggml_set_name(pos_freqs, "pos_freqs"); + ggml_set_input(pos_freqs); + + struct ggml_tensor * rel_positions = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, window_size); + ggml_set_name(rel_positions, "rel_positions"); + ggml_set_input(rel_positions); + + struct ggml_tensor * freqs = ggml_repeat_4d(ctx0, pos_freqs, d_half, window_size, 1, 1); + struct ggml_tensor * theta = ggml_mul(ctx0, freqs, rel_positions); + + struct ggml_tensor * sin = ggml_reshape_3d(ctx0, ggml_sin(ctx0, theta), 1, d_half, window_size); + struct ggml_tensor * cos = ggml_reshape_3d(ctx0, ggml_cos(ctx0, theta), 1, d_half, window_size); + struct ggml_tensor * pos_emb = ggml_reshape_2d(ctx0, ggml_cont(ctx0, ggml_concat(ctx0, sin, cos, 0)), n_state, window_size); + ggml_set_name(pos_emb, "pos_emb"); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + // FFN1 + { + struct ggml_tensor * residual = cur; + ggml_format_name(cur, "enc_%d_res", il); + + // norm + cur = ggml_norm(ctx0, cur, hparams.eps); + cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.ff_norm_w), layer.ff_norm_b); + ggml_format_name(cur, "enc_%d_ffn_norm_1", il); + + cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_SILU, il); + ggml_format_name(cur, "enc_%d_ffn_1", il); + + cur = ggml_add(ctx0, residual, ggml_scale(ctx0, cur, fc_factor)); + ggml_format_name(cur, "enc_%d_res_ffn", il); + } + + // self attention block using relative positional encoding from model.position_embedding. + { + // [feat, time_frames, 1, 1] + struct ggml_tensor * residual = cur; + + cur = ggml_norm(ctx0, cur, hparams.eps); + cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.ln_1_w), layer.ln_1_b); + ggml_format_name(cur, "enc_%d_attn_norm", il); + + const int n_head = hparams.n_head; + const int d_head = n_state / n_head; + + // [feat, time_frames, 1, 1] + struct ggml_tensor * Q_cur = build_mm(layer.q_w, cur); + struct ggml_tensor * K_cur = build_mm(layer.k_w, cur); + struct ggml_tensor * V_cur = build_mm(layer.v_w, cur); + + // [d_head, n_heads, n_time, 1] + Q_cur = ggml_reshape_3d(ctx0, Q_cur, d_head, n_head, n_time); + K_cur = ggml_reshape_3d(ctx0, K_cur, d_head, n_head, n_time); + V_cur = ggml_reshape_3d(ctx0, V_cur, d_head, n_head, n_time); + + // [n_state, window_size] + struct ggml_tensor * pos = build_mm(layer.linear_pos_w, pos_emb); + // [feat, head, window_size, 1] + pos = ggml_reshape_3d(ctx0, pos, d_head, n_head, pos_emb->ne[1]); + // [feat, window_size, head, 1] + pos = ggml_cont(ctx0, ggml_permute(ctx0, pos, 0, 2, 1, 3)); + ggml_format_name(pos, "enc_%d_attn_pos", il); + + if (local_attn) { + const int chunk = att_left + att_right; + const int n_group = (n_time + chunk - 1) / chunk; + const int n_time_padded = n_group * chunk; + const int n_kv_chunk = chunk + window_size - 1; + const int n_kv_dense = n_kv_chunk * n_group; + const bool need_padding = n_time_padded > n_time; + + Q_cur = ggml_cont(ctx0, ggml_permute(ctx0, Q_cur, 0, 2, 1, 3)); + K_cur = ggml_cont(ctx0, ggml_permute(ctx0, K_cur, 0, 2, 1, 3)); + V_cur = ggml_cont(ctx0, ggml_permute(ctx0, V_cur, 0, 2, 1, 3)); + + // content bias + struct ggml_tensor * bias_u = ggml_reshape_3d(ctx0, layer.pos_bias_u, d_head, 1, n_head); + struct ggml_tensor * Q_u = ggml_add(ctx0, Q_cur, bias_u); + + // position bias + struct ggml_tensor * bias_v = ggml_reshape_3d(ctx0, layer.pos_bias_v, d_head, 1, n_head); + struct ggml_tensor * Q_v = ggml_add(ctx0, Q_cur, bias_v); + + // right pad the time dimension + struct ggml_tensor * Q_u_padded = need_padding ? + ggml_pad_ext(ctx0, Q_u, 0, 0, 0, n_time_padded - n_time, 0, 0, 0, 0) : Q_u; + Q_u_padded = ggml_reshape_4d(ctx0, Q_u_padded, d_head, chunk, n_group, n_head); + + // pad front and back for the first and last time frames + struct ggml_tensor * K_padded = ggml_pad_ext(ctx0, K_cur, 0, 0, att_left, att_right, 0, 0, 0, 0); + if (n_kv_dense > K_padded->ne[1]) { + K_padded = ggml_pad_ext(ctx0, K_padded, 0, 0, 0, n_kv_dense - K_padded->ne[1], 0, 0, 0, 0); + } + + // sliding window view: each group spans n_kv_chunk keys but steps by chunk + struct ggml_tensor * K_chunk = ggml_view_4d(ctx0, K_padded, + d_head, n_kv_chunk, n_group, n_head, + K_padded->nb[1], + (size_t) chunk * K_padded->nb[1], + K_padded->nb[2], + 0); + K_chunk = ggml_cont(ctx0, K_chunk); + + struct ggml_tensor * content_scores = ggml_mul_mat(ctx0, K_chunk, Q_u_padded); + + // trim the dense output down to window_size scores per query + content_scores = ggml_view_4d(ctx0, content_scores, + window_size, chunk, n_group, n_head, + (size_t) (chunk + window_size) * content_scores->nb[0], + content_scores->nb[2], + content_scores->nb[3], + 0); + content_scores = ggml_cont(ctx0, content_scores); + + // ungroup: [window_size, n_time_padded, n_head] + content_scores = ggml_reshape_3d(ctx0, content_scores, window_size, n_time_padded, n_head); + if (need_padding) { + content_scores = ggml_view_3d(ctx0, content_scores, + window_size, n_time, n_head, + content_scores->nb[1], + content_scores->nb[2], + 0); + } + + // Q_v: [d_head, time, head] + Q_v = ggml_cont(ctx0, ggml_permute(ctx0, Q_v, 0, 2, 1, 3)); + struct ggml_tensor * rel_pos_scores = ggml_mul_mat(ctx0, pos, Q_v); + + struct ggml_tensor * attn_scores = ggml_add(ctx0, content_scores, rel_pos_scores); + attn_scores = ggml_soft_max_ext(ctx0, attn_scores, attn_mask, 1.0f / std::sqrt(d_head), 0.0f); + ggml_format_name(attn_scores, "enc_%d_attn_probs", il); + + // expand probs back to n_kv_chunk width for the V matmul + struct ggml_tensor * probs_padded = need_padding ? + ggml_pad_ext(ctx0, attn_scores, 0, 0, 0, n_time_padded - n_time, 0, 0, 0, 0) : attn_scores; + + probs_padded = ggml_reshape_4d(ctx0, probs_padded, window_size, chunk, n_group, n_head); + probs_padded = ggml_pad_ext(ctx0, probs_padded, 0, chunk, 0, 0, 0, 0, 0, 0); + probs_padded = ggml_view_4d(ctx0, probs_padded, + n_kv_chunk, chunk, n_group, n_head, + (size_t) n_kv_chunk * probs_padded->nb[0], + probs_padded->nb[2], + probs_padded->nb[3], + 0); + probs_padded = ggml_cont(ctx0, probs_padded); + probs_padded = ggml_mul(ctx0, probs_padded, local_mask); + + struct ggml_tensor * V_padded = ggml_pad_ext(ctx0, V_cur, 0, 0, att_left, att_right, 0, 0, 0, 0); + if (n_kv_dense > V_padded->ne[1]) { + V_padded = ggml_pad_ext(ctx0, V_padded, 0, 0, 0, n_kv_dense - V_padded->ne[1], 0, 0, 0, 0); + } + V_padded = ggml_cont(ctx0, ggml_transpose(ctx0, V_padded)); + + struct ggml_tensor * V_chunk = ggml_view_4d(ctx0, V_padded, + n_kv_chunk, d_head, n_group, n_head, + V_padded->nb[1], + (size_t) chunk * V_padded->nb[0], + V_padded->nb[2], + 0); + V_chunk = ggml_cont(ctx0, V_chunk); + + cur = ggml_mul_mat(ctx0, V_chunk, probs_padded); + cur = ggml_reshape_3d(ctx0, cur, d_head, n_time_padded, n_head); + if (need_padding) { + cur = ggml_view_3d(ctx0, cur, d_head, n_time, n_head, cur->nb[1], cur->nb[2], 0); + } + cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 0, 2, 1, 3)); + cur = ggml_reshape_2d(ctx0, cur, n_state, n_time); + cur = build_mm(layer.o_w, cur); + } else { + // full attention + struct ggml_tensor * Q_u = ggml_add(ctx0, Q_cur, layer.pos_bias_u); + ggml_format_name(Q_u, "enc_%d_attn_q_u", il); + + struct ggml_tensor * K_prep = ggml_permute(ctx0, K_cur, 0, 2, 1, 3); + struct ggml_tensor * Q_prep = ggml_permute(ctx0, Q_u, 0, 2, 1, 3); + struct ggml_tensor * content_scores = ggml_mul_mat(ctx0, K_prep, Q_prep); + ggml_format_name(content_scores, "enc_%d_attn_content_scores", il); + + struct ggml_tensor * Q_v = ggml_add(ctx0, Q_cur, layer.pos_bias_v); + ggml_format_name(Q_v, "enc_%d_attn_q_v", il); + + Q_v = ggml_permute(ctx0, Q_v, 0, 2, 1, 3); + Q_v = ggml_cont(ctx0, Q_v); + ggml_format_name(Q_v, "enc_%d_attn_q_v_perm", il); + + struct ggml_tensor * rel_pos_scores = ggml_mul_mat(ctx0, pos, Q_v); + ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos", il); + + // Relative positional shift + { + const auto pos_window = rel_pos_scores->ne[0]; + const auto n_frame = rel_pos_scores->ne[1]; + const auto n_head = rel_pos_scores->ne[2]; + + rel_pos_scores = ggml_pad(ctx0, rel_pos_scores, 1, 0, 0, 0); + rel_pos_scores = ggml_roll(ctx0, rel_pos_scores, 1, 0, 0, 0); + + rel_pos_scores = ggml_reshape_3d(ctx0, rel_pos_scores, n_frame, pos_window + 1, n_head); + rel_pos_scores = ggml_cont(ctx0, rel_pos_scores); + ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos_reshaped", il); + + int center = pos_window / 2; + size_t offset = rel_pos_scores->nb[0] * (center+1); + + rel_pos_scores = ggml_view_3d(ctx0, rel_pos_scores, + n_frame, pos_window, n_head, + (pos_window) * 4, + rel_pos_scores->nb[2], + offset); + rel_pos_scores = ggml_cont(ctx0, rel_pos_scores); + ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos_shifted", il); + + rel_pos_scores = ggml_view_3d(ctx0, rel_pos_scores, + content_scores->ne[0], + content_scores->ne[1], + rel_pos_scores->ne[2], + rel_pos_scores->nb[1], + rel_pos_scores->nb[2], + 0); + rel_pos_scores = ggml_cont(ctx0, rel_pos_scores); + ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos_shifted_view", il); + } + + struct ggml_tensor * attn_scores = ggml_add(ctx0, content_scores, rel_pos_scores); + ggml_format_name(attn_scores, "enc_%d_attn_scores", il); + attn_scores = ggml_scale(ctx0, attn_scores, 1.0f / std::sqrt(d_head)); + attn_scores = ggml_add(ctx0, attn_scores, attn_mask); + ggml_format_name(attn_scores, "enc_%d_attn_scores_scaled", il); + + struct ggml_tensor * probs = ggml_soft_max(ctx0, attn_scores); + ggml_format_name(probs, "enc_%d_attn_probs", il); + + V_cur = ggml_cont(ctx0, ggml_permute(ctx0, V_cur, 1, 2, 0, 3)); + ggml_format_name(V_cur, "enc_%d_attn_v_cur", il); + cur = ggml_mul_mat(ctx0, probs, V_cur); + ggml_format_name(cur, "enc_%d_attn_inp", il); + + cur = ggml_permute(ctx0, cur, 2, 0, 1, 3); + cur = ggml_cont_2d(ctx0, cur, n_state, n_time); + cur = build_mm(layer.o_w, cur); + } + ggml_format_name(cur, "enc_%d_attn_out", il); + + cur = ggml_add(ctx0, residual, cur); + ggml_format_name(cur, "enc_%d_attn_res", il); + } + + // Convolution + { + struct ggml_tensor * residual = cur; + ggml_format_name(cur, "enc_%d_residual_conv", il); + + cur = ggml_norm(ctx0, cur, hparams.eps); + cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.norm_conv_w), layer.norm_conv_b); + ggml_format_name(cur, "enc_%d_norm_conv", il); + + // pointwise 1d convolution: + cur = build_mm(layer.conv_pw1_w, cur); + ggml_format_name(cur, "enc_%d_conv_pw1", il); + + { + int64_t d = cur->ne[0] / 2; + struct ggml_tensor * signal = ggml_view_2d(ctx0, cur, d, cur->ne[1], cur->nb[1], 0); + struct ggml_tensor * gate = ggml_view_2d(ctx0, cur, d, cur->ne[1], cur->nb[1], d * cur->nb[0]); + + cur = ggml_mul(ctx0, signal, ggml_sigmoid(ctx0, gate)); + ggml_format_name(cur, "enc_%d_conv_glu", il); + } + + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + + // use ggml_ssm_conv for f32 precision + const int dw_pad = (hparams.audio_conv_kernel_size - 1) / 2; + cur = ggml_pad(ctx0, cur, dw_pad, 0, 0, 0); + cur = ggml_roll(ctx0, cur, dw_pad, 0, 0, 0); + cur = ggml_pad(ctx0, cur, dw_pad, 0, 0, 0); + ggml_format_name(cur, "enc_%d_conv_dw_pad", il); + + cur = ggml_ssm_conv(ctx0, cur, layer.conv_dw_w); + ggml_format_name(cur, "enc_%d_conv_1d_dw", il); + + cur = ggml_sub(ctx0, cur, layer.conv_norm_mean); + struct ggml_tensor * std = ggml_sqrt(ctx0, layer.conv_norm_var); + cur = ggml_div(ctx0, cur, std); + cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.conv_norm_w), layer.conv_norm_b); + ggml_format_name(cur, "enc_%d_conv_bn", il); + + cur = ggml_silu(ctx0, cur); + ggml_format_name(cur, "enc_%d_conv_silu", il); + + cur = build_mm(layer.conv_pw2_w, cur); + ggml_format_name(cur, "enc_%d_conv_pw2", il); + + cur = ggml_add(ctx0, residual, cur); + ggml_format_name(cur, "enc_%d_conv_res", il); + } + + // FFN2 + { + struct ggml_tensor * residual = cur; + cur = ggml_norm(ctx0, cur, hparams.eps); + cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.ff_norm_1_w), layer.ff_norm_1_b); + ggml_format_name(cur, "enc_%d_ffn_norm_2", il); + + cur = build_ffn(cur, layer.ff_up_1_w, nullptr, nullptr, nullptr, layer.ff_down_1_w, nullptr, FFN_SILU, il); + cur = ggml_add(ctx0, residual, ggml_scale(ctx0, cur, 0.5)); + ggml_format_name(cur, "enc_%d_ffn_res", il); + } + + cur = ggml_norm(ctx0, cur, hparams.eps); + cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.ln_2_w), layer.ln_2_b); + } + + cb(cur, "encoder_out", -1); + + cur = ggml_rms_norm(ctx0, cur, 1e-6); + cur = ggml_mul(ctx0, cur, model.mm_norm_pre_w); + cb(cur, "sound_projection.norm", -1); + + cur = build_ffn(cur, model.mm_0_w, model.mm_0_b, nullptr, nullptr, model.mm_1_w, model.mm_1_b, FFN_RELU_SQR, -1); + cb(cur, "projected", -1); + + ggml_build_forward_expand(gf, cur); + + return gf; +} diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index ed48329fca..d79e606ac9 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -1086,6 +1086,209 @@ bool mtmd_audio_preprocessor_gemma4a::preprocess(const float * s } // +// mtmd_audio_preprocessor_parakeet implementation +// + +void mtmd_audio_preprocessor_parakeet::worker_thread( + int ith, + const float * window_func, + int window_size, + const std::vector & samples, + int n_samples, + int frame_size, + int frame_step, + int n_threads, + int n_fft_bins, + const mtmd_audio_cache & cache, + mtmd_audio_mel & mel) { + std::vector fft_in(frame_size * 2, 0.0); + std::vector fft_out(frame_size * 2 * 2 * 2); + + int n_fb = n_fft_bins; + int i = ith; + + GGML_ASSERT(n_fb == 1 + (frame_size / 2)); + + const double eps = 5.960464477539063e-08; + + for (; i < std::min(n_samples / frame_step + 1, (int) mel.n_len); i += n_threads) { + const int offset = i * frame_step; + const int window_pad_left = (frame_size - window_size) / 2; + + // Zero-pad left. + std::fill(fft_in.begin(), fft_in.begin() + window_pad_left, 0.0f); + + // Apply windowed samples in the center. + const int n_to_process = std::min({window_size, n_samples - offset}); + for (int j = 0; j < n_to_process; j++) { + fft_in[window_pad_left + j] = window_func[j] * samples[offset + window_pad_left + j]; + } + + // Zero-pad right. + std::fill(fft_in.begin() + window_pad_left + n_to_process, fft_in.begin() + frame_size, 0.0f); + + // FFT. + fft(cache, fft_in.data(), frame_size, fft_out.data()); + + // Calculate modulus^2 of complex numbers. + for (int j = 0; j < n_fb; j++) { + fft_out[j] = (fft_out[2 * j + 0] * fft_out[2 * j + 0] + fft_out[2 * j + 1] * fft_out[2 * j + 1]); + } + + // mel spectrogram. + for (int j = 0; j < mel.n_mel; j++) { + double sum = 0.0; + int k = 0; + for (k = 0; k < n_fb - 3; k += 4) { + sum += + fft_out[k + 0] * cache.filters.data[j * n_fb + k + 0] + + fft_out[k + 1] * cache.filters.data[j * n_fb + k + 1] + + fft_out[k + 2] * cache.filters.data[j * n_fb + k + 2] + + fft_out[k + 3] * cache.filters.data[j * n_fb + k + 3]; + } + for (; k < n_fb; k++) { + sum += fft_out[k] * cache.filters.data[j * n_fb + k]; + } + mel.data[j * mel.n_len + i] = std::log(sum + eps); + } + } + + // Otherwise fft_out are all zero. + const double empty_sum = std::log(eps); + for (; i < mel.n_len; i += n_threads) { + for (int j = 0; j < mel.n_mel; j++) { + mel.data[j * mel.n_len + i] = empty_sum; + } + } +} + +void mtmd_audio_preprocessor_parakeet::initialize() { + cache.fill_sin_cos_table(hparams.audio_n_fft); + + const size_t n_fft = hparams.audio_n_fft / 2 + 1; + GGML_ASSERT(hparams.mel_filters.size() == (size_t)hparams.n_mel_bins * n_fft); + cache.filters.n_mel = hparams.n_mel_bins; + cache.filters.n_fft = n_fft; + cache.filters.data = hparams.mel_filters; + + GGML_ASSERT(hparams.window.size() == (size_t)hparams.audio_window_len); + GGML_ASSERT(hparams.window.size() <= (size_t) hparams.audio_n_fft); + cache.hann_window = hparams.window; +} + +bool mtmd_audio_preprocessor_parakeet::preprocess(const float * samples, + size_t n_samples_in, + std::vector & output) { + if (n_samples_in == 0) { + return false; + } + + filter_params params; + params.n_mel = hparams.n_mel_bins; + params.n_fft_bins = 1 + (hparams.audio_n_fft / 2); + params.hann_window_size = hparams.audio_window_len; + params.hop_length = hparams.audio_hop_len; + params.sample_rate = hparams.audio_sample_rate; + + GGML_ASSERT(!cache.sin_vals.empty()); + GGML_ASSERT(!cache.cos_vals.empty()); + GGML_ASSERT(!cache.filters.data.empty()); + + const float * window_func = cache.hann_window.data(); + const int window_size = params.hann_window_size; + const int frame_size = (params.n_fft_bins - 1) * 2; + const int frame_step = params.hop_length; + + // Apply preemphasis filter (high-pass): x[i] = x[i] - 0.97 * x[i-1] + std::vector samples_preprocessed(samples, samples + n_samples_in); + { + const float preemph = 0.97f; + for (int i = n_samples_in - 1; i > 0; i--) { + samples_preprocessed[i] = samples_preprocessed[i] - preemph * samples_preprocessed[i - 1]; + } + } + + // Parakeet uses centered constant padding + const size_t pad = (size_t)(frame_size / 2); + std::vector samples_padded(n_samples_in + 2 * pad, 0.0f); + std::copy(samples_preprocessed.begin(), samples_preprocessed.end(), samples_padded.begin() + pad); + + mtmd_audio_mel out_full; + out_full.n_mel = params.n_mel; + out_full.n_len = (samples_padded.size() - frame_size) / frame_step + 1; + out_full.n_len_org = out_full.n_len; + out_full.data.resize(out_full.n_mel * out_full.n_len); + + const int n_threads = 4; + std::vector workers(n_threads - 1); + for (int iw = 0; iw < n_threads - 1; ++iw) { + workers[iw] = std::thread( + worker_thread, iw + 1, + window_func, + window_size, + std::cref(samples_padded), + samples_padded.size(), + frame_size, + frame_step, + n_threads, + params.n_fft_bins, + std::cref(cache), + std::ref(out_full) + ); + } + + worker_thread(0, + window_func, + window_size, + samples_padded, + samples_padded.size(), + frame_size, + frame_step, + n_threads, + params.n_fft_bins, + cache, + out_full); + + for (int iw = 0; iw < n_threads - 1; ++iw) { + workers[iw].join(); + } + + // Per-feature normalization (only on valid frames) + { + const double eps = 1e-5; + int valid_frames = n_samples_in / frame_step; + + for (int j = 0; j < out_full.n_mel; j++) { + double sum = 0.0; + double sq_diff_sum = 0.0; + + // Calculate Mean ONLY on valid audio frames + for (int i = 0; i < valid_frames; i++) { + sum += (double)out_full.data[j * out_full.n_len + i]; + } + double mean = sum / valid_frames; + + // Calculate Variance ONLY on valid audio frames + for (int i = 0; i < valid_frames; i++) { + double diff = (double)out_full.data[j * out_full.n_len + i] - mean; + sq_diff_sum += diff * diff; + } + + double std_dev = std::sqrt(sq_diff_sum / (valid_frames - 1.0)); + double denominator = std_dev + eps; + + // Apply to ALL frames (including the padded ones) + for (int i = 0; i < out_full.n_len; i++) { + out_full.data[j * out_full.n_len + i] = (float)((out_full.data[j * out_full.n_len + i] - mean) / denominator); + } + } + } + + output.push_back(std::move(out_full)); + return true; +} + + // mtmd_audio_preprocessor_gemma4ua // diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h index 2011d95517..b4d6f72598 100644 --- a/tools/mtmd/mtmd-audio.h +++ b/tools/mtmd/mtmd-audio.h @@ -129,6 +129,21 @@ struct mtmd_audio_preprocessor_qwen3tts_spk : mtmd_audio_preprocessor { mtmd_audio_cache cache; }; +struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor { + mtmd_audio_preprocessor_parakeet(clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) { } + void initialize() override; + bool preprocess(const float * samples, size_t n_samples, std::vector & output) override; + + private: + mtmd_audio_cache cache; + + static void worker_thread(int ith, const float * window_func, int window_size, + const std::vector & samples, int n_samples, + int frame_size, int frame_step, int n_threads, + int n_fft_bins, + const mtmd_audio_cache & cache, mtmd_audio_mel & mel); +}; + // // streaming ISTFT - converts spectrogram frames back to audio one frame at a time // diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 991c177b60..4841cde564 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -724,6 +724,10 @@ struct mtmd_context { aud_end = ""; audio_preproc = std::make_unique(ctx_a); } break; + case PROJECTOR_TYPE_PARAKEET: + { + audio_preproc = std::make_unique(ctx_a); + } break; case PROJECTOR_TYPE_GEMMA4UA: { aud_beg = "<|audio>"; diff --git a/tools/server/README.md b/tools/server/README.md index 25aacf9f51..f45c018972 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -259,7 +259,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `--spec-draft-device, -devd, --device-draft ` | comma-separated list of devices to use for offloading the draft model (none = don't offload)
use --list-devices to see a list of available devices | | `--spec-draft-ngl, -ngld, --gpu-layers-draft, --n-gpu-layers-draft N` | max. number of draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)
(env: LLAMA_ARG_N_GPU_LAYERS_DRAFT) | | `--spec-draft-model, -md, --model-draft FNAME` | draft model for speculative decoding (default: unused)
(env: LLAMA_ARG_SPEC_DRAFT_MODEL) | -| `--spec-type none,draft-simple,draft-eagle3,draft-mtp,draft-dflash,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)

(env: LLAMA_ARG_SPEC_TYPE) | +| `--spec-type none,draft-simple,draft-eagle3,draft-mtp,draft-dflash,draft-dspark,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)

(env: LLAMA_ARG_SPEC_TYPE) | | `--spec-ngram-mod-n-min N` | minimum number of ngram tokens to use for ngram-based speculative decoding (default: 48) | | `--spec-ngram-mod-n-max N` | maximum number of ngram tokens to use for ngram-based speculative decoding (default: 64) | | `--spec-ngram-mod-n-match N` | ngram-mod lookup length (default: 24) | diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 744593c760..dba15d426f 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -164,6 +164,8 @@ struct server_slot { llama_context * ctx_tgt = nullptr; llama_context * ctx_dft = nullptr; + common_memory mem; + // multimodal mtmd_context * mctx = nullptr; mtmd::batch_ptr mbatch = nullptr; @@ -253,10 +255,7 @@ struct server_slot { void prompt_clear() { SLT_TRC(*this, "clearing prompt with %zu tokens\n", prompt.tokens.size()); - common_context_seq_rm(ctx_tgt, id, -1, -1); - if (ctx_dft) { - common_context_seq_rm(ctx_dft, id, -1, -1); - } + mem.seq_rm(id, -1, -1); prompt.clear(); } @@ -668,13 +667,8 @@ struct server_slot { void copy_state_to(server_slot & other) const { GGML_ASSERT(state == SLOT_STATE_DONE_PROMPT); - common_context_seq_rm(ctx_tgt, other.id, -1, -1); - common_context_seq_cp(ctx_tgt, id, other.id, -1, -1); - - if (ctx_dft) { - common_context_seq_rm(ctx_dft, other.id, -1, -1); - common_context_seq_cp(ctx_dft, id, other.id, -1, -1); - } + mem.seq_rm(other.id, -1, -1); + mem.seq_cp(id, other.id, -1, -1); other.n_decoded = n_decoded; other.n_remaining = n_remaining; @@ -1302,6 +1296,7 @@ private: slot.id = i; slot.ctx_tgt = ctx_tgt; slot.ctx_dft = ctx_dft; + slot.mem.init(ctx_tgt, ctx_dft); slot.spec = spec.get(); slot.n_ctx = n_ctx_slot; @@ -2881,13 +2876,8 @@ private: SLT_WRN(slot, "slot context shift, n_keep = %d, n_left = %d, n_discard = %d\n", n_keep, n_left, n_discard); - common_context_seq_rm (ctx_tgt, slot.id, n_keep , n_keep + n_discard); - common_context_seq_add(ctx_tgt, slot.id, n_keep + n_discard, slot.prompt.n_tokens(), -n_discard); - - if (ctx_dft) { - common_context_seq_rm (ctx_dft, slot.id, n_keep , n_keep + n_discard); - common_context_seq_add(ctx_dft, slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); - } + slot.mem.seq_rm (slot.id, n_keep , n_keep + n_discard); + slot.mem.seq_add(slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); // add generated tokens to cache // ref: https://github.com/ggml-org/llama.cpp/pull/16818#discussion_r2473269481 @@ -2998,7 +2988,9 @@ private: ckpt.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } - common_context_seq_rm(ctx_dft, slot.id, ckpt.pos_max + 1, -1); + if (!llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, ckpt.pos_max + 1, -1)) { + GGML_ABORT("failed to remove sequence %d\n", slot.id); + } } if (!draft.empty()) { @@ -3201,13 +3193,8 @@ private: const int64_t kv_shift = (int64_t) head_p - (int64_t) head_c; - common_context_seq_rm (ctx_tgt, slot.id, head_p, head_c); - common_context_seq_add(ctx_tgt, slot.id, head_c, head_c + n_match, kv_shift); - - if (ctx_dft) { - common_context_seq_rm (ctx_dft, slot.id, head_p, head_c); - common_context_seq_add(ctx_dft, slot.id, head_c, head_c + n_match, kv_shift); - } + slot.mem.seq_rm (slot.id, head_p, head_c); + slot.mem.seq_add(slot.id, head_c, head_c + n_match, kv_shift); for (size_t i = 0; i < n_match; i++) { slot.prompt.tokens.set_token(head_p + i, slot.prompt.tokens[head_c + i]); @@ -3379,10 +3366,7 @@ private: SLT_TRC(slot, "cached n_tokens = %d, memory_seq_rm [%d, end)\n", slot.prompt.n_tokens(), p0); - common_context_seq_rm(ctx_tgt, slot.id, p0, -1); - if (ctx_dft) { - common_context_seq_rm(ctx_dft, slot.id, p0, -1); - } + slot.mem.seq_rm(slot.id, p0, -1); // If using an alora, there may be uncached tokens that come // before the invocation sequence. When this happens, the @@ -3837,18 +3821,14 @@ private: SLT_DBG(slot, "restoring speculative checkpoint (pos_min = %d, pos_max = %d, size = %zu)\n", ckpt.pos_min, ckpt.pos_max, ckpt.size()); - { - ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - - common_context_seq_rm(slot.ctx_tgt, slot.id, ckpt.pos_max + 1, -1); - } + ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); if (slot.ctx_dft) { ckpt.load_dft(slot.ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - - common_context_seq_rm(slot.ctx_dft, slot.id, ckpt.pos_max + 1, -1); } + slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1); + slot.prompt.tokens.keep_first(ckpt.n_tokens); slot.smpl = std::move(smpl_save); @@ -3889,10 +3869,7 @@ private: slot.sampled = ids.back(); // last accepted token SLT_DBG(slot, "add accepted tokens: sampled=%d, ids.size=%zu, n_draft=%zu\n", slot.sampled, ids.size(), n_draft); - common_context_seq_rm(slot.ctx_tgt, slot.id, slot.prompt.tokens.pos_next(), -1); - if (slot.ctx_dft) { - common_context_seq_rm(slot.ctx_dft, slot.id, slot.prompt.tokens.pos_next(), -1); - } + slot.mem.seq_rm(slot.id, slot.prompt.tokens.pos_next(), -1); for (size_t i = 0; i < ids.size(); ++i) { completion_token_output result; diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp index e880f4ca72..674d3ba337 100644 --- a/tools/server/server-schema.cpp +++ b/tools/server/server-schema.cpp @@ -209,6 +209,7 @@ std::vector> make_llama_cmpl_schema(const common_params & ->set_hard_limits(0.0f, 1.0f) ->set_desc("Minimum speculative decoding probability for draft tokens (0 = greedy)")); + add((new field_str("speculative.type")) ->set_desc("Speculative decoding method (for debugging and research purposes)") ->set_handler([&](field_eval_context & ctx, const json & data) { diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockDefault.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockDefault.svelte index acf2de12ac..92652a0a86 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockDefault.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockDefault.svelte @@ -11,8 +11,7 @@ classifyToolResult, formatJsonPretty, parseToolResultWithImages, - type AgenticSection, - type ToolResultLine + type AgenticSection } from '$lib/utils'; import { getBuiltinToolUi } from '$lib/constants/built-in-tools'; import type { DatabaseMessageExtra } from '$lib/types'; @@ -29,11 +28,10 @@ let { section, open, isStreaming, attachments, onToggle }: Props = $props(); const title = $derived(getBuiltinToolUi(section.toolName)?.label ?? section.toolName ?? ''); - - const parsedLines: ToolResultLine[] = $derived( + const outputKind = $derived(classifyToolResult(section.toolResult)); + const parsedLines = $derived( section.toolResult ? parseToolResultWithImages(section.toolResult, attachments) : [] ); - const outputKind = $derived(classifyToolResult(section.toolResult)); diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockEditFile.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockEditFile.svelte index 6f30060f54..b990c3898b 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockEditFile.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockEditFile.svelte @@ -15,7 +15,6 @@ let { section, open, isStreaming, onToggle }: Props = $props(); const editFileMeta = $derived(parseEditFileMeta(section)); - const editDiffs = $derived( (editFileMeta?.edits ?? []).map((edit) => computeLineDiff(edit.oldText, edit.newText)) ); diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockSearchResults.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockSearchResults.svelte index e4b4adf151..60862dd063 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockSearchResults.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockSearchResults.svelte @@ -27,7 +27,7 @@ const isStreamingCall = $derived(section.type === AgenticSectionType.TOOL_CALL_STREAMING); const showSpinner = $derived(isPending || (isStreamingCall && isStreaming)); - const results: SearchResult[] = $derived(extractSearchResults(section.toolResult)); + const results = $derived(extractSearchResults(section.toolResult)); const query = $derived(extractSearchQuery(section.toolArgs)); // Same icon-resolution chain as ChatMessageToolCallBlockDefault so diff --git a/tools/ui/src/lib/constants/latex-protection.ts b/tools/ui/src/lib/constants/latex-protection.ts index da27c00837..c42aec41ac 100644 --- a/tools/ui/src/lib/constants/latex-protection.ts +++ b/tools/ui/src/lib/constants/latex-protection.ts @@ -28,6 +28,18 @@ export const LATEX_MATH_AND_CODE_PATTERN = /** Regex to capture the content of a $$...\\\\...$$ block (display-formula with line-break) */ export const LATEX_LINEBREAK_REGEXP = /\$\$([\s\S]*?\\\\[\s\S]*?)\$\$/; +/** + * Matches the unescaped `\[...\]` display-math delimiter and surrounding + * context so callers can insert line-breaks around the placeholder or convert + * to inline when the formula has a non-empty trailing context (e.g. a table + * cell that opens with `\[` and closes with content after `\]`). + * + * group 1: prefix before `\[` + * group 2: formula body + * group 3: trailing context after `\]` + */ +export const LATEX_DISPLAY_BLOCK_REGEXP = /([\S].*?)\\\[([\s\S]*?)\\\](.*)/g; + /** * Cheap gate for `preprocessLaTeX`. Every transformation it performs is triggered * by a `$` (inline/display math, currency escaping) or a backslash escape @@ -36,6 +48,76 @@ export const LATEX_LINEBREAK_REGEXP = /\$\$([\s\S]*?\\\\[\s\S]*?)\$\$/; */ export const LATEX_TRIGGER_REGEXP = /[$\\]/; +/** Inline LaTeX math delimiter (the dollar sign). */ +export const LATEX_INLINE_DELIMITER = '$'; + +/** Display LaTeX math delimiter (paired dollar signs). */ +export const LATEX_DISPLAY_DELIMITER = '$$'; + +/** Matches a single non-whitespace character. */ +export const LATEX_NON_WHITESPACE_REGEXP = /\S/; + +/** Matches a character that may appear adjacent to `$`, indicating a non-TeX + * context such as an identifier (`var$`, `$var`), currency ($5), or code. */ +export const LATEX_NEIGHBOR_CHAR_REGEXP = /[A-Za-z0-9_$-]/; + +/** Matches a single digit (used to detect currency-like `$5`). */ +export const LATEX_DIGIT_REGEXP = /[0-9]/; + +/** Matches the leading blockquote prefix (`> ` or `>`) on a markdown line. */ +export const LATEX_BLOCKQUOTE_PREFIX_REGEXP = /^(>\s*)/; + +/** Matches the placeholder inserted by the protect/restore pipeline for a + * protected LaTeX expression. Group 1 is the index into `latexExpressions`. */ +export const LATEX_PLACEHOLDER_REGEXP = /<>/g; + +/** Matches the placeholder inserted by the protect/restore pipeline for a + * protected code block. Group 1 is the index into `codeBlocks`. */ +export const CODE_BLOCK_PLACEHOLDER_REGEXP = /<>/g; + +/** Matches a `$` immediately followed by a digit, which is treated as a + * currency amount (e.g. `$5`) and escaped to `\$5` so it isn't parsed as math. */ +export const LATEX_CURRENCY_DOLLAR_REGEXP = /\$(?=\d)/g; + +/** Captures remaining `$$...$$`, `\[...\]`, `\(...\)` (only unescaped via + * `(?(); + for (const tm of toolMessages) { + if (tm.toolCallId && !toolMsgById.has(tm.toolCallId)) { + toolMsgById.set(tm.toolCallId, tm); + } + } + for (const tc of toolCalls) { - const resultMsg = toolMessages.find((m) => m.toolCallId === tc.id); + const resultMsg = tc.id ? toolMsgById.get(tc.id) : undefined; // Only show as pending/loading if we're actively streaming; otherwise it's just a tool call without result const type = resultMsg ? AgenticSectionType.TOOL_CALL @@ -112,9 +121,10 @@ function deriveSingleTurnSections( } // 4. Streaming tool calls (not yet persisted - currently being received) + const persistedIds = new Set(toolCalls.map((t) => t.id).filter(Boolean)); for (const tc of streamingToolCalls) { // Skip if already in persisted tool calls - if (tc.id && toolCalls.find((t) => t.id === tc.id)) continue; + if (tc.id && persistedIds.has(tc.id)) continue; sections.push({ type: AgenticSectionType.TOOL_CALL_STREAMING, content: '', @@ -281,15 +291,31 @@ export function splitSearchSummaryList( return { lines }; } +/** Bounded cache for parseToolResultWithImages results. */ +const TOOL_RESULT_LINES_CACHE_MAX_SIZE = 32; +const toolResultLinesCache = new Map(); + /** * Parse tool result text into lines, matching image attachments by name. + * Memoized: called per render during streaming on unchanged tool result + * strings with unchanged extras. */ export function parseToolResultWithImages( toolResult: string, extras?: DatabaseMessageExtra[] ): ToolResultLine[] { + // Cache key includes image attachment names so we recompute when + // attachments change, even if the count stays the same. + const imageNames = (extras ?? []) + .filter((e): e is DatabaseMessageExtraImageFile => e.type === AttachmentType.IMAGE) + .map((e) => e.name) + .join(NEWLINE); + const cacheKey = `${imageNames}:${toolResult}`; + const cached = toolResultLinesCache.get(cacheKey); + if (cached !== undefined) return cached; + const lines = toolResult.split(NEWLINE); - return lines.map((line) => { + const result = lines.map((line) => { const match = line.match(ATTACHMENT_SAVED_REGEX); if (!match || !extras) return { text: line }; @@ -301,8 +327,19 @@ export function parseToolResultWithImages( return { text: line, image }; }); + + if (toolResultLinesCache.size >= TOOL_RESULT_LINES_CACHE_MAX_SIZE) { + toolResultLinesCache.delete(toolResultLinesCache.keys().next().value!); + } + toolResultLinesCache.set(cacheKey, result); + + return result; } +/** Bounded cache for classifyToolResult results. */ +const CLASSIFY_CACHE_MAX_SIZE = 32; +const classifyCache = new Map(); + /** * Pick a renderer tier for a tool's result content. * @@ -312,25 +349,39 @@ export function parseToolResultWithImages( * through MarkdownContent for proper formatting. * text - everything else, rendered as plain text lines (with image * attachment resolution as a side effect). + * Memoized: called per render during streaming on unchanged content. */ export function classifyToolResult(content: string | undefined): ToolResultKind { if (!content) return ToolResultKind.TEXT; + + const cached = classifyCache.get(content); + if (cached !== undefined) return cached; + const trimmed = content.trim(); if (!trimmed) return ToolResultKind.TEXT; + let result: ToolResultKind = ToolResultKind.TEXT; + // Strongest signal: JSON object/array round-trips through JSON.parse. if (TOOL_RESULT_JSON_OPEN_REGEX.test(trimmed)) { try { JSON.parse(trimmed); - return ToolResultKind.JSON; + result = ToolResultKind.JSON; } catch (error) { console.error('[agentic] tool result looked like JSON but failed to parse:', error); } } - if (looksLikeMarkdown(trimmed)) return ToolResultKind.MARKDOWN; + if (result === ToolResultKind.TEXT && looksLikeMarkdown(trimmed)) { + result = ToolResultKind.MARKDOWN; + } - return ToolResultKind.TEXT; + if (classifyCache.size >= CLASSIFY_CACHE_MAX_SIZE) { + classifyCache.delete(classifyCache.keys().next().value!); + } + classifyCache.set(content, result); + + return result; } /** @@ -370,19 +421,35 @@ function looksLikeMarkdown(content: string): boolean { return false; } +/** Bounded cache for parsed tool-call JSON blobs. */ +const TOOL_CALLS_CACHE_MAX_SIZE = 64; +const toolCallsParseCache = new Map(); + /** * Safely parse the toolCalls JSON string from a DatabaseMessage. + * Memoized: the same JSON string is re-parsed on every render during + * streaming, which is wasted CPU since tool calls don't change mid-stream. */ function parseToolCalls(toolCallsJson?: string): ApiChatCompletionToolCall[] { if (!toolCallsJson) return []; + const cached = toolCallsParseCache.get(toolCallsJson); + if (cached) return cached; + + let result: ApiChatCompletionToolCall[]; try { const parsed = JSON.parse(toolCallsJson); - - return Array.isArray(parsed) ? parsed : []; + result = Array.isArray(parsed) ? parsed : []; } catch { - return []; + result = []; } + + if (toolCallsParseCache.size >= TOOL_CALLS_CACHE_MAX_SIZE) { + toolCallsParseCache.delete(toolCallsParseCache.keys().next().value!); + } + toolCallsParseCache.set(toolCallsJson, result); + + return result; } /** diff --git a/tools/ui/src/lib/utils/code.ts b/tools/ui/src/lib/utils/code.ts index 44b9b48415..35f3877f4a 100644 --- a/tools/ui/src/lib/utils/code.ts +++ b/tools/ui/src/lib/utils/code.ts @@ -34,6 +34,10 @@ function escapeCode(code: string): string { return code.replace(AMPERSAND_REGEX, '&').replace(LT_REGEX, '<').replace(GT_REGEX, '>'); } +/** Bounded cache for highlightCode results. */ +const HIGHLIGHT_CACHE_MAX_SIZE = 64; +const highlightCache = new Map(); + /** * Highlights code using highlight.js * @param code - The code to highlight @@ -47,23 +51,37 @@ function escapeCode(code: string): string { export function highlightCode(code: string, language: string, autoDetect = true): string { if (!code) return ''; + // Cache key includes language and autoDetect flag since results differ. + // During streaming, the same code string may be highlighted repeatedly + // (e.g., when text after a code block changes but the code itself doesn't). + const cacheKey = `${language}:${autoDetect}:${code}`; + const cached = highlightCache.get(cacheKey); + if (cached) return cached; + const trimmed = trimCodePadding(code); + let result: string; try { const lang = language.toLowerCase(); const isSupported = hljs.getLanguage(lang); if (isSupported) { - return hljs.highlight(trimmed, { language: lang }).value; + result = hljs.highlight(trimmed, { language: lang }).value; } else if (autoDetect) { - return hljs.highlightAuto(trimmed).value; + result = hljs.highlightAuto(trimmed).value; } else { - return escapeCode(trimmed); + result = escapeCode(trimmed); } } catch { - // Fallback to escaped plain text - return escapeCode(trimmed); + result = escapeCode(trimmed); } + + if (highlightCache.size >= HIGHLIGHT_CACHE_MAX_SIZE) { + highlightCache.delete(highlightCache.keys().next().value!); + } + highlightCache.set(cacheKey, result); + + return result; } export { trimCodePadding }; diff --git a/tools/ui/src/lib/utils/latex-protection.ts b/tools/ui/src/lib/utils/latex-protection.ts index 573eb92978..bbeed82500 100644 --- a/tools/ui/src/lib/utils/latex-protection.ts +++ b/tools/ui/src/lib/utils/latex-protection.ts @@ -1,9 +1,31 @@ import { + CODE_BLOCK_PLACEHOLDER_REGEXP, CODE_BLOCK_REGEXP, + LATEX_BACKSLASH, + LATEX_BLOCKQUOTE_PREFIX_REGEXP, + LATEX_CURRENCY_DOLLAR_REGEXP, + LATEX_CURRENCY_ESCAPE, + LATEX_DIGIT_REGEXP, + LATEX_DISPLAY_BLOCK_REGEXP, + LATEX_DISPLAY_CLOSE, + LATEX_DISPLAY_CONVERT_REGEXP, + LATEX_DISPLAY_DELIMITER, + LATEX_DISPLAY_OPEN, + LATEX_INLINE_CLOSE, + LATEX_INLINE_CONVERT_REGEXP, + LATEX_INLINE_DELIMITER, + LATEX_INLINE_OPEN, LATEX_MATH_AND_CODE_PATTERN, + LATEX_MHCHEM_CE, + LATEX_MHCHEM_PU, LATEX_LINEBREAK_REGEXP, + LATEX_NEIGHBOR_CHAR_REGEXP, + LATEX_NON_WHITESPACE_REGEXP, + LATEX_PLACEHOLDER_REGEXP, + LATEX_PROTECT_REGEXP, LATEX_TRIGGER_REGEXP, - MHCHEM_PATTERN_MAP + MHCHEM_PATTERN_MAP, + NEWLINE } from '$lib/constants'; /** @@ -20,13 +42,13 @@ import { * @returns The processed string with LaTeX replaced by placeholders. */ export function maskInlineLaTeX(content: string, latexExpressions: string[]): string { - if (!content.includes('$')) { + if (!content.includes(LATEX_INLINE_DELIMITER)) { return content; } return content - .split('\n') + .split(NEWLINE) .map((line) => { - if (line.indexOf('$') == -1) { + if (line.indexOf(LATEX_INLINE_DELIMITER) == -1) { return line; } @@ -34,7 +56,7 @@ export function maskInlineLaTeX(content: string, latexExpressions: string[]): st let currentPosition = 0; while (currentPosition < line.length) { - const openDollarIndex = line.indexOf('$', currentPosition); + const openDollarIndex = line.indexOf(LATEX_INLINE_DELIMITER, currentPosition); if (openDollarIndex == -1) { processedLine += line.slice(currentPosition); @@ -42,7 +64,7 @@ export function maskInlineLaTeX(content: string, latexExpressions: string[]): st } // Is there a next $-sign? - const closeDollarIndex = line.indexOf('$', openDollarIndex + 1); + const closeDollarIndex = line.indexOf(LATEX_INLINE_DELIMITER, openDollarIndex + 1); if (closeDollarIndex == -1) { processedLine += line.slice(currentPosition); @@ -62,14 +84,14 @@ export function maskInlineLaTeX(content: string, latexExpressions: string[]): st shouldSkipAsNonLatex = true; } - if (/[A-Za-z0-9_$-]/.test(charBeforeOpen)) { + if (LATEX_NEIGHBOR_CHAR_REGEXP.test(charBeforeOpen)) { // Character, digit, $, _ or - before first '$', no TeX. shouldSkipAsNonLatex = true; } if ( - /[0-9]/.test(charAfterOpen) && - (/[A-Za-z0-9_$-]/.test(charAfterClose) || ' ' == charBeforeClose) + LATEX_DIGIT_REGEXP.test(charAfterOpen) && + (LATEX_NEIGHBOR_CHAR_REGEXP.test(charAfterClose) || ' ' == charBeforeClose) ) { // First $ seems to belong to an amount. shouldSkipAsNonLatex = true; @@ -92,7 +114,7 @@ export function maskInlineLaTeX(content: string, latexExpressions: string[]): st return processedLine; }) - .join('\n'); + .join(NEWLINE); } function escapeBrackets(text: string): string { @@ -107,9 +129,9 @@ function escapeBrackets(text: string): string { if (codeBlock != null) { return codeBlock; } else if (squareBracket != null) { - return `$$${squareBracket}$$`; + return `${LATEX_DISPLAY_DELIMITER}${squareBracket}${LATEX_DISPLAY_DELIMITER}`; } else if (roundBracket != null) { - return `$${roundBracket}$`; + return `${LATEX_INLINE_DELIMITER}${roundBracket}${LATEX_INLINE_DELIMITER}`; } return match; @@ -145,32 +167,49 @@ const doEscapeMhchem = false; * preprocessLaTeX("Price: $10. The equation is \\(x^2\\).") * // → "Price: $10. The equation is $x^2$." */ +/** Bounded cache for preprocessLaTeX results. */ +const LATEX_CACHE_MAX_SIZE = 64; +const latexCache = new Map(); + export function preprocessLaTeX(content: string): string { // See also: // https://github.com/danny-avila/LibreChat/blob/main/client/src/utils/latex.ts + // Memoize on the input string. During streaming the prefix before an + // incomplete code block stays the same across multiple tokens, so the + // full protect/restore pipeline would re-run unnecessarily. + const cached = latexCache.get(content); + if (cached !== undefined) return cached; + + // Save original before the function mutates `content` through steps 0-8 + const originalContent = content; + // Every step below keys off a `$` or a backslash escape (\[ \] \( \) \ce{ \pu{). // With neither present the protect/restore passes round-trip the input // unchanged, so skip them: the step 2 scan is O(n^2) in line length and costs // ~90ms on a 26KB single-line message that contains no math at all. This // matters during streaming, where the whole message is reprocessed per frame. if (!LATEX_TRIGGER_REGEXP.test(content)) { + if (latexCache.size >= LATEX_CACHE_MAX_SIZE) { + latexCache.delete(latexCache.keys().next().value!); + } + latexCache.set(originalContent, content); return content; } // Step 0: Temporarily remove blockquote markers (>) to process LaTeX correctly // Store the structure so we can restore it later const blockquoteMarkers: Map = new Map(); - const lines = content.split('\n'); + const lines = content.split(NEWLINE); const processedLines = lines.map((line, index) => { - const match = line.match(/^(>\s*)/); + const match = line.match(LATEX_BLOCKQUOTE_PREFIX_REGEXP); if (match) { blockquoteMarkers.set(index, match[1]); return line.slice(match[1].length); } return line; }); - content = processedLines.join('\n'); + content = processedLines.join(NEWLINE); // Step 1: Protect code blocks const codeBlocks: string[] = []; @@ -187,58 +226,52 @@ export function preprocessLaTeX(content: string): string { // Match \S...\[...\] and protect them and insert a line-break. // Guarded: with no `\[` present this pattern still probes every start offset, // expanding `.*?` to the end of each line before failing - O(n^2) for nothing. - if (content.includes('\\[')) { - content = content.replace( - /([\S].*?)\\\[([\s\S]*?)\\\](.*)/g, - (match, group1, group2, group3) => { - // Check if there are characters following the formula (display-formula in a table-cell?) - if (group1.endsWith('\\')) { - return match; // Backslash before \[, do nothing. - } - const hasSuffix = /\S/.test(group3); - let optBreak; - - if (hasSuffix) { - latexExpressions.push(`\\(${group2.trim()}\\)`); // Convert into inline. - optBreak = ''; - } else { - latexExpressions.push(`\\[${group2}\\]`); - optBreak = '\n'; - } - - return `${group1}${optBreak}<>${optBreak}${group3}`; + if (content.includes(LATEX_DISPLAY_OPEN)) { + content = content.replace(LATEX_DISPLAY_BLOCK_REGEXP, (match, group1, group2, group3) => { + // Check if there are characters following the formula (display-formula in a table-cell?) + if (group1.endsWith(LATEX_BACKSLASH)) { + return match; // Backslash before \[, do nothing. } - ); + const hasSuffix = LATEX_NON_WHITESPACE_REGEXP.test(group3); + let optBreak; + + if (hasSuffix) { + latexExpressions.push(`${LATEX_INLINE_OPEN}${group2.trim()}${LATEX_INLINE_CLOSE}`); // Convert into inline. + optBreak = ''; + } else { + latexExpressions.push(`${LATEX_DISPLAY_OPEN}${group2}${LATEX_DISPLAY_CLOSE}`); + optBreak = NEWLINE; + } + + return `${group1}${optBreak}<>${optBreak}${group3}`; + }); } // Match \(...\), \[...\], $$...$$ and protect them - content = content.replace( - /(\$\$[\s\S]*?\$\$|(? { - latexExpressions.push(match); + content = content.replace(LATEX_PROTECT_REGEXP, (match) => { + latexExpressions.push(match); - return `<>`; - } - ); + return `<>`; + }); // Protect inline $...$ but NOT if it looks like money (e.g., $10, $3.99) content = maskInlineLaTeX(content, latexExpressions); // Step 3: Escape standalone $ before digits (currency like $5 → \$5) // (Now that inline math is protected, this will only escape dollars not already protected) - content = content.replace(/\$(?=\d)/g, '\\$'); + content = content.replace(LATEX_CURRENCY_DOLLAR_REGEXP, LATEX_CURRENCY_ESCAPE); // Step 4: Restore protected LaTeX expressions (they are valid) - content = content.replace(/<>/g, (_, index) => { + content = content.replace(LATEX_PLACEHOLDER_REGEXP, (_, index) => { let expr = latexExpressions[parseInt(index)]; const match = expr.match(LATEX_LINEBREAK_REGEXP); if (match) { // Katex: The $$-delimiters should be in their own line // if there are \\-line-breaks. const formula = match[1]; - const prefix = formula.startsWith('\n') ? '' : '\n'; - const suffix = formula.endsWith('\n') ? '' : '\n'; - expr = '$$' + prefix + formula + suffix + '$$'; + const prefix = formula.startsWith(NEWLINE) ? '' : NEWLINE; + const suffix = formula.endsWith(NEWLINE) ? '' : NEWLINE; + expr = LATEX_DISPLAY_DELIMITER + prefix + formula + suffix + LATEX_DISPLAY_DELIMITER; } return expr; }); @@ -247,7 +280,7 @@ export function preprocessLaTeX(content: string): string { // This must happen BEFORE restoring code blocks to avoid affecting code content content = escapeBrackets(content); - if (doEscapeMhchem && (content.includes('\\ce{') || content.includes('\\pu{'))) { + if (doEscapeMhchem && (content.includes(LATEX_MHCHEM_CE) || content.includes(LATEX_MHCHEM_PU))) { content = escapeMhchem(content); } @@ -257,31 +290,38 @@ export function preprocessLaTeX(content: string): string { // Using the look‑behind pattern `(? { + return `${LATEX_INLINE_DELIMITER}${formula}${LATEX_INLINE_DELIMITER}`; + }) // inline .replace( // Using the look‑behind pattern `(? { - return `$$${content}$$`; + LATEX_DISPLAY_CONVERT_REGEXP, // display, see also PR #16599 + (_, formula: string) => { + return `${LATEX_DISPLAY_DELIMITER}${formula}${LATEX_DISPLAY_DELIMITER}`; } ); // Step 7: Restore code blocks // This happens AFTER all LaTeX conversions to preserve code content - content = content.replace(/<>/g, (_, index) => { + content = content.replace(CODE_BLOCK_PLACEHOLDER_REGEXP, (_, index) => { return codeBlocks[parseInt(index)]; }); // Step 8: Restore blockquote markers if (blockquoteMarkers.size > 0) { - const finalLines = content.split('\n'); + const finalLines = content.split(NEWLINE); const restoredLines = finalLines.map((line, index) => { const marker = blockquoteMarkers.get(index); return marker ? marker + line : line; }); - content = restoredLines.join('\n'); + content = restoredLines.join(NEWLINE); } + if (latexCache.size >= LATEX_CACHE_MAX_SIZE) { + latexCache.delete(latexCache.keys().next().value!); + } + latexCache.set(originalContent, content); + return content; } diff --git a/tools/ui/src/lib/utils/parse-partial-json-args.ts b/tools/ui/src/lib/utils/parse-partial-json-args.ts index 09c1122890..58439bd6e3 100644 --- a/tools/ui/src/lib/utils/parse-partial-json-args.ts +++ b/tools/ui/src/lib/utils/parse-partial-json-args.ts @@ -14,70 +14,94 @@ const JSON_ARRAY_CLOSE = ']'; // comma when the model cut off mid-key. const TRAILING_JSON_PUNCTUATION_REGEX = /,?\s*$/; +/** Bounded cache for parsePartialJsonArgs results. */ +const PARTIAL_JSON_CACHE_MAX_SIZE = 32; +const partialJsonCache = new Map | null>(); + +function cacheResult(input: string, result: Record | null): void { + if (partialJsonCache.size >= PARTIAL_JSON_CACHE_MAX_SIZE) { + partialJsonCache.delete(partialJsonCache.keys().next().value!); + } + partialJsonCache.set(input, result); +} + // Parse partial tool-arg JSON streamed token-by-token. Closes any // unterminated string and dangling open containers (in reverse order), // so parsers can still surface keys already received while the call -// is still in flight. +// is still in flight. Memoized: the char-by-char scanner runs on every +// render during streaming even when toolArgs hasn't changed. export function parsePartialJsonArgs(toolArgsString: string): Record | null { + const cached = partialJsonCache.get(toolArgsString); + if (cached !== undefined) return cached; + + let result: Record | null; + try { const parsed: unknown = JSON.parse(toolArgsString); - if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) { - return parsed as Record; - } - return null; + result = + parsed && typeof parsed === 'object' && !Array.isArray(parsed) + ? (parsed as Record) + : null; } catch { - let inString = false; - let escape = false; - const stack: ('{' | '[')[] = []; + result = scanPartialJson(toolArgsString); + } - for (let i = 0; i < toolArgsString.length; i++) { - const ch = toolArgsString[i]; - if (escape) { - escape = false; - continue; - } - if (ch === JSON_BACKSLASH && inString) { - escape = true; - continue; - } - if (ch === JSON_QUOTE) { - inString = !inString; - continue; - } - if (inString) continue; - if (ch === JSON_OBJECT_OPEN) stack.push(JSON_OBJECT_OPEN); - else if (ch === JSON_OBJECT_CLOSE) { - if (stack.length === 0 || stack[stack.length - 1] !== JSON_OBJECT_OPEN) return null; - stack.pop(); - } else if (ch === JSON_ARRAY_OPEN) stack.push(JSON_ARRAY_OPEN); - else if (ch === JSON_ARRAY_CLOSE) { - if (stack.length === 0 || stack[stack.length - 1] !== JSON_ARRAY_OPEN) return null; - stack.pop(); - } - } + cacheResult(toolArgsString, result); + return result; +} - let completed = toolArgsString; +/** Char-by-char scanner for unterminated partial JSON. */ +function scanPartialJson(toolArgsString: string): Record | null { + let inString = false; + let escape = false; + const stack: ('{' | '[')[] = []; + + for (let i = 0; i < toolArgsString.length; i++) { + const ch = toolArgsString[i]; if (escape) { - // Dangling escape at end of partial JSON: escape the trailing - // backslash as a literal so we can close the string cleanly. - completed += JSON_BACKSLASH; + escape = false; + continue; } - if (inString) completed += JSON_QUOTE; - if (!inString) completed = completed.replace(TRAILING_JSON_PUNCTUATION_REGEX, ''); - - // Close in reverse nesting order: innermost container first. - for (let i = stack.length - 1; i >= 0; i--) { - completed += stack[i] === JSON_OBJECT_OPEN ? JSON_OBJECT_CLOSE : JSON_ARRAY_CLOSE; + if (ch === JSON_BACKSLASH && inString) { + escape = true; + continue; } - - try { - const parsed: unknown = JSON.parse(completed); - if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) { - return parsed as Record; - } - return null; - } catch { - return null; + if (ch === JSON_QUOTE) { + inString = !inString; + continue; + } + if (inString) continue; + if (ch === JSON_OBJECT_OPEN) stack.push(JSON_OBJECT_OPEN); + else if (ch === JSON_OBJECT_CLOSE) { + if (stack.length === 0 || stack[stack.length - 1] !== JSON_OBJECT_OPEN) return null; + stack.pop(); + } else if (ch === JSON_ARRAY_OPEN) stack.push(JSON_ARRAY_OPEN); + else if (ch === JSON_ARRAY_CLOSE) { + if (stack.length === 0 || stack[stack.length - 1] !== JSON_ARRAY_OPEN) return null; + stack.pop(); } } + + let completed = toolArgsString; + if (escape) { + // Dangling escape at end of partial JSON: escape the trailing + // backslash as a literal so we can close the string cleanly. + completed += JSON_BACKSLASH; + } + if (inString) completed += JSON_QUOTE; + if (!inString) completed = completed.replace(TRAILING_JSON_PUNCTUATION_REGEX, ''); + + // Close in reverse nesting order: innermost container first. + for (let i = stack.length - 1; i >= 0; i--) { + completed += stack[i] === JSON_OBJECT_OPEN ? JSON_OBJECT_CLOSE : JSON_ARRAY_CLOSE; + } + + try { + const parsed: unknown = JSON.parse(completed); + return parsed && typeof parsed === 'object' && !Array.isArray(parsed) + ? (parsed as Record) + : null; + } catch { + return null; + } } diff --git a/tools/ui/src/lib/utils/search-results.ts b/tools/ui/src/lib/utils/search-results.ts index 861b0d2d32..d090e6dc4b 100644 --- a/tools/ui/src/lib/utils/search-results.ts +++ b/tools/ui/src/lib/utils/search-results.ts @@ -155,41 +155,71 @@ function parseChunk(chunk: string): SearchResult | null { return result; } +/** Bounded cache for extractSearchResults results. */ +const SEARCH_RESULTS_CACHE_MAX_SIZE = 32; +const searchResultsCache = new Map(); + /** * Extract a SearchResult[] from a tool-result string. Returns `[]` when * the input does not match the expected shape — useful for branching * between dedicated search-results rendering and the generic tool-call - * block. + * block. Memoized: called per render during streaming on unchanged + * tool result strings. */ export function extractSearchResults(text: string | undefined | null): SearchResult[] { if (!text) return []; + const cached = searchResultsCache.get(text); + if (cached) return cached; + const results: SearchResult[] = []; for (const chunk of splitChunks(text)) { const parsed = parseChunk(chunk); if (parsed) results.push(parsed); } + + if (searchResultsCache.size >= SEARCH_RESULTS_CACHE_MAX_SIZE) { + searchResultsCache.delete(searchResultsCache.keys().next().value!); + } + searchResultsCache.set(text, results); + return results; } +/** Bounded cache for extractSearchQuery results. */ +const SEARCH_QUERY_CACHE_MAX_SIZE = 32; +const searchQueryCache = new Map(); + /** * Best-effort extraction of the search query out of a tool call's JSON * argument blob. Currently looks for a `query` field (the convention * used by Exa and most web-search MCP servers); returns an empty string - * if it cannot be located. + * if it cannot be located. Memoized: called per render during streaming + * on unchanged tool args strings. */ export function extractSearchQuery(toolArgs: string | undefined | null): string { if (!toolArgs) return ''; + + const cached = searchQueryCache.get(toolArgs); + if (cached !== undefined) return cached; + + let result = ''; try { const parsed: unknown = JSON.parse(toolArgs); if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) { const candidate = (parsed as Record)[SEARCH_TOOL_QUERY_FIELD]; - if (typeof candidate === 'string') return candidate.trim(); + if (typeof candidate === 'string') result = candidate.trim(); } } catch { - return ''; + result = ''; } - return ''; + + if (searchQueryCache.size >= SEARCH_QUERY_CACHE_MAX_SIZE) { + searchQueryCache.delete(searchQueryCache.keys().next().value!); + } + searchQueryCache.set(toolArgs, result); + + return result; } /** diff --git a/tools/ui/tests/unit/parse-toolcalls-memo.test.ts b/tools/ui/tests/unit/parse-toolcalls-memo.test.ts new file mode 100644 index 0000000000..85187febcb --- /dev/null +++ b/tools/ui/tests/unit/parse-toolcalls-memo.test.ts @@ -0,0 +1,146 @@ +// Tests for the memoized parseToolCalls and O(1) tool message lookup in +// deriveAgenticSections. These were added to prevent regressions where +// streaming text tokens trigger redundant JSON.parse calls on unchanged +// tool call data. + +import { describe, it, expect, vi } from 'vitest'; +import { deriveAgenticSections } from '$lib/utils/agentic'; +import type { ApiChatCompletionToolCall } from '$lib/types/api'; +import type { DatabaseMessage } from '$lib/types/database'; +import { MessageRole, AgenticSectionType } from '$lib/enums'; + +function makeMessage(overrides: Partial): DatabaseMessage { + return { + id: 'm1', + convId: 'c1', + type: 'text', + timestamp: 0, + role: MessageRole.ASSISTANT, + content: '', + parent: null, + children: [], + ...overrides + } as DatabaseMessage; +} + +describe('parseToolCalls memoization', () => { + it('returns the same array reference for the same JSON string', () => { + // parseToolCalls is not exported, but deriveAgenticSections uses it + // internally. We verify memoization through behavior: calling + // deriveAgenticSections twice with the same toolCalls should not + // re-parse (which we verify by checking the returned sections + // are equivalent). + const toolCallsJson = JSON.stringify([ + { id: 'call_1', type: 'function', function: { name: 'test', arguments: '{}' } } + ]); + + const msg = makeMessage({ content: 'hello', toolCalls: toolCallsJson }); + const sections1 = deriveAgenticSections(msg, [], [], false); + const sections2 = deriveAgenticSections(msg, [], [], false); + + expect(sections1).toHaveLength(sections2.length); + expect(sections1[0].type).toBe(sections2[0].type); + }); + + it('does not re-parse JSON on cache hit', () => { + const toolCallsJson = JSON.stringify([ + { id: 'call_1', type: 'function', function: { name: 'test', arguments: '{}' } } + ]); + + const msg = makeMessage({ content: 'hello', toolCalls: toolCallsJson }); + const spy = vi.spyOn(JSON, 'parse'); + + deriveAgenticSections(msg, [], [], false); + const callsAfterFirst = spy.mock.calls.length; + + deriveAgenticSections(msg, [], [], false); + expect(spy.mock.calls.length).toBe(callsAfterFirst); + + spy.mockRestore(); + }); + + it('handles empty/undefined toolCalls without error', () => { + const msg = makeMessage({ content: 'hello' }); + const sections = deriveAgenticSections(msg, [], [], false); + + expect(sections).toHaveLength(1); + expect(sections[0].type).toBe(AgenticSectionType.TEXT); + }); + + it('handles invalid JSON gracefully', () => { + const msg = makeMessage({ content: 'hello', toolCalls: '{invalid' }); + const sections = deriveAgenticSections(msg, [], [], false); + + // Should return just the text section, no tool call sections + expect(sections).toHaveLength(1); + expect(sections[0].type).toBe(AgenticSectionType.TEXT); + }); +}); + +describe('deriveAgenticSections O(1) tool message lookup', () => { + it('matches tool messages to tool calls by toolCallId', () => { + const toolCallsJson = JSON.stringify([ + { id: 'call_1', type: 'function', function: { name: 'test_1', arguments: '{}' } }, + { id: 'call_2', type: 'function', function: { name: 'test_2', arguments: '{}' } } + ]); + + const toolMessages = [ + makeMessage({ role: MessageRole.TOOL, toolCallId: 'call_1', content: 'result_1' }), + makeMessage({ role: MessageRole.TOOL, toolCallId: 'call_2', content: 'result_2' }) + ]; + + const msg = makeMessage({ content: 'hello', toolCalls: toolCallsJson }); + const sections = deriveAgenticSections(msg, toolMessages, [], false); + + // Expect: TEXT + 2 TOOL_CALL sections + const toolCallSections = sections.filter((s) => s.type === AgenticSectionType.TOOL_CALL); + expect(toolCallSections).toHaveLength(2); + expect(toolCallSections[0].toolResult).toBe('result_1'); + expect(toolCallSections[1].toolResult).toBe('result_2'); + }); + + it('handles missing tool messages (pending calls during streaming)', () => { + const toolCallsJson = JSON.stringify([ + { id: 'call_1', type: 'function', function: { name: 'test', arguments: '{}' } } + ]); + + const msg = makeMessage({ content: '', toolCalls: toolCallsJson }); + const sections = deriveAgenticSections(msg, [], [], true); + + const toolCallSection = sections.find((s) => s.type === AgenticSectionType.TOOL_CALL_PENDING); + expect(toolCallSection).toBeDefined(); + expect(toolCallSection?.content).toBe(''); + }); + + it('scales with many tool calls (no O(n^2) blowup)', () => { + const N = 100; + const toolCalls = Array.from( + { length: N }, + (_, i): ApiChatCompletionToolCall => ({ + id: `call_${i}`, + type: 'function', + function: { name: `tool_${i}`, arguments: '{}' } + }) + ); + const toolCallsJson = JSON.stringify(toolCalls); + + const toolMessages = Array.from({ length: N }, (_, i) => + makeMessage({ + role: MessageRole.TOOL, + toolCallId: `call_${i}`, + content: `result_${i}` + }) + ); + + const msg = makeMessage({ content: 'hello', toolCalls: toolCallsJson }); + + // If the lookup were still O(n^2), this would be noticeably slow + const start = Date.now(); + const sections = deriveAgenticSections(msg, toolMessages, [], false); + const elapsed = Date.now() - start; + + const toolCallSections = sections.filter((s) => s.type === AgenticSectionType.TOOL_CALL); + expect(toolCallSections).toHaveLength(N); + expect(elapsed).toBeLessThan(100); // Should be fast with O(1) lookup + }); +});