llama : respect no_alloc in all memory caches (#31087)

* llama : respect no_alloc in memory caches

Assisted-by: dsh:llama.cpp/ggml-org/DeepSeek-V4-Flash-Vision-Exp-GGUF (Mac Studio, M2 Ultra, 192GB)

* llama : TODO: move no_alloc out of llama_hparams

Assisted-by: dsh:llama.cpp/ggml-org/DeepSeek-V4-Flash-Vision-Exp-GGUF (Mac Studio, M2 Ultra, 192GB)

* llama : TODO: move no_alloc out of llama_hparams

Assisted-by: dsh:llama.cpp/ggml-org/DeepSeek-V4-Flash-Vision-Exp-GGUF (Mac Studio, M2 Ultra, 192GB)

.github/workflows/build-cmake-pkg.yml

@@ -9,7 +9,7 @@ permissions:
 
 jobs:
   linux:
-    runs-on: [self-hosted, Linux, CPU]
+    runs-on: [self-hosted, Linux, fast]
     steps:
       - uses: actions/checkout@v6
         with:

.github/workflows/build-sanitize.yml

@@ -49,7 +49,7 @@ jobs:
         - sanitizer: THREAD
           machine: ubuntu-24.04
         - sanitizer: UNDEFINED
-          machine: [self-hosted, X64, Linux]
+          machine: [self-hosted, X64, Linux, CPU]
 
     runs-on: ${{ matrix.machine }}
 

.github/workflows/ci-self-hosted-cpu.yml

@@ -46,7 +46,7 @@ env:
 
 jobs:
   cpu-x64-high-perf:
-    runs-on: [self-hosted, Linux, X64]
+    runs-on: [self-hosted, Linux, X64, CPU]
 
     steps:
       - name: Clone
@@ -56,7 +56,7 @@ jobs:
       - name: Test
         id: ggml-ci
         run: |
-          LLAMA_ARG_THREADS=$(nproc) GG_BUILD_HIGH_PERF=1 GG_BUILD_EXTRA_TESTS_0=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
+          LLAMA_ARG_THREADS=$(nproc) GG_BUILD_HIGH_PERF=1 GG_BUILD_EXTRA_TESTS_0=1 bash ./ci/run.sh ~/results/llama.cpp/$RUNNER_NAME ~/mnt/llama.cpp
 
   cpu-arm64-high-perf-graviton4:
     runs-on: ah-ubuntu_24_04-c8g_8x

.github/workflows/ci-self-hosted-vulkan.yml

@@ -79,7 +79,7 @@ jobs:
         id: ggml-ci
         run: |
           vulkaninfo --summary
-          GG_BUILD_VULKAN=1 GGML_VK_DISABLE_COOPMAT2=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
+          GG_BUILD_VULKAN=1 GGML_VK_DISABLE_COOPMAT2=1 bash ./ci/run.sh ~/results/llama.cpp/$RUNNER_NAME ~/mnt/llama.cpp
 
       # - name: ccache-buckets-save
       #   if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
@@ -124,7 +124,7 @@ jobs:
         id: ggml-ci
         run: |
           vulkaninfo --summary
-          GG_BUILD_VULKAN=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
+          GG_BUILD_VULKAN=1 bash ./ci/run.sh ~/results/llama.cpp/$RUNNER_NAME ~/mnt/llama.cpp
 
       # - name: ccache-buckets-save
       #   if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}

common/chat.cpp

@@ -428,6 +428,41 @@ std::vector<common_chat_msg> common_chat_msgs_parse_oaicompat(const json & messa
             if (message.contains("tool_call_id")) {
                 msg.tool_call_id = message.at("tool_call_id");
             }
+            {
+                json extra = json::object();
+                for (const auto & [k, v] : message.items()) {
+                    if (k != "role" && k != "content" && k != "tool_calls" && k != "reasoning_content" &&
+                        k != "name" && k != "tool_call_id") {
+                        extra[k] = v;
+                    }
+                }
+                if (has_tool_calls) {
+                    // per tool call: keys besides type/function/id, and function keys besides name/arguments
+                    json tc_extra = json::array();
+                    bool any      = false;
+                    for (const auto & tool_call : message.at("tool_calls")) {
+                        json e = json::object();
+                        for (const auto & [k, v] : tool_call.items()) {
+                            if (k != "type" && k != "function" && k != "id") {
+                                e[k] = v;
+                            }
+                        }
+                        for (const auto & [k, v] : tool_call.at("function").items()) {
+                            if (k != "name" && k != "arguments") {
+                                e["function"][k] = v;
+                            }
+                        }
+                        any = any || !e.empty();
+                        tc_extra.push_back(e);
+                    }
+                    if (any) {
+                        extra["__tool_call_extras"] = tc_extra;
+                    }
+                }
+                if (!extra.empty()) {
+                    msg.extra_json = extra.dump();
+                }
+            }
 
             msgs.push_back(msg);
         }
@@ -516,6 +551,50 @@ static json render_message_to_json(const std::vector<common_chat_msg> & msgs, co
     return messages_inp_normalizer(c).normalize(messages);
 }
 
+// DeepSeek V4.1 reads non-OpenAI message keys and joins content blocks with "\n\n" around images,
+// so keep media markers as separate text parts and restore the extra keys
+static json render_message_to_json_deepseek_v4_1(const std::vector<common_chat_msg> & msgs, const jinja::caps & c) {
+    json messages = json::array();
+    for (const auto & msg : msgs) {
+        json jmsg;
+        if (msg.contains_media() && msg.content.empty()) {
+            common_chat_msg typed = msg;
+            for (auto & part : typed.content_parts) {
+                part.type = "text";
+            }
+            jmsg = typed.to_json_oaicompat(/* concat_typed_text= */ false);
+        } else {
+            jmsg = msg.to_json_oaicompat(/* concat_typed_text= */ false);
+        }
+        if (!msg.extra_json.empty()) {
+            const json extra = json::parse(msg.extra_json);
+            for (const auto & [k, v] : extra.items()) {
+                if (k == "__tool_call_extras") {
+                    if (!jmsg.contains("tool_calls")) {
+                        continue;
+                    }
+                    auto & tcs = jmsg.at("tool_calls");
+                    for (size_t i = 0; i < tcs.size() && i < v.size(); i++) {
+                        for (const auto & [tk, tv] : v[i].items()) {
+                            if (tk == "function") {
+                                for (const auto & [fk, fv] : tv.items()) {
+                                    tcs[i]["function"][fk] = fv;
+                                }
+                            } else {
+                                tcs[i][tk] = tv;
+                            }
+                        }
+                    }
+                } else if (!jmsg.contains(k)) {
+                    jmsg[k] = v;
+                }
+            }
+        }
+        messages.push_back(std::move(jmsg));
+    }
+    return messages_inp_normalizer(c).normalize(messages);
+}
+
 // DEPRECATED: only used in tests
 json common_chat_msgs_to_json_oaicompat(const std::vector<common_chat_msg> & msgs, bool concat_typed_text) {
     jinja::caps c;
@@ -1208,11 +1287,13 @@ static common_chat_params_init_fn common_chat_template_detect_params_init(const
     // DeepSeek V3.2/V4 format detection: template defines dsml_token and uses it for tool calls.
     // The template source contains the token as a variable assignment, not as a literal in markup.
     // V3.2 names the tool call block "function_calls", V4 names it "tool_calls".
-    if (src.find("dsml_token") != std::string::npos &&
-        src.find("DSML") != std::string::npos &&
-        (src.find("function_calls") != std::string::npos ||
-         src.find("tool_calls") != std::string::npos)) {
-        LOG_DBG("Using specialized template: DeepSeek V3.2/V4\n");
+    // V4.1 tag names have a leading space and its template does not name the dsml_token variable.
+    if ((src.find("dsml_token") != std::string::npos &&
+         src.find("DSML") != std::string::npos &&
+         (src.find("function_calls") != std::string::npos ||
+          src.find("tool_calls") != std::string::npos)) ||
+        src.find("｜DSML｜ calls>") != std::string::npos) {
+        LOG_DBG("Using specialized template: DeepSeek V3.2/V4/V4.1\n");
         return common_chat_params_init_deepseek_v3_2;
     }
 
@@ -1313,7 +1394,9 @@ static common_chat_params common_chat_templates_apply_jinja(const struct common_
         workaround::trim_all_content(trimmed_messages);
         messages_to_render = &trimmed_messages;
     }
-    params.messages              = render_message_to_json(*messages_to_render, tmpl.original_caps());
+    params.messages              = src.find("｜DSML｜ calls>") != std::string::npos ?
+                                       render_message_to_json_deepseek_v4_1(*messages_to_render, tmpl.original_caps()) :
+                                       render_message_to_json(*messages_to_render, tmpl.original_caps());
     params.tool_choice           = inputs.tool_choice;
     params.reasoning_format      = inputs.reasoning_format;
     params.enable_thinking       = inputs.enable_thinking;

common/chat.h

@@ -94,6 +94,8 @@ struct common_chat_msg {
     std::string                               reasoning_content;
     std::string                               tool_name;
     std::string                               tool_call_id;
+    // non-OpenAI message keys as a JSON object string, only for templates that read them (DeepSeek V4.1)
+    std::string                               extra_json;
 
     common_json to_json_oaicompat(bool concat_typed_text = false) const;
 

common/jinja/value.cpp

@@ -6,6 +6,7 @@
 #include "json.h"
 
 #include <sstream>
+#include <stdexcept>
 #include <string>
 #include <cctype>
 #include <cfloat>
@@ -299,74 +300,53 @@ static value & get_attribute(const value & val, const value & attr, value & defa
     return default_val;
 }
 
-template<bool is_reject>
-static value selectattr(const func_args & args) {
-    args.ensure_count(2, 4);
-    args.ensure_vals<value_array, value_string, value_string, value_string>(true, true, false, false);
+template<bool is_reject, bool lookup_attr>
+static value select_or_reject(const func_args & args) {
+    // example: array | select("odd") -> test_is_odd(item)
+    // example: array | selectattr("status", "equalto", "active") -> test_is_equalto(item.status, "active")
+    constexpr size_t test_pos = lookup_attr ? 2 : 1;
+    args.ensure_count(test_pos);
+    args.ensure_val<value_array>(args.get_pos(0));
+    value attr;
+    if constexpr (lookup_attr) {
+        attr = args.get_pos(1);
+        args.ensure_val<value_string>(attr);
+    }
+
+    auto & builtins = global_builtins();
+    auto test_it = builtins.end();
+    if (args.count() > test_pos) {
+        args.ensure_val<value_string>(args.get_pos(test_pos));
+        std::string test_name = args.get_pos(test_pos)->as_string().str();
+        test_it = builtins.find("test_is_" + test_name);
+        if (test_it == builtins.end()) {
+            throw raised_exception(args.func_name + ": unknown test '" + test_name + "'");
+        }
+    }
 
     auto arr = args.get_pos(0)->as_array();
-    auto attribute = args.get_pos(1);
     auto out = mk_val<value_array>();
     value val_default = mk_val<value_undefined>();
 
-    if (args.count() == 2) {
-        // example: array | selectattr("active")
-        for (const auto & item : arr) {
-            value attr_val = get_attribute(item, attribute, val_default);
-            bool is_selected = attr_val->as_bool();
-            if constexpr (is_reject) is_selected = !is_selected;
-            if (is_selected) out->push_back(item);
+    for (const auto & item : arr) {
+        value val = item;
+        if constexpr (lookup_attr) {
+            val = get_attribute(item, attr, val_default);
         }
-        return out;
-
-    } else if (args.count() == 3) {
-        // example: array | selectattr("equalto", "text")
-        // translated to: test_is_equalto(item, "text")
-        std::string test_name = args.get_pos(1)->as_string().str();
-        value test_val = args.get_pos(2);
-        auto & builtins = global_builtins();
-        auto it = builtins.find("test_is_" + test_name);
-        if (it == builtins.end()) {
-            throw raised_exception("selectattr: unknown test '" + test_name + "'");
-        }
-        auto test_fn = it->second;
-        for (const auto & item : arr) {
+        bool is_selected;
+        if (test_it == builtins.end()) {
+            is_selected = val->as_bool();
+        } else {
             func_args test_args(args.ctx);
-            test_args.push_back(item); // current object
-            test_args.push_back(test_val); // extra argument
-            value test_result = test_fn(test_args);
-            bool is_selected = test_result->as_bool();
-            if constexpr (is_reject) is_selected = !is_selected;
-            if (is_selected) out->push_back(item);
+            test_args.push_back(val);
+            for (size_t i = test_pos + 1; i < args.count(); i++) {
+                test_args.push_back(args.get_pos(i));
+            }
+            is_selected = test_it->second(test_args)->as_bool();
         }
-        return out;
-
-    } else if (args.count() == 4) {
-        // example: array | selectattr("status", "equalto", "active")
-        // translated to: test_is_equalto(item.status, "active")
-        std::string test_name = args.get_pos(2)->as_string().str();
-        auto extra_arg = args.get_pos(3);
-        auto & builtins = global_builtins();
-        auto it = builtins.find("test_is_" + test_name);
-        if (it == builtins.end()) {
-            throw raised_exception("selectattr: unknown test '" + test_name + "'");
-        }
-        auto test_fn = it->second;
-        for (const auto & item : arr) {
-            value attr_val = get_attribute(item, attribute, val_default);
-            func_args test_args(args.ctx);
-            test_args.push_back(attr_val); // attribute value
-            test_args.push_back(extra_arg); // extra argument
-            value test_result = test_fn(test_args);
-            bool is_selected = test_result->as_bool();
-            if constexpr (is_reject) is_selected = !is_selected;
-            if (is_selected) out->push_back(item);
-        }
-        return out;
-    } else {
-        throw raised_exception("selectattr: invalid number of arguments");
+        if constexpr (is_reject) is_selected = !is_selected;
+        if (is_selected) out->push_back(item);
     }
-
     return out;
 }
 
@@ -1095,10 +1075,10 @@ const func_builtins & value_array_t::get_builtins() const {
             auto arr = slice(val->as_array(), start, stop, step);
             return is_val<value_tuple>(val) ? mk_val<value_tuple>(std::move(arr)) : mk_val<value_array>(std::move(arr));
         }},
-        {"selectattr", selectattr<false>},
-        {"select", selectattr<false>},
-        {"rejectattr", selectattr<true>},
-        {"reject", selectattr<true>},
+        {"selectattr", select_or_reject<false, true>},
+        {"select", select_or_reject<false, false>},
+        {"rejectattr", select_or_reject<true, true>},
+        {"reject", select_or_reject<true, false>},
         {"join", [](const func_args & args) -> value {
             args.ensure_count(1, 3);
             if (!is_val<value_array>(args.get_pos(0))) {
@@ -1460,7 +1440,10 @@ const func_builtins & value_undefined_t::get_builtins() const {
 //////////////////////////////////
 
 
-static value from_json(const common_json & j, bool mark_input) {
+// the runtime walks values with recursion, a deeper input overflows the stack
+static constexpr int FROM_JSON_MAX_DEPTH = 128;
+
+static value from_json(const common_json & j, bool mark_input, int depth = 0) {
     if (j.is_null()) {
         return mk_val<value_none>();
     } else if (j.is_boolean()) {
@@ -1475,16 +1458,18 @@ static value from_json(const common_json & j, bool mark_input) {
             str->mark_input();
         }
         return str;
+    } else if ((j.is_array() || j.is_object()) && depth >= FROM_JSON_MAX_DEPTH) {
+        throw std::invalid_argument("JSON is nested more than " + std::to_string(FROM_JSON_MAX_DEPTH) + " levels deep");
     } else if (j.is_array()) {
         auto arr = mk_val<value_array>();
         for (const auto & item : j) {
-            arr->push_back(from_json(item, mark_input));
+            arr->push_back(from_json(item, mark_input, depth + 1));
         }
         return arr;
     } else if (j.is_object()) {
         auto obj = mk_val<value_object>();
         for (auto it = j.begin(); it != j.end(); ++it) {
-            obj->insert(it.key(), from_json(it.value(), mark_input));
+            obj->insert(it.key(), from_json(it.value(), mark_input, depth + 1));
         }
         return obj;
     } else {

common/json-schema.cpp

@@ -224,7 +224,21 @@ class common_chat_schema_builder {
         return node;
     }
 
+    // the builder and the code that walks the built schema recurse once per level, a deeper schema overflows the stack
+    static constexpr int max_depth = 128;
+    int                  depth_    = 0;
+
     common_chat_schema_ptr build_node(const common_json & schema, const std::string & path) {
+        if (depth_ >= max_depth) {
+            throw std::invalid_argument("JSON schema is nested more than " + std::to_string(max_depth) + " levels deep");
+        }
+        depth_++;
+        auto node = build_node_impl(schema, path);
+        depth_--;
+        return node;
+    }
+
+    common_chat_schema_ptr build_node_impl(const common_json & schema, const std::string & path) {
         if (!schema.is_object()) {
             fail(path, "schema must be an object");
         }

common/parsers/deepseek.cpp

@@ -1,4 +1,5 @@
 #include "parsers.h"
+#include "log.h"
 
 // The DeepSeek V4 reference implementation renders consecutive tool results into a single
 // user block, ordered by the tool call order of the preceding assistant message (matched
@@ -75,13 +76,30 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
     // non-thinking generation prompt ends with a bare </think> instead of an empty
     // <think></think> pair.
     const bool is_v4 = tmpl.source().find("function_calls") == std::string::npos;
+    // V4.1: DSML tag names carry a leading space, e.g. <｜DSML｜ calls>
+    const bool is_v41 = is_v4 && tmpl.source().find("｜DSML｜ calls>") != std::string::npos;
 
     std::optional<json> adjusted_messages;
     if (is_v4) {
         adjusted_messages = deepseek_v4_sort_tool_results(inputs.messages);
     }
 
-    auto has_tools           = inputs.tools.is_array() && !inputs.tools.empty();
+    // V4.1 can declare tools on a system message instead of top-level "tools": parse calls to those too.
+    json tools = inputs.tools;
+    if (is_v41 && !(tools.is_array() && !tools.empty()) && inputs.messages.is_array()) {
+        for (const auto & msg : inputs.messages) {
+            if (msg.value("role", "") == "system" && msg.contains("tools") && msg.at("tools").is_array()) {
+                if (!tools.is_array()) {
+                    tools = json::array();
+                }
+                for (const auto & tool : msg.at("tools")) {
+                    tools.push_back(tool);
+                }
+            }
+        }
+    }
+
+    auto has_tools           = tools.is_array() && !tools.empty();
     auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object();
     auto extract_reasoning   = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE;
     auto include_grammar     = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE);
@@ -90,17 +108,39 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
     if (is_v4 && has_response_format) {
         additional_context = json{ { "response_format", inputs.json_schema } };
     }
+    // V4.1 raises on anything but low/high/max or 1-100: map minimal -> low, xhigh -> max, medium -> default
+    if (is_v41 && inputs.extra_context.is_object() && inputs.extra_context.contains("reasoning_effort") &&
+        inputs.extra_context.at("reasoning_effort").is_string()) {
+        const std::string effort = inputs.extra_context.at("reasoning_effort").get<std::string>();
+        json mapped;
+        if (effort == "minimal") {
+            mapped = "low";
+        } else if (effort == "xhigh") {
+            mapped = "max";
+        } else if (effort == "medium") {
+            mapped = nullptr;
+        }
+        if (effort == "minimal" || effort == "xhigh" || effort == "medium") {
+            LOG_DBG("%s: DeepSeek V4.1 reasoning_effort '%s' -> %s\n", __func__, effort.c_str(),
+                    mapped.is_null() ? "unset (template default)" : mapped.get<std::string>().c_str());
+            if (!additional_context.has_value()) {
+                additional_context = json::object();
+            }
+            (*additional_context)["reasoning_effort"] = mapped;
+        }
+    }
 
     const std::string DSML         = "｜DSML｜";
     const std::string THINK_START  = "<think>";
     const std::string THINK_END    = "</think>";
-    const std::string TC_BLOCK     = is_v4 ? "tool_calls" : "function_calls";
-    const std::string FC_START     = "<" + DSML + TC_BLOCK + ">";
-    const std::string FC_END       = "</" + DSML + TC_BLOCK + ">";
-    const std::string INVOKE_START = "<" + DSML + "invoke";
-    const std::string INVOKE_END   = "</" + DSML + "invoke>";
-    const std::string PARAM_START  = "<" + DSML + "parameter";
-    const std::string PARAM_END    = "</" + DSML + "parameter>";
+    const std::string TAG_SP       = is_v41 ? " " : "";
+    const std::string TC_BLOCK     = is_v41 ? "calls" : is_v4 ? "tool_calls" : "function_calls";
+    const std::string FC_START     = "<" + DSML + TAG_SP + TC_BLOCK + ">";
+    const std::string FC_END       = "</" + DSML + TAG_SP + TC_BLOCK + ">";
+    const std::string INVOKE_START = "<" + DSML + TAG_SP + "invoke";
+    const std::string INVOKE_END   = "</" + DSML + TAG_SP + "invoke>";
+    const std::string PARAM_START  = "<" + DSML + TAG_SP + "parameter";
+    const std::string PARAM_END    = "</" + DSML + TAG_SP + "parameter>";
     const std::string GEN_PROMPT   = "<｜Assistant｜>";
     const std::string TC_SEPARATOR = "\n\n";
 
@@ -152,7 +192,7 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
         // build tool call section first since we might need it in reasoning
         auto tool_choice = p.choice();
         if (has_tool_calls) {
-            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
+            foreach_function(tools, [&](size_t tool_index, const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
 
@@ -223,7 +263,10 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
         }
 
         if (extract_reasoning && inputs.enable_thinking) {
-            reasoning = p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END);
+            // V4.1 renders a continuation without reasoning as "<｜Assistant｜></think>" even in thinking mode,
+            // so consume the bare THINK_END too or it leaks into the content
+            reasoning = is_v41 ? p.optional((THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END) | p.literal(THINK_END))
+                               : p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END);
             reasoning_with_tc = THINK_START +
                 p.reasoning(p.until_one_of({ TC_SEPARATOR + FC_START, FC_START, THINK_END })) +
                 p.space() + obligatory_tool_calls;

conversion/__init__.py

@@ -70,6 +70,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "LingDSparkModel": "qwen",
     "DeepseekV4ForCausalLM": "deepseek",
     "DeepseekV4DSparkModel": "deepseek",
+    "DeepseekV41ForCausalLM": "deepseek41",
     "DistilBertForMaskedLM": "bert",
     "DistilBertForSequenceClassification": "bert",
     "DistilBertModel": "bert",
@@ -326,6 +327,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
     "VllmD3Model": "vllm_d3",
     "DeepseekOCR2ForCausalLM": "deepseek",
     "DeepseekOCRForCausalLM": "deepseek",
+    "DeepseekV41ForCausalLM": "deepseek41",
     "DeepseekV4ForCausalLM": "deepseek",
     "Dots3NoteForCausalLM": "dots3",
     "Dots3NoteForConditionalGeneration": "dots3",

conversion/deepseek41.py

@@ -0,0 +1,467 @@
+from __future__ import annotations
+
+import json
+import math
+import re
+from pathlib import Path
+from typing import Any, Callable, Iterable, TYPE_CHECKING, cast
+
+import numpy as np
+import torch
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import LazyTorchTensor, ModelBase, TextModel, gguf
+from .deepseek import DeepseekV4FlashVisionModel, DeepseekV4Model
+
+
+# Engram hash constants from inference/engram.py, stored in the GGUF: the token map needs the HF tokenizers
+# normalizers and the multipliers numpy's PCG64 stream, neither of which the runtime has.
+
+def _is_prime(n: int) -> bool:
+    if n < 2 or n % 2 == 0:
+        return n == 2
+    return all(n % d for d in range(3, math.isqrt(n) + 1, 2))
+
+
+def engram_primes(n_layers: int, max_ngram_size: int, n_heads: int, vocab_size: int) -> list[list[list[int]]]:
+    """EngramLayout.from_args: [layer][ngram order 2..max][head] bucket moduli, never reused across layers."""
+    seen: set[int] = set()
+    out: list[list[list[int]]] = []
+    for _ in range(n_layers):
+        per_ngram = []
+        for _ in range(max_ngram_size - 1):
+            sizes, cur = [], vocab_size - 1
+            for _ in range(n_heads):
+                cand = cur + 1
+                while not _is_prime(cand) or cand in seen:
+                    cand += 1
+                cur = cand
+                seen.add(cur)
+                sizes.append(cur)
+            per_ngram.append(sizes)
+        out.append(per_ngram)
+    return out
+
+
+def engram_multipliers(layer_ids: list[int], max_ngram_size: int, compressed_vocab_size: int) -> list[list[int]]:
+    """compute_hash_multipliers: odd, bounded so that token_id * multiplier cannot overflow int64."""
+    bound = max(1, (np.iinfo(np.int64).max // compressed_vocab_size) // 2)
+    rows = []
+    for layer_id in layer_ids:
+        rng = np.random.default_rng(10007 * layer_id)
+        values = rng.integers(low=0, high=bound, size=(max_ngram_size,), dtype=np.int64)
+        rows.append([int(v) * 2 + 1 for v in values])
+    return rows
+
+
+def engram_compressed_token_map(tokenizer_dir: Path) -> tuple[list[int], int]:
+    """build_compressed_token_map: token id -> compressed id (tokens that normalize alike collapse)."""
+    from tokenizers import Regex, normalizers
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)
+    sentinel = "\ue000"
+    normalizer = normalizers.Sequence([
+        normalizers.NFKC(),
+        normalizers.NFD(),
+        normalizers.StripAccents(),
+        normalizers.Lowercase(),
+        normalizers.Replace(Regex(r"[ \t\r\n]+"), " "),
+        normalizers.Replace(Regex(r"^ $"), sentinel),
+        normalizers.Strip(),
+        normalizers.Replace(sentinel, " "),
+    ])
+    backend = tokenizer.backend_tokenizer
+    key_to_new: dict[str, int] = {}
+    lookup = [0] * len(tokenizer)
+    for token_id in range(len(tokenizer)):
+        text = backend.decode([token_id], skip_special_tokens=False)
+        if "\ufffd" in text:
+            key = backend.id_to_token(token_id)
+        else:
+            normalized = normalizer.normalize_str(text)
+            key = normalized if normalized else text
+        lookup[token_id] = key_to_new.setdefault(key, len(key_to_new))
+    return lookup, len(key_to_new)
+
+
+class _StreamedBytes:
+    """Raw tensor bytes produced chunk by chunk at write time (an Engram table is ~200 GB in BF16)."""
+    def __init__(self, shape: tuple[int, ...], chunks: Callable[[], Iterable[np.ndarray]]):
+        self.shape = shape
+        self.dtype = np.dtype(np.uint8)
+        self.nbytes = int(np.prod(shape))
+        self._chunks = chunks
+
+    def tofile(self, fout) -> None:
+        written = 0
+        for chunk in self._chunks():
+            chunk = np.ascontiguousarray(chunk, dtype=np.uint8)
+            chunk.tofile(fout)
+            written += chunk.nbytes
+        if written != self.nbytes:
+            raise ValueError(f"streamed {written} bytes, expected {self.nbytes}")
+
+    def byteswap(self, inplace: bool = False):
+        raise NotImplementedError("streamed tensors are written little-endian only")
+
+
+def _e8m0_check(scale_u8: np.ndarray, name: str) -> None:
+    # e4m3 (|x| in [2^-9, 448]) * 2^(e - 127) is exact in BF16 for e in [3, 246]
+    if scale_u8.size and (int(scale_u8.min()) < 3 or int(scale_u8.max()) > 246):
+        raise ValueError(f"{name}: E8M0 scale exponents outside [3, 246], the dequantized values would not be exact")
+
+
+@ModelBase.register("DeepseekV41ForCausalLM")
+class DeepseekV41Model(DeepseekV4Model):
+    """DeepSeek-V4.1 text model, lossless: FP8 (e4m3, 32x32 ue8m0 blocks) linears dequantized to BF16 (F32 with
+    --outtype f32), routed experts copied bit-exact as MXFP4, Engram tables (e4m3 rows, one e8m0 scale per 32) as the
+    exact BF16 lookup (F32 with --outtype f32), router and Engram gate F32. DSpark (mtp.*) is not converted here
+    (--dspark, DeepseekV41DSparkModel); the vision tower goes to an mmproj (--mmproj, DeepseekV41VisionModel).
+    The Engram tables of all layers are written as one per_layer_token_embd (rows concatenated in layer order) with
+    the PLE keys of qwen4exp: absolute head offsets and bucket counts, so the runtime can read it lazily."""
+    model_arch = gguf.MODEL_ARCH.DEEPSEEK41
+    supports_mtp_export = False
+
+    _ENGRAM_EMBED_RE = re.compile(r"layers\.(\d+)\.engram\.embed\.(weight|scale)$")
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.hparams.setdefault("num_hash_layers", 0)  # no hash-routed layers in V4.1
+        self._f32 = self.ftype == gguf.LlamaFileType.ALL_F32
+        qc = self.hparams.get("quantization_config") or {}
+        self._fp8_block = tuple(qc.get("weight_block_size") or (32, 32))
+        if qc.get("scale_fmt", "ue8m0") != "ue8m0":
+            raise ValueError(f"unsupported FP8 scale format {qc.get('scale_fmt')!r}")
+        if self.hparams.get("gate_temp", 1.0) != 1.0:
+            raise ValueError("gate_temp != 1 is not supported (the router logits would need dividing by it)")
+
+        # V4.1 does not use the DeepSeek-V4 template; a template shipped with the model is still added by set_vocab
+        self.gguf_writer.remove_key(gguf.Keys.Tokenizer.CHAT_TEMPLATE)
+
+    def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
+        # num_hidden_layers lives in text_config, which is merged only after indexing; MTP stages are skipped anyway
+        return TextModel.index_tensors(self, remote_hf_model_id=remote_hf_model_id)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        if item[0].startswith("vision."):
+            return None
+        return super().filter_tensors(item)
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        hp = self.hparams
+        w = self.gguf_writer
+        w.add_attention_kv_source_layers(hp["kv_source_layer_ids"])
+        w.add_indexer_source_layers(hp["index_source_layer_ids"])
+        if hp.get("candidate_source_layer_id", -1) >= 0:
+            w.add_indexer_candidate_source_layer(hp["candidate_source_layer_id"])
+            w.add_indexer_candidate_top_k_blocks(hp["candidate_topk_blocks"])
+            w.add_indexer_candidate_block_size(hp["candidate_block_size"])
+
+        # GGUF cannot hold empty arrays: the Engram keys are omitted when the model has no Engram layers
+        layer_ids = list(hp.get("engram_layer_ids") or [])
+        if layer_ids:
+            max_ngram, n_heads = hp["engram_max_ngram_size"], hp["engram_n_heads"]
+            primes = engram_primes(len(layer_ids), max_ngram, n_heads, hp["engram_vocab_size"])
+            for li, layer_primes in enumerate(primes):
+                if sum(p for per in layer_primes for p in per) != hp["engram_num_embeddings"][li]:
+                    raise ValueError(f"engram layer {layer_ids[li]}: prime sum != num_embeddings")
+            token_map, cvocab = engram_compressed_token_map(self.dir_model)
+            if cvocab != hp["engram_compressed_vocab_size"] or len(token_map) != hp["vocab_size"]:
+                raise ValueError(f"compressed token map ({len(token_map)} -> {cvocab}) does not match the config")
+            # absolute row offsets into the concatenated table
+            sizes = [p for layer in primes for per in layer for p in per]
+            offsets = [sum(sizes[:i]) for i in range(len(sizes))]
+            w.add_key_value(gguf.Keys.PerLayerEmbedding.LAYERS.format(arch=w.arch), layer_ids,
+                            gguf.GGUFValueType.ARRAY, gguf.GGUFValueType.UINT32)
+            w.add_ple_ngram_size(max_ngram)
+            w.add_ple_heads_per_ngram(n_heads)
+            w.add_embedding_length_per_layer_input(hp["engram_head_dim"])
+            w.add_ple_layer_multipliers([m for row in engram_multipliers(layer_ids, max_ngram, cvocab) for m in row])
+            w.add_ple_head_offsets(offsets)
+            w.add_ple_head_vocab_sizes(sizes)
+            w.add_engram_pad_token_id(hp["engram_pad_token_id"])
+            w.add_engram_token_map(token_map)
+
+    def dequant_model(self):
+        bo, bi = self._fp8_block
+
+        def dequant(weight: Tensor, scale: Tensor) -> Tensor:
+            out_f, in_f = weight.shape
+            s = self._e8m0_to_float(scale).repeat_interleave(bo, 0)[:out_f].repeat_interleave(bi, 1)[:, :in_f]
+            return weight.float() * s
+
+        for name in [n for n in self.model_tensors if n.endswith(".scale") and not self._ENGRAM_EMBED_RE.match(n)]:
+            weight_name = name.removesuffix(".scale") + ".weight"
+            if weight_name not in self.model_tensors or self.model_tensors[weight_name]().dtype not in self._float8_dtypes():
+                continue
+            w, s = self.model_tensors[weight_name], self.model_tensors.pop(name)
+            out_f, in_f = w().shape
+            if tuple(s().shape) != (-(-out_f // bo), -(-in_f // bi)):
+                raise ValueError(f"{weight_name}: scale shape {tuple(s().shape)} does not match {(bo, bi)} blocks")
+            _e8m0_check(LazyTorchTensor.to_eager(s()).view(torch.uint8).numpy(), name)
+            self.model_tensors[weight_name] = lambda w=w, s=s: dequant(w(), s())
+            self._dsv4_fp8_dequantized.add(weight_name)
+
+    def _tensor_part(self, name: str) -> Path:
+        index = self.dir_model / "model.safetensors.index.json"
+        if not index.is_file():
+            return self.dir_model / "model.safetensors"
+        with open(index, "r", encoding="utf-8") as f:
+            return self.dir_model / json.load(f)["weight_map"][name]
+
+    def _write_mxfp4_experts(self, bid: int, proj: str, key: gguf.MODEL_TENSOR, n_expert: int | None = None) -> None:
+        n_expert = self.hparams["n_routed_experts"] if n_expert is None else n_expert
+        names = [f"layers.{bid}.ffn.experts.{e}.{proj}" for e in range(n_expert)]
+        pairs = [(self.model_tensors.pop(n + ".weight"), self.model_tensors.pop(n + ".scale")) for n in names]
+        rows, packed_cols = pairs[0][0]().shape
+
+        def chunks() -> Iterable[np.ndarray]:
+            for w, s in pairs:
+                yield self.repack_mxfp4_blocks(LazyTorchTensor.to_eager(w()), LazyTorchTensor.to_eager(s()))
+        data = cast(np.ndarray, _StreamedBytes((len(pairs), rows, packed_cols * 2 // 32 * 17), chunks))
+        self.gguf_writer.add_tensor(self.format_tensor_name(key, bid), data, raw_dtype=gguf.GGMLQuantizationType.MXFP4)
+
+    def _write_engram_tables(self, layer_ids: list[int]) -> None:
+        # official lookup e4m3 * 2^(e8m0 - 127), exact in BF16
+        parts, dim = [], None
+        for bid in layer_ids:
+            wn, sn = f"layers.{bid}.engram.embed.weight", f"layers.{bid}.engram.embed.scale"
+            rows, d = self.model_tensors.pop(wn)().shape
+            if d % 32 != 0 or tuple(self.model_tensors.pop(sn)().shape) != (rows, d // 32) or d != (dim or d):
+                raise ValueError(f"{wn}: expected one e8m0 scale per 32 elements and the same width in every layer")
+            if rows != self.hparams["engram_num_embeddings"][len(parts)]:
+                raise ValueError(f"{wn}: {rows} rows, config says {self.hparams['engram_num_embeddings'][len(parts)]}")
+            dim = d
+            parts.append((wn, sn, rows, self._tensor_part(wn), self._tensor_part(sn)))
+        assert dim is not None
+        name = self.format_tensor_name(gguf.MODEL_TENSOR.PER_LAYER_TOKEN_EMBD)
+        n_rows = sum(p[2] for p in parts)
+
+        def chunks() -> Iterable[np.ndarray]:
+            from safetensors import safe_open
+            for wn, sn, rows, wpart, spart in parts:
+                with safe_open(wpart, framework="pt") as fw, safe_open(spart, framework="pt") as fs:
+                    ws, ss = fw.get_slice(wn), fs.get_slice(sn)
+                    for r0 in range(0, rows, 1 << 18):
+                        q = ws[r0:r0 + (1 << 18)].view(torch.uint8).reshape(-1, 32)
+                        e = ss[r0:r0 + (1 << 18)].view(torch.uint8).reshape(-1, 1)
+                        _e8m0_check(e.numpy(), sn)
+                        v = (q.view(torch.float8_e4m3fn).float() * torch.exp2(e.float() - 127)).reshape(-1, dim)
+                        yield (v if self._f32 else v.bfloat16()).view(torch.uint8).numpy()
+        qt = gguf.GGMLQuantizationType.F32 if self._f32 else gguf.GGMLQuantizationType.BF16
+        data = cast(np.ndarray, _StreamedBytes((n_rows, dim * (4 if self._f32 else 2)), chunks))
+        self.gguf_writer.add_tensor(name, data, raw_dtype=qt)
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        if not self._dsv4_mxfp4_generated:
+            for bid in range(self.block_count):
+                self._write_mxfp4_experts(bid, "w1", gguf.MODEL_TENSOR.FFN_GATE_EXP)
+                self._write_mxfp4_experts(bid, "w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP)
+                self._write_mxfp4_experts(bid, "w3", gguf.MODEL_TENSOR.FFN_UP_EXP)
+            if self.hparams.get("engram_layer_ids"):
+                self._write_engram_tables(list(self.hparams["engram_layer_ids"]))
+            self._dsv4_mxfp4_generated = True
+        return ()
+
+    _LAYER_MAP: dict[str, tuple[gguf.MODEL_TENSOR, str]] = {
+        "attn.indexer.wk.weight": (gguf.MODEL_TENSOR.INDEXER_ATTN_K, ".weight"),
+        "attn.indexer.k_norm.weight": (gguf.MODEL_TENSOR.INDEXER_K_NORM, ".weight"),
+        "engram.wkv.weight": (gguf.MODEL_TENSOR.ENGRAM_KV, ".weight"),
+    }
+
+    def _map_dsv4_tensor_name(self, name: str, bid: int | None) -> tuple[gguf.MODEL_TENSOR, str]:
+        m = re.match(r"layers\.\d+\.(.+)$", name)
+        if m is not None and m.group(1) in self._LAYER_MAP:
+            return self._LAYER_MAP[m.group(1)]
+        return super()._map_dsv4_tensor_name(name, bid)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name.endswith(".engram.k_weight"):
+            return []  # folded into engram_gate
+        if name.endswith(".engram.q_weight"):
+            # Engram gate weight = q_weight * k_weight (bf16 * bf16 is exact in f32)
+            k = self.model_tensors[name.removesuffix("q_weight") + "k_weight"]
+            return [(self.format_tensor_name(gguf.MODEL_TENSOR.ENGRAM_GATE, bid), data_torch.float() * k().float())]
+        return super().modify_tensors(data_torch, name, bid)
+
+    def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
+        if new_name.endswith((".engram_gate.weight", ".ffn_gate_inp.weight")) or name in self._dsv4_f32_tensors:
+            return gguf.GGMLQuantizationType.F32
+        if n_dims >= 2 and (name in self._dsv4_fp8_dequantized or name in self._dsv4_bf16_tensors):
+            return gguf.GGMLQuantizationType.F32 if self._f32 else gguf.GGMLQuantizationType.BF16
+        return False
+
+
+class DeepseekV41DSparkModel(DeepseekV41Model):
+    """DeepSeek-V4.1 DSpark speculative draft (mtp.0..2 of the checkpoint) as a DFLASH GGUF (`--dspark`).
+
+    Layout = the V4 DSpark draft (conversion/deepseek.py DeepseekV4DSparkModel) with V4.1's differences flagged by
+    `dflash.hyper_connection.single_pass`: stages blk.0..2 (MXFP4 experts bit-exact, dense BF16), mtp.0.main_proj /
+    main_norm -> fc / enc.output_norm, mtp.2.norm -> output_norm, markov_head.embed / head -> markov_w1 / markov_w2,
+    confidence_head.proj -> conf_proj (F32), no hc_head. Token embeddings and the lm head are the target's (shared at
+    runtime through ctx_other). target_layers are the V4.1 ids as-is: the draft reads the attention input of those
+    layers, which is what llama_set_embeddings_layer_inp extracts.
+    """
+    model_arch = gguf.MODEL_ARCH.DFLASH
+    mtp_only = True
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        hp = self.hparams
+        self._n_main = hp["num_hidden_layers"]
+        self._n_stage = hp["num_nextn_predict_layers"]
+        if self._n_stage <= 0 or "dspark_block_size" not in hp:
+            raise ValueError("DeepSeek-V4.1 DSpark export needs num_nextn_predict_layers > 0 and the dspark_* config")
+        # stages keep main-model block ids (mtp.M -> layers.n_main+M); format_tensor_name() maps them back
+        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self._n_stage)
+
+    def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
+        hp = {**self.hparams, **self.hparams.get("text_config", {})}  # text_config is merged only after indexing
+        type(self)._dsv4_main_layers = hp["num_hidden_layers"]
+        type(self)._dsv4_nextn_layers = hp.get("num_nextn_predict_layers", 0)
+        return TextModel.index_tensors(self, remote_hf_model_id=remote_hf_model_id)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+        m = re.match(r"mtp\.(\d+)\.(.+)$", name)
+        if m is None:
+            return None  # shared with the target at runtime
+        stage, rest = int(m.group(1)), m.group(2)
+        assert cls._dsv4_main_layers is not None
+        if stage >= cls._dsv4_nextn_layers:
+            raise ValueError(f"Unexpected DSpark stage {stage}")
+        if rest.startswith(("markov_head.", "confidence_head.")):
+            return f"dspark.{rest}", gen
+        return f"layers.{cls._dsv4_main_layers + stage}.{rest}", gen
+
+    def format_tensor_name(self, key: gguf.MODEL_TENSOR, bid: int | None = None, suffix: str = ".weight") -> str:
+        if bid is not None and bid >= self._n_main:
+            bid -= self._n_main
+        return super().format_tensor_name(key, bid, suffix)
+
+    def set_vocab(self):
+        super().set_vocab()
+        self.gguf_writer.add_mask_token_id(self.hparams["dspark_noise_token_id"])
+
+    def set_gguf_parameters(self):
+        hp = self.hparams
+        n = self._n_stage
+        w = self.gguf_writer
+        w.add_block_count(n)
+        w.add_context_length(hp["max_position_embeddings"])
+        w.add_embedding_length(hp["hidden_size"])
+        w.add_feed_forward_length(hp["moe_intermediate_size"])
+        w.add_head_count(hp["num_attention_heads"])
+        w.add_head_count_kv(1)
+        w.add_key_length(hp["head_dim"])
+        w.add_value_length(hp["head_dim"])
+        w.add_rope_dimension_count(hp["qk_rope_head_dim"])
+        # compress-ratio-0 layers: plain rope at rope_theta, no YaRN
+        w.add_rope_freq_base(float(hp["rope_theta"]))
+        w.add_layer_norm_rms_eps(hp["rms_norm_eps"])
+        w.add_q_lora_rank(hp["q_lora_rank"])
+        w.add_sliding_window(hp["sliding_window"])
+        w.add_expert_count(hp["dspark_n_routed_experts"])
+        w.add_expert_used_count(hp["dspark_num_experts_per_tok"])
+        w.add_expert_feed_forward_length(hp["moe_intermediate_size"])
+        w.add_expert_shared_count(hp["n_shared_experts"])
+        w.add_expert_weights_scale(hp["routed_scaling_factor"])
+        w.add_expert_weights_norm(hp["norm_topk_prob"])
+        w.add_expert_gating_func(gguf.ExpertGatingFuncType.SQRTSOFTPLUS)
+        w.add_swiglu_clamp_exp([hp["swiglu_limit"]] * n)
+        w.add_swiglu_clamp_shexp([hp["swiglu_limit"]] * n)
+        w.add_attention_output_group_count(hp["o_groups"])
+        w.add_attention_output_lora_rank(hp["o_lora_rank"])
+        w.add_attention_compress_ratios([0] * n)
+        w.add_hyper_connection_count(hp["hc_mult"])
+        w.add_hyper_connection_sinkhorn_iterations(hp["hc_sinkhorn_iters"])
+        w.add_hyper_connection_epsilon(hp["hc_eps"])
+        w.add_hyper_connection_single_pass(True)
+        w.add_block_size(hp["dspark_block_size"])
+        w.add_target_layers(hp["dspark_target_layer_ids"])
+        w.add_sample_from_anchor(True)
+        w.add_has_confidence_head(True)
+        w.add_file_type(self.ftype)
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        # no Engram tables in the draft
+        if not self._dsv4_mxfp4_generated:
+            n_expert = self.hparams["dspark_n_routed_experts"]
+            for bid in range(self._n_main, self._n_main + self._n_stage):
+                self._write_mxfp4_experts(bid, "w1", gguf.MODEL_TENSOR.FFN_GATE_EXP, n_expert)
+                self._write_mxfp4_experts(bid, "w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP, n_expert)
+                self._write_mxfp4_experts(bid, "w3", gguf.MODEL_TENSOR.FFN_UP_EXP, n_expert)
+            self._dsv4_mxfp4_generated = True
+        return ()
+
+    _DSPARK_ROOT: dict[str, gguf.MODEL_TENSOR] = {
+        "main_proj.weight": gguf.MODEL_TENSOR.FC,
+        "main_norm.weight": gguf.MODEL_TENSOR.ENC_OUTPUT_NORM,
+        "norm.weight": gguf.MODEL_TENSOR.OUTPUT_NORM,
+        "dspark.markov_head.embed.weight": gguf.MODEL_TENSOR.DSPARK_MARKOV_W1,
+        "dspark.markov_head.head.weight": gguf.MODEL_TENSOR.DSPARK_MARKOV_W2,
+        "dspark.confidence_head.proj.weight": gguf.MODEL_TENSOR.DSPARK_CONF_PROJ,
+    }
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name.endswith(".ffn.gate.bias_vl"):
+            return []
+        if name in self._DSPARK_ROOT:
+            return [(self.format_tensor_name(self._DSPARK_ROOT[name], None), data_torch)]
+        m = re.match(r"layers\.\d+\.(main_proj\.weight|main_norm\.weight|norm\.weight)$", name)
+        if m is not None:
+            return [(self.format_tensor_name(self._DSPARK_ROOT[m.group(1)], None), data_torch)]
+        return super().modify_tensors(data_torch, name, bid)
+
+    def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
+        if new_name == "conf_proj.weight":
+            return gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+
+@ModelBase.register("DeepseekV41ForCausalLM")
+@ModelBase.example("deepseek-ai/DeepSeek-V4.1-Flash-1017")
+class DeepseekV41VisionModel(DeepseekV4FlashVisionModel):
+    """DeepSeek-V4.1 vision tower (projector deepseek41v): the V4-Flash-Vision ViT + 3x3 aligner, with the V4.1 token
+    block ([START] ([IMAGE]*w [NEWLINE])*h [END], no pads) and resize plan. ref: inference/vision.py,
+    inference/image_processor.py of deepseek-ai/DeepSeek-V4.1-Flash-1017. The checkpoint config is HF-style
+    (vision_config); image_start / image_end / image_newline are written like V4's sentinels, there is no image_pad."""
+
+    def get_vision_config(self) -> dict[str, Any] | None:
+        vc = self.global_config.get("vision_config")
+        if not vc or vc.get("num_hidden_layers", 0) == 0:
+            raise ValueError("DeepseekV41VisionModel requires vision_config.num_hidden_layers > 0")
+        return {
+            "num_hidden_layers":   vc["num_hidden_layers"],
+            "hidden_size":         vc["hidden_size"],
+            "num_attention_heads": vc["num_attention_heads"],
+            "intermediate_size":   vc["intermediate_size"],
+            "patch_size":          vc["patch_size"],
+            # dynamic resolution; only used for compat / warmup
+            "image_size":          vc["patch_size"] * vc["downsample_ratio"] * 16,
+            "rope_theta":          vc.get("rope_theta", 10000.0),
+            "downsample_ratio":    vc["downsample_ratio"],
+            "min_pixels":          vc["min_pixels"],
+            "max_image_tokens":    vc["max_image_tokens"],
+            "max_wh_ratio":        vc.get("max_wh_ratio"),
+        }
+
+    def set_gguf_parameters(self):
+        assert self.hparams_vision is not None
+        super(DeepseekV4FlashVisionModel, self).set_gguf_parameters()
+        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.DEEPSEEK41V)
+        # pytorch default eps (vision.py), not the LLM rms_norm_eps (1e-20)
+        self.gguf_writer.add_vision_attention_layernorm_eps(1e-6)
+        self.gguf_writer.add_vision_use_silu(True)
+        self.gguf_writer.add_vision_projector_scale_factor(self.hparams_vision["downsample_ratio"])
+        self.gguf_writer.add_vision_min_pixels(self.hparams_vision["min_pixels"])
+        # hardcoded in clip.cpp (PROJECTOR_TYPE_DEEPSEEK41V); add GGUF keys if a later model differs
+        assert self.hparams_vision["max_image_tokens"] == 1024, self.hparams_vision["max_image_tokens"]
+        assert self.hparams_vision["max_wh_ratio"] is None, self.hparams_vision["max_wh_ratio"]
+        assert self.hparams_vision["rope_theta"] == 10000, self.hparams_vision["rope_theta"]

convert_hf_to_gguf.py

@@ -267,11 +267,15 @@ def main() -> None:
             sys.exit(1)
 
         if args.dspark:
-            if is_mistral_format or model_architecture != "DeepseekV4ForCausalLM":
-                logger.error("--dspark is only supported for DeepseekV4ForCausalLM")
+            if is_mistral_format or model_architecture not in ("DeepseekV4ForCausalLM", "DeepseekV41ForCausalLM"):
+                logger.error("--dspark is only supported for DeepseekV4ForCausalLM and DeepseekV41ForCausalLM")
                 sys.exit(1)
-            from conversion.deepseek import DeepseekV4DSparkModel
-            model_class = DeepseekV4DSparkModel
+            if model_architecture == "DeepseekV41ForCausalLM":
+                from conversion.deepseek41 import DeepseekV41DSparkModel
+                model_class = DeepseekV41DSparkModel
+            else:
+                from conversion.deepseek import DeepseekV4DSparkModel
+                model_class = DeepseekV4DSparkModel
 
         if args.mtp or args.no_mtp:
             if not model_class.supports_mtp_export:

ggml/src/ggml-cpu/unary-ops.cpp

@@ -78,7 +78,7 @@ static inline float op_expm1(float x) {
 }
 
 static inline float op_softplus(float x) {
-    return (x > 20.0f) ? x : logf(1.0f + expf(x));
+    return (x > 20.0f) ? x : log1pf(expf(x));
 }
 
 static inline float op_floor(float x) {

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

@@ -1973,6 +1973,61 @@ static bool ggml_cuda_match_shared_expert(const ggml_cgraph * graph, int routed_
     return true;
 }
 
+// Buffer- and batch-size independent, so graph_optimize produces the same graph topology for every ubatch size (see #29986).
+static bool ggml_cuda_mul_mat_q_fusion_matches(const ggml_tensor * up, const ggml_tensor * gate, const ggml_tensor * glu, const int device) {
+    const ggml_tensor * x_up = up->src[0];
+    const ggml_tensor * y    = up->src[1];
+
+    if (up->op != GGML_OP_MUL_MAT || !ggml_cuda_mmq_fusion_supported(x_up->type) || y->type != GGML_TYPE_F32) {
+        return false;
+    }
+
+    // The fused kernel is only compiled for weights whose rows fill whole tiles (no fallback).
+    if (ggml_cuda_mmq_needs_fallback(x_up->ne[1])) {
+        return false;
+    }
+
+    if (!ggml_cuda_should_fuse_mul_mat(up, gate, glu)) {
+        return false;
+    }
+
+    // The fused write-back applies SWIGLU_OAI with the default alpha and limit of ggml_cuda_op_swiglu_oai_single.
+    if (ggml_get_glu_op(glu) == GGML_GLU_OP_SWIGLU_OAI && (ggml_get_op_params_f32(glu, 2) != 1.702f || ggml_get_op_params_f32(glu, 3) != 7.0f)) {
+        return false;
+    }
+
+    // NVIDIA only: the fused write-back relies on the MMA accumulator layout.
+    return turing_mma_available(ggml_cuda_info().devices[device].cc);
+}
+
+static bool ggml_cuda_should_fuse_mul_mat_q(const ggml_tensor * up, const ggml_tensor * gate, const ggml_tensor * glu) {
+    // A hinted matmul can be fused, but its unfused path may be faster (e.g. FWHT for GGML_HINT_SRC0_IS_HADAMARD).
+    if (ggml_get_op_params_i32(up, 1) != GGML_HINT_NONE || ggml_get_op_params_i32(gate, 1) != GGML_HINT_NONE) {
+        return false;
+    }
+
+    // The padding of a compute-buffer weight that is a view cannot be cleared, see ggml_cuda_mul_mat.
+    const auto bad_padding_clear = [](const ggml_tensor * x) {
+        return ggml_backend_buffer_get_usage(x->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE &&
+            ggml_nbytes(x) != ggml_backend_buffer_get_alloc_size(x->buffer, x) && x->view_src;
+    };
+    const int device = ggml_cuda_get_device();
+    if (!ggml_cuda_mul_mat_q_fusion_matches(up, gate, glu, device) || bad_padding_clear(up->src[0]) || bad_padding_clear(gate->src[0])) {
+        return false;
+    }
+
+    // The batch size checks are done here, not in ggml_cuda_mul_mat_q_fusion_matches.
+    const ggml_tensor * x_up = up->src[0];
+    const ggml_tensor * y    = up->src[1];
+    const int           cc   = ggml_cuda_info().devices[device].cc;
+    if (ggml_cuda_should_use_mmvq(x_up->type, cc, y->ne[1]) || !ggml_cuda_should_use_mmq(x_up->type, cc, y->ne[1], 0)) {
+        return false;
+    }
+
+    // The fused kernel needs one block per tile (no stream-k split), so only fuse if the tiling is efficient.
+    return ggml_cuda_mmq_fusion_is_efficient(x_up->type, x_up->ne[1], y->ne[1], y->ne[2]*y->ne[3], device);
+}
+
 static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
     GGML_TENSOR_BINARY_OP_LOCALS
 
@@ -2881,9 +2936,10 @@ static bool ggml_cuda_should_fuse_rope_set_rows(const ggml_tensor * rope,
         return false;
     }
 
-    // Only norm/neox shaders have the fusion code
+    // Only norm/neox/multi shaders have the fusion code
     const int mode = ((const int32_t *) rope->op_params)[2];
-    if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX) {
+    if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX &&
+        mode != GGML_ROPE_TYPE_MROPE && mode != GGML_ROPE_TYPE_IMROPE) {
         return false;
     }
 
@@ -4377,6 +4433,18 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
                 fused_node_count  = 3;
                 break;
             }
+
+            if (ggml_cuda_should_fuse_mul_mat_q(up, gate, glu)) {
+                ggml_cuda_mm_fusion_args_host fusion_data{};
+                fusion_data.gate      = gate->src[0];
+                fusion_data.glu_op    = ggml_get_glu_op(glu);
+                fusion_data.glu_limit = ggml_get_op_params_f32(glu, 3);
+
+                ggml_cuda_mul_mat_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data);
+                fused_mul_mat_vec = true;
+                fused_node_count  = 3;
+                break;
+            }
         }
     }
 
@@ -4973,6 +5041,19 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph
             }
         }
         for (int i = 0; i < cgraph->n_nodes; ++i) {
+            // Keep the inputs of the fused gate/up/GLU MMQ kernel alive until the GLU so that they cannot alias its output (under op offload the weights are buffer copies).
+            if (i + 2 < cgraph->n_nodes && cgraph->nodes[i + 2]->op == GGML_OP_GLU) {
+                ggml_tensor * glu  = cgraph->nodes[i + 2];
+                ggml_tensor * gate = glu->src[0];
+                ggml_tensor * up   = glu->src[1];
+                const bool ok = (gate == cgraph->nodes[i] && up == cgraph->nodes[i + 1]) || (gate == cgraph->nodes[i + 1] && up == cgraph->nodes[i]);
+                if (ok && ggml_cuda_mul_mat_q_fusion_matches(up, gate, glu, cuda_ctx->device)) {
+                    add_alloc_deps(i, i + 2);
+                    i += 2;
+                    continue;
+                }
+            }
+
             ggml_cuda_moe_weighted_reduction_match match;
             if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) {
                 params->add_alloc_dep(params->user_data, const_cast<ggml_tensor *>(match.experts), match.dst);

ggml/src/ggml-cuda/mmq-load-tiles.cuh

@@ -5,7 +5,9 @@
 #include "mmq.cuh"
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q1_0(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -116,7 +118,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q2_0(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -205,7 +209,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_0(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -268,7 +274,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_1(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -331,7 +339,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_0(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -411,7 +421,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_1(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -488,8 +500,8 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     }
 }
 
-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q8_0(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+template <ggml_type type, int J, bool fallback, bool has_fusion> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q8_0(
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -519,7 +531,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
             i = min(i, i_max);
         }
 
-        const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbx;
+        const block_q8_0 * bxi = ggml_cuda_mmq_x_row<block_q8_0, has_fusion>(x, x_gate, kbx0, i, stride) + kbx;
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
         x_qs[i*sram_stride + 0             + txi] = get_int_b2(bxi[0].qs,                   kqsx);
@@ -542,7 +554,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
             i = min(i, i_max);
         }
 
-        const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbxd;
+        const block_q8_0 * bxi = ggml_cuda_mmq_x_row<block_q8_0, has_fusion>(x, x_gate, kbx0, i, stride) + kbxd;
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
         x_df[i*sram_stride                           + kbxd] = bxi->d;
@@ -555,7 +567,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 // ---------------------------------------------------------------------------------------------
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q2_K(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -616,7 +630,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q3_K(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -728,8 +744,8 @@ static __device__ __forceinline__ int unpack_scales_q45_K(const int * scales, co
            ((scales[ksc/2]              >> (2 * (ksc % 2)))       & 0x30303030);  // upper 2 bits
 }
 
-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_K(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+template <ggml_type type, int J, bool fallback, bool has_fusion> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_K(
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -757,7 +773,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
             i = min(i, i_max);
         }
 
-        const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
+        const block_q4_K * bxi = ggml_cuda_mmq_x_row<block_q4_K, has_fusion>(x, x_gate, kbx0, i, stride);
         const int qs0 = get_int_b4(bxi->qs, txi);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
@@ -786,7 +802,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
                 i = min(i, i_max);
             }
 
-            const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
+            const block_q4_K * bxi = ggml_cuda_mmq_x_row<block_q4_K, has_fusion>(x, x_gate, kbx0, i, stride);
 
             const int * scales = (const int *) bxi->scales;
             const int ksc = threadIdx.x % 2;
@@ -814,7 +830,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
             i = min(i, i_max);
         }
 
-        const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
+        const block_q4_K * bxi = ggml_cuda_mmq_x_row<block_q4_K, has_fusion>(x, x_gate, kbx0, i, stride);
 
         x_dm[i] = bxi->dm;
     }
@@ -827,7 +843,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
             i = min(i, i_max);
         }
 
-        const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride + (threadIdx.x % (MMQ_TILE_NE_K/8)) / (QI4_K/8);
+        const block_q4_K * bxi = ggml_cuda_mmq_x_row<block_q4_K, has_fusion>(x, x_gate, kbx0, i, stride) + (threadIdx.x % (MMQ_TILE_NE_K/8)) / (QI4_K/8);
 
         const int * scales = (const int *) bxi->scales;
 
@@ -839,8 +855,8 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 #endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
 }
 
-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_K(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+template <ggml_type type, int J, bool fallback, bool has_fusion> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_K(
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -868,7 +884,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
             i = min(i, i_max);
         }
 
-        const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
+        const block_q5_K * bxi = ggml_cuda_mmq_x_row<block_q5_K, has_fusion>(x, x_gate, kbx0, i, stride);
         const int ky = QR5_K*txi;
 
         const int ql = get_int_b4(bxi->qs, txi);
@@ -909,7 +925,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
                 i = min(i, i_max);
             }
 
-            const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
+            const block_q5_K * bxi = ggml_cuda_mmq_x_row<block_q5_K, has_fusion>(x, x_gate, kbx0, i, stride);
 
             const int * scales = (const int *) bxi->scales;
             const int ksc = threadIdx.x % 2;
@@ -937,7 +953,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
             i = min(i, i_max);
         }
 
-        const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
+        const block_q5_K * bxi = ggml_cuda_mmq_x_row<block_q5_K, has_fusion>(x, x_gate, kbx0, i, stride);
 
         x_dm[i] = bxi->dm;
     }
@@ -951,7 +967,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
             i = min(i, i_max);
         }
 
-        const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
+        const block_q5_K * bxi = ggml_cuda_mmq_x_row<block_q5_K, has_fusion>(x, x_gate, kbx0, i, stride);
 
         const int * scales = (const int *) bxi->scales;
 
@@ -963,8 +979,8 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 #endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
 }
 
-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q6_K(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+template <ggml_type type, int J, bool fallback, bool has_fusion> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q6_K(
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -993,7 +1009,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
             i = min(i, i_max);
         }
 
-        const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride;
+        const block_q6_K * bxi = ggml_cuda_mmq_x_row<block_q6_K, has_fusion>(x, x_gate, kbx0, i, stride);
 
         const int ql = get_int_b2(bxi->ql, txi);
         const int ql0 = (ql >> 0) & 0x0F0F0F0F;
@@ -1023,7 +1039,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
             i = min(i, i_max);
         }
 
-        const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride;
+        const block_q6_K * bxi = ggml_cuda_mmq_x_row<block_q6_K, has_fusion>(x, x_gate, kbx0, i, stride);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
         x_df[i*sram_stride]                     = bxi->d;
@@ -1041,7 +1057,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
             i = min(i, i_max);
         }
 
-        const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride + (threadIdx.x % (MMQ_TILE_NE_K/8)) / 4;
+        const block_q6_K * bxi = ggml_cuda_mmq_x_row<block_q6_K, has_fusion>(x, x_gate, kbx0, i, stride) + (threadIdx.x % (MMQ_TILE_NE_K/8)) / 4;
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
         x_sc[i*sram_stride + threadIdx.x%4] = get_int_b2(bxi->scales, threadIdx.x % (MMQ_TILE_NE_K/8));
@@ -1054,7 +1070,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 // ---------------------------------------------------------------------------------------------
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq1_s(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1116,7 +1134,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_xxs(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1180,7 +1200,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_xs(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1245,7 +1267,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_s(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1313,7 +1337,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq3_xxs(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1377,7 +1403,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq3_s(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1445,8 +1473,8 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     }
 }
 
-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq4_xs(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+template <ggml_type type, int J, bool fallback, bool has_fusion> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq4_xs(
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1473,7 +1501,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
             i = min(i, i_max);
         }
 
-        const block_iq4_xs * bxi = (const block_iq4_xs *) x + kbx0 + i*stride;
+        const block_iq4_xs * bxi = ggml_cuda_mmq_x_row<block_iq4_xs, has_fusion>(x, x_gate, kbx0, i, stride);
 
         const int aux_q4 = get_int_b4(bxi->qs, kqsx);
         const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl);
@@ -1497,7 +1525,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
             i = min(i, i_max);
         }
 
-        const block_iq4_xs * bxi = (const block_iq4_xs *) x + kbx0 + i*stride;
+        const block_iq4_xs * bxi = ggml_cuda_mmq_x_row<block_iq4_xs, has_fusion>(x, x_gate, kbx0, i, stride);
 
         const float d = __half2float(bxi->d);
 
@@ -1513,7 +1541,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq4_nl(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1582,7 +1612,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 // ---------------------------------------------------------------------------------------------
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_mxfp4(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1649,7 +1681,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_mxfp4_fp4(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q4) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q4);
@@ -1691,7 +1725,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }
 
 template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kb0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kb0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
@@ -1747,7 +1783,9 @@ template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __de
 }
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4_nvfp4(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+    GGML_UNUSED(x_gate);
+
     constexpr int warp_size       = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps          = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q4) / warp_size;
     constexpr int I               = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q4);

ggml/src/ggml-cuda/mmq.cu

@@ -131,12 +131,51 @@ static ggml_prec ggml_cuda_mmq_get_prec_src1(const ggml_tensor * src0, const ggm
     return GGML_PREC_Q4;
 }
 
+// Tile width in ne11 (dense) / ne12 (MoE) direction with the fewest tiles that fits into shared memory, 0 if there is none.
+static int ggml_cuda_mmq_get_J_best(
+        const ggml_type type, const int64_t ncols, const bool fallback, const int cc, const size_t smpbo, const ggml_prec prec_src1,
+        int * nthreads_best = nullptr) {
+    int J_best        = 0;
+    int ntiles_J_best = INT_MAX;
+
+    for (int J = 8; J <= 128 && ntiles_J_best > 1; J += 8) {
+        const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1);
+        if (config.type == GGML_TYPE_COUNT) {
+            continue;
+        }
+
+        if (mmq_get_nbytes_shared(config, cc) > smpbo) {
+            continue;
+        }
+
+        const int ntiles_x = (ncols + config.J - 1) / config.J;
+
+        if (ntiles_x < ntiles_J_best) {
+            J_best = J;
+            if (nthreads_best) {
+                *nthreads_best = config.nthreads;
+            }
+            ntiles_J_best = ntiles_x;
+        }
+    }
+    return J_best;
+}
+
 void ggml_cuda_mul_mat_q(
-        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) {
+        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst,
+        const ggml_cuda_mm_fusion_args_host * fusion) {
     GGML_ASSERT(        src1->type == GGML_TYPE_F32);
     GGML_ASSERT(        dst->type  == GGML_TYPE_F32);
     GGML_ASSERT(!ids || ids->type  == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID.
 
+    // Fused gate/up/GLU: src0 is the up weight, dst the GLU output.
+    const ggml_tensor * gate = fusion ? fusion->gate : nullptr;
+    if (fusion) {
+        // TODO MoE support
+        GGML_ASSERT(gate && !ids && !fusion->x_bias && !fusion->gate_bias && !fusion->x_scale && !fusion->gate_scale);
+        GGML_ASSERT(gate->type == src0->type && ggml_are_same_shape(gate, src0) && ggml_are_same_stride(gate, src0));
+    }
+
     GGML_TENSOR_BINARY_OP_LOCALS;
 
     cudaStream_t stream = ctx.stream();
@@ -158,14 +197,16 @@ void ggml_cuda_mul_mat_q(
     const float * src1_d = (const float *) src1->data;
     float       *  dst_d = (float       *)  dst->data;
 
-    // If src0 is a temporary compute buffer, clear any potential padding.
-    if (ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
-        const size_t size_data  = ggml_nbytes(src0);
-        const size_t size_alloc = ggml_backend_buffer_get_alloc_size(src0->buffer, src0);
-        if (size_alloc > size_data) {
-            GGML_ASSERT(ggml_is_contiguously_allocated(src0));
-            GGML_ASSERT(!src0->view_src);
-            CUDA_CHECK(cudaMemsetAsync((char *) src0->data + size_data, 0, size_alloc - size_data, stream));
+    // If src0 (or the fused gate weight) is a temporary compute buffer, clear any potential padding.
+    for (const ggml_tensor * t : { src0, gate }) {
+        if (t && ggml_backend_buffer_get_usage(t->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
+            const size_t size_data  = ggml_nbytes(t);
+            const size_t size_alloc = ggml_backend_buffer_get_alloc_size(t->buffer, t);
+            if (size_alloc > size_data) {
+                GGML_ASSERT(ggml_is_contiguously_allocated(t));
+                GGML_ASSERT(!t->view_src);
+                CUDA_CHECK(cudaMemsetAsync((char *) t->data + size_data, 0, size_alloc - size_data, stream));
+            }
         }
     }
 
@@ -180,7 +221,9 @@ void ggml_cuda_mul_mat_q(
 
     const bool fallback = ggml_cuda_mmq_needs_fallback(ne01);
 
-    const ggml_prec prec_src1 = ggml_cuda_mmq_get_prec_src1(src0, dst, cc);
+    // With fusion dst is the GLU output, not a MUL_MAT; the fused types always quantize src1 to Q8_1.
+    // TODO support for Q4 precision if both tensors allow it
+    const ggml_prec prec_src1 = gate ? GGML_PREC_Q8 : ggml_cuda_mmq_get_prec_src1(src0, dst, cc);
 
     const bool use_native_fp4 = prec_src1 == GGML_PREC_Q4;
     const size_t y_block_size       = use_native_fp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq);
@@ -201,26 +244,7 @@ void ggml_cuda_mul_mat_q(
             }
         }
 
-        int ntiles_J_best = INT_MAX;
-
-        for (int J = 8; J <= 128 && ntiles_J_best > 1; J += 8) {
-            const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(src0->type, J, fallback, cc, prec_src1);
-            if (config.type == GGML_TYPE_COUNT) {
-                continue;
-            }
-
-            if (mmq_get_nbytes_shared(config, cc) > smpbo) {
-                continue;
-            }
-
-            const int ntiles_x = (ncols_opt + config.J - 1) / config.J;
-
-            if (ntiles_x < ntiles_J_best) {
-                J_best = J;
-                nthreads_best = config.nthreads;
-                ntiles_J_best = ntiles_x;
-            }
-        }
+        J_best = ggml_cuda_mmq_get_J_best(src0->type, ncols_opt, fallback, cc, smpbo, prec_src1, &nthreads_best);
     }
     GGML_ASSERT(J_best > 0);
 
@@ -262,13 +286,20 @@ void ggml_cuda_mul_mat_q(
                                 ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
         const int64_t s13 = ne12*s12;
 
+        ggml_cuda_mm_fusion_args_device fusion_device = {};
+        if (gate) {
+            fusion_device.gate      = gate->data;
+            fusion_device.glu_op    = fusion->glu_op;
+            fusion_device.glu_limit = fusion->glu_limit;
+        }
+
         const mmq_args args = {
             src0_d, src0->type, (const int *) src1_q8_1.ptr, nullptr, nullptr, dst_d,
             src0->type == GGML_TYPE_NVFP4 && use_native_fp4 ? src1_scale.ptr : nullptr,
             ne00, ne01, ne1, s01, ne11, s1,
             ne02, ne12, s02, s12, s2,
             ne03, ne13, s03, s13, s3,
-            ne1, J_best};
+            ne1, J_best, fusion_device};
         ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1);
         return;
     }
@@ -485,3 +516,22 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t
 
     return (!GGML_CUDA_CC_IS_CDNA(cc)) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE;
 }
+
+bool ggml_cuda_mmq_fusion_is_efficient(const ggml_type type, const int64_t nrows_x, const int64_t ncols, const int64_t nchannels, const int device) {
+    const int    cc    = ggml_cuda_info().devices[device].cc;
+    const int    nsm   = ggml_cuda_info().devices[device].nsm;
+    const size_t smpbo = ggml_cuda_info().devices[device].smpbo;
+
+    // Same tile width as ggml_cuda_mul_mat_q; the fused kernel has no fallback, quantizes src1 to Q8_1 and writes I/2 output rows per tile.
+    const bool fallback = false;
+    const int  J_best   = ggml_cuda_mmq_get_J_best(type, ncols, fallback, cc, smpbo, GGML_PREC_Q8);
+    if (J_best == 0) {
+        return false;
+    }
+    const int64_t nrows_tile = ggml_cuda_mmq_get_config(type, J_best, fallback, cc, GGML_PREC_Q8).I / 2;
+    const int64_t ntiles     = (ncols + J_best - 1)/J_best * ((nrows_x + nrows_tile - 1)/nrows_tile) * nchannels;
+    // Same condition as launch_mul_mat_q uses to run the stream-k kernel with one block per tile.
+    const int64_t tiles_nwaves             = (ntiles + nsm - 1) / nsm;
+    const int64_t tiles_efficiency_percent = 100 * ntiles / (nsm*tiles_nwaves);
+    return tiles_efficiency_percent >= 90;
+}

ggml/src/ggml-cuda/mmq.cuh

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "common.cuh"
+#include "unary.cuh"
 
 #include <climits>
 #include <cstdint>
@@ -10,10 +11,12 @@
 #define MMQ_ITER_K_FP4         512
 #define MMQ_NWARPS               8
 
-typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, int * x_tile, const int kbx0, const int i_max, const int stride);
+typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, const char * __restrict__ x_gate, int * x_tile,
+    const int kbx0, const int i_max, const int stride);
 typedef void (*ggml_cuda_mmq_vec_dot_t)(const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00);
 typedef void (*ggml_cuda_mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ get_rows_to_sorted,
-    float * __restrict__ dst, const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max);
+    float * __restrict__ dst, const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max,
+    const ggml_glu_op glu_op, const float glu_limit);
 
 enum mmq_q8_1_ds_layout {
     MMQ_Q8_1_DS_LAYOUT_D4,
@@ -426,12 +429,37 @@ static __host__ int ggml_cuda_mmq_get_nbytes_shared_x(const ggml_cuda_mmq_config
 
 // ------------------------------------------------------------
 
+// Weight types with a fused gate/up/GLU path (see ggml_cuda_mmq_get_util_funcs).
+static constexpr __host__ __device__ bool ggml_cuda_mmq_fusion_supported(const ggml_type type) {
+    return type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K || type == GGML_TYPE_Q6_K || type == GGML_TYPE_Q8_0 || type == GGML_TYPE_IQ4_XS;
+}
+
+// Fused gate/up/GLU: the weight tile alternates chunks of this many up and gate rows, so rows r and r + 8 of an MMA tile (same thread) pair up.
+static constexpr int GGML_CUDA_MMQ_FUSION_ROWS = 8;
+
+// Source of weight tile row i; with has_fusion the up or gate row that the interleaved tile holds there.
+template <typename block_t, bool has_fusion>
+static __device__ __forceinline__ const block_t * ggml_cuda_mmq_x_row(
+        const char * __restrict__ x, const char * __restrict__ x_gate, const int kbx0, const int i, const int stride) {
+    if constexpr (has_fusion) {
+        constexpr int R = GGML_CUDA_MMQ_FUSION_ROWS;
+        const char * x_src = (i / R) % 2 == 0 ? x : x_gate;
+        return (const block_t *) x_src + kbx0 + ((i / (2*R))*R + i % R)*stride;
+    } else {
+        GGML_UNUSED(x_gate);
+        return (const block_t *) x + kbx0 + i*stride;
+    }
+}
+
 #include "mmq-load-tiles.cuh"
 #include "mmq-vec-dot.cuh"
 
 template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __device__ __forceinline__ void ggml_cuda_mmq_write_back_dp4a(
         const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, float * __restrict__ dst,
-        const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {
+        const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max,
+        const ggml_glu_op glu_op, const float glu_limit) {
+    GGML_UNUSED_VARS(glu_op, glu_limit);
+
     constexpr int warp_size = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
     constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
@@ -469,16 +497,19 @@ template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __de
     }
 }
 
-template<ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+// has_fusion: accumulators l and l + ne/2 hold the up and gate values of one output element, the GLU of the two is written.
+template<ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
 static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
             const float * __restrict__ sum, const int * __restrict__ ids_dst, float * __restrict__ dst,
-            const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {
+            const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max,
+            const ggml_glu_op glu_op, const float glu_limit) {
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
     typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
 #else
     typedef tile<16,  8, int> tile_C;
 #endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    static_assert(!has_fusion || tile_C::I == 2*GGML_CUDA_MMQ_FUSION_ROWS, "each C tile must hold one chunk of up rows and one of gate rows");
 
     constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, prec_src1);
     constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
@@ -492,20 +523,23 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
 #pragma unroll
         for (int n = 0; n < ntx; ++n) {
 #pragma unroll
-            for (int l = 0; l < tile_C::ne; ++l) {
+            for (int l = 0; l < (has_fusion ? tile_C::ne/2 : tile_C::ne); ++l) {
                 const int j = j0 + (threadIdx.y % ntx) * tile_C::J + tile_C::get_j(l);
 
                 if (j > j_max) {
                     continue;
                 }
 
-                const int i = i0 + n*tile_C::I + tile_C::get_i(l);
+                const int i = (i0 + n*tile_C::I)/(has_fusion ? 2 : 1) + tile_C::get_i(l);
 
                 if (fallback && i > i_max) {
                     continue;
                 }
 
-                if constexpr (type == GGML_TYPE_NVFP4) {
+                if constexpr (has_fusion) {
+                    const int sum_index = (j0/tile_C::J + n)*tile_C::ne + l;
+                    dst[ids_dst[j]*stride + i] = ggml_cuda_op_glu_single(glu_op, sum[sum_index + tile_C::ne/2], sum[sum_index], glu_limit);
+                } else if constexpr (type == GGML_TYPE_NVFP4) {
                     if (y_scale_used) {
                         dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l];
                     } else {
@@ -518,6 +552,8 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
             }
         }
     }
+
+    GGML_UNUSED_VARS(y_scale_used, glu_op, glu_limit);
 }
 
 // -------------------------------------------------------------------------------------------------------------------------------------
@@ -534,8 +570,9 @@ struct ggml_cuda_mmq_util_funcs {
         vdr(vdr), load_tiles(load_tiles), vec_dot(vec_dot), write_back(write_back) {}
 };
 
-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
 static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_funcs() {
+    static_assert(!has_fusion || ggml_cuda_mmq_fusion_supported(type), "no fused loader for this type");
     if (!ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).use_mma_data_layout()) {
         switch (type) {
             case GGML_TYPE_Q1_0:
@@ -577,7 +614,7 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
             case GGML_TYPE_Q8_0:
                 return ggml_cuda_mmq_util_funcs(
                     VDR_Q8_0_Q8_1_MMQ,
-                    ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback>,
+                    ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback, false>,
                     ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
                     ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
 // ---------------------------------------------------------------------------------------------
@@ -596,19 +633,19 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
             case GGML_TYPE_Q4_K:
                 return ggml_cuda_mmq_util_funcs(
                     VDR_Q4_K_Q8_1_MMQ,
-                    ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback>,
+                    ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback, false>,
                     ggml_cuda_mmq_vec_dot_q4_K_q8_1_dp4a<type, J, fallback>,
                     ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
             case GGML_TYPE_Q5_K:
                 return ggml_cuda_mmq_util_funcs(
                     VDR_Q5_K_Q8_1_MMQ,
-                    ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback>,
+                    ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback, false>,
                     ggml_cuda_mmq_vec_dot_q5_K_q8_1_dp4a<type, J, fallback>,
                     ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
             case GGML_TYPE_Q6_K:
                 return ggml_cuda_mmq_util_funcs(
                     VDR_Q6_K_Q8_1_MMQ,
-                    ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback>,
+                    ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback, false>,
                     ggml_cuda_mmq_vec_dot_q6_K_q8_1_dp4a<type, J, fallback>,
                     ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
 // ---------------------------------------------------------------------------------------------
@@ -651,7 +688,7 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
             case GGML_TYPE_IQ4_XS:
                 return ggml_cuda_mmq_util_funcs(
                     VDR_IQ4_XS_Q8_1_MMQ,
-                    ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback>,
+                    ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback, false>,
                     ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
                     ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
             case GGML_TYPE_IQ4_NL:
@@ -688,7 +725,7 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
                     -1,
                     ggml_cuda_mmq_load_tiles_mxfp4_fp4<type, J, fallback>,
                     ggml_cuda_mmq_vec_dot_fp4_fp4_mma<type, J, fallback>,
-                    ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                    ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
             }
             break;
         case GGML_TYPE_NVFP4:
@@ -697,7 +734,7 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
                     -1,
                     ggml_cuda_mmq_load_tiles_nvfp4_nvfp4<type, J, fallback>,
                     ggml_cuda_mmq_vec_dot_fp4_fp4_mma<type, J, fallback>,
-                    ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                    ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
             }
             break;
         default:
@@ -713,136 +750,136 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
                 -1,
                 ggml_cuda_mmq_load_tiles_q1_0<type, J, fallback>,
                 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
         case GGML_TYPE_Q2_0:
             return ggml_cuda_mmq_util_funcs(
                 -1,
                 ggml_cuda_mmq_load_tiles_q2_0<type, J, fallback>,
                 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
         case GGML_TYPE_Q4_0:
             return ggml_cuda_mmq_util_funcs(
                 -1,
                 ggml_cuda_mmq_load_tiles_q4_0<type, J, fallback>,
                 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_DS4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
         case GGML_TYPE_Q4_1:
             return ggml_cuda_mmq_util_funcs(
                 -1,
                 ggml_cuda_mmq_load_tiles_q4_1<type, J, fallback>,
                 ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
         case GGML_TYPE_Q5_0:
             return ggml_cuda_mmq_util_funcs(
                 -1,
                 ggml_cuda_mmq_load_tiles_q5_0<type, J, fallback>,
                 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
         case GGML_TYPE_Q5_1:
             return ggml_cuda_mmq_util_funcs(
                 -1,
                 ggml_cuda_mmq_load_tiles_q5_1<type, J, fallback>,
                 ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
         case GGML_TYPE_Q8_0:
             return ggml_cuda_mmq_util_funcs(
                 -1,
-                ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback>,
+                ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback, has_fusion>,
                 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, has_fusion>);
 // ---------------------------------------------------------------------------------------------
         case GGML_TYPE_Q2_K:
             return ggml_cuda_mmq_util_funcs(
                 -1,
                 ggml_cuda_mmq_load_tiles_q2_K<type, J, fallback>,
                 ggml_cuda_mmq_vec_dot_q2_K_q8_1_mma<type, J, fallback>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
         case GGML_TYPE_Q3_K:
             return ggml_cuda_mmq_util_funcs(
                 -1,
                 ggml_cuda_mmq_load_tiles_q3_K<type, J, fallback>,
                 ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback, prec_src1>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
         case GGML_TYPE_Q4_K:
             return ggml_cuda_mmq_util_funcs(
                 -1,
-                ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback>,
+                ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback, has_fusion>,
                 ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, has_fusion>);
         case GGML_TYPE_Q5_K:
             return ggml_cuda_mmq_util_funcs(
                 -1,
-                ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback>,
+                ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback, has_fusion>,
                 ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, has_fusion>);
         case GGML_TYPE_Q6_K:
             return ggml_cuda_mmq_util_funcs(
                 -1,
-                ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback>,
+                ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback, has_fusion>,
                 ggml_cuda_mmq_vec_dot_q6_K_q8_1_mma<type, J, fallback>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, has_fusion>);
 // ---------------------------------------------------------------------------------------------
         case GGML_TYPE_IQ1_S:
             return ggml_cuda_mmq_util_funcs(
                 -1,
                 ggml_cuda_mmq_load_tiles_iq1_s<type, J, fallback>,
                 ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
         case GGML_TYPE_IQ2_XXS:
             return ggml_cuda_mmq_util_funcs(
                 -1,
                 ggml_cuda_mmq_load_tiles_iq2_xxs<type, J, fallback>,
                 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
         case GGML_TYPE_IQ2_XS:
             return ggml_cuda_mmq_util_funcs(
                 -1,
                 ggml_cuda_mmq_load_tiles_iq2_xs<type, J, fallback>,
                 ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback, prec_src1>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
         case GGML_TYPE_IQ2_S:
             return ggml_cuda_mmq_util_funcs(
                 -1,
                 ggml_cuda_mmq_load_tiles_iq2_s<type, J, fallback>,
                 ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback, prec_src1>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
         case GGML_TYPE_IQ3_XXS:
             return ggml_cuda_mmq_util_funcs(
                 -1,
                 ggml_cuda_mmq_load_tiles_iq3_xxs<type, J, fallback>,
                 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
         case GGML_TYPE_IQ3_S:
             return ggml_cuda_mmq_util_funcs(
                 -1,
                 ggml_cuda_mmq_load_tiles_iq3_s<type, J, fallback>,
                 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
         case GGML_TYPE_IQ4_XS:
             return ggml_cuda_mmq_util_funcs(
                 -1,
-                ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback>,
+                ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback, has_fusion>,
                 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, has_fusion>);
         case GGML_TYPE_IQ4_NL:
             return ggml_cuda_mmq_util_funcs(
                 -1,
                 ggml_cuda_mmq_load_tiles_iq4_nl<type, J, fallback>,
                 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
 // ---------------------------------------------------------------------------------------------
         case GGML_TYPE_MXFP4:
             return ggml_cuda_mmq_util_funcs(
                 -1,
                 ggml_cuda_mmq_load_tiles_mxfp4<type, J, fallback>,
                 ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
         case GGML_TYPE_NVFP4:
             return ggml_cuda_mmq_util_funcs(
                 -1,
                 ggml_cuda_mmq_load_tiles_nvfp4<type, J, fallback, prec_src1>,
                 ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback, prec_src1>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
         default:
             return ggml_cuda_mmq_util_funcs(1, nullptr, nullptr, nullptr);
     }
@@ -850,41 +887,44 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
 
 template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
 static constexpr __device__ int ggml_cuda_mmq_get_vdr() {
-    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().vdr;
+    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1, false>().vdr;
 }
 
-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
 static constexpr __device__ ggml_cuda_mmq_load_tiles_t ggml_cuda_mmq_get_load_tiles() {
-    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().load_tiles;
+    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1, has_fusion>().load_tiles;
 }
 
 template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
 static constexpr __device__ ggml_cuda_mmq_vec_dot_t ggml_cuda_mmq_get_vec_dot() {
-    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().vec_dot;
+    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1, false>().vec_dot;
 }
 
-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
 static constexpr __device__ ggml_cuda_mmq_write_back_t ggml_cuda_mmq_get_write_back() {
-    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().write_back;
+    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1, has_fusion>().write_back;
 }
 
 // ---------------------------------------------------------------------------------------------
 
-template <ggml_type type, int J, bool fallback, bool fixup, ggml_prec prec_src1>
+// has_fusion: x and x_gate are the up and gate weights.
+template <ggml_type type, int J, bool fallback, bool fixup, ggml_prec prec_src1, bool has_fusion>
 static __device__ __forceinline__ void mul_mat_q_process_tile(
-        const char * __restrict__ x, const int offset_x, const int * __restrict__ y,
+        const char * __restrict__ x, const char * __restrict__ x_gate, const int offset_x, const int * __restrict__ y,
         const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup,
-        const float * __restrict__ y_scale,
+        const float * __restrict__ y_scale, const ggml_glu_op glu_op, const float glu_limit,
         const int stride_row_x, const int ncols_y, const int stride_col_dst,
         const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) {
+    static_assert(!(fixup && has_fusion), "the fused gate/up path does not split K");
+    static_assert(!(fallback && has_fusion), "the fused gate/up path needs nrows_x % 128 == 0");
 
     constexpr int              warp_size  = ggml_cuda_get_physical_warp_size();
     constexpr int              nwarps     = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
     constexpr int              qk         = ggml_cuda_type_traits<type>::qk;
     constexpr int              I          = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
-    constexpr ggml_cuda_mmq_load_tiles_t load_tiles = ggml_cuda_mmq_get_load_tiles<type, J, fallback, prec_src1>();
+    constexpr ggml_cuda_mmq_load_tiles_t load_tiles = ggml_cuda_mmq_get_load_tiles<type, J, fallback, prec_src1, has_fusion>();
     constexpr ggml_cuda_mmq_vec_dot_t    vec_dot    = ggml_cuda_mmq_get_vec_dot<type, J, fallback, prec_src1>();
-    constexpr ggml_cuda_mmq_write_back_t write_back = ggml_cuda_mmq_get_write_back<type, J, fallback, prec_src1>();
+    constexpr ggml_cuda_mmq_write_back_t write_back = ggml_cuda_mmq_get_write_back<type, J, fallback, prec_src1, has_fusion>();
 
     extern __shared__ int data_mul_mat_q[];
     int * tile_y = data_mul_mat_q + J;
@@ -907,7 +947,7 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
     constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int);
 
     for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) {
-        load_tiles(x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x);
+        load_tiles(x, x_gate, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x);
         {
             const int * by0 = y + ncols_y * (kb0 * qk / ne_block) * sz;
 #pragma unroll
@@ -942,28 +982,34 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
     }
 
     if (fixup) {
-        write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J);
+        write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J, glu_op, glu_limit);
     } else {
-        write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j);
+        write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j, glu_op, glu_limit);
     }
 }
 
 
 // The mul_mat_q kernel implements "stream-k" work partitioning as described in https://arxiv.org/abs/2301.03598
 
-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
 __launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1), ggml_cuda_mmq_get_occupancy(type, J, fallback, prec_src1))
 static __global__ void mul_mat_q(
-        const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst,
+        const char * __restrict__ x, const char * __restrict__ x_gate, const int * __restrict__ y, const int32_t * __restrict__ ids_dst,
         const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, float * __restrict__ tmp_fixup,
-        const float * __restrict__ y_scale,
+        const float * __restrict__ y_scale, const ggml_glu_op glu_op, const float glu_limit,
         const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, const int stride_row_x, const int ncols_y, const int stride_col_dst,
         const uint3 channel_ratio, const uint3 nchannels_y, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
         const uint3 sample_ratio, const uint3 nsamples_y, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst,
         const uint3 ntx) {
 
+#ifdef TURING_MMA_AVAILABLE
+    constexpr bool fusion_available = true;
+#else
+    constexpr bool fusion_available = false;
+#endif // TURING_MMA_AVAILABLE
+
     // Skip unused template specializations for faster compilation:
-    if (ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).type == GGML_TYPE_COUNT) {
+    if (ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).type == GGML_TYPE_COUNT || (has_fusion && !fusion_available)) {
         NO_DEVICE_CODE;
         return;
     }
@@ -991,12 +1037,16 @@ static __global__ void mul_mat_q(
     }
     __syncthreads();
 
-    if constexpr (!ggml_cuda_mmq_get_stream_k(type, J, fallback, prec_src1)) {
+    // The fused kernel needs the whole K range of a tile in one block, so it always uses tiling.
+    if constexpr (has_fusion || !ggml_cuda_mmq_get_stream_k(type, J, fallback, prec_src1)) {
+        // The fused kernel writes the GLU of I/2 up rows and I/2 gate rows per tile.
+        constexpr int I_dst = has_fusion ? I/2 : I;
+
         const uint2 tmp2 = fast_div_modulo(blockIdx.z, nchannels_y);
         const int wt = tmp2.x;
         const int zt = tmp2.y;
-        const int jt = blockIdx.y;
-        const int it = blockIdx.x;
+        const int jt = blockIdx.x;
+        const int it = blockIdx.y;
 
         // Defaults for regular matrix multiplication:
         int col_low    = 0;
@@ -1041,22 +1091,22 @@ static __global__ void mul_mat_q(
         }
 
         offset_y   += (col_low + jt*J)*(sizeof(block_q8_1_mmq)/sizeof(int));
-        offset_dst += it*I;
+        offset_dst += it*I_dst;
         const float * y_scale_tile = nullptr;
         if constexpr (type == GGML_TYPE_NVFP4) {
             offset_y_scale += col_low + jt*J;
             y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
         }
 
-        const int tile_x_max_i = nrows_x  - it*I - 1;
+        const int tile_x_max_i = nrows_x  - it*I_dst - 1;
         const int tile_y_max_j = col_diff - jt*J - 1;
 
-        const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
+        const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I_dst*stride_row_x;
 
         constexpr bool fixup = false;
-        mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1>
-            (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
-             stride_row_x, ncols_y, stride_col_dst,
+        mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1, has_fusion>
+            (x, x_gate, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+             glu_op, glu_limit, stride_row_x, ncols_y, stride_col_dst,
              tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z);
         return;
     }
@@ -1148,9 +1198,9 @@ static __global__ void mul_mat_q(
         const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
 
         constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer.
-        mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1>
-            (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
-             stride_row_x, ncols_y, stride_col_dst,
+        mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1, false>
+            (x, x_gate, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+             glu_op, glu_limit, stride_row_x, ncols_y, stride_col_dst,
              tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
 
         kbc += blocks_per_ne00.z;
@@ -1232,9 +1282,9 @@ static __global__ void mul_mat_q(
     const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
 
     constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks.
-    mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1>
-        (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
-         stride_row_x, ncols_y, stride_col_dst,
+    mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1, false>
+        (x, x_gate, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+         glu_op, glu_limit, stride_row_x, ncols_y, stride_col_dst,
          tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
 }
 
@@ -1384,6 +1434,7 @@ struct mmq_args {
     int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst;
     int64_t ncols_max;
     int J_best; // Tile width in ne11(dense)/ne12(MoE) direction to use for optimal performance.
+    ggml_cuda_mm_fusion_args_device fusion = {}; // gate != nullptr selects the fused gate/up/GLU path
 };
 
 static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const int cc) {
@@ -1393,7 +1444,7 @@ static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const i
     return nbs_ids + nbs_x + GGML_PAD(nbs_y, config.nthreads*sizeof(int));
 }
 
-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
 static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
     const int id = ggml_cuda_get_device();
     const int cc = ggml_cuda_info().devices[id].cc;
@@ -1407,13 +1458,17 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
 
     const dim3 block_dims(warp_size, nwarps, 1);
 
-    CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, false, prec_src1>), nbytes_shared);
-    CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J,  true, prec_src1>), nbytes_shared);
+    CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, /*fallback =*/ false, prec_src1, has_fusion>), nbytes_shared);
+    if constexpr (!has_fusion) {
+        CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, /*fallback =*/ true, prec_src1, false>), nbytes_shared);
+    }
 
-    const int nty  = (args.nrows_x   + config.I - 1) / config.I;
+    const int I_dst = has_fusion ? config.I/2 : config.I; // Rows of dst per tile, see mul_mat_q.
+    const int nty  = (args.nrows_x   + I_dst - 1) / I_dst;
     const int ntx  = (args.ncols_max + config.J - 1) / config.J;
     const int ntzw = args.nchannels_y * args.nsamples_y;
-    const dim3 block_nums_xy_tiling(nty, ntx, ntzw);
+    // Column tiles innermost so that the concurrently running blocks share the same tile of x and read it from L2.
+    const dim3 block_nums_xy_tiling(ntx, nty, ntzw);
 
     GGML_ASSERT(args.nchannels_y % args.nchannels_x == 0);
     GGML_ASSERT(args.nsamples_y  % args.nsamples_x  == 0);
@@ -1427,9 +1482,10 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
     const uint3 channel_ratio_fd   = init_fastdiv_values(channel_ratio);
     const uint3 sample_ratio_fd    = init_fastdiv_values(sample_ratio);
 
-    if (!config.stream_k) {
-        mul_mat_q<type, J, fallback, prec_src1><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
-            (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale,
+    if (has_fusion || !config.stream_k) {
+        GGML_ASSERT(nty <= 65535); // gridDim.y limit
+        mul_mat_q<type, J, fallback, prec_src1, has_fusion><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
+            (args.x, (const char *) args.fusion.gate, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale, args.fusion.glu_op, args.fusion.glu_limit,
              blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
              channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
              sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
@@ -1459,8 +1515,8 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
     const dim3 block_nums_fixup(block_nums_stream_k.x, config.I/warp_size, 1);
     const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z);
 
-    mul_mat_q<type, J, fallback, prec_src1><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
-        (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale,
+    mul_mat_q<type, J, fallback, prec_src1, has_fusion><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
+        (args.x, (const char *) args.fusion.gate, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale, args.fusion.glu_op, args.fusion.glu_limit,
          blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
          channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
          sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
@@ -1477,56 +1533,56 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
          ntx_fd);
 }
 
-template <ggml_type type, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, bool fallback, ggml_prec prec_src1, bool has_fusion>
 void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
     switch (args.J_best) {
         case   8:
-            launch_mul_mat_q<type,   8, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,   8, fallback, prec_src1, has_fusion>(ctx, args, stream);
             break;
         case  16:
-            launch_mul_mat_q<type,  16, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  16, fallback, prec_src1, has_fusion>(ctx, args, stream);
             break;
         case  24:
-            launch_mul_mat_q<type,  24, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  24, fallback, prec_src1, has_fusion>(ctx, args, stream);
             break;
         case  32:
-            launch_mul_mat_q<type,  32, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  32, fallback, prec_src1, has_fusion>(ctx, args, stream);
             break;
         case  40:
-            launch_mul_mat_q<type,  40, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  40, fallback, prec_src1, has_fusion>(ctx, args, stream);
             break;
         case  48:
-            launch_mul_mat_q<type,  48, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  48, fallback, prec_src1, has_fusion>(ctx, args, stream);
             break;
         case  56:
-            launch_mul_mat_q<type,  56, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  56, fallback, prec_src1, has_fusion>(ctx, args, stream);
             break;
         case  64:
-            launch_mul_mat_q<type,  64, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  64, fallback, prec_src1, has_fusion>(ctx, args, stream);
             break;
         case  72:
-            launch_mul_mat_q<type,  72, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  72, fallback, prec_src1, has_fusion>(ctx, args, stream);
             break;
         case  80:
-            launch_mul_mat_q<type,  80, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  80, fallback, prec_src1, has_fusion>(ctx, args, stream);
             break;
         case  88:
-            launch_mul_mat_q<type,  88, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  88, fallback, prec_src1, has_fusion>(ctx, args, stream);
             break;
         case  96:
-            launch_mul_mat_q<type,  96, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  96, fallback, prec_src1, has_fusion>(ctx, args, stream);
             break;
         case 104:
-            launch_mul_mat_q<type, 104, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type, 104, fallback, prec_src1, has_fusion>(ctx, args, stream);
             break;
         case 112:
-            launch_mul_mat_q<type, 112, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type, 112, fallback, prec_src1, has_fusion>(ctx, args, stream);
             break;
         case 120:
-            launch_mul_mat_q<type, 120, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type, 120, fallback, prec_src1, has_fusion>(ctx, args, stream);
             break;
         case 128:
-            launch_mul_mat_q<type, 128, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type, 128, fallback, prec_src1, has_fusion>(ctx, args, stream);
             break;
         default:
             fprintf(stderr, "J_best=%d\n", args.J_best);
@@ -1534,14 +1590,26 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args,
     }
 }
 
+template <ggml_type type, bool fallback, ggml_prec prec_src1>
+static void mul_mat_q_switch_fusion(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
+    if constexpr (!fallback && ggml_cuda_mmq_fusion_supported(type)) {
+        if (args.fusion.gate) {
+            mul_mat_q_switch_J<type, fallback, prec_src1, true>(ctx, args, stream);
+            return;
+        }
+    }
+    GGML_ASSERT(!args.fusion.gate);
+    mul_mat_q_switch_J<type, fallback, prec_src1, false>(ctx, args, stream);
+}
+
 template <ggml_type type, ggml_prec prec_src1>
 void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
     if (ggml_cuda_mmq_needs_fallback(args.nrows_x)) {
         constexpr bool fallback = true;
-        mul_mat_q_switch_J<type, fallback, prec_src1>(ctx, args, stream);
+        mul_mat_q_switch_fusion<type, fallback, prec_src1>(ctx, args, stream);
     } else {
         constexpr bool fallback = false;
-        mul_mat_q_switch_J<type, fallback, prec_src1>(ctx, args, stream);
+        mul_mat_q_switch_fusion<type, fallback, prec_src1>(ctx, args, stream);
     }
 }
 
@@ -1583,6 +1651,9 @@ extern DECL_MMQ_CASE_W4A4(GGML_TYPE_NVFP4);
 // -------------------------------------------------------------------------------------------------------------------------
 
 void ggml_cuda_mul_mat_q(
-        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst);
+        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst,
+        const ggml_cuda_mm_fusion_args_host * fusion = nullptr);
 
 bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts);
+
+bool ggml_cuda_mmq_fusion_is_efficient(ggml_type type, int64_t nrows_x, int64_t ncols, int64_t nchannels, int device);

ggml/src/ggml-cuda/rope.cu

@@ -196,9 +196,9 @@ static __global__ void rope_neox(const T *            x,
     dst[idst + n_offs/2 + n_dims / 2] = ggml_cuda_cast<D>(x0 * sin_theta + x1 * cos_theta);
 }
 
-template <bool forward, bool has_ff, typename T>
+template <bool forward, bool has_ff, typename T, typename D>
 static __global__ void rope_multi(const T *            x,
-                                  T *                  dst,
+                                  D *                  dst,
                                   const int            ne00,
                                   const int            ne01,
                                   const int            ne02,
@@ -219,6 +219,8 @@ static __global__ void rope_multi(const T *            x,
                                   const float *        freq_factors,
                                   const mrope_sections sections,
                                   const bool           is_imrope,
+                                  const int64_t *      row_indices,
+                                  const int            set_rows_stride,
                                   const bool           inplace) {
     const int i0 = 2 * (blockDim.y * blockIdx.y + threadIdx.y);
 
@@ -236,12 +238,20 @@ static __global__ void rope_multi(const T *            x,
     const int ix   = i0 / 2 + i1 * s01 + i2 * s02 + i3 * s03;
 
     ggml_cuda_pdl_sync();
+
+    // Fusion optimization: ROPE + VIEW + SET_ROWS.
+    // The rope output is viewed as a 1D tensor and offset based on a row index in row_indices.
+    if (set_rows_stride != 0) {
+        idst = i1 * s1 + i0 / 2;
+        idst += row_indices[i2] * set_rows_stride;
+    }
+
     if (i0 < n_offs || i0 >= n_offs + n_dims) {
         if (inplace) {
             return;
         }
-        dst[idst + i0/2 + 0] = x[ix + i0/2 + 0];
-        dst[idst + i0/2 + 1] = x[ix + i0/2 + 1];
+        dst[idst + i0/2 + 0] = ggml_cuda_cast<D>(x[ix + i0/2 + 0]);
+        dst[idst + i0/2 + 1] = ggml_cuda_cast<D>(x[ix + i0/2 + 1]);
 
         return;
     }
@@ -286,8 +296,8 @@ static __global__ void rope_multi(const T *            x,
     const float x0 = x[ix + n_offs/2 + 0];
     const float x1 = x[ix + n_offs/2 + n_dims/2];
 
-    dst[idst + n_offs/2 + 0]        = x0*cos_theta - x1*sin_theta;
-    dst[idst + n_offs/2 + n_dims/2] = x0*sin_theta + x1*cos_theta;
+    dst[idst + n_offs/2 + 0]        = ggml_cuda_cast<D>(x0*cos_theta - x1*sin_theta);
+    dst[idst + n_offs/2 + n_dims/2] = ggml_cuda_cast<D>(x0*sin_theta + x1*cos_theta);
 }
 
 template <bool forward, bool has_ff, typename T>
@@ -443,9 +453,9 @@ static void rope_neox_cuda(const T *            x,
     }
 }
 
-template <bool forward, typename T>
+template <bool forward, typename T, typename D>
 static void rope_multi_cuda(const T *            x,
-                            T *                  dst,
+                            D *                  dst,
                             const int            ne00,
                             const int            ne01,
                             const int            ne02,
@@ -467,6 +477,8 @@ static void rope_multi_cuda(const T *            x,
                             const float *        freq_factors,
                             const mrope_sections sections,
                             const bool           is_imrope,
+                            const int64_t *      row_indices,
+                            const int            set_rows_stride,
                             const bool           inplace,
                             cudaStream_t         stream) {
     GGML_ASSERT(ne00 % 2 == 0);
@@ -478,14 +490,14 @@ static void rope_multi_cuda(const T *            x,
 
     if (freq_factors == nullptr) {
         const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
-        ggml_cuda_kernel_launch(rope_multi<forward, false, T>, launch_params,
+        ggml_cuda_kernel_launch(rope_multi<forward, false, T, D>, launch_params,
             x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor,
-            attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, inplace);
+            attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, row_indices, set_rows_stride, inplace);
     } else {
         const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
-        ggml_cuda_kernel_launch(rope_multi<forward, true, T>, launch_params,
+        ggml_cuda_kernel_launch(rope_multi<forward, true, T, D>, launch_params,
             x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor,
-            attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, inplace);
+            attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, row_indices, set_rows_stride, inplace);
     }
 }
 
@@ -648,14 +660,21 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx,
             GGML_ABORT("fatal error");
         }
     } else if (is_mrope && !is_vision) {
-        if (src0->type == GGML_TYPE_F32) {
-            rope_multi_cuda<forward>((const float *) src0_d, (float *) dst_d, ne00, ne01, ne02, s01, s02, s03, s1,
-                                     s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor,
-                                     corr_dims, freq_factors, sections, is_imrope, inplace, stream);
-        } else if (src0->type == GGML_TYPE_F16) {
-            rope_multi_cuda<forward>((const half *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02, s03, s1,
-                                     s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor,
-                                     corr_dims, freq_factors, sections, is_imrope, inplace, stream);
+        if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) {
+            rope_multi_cuda<forward, float, float>((const float *) src0_d, (float *) dst_d, ne00, ne01, ne02, s01, s02,
+                                                   s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
+                                                   ext_factor, attn_factor, corr_dims, freq_factors, sections, is_imrope,
+                                                   row_indices, set_rows_stride, inplace, stream);
+        } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) {
+            rope_multi_cuda<forward, float, half>((const float *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02,
+                                                  s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
+                                                  ext_factor, attn_factor, corr_dims, freq_factors, sections, is_imrope,
+                                                  row_indices, set_rows_stride, inplace, stream);
+        } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) {
+            rope_multi_cuda<forward, half, half>((const half *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02,
+                                                 s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
+                                                 ext_factor, attn_factor, corr_dims, freq_factors, sections, is_imrope,
+                                                 row_indices, set_rows_stride, inplace, stream);
         } else {
             GGML_ABORT("fatal error");
         }

ggml/src/ggml-cuda/topk-moe.cu

@@ -76,7 +76,7 @@ __device__ void sqrt_softplus_warp_inplace(float (&vals)[experts_per_thread], co
     for (int i = 0; i < experts_per_thread; i++) {
         const int  idx    = lane + i * WARP_SIZE;
         const bool active = !use_limit || (idx < limit);
-        vals[i]           = active ? sqrtf(vals[i] > 20.0f ? vals[i] : logf(1.0f + expf(vals[i]))) : -INFINITY;
+        vals[i]           = active ? sqrtf(vals[i] > 20.0f ? vals[i] : log1pf(expf(vals[i]))) : -INFINITY;
     }
 }
 

ggml/src/ggml-cuda/unary.cu

@@ -91,7 +91,7 @@ static __device__ __forceinline__ float op_expm1(float x) {
 }
 
 static __device__ __forceinline__ float op_softplus(float x) {
-    return (x > 20.0f) ? x : logf(1.0f + expf(x));
+    return (x > 20.0f) ? x : log1pf(expf(x));
 }
 
 static __device__ __forceinline__ float op_elu(float x) {

ggml/src/ggml-cuda/unary.cuh

@@ -121,3 +121,18 @@ __device__ __forceinline__ float ggml_cuda_op_swiglu_clamp_single(float gate, fl
 
     return ggml_cuda_op_silu_single(gate) * up;
 }
+
+__device__ __forceinline__ float ggml_cuda_op_glu_single(ggml_glu_op op, float gate, float up, float limit) {
+    switch (op) {
+        case GGML_GLU_OP_SWIGLU:
+            return up * ggml_cuda_op_silu_single(gate);
+        case GGML_GLU_OP_GEGLU:
+            return up * ggml_cuda_op_gelu_single(gate);
+        case GGML_GLU_OP_SWIGLU_OAI:
+            return ggml_cuda_op_swiglu_oai_single(gate, up);
+        case GGML_GLU_OP_SWIGLU_CLAMP:
+            return ggml_cuda_op_swiglu_clamp_single(gate, up, limit);
+        default:
+            return up * gate;
+    }
+}

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

@@ -133,7 +133,7 @@ size_t opt_mbuf_total  = 0;                        // total buffer space limit (
 
 int    opt_mm_select     = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
 int    opt_fa_select     = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
-int    opt_fa_head_split = 1; // 1 = partition flash_attn by KV heads in multicore (default on), 0 = token-based (original)
+int    opt_fa_split      = 1; // 1 = head-split + team-split when n_kv_heads < N_cores (default), 0 = off
 int    opt_gdn_select    = 2; // 2 = HMX -> HVX, 1 = HVX, 0 = CPU (unsupported)
 int    opt_ar_select     = 2; // 2 = fused ALLREDUCE+ADD (default), 1 = unfused ALLREDUCE, 0 = fallback to CPY+FENCE
 int    opt_ar_scatter    = 1; // 1 = reduce-scatter the fused ALLREDUCE+ADD (default), 0 = full reduction
@@ -2397,7 +2397,7 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg
                 }
             }
 
-            GGML_ASSERT(htp_kparams_solve(sess, node.node, node.kernel_params));
+            GGML_ASSERT(htp_kparams_solve(sess, node));
             if (node.opcode == HTP_OP_CPY && node.kparams<htp_copy_kernel_params>()->total_elems == 0) {
                 continue;
             }
@@ -3393,7 +3393,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
     const char * str_nhmx     = getenv("GGML_HEXAGON_NHMX");
     const char * str_mm_select = getenv("GGML_HEXAGON_MM_SELECT");
     const char * str_fa_select = getenv("GGML_HEXAGON_FA_SELECT");
-    const char * str_fa_head_split = getenv("GGML_HEXAGON_FA_HEAD_SPLIT");
+    const char * str_fa_split  = getenv("GGML_HEXAGON_FA_SPLIT");
     const char * str_gdn_select = getenv("GGML_HEXAGON_GDN_SELECT");
     const char * str_ar_select = getenv("GGML_HEXAGON_AR_SELECT");
     const char * str_ar_scatter = getenv("GGML_HEXAGON_AR_SCATTER");
@@ -3458,7 +3458,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
     opt_nhmx      = str_nhmx     ? atoi(str_nhmx)                         : opt_nhmx;
     opt_mm_select = str_mm_select ? atoi(str_mm_select)                   : opt_mm_select;
     opt_fa_select = str_fa_select ? atoi(str_fa_select)                   : opt_fa_select;
-    opt_fa_head_split = str_fa_head_split ? atoi(str_fa_head_split)       : opt_fa_head_split;
+    opt_fa_split  = str_fa_split  ? atoi(str_fa_split)                   : opt_fa_split;
     opt_gdn_select    = str_gdn_select    ? atoi(str_gdn_select)          : opt_gdn_select;
     opt_ar_select     = str_ar_select     ? atoi(str_ar_select)           : opt_ar_select;
     opt_ar_scatter    = str_ar_scatter    ? atoi(str_ar_scatter)          : opt_ar_scatter;

ggml/src/ggml-hexagon/htp-common.h

@@ -28,7 +28,7 @@ extern size_t opt_mbuf_total;
 
 extern int    opt_mm_select;
 extern int    opt_fa_select;
-extern int    opt_fa_head_split;
+extern int    opt_fa_split;
 extern int    opt_gdn_select;
 extern int    opt_ar_select;
 extern int    opt_ar_scatter;

ggml/src/ggml-hexagon/htp-fusion.cpp

@@ -192,15 +192,20 @@ static bool fuse_mul_mat_add(const ggml_hexagon_session * sess, htp_opnode & pre
     if (!htp_match_consumer(prev, cur, /*commutative=*/ true, &src2)) return false;
     if (src2->type != GGML_TYPE_F32) return false;
 
-    const ggml_tensor * src0 = prev.src0();
-    const ggml_tensor * src1 = prev.src1();
-    const ggml_tensor * dst  = cur.dst();
-    if (!ggml_are_same_shape(prev.dst(), dst)) return false;
+    const ggml_tensor * src0 = prev.node->src[0];
+    const ggml_tensor * src1 = prev.node->src[1];
+    const ggml_tensor * dst  = cur.node;
+    if (!ggml_are_same_shape(prev.node, dst)) return false;
+
+    const bool prev_collapsed = htp_matmul_can_collapse(src0, src1, prev.node);
+    const bool cur_collapsed  = htp_matmul_add_can_collapse(src0, src1, src2, dst);
+    if (prev_collapsed != cur_collapsed) {
+        return false;
+    }
 
     const auto * orig_kparams = prev.kparams<htp_mm_kernel_params>();
     struct htp_mm_kernel_params kparams;
-    if (!htp_kparams_solve_matmul_add(sess, src0, src1, src2, dst, &kparams) ||
-        orig_kparams->collapse != kparams.collapse) {
+    if (!htp_kparams_solve_matmul_add(sess, src0, src1, src2, dst, &kparams)) {
         return false;
     }
 
@@ -222,6 +227,7 @@ static bool fuse_mul_mat_add(const ggml_hexagon_session * sess, htp_opnode & pre
     prev.opcode  = HTP_OP_MUL_MAT_ADD;
     prev.inputs  = { src0, src1, src2 };
     prev.outputs = { dst };
+    htp_opnode_collapse(sess, prev);
     prev.fused.push_back(cur.node);
     prev.set_kparams(kparams);
     return true;
@@ -233,36 +239,31 @@ static bool fuse_mul_mat_nx(const ggml_hexagon_session * sess, htp_opnode & prev
     const uint32_t curr_n = (uint32_t) prev.outputs.size();
     if (curr_n >= HTP_OP_MAX_OUTPUTS || curr_n + 1 >= HTP_OP_MAX_INPUTS) return false;
 
-    const ggml_tensor * w0   = prev.inputs[0];
-    const ggml_tensor * x    = prev.inputs[curr_n];
-    const ggml_tensor * w_in = cur.src0();
-    const ggml_tensor * x_in = cur.src1();
-    const ggml_tensor * d_in = cur.dst();
+    const ggml_tensor * w0     = prev.inputs[0];
+    const ggml_tensor * x_orig = prev.node->src[1];
+    const ggml_tensor * w_in   = cur.node->src[0];
+    const ggml_tensor * x_in   = cur.node->src[1];
+    const ggml_tensor * d_in   = cur.node;
 
     auto * orig_kparams = prev.kparams<htp_mm_kernel_params>();
     const auto * cur_kparams = cur.kparams<htp_mm_kernel_params>();
 
-    if (x_in != x || w_in->type != w0->type || w_in->ne[0] != w0->ne[0] ||
+    if (x_in != x_orig || w_in->type != w0->type || w_in->ne[0] != w0->ne[0] ||
         w0->ne[2] != 1 || w0->ne[3] != 1 || w_in->ne[2] != 1 || w_in->ne[3] != 1 ||
         htp_tensor_is_tiled(w_in) != htp_tensor_is_tiled(w0) ||
         orig_kparams->n_hmx != cur_kparams->n_hmx) {
         return false;
     }
 
-    const bool collapse = orig_kparams->collapse && htp_matmul_can_collapse(w_in, x, d_in);
-    if ((x->ne[2] > 1 || x->ne[3] > 1) && !collapse) return false;
+    const bool collapse = htp_matmul_can_collapse(w0, x_orig, prev.node) && htp_matmul_can_collapse(w_in, x_in, d_in);
+    if ((x_orig->ne[2] > 1 || x_orig->ne[3] > 1) && !collapse) return false;
 
     if (prev.opcode == HTP_OP_MUL_MAT_NX) {
         orig_kparams->n_weights = curr_n + 1;
     } else {
         struct htp_mm_kernel_params kparams;
-        if (collapse) {
-            const ggml_tensor x_collapsed = htp_tensor_collapse_rows(x);
-            if (!htp_kparams_solve_mmnx(sess, w0, &x_collapsed, curr_n + 1, &kparams)) return false;
-            kparams.collapse = 1;
-        } else if (!htp_kparams_solve_mmnx(sess, w0, x, curr_n + 1, &kparams)) {
-            return false;
-        }
+        const ggml_tensor x_collapsed = collapse ? htp_tensor_collapse_rows(x_orig) : *x_orig;
+        if (!htp_kparams_solve_mmnx(sess, w0, &x_collapsed, curr_n + 1, &kparams)) return false;
         if (!is_supported_mul_mat_nx_kernel(w0, &kparams)) return false;
 
         prev.set_kparams(kparams);
@@ -270,7 +271,7 @@ static bool fuse_mul_mat_nx(const ggml_hexagon_session * sess, htp_opnode & prev
     }
 
     prev.inputs.insert(prev.inputs.begin() + curr_n, w_in);
-    prev.outputs.push_back(d_in);
+    prev.outputs.push_back(collapse ? prev.add_dummy(htp_tensor_collapse_rows(d_in)) : d_in);
     prev.fused.push_back(cur.node);
     return true;
 }
@@ -415,15 +416,13 @@ static bool fuse_unary_mul_glu(const ggml_hexagon_session * sess, htp_opnode & p
     const ggml_tensor * dst  = cur.dst();
     if (htp_inputs_overlap_unsafe(prev, dst, up)) return false;
 
-    struct htp_glu_kernel_params kparams;
-    if (!htp_kparams_solve_unary_mul_glu(sess, glu_op, gate, up, dst, &kparams)) return false;
-
     prev.opcode  = glu_op;
     prev.inputs  = { gate, up };
     prev.outputs = { dst };
+    if (!htp_kparams_solve(sess, prev)) return false;
+
     prev.fused.push_back(cur.node);
     memset(prev.params, 0, sizeof(prev.params));
-    prev.set_kparams(kparams);
     return true;
 }
 

ggml/src/ggml-hexagon/htp-kparams.cpp

@@ -1,6 +1,7 @@
 #define _USE_MATH_DEFINES
 
 #include "htp-kparams.h"
+#include "htp-opnode.h"
 #include "ggml-impl.h"
 
 #include <algorithm>
@@ -397,9 +398,20 @@ static bool htp_kparams_solve_flash_attn(
     kparams->scale         = scale;
     kparams->max_bias      = max_bias;
     kparams->logit_softcap = logit_softcap;
-    kparams->head_split    = (opt_fa_head_split != 0) ? 1 : 0;
+    kparams->head_split    = (opt_fa_split != 0) ? 1 : 0;
+    kparams->flags         = 0;
     kparams->G             = G;
 
+    // team-split only when a team has >1 member (n_kv_heads < N_cores)
+    uint32_t fa_teams = 1, fa_team_size = 1;
+    const uint32_t mdev_count = htp_session_mdev_count(sess);
+    if (opt_fa_split && mdev_count > 1) {
+        htp_fa_mdev_teams(n_kv_heads, mdev_count, kparams->head_split != 0, &fa_teams, &fa_team_size);
+        if (fa_team_size > 1) {
+            kparams->flags |= HTP_FA_FLAG_TEAM_SPLIT;
+        }
+    }
+
     const uint32_t n_head = q->ne[2];
     kparams->n_head_log2 = 1u << (uint32_t) std::floor(std::log2(n_head));
     kparams->m0 = std::pow(2.0f, -(max_bias) / kparams->n_head_log2);
@@ -409,7 +421,7 @@ static bool htp_kparams_solve_flash_attn(
         const uint32_t DK_pad = hex_round_up(DK, 64);
         const uint32_t DV_pad = hex_round_up(DV, 64);
         size_t Br = 0, Bc = 0;
-        int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK_pad, DV_pad, neq1, nek1, vtcm_size, n_threads, (q->type == GGML_TYPE_F32), sinks != nullptr, n_head);
+        int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK_pad, DV_pad, neq1, nek1, vtcm_size, n_threads, (q->type == GGML_TYPE_F32), sinks != nullptr, n_head, fa_team_size);
         if (ret == 0) {
             kparams->kernel_type = HTP_FA_KERNEL_HMX;
             if (logit_softcap == 0.0f) {
@@ -1012,14 +1024,12 @@ static bool htp_kparams_solve_matmul(
     if (htp_matmul_can_collapse(src0, src1, dst)) {
         const ggml_tensor src1_collapsed = htp_tensor_collapse_rows(src1);
         const ggml_tensor dst_collapsed  = htp_tensor_collapse_rows(dst);
-        const bool ok = htp_kparams_solve_matmul_impl(sess, src0, &src1_collapsed, &dst_collapsed, 0, 0, kparams);
-        kparams->collapse = 1;
-        return ok;
+        return htp_kparams_solve_matmul_impl(sess, src0, &src1_collapsed, &dst_collapsed, 0, 0, kparams);
     }
     return htp_kparams_solve_matmul_impl(sess, src0, src1, dst, 0, 0, kparams);
 }
 
-static bool htp_matmul_add_can_collapse(
+bool htp_matmul_add_can_collapse(
     const struct ggml_tensor * src0,
     const struct ggml_tensor * src1,
     const struct ggml_tensor * src2,
@@ -1057,9 +1067,7 @@ bool htp_kparams_solve_matmul_add(
                                            : (src2 ? *src2 : ggml_tensor{});
         const struct ggml_tensor * p_src2 = src2 ? &src2_collapsed : nullptr;
         const size_t src2_size = p_src2 ? hex_round_up(ggml_nbytes(p_src2), 128) : 0;
-        const bool ok = htp_kparams_solve_matmul_impl(sess, src0, &src1_collapsed, &dst_collapsed, p_src2 ? p_src2->nb[1] : 0, src2_size, kparams);
-        kparams->collapse = 1;
-        return ok;
+        return htp_kparams_solve_matmul_impl(sess, src0, &src1_collapsed, &dst_collapsed, p_src2 ? p_src2->nb[1] : 0, src2_size, kparams);
     }
     const size_t src2_size = src2 ? hex_round_up(ggml_nbytes(src2), 128) : 0;
     return htp_kparams_solve_matmul_impl(sess, src0, src1, dst, src2 ? src2->nb[1] : 0, src2_size, kparams);
@@ -1394,47 +1402,103 @@ static bool htp_kparams_solve_binary(
 }
 
 // collapse contiguous dims of elementwise ops so that rows fill at least one HVX vector
-static void htp_unary_collapse_dims(struct ggml_tensor * src, struct ggml_tensor * dst, size_t elem_size) {
-    if (src->nb[0] != elem_size || dst->nb[0] != elem_size) {
+static void htp_collapse_dims(struct ggml_tensor ** tensors, size_t n_tensors, size_t elem_size, uint32_t n_threads) {
+    if (n_tensors == 0 || !tensors[0]) {
         return;
     }
+    for (size_t t = 0; t < n_tensors; t++) {
+        if (!tensors[t] || tensors[t]->nb[0] != elem_size) {
+            return;
+        }
+        if (t > 0 && !ggml_are_same_shape(tensors[0], tensors[t])) {
+            return;
+        }
+    }
 
-    while (src->ne[1] > 1 || src->ne[2] > 1 || src->ne[3] > 1) {
-        const size_t row_size    = (size_t) src->ne[0] * elem_size;
-        const bool   rows_contig = src->nb[1] == row_size && dst->nb[1] == (size_t) dst->ne[0] * elem_size;
-        const bool   can_merge   = (src->ne[1] == 1 && dst->ne[1] == 1) ||
-                                   (rows_contig && (src->ne[0] == 1 || (row_size < 128 && row_size * src->ne[1] <= 2048)));
+    struct ggml_tensor * t0 = tensors[0];
+
+    while (t0->ne[1] > 1 || t0->ne[2] > 1 || t0->ne[3] > 1) {
+        const size_t row_size = (size_t) t0->ne[0] * elem_size;
+        bool rows_contig = true;
+        for (size_t t = 0; t < n_tensors; t++) {
+            if (tensors[t]->nb[1] != (size_t) tensors[t]->ne[0] * elem_size) {
+                rows_contig = false;
+                break;
+            }
+        }
+        bool all_ne1_one = true;
+        for (size_t t = 0; t < n_tensors; t++) {
+            if (tensors[t]->ne[1] != 1) {
+                all_ne1_one = false;
+                break;
+            }
+        }
+        const bool can_merge = all_ne1_one ||
+                               (rows_contig && (t0->ne[0] == 1 || (row_size < 128 && row_size * t0->ne[1] <= 2048)));
         if (!can_merge) {
             break;
         }
-        src->ne[0] *= src->ne[1];
-        src->ne[1]  = src->ne[2];
-        src->ne[2]  = src->ne[3];
-        src->ne[3]  = 1;
-        dst->ne[0] *= dst->ne[1];
-        dst->ne[1]  = dst->ne[2];
-        dst->ne[2]  = dst->ne[3];
-        dst->ne[3]  = 1;
-        src->nb[1]  = src->nb[2];
-        src->nb[2]  = src->nb[3];
-        dst->nb[1]  = dst->nb[2];
-        dst->nb[2]  = dst->nb[3];
+        for (size_t t = 0; t < n_tensors; t++) {
+            tensors[t]->ne[0] *= tensors[t]->ne[1];
+            tensors[t]->ne[1]  = tensors[t]->ne[2];
+            tensors[t]->ne[2]  = tensors[t]->ne[3];
+            tensors[t]->ne[3]  = 1;
+            tensors[t]->nb[1]  = tensors[t]->nb[2];
+            tensors[t]->nb[2]  = tensors[t]->nb[3];
+            tensors[t]->nb[3]  = tensors[t]->nb[2] * tensors[t]->ne[2];
+        }
     }
 
     // ne1 too big to merge in full: merge it in chunks of k rows (k divides ne1)
-    const size_t row_size = (size_t) src->ne[0] * elem_size;
-    if (src->ne[1] > 1 && row_size < 128 && src->nb[1] == row_size && dst->nb[1] == row_size) {
+    const size_t row_size = (size_t) t0->ne[0] * elem_size;
+    bool rows_contig = true;
+    for (size_t t = 0; t < n_tensors; t++) {
+        if (tensors[t]->nb[1] != row_size) {
+            rows_contig = false;
+            break;
+        }
+    }
+    if (t0->ne[1] > 1 && row_size < 128 && rows_contig) {
         int64_t k = 2048 / row_size;
-        while (k > 1 && src->ne[1] % k != 0) {
+        while (k > 1 && t0->ne[1] % k != 0) {
             k--;
         }
         if (k > 1) {
-            src->ne[0] *= k;
-            src->ne[1] /= k;
-            src->nb[1] *= k;
-            dst->ne[0] *= k;
-            dst->ne[1] /= k;
-            dst->nb[1] *= k;
+            for (size_t t = 0; t < n_tensors; t++) {
+                tensors[t]->ne[0] *= k;
+                tensors[t]->ne[1] /= k;
+                tensors[t]->nb[1] *= k;
+                tensors[t]->nb[2]  = tensors[t]->nb[1] * tensors[t]->ne[1];
+                tensors[t]->nb[3]  = tensors[t]->nb[2] * tensors[t]->ne[2];
+            }
+        }
+    }
+
+    // nrows too small to saturate threads: split ne0 into k rows (k divides ne0)
+    const size_t   row_size_curr = (size_t) t0->ne[0] * elem_size;
+    const uint32_t nrows         = (uint32_t) (t0->ne[1] * t0->ne[2] * t0->ne[3]);
+    rows_contig = true;
+    for (size_t t = 0; t < n_tensors; t++) {
+        if (tensors[t]->ne[1] != 1 && tensors[t]->nb[1] != row_size_curr) {
+            rows_contig = false;
+            break;
+        }
+    }
+    if (nrows < n_threads && rows_contig && row_size_curr >= 4096) {
+        const size_t vec_elems = 128 / elem_size;
+        int64_t k = std::max<int64_t>(row_size_curr / 16384, (n_threads * 4 + nrows - 1) / nrows);
+        k = std::min<int64_t>(k, row_size_curr / 2048);
+        while (k > 1 && t0->ne[0] % (k * vec_elems) != 0) {
+            k--;
+        }
+        if (k > 1) {
+            for (size_t t = 0; t < n_tensors; t++) {
+                tensors[t]->ne[0] /= k;
+                tensors[t]->ne[1] *= k;
+                tensors[t]->nb[1]  = tensors[t]->ne[0] * elem_size;
+                tensors[t]->nb[2]  = tensors[t]->nb[1] * tensors[t]->ne[1];
+                tensors[t]->nb[3]  = tensors[t]->nb[2] * tensors[t]->ne[2];
+            }
         }
     }
 }
@@ -1496,19 +1560,6 @@ static bool htp_kparams_solve_unary_impl(
 
     const size_t elem_size = ggml_type_size(src0->type);
 
-    struct ggml_tensor src0_c = *src0;
-    struct ggml_tensor dst_c  = *dst;
-    if (htp_unary_is_elementwise(op)) {
-        htp_unary_collapse_dims(&src0_c, &dst_c, elem_size);
-        src0 = &src0_c;
-        dst  = &dst_c;
-        for (int i = 0; i < HTP_OP_MAX_DIMS; i++) {
-            kparams->ne[i]      = (uint32_t) src0->ne[i];
-            kparams->src0_nb[i] = (uint32_t) src0->nb[i];
-            kparams->dst_nb[i]  = (uint32_t) dst->nb[i];
-        }
-    }
-
     const size_t dst_data_row_size  = dst->ne[0]  * ggml_type_size(dst->type);
 
     const uint32_t src0_nrows          = src0->ne[1] * src0->ne[2] * src0->ne[3];
@@ -1522,9 +1573,14 @@ static bool htp_kparams_solve_unary_impl(
     const bool broadcast_weight = (op == HTP_OP_RMS_NORM_MUL) && (src1->ne[1] * src1->ne[2] * src1->ne[3] == 1);
     kparams->broadcast_weight   = broadcast_weight;
 
-    const bool     src0_contig = ggml_is_contiguous_1(src0);
-    const bool     src1_contig = (op == HTP_OP_RMS_NORM_MUL && !broadcast_weight) ? ggml_is_contiguous_1(src1) : true;
-    const uint32_t max_block   = (src0_contig && src1_contig) ? src0_nrows : (uint32_t) src0->ne[1];
+    const bool     src0_contig           = ggml_is_contiguous_1(src0);
+    const bool     src1_contig           = (op == HTP_OP_RMS_NORM_MUL && !broadcast_weight) ? ggml_is_contiguous_1(src1) : true;
+    const uint32_t src0_nrows_per_thread = (src0_nrows + n_threads - 1) / n_threads;
+    const uint32_t min_block_rows        = (uint32_t) (std::max<size_t>(1, 2048 / dst_data_row_size));
+    uint32_t       max_block             = (src0_contig && src1_contig) ? src0_nrows : (uint32_t) src0->ne[1];
+    if (src0_nrows_per_thread >= 2 * min_block_rows) {
+        max_block = (std::min)(max_block, (std::max)(min_block_rows, src0_nrows_per_thread / 4));
+    }
 
     kparams->src0_contig = src0_contig ? 1 : 0;
     kparams->src1_contig = src1_contig ? 1 : 0;
@@ -1989,11 +2045,17 @@ bool htp_kparams_solve_unary_mul_glu(
     const size_t   vtcm_size           = htp_session_vtcm_size(sess);
 
     const size_t   row_size_aligned    = hex_round_up(dst_row_size, 128);
-    const uint32_t vtcm_row_per_thread = (uint32_t) (vtcm_size / (n_threads * 3 * row_size_aligned));
+    uint32_t       vtcm_row_per_thread = (uint32_t) (vtcm_size / (n_threads * 3 * row_size_aligned));
     if (vtcm_row_per_thread < 2) {
         return false;
     }
 
+    const uint32_t src0_nrows_per_thread = (src0_nrows + n_threads - 1) / n_threads;
+    if (src0_nrows_per_thread >= 2 && vtcm_row_per_thread > src0_nrows_per_thread) {
+        uint32_t target_rpt = (src0_nrows_per_thread >= 4) ? (src0_nrows_per_thread / 2) & ~1u : 2;
+        vtcm_row_per_thread = (std::max)(2u, (std::min)(vtcm_row_per_thread, target_rpt));
+    }
+
     struct htp_glu_vtcm_layout L;
     htp_glu_vtcm_layout_build(&L, dst_row_size, n_threads, vtcm_row_per_thread);
 
@@ -3235,12 +3297,131 @@ static inline bool htp_kparams_invoke(
     return SolveFn(sess, op, kparams ? (KParams *) kparams : &tmp);
 }
 
+void htp_opnode_collapse(const struct ggml_hexagon_session * sess, htp_opnode & node) {
+    const uint32_t n_threads = htp_session_n_threads(sess);
+
+    if (htp_op_is_unary(node.opcode) && htp_unary_is_elementwise(node.opcode)) {
+        if (node.inputs.empty() || !node.inputs[0] || node.outputs.empty() || !node.outputs[0]) {
+            return;
+        }
+        if (!ggml_are_same_shape(node.inputs[0], node.outputs[0])) {
+            return;
+        }
+        const size_t elem_size = ggml_type_size(node.inputs[0]->type);
+        struct ggml_tensor src0_c = *node.inputs[0];
+        struct ggml_tensor dst_c  = *node.outputs[0];
+        struct ggml_tensor * tensors[2] = { &src0_c, &dst_c };
+        htp_collapse_dims(tensors, 2, elem_size, n_threads);
+        if (src0_c.ne[0] != node.inputs[0]->ne[0] || src0_c.ne[1] != node.inputs[0]->ne[1]) {
+            node.inputs[0]  = node.add_dummy(src0_c);
+            node.outputs[0] = node.add_dummy(dst_c);
+        }
+        return;
+    }
+
+    if (htp_op_is_glu(node.opcode)) {
+        if (node.inputs.size() < 2 || !node.inputs[0] || !node.inputs[1] || node.outputs.empty() || !node.outputs[0]) {
+            return;
+        }
+        if (!ggml_are_same_shape(node.inputs[0], node.outputs[0]) ||
+            !ggml_are_same_shape(node.inputs[1], node.outputs[0])) {
+            return;
+        }
+        const size_t elem_size = ggml_type_size(node.inputs[0]->type);
+        struct ggml_tensor src0_c = *node.inputs[0];
+        struct ggml_tensor src1_c = *node.inputs[1];
+        struct ggml_tensor dst_c  = *node.outputs[0];
+        struct ggml_tensor * tensors[3] = { &src0_c, &src1_c, &dst_c };
+        htp_collapse_dims(tensors, 3, elem_size, n_threads);
+        if (src0_c.ne[0] != node.inputs[0]->ne[0] || src0_c.ne[1] != node.inputs[0]->ne[1]) {
+            node.inputs[0]  = node.add_dummy(src0_c);
+            node.inputs[1]  = node.add_dummy(src1_c);
+            node.outputs[0] = node.add_dummy(dst_c);
+        }
+        return;
+    }
+
+    if (node.opcode == HTP_OP_MUL_MAT) {
+        if (node.inputs.size() < 2 || !node.inputs[0] || !node.inputs[1] || node.outputs.empty() || !node.outputs[0]) {
+            return;
+        }
+        const struct ggml_tensor * src0 = node.inputs[0];
+        const struct ggml_tensor * src1 = node.inputs[1];
+        const struct ggml_tensor * dst  = node.outputs[0];
+        if (htp_matmul_can_collapse(src0, src1, dst)) {
+            node.inputs[1]  = node.add_dummy(htp_tensor_collapse_rows(src1));
+            node.outputs[0] = node.add_dummy(htp_tensor_collapse_rows(dst));
+        }
+        return;
+    }
+
+    if (node.opcode == HTP_OP_MUL_MAT_ADD) {
+        if (node.inputs.size() < 2 || !node.inputs[0] || !node.inputs[1] || node.outputs.empty() || !node.outputs[0]) {
+            return;
+        }
+        const struct ggml_tensor * src0 = node.inputs[0];
+        const struct ggml_tensor * src1 = node.inputs[1];
+        const struct ggml_tensor * src2 = node.inputs.size() > 2 ? node.inputs[2] : nullptr;
+        const struct ggml_tensor * dst  = node.outputs[0];
+        if (htp_matmul_add_can_collapse(src0, src1, src2, dst)) {
+            node.inputs[1]  = node.add_dummy(htp_tensor_collapse_rows(src1));
+            node.outputs[0] = node.add_dummy(htp_tensor_collapse_rows(dst));
+            if (src2 && (src2->ne[1] * src2->ne[2] * src2->ne[3] > 1)) {
+                node.inputs[2] = node.add_dummy(htp_tensor_collapse_rows(src2));
+            }
+        }
+        return;
+    }
+}
+
+bool htp_kparams_solve(
+    const struct ggml_hexagon_session * sess,
+    htp_opnode & node
+) {
+    htp_opnode_collapse(sess, node);
+
+    const ggml_tensor * src0 = node.inputs.empty() ? nullptr : node.inputs[0];
+    const ggml_tensor * src1 = node.inputs.size() > 1 ? node.inputs[1] : nullptr;
+    const ggml_tensor * dst  = node.dst();
+
+    if (htp_op_is_unary(node.opcode)) {
+        return htp_kparams_solve_unary_impl(sess, node.opcode, src0, src1, dst, node.kparams<htp_unary_kernel_params>());
+    }
+    if (htp_op_is_glu(node.opcode)) {
+        return htp_kparams_solve_unary_mul_glu(sess, node.opcode, src0, src1, dst, node.kparams<htp_glu_kernel_params>());
+    }
+    if (node.opcode == HTP_OP_MUL_MAT) {
+        if (!htp_matmul_check_types(src0, src1, dst)) {
+            return false;
+        }
+        return htp_kparams_solve_matmul_impl(sess, src0, src1, dst, 0, 0, node.kparams<htp_mm_kernel_params>());
+    }
+    if (node.opcode == HTP_OP_MUL_MAT_ADD) {
+        const ggml_tensor * src2 = node.inputs.size() > 2 ? node.inputs[2] : nullptr;
+        const size_t src2_size = src2 ? hex_round_up(ggml_nbytes(src2), 128) : 0;
+        return htp_kparams_solve_matmul_impl(sess, src0, src1, dst, src2 ? src2->nb[1] : 0, src2_size, node.kparams<htp_mm_kernel_params>());
+    }
+
+    return htp_kparams_solve(sess, node.node, node.kernel_params);
+}
+
 bool htp_kparams_solve(
     const struct ggml_hexagon_session * sess,
     const struct ggml_tensor * op,
     void * kparams
 ) {
     const htp_op_code opcode = htp_op_remap(op);
+    if (htp_op_is_unary(opcode) || htp_op_is_glu(opcode)) {
+        htp_opnode tmp(opcode, op);
+        if (!htp_kparams_solve(sess, tmp)) {
+            return false;
+        }
+        if (kparams) {
+            memcpy(kparams, tmp.kernel_params, sizeof(tmp.kernel_params));
+        }
+        return true;
+    }
+
     switch (opcode) {
         case HTP_OP_MUL_MAT:
         case HTP_OP_MUL_MAT_ID:

ggml/src/ggml-hexagon/htp-kparams.h

@@ -34,6 +34,7 @@
 
 bool        htp_rows_stride(const int64_t * ne, const size_t * nb, size_t * stride);
 bool        htp_matmul_can_collapse(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst);
+bool        htp_matmul_add_can_collapse(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * src2, const struct ggml_tensor * dst);
 ggml_tensor htp_tensor_collapse_rows(const struct ggml_tensor * t);
 
 template <typename... Args>
@@ -45,6 +46,18 @@ static inline bool htp_kparams_check(Args &&... args) {
     return SolveFn(std::forward<Args>(args)..., &kparams);
 }
 
+struct htp_opnode;
+
+bool htp_kparams_solve(
+    const struct ggml_hexagon_session * sess,
+    struct htp_opnode & node
+);
+
+void htp_opnode_collapse(
+    const struct ggml_hexagon_session * sess,
+    struct htp_opnode & node
+);
+
 bool htp_kparams_solve(
     const struct ggml_hexagon_session * sess,
     const struct ggml_tensor * op,

ggml/src/ggml-hexagon/htp/flash-attn-ops.c

@@ -2094,7 +2094,29 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
         const uint32_t dst_e_size = (dst->type == HTP_TYPE_F32) ? sizeof(float) : sizeof(__fp16);
         const bool can_split      = htp_tensor_can_row_partition(dst, dst_e_size);
 
-        if (kparams->head_split && can_split && n_kv_heads >= mdev_count && n_kv_heads % mdev_count == 0) {
+        if ((kparams->flags & HTP_FA_FLAG_TEAM_SPLIT) && can_split) {
+            // KV heads split across teams, Q-blocks split within a team
+            uint32_t n_teams, team_size;
+            htp_fa_mdev_teams(n_kv_heads, mdev_count, kparams->head_split != 0, &n_teams, &team_size);
+
+            const uint32_t team        = mdev_idx / team_size;
+            const uint32_t member      = mdev_idx - team * team_size;
+            const uint32_t kv_per_team = n_kv_heads / n_teams;
+            kv_head_min = team * kv_per_team;
+            kv_head_max = kv_head_min + kv_per_team;
+
+            const uint32_t base  = n_q_blocks / team_size;
+            const uint32_t rem   = n_q_blocks - base * team_size;
+            const uint32_t b0    = member * base + MIN(member, rem);
+            const uint32_t nblk  = base + (member < rem ? 1 : 0);
+
+            if (nblk == 0) {
+                return HTP_STATUS_OK;
+            }
+
+            q_start_min = b0 * Br;
+            q_start_max = MIN((b0 + nblk) * Br, neq1);
+        } else if (kparams->head_split && can_split && n_kv_heads >= mdev_count && n_kv_heads % mdev_count == 0) {
             const uint32_t kv_per_core = n_kv_heads / mdev_count;
             kv_head_min = mdev_idx * kv_per_core;
             kv_head_max = kv_head_min + kv_per_core;

ggml/src/ggml-hexagon/htp/flash-attn-ops.h

@@ -32,9 +32,20 @@ enum htp_fa_kernel_type {
     HTP_FA_KERNEL_HMX
 };
 
+// Multicore: gcd(n_kv_heads, N) teams own the KV heads, team members share the Q-blocks
+#define HTP_FA_FLAG_TEAM_SPLIT 0x01
+
+static inline void htp_fa_mdev_teams(uint32_t n_kv_heads, uint32_t mdev_count, bool head_split,
+                                     uint32_t * n_teams, uint32_t * team_size) {
+    const uint32_t t = (head_split && n_kv_heads > 0 && mdev_count > 1) ? hex_gcd_u32(n_kv_heads, mdev_count) : 1;
+    *n_teams   = t;
+    *team_size = mdev_count > 1 ? mdev_count / t : 1;
+}
+
 struct htp_fa_kernel_params {
     uint8_t  kernel_type;        // enum htp_fa_kernel_type
     uint8_t  head_split;         // 1 = partition by KV heads in multicore, 0 = token partition
+    uint8_t  flags;              // HTP_FA_FLAG_*
     uint8_t  n_threads;          // Number of threads to run
 
     // Common parameters
@@ -321,7 +332,8 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
                                          size_t   n_threads,
                                          bool     is_q_fp32,
                                          bool     has_sinks,
-                                         size_t   n_heads) {
+                                         size_t   n_heads,
+                                         size_t   q_split) {
     const size_t T       = HMX_FP16_TILE_N_ROWS;  // 32
     const size_t br_unit = hmx_ceil_div(T, gqa_factor);
     const size_t bc_unit = HMX_FP16_TILE_N_COLS * 2;  // 64
@@ -361,7 +373,8 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
                 const size_t vecs_per_t     = n_use > 0 ? (n_row_vec_cnt + n_use - 1) / n_use : 1;
 
                 const size_t c_iter_actual  = c_iter_base + c_softmax * vecs_per_t;
-                const size_t cost           = q_blocks * (c_q_fixed + kv_blocks * c_iter_actual);
+                const size_t q_blocks_nsp   = (q_blocks + q_split - 1) / (q_split ? q_split : 1);
+                const size_t cost           = q_blocks_nsp * (c_q_fixed + kv_blocks * c_iter_actual);
                 const size_t mn             = Br * Bc;
 
                 if (cost < best_cost || (cost == best_cost && mn > best_mn)) {

ggml/src/ggml-hexagon/htp/matmul-ops.c

@@ -3775,16 +3775,6 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k
     return HTP_STATUS_OK;
 }
 
-static inline void htp_mm_tensor_collapse_rows(struct htp_tensor * c, const struct htp_tensor * t, uint32_t stride) {
-    *c = *t;
-    c->ne[1] = t->ne[1] * t->ne[2] * t->ne[3];
-    c->ne[2] = 1;
-    c->ne[3] = 1;
-    c->nb[1] = stride;
-    c->nb[2] = c->nb[1] * c->ne[1];
-    c->nb[3] = c->nb[2];
-}
-
 static int op_matmul_impl(struct htp_ops_context * octx) {
     const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
 
@@ -3801,36 +3791,6 @@ static int op_matmul_impl(struct htp_ops_context * octx) {
 }
 
 int op_matmul(struct htp_ops_context * octx) {
-    const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
-
-    if (kparams->collapse) {
-        const struct htp_tensor * act  = octx->src[1];
-        const struct htp_tensor * dst  = octx->dst;
-        const uint32_t s_act = (act->ne[1] > 1) ? act->nb[1] : ((act->ne[2] > 1) ? act->nb[2] : act->nb[3]);
-        const uint32_t sd    = (dst->ne[1] > 1) ? dst->nb[1] : ((dst->ne[2] > 1) ? dst->nb[2] : dst->nb[3]);
-        struct htp_tensor act_collapsed, dst_collapsed;
-        htp_mm_tensor_collapse_rows(&act_collapsed, act, s_act);
-        htp_mm_tensor_collapse_rows(&dst_collapsed, dst, sd);
-        octx->src[1] = &act_collapsed;
-        octx->dst    = &dst_collapsed;
-
-        struct htp_tensor src2_collapsed;
-        const struct htp_tensor * src2 = octx->src[2];
-        if (src2 && (src2->ne[1] * src2->ne[2] * src2->ne[3] > 1)) {
-            const uint32_t s2 = (src2->ne[1] > 1) ? src2->nb[1] : ((src2->ne[2] > 1) ? src2->nb[2] : src2->nb[3]);
-            htp_mm_tensor_collapse_rows(&src2_collapsed, src2, s2);
-            octx->src[2] = &src2_collapsed;
-        }
-
-        const int status = op_matmul_impl(octx);
-        octx->src[1] = act;
-        octx->dst    = dst;
-        if (src2) {
-            octx->src[2] = src2;
-        }
-        return status;
-    }
-
     return op_matmul_impl(octx);
 }
 
@@ -4562,35 +4522,5 @@ static int op_matmul_nx_impl(struct htp_ops_context * octx) {
 }
 
 int op_matmul_nx(struct htp_ops_context * octx) {
-    const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
-
-    if (kparams->collapse) {
-        const uint32_t n_weights = kparams->n_weights;
-        const struct htp_tensor * act = octx->src[n_weights];
-        const uint32_t s1 = (act->ne[1] > 1) ? act->nb[1] : ((act->ne[2] > 1) ? act->nb[2] : act->nb[3]);
-        struct htp_tensor act_collapsed;
-        struct htp_tensor dsts_collapsed[HTP_OP_MAX_OUTPUTS];
-        const struct htp_tensor * orig_dsts[HTP_OP_MAX_OUTPUTS];
-
-        htp_mm_tensor_collapse_rows(&act_collapsed, act, s1);
-        octx->src[n_weights] = &act_collapsed;
-
-        for (uint32_t p = 0; p < n_weights; p++) {
-            orig_dsts[p] = octx->dsts[p];
-            const struct htp_tensor * d = octx->dsts[p];
-            const uint32_t sd = (d->ne[1] > 1) ? d->nb[1] : ((d->ne[2] > 1) ? d->nb[2] : d->nb[3]);
-            htp_mm_tensor_collapse_rows(&dsts_collapsed[p], d, sd);
-            octx->dsts[p] = &dsts_collapsed[p];
-        }
-
-        const int status = op_matmul_nx_impl(octx);
-
-        octx->src[n_weights] = act;
-        for (uint32_t p = 0; p < n_weights; p++) {
-            octx->dsts[p] = orig_dsts[p];
-        }
-        return status;
-    }
-
     return op_matmul_nx_impl(octx);
 }

ggml/src/ggml-hexagon/htp/matmul-ops.h

@@ -61,7 +61,7 @@ enum htp_mm_kernel_type {
 struct htp_mm_kernel_params {
     uint8_t  kernel_type;        // enum htp_mm_kernel_type
     uint8_t  pipeline;           // 1 = pipelined execution, 0 = standard
-    uint8_t  collapse;           // 1 = collapse outer dims into 2D, 0 = standard
+    uint8_t  unused;             // reserved
     uint8_t  n_hmx;              // 1 = use HMX, 0 = use HVX
 
     uint8_t  n_threads;          // Number of threads to spawn

ggml/src/ggml-hexagon/htp/unary-ops.c

@@ -1984,29 +1984,8 @@ static int execute_op_unary(struct htp_ops_context * octx) {
 int op_unary(struct htp_ops_context * octx) {
     switch (octx->src[0]->type) {
         case HTP_TYPE_F32:
-        case HTP_TYPE_F16: {
-            if (!htp_unary_is_elementwise(octx->op)) {
-                return execute_op_unary(octx);
-            }
-
-            const struct htp_unary_kernel_params * kparams = (const struct htp_unary_kernel_params *) octx->kernel_params;
-
-            struct htp_tensor src0_c = *octx->src[0];
-            struct htp_tensor dst_c  = *octx->dst;
-            memcpy(src0_c.ne, kparams->ne,      sizeof(src0_c.ne));
-            memcpy(dst_c.ne,  kparams->ne,      sizeof(dst_c.ne));
-            memcpy(src0_c.nb, kparams->src0_nb, sizeof(src0_c.nb));
-            memcpy(dst_c.nb,  kparams->dst_nb,  sizeof(dst_c.nb));
-
-            const struct htp_tensor * orig_src0 = octx->src[0];
-            const struct htp_tensor * orig_dst  = octx->dst;
-            octx->src[0] = &src0_c;
-            octx->dst    = &dst_c;
-            const int status = execute_op_unary(octx);
-            octx->src[0] = orig_src0;
-            octx->dst    = orig_dst;
-            return status;
-        }
+        case HTP_TYPE_F16:
+            return execute_op_unary(octx);
 
         default:
             return HTP_STATUS_NO_SUPPORT;

ggml/src/ggml-hexagon/htp/unary-ops.h

@@ -24,11 +24,6 @@ struct htp_unary_kernel_params {
     struct fastdiv_values div_ne02;
     struct fastdiv_values div_ne012;
     struct fastdiv_values div_tpr;
-
-    // collapsed shape of elementwise ops
-    uint32_t  ne[HTP_OP_MAX_DIMS];
-    uint32_t  src0_nb[HTP_OP_MAX_DIMS];
-    uint32_t  dst_nb[HTP_OP_MAX_DIMS];
 };
 
 #if defined(__cplusplus)
@@ -72,6 +67,9 @@ static inline bool htp_op_is_unary(uint32_t opcode) {
 }
 
 static inline bool htp_unary_is_elementwise(uint32_t opcode) {
+    if (!htp_op_is_unary(opcode)) {
+        return false;
+    }
     switch (opcode) {
         case HTP_OP_NORM:
         case HTP_OP_RMS_NORM:

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -1225,40 +1225,6 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta
     return res;
 }
 
-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax_part(ggml_metal_library_t lib) {
-    char base[256];
-    char name[256];
-
-    snprintf(base, 256, "kernel_mul_mm_id_amax_part_f32");
-    snprintf(name, 256, "%s", base);
-
-    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
-    if (!res.pipeline) {
-        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
-    }
-
-    res.smem = 32*sizeof(float);
-
-    return res;
-}
-
-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax(ggml_metal_library_t lib) {
-    char base[256];
-    char name[256];
-
-    snprintf(base, 256, "kernel_mul_mm_id_amax_f32");
-    snprintf(name, 256, "%s", base);
-
-    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
-    if (!res.pipeline) {
-        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
-    }
-
-    res.smem = 32*sizeof(float);
-
-    return res;
-}
-
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0(ggml_metal_library_t lib, int ne02, int ne20) {
     char base[256];
     char name[256];
@@ -1301,18 +1267,14 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id(ggml_m
 
     const bool bc_inp = op->src[0]->ne[0] % 32 != 0;
 
-    // src1 prec [TAG_GGML_PREC]
-    const bool amax = ggml_get_op_params_i32(op, 3) == GGML_PREC_F32;
-
     snprintf(base, 256, "kernel_mul_mm_id_%s_%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1));
-    snprintf(name, 256, "%s_bci=%d_amax=%d", base, bc_inp, amax);
+    snprintf(name, 256, "%s_bci=%d", base, bc_inp);
 
     ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
     if (!res.pipeline) {
         ggml_metal_cv_t cv = ggml_metal_cv_init();
 
         ggml_metal_cv_set_bool(cv, bc_inp, FC_MUL_MM + 0);
-        ggml_metal_cv_set_bool(cv, amax,   FC_MUL_MM + 6);
 
         res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
 

ggml/src/ggml-metal/ggml-metal-device.h

@@ -141,8 +141,6 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv            (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0    (ggml_metal_library_t lib, int ne02, int ne20);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_zero    (ggml_metal_library_t lib);
-struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax(ggml_metal_library_t lib);
-struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax_part(ggml_metal_library_t lib);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id         (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id         (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax            (ggml_metal_library_t lib, const struct ggml_tensor * op);

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

@@ -14,8 +14,6 @@
 #define N_MM_SIMD_GROUP_X 2
 #define N_MM_SIMD_GROUP_Y 2
 
-#define N_MM_NPART_AMAX 256
-
 #define N_ARGMAX_NPART_MAX 256
 #define N_ARGMAX_CHUNK 4096
 
@@ -579,14 +577,6 @@ typedef struct {
     uint64_t nb21;
 } ggml_metal_kargs_mul_mm_id_zero;
 
-typedef struct {
-    int32_t  ne00;
-    int32_t  ne01;
-    int32_t  ne02;
-    uint64_t nb01;
-    uint64_t nb02;
-} ggml_metal_kargs_mul_mm_id_amax;
-
 typedef struct {
     int32_t  ne00;
     int32_t  ne02;

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

@@ -2812,15 +2812,6 @@ size_t ggml_metal_op_mul_mat_id_extra_ids(const ggml_tensor * op) {
     return ggml_type_size(GGML_TYPE_I32)*(ne20*ne21 + ne02);
 }
 
-size_t ggml_metal_op_mul_mat_id_extra_amax(const ggml_tensor * op) {
-    assert(op->op == GGML_OP_MUL_MAT_ID);
-
-    GGML_UNUSED(op);
-
-    // 2 scaling factors (8 bytes) + N_MM_NPART_AMAX per-threadgroup scales for stage-1
-    return 8 + N_MM_NPART_AMAX*sizeof(float);
-}
-
 int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
     ggml_tensor * op = ctx->node(idx);
 
@@ -2872,39 +2863,6 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
         ggml_metal_buffer_id bid_ids = bid_tpe;
         bid_ids.offs += ggml_metal_op_mul_mat_id_extra_tpe(op);
 
-        ggml_metal_buffer_id bid_amax = bid_ids;
-        bid_amax.offs += ggml_metal_op_mul_mat_id_extra_ids(op);
-
-        // src1 prec [TAG_GGML_PREC]
-        const bool use_amax = ggml_get_op_params_i32(op, 3) == GGML_PREC_F32;
-
-        // src1 rescale factors, computed before the matmul
-        // ref: https://github.com/ggml-org/llama.cpp/pull/26223
-        if (use_amax) {
-            ggml_metal_kargs_mul_mm_id_amax args = {
-                /*.ne00 =*/ ne10,
-                /*.ne01 =*/ ne11,
-                /*.ne02 =*/ ne12,
-                /*.nb01 =*/ nb11,
-                /*.nb02 =*/ nb12,
-            };
-
-            auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id_amax_part(lib);
-
-            const size_t smem = pipeline.smem;
-
-            GGML_ASSERT(smem <= props_dev->max_theadgroup_memory_size);
-
-            ggml_metal_encoder_set_pipeline(enc, pipeline);
-            ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
-            ggml_metal_encoder_set_buffer  (enc, bid_src1, 1);
-            ggml_metal_encoder_set_buffer  (enc, bid_amax, 2);
-
-            ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
-
-            ggml_metal_encoder_dispatch_threadgroups(enc, N_MM_NPART_AMAX, 1, 1, 256, 1, 1);
-        }
-
         {
             ggml_metal_kargs_mul_mm_id_map0 args = {
                 ne10,
@@ -2954,20 +2912,9 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
             ggml_metal_encoder_dispatch_threadgroups(enc, ne20, ne21, 1, 32, 1, 1);
         }
 
+        // this barrier is always needed because the next kernel has to wait for the id maps to be computed
         ggml_metal_op_concurrency_reset(ctx);
 
-        if (use_amax) {
-            auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id_amax(lib);
-
-            ggml_metal_encoder_set_pipeline(enc, pipeline);
-            ggml_metal_encoder_set_buffer  (enc, bid_amax, 0);
-
-            ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, 32, 1, 1);
-
-            // the next kernel has to wait for the amax data
-            ggml_metal_op_concurrency_reset(ctx);
-        }
-
         {
             auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id(lib, op);
 
@@ -2996,7 +2943,6 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
             ggml_metal_encoder_set_buffer  (enc, bid_tpe,  3);
             ggml_metal_encoder_set_buffer  (enc, bid_ids,  4);
             ggml_metal_encoder_set_buffer  (enc, bid_dst,  5);
-            ggml_metal_encoder_set_buffer  (enc, bid_amax, 6);
 
             const size_t smem = pipeline.smem;
 

ggml/src/ggml-metal/ggml-metal-ops.h

@@ -36,7 +36,6 @@ size_t ggml_metal_op_mul_mat_id_extra_tpe(const struct ggml_tensor * op);
 
 // id map [n_tokens, n_expert]
 size_t ggml_metal_op_mul_mat_id_extra_ids(const struct ggml_tensor * op);
-size_t ggml_metal_op_mul_mat_id_extra_amax(const struct ggml_tensor * op);
 
 // return true if we should use the FA vector kernel for this op
 bool ggml_metal_op_flash_attn_ext_use_vec(const struct ggml_tensor * op);

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

@@ -226,7 +226,6 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_
             {
                 res += ggml_metal_op_mul_mat_id_extra_tpe(tensor);
                 res += ggml_metal_op_mul_mat_id_extra_ids(tensor);
-                res += ggml_metal_op_mul_mat_id_extra_amax(tensor);
             } break;
         case GGML_OP_FLASH_ATTN_EXT:
             {

ggml/src/ggml-metal/kernels/mul_mm.metal

@@ -7,7 +7,6 @@ constant short FC_mul_mm_ne12  [[function_constant(FC_MUL_MM + 2)]];
 constant short FC_mul_mm_ne13  [[function_constant(FC_MUL_MM + 3)]];
 constant short FC_mul_mm_r2    [[function_constant(FC_MUL_MM + 4)]];
 constant short FC_mul_mm_r3    [[function_constant(FC_MUL_MM + 5)]];
-constant bool FC_mul_mm_id_amax [[function_constant(FC_MUL_MM + 6)]];
 constant int32_t FC_mul_mm_id_map0_ne02 [[function_constant(FC_MUL_MM + 7)]];
 
 // each block_q contains 16*nl weights
@@ -532,85 +531,6 @@ kernel void kernel_mul_mm_id_zero(
     }
 }
 
-kernel void kernel_mul_mm_id_amax_part_f32(
-        constant ggml_metal_kargs_mul_mm_id_amax & args,
-        device   const char * src1,
-        device         char * dst,
-        threadgroup    char * shmem [[threadgroup(0)]],
-        uint  tgpig[[threadgroup_position_in_grid]],
-        ushort tiitg[[thread_index_in_threadgroup]],
-        ushort tiisg[[thread_index_in_simdgroup]],
-        ushort sgitg[[simdgroup_index_in_threadgroup]],
-        ushort   ntg[[threads_per_threadgroup]]) {
-    const int nrow = args.ne01*args.ne02;
-
-    float lmax = 0.0f;
-
-    for (int ir = tgpig; ir < nrow; ir += N_MM_NPART_AMAX) {
-        const int i01 = ir % args.ne01;
-        const int i02 = ir / args.ne01;
-
-        device const float * row = (device const float *) (src1 + i02*args.nb02 + i01*args.nb01);
-
-        for (int i00 = tiitg; i00 < args.ne00; i00 += ntg) {
-            lmax = max(lmax, fabs(row[i00]));
-        }
-    }
-
-    float amax = simd_max(lmax);
-
-    threadgroup float * shared_amax = (threadgroup float *) shmem;
-
-    if (ntg > N_SIMDWIDTH) {
-        if (sgitg == 0) {
-            shared_amax[tiisg] = 0.0f;
-        }
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        if (tiisg == 0) {
-            shared_amax[sgitg] = amax;
-        }
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        amax = shared_amax[tiisg];
-        amax = simd_max(amax);
-    }
-
-    if (tiitg == 0) {
-        ((device float *) (dst + 8))[tgpig] = amax;
-    }
-}
-
-kernel void kernel_mul_mm_id_amax_f32(
-        device char * dst,
-        ushort tiitg[[thread_index_in_threadgroup]]) {
-    device const float * part = (device const float *) (dst + 8);
-
-    float amax = 0.0f;
-
-    for (int i = tiitg; i < N_MM_NPART_AMAX; i += N_SIMDWIDTH) {
-        amax = max(amax, part[i]);
-    }
-
-    amax = simd_max(amax);
-
-    if (tiitg == 0) {
-        // leave a comfortable margin below the f16 max of 65504
-        float scale = 1.0f;
-
-        // isfinite: src1 already inf/nan is not ours to fix - keep the
-        // scale at 1.0 instead of turning it into a different failure
-        if (isfinite(amax) && amax > 32768.0f) {
-            scale = exp2(ceil(log2(amax)) - 15.0f);
-        }
-
-        device float * d = (device float *) dst;
-
-        d[0] = 1.0f/scale; // exact: scale is a power of two
-        d[1] = scale;
-    }
-}
-
 typedef decltype(kernel_mul_mm_id_map0<1>) kernel_mul_mm_id_map0_t;
 
 template [[host_name("kernel_mul_mm_id_map0_ne20_1" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<1>;
@@ -633,7 +553,6 @@ kernel void kernel_mul_mm_id(
         device const char * htpe,
         device const char * hids,
         device       char * dst,
-        device const char * amax,
         threadgroup  char * shmem [[threadgroup(0)]],
         uint3  tgpig[[threadgroup_position_in_grid]],
         ushort tiitg[[thread_index_in_threadgroup]],
@@ -689,10 +608,6 @@ kernel void kernel_mul_mm_id(
     // skip the upper half of the token tile when the expert did not fill it
     constexpr short NR1H = NR1/2;
 
-    // power-of-two rescaling
-    const float s1_inv   = FC_mul_mm_id_amax ? ((device const float *) amax)[0] : 1.0f;
-    const float s1_scale = FC_mul_mm_id_amax ? ((device const float *) amax)[1] : 1.0f;
-
 #ifndef GGML_METAL_HAS_TENSOR
     S0_8x8 ma[4];
     S1_8x8 mb[2];
@@ -798,7 +713,7 @@ kernel void kernel_mul_mm_id(
 
                 const short ib = 4*sx + sy;
 
-                *(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) (*((device T1 *) y + i) * (T1) s1_inv) : 0;
+                *(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
             }
         } else {
             const short sx = (tiitg%NL1);
@@ -811,7 +726,7 @@ kernel void kernel_mul_mm_id(
 
             const short ib = 4*sx + sy;
 
-            *(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)((*((device T1_2x4 *) y)) * (T1) s1_inv);
+            *(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)(*((device T1_2x4 *) y));
         }
 #else
         // load data and store to threadgroup memory
@@ -859,7 +774,7 @@ kernel void kernel_mul_mm_id(
                 //const short lx = (tiitg/NL1)%8;
                 //const short ly = i;
 
-                *(sb + NK*(8*sy + ly) + 8*sx + lx) = loop_k + iy + i < args.ne00 ? (S1) (*((device T1 *) y + i) * (T1) s1_inv) : 0;
+                *(sb + NK*(8*sy + ly) + 8*sx + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
             }
         } else {
             const short sx = (tiitg%NL1);
@@ -870,7 +785,7 @@ kernel void kernel_mul_mm_id(
             //const short lx = (tiitg/NL1)%8;
             //const short ly = i;
 
-            *(threadgroup S1_2x4 *)(sb + NK*(8*sy + ly) + 8*sx) = (S1_2x4)((*((device T1_2x4 *) y)) * (T1) s1_inv);
+            *(threadgroup S1_2x4 *)(sb + NK*(8*sy + ly) + 8*sx) = (S1_2x4)(*((device T1_2x4 *) y));
         }
 #endif
 
@@ -961,12 +876,12 @@ kernel void kernel_mul_mm_id(
 
         int i = tiisg;
         for (; i < nr0/4; i += 32) {
-            *(D4 + i) = *(C4 + i) * s1_scale;
+            *(D4 + i) = *(C4 + i);
         }
 
         i = (4*(nr0/4)) + tiisg;
         for (; i < nr0; i += 32) {
-            *(D + i) = *(C + i) * s1_scale;
+            *(D + i) = *(C + i);
         }
     }
 }

ggml/src/ggml-opencl/CMakeLists.txt

@@ -66,6 +66,8 @@ set(GGML_OPENCL_KERNELS
     moe_add_id_glu
     argmax
     argsort
+    top_k
+    lightning_indexer
     tri
     fill
     clamp
@@ -99,6 +101,12 @@ set(GGML_OPENCL_KERNELS
     mul_mv_q4_0_f32_1d_16x_flat
     mul_mv_q4_1_f32
     mul_mv_q4_1_f32_flat
+    gemm_noshuffle_q2_k_q8_1_dp4a
+    gemv_noshuffle_q2_k_f32
+    gemm_noshuffle_q2_k_f32
+    mul_mv_q2_k_f32
+    mul_mv_q2_k_f32_flat
+    mul_mm_q2_k_f32_l4_lm
     mul_mv_q4_k_f32
     mul_mv_q4_k_f32_flat
     mul_mv_q5_0_f32

ggml/src/ggml-opencl/kernels/cvt.cl

@@ -2661,3 +2661,115 @@ kernel void kernel_moe_expand_scale_q5_K(
         dst_min[((long)e*ne01 + row)*nblk32 + sub] = (half)(dm * (float)mn);
     }
 }
+
+struct block_q2_K {
+    uint8_t scales[QK_K/16]; // low nibble scale, high nibble min, one per 16 weights
+    uint8_t qs[QK_K/4];      // quants, 2 bits each
+    half d;                  // super-block scale for the quantized scales
+    half dmin;               // super-block scale for the quantized mins
+};
+
+kernel void kernel_convert_block_q2_K(
+    global struct block_q2_K * src0,
+    global uchar * dst_qs,
+    global uchar * dst_sc,
+    global half * dst_dm,
+    ulong n_blk
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blk) {
+        return;
+    }
+    global struct block_q2_K * b = src0 + i;
+    for (int j = 0; j < 64; ++j) {
+        dst_qs[64*i + j] = b->qs[j];
+    }
+    for (int j = 0; j < 16; ++j) {
+        dst_sc[16*i + j] = b->scales[j];
+    }
+    dst_dm[2*i + 0] = b->d;
+    dst_dm[2*i + 1] = b->dmin;
+}
+
+kernel void kernel_restore_block_q2_K(
+    global uchar * src_qs,
+    global uchar * src_sc,
+    global half * src_dm,
+    global struct block_q2_K * dst,
+    ulong n_blk
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blk) {
+        return;
+    }
+    global struct block_q2_K * b = dst + i;
+    for (int j = 0; j < 64; ++j) {
+        b->qs[j] = src_qs[64*i + j];
+    }
+    for (int j = 0; j < 16; ++j) {
+        b->scales[j] = src_sc[16*i + j];
+    }
+    b->d = src_dm[2*i + 0];
+    b->dmin = src_dm[2*i + 1];
+}
+
+kernel void kernel_convert_block_q2_k_noshuffle(
+    global struct block_q2_K * src0,
+    global uchar * dst_qs,      // QK_K/4 per block
+    global uchar * dst_sc,      // QK_K/16 per block
+    global half  * dst_dm,      // 2 per block (d, dmin)
+    ulong          n_blk
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blk) {
+        return;
+    }
+    global struct block_q2_K * b = (global struct block_q2_K *) src0 + i;
+
+    dst_dm[2*i + 0] = b->d;
+    dst_dm[2*i + 1] = b->dmin;
+
+    for (int j = 0; j < QK_K/16; ++j) { dst_sc[(QK_K/16) * i + j] = b->scales[j]; }
+
+    for (int g = 0; g < QK_K/4; ++g) {
+        uint pk = 0;
+        for (int t = 0; t < 4; ++t) {
+            const int e  = 4*g + t;
+            const int sb = e >> 5;
+            const int ee = e & 31;
+            pk |= (uint)(((b->qs[32*(sb >> 2) + ee] >> (2*(sb & 3))) & 3) << (2*t));
+        }
+        dst_qs[(QK_K/4) * i + g] = (uchar)pk;
+    }
+}
+
+kernel void kernel_restore_block_q2_k_noshuffle(
+    global uchar * src_qs,
+    global uchar * src_sc,
+    global half  * src_dm,
+    global struct block_q2_K * dst,
+    ulong          n_blk
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blk) {
+        return;
+    }
+    global struct block_q2_K * b = (global struct block_q2_K *) dst + i;
+
+    b->d    = src_dm[2*i + 0];
+    b->dmin = src_dm[2*i + 1];
+
+    for (int j = 0; j < QK_K/16; ++j) { b->scales[j] = src_sc[(QK_K/16) * i + j]; }
+
+    for (int n2 = 0; n2 < 2; ++n2) {
+        for (int ee = 0; ee < 32; ++ee) {
+            uint v = 0;
+            for (int j = 0; j < 4; ++j) {
+                const int sb = 4*n2 + j;
+                const int e  = 32*sb + ee;
+                v |= ((uint)(src_qs[(QK_K/4) * i + (e >> 2)] >> (2*(e & 3))) & 3u) << (2*j);
+            }
+            b->qs[32*n2 + ee] = (uchar)v;
+        }
+    }
+}

ggml/src/ggml-opencl/kernels/gemm_noshuffle_q2_k_f32.cl

@@ -0,0 +1,83 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_qcom_reqd_sub_group_size
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+__attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+kernel void kernel_gemm_noshuffle_q2_k_f32(
+    global const uchar * src0_qs,
+    global const uchar * src0_sc,
+    global const half  * src0_dm,
+    read_only image1d_buffer_t src1,
+    global float * dst,
+    ulong offsetd,
+    int m,
+    int n,
+    int k,
+    int n_no_padding
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+    const int gy = get_global_id(0);
+    const int row = get_global_id(1) * 4;
+    const int n_4 = n / 4;
+    if (row >= m) {
+        return;
+    }
+
+    half8 c0 = 0, c1 = 0, c2 = 0, c3 = 0;
+    for (int i = 0; i < k; i += 16) {
+        const half8 dm = vload8(0, src0_dm + 2 * (row + (i / 256) * m));
+        const uchar4 sc = vload4(0, src0_sc + row + (i / 16) * m);
+        const half4 scale = convert_half4(convert_float4(dm.s0246) * convert_float4(sc & (uchar4)(15)));
+        const half4 mval  = convert_half4(convert_float4(dm.s1357) * convert_float4(sc >> (uchar4)(4)));
+
+        for (int l = 0; l < 16; l += 4) {
+            const int ki = i + l;
+            const uchar4 bits = vload4(0, src0_qs + row + (ki / 4) * m);
+            #pragma unroll
+            for (int j = 0; j < 4; ++j) {
+                half8 b;
+                b.s0123 = read_imageh(src1, gy * 2     + (ki + j) * n_4);
+                b.s4567 = read_imageh(src1, gy * 2 + 1 + (ki + j) * n_4);
+                const half4 w = convert_half4((bits >> (uchar4)(2 * j)) & (uchar4)(3)) * scale - mval;
+                c0 += b * w.s0;
+                c1 += b * w.s1;
+                c2 += b * w.s2;
+                c3 += b * w.s3;
+            }
+        }
+    }
+
+    int idx = gy * 8 * m + row;
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s0, c1.s0, c2.s0, c3.s0), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s1, c1.s1, c2.s1, c3.s1), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s2, c1.s2, c2.s2, c3.s2), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s3, c1.s3, c2.s3, c3.s3), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s4, c1.s4, c2.s4, c3.s4), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s5, c1.s5, c2.s5, c3.s5), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s6, c1.s6, c2.s6, c3.s6), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s7, c1.s7, c2.s7, c3.s7), 0, dst + idx);
+    }
+}

ggml/src/ggml-opencl/kernels/gemm_noshuffle_q2_k_q8_1_dp4a.cl

@@ -0,0 +1,146 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+#ifdef cl_khr_integer_dot_product
+#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
+#endif
+
+#define QK_K 256
+
+#ifndef TILESIZE_N
+#define TILESIZE_N 32
+#endif
+
+inline uint q2_k_pack(uint pk) {
+    return ( (pk      ) & 3u)
+         | (((pk >> 2) & 3u) <<  8)
+         | (((pk >> 4) & 3u) << 16)
+         | (((pk >> 6) & 3u) << 24);
+}
+
+// Pass staged operands by value; local dp4a operands miscompile on X2.
+#define KQ_STAGE(p) vload4(0, (p))
+
+inline int dot4_q8a_v(uint4 qw, uint4 a) {
+    int r = 0;
+    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a.w, r);
+    return r;
+}
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_q2_k_q8_1_dp4a(
+        __global const uchar  * src0_qs,
+        __global const uchar  * src0_sc,
+        __global const half   * src0_dm,   // half pairs: d, dmin
+        __global const uint   * src1_qa,   // q8_1 activations int8 (as uint, 4/elem) [N, K]
+        __global const half   * src1_da,   // q8_1 per-block scale [N, K/32]
+        __global       float  * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local float sh_s [TILESIZE_N][2];   // per-16 activation sums, in qa units
+    __local half  sh_d [TILESIZE_N];
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub = step >> 5;
+        const uint ib  = sub >> 3;
+
+        const half2 dm  = vload2(rrow + ib * (uint)m, src0_dm);
+        const float dv  = (float)dm.s0;
+        const float dmv = (float)dm.s1;
+
+        const uint  scb = rrow + (sub << 1) * (uint)m;
+        const uint  sc0 = (uint)src0_sc[scb + 0u * (uint)m];
+        const uint  sc1 = (uint)src0_sc[scb + 1u * (uint)m];
+        const float dl0 = dv  * (float)( sc0       & 0xFu);
+        const float ml0 = dmv * (float)( sc0 >> 4);
+        const float dl1 = dv  * (float)( sc1       & 0xFu);
+        const float ml1 = dmv * (float)( sc1 >> 4);
+
+        const uint qsb = rrow + (step >> 2) * (uint)m;
+        uint4 qlo, qhi;
+        qlo.s0 = q2_k_pack((uint)src0_qs[qsb + 0u * (uint)m]);
+        qlo.s1 = q2_k_pack((uint)src0_qs[qsb + 1u * (uint)m]);
+        qlo.s2 = q2_k_pack((uint)src0_qs[qsb + 2u * (uint)m]);
+        qlo.s3 = q2_k_pack((uint)src0_qs[qsb + 3u * (uint)m]);
+        qhi.s0 = q2_k_pack((uint)src0_qs[qsb + 4u * (uint)m]);
+        qhi.s1 = q2_k_pack((uint)src0_qs[qsb + 5u * (uint)m]);
+        qhi.s2 = q2_k_pack((uint)src0_qs[qsb + 6u * (uint)m]);
+        qhi.s3 = q2_k_pack((uint)src0_qs[qsb + 7u * (uint)m]);
+
+        for (uint idx = lid; idx < TILESIZE_N * 2u; idx += 64u) {
+            const uint t  = idx >> 1;
+            const uint h  = idx & 1u;
+            const uint c  = col_base + t;
+            const bool ok = c < (uint)n_no_padding;
+            const uint4 w = ok ? KQ_STAGE(src1_qa + c * k_u + (step >> 2) + (h << 2))
+                               : (uint4)(0u);
+            sh_qa4[t][h] = w;
+            int s = 0;
+            s = dot_acc_sat_4x8packed_ss_int(w.x, 0x01010101u, s);
+            s = dot_acc_sat_4x8packed_ss_int(w.y, 0x01010101u, s);
+            s = dot_acc_sat_4x8packed_ss_int(w.z, 0x01010101u, s);
+            s = dot_acc_sat_4x8packed_ss_int(w.w, 0x01010101u, s);
+            sh_s[t][h] = (float)s;
+            if (h == 0u) {
+                sh_d[t] = ok ? src1_da[c * k_b + sub] : (half)0;
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
+#define Q2_K_COL(b) (dl0 * (float)dot4_q8a_v(qlo, (uint4)(sh_qa4[b][0])) - ml0 * sh_s[b][0] + \
+                    dl1 * (float)dot4_q8a_v(qhi, (uint4)(sh_qa4[b][1])) - ml1 * sh_s[b][1])
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+            rf.s0 = Q2_K_COL(b+0);  rf.s1 = Q2_K_COL(b+1);
+            rf.s2 = Q2_K_COL(b+2);  rf.s3 = Q2_K_COL(b+3);
+            acc[g] += LD4(sh_d, b) * rf;
+        }
+#undef Q2_K_COL
+#undef LD4
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
+        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
+        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
+        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
+    }
+#undef NGROUPS
+}

ggml/src/ggml-opencl/kernels/gemv_noshuffle_q2_k_f32.cl

@@ -0,0 +1,231 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#define QK_K 256
+
+#ifndef Q2_K_MV_NSG
+#define Q2_K_MV_NSG 8
+#endif
+
+#ifndef Q2_K_MV_R
+#define Q2_K_MV_R 4
+#endif
+
+inline float4 q2_k_vals(uint pk) {
+    return (float4)((float)( pk       & 3u), (float)((pk >> 2) & 3u),
+                    (float)((pk >> 4) & 3u), (float)((pk >> 6) & 3u));
+}
+
+kernel void kernel_gemv_noshuffle_q2_k_f32(
+        global const uchar * src0_qs,
+        global const uchar * src0_sc,
+        global const half  * src0_dm,
+        global const float * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,      // K
+        int ne01,      // M
+        int ne10,      // activation row stride, == K
+        int ne0        // dst row stride
+) {
+    src1 = (global const float *)((global const char *)src1 + offset1);
+    dst  = (global float       *)((global char       *)dst  + offsetd);
+
+    const uint m   = (uint)ne01;
+    const uint K   = (uint)ne00;
+    const uint nsb = K / QK_K;
+
+    const uint lid = get_local_id(0);
+    const uint sgi = get_local_id(1);
+    const uint col = get_group_id(1);
+
+    global const float * y = src1 + (ulong)col * (uint)ne10;
+
+    const uint mr  = m / Q2_K_MV_R;                  // row groups
+    const uint j   = get_group_id(0) * 64u + lid;   // row group this lane owns
+    const uint row = j * Q2_K_MV_R;
+
+#if Q2_K_MV_R == 4
+    float sumf0 = 0.f;
+    float sumf1 = 0.f;
+    float sumf2 = 0.f;
+    float sumf3 = 0.f;
+
+    if (j < mr) {
+        global const uint * qsu = (global const uint *)src0_qs;
+        global const uint * scu = (global const uint *)src0_sc;
+
+        for (uint ib = sgi; ib < nsb; ib += Q2_K_MV_NSG) {
+            const half4 dma = vload4(2u*(j + ib*mr) + 0u, src0_dm);  // d0, dmin0, d1, dmin1
+            const half4 dmb = vload4(2u*(j + ib*mr) + 1u, src0_dm);  // d2, dmin2, d3, dmin3
+
+            float ad0 = 0.f, am0 = 0.f, ad1 = 0.f, am1 = 0.f;
+            float ad2 = 0.f, am2 = 0.f, ad3 = 0.f, am3 = 0.f;
+
+            for (uint sb = 0; sb < 8u; ++sb) {
+                const uint grp = ib * 64u + sb * 8u;
+                const uint qsb = j + grp * mr;
+
+                for (uint h = 0; h < 2u; ++h) {
+                    const uint scv = scu[j + (2u * (ib * 8u + sb) + h) * mr];
+
+                    float a0 = 0.f, a1 = 0.f, a2 = 0.f, a3 = 0.f;
+                    float4 as = (float4)(0.f);
+                    for (uint u = 0; u < 4u; ++u) {
+                        const uint gg  = 4u*h + u;
+                        const uint qsv = qsu[qsb + gg * mr];   // four rows, one load
+                        const float4 yv = vload4(grp + gg, y);
+                        as += yv;
+                        a0 += dot(yv, q2_k_vals( qsv        & 0xFFu));
+                        a1 += dot(yv, q2_k_vals((qsv >>  8) & 0xFFu));
+                        a2 += dot(yv, q2_k_vals((qsv >> 16) & 0xFFu));
+                        a3 += dot(yv, q2_k_vals((qsv >> 24) & 0xFFu));
+                    }
+                    const float asum = as.s0 + as.s1 + as.s2 + as.s3;
+
+                    const uint s0 =  scv        & 0xFFu;
+                    const uint s1 = (scv >>  8) & 0xFFu;
+                    const uint s2 = (scv >> 16) & 0xFFu;
+                    const uint s3 = (scv >> 24) & 0xFFu;
+                    ad0 += (float)(s0 & 0xFu) * a0;   am0 += (float)(s0 >> 4) * asum;
+                    ad1 += (float)(s1 & 0xFu) * a1;   am1 += (float)(s1 >> 4) * asum;
+                    ad2 += (float)(s2 & 0xFu) * a2;   am2 += (float)(s2 >> 4) * asum;
+                    ad3 += (float)(s3 & 0xFu) * a3;   am3 += (float)(s3 >> 4) * asum;
+                }
+            }
+            sumf0 += (float)dma.s0 * ad0 - (float)dma.s1 * am0;
+            sumf1 += (float)dma.s2 * ad1 - (float)dma.s3 * am1;
+            sumf2 += (float)dmb.s0 * ad2 - (float)dmb.s1 * am2;
+            sumf3 += (float)dmb.s2 * ad3 - (float)dmb.s3 * am3;
+        }
+    }
+
+#if Q2_K_MV_NSG > 1
+    __local float4 part[Q2_K_MV_NSG][64];
+    part[sgi][lid] = (float4)(sumf0, sumf1, sumf2, sumf3);
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (sgi != 0) {
+        return;
+    }
+    for (uint s = 1; s < Q2_K_MV_NSG; ++s) {
+        const float4 p = part[s][lid];
+        sumf0 += p.s0;
+        sumf1 += p.s1;
+        sumf2 += p.s2;
+        sumf3 += p.s3;
+    }
+#endif
+
+    if (j < mr) {
+        global float * o = dst + (ulong)col * (uint)ne0 + row;
+        o[0] = sumf0; o[1] = sumf1; o[2] = sumf2; o[3] = sumf3;
+    }
+
+#elif Q2_K_MV_R == 2
+    float sumf0 = 0.f;
+    float sumf1 = 0.f;
+
+    if (j < mr) {
+        global const ushort * qsu = (global const ushort *)src0_qs;
+        global const ushort * scu = (global const ushort *)src0_sc;
+
+        for (uint ib = sgi; ib < nsb; ib += Q2_K_MV_NSG) {
+            const half4 dm = vload4(j + ib * mr, src0_dm);  // d0, dmin0, d1, dmin1
+
+            float ad0 = 0.f, am0 = 0.f, ad1 = 0.f, am1 = 0.f;
+            for (uint sb = 0; sb < 8u; ++sb) {
+                const uint grp = ib * 64u + sb * 8u;
+                const uint qsb = j + grp * mr;
+
+                for (uint h = 0; h < 2u; ++h) {
+                    const uint scv = (uint)scu[j + (2u * (ib * 8u + sb) + h) * mr];
+                    const uint s0  = scv & 0xFFu;
+                    const uint s1  = scv >> 8;
+
+                    float a0 = 0.f, a1 = 0.f;
+                    float4 as = (float4)(0.f);
+                    for (uint u = 0; u < 4u; ++u) {
+                        const uint gg  = 4u*h + u;
+                        const uint qsv = (uint)qsu[qsb + gg * mr];
+                        const float4 yv = vload4(grp + gg, y);
+                        as += yv;
+                        a0 += dot(yv, q2_k_vals( qsv       & 0xFFu));
+                        a1 += dot(yv, q2_k_vals((qsv >> 8) & 0xFFu));
+                    }
+                    const float asum = as.s0 + as.s1 + as.s2 + as.s3;
+                    ad0 += (float)(s0 & 0xFu) * a0;   am0 += (float)(s0 >> 4) * asum;
+                    ad1 += (float)(s1 & 0xFu) * a1;   am1 += (float)(s1 >> 4) * asum;
+                }
+            }
+            sumf0 += (float)dm.s0 * ad0 - (float)dm.s1 * am0;
+            sumf1 += (float)dm.s2 * ad1 - (float)dm.s3 * am1;
+        }
+    }
+
+#if Q2_K_MV_NSG > 1
+    __local float2 part[Q2_K_MV_NSG][64];
+    part[sgi][lid] = (float2)(sumf0, sumf1);
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (sgi != 0) {
+        return;
+    }
+    for (uint s = 1; s < Q2_K_MV_NSG; ++s) {
+        const float2 p = part[s][lid];
+        sumf0 += p.s0;
+        sumf1 += p.s1;
+    }
+#endif
+
+    if (j < mr) {
+        vstore2((float2)(sumf0, sumf1), 0, dst + (ulong)col * (uint)ne0 + row);
+    }
+
+#else
+    float sumf0 = 0.f;
+
+    if (j < mr) {
+        for (uint ib = sgi; ib < nsb; ib += Q2_K_MV_NSG) {
+            const half2 dm = vload2(row + ib * m, src0_dm);
+
+            float ad = 0.f, am = 0.f;
+            for (uint sb = 0; sb < 8u; ++sb) {
+                const uint grp = ib * 64u + sb * 8u;
+                const uint qsb = row + grp * m;
+
+                for (uint h = 0; h < 2u; ++h) {
+                    const uint sc = (uint)src0_sc[row + (2u * (ib * 8u + sb) + h) * m];
+
+                    float a = 0.f;
+                    float4 as = (float4)(0.f);
+                    for (uint u = 0; u < 4u; ++u) {
+                        const uint gg = 4u*h + u;
+                        const float4 yv = vload4(grp + gg, y);
+                        as += yv;
+                        const uint pkv = (uint)src0_qs[qsb + gg * m];
+                        a += dot(yv, q2_k_vals(pkv));
+                    }
+                    const float asum = as.s0 + as.s1 + as.s2 + as.s3;
+                    ad += (float)(sc & 0xFu) * a;   am += (float)(sc >> 4) * asum;
+                }
+            }
+            sumf0 += (float)dm.s0 * ad - (float)dm.s1 * am;
+        }
+    }
+
+#if Q2_K_MV_NSG > 1
+    __local float part[Q2_K_MV_NSG][64];
+    part[sgi][lid] = sumf0;
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (sgi != 0) {
+        return;
+    }
+    for (uint s = 1; s < Q2_K_MV_NSG; ++s) {
+        sumf0 += part[s][lid];
+    }
+#endif
+
+    if (j < mr) {
+        dst[(ulong)col * (uint)ne0 + row] = sumf0;
+    }
+#endif
+}

ggml/src/ggml-opencl/kernels/lightning_indexer.cl

@@ -0,0 +1,111 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#ifdef cl_khr_subgroup_shuffle
+#pragma OPENCL EXTENSION cl_khr_subgroup_shuffle : enable
+#endif
+
+kernel void kernel_lightning_indexer(
+        global const char * q,
+        ulong offset_q,
+        global const char * k,
+        ulong offset_k,
+        global const char * w,
+        ulong offset_w,
+        global const char * mask,
+        ulong offset_m,
+        global float * dst,
+        ulong offset_d,
+        int head_size,
+        int n_heads,
+        int n_tokens,
+        int n_keys,
+        int n_masks,
+        int type_k,
+        ulong nbq1,
+        ulong nbq2,
+        ulong nbq3,
+        ulong nbk2,
+        ulong nbk3,
+        ulong nbw1,
+        ulong nbw3,
+        ulong nbm1,
+        ulong nbm3,
+        int use_shuffle
+) {
+    const int tid = get_local_id(0);
+    const int lane = tid & (use_shuffle ? 3 : 15);
+    const int key = use_shuffle ? get_group_id(0)*16 + tid/4 : get_group_id(0)*4 + tid/16;
+    const int token = get_group_id(1);
+    const int stream = get_group_id(2);
+
+    q += offset_q + token*nbq2 + stream*nbq3;
+    k += offset_k + (ulong) key*nbk2 + stream*nbk3;
+    w += offset_w + token*nbw1 + stream*nbw3;
+    mask += offset_m + token*nbm1 + (stream % n_masks)*nbm3;
+    dst = (global float *) ((global char *) dst + offset_d);
+
+#ifdef cl_khr_subgroup_shuffle
+    if (use_shuffle) {
+        // Each group of four lanes keeps one key for all query heads.
+        float4 cached[8];
+        #pragma unroll
+        for (int j = 0; j < 8; ++j) {
+            const int i = 4*(lane + j*4);
+            cached[j] = key < n_keys ? (type_k == 0 ? vload4(0, (global const float *) k + i) : convert_float4(vload4(0, (global const half *) k + i))) : (float4) (0.0f);
+        }
+
+        float score = 0.0f;
+        for (int h = 0; h < n_heads; ++h) {
+            global const float * q_head = (global const float *) (q + h*nbq1);
+            float4 sum = (float4) (0.0f);
+            #pragma unroll
+            for (int j = 0; j < 8; ++j) {
+                sum = mad(vload4(0, q_head + 4*(lane + j*4)), cached[j], sum);
+            }
+            float value = sum.s0 + sum.s1 + sum.s2 + sum.s3;
+            value += sub_group_shuffle_xor(value, 2);
+            value += sub_group_shuffle_xor(value, 1);
+            if (lane == 0) {
+                score = mad(fmax(value, 0.0f), ((global const float *) w)[h], score);
+            }
+        }
+        if (lane == 0 && key < n_keys) {
+            dst[((ulong) stream*n_tokens + token)*n_keys + key] = score + (float) ((global const half *) mask)[key];
+        }
+        return;
+    }
+#endif
+
+    local float sums[64];
+    float score = 0.0f;
+    for (int h = 0; h < n_heads; ++h) {
+        global const float * q_head = (global const float *) (q + h*nbq1);
+        float sum = 0.0f;
+        if (key < n_keys) {
+            for (int i = lane; i < head_size; i += 16) {
+                float kv;
+                if (type_k == 0) {
+                    kv = ((global const float *) k)[i];
+                } else {
+                    kv = (float) ((global const half *) k)[i];
+                }
+                sum = mad(q_head[i], kv, sum);
+            }
+        }
+        sums[tid] = sum;
+        barrier(CLK_LOCAL_MEM_FENCE);
+        for (int stride = 8; stride > 0; stride /= 2) {
+            if (lane < stride) {
+                sums[tid] += sums[tid + stride];
+            }
+            barrier(CLK_LOCAL_MEM_FENCE);
+        }
+        if (lane == 0) {
+            score = mad(fmax(sums[tid], 0.0f), ((global const float *) w)[h], score);
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (lane == 0 && key < n_keys) {
+        dst[((ulong) stream*n_tokens + token)*n_keys + key] = score + (float) ((global const half *) mask)[key];
+    }
+}

ggml/src/ggml-opencl/kernels/mul_mm_q2_k_f32_l4_lm.cl

@@ -0,0 +1,203 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_intel_required_subgroup_size
+#define INTEL_GPU 1
+#endif
+
+#define QK_K 256
+
+struct block_q2_K {
+    uchar scales[QK_K/16];
+    uchar qs[QK_K/4];
+    half  d;
+    half  dmin;
+};
+
+#define LOAD_VEC_A 4
+#define LOAD_VEC_B 4
+
+#define BM 64
+#define BN 64
+#ifndef BK
+#define BK 32
+#endif
+#ifdef INTEL_GPU
+#define TM 8
+#else
+#define TM 4
+#endif
+#define TN 8
+
+#ifndef LM_HALF
+#define LM_HALF 0
+#endif
+#if LM_HALF
+typedef half lm_st;
+#define LM_LD4(p, i) convert_float4(vload4((i), (p)))
+#else
+typedef float lm_st;
+#define LM_LD4(p, i) vload4((i), (p))
+#endif
+
+kernel void kernel_mul_mm_q2_k_f32_l4_lm(
+    global uchar * src0_qs,
+    global uchar * src0_sc,
+    global half * src0_dm,
+    ulong offset0,
+    global float4 * src1,
+    ulong offset1,
+    global float  * dst,
+    ulong offsetd,
+
+    int ne00,
+    int ne01,
+    int ne02,
+    int ne11,
+    int ne12,
+
+    int stride_a,
+    int stride_b,
+    int stride_d,
+
+    int batch_stride_a,
+    int batch_stride_b,
+    int batch_stride_d,
+
+    int r2,
+    int r3
+) {
+    src1 = (global float4*)((global char*)src1 + offset1);
+    dst  = (global float *)((global char*)dst  + offsetd);
+
+    local lm_st buf_a[BM * BK];
+    local lm_st buf_b[BN * BK];
+
+    const int batch_idx = get_global_id(2);
+
+    const int i13 = batch_idx / ne12;
+    const int i12 = batch_idx % ne12;
+
+    const int i03 = i13 / r3;
+    const int i02 = i12 / r2;
+
+    const int batch_idx_a = i03 * ne02 + i02;
+
+    const int ir = get_group_id(0);
+    const int ic = get_group_id(1);
+
+    const int tid = get_local_id(0);
+    const int th_r  = tid % (BM / TM);
+    const int th_c  = tid / (BM / TM);
+
+    const int loadr_a = get_local_id(0) % (BK / LOAD_VEC_A);
+    const int loadc_a = get_local_id(0) / (BK / LOAD_VEC_A);
+    const int loadr_b = get_local_id(0) % (BK / LOAD_VEC_B);
+    const int loadc_b = get_local_id(0) / (BK / LOAD_VEC_B);
+
+    const int loadstride_a = get_local_size(0) * LOAD_VEC_A / BK;
+    const int loadstride_b = get_local_size(0) * LOAD_VEC_B / BK;
+
+    // pos_a counts elements, not blocks - one q2_K super block holds QK_K of them
+    int pos_a = batch_idx_a * batch_stride_a + ir * BM * stride_a;
+    int pos_b = (batch_idx   * batch_stride_b + ic * BN * stride_b) / LOAD_VEC_B;
+
+    float4 sums4[(TM/4) * TN];
+    float4 cache_a4[TM/4];
+
+    for (int i = 0; i < (TM/4) * TN; i++) {
+        sums4[i] = (float4)(0.0f);
+    }
+
+    for (int block = 0; block < ne00; block += BK) {
+        for (int l = 0; l < BM; l += loadstride_a) {
+            if (ir*BM + loadc_a + l < ne01) {
+                int idx = pos_a + (loadc_a + l) * stride_a + loadr_a * LOAD_VEC_A;
+                int ib  = idx / QK_K;
+                ib += offset0/sizeof(struct block_q2_K);
+                int e   = idx % QK_K;
+
+                // dequantize_row_q2_K walks 2 halves of 128, each as 4 shifts of
+                // 2 groups of 16. Recover that position from the element index.
+                int n     = e >> 7;
+                int rem   = e & 127;
+                int j     = rem >> 5;
+                int s     = (rem >> 4) & 1;
+                int l16   = rem & 15;
+                int isc   = 8*n + 2*j + s;
+                int shift = 2*j;
+
+                uchar sc = src0_sc[16*ib + isc];
+                float d  =  (float)src0_dm[2*ib]    * (float)(sc & 0xF);
+                float m  = -(float)src0_dm[2*ib + 1] * (float)(sc >> 4);
+
+                uchar4 q = vload4(0, src0_qs + 64*ib + 32*n + 16*s + l16);
+                float4 v1 = convert_float4((uchar4)((q.s0 >> shift) & 3,
+                                                    (q.s1 >> shift) & 3,
+                                                    (q.s2 >> shift) & 3,
+                                                    (q.s3 >> shift) & 3)) * d + m;
+
+                buf_a[(loadr_a * LOAD_VEC_A + 0) * BM + loadc_a + l] = v1.s0;
+                buf_a[(loadr_a * LOAD_VEC_A + 1) * BM + loadc_a + l] = v1.s1;
+                buf_a[(loadr_a * LOAD_VEC_A + 2) * BM + loadc_a + l] = v1.s2;
+                buf_a[(loadr_a * LOAD_VEC_A + 3) * BM + loadc_a + l] = v1.s3;
+            } else {
+                buf_a[(loadr_a * LOAD_VEC_A + 0) * BM + loadc_a + l] = 0.0f;
+                buf_a[(loadr_a * LOAD_VEC_A + 1) * BM + loadc_a + l] = 0.0f;
+                buf_a[(loadr_a * LOAD_VEC_A + 2) * BM + loadc_a + l] = 0.0f;
+                buf_a[(loadr_a * LOAD_VEC_A + 3) * BM + loadc_a + l] = 0.0f;
+            }
+        }
+
+        for (int l = 0; l < BN; l += loadstride_b) {
+            if (ic*BN + loadc_b + l < ne11) {
+                int idx = pos_b + (loadc_b + l) * stride_b / LOAD_VEC_B + loadr_b;
+                buf_b[(loadr_b * LOAD_VEC_B + 0) * BN + loadc_b + l] = src1[idx].s0;
+                buf_b[(loadr_b * LOAD_VEC_B + 1) * BN + loadc_b + l] = src1[idx].s1;
+                buf_b[(loadr_b * LOAD_VEC_B + 2) * BN + loadc_b + l] = src1[idx].s2;
+                buf_b[(loadr_b * LOAD_VEC_B + 3) * BN + loadc_b + l] = src1[idx].s3;
+            } else {
+                buf_b[(loadr_b * LOAD_VEC_B + 0) * BN + loadc_b + l] = 0.0f;
+                buf_b[(loadr_b * LOAD_VEC_B + 1) * BN + loadc_b + l] = 0.0f;
+                buf_b[(loadr_b * LOAD_VEC_B + 2) * BN + loadc_b + l] = 0.0f;
+                buf_b[(loadr_b * LOAD_VEC_B + 3) * BN + loadc_b + l] = 0.0f;
+            }
+        }
+
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+        pos_a += BK;
+        pos_b += BK / LOAD_VEC_B;
+
+        for (int i = 0; i < BK; i++) {
+            for (int a = 0; a < TM/4; a++) {
+                cache_a4[a] = LM_LD4(buf_a + (i) * BM + th_r * TM, a);
+            }
+
+            for (int cc = 0; cc < TN; cc++) {
+                const float cache_b = buf_b[(i) * BN + th_c * TN + cc];
+                for (int a = 0; a < TM/4; a++) {
+                    const int sums_idx = cc*(TM/4) + a;
+                    sums4[sums_idx] = mad(cache_a4[a], (float4)cache_b, sums4[sums_idx]);
+                }
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    const int dr = ir * BM + th_r * TM;
+    const int dc = ic * BN + th_c * TN;
+
+    const ulong offsets = batch_idx * batch_stride_d;
+
+    for (int cc = 0; cc < TN; cc++) {
+        for (int a = 0; a < TM/4; a++) {
+            const float4 v = sums4[cc * (TM/4) + a];
+            const float  vs[4] = { v.s0, v.s1, v.s2, v.s3 };
+            for (int k = 0; k < 4; k++) {
+                if (dr + 4*a + k < ne01 && dc + cc < ne11) {
+                    dst[offsets + (dc + cc) * stride_d + dr + 4*a + k] = vs[k];
+                }
+            }
+        }
+    }
+}

ggml/src/ggml-opencl/kernels/mul_mv_q2_k_f32.cl

@@ -0,0 +1,161 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_intel_required_subgroup_size
+#pragma OPENCL EXTENSION cl_intel_required_subgroup_size : enable
+#define INTEL_GPU 1
+#define REQD_SUBGROUP_SIZE_16 __attribute__((intel_reqd_sub_group_size(16)))
+#define REQD_SUBGROUP_SIZE_32 __attribute__((intel_reqd_sub_group_size(32)))
+#elif defined(cl_qcom_reqd_sub_group_size)
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+#define ADRENO_GPU 1
+#define REQD_SUBGROUP_SIZE_64  __attribute__((qcom_reqd_sub_group_size("half")))
+#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+
+//------------------------------------------------------------------------------
+// block_q2_K
+//------------------------------------------------------------------------------
+#define QK_K 256
+
+// 16 blocks of 16 elements each
+// weight is represented as x = a * q + b
+typedef struct {
+    uchar scales[QK_K/16]; // scales and mins, quantized with 4 bits
+    uchar qs[QK_K/4];      // 2-bit quants
+    half  d;              // super-block scale for quantized scales
+    half  dmin;            // super-block scale for quantized mins
+} block_q2_K;
+
+#undef N_DST
+#undef N_SIMDGROUP
+#undef N_SIMDWIDTH
+
+#ifdef INTEL_GPU
+#define N_DST 4 // number of rows each SIMD group works on
+#define N_SIMDGROUP 1 // number of SIMD groups in a thread group
+#define N_SIMDWIDTH 16 // SIMD group size
+#elif defined (ADRENO_GPU)
+#define N_DST 4
+#define N_SIMDGROUP 1
+#define N_SIMDWIDTH 64
+#endif
+
+#undef  BLOCK_STRIDE
+// number of (super) blocks each subgroup processes
+// 8 threads cover one super block, so a wave covers N_SIMDWIDTH/8 of them
+#define BLOCK_STRIDE (N_SIMDWIDTH/8)
+
+#ifdef INTEL_GPU
+REQD_SUBGROUP_SIZE_16
+#elif defined (ADRENO_GPU)
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mv_q2_K_f32(
+        global char * src0,
+        ulong offset0,
+        global char * src1,
+        ulong offset1,
+        global char * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne12,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src0 = src0 + offset0;
+    src1 = src1 + offset1;
+    dst  = dst  + offsetd;
+
+    int ix = get_sub_group_local_id()/8;  // super block index
+    int it = get_sub_group_local_id()%8;  // block index (inside super block)
+    int iq = it/4;      // 0 or 1 - first or second half of the super block
+    int ir = it%4;      // 0...3 - block index in the half super block
+    int is = (8*ir)/16; // 0 or 1 - scale pair
+
+    int nb = ne00/QK_K;
+
+    int r0 = get_group_id(0);
+    int r1 = get_group_id(1);
+    int im = get_group_id(2);
+    int first_row = (r0 * N_SIMDGROUP + get_sub_group_id()) * N_DST;
+
+    int i12 = im%ne12;
+    int i13 = im/ne12;
+
+    ulong offset_src0 = first_row*nb01 + (i12/r2)*nb02 + (i13/r3)*nb03;
+    ulong offset_src1 =        r1*nb11 + (i12   )*nb12 + (i13   )*nb13;
+
+    global block_q2_K * x = (global block_q2_K *) (src0 + offset_src0);
+    global float      * y = (global float      *) (src1 + offset_src1);
+
+    float yl[32];
+    float sumf[N_DST] = {0.f};
+    float all_sum;
+
+    global float * y4 = y + ix * QK_K + 128 * iq + 8 * ir;
+
+    for (int ib = ix; ib < nb; ib += BLOCK_STRIDE) {
+        float4 sumy = {0.f, 0.f, 0.f, 0.f};
+        for (int i = 0; i < 8; ++i) {
+            yl[i+ 0] = y4[i+ 0]; sumy.s0 += yl[i+ 0];
+            yl[i+ 8] = y4[i+32]; sumy.s1 += yl[i+ 8];
+            yl[i+16] = y4[i+64]; sumy.s2 += yl[i+16];
+            yl[i+24] = y4[i+96]; sumy.s3 += yl[i+24];
+        }
+
+        global uchar  * sc = (global uchar  *)x[ib].scales + 8*iq + is;
+        global ushort * qs = (global ushort *)x[ib].qs + 16 * iq + 4 * ir;
+        global half   * dh = &x[ib].d;
+
+        // only ne01 output rows exist; reading past them can pull in an inf scale
+        for (int row = 0; row < N_DST && first_row + row < ne01; row++) {
+            float4 acc1 = {0.f, 0.f, 0.f, 0.f};
+            float4 acc2 = {0.f, 0.f, 0.f, 0.f};
+            for (int i = 0; i < 8; i += 2) {
+                acc1.s0 += yl[i+ 0] * (qs[i/2] & 0x0003);
+                acc2.s0 += yl[i+ 1] * (qs[i/2] & 0x0300);
+                acc1.s1 += yl[i+ 8] * (qs[i/2] & 0x000c);
+                acc2.s1 += yl[i+ 9] * (qs[i/2] & 0x0c00);
+                acc1.s2 += yl[i+16] * (qs[i/2] & 0x0030);
+                acc2.s2 += yl[i+17] * (qs[i/2] & 0x3000);
+                acc1.s3 += yl[i+24] * (qs[i/2] & 0x00c0);
+                acc2.s3 += yl[i+25] * (qs[i/2] & 0xc000);
+            }
+
+            float dall = dh[0];
+            float dmin = dh[1] * 1.f/16.f;
+            sumf[row] += dall * ((acc1.s0 + 1.f/256.f * acc2.s0) * (sc[0] & 0xF) * 1.f/ 1.f +
+                                 (acc1.s1 + 1.f/256.f * acc2.s1) * (sc[2] & 0xF) * 1.f/ 4.f +
+                                 (acc1.s2 + 1.f/256.f * acc2.s2) * (sc[4] & 0xF) * 1.f/16.f +
+                                 (acc1.s3 + 1.f/256.f * acc2.s3) * (sc[6] & 0xF) * 1.f/64.f) -
+                        dmin * (sumy.s0 * (sc[0] & 0xF0) + sumy.s1 * (sc[2] & 0xF0) +
+                                sumy.s2 * (sc[4] & 0xF0) + sumy.s3 * (sc[6] & 0xF0));
+
+            qs += nb01/2;
+            sc += nb01;
+            dh += nb01/2;
+        }
+
+        y4 += BLOCK_STRIDE * QK_K;
+    }
+
+    global float * dst_f32 = (global float *) dst + im*ne0*ne1 + r1*ne0;
+
+    for (int row = 0; row < N_DST; ++row) {
+        all_sum = sub_group_reduce_add(sumf[row]);
+        if (first_row + row < ne01) {
+            if (get_sub_group_local_id() == 0) {
+                dst_f32[first_row + row] = all_sum;
+            }
+        }
+    }
+}

ggml/src/ggml-opencl/kernels/mul_mv_q2_k_f32_flat.cl

@@ -0,0 +1,162 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_intel_required_subgroup_size
+#pragma OPENCL EXTENSION cl_intel_required_subgroup_size : enable
+#define INTEL_GPU 1
+#define REQD_SUBGROUP_SIZE_16 __attribute__((intel_reqd_sub_group_size(16)))
+#define REQD_SUBGROUP_SIZE_32 __attribute__((intel_reqd_sub_group_size(32)))
+#elif defined(cl_qcom_reqd_sub_group_size)
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+#define ADRENO_GPU 1
+#define REQD_SUBGROUP_SIZE_64  __attribute__((qcom_reqd_sub_group_size("half")))
+#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+
+//------------------------------------------------------------------------------
+// struct block_q2_K
+//------------------------------------------------------------------------------
+#define QK_K 256
+
+// 16 blocks of 16 elements each
+// weight is represented as x = a * q + b
+struct block_q2_K {
+    uchar scales[QK_K/16]; // scales and mins, quantized with 4 bits
+    uchar qs[QK_K/4];      // 2-bit quants
+    half  d;              // super-block scale for quantized scales
+    half  dmin;            // super-block scale for quantized mins
+};
+
+#undef N_DST
+#undef N_SIMDGROUP
+#undef N_SIMDWIDTH
+
+#ifdef INTEL_GPU
+#define N_DST 4 // number of rows each SIMD group works on
+#define N_SIMDGROUP 1 // number of SIMD groups in a thread group
+#define N_SIMDWIDTH 16 // SIMD group size
+#elif defined (ADRENO_GPU)
+#define N_DST 4
+#define N_SIMDGROUP 1
+#define N_SIMDWIDTH 64
+#endif
+
+#undef  BLOCK_STRIDE
+// number of (super) blocks each subgroup processes
+// 8 threads cover one super block, so a wave covers N_SIMDWIDTH/8 of them
+#define BLOCK_STRIDE (N_SIMDWIDTH/8)
+
+#ifdef INTEL_GPU
+REQD_SUBGROUP_SIZE_16
+#elif defined (ADRENO_GPU)
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mv_q2_K_f32_flat(
+        global uchar * src0_qs,
+        global uchar * src0_sc,
+        global half * src0_dm,
+        ulong offset0,
+        global char * src1,
+        ulong offset1,
+        global char * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne12,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src1 = src1 + offset1;
+    dst  = dst  + offsetd;
+
+    int ix = get_sub_group_local_id()/8;  // super block index
+    int it = get_sub_group_local_id()%8;  // block index (inside super block)
+    int iq = it/4;      // 0 or 1 - first or second half of the super block
+    int ir = it%4;      // 0...3 - block index in the half super block
+    int is = (8*ir)/16; // 0 or 1 - scale pair
+
+    int nb = ne00/QK_K;
+
+    int r0 = get_group_id(0);
+    int r1 = get_group_id(1);
+    int im = get_group_id(2);
+    int first_row = (r0 * N_SIMDGROUP + get_sub_group_id()) * N_DST;
+
+    int i12 = im%ne12;
+    int i13 = im/ne12;
+
+    ulong offset_src0 = (offset0 + first_row*nb01 + (i12/r2)*nb02 + (i13/r3)*nb03)/sizeof(struct block_q2_K);
+    ulong row_stride = nb01/sizeof(struct block_q2_K);
+    ulong offset_src1 =        r1*nb11 + (i12   )*nb12 + (i13   )*nb13;
+
+    global float      * y = (global float      *) (src1 + offset_src1);
+
+    float yl[32];
+    float sumf[N_DST] = {0.f};
+    float all_sum;
+
+    global float * y4 = y + ix * QK_K + 128 * iq + 8 * ir;
+
+    for (int ib = ix; ib < nb; ib += BLOCK_STRIDE) {
+        float4 sumy = {0.f, 0.f, 0.f, 0.f};
+        for (int i = 0; i < 8; ++i) {
+            yl[i+ 0] = y4[i+ 0]; sumy.s0 += yl[i+ 0];
+            yl[i+ 8] = y4[i+32]; sumy.s1 += yl[i+ 8];
+            yl[i+16] = y4[i+64]; sumy.s2 += yl[i+16];
+            yl[i+24] = y4[i+96]; sumy.s3 += yl[i+24];
+        }
+
+        global uchar  * sc = (src0_sc + 16*(offset_src0 + ib)) + 8*iq + is;
+        global ushort * qs = (global ushort *)(src0_qs + 64*(offset_src0 + ib)) + 16 * iq + 4 * ir;
+        global half   * dh = src0_dm + 2*(offset_src0 + ib);
+
+        // only ne01 output rows exist; reading past them can pull in an inf scale
+        for (int row = 0; row < N_DST && first_row + row < ne01; row++) {
+            float4 acc1 = {0.f, 0.f, 0.f, 0.f};
+            float4 acc2 = {0.f, 0.f, 0.f, 0.f};
+            for (int i = 0; i < 8; i += 2) {
+                acc1.s0 += yl[i+ 0] * (qs[i/2] & 0x0003);
+                acc2.s0 += yl[i+ 1] * (qs[i/2] & 0x0300);
+                acc1.s1 += yl[i+ 8] * (qs[i/2] & 0x000c);
+                acc2.s1 += yl[i+ 9] * (qs[i/2] & 0x0c00);
+                acc1.s2 += yl[i+16] * (qs[i/2] & 0x0030);
+                acc2.s2 += yl[i+17] * (qs[i/2] & 0x3000);
+                acc1.s3 += yl[i+24] * (qs[i/2] & 0x00c0);
+                acc2.s3 += yl[i+25] * (qs[i/2] & 0xc000);
+            }
+
+            float dall = dh[0];
+            float dmin = dh[1] * 1.f/16.f;
+            sumf[row] += dall * ((acc1.s0 + 1.f/256.f * acc2.s0) * (sc[0] & 0xF) * 1.f/ 1.f +
+                                 (acc1.s1 + 1.f/256.f * acc2.s1) * (sc[2] & 0xF) * 1.f/ 4.f +
+                                 (acc1.s2 + 1.f/256.f * acc2.s2) * (sc[4] & 0xF) * 1.f/16.f +
+                                 (acc1.s3 + 1.f/256.f * acc2.s3) * (sc[6] & 0xF) * 1.f/64.f) -
+                        dmin * (sumy.s0 * (sc[0] & 0xF0) + sumy.s1 * (sc[2] & 0xF0) +
+                                sumy.s2 * (sc[4] & 0xF0) + sumy.s3 * (sc[6] & 0xF0));
+
+            qs += row_stride*32;
+            sc += row_stride*16;
+            dh += row_stride*2;
+        }
+
+        y4 += BLOCK_STRIDE * QK_K;
+    }
+
+    global float * dst_f32 = (global float *) dst + im*ne0*ne1 + r1*ne0;
+
+    for (int row = 0; row < N_DST; ++row) {
+        all_sum = sub_group_reduce_add(sumf[row]);
+        if (first_row + row < ne01) {
+            if (get_sub_group_local_id() == 0) {
+                dst_f32[first_row + row] = all_sum;
+            }
+        }
+    }
+}

ggml/src/ggml-opencl/kernels/top_k.cl

@@ -0,0 +1,207 @@
+#ifdef cl_khr_subgroups
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+#endif
+
+uint top_k_key(float value) {
+    uint bits = as_uint(value);
+    if ((bits & 0x7fffffffu) == 0) {
+        bits = 0;
+    }
+    return bits & 0x80000000u ? ~bits : bits ^ 0x80000000u;
+}
+
+uint4 top_k_keys(global const char * src, int i, int n, ulong nb0) {
+    if (nb0 == sizeof(float) && i + 3 < n) {
+        const float4 v = vload4(0, (global const float *) src + i);
+        return (uint4) (top_k_key(v.s0), top_k_key(v.s1), top_k_key(v.s2), top_k_key(v.s3));
+    }
+    return (uint4) (
+        i     < n ? top_k_key(*(global const float *) (src + (ulong)  i     *nb0)) : 0,
+        i + 1 < n ? top_k_key(*(global const float *) (src + (ulong) (i + 1)*nb0)) : 0,
+        i + 2 < n ? top_k_key(*(global const float *) (src + (ulong) (i + 2)*nb0)) : 0,
+        i + 3 < n ? top_k_key(*(global const float *) (src + (ulong) (i + 3)*nb0)) : 0);
+}
+
+void top_k_store(global int * dst, global uint2 * candidates, global const char * src,
+        ulong nb0, int indexed, int base, int i, uint pos) {
+    const int index = indexed ? *(global const int *) (src + (ulong) i*nb0 + sizeof(float)) : base + i;
+    if (candidates) {
+        candidates[pos] = (uint2) (as_uint(*(global const float *) (src + (ulong) i*nb0)), index);
+    } else {
+        dst[pos] = index;
+    }
+}
+
+void top_k_select(
+        global const char * src,
+        global int * dst,
+        global uint2 * candidates,
+        ulong nb00,
+        int ne00,
+        int k,
+        int base,
+        int indexed,
+        local uint * hist
+) {
+    const int tid = get_local_id(0);
+    const int nth = get_local_size(0);
+    if (k == ne00) {
+        for (int i = tid; i < ne00; i += nth) {
+            top_k_store(dst, candidates, src, nb00, indexed, base, i, i);
+        }
+        return;
+    }
+
+    uint prefix = 0;
+    uint mask = 0;
+    uint rank = k;
+    for (int shift = 24; shift >= 0; shift -= 8) {
+        for (int i = tid; i < 256; i += nth) {
+            hist[i] = 0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+        for (int i = 4*tid; i < ne00; i += 4*nth) {
+            const uint4 v = top_k_keys(src, i, ne00, nb00);
+            const uint keys[4] = {v.s0, v.s1, v.s2, v.s3};
+            #pragma unroll
+            for (int j = 0; j < 4; ++j) {
+                const uint key = keys[j];
+                if (i + j < ne00 && (key & mask) == prefix) {
+                    atomic_inc(hist + ((key >> shift) & 255));
+                }
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+        if (tid == 0) {
+            uint above = 0;
+            for (int b = 255; b >= 0; --b) {
+                uint count = hist[b];
+                if (above + count >= rank) {
+                    hist[256] = b;
+                    hist[257] = rank - above;
+                    hist[258] = count;
+                    break;
+                }
+                above += count;
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+        prefix |= hist[256] << shift;
+        rank = hist[257];
+        const bool done = hist[258] == rank;
+        barrier(CLK_LOCAL_MEM_FENCE);
+        if (done) {
+            // The whole bucket is selected, so its low bits do not matter.
+            break;
+        }
+        mask |= 255u << shift;
+    }
+
+    uint greater = 0;
+    uint equal = 0;
+    for (int i = 4*tid; i < ne00; i += 4*nth) {
+        const uint4 v = top_k_keys(src, i, ne00, nb00);
+        const uint keys[4] = {v.s0, v.s1, v.s2, v.s3};
+        #pragma unroll
+        for (int j = 0; j < 4; ++j) {
+            greater += i + j < ne00 && keys[j] > prefix;
+            equal += i + j < ne00 && keys[j] == prefix;
+        }
+    }
+    uint pos_greater;
+    uint pos_equal;
+#ifdef cl_khr_subgroups
+    const uint sg = get_sub_group_id();
+    const uint nsg = get_num_sub_groups();
+    const uint total_greater = sub_group_reduce_add(greater);
+    const uint total_equal = sub_group_reduce_add(equal);
+    pos_greater = sub_group_scan_exclusive_add(greater);
+    pos_equal = sub_group_scan_exclusive_add(equal);
+    if (get_sub_group_local_id() == 0) {
+        hist[sg] = total_greater;
+        hist[nsg + sg] = total_equal;
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+    for (uint i = 0; i < nsg; ++i) {
+        pos_equal += hist[i];
+        if (i < sg) {
+            pos_greater += hist[i];
+            pos_equal += hist[nsg + i];
+        }
+    }
+#else
+    hist[tid] = greater;
+    hist[nth + tid] = equal;
+    barrier(CLK_LOCAL_MEM_FENCE);
+    pos_greater = 0;
+    pos_equal = 0;
+    for (int i = 0; i < nth; ++i) {
+        pos_equal += hist[i];
+        if (i < tid) {
+            pos_greater += hist[i];
+            pos_equal += hist[nth + i];
+        }
+    }
+#endif
+    for (int i = 4*tid; i < ne00; i += 4*nth) {
+        const uint4 v = top_k_keys(src, i, ne00, nb00);
+        const uint keys[4] = {v.s0, v.s1, v.s2, v.s3};
+        #pragma unroll
+        for (int j = 0; j < 4; ++j) {
+            if (i + j < ne00) {
+                if (keys[j] > prefix) {
+                    top_k_store(dst, candidates, src, nb00, indexed, base, i + j, pos_greater++);
+                } else if (keys[j] == prefix && pos_equal < k) {
+                    top_k_store(dst, candidates, src, nb00, indexed, base, i + j, pos_equal++);
+                }
+            }
+        }
+    }
+}
+
+kernel void kernel_top_k_f32(
+        global const char * src,
+        ulong offset0,
+        global int * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        int ne02,
+        ulong nb00,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int k_max,
+        local uint * hist,
+        global uint2 * candidates,
+        int tile_width
+) {
+    const ulong row = get_group_id(1);
+    const int part = get_group_id(0);
+    const int parts = get_num_groups(0);
+    const int base = part*tile_width;
+    const int n = min(ne00 - base, tile_width);
+    const int k = min(k_max, n);
+    const int out_width = (parts - 1)*k_max + min(k_max, ne00 - (parts - 1)*tile_width);
+    src += offset0 + (row % ne01)*nb01 + (row / ne01 % ne02)*nb02 + (row / ne01 / ne02)*nb03 + (ulong) base*nb00;
+    if (candidates) {
+        candidates += row*out_width + part*k_max;
+    } else {
+        dst = (global int *) ((global char *) dst + offsetd) + row*out_width + part*k_max;
+    }
+    top_k_select(src, dst, candidates, nb00, n, k, base, 0, hist);
+}
+
+kernel void kernel_top_k_merge_f32(
+        global const uint2 * candidates,
+        global int * dst,
+        ulong offsetd,
+        int n_candidates,
+        int k,
+        local uint * hist
+) {
+    const ulong row = get_group_id(1);
+    candidates += row*n_candidates;
+    dst = (global int *) ((global char *) dst + offsetd) + row*k;
+    top_k_select((global const char *) candidates, dst, 0, sizeof(uint2), n_candidates, k, 0, 1, hist);
+}

ggml/src/ggml-rpc/transport.cpp

@@ -517,6 +517,11 @@ bool socket_t::impl::rdma_recv(void * data, size_t size) {
 
         int slot = (int)wc.wr_id;
         size_t got = wc.byte_len;
+        if (got > rem) {
+            LOG_ERROR("RDMA recv overrun: got=%zu rem=%zu\n", got, rem);
+            c->post_rx(slot);
+            return false;
+        }
         memcpy(dst, c->rx_slot(slot), got);
 
         if (!c->post_rx(slot)) return false;

ggml/src/ggml-vulkan/ggml-vulkan-types.h

@@ -382,6 +382,7 @@ enum vk_conv_shapes {
     CONV_SHAPE_64x32,
     CONV_SHAPE_32x256,
     CONV_SHAPE_64x128,
+    CONV_SHAPE_256x128,
     CONV_SHAPE_COUNT,
 };
 
@@ -397,6 +398,7 @@ inline vk_conv_block_size vk_conv_block_sizes[CONV_SHAPE_COUNT] = {
     {  64,  32, 32 }, // CONV_SHAPE_64x32
     {  32, 256, 16 }, // CONV_SHAPE_32x256
     {  64, 128, 16 }, // CONV_SHAPE_64x128
+    { 256, 128, 16 }, // CONV_SHAPE_256x128
 };
 
 enum dmmv_wg_sizes {

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

@@ -1690,6 +1690,11 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     const uint32_t mul_mat_mm_warp_8  = std::min(mul_mat_subgroup_size_8,  64u);
     const uint32_t mul_mat_mm_warp_16 = std::min(mul_mat_subgroup_size_16, 64u);
 
+    // small warptiles split BM=32 between BLOCK_SIZE/WARP subgroups, so each subgroup needs its own warp tile
+    const uint32_t s_warptile_wm_8          = 32 * subgroup_size_8 / subgroup_size_32;
+    const uint32_t mul_mat_s_warptile_wm_8  = 32 * mul_mat_subgroup_size_8 / mul_mat_subgroup_size_32;
+    const uint32_t mul_mat_s_warptile_wm_16 = 32 * mul_mat_subgroup_size_16 / mul_mat_subgroup_size_32;
+
     const bool subgroup_min_size_16 = (!device->subgroup_size_control && device->subgroup_size >= 16) ||
                                       (device->subgroup_size_control && device->subgroup_max_size >= 16);
 
@@ -1784,20 +1789,18 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
         const uint32_t itn = device->coopmat_int_n;
         const uint32_t itk = device->coopmat_int_k;
 
-        const uint32_t s_warptile_wm = device->subgroup_size == 8 ? 8 : 32;
-
         l_warptile = { 128,             128, 128, 16, mm_warp_8 * 2, 64, 2, tm_l, tn_l, tk_l, mm_warp_8 };
         m_warptile = { 128,              64,  64, 16, mm_warp_8,     32, 2, tm_m, tn_m, tk_m, mm_warp_8 };
-        s_warptile = { subgroup_size_32, 32,  32, 16, s_warptile_wm, 32, 2, tm_s, tn_s, tk_s, subgroup_size_8 };
+        s_warptile = { subgroup_size_32, 32,  32, 16, s_warptile_wm_8, 32, 2, tm_s, tn_s, tk_s, subgroup_size_8 };
 
         l_warptile_mmq = { 128,             128, 128, 32, mm_warp_8 * 2, 64, 2, tm_l, tn_l, tk_l, mm_warp_8 };
         m_warptile_mmq = { 128,              64,  64, 32, mm_warp_8,     32, 2, tm_m, tn_m, tk_m, mm_warp_8 };
-        s_warptile_mmq = { subgroup_size_32, 32,  32, 32, s_warptile_wm, 32, 2, tm_s, tn_s, tk_s, subgroup_size_8 };
+        s_warptile_mmq = { subgroup_size_32, 32,  32, 32, s_warptile_wm_8, 32, 2, tm_s, tn_s, tk_s, subgroup_size_8 };
 
         // Integer MMQ has a smaller shared memory profile, but heavier register use
         l_warptile_mmq_int = { 128,             128, 128, 32, mm_warp_8 * 2, 64, 2, 4, 4, 1, mm_warp_8 };
         m_warptile_mmq_int = { 128,              64,  64, 32, mm_warp_8,     32, 2, 2, 2, 1, mm_warp_8 };
-        s_warptile_mmq_int = { subgroup_size_32, 32,  32, 32, s_warptile_wm, 32, 2, 2, 1, 1, subgroup_size_8 };
+        s_warptile_mmq_int = { subgroup_size_32, 32,  32, 32, s_warptile_wm_8, 32, 2, 2, 1, 1, subgroup_size_8 };
 
         const auto cm1_bs = [cm1_sg](uint32_t bm, uint32_t bn) {
             return cm1_sg * (bm / std::min(cm1_sg, bm)) * (bn / 32);
@@ -1818,23 +1821,23 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
         // K-quants use even more registers, mitigate by setting WMITER to 1
         l_warptile_mmq_int_k = { 128,               128, 128, 32, mm_warp_8 * 2, 64, 1, 4, 4, 1, mm_warp_8 };
         m_warptile_mmq_int_k = { 128,                64,  64, 32, mm_warp_8,     32, 1, 2, 2, 1, mm_warp_8 };
-        s_warptile_mmq_int_k = { subgroup_size_32,   32,  32, 32, s_warptile_wm, 32, 1, 2, 1, 1, subgroup_size_8 };
+        s_warptile_mmq_int_k = { subgroup_size_32,   32,  32, 32, s_warptile_wm_8, 32, 1, 2, 1, 1, subgroup_size_8 };
 
         l_warptile_id = { 128,                      128, 128, 16, mul_mat_mm_warp_16 * 2, 64, 2, tm_l, tn_l, tk_l, mul_mat_mm_warp_16 };
         m_warptile_id = { 128,                       64,  64, 16, mul_mat_mm_warp_16,     32, 2, tm_m, tn_m, tk_m, mul_mat_mm_warp_16 };
-        s_warptile_id = { mul_mat_subgroup_size_16,  32,  32, 16, s_warptile_wm,          32, 2, tm_s, tn_s, tk_s, mul_mat_subgroup_size_16 };
+        s_warptile_id = { mul_mat_subgroup_size_16,  32,  32, 16, 32,                     32, 2, tm_s, tn_s, tk_s, mul_mat_subgroup_size_16 };
 
         l_warptile_mmqid = { 128,                       128, 128, 32, mul_mat_mm_warp_8 * 2, 64, 2, tm_l, tn_l, tk_l, mul_mat_mm_warp_8 };
         m_warptile_mmqid = { 128,                        64,  64, 32, mul_mat_mm_warp_8,     32, 2, tm_m, tn_m, tk_m, mul_mat_mm_warp_8 };
-        s_warptile_mmqid = { mul_mat_subgroup_size_32,   32,  32, 32, s_warptile_wm,         32, 2, tm_s, tn_s, tk_s, mul_mat_subgroup_size_8 };
+        s_warptile_mmqid = { mul_mat_subgroup_size_32,   32,  32, 32, mul_mat_s_warptile_wm_8, 32, 2, tm_s, tn_s, tk_s, mul_mat_subgroup_size_8 };
 
         l_warptile_mmqid_int = { 128,                       128, 128, 32, mul_mat_mm_warp_8 * 2, 64, 2, 4, 4, 1, mul_mat_mm_warp_8 };
         m_warptile_mmqid_int = { 128,                        64,  64, 32, mul_mat_mm_warp_8,     32, 2, 2, 2, 1, mul_mat_mm_warp_8 };
-        s_warptile_mmqid_int = { mul_mat_subgroup_size_32,   32,  32, 32, s_warptile_wm,         32, 2, 2, 1, 1, mul_mat_subgroup_size_8 };
+        s_warptile_mmqid_int = { mul_mat_subgroup_size_32,   32,  32, 32, mul_mat_s_warptile_wm_8, 32, 2, 2, 1, 1, mul_mat_subgroup_size_8 };
 
         l_warptile_mmqid_int_k = { 128,                     128, 128, 32, mul_mat_mm_warp_16 * 2, 64, 1, 4, 4, 1, mul_mat_mm_warp_16 };
         m_warptile_mmqid_int_k = { 128,                      64,  64, 32, mul_mat_mm_warp_16,     32, 1, 2, 2, 1, mul_mat_mm_warp_16 };
-        s_warptile_mmqid_int_k = { mul_mat_subgroup_size_32, 32,  32, 32, s_warptile_wm,          32, 1, 2, 1, 1, mul_mat_subgroup_size_16 };
+        s_warptile_mmqid_int_k = { mul_mat_subgroup_size_32, 32,  32, 32, mul_mat_s_warptile_wm_16, 32, 1, 2, 1, 1, mul_mat_subgroup_size_16 };
 
         // chip specific tuning
         if ((device->architecture == AMD_GCN) && (device->driver_id != vk::DriverId::eAmdProprietary)) {
@@ -2793,9 +2796,8 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
         && !device->coopmat_bf16_support
 #endif
         ) {
-        const uint32_t s_warptile_wm_bf16 = device->subgroup_size == 8 ? 8 : 32;
         std::vector<vk_tile_config> tc_bf16_fb = {
-            {{ subgroup_size_32, 32, 32, 16, s_warptile_wm_bf16, 32, 2, 2, 2, 1, subgroup_size_8 }, {32, 32, 1}, s_align},
+            {{ subgroup_size_32, 32, 32, 16, s_warptile_wm_8, 32, 2, 2, 2, 1, subgroup_size_8 }, {32, 32, 1}, s_align},
             {{ 128, 64, 64, 16, mm_warp_8, 32, 2, 4, 2, 1, mm_warp_8 }, {64, 64, 1}, m_align},
             {{ 128, 128, 128, 16, mm_warp_8 * 2, 64, 2, 4, 4, 1, mm_warp_8 }, {128, 128, 1}, l_align},
         };
@@ -3808,7 +3810,7 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
 
         auto shmem_req = [&](uint32_t pad, bool csh_store, bool fp16_shmem) {
             const uint32_t elem_size = fp16_shmem ? (uint32_t)sizeof(uint16_t) : (uint32_t)sizeof(float);
-            const uint32_t csh_elems = csh_store ? conv2d_BS.K * conv2d_BS.NPQ : 0u;
+            const uint32_t csh_elems = csh_store ? conv2d_BS.K * conv2d_BS.NPQ : (fp16_shmem ? 8u : 0u);
             return (conv2d_BS.K * (conv2d_BS.CRS + pad) + conv2d_BS.CRS * (conv2d_BS.NPQ + pad) + csh_elems) * elem_size;
         };
 
@@ -3836,9 +3838,9 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
         }
 
         // stage cm2 accumulator through shmem for coalesced global stores;
-        // skipped on 128x128 where the extra Csh footprint hurts occupancy.
+        // skipped on large tiles where the extra Csh footprint hurts occupancy.
         // cm1 always uses the staged path.
-        uint32_t conv2d_csh_store = (device->coopmat2 && s != CONV_SHAPE_128x128) ? 1u : 0u;
+        uint32_t conv2d_csh_store = (device->coopmat2 && s != CONV_SHAPE_128x128 && s != CONV_SHAPE_256x128) ? 1u : 0u;
         if (conv2d_use_cm1) {
             conv2d_csh_store = 1;
         }
@@ -8718,6 +8720,26 @@ static vk_conv_shapes ggml_vk_conv_select_shape(ggml_backend_vk_context * ctx, u
     }
 }
 
+static vk_conv_shapes ggml_vk_conv3d_select_shape(ggml_backend_vk_context * ctx, uint32_t K, uint32_t NPQ) {
+    const auto & block = vk_conv_block_sizes[CONV_SHAPE_256x128];
+    const uint32_t shader_core_count = ctx->device->shader_core_count;
+    // Slowdowns were observed on DGX Spark with drivers older than R615.
+    const bool driver_supported = ctx->device->driver_id != vk::DriverId::eNvidiaProprietary ||
+                                  (ctx->device->properties.driverVersion >> 22) >= 615;
+    if (ctx->device->coopmat2 && shader_core_count > 0 && K % block.K == 0 && driver_supported) {
+        const uint64_t n_tiles = (uint64_t)(K / block.K) * CEIL_DIV((uint64_t)NPQ, block.NPQ);
+        // Use the largest cm2 padding and include the dummy Csh[8].
+        const uint32_t pad = 8, csh_elems = 8;
+        const uint32_t shmem = (block.K * (block.CRS + pad) + block.CRS * (block.NPQ + pad) + csh_elems) * sizeof(uint16_t);
+        // Reuse each loaded input across 256 output channels instead of 128.
+        // Require full channel tiles and at least as many workgroups as GPU cores.
+        if (n_tiles >= shader_core_count && shmem <= ctx->device->properties.limits.maxComputeSharedMemorySize) {
+            return CONV_SHAPE_256x128;
+        }
+    }
+    return ggml_vk_conv_select_shape(ctx, K, NPQ);
+}
+
 static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * dst, ggml_op op) {
     switch (op) {
     case GGML_OP_GET_ROWS:
@@ -9387,7 +9409,7 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
             const uint32_t IC = (uint32_t)ggml_get_op_params_i32(dst, 9);
             const uint32_t N  = (uint32_t)ggml_get_op_params_i32(dst, 10);
             const uint32_t NPQ = N * dst->ne[2] * dst->ne[1] * dst->ne[0];
-            const vk_conv_shapes shape = ggml_vk_conv_select_shape(ctx, OC, NPQ);
+            const vk_conv_shapes shape = ggml_vk_conv3d_select_shape(ctx, OC, NPQ);
 
             const uint32_t KW = (uint32_t)src0->ne[0];
             const uint32_t KH = (uint32_t)src0->ne[1];
@@ -9700,7 +9722,7 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
     case GGML_OP_CONV_3D:
         if constexpr (std::is_same_v<PC, vk_op_conv3d_push_constants>) {
             const uint32_t NPQ = pc.N * pc.OD * pc.OH * pc.OW;
-            const vk_conv_shapes shape = ggml_vk_conv_select_shape(ctx, pc.OC, NPQ);
+            const vk_conv_shapes shape = ggml_vk_conv3d_select_shape(ctx, pc.OC, NPQ);
             const uint32_t NPQ_blocks = CEIL_DIV(NPQ, vk_conv_block_sizes[shape].NPQ);
 
             elements = { pc.OC, NPQ_blocks, 1 };

gguf-py/gguf/constants.py

@@ -211,6 +211,7 @@ class Keys:
         VALUE_SCALE                  = "{arch}.attention.value_scale"
         COMPRESS_RATIOS              = "{arch}.attention.compress_ratios"
         COMPRESS_ROPE_FREQ_BASE      = "{arch}.attention.compress_rope_freq_base"
+        KV_SOURCE_LAYERS             = "{arch}.attention.kv_source_layers"
         TEMPERATURE_LENGTH           = "{arch}.attention.temperature_length"
         KEY_LENGTH_MLA               = "{arch}.attention.key_length_mla"
         VALUE_LENGTH_MLA             = "{arch}.attention.value_length_mla"
@@ -236,6 +237,10 @@ class Keys:
             TYPES      = "{arch}.attention.indexer.types"
             KPOOL             = "{arch}.attention.indexer.kpool"              # GLM5-Next
             KPOOL_SELECT_TAIL = "{arch}.attention.indexer.kpool_select_tail"  # GLM5-Next
+            SOURCE_LAYERS          = "{arch}.attention.indexer.source_layers"
+            CANDIDATE_SOURCE_LAYER = "{arch}.attention.indexer.candidate_source_layer"
+            CANDIDATE_TOP_K_BLOCKS = "{arch}.attention.indexer.candidate_top_k_blocks"
+            CANDIDATE_BLOCK_SIZE   = "{arch}.attention.indexer.candidate_block_size"
 
     class HyperConnection:
         COUNT                = "{arch}.hyper_connection.count"
@@ -245,6 +250,12 @@ class Keys:
         MAGNITUDE            = "{arch}.hyper_connection.magnitude"
         # absent means the mix projection is full rank (DeepSeek-V4 behaviour)
         LOW_RANK             = "{arch}.hyper_connection.low_rank"
+        # single-pass mHC: each sublayer collapses with the previous pre-mix (DeepSeek-V4.1)
+        SINGLE_PASS          = "{arch}.hyper_connection.single_pass"
+
+    class Engram:  # deepseek41; the table uses the PerLayerEmbedding keys
+        PAD_TOKEN_ID          = "{arch}.engram.pad_token_id"
+        TOKEN_MAP             = "{arch}.engram.token_map"
 
     class PerLayerEmbedding:
         LAYERS             = "{arch}.ple.layers"
@@ -636,6 +647,7 @@ class MODEL_ARCH(IntEnum):
     DEEPSEEK2OCR     = auto()
     DEEPSEEK32       = auto()
     DEEPSEEK4        = auto()
+    DEEPSEEK41       = auto()
     CHATGLM          = auto()
     GLM4             = auto()
     GLM4_MOE         = auto()
@@ -1317,6 +1329,8 @@ class MODEL_TENSOR(IntEnum):
     DSPARK_MARKOV_W1       = auto()  # markov head: prev-token embed
     DSPARK_MARKOV_W2       = auto()  # markov head: bias projection
     DSPARK_CONF_PROJ       = auto()  # confidence head
+    ENGRAM_KV              = auto()
+    ENGRAM_GATE            = auto()
     DFLASH_ATTN_CONV_BASE  = auto()
     DFLASH_ATTN_CONV_PROJ  = auto()
     DFLASH_FFN_CONV_BASE   = auto()
@@ -1460,6 +1474,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
     MODEL_ARCH.DEEPSEEK2OCR:     "deepseek2-ocr",
     MODEL_ARCH.DEEPSEEK32:       "deepseek32",
     MODEL_ARCH.DEEPSEEK4:        "deepseek4",
+    MODEL_ARCH.DEEPSEEK41:       "deepseek41",
     MODEL_ARCH.CHATGLM:          "chatglm",
     MODEL_ARCH.GLM4:             "glm4",
     MODEL_ARCH.GLM4_MOE:         "glm4moe",
@@ -2164,6 +2179,8 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
     MODEL_TENSOR.DSPARK_MARKOV_W1:          "markov_w1",
     MODEL_TENSOR.DSPARK_MARKOV_W2:          "markov_w2",
     MODEL_TENSOR.DSPARK_CONF_PROJ:          "conf_proj",
+    MODEL_TENSOR.ENGRAM_KV:                 "blk.{bid}.engram_kv",
+    MODEL_TENSOR.ENGRAM_GATE:               "blk.{bid}.engram_gate",
     MODEL_TENSOR.DFLASH_ATTN_CONV_BASE:     "blk.{bid}.attn_conv_base",
     MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ:     "blk.{bid}.attn_conv_proj",
     MODEL_TENSOR.DFLASH_FFN_CONV_BASE:      "blk.{bid}.ffn_conv_base",
@@ -4358,6 +4375,46 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
         MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
     ],
+    MODEL_ARCH.DEEPSEEK41: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_SINKS,
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_KV,
+        MODEL_TENSOR.ATTN_KV_NORM,
+        MODEL_TENSOR.ATTN_OUT_A,
+        MODEL_TENSOR.ATTN_OUT_B,
+        MODEL_TENSOR.HC_ATTN_FN,
+        MODEL_TENSOR.HC_ATTN_BASE,
+        MODEL_TENSOR.HC_ATTN_SCALE,
+        MODEL_TENSOR.HC_FFN_FN,
+        MODEL_TENSOR.HC_FFN_BASE,
+        MODEL_TENSOR.HC_FFN_SCALE,
+        MODEL_TENSOR.ATTN_COMPRESSOR_WKV,
+        MODEL_TENSOR.ATTN_COMPRESSOR_WGATE,
+        MODEL_TENSOR.ATTN_COMPRESSOR_NORM,
+        MODEL_TENSOR.INDEXER_ATTN_K,
+        MODEL_TENSOR.INDEXER_K_NORM,
+        MODEL_TENSOR.INDEXER_ATTN_Q_B,
+        MODEL_TENSOR.INDEXER_PROJ,
+        MODEL_TENSOR.PER_LAYER_TOKEN_EMBD,
+        MODEL_TENSOR.ENGRAM_KV,
+        MODEL_TENSOR.ENGRAM_GATE,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_EXP_PROBS_B_VL,
+        MODEL_TENSOR.FFN_NORM,
+        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_ARCH.ERNIE4_5_MOE: [
         MODEL_TENSOR.TOKEN_EMBD,
         MODEL_TENSOR.OUTPUT_NORM,
@@ -6449,6 +6506,7 @@ class VisionProjectorType:
     DEEPSEEKOCR = "deepseekocr"
     DEEPSEEKOCR2 = "deepseekocr2"
     DEEPSEEK4V = "deepseek4v"
+    DEEPSEEK41V = "deepseek41v"
     LFM2A = "lfm2a" # audio
     MUSIC_FLAMINGO = "musicflamingo" # audio
     GLM4V = "glm4v"

gguf-py/gguf/gguf_writer.py

@@ -1077,6 +1077,28 @@ class GGUFWriter:
     def add_attention_compress_rope_freq_base(self, value: float) -> None:
         self.add_float32(Keys.Attention.COMPRESS_ROPE_FREQ_BASE.format(arch=self.arch), value)
 
+    # deepseek41
+    def add_attention_kv_source_layers(self, layers: Sequence[int]) -> None:
+        self.add_key_value(Keys.Attention.KV_SOURCE_LAYERS.format(arch=self.arch), list(layers), GGUFValueType.ARRAY, GGUFValueType.UINT32)
+
+    def add_indexer_source_layers(self, layers: Sequence[int]) -> None:
+        self.add_key_value(Keys.Attention.Indexer.SOURCE_LAYERS.format(arch=self.arch), list(layers), GGUFValueType.ARRAY, GGUFValueType.UINT32)
+
+    def add_indexer_candidate_source_layer(self, layer: int) -> None:
+        self.add_int32(Keys.Attention.Indexer.CANDIDATE_SOURCE_LAYER.format(arch=self.arch), layer)
+
+    def add_indexer_candidate_top_k_blocks(self, value: int) -> None:
+        self.add_uint32(Keys.Attention.Indexer.CANDIDATE_TOP_K_BLOCKS.format(arch=self.arch), value)
+
+    def add_indexer_candidate_block_size(self, value: int) -> None:
+        self.add_uint32(Keys.Attention.Indexer.CANDIDATE_BLOCK_SIZE.format(arch=self.arch), value)
+
+    def add_engram_pad_token_id(self, value: int) -> None:
+        self.add_uint32(Keys.Engram.PAD_TOKEN_ID.format(arch=self.arch), value)
+
+    def add_engram_token_map(self, values: Sequence[int]) -> None:
+        self.add_key_value(Keys.Engram.TOKEN_MAP.format(arch=self.arch), list(values), GGUFValueType.ARRAY, GGUFValueType.INT32)
+
     def add_hyper_connection_count(self, count: int) -> None:
         self.add_uint32(Keys.HyperConnection.COUNT.format(arch=self.arch), count)
 
@@ -1092,6 +1114,9 @@ class GGUFWriter:
     def add_hyper_connection_low_rank(self, value: int) -> None:
         self.add_uint32(Keys.HyperConnection.LOW_RANK.format(arch=self.arch), value)
 
+    def add_hyper_connection_single_pass(self, value: bool) -> None:
+        self.add_bool(Keys.HyperConnection.SINGLE_PASS.format(arch=self.arch), value)
+
     def add_ple_layers(self, values: Sequence[int]) -> None:
         self.add_array(Keys.PerLayerEmbedding.LAYERS.format(arch=self.arch), values)
 

scripts/snapdragon/ggml-hexagon-profile.py

@@ -215,6 +215,7 @@ def parse_log(file_path, pmu_index=None, limit=None, device_filter=None, op_filt
                 'unwrapped_cycles_start': unwrapped_cycles_start,
                 'pmu_val':      pmu_val,
                 'evt_val':      evt_val,
+                'mhz':          float(op_match.group('mhz')) if op_match.group('mhz') else None,
                 'abs_usec':     abs_usec,
                 'trace_events': [],
                 'device':       device
@@ -630,6 +631,165 @@ def generate_report(ops, top_n, width_overrides, sort_col, pmu_name=None):
         logger.info("| " + " | ".join(row_vals) + " |")
 
 
+def has_perf_print(file_path):
+    if file_path == "-":
+        return True
+    with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
+        return any("common_perf_print" in line for line in f)
+
+
+def pct_delta(a, b):
+    return (b - a) * 100.0 / a if a else None
+
+
+def fmt_pct(p):
+    return f"{p:+.2f}%" if p is not None else "n/a"
+
+
+def print_table(title, cols, rows, width_overrides):
+    # cols: list of (cli name, header); rows: list of dicts keyed by cli name
+    logger.info(f"\n# {title}\n")
+    if not rows:
+        logger.info("None.")
+        return
+
+    headers, keys, widths = [], [], []
+    for key, header in cols:
+        natural_width = max([len(str(row[key])) for row in rows] + [len(header)])
+        target_width  = width_overrides.get(key, natural_width)
+        if target_width == 0:
+            continue
+        headers.append(header)
+        keys.append(key)
+        widths.append(target_width)
+
+    logger.info("| " + " | ".join(f"{h:<{widths[i]}}" for i, h in enumerate(headers)) + " |")
+    logger.info("| " + " | ".join("-" * widths[i] for i in range(len(headers))) + " |")
+    for row in rows:
+        row_vals = []
+        for i, key in enumerate(keys):
+            val = str(row[key])
+            if len(val) > widths[i]:
+                val = val[:widths[i] - 3] + "..."
+            row_vals.append(f"{val:<{widths[i]}}")
+        logger.info("| " + " | ".join(row_vals) + " |")
+
+
+def generate_compare_report(ops_a, ops_b, file_a, file_b, top_n, width_overrides, min_delta):
+    def summarize(ops):
+        batches = [o for o in ops if o['name'] == "OPBATCH"]
+        others  = [o for o in ops if o['name'] != "OPBATCH"]
+        mhz     = [o['mhz'] for o in others if o['mhz'] is not None]
+        return {
+            'n_ops':     len(others),
+            'tot_usec':  sum(o['usec'] for o in others),
+            'n_batch':   len(batches),
+            'batch_usec': sum(o['usec'] for o in batches),
+            'avg_mhz':   statistics.mean(mhz) if mhz else 0.0,
+        }
+
+    sa, sb = summarize(ops_a), summarize(ops_b)
+
+    logger.info("\n# Compare Summary (OPBATCH excluded from op totals)\n")
+    logger.info(f"  A (baseline) : {file_a}")
+    logger.info(f"  B (new)      : {file_b}")
+    logger.info("  Deltas are B - A: negative means B is faster\n")
+    logger.info(f"  Ops          : A {sa['n_ops']}  B {sb['n_ops']}")
+    logger.info(f"  Tot op usec  : A {sa['tot_usec']}  B {sb['tot_usec']}  delta {sb['tot_usec'] - sa['tot_usec']:+d} ({fmt_pct(pct_delta(sa['tot_usec'], sb['tot_usec']))})")
+    logger.info(f"  OPBATCH      : A {sa['n_batch']} / {sa['batch_usec']} usec  B {sb['n_batch']} / {sb['batch_usec']} usec")
+    logger.info(f"  Avg op mhz   : A {sa['avg_mhz']:.1f}  B {sb['avg_mhz']:.1f}")
+
+    if sa['n_batch'] != sb['n_batch']:
+        logger.warning(f"  WARNING: OPBATCH counts differ (A {sa['n_batch']}, B {sb['n_batch']}), one log may be incomplete")
+    for label, path in (("A", file_a), ("B", file_b)):
+        if not has_perf_print(path):
+            logger.warning(f"  WARNING: log {label} has no common_perf_print line, run may be incomplete")
+
+    # Per-op-name rollup
+    by_name = defaultdict(lambda: [0, 0, 0, 0])
+    for idx, ops in ((0, ops_a), (1, ops_b)):
+        for op in ops:
+            if op['name'] == "OPBATCH":
+                continue
+            by_name[op['name']][idx] += 1
+            by_name[op['name']][idx + 2] += op['usec']
+
+    rows = []
+    for name, (cnt_a, cnt_b, tot_a, tot_b) in by_name.items():
+        p = pct_delta(tot_a, tot_b)
+        if min_delta and p is not None and abs(p) < min_delta:
+            continue
+        rows.append({'op': name, 'count-a': cnt_a, 'count-b': cnt_b, 'tot-usec-a': tot_a, 'tot-usec-b': tot_b,
+                     'delta-usec': f"{tot_b - tot_a:+d}", 'delta-pct': fmt_pct(p), '_sort': abs(tot_b - tot_a)})
+    rows.sort(key=lambda r: r['_sort'], reverse=True)
+
+    print_table("Per-Op Rollup (delta = B - A, sorted by abs delta usec)", [
+        ("op", "Op"), ("count-a", "Count A"), ("count-b", "Count B"), ("tot-usec-a", "Tot usec A"),
+        ("tot-usec-b", "Tot usec B"), ("delta-usec", "Delta usec"), ("delta-pct", "Delta %")], rows, width_overrides)
+
+    # Per-group compare keyed on (op, dims, dtypes)
+    def group(ops):
+        grouped = defaultdict(list)
+        for op in ops:
+            if op['name'] != "OPBATCH":
+                grouped[(op['name'], op['dims'], op['types'])].append(op['usec'])
+        return grouped
+
+    ga, gb = group(ops_a), group(ops_b)
+
+    rows = []
+    for key in ga.keys() & gb.keys():
+        ua, ub = ga[key], gb[key]
+        tot_a, tot_b = sum(ua), sum(ub)
+        avg_a, avg_b = tot_a / len(ua), tot_b / len(ub)
+        p = pct_delta(avg_a, avg_b)
+        if min_delta and p is not None and abs(p) < min_delta:
+            continue
+        rows.append({'op': key[0], 'dims': key[1], 'dtypes': key[2], 'count-a': len(ua), 'count-b': len(ub),
+                     'avg-usec-a': f"{avg_a:.2f}", 'avg-usec-b': f"{avg_b:.2f}", 'delta-pct': fmt_pct(p),
+                     'tot-usec-a': tot_a, 'tot-usec-b': tot_b, 'delta-usec': f"{tot_b - tot_a:+d}", '_sort': abs(tot_b - tot_a)})
+    rows.sort(key=lambda r: r['_sort'], reverse=True)
+
+    print_table(f"Per-Group Compare (delta = B - A, top {top_n} sorted by abs delta usec)", [
+        ("op", "Op"), ("dims", "Dims"), ("dtypes", "DTypes"), ("count-a", "Count A"), ("count-b", "Count B"),
+        ("avg-usec-a", "Avg usec A"), ("avg-usec-b", "Avg usec B"), ("delta-pct", "Delta %"),
+        ("tot-usec-a", "Tot usec A"), ("tot-usec-b", "Tot usec B"), ("delta-usec", "Delta usec")], rows[:top_n], width_overrides)
+
+    # Groups present in only one log
+    for label, g, other in (("A", ga, gb), ("B", gb, ga)):
+        rows = [{'op': k[0], 'dims': k[1], 'dtypes': k[2], 'count': len(v), 'tot-usec': sum(v)} for k, v in g.items() if k not in other]
+        rows.sort(key=lambda r: r['tot-usec'], reverse=True)
+        print_table(f"Groups only in {label} (Top {top_n} sorted by tot usec)", [
+            ("op", "Op"), ("dims", "Dims"), ("dtypes", "DTypes"), ("count", "Count"), ("tot-usec", "Tot usec")], rows[:top_n], width_overrides)
+
+
+def filter_ops(ops, args):
+    if args.filter:
+        try:
+            filter_re = re.compile(args.filter)
+        except re.error as e:
+            logger.error(f"Invalid regex filter: {e}")
+            sys.exit(1)
+        ops = [op for op in ops if filter_re.search(op['op_text'])]
+
+    if args.head is not None or args.tail is not None:
+        ops_by_dev = defaultdict(list)
+        for op in ops:
+            ops_by_dev[op['device']].append(op)
+
+        filtered_ops = []
+        for dev in sorted(ops_by_dev.keys()):
+            dev_ops = ops_by_dev[dev]
+            if args.head is not None:
+                dev_ops = dev_ops[:args.head]
+            elif args.tail is not None:
+                dev_ops = dev_ops[-args.tail:]
+            filtered_ops.extend(dev_ops)
+        ops = filtered_ops
+
+    return ops
+
+
 def main():
     parser = argparse.ArgumentParser(description="Post-process Op profile info.")
     parser.add_argument("logfile")
@@ -642,6 +802,8 @@ def main():
                         help="Output ASCII art event summary or thread idle bubble analysis (default: summary)")
     parser.add_argument("--filter", type=str, help="Regex filter matching against the original profile-op line")
     parser.add_argument("--device", type=str, help="Device to filter by (e.g. HTP0, HTP0:0) or 'split' to generate separate reports per device")
+    parser.add_argument("--compare", type=str, metavar="LOGFILE", help="Compare against another log (logfile is baseline A, this is B)")
+    parser.add_argument("--min-delta", type=float, default=0, metavar="PCT", help="In compare mode hide rows with abs(Delta %%) below PCT")
 
     group = parser.add_mutually_exclusive_group()
     group.add_argument("--head", type=int, help="Limit to first N ops")
@@ -655,6 +817,10 @@ def main():
         logger.error(f"Cannot sort by '{args.sort}' without --pmu-index.")
         sys.exit(1)
 
+    if args.compare and (args.timeline or args.device == "split"):
+        logger.error("--compare cannot be used with --timeline or --device split.")
+        sys.exit(1)
+
     overrides = {}
     if args.width:
         for w in args.width:
@@ -681,6 +847,13 @@ def main():
     if args.device and args.device != "split":
         ops = [op for op in ops if device_matches(op['device'], args.device)]
 
+    if args.compare:
+        ops_b = parse_log(args.compare, pmu_index=args.pmu_index, limit=limit, device_filter=device_filter, op_filter_re=op_filter_re)
+        if device_filter:
+            ops_b = [op for op in ops_b if device_matches(op['device'], device_filter)]
+        generate_compare_report(filter_ops(ops, args), filter_ops(ops_b, args), args.logfile, args.compare, args.top, overrides, args.min_delta)
+        return
+
     if args.device == "split":
         unique_devices = sorted(list(set(op['device'] for op in ops)))
         for dev in unique_devices:
@@ -712,28 +885,7 @@ def main():
             else:
                 generate_report(dev_ops, args.top, overrides, args.sort, pmu_name=final_pmu_name)
     else:
-        if args.filter:
-            try:
-                filter_re = re.compile(args.filter)
-            except re.error as e:
-                logger.error(f"Invalid regex filter: {e}")
-                sys.exit(1)
-            ops = [op for op in ops if filter_re.search(op['op_text'])]
-
-        if args.head is not None or args.tail is not None:
-            ops_by_dev = defaultdict(list)
-            for op in ops:
-                ops_by_dev[op['device']].append(op)
-
-            filtered_ops = []
-            for dev in sorted(ops_by_dev.keys()):
-                dev_ops = ops_by_dev[dev]
-                if args.head is not None:
-                    dev_ops = dev_ops[:args.head]
-                elif args.tail is not None:
-                    dev_ops = dev_ops[-args.tail:]
-                filtered_ops.extend(dev_ops)
-            ops = filtered_ops
+        ops = filter_ops(ops, args)
 
         if args.timeline:
             for op in ops:

src/llama-arch.cpp

@@ -87,6 +87,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
     { LLM_ARCH_DEEPSEEK2OCR,     "deepseek2-ocr"    },
     { LLM_ARCH_DEEPSEEK32,       "deepseek32"       },
     { LLM_ARCH_DEEPSEEK4,        "deepseek4"        },
+    { LLM_ARCH_DEEPSEEK41,       "deepseek41"       },
     { LLM_ARCH_CHATGLM,          "chatglm"          },
     { LLM_ARCH_GLM4,             "glm4"             },
     { LLM_ARCH_GLM4_MOE,         "glm4moe"          },
@@ -311,6 +312,11 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
     { LLM_KV_ATTENTION_INDEXER_TYPES,                "%s.attention.indexer.types"                },
     { LLM_KV_ATTENTION_INDEXER_KPOOL,                "%s.attention.indexer.kpool"                },
     { LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL,    "%s.attention.indexer.kpool_select_tail"    },
+    { LLM_KV_ATTENTION_INDEXER_SOURCE_LAYERS,          "%s.attention.indexer.source_layers"          },
+    { LLM_KV_ATTENTION_INDEXER_CANDIDATE_SOURCE_LAYER, "%s.attention.indexer.candidate_source_layer" },
+    { LLM_KV_ATTENTION_INDEXER_CANDIDATE_TOP_K_BLOCKS, "%s.attention.indexer.candidate_top_k_blocks" },
+    { LLM_KV_ATTENTION_INDEXER_CANDIDATE_BLOCK_SIZE,   "%s.attention.indexer.candidate_block_size"   },
+    { LLM_KV_ATTENTION_KV_SOURCE_LAYERS,               "%s.attention.kv_source_layers"               },
     { LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,           "%s.attention.output_group_count"           },
     { LLM_KV_ATTENTION_OUTPUT_LORA_RANK,             "%s.attention.output_lora_rank"             },
     { LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE,      "%s.attention.compress_rope_freq_base"      },
@@ -323,6 +329,10 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
     { LLM_KV_HYPER_CONNECTION_EPSILON,               "%s.hyper_connection.epsilon"               },
     { LLM_KV_HYPER_CONNECTION_MAGNITUDE,             "%s.hyper_connection.magnitude"             },
     { LLM_KV_HYPER_CONNECTION_LOW_RANK,              "%s.hyper_connection.low_rank"              },
+    { LLM_KV_HYPER_CONNECTION_SINGLE_PASS,           "%s.hyper_connection.single_pass"           },
+
+    { LLM_KV_ENGRAM_PAD_TOKEN_ID,                    "%s.engram.pad_token_id"                    },
+    { LLM_KV_ENGRAM_TOKEN_MAP,                       "%s.engram.token_map"                       },
 
     { LLM_KV_PLE_LAYERS,                             "%s.ple.layers"                             },
     { LLM_KV_PLE_NGRAM_SIZE,                         "%s.ple.ngram_size"                         },
@@ -782,6 +792,8 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
     { LLM_TENSOR_DSPARK_MARKOV_W1,                       "markov_w1" },
     { LLM_TENSOR_DSPARK_MARKOV_W2,                       "markov_w2" },
     { LLM_TENSOR_DSPARK_CONF_PROJ,                       "conf_proj" },
+    { LLM_TENSOR_ENGRAM_KV,                              "blk.%d.engram_kv" },
+    { LLM_TENSOR_ENGRAM_GATE,                            "blk.%d.engram_gate" },
     { LLM_TENSOR_DFLASH_ATTN_CONV_BASE,                  "blk.%d.attn_conv_base" },
     { LLM_TENSOR_DFLASH_ATTN_CONV_PROJ,                  "blk.%d.attn_conv_proj" },
     { LLM_TENSOR_DFLASH_FFN_CONV_BASE,                   "blk.%d.ffn_conv_base" },
@@ -1120,6 +1132,8 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
     {LLM_TENSOR_DSPARK_MARKOV_W1,           {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_GET_ROWS}},
     {LLM_TENSOR_DSPARK_MARKOV_W2,           {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_MUL_MAT}},
     {LLM_TENSOR_DSPARK_CONF_PROJ,           {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_ENGRAM_KV,                  {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_ENGRAM_GATE,                {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
     {LLM_TENSOR_DFLASH_ATTN_CONV_BASE,      {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
     {LLM_TENSOR_DFLASH_ATTN_CONV_PROJ,      {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_DFLASH_FFN_CONV_BASE,       {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
@@ -1239,6 +1253,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) {
         case LLM_ARCH_QWEN35MOE:
         case LLM_ARCH_QWEN4EXP:
         case LLM_ARCH_DEEPSEEK4:
+        case LLM_ARCH_DEEPSEEK41:
         case LLM_ARCH_MINIMAX_01:
             return true;
         default:
@@ -1265,6 +1280,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) {
         case LLM_ARCH_QWEN35MOE:
         case LLM_ARCH_QWEN4EXP:
         case LLM_ARCH_DEEPSEEK4:
+        case LLM_ARCH_DEEPSEEK41:
         case LLM_ARCH_GLM5_NEXT:
         case LLM_ARCH_NEMOTRON_H:
         case LLM_ARCH_NEMOTRON_H_MOE:
@@ -1282,6 +1298,7 @@ bool llm_arch_supports_mixed_batch(const llm_arch & arch) {
     switch (arch) {
         case LLM_ARCH_COGVLM:
         case LLM_ARCH_DEEPSEEK4:
+        case LLM_ARCH_DEEPSEEK41:
         case LLM_ARCH_GRANITE_SWITCH:
         case LLM_ARCH_EAGLE3:
         case LLM_ARCH_DFLASH:

src/llama-arch.h

@@ -92,6 +92,7 @@ enum llm_arch {
     LLM_ARCH_DEEPSEEK2OCR,
     LLM_ARCH_DEEPSEEK32,
     LLM_ARCH_DEEPSEEK4,
+    LLM_ARCH_DEEPSEEK41,
     LLM_ARCH_CHATGLM,
     LLM_ARCH_GLM4,
     LLM_ARCH_GLM4_MOE,
@@ -316,6 +317,11 @@ enum llm_kv {
     LLM_KV_ATTENTION_INDEXER_TYPES,
     LLM_KV_ATTENTION_INDEXER_KPOOL,
     LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL,
+    LLM_KV_ATTENTION_INDEXER_SOURCE_LAYERS,
+    LLM_KV_ATTENTION_INDEXER_CANDIDATE_SOURCE_LAYER,
+    LLM_KV_ATTENTION_INDEXER_CANDIDATE_TOP_K_BLOCKS,
+    LLM_KV_ATTENTION_INDEXER_CANDIDATE_BLOCK_SIZE,
+    LLM_KV_ATTENTION_KV_SOURCE_LAYERS,
     LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,
     LLM_KV_ATTENTION_OUTPUT_LORA_RANK,
     LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE,
@@ -328,6 +334,10 @@ enum llm_kv {
     LLM_KV_HYPER_CONNECTION_EPSILON,
     LLM_KV_HYPER_CONNECTION_MAGNITUDE,
     LLM_KV_HYPER_CONNECTION_LOW_RANK,
+    LLM_KV_HYPER_CONNECTION_SINGLE_PASS,
+
+    LLM_KV_ENGRAM_PAD_TOKEN_ID,
+    LLM_KV_ENGRAM_TOKEN_MAP,
 
     LLM_KV_PLE_LAYERS,
     LLM_KV_PLE_NGRAM_SIZE,
@@ -788,6 +798,8 @@ enum llm_tensor {
     LLM_TENSOR_DSPARK_MARKOV_W1,
     LLM_TENSOR_DSPARK_MARKOV_W2,
     LLM_TENSOR_DSPARK_CONF_PROJ,
+    LLM_TENSOR_ENGRAM_KV,
+    LLM_TENSOR_ENGRAM_GATE,
     LLM_TENSOR_DFLASH_ATTN_CONV_BASE,
     LLM_TENSOR_DFLASH_ATTN_CONV_PROJ,
     LLM_TENSOR_DFLASH_FFN_CONV_BASE,

src/llama-context.cpp

@@ -1339,8 +1339,9 @@ void llama_context::set_abort_callback(bool (*abort_callback)(void * data), void
 void llama_context::set_embeddings(bool value) {
     LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value);
 
-    // a pooled context computes the logits only without embeddings, so the toggle changes the graph shape
-    if (cparams.embeddings != value && cparams.pooling_type != LLAMA_POOLING_TYPE_NONE) {
+    // the toggle changes the graph shape (pooled: logits only; pooling NONE: the reserve graph outputs all tokens,
+    // which crashed ggml_gallocr on the first /completion), so re-reserve on every toggle
+    if (cparams.embeddings != value) {
         sched_need_reserve = true;
     }
 
@@ -2536,6 +2537,7 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
         model.arch == LLM_ARCH_CLEF ||
         model.arch == LLM_ARCH_QWEN4EXP ||
         model.arch == LLM_ARCH_DEEPSEEK4 ||
+        model.arch == LLM_ARCH_DEEPSEEK41 ||
         (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) ||
         model.arch == LLM_ARCH_NANBEIGE ||
         model.arch == LLM_ARCH_MINIMAX_01 ||
@@ -4071,7 +4073,7 @@ llama_context * llama_init_from_model(
         }
     }
 
-    if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4) && params.type_k != params.type_v) {
+    if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4 || model->arch == LLM_ARCH_DEEPSEEK41) && params.type_k != params.type_v) {
         LLAMA_LOG_ERROR("%s: model does not support different K (%s) and V (%s) cache types\n", __func__, ggml_type_name(params.type_k), ggml_type_name(params.type_v));
         return nullptr;
     }

src/llama-graph.cpp

@@ -2010,7 +2010,7 @@ ggml_tensor * llm_graph_context::build_ffn(
                     const float limit = hparams.swiglu_clamp_shexp[il];
                     constexpr float eps = 1e-6f;
                     if (limit > eps) {
-                        if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) {
+                        if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) {
                             cur = ggml_swiglu_clamp(ctx0, cur, tmp, limit);
                         } else {
                             tmp = ggml_clamp(ctx0, tmp, -limit, limit);
@@ -2340,8 +2340,13 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
         ggml_tensor * weights_sum = ggml_sum_rows(ctx0, weights); // [1, n_tokens]
         cb(weights_sum, "ffn_moe_weights_sum", il);
 
-        // Avoid division by zero, clamp to smallest number representable by F16
-        weights_sum = ggml_clamp(ctx0, weights_sum, 6.103515625e-5, INFINITY);
+        if (arch == LLM_ARCH_DEEPSEEK41) {
+            // model.py Gate.forward: weights / (sum + 1e-20); the sum can fall below the F16 clamp below
+            weights_sum = ggml_scale_bias(ctx0, weights_sum, 1.0f, 1e-20f);
+        } else {
+            // Avoid division by zero, clamp to smallest number representable by F16
+            weights_sum = ggml_clamp(ctx0, weights_sum, 6.103515625e-5, INFINITY);
+        }
         cb(weights_sum, "ffn_moe_weights_sum_clamped", il);
 
         weights = ggml_div(ctx0, weights, weights_sum); // [n_expert_used, n_tokens]
@@ -2431,7 +2436,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
                     const float limit = hparams.swiglu_clamp_exp[il];
                     constexpr float eps = 1e-6f;
                     if (limit > eps) {
-                        if (arch == LLM_ARCH_MAPLE || arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0) || arch == LLM_ARCH_HY_V4) {
+                        if (arch == LLM_ARCH_MAPLE || arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0) || arch == LLM_ARCH_HY_V4) {
                             cur = ggml_swiglu_clamp(ctx0, cur, up, limit);
                         } else {
                             up = ggml_clamp(ctx0, up, -limit, limit);

src/llama-hparams.h

@@ -10,8 +10,8 @@
 // bump if necessary
 #define LLAMA_MAX_LAYERS  512
 #define LLAMA_MAX_EXPERTS 1024 // Kimi K3
-#define LLAMA_MAX_PLE_NGRAM 8  // qwen4exp
-#define LLAMA_MAX_PLE_HEADS 64 // qwen4exp
+#define LLAMA_MAX_PLE_NGRAM 8  // qwen4exp, deepseek41 (all hash modules)
+#define LLAMA_MAX_PLE_HEADS 64 // qwen4exp, deepseek41 (all hash modules)
 
 enum llama_expert_gating_func_type {
     LLAMA_EXPERT_GATING_FUNC_TYPE_NONE           = 0,
@@ -61,7 +61,7 @@ struct llama_hparams {
     //       for example: n_embd_out() vs n_embd_out_impl
 
     bool vocab_only;
-    bool no_alloc;
+    bool no_alloc; // TODO: no_alloc is not a model hyperparameter; move it out of llama_hparams
     bool rope_finetuned;
     bool use_par_res;
     bool swin_norm;
@@ -326,12 +326,20 @@ struct llama_hparams {
     uint32_t dsv4_o_group_count        = 0;
     uint32_t dsv4_o_lora_rank          = 0;
     uint32_t dsv4_hc_mult              = 0;
+    bool     dsv4_hc_single_pass       = false; // DeepSeek-V4.1 DSpark draft: lagged pre-mix, no hc_head
     uint32_t dsv4_hc_sinkhorn_iters    = 0;
     uint32_t dsv4_hash_layer_count     = 0;
     float    dsv4_compress_rope_base   = 0.0f;
     float    dsv4_hc_eps               = 0.0f;
     std::array<uint32_t, LLAMA_MAX_LAYERS> dsv4_compress_ratios;
 
+    // deepseek41
+    std::array<int32_t, LLAMA_MAX_LAYERS> dsv41_kv_src;   // layer whose compressed KV this layer reads, -1 = none
+    std::array<int32_t, LLAMA_MAX_LAYERS> dsv41_topk_src; // layer whose top-k this layer uses, -1 = none
+    int32_t  dsv41_cand_src         = -1;    // -1 = no candidate block selection
+    uint32_t dsv41_cand_topk_blocks = 0;
+    uint32_t dsv41_cand_block_size  = 0;
+
     // 0 = full rank (DeepSeek-V4)
     uint32_t hc_low_rank = 0;
 
@@ -346,6 +354,10 @@ struct llama_hparams {
     uint32_t ple_eos_token_id    = 0;
     // the id the PLE hash stands in at image positions; 0 makes the loader fall back to EOS
     uint32_t ple_image_token_id  = 0;
+    // the id the KV cells record at embedding (image) positions for the n-gram history; LLAMA_TOKEN_NULL = none
+    llama_token ple_embd_token   = LLAMA_TOKEN_NULL;
+    // hash constant sets, stored back to back: 1 shared by all PLE layers (qwen4exp), or 1 per PLE layer in layer order (deepseek41)
+    uint32_t ple_n_modules       = 1;
     // the file lists PLE layer indices, so this is never a per-layer gguf array and can hold one bit per layer
     std::bitset<LLAMA_MAX_LAYERS> is_ple_impl;
     // the hash multipliers reach ~2e13 and have to stay 64-bit

src/llama-kv-cache-dsv4.cpp

@@ -333,12 +333,24 @@ static void dsv4_state_read_tensor_streams(
     }
 }
 
+// serialized compressed-KV layers: all except per-device mirrors (kv_src[il] != il), refilled from their source on read
+static std::vector<uint32_t> dsv4_k_cache_state_layers(const llama_kv_cache * kv, const std::vector<int32_t> & kv_src) {
+    std::vector<uint32_t> res;
+    for (uint32_t il : kv->get_layer_ids()) {
+        if (kv_src[il] == (int32_t) il) {
+            res.push_back(il);
+        }
+    }
+    return res;
+}
+
 static void dsv4_state_write_k_cache(
         llama_io_write_i    & io,
         const llama_kv_cache * kv,
         llama_seq_id          seq_id,
         llama_state_seq_flags flags,
-        uint32_t              n_rows) {
+        uint32_t              n_rows,
+        const std::vector<int32_t> & kv_src) {
     GGML_UNUSED(flags);
 
     uint32_t s0;
@@ -347,7 +359,8 @@ static void dsv4_state_write_k_cache(
 
     const uint32_t version = DSV4_K_CACHE_STATE_VER;
     const uint32_t kv_size = kv->get_size();
-    const auto layer_ids = kv->get_layer_ids();
+    // mirrors are never serialized, so the state format does not depend on the split
+    const auto layer_ids = dsv4_k_cache_state_layers(kv, kv_src);
     const uint32_t n_layer = layer_ids.size();
 
     if (n_rows > kv_size) {
@@ -369,7 +382,8 @@ static void dsv4_state_read_k_cache(
         llama_io_read_i  & io,
         llama_kv_cache   * kv,
         llama_seq_id       seq_id,
-        llama_state_seq_flags flags) {
+        llama_state_seq_flags flags,
+        const std::vector<int32_t> & kv_src) {
     GGML_UNUSED(flags);
 
     uint32_t version;
@@ -399,7 +413,7 @@ static void dsv4_state_read_k_cache(
     uint32_t s0;
     dsv4_state_dst_stream_range(kv->get_n_stream(), seq_id, ns, s0);
 
-    const auto layer_ids = kv->get_layer_ids();
+    const auto layer_ids = dsv4_k_cache_state_layers(kv, kv_src);
     if (n_layer_ref != layer_ids.size()) {
         throw std::runtime_error("DSV4 K-cache layer count mismatch");
     }
@@ -413,6 +427,31 @@ static void dsv4_state_read_k_cache(
 
         dsv4_state_read_tensor_streams(io, kv->get_k_storage(il), kv_size, n_rows_ref, s0, ns);
     }
+
+    // refill the mirrors from their source rows just read
+    std::vector<uint8_t> buf;
+    for (uint32_t il : kv->get_layer_ids()) {
+        const int32_t src = kv_src[il];
+        if (src == (int32_t) il) {
+            continue;
+        }
+        ggml_tensor * t_src = kv->get_k_storage(src);
+        ggml_tensor * t_dst = kv->get_k_storage(il);
+        GGML_ASSERT(t_src->type == t_dst->type && t_src->ne[0] == t_dst->ne[0]);
+
+        const size_t row_size      = ggml_row_size(t_src->type, t_src->ne[0]);
+        const size_t stream_stride = (size_t) kv_size*row_size;
+        const size_t size          = (size_t) n_rows_ref*row_size;
+        if (size == 0) {
+            continue;
+        }
+        buf.resize(size);
+        for (uint32_t s = 0; s < ns; ++s) {
+            const size_t offset = (size_t) (s0 + s)*stream_stride;
+            ggml_backend_tensor_get(t_src, buf.data(), offset, size);
+            ggml_backend_tensor_set(t_dst, buf.data(), offset, size);
+        }
+    }
 }
 
 static std::string dsv4_plan_positions(const std::vector<int32_t> & values) {
@@ -432,6 +471,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
         const llama_ubatch & ubatch,
         uint32_t ratio,
         bool overlap,
+        bool pad_blocks,
         uint32_t state_size,
         uint32_t kv_size,
         uint32_t n_stream,
@@ -548,7 +588,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
         }
     }
 
-    if (ratio == DSV4_CSA_RATIO && !plan.state_pos.empty()) {
+    if (pad_blocks && !plan.state_pos.empty()) {
         assert(kv_size > 0);
 
         // Pad each stream to the reserve plan's block count.
@@ -603,7 +643,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
         }
     }
 
-    if (ratio == DSV4_HCA_RATIO && !plan.state_pos.empty() && plan.state_write_idxs.empty()) {
+    if (!pad_blocks && !plan.state_pos.empty() && plan.state_write_idxs.empty()) {
         assert(kv_size > 0);
         // the last slot must not be live, or the dummy write would corrupt it;
         // a full stream implies a completed block, which implies real writes
@@ -744,6 +784,7 @@ static std::vector<llama_kv_cache_dsv4_context::comp_plan> dsv4_build_comp_plans
         const std::vector<llama_ubatch> & ubatches,
         uint32_t ratio,
         bool overlap,
+        bool pad_blocks,
         uint32_t state_size,
         uint32_t kv_size,
         uint32_t n_stream,
@@ -755,7 +796,7 @@ static std::vector<llama_kv_cache_dsv4_context::comp_plan> dsv4_build_comp_plans
     // the first ubatch touching a seq consumes its rollback restore
     std::vector<uint32_t> rs(rs_idx);
     for (const llama_ubatch & ubatch : ubatches) {
-        plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream, n_rs_seq, rs));
+        plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, pad_blocks, state_size, kv_size, n_stream, n_rs_seq, rs));
 
         for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
             const llama_seq_id seq_id = ubatch.seq_id_unq[s];
@@ -905,13 +946,12 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
             uint32_t        n_rs_seq,
         const char    * name,
         const llama_memory_i::layer_filter_cb & filter) :
+    hparams(model.hparams),
     ratio(ratio),
     state_size(state_size),
     n_embd_state(n_embd_state),
     n_stream(unified ? 1 : n_seq_max),
     n_rs_seq(n_rs_seq) {
-    const llama_hparams & hparams = model.hparams;
-
     struct ggml_backend_buft_comparator {
         bool operator()(const ggml_backend_buffer_type_t & lhs, const ggml_backend_buffer_type_t & rhs) const {
             return strcmp(ggml_backend_buft_name(lhs), ggml_backend_buft_name(rhs)) < 0;
@@ -986,7 +1026,15 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
     }
 
     for (auto & [buft, ctx] : ctx_map) {
-        ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft);
+        ggml_backend_buffer_t buf;
+        if (hparams.no_alloc) {
+            buf = ggml_backend_buft_alloc_buffer(buft, /*size =*/ 0); // dummy buffer
+            for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) {
+                t->buffer = buf; // set dummy buffer for DSV4 state so that the backend scheduler won't try to allocate it
+            }
+        } else {
+            buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); // real buffer
+        }
         if (!buf) {
             throw std::runtime_error("failed to allocate buffer for DSV4 compressor state");
         }
@@ -1056,6 +1104,10 @@ uint32_t llama_dsv4_comp_state::get_ratio() const {
     return ratio;
 }
 
+bool llama_dsv4_comp_state::is_overlap() const {
+    return state_size == 2*ratio;
+}
+
 uint32_t llama_dsv4_comp_state::get_state_size() const {
     return state_size;
 }
@@ -1074,9 +1126,15 @@ uint32_t llama_dsv4_comp_state::get_n_rows() const {
 
 std::map<ggml_backend_buffer_type_t, size_t> llama_dsv4_comp_state::memory_breakdown() const {
     std::map<ggml_backend_buffer_type_t, size_t> ret;
-    for (const auto & [_, buf] : ctxs_bufs) {
+    for (const auto & [ctx, buf] : ctxs_bufs) {
         ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get());
-        ret[buft] += ggml_backend_buffer_get_size(buf.get());
+
+        if (hparams.no_alloc) {
+            GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) == nullptr);
+            ret[buft] += ggml_backend_alloc_ctx_tensors_from_buft_size(ctx.get(), buft);
+        } else {
+            ret[buft] += ggml_backend_buffer_get_size(buf.get());
+        }
     }
     return ret;
 }
@@ -1261,6 +1319,23 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
             v_trans, offload, swa_full, unified_raw, kv_size, n_seq_max, n_ubatch, n_pad,
             nullptr, filter_raw, reuse, nullptr);
 
+    // V4.1: CSA slot = ratio 2, HCA slot = ratio 1, no overlap and no indexer compressor; the index K is stored
+    // next to the compressed KV ([kv | index K]) and only the KV source layers own rows
+    const bool is_v41 = model.arch == LLM_ARCH_DEEPSEEK41;
+    const uint32_t ratio_csa = is_v41 ? 2u : DSV4_CSA_RATIO;
+    const uint32_t ratio_hca = is_v41 ? 1u : DSV4_HCA_RATIO;
+
+    if (is_v41) {
+        const uint32_t n_embd_comp = model.hparams.n_embd_head_k() + model.hparams.indexer_head_size;
+        for (llama_hparams * hp : { &hparams_csa, &hparams_hca }) {
+            std::fill(hp->n_head_kv_arr.begin(), hp->n_head_kv_arr.end(), 1);
+            hp->n_embd_head_k_full = n_embd_comp;
+            hp->n_embd_head_v_full = n_embd_comp;
+            hp->n_embd_head_k_swa  = n_embd_comp;
+            hp->n_embd_head_v_swa  = n_embd_comp;
+        }
+    }
+
     dsv4_make_k_only(hparams_csa);
     dsv4_make_k_only(hparams_hca);
 
@@ -1272,12 +1347,50 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
     hparams_lid.rope_type          = LLAMA_ROPE_TYPE_NEOX;
     dsv4_make_k_only(hparams_lid);
 
+    // V4.1 cross-layer KV sharing: a compressed layer reads the storage of its kv source or, when the source is on
+    // another device (layer split), of a per-device mirror fed with the source's new rows, so the whole cache never
+    // crosses devices. no mirrors without offload
+    const uint32_t n_layer = model.hparams.n_layer();
+    comp_read.assign(n_layer, -1);
+    comp_kv_src.resize(n_layer);
+    for (uint32_t il = 0; il < n_layer; ++il) {
+        const int32_t src = is_v41 ? model.hparams.dsv41_kv_src[il] : (int32_t) il;
+        comp_kv_src[il] = src >= 0 ? src : (int32_t) il;
+        if (model.hparams.dsv4_compress_ratios[il] == 0) {
+            continue;
+        }
+        comp_read[il] = src;
+        if (src == (int32_t) il || !offload || model.dev_layer(src) == model.dev_layer(il)) {
+            continue;
+        }
+        // the first layer of this source on this device holds the mirror
+        comp_read[il] = il;
+        for (int32_t j = src + 1; j < (int32_t) il; ++j) {
+            if (comp_read[j] == j && comp_kv_src[j] == src && model.dev_layer(j) == model.dev_layer(il)) {
+                comp_read[il] = j;
+                break;
+            }
+        }
+    }
+
+    // compressed-KV storage: kv sources and mirrors; compressor state: kv sources only
+    const auto is_comp_owner = [&](int32_t il, uint32_t ratio) {
+        return model.hparams.dsv4_compress_ratios[il] == ratio && comp_read[il] == (int32_t) il;
+    };
+
+    const auto is_state_owner = [&](int32_t il, uint32_t ratio) {
+        if (model.hparams.dsv4_compress_ratios[il] != ratio) {
+            return false;
+        }
+        return !is_v41 || model.hparams.dsv41_kv_src[il] == il;
+    };
+
     const layer_filter_cb filter_csa = [&](int32_t il) {
         if (filter && !filter(il)) {
             return false;
         }
 
-        return model.hparams.dsv4_compress_ratios[il] == DSV4_CSA_RATIO;
+        return is_comp_owner(il, ratio_csa);
     };
 
     const layer_filter_cb filter_hca = [&](int32_t il) {
@@ -1285,52 +1398,67 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
             return false;
         }
 
-        return model.hparams.dsv4_compress_ratios[il] == DSV4_HCA_RATIO;
+        return is_comp_owner(il, ratio_hca);
+    };
+
+    const layer_filter_cb filter_csa_state = [&](int32_t il) {
+        return (!filter || filter(il)) && is_state_owner(il, ratio_csa);
+    };
+
+    const layer_filter_cb filter_hca_state = [&](int32_t il) {
+        return (!filter || filter(il)) && is_state_owner(il, ratio_hca);
+    };
+
+    const layer_filter_cb filter_lid = [&](int32_t il) {
+        return !is_v41 && filter_csa(il);
     };
 
     const bool unified_compressed = false;
 
     LLAMA_LOG_INFO("%s: creating DSV4 CSA compressed KV cache, size = %u cells\n",
-            __func__, dsv4_comp_size(kv_size, DSV4_CSA_RATIO));
+            __func__, dsv4_comp_size(kv_size, ratio_csa));
 
     kv_csa = std::make_unique<llama_kv_cache>(
             model, hparams_csa, type_k, type_v,
-            v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, DSV4_CSA_RATIO), 256u), n_seq_max, n_pad,
+            v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, ratio_csa), 256u), n_seq_max, n_pad,
             0, LLAMA_SWA_TYPE_NONE, nullptr, filter_csa, nullptr, nullptr);
 
     LLAMA_LOG_INFO("%s: creating DSV4 HCA compressed KV cache, size = %u cells\n",
-            __func__, dsv4_comp_size(kv_size, DSV4_HCA_RATIO));
+            __func__, dsv4_comp_size(kv_size, ratio_hca));
 
     kv_hca = std::make_unique<llama_kv_cache>(
             model, hparams_hca, type_k, type_v,
-            v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, DSV4_HCA_RATIO), 256u), n_seq_max, n_pad,
+            v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, ratio_hca), 256u), n_seq_max, n_pad,
             0, LLAMA_SWA_TYPE_NONE, nullptr, filter_hca, nullptr, nullptr);
 
     LLAMA_LOG_INFO("%s: creating DSV4 lightning-indexer KV cache, size = %u cells\n",
-            __func__, dsv4_comp_size(kv_size, DSV4_CSA_RATIO));
+            __func__, dsv4_comp_size(kv_size, ratio_csa));
 
     kv_lid = std::make_unique<llama_kv_cache>(
             model, hparams_lid, type_k, type_v,
-            v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, DSV4_CSA_RATIO), 256u), n_seq_max, n_pad,
-            0, LLAMA_SWA_TYPE_NONE, nullptr, filter_csa, nullptr, nullptr);
+            v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, ratio_csa), 256u), n_seq_max, n_pad,
+            0, LLAMA_SWA_TYPE_NONE, nullptr, filter_lid, nullptr, nullptr);
 
     LLAMA_LOG_INFO("%s: creating DSV4 CSA compressor state\n", __func__);
 
+    const uint32_t csa_state_size = is_v41 ? ratio_csa : 2*ratio_csa;
+    const uint32_t csa_state_embd = is_v41 ? model.hparams.n_embd_head_k() : 2*model.hparams.n_embd_head_k();
+
     csa_state = std::make_unique<llama_dsv4_comp_state>(
-            model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO,
-            2*model.hparams.n_embd_head_k(), n_rs_seq, "csa", filter_csa);
+            model, offload, unified_compressed, n_seq_max, ratio_csa, csa_state_size,
+            csa_state_embd, n_rs_seq, "csa", filter_csa_state);
 
     LLAMA_LOG_INFO("%s: creating DSV4 HCA compressor state\n", __func__);
 
     hca_state = std::make_unique<llama_dsv4_comp_state>(
-            model, offload, unified_compressed, n_seq_max, DSV4_HCA_RATIO, DSV4_HCA_RATIO,
-            model.hparams.n_embd_head_k(), n_rs_seq, "hca", filter_hca);
+            model, offload, unified_compressed, n_seq_max, ratio_hca, ratio_hca,
+            model.hparams.n_embd_head_k(), n_rs_seq, "hca", filter_hca_state);
 
     LLAMA_LOG_INFO("%s: creating DSV4 lightning-indexer compressor state\n", __func__);
 
     lid_state = std::make_unique<llama_dsv4_comp_state>(
             model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO,
-            2*model.hparams.indexer_head_size, n_rs_seq, "lid", filter_csa);
+            2*model.hparams.indexer_head_size, n_rs_seq, "lid", filter_lid);
 
     // DSV4 attention reads compressed-K / compressor-state rows that the current
     // graph does not necessarily overwrite; uninitialized buffer contents would
@@ -1339,12 +1467,76 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
     clear_compressed(-1, true);
 }
 
+// Engram hashes the n-gram ending at each token, so a token shared by several sequences (llama-batched prompt)
+// is well-defined only if all of them agree on its predecessors; returns false (and logs) otherwise
+static bool dsv41_shared_history_ok(const llama_batch & batch, const llama_kv_cache * kv, uint32_t n_hist) {
+    if (n_hist == 0 || !batch.token || !batch.n_seq_id) {
+        return true;
+    }
+
+    bool any_shared = false;
+    for (int32_t i = 0; i < batch.n_tokens && !any_shared; ++i) {
+        any_shared = batch.n_seq_id[i] > 1;
+    }
+    if (!any_shared) {
+        return true;
+    }
+
+    std::map<std::pair<llama_seq_id, llama_pos>, llama_token> in_batch;
+    for (int32_t i = 0; i < batch.n_tokens; ++i) {
+        for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) {
+            in_batch[{ batch.seq_id[i][s], batch.pos[i] }] = batch.token[i];
+        }
+    }
+
+    const auto tok_at = [&](llama_seq_id seq_id, llama_pos p) -> llama_token {
+        // batch first: the batch continues the cache
+        auto it = in_batch.upper_bound({ seq_id, p });
+        if (it != in_batch.begin()) {
+            --it;
+            if (it->first.first == seq_id) {
+                return it->second;
+            }
+        }
+        return kv ? kv->get_cells(seq_id).seq_pos_tok_le(seq_id, p) : LLAMA_TOKEN_NULL;
+    };
+
+    for (int32_t i = 0; i < batch.n_tokens; ++i) {
+        if (batch.n_seq_id[i] <= 1) {
+            continue;
+        }
+        for (uint32_t d = 1; d <= n_hist; ++d) {
+            const llama_pos p = batch.pos[i] - (llama_pos) d;
+            if (p < 0) {
+                break;
+            }
+            const llama_token t0 = tok_at(batch.seq_id[i][0], p);
+            for (int32_t s = 1; s < batch.n_seq_id[i]; ++s) {
+                const llama_token ts = tok_at(batch.seq_id[i][s], p);
+                if (ts != t0) {
+                    LLAMA_LOG_ERROR("%s: token %d at pos %d is shared by seq %d and seq %d, whose Engram n-gram histories "
+                            "differ at pos %d (%d vs %d); decode such tokens per sequence\n", __func__,
+                            i, batch.pos[i], batch.seq_id[i][0], batch.seq_id[i][s], p, t0, ts);
+                    return false;
+                }
+            }
+        }
+    }
+
+    return true;
+}
+
 llama_memory_context_ptr llama_kv_cache_dsv4::init_batch(
             llama_batch_allocr & balloc,
             uint32_t n_ubatch,
             bool embd_all) {
     GGML_UNUSED(embd_all);
 
+    if (hparams_raw.ple_n_heads > 0 &&
+        !dsv41_shared_history_ok(balloc.get_batch(), kv_raw->get_swa(), hparams_raw.ple_ngram_size - 1)) {
+        return std::make_unique<llama_kv_cache_dsv4_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
+    }
+
     const bool raw_per_seq  = kv_raw->get_base()->get_n_stream() != 1;
     const bool comp_per_seq = csa_state->get_n_stream() > 1;
     const bool has_coupled = dsv4_batch_has_coupled(balloc.get_batch());
@@ -1475,9 +1667,9 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
             bool res = true;
 
             res = res & kv_raw->seq_rm(seq_id, p0, -1);
-            res = res & kv_csa->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
-            res = res & kv_hca->seq_rm(seq_id, p0/DSV4_HCA_RATIO, -1);
-            res = res & kv_lid->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
+            res = res & kv_csa->seq_rm(seq_id, p0/csa_state->get_ratio(), -1);
+            res = res & kv_hca->seq_rm(seq_id, p0/hca_state->get_ratio(), -1);
+            res = res & kv_lid->seq_rm(seq_id, p0/lid_state->get_ratio(), -1);
 
             return res;
         }
@@ -1613,15 +1805,15 @@ void llama_kv_cache_dsv4::state_write(llama_io_write_i & io, llama_seq_id seq_id
 
         //FIXME : note that we conflate token positions with rows, which is not true for multi-modal case.
         const uint32_t n_rows_csa = seq_id >= 0 ?
-            dsv4_state_n_used_k_rows(pos_max, DSV4_CSA_RATIO, kv_csa->get_size()) : kv_csa->get_size();
+            dsv4_state_n_used_k_rows(pos_max, csa_state->get_ratio(), kv_csa->get_size()) : kv_csa->get_size();
         const uint32_t n_rows_hca = seq_id >= 0 ?
-            dsv4_state_n_used_k_rows(pos_max, DSV4_HCA_RATIO, kv_hca->get_size()) : kv_hca->get_size();
+            dsv4_state_n_used_k_rows(pos_max, hca_state->get_ratio(), kv_hca->get_size()) : kv_hca->get_size();
         const uint32_t n_rows_lid = seq_id >= 0 ?
-            dsv4_state_n_used_k_rows(pos_max, DSV4_CSA_RATIO, kv_lid->get_size()) : kv_lid->get_size();
+            dsv4_state_n_used_k_rows(pos_max, lid_state->get_ratio(), kv_lid->get_size()) : kv_lid->get_size();
 
-        dsv4_state_write_k_cache(io, kv_csa.get(), seq_id, flags, n_rows_csa);
-        dsv4_state_write_k_cache(io, kv_hca.get(), seq_id, flags, n_rows_hca);
-        dsv4_state_write_k_cache(io, kv_lid.get(), seq_id, flags, n_rows_lid);
+        dsv4_state_write_k_cache(io, kv_csa.get(), seq_id, flags, n_rows_csa, comp_kv_src);
+        dsv4_state_write_k_cache(io, kv_hca.get(), seq_id, flags, n_rows_hca, comp_kv_src);
+        dsv4_state_write_k_cache(io, kv_lid.get(), seq_id, flags, n_rows_lid, comp_kv_src);
     }
 
     csa_state->state_write(io, seq_id, flags, rs_idx);
@@ -1659,9 +1851,9 @@ void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id,
     if (!partial_only) {
         clear_compressed(seq_id, true);
 
-        dsv4_state_read_k_cache(io, kv_csa.get(), seq_id, flags);
-        dsv4_state_read_k_cache(io, kv_hca.get(), seq_id, flags);
-        dsv4_state_read_k_cache(io, kv_lid.get(), seq_id, flags);
+        dsv4_state_read_k_cache(io, kv_csa.get(), seq_id, flags, comp_kv_src);
+        dsv4_state_read_k_cache(io, kv_hca.get(), seq_id, flags, comp_kv_src);
+        dsv4_state_read_k_cache(io, kv_lid.get(), seq_id, flags, comp_kv_src);
     }
 
     csa_state->state_read(io, seq_id, flags);
@@ -1704,6 +1896,10 @@ llama_dsv4_comp_state * llama_kv_cache_dsv4::get_lid_state() const {
     return lid_state.get();
 }
 
+int32_t llama_kv_cache_dsv4::get_comp_read_layer(int32_t il) const {
+    return comp_read[il];
+}
+
 uint32_t llama_kv_cache_dsv4::get_n_rs_seq() const {
     return n_rs_seq;
 }
@@ -1932,6 +2128,11 @@ void llama_kv_cache_dsv4_raw_context::set_input_kq_mask(ggml_tensor * dst, const
     kv_swa->set_input_kq_mask(dst, ubatch, causal_attn);
 }
 
+void llama_kv_cache_dsv4_raw_context::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const {
+    GGML_ASSERT(kv_swa != nullptr);
+    kv_swa->get_prev_tokens(ubatch, n, res);
+}
+
 void llama_kv_cache_dsv4_raw_context::set_input_k_rot(ggml_tensor * dst) const {
     kv_swa->set_input_k_rot(dst);
 }
@@ -2016,7 +2217,8 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
     reserve_plans(true),
     status(llama_memory_status_combine(
                 llama_memory_status_combine(ctx_raw->get_status(), ctx_csa_mem->get_status()),
-                llama_memory_status_combine(ctx_hca_mem->get_status(), ctx_lid_mem->get_status()))) {
+                llama_memory_status_combine(ctx_hca_mem->get_status(), ctx_lid_mem->get_status()))),
+    kv(kv) {
 }
 
 llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
@@ -2041,7 +2243,8 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
                     llama_memory_status_combine(ctx_raw->get_status(), ctx_csa_mem->get_status()),
                     llama_memory_status_combine(ctx_hca_mem->get_status(), ctx_lid_mem->get_status())),
                 this->sc_info_csa.empty() && this->sc_info_hca.empty() && this->sc_info_lid.empty() ?
-                    LLAMA_MEMORY_STATUS_NO_UPDATE : LLAMA_MEMORY_STATUS_SUCCESS)) {
+                    LLAMA_MEMORY_STATUS_NO_UPDATE : LLAMA_MEMORY_STATUS_SUCCESS)),
+    kv(kv) {
 }
 
 llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
@@ -2052,13 +2255,13 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
         std::vector<llama_ubatch> ubatches,
         std::vector<llama_ubatch> ubatches_raw) :
     ubatches(std::move(ubatches)),
-    plans_csa(dsv4_build_comp_plans(this->ubatches, DSV4_CSA_RATIO, true,
+    plans_csa(dsv4_build_comp_plans(this->ubatches, kv->get_csa_state()->get_ratio(), kv->get_csa_state()->is_overlap(), true,
                 kv->get_csa_state()->get_state_size(), kv->get_csa()->get_size(), kv->get_csa_state()->get_n_stream(),
                 kv->get_n_rs_seq(), kv->get_rs_idx())),
-    plans_hca(dsv4_build_comp_plans(this->ubatches, DSV4_HCA_RATIO, false,
+    plans_hca(dsv4_build_comp_plans(this->ubatches, kv->get_hca_state()->get_ratio(), kv->get_hca_state()->is_overlap(), false,
                 kv->get_hca_state()->get_state_size(), kv->get_hca()->get_size(), kv->get_hca_state()->get_n_stream(),
                 kv->get_n_rs_seq(), kv->get_rs_idx())),
-    plans_lid(dsv4_build_comp_plans(this->ubatches, DSV4_CSA_RATIO, true,
+    plans_lid(dsv4_build_comp_plans(this->ubatches, kv->get_lid_state()->get_ratio(), kv->get_lid_state()->is_overlap(), true,
                 kv->get_lid_state()->get_state_size(), kv->get_lid()->get_size(), kv->get_lid_state()->get_n_stream(),
                 kv->get_n_rs_seq(), kv->get_rs_idx())),
     ctx_raw(std::make_unique<llama_kv_cache_dsv4_raw_context>(
@@ -2086,7 +2289,8 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
     csa_state(kv->get_csa_state()),
     hca_state(kv->get_hca_state()),
     lid_state(kv->get_lid_state()),
-    status(ctx_raw->get_status()) {
+    status(ctx_raw->get_status()),
+    kv(kv) {
     kv->reset_rs_idx_for_ubatches(this->ubatches);
 }
 
@@ -2181,6 +2385,10 @@ const llama_dsv4_comp_state * llama_kv_cache_dsv4_context::get_lid_state() const
     return lid_state;
 }
 
+int32_t llama_kv_cache_dsv4_context::get_comp_read_layer(int32_t il) const {
+    return kv->get_comp_read_layer(il);
+}
+
 const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_csa_plan() const {
     assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
 
@@ -2222,7 +2430,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_
     }
 
     reserve_plan_csa = dsv4_build_reserve_comp_plan(
-            ubatch, DSV4_CSA_RATIO, true,
+            ubatch, csa_state->get_ratio(), csa_state->is_overlap(),
             csa_state->get_state_size(), get_csa()->get_n_kv(), csa_state->get_n_stream(), csa_state->get_n_rs_seq());
 
     return reserve_plan_csa;
@@ -2236,7 +2444,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_
     }
 
     reserve_plan_hca = dsv4_build_reserve_comp_plan(
-            ubatch, DSV4_HCA_RATIO, false,
+            ubatch, hca_state->get_ratio(), hca_state->is_overlap(),
             hca_state->get_state_size(), get_hca()->get_n_kv(), hca_state->get_n_stream(), hca_state->get_n_rs_seq());
 
     return reserve_plan_hca;
@@ -2250,7 +2458,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_
     }
 
     reserve_plan_lid = dsv4_build_reserve_comp_plan(
-            ubatch, DSV4_CSA_RATIO, true,
+            ubatch, lid_state->get_ratio(), lid_state->is_overlap(),
             lid_state->get_state_size(), get_lid()->get_n_kv(), lid_state->get_n_stream(), lid_state->get_n_rs_seq());
 
     return reserve_plan_lid;

src/llama-kv-cache-dsv4.h

@@ -31,6 +31,7 @@ public:
     void apply_copies(const stream_copy_info & sc_info) const;
 
     uint32_t get_ratio()      const;
+    bool     is_overlap()     const; // V4 CSA: state holds 2*ratio rows of 2x width
     uint32_t get_state_size() const;
     uint32_t get_n_stream()   const;
     uint32_t get_n_rs_seq()   const;
@@ -60,6 +61,8 @@ private:
         std::vector<ggml_tensor *> score_stream;
     };
 
+    const llama_hparams & hparams;
+
     const uint32_t ratio;
     const uint32_t state_size;
     const uint32_t n_embd_state;
@@ -148,6 +151,9 @@ public:
     llama_dsv4_comp_state * get_hca_state() const;
     llama_dsv4_comp_state * get_lid_state() const;
 
+    // layer whose compressed-KV storage `il` reads, -1 = no compressed KV
+    int32_t get_comp_read_layer(int32_t il) const;
+
     uint32_t get_n_rs_seq() const;
     const std::vector<uint32_t> & get_rs_idx() const;
     void reset_rs_idx_for_ubatches(const std::vector<llama_ubatch> & ubatches);
@@ -163,6 +169,12 @@ private:
 
     std::vector<uint32_t> rs_idx;
 
+    // [n_layer] storage layer -> layer whose rows it holds (itself, or the source of a mirror)
+    std::vector<int32_t> comp_kv_src;
+
+    // [n_layer] layer -> layer whose compressed-KV storage it reads (itself, its V4.1 kv source or a mirror), -1 = none
+    std::vector<int32_t> comp_read;
+
     std::unique_ptr<llama_kv_cache_iswa> kv_raw;
     std::unique_ptr<llama_kv_cache>      kv_csa;
     std::unique_ptr<llama_kv_cache>      kv_hca;
@@ -214,6 +226,9 @@ public:
     void set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const;
     void set_input_k_rot(ggml_tensor * dst) const;
 
+    // token history of the raw (window) cells, for host-side n-gram hashing (deepseek41 Engram)
+    void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const;
+
 private:
     size_t i_next = 0;
 
@@ -371,6 +386,8 @@ public:
     const comp_plan & get_hca_plan(const llama_ubatch & ubatch) const;
     const comp_plan & get_lid_plan(const llama_ubatch & ubatch) const;
 
+    int32_t get_comp_read_layer(int32_t il) const;
+
 private:
     size_t i_next = 0;
 
@@ -403,4 +420,6 @@ private:
     mutable comp_plan reserve_plan_lid;
 
     const llama_memory_status status;
+
+    const llama_kv_cache_dsv4 * kv = nullptr;
 };

src/llama-kv-cache.cpp

@@ -1168,9 +1168,7 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
                 } else if (hparams.ple_n_heads > 0) {
                     // embd batch (multimodal input) has no token ids, need to pad it with the correct ID for PLE layers
                     // TODO @ngxson : check if we can do the same as gemma 3n / gemma 4
-                    ext.tok = hparams.ple_image_token_id != 0
-                        ? (llama_token) hparams.ple_image_token_id
-                        : (llama_token) hparams.ple_eos_token_id;
+                    ext.tok = hparams.ple_embd_token;
                 }
 
                 cells.ext_set(idx, ext);
@@ -1849,7 +1847,7 @@ void llama_kv_cache::set_input_v_rot(ggml_tensor * dst) const {
 }
 
 bool llama_kv_cache::has_cell_ext() const {
-    // M-RoPE needs the 2D position, the PLE n-gram hash needs the token id
+    // M-RoPE needs the 2D position, the PLE / Engram n-gram hashes need the token id
     return hparams.n_pos_per_embd() > 1 || hparams.ple_n_heads > 0;
 }
 

src/llama-memory-recurrent.cpp

@@ -116,7 +116,15 @@ llama_memory_recurrent::llama_memory_recurrent(
 
     // allocate tensors and initialize the buffers to avoid NaNs in the padding
     for (auto & [buft, ctx] : ctx_map) {
-        ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft);
+        ggml_backend_buffer_t buf;
+        if (hparams.no_alloc) {
+            buf = ggml_backend_buft_alloc_buffer(buft, /*size =*/ 0); // dummy buffer
+            for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) {
+                t->buffer = buf; // set dummy buffer for RS cache so that the backend scheduler won't try to allocate it
+            }
+        } else {
+            buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); // real buffer
+        }
         if (!buf) {
             throw std::runtime_error("failed to allocate buffer for rs cache");
         }
@@ -435,8 +443,15 @@ void llama_memory_recurrent::set_rs_idx(llama_seq_id seq_id, uint32_t idx) {
 
 std::map<ggml_backend_buffer_type_t, size_t> llama_memory_recurrent::memory_breakdown() const {
     std::map<ggml_backend_buffer_type_t, size_t> ret;
-    for (const auto & [_, buf] : ctxs_bufs) {
-        ret[ggml_backend_buffer_get_type(buf.get())] += ggml_backend_buffer_get_size(buf.get());
+    for (const auto & [ctx, buf] : ctxs_bufs) {
+        ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get());
+
+        if (hparams.no_alloc) {
+            GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) == nullptr);
+            ret[buft] += ggml_backend_alloc_ctx_tensors_from_buft_size(ctx.get(), buft);
+        } else {
+            ret[buft] += ggml_backend_buffer_get_size(buf.get());
+        }
     }
     return ret;
 }

src/llama-model-loader.cpp

@@ -414,6 +414,7 @@ namespace GGUFMeta {
     template bool llama_model_loader::get_arr<std::vector<int32_t>>(enum llm_kv kid, std::vector<int32_t> & result, bool required);
     template bool llama_model_loader::get_arr<std::array<uint32_t, LLAMA_MAX_LAYERS>>(enum llm_kv kid, std::array<uint32_t, LLAMA_MAX_LAYERS> & result, bool required);
     template bool llama_model_loader::get_arr<std::vector<uint32_t>>(enum llm_kv kid, std::vector<uint32_t> & result, bool required);
+    template bool llama_model_loader::get_arr<std::vector<uint64_t>>(enum llm_kv kid, std::vector<uint64_t> & result, bool required);
     template bool llama_model_loader::get_arr<std::array<uint64_t, LLAMA_MAX_PLE_NGRAM>>(enum llm_kv kid, std::array<uint64_t, LLAMA_MAX_PLE_NGRAM> & result, bool required);
     template bool llama_model_loader::get_arr<std::array<uint64_t, LLAMA_MAX_PLE_HEADS>>(enum llm_kv kid, std::array<uint64_t, LLAMA_MAX_PLE_HEADS> & result, bool required);
 
@@ -441,6 +442,7 @@ namespace GGUFMeta {
     template bool llama_model_loader::get_key<bool>       (enum llm_kv kid, bool & result,        bool required);
     template bool llama_model_loader::get_key<float>      (enum llm_kv kid, float & result,       bool required);
     template bool llama_model_loader::get_key<uint32_t>   (enum llm_kv kid, uint32_t & result,    bool required);
+    template bool llama_model_loader::get_key<int32_t>    (enum llm_kv kid, int32_t & result,     bool required);
     template bool llama_model_loader::get_key<std::string>(enum llm_kv kid, std::string & result, bool required);
 
     template<>

src/llama-model-saver.cpp

@@ -8,6 +8,7 @@
 #include "llama-hparams.h"
 #include "llama-model.h"
 #include "llama-vocab.h"
+#include "models/models.h"
 
 #include <cstdint>
 #include <string>
@@ -122,6 +123,7 @@ void llama_model_saver::add_kv(const enum llm_kv key, const Container & value, c
 template void llama_model_saver::add_kv<std::vector<uint32_t>>(const enum llm_kv, const std::vector<uint32_t> &, const bool);
 template void llama_model_saver::add_kv<std::vector<float>>(const enum llm_kv, const std::vector<float> &, const bool);
 template void llama_model_saver::add_kv<std::vector<uint64_t>>(const enum llm_kv, const std::vector<uint64_t> &, const bool);
+template void llama_model_saver::add_kv<std::vector<int32_t>>(const enum llm_kv, const std::vector<int32_t> &, const bool);
 
 void llama_model_saver::add_kv(const enum llm_kv key, const std::vector<std::string> & value) {
     std::vector<const char *> tmp(value.size());
@@ -339,6 +341,32 @@ void llama_model_saver::add_kv_from_model() {
     add_kv(LLM_KV_HASH_LAYER_COUNT,                     hparams.dsv4_hash_layer_count);
     add_kv(LLM_KV_HYPER_CONNECTION_LOW_RANK,             hparams.hc_low_rank);
 
+    if (model->arch == LLM_ARCH_DEEPSEEK41) {
+        const auto * m41 = dynamic_cast<const llama_model_deepseek41 *>(model);
+        GGML_ASSERT(m41 != nullptr);
+        std::vector<uint32_t> kv_src;
+        std::vector<uint32_t> idx_src;
+        for (uint32_t il = 0; il < hparams.n_layer(); ++il) {
+            if (hparams.dsv41_kv_src[il] == (int32_t) il) {
+                kv_src.push_back(il);
+            }
+            if (hparams.dsv41_topk_src[il] == (int32_t) il) {
+                idx_src.push_back(il);
+            }
+        }
+        add_kv(LLM_KV_ATTENTION_KV_SOURCE_LAYERS,      kv_src);
+        add_kv(LLM_KV_ATTENTION_INDEXER_SOURCE_LAYERS, idx_src);
+        if (hparams.dsv41_cand_src >= 0) {
+            add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_SOURCE_LAYER, hparams.dsv41_cand_src);
+            add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_TOP_K_BLOCKS, hparams.dsv41_cand_topk_blocks);
+            add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_BLOCK_SIZE,   hparams.dsv41_cand_block_size);
+        }
+        if (hparams.ple_n_heads > 0) {
+            add_kv(LLM_KV_ENGRAM_PAD_TOKEN_ID, m41->engram_pad_id);
+            add_kv(LLM_KV_ENGRAM_TOKEN_MAP,    m41->engram_token_map);
+        }
+    }
+
     // the PLE group only means anything whole: write all of it or none
     if (hparams.ple_n_heads > 0) {
         std::vector<uint32_t> ple_layers;
@@ -355,13 +383,13 @@ void llama_model_saver::add_kv_from_model() {
         add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.ple_head_dim);
         add_kv(LLM_KV_PLE_LAYER_MULTIPLIERS, std::vector<uint64_t>(
                 hparams.ple_layer_multipliers.begin(),
-                hparams.ple_layer_multipliers.begin() + hparams.ple_ngram_size));
+                hparams.ple_layer_multipliers.begin() + hparams.ple_n_modules*hparams.ple_ngram_size));
         add_kv(LLM_KV_PLE_HEAD_OFFSETS,      std::vector<uint64_t>(
                 hparams.ple_head_offsets.begin(),
-                hparams.ple_head_offsets.begin() + hparams.ple_n_heads));
+                hparams.ple_head_offsets.begin() + hparams.ple_n_modules*hparams.ple_n_heads));
         add_kv(LLM_KV_PLE_HEAD_VOCAB_SIZES,  std::vector<uint64_t>(
                 hparams.ple_head_vocab_sizes.begin(),
-                hparams.ple_head_vocab_sizes.begin() + hparams.ple_n_heads));
+                hparams.ple_head_vocab_sizes.begin() + hparams.ple_n_modules*hparams.ple_n_heads));
     }
 
     const float rope_scaling_factor = hparams.rope_freq_scale_train == 1.0f ? 0.0f : 1.0f/hparams.rope_freq_scale_train;

src/llama-model.cpp

@@ -212,6 +212,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
             return new llama_model_dots3note(params);
         case LLM_ARCH_DEEPSEEK4:
             return new llama_model_deepseek4(params);
+        case LLM_ARCH_DEEPSEEK41:
+            return new llama_model_deepseek41(params);
         case LLM_ARCH_GLM_DSA:
             return new llama_model_glm_dsa(params);
         case LLM_ARCH_MISTRAL4:
@@ -398,7 +400,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
     const llama_meta_device_get_split_state_userdata * ud = (const llama_meta_device_get_split_state_userdata *) userdata;
     const llama_hparams & hparams = ud->model->hparams;
     const std::string tensor_name = tensor->name;
-    const bool is_dsv4 = ud->model->arch == LLM_ARCH_DEEPSEEK4 ||
+    const bool is_dsv4 = ud->model->arch == LLM_ARCH_DEEPSEEK4 || ud->model->arch == LLM_ARCH_DEEPSEEK41 ||
         (ud->model->arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0);
 
     static const std::regex pattern_q_weight        ("blk\\.\\d*\\.attn_q.weight");
@@ -3032,6 +3034,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                             nullptr);
                 }
             } break;
+        case LLM_ARCH_DEEPSEEK41:
         case LLM_ARCH_DEEPSEEK4:
             {
                 GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE);
@@ -3441,7 +3444,7 @@ int32_t llama_model_n_head_kv(const llama_model * model) {
 int32_t llama_model_n_swa(const llama_model * model) {
     // dsv4 kv-cache has SWA but it cannot be used as a rollback because of
     // other compression ratios, so we return 0 here
-    if (model->arch == LLM_ARCH_DEEPSEEK4) {
+    if (model->arch == LLM_ARCH_DEEPSEEK4 || model->arch == LLM_ARCH_DEEPSEEK41) {
         return 0;
     }
     return model->hparams.n_swa;
@@ -3536,6 +3539,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
         case LLM_ARCH_DEEPSEEK2OCR:
         case LLM_ARCH_DEEPSEEK32:
         case LLM_ARCH_DEEPSEEK4:
+        case LLM_ARCH_DEEPSEEK41:
         case LLM_ARCH_MUSE_GLIMMER:
         case LLM_ARCH_MUSE_SPARK:
         case LLM_ARCH_PLM:

src/llama-model.h

@@ -605,6 +605,10 @@ struct llama_layer {
     struct ggml_tensor * indexer_kpool_gate = nullptr;
     struct ggml_tensor * indexer_kpool_ape  = nullptr;
 
+    // deepseek41
+    struct ggml_tensor * engram_kv   = nullptr;
+    struct ggml_tensor * engram_gate = nullptr; // q_weight*k_weight folded
+
     // MSA
     struct ggml_tensor * index_q_proj = nullptr;
     struct ggml_tensor * index_k_proj = nullptr;

src/llama-quant.cpp

@@ -316,6 +316,9 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param
     // do not quantize the i32 token-id -> expert-id routing table (DeepSeek-V4)
     quantize &= name.find("ffn_gate_tid2eid.weight") == std::string::npos;
 
+    // Engram gate is an elementwise weight (ggml_mul), not a matmul operand
+    quantize &= name.find("engram_gate.weight") == std::string::npos;
+
     // these are very small (e.g. 4x4)
     quantize &= name.find("altup")  == std::string::npos;
     quantize &= name.find("laurel") == std::string::npos;

src/models/deepseek4.cpp

@@ -8,14 +8,6 @@
 #include <stdexcept>
 #include <string>
 
-static float dsv4_rope_attn_factor(float freq_scale, float ext_factor) {
-    if (ext_factor == 0.0f) {
-        return 1.0f;
-    }
-
-    return 1.0f / (1.0f + 0.1f*logf(1.0f/freq_scale));
-}
-
 void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) {
     if (hparams.n_layer_nextn > 0) {
         const uint32_t n_layer_main = hparams.n_layer_all - hparams.n_layer_nextn;
@@ -358,6 +350,18 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_pre(
         ggml_tensor ** post,
         ggml_tensor ** comb,
         int il) const {
+    return build_hc_pre(x, build_hc_mixes(x, hc_fn, hc_scale, hc_base, post, comb, false, il), il);
+}
+
+ggml_tensor * llama_model_deepseek4::graph::build_hc_mixes(
+        ggml_tensor * x,
+        ggml_tensor * hc_fn,
+        ggml_tensor * hc_scale,
+        ggml_tensor * hc_base,
+        ggml_tensor ** post,
+        ggml_tensor ** comb,
+        bool f32_src,
+        int il) const {
     const int64_t hc         = hparams.dsv4_hc_mult;
     const int64_t hc_dim     = hc*n_embd;
     const int64_t hc_mix_dim = (2 + hc)*hc;
@@ -369,6 +373,10 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_pre(
     ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc_dim, nt);
     ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, norm_rms_eps);
     ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm);
+    if (f32_src && !ggml_is_quantized(hc_fn->type)) {
+        // F32 weights (model.py hc_pre runs in fp32): request F32 inputs, the mixes are near-tie sensitive
+        ggml_prec_set_src(mixes, GGML_PREC_F32, 1);
+    }
     cb(mixes, "hc_mixes", il);
 
     ggml_tensor * scale_pre  = dsv4_view_1d(ctx0, hc_scale, 1, 0);
@@ -404,8 +412,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_pre(
     }
     cb(*comb, "hc_comb", il);
 
-    ggml_tensor * result = build_hc_pre(x, pre, il);
-    return result;
+    return pre;
 }
 
 ggml_tensor * llama_model_deepseek4::graph::build_hc_post(

src/models/deepseek41.cpp

@@ -0,0 +1,950 @@
+#include "models.h"
+#include "llama-impl.h"
+
+#include "llama-kv-cache-dsv4.h"
+
+#include <algorithm>
+#include <cinttypes>
+#include <cmath>
+#include <cstdlib>
+#include <stdexcept>
+#include <string>
+
+// ref: deepseek-ai/DeepSeek-V4.1-Flash-1017 inference/model.py, engram.py
+// compressed layers attend over [window | all compressed rows] with top-k as a mask. with flash attention the default
+// is sparse: fused indexer, FA bounded to window + top-k, decode-sized ubatches gather only the selected rows (same
+// keys and mask, different key order). LLAMA_DSV41_SPARSE=0 restores the reference path
+
+static constexpr int64_t DSV41_ENGRAM_DEAD = -1;
+
+static bool dsv41_env_on(const char * name) {
+    const char * v = getenv(name);
+    return !v || atoi(v) != 0;
+}
+
+void llama_model_deepseek41::load_arch_hparams(llama_model_loader & ml) {
+    llama_model_deepseek4::load_arch_hparams(ml);
+
+    hparams.n_embd_out_impl = hparams.n_embd; // no hc_head
+    if (hparams.indexer_top_k == 0) {
+        throw std::runtime_error("DeepSeek-V4.1 indexer top_k must be > 0");
+    }
+
+    // each compressed layer reads the KV / top-k of the latest source layer at or before it
+    std::vector<uint32_t> kv_src;
+    std::vector<uint32_t> idx_src;
+    ml.get_arr(LLM_KV_ATTENTION_KV_SOURCE_LAYERS,      kv_src);
+    ml.get_arr(LLM_KV_ATTENTION_INDEXER_SOURCE_LAYERS, idx_src);
+
+    const uint32_t n_layer = hparams.n_layer();
+    const auto out_of_range = [&](const std::vector<uint32_t> & v) {
+        return std::any_of(v.begin(), v.end(), [&](uint32_t il) { return il >= n_layer; });
+    };
+    if (out_of_range(kv_src) || out_of_range(idx_src)) {
+        throw std::runtime_error("DeepSeek-V4.1 kv/index source layer out of range");
+    }
+    hparams.dsv41_kv_src.fill(-1);
+    hparams.dsv41_topk_src.fill(-1);
+    int32_t last_kv  = -1;
+    int32_t last_idx = -1;
+    for (uint32_t il = 0; il < n_layer; ++il) {
+        const uint32_t ratio = hparams.dsv4_compress_ratios[il];
+        const bool is_kv  = std::find(kv_src.begin(),  kv_src.end(),  il) != kv_src.end();
+        const bool is_idx = std::find(idx_src.begin(), idx_src.end(), il) != idx_src.end();
+        if ((is_kv || is_idx) && ratio == 0) {
+            throw std::runtime_error(format("DeepSeek-V4.1 source layer %u has compress ratio 0", il));
+        }
+        last_kv  = is_kv  ? (int32_t) il : last_kv;
+        last_idx = is_idx ? (int32_t) il : last_idx;
+        if (ratio == 0) {
+            continue;
+        }
+        if (ratio > 2 || last_kv < 0 || last_idx < last_kv ||
+                hparams.dsv4_compress_ratios[last_kv] != ratio || hparams.dsv4_compress_ratios[last_idx] != ratio) {
+            throw std::runtime_error(format("DeepSeek-V4.1 layer %u: ratio %u without a matching kv/index source", il, ratio));
+        }
+        hparams.dsv41_kv_src[il]   = last_kv;
+        hparams.dsv41_topk_src[il] = last_idx;
+    }
+
+    ml.get_key(LLM_KV_ATTENTION_INDEXER_CANDIDATE_SOURCE_LAYER, hparams.dsv41_cand_src, false);
+    if (hparams.dsv41_cand_src >= 0) {
+        ml.get_key(LLM_KV_ATTENTION_INDEXER_CANDIDATE_TOP_K_BLOCKS, hparams.dsv41_cand_topk_blocks);
+        ml.get_key(LLM_KV_ATTENTION_INDEXER_CANDIDATE_BLOCK_SIZE,   hparams.dsv41_cand_block_size);
+        const int32_t cs = hparams.dsv41_cand_src;
+        if ((uint32_t) cs >= n_layer || hparams.dsv41_topk_src[cs] != cs || hparams.dsv41_cand_block_size == 0 ||
+                hparams.dsv41_cand_topk_blocks == 0) {
+            throw std::runtime_error("DeepSeek-V4.1 invalid candidate selection (source must be an index source, block size and top-k blocks > 0)");
+        }
+        for (uint32_t il = cs + 1; il < n_layer; ++il) {
+            if (hparams.dsv41_topk_src[il] == (int32_t) il && hparams.dsv4_compress_ratios[il] != hparams.dsv4_compress_ratios[cs]) {
+                throw std::runtime_error("DeepSeek-V4.1 candidate consumers must share the source's compress ratio");
+            }
+        }
+    }
+
+    // Engram: the PLE n-gram hash of qwen4exp with one set of hash constants per Engram layer, stored back to back
+    std::vector<uint32_t> e_layers;
+    ml.get_arr(LLM_KV_PLE_LAYERS, e_layers, false);
+    if (!e_layers.empty()) {
+        std::vector<uint64_t> mult;
+        std::vector<uint64_t> offsets;
+        std::vector<uint64_t> primes;
+        ml.get_key(LLM_KV_PLE_NGRAM_SIZE,             hparams.ple_ngram_size);
+        ml.get_key(LLM_KV_PLE_HEADS_PER_NGRAM,        hparams.ple_heads_per_ngram);
+        ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.ple_head_dim);
+        ml.get_arr(LLM_KV_PLE_LAYER_MULTIPLIERS,      mult);
+        ml.get_arr(LLM_KV_PLE_HEAD_OFFSETS,           offsets);
+        ml.get_arr(LLM_KV_PLE_HEAD_VOCAB_SIZES,       primes);
+        ml.get_key(LLM_KV_ENGRAM_PAD_TOKEN_ID,        engram_pad_id);
+        ml.get_arr(LLM_KV_ENGRAM_TOKEN_MAP,           engram_token_map);
+
+        const size_t n_e    = e_layers.size();
+        const size_t n_gram = hparams.ple_ngram_size;
+        const size_t n_cols = n_gram >= 2 ? (n_gram - 1)*hparams.ple_heads_per_ngram : 0;
+        if (n_cols == 0 || hparams.ple_head_dim == 0 ||
+                mult.size() != n_e*n_gram || primes.size() != n_e*n_cols || offsets.size() != n_e*n_cols ||
+                mult.size() > LLAMA_MAX_PLE_NGRAM || primes.size() > LLAMA_MAX_PLE_HEADS ||
+                engram_pad_id >= engram_token_map.size()) {
+            throw std::runtime_error("DeepSeek-V4.1 engram metadata is inconsistent");
+        }
+        for (size_t e = 0; e < n_e; ++e) {
+            // module e belongs to the e-th Engram layer
+            if (e_layers[e] >= n_layer || (e > 0 && e_layers[e] <= e_layers[e - 1])) {
+                throw std::runtime_error(format("DeepSeek-V4.1 engram layer %u is out of range or out of order", e_layers[e]));
+            }
+            hparams.is_ple_impl.set(e_layers[e]);
+        }
+        for (size_t c = 0; c < n_e*n_cols; ++c) {
+            // a bucket count of 0 traps in the hash; row ids are int32
+            if (primes[c] < 2 || offsets[c] + primes[c] > (uint64_t) INT32_MAX) {
+                throw std::runtime_error(format("DeepSeek-V4.1 engram head %zu: bucket range out of range", c));
+            }
+            hparams.ple_head_offsets[c]     = (uint32_t) offsets[c];
+            hparams.ple_head_vocab_sizes[c] = (uint32_t) primes[c];
+        }
+        std::copy(mult.begin(), mult.end(), hparams.ple_layer_multipliers.begin());
+        hparams.ple_n_heads   = (uint32_t) n_cols;
+        hparams.ple_n_modules = (uint32_t) n_e;
+    }
+}
+
+void llama_model_deepseek41::load_arch_tensors(llama_model_loader & ml) {
+    LLAMA_LOAD_LOCALS;
+
+    const int64_t q_lora_rank     = hparams.n_lora_q;
+    const int64_t n_ff_exp        = hparams.n_ff_exp();
+    const int64_t n_expert_shared = hparams.n_expert_shared;
+    const int64_t n_embd_head     = hparams.n_embd_head_k();
+    const int64_t o_groups        = hparams.dsv4_o_group_count;
+    const int64_t o_lora_rank     = hparams.dsv4_o_lora_rank;
+    const int64_t hc_mult         = hparams.dsv4_hc_mult;
+    const int64_t hc_dim          = hc_mult * n_embd;
+    const int64_t hc_mix_dim      = (2 + hc_mult) * hc_mult;
+    const int64_t n_idx_embd      = hparams.indexer_head_size;
+    const int64_t n_idx_head      = hparams.indexer_n_head;
+
+    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);
+
+    // all Engram tables in one lazily read tensor, like the qwen4exp n-gram table
+    if (hparams.ple_n_heads > 0) {
+        int64_t n_rows = 0;
+        for (size_t c = 0; c < (size_t) hparams.ple_n_modules*hparams.ple_n_heads; ++c) {
+            n_rows = std::max(n_rows, (int64_t) hparams.ple_head_offsets[c] + hparams.ple_head_vocab_sizes[c]);
+        }
+        const std::string name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").str();
+        if (const auto * w = ml.get_weight(name.c_str())) {
+            if (w->tensor->ne[1] < n_rows) {
+                throw std::runtime_error(format("%s has %" PRId64 " rows, too few for the Engram heads (%" PRId64 ")",
+                                                name.c_str(), w->tensor->ne[1], n_rows));
+            }
+            n_rows = w->tensor->ne[1];
+        }
+        per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), {(int64_t) hparams.ple_head_dim, n_rows}, TENSOR_READ_LAZY);
+    }
+
+    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.attn_sinks    = create_tensor(tn(LLM_TENSOR_ATTN_SINKS,    "weight", i), {n_head}, 0);
+        layer.wq_a          = create_tensor(tn(LLM_TENSOR_ATTN_Q_A,      "weight", i), {n_embd, q_lora_rank}, 0);
+        layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {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}, 0);
+        layer.wkv           = create_tensor(tn(LLM_TENSOR_ATTN_KV,       "weight", i), {n_embd, n_embd_head}, 0);
+        layer.attn_kv_norm  = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM,  "weight", i), {n_embd_head}, 0);
+        layer.wo_a          = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A,    "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank, o_groups}, TENSOR_ALLOW_RESHAPE);
+        layer.wo_b          = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B,    "weight", i), {o_groups * o_lora_rank, n_embd}, 0);
+
+        layer.hc_attn_fn    = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN,    "weight", i), {hc_dim, hc_mix_dim}, 0);
+        layer.hc_attn_base  = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE,  "weight", i), {hc_mix_dim}, 0);
+        layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0);
+        layer.hc_ffn_fn     = create_tensor(tn(LLM_TENSOR_HC_FFN_FN,     "weight", i), {hc_dim, hc_mix_dim}, 0);
+        layer.hc_ffn_base   = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE,   "weight", i), {hc_mix_dim}, 0);
+        layer.hc_ffn_scale  = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE,  "weight", i), {3}, 0);
+
+        if (hparams.dsv41_kv_src[i] == i) {
+            layer.attn_comp_wkv  = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WKV,  "weight", i), {n_embd, n_embd_head}, 0);
+            layer.attn_comp_norm = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_NORM, "weight", i), {n_embd_head}, 0);
+            if (hparams.dsv4_compress_ratios[i] > 1) {
+                layer.attn_comp_wgate = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "weight", i), {n_embd, n_embd_head}, 0);
+            }
+            layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd_head, n_idx_embd}, 0);
+            layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {n_idx_embd}, 0);
+        }
+        if (hparams.dsv41_topk_src[i] == i) {
+            layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, n_idx_head * n_idx_embd}, 0);
+            layer.indexer_proj     = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ,     "weight", i), {n_embd, n_idx_head}, 0);
+        }
+
+        if (engram_index(i) >= 0) {
+            layer.engram_kv   = create_tensor(tn(LLM_TENSOR_ENGRAM_KV,   "weight", i), {(int64_t) hparams.ple_n_heads * hparams.ple_head_dim, n_embd * (hc_mult + 1)}, 0);
+            layer.engram_gate = create_tensor(tn(LLM_TENSOR_ENGRAM_GATE, "weight", i), {n_embd, hc_mult}, 0);
+        }
+
+        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);
+        layer.ffn_exp_probs_b_vl = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B_VL, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED);
+        layer.ffn_norm        = create_tensor(tn(LLM_TENSOR_FFN_NORM,        "weight", i), {n_embd}, 0);
+
+        layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd,   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,   n_expert}, 0);
+        layer.ffn_up_exps   = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,   "weight", i), {n_embd,   n_ff_exp, n_expert}, 0);
+
+        layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
+        layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd}, 0);
+        layer.ffn_up_shexp   = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,   "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
+    }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_deepseek41::build_arch_graph(const llm_graph_params & params) const {
+    if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
+        throw std::runtime_error("DeepSeek-V4.1 has no MTP graph: its DSpark draft is a separate DFLASH GGUF "
+                "(convert_hf_to_gguf.py --dspark), used with --spec-type draft-dspark");
+    }
+    return std::make_unique<graph>(*this, params);
+}
+
+// Engram n-gram hash (engram.py NgramHashState.forward)
+int llama_model_deepseek41::engram_index(int il) const {
+    if (hparams.ple_n_heads == 0 || !hparams.is_ple(il)) {
+        return -1;
+    }
+    int e = 0;
+    for (int j = 0; j < il; ++j) {
+        e += hparams.is_ple(j);
+    }
+    return e;
+}
+
+void llama_model_deepseek41::engram_hash(int e, const llama_ubatch & ubatch, const std::vector<llama_token> & prev, std::vector<int32_t> & out) const {
+    const int64_t n_gram = hparams.ple_ngram_size;
+    const int64_t n_prev = n_gram - 1;
+    const int64_t n_head = hparams.ple_heads_per_ngram;
+    const int64_t n_cols = hparams.ple_n_heads;
+
+    const uint64_t * mult   = hparams.ple_layer_multipliers.data() + (size_t) e*n_gram;
+    const uint32_t * primes = hparams.ple_head_vocab_sizes.data()  + (size_t) e*n_cols;
+    const int64_t    pad    = engram_token_map[engram_pad_id];
+    const uint32_t * offs   = hparams.ple_head_offsets.data()      + (size_t) e*n_cols;
+
+    const auto cid = [&](llama_token t) -> int64_t {
+        return t < 0 || t >= (llama_token) engram_token_map.size() ? DSV41_ENGRAM_DEAD : engram_token_map[t];
+    };
+
+    out.resize(ubatch.n_tokens*n_cols);
+    std::vector<uint64_t> toks(n_gram);
+    for (int64_t i = 0; i < ubatch.n_tokens; ++i) {
+        // shift 0 is the token itself; once a shift is blocked (sequence start or dead token) every longer one is too
+        const bool is_embd = !ubatch.token || (ubatch.is_mixed() && ubatch.type[i]);
+        int64_t c = is_embd ? DSV41_ENGRAM_DEAD : cid(ubatch.token[i]);
+        bool blocked = c == DSV41_ENGRAM_DEAD;
+        toks[0] = blocked ? pad : c;
+        for (int64_t s = 1; s < n_gram; ++s) {
+            c = ubatch.pos[i] - s < 0 ? DSV41_ENGRAM_DEAD : cid(prev[i*n_prev + (n_prev - s)]);
+            blocked = blocked || c == DSV41_ENGRAM_DEAD;
+            toks[s] = blocked ? pad : c;
+        }
+
+        uint64_t rolling = toks[0]*mult[0];
+        for (int64_t o = 1; o < n_gram; ++o) {
+            rolling ^= toks[o]*mult[o];
+            for (int64_t h = 0; h < n_head; ++h) {
+                const int64_t col = (o - 1)*n_head + h;
+                out[i*n_cols + col] = (int32_t) (rolling % primes[col] + offs[col]);
+            }
+        }
+    }
+}
+
+// cand_pin: +1e30 (the reference uses +inf) on the block holding each query's newest compressed row, so a partly
+// filled block always survives; a finite pin keeps fully masked blocks at -inf instead of NaN
+class llm_graph_input_dsv41 : public llm_graph_input_i {
+public:
+    llm_graph_input_dsv41(const llama_model_deepseek41 & model, const llama_kv_cache_dsv4_context * mctx) :
+        model(model), mctx(mctx) {}
+
+    void set_input(const llama_ubatch * ubatch) override;
+    bool can_reuse(const llm_graph_params & params) override;
+
+    std::vector<ggml_tensor *> engram_ids;
+    ggml_tensor * engram_keep = nullptr;    // 0 for embedding (image) tokens; only in mixed ubatches
+    ggml_tensor * cand_pin = nullptr;
+    bool cand_csa = false;
+
+    const llama_model_deepseek41 & model;
+    const llama_kv_cache_dsv4_context * mctx;
+
+private:
+    const llama_kv_cache_dsv4_context::comp_plan & cand_plan(const llama_ubatch & ubatch) const {
+        return cand_csa ? mctx->get_csa_plan(ubatch) : mctx->get_hca_plan(ubatch);
+    }
+};
+
+void llm_graph_input_dsv41::set_input(const llama_ubatch * ubatch) {
+    if (!engram_ids.empty()) {
+        // a token shared by several sequences has the same history in each, so its first sequence resolves it
+        std::vector<llama_token> prev;
+        mctx->get_raw()->get_prev_tokens(*ubatch, model.hparams.ple_ngram_size - 1, prev);
+        // hash all Engram layers first so a lazily read table gets one batched readahead before the first gather
+        std::vector<std::vector<int32_t>> ids(engram_ids.size());
+        std::vector<int32_t> all;
+        for (size_t e = 0; e < engram_ids.size(); ++e) {
+            model.engram_hash((int) e, *ubatch, prev, ids[e]);
+            all.insert(all.end(), ids[e].begin(), ids[e].end());
+        }
+        if (model.can_prefetch.count(model.per_layer_tok_embd)) {
+            llama_prefetch_rows(model.per_layer_tok_embd, all.data(), all.size());
+        }
+        for (size_t e = 0; e < engram_ids.size(); ++e) {
+            if (engram_ids[e]->buffer) {
+                ggml_backend_tensor_set(engram_ids[e], ids[e].data(), 0, ids[e].size()*sizeof(int32_t));
+            }
+        }
+    }
+
+    // embedding (image) tokens are hashed DEAD and get no Engram contribution (gate 0)
+    if (engram_keep && engram_keep->buffer) {
+        std::vector<float> keep(ubatch->n_tokens);
+        for (uint32_t i = 0; i < ubatch->n_tokens; ++i) {
+            const bool is_embd = !ubatch->token || (ubatch->is_mixed() && ubatch->type[i]);
+            keep[i] = is_embd ? 0.0f : 1.0f;
+        }
+        ggml_backend_tensor_set(engram_keep, keep.data(), 0, keep.size()*sizeof(float));
+    }
+
+    if (cand_pin && cand_pin->buffer) {
+        const auto & plan = cand_plan(*ubatch);
+        const int64_t n_blk = cand_pin->ne[0];
+        const int64_t bs    = model.hparams.dsv41_cand_block_size;
+        GGML_ASSERT((int64_t) plan.n_visible.size() == (int64_t) ubatch->n_tokens);
+        std::vector<float> pin(n_blk*ubatch->n_tokens, 0.0f);
+        for (int64_t i = 0; i < (int64_t) ubatch->n_tokens; ++i) {
+            const int64_t nv = plan.n_visible[i];
+            if (nv > 0) {
+                GGML_ASSERT((nv - 1)/bs < n_blk);
+                pin[i*n_blk + (nv - 1)/bs] = 1e30f;
+            }
+        }
+        ggml_backend_tensor_set(cand_pin, pin.data(), 0, pin.size()*sizeof(float));
+    }
+}
+
+bool llm_graph_input_dsv41::can_reuse(const llm_graph_params & params) {
+    mctx = static_cast<const llama_kv_cache_dsv4_context *>(params.mctx);
+
+    bool res = true;
+    for (auto * t : engram_ids) {
+        res &= t->ne[0] == (int64_t) model.hparams.ple_n_heads*params.ubatch.n_tokens;
+    }
+    res &= (engram_keep != nullptr) == (!engram_ids.empty() && params.ubatch.embd != nullptr);
+    if (engram_keep) {
+        res &= engram_keep->ne[2] == (int64_t) params.ubatch.n_tokens;
+    }
+    if (cand_pin) {
+        const auto & plan = cand_plan(params.ubatch);
+        const int64_t bs = model.hparams.dsv41_cand_block_size;
+        res &= plan.n_kv > 0 && cand_pin->ne[0] == (plan.n_kv + bs - 1)/bs;
+        res &= cand_pin->ne[1]*cand_pin->ne[3] == (int64_t) params.ubatch.n_tokens;
+    }
+    return res;
+}
+
+ggml_tensor * llama_model_deepseek41::graph::build_engram(
+        const llama_model_deepseek41 & model,
+        ggml_tensor * x,
+        ggml_tensor * rows,
+        ggml_tensor * keep,
+        int il) const {
+    const auto & layer = model.layers[il];
+    const int64_t hc = hparams.dsv4_hc_mult;
+    const int64_t nt = x->ne[2];
+
+    GGML_ASSERT(rows->ne[0] == (int64_t) (hparams.ple_head_dim*hparams.ple_n_heads) && rows->ne[1] == nt);
+
+    ggml_tensor * kv = build_lora_mm(layer.engram_kv, rows);
+    cb(kv, "engram_kv", il);
+
+    ggml_tensor * key   = ggml_cont(ctx0, ggml_view_3d(ctx0, kv, n_embd, hc, nt, ggml_row_size(kv->type, n_embd), kv->nb[1], 0));
+    ggml_tensor * value = ggml_cont(ctx0, ggml_view_3d(ctx0, kv, n_embd, 1,  nt, ggml_row_size(kv->type, n_embd), kv->nb[1],
+            ggml_row_size(kv->type, hc*n_embd)));
+
+    // quantized copy has no binary-op kernel
+    ggml_tensor * egate = ggml_is_quantized(layer.engram_gate->type) ? ggml_cast(ctx0, layer.engram_gate, GGML_TYPE_F32) : layer.engram_gate;
+    ggml_tensor * dot = ggml_mul(ctx0, ggml_mul(ctx0, ggml_rms_norm(ctx0, x, norm_rms_eps), ggml_rms_norm(ctx0, key, norm_rms_eps)), egate);
+    dot = ggml_scale(ctx0, ggml_sum_rows(ctx0, dot), 1.0f/sqrtf((float) n_embd));
+    cb(dot, "engram_dot", il);
+
+    // sigmoid(copysign(sqrt(max(|dot|, 1e-6)), dot)); sign(0) = +1 as in copysign
+    ggml_tensor * mag  = ggml_sqrt(ctx0, ggml_clamp(ctx0, ggml_abs(ctx0, dot), 1e-6f, INFINITY));
+    ggml_tensor * sign = ggml_scale_bias(ctx0, ggml_step(ctx0, ggml_neg(ctx0, dot)), -2.0f, 1.0f);
+    ggml_tensor * gate = ggml_sigmoid(ctx0, ggml_mul(ctx0, mag, sign));
+    if (keep) {
+        gate = ggml_mul(ctx0, gate, keep);
+    }
+    cb(gate, "engram_gate", il);
+
+    return ggml_add(ctx0, x, ggml_mul(ctx0, ggml_repeat_4d(ctx0, value, n_embd, hc, nt, 1), gate));
+}
+
+ggml_tensor * llama_model_deepseek41::graph::build_attention41(
+        const llama_model_deepseek41 & model,
+        llm_graph_input_dsv4 * inp_dsv4,
+        ggml_tensor * cur,
+        ggml_tensor * inp_pos,
+        ggml_tensor * cand_pin,
+        int il) {
+    const auto & layer = model.layers[il];
+    llm_graph_input_dsv4_raw * inp_attn = inp_dsv4->get_raw();
+
+    const int64_t n_embd_head      = hparams.n_embd_head_k();
+    const int64_t n_embd_head_rope = hparams.n_rot();
+    const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;
+    const int64_t n_groups         = hparams.dsv4_o_group_count;
+    const int64_t o_lora_rank      = hparams.dsv4_o_lora_rank;
+    const int64_t n_idx_embd       = hparams.indexer_head_size;
+    const int64_t n_idx_head       = hparams.indexer_n_head;
+    const int64_t nt               = cur->ne[1];
+    const float   kq_scale         = 1.0f/sqrtf(float(n_embd_head));
+    const int64_t ratio            = hparams.dsv4_compress_ratios[il];
+
+    const bool    comp_rope     = ratio != 0;
+    const float   freq_base_l   = comp_rope ? hparams.dsv4_compress_rope_base : freq_base;
+    const float   freq_scale_l  = comp_rope ? freq_scale : 1.0f;
+    const float   ext_factor_l  = comp_rope ? ext_factor : 0.0f;
+    const float   attn_factor_l = dsv4_rope_attn_factor(freq_scale_l, ext_factor_l);
+    const float   beta_fast_l   = comp_rope ? beta_fast : 0.0f;
+    const float   beta_slow_l   = comp_rope ? beta_slow : 0.0f;
+    const int32_t n_ctx_orig_l  = comp_rope ? n_ctx_orig : 0;
+
+    const auto rope = [&](ggml_tensor * t, ggml_tensor * pos, int64_t n_nope) {
+        t = ggml_rope_ext(ctx0, t, pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l,
+                freq_base_l, freq_scale_l, ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l);
+        return ggml_rope_set_offset(t, n_nope);
+    };
+
+    ggml_tensor * qr = build_norm(build_lora_mm(layer.wq_a, cur), layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
+    cb(qr, "qr_norm", il);
+
+    ggml_tensor * q = ggml_reshape_3d(ctx0, build_lora_mm(layer.wq_b, qr), n_embd_head, n_head, nt);
+    q = rope(q, inp_pos, n_embd_head_nope);
+    cb(q, "q", il);
+
+    ggml_tensor * kv = build_norm(build_lora_mm(layer.wkv, cur), layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il);
+    kv = rope(ggml_reshape_3d(ctx0, kv, n_embd_head, 1, nt), inp_pos, n_embd_head_nope);
+    cb(kv, "kv", il);
+
+    const auto & ci         = ratio == 2 ? inp_dsv4->get_csa()             : inp_dsv4->get_hca();
+    const auto * comp_ctx   = ratio == 2 ? inp_dsv4->mctx->get_csa()       : inp_dsv4->mctx->get_hca();
+    const auto * comp_state = ratio == 2 ? inp_dsv4->mctx->get_csa_state() : inp_dsv4->mctx->get_hca_state();
+
+    const int32_t rd = inp_dsv4->mctx->get_comp_read_layer(il);
+
+    if (ratio != 0) {
+        GGML_ASSERT(hparams.dsv41_kv_src[il] >= 0 && (int64_t) comp_state->get_ratio() == ratio);
+
+        if (hparams.dsv41_kv_src[il] == il && ci.state_pos) {
+            ggml_tensor * state_kv    = build_lora_mm(layer.attn_comp_wkv, cur);
+            ggml_tensor * state_score = ratio > 1 ? build_lora_mm(layer.attn_comp_wgate, cur) : ggml_fill(ctx0, state_kv, 0.0f);
+
+            ggml_tensor * restored_kv    = comp_state->get_kv_all(ctx0, il);
+            ggml_tensor * restored_score = comp_state->get_score_all(ctx0, il);
+            if (ci.state_restore_src_idxs && ci.state_restore_dst_idxs) {
+                restored_kv    = comp_state->cpy_kv   (ctx0, ggml_get_rows(ctx0, restored_kv,    ci.state_restore_src_idxs), ci.state_restore_dst_idxs, il);
+                restored_score = comp_state->cpy_score(ctx0, ggml_get_rows(ctx0, restored_score, ci.state_restore_src_idxs), ci.state_restore_dst_idxs, il);
+            }
+
+            if (ci.state_write_idxs) {
+                const int64_t n_blocks = ci.state_write_pos->ne[0];
+                const int64_t n_state  = comp_state->get_n_rows();
+                GGML_ASSERT(ci.state_read_idxs->ne[0] == ratio*n_blocks);
+
+                ggml_tensor * src_kv    = ggml_concat(ctx0, ggml_view_2d(ctx0, restored_kv,    restored_kv->ne[0],    n_state, restored_kv->nb[1],    0), state_kv,    1);
+                ggml_tensor * src_score = ggml_concat(ctx0, ggml_view_2d(ctx0, restored_score, restored_score->ne[0], n_state, restored_score->nb[1], 0), state_score, 1);
+
+                ggml_tensor * v = ggml_get_rows(ctx0, src_kv,    ci.state_read_idxs);
+                ggml_tensor * s = ggml_get_rows(ctx0, src_score, ci.state_read_idxs);
+                v = ggml_cont(ctx0, ggml_permute(ctx0, ggml_reshape_3d(ctx0, v, n_embd_head, ratio, n_blocks), 1, 0, 2, 3));
+                s = ggml_cont(ctx0, ggml_permute(ctx0, ggml_reshape_3d(ctx0, s, n_embd_head, ratio, n_blocks), 1, 0, 2, 3));
+
+                ggml_tensor * comp = ggml_sum_rows(ctx0, ggml_mul(ctx0, v, ggml_soft_max(ctx0, s)));
+                comp = ggml_cont(ctx0, ggml_permute(ctx0, comp, 1, 0, 2, 3));
+                comp = build_norm(comp, layer.attn_comp_norm, nullptr, LLM_NORM_RMS, il);
+                cb(comp, "comp_latent", il);
+
+                // index K from the RoPE-free latent; a latent stands for the first token of its group
+                ggml_tensor * ik = build_norm(build_lora_mm(layer.indexer_attn_k, comp), layer.indexer_k_norm, nullptr, LLM_NORM_RMS, il);
+                ik = rope(ik, ci.state_write_pos, n_idx_embd - n_embd_head_rope);
+                comp = rope(comp, ci.state_write_pos, n_embd_head_nope);
+
+                ggml_tensor * row = ggml_concat(ctx0, comp, ik, 0);
+                ggml_build_forward_expand(gf, comp_ctx->cpy_k(ctx0, row, ci.state_write_idxs, il));
+                comp_row_of[il] = row;
+            }
+
+            if (ci.state_snapshot_src_idxs && ci.state_snapshot_dst_idxs) {
+                ggml_tensor * snap_kv    = ggml_concat(ctx0, restored_kv,    state_kv,    1);
+                ggml_tensor * snap_score = ggml_concat(ctx0, restored_score, state_score, 1);
+                ggml_build_forward_expand(gf, comp_state->cpy_kv   (ctx0, ggml_get_rows(ctx0, snap_kv,    ci.state_snapshot_src_idxs), ci.state_snapshot_dst_idxs, il));
+                ggml_build_forward_expand(gf, comp_state->cpy_score(ctx0, ggml_get_rows(ctx0, snap_score, ci.state_snapshot_src_idxs), ci.state_snapshot_dst_idxs, il));
+            }
+            ggml_build_forward_expand(gf, comp_state->cpy_kv   (ctx0, ggml_get_rows(ctx0, state_kv,    ci.state_persist_src_idxs), ci.state_persist_dst_idxs, il));
+            ggml_build_forward_expand(gf, comp_state->cpy_score(ctx0, ggml_get_rows(ctx0, state_score, ci.state_persist_src_idxs), ci.state_persist_dst_idxs, il));
+        }
+
+        if (rd == il && hparams.dsv41_kv_src[il] != il && comp_row_of[hparams.dsv41_kv_src[il]]) {
+            // first consumer on this device: append the source's new rows to the local mirror
+            ggml_build_forward_expand(gf, comp_ctx->cpy_k(ctx0, comp_row_of[hparams.dsv41_kv_src[il]], ci.state_write_idxs, il));
+        }
+    }
+
+    ggml_tensor * out = nullptr;
+    if (ratio == 0 || !ci.kq_mask) {
+        // no compressed row visible yet: window only
+        out = build_raw_attention(inp_attn, q, kv, layer.attn_sinks, kq_scale, il);
+    } else {
+        const int64_t n_comp = ci.kq_mask->ne[0];
+
+        ggml_tensor * comp_all = comp_ctx->get_k(ctx0, rd);
+        GGML_ASSERT(comp_all->ne[0] == n_embd_head + n_idx_embd && n_comp <= comp_all->ne[2]);
+        const int64_t n_stream = comp_all->ne[3];
+
+        ggml_tensor * top_k = topk_of[hparams.dsv41_topk_src[il]];
+        if (hparams.dsv41_topk_src[il] == il) {
+            ggml_tensor * ik = ggml_view_4d(ctx0, comp_all, n_idx_embd, 1, n_comp, n_stream,
+                    comp_all->nb[1], comp_all->nb[2], comp_all->nb[3], ggml_row_size(comp_all->type, n_embd_head));
+            // the fused indexer reads the strided view in place, the unfused path needs it contiguous; quantized
+            // caches are dequantized straight from the strided view (CUDA has no same-type copy of quantized views)
+            const bool fused_idx = sparse && cparams.fused_lid && ci.kq_mask->type == GGML_TYPE_F16;
+            if (ggml_is_quantized(ik->type)) {
+                ik = ggml_cast(ctx0, ik, GGML_TYPE_F32);
+            } else if (!fused_idx) {
+                ik = ggml_cont(ctx0, ik);
+            }
+
+            ggml_tensor * iq = ggml_reshape_3d(ctx0, build_lora_mm(layer.indexer_attn_q_b, qr), n_idx_embd, n_idx_head, nt);
+            iq = rope(iq, inp_pos, n_idx_embd - n_embd_head_rope);
+            iq = ggml_view_4d(ctx0, iq, iq->ne[0], iq->ne[1], nt/n_stream, n_stream, iq->nb[1], iq->nb[2], iq->nb[2]*(nt/n_stream), 0);
+
+            ggml_tensor * iw = ggml_scale(ctx0, build_lora_mm(layer.indexer_proj, cur), 1.0f/sqrtf(float(n_idx_embd*n_idx_head)));
+            iw = ggml_view_4d(ctx0, iw, iw->ne[0], nt/n_stream, 1, n_stream, iw->nb[1], iw->nb[1]*(nt/n_stream), iw->nb[1]*(nt/n_stream), 0);
+
+            // chunk queries so the [n_ih, tokens, n_comp] scores stay near 2^26 elements (at most 64 chunks); the fused
+            // indexer has no such intermediate and runs as one chunk
+            const int64_t n_tok = nt/n_stream;
+            GGML_ASSERT(ci.kq_mask->ne[1] == n_tok);
+            const int64_t chunk = fused_idx ? n_tok :
+                std::max({ (n_tok + 63)/64, (INT64_C(1) << 26)/(n_idx_head*n_comp*n_stream), INT64_C(1) });
+
+            const int32_t cand_src = hparams.dsv41_cand_src;
+            const bool    pick_blk = cand_src == il && cand_pin;              // level one: pick the candidate blocks
+            const bool    use_blk  = cand_src >= 0 && cand_src < il && cand_sel; // level two: score inside them
+            const int64_t bs       = hparams.dsv41_cand_block_size;
+            const int64_t n_blk    = pick_blk || use_blk ? (n_comp + bs - 1)/bs : 0;
+            GGML_ASSERT(!pick_blk || cand_pin->ne[0] == n_blk);
+            GGML_ASSERT(!use_blk  || cand_n_comp == n_comp);
+
+            ggml_tensor * ik_p = fused_idx ? nullptr : ggml_permute(ctx0, ik, 0, 2, 1, 3);
+            std::vector<ggml_tensor *> parts_top_k;
+            std::vector<ggml_tensor *> parts_sel;
+            for (int64_t t0 = 0; t0 < n_tok; t0 += chunk) {
+                const int64_t nc = std::min(chunk, n_tok - t0);
+                const auto rows = [&](ggml_tensor * t) {
+                    return ggml_view_4d(ctx0, t, t->ne[0], nc, 1, n_stream, t->nb[1], t->nb[2], t->nb[3], t0*t->nb[1]);
+                };
+                ggml_tensor * score = nullptr;
+                if (fused_idx) {
+                    score = ggml_lightning_indexer(ctx0, iq, ik, iw, ci.kq_mask);
+                    res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, score, il});
+                } else {
+                    ggml_tensor * iqc = ggml_view_4d(ctx0, iq, iq->ne[0], iq->ne[1], nc, n_stream, iq->nb[1], iq->nb[2], iq->nb[3], t0*iq->nb[2]);
+
+                    ggml_tensor * kq = ggml_mul_mat(ctx0, ik_p, ggml_permute(ctx0, iqc, 0, 2, 1, 3));
+                    if (!ggml_is_quantized(ik->type)) {
+                        ggml_prec_set_src(kq, GGML_PREC_F32, 1);
+                    }
+                    kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 2, 1, 0, 3));
+                    score = ggml_sum_rows(ctx0, ggml_mul(ctx0, ggml_relu(ctx0, kq), rows(iw)));
+                    ggml_tensor * mask  = rows(ci.kq_mask);
+                    score = ggml_add(ctx0, ggml_cont(ctx0, ggml_permute(ctx0, score, 2, 1, 0, 3)),
+                            mask->type == GGML_TYPE_F32 ? mask : ggml_cast(ctx0, mask, GGML_TYPE_F32));
+                }
+                cb(score, "idx_score", il);
+
+                if (pick_blk) {
+                    // best position per block, pin the newest block, keep the top blocks
+                    ggml_tensor * s2 = ggml_reshape_2d(ctx0, score, n_comp, nc*n_stream);
+                    if (n_blk*bs != n_comp) {
+                        s2 = ggml_concat(ctx0, s2, ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_blk*bs - n_comp, nc*n_stream), -INFINITY), 0);
+                    }
+                    ggml_tensor * bmax = ggml_pool_1d(ctx0, s2, GGML_OP_POOL_MAX, bs, bs, 0);
+                    bmax = ggml_add(ctx0, ggml_reshape_4d(ctx0, bmax, n_blk, nc, 1, n_stream), rows(cand_pin));
+                    cb(bmax, "cand_block_score", il);
+                    parts_sel.push_back(ggml_top_k(ctx0, bmax, std::min<int64_t>(hparams.dsv41_cand_topk_blocks, n_blk)));
+                } else if (use_blk) {
+                    // own scores, restricted to the source's candidate blocks: 0 inside them, -inf elsewhere
+                    ggml_tensor * bsel = ggml_cont(ctx0, rows(cand_sel));
+                    const int64_t k_b = bsel->ne[0];
+                    ggml_tensor * bm    = ggml_fill(ctx0, ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_blk, nc*n_stream, 1), -INFINITY);
+                    ggml_tensor * zeros = ggml_fill(ctx0, ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, k_b, nc*n_stream, 1), 0.0f);
+                    bm = ggml_set_rows(ctx0, bm, zeros, ggml_reshape_3d(ctx0, bsel, k_b, nc*n_stream, 1));
+                    bm = ggml_reshape_2d(ctx0, ggml_repeat_4d(ctx0, bm, bs, n_blk, nc*n_stream, 1), bs*n_blk, nc*n_stream);
+                    bm = ggml_reshape_4d(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, bm, n_comp, nc*n_stream, bm->nb[1], 0)), n_comp, nc, 1, n_stream);
+                    cb(bm, "cand_mask", il);
+                    score = ggml_add(ctx0, score, bm);
+                }
+
+                parts_top_k.push_back(ggml_top_k(ctx0, score, std::min<int64_t>(n_comp, hparams.indexer_top_k)));
+            }
+
+            const auto join = [&](std::vector<ggml_tensor *> v) {
+                while (v.size() > 1) {
+                    std::vector<ggml_tensor *> w;
+                    for (size_t i = 0; i < v.size(); i += 2) {
+                        w.push_back(i + 1 < v.size() ? ggml_concat(ctx0, v[i], v[i + 1], 1) : v[i]);
+                    }
+                    v.swap(w);
+                }
+                return v[0];
+            };
+            if (pick_blk) {
+                cand_sel    = join(parts_sel);
+                cand_n_comp = n_comp;
+            }
+            top_k = ggml_cont(ctx0, join(parts_top_k));
+            cb(top_k, "idx_top_k", il);
+            topk_of[il] = top_k;
+        }
+        GGML_ASSERT(top_k && "DeepSeek-V4.1: top-k source layer produced no top-k");
+
+        // a quantized K cache stores the window Hadamard-rotated (as build_raw_attention)
+        ggml_tensor * k_rot = inp_attn->self_k_rot;
+        if (k_rot) {
+            q  = llama_mul_mat_hadamard(ctx0, q, k_rot);
+            kv = llama_mul_mat_hadamard(ctx0, kv, k_rot);
+        }
+        ggml_build_forward_expand(gf, q);
+        ggml_build_forward_expand(gf, kv);
+        ggml_build_forward_expand(gf, inp_attn->mctx->cpy_k(ctx0, kv, inp_attn->get_k_idxs(), il));
+
+        ggml_tensor * raw_k    = inp_attn->mctx->get_k(ctx0, il);
+        ggml_tensor * raw_mask = inp_attn->get_kq_mask();
+
+        const int32_t topk_src = hparams.dsv41_topk_src[il];
+        const int64_t k_sel    = top_k->ne[0];
+        const int64_t n_raw    = raw_k->ne[2];
+
+        // few queries per stream: gather each query's top-k rows and attend over [window | selected], one FA stream
+        // per query, instead of a full n_comp mask per layer. same keys and mask as dense, different key order
+        const bool gather = sparse && cparams.flash_attn && raw_mask->type == GGML_TYPE_F16 && ci.kq_mask->type == GGML_TYPE_F16 &&
+                ((n_stream == 1 && nt <= 32) || nt == n_stream);
+
+        if (gather) {
+            // ggml_top_k returns the rows in no particular order (CUB DeviceTopK can vary between runs) and FA sums the
+            // gathered keys in that order: sort the rows once per top-k source so the result is deterministic
+            ggml_tensor *& sel = gather_sel_of[topk_src];
+            if (!sel) {
+                GGML_ASSERT(n_comp <= (INT64_C(1) << 24)); // row ids are exact in F32
+                ggml_tensor * ids = ggml_reshape_2d(ctx0, top_k, k_sel, nt);
+                ggml_tensor * ord = ggml_argsort(ctx0, ggml_cast(ctx0, ids, GGML_TYPE_F32), GGML_SORT_ORDER_ASC);
+                sel = ggml_reshape_2d(ctx0, ggml_get_rows(ctx0, ggml_reshape_3d(ctx0, ids, 1, k_sel, nt), ord), k_sel, nt);
+                cb(sel, "comp_top_k_sorted", il);
+            }
+
+            // built once per (storage, top-k) pair
+            ggml_tensor *& comp_k = gather_k_of[{rd, topk_src}];
+            if (!comp_k) {
+                if (n_stream == 1) {
+                    ggml_tensor * src_rows = ggml_view_3d(ctx0, comp_all, n_embd_head, n_comp, 1, comp_all->nb[2], comp_all->nb[3], 0);
+                    comp_k = ggml_get_rows(ctx0, src_rows, ggml_reshape_1d(ctx0, sel, k_sel*nt));
+                } else {
+                    ggml_tensor * src_rows = ggml_view_3d(ctx0, comp_all, n_embd_head, n_comp, n_stream, comp_all->nb[2], comp_all->nb[3], 0);
+                    comp_k = ggml_get_rows(ctx0, src_rows, sel);
+                }
+                comp_k = ggml_reshape_4d(ctx0, comp_k, n_embd_head, 1, k_sel, nt); // F32 (get_rows dequantizes)
+                if (k_rot) {
+                    // quantized K cache: the window is stored rotated, the compressed rows are not
+                    comp_k = llama_mul_mat_hadamard(ctx0, comp_k, k_rot);
+                }
+                if (comp_k->type != raw_k->type) {
+                    comp_k = ggml_cast(ctx0, comp_k, raw_k->type); // exact: rows are raw_k->type values
+                }
+                cb(comp_k, "comp_k_sel", il);
+            }
+
+            if (n_stream == 1 && nt > 1) {
+                if (ggml_is_quantized(raw_k->type)) {
+                    // no REPEAT kernel for quantized types: byte-exact copies by doubling concats
+                    ggml_tensor * rep = nullptr;
+                    ggml_tensor * pw  = raw_k;
+                    for (int64_t n = nt; n > 0; n >>= 1) {
+                        if (n & 1) {
+                            rep = rep ? ggml_concat(ctx0, rep, pw, 3) : pw;
+                        }
+                        if (n > 1) {
+                            pw = ggml_concat(ctx0, pw, pw, 3);
+                        }
+                    }
+                    raw_k = rep;
+                } else {
+                    raw_k = ggml_repeat_4d(ctx0, raw_k, n_embd_head, 1, n_raw, nt);
+                }
+            }
+            ggml_tensor * k_all = ggml_concat(ctx0, raw_k, comp_k, 2);
+
+            // mask: window row, then comp mask at the selected rows (-inf where top-k picked an invisible row);
+            // one per top-k source, shared by its layers
+            ggml_tensor *& kq_mask = gather_mask_of[topk_src];
+            if (!kq_mask) {
+                ggml_tensor * sel_mask = ggml_get_rows(ctx0, ggml_reshape_3d(ctx0, ci.kq_mask, 1, n_comp, nt), sel);
+                sel_mask = ggml_cast(ctx0, ggml_reshape_4d(ctx0, sel_mask, k_sel, 1, 1, nt), GGML_TYPE_F16);
+                kq_mask = ggml_concat(ctx0, ggml_reshape_4d(ctx0, raw_mask, raw_mask->ne[0], 1, 1, nt), sel_mask, 0);
+                cb(kq_mask, "comp_kq_mask_sel", il);
+            }
+
+            out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, layer.attn_sinks, nullptr, 0, kq_scale, il);
+        } else {
+            // dense attention over [window | all compressed rows], selection applied as a mask
+            ggml_tensor * comp_k = ggml_view_4d(ctx0, comp_all, n_embd_head, 1, n_comp, n_stream,
+                    comp_all->nb[1], comp_all->nb[2], comp_all->nb[3], 0);
+            if (k_rot) {
+                // compressed rows are stored unrotated: rotate in F32, then cast to the cache type (dequantized from
+                // the strided view: CUDA has no same-type quantized copy)
+                comp_k = comp_k->type != GGML_TYPE_F32 ? ggml_cast(ctx0, comp_k, GGML_TYPE_F32) : ggml_cont(ctx0, comp_k);
+                comp_k = llama_mul_mat_hadamard(ctx0, comp_k, k_rot);
+                if (comp_k->type != raw_k->type) {
+                    comp_k = ggml_cast(ctx0, comp_k, raw_k->type);
+                }
+            } else if (ggml_is_quantized(comp_k->type)) {
+                // quantized cache without rotation (LLAMA_ATTN_ROT_DISABLE=1): via F32 for the same reason
+                comp_k = ggml_cast(ctx0, ggml_cast(ctx0, comp_k, GGML_TYPE_F32), raw_k->type);
+            } else {
+                comp_k = comp_k->type != raw_k->type ? ggml_cast(ctx0, comp_k, raw_k->type) : ggml_cont(ctx0, comp_k);
+            }
+            cb(comp_k, "comp_k", il);
+            ggml_tensor * k_all = ggml_concat(ctx0, raw_k, comp_k, 2);
+
+            ggml_tensor * sel_mask = build_top_k_mask(ci.kq_mask, top_k, "comp_top_k_mask", il);
+            ggml_tensor * kq_mask  = ggml_concat(ctx0, raw_mask, sel_mask, 0);
+            cb(kq_mask, "comp_kq_mask", il);
+
+            // at most window + top-k visible columns per query: lets CUDA FA gather only those
+            const int64_t n_kv_max = sparse ? std::min<int64_t>(raw_mask->ne[0], hparams.n_swa) + k_sel : 0;
+            out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, layer.attn_sinks, nullptr, n_kv_max, kq_scale, il);
+        }
+        if (k_rot) {
+            out = llama_mul_mat_hadamard(ctx0, out, k_rot);
+        }
+        cb(out, "attn_comp", il);
+    }
+
+    out = ggml_reshape_3d(ctx0, out, n_embd_head, n_head, nt);
+    out = ggml_rope_ext_back(ctx0, out, inp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l,
+            freq_base_l, freq_scale_l, ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l);
+    out = ggml_rope_set_offset(out, n_embd_head_nope);
+    cb(out, "attn_derope", il);
+
+    out = ggml_permute(ctx0, ggml_reshape_3d(ctx0, out, n_head/n_groups*n_embd_head, n_groups, nt), 0, 2, 1, 3);
+    ggml_tensor * oa = ggml_mul_mat(ctx0, layer.wo_a, out);
+    cb(oa, "attn_wo_a", il);
+    oa = ggml_cont_2d(ctx0, ggml_permute(ctx0, oa, 0, 2, 1, 3), o_lora_rank*n_groups, nt);
+
+    out = build_lora_mm(layer.wo_b, oa);
+    cb(out, "attn_out", il);
+
+    return out;
+}
+
+llama_model_deepseek41::graph::graph(const llama_model & model_base, const llm_graph_params & params) :
+    llama_model_deepseek4::graph(params) {
+    const auto & model = static_cast<const llama_model_deepseek41 &>(model_base);
+
+    static const bool sparse_env = dsv41_env_on("LLAMA_DSV41_SPARSE");
+    sparse = sparse_env;
+
+    topk_of.assign(n_layer, nullptr);
+    comp_row_of.assign(n_layer, nullptr);
+
+    ggml_tensor * cur;
+
+    ggml_tensor * inp = build_inp_embd(model.tok_embd);
+    ggml_tensor * inp_pos = build_inp_pos();
+    ggml_tensor * inp_out_ids = build_inp_out_ids();
+    llm_graph_input_dsv4 * inp_dsv4 = build_inp_dsv4();
+    ggml_build_forward_expand(gf, inp_dsv4->get_raw()->self_kq_mask);
+
+    auto inp41 = std::make_unique<llm_graph_input_dsv41>(model, inp_dsv4->mctx);
+    // embedding input has no token ids: it skips Engram, like the reference's masked (image) positions
+    // the gathers go right after the token embeddings, so they share the first graph split (the table stays on the host)
+    std::vector<ggml_tensor *> engram_rows;
+    for (int il = 0; ubatch.token && hparams.ple_n_heads > 0 && il < n_layer; ++il) {
+        if (!hparams.is_ple(il)) {
+            continue;
+        }
+        ggml_tensor * ids = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, (int64_t) hparams.ple_n_heads*n_tokens);
+        ggml_set_name(ids, ("engram_ids-" + std::to_string(il)).c_str());
+        ggml_set_input(ids);
+        inp41->engram_ids.push_back(ids);
+
+        ggml_tensor * rows = ggml_get_rows(ctx0, model.per_layer_tok_embd, ids);
+        rows = ggml_reshape_2d(ctx0, rows, hparams.ple_head_dim*hparams.ple_n_heads, n_tokens);
+        cb(rows, "engram_rows", il);
+        ggml_build_forward_expand(gf, rows);
+        engram_rows.push_back(rows);
+    }
+    ggml_tensor * engram_keep = nullptr;
+    if (hparams.ple_n_heads > 0 && ubatch.token && ubatch.embd) {
+        engram_keep = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, 1, n_tokens);
+        ggml_set_input(engram_keep);
+        ggml_set_name(engram_keep, "engram_keep");
+        inp41->engram_keep = engram_keep;
+    }
+    ggml_tensor * cand_pin = nullptr;
+    if (hparams.dsv41_cand_src >= 0) {
+        inp41->cand_csa = hparams.dsv4_compress_ratios[hparams.dsv41_cand_src] == 2;
+        const auto & ci = inp41->cand_csa ? inp_dsv4->get_csa() : inp_dsv4->get_hca();
+        if (ci.kq_mask) {
+            const int64_t bs = hparams.dsv41_cand_block_size;
+            cand_pin = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, (ci.kq_mask->ne[0] + bs - 1)/bs, ci.kq_mask->ne[1], 1, ci.kq_mask->ne[3]);
+            ggml_set_input(cand_pin);
+            ggml_set_name(cand_pin, "cand_pin");
+            inp41->cand_pin = cand_pin;
+        }
+    }
+    res->add_input(std::move(inp41));
+
+    const int64_t hc = hparams.dsv4_hc_mult;
+    ggml_tensor * inpL = ggml_repeat_4d(ctx0, ggml_reshape_3d(ctx0, inp, n_embd, 1, n_tokens), n_embd, hc, n_tokens, 1);
+    cb(inpL, "hc_init", -1);
+
+    // single-pass mHC: the pre-mix used to collapse the stream lags one sublayer
+    ggml_tensor * pre_prev = nullptr;
+
+    for (int il = 0; il < n_layer; ++il) {
+        const auto & layer = model.layers[il];
+
+        const int e = model.engram_index(il);
+        if (e >= 0 && !engram_rows.empty()) {
+            inpL = build_engram(model, inpL, engram_rows[e], engram_keep, il);
+            cb(inpL, "engram_out", il);
+        }
+
+        // DSpark target features: mean over the hc copies of the layer input
+        if ((size_t) il < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[il]) {
+            res->t_layer_inp[il] = build_hc_mean(inpL);
+            cb(res->t_layer_inp[il], "layer_inp", il);
+            ggml_build_forward_expand(gf, res->t_layer_inp[il]);
+        }
+
+        ggml_tensor * residual = inpL;
+        ggml_tensor * post = nullptr;
+        ggml_tensor * comb = nullptr;
+
+        ggml_tensor * attn_pre = build_hc_mixes(inpL, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, &post, &comb, true, il);
+
+        // the first layer's attention collapses with the identity pre-mix: copy 0
+        cur = pre_prev ? build_hc_pre(inpL, pre_prev, il) : ggml_cont(ctx0, ggml_view_2d(ctx0, inpL, n_embd, n_tokens, inpL->nb[2], 0));
+        cb(cur, "hc_attn_pre", il);
+
+        cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+        cb(cur, "attn_norm", il);
+
+        cur = build_attention41(model, inp_dsv4, cur, inp_pos, cand_pin, il);
+
+        inpL = build_hc_post(cur, residual, post, comb, il);
+        cb(inpL, "hc_attn_post", il);
+
+        residual = inpL;
+        pre_prev = build_hc_mixes(inpL, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, &post, &comb, true, il);
+
+        cur = build_hc_pre(inpL, attn_pre, il);
+        cb(cur, "hc_ffn_pre", il);
+
+        ggml_build_forward_expand(gf, residual);
+        ggml_build_forward_expand(gf, post);
+        ggml_build_forward_expand(gf, comb);
+
+        cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
+        cb(cur, "ffn_norm", il);
+
+        // F32 router inputs: expert selection is near-tie sensitive
+        ggml_tensor * logits = build_lora_mm(layer.ffn_gate_inp, cur);
+        ggml_prec_set_acc(logits, GGML_PREC_F32);
+        if (!ggml_is_quantized(layer.ffn_gate_inp->type)) {
+            ggml_prec_set_src(logits, GGML_PREC_F32, 1);
+        }
+        cb(logits, "ffn_moe_logits", il);
+
+        // image (embedding) tokens use the vision router bias, as DeepSeek-V4
+        ggml_tensor * exp_probs_b = layer.ffn_exp_probs_b;
+        if (ubatch.embd != nullptr && layer.ffn_exp_probs_b_vl) {
+            exp_probs_b = layer.ffn_exp_probs_b_vl;
+        }
+
+        // the reference normalizes only when it routes to more than one expert
+        ggml_tensor * moe_out = build_moe_ffn(cur,
+                layer.ffn_gate_inp,
+                layer.ffn_up_exps,
+                layer.ffn_gate_exps,
+                layer.ffn_down_exps,
+                exp_probs_b,
+                n_expert, hparams.n_expert_used(),
+                LLM_FFN_SILU, hparams.expert_weights_norm && hparams.n_expert_used() > 1,
+                hparams.expert_weights_scale,
+                (llama_expert_gating_func_type) hparams.expert_gating_func,
+                il, logits);
+        cb(moe_out, "ffn_moe_out", il);
+
+        ggml_tensor * ffn_shexp = build_ffn(cur,
+                layer.ffn_up_shexp, nullptr, nullptr,
+                layer.ffn_gate_shexp, nullptr, nullptr,
+                layer.ffn_down_shexp, nullptr, nullptr,
+                nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
+        cb(ffn_shexp, "ffn_shexp", il);
+
+        cur = ggml_add(ctx0, moe_out, ffn_shexp);
+        cb(cur, "ffn_out", il);
+
+        inpL = build_hc_post(cur, residual, post, comb, il);
+        inpL = build_cvec(inpL, il);
+        cb(inpL, "l_out", il);
+    }
+
+    // no hc_head: collapse with the last FFN's pre-mix
+    if (inp_out_ids) {
+        inpL     = ggml_reshape_3d(ctx0, ggml_get_rows(ctx0, ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens), inp_out_ids), n_embd, hc, inp_out_ids->ne[0]);
+        pre_prev = ggml_get_rows(ctx0, pre_prev, inp_out_ids);
+    }
+    cur = build_hc_pre(inpL, pre_prev, -1);
+    cb(cur, "hc_out", -1);
+
+    cur = build_norm(cur, model.output_norm, nullptr, 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;
+}

src/models/dflash.cpp

@@ -81,6 +81,7 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
         ml.get_key(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters);
         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);
+        ml.get_key(LLM_KV_HYPER_CONNECTION_SINGLE_PASS,         hparams.dsv4_hc_single_pass, false);
 
         GGML_ASSERT(hparams.dsv4_o_group_count > 0); // avoid div by zero
 
@@ -94,6 +95,11 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
         }
 
         GGML_ASSERT(hparams.n_swa > 0);
+        if (hparams.dsv4_hc_single_pass) {
+            // DeepSeek-V4.1 DSpark: block queries see the main-KV ring plus the block via the mask row of the
+            // first query, so that row's SWA window must reach window_size positions before the anchor
+            hparams.n_swa += 1;
+        }
         hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
         hparams.set_swa_pattern(0);
         for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {
@@ -257,9 +263,12 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
         const int64_t hc_dim          = hc_mult * n_embd;
         const int64_t hc_mix_dim      = (2 + hc_mult) * hc_mult;
 
-        hc_head_fn    = create_tensor(tn(LLM_TENSOR_HC_HEAD_FN,    "weight"), {hc_dim, hc_mult}, 0);
-        hc_head_base  = create_tensor(tn(LLM_TENSOR_HC_HEAD_BASE,  "weight"), {hc_mult}, 0);
-        hc_head_scale = create_tensor(tn(LLM_TENSOR_HC_HEAD_SCALE, "weight"), {1}, 0);
+        // single-pass mHC: collapse with the last FFN pre-mix, no hc_head
+        if (!hparams.dsv4_hc_single_pass) {
+            hc_head_fn    = create_tensor(tn(LLM_TENSOR_HC_HEAD_FN,    "weight"), {hc_dim, hc_mult}, 0);
+            hc_head_base  = create_tensor(tn(LLM_TENSOR_HC_HEAD_BASE,  "weight"), {hc_mult}, 0);
+            hc_head_scale = create_tensor(tn(LLM_TENSOR_HC_HEAD_SCALE, "weight"), {1}, 0);
+        }
 
         for (int i = 0; i < n_layer; ++i) {
             auto & layer = layers[i];
@@ -1280,12 +1289,260 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
     }
 }
 
+// DeepSeek-V4.1 DSpark decoder (mtp.0..2), dual-mode by batch type like graph_dsv4:
+//   * embd batch  -> writes every stage's kv_norm(wkv(main_proj(target features))) into its window ring
+//   * token batch -> [anchor, noise...] block through the stages, single-pass mHC, no per-head q norm
+llama_model_dflash::graph_dsv41::graph_dsv41(const llama_model & model, const llm_graph_params & params) :
+    llama_model_deepseek41::graph(params) {
+    const int64_t n_embd_inp       = hparams.n_embd_inp_enc();
+    const int64_t n_embd_head      = hparams.n_embd_head_k();
+    const int64_t n_embd_head_rope = hparams.n_rot();
+    const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;
+
+    GGML_ASSERT(hparams.dsv4_hc_single_pass);
+
+    ggml_tensor * inp_pos = build_inp_pos();
+
+    llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa();
+
+    // ratio-0 layers: plain rope at rope_theta, no YaRN (original_seq_len 0)
+    const auto rope = [&](ggml_tensor * t) {
+        t = ggml_rope_ext(ctx0, t, inp_pos, nullptr, n_embd_head_rope, rope_type, 0,
+                freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+        return ggml_rope_set_offset(t, n_embd_head_nope);
+    };
+
+    ASSERT_EMBD_OR_TOKEN(ubatch);
+    if (ubatch.embd) {
+        auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
+
+        inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, n_tokens);
+        ggml_set_input(inp->embd);
+
+        ggml_tensor * inp_target = inp->embd;
+        cb(inp_target, "inp_target_features", -1);
+
+        res->add_input(std::move(inp));
+
+        ggml_tensor * inp_g = build_lora_mm(model.fc, inp_target, model.fc_s);
+        inp_g = build_norm(inp_g, model.output_norm_enc, nullptr, LLM_NORM_RMS, -1);
+        cb(inp_g, "main_x", -1);
+
+        for (int il = 0; il < n_layer; ++il) {
+            const auto & layer = model.layers[il];
+
+            ggml_tensor * kv = build_lora_mm(layer.wkv, inp_g);
+            kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il);
+            kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, n_tokens);
+            kv = rope(kv);
+            cb(kv, "kv_injected", il);
+
+            if (inp_attn->self_k_rot_swa) {
+                kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa);
+            }
+            ggml_build_forward_expand(gf, inp_attn->mctx->get_swa()->cpy_k(ctx0, kv, inp_attn->get_k_idxs_swa(), il));
+        }
+
+        res->t_embd = inp_g;
+
+        ggml_build_forward_expand(gf, inp_g);
+        return;
+    }
+
+    auto * tok_embd = model.tok_embd;
+    if (tok_embd == nullptr) {
+        GGML_ASSERT(cparams.ctx_other != nullptr);
+        const auto * model_other = llama_get_model(cparams.ctx_other);
+        GGML_ASSERT(model_other->tok_embd != nullptr && "DSpark decoder requires the target model's token embeddings");
+        tok_embd = model_other->tok_embd;
+    }
+
+    auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
+
+    inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+    ggml_set_input(inp->tokens);
+
+    ggml_tensor * inp_tokens = inp->tokens;
+
+    ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens);
+    cb(inpL, "inp_noise_embd", -1);
+
+    res->add_input(std::move(inp));
+
+    const int64_t hc            = hparams.dsv4_hc_mult;
+    const int64_t n_groups      = hparams.dsv4_o_group_count;
+    const int64_t n_heads_group = n_head / n_groups;
+    const int64_t o_lora_rank   = hparams.dsv4_o_lora_rank;
+    const int64_t o_group_dim   = n_heads_group*n_embd_head;
+    const float   kq_scale      = 1.0f/sqrtf(float(n_embd_head));
+
+    GGML_ASSERT(n_embd_head == n_embd_head_v);
+    GGML_ASSERT(n_head % n_groups == 0);
+
+    inpL = ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens);
+    inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1);
+    cb(inpL, "hc_init", -1);
+
+    // window of a block = window of its first query, see n_swa
+    {
+        ggml_tensor * m = inp_attn->self_kq_mask_swa_cnv;
+        const int64_t n_blocks = ubatch.n_seqs_unq;
+        const int64_t bd       = n_blocks > 0 ? n_tokens / n_blocks : 0;
+        if (m && bd > 1 && n_tokens % n_blocks == 0 && m->ne[2] == 1 && m->ne[1] % bd == 0 && ggml_is_contiguous(m)) {
+            ggml_tensor * m4 = ggml_reshape_4d(ctx0, m, m->ne[0], bd, m->ne[1]/bd, m->ne[3]);
+            ggml_tensor * r0 = ggml_view_4d(ctx0, m4, m4->ne[0], 1, m4->ne[2], m4->ne[3], m4->nb[1], m4->nb[2], m4->nb[3], 0);
+            m4 = ggml_repeat(ctx0, r0, m4);
+            inp_attn->self_kq_mask_swa_cnv = ggml_reshape_4d(ctx0, m4, m->ne[0], m->ne[1], 1, m->ne[3]);
+            cb(inp_attn->self_kq_mask_swa_cnv, "dspark_block_mask", -1);
+        }
+    }
+
+    ggml_tensor * pre_prev = nullptr;
+
+    for (int il = 0; il < n_layer; ++il) {
+        const auto & layer = model.layers[il];
+
+        ggml_tensor * residual = inpL;
+        ggml_tensor * post = nullptr;
+        ggml_tensor * comb = nullptr;
+
+        ggml_tensor * attn_pre = build_hc_mixes(inpL, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, &post, &comb, false, il);
+
+        ggml_tensor * cur;
+        if (pre_prev == nullptr) {
+            // identity pre-mix of the first stage: copy 0
+            cur = ggml_cont(ctx0, ggml_view_2d(ctx0, inpL, n_embd, n_tokens, inpL->nb[2], 0));
+        } else {
+            cur = build_hc_pre(inpL, pre_prev, il);
+        }
+        cb(cur, "hc_attn_pre", il);
+
+        cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+        cb(cur, "attn_norm", il);
+
+        // non-causal draft context: block queries see the main-KV ring plus every block key
+        {
+            ggml_tensor * qr = build_lora_mm(layer.wq_a, cur);
+            qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
+            cb(qr, "qr_norm", il);
+
+            ggml_tensor * q = build_lora_mm(layer.wq_b, qr);
+            q = ggml_reshape_3d(ctx0, q, n_embd_head, n_head, n_tokens);
+            q = rope(q);
+            cb(q, "q", il);
+
+            ggml_tensor * kv = build_lora_mm(layer.wkv, cur);
+            kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il);
+            kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, n_tokens);
+            kv = rope(kv);
+            cb(kv, "kv", il);
+
+            ggml_tensor * out = build_attn(inp_attn,
+                    nullptr, nullptr, nullptr,
+                    q, kv, kv,
+                    nullptr, layer.attn_sinks, nullptr,
+                    kq_scale, il);
+            cb(out, "attn_raw", il);
+
+            out = ggml_reshape_3d(ctx0, out, n_embd_head, n_head, n_tokens);
+            out = ggml_rope_ext_back(ctx0, out, inp_pos, nullptr, n_embd_head_rope, rope_type, 0,
+                    freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+            out = ggml_rope_set_offset(out, n_embd_head_nope);
+            cb(out, "attn_derope", il);
+
+            out = ggml_reshape_3d(ctx0, out, o_group_dim, n_groups, n_tokens);
+            out = ggml_permute(ctx0, out, 0, 2, 1, 3);
+            ggml_tensor * oa = ggml_mul_mat(ctx0, layer.wo_a, out);
+            cb(oa, "attn_wo_a", il);
+            oa = ggml_permute(ctx0, oa, 0, 2, 1, 3);
+            oa = ggml_cont_2d(ctx0, oa, o_lora_rank*n_groups, n_tokens);
+
+            cur = build_lora_mm(layer.wo_b, oa);
+            cb(cur, "attn_out", il);
+        }
+
+        inpL = build_hc_post(cur, residual, post, comb, il);
+        cb(inpL, "hc_attn_post", il);
+
+        residual = inpL;
+        ggml_tensor * ffn_pre = build_hc_mixes(inpL, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, &post, &comb, false, il);
+
+        cur = build_hc_pre(inpL, attn_pre, il);
+        cb(cur, "hc_ffn_pre", il);
+
+        cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
+        cb(cur, "ffn_norm", il);
+
+        ggml_tensor * moe_out = build_moe_ffn(cur,
+                layer.ffn_gate_inp,
+                layer.ffn_up_exps,
+                layer.ffn_gate_exps,
+                layer.ffn_down_exps,
+                layer.ffn_exp_probs_b,
+                n_expert, hparams.n_expert_used(),
+                LLM_FFN_SILU, hparams.expert_weights_norm,
+                hparams.expert_weights_scale,
+                (llama_expert_gating_func_type) hparams.expert_gating_func,
+                il);
+        cb(moe_out, "ffn_moe_out", il);
+
+        ggml_tensor * ffn_shexp = build_ffn(cur,
+                layer.ffn_up_shexp, nullptr, nullptr,
+                layer.ffn_gate_shexp, nullptr, nullptr,
+                layer.ffn_down_shexp, nullptr, nullptr,
+                nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
+        cb(ffn_shexp, "ffn_shexp", il);
+
+        cur = ggml_add(ctx0, moe_out, ffn_shexp);
+        cb(cur, "ffn_out", il);
+
+        inpL = build_hc_post(cur, residual, post, comb, il);
+
+        ggml_tensor * l_out = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
+        cb(l_out, "l_out", il);
+        ggml_build_forward_expand(gf, l_out);
+
+        pre_prev = ffn_pre;
+    }
+
+    // no hc_head: collapse with the last FFN's pre-mix
+    ggml_tensor * cur = build_hc_pre(inpL, pre_prev, -1);
+    cb(cur, "hc_out", -1);
+
+    // confidence head takes the un-normed collapsed state
+    res->t_embd = cur;
+
+    cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
+    cb(cur, "result_norm", -1);
+
+    auto * output   = model.output;
+    auto * output_s = model.output_s;
+    if (output == nullptr) {
+        GGML_ASSERT(cparams.ctx_other != nullptr);
+        const auto * model_other = llama_get_model(cparams.ctx_other);
+        GGML_ASSERT(model_other->output != nullptr && "DSpark decoder requires the target model's output projection");
+        output   = model_other->output;
+        output_s = model_other->output_s;
+    }
+
+    cur = build_lora_mm(output, cur, output_s);
+    cb(cur, "result_output", -1);
+    res->t_logits = cur;
+
+    if (model.dspark_markov_w1) {
+        build_dspark_markov_head(*this, model, inp_tokens);
+    }
+}
+
 std::unique_ptr<llm_graph_context> llama_model_dflash::build_arch_graph(const llm_graph_params & params) const {
     switch (params.gtype) {
         case LLM_GRAPH_TYPE_ENCODER:
             return std::make_unique<graph<true>>(*this, params);
         case LLM_GRAPH_TYPE_DEFAULT:
         case LLM_GRAPH_TYPE_DECODER:
+            if (hparams.dsv4_hc_mult > 0 && hparams.dsv4_hc_single_pass) {
+                return std::make_unique<graph_dsv41>(*this, params);
+            }
             if (hparams.dsv4_hc_mult > 0) {
                 return std::make_unique<graph_dsv4>(*this, params);
             }

src/models/minimax-m3.cpp

@@ -550,9 +550,9 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
         if ((uint32_t) il < hparams.n_layer_dense_lead) {
             // leading dense FFN (swigluoai)
             cur = build_ffn(cur,
-                    model.layers[il].ffn_up,   NULL, NULL,
-                    model.layers[il].ffn_gate, NULL, NULL,
-                    model.layers[il].ffn_down, NULL, NULL,
+                    model.layers[il].ffn_up,   NULL, model.layers[il].ffn_up_s,
+                    model.layers[il].ffn_gate, NULL, model.layers[il].ffn_gate_s,
+                    model.layers[il].ffn_down, NULL, model.layers[il].ffn_down_s,
                     NULL,
                     LLM_FFN_SWIGLU_OAI_MOE, LLM_FFN_PAR, il);
             cb(cur, "ffn_out", il);
@@ -568,14 +568,17 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
                     LLM_FFN_SWIGLU_OAI_MOE, hparams.expert_weights_norm,
                     hparams.expert_weights_scale,
                     (llama_expert_gating_func_type) hparams.expert_gating_func,
-                    il);
+                    il, nullptr, nullptr,
+                    model.layers[il].ffn_up_exps_s,
+                    model.layers[il].ffn_gate_exps_s,
+                    model.layers[il].ffn_down_exps_s);
             cb(moe_out, "ffn_moe_out", il);
 
             // shared expert (swigluoai)
             ggml_tensor * ffn_shexp = build_ffn(cur,
-                    model.layers[il].ffn_up_shexp,   NULL, NULL,
-                    model.layers[il].ffn_gate_shexp, NULL, NULL,
-                    model.layers[il].ffn_down_shexp, NULL, NULL,
+                    model.layers[il].ffn_up_shexp,   NULL, model.layers[il].ffn_up_shexp_s,
+                    model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s,
+                    model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s,
                     NULL,
                     LLM_FFN_SWIGLU_OAI_MOE, LLM_FFN_PAR, il);
             cb(ffn_shexp, "ffn_shexp", il);

src/models/models.h

@@ -1331,6 +1331,15 @@ struct llama_model_dots3note : public llama_model_base {
     std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
 };
 
+// YaRN attention factor that cancels ggml's rope mscale (the DeepSeek-V4 reference applies none)
+static inline float dsv4_rope_attn_factor(float freq_scale, float ext_factor) {
+    if (ext_factor == 0.0f) {
+        return 1.0f;
+    }
+
+    return 1.0f / (1.0f + 0.1f*logf(1.0f/freq_scale));
+}
+
 struct llama_model_deepseek4 : public llama_model_base {
     llama_model_deepseek4(const struct llama_model_params & params) : llama_model_base(params) {}
     void load_arch_hparams(llama_model_loader & ml) override;
@@ -1359,6 +1368,17 @@ struct llama_model_deepseek4 : public llama_model_base {
                 ggml_tensor ** comb,
                 int il) const;
 
+        // returns the pre weights and fills post / comb
+        ggml_tensor * build_hc_mixes(
+                ggml_tensor * x,
+                ggml_tensor * hc_fn,
+                ggml_tensor * hc_scale,
+                ggml_tensor * hc_base,
+                ggml_tensor ** post,
+                ggml_tensor ** comb,
+                bool f32_src,
+                int il) const;
+
         ggml_tensor * build_hc_post(
                 ggml_tensor * x,
                 ggml_tensor * residual,
@@ -1473,6 +1493,57 @@ struct llama_model_deepseek4 : public llama_model_base {
 };
 
 
+struct llama_model_deepseek41 : public llama_model_deepseek4 {
+    llama_model_deepseek41(const struct llama_model_params & params) : llama_model_deepseek4(params) {}
+    void load_arch_hparams(llama_model_loader & ml) override;
+    void load_arch_tensors(llama_model_loader & ml) override;
+
+    // Engram: the hash constants are hparams.ple_* (one module per Engram layer), hashed on the host (ggml has no
+    // int64 mul/xor/mod); the token map and its pad id are V4.1 only
+    uint32_t engram_pad_id = 0;
+    std::vector<int32_t> engram_token_map;
+
+    int engram_index(int il) const;
+
+    void engram_hash(int e, const llama_ubatch & ubatch, const std::vector<llama_token> & prev, std::vector<int32_t> & out) const;
+
+    struct graph : public llama_model_deepseek4::graph {
+        graph(const llama_model & model, const llm_graph_params & params);
+        // helpers only (no graph built): base of graph_dsv41
+        explicit graph(const llm_graph_params & params) : llama_model_deepseek4::graph(params) {}
+
+        ggml_tensor * build_engram(
+                const llama_model_deepseek41 & model,
+                ggml_tensor * x,
+                ggml_tensor * rows,
+                ggml_tensor * keep,
+                int il) const;
+
+        ggml_tensor * build_attention41(
+                const llama_model_deepseek41 & model,
+                llm_graph_input_dsv4 * inp_dsv4,
+                ggml_tensor * cur,
+                ggml_tensor * inp_pos,
+                ggml_tensor * cand_pin,
+                int il);
+
+        std::vector<ggml_tensor *> topk_of; // top-k of each index source layer
+        std::vector<ggml_tensor *> comp_row_of; // [n_layer] new compressed rows of each kv source (for mirrors)
+        ggml_tensor * cand_sel    = nullptr; // candidate blocks of the candidate source layer
+        int64_t       cand_n_comp = 0;       // compressed rows the candidate source scored
+
+        // sparse path (LLAMA_DSV41_SPARSE=0 disables): fused indexer, FA bounded to window + top-k, decode gather
+        bool sparse = true;
+        // gather path, shared by layers with the same top-k: selected K rows per (storage layer, top-k source),
+        // the sorted top-k and the mask per top-k source (one input copy per split)
+        std::map<std::pair<int32_t, int32_t>, ggml_tensor *> gather_k_of;
+        std::map<int32_t, ggml_tensor *> gather_sel_of;
+        std::map<int32_t, ggml_tensor *> gather_mask_of;
+    };
+
+    std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
+};
+
 struct llama_model_deepseek2ocr : public llama_model_base {
     llama_model_deepseek2ocr(const struct llama_model_params & params) : llama_model_base(params) {}
     void load_arch_hparams(llama_model_loader & ml) override;
@@ -1532,6 +1603,11 @@ struct llama_model_dflash : public llama_model_base {
         graph_dsv4(const llama_model & model, const llm_graph_params & params);
     };
 
+    // DeepSeek-V4.1 DSpark: single-pass mHC (lagged pre-mix, no hc_head), no per-head q norm
+    struct graph_dsv41 : public llama_model_deepseek41::graph {
+        graph_dsv41(const llama_model & model, const llm_graph_params & params);
+    };
+
     std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
 };
 

src/models/qwen4exp.cpp

@@ -115,6 +115,9 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
 
         hparams.ple_n_heads  = (hparams.ple_ngram_size - 1) * hparams.ple_heads_per_ngram;
         hparams.ple_head_dim = hparams.n_embd_per_layer;
+        hparams.ple_embd_token = hparams.ple_image_token_id != 0
+            ? (llama_token) hparams.ple_image_token_id
+            : (llama_token) hparams.ple_eos_token_id;
         if (hparams.ple_ngram_size < 2 || hparams.ple_ngram_size > LLAMA_MAX_PLE_NGRAM) {
             throw std::runtime_error(format("PLE n-gram size %u is out of range", hparams.ple_ngram_size));
         }

tests/fusion/MTL.csv

@@ -49,6 +49,8 @@ deepseek32          ,0   ,any     ,NORM+MUL+ADD                ,      2
 deepseek32          ,0   ,any     ,RMS_NORM+MUL                ,      9
 deepseek4           ,0   ,any     ,MUL+ADD                     ,      8
 deepseek4           ,0   ,any     ,RMS_NORM+MUL                ,     20
+deepseek41          ,0   ,any     ,MUL+ADD                     ,     12
+deepseek41          ,0   ,any     ,RMS_NORM+MUL                ,     29
 dots1               ,0   ,any     ,ADD+ADD                     ,      1
 dots1               ,0   ,any     ,MUL+ADD                     ,      2
 dots1               ,0   ,any     ,RMS_NORM+MUL                ,      9

tests/test-backend-ops.cpp

@@ -2798,11 +2798,11 @@ struct test_rope_set_rows : public test_case {
     const ggml_type type_idx;
     const std::array<int64_t, 4> ne_a;
     int mode;
+    const int n_dims;
     const int n_ctx{512};
-    const int n_dims{128};
 
     std::string vars() override {
-        return VARS_TO_STR4(type, type_idx, ne_a, mode);
+        return VARS_TO_STR5(type, type_idx, ne_a, mode, n_dims);
     }
 
     std::string op_desc(ggml_tensor * t) override {
@@ -2815,8 +2815,9 @@ struct test_rope_set_rows : public test_case {
     test_rope_set_rows(ggml_type type,
             ggml_type type_idx,
             std::array<int64_t, 4> ne_a,
-            int mode)
-        : type(type), type_idx(type_idx), ne_a(ne_a), mode(mode) {}
+            int mode,
+            int n_dims = 128)
+        : type(type), type_idx(type_idx), ne_a(ne_a), mode(mode), n_dims(n_dims) {}
 
     ggml_tensor * build_graph(ggml_context * ctx) override {
         ggml_tensor * a = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne_a[0], ne_a[1], ne_a[2], 1);
@@ -2850,7 +2851,7 @@ struct test_rope_set_rows : public test_case {
                 rope = ggml_rope_multi(ctx, a, pos, freq, n_dims, rope_sections, mode, 0, 10000.0f, fs, ef, af, 1.0f, 1.0f);
             }
         } else {
-            rope = ggml_rope(ctx, a, pos, ne_a[0], mode);
+            rope = ggml_rope(ctx, a, pos, n_dims, mode);
         }
 
         ggml_tensor * view = ggml_view_2d(ctx, rope, ne_a[0] * ne_a[1], ne_a[2], rope->nb[2], 0);
@@ -5617,14 +5618,9 @@ static void init_mul_mat_id_ids(ggml_context * ctx, int n_mats, bool skip = fals
     }
 }
 
-static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats, bool skip = false, float amax = 1.0f) {
+static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats, bool skip = false) {
     for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
-        if (t->type == GGML_TYPE_I32) {
-            continue;
-        }
-        if (amax != 1.0f && t->type == GGML_TYPE_F32) {
-            init_tensor_uniform(t, -amax, amax);
-        } else {
+        if (t->type != GGML_TYPE_I32) {
             init_tensor_uniform(t);
         }
     }
@@ -5641,13 +5637,12 @@ struct test_mul_mat_id : public test_case {
     const int64_t m;
     const int64_t n;
     const int64_t k;
-    const float amax; // magnitude of src1
     const int64_t m_v; // rows of as in memory, the experts of as are strided for m_v > m, no view for m_v == 0
     const bool ids_offset;
     const bool skip; // set some ids to a negative value
 
     std::string vars() override {
-        return VARS_TO_STR12(type_a, type_b, n_mats, n_used, b, m, n, k, amax, m_v, ids_offset, skip);
+        return VARS_TO_STR11(type_a, type_b, n_mats, n_used, b, m, n, k, m_v, ids_offset, skip);
     }
 
     double max_nmse_err() override {
@@ -5672,9 +5667,9 @@ struct test_mul_mat_id : public test_case {
     test_mul_mat_id(ggml_type type_a = GGML_TYPE_F32, ggml_type type_b = GGML_TYPE_F32,
             int n_mats = 8, int n_used = 2, bool b = false,
             int64_t m = 32, int64_t n = 32, int64_t k = 32,
-            float amax = 1.0f, int64_t m_v = 0, bool ids_offset = false, bool skip = false)
+            int64_t m_v = 0, bool ids_offset = false, bool skip = false)
         : type_a(type_a), type_b(type_b), n_mats(n_mats), n_used(n_used), b(b),
-            m(m), n(n), k(k), amax(amax), m_v(m_v), ids_offset(ids_offset), skip(skip) {
+            m(m), n(n), k(k), m_v(m_v), ids_offset(ids_offset), skip(skip) {
             GGML_ASSERT(n_used <= n_mats);
             GGML_ASSERT(m_v == 0 || m_v > m);
         }
@@ -5700,16 +5695,11 @@ struct test_mul_mat_id : public test_case {
         ggml_tensor * out = ggml_mul_mat_id(ctx, as, b, ids);
         ggml_set_name(out, "out");
 
-        if (amax > 65504.0f) {
-            // src1 exceeds F16 range
-            ggml_prec_set_src(out, GGML_PREC_F32, 1);
-        }
-
         return out;
     }
 
     void initialize_tensors(ggml_context * ctx) override {
-        init_mul_mat_id_tensors(ctx, n_mats, skip, amax);
+        init_mul_mat_id_tensors(ctx, n_mats, skip);
     }
 
     void reinit_perf_iter(ggml_context * ctx) override {
@@ -10244,6 +10234,12 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         }
     }
     test_cases.emplace_back(new test_rope_set_rows(GGML_TYPE_F32, GGML_TYPE_I32, { 128, 32, 8, 1 }, GGML_ROPE_TYPE_IMROPE));
+    // partial rotation
+    for (int mode : { GGML_ROPE_TYPE_NORMAL, GGML_ROPE_TYPE_NEOX, GGML_ROPE_TYPE_MROPE, GGML_ROPE_TYPE_IMROPE }) {
+        for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) {
+            test_cases.emplace_back(new test_rope_set_rows(type, GGML_TYPE_I64, { 256, 4, 8, 1 }, mode, 64));
+        }
+    }
 
     for (ggml_type type_input : {GGML_TYPE_F32}) {
         for (ggml_op_pool pool_type : {GGML_OP_POOL_AVG, GGML_OP_POOL_MAX}) {
@@ -10294,6 +10290,8 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
 
     // im2col 1D
     test_cases.emplace_back(new test_im2col(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, {3000, 128, 1, 1}, {3, 128, 1280, 1}, 1, 0, 1, 0, 1, 0, false));
+    // deepseek4v / deepseek41v aligner: 3x3 stride-3 unfold (F.unfold == im2col)
+    test_cases.emplace_back(new test_im2col(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, {48, 33, 1024, 1}, {3, 3, 1024, 1}, 3, 3, 0, 0, 1, 1, true));
     test_cases.emplace_back(new test_im2col(GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_F32, {3000, 128, 1, 1}, {3, 128, 1280, 1}, 1, 0, 1, 0, 1, 0, false));
     test_cases.emplace_back(new test_im2col(GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_F16, {3000, 128, 1, 1}, {3, 128, 1280, 1}, 1, 0, 1, 0, 1, 0, false));
     test_cases.emplace_back(new test_im2col(GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_F16, {3000, 384, 1, 1}, {3, 384,  384, 1}, 1, 0, 1, 0, 1, 0, false));
@@ -10502,6 +10500,8 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     };
 
     for (ggml_type kernel_type : {GGML_TYPE_F32, GGML_TYPE_F16}) {
+        test_cases.emplace_back(new test_conv_3d(1, 4, 6, 78, 54, 512, 3, 3, 3, 1, 1, 1, 0, 0, 0, 1, 1, 1, kernel_type));
+        test_cases.emplace_back(new test_conv_3d(2, 16, 4, 64, 64, 256, 1, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, kernel_type));
         for (int N : {1, 2}) {
             for (int IC : {1, 3}) {
                 for (int OC : {1, 4}) {
@@ -11337,8 +11337,8 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 8, 1, 64, {8, 1}, {1, 1}, {0, 1, 2, 3}, 0, 1, false, 16, 16));
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 8, 16, 64, {8, 1}, {1, 1}, {0, 1, 2, 3}, 0, 1, false, 16, 16));
     // as is a view whose experts are strided by more rows than it uses
-    test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F32, GGML_TYPE_F32, 4, 2, false, 8,  1, 64, 1.0f, 64));
-    test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 4, 2, false, 8, 16, 64, 1.0f, 64));
+    test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F32, GGML_TYPE_F32, 4, 2, false, 8,  1, 64, 64));
+    test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 4, 2, false, 8, 16, 64, 64));
 
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 576, 512, 576, {1,1}, {1,1}));
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 1, 2048, 8192, {1,  1}, {1, 1}));
@@ -11452,7 +11452,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     }
 
     for (ggml_type ta : {GGML_TYPE_F16, GGML_TYPE_Q4_0}) {
-        test_cases.emplace_back(new test_mul_mat_id_dup(ta, GGML_TYPE_F32, 1025, 10, false, 64, 33, 256, 1.0f, 0, true));
+        test_cases.emplace_back(new test_mul_mat_id_dup(ta, GGML_TYPE_F32, 1025, 10, false, 64, 33, 256, 0, true));
         test_cases.emplace_back(new test_mul_mat_id_dup(ta, GGML_TYPE_F32, 2050, 10, true, 64, 33, 256));
     }
 #endif
@@ -11467,12 +11467,6 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 4, 4, false, 16, 10, 256));
     }
 
-    // test src1 f16 overflow
-    for (int n : {16, 32, 64}) {
-        test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q4_K, GGML_TYPE_F32, 128, 4, false, 4096, n, 2048, 1e5f));
-        test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q8_0, GGML_TYPE_F32, 8,   2, false, 512,  n, 256,  1e5f));
-    }
-
     for (ggml_type type_a : base_types) {
         for (ggml_type type_b : {GGML_TYPE_F32 /*, GGML_TYPE_F16 */}) {
             for (int n_mats : {4, 8}) {
@@ -11510,14 +11504,14 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         for (int n_used : {1, 2, 4}) {
             for (bool b : {false, true}) {
                 for (int n : {1, 4, 33, 129}) {
-                    test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 8, n_used, b, 512, n, 256, 1.0f, 0, false, true));
+                    test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 8, n_used, b, 512, n, 256, 0, false, true));
                 }
             }
         }
     }
     // n_used == n_mats, so every expert can be skipped
-    test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 4, 4, false, 512, 64, 256, 1.0f, 0, false, true));
-    test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 1, 1, false, 8, 16, 1, 1.0f, 0, false, true));
+    test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 4, 4, false, 512, 64, 256, 0, false, true));
+    test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 1, 1, false, 8, 16, 1, 0, false, true));
     test_cases.emplace_back(new test_mul_mat_id_fusion(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, false, 32, 32, 32, 3, false, true));
 
     for (int bs : {1, 4, 512}) {
@@ -12068,6 +12062,17 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         }
     }
 
+    // KV-head counts that do not divide the device count
+    for (int hs : { 128, 256, 512, }) {
+        for (int nh : { 1, 2, 6, }) {
+            for (int nr2 : { 4, 8, }) {
+                for (int nb : { 75, 1024, }) {
+                    test_cases.emplace_back(new test_flash_attn_ext(hs, hs, nh, {nr2, 1}, 1024, nb, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+                }
+            }
+        }
+    }
+
     for (int hsk : { 40, 64, 72, 80, 96, 128, 192, 256, 320, 512, 576 }) {
         for (int hsv : { 40, 64, 72, 80, 96, 128, 192, 256, 512 }) {
             if (hsk != 96 && hsk != 192 && hsk != 320 && hsk != 576 && hsk != hsv) continue;
@@ -12360,6 +12365,13 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
             false, 16, 8, false, false, true, false, { 1, 1 }));
     }
 
+    // Fused gate/up + GLU on the MMQ path, only used if its tiling is efficient: ~1280 tiles per case, one case per type and GLU op.
+    test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_Q4_K,   GGML_GLU_OP_SWIGLU,        17, 81920, 256, false, 1, 1, false, false, true, false, { 1, 1 }));
+    test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_Q5_K,   GGML_GLU_OP_GEGLU,        150, 40960, 256, false, 1, 1, false, false, true, false, { 1, 1 }));
+    test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_Q6_K,   GGML_GLU_OP_SWIGLU_OAI,   150, 40960, 256, false, 1, 1, false, false, true, false, { 1, 1 }));
+    test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_Q8_0,   GGML_GLU_OP_SWIGLU_CLAMP, 150,  6912, 256, false, 1, 1, false, false, true, false, { 2, 3 }));
+    test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_IQ4_XS, GGML_GLU_OP_SWIGLU,       150, 40960, 512, false, 1, 1, false, false, true, false, { 1, 1 }));
+
     for (auto gate : {GATING_FUNC_SOFTMAX, GATING_FUNC_SIGMOID, GATING_FUNC_SOFTMAX_WEIGHT, GATING_FUNC_SQRT_SOFTPLUS}) {
         for (bool with_norm : {false, true}) {
             for (bool bias_probs : {false, true}) {
@@ -12597,6 +12609,15 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
         {1,  320, 8,  76,  52, 1280, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1},
         {1, 1280, 8,  76,  52, 1280, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1},
         {1,  320, 8, 152, 104, 1280, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1},
+        // Wan VAE pads spatial borders and temporal history before conv3d; operator padding is zero.
+        {1,  384, 3,  40,  28,  384, 3, 3, 3, 1, 1, 1, 0, 0, 0, 1, 1, 1},
+        {1,  192, 4,  78,  54,  192, 3, 3, 3, 1, 1, 1, 0, 0, 0, 1, 1, 1},
+        {1,   96, 6, 154, 106,   96, 3, 3, 3, 1, 1, 1, 0, 0, 0, 1, 1, 1},
+        {1,   96, 6, 306, 210,    3, 3, 3, 3, 1, 1, 1, 0, 0, 0, 1, 1, 1},
+        {1, 1024, 3,  40,  28, 1024, 3, 3, 3, 1, 1, 1, 0, 0, 0, 1, 1, 1},
+        {1,  256, 6, 154, 106,  256, 3, 3, 3, 1, 1, 1, 0, 0, 0, 1, 1, 1},
+        {1,  384, 3,  38,  26,  768, 3, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1},
+        {1,  192, 2,  76,  52,   96, 1, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1},
 #if 0
         // too slow on some devices
         {1, 1280, 8, 152, 104, 1280, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1},

tests/test-chat.cpp

@@ -10,6 +10,8 @@
 #include "../tools/server/server-chat.h"
 #include "chat-auto-parser.h"
 #include "chat.h"
+#include "json-schema-to-grammar.h"
+#include "json-schema.h"
 #include "common.h"
 #include "ggml.h"
 #include "log.h"
@@ -149,6 +151,21 @@ static common_chat_templates_ptr read_templates(const std::string & path) {
     return common_chat_templates_ptr(common_chat_templates_init(/* model= */ nullptr, read_file(path)));
 }
 
+// DeepSeek V4.1 stand-in: the real template ships in the GGUF; this keeps only the turn markers,
+// the generation prompt, the V4.1 DSML block tag used for detection and the extra keys it reads
+static const char * DEEPSEEK_V4_1_MIN_TEMPLATE = R"({%- for m in messages -%}
+{%- if m.role == 'system' -%}<｜System｜>{{ m.content }}
+{%- elif m.role == 'user' -%}<｜User｜>
+{%- if m.content is string -%}{{ m.content }}{%- else -%}{%- for b in m.content -%}{%- if not loop.first -%}{{ '\n\n' }}{%- endif -%}{{ b.text }}{%- endfor -%}{%- endif -%}
+{%- if m.task -%}<｜{{ m.task }}｜>{%- endif -%}
+{%- elif m.role == 'assistant' -%}<｜Assistant｜></think>{{ m.content }}
+{%- for tc in m.tool_calls or [] -%}<｜DSML｜ calls><｜DSML｜ invoke name="{% if tc.namespace %}{{ tc.namespace }}::{% endif %}{{ tc.function.name }}"></｜DSML｜ invoke></｜DSML｜ calls>{%- endfor -%}
+<｜end▁of▁sentence｜>
+{%- endif -%}
+{%- endfor -%}
+{%- if reasoning_effort is defined and reasoning_effort is not none -%}<｜effort={{ reasoning_effort }}｜>{%- endif -%}
+{%- if add_generation_prompt -%}<｜Assistant｜>{%- if enable_thinking -%}<think>{%- else -%}</think>{%- endif -%}{%- endif -%})";
+
 static std::unique_ptr<llama_grammar> build_grammar(const std::string & grammar_str) {
     return std::unique_ptr<llama_grammar>(
         llama_grammar_init_impl(nullptr, grammar_str.c_str(), "root", false, nullptr, 0, nullptr, 0));
@@ -961,6 +978,7 @@ const common_chat_msg message_user{
     /* .reasoning_content = */ "",
     /* .tool_name = */ "",
     /* .tool_call_id = */ "",
+    /* .extra_json = */ "",
 };
 
 const common_chat_msg message_user_parts{
@@ -977,6 +995,7 @@ const common_chat_msg message_user_parts{
     "",
     /* .tool_name = */ "",
     /* .tool_call_id = */ "",
+    /* .extra_json = */ "",
 };
 
 static common_chat_msg simple_assist_msg(const std::string & content,
@@ -1486,6 +1505,11 @@ class peg_tester {
         template_path_(template_path),
         detailed_debug_(detailed_debug) {}
 
+    peg_tester(const std::string & name, common_chat_templates_ptr tmpls, const bool detailed_debug = false) :
+        tmpls_(std::move(tmpls)),
+        template_path_(name),
+        detailed_debug_(detailed_debug) {}
+
     const std::string & template_path() const { return template_path_; }
 
     peg_test_builder test(const std::string & input);
@@ -4381,6 +4405,103 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .run();
     }
 
+    // DeepSeek V4.1: V4 behaviour, DSML tag names with a leading space
+    {
+        auto tst = peg_tester("deepseek-v4.1-min",
+                              common_chat_templates_ptr(common_chat_templates_init(nullptr, DEEPSEEK_V4_1_MIN_TEMPLATE)),
+                              detailed_debug);
+
+        tst.test("Hello, world!\nWhat's up?")
+            .enable_thinking(false)
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .expect(message_assist)
+            .run();
+
+        tst.test("I'm\nthinking</think>Hello, world!\nWhat's up?")
+            .enable_thinking(true)
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .expect(message_assist_thoughts)
+            .run();
+
+        tst.test(
+               "Let me check the time</think>\n\n"
+               "<｜DSML｜ calls>\n"
+               "<｜DSML｜ invoke name=\"get_time\">\n"
+               "<｜DSML｜ parameter name=\"city\" string=\"true\">Tokyo</｜DSML｜ parameter>\n"
+               "</｜DSML｜ invoke>\n"
+               "</｜DSML｜ calls>")
+            .enable_thinking(true)
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .tools({ get_time_tool })
+            .expect(message_with_tool_calls_and_reasoning("get_time", R"({"city": "Tokyo"})", "Let me check the time"))
+            .run();
+
+        tst.test(
+               "<｜DSML｜ calls>\n"
+               "<｜DSML｜ invoke name=\"special_function\">\n"
+               "<｜DSML｜ parameter name=\"arg1\" string=\"false\">1</｜DSML｜ parameter>\n"
+               "</｜DSML｜ invoke>\n"
+               "</｜DSML｜ calls>")
+            .enable_thinking(false)
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .tools({ special_function_tool })
+            .expect(message_assist_call)
+            .run();
+
+        tst.test(
+               "Calling both</think>\n\n"
+               "<｜DSML｜ calls>\n"
+               "<｜DSML｜ invoke name=\"get_time\">\n"
+               "<｜DSML｜ parameter name=\"city\" string=\"true\">Paris</｜DSML｜ parameter>\n"
+               "</｜DSML｜ invoke>\n"
+               "<｜DSML｜ invoke name=\"get_weather\">\n"
+               "<｜DSML｜ parameter name=\"city\" string=\"true\">Paris</｜DSML｜ parameter>\n"
+               "</｜DSML｜ invoke>\n"
+               "</｜DSML｜ calls>")
+            .enable_thinking(true)
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .parallel_tool_calls(true)
+            .tools({ get_time_tool, get_weather_tool })
+            .expect(message_with_reasoning_content_and_multiple_tool_calls(
+                "Calling both", "",
+                { { "get_time", R"({"city": "Paris"})" }, { "get_weather", R"({"city": "Paris"})" } }))
+            .run();
+
+        tst.test("world!\nWhat's up?")
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .enable_thinking(true)
+            .messages({ message_user, message_assist_prefill_content })
+            .add_generation_prompt(false)
+            .continue_final_message(COMMON_CHAT_CONTINUATION_CONTENT)
+            .expect_reasoning("I'm thinking")
+            .expect_content("Hello, world!\nWhat's up?")
+            .run();
+
+        tst.test(" thinking</think>Hello, world!\nWhat's up?")
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .enable_thinking(true)
+            .messages({ message_user, message_assist_prefill_reasoning })
+            .add_generation_prompt(false)
+            .continue_final_message(COMMON_CHAT_CONTINUATION_REASONING)
+            .expect_reasoning("I'm thinking")
+            .expect_content("Hello, world!\nWhat's up?")
+            .run();
+
+        tst.test(
+            "Let me check the time\n\n"
+            "<｜DSML｜ calls>\n"
+            "<｜DSML｜ invoke name=\"get_time\">\n"
+            "<｜DSML｜ parameter name=\"city\" string=\"true\">Tokyo</｜DSML｜ parameter>\n"
+            "</｜DSML｜ invoke>\n"
+            "</｜DSML｜ calls>") // no </think> after the TC close because the grammar will immediately constrain it to end
+            .enable_thinking(true)
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .tools({ get_time_tool })
+            .expect_reasoning("Let me check the time")
+            .expect_tool_calls({ { "get_time", R"({"city": "Tokyo"})", {} } })
+            .run();
+    }
+
     {
         // The DSML separator belongs to the tool call block, not assistant content.
         auto tst = peg_tester("models/templates/deepseek-ai-DeepSeek-V4-Flash-0731.jinja", detailed_debug);
@@ -7825,6 +7946,62 @@ static void test_deepseek_v4_tool_result_ordering() {
     }
 }
 
+// DeepSeek V4.1: extra message keys, "\n\n" around media markers, reasoning_effort mapping, system-message tools
+static void test_deepseek_v4_1_message_fields() {
+    LOG_DBG("%s\n", __func__);
+
+    common_chat_templates_ptr tmpls(common_chat_templates_init(nullptr, DEEPSEEK_V4_1_MIN_TEMPLATE));
+    auto render = [&](const char * messages_json, bool add_generation_prompt) {
+        common_chat_templates_inputs inputs;
+        inputs.messages = common_chat_msgs_parse_oaicompat(json::parse(messages_json));
+        inputs.add_generation_prompt = add_generation_prompt;
+        inputs.enable_thinking = false;
+        return common_chat_templates_apply(tmpls.get(), inputs).prompt.str();
+    };
+
+    assert_equals(std::string("<｜User｜>q<｜action｜>"),
+                  render(R"([{"role": "user", "content": "q", "task": "action"}])", false));
+
+    assert_contains(render(R"([
+            {"role": "user", "content": "q"},
+            {"role": "assistant", "content": "", "tool_calls": [{"type": "function", "namespace": "ns", "function": {"name": "f", "arguments": {"x": 1}}}]}
+        ])", false),
+        "<｜DSML｜ invoke name=\"ns::f\">");
+
+    assert_equals(std::string("<｜User｜>a\n\n<__media__>\n\nb<｜Assistant｜></think>"),
+                  render(R"([{"role": "user", "content": [{"type": "text", "text": "a"}, {"type": "media_marker", "text": "<__media__>"}, {"type": "text", "text": "b"}]}])", true));
+
+    for (const auto & [effort, expected] : std::vector<std::pair<std::string, std::string>>{
+             { "minimal", "<｜effort=low｜>" }, { "low", "<｜effort=low｜>" }, { "medium", "" },
+             { "high", "<｜effort=high｜>" }, { "xhigh", "<｜effort=max｜>" }, { "max", "<｜effort=max｜>" } }) {
+        common_chat_templates_inputs inputs;
+        inputs.messages = { message_user };
+        inputs.add_generation_prompt = true;
+        inputs.chat_template_kwargs["reasoning_effort"] = json(effort).dump();
+        assert_equals("<｜User｜>Hey there!" + expected + "<｜Assistant｜><think>",
+                      common_chat_templates_apply(tmpls.get(), inputs).prompt.str());
+    }
+
+    {
+        common_chat_templates_inputs inputs;
+        inputs.messages = common_chat_msgs_parse_oaicompat(json::parse(R"([
+            {"role": "system", "content": "s", "tools": [{"type": "function", "function": {"name": "get_time",
+                "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}]},
+            {"role": "user", "content": "q"}
+        ])"));
+        inputs.add_generation_prompt = true;
+        inputs.enable_thinking = false;
+        inputs.reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK;
+        common_chat_session session(tmpls.get(), nullptr, inputs, {});
+        const auto & msg = session.finish(common_chat_input(std::string(
+            "\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"get_time\">\n"
+            "<｜DSML｜ parameter name=\"city\" string=\"true\">Tokyo</｜DSML｜ parameter>\n</｜DSML｜ invoke>\n</｜DSML｜ calls>")));
+        assert_equals(size_t(1), msg.tool_calls.size());
+        assert_equals(std::string("get_time"), msg.tool_calls[0].name);
+        assert_equals(std::string("{\"city\":\"Tokyo\"}"), msg.tool_calls[0].arguments);
+    }
+}
+
 static void test_prompt_parts() {
     LOG_DBG("%s\n", __func__);
     const std::string tmpl_src =
@@ -7931,6 +8108,58 @@ static void test_chat_session() {
     }
 }
 
+static void test_chat_json_depth() {
+    LOG_DBG("%s\n", __func__);
+    const auto arrays = [](size_t n) {
+        return std::string(n, '[') + std::string(n, ']');
+    };
+    // a schema of n nested array types
+    const auto schema = [](size_t n) {
+        std::string s;
+        for (size_t i = 0; i < n; i++) {
+            s += R"({"type": "array", "items": )";
+        }
+        return s + R"({"type": "string"})" + std::string(n, '}');
+    };
+    const auto rejected = [](const std::function<void()> & fn) {
+        try {
+            fn();
+        } catch (const std::invalid_argument & e) {
+            return std::string(e.what()).find("nested more than 128 levels deep") != std::string::npos;
+        }
+        return false;
+    };
+
+    // the schema builder stops at 128 levels, the root schema is one of them
+    assert_equals(false, rejected([&] { common_chat_schema_from_json(json::parse(schema(127))); }));
+    assert_equals(true, rejected([&] { common_chat_schema_from_json(json::parse(schema(128))); }));
+    assert_equals(true, rejected([&] { json_schema_to_grammar(json::parse(schema(1000))); }));
+
+    // the jinja engine and the schema builder reject the inputs before they recurse too deep
+    auto         tmpls = read_templates("models/templates/Qwen3.5-4B.jinja");
+    const size_t deep  = 1000;
+
+    common_chat_templates_inputs tool_schema;
+    tool_schema.messages = { message_user };
+    tool_schema.tools    = { { "f", "", R"({"type": "object", "properties": {"x": )" + schema(deep) + "}}" } };
+    assert_equals(true, rejected([&] { common_chat_templates_apply(tmpls.get(), tool_schema); }));
+
+    common_chat_templates_inputs tool_args;
+    tool_args.messages = { message_user, simple_assist_msg("", "", "special_function", R"({"arg1": )" + arrays(deep) + "}") };
+    tool_args.tools    = { special_function_tool };
+    assert_equals(true, rejected([&] { common_chat_templates_apply(tmpls.get(), tool_args); }));
+
+    common_chat_templates_inputs kwargs;
+    kwargs.messages                     = { message_user };
+    kwargs.chat_template_kwargs["deep"] = arrays(deep);
+    assert_equals(true, rejected([&] { common_chat_templates_apply(tmpls.get(), kwargs); }));
+
+    common_chat_templates_inputs json_schema;
+    json_schema.messages    = { message_user };
+    json_schema.json_schema = schema(deep);
+    assert_equals(true, rejected([&] { common_chat_templates_apply(tmpls.get(), json_schema); }));
+}
+
 static void test_reasoning_budget_tokens_per_request() {
     LOG_DBG("%s\n", __func__);
     // Use Qwen3 template which has <think>...</think> reasoning markers.
@@ -8175,10 +8404,12 @@ int main(int argc, char ** argv) {
         test_translate_gemma();
         test_deepseek_v4_thinking_retention();
         test_deepseek_v4_tool_result_ordering();
+        test_deepseek_v4_1_message_fields();
         test_template_generation_prompt();
         test_reasoning_effort_caps();
         test_prompt_parts();
         test_chat_session();
+        test_chat_json_depth();
         test_reasoning_budget_tokens_per_request();
         test_reasoning_budget_message_per_request();
         test_template_output_peg_parsers(detailed_debug);

tests/test-jinja.cpp

@@ -1804,6 +1804,54 @@ static void test_array_methods(testing & t) {
         "z y "
     );
 
+    test_template(t, "array|select with test",
+        "{{ items|select('odd')|join(',') }}",
+        {{"items", json::array({1, 2, 3, 4, 5})}},
+        "1,3,5"
+    );
+
+    test_template(t, "array|reject with test",
+        "{{ items|reject('odd')|join(',') }}",
+        {{"items", json::array({1, 2, 3, 4, 5})}},
+        "2,4"
+    );
+
+    test_template(t, "array|select without test",
+        "{{ items|select|join(',') }}",
+        {{"items", json::array({0, 1, "", "a", nullptr, false})}},
+        "1,a"
+    );
+
+    test_template(t, "array|selectattr with test",
+        "{% for p in items|selectattr('type', 'defined') %}{{ p.type }} {% endfor %}",
+        {{"items", json::array({
+            {{"type", "string"}},
+            {{"enum", json::array({1})}},
+            {{"type", "integer"}}
+        })}},
+        "string integer "
+    );
+
+    test_template(t, "array|rejectattr with test",
+        "{{ items|rejectattr('type', 'defined')|list|length }}",
+        {{"items", json::array({
+            {{"type", "string"}},
+            {{"enum", json::array({1})}},
+            {{"type", "integer"}}
+        })}},
+        "1"
+    );
+
+    test_template(t, "array|selectattr with attr named like a test",
+        "{% for p in items|selectattr('odd', 'even') %}{{ p.odd }} {% endfor %}",
+        {{"items", json::array({
+            {{"odd", 1}},
+            {{"odd", 2}},
+            {{"odd", 4}}
+        })}},
+        "2 4 "
+    );
+
     test_template(t, "array|tojson",
         "{{ arr|tojson }}",
         {{"arr", json::array({1, 2, 3})}},

tests/test-llama-archs.cpp

@@ -195,6 +195,12 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
         n_head  = 8;
         n_ff    = 1024;
         n_layer = 4;
+    } else if (arch == LLM_ARCH_DEEPSEEK41) {
+        // window-only, Engram, ratio-2 and ratio-1 sources, top-k reuse and a candidate-masked index layer
+        n_embd  = 512;
+        n_head  = 8;
+        n_ff    = 1024;
+        n_layer = 6;
     } else if (arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_LAGUNA) {
         n_embd = 160; // exercise per-head tensor split granularity with head size 80
     } else if (arch == LLM_ARCH_QWEN3 || arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE) {
@@ -285,11 +291,11 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
         ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head_per_layer);
     } else {
         ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head);
-        ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(1) : n_head_kv);
+        ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41) ? uint32_t(1) : n_head_kv);
     }
 
     ms.add_kv(LLM_KV_ATTENTION_MAX_ALIBI_BIAS, 8.0f);
-    if (arch == LLM_ARCH_DEEPSEEK4) {
+    if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41) {
         ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH,   n_embd_head);
         ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, n_embd_head);
         ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT,   n_embd_head/2);
@@ -337,7 +343,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
     ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS,      1e-5f);
     ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_EPS,          1e-5f);
     ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_GROUPS,       uint32_t(8));
-    ms.add_kv(LLM_KV_ATTENTION_Q_LORA_RANK,            arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(64) : uint32_t(512));
+    ms.add_kv(LLM_KV_ATTENTION_Q_LORA_RANK,            (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41) ? uint32_t(64) : uint32_t(512));
     ms.add_kv(LLM_KV_ATTENTION_KV_LORA_RANK,           uint32_t(512));
     ms.add_kv(LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT, uint32_t(8));
     ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW,         n_ctx/8);
@@ -476,6 +482,39 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
         ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM,                   true);
     }
 
+    if (arch == LLM_ARCH_DEEPSEEK41) {
+        ms.add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,          uint32_t(8));
+        ms.add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK,            uint32_t(32));
+        ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS,             std::vector<uint32_t>({0, 0, 2, 2, 1, 1}));
+        ms.add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE,     160000.0f);
+        ms.add_kv(LLM_KV_ATTENTION_KV_SOURCE_LAYERS,            std::vector<uint32_t>({2, 4}));
+        ms.add_kv(LLM_KV_ATTENTION_INDEXER_SOURCE_LAYERS,       std::vector<uint32_t>({2, 4, 5}));
+        ms.add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_SOURCE_LAYER, int32_t(4));
+        ms.add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_TOP_K_BLOCKS, uint32_t(4096));
+        ms.add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_BLOCK_SIZE,   uint32_t(4));
+        ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT,                uint32_t(4));
+        ms.add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS,  uint32_t(2));
+        ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON,              1.0e-6f);
+        ms.add_kv(LLM_KV_HASH_LAYER_COUNT,                      uint32_t(0));
+        ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP,                      10.0f);
+        ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE,                  1.0f);
+        ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM,                   true);
+        std::vector<int32_t> token_map(n_vocab);
+        for (uint32_t i = 0; i < n_vocab; ++i) {
+            token_map[i] = (int32_t) (i % 61);
+        }
+        ms.add_kv(LLM_KV_PLE_LAYERS,                   std::vector<uint32_t>({1, 3}));
+        ms.add_kv(LLM_KV_PLE_NGRAM_SIZE,               uint32_t(3));
+        ms.add_kv(LLM_KV_PLE_HEADS_PER_NGRAM,          uint32_t(2));
+        ms.add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER,   uint32_t(64));
+        ms.add_kv(LLM_KV_PLE_LAYER_MULTIPLIERS,        std::vector<uint64_t>({ 76632096046245ull, 4839876093313ull, 35959672319349ull,
+                                                                               67716810739261ull, 51510806800915ull, 30921347202721ull }));
+        ms.add_kv(LLM_KV_PLE_HEAD_OFFSETS,             std::vector<uint64_t>({0, 101, 204, 311, 420, 533, 660, 791}));
+        ms.add_kv(LLM_KV_PLE_HEAD_VOCAB_SIZES,         std::vector<uint64_t>({101, 103, 107, 109, 113, 127, 131, 137}));
+        ms.add_kv(LLM_KV_ENGRAM_PAD_TOKEN_ID,          uint32_t(2));
+        ms.add_kv(LLM_KV_ENGRAM_TOKEN_MAP,             token_map);
+    }
+
     if (arch == LLM_ARCH_HRM_TEXT) {
         // 6 cache slots alias 2 physical blocks: 1 low-stack layer + 1 high-stack layer
         ms.add_kv(LLM_KV_HRM_LAYERS_PER_STACK, uint32_t(1));
@@ -519,6 +558,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
                 gating_func = 1;
                 break;
             case LLM_ARCH_DEEPSEEK4:
+            case LLM_ARCH_DEEPSEEK41:
                 gating_func = 4;
                 break;
             default:
@@ -780,6 +820,7 @@ static bool moe_mandatory(const llm_arch arch) {
         case LLM_ARCH_DEEPSEEK32:
         case LLM_ARCH_DOTS3NOTE:
         case LLM_ARCH_DEEPSEEK4:
+        case LLM_ARCH_DEEPSEEK41:
         case LLM_ARCH_GLM4_MOE:
         case LLM_ARCH_GLM_DSA:
         case LLM_ARCH_EXAONE_MOE:

tests/test-server-component.cpp

@@ -464,6 +464,32 @@ static void test_rerank_template(testing & t) {
     }
 }
 
+//
+// server_decision_context: the recursive functions stop at RECURSIVE_LIMIT
+//
+
+static void test_decision_depth(testing & t) {
+    const auto prepare = [](size_t depth) {
+        server_decision_context ctx;
+        ctx.type = COMMON_DECISION_TYPE_H2O_LIGHTNING; // the state goes through decision_py_dump
+        std::string error;
+        try {
+            ctx.prepare_state(json::parse(std::string(depth, '[') + std::string(depth, ']')));
+        } catch (const std::invalid_argument & e) {
+            error = e.what();
+        }
+        return error;
+    };
+    t.test("state at the limit", [&](testing & t) {
+        const std::string error = prepare(128);
+        t.assert_true("accepted: " + error, error.empty());
+    });
+    t.test("deep state", [&](testing & t) {
+        const std::string error = prepare(1000);
+        t.assert_true("rejected as nested: " + error, error.find("nested more than 128 levels deep") != std::string::npos);
+    });
+}
+
 int main(int argc, char ** argv) {
     testing t;
 
@@ -484,6 +510,7 @@ int main(int argc, char ** argv) {
     t.test("server_pipe",   test_server_pipe);
     t.test("decision_escape", test_decision_escape);
     t.test("rerank_template", test_rerank_template);
+    t.test("decision_depth",  test_decision_depth);
 
     return t.summary();
 }

tools/gguf-split/gguf-split.cpp

@@ -553,6 +553,10 @@ static void gguf_merge(const split_params & split_params) {
             auto offset = gguf_get_data_offset(ctx_gguf) + gguf_get_tensor_offset(ctx_gguf, i_tensor);
             f_input.seekg(offset);
             f_input.read((char *)read_data.data(), n_bytes);
+            if (!f_input) {
+                fprintf(stderr, "%s: failed to read %zu bytes at offset %zu, input file may be truncated\n", __func__, n_bytes, offset);
+                exit(EXIT_FAILURE);
+            }
             if (!split_params.dry_run) {
                 // write tensor data + padding
                 fout.write((const char *)read_data.data(), n_bytes);

tools/mtmd/clip-impl.h

@@ -529,6 +529,7 @@ enum projector_type {
     PROJECTOR_TYPE_DEEPSEEKOCR,
     PROJECTOR_TYPE_DEEPSEEKOCR2,
     PROJECTOR_TYPE_DEEPSEEK4V,
+    PROJECTOR_TYPE_DEEPSEEK41V,
     PROJECTOR_TYPE_LFM2A,
     PROJECTOR_TYPE_D1OMNI_A,
     PROJECTOR_TYPE_GLM4V,
@@ -606,6 +607,7 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
     { PROJECTOR_TYPE_DEEPSEEKOCR,       "deepseekocr"},
     { PROJECTOR_TYPE_DEEPSEEKOCR2,      "deepseekocr2"},
     { PROJECTOR_TYPE_DEEPSEEK4V,        "deepseek4v"},
+    { PROJECTOR_TYPE_DEEPSEEK41V,       "deepseek41v"},
     { PROJECTOR_TYPE_LFM2A,             "lfm2a"},
     { PROJECTOR_TYPE_D1OMNI_A,          "d1omni_a"},
     { PROJECTOR_TYPE_GLM4V,             "glm4v"},
@@ -879,6 +881,11 @@ static inline void clip_anyres_unpad(int cur_w, int cur_h, int orig_w, int orig_
     }
 }
 
+// deepseek41v: [START] ([IMAGE]*n_llm_w [NEWLINE])*n_llm_h [END], ref: num_image_tokens() in image_processor.py
+static inline int dsv41_n_block_tokens(int n_llm_w, int n_llm_h) {
+    return n_llm_h * (n_llm_w + 1) + 2;
+}
+
 // deepseek4v: layout of the LLM token block built from the aligner grid
 struct dsv4_block_layout {
     int rows;     // grid rows, padded to an even count

tools/mtmd/clip.cpp

@@ -1111,6 +1111,7 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
                 builder = std::make_unique<clip_graph_kimik25>(ctx, img);
             } break;
         case PROJECTOR_TYPE_DEEPSEEK4V:
+        case PROJECTOR_TYPE_DEEPSEEK41V:
             {
                 builder = std::make_unique<clip_graph_deepseek4v>(ctx, img);
             } break;
@@ -1773,6 +1774,30 @@ struct clip_model_loader {
                         const int warmup_side = (int) std::sqrt((double) std::min(256, hparams.dsv4_max_n_token));
                         hparams.set_warmup_n_tokens(warmup_side * warmup_side);
                     } break;
+                case PROJECTOR_TYPE_DEEPSEEK41V:
+                    {
+                        // ref: load_image in inference/image_processor.py (V4.1), PIL ImageOps.pad on gray 127
+                        hparams.image_resize_algo = RESIZE_ALGO_BICUBIC;
+                        hparams.image_pad_color   = {127, 127, 127};
+                        hparams.rope_theta = 10000.0f;
+                        get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge);
+                        get_u32(KEY_IMAGE_MIN_PIXELS,  hparams.image_min_pixels);
+                        // vision_max_n_token = 1024, vision_max_wh_ratio = null (converter asserts both)
+                        hparams.dsv4_max_n_token  = 1024;
+                        hparams.dsv4_max_wh_ratio = 0;
+                        const int patch_area = hparams.patch_size * hparams.patch_size * hparams.n_merge * hparams.n_merge;
+                        if (hparams.custom_image_min_tokens > 0) {
+                            hparams.image_min_pixels = hparams.custom_image_min_tokens * patch_area;
+                        }
+                        if (hparams.custom_image_max_tokens > 0) {
+                            // the cap is on the whole token block (newlines + start/end included)
+                            hparams.dsv4_max_n_token = std::max(hparams.custom_image_max_tokens, 16);
+                        }
+                        hparams.image_max_pixels = hparams.dsv4_max_n_token * patch_area;
+                        hparams.image_min_pixels = std::min(hparams.image_min_pixels, hparams.image_max_pixels);
+                        const int warmup_side = (int) std::sqrt((double) std::min(256, hparams.dsv4_max_n_token));
+                        hparams.set_warmup_n_tokens(warmup_side * warmup_side);
+                    } break;
                 case PROJECTOR_TYPE_GEMMA3:
                     {
                         // default value (used by all model sizes in gemma 3 family)
@@ -3065,6 +3090,17 @@ struct clip_model_loader {
                     model.token_embd_img_end   = get_tensor(TN_TOK_IMG_END);
                     model.token_embd_img_pad   = get_tensor(TN_TOK_IMG_PAD);
                 } break;
+            case PROJECTOR_TYPE_DEEPSEEK41V:
+                {
+                    model.mm_1_w = get_tensor(string_format(TN_LLAVA_PROJ, 1, "weight"));
+                    model.mm_1_b = get_tensor(string_format(TN_LLAVA_PROJ, 1, "bias"));
+                    model.mm_2_w = get_tensor(string_format(TN_LLAVA_PROJ, 2, "weight"));
+                    model.mm_2_b = get_tensor(string_format(TN_LLAVA_PROJ, 2, "bias"));
+                    // image_start / image_end / image_newline, written into the output block
+                    model.image_newline        = get_tensor(TN_IMAGE_NEWLINE);
+                    model.token_embd_img_start = get_tensor(TN_TOK_IMG_START);
+                    model.token_embd_img_end   = get_tensor(TN_TOK_IMG_END);
+                } break;
             case PROJECTOR_TYPE_PIXTRAL:
                 {
                     model.mm_1_w = get_tensor(string_format(TN_LLAVA_PROJ, 1, "weight"));
@@ -4377,7 +4413,7 @@ struct clip_cap clip_get_cap(const char * fname) {
 
 int clip_get_image_max_tokens(const clip_ctx * ctx) {
     const auto & hparams = ctx->model.hparams;
-    if (ctx->proj_type() == PROJECTOR_TYPE_DEEPSEEK4V) {
+    if (ctx->proj_type() == PROJECTOR_TYPE_DEEPSEEK4V || ctx->proj_type() == PROJECTOR_TYPE_DEEPSEEK41V) {
         return hparams.dsv4_max_n_token;
     }
     if (hparams.image_max_pixels <= 0) {
@@ -4585,6 +4621,13 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
                 const int n_llm_h = CLIP_ALIGN(img->ny(), out_patch_size) / out_patch_size;
                 n_patches = dsv4_get_block_layout(n_llm_w, n_llm_h, img->lead_pad).n_out;
             } break;
+        case PROJECTOR_TYPE_DEEPSEEK41V:
+            {
+                const int out_patch_size = params.patch_size * params.n_merge;
+                const int n_llm_w = CLIP_ALIGN(img->nx(), out_patch_size) / out_patch_size;
+                const int n_llm_h = CLIP_ALIGN(img->ny(), out_patch_size) / out_patch_size;
+                n_patches = dsv41_n_block_tokens(n_llm_w, n_llm_h);
+            } break;
         case PROJECTOR_TYPE_PADDLEOCR:
         case PROJECTOR_TYPE_DOTS_OCR:
         case PROJECTOR_TYPE_DOTS3NOTE_V:
@@ -5692,6 +5735,34 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
                 idx.push_back(idx_end);
                 set_input_i32("layout_idx", idx);
             } break;
+        case PROJECTOR_TYPE_DEEPSEEK41V:
+            {
+                // mrope layout, only the first 2 channels are used
+                const int n_patches_per_row = image_size_width / patch_size;
+                std::vector<int32_t> positions(n_pos * 4, 0);
+                for (int i = 0; i < n_pos; i++) {
+                    positions[i]         = i / n_patches_per_row;
+                    positions[n_pos + i] = i % n_patches_per_row;
+                }
+                set_input_i32("positions", positions);
+
+                // reading-order block: rows [0, n_grid) are the aligner output, then start / end / newline
+                const int n_merge = hparams.n_merge;
+                const int n_llm_w = CLIP_ALIGN(pos_w, n_merge) / n_merge;
+                const int n_llm_h = CLIP_ALIGN(pos_h, n_merge) / n_merge;
+                const int n_grid  = n_llm_w * n_llm_h;
+                std::vector<int32_t> idx;
+                idx.reserve(dsv41_n_block_tokens(n_llm_w, n_llm_h));
+                idx.push_back(n_grid);         // IMAGE_START
+                for (int r = 0; r < n_llm_h; r++) {
+                    for (int c = 0; c < n_llm_w; c++) {
+                        idx.push_back(r * n_llm_w + c);
+                    }
+                    idx.push_back(n_grid + 2); // IMAGE_NEW_LINE
+                }
+                idx.push_back(n_grid + 1);     // IMAGE_END
+                set_input_i32("layout_idx", idx);
+            } break;
         case PROJECTOR_TYPE_GLM_EDGE:
         {
             // llava and other models
@@ -6583,6 +6654,7 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
         case PROJECTOR_TYPE_KIMIK25:
         case PROJECTOR_TYPE_YASA2:
         case PROJECTOR_TYPE_DEEPSEEK4V:
+        case PROJECTOR_TYPE_DEEPSEEK41V:
         case PROJECTOR_TYPE_COHERE2V:
             return ctx->model.mm_2_w->ne[1];
         case PROJECTOR_TYPE_HUNYUANVL:

tools/mtmd/models/deepseek4v.cpp

@@ -14,6 +14,9 @@
 // the mapping is precomputed on CPU as the "layout_idx" input (see set_input in clip.cpp)
 //
 // ref: inference/vision.py and inference/image_processor.py in the HF repo
+//
+// DeepSeek-V4.1-Flash (deepseek41v) shares the encoder and aligner; its block is plain reading order
+//   [START] ([IMAGE]*n_llm_w [NEWLINE])*n_llm_h [END], no lead pads, no pad row, no IMAGE_PAD tensor
 
 ggml_cgraph * clip_graph_deepseek4v::build() {
     const int n_merge = hparams.n_merge;
@@ -52,7 +55,10 @@ ggml_cgraph * clip_graph_deepseek4v::build() {
             cur = ggml_pad(ctx0, cur, pad_x, pad_y, 0, 0);
         }
 
-        ggml_tensor * kernel = ggml_view_3d(ctx0, cur, n_merge, n_merge, cur->ne[2], 0, 0, 0);
+        // contiguous kernel view: CUDA rejects a zero-stride one and im2col falls back to CPU
+        const size_t esz = ggml_element_size(cur);
+        ggml_tensor * kernel = ggml_view_3d(ctx0, cur, n_merge, n_merge, cur->ne[2],
+                                            n_merge*esz, n_merge*n_merge*esz, 0);
         cur = ggml_im2col(ctx0, kernel, cur, n_merge, n_merge, 0, 0, 1, 1, true, inp->type);
         cur = ggml_reshape_2d(ctx0, cur, cur->ne[0], cur->ne[1] * cur->ne[2]);
 
@@ -73,19 +79,23 @@ ggml_cgraph * clip_graph_deepseek4v::build() {
         const int64_t n_grid     = cur->ne[1]; // n_llm_w * n_llm_h
 
         // rows n_grid + 0..3, keep in sync with the index computation in set_input
+        const bool is_v41 = proj_type == PROJECTOR_TYPE_DEEPSEEK41V;
         ggml_tensor * sentinels[] = {
             model.token_embd_img_start,
             model.token_embd_img_end,
             model.image_newline,
-            model.token_embd_img_pad,
+            is_v41 ? nullptr : model.token_embd_img_pad,
         };
         for (ggml_tensor * tok : sentinels) {
-            cur = ggml_concat(ctx0, cur, ggml_reshape_2d(ctx0, tok, n_embd_out, 1), 1);
+            if (tok) {
+                cur = ggml_concat(ctx0, cur, ggml_reshape_2d(ctx0, tok, n_embd_out, 1), 1);
+            }
         }
 
         const int n_llm_w = CLIP_ALIGN(n_patches_x, n_merge) / n_merge;
         const int n_llm_h = CLIP_ALIGN(n_patches_y, n_merge) / n_merge;
-        const int n_out   = dsv4_get_block_layout(n_llm_w, n_llm_h, img.lead_pad).n_out;
+        const int n_out   = is_v41 ? dsv41_n_block_tokens(n_llm_w, n_llm_h)
+                                   : dsv4_get_block_layout(n_llm_w, n_llm_h, img.lead_pad).n_out;
         GGML_ASSERT(n_grid == n_llm_w * n_llm_h);
 
         ggml_tensor * layout_idx = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_out);

tools/mtmd/mtmd-image.cpp

@@ -1454,6 +1454,100 @@ mtmd_image_preproc_out mtmd_image_preprocessor_deepseek4v::preprocess(const clip
     return out;
 }
 
+//
+// DeepSeek-V4.1-Flash (deepseek41v)
+//
+// port of plan_image_grid / load_image from inference/image_processor.py (V4.1)
+// differs from deepseek4v: block = h * (w + 1) + 2 tokens, no even-row rounding, a single solve, no stretch path
+//
+
+void mtmd_image_preprocessor_deepseek41v::plan_grid(int width, int height, int p, int r, int max_n_token, int min_pixels,
+                                                    int & best_width, int & best_height) {
+    auto n_tokens = [&](int bh, int bw) {
+        const int n_llm_h = ((bh / p) + r - 1) / r;
+        const int n_llm_w = ((bw / p) + r - 1) / r;
+        return n_llm_h * (n_llm_w + 1) + 2;
+    };
+    if (min_pixels > 0 && (int64_t) width * height > 0 && (int64_t) width * height < min_pixels) {
+        // python: ratio = (min_pixels / (w * h)) ** 0.5; int(w * ratio)
+        const double ratio = std::sqrt((double) min_pixels / ((double) width * (double) height));
+        width  = (int) (width  * ratio);
+        height = (int) (height * ratio);
+    }
+    best_width  = CLIP_ALIGN(width,  p);
+    best_height = CLIP_ALIGN(height, p);
+    if (n_tokens(best_height, best_width) > max_n_token) {
+        // solve_resize_ratio: largest aspect-preserving size whose token block fits
+        const double ratio   = (double) height / width;
+        const double max_w_f = std::sqrt((max_n_token - 2) / ratio + 0.25) - 0.5;
+        const double max_h_f = max_w_f * ratio;
+        const int cell = p * r;
+        if (max_w_f < 1.0) {
+            best_height = (max_n_token - 2) / 2 * cell;
+            best_width  = cell;
+        } else if (max_h_f < 1.0) {
+            best_height = cell;
+            best_width  = (max_n_token - 3) * cell;
+        } else {
+            const double beta = std::min(std::floor(max_w_f) * cell / width, std::floor(max_h_f) * cell / height);
+            best_height = (int) std::floor(height * beta / p) * p;
+            best_width  = (int) std::floor(width  * beta / p) * p;
+        }
+        GGML_ASSERT(n_tokens(best_height, best_width) <= max_n_token);
+    }
+}
+
+mtmd_image_preproc_out mtmd_image_preprocessor_deepseek41v::preprocess(const clip_image_u8 & img) const {
+    mtmd_image_preproc_out out;
+
+    const int p = hparams.patch_size;
+    const clip_image_size orig = img.get_size();
+    int best_width = 0, best_height = 0;
+    plan_grid(orig.width, orig.height, p, hparams.n_merge, hparams.dsv4_max_n_token, hparams.image_min_pixels,
+              best_width, best_height);
+
+    // PIL ImageOps.pad(color=127): ImageOps.contain then a centered paste; python round is half to even
+    auto round_half_even = [](double x) { return (int) std::nearbyint(x); };
+    int new_w = best_width;
+    int new_h = best_height;
+    const double im_ratio   = (double) orig.width / orig.height;
+    const double dest_ratio = (double) best_width / best_height;
+    if (im_ratio != dest_ratio) {
+        if (im_ratio > dest_ratio) {
+            new_h = round_half_even((double) orig.height / orig.width * best_width);
+        } else {
+            new_w = round_half_even((double) orig.width / orig.height * best_height);
+        }
+    }
+    clip_image_u8 resized;
+    img_tool::resize(img, resized, {new_w, new_h}, hparams.image_resize_algo, PAD_NONE);
+
+    clip_image_u8 padded;
+    if (new_w == best_width && new_h == best_height) {
+        padded = std::move(resized);
+    } else {
+        padded.set_size({best_width, best_height}, img.is_placeholder());
+        if (!img.is_placeholder()) {
+            img_tool::fill(padded, hparams.image_pad_color);
+            const int off_x = round_half_even((best_width  - new_w) * 0.5);
+            const int off_y = round_half_even((best_height - new_h) * 0.5);
+            img_tool::composite(padded, resized, off_x, off_y);
+        }
+    }
+
+    out.append(hparams, padded);
+
+    // load_image rounds the normalized input to bf16 (round to nearest even)
+    for (auto & e : out.entries) {
+        std::vector<float> buf = e.get_ro_buf();
+        for (float & v : buf) {
+            v = ggml_bf16_to_fp32(ggml_fp32_to_bf16(v));
+        }
+        e.cpy_buf(buf);
+    }
+    return out;
+}
+
 mtmd_image_preproc_out mtmd_image_preprocessor_deepseekocr::preprocess(const clip_image_u8 & img) const {
     mtmd_image_preproc_out output;
     int grid_w = 0;

tools/mtmd/mtmd-image.h

@@ -166,6 +166,15 @@ private:
     static void safe_resize(int height, int width, int & best_height, int & best_width, int p, int r, int max_n_token);
 };
 
+// ref: inference/image_processor.py in the HF repo (DeepSeek-V4.1-Flash): plan_image_grid / load_image
+struct mtmd_image_preprocessor_deepseek41v : mtmd_image_preprocessor {
+    mtmd_image_preprocessor_deepseek41v(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {}
+    mtmd_image_preproc_out preprocess(const clip_image_u8 & img) const override;
+    // target pixel size (multiples of patch_size) for an image of the given size, exposed for tests
+    static void plan_grid(int width, int height, int p, int r, int max_n_token, int min_pixels,
+                          int & best_width, int & best_height);
+};
+
 // custom llava-uhd slicing logic for MiniCPM-V
 struct mtmd_image_preprocessor_minicpmv : mtmd_image_preprocessor_llava_uhd {
     using mtmd_image_preprocessor_llava_uhd::mtmd_image_preprocessor_llava_uhd;

tools/mtmd/mtmd.cpp

@@ -903,6 +903,11 @@ struct mtmd_context {
                     // no vocab tokens are added; the start/end/newline markers are learned embeddings emitted by the encoder
                     image_preproc = std::make_unique<mtmd_image_preprocessor_deepseek4v>(ctx_v);
                 } break;
+            case PROJECTOR_TYPE_DEEPSEEK41V:
+                {
+                    // as deepseek4v: markers come from the encoder, the block is decoded causally
+                    image_preproc = std::make_unique<mtmd_image_preprocessor_deepseek41v>(ctx_v);
+                } break;
             case PROJECTOR_TYPE_DOTS_OCR:
             case PROJECTOR_TYPE_DOTS3NOTE_V:
                 {

tools/server/README-dev.md

@@ -286,7 +286,7 @@ Instead of building everything from the ground up (like what most AI agents will
 
 The flow for downloading a new model:
 - POST request comes in --> `post_router_models` --> validation
-- A new `llama-server` subprocess will be spawned with special `SERVER_CHILD_MODE_DOWNLOAD`
+- A new `llama-server` subprocess will be spawned with the `SERVER_ROLE_DOWNLOAD` role (see `server_role_type` in `server-common.h`)
 - Child process runs the download and report status back to router via stdin/out
 - If a stop request comes in, the router asks the child process to stop (same mechanism as running a model in child process)
 - Otherwise, upon completion, we call `load_models()` to refresh the list of models

tools/server/server-common.cpp

@@ -16,7 +16,9 @@
 #include <limits>
 #include <cstring>
 #include <type_traits>
+#include <atomic>
 #include <chrono>
+#include <iostream>
 #include <thread>
 
 #ifdef _WIN32
@@ -2210,3 +2212,64 @@ void server_subproc::waiter::wait(const std::vector<server_subproc *> & procs, s
     }
 #endif
 }
+
+//
+// server_role
+//
+
+server_role::server_role() {
+    if (is_managed_by_router()) {
+        cmd_out = server_reserve_stdout();
+    }
+}
+
+server_role::~server_role() {
+    if (cmd_out) {
+        fclose(cmd_out);
+    }
+}
+
+bool server_role::is_managed_by_router() {
+    return std::getenv(SERVER_ENV_ROUTER_PORT) != nullptr;
+}
+
+server_role_type server_role::get_role() {
+    const char * role = std::getenv(SERVER_ENV_CHILD_ROLE);
+    return server_role_type_from_str(role ? role : "");
+}
+
+std::thread server_role::setup(const std::function<void(int)> & shutdown_handler) {
+    // setup thread for monitoring stdin
+    return std::thread([shutdown_handler]() {
+        // wait for EOF on stdin
+        SRV_INF("%s", "child server monitoring thread started, waiting for EOF on stdin...\n");
+        bool eof = false;
+        while (true) {
+            std::string line;
+            if (!std::getline(std::cin, line)) {
+                // EOF detected, that means the router server is unexpectedly exit or killed
+                eof = true;
+                break;
+            }
+            if (line.find(CMD_ROUTER_TO_CHILD_EXIT) != std::string::npos) {
+                SRV_INF("%s", "exit command received, exiting...\n");
+                shutdown_handler(0);
+                break;
+            }
+        }
+        if (eof) {
+            SRV_INF("%s", "EOF on stdin detected, forcing shutdown...\n");
+            exit(1);
+        }
+    });
+}
+
+void server_role::notify_to_router(const std::string & state, const json & payload) {
+    json data = {
+        {"state", state},
+        {"payload", payload},
+    };
+    std::lock_guard<std::mutex> lk(mtx_stdout);
+    fprintf(cmd_out, "%s%s\n", CMD_CHILD_TO_ROUTER_STATE, safe_json_to_str(data).c_str());
+    fflush(cmd_out);
+}

tools/server/server-common.h

@@ -18,6 +18,7 @@
 #include <filesystem>
 #include <functional>
 #include <mutex>
+#include <thread>
 #include <queue>
 #include <string>
 #include <vector>
@@ -725,3 +726,97 @@ private:
     intptr_t out_handles[SERVER_SUBPROC_STREAMS] = { -1, -1 }; // fd on POSIX, HANDLE on Windows; taken lazily from sproc
     bool     out_closed [SERVER_SUBPROC_STREAMS] = { false, false };
 };
+
+//
+// router <-> child protocol
+//
+
+// what this llama-server process does. a router spawns the other roles as child processes
+// and talks to them over their stdin/stdout
+enum server_role_type {
+    SERVER_ROLE_ROUTER,   // no model: lists the models, spawns instances on demand and proxies requests to them
+    SERVER_ROLE_INSTANCE, // serves one model, standalone or as a child of a router
+    SERVER_ROLE_DOWNLOAD, // child of a router: download a model, report the progress, exit
+};
+
+static inline const char * server_role_type_to_str(server_role_type role) {
+    switch (role) {
+        case SERVER_ROLE_ROUTER:   return "router";
+        case SERVER_ROLE_INSTANCE: return "instance";
+        case SERVER_ROLE_DOWNLOAD: return "download";
+        default: GGML_ASSERT(false && "invalid server_role_type");
+    }
+}
+
+// the role a router asks of a child; anything it does not name is an instance
+static inline server_role_type server_role_type_from_str(const std::string & str) {
+    if (str == "download") {
+        return SERVER_ROLE_DOWNLOAD;
+    }
+    return SERVER_ROLE_INSTANCE;
+}
+
+// set by the router in the environment of every child
+#define SERVER_ENV_ROUTER_PORT "LLAMA_SERVER_ROUTER_PORT" // its presence makes the process a child
+#define SERVER_ENV_CHILD_ROLE  "LLAMA_SERVER_CHILD_ROLE"  // server_role_type_to_str(), an instance when unset
+
+// lines on the child's stdin (router -> child) and stdout (child -> router),
+// see server_role::setup() and server_role::notify_to_router()
+#define CMD_ROUTER_TO_CHILD_EXIT  "cmd_router_to_child:exit"
+#define CMD_CHILD_TO_ROUTER_STATE "cmd_child_to_router:state:" // followed by json string
+
+// address a child binds to: the router may listen on 0.0.0.0, but reaches its children locally
+// ref: https://github.com/ggml-org/llama.cpp/issues/17862
+#define CHILD_ADDR "127.0.0.1"
+
+// the states a child reports to the router
+enum server_state {
+    SERVER_STATE_DOWNLOADING,
+    SERVER_STATE_LOADING,
+    SERVER_STATE_READY,
+    SERVER_STATE_SLEEPING,
+};
+
+static inline std::string server_state_to_str(server_state state) {
+    switch (state) {
+        case SERVER_STATE_DOWNLOADING: return "downloading";
+        case SERVER_STATE_LOADING:     return "loading";
+        case SERVER_STATE_READY:       return "ready";
+        case SERVER_STATE_SLEEPING:    return "sleeping";
+        default: GGML_ASSERT(false && "invalid server_state");
+    }
+}
+
+static inline server_state server_state_from_str(const std::string & str) {
+    if (str == "downloading") return SERVER_STATE_DOWNLOADING;
+    if (str == "loading")     return SERVER_STATE_LOADING;
+    if (str == "ready")       return SERVER_STATE_READY;
+    if (str == "sleeping")    return SERVER_STATE_SLEEPING;
+    GGML_ASSERT(false && "invalid server_state string");
+}
+
+// the role of this process and, when it is a child of a router, its side of the protocol
+struct server_role {
+    // serializes the notify_to_router writes
+    std::mutex mtx_stdout;
+
+    // in a child, keeps stdout for the commands to the router, so it is created before anything is written;
+    // everything else written to stdout goes to stderr with the logs
+    server_role();
+    ~server_role();
+
+    // true if the current process was spawned by a router
+    bool is_managed_by_router();
+    server_role_type get_role();
+
+    // register the shutdown_handler to be called by the router
+    // return the monitoring thread (to be joined by the caller)
+    std::thread setup(const std::function<void(int)> & shutdown_handler);
+
+    // notify router server for status changes (e.g. loading, downloading, sleeping, etc.)
+    // message will be handled by server_models::handle_child_state() on the router side
+    void notify_to_router(const std::string & state_name, const json & payload);
+
+private:
+    FILE * cmd_out = nullptr; // the stdout the router reads the commands from
+};

tools/server/server-context.h

@@ -58,31 +58,6 @@ struct server_context_meta {
     std::string model_ftype;
 };
 
-enum server_state {
-    SERVER_STATE_DOWNLOADING,
-    SERVER_STATE_LOADING,
-    SERVER_STATE_READY,
-    SERVER_STATE_SLEEPING,
-};
-
-static std::string server_state_to_str(server_state state) {
-    switch (state) {
-        case SERVER_STATE_DOWNLOADING: return "downloading";
-        case SERVER_STATE_LOADING:     return "loading";
-        case SERVER_STATE_READY:       return "ready";
-        case SERVER_STATE_SLEEPING:    return "sleeping";
-        default: GGML_ASSERT(false && "invalid server_state");
-    }
-}
-
-static server_state server_state_from_str(const std::string & str) {
-    if (str == "downloading") return SERVER_STATE_DOWNLOADING;
-    if (str == "loading")     return SERVER_STATE_LOADING;
-    if (str == "ready")       return SERVER_STATE_READY;
-    if (str == "sleeping")    return SERVER_STATE_SLEEPING;
-    GGML_ASSERT(false && "invalid server_state string");
-}
-
 using server_state_callback_t = std::function<void(server_state, json /* payload */)>;
 
 struct server_context {

tools/server/server-decision.cpp

@@ -23,6 +23,9 @@ static const size_t DECISION_LEV_N_RATINGS = 9;
 // a longer state keeps its head and its tail
 static const size_t DECISION_H2O_MAX_STATE_TOKENS = 32000;
 
+// the recursive functions of this file walk request JSON, a deeper value would overflow the stack
+static const int RECURSIVE_LIMIT = 128;
+
 static std::string decision_meta_str(const llama_model * model, const std::string & key) {
     char buf[256];
     const int32_t n = llama_model_meta_val_str(model, key.c_str(), buf, sizeof(buf));
@@ -417,7 +420,10 @@ json server_decision_context::parse_state(const json & body, std::vector<raw_buf
 //
 
 // replace text in all strings of a JSON value
-static json decision_replace_text(const json & val, const std::string & search, const std::string & replace) {
+static json decision_replace_text(const json & val, const std::string & search, const std::string & replace, int depth = 0) {
+    if (depth >= RECURSIVE_LIMIT) {
+        throw std::invalid_argument("JSON is nested more than " + std::to_string(RECURSIVE_LIMIT) + " levels deep");
+    }
     if (val.is_string()) {
         std::string str = val.get<std::string>();
         string_replace_all(str, search, replace);
@@ -426,14 +432,14 @@ static json decision_replace_text(const json & val, const std::string & search,
     if (val.is_array()) {
         json out = json::array();
         for (const auto & item : val) {
-            out.push_back(decision_replace_text(item, search, replace));
+            out.push_back(decision_replace_text(item, search, replace, depth + 1));
         }
         return out;
     }
     if (val.is_object()) {
         json out = json::object();
         for (const auto & [key, item] : val.items()) {
-            out[key] = decision_replace_text(item, search, replace);
+            out[key] = decision_replace_text(item, search, replace, depth + 1);
         }
         return out;
     }
@@ -441,18 +447,21 @@ static json decision_replace_text(const json & val, const std::string & search,
 }
 
 // sort the keys of all objects of a JSON value
-static json decision_sort_keys(const json & val) {
+static json decision_sort_keys(const json & val, int depth = 0) {
+    if (depth >= RECURSIVE_LIMIT) {
+        throw std::invalid_argument("JSON is nested more than " + std::to_string(RECURSIVE_LIMIT) + " levels deep");
+    }
     if (val.is_array()) {
         json out = json::array();
         for (const auto & item : val) {
-            out.push_back(decision_sort_keys(item));
+            out.push_back(decision_sort_keys(item, depth + 1));
         }
         return out;
     }
     if (val.is_object()) {
         std::map<std::string, json> sorted;
         for (const auto & [key, item] : val.items()) {
-            sorted[key] = decision_sort_keys(item);
+            sorted[key] = decision_sort_keys(item, depth + 1);
         }
         json out = json::object();
         for (const auto & [key, item] : sorted) {
@@ -464,7 +473,10 @@ static json decision_sort_keys(const json & val) {
 }
 
 // kev flattens a JSON value into text, the keys of an object are kept as labels (kev/api.py: render)
-static std::string decision_kev_render(const json & val, int indent = 0) {
+static std::string decision_kev_render(const json & val, int indent = 0, int depth = 0) {
+    if (depth >= RECURSIVE_LIMIT) {
+        throw std::invalid_argument("JSON is nested more than " + std::to_string(RECURSIVE_LIMIT) + " levels deep");
+    }
     const std::string pad(2 * indent, ' ');
     if (val.is_null()) {
         return "";
@@ -478,7 +490,7 @@ static std::string decision_kev_render(const json & val, int indent = 0) {
     if (val.is_array()) {
         std::string out;
         for (const auto & item : val) {
-            const std::string text = decision_kev_render(item, indent + 1);
+            const std::string text = decision_kev_render(item, indent + 1, depth + 1);
             out += (out.empty() ? "" : "\n") + pad + "- " + text.substr(std::min(text.size(), text.find_first_not_of(" \t\n\r")));
         }
         return out;
@@ -487,7 +499,7 @@ static std::string decision_kev_render(const json & val, int indent = 0) {
         std::string out;
         for (const auto & [key, item] : val.items()) {
             const bool is_nested = item.is_object() || item.is_array();
-            out += (out.empty() ? "" : "\n") + pad + key + (is_nested ? ":\n" : ": ") + decision_kev_render(item, is_nested ? indent + 1 : 0);
+            out += (out.empty() ? "" : "\n") + pad + key + (is_nested ? ":\n" : ": ") + decision_kev_render(item, is_nested ? indent + 1 : 0, depth + 1);
         }
         return out;
     }
@@ -528,21 +540,24 @@ static std::string decision_kev_text(const json & val) {
 }
 
 // lfm2-d1-omni: special tokens written in the input must not be parsed as such, in keys too (d1-omni prompt.py: escape)
-static json decision_d1omni_escape(const json & val) {
+static json decision_d1omni_escape(const json & val, int depth = 0) {
+    if (depth >= RECURSIVE_LIMIT) {
+        throw std::invalid_argument("JSON is nested more than " + std::to_string(RECURSIVE_LIMIT) + " levels deep");
+    }
     if (val.is_string()) {
         return server_decision_escape_special_tokens(val.get<std::string>());
     }
     if (val.is_array()) {
         json out = json::array();
         for (const auto & item : val) {
-            out.push_back(decision_d1omni_escape(item));
+            out.push_back(decision_d1omni_escape(item, depth + 1));
         }
         return out;
     }
     if (val.is_object()) {
         json out = json::object();
         for (const auto & [key, item] : val.items()) {
-            out[server_decision_escape_special_tokens(key)] = decision_d1omni_escape(item);
+            out[server_decision_escape_special_tokens(key)] = decision_d1omni_escape(item, depth + 1);
         }
         return out;
     }
@@ -669,6 +684,9 @@ static void decision_py_str(std::string & out, const std::string & s, bool ensur
 
 // json.dumps() of Python: indent < 0 is no indent, keys stay in their order
 static void decision_py_dump(std::string & out, const json & val, int indent, int level, const std::string & item_sep, const std::string & key_sep, bool ensure_ascii) {
+    if (level >= RECURSIVE_LIMIT) {
+        throw std::invalid_argument("JSON is nested more than " + std::to_string(RECURSIVE_LIMIT) + " levels deep");
+    }
     const auto newline = [&](int lvl) {
         return indent < 0 ? std::string() : "\n" + std::string(indent * lvl, ' ');
     };

tools/server/server-models.cpp

@@ -37,9 +37,6 @@ extern char **environ;
 
 #define DEFAULT_STOP_TIMEOUT 10 // seconds
 
-#define CMD_ROUTER_TO_CHILD_EXIT  "cmd_router_to_child:exit"
-#define CMD_CHILD_TO_ROUTER_STATE "cmd_child_to_router:state:" // followed by json string
-
 // note: SIGPIPE is ignored by the server
 static void request_child_exit(server_subproc & proc) {
     FILE * stdin_file = proc.sproc.stdin_file();
@@ -49,10 +46,6 @@ static void request_child_exit(server_subproc & proc) {
     }
 }
 
-// address for child process, this is needed because router may run on 0.0.0.0
-// ref: https://github.com/ggml-org/llama.cpp/issues/17862
-#define CHILD_ADDR "127.0.0.1"
-
 // single-threaded, watching all child processes at once
 struct server_monitor {
     server_monitor(server_models & models) : models(models) {
@@ -65,11 +58,11 @@ struct server_monitor {
     }
 
     // thread-safe
-    void watch(const std::string & name, std::shared_ptr<server_subproc> proc, server_child_mode mode, int port) {
+    void watch(const std::string & name, std::shared_ptr<server_subproc> proc, server_role_type role, int port) {
         child_t c;
         c.name = name;
         c.proc = std::move(proc);
-        c.mode = mode;
+        c.role = role;
         c.port = port;
         if (!c.proc->has_output()) {
             SRV_ERR("failed to get stdout/stderr of child process for name=%s\n", name.c_str());
@@ -87,7 +80,7 @@ private:
     struct child_t {
         std::string name;
         std::shared_ptr<server_subproc> proc;
-        server_child_mode mode = SERVER_CHILD_MODE_NORMAL;
+        server_role_type role = SERVER_ROLE_INSTANCE;
         int port = 0;
         std::string buf[SERVER_SUBPROC_STREAMS]; // partial line of each pipe
         bool eof = false;     // output closed, waiting for the process to be reaped
@@ -241,7 +234,7 @@ private:
                 if (it->eof && !it->proc->is_alive()) {
                     int exit_code = it->proc->join();
                     it->proc->stopped.store(true, std::memory_order_release);
-                    models.on_child_exit(it->name, it->proc, it->mode, exit_code);
+                    models.on_child_exit(it->name, it->proc, it->role, exit_code);
                     SRV_INF("instance name=%s exited with status %d\n", it->name.c_str(), exit_code);
                     it = children.erase(it);
                 } else {
@@ -1324,7 +1317,7 @@ void server_models::load(const std::string & name) {
 
 void server_models::load(const std::string & name, const load_options & opts) {
     // the fake load time of a model; a download worker loads none
-    if (debug_fake_timing && opts.mode == SERVER_CHILD_MODE_NORMAL) {
+    if (debug_fake_timing && opts.role == SERVER_ROLE_INSTANCE) {
         // do not hold the mutex here, other requests must keep making progress
         std::this_thread::sleep_for(std::chrono::seconds(2));
     }
@@ -1345,7 +1338,7 @@ void server_models::load(const std::string & name, const load_options & opts) {
     // same model cannot start in between: a paused download or an idle cached model
     // is downloaded again with its files kept, a download on its way and a running
     // model are refused, and a preset or local model is not a download
-    if (opts.mode == SERVER_CHILD_MODE_DOWNLOAD) {
+    if (opts.role == SERVER_ROLE_DOWNLOAD) {
         auto it = mapping.find(name);
         if (it != mapping.end()) {
             const auto & cur = it->second.meta;
@@ -1389,7 +1382,7 @@ void server_models::load(const std::string & name, const load_options & opts) {
     // releasing its lock and this lock_guard acquiring allows multiple
     // threads to each observe capacity and all proceed to load.
     // Download workers do not use models_max slots.
-    if (opts.mode == SERVER_CHILD_MODE_NORMAL && base_params.models_max > 0) {
+    if (opts.role == SERVER_ROLE_INSTANCE && base_params.models_max > 0) {
         size_t count_active = 0;
         for (const auto & m : mapping) {
             if (m.second.meta.is_running()) {
@@ -1421,11 +1414,11 @@ void server_models::load(const std::string & name, const load_options & opts) {
 
         std::vector<std::string> child_args = inst.meta.args; // copy
         std::vector<std::string> child_env  = base_env; // copy
-        child_env.push_back("LLAMA_SERVER_ROUTER_PORT=" + std::to_string(base_params.port));
+        child_env.push_back(std::string(SERVER_ENV_ROUTER_PORT) + "=" + std::to_string(base_params.port));
 
-        if (opts.mode == SERVER_CHILD_MODE_DOWNLOAD) {
+        if (opts.role == SERVER_ROLE_DOWNLOAD) {
             inst.meta.status = SERVER_MODEL_STATUS_DOWNLOADING;
-            child_env.push_back("LLAMA_SERVER_CHILD_MODE=download");
+            child_env.push_back(std::string(SERVER_ENV_CHILD_ROLE) + "=" + server_role_type_to_str(SERVER_ROLE_DOWNLOAD));
             child_env.push_back("LLAMA_ARG_HF_REPO=" + name);
 
             // the sidecars as -hf would pull them: the draft types, and the mmproj switch
@@ -1460,7 +1453,7 @@ void server_models::load(const std::string & name, const load_options & opts) {
     }
 
     // remember the download: the partial files carry a content hash, not the name
-    if (opts.mode == SERVER_CHILD_MODE_DOWNLOAD) {
+    if (opts.role == SERVER_ROLE_DOWNLOAD) {
         unfinished_downloads.add({ name, opts.sidecars, opts.partials });
     }
 
@@ -1480,7 +1473,7 @@ void server_models::load(const std::string & name, const load_options & opts) {
     auto proc = inst.subproc;
     int  port = inst.meta.port;
     mapping[name] = std::move(inst);
-    monitor->watch(name, proc, opts.mode, port);
+    monitor->watch(name, proc, opts.role, port);
     cv.notify_all();
 }
 
@@ -1493,7 +1486,7 @@ void server_models::request_stop(const std::string & name, bool send_exit) {
     monitor->stop(name, it->second.meta.stop_timeout, send_exit);
 }
 
-void server_models::on_child_exit(const std::string & name, const std::shared_ptr<server_subproc> & proc, server_child_mode mode, int exit_code) {
+void server_models::on_child_exit(const std::string & name, const std::shared_ptr<server_subproc> & proc, server_role_type role, int exit_code) {
     {
         std::lock_guard<std::mutex> lk(mutex);
         auto it = mapping.find(name);
@@ -1502,7 +1495,7 @@ void server_models::on_child_exit(const std::string & name, const std::shared_pt
             return; // entry erased, or a newer instance took the name
         }
     }
-    if (mode == SERVER_CHILD_MODE_DOWNLOAD) {
+    if (role == SERVER_ROLE_DOWNLOAD) {
         // instance will be cleaned up on next load_models() call
         std::lock_guard<std::mutex> lk(mutex);
         stopping_models.erase(name);
@@ -2046,149 +2039,6 @@ void server_models::handle_child_state(const std::string & name, const std::stri
     }
 }
 
-//
-// server_child
-//
-
-server_child::server_child() {
-    if (is_child()) {
-        cmd_out = server_reserve_stdout();
-    }
-}
-
-server_child::~server_child() {
-    if (cmd_out) {
-        fclose(cmd_out);
-    }
-}
-
-bool server_child::is_child() {
-    const char * router_port = std::getenv("LLAMA_SERVER_ROUTER_PORT");
-    return router_port != nullptr;
-}
-
-server_child_mode server_child::get_mode() {
-    const char * mode = std::getenv("LLAMA_SERVER_CHILD_MODE");
-    std::string mode_str(mode ? mode : "");
-    if (mode_str == "download") {
-        return SERVER_CHILD_MODE_DOWNLOAD;
-    } else {
-        return SERVER_CHILD_MODE_NORMAL;
-    }
-}
-
-struct server_download_state : public common_download_callback {
-    server_child * self;
-    std::function<bool()> should_stop;
-    std::atomic<int64_t> last_progress_time{0}; // multiple files downloading in different threads
-    bool is_ok = false;
-
-    server_download_state(server_child * s) : self(s) {}
-
-    bool run(common_params & params) {
-        try {
-            common_models_handler handler = common_models_handler_init(params, LLAMA_EXAMPLE_SERVER);
-            common_models_handler_apply(handler, params, this);
-            is_ok = true;
-        } catch (const std::exception & e) {
-            auto model_name = params.model.get_name();
-            SRV_ERR("download failed for model name=%s: %s\n", model_name.c_str(), e.what());
-            is_ok = false;
-        }
-        return is_ok;
-    }
-    void on_progress(const common_download_progress & p) {
-        json data = {
-            {"url", p.url},
-            {"downloaded", p.downloaded},
-            {"total", p.total},
-        };
-        self->notify_to_router(server_state_to_str(SERVER_STATE_DOWNLOADING), data);
-    }
-    void on_start(const common_download_progress & p) override {
-        on_progress(p);
-    }
-    void on_update(const common_download_progress & p) override {
-        int64_t now = ggml_time_ms();
-        // throttle progress updates to avoid flooding logs
-        if (now - last_progress_time.load(std::memory_order_relaxed) >= 100) {
-            on_progress(p);
-            last_progress_time.store(now, std::memory_order_relaxed);
-        }
-    }
-    void on_done(const common_download_progress & p, bool) override {
-        on_progress(p);
-    }
-    bool is_cancelled() const override {
-        return should_stop ? should_stop() : false;
-    }
-};
-
-int server_child::run_download(common_params & params) {
-    auto cancelled = std::make_shared<std::atomic<bool>>(false);
-
-    // monitor stdin for cancellation command from the router
-    std::thread signal_thread = setup([cancelled](int) {
-        cancelled->store(true, std::memory_order_relaxed);
-    });
-
-    server_download_state dl(this);
-    dl.should_stop = [cancelled]() {
-        return cancelled->load(std::memory_order_relaxed);
-    };
-
-    bool ok = dl.run(params);
-
-    notify_to_router(server_state_to_str(SERVER_STATE_DOWNLOADING), {
-        {"result", ok ? "download_finished" : "download_failed"},
-    });
-
-    // router should send CMD_ROUTER_TO_CHILD_EXIT after receiving the result
-    if (signal_thread.joinable()) {
-        signal_thread.join();
-    }
-
-    SRV_INF("download completed %s\n", ok ? "successfully" : "with errors");
-    return 0;
-}
-
-std::thread server_child::setup(const std::function<void(int)> & shutdown_handler) {
-    // setup thread for monitoring stdin
-    return std::thread([shutdown_handler]() {
-        // wait for EOF on stdin
-        SRV_INF("%s", "child server monitoring thread started, waiting for EOF on stdin...\n");
-        bool eof = false;
-        while (true) {
-            std::string line;
-            if (!std::getline(std::cin, line)) {
-                // EOF detected, that means the router server is unexpectedly exit or killed
-                eof = true;
-                break;
-            }
-            if (line.find(CMD_ROUTER_TO_CHILD_EXIT) != std::string::npos) {
-                SRV_INF("%s", "exit command received, exiting...\n");
-                shutdown_handler(0);
-                break;
-            }
-        }
-        if (eof) {
-            SRV_INF("%s", "EOF on stdin detected, forcing shutdown...\n");
-            exit(1);
-        }
-    });
-}
-
-void server_child::notify_to_router(const std::string & state, const json & payload) {
-    json data = {
-        {"state", state},
-        {"payload", payload},
-    };
-    std::lock_guard<std::mutex> lk(mtx_stdout);
-    fprintf(cmd_out, "%s%s\n", CMD_CHILD_TO_ROUTER_STATE, safe_json_to_str(data).c_str());
-    fflush(cmd_out);
-}
-
-
 //
 // server_models_routes
 //
@@ -2617,7 +2467,7 @@ void server_models_routes::init_routes() {
         SRV_INF("starting download for model '%s'\n", name.c_str());
         {
             server_models::load_options load_opts;
-            load_opts.mode = SERVER_CHILD_MODE_DOWNLOAD;
+            load_opts.role = SERVER_ROLE_DOWNLOAD;
             load_opts.custom_meta = server_model_meta{};
             load_opts.custom_meta->source = SERVER_MODEL_SOURCE_CACHE;
             load_opts.custom_meta->name   = name;

tools/server/server-models.h

@@ -46,11 +46,6 @@ enum server_model_source {
     SERVER_MODEL_SOURCE_CACHE,
 };
 
-enum server_child_mode {
-    SERVER_CHILD_MODE_NORMAL,   // load the model and run normally
-    SERVER_CHILD_MODE_DOWNLOAD, // download the model and exit
-};
-
 static std::string server_model_status_to_string(server_model_status status) {
     switch (status) {
         case SERVER_MODEL_STATUS_DOWNLOADING: return "downloading";
@@ -171,7 +166,7 @@ private:
         server_model_meta meta;
         int req_count = 0; // number of active proxy requests
 
-        // ask the child to exit (it handles the command on its stdin, see server_child::setup)
+        // ask the child to exit (it handles the command on its stdin, see server_role::setup)
         void request_exit() const;
     };
 
@@ -276,7 +271,7 @@ private:
     void request_stop(const std::string & name, bool send_exit = true);
 
     // called by the monitor once a child exited and was reaped
-    void on_child_exit(const std::string & name, const std::shared_ptr<server_subproc> & proc, server_child_mode mode, int exit_code);
+    void on_child_exit(const std::string & name, const std::shared_ptr<server_subproc> & proc, server_role_type role, int exit_code);
 
     // notify SSE clients
     void notify_sse(const std::string & event, const std::string & model_id, const json & data = nullptr);
@@ -320,7 +315,7 @@ public:
     std::vector<server_model_meta> get_all_meta();
 
     struct load_options {
-        server_child_mode mode = SERVER_CHILD_MODE_NORMAL;
+        server_role_type role = SERVER_ROLE_INSTANCE;
         // used for spawning a downloading child process
         std::optional<server_model_meta> custom_meta = std::nullopt;
         // download mode: the sidecar kinds pulled with the model, as -hf would, and
@@ -369,7 +364,7 @@ public:
     // forward a websocket connection to the child, returns when either side closes
     void proxy_ws(const server_http_req & req, server_http_context::ws_conn & conn, const std::string & name);
 
-    // handle message sent from server_child::notify_to_router()
+    // handle message sent from server_role::notify_to_router()
     // raw input must starts with CMD_CHILD_TO_ROUTER_STATE, followed by a JSON string
     // called from the monitor thread
     // payload per state:
@@ -383,33 +378,6 @@ private:
     std::unique_ptr<server_monitor> monitor;
 };
 
-struct server_child {
-    // serializes the notify_to_router writes
-    std::mutex mtx_stdout;
-    std::atomic<bool> is_finished_downloading = false; // set by run_download
-
-    // in a child, keeps stdout for the commands to the router, so it is created before anything is written;
-    // everything else written to stdout goes to stderr with the logs
-    server_child();
-    ~server_child();
-
-    // return true if the current process is a child server instance
-    bool is_child();
-    server_child_mode get_mode();
-    int run_download(common_params & params);
-
-    // register the shutdown_handler to be called by the router
-    // return the monitoring thread (to be joined by the caller)
-    std::thread setup(const std::function<void(int)> & shutdown_handler);
-
-    // notify router server for status changes (e.g. loading, downloading, sleeping, etc.)
-    // message will be handled by server_models::handle_child_state() on the router side
-    void notify_to_router(const std::string & state_name, const json & payload);
-
-private:
-    FILE * cmd_out = nullptr; // the stdout the router reads the commands from
-};
-
 struct server_models_routes {
     common_params params;
     json ui_settings = json::object();     // Primary: new name

tools/server/server.cpp

@@ -37,12 +37,17 @@ static inline void signal_handler(int signal) {
     shutdown_handler(signal);
 }
 
+// how the server is run; everything is on when it is the llama-server binary
+struct server_run_opts {
+    bool install_signal_handlers = true; // off when embedded in llama-cli, which handles them itself
+    bool resolve_model           = true; // download / resolve -hf and docker models; off when the caller did it already
+};
+
 // satisfies -Wmissing-declarations (used by llama command)
 int llama_server(int argc, char ** argv);
 
 // to be used via CLI (argc / argv are used by router mode only)
 int llama_server(common_params & params, int argc, char ** argv);
-int llama_server(common_params & params, int argc, char ** argv, server_child & child);
 void llama_server_terminate();
 void llama_server_terminate() {
     if (shutdown_handler) {
@@ -91,72 +96,24 @@ static server_http_context::handler_t ex_wrapper(server_http_context::handler_t
     };
 }
 
-int llama_server(int argc, char ** argv) {
-    server_child child;
+//
+// shared by every role
+//
 
-    std::setlocale(LC_NUMERIC, "C");
-
-#ifndef _WIN32
-    // Ignore SIGPIPE so the server does not crash if a child (MCP server, tools runtime) exits while we are writing to its stdin
-    signal(SIGPIPE, SIG_IGN);
-#endif
-
-    // own arguments required by this example
-    common_params params;
-
-    common_init();
-
-    // start the stream session manager GC right after common init, before any HTTP route can
-    // touch it. lifecycle is symmetric, stop_gc() runs in clean_up() before backend free
-    server_stream_session_manager_start();
-
-    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SERVER)) {
-        return 1;
+static server_role_type server_role_detect(const common_params & params, server_role & role) {
+    if (role.is_managed_by_router()) {
+        return role.get_role();
     }
-
-    SRV_INF("%s", "initializing ...\n");
-
-    llama_backend_init();
-    llama_numa_init(params.numa);
-
-    const int result = llama_server(params, argc, argv, child);
-    common_log_flush(common_log_main());
-    return result;
+    // a server given no model is a router
+    const bool has_model = !params.model.path.empty()
+                        || !params.model.hf_repo.empty()
+                        || !params.model.docker_repo.empty();
+    return has_model ? SERVER_ROLE_INSTANCE : SERVER_ROLE_ROUTER;
 }
 
-int llama_server(common_params & params, int argc, char ** argv) {
-    server_child child;
-    return llama_server(params, argc, argv, child);
-}
-
-int llama_server(common_params & params, int argc, char ** argv, server_child & child) {
-    bool is_run_by_cli = (argv == nullptr);
-
-    common_models_handler models_handler;
-
-    // note: router mode also accepts -hf remote-preset, so we need to check that first
-    if (!is_run_by_cli && !params.model.hf_repo.empty()) {
-        try {
-            models_handler = common_models_handler_init(params, LLAMA_EXAMPLE_SERVER);
-            if (common_models_handler_is_preset_repo(models_handler)) {
-                // apply the preset and start the server in router mode
-                common_models_handler_apply(models_handler, params);
-            }
-        } catch (const std::exception & e) {
-            SRV_ERR("failed to fetch model metadata: %s\n", e.what());
-            return 1;
-        }
-    }
-
-    // router server never loads a model and must not touch the GPU
-    const bool is_router_server = params.model.path.empty()
-                               && params.model.hf_repo.empty()
-                               && params.model.docker_repo.empty();
-
-    // skip device enumeration so the CUDA primary context stays uncreated
-    common_params_print_info(params, !is_router_server);
-
-    if (!is_router_server) {
+// adjust the params for the role; returns false on a configuration error
+static bool server_prepare_params(common_params & params, server_role_type role_type) {
+    if (role_type != SERVER_ROLE_ROUTER) {
         // validate batch size for embeddings
         // embeddings require all tokens to be processed in a single ubatch
         // see https://github.com/ggml-org/llama.cpp/issues/12836
@@ -197,7 +154,7 @@ int llama_server(common_params & params, int argc, char ** argv, server_child &
         const std::string reason = server_connect::unavailable_reason(params);
         if (!reason.empty()) {
             SRV_ERR("--connect is not available: %s\n", reason.c_str());
-            return 1;
+            return false;
         }
         params.server_connect_code = server_connect::resolve_code(params.server_connect_code);
     } else if (!params.server_connect_code.empty()) {
@@ -206,77 +163,11 @@ int llama_server(common_params & params, int argc, char ** argv, server_child &
         params.server_connect_code.clear();
     }
 
-    // note: this is guaranteed to out-live ctx_http and tools
-    server_mcp mcp_mgr;
-
-    // struct that contains llama context and inference
-    server_context ctx_server;
-
-    //
-    // Router
-    //
-
-    // register API routes
-    server_routes routes(params, ctx_server);
-    server_tools tools;
-
-    std::optional<server_models_routes> models_routes{};
-
-    server_http_context ctx_http;
-    if (!ctx_http.init(params)) {
-        SRV_ERR("%s", "failed to initialize HTTP server\n");
-        return 1;
-    }
-
-    if (is_router_server) {
-        // setup server instances manager
-        try {
-            models_routes.emplace(params, argc, argv);
-        } catch (const std::exception & e) {
-            SRV_ERR("failed to initialize router models: %s\n", e.what());
-            return 1;
-        }
-
-        // proxy handlers
-        // note: routes.get_health stays the same
-        routes.get_metrics                 = models_routes->proxy_get;
-        routes.post_props                  = models_routes->proxy_post;
-        routes.post_completions            = models_routes->proxy_post;
-        routes.post_completions_oai        = models_routes->proxy_post;
-        routes.post_chat_completions       = models_routes->proxy_post;
-        routes.post_control                = models_routes->proxy_post;
-        routes.post_responses_oai          = models_routes->proxy_post;
-        routes.post_transcriptions_oai     = models_routes->proxy_post;
-        routes.post_tts                    = models_routes->proxy_post;
-        routes.ws_realtime                 = models_routes->proxy_ws;
-        routes.post_anthropic_messages     = models_routes->proxy_post;
-        routes.post_anthropic_count_tokens = models_routes->proxy_post;
-        routes.post_infill                 = models_routes->proxy_post;
-        routes.post_embeddings             = models_routes->proxy_post;
-        routes.post_embeddings_oai         = models_routes->proxy_post;
-        routes.post_rerank                 = models_routes->proxy_post;
-        routes.post_systemone              = models_routes->proxy_post;
-        routes.post_tokenize               = models_routes->proxy_post;
-        routes.post_detokenize             = models_routes->proxy_post;
-        routes.post_apply_template         = models_routes->proxy_post;
-        routes.post_chat_completions_tok   = models_routes->proxy_post;
-        routes.post_responses_tok_oai      = models_routes->proxy_post;
-        routes.get_lora_adapters           = models_routes->proxy_get;
-        routes.post_lora_adapters          = models_routes->proxy_post;
-        routes.get_slots                   = models_routes->proxy_get;
-        routes.post_slots                  = models_routes->proxy_post;
-
-        // custom routes for router
-        routes.get_props                   = models_routes->get_router_props;
-        routes.get_models                  = models_routes->get_router_models;
-
-        ctx_http.post("/models",               ex_wrapper(models_routes->post_router_models));
-        ctx_http.post("/models/load",          ex_wrapper(models_routes->post_router_models_load));
-        ctx_http.post("/models/unload",        ex_wrapper(models_routes->post_router_models_unload));
-        ctx_http.get ("/models/sse",           ex_wrapper(models_routes->get_router_models_sse));
-        ctx_http.del ("/models",               ex_wrapper(models_routes->del_router_models));
-    }
+    return true;
+}
 
+// the API of every server; a router points the handlers in routes at its proxies before calling this
+static void server_register_routes(server_http_context & ctx_http, server_routes & routes) {
     ctx_http.get ("/health",                   ex_wrapper(routes.get_health)); // public endpoint (no API key check)
     ctx_http.get ("/v1/health",                ex_wrapper(routes.get_health)); // public endpoint (no API key check)
     ctx_http.get ("/metrics",                  ex_wrapper(routes.get_metrics));
@@ -322,27 +213,30 @@ int llama_server(common_params & params, int argc, char ** argv, server_child &
     ctx_http.get ("/slots",                    ex_wrapper(routes.get_slots));
     ctx_http.post("/slots/:id_slot",           ex_wrapper(routes.post_slots));
 
-    // resumable streaming: a child binds the local session factories, the router binds
-    // proxies that resolve the owning child, see server-stream.h
-    server_http_context::handler_t stream_get_h;
-    server_http_context::handler_t streams_lookup_h;
-    server_http_context::handler_t stream_delete_h;
-    if (is_router_server) {
-        stream_get_h     = models_routes->router_stream_get;
-        streams_lookup_h = models_routes->router_streams_lookup;
-        stream_delete_h  = models_routes->router_stream_delete;
-    } else {
-        stream_get_h     = server_stream_make_get_handler();
-        streams_lookup_h = server_stream_make_lookup_handler();
-        stream_delete_h  = server_stream_make_delete_handler();
-    }
-    ctx_http.get ("/v1/stream",                ex_wrapper(stream_get_h));
-    ctx_http.post("/v1/streams/lookup",        ex_wrapper(streams_lookup_h));
-    ctx_http.del ("/v1/stream",                ex_wrapper(stream_delete_h));
-
     // Google Cloud Platform (Vertex AI) compat
     ctx_http.register_gcp_compat();
+}
 
+// resumable streaming, see server-stream.h: an instance binds the local session factories,
+// a router binds proxies that resolve the owning child
+static void server_register_stream_routes(
+        server_http_context & ctx_http,
+        const server_http_context::handler_t & stream_get,
+        const server_http_context::handler_t & streams_lookup,
+        const server_http_context::handler_t & stream_delete) {
+    ctx_http.get ("/v1/stream",                ex_wrapper(stream_get));
+    ctx_http.post("/v1/streams/lookup",        ex_wrapper(streams_lookup));
+    ctx_http.del ("/v1/stream",                ex_wrapper(stream_delete));
+}
+
+// the features behind a flag, and the warning about the ones that must not face untrusted clients
+// returns false when one of them fails to start
+static bool server_setup_features(
+        server_http_context & ctx_http,
+        const common_params & params,
+        server_mcp & mcp_mgr,
+        server_tools & tools,
+        server_role_type role_type) {
     // return 403 for disabled features
     server_http_context::handler_t res_403 = [](const server_http_req &) {
         auto res = std::make_unique<server_http_res>();
@@ -360,9 +254,8 @@ int llama_server(common_params & params, int argc, char ** argv, server_child &
         SRV_WRN("%s", "security: no API key is set and CORS allows all origins (see https://github.com/ggml-org/llama.cpp/pull/25655)\n");
     }
 
-    // CORS proxy (EXPERIMENTAL, only used by the Web UI for MCP)
     std::vector<std::string> warn_names;
-    if (is_router_server) {
+    if (role_type == SERVER_ROLE_ROUTER) {
         warn_names.push_back("router mode");
     }
 
@@ -370,6 +263,7 @@ int llama_server(common_params & params, int argc, char ** argv, server_child &
         warn_names.push_back("peer-to-peer tunnel (llama-connect, experimental)");
     }
 
+    // CORS proxy (EXPERIMENTAL, only used by the Web UI for MCP)
     if (params.ui_mcp_proxy) {
         ctx_http.get ("/cors-proxy",      ex_wrapper(proxy_handler_get));
         ctx_http.post("/cors-proxy",      ex_wrapper(proxy_handler_post));
@@ -383,7 +277,7 @@ int llama_server(common_params & params, int argc, char ** argv, server_child &
         mcp_mgr.start(params);
     } catch (const std::exception & e) {
         SRV_ERR("MCP starting failed: %s\n", e.what());
-        return 1;
+        return false;
     }
 
     if (!params.server_tools.empty() || !mcp_mgr.empty()) {
@@ -391,7 +285,7 @@ int llama_server(common_params & params, int argc, char ** argv, server_child &
             tools.setup(params.server_tools, mcp_mgr, params.server_tools_runtime);
         } catch (const std::exception & e) {
             SRV_ERR("tools setup failed: %s\n", e.what());
-            return 1;
+            return false;
         }
         ctx_http.get ("/tools",           ex_wrapper(tools.handle_get));
         ctx_http.post("/tools",           ex_wrapper(tools.handle_post));
@@ -418,15 +312,267 @@ int llama_server(common_params & params, int argc, char ** argv, server_child &
         SRV_WRN("security: %s enabled - do not expose to untrusted environments\n", features.c_str());
     }
 
-    //
-    // Handle downloading model
-    //
+    return true;
+}
 
-    if (child.is_child() && child.get_mode() == SERVER_CHILD_MODE_DOWNLOAD) {
-        return child.run_download(params);
-    } else if (!is_router_server && !is_run_by_cli) {
-        // single-model mode (NOT spawned by router)
-        // if this is invoked by CLI, model downloading should be already handled
+static void server_install_signal_handlers() {
+#if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__))
+    struct sigaction sigint_action;
+    sigint_action.sa_handler = signal_handler;
+    sigemptyset (&sigint_action.sa_mask);
+    sigint_action.sa_flags = 0;
+    sigaction(SIGINT, &sigint_action, NULL);
+    sigaction(SIGTERM, &sigint_action, NULL);
+#elif defined (_WIN32)
+    auto console_ctrl_handler = +[](DWORD ctrl_type) -> BOOL {
+        return (ctrl_type == CTRL_C_EVENT) ? (signal_handler(SIGINT), true) : false;
+    };
+    SetConsoleCtrlHandler(reinterpret_cast<PHANDLER_ROUTINE>(console_ctrl_handler), true);
+#endif
+}
+
+// once the HTTP server listens: log the addresses and spawn llama-connect if asked, so the port is
+// bound and the child startup health check reaches us. the destructor of connect_proc stops it
+static bool server_on_listening(server_http_context & ctx_http, const common_params & params, server_connect & connect_proc) {
+    for (const auto & address : ctx_http.listening_addresses) {
+        SRV_INF("listening on %s\n", address.c_str());
+    }
+
+    if (params.server_connect && !connect_proc.start(params, ctx_http.port)) {
+        SRV_ERR("%s", "exiting due to llama-connect error\n");
+        return false;
+    }
+    return true;
+}
+
+//
+// download worker: a child of the router that pulls a model into the cache and reports the progress
+//
+
+struct server_download_state : public common_download_callback {
+    server_role * role;
+    std::function<bool()> should_stop;
+    std::atomic<int64_t> last_progress_time{0}; // multiple files downloading in different threads
+    bool is_ok = false;
+
+    server_download_state(server_role * r) : role(r) {}
+
+    bool run(common_params & params) {
+        try {
+            common_models_handler handler = common_models_handler_init(params, LLAMA_EXAMPLE_SERVER);
+            common_models_handler_apply(handler, params, this);
+            is_ok = true;
+        } catch (const std::exception & e) {
+            auto model_name = params.model.get_name();
+            SRV_ERR("download failed for model name=%s: %s\n", model_name.c_str(), e.what());
+            is_ok = false;
+        }
+        return is_ok;
+    }
+    void on_progress(const common_download_progress & p) {
+        json data = {
+            {"url", p.url},
+            {"downloaded", p.downloaded},
+            {"total", p.total},
+        };
+        role->notify_to_router(server_state_to_str(SERVER_STATE_DOWNLOADING), data);
+    }
+    void on_start(const common_download_progress & p) override {
+        on_progress(p);
+    }
+    void on_update(const common_download_progress & p) override {
+        int64_t now = ggml_time_ms();
+        // throttle progress updates to avoid flooding logs
+        if (now - last_progress_time.load(std::memory_order_relaxed) >= 100) {
+            on_progress(p);
+            last_progress_time.store(now, std::memory_order_relaxed);
+        }
+    }
+    void on_done(const common_download_progress & p, bool) override {
+        on_progress(p);
+    }
+    bool is_cancelled() const override {
+        return should_stop ? should_stop() : false;
+    }
+};
+
+static int server_run_download(common_params & params, server_role & role) {
+    auto cancelled = std::make_shared<std::atomic<bool>>(false);
+
+    // monitor stdin for cancellation command from the router
+    std::thread signal_thread = role.setup([cancelled](int) {
+        cancelled->store(true, std::memory_order_relaxed);
+    });
+
+    server_download_state dl(&role);
+    dl.should_stop = [cancelled]() {
+        return cancelled->load(std::memory_order_relaxed);
+    };
+
+    bool ok = dl.run(params);
+
+    role.notify_to_router(server_state_to_str(SERVER_STATE_DOWNLOADING), {
+        {"result", ok ? "download_finished" : "download_failed"},
+    });
+
+    // router should send CMD_ROUTER_TO_CHILD_EXIT after receiving the result
+    if (signal_thread.joinable()) {
+        signal_thread.join();
+    }
+
+    SRV_INF("download completed %s\n", ok ? "successfully" : "with errors");
+    return 0;
+}
+
+//
+// router
+//
+
+static int server_run_router(common_params & params, int argc, char ** argv, const server_run_opts & opts) {
+    SRV_INF("%s", "starting server in router mode. models will be automatically loaded on-demand\n");
+
+    // note: this is guaranteed to out-live ctx_http and tools
+    server_mcp mcp_mgr;
+
+    // the router loads no model, its context stays empty; the routes it does not proxy still need one
+    server_context ctx_server;
+    server_routes  routes(params, ctx_server);
+    server_tools   tools;
+
+    server_http_context ctx_http;
+    if (!ctx_http.init(params)) {
+        SRV_ERR("%s", "failed to initialize HTTP server\n");
+        return 1;
+    }
+
+    // setup server instances manager
+    std::optional<server_models_routes> models_routes;
+    try {
+        models_routes.emplace(params, argc, argv);
+    } catch (const std::exception & e) {
+        SRV_ERR("failed to initialize router models: %s\n", e.what());
+        return 1;
+    }
+
+    // proxy handlers
+    // note: routes.get_health stays the same
+    routes.get_metrics                 = models_routes->proxy_get;
+    routes.post_props                  = models_routes->proxy_post;
+    routes.post_completions            = models_routes->proxy_post;
+    routes.post_completions_oai        = models_routes->proxy_post;
+    routes.post_chat_completions       = models_routes->proxy_post;
+    routes.post_control                = models_routes->proxy_post;
+    routes.post_responses_oai          = models_routes->proxy_post;
+    routes.post_transcriptions_oai     = models_routes->proxy_post;
+    routes.post_tts                    = models_routes->proxy_post;
+    routes.ws_realtime                 = models_routes->proxy_ws;
+    routes.post_anthropic_messages     = models_routes->proxy_post;
+    routes.post_anthropic_count_tokens = models_routes->proxy_post;
+    routes.post_infill                 = models_routes->proxy_post;
+    routes.post_embeddings             = models_routes->proxy_post;
+    routes.post_embeddings_oai         = models_routes->proxy_post;
+    routes.post_rerank                 = models_routes->proxy_post;
+    routes.post_systemone              = models_routes->proxy_post;
+    routes.post_tokenize               = models_routes->proxy_post;
+    routes.post_detokenize             = models_routes->proxy_post;
+    routes.post_apply_template         = models_routes->proxy_post;
+    routes.post_chat_completions_tok   = models_routes->proxy_post;
+    routes.post_responses_tok_oai      = models_routes->proxy_post;
+    routes.get_lora_adapters           = models_routes->proxy_get;
+    routes.post_lora_adapters          = models_routes->proxy_post;
+    routes.get_slots                   = models_routes->proxy_get;
+    routes.post_slots                  = models_routes->proxy_post;
+
+    // custom routes for router
+    routes.get_props                   = models_routes->get_router_props;
+    routes.get_models                  = models_routes->get_router_models;
+
+    ctx_http.post("/models",               ex_wrapper(models_routes->post_router_models));
+    ctx_http.post("/models/load",          ex_wrapper(models_routes->post_router_models_load));
+    ctx_http.post("/models/unload",        ex_wrapper(models_routes->post_router_models_unload));
+    ctx_http.get ("/models/sse",           ex_wrapper(models_routes->get_router_models_sse));
+    ctx_http.del ("/models",               ex_wrapper(models_routes->del_router_models));
+
+    server_register_routes(ctx_http, routes);
+    server_register_stream_routes(ctx_http,
+        models_routes->router_stream_get,
+        models_routes->router_streams_lookup,
+        models_routes->router_stream_delete);
+
+    if (!server_setup_features(ctx_http, params, mcp_mgr, tools, SERVER_ROLE_ROUTER)) {
+        return 1;
+    }
+
+    auto clean_up = [&models_routes, &mcp_mgr]() {
+        SRV_INF("%s: cleaning up before exit...\n", __func__);
+        // stop the session GC first, it finalizes live sessions and wakes pending readers
+        server_stream_session_manager_stop();
+        if (models_routes.has_value()) {
+            models_routes->stopping.store(true); // maybe redundant, but just to be safe
+            models_routes->models.unload_all();
+        }
+        mcp_mgr.shutdown();
+        llama_backend_free();
+    };
+
+    if (!ctx_http.start()) {
+        clean_up();
+        SRV_ERR("%s", "exiting due to HTTP server error\n");
+        return 1;
+    }
+    ctx_http.is_ready.store(true);
+
+    shutdown_handler = [&](int) {
+        if (models_routes.has_value()) {
+            // important to disconnect any SSE clients
+            models_routes->stopping.store(true);
+        }
+        mcp_mgr.shutdown();
+        ctx_http.stop();
+    };
+
+    try {
+        models_routes->models.load_startup_models();
+    } catch (const std::exception & e) {
+        SRV_ERR("failed to load models on startup: %s\n", e.what());
+        ctx_http.stop();
+        ctx_http.join();
+        clean_up();
+        return 1;
+    }
+
+    if (opts.install_signal_handlers) {
+        server_install_signal_handlers();
+    }
+
+    server_connect connect_proc;
+    if (!server_on_listening(ctx_http, params, connect_proc)) {
+        ctx_http.stop();
+        ctx_http.join();
+        clean_up();
+        return 1;
+    }
+
+    if (!params.models_preset_hf.empty()) {
+        SRV_WRN(      "NOTE: using preset.ini from HF repo '%s'\n", params.models_preset_hf.c_str());
+        SRV_WRN("%s", "      please only use presets that you can trust! Unknown presets may be unsafe\n");
+    }
+
+    ctx_http.join(); // keep the main thread alive
+
+    // when the HTTP server stops, clean up and exit
+    clean_up();
+
+    return 0;
+}
+
+//
+// instance
+//
+
+static int server_run_instance(common_params & params, server_role & role, common_models_handler & models_handler, const server_run_opts & opts) {
+    // -hf / docker: download if needed and resolve the local path. a child resolves what the router downloaded
+    if (opts.resolve_model) {
         try {
             common_models_handler_apply(models_handler, params);
         } catch (const std::exception & e) {
@@ -435,161 +581,188 @@ int llama_server(common_params & params, int argc, char ** argv, server_child &
         }
     }
 
-    //
-    // Start the server
-    //
+    // note: this is guaranteed to out-live ctx_http and tools
+    server_mcp mcp_mgr;
 
-    std::function<void()> clean_up;
+    // struct that contains llama context and inference
+    server_context ctx_server;
+    server_routes  routes(params, ctx_server);
+    server_tools   tools;
 
-    if (is_router_server) {
-        SRV_INF("%s", "starting server in router mode. models will be automatically loaded on-demand\n");
-
-        clean_up = [&models_routes, &mcp_mgr]() {
-            SRV_INF("%s: cleaning up before exit...\n", __func__);
-            // stop the session GC first, it finalizes live sessions and wakes pending readers
-            server_stream_session_manager_stop();
-            if (models_routes.has_value()) {
-                models_routes->stopping.store(true); // maybe redundant, but just to be safe
-                models_routes->models.unload_all();
-            }
-            mcp_mgr.shutdown();
-            llama_backend_free();
-        };
-
-        if (!ctx_http.start()) {
-            clean_up();
-            SRV_ERR("%s", "exiting due to HTTP server error\n");
-            return 1;
-        }
-        ctx_http.is_ready.store(true);
-
-        shutdown_handler = [&](int) {
-            if (models_routes.has_value()) {
-                // important to disconnect any SSE clients
-                models_routes->stopping.store(true);
-            }
-            mcp_mgr.shutdown();
-            ctx_http.stop();
-        };
-
-        try {
-            models_routes->models.load_startup_models();
-        } catch (const std::exception & e) {
-            SRV_ERR("failed to load models on startup: %s\n", e.what());
-            ctx_http.stop();
-            ctx_http.join();
-            clean_up();
-            return 1;
-        }
-
-    } else {
-        // setup clean up function, to be called before exit
-        clean_up = [&ctx_http, &ctx_server, &mcp_mgr]() {
-            SRV_INF("%s: cleaning up before exit...\n", __func__);
-            // stop the session GC first, it finalizes live sessions and wakes pending readers
-            server_stream_session_manager_stop();
-            ctx_http.stop();
-            ctx_server.terminate();
-            mcp_mgr.shutdown();
-            llama_backend_free();
-        };
-
-        // start the HTTP server before loading the model to be able to serve /health requests
-        if (!ctx_http.start()) {
-            clean_up();
-            SRV_ERR("%s", "exiting due to HTTP server error\n");
-            return 1;
-        }
-
-        // setup communication child --> router if necessary
-        if (child.is_child()) {
-            ctx_server.set_state_callback([&](server_state state, json payload) {
-                child.notify_to_router(server_state_to_str(state), payload);
-            });
-        }
-
-        if (!ctx_server.load_model(params)) {
-            clean_up();
-            ctx_http.join();
-            SRV_ERR("%s", "exiting due to model loading error\n");
-            return 1;
-        }
-
-        routes.update_meta(ctx_server);
-        ctx_http.is_ready.store(true);
-
-        SRV_INF("%s", "model loaded\n");
-
-        shutdown_handler = [&](int) {
-            mcp_mgr.shutdown();
-            // this will unblock start_loop()
-            ctx_server.terminate();
-        };
+    server_http_context ctx_http;
+    if (!ctx_http.init(params)) {
+        SRV_ERR("%s", "failed to initialize HTTP server\n");
+        return 1;
     }
 
-    // register signal handler if not running by CLI
-    if (!is_run_by_cli) {
-#if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__))
-        struct sigaction sigint_action;
-        sigint_action.sa_handler = signal_handler;
-        sigemptyset (&sigint_action.sa_mask);
-        sigint_action.sa_flags = 0;
-        sigaction(SIGINT, &sigint_action, NULL);
-        sigaction(SIGTERM, &sigint_action, NULL);
-#elif defined (_WIN32)
-        auto console_ctrl_handler = +[](DWORD ctrl_type) -> BOOL {
-            return (ctrl_type == CTRL_C_EVENT) ? (signal_handler(SIGINT), true) : false;
-        };
-        SetConsoleCtrlHandler(reinterpret_cast<PHANDLER_ROUTINE>(console_ctrl_handler), true);
-#endif
+    server_register_routes(ctx_http, routes);
+    server_register_stream_routes(ctx_http,
+        server_stream_make_get_handler(),
+        server_stream_make_lookup_handler(),
+        server_stream_make_delete_handler());
+
+    if (!server_setup_features(ctx_http, params, mcp_mgr, tools, SERVER_ROLE_INSTANCE)) {
+        return 1;
     }
 
-    for (const auto & address : ctx_http.listening_addresses) {
-        SRV_INF("listening on %s\n", address.c_str());
+    // setup clean up function, to be called before exit
+    auto clean_up = [&ctx_http, &ctx_server, &mcp_mgr]() {
+        SRV_INF("%s: cleaning up before exit...\n", __func__);
+        // stop the session GC first, it finalizes live sessions and wakes pending readers
+        server_stream_session_manager_stop();
+        ctx_http.stop();
+        ctx_server.terminate();
+        mcp_mgr.shutdown();
+        llama_backend_free();
+    };
+
+    // start the HTTP server before loading the model to be able to serve /health requests
+    if (!ctx_http.start()) {
+        clean_up();
+        SRV_ERR("%s", "exiting due to HTTP server error\n");
+        return 1;
+    }
+
+    // setup communication child --> router if necessary
+    if (role.is_managed_by_router()) {
+        ctx_server.set_state_callback([&](server_state state, json payload) {
+            role.notify_to_router(server_state_to_str(state), payload);
+        });
+    }
+
+    if (!ctx_server.load_model(params)) {
+        clean_up();
+        ctx_http.join();
+        SRV_ERR("%s", "exiting due to model loading error\n");
+        return 1;
+    }
+
+    routes.update_meta(ctx_server);
+    ctx_http.is_ready.store(true);
+
+    SRV_INF("%s", "model loaded\n");
+
+    shutdown_handler = [&](int) {
+        mcp_mgr.shutdown();
+        // this will unblock start_loop()
+        ctx_server.terminate();
+    };
+
+    if (opts.install_signal_handlers) {
+        server_install_signal_handlers();
     }
 
-    // spawn only once listening, so the port is bound and the child startup health check reaches us. the destructor stops it
     server_connect connect_proc;
-    if (params.server_connect && !connect_proc.start(params, ctx_http.port)) {
-        SRV_ERR("%s", "exiting due to llama-connect error\n");
+    if (!server_on_listening(ctx_http, params, connect_proc)) {
         ctx_http.stop();
         ctx_http.join();
         clean_up();
         return 1;
     }
 
-    if (is_router_server) {
-        if (!params.models_preset_hf.empty()) {
-            SRV_WRN(      "NOTE: using preset.ini from HF repo '%s'\n", params.models_preset_hf.c_str());
-            SRV_WRN("%s", "      please only use presets that you can trust! Unknown presets may be unsafe\n");
-        }
+    // optionally, notify router server that this instance is ready
+    std::thread monitor_thread;
+    if (role.is_managed_by_router()) {
+        monitor_thread = role.setup(shutdown_handler);
+        role.notify_to_router(server_state_to_str(SERVER_STATE_READY), routes.get_model_info());
+    }
 
-        ctx_http.join(); // keep the main thread alive
+    // this call blocks the main thread until queue_tasks.terminate() is called
+    ctx_server.start_loop();
 
-        // when the HTTP server stops, clean up and exit
-        clean_up();
-    } else {
-        // optionally, notify router server that this instance is ready
-        std::thread monitor_thread;
-        if (child.is_child()) {
-            monitor_thread = child.setup(shutdown_handler);
-            child.notify_to_router(server_state_to_str(SERVER_STATE_READY), routes.get_model_info());
-        }
+    clean_up();
+    ctx_http.join();
+    if (monitor_thread.joinable()) {
+        monitor_thread.join();
+    }
 
-        // this call blocks the main thread until queue_tasks.terminate() is called
-        ctx_server.start_loop();
-
-        clean_up();
-        ctx_http.join();
-        if (monitor_thread.joinable()) {
-            monitor_thread.join();
-        }
-
-        auto * ll_ctx = ctx_server.get_llama_context();
-        if (ll_ctx != nullptr) {
-            common_memory_breakdown_print(ll_ctx);
-        }
+    auto * ll_ctx = ctx_server.get_llama_context();
+    if (ll_ctx != nullptr) {
+        common_memory_breakdown_print(ll_ctx);
     }
 
     return 0;
 }
+
+//
+// entry points
+//
+
+static int server_run(common_params & params, int argc, char ** argv, server_role & role, const server_run_opts & opts) {
+    common_models_handler models_handler;
+
+    // -hf may name a preset repo, which makes this a router: settle that before deciding the role
+    if (opts.resolve_model && !params.model.hf_repo.empty()) {
+        try {
+            models_handler = common_models_handler_init(params, LLAMA_EXAMPLE_SERVER);
+            if (common_models_handler_is_preset_repo(models_handler)) {
+                // apply the preset and start the server in router mode
+                common_models_handler_apply(models_handler, params);
+            }
+        } catch (const std::exception & e) {
+            SRV_ERR("failed to fetch model metadata: %s\n", e.what());
+            return 1;
+        }
+    }
+
+    const server_role_type role_type = server_role_detect(params, role);
+
+    // only an instance loads a model; the others must not touch the GPU, so skip the
+    // device enumeration for them and the CUDA primary context stays uncreated
+    common_params_print_info(params, role_type == SERVER_ROLE_INSTANCE);
+
+    if (!server_prepare_params(params, role_type)) {
+        return 1;
+    }
+
+    switch (role_type) {
+        case SERVER_ROLE_DOWNLOAD: return server_run_download(params, role);
+        case SERVER_ROLE_ROUTER:   return server_run_router(params, argc, argv, opts);
+        case SERVER_ROLE_INSTANCE: return server_run_instance(params, role, models_handler, opts);
+    }
+
+    GGML_ABORT("invalid server role");
+}
+
+int llama_server(int argc, char ** argv) {
+    server_role role;
+
+    std::setlocale(LC_NUMERIC, "C");
+
+#ifndef _WIN32
+    // Ignore SIGPIPE so the server does not crash if a child (MCP server, tools runtime) exits while we are writing to its stdin
+    signal(SIGPIPE, SIG_IGN);
+#endif
+
+    // own arguments required by this example
+    common_params params;
+
+    common_init();
+
+    // start the stream session manager GC right after common init, before any HTTP route can
+    // touch it. lifecycle is symmetric, stop_gc() runs in clean_up() before backend free
+    server_stream_session_manager_start();
+
+    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SERVER)) {
+        return 1;
+    }
+
+    SRV_INF("%s", "initializing ...\n");
+
+    llama_backend_init();
+    llama_numa_init(params.numa);
+
+    const int result = server_run(params, argc, argv, role, server_run_opts{});
+    common_log_flush(common_log_main());
+    return result;
+}
+
+int llama_server(common_params & params, int argc, char ** argv) {
+    // embedded in llama-cli: it handles the signals, and the model is resolved already
+    server_run_opts opts;
+    opts.install_signal_handlers = false;
+    opts.resolve_model           = false;
+
+    server_role role;
+    return server_run(params, argc, argv, role, opts);
+}

tools/ui/src/lib/components/app/chat/ChatForm/ChatForm.svelte

@@ -477,12 +477,9 @@
 		queueMicrotask(() => inputRef?.focus());
 	}
 
-	// A flow outside the chat, e.g. the models manager closing onto it, asks for the
-	// composer back. The dialog's focus scope is still tearing down, so focus on a
-	// frame of its own.
+	// A flow outside the chat, e.g. the models page returning to it, asks for the
+	// composer back. Focus on a frame of its own, once the route has settled.
 	$effect(() => {
-		if (uiStore.manageModelsOpen) return;
-
 		if (!uiStore.composerFocusRequested) return;
 
 		uiStore.composerFocusRequested = false;

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugePopup.svelte

@@ -129,13 +129,13 @@
 {:else if gaugePopup.open}
 	<div
 		bind:this={cardEl}
-		class="absolute z-50 w-64 -translate-x-1/2 rounded-lg border border-border/50 bg-popover p-3 text-sm text-popover-foreground shadow-lg ring-1 ring-foreground/10"
+		class="absolute z-50 w-64 -translate-x-1/2 rounded-xl border border-border bg-popover p-1.5 text-sm text-popover-foreground shadow-md dark:border-border/20"
 		onpointerenter={gaugeCardEnter}
 		onpointerleave={gaugeCardLeave}
 		role="status"
 		style="left: {gaugePopup.centerX}px; bottom: {gaugePopup.bottom}px"
 	>
-		<div class="flex flex-col gap-2">
+		<div class="flex flex-col gap-2 px-2 py-1.5">
 			<div class="flex items-center gap-2">
 				<span class="font-medium">Context</span>
 

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockSearchResults.svelte

@@ -111,13 +111,8 @@
 			{#if showHoverCard}
 				{@const publishDate = formatPublishDate(result.published)}
 				{@const host = hostFor(safeUrl)}
-				<HoverCard.Content
-					align="start"
-					class="bg-popover text-popover-foreground z-50 w-80 max-w-[90vw] rounded-lg border p-0 shadow-lg"
-					side="top"
-					sideOffset={6}
-				>
-					<div class="flex flex-col gap-2 p-3">
+				<HoverCard.Content align="start" class="w-80 max-w-[90vw]" side="top" sideOffset={6}>
+					<div class="flex flex-col gap-2 px-2 py-1.5">
 						<a
 							class="line-clamp-3 text-sm font-medium leading-snug hover:underline"
 							href={safeUrl}