model: add Kimi-K3 text model (#26185)

* model: add Kimi-K3 text model

Hybrid KDA (linear) + MLA (full) attention as in Kimi-Linear-48B, plus five
things that architecture 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)

K3's text_config reports KimiLinearForCausalLM - the older 48B architecture -
so get_model_architecture routes on the top-level name instead.

The KDA decay gate has two forms, selected by linear_attn_config's
gate_lower_bound. It is not a clamp: when set it swaps the activation entirely
(fla/ops/kda/gate.py), from -exp(A_log)*softplus(x) to
lower_bound*sigmoid(exp(A_log)*x). K3 sets it to -5.0; kimi-linear leaves it
unset, so that path is unchanged.

Cross-layer residuals reuse ggml_dsv4_hc_pre for the weighted sum. That op is
CPU + CUDA only, so Metal/Vulkan will fall back per-node until those kernels
exist.

The routed experts ship as compressed-tensors "mxfp4-pack-quantized". That is
bit-compatible with ggml's MXFP4 - same E2M1 code assignment, same E8M0 scale
byte, only the nibble positions within a block differ - so they are repacked
rather than dequantized, losslessly and without a ~5.5 TB bf16 round-trip.
The repack is built lazily because gguf_writer holds every added tensor until
the final write. DeepSeek-V4 was already doing the identical bit-shuffling, so
it now shares the helper.

Verified against Moonshot's own code path (transformers + fla's Triton KDA
kernels) on a tiny model exercising every K3-specific feature. Final-position
logits vs the fp32 reference: 6.7e-05 rel / corr 1.00000000 for both the
chunked and the recurrent delta-net path. MXFP4 blocks dequantize to the source
weights with 0.0e+00 error.

Assisted-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* model: fix ty errors in the Kimi-K3 converter

- `_res_parts` buffers (kind, tensor) pairs, not bare tensors
- `get_tensors` must return an Iterator, matching ModelBase
- LazyBase's `func` takes one argument, so pass the expert loaders through
  `args` instead of the closure
- borrowing KimiLinearModel.set_vocab from an unrelated TextModel is
  deliberate and safe, but not expressible in the signature

No behaviour change: the MXFP4 repack still dequantizes to the source weights
with 0.0e+00 error and end-to-end logits are unchanged (8.386e-03 rel,
corr 0.99996630).

Assisted-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* Update conversion/kimi_k3.py

Co-authored-by: Boris Dvorkin  <b_dvorkin@niuitmo.ru>

* Increase LLAMA_MAX_EXPERTS from 512 to 1024

* tests : support for Kimi K3 in archs test

* chat : add Kimi K3 chat format (reasoning, content, typed tool calls)

K3's assistant output is an XTML-ish tagged format built by the template's
open_tag/close_tag macros. Two properties break generic parsing:

1. The generation prompt ends with open_tag('think'), so the completion
   starts inside the think section with no opening marker in the output
   (thinking_forced_open).
2. Only <|open|>/<|close|>/<|sep|>/<|end_of_msg|> are special tokens; tag
   names ("think", "response", "message") are ordinary text tokens.

Adds common_chat_params_init_kimi_k3 (PEG_NATIVE) with detection on the
marker trio, reasoning extraction, response unwrapping, and tool-call
parsing of the tools/call/argument tag structure with argument types
taken from the tool schema. Includes the K3 chat template fixture and 9
test-chat cases derived from real generations of the full 2.8T model.

Verified end-to-end against Kimi-K3-Q2_K (GrEarl/Kimi-K3-GGUF) on 8x B200:
content, reasoning_content, streaming deltas, and tool_calls all correct;
finish_reason stop/tool_calls as appropriate.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* chat : add message_delimiters for Kimi K3

Per-role message-start markers for token-level span splitting. User and
assistant messages carry only the role attribute, so their full opener
(through <|sep|>) is used; system and tool messages continue with more
attributes (type=/tool=/index=), so those delimiters stop after the
role's closing quote. Verified against the K3 tiktoken vocabulary that
the closing quote is always a standalone token across all attribute
variants, so the token-level prefix match stays exact.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* fix: apply nits from @ngxson and text fixes from @danielhanchen

* tests : added missing hyperparameters and tensors for Kimi K3 in test-llama-archs

* chore : move overly verbose header file comments to Kimi K3 source file

* tests : re-enabled KIMI_K3 in test-llama-archs for WebGPU backend

* model-saver : emit kda_gate_lower_bound for Kimi K3

Quick fix. The Kimi K3 loader reads kda_gate_lower_bound and gates a graph branch on it (it scales the KDA gate when the bound is above -INFINITY), but the model
saver never wrote the key, so a save->load roundtrip silently dropped it back to the -INFINITY default and changed the model's output. The real K3 config sets gate_lower_bound = -5.0.

I propose to emit it from the saver, and set it to -5.0 in the test-llama-archs K3 case so the roundtrip check exercises it (the roundtrip fails without the saver line).

* Refactor conditional for model architecture check

* tests : re-enabled (again) KIMI_K3 and MINIMAX_M3 in test-llama-archs for WebGPU backend

* fix code comments

* add template on conversion

* move repack_mxfp4_blocks to model base

* nits

* add_value_length

* optimize res_stack construction

* nits

---------

Co-authored-by: Boris Dvorkin <b_dvorkin@niuitmo.ru>
Co-authored-by: Stanisław Szymczyk <sszymczy@gmail.com>
Co-authored-by: Deepankar Singh <singh.deepankar39@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Caleb DeLeeuw <caleb.deleeuw@gmail.com>
Co-authored-by: Xuan Son Nguyen <son@huggingface.co>

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/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/ggml-sycl.cpp

@@ -4561,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,
@@ -5591,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

@@ -43,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) |

tools/completion/README.md

@@ -251,6 +251,7 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1
 | `--jinja, --no-jinja` | whether to use jinja template engine for chat (default: disabled)<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) |
@@ -523,13 +524,15 @@ These options help improve the performance and memory usage of the LLaMA models.
 -   `-t N, --threads N`: Set the number of threads to use during generation. For optimal performance, it is recommended to set this value to the number of physical CPU cores your system has (as opposed to the logical number of cores). Using the correct number of threads can greatly improve performance.
 -   `-tb N, --threads-batch N`: Set the number of threads to use during batch and prompt processing. In some systems, it is beneficial to use a higher number of threads during batch processing than during generation. If not specified, the number of threads used for batch processing will be the same as the number of threads used for generation.
 
-### Mlock
+### Model Loading Mode
 
--   `--mlock`: Lock the model in memory, preventing it from being swapped out when memory-mapped. This can improve performance but trades away some of the advantages of memory-mapping by requiring more RAM to run and potentially slowing down load times as the model loads into RAM.
-
-### No Memory Mapping
-
--   `--no-mmap`: Do not memory-map the model. By default, models are mapped into memory, which allows the system to load only the necessary parts of the model as needed. However, if the model is larger than your total amount of RAM or if your system is low on available memory, using mmap might increase the risk of pageouts, negatively impacting performance. Disabling mmap results in slower load times but may reduce pageouts if you're not using `--mlock`. Note that if the model is larger than the total amount of RAM, turning off mmap would prevent the model from loading at all.
+-   `-lm MODE, --load-mode MODE`: Specify the model loading mode (default: `auto`).
+    -   `auto`: Memory-map the model, unless the device does not support it.
+    -   `none`: No special loading mode. Disabling mmap results in slower load times but may reduce pageouts if you're not using `mlock`. Note that if the model is larger than the total amount of RAM, turning off mmap would prevent the model from loading at all.
+    -   `mmap`: Memory-map the model.
+    -   `mlock`: Lock the model in memory, preventing it from being swapped out when memory-mapped. This can improve performance but trades away some of the advantages of memory-mapping by requiring more RAM to run and potentially slowing down load times as the model loads into RAM.
+    -   `mmap+mlock`: Memory-map the model and lock it in memory.
+    -   `dio`: Use DirectIO if available.
 
 ### NUMA support
 

tools/llama-bench/README.md

@@ -67,8 +67,8 @@ test parameters:
   -nkvo, --no-kv-offload <0|1>              (default: 0)
   -fa, --flash-attn <on|off|auto>           (default: auto)
   -dev, --device <dev0/dev1/...>            (default: auto)
-  -mmp, --mmap <0|1>                        (default: 1)
-  -dio, --direct-io <0|1>                   (default: 0)
+  -mmp, --mmap <0|1>                        (DEPRECATED IN FAVOUR OF --load-mode)
+  -dio, --direct-io <0|1>                   (DEPRECATED IN FAVOUR OF --load-mode)
   -embd, --embeddings <0|1>                 (default: 0)
   -ts, --tensor-split <ts0/ts1/..>          (default: 0)
   -ot --override-tensor <tensor name pattern>=<buffer type>;...

tools/mtmd/clip-impl.h

@@ -6,6 +6,7 @@
 
 #include <array>
 #include <climits>
+#include <cmath>
 #include <cstdarg>
 #include <cinttypes>
 #include <string>
@@ -603,7 +604,7 @@ struct clip_image_u8 {
             // return a dummy value, so that legacy code can still process image without errors
             return { 0, 0, 0 };
         }
-        int idx = (y * nx + x) * 3;
+        size_t idx = ((size_t) y * (size_t) nx + (size_t) x) * 3;
         return { buf[idx], buf[idx + 1], buf[idx + 2] };
     }
 
@@ -611,8 +612,8 @@ struct clip_image_u8 {
         if (is_placeholder()) {
             return; // no-op
         }
-        int idx = (y * nx + x) * 3;
-        buf[idx] = rgb[0];
+        size_t idx = ((size_t) y * (size_t) nx + (size_t) x) * 3;
+        buf[idx]     = rgb[0];
         buf[idx + 1] = rgb[1];
         buf[idx + 2] = rgb[2];
     }
@@ -642,9 +643,25 @@ struct mtmd_serialization; // forward declaration
 struct clip_image_f32 {
     // marks the global view in e.g., DeepSeek-OCR Models
     bool add_viewsep = false;
-    // whether a learned newline (or EOI) token should be appended after the image (eg Granite4 Vision)
+    // appends a learned newline (or EOI) token after the image
+    // no model uses it now (Granite4 Vision moved to anyres), kept for future models
     bool add_newline = false;
 
+    // llava-next "anyres" tiling, used by Granite4 Vision
+    // the whole grid is encoded and assembled in a single graph
+    // NOTE: excluded from serialized: a deserialized image is always a placeholder, which is never encoded
+    struct anyres_info {
+        int grid_x = 0; // tiles per row, 0 means the image is not tiled
+        int grid_y = 0; // tiles per column
+        int orig_nx = 0; // size of the source image, used to drop the padding tokens
+        int orig_ny = 0;
+
+        bool is_tiled() const {
+            return grid_x > 0 && grid_y > 0;
+        }
+    };
+    anyres_info anyres;
+
     clip_image_size get_size() const {
         return { nx_, ny_ };
     }
@@ -726,6 +743,25 @@ struct clip_image_f32 {
     }
 };
 
+// token area kept after removing the padding added by the anyres resize
+// ref: https://github.com/huggingface/transformers/blob/v5.0.0/src/transformers/models/llava_next/modeling_llava_next.py#L109
+static inline void clip_anyres_unpad(int cur_w, int cur_h, int orig_w, int orig_h,
+                                     int & off_x, int & off_y, int & out_w, int & out_h) {
+    off_x = 0;
+    off_y = 0;
+    out_w = cur_w;
+    out_h = cur_h;
+    if ((float) orig_w / orig_h > (float) cur_w / cur_h) {
+        const int new_h = (int) std::floor((double) orig_h * cur_w / orig_w + 1e-7);
+        off_y = (cur_h - new_h) / 2;
+        out_h = cur_h - 2 * off_y;
+    } else {
+        const int new_w = (int) std::floor((double) orig_w * cur_h / orig_h + 1e-7);
+        off_x = (cur_w - new_w) / 2;
+        out_w = cur_w - 2 * off_x;
+    }
+}
+
 //
 // logging
 //

tools/mtmd/clip.cpp

@@ -1595,6 +1595,9 @@ struct clip_model_loader {
                         hparams.image_resize_algo = RESIZE_ALGO_BICUBIC_PILLOW;
                         hparams.image_resize_pad  = PAD_NONE;
                         get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false);
+                        // n_merge is used as a divisor in clip_image_batch_encode
+                        // (gh / n_merge); reject 0 to avoid int div-by-zero (DoS).
+                        GGML_ASSERT(hparams.n_merge > 0);
                         hparams.rope_theta = 10000.0f; // vision_config.rope_theta
                         // MiniMax-M3: max_pixels 451584 (=672^2) -> 576 merged tokens (image_seq_length)
                         hparams.set_limit_image_tokens(8, 576);
@@ -1823,7 +1826,9 @@ struct clip_model_loader {
                         // unlimited-ocr shares the v1 projector but tiles up to 32
                         get_u32(KEY_PREPROC_MIN_TILES, hparams.preproc_min_tiles, false);
                         get_u32(KEY_PREPROC_MAX_TILES, hparams.preproc_max_tiles, false);
-                        GGML_ASSERT(hparams.preproc_min_tiles <= hparams.preproc_max_tiles);
+                        GGML_ASSERT(hparams.preproc_min_tiles >= 0
+                                    && hparams.preproc_min_tiles <= hparams.preproc_max_tiles
+                                    && hparams.preproc_max_tiles <= 256);
                      } break;
                 case PROJECTOR_TYPE_HUNYUANVL:
                     {
@@ -1888,6 +1893,9 @@ struct clip_model_loader {
                         hparams.audio_window_len       = 400;
                         hparams.audio_hop_len          = 160;
                         get_u32(KEY_A_CHUNK_SIZE,           hparams.audio_chunk_size);
+                        // context_size is squared for the attn_dists/mask buffers; cap to prevent int32 overflow
+                        // (legitimate values are small, e.g. 12-200; 8192^2 = 67M still fits int32)
+                        GGML_ASSERT(hparams.audio_chunk_size > 0 && hparams.audio_chunk_size <= 8192);
                         get_u32(KEY_A_CONV_KERNEL_SIZE,     hparams.audio_conv_kernel_size);
                         get_u32(KEY_A_MAX_POS_EMB,          hparams.audio_max_pos_emb);
                         get_u32(KEY_A_PROJ_WINDOW_SIZE,     hparams.audio_proj_window_size);
@@ -1927,8 +1935,9 @@ struct clip_model_loader {
                     // note: some models having hparams.image_size == 0, which means the image size is dynamic
                     throw std::runtime_error(string_format("%s: image_size (%d) cannot be negative\n", __func__, hparams.image_size));
                 }
-                if (hparams.image_size > 65536) {
-                    throw std::runtime_error(string_format("%s: image_size (%d) is too large (max 65536)\n", __func__, hparams.image_size));
+                if (hparams.image_size > 8192) {
+                    // cap prevents int32 overflow in n_patches = (image_size/patch_size)^2
+                    throw std::runtime_error(string_format("%s: image_size (%d) is too large (max 8192)\n", __func__, hparams.image_size));
                 }
                 if (hparams.patch_size <= 0 || hparams.patch_size >= 65536) {
                     throw std::runtime_error(string_format("%s: patch_size (%d) must be positive and less than 65536\n", __func__, hparams.patch_size));
@@ -1939,9 +1948,12 @@ struct clip_model_loader {
                 if (hparams.image_max_pixels < hparams.image_min_pixels) {
                     throw std::runtime_error(string_format("%s: image_max_pixels (%d) is less than image_min_pixels (%d)\n", __func__, hparams.image_max_pixels, hparams.image_min_pixels));
                 }
-                if (hparams.n_merge < 0 || hparams.n_merge >= 65536) {
+                if (hparams.n_merge <= 0 || hparams.n_merge >= 65536) {
                     throw std::runtime_error(string_format("%s: n_merge (%d) must be greater than 0 and less than 65536\n", __func__, hparams.n_merge));
                 }
+                if (hparams.attn_window_size > 4096) {
+                    throw std::runtime_error(string_format("%s: attn_window_size (%d) is too large (max 4096)\n", __func__, hparams.attn_window_size));
+                }
             }
 
             LOG_INF("%s: projector:          %s\n", __func__, proj_type.c_str());
@@ -3734,6 +3746,9 @@ struct clip_model_loader {
             }
             return;
         }
+        if (gguf_get_kv_type(ctx_gguf.get(), i) != GGUF_TYPE_ARRAY) {
+            throw std::runtime_error(string_format("%s: key '%s' is not an array\n", __func__, key.c_str()));
+        }
         const auto type = gguf_get_arr_type(ctx_gguf.get(), i);
         if (type != GGUF_TYPE_FLOAT32) {
             throw std::runtime_error(string_format("%s: array '%s' has type %d, expected %d (GGUF_TYPE_FLOAT32)\n", __func__, key.c_str(), type, GGUF_TYPE_FLOAT32));
@@ -3768,6 +3783,9 @@ struct clip_model_loader {
             }
             return;
         }
+        if (gguf_get_kv_type(ctx_gguf.get(), i) != GGUF_TYPE_ARRAY) {
+            throw std::runtime_error(string_format("%s: key '%s' is not an array\n", __func__, key.c_str()));
+        }
         const auto type = gguf_get_arr_type(ctx_gguf.get(), i);
         if (type != GGUF_TYPE_INT32) {
             throw std::runtime_error(string_format("%s: array '%s' has type %d, expected %d (GGUF_TYPE_INT32)\n", __func__, key.c_str(), type, GGUF_TYPE_INT32));
@@ -4217,18 +4235,20 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
         case PROJECTOR_TYPE_GRANITE4_VISION:
             {
                 // Per-tile output token count: each projector block outputs
-                // query_side^2 tokens per window × n^2 windows.
-                // For 384×384 input: n = 24/8 = 3, query_side = 4 → 144.
+                // query_side^2 tokens per window x n^2 windows.
+                // For 384x384 input: n = 24/8 = 3, query_side = 4 -> 144.
                 const int window_side = ctx->model.hparams.downsample_window_side;
                 const int query_side  = ctx->model.hparams.downsample_query_side;
                 const int side        = img->nx() / params.patch_size;
                 const int n           = side / window_side;
-                n_patches             = (query_side * n) * (query_side * n);
-                if (img->add_newline) {
-                    // For single-tile case: append 1 newline row.
-                    // For multi-tile rowwise: handled by caller, but here we
-                    // report the per-tile count including one trailing newline.
-                    n_patches += 1;
+                const int out_side    = query_side * n;
+                n_patches             = out_side * out_side;
+                if (img->anyres.is_tiled()) {
+                    // overview tile, then the unpadded tile grid with one newline per row
+                    int off_x, off_y, w, h;
+                    clip_anyres_unpad(img->anyres.grid_x * out_side, img->anyres.grid_y * out_side,
+                                      img->anyres.orig_nx, img->anyres.orig_ny, off_x, off_y, w, h);
+                    n_patches += h * (w + 1);
                 }
             } break;
         default:
@@ -5408,13 +5428,13 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
                 const int context_size = ctx->model.hparams.audio_chunk_size;
                 const int max_pos_emb  = ctx->model.hparams.audio_max_pos_emb;
 
-                std::vector<int32_t> dists(context_size * context_size);
+                std::vector<int32_t> dists((size_t) context_size * (size_t) context_size);
                 for (int i = 0; i < context_size; i++) {
                     for (int j = 0; j < context_size; j++) {
                         int d = i - j;
                         if (d < -context_size) d = -context_size;
                         if (d >  context_size) d =  context_size;
-                        dists[i * context_size + j] = d + max_pos_emb;
+                        dists[(size_t) i * (size_t) context_size + (size_t) j] = d + max_pos_emb;
                     }
                 }
                 set_input_i32("attn_dists", dists);
@@ -5423,13 +5443,13 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
                 const int remainder  = n_frames % context_size;
                 if (remainder > 0) {
                     const int num_blocks = (n_frames + context_size - 1) / context_size;
-                    std::vector<float> mask(context_size * context_size * num_blocks, 0.0f);
+                    std::vector<float> mask((size_t) context_size * (size_t) context_size * (size_t) num_blocks, 0.0f);
                     const float neg_inf = -INFINITY;
-                    const int last_block_offset = (num_blocks - 1) * context_size * context_size;
+                    const size_t last_block_offset = (size_t) (num_blocks - 1) * (size_t) context_size * (size_t) context_size;
                     for (int q = 0; q < context_size; q++) {
                         for (int k = 0; k < context_size; k++) {
                             if (q >= remainder || k >= remainder) {
-                                mask[last_block_offset + q * context_size + k] = neg_inf;
+                                mask[last_block_offset + (size_t) q * (size_t) context_size + (size_t) k] = neg_inf;
                             }
                         }
                     }
@@ -5493,10 +5513,18 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
                     return idx;
                 };
 
+                // the same permutation is applied to every tile of the stacked image
                 auto upload = [&](const std::string & name, const std::vector<int32_t> & idx) {
                     ggml_tensor * t = ggml_graph_get_tensor(gf, name.c_str());
                     GGML_ASSERT(t);
-                    ggml_backend_tensor_set(t, idx.data(), 0, idx.size() * sizeof(int32_t));
+                    GGML_ASSERT(ggml_nelements(t) % (int64_t) idx.size() == 0);
+                    const int n_rep = ggml_nelements(t) / idx.size();
+                    std::vector<int32_t> buf;
+                    buf.reserve(idx.size() * n_rep);
+                    for (int i = 0; i < n_rep; ++i) {
+                        buf.insert(buf.end(), idx.begin(), idx.end());
+                    }
+                    ggml_backend_tensor_set(t, buf.data(), 0, ggml_nbytes(t));
                 };
 
                 // Stage 1b only uses block 0's permutations; future stages