ui: mask API Key field in settings and error splash to stop browser a… (#26562)

* ui: mask API Key field in settings and error splash to stop browser autofill

* ui: set autocomplete=new-password on private fields

The password input type makes browsers offer to save the API key
in the password manager and autofill saved site credentials into
the field. The new-password autocomplete value disables both.

---------

Co-authored-by: Pascal <admin@serveurperso.com>

common/arg.cpp

@@ -3646,6 +3646,18 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             }
         }
     ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_REASONING"));
+    add_opt(common_arg(
+        {"--reasoning-effort"}, "LEVEL",
+        "reasoning effort level given to the chat template: 'default' to keep the template default,\n"
+        "or a level such as 'minimal', 'low', 'medium', 'high', 'xhigh' or 'max' (default: default)",
+        [](common_params & params, const std::string & value) {
+            if (value == "default") {
+                params.default_template_kwargs.erase("reasoning_effort");
+            } else {
+                params.default_template_kwargs["reasoning_effort"] = json(value).dump();
+            }
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_REASONING_EFFORT"));
     add_opt(common_arg(
         {"--reasoning-budget"}, "N",
         "token budget for thinking: -1 for unrestricted, 0 for immediate end, N>0 for token budget (default: -1)",
@@ -4065,6 +4077,9 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         {"--spec-draft-n-max"}, "N",
         string_format("number of tokens to draft for speculative decoding (default: %d)", params.speculative.draft.n_max),
         [](common_params & params, int value) {
+            if (value < 0) {
+                throw std::invalid_argument("invalid value");
+            }
             params.speculative.draft.n_max = value;
         }
     ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MAX"));

common/chat.cpp

@@ -920,6 +920,10 @@ static std::string common_chat_template_direct_apply_impl(
         bool enabled = inp["preserve_reasoning"].get<bool>();
         jinja::caps_apply_preserve_reasoning(ctx, enabled);
     }
+    if (inp.contains("reasoning_effort") && inp["reasoning_effort"].is_string() && !inp["reasoning_effort"].empty()) {
+        std::string reasoning_effort = inp["reasoning_effort"].get<std::string>();
+        jinja::caps_apply_reasoning_effort(ctx, reasoning_effort);
+    }
 
     jinja::global_from_json(ctx, inp, inputs.mark_input);
 
@@ -2321,6 +2325,179 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha
     return data;
 }
 
+// Kimi K3 - XTML tagged format, built by open_tag/close_tag macros:
+//   open_tag(t, attrs) = <|open|>t k="v"...<|sep|>   close_tag(t) = <|close|>t<|sep|>
+//   assistant := [think] [response] [tools] close_tag(message) <|end_of_msg|>
+// the generation prompt already opens the think (or response) section, so the
+// section opener is optional here - same as Kimi K2 Thinking
+static common_chat_params common_chat_params_init_kimi_k3(const common_chat_template &          tmpl,
+                                                          const autoparser::generation_params & inputs) {
+    common_chat_params data;
+
+    data.prompt            = common_chat_template_direct_apply_impl(tmpl, inputs);
+    data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs);
+    data.format            = COMMON_CHAT_FORMAT_PEG_NATIVE;
+    data.supports_thinking = true;
+
+    const std::string SEP         = "<|sep|>";
+    const std::string MSG_START   = "<|open|>message role=\"assistant\"<|sep|>";
+    const std::string THINK_START = "<|open|>think<|sep|>";
+    const std::string THINK_END   = "<|close|>think<|sep|>";
+    const std::string RESP_START  = "<|open|>response<|sep|>";
+    const std::string RESP_END    = "<|close|>response<|sep|>";
+    const std::string TOOLS_START = "<|open|>tools<|sep|>";
+    const std::string TOOLS_END   = "<|close|>tools<|sep|>";
+    const std::string CALL_START  = "<|open|>call tool=\"";
+    const std::string CALL_END    = "<|close|>call<|sep|>";
+    const std::string ARG_START   = "<|open|>argument key=\"";
+    const std::string ARG_END     = "<|close|>argument<|sep|>";
+    const std::string MSG_END     = "<|close|>message<|sep|>";
+    const std::string EOM_TOKEN   = "<|end_of_msg|>";
+
+    // only the markers are special tokens. tag names ("think", "response", ...) are
+    // normal tokens and must not be preserved, or prose with those words is broken
+    data.preserved_tokens = {
+        "<|open|>",
+        "<|close|>",
+        "<|sep|>",
+        "<|end_of_msg|>",
+    };
+
+    data.thinking_start_tag = THINK_START;
+    data.thinking_end_tags  = { THINK_END };
+
+    // per-role message-start delimiters. user/assistant messages only have the role
+    // attribute, so the full opener is used. system and tool messages have more
+    // attributes, so those delimiters stop after the closing quote of the role
+    data.message_delimiters = {
+        { COMMON_CHAT_ROLE_ASSISTANT, "<|open|>message role=\"assistant\"<|sep|>" },
+        { COMMON_CHAT_ROLE_USER,      "<|open|>message role=\"user\"<|sep|>"      },
+        { COMMON_CHAT_ROLE_TOOL,      "<|open|>message role=\"tool\""             },
+        { COMMON_CHAT_ROLE_SYSTEM,    "<|open|>message role=\"system\""           },
+    };
+
+    auto has_tools         = inputs.tools.is_array() && !inputs.tools.empty();
+    auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE;
+    auto include_grammar   = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE;
+
+    if (inputs.has_continuation()) {
+        const auto & msg = inputs.continue_msg;
+
+        data.generation_prompt = MSG_START + THINK_START + msg.reasoning_content;
+        if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) {
+            data.generation_prompt += THINK_END + RESP_START + msg.render_content();
+        }
+
+        data.prompt += data.generation_prompt;
+    }
+
+    auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
+        auto end = p.end();
+
+        auto start = p.optional(p.literal(MSG_START));
+
+        // the think section is always consumed, even with reasoning extraction off:
+        // the generation prompt ends with open_tag('think'), so it is always present.
+        // reasoning stops at its own closer, or at the response opener if the model
+        // skips the closer
+        auto think_body = extract_reasoning ? p.reasoning(p.until_one_of({ THINK_END, RESP_START })) :
+                                              p.content(p.until_one_of({ THINK_END, RESP_START }));
+
+        auto reasoning = p.optional(p.optional(p.literal(THINK_START)) + think_body +
+                                    p.optional(p.literal(THINK_END)));
+
+        // content runs to the response closer, or to the next section if truncated
+        auto response = p.optional(p.literal(RESP_START)) +
+                        p.content(p.until_one_of({ RESP_END, TOOLS_START, MSG_END })) +
+                        p.optional(p.literal(RESP_END));
+
+        // the EOG token after the message closer reaches the parser as text,
+        // so it must be consumed or the parse stays incomplete
+        auto trailer = p.optional(p.literal(MSG_END)) + p.optional(p.literal(EOM_TOKEN));
+
+        if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) {
+            return start + reasoning + response + trailer + end;
+        }
+
+        auto tool_choices = p.choice();
+        foreach_function(inputs.tools, [&](const json & tool) {
+            const auto & function = tool.at("function");
+            std::string  name     = function.at("name");
+            const json   schema   = function.contains("parameters") ? function.at("parameters") : json::object();
+
+            // arguments come one tag per key, with the JSON type in a type="..."
+            // attribute. the type is taken from the tool schema instead, as it tells
+            // us if the value is JSON or a literal string
+            auto args = p.eps();
+            if (schema.contains("properties") && !schema.at("properties").empty()) {
+                auto arg_choices = p.choice();
+                for (const auto & prop : schema.at("properties").items()) {
+                    const std::string & key = prop.key();
+
+                    std::string type = "string";
+                    if (prop.value().is_object() && prop.value().contains("type") &&
+                        prop.value().at("type").is_string()) {
+                        type = prop.value().at("type").get<std::string>();
+                    }
+
+                    auto value = type == "string" ? p.tool_arg_string_value(p.until(ARG_END)) :
+                                                    p.tool_arg_value(p.until(ARG_END));
+
+                    // skip the trailing type="..." attribute: anything up to <|sep|>
+                    arg_choices |= p.rule("kimi-k3-arg-" + name + "-" + key,
+                                          p.tool_arg(p.tool_arg_open(p.literal(ARG_START)) +
+                                                     p.tool_arg_name(p.literal(key)) + p.literal("\"") +
+                                                     p.until(SEP) + p.literal(SEP) + value +
+                                                     p.tool_arg_close(p.literal(ARG_END))));
+                }
+                args = p.zero_or_more(arg_choices);
+            }
+
+            // skip the trailing index="N" attribute the same way
+            auto call = p.tool(p.tool_open(p.literal(CALL_START) + p.tool_name(p.literal(name)) + p.literal("\"") +
+                                           p.until(SEP) + p.literal(SEP)) +
+                               p.tool_args(args) + p.tool_close(p.literal(CALL_END)));
+
+            tool_choices |= p.rule("kimi-k3-tool-" + name, call);
+        });
+
+        // all calls go inside one tools section, then the message is closed. the
+        // message closer is part of the trigger rule, or else the lazy grammar
+        // rejects it once tool calls have started
+        auto tools_section =
+            p.trigger_rule("kimi-k3-tool-call", p.literal(TOOLS_START) + p.one_or_more(tool_choices) +
+                                                    p.literal(TOOLS_END) + p.optional(p.literal(MSG_END)) +
+                                                    p.optional(p.literal(EOM_TOKEN)));
+
+        auto tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? tools_section :
+                                                                              p.optional(tools_section);
+
+        return start + reasoning + response + tools + trailer + end;
+    });
+
+    data.parser = parser.save();
+
+    if (include_grammar) {
+        data.grammar_lazy = inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_REQUIRED;
+        data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
+            foreach_function(inputs.tools, [&](const json & tool) {
+                const auto & function = tool.at("function");
+                if (function.contains("parameters")) {
+                    auto schema = function.at("parameters");
+                    builder.resolve_refs(schema);
+                }
+            });
+            parser.build_grammar(builder, data.grammar_lazy);
+        });
+
+        data.grammar_triggers = {
+            { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, TOOLS_START },
+        };
+    }
+
+    return data;
+}
+
 // Cohere2 MoE (a.k.a. "North Code") parser.
 //
 // The assistant turn is fully marker-wrapped:
@@ -3289,6 +3466,13 @@ std::optional<common_chat_params> common_chat_try_specialized_template(
         return common_chat_params_init_kimi_k2(tmpl, params);
     }
 
+    // Kimi K3 - the <|open|>/<|close|>/<|end_of_msg|> markers are unique to it
+    if (src.find("<|open|>") != std::string::npos && src.find("<|close|>") != std::string::npos &&
+        src.find("<|end_of_msg|>") != std::string::npos) {
+        LOG_DBG("Using specialized template: Kimi K3\n");
+        return common_chat_params_init_kimi_k3(tmpl, params);
+    }
+
     // Cohere2 MoE / North Code - marker-wrapped format with <|START_TEXT|> content and
     // <|START_ACTION|> JSON tool calls. <|START_TEXT|> is unique to this template (the older
     // Command-R templates use <|START_RESPONSE|>).

common/common.cpp

@@ -1275,6 +1275,8 @@ struct common_init_result::impl {
 
     // note: the order in which model, context, etc. are declared matters because their destructors will be called bottom-to-top
 
+    common_threadpools threadpools;
+
     llama_model_ptr   model;
     llama_context_ptr context;
 
@@ -1376,6 +1378,10 @@ common_init_result::common_init_result(common_params & params, bool model_only)
     }
 
     pimpl->context.reset(lctx);
+
+    set_process_priority(params.cpuparams.priority);
+
+    pimpl->threadpools.init(lctx, params);
 }
 
 llama_model * common_init_result::model() {
@@ -1724,6 +1730,10 @@ struct llama_context_params common_context_params_to_llama(const common_params &
     return cparams;
 }
 
+//
+// Threadpool utils
+//
+
 struct ggml_threadpool_params ggml_threadpool_params_from_cpu_params(const common_cpu_params & params) {
     struct ggml_threadpool_params tpp;
 
@@ -1740,6 +1750,56 @@ struct ggml_threadpool_params ggml_threadpool_params_from_cpu_params(const commo
     return tpp;
 }
 
+common_threadpools::~common_threadpools() {
+    if (!free_fn) {
+        return;
+    }
+    free_fn(threadpool);
+    free_fn(threadpool_batch);
+}
+
+void common_threadpools::init(llama_context * ctx, const common_params & params) {
+    GGML_ASSERT(!threadpool);
+    GGML_ASSERT(!threadpool_batch);
+
+    COM_INF("llama threadpool init, n_threads = %d\n", (int) params.cpuparams.n_threads);
+
+    auto * cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
+    if (!cpu_dev) {
+        COM_WRN("%s", "no CPU backend found\n");
+        return;
+    }
+    auto * reg = ggml_backend_dev_backend_reg(cpu_dev);
+    auto * ggml_threadpool_new_fn = (decltype(ggml_threadpool_new) *) ggml_backend_reg_get_proc_address(reg, "ggml_threadpool_new");
+    free_fn = (decltype(ggml_threadpool_free) *) ggml_backend_reg_get_proc_address(reg, "ggml_threadpool_free");
+
+    struct ggml_threadpool_params tpp_batch =
+            ggml_threadpool_params_from_cpu_params(params.cpuparams_batch);
+    struct ggml_threadpool_params tpp =
+            ggml_threadpool_params_from_cpu_params(params.cpuparams);
+
+    if (!ggml_threadpool_params_match(&tpp, &tpp_batch)) {
+        threadpool_batch = ggml_threadpool_new_fn(&tpp_batch);
+        if (!threadpool_batch) {
+            COM_WRN("batch threadpool create failed : n_threads %d\n", tpp_batch.n_threads);
+            return;
+        }
+
+        // start the non-batch threadpool in the paused state
+        tpp.paused = true;
+    }
+
+    threadpool = ggml_threadpool_new_fn(&tpp);
+    if (!threadpool) {
+        COM_WRN("threadpool create failed : n_threads %d\n", tpp.n_threads);
+        free_fn(threadpool_batch);
+        threadpool_batch = nullptr;
+        return;
+    }
+
+    llama_attach_threadpool(ctx, threadpool, threadpool_batch);
+}
+
 //
 // Batch utils
 //

common/common.h

@@ -929,9 +929,8 @@ using common_init_result_ptr = std::unique_ptr<common_init_result>;
 
 common_init_result_ptr common_init_from_params(common_params & params, bool model_only = false);
 
-struct llama_model_params     common_model_params_to_llama  (      common_params & params);
-struct llama_context_params   common_context_params_to_llama(const common_params & params);
-struct ggml_threadpool_params ggml_threadpool_params_from_cpu_params(const common_cpu_params & params);
+struct llama_model_params   common_model_params_to_llama  (      common_params & params);
+struct llama_context_params common_context_params_to_llama(const common_params & params);
 
 // clear LoRA adapters from context, then apply new list of adapters
 void common_set_adapter_lora(struct llama_context * ctx, std::vector<common_adapter_lora_info> & lora);
@@ -942,6 +941,28 @@ std::string common_get_model_endpoint();
 // for testing purposes
 char * common_get_model_or_exit(int, char*[]);
 
+//
+// Threadpool utils
+//
+
+struct ggml_threadpool_params ggml_threadpool_params_from_cpu_params(const common_cpu_params & params);
+
+struct common_threadpools {
+    common_threadpools() = default;
+    ~common_threadpools();
+
+    common_threadpools(const common_threadpools &) = delete;
+    common_threadpools & operator=(const common_threadpools &) = delete;
+
+    void init(llama_context * ctx, const common_params & params);
+
+private:
+    ggml_threadpool * threadpool       = nullptr;
+    ggml_threadpool * threadpool_batch = nullptr;
+
+    decltype(ggml_threadpool_free) * free_fn = nullptr;
+};
+
 //
 // Context utils
 //

common/imatrix-loader.cpp

@@ -102,7 +102,8 @@ bool common_imatrix_load(const std::string & fname, common_imatrix & imatrix) {
     const int64_t chunk_count_key = gguf_find_key(ctx_gguf, LLM_KV_IMATRIX_CHUNK_COUNT);
     const int64_t chunk_size_key  = gguf_find_key(ctx_gguf, LLM_KV_IMATRIX_CHUNK_SIZE);
 
-    if (datasets_key != -1 && gguf_get_arr_type(ctx_gguf, datasets_key) == GGUF_TYPE_STRING) {
+    if (datasets_key != -1 && gguf_get_kv_type(ctx_gguf, datasets_key) == GGUF_TYPE_ARRAY &&
+        gguf_get_arr_type(ctx_gguf, datasets_key) == GGUF_TYPE_STRING) {
         const int64_t n = gguf_get_arr_n(ctx_gguf, datasets_key);
         imatrix.datasets.reserve(imatrix.datasets.size() + n);
         for (int64_t i = 0; i < n; ++i) {
@@ -143,6 +144,13 @@ bool common_imatrix_load(const std::string & fname, common_imatrix & imatrix) {
             return false;
         }
 
+        if (in_sum2->type != GGML_TYPE_F32 || counts->type != GGML_TYPE_F32) {
+            LOG_ERR("%s: sums and counts for %s must be F32\n", __func__, name.c_str());
+            gguf_free(ctx_gguf);
+            ggml_free(ctx);
+            return false;
+        }
+
         auto & e = imatrix.entries[name];
 
         const int64_t nval    = ggml_nelements(in_sum2);

common/jinja/caps.cpp

@@ -17,7 +17,7 @@ namespace jinja {
 
 using caps_json_fn = std::function<json()>;
 using caps_ctx_fn = std::function<void(context &)>;
-using caps_analyze_fn = std::function<void(bool, value &, value &, const std::string &)>;
+using caps_analyze_fn = std::function<void(context &, bool, value &, value &, const std::string &)>;
 
 void caps_apply_preserve_reasoning(jinja::context & ctx, bool enabled) {
     ctx.set_val("preserve_thinking",         mk_val<value_bool>(enabled));
@@ -26,6 +26,12 @@ void caps_apply_preserve_reasoning(jinja::context & ctx, bool enabled) {
     ctx.set_val("drop_thinking",            mk_val<value_bool>(!enabled));
 }
 
+void caps_apply_reasoning_effort(jinja::context & ctx, const std::string & effort) {
+    value var = mk_val<value_string>(effort); // bind to the same value for stats
+    ctx.set_val("reasoning_effort",   var);
+    ctx.set_val("reasoning_strength", var);
+}
+
 static void caps_try_execute(jinja::program & prog,
                              const caps_json_fn & messages_fn,
                              const caps_ctx_fn & ctx_fn,
@@ -62,7 +68,7 @@ static void caps_try_execute(jinja::program & prog,
         // ignore exceptions during capability analysis
     }
 
-    analyze_fn(success, messages, tools, result);
+    analyze_fn(ctx, success, messages, tools, result);
 }
 
 // for debugging only
@@ -87,6 +93,7 @@ std::map<std::string, bool> caps::to_map() const {
         {"supports_parallel_tool_calls", supports_parallel_tool_calls},
         {"supports_system_role", supports_system_role},
         {"supports_preserve_reasoning", supports_preserve_reasoning},
+        {"supports_reasoning_effort", supports_reasoning_effort},
         {"supports_object_arguments", supports_object_arguments},
     };
 }
@@ -124,7 +131,7 @@ caps caps_get(jinja::program & prog) {
         },
         nullptr, // ctx_fn
         nullptr, // tools_fn
-        [&](bool success, value & messages, value &, const std::string &) {
+        [&](context &, bool success, value & messages, value &, const std::string &) {
             auto & content = messages->at(0)->at("content");
             caps_print_stats(content, "messages[0].content");
             if (has_op(content, "selectattr") || has_op(content, "array_access")) {
@@ -158,7 +165,7 @@ caps caps_get(jinja::program & prog) {
         },
         nullptr, // ctx_fn
         nullptr, // tools_fn
-        [&](bool, value & messages, value &, const std::string &) {
+        [&](context &, bool, value & messages, value &, const std::string &) {
             auto & content = messages->at(0)->at("content");
             caps_print_stats(content, "messages[0].content");
             if (!content->stats.used) {
@@ -234,7 +241,7 @@ caps caps_get(jinja::program & prog) {
                 },
             });
         },
-        [&](bool success, value & messages, value & tools, const std::string &) {
+        [&](context &, bool success, value & messages, value & tools, const std::string &) {
             if (!success) {
                 return; // Nothing can be inferred
             }
@@ -327,7 +334,7 @@ caps caps_get(jinja::program & prog) {
                     },
                 });
             },
-            [&](bool success, value & messages, value & tools, const std::string &) {
+            [&](context &, bool success, value & messages, value & tools, const std::string &) {
                 if (!success) {
                     result.supports_tool_calls = false;
                     result.supports_tools = false;
@@ -429,7 +436,7 @@ caps caps_get(jinja::program & prog) {
                 },
             });
         },
-        [&](bool success, value & messages, value &, const std::string &) {
+        [&](context &, bool success, value & messages, value &, const std::string &) {
             if (!success) {
                 result.supports_parallel_tool_calls = false;
                 return;
@@ -486,7 +493,7 @@ caps caps_get(jinja::program & prog) {
             caps_apply_preserve_reasoning(ctx, true);
         },
         nullptr, // tools_fn
-        [&](bool, value &, value &, const std::string & output) {
+        [&](context &, bool, value &, value &, const std::string & output) {
             // note: we cannot use stats here because the reasoning_content may be used for "if" condition test, but not actually outputted in the final result
             if (output.find(reasoning_placeholder) != std::string::npos) {
                 result.supports_preserve_reasoning = true;
@@ -494,6 +501,32 @@ caps caps_get(jinja::program & prog) {
         }
     );
 
+    JJ_DEBUG("%s\n", ">>> Running capability check: reasoning effort");
+
+    // case: reasoning effort level
+    caps_try_execute(
+        prog,
+        [&]() {
+            // messages
+            return json::array({
+                {
+                    {"role", "user"},
+                    {"content", "User message"}
+                },
+            });
+        },
+        [&](context & ctx) {
+            ctx.set_val("enable_thinking", mk_val<value_bool>(true));
+            caps_apply_reasoning_effort(ctx, "low");
+        },
+        nullptr, // tools_fn
+        [&](context & ctx, bool, value &, value &, const std::string &) {
+            value effort = ctx.get_val("reasoning_effort");
+            caps_print_stats(effort, "reasoning_effort");
+            result.supports_reasoning_effort = effort->stats.used;
+        }
+    );
+
     JJ_DEBUG("%s\n", result.to_string().c_str());
 
     return result;

common/jinja/caps.h

@@ -16,6 +16,9 @@ struct caps {
     // supports preserve reasoning trace in the full history, not just the last assistant message
     bool supports_preserve_reasoning = false;
 
+    // supports reasoning effort levels
+    bool supports_reasoning_effort = false;
+
     // one of the 2 content capabilities must be true
     bool supports_string_content = true;
     bool supports_typed_content = false;
@@ -32,5 +35,6 @@ struct caps {
 caps caps_get(jinja::program & prog);
 
 void caps_apply_preserve_reasoning(jinja::context & ctx, bool enabled);
+void caps_apply_reasoning_effort(jinja::context & ctx, const std::string & effort);
 
 } // namespace jinja

common/jinja/runtime.cpp

@@ -263,7 +263,7 @@ value binary_expression::execute_impl(context & ctx) {
             return res;
         }
         for (int64_t i = 0; i < repeat; ++i) {
-            res->val_str = res->val_str.append(str);
+            res->val_str.append(str);
         }
         return res;
     }

common/jinja/runtime.h

@@ -763,14 +763,22 @@ struct runtime {
         gather_string_parts_recursive(val, parts);
         // join consecutive parts with the same type
         auto & p = parts->val_str.parts;
-        for (size_t i = 1; i < p.size(); ) {
-            if (p[i].is_input == p[i - 1].is_input) {
-                p[i - 1].val += p[i].val;
-                p.erase(p.begin() + i);
+        if (p.empty()) {
+            return parts;
+        }
+        size_t w = 0;
+        for (size_t r = 1; r < p.size(); r++) {
+            if (p[w].is_input == p[r].is_input) {
+                p[w].val += p[r].val;
             } else {
-                i++;
+                w++;
+                if (w != r) {
+                    // the guard is needed, self-move leaves the string in an unspecified state
+                    p[w] = std::move(p[r]);
+                }
             }
         }
+        p.resize(w + 1);
         return parts;
     }
 

common/jinja/string.cpp

@@ -103,7 +103,7 @@ void string::mark_input_based_on(const string & other) {
     }
 }
 
-string string::append(const string & other) {
+string & string::append(const string & other) {
     for (const auto & part : other.parts) {
         parts.push_back(part);
     }

common/jinja/string.h

@@ -47,7 +47,7 @@ struct string {
     // mark this string as input if other has ALL parts as input
     void mark_input_based_on(const string & other);
 
-    string append(const string & other);
+    string & append(const string & other);
 
     // in-place transformations
 

common/preset.cpp

@@ -365,8 +365,25 @@ struct local_model {
     std::string name;
     std::string path;
     std::string path_mmproj;
+    std::string path_draft;
 };
 
+// TODO @ngxson: handle "eagle3-" when it's supported by common_speculative_types_from_gguf()
+static const char * draft_prefixes[] = { "mtp-", "dspark-", "dflash-" };
+
+static bool is_mmproj_file(const std::string & fname) {
+    return fname.find("mmproj") != std::string::npos;
+}
+
+static bool is_draft_file(const std::string & fname) {
+    for (const auto & prefix : draft_prefixes) {
+        if (fname.rfind(prefix, 0) == 0) {
+            return true;
+        }
+    }
+    return false;
+}
+
 common_presets common_preset_context::load_from_models_dir(const std::string & models_dir) const {
     if (!std::filesystem::exists(models_dir) || !std::filesystem::is_directory(models_dir)) {
         throw std::runtime_error(string_format("error: '%s' does not exist or is not a directory\n", models_dir.c_str()));
@@ -378,10 +395,15 @@ common_presets common_preset_context::load_from_models_dir(const std::string & m
         common_file_info model_file;
         common_file_info first_shard_file;
         common_file_info mmproj_file;
+        common_file_info draft_file;
         for (const auto & file : files) {
             if (string_ends_with(file.name, ".gguf")) {
-                if (file.name.find("mmproj") != std::string::npos) {
+                if (is_mmproj_file(file.name)) {
                     mmproj_file = file;
+                } else if (is_draft_file(file.name)) {
+                    if (draft_file.path.empty()) {
+                        draft_file = file; // first sidecar found wins
+                    }
                 } else if (file.name.find("-00001-of-") != std::string::npos) {
                     first_shard_file = file;
                 } else {
@@ -393,7 +415,8 @@ common_presets common_preset_context::load_from_models_dir(const std::string & m
         local_model model{
             /* name        */ name,
             /* path        */ first_shard_file.path.empty() ? model_file.path : first_shard_file.path,
-            /* path_mmproj */ mmproj_file.path // can be empty
+            /* path_mmproj */ mmproj_file.path, // can be empty
+            /* path_draft  */ draft_file.path   // can be empty
         };
         if (!model.path.empty()) {
             models.push_back(model);
@@ -405,13 +428,17 @@ common_presets common_preset_context::load_from_models_dir(const std::string & m
         if (file.is_dir) {
             scan_subdir(file.path, file.name);
         } else if (string_ends_with(file.name, ".gguf")) {
+            if (is_mmproj_file(file.name) || is_draft_file(file.name)) {
+                continue; // companion file, cannot be loaded as a model on its own
+            }
             // single file model
             std::string name = file.name;
             string_replace_all(name, ".gguf", "");
             local_model model{
                 /* name        */ name,
                 /* path        */ file.path,
-                /* path_mmproj */ ""
+                /* path_mmproj */ "",
+                /* path_draft  */ ""
             };
             models.push_back(model);
         }
@@ -426,6 +453,9 @@ common_presets common_preset_context::load_from_models_dir(const std::string & m
         if (!model.path_mmproj.empty()) {
             preset.set_option(*this, "LLAMA_ARG_MMPROJ", model.path_mmproj);
         }
+        if (!model.path_draft.empty()) {
+            preset.set_option(*this, "LLAMA_ARG_SPEC_DRAFT_MODEL", model.path_draft);
+        }
         out[preset.name] = preset;
     }
 

conversion/__init__.py

@@ -125,6 +125,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "JinaEmbeddingsV5Model": "bert",
     "KORMoForCausalLM": "qwen",
     "KimiK25ForConditionalGeneration": "deepseek",
+    "KimiK3ForConditionalGeneration": "kimi_k3",
     "KimiLinearForCausalLM": "kimi_linear",
     "KimiLinearModel": "kimi_linear",
     "KimiVLForConditionalGeneration": "deepseek",
@@ -161,6 +162,8 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "MiniCPM3ForCausalLM": "minicpm",
     "MiniCPMForCausalLM": "minicpm",
     "MiniCPMV4_6ForConditionalGeneration": "minicpm",
+    "MiniMaxText01ForCausalLM": "minimax",
+    "MiniMaxM1ForCausalLM": "minimax",
     "MiniMaxM2ForCausalLM": "minimax",
     "MiniMaxM3SparseForCausalLM": "minimax",
     "MiniMaxM3SparseForConditionalGeneration": "minimax",

conversion/base.py

@@ -658,6 +658,43 @@ class ModelBase:
     def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
         return ()
 
+    @staticmethod
+    def repack_mxfp4_blocks(packed: Tensor, scale: Tensor) -> np.ndarray:
+        """
+        Repack 4-bit MX weights into ggml `block_mxfp4`. Lossless - only moves bits.
+
+        Source (compressed-tensors "mxfp4-pack-quantized", also used by DeepSeek-V4):
+          packed  uint8 [rows, cols/2]  element 2i in the low nibble, 2i+1 in the high one
+          scale   uint8 [rows, cols/32] one E8M0 biased exponent per 32-element group
+
+        Destination, per group: one scale byte then 16 code bytes, where byte j holds
+        element j in the low nibble and element j+16 in the high one.
+
+        The 4-bit codes need no remapping: both sides index into ggml's kvalues_mxfp4
+        order. ggml doubles the kvalues and halves the scale, so the value is the same.
+        """
+        p = packed.contiguous().view(torch.uint8)
+        s = scale.contiguous().view(torch.uint8)
+
+        rows, packed_cols = p.shape
+        cols = packed_cols * 2
+        if cols % 32 != 0:
+            raise ValueError(f"MXFP4 source row has {cols} values, expected a multiple of 32")
+
+        n_blocks = cols // 32
+        if tuple(s.shape) != (rows, n_blocks):
+            raise ValueError(f"MXFP4 scale shape {tuple(s.shape)} does not match {(rows, n_blocks)}")
+
+        src = p.reshape(rows, n_blocks, 16)
+        lo = src & 0x0F           # elements 0, 2, 4, ...
+        hi = (src >> 4) & 0x0F    # elements 1, 3, 5, ...
+
+        vals = torch.stack((lo, hi), dim=-1).reshape(rows, n_blocks, 32)
+        qs = vals[:, :, :16] | (vals[:, :, 16:] << 4)
+
+        raw = torch.cat((s.unsqueeze(-1), qs.to(torch.uint8)), dim=-1)
+        return raw.reshape(rows, n_blocks * 17).cpu().numpy()
+
     @staticmethod
     def _nvfp4_pack(weight: Tensor, scale: Tensor) -> tuple[np.ndarray, list[int]]:
         """Repack NVFP4 ModelOpt tensors into ggml super-block layout.
@@ -2661,7 +2698,10 @@ def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> st
     # Step3-VL keeps text config under text_config but uses a custom top-level architecture.
     # For text conversion we route to a dedicated text-only class.
     # TODO: refactor this later to avoid adding exception here
-    if model_type == ModelType.TEXT and arch in ("StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration"):
+    # Kimi-K3's text_config reports "KimiLinearForCausalLM", which is the older
+    # Kimi-Linear-48B architecture and cannot load K3 (no attention residuals,
+    # latent MoE, situ, ...). Route on the top-level architecture instead.
+    if model_type == ModelType.TEXT and arch in ("StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration"):
         return arch
 
     # if "architectures" is found in the sub-config, use that instead

conversion/deepseek.py

@@ -709,31 +709,6 @@ class DeepseekV4Model(TextModel):
         for name in tensors_to_remove:
             del self.model_tensors[name]
 
-    @staticmethod
-    def _pack_mxfp4_blocks(weight: Tensor, scale: Tensor) -> np.ndarray:
-        packed = weight.contiguous().view(torch.uint8)
-        scale_u8 = scale.contiguous().view(torch.uint8)
-
-        out_features, packed_cols = packed.shape
-        logical_cols = packed_cols * 2
-        if logical_cols % 32 != 0:
-            raise ValueError(f"MXFP4 source row has {logical_cols} values, expected a multiple of 32")
-
-        n_blocks = logical_cols // 32
-        if tuple(scale_u8.shape) != (out_features, n_blocks):
-            raise ValueError(f"MXFP4 scale shape {tuple(scale_u8.shape)} does not match {(out_features, n_blocks)}")
-
-        src = packed.reshape(out_features, n_blocks, 16)
-        low = src & 0x0F
-        high = (src >> 4) & 0x0F
-
-        # The safetensors bytes store adjacent values as low/high nibbles.
-        # ggml MXFP4 blocks store values 0..15 in low nibbles and 16..31 in high nibbles.
-        vals = torch.stack((low, high), dim=-1).reshape(out_features, n_blocks, 32)
-        qs = vals[:, :, :16] | (vals[:, :, 16:] << 4)
-        raw = torch.cat((scale_u8.unsqueeze(-1), qs.to(torch.uint8)), dim=-1)
-        return raw.reshape(out_features, n_blocks * 17).cpu().numpy()
-
     def _write_mxfp4_expert_tensor(self, bid: int, proj: str, tensor_key: gguf.MODEL_TENSOR) -> list[str]:
         n_experts = self.hparams["n_routed_experts"]
         data: np.ndarray | None = None
@@ -747,7 +722,7 @@ class DeepseekV4Model(TextModel):
 
             weight = LazyTorchTensor.to_eager(self.model_tensors[weight_name]())
             scale = LazyTorchTensor.to_eager(self.model_tensors[scale_name]())
-            packed = self._pack_mxfp4_blocks(weight, scale)
+            packed = self.repack_mxfp4_blocks(weight, scale)
             if data is None:
                 data = np.empty((n_experts, *packed.shape), dtype=packed.dtype)
             data[eid] = packed

conversion/kimi_k3.py

@@ -0,0 +1,375 @@
+from __future__ import annotations
+
+import re
+from pathlib import Path
+from typing import Callable, Iterable, Iterator, TYPE_CHECKING
+
+import numpy as np
+import torch
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import LazyTorchTensor, ModelBase, TextModel, gguf, logger
+
+from .kimi_linear import KimiLinearModel
+
+
+@ModelBase.register("KimiK3ForConditionalGeneration")
+class KimiK3Model(TextModel):
+    """
+    Kimi-K3 text model (KimiLinearForCausalLM under a `language_model.` prefix).
+
+    Shares the hybrid MLA + KDA skeleton with kimi-linear, but that converter
+    cannot load it: K3 adds cross-layer attention residuals, a latent MoE, the
+    situ activation, an MLA output gate and a full-rank KDA gate.
+
+    The vision tower and mm_projector are skipped - text only for now.
+    """
+
+    model_arch = gguf.MODEL_ARCH.KIMI_K3
+
+    _experts: list[dict[str, Tensor]] | None = None
+
+    # `<x>_res_norm.weight` and `<x>_res_proj.weight` are only used as their
+    # elementwise product, so they are fused into one [n_embd] vector here.
+    # they arrive apart, so buffer the first one and tag it with its kind.
+    _res_parts: dict[str, tuple[str, Tensor]]
+
+    # HF suffix -> (gguf tensor, per-layer?)
+    _RES_FUSIONS = {
+        "self_attention_res": (gguf.MODEL_TENSOR.ATTN_RES_SCORE, True),
+        "mlp_res":            (gguf.MODEL_TENSOR.FFN_RES_SCORE,   True),
+        "output_attn_res":    (gguf.MODEL_TENSOR.OUTPUT_RES_SCORE, False),
+    }
+
+    # compressed-tensors MXFP4. the `language_model.` prefix is still there, as
+    # self.model_tensors is keyed by the raw checkpoint names
+    _MXFP4_FORMAT = "mxfp4-pack-quantized"
+    _MXFP4_EXPERT_RE = re.compile(
+        r"^(?:language_model\.)?model\.layers\.(\d+)"
+        r"\.block_sparse_moe\.experts\.(\d+)\.(w[123])\.weight_packed$"
+    )
+    _MXFP4_PROJ = {
+        "w1": gguf.MODEL_TENSOR.FFN_GATE_EXP,
+        "w2": gguf.MODEL_TENSOR.FFN_DOWN_EXP,
+        "w3": gguf.MODEL_TENSOR.FFN_UP_EXP,
+    }
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._res_parts = {}
+
+    def set_vocab(self):
+        # K3 has the same TikToken vocab as K2, so kimi-linear's vocab handling works.
+        # borrowed, not inherited: the method only touches TextModel members, and K3
+        # shares none of kimi-linear's tensor layout.
+        KimiLinearModel.set_vocab(self)  # ty: ignore[invalid-argument-type]
+
+        # ...but that forces eos to the tokenizer's eos_id, which is [EOS], the
+        # document terminator. K3's config says <|end_of_msg|>, the turn terminator;
+        # with [EOS] the generation never stops at the end of a turn.
+        if (eos := self.hparams.get("eos_token_id")) is not None:
+            logger.info(f"restoring configured eos_token_id {eos} (kimi-linear forces the tokenizer's)")
+            self.gguf_writer.add_eos_token_id(eos)
+
+        # K3 renders chats in python (encoding_k3.py) and ships no jinja template,
+        # so add the bundled one when the model has none
+        if gguf.SpecialVocab(self.dir_model, load_merges=False).chat_template is None:
+            template_path = Path(__file__).parent.parent / "models" / "templates" / "Kimi-K3.jinja"
+            logger.info(f"gguf: model has no chat template, using {template_path.name}")
+            self.gguf_writer.add_chat_template(template_path.read_text(encoding="utf-8"))
+
+    #
+    # compressed-tensors MXFP4 -> ggml MXFP4
+    #
+
+    def _is_mxfp4_packed(self) -> bool:
+        quant_config = self.hparams.get("quantization_config") or {}
+        return (quant_config.get("quant_method") == "compressed-tensors"
+                and quant_config.get("format") == self._MXFP4_FORMAT)
+
+    def dequant_model(self):
+        if not self._is_mxfp4_packed():
+            return super().dequant_model()
+
+        # skipping base.py's dequant is only safe if the experts are the only
+        # quantized tensors, so check it
+        stray = [n for n in self.model_tensors
+                 if n.endswith(".weight_packed") and not self._MXFP4_EXPERT_RE.match(n)]
+        if stray:
+            raise NotImplementedError(
+                f"{len(stray)} MXFP4 tensor(s) outside the routed experts, e.g. {stray[0]!r}; "
+                "only the routed experts have a repack path"
+            )
+
+    def _mxfp4_expert_tensor(self, loaders: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]):
+        """
+        One stacked [n_expert, rows, cols] MXFP4 tensor, built lazily.
+
+        gguf_writer holds every added tensor until the final write, so building
+        this eagerly (like the DeepSeek-V4 path does) keeps all ~1.38 TB of
+        experts in memory. lazy means only the tensor being written is resident.
+        """
+        # meta shapes, so this does not read any weights
+        rows, packed_cols = loaders[0][0]().shape
+        n_blocks = (packed_cols * 2) // 32
+        byte_shape = (len(loaders), rows, n_blocks * 17)
+
+        def load(fns: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]) -> np.ndarray:
+            out = np.empty(byte_shape, dtype=np.uint8)
+            for eid, (packed_fn, scale_fn) in enumerate(fns):
+                out[eid] = self.repack_mxfp4_blocks(
+                    LazyTorchTensor.to_eager(packed_fn()),
+                    LazyTorchTensor.to_eager(scale_fn()),
+                )
+            return out
+
+        # loaders goes through args, not the closure, so that `func` matches
+        # LazyBase's single-argument shape
+        return gguf.LazyNumpyTensor(
+            meta=gguf.LazyNumpyTensor.meta_with_dtype_and_shape(np.uint8, byte_shape),
+            args=(loaders,),
+            func=load,
+        )
+
+    def _write_mxfp4_experts(self) -> None:
+        n_experts = self.hparams["num_experts"]
+
+        # (bid, wid) -> {expert id: (packed name, scale name)}
+        groups: dict[tuple[int, str], dict[int, tuple[str, str]]] = {}
+        for name in self.model_tensors:
+            m = self._MXFP4_EXPERT_RE.match(name)
+            if m is None:
+                continue
+            bid, eid, wid = int(m.group(1)), int(m.group(2)), m.group(3)
+            scale_name = name.removesuffix("_packed") + "_scale"
+            if scale_name not in self.model_tensors:
+                raise KeyError(f"missing {scale_name} for {name}")
+            groups.setdefault((bid, wid), {})[eid] = (name, scale_name)
+
+        consumed: list[str] = []
+        for (bid, wid), experts in sorted(groups.items()):
+            missing = [e for e in range(n_experts) if e not in experts]
+            if missing:
+                raise KeyError(
+                    f"layer {bid} {wid}: {len(missing)} of {n_experts} experts missing, "
+                    f"first is {missing[0]}"
+                )
+            if len(experts) != n_experts:
+                raise KeyError(f"layer {bid} {wid}: {len(experts)} experts, expected {n_experts}")
+
+            loaders = []
+            for eid in range(n_experts):
+                packed_name, scale_name = experts[eid]
+                loaders.append((self.model_tensors[packed_name], self.model_tensors[scale_name]))
+                consumed += [packed_name, scale_name]
+
+            data = self._mxfp4_expert_tensor(loaders)
+            new_name = self.format_tensor_name(self._MXFP4_PROJ[wid], bid)
+            shape = gguf.quant_shape_from_byte_shape(data.shape, gguf.GGMLQuantizationType.MXFP4)
+            logger.info(
+                f"{new_name}: repacked {n_experts} experts to MXFP4, "
+                f"shape = {{{', '.join(str(n) for n in reversed(shape))}}}"
+            )
+            self.gguf_writer.add_tensor(new_name, data, raw_dtype=gguf.GGMLQuantizationType.MXFP4)
+
+        for name in consumed:
+            del self.model_tensors[name]
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        # not a generator on purpose: base.py chains this with get_tensors(), so the
+        # tensors used here must be removed from model_tensors before that starts
+        if self._is_mxfp4_packed():
+            self._write_mxfp4_experts()
+        return ()
+
+    def get_tensors(self) -> Iterator[tuple[str, Tensor]]:
+        for name, data in super().get_tensors():
+            if name.startswith(("vision_tower.", "mm_projector.")):
+                continue  # text only
+            if name.startswith("language_model."):
+                name = name[len("language_model."):]
+            yield name, data
+
+    def set_gguf_parameters(self):
+        # MLA is served as MQA with a single large head, then decompressed
+        self.hparams["num_key_value_heads"] = 1
+
+        super().set_gguf_parameters()
+        self.gguf_writer.add_vocab_size(self.hparams["vocab_size"])
+
+        linear_attn_config = self.hparams["linear_attn_config"]
+
+        # n_head_kv == 0 marks a KDA (recurrent) layer. the layer lists are 1-indexed,
+        # as KimiLinearConfig.is_kda_layer uses (layer_idx + 1)
+        full_attn_layers = linear_attn_config["full_attn_layers"]
+        n_kv_heads = [
+            self.hparams["num_key_value_heads"] if (il + 1) in full_attn_layers else 0
+            for il in range(self.hparams["num_hidden_layers"])
+        ]
+        assert len(n_kv_heads) == self.hparams["num_hidden_layers"]
+        self.gguf_writer.add_head_count_kv(n_kv_heads)
+
+        # --- KDA ---
+        self.gguf_writer.add_ssm_conv_kernel(linear_attn_config["short_conv_kernel_size"])
+        self.gguf_writer.add_kda_head_dim(linear_attn_config["head_dim"])
+        if (lb := linear_attn_config.get("gate_lower_bound")) is not None:
+            self.gguf_writer.add_kda_gate_lower_bound(lb)
+
+        # --- MLA ---
+        if (q_lora_rank := self.hparams.get("q_lora_rank")) is not None:
+            self.gguf_writer.add_q_lora_rank(q_lora_rank)
+        kv_lora_rank = self.hparams["kv_lora_rank"]
+        self.gguf_writer.add_kv_lora_rank(kv_lora_rank)
+
+        qk_nope_head_dim = self.hparams["qk_nope_head_dim"]
+        qk_rope_head_dim = self.hparams["qk_rope_head_dim"]
+        v_head_dim = self.hparams["v_head_dim"]
+        # K3 is nope-only; qk_rope_head_dim still sizes the un-absorbed part of K
+        assert self.hparams.get("mla_use_nope"), "K3 MLA is expected to be nope-only"
+        self.gguf_writer.add_rope_dimension_count(qk_rope_head_dim)
+        # MLA is served as MQA, so the cache holds the compressed latent
+        self.gguf_writer.add_key_length(kv_lora_rank + qk_rope_head_dim)
+        self.gguf_writer.add_value_length(kv_lora_rank)
+        self.gguf_writer.add_key_length_mla(qk_nope_head_dim + qk_rope_head_dim)
+        self.gguf_writer.add_value_length_mla(v_head_dim)
+
+        # --- MoE ---
+        self.gguf_writer.add_expert_feed_forward_length(self.hparams["moe_intermediate_size"])
+        self.gguf_writer.add_expert_shared_count(self.hparams["num_shared_experts"])
+        self.gguf_writer.add_leading_dense_block_count(self.hparams["first_k_dense_replace"])
+        self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"])
+        self.gguf_writer.add_expert_weights_norm(self.hparams["moe_renormalize"])
+        assert self.hparams["moe_router_activation_func"] == "sigmoid"
+        self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)
+        # latent MoE: routed experts live in a down-projected space
+        if (latent := self.hparams.get("routed_expert_hidden_size")) is not None:
+            self.gguf_writer.add_expert_latent_length(latent)
+
+        # --- situ activation ---
+        assert self.hparams["hidden_act"] == "situ", \
+            f"unexpected hidden_act {self.hparams['hidden_act']!r}"
+        self.gguf_writer.add_activation_situ_beta(self.hparams["activation_situ_beta"])
+        self.gguf_writer.add_activation_situ_linear_beta(self.hparams["activation_situ_linear_beta"])
+
+        # --- cross-layer attention residuals ---
+        self.gguf_writer.add_attn_res_block_size(self.hparams["attn_res_block_size"])
+
+    def prepare_tensors(self):
+        super().prepare_tensors()
+        if self._experts is not None:
+            leftover = [k for d in self._experts for k in d.keys()]
+            if leftover:
+                raise ValueError(f"Unprocessed experts: {leftover}")
+        if self._res_parts:
+            raise ValueError(f"Unpaired attention-residual tensors: {sorted(self._res_parts)}")
+        if self._is_mxfp4_packed():
+            # label the file for what it is; prepare_metadata runs after this
+            self._is_mxfp4 = True
+            self.ftype = gguf.LlamaFileType.MOSTLY_MXFP4_MOE
+
+    def _try_fuse_res(self, data_torch: Tensor, name: str, bid: int | None):
+        """
+        Pair <x>_res_norm.weight with <x>_res_proj.weight and emit their product.
+
+        Returns None if this is not a res tensor, [] if buffered until its pair.
+        """
+        for prefix, (tensor_id, per_layer) in self._RES_FUSIONS.items():
+            for kind in ("norm", "proj"):
+                if not name.endswith(f"{prefix}_{kind}.weight"):
+                    continue
+                key = f"{prefix}.{bid}"
+                other = self._res_parts.pop(key, None)
+                if other is None:
+                    self._res_parts[key] = (kind, data_torch)
+                    return []
+                other_kind, other_data = other
+                assert other_kind != kind, f"duplicate {kind} for {key}"
+                norm = data_torch if kind == "norm" else other_data
+                proj = data_torch if kind == "proj" else other_data
+                fused = norm.float().flatten() * proj.float().flatten()
+                # ".weight" suffix matches the convention map_tensor_name applies
+                new_name = (self.format_tensor_name(tensor_id, bid) if per_layer
+                            else gguf.TENSOR_NAMES[tensor_id] + ".weight")
+                logger.info(f"fused {prefix}_norm * {prefix}_proj -> {new_name}")
+                return [(new_name, fused)]
+        return None
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # --- cross-layer attention residuals: fuse norm * proj ---
+        fused = self._try_fuse_res(data_torch, name, bid)
+        if fused is not None:
+            yield from fused
+            return
+
+        # --- KDA conv1d: HF [d_inner, 1, d_conv] -> ggml ne [d_conv, 1, d_inner, 1] ---
+        # GGUF reverses the numpy shape on write, so target numpy (1, d_inner, 1, d_conv).
+        # conv_step varies fastest in both layouts, so this is a pure reshape.
+        if name.endswith((".q_conv1d.weight", ".k_conv1d.weight", ".v_conv1d.weight")):
+            if data_torch.ndim == 3:      # [d_inner, 1, d_conv]
+                d_inner, _, d_conv = data_torch.shape
+            elif data_torch.ndim == 2:    # [d_inner, d_conv]
+                d_inner, d_conv = data_torch.shape
+            else:
+                raise ValueError(f"unexpected conv1d rank {data_torch.ndim} for {name}")
+            data_torch = data_torch.reshape(1, d_inner, 1, d_conv)
+
+        # -exp(A_log) is folded here so the graph does not have to
+        if name.endswith(".A_log"):
+            n_head = self.hparams["num_attention_heads"]
+            data_torch = -torch.exp(data_torch.float()[:n_head])
+
+        # dt_bias -> the name SSM_DT's mapping expects
+        if name.endswith(".dt_bias"):
+            name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias"
+
+        # --- g_proj is two different tensors sharing one HF name ---
+        # KDA layers: full-rank gate, [d_inner, n_embd]  (replaces g_a/g_b)
+        # MLA layers: output gate,    [n_head*v_head_dim, n_embd]
+        # Name-based mapping cannot tell them apart, so resolve by layer type.
+        if name.endswith(".self_attn.g_proj.weight"):
+            assert bid is not None
+            is_kda = (bid + 1) not in self.hparams["linear_attn_config"]["full_attn_layers"]
+            tensor_id = gguf.MODEL_TENSOR.SSM_G if is_kda else gguf.MODEL_TENSOR.ATTN_GATE
+            yield self.format_tensor_name(tensor_id, bid), data_torch
+            return
+
+        # --- routed experts: stack per-expert 2D weights into one 3D tensor ---
+        if ".block_sparse_moe.experts." in name:
+            n_experts = self.hparams["num_experts"]
+            assert bid is not None
+
+            if self._experts is None:
+                self._experts = [{} for _ in range(self.block_count)]
+            self._experts[bid][name] = data_torch
+
+            if len(self._experts[bid]) < n_experts * 3:
+                return
+
+            # w1: gate, w2: down, w3: up
+            for wid, tensor_id in (("w1", gguf.MODEL_TENSOR.FFN_GATE_EXP),
+                                   ("w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP),
+                                   ("w3", gguf.MODEL_TENSOR.FFN_UP_EXP)):
+                datas = []
+                for xid in range(n_experts):
+                    ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight"
+                    datas.append(self._experts[bid].pop(ename))
+                stacked = torch.stack(datas, dim=0)
+                yield from super().modify_tensors(stacked, self.format_tensor_name(tensor_id, bid), bid)
+            return
+
+        # --- MLA absorption: split kv_b into k_b (transposed) and v_b ---
+        if name.endswith("kv_b_proj.weight"):
+            n_head_kv = self.hparams["num_key_value_heads"]
+            v_head_dim = self.hparams["v_head_dim"]
+            qk_nope_head_dim = self.hparams["qk_nope_head_dim"]
+            assert data_torch.shape[0] == n_head_kv * (v_head_dim + qk_nope_head_dim)
+            kv_b = data_torch.view(n_head_kv, v_head_dim + qk_nope_head_dim, data_torch.shape[-1])
+            k_b, v_b = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1)
+            k_b = k_b.transpose(1, 2)
+            yield from super().modify_tensors(k_b, name.replace("kv_b_proj", "k_b_proj"), bid)
+            yield from super().modify_tensors(v_b, name.replace("kv_b_proj", "v_b_proj"), bid)
+            return
+
+        yield from super().modify_tensors(data_torch, name, bid)

conversion/minimax.py

@@ -1,13 +1,121 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import Iterable, Sequence, TYPE_CHECKING
 
 import torch
 
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, TextModel, MmprojModel, gguf
+from .base import ModelBase, TextModel, MmprojModel, gguf, logger
+
+
+@ModelBase.register("MiniMaxText01ForCausalLM")
+@ModelBase.register("MiniMaxM1ForCausalLM")
+class MiniMaxText01Model(TextModel):
+    model_arch = gguf.MODEL_ARCH.MINIMAX01
+
+    def _get_suppress_tokens(self) -> Sequence[int] | None:
+        import json
+        from transformers import AutoTokenizer
+        from .base import LazyTorchTensor
+
+        # check added tokens embeddings in embeddings tensor for zero-valued embeddings
+        # they get in the way of the token sampling process and must be suppressed
+
+        tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True)
+        tokenizer_vocab_size = tokenizer.vocab_size
+
+        with open(self.dir_model / "model.safetensors.index.json", "r", encoding="utf-8") as f:
+            weight_map = json.load(f)["weight_map"]
+
+        embeddings_tensor_name = "model.embed_tokens.weight"
+        embeddings_shard_name = weight_map[embeddings_tensor_name]
+        with gguf.utility.SafetensorsLocal(self.dir_model / embeddings_shard_name) as model_shard:
+            embeddings_data = model_shard[embeddings_tensor_name]
+
+        embeddings_weights_dtype = LazyTorchTensor._dtype_str_map[embeddings_data.dtype]
+        embeddings_weights = torch.from_numpy(embeddings_data.mmap_bytes()).view(embeddings_weights_dtype).reshape(embeddings_data.shape)
+        embeddings_vocab_size = embeddings_weights.shape[0]
+
+        embeddings_added_tokens = embeddings_weights[tokenizer_vocab_size:embeddings_vocab_size]
+        embeddings_zero_rows = torch.all(embeddings_added_tokens == 0, dim=1)
+        tokens_zero_embeddings_ids = (torch.nonzero(embeddings_zero_rows, as_tuple=False).flatten() + tokenizer_vocab_size).tolist()
+
+        return tokens_zero_embeddings_ids
+
+    def set_vocab(self) -> None:
+        from pathlib import Path
+
+        self._set_vocab_gpt2()
+
+        for tmpl_file in [
+            self.dir_model / "chat_template.jinja",
+            Path(__file__).parent.parent / "models" / "templates" / "MiniMax-M1.jinja"
+        ]:
+            if tmpl_file.is_file():
+                self.gguf_writer.add_chat_template(tmpl_file.read_text(encoding="utf-8"))
+                logger.info(f"Chat template overridden with {tmpl_file}.")
+                break
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+
+        suppress_tokens = self._get_suppress_tokens()
+        if suppress_tokens:
+            logger.info(f"Suppressing tokens with zero embeddings {suppress_tokens}")
+            self.gguf_writer.add_suppress_tokens(suppress_tokens)
+
+        layernorm_full_attention_alpha = self.hparams["layernorm_full_attention_alpha"]
+        layernorm_full_attention_beta = self.hparams["layernorm_full_attention_beta"]
+        layernorm_linear_attention_alpha = self.hparams["layernorm_linear_attention_alpha"]
+        layernorm_linear_attention_beta = self.hparams["layernorm_linear_attention_beta"]
+        layernorm_mlp_alpha = self.hparams["layernorm_mlp_alpha"]
+        layernorm_mlp_beta = self.hparams["layernorm_mlp_beta"]
+        assert layernorm_full_attention_alpha == layernorm_linear_attention_alpha == layernorm_mlp_alpha
+        assert layernorm_full_attention_beta == layernorm_linear_attention_beta == layernorm_mlp_beta == 1.0
+        # we do not store the layernorm betas as they are all 1.0
+        # layernorm alphas are stored as single residual_scale hparam
+        self.gguf_writer.add_residual_scale(layernorm_full_attention_alpha)
+
+        self.gguf_writer.add_rope_dimension_count(self.hparams["rotary_dim"])
+
+    _experts: list[dict[str, Tensor]] | None = None
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # process the experts separately
+        if name.find("block_sparse_moe.experts") != -1:
+            n_experts = self.hparams["num_local_experts"]
+
+            assert bid is not None
+
+            if self._experts is None:
+                self._experts = [{} for _ in range(self.block_count)]
+
+            self._experts[bid][name] = data_torch
+
+            if len(self._experts[bid]) >= n_experts * 3:
+                # merge the experts into a single 3d tensor
+                for wid in ["w1", "w2", "w3"]:
+                    datas: list[Tensor] = []
+
+                    for xid in range(n_experts):
+                        ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight"
+                        datas.append(self._experts[bid][ename])
+                        del self._experts[bid][ename]
+
+                    data_torch = torch.stack(datas, dim=0)
+
+                    merged_name = f"layers.{bid}.feed_forward.experts.{wid}.weight"
+
+                    new_name = self.map_tensor_name(merged_name)
+
+                    yield from super().modify_tensors(data_torch, new_name, bid)
+                return
+            else:
+                return
+
+        yield from super().modify_tensors(data_torch, name, bid)
 
 
 @ModelBase.register("MiniMaxM2ForCausalLM")

docs/backend/SYCL.md

@@ -428,13 +428,13 @@ Examples:
 - Use device 0:
 
 ```sh
-ZES_ENABLE_SYSMAN=1 ./build/bin/llama-completion -no-cnv -m models/llama-2-7b.Q4_0.gguf -p "Building a website can be done in 10 simple steps:" -n 400 -e -ngl 99 -sm none -mg 0 --mmap
+ZES_ENABLE_SYSMAN=1 ./build/bin/llama-completion -no-cnv -m models/llama-2-7b.Q4_0.gguf -p "Building a website can be done in 10 simple steps:" -n 400 -e -ngl 99 -sm none -mg 0 --load-mode auto
 ```
 
 - Use multiple devices:
 
 ```sh
-ZES_ENABLE_SYSMAN=1 ./build/bin/llama-completion -no-cnv -m models/llama-2-7b.Q4_0.gguf -p "Building a website can be done in 10 simple steps:" -n 400 -e -ngl 99 -sm layer --mmap
+ZES_ENABLE_SYSMAN=1 ./build/bin/llama-completion -no-cnv -m models/llama-2-7b.Q4_0.gguf -p "Building a website can be done in 10 simple steps:" -n 400 -e -ngl 99 -sm layer --load-mode auto
 ```
 
 *Notes:*
@@ -741,13 +741,13 @@ Examples:
 - Use device 0:
 
 ```
-build\bin\llama-completion.exe -no-cnv -m models\llama-2-7b.Q4_0.gguf -p "Building a website can be done in 10 simple steps:\nStep 1:" -n 400 -e -ngl 99 -sm none -mg 0 --mmap
+build\bin\llama-completion.exe -no-cnv -m models\llama-2-7b.Q4_0.gguf -p "Building a website can be done in 10 simple steps:\nStep 1:" -n 400 -e -ngl 99 -sm none -mg 0 --load-mode auto
 ```
 
 - Use multiple devices:
 
 ```
-build\bin\llama-completion.exe -no-cnv -m models\llama-2-7b.Q4_0.gguf -p "Building a website can be done in 10 simple steps:\nStep 1:" -n 400 -e -ngl 99 -sm layer --mmap
+build\bin\llama-completion.exe -no-cnv -m models\llama-2-7b.Q4_0.gguf -p "Building a website can be done in 10 simple steps:\nStep 1:" -n 400 -e -ngl 99 -sm layer --load-mode auto
 ```
 
 
@@ -804,7 +804,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_ENABLE_MKL_FA | 1 (default) or 0 | Enable oneMKL GEMM flash attention for XMX-accelerated prompt processing with quantized KV cache. Automatically activates during prefill (prompt processing) when all conditions are met: (1) flash-attn enabled (`-fa` or `--flash-attn on`), (2) KV cache quantized (`--cache-type-k q8_0 --cache-type-v q8_0` or other `*_0/*_1` types), (3) batch size ≥ 1024 (`--batch-size 1024`), (4) prompt length ≥ 1024 tokens. Set to 0 to force the TILE kernel for A/B testing. Example minimum command: `llama-cli -m model.gguf -fa -ngl 99 --cache-type-k q8_0 --cache-type-v q8_0 --batch-size 1024 -p "your prompt"` |
 | GGML_SYCL_MKL_FA_DEBUG | 0 (default) or 1 | Enable per-call diagnostic logging for MKL flash attention: GEMM/softmax timings, interleaved-head detection, and buffer memory usage. |
 | GGML_SYCL_MKL_FA_DIAG | 0 (default) or 1 | Enable output fingerprinting for MKL flash attention. Dumps the first 64 float output values for the first 6 FA calls with n_kv ≥ 1024, labeled with kernel type (MKL/TILE/VEC) for cross-kernel comparison. |
-| GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute (currently top-k MoE gating). |
+| GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute. |
 | GGML_SYCL_ENABLE_ESIMD | 0 or 1 (default)| Enable ESIMD kernels when available. |
 | ZES_ENABLE_SYSMAN | 0 (default) or 1 | Support to get free memory of GPU by sycl::aspect::ext_intel_free_memory.<br>Recommended to use when --split-mode = layer |
 | UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS | 0 (default) or 1 | Allow SYCL/Unified Runtime Level Zero device allocations larger than 4 GiB. llama.cpp's direct Level Zero allocation path requests the relaxed maximum-size limit itself when GGML_SYCL_ENABLE_LEVEL_ZERO=1. |

docs/backend/snapdragon/developer.md

@@ -53,7 +53,7 @@ M=gpt-oss-20b-Q4_0.gguf NDEV=4 D=HTP0,HTP1,HTP2,HTP3 P=surfing.txt scripts/snapd
 ...
 LD_LIBRARY_PATH=/data/local/tmp/llama.cpp/lib
 ADSP_LIBRARY_PATH=/data/local/tmp/llama.cpp/lib
-GGML_HEXAGON_NDEV=4 ./bin/llama-cli --no-mmap -m /data/local/tmp/llama.cpp/../gguf/gpt-oss-20b-Q4_0.gguf
+GGML_HEXAGON_NDEV=4 ./bin/llama-cli --load-mode none -m /data/local/tmp/llama.cpp/../gguf/gpt-oss-20b-Q4_0.gguf
       -t 4 --ctx-size 8192 --batch-size 128 -ctk q8_0 -ctv q8_0 -fa on -ngl 99 --device HTP0,HTP1,HTP2,HTP3 -no-cnv -f surfing.txt
 ...
 llama_model_loader: - type  f32:  289 tensors

examples/convert-llama2c-to-ggml/convert-llama2c-to-ggml.cpp

@@ -549,20 +549,34 @@ static void load_vocab(const char * filename, const Config * config, struct my_l
 
         const int token_idx = gguf_find_key(ctx, KV_TOKENIZER_LIST);
         GGML_ASSERT(token_idx >= 0);
-
-        const int score_idx = gguf_find_key(ctx, KV_TOKENIZER_SCORES);
-        GGML_ASSERT(score_idx >= 0);
-        const float * scores = (const float * ) gguf_get_arr_data(ctx, score_idx);
-
-        const int toktype_idx = gguf_find_key(ctx, KV_TOKENIZER_TOKEN_TYPE);
-        GGML_ASSERT(toktype_idx >= 0);
-        const int * toktypes = (const int * ) gguf_get_arr_data(ctx, toktype_idx);
+        if (gguf_get_kv_type(ctx, token_idx) != GGUF_TYPE_ARRAY ||
+            gguf_get_arr_type(ctx, token_idx) != GGUF_TYPE_STRING) {
+            die_fmt("invalid gguf type for %s", KV_TOKENIZER_LIST);
+        }
 
         const uint32_t n_vocab = gguf_get_arr_n(ctx, token_idx);
         if (n_vocab != static_cast<uint32_t>(config->vocab_size)) {
             die_fmt("vocab size mismatch: (gguf) %u != (llama2c) %d", n_vocab, config->vocab_size);
         }
 
+        const int score_idx = gguf_find_key(ctx, KV_TOKENIZER_SCORES);
+        GGML_ASSERT(score_idx >= 0);
+        if (gguf_get_kv_type(ctx, score_idx) != GGUF_TYPE_ARRAY ||
+            gguf_get_arr_type(ctx, score_idx) != GGUF_TYPE_FLOAT32 ||
+            gguf_get_arr_n(ctx, score_idx) < n_vocab) {
+            die_fmt("invalid gguf type or size for %s", KV_TOKENIZER_SCORES);
+        }
+        const float * scores = (const float * ) gguf_get_arr_data(ctx, score_idx);
+
+        const int toktype_idx = gguf_find_key(ctx, KV_TOKENIZER_TOKEN_TYPE);
+        GGML_ASSERT(toktype_idx >= 0);
+        if (gguf_get_kv_type(ctx, toktype_idx) != GGUF_TYPE_ARRAY ||
+            gguf_get_arr_type(ctx, toktype_idx) != GGUF_TYPE_INT32 ||
+            gguf_get_arr_n(ctx, toktype_idx) < n_vocab) {
+            die_fmt("invalid gguf type or size for %s", KV_TOKENIZER_TOKEN_TYPE);
+        }
+        const int * toktypes = (const int * ) gguf_get_arr_data(ctx, toktype_idx);
+
         vocab->id_to_token.resize(n_vocab);
 
         for (uint32_t i = 0; i < n_vocab; i++) {

examples/sycl/run-llama2.sh

@@ -18,7 +18,7 @@ CONTEXT=4096
 #support malloc device memory more than 4GB.
 export UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS=1
 
-LOAD_MODE='--mmap'
+LOAD_MODE='--load-mode auto'
 if [ $# -gt 0 ]; then
     GGML_SYCL_DEVICE=$1
     echo "use $GGML_SYCL_DEVICE as main GPU"

examples/sycl/start-svr.sh

@@ -124,7 +124,7 @@ else
     GPUS_SETTING="-sm ${SPLIT_MODE}"
 fi
 
-echo "run cmd: ZES_ENABLE_SYSMAN=1 ${BIN_FILE} -m ${MODEL_FILE} -ngl ${NGL} -s ${SEED} -c ${CONTEXT} ${GPUS_SETTING} -lv ${LOG_VERBOSE} --device ${SYCL_DEVICES} --mmap --host 0.0.0.0 --port 8000"
-ZES_ENABLE_SYSMAN=1 ${BIN_FILE} -m ${MODEL_FILE} -ngl ${NGL} -s ${SEED} -c ${CONTEXT} ${GPUS_SETTING} -lv ${LOG_VERBOSE} --device ${SYCL_DEVICES} --mmap --host 0.0.0.0 --port 8000
+echo "run cmd: ZES_ENABLE_SYSMAN=1 ${BIN_FILE} -m ${MODEL_FILE} -ngl ${NGL} -s ${SEED} -c ${CONTEXT} ${GPUS_SETTING} -lv ${LOG_VERBOSE} --device ${SYCL_DEVICES} --load-mode auto --host 0.0.0.0 --port 8000"
+ZES_ENABLE_SYSMAN=1 ${BIN_FILE} -m ${MODEL_FILE} -ngl ${NGL} -s ${SEED} -c ${CONTEXT} ${GPUS_SETTING} -lv ${LOG_VERBOSE} --device ${SYCL_DEVICES} --load-mode auto --host 0.0.0.0 --port 8000
 
 

examples/sycl/test.sh

@@ -133,6 +133,6 @@ else
     GPUS_SETTING="-sm ${SPLIT_MODE}"
  fi
 
-echo "run cmd: ZES_ENABLE_SYSMAN=1 ${BIN_FILE} -m ${MODEL_FILE} -no-cnv -p "${INPUT_PROMPT}" -n 200 -e -ngl ${NGL} -s ${SEED} -c ${CONTEXT} ${GPUS_SETTING} -lv ${LOG_VERBOSE} --device ${SYCL_DEVICES} --mmap "
-ZES_ENABLE_SYSMAN=1 ${BIN_FILE} -m ${MODEL_FILE} -no-cnv -p "${INPUT_PROMPT}" -n 200 -e -ngl ${NGL} -s ${SEED} -c ${CONTEXT} ${GPUS_SETTING} -lv ${LOG_VERBOSE} --device ${SYCL_DEVICES} --mmap
+echo "run cmd: ZES_ENABLE_SYSMAN=1 ${BIN_FILE} -m ${MODEL_FILE} -no-cnv -p "${INPUT_PROMPT}" -n 200 -e -ngl ${NGL} -s ${SEED} -c ${CONTEXT} ${GPUS_SETTING} -lv ${LOG_VERBOSE} --device ${SYCL_DEVICES} --load-mode auto "
+ZES_ENABLE_SYSMAN=1 ${BIN_FILE} -m ${MODEL_FILE} -no-cnv -p "${INPUT_PROMPT}" -n 200 -e -ngl ${NGL} -s ${SEED} -c ${CONTEXT} ${GPUS_SETTING} -lv ${LOG_VERBOSE} --device ${SYCL_DEVICES} --load-mode auto
 

examples/sycl/win-run-llama2.bat

@@ -7,5 +7,5 @@ set INPUT2="Building a website can be done in 10 simple steps:\nStep 1:"
 
 :: support malloc device memory more than 4GB.
 set UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS=1
-set LOAD_MODE="--mmap"
+set LOAD_MODE="--load-mode auto"
 .\build\bin\llama-completion.exe -m models\llama-2-7b.Q4_0.gguf -no-cnv -p %INPUT2% -n 400 -e -ngl 99 -s 0 %LOAD_MODE%

examples/sycl/win-start-svr.bat

@@ -188,9 +188,9 @@ if not "%GGML_SYCL_DEVICE%"=="-1" (
   set "GPUS_SETTING=-sm %SPLIT_MODE%"
 )
 
-echo run cmd: ZES_ENABLE_SYSMAN=1 %BIN_FILE% -m "%MODEL_FILE%" -ngl %NGL% -s %SEED% -c %CONTEXT% %GPUS_SETTING% -lv %LOG_VERBOSE% --device %SYCL_DEVICES% --mmap --host 0.0.0.0 --port 8000
+echo run cmd: ZES_ENABLE_SYSMAN=1 %BIN_FILE% -m "%MODEL_FILE%" -ngl %NGL% -s %SEED% -c %CONTEXT% %GPUS_SETTING% -lv %LOG_VERBOSE% --device %SYCL_DEVICES% --load-mode auto --host 0.0.0.0 --port 8000
 set "ZES_ENABLE_SYSMAN=1"
-%BIN_FILE% -m "%MODEL_FILE%" -ngl %NGL% -s %SEED% -c %CONTEXT% %GPUS_SETTING% -lv %LOG_VERBOSE% --device "%SYCL_DEVICES%" --mmap --host 0.0.0.0 --port 8000
+%BIN_FILE% -m "%MODEL_FILE%" -ngl %NGL% -s %SEED% -c %CONTEXT% %GPUS_SETTING% -lv %LOG_VERBOSE% --device "%SYCL_DEVICES%" --load-mode auto --host 0.0.0.0 --port 8000
 
 endlocal
 

examples/sycl/win-test.bat

@@ -211,9 +211,9 @@ else (
   set "GPUS_SETTING=-sm %SPLIT_MODE%"
 )
 
-echo run cmd: ZES_ENABLE_SYSMAN=1 %BIN_FILE% -m %MODEL_FILE% -no-cnv -p "%INPUT_PROMPT%" -n 200 -e -ngl %NGL% -s %SEED% -c %CONTEXT% %GPUS_SETTING% -lv %LOG_VERBOSE% --device %SYCL_DEVICES% --mmap
+echo run cmd: ZES_ENABLE_SYSMAN=1 %BIN_FILE% -m %MODEL_FILE% -no-cnv -p "%INPUT_PROMPT%" -n 200 -e -ngl %NGL% -s %SEED% -c %CONTEXT% %GPUS_SETTING% -lv %LOG_VERBOSE% --device %SYCL_DEVICES% --load-mode auto
 set "ZES_ENABLE_SYSMAN=1"
-%BIN_FILE% -m "%MODEL_FILE%" -no-cnv -p "%INPUT_PROMPT%" -n 200 -e -ngl %NGL% -s %SEED% -c %CONTEXT% %GPUS_SETTING% -lv %LOG_VERBOSE% --device "%SYCL_DEVICES%" --mmap
+%BIN_FILE% -m "%MODEL_FILE%" -no-cnv -p "%INPUT_PROMPT%" -n 200 -e -ngl %NGL% -s %SEED% -c %CONTEXT% %GPUS_SETTING% -lv %LOG_VERBOSE% --device "%SYCL_DEVICES%" --load-mode auto
 
 endlocal
 

examples/test-cmake/README.md

@@ -27,10 +27,10 @@ Build/run this project using the installation created above:
 (venv) $ ./build.sh
 -- Configuring done (0.0s)
 -- Generating done (0.0s)
--- Build files have been written to: /home/danbev/work/ai/llama.cpp/examples/test-cmake/build
+-- Build files have been written to: /path/to/llama.cpp/examples/test-cmake/build
 [100%] Built target test-cmake
 [test-cmake] Using llama.cpp version 0.1.0-dev-b10335
 [test-cmake] Initializing backend...
-load_backend: loaded CPU backend from /home/danbev/work/ai/llama.cpp/examples/test-cmake/install/lib/llama.cpp/libggml-cpu-alderlake.so
+load_backend: loaded CPU backend from /path/to/llama.cpp/examples/test-cmake/install/lib/llama.cpp/libggml-cpu-alderlake.so
 [test-cmake] Backend initialized.
 ```

ggml/CMakeLists.txt

@@ -4,7 +4,7 @@ project("ggml" C CXX ASM)
 
 ### GGML Version
 set(GGML_VERSION_MAJOR 0)
-set(GGML_VERSION_MINOR 19)
+set(GGML_VERSION_MINOR 20)
 set(GGML_VERSION_PATCH 0)
 set(GGML_VERSION_BASE "${GGML_VERSION_MAJOR}.${GGML_VERSION_MINOR}.${GGML_VERSION_PATCH}")
 

ggml/include/ggml.h

@@ -2459,7 +2459,8 @@ extern "C" {
             struct ggml_tensor  * A,
             struct ggml_tensor  * B,
             struct ggml_tensor  * C,
-            struct ggml_tensor  * ids);
+            struct ggml_tensor  * ids,
+            int64_t               K);
 
     // partition into non-overlapping windows with padding if needed
     // example:

ggml/src/ggml-cpu/ggml-cpu.c

@@ -2795,6 +2795,11 @@ struct ggml_cplan ggml_graph_plan(
     n_threads = 1;
 #endif
 
+#if defined(__wasi__)
+    // WASI doesn't support parallelism yet
+    n_threads = 1;
+#endif
+
     size_t work_size = 0;
 
     struct ggml_cplan cplan;

ggml/src/ggml-cpu/ggml-cpu.cpp

@@ -472,6 +472,8 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st
                 src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;
         case GGML_OP_CONV_2D:
             return ggml_is_contiguous(op->src[0]);
+        case GGML_OP_SSM_SCAN:
+            return ggml_get_op_params_i32(op, 0) == 1 || op->src[3]->ne[0] == 1;
         default:
             return true;
     }

ggml/src/ggml-cpu/ops.cpp

@@ -9644,11 +9644,13 @@ static void ggml_compute_forward_ssm_scan_f32(
     const int64_t ng = src4->ne[1];
     const int64_t nt = src1->ne[2]; // number of tokens per sequence
     const int64_t ns = src1->ne[3]; // number of sequences in the batch
+    const int64_t K  = ggml_get_op_params_i32(dst, 0);
 
     // can't use ggml_nbytes because src1 is not necessarily contiguous
     const int64_t s_off = ggml_nelements(src1) * ggml_element_size(src1);
 
-    GGML_ASSERT(ggml_nelements(src1) + nc*nr*nh*ns == ggml_nelements(dst));
+    GGML_ASSERT(K >= 1);
+    GGML_ASSERT(ggml_nelements(src1) + K*nc*nr*nh*ns == ggml_nelements(dst));
     GGML_ASSERT(src0->nb[0] == sizeof(float));
     GGML_ASSERT(src1->nb[0] == sizeof(float));
     GGML_ASSERT(src2->nb[0] == sizeof(float));
@@ -9657,6 +9659,7 @@ static void ggml_compute_forward_ssm_scan_f32(
     GGML_ASSERT(src5->nb[0] == sizeof(float));
     GGML_ASSERT(src6->nb[0] == sizeof(int32_t));
     GGML_ASSERT(nh % ng == 0);
+    GGML_ASSERT(src3->ne[0] == 1 || K == 1);
 
     // heads per thread
     const int dh = (nh + nth - 1)/nth;
@@ -9831,6 +9834,13 @@ static void ggml_compute_forward_ssm_scan_f32(
                     }
                 }
             }
+            const int64_t slot = nt - 1 - i2;
+            if (K > 1 && slot > 0 && slot < K) {
+                float * s_snapshot = (float *) ((char *) dst->data + s_off + (slot*ns + i3)*(src0->nb[3]));
+                for (int h = ih0; h < ih1; ++h) {
+                    memcpy((char *) s_snapshot + h*src0->nb[2], (char *) s + h*src0->nb[2], src0->nb[2]);
+                }
+            }
             // use the output as the source when it's not the first token-wise iteration
             s0 = s;
         }

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -5189,11 +5189,17 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                    (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) &&
                    (op->type         == GGML_TYPE_F32 || op->type         == GGML_TYPE_F16);
         case GGML_OP_SSM_SCAN: {
+            const int32_t K = ggml_get_op_params_i32(op, 0);
+
             if (op->src[3]->ne[0] == 1) {
                 // Mamba2
                 // (kernel only supports (d_state == 128 || d_state == 256) && d_head % 16 == 0)
                 return (op->src[0]->ne[0] == 128 || op->src[0]->ne[0] == 256) && op->src[0]->ne[1] % 16 == 0;
             } else {
+                if (K > 1) {
+                    return false;
+                }
+
                 // Mamba
                 // (kernel only supports d_state == 16, d_head == 1, n_head % 128 == 0, n_group == 1)
                 return op->src[0]->ne[0] == 16 && op->src[0]->ne[1] == 1 && op->src[0]->ne[2] % 128 == 0 && op->src[4]->ne[1] == 1;

ggml/src/ggml-cuda/ssm-scan.cu

@@ -149,7 +149,7 @@ __global__ void __launch_bounds__(d_state, 1)
         const int src0_nb2, const int src0_nb3, const int src1_nb2, const int src1_nb3,
         const int src2_nb1, const int src2_nb2, const int src3_nb1,
         const int src4_nb2, const int src4_nb3, const int src5_nb2, const int src5_nb3,
-        const int64_t s_off, const int64_t n_head, const int64_t d_head, const int64_t n_group, const int64_t n_tok) {
+        const int64_t s_off, const int64_t n_head, const int64_t d_head, const int64_t n_group, const int64_t n_tok, const int64_t K) {
     const float   * GGML_CUDA_RESTRICT src0 = src0_ptr;
     const float   * GGML_CUDA_RESTRICT src1 = src1_ptr;
     const float   * GGML_CUDA_RESTRICT src2 = src2_ptr;
@@ -217,6 +217,16 @@ __global__ void __launch_bounds__(d_state, 1)
         if (lane == 0) {
             y_warp[i * stride_y] = state_sum;
         }
+
+        // Slot 0 is the final state written below; slots 1..K-1 are rollback snapshots.
+        const int64_t slot = n_tok - 1 - i;
+        if (K > 1 && slot > 0 && slot < K) {
+            float * s_snapshot_warp = (float *) ((char *) dst + s_off + (slot * gridDim.y + seq_idx) * src0_nb3 + head_idx * src0_nb2 + head_off * d_state);
+#pragma unroll
+            for (int j = 0; j < c_factor; j++) {
+                s_snapshot_warp[WARP_SIZE * j + lane] = state[j];
+            }
+        }
     }
 
     // write back the state
@@ -232,7 +242,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
                               const int src2_nb2, const int src3_nb1, const int src4_nb2, const int src4_nb3, const int src5_nb2,
                               const int src5_nb3, const int64_t s_off, const int64_t d_state, const int64_t head_dim,
                               const int64_t n_head, const int64_t n_group, const int64_t n_tok, const int64_t n_seq,
-                              cudaStream_t stream) {
+                              const int64_t K, cudaStream_t stream) {
     // NOTE: if you change conditions here, be sure to update the corresponding supports_op condition!
     if (src3_nb1 == sizeof(float)) {
         // Mamba-2
@@ -245,7 +255,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
             ggml_cuda_kernel_launch(ssm_scan_f32_group<128/WARP_SIZE, 128>, launch_params,
                     src0, src1, src2, src3, src4, src5, src6, dst,
                     src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
-                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok);
+                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K);
         } else if (d_state == 256) { // Falcon-H1
             constexpr int threads   = 256;
             constexpr int num_warps = threads/WARP_SIZE;
@@ -255,12 +265,13 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
             ggml_cuda_kernel_launch(ssm_scan_f32_group<256/WARP_SIZE, 256>, launch_params,
                     src0, src1, src2, src3, src4, src5, src6, dst,
                     src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
-                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok);
+                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K);
         } else {
             GGML_ABORT("doesn't support d_state!=(128 or 256).");
         }
     } else {
         // Mamba-1
+        GGML_ASSERT(K == 1);
         constexpr int threads = 128;
         GGML_ASSERT(n_head % threads == 0);
         GGML_ASSERT(head_dim == 1);
@@ -769,10 +780,12 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const int64_t ng  = src4->ne[1];  // n_group
     const int64_t n_t = src1->ne[2];  // number of tokens per sequence
     const int64_t n_s = src1->ne[3];  // number of sequences in the batch
+    const int32_t K_param = ggml_get_op_params_i32(dst, 0);
+    const int64_t K = K_param > 0 ? K_param : 1;
 
     const int64_t s_off = ggml_nelements(src1) * sizeof(float);
 
-    GGML_ASSERT(ggml_nelements(src1) + nc*nr*nh*n_s == ggml_nelements(dst));
+    GGML_ASSERT(ggml_nelements(src1) + K*nc*nr*nh*n_s == ggml_nelements(dst));
     GGML_ASSERT(src0->nb[0] == sizeof(float));
     GGML_ASSERT(src1->nb[0] == sizeof(float));
     GGML_ASSERT(src2->nb[0] == sizeof(float));
@@ -780,6 +793,7 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     GGML_ASSERT(src4->nb[0] == sizeof(float));
     GGML_ASSERT(src5->nb[0] == sizeof(float));
     GGML_ASSERT(src6->nb[0] == sizeof(int32_t));
+    GGML_ASSERT(src3->ne[0] == 1 || K == 1);
 
     const float * src0_d = (const float *) src0->data;
     const float * src1_d = (const float *) src1->data;
@@ -814,6 +828,7 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const bool is_mamba2 = (src3->nb[1] == sizeof(float));
     const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
     const bool use_ssd = is_mamba2 && n_t > SSM_SSD_MIN_TOKENS
+                      && K == 1
                       && n_t <= SSM_SSD_MAX_TOKENS
                       && GGML_CUDA_CC_IS_NVIDIA(cc)
                       && cc >= GGML_CUDA_CC_TURING
@@ -841,5 +856,5 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     ssm_scan_f32_cuda(src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d,
                       src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2],
                       src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3],
-                      s_off, nc, nr, nh, ng, n_t, n_s, stream);
+                      s_off, nc, nr, nh, ng, n_t, n_s, K, stream);
 }

ggml/src/ggml-et/et-kernels/src/ssm_scan_f32.c

@@ -12,7 +12,8 @@ struct ggml_et_ssm_scan_params {
     struct ggml_tensor src4;  // B:   [d_state, n_group, n_seq_tokens, n_seqs]
     struct ggml_tensor src5;  // C:   [d_state, n_group, n_seq_tokens, n_seqs]
     struct ggml_tensor src6;  // ids: [n_seqs] i32
-    struct ggml_tensor dst;   // packed [y, final_state]
+    struct ggml_tensor dst;   // packed [y, states]
+    int32_t            K;
 };
 
 static inline float softplus_f32(float x) {
@@ -72,6 +73,7 @@ int entry_point(struct ggml_et_ssm_scan_params * params, void * env) {
     const int64_t n_seq_tokens = src1->ne[2];
     const int64_t n_seqs       = src1->ne[3];
     const int64_t y_elems      = src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3];
+    const int64_t K            = params->K;
 
     if (src0->nb[0] != sizeof(float) || src1->nb[0] != sizeof(float) || src2->nb[0] != sizeof(float) ||
         src3->nb[0] != sizeof(float) || src4->nb[0] != sizeof(float) || src5->nb[0] != sizeof(float) ||
@@ -79,7 +81,7 @@ int entry_point(struct ggml_et_ssm_scan_params * params, void * env) {
         return -1;
     }
 
-    if (n_group <= 0 || n_head % n_group != 0) {
+    if (K < 1 || n_group <= 0 || n_head % n_group != 0) {
         return -1;
     }
 
@@ -260,6 +262,15 @@ int entry_point(struct ggml_et_ssm_scan_params * params, void * env) {
                         sumf += st * C_row[state_idx];
                     }
 
+                    const int64_t slot = n_seq_tokens - 1 - token_idx;
+                    if (slot > 0 && slot < K) {
+                        float * state_snapshot =
+                            (float *) ((char *) state_dst + (size_t) slot * n_seqs * src0->nb[3]);
+                        for (int64_t i = 0; i < d_state; ++i) {
+                            state_snapshot[i] = state_dst[i];
+                        }
+                    }
+
                     dst_data[seq_idx * (n_seq_tokens * n_head * head_dim) + token_idx * (n_head * head_dim) +
                              head_idx * head_dim + dim_idx] = sumf;
                 }

ggml/src/ggml-et/ggml-et-ops.cpp

@@ -2064,6 +2064,7 @@ bool ggml_et_op_ssm_scan(ggml_backend_et_device_context * dev_ctx, const ggml_te
     params.src5 = *node->src[5];
     params.src6 = *node->src[6];
     params.dst  = *node;
+    params.K    = ggml_get_op_params_i32(node, 0);
 
     bool kernel_result = ggml_et_launch_kernel(dev_ctx, "ssm_scan_f32", &params, sizeof(params), 0xFFFFFFFF);
 

ggml/src/ggml-et/ggml-et-ops.h

@@ -218,7 +218,8 @@ struct ggml_et_ssm_scan_params {
     ggml_tensor src4;  // B:   [d_state, n_group, n_seq_tokens, n_seqs]
     ggml_tensor src5;  // C:   [d_state, n_group, n_seq_tokens, n_seqs]
     ggml_tensor src6;  // ids: [n_seqs] i32
-    ggml_tensor dst;   // [y, final_state] packed output from ggml_ssm_scan()
+    ggml_tensor dst;   // [y, states] packed output from ggml_ssm_scan()
+    int32_t     K;
 };
 
 struct ggml_et_rwkv_wkv6_params {

ggml/src/ggml-metal/ggml-metal-device.m

@@ -1376,9 +1376,10 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
                 ggml_is_contiguous_rows(op->src[1]) &&
                 ggml_is_contiguous_rows(op->src[2]) &&
                 ggml_is_contiguous_rows(op->src[3]);
-        case GGML_OP_SSM_CONV:
         case GGML_OP_SSM_SCAN:
             return has_simdgroup_reduction;
+        case GGML_OP_SSM_CONV:
+            return has_simdgroup_reduction;
         case GGML_OP_RWKV_WKV6:
         case GGML_OP_RWKV_WKV7:
             return true;

ggml/src/ggml-metal/ggml-metal-impl.h

@@ -880,6 +880,7 @@ typedef struct {
     int64_t  n_group;
     int64_t  n_seq_tokens;
     int64_t  n_seqs;
+    int64_t  K;
     uint64_t s_off;
     uint64_t nb00;
     uint64_t nb01;

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

@@ -1710,6 +1710,10 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) {
     const int64_t n_group      = ne41;
     const int64_t n_seq_tokens = ne12;
     const int64_t n_seqs       = ne13;
+    const int64_t K            = ggml_get_op_params_i32(op, 0);
+
+    GGML_ASSERT(K >= 1);
+    GGML_ASSERT(ggml_nelements(op->src[1]) + K*d_state*d_inner*n_head*n_seqs == ggml_nelements(op));
 
     ggml_metal_kargs_ssm_scan args = {
         /*.d_state      =*/ d_state,
@@ -1718,6 +1722,7 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) {
         /*.n_group      =*/ n_group,
         /*.n_seq_tokens =*/ n_seq_tokens,
         /*.n_seqs       =*/ n_seqs,
+        /*.K            =*/ K,
         /*.s_off        =*/ ggml_nelements(op->src[1]) * sizeof(float),
         /*.nb00         =*/ nb00,
         /*.nb01         =*/ nb01,

ggml/src/ggml-metal/ggml-metal.metal

@@ -2429,6 +2429,8 @@ kernel void kernel_ssm_scan_f32(
     const int32_t nh  = args.n_head;
     const int32_t ng  = args.n_group;
     const int32_t n_t = args.n_seq_tokens;
+    const int32_t n_s = args.n_seqs;
+    const int32_t K   = args.K;
 
     const int32_t s_off = args.s_off;
 
@@ -2487,6 +2489,12 @@ kernel void kernel_ssm_scan_f32(
             // recurse
             s0 = s;
 
+            const int32_t slot = n_t - 1 - (i2 + t);
+            if (slot > 0 && slot < K) {
+                device float * s_snapshot = (device float *) ((device char *) s_buff + (int64_t) slot*n_s*args.nb03);
+                s_snapshot[i] = s;
+            }
+
             B  += args.ns42;
             C  += args.ns52;
         }

ggml/src/ggml-sycl/element_wise.cpp

@@ -81,43 +81,6 @@ static __dpct_inline__ T op_elu(T x) {
     return (x > static_cast<T>(0.f)) ? x : op_expm1(x);
 }
 
-template<typename T>
-static __dpct_inline__ T op_tanh(T x) {
-    if constexpr (std::is_same_v<T, sycl::ext::oneapi::bfloat16>) {
-        constexpr int ver = __INTEL_LLVM_COMPILER;
-#if defined(__INTEL_LLVM_COMPILER) && (__INTEL_LLVM_COMPILER >= 20260000)
-            return sycl::ext::oneapi::experimental::tanh(x);
-#else
-            return static_cast<T>(sycl::tanh(static_cast<float>(x)));
-#endif
-    } else {
-        return sycl::tanh(x);
-    }
-}
-
-template<typename T>
-static __dpct_inline__ T op_gelu(T x) {
-    const T GELU_COEF_A    = static_cast<T>(0.044715f);
-    const T SQRT_2_OVER_PI = static_cast<T>(0.79788456080286535587989211986876f);
-    return static_cast<T>(0.5f) * x *
-           (static_cast<T>(1.0f) +
-            op_tanh(SQRT_2_OVER_PI * x * (static_cast<T>(1.0f) + GELU_COEF_A * x * x)));
-}
-
-template<typename T>
-static __dpct_inline__ T op_exp(T x) {
-    if constexpr (std::is_same_v<T, sycl::ext::oneapi::bfloat16>) {
-        return sycl::ext::oneapi::experimental::exp(x);
-    } else {
-        return sycl::exp(x);
-    }
-}
-
-template<typename T>
-static __dpct_inline__ T op_silu(T x) {
-    return x / (static_cast<T>(1.0f) + op_exp(-x));
-}
-
 template<typename T>
 static __dpct_inline__ T op_erf(T x) {
     if constexpr (std::is_same_v<T, sycl::ext::oneapi::bfloat16>) {

ggml/src/ggml-sycl/element_wise.hpp

@@ -28,6 +28,39 @@ typed_data<T_Dst, T_Src> cast_data(ggml_tensor * dst) {
 
 const float GELU_QUICK_COEF = -1.702f;
 
+// Single-element activations, shared with the mat-vec kernels that fuse a GLU epilogue
+// (mmvq.cpp), so both apply the same formula.
+template <typename T> static __dpct_inline__ T op_tanh(T x) {
+    if constexpr (std::is_same_v<T, sycl::ext::oneapi::bfloat16>) {
+#if defined(__INTEL_LLVM_COMPILER) && (__INTEL_LLVM_COMPILER >= 20260000)
+        return sycl::ext::oneapi::experimental::tanh(x);
+#else
+        return static_cast<T>(sycl::tanh(static_cast<float>(x)));
+#endif
+    } else {
+        return sycl::tanh(x);
+    }
+}
+
+template <typename T> static __dpct_inline__ T op_gelu(T x) {
+    const T GELU_COEF_A    = static_cast<T>(0.044715f);
+    const T SQRT_2_OVER_PI = static_cast<T>(0.79788456080286535587989211986876f);
+    return static_cast<T>(0.5f) * x *
+           (static_cast<T>(1.0f) +
+            op_tanh(SQRT_2_OVER_PI * x * (static_cast<T>(1.0f) + GELU_COEF_A * x * x)));
+}
+
+template <typename T> static __dpct_inline__ T op_exp(T x) {
+    if constexpr (std::is_same_v<T, sycl::ext::oneapi::bfloat16>) {
+        return sycl::ext::oneapi::experimental::exp(x);
+    } else {
+        return sycl::exp(x);
+    }
+}
+
+template <typename T> static __dpct_inline__ T op_silu(T x) {
+    return x / (static_cast<T>(1.0f) + op_exp(-x));
+}
 
 void ggml_sycl_sqrt(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
 

ggml/src/ggml-sycl/fusion.cpp

@@ -2,6 +2,61 @@
 
 #include <algorithm>
 
+// mul_mat(gate) + mul_mat(up) + GLU: graph shape and tensor properties only. Backend state
+// (weight layout, split buffers, DMMV) is checked by ggml_sycl_mul_mat_glu_mmvq_fused().
+static bool ggml_sycl_should_fuse_mul_mat_glu(const ggml_tensor * gate, const ggml_tensor * up,
+                                              const ggml_tensor * glu) {
+    // the fused epilogue implements these two; the rest fall back to the standalone GLU kernels
+    const ggml_glu_op glu_op = ggml_get_glu_op(glu);
+    if (glu_op != GGML_GLU_OP_SWIGLU && glu_op != GGML_GLU_OP_GEGLU) {
+        return false;
+    }
+
+    // the kernel always treats src[0] as the activated operand and src[1] as the multiplier
+    if (ggml_get_op_params_i32(glu, 1) /* swapped */) {
+        return false;
+    }
+
+    const ggml_tensor * wu  = up->src[0];
+    const ggml_tensor * wg  = gate->src[0];
+    const ggml_tensor * act = up->src[1];
+
+    // one set of block offsets and one quantized activation must serve both weights
+    if (wu->type != wg->type || !ggml_are_same_shape(wu, wg) || !ggml_are_same_stride(wu, wg)) {
+        return false;
+    }
+    if (act != gate->src[1]) {
+        return false;
+    }
+
+    // only q4_K has a fused reorder GEMV so far, and it walks whole super-blocks
+    if (wu->type != GGML_TYPE_Q4_K || wu->ne[0] % QK_K != 0) {
+        return false;
+    }
+
+    // one 2D reorder-layout matrix in, a plain column stride out: no broadcast or padding
+    if (!ggml_is_contiguous(wu) || !ggml_is_contiguous(wg) || !ggml_is_contiguous(act) ||
+        !ggml_is_contiguous(glu)) {
+        return false;
+    }
+    if (act->type != GGML_TYPE_F32 || glu->type != GGML_TYPE_F32) {
+        return false;
+    }
+    if (act->ne[2] != 1 || act->ne[3] != 1 || wu->ne[2] != 1 || wu->ne[3] != 1) {
+        return false;
+    }
+    // the kernel writes rows [0, wu->ne[1]) of each glu column, strided by glu->ne[0]
+    if (glu->ne[0] != wu->ne[1] || glu->ne[1] != act->ne[1]) {
+        return false;
+    }
+    // mat-vec only: one column per decoded token, up to the batch the reorder kernels cover
+    if (act->ne[1] > MMVQ_MAX_BATCH_SIZE) {
+        return false;
+    }
+
+    return true;
+}
+
 bool ggml_sycl_can_fuse(const ggml_cgraph * cgraph, int node_idx, std::initializer_list<enum ggml_op> ops,
                         std::initializer_list<enum ggml_unary_op> unary_ops) {
 #ifndef NDEBUG
@@ -13,6 +68,28 @@ bool ggml_sycl_can_fuse(const ggml_cgraph * cgraph, int node_idx, std::initializ
         return false;
     }
 
+    // gate and up are siblings, not a chain, so ggml_can_fuse cannot express this: use the
+    // subgraph form with the GLU as the only materialised output.
+    if (ops.size() == 3 && ops.begin()[0] == GGML_OP_MUL_MAT && ops.begin()[1] == GGML_OP_MUL_MAT &&
+        ops.begin()[2] == GGML_OP_GLU) {
+        if (!ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 2 })) {
+            return false;
+        }
+
+        const ggml_tensor * glu  = cgraph->nodes[node_idx + 2];
+        const ggml_tensor * gate = glu->src[0];
+        const ggml_tensor * up   = glu->src[1];
+
+        // don't assume which of the two mat-muls is the gate; infer it from the GLU's operands
+        const bool ok = (gate == cgraph->nodes[node_idx] && up == cgraph->nodes[node_idx + 1]) ||
+                        (gate == cgraph->nodes[node_idx + 1] && up == cgraph->nodes[node_idx]);
+        if (!ok) {
+            return false;
+        }
+
+        return ggml_sycl_should_fuse_mul_mat_glu(gate, up, glu);
+    }
+
     if (!ggml_can_fuse(cgraph, node_idx, ops)) {
         return false;
     }

ggml/src/ggml-sycl/gated_delta_net.cpp

@@ -14,9 +14,9 @@ void gated_delta_net_sycl(const float *     q,
                           const float *     beta,
                           const float *     curr_state,
                           float *           dst,
+                          float *           state,
                           int64_t           H,
                           int64_t           n_tokens,
-                          int64_t           n_seqs,
                           int64_t           sq1,
                           int64_t           sq2,
                           int64_t           sq3,
@@ -29,6 +29,7 @@ void gated_delta_net_sycl(const float *     q,
                           const sycl::uint3 neqk1_magic,
                           const sycl::uint3 rq3_magic,
                           float             scale,
+                          int64_t           state_slot_stride,
                           int               K) {
     auto           item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
     const uint32_t h_idx    = item_ct1.get_group(2);
@@ -40,15 +41,12 @@ void gated_delta_net_sycl(const float *     q,
     const uint32_t iq1 = fastmodulo(h_idx, neqk1_magic);
     const uint32_t iq3 = fastdiv(sequence, rq3_magic);
 
-    const int64_t attn_score_elems = S_v * H * n_tokens * n_seqs;
     float *       attn_data        = dst;
-    float *       state            = dst + attn_score_elems;
 
     // input state holds s0 only [S_v, S_v, H, n_seqs] — seq stride is D = H * S_v * S_v.
     // output state layout (per-slot D * n_seqs) — same per-(seq,head) offset as before.
     const int64_t state_in_offset      = sequence * H * S_v * S_v + h_idx * S_v * S_v;
     const int64_t state_out_offset     = (sequence * H + h_idx) * S_v * S_v;
-    const int64_t state_size_per_token = S_v * S_v * H * n_seqs; // per-slot stride in output
     state += state_out_offset;
     curr_state += state_in_offset + col * S_v;
     attn_data += (sequence * n_tokens * H + h_idx) * S_v;
@@ -145,7 +143,7 @@ void gated_delta_net_sycl(const float *     q,
         if constexpr (keep_rs_t) {
             const int target_slot = (int) n_tokens - 1 - t;
             if (target_slot >= 0 && target_slot < K) {
-                float * curr_state = (dst + attn_score_elems) + target_slot * state_size_per_token + state_out_offset;
+                float * curr_state = state + target_slot * state_slot_stride;
 #pragma unroll
                 for (int r = 0; r < rows_per_lane; r++) {
                     const int i = r * warp_size + lane;
@@ -172,6 +170,7 @@ static void launch_gated_delta_net(const float *   q_d,
                                    const float *   b_d,
                                    const float *   s_d,
                                    float *         dst_d,
+                                   float *         state_d,
                                    int64_t         S_v,
                                    int64_t         H,
                                    int64_t         n_tokens,
@@ -188,6 +187,7 @@ static void launch_gated_delta_net(const float *   q_d,
                                    int64_t         neqk1,
                                    int64_t         rq3,
                                    float           scale,
+                                   int64_t         state_slot_stride,
                                    int             K,
                                    dpct::queue_ptr stream) {
     //TODO: Add chunked kernel for even faster pre-fill
@@ -206,9 +206,9 @@ static void launch_gated_delta_net(const float *   q_d,
                 constexpr int sv = 16;
                 stream->parallel_for(sycl::nd_range<3>(grid_dims * block_dims, block_dims),
                                      [=](sycl::nd_item<3> /*item_ct1*/) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                                         gated_delta_net_sycl<sv, KDA, keep_rs_t>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H, n_tokens,
-                                                                       n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2,
-                                                                       sb3, neqk1_magic, rq3_magic, scale, K);
+                                         gated_delta_net_sycl<sv, KDA, keep_rs_t>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens,
+                                                                       sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2,
+                                                                       sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
                                      });
             }
             break;
@@ -217,9 +217,9 @@ static void launch_gated_delta_net(const float *   q_d,
                 constexpr int sv = 32;
                 stream->parallel_for(sycl::nd_range<3>(grid_dims * block_dims, block_dims),
                                      [=](sycl::nd_item<3> /*item_ct1*/) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                                         gated_delta_net_sycl<sv, KDA, keep_rs_t>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H, n_tokens,
-                                                                       n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2,
-                                                                       sb3, neqk1_magic, rq3_magic, scale, K);
+                                         gated_delta_net_sycl<sv, KDA, keep_rs_t>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens,
+                                                                       sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2,
+                                                                       sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
                                      });
             }
             break;
@@ -229,8 +229,8 @@ static void launch_gated_delta_net(const float *   q_d,
                 stream->parallel_for(sycl::nd_range<3>(grid_dims * block_dims, block_dims),
                                         [=](sycl::nd_item<3> /*item_ct1*/) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
                                             gated_delta_net_sycl<sv, KDA, keep_rs_t>(
-                                                q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H, n_tokens, n_seqs, sq1, sq2,
-                                                sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
+                                                q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens, sq1, sq2,
+                                                sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
                                         });
             }
             break;
@@ -241,8 +241,8 @@ static void launch_gated_delta_net(const float *   q_d,
                 stream->parallel_for(sycl::nd_range<3>(grid_dims * block_dims, block_dims),
                                         [=](sycl::nd_item<3> /*item_ct1*/) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
                                             gated_delta_net_sycl<sv, KDA, keep_rs_t>(
-                                                q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H, n_tokens, n_seqs, sq1, sq2,
-                                                sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
+                                                q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens, sq1, sq2,
+                                                sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
                                         });
             }
             break;
@@ -253,7 +253,8 @@ static void launch_gated_delta_net(const float *   q_d,
     }
 }
 
-void ggml_sycl_op_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
+static void ggml_sycl_op_gated_delta_net_impl(ggml_backend_sycl_context & ctx, ggml_tensor * dst,
+                                              const ggml_sycl_gated_delta_net_fused_cache * cache) {
     ggml_tensor * src_q     = dst->src[0];
     ggml_tensor * src_k     = dst->src[1];
     ggml_tensor * src_v     = dst->src[2];
@@ -318,30 +319,48 @@ void ggml_sycl_op_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor *
     const int K = ggml_get_op_params_i32(dst, 0);
     const bool keep_rs = K > 1;
 
+    // recurrent state -> dst tail (after attention scores), or the cache when fusing
+    float * state_d           = dst_d + S_v * H * n_tokens * n_seqs;
+    int64_t state_slot_stride = S_v * S_v * H * n_seqs;
+    if (cache != nullptr) {
+        state_d           = cache->data;
+        state_slot_stride = cache->slot_stride;
+    }
+
     if (kda) {
         if (keep_rs) {
-            launch_gated_delta_net<true, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
+            launch_gated_delta_net<true, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
                 S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
+                sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream);
         } else {
-            launch_gated_delta_net<true, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
+            launch_gated_delta_net<true, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
                 S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
+                sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream);
         }
     } else {
         if (keep_rs) {
-            launch_gated_delta_net<false, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
+            launch_gated_delta_net<false, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
                 S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
+                sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream);
         } else {
-            launch_gated_delta_net<false, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
+            launch_gated_delta_net<false, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
                 S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
+                sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream);
         }
     }
 }
 
+void ggml_sycl_op_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
+    ggml_sycl_op_gated_delta_net_impl(ctx, dst, nullptr);
+}
+
 void ggml_sycl_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/6);
     ggml_sycl_op_gated_delta_net(ctx, dst);
 }
+
+void ggml_sycl_op_gated_delta_net_fused_cache(ggml_backend_sycl_context & ctx, ggml_tensor * dst,
+                                              ggml_sycl_gated_delta_net_fused_cache cache) {
+    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/6);
+    ggml_sycl_op_gated_delta_net_impl(ctx, dst, &cache);
+}

ggml/src/ggml-sycl/gated_delta_net.hpp

@@ -5,5 +5,15 @@
 #include "common.hpp"
 #include "ggml.h"
 
+// fused-kernel recurrent-state output; strides in elements (per-seq stride is always D, set in-kernel)
+struct ggml_sycl_gated_delta_net_fused_cache {
+    float * data;        // rollback slot 0
+    int64_t slot_stride; // between rollback slots (0 when K==1)
+};
+
 void ggml_sycl_op_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
 void ggml_sycl_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
+
+// same op, but writes the snapshot(s) into the cache instead of dst (see ggml_sycl_try_gdn_cache_fusion)
+void ggml_sycl_op_gated_delta_net_fused_cache(ggml_backend_sycl_context & ctx, ggml_tensor * dst,
+                                              ggml_sycl_gated_delta_net_fused_cache cache);

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

@@ -11,6 +11,7 @@
 //
 
 #include <algorithm>
+#include <array>
 #include <assert.h>
 #include <atomic>
 #include <cinttypes>
@@ -4560,6 +4561,66 @@ static void ggml_sycl_mul_mat(ggml_backend_sycl_context & ctx, const ggml_tensor
     }
 }
 
+// Fused dense-FFN mat-vec for the {mul_mat(gate), mul_mat(up), GLU} subgraph at node_idx.
+// Returns false if it declined, in which case the caller runs the three nodes normally.
+static bool ggml_sycl_mul_mat_glu_mmvq_fused(ggml_backend_sycl_context & ctx, ggml_cgraph * cgraph, int node_idx) {
+    if (!ggml_sycl_can_fuse(cgraph, node_idx, { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU }, {})) {
+        return false;
+    }
+
+    ggml_tensor *       glu  = cgraph->nodes[node_idx + 2];
+    ggml_tensor *       gate = glu->src[0];
+    ggml_tensor *       up   = glu->src[1];
+    const ggml_tensor * wu   = up->src[0];
+    const ggml_tensor * wg   = gate->src[0];
+    const ggml_tensor * act  = up->src[1];
+
+    // this writes glu->data directly rather than the per-device row slices that
+    // ggml_sycl_op_mul_mat() stitches back together, so it cannot serve split weights
+    if (ggml_backend_buffer_is_sycl_split(wu->buffer) || ggml_backend_buffer_is_sycl_split(wg->buffer)) {
+        return false;
+    }
+
+    // with DMMV prioritised the unfused path would not have gone through mmvq at all
+    if (g_ggml_sycl_prioritize_dmmv) {
+        return false;
+    }
+
+    // install the reorder (SoA) layout the fused kernel needs, as the unfused mmvq path would;
+    // a no-op once done. after the bail checks so a declined op does not pay for it.
+    opt_for_reorder(&ctx, wu, act, up, mul_mat_algo::MMVQ);
+    opt_for_reorder(&ctx, wg, act, gate, mul_mat_algo::MMVQ);
+
+    const auto * extra_u = static_cast<const ggml_tensor_extra_gpu *>(wu->extra);
+    const auto * extra_g = static_cast<const ggml_tensor_extra_gpu *>(wg->extra);
+    if (!extra_u || !extra_g || !extra_u->optimized_feature.reorder || !extra_g->optimized_feature.reorder) {
+        return false;
+    }
+
+    // log the up mat-mul: glu's own srcs are the two intermediates the fusion never materialises
+    scope_op_debug_print scope_dbg_print(__func__, up, /*num_src=*/2, " : fused with gate + GLU");
+
+    const int64_t ne00 = wu->ne[0];
+    const int64_t ne11 = act->ne[1];
+
+    const queue_ptr stream           = ctx.stream();
+    const int       src1_padded_cols = GGML_PAD((int) ne00, MATRIX_ROW_PADDING);
+
+    // one activation, quantized once and fully consumed into src1_ddq before the GEMV on this
+    // in-order queue, so glu->data aliasing the dead activation needs no memory-range check
+    ggml_sycl_pool_alloc<char> src1_q8_alloc(ctx.pool(),
+                                             (size_t) ne11 * src1_padded_cols * sizeof(block_q8_1) / QK8_1);
+    char *                     src1_ddq = src1_q8_alloc.get();
+
+    quantize_row_q8_1_sycl<quantize_and_reorder_q8_1_soa>((const float *) act->data, src1_ddq, (int) ne00, (int) ne11,
+                                                          src1_padded_cols, stream);
+
+    return ggml_sycl_mul_mat_vec_q_glu_reorder(wu->type, ggml_get_glu_op(glu), wu->data, wg->data, src1_ddq,
+                                               (float *) glu->data, (int) ne00, (int) wu->ne[1], (int) ne11,
+                                               /*stride_col_y_bytes=*/src1_padded_cols * (int) sizeof(block_q8_1) /
+                                                   QK8_1,
+                                               /*stride_col_dst=*/(int) glu->ne[0], stream);
+}
 
 __dpct_inline__ static void k_copy_src1_to_contiguous(
     const char *__restrict__ src1_original, char *__restrict__ src1_contiguous,
@@ -5464,12 +5525,90 @@ catch (sycl::exception const &exc) {
   std::exit(1);
 }
 
+static bool ggml_sycl_is_view_or_noop(const ggml_tensor * t) {
+    return ggml_is_empty(t) || t->op == GGML_OP_RESHAPE || t->op == GGML_OP_TRANSPOSE ||
+           t->op == GGML_OP_VIEW || t->op == GGML_OP_PERMUTE || t->op == GGML_OP_NONE;
+}
+
+// match gated_delta_net + the strided cpy that scatters its state snapshots into the cache
+// (slot i -> rollback group i, slot 0 newest), so the kernel can write them and skip the cpy.
+// returns the number of following nodes to skip (0 = no fusion)
+// ported from ggml_cuda_try_gdn_cache_fusion - pure graph inspection, backend-agnostic
+static int ggml_sycl_try_gdn_cache_fusion(const ggml_cgraph * cgraph, int node_idx,
+                                          ggml_sycl_gated_delta_net_fused_cache & fused_state_cpy) {
+    if (!g_ggml_sycl_enable_fusion) {
+        return 0;
+    }
+
+    const ggml_tensor * gdn = cgraph->nodes[node_idx];
+    // the kernel skips the snapshot tail, so the gdn output must not be a graph output, and the cpy
+    // found below is taken to be its only reader, as it is in every graph that builds this op
+    if (gdn->op != GGML_OP_GATED_DELTA_NET || gdn->type != GGML_TYPE_F32 ||
+        (gdn->flags & GGML_TENSOR_FLAG_OUTPUT)) {
+        return 0;
+    }
+
+    const ggml_tensor * src_v     = gdn->src[2];
+    const int64_t       S_v       = src_v->ne[0];
+    const int64_t       H         = src_v->ne[1];
+    const int64_t       n_tokens  = src_v->ne[2];
+    const int64_t       n_seqs    = src_v->ne[3];
+    const int64_t       D         = S_v * S_v * H;
+    const int64_t       K         = ggml_get_op_params_i32(gdn, 0); // snapshot slot count
+    const int64_t       n_written = std::min<int64_t>(n_tokens, K); // newest n_written slots are written
+
+    // snapshot tail starts right after the attention scores
+    const size_t tail_off = ggml_row_size(GGML_TYPE_F32, S_v * H * n_tokens * n_seqs);
+
+    // the cpy must be the first node the compute loop below runs, so nothing can read the cache first.
+    // skip exactly what that loop skips: views, no-ops, and nodes the graph does not compute.
+    const ggml_tensor * cpy  = nullptr;
+    int                 skip = 0;
+    for (int j = node_idx + 1; j < cgraph->n_nodes && cpy == nullptr; ++j) {
+        const ggml_tensor * n = cgraph->nodes[j];
+        if (ggml_sycl_is_view_or_noop(n) || (n->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
+            continue;
+        }
+        if (n->op != GGML_OP_CPY || (n->flags & GGML_TENSOR_FLAG_OUTPUT)) {
+            return 0;
+        }
+        cpy  = n;
+        skip = j - node_idx;
+    }
+    if (cpy == nullptr) {
+        return 0;
+    }
+
+    const ggml_tensor * src = cpy->src[0]; // view of the gdn snapshot tail
+    const ggml_tensor * dst = cpy->src[1]; // cache view the kernel writes to
+
+    // src must be this gdn's snapshot tail (contiguous, at the tail offset)
+    if (src->op != GGML_OP_VIEW || src->view_src != gdn || src->view_offs != tail_off ||
+        !ggml_is_contiguous(src)) {
+        return 0;
+    }
+
+    // dst is the [D, n_seqs, n_written] cache view, with the per-seq stride D that the kernel assumes.
+    // ggml_cpy pins src to the same element count, so src needs no shape check of its own.
+    const std::array<int64_t, GGML_MAX_DIMS> expected_ne = { D, n_seqs, n_written, 1 };
+    if (dst->op != GGML_OP_VIEW || dst->type != GGML_TYPE_F32 || dst->data == nullptr ||
+        !std::equal(expected_ne.begin(), expected_ne.end(), dst->ne) ||
+        dst->nb[0] != ggml_type_size(GGML_TYPE_F32) ||
+        dst->nb[1] != (size_t) ggml_row_size(GGML_TYPE_F32, D)) {
+        return 0;
+    }
+
+    fused_state_cpy.data        = (float *) dst->data; // rollback group 0 (newest)
+    fused_state_cpy.slot_stride = K > 1 ? (int64_t) (dst->nb[2] / sizeof(float)) : 0;
+    return skip;
+}
+
 static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * sycl_ctx, ggml_cgraph * cgraph) {
     ggml_sycl_set_main_device(sycl_ctx->device);
 
     for (int i = 0; i < cgraph->n_nodes; i++) {
         ggml_tensor * node = cgraph->nodes[i];
-        if (ggml_is_empty(node) || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_NONE) {
+        if (ggml_sycl_is_view_or_noop(node)) {
             continue;
         }
         if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
@@ -5489,6 +5628,16 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc
             }
         }
 #endif
+        // gated_delta_net -> cpy: scatter recurrent-state snapshots into the cache
+        if (node->op == GGML_OP_GATED_DELTA_NET) {
+            ggml_sycl_gated_delta_net_fused_cache fused_state_cpy;
+            const int gdn_nodes_to_skip = ggml_sycl_try_gdn_cache_fusion(cgraph, i, fused_state_cpy);
+            if (gdn_nodes_to_skip > 0) {
+                ggml_sycl_op_gated_delta_net_fused_cache(*sycl_ctx, node, fused_state_cpy);
+                i += gdn_nodes_to_skip;
+                continue;
+            }
+        }
         if (node->op == GGML_OP_RMS_NORM &&
             ggml_sycl_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL }, {})) {
             ggml_sycl_op_rms_norm_fused(*sycl_ctx, node, cgraph->nodes[i + 1]);
@@ -5502,6 +5651,11 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc
             continue;
         }
 
+        if (node->op == GGML_OP_MUL_MAT && ggml_sycl_mul_mat_glu_mmvq_fused(*sycl_ctx, cgraph, i)) {
+            i += 2;
+            continue;
+        }
+
         bool ok = ggml_sycl_compute_forward(*sycl_ctx, node);
         if (!ok) {
             GGML_LOG_ERROR("%s: error: op not supported %s (%s)\n", __func__, node->name, ggml_op_name(node->op));

ggml/src/ggml-sycl/mmvq.cpp

@@ -2,6 +2,7 @@
 
 #include "ggml.h"
 #include "common.hpp"
+#include "element_wise.hpp"
 #include "quants.hpp"
 #include "vecdotq.hpp"
 
@@ -56,11 +57,13 @@ static void mul_mat_vec_q_reorder(const void * __restrict__ vx, const void * __r
     }
 }
 
-template <typename reorder_vec_dot_q_sycl, int ncols_dst>
-static void mul_mat_vec_q_reorder_ncols(const void * __restrict__ vx, const void * __restrict__ vy,
-                                        float * __restrict__ dst, const int ncols, const int nrows,
-                                        const int stride_col_y_bytes, const int stride_col_dst,
-                                        const sycl::nd_item<3> & nd_item) {
+// With has_fusion, `vgate` is a second weight matrix sharing vx's shape, stride and reorder
+// layout: one pass computes both row dot products and the epilogue writes glu(gate, up).
+template <typename reorder_vec_dot_q_sycl, int ncols_dst, bool has_fusion = false>
+static void mul_mat_vec_q_reorder_ncols(const void * __restrict__ vx, const void * __restrict__ vgate,
+                                        const void * __restrict__ vy, float * __restrict__ dst, const int ncols,
+                                        const int nrows, const int stride_col_y_bytes, const int stride_col_dst,
+                                        const ggml_glu_op glu_op, const sycl::nd_item<3> & nd_item) {
     using block_type   = ggml_sycl_reordered::block_q_t<reorder_vec_dot_q_sycl::gtype>;
     using block_traits = typename block_type::traits;
 
@@ -70,6 +73,8 @@ static void mul_mat_vec_q_reorder_ncols(const void * __restrict__ vx, const void
     const int  sg_id        = sg.get_group_linear_id();
     const int  row          = workgroup_id * sg_range + sg_id;
 
+    // row is sub-group uniform, so this retires whole sub-groups and the collectives below
+    // stay convergent
     if (row >= nrows) {
         return;
     }
@@ -82,10 +87,15 @@ static void mul_mat_vec_q_reorder_ncols(const void * __restrict__ vx, const void
     static_assert(blocks_per_subgroup > 0);
     static_assert(block_elements_per_subgroup > 0);
 
-    float partial_sum[ncols_dst] = {0.0f};
+    float partial_sum[ncols_dst] = { 0.0f };
+    // sized 1 rather than 0 when unused: zero-length arrays are not standard C++, and the
+    // array is dead and eliminated in that case
+    [[maybe_unused]] float partial_gate[has_fusion ? ncols_dst : 1] = { 0.0f };
     for (int i = sg.get_local_linear_id() / block_elements_per_subgroup; i < blocks_per_row; i += blocks_per_subgroup) {
         const int ibx = row * blocks_per_row + i;
 
+        // the offsets depend only on the block index and the matrix shape, never on the base
+        // pointer, which is what lets vgate reuse them
         const auto bx_offset = block_type::get_block_offset(ibx, nblocks);
         const auto d_offset  = block_type::get_d_offset(nrows, ncols, ibx);
         const int  iby       = i * block_type::block_to_q8_1_ratio();
@@ -96,11 +106,16 @@ static void mul_mat_vec_q_reorder_ncols(const void * __restrict__ vx, const void
 
 #pragma unroll
             for (int j = 0; j < ncols_dst; ++j) {
-                const char       * vy_j           = (const char *)vy + j * stride_col_y_bytes;
-                const int8_t     * q8_1_quant_ptr = (const int8_t *)vy_j + iby * QK8_1;
-                const sycl::half2* q8_1_ds_ptr    = (const sycl::half2 *)(vy_j + ncols + iby * sizeof(sycl::half2));
+                const char        * vy_j           = (const char *) vy + j * stride_col_y_bytes;
+                const int8_t      * q8_1_quant_ptr = (const int8_t *) vy_j + iby * QK8_1;
+                const sycl::half2 * q8_1_ds_ptr    = (const sycl::half2 *) (vy_j + ncols + iby * sizeof(sycl::half2));
 
                 partial_sum[j] += reorder_vec_dot_q_sycl()(vx, bx_offset, d_offset, q8_1_quant_ptr, q8_1_ds_ptr, iqs);
+
+                if constexpr (has_fusion) {
+                    partial_gate[j] +=
+                        reorder_vec_dot_q_sycl()(vgate, bx_offset, d_offset, q8_1_quant_ptr, q8_1_ds_ptr, iqs);
+                }
             }
         }
     }
@@ -109,6 +124,13 @@ static void mul_mat_vec_q_reorder_ncols(const void * __restrict__ vx, const void
     for (int j = 0; j < ncols_dst; ++j) {
         float sum = sycl::reduce_over_group(nd_item.get_sub_group(), partial_sum[j], std::plus<>());
 
+        if constexpr (has_fusion) {
+            const float gate = sycl::reduce_over_group(nd_item.get_sub_group(), partial_gate[j], std::plus<>());
+
+            // uniform across the launch; the launcher only instantiates SWIGLU and GEGLU
+            sum *= glu_op == GGML_GLU_OP_SWIGLU ? op_silu(gate) : op_gelu(gate);
+        }
+
         if (sg.leader()) {
             dst[j * stride_col_dst + row] = sum;
         }
@@ -691,7 +713,8 @@ static void reorder_mul_mat_vec_q4_0_q8_1_sycl_ncols(
         cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
                          [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
                              mul_mat_vec_q_reorder_ncols<reorder_vec_dot_q_sycl<GGML_TYPE_Q4_0>, ncols_dst>(
-                                 vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, nd_item);
+                                 vx, /*vgate=*/ nullptr, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst,
+                                 /*glu_op=*/ GGML_GLU_OP_SWIGLU, nd_item);
                          });
     });
 }
@@ -1108,7 +1131,8 @@ static void reorder_mul_mat_vec_q8_0_q8_1_sycl_ncols(
         cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
                          [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
                              mul_mat_vec_q_reorder_ncols<reorder_vec_dot_q_sycl<GGML_TYPE_Q8_0>, ncols_dst>(
-                                 vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, nd_item);
+                                 vx, /*vgate=*/ nullptr, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst,
+                                 /*glu_op=*/ GGML_GLU_OP_SWIGLU, nd_item);
                          });
     });
 }
@@ -1436,7 +1460,8 @@ static void reorder_mul_mat_vec_q3_k_q8_1_sycl_ncols(
         cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
                          [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
                              mul_mat_vec_q_reorder_ncols<reorder_vec_dot_q_sycl<GGML_TYPE_Q3_K>, ncols_dst>(
-                                 vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, nd_item);
+                                 vx, /*vgate=*/ nullptr, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst,
+                                 /*glu_op=*/ GGML_GLU_OP_SWIGLU, nd_item);
                          });
     });
 }
@@ -1604,7 +1629,8 @@ static void reorder_mul_mat_vec_q4_k_q8_1_sycl_ncols(
         cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
                          [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
                              mul_mat_vec_q_reorder_ncols<reorder_vec_dot_q_sycl<GGML_TYPE_Q4_K>, ncols_dst>(
-                                 vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, nd_item);
+                                 vx, /*vgate=*/ nullptr, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst,
+                                 /*glu_op=*/ GGML_GLU_OP_SWIGLU, nd_item);
                          });
     });
 }
@@ -1731,7 +1757,8 @@ static void reorder_mul_mat_vec_q5_k_q8_1_sycl_ncols(
         cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
                          [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
                              mul_mat_vec_q_reorder_ncols<reorder_vec_dot_q_sycl<GGML_TYPE_Q5_K>, ncols_dst>(
-                                 vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, nd_item);
+                                 vx, /*vgate=*/ nullptr, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst,
+                                 /*glu_op=*/ GGML_GLU_OP_SWIGLU, nd_item);
                          });
     });
 }
@@ -1789,7 +1816,8 @@ static void reorder_mul_mat_vec_q6_k_q8_1_sycl_ncols(
         cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
                          [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
                              mul_mat_vec_q_reorder_ncols<reorder_vec_dot_q_sycl<GGML_TYPE_Q6_K>, ncols_dst>(
-                                 vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, nd_item);
+                                 vx, /*vgate=*/ nullptr, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst,
+                                 /*glu_op=*/ GGML_GLU_OP_SWIGLU, nd_item);
                          });
     });
 }
@@ -2736,3 +2764,77 @@ bool ggml_sycl_mul_mat_vec_q_id_reorder(
             return false;
     }
 }
+
+template <typename reorder_vec_dot_q_sycl, int ncols_dst>
+static void launch_mul_mat_vec_q_reorder_glu(const void * vx, const void * vgate, const void * vy, float * dst,
+                                             const int ncols, const int nrows, const int stride_col_y_bytes,
+                                             const int stride_col_dst, const ggml_glu_op glu_op,
+                                             dpct::queue_ptr stream) {
+    GGML_ASSERT(ncols % QK_K == 0);
+
+    constexpr size_t num_subgroups = WARP_SIZE;
+
+    const int            block_num_y = ceil_div(nrows, GGML_SYCL_MMV_Y * (int) num_subgroups);
+    const sycl::range<3> block_nums(1, 1, block_num_y);
+    const sycl::range<3> block_dims(1, GGML_SYCL_MMV_Y, num_subgroups * WARP_SIZE);
+
+    stream->submit([&](sycl::handler & cgh) {
+        cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
+                         [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                             mul_mat_vec_q_reorder_ncols<reorder_vec_dot_q_sycl, ncols_dst, /*has_fusion=*/ true>(
+                                 vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, glu_op,
+                                 nd_item);
+                         });
+    });
+}
+
+bool ggml_sycl_mul_mat_vec_q_glu_reorder(enum ggml_type src0_type, enum ggml_glu_op glu_op, const void * vx,
+                                         const void * vgate, const void * vy, float * dst, int ncols, int nrows,
+                                         int ncols_dst, int stride_col_y_bytes, int stride_col_dst,
+                                         dpct::queue_ptr stream) {
+    if (src0_type != GGML_TYPE_Q4_K) {
+        return false;
+    }
+    if (glu_op != GGML_GLU_OP_SWIGLU && glu_op != GGML_GLU_OP_GEGLU) {
+        return false;
+    }
+
+    using vec_dot = reorder_vec_dot_q_sycl<GGML_TYPE_Q4_K>;
+
+    switch (ncols_dst) {
+        case 1:
+            launch_mul_mat_vec_q_reorder_glu<vec_dot, 1>(vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes,
+                                                         stride_col_dst, glu_op, stream);
+            return true;
+        case 2:
+            launch_mul_mat_vec_q_reorder_glu<vec_dot, 2>(vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes,
+                                                         stride_col_dst, glu_op, stream);
+            return true;
+        case 3:
+            launch_mul_mat_vec_q_reorder_glu<vec_dot, 3>(vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes,
+                                                         stride_col_dst, glu_op, stream);
+            return true;
+        case 4:
+            launch_mul_mat_vec_q_reorder_glu<vec_dot, 4>(vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes,
+                                                         stride_col_dst, glu_op, stream);
+            return true;
+        case 5:
+            launch_mul_mat_vec_q_reorder_glu<vec_dot, 5>(vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes,
+                                                         stride_col_dst, glu_op, stream);
+            return true;
+        case 6:
+            launch_mul_mat_vec_q_reorder_glu<vec_dot, 6>(vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes,
+                                                         stride_col_dst, glu_op, stream);
+            return true;
+        case 7:
+            launch_mul_mat_vec_q_reorder_glu<vec_dot, 7>(vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes,
+                                                         stride_col_dst, glu_op, stream);
+            return true;
+        case 8:
+            launch_mul_mat_vec_q_reorder_glu<vec_dot, 8>(vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes,
+                                                         stride_col_dst, glu_op, stream);
+            return true;
+        default:
+            return false;
+    }
+}

ggml/src/ggml-sycl/mmvq.hpp

@@ -57,4 +57,20 @@ bool ggml_sycl_mul_mat_vec_q_id_reorder(
     size_t             src1_row_stride,
     dpct::queue_ptr    stream);
 
+// Fused dense-FFN GEMV: writes glu(gate . y, up . y) instead of the two mat-vec results.
+// vx / vgate must share shape, stride and reorder layout. Returns false if unhandled.
+bool ggml_sycl_mul_mat_vec_q_glu_reorder(
+    enum ggml_type     src0_type,
+    enum ggml_glu_op   glu_op,
+    const void *       vx,
+    const void *       vgate,
+    const void *       vy,
+    float *            dst,
+    int                ncols,                // K, shared by both weights
+    int                nrows,                // output rows, i.e. weight ne[1]
+    int                ncols_dst,            // activation columns, 1..MMVQ_MAX_BATCH_SIZE
+    int                stride_col_y_bytes,   // bytes between activation columns in vy
+    int                stride_col_dst,       // floats between output columns in dst
+    dpct::queue_ptr    stream);
+
 #endif // GGML_SYCL_MMVQ_HPP

ggml/src/ggml-sycl/ssm_scan.cpp

@@ -10,6 +10,7 @@ static void ssm_scan_f32_group(
         const int src2_nb1, const int src2_nb2, const int src3_nb1,
         const int src4_nb2, const int src4_nb3, const int src5_nb2, const int src5_nb3,
         const int64_t s_off, const int64_t n_head, const int64_t d_head, const int64_t n_group, const int64_t n_tok,
+        const int64_t K,
         const sycl::nd_item<2> & item) {
 
     const int lane     = item.get_local_id(1) % WARP_SIZE;
@@ -64,6 +65,15 @@ static void ssm_scan_f32_group(
         if (lane == 0) {
             y_warp[i * stride_y] = state_sum;
         }
+
+        const int64_t slot = n_tok - 1 - i;
+        if (K > 1 && slot > 0 && slot < K) {
+            float * s_snapshot_warp = (float *) ((char *) dst + s_off + (slot * item.get_group_range(0) + seq_idx) * src0_nb3 + head_idx * src0_nb2 + head_off * d_state);
+#pragma unroll
+            for (int j = 0; j < c_factor; j++) {
+                s_snapshot_warp[WARP_SIZE * j + lane] = state[j];
+            }
+        }
     }
 
 #pragma unroll
@@ -79,6 +89,7 @@ static void ssm_scan_f32_sycl(
         const int src2_nb2, const int src3_nb1, const int src4_nb2, const int src4_nb3, const int src5_nb2,
         const int src5_nb3, const int64_t s_off, const int64_t d_state, const int64_t head_dim,
         const int64_t n_head, const int64_t n_group, const int64_t n_tok, const int64_t n_seq,
+        const int64_t K,
         dpct::queue_ptr stream) {
 
     // NOTE: if you change conditions here, be sure to update the corresponding supports_op condition!
@@ -94,7 +105,7 @@ static void ssm_scan_f32_sycl(
                 ssm_scan_f32_group<128 / WARP_SIZE, 128>(
                     src0, src1, src2, src3, src4, src5, src6, dst,
                     src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
-                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, item);
+                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K, item);
             });
     } else if (d_state == 256) {
         constexpr int threads   = 256;
@@ -107,7 +118,7 @@ static void ssm_scan_f32_sycl(
                 ssm_scan_f32_group<256 / WARP_SIZE, 256>(
                     src0, src1, src2, src3, src4, src5, src6, dst,
                     src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
-                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, item);
+                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K, item);
             });
     } else {
         GGML_ABORT("ssm_scan: unsupported d_state (must be 128 or 256)");
@@ -133,9 +144,12 @@ inline void ggml_sycl_op_ssm_scan(ggml_backend_sycl_context & ctx, ggml_tensor *
     const int64_t ng  = src4->ne[1];
     const int64_t n_t = src1->ne[2];
     const int64_t n_s = src1->ne[3];
+    const int64_t K   = ggml_get_op_params_i32(dst, 0);
     const int64_t s_off = ggml_nelements(src1) * sizeof(float);
 
-    GGML_ASSERT(ggml_nelements(src1) + nc * nr * nh * n_s == ggml_nelements(dst));
+    GGML_ASSERT(K >= 1);
+    GGML_ASSERT(ggml_nelements(src1) + K * nc * nr * nh * n_s == ggml_nelements(dst));
+    GGML_ASSERT(src3->ne[0] == 1 || K == 1);
 
     dpct::queue_ptr stream = ctx.stream();
     SYCL_CHECK(ggml_sycl_set_device(ctx.device));
@@ -147,7 +161,7 @@ inline void ggml_sycl_op_ssm_scan(ggml_backend_sycl_context & ctx, ggml_tensor *
         static_cast<const int32_t *>(src6->data), static_cast<float *>(dst->data),
         src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2],
         src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3],
-        s_off, nc, nr, nh, ng, n_t, n_s, stream);
+        s_off, nc, nr, nh, ng, n_t, n_s, K, stream);
 }
 
 void ggml_sycl_ssm_scan(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -1861,6 +1861,7 @@ struct vk_op_ssm_scan_push_constants {
     uint32_t nb42, nb43, nb52, nb53;
     uint32_t s_off;
     uint32_t n_head, d_head, n_group, n_tok;
+    uint32_t n_seq, K;
 };
 struct vk_op_ssm_conv_push_constants {
     uint32_t nb01, nb02;
@@ -2065,7 +2066,7 @@ struct ggml_vk_garbage_collector {
 static void ggml_vk_preallocate_buffers(ggml_backend_vk_context * ctx, vk_context subctx);
 static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested = nullptr);
 static void ggml_pipeline_allocate_descriptor_sets(ggml_backend_vk_context * ctx);
-static bool ggml_vk_intel_windows_driver_equals_or_newer_than(uint32_t driver_version, uint32_t threshold_major, uint32_t threshold_minor);
+static bool ggml_vk_intel_windows_driver_in_range(uint32_t driver_version, uint32_t lower_major, uint32_t lower_minor, uint32_t upper_major, uint32_t upper_minor);
 
 static bool vk_memory_logger_enabled = false;
 
@@ -3961,7 +3962,10 @@ static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vec
     }
 
     // Needs to be kept up to date on shader changes
-    const uint32_t bank_conflict_offset = device->coopmat_support ? 8 : 1;
+    // Needs to stay aligned with ggml_vk_mul_mm_spec.
+    const bool intel_shmem_stride_pad_zero = device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support &&
+                                              device->driver_id == vk::DriverId::eIntelProprietaryWindows;
+    const uint32_t bank_conflict_offset = intel_shmem_stride_pad_zero ? 0 : (device->coopmat_support ? 8 : 1);
     const uint32_t type_size = device->fp16 ? sizeof(ggml_fp16_t) : sizeof(float);
     const uint32_t warps = warptile[0] / warptile[10];
 
@@ -4578,8 +4582,13 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     }
 #endif
 
-    auto const &ggml_vk_mul_mm_spec = [](std::vector<uint32_t> spec, bool aligned) {
-        spec.push_back(aligned ? 1u : 0u);
+    auto const &ggml_vk_mul_mm_spec = [&device](std::vector<uint32_t> spec, bool aligned) {
+        spec.push_back(aligned ? 1u : 0u);  // constantID=11: ALIGNED
+        if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support &&
+            device->driver_id == vk::DriverId::eIntelProprietaryWindows) {
+            spec.push_back(0u);  // constantID=12: SHMEM_STRIDE_PAD = 0
+            spec.push_back(1u);  // constantID=13: APPLY_SLM_A_RESHAPE = true
+        }
         return spec;
     };
 
@@ -5741,10 +5750,9 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     ggml_vk_create_pipeline(device, device->pipeline_argmax_f32, "argmax_f32", argmax_f32_len, argmax_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
 
     ggml_vk_create_pipeline(device, device->pipeline_sum_rows_f32, "sum_rows_f32", sum_rows_f32_len, sum_rows_f32_data, "main", 2, sizeof(vk_op_sum_rows_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
-    // Intel Windows driver older than 32.0.101.8860 will crash when using fwht kernels on Xe2+ GPUS so we gate that here
+    // Intel Windows driver in range [32.0.101.8509, 32.0.101.8860) will crash when using fwht kernels so we gate that here
     const bool can_use_fwht = device->driver_id != vk::DriverId::eIntelProprietaryWindows ||
-        device->architecture != vk_device_architecture::INTEL_XE2 ||
-        (device->architecture == vk_device_architecture::INTEL_XE2 && ggml_vk_intel_windows_driver_equals_or_newer_than(device->properties.driverVersion, 101, 8860));
+        !ggml_vk_intel_windows_driver_in_range(device->properties.driverVersion, 101, 8509, 101, 8860);
     if (can_use_fwht && device->subgroup_basic && device->subgroup_shuffle) {
         int idx = 0;
         for (uint32_t n : {64, 128, 256, 512}) {
@@ -12731,7 +12739,8 @@ static void ggml_vk_ssm_scan(ggml_backend_vk_context * ctx, vk_context& subctx,
         (uint32_t)src4->nb[2], (uint32_t)src4->nb[3],
         (uint32_t)src5->nb[2], (uint32_t)src5->nb[3],
         (uint32_t)s_off,
-        n_head, head_dim, n_group, n_tok
+        n_head, head_dim, n_group, n_tok,
+        n_seq, (uint32_t) ggml_get_op_params_i32(dst, 0)
     };
 
     vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst);
@@ -18869,17 +18878,23 @@ static uint32_t ggml_vk_intel_shader_core_count(const vk::PhysicalDevice& vkdev)
     }
 }
 
-static bool ggml_vk_intel_windows_driver_equals_or_newer_than(uint32_t driver_version, uint32_t threshold_major, uint32_t threshold_minor) {
+// checks whether lower <= driver_version < upper, with each bound given as xxx.yyyy
+static bool ggml_vk_intel_windows_driver_in_range(uint32_t driver_version, uint32_t lower_major, uint32_t lower_minor, uint32_t upper_major, uint32_t upper_minor) {
 #if defined(_WIN32)
     // Intel Windows encodes xxx.yyyy as [31:14].[13:0].
     const uint32_t major = driver_version >> 14;
     const uint32_t minor = driver_version & 0x3fff;
 
-    return major > threshold_major || (major == threshold_major && minor >= threshold_minor);
+    const bool ge_lower = major > lower_major || (major == lower_major && minor >= lower_minor);
+    const bool lt_upper = major < upper_major || (major == upper_major && minor < upper_minor);
+
+    return ge_lower && lt_upper;
 #else
     GGML_UNUSED(driver_version);
-    GGML_UNUSED(threshold_major);
-    GGML_UNUSED(threshold_minor);
+    GGML_UNUSED(lower_major);
+    GGML_UNUSED(lower_minor);
+    GGML_UNUSED(upper_major);
+    GGML_UNUSED(upper_minor);
     return true;
 #endif
 }
@@ -19417,8 +19432,9 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
         } else if (tensor->op == GGML_OP_ADD_ID) {
             tensor_clone = ggml_add_id(ggml_ctx, src_clone[0], src_clone[1], src_clone[2]);
         } else if (tensor->op == GGML_OP_SSM_SCAN) {
+            const int32_t K = ggml_get_op_params_i32(tensor, 0);
             tensor_clone = ggml_ssm_scan(ggml_ctx, src_clone[0], src_clone[1], src_clone[2],
-                                         src_clone[3], src_clone[4], src_clone[5], src_clone[6]);
+                                         src_clone[3], src_clone[4], src_clone[5], src_clone[6], K);
         } else if (tensor->op == GGML_OP_SSM_CONV) {
             tensor_clone = ggml_ssm_conv(ggml_ctx, src_clone[0], src_clone[1]);
         } else if (tensor->op == GGML_OP_ROLL) {

ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp

@@ -119,10 +119,13 @@ layout (constant_id = 3) const uint BK = 16;  // Assumed to be 32 if working wit
 #endif
 
 #ifdef COOPMAT
-#define SHMEM_STRIDE (BK / 2 + 4)
+layout(constant_id = 12) const uint SHMEM_STRIDE_PAD = 4;
+layout(constant_id = 13) const bool APPLY_SLM_A_RESHAPE = false;
 #else
-#define SHMEM_STRIDE (BK / 2 + 1)
+const uint SHMEM_STRIDE_PAD = 1;
+const bool APPLY_SLM_A_RESHAPE = false;
 #endif
+#define SHMEM_STRIDE (BK / 2 + SHMEM_STRIDE_PAD)
 
 shared FLOAT_TYPEV2 buf_a[BM * SHMEM_STRIDE];
 shared FLOAT_TYPEV2 buf_b[BN * SHMEM_STRIDE];
@@ -302,7 +305,7 @@ void main() {
         [[unroll]] for (uint i = 0; i < BK; i += TK) {
             [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) {
                 // Load from shared into cache
-                coopMatLoad(cache_a, buf_a, (warp_r * WM + cm_row * TM) * SHMEM_STRIDE + i / 2, SHMEM_STRIDE, gl_CooperativeMatrixLayoutRowMajor);
+                coopMatLoad(cache_a, buf_a, a_shmem_index(warp_r * WM + cm_row * TM, i / 2), a_shmem_stride(), gl_CooperativeMatrixLayoutRowMajor);
 
                 [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) {
                     coopMatLoad(cache_b, buf_b, (warp_c * WN + cm_col * TN) * SHMEM_STRIDE + i / 2, SHMEM_STRIDE, gl_CooperativeMatrixLayoutColumnMajor);

ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl

@@ -1,60 +1,76 @@
+// k_pair is the K coordinate measured in FLOAT_TYPEV2 elements.
+uint a_shmem_index(uint m, uint k_pair) {
+    if (APPLY_SLM_A_RESHAPE) {
+        const uint tile_width = TK / 2;
+        return (k_pair / tile_width) * BM * tile_width
+             + m * tile_width
+             + k_pair % tile_width;
+    }
+    return m * SHMEM_STRIDE + k_pair;
+}
+
+uint a_shmem_stride() {
+    return APPLY_SLM_A_RESHAPE ? TK / 2 : SHMEM_STRIDE;
+}
+
+void store_a(uint m, uint k_pair, FLOAT_TYPEV2 value) {
+    buf_a[a_shmem_index(m, k_pair)] = value;
+}
+
 void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uint idx_m, const uint block, const uint end_k) {
 #if defined(DATA_A_F32) || defined(DATA_A_F16)
 #if LOAD_VEC_A == 8
             if (ALIGNED != 0) {
                 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-                const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
+                const uint k_pair = row * LOAD_VEC_A / 2;
                 FLOAT_TYPEV8 aa = FLOAT_TYPEV8(data_a[idx]);
-                buf_a[buf_idx    ] = aa[0].xy;
-                buf_a[buf_idx + 1] = aa[0].zw;
-                buf_a[buf_idx + 2] = aa[1].xy;
-                buf_a[buf_idx + 3] = aa[1].zw;
+                store_a(col, k_pair,     aa[0].xy);
+                store_a(col, k_pair + 1, aa[0].zw);
+                store_a(col, k_pair + 2, aa[1].xy);
+                store_a(col, k_pair + 3, aa[1].zw);
                 return;
             }
 #elif LOAD_VEC_A == 4
             if (ALIGNED != 0) {
                 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-                const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
+                const uint k_pair = row * LOAD_VEC_A / 2;
                 FLOAT_TYPEV4 aa = FLOAT_TYPEV4(data_a[idx]);
-                buf_a[buf_idx    ] = aa.xy;
-                buf_a[buf_idx + 1] = aa.zw;
+                store_a(col, k_pair,     aa.xy);
+                store_a(col, k_pair + 1, aa.zw);
                 return;
             }
 #endif
             const uint idx = pos_a + col * p.stride_a + row * 2;
-            const uint buf_idx = col * SHMEM_STRIDE + row;
             if (idx_m < p.M && block + row * 2 + 1 < end_k) {
-                buf_a[buf_idx] = FLOAT_TYPEV2(data_a_scalar[idx],
-                                              data_a_scalar[idx + 1]);
+                store_a(col, row, FLOAT_TYPEV2(data_a_scalar[idx],
+                                               data_a_scalar[idx + 1]));
             } else if (idx_m < p.M && block + row * 2 < end_k) {
-                buf_a[buf_idx] = FLOAT_TYPEV2(data_a_scalar[idx], 0.0f);
+                store_a(col, row, FLOAT_TYPEV2(data_a_scalar[idx], 0.0f));
             } else {
-                buf_a[buf_idx] = FLOAT_TYPEV2(0.0f);
+                store_a(col, row, FLOAT_TYPEV2(0.0f));
             }
 #elif defined(DATA_A_BF16)
 #if LOAD_VEC_A == 4
             if (ALIGNED != 0) {
                 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-                const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
+                const uint k_pair = row * LOAD_VEC_A / 2;
                 FLOAT_TYPEV4 aa = FLOAT_TYPEV4(TO_FLOAT_TYPE(data_a[idx]));
-                buf_a[buf_idx    ] = aa.xy;
-                buf_a[buf_idx + 1] = aa.zw;
+                store_a(col, k_pair,     aa.xy);
+                store_a(col, k_pair + 1, aa.zw);
                 return;
             }
 #endif
             const uint idx = pos_a + col * p.stride_a + row * 2;
-            const uint buf_idx = col * SHMEM_STRIDE + row;
             if (idx_m < p.M && block + row * 2 + 1 < end_k) {
-                buf_a[buf_idx] = FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]),
-                                              TO_FLOAT_TYPE(data_a_scalar[idx + 1]));
+                store_a(col, row, FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]),
+                                               TO_FLOAT_TYPE(data_a_scalar[idx + 1])));
             } else if (idx_m < p.M && block + row * 2 < end_k) {
-                buf_a[buf_idx] = FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]), 0.0f);
+                store_a(col, row, FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]), 0.0f));
             } else {
-                buf_a[buf_idx] = FLOAT_TYPEV2(0.0f);
+                store_a(col, row, FLOAT_TYPEV2(0.0f));
             }
 #elif defined(DATA_A_Q4_0)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
 
             const uint ib = idx / 4;
             const uint iqs = idx & 0x03;
@@ -64,13 +80,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const vec4 v0 = (vec4(unpack8(vui & 0x0F0F0F0F)) - 8.0f) * d;
             const vec4 v1 = (vec4(unpack8((vui >> 4) & 0x0F0F0F0F)) - 8.0f) * d;
 
-            buf_a[buf_idx    ] = FLOAT_TYPEV2(v0.xy);
-            buf_a[buf_idx + 1] = FLOAT_TYPEV2(v0.zw);
-            buf_a[buf_idx + 8] = FLOAT_TYPEV2(v1.xy);
-            buf_a[buf_idx + 9] = FLOAT_TYPEV2(v1.zw);
+            const uint k_pair = row * LOAD_VEC_A / 4;
+            store_a(col, k_pair,     FLOAT_TYPEV2(v0.xy));
+            store_a(col, k_pair + 1, FLOAT_TYPEV2(v0.zw));
+            store_a(col, k_pair + 8, FLOAT_TYPEV2(v1.xy));
+            store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.zw));
 #elif defined(DATA_A_Q4_1)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
 
             const uint ib = idx / 4;
             const uint iqs = idx & 0x03;
@@ -80,13 +96,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const vec4 v0 = vec4(unpack8(vui & 0x0F0F0F0F)) * dm.x + dm.y;
             const vec4 v1 = vec4(unpack8((vui >> 4) & 0x0F0F0F0F)) * dm.x + dm.y;
 
-            buf_a[buf_idx     ] = FLOAT_TYPEV2(v0.xy);
-            buf_a[buf_idx + 1 ] = FLOAT_TYPEV2(v0.zw);
-            buf_a[buf_idx + 8 ] = FLOAT_TYPEV2(v1.xy);
-            buf_a[buf_idx + 9 ] = FLOAT_TYPEV2(v1.zw);
+            const uint k_pair = row * LOAD_VEC_A / 4;
+            store_a(col, k_pair,     FLOAT_TYPEV2(v0.xy));
+            store_a(col, k_pair + 1, FLOAT_TYPEV2(v0.zw));
+            store_a(col, k_pair + 8, FLOAT_TYPEV2(v1.xy));
+            store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.zw));
 #elif defined(DATA_A_Q5_0)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
 
             const uint ib = idx / 8;
             const uint iqs = idx & 0x07;
@@ -98,12 +114,10 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
 
             const uint vui = uint(data_a_packed16[ib].qs[iqs]);
             const vec4 v = (vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y) - 16.0f) * d;
-
-            buf_a[buf_idx    ] = FLOAT_TYPEV2(v.xz);
-            buf_a[buf_idx + 8] = FLOAT_TYPEV2(v.yw);
+            store_a(col, row,     FLOAT_TYPEV2(v.xz));
+            store_a(col, row + 8, FLOAT_TYPEV2(v.yw));
 #elif defined(DATA_A_Q5_1)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
 
             const uint ib = idx / 4;
             const uint iqs = idx & 0x03;
@@ -119,13 +133,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const vec4 v0 = vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, ((vui >> 12) & 0xF) | qh1.y) * dm.x + dm.y;
             const vec4 v1 = vec4(((vui >> 16) & 0xF) | qh2.x, ((vui >> 20) & 0xF) | qh2.y, ((vui >> 24) & 0xF) | qh3.x, ((vui >> 28) & 0xF) | qh3.y) * dm.x + dm.y;
 
-            buf_a[buf_idx    ] = FLOAT_TYPEV2(v0.xz);
-            buf_a[buf_idx + 1] = FLOAT_TYPEV2(v1.xz);
-            buf_a[buf_idx + 8] = FLOAT_TYPEV2(v0.yw);
-            buf_a[buf_idx + 9] = FLOAT_TYPEV2(v1.yw);
+            const uint k_pair = row * LOAD_VEC_A / 4;
+            store_a(col, k_pair,     FLOAT_TYPEV2(v0.xz));
+            store_a(col, k_pair + 1, FLOAT_TYPEV2(v1.xz));
+            store_a(col, k_pair + 8, FLOAT_TYPEV2(v0.yw));
+            store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.yw));
 #elif defined(DATA_A_Q8_0)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
 
             const uint ib = idx / 8;
             const uint iqs = idx & 0x07;
@@ -135,11 +149,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const i8vec2 v1 = unpack8(int32_t(data_a_packed16[ib].qs[2*iqs + 1])).xy;
             const vec4 v = vec4(v0.x, v0.y, v1.x, v1.y) * d;
 
-            buf_a[buf_idx    ] = FLOAT_TYPEV2(v.xy);
-            buf_a[buf_idx + 1] = FLOAT_TYPEV2(v.zw);
+            const uint k_pair = row * LOAD_VEC_A / 2;
+            store_a(col, k_pair,     FLOAT_TYPEV2(v.xy));
+            store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw));
 #elif defined(DATA_A_Q1_0)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
 
             const uint ib = idx / 16;
             const uint iqs = idx & 0xfu;
@@ -147,13 +161,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const float d = float(data_a[ib].d);
             const uint bits = uint(data_a[ib].qs[iqs]);
 
-            buf_a[buf_idx    ] = FLOAT_TYPEV2((bits & 0x01u) != 0u ? d : -d, (bits & 0x02u) != 0u ? d : -d);
-            buf_a[buf_idx + 1] = FLOAT_TYPEV2((bits & 0x04u) != 0u ? d : -d, (bits & 0x08u) != 0u ? d : -d);
-            buf_a[buf_idx + 2] = FLOAT_TYPEV2((bits & 0x10u) != 0u ? d : -d, (bits & 0x20u) != 0u ? d : -d);
-            buf_a[buf_idx + 3] = FLOAT_TYPEV2((bits & 0x40u) != 0u ? d : -d, (bits & 0x80u) != 0u ? d : -d);
+            const uint k_pair = row * LOAD_VEC_A / 2;
+            store_a(col, k_pair,     FLOAT_TYPEV2((bits & 0x01u) != 0u ? d : -d, (bits & 0x02u) != 0u ? d : -d));
+            store_a(col, k_pair + 1, FLOAT_TYPEV2((bits & 0x04u) != 0u ? d : -d, (bits & 0x08u) != 0u ? d : -d));
+            store_a(col, k_pair + 2, FLOAT_TYPEV2((bits & 0x10u) != 0u ? d : -d, (bits & 0x20u) != 0u ? d : -d));
+            store_a(col, k_pair + 3, FLOAT_TYPEV2((bits & 0x40u) != 0u ? d : -d, (bits & 0x80u) != 0u ? d : -d));
 #elif defined(DATA_A_Q2_0)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
 
             const uint ib = idx / 16;
             const uint iqs = idx & 0xfu;
@@ -161,11 +175,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib].d);
             const uint bits = uint(data_a[ib].qs[iqs]);
 
-            buf_a[buf_idx    ] = d * (FLOAT_TYPEV2(bits & 3u, (bits >> 2u) & 3u) - FLOAT_TYPEV2(1.0f));
-            buf_a[buf_idx + 1] = d * (FLOAT_TYPEV2((bits >> 4u) & 3u, bits >> 6u) - FLOAT_TYPEV2(1.0f));
+            const uint k_pair = row * LOAD_VEC_A / 2;
+            store_a(col, k_pair,     d * (FLOAT_TYPEV2(bits & 3u, (bits >> 2u) & 3u) - FLOAT_TYPEV2(1.0f)));
+            store_a(col, k_pair + 1, d * (FLOAT_TYPEV2((bits >> 4u) & 3u, bits >> 6u) - FLOAT_TYPEV2(1.0f)));
 #elif defined(DATA_A_Q2_K)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
 
             const uint ib = idx / 64;                          // 4 values per idx
             const uint iqs = (idx % 64) * 2;                   // 0,2,4..126
@@ -180,11 +194,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
 
             const vec4 v = dm.x * float(scales & 0xF) * qs - dm.y * float(scales >> 4);
 
-            buf_a[buf_idx    ] = FLOAT_TYPEV2(v.xy);
-            buf_a[buf_idx + 1] = FLOAT_TYPEV2(v.zw);
+            const uint k_pair = row * LOAD_VEC_A / 2;
+            store_a(col, k_pair,     FLOAT_TYPEV2(v.xy));
+            store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw));
 #elif defined(DATA_A_TQ2_0)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
 
             const uint ib = idx / 128;                         // 2 values per idx
             const uint iqs = (idx % 128) * 2;                  // elem 0,2,4..254
@@ -197,10 +211,10 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
 
             const vec2 v = d * (vec2((qs >> shift) & 3) - 1.0);
 
-            buf_a[buf_idx] = FLOAT_TYPEV2(v.xy);
+            const uint k_pair = row * LOAD_VEC_A / 2;
+            store_a(col, k_pair, FLOAT_TYPEV2(v.xy));
 #elif defined(DATA_A_Q3_K)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
 
             const uint ib = idx / 128;                   // 2 values per idx
             const uint iqs = idx % 128;                  // 0..127
@@ -220,11 +234,10 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const vec2 qs = vec2(unpack8((uint(data_a_packed16[ib].qs[qsi / 2]) >> qsshift) & 0x0303).xy);
             const vec2 hm = vec2(unpack8(((uint(data_a_packed16[ib].hmask[hmi / 2]) >> (4 * n + halfsplit)) & 0x0101 ^ 0x0101) << 2).xy);
 
-            buf_a[buf_idx] = FLOAT_TYPEV2(dl * (qs.x - hm.x),
-                                          dl * (qs.y - hm.y));
+            store_a(col, row * LOAD_VEC_A / 2, FLOAT_TYPEV2(dl * (qs.x - hm.x),
+                                                              dl * (qs.y - hm.y)));
 #elif defined(DATA_A_Q4_K)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
 
             const uint ib = idx / 64;                  // 4 values per idx
             const uint iqs = (idx % 64) * 2;           // 0,2,4..126
@@ -256,11 +269,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
 
             const vec4 q = vec4(unpack8((data_a_packed32[ib].qs[qsi / 4] >> (b * 4)) & 0x0F0F0F0F));
 
-            buf_a[buf_idx    ] = FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m));
-            buf_a[buf_idx + 1] = FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m));
+            const uint k_pair = row * LOAD_VEC_A / 2;
+            store_a(col, k_pair,     FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m)));
+            store_a(col, k_pair + 1, FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m)));
 #elif defined(DATA_A_Q5_K)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
 
             const uint ib = idx / 64;                  // 4 values per idx
             const uint iqs = (idx % 64) * 2;           // 0,2,4..126
@@ -295,11 +308,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const uint qh = ((data_a_packed32[ib].qh[qhi / 4] >> (iqs / 16)) & 0x01010101) << 4;
             const vec4 q = vec4(unpack8(qs | qh));
 
-            buf_a[buf_idx    ] = FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m));
-            buf_a[buf_idx + 1] = FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m));
+            const uint k_pair = row * LOAD_VEC_A / 2;
+            store_a(col, k_pair,     FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m)));
+            store_a(col, k_pair + 1, FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m)));
 #elif defined(DATA_A_Q6_K)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
 
             const uint ib = idx / 128;                  // 2 values per idx
             const uint iqs = idx % 128;                 // 0..127
@@ -318,10 +331,9 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const uint qh = (uint(data_a_packed16[ib].qh[qhi]) >> qhshift) & 0x0303;
             const vec2 q = (vec2(unpack8(ql | (qh << 4)).xy) - 32) * dscale;
 
-            buf_a[buf_idx] = FLOAT_TYPEV2(q.x, q.y);
+            store_a(col, row * LOAD_VEC_A / 2, FLOAT_TYPEV2(q.x, q.y));
 #elif defined(DATA_A_IQ1_S)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
 
             const uint ib = idx / 32;                  // 8 values per idx
             const uint ib32 = (idx % 32) / 4;         // 0..7
@@ -334,13 +346,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const float delta = ((qh & 0x8000) != 0) ? -IQ1S_DELTA : IQ1S_DELTA;
             const int16_t grid = int16_t(iq1s_grid[qs | (bitfieldExtract(qh, 3 * int(ib8 & 3), 3) << 8)]);
 
+            const uint k_pair = row * LOAD_VEC_A / 2;
             [[unroll]] for (int k = 0; k < 4; ++k) {
-                buf_a[buf_idx + k] = FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k    , 2) + delta),
-                                                  dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta));
+                store_a(col, k_pair + k, FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k    , 2) + delta),
+                                                       dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta)));
             }
 #elif defined(DATA_A_IQ1_M)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
 
             const uint ib = idx / 32;  // 8 values per idx
             const uint ib8 = idx % 32;
@@ -356,13 +368,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const float delta = ((qh & 8) != 0) ? -IQ1M_DELTA : IQ1M_DELTA;
             const int16_t grid = int16_t(iq1s_grid[qs | ((qh & 7) << 8)]);
 
+            const uint k_pair = row * LOAD_VEC_A / 2;
             [[unroll]] for (int k = 0; k < 4; ++k) {
-                buf_a[buf_idx + k] = FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k    , 2) + delta),
-                                                  dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta));
+                store_a(col, k_pair + k, FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k    , 2) + delta),
+                                                       dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta)));
             }
 #elif defined(DATA_A_IQ2_XXS)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
 
             const uint ib = idx / 32;                 // 8 values per idx
             const uint ib32 = (idx % 32) / 4;         // 0..7
@@ -383,17 +395,17 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const vec4 grid0 = vec4(unpack8(grid.x));
             const vec4 grid1 = vec4(unpack8(grid.y));
 
-            buf_a[buf_idx    ] = db * FLOAT_TYPEV2((sign &   1) != 0 ? -grid0.x : grid0.x,
-                                                   (sign &   2) != 0 ? -grid0.y : grid0.y);
-            buf_a[buf_idx + 1] = db * FLOAT_TYPEV2((sign &   4) != 0 ? -grid0.z : grid0.z,
-                                                   (sign &   8) != 0 ? -grid0.w : grid0.w);
-            buf_a[buf_idx + 2] = db * FLOAT_TYPEV2((sign &  16) != 0 ? -grid1.x : grid1.x,
-                                                   (sign &  32) != 0 ? -grid1.y : grid1.y);
-            buf_a[buf_idx + 3] = db * FLOAT_TYPEV2((sign &  64) != 0 ? -grid1.z : grid1.z,
-                                                   (sign & 128) != 0 ? -grid1.w : grid1.w);
+            const uint k_pair = row * LOAD_VEC_A / 2;
+            store_a(col, k_pair,     db * FLOAT_TYPEV2((sign &   1) != 0 ? -grid0.x : grid0.x,
+                                                        (sign &   2) != 0 ? -grid0.y : grid0.y));
+            store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign &   4) != 0 ? -grid0.z : grid0.z,
+                                                        (sign &   8) != 0 ? -grid0.w : grid0.w));
+            store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign &  16) != 0 ? -grid1.x : grid1.x,
+                                                        (sign &  32) != 0 ? -grid1.y : grid1.y));
+            store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign &  64) != 0 ? -grid1.z : grid1.z,
+                                                        (sign & 128) != 0 ? -grid1.w : grid1.w));
 #elif defined(DATA_A_IQ2_XS)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
 
             const uint ib = idx / 32;            // 8 values per idx
             const uint ib32 = (idx % 32) / 4;    // 0..7
@@ -409,17 +421,17 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const vec4 grid0 = vec4(unpack8(grid.x));
             const vec4 grid1 = vec4(unpack8(grid.y));
 
-            buf_a[buf_idx    ] = db * FLOAT_TYPEV2((sign &   1) != 0 ? -grid0.x : grid0.x,
-                                                   (sign &   2) != 0 ? -grid0.y : grid0.y);
-            buf_a[buf_idx + 1] = db * FLOAT_TYPEV2((sign &   4) != 0 ? -grid0.z : grid0.z,
-                                                   (sign &   8) != 0 ? -grid0.w : grid0.w);
-            buf_a[buf_idx + 2] = db * FLOAT_TYPEV2((sign &  16) != 0 ? -grid1.x : grid1.x,
-                                                   (sign &  32) != 0 ? -grid1.y : grid1.y);
-            buf_a[buf_idx + 3] = db * FLOAT_TYPEV2((sign &  64) != 0 ? -grid1.z : grid1.z,
-                                                   (sign & 128) != 0 ? -grid1.w : grid1.w);
+            const uint k_pair = row * LOAD_VEC_A / 2;
+            store_a(col, k_pair,     db * FLOAT_TYPEV2((sign &   1) != 0 ? -grid0.x : grid0.x,
+                                                        (sign &   2) != 0 ? -grid0.y : grid0.y));
+            store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign &   4) != 0 ? -grid0.z : grid0.z,
+                                                        (sign &   8) != 0 ? -grid0.w : grid0.w));
+            store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign &  16) != 0 ? -grid1.x : grid1.x,
+                                                        (sign &  32) != 0 ? -grid1.y : grid1.y));
+            store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign &  64) != 0 ? -grid1.z : grid1.z,
+                                                        (sign & 128) != 0 ? -grid1.w : grid1.w));
 #elif defined(DATA_A_IQ2_S)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
 
             const uint ib = idx / 32;  // 8 values per idx
             const uint ib8 = idx % 32; // 0..31
@@ -437,17 +449,17 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const vec4 grid0 = vec4(unpack8(grid.x));
             const vec4 grid1 = vec4(unpack8(grid.y));
 
-            buf_a[buf_idx    ] = db * FLOAT_TYPEV2((sign &   1) != 0 ? -grid0.x : grid0.x,
-                                                   (sign &   2) != 0 ? -grid0.y : grid0.y);
-            buf_a[buf_idx + 1] = db * FLOAT_TYPEV2((sign &   4) != 0 ? -grid0.z : grid0.z,
-                                                   (sign &   8) != 0 ? -grid0.w : grid0.w);
-            buf_a[buf_idx + 2] = db * FLOAT_TYPEV2((sign &  16) != 0 ? -grid1.x : grid1.x,
-                                                   (sign &  32) != 0 ? -grid1.y : grid1.y);
-            buf_a[buf_idx + 3] = db * FLOAT_TYPEV2((sign &  64) != 0 ? -grid1.z : grid1.z,
-                                                   (sign & 128) != 0 ? -grid1.w : grid1.w);
+            const uint k_pair = row * LOAD_VEC_A / 2;
+            store_a(col, k_pair,     db * FLOAT_TYPEV2((sign &   1) != 0 ? -grid0.x : grid0.x,
+                                                        (sign &   2) != 0 ? -grid0.y : grid0.y));
+            store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign &   4) != 0 ? -grid0.z : grid0.z,
+                                                        (sign &   8) != 0 ? -grid0.w : grid0.w));
+            store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign &  16) != 0 ? -grid1.x : grid1.x,
+                                                        (sign &  32) != 0 ? -grid1.y : grid1.y));
+            store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign &  64) != 0 ? -grid1.z : grid1.z,
+                                                        (sign & 128) != 0 ? -grid1.w : grid1.w));
 #elif defined(DATA_A_IQ3_XXS)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
 
             const uint ib = idx / 64;            // 4 values per idx
             const uint iqs = idx % 64;           // 0..63
@@ -465,13 +477,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const uint grid = iq3xxs_grid[qs];
             const vec4 v = db * vec4(unpack8(grid));
 
-            buf_a[buf_idx    ] = FLOAT_TYPEV2((sign &   1) != 0 ? -v.x : v.x,
-                                              (sign &   2) != 0 ? -v.y : v.y);
-            buf_a[buf_idx + 1] = FLOAT_TYPEV2((sign &   4) != 0 ? -v.z : v.z,
-                                              (sign &   8) != 0 ? -v.w : v.w);
+            const uint k_pair = row * LOAD_VEC_A / 2;
+            store_a(col, k_pair,     FLOAT_TYPEV2((sign &   1) != 0 ? -v.x : v.x,
+                                                   (sign &   2) != 0 ? -v.y : v.y));
+            store_a(col, k_pair + 1, FLOAT_TYPEV2((sign &   4) != 0 ? -v.z : v.z,
+                                                   (sign &   8) != 0 ? -v.w : v.w));
 #elif defined(DATA_A_IQ3_S)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
 
             const uint ib = idx / 64;            // 4 values per idx
             const uint iqs = idx % 64;           // 0..63
@@ -487,13 +499,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const uint32_t grid = iq3s_grid[qs | ((qh << (8 - (iqs % 8))) & 256)];
             const vec4 v = db * vec4(unpack8(grid));
 
-            buf_a[buf_idx    ] = FLOAT_TYPEV2((sign &   1) != 0 ? -v.x : v.x,
-                                              (sign &   2) != 0 ? -v.y : v.y);
-            buf_a[buf_idx + 1] = FLOAT_TYPEV2((sign &   4) != 0 ? -v.z : v.z,
-                                              (sign &   8) != 0 ? -v.w : v.w);
+            const uint k_pair = row * LOAD_VEC_A / 2;
+            store_a(col, k_pair,     FLOAT_TYPEV2((sign &   1) != 0 ? -v.x : v.x,
+                                                   (sign &   2) != 0 ? -v.y : v.y));
+            store_a(col, k_pair + 1, FLOAT_TYPEV2((sign &   4) != 0 ? -v.z : v.z,
+                                                   (sign &   8) != 0 ? -v.w : v.w));
 #elif defined(DATA_A_IQ4_XS)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
 
             const uint ib = idx / 64;            // 4 values per idx
             const uint ib32 = (idx % 64) / 8;    // 0..7
@@ -507,11 +519,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const float d = float(data_a[ib].d);
             const vec4 v = d * float(int(sl | (sh << 4)) - 32) * vec4(kvalues_iq4nl[qs.x], kvalues_iq4nl[qs.y], kvalues_iq4nl[qs.z], kvalues_iq4nl[qs.w]);
 
-            buf_a[buf_idx    ] = FLOAT_TYPEV2(v.xy);
-            buf_a[buf_idx + 1] = FLOAT_TYPEV2(v.zw);
+            const uint k_pair = row * LOAD_VEC_A / 2;
+            store_a(col, k_pair,     FLOAT_TYPEV2(v.xy));
+            store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw));
 #elif defined(DATA_A_IQ4_NL)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
 
             const uint ib = idx / 8;
             const uint iqs = idx & 0x07;
@@ -519,13 +531,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             const FLOAT_TYPE d = FLOAT_TYPE(data_a_packed16[ib].d);
             const uint vui = uint(data_a_packed16[ib].qs[iqs]);
 
-            buf_a[buf_idx    ] = d * FLOAT_TYPEV2(kvalues_iq4nl[vui & 0xF],
-                                                  kvalues_iq4nl[bitfieldExtract(vui, 8, 4)]);
-            buf_a[buf_idx + 8] = d * FLOAT_TYPEV2(kvalues_iq4nl[bitfieldExtract(vui, 4, 4)],
-                                                  kvalues_iq4nl[vui >> 12]);
+            const uint k_pair = row * LOAD_VEC_A / 4;
+            store_a(col, k_pair,     d * FLOAT_TYPEV2(kvalues_iq4nl[vui & 0xF],
+                                                       kvalues_iq4nl[bitfieldExtract(vui, 8, 4)]));
+            store_a(col, k_pair + 8, d * FLOAT_TYPEV2(kvalues_iq4nl[bitfieldExtract(vui, 4, 4)],
+                                                       kvalues_iq4nl[vui >> 12]));
 #elif defined(DATA_A_MXFP4)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
 
             const uint ib = idx / 8;
             const uint iqs = (idx & 0x07) * 2;
@@ -536,38 +548,37 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
 #ifdef USE_OCP_FP4
             const float d = e8m0_to_fp32(data_a[ib].e);
             const u8vec2 packed = u8vec2(vui, vui2);
-            buf_a[buf_idx    ] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * FLOAT_TYPE(d);
-            buf_a[buf_idx + 8] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * FLOAT_TYPE(d);
+            store_a(col, row,     FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * FLOAT_TYPE(d));
+            store_a(col, row + 8, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * FLOAT_TYPE(d));
 #else
             const float d = e8m0_to_fp32(data_a[ib].e) * 0.5;
-            buf_a[buf_idx    ] = FLOAT_TYPEV2(kvalues_mxfp4[vui  & 0xF] * d,
-                                              kvalues_mxfp4[vui2 & 0xF] * d);
-            buf_a[buf_idx + 8] = FLOAT_TYPEV2(kvalues_mxfp4[vui  >>  4] * d,
-                                              kvalues_mxfp4[vui2 >>  4] * d);
+            store_a(col, row,     FLOAT_TYPEV2(kvalues_mxfp4[vui  & 0xF] * d,
+                                              kvalues_mxfp4[vui2 & 0xF] * d));
+            store_a(col, row + 8, FLOAT_TYPEV2(kvalues_mxfp4[vui  >>  4] * d,
+                                              kvalues_mxfp4[vui2 >>  4] * d));
 #endif
 #elif defined(DATA_A_NVFP4)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
-            // lo and hi nibbles are 8 elements apart, which doesn't quite line up with
-            // how the thread mapping and buf_idx calculation works for other types.
-            const uint buf_idx = col * SHMEM_STRIDE + (row & 3) + (row & ~3) * 2;
-
             const uint ib = idx / 16u;
             const uint sub = (idx & 0xC) >> 2;
             const uint iqs = (idx & 0xF) * 2;
             const uint vui = uint(data_a[ib].qs[iqs]);
             const uint vui2 = uint(data_a[ib].qs[iqs+1]);
 
+            // lo and hi nibbles are 8 elements apart, which doesn't quite line up with
+            // how the thread mapping and buf_idx calculation works for other types.
+            const uint eff_row = (row & 3) + (row & ~3) * 2;
 #ifdef USE_OCP_FP4
             const FLOAT_TYPE d = FLOAT_TYPE(ue4m3_from_bits(data_a[ib].d[sub]));
             const u8vec2 packed = u8vec2(vui, vui2);
-            buf_a[buf_idx    ] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * d;
-            buf_a[buf_idx + 4] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * d;
+            store_a(col, eff_row,     FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * d);
+            store_a(col, eff_row + 4, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * d);
 #else
             const float d = ue4m3_to_fp32(data_a[ib].d[sub]) * 0.5;
-            buf_a[buf_idx    ] = FLOAT_TYPEV2(kvalues_mxfp4[vui  & 0xF] * d,
-                                              kvalues_mxfp4[vui2 & 0xF] * d);
-            buf_a[buf_idx + 4] = FLOAT_TYPEV2(kvalues_mxfp4[vui  >>  4] * d,
-                                              kvalues_mxfp4[vui2 >>  4] * d);
+            store_a(col, eff_row,     FLOAT_TYPEV2(kvalues_mxfp4[vui  & 0xF] * d,
+                                                    kvalues_mxfp4[vui2 & 0xF] * d));
+            store_a(col, eff_row + 4, FLOAT_TYPEV2(kvalues_mxfp4[vui  >>  4] * d,
+                                                    kvalues_mxfp4[vui2 >>  4] * d));
 #endif
 #endif
 }

ggml/src/ggml-vulkan/vulkan-shaders/ssm_scan.comp

@@ -33,6 +33,8 @@ layout(push_constant) uniform PushConstants {
     uint d_head;
     uint n_group;
     uint n_tok;
+    uint n_seq;
+    uint K;
 };
 
 float softplus(float x) {
@@ -114,6 +116,14 @@ void main() {
         if (lane == 0) {
             d[y_base_idx + i * stride_y] = state_sum;
         }
+
+        const uint slot = n_tok - 1u - i;
+        if (slot > 0u && slot < K) {
+            const uint snapshot_base_idx = s_base_idx + slot * n_seq * (nb03 / 4u);
+            [[unroll]] for (uint j = 0; j < c_factor; j++) {
+                d[snapshot_base_idx + SUBGROUP_SIZE * j + lane] = state[j];
+            }
+        }
     }
 
     // write back the state

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

@@ -1327,6 +1327,7 @@ static webgpu_encoded_op ggml_webgpu_ssm_scan(webgpu_context & ctx,
         (uint32_t) src4->ne[1],
         (uint32_t) src1->ne[2],
         (uint32_t) ggml_nelements(src1),
+        (uint32_t) ggml_get_op_params_i32(dst, 0),
     };
 
     std::vector<wgpu::BindGroupEntry> entries = {

ggml/src/ggml-webgpu/wgsl-shaders/ssm_scan.wgsl

@@ -41,6 +41,7 @@ struct Params {
     n_seq_tokens: u32,
 
     y_elems: u32,
+    K: u32,
 };
 
 @group(0) @binding(0) var<storage, read_write> s_in: array<f32>;
@@ -123,6 +124,7 @@ fn main(
     let head_seq = wg_linear / params.d_inner;
     let ir = head_seq % params.n_head;
     let i3 = head_seq / params.n_head;
+    let n_seqs = params.y_elems / (params.n_seq_tokens * params.n_head * params.d_inner);
 
     let state_slot = read_state_slot(i3);
     let g = ir / (params.n_head / params.n_group);
@@ -179,6 +181,15 @@ fn main(
 #endif
             s_prev = s;
 
+            let slot = params.n_seq_tokens - 1u - token;
+            if (slot > 0u && slot < params.K) {
+                let snapshot_idx =
+                    params.offset_dst + params.y_elems + tid + i1 * params.d_state +
+                    ir * (params.d_state * params.d_inner) +
+                    (slot * n_seqs + i3) * (params.d_state * params.d_inner * params.n_head);
+                dst[snapshot_idx] = s;
+            }
+
 #ifdef USE_SUBGROUP_REDUCTION
 #ifdef XBC_OVERLAP
             let subgroup_partial = subgroupAdd(s * read_merged_f32(c_idx));

ggml/src/ggml.c

@@ -5588,7 +5588,10 @@ struct ggml_tensor * ggml_ssm_scan(
         struct ggml_tensor  * A,
         struct ggml_tensor  * B,
         struct ggml_tensor  * C,
-        struct ggml_tensor  * ids) {
+        struct ggml_tensor  * ids,
+        int64_t               K) {
+    GGML_ASSERT(K >= 1);
+    GGML_ASSERT(K <= INT32_MAX);
     GGML_ASSERT(ggml_is_contiguous(s));
     GGML_ASSERT(ggml_is_contiguous(dt));
     GGML_ASSERT(ggml_is_contiguous(A));
@@ -5625,11 +5628,12 @@ struct ggml_tensor * ggml_ssm_scan(
         if (A->ne[0] != 1) {
             // Mamba-1 has more granular decay factors
             GGML_ASSERT(A->ne[0] == d_state);
+            GGML_ASSERT(K == 1);
         }
     }
 
     // concatenated y + ssm_states
-    struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ggml_nelements(x) + s->ne[0]*s->ne[1]*s->ne[2]*ids->ne[0]);
+    struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ggml_nelements(x) + K*s->ne[0]*s->ne[1]*s->ne[2]*ids->ne[0]);
 
     result->op   = GGML_OP_SSM_SCAN;
     result->src[0] = s;
@@ -5640,6 +5644,8 @@ struct ggml_tensor * ggml_ssm_scan(
     result->src[5] = C;
     result->src[6] = ids;
 
+    ggml_set_op_params_i32(result, 0, (int32_t) K);
+
     return result;
 }
 

gguf-py/gguf/constants.py

@@ -124,6 +124,7 @@ class Keys:
         EXPERT_WEIGHTS_NORM               = "{arch}.expert_weights_norm"
         EXPERT_GATING_FUNC                = "{arch}.expert_gating_func"
         EXPERT_GROUP_SCALE                = "{arch}.expert_group_scale"
+        EXPERT_LATENT_LENGTH              = "{arch}.expert_latent_length"
         EXPERTS_PER_GROUP                 = "{arch}.experts_per_group"
         MOE_EVERY_N_LAYERS                = "{arch}.moe_every_n_layers"
         MOE_LATENT_SIZE                   = "{arch}.moe_latent_size"
@@ -238,6 +239,13 @@ class Keys:
         SCALING_YARN_BETA_FAST    = "{arch}.rope.scaling.yarn_beta_fast"
         SCALING_YARN_BETA_SLOW    = "{arch}.rope.scaling.yarn_beta_slow"
 
+    class Activation:
+        SITU_BETA        = "{arch}.activation.situ_beta"
+        SITU_LINEAR_BETA = "{arch}.activation.situ_linear_beta"
+
+    class AttnRes:
+        BLOCK_SIZE = "{arch}.attn_res.block_size"
+
     class Split:
         LLM_KV_SPLIT_NO            = "split.no"
         LLM_KV_SPLIT_COUNT         = "split.count"
@@ -252,7 +260,8 @@ class Keys:
         DT_B_C_RMS     = "{arch}.ssm.dt_b_c_rms"
 
     class KDA:
-        HEAD_DIM = "{arch}.kda.head_dim"
+        HEAD_DIM         = "{arch}.kda.head_dim"
+        GATE_LOWER_BOUND = "{arch}.kda.gate_lower_bound"
 
     class WKV:
         HEAD_SIZE = "{arch}.wkv.head_size"
@@ -565,6 +574,7 @@ class MODEL_ARCH(IntEnum):
     GROVEMOE         = auto()
     APERTUS          = auto()
     COGVLM           = auto()
+    MINIMAX01        = auto()
     MINIMAXM2        = auto()
     MINIMAXM3        = auto()
     RND1             = auto()
@@ -579,6 +589,7 @@ class MODEL_ARCH(IntEnum):
     LLAMA_EMBED      = auto()
     MAINCODER        = auto()
     KIMI_LINEAR      = auto()
+    KIMI_K3          = auto()
     TALKIE           = auto()
     MELLUM           = auto()
     NANBEIGE         = auto()
@@ -697,6 +708,13 @@ class MODEL_TENSOR(IntEnum):
     SSM_BETA             = auto() # Kimi Linear qwen3.5
     SSM_G_A              = auto() # Kimi Linear
     SSM_G_B              = auto() # Kimi Linear
+    SSM_G                = auto() # Kimi K3 (full-rank KDA gate, replaces SSM_G_A/SSM_G_B)
+    ATTN_RES_SCORE       = auto() # Kimi K3 (fused res_norm * res_proj, pre-attention)
+    FFN_RES_SCORE        = auto() # Kimi K3 (fused res_norm * res_proj, pre-FFN)
+    OUTPUT_RES_SCORE     = auto() # Kimi K3 (fused res_norm * res_proj, final)
+    FFN_ROUTED_DOWN      = auto() # Kimi K3 (latent MoE: hidden -> latent)
+    FFN_ROUTED_UP        = auto() # Kimi K3 (latent MoE: latent -> hidden)
+    FFN_ROUTED_NORM      = auto() # Kimi K3 (latent MoE: norm on expert output)
     TIME_MIX_W0          = auto()
     TIME_MIX_W1          = auto()
     TIME_MIX_W2          = auto()
@@ -1271,6 +1289,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
     MODEL_ARCH.SEED_OSS:         "seed_oss",
     MODEL_ARCH.GROVEMOE:         "grovemoe",
     MODEL_ARCH.APERTUS:          "apertus",
+    MODEL_ARCH.MINIMAX01:        "minimax-01",
     MODEL_ARCH.MINIMAXM2:        "minimax-m2",
     MODEL_ARCH.MINIMAXM3:        "minimax-m3",
     MODEL_ARCH.COGVLM:           "cogvlm",
@@ -1286,6 +1305,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
     MODEL_ARCH.LLAMA_EMBED:      "llama-embed",
     MODEL_ARCH.MAINCODER:        "maincoder",
     MODEL_ARCH.KIMI_LINEAR:      "kimi-linear",
+    MODEL_ARCH.KIMI_K3:          "kimi-k3",
     MODEL_ARCH.TALKIE:           "talkie",
     MODEL_ARCH.MELLUM:           "mellum",
     MODEL_ARCH.NANBEIGE:         "nanbeige",
@@ -1402,6 +1422,13 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
     MODEL_TENSOR.SSM_BETA:                  "blk.{bid}.ssm_beta",             # Kimi Linear qwen3.5
     MODEL_TENSOR.SSM_G_A:                   "blk.{bid}.ssm_g_a",              # Kimi Linear
     MODEL_TENSOR.SSM_G_B:                   "blk.{bid}.ssm_g_b",              # Kimi Linear
+    MODEL_TENSOR.SSM_G:                     "blk.{bid}.ssm_g",                # Kimi K3
+    MODEL_TENSOR.ATTN_RES_SCORE:            "blk.{bid}.attn_res_score",       # Kimi K3
+    MODEL_TENSOR.FFN_RES_SCORE:             "blk.{bid}.ffn_res_score",        # Kimi K3
+    MODEL_TENSOR.OUTPUT_RES_SCORE:          "output_res_score",               # Kimi K3
+    MODEL_TENSOR.FFN_ROUTED_DOWN:           "blk.{bid}.ffn_routed_down",      # Kimi K3
+    MODEL_TENSOR.FFN_ROUTED_UP:             "blk.{bid}.ffn_routed_up",        # Kimi K3
+    MODEL_TENSOR.FFN_ROUTED_NORM:           "blk.{bid}.ffn_routed_norm",      # Kimi K3
     MODEL_TENSOR.TIME_MIX_W0:               "blk.{bid}.time_mix_w0",
     MODEL_TENSOR.TIME_MIX_W1:               "blk.{bid}.time_mix_w1",
     MODEL_TENSOR.TIME_MIX_W2:               "blk.{bid}.time_mix_w2",
@@ -4592,6 +4619,24 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.FFN_DOWN_CHEXP,
         MODEL_TENSOR.FFN_UP_CHEXP,
     ],
+    MODEL_ARCH.MINIMAX01: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_NORM_2,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+    ],
     MODEL_ARCH.MINIMAXM2: [
         MODEL_TENSOR.TOKEN_EMBD,
         MODEL_TENSOR.OUTPUT_NORM,
@@ -4940,6 +4985,56 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.FFN_DOWN_SHEXP,
         MODEL_TENSOR.FFN_UP_SHEXP,
     ],
+    MODEL_ARCH.KIMI_K3: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.OUTPUT_RES_SCORE,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_RES_SCORE,
+        MODEL_TENSOR.FFN_RES_SCORE,
+        # MLA (full-attention layers)
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_KV_A_MQA,
+        MODEL_TENSOR.ATTN_KV_B,
+        MODEL_TENSOR.ATTN_K_B,
+        MODEL_TENSOR.ATTN_V_B,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_KV_A_NORM,
+        # KDA (linear-attention layers)
+        MODEL_TENSOR.SSM_CONV1D_Q,
+        MODEL_TENSOR.SSM_CONV1D_K,
+        MODEL_TENSOR.SSM_CONV1D_V,
+        MODEL_TENSOR.SSM_F_A,
+        MODEL_TENSOR.SSM_F_B,
+        MODEL_TENSOR.SSM_BETA,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_G,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_NORM,
+        # FFN
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_ROUTED_DOWN,
+        MODEL_TENSOR.FFN_ROUTED_UP,
+        MODEL_TENSOR.FFN_ROUTED_NORM,
+    ],
     MODEL_ARCH.TALKIE: [
         MODEL_TENSOR.TOKEN_EMBD,
         MODEL_TENSOR.OUTPUT,

gguf-py/gguf/gguf_writer.py

@@ -1103,6 +1103,21 @@ class GGUFWriter:
     def add_ssm_dt_b_c_rms(self, value: bool) -> None:
         self.add_bool(Keys.SSM.DT_B_C_RMS.format(arch=self.arch), value)
 
+    def add_kda_gate_lower_bound(self, value: float) -> None:
+        self.add_float32(Keys.KDA.GATE_LOWER_BOUND.format(arch=self.arch), value)
+
+    def add_expert_latent_length(self, value: int) -> None:
+        self.add_uint32(Keys.LLM.EXPERT_LATENT_LENGTH.format(arch=self.arch), value)
+
+    def add_activation_situ_beta(self, value: float) -> None:
+        self.add_float32(Keys.Activation.SITU_BETA.format(arch=self.arch), value)
+
+    def add_activation_situ_linear_beta(self, value: float) -> None:
+        self.add_float32(Keys.Activation.SITU_LINEAR_BETA.format(arch=self.arch), value)
+
+    def add_attn_res_block_size(self, value: int) -> None:
+        self.add_uint32(Keys.AttnRes.BLOCK_SIZE.format(arch=self.arch), value)
+
     def add_kda_head_dim(self, value: int) -> None:
         self.add_uint32(Keys.KDA.HEAD_DIM.format(arch=self.arch), value)
 

gguf-py/gguf/tensor_mapping.py

@@ -225,6 +225,7 @@ class TensorNameMap:
             "rwkv.blocks.{bid}.ln2",                        # rwkv6
             "model.layers.{bid}.ln2",                       # rwkv7
             "model.layers.{bid}.post_attention_layernorm",  # cogvlm
+            "model.layers.{bid}.self_attn.norm",            # minimax-01
         ),
 
         # Attention query-key-value
@@ -321,7 +322,7 @@ class TensorNameMap:
             "h.{bid}.self_attention.dense",                                 # bloom
             "model.layers.{bid}.self_attn.o_proj",                          # llama-hf nemotron olmoe olmo2 phimoe
             "layers.{bid}.self_attn.o_proj",                                # embeddinggemma
-            "model.layers.{bid}.self_attn.out_proj",                        # lfm2
+            "model.layers.{bid}.self_attn.out_proj",                        # lfm2 minimax-01
             "model.layers.{bid}.self_attn.linear_attn",                     # deci
             "layers.{bid}.attention.wo",                                    # llama-pth
             "encoder.layer.{bid}.attention.output.dense",                   # bert
@@ -385,6 +386,7 @@ class TensorNameMap:
             "model.layers.{bid}.self_attn.gate_proj", # afmoe muse-glimmer
             "model.layers.{bid}.linear_attn.in_proj_z",  # qwen3.5
             "model.layers.{bid}.self_attn.g_proj",    # step3.5 head-wise attention gate
+            "model.layers.{bid}.self_attn.output_gate",  # minimax-01
         ),
 
         # Feed-forward norm
@@ -910,6 +912,19 @@ class TensorNameMap:
             "model.layers.{bid}.linear_attn.in_proj_b",  # qwen3.5
             "model.layers.{bid}.self_attn.b_proj",       # Kimi Linear
         ),
+        # Kimi K3 latent MoE: routed experts operate in a down-projected space
+        MODEL_TENSOR.FFN_ROUTED_DOWN: (
+            "model.layers.{bid}.block_sparse_moe.routed_expert_down_proj",
+        ),
+
+        MODEL_TENSOR.FFN_ROUTED_UP: (
+            "model.layers.{bid}.block_sparse_moe.routed_expert_up_proj",
+        ),
+
+        MODEL_TENSOR.FFN_ROUTED_NORM: (
+            "model.layers.{bid}.block_sparse_moe.routed_expert_norm",
+        ),
+
         MODEL_TENSOR.SSM_G_A: (
             "model.layers.{bid}.self_attn.g_a_proj",
         ),

models/templates/Kimi-K3.jinja

@@ -0,0 +1,324 @@
+{%- macro escape_attr(value) -%}
+{{- value|string|replace('&', '&')|replace('"', '"') -}}
+{%- endmacro -%}
+
+{%- macro open_tag(tag, attrs=[]) -%}
+{{- '<|open|>' + tag -}}
+{%- for attr in attrs -%}
+{{- ' ' + attr[0] + '="' -}}{{- escape_attr(attr[1]) -}}{{- '"' -}}
+{%- endfor -%}
+{{- '<|sep|>' -}}
+{%- endmacro -%}
+
+{%- macro close_tag(tag) -%}
+{{- '<|close|>' + tag + '<|sep|>' -}}
+{%- endmacro -%}
+
+{%- macro next_image(state) -%}
+{%- if image_prompts is defined and image_prompts is not none -%}
+    {%- if state.image_index >= image_prompts|length -%}
+        {{- raise_exception('More image placeholders than image prompts.') -}}
+    {%- endif -%}
+    {{- image_prompts[state.image_index] -}}
+    {%- set state.image_index = state.image_index + 1 -%}
+{%- else -%}
+    {{- '<|kimi_image_placeholder|>' -}}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_text(text, state) -%}
+{%- set text = text|string -%}
+{%- if image_prompts is defined and image_prompts is not none and '<|kimi_image_placeholder|>' in text -%}
+    {%- set parts = text.split('<|kimi_image_placeholder|>') -%}
+    {%- for part in parts -%}
+        {{- part -}}
+        {%- if not loop.last -%}{{- next_image(state) -}}{%- endif -%}
+    {%- endfor -%}
+{%- else -%}
+    {{- text -}}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_content(content, state) -%}
+{%- if content is string -%}
+    {{- render_text(content, state) -}}
+{%- elif content is not none and content is defined -%}
+    {%- for part in content -%}
+        {%- if part.type in ['image', 'image_url'] -%}
+            {{- next_image(state) -}}
+        {%- else -%}
+            {{- render_text(part.text, state) -}}
+        {%- endif -%}
+    {%- endfor -%}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro internal_system_message(message_type, body) -%}
+{{- open_tag('message', [('role', 'system'), ('type', message_type)]) -}}
+{{- body|trim -}}
+{{- close_tag('message') -}}
+{{- '<|end_of_msg|>' -}}
+{%- endmacro -%}
+
+{%- macro json_sorted(value) -%}
+{#- tojson has no sort_keys, so sort each mapping level with dictsort to match the
+    reference implementation. Array order is kept as-is. -#}
+{%- if value is mapping -%}
+{{- '{' -}}
+{%- for key, item in value|dictsort -%}
+{%- if not loop.first -%}{{- ',' -}}{%- endif -%}
+{{- key|tojson(ensure_ascii=false) -}}{{- ':' -}}{{- json_sorted(item) -}}
+{%- endfor -%}
+{{- '}' -}}
+{%- elif value is string or value is number or value is boolean or value is none -%}
+{{- value|tojson(ensure_ascii=false) -}}
+{%- else -%}
+{{- '[' -}}
+{%- for item in value -%}
+{%- if not loop.first -%}{{- ',' -}}{%- endif -%}
+{{- json_sorted(item) -}}
+{%- endfor -%}
+{{- ']' -}}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_tool_declare(tool_list, dynamic=false) -%}
+{{- open_tag('message', [('role', 'system'), ('type', 'tool-declare')]) -}}
+{%- if dynamic -%}
+{{- '## New Tools Available\nThe system dynamically extends the toolset via lazy-loading.\nYou have access to all existing and extended tools.\nHere are the specs for the extended tools.\n\n```json\n' -}}
+{%- else -%}
+{{- '# Tools\nHere are the available tools, described in JSONSchema.\n\n```json\n' -}}
+{%- endif -%}
+{{- json_sorted(tool_list) -}}
+{{- '\n```' -}}
+{{- close_tag('message') -}}
+{{- '<|end_of_msg|>' -}}
+{%- endmacro -%}
+
+{%- macro xtml_type(value) -%}
+{%- if value is boolean -%}boolean
+{%- elif value is none -%}null
+{%- elif value is number -%}number
+{%- elif value is string -%}string
+{%- elif value is mapping -%}object
+{%- else -%}array
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro xtml_value(value) -%}
+{%- if value is string -%}
+{{- value -}}
+{%- else -%}
+{{- value|tojson(ensure_ascii=false) -}}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_assistant(message, state) -%}
+{%- if thinking -%}
+    {%- set reasoning_content = message.get('reasoning_content') or message.get('reasoning') -%}
+    {{- open_tag('think') -}}
+    {%- if reasoning_content is not none and reasoning_content|string|trim -%}
+        {{- render_text(reasoning_content, state) -}}
+    {%- endif -%}
+    {{- close_tag('think') -}}
+{%- endif -%}
+{{- open_tag('response') -}}
+{{- render_content(message.get('content'), state) -}}
+{{- close_tag('response') -}}
+{%- set tool_calls = message.get('tool_calls') -%}
+{%- if tool_calls -%}
+    {{- open_tag('tools') -}}
+    {%- for tool_call in tool_calls -%}
+        {%- if tool_call is not mapping -%}
+            {{- raise_exception('Kimi K3 tool calls must be mappings.') -}}
+        {%- endif -%}
+        {%- set fn = tool_call.function if tool_call.function is defined and tool_call.function is mapping else tool_call -%}
+        {%- if fn.get('name') is none -%}
+            {{- raise_exception('Kimi K3 tool calls require a function name.') -}}
+        {%- endif -%}
+        {{- open_tag('call', [('tool', fn.name), ('index', loop.index)]) -}}
+        {%- set arguments = fn.get('arguments', {}) -%}
+        {%- set json_block = fn.get('_xtml_json_block') -%}
+        {%- if json_block is not none -%}
+            {{- open_tag('json', [('type', 'object')]) -}}
+            {{- render_text(json_block, state) -}}
+            {{- close_tag('json') -}}
+        {%- elif arguments is mapping -%}
+            {%- for key, value in arguments.items() -%}
+                {{- open_tag('argument', [('key', key), ('type', xtml_type(value))]) -}}
+                {{- render_text(xtml_value(value), state) -}}
+                {{- close_tag('argument') -}}
+            {%- endfor -%}
+        {%- elif arguments is string and arguments|trim -%}
+            {{- open_tag('json', [('type', 'object')]) -}}
+            {{- render_text(arguments, state) -}}
+            {{- close_tag('json') -}}
+        {%- elif arguments is not none and arguments is not string -%}
+            {{- raise_exception('Kimi K3 tool call arguments must be a mapping or a JSON object string.') -}}
+        {%- endif -%}
+        {{- close_tag('call') -}}
+    {%- endfor -%}
+    {{- close_tag('tools') -}}
+{%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_tool_message(message, state, resolved_name=none) -%}
+{%- set state.tool_index = state.tool_index + 1 -%}
+{%- if resolved_name is not none -%}
+    {%- set tool_name = resolved_name -%}
+{%- elif 'tool' in message -%}
+    {%- set tool_name = message.get('tool') -%}
+{%- else -%}
+    {%- set tool_name = message.get('name') -%}
+{%- endif -%}
+{%- if tool_name is none and state.tool_calls is not none and state.tool_index <= state.tool_calls|length -%}
+    {%- set fallback_call = state.tool_calls[state.tool_index - 1] -%}
+    {%- set fallback_fn = fallback_call.function if fallback_call.function is defined and fallback_call.function is mapping else fallback_call -%}
+    {%- set tool_name = fallback_fn.name -%}
+{%- endif -%}
+{%- if tool_name is none -%}
+    {{- raise_exception('Kimi K3 tool messages need a resolvable tool name: carry `tool`/`name`, or match a preceding assistant tool_call by order.') -}}
+{%- endif -%}
+{{- open_tag('message', [('role', 'tool'), ('tool', tool_name), ('index', state.tool_index)]) -}}
+{{- render_content(message.get('content'), state) -}}
+{{- close_tag('message') -}}
+{{- '<|end_of_msg|>' -}}
+{%- endmacro -%}
+
+{%- if thinking is undefined -%}
+    {%- set thinking = true -%}
+{%- endif -%}
+{%- if thinking_effort is undefined -%}
+    {%- set thinking_effort = 'max' -%}
+{%- endif -%}
+{%- if thinking and thinking_effort is not none and thinking_effort not in ['low', 'high', 'max'] -%}
+    {{- raise_exception('Unsupported thinking_effort=' + thinking_effort|string + '; supported values are low, high, and max.') -}}
+{%- endif -%}
+
+{%- set state = namespace(image_index=0, tool_calls=none, tool_index=0, response_schema=none) -%}
+
+{%- if tools is defined and tools -%}
+    {{- render_tool_declare(tools) -}}
+{%- endif -%}
+
+{%- if thinking and thinking_effort in ['low', 'high', 'max'] -%}
+    {{- internal_system_message(
+        'thinking-effort',
+        '`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`.\nNow the system is invoked with `thinking_effort=' + thinking_effort|string + '`.'
+    ) -}}
+{%- endif -%}
+
+{%- for message in messages -%}
+    {%- if message is mapping -%}
+        {%- if 'role' not in message -%}
+            {{- raise_exception('Kimi K3 messages require a role.') -}}
+        {%- elif message.role == 'user' -%}
+            {%- set attrs = [('role', 'user')] -%}
+            {%- if message.get('name') -%}{%- set attrs = attrs + [('name', message.name)] -%}{%- endif -%}
+            {{- open_tag('message', attrs) -}}
+            {{- render_content(message.get('content'), state) -}}
+            {{- close_tag('message') -}}
+            {{- '<|end_of_msg|>' -}}
+        {%- elif message.role == 'system' and message.get('tools') -%}
+            {{- render_tool_declare(message.tools, dynamic=true) -}}
+        {%- elif message.role == 'system' -%}
+            {%- set attrs = [('role', 'system')] -%}
+            {%- if message.get('name') -%}{%- set attrs = attrs + [('name', message.name)] -%}{%- endif -%}
+            {{- open_tag('message', attrs) -}}
+            {{- render_content(message.get('content'), state) -}}
+            {{- close_tag('message') -}}
+            {{- '<|end_of_msg|>' -}}
+        {%- elif message.role == 'assistant' -%}
+            {%- set state.tool_calls = message.get('tool_calls') -%}
+            {%- set state.tool_index = 0 -%}
+            {%- set attrs = [('role', 'assistant')] -%}
+            {%- if message.get('name') -%}{%- set attrs = attrs + [('name', message.name)] -%}{%- endif -%}
+            {{- open_tag('message', attrs) -}}
+            {{- render_assistant(message, state) -}}
+            {{- close_tag('message') -}}
+            {{- '<|end_of_msg|>' -}}
+        {%- elif message.role == 'tool' and (loop.first or messages[loop.index0 - 1].role != 'tool') -%}
+            {%- set run = namespace(tool_messages=[], resolved_count=0) -%}
+            {%- for candidate in messages[loop.index0:] -%}
+                {%- if candidate is not mapping or candidate.role != 'tool' -%}{%- break -%}{%- endif -%}
+                {%- set run.tool_messages = run.tool_messages + [candidate] -%}
+                {%- set call_id = candidate.get('tool_call_id', candidate.get('id')) -%}
+                {%- set match = namespace(found=false) -%}
+                {%- if call_id is not none and state.tool_calls is not none -%}
+                    {%- for tool_call in state.tool_calls -%}
+                        {%- if not match.found and tool_call is mapping and tool_call.get('id') is not none and tool_call.get('id')|string == call_id|string -%}
+                            {%- set match.found = true -%}
+                        {%- endif -%}
+                    {%- endfor -%}
+                {%- endif -%}
+                {%- if match.found -%}{%- set run.resolved_count = run.resolved_count + 1 -%}{%- endif -%}
+            {%- endfor -%}
+            {%- if run.tool_messages|length > 0 and run.resolved_count == run.tool_messages|length -%}
+                {%- set emitted = namespace(ids=[]) -%}
+                {%- for tool_call in state.tool_calls -%}
+                    {%- if tool_call is mapping and tool_call.get('id') is not none and tool_call.get('id')|string not in emitted.ids -%}
+                        {%- set emitted.ids = emitted.ids + [tool_call.get('id')|string] -%}
+                        {%- set fn = tool_call.function if tool_call.function is defined and tool_call.function is mapping else tool_call -%}
+                        {%- for tool_message in run.tool_messages -%}
+                            {%- set result_id = tool_message.get('tool_call_id', tool_message.get('id')) -%}
+                            {%- if result_id is not none and result_id|string == tool_call.get('id')|string -%}
+                                {{- render_tool_message(tool_message, state, fn.get('name')) -}}
+                            {%- endif -%}
+                        {%- endfor -%}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- else -%}
+                {%- for tool_message in run.tool_messages -%}
+                    {{- render_tool_message(tool_message, state) -}}
+                {%- endfor -%}
+            {%- endif -%}
+        {%- endif -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{%- if tool_choice is defined and tool_choice == 'required' -%}
+    {{- internal_system_message('tool-choice', 'The system is invoked with `tool_choice=required`.\nYou MUST call tools in the next message.') -}}
+{%- elif tool_choice is defined and tool_choice == 'none' -%}
+    {{- internal_system_message('tool-choice', 'The system is invoked with `tool_choice=none`.\nYou MUST NOT call any tools in the next message.') -}}
+{%- endif -%}
+
+{%- if response_schema is defined -%}
+    {%- set state.response_schema = response_schema -%}
+{%- elif response_format is defined and response_format is mapping and response_format.get('json_schema') is not none -%}
+    {%- set schema_wrapper = response_format.get('json_schema') -%}
+    {%- if schema_wrapper is mapping and 'schema' in schema_wrapper -%}
+        {%- set state.response_schema = schema_wrapper.get('schema') -%}
+    {%- elif schema_wrapper is mapping and 'json_schema' in schema_wrapper -%}
+        {%- set state.response_schema = schema_wrapper.get('json_schema') -%}
+    {%- else -%}
+        {%- set state.response_schema = schema_wrapper -%}
+    {%- endif -%}
+{%- endif -%}
+
+{%- set response_format_type = none -%}
+{%- if response_format is defined and response_format is mapping -%}
+    {%- set response_format_type = response_format.get('type') -%}
+{%- elif response_format is defined -%}
+    {%- set response_format_type = response_format -%}
+{%- endif -%}
+{%- if response_format_type == 'json_object' -%}
+    {{- internal_system_message(
+        'response-format',
+        'The system is invoked with `response_format=json_object`.\nYour response must be raw JSON data without markdown code blocks (```json) or any additional formatting.'
+    ) -}}
+{%- elif response_format_type == 'json_schema' -%}
+    {{- internal_system_message(
+        'response-format',
+        'The system is invoked with `response_format=json_schema`.\nYour response must be raw JSON data without markdown code blocks (```json) or any additional formatting.\nThe JSON data must match the following schema:\n```json\n' + json_sorted(state.response_schema) + '\n```'
+    ) -}}
+{%- endif -%}
+
+{%- if add_generation_prompt -%}
+    {{- open_tag('message', [('role', 'assistant')]) -}}
+    {{- open_tag('think' if thinking else 'response') -}}
+{%- endif -%}
+
+{%- if image_prompts is defined and image_prompts is not none and state.image_index != image_prompts|length -%}
+    {{- raise_exception('image prompt count ' + image_prompts|length|string + ' != consumed placeholder count ' + state.image_index|string) -}}
+{%- endif -%}
+

models/templates/MiniMax-M1.jinja

@@ -0,0 +1,91 @@
+{{ '<begin_of_document>' -}}
+{%- if custom_tools is defined %}
+    {%- set tools = custom_tools %}
+{%- endif %}
+{%- if not tools is defined %}
+    {%- set tools = none %}
+{%- endif %}
+
+{#- Extract system message #}
+{% set ns = namespace(system_prompt='') -%}
+{%- if messages[0]['role'] == 'system' %}
+    {%- if messages[0]['content'] is string %}
+        {%- set ns.system_prompt = messages[0]['content']|trim %}
+    {%- else %}
+        {%- set ns.system_prompt = messages[0]['content'][0]['text']|trim %}
+    {%- endif %}
+    {%- set messages = messages[1:] %}
+{%- else %}
+    {%- if tools is not none %}
+        {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
+    {%- else %}
+        {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
+    {%- endif %}
+{%- endif %}
+
+{#- System message #}
+{%- if ns.system_prompt != '' %}
+{{ '<beginning_of_sentence>system ai_setting=assistant\n' + ns.system_prompt + '<end_of_sentence>\n' -}}
+{%- endif %}
+
+{#- Tools configuration #}
+{%- if tools is not none %}
+{{ '<beginning_of_sentence>system tool_setting=tools\nYou are provided with these tools:\n<tools>\n' -}}
+{%- for tool in tools %}
+{{ tool | tojson ~ '\n' -}}
+{%- endfor %}
+{{ '</tools>\n\nIf you need to call tools, please respond with <tool_calls></tool_calls> XML tags, and provide tool-name and json-object of arguments, following the format below:\n<tool_calls>\n{"name": <tool-name>, "arguments": <args-json-object>}\n...\n</tool_calls><end_of_sentence>\n' -}}
+{%- endif %}
+
+{#- Process messages #}
+{%- for message in messages %}
+    {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
+        {%- if message['role'] == 'user' %}
+{{ '<beginning_of_sentence>user name=user\n' -}}
+{%- if message['content'] is string %}
+{{ message['content']|trim -}}
+{%- else %}
+{%- for content in message['content'] %}
+{%- if content['type'] == 'text' %}
+{{ content['text']|trim -}}
+{%- endif %}
+{%- endfor %}
+{%- endif %}
+{{ '<end_of_sentence>\n' -}}
+        {%- elif message['role'] == 'assistant' %}
+{{ '<beginning_of_sentence>ai name=assistant\n' -}}
+{%- if message['content'] is string %}
+{{ message['content']|trim -}}
+{%- else %}
+{%- for content in message['content'] | selectattr('type', 'equalto', 'text') %}
+{{ content['text']|trim -}}
+{%- endfor %}
+{%- endif %}
+{{ '<end_of_sentence>\n' -}}
+        {%- endif %}
+    {%- elif 'tool_calls' in message %}
+{{ '<beginning_of_sentence>ai name=assistant\n<tool_calls>\n' -}}
+{%- for tool_call in message.tool_calls %}
+{{ '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}\n' -}}
+{%- endfor %}
+{{ '</tool_calls><end_of_sentence>\n' -}}
+    {%- elif message.role == "tool" or message.role == "ipython" %}
+{{ '<beginning_of_sentence>tool name=tools\n' -}}
+{%- if message.content is string %}
+{{ 'tool result: ' + message.content + '\n\n' -}}
+{%- else %}
+{%- for content in message['content'] %}
+{%- if content['type'] == 'text' %}
+{{ 'tool result: ' + content['text'] + '\n\n' -}}
+{%- elif content.get('name') %}
+{{ 'tool name: ' + content['name'] + '\ntool result: ' + content['text'] + '\n\n' -}}
+{%- endif %}
+{%- endfor %}
+{%- endif %}
+{{ '<end_of_sentence>\n' -}}
+    {%- endif %}
+{%- endfor %}
+
+{%- if add_generation_prompt %}
+{{ '<beginning_of_sentence>ai name=assistant\n' -}}
+{%- endif %}

scripts/bench-models.sh

@@ -22,8 +22,8 @@ if (( QUICK )); then
 fi
 
 if (( DIO )); then
-    ARGS_BB="${ARGS_BB} --no-mmap --direct-io"
-    ARGS_B="${ARGS_B} -mmp 0 -dio 1"
+    ARGS_BB="${ARGS_BB} --load-mode dio"
+    ARGS_B="${ARGS_B} --load-mode dio"
 fi
 
 run_model() {

scripts/snapdragon/adb/run-bench.sh

@@ -43,7 +43,7 @@ adb $adbserial $adbhost shell " \
   cd $basedir;         \
   LD_LIBRARY_PATH=$basedir/$branch/lib   \
   ADSP_LIBRARY_PATH=$basedir/$branch/lib \
-    $ndev $nhvx $opmask $verbose $profile $hb ./$branch/bin/llama-bench --device $device --mmap 0 -m $basedir/../gguf/$model \
+    $ndev $nhvx $opmask $verbose $profile $hb ./$branch/bin/llama-bench --device $device --load-mode none -m $basedir/../gguf/$model \
         --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 \
         --ubatch-size 1024 -fa 1 -ngl 99 $cli_opts $@   \
 "

scripts/snapdragon/adb/run-cli.sh

@@ -71,7 +71,7 @@ adb $adbserial $adbhost shell " \
     LD_LIBRARY_PATH=$basedir/$branch/lib   \
     ADSP_LIBRARY_PATH=$basedir/$branch/lib \
     $verbose $sched $opmask $profile $nhvx $hmx $ndev $hb $opbatch $opqueue $opflt $vmem $mbuf \
-      ./$branch/bin/llama-cli --no-mmap -m $basedir/../gguf/$model \
+      ./$branch/bin/llama-cli --load-mode none -m $basedir/../gguf/$model \
          --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1           \
          --ctx-size 8192 --ubatch-size 1024 -fa on                 \
          -ngl 99 --device $device $cli_opts $@                     \

scripts/snapdragon/adb/run-completion.sh

@@ -79,7 +79,7 @@ adb $adbserial $adbhost shell " \
     LD_LIBRARY_PATH=$basedir/$branch/lib   \
     ADSP_LIBRARY_PATH=$basedir/$branch/lib \
     $verbose $sched $opmask $profile $nhvx $hmx $ndev $hb $opbatch $opqueue $oppoll $opflt $opfuse $vmem $mbuf $mmsel $fasel \
-      ./$branch/bin/llama-completion --no-mmap -m $basedir/../gguf/$model \
+      ./$branch/bin/llama-completion --load-mode none -m $basedir/../gguf/$model \
          --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1                  \
          --ctx-size 8192 --ubatch-size 1024 -fa on                        \
          -ngl 99 --device $device $cli_opts $@                            \

scripts/snapdragon/adb/run-mtmd.sh

@@ -62,7 +62,7 @@ adb $adbserial $adbhost shell " \
     LD_LIBRARY_PATH=$basedir/$branch/lib   \
     ADSP_LIBRARY_PATH=$basedir/$branch/lib \
     $verbose $experimental $sched $opmask $profile $hmx $nhvx $ndev $mtmd_backend \
-      ./$branch/bin/llama-mtmd-cli --no-mmap -m $basedir/../gguf/$model      \
+      ./$branch/bin/llama-mtmd-cli --load-mode none -m $basedir/../gguf/$model      \
          --mmproj $basedir/../gguf/$mmproj                                   \
          --image $basedir/../gguf/$image                                     \
          --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1                     \

scripts/snapdragon/windows/run-bench.ps1

@@ -43,6 +43,6 @@ if ($null -ne $env:HB) {
 $env:ADSP_LIBRARY_PATH="$basedir\lib"
 
 & "$basedir\bin\llama-bench.exe" `
-    --mmap 0 -m $basedir\..\..\gguf\$model `
+    --load-mode none -m $basedir\..\..\gguf\$model `
     --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 `
     --ubatch-size 1024 -ngl 99 --device $device $cli_opts

scripts/snapdragon/windows/run-cli.ps1

@@ -47,7 +47,7 @@ if ($null -ne $env:HB) {
 $env:ADSP_LIBRARY_PATH="$basedir\lib"
 
 & "$basedir\bin\llama-cli.exe" `
-    --no-mmap -m $basedir\..\..\gguf\$model `
+    --load-mode none -m $basedir\..\..\gguf\$model `
     --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 `
     --ctx-size 8192 --ubatch-size 1024 -fa on `
     -ngl 99 --device $device $cli_opts

scripts/snapdragon/windows/run-completion.ps1

@@ -47,7 +47,7 @@ if ($null -ne $env:HB) {
 $env:ADSP_LIBRARY_PATH="$basedir\lib"
 
 & "$basedir\bin\llama-completion.exe" `
-    --no-mmap -m $basedir\..\..\gguf\$model `
+    --load-mode none -m $basedir\..\..\gguf\$model `
     --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 `
     --ctx-size 8192 --ubatch-size 1024 -fa on `
     -ngl 99 -no-cnv --device $device $cli_opts

scripts/snapdragon/windows/run-mtmd.ps1

@@ -60,7 +60,7 @@ if ($null -ne $env:MTMD_DEVICE) {
 $env:ADSP_LIBRARY_PATH="$basedir\lib"
 
 & "$basedir\bin\llama-mtmd-cli.exe" `
-    --no-mmap -m $basedir\..\..\gguf\$model `
+    --load-mode none -m $basedir\..\..\gguf\$model `
     --mmproj $basedir\..\..\gguf\$mmproj `
     --image $basedir\..\..\gguf\$image `
     --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 `

scripts/sync-ggml.last

@@ -1 +1 @@
-8846b79e66747bb9f68597420e95114c177315ce
+2d191b5dee1a591c41ee8a653ce42bfcd9c8716d

scripts/sync_vendor.py

@@ -5,7 +5,7 @@ import os
 import sys
 import subprocess
 
-HTTPLIB_VERSION = "refs/tags/v0.53.0"
+HTTPLIB_VERSION = "refs/tags/v0.53.1"
 
 vendor = {
     "https://github.com/nlohmann/json/releases/latest/download/json.hpp":     "vendor/nlohmann/json.hpp",

skills/code-review/SKILL.md

@@ -46,6 +46,8 @@ Mandatory on every review; any finding here is **blocking**. Rule of thumb: GGUF
 
 - **Sizes/counts from tensor dims:** validate before allocating. Products like `ne[i]*nb[i]`/nbytes can overflow on crafted dims into an undersized alloc then heap overflow. Overflow checks must run BEFORE the arithmetic they guard - padding/alignment macros wrap to 0 near `SIZE_MAX`, so a guard after the pad passes.
 - **GGUF strings/arrays:** cap declared lengths and element counts before using them to size a loop or buffer; validate element type and length before casting an array to a pointer or reading fixed indices (`[i+1]`, `[0..2]`).
+- **Element-type confusion:** casting `gguf_get_arr_data()` or `tensor->data` to `float *`/`int32_t *` needs an element-type check first (`gguf_get_kv_type() == GGUF_TYPE_ARRAY` then `gguf_get_arr_type()`; `type == GGML_TYPE_F32` for tensors). A `UINT8` array or `I8` tensor passes every length check, then gets read 4 bytes per element - a nearby length check is not a type check.
+- **Loaders:** `GGML_ASSERT` on a file-derived value aborts the process; throw instead where the caller already catches (vocab, model loader, clip).
 - **File-supplied counts indexing fixed arrays:** bound any count (e.g. layer/block count into a `LLAMA_MAX_*` array) before indexing; watch checks that only fire when an optional key is present.
 - **Declared vs actual array length:** check the declared length of a GGUF array against the count actually read, not just against a buffer size.
 - **Bounds comparisons:** flag narrowing casts (`size_t`->`int32_t`) and signed/unsigned mixing that can bypass a length check and copy past a buffer.

src/llama-arch.cpp

@@ -128,6 +128,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
     { LLM_ARCH_SEED_OSS,         "seed_oss"         },
     { LLM_ARCH_GROVEMOE,         "grovemoe"         },
     { LLM_ARCH_APERTUS,          "apertus"          },
+    { LLM_ARCH_MINIMAX_01,       "minimax-01"       },
     { LLM_ARCH_MINIMAX_M2,       "minimax-m2"       },
     { LLM_ARCH_MINIMAX_M3,       "minimax-m3"       },
     { LLM_ARCH_COGVLM,           "cogvlm"           },
@@ -143,6 +144,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
     { LLM_ARCH_LLAMA_EMBED,      "llama-embed"      },
     { LLM_ARCH_MAINCODER,        "maincoder"        },
     { LLM_ARCH_KIMI_LINEAR,      "kimi-linear"      },
+    { LLM_ARCH_KIMI_K3,          "kimi-k3"          },
     { LLM_ARCH_TALKIE,           "talkie"           },
     { LLM_ARCH_MELLUM,           "mellum"           },
     { LLM_ARCH_NANBEIGE,         "nanbeige"         },
@@ -186,6 +188,9 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
     { LLM_KV_FEATURES_LENGTH,                   "%s.features_length"                   },
     { LLM_KV_BLOCK_COUNT,                       "%s.block_count"                       },
     { LLM_KV_LEADING_DENSE_BLOCK_COUNT,         "%s.leading_dense_block_count"         },
+    { LLM_KV_ATTN_RES_BLOCK_SIZE,               "%s.attn_res.block_size"               },
+    { LLM_KV_ACTIVATION_SITU_BETA,              "%s.activation.situ_beta"              },
+    { LLM_KV_ACTIVATION_SITU_LINEAR_BETA,       "%s.activation.situ_linear_beta"       },
     { LLM_KV_FEED_FORWARD_LENGTH,               "%s.feed_forward_length"               },
     { LLM_KV_EXPERT_FEED_FORWARD_LENGTH,        "%s.expert_feed_forward_length"        },
     { LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, "%s.expert_shared_feed_forward_length" },
@@ -201,6 +206,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
     { LLM_KV_EXPERT_GROUP_USED_COUNT,           "%s.expert_group_used_count"           },
     { LLM_KV_EXPERT_WEIGHTS_SCALE,              "%s.expert_weights_scale"              },
     { LLM_KV_EXPERT_WEIGHTS_NORM,               "%s.expert_weights_norm"               },
+    { LLM_KV_EXPERT_LATENT_LENGTH,              "%s.expert_latent_length"              },
     { LLM_KV_EXPERT_GATING_FUNC,                "%s.expert_gating_func"                },
     { LLM_KV_EXPERT_GROUP_SCALE,                "%s.expert_group_scale"                },
     { LLM_KV_EXPERTS_PER_GROUP,                 "%s.experts_per_group"                 },
@@ -312,6 +318,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
     { LLM_KV_SSM_DT_B_C_RMS,     "%s.ssm.dt_b_c_rms"     },
 
     { LLM_KV_KDA_HEAD_DIM, "%s.kda.head_dim" },
+    { LLM_KV_KDA_GATE_LOWER_BOUND, "%s.kda.gate_lower_bound" },
 
     { LLM_KV_WKV_HEAD_SIZE, "%s.wkv.head_size" },
 
@@ -462,6 +469,13 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
     { LLM_TENSOR_SSM_F_B,                                "blk.%d.ssm_f_b" },
     { LLM_TENSOR_SSM_BETA,                               "blk.%d.ssm_beta" },
     { LLM_TENSOR_SSM_G_A,                                "blk.%d.ssm_g_a" },
+    { LLM_TENSOR_SSM_G,                                  "blk.%d.ssm_g" },
+    { LLM_TENSOR_ATTN_RES_SCORE,                         "blk.%d.attn_res_score" },
+    { LLM_TENSOR_FFN_RES_SCORE,                          "blk.%d.ffn_res_score" },
+    { LLM_TENSOR_OUTPUT_RES_SCORE,                       "output_res_score" },
+    { LLM_TENSOR_FFN_ROUTED_DOWN,                        "blk.%d.ffn_routed_down" },
+    { LLM_TENSOR_FFN_ROUTED_UP,                          "blk.%d.ffn_routed_up" },
+    { LLM_TENSOR_FFN_ROUTED_NORM,                        "blk.%d.ffn_routed_norm" },
     { LLM_TENSOR_SSM_G_B,                                "blk.%d.ssm_g_b" },
     { LLM_TENSOR_SSM_NORM,                               "blk.%d.ssm_norm" },
     { LLM_TENSOR_ATTN_Q_A_NORM,                          "blk.%d.attn_q_a_norm" },
@@ -755,6 +769,13 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
     {LLM_TENSOR_SSM_F_B,                    {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_SSM_BETA,                   {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_SSM_G_A,                    {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_SSM_G,                      {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_ATTN_RES_SCORE,             {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
+    {LLM_TENSOR_FFN_RES_SCORE,              {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
+    {LLM_TENSOR_OUTPUT_RES_SCORE,           {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_MUL}},
+    {LLM_TENSOR_FFN_ROUTED_DOWN,            {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_FFN_ROUTED_UP,              {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_FFN_ROUTED_NORM,            {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
     {LLM_TENSOR_SSM_G_B,                    {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_TIME_MIX_LERP_X,            {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
     {LLM_TENSOR_TIME_MIX_LN,                {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
@@ -975,9 +996,11 @@ bool llm_arch_is_hybrid(const llm_arch & arch) {
         case LLM_ARCH_NEMOTRON_H_MOE:
         case LLM_ARCH_QWEN3NEXT:
         case LLM_ARCH_KIMI_LINEAR:
+        case LLM_ARCH_KIMI_K3:
         case LLM_ARCH_QWEN35:
         case LLM_ARCH_QWEN35MOE:
         case LLM_ARCH_DEEPSEEK4:
+        case LLM_ARCH_MINIMAX_01:
             return true;
         default:
             return false;
@@ -1001,6 +1024,8 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) {
         case LLM_ARCH_QWEN35:
         case LLM_ARCH_QWEN35MOE:
         case LLM_ARCH_DEEPSEEK4:
+        case LLM_ARCH_NEMOTRON_H:
+        case LLM_ARCH_NEMOTRON_H_MOE:
             return true;
         default:
             return false;
@@ -1031,10 +1056,12 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
         case LLM_ARCH_GRANITE_HYBRID:
         case LLM_ARCH_LFM2:
         case LLM_ARCH_LFM2MOE:
+        case LLM_ARCH_MINIMAX_01:
         case LLM_ARCH_MINIMAX_M2:
         case LLM_ARCH_MINIMAX_M3:
         case LLM_ARCH_MISTRAL4:
         case LLM_ARCH_KIMI_LINEAR:
+        case LLM_ARCH_KIMI_K3:
         case LLM_ARCH_QWEN3TTS:
             return false;
         default:

src/llama-arch.h

@@ -145,6 +145,7 @@ enum llm_arch {
     LLM_ARCH_LLAMA_EMBED,
     LLM_ARCH_MAINCODER,
     LLM_ARCH_KIMI_LINEAR,
+    LLM_ARCH_KIMI_K3,
     LLM_ARCH_TALKIE,
     LLM_ARCH_MELLUM,
     LLM_ARCH_EAGLE3,
@@ -153,6 +154,7 @@ enum llm_arch {
     LLM_ARCH_NANBEIGE,
     LLM_ARCH_QWEN3TTS,
     LLM_ARCH_POCKETTTS,
+    LLM_ARCH_MINIMAX_01,
     LLM_ARCH_UNKNOWN,
 };
 
@@ -191,6 +193,9 @@ enum llm_kv {
     LLM_KV_FEATURES_LENGTH,
     LLM_KV_BLOCK_COUNT,
     LLM_KV_LEADING_DENSE_BLOCK_COUNT,
+    LLM_KV_ATTN_RES_BLOCK_SIZE,
+    LLM_KV_ACTIVATION_SITU_BETA,
+    LLM_KV_ACTIVATION_SITU_LINEAR_BETA,
     LLM_KV_FEED_FORWARD_LENGTH,
     LLM_KV_EXPERT_FEED_FORWARD_LENGTH,
     LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH,
@@ -206,6 +211,7 @@ enum llm_kv {
     LLM_KV_EXPERT_GROUP_USED_COUNT,
     LLM_KV_EXPERT_WEIGHTS_SCALE,
     LLM_KV_EXPERT_WEIGHTS_NORM,
+    LLM_KV_EXPERT_LATENT_LENGTH,
     LLM_KV_EXPERT_GATING_FUNC,
     LLM_KV_EXPERT_GROUP_SCALE,
     LLM_KV_EXPERTS_PER_GROUP,
@@ -317,6 +323,7 @@ enum llm_kv {
     LLM_KV_SSM_DT_B_C_RMS,
 
     LLM_KV_KDA_HEAD_DIM,
+    LLM_KV_KDA_GATE_LOWER_BOUND,
 
     LLM_KV_WKV_HEAD_SIZE,
 
@@ -491,6 +498,13 @@ enum llm_tensor {
     LLM_TENSOR_SSM_BETA,            // kimi: beta mixing coefficient and qwen3.5
     LLM_TENSOR_SSM_G_A,             // kimi: output gate projection A
     LLM_TENSOR_SSM_G_B,             // kimi: output gate projection B
+    LLM_TENSOR_SSM_G,               // kimi-k3: full-rank KDA gate
+    LLM_TENSOR_ATTN_RES_SCORE,      // kimi-k3: fused res_norm*res_proj (pre-attn)
+    LLM_TENSOR_FFN_RES_SCORE,       // kimi-k3: fused res_norm*res_proj (pre-ffn)
+    LLM_TENSOR_OUTPUT_RES_SCORE,    // kimi-k3: fused res_norm*res_proj (final)
+    LLM_TENSOR_FFN_ROUTED_DOWN,     // kimi-k3: latent MoE down
+    LLM_TENSOR_FFN_ROUTED_UP,       // kimi-k3: latent MoE up
+    LLM_TENSOR_FFN_ROUTED_NORM,     // kimi-k3: latent MoE norm
     LLM_TENSOR_TIME_MIX_W0,
     LLM_TENSOR_TIME_MIX_W1,
     LLM_TENSOR_TIME_MIX_W2,

src/llama-context.cpp

@@ -103,7 +103,7 @@ llama_context::llama_context(
 
     cparams.n_rs_seq = params.n_rs_seq;
     if (cparams.n_rs_seq > 0 && !llm_arch_supports_rs_rollback(model.arch)) {
-        LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model arch does not support recurrent partial rollback; clamping to 0\n",
+        LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model does not support recurrent partial rollback; clamping to 0\n",
                         __func__, cparams.n_rs_seq);
         cparams.n_rs_seq = 0;
     }
@@ -2293,13 +2293,17 @@ void llama_context::output_reorder() {
 
 uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
     uint32_t res;
-    if (model.arch == LLM_ARCH_QWEN3NEXT ||
+    if (model.arch == LLM_ARCH_KIMI_K3) {
+        // the n_tokens*40 budget below is exhausted at ubatch 3840
+        res = std::max<uint32_t>(n_tokens * 160, 64u * model.n_tensors());
+    } else if (model.arch == LLM_ARCH_QWEN3NEXT ||
         model.arch == LLM_ARCH_KIMI_LINEAR ||
         model.arch == LLM_ARCH_QWEN35 ||
         model.arch == LLM_ARCH_QWEN35MOE ||
         model.arch == LLM_ARCH_DEEPSEEK4 ||
         (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) ||
         model.arch == LLM_ARCH_NANBEIGE ||
+        model.arch == LLM_ARCH_MINIMAX_01 ||
         model.arch == LLM_ARCH_MINIMAX_M3) {
         res = std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
     } else {

src/llama-graph.cpp

@@ -1835,6 +1835,8 @@ ggml_tensor * llm_graph_context::build_ffn(
                 cur = ggml_reglu(ctx0, cur);
                 cb(cur, "ffn_reglu", il);
             } break;
+        case LLM_FFN_SITU:
+            GGML_ABORT("not yet supported");
         default:
             GGML_ABORT("fatal error");
     }
@@ -2174,6 +2176,21 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
                 cur = ggml_silu(ctx0, cur);
                 cb(cur, "ffn_moe_silu", il);
             } break;
+        case LLM_FFN_SITU:
+            {
+                // situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * lb*tanh(up/lb)
+                GGML_ASSERT(has_gate);
+                const float beta = hparams.situ_beta;
+                const float lb   = hparams.situ_linear_beta;
+
+                ggml_tensor * act = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, cur, 1.0f/beta)), beta);
+                act = ggml_mul(ctx0, act, ggml_sigmoid(ctx0, cur));
+                if (lb > 0.0f) {
+                    up = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, up, 1.0f/lb)), lb);
+                }
+                cur = ggml_mul(ctx0, act, up);
+                cb(cur, "ffn_moe_situ", il);
+            } break;
         case LLM_FFN_GELU:
             if (has_gate) {
                 cur = ggml_geglu_split(ctx0, cur, up);

src/llama-graph.h

@@ -59,6 +59,7 @@ enum llm_ffn_op_type : int {
     LLM_FFN_GEGLU,
     LLM_FFN_REGLU,
     LLM_FFN_SWIGLU_OAI_MOE,
+    LLM_FFN_SITU,           // kimi-k3
 };
 
 enum llm_ffn_gate_type {

src/llama-hparams.cpp

@@ -217,6 +217,13 @@ uint32_t llama_hparams::n_embd_s() const {
         return n_embd_head_kda * n_embd_head_kda * n_head();  // 128 * 128 * 32 = 524288
     }
 
+    if (n_embd_head_la != 0) {
+        // for MiniMax-Text-01 linear attention layers
+        // Full recurrent state: head_dim * head_dim * n_head
+        // tensor shape for linear attention: [head_dim, head_dim, n_head]
+        return n_embd_head_la * n_embd_head_la * n_head();  // 128 * 128 * 64 = 1048576
+    }
+
     // corresponds to Mamba's ssm_states size
     return ssm_d_state * ssm_d_inner;
 }

src/llama-hparams.h

@@ -4,10 +4,11 @@
 
 #include <array>
 #include <cassert>
+#include <cmath>
 
 // bump if necessary
 #define LLAMA_MAX_LAYERS  512
-#define LLAMA_MAX_EXPERTS 512 // Qwen3 Next
+#define LLAMA_MAX_EXPERTS 1024 // Kimi K3
 
 enum llama_expert_gating_func_type {
     LLAMA_EXPERT_GATING_FUNC_TYPE_NONE           = 0,
@@ -164,9 +165,19 @@ struct llama_hparams {
     uint32_t ssm_dt_rank = 0;
     uint32_t ssm_n_group = 0;
 
+    // for MiniMax-Text-01 linear attention
+    uint32_t n_embd_head_la = 0;
+
     // for Kimi Linear KDA
     uint32_t n_embd_head_kda = 0;
 
+    // kimi-k3
+    uint32_t n_expert_latent      = 0;      // routed_expert_hidden_size (0 = experts run at n_embd)
+    uint32_t attn_res_block_size  = 0;      // 0 = no cross-layer attention residuals
+    float    kda_gate_lower_bound = -INFINITY;
+    float    situ_beta            = 1.0f;
+    float    situ_linear_beta     = 0.0f;   // 0 = no linear-beta transform on the up branch
+
     bool ssm_dt_b_c_rms = false;
 
     float f_clamp_kqv      = 0.0f;

src/llama-model-loader.cpp

@@ -316,15 +316,19 @@ namespace GGUFMeta {
         struct GGUFMeta::ArrayInfo arr_info =
             GGUFMeta::GKV<GGUFMeta::ArrayInfo>::get_kv(ctx, kid);
 
+        bool type_ok = false;
         switch (arr_info.gt) {
             case GGUF_TYPE_UINT32:
-            case GGUF_TYPE_INT32:   GGML_ASSERT((std::is_same<T,     int32_t>::value) ||
-                                                (std::is_same<T,    uint32_t>::value)); break;
-            case GGUF_TYPE_FLOAT32: GGML_ASSERT((std::is_same<T,       float>::value)); break;
-            case GGUF_TYPE_STRING:  GGML_ASSERT((std::is_same<T, std::string>::value)); break;
+            case GGUF_TYPE_INT32:   type_ok = (std::is_same<T,     int32_t>::value) ||
+                                              (std::is_same<T,    uint32_t>::value); break;
+            case GGUF_TYPE_FLOAT32: type_ok = (std::is_same<T,       float>::value); break;
+            case GGUF_TYPE_STRING:  type_ok = (std::is_same<T, std::string>::value); break;
             default:
                 throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str()));
         }
+        if (!type_ok) {
+            throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt)));
+        }
 
         if constexpr (std::is_same<T, std::string>::value) {
             const size_t n_items = gguf_get_arr_n(ctx, kid);
@@ -357,16 +361,20 @@ namespace GGUFMeta {
         struct GGUFMeta::ArrayInfo arr_info =
             GGUFMeta::GKV<GGUFMeta::ArrayInfo>::get_kv(ctx, kid);
 
+        bool type_ok = false;
         switch (arr_info.gt) {
             case GGUF_TYPE_BOOL:
             case GGUF_TYPE_UINT32:
-            case GGUF_TYPE_INT32:   GGML_ASSERT((std::is_same<T,     int32_t>::value) ||
-                                                (std::is_same<T,    uint32_t>::value)); break;
-            case GGUF_TYPE_FLOAT32: GGML_ASSERT((std::is_same<T,       float>::value)); break;
-            case GGUF_TYPE_STRING:  GGML_ASSERT((std::is_same<T, std::string>::value)); break;
+            case GGUF_TYPE_INT32:   type_ok = (std::is_same<T,     int32_t>::value) ||
+                                              (std::is_same<T,    uint32_t>::value); break;
+            case GGUF_TYPE_FLOAT32: type_ok = (std::is_same<T,       float>::value); break;
+            case GGUF_TYPE_STRING:  type_ok = (std::is_same<T, std::string>::value); break;
             default:
                 throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str()));
         }
+        if (!type_ok) {
+            throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt)));
+        }
 
         if (arr_info.length > N_MAX) {
             throw std::runtime_error(format("array length %u for key %s exceeds max %u", (uint32_t) arr_info.length, key.c_str(), (uint32_t) N_MAX));
@@ -1002,7 +1010,7 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w
                 ggml_tensor * B   = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d_state, n_group, n_seq_tokens, n_seqs);
                 ggml_tensor * C   = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d_state, n_group, n_seq_tokens, n_seqs);
                 ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_seqs);
-                op_tensor = ggml_ssm_scan(ctx, s, x, dt, w, B, C, ids);
+                op_tensor = ggml_ssm_scan(ctx, s, x, dt, w, B, C, ids, /*K=*/1);
             } break;
         case GGML_OP_RWKV_WKV6:
             {
@@ -1178,7 +1186,7 @@ struct ggml_tensor * llama_model_loader::create_tensor(
                         if (use_mmap) {
                             static std::once_flag once;
                             std::call_once(once, [] {
-                                LLAMA_LOG_WARN("llama_model_loader: tensor overrides to CPU are used with mmap enabled - consider using --no-mmap for better performance\n");
+                                LLAMA_LOG_WARN("llama_model_loader: tensor overrides to CPU are used with mmap enabled - consider using --load-mode none for better performance\n");
                             });
                         }
                     } else {

src/llama-model-saver.cpp

@@ -213,6 +213,7 @@ void llama_model_saver::add_kv_from_model() {
     add_kv(LLM_KV_LEADING_DENSE_BLOCK_COUNT,         hparams.n_layer_dense_lead);
     add_kv(LLM_KV_FEED_FORWARD_LENGTH,               hparams.n_ff_arr, true);
     add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH,        hparams.n_ff_exp);
+    add_kv(LLM_KV_EXPERT_LATENT_LENGTH,              hparams.n_expert_latent);
     add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp);
     add_kv(LLM_KV_EXPERT_CHUNK_FEED_FORWARD_LENGTH,  hparams.n_ff_chexp);
     add_kv(LLM_KV_SWIGLU_CLAMP_EXP,                  hparams.swiglu_clamp_exp);
@@ -319,6 +320,7 @@ void llama_model_saver::add_kv_from_model() {
     add_kv(LLM_KV_SSM_DT_B_C_RMS,                    hparams.ssm_dt_b_c_rms);
 
     add_kv(LLM_KV_KDA_HEAD_DIM,                      hparams.n_embd_head_kda);
+    add_kv(LLM_KV_KDA_GATE_LOWER_BOUND,              hparams.kda_gate_lower_bound);
 
     add_kv(LLM_KV_WKV_HEAD_SIZE,                     hparams.wkv_head_size);
 
@@ -376,6 +378,10 @@ void llama_model_saver::add_kv_from_model() {
     add_kv(LLM_KV_XIELU_BETA,                        hparams.xielu_beta);
     add_kv(LLM_KV_XIELU_EPS,                         hparams.xielu_eps);
 
+    add_kv(LLM_KV_ATTN_RES_BLOCK_SIZE,               hparams.attn_res_block_size);
+    add_kv(LLM_KV_ACTIVATION_SITU_BETA,              hparams.situ_beta);
+    add_kv(LLM_KV_ACTIVATION_SITU_LINEAR_BETA,       hparams.situ_linear_beta);
+
     // deprecated
     // add_kv(LLM_KV_TOKENIZER_PREFIX_ID,               ???);
     // add_kv(LLM_KV_TOKENIZER_SUFFIX_ID,               ???);
@@ -403,6 +409,7 @@ void llama_model_saver::add_tensors_from_model() {
     add_tensor(model->output_norm_enc);
     add_tensor(model->output_s);
     add_tensor(model->output_in_s);
+    add_tensor(model->output_res_score);
     add_tensor(model->cls);
     add_tensor(model->cls_b);
     add_tensor(model->cls_out);

src/llama-model.cpp

@@ -296,6 +296,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
             return new llama_model_grovemoe(params);
         case LLM_ARCH_APERTUS:
             return new llama_model_apertus(params);
+        case LLM_ARCH_MINIMAX_01:
+            return new llama_model_minimax_01(params);
         case LLM_ARCH_MINIMAX_M2:
             return new llama_model_minimax_m2(params);
         case LLM_ARCH_MINIMAX_M3:
@@ -320,6 +322,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
             return new llama_model_mimo2(params);
         case LLM_ARCH_KIMI_LINEAR:
             return new llama_model_kimi_linear(params);
+        case LLM_ARCH_KIMI_K3:
+            return new llama_model_kimi_k3(params);
         case LLM_ARCH_STEP35:
             return new llama_model_step35(params);
         default:
@@ -798,6 +802,7 @@ const char * llm_type_name(llm_type type) {
         case LLM_TYPE_290B:          return "290B";
         case LLM_TYPE_314B:          return "314B";
         case LLM_TYPE_405B:          return "405B";
+        case LLM_TYPE_456B:          return "456B";
         case LLM_TYPE_671B:          return "671B";
         case LLM_TYPE_SMALL:         return "0.1B";
         case LLM_TYPE_MEDIUM:        return "0.4B";
@@ -842,6 +847,7 @@ const char * llm_type_name(llm_type type) {
         case LLM_TYPE_397B_A17B:     return "397B.A17B";
         case LLM_TYPE_685B_A37B:     return "685B.A37B";
         case LLM_TYPE_744B_A40B:     return "744B.A40B";
+        case LLM_TYPE_2_8T_A50B:     return "2.8T.A50B";
         case LLM_TYPE_E2B:           return "E2B";
         case LLM_TYPE_E4B:           return "E4B";
         default:                     return "?B";
@@ -2283,7 +2289,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                         filter_recr = [&](uint32_t il) {
                             return hparams.is_recr(il) && hparams.n_ff(il) == 0;
                         };
-                    } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE) {
+                    } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_MINIMAX_01) {
                         filter_attn = [&](uint32_t il) {
                             return il < hparams.n_layer() && !hparams.is_recr(il);
                         };
@@ -2599,6 +2605,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
         case LLM_ARCH_NEMOTRON_H:
         case LLM_ARCH_NEMOTRON_H_MOE:
         case LLM_ARCH_KIMI_LINEAR:
+        case LLM_ARCH_KIMI_K3:
             return LLAMA_ROPE_TYPE_NONE;
 
         // use what we call a normal RoPE, operating on pairs of consecutive head values
@@ -2704,6 +2711,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
         case LLM_ARCH_SEED_OSS:
         case LLM_ARCH_GROVEMOE:
         case LLM_ARCH_APERTUS:
+        case LLM_ARCH_MINIMAX_01:
         case LLM_ARCH_MINIMAX_M2:
         case LLM_ARCH_MINIMAX_M3:
         case LLM_ARCH_COGVLM:

src/llama-model.h

@@ -99,6 +99,7 @@ enum llm_type {
     LLM_TYPE_290B,
     LLM_TYPE_314B,
     LLM_TYPE_405B,
+    LLM_TYPE_456B,
     LLM_TYPE_671B,
     LLM_TYPE_SMALL,
     LLM_TYPE_MEDIUM,
@@ -143,6 +144,7 @@ enum llm_type {
     LLM_TYPE_397B_A17B, // Qwen3.5
     LLM_TYPE_685B_A37B, // DeepSeek V3.2
     LLM_TYPE_744B_A40B, // GLM-5
+    LLM_TYPE_2_8T_A50B, // Kimi-K3
     LLM_TYPE_E2B,
     LLM_TYPE_E4B,
 };
@@ -271,6 +273,7 @@ struct llama_layer {
     struct ggml_tensor * wv        = nullptr;
     struct ggml_tensor * wo        = nullptr;
     struct ggml_tensor * wqkv      = nullptr;
+    struct ggml_tensor * wg        = nullptr;
     struct ggml_tensor * wq_a      = nullptr;
     struct ggml_tensor * wq_b      = nullptr;
     struct ggml_tensor * wkv_a_mqa = nullptr;
@@ -528,6 +531,14 @@ struct llama_layer {
     struct ggml_tensor * ssm_g_b    = nullptr;
     struct ggml_tensor * ssm_o_norm = nullptr;
 
+    // kimi-k3
+    struct ggml_tensor * ssm_g           = nullptr; // full-rank KDA gate (replaces ssm_g_a/ssm_g_b)
+    struct ggml_tensor * attn_res_score  = nullptr; // fused res_norm*res_proj, pre-attention
+    struct ggml_tensor * ffn_res_score   = nullptr; // fused res_norm*res_proj, pre-FFN
+    struct ggml_tensor * ffn_routed_down = nullptr; // latent MoE: n_embd -> n_expert_latent
+    struct ggml_tensor * ffn_routed_up   = nullptr; // latent MoE: n_expert_latent -> n_embd
+    struct ggml_tensor * ffn_routed_norm = nullptr;
+
     // DSA (deepseek sparse attention)
     struct ggml_tensor * indexer_k_norm   = nullptr;
     struct ggml_tensor * indexer_k_norm_b = nullptr;
@@ -587,6 +598,7 @@ struct llama_model {
     struct ggml_tensor * tok_norm_b = nullptr;
 
     struct ggml_tensor * output_norm     = nullptr;
+    struct ggml_tensor * output_res_score = nullptr; // kimi-k3: final cross-layer residual mix
     struct ggml_tensor * output_norm_b   = nullptr;
     struct ggml_tensor * output          = nullptr;
     struct ggml_tensor * output_b        = nullptr;

src/llama-vocab.cpp

@@ -1989,6 +1989,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                 // Kimi-K2 doesn't need merges, skip
                 LLAMA_LOG_INFO("%s: Kimi-K2 tokenizer detected, skipping BPE merges\n", __func__);
             } else {
+                if (gguf_get_kv_type(ctx, merges_keyidx) != GGUF_TYPE_ARRAY ||
+                    gguf_get_arr_type(ctx, merges_keyidx) != GGUF_TYPE_STRING) {
+                    throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_MERGES).c_str()));
+                }
                 const int n_merges = gguf_get_arr_n(ctx, merges_keyidx);
                 for (int i = 0; i < n_merges; i++) {
                     const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i);
@@ -2028,8 +2032,13 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
 
             const int precompiled_charsmap_keyidx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str());
             if (precompiled_charsmap_keyidx != -1) {
+                if (gguf_get_kv_type(ctx, precompiled_charsmap_keyidx) != GGUF_TYPE_ARRAY) {
+                    throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str()));
+                }
                 const gguf_type pc_type = gguf_get_arr_type(ctx, precompiled_charsmap_keyidx);
-                GGML_ASSERT(pc_type == GGUF_TYPE_INT8 || pc_type == GGUF_TYPE_UINT8);
+                if (pc_type != GGUF_TYPE_INT8 && pc_type != GGUF_TYPE_UINT8) {
+                    throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str()));
+                }
 
                 const size_t n_precompiled_charsmap = gguf_get_arr_n(ctx, precompiled_charsmap_keyidx);
                 const char * pc = (const char *) gguf_get_arr_data(ctx, precompiled_charsmap_keyidx);
@@ -2081,6 +2090,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                 throw std::runtime_error("cannot find tokenizer merges in model file\n");
             }
             {
+                if (gguf_get_kv_type(ctx, merges_keyidx) != GGUF_TYPE_ARRAY ||
+                    gguf_get_arr_type(ctx, merges_keyidx) != GGUF_TYPE_STRING) {
+                    throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_MERGES).c_str()));
+                }
                 const int n_merges = gguf_get_arr_n(ctx, merges_keyidx);
                 for (int i = 0; i < n_merges; i++) {
                     const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i);
@@ -2407,11 +2420,20 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
         throw std::runtime_error("cannot find tokenizer vocab in model file\n");
     }
 
+    if (gguf_get_kv_type(ctx, token_idx) != GGUF_TYPE_ARRAY ||
+        gguf_get_arr_type(ctx, token_idx) != GGUF_TYPE_STRING) {
+        throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_LIST).c_str()));
+    }
+
     const uint32_t n_tokens = gguf_get_arr_n(ctx, token_idx);
 
     const float * scores = nullptr;
     const int score_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_SCORES).c_str());
     if (score_idx != -1) {
+        if (gguf_get_kv_type(ctx, score_idx) != GGUF_TYPE_ARRAY ||
+            gguf_get_arr_type(ctx, score_idx) != GGUF_TYPE_FLOAT32) {
+            throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_SCORES).c_str()));
+        }
         const uint32_t n_scores = gguf_get_arr_n(ctx, score_idx);
         if (n_scores < n_tokens) {
             throw std::runtime_error("Index out of array bounds for scores (" + std::to_string(n_scores) + " < " + std::to_string(n_tokens) + ")\n");
@@ -2422,6 +2444,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
     const int * toktypes = nullptr;
     const int toktype_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_TOKEN_TYPE).c_str());
     if (toktype_idx != -1) {
+        if (gguf_get_kv_type(ctx, toktype_idx) != GGUF_TYPE_ARRAY ||
+            gguf_get_arr_type(ctx, toktype_idx) != GGUF_TYPE_INT32) {
+            throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_TOKEN_TYPE).c_str()));
+        }
         const uint32_t n_toktypes = gguf_get_arr_n(ctx, toktype_idx);
         if (n_toktypes < n_tokens) {
             throw std::runtime_error("Index out of array bounds for toktypes (" + std::to_string(n_toktypes) + " < " + std::to_string(n_tokens) + ")\n");
@@ -2584,6 +2610,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
         {
             const int suppress_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_SUPPRESS_TOKENS).c_str());
             if (suppress_idx != -1) {
+                if (gguf_get_kv_type(ctx, suppress_idx) != GGUF_TYPE_ARRAY ||
+                    gguf_get_arr_type(ctx, suppress_idx) != GGUF_TYPE_INT32) {
+                    throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_SUPPRESS_TOKENS).c_str()));
+                }
                 const int n = gguf_get_arr_n(ctx, suppress_idx);
                 const int32_t * data = (const int32_t *) gguf_get_arr_data(ctx, suppress_idx);
                 // drop out-of-range ids

src/llama.cpp

@@ -257,7 +257,11 @@ static bool llama_prepare_model_devices(const llama_model_params & params, llama
                     }
 
                     case GGML_BACKEND_DEVICE_TYPE_IGPU:
-                        if (igpus.empty()) {
+                        // igpus.empty() - workaround for integrated devices seen by multiple backends
+                        // ref: https://github.com/ggml-org/llama.cpp/pull/23897
+                        // ggml_backend_dev_backend_reg - allow devices of the same backend regardless if integrated
+                        // ref: https://github.com/ggml-org/llama.cpp/pull/23897#issuecomment-5264222997
+                        if (igpus.empty() || ggml_backend_dev_backend_reg(dev) == ggml_backend_dev_backend_reg(igpus.back().dev)) {
                             igpus.push_back({false, dev});
                         }
                         break;

src/models/dflash.cpp

@@ -14,11 +14,14 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
 
     hparams.n_embd_inp_enc_impl = (uint32_t) target_layer_ids.size() * hparams.n_embd;
 
-    LLAMA_LOG_INFO("%s: DFlash extract_layers = [", __func__);
-    for (size_t i = 0; i < target_layer_ids.size(); ++i) {
-        LLAMA_LOG_INFO("%d%s", target_layer_ids[i], i + 1 < target_layer_ids.size() ? ", " : "");
+    std::string layers;
+    const char * sep = "";
+    for (const auto id : target_layer_ids) {
+        layers += sep;
+        layers += std::to_string(id);
+        sep = ", ";
     }
-    LLAMA_LOG_INFO("]\n");
+    LLAMA_LOG_INFO("%s: DFlash extract_layers = [%s]\n", __func__, layers.c_str());
 
     // DeepSeek-V4 DSpark backbone: stages are full DSV4 blocks, uniform sliding window (the draft KV ring)
     ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult, false);
@@ -40,6 +43,8 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
         ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON,             hparams.dsv4_hc_eps);
         ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS,            hparams.dsv4_compress_ratios, false);
 
+        GGML_ASSERT(hparams.dsv4_o_group_count > 0); // avoid div by zero
+
         if (hparams.expert_gating_func != LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) {
             throw std::runtime_error("DSpark DSV4 draft expects sqrtsoftplus MoE scoring");
         }

src/models/kimi-k3.cpp

@@ -0,0 +1,614 @@
+#include "models.h"
+#include "llama-memory-recurrent.h"
+
+//
+// Kimi-K3 text model: hybrid KDA (linear) + MLA (full) attention, as in kimi-linear.
+// Parts that kimi-linear does not have:
+//   1. cross-layer residual attention  (attn_res_block_size)
+//   2. latent MoE                      (routed experts run at n_expert_latent)
+//   3. situ activation                 (replaces SwiGLU everywhere)
+//   4. MLA output gate                 (sigmoid gate before o_proj)
+//   5. full-rank KDA gate              (single ssm_g instead of ssm_g_a/ssm_g_b)
+//
+
+void llama_model_kimi_k3::load_arch_hparams(llama_model_loader & ml) {
+    ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+    ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA,    hparams.n_embd_head_k_mla_impl);
+    ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA,  hparams.n_embd_head_v_mla_impl);
+    ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK,       hparams.n_lora_q, false);
+    ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK,      hparams.n_lora_kv);
+    ml.get_key(LLM_KV_SSM_CONV_KERNEL,             hparams.ssm_d_conv);
+    ml.get_key(LLM_KV_KDA_HEAD_DIM,                hparams.n_embd_head_kda);
+    ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND,        hparams.kda_gate_lower_bound, false);
+
+    // the MLA cache holds the compressed latent
+    // set it here too, as older GGUFs have no value_length key
+    hparams.n_embd_head_v_full = hparams.n_lora_kv;
+
+    // n_head_kv == 0 marks a KDA (recurrent) layer, as in kimi-linear
+    for (uint32_t i = 0; i < hparams.n_layer(); ++i) {
+        hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0;
+    }
+
+    ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
+    ml.get_key(LLM_KV_EXPERT_SHARED_COUNT,        hparams.n_expert_shared);
+    ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT,  hparams.n_layer_dense_lead, false);
+    ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE,       hparams.expert_weights_scale, false);
+    ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM,        hparams.expert_weights_norm, false);
+    ml.get_key(LLM_KV_EXPERT_GATING_FUNC,         hparams.expert_gating_func);
+    ml.get_key(LLM_KV_EXPERT_LATENT_LENGTH,       hparams.n_expert_latent, false);
+
+    ml.get_key(LLM_KV_ATTN_RES_BLOCK_SIZE,          hparams.attn_res_block_size);
+    ml.get_key(LLM_KV_ACTIVATION_SITU_BETA,         hparams.situ_beta);
+    ml.get_key(LLM_KV_ACTIVATION_SITU_LINEAR_BETA,  hparams.situ_linear_beta);
+
+    switch (hparams.n_layer()) {
+        case 93: type = LLM_TYPE_2_8T_A50B; break; // Kimi-K3
+        default: type = LLM_TYPE_UNKNOWN;
+    }
+}
+
+void llama_model_kimi_k3::load_arch_tensors(llama_model_loader &) {
+    LLAMA_LOAD_LOCALS;
+
+    const int64_t n_embd_latent = hparams.n_expert_latent > 0 ? hparams.n_expert_latent : n_embd;
+
+    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
+
+    output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
+    output      = create_tensor(tn(LLM_TENSOR_OUTPUT,      "weight"), {n_embd, n_vocab}, 0);
+
+    if (hparams.attn_res_block_size > 0) {
+        output_res_score = create_tensor(tn(LLM_TENSOR_OUTPUT_RES_SCORE, "weight"), {n_embd}, 0);
+    }
+
+    for (int i = 0; i < n_layer; ++i) {
+        auto & layer = layers[i];
+
+        layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+        layer.ffn_norm  = create_tensor(tn(LLM_TENSOR_FFN_NORM,  "weight", i), {n_embd}, 0);
+
+        if (hparams.attn_res_block_size > 0) {
+            layer.attn_res_score = create_tensor(tn(LLM_TENSOR_ATTN_RES_SCORE, "weight", i), {n_embd}, 0);
+            layer.ffn_res_score  = create_tensor(tn(LLM_TENSOR_FFN_RES_SCORE,  "weight", i), {n_embd}, 0);
+        }
+
+        const int64_t head_dim = hparams.n_embd_head_kda;
+        const int64_t d_conv   = hparams.ssm_d_conv;
+        const int64_t d_inner  = head_dim * n_head;
+
+        if (hparams.is_recr(i)) {
+            // conv1d may be stored 4D [d_conv, 1, d_inner, 1] or 3D (quantization drops the trailing 1)
+            auto conv = [&](llm_tensor tid) {
+                ggml_tensor * t = create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner, 1}, TENSOR_NOT_REQUIRED);
+                return t ? t : create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner}, 0);
+            };
+            layer.ssm_q_conv = conv(LLM_TENSOR_SSM_CONV1D_Q);
+            layer.ssm_k_conv = conv(LLM_TENSOR_SSM_CONV1D_K);
+            layer.ssm_v_conv = conv(LLM_TENSOR_SSM_CONV1D_V);
+
+            create_tensor_qkv(layer, i, n_embd, d_inner, d_inner, d_inner, 0);
+
+            layer.ssm_f_a  = create_tensor(tn(LLM_TENSOR_SSM_F_A,  "weight", i), {n_embd, head_dim}, 0);
+            layer.ssm_f_b  = create_tensor(tn(LLM_TENSOR_SSM_F_B,  "weight", i), {head_dim, d_inner}, 0);
+            layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0);
+
+            // K3's A_log is a plain 1-D [n_head] tensor (kimi-linear's is padded)
+            layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {n_head}, 0);
+            layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, 0);
+
+            // K3 uses a single full-rank gate instead of kimi-linear's g_a/g_b pair
+            layer.ssm_g      = create_tensor(tn(LLM_TENSOR_SSM_G,    "weight", i), {n_embd, d_inner}, 0);
+            layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_dim}, 0);
+            layer.wo         = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {d_inner, n_embd}, 0);
+        } else {
+            const int64_t q_lora_rank      = hparams.n_lora_q;
+            const int64_t kv_lora_rank     = hparams.n_lora_kv;
+            const int64_t n_embd_head_k    = hparams.n_embd_head_k_mla();
+            const int64_t n_embd_head_v    = hparams.n_embd_head_v_mla();
+            const int64_t qk_rope_head_dim = hparams.n_rot();
+            const int64_t qk_nope_head_dim = n_embd_head_k - qk_rope_head_dim;
+
+            layer.attn_q_a_norm  = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM,  "weight", i), {q_lora_rank}, TENSOR_NOT_REQUIRED);
+            layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, 0);
+
+            if (layer.attn_q_a_norm) {
+                layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0);
+                layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k}, 0);
+            } else {
+                layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_head * n_embd_head_k}, 0);
+            }
+
+            layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + qk_rope_head_dim}, 0);
+            layer.wkv_b     = create_tensor(tn(LLM_TENSOR_ATTN_KV_B, "weight", i),
+                                            {kv_lora_rank, n_head * (qk_nope_head_dim + n_embd_head_v)},
+                                            TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL);
+            if (!layer.wkv_b) {
+                layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {qk_nope_head_dim, kv_lora_rank, n_head}, 0);
+                layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v, n_head}, 0);
+            }
+
+            // K3: sigmoid output gate applied to the attention output before o_proj
+            layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head * n_embd_head_v}, TENSOR_NOT_REQUIRED);
+
+            layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v, n_embd}, 0);
+        }
+
+        if (i < (int) hparams.n_layer_dense_lead) {
+            layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
+            layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
+            layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd, n_ff}, 0);
+        } else {
+            const int64_t n_ff_exp = hparams.n_ff_exp;
+
+            layer.ffn_gate_inp    = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,    "weight", i), {n_embd, n_expert}, 0);
+            layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias",   i), {n_expert}, 0);
+
+            // routed experts live in the latent space
+            layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd_latent, n_ff_exp, n_expert}, 0);
+            layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd_latent, n_expert}, 0);
+            layer.ffn_up_exps   = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,   "weight", i), {n_embd_latent, n_ff_exp, n_expert}, 0);
+
+            if (hparams.n_expert_latent > 0) {
+                layer.ffn_routed_down = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_DOWN, "weight", i), {n_embd, n_embd_latent}, 0);
+                layer.ffn_routed_up   = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_UP,   "weight", i), {n_embd_latent, n_embd}, 0);
+                layer.ffn_routed_norm = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_NORM, "weight", i), {n_embd_latent}, TENSOR_NOT_REQUIRED);
+            }
+
+            // shared experts stay at n_embd, width = moe_intermediate_size * n_expert_shared
+            const int64_t n_ff_shexp = n_ff_exp * (hparams.n_expert_shared > 0 ? hparams.n_expert_shared : 1);
+            layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, TENSOR_NOT_REQUIRED);
+            layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, TENSOR_NOT_REQUIRED);
+            layer.ffn_up_shexp   = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,   "weight", i), {n_embd, n_ff_shexp}, TENSOR_NOT_REQUIRED);
+        }
+    }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_kimi_k3::build_arch_graph(const llm_graph_params & params) const {
+    return std::make_unique<graph>(*this, params);
+}
+
+// situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta)
+// linear_beta <= 0 disables the transform on the up branch
+static ggml_tensor * kimi_k3_situ(ggml_context * ctx0, ggml_tensor * gate, ggml_tensor * up,
+                                  float beta, float linear_beta) {
+    ggml_tensor * a = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, gate, 1.0f/beta)), beta);
+    a = ggml_mul(ctx0, a, ggml_sigmoid(ctx0, gate));
+
+    if (linear_beta > 0.0f) {
+        up = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, up, 1.0f/linear_beta)), linear_beta);
+    }
+    return ggml_mul(ctx0, a, up);
+}
+
+//
+// cross-layer residual attention
+//
+
+// layout is [n_embd, n_ckpt, n_tokens]: rms_norm reduces over ne0, dsv4_hc_pre over ne1
+// append the new checkpoint, do not re-fold the whole chain
+void llama_model_kimi_k3::graph::res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens) {
+    ggml_tensor * ckpt = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens);
+
+    resi_stack = resi_stack ? ggml_concat(ctx0, resi_stack, ckpt, 1) : ckpt;
+}
+
+ggml_tensor * llama_model_kimi_k3::graph::res_mix(ggml_tensor * cur, ggml_tensor * score_w,
+                                                  int64_t n_tokens, int il) {
+    if (!resi_stack) {
+        return cur; // layer 0: nothing banked yet
+    }
+
+    const int   n_ckpt = (int) resi_stack->ne[1];
+    const float eps    = hparams.f_norm_rms_eps;
+
+    ggml_tensor * src = resi_stack;   // [n_embd, n_ckpt, n_tokens]
+
+    // one rms_norm scores all checkpoints at once
+    // note: the scores use the normalized values, but the sum below uses the raw ones
+    ggml_tensor * sc_src = ggml_rms_norm(ctx0, src, eps);
+    sc_src = ggml_mul(ctx0, sc_src, score_w);
+    sc_src = ggml_sum_rows(ctx0, sc_src);                          // [1, n_ckpt, n_tokens]
+    sc_src = ggml_reshape_2d(ctx0, sc_src, n_ckpt, n_tokens);
+
+    // the current residual stream is scored apart, so the stack stays append-only
+    ggml_tensor * sc_cur = ggml_rms_norm(ctx0, cur, eps);
+    sc_cur = ggml_mul(ctx0, sc_cur, score_w);
+    sc_cur = ggml_sum_rows(ctx0, sc_cur);                          // [1, n_tokens]
+
+    ggml_tensor * scores = ggml_concat(ctx0, sc_src, sc_cur, 0);   // [n_ckpt+1, n_tokens]
+    ggml_tensor * probs  = ggml_soft_max(ctx0, scores);            // over ne0 = n_ckpt+1
+    cb(probs, "res_probs", il);
+
+    // split the sum: hc_pre handles the stack, a broadcast-multiply the current stream
+    ggml_tensor * p_src = ggml_cont(ctx0, ggml_view_2d(ctx0, probs, n_ckpt, n_tokens, probs->nb[1], 0));
+    ggml_tensor * p_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, probs, 1, n_tokens, probs->nb[1],
+                                                       probs->nb[0] * n_ckpt));
+
+    ggml_tensor * out = ggml_dsv4_hc_pre(ctx0, src, p_src);
+    out = ggml_add(ctx0, out, ggml_mul(ctx0, cur, p_cur));
+
+    return out;
+}
+
+llama_model_kimi_k3::graph::graph(const llama_model & model, const llm_graph_params & params) :
+    llm_build_delta_net_base(params), model(model) {
+
+    ggml_tensor * cur;
+    ggml_tensor * inpL;
+
+    inpL = build_inp_embd(model.tok_embd);
+    cb(inpL, "inp_embd", -1);
+
+    // K3 MLA is nope-only, so there is no position input
+
+    auto * inp_kv      = !hparams.is_mla() ? build_inp_mem_hybrid()   : nullptr;
+    auto * inp_k       =  hparams.is_mla() ? build_inp_mem_hybrid_k() : nullptr;
+    auto * inp_rs      =  hparams.is_mla() ? inp_k->get_recr() : inp_kv->get_recr();
+    auto * inp_attn_kv = !hparams.is_mla() ? inp_kv->get_attn() : nullptr;
+    auto * inp_attn_k  =  hparams.is_mla() ? inp_k->get_attn()  : nullptr;
+
+    ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+    const int64_t n_head_kda = hparams.n_head();
+    const int64_t head_dim   = hparams.n_embd_head_kda;
+    const int64_t d_conv     = hparams.ssm_d_conv;
+    const int64_t d_inner    = n_head_kda * head_dim;
+    const int64_t n_seqs     = ubatch.n_seqs;
+    const int64_t n_seq_tokens = ubatch.n_seq_tokens;
+
+    GGML_ASSERT(n_seqs != 0);
+    GGML_ASSERT(ubatch.equal_seqs());
+    GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);
+
+    const int64_t n_embd_head_k_mla   = hparams.n_embd_head_k_mla();
+    const int64_t n_embd_head_v_mla   = hparams.n_embd_head_v_mla();
+    const int64_t kv_lora_rank        = hparams.n_lora_kv;
+    const int64_t n_embd_head_qk_rope = hparams.n_rot();
+    const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope;
+    const float   kq_scale_mla        = 1.0f / sqrtf((float) n_embd_head_k_mla);
+
+    const uint32_t res_bs        = hparams.attn_res_block_size;
+    const bool     use_attn_res  = res_bs > 0;
+    const int64_t  n_embd_latent = hparams.n_expert_latent > 0 ? hparams.n_expert_latent : n_embd;
+
+    for (int il = 0; il < n_layer; ++il) {
+        const auto & layer = model.layers[il];
+
+        // the residual stream, banked on checkpoint layers and then restarted
+        // from the attention output alone
+        ggml_tensor * prefix_sum = inpL;
+
+        cur = use_attn_res ? res_mix(prefix_sum, layer.attn_res_score, n_tokens, il)
+                           : prefix_sum;
+
+        bool banked = false;
+        if (use_attn_res && (uint32_t) il % res_bs == 0) {
+            res_push(prefix_sum, n_embd, n_tokens);  // banks the RAW layer input, not `cur`
+            banked = true;
+        }
+
+        cur = build_norm(cur, layer.attn_norm, NULL, LLM_NORM_RMS, il);
+        cb(cur, "attn_norm", il);
+        ggml_build_forward_expand(gf, cur);
+
+        if (hparams.is_recr(il)) {
+            cur = build_kda_layer(cur, layer, inp_rs, d_conv, head_dim, n_head_kda,
+                                  d_inner, n_seq_tokens, n_seqs, il);
+        } else {
+            cur = build_mla_layer(cur, layer, inp_attn_k, inp_attn_kv,
+                                  n_embd_head_k_mla, n_embd_head_v_mla, kv_lora_rank,
+                                  n_embd_head_qk_rope, n_embd_head_qk_nope, kq_scale_mla, il);
+        }
+
+        prefix_sum = banked ? cur : ggml_add(ctx0, prefix_sum, cur);
+        cb(prefix_sum, "prefix_sum_attn", il);
+
+        cur = use_attn_res ? res_mix(prefix_sum, layer.ffn_res_score, n_tokens, il)
+                           : prefix_sum;
+
+        cur = build_norm(cur, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
+        cb(cur, "ffn_norm", il);
+
+        if ((uint32_t) il < hparams.n_layer_dense_lead) {
+            ggml_tensor * g = ggml_mul_mat(ctx0, layer.ffn_gate, cur);
+            ggml_tensor * u = ggml_mul_mat(ctx0, layer.ffn_up,   cur);
+            cur = kimi_k3_situ(ctx0, g, u, hparams.situ_beta, hparams.situ_linear_beta);
+            cur = ggml_mul_mat(ctx0, layer.ffn_down, cur);
+            cb(cur, "ffn_out", il);
+        } else {
+            cur = build_latent_moe(cur, layer, n_embd_latent, il);
+        }
+
+        prefix_sum = ggml_add(ctx0, prefix_sum, cur);
+        prefix_sum = build_cvec(prefix_sum, il);
+        cb(prefix_sum, "l_out", il);
+
+        inpL = prefix_sum;
+    }
+
+    cur = inpL;
+
+    // final mix, then narrow to the output tokens
+    if (use_attn_res) {
+        cur = res_mix(cur, model.output_res_score, n_tokens, -1);
+    }
+    if (inp_out_ids) {
+        cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+    }
+
+    cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
+    cb(cur, "result_norm", -1);
+    res->t_embd = cur;
+
+    cur = ggml_mul_mat(ctx0, model.output, cur);
+    cb(cur, "result_output", -1);
+    res->t_logits = cur;
+
+    ggml_build_forward_expand(gf, cur);
+}
+
+//
+// KDA layer
+//
+
+// causal conv1d over one of Q/K/V. `qkv` selects which third of the conv state to use
+static ggml_tensor * kimi_k3_conv1d(ggml_cgraph * gf, ggml_context * ctx0,
+                                    ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
+                                    int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
+                                    int64_t d_conv, int64_t head_dim, int64_t n_head,
+                                    int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head) {
+    const int64_t d_inner         = head_dim * n_head;
+    const int64_t conv_state_size = (d_conv - 1) * d_inner;
+    const int64_t n_embd_r_total  = 3 * conv_state_size;
+
+    ggml_tensor * conv_state_x = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs,
+        (d_conv - 1)   * ggml_element_size(conv_state_all),
+        n_embd_r_total * ggml_element_size(conv_state_all),
+        qkv * conv_state_size * ggml_element_size(conv_state_all));
+
+    ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x);
+    ggml_tensor * x_3d   = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
+    ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0);
+
+    ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
+        conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]);
+    ggml_build_forward_expand(gf,
+        ggml_cpy(ctx0, last_conv_x,
+            ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,
+                (d_conv - 1)   * ggml_element_size(conv_states_all),
+                n_embd_r_total * ggml_element_size(conv_states_all),
+                (kv_head * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
+
+    ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner);
+    ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight);
+    Xcur = ggml_reshape_2d(ctx0, Xcur, d_inner, n_tokens);
+    Xcur = ggml_silu(ctx0, Xcur);
+
+    return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head, n_seq_tokens, n_seqs);
+}
+
+ggml_tensor * llama_model_kimi_k3::graph::build_kda_layer(
+        ggml_tensor * cur, const llama_layer & layer, llm_graph_input_rs * inp_rs,
+        int64_t d_conv, int64_t head_dim, int64_t n_head_kda,
+        int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il) {
+
+    const auto * mctx_cur = inp_rs->mctx;
+    const auto   kv_head  = mctx_cur->get_head();
+
+    ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
+    ggml_tensor * conv_state_all  = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
+
+    ggml_tensor * Qcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head);
+    ggml_tensor * Kcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head);
+    ggml_tensor * Vcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head);
+    cb(Qcur, "kda_q_conv", il);
+    cb(Kcur, "kda_k_conv", il);
+    cb(Vcur, "kda_v_conv", il);
+
+    // gate_lower_bound is not a clamp - when set, it swaps the decay gate activation:
+    //   unset (kimi-linear):  g = -exp(A_log) * softplus(f_b(f_a(x)) + dt_bias)
+    //   set   (K3, -5.0):     g = lower_bound * sigmoid(exp(A_log) * (f_b(f_a(x)) + dt_bias))
+    // ssm_a holds -exp(A_log) (folded at conversion time), so exp(A_log) == -ssm_a
+    ggml_tensor * f_a = ggml_mul_mat(ctx0, layer.ssm_f_a, cur);
+    ggml_tensor * g1  = ggml_mul_mat(ctx0, layer.ssm_f_b, f_a);
+    g1 = ggml_add(ctx0, g1, layer.ssm_dt_b);
+
+    ggml_tensor * A = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head_kda, 1);
+
+    if (hparams.kda_gate_lower_bound > -INFINITY) {
+        g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens);
+        g1 = ggml_mul(ctx0, g1, A);                                    // -exp(A_log) * (...)
+        g1 = ggml_sigmoid(ctx0, ggml_scale(ctx0, g1, -1.0f));
+        g1 = ggml_scale(ctx0, g1, hparams.kda_gate_lower_bound);
+    } else {
+        g1 = ggml_softplus(ctx0, g1);
+        g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens);
+        g1 = ggml_mul(ctx0, g1, A);
+    }
+    cb(g1, "kda_g1", il);
+
+    g1 = ggml_reshape_4d(ctx0, g1, head_dim, n_head_kda, n_seq_tokens, n_seqs);
+
+    ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur);
+    beta = ggml_reshape_4d(ctx0, beta, 1, n_head_kda, n_seq_tokens, n_seqs);
+    beta = ggml_sigmoid(ctx0, beta);
+    cb(beta, "kda_beta", il);
+
+    ggml_tensor * cur_3d = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs);
+
+    ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il);
+    ggml_tensor * state = build_rs(inp_rs, ssm_states_all, hparams.n_embd_s(), n_seqs);
+    state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head_kda, n_seqs);
+
+    const float eps = hparams.f_norm_rms_eps;
+    Qcur = ggml_l2_norm(ctx0, Qcur, eps);
+    Kcur = ggml_l2_norm(ctx0, Kcur, eps);
+
+    auto attn_out = build_delta_net(Qcur, Kcur, Vcur, g1, beta, state, il);
+
+    ggml_tensor * output    = ggml_cont(ctx0, attn_out.first);
+    cb(output, "kda_scan_out", il);
+    ggml_tensor * new_state = attn_out.second;
+
+    ggml_build_forward_expand(gf,
+        ggml_cpy(ctx0, new_state,
+            ggml_view_1d(ctx0, ssm_states_all, hparams.n_embd_s() * n_seqs,
+                         kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all))));
+
+    // K3: single full-rank gate (kimi-linear factors this as g_b(g_a(x)))
+    ggml_tensor * cur_2d = ggml_reshape_2d(ctx0, cur_3d, cur_3d->ne[0], n_seq_tokens * n_seqs);
+    ggml_tensor * g2     = ggml_mul_mat(ctx0, layer.ssm_g, cur_2d);
+    g2 = ggml_reshape_3d(ctx0, g2, head_dim, n_head_kda, n_seq_tokens * n_seqs);
+
+    ggml_tensor * o      = ggml_reshape_3d(ctx0, output, head_dim, n_head_kda, n_seq_tokens * n_seqs);
+    ggml_tensor * normed = build_norm(o, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il);
+    cb(g2, "kda_g2", il);
+    cb(normed, "kda_normed", il);
+    ggml_tensor * gated  = ggml_mul(ctx0, normed, ggml_sigmoid(ctx0, g2));
+
+    gated = ggml_cont_2d(ctx0, gated, d_inner, n_tokens);
+    cur   = ggml_mul_mat(ctx0, layer.wo, gated);
+    cb(cur, "kda_out", il);
+
+    return cur;
+}
+
+//
+// MLA layer (nope-only, with K3's sigmoid output gate)
+//
+
+ggml_tensor * llama_model_kimi_k3::graph::build_mla_layer(
+        ggml_tensor * cur, const llama_layer & layer,
+        llm_graph_input_attn_k * inp_attn_k, llm_graph_input_attn_kv * inp_attn_kv,
+        int64_t n_embd_head_k_mla, int64_t n_embd_head_v_mla, int64_t kv_lora_rank,
+        int64_t n_embd_head_qk_rope, int64_t n_embd_head_qk_nope, float kq_scale, int il) {
+
+    ggml_tensor * inp_gate = cur; // the output gate reads the *normed* layer input
+
+    ggml_tensor * Qcur;
+    if (layer.wq_a) {
+        Qcur = ggml_mul_mat(ctx0, layer.wq_a, cur);
+        Qcur = build_norm(Qcur, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
+        Qcur = ggml_mul_mat(ctx0, layer.wq_b, Qcur);
+    } else {
+        Qcur = ggml_mul_mat(ctx0, layer.wq, cur);
+    }
+
+    ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);
+
+    ggml_tensor * kv_cmpr = ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,
+        ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0);
+    ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens,
+        ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),
+        ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),
+        ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));
+
+    // no RoPE: mla_use_nope is asserted at conversion time
+    kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
+
+    ggml_tensor * out;
+    if (layer.wk_b && layer.wv_b) {
+        ggml_tensor * q_nope = ggml_view_3d(ctx0, Qcur, n_embd_head_qk_nope, n_head, n_tokens,
+            ggml_row_size(Qcur->type, n_embd_head_k_mla),
+            ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head, 0);
+        ggml_tensor * q_pe = ggml_view_3d(ctx0, Qcur, n_embd_head_qk_rope, n_head, n_tokens,
+            ggml_row_size(Qcur->type, n_embd_head_k_mla),
+            ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head,
+            ggml_row_size(Qcur->type, n_embd_head_qk_nope));
+
+        q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
+        ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope);
+        q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);
+
+        ggml_tensor * Q = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);
+        ggml_tensor * kv_cmpr_3d = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);
+        ggml_tensor * K = ggml_concat(ctx0, kv_cmpr_3d, k_pe, 0);
+        ggml_tensor * V = kv_cmpr_3d;
+
+        // wo == NULL: the output projection is applied after the gate below
+        out = build_attn(inp_attn_k, nullptr, NULL, nullptr, Q, K, V, nullptr, nullptr, layer.wv_b, kq_scale, il);
+    } else {
+        ggml_tensor * Q = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k_mla, n_head, n_tokens);
+        ggml_tensor * kv = ggml_mul_mat(ctx0, layer.wkv_b, kv_cmpr);
+        const int64_t kv_per_head = n_embd_head_qk_nope + n_embd_head_v_mla;
+
+        ggml_tensor * k_nope = ggml_view_3d(ctx0, kv, n_embd_head_qk_nope, n_head, n_tokens,
+            ggml_row_size(kv->type, kv_per_head), ggml_row_size(kv->type, kv_per_head * n_head), 0);
+        ggml_tensor * V = ggml_cont(ctx0, ggml_view_3d(ctx0, kv, n_embd_head_v_mla, n_head, n_tokens,
+            ggml_row_size(kv->type, kv_per_head), ggml_row_size(kv->type, kv_per_head * n_head),
+            ggml_row_size(kv->type, n_embd_head_qk_nope)));
+
+        ggml_tensor * k_pe_t = ggml_new_tensor_3d(ctx0, k_pe->type, n_embd_head_qk_rope, n_head, n_tokens);
+        ggml_tensor * K = ggml_concat(ctx0, ggml_repeat(ctx0, k_pe, k_pe_t), k_nope, 0);
+
+        out = build_attn(inp_attn_kv, nullptr, NULL, nullptr, Q, K, V, nullptr, nullptr, nullptr, kq_scale, il);
+    }
+
+    // K3: attn_output *= sigmoid(g_proj(x)), then o_proj
+    if (layer.wqkv_gate) {
+        ggml_tensor * g = ggml_sigmoid(ctx0, ggml_mul_mat(ctx0, layer.wqkv_gate, inp_gate));
+        out = ggml_mul(ctx0, out, g);
+        cb(out, "mla_gated", il);
+    }
+
+    out = ggml_mul_mat(ctx0, layer.wo, out);
+    cb(out, "mla_out", il);
+
+    return out;
+}
+
+//
+// latent MoE: down-project, run the routed experts in the latent space, norm, up-project;
+// shared experts stay at n_embd and read the un-projected input.
+//
+
+ggml_tensor * llama_model_kimi_k3::graph::build_latent_moe(
+        ggml_tensor * cur, const llama_layer & layer, int64_t n_embd_latent, int il) {
+
+    ggml_tensor * identity = cur;
+
+    ggml_tensor * routed_in = layer.ffn_routed_down
+        ? ggml_mul_mat(ctx0, layer.ffn_routed_down, cur)
+        : cur;
+
+    // the router scores the full-width input while the experts take the latent one,
+    // so the logits are computed here and passed to build_moe_ffn
+    ggml_tensor * logits = ggml_mul_mat(ctx0, layer.ffn_gate_inp, identity);
+    cb(logits, "ffn_moe_logits", il);
+
+    ggml_tensor * moe_out = build_moe_ffn(routed_in,
+        nullptr, // gate_inp unused: the logits above are passed instead
+        layer.ffn_up_exps,
+        layer.ffn_gate_exps,
+        layer.ffn_down_exps,
+        layer.ffn_exp_probs_b,
+        hparams.n_expert,
+        hparams.n_expert_used,
+        LLM_FFN_SITU, hparams.expert_weights_norm,
+        hparams.expert_weights_scale,
+        (llama_expert_gating_func_type) hparams.expert_gating_func,
+        il,
+        logits);
+    cb(moe_out, "ffn_moe_out", il);
+
+    if (layer.ffn_routed_norm) {
+        moe_out = build_norm(moe_out, layer.ffn_routed_norm, NULL, LLM_NORM_RMS, il);
+    }
+    if (layer.ffn_routed_up) {
+        moe_out = ggml_mul_mat(ctx0, layer.ffn_routed_up, moe_out);
+    }
+    GGML_UNUSED(n_embd_latent);
+
+    if (layer.ffn_gate_shexp) {
+        ggml_tensor * g = ggml_mul_mat(ctx0, layer.ffn_gate_shexp, identity);
+        ggml_tensor * u = ggml_mul_mat(ctx0, layer.ffn_up_shexp,   identity);
+        ggml_tensor * sh = kimi_k3_situ(ctx0, g, u, hparams.situ_beta, hparams.situ_linear_beta);
+        sh = ggml_mul_mat(ctx0, layer.ffn_down_shexp, sh);
+        cb(sh, "ffn_shexp", il);
+        moe_out = ggml_add(ctx0, moe_out, sh);
+    }
+
+    cb(moe_out, "ffn_out", il);
+    return moe_out;
+}

src/models/mamba-base.cpp

@@ -2,6 +2,8 @@
 
 #include "llama-memory-recurrent.h"
 
+#include <algorithm>
+
 llm_build_mamba_base::llm_build_mamba_base(const llm_graph_params & params) : llm_graph_context(params) {}
 
 ggml_tensor * llm_build_mamba_base::build_mamba_layer(llm_graph_input_rs * inp,
@@ -118,7 +120,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba_layer(llm_graph_input_rs * inp,
             // Custom operator to optimize the parallel associative scan
             // as described in the Annex D of the Mamba paper.
             // => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}
-            return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids);
+            return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, /*K=*/1);
         };
 
         ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows);
@@ -153,7 +155,8 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
                                                           int                  il) const {
     const auto * mctx_cur = inp->mctx;
 
-    const auto kv_head = mctx_cur->get_head();
+    const auto kv_head  = mctx_cur->get_head();
+    const auto mem_size = mctx_cur->get_size();
 
     const int64_t d_conv   = hparams.ssm_d_conv;
     const int64_t d_inner  = hparams.ssm_d_inner;
@@ -164,6 +167,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
     const int64_t n_seqs   = ubatch.n_seqs;
 
     const int64_t n_seq_tokens = ubatch.n_seq_tokens;
+    const int64_t K            = cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1;
 
     GGML_ASSERT(n_seqs != 0);
     GGML_ASSERT(ubatch.equal_seqs());
@@ -173,6 +177,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
 
     ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
     ggml_tensor * ssm_states_all  = mctx_cur->get_s_l(il);
+    const int64_t state_slots     = ssm_states_all->ne[1];
 
     ggml_tensor * conv = build_rs(inp, conv_states_all, hparams.n_embd_r(), n_seqs);
     conv               = ggml_reshape_3d(ctx0, conv, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs);
@@ -198,15 +203,19 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
         // => {d_conv - 1 + n_seq_tokens, d_inner + 2*n_group*d_state, n_seqs}
         ggml_tensor * conv_x = ggml_concat(ctx0, conv, ggml_transpose(ctx0, xBC), 0);
 
-        // copy last (d_conv - 1) columns back into the state cache
-        ggml_tensor * last_conv = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs,
-                                               conv_x->nb[1], conv_x->nb[2], n_seq_tokens * (conv_x->nb[0]));
+        const int64_t row_count = (d_conv - 1) * (d_inner + 2 * n_group * d_state);
+        const size_t  row_size  = ggml_row_size(conv_states_all->type, row_count);
+        const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
 
-        ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv,
-                                               ggml_view_1d(ctx0, conv_states_all,
-                                                            (d_conv - 1) * (d_inner + 2 * n_group * d_state) * (n_seqs),
-                                                            kv_head * (d_conv - 1) * (d_inner + 2 * n_group * d_state) *
-                                                                ggml_element_size(conv_states_all))));
+        for (int64_t slot = 0; slot < n_written; ++slot) {
+            ggml_tensor * last_conv = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs,
+                                                   conv_x->nb[1], conv_x->nb[2], (n_seq_tokens - slot) * conv_x->nb[0]);
+
+            ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv,
+                                                   ggml_view_2d(ctx0, conv_states_all, row_count, n_seqs,
+                                                                conv_states_all->nb[1],
+                                                                ((size_t) slot * mem_size + kv_head) * row_size)));
+        }
 
         // 1D convolution
         // The equivalent is to make a self-overlapping view of conv_x
@@ -244,20 +253,27 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
         // (this is necessary in order to properly use the states before they are overwritten,
         //  while avoiding to make unnecessary copies of the states)
         auto get_ssm_rows = [&](ggml_context * ctx, ggml_tensor * states, ggml_tensor * ids) {
-            ggml_tensor * ssm = ggml_reshape_4d(ctx, states, d_state, head_dim, n_head, mctx_cur->get_size());
+            ggml_tensor * ssm = ggml_reshape_4d(ctx, states, d_state, head_dim, n_head, state_slots);
 
             // TODO: use semistructured matrices to implement state-space duality
             // => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}
-            return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids);
+            // K > 1 asks the backend to return rollback snapshots in addition to the final state.
+            return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, K);
         };
 
         ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows);
+        const int64_t D            = d_state * d_inner;
+        const int64_t n_written    = std::min<int64_t>(n_seq_tokens, K);
+        const size_t  row_size     = ggml_row_size(ssm_states_all->type, D);
+        const size_t  y_row_size   = ggml_row_size(y_ssm->type, D);
+        const size_t  state_offset = ggml_nelements(x) * ggml_element_size(x);
 
-        // store last states
         ggml_build_forward_expand(
-            gf, ggml_cpy(ctx0, ggml_view_1d(ctx0, y_ssm, d_state * d_inner * n_seqs, ggml_nelements(x) * x->nb[0]),
-                         ggml_view_1d(ctx0, ssm_states_all, d_state * d_inner * n_seqs,
-                                      kv_head * d_state * d_inner * ggml_element_size(ssm_states_all))));
+            gf, ggml_cpy(ctx0,
+                         ggml_view_3d(ctx0, y_ssm, D, n_seqs, n_written,
+                                      y_row_size, y_row_size * n_seqs, state_offset),
+                         ggml_view_3d(ctx0, ssm_states_all, D, n_seqs, n_written,
+                                      ssm_states_all->nb[1], (size_t) mem_size * row_size, kv_head * row_size)));
 
         ggml_tensor * y = ggml_view_4d(ctx0, y_ssm, head_dim, n_head, n_seq_tokens, n_seqs, x->nb[1], n_head * x->nb[1],
                                        n_seq_tokens * n_head * x->nb[1], 0);

src/models/minimax-01.cpp

@@ -0,0 +1,520 @@
+#include "models.h"
+#include "llama-memory-recurrent.h"
+
+void llama_model_minimax_01::load_arch_hparams(llama_model_loader & ml) {
+    ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+    ml.get_key(LLM_KV_RESIDUAL_SCALE,              hparams.f_residual_scale);
+
+    // we use n_embd_head_la to set recurrent memory n_embd_s
+    hparams.n_embd_head_la = hparams.n_embd_head_k_full;
+
+    // Mark recurrent layers (lightning attention layers).
+    if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) {
+        uint32_t full_attn_interval = 8;
+        ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false);
+        for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {
+            hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0);
+        }
+    }
+
+    switch (hparams.n_layer()) {
+        case 80: type = LLM_TYPE_456B; break;
+        default: type = LLM_TYPE_UNKNOWN;
+    }
+}
+
+void llama_model_minimax_01::load_arch_tensors(llama_model_loader &) {
+    LLAMA_LOAD_LOCALS;
+
+    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
+
+    // output
+    output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
+    output      = create_tensor(tn(LLM_TENSOR_OUTPUT,      "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
+
+    // if output is NULL, init from the input tok embed
+    if (output == NULL) {
+        output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
+    }
+
+    for (int i = 0; i < n_layer; ++i) {
+        auto & layer = layers[i];
+
+        layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+
+        if (!hparams.is_recr(i)) {
+            create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);
+        } else {
+            layer.attn_norm_2 = create_tensor(tn(LLM_TENSOR_ATTN_NORM_2, "weight", i), {n_embd_head_k * n_head}, 0);
+            layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, 3 * n_embd_head_k * n_head}, 0);
+            layer.wg = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_embd_head_k * n_head}, 0);
+        }
+        layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);
+
+        layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+
+        layer.ffn_gate_inp  = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,  "weight", i), {n_embd, n_expert}, 0);
+        layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd,   n_ff, n_expert}, TENSOR_NOT_REQUIRED);
+        layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {  n_ff, n_embd, n_expert}, 0);
+        layer.ffn_up_exps   = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,   "weight", i), {n_embd,   n_ff, n_expert}, 0);
+    }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_minimax_01::build_arch_graph(const llm_graph_params & params) const {
+    return std::make_unique<graph>(*this, params);
+}
+
+class llm_graph_input_la : public llm_graph_input_i {
+public:
+    llm_graph_input_la(const llama_hparams & hparams) : hparams(hparams) {}
+
+    void set_input(const llama_ubatch * ubatch) override {
+        // this operates on assumption that we have an equal ubatch split
+
+        const int64_t n_head = hparams.n_head();
+        const int32_t n_seqs = ubatch->n_seqs;
+        const int32_t n_seqs_unq = ubatch->n_seqs_unq;
+        const int32_t n_tokens = ubatch->n_tokens;
+        const int32_t n_seq_tokens = ubatch->n_seq_tokens;
+
+        std::vector<llama_pos> p0(n_seqs_unq);
+        std::fill(p0.begin(), p0.end(), std::numeric_limits<llama_pos>::max());
+
+        // get lowest token position in a ubatch for each stream
+        for (int i = 0; i < n_tokens; ++i) {
+            llama_seq_id seq_id = ubatch->seq_id[i][0];
+            int32_t seq_idx = ubatch->seq_idx[seq_id];
+            llama_pos pos = ubatch->pos[i];
+            if (p0[seq_idx] > pos) {
+                p0[seq_idx] = pos;
+            }
+        }
+
+        if (inp_slopes) {
+            GGML_ASSERT(ggml_backend_buffer_is_host(inp_slopes->buffer));
+
+            float * data = (float *) inp_slopes->data;
+
+            float start = powf(2, -powf(2, -(log2f(n_head) - 3)));
+            float ratio = start;
+
+            for (int h = 0; h < n_head; ++h) {
+                data[h] = start * powf(ratio, h);
+            }
+        }
+
+        if (inp_q_decay) {
+            GGML_ASSERT(ggml_backend_buffer_is_host(inp_q_decay->buffer));
+
+            float * slopes = (float *) inp_slopes->data;
+            float * data = (float *) inp_q_decay->data;
+
+            for (int s = 0; s < n_seqs; ++s) {
+                for (int i = 0; i < n_seq_tokens; ++i) {
+                    llama_seq_id seq_id = ubatch->seq_id[s * n_seq_tokens + i][0];
+                    int32_t seq_idx = ubatch->seq_idx[seq_id];
+                    llama_pos pos = ubatch->pos[s * n_seq_tokens + i];
+                    int pos_rel = pos - p0[seq_idx];
+
+                    for (int h = 0; h < n_head; ++h) {
+                        data[seq_idx * n_head * n_seq_tokens + i * n_head + h] = -slopes[h] * (pos_rel + 1);
+                    }
+                }
+            }
+        }
+
+        if (inp_k_decay) {
+            GGML_ASSERT(ggml_backend_buffer_is_host(inp_k_decay->buffer));
+
+            float * slopes = (float *) inp_slopes->data;
+            float * data = (float *) inp_k_decay->data;
+
+            for (int s = 0; s < n_seqs; ++s) {
+                for (int i = 0; i < n_seq_tokens; ++i) {
+                    llama_seq_id seq_id = ubatch->seq_id[s * n_seq_tokens + i][0];
+                    int32_t seq_idx = ubatch->seq_idx[seq_id];
+                    llama_pos pos = ubatch->pos[s * n_seq_tokens + i];
+                    int pos_rel = pos - p0[seq_idx];
+
+                    for (int h = 0; h < n_head; ++h) {
+                        data[seq_idx * n_head * n_seq_tokens + i * n_head + h] = -slopes[h] * (n_seq_tokens - pos_rel - 1);
+                    }
+                }
+            }
+        }
+
+        if (inp_diag_decay) {
+            GGML_ASSERT(ggml_backend_buffer_is_host(inp_diag_decay->buffer));
+
+            float * slopes = (float *) inp_slopes->data;
+            float * data = (float *) inp_diag_decay->data;
+
+            for (int s = 0; s < n_seqs; ++s) {
+                for (int h = 0; h < n_head; ++h) {
+                    for (int j = 0; j < n_seq_tokens; ++j) {
+                        llama_seq_id seq_id = ubatch->seq_id[s * n_seq_tokens + j][0];
+                        int32_t seq_idx = ubatch->seq_idx[seq_id];
+                        llama_pos pos_j = ubatch->pos[s * n_seq_tokens + j];
+                        int pos_rel_j = pos_j - p0[seq_idx];
+
+                        for (int i = 0; i < n_seq_tokens; ++i) {
+                            llama_pos pos_i = ubatch->pos[s * n_seq_tokens + i];
+                            int pos_rel_i = pos_i - p0[seq_idx];
+
+                            int index = pos_rel_j - pos_rel_i;
+                            float s_index = index >= 0 ? -slopes[h] * index : -INFINITY;
+                            data[seq_idx * n_head * n_seq_tokens * n_seq_tokens + h * n_seq_tokens * n_seq_tokens + j * n_seq_tokens + i] = s_index;
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    bool can_reuse(const llm_graph_params & params) override {
+        bool res = true;
+
+        if (params.ubatch.n_seq_tokens > 1) {
+            res &= (   inp_q_decay &&    inp_q_decay->ne[2] == params.ubatch.n_seq_tokens);
+            res &= (   inp_k_decay &&    inp_k_decay->ne[2] == params.ubatch.n_seq_tokens);
+            res &= (inp_diag_decay && inp_diag_decay->ne[1] == params.ubatch.n_seq_tokens);
+        }
+
+        return res;
+    }
+
+    const llama_hparams & hparams;
+
+    ggml_tensor * inp_slopes     = nullptr; // F32 [n_head]
+    ggml_tensor * inp_q_decay    = nullptr; // F32 [1, n_head, n_batch]
+    ggml_tensor * inp_k_decay    = nullptr; // F32 [1, n_head, n_batch]
+    ggml_tensor * inp_diag_decay = nullptr; // F32 [n_batch, n_batch, n_head]
+};
+
+llama_model_minimax_01::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
+    const int64_t n_embd_head = hparams.n_embd_head_v();
+
+    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
+    // GGML_ASSERT(n_embd_head == n_rot); this is wrong in case of minimax, head_dim = 128, n_rot = 64
+
+    const int64_t n_seqs  = ubatch.n_seqs;
+    const int64_t n_seq_tokens = ubatch.n_seq_tokens;
+
+    GGML_ASSERT(n_seqs != 0);
+    GGML_ASSERT(ubatch.equal_seqs());
+    GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);
+
+    ggml_tensor * cur;
+    ggml_tensor * inpL;
+
+    inpL = build_inp_embd(model.tok_embd);
+
+    auto * inp_hybrid = build_inp_mem_hybrid();
+    auto * inp_rs = inp_hybrid->get_recr();
+
+    ggml_tensor * inp_pos = build_inp_pos();
+    ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+    llm_graph_input_la * la = nullptr;
+
+    auto inp = std::make_unique<llm_graph_input_la>(hparams);
+
+    inp->inp_slopes = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_head);
+    ggml_set_input(inp->inp_slopes);
+    cb(inp->inp_slopes, "slopes", -1);
+
+    if (n_seq_tokens != 1) {
+        inp->inp_q_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs);
+        ggml_set_input(inp->inp_q_decay);
+        cb(inp->inp_q_decay, "q_decay_exp", -1);
+
+        inp->inp_k_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs);
+        ggml_set_input(inp->inp_k_decay);
+        cb(inp->inp_k_decay, "k_decay_exp", -1);
+
+        inp->inp_diag_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_seq_tokens, n_seq_tokens, n_head, n_seqs);
+        ggml_set_input(inp->inp_diag_decay);
+        cb(inp->inp_diag_decay, "diag_decay_exp", -1);
+    }
+
+    la = (llm_graph_input_la *) res->add_input(std::move(inp));
+
+    ggml_tensor * slopes = la->inp_slopes;
+
+    for (int il = 0; il < n_layer; ++il) {
+        res->t_layer_inp[il] = inpL;
+
+        ggml_tensor * inpSA = inpL;
+
+        cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
+        cb(cur, "attn_norm", il);
+
+        ggml_tensor * residual = cur;
+
+        // self_attention
+        if (!hparams.is_recr(il)) {
+            // softmax attention layer
+
+            auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
+                    n_embd_head, n_head, n_head_kv, il);
+
+            Qcur = ggml_rope_ext(
+                ctx0, Qcur, inp_pos, nullptr,
+                n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+                ext_factor, attn_factor, beta_fast, beta_slow
+                );
+
+            Kcur = ggml_rope_ext(
+                ctx0, Kcur, inp_pos, nullptr,
+                n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+                ext_factor, attn_factor, beta_fast, beta_slow
+                );
+
+            cb(Qcur, "Qcur", il);
+            cb(Kcur, "Kcur", il);
+            cb(Vcur, "Vcur", il);
+
+            cur = build_attn(inp_hybrid->get_attn(),
+                    model.layers[il].wo, NULL, model.layers[il].wo_s,
+                    Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);
+        } else {
+            // lightning attention layer
+
+            const auto * mctx_cur = inp_rs->mctx;
+            const auto kv_head = mctx_cur->get_head();
+
+            // TODO unneeded - any way to make conv states optional in recurrent memory?
+            ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
+            ggml_tensor * conv_state_all  = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
+            ggml_build_forward_expand(gf, conv_state_all);
+
+            float slope_scale = 1.0 - 1.0 * il / (n_layer - 1) + 1e-5;
+            ggml_tensor * slope_rate = ggml_scale(ctx0, slopes, slope_scale);
+            cb(slope_rate, "slope_rate", il);
+
+            cur = ggml_reshape_4d(ctx0, cur, cur->ne[0], n_seq_tokens, 1, n_seqs);
+
+            ggml_tensor * QKVcur = build_lora_mm(model.layers[il].wqkv, cur);
+            cb(QKVcur, "QKVcur", il);
+
+            QKVcur = ggml_silu(ctx0, QKVcur);
+            cb(QKVcur, "QKVcur_silu", il);
+
+            QKVcur = ggml_reshape_4d(ctx0, QKVcur, n_embd_head * 3, n_head, n_seq_tokens, n_seqs);
+
+            ggml_tensor * Qcur = ggml_view_4d(ctx0, QKVcur, n_embd_head, n_head, n_seq_tokens, n_seqs, QKVcur->nb[1], QKVcur->nb[2], QKVcur->nb[3], 0*ggml_element_size(QKVcur)*n_embd_head);
+            ggml_tensor * Kcur = ggml_view_4d(ctx0, QKVcur, n_embd_head, n_head, n_seq_tokens, n_seqs, QKVcur->nb[1], QKVcur->nb[2], QKVcur->nb[3], 1*ggml_element_size(QKVcur)*n_embd_head);
+            ggml_tensor * Vcur = ggml_view_4d(ctx0, QKVcur, n_embd_head, n_head, n_seq_tokens, n_seqs, QKVcur->nb[1], QKVcur->nb[2], QKVcur->nb[3], 2*ggml_element_size(QKVcur)*n_embd_head);
+
+            cb(Qcur, "Qcur", il);
+            cb(Kcur, "Kcur", il);
+            cb(Vcur, "Vcur", il);
+
+            // get previous KV
+            ggml_tensor * la_states_all = mctx_cur->get_s_l(il);
+            ggml_tensor * state = build_rs(inp_rs, la_states_all, hparams.n_embd_s(), n_seqs);
+
+            ggml_tensor * kv_old = ggml_reshape_4d(ctx0, state, n_embd_head, n_embd_head, n_head, n_seqs);
+            cb(kv_old, "kv_old", il);
+
+            ggml_tensor * qkv = nullptr;
+            ggml_tensor * kv_new = nullptr;
+
+            if (n_seq_tokens == 1) {
+                // lightning attention - optimized single token case for TG
+
+                ggml_tensor * slopes_neg = ggml_scale(ctx0, slope_rate, -1.0);
+                cb(slopes_neg, "slopes_neg", il);
+
+                ggml_tensor * ratio = ggml_exp(ctx0, slopes_neg);
+                cb(ratio, "ratio", il);
+
+                ggml_tensor * ratio_3d = ggml_reshape_3d(ctx0, ratio, 1, 1, n_head);
+                cb(ratio_3d, "ratio3d", il);
+
+                ggml_tensor * v_trans = ggml_cont(ctx0, ggml_permute(ctx0, Vcur, 1, 2, 0, 3));
+                cb(v_trans, "v_trans", il);
+
+                ggml_tensor * k_trans = ggml_cont(ctx0, ggml_permute(ctx0, Kcur, 1, 2, 0, 3));
+                cb(k_trans, "k_trans", il);
+
+                ggml_tensor * kv_cur = ggml_mul_mat(ctx0, k_trans, v_trans);
+                cb(kv_cur, "kv_cur", il);
+
+                ggml_tensor * kv_old_s = ggml_mul(ctx0, kv_old, ratio_3d);
+                cb(kv_old_s, "kv_old_s", il);
+
+                kv_new = ggml_add(ctx0, kv_old_s, kv_cur);
+                cb(kv_new, "kv_new", il);
+
+                ggml_tensor * q_trans = ggml_permute(ctx0, Qcur, 0, 2, 1, 3);
+                cb(q_trans, "q_trans", il);
+
+                qkv = ggml_mul_mat(ctx0, kv_new, q_trans);
+                cb(qkv, "qkv", il);
+            } else if(n_seq_tokens > 1) {
+                // lightning attention - general multi token case for PP
+
+                ggml_tensor *    q_decay_exp = la->inp_q_decay;
+                ggml_tensor *    k_decay_exp = la->inp_k_decay;
+                ggml_tensor * diag_decay_exp = la->inp_diag_decay;
+
+                ggml_tensor *    q_decay = ggml_exp(ctx0, ggml_scale(ctx0, q_decay_exp, slope_scale));
+                cb(q_decay, "q_decay", il);
+                ggml_tensor *    k_decay = ggml_exp(ctx0, ggml_scale(ctx0, k_decay_exp, slope_scale));
+                cb(k_decay, "k_decay", il);
+                ggml_tensor * diag_decay = ggml_exp(ctx0, ggml_scale(ctx0, diag_decay_exp, slope_scale));
+                cb(diag_decay, "diag_decay", il);
+
+                ggml_tensor * q_s = ggml_mul(ctx0, Qcur, q_decay);
+                cb(q_s, "q_s", il);
+
+                ggml_tensor * q_s_trans = ggml_permute(ctx0, q_s, 0, 2, 1, 3);
+                cb(q_s_trans, "q_s_trans", il);
+
+                ggml_tensor * qkv_none_diag = ggml_mul_mat(ctx0, kv_old, q_s_trans);
+                cb(qkv_none_diag, "qkv_none_diag", il);
+
+                ggml_tensor * q_trans = ggml_permute(ctx0, Qcur, 0, 2, 1, 3);
+                cb(q_trans, "q_trans", il);
+
+                ggml_tensor * k_trans = ggml_permute(ctx0, Kcur, 0, 2, 1, 3);
+                cb(k_trans, "k_trans", il);
+
+                ggml_tensor * qk = ggml_mul_mat(ctx0, k_trans, q_trans);
+                cb(qk, "qk", il);
+
+                qk = ggml_mul(ctx0, qk, diag_decay);
+                cb(qk, "qk_s", il);
+
+                ggml_tensor * v_trans = ggml_cont(ctx0, ggml_permute(ctx0, Vcur, 1, 2, 0, 3));
+                cb(v_trans, "v_trans", il);
+
+                ggml_tensor * qkv_diag = ggml_mul_mat(ctx0, v_trans, qk);
+                cb(qkv_diag, "qkv_diag", il);
+
+                qkv = ggml_add(ctx0, qkv_none_diag, qkv_diag);
+                cb(qkv, "qkv", il);
+
+                ggml_build_forward_expand(gf, qkv);
+
+                ggml_tensor * slopes_neg = ggml_scale(ctx0, slope_rate, -1.0*n_seq_tokens);
+                cb(slopes_neg, "slopes_neg", il);
+
+                ggml_tensor * block_decay = ggml_exp(ctx0, slopes_neg);
+                cb(block_decay, "block_decay", il);
+
+                ggml_tensor * block_decay_3d = ggml_reshape_3d(ctx0, block_decay, 1, 1, n_head);
+                cb(block_decay_3d, "block_decay_3d", il);
+
+                ggml_tensor * kv_old_s = ggml_mul(ctx0, kv_old, block_decay_3d);
+                cb(kv_old_s, "kv_old_s", il);
+
+                ggml_tensor * k_after_decay = ggml_mul(ctx0, Kcur, k_decay);
+                cb(k_after_decay, "k_after_decay", il);
+
+                ggml_tensor * k_after_decay_trans = ggml_cont(ctx0, ggml_permute(ctx0, k_after_decay, 1, 2, 0, 3));
+                cb(k_after_decay_trans, "k_after_decay_trans", il);
+
+                ggml_tensor * kv_cur = ggml_mul_mat(ctx0, k_after_decay_trans, v_trans);
+                cb(kv_cur, "kv_cur", il);
+
+                kv_new = ggml_add(ctx0, kv_old_s, kv_cur);
+                cb(kv_new, "kv_new", il);
+            }
+
+            // store new KV
+            ggml_build_forward_expand(gf,
+                                     ggml_cpy(ctx0, kv_new,
+                                              ggml_view_1d(ctx0, la_states_all, hparams.n_embd_s() * n_seqs,
+                                                           kv_head * hparams.n_embd_s() * ggml_element_size(la_states_all))));
+
+            qkv = ggml_cont(ctx0, ggml_permute(ctx0, qkv, 0, 2, 1, 3));
+            cb(qkv, "qkv_permuted", il);
+
+            qkv = ggml_reshape_4d(ctx0, qkv, qkv->ne[0]*qkv->ne[1], qkv->ne[2], 1, qkv->ne[3]);
+
+            // norm
+            ggml_tensor * qkv_norm = build_norm(qkv,
+                    model.layers[il].attn_norm_2, NULL,
+                    LLM_NORM_RMS, il);
+            cb(qkv_norm, "qkv_norm", il);
+
+            ggml_tensor * g = build_lora_mm(model.layers[il].wg, cur);
+            cb(g, "g", il);
+
+            g = ggml_sigmoid(ctx0, g);
+            cb(g, "g_sigm", il);
+
+            cur = ggml_mul(ctx0, g, qkv_norm);
+
+            cur = build_lora_mm(model.layers[il].wo, cur);
+            cb(cur, "attn_out", il);
+
+            cur = ggml_reshape_2d(ctx0, cur, cur->ne[0], n_seq_tokens*n_seqs);
+            cb(cur, "attn_out", il);
+        }
+
+        if (il == n_layer - 1 && inp_out_ids) {
+            cur   = ggml_get_rows(ctx0,   cur, inp_out_ids);
+            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
+            residual = ggml_get_rows(ctx0, residual, inp_out_ids);
+        }
+
+        residual = ggml_scale(ctx0, residual, hparams.f_residual_scale);
+        cb(residual, "residual_scaled_attn", il);
+
+        ggml_tensor * ffn_inp = ggml_add(ctx0, cur, residual);
+        cb(ffn_inp, "ffn_inp", il);
+
+        // MoE branch
+        cur = build_norm(ffn_inp,
+                model.layers[il].ffn_norm, NULL,
+                LLM_NORM_RMS, il);
+        cb(cur, "ffn_norm", il);
+
+        residual = cur;
+
+        cur = build_moe_ffn(cur,
+                model.layers[il].ffn_gate_inp,
+                model.layers[il].ffn_up_exps,
+                model.layers[il].ffn_gate_exps,
+                model.layers[il].ffn_down_exps,
+                model.layers[il].ffn_exp_probs_b,
+                n_expert, n_expert_used,
+                LLM_FFN_SILU, true,
+                hparams.expert_weights_scale,
+                LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX,
+                il);
+        cb(cur, "ffn_moe_out", il);
+
+        residual = ggml_scale(ctx0, residual, hparams.f_residual_scale);
+        cb(residual, "residual_scaled_ffn", il);
+
+        cur = ggml_add(ctx0, cur, residual);
+        cb(cur, "ffn_out", il);
+
+        cur = build_cvec(cur, il);
+        cb(cur, "l_out", il);
+
+        // input for next layer
+        inpL = cur;
+    }
+
+    cur = inpL;
+
+    cur = build_norm(cur,
+            model.output_norm, NULL,
+            LLM_NORM_RMS, -1);
+
+    cb(cur, "result_norm", -1);
+    res->t_embd = cur;
+
+    // lm_head
+    cur = build_lora_mm(model.output, cur, model.output_s);
+
+    cb(cur, "result_output", -1);
+    res->t_logits = cur;
+
+    ggml_build_forward_expand(gf, cur);
+}

src/models/minimax-m3.cpp

@@ -25,6 +25,8 @@ void llama_model_minimax_m3::load_arch_hparams(llama_model_loader & ml) {
     ml.get_key(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS,  hparams.indexer_local_blocks);
     msa_p = { (int) hparams.indexer_block_size, (int) hparams.indexer_top_k, (int) hparams.indexer_local_blocks };
 
+    GGML_ASSERT(hparams.indexer_block_size > 0); // avoid div by zero
+
     switch (hparams.n_layer()) {
         case 60: type = LLM_TYPE_428B_A23B; break;
         default: type = LLM_TYPE_UNKNOWN;

src/models/models.h

@@ -2043,6 +2043,19 @@ struct llama_model_apertus : public llama_model_base {
 };
 
 
+struct llama_model_minimax_01 : public llama_model_base {
+    llama_model_minimax_01(const struct llama_model_params & params) : llama_model_base(params) {}
+    void load_arch_hparams(llama_model_loader & ml) override;
+    void load_arch_tensors(llama_model_loader & ml) override;
+
+    struct graph : public llm_graph_context {
+        graph(const llama_model & model, const llm_graph_params & params);
+    };
+
+    std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
+};
+
+
 struct llama_model_minimax_m2 : public llama_model_base {
     llama_model_minimax_m2(const struct llama_model_params & params) : llama_model_base(params) {}
     void load_arch_hparams(llama_model_loader & ml) override;
@@ -2272,6 +2285,42 @@ struct llama_model_mimo2 : public llama_model_base {
 };
 
 
+struct llama_model_kimi_k3 : public llama_model_base {
+    llama_model_kimi_k3(const struct llama_model_params & params) : llama_model_base(params) {}
+    void load_arch_hparams(llama_model_loader & ml) override;
+    void load_arch_tensors(llama_model_loader & ml) override;
+
+    struct graph : public llm_build_delta_net_base {
+        graph(const llama_model & model, const llm_graph_params & params);
+
+        const llama_model & model;
+
+        // Cross-layer residual attention (K3's `_apply_attn_res`).
+        ggml_tensor * resi_stack = nullptr;
+
+        void          res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens);
+        ggml_tensor * res_mix(ggml_tensor * cur, ggml_tensor * score_w,
+                              int64_t n_tokens, int il);
+
+        ggml_tensor * build_kda_layer(ggml_tensor * cur, const llama_layer & layer,
+                                      llm_graph_input_rs * inp_rs,
+                                      int64_t d_conv, int64_t head_dim, int64_t n_head_kda,
+                                      int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il);
+
+        ggml_tensor * build_mla_layer(ggml_tensor * cur, const llama_layer & layer,
+                                      llm_graph_input_attn_k  * inp_attn_k,
+                                      llm_graph_input_attn_kv * inp_attn_kv,
+                                      int64_t n_embd_head_k_mla, int64_t n_embd_head_v_mla,
+                                      int64_t kv_lora_rank, int64_t n_embd_head_qk_rope,
+                                      int64_t n_embd_head_qk_nope, float kq_scale, int il);
+
+        ggml_tensor * build_latent_moe(ggml_tensor * cur, const llama_layer & layer,
+                                       int64_t n_embd_latent, int il);
+    };
+
+    std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
+};
+
 struct llama_model_kimi_linear : public llama_model_base {
     llama_model_kimi_linear(const struct llama_model_params & params) : llama_model_base(params) {}
     void load_arch_hparams(llama_model_loader & ml) override;

src/models/plamo2.cpp

@@ -382,7 +382,7 @@ ggml_tensor * llama_model_plamo2::graph::build_plamo2_mamba_layer(llm_graph_inpu
             // Custom operator to optimize the parallel associative scan
             // as described in the Annex D of the Mamba paper.
             // => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}
-            return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids);
+            return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, /*K=*/1);
         };
 
         ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows);

tests/CMakeLists.txt

@@ -217,6 +217,16 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
     set_tests_properties(test-recurrent-state-rollback PROPERTIES
         FIXTURES_REQUIRED generate-models
     )
+
+    llama_test(
+        test-recurrent-state-rollback
+        NAME test-recurrent-state-rollback-nemotron-h
+        LABEL main
+        ARGS -m "${MODEL_DIR}/nemotron_h-dense.gguf"
+    )
+    set_tests_properties(test-recurrent-state-rollback-nemotron-h PROPERTIES
+        FIXTURES_REQUIRED generate-models
+    )
 endif()
 
 llama_build_and_test(test-chat-peg-parser.cpp peg-parser/simple-tokenize.cpp)

tests/test-backend-ops.cpp

@@ -4111,9 +4111,10 @@ struct test_ssm_scan : public test_case {
     const int64_t n_seq_tokens;
     const int64_t n_seqs;
     const bool    xbc_overlap;
+    const int64_t K;
 
     std::string vars() override {
-        return VARS_TO_STR8(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs, xbc_overlap);
+        return VARS_TO_STR9(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs, xbc_overlap, K);
     }
 
     test_ssm_scan(ggml_type type = GGML_TYPE_F32,
@@ -4123,8 +4124,9 @@ struct test_ssm_scan : public test_case {
             int64_t n_group = 1,
             int64_t n_seq_tokens = 32,
             int64_t n_seqs = 32,
-            bool xbc_overlap = false)
-        : type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), xbc_overlap(xbc_overlap) {}
+            bool xbc_overlap = false,
+            int64_t K = 1)
+        : type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), xbc_overlap(xbc_overlap), K(K) {}
 
     double max_nmse_err() override {
         // SSD path (head_dim > 1) uses FP16 intermediates (M matrix, X_dt); Mamba-1 is pure FP32.
@@ -4153,7 +4155,7 @@ struct test_ssm_scan : public test_case {
             C = ggml_new_tensor_4d(ctx, type, d_state,  n_group, n_seq_tokens, n_seqs);
         }
         ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32,  n_seqs);
-        ggml_tensor * out = ggml_ssm_scan(ctx, s, x, dt, A, B, C, ids);
+        ggml_tensor * out = ggml_ssm_scan(ctx, s, x, dt, A, B, C, ids, K);
         return out;
     }
 
@@ -4185,6 +4187,114 @@ struct test_ssm_scan : public test_case {
     }
 };
 
+struct test_ssm_scan_rollback : public test_case {
+    const ggml_type type;
+
+    const int64_t d_state;
+    const int64_t head_dim;
+    const int64_t n_head;
+    const int64_t n_group;
+    const int64_t n_seq_tokens;
+    const int64_t n_seqs;
+    const int64_t K;
+
+    std::string vars() override {
+        return VARS_TO_STR8(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs, K);
+    }
+
+    std::string op_desc(ggml_tensor * t) override {
+        GGML_UNUSED(t);
+        return "SSM_SCAN_ROLLBACK";
+    }
+
+    bool run_whole_graph() override {
+        return true;
+    }
+
+    double max_err() override {
+        return 1e-6;
+    }
+
+    double err(const float * a, const float * b, size_t n) override {
+        double result = 0.0;
+        for (size_t i = 0; i < n; ++i) {
+            result = std::max(result, (double) fabsf(a[i]));
+            result = std::max(result, (double) fabsf(b[i]));
+        }
+        return result;
+    }
+
+    test_ssm_scan_rollback(ggml_type type = GGML_TYPE_F32,
+            int64_t d_state = 32,
+            int64_t head_dim = 64,
+            int64_t n_head  = 16,
+            int64_t n_group = 2,
+            int64_t n_seq_tokens = 8,
+            int64_t n_seqs = 2,
+            int64_t K = 3)
+        : type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group),
+          n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), K(K) {}
+
+    ggml_tensor * build_graph(ggml_context * ctx) override {
+        ggml_tensor * s   = ggml_new_tensor_4d(ctx, type, d_state,  head_dim,     n_head,       n_seqs);
+        ggml_tensor * x   = ggml_new_tensor_4d(ctx, type, head_dim, n_head,       n_seq_tokens, n_seqs);
+        ggml_tensor * dt  = ggml_new_tensor_3d(ctx, type, n_head,   n_seq_tokens, n_seqs);
+        ggml_tensor * A   = ggml_new_tensor_2d(ctx, type, 1,        n_head);
+        ggml_tensor * B   = ggml_new_tensor_4d(ctx, type, d_state,  n_group,      n_seq_tokens, n_seqs);
+        ggml_tensor * C   = ggml_new_tensor_4d(ctx, type, d_state,  n_group,      n_seq_tokens, n_seqs);
+        ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32,  n_seqs);
+
+        ggml_tensor * full = ggml_ssm_scan(ctx, s, x, dt, A, B, C, ids, K);
+
+        const int64_t y_elems     = head_dim * n_head * n_seq_tokens * n_seqs;
+        const int64_t state_elems = d_state  * head_dim * n_head      * n_seqs;
+
+        ggml_tensor * out = nullptr;
+        for (int64_t slot = 0; slot < K; ++slot) {
+            const int64_t prefix_tokens = n_seq_tokens - slot;
+
+            ggml_tensor * x_prefix  = ggml_cont(ctx, ggml_view_4d(ctx, x,  head_dim, n_head,  prefix_tokens, n_seqs, x->nb[1],  x->nb[2],  x->nb[3],  0));
+            ggml_tensor * dt_prefix = ggml_cont(ctx, ggml_view_3d(ctx, dt, n_head,   prefix_tokens, n_seqs, dt->nb[1], dt->nb[2], 0));
+            ggml_tensor * B_prefix  = ggml_cont(ctx, ggml_view_4d(ctx, B,  d_state,  n_group, prefix_tokens, n_seqs, B->nb[1],  B->nb[2],  B->nb[3],  0));
+            ggml_tensor * C_prefix  = ggml_cont(ctx, ggml_view_4d(ctx, C,  d_state,  n_group, prefix_tokens, n_seqs, C->nb[1],  C->nb[2],  C->nb[3],  0));
+
+            ggml_tensor * prefix = ggml_ssm_scan(ctx, s, x_prefix, dt_prefix, A, B_prefix, C_prefix, ids, /*K=*/1);
+
+            ggml_tensor * full_state   = ggml_view_1d(ctx, full,   state_elems, (y_elems + slot*state_elems)*ggml_element_size(full));
+            ggml_tensor * prefix_state = ggml_view_1d(ctx, prefix, state_elems, (head_dim*n_head*prefix_tokens*n_seqs)*ggml_element_size(prefix));
+            ggml_tensor * diff         = ggml_sum(ctx, ggml_sqr(ctx, ggml_sub(ctx, full_state, prefix_state)));
+
+            out = out == nullptr ? diff : ggml_add(ctx, out, diff);
+        }
+
+        return out;
+    }
+
+    void initialize_tensors(ggml_context * ctx) override {
+        std::random_device rd;
+        std::default_random_engine rng(rd());
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+            if (t->type == GGML_TYPE_I32) {
+                if (ggml_is_view_op(t->op)) { continue; }
+                for (int64_t r = 0; r < ggml_nrows(t); r++) {
+                    std::vector<int32_t> data(t->ne[0]);
+                    for (int i = 0; i < t->ne[0]; i++) {
+                        data[i] = i;
+                    }
+                    std::shuffle(data.begin(), data.end(), rng);
+                    ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t));
+                }
+            } else if (ggml_is_view_op(t->op)) {
+                continue;
+            } else if (t->ne[1] == n_head && t->ne[2] == 1) {
+                init_tensor_uniform(t, -1.0f, -0.5f);
+            } else {
+                init_tensor_uniform(t);
+            }
+        }
+    }
+};
+
 // GGML_OP_RWKV_WKV6
 struct test_rwkv_wkv6 : public test_case {
     const ggml_type type;
@@ -8952,6 +9062,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 256, 1)); // Nemotron-9B SSD path
     test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 512, 1)); // Nemotron-9B SSD multi-chunk (2 aligned chunks)
     test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 80, 8, 300, 2)); // Mamba-2 SSD multi-chunk (partial 2nd chunk, 2 seqs)
+    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 4, 2, false, /*K=*/4)); // Mamba-2 rollback snapshots
+    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 8, 2, false, /*K=*/3)); // Mamba-2 rollback overflow
+    test_cases.emplace_back(new test_ssm_scan_rollback(GGML_TYPE_F32, 128, 64, 16, 2, 8, 2, /*K=*/3)); // rollback snapshots match prefix states
 
     test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 1, 1));
     test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 32, 1));
@@ -9804,6 +9917,13 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
                                     use_id, 16, 8, b, with_bias, with_gate, with_lane_scale));
                                 test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, 1, 32, 256,
                                     use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, {1, 1}));
+                                if (!use_id && with_gate && !with_bias) {
+                                    // small multi-token batches (speculative decoding / MTP verify)
+                                    for (int64_t m_batch : { 2, 4, 8 }) {
+                                        test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, m_batch, 32, 256,
+                                            use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, {1, 1}));
+                                    }
+                                }
                             }
                         }
                     }

tests/test-chat.cpp

@@ -4462,6 +4462,109 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
         }
     }
 
+    // Kimi-K3 tests - custom parser
+    // Unique feature: XTML tags built from <|open|>/<|close|>/<|sep|>, and a
+    // generation prompt that leaves the think section already open.
+    {
+        auto tst = peg_tester("models/templates/Kimi-K3.jinja", detailed_debug);
+
+        // Content only. The response section is explicit even with no reasoning.
+        tst.test("<|open|>response<|sep|>Hello, world!\nWhat's up?<|close|>response<|sep|>"
+                 "<|close|>message<|sep|>")
+            .expect(message_assist)
+            .run();
+
+        // Reasoning with no opening tag - the generation prompt already opened it
+        tst.test("I'm thinking about this<|close|>think<|sep|>"
+                 "<|open|>response<|sep|>Hello, world!\nWhat's up?<|close|>response<|sep|>"
+                 "<|close|>message<|sep|>")
+            .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
+            .expect(simple_assist_msg("Hello, world!\nWhat's up?", "I'm thinking about this"))
+            .run();
+
+        // Prose that mentions the tag names must survive intact.
+        tst.test("<|open|>response<|sep|>Use the response tag, then message the handler."
+                 "<|close|>response<|sep|><|close|>message<|sep|>")
+            .expect(simple_assist_msg("Use the response tag, then message the handler."))
+            .run();
+
+        // Truncated mid-reasoning (hit the token budget): keep the reasoning.
+        tst.test("I was still thinking when the budget ran out")
+            .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
+            .expect_reasoning("I was still thinking when the budget ran out")
+            .run();
+
+        // Single tool call, one argument.
+        tst.test("<|open|>response<|sep|><|close|>response<|sep|>"
+                 "<|open|>tools<|sep|>"
+                 "<|open|>call tool=\"special_function\" index=\"1\"<|sep|>"
+                 "<|open|>argument key=\"arg1\" type=\"number\"<|sep|>1<|close|>argument<|sep|>"
+                 "<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>")
+            .tools({ special_function_tool })
+            .expect_tool_calls({
+                { "special_function", R"({"arg1":1})", "" },
+            })
+            .run();
+
+        // Tool call preceded by reasoning (no opening think tag) and content.
+        tst.test("I should call it<|close|>think<|sep|>"
+                 "<|open|>response<|sep|>On it.<|close|>response<|sep|>"
+                 "<|open|>tools<|sep|>"
+                 "<|open|>call tool=\"special_function\" index=\"1\"<|sep|>"
+                 "<|open|>argument key=\"arg1\" type=\"number\"<|sep|>1<|close|>argument<|sep|>"
+                 "<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>")
+            .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
+            .tools({ special_function_tool })
+            .expect(simple_assist_msg("On it.", "I should call it", "special_function",
+                                      R"({"arg1":1})", ""))
+            .run();
+
+        // Multiple typed arguments: values must come back as JSON numbers, not strings
+        tst.test("<|open|>response<|sep|><|close|>response<|sep|>"
+                 "<|open|>tools<|sep|>"
+                 "<|open|>call tool=\"special_function_with_opt\" index=\"1\"<|sep|>"
+                 "<|open|>argument key=\"arg1\" type=\"number\"<|sep|>1<|close|>argument<|sep|>"
+                 "<|open|>argument key=\"arg2\" type=\"number\"<|sep|>2<|close|>argument<|sep|>"
+                 "<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>")
+            .tools({ special_function_tool_with_optional_param })
+            .expect_tool_calls({
+                { "special_function_with_opt", R"({"arg1":1,"arg2":2})", "" },
+            })
+            .run();
+
+        // Parallel tool calls in one <|open|>tools<|sep|> section.
+        tst.test("<|open|>response<|sep|><|close|>response<|sep|>"
+                 "<|open|>tools<|sep|>"
+                 "<|open|>call tool=\"special_function\" index=\"1\"<|sep|>"
+                 "<|open|>argument key=\"arg1\" type=\"number\"<|sep|>1<|close|>argument<|sep|>"
+                 "<|close|>call<|sep|>"
+                 "<|open|>call tool=\"special_function_with_opt\" index=\"2\"<|sep|>"
+                 "<|open|>argument key=\"arg1\" type=\"number\"<|sep|>1<|close|>argument<|sep|>"
+                 "<|open|>argument key=\"arg2\" type=\"number\"<|sep|>2<|close|>argument<|sep|>"
+                 "<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>")
+            .parallel_tool_calls(true)
+            .tools({ special_function_tool, special_function_tool_with_optional_param })
+            .expect_tool_calls({
+                { "special_function", R"({"arg1":1})", "" },
+                { "special_function_with_opt", R"({"arg1":1,"arg2":2})", "" },
+            })
+            .run();
+
+        // String-typed argument keeps its literal text (no JSON coercion).
+        tst.test("<|open|>response<|sep|><|close|>response<|sep|>"
+                 "<|open|>tools<|sep|>"
+                 "<|open|>call tool=\"python\" index=\"1\"<|sep|>"
+                 "<|open|>argument key=\"code\" type=\"string\"<|sep|>print('hey')"
+                 "<|close|>argument<|sep|>"
+                 "<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>")
+            .tools({ python_tool })
+            .expect_tool_calls({
+                // custom delimiter: the payload itself contains )"
+                { "python", R"JSON({"code":"print('hey')"})JSON", "" },
+            })
+            .run();
+    }
+
     // Kimi-K2-Thinking tests - custom parser
     // Unique feature: tool call ID embeds function name as functions.<name>:<counter>
     {
@@ -4618,7 +4721,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
 
         // Real life test - execute_command
         tst.test("<|tool_call_begin|>functions.execute_command:0<|tool_call_argument_begin|>{\"command\": \"ls -lah\""
-            ", \"cwd\": \"/home/jarvis/development/exllamav3\", \"timeout\": 10}")
+            ", \"cwd\": \"/home/user/development/exllamav3\", \"timeout\": 10}")
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
             .parallel_tool_calls(true)
             .tools({
@@ -4648,7 +4751,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             expect_tool_calls({
                 {
                     "execute_command",
-                    R"({"command": "ls -lah", "cwd": "/home/jarvis/development/exllamav3", "timeout": 10})",
+                    R"({"command": "ls -lah", "cwd": "/home/user/development/exllamav3", "timeout": 10})",
                     "functions.execute_command:0"
                 }
             })
@@ -6955,6 +7058,24 @@ static void test_reasoning_budget_message_per_request() {
     }
 }
 
+static void test_reasoning_effort_caps() {
+    LOG_DBG("%s\n", __func__);
+
+    auto assert_supports_effort = [](const std::string & path, bool expected) {
+        auto tmpls = read_templates(path);
+        assert_equals(expected, common_chat_templates_get_caps(tmpls.get()).at("supports_reasoning_effort"));
+    };
+
+    assert_supports_effort("models/templates/deepseek-ai-DeepSeek-V4.jinja", true);
+    assert_supports_effort("models/templates/muse-glimmer.jinja", true);
+    assert_supports_effort("models/templates/tencent-Hy3.jinja", true);
+    assert_supports_effort("models/templates/openai-gpt-oss-120b.jinja", true);
+    assert_supports_effort("models/templates/upstage-Solar-Open-100B.jinja", true);
+    assert_supports_effort("models/templates/Cohere2MoE.jinja", true);
+    assert_supports_effort("models/templates/meta-llama-Llama-3.1-8B-Instruct.jinja", false);
+    assert_supports_effort("models/templates/Qwen-Qwen3-0.6B.jinja", false);
+}
+
 static void test_msg_diffs_compute() {
     LOG_DBG("%s\n", __func__);
     {
@@ -7114,6 +7235,7 @@ int main(int argc, char ** argv) {
         test_deepseek_v4_thinking_retention();
         test_deepseek_v4_tool_result_ordering();
         test_template_generation_prompt();
+        test_reasoning_effort_caps();
         test_reasoning_budget_tokens_per_request();
         test_reasoning_budget_message_per_request();
         test_template_output_peg_parsers(detailed_debug);

tests/test-jinja.cpp

@@ -33,6 +33,7 @@ static void test_array_methods(testing & t);
 static void test_object_methods(testing & t);
 static void test_hasher(testing & t);
 static void test_stats(testing & t);
+static void test_string_parts(testing & t);
 static void test_fuzzing(testing & t);
 
 static bool g_python_mode = false;
@@ -72,6 +73,7 @@ int main(int argc, char *argv[]) {
     if (!g_python_mode) {
         t.test("hasher", test_hasher);
         t.test("stats", test_stats);
+        t.test("string parts", test_string_parts);
         t.test("fuzzing", test_fuzzing);
     }
 
@@ -2057,6 +2059,36 @@ static void test_stats(testing & t) {
     });
 }
 
+static void test_string_parts(testing & t) {
+    static auto render = [](const std::string & tmpl, const json & vars) -> jinja::string {
+        jinja::lexer lexer;
+        auto lexer_res = lexer.tokenize(tmpl);
+
+        jinja::program ast = jinja::parse_from_tokens(lexer_res);
+
+        jinja::context ctx(tmpl);
+        jinja::global_from_json(ctx, vars, true);
+
+        jinja::runtime runtime(ctx);
+        return runtime.gather_string_parts(runtime.execute(ast))->as_string();
+    };
+
+    t.test("merge joins only the neighbours with the same type", [](testing & t) {
+        // "AB" comes from the input and merges, "-" comes from the template and must not
+        jinja::string res = render("{{ val.a }}{{ val.b }}-{{ val.c }}",
+                                   json{{"val", json{{"a", "A"}, {"b", "B"}, {"c", "C"}}}});
+
+        if (t.assert_true("3 parts after the merge", res.parts.size() == 3)) {
+            t.assert_true("part 0 is the merged input", res.parts[0].val == "AB" && res.parts[0].is_input);
+            t.assert_true("part 1 is from the template", res.parts[1].val == "-" && !res.parts[1].is_input);
+            t.assert_true("part 2 is input",             res.parts[2].val == "C" && res.parts[2].is_input);
+        } else {
+            t.log("parts: " + std::to_string(res.parts.size()) + ", rendered: " + json(res.str()).dump());
+        }
+    });
+
+}
+
 static void test_template_cpp(testing & t, const std::string & name, const std::string & tmpl, const json & vars, const std::string & expect) {
     t.test(name, [&tmpl, &vars, &expect](testing & t) {
         jinja::lexer lexer;

tests/test-llama-archs.cpp

@@ -105,6 +105,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
             || arch == LLM_ARCH_DEEPSEEK32
             || arch == LLM_ARCH_GLM_DSA
             || arch == LLM_ARCH_KIMI_LINEAR
+            || arch == LLM_ARCH_KIMI_K3
             || arch == LLM_ARCH_MISTRAL4) {
         n_embd = 128;
         n_head = 1;
@@ -145,7 +146,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
     ms.add_kv(LLM_KV_FULL_ATTENTION_INTERVAL, uint32_t(2));
 
     if (arch == LLM_ARCH_PLAMO2 || arch == LLM_ARCH_JAMBA || arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE ||
-            arch == LLM_ARCH_GRANITE_HYBRID || arch == LLM_ARCH_LFM2 || arch == LLM_ARCH_LFM2MOE || arch == LLM_ARCH_KIMI_LINEAR) {
+            arch == LLM_ARCH_GRANITE_HYBRID || arch == LLM_ARCH_LFM2 || arch == LLM_ARCH_LFM2MOE || arch == LLM_ARCH_KIMI_LINEAR || arch == LLM_ARCH_KIMI_K3) {
         GGML_ASSERT(n_layer >= 2);
         std::vector<uint32_t> n_head_per_layer;
         n_head_per_layer.reserve(n_layer);
@@ -164,6 +165,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
             || arch == LLM_ARCH_DEEPSEEK32
             || arch == LLM_ARCH_GLM_DSA
             || arch == LLM_ARCH_KIMI_LINEAR
+            || arch == LLM_ARCH_KIMI_K3
             || arch == LLM_ARCH_MISTRAL4) {
         ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH,       uint32_t(576));
         ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH,     uint32_t(512));
@@ -218,6 +220,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
     if (moe) {
         ms.add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, n_ff);
         ms.add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, n_ff / 2);  // distinct from n_ff so a saver key-clobber surfaces on reload
+        ms.add_kv(LLM_KV_EXPERT_LATENT_LENGTH,       n_ff);
         ms.add_kv(LLM_KV_INTERLEAVE_MOE_LAYER_STEP,  uint32_t(2));
         ms.add_kv(LLM_KV_EXPERT_COUNT,               uint32_t(2));
         ms.add_kv(LLM_KV_EXPERT_USED_COUNT,          uint32_t(1));
@@ -243,6 +246,11 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
     ms.add_kv(LLM_KV_KDA_HEAD_DIM,              uint32_t(128));
     ms.add_kv(LLM_KV_WKV_HEAD_SIZE,             n_embd/n_head);
     ms.add_kv(LLM_KV_SHORTCONV_L_CACHE,         uint32_t(3));
+    ms.add_kv(LLM_KV_RESIDUAL_SCALE,            3.5565588200778455f);
+    ms.add_kv(LLM_KV_ATTN_RES_BLOCK_SIZE,       uint32_t(12));
+    ms.add_kv(LLM_KV_ACTIVATION_SITU_BETA,      4.0f);
+    ms.add_kv(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, 25.0f);
+    ms.add_kv(LLM_KV_KDA_GATE_LOWER_BOUND,      -5.0f);
 
     for (uint32_t il = 0; il < n_layer; il++) {
         ggml_tensor t;
@@ -364,12 +372,14 @@ static bool moe_mandatory(const llm_arch arch) {
         case LLM_ARCH_SMALLTHINKER:
         case LLM_ARCH_LLADA_MOE:
         case LLM_ARCH_GROVEMOE:
+        case LLM_ARCH_MINIMAX_01:
         case LLM_ARCH_MINIMAX_M2:
         case LLM_ARCH_MINIMAX_M3:
         case LLM_ARCH_RND1:
         case LLM_ARCH_PADDLEOCR:
         case LLM_ARCH_MIMO2:
         case LLM_ARCH_KIMI_LINEAR:
+        case LLM_ARCH_KIMI_K3:
         case LLM_ARCH_STEP35:
         case LLM_ARCH_MISTRAL4:
         case LLM_ARCH_MELLUM:
@@ -436,7 +446,7 @@ static bool arch_supported(const llm_arch arch) {
 
     // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI.
 #ifdef GGML_USE_WEBGPU
-    if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA) {
+    if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_MINIMAX_01) {
         return false;
     }
 #endif // GGML_USE_WEBGPU

tools/cli/README.md

@@ -170,6 +170,7 @@
 | `--jinja, --no-jinja` | whether to use jinja template engine for chat (default: enabled)<br/>(env: LLAMA_ARG_JINJA) |
 | `--reasoning-format FORMAT` | controls whether thought tags are allowed and/or extracted from the response, and in which format they're returned; one of:<br/>- none: leaves thoughts unparsed in `message.content`<br/>- deepseek: puts thoughts in `message.reasoning_content`<br/>- deepseek-legacy: keeps `<think>` tags in `message.content` while also populating `message.reasoning_content`<br/>(default: auto)<br/>(env: LLAMA_ARG_THINK) |
 | `-rea, --reasoning [on\|off\|auto]` | Use reasoning/thinking in the chat ('on', 'off', or 'auto', default: 'auto' (detect from template))<br/>(env: LLAMA_ARG_REASONING) |
+| `--reasoning-effort LEVEL` | reasoning effort level given to the chat template: 'default' to keep the template default,<br/>or a level such as 'minimal', 'low', 'medium', 'high', 'xhigh' or 'max' (default: default)<br/>(env: LLAMA_ARG_REASONING_EFFORT) |
 | `--reasoning-budget N` | token budget for thinking: -1 for unrestricted, 0 for immediate end, N>0 for token budget (default: -1)<br/>(env: LLAMA_ARG_THINK_BUDGET) |
 | `--reasoning-budget-message MESSAGE` | message injected before the end-of-thinking tag when reasoning budget is exhausted (default: none)<br/>(env: LLAMA_ARG_THINK_BUDGET_MESSAGE) |
 | `--reasoning-preserve, --no-reasoning-preserve` | preserve reasoning trace in the full history, not just the last assistant message (default: template default)<br/>compatible with certain templates having 'supports_preserve_reasoning' capability<br/>example: https://docs.z.ai/guides/capabilities/thinking-mode#preserved-thinking<br/>(env: LLAMA_ARG_REASONING_PRESERVE) |