Update scripts/snapdragon/run.py

Co-authored-by: Jhen-Jie Hong <iainst0409@gmail.com>

common/arg.cpp

@@ -2776,6 +2776,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             llm_add_n_cpu_ffn_overrides(value, LLM_FFN_EXPS_REGEX, params.tensor_buft_overrides);
         }
     ).set_env("LLAMA_ARG_N_CPU_MOE"));
+    add_opt(common_arg(
+        {"--moe-cache-mib"}, "N",
+        "GPU cache size in MiB for the MoE experts kept in the CPU (default: 0, disabled)",
+        [](common_params & params, int value) {
+            if (value < 0) {
+                throw std::invalid_argument("invalid value");
+            }
+            params.moe_cache_size = (size_t) value*1024*1024;
+        }
+    ).set_env("LLAMA_ARG_MOE_CACHE_MIB"));
     add_opt(common_arg(
         {"-ncffn", "--n-cpu-ffn"}, "N",
         "keep the dense FFN weights of the first N layers in the CPU\n"

common/chat-auto-parser-generator.cpp

@@ -291,7 +291,7 @@ common_peg_parser analyze_tools::build_tool_parser_tag_json(parser_build_context
 
     common_peg_parser tool_choice = p.choice();
 
-    foreach_function(inputs.tools, [&](const json & tool) {
+    foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
         const auto & func   = tool.at("function");
         std::string  name   = func.at("name");
         const auto   schema = common_chat_tool_parameters(func);
@@ -308,7 +308,7 @@ common_peg_parser analyze_tools::build_tool_parser_tag_json(parser_build_context
             }
             have_call_id = true;
         }
-        auto args_parser = p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema));
+        auto args_parser = p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", schema));
         if (!arguments.start.empty()) {
             args_parser = p.literal(arguments.start) + args_parser;
         }
@@ -318,7 +318,7 @@ common_peg_parser analyze_tools::build_tool_parser_tag_json(parser_build_context
 
         auto atomic_peek = !arguments.start.empty() ? std::optional(p.peek(p.literal(arguments.start))) : std::nullopt;
         auto func_parser = build_func_parser(p, name, call_id_section, have_call_id, args_parser, atomic_peek);
-        tool_choice |= p.rule("tool-" + name, func_parser);
+        tool_choice |= p.rule("tool-" + std::to_string(tool_index), func_parser);
     });
 
     auto require_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED;
@@ -364,14 +364,14 @@ common_peg_parser analyze_tools::build_tool_parser_tag_tagged(parser_build_conte
 
     common_peg_parser tool_choice = p.choice();
 
-    foreach_function(inputs.tools, [&](const json & tool) {
+    foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
         const auto & func = tool.at("function");
         std::string  name = func.at("name");
 
         // Build parser for each argument, separating required and optional
         std::vector<common_peg_parser> required_parsers;
         std::vector<common_peg_parser> optional_parsers;
-        foreach_parameter(func, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
+        foreach_parameter(func, [&](size_t param_index, const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
             auto arg =
                 p.tool_arg(p.tool_arg_open(arguments.name_prefix + p.tool_arg_name(p.literal(param.name)) +
                                            arguments.name_suffix) +
@@ -380,10 +380,10 @@ common_peg_parser analyze_tools::build_tool_parser_tag_tagged(parser_build_conte
                                 p.ac(p.tool_arg_string_value(until_suffix) +
                                     p.tool_arg_close(p.literal(arguments.value_suffix)), arguments.value_suffix) :
                                 (p.tool_arg_json_value(p.schema(
-                                    p.json(), "tool-" + name + "-arg-" + param.name + "-schema", doc, *param.schema)) +
+                                    p.json(), "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index) + "-schema", doc, *param.schema)) +
                                     p.tool_arg_close(p.literal(arguments.value_suffix)))));
 
-            auto named_arg = p.rule("tool-" + name + "-arg-" + param.name, arg);
+            auto named_arg = p.rule("tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index), arg);
             if (param.required) {
                 required_parsers.push_back(named_arg);
             } else {
@@ -434,7 +434,7 @@ common_peg_parser analyze_tools::build_tool_parser_tag_tagged(parser_build_conte
         auto atomic_peek = (!arguments.name_prefix.empty() && !required_parsers.empty()) ?
             std::optional(p.peek(p.literal(arguments.name_prefix))) : std::nullopt;
         auto func_parser = build_func_parser(p, name, call_id_section, have_call_id, args_seq, atomic_peek);
-        tool_choice |= p.rule("tool-" + name, func_parser);
+        tool_choice |= p.rule("tool-" + std::to_string(tool_index), func_parser);
     });
 
     auto require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED;

common/chat-peg-parser.cpp

@@ -483,7 +483,9 @@ common_peg_parser common_chat_peg_builder::standard_constructed_tools(
     // Build tool choices for tagged format
     auto tool_choices = choice();
 
-    for (const auto & tool_def : tools) {
+    for (size_t i = 0; i < tools.size(); i++) {
+        const auto & tool_def = tools[i];
+
         if (!tool_def.contains("function")) {
             continue;
         }
@@ -513,7 +515,7 @@ common_peg_parser common_chat_peg_builder::standard_constructed_tools(
         auto tool_parser = tool(tool_open(literal(func_opener) + tool_name(literal(name)) + literal(func_name_suffix)) +
                                 space() + tool_args(args) + space() + tool_close(literal(func_closer)));
 
-        tool_choices |= rule("tool-" + name, tool_parser);
+        tool_choices |= rule("tool-" + std::to_string(i), tool_parser);
     }
 
     // Build the section with markers
@@ -560,7 +562,8 @@ common_peg_parser common_chat_peg_builder::python_style_tool_calls(
 
     auto tool_choices = choice();
 
-    for (const auto & tool_def : tools) {
+    for (size_t i = 0; i < tools.size(); i++) {
+        const auto & tool_def = tools[i];
         if (!tool_def.contains("function")) {
             continue;
         }
@@ -607,7 +610,7 @@ common_peg_parser common_chat_peg_builder::python_style_tool_calls(
             space() + tool_args(args) + space() + tool_close(literal(")"))
         );
 
-        tool_choices |= rule("tool-" + name, tool_parser);
+        tool_choices |= rule("tool-" + std::to_string(i), tool_parser);
     }
 
     if (parallel_tool_calls) {
@@ -635,7 +638,8 @@ common_peg_parser common_chat_peg_builder::build_json_tools_function_is_key(
 
     auto tool_choices = choice();
 
-    for (const auto & tool_def : tools) {
+    for (size_t i = 0; i < tools.size(); i++) {
+        const auto & tool_def = tools[i];
         if (!tool_def.contains("function")) {
             continue;
         }
@@ -668,10 +672,10 @@ common_peg_parser common_chat_peg_builder::build_json_tools_function_is_key(
         // Arguments — either wrapped in args_key or parsed directly
         common_peg_parser args_parser = eps();
         if (args_key.empty()) {
-            args_parser = tool_args(schema(json(), "tool-" + name + "-schema", params));
+            args_parser = tool_args(schema(json(), "tool-" + std::to_string(i) + "-schema", params));
         } else {
             args_parser = literal("\"" + effective_args_key + "\"") + space() + literal(":") + space() +
-                          tool_args(schema(json(), "tool-" + name + "-schema", params));
+                          tool_args(schema(json(), "tool-" + std::to_string(i) + "-schema", params));
         }
         inner_fields.push_back(args_parser);
 
@@ -698,7 +702,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_function_is_key(
             space() + tool_close(literal("}"))
         );
 
-        tool_choices |= rule("tool-" + name, tool_parser);
+        tool_choices |= rule("tool-" + std::to_string(i), tool_parser);
     }
 
     return tool_choices;
@@ -721,7 +725,8 @@ common_peg_parser common_chat_peg_builder::build_json_tools_nested_keys(
     std::string nested_name_field = !name_spec.first.empty() ? name_spec.second  : effective_name_key;
     std::string nested_args_field = !args_spec.first.empty() ? args_spec.second  : effective_args_key;
 
-    for (const auto & tool_def : tools) {
+    for (size_t i = 0; i < tools.size(); i++) {
+        const auto & tool_def = tools[i];
         if (!tool_def.contains("function")) {
             continue;
         }
@@ -732,7 +737,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_nested_keys(
         auto nested_name = literal("\"" + nested_name_field + "\"") + space() + literal(":") + space() +
                           atomic(literal("\"") + tool_name(literal(name)) + literal("\""));
         auto nested_args = literal("\"" + nested_args_field + "\"") + space() + literal(":") + space() +
-                          tool_args(schema(json(), "tool-" + name + "-schema", params));
+                          tool_args(schema(json(), "tool-" + std::to_string(i) + "-schema", params));
 
         auto nested_object = literal("{") + space() +
                             nested_name + space() + literal(",") + space() +
@@ -770,7 +775,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_nested_keys(
         auto nested_field = literal("\"" + nested_prefix + "\"") + space() + literal(":") + space() + nested_object;
         tool_parser_body = tool_parser_body + nested_field + space() + tool_close(literal("}"));
 
-        tool_choices |= rule("tool-" + name, tool(tool_parser_body));
+        tool_choices |= rule("tool-" + std::to_string(i), tool(tool_parser_body));
     }
 
     return tool_choices;
@@ -790,7 +795,8 @@ common_peg_parser common_chat_peg_builder::build_json_tools_flat_keys(
     auto name_key_parser = literal("\"" + effective_name_key + "\"");
     auto args_key_parser = literal("\"" + effective_args_key + "\"");
 
-    for (const auto & tool_def : tools) {
+    for (size_t i = 0; i < tools.size(); i++) {
+        const auto & tool_def = tools[i];
         if (!tool_def.contains("function")) {
             continue;
         }
@@ -801,7 +807,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_flat_keys(
         auto tool_name_ = name_key_parser + space() + literal(":") + space() +
                          atomic(literal("\"") + tool_name(literal(name)) + literal("\""));
         auto tool_args_ = args_key_parser + space() + literal(":") + space() +
-                         tool_args(schema(json(), "tool-" + name + "-schema", params));
+                         tool_args(schema(json(), "tool-" + std::to_string(i) + "-schema", params));
 
         // Build ID parsers if keys are provided
         common_peg_parser id_parser = eps();
@@ -861,7 +867,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_flat_keys(
         }
         ordered_body = ordered_body + space() + tool_close(literal("}"));
 
-        tool_choices |= rule("tool-" + name, tool(ordered_body));
+        tool_choices |= rule("tool-" + std::to_string(i), tool(ordered_body));
     }
 
     return tool_choices;

common/chat.cpp

@@ -1223,6 +1223,13 @@ std::optional<common_chat_params> common_chat_try_specialized_template(
         return common_chat_params_init_minicpm5(tmpl, params);
     }
 
+    // TranslateGemma - user content must follow a custom schema with language codes
+    if (src.find("[source_lang_code]") != std::string::npos &&
+        src.find("[target_lang_code]") != std::string::npos) {
+        LOG_DBG("Using specialized template: TranslateGemma\n");
+        return common_chat_params_init_translate_gemma(tmpl, params);
+    }
+
     // Qwen3-Coder XML tool calls, also used by Nemotron Nano 3, Qwen3.5 and StepFun-3.5-Flash
     if (src.find("<tool_call>") != std::string::npos &&
         src.find("<function=") != std::string::npos &&
@@ -1452,14 +1459,70 @@ common_chat_params common_chat_templates_apply(const struct common_chat_template
                               common_chat_templates_apply_legacy(tmpls, inputs);
 }
 
-common_chat_msg common_chat_parse(const std::string &               input,
+void common_chat_input::append(const std::string & piece, llama_token token) {
+    if (piece.empty()) {
+        return;
+    }
+    tokens.push_back(token);
+    tokens.resize(tokens.size() + piece.size() - 1, LLAMA_TOKEN_NULL);
+    text += piece;
+}
+
+void common_chat_input::append(const common_chat_input & chunk) {
+    tokens.insert(tokens.end(), chunk.tokens.begin(), chunk.tokens.end());
+    text += chunk.text;
+}
+
+void common_chat_input::truncate(size_t pos) {
+    if (pos < text.size()) {
+        text.erase(pos);
+        tokens.resize(pos);
+    }
+}
+
+common_chat_input common_chat_input::substr(size_t pos, size_t n) const {
+    common_chat_input out;
+    out.text = text.substr(pos, n);
+    out.tokens.assign(tokens.begin() + pos, tokens.begin() + pos + out.size());
+    return out;
+}
+
+void common_chat_input::prepend(const std::string & prefix) {
+    tokens.insert(tokens.begin(), prefix.size(), LLAMA_TOKEN_NULL);
+    text = prefix + text;
+}
+
+void common_chat_input::prepend(const common_chat_input & prefix) {
+    tokens.insert(tokens.begin(), prefix.tokens.begin(), prefix.tokens.end());
+    text = prefix.text + text;
+}
+
+common_chat_input common_chat_input_tokenize(const llama_vocab * vocab, const std::string & text) {
+    common_chat_input input;
+    auto tokens = common_tokenize(vocab, text, false, true);
+    for (size_t i = 0; i < tokens.size(); i++) {
+        std::string piece = common_token_to_piece(vocab, tokens[i], true);
+        if (i == 0 && std::isspace(piece[0]) && !std::isspace(text[0])) {
+            // Some tokenizers will add a space before the first special token, need to exclude
+            continue;
+        }
+        input.append(piece, tokens[i]);
+    }
+    if (input.text != text) {
+        // the pieces do not give back the same text, keep the text without tokens
+        return common_chat_input(text);
+    }
+    return input;
+}
+
+common_chat_msg common_chat_parse(const common_chat_input &         input,
                                   bool                              is_partial,
                                   const common_chat_parser_params & params) {
     return common_chat_peg_parse(params.parser, input, is_partial, params);
 }
 
 common_chat_msg common_chat_peg_parse(const common_peg_arena &          src_parser,
-                                      const std::string &               input,
+                                      const common_chat_input &         input,
                                       bool                              is_partial,
                                       const common_chat_parser_params & params) {
     const common_peg_arena & parser = src_parser.empty() ?
@@ -1470,18 +1533,17 @@ common_chat_msg common_chat_peg_parse(const common_peg_arena &          src_pars
         LOG_DBG("No parser definition detected, assuming pure content parser.");
     }
 
-    const std::string effective_input = params.generation_prompt.empty()
-        ? input
-        : params.generation_prompt + input;
+    common_chat_input effective_input = input;
+    effective_input.prepend(params.generation_prompt);
 
-    //LOG_DBG("Parsing PEG input with format %s: %s\n", common_chat_format_name(params.format), effective_input.c_str());
+    //LOG_DBG("Parsing PEG input with format %s: %s\n", common_chat_format_name(params.format), effective_input.text.c_str());
 
     common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_LENIENT;
     if (params.debug) {
         flags |= COMMON_PEG_PARSE_FLAG_DEBUG;
     }
 
-    common_peg_parse_context ctx(effective_input, flags);
+    common_peg_parse_context ctx(std::move(effective_input.text), std::move(effective_input.tokens), flags);
     auto result = parser.parse(ctx);
 
     if (result.fail()) {
@@ -1507,8 +1569,8 @@ common_chat_msg common_chat_peg_parse(const common_peg_arena &          src_pars
             }
             return msg;
         }
-        LOG_WRN("%s: unparsed %s output: %s\n", __func__, common_chat_format_name(params.format), effective_input.substr(result.end).c_str());
-        LOG_DBG("%s: full %s output triggering error:\n=== BEGIN ===\n%s\n=== END ===\n", __func__, common_chat_format_name(params.format), effective_input.c_str());
+        LOG_WRN("%s: unparsed %s output: %s\n", __func__, common_chat_format_name(params.format), ctx.input.substr(result.end).c_str());
+        LOG_DBG("%s: full %s output triggering error:\n=== BEGIN ===\n%s\n=== END ===\n", __func__, common_chat_format_name(params.format), ctx.input.c_str());
         throw std::runtime_error(std::string("The model produced output that does not match the expected ") + common_chat_format_name(params.format) + " format");
     }
 

common/chat.h

@@ -282,6 +282,31 @@ struct common_chat_params {
     common_chat_msg_delimiters          message_delimiters;
 };
 
+struct common_chat_input {
+    std::string              text;
+    std::vector<llama_token> tokens;
+
+    common_chat_input() = default;
+
+    // plain text, with no tokens
+    explicit common_chat_input(std::string text) : text(std::move(text)), tokens(this->text.size(), LLAMA_TOKEN_NULL) {}
+
+    size_t size() const { return text.size(); }
+    bool empty() const { return text.empty(); }
+
+    void append(const std::string & piece, llama_token token);
+    void append(const common_chat_input & chunk);
+
+    void prepend(const std::string & prefix);
+    void prepend(const common_chat_input & prefix);
+
+    void truncate(size_t pos);
+
+    common_chat_input substr(size_t pos, size_t n = std::string::npos) const;
+};
+
+common_chat_input common_chat_input_tokenize(const llama_vocab * vocab, const std::string & text);
+
 // per-message parsing syntax
 // should be derived from common_chat_params
 struct common_chat_parser_params {
@@ -289,7 +314,7 @@ struct common_chat_parser_params {
     common_reasoning_format reasoning_format     = COMMON_REASONING_FORMAT_NONE; // TODO: refactor this to "bool parse_reasoning"
     // Whether reasoning_content should be inlined in the content (e.g. for reasoning_format=deepseek in stream mode)
     bool                    reasoning_in_content = false;
-    std::string             generation_prompt;
+    common_chat_input       generation_prompt;
     bool                    parse_tool_calls     = true;
     bool                    is_continuation      = false;
     bool                    echo                 = false;  // Include assistant prefilled msg in output
@@ -298,7 +323,7 @@ struct common_chat_parser_params {
     common_chat_parser_params() = default;
     common_chat_parser_params(const common_chat_params & chat_params) {
         format  = chat_params.format;
-        generation_prompt = chat_params.generation_prompt;
+        generation_prompt = common_chat_input(chat_params.generation_prompt);
     }
 };
 
@@ -337,8 +362,8 @@ std::string common_chat_format_example(const struct common_chat_templates *
                                        const std::map<std::string, std::string> & chat_template_kwargs);
 
 const char *    common_chat_format_name(common_chat_format format);
-common_chat_msg common_chat_parse(const std::string & input, bool is_partial, const common_chat_parser_params & params);
-common_chat_msg common_chat_peg_parse(const common_peg_arena & src_parser, const std::string & input, bool is_partial, const common_chat_parser_params & params);
+common_chat_msg common_chat_parse(const common_chat_input & input, bool is_partial, const common_chat_parser_params & params);
+common_chat_msg common_chat_peg_parse(const common_peg_arena & src_parser, const common_chat_input & input, bool is_partial, const common_chat_parser_params & params);
 
 // used by arg and server
 const char *            common_reasoning_format_name(common_reasoning_format format);

common/common.cpp

@@ -1170,6 +1170,7 @@ static const std::map<common_decision_type, std::string> COMMON_DECISION_TYPE_NA
     { COMMON_DECISION_TYPE_LAYA,           "laya"          },
     { COMMON_DECISION_TYPE_CLEF,           "clef"          },
     { COMMON_DECISION_TYPE_PPLX_DECIDER,   "pplx-decider"  },
+    { COMMON_DECISION_TYPE_LFM2_D1,        "lfm2-d1"       },
 };
 
 static common_decision_type common_decision_type_from_string(const std::string & str) {
@@ -1722,6 +1723,8 @@ struct llama_context_params common_context_params_to_llama(const common_params &
     cparams.type_k = params.cache_type_k;
     cparams.type_v = params.cache_type_v;
 
+    cparams.moe_cache_size = params.moe_cache_size;
+
     return cparams;
 }
 

common/common.h

@@ -593,6 +593,8 @@ struct common_params {
     ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K
     ggml_type cache_type_v = GGML_TYPE_F16; // KV cache data type for the V
 
+    size_t moe_cache_size = 0; // GPU cache size in bytes for the MoE experts kept in the CPU
+
     common_conversation_mode conversation_mode = COMMON_CONVERSATION_MODE_AUTO;
 
     // multimodal models (see tools/mtmd)
@@ -961,6 +963,7 @@ enum common_decision_type {
     COMMON_DECISION_TYPE_LAYA,    // score of one marker token per option, read from the embeddings output
     COMMON_DECISION_TYPE_CLEF,    // all questions in one prompt, score of option i read from the embeddings output at row i
     COMMON_DECISION_TYPE_PPLX_DECIDER, // same as openjev, label codes of 1 or 2 letters
+    COMMON_DECISION_TYPE_LFM2_D1, // same as openjev, the labels depend on the question type
     COMMON_DECISION_TYPE_UNKNOWN, // a decision model of a type that is not supported
 };
 

common/parsers/deepseek.cpp

@@ -152,13 +152,13 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
         // build tool call section first since we might need it in reasoning
         auto tool_choice = p.choice();
         if (has_tool_calls) {
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
 
                 std::vector<common_peg_parser> required_parsers;
                 std::vector<common_peg_parser> optional_parsers;
-                foreach_parameter(function, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
+                foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
                     bool is_string = param.schema->may_be_string();
 
                     auto arg = p.tool_arg(
@@ -166,11 +166,11 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
                                         p.literal("\" string=\"" + std::string(is_string ? "true" : "false") + "\">")) +
                         (is_string ?
                              p.tool_arg_string_value(p.until(PARAM_END)) :
-                             p.tool_arg_json_value(p.schema(p.json(), "tool-" + name + "-arg-" + param.name + "-schema",
+                             p.tool_arg_json_value(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index) + "-schema",
                                                             doc, *param.schema))) +
                         p.tool_arg_close(p.literal(PARAM_END)));
 
-                    auto named_arg = p.rule("tool-" + name + "-arg-" + param.name, arg);
+                    auto named_arg = p.rule("tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index), arg);
                     if (param.required) {
                         required_parsers.push_back(named_arg);
                     } else {
@@ -199,7 +199,7 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
                                                                    p.tool_name(p.literal(name)) + p.literal("\">\n")) +
                                                        invoke_body + p.space() + p.tool_close(p.literal(INVOKE_END)));
 
-                tool_choice |= p.rule("tool-" + name, func_parser);
+                tool_choice |= p.rule("tool-" + std::to_string(tool_index), func_parser);
             });
         }
 

common/parsers/functionary-v3-2.cpp

@@ -42,7 +42,7 @@ common_chat_params common_chat_params_init_functionary_v3_2(const common_chat_te
 
         // Build tool call parsers for each available function
         auto tool_choice = p.choice();
-        foreach_function(inputs.tools, [&](const json & tool) {
+        foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
             const auto & function = tool.at("function");
             std::string  name     = function.at("name");
             const auto   schema   = common_chat_tool_parameters(function);
@@ -50,10 +50,10 @@ common_chat_params common_chat_params_init_functionary_v3_2(const common_chat_te
             // Tool format: >>>function_name\n{json_args}
             auto tool_parser = p.tool(
                 p.tool_open(p.tool_name(p.literal(name)) + p.literal("\n")) +
-                p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema))
+                p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", schema))
             );
 
-            tool_choice |= p.rule("tool-" + name, tool_parser);
+            tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_parser);
         });
 
         auto content_only = content_until_end;

common/parsers/gemma4.cpp

@@ -254,13 +254,13 @@ common_chat_params common_chat_params_init_gemma4(const common_chat_template &
 
             auto tool_choice = p.choice();
 
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
                 // TODO @aldehir : need to extend json-schema-to-grammar to produce more than JSON rules
                 // const auto & params   = function.at("parameters");
 
-                tool_choice |= p.rule("tool-" + name, p.tool(p.sequence({
+                tool_choice |= p.rule("tool-" + std::to_string(tool_index), p.tool(p.sequence({
                     p.tool_open(p.tool_name(p.literal(name)) + p.peek(p.literal("{"))),
                     p.tool_args(p.ref("gemma4-dict")),
                 })));

common/parsers/gigachat-v3.cpp

@@ -30,17 +30,18 @@ common_chat_params common_chat_params_init_gigachat_v3(
         if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
             // Build a choice of all available tools
             auto tool_choice = p.choice();
-            for (const auto & tool : inputs.tools) {
+            for (size_t i = 0; i < inputs.tools.size(); i++) {
+                const auto & tool = inputs.tools[i];
                 const auto & function = tool.at("function");
                 std::string name = function.at("name");
                 const auto  schema = common_chat_tool_parameters(function);
 
                 auto tool_name = p.json_member("name", "\"" + p.tool_name(p.literal(name)) + "\"");
-                auto tool_args = p.json_member("arguments", p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema)));
+                auto tool_args = p.json_member("arguments", p.tool_args(p.schema(p.json(), "tool-" + std::to_string(i) + "-schema", schema)));
 
                 auto tool_open = p.tool_open(p.literal("{") << tool_name);
 
-                tool_choice |= p.rule("tool-" + name, tool_open << "," << tool_args << "}");
+                tool_choice |= p.rule("tool-" + std::to_string(i), tool_open << "," << tool_args << "}");
             }
 
             // Define the tool call structure

common/parsers/gpt-oss.cpp

@@ -106,14 +106,14 @@ common_chat_params common_chat_params_init_gpt_oss(const common_chat_template &
         if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
             auto tool_choice = p.choice();
 
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
                 const auto   params   = common_chat_tool_parameters(function);
 
                 auto func_name  = p.literal(" to=functions.") + p.tool_name(p.literal(name));
                 auto constraint = p.optional(p.space() + p.optional(p.literal("<|constrain|>")) + constrain_type);
-                auto args       = p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", params));
+                auto args       = p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", params));
 
                 // recipient in role header
                 //   <|start|>assistant to=functions.NAME<|channel|>(commentary|analysis)[constraint]<|message|>ARGS
@@ -123,7 +123,7 @@ common_chat_params common_chat_params_init_gpt_oss(const common_chat_template &
                 //   <|channel|>(commentary|analysis) to=functions.NAME[constraint]<|message|>ARGS
                 auto tool_in_channel = p.tool(p.tool_open(channel + func_name + constraint + p.literal("<|message|>")) + args);
 
-                tool_choice |= p.rule("tool-" + name, tool_in_role | tool_in_channel);
+                tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_in_role | tool_in_channel);
             });
 
             auto tool_call  = p.trigger_rule("tool-call", tool_choice);

common/parsers/k2-horizon.cpp

@@ -110,14 +110,14 @@ common_chat_params common_chat_params_init_k2_horizon(const common_chat_template
             // The models leave out <ifm|arg_type> even when asked for xml_typed
             auto arg_type = call_format == "xml_typed" ? p.optional(ARG_TYPE + p.until(ARG_TYPE_END) + ARG_TYPE_END + p.space()) : p.eps();
 
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
 
                 std::vector<common_peg_parser> required_args;
                 std::vector<common_peg_parser> optional_args;
-                foreach_parameter(function, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
-                    auto rule_name = "tool-" + name + "-arg-" + param.name;
+                foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
+                    auto rule_name = "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index);
                     auto types     = param.schema->value_types();
                     auto arg_value = arg_string;
                     if (!types.has(common_chat_schema::TYPE_STRING)) {
@@ -149,12 +149,12 @@ common_chat_params common_chat_params_init_k2_horizon(const common_chat_template
                     (param.required ? required_args : optional_args).push_back(p.rule(rule_name, arg));
                 });
 
-                auto args = p.permute("tool-" + name + "-args", required_args);
+                auto args = p.permute("tool-" + std::to_string(tool_index) + "-args", required_args);
                 if (!optional_args.empty()) {
                     args = args + p.zero_or_more(p.choice(optional_args));
                 }
 
-                tool_choice |= p.rule("tool-" + name, p.tool(
+                tool_choice |= p.rule("tool-" + std::to_string(tool_index), p.tool(
                     p.tool_open(CALL_START + p.tool_name(p.literal(name)) + "\n") + p.tool_args(args) << p.tool_close(p.literal(CALL_END))));
             });
         }

common/parsers/kimi-k2.cpp

@@ -79,7 +79,7 @@ common_chat_params common_chat_params_init_kimi_k2(const common_chat_template &
         // The ID format is: functions.<name>:<index>
         // We need to match: functions.<name>:<digits>
         auto tool_choice = p.choice();
-        foreach_function(inputs.tools, [&](const json & tool) {
+        foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
             const auto & function = tool.at("function");
             std::string  name     = function.at("name");
             const auto   schema   = common_chat_tool_parameters(function);
@@ -89,11 +89,11 @@ common_chat_params common_chat_params_init_kimi_k2(const common_chat_template &
             auto tool_id = p.tool_id(p.literal("functions.") + p.tool_name(p.literal(name)) + p.literal(":") + p.chars("[0-9]", 1, -1));
             auto tool_parser = p.tool(
                 p.tool_open(tool_id + p.literal(ARGS_BEGIN)) +
-                p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema)) +
+                p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", schema)) +
                 p.tool_close(p.optional((p.literal(CALL_END))))
             );
 
-            tool_choice |= p.rule("tool-" + name, tool_parser);
+            tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_parser);
         });
 
         // Tool calls section: <|tool_calls_section_begin|> tool_calls <|tool_calls_section_end|>

common/parsers/kimi-k3.cpp

@@ -95,7 +95,7 @@ common_chat_params common_chat_params_init_kimi_k3(const common_chat_template &
         }
 
         auto tool_choices = p.choice();
-        foreach_function(inputs.tools, [&](const json & tool) {
+        foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
             const auto & function = tool.at("function");
             std::string  name     = function.at("name");
             const json   schema   = common_chat_tool_parameters(function);
@@ -106,6 +106,7 @@ common_chat_params common_chat_params_init_kimi_k3(const common_chat_template &
             auto args = p.eps();
             if (schema.contains("properties") && !schema.at("properties").empty()) {
                 auto arg_choices = p.choice();
+                size_t param_index = 0;
                 for (const auto & prop : schema.at("properties").items()) {
                     const std::string & key = prop.key();
 
@@ -119,7 +120,7 @@ common_chat_params common_chat_params_init_kimi_k3(const common_chat_template &
                                                     p.tool_arg_value(p.until(ARG_END));
 
                     // skip the trailing type="..." attribute: anything up to <|sep|>
-                    arg_choices |= p.rule("kimi-k3-arg-" + name + "-" + key,
+                    arg_choices |= p.rule("kimi-k3-arg-" + std::to_string(tool_index) + "-" + std::to_string(param_index++),
                                           p.tool_arg(p.tool_arg_open(p.literal(ARG_START)) +
                                                      p.tool_arg_name(p.literal(key)) + p.literal("\"") +
                                                      p.until(SEP) + p.literal(SEP) + value +
@@ -133,7 +134,7 @@ common_chat_params common_chat_params_init_kimi_k3(const common_chat_template &
                                            p.until(SEP) + p.literal(SEP)) +
                                p.tool_args(args) + p.tool_close(p.literal(CALL_END)));
 
-            tool_choices |= p.rule("kimi-k3-tool-" + name, call);
+            tool_choices |= p.rule("kimi-k3-tool-" + std::to_string(tool_index), call);
         });
 
         // all calls go inside one tools section, then the message is closed. the

common/parsers/ling3.cpp

@@ -118,7 +118,7 @@ common_chat_params common_chat_params_init_ling3(const common_chat_template &
         auto arg_string   = p.rule("ling3-arg-string",
                                    p.tool_arg_string_value(p.until(ARG_VAL_END)) + arg_close);
 
-        foreach_function(inputs.tools, [&](const json & tool) {
+        foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
             const auto & function = tool.at("function");
             std::string  name     = function.at("name");
 
@@ -127,8 +127,8 @@ common_chat_params common_chat_params_init_ling3(const common_chat_template &
 
             // each argument may be preceded by whitespace: the model emits
             // newlines between arguments, the template history does not
-            foreach_parameter(function, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
-                auto rule_name = "ling3-arg-" + name + "-" + param.name;
+            foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
+                auto rule_name = "ling3-arg-" + std::to_string(tool_index) + "-" + std::to_string(param_index);
 
                 auto types = param.schema->value_types();
 
@@ -159,7 +159,7 @@ common_chat_params common_chat_params_init_ling3(const common_chat_template &
 
             // required arguments in any order (as Qwen3-Coder does), then
             // optional ones in any order and number
-            auto args = p.permute("ling3-" + name + "-args", required_args);
+            auto args = p.permute("ling3-" + std::to_string(tool_index) + "-args", required_args);
             if (!optional_args.empty()) {
                 args = args + p.zero_or_more(p.choice(optional_args));
             }
@@ -169,7 +169,7 @@ common_chat_params common_chat_params_init_ling3(const common_chat_template &
                                p.tool_args(args) +
                                p.tool_close(p.optional(p.space()) + p.literal(CALL_END)));
 
-            tool_choices |= p.rule("ling3-tool-" + name, call);
+            tool_choices |= p.rule("ling3-tool-" + std::to_string(tool_index), call);
         });
 
         auto calls = inputs.parallel_tool_calls ?

common/parsers/llm-jp-harmony.cpp

@@ -109,13 +109,13 @@ common_chat_params common_chat_params_init_llm_jp_harmony(const common_chat_temp
         if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
             auto tool_choice = p.choice();
 
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
                 const auto   params   = common_chat_tool_parameters(function);
 
                 auto func_name = p.literal(" to=functions.") + p.tool_name(p.literal(name));
-                auto args      = p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", params));
+                auto args      = p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", params));
 
                 // recipient in role header
                 //   <|start|>assistant to=functions.NAME<|channel|>(commentary|analysis)[constraint]<|message|>ARGS
@@ -125,7 +125,7 @@ common_chat_params common_chat_params_init_llm_jp_harmony(const common_chat_temp
                 //   <|channel|>(commentary|analysis) to=functions.NAME[constraint]<|message|>ARGS
                 auto tool_in_channel = p.tool(p.tool_open(channel + func_name + constraint + message) + args);
 
-                tool_choice |= p.rule("tool-" + name, tool_in_role | tool_in_channel);
+                tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_in_role | tool_in_channel);
             });
 
             // parallel calls are separated by <|end|>; inside the trigger rule so the lazy grammar covers all of them

common/parsers/minicpm5.cpp

@@ -68,18 +68,18 @@ common_chat_params common_chat_params_init_minicpm5(const common_chat_template &
             });
 
             auto tool_choice = p.choice();
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto &      function = tool.at("function");
                 const std::string name     = function.at("name");
 
                 std::vector<common_peg_parser> arg_rules;
-                foreach_parameter(function, [&](const common_chat_schema_property & prop, const common_chat_schema_document_ptr & doc) {
+                foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & prop, const common_chat_schema_document_ptr & doc) {
                     auto value_parser = p.eps();
                     if (prop.schema->may_be_string()) {
                         value_parser = string_value;
                     } else {
                         value_parser = p.tool_arg_json_value(
-                                p.schema(p.json(), "tool-" + name + "-arg-" + prop.name + "-schema", doc, *prop.schema)
+                                p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index) + "-schema", doc, *prop.schema)
                             ) + p.tool_arg_close(p.literal("</param>"));
                     }
 
@@ -99,7 +99,7 @@ common_chat_params common_chat_params_init_minicpm5(const common_chat_template &
                     << p.tool_args(args)
                     << p.tool_close(p.literal("</function>")));
 
-                tool_choice |= p.rule("tool-" + name, tool_parser);
+                tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_parser);
             });
 
             auto max_calls  = inputs.parallel_tool_calls ? -1 : 1;

common/parsers/minimax-m3.cpp

@@ -85,7 +85,7 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template
         }
 
         auto tool_choice = p.choice();
-        foreach_function(inputs.tools, [&](const json & tool) {
+        foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
             const auto & function = tool.at("function");
             std::string  name     = function.at("name");
             auto         params   = common_chat_tool_parameters(function);
@@ -154,8 +154,9 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template
             members_of = [&](const common_chat_schema_object & object, const std::string & rule_prefix) -> common_peg_parser {
                 std::vector<common_peg_parser> required_elements;
                 std::vector<common_peg_parser> optional_elements;
-                for (const auto & prop : object.properties) {
-                    auto element = element_of(prop.name, *prop.schema, rule_prefix + "-" + prop.name);
+                for (size_t i = 0; i < object.properties.size(); i++) {
+                    const auto & prop    = object.properties[i];
+                    auto         element = element_of(prop.name, *prop.schema, rule_prefix + "-" + std::to_string(i));
                     (prop.required ? required_elements : optional_elements).push_back(element);
                 }
 
@@ -180,7 +181,7 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template
 
             common_peg_parser invoke_body = p.eps();
             if (doc->root->kind() == common_chat_schema::KIND_OBJECT) {
-                invoke_body = members_of(static_cast<const common_chat_schema_object &>(*doc->root), "tool-" + name + "-arg");
+                invoke_body = members_of(static_cast<const common_chat_schema_object &>(*doc->root), "tool-" + std::to_string(tool_index) + "-arg");
             }
 
             auto func_parser = p.tool(
@@ -189,7 +190,7 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template
                 p.space() + invoke_body + p.space() +
                 p.tool_close(p.literal(INVOKE_END)));
 
-            tool_choice |= p.rule("tool-" + name, func_parser);
+            tool_choice |= p.rule("tool-" + std::to_string(tool_index), func_parser);
         });
 
         auto require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED;

common/parsers/ministral3.cpp

@@ -86,14 +86,14 @@ common_chat_params common_chat_params_init_ministral_3(const common_chat_templat
         // Tool call parser
         if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
             auto tool_choice = p.choice();
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
                 const auto   schema   = common_chat_tool_parameters(function);
 
                 tool_choice |=
-                    p.rule("tool-" + name, p.tool_open(p.tool_name(p.literal(name)) + "[ARGS]") +
-                                               p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema)));
+                    p.rule("tool-" + std::to_string(tool_index), p.tool_open(p.tool_name(p.literal(name)) + "[ARGS]") +
+                                               p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", schema)));
             });
 
             auto min_calls  = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0;

common/parsers/muse-glimmer.cpp

@@ -81,18 +81,18 @@ common_chat_params common_chat_params_init_muse_glimmer(const common_chat_templa
                 "</atem:parameter>");
 
             auto tool_choice = p.choice();
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto &      function = tool.at("function");
                 const std::string name     = function.at("name");
 
                 std::vector<common_peg_parser> arg_rules;
-                foreach_parameter(function, [&](const common_chat_schema_property & prop, const common_chat_schema_document_ptr & doc) {
+                foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & prop, const common_chat_schema_document_ptr & doc) {
                     auto value_parser = p.eps();
                     if (prop.schema->may_be_string()) {
                         value_parser = string_value;
                     } else {
                         value_parser = p.tool_arg_json_value(
-                                p.schema(p.json(), "tool-" + name + "-arg-" + prop.name + "-schema", doc, *prop.schema))
+                                p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index) + "-schema", doc, *prop.schema))
                             + p.tool_arg_close(p.literal("</atem:parameter>"));
                     }
 
@@ -113,7 +113,7 @@ common_chat_params common_chat_params_init_muse_glimmer(const common_chat_templa
                     << p.tool_args(args)
                     << p.tool_close(p.literal("</atem:invoke>") + p.space() + p.literal("</atem:function_calls>")));
 
-                tool_choice |= p.rule("tool-" + name, tool_parser);
+                tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_parser);
             });
 
             auto tool_calls = inputs.parallel_tool_calls

common/parsers/parsers.cpp

@@ -2,24 +2,25 @@
 
 #include "log.h"
 
-void foreach_function(const json & tools, const std::function<void(const json &)> & fn) {
-    for (const auto & tool : tools) {
+void foreach_function(const json & tools, const std::function<void(size_t, const json &)> & fn) {
+    for (size_t i = 0; i < tools.size(); i++) {
+        const auto & tool = tools[i];
         if (!tool.contains("type") || tool.at("type") != "function" || !tool.contains("function")) {
             LOG_INF("Skipping tool without function: %s", tool.dump(2).c_str());
             continue;
         }
-        fn(tool);
+        fn(i, tool);
     }
 }
 
-void foreach_parameter(const json & function, const std::function<void(const common_chat_schema_property &, const common_chat_schema_document_ptr &)> & fn) {
+void foreach_parameter(const json & function, const std::function<void(size_t, const common_chat_schema_property &, const common_chat_schema_document_ptr &)> & fn) {
     auto         params = common_chat_tool_parameters(function);
     auto         doc    = std::make_shared<const common_chat_schema_document>(common_chat_schema_from_json(params));
     const auto * object = dynamic_cast<const common_chat_schema_object *>(doc->root.get());
     if (!object) {
         return;
     }
-    for (const auto & prop : object->properties) {
-        fn(prop, doc);
+    for (size_t i = 0; i < object->properties.size(); i++) {
+        fn(i, object->properties[i], doc);
     }
 }

common/parsers/parsers.h

@@ -17,11 +17,11 @@
 
 using json = common_json;
 
-// iterate over the function tools of an OpenAI-style tools array
-void foreach_function(const json & tools, const std::function<void(const json &)> & fn);
+// iterate over the function tools of an OpenAI-style tools array, passing each tool with its index in the array
+void foreach_function(const json & tools, const std::function<void(size_t, const json &)> & fn);
 
-// iterate over the parameters of a function tool, with the document that owns them
-void foreach_parameter(const json & function, const std::function<void(const common_chat_schema_property &, const common_chat_schema_document_ptr &)> & fn);
+// iterate over the parameters of a function tool, passing each parameter with its index and the document that owns it
+void foreach_parameter(const json & function, const std::function<void(size_t, const common_chat_schema_property &, const common_chat_schema_document_ptr &)> & fn);
 
 // render a template; the override arguments let a parser feed in messages, tools or context it has rewritten
 std::string common_chat_template_direct_apply_impl(
@@ -81,3 +81,5 @@ common_chat_params common_chat_params_init_ministral_3(const common_chat_templat
 common_chat_params common_chat_params_init_muse_glimmer(const common_chat_template & tmpl, const autoparser::generation_params & inputs);
 
 common_chat_params common_chat_params_init_qwen3_coder(const common_chat_template & tmpl, const autoparser::generation_params & inputs);
+
+common_chat_params common_chat_params_init_translate_gemma(const common_chat_template & tmpl, const autoparser::generation_params & inputs);

common/parsers/qwen3-coder.cpp

@@ -65,7 +65,7 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat
         // Match complete <function=name> opener for Qwen3-Coder models that occasionally omit the
         // starting <tool_call>. The model may hallucinate a tool name, but it is preferable over
         // constraining on <function which may occur in valid content generation, e.g. #include <functional>
-        foreach_function(inputs.tools, [&](const json & tool) {
+        foreach_function(inputs.tools, [&](size_t, const json & tool) {
             const std::string name = tool.at("function").at("name");
             tool_call_starts.push_back("<function=" + name + ">");
         });
@@ -93,15 +93,15 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat
                 p.ac(p.tool_arg_string_value(p.until("\n</parameter>\n")) + arg_close, "\n</parameter>\n"));
 
             auto tool_choice = p.choice();
-            foreach_function(inputs.tools, [&](const json & tool) {
+            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
 
                 std::vector<common_peg_parser> required_args;
                 std::vector<common_peg_parser> optional_args;
 
-                foreach_parameter(function, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
-                    auto rule_name = "tool-" + name + "-arg-" + param.name;
+                foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
+                    auto rule_name = "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index);
 
                     auto arg_open = p.tool_arg_open("<parameter=" + p.tool_arg_name(p.literal(param.name)) + ">\n");
 
@@ -141,7 +141,7 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat
 
                 // Accept required arguments in any order, as Qwen does not always adhere to the
                 // order provided.
-                auto args = p.permute("tool-" + name + "-args", required_args);
+                auto args = p.permute("tool-" + std::to_string(tool_index) + "-args", required_args);
                 if (!optional_args.empty()) {
                     args = args + p.zero_or_more(p.choice(optional_args));
                 }
@@ -150,7 +150,7 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat
                                    p.tool_args(args) +
                                    p.tool_close(p.literal("</function>\n")));
 
-                tool_choice |= p.rule("tool-" + name, func);
+                tool_choice |= p.rule("tool-" + std::to_string(tool_index), func);
             });
 
             auto min_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0;

common/parsers/sources.cmake

@@ -20,4 +20,5 @@ set(LLAMA_CHAT_PARSERS_SOURCES
     ${CMAKE_CURRENT_LIST_DIR}/ministral3.cpp
     ${CMAKE_CURRENT_LIST_DIR}/muse-glimmer.cpp
     ${CMAKE_CURRENT_LIST_DIR}/qwen3-coder.cpp
+    ${CMAKE_CURRENT_LIST_DIR}/translate-gemma.cpp
 )

common/parsers/translate-gemma.cpp

@@ -0,0 +1,63 @@
+#include "parsers.h"
+#include "log.h"
+
+// TranslateGemma does not support tools or reasoning, it only needs user messages in its own content schema
+common_chat_params common_chat_params_init_translate_gemma(
+        const common_chat_template & tmpl,
+        const autoparser::generation_params & inputs) {
+
+    common_chat_params data;
+
+    // default to chat_template_kwargs, or en-GB if not specified
+    std::string src_lang = inputs.extra_context.value("source_lang_code", "en-GB");
+    std::string tgt_lang = inputs.extra_context.value("target_lang_code", "en-GB");
+    for (const char * key : { "source_lang_code", "target_lang_code" }) {
+        if (!inputs.extra_context.contains(key)) {
+            LOG_WRN("TranslateGemma: %s not set in chat_template_kwargs, defaulting to en-GB\n", key);
+        }
+    }
+
+    json messages = inputs.messages;
+    for (auto & message : messages) {
+        if (message.value("role", "") != "user") {
+            continue;
+        }
+        std::string text;
+        const auto & content = message.contains("content") ? message.at("content") : json();
+        if (content.is_string()) {
+            text = content.get<std::string>();
+        } else if (content.is_array()) {
+            for (const auto & part : content) {
+                if (!text.empty()) {
+                    text += "\n";
+                }
+                text += part.value("text", "");
+            }
+        }
+        message["content"] = json::array({
+            json{
+                {"type", "text"},
+                {"text", text},
+                {"source_lang_code", src_lang},
+                {"target_lang_code", tgt_lang},
+            }
+        });
+    }
+
+    data.prompt            = common_chat_template_direct_apply_impl(tmpl, inputs, messages);
+    data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, messages);
+    data.format            = COMMON_CHAT_FORMAT_PEG_NATIVE;
+    data.supports_thinking = false;
+
+    if (inputs.has_continuation()) {
+        data.generation_prompt = "<start_of_turn>model\n" + inputs.continue_msg.render_content();
+        data.prompt += data.generation_prompt;
+    }
+
+    auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
+        return p.literal(data.generation_prompt) << p.content(p.rest());
+    });
+    data.parser = parser.save();
+
+    return data;
+}

common/peg-parser.h

@@ -2,6 +2,7 @@
 
 #include "json-schema.h"
 #include "json.h"
+#include "llama.h"
 
 #include <memory>
 #include <set>
@@ -182,7 +183,8 @@ inline common_peg_parse_flags operator~(common_peg_parse_flags a) {
 }
 
 struct common_peg_parse_context {
-    std::string input;
+    std::string input;               // [h,   e,  l,  l,  o,  _,  w,  o,  r,  l,  d]
+    std::vector<llama_token> tokens; // [id, -1, -1, -1, -1, id, -1, -1, -1, -1, -1]
     common_peg_parse_flags flags;
     common_peg_ast_arena ast;
 
@@ -194,6 +196,11 @@ struct common_peg_parse_context {
     common_peg_parse_context(const std::string & input, common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_NONE)
         : input(input), flags(flags), parse_depth(0) {}
 
+    common_peg_parse_context(std::string input, std::vector<llama_token> tokens, common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_NONE)
+        : input(std::move(input)), tokens(std::move(tokens)), flags(flags), parse_depth(0) {
+        GGML_ASSERT(this->tokens.empty() || this->tokens.size() == this->input.size());
+    }
+
     bool is_lenient() const { return flags & COMMON_PEG_PARSE_FLAG_LENIENT; }
     bool is_debug() const { return flags & COMMON_PEG_PARSE_FLAG_DEBUG; }
 };

common/sampling.cpp

@@ -399,8 +399,11 @@ struct common_sampler * common_sampler_init(
             // only if user explicitly included adaptive-p sampler
             samplers.push_back(llama_sampler_init_adaptive_p(params.adaptive_target, params.adaptive_decay, params.seed));
         } else {
-            // default: sample from distribution
-            samplers.push_back(llama_sampler_init_dist(params.seed));
+            // Keep distribution sampling when callers request probabilities.
+            const bool greedy = params.n_probs == 0 && !params.samplers.empty() &&
+                ((params.samplers.back() == COMMON_SAMPLER_TYPE_TEMPERATURE && params.temp == 0.0f && params.dynatemp_range == 0.0f) ||
+                 (params.samplers.back() == COMMON_SAMPLER_TYPE_TOP_K && params.top_k == 1));
+            samplers.push_back(greedy ? llama_sampler_init_greedy() : llama_sampler_init_dist(params.seed));
         }
     } else if (params.mirostat == 1) {
         samplers.push_back(llama_sampler_init_temp(params.temp));

common/speculative.cpp

@@ -2561,6 +2561,9 @@ common_params common_base_params_to_speculative(const common_params & params) {
     result.n_outputs_max = params.n_parallel;
     result.n_outputs_max_per_seq = 1;
 
+    // the MoE cache is only used by the target context
+    result.moe_cache_size = 0;
+
     // dflash/dspark decode the whole noise block in a single pass and sample every block position on the backend
     // TODO: refactor such properties to be announced by the speculative types
     //       something like `struct common_speculative_type_props common_speculative_type_get_props(...);`

conversion/__init__.py

@@ -161,6 +161,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "Lfm2BidirectionalForMaskedLM": "lfm2",
     "Lfm2BidirectionalModel": "lfm2",
     "Lfm2ForCausalLM": "lfm2",
+    "D1Model": "lfm2",
     "Lfm2Model": "lfm2",
     "Lfm2MoeForCausalLM": "lfm2",
     "Llama4ForCausalLM": "llama",
@@ -303,6 +304,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
     "AudioFlamingo3ForConditionalGeneration": "ultravox",
     "ClefModel": "clef",
     "CogVLMForCausalLM": "cogvlm",
+    "Cohere2VisionForConditionalGeneration": "command_r",
     "PplxDeciderModel": "pplx_decider",
     "DeepseekOCR2ForCausalLM": "deepseek",
     "DeepseekOCRForCausalLM": "deepseek",

conversion/base.py

@@ -2499,7 +2499,11 @@ class TextModel(ModelBase):
         if template is not None:
             self.gguf_writer.add_chat_template(template)
 
-    def _set_vocab_plamo(self):
+    def _set_vocab_plamo(
+        self,
+        eot_token: str,
+        normal_tokens: Iterable[str] = (),
+    ):
         # PLaMo models use a custom tokenizer with a .jsonl file
         tokenizer_jsonl_path = self.dir_model / "tokenizer.jsonl"
         tokenizer_config_path = self.dir_model / "tokenizer_config.json"
@@ -2523,27 +2527,30 @@ class TextModel(ModelBase):
         tokens = []
         scores = []
         toktypes = []
+        normal_tokens = set(normal_tokens)
 
         with open(tokenizer_jsonl_path, "r", encoding="utf-8") as f:
             for line_num, line in enumerate(f):
                 if line.strip():
                     token_data = json.loads(line)
                     # Format: [token, score, type, ?, ?, ?, ?]
-                    token = token_data[0].encode("utf-8")
+                    token_str = token_data[0]
+                    token = token_str.encode("utf-8")
                     score = float(token_data[1])
                     token_type_str = token_data[2] if len(token_data) > 2 else "NORMAL"
 
                     tokens.append(token)
                     scores.append(score)
 
-                    if token_type_str == "UNKNOWN":
+                    if token_str in normal_tokens:
+                        toktypes.append(gguf.TokenType.NORMAL)
+                    elif token_type_str == "UNKNOWN":
                         toktypes.append(gguf.TokenType.UNKNOWN)
                     elif token_type_str == "CONTROL":
                         toktypes.append(gguf.TokenType.CONTROL)
                     elif token_type_str == "BYTE":
                         toktypes.append(gguf.TokenType.BYTE)
                     else:
-                        token_str = token_data[0]
                         if token_str.startswith("<|plamo:") and token_str.endswith("|>"):
                             toktypes.append(gguf.TokenType.CONTROL)
                         else:
@@ -2580,8 +2587,7 @@ class TextModel(ModelBase):
             token_id = tokens.index(tokenizer_config["unk_token"].encode("utf-8"))
             self.gguf_writer.add_unk_token_id(token_id)
 
-        # Add <|plamo:op|> as EOT to ensure appropriate end of generation
-        self.gguf_writer.add_eot_token_id(4)
+        self.gguf_writer.add_eot_token_id(tokens.index(eot_token.encode("utf-8")))
 
         self.gguf_writer.add_add_space_prefix(False)
 

conversion/command_r.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 
 import re
-from typing import Iterable, TYPE_CHECKING
+from typing import Callable, Iterable, TYPE_CHECKING
 
 import torch
 
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, TextModel, gguf, logger
+from .base import MmprojModel, ModelBase, TextModel, gguf, logger
 
 
 @ModelBase.register("CohereForCausalLM")
@@ -180,3 +180,28 @@ class Cohere2MoeModel(TextModel):
         experts = [k for d in self._experts for k in d.keys()]
         if len(experts) > 0:
             raise ValueError(f"Unprocessed experts: {experts}")
+
+
+@ModelBase.register("Cohere2VisionForConditionalGeneration")
+# [TAG_HF_EXAMPLE_GATED] CohereLabs/command-a-vision-07-2025 is gated
+@ModelBase.example("CohereLabs/command-a-plus-05-2026-bf16")
+class Cohere2VisionModel(MmprojModel):
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.COHERE2V)
+        self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams["layer_norm_eps"])
+        self.gguf_writer.add_vision_projector_scale_factor(self.global_config["downsample_factor"])
+        self.gguf_writer.add_vision_preproc_max_tiles(self.preprocessor_config["max_patches"])
+        self.gguf_writer.add_vision_use_gelu(True)
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        if ".embeddings." in 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 not name.startswith(("model.vision_tower.", "model.multi_modal_projector.")):
+            return None
+        return super().filter_tensors((name, gen))

conversion/lfm2.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+import json
+
+from pathlib import Path
 from typing import Any, Callable, Iterable, TYPE_CHECKING
 
 import torch
@@ -7,7 +10,7 @@ import torch
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import MmprojModel, ModelBase, TextModel, gguf
+from .base import MmprojModel, ModelBase, TextModel, gguf, jinja_str_or_json, logger
 
 from .gemma import ConformerAudioModel
 
@@ -65,6 +68,68 @@ class LFM2Model(TextModel):
         yield from super().modify_tensors(data_torch, name, bid)
 
 
+def _is_d1_checkpoint(dir_model: Path) -> bool:
+    if not (dir_model / "config.json").is_file():
+        return False
+    with open(dir_model / "config.json", encoding="utf-8") as f:
+        return json.load(f).get("auto_map", {}).get("AutoModel", "").endswith(".D1Model")
+
+
+@ModelBase.register_hparams_loader(_is_d1_checkpoint)
+def _load_d1_hparams(dir_model: Path) -> dict[str, Any]:
+    logger.info("gguf: detected d1 checkpoint")
+    hparams = ModelBase.load_hparams(dir_model, False, guess=False)
+    # the mmproj stays LFM2-VL
+    hparams["text_config"]["architectures"] = ["D1Model"]
+    return hparams
+
+
+@ModelBase.register("D1Model")
+@ModelBase.example("LiquidAI/d1-3b")
+class D1Model(LFM2Model):
+    model_arch = gguf.MODEL_ARCH.LFM2
+
+    def set_vocab(self):
+        super().set_vocab()
+        self.gguf_writer.add_chat_template([{"name": "systemone", "template": self._systemone_template()}])
+
+    @staticmethod
+    def _systemone_template() -> str:
+        # follows prompt.py of the model repo
+        description = jinja_str_or_json("o.description")
+        choice = (
+            "{{ '\\n\\nOptions:\\n' }}"
+            "{% for o in options %}{{ o.label }} {% if o.description %}" + description + "{% else %}{{ o.key | replace('_', ' ') }}{% endif %}"
+            "{% if not loop.last %}{{ '\\n' }}{% endif %}{% endfor %}"
+            "{{ '\\n\\nReply with the option code only.' }}"
+        )
+        # with criteria, a missing description is written as None
+        noul = (
+            "{% set ns = namespace(criteria=false) %}{% for o in options %}{% if o.description is not none %}{% set ns.criteria = true %}{% endif %}{% endfor %}"
+            "{% if ns.criteria %}"
+            "{% for o in options %}{{ '\\nYes: ' if o.key == 'true' else '\\nNo: ' }}"
+            "{% if o.description is none %}None{% else %}" + description + "{% endif %}{% endfor %}{% endif %}"
+            "{{ '\\n\\nReply with yes or no only.' }}"
+        )
+        score = (
+            "{{ '\\n\\n' }}{% for o in options %}{{ o.key }} " + description + "{{ '\\n' }}{% endfor %}"
+            "{{ '\\nReply with a single digit 0-' }}{{ options | length - 1 }}{{ ' only.' }}"
+        )
+        return (
+            "<|startoftext|><|im_start|>user\n"
+            "{% for image in images %}{{ image }}{% endfor %}"
+            "{% if state is not none %}{% if state is string %}{{ state }}{% else %}{{ state | tojson(indent=2) }}{% endif %}"
+            "{{ '\\n\\n\\nQUESTION:\\n' }}{% endif %}"
+            + jinja_str_or_json("instructions")
+            + "{% if type == 'choice' %}" + choice + "{% elif type == 'noul' %}" + noul + "{% else %}" + score + "{% endif %}"
+            "{{ '<|im_end|>\\n<|im_start|>assistant\\n' }}"
+        )
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_decision_type(gguf.DecisionType.LFM2_D1)
+
+
 @ModelBase.register("Lfm2Model", "Lfm2BidirectionalModel", "Lfm2BidirectionalForMaskedLM")
 @ModelBase.example("LiquidAI/LFM2.5-ColBERT-350M", "LiquidAI/LFM2.5-Embedding-350M", "LiquidAI/LFM2.5-Encoder-350M", "LiquidAI/LFM2.5-Encoder-230M")
 class LFM2ColBertModel(LFM2Model):
@@ -188,6 +253,12 @@ class LFM2VLModel(MmprojModel):
         # python notation, e.g. for vision_feature_layer == -1, we pick last layer -> vision_feature_layers_to_drop = 0
         vision_feature_layers_to_drop = -(self.global_config.get("vision_feature_layer", -1) + 1)
         self.gguf_writer.add_vision_block_count(self.find_vparam(self.n_block_keys) - vision_feature_layers_to_drop)
+        # PIL resample enum
+        if (resample := self.preprocessor_config.get("resample")) is not None:
+            resize_algo = {1: "lanczos", 2: "bilinear", 3: "bicubic"}.get(resample)
+            if resize_algo is None:
+                raise ValueError(f"unsupported resample: {resample}")
+            self.gguf_writer.add_vision_image_resize_algo(resize_algo)
 
     @classmethod
     def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:

conversion/plamo.py

@@ -64,7 +64,7 @@ class Plamo2Model(TextModel):
     model_arch = gguf.MODEL_ARCH.PLAMO2
 
     def set_vocab(self):
-        self._set_vocab_plamo()
+        self._set_vocab_plamo(eot_token="<|plamo:op|>")
 
     def set_gguf_parameters(self):
         hparams = self.hparams
@@ -170,7 +170,10 @@ class Plamo3Model(TextModel):
             })
 
     def set_vocab(self):
-        self._set_vocab_plamo()
+        self._set_vocab_plamo(
+            eot_token="<|plamo:tag|>",
+            normal_tokens=("<|plamo:begin_", "<|plamo:end_", ":plamo|>"),
+        )
 
         tokenizer_config_path = self.dir_model / "tokenizer_config.json"
         tokenizer_config = {}

conversion/qwen3tts.py

@@ -218,8 +218,10 @@ class Qwen3TTSSpeakerEncoderModel(MmprojModel):
         if hparams is None:
             hparams = ModelBase.load_hparams(dir_model, is_mistral_format=False)
         hparams["text_config"] = {"hidden_size": hparams["talker_config"]["hidden_size"]}
-        # ECAPA-TDNN has a fixed 4-stage backbone, but MmprojModel.__init__ needs a n_block_keys
-        hparams["speaker_encoder_config"]["n_layers"] = 4
+        # ECAPA-TDNN has a fixed 4-stage backbone, but MmprojModel.__init__ needs a n_block_keys.
+        # The CustomVoice variant ships no speaker encoder, so its config lacks this key entirely.
+        if "speaker_encoder_config" in hparams:
+            hparams["speaker_encoder_config"]["n_layers"] = 4
         super().__init__(dir_model, *args, hparams=hparams, **kwargs)
         self._wav_config_cache = None
 

docs/development/parsing.md

@@ -25,15 +25,16 @@ output from a model that emits arguments as JSON.
 auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
     // Build a choice of all available tools
     auto tool_choice = p.choice();
-    for (const auto & tool : tools) {
+    for (size_t i = 0; i < tools.size(); i++) {
+        const auto & tool = tools[i];
         const auto & function = tool.at("function");
         std::string name = function.at("name");
         const auto   schema = common_chat_tool_parameters(function);
 
         auto tool_name = p.json_member("name", "\"" + p.literal(name) + "\"");
-        auto tool_args = p.json_member("arguments", p.schema(p.json(), "tool-" + name + "-schema", schema));
+        auto tool_args = p.json_member("arguments", p.schema(p.json(), "tool-" + std::to_string(i) + "-schema", schema));
 
-        tool_choice |= p.rule("tool-" + name, "{" << tool_name << "," << tool_args << "}");
+        tool_choice |= p.rule("tool-" + std::to_string(i), "{" << tool_name << "," << tool_args << "}");
     }
 
     // Define the tool call structure: <tool_call>[{tool}]</tool_call>

ggml/src/ggml-cuda/fwht.cu

@@ -2,6 +2,9 @@
 #include "convert.cuh"
 #include "fwht.cuh"
 
+// wide FWHT blocks use one row per thread block with this many threads
+#define GGML_CUDA_FWHT_BLOCK_NT 256
+
 template <int N, typename T>
 __launch_bounds__(4*ggml_cuda_get_physical_warp_size(), 1)
 __global__ void fwht_cuda(const T * src, float * dst, const int64_t n_rows, const float scale) {
@@ -59,6 +62,87 @@ __global__ void fwht_cuda(const T * src, float * dst, const int64_t n_rows, cons
     }
 }
 
+// Wide blocks: one row per thread block instead of per warp, so each thread keeps N/NT
+// values rather than N/32. Stages below the warp width still shuffle, those up to the
+// block width go through shared memory, and the rest stay in registers.
+template <int N, int NT, typename T>
+__launch_bounds__(NT, 1)
+__global__ void fwht_cuda_block(const T * src, float * dst, const int64_t n_rows, const float scale) {
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    constexpr int NE        = N / NT;
+    static_assert(NE >= 1 && N % NT == 0 && NT % warp_size == 0, "bad FWHT block shape");
+
+    __shared__ float s[N];
+
+    const int64_t r = blockIdx.x;
+    if (r >= n_rows) {
+        return;
+    }
+
+    src += r * N;
+    dst += r * N;
+
+    const int tid  = threadIdx.x;
+    const int lane = tid % warp_size;
+
+    ggml_cuda_pdl_sync();
+
+    float reg[NE];
+#pragma unroll
+    for (int i = 0; i < NE; ++i) {
+        reg[i] = ggml_cuda_cast<float>(src[i * NT + tid]) * scale;
+    }
+
+    // stages within a warp: partner differs in the lane bits
+#pragma unroll
+    for (int h = 1; h < warp_size; h *= 2) {
+#pragma unroll
+        for (int j = 0; j < NE; j++) {
+            const float val  = reg[j];
+            const float val2 = __shfl_xor_sync(0xFFFFFFFF, val, h, warp_size);
+            reg[j] = (lane & h) == 0 ? val + val2 : val2 - val;
+        }
+    }
+
+    // stages across warps: partner differs in the thread-index bits above the lane
+#pragma unroll
+    for (int h = warp_size; h < NT; h *= 2) {
+#pragma unroll
+        for (int j = 0; j < NE; j++) {
+            s[j * NT + tid] = reg[j];
+        }
+        __syncthreads();
+#pragma unroll
+        for (int j = 0; j < NE; j++) {
+            const float val  = reg[j];
+            const float val2 = s[j * NT + (tid ^ h)];
+            reg[j] = (tid & h) == 0 ? val + val2 : val2 - val;
+        }
+        __syncthreads();
+    }
+
+    // stages above the block width: partner is another register of the same thread
+#pragma unroll
+    for (int h = NT; h < N; h *= 2) {
+        const int step = h / NT;
+#pragma unroll
+        for (int j = 0; j < NE; j += 2 * step) {
+#pragma unroll
+            for (int k = 0; k < step; k++) {
+                const float x = reg[j + k];
+                const float y = reg[j + k + step];
+                reg[j + k]        = x + y;
+                reg[j + k + step] = x - y;
+            }
+        }
+    }
+
+#pragma unroll
+    for (int i = 0; i < NE; ++i) {
+        dst[i * NT + tid] = reg[i];
+    }
+}
+
 template <typename T>
 static bool ggml_cuda_op_fwht_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst) {
     const int     n    = src->ne[0];
@@ -94,7 +178,34 @@ static bool ggml_cuda_op_fwht_impl(ggml_backend_cuda_context & ctx, const ggml_t
             ggml_cuda_kernel_launch(fwht_cuda<512, T>, launch_params, src_d, dst_d, rows, scale);
             return true;
         default:
-            return false;
+            break;
+    }
+
+    // wide blocks: one row per thread block
+    {
+        constexpr int nt = GGML_CUDA_FWHT_BLOCK_NT;
+
+        dim3 grid_dims_w(rows, 1, 1);
+        dim3 block_dims_w(nt, 1, 1);
+        const ggml_cuda_kernel_launch_params launch_params_w =
+            ggml_cuda_kernel_launch_params(grid_dims_w, block_dims_w, 0, stream);
+
+        switch (n) {
+            case 1024:
+                ggml_cuda_kernel_launch(fwht_cuda_block<1024, nt, T>, launch_params_w, src_d, dst_d, rows, scale);
+                return true;
+            case 2048:
+                ggml_cuda_kernel_launch(fwht_cuda_block<2048, nt, T>, launch_params_w, src_d, dst_d, rows, scale);
+                return true;
+            case 4096:
+                ggml_cuda_kernel_launch(fwht_cuda_block<4096, nt, T>, launch_params_w, src_d, dst_d, rows, scale);
+                return true;
+            case 8192:
+                ggml_cuda_kernel_launch(fwht_cuda_block<8192, nt, T>, launch_params_w, src_d, dst_d, rows, scale);
+                return true;
+            default:
+                return false;
+        }
     }
 }
 

ggml/src/ggml-cuda/lightning-indexer.cu

@@ -239,6 +239,14 @@ static __global__ void lightning_indexer_kernel_wmma(
 // tokens scored per block by the tile kernel
 #define LIGHTNING_INDEXER_TILE_TOKENS 8
 
+// heads whose queries the tile kernel stages per pass, MUSA arch 21 caps static shared memory
+// at 28 KB and the queries of four heads do not fit there next to the key tile
+#if defined(GGML_USE_MUSA) && defined(__MUSA_ARCH__) && __MUSA_ARCH__ < 220
+#define LIGHTNING_INDEXER_TILE_HEADS_PER_PASS 2
+#else
+#define LIGHTNING_INDEXER_TILE_HEADS_PER_PASS 4
+#endif
+
 // TODO there is one ugly assumption used in this kernel - that WARP_SIZE is equal to 32
 // thanks to that one warp operating on float4 processes whole indexer K/Q vectors
 // 32 * 4 = 128 (N_EMBD)
@@ -406,9 +414,11 @@ static __global__ void lightning_indexer_kernel_tile(
     constexpr int KEY_LANES         = THREADS_PER_BLOCK / TOKENS_PER_BLOCK;
     constexpr int KEYS_PER_THREAD   = K_VECS_PER_BLOCK / KEY_LANES;
     constexpr int N_EMBD_H2         = N_EMBD / 2;
+    constexpr int HEADS_PER_PASS    = N_HEAD < LIGHTNING_INDEXER_TILE_HEADS_PER_PASS ? N_HEAD : LIGHTNING_INDEXER_TILE_HEADS_PER_PASS;
 
     static_assert(THREADS_PER_BLOCK % TOKENS_PER_BLOCK == 0, "threads must cover the token tile");
     static_assert(K_VECS_PER_BLOCK % KEY_LANES == 0, "key lanes must cover the key tile");
+    static_assert(N_HEAD % HEADS_PER_PASS == 0, "head passes must cover the heads");
 
     const int tid         = threadIdx.y * WARP_SIZE + threadIdx.x;
     const int start_kv    = blockIdx.x * K_VECS_PER_BLOCK;
@@ -417,7 +427,7 @@ static __global__ void lightning_indexer_kernel_tile(
 
     // the row padding keeps the keys of consecutive threads in distinct banks
     __shared__ half2 k_shared[K_VECS_PER_BLOCK][N_EMBD_H2 + 1];
-    __shared__ float2 q_shared[N_HEAD][TOKENS_PER_BLOCK][N_EMBD_H2];
+    __shared__ float2 q_shared[HEADS_PER_PASS][TOKENS_PER_BLOCK][N_EMBD_H2];
     __shared__ float w_shared[N_HEAD][TOKENS_PER_BLOCK];
 
     // phase 1 - stage the key tile four elements at a time, rows past n_kv are zero
@@ -451,22 +461,7 @@ static __global__ void lightning_indexer_kernel_tile(
         k_shared[r][2*c4 + 1] = hi;
     }
 
-    // phase 2 - stage the queries and weights of every head, tokens past n_batch are zero
-
-#pragma unroll
-    for (int i = tid; i < N_HEAD * TOKENS_PER_BLOCK * (N_EMBD / 4); i += THREADS_PER_BLOCK) {
-        const int h  = i / (TOKENS_PER_BLOCK * (N_EMBD / 4));
-        const int r  = i / (N_EMBD / 4) % TOKENS_PER_BLOCK;
-        const int c4 = i % (N_EMBD / 4);
-
-        float4 v = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
-        if (start_batch + r < n_batch) {
-            v = *(const float4 *) ((const char *) Q + h*nbq1 + (start_batch + r)*nbq2 + i_stream*nbq3 + c4*sizeof(float4));
-        }
-
-        q_shared[h][r][2*c4 + 0] = make_float2(v.x, v.y);
-        q_shared[h][r][2*c4 + 1] = make_float2(v.z, v.w);
-    }
+    // phase 2 - stage the weights of every head, tokens past n_batch are zero
 
     if (tid < N_HEAD * TOKENS_PER_BLOCK) {
         const int h = tid / TOKENS_PER_BLOCK;
@@ -475,33 +470,60 @@ static __global__ void lightning_indexer_kernel_tile(
             ((const float *) ((const char *) W + (start_batch + r)*nbw1 + i_stream*nbw3))[h] : 0.0f;
     }
 
-    __syncthreads();
-
-    // phase 3 - float products of the widened keys for every head, ReLU, weight
-
     const int kl = tid % KEY_LANES;
     const int tl = tid / KEY_LANES;
 
     float qk[N_HEAD][KEYS_PER_THREAD] = { { 0.0f } };
 
-#pragma unroll 8
-    for (int c = 0; c < N_EMBD_H2; ++c) {
-        float2 k_val[KEYS_PER_THREAD];
 #pragma unroll
-        for (int j = 0; j < KEYS_PER_THREAD; ++j) {
-            k_val[j] = __half22float2(k_shared[kl + j*KEY_LANES][c]);
+    for (int h0 = 0; h0 < N_HEAD; h0 += HEADS_PER_PASS) {
+        // the previous pass is fully consumed before its queries are replaced
+        if (h0 > 0) {
+            __syncthreads();
         }
+
+        // phase 3 - stage the queries of the heads of this pass, tokens past n_batch are zero
+
 #pragma unroll
-        for (int h = 0; h < N_HEAD; ++h) {
-            const float2 q_val = q_shared[h][tl][c];
+        for (int i = tid; i < HEADS_PER_PASS * TOKENS_PER_BLOCK * (N_EMBD / 4); i += THREADS_PER_BLOCK) {
+            const int h  = i / (TOKENS_PER_BLOCK * (N_EMBD / 4));
+            const int r  = i / (N_EMBD / 4) % TOKENS_PER_BLOCK;
+            const int c4 = i % (N_EMBD / 4);
+
+            float4 v = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
+            if (start_batch + r < n_batch) {
+                v = *(const float4 *) ((const char *) Q + (h0 + h)*nbq1 + (start_batch + r)*nbq2 + i_stream*nbq3 + c4*sizeof(float4));
+            }
+
+            q_shared[h][r][2*c4 + 0] = make_float2(v.x, v.y);
+            q_shared[h][r][2*c4 + 1] = make_float2(v.z, v.w);
+        }
+
+        __syncthreads();
+
+        // phase 4 - float products of the widened keys for the heads of this pass
+
+#pragma unroll 8
+        for (int c = 0; c < N_EMBD_H2; ++c) {
+            float2 k_val[KEYS_PER_THREAD];
 #pragma unroll
             for (int j = 0; j < KEYS_PER_THREAD; ++j) {
-                qk[h][j] = fmaf(k_val[j].x, q_val.x, qk[h][j]);
-                qk[h][j] = fmaf(k_val[j].y, q_val.y, qk[h][j]);
+                k_val[j] = __half22float2(k_shared[kl + j*KEY_LANES][c]);
+            }
+#pragma unroll
+            for (int h = 0; h < HEADS_PER_PASS; ++h) {
+                const float2 q_val = q_shared[h][tl][c];
+#pragma unroll
+                for (int j = 0; j < KEYS_PER_THREAD; ++j) {
+                    qk[h0 + h][j] = fmaf(k_val[j].x, q_val.x, qk[h0 + h][j]);
+                    qk[h0 + h][j] = fmaf(k_val[j].y, q_val.y, qk[h0 + h][j]);
+                }
             }
         }
     }
 
+    // phase 5 - ReLU, weight, add the mask and write, consecutive threads write consecutive keys
+
     float score[KEYS_PER_THREAD] = { 0.0f };
 
 #pragma unroll
@@ -512,8 +534,6 @@ static __global__ void lightning_indexer_kernel_tile(
         }
     }
 
-    // phase 4 - add the mask and write, consecutive threads write consecutive keys
-
     const int i_batch = start_batch + tl;
     if (i_batch >= n_batch) {
         return;
@@ -677,8 +697,6 @@ void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor *
             LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_F32)
             GGML_ABORT("fatal error");
         }
-#ifndef GGML_USE_MUSA
-    // MUSA archs 21 and 22 cap static shared memory at 28 KB, below what the tile kernel stages
     } else if (n_embd == 128 && n_head == 4 && n_batch >= LIGHTNING_INDEXER_TILE_TOKENS) {
         // too few heads for a wmma tile, the tile kernel shares the keys across the tokens
         constexpr int WARPS_PER_BLOCK = 8;
@@ -698,9 +716,8 @@ void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor *
         LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_BF16)
         LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_F32)
         GGML_ABORT("fatal error");
-#endif // GGML_USE_MUSA
     } else if (n_embd == 128 && n_head == 4) {
-        // a batch smaller than a token tile, or MUSA, use vector kernel
+        // a batch smaller than a token tile, use vector kernel
         constexpr int K_VECS_PER_WARP = 8;
         constexpr int WARPS_PER_BLOCK = 8;
         constexpr int K_VECS_PER_BLOCK = K_VECS_PER_WARP * WARPS_PER_BLOCK;

ggml/src/ggml-cuda/pool2d.cu

@@ -50,6 +50,65 @@ static  __global__ void pool2d_nchw_kernel(
     o_ptr[cur_oh * ow + cur_ow] = res;
 }
 
+template <typename Ti, typename To>
+static __global__ void pool2d_nchw_kernel_warp(
+        const int ih, const int iw, const int oh, const int ow,
+        const int kh, const int kw, const int sh, const int sw,
+        const int ph, const int pw, const int parallel_elements,
+        const Ti * __restrict__ src, To * __restrict__ dst, const enum ggml_op_pool op) {
+    const int warp_id = (threadIdx.x + blockIdx.x * blockDim.x) / WARP_SIZE;
+    const int lane     = threadIdx.x % WARP_SIZE;
+    if (warp_id >= parallel_elements) {
+        return;
+    }
+
+    const int I_HW = ih * iw;
+    const int O_HW = oh * ow;
+    const int nc     = warp_id / O_HW;
+    const int cur_oh = warp_id % O_HW / ow;
+    const int cur_ow = warp_id % O_HW % ow;
+    const Ti* i_ptr = src + nc * I_HW;
+
+    const int start_h = cur_oh * sh - ph;
+    const int bh = max(0, start_h);
+    const int eh = min(ih, start_h + kh);
+    const int start_w = cur_ow * sw - pw;
+    const int bw = max(0, start_w);
+    const int ew = min(iw, start_w + kw);
+
+    const int win_w     = ew - bw;
+    const int win_elems = (eh - bh) * win_w;
+    const To scale = 1. / (kh * kw);
+
+    To res;
+    switch (op) {
+        case GGML_OP_POOL_AVG: res = 0; break;
+        case GGML_OP_POOL_MAX: res = -FLT_MAX; break;
+        default: res = 0; assert(false);
+    }
+
+    for (int t = lane; t < win_elems; t += WARP_SIZE) {
+        const int i = bh + t / win_w;
+        const int j = bw + t % win_w;
+        const Ti cur = i_ptr[i * iw + j];
+        switch (op) {
+            case GGML_OP_POOL_AVG: res += cur * scale; break;
+            case GGML_OP_POOL_MAX: res = max(res, (To)cur); break;
+            default: break;
+        }
+    }
+
+#pragma unroll
+    for (int offset = WARP_SIZE/2; offset > 0; offset >>= 1) {
+        const To other = __shfl_xor_sync(0xFFFFFFFF, res, offset, WARP_SIZE);
+        res = (op == GGML_OP_POOL_MAX) ? max(res, other) : res + other;
+    }
+
+    if (lane == 0) {
+        dst[nc * O_HW + cur_oh * ow + cur_ow] = res;
+    }
+}
+
 static void pool2d_nchw_kernel_f32_f32_cuda(
         const int ih, const int iw, const int oh, const int ow,
         const int kh, const int kw, const int sh, const int sw,
@@ -57,6 +116,13 @@ static void pool2d_nchw_kernel_f32_f32_cuda(
         const float * src, float * dst, const enum ggml_op_pool op,
         cudaStream_t stream) {
 
+    if (kh * kw >= CUDA_POOL2D_WARP_KERNEL_MIN_WINDOW) {
+        const int warps_per_block = CUDA_POOL2D_BLOCK_SIZE / WARP_SIZE;
+        const int num_blocks = (parallel_elements + warps_per_block - 1) / warps_per_block;
+        pool2d_nchw_kernel_warp<<<num_blocks, CUDA_POOL2D_BLOCK_SIZE, 0, stream>>>(ih, iw, oh, ow, kh, kw, sh, sw, ph, pw, parallel_elements, src, dst, op);
+        return;
+    }
+
     const int num_blocks = (parallel_elements + CUDA_POOL2D_BLOCK_SIZE - 1) / CUDA_POOL2D_BLOCK_SIZE;
     dim3 block_nums(num_blocks);
     pool2d_nchw_kernel<<<block_nums, CUDA_POOL2D_BLOCK_SIZE, 0, stream>>>(ih, iw, oh, ow, kh, kw, sh, sw, ph, pw, parallel_elements, src, dst, op);

ggml/src/ggml-cuda/pool2d.cuh

@@ -1,5 +1,6 @@
 #include "common.cuh"
 
 #define CUDA_POOL2D_BLOCK_SIZE 256
+#define CUDA_POOL2D_WARP_KERNEL_MIN_WINDOW 32
 
 void ggml_cuda_op_pool2d(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

ggml/src/ggml-hexagon/ggml-hexagon.cpp

@@ -53,6 +53,7 @@
 
 #define GGML_COMMON_IMPL_CPP
 #include "ggml-backend-impl.h"
+#include "ggml-alloc.h"
 #include "ggml-common.h"
 #include "ggml-hexagon.h"
 #include "ggml-impl.h"
@@ -101,13 +102,16 @@ static size_t opt_ndev    = 1;
 static size_t opt_nhvx    = 0; // use all
 static int    opt_nhmx    = 1; // when set, enable HMX; when 0, use HVX only
 static size_t opt_vmem    = HTP_OP_MAX_VMEM_DEFAULT;  // max available va space for buffer mappings
-static size_t opt_mbuf    = 1ul * 1024 * 1024 * 1024; // max buffer size
 static int    opt_etm     = 0;
 static int    opt_verbose = 0;
 static int    opt_profile = 0; // profiling mode (0-disabled, 1-basic, 2-pmu)
 static bool   opt_hostbuf = false;
 static bool   opt_dma64   = false;
 
+static size_t opt_mbuf_dyn    = 512ul * 1024 * 1024;      // max dynamic (compute) buffer size
+static size_t opt_mbuf_static = 1ul * 1024 * 1024 * 1024; // max static (weight/KV) buffer size
+static size_t opt_mbuf_total  = 0;                        // total buffer space limit (0 = unconstrained)
+
 static int    opt_mm_select  = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
 static int    opt_fa_select  = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
 static int    opt_fa_head_split = 1; // 1 = partition flash_attn by KV heads in multicore (default on), 0 = token-based (original)
@@ -121,7 +125,7 @@ static int    opt_ar_scatter = 1; // 1 = reduce-scatter the fused ALLREDUCE+ADD
 static u32vec opt_pmu_evt { 0x3, 0x111, 0x100, 0x105, 0x240, 0x256, 0x7D, 0x8C };
 
 static int opt_opbatch  = 1280; // max number of ops in a batch
-static int opt_opqueue  = 32;   // max number of pending batches
+static int opt_opqueue  = 8;   // max number of pending batches
 static int opt_optrace  = 0;    // trace buffer size per thread (0 means default)
 static int opt_oppoll   = 0;    // polling for batch completions
 static int opt_opfusion = 1;    // enable/disable op fusion
@@ -2820,8 +2824,30 @@ static size_t ggml_backend_hexagon_buffer_type_get_alloc_size(ggml_backend_buffe
     GGML_UNUSED(buft);
 }
 
+static size_t parse_size(const char * str, size_t default_unit = 1024 * 1024) {
+    if (!str || str[0] == '\0') {
+        return 0;
+    }
+    char * end = NULL;
+    double val = strtod(str, &end);
+    if (val < 0) {
+        return 0;
+    }
+    if (end && *end) {
+        while (*end == ' ') end++;
+        if (*end == 'k' || *end == 'K') {
+            return (size_t) (val * 1024);
+        } else if (*end == 'm' || *end == 'M') {
+            return (size_t) (val * 1024 * 1024);
+        } else if (*end == 'g' || *end == 'G') {
+            return (size_t) (val * 1024 * 1024 * 1024);
+        }
+    }
+    return (size_t) (val * default_unit);
+}
+
 static size_t ggml_backend_hexagon_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) {
-    return opt_mbuf;
+    return opt_mbuf_dyn;
     GGML_UNUSED(buft);
 }
 
@@ -2835,25 +2861,199 @@ static bool ggml_backend_hexagon_host_buffer_type_is_host(ggml_backend_buffer_ty
     GGML_UNUSED(buft);
 }
 
+struct ggml_backend_hexagon_alloc_buffer_n_plan_item {
+    size_t size;
+    int    first;
+    int    last;
+};
+
+using ggml_backend_hexagon_alloc_buffer_n_plan_t = std::vector<ggml_backend_hexagon_alloc_buffer_n_plan_item>;
+
+static const char * ggml_hexagon_kv_layer_suffix(const struct ggml_tensor * t) {
+    if (strncmp(t->name, "cache_", 6) != 0) {
+        return NULL;
+    }
+    const char * p = strstr(t->name, "_l");
+    if (!p || !isdigit((unsigned char)p[2])) {
+        return NULL;
+    }
+    return p;
+}
+
+struct ggml_backend_hexagon_alloc_unit {
+    size_t size;
+    int    first;
+    int    last;
+};
+
+static ggml_backend_hexagon_alloc_buffer_n_plan_t ggml_backend_hexagon_alloc_buffer_n_plan(
+        ggml_backend_buffer_type_t buft, struct ggml_tensor ** tensors, int n_tensors) {
+    ggml_backend_hexagon_alloc_buffer_n_plan_t plan;
+
+    const size_t alignment = ggml_backend_buft_get_alignment(buft);
+    const size_t max_size  = opt_mbuf_static > 0 ? opt_mbuf_static : SIZE_MAX;
+
+    std::vector<ggml_backend_hexagon_alloc_unit> units;
+
+    int i = 0;
+    while (i < n_tensors) {
+        struct ggml_tensor * t = tensors[i];
+        size_t unit_size = 0;
+        int unit_first = i;
+        int unit_last = i + 1;
+
+        if (t->data == NULL && t->view_src == NULL) {
+            unit_size += GGML_PAD(ggml_backend_buft_get_alloc_size(buft, t), alignment);
+        }
+
+        const char * layer_suffix = ggml_hexagon_kv_layer_suffix(t);
+
+        while (unit_last < n_tensors) {
+            struct ggml_tensor * next = tensors[unit_last];
+
+            if (next->view_src != NULL) {
+                unit_last++;
+                continue;
+            }
+
+            if (layer_suffix != NULL) {
+                const char * next_suffix = ggml_hexagon_kv_layer_suffix(next);
+                if (next_suffix != NULL && strcmp(layer_suffix, next_suffix) == 0) {
+                    if (next->data == NULL) {
+                        unit_size += GGML_PAD(ggml_backend_buft_get_alloc_size(buft, next), alignment);
+                    }
+                    unit_last++;
+                    continue;
+                }
+            }
+
+            break;
+        }
+
+        units.push_back({ unit_size, unit_first, unit_last });
+        i = unit_last;
+    }
+
+    size_t cur_buf_size  = 0;
+    int    cur_buf_first = 0;
+
+    for (const auto & unit : units) {
+        if (unit.size == 0) {
+            continue;
+        }
+
+        if (cur_buf_size > 0 && (cur_buf_size + unit.size) > max_size) {
+            plan.push_back({ cur_buf_size, cur_buf_first, unit.first });
+            cur_buf_size  = 0;
+            cur_buf_first = unit.first;
+        }
+
+        cur_buf_size += unit.size;
+    }
+
+    if (cur_buf_size > 0) {
+        plan.push_back({ cur_buf_size, cur_buf_first, n_tensors });
+    }
+
+    return plan;
+}
+
+static ggml_backend_buffer_t ggml_backend_hexagon_buffer_type_alloc_buffer_n(
+        ggml_backend_buffer_type_t buft, struct ggml_tensor ** tensors, int n_tensors) {
+    const ggml_backend_hexagon_alloc_buffer_n_plan_t plan = ggml_backend_hexagon_alloc_buffer_n_plan(buft, tensors, n_tensors);
+
+    std::vector<ggml_backend_buffer_t> buffers;
+    buffers.reserve(plan.size());
+
+    for (const auto & item : plan) {
+        ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, item.size);
+        if (buffer == NULL) {
+            GGML_LOG_ERROR("%s: failed to allocate %s buffer of size %zu\n", __func__, ggml_backend_buft_name(buft), item.size);
+            for (ggml_backend_buffer_t b : buffers) {
+                ggml_backend_buffer_free(b);
+            }
+            return NULL;
+        }
+
+        struct ggml_tallocr tallocr = ggml_tallocr_new(buffer);
+
+        struct ggml_tensor * t_failed = NULL;
+        for (int j = item.first; j < item.last; j++) {
+            struct ggml_tensor * t = tensors[j];
+            if (t->data == NULL) {
+                if (t->view_src == NULL) {
+                    if (ggml_tallocr_alloc(&tallocr, t) != GGML_STATUS_SUCCESS) {
+                        t_failed = t;
+                        break;
+                    }
+                } else if (t->buffer == NULL) {
+                    if (ggml_backend_view_init(t) != GGML_STATUS_SUCCESS) {
+                        t_failed = t;
+                        break;
+                    }
+                }
+            } else {
+                if (t->view_src != NULL && t->buffer == NULL) {
+                    if (ggml_backend_view_init(t) != GGML_STATUS_SUCCESS) {
+                        t_failed = t;
+                        break;
+                    }
+                }
+            }
+        }
+        if (t_failed != NULL) {
+            GGML_LOG_ERROR("%s: failed to initialize tensor %s\n", __func__, t_failed->name);
+            for (ggml_backend_buffer_t b : buffers) {
+                ggml_backend_buffer_free(b);
+            }
+            ggml_backend_buffer_free(buffer);
+            return NULL;
+        }
+
+        buffers.push_back(buffer);
+    }
+
+    if (buffers.empty()) {
+        return NULL;
+    }
+
+    if (buffers.size() == 1) {
+        return buffers[0];
+    }
+
+    return ggml_backend_multi_buffer_alloc_buffer(buffers.data(), buffers.size());
+}
+
+static size_t ggml_backend_hexagon_buffer_type_get_alloc_size_n(
+        ggml_backend_buffer_type_t buft, struct ggml_tensor ** tensors, int n_tensors) {
+    const ggml_backend_hexagon_alloc_buffer_n_plan_t plan = ggml_backend_hexagon_alloc_buffer_n_plan(buft, tensors, n_tensors);
+
+    size_t total = 0;
+    for (const auto & item : plan) {
+        total += item.size;
+    }
+    return total;
+}
+
 static ggml_backend_buffer_type_i ggml_backend_hexagon_buffer_type_interface = {
     /* .get_name            = */ ggml_backend_hexagon_buffer_type_name,
     /* .alloc_buffer        = */ ggml_backend_hexagon_buffer_type_alloc_buffer,
-    /* .alloc_buffer_n      = */ NULL,
+    /* .alloc_buffer_n      = */ ggml_backend_hexagon_buffer_type_alloc_buffer_n,
     /* .get_alignment       = */ ggml_backend_hexagon_buffer_type_get_alignment,
     /* .get_max_size        = */ ggml_backend_hexagon_buffer_type_get_max_size,
     /* .get_alloc_size      = */ ggml_backend_hexagon_buffer_type_get_alloc_size,
-    /* .get_alloc_size_n    = */ NULL,
+    /* .get_alloc_size_n    = */ ggml_backend_hexagon_buffer_type_get_alloc_size_n,
     /* .is_host             = */ ggml_backend_hexagon_buffer_type_is_host,
 };
 
 static ggml_backend_buffer_type_i ggml_backend_hexagon_host_buffer_type_interface = {
     /* .get_name            = */ ggml_backend_hexagon_buffer_type_name,
     /* .alloc_buffer        = */ ggml_backend_hexagon_host_buffer_type_alloc_buffer,
-    /* .alloc_buffer_n      = */ NULL,
+    /* .alloc_buffer_n      = */ ggml_backend_hexagon_buffer_type_alloc_buffer_n,
     /* .get_alignment       = */ ggml_backend_hexagon_buffer_type_get_alignment,
     /* .get_max_size        = */ ggml_backend_hexagon_buffer_type_get_max_size,
     /* .get_alloc_size      = */ ggml_backend_hexagon_buffer_type_get_alloc_size,
-    /* .get_alloc_size_n    = */ NULL,
+    /* .get_alloc_size_n    = */ ggml_backend_hexagon_buffer_type_get_alloc_size_n,
     /* .is_host             = */ ggml_backend_hexagon_host_buffer_type_is_host,
 };
 
@@ -8541,8 +8741,8 @@ static const char * ggml_backend_hexagon_device_get_description(ggml_backend_dev
 }
 
 static void ggml_backend_hexagon_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) {
-    *free  = 0;
-    *total = *free;
+    *free  = opt_mbuf_total;
+    *total = opt_mbuf_total;
 
     GGML_UNUSED(dev);
 }
@@ -9291,7 +9491,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
     size_t MiB = 1024 * 1024;
 
     // Update vmem default
-    opt_vmem = opt_arch >= 75 ? HTP_OP_MAX_VMEM_DEFAULT : 3000 * MiB;
+    opt_vmem  = opt_arch >= 75 ? HTP_OP_MAX_VMEM_DEFAULT : 3000 * MiB;
     opt_dma64 = opt_arch > 79 && (!str_dma64 || atoi(str_dma64) != 0);
 
     auto RE_ICASE = std::regex_constants::icase;
@@ -9309,13 +9509,30 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
     opt_nhmx      = str_nhmx     ? atoi(str_nhmx)                         : opt_nhmx;
     opt_mm_select = str_mm_select ? atoi(str_mm_select)                   : opt_mm_select;
     opt_fa_select = str_fa_select ? atoi(str_fa_select)                   : opt_fa_select;
-    opt_fa_head_split = str_fa_head_split ? atoi(str_fa_head_split)        : opt_fa_head_split;
-    opt_gdn_select = str_gdn_select ? atoi(str_gdn_select)                 : opt_gdn_select;
-    opt_ar_select = str_ar_select ? atoi(str_ar_select)                   : opt_ar_select;
-    opt_ar_scatter = str_ar_scatter ? atoi(str_ar_scatter)                : opt_ar_scatter;
-    opt_mbuf      = str_mbuf     ? strtoul(str_mbuf, NULL, 0) * MiB       : opt_mbuf;
-    opt_vmem      = str_vmem     ? strtoul(str_vmem, NULL, 0) * MiB       : opt_vmem;
-    opt_hostbuf   = str_hostbuf  ? atoi(str_hostbuf) != 0                 : opt_hostbuf;
+    opt_fa_head_split = str_fa_head_split ? atoi(str_fa_head_split)       : opt_fa_head_split;
+    opt_gdn_select    = str_gdn_select    ? atoi(str_gdn_select)          : opt_gdn_select;
+    opt_ar_select     = str_ar_select     ? atoi(str_ar_select)           : opt_ar_select;
+    opt_ar_scatter    = str_ar_scatter    ? atoi(str_ar_scatter)          : opt_ar_scatter;
+
+    if (str_mbuf) {
+        const char * p = str_mbuf;
+        for (int idx = 0; idx < 3 && p && *p; idx++) {
+            while (*p == ' ') p++;
+            const char * comma = strchr(p, ',');
+            size_t len = comma ? (size_t)(comma - p) : strlen(p);
+            while (len > 0 && p[len - 1] == ' ') len--;
+            if (len > 0) {
+                std::string token(p, len);
+                if (idx == 0) opt_mbuf_dyn    = parse_size(token.c_str());
+                if (idx == 1) opt_mbuf_static = parse_size(token.c_str());
+                if (idx == 2) opt_mbuf_total  = parse_size(token.c_str());
+            }
+            if (!comma) break;
+            p = comma + 1;
+        }
+    }
+    opt_vmem    = str_vmem    ? parse_size(str_vmem)   : opt_vmem;
+    opt_hostbuf = str_hostbuf ? atoi(str_hostbuf) != 0 : opt_hostbuf;
 
     // Parse device configuration
     const char * str_devices  = getenv("GGML_HEXAGON_DEVICES");

ggml/src/ggml-metal/ggml-metal-common.cpp

@@ -80,6 +80,22 @@ static bool ggml_metal_mul_mv_mma_type_supported(enum ggml_type type) {
         case GGML_TYPE_Q4_K:
         case GGML_TYPE_Q5_K:
         case GGML_TYPE_Q6_K:
+        case GGML_TYPE_BF16:
+        case GGML_TYPE_Q1_0:
+        case GGML_TYPE_Q2_0:
+        case GGML_TYPE_MXFP4:
+        case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q3_K:
+        case GGML_TYPE_IQ2_XXS:
+        case GGML_TYPE_IQ2_XS:
+        case GGML_TYPE_IQ2_S:
+        case GGML_TYPE_IQ3_XXS:
+        case GGML_TYPE_IQ3_S:
+        case GGML_TYPE_IQ1_S:
+        case GGML_TYPE_IQ1_M:
+        case GGML_TYPE_IQ4_NL:
+        case GGML_TYPE_IQ4_XS:
+        case GGML_TYPE_TQ2_0:
             return true;
         default:
             return false;
@@ -105,10 +121,18 @@ static int64_t ggml_metal_mul_mv_mma_rows_min(enum ggml_type type) {
     switch (type) {
         case GGML_TYPE_F32:
             return 6;
+        case GGML_TYPE_TQ2_0:
+            return 5;
+        case GGML_TYPE_BF16:
+            return 4;
         case GGML_TYPE_F16:
+        case GGML_TYPE_Q2_0:
         case GGML_TYPE_Q4_K:
         case GGML_TYPE_Q5_0:
         case GGML_TYPE_Q5_1:
+        case GGML_TYPE_Q2_K:
+        case GGML_TYPE_IQ4_NL:
+        case GGML_TYPE_MXFP4:
             return 3;
         default:
             return 2;

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -2476,7 +2476,8 @@ static int ggml_metal_op_mul_mat_mma(ggml_metal_op_t ctx, int idx) {
 
     if (fuse_add) {
         dst = ctx->node(idx + n_fuse - 1);
-        res = dst->src[0]->op == GGML_OP_MUL_MAT ? dst->src[1] : dst->src[0];
+        // the residual is the other operand of the ADD, by identity: it can itself be a MUL_MAT output
+        res = dst->src[0] == op ? dst->src[1] : dst->src[0];
     }
 
     auto pipeline = ggml_metal_library_get_pipeline_mul_mv_mma_auto(lib, op, fuse_add);

ggml/src/ggml-metal/kernels/mul_mv_mma.metal

@@ -605,4 +605,23 @@ MUL_MV_MMA_GEN("q5_1", block_q5_1, 2,     dequantize_q5_1)
 MUL_MV_MMA_GEN("q4_K", block_q4_K, QK_NL, dequantize_q4_K)
 MUL_MV_MMA_GEN("q6_K", block_q6_K, QK_NL, dequantize_q6_K)
 
+#if defined(GGML_METAL_HAS_BF16)
+MUL_MV_MMA_GEN("bf16",    bfloat4x4,     1,     dequantize_bf16)
+#endif
+MUL_MV_MMA_GEN("q1_0",    block_q1_0,    8,     dequantize_q1_0)
+MUL_MV_MMA_GEN("q2_0",    block_q2_0,    4,     dequantize_q2_0)
+MUL_MV_MMA_GEN("mxfp4",   block_mxfp4,   2,     dequantize_mxfp4)
+MUL_MV_MMA_GEN("q2_K",    block_q2_K,    QK_NL, dequantize_q2_K)
+MUL_MV_MMA_GEN("q3_K",    block_q3_K,    QK_NL, dequantize_q3_K)
+MUL_MV_MMA_GEN("iq2_xxs", block_iq2_xxs, QK_NL, dequantize_iq2_xxs)
+MUL_MV_MMA_GEN("iq2_xs",  block_iq2_xs,  QK_NL, dequantize_iq2_xs)
+MUL_MV_MMA_GEN("iq2_s",   block_iq2_s,   QK_NL, dequantize_iq2_s)
+MUL_MV_MMA_GEN("iq3_xxs", block_iq3_xxs, QK_NL, dequantize_iq3_xxs)
+MUL_MV_MMA_GEN("iq3_s",   block_iq3_s,   QK_NL, dequantize_iq3_s)
+MUL_MV_MMA_GEN("iq1_s",   block_iq1_s,   QK_NL, dequantize_iq1_s)
+MUL_MV_MMA_GEN("iq1_m",   block_iq1_m,   QK_NL, dequantize_iq1_m)
+MUL_MV_MMA_GEN("iq4_nl",  block_iq4_nl,  2,     dequantize_iq4_nl)
+MUL_MV_MMA_GEN("iq4_xs",  block_iq4_xs,  QK_NL, dequantize_iq4_xs)
+MUL_MV_MMA_GEN("tq2_0",   block_tq2_0,   QK_NL, dequantize_tq2_0)
+
 #undef MUL_MV_MMA_GEN

ggml/src/ggml-sycl/common.hpp

@@ -26,7 +26,6 @@
 #include "presets.hpp"
 #include "type.hpp"
 #include "sycl_hw.hpp"
-#include "fattn-buffers.hpp"
 #include "memtrace.hpp"
 
 namespace syclexp = sycl::ext::oneapi::experimental;
@@ -408,8 +407,6 @@ struct ggml_backend_sycl_context {
     // pool
     std::unique_ptr<ggml_sycl_pool> pools[GGML_SYCL_MAX_DEVICES];
 
-    std::unique_ptr<ggml_sycl_fattn_kv_buffers> fattn_bufs[GGML_SYCL_MAX_DEVICES];
-
     std::unique_ptr<ggml_sycl_pool> host_pools[GGML_SYCL_MAX_DEVICES];
 
     std::vector<mmid_row_mapping> mmid_row_mapping_host;
@@ -418,8 +415,6 @@ struct ggml_backend_sycl_context {
 
     static std::unique_ptr<ggml_sycl_pool> new_pool_for_host(queue_ptr qptr, int device);
 
-    static std::unique_ptr<ggml_sycl_fattn_kv_buffers> new_fattn_kv_buffers(queue_ptr qptr, int device);
-
     ggml_sycl_pool & pool(int device) {
         if (pools[device] == nullptr) {
             pools[device] = new_pool_for_device(stream(device,0), device);
@@ -431,17 +426,6 @@ struct ggml_backend_sycl_context {
         return pool(device);
     }
 
-    ggml_sycl_fattn_kv_buffers & fattn_buffers(int device) {
-        if (fattn_bufs[device] == nullptr) {
-            fattn_bufs[device] = new_fattn_kv_buffers(stream(device, 0), device);
-        }
-        return *fattn_bufs[device];
-    }
-
-    ggml_sycl_fattn_kv_buffers & fattn_buffers() {
-        return fattn_buffers(device);
-    }
-
 #ifdef GGML_SYCL_GRAPH
     std::unique_ptr<sycl_ex::command_graph<sycl_ex::graph_state::executable>> exec_graph = nullptr;
 #endif

ggml/src/ggml-sycl/fattn-buffers.cpp

@@ -1,60 +0,0 @@
-//
-// MIT license
-// Copyright (C) 2025 Intel Corporation
-// SPDX-License-Identifier: MIT
-//
-
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-
-#include "common.hpp"
-
-sycl::half * ggml_sycl_fattn_kv_buffers::kv_buffer::ensure_half(size_t n_elems) {
-    const size_t need_bytes = n_elems * sizeof(sycl::half);
-
-    if (capacity >= need_bytes) {
-        return ptr;
-    }
-
-    if (ptr) {
-        SYCL_CHECK(CHECK_TRY_ERROR(qptr->wait()));
-        ggml_sycl_memtrace_del(ptr);
-        SYCL_CHECK(CHECK_TRY_ERROR(sycl::free(ptr, *qptr)));
-        ptr = nullptr;
-        capacity = 0;
-    }
-
-    size_t cap = 0;
-    while (cap < need_bytes) {
-        cap += CHUNK_SIZE;
-    }
-
-    void * dev_ptr;
-    SYCL_CHECK(
-        CHECK_TRY_ERROR(dev_ptr = sycl::malloc_device(
-                        cap, *qptr)));
-
-    if (!dev_ptr) {
-        GGML_LOG_ERROR("%s: can't allocate %lu Bytes of memory on device\n", __func__, cap);
-        ggml_sycl_memtrace_fail(GGML_SYCL_MEM_FATTN_KV, cap);
-        GGML_ABORT("fattn buffer alloc failed");
-    }
-
-    ptr = static_cast<sycl::half *>(dev_ptr);
-    capacity = cap;
-    ggml_sycl_memtrace_add(GGML_SYCL_MEM_FATTN_KV, ptr, cap);
-    return ptr;
-}
-
-ggml_sycl_fattn_kv_buffers::kv_buffer::~kv_buffer() {
-#ifdef DEBUG_SYCL_POOL
-    GGML_LOG_INFO("ggml_sycl_fattn_kv_buffer[%d]: %.2f MiB\n", device, capacity / 1024.0 / 1024.0);
-#endif
-    if (ptr) {
-        ggml_sycl_memtrace_del(ptr);
-        SYCL_CHECK(CHECK_TRY_ERROR(sycl::free(ptr, *qptr)));
-    }
-}

ggml/src/ggml-sycl/fattn-buffers.hpp

@@ -1,63 +0,0 @@
-//
-// MIT license
-// Copyright (C) 2025 Intel Corporation
-// SPDX-License-Identifier: MIT
-//
-
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-
-#ifndef GGML_SYCL_FATTN_BUFFERS_HPP
-#define GGML_SYCL_FATTN_BUFFERS_HPP
-
-#include <sycl/sycl.hpp>
-
-typedef sycl::queue *queue_ptr;
-
-struct ggml_sycl_fattn_kv_buffers {
-    // buffers grow in chunks of this size
-    static constexpr size_t CHUNK_SIZE = 16ull << 20; // 16 MiB
-
-    struct kv_buffer {
-        kv_buffer(queue_ptr qptr_, int device_) : qptr(qptr_), device(device_) {}
-        ~kv_buffer();
-
-        kv_buffer(const kv_buffer &) = delete;
-        kv_buffer & operator=(const kv_buffer &) = delete;
-
-        sycl::half * ensure_half(size_t n_elems);
-
-    private:
-        sycl::half * ptr      = nullptr;
-        size_t       capacity = 0;
-        queue_ptr    qptr     = nullptr;
-        [[maybe_unused]] int device = 0;
-    };
-
-    kv_buffer K;
-    kv_buffer V;
-
-    ggml_sycl_fattn_kv_buffers(queue_ptr qptr, int device) : K(qptr, device), V(qptr, device) {}
-
-    ggml_sycl_fattn_kv_buffers(const ggml_sycl_fattn_kv_buffers &) = delete;
-    ggml_sycl_fattn_kv_buffers & operator=(const ggml_sycl_fattn_kv_buffers &) = delete;
-};
-
-/**
- * Imitates `ggml_sycl_pool_alloc` to keep the code calling alloc unchanged.
- */
-struct ggml_sycl_fattn_alloc {
-    ggml_sycl_fattn_kv_buffers::kv_buffer & buf;
-    sycl::half *                         ptr = nullptr;
-
-    explicit ggml_sycl_fattn_alloc(ggml_sycl_fattn_kv_buffers::kv_buffer & buf_) : buf(buf_) {}
-
-    sycl::half * alloc(size_t n_elems) {
-        ptr = buf.ensure_half(n_elems);
-        return ptr;
-    }
-};
-#endif

ggml/src/ggml-sycl/fattn-common.hpp

@@ -6,7 +6,6 @@
 #include "common.hpp"
 #include "convert.hpp"
 #include "vecdotq.hpp"
-#include "fattn-buffers.hpp"
 #include "fattn.hpp"
 
 #include "ggml.h"
@@ -933,13 +932,10 @@ void launch_fattn(
     GGML_ASSERT(!mask || mask->type == GGML_TYPE_F16);
 
     ggml_sycl_pool & pool = ctx.pool();
-    ggml_sycl_fattn_kv_buffers & fbuf = ctx.fattn_buffers();
     dpct::queue_ptr  main_stream = ctx.stream();
     const int id  = ggml_sycl_get_device();
     const int nsm = ggml_sycl_info().devices[id].nsm;
 
-    ggml_sycl_fattn_alloc        K_f16(fbuf.K);
-    ggml_sycl_fattn_alloc        V_f16(fbuf.V);
     const ggml_sycl_fattn_extra  extra = ggml_sycl_fattn_get_extra(dst);
     ggml_sycl_pool_alloc<int>    KV_max(pool);
     ggml_sycl_pool_alloc<float>  dst_tmp(pool);
@@ -959,8 +955,8 @@ void launch_fattn(
         const size_t bs = ggml_blck_size(K->type);
         const size_t ts = ggml_type_size(K->type);
 
-        sycl::half * K_f16_ptr = extra.K_buffer_ptr ? (sycl::half *) extra.K_buffer_ptr
-                                                    : K_f16.alloc(ggml_nelements(K));
+        GGML_ASSERT(extra.K_buffer_ptr);
+        sycl::half * K_f16_ptr = (sycl::half *) extra.K_buffer_ptr;
         if (ggml_is_contiguously_allocated(K)) {
             to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(K->type, dst);
             to_fp16(K_data, K_f16_ptr, ggml_nelements(K), main_stream);
@@ -993,8 +989,8 @@ void launch_fattn(
             const size_t bs = ggml_blck_size(V->type);
             const size_t ts = ggml_type_size(V->type);
 
-            sycl::half * V_f16_ptr = extra.V_buffer_ptr ? (sycl::half *) extra.V_buffer_ptr
-                                                        : V_f16.alloc(ggml_nelements(V));
+            GGML_ASSERT(extra.V_buffer_ptr);
+            sycl::half * V_f16_ptr = (sycl::half *) extra.V_buffer_ptr;
             if (ggml_is_contiguously_allocated(V)) {
                 to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(V->type, dst);
                 to_fp16(V_data, V_f16_ptr, ggml_nelements(V), main_stream);

ggml/src/ggml-sycl/fattn-mkl.cpp

@@ -8,7 +8,6 @@
 
 #include "common.hpp"
 #include "fattn-common.hpp"
-#include "fattn-buffers.hpp"
 #include "convert.hpp"
 #include "fattn.hpp"
 
@@ -283,8 +282,10 @@ static mkl_fa_kv_desc mkl_fa_make_desc(const ggml_tensor * T, bool interleaved,
     d.ts   = (int64_t)ggml_type_size(T->type);
 
     if (T->type == GGML_TYPE_F16) {
-        d.mode = interleaved ? MKL_FA_KV_MODE_F16_INTERLEAVED
-                             : MKL_FA_KV_MODE_F16_DENSE;
+        // MLA's V cache is a 512-wide view of 576-wide K rows. Treat any
+        // padded row stride as strided even when there is only one KV head.
+        d.mode = interleaved || d.nb1 != d.D * (int64_t)sizeof(sycl::half)
+            ? MKL_FA_KV_MODE_F16_INTERLEAVED : MKL_FA_KV_MODE_F16_DENSE;
     } else if (ggml_is_contiguously_allocated(T) && !interleaved) {
         d.mode = MKL_FA_KV_MODE_QUANT_CONTIG;
     } else {
@@ -413,7 +414,9 @@ void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor *
     const int64_t q_row_stride  = Q->nb[1] / sizeof(float);
     const int64_t q_head_stride = Q->nb[2] / sizeof(float);
 
-    const bool V_is_K_view = V->view_src
+    // Alias the dequantized buffers only when K and V expose the same values.
+    // MLA V is a narrower view of K and needs its own strided dequantization.
+    const bool V_is_K_view = V->ne[0] == K->ne[0] && V->view_src
         && (V->view_src == K || (V->view_src == K->view_src
             && V->view_offs == K->view_offs));
 

ggml/src/ggml-sycl/fattn-vec.hpp

@@ -589,10 +589,10 @@ void ggml_sycl_flash_attn_ext_vec_case_impl(ggml_backend_sycl_context & ctx, ggm
     const bool need_f16_V = type_V == GGML_TYPE_F16;
     constexpr size_t nbytes_shared = 0;
 
+    const auto arch = ggml_sycl_info().devices[ggml_sycl_get_device()].hw_info.arch;
     // D=512 does not fit the default register file; it spills up to 343 bytes per thread, against at most 57 for D <= 256. This kernel is decode only, so thread occupancy is not the limit. It is 1.9x faster at every KV depth on Battlemage.
     constexpr bool use_large_grf = D >= 512;
 
-    const auto arch = ggml_sycl_info().devices[ctx.device].hw_info.arch;
     const int nthreads = ggml_sycl_fattn_vec_get_nthreads_device(arch);
     if constexpr (D <= 256) {
         if (nthreads == 256) {

ggml/src/ggml-sycl/fattn.cpp

@@ -146,15 +146,17 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const
     // Set GGML_SYCL_ENABLE_MKL_FA=0 to force TILE/VEC path for A/B testing.
     // Example: GGML_SYCL_ENABLE_MKL_FA=0 llama-cli -m model.gguf -fa -ngl 99 ...
     // Note: MKL GEMM calls are incompatible with SYCL graph capture replay.
-    // MKL is validated for the mainstream GQA envelope: grouped-query
-    // (gqa_ratio >= 2), head_dim a multiple of 64 in [64,512] with matching
-    // K/V head size, mask, no sinks/ALiBi/softcap. Gemma's global layers use
-    // head_dim 512, so the cap must include it. Head sizes not a multiple of
-    // 64 (72/80/96), MHA (gqa_ratio == 1), and MLA (DKQ != DV, e.g. 576/512)
-    // fall through to TILE/VEC; see follow-up work.
+    const bool standard_shape = Q->ne[0] >= 64 && Q->ne[0] <= 512 &&
+        Q->ne[0] % 64 == 0 && Q->ne[0] == V->ne[0];
+    // GLM-4.7 Flash's MLA shape is already expressible by the MKL pipeline:
+    // KQ GEMM uses DKQ=576 while VKQ and output use DV=512. Keep this narrow
+    // until other mismatched K/V shapes have independent correctness data.
+    const bool glm_mla_shape = Q->ne[0] == 576 && K->ne[0] == 576 &&
+        V->ne[0] == 512 && gqa_ratio == 20 &&
+        K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16;
+
     if (g_ggml_sycl_enable_mkl_fa == 1 && mask && !sinks && gqa_ratio >= 2 &&
-        Q->ne[0] >= 64 && Q->ne[0] <= 512 && Q->ne[0] % 64 == 0 &&
-        Q->ne[0] == V->ne[0] &&
+        (standard_shape || glm_mla_shape) &&
         Q->ne[1] >= 32 && K->ne[1] >= 1024 &&
         max_bias == 0.0f && logit_softcap == 0.0f &&
         (Q->ne[3] == K->ne[3] || K->ne[3] == 1)) {
@@ -164,7 +166,10 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const
         // nb1=75 for ne0=40 fall through to TILE.
         bool kv_strides_ok = true;
         for (const ggml_tensor * t : {K, V}) {
-            if (t->type == GGML_TYPE_F16 && t->nb[1] % (t->ne[0] * 2) != 0) {
+            const bool glm_v_stride = glm_mla_shape && t == V &&
+                V->view_src && V->nb[1] == K->nb[1];
+            if (!glm_v_stride && t->type == GGML_TYPE_F16 &&
+                    t->nb[1] % (t->ne[0] * 2) != 0) {
                 kv_strides_ok = false;
                 break;
             }

ggml/src/ggml-sycl/ggml-sycl.cpp

@@ -792,6 +792,14 @@ static bool ggml_sycl_is_l0_discrete_gpu(int device) {
 }
 #endif
 
+static void memcpy_host_forward(sycl::queue &q_dst, sycl::queue &q_src, void *ptr_dst,
+                         const void *ptr_src, size_t size) {
+    char *host_buf = (char *)malloc(size);
+    q_src.memcpy(host_buf, (const char *)ptr_src, size).wait();
+    q_dst.memcpy((char *)ptr_dst, host_buf, size).wait();
+    free(host_buf);
+}
+
 static void dev2dev_memcpy(int device_dst, sycl::queue &q_dst, int device_src, sycl::queue &q_src, void *ptr_dst,
                     const void *ptr_src, size_t size) {
 
@@ -835,10 +843,7 @@ static void dev2dev_memcpy(int device_dst, sycl::queue &q_dst, int device_src, s
     } else {
         GGML_SYCL_DEBUG("[SYCL] dev2dev memcpy by host forward for SYCL/L0 fallback\n");
     }
-    char *host_buf = (char *)malloc(size);
-    q_src.memcpy(host_buf, (const char *)ptr_src, size).wait();
-    q_dst.memcpy((char *)ptr_dst, host_buf, size).wait();
-    free(host_buf);
+    memcpy_host_forward(q_dst, q_src, ptr_dst, ptr_src, size);
 }
 
 static bool
@@ -2064,11 +2069,6 @@ std::unique_ptr<ggml_sycl_pool> ggml_backend_sycl_context::new_pool_for_device(q
     return std::unique_ptr<ggml_sycl_pool>(new ggml_sycl_pool_leg(qptr, device));
 }
 
-
-std::unique_ptr<ggml_sycl_fattn_kv_buffers> ggml_backend_sycl_context::new_fattn_kv_buffers(queue_ptr qptr, int device) {
-    return std::unique_ptr<ggml_sycl_fattn_kv_buffers>(new ggml_sycl_fattn_kv_buffers(qptr, device));
-}
-
 /// kernels
 typedef void (*ggml_sycl_op_mul_mat_t)(
     ggml_backend_sycl_context & ctx,

ggml/src/ggml-sycl/memtrace.cpp

@@ -15,7 +15,6 @@ static const char * mem_type_name(ggml_sycl_mem_type type) {
         case GGML_SYCL_MEM_POOL_LEG: return "pool_leg";
         case GGML_SYCL_MEM_POOL_VMM: return "pool_vmm";
         case GGML_SYCL_MEM_ASYNC:    return "async";
-        case GGML_SYCL_MEM_FATTN_KV: return "fattn_kv";
         case GGML_SYCL_MEM_DIRECT:   return "direct";
         default:                     GGML_ABORT("[%s] The type value %d is not supported\n", __func__, (int) type);
     }

ggml/src/ggml-sycl/memtrace.hpp

@@ -10,7 +10,6 @@ enum ggml_sycl_mem_type {
     GGML_SYCL_MEM_POOL_LEG,
     GGML_SYCL_MEM_POOL_VMM,
     GGML_SYCL_MEM_ASYNC,
-    GGML_SYCL_MEM_FATTN_KV,
     GGML_SYCL_MEM_DIRECT,
 
     GGML_SYCL_MEM_TYPE_COUNT,

ggml/src/ggml-webgpu/ggml-webgpu.cpp

@@ -4544,7 +4544,7 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
                     break;
                 }
                 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)) {
+                    src1->type != src2->type && !(ggml_is_quantized(src1->type) && ggml_is_quantized(src2->type))) {
                     supports_op = false;
                     break;
                 }

gguf-py/gguf/constants.py

@@ -408,6 +408,7 @@ class Keys:
         BLOCK_COUNT           = "clip.vision.block_count"
         IMAGE_MEAN            = "clip.vision.image_mean"
         IMAGE_STD             = "clip.vision.image_std"
+        IMAGE_RESIZE_ALGO     = "clip.vision.image_resize_algo"
         SPATIAL_MERGE_SIZE    = "clip.vision.spatial_merge_size"
         SWIGLU_CLAMP          = "clip.vision.swiglu_clamp"
         EXPERT_COUNT_PER_LAYER = "clip.vision.expert_count_per_layer" # dots3note pyramid MoE, 0 = dense layer
@@ -6070,6 +6071,7 @@ class DecisionType:
     NIMBLE  = "nimble"   # same as openjev, the prompt lists all the questions of the request
     CLEF    = "clef"     # joint head over all questions, one score per option
     PPLX_DECIDER = "pplx-decider"  # same as openjev, label codes of 1 or 2 letters
+    LFM2_D1 = "lfm2-d1"  # same as openjev, the labels depend on the question type
 
 
 class VisionProjectorType:
@@ -6130,6 +6132,7 @@ class VisionProjectorType:
     MIMO_AUDIO     = "mimo_audio"
     GRANITE4_VISION = "granite4_vision"
     MUSE_GLIMMER   = "muse-glimmer"
+    COHERE2V       = "cohere2v"
 
 
 # Items here are (block size, type size)

gguf-py/gguf/gguf_writer.py

@@ -1432,6 +1432,9 @@ class GGUFWriter:
     def add_vision_image_mean(self, values: Sequence[float]) -> None:
         self.add_array(Keys.ClipVision.IMAGE_MEAN, values)
 
+    def add_vision_image_resize_algo(self, value: str) -> None:
+        self.add_string(Keys.ClipVision.IMAGE_RESIZE_ALGO, value)
+
     def add_vision_image_std(self, values: Sequence[float]) -> None:
         self.add_array(Keys.ClipVision.IMAGE_STD, values)
 

gguf-py/gguf/tensor_mapping.py

@@ -1622,6 +1622,7 @@ class TensorNameMap:
         MODEL_TENSOR.V_MMPROJ: (
             "aligner.w{bid}", # deepseek4v (w1 -> mm.1, w2 -> mm.2)
             "multi_modal_projector.linear_{bid}",
+            "model.multi_modal_projector.linear_{bid}", # cohere2v
             "mm_projector.proj.linear_{bid}", # Kimi-K2.5
             "visual.merger.mlp.{bid}", # qwen2vl
             "mlp_AR.linear_{bid}", # PaddleOCR-VL

include/llama.h

@@ -396,6 +396,8 @@ extern "C" {
         enum ggml_type type_k; // data type for K cache [EXPERIMENTAL]
         enum ggml_type type_v; // data type for V cache [EXPERIMENTAL]
 
+        size_t moe_cache_size; // device cache in bytes for the experts kept in host memory, 0 = disabled [EXPERIMENTAL]
+
         // Abort callback
         // if it returns true, execution of llama_decode() will be aborted
         // currently works only with CPU execution

models/templates/google-translategemma-4b-it.jinja

@@ -0,0 +1,641 @@
+{%- set languages = {
+    "aa": "Afar",
+    "aa-DJ": "Afar",
+    "aa-ER": "Afar",
+    "ab": "Abkhazian",
+    "af": "Afrikaans",
+    "af-NA": "Afrikaans",
+    "ak": "Akan",
+    "am": "Amharic",
+    "an": "Aragonese",
+    "ar": "Arabic",
+    "ar-AE": "Arabic",
+    "ar-BH": "Arabic",
+    "ar-DJ": "Arabic",
+    "ar-DZ": "Arabic",
+    "ar-EG": "Arabic",
+    "ar-EH": "Arabic",
+    "ar-ER": "Arabic",
+    "ar-IL": "Arabic",
+    "ar-IQ": "Arabic",
+    "ar-JO": "Arabic",
+    "ar-KM": "Arabic",
+    "ar-KW": "Arabic",
+    "ar-LB": "Arabic",
+    "ar-LY": "Arabic",
+    "ar-MA": "Arabic",
+    "ar-MR": "Arabic",
+    "ar-OM": "Arabic",
+    "ar-PS": "Arabic",
+    "ar-QA": "Arabic",
+    "ar-SA": "Arabic",
+    "ar-SD": "Arabic",
+    "ar-SO": "Arabic",
+    "ar-SS": "Arabic",
+    "ar-SY": "Arabic",
+    "ar-TD": "Arabic",
+    "ar-TN": "Arabic",
+    "ar-YE": "Arabic",
+    "as": "Assamese",
+    "az": "Azerbaijani",
+    "az-Arab": "Azerbaijani",
+    "az-Arab-IQ": "Azerbaijani",
+    "az-Arab-TR": "Azerbaijani",
+    "az-Cyrl": "Azerbaijani",
+    "az-Latn": "Azerbaijani",
+    "ba": "Bashkir",
+    "be": "Belarusian",
+    "be-tarask": "Belarusian",
+    "bg": "Bulgarian",
+    "bg-BG": "Bulgarian",
+    "bm": "Bambara",
+    "bm-Nkoo": "Bambara",
+    "bn": "Bengali",
+    "bn-IN": "Bengali",
+    "bo": "Tibetan",
+    "bo-IN": "Tibetan",
+    "br": "Breton",
+    "bs": "Bosnian",
+    "bs-Cyrl": "Bosnian",
+    "bs-Latn": "Bosnian",
+    "ca": "Catalan",
+    "ca-AD": "Catalan",
+    "ca-ES": "Catalan",
+    "ca-FR": "Catalan",
+    "ca-IT": "Catalan",
+    "ce": "Chechen",
+    "co": "Corsican",
+    "cs": "Czech",
+    "cs-CZ": "Czech",
+    "cv": "Chuvash",
+    "cy": "Welsh",
+    "da": "Danish",
+    "da-DK": "Danish",
+    "da-GL": "Danish",
+    "de": "German",
+    "de-AT": "German",
+    "de-BE": "German",
+    "de-CH": "German",
+    "de-DE": "German",
+    "de-IT": "German",
+    "de-LI": "German",
+    "de-LU": "German",
+    "dv": "Divehi",
+    "dz": "Dzongkha",
+    "ee": "Ewe",
+    "ee-TG": "Ewe",
+    "el": "Greek",
+    "el-CY": "Greek",
+    "el-GR": "Greek",
+    "el-polyton": "Greek",
+    "en": "English",
+    "en-AE": "English",
+    "en-AG": "English",
+    "en-AI": "English",
+    "en-AS": "English",
+    "en-AT": "English",
+    "en-AU": "English",
+    "en-BB": "English",
+    "en-BE": "English",
+    "en-BI": "English",
+    "en-BM": "English",
+    "en-BS": "English",
+    "en-BW": "English",
+    "en-BZ": "English",
+    "en-CA": "English",
+    "en-CC": "English",
+    "en-CH": "English",
+    "en-CK": "English",
+    "en-CM": "English",
+    "en-CX": "English",
+    "en-CY": "English",
+    "en-CZ": "English",
+    "en-DE": "English",
+    "en-DG": "English",
+    "en-DK": "English",
+    "en-DM": "English",
+    "en-ER": "English",
+    "en-ES": "English",
+    "en-FI": "English",
+    "en-FJ": "English",
+    "en-FK": "English",
+    "en-FM": "English",
+    "en-FR": "English",
+    "en-GB": "English",
+    "en-GD": "English",
+    "en-GG": "English",
+    "en-GH": "English",
+    "en-GI": "English",
+    "en-GM": "English",
+    "en-GS": "English",
+    "en-GU": "English",
+    "en-GY": "English",
+    "en-HK": "English",
+    "en-HU": "English",
+    "en-ID": "English",
+    "en-IE": "English",
+    "en-IL": "English",
+    "en-IM": "English",
+    "en-IN": "English",
+    "en-IO": "English",
+    "en-IT": "English",
+    "en-JE": "English",
+    "en-JM": "English",
+    "en-KE": "English",
+    "en-KI": "English",
+    "en-KN": "English",
+    "en-KY": "English",
+    "en-LC": "English",
+    "en-LR": "English",
+    "en-LS": "English",
+    "en-MG": "English",
+    "en-MH": "English",
+    "en-MO": "English",
+    "en-MP": "English",
+    "en-MS": "English",
+    "en-MT": "English",
+    "en-MU": "English",
+    "en-MV": "English",
+    "en-MW": "English",
+    "en-MY": "English",
+    "en-NA": "English",
+    "en-NF": "English",
+    "en-NG": "English",
+    "en-NL": "English",
+    "en-NO": "English",
+    "en-NR": "English",
+    "en-NU": "English",
+    "en-NZ": "English",
+    "en-PG": "English",
+    "en-PH": "English",
+    "en-PK": "English",
+    "en-PL": "English",
+    "en-PN": "English",
+    "en-PR": "English",
+    "en-PT": "English",
+    "en-PW": "English",
+    "en-RO": "English",
+    "en-RW": "English",
+    "en-SB": "English",
+    "en-SC": "English",
+    "en-SD": "English",
+    "en-SE": "English",
+    "en-SG": "English",
+    "en-SH": "English",
+    "en-SI": "English",
+    "en-SK": "English",
+    "en-SL": "English",
+    "en-SS": "English",
+    "en-SX": "English",
+    "en-SZ": "English",
+    "en-TC": "English",
+    "en-TK": "English",
+    "en-TO": "English",
+    "en-TT": "English",
+    "en-TV": "English",
+    "en-TZ": "English",
+    "en-UG": "English",
+    "en-UM": "English",
+    "en-VC": "English",
+    "en-VG": "English",
+    "en-VI": "English",
+    "en-VU": "English",
+    "en-WS": "English",
+    "en-ZA": "English",
+    "en-ZM": "English",
+    "en-ZW": "English",
+    "eo": "Esperanto",
+    "es": "Spanish",
+    "es-AR": "Spanish",
+    "es-BO": "Spanish",
+    "es-BR": "Spanish",
+    "es-BZ": "Spanish",
+    "es-CL": "Spanish",
+    "es-CO": "Spanish",
+    "es-CR": "Spanish",
+    "es-CU": "Spanish",
+    "es-DO": "Spanish",
+    "es-EA": "Spanish",
+    "es-EC": "Spanish",
+    "es-ES": "Spanish",
+    "es-GQ": "Spanish",
+    "es-GT": "Spanish",
+    "es-HN": "Spanish",
+    "es-IC": "Spanish",
+    "es-MX": "Spanish",
+    "es-NI": "Spanish",
+    "es-PA": "Spanish",
+    "es-PE": "Spanish",
+    "es-PH": "Spanish",
+    "es-PR": "Spanish",
+    "es-PY": "Spanish",
+    "es-SV": "Spanish",
+    "es-US": "Spanish",
+    "es-UY": "Spanish",
+    "es-VE": "Spanish",
+    "et": "Estonian",
+    "et-EE": "Estonian",
+    "eu": "Basque",
+    "fa": "Persian",
+    "fa-AF": "Persian",
+    "fa-IR": "Persian",
+    "ff": "Fulah",
+    "ff-Adlm": "Fulah",
+    "ff-Adlm-BF": "Fulah",
+    "ff-Adlm-CM": "Fulah",
+    "ff-Adlm-GH": "Fulah",
+    "ff-Adlm-GM": "Fulah",
+    "ff-Adlm-GW": "Fulah",
+    "ff-Adlm-LR": "Fulah",
+    "ff-Adlm-MR": "Fulah",
+    "ff-Adlm-NE": "Fulah",
+    "ff-Adlm-NG": "Fulah",
+    "ff-Adlm-SL": "Fulah",
+    "ff-Adlm-SN": "Fulah",
+    "ff-Latn": "Fulah",
+    "ff-Latn-BF": "Fulah",
+    "ff-Latn-CM": "Fulah",
+    "ff-Latn-GH": "Fulah",
+    "ff-Latn-GM": "Fulah",
+    "ff-Latn-GN": "Fulah",
+    "ff-Latn-GW": "Fulah",
+    "ff-Latn-LR": "Fulah",
+    "ff-Latn-MR": "Fulah",
+    "ff-Latn-NE": "Fulah",
+    "ff-Latn-NG": "Fulah",
+    "ff-Latn-SL": "Fulah",
+    "fi": "Finnish",
+    "fi-FI": "Finnish",
+    "fil-PH": "Filipino",
+    "fo": "Faroese",
+    "fo-DK": "Faroese",
+    "fr": "French",
+    "fr-BE": "French",
+    "fr-BF": "French",
+    "fr-BI": "French",
+    "fr-BJ": "French",
+    "fr-BL": "French",
+    "fr-CA": "French",
+    "fr-CD": "French",
+    "fr-CF": "French",
+    "fr-CG": "French",
+    "fr-CH": "French",
+    "fr-CI": "French",
+    "fr-CM": "French",
+    "fr-DJ": "French",
+    "fr-DZ": "French",
+    "fr-FR": "French",
+    "fr-GA": "French",
+    "fr-GF": "French",
+    "fr-GN": "French",
+    "fr-GP": "French",
+    "fr-GQ": "French",
+    "fr-HT": "French",
+    "fr-KM": "French",
+    "fr-LU": "French",
+    "fr-MA": "French",
+    "fr-MC": "French",
+    "fr-MF": "French",
+    "fr-MG": "French",
+    "fr-ML": "French",
+    "fr-MQ": "French",
+    "fr-MR": "French",
+    "fr-MU": "French",
+    "fr-NC": "French",
+    "fr-NE": "French",
+    "fr-PF": "French",
+    "fr-PM": "French",
+    "fr-RE": "French",
+    "fr-RW": "French",
+    "fr-SC": "French",
+    "fr-SN": "French",
+    "fr-SY": "French",
+    "fr-TD": "French",
+    "fr-TG": "French",
+    "fr-TN": "French",
+    "fr-VU": "French",
+    "fr-WF": "French",
+    "fr-YT": "French",
+    "fy": "Western Frisian",
+    "ga": "Irish",
+    "ga-GB": "Irish",
+    "gd": "Scottish Gaelic",
+    "gl": "Galician",
+    "gn": "Guarani",
+    "gu": "Gujarati",
+    "gu-IN": "Gujarati",
+    "gv": "Manx",
+    "ha": "Hausa",
+    "ha-Arab": "Hausa",
+    "ha-Arab-SD": "Hausa",
+    "ha-GH": "Hausa",
+    "ha-NE": "Hausa",
+    "he": "Hebrew",
+    "he-IL": "Hebrew",
+    "hi": "Hindi",
+    "hi-IN": "Hindi",
+    "hi-Latn": "Hindi",
+    "hr": "Croatian",
+    "hr-BA": "Croatian",
+    "hr-HR": "Croatian",
+    "ht": "Haitian",
+    "hu": "Hungarian",
+    "hu-HU": "Hungarian",
+    "hy": "Armenian",
+    "ia": "Interlingua",
+    "id": "Indonesian",
+    "id-ID": "Indonesian",
+    "ie": "Interlingue",
+    "ig": "Igbo",
+    "ii": "Sichuan Yi",
+    "ik": "Inupiaq",
+    "io": "Ido",
+    "is": "Icelandic",
+    "it": "Italian",
+    "it-CH": "Italian",
+    "it-IT": "Italian",
+    "it-SM": "Italian",
+    "it-VA": "Italian",
+    "iu": "Inuktitut",
+    "iu-Latn": "Inuktitut",
+    "ja": "Japanese",
+    "ja-JP": "Japanese",
+    "jv": "Javanese",
+    "ka": "Georgian",
+    "ki": "Kikuyu",
+    "kk": "Kazakh",
+    "kk-Arab": "Kazakh",
+    "kk-Cyrl": "Kazakh",
+    "kk-KZ": "Kazakh",
+    "kl": "Kalaallisut",
+    "km": "Central Khmer",
+    "kn": "Kannada",
+    "kn-IN": "Kannada",
+    "ko": "Korean",
+    "ko-CN": "Korean",
+    "ko-KP": "Korean",
+    "ko-KR": "Korean",
+    "ks": "Kashmiri",
+    "ks-Arab": "Kashmiri",
+    "ks-Deva": "Kashmiri",
+    "ku": "Kurdish",
+    "kw": "Cornish",
+    "ky": "Kyrgyz",
+    "la": "Latin",
+    "lb": "Luxembourgish",
+    "lg": "Ganda",
+    "ln": "Lingala",
+    "ln-AO": "Lingala",
+    "ln-CF": "Lingala",
+    "ln-CG": "Lingala",
+    "lo": "Lao",
+    "lt": "Lithuanian",
+    "lt-LT": "Lithuanian",
+    "lu": "Luba-Katanga",
+    "lv": "Latvian",
+    "lv-LV": "Latvian",
+    "mg": "Malagasy",
+    "mi": "Maori",
+    "mk": "Macedonian",
+    "ml": "Malayalam",
+    "ml-IN": "Malayalam",
+    "mn": "Mongolian",
+    "mn-Mong": "Mongolian",
+    "mn-Mong-MN": "Mongolian",
+    "mr": "Marathi",
+    "mr-IN": "Marathi",
+    "ms": "Malay",
+    "ms-Arab": "Malay",
+    "ms-Arab-BN": "Malay",
+    "ms-BN": "Malay",
+    "ms-ID": "Malay",
+    "ms-SG": "Malay",
+    "mt": "Maltese",
+    "my": "Burmese",
+    "nb": "Norwegian Bokmål",
+    "nb-SJ": "Norwegian Bokmål",
+    "nd": "North Ndebele",
+    "ne": "Nepali",
+    "ne-IN": "Nepali",
+    "nl": "Dutch",
+    "nl-AW": "Dutch",
+    "nl-BE": "Dutch",
+    "nl-BQ": "Dutch",
+    "nl-CW": "Dutch",
+    "nl-NL": "Dutch",
+    "nl-SR": "Dutch",
+    "nl-SX": "Dutch",
+    "nn": "Norwegian Nynorsk",
+    "no": "Norwegian",
+    "no-NO": "Norwegian",
+    "nr": "South Ndebele",
+    "nv": "Navajo",
+    "ny": "Chichewa",
+    "oc": "Occitan",
+    "oc-ES": "Occitan",
+    "om": "Oromo",
+    "om-KE": "Oromo",
+    "or": "Oriya",
+    "os": "Ossetian",
+    "os-RU": "Ossetian",
+    "pa": "Punjabi",
+    "pa-IN": "Punjabi",
+    "pa-Arab": "Punjabi",
+    "pa-Guru": "Punjabi",
+    "pl": "Polish",
+    "pl-PL": "Polish",
+    "ps": "Pashto",
+    "ps-PK": "Pashto",
+    "pt": "Portuguese",
+    "pt-AO": "Portuguese",
+    "pt-BR": "Portuguese",
+    "pt-CH": "Portuguese",
+    "pt-CV": "Portuguese",
+    "pt-GQ": "Portuguese",
+    "pt-GW": "Portuguese",
+    "pt-LU": "Portuguese",
+    "pt-MO": "Portuguese",
+    "pt-MZ": "Portuguese",
+    "pt-PT": "Portuguese",
+    "pt-ST": "Portuguese",
+    "pt-TL": "Portuguese",
+    "qu": "Quechua",
+    "qu-BO": "Quechua",
+    "qu-EC": "Quechua",
+    "rm": "Romansh",
+    "rn": "Rundi",
+    "ro": "Romanian",
+    "ro-MD": "Romanian",
+    "ro-RO": "Romanian",
+    "ru": "Russian",
+    "ru-BY": "Russian",
+    "ru-KG": "Russian",
+    "ru-KZ": "Russian",
+    "ru-MD": "Russian",
+    "ru-RU": "Russian",
+    "ru-UA": "Russian",
+    "rw": "Kinyarwanda",
+    "sa": "Sanskrit",
+    "sc": "Sardinian",
+    "sd": "Sindhi",
+    "sd-Arab": "Sindhi",
+    "sd-Deva": "Sindhi",
+    "se": "Northern Sami",
+    "se-FI": "Northern Sami",
+    "se-SE": "Northern Sami",
+    "sg": "Sango",
+    "si": "Sinhala",
+    "sk": "Slovak",
+    "sk-SK": "Slovak",
+    "sl": "Slovenian",
+    "sl-SI": "Slovenian",
+    "sn": "Shona",
+    "so": "Somali",
+    "so-DJ": "Somali",
+    "so-ET": "Somali",
+    "so-KE": "Somali",
+    "sq": "Albanian",
+    "sq-MK": "Albanian",
+    "sq-XK": "Albanian",
+    "sr": "Serbian",
+    "sr-RS": "Serbian",
+    "sr-Cyrl": "Serbian",
+    "sr-Cyrl-BA": "Serbian",
+    "sr-Cyrl-ME": "Serbian",
+    "sr-Cyrl-XK": "Serbian",
+    "sr-Latn": "Serbian",
+    "sr-Latn-BA": "Serbian",
+    "sr-Latn-ME": "Serbian",
+    "sr-Latn-XK": "Serbian",
+    "ss": "Swati",
+    "ss-SZ": "Swati",
+    "st": "Southern Sotho",
+    "st-LS": "Southern Sotho",
+    "su": "Sundanese",
+    "su-Latn": "Sundanese",
+    "sv": "Swedish",
+    "sv-AX": "Swedish",
+    "sv-FI": "Swedish",
+    "sv-SE": "Swedish",
+    "sw": "Swahili",
+    "sw-CD": "Swahili",
+    "sw-KE": "Swahili",
+    "sw-TZ": "Swahili",
+    "sw-UG": "Swahili",
+    "ta": "Tamil",
+    "ta-IN": "Tamil",
+    "ta-LK": "Tamil",
+    "ta-MY": "Tamil",
+    "ta-SG": "Tamil",
+    "te": "Telugu",
+    "te-IN": "Telugu",
+    "tg": "Tajik",
+    "th": "Thai",
+    "th-TH": "Thai",
+    "ti": "Tigrinya",
+    "ti-ER": "Tigrinya",
+    "tk": "Turkmen",
+    "tl": "Tagalog",
+    "tn": "Tswana",
+    "tn-BW": "Tswana",
+    "to": "Tonga",
+    "tr": "Turkish",
+    "tr-CY": "Turkish",
+    "tr-TR": "Turkish",
+    "ts": "Tsonga",
+    "tt": "Tatar",
+    "ug": "Uyghur",
+    "uk": "Ukrainian",
+    "uk-UA": "Ukrainian",
+    "ur": "Urdu",
+    "ur-IN": "Urdu",
+    "ur-PK": "Urdu",
+    "uz": "Uzbek",
+    "uz-Arab": "Uzbek",
+    "uz-Cyrl": "Uzbek",
+    "uz-Latn": "Uzbek",
+    "ve": "Venda",
+    "vi": "Vietnamese",
+    "vi-VN": "Vietnamese",
+    "vo": "Volapük",
+    "wa": "Walloon",
+    "wo": "Wolof",
+    "xh": "Xhosa",
+    "yi": "Yiddish",
+    "yo": "Yoruba",
+    "yo-BJ": "Yoruba",
+    "za": "Zhuang",
+    "zh": "Chinese",
+    "zh-CH": "Chinese",
+    "zh-TW": "Chinese",
+    "zh-Hans": "Chinese",
+    "zh-Hans-HK": "Chinese",
+    "zh-Hans-MO": "Chinese",
+    "zh-Hans-MY": "Chinese",
+    "zh-Hans-SG": "Chinese",
+    "zh-Hant": "Chinese",
+    "zh-Hant-HK": "Chinese",
+    "zh-Hant-MO": "Chinese",
+    "zh-Hant-MY": "Chinese",
+    "zh-Latn": "Chinese",
+    "zu": "Zulu",
+    "zu-ZA": "Zulu",
+}
+-%}
+{{ bos_token }}
+{%- if (messages[0]['role'] != 'user') -%}
+    {{ raise_exception("Conversations must start with a user prompt.") }}
+{%- endif -%}
+{%- for message in messages -%}
+    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%}
+        {{ raise_exception("Conversation roles must alternate user/assistant/user/assistant/...") }}
+    {%- endif -%}
+    {%- if (message['role'] == 'assistant') -%}
+        {%- if message['content'] is none or message['content'] is not string -%}
+            {{ raise_exception("Assistant role must provide content as a string") }}
+        {%- endif -%}
+        {{ '<start_of_turn>model\n'}}
+        {{ message["content"] | trim }}
+    {%- elif (message['role'] == 'user') -%}
+        {%- if message['content'] is none or message['content'] is not iterable or message['content'] | length != 1 -%}
+            {{ raise_exception(
+                "User role must provide `content` as an iterable with exactly one item. That item must be a " +
+                "`mapping(type:'text' | 'image', source_lang_code:string, target_lang_code:string, " +
+                "text:string | none, image:string | none)`."
+            ) }}
+        {%- endif -%}
+        {%- set content = message["content"][0] -%}
+        {%- set source_lang_code = content["source_lang_code"] | replace("_", "-")  -%}
+        {%- set source_lang = languages[source_lang_code] -%}
+        {%- set target_lang_code = content["target_lang_code"] | replace("_", "-") -%}
+        {%- set target_lang = languages[target_lang_code] -%}
+        {{ '<start_of_turn>user\nYou are a professional ' + source_lang + ' (' + source_lang_code + ') to ' +
+           target_lang + ' (' + target_lang_code + ') translator. Your goal is to accurately convey the meaning and '
+           'nuances of the original ' + source_lang + ' text while adhering to ' + target_lang + ' grammar, '
+           'vocabulary, and cultural sensitivities.\n'
+        }}
+        {%- if content["type"] == 'text' -%}
+            {{
+                'Produce only the ' + target_lang + ' translation, without any additional explanations or ' +
+                'commentary. Please translate the following ' + source_lang + ' text into ' + target_lang + ':\n\n\n' +
+                content["text"] | trim
+            }}
+        {%- elif content["type"] == 'image' -%}
+            {{
+                'Please translate the ' + source_lang + ' text in the provided image into ' + target_lang + '. ' +
+                'Produce only the ' + target_lang + ' translation, without any additional explanations, ' +
+                'alternatives or commentary. Focus only on the text, do not output where the text is located, ' +
+                'surrounding objects or any other explanation about the picture. Ignore symbols, pictogram, and ' +
+                'arrows!\n\n\n<start_of_image>'
+            }}
+        {%- else -%}
+            {{ raise_exception("User role `content` can contain either `text` or `image` data.") }}
+        {%- endif -%}
+    {%- else -%}
+        {{ raise_exception("Conversations must only contain user or assistant roles.") }}
+    {%- endif -%}
+    {{ '<end_of_turn>\n' }}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    {{'<start_of_turn>model\n'}}
+{%- endif -%}

scripts/snapdragon/run.py

@@ -152,6 +152,7 @@ def main():
     parser.add_argument("--profile", help="Profiling flag (enables Hexagon profiling and OpenCL autotuning)")
     parser.add_argument("--sched-debug", action="store_true", help="Enable GGML/llama.cpp scheduler debug output (GGML_SCHED_DEBUG=2)")
     parser.add_argument("--mtmd-device", help="Specify the backend device ID for Multi-Threaded Multi-Device setup (MTMD_BACKEND_DEVICE)")
+    parser.add_argument("--no-embd-offload", action="store_true", help="Keep token embeddings and output projection on CPU (-ot token_embd.weight=CPU,output.weight=CPU)")
 
     # Hexagon specific parameters
     parser.add_argument("--hex-verbose", help="Enable verbose logging (GGML_HEXAGON_VERBOSE)")
@@ -166,7 +167,7 @@ def main():
     parser.add_argument("--hex-opfilter", help="Regex pattern to filter/select which operators are offloaded to NPU (GGML_HEXAGON_OPFILTER)")
     parser.add_argument("--hex-opfusion", help="NPU graph node fusion optimization level (0: disabled, 1: enabled) (GGML_HEXAGON_OPFUSION)")
     parser.add_argument("--hex-vmem", help="Maximum NPU VMEM size limit in MB to allocate (GGML_HEXAGON_VMEM)")
-    parser.add_argument("--hex-mbuf", help="Maximum host buffer size limit in MB to allocate (GGML_HEXAGON_MBUF)")
+    parser.add_argument("--hex-mbuf", help="Host buffer size limits in MB (supports K/M/G suffix): <dyn>[,<static>[,<total>]] (default: 512,1024,0) (GGML_HEXAGON_MBUF)")
     parser.add_argument("--hex-mm-select", help="Select MUL_MAT and MUL_MAT_ID kernel (GGML_HEXAGON_MM_SELECT) 2:HMX,1:HVX,0:disable")
     parser.add_argument("--hex-fa-select", help="Select Flash Attention kernel (GGML_HEXAGON_FA_SELECT) 2:HMX,1:HVX,0:disable")
     parser.add_argument("--hex-fa-head-split", help="Enable (1) or disable (0) head-parallel flash_attn partitioning (GGML_HEXAGON_FA_HEAD_SPLIT)")
@@ -422,6 +423,9 @@ def main():
     if basename in ("llama-cli", "llama-completion", "llama-server", "llama-bench"):
         if "-t" not in cmd_args and "--threads" not in cmd_args:
             cmd_args += ["-t", "6"]
+        if getattr(args, "no_embd_offload", False):
+            if not any("token_embd" in arg or "output.weight" in arg for arg in cmd_args):
+                cmd_args += ["-ot", r"^(token_embd|output)\.weight$=CPU"]
 
     # Resolve target directory on device
     target_dir = args.target_dir

scripts/sync_vendor.py

@@ -5,7 +5,7 @@ import os
 import sys
 import subprocess
 
-HTTPLIB_VERSION = "refs/tags/v0.59.0"
+HTTPLIB_VERSION = "refs/tags/v0.60.0"
 
 # used by examples/gguf-hash, these repos have no release tag, so we pin a commit
 XXHASH_COMMIT      = "9f465f1ea932d6ad9a26cd77496311ffa544cd68"

src/CMakeLists.txt

@@ -28,6 +28,7 @@ set(LLAMA_CORE_SOURCES
     llama-kv-cache-msa.cpp
     llama-kv-cache-dsv4.cpp
     llama-memory.cpp
+    llama-moe-cache.cpp
     llama-memory-hybrid.cpp
     llama-memory-hybrid-iswa.cpp
     llama-memory-hybrid-idx.cpp

src/llama-context.cpp

@@ -9,6 +9,7 @@
 #include "llama-memory.h"
 #include "llama-mmap.h"
 #include "llama-model.h"
+#include "llama-moe-cache.h"
 #include "llama-ext.h"
 #include "llama-sampler.h"
 #include "llama.h"
@@ -272,8 +273,9 @@ llama_context::llama_context(
         }
     }
 
-    cparams.op_offload = params.op_offload;
-    cparams.kv_unified = params.kv_unified;
+    cparams.op_offload     = params.op_offload;
+    cparams.kv_unified     = params.kv_unified;
+    cparams.moe_cache_size = params.moe_cache_size;
 
     // initialized later
     cparams.pipeline_parallel = false;
@@ -462,6 +464,22 @@ llama_context::llama_context(
             LLAMA_LOG_INFO("%s: pipeline parallelism enabled\n", __func__);
         }
 
+        if (cparams.moe_cache_size > 0) {
+            if (cparams.pipeline_parallel || model.n_devices() > 1) {
+                throw std::runtime_error("MoE cache does not support multiple devices");
+            }
+            for (size_t i = 0; i < backend_ptrs.size(); ++i) {
+                const auto type = ggml_backend_dev_type(ggml_backend_get_device(backend_ptrs[i]));
+                if (type == GGML_BACKEND_DEVICE_TYPE_GPU || type == GGML_BACKEND_DEVICE_TYPE_IGPU) {
+                    moe_cache = std::make_unique<llama_moe_cache>(model, backend_ptrs[i], backend_buft[i], cparams.moe_cache_size);
+                    break;
+                }
+            }
+            if (!moe_cache) {
+                throw std::runtime_error("MoE cache requires a GPU backend");
+            }
+        }
+
         sched_reserve();
 
         if (!cparams.flash_attn) {
@@ -2605,6 +2623,7 @@ llm_graph_params llama_context::graph_params(
         /*.loras       =*/ loras.get(),
         /*.mctx        =*/ mctx,
         /*.cross       =*/ &cross,
+        /*.moe_cache   =*/ moe_cache.get(),
         /*.prec_policy =*/ &model.prec_policy,
         /*.samplers    =*/ sampling.samplers,
         /*.n_outputs   =*/ n_outputs,
@@ -2645,7 +2664,14 @@ ggml_status llama_context::graph_compute(
 }
 
 bool llama_context::sched_copy_experts(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph, void * user_data) {
-    auto & st = static_cast<llama_context *>(user_data)->copy_experts;
+    auto * lctx = static_cast<llama_context *>(user_data);
+
+    // the slot maps of the MoE cache
+    if (lctx->moe_cache && lctx->moe_cache->copy(backend, src, dst, graph)) {
+        return true;
+    }
+
+    auto & st = lctx->copy_experts;
 
     // the ids must be computed before the split starts, so only the first node of the split is considered
     if (ggml_graph_n_nodes(graph) == 0) {
@@ -2692,10 +2718,28 @@ bool llama_context::sched_copy_experts(ggml_backend_t backend, const ggml_tensor
             last++;
         }
 
+        // the experts in the MoE cache are copied from device memory, the others are uploaded
+        int64_t next = first;
+        for (int64_t e = first; e <= last && lctx->moe_cache; ) {
+            const int64_t n = lctx->moe_cache->copy_experts(backend, src, dst, e, last);
+            if (n == 0) {
+                e++;
+                continue;
+            }
+            if (next < e) {
+                ggml_backend_tensor_set_async(backend, dst, (const uint8_t *) src->data + next*expert_size, next*expert_size, (e - next)*expert_size);
+            }
+            e   += n;
+            next = e;
+        }
+
         // copy a bit extra to ensure there are no NaNs in the padding of the last expert, this is necessary for MMQ in the CUDA backend
-        const size_t offset  = first*expert_size;
+        const size_t offset  = next*expert_size;
         const size_t padding = last < n_expert - 1 ? std::min<size_t>(expert_size, 512) : 0;
-        ggml_backend_tensor_set_async(backend, dst, (const uint8_t *) src->data + offset, offset, (last - first + 1)*expert_size + padding);
+        const size_t size    = (last + 1 - next)*expert_size + padding;
+        if (size > 0) {
+            ggml_backend_tensor_set_async(backend, dst, (const uint8_t *) src->data + offset, offset, size);
+        }
 
         first = last + 1;
     }
@@ -3562,6 +3606,11 @@ llama_memory_breakdown llama_context::memory_breakdown() const {
             ret[buft].context += size;
         }
     }
+    if (moe_cache) {
+        for (const auto & [buft, size] : moe_cache->memory_breakdown()) {
+            ret[buft].context += size;
+        }
+    }
     if (model.hparams.no_alloc) {
         for (size_t i = 0; i < backends.size(); ++i) {
             ggml_backend_t             backend = backends[i].get();
@@ -3851,6 +3900,7 @@ llama_context_params llama_context_default_params() {
         /*.cb_eval_user_data           =*/ nullptr,
         /*.type_k                      =*/ GGML_TYPE_F16,
         /*.type_v                      =*/ GGML_TYPE_F16,
+        /*.moe_cache_size              =*/ 0,
         /*.abort_callback              =*/ nullptr,
         /*.abort_callback_data         =*/ nullptr,
         /*.embeddings                  =*/ false,

src/llama-context.h

@@ -7,6 +7,7 @@
 #include "llama-adapter.h"
 #include "llama-impl.h"
 #include "llama-memory.h"
+#include "llama-moe-cache.h"
 
 #include "ggml-cpp.h"
 #include "ggml-opt.h"
@@ -17,6 +18,7 @@
 
 struct llama_model;
 class llama_batch_allocr;
+class llama_moe_cache;
 
 class llama_io_read_i;
 class llama_io_write_i;
@@ -271,7 +273,7 @@ private:
 
     llm_graph_cb graph_get_cb() const;
 
-    // ggml_backend_sched copy callback, copies only the experts used by MUL_MAT_ID
+    // ggml_backend_sched copy callback, copies only the experts used by MUL_MAT_ID and updates the MoE cache
     static bool sched_copy_experts(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph, void * user_data);
 
     // disable auto fused ops (Flash Attention, Gated Delta Net) whose op lands on a device
@@ -299,6 +301,7 @@ private:
     llama_cross cross; // TODO: tmp for handling cross-attention - need something better probably
 
     llama_memory_ptr memory;
+    llama_moe_cache_ptr moe_cache;
 
     // decode output (2-dimensional array: [n_outputs][n_vocab])
     buffer_view<float> logits = {nullptr, 0};

src/llama-cparams.h

@@ -55,6 +55,8 @@ struct llama_cparams {
     bool pipeline_parallel;
     bool training;           // set by llama_opt_init()
 
+    size_t moe_cache_size;
+
     std::vector<bool> embeddings_layer_inp; // [n_layer()] extract input embeddings for layer
 
     enum llama_context_type ctx_type;

src/llama-graph.cpp

@@ -2,6 +2,7 @@
 
 #include "llama-impl.h"
 #include "llama-model.h"
+#include "llama-moe-cache.h"
 #include "llama-batch.h"
 #include "llama-cparams.h"
 #include "llama-sampler.h"
@@ -1523,6 +1524,7 @@ llm_graph_context::llm_graph_context(const llm_graph_params & params) :
     loras            (params.loras),
     mctx             (params.mctx),
     cross            (params.cross),
+    moe_cache        (params.moe_cache),
     prec_policy      (params.prec_policy),
     samplers         (params.samplers),
     cb_func          (params.cb),
@@ -1589,8 +1591,12 @@ ggml_tensor * llm_graph_context::build_lora_mm_id(
           ggml_tensor * w,   // ggml_tensor * as
           ggml_tensor * cur, // ggml_tensor * b
           ggml_tensor * ids,
-          ggml_tensor * w_s) const {
-    ggml_tensor * res = ggml_mul_mat_id(ctx0, w, cur, ids);
+          ggml_tensor * w_s,
+          ggml_tensor * slots) const {
+    // the experts in the MoE cache are selected by their slots
+    ggml_tensor * res = slots == nullptr ?
+        ggml_mul_mat_id(ctx0, w, cur, ids) :
+        ggml_mul_mat_id(ctx0, moe_cache->get_experts(w), cur, slots);
 
     if (prec_policy) {
         prec_policy->apply(res);
@@ -2205,6 +2211,9 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
     //call early so that topk-moe can be used
     ggml_build_forward_expand(gf, weights);
 
+    // the experts of host-resident layers may be read from the MoE cache
+    ggml_tensor * slots = build_moe_cache_slots(selected_experts, up_exps, gate_exps, down_exps, gate_up_exps, il);
+
     cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens);
 
     if (weight_before_ffn) {
@@ -2219,7 +2228,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
 
     if (gate_up_exps) {
         // merged gate_up path: one mul_mat_id, then split into gate and up views
-        ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, selected_experts, up_exps_s); // [n_ff*2, n_expert_used, n_tokens]
+        ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, selected_experts, up_exps_s, slots); // [n_ff*2, n_expert_used, n_tokens]
         cb(gate_up, "ffn_moe_gate_up", il);
 
         if (up_exps_s) {
@@ -2238,7 +2247,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
         cb(up, "ffn_moe_up", il);
     } else {
         // separate gate and up path
-        up = build_lora_mm_id(up_exps, cur, selected_experts, up_exps_s); // [n_ff, n_expert_used, n_tokens]
+        up = build_lora_mm_id(up_exps, cur, selected_experts, up_exps_s, slots); // [n_ff, n_expert_used, n_tokens]
         cb(up, "ffn_moe_up", il);
 
         if (up_exps_s) {
@@ -2251,7 +2260,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
         }
 
         if (gate_exps) {
-            cur = build_lora_mm_id(gate_exps, cur, selected_experts, gate_exps_s); // [n_ff, n_expert_used, n_tokens]
+            cur = build_lora_mm_id(gate_exps, cur, selected_experts, gate_exps_s, slots); // [n_ff, n_expert_used, n_tokens]
             cb(cur, "ffn_moe_gate", il);
         } else {
             cur = up;
@@ -2352,7 +2361,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
             GGML_ABORT("fatal error");
     }
 
-    experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s); // [n_embd, n_expert_used, n_tokens]
+    experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s, slots); // [n_embd, n_expert_used, n_tokens]
     if (arch == LLM_ARCH_MISTRAL4) {
         // src1 can exceed F16 range
         ggml_prec_set_src(experts, GGML_PREC_F32, 1);
@@ -2409,6 +2418,45 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
     return moe_out;
 }
 
+ggml_tensor * llm_graph_context::build_moe_cache_slots(
+         ggml_tensor * selected_experts,
+         ggml_tensor * up_exps,
+         ggml_tensor * gate_exps,
+         ggml_tensor * down_exps,
+         ggml_tensor * gate_up_exps,
+                 int   il) const {
+    if (moe_cache == nullptr) {
+        return nullptr;
+    }
+
+    ggml_tensor * slot_map = moe_cache->get_slot_map(il, selected_experts->ne[1], selected_experts->ne[0]);
+    if (slot_map == nullptr) {
+        return nullptr;
+    }
+    for (ggml_tensor * w : { up_exps, gate_exps, down_exps, gate_up_exps }) {
+        if (w != nullptr && moe_cache->get_experts(w) == nullptr) {
+            return nullptr;
+        }
+    }
+
+    ggml_tensor * ids = selected_experts;
+    if (!ggml_is_contiguous(ids)) {
+        ids = ggml_cont(ctx0, ids);
+    }
+    ids = ggml_reshape_1d(ctx0, ids, ggml_nelements(ids));
+
+    // the slot map is a host weight, so the scheduler starts a new split here and copies it with the copy callback
+    // the callback reads the selected experts, uploads the missing ones and updates the slot map
+    ggml_tensor * slots = ggml_get_rows(ctx0, slot_map, ids); // [1, n_expert_used*n_tokens]
+    if (!ggml_backend_supports_op(moe_cache->backend(), slots)) {
+        return nullptr;
+    }
+    ggml_backend_sched_set_tensor_backend(sched, slots, moe_cache->backend());
+    cb(slots, "ffn_moe_slots", il);
+
+    return ggml_reshape_2d(ctx0, slots, selected_experts->ne[0], selected_experts->ne[1]); // [n_expert_used, n_tokens]
+}
+
 // input embeddings with optional lora
 ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd, float tok_scale) const {
     const int64_t n_embd_inp = hparams.n_embd_inp();

src/llama-graph.h

@@ -21,6 +21,8 @@ struct llama_cparams;
 struct llama_layer;
 struct llama_prec_policy;
 
+class llama_moe_cache;
+
 struct llama_memory_context_i;
 
 class llama_kv_cache_context;
@@ -793,6 +795,7 @@ struct llm_graph_params {
     const llama_adapter_loras    * loras;
     const llama_memory_context_i * mctx;
     const llama_cross            * cross;
+    const llama_moe_cache        * moe_cache;
 
     const llama_prec_policy * prec_policy = nullptr;
 
@@ -1036,6 +1039,7 @@ struct llm_graph_context {
     const llama_adapter_loras    * loras;
     const llama_memory_context_i * mctx;
     const llama_cross            * cross;
+    const llama_moe_cache        * moe_cache;
 
     const llama_prec_policy * prec_policy;
 
@@ -1078,11 +1082,13 @@ struct llm_graph_context {
               ggml_tensor * w_s = nullptr) const;
 
     // do mat_mul_id, while optionally apply lora and per-expert scale
+    // if slots is set, the experts are read from the MoE cache at these slots (see build_moe_cache_slots)
     ggml_tensor * build_lora_mm_id(
               ggml_tensor * w,   // ggml_tensor * as
               ggml_tensor * cur, // ggml_tensor * b
               ggml_tensor * ids,
-              ggml_tensor * w_s = nullptr) const;
+              ggml_tensor * w_s   = nullptr,
+              ggml_tensor * slots = nullptr) const;
 
     ggml_tensor * build_norm(
              ggml_tensor * cur,
@@ -1179,6 +1185,15 @@ struct llm_graph_context {
              ggml_tensor * down_exps_s = nullptr,
              ggml_tensor * selected_experts_in = nullptr) const;
 
+    // the slots of the selected experts in the MoE cache, nullptr if the experts of the layer are not read from the cache
+    ggml_tensor * build_moe_cache_slots(
+             ggml_tensor * selected_experts,
+             ggml_tensor * up_exps,
+             ggml_tensor * gate_exps,
+             ggml_tensor * down_exps,
+             ggml_tensor * gate_up_exps,
+                     int   il) const;
+
     //
     // inputs
     //

src/llama-memory-recurrent.cpp

@@ -199,6 +199,10 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos
 
             // partial rollback via per-token snapshot index (bounded by n_rs_seq)
             if (0 < p0 && p0 <= cell.pos && p1 > cell.pos) {
+                // a cell shared by several sequences cannot move back for only one of them
+                if (cell.seq_id.size() > 1) {
+                    return false;
+                }
                 // the filter kept no layer (e.g. an MTP draft context), so only the position moves back
                 if (is_empty()) {
                     cell.pos = p0 - 1;

src/llama-model.cpp

@@ -3461,6 +3461,27 @@ void llama_model_base::create_tensor_qkv(llama_layer & layer, int bid,
     }
 }
 
+llama_model_base::nextn_flags_t llama_model_base::nextn_flags(llama_model_loader & ml, llm_tensor trunk_probe) const {
+    int trunk = 0;
+    int mtp   = 0;
+
+    // a file without the first trunk layer is MTP-only, a file without the first NextN layer is trunk-only
+    if (hparams.n_layer_nextn > 0) {
+        if (ml.get_weight(tn(trunk_probe, "weight", 0).str().c_str()) == nullptr) {
+            trunk = TENSOR_NOT_REQUIRED;
+        }
+        if (ml.get_weight(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", hparams.n_layer()).str().c_str()) == nullptr) {
+            mtp = TENSOR_NOT_REQUIRED;
+        }
+    }
+
+    if (!ml.load_mtp) {
+        mtp |= TENSOR_SKIP;
+    }
+
+    return { trunk, mtp };
+}
+
 void llama_model_base::load_swa_pattern(llama_model_loader & ml, uint32_t n_pattern, bool dense_first) {
     if (ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, false)) {
         return;

src/llama-model.h

@@ -849,6 +849,13 @@ struct llama_model_base : public llama_model {
                 int64_t n_embd_, int64_t n_embd_q_, int64_t n_embd_k_, int64_t n_embd_v_,
                 int flags);
 
+    // helper: tensor flags for a file that holds the trunk, the NextN layers, or both
+    struct nextn_flags_t {
+        int trunk; // TENSOR_NOT_REQUIRED when the file holds only the NextN layers
+        int mtp;   // TENSOR_NOT_REQUIRED when the file holds only the trunk, TENSOR_SKIP when MTP is not loaded
+    };
+    nextn_flags_t nextn_flags(llama_model_loader & ml, llm_tensor trunk_probe = LLM_TENSOR_ATTN_NORM) const;
+
     // helper: read the SWA pattern as one flag per layer, or as a period expanded by set_swa_pattern
     void load_swa_pattern(llama_model_loader & ml, uint32_t n_pattern, bool dense_first = false);
 

src/llama-moe-cache.cpp

@@ -0,0 +1,550 @@
+#include "llama-moe-cache.h"
+
+#include "llama-impl.h"
+#include "llama-model.h"
+
+#include "ggml-cpp.h"
+
+#include <algorithm>
+#include <stdexcept>
+#include <unordered_map>
+#include <vector>
+
+namespace {
+
+// LRU of the experts of a group of layers, the slot of each expert is kept in the slot map of its layer
+struct moe_cache_lru {
+    int32_t n_expert = 0;
+    int32_t n_slots  = 0;
+
+    std::vector<int32_t *> slot_map; // [n_layer] data of the slot maps, -1 if the expert is not cached
+    std::vector<int32_t>   key_of;   // [n_slots] il*n_expert + expert, -1 if empty
+
+    // doubly linked list of the slots, head is the least recently used
+    std::vector<int32_t> prev;
+    std::vector<int32_t> next;
+    int32_t head = -1;
+    int32_t tail = -1;
+
+    std::vector<uint32_t> seen; // [n_expert]
+    uint32_t seen_gen = 0;
+    std::vector<int32_t> uniq;
+
+    void init(int32_t n_layer, int32_t n_expert, int32_t n_slots) {
+        this->n_expert = n_expert;
+        this->n_slots  = n_slots;
+        slot_map.assign(n_layer, nullptr);
+        key_of.assign(n_slots, -1);
+        prev.resize(n_slots);
+        next.resize(n_slots);
+        for (int32_t s = 0; s < n_slots; ++s) {
+            prev[s] = s - 1;
+            next[s] = s + 1 < n_slots ? s + 1 : -1;
+        }
+        head = 0;
+        tail = n_slots - 1;
+        seen.assign(n_expert, 0);
+    }
+
+    // move slot s to the tail (most recently used)
+    void touch(int32_t s) {
+        if (s == tail) {
+            return;
+        }
+        if (prev[s] >= 0) {
+            next[prev[s]] = next[s];
+        } else {
+            head = next[s];
+        }
+        prev[next[s]] = prev[s];
+
+        prev[s] = tail;
+        next[s] = -1;
+        next[tail] = s;
+        tail = s;
+    }
+
+    struct fill {
+        int32_t expert;
+        int32_t slot;
+    };
+
+    // give a slot to each expert selected by ids in layer il, the misses evict the least recently used experts
+    // returns false if the ids select more distinct experts than there are slots
+    bool plan(int32_t il, const int32_t * ids, size_t n_ids, std::vector<fill> & fills, size_t & n_hit) {
+        fills.clear();
+        n_hit = 0;
+
+        if (++seen_gen == 0) {
+            std::fill(seen.begin(), seen.end(), 0);
+            seen_gen = 1;
+        }
+        uniq.clear();
+        for (size_t i = 0; i < n_ids; ++i) {
+            GGML_ASSERT(ids[i] >= 0 && ids[i] < n_expert);
+            if (seen[ids[i]] != seen_gen) {
+                seen[ids[i]] = seen_gen;
+                uniq.push_back(ids[i]);
+            }
+        }
+        if (uniq.size() > (size_t) n_slots) {
+            return false;
+        }
+
+        int32_t * slots = slot_map[il];
+
+        // hits go to the tail first, so the head can be evicted below
+        for (int32_t e : uniq) {
+            if (slots[e] >= 0) {
+                touch(slots[e]);
+                n_hit++;
+            }
+        }
+        // sorted misses usually get consecutive slots, so the uploads can be merged
+        std::sort(uniq.begin(), uniq.end());
+        for (int32_t e : uniq) {
+            if (slots[e] >= 0) {
+                continue;
+            }
+            const int32_t s = head;
+            if (key_of[s] >= 0) {
+                slot_map[key_of[s] / n_expert][key_of[s] % n_expert] = -1;
+            }
+            key_of[s] = il*n_expert + e;
+            slots[e] = s;
+            touch(s);
+            fills.push_back({ e, s });
+        }
+        return true;
+    }
+};
+
+// gate, up, down or gate_up, down
+static std::vector<ggml_tensor *> llama_moe_cache_layer_experts(const llama_layer & layer) {
+    std::vector<ggml_tensor *> res;
+    for (ggml_tensor * t : { layer.ffn_gate_up_exps, layer.ffn_gate_exps, layer.ffn_up_exps, layer.ffn_down_exps }) {
+        if (t != nullptr) {
+            res.push_back(t);
+        }
+    }
+    return res;
+}
+
+static bool llama_moe_cache_same_layout(const std::vector<ggml_tensor *> & a, const std::vector<ggml_tensor *> & b) {
+    if (a.size() != b.size()) {
+        return false;
+    }
+    for (size_t i = 0; i < a.size(); ++i) {
+        if (a[i]->type != b[i]->type || !ggml_are_same_shape(a[i], b[i]) || a[i]->nb[2] != b[i]->nb[2]) {
+            return false;
+        }
+    }
+    return true;
+}
+
+static bool llama_moe_cache_is_host_weight(const ggml_tensor * t) {
+    return t->buffer != nullptr &&
+        ggml_backend_buffer_get_usage(t->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS &&
+        ggml_backend_buffer_is_host(t->buffer);
+}
+
+}
+
+struct llama_moe_cache::impl {
+    // layers with the same expert tensor layout share the banks and the LRU of a group
+    struct group {
+        std::vector<ggml_tensor *> ref;   // expert tensors of the first layer
+        std::vector<int32_t> layers;
+        std::vector<ggml_tensor *> banks; // device storage of all slots, one per expert tensor
+        size_t host_bytes = 0;
+        int32_t n_slots = 0;
+        moe_cache_lru lru;
+    };
+
+    struct layer {
+        int32_t ig = -1;                    // -1 if the layer is not cached
+        ggml_tensor * slot_map = nullptr;   // I32 [1, n_expert] in host memory
+        std::vector<ggml_tensor *> experts; // host expert tensors, in the order of the banks
+    };
+
+    struct binding {
+        int32_t il;
+        int32_t ip;           // index of the bank
+        ggml_tensor * cached; // view of the bank used in place of the host experts
+    };
+
+    struct stats {
+        size_t hits   = 0;
+        size_t misses = 0;
+        size_t bytes  = 0;
+    };
+
+    static constexpr int64_t max_batch = 32;
+
+    ggml_backend_t backend;
+    int32_t n_expert_used;
+
+    stats stats_small; // up to 8 tokens per ubatch
+    stats stats_large;
+    stats stats_copy;  // experts copied from the cache for large batches
+
+    std::vector<group> groups;
+    std::vector<layer> layers;
+    std::unordered_map<const ggml_tensor *, binding> bindings; // host experts -> cached experts
+    std::unordered_map<const ggml_tensor *, int32_t> layer_of; // slot map -> layer
+
+    std::vector<int32_t> ids;
+    std::vector<moe_cache_lru::fill> fills;
+
+    // banks and their views on the device
+    ggml_context_ptr ctx;
+    ggml_backend_buffer_ptr buf;
+    size_t buf_size = 0;
+
+    // slot maps in host memory
+    ggml_context_ptr ctx_host;
+    ggml_backend_buffer_ptr buf_host;
+    size_t buf_host_size = 0;
+
+    // views used by copy_experts
+    ggml_context_ptr ctx_views;
+
+    impl(const llama_model & model, ggml_backend_t backend, ggml_backend_buffer_type_t buft, size_t size) :
+            backend(backend), n_expert_used(model.hparams.n_expert_used_max()), layers(model.layers.size()) {
+        ggml_backend_dev_t dev = ggml_backend_get_device(backend);
+        const auto dev_type = ggml_backend_dev_type(dev);
+        if (dev_type != GGML_BACKEND_DEVICE_TYPE_GPU && dev_type != GGML_BACKEND_DEVICE_TYPE_IGPU) {
+            throw std::runtime_error("MoE cache requires a GPU backend");
+        }
+        if (model.split_mode() == LLAMA_SPLIT_MODE_TENSOR) {
+            throw std::runtime_error("MoE cache does not support tensor parallelism");
+        }
+        if (model.hparams.n_expert == 0 || n_expert_used == 0) {
+            throw std::runtime_error("MoE cache requires a MoE model");
+        }
+
+        // only cache layers that keep all of their experts in host memory
+        size_t host_bytes = 0;
+        for (size_t il = 0; il < model.layers.size(); ++il) {
+            auto experts = llama_moe_cache_layer_experts(model.layers[il]);
+            if (experts.empty() || model.dev_layer(il) != dev ||
+                !std::all_of(experts.begin(), experts.end(), llama_moe_cache_is_host_weight)) {
+                continue;
+            }
+            auto it = std::find_if(groups.begin(), groups.end(), [&](const group & g) { return llama_moe_cache_same_layout(g.ref, experts); });
+            if (it == groups.end()) {
+                groups.emplace_back();
+                it = groups.end() - 1;
+                it->ref = experts;
+            }
+            it->layers.push_back(il);
+            for (const ggml_tensor * t : experts) {
+                it->host_bytes += ggml_nbytes(t);
+                host_bytes     += ggml_nbytes(t);
+            }
+        }
+        if (groups.empty()) {
+            LLAMA_LOG_WARN("%s: no layer has all of its experts in host memory, MoE cache is disabled\n", __func__);
+            return;
+        }
+
+        // one extra slot at the end, CUDA MMQ can read past the last expert
+        const size_t alignment = ggml_backend_buft_get_alignment(buft);
+        auto alloc_size = [&](const group & g, int32_t n_slots) {
+            size_t res = 0;
+            for (const ggml_tensor * t : g.ref) {
+                res += GGML_PAD(t->nb[2]*(n_slots + 1), alignment);
+            }
+            return res;
+        };
+
+        // split the budget by the size of the experts, so each group caches the same fraction of its experts
+        size_t n_tensors      = 0;
+        size_t n_tensors_host = 0;
+        for (group & g : groups) {
+            const int32_t n_expert  = g.ref[0]->ne[2];
+            const size_t  budget    = (size_t) ((double) size*g.host_bytes/host_bytes);
+            const int32_t max_slots = g.layers.size()*n_expert;
+            while (g.n_slots < max_slots && alloc_size(g, g.n_slots + 1) <= budget) {
+                g.n_slots++;
+            }
+            if (g.n_slots < n_expert_used) {
+                LLAMA_LOG_WARN("%s: MoE cache budget is too small for %zu layers, they are not cached\n", __func__, g.layers.size());
+                g.n_slots = 0;
+                continue;
+            }
+            g.lru.init(model.layers.size(), n_expert, g.n_slots);
+            n_tensors      += g.ref.size()*(1 + g.layers.size());
+            n_tensors_host += g.layers.size();
+        }
+        if (n_tensors == 0) {
+            throw std::runtime_error("MoE cache is too small to hold the experts of one token");
+        }
+
+        auto init_ctx = [](size_t n_tensors) {
+            ggml_init_params params = {
+                /*.mem_size   =*/ n_tensors*ggml_tensor_overhead(),
+                /*.mem_buffer =*/ nullptr,
+                /*.no_alloc   =*/ true,
+            };
+            ggml_context_ptr res(ggml_init(params));
+            if (!res) {
+                throw std::runtime_error("failed to create the MoE cache context");
+            }
+            return res;
+        };
+        ctx       = init_ctx(n_tensors);
+        ctx_host  = init_ctx(n_tensors_host);
+        ctx_views = init_ctx(2);
+
+        ggml_backend_buffer_type_t buft_host = ggml_backend_cpu_buffer_type();
+        const size_t alignment_host = ggml_backend_buft_get_alignment(buft_host);
+
+        for (size_t ig = 0; ig < groups.size(); ++ig) {
+            group & g = groups[ig];
+            if (g.n_slots == 0) {
+                continue;
+            }
+            for (const ggml_tensor * t : g.ref) {
+                ggml_tensor * bank = ggml_new_tensor_3d(ctx.get(), t->type, t->ne[0], t->ne[1], g.n_slots + 1);
+                GGML_ASSERT(bank->nb[2] == t->nb[2]);
+                ggml_format_name(bank, "moe_cache.%zu.%s", ig, t->name);
+                g.banks.push_back(bank);
+            }
+            for (int32_t il : g.layers) {
+                layer & l = layers[il];
+                l.ig      = (int32_t) ig;
+                l.experts = llama_moe_cache_layer_experts(model.layers[il]);
+                for (size_t ip = 0; ip < l.experts.size(); ++ip) {
+                    ggml_tensor * bank   = g.banks[ip];
+                    ggml_tensor * cached = ggml_view_3d(ctx.get(), bank, bank->ne[0], bank->ne[1], g.n_slots, bank->nb[1], bank->nb[2], 0);
+                    ggml_format_name(cached, "moe_cache.%s", l.experts[ip]->name);
+                    bindings[l.experts[ip]] = { il, (int32_t) ip, cached };
+                }
+                l.slot_map = ggml_new_tensor_2d(ctx_host.get(), GGML_TYPE_I32, 1, g.ref[0]->ne[2]);
+                ggml_format_name(l.slot_map, "moe_cache.slot_map-%d", il);
+                layer_of[l.slot_map] = il;
+                buf_host_size += GGML_PAD(ggml_nbytes(l.slot_map), alignment_host);
+            }
+            buf_size += alloc_size(g, g.n_slots);
+        }
+
+        if (model.hparams.no_alloc) {
+            // only used to measure the memory use, see llama_context::memory_breakdown
+            buf.reset(ggml_backend_buft_alloc_buffer(buft, 0));
+            buf_host.reset(ggml_backend_buft_alloc_buffer(buft_host, 0));
+            for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) {
+                t->buffer = buf.get();
+            }
+            for (ggml_tensor * t = ggml_get_first_tensor(ctx_host.get()); t != nullptr; t = ggml_get_next_tensor(ctx_host.get(), t)) {
+                t->buffer = buf_host.get();
+            }
+        } else {
+            buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft));
+            buf_host.reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx_host.get(), buft_host));
+            if (!buf || !buf_host) {
+                throw std::runtime_error("failed to allocate the MoE cache buffers");
+            }
+            ggml_backend_buffer_clear(buf.get(), 0);
+            ggml_backend_buffer_clear(buf_host.get(), 0xff); // all slots are -1
+            buf_size      = ggml_backend_buffer_get_size(buf.get());
+            buf_host_size = ggml_backend_buffer_get_size(buf_host.get());
+
+            for (group & g : groups) {
+                for (int32_t il : g.layers) {
+                    if (layers[il].slot_map != nullptr) {
+                        g.lru.slot_map[il] = (int32_t *) layers[il].slot_map->data;
+                    }
+                }
+            }
+        }
+
+        // as weights, the ops that read the banks run on the device and the slot maps are copied with the copy callback
+        ggml_backend_buffer_set_usage(buf.get(),      GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+        ggml_backend_buffer_set_usage(buf_host.get(), GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+
+        LLAMA_LOG_INFO("%s: %10s MoE cache size = %8.2f MiB for %.2f MiB of host experts\n", __func__,
+            ggml_backend_buft_name(buft), buf_size/1024.0/1024.0, host_bytes/1024.0/1024.0);
+        for (const group & g : groups) {
+            LLAMA_LOG_INFO("%s: %2zu layers, %s: %5d slots (%.1f%%)\n", __func__,
+                g.layers.size(), ggml_type_name(g.ref.back()->type), g.n_slots, 100.0*g.n_slots/(g.layers.size()*g.ref[0]->ne[2]));
+        }
+    }
+
+    ~impl() {
+        log_stats();
+    }
+
+    ggml_tensor * get_slot_map(int32_t il, int64_t n_tokens, int64_t n_expert_used) const {
+        if (il < 0 || il >= (int32_t) layers.size() || layers[il].ig < 0) {
+            return nullptr;
+        }
+        const layer & l = layers[il];
+
+        // large batches use most experts of a layer, so they gain little from the cache and would evict the experts used in generation
+        if (n_tokens == 0 || n_tokens > max_batch || std::min(n_tokens*n_expert_used, l.slot_map->ne[1]) > groups[l.ig].n_slots) {
+            return nullptr;
+        }
+        return l.slot_map;
+    }
+
+    ggml_tensor * get_experts(const ggml_tensor * w) const {
+        const auto it = bindings.find(w);
+        return it != bindings.end() ? it->second.cached : nullptr;
+    }
+
+    int64_t copy_experts(ggml_backend_t backend, const ggml_tensor * w, ggml_tensor * dst, int64_t e, int64_t last) {
+        const auto it = bindings.find(w);
+        if (it == bindings.end() || backend != this->backend) {
+            return 0;
+        }
+        const binding & b = it->second;
+        const group & g = groups[layers[b.il].ig];
+
+        // large batches only read the cache, so the experts used in generation stay in it
+        const int32_t * slots = g.lru.slot_map[b.il];
+        if (slots == nullptr || slots[e] < 0) {
+            return 0;
+        }
+        int64_t n = 1;
+        while (e + n <= last && slots[e + n] == slots[e] + n) {
+            n++;
+        }
+
+        ggml_tensor * bank = g.banks[b.ip];
+        ggml_reset(ctx_views.get());
+        ggml_tensor * src_view = ggml_view_3d(ctx_views.get(), bank, bank->ne[0], bank->ne[1], n, bank->nb[1], bank->nb[2], slots[e]*bank->nb[2]);
+        ggml_tensor * dst_view = ggml_view_3d(ctx_views.get(), dst,  dst->ne[0],  dst->ne[1],  n, dst->nb[1],  dst->nb[2],  e*dst->nb[2]);
+        ggml_backend_view_init(src_view);
+        ggml_backend_view_init(dst_view);
+        ggml_backend_tensor_copy_async(backend, backend, src_view, dst_view);
+
+        stats_copy.hits  += n;
+        stats_copy.bytes += ggml_nbytes(src_view);
+
+        return n;
+    }
+
+    bool copy(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph) {
+        const auto it = layer_of.find(src);
+        if (it == layer_of.end()) {
+            return false;
+        }
+        const int32_t il = it->second;
+        const layer & l = layers[il];
+        group & g = groups[l.ig];
+
+        GGML_ASSERT(backend == this->backend);
+
+        // the get_rows that looks up the slots of the selected experts
+        const int n_nodes = ggml_graph_n_nodes(graph);
+        const ggml_tensor * lookup = nullptr;
+        for (int i = 0; i < n_nodes && lookup == nullptr; ++i) {
+            const ggml_tensor * node = ggml_graph_node(graph, i);
+            if (node->op == GGML_OP_GET_ROWS && node->src[0] == dst) {
+                lookup = node;
+            }
+        }
+        GGML_ASSERT(lookup != nullptr);
+
+        // the selected experts must be computed in an earlier split
+        // the scheduler starts a new split at the lookup because it reads a host weight, but only if the split already has inputs
+        const ggml_tensor * sel = lookup->src[1];
+        for (int i = 0; i < n_nodes; ++i) {
+            const ggml_tensor * node = ggml_graph_node(graph, i);
+            if (node == sel || node == sel->view_src) {
+                GGML_ABORT("the experts of layer %d are selected in the same split as their MoE cache lookup", il);
+            }
+        }
+        GGML_ASSERT(ggml_is_contiguous(sel));
+
+        ids.resize(ggml_nelements(sel));
+        ggml_backend_tensor_get_async(backend, sel, ids.data(), 0, ggml_nbytes(sel));
+        ggml_backend_synchronize(backend);
+
+        size_t n_hit = 0;
+        if (!g.lru.plan(il, ids.data(), ids.size(), fills, n_hit)) {
+            GGML_ABORT("the MoE cache is too small for the experts selected in layer %d", il);
+        }
+
+        // upload the missing experts, consecutive experts going to consecutive slots are uploaded together
+        size_t bytes = 0;
+        for (size_t ip = 0; ip < l.experts.size(); ++ip) {
+            const ggml_tensor * w    = l.experts[ip];
+            ggml_tensor       * bank = g.banks[ip];
+            const size_t expert_size = w->nb[2];
+            for (size_t i = 0; i < fills.size();) {
+                size_t n = 1;
+                while (i + n < fills.size() && fills[i + n].expert == fills[i].expert + (int32_t) n && fills[i + n].slot == fills[i].slot + (int32_t) n) {
+                    n++;
+                }
+                ggml_backend_tensor_set_async(backend, bank, (const uint8_t *) w->data + fills[i].expert*expert_size, fills[i].slot*expert_size, n*expert_size);
+                bytes += n*expert_size;
+                i += n;
+            }
+        }
+
+        stats & st = ids.size() <= (size_t) 8*n_expert_used ? stats_small : stats_large;
+        st.hits   += n_hit;
+        st.misses += fills.size();
+        st.bytes  += bytes;
+
+        // the next copy synchronizes the backend before it changes the slot map again
+        ggml_backend_tensor_set_async(backend, dst, src->data, 0, ggml_nbytes(src));
+
+        return true;
+    }
+
+    void log_stats() const {
+        auto log = [](const char * name, const stats & st) {
+            const size_t n = st.hits + st.misses;
+            if (n == 0) {
+                return;
+            }
+            LLAMA_LOG_INFO("llama_moe_cache: %s: hits = %zu, misses = %zu, hit rate = %.2f%%, uploaded = %.2f MiB\n",
+                name, st.hits, st.misses, 100.0*st.hits/n, st.bytes/1024.0/1024.0);
+        };
+        log("ubatch <= 8", stats_small);
+        log("ubatch  > 8", stats_large);
+        if (stats_copy.hits > 0) {
+            LLAMA_LOG_INFO("llama_moe_cache: large batches: %zu experts copied from the cache, %.2f MiB\n", stats_copy.hits, stats_copy.bytes/1024.0/1024.0);
+        }
+    }
+};
+
+llama_moe_cache::llama_moe_cache(const llama_model & model, ggml_backend_t backend, ggml_backend_buffer_type_t buft, size_t size) :
+    pimpl(new impl(model, backend, buft, size)) {
+}
+
+llama_moe_cache::~llama_moe_cache() = default;
+
+ggml_backend_t llama_moe_cache::backend() const {
+    return pimpl->backend;
+}
+
+ggml_tensor * llama_moe_cache::get_slot_map(int32_t il, int64_t n_tokens, int64_t n_expert_used) const {
+    return pimpl->get_slot_map(il, n_tokens, n_expert_used);
+}
+
+ggml_tensor * llama_moe_cache::get_experts(const ggml_tensor * w) const {
+    return pimpl->get_experts(w);
+}
+
+bool llama_moe_cache::copy(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph) {
+    return pimpl->copy(backend, src, dst, graph);
+}
+
+int64_t llama_moe_cache::copy_experts(ggml_backend_t backend, const ggml_tensor * w, ggml_tensor * dst, int64_t e, int64_t last) {
+    return pimpl->copy_experts(backend, w, dst, e, last);
+}
+
+std::map<ggml_backend_buffer_type_t, size_t> llama_moe_cache::memory_breakdown() const {
+    std::map<ggml_backend_buffer_type_t, size_t> res;
+    if (pimpl->buf) {
+        res[ggml_backend_buffer_get_type(pimpl->buf.get())] += pimpl->buf_size;
+    }
+    if (pimpl->buf_host) {
+        res[ggml_backend_buffer_get_type(pimpl->buf_host.get())] += pimpl->buf_host_size;
+    }
+    return res;
+}

src/llama-moe-cache.h

@@ -0,0 +1,39 @@
+#pragma once
+
+#include "ggml-backend.h"
+
+#include <map>
+#include <memory>
+
+struct llama_model;
+
+// keeps the most recently used experts of host-resident MoE layers in a device buffer
+// each layer has a slot map in host memory: when the scheduler copies it to the device, the copy callback uploads the missing experts
+class llama_moe_cache {
+public:
+    llama_moe_cache(const llama_model & model, ggml_backend_t backend, ggml_backend_buffer_type_t buft, size_t size);
+    ~llama_moe_cache();
+
+    ggml_backend_t backend() const;
+
+    // the slot map of layer il, if its experts can be read from the cache for n_tokens tokens, nullptr otherwise
+    ggml_tensor * get_slot_map(int32_t il, int64_t n_tokens, int64_t n_expert_used) const;
+
+    // the experts of w in the cache, nullptr if w is not cached
+    ggml_tensor * get_experts(const ggml_tensor * w) const;
+
+    // ggml_backend_sched copy callback, returns false if src is not a slot map
+    bool copy(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph);
+
+    // for large batches: copy the experts of w that are in the cache, starting at expert e and up to expert last, to the copy dst of w
+    // returns the number of experts copied, 0 if expert e is not in the cache
+    int64_t copy_experts(ggml_backend_t backend, const ggml_tensor * w, ggml_tensor * dst, int64_t e, int64_t last);
+
+    std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const;
+
+private:
+    struct impl;
+    std::unique_ptr<impl> pimpl;
+};
+
+using llama_moe_cache_ptr = std::unique_ptr<llama_moe_cache>;

src/llama-sampler.cpp

@@ -1086,6 +1086,12 @@ static void llama_sampler_greedy_backend_apply(
     struct ggml_tensor * curl = ggml_argmax(ctx, logits);
     ggml_set_name(curl, "greedy_argmax");
 
+    if (data->candidates != nullptr) {
+        struct ggml_tensor * candidates = ggml_reshape_2d(ctx, data->candidates, 1, ggml_nelements(data->candidates));
+        curl = ggml_get_rows(ctx, candidates, curl);
+        ggml_set_name(curl, "greedy_sampled_token");
+    }
+
     data->sampled = curl;
 }
 

src/llama-vocab.cpp

@@ -2860,6 +2860,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                         || t.first == "▁<PRE>"          // CodeLlama
                         || t.first == "<|code_prefix|>" // GLM-4.5
                         || t.first == "<|prefix|>"      // Falcon-H1-Tiny-Coder
+                        || t.first == "<|plamo:fim_prefix|>" // PLaMo-3
                         ) {
                     special_fim_pre_id = t.second;
                     if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
@@ -2881,6 +2882,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                         || t.first == "▁<SUF>"         // CodeLlama
                         || t.first == "<|code_suffix|>" // GLM-4.5
                         || t.first == "<|suffix|>"      // Falcon-H1-Tiny-Coder
+                        || t.first == "<|plamo:fim_suffix|>" // PLaMo-3
                         ) {
                     special_fim_suf_id = t.second;
                     if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
@@ -2902,6 +2904,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                         || t.first == "▁<MID>"         // CodeLlama
                         || t.first == "<|code_middle|>" // GLM-4.5
                         || t.first == "<|middle|>"      // Falcon-H1-Tiny-Coder
+                        || t.first == "<|plamo:fim_middle|>" // PLaMo-3
                         ) {
                     special_fim_mid_id = t.second;
                     if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
@@ -2952,6 +2955,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
             if (special_fim_sep_id == LLAMA_TOKEN_NULL) {
                 if (false
                         || t.first == "<|file_sep|>" // Qwen
+                        || t.first == "<|plamo:file_separator|>" // PLaMo-2/3
                         ) {
                     special_fim_sep_id = t.second;
                     if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {

src/models/bailingmoe3.cpp

@@ -64,15 +64,9 @@ void llama_model_bailingmoe3::load_arch_tensors(llama_model_loader & ml) {
     const int64_t qk_head_dim = hparams.n_embd_head_k_mla();
     const int64_t v_head_dim = hparams.n_embd_head_v_mla();
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
-    int       mtp_flags   = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     for (int il = 0; il < n_layer; ++il) {
         auto & layer = layers[il];

src/models/cohere2moe.cpp

@@ -40,18 +40,9 @@ void llama_model_cohere2moe::load_arch_hparams(llama_model_loader & ml) {
 void llama_model_cohere2moe::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    // Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP
-    // tensors live in a separate file. Mark MTP tensors NOT_REQUIRED so the
-    // trunk loads cleanly.
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only  ? TENSOR_NOT_REQUIRED : 0;
-    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
 

src/models/deepseek2.cpp

@@ -56,15 +56,9 @@ void llama_model_deepseek2::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
     const int64_t n_expert_shared = hparams.n_expert_shared;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only  ? TENSOR_NOT_REQUIRED : 0;
-    int       mtp_flags   = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     const bool is_mla = hparams.is_mla();
 

src/models/deepseek32.cpp

@@ -46,15 +46,9 @@ void llama_model_deepseek32::load_arch_hparams(llama_model_loader & ml) {
 void llama_model_deepseek32::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only   ? TENSOR_NOT_REQUIRED : 0;
-    int       mtp_flags   = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     const bool is_mla = hparams.is_mla();
     if (!is_mla) {

src/models/deepseek4.cpp

@@ -93,9 +93,9 @@ void llama_model_deepseek4::load_arch_tensors(llama_model_loader & ml) {
     const int64_t hc_dim      = hc_mult * n_embd;
     const int64_t hc_mix_dim  = (2 + hc_mult) * hc_mult;
 
-    const bool mtp_only = (n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const int trunk_flags = mtp_only    ? TENSOR_NOT_REQUIRED : 0;
-    const int mtp_flags   = ml.load_mtp ? 0 : TENSOR_SKIP;
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
 

src/models/glm-dsa.cpp

@@ -75,19 +75,9 @@ void llama_model_glm_dsa::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
     const int64_t n_expert_shared = hparams.n_expert_shared;
 
-    // MTP-only: the GGUF carries only the NextN/MTP block(s) (user split target/draft).
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    // Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP
-    // tensors live in a separate file (or were stripped at conversion). Mark
-    // MTP tensors NOT_REQUIRED so the trunk loads cleanly.
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only   ? TENSOR_NOT_REQUIRED : 0;
-    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     const bool is_mla = hparams.is_mla();
     if (!is_mla) {

src/models/glm4-moe.cpp

@@ -29,15 +29,9 @@ void llama_model_glm4_moe::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
     const int64_t n_expert_shared = hparams.n_expert_shared;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only  ? TENSOR_NOT_REQUIRED : 0;
-    int       mtp_flags   = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     GGML_ASSERT(hparams.n_expert > 0 && "n_expert must be > 0 for GLM4_MOE MoE layers");
     GGML_ASSERT(hparams.n_expert_used() > 0 && "n_expert_used must be > 0 for GLM4_MOE MoE layers");

src/models/glm5-next.cpp

@@ -64,12 +64,9 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) {
     const int64_t hc         = hparams.dsv4_hc_mult;
     const int64_t hc_mix_dim = (2 + hc)*hc;
 
-    // the NextN block is loaded but only used by the MTP graph.
-    // Separated trunk_only/mtp_only handling TODO with DECODER_MTP graph in the MTP follow up
-    int mtp_flags = 0;
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
 
@@ -79,18 +76,19 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) {
     for (int i = 0; i < n_layer_all; ++i) {
         auto & layer = layers[i];
 
-        const int flags = (i >= n_layer) ? mtp_flags : 0;
+        const bool is_nextn = i >= n_layer;
+        const int flags = is_nextn ? mtp_flags : trunk_flags;
 
         layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags);
         layer.ffn_norm  = create_tensor(tn(LLM_TENSOR_FFN_NORM,  "weight", i), {n_embd}, flags);
 
         if (i < n_layer) {
-            layer.hc_attn_fn    = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN,    "weight", i), {hc*n_embd, hc_mix_dim}, 0);
-            layer.hc_attn_base  = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE,  "weight", i), {hc_mix_dim}, 0);
-            layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0);
-            layer.hc_ffn_fn     = create_tensor(tn(LLM_TENSOR_HC_FFN_FN,     "weight", i), {hc*n_embd, hc_mix_dim}, 0);
-            layer.hc_ffn_base   = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE,   "weight", i), {hc_mix_dim}, 0);
-            layer.hc_ffn_scale  = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE,  "weight", i), {3}, 0);
+            layer.hc_attn_fn    = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN,    "weight", i), {hc*n_embd, hc_mix_dim}, flags);
+            layer.hc_attn_base  = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE,  "weight", i), {hc_mix_dim}, flags);
+            layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, flags);
+            layer.hc_ffn_fn     = create_tensor(tn(LLM_TENSOR_HC_FFN_FN,     "weight", i), {hc*n_embd, hc_mix_dim}, flags);
+            layer.hc_ffn_base   = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE,   "weight", i), {hc_mix_dim}, flags);
+            layer.hc_ffn_scale  = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE,  "weight", i), {3}, flags);
         }
 
         const int64_t head_dim = hparams.n_embd_head_kda;
@@ -100,25 +98,25 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) {
         if (hparams.is_recr(i)) {
             auto conv = [&](llm_tensor tid) {
                 ggml_tensor * t = create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner, 1}, TENSOR_NOT_REQUIRED);
-                return t ? t : create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner}, 0);
+                return t ? t : create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner}, flags);
             };
             layer.ssm_q_conv = conv(LLM_TENSOR_SSM_CONV1D_Q);
             layer.ssm_k_conv = conv(LLM_TENSOR_SSM_CONV1D_K);
             layer.ssm_v_conv = conv(LLM_TENSOR_SSM_CONV1D_V);
 
-            create_tensor_qkv(layer, i, n_embd, d_inner, d_inner, d_inner, 0);
+            create_tensor_qkv(layer, i, n_embd, d_inner, d_inner, d_inner, flags);
 
-            layer.ssm_f_a  = create_tensor(tn(LLM_TENSOR_SSM_F_A,  "weight", i), {n_embd, head_dim}, 0);
-            layer.ssm_f_b  = create_tensor(tn(LLM_TENSOR_SSM_F_B,  "weight", i), {head_dim, d_inner}, 0);
-            layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0);
+            layer.ssm_f_a  = create_tensor(tn(LLM_TENSOR_SSM_F_A,  "weight", i), {n_embd, head_dim}, flags);
+            layer.ssm_f_b  = create_tensor(tn(LLM_TENSOR_SSM_F_B,  "weight", i), {head_dim, d_inner}, flags);
+            layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, flags);
 
-            layer.ssm_a    = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), {n_head}, 0);
-            layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, 0);
+            layer.ssm_a    = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), {n_head}, flags);
+            layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, flags);
 
-            layer.ssm_g_a    = create_tensor(tn(LLM_TENSOR_SSM_G_A,  "weight", i), {n_embd, head_dim}, 0);
-            layer.ssm_g_b    = create_tensor(tn(LLM_TENSOR_SSM_G_B,  "weight", i), {head_dim, d_inner}, 0);
-            layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_dim}, 0);
-            layer.wo         = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {d_inner, n_embd}, 0);
+            layer.ssm_g_a    = create_tensor(tn(LLM_TENSOR_SSM_G_A,  "weight", i), {n_embd, head_dim}, flags);
+            layer.ssm_g_b    = create_tensor(tn(LLM_TENSOR_SSM_G_B,  "weight", i), {head_dim, d_inner}, flags);
+            layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_dim}, flags);
+            layer.wo         = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {d_inner, n_embd}, flags);
         } else {
             const int64_t q_lora_rank      = hparams.n_lora_q;
             const int64_t kv_lora_rank     = hparams.n_lora_kv;
@@ -155,9 +153,9 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) {
         }
 
         if (i < (int) hparams.n_layer_dense_lead) {
-            layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
-            layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
-            layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd, n_ff}, 0);
+            layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, flags);
+            layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, flags);
+            layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd, n_ff}, flags);
         } else {
             const int64_t n_ff_exp        = hparams.n_ff_exp(i);
             const int64_t n_expert_shared = hparams.n_expert_shared;
@@ -187,7 +185,7 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) {
 
 std::unique_ptr<llm_graph_context> llama_model_glm5_next::build_arch_graph(const llm_graph_params & params) const {
     if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
-        throw std::runtime_error("GLM5-Next NextN graph not implemented yet");
+        return std::make_unique<graph_mtp>(*this, params);
     }
     return std::make_unique<graph>(*this, params);
 }
@@ -537,6 +535,136 @@ ggml_tensor * llama_model_glm5_next::graph::build_hc_post(
     return out;
 }
 
+// construct the graph helpers without building the trunk, so graph_mtp can reuse them
+llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_params & params, no_build) :
+    llm_build_delta_net_base(params), model(model) {
+}
+
+llama_model_glm5_next::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) :
+    graph(model, params, no_build{}) {
+
+    GGML_ASSERT(hparams.n_layer_nextn > 0 && "GLM5-Next MTP requires n_layer_nextn > 0");
+    GGML_ASSERT(hparams.n_layer_nextn == 1 && "GLM5-Next MTP currently supports a single NextN block");
+
+    const int il = hparams.n_layer() + cparams.nextn_layer_offset;
+    GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
+                cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
+                "nextn_layer_offset out of range [0, n_layer_nextn)");
+
+    const auto & layer = model.layers[il];
+
+    GGML_ASSERT(layer.nextn.eh_proj && layer.nextn.enorm && layer.nextn.hnorm &&
+                "GLM5-Next MTP block is missing - load the model with MTP enabled");
+
+    GGML_ASSERT(hparams.n_embd_out() == (uint32_t) n_embd && "GLM5-Next MTP hidden width mismatch");
+
+    const auto * mctx_hyb = static_cast<const llama_memory_hybrid_idx_context *>(mctx);
+
+    auto * inp_hyb   = build_inp_mem_hybrid_k();
+    auto * inp_attn  = inp_hyb->get_attn();
+    auto * inp_kpool = build_inp_kpool(mctx_hyb);
+
+    ggml_build_forward_expand(gf, inp_hyb->get_recr()->s_copy);
+
+    ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+    auto inp = std::make_unique<llm_graph_input_embd_h>(n_embd);
+
+    inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+    ggml_set_input(inp->tokens);
+
+    inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
+    ggml_set_input(inp->embd);
+
+    inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
+    ggml_set_input(inp->h);
+    ggml_set_name(inp->h, "mtp_h_input");
+
+    ggml_tensor * tok_embd = ggml_get_rows(ctx0,
+            layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd, inp->tokens);
+    cb(tok_embd, "mtp_tok_embd", il);
+
+    ggml_tensor * h = inp->h;
+
+    res->add_input(std::move(inp));
+
+    ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
+    cb(e_norm, "mtp_enorm", il);
+
+    ggml_tensor * h_norm = build_norm(h, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
+    cb(h_norm, "mtp_hnorm", il);
+
+    ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, ggml_concat(ctx0, e_norm, h_norm, 0), layer.nextn.eh_proj_s);
+    cb(cur, "mtp_eh_proj", il);
+
+    ggml_tensor * inpSA = cur;
+
+    cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+    cb(cur, "mtp_attn_norm", il);
+
+    ggml_tensor * prev_sel = nullptr;
+    cur = build_dsa_layer(cur, layer, mctx_hyb, inp_attn, inp_kpool, &prev_sel, il);
+    cb(cur, "mtp_attn_out", il);
+
+    // narrow to the output tokens before the position-wise FFN; unmasked nextn embeddings need all rows
+    if (crop_before_nextn(inp_out_ids)) {
+        cur   = ggml_get_rows(ctx0, cur,   inp_out_ids);
+        inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
+    }
+
+    ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+    cb(ffn_inp, "mtp_ffn_inp", il);
+
+    cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
+    cb(cur, "mtp_ffn_norm", il);
+
+    ggml_tensor * moe_out = build_moe_ffn(cur,
+            layer.ffn_gate_inp,
+            layer.ffn_up_exps,
+            layer.ffn_gate_exps,
+            layer.ffn_down_exps,
+            layer.ffn_exp_probs_b,
+            n_expert, n_expert_used,
+            LLM_FFN_SILU, hparams.expert_weights_norm,
+            hparams.expert_weights_scale,
+            (llama_expert_gating_func_type) hparams.expert_gating_func,
+            il);
+    cb(moe_out, "mtp_ffn_moe_out", il);
+
+    ggml_tensor * ffn_shexp = build_ffn(cur,
+            layer.ffn_up_shexp,   nullptr, nullptr,
+            layer.ffn_gate_shexp, nullptr, nullptr,
+            layer.ffn_down_shexp, nullptr, nullptr,
+            nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
+    cb(ffn_shexp, "mtp_ffn_shexp", il);
+
+    cur = ggml_add(ctx0, moe_out, ffn_shexp);
+    cb(cur, "mtp_ffn_out", il);
+
+    cur = ggml_add(ctx0, cur, ffn_inp);
+    cb(cur, "mtp_post_ffn", il);
+
+    ggml_tensor * head_norm = layer.nextn.shared_head_norm ? layer.nextn.shared_head_norm : model.output_norm;
+    GGML_ASSERT(head_norm && "GLM5-Next MTP: missing both nextn.shared_head_norm and output_norm");
+    cur = build_norm(cur, head_norm, nullptr, LLM_NORM_RMS, -1);
+    cb(cur, "h_nextn", -1);
+    res->t_h_nextn = cur;
+
+    if (crop_after_nextn(inp_out_ids)) {
+        cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+    }
+    cb(cur, "mtp_shared_head_norm", -1);
+
+    ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
+    ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;
+    GGML_ASSERT(head_w && "GLM5-Next MTP: missing both nextn.shared_head_head and output");
+    cur = build_lora_mm(head_w, cur, head_s);
+    cb(cur, "result_output", -1);
+
+    res->t_logits = cur;
+    ggml_build_forward_expand(gf, cur);
+}
+
 llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_params & params) :
     llm_build_delta_net_base(params), model(model) {
 

src/models/hy-v3.cpp

@@ -22,18 +22,9 @@ void llama_model_hy_v3::load_arch_hparams(llama_model_loader & ml) {
 void llama_model_hy_v3::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    // Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP
-    // tensors live in a separate file (e.g. user split target/draft). Mark
-    // MTP tensors NOT_REQUIRED so the trunk loads cleanly.
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only   ? TENSOR_NOT_REQUIRED : 0;
-    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
 

src/models/mimo2.cpp

@@ -25,15 +25,9 @@ void llama_model_mimo2::load_arch_hparams(llama_model_loader & ml) {
 void llama_model_mimo2::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
-    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
 

src/models/models.h

@@ -2741,6 +2741,10 @@ struct llama_model_glm5_next : public llama_model_base {
     struct graph : public llm_build_delta_net_base {
         graph(const llama_model & model, const llm_graph_params & params);
 
+        // build the helpers without the trunk, so graph_mtp can reuse them
+        struct no_build {};
+        graph(const llama_model & model, const llm_graph_params & params, no_build);
+
         // collapse the hc streams with per-stream weights
         ggml_tensor * build_hc_pre(
                 ggml_tensor * x,
@@ -2789,6 +2793,11 @@ struct llama_model_glm5_next : public llama_model_base {
 
     };
 
+    // the NextN / MTP draft head: one DSA block appended after the trunk
+    struct graph_mtp : public graph {
+        graph_mtp(const llama_model & model, const llm_graph_params & params);
+    };
+
     std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
 };
 

src/models/nemotron-h.cpp

@@ -49,9 +49,9 @@ void llama_model_nemotron_h::load_arch_hparams(llama_model_loader & ml) {
 void llama_model_nemotron_h::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
-    const bool mtp_only    = hparams.n_layer_nextn > 0 && ml.get_weight("blk.0.attn_norm.weight") == nullptr;
-    const int  trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
-    const int  mtp_flags   = !ml.load_mtp ? TENSOR_SKIP : 0;
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     // mamba2 Mixer SSM params
     // NOTE: int64_t for tensor dimensions

src/models/qwen35.cpp

@@ -33,9 +33,9 @@ void llama_model_qwen35::load_arch_hparams(llama_model_loader & ml) {
 void llama_model_qwen35::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
-    int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
 

src/models/qwen35moe.cpp

@@ -36,9 +36,9 @@ void llama_model_qwen35moe::load_arch_hparams(llama_model_loader & ml) {
 void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
-    int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
 

src/models/qwen3next.cpp

@@ -35,9 +35,9 @@ void llama_model_qwen3next::load_arch_tensors(llama_model_loader & ml) {
         throw std::runtime_error(arch_name() + " model cannot have zero experts");
     }
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
-    int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
 

src/models/qwen4exp.cpp

@@ -175,10 +175,9 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) {
     const int64_t hc_dim = hc * n_embd;
     const int64_t hc_lr  = hparams.hc_low_rank;
 
-    // an MTP-only file carries the MTP block, the embeddings and the LM head, but no trunk
-    const bool mtp_only    = n_layer_nextn > 0 && ml.get_weight(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", 0).str().c_str()) == nullptr;
-    const int  trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
-    const int  mtp_flags   = ml.load_mtp ? 0 : TENSOR_SKIP;
+    const auto nf = nextn_flags(ml, LLM_TENSOR_HC_ATTN_NORM);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
 

src/models/step35.cpp

@@ -37,18 +37,9 @@ void llama_model_step35::load_arch_hparams(llama_model_loader & ml) {
 void llama_model_step35::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
-    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
-    // Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP
-    // tensors live in a separate file (e.g. user split target/draft). Mark
-    // MTP tensors NOT_REQUIRED so the trunk loads cleanly.
-    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
-    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int trunk_flags = mtp_only  ? TENSOR_NOT_REQUIRED : 0;
-    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
-
-    if (!ml.load_mtp) {
-        mtp_flags |= TENSOR_SKIP;
-    }
+    const auto nf = nextn_flags(ml);
+    const int trunk_flags = nf.trunk;
+    const int mtp_flags   = nf.mtp;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
 

tests/test-backend-ops.cpp

@@ -7233,6 +7233,7 @@ enum mul_mat_add_mode {
     MUL_MAT_ADD_ROW,        // mm + a one-row res, which a same-shape fusion must leave alone
     MUL_MAT_ADD_RES_INPLACE, // res += mm
     MUL_MAT_ADD_B_INPLACE,   // b += mm: the sum overwrites the mat-mul input (m == k)
+    MUL_MAT_ADD_MM_MM,       // mm2 + mm: the residual is itself a mat-mul output, so both operands are MUL_MAT
 };
 
 static std::string var_to_str(mul_mat_add_mode mode) {
@@ -7242,6 +7243,7 @@ static std::string var_to_str(mul_mat_add_mode mode) {
         case MUL_MAT_ADD_ROW:         return "mm+row";
         case MUL_MAT_ADD_RES_INPLACE: return "res+=mm";
         case MUL_MAT_ADD_B_INPLACE:   return "b+=mm";
+        case MUL_MAT_ADD_MM_MM:       return "mm2+mm";
     }
     return "unknown";
 }
@@ -7287,6 +7289,12 @@ struct test_mul_mat_add : public test_case {
             case MUL_MAT_ADD_RES_MM:      out = ggml_add(ctx, res, mm);         break;
             case MUL_MAT_ADD_RES_INPLACE: out = ggml_add_inplace(ctx, res, mm); break;
             case MUL_MAT_ADD_B_INPLACE:   out = ggml_add_inplace(ctx, b, mm);   break;
+            case MUL_MAT_ADD_MM_MM:
+                {
+                    ggml_tensor * a2  = ggml_new_tensor_2d(ctx, type_a, k, m);
+                    ggml_tensor * mm2 = ggml_mul_mat(ctx, a2, b);
+                    out = ggml_add(ctx, mm2, mm);
+                } break;
         }
         ggml_set_name(out, "out");
 
@@ -9550,6 +9558,12 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
                     }
                 }
             }
+            // MSA indexer block-max shape: wide 1-D window along ne0, no padding, non-divisible width
+            for (int64_t iw : {4096, 4160}) {
+                for (int blk : {32, 64}) {
+                    test_cases.emplace_back(new test_pool2d(pool_type, type_input, {iw, 2, 1, 1}, blk, 1, blk, 1, 0, 0));
+                }
+            }
         }
     }
 
@@ -10396,8 +10410,11 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
             test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, m, n, k, {1, 1}, {1, 1}));
         }
         test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 64, n, 256, {3, 2}, {2, 1}));
-        for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, GGML_TYPE_Q8_0,
-                                 GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K}) {
+        for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1,
+                                 GGML_TYPE_Q8_0, GGML_TYPE_Q1_0, GGML_TYPE_Q2_0, GGML_TYPE_MXFP4, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K,
+                                 GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_TQ2_0,
+                                 GGML_TYPE_IQ2_XXS, GGML_TYPE_IQ2_XS, GGML_TYPE_IQ2_S, GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ3_S,
+                                 GGML_TYPE_IQ1_S, GGML_TYPE_IQ1_M, GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_XS}) {
             test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 48,  n, 2560, {1, 1}, {1, 1}));
             test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 1000, n, 1024, {1, 1}, {1, 1}));
             test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 3000, n, 512,  {1, 1}, {1, 1}));
@@ -10409,7 +10426,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
             }
         }
         for (ggml_type type_a : {GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q5_K, GGML_TYPE_F16}) {
-            for (mul_mat_add_mode mode : {MUL_MAT_ADD_MM_RES, MUL_MAT_ADD_RES_MM, MUL_MAT_ADD_ROW, MUL_MAT_ADD_RES_INPLACE}) {
+            for (mul_mat_add_mode mode : {MUL_MAT_ADD_MM_RES, MUL_MAT_ADD_RES_MM, MUL_MAT_ADD_ROW, MUL_MAT_ADD_RES_INPLACE, MUL_MAT_ADD_MM_MM}) {
                 test_cases.emplace_back(new test_mul_mat_add(type_a, 1000, n, 1024, mode));
             }
             test_cases.emplace_back(new test_mul_mat_add(type_a, 2048, n, 2048, MUL_MAT_ADD_B_INPLACE));
@@ -10420,7 +10437,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
             test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 1000, n, 1024, {1, 1}, {1, 1}, {0, 1, 2, 3}, 1280));
         }
         // a src1 with a nonzero mean, for the zero points and mins of quantized src0 types
-        for (ggml_type type_a : {GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q8_0, GGML_TYPE_Q4_K, GGML_TYPE_Q5_K}) {
+        for (ggml_type type_a : {GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q8_0, GGML_TYPE_Q2_K, GGML_TYPE_Q4_K, GGML_TYPE_Q5_K}) {
             test_cases.emplace_back(new test_mul_mat_pos(type_a, GGML_TYPE_F32, 1000, n, 1024, {1, 1}, {1, 1}));
         }
     }
@@ -11641,6 +11658,15 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
             GGML_TYPE_F32, {n_kv, 512, 64, 1}, false, {2, 1, 0, 3}));
     }
 
+    // POOL_2D max over MiniMax-M3 indexer block scores:
+    // sc is [n_ps, n_head=4, n_tokens] and gets ggml_pool_2d(.., MAX, blk=128, 1, blk=128, 1, 0, 0).
+    for (int64_t n_ps : { 8192, 32768 }) {
+        for (int64_t n_tokens : { 1, 512 }) {
+            test_cases.emplace_back(new test_pool2d(
+                GGML_OP_POOL_MAX, GGML_TYPE_F32, {n_ps, 4, n_tokens, 1}, 128, 1, 128, 1, 0, 0));
+        }
+    }
+
     // LEAKY_RELU at FFN activation width, for direct comparison with RELU
     for (int64_t n_tokens : {512, 2048}) {
         test_cases.emplace_back(new test_leaky_relu(GGML_TYPE_F32, { 17408, n_tokens, 1, 1 }, 0.1f));

tests/test-backend-sampler.cpp

@@ -2,6 +2,7 @@
 #include "llama.h"
 #include "llama-cpp.h"
 #include "common.h"
+#include "sampling.h"
 
 #ifdef NDEBUG
 #undef NDEBUG
@@ -316,7 +317,7 @@ static llama_sampler * test_single_output_backend_sampler_init(
     return llama_sampler_init(&test_single_output_backend_sampler_i, ctx);
 }
 
-static void test_backend_greedy_sampling(const test_params & params) {
+static void test_greedy(const test_params & params) {
     const int seq_id = 0;
 
     struct llama_sampler_chain_params backend_sampler_params = llama_sampler_chain_default_params();
@@ -351,7 +352,95 @@ static void test_backend_greedy_sampling(const test_params & params) {
     }
 }
 
-static void test_backend_top_k_sampling(const test_params & params) {
+
+static void test_greedy_filtered_common(const test_params & params) {
+    auto check = [&](common_params_sampling sp, bool greedy) {
+        common_sampler_ptr sampler(common_sampler_init(params.model.get(), sp));
+        auto * chain = common_sampler_get(sampler.get());
+        auto * last = llama_sampler_chain_get(chain, llama_sampler_chain_n(chain) - 1);
+        GGML_ASSERT(strcmp(llama_sampler_name(last), greedy ? "greedy" : "dist") == 0);
+
+        llama_token_data tokens[] = {{0, -2.0f, 0.0f}, {1, -1.0f, 0.0f}, {2, 0.0f, 0.0f}, {3, 1.0f, 0.0f}};
+        llama_token_data_array candidates = {tokens, 4, -1, false};
+        llama_sampler_apply(chain, &candidates);
+        GGML_ASSERT(candidates.selected >= 0);
+        if (greedy || sp.n_probs > 0) {
+            GGML_ASSERT(candidates.data[candidates.selected].id == 3);
+        }
+        if (sp.n_probs > 0) {
+            GGML_ASSERT(candidates.data[candidates.selected].p == 1.0f);
+        }
+        if (sp.dynatemp_range > 0.0f) {
+            int positive = 0;
+            for (size_t i = 0; i < candidates.size; ++i) {
+                positive += candidates.data[i].p > 0.0f;
+            }
+            GGML_ASSERT(positive > 1);
+        }
+    };
+
+    common_params_sampling sp;
+    sp.temp = 0.0f;
+    sp.samplers = {COMMON_SAMPLER_TYPE_TOP_K, COMMON_SAMPLER_TYPE_TEMPERATURE};
+    for (bool backend : {false, true}) {
+        sp.backend_sampling = backend;
+        check(sp, true);
+        auto probabilities = sp;
+        probabilities.n_probs = 4;
+        check(probabilities, false);
+        auto dynamic = sp;
+        dynamic.dynatemp_range = 1.0f;
+        check(dynamic, false);
+    }
+    auto grammar = sp;
+    grammar.grammar = {COMMON_GRAMMAR_TYPE_USER, "root ::= [a-z]+"};
+    check(grammar, true);
+    auto budget = sp;
+    budget.reasoning_budget_start = {1};
+    budget.reasoning_budget_end = {{2}};
+    budget.reasoning_budget_tokens = 1;
+    check(budget, true);
+    sp.samplers = {COMMON_SAMPLER_TYPE_TEMPERATURE, COMMON_SAMPLER_TYPE_TOP_K};
+    check(sp, false);
+    sp.samplers.clear();
+    check(sp, false);
+
+    sp.temp = 0.8f;
+    sp.samplers = {COMMON_SAMPLER_TYPE_TEMPERATURE, COMMON_SAMPLER_TYPE_TOP_K};
+    for (bool backend : {false, true}) {
+        sp.backend_sampling = backend;
+        sp.top_k = 1;
+        check(sp, true);
+        auto probabilities = sp;
+        probabilities.n_probs = 4;
+        check(probabilities, false);
+        sp.top_k = 8;
+        check(sp, false);
+    }
+}
+
+static void test_greedy_filtered(const test_params & params) {
+    for (int k : {1, 8}) {
+        llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params()));
+        llama_sampler_chain_add(chain.get(), llama_sampler_init_top_k(k));
+        llama_sampler_chain_add(chain.get(), llama_sampler_init_greedy());
+        std::vector<llama_sampler_seq_config> configs = {{0, chain.get()}};
+        test_context test_ctx(params, configs);
+        GGML_ASSERT(test_ctx.decode({{0, "Write a Python function"}}));
+        for (int step = 0; step < 4; ++step) {
+            const int idx = test_ctx.idx_for_seq(0);
+            const auto * logits = llama_get_sampled_logits_ith(test_ctx.ctx.get(), idx);
+            const auto * ids = llama_get_sampled_candidates_ith(test_ctx.ctx.get(), idx);
+            GGML_ASSERT(llama_get_sampled_logits_count_ith(test_ctx.ctx.get(), idx) == (uint32_t) k);
+            GGML_ASSERT(llama_get_sampled_candidates_count_ith(test_ctx.ctx.get(), idx) == (uint32_t) k);
+            const auto expected = ids[std::max_element(logits, logits + k) - logits];
+            GGML_ASSERT(llama_get_sampled_token_ith(test_ctx.ctx.get(), idx) == expected);
+            GGML_ASSERT(test_ctx.decode_token(expected));
+        }
+    }
+}
+
+static void test_top_k(const test_params & params) {
     const int seq_id = 0;
     const int32_t k = 8;
     struct llama_sampler_chain_params backend_chain_params = llama_sampler_chain_default_params();
@@ -393,7 +482,7 @@ static void test_backend_top_k_sampling(const test_params & params) {
     printf("backend top-k hybrid sampling test PASSED\n");
 }
 
-static void test_backend_temp_sampling(const test_params & params) {
+static void test_temp(const test_params & params) {
     {
         const float temp_0 = 0.8f;
         struct llama_sampler_chain_params backend_chain_params_0 = llama_sampler_chain_default_params();
@@ -481,7 +570,7 @@ static void test_backend_temp_sampling(const test_params & params) {
     printf("backend temp sampling test PASSED\n");
 }
 
-static void test_backend_temp_ext_sampling(const test_params & params) {
+static void test_temp_ext(const test_params & params) {
     {
         int seq_id = 0;
         const float temp = 0.8f;
@@ -546,7 +635,7 @@ static void test_backend_temp_ext_sampling(const test_params & params) {
     printf("backend temp_ext sampling test PASSED\n");
 }
 
-static void test_backend_min_p_sampling(const test_params & params) {
+static void test_min_p(const test_params & params) {
     const int seq_id = 0;
     const float p = 0.1;
     struct llama_sampler_chain_params backend_chain_params = llama_sampler_chain_default_params();
@@ -598,7 +687,7 @@ static void test_backend_min_p_sampling(const test_params & params) {
     printf("min-p sampling test PASSED\n");
 }
 
-static void test_backend_top_p_sampling(const test_params & params) {
+static void test_top_p(const test_params & params) {
     const int seq_id = 0;
     const float p = 0.9;
     struct llama_sampler_chain_params backend_chain_params = llama_sampler_chain_default_params();
@@ -648,7 +737,7 @@ static void test_backend_top_p_sampling(const test_params & params) {
     printf("top-p sampling test PASSED\n");
 }
 
-static void test_backend_multi_sequence_sampling(const test_params & params) {
+static void test_multi_sequence(const test_params & params) {
     struct llama_sampler_chain_params chain_params_0 = llama_sampler_chain_default_params();
     llama_sampler_ptr sampler_chain_0(llama_sampler_chain_init(chain_params_0));
     llama_sampler_chain_add(sampler_chain_0.get(), llama_sampler_init_greedy());
@@ -714,7 +803,7 @@ static void test_backend_multi_sequence_sampling(const test_params & params) {
     printf("backend multi-sequence sampling test PASSED\n");
 }
 
-static void test_backend_dist_sampling(const test_params & params) {
+static void test_dist(const test_params & params) {
     const int seq_id = 0;
     const int32_t seed = 88;
 
@@ -742,7 +831,7 @@ static void test_backend_dist_sampling(const test_params & params) {
     printf("backend dist sampling test PASSED\n");
 }
 
-static void test_backend_dist_sampling_and_cpu(const test_params & params) {
+static void test_dist_and_cpu(const test_params & params) {
     const int seq_id = 0;
     const int32_t seed = 88;
 
@@ -772,7 +861,7 @@ static void test_backend_dist_sampling_and_cpu(const test_params & params) {
     printf("backend dist & cpu sampling test PASSED\n");
 }
 
-static void test_backend_logit_bias_sampling(const test_params & params) {
+static void test_logit_bias(const test_params & params) {
     const auto * model = params.model.get();
     const auto * vocab = llama_model_get_vocab(model);
 
@@ -1093,7 +1182,7 @@ static void compare_penalties_logits(
     GGML_ASSERT(stats.n_mismatch == 0);
 }
 
-static void test_penalty_parameter_values(const test_params & params) {
+static void check_penalty_parameter_values(const test_params & params) {
     struct penalty_test_case {
         const char * name;
         float repeat;
@@ -1292,7 +1381,7 @@ static void compare_masking_penalties_logits(
     GGML_ASSERT(stats.n_mismatch == 0);
 }
 
-static void test_backend_penalties_sampling(const test_params & params) {
+static void test_penalties(const test_params & params) {
     printf("Testing backend penalties (repeat + freq + presence)\n");
     compare_penalties_logits(params, 64, 1.1f, 0.5f, 0.25f, "Hello Hello world");
 
@@ -1371,14 +1460,14 @@ static void test_backend_penalties_sampling(const test_params & params) {
     }, 64, 1.0f, 0.0f, 0.25f, "Hello", penalties_position::after_filter, true);
 
     printf("Testing backend penalty parameter values\n");
-    test_penalty_parameter_values(params);
+    check_penalty_parameter_values(params);
 
     printf("backend penalties sampling test PASSED\n");
 }
 
 // This test verifies that it is possible to have two different backend samplers,
 // one that uses the backend dist sampler, and another that uses CPU dist sampler.
-static void test_backend_mixed_sampling(const test_params & params) {
+static void test_mixed(const test_params & params) {
     struct llama_sampler_chain_params chain_params_0 = llama_sampler_chain_default_params();
     llama_sampler_ptr sampler_chain_0(llama_sampler_chain_init(chain_params_0));
     llama_sampler_chain_add(sampler_chain_0.get(), llama_sampler_init_dist(88));
@@ -1428,7 +1517,7 @@ static void test_backend_mixed_sampling(const test_params & params) {
     printf("backend mixed sampling test PASSED\n");
 }
 
-static void test_backend_set_sampler(const test_params & params) {
+static void test_set_sampler(const test_params & params) {
     const int seq_id = 0;
     const int32_t seed = 88;
 
@@ -1493,7 +1582,7 @@ static void test_backend_set_sampler(const test_params & params) {
     printf("backend set sampler test PASSED\n");
 }
 
-static void test_backend_cpu_mixed_batch(const test_params & params) {
+static void test_cpu_mixed(const test_params & params) {
     // Sequence 0 uses backend sampling
     struct llama_sampler_chain_params chain_params_0 = llama_sampler_chain_default_params();
     llama_sampler_ptr sampler_chain_0(llama_sampler_chain_init(chain_params_0));
@@ -1581,7 +1670,7 @@ static void test_backend_cpu_mixed_batch(const test_params & params) {
     printf("backend-cpu mixed batch test PASSED\n");
 }
 
-static void test_backend_multi_output_limit(const test_params & params) {
+static void test_multi_output_limit(const test_params & params) {
     const llama_seq_id seq_id = 0;
 
     llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params()));
@@ -1594,15 +1683,15 @@ static void test_backend_multi_output_limit(const test_params & params) {
         batch.add(llama_vocab_bos(test_ctx.vocab), i, seq_id, true);
     }
 
-    printf(">>> test_backend_multi_output_limit expected error start:\n");
+    printf(">>> test_multi_output_limit expected error start:\n");
     const int ret = llama_process(test_ctx.ctx.get(), LLAMA_PROCESS_TYPE_DECODE, batch.get());
     GGML_ASSERT(ret != 0 && "llama_decode should reject outputs above the per-sequence limit");
-    printf("<<< test_backend_multi_output_limit expected error end.\n");
+    printf("<<< test_multi_output_limit expected error end.\n");
 
     printf("backend multi-output limit test PASSED\n");
 }
 
-static void test_backend_multi_sequence_multi_output_dist(const test_params & params) {
+static void test_multi_output_multi_sequence_dist(const test_params & params) {
     const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
     const int32_t n_vocab = llama_vocab_n_tokens(vocab);
     const uint32_t seeds[] = { 88, 1337 };
@@ -1697,7 +1786,7 @@ static void test_backend_multi_sequence_multi_output_dist(const test_params & pa
     printf("backend multi-sequence multi-output dist test PASSED\n");
 }
 
-static void test_backend_multi_output_dist_transaction(const test_params & params) {
+static void test_multi_output_dist_transaction(const test_params & params) {
     const llama_seq_id seq_id = 0;
     const uint32_t seed = 95;
     const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
@@ -1762,7 +1851,7 @@ static void test_backend_multi_output_dist_transaction(const test_params & param
     printf("backend multi-output dist transaction test PASSED\n");
 }
 
-static void test_backend_multi_output_sampling_chain(const test_params & params) {
+static void test_multi_output_sampling_chain(const test_params & params) {
     const llama_seq_id seq_id = 0;
     const uint32_t seed = 88;
     const float p = 0.9f;
@@ -1913,7 +2002,7 @@ static void test_backend_multi_output_sampling_chain(const test_params & params)
     printf("backend multi-output sampling chain test PASSED\n");
 }
 
-static void test_backend_multi_output_cpu_suffix(const test_params & params) {
+static void test_multi_output_cpu_suffix(const test_params & params) {
     const llama_seq_id seq_id = 0;
     const int32_t k = 8;
     const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
@@ -1977,28 +2066,35 @@ struct backend_test_case {
     bool enabled_by_default;
 };
 
+// note: test names are "test_<suffix>" and match the function implementing them
 static const backend_test_case BACKEND_TESTS[] = {
-    { "greedy",          test_backend_greedy_sampling,         true  },
-    { "logit_bias",      test_backend_logit_bias_sampling,     true  },
-    { "penalties",       test_backend_penalties_sampling,      true  },
-    { "temp",            test_backend_temp_sampling,           true  },
-    { "temp_ext",        test_backend_temp_ext_sampling,       true  },
-    { "top_k",           test_backend_top_k_sampling,          true  },
-    { "multi_sequence",  test_backend_multi_sequence_sampling, true  },
-    { "dist",            test_backend_dist_sampling,           true  },
-    { "dist_and_cpu",    test_backend_dist_sampling_and_cpu,   true  },
-    { "set_sampler",     test_backend_set_sampler,             true  },
-    { "multi_output_limit",    test_backend_multi_output_limit,      true },
-    { "multi_sequence_multi_output_dist", test_backend_multi_sequence_multi_output_dist, true },
-    { "multi_output_dist_transaction", test_backend_multi_output_dist_transaction, true },
-    { "multi_output_sampling_chain", test_backend_multi_output_sampling_chain, true },
-    { "multi_output_cpu",      test_backend_multi_output_cpu_suffix, true },
-    { "mixed",           test_backend_mixed_sampling,          true  },
-    { "min_p",           test_backend_min_p_sampling,          true  },
-    { "cpu_mixed",       test_backend_cpu_mixed_batch,         true  },
-    { "top_p",           test_backend_top_p_sampling,          true  },
+    // single sampler
+    { "test_greedy",                 test_greedy,                 true },
+    { "test_greedy_filtered",        test_greedy_filtered,        true },
+    { "test_greedy_filtered_common", test_greedy_filtered_common, true },
+    { "test_temp",                   test_temp,                   true },
+    { "test_temp_ext",               test_temp_ext,               true },
+    { "test_top_k",                  test_top_k,                  true },
+    { "test_top_p",                  test_top_p,                  true },
+    { "test_min_p",                  test_min_p,                  true },
+    { "test_logit_bias",             test_logit_bias,             true },
+    { "test_penalties",              test_penalties,              true },
+    // multiple sequences
+    { "test_multi_sequence",         test_multi_sequence,         true },
+    { "test_dist",                   test_dist,                   true },
+    { "test_dist_and_cpu",           test_dist_and_cpu,           true },
+    { "test_mixed",                  test_mixed,                  true },
+    { "test_cpu_mixed",              test_cpu_mixed,              true },
+    { "test_set_sampler",            test_set_sampler,            true },
+    // multiple outputs per sequence
+    { "test_multi_output_limit",                  test_multi_output_limit,               true },
+    { "test_multi_output_multi_sequence_dist",    test_multi_output_multi_sequence_dist, true },
+    { "test_multi_output_dist_transaction",       test_multi_output_dist_transaction,    true },
+    { "test_multi_output_sampling_chain",         test_multi_output_sampling_chain,      true },
+    { "test_multi_output_cpu_suffix",             test_multi_output_cpu_suffix,          true },
 };
 
+// TODO: add usage, examples
 static test_args parse_cli(int argc, char ** argv) {
     test_args out;
 
@@ -2082,10 +2178,12 @@ static std::vector<const backend_test_case *> collect_tests_to_run(const std::st
             }
 #ifdef GGML_USE_HIP
             // TODO: remove this when https://github.com/ggml-org/llama.cpp/pull/26592 is merged
-            if (test.name == "penalties" || test.name == "set_sampler" ||
-                test.name == "mixed"     || test.name == "top_p"       ||
-                test.name == "multi_output_sampling_chain" ||
-                test.name == "multi_output_cpu") {
+            if (test.name == "test_penalties"                   ||
+                test.name == "test_set_sampler"                 ||
+                test.name == "test_mixed"                       ||
+                test.name == "test_top_p"                       ||
+                test.name == "test_multi_output_sampling_chain" ||
+                test.name == "test_multi_output_cpu_suffix") {
                 fprintf(stderr, "Skipping test '%s' on HIP backend (no backend TOP_K support)\n", test.name.c_str());
                 continue;
             }

tests/test-chat.cpp

@@ -1087,6 +1087,7 @@ struct peg_test_case {
     common_chat_msg              expect;
     bool                         is_partial            = false;
     bool                         expect_reconstruction = false;
+    std::vector<std::string>     expect_rules;
 };
 
 struct make_peg_parser {
@@ -1105,7 +1106,7 @@ struct make_peg_parser {
     common_chat_msg parse(const std::string & msg, bool is_partial) const {
         common_chat_parser_params parser_params(params_);
         parser_params.debug = detailed_debug_;
-        return common_chat_peg_parse(arena_, msg, is_partial, parser_params);
+        return common_chat_peg_parse(arena_, common_chat_input(msg), is_partial, parser_params);
     }
 };
 
@@ -1166,6 +1167,14 @@ static void test_peg_parser(common_chat_templates *                      tmpls,
         LOG_DBG("Generation prompt: '%s'\n", parser.params_.generation_prompt.c_str());
     }
 
+    for (const auto & rule : tc.expect_rules) {
+        if (!parser.arena_.has_rule(rule)) {
+            LOG_ERR("Missing rule: %s\n", rule.c_str());
+            common_log_flush(common_log_main());
+            throw std::runtime_error("Test failed");
+        }
+    }
+
     common_chat_msg msg_accum;
     common_chat_msg msg_prev;
     msg_accum.role = msg_prev.role = "assistant";
@@ -1573,6 +1582,11 @@ class peg_test_builder {
         return *this;
     }
 
+    peg_test_builder & expect_rules(std::vector<std::string> rules) {
+        tc_.expect_rules = std::move(rules);
+        return *this;
+    }
+
     peg_test_builder & tool_choice(common_chat_tool_choice choice) {
         tc_.params.tool_choice = choice;
         return *this;
@@ -2174,6 +2188,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
             .tools({ special_function_tool })
             .expect(message_assist_call)
+            .expect_rules({ "tool-0", "tool-0-arg-0" })
             .run();
 
         tst.test(
@@ -2625,6 +2640,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .enable_thinking(true)
             .tools({ special_function_tool })
             .expect(message_assist_call)
+            .expect_rules({ "tool-0" })
             .run();
 
         tst.test(
@@ -3024,6 +3040,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
                 "<|tool_call>call:get_time{city:<|\"|>London<|\"|>}<tool_call|>")
             .tools({ get_time_tool })
             .expect(message_with_tool_calls("get_time", R"({"city": "London"})"))
+            .expect_rules({ "tool-0" })
             .run();
 
         // Tool call with string argument containing special chars
@@ -3340,6 +3357,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
                "</tool_call>")
             .tools({ special_function_tool })
             .expect(message_assist_call)
+            .expect_rules({ "tool-0" })
             .run();
 
         tst.test(
@@ -3624,6 +3642,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
                 { "special_function", R"({"arg1": 1})", {} },
                 { "special_function_with_opt", R"({"arg1": 1, "arg2": 2})", {} },
             })
+            .expect_rules({ "tool-0", "tool-0-arg-0", "tool-1", "tool-1-arg-0", "tool-1-arg-1" })
             .run();
 
         tst.test(
@@ -4035,6 +4054,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
             .tools({ get_time_tool })
             .expect(message_with_tool_calls_and_reasoning("get_time", R"({"city": "Tokyo"})", "Let me check the time"))
+            .expect_rules({ "tool-0", "tool-0-arg-0" })
             .run();
 
         // Tool call without reasoning (non-thinking mode), integer param (string="false")
@@ -4643,7 +4663,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             bool got_out_of_range = false;
             std::string error_msg;
             try {
-                common_chat_peg_parse(arena, bad_input, /*is_partial=*/false, pp);
+                common_chat_peg_parse(arena, common_chat_input(bad_input), /*is_partial=*/false, pp);
             } catch (const std::out_of_range & e) {
                 got_out_of_range = true;
                 error_msg = e.what();
@@ -4673,6 +4693,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ get_time_tool })
             .expect_reasoning("I need to check the time first.\n")
             .expect_tool_calls({ { "get_time", R"({"city": "Paris"})", "" } })
+            .expect_rules({ "ling3-tool-0", "ling3-arg-0-0" })
             .run();
 
         // Closed think block, prose, then a tool call.
@@ -4911,6 +4932,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ special_function_tool })
             .expect(message_assist_call_thoughts)
             .expect_reconstruction()
+            .expect_rules({ "tool-0", "tool-0-arg-0" })
             .run();
 
         tst.test(
@@ -5123,6 +5145,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .expect_tool_calls({
                 { "special_function", R"({"arg1":1})", "" },
             })
+            .expect_rules({ "kimi-k3-tool-0", "kimi-k3-arg-0-0" })
             .run();
 
         // Tool call preceded by reasoning (no opening think tag) and content.
@@ -5411,6 +5434,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ special_function_tool })
             .expect(kimi_id_special_func_tool_call)
             .expect_reconstruction()
+            .expect_rules({ "tool-0" })
             .run();
 
         // Kimi-K2-Instruct
@@ -5465,6 +5489,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
             .tools({ special_function_tool })
             .expect(message_assist_call_thoughts)
+            .expect_rules({ "tool-0" })
             .run();
 
         // Tool call with reasoning and content
@@ -5742,6 +5767,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ get_time_tool })
             .expect(message_with_tool_calls_and_reasoning("get_time", R"({"city": "Tokyo"})", "Let me check the time"))
             .expect_reconstruction()
+            .expect_rules({ "tool-0", "tool-0-arg-0" })
             .run();
 
         // Tool call without reasoning, integer param
@@ -6140,6 +6166,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ special_function_tool })
             .expect(message_assist_call_id)
             .expect_reconstruction()
+            .expect_rules({ "tool-0" })
             .run();
 
         // Continuation tests
@@ -6175,6 +6202,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ special_function_tool })
             .expect(message_assist_call)
             .expect_reconstruction()
+            .expect_rules({ "tool-0" })
             .run();
 
         // Continuation tests
@@ -6547,7 +6575,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
     {
         // Llama 3.2
         auto tst = peg_tester("models/templates/meta-llama-Llama-3.2-3B-Instruct.jinja", detailed_debug);
-        tst.test("Hello, world!\nWhat's up?").tools({ special_function_tool }).expect(message_assist).expect_reconstruction().run();
+        tst.test("Hello, world!\nWhat's up?").tools({ special_function_tool }).expect(message_assist).expect_reconstruction().expect_rules({ "tool-0" }).run();
 
         // Continuation tests
         tst.test("world!\nWhat's up?")
@@ -6595,6 +6623,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ special_function_tool })
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
             .expect(message_assist_call)
+            .expect_rules({ "tool-0" })
             .run();
 
         // "Inform then act": the model answers the user and calls a tool in ONE generation,
@@ -6680,6 +6709,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
             .tools({ special_function_tool })
             .expect(message_assist_call)
+            .expect_rules({ "tool-0" })
             .run();
 
         // Tool call with recipient in channel header: "<|channel|>analysis to=functions.NAME<|message|>JSON"
@@ -6855,6 +6885,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
             .tools({ special_function_tool })
             .expect(message_assist_call_thoughts)
+            .expect_rules({ "tool-0" })
             .run();
 
         // Tool call, recipient in channel header
@@ -7111,6 +7142,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ special_function_tool })
             .expect(message_assist_call)
             .expect_reconstruction()
+            .expect_rules({ "tool-0" })
             .run();
 
         tst.test(
@@ -7139,6 +7171,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .tools({ special_function_tool })
             .expect(message_assist_call)
             .expect_reconstruction()
+            .expect_rules({ "tool-0" })
             .run();
 
         tst.test(
@@ -7159,6 +7192,13 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .run();
     }
 
+    // TranslateGemma
+    {
+        // no reconstruction check, the template adds whitespace around assistant content
+        auto tst = peg_tester("models/templates/google-translategemma-4b-it.jinja", detailed_debug);
+        tst.test("Hello, world!\nWhat's up?").expect(message_assist).run();
+    }
+
     // MiniCPM5 - XML tool calls with <function name="..."><param name="...">...</param></function>
     {
         auto tst = peg_tester("models/templates/openbmb-MiniCPM5-1B.jinja", detailed_debug);
@@ -7174,6 +7214,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
             .tools({ python_tool })
             .expect_tool_calls({ { "python", R"#({"code": "print('Hello, World!')"})#", {} } })
+            .expect_rules({ "tool-0" })
             .run();
 
         tst.test(R"(<function name="empty_args"></function>)")
@@ -7575,6 +7616,43 @@ static void test_developer_role_to_system_workaround() {
     }
 }
 
+// TranslateGemma raises on plain string user content, the specialized handler must rewrite it
+static void test_translate_gemma() {
+    LOG_DBG("%s\n", __func__);
+
+    auto tmpls = read_templates("models/templates/google-translategemma-4b-it.jinja");
+
+    // server startup renders this example, it must not throw
+    auto example = common_chat_format_example(tmpls.get(), /* use_jinja= */ true, {});
+    assert_contains(example, "English (en-GB) to English (en-GB) translator");
+    assert_contains(example, "How are you?");
+
+    common_chat_templates_inputs inputs;
+    inputs.messages                                 = { message_user };
+    inputs.add_generation_prompt                    = true;
+    inputs.chat_template_kwargs["source_lang_code"] = R"("en")";
+    inputs.chat_template_kwargs["target_lang_code"] = R"("fr")";
+
+    auto params = common_chat_templates_apply(tmpls.get(), inputs);
+    assert_contains(params.prompt, "English (en) to French (fr) translator");
+    assert_contains(params.prompt, "into French:\n\n\nHey there!<end_of_turn>\n");
+    assert_equals(std::string("<start_of_turn>model\n"), params.generation_prompt);
+    assert_ends_with(params.prompt, params.generation_prompt);
+
+    // typed text parts are joined into one item
+    inputs.messages = { message_user_parts };
+    params = common_chat_templates_apply(tmpls.get(), inputs);
+    assert_contains(params.prompt, "into French:\n\n\nHey\nthere<end_of_turn>\n");
+
+    // assistant prefill is appended after the generation prompt
+    inputs.messages               = { message_user, message_assist_prefill_content };
+    inputs.add_generation_prompt  = false;
+    inputs.continue_final_message = COMMON_CHAT_CONTINUATION_CONTENT;
+    params = common_chat_templates_apply(tmpls.get(), inputs);
+    assert_equals(std::string("<start_of_turn>model\nHello, "), params.generation_prompt);
+    assert_ends_with(params.prompt, "Hey there!<end_of_turn>\n<start_of_turn>model\nHello, ");
+}
+
 // Verify reasoning-trace retention rules in the DeepSeek-V4 template:
 // all traces are retained unless drop_thinking is true AND the conversation
 // has no tool calls, in which case only the last (after-final-user) trace is
@@ -7952,6 +8030,7 @@ int main(int argc, char ** argv) {
         test_tools_oaicompat_json_conversion();
         test_convert_responses_to_chatcmpl();
         test_developer_role_to_system_workaround();
+        test_translate_gemma();
         test_deepseek_v4_thinking_retention();
         test_deepseek_v4_tool_result_ordering();
         test_template_generation_prompt();

tests/test-llama-archs.cpp

@@ -495,7 +495,7 @@ static std::pair<llama_model_ptr, llama_context_ptr> get_model_and_ctx(
         struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const float stdev,
         const std::vector<ggml_backend_dev_t> & devs,
         const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false,
-        const llama_model_tensor_buft_override * tensor_buft_overrides = nullptr) {
+        const llama_model_tensor_buft_override * tensor_buft_overrides = nullptr, const size_t moe_cache_size = 0) {
     GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr));
     llama_model_params model_params = llama_model_default_params();
     model_params.progress_callback = silent_model_load_progress;
@@ -512,6 +512,11 @@ static std::pair<llama_model_ptr, llama_context_ptr> get_model_and_ctx(
     if (!encode) {
         ctx_params.n_ubatch = 64;
     }
+    if (moe_cache_size > 0) {
+        // the MoE cache is only used for small ubatches
+        ctx_params.moe_cache_size = moe_cache_size;
+        ctx_params.n_ubatch = 2;
+    }
 
     tensor_data_params tensor_params = { seed, stdev };
     llama_model_ptr model(gguf_ctx != nullptr ?
@@ -865,9 +870,10 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con
         std::string                     label;
         llama_split_mode                split_mode;
         bool                            host_experts; // keep the experts in host memory, see host_experts_test
+        size_t                          moe_cache_size;
 
-        device_config(std::vector<ggml_backend_dev_t> devs, std::string name, llama_split_mode split_mode, bool host_experts = false)
-            : devs(std::move(devs)), label(std::move(name)), split_mode(split_mode), host_experts(host_experts) {}
+        device_config(std::vector<ggml_backend_dev_t> devs, std::string name, llama_split_mode split_mode, bool host_experts = false, size_t moe_cache_size = 0)
+            : devs(std::move(devs)), label(std::move(name)), split_mode(split_mode), host_experts(host_experts), moe_cache_size(moe_cache_size) {}
     };
 
     const llama_model_tensor_buft_override host_experts_overrides[] = {
@@ -905,6 +911,16 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con
             dev_configs.emplace_back(devices_meta, "Host experts", LLAMA_SPLIT_MODE_LAYER, true);
             max_device_label_length = std::max(max_device_label_length, dev_configs.back().label.length());
         }
+
+        // the ops that use the host experts run on a GPU and read the experts from a cache
+        // the cache has only a few slots (4 for 288 KiB experts), so the experts are evicted and uploaded again
+        if (!devices_meta.empty()) {
+            const enum ggml_backend_dev_type type = ggml_backend_dev_type(devices_meta[0]);
+            if (type == GGML_BACKEND_DEVICE_TYPE_GPU || type == GGML_BACKEND_DEVICE_TYPE_IGPU) {
+                dev_configs.emplace_back(std::vector<ggml_backend_dev_t>{devices_meta[0]}, "MoE cache", LLAMA_SPLIT_MODE_LAYER, true, 1536*1024);
+                max_device_label_length = std::max(max_device_label_length, dev_configs.back().label.length());
+            }
+        }
     }
 
     size_t max_arch_name_length = 0;
@@ -987,7 +1003,7 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con
                     }
                     if (dc.split_mode != LLAMA_SPLIT_MODE_TENSOR || llm_arch_supports_sm_tensor(arch)) {
                         test_executed = true;
-                        model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, stdev, dc.devs, dc.split_mode, encode, overrides);
+                        model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, stdev, dc.devs, dc.split_mode, encode, overrides, dc.moe_cache_size);
                         logits_dev = get_logits(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), tokens, encode);
                         const double nmse_val = nmse(logits_cpu, logits_dev);
                         snprintf(nmse_str, sizeof(nmse_str), "(%.2e)", nmse_val);
@@ -1053,7 +1069,7 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con
                         ms.save(file);
                         rewind(file);
 
-                        auto model_and_ctx_roundtrip = get_model_and_ctx(nullptr, file, seed, stdev, dc.devs, dc.split_mode, encode, overrides);
+                        auto model_and_ctx_roundtrip = get_model_and_ctx(nullptr, file, seed, stdev, dc.devs, dc.split_mode, encode, overrides, dc.moe_cache_size);
                         const std::vector<float> logits_roundtrip = get_logits(
                             model_and_ctx_roundtrip.first.get(), model_and_ctx_roundtrip.second.get(), tokens, encode);
                         status_roundtrip = "\033[1;32mOK\033[0m";

tests/test-recurrent-state-rollback.cpp

@@ -131,7 +131,7 @@ static test_status test_multi_seq_split_replay(const common_params & params, lla
 
     constexpr uint32_t  n_seqs     = 2;
     constexpr uint32_t  n_ubatch   = 16;
-    constexpr uint32_t  n_prompt   = 19;
+    constexpr uint32_t  n_prompt   = 20;
     constexpr uint32_t  n_rollback = 3;
     constexpr uint32_t  n_replay   = 40; // > n_ubatch so each seq spans multiple ubatches
     constexpr llama_pos p0         = n_prompt - n_rollback;
@@ -174,13 +174,13 @@ static test_status test_multi_seq_split_replay(const common_params & params, lla
         return llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) == 0;
     };
 
-    // both contexts decode the identical [0, p0) prefill; only ctx_roll decodes
-    // the tail, which is then rolled back so its restore is pending at replay
+    // Keep the anchor at p0 - 1 when rolling back the tail.
     for (uint32_t s = 0; s < n_seqs && ok; ++s) {
-        ok = ok && decode_range(ctx_roll.get(), s, 0, (llama_pos) p0);
-        ok = ok && decode_range(ctx_ref.get(),  s, 0, (llama_pos) p0);
+        ok = ok && decode_range(ctx_roll.get(), s, 0, (llama_pos) p0 - 1);
+        ok = ok && decode_range(ctx_ref.get(),  s, 0, (llama_pos) p0 - 1);
 
-        ok = ok && decode_range(ctx_roll.get(), s, (llama_pos) p0, (llama_pos) n_prompt);
+        ok = ok && decode_range(ctx_roll.get(), s, (llama_pos) p0 - 1, (llama_pos) n_prompt);
+        ok = ok && decode_range(ctx_ref.get(),  s, (llama_pos) p0 - 1, (llama_pos) p0);
 
         ok = ok && llama_memory_seq_rm(llama_get_memory(ctx_roll.get()), (llama_seq_id) s, p0, -1);
 

tools/imatrix/imatrix.cpp

@@ -10,6 +10,7 @@
 
 #include <chrono>
 #include <climits>
+#include <clocale>
 #include <cmath>
 #include <cstring>
 #include <fstream>

tools/mtmd/clip-impl.h

@@ -58,6 +58,7 @@
 #define KEY_PATCH_SIZE              "clip.vision.patch_size"
 #define KEY_IMAGE_MEAN              "clip.vision.image_mean"
 #define KEY_IMAGE_STD               "clip.vision.image_std"
+#define KEY_IMAGE_RESIZE_ALGO       "clip.vision.image_resize_algo"
 #define KEY_PROJ_SCALE_FACTOR       "clip.vision.projector.scale_factor"
 #define KEY_PROJ_SAMPLE_QUERY_SIDE  "clip.vision.projector.query_side"
 #define KEY_PROJ_SAMPLE_WINDOW_SIDE "clip.vision.projector.window_side"
@@ -507,6 +508,7 @@ enum projector_type {
     PROJECTOR_TYPE_POCKETTTS_SPKENC,
     PROJECTOR_TYPE_POCKETTTS_GEN,
     PROJECTOR_TYPE_MUSE_GLIMMER,
+    PROJECTOR_TYPE_COHERE2V,
     PROJECTOR_TYPE_UNKNOWN,
 };
 
@@ -574,6 +576,7 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
     { PROJECTOR_TYPE_POCKETTTS_SPKENC,  "pockettts_spkenc"},
     { PROJECTOR_TYPE_POCKETTTS_GEN,     "pockettts_gen"},
     { PROJECTOR_TYPE_MUSE_GLIMMER,      "muse-glimmer"},
+    { PROJECTOR_TYPE_COHERE2V,          "cohere2v"},
 };
 
 static projector_type clip_projector_type_from_string(const std::string & str) {

tools/mtmd/clip.cpp

@@ -937,6 +937,7 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
     switch (ctx->proj_type()) {
         case PROJECTOR_TYPE_GEMMA3:
         case PROJECTOR_TYPE_IDEFICS3:
+        case PROJECTOR_TYPE_COHERE2V:
         case PROJECTOR_TYPE_LFM2:
         case PROJECTOR_TYPE_JANUS_PRO:
         case PROJECTOR_TYPE_PHI4:
@@ -1514,11 +1515,31 @@ struct clip_model_loader {
                         get_u32(KEY_PREPROC_IMAGE_SIZE, hparams.image_longest_edge, false);
                         hparams.set_limit_image_tokens();
                     } break;
+                case PROJECTOR_TYPE_COHERE2V:
+                    {
+                        hparams.image_pad_rf = PAD_NONE;
+                        get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge);
+                        get_u32(KEY_PREPROC_MAX_TILES, hparams.preproc_max_tiles);
+                        if (hparams.preproc_max_tiles <= 0 || hparams.preproc_max_tiles > 256) {
+                            throw std::runtime_error(string_format("%s: preproc_max_tiles (%d) must be in range [1, 256]\n", __func__, hparams.preproc_max_tiles));
+                        }
+                    } break;
                 case PROJECTOR_TYPE_LFM2:
                     {
-                        hparams.image_resize_algo    = RESIZE_ALGO_BILINEAR;
-                        hparams.image_resize_algo_rf = RESIZE_ALGO_BILINEAR;
-                        hparams.image_resize_algo_ov = RESIZE_ALGO_BILINEAR;
+                        // default for older GGUFs
+                        std::string resize_algo = "bilinear";
+                        get_string(KEY_IMAGE_RESIZE_ALGO, resize_algo, false);
+                        if (resize_algo == "bilinear") {
+                            hparams.image_resize_algo = RESIZE_ALGO_BILINEAR;
+                        } else if (resize_algo == "bicubic") {
+                            hparams.image_resize_algo = RESIZE_ALGO_BICUBIC;
+                        } else if (resize_algo == "lanczos") {
+                            hparams.image_resize_algo = RESIZE_ALGO_LANCZOS;
+                        } else {
+                            throw std::runtime_error(string_format("%s: unsupported image resize algo: %s\n", __func__, resize_algo.c_str()));
+                        }
+                        hparams.image_resize_algo_rf = hparams.image_resize_algo;
+                        hparams.image_resize_algo_ov = hparams.image_resize_algo;
                         get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge, false);
                         // ref: https://huggingface.co/LiquidAI/LFM2.5-VL-1.6B/blob/main/processor_config.json
                         hparams.set_limit_image_tokens(64, 256);
@@ -2770,6 +2791,13 @@ struct clip_model_loader {
                 {
                     model.mm_fc_w = get_tensor(string_format(TN_MM_PROJECTOR, "weight"));
                 } break;
+            case PROJECTOR_TYPE_COHERE2V:
+                {
+                    model.mm_1_w = get_tensor(string_format(TN_LLAVA_PROJ, 1, "weight"));
+                    model.mm_1_b = get_tensor(string_format(TN_LLAVA_PROJ, 1, "bias"));
+                    model.mm_2_w = get_tensor(string_format(TN_LLAVA_PROJ, 2, "weight"));
+                    model.mm_2_b = get_tensor(string_format(TN_LLAVA_PROJ, 2, "bias"));
+                } break;
             case PROJECTOR_TYPE_LFM2:
                 {
                     model.mm_input_norm_w = get_tensor(TN_MM_INP_NORM, false);
@@ -4233,6 +4261,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
         case PROJECTOR_TYPE_GEMMA4V:
         case PROJECTOR_TYPE_GEMMA4UV:
         case PROJECTOR_TYPE_IDEFICS3:
+        case PROJECTOR_TYPE_COHERE2V:
         case PROJECTOR_TYPE_INTERNVL:
         case PROJECTOR_TYPE_NEMOTRON_V2_VL:
         case PROJECTOR_TYPE_LLAMA4:
@@ -5321,6 +5350,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
         case PROJECTOR_TYPE_GEMMA3:
         case PROJECTOR_TYPE_GEMMA3NV:
         case PROJECTOR_TYPE_IDEFICS3:
+        case PROJECTOR_TYPE_COHERE2V:
         case PROJECTOR_TYPE_INTERNVL:
         case PROJECTOR_TYPE_NEMOTRON_V2_VL:
         case PROJECTOR_TYPE_QWEN2A:
@@ -6055,6 +6085,7 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
         case PROJECTOR_TYPE_KIMIK25:
         case PROJECTOR_TYPE_YASA2:
         case PROJECTOR_TYPE_DEEPSEEK4V:
+        case PROJECTOR_TYPE_COHERE2V:
             return ctx->model.mm_2_w->ne[1];
         case PROJECTOR_TYPE_HUNYUANVL:
             return ctx->model.mm_model_proj->ne[1];