metal: implement DeepSeek V4 hyper-connections (#26459)

- Implement GGML_OP_DSV4_HC_COMB, GGML_OP_DSV4_HC_PRE, and
  GGML_OP_DSV4_HC_POST with SIMDgroup register and shuffle optimized kernels.
- Add Metal dispatch and support plumbing and test the production Sinkhorn
  iteration count and embedding width.

Assisted-by: Codex

Co-authored-by: Thiago Padilha <thiago@padilha.cc>

.github/workflows/build-cpu.yml

@@ -121,7 +121,7 @@ jobs:
     env:
       OPENBLAS_VERSION: 0.3.23
       SDE_VERSION: 9.33.0-2024-01-07
-      VULKAN_VERSION: 1.4.313.2
+      VULKAN_VERSION: 1.4.357.0
 
     strategy:
       matrix:

.github/workflows/release.yml

@@ -759,7 +759,7 @@ jobs:
 
     env:
       OPENBLAS_VERSION: 0.3.23
-      VULKAN_VERSION: 1.4.313.2
+      VULKAN_VERSION: 1.4.357.0
 
     strategy:
       matrix:

AGENTS.md

@@ -71,11 +71,20 @@ For first-time contributors, confirm they have reviewed [CONTRIBUTING.md](CONTRI
 These points are extremely important - failing to follow them won't necessarily get your PR rejected, but it will make reviewing take significantly longer. Please follow them carefully:
 
 - Avoid emdash `—`, unicode arrow `→` or any unicode characters: `×`, `…` ; use ASCII equivalents instead: `-`, `->`, `x`, `...`
-- Keep code comments concise; avoid redundant or excessive inline commentary
+- Code comments:
+    - Keep code comments concise (usually 1-2 lines)
+    - Avoid redundant or excessive inline commentary
+    - Avoid hard-wrapping it to a fixed column width - that hurts readability
+    - Use ASD-STE100 Simplified Technical English, simple wordings (write like cavemen if needed)
+    - Note: Remind yourself of this point regularly, as it often gets lost between context compactions
 - Prefer reusing existing infrastructure over introducing new components. Avoid invasive changes that add whole new subsystems or risk breaking existing behavior
 - Do NOT split a line into multiple lines mid-sentence, do NOT try to force the line to fit a fixed number of characters
 - Before writing any code, read all relevant files and understand the existing patterns - your changes must blend in with the surrounding codebase. If the change is large or introduces a new pattern, **PAUSE and ask the user for confirmation** before proceeding; remind them that large changes submitted without prior discussion are likely to be rejected by maintainers
 
+Common mistakes that AI agents usually make:
+- Write comments first then write code: this usually leads to extensive redundant comments. Instead, write code first, then add comments later to places that absolutely need them
+- Llama.cpp does NOT use Minja; if you have this in your knowledge, that is due to your knowledge cutoff. Llama.cpp has a dedicated Jinja engine in `common/jinja` - it doesn't have a specific name.
+
 ### Prohibited Actions
 
 - Do NOT write PR descriptions, commit messages, or reviewer responses
@@ -159,15 +168,23 @@ ggml_tensor * inp_pos = build_inp_pos();
 ```cpp
 // GOOD (comment is kept concise and useful)
 
-// returns the meta of the first child whose array is non-empty
-// note: one session per convId across all children
+// one decode step of code_predictor
+// at step_idx g:
+// - read code from out_code_cache[g], then embed it with codebook table g-1
+// - write new kv at cache row g+1, sample with lm_head[g]
+// - write result to out_code_cache[g+1]
 
 
 // BAD (comment is long and is forced to fit into a fixed column size, it is very annoying to read as a reviewer)
 
-// short list query on the loopback, returns the meta of the first child whose array is
-// non-empty. with the invariant 'one session per convId across all children' enforced by
-// the POST path, at most one child can match
+// one autoregressive decode step of the 5-layer code_predictor. See the
+// comment in models.h for the cache/tensor conventions this relies on.
+//
+// index mapping (derived from the reference pipeline-tts.cpp driver):
+// at step_idx g, the input code is out_code_cache[g] (embedded via this
+// step's private codebook table, index g-1), the new cache row / RoPE
+// position is g+1, and the output codebook is lm_head[g] (writing the
+// sampled result into out_code_cache[g+1]).
 ```
 
 Commit message:

common/arg.cpp

@@ -374,6 +374,10 @@ common_models_handler common_models_handler_init(const common_params & params, l
                                            params.speculative.types.end(),
                                            COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3) != params.speculative.types.end();
 
+    const bool spec_type_draft_dspark = std::find(params.speculative.types.begin(),
+                                           params.speculative.types.end(),
+                                           COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) != params.speculative.types.end();
+
     // only download mmproj if the current example is using it
     bool use_mmproj = false;
     for (const auto & ex : mmproj_examples) {
@@ -388,6 +392,7 @@ common_models_handler common_models_handler_init(const common_params & params, l
     opts.download_mtp    = spec_type_draft_mtp;
     opts.download_eagle3 = spec_type_draft_eagle3;
     opts.download_dflash = spec_type_draft_dflash;
+    opts.download_dspark = spec_type_draft_dspark;
     opts.download_mmproj = use_mmproj && !params.no_mmproj
                         && params.mmproj.path.empty() && params.mmproj.url.empty();
 
@@ -402,6 +407,7 @@ common_models_handler common_models_handler_init(const common_params & params, l
             opts_spec.download_mtp    = true;
             opts_spec.download_dflash = true;
             opts_spec.download_eagle3 = true;
+            opts_spec.download_dspark = true;
         }
         plan_spec = common_download_get_hf_plan(params.speculative.draft.mparams, opts_spec);
     }
@@ -544,12 +550,19 @@ void common_models_handler_apply(common_models_handler & handler, common_params
         plan_spec.mtp    = {};
         plan_spec.dflash = {};
         plan_spec.eagle3 = {};
+        plan_spec.dspark = {};
     }
 
     // infer the speculative type from the sidecar shipped by the draft repo when none is requested
     if (spec_types_is_default(params)) {
         if (!plan_spec.mtp.local_path.empty()) {
             params.speculative.types = { COMMON_SPECULATIVE_TYPE_DRAFT_MTP };
+            plan_spec.dspark = {};
+            plan_spec.dflash = {};
+            plan_spec.eagle3 = {};
+        } else if (!plan_spec.dspark.local_path.empty()) {
+            // dspark outranks dflash, its sidecar carries the extra Markov head
+            params.speculative.types = { COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK };
             plan_spec.dflash = {};
             plan_spec.eagle3 = {};
         } else if (!plan_spec.dflash.local_path.empty()) {
@@ -563,7 +576,8 @@ void common_models_handler_apply(common_models_handler & handler, common_params
     // when a sidecar type is requested, the draft repo resolves to its sidecar instead of a full model
     const bool spec_sidecar_found = !plan_spec.mtp.local_path.empty() ||
                                     !plan_spec.dflash.local_path.empty() ||
-                                    !plan_spec.eagle3.local_path.empty();
+                                    !plan_spec.eagle3.local_path.empty() ||
+                                    !plan_spec.dspark.local_path.empty();
     if (!plan_spec.mtp.local_path.empty() && !had_spec_url) {
         tasks.emplace_back(plan_spec.mtp, opts, [&]() {
             // only use the discovered MTP head when no draft path is set yet
@@ -594,6 +608,16 @@ void common_models_handler_apply(common_models_handler & handler, common_params
             }
         });
     }
+    if (!plan_spec.dspark.local_path.empty() && !had_spec_url) {
+        tasks.emplace_back(plan_spec.dspark, opts, [&]() {
+            // only use the discovered DSpark sidecar when no draft path is set yet
+            if (params.speculative.draft.mparams.path.empty()) {
+                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.dspark);
+            } else {
+                hf_cache::finalize_file(plan_spec.dspark);
+            }
+        });
+    }
 
     // a wired draft sidecar counts as an explicit draft for the main plan fallback below
     if (spec_sidecar_found) {
@@ -649,6 +673,16 @@ void common_models_handler_apply(common_models_handler & handler, common_params
             }
         });
     }
+    if (!plan.dspark.local_path.empty() && !had_spec_url) {
+        tasks.emplace_back(plan.dspark, opts, [&]() {
+            // only fall back to the discovered DSpark sidecar when no draft was explicitly provided
+            if (params.speculative.draft.mparams.empty()) {
+                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan.dspark);
+            } else {
+                hf_cache::finalize_file(plan.dspark);
+            }
+        });
+    }
     if (!plan.preset.local_path.empty()) {
         tasks.emplace_back(plan.preset, opts, [&]() {
             // if HF repo is a preset repo, we simply run server in router mode with the preset.ini file

common/chat-peg-parser.cpp

@@ -6,6 +6,9 @@
 
 #include <nlohmann/json.hpp>
 
+#include <cstdint>
+#include <functional>
+
 using ordered_json = nlohmann::ordered_json;
 
 static std::string_view trim_trailing_space(std::string_view sv, int max = -1) {
@@ -235,6 +238,43 @@ common_peg_parser common_chat_peg_builder::tag_with_safe_content(const std::stri
     return zero_or_more(choice({ p, content_chunk }));
 }
 
+common_peg_parser common_chat_peg_builder::permute(const std::string &                    rule_prefix,
+                                                   const std::vector<common_peg_parser> & parsers) {
+    if (parsers.empty()) {
+        return eps();
+    }
+
+    if (parsers.size() == 1 || parsers.size() > COMMON_CHAT_MAX_PERMUTE) {
+        return sequence(parsers);
+    }
+
+    std::map<uint32_t, common_peg_parser>      rules;
+    std::function<common_peg_parser(uint32_t)> remaining_of;
+
+    remaining_of = [&](uint32_t remaining) -> common_peg_parser {
+        if (remaining == 0) {
+            return eps();
+        }
+
+        auto cached = rules.find(remaining);
+        if (cached != rules.end()) {
+            return cached->second;
+        }
+
+        auto alternatives = choice();
+        for (size_t i = 0; i < parsers.size(); i++) {
+            const uint32_t bit = 1u << i;
+            if (remaining & bit) {
+                alternatives |= parsers[i] + remaining_of(remaining & ~bit);
+            }
+        }
+
+        return rules.emplace(remaining, rule(rule_prefix + "-" + std::to_string(remaining), alternatives)).first->second;
+    };
+
+    return remaining_of((1u << parsers.size()) - 1);
+}
+
 std::string & common_chat_peg_mapper::args_target() {
     return (current_tool && !current_tool->name.empty()) ? current_tool->arguments : args_buffer;
 }

common/chat-peg-parser.h

@@ -55,6 +55,8 @@ class common_chat_peg_minimax_m3_mapper : public common_chat_peg_mapper {
 struct content_structure;
 struct tool_call_structure;
 
+constexpr size_t COMMON_CHAT_MAX_PERMUTE = 6;
+
 class common_chat_peg_builder : public common_peg_parser_builder {
   public:
     // Tag constants (from former common_chat_peg_base_builder)
@@ -105,6 +107,9 @@ class common_chat_peg_builder : public common_peg_parser_builder {
     common_peg_parser tool_arg_json_value(const common_peg_parser & p) { return tag(TOOL_ARG_VALUE, p); }
 
 
+    // Matches every parser exactly once, in any order.
+    common_peg_parser permute(const std::string & rule_prefix, const std::vector<common_peg_parser> & parsers);
+
     // Return a parser that parses the prefix of a string, up to a given delimiter.
     common_peg_parser prefix(const std::string & s, const std::string & delimiter = {});
 

common/chat.cpp

@@ -1110,6 +1110,172 @@ static common_chat_params common_chat_params_init_ministral_3(const common_chat_
     return data;
 }
 
+static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_template &          tmpl,
+                                                              const autoparser::generation_params & inputs) {
+    common_chat_params data;
+
+    const std::string GEN_PREFIX = "<|im_start|>assistant\n";
+
+    data.prompt            = common_chat_template_direct_apply_impl(tmpl, inputs);
+    data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs);
+    data.format            = COMMON_CHAT_FORMAT_PEG_NATIVE;
+
+    auto supports_reasoning = tmpl.source().find("<think>") != std::string::npos;
+
+    data.supports_thinking = supports_reasoning;
+    data.preserved_tokens  = {
+        "<tool_call>",
+        "</tool_call>",
+    };
+
+    if (supports_reasoning) {
+        data.thinking_start_tag = "<think>";
+        // Support both </think> and <tool_call> as reasoning end sequences.
+        // <function= is omitted, as it is a workaround for Qwen3-Coder which is not a thinking model
+        data.thinking_end_tags = { "</think>", "<tool_call>" };
+        data.preserved_tokens.insert(data.preserved_tokens.end(), { "<think>", "</think>" });
+    }
+
+    data.message_delimiters = {
+        { COMMON_CHAT_ROLE_ASSISTANT, "<|im_start|>assistant"             },
+        { COMMON_CHAT_ROLE_TOOL,      "<|im_start|>user\n<tool_response>" }, // Qwen3-Coder, Qwen3.5, Nemotron Nano 3
+        { COMMON_CHAT_ROLE_TOOL,      "<|im_start|>tool_response"         }, // StepFun-3.5-Flash
+        { COMMON_CHAT_ROLE_USER,      "<|im_start|>user"                  },
+        { COMMON_CHAT_ROLE_SYSTEM,    "<|im_start|>system"                },
+    };
+
+    auto has_tools           = inputs.tools.is_array() && !inputs.tools.empty();
+    auto has_response_format = inputs.json_schema.is_object() && !inputs.json_schema.empty();
+    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);
+
+    if (inputs.has_continuation()) {
+        const auto & msg = inputs.continue_msg;
+
+        data.generation_prompt = GEN_PREFIX;
+        if (supports_reasoning) {
+            data.generation_prompt += "<think>\n" + msg.reasoning_content;
+            if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) {
+                data.generation_prompt += "\n</think>\n\n";
+            }
+        }
+        if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) {
+            data.generation_prompt += msg.render_content();
+        }
+
+        data.prompt += data.generation_prompt;
+    }
+
+    auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
+        auto generation_prompt = p.literal(GEN_PREFIX);
+
+        auto reasoning = p.eps();
+        if (supports_reasoning && extract_reasoning) {
+            reasoning = p.optional("<think>" + p.space() +
+                                   p.reasoning(p.until_one_of({ "</think>", "<tool_call>" })) +
+                                   (p.literal("</think>") | p.peek(p.literal("<tool_call>"))));
+        }
+
+        // Response format parser
+        if (has_response_format) {
+            return generation_prompt + (reasoning << p.content(p.schema(p.json(), "response-format", inputs.json_schema)));
+        }
+
+        // Tool call parser
+        if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
+            auto arg_close  = p.tool_arg_close(p.literal("\n</parameter>\n"));
+            auto arg_string = p.rule("xml-arg-string",
+                p.ac(p.tool_arg_string_value(p.until("\n</parameter>\n")) + arg_close, "\n</parameter>\n"));
+
+            auto tool_choice = p.choice();
+            foreach_function(inputs.tools, [&](const json & tool) {
+                const auto & function   = tool.at("function");
+                std::string  name       = function.at("name");
+                auto         parameters = function.contains("parameters") ? function.at("parameters") : json::object();
+
+                auto schema_info = common_schema_info();
+                schema_info.resolve_refs(parameters);
+
+                std::vector<common_peg_parser> required_args;
+                std::vector<common_peg_parser> optional_args;
+
+                foreach_parameter(function, [&](const std::string & param_name, const json & param_schema, bool is_required) {
+                    auto rule_name = "tool-" + name + "-arg-" + param_name;
+
+                    auto arg_open = p.tool_arg_open("<parameter=" + p.tool_arg_name(p.literal(param_name)) + ">\n");
+
+                    auto arg_value = schema_info.resolves_to_string(param_schema) ?
+                        arg_string :
+                        p.tool_arg_json_value(p.schema(p.json(), rule_name + "-schema", param_schema)) + arg_close;
+
+                    auto arg_rule = p.rule(rule_name, p.tool_arg(arg_open + arg_value));
+
+                    (is_required ? required_args : optional_args).push_back(arg_rule);
+                });
+
+                // Accept required arguments in any order, as Qwen does not always adhere to the
+                // order provided.
+                auto args = p.permute("tool-" + name + "-args", required_args);
+                if (!optional_args.empty()) {
+                    args = args + p.zero_or_more(p.choice(optional_args));
+                }
+
+                auto func = p.tool(p.tool_open("<function=" + p.tool_name(p.literal(name)) + ">\n") +
+                                   p.tool_args(args) +
+                                   p.tool_close(p.literal("</function>\n")));
+
+                tool_choice |= p.rule("tool-" + name, func);
+            });
+
+            auto min_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0;
+
+            // Qwen3-Coder models may occasionally omit the <tool_call> token.
+            auto tool_call_body  = tool_choice + "</tool_call>" + p.space();
+            auto tool_call_first = p.rule("tool-call-first", p.optional(p.literal("<tool_call>\n")) + tool_call_body);
+            auto tool_call       = p.rule("tool-call", "<tool_call>\n" + tool_call_body);
+
+            auto calls      = inputs.parallel_tool_calls ? tool_call_first + p.zero_or_more(tool_call) : tool_call_first;
+            auto tool_calls = p.trigger_rule("tool-call-root", p.repeat(calls, min_calls, 1));
+
+            return generation_prompt +
+                   (reasoning << p.content(p.until_one_of({ "<tool_call>", "<function=" })) << tool_calls);
+        }
+
+        // Content only parser
+        return generation_prompt + (reasoning << p.content(p.rest()));
+    });
+
+    data.parser = parser.save();
+
+    if (include_grammar) {
+        data.grammar_lazy = has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;
+
+        data.grammar = build_grammar([&](const common_grammar_builder & builder) {
+            foreach_function(inputs.tools, [&](const json & tool) {
+                const auto & function = tool.at("function");
+                auto         schema   = function.contains("parameters") ? function.at("parameters") : json::object();
+                builder.resolve_refs(schema);
+            });
+            if (has_response_format) {
+                auto schema = inputs.json_schema;
+                builder.resolve_refs(schema);
+            }
+            parser.build_grammar(builder, data.grammar_lazy);
+        });
+
+        if (data.grammar_lazy) {
+            data.grammar_triggers = {
+                { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, "<tool_call>" },
+                // Trigger on "<function" and not "<function=" because the trailing "=" is part of
+                // the token with the function name e.g. "=read"
+                { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, "<function"   },
+            };
+        }
+    }
+
+    return data;
+}
+
 static common_chat_params common_chat_params_init_gpt_oss(const common_chat_template &    tmpl,
                                                           const autoparser::generation_params & inputs) {
     common_chat_params data;
@@ -1943,18 +2109,6 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha
         adjusted_messages = deepseek_v4_sort_tool_results(inputs.messages);
     }
 
-    data.prompt             = common_chat_template_direct_apply_impl(tmpl, inputs, adjusted_messages);
-    data.generation_prompt  = common_chat_template_generation_prompt_impl(tmpl, inputs, adjusted_messages);
-    data.format             = COMMON_CHAT_FORMAT_PEG_NATIVE;
-    data.supports_thinking  = true;
-    data.thinking_start_tag = "<think>";
-    data.thinking_end_tags  = {"</think>"};
-    data.preserved_tokens   = {
-        "｜DSML｜",
-        "<think>",
-        "</think>",
-    };
-
     auto has_tools           = inputs.tools.is_array() && !inputs.tools.empty();
     auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object();
     auto extract_reasoning   = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE;
@@ -1972,6 +2126,18 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha
     const std::string PARAM_END    = "</" + DSML + "parameter>";
     const std::string GEN_PROMPT   = "<｜Assistant｜>";
 
+    data.prompt             = common_chat_template_direct_apply_impl(tmpl, inputs, adjusted_messages);
+    data.generation_prompt  = common_chat_template_generation_prompt_impl(tmpl, inputs, adjusted_messages);
+    data.format             = COMMON_CHAT_FORMAT_PEG_NATIVE;
+    data.supports_thinking  = true;
+    data.thinking_start_tag = THINK_START;
+    data.thinking_end_tags  = {THINK_END, FC_START};
+    data.preserved_tokens   = {
+        DSML,
+        THINK_START,
+        THINK_END,
+    };
+
     if (inputs.has_continuation()) {
         const auto & msg = inputs.continue_msg;
 
@@ -1983,13 +2149,101 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha
         data.prompt += data.generation_prompt;
     }
 
+    bool require_tools   = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED;
+    bool has_tool_calls = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE;
+
     auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
         auto generation_prompt = p.literal(GEN_PROMPT);
-        auto end = p.end();
+        auto end               = p.end();
+
+        // build tool call section first since we might need it in reasoning
+        auto tool_choice = p.choice();
+        if (has_tool_calls) {
+            foreach_function(inputs.tools, [&](const json & tool) {
+                const auto & function = tool.at("function");
+                std::string  name     = function.at("name");
+                auto         params   = function.contains("parameters") ? function.at("parameters") : json::object();
+                const auto & props    = params.contains("properties") ? params.at("properties") : json::object();
+
+                std::set<std::string> required;
+                if (params.contains("required")) {
+                    params.at("required").get_to(required);
+                }
+
+                auto schema_info = common_schema_info();
+                schema_info.resolve_refs(params);
+
+                std::vector<common_peg_parser> required_parsers;
+                std::vector<common_peg_parser> optional_parsers;
+                for (const auto & [param_name, param_schema] : props.items()) {
+                    bool is_required = required.find(param_name) != required.end();
+                    bool is_string   = schema_info.resolves_to_string(param_schema);
+
+                    auto arg = p.tool_arg(
+                        p.tool_arg_open(p.literal(PARAM_START + " name=\"") + p.tool_arg_name(p.literal(param_name)) +
+                                        p.literal("\" string=\"" + std::string(is_string ? "true" : "false") + "\">")) +
+                        (is_string ?
+                             p.tool_arg_string_value(p.until(PARAM_END)) :
+                             p.tool_arg_json_value(p.schema(p.json(), "tool-" + name + "-arg-" + param_name + "-schema",
+                                                            param_schema, false))) +
+                        p.tool_arg_close(p.literal(PARAM_END)));
+
+                    auto named_arg = p.rule("tool-" + name + "-arg-" + param_name, arg);
+                    if (is_required) {
+                        required_parsers.push_back(named_arg);
+                    } else {
+                        optional_parsers.push_back(named_arg);
+                    }
+                }
+
+                common_peg_parser args_seq = p.eps();
+                for (size_t i = 0; i < required_parsers.size(); i++) {
+                    if (i > 0) {
+                        args_seq = args_seq + p.space();
+                    }
+                    args_seq = args_seq + required_parsers[i];
+                }
+
+                if (!optional_parsers.empty()) {
+                    common_peg_parser any_opt = p.choice();
+                    for (const auto & opt : optional_parsers) {
+                        any_opt |= opt;
+                    }
+                    args_seq = args_seq + p.repeat(p.space() + any_opt, 0, -1);
+                }
+
+                common_peg_parser invoke_body = args_seq;
+                auto              func_parser = p.tool(p.tool_open(p.literal(INVOKE_START + " name=\"") +
+                                                                   p.tool_name(p.literal(name)) + p.literal("\">\n")) +
+                                                       invoke_body + p.space() + p.tool_close(p.literal(INVOKE_END)));
+
+                tool_choice |= p.rule("tool-" + name, func_parser);
+            });
+        }
+
+        common_peg_parser tool_calls = p.eps();
+        if (inputs.parallel_tool_calls) {
+            tool_calls = p.trigger_rule("tool-call",
+                p.literal(FC_START) + p.space() + tool_choice +
+                p.zero_or_more(p.space() + tool_choice) + p.space() + p.literal(FC_END));
+        } else {
+            tool_calls = p.trigger_rule("tool-call",
+                p.literal(FC_START) + p.space() + tool_choice + p.space() + p.literal(FC_END));
+        }
 
         auto reasoning = p.eps();
+        auto reasoning_with_tc = p.eps();
+        auto obligatory_tool_calls = tool_calls;
+        bool allow_reasoning_with_tc = false;
+
+        if (!require_tools) {
+            tool_calls = p.optional(tool_calls);
+        }
+
         if (extract_reasoning && inputs.enable_thinking) {
             reasoning = p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END);
+            reasoning_with_tc = THINK_START + p.reasoning(p.until_one_of({ FC_START, THINK_END })) + obligatory_tool_calls;
+            allow_reasoning_with_tc = true;
         } else if (extract_reasoning) {
             // Thinking disabled but reasoning extraction requested: the generation prompt
             // contains an empty <think></think> pair (V3.2) or a bare </think> (V4) that
@@ -2007,101 +2261,19 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha
             return generation_prompt + reasoning + response_format + end;
         }
 
-        if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) {
+        if (!has_tool_calls) {
             return generation_prompt + reasoning + p.content(p.rest()) + end;
         }
 
-        auto tool_choice = p.choice();
-        foreach_function(inputs.tools, [&](const json & tool) {
-            const auto & function = tool.at("function");
-            std::string  name     = function.at("name");
-            auto params   = function.contains("parameters") ? function.at("parameters") : json::object();
-            const auto & props    = params.contains("properties") ? params.at("properties") : json::object();
-
-            std::set<std::string> required;
-            if (params.contains("required")) {
-                params.at("required").get_to(required);
-            }
-
-            auto schema_info = common_schema_info();
-            schema_info.resolve_refs(params);
-
-            std::vector<common_peg_parser> required_parsers;
-            std::vector<common_peg_parser> optional_parsers;
-            for (const auto & [param_name, param_schema] : props.items()) {
-                bool is_required = required.find(param_name) != required.end();
-                bool is_string   = schema_info.resolves_to_string(param_schema);
-
-                auto arg = p.tool_arg(
-                    p.tool_arg_open(
-                        p.literal(PARAM_START + " name=\"") +
-                        p.tool_arg_name(p.literal(param_name)) +
-                        p.literal("\" string=\"" + std::string(is_string ? "true" : "false") + "\">")) +
-                    (is_string
-                         ? p.tool_arg_string_value(p.until(PARAM_END))
-                         : p.tool_arg_json_value(p.schema(p.json(),
-                                                          "tool-" + name + "-arg-" + param_name + "-schema",
-                                                          param_schema, false))) +
-                    p.tool_arg_close(p.literal(PARAM_END)));
-
-                auto named_arg = p.rule("tool-" + name + "-arg-" + param_name, arg);
-                if (is_required) {
-                    required_parsers.push_back(named_arg);
-                } else {
-                    optional_parsers.push_back(named_arg);
-                }
-            }
-
-            common_peg_parser args_seq = p.eps();
-            for (size_t i = 0; i < required_parsers.size(); i++) {
-                if (i > 0) {
-                    args_seq = args_seq + p.space();
-                }
-                args_seq = args_seq + required_parsers[i];
-            }
-
-            if (!optional_parsers.empty()) {
-                common_peg_parser any_opt = p.choice();
-                for (const auto & opt : optional_parsers) {
-                    any_opt |= opt;
-                }
-                args_seq = args_seq + p.repeat(p.space() + any_opt, 0, -1);
-            }
-
-            common_peg_parser invoke_body = args_seq;
-            auto func_parser = p.tool(
-                p.tool_open(p.literal(INVOKE_START + " name=\"") +
-                            p.tool_name(p.literal(name)) + p.literal("\">\n")) +
-                invoke_body + p.space() +
-                p.tool_close(p.literal(INVOKE_END)));
-
-            tool_choice |= p.rule("tool-" + name, func_parser);
-        });
-
-        auto require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED;
-
-        common_peg_parser tool_calls = p.eps();
-        if (inputs.parallel_tool_calls) {
-            tool_calls = p.trigger_rule("tool-call",
-                p.literal(FC_START) + p.space() + tool_choice +
-                p.zero_or_more(p.space() + tool_choice) + p.space() + p.literal(FC_END));
-        } else {
-            tool_calls = p.trigger_rule("tool-call",
-                p.literal(FC_START) + p.space() + tool_choice + p.space() + p.literal(FC_END));
-        }
-
-        if (!require_tools) {
-            tool_calls = p.optional(tool_calls);
-        }
-
-        auto content_before_tools = p.content(p.until(FC_START));
-        return generation_prompt + reasoning + content_before_tools + tool_calls + end;
+        auto content_before_tools = p.negate(p.literal(THINK_START)) + p.content(p.until(FC_START));
+        return allow_reasoning_with_tc ? generation_prompt + (reasoning_with_tc | (reasoning + content_before_tools + tool_calls)) + end :
+            generation_prompt + reasoning + content_before_tools + tool_calls + end;
     });
 
     data.parser = parser.save();
 
     if (include_grammar) {
-        data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED));
+        data.grammar_lazy = has_tools && !require_tools;
         data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
             foreach_function(inputs.tools, [&](const json & tool) {
                 const auto & function = tool.at("function");
@@ -3006,6 +3178,14 @@ std::optional<common_chat_params> common_chat_try_specialized_template(
         return common_chat_params_init_minicpm5(tmpl, params);
     }
 
+    // Qwen3-Coder XML tool calls, also used by Nemotron Nano 3, Qwen3.5 and StepFun-3.5-Flash
+    if (src.find("<tool_call>") != std::string::npos &&
+        src.find("<function=") != std::string::npos &&
+        src.find("<parameter=") != std::string::npos) {
+        LOG_DBG("Using specialized template: Qwen3-Coder\n");
+        return common_chat_params_init_qwen3_coder(tmpl, params);
+    }
+
     return std::nullopt;
 }
 

common/common.cpp

@@ -1620,6 +1620,7 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
     mparams.progress_callback           = params.load_progress_callback;
     mparams.progress_callback_user_data = params.load_progress_callback_user_data;
     mparams.no_alloc                    = params.no_alloc;
+    mparams.load_mtp                    = std::find(params.speculative.types.begin(), params.speculative.types.end(), COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end();
 
     return mparams;
 }

common/download.cpp

@@ -656,6 +656,12 @@ static hf_cache::hf_file find_best_dflash(const hf_cache::hf_files & files,
     return find_best_sibling(files, model, "dflash-", tag);
 }
 
+static hf_cache::hf_file find_best_dspark(const hf_cache::hf_files & files,
+                                          const std::string        & model,
+                                          const std::string        & tag = "") {
+    return find_best_sibling(files, model, "dspark-", tag);
+}
+
 static bool gguf_filename_is_model(const std::string & filepath) {
     if (!string_ends_with(filepath, ".gguf")) {
         return false;
@@ -670,7 +676,8 @@ static bool gguf_filename_is_model(const std::string & filepath) {
            filename.find("imatrix") == std::string::npos &&
            filename.find("mtp-")    == std::string::npos &&
            filename.find("eagle3-") == std::string::npos &&
-           filename.find("dflash-") == std::string::npos;
+           filename.find("dflash-") == std::string::npos &&
+           filename.find("dspark-") == std::string::npos;
 }
 
 static hf_cache::hf_file find_best_model(const hf_cache::hf_files & files,
@@ -763,7 +770,7 @@ common_download_hf_plan common_download_get_hf_plan(const common_params_model &
     } else {
         primary = find_best_model(all, tag);
         // a requested sidecar can resolve on its own, without a full model of the same tag
-        if (primary.path.empty() && !opts.download_mtp && !opts.download_dflash && !opts.download_eagle3) {
+        if (primary.path.empty() && !opts.download_mtp && !opts.download_dflash && !opts.download_eagle3 && !opts.download_dspark) {
             LOG_ERR("%s: no GGUF files found in repository %s\n", __func__, repo.c_str());
             list_available_gguf_files(all);
             return plan;
@@ -787,9 +794,12 @@ common_download_hf_plan common_download_get_hf_plan(const common_params_model &
     if (opts.download_eagle3) {
         plan.eagle3 = find_best_eagle3(all, primary.path, tag);
     }
+    if (opts.download_dspark) {
+        plan.dspark = find_best_dspark(all, primary.path, tag);
+    }
 
     if (primary.path.empty() &&
-        plan.mtp.local_path.empty() && plan.dflash.local_path.empty() && plan.eagle3.local_path.empty()) {
+        plan.mtp.local_path.empty() && plan.dflash.local_path.empty() && plan.eagle3.local_path.empty() && plan.dspark.local_path.empty()) {
         LOG_ERR("%s: no GGUF files found in repository %s\n", __func__, repo.c_str());
         list_available_gguf_files(all);
     }
@@ -967,7 +977,8 @@ std::vector<common_cached_model_info> common_list_cached_models() {
             split.prefix.find("mmproj")  != std::string::npos ||
             split.prefix.find("mtp-")    != std::string::npos ||
             split.prefix.find("eagle3-") != std::string::npos ||
-            split.prefix.find("dflash-") != std::string::npos) {
+            split.prefix.find("dflash-") != std::string::npos ||
+            split.prefix.find("dspark-") != std::string::npos) {
             continue;
         }
         if (seen.insert(f.repo_id + ":" + split.tag).second) {

common/download.h

@@ -59,6 +59,7 @@ struct common_download_opts {
     bool download_mtp     = false;
     bool download_eagle3  = false;
     bool download_dflash  = false;
+    bool download_dspark  = false;
     common_download_callback * callback = nullptr;
 };
 
@@ -110,6 +111,7 @@ struct common_download_hf_plan {
     hf_cache::hf_file mtp;
     hf_cache::hf_file eagle3;
     hf_cache::hf_file dflash;
+    hf_cache::hf_file dspark;
     hf_cache::hf_file preset; // if set, only this file is downloaded
 };
 common_download_hf_plan common_download_get_hf_plan(const common_params_model & model, const common_download_opts & opts);

common/speculative.cpp

@@ -1291,7 +1291,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
         GGML_ASSERT(ctx_tgt && ctx_dft && "MTP requires ctx_tgt and ctx_dft to be set");
 
         n_embd = llama_model_n_embd_out(llama_get_model(ctx_dft));
-        GGML_ASSERT(n_embd == llama_model_n_embd(llama_get_model(ctx_tgt)) &&
+        GGML_ASSERT(n_embd == llama_model_n_embd_out(llama_get_model(ctx_tgt)) &&
                 "MTP input row width must match the target h_nextn width");
         n_mtp_layers = std::max(1, (int) llama_model_n_layer_nextn(llama_get_model(ctx_dft)));
 

conversion/__init__.py

@@ -55,6 +55,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "DFlashDraftModel": "qwen",
     "Qwen3DSparkModel": "qwen",
     "DeepseekV4ForCausalLM": "deepseek",
+    "DeepseekV4DSparkModel": "deepseek",
     "DistilBertForMaskedLM": "bert",
     "DistilBertForSequenceClassification": "bert",
     "DistilBertModel": "bert",

conversion/deepseek.py

@@ -475,7 +475,10 @@ class DeepseekV32Model(DeepseekV2Model):
 @ModelBase.register("DeepseekV4ForCausalLM")
 class DeepseekV4Model(TextModel):
     model_arch = gguf.MODEL_ARCH.DEEPSEEK4
+    supports_mtp_export = True
     _skipped_mtp_tensors = 0
+    _dsv4_main_layers: int | None = None
+    _dsv4_nextn_layers: int = 0
 
     def __init__(self, *args, **kwargs):
         type(self)._skipped_mtp_tensors = 0
@@ -487,6 +490,8 @@ class DeepseekV4Model(TextModel):
             self.hparams.setdefault(key, value)
 
         self.block_count = self.hparams["num_hidden_layers"]
+        if self.mtp_only:
+            self.block_count += self.hparams.get("num_nextn_predict_layers", 0)
         self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
 
         self._dsv4_fp8_dequantized: set[str] = set()
@@ -504,13 +509,63 @@ class DeepseekV4Model(TextModel):
             with open(template_path, "r", encoding="utf-8") as f:
                 self.gguf_writer.add_chat_template(f.read())
 
+    def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
+        type(self)._dsv4_main_layers = self.hparams["num_hidden_layers"]
+        type(self)._dsv4_nextn_layers = self.hparams.get("num_nextn_predict_layers", 0)
+        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)
+
     @classmethod
     def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
-        name, _ = item
+        name, gen = item
         if name.startswith("mtp."):
-            cls._skipped_mtp_tensors += 1
-            return None
-        return super().filter_tensors(item)
+            if not cls.mtp_only:
+                cls._skipped_mtp_tensors += 1
+                return None
+
+            assert cls._dsv4_main_layers is not None
+            parts = name.split(".", 2)
+            if len(parts) < 3 or not parts[1].isdecimal():
+                raise ValueError(f"Unexpected DeepSeek-V4 MTP tensor {name!r}")
+
+            mtp_idx = int(parts[1])
+            if mtp_idx >= cls._dsv4_nextn_layers:
+                raise ValueError(f"Unexpected DeepSeek-V4 MTP layer {mtp_idx}")
+
+            bid = cls._dsv4_main_layers + mtp_idx
+            suffix = parts[2]
+            root_hc_head = {
+                "hc_head_fn",
+                "hc_head_base",
+                "hc_head_scale",
+            }
+            if suffix in root_hc_head:
+                name = suffix
+            elif suffix in (
+                "e_proj.weight", "e_proj.scale",
+                "h_proj.weight", "h_proj.scale",
+            ):
+                name = f"layers.{bid}.nextn.{suffix}"
+            elif suffix == "enorm.weight":
+                name = f"layers.{bid}.nextn.enorm.weight"
+            elif suffix == "hnorm.weight":
+                name = f"layers.{bid}.nextn.hnorm.weight"
+            elif suffix == "norm.weight":
+                name = f"layers.{bid}.nextn.shared_head_norm.weight"
+            else:
+                name = f"layers.{bid}.{suffix}"
+            return name, gen
+
+        if cls.mtp_only:
+            keep = name in (
+                "embed.weight",
+                "norm.weight",
+                "head.weight",
+                "head.scale",
+            )
+            if not keep:
+                return None
+
+        return super().filter_tensors((name, gen))
 
     @staticmethod
     def _float8_dtypes() -> tuple[torch.dtype, ...]:
@@ -565,6 +620,10 @@ class DeepseekV4Model(TextModel):
         self.gguf_writer.add_hyper_connection_sinkhorn_iterations(hparams["hc_sinkhorn_iters"])
         self.gguf_writer.add_hyper_connection_epsilon(hparams["hc_eps"])
         self.gguf_writer.add_hash_layer_count(hparams["num_hash_layers"])
+        if self.model_arch == gguf.MODEL_ARCH.DEEPSEEK4:
+            self.gguf_writer.add_embedding_length_out(hparams["hidden_size"] * hparams["hc_mult"])
+        if self.mtp_only and (num_nextn_predict_layers := hparams.get("num_nextn_predict_layers", 0)) > 0:
+            self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers)
 
     def dequant_model(self):
         fp8_dtypes = self._float8_dtypes()
@@ -669,12 +728,37 @@ class DeepseekV4Model(TextModel):
         if self._dsv4_mxfp4_generated:
             return ()
 
-        consumed: list[str] = self._write_hash_routing_tensors()
+        consumed: list[str] = []
+        main_layers = self.hparams["num_hidden_layers"]
+        if not self.mtp_only:
+            consumed.extend(self._write_hash_routing_tensors())
+        elif self.hparams["num_hash_layers"] > 0:
+            for bid in range(self.hparams["num_hash_layers"]):
+                name = f"layers.{bid}.ffn.gate.tid2eid"
+                if name in self.model_tensors:
+                    consumed.extend(self._write_hash_routing_tensors())
+                    break
+
         for bid in range(self.block_count):
+            if self.mtp_only and bid < main_layers:
+                continue
             consumed.extend(self._write_mxfp4_expert_tensor(bid, "w1", gguf.MODEL_TENSOR.FFN_GATE_EXP))
             consumed.extend(self._write_mxfp4_expert_tensor(bid, "w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP))
             consumed.extend(self._write_mxfp4_expert_tensor(bid, "w3", gguf.MODEL_TENSOR.FFN_UP_EXP))
 
+        for bid in range(main_layers, self.block_count):
+            e_name = f"layers.{bid}.nextn.e_proj.weight"
+            h_name = f"layers.{bid}.nextn.h_proj.weight"
+            if e_name not in self.model_tensors and h_name not in self.model_tensors:
+                continue
+            if e_name not in self.model_tensors or h_name not in self.model_tensors:
+                raise KeyError(f"Missing DeepSeek-V4 MTP e/h projection pair for block {bid}")
+
+            e_proj = LazyTorchTensor.to_eager(self.model_tensors[e_name]())
+            h_proj = LazyTorchTensor.to_eager(self.model_tensors[h_name]())
+            yield (f"layers.{bid}.nextn.eh_proj.weight", torch.cat((e_proj, h_proj), dim=1).contiguous())
+            consumed.extend((e_name, h_name))
+
         for name in consumed:
             del self.model_tensors[name]
 
@@ -737,6 +821,12 @@ class DeepseekV4Model(TextModel):
             "ffn.shared_experts.w1.weight": (gguf.MODEL_TENSOR.FFN_GATE_SHEXP, ".weight"),
             "ffn.shared_experts.w2.weight": (gguf.MODEL_TENSOR.FFN_DOWN_SHEXP, ".weight"),
             "ffn.shared_experts.w3.weight": (gguf.MODEL_TENSOR.FFN_UP_SHEXP, ".weight"),
+            "nextn.eh_proj.weight": (gguf.MODEL_TENSOR.NEXTN_EH_PROJ, ".weight"),
+            "nextn.enorm.weight": (gguf.MODEL_TENSOR.NEXTN_ENORM, ".weight"),
+            "nextn.hnorm.weight": (gguf.MODEL_TENSOR.NEXTN_HNORM, ".weight"),
+            "nextn.shared_head_norm.weight": (gguf.MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ".weight"),
+            "nextn.embed_tokens.weight": (gguf.MODEL_TENSOR.NEXTN_EMBED_TOKENS, ".weight"),
+            "nextn.shared_head_head.weight": (gguf.MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, ".weight"),
         }
 
         tensor_name = match.group(2)
@@ -759,10 +849,12 @@ class DeepseekV4Model(TextModel):
         return [(self._format_dsv4_tensor_name(tensor_key, bid, suffix), data_torch)]
 
     def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
-        del new_name, bid  # unused
+        del bid  # unused
 
         if name in self._dsv4_fp8_dequantized and n_dims >= 2:
             return gguf.GGMLQuantizationType.Q8_0
+        if new_name.endswith(".nextn.eh_proj.weight"):
+            return gguf.GGMLQuantizationType.Q8_0
         if name in self._dsv4_f32_tensors:
             return gguf.GGMLQuantizationType.F32
         if name in self._dsv4_bf16_tensors and n_dims >= 2:
@@ -770,7 +862,122 @@ class DeepseekV4Model(TextModel):
 
         return False
 
+    def prepare_metadata(self, vocab_only: bool):
+        from_dir = self.fname_out.is_dir()
+        super().prepare_metadata(vocab_only=vocab_only)
+
+        if not self.mtp_only or not from_dir:
+            return
+
+        output_type: str = self.ftype.name.partition("_")[2]
+        fname_default: str = gguf.naming_convention(
+            self.metadata.name, self.metadata.basename, self.metadata.finetune,
+            self.metadata.version, size_label=None, output_type=output_type, model_type=None)
+        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
+
     def prepare_tensors(self):
         super().prepare_tensors()
         self._is_mxfp4 = True
         self.ftype = gguf.LlamaFileType.MOSTLY_MXFP4_MOE
+
+
+@ModelBase.register("DeepseekV4DSparkModel")
+class DeepseekV4DSparkModel(DeepseekV4Model):
+    model_arch = gguf.MODEL_ARCH.DFLASH
+
+    _DSPARK_ROOT_MAP: dict[str, tuple[gguf.MODEL_TENSOR, str]] = {
+        "main_proj.weight": (gguf.MODEL_TENSOR.FC, ".weight"),
+        "main_norm.weight": (gguf.MODEL_TENSOR.ENC_OUTPUT_NORM, ".weight"),
+        "markov_head.markov_w1.weight": (gguf.MODEL_TENSOR.DSPARK_MARKOV_W1, ".weight"),
+        "markov_head.markov_w2.weight": (gguf.MODEL_TENSOR.DSPARK_MARKOV_W2, ".weight"),
+        "confidence_head.proj.weight": (gguf.MODEL_TENSOR.DSPARK_CONF_PROJ, ".weight"),
+    }
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.block_count = 1 + max(
+            int(match.group(1)) for name in self.model_tensors
+            if (match := re.match(r"layers\.(\d+)\.", name))
+        )
+        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
+
+        self.hparams["compress_ratios"] = [0] * self.block_count
+        self.hparams["num_hash_layers"] = 0
+
+    def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
+        if remote_hf_model_id is None:
+            return super().index_tensors()
+
+        with open(self.dir_model / "model.safetensors.index.json", "r", encoding="utf-8") as f:
+            weight_map = json.load(f)["weight_map"]
+
+        part_names = sorted({
+            part_name for name, part_name in weight_map.items()
+            if name.startswith("mtp.")
+        })
+        tensors: dict[str, Callable[[], Tensor]] = {}
+
+        for part_name in part_names:
+            from huggingface_hub import hf_hub_download
+
+            logger.info("gguf: caching remote DSpark part '%s'", part_name)
+            part_path = Path(hf_hub_download(repo_id=remote_hf_model_id, filename=part_name))
+            with gguf.utility.SafetensorsLocal(part_path) as model_part:
+                for name in model_part:
+                    data = model_part[name]
+                    data_gen = lambda data=data: LazyTorchTensor.from_local_tensor(data)  # noqa: E731
+                    if titem := self.filter_tensors((name, data_gen)):
+                        tensor_name, tensor_gen = titem
+                        tensors[tensor_name] = tensor_gen
+
+        return tensors
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+        if not name.startswith("mtp."):
+            return None
+        return super().filter_tensors((cls._rekey_mtp_tensor_name(name), gen))
+
+    @staticmethod
+    def _rekey_mtp_tensor_name(name: str) -> str:
+        match = re.match(r"mtp\.(\d+)\.(.+)$", name)
+        if match is None:
+            raise ValueError(f"Unexpected DSpark tensor {name!r}")
+
+        stage, rest = match.group(1), match.group(2)
+        root_names = (
+            "main_proj.scale",
+            "norm.weight",
+            "hc_head_fn",
+            "hc_head_base",
+            "hc_head_scale",
+        )
+        if rest in DeepseekV4DSparkModel._DSPARK_ROOT_MAP or rest in root_names:
+            return rest
+        return f"layers.{stage}.{rest}"
+
+    def _map_dsv4_tensor_name(self, name: str, bid: int | None) -> tuple[gguf.MODEL_TENSOR, str]:
+        if name in self._DSPARK_ROOT_MAP:
+            return self._DSPARK_ROOT_MAP[name]
+        return super()._map_dsv4_tensor_name(name, bid)
+
+    def set_vocab(self):
+        if self.target_model_dir is None:
+            raise ValueError("DeepSeek-V4 DSpark requires --target-model-dir with the target tokenizer")
+
+        original_dir = self.dir_model
+        try:
+            self.dir_model = self.target_model_dir
+            super().set_vocab()
+        finally:
+            self.dir_model = original_dir
+
+        self.gguf_writer.add_mask_token_id(self.hparams["dspark_noise_token_id"])
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+
+        self.gguf_writer.add_block_size(self.hparams["dspark_block_size"])
+        self.gguf_writer.add_target_layers([layer + 1 for layer in self.hparams["dspark_target_layer_ids"]])

conversion/minicpm.py

@@ -137,6 +137,15 @@ class MiniCPMV4_6TextModel(Qwen3_5TextModel):
 class MiniCPMV4_6VisionModel(MmprojModel):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self.downsample_mode = self.preprocessor_config.get("downsample_mode", "16x")
+        if self.downsample_mode not in {"4x", "16x"}:
+            raise ValueError(f"Unsupported downsample mode: {self.downsample_mode}")
+        if self.downsample_mode == "4x":
+            self.model_tensors = {
+                name: tensor for name, tensor in self.model_tensors.items()
+                if ".vit_merger." not in name
+            }
+
         if self.hparams_vision is not None:
             # In MiniCPM-V 4.6 `vision_config.image_size` (980) describes the SigLIP
             # positional embedding bucket grid (70 x 70), while the per-slice processing
@@ -156,8 +165,8 @@ class MiniCPMV4_6VisionModel(MmprojModel):
         # (mapped to PROJECTOR_TYPE_MINICPMV4_6).
         self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.MINICPMV4_6)
 
-        # ViT merger 2x2 + final merger 2x2 = 4x spatial merge per dimension; used for slice alignment
-        self.gguf_writer.add_vision_projector_scale_factor(4)
+        self.gguf_writer.add_vision_projector_scale_factor(
+            2 if self.downsample_mode == "4x" else 4)
 
         # borrow wa_layer_indexes for vit_merger insertion point
         insert_layer_id = int(self.global_config.get(

convert_hf_to_gguf.py

@@ -122,8 +122,12 @@ def parse_args() -> argparse.Namespace:
         help="Export only the multi-token prediction (MTP) head as a separate GGUF, suitable for use as a speculative draft. An 'mtp-' prefix will be added to the output file name.",
     )
     parser.add_argument(
-        "--no-mtp", action="store_true",
-        help="Exclude the multi-token prediction (MTP) head from the converted GGUF. Pair with --mtp on a second run to publish trunk and MTP as two files. Note: the split form duplicates embeddings, but even though the bundled default is more space-efficient overall, this allows differing quantization which may be more performant.",
+        "--no-nextn", "--no-mtp", dest="no_mtp", action="store_true",
+        help="Exclude NextN speculative draft tensors from the converted GGUF. Pair with --mtp or --dspark on a second run to publish target and draft as two files.",
+    )
+    parser.add_argument(
+        "--dspark", action="store_true",
+        help="Export only the DeepSeek-V4 DSpark draft tensors as a separate GGUF.",
     )
     parser.add_argument(
         "--mistral-format", action="store_true",
@@ -254,13 +258,20 @@ def main() -> None:
             from conversion.mistral import MistralModel
             model_class = MistralModel
 
-        if args.mtp and args.no_mtp:
-            logger.error("--mtp and --no-mtp are mutually exclusive")
+        if sum((args.mtp, args.no_mtp, args.dspark)) > 1:
+            logger.error("--mtp, --no-nextn, and --dspark are mutually exclusive")
             sys.exit(1)
 
+        if args.dspark:
+            if is_mistral_format or model_architecture != "DeepseekV4ForCausalLM":
+                logger.error("--dspark is only supported for DeepseekV4ForCausalLM")
+                sys.exit(1)
+            from conversion.deepseek import DeepseekV4DSparkModel
+            model_class = DeepseekV4DSparkModel
+
         if args.mtp or args.no_mtp:
             if not model_class.supports_mtp_export:
-                logger.error("--mtp / --no-mtp are not supported for %s", model_architecture)
+                logger.error("--mtp / --no-nextn are not supported for %s", model_architecture)
                 sys.exit(1)
             if args.no_mtp:
                 model_class.no_mtp = True

docs/backend/SYCL.md

@@ -788,7 +788,7 @@ use 1 SYCL GPUs: [0] with Max compute units:512
 | Name              | Value            | Function                                                                                                                  |
 |-------------------|------------------|---------------------------------------------------------------------------------------------------------------------------|
 | GGML_SYCL_DEBUG   | 0 (default) or 1 | Enable log function by macro: GGML_SYCL_DEBUG                                                                             |
-| GGML_SYCL_DEV2DEV_MEMCPY | 0 (default) or 1 | Choose the SYCL or L0 API in dev2dev memory copy.<br>Value: <br>*  0: SYCL API (default)<br>* 1: L0 API -- L0 API is found to lead to abnormal crash in some case. This debug flag is used to check the issue.|
+| GGML_SYCL_DEV2DEV_MEMCPY | 0 (default), 1, 2 | Choose the method of dev2dev memory copy.<br>Value: <br>*  0: SYCL API (default), only support dGPUs.<br>* 1: L0 API -- Better performance, only support dGPUs, found to lead to abnormal crash in some case. <br>* 2: Host Forward -- Most stable method for all cases (including iGPU + dGPU*N), but with lower performance (-2% to -5%).<br>SYCL & L0 API are easy to be impacted by Intel GPU driver issue. When you meet the garbled output or crash issues in multiple GPUs case, try with this debug flag to work around or check the issue.|
 | GGML_SYCL_ENABLE_FLASH_ATTN | 1 (default) or 0| Enable Flash-Attention. It can reduce memory usage. The performance impact depends on the LLM.|
 | GGML_SYCL_ENABLE_OPT | 0 or 1 (default)| Enable optimize features for Intel GPUs. (Recommended to 0 for Intel devices older than Gen 10) |
 | GGML_SYCL_ENABLE_GRAPH | 0 (default) or 1 | Enable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still on development, no better performance. |
@@ -797,6 +797,9 @@ use 1 SYCL GPUs: [0] with Max compute units:512
 | GGML_SYCL_FA_ONEDNN | 1 (default) or 0 | Enable the oneDNN fused SDPA (flash-attention) path on supported GPUs. Set to 0 to always use the native SYCL flash-attention kernel. |
 | GGML_SYCL_FA_ONEDNN_MAX_KV | 0 (default, disabled) or positive integer | By default (0), all sequences are handled by the oneDNN fused SDPA path, regardless of KV length; a positive value caps that length, past which sequences fall back to the native kernel. If GPU driver watchdog resets (DEVICE_LOST) occur during long-context inference, set this near the context depth where they start, e.g. 24576. |
 | GGML_SYCL_ENABLE_VMM | 0 or 1 (default) | Enable the virtual-memory device pool. |
+| GGML_SYCL_ENABLE_MKL_FA | 1 (default) or 0 | Enable oneMKL GEMM flash attention for XMX-accelerated prompt processing with quantized KV cache. Automatically activates during prefill (prompt processing) when all conditions are met: (1) flash-attn enabled (`-fa` or `--flash-attn on`), (2) KV cache quantized (`--cache-type-k q8_0 --cache-type-v q8_0` or other `*_0/*_1` types), (3) batch size ≥ 1024 (`--batch-size 1024`), (4) prompt length ≥ 1024 tokens. Set to 0 to force the TILE kernel for A/B testing. Example minimum command: `llama-cli -m model.gguf -fa -ngl 99 --cache-type-k q8_0 --cache-type-v q8_0 --batch-size 1024 -p "your prompt"` |
+| GGML_SYCL_MKL_FA_DEBUG | 0 (default) or 1 | Enable per-call diagnostic logging for MKL flash attention: GEMM/softmax timings, interleaved-head detection, and buffer memory usage. |
+| GGML_SYCL_MKL_FA_DIAG | 0 (default) or 1 | Enable output fingerprinting for MKL flash attention. Dumps the first 64 float output values for the first 6 FA calls with n_kv ≥ 1024, labeled with kernel type (MKL/TILE/VEC) for cross-kernel comparison. |
 | GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute (currently top-k MoE gating). |
 | ZES_ENABLE_SYSMAN | 0 (default) or 1 | Support to get free memory of GPU by sycl::aspect::ext_intel_free_memory.<br>Recommended to use when --split-mode = layer |
 | UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS | 0 (default) or 1 | Allow SYCL/Unified Runtime Level Zero device allocations larger than 4 GiB. llama.cpp's direct Level Zero allocation path requests the relaxed maximum-size limit itself when GGML_SYCL_ENABLE_LEVEL_ZERO=1. |

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

@@ -130,43 +130,27 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
         }
 
         const block_q2_0 * bxi = (const block_q2_0 *) x + kbx0 + i*stride + kbx;
-        // Each 32-element chunk occupies 8 bytes of qs (32 elements * 2 bits = 64 bits)
-        const int qs_offset = 8*kqsx;
-        const int qs0 = bxi->qs[qs_offset + 0] | (bxi->qs[qs_offset + 1] << 8) |
-                        (bxi->qs[qs_offset + 2] << 16) | (bxi->qs[qs_offset + 3] << 24);
-        const int qs1 = bxi->qs[qs_offset + 4] | (bxi->qs[qs_offset + 5] << 8) |
-                        (bxi->qs[qs_offset + 6] << 16) | (bxi->qs[qs_offset + 7] << 24);
-
-        // Unpack 32 2-bit codes into 8 int32s, each holding 4 signed int8s in {-1,0,1,2}.
-        int unpacked_bytes[8];
-#pragma unroll
-        for (int j = 0; j < 4; ++j) {
-            const int shift = j * 8;
-            const int codes = (qs0 >> shift) & 0xFF;
-            const int c0 = ((codes >> 0) & 0x3) - 1;
-            const int c1 = ((codes >> 2) & 0x3) - 1;
-            const int c2 = ((codes >> 4) & 0x3) - 1;
-            const int c3 = ((codes >> 6) & 0x3) - 1;
-            unpacked_bytes[j] = (c0 & 0xFF) | ((c1 & 0xFF) << 8) | ((c2 & 0xFF) << 16) | ((c3 & 0xFF) << 24);
-        }
-#pragma unroll
-        for (int j = 0; j < 4; ++j) {
-            const int shift = j * 8;
-            const int codes = (qs1 >> shift) & 0xFF;
-            const int c0 = ((codes >> 0) & 0x3) - 1;
-            const int c1 = ((codes >> 2) & 0x3) - 1;
-            const int c2 = ((codes >> 4) & 0x3) - 1;
-            const int c3 = ((codes >> 6) & 0x3) - 1;
-            unpacked_bytes[4 + j] = (c0 & 0xFF) | ((c1 & 0xFF) << 8) | ((c2 & 0xFF) << 16) | ((c3 & 0xFF) << 24);
-        }
+        const int16_t    * qxi = (const int16_t *) bxi->qs + kqsx * 4;
 
         const int dst_offset = kbx*(scale_entries_per_block*QI8_0) + kqsx*QI8_0;
+
 #pragma unroll
-        for (int j = 0; j < 8; ++j) {
+        for (int j = 0; j < 4; ++j) {
+            const int q  = qxi[j];
+
+            // unpack even and odd crumbs into byte values
+            const int qe = __byte_perm(0x020100FF, 0x020100FF, q >> 0);
+            const int qo = __byte_perm(0x020100FF, 0x020100FF, q >> 2);
+            // unshuffle values
+            const int qx = __byte_perm(qe, qo, 0x5140);
+            const int qy = __byte_perm(qe, qo, 0x7362);
+
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-            x_qs[i*sram_stride           + dst_offset + j] = unpacked_bytes[j];
+            x_qs[i*sram_stride           + dst_offset + j*2+0] = qx;
+            x_qs[i*sram_stride           + dst_offset + j*2+1] = qy;
 #else
-            x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j] = unpacked_bytes[j];
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*2+0] = qx;
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*2+1] = qy;
 #endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
         }
     }

ggml/src/ggml-cuda/vecdotq.cuh

@@ -734,48 +734,28 @@ static __device__ __forceinline__ float vec_dot_q2_0_q8_1(
     // Q8_1: 32 elements per block with individual scales
     // iqs selects which of the 2 chunks of 32 elements to process (0-1)
 
-    const float d2 = bq2_0->d;
+    const float     d2 = bq2_0->d;
+    const int16_t * qs = (const int16_t *) bq2_0->qs + iqs * 4;
 
     // Process only the chunk specified by iqs
     const block_q8_1 * bq8_1_chunk = bq8_1 + iqs;
 
-    // Load 64 bits (8 bytes) for this chunk from Q2_0: bytes [8*iqs, 8*iqs+8)
-    const int offset = iqs * 8;
-    const int v0 = bq2_0->qs[offset + 0] | (bq2_0->qs[offset + 1] << 8) |
-                   (bq2_0->qs[offset + 2] << 16) | (bq2_0->qs[offset + 3] << 24);
-    const int v1 = bq2_0->qs[offset + 4] | (bq2_0->qs[offset + 5] << 8) |
-                   (bq2_0->qs[offset + 6] << 16) | (bq2_0->qs[offset + 7] << 24);
-
-    // Unpack 32 2-bit codes into 8 int32s, each holding 4 signed int8 symbols in {-1,0,1,2}.
-    // Stored code c in {0,1,2,3} -> symbol s = c - 1.
-    int vi_bytes[8];
-#pragma unroll
-    for (int j = 0; j < 4; ++j) {
-        const int shift = j * 8;
-        const int codes = (v0 >> shift) & 0xFF;
-        const int c0 = ((codes >> 0) & 0x3) - 1;
-        const int c1 = ((codes >> 2) & 0x3) - 1;
-        const int c2 = ((codes >> 4) & 0x3) - 1;
-        const int c3 = ((codes >> 6) & 0x3) - 1;
-        vi_bytes[j] = (c0 & 0xFF) | ((c1 & 0xFF) << 8) | ((c2 & 0xFF) << 16) | ((c3 & 0xFF) << 24);
-    }
-#pragma unroll
-    for (int j = 0; j < 4; ++j) {
-        const int shift = j * 8;
-        const int codes = (v1 >> shift) & 0xFF;
-        const int c0 = ((codes >> 0) & 0x3) - 1;
-        const int c1 = ((codes >> 2) & 0x3) - 1;
-        const int c2 = ((codes >> 4) & 0x3) - 1;
-        const int c3 = ((codes >> 6) & 0x3) - 1;
-        vi_bytes[4 + j] = (c0 & 0xFF) | ((c1 & 0xFF) << 8) | ((c2 & 0xFF) << 16) | ((c3 & 0xFF) << 24);
-    }
-
-    // Compute dot product for this 32-element chunk
     int sumi = 0;
 #pragma unroll
-    for (int j = 0; j < 8; ++j) {
-        const int u = get_int_b4(bq8_1_chunk->qs, j);
-        sumi = ggml_cuda_dp4a(vi_bytes[j], u, sumi);
+    for (int j = 0; j < 4; ++j) {
+        const int q  = qs[j];
+        const int u  = get_int_b4(bq8_1_chunk->qs, j*2+0);
+        const int v  = get_int_b4(bq8_1_chunk->qs, j*2+1);
+
+        // unpack even and odd crumbs into byte values
+        const int qe = __byte_perm(0x020100FF, 0x020100FF, q >> 0);
+        const int qo = __byte_perm(0x020100FF, 0x020100FF, q >> 2);
+        // unshuffle values
+        const int qx = __byte_perm(qe, qo, 0x5140);
+        const int qy = __byte_perm(qe, qo, 0x7362);
+
+        sumi = ggml_cuda_dp4a(u, qx, sumi);
+        sumi = ggml_cuda_dp4a(v, qy, sumi);
     }
 
     // Apply Q2_0's single scale and this chunk's Q8_1 scale

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

@@ -477,6 +477,24 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_soft_max(ggml_me
     return res;
 }
 
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc(ggml_metal_library_t lib, ggml_op op) {
+    const char * name = nullptr;
+
+    switch (op) {
+        case GGML_OP_DSV4_HC_COMB: name = "kernel_dsv4_hc_comb_f32"; break;
+        case GGML_OP_DSV4_HC_PRE:  name = "kernel_dsv4_hc_pre_f32";  break;
+        case GGML_OP_DSV4_HC_POST: name = "kernel_dsv4_hc_post_f32"; break;
+        default: GGML_ABORT("fatal error");
+    }
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr);
+    }
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv(ggml_metal_library_t lib, const ggml_tensor * op) {
     GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32);
     GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32);

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

@@ -124,6 +124,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_bl
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_add        (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_tri               (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_soft_max          (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc           (ggml_metal_library_t lib, enum ggml_op op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv          (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched  (ggml_metal_library_t lib, const struct ggml_tensor * op, int ssm_conv_bs);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan          (ggml_metal_library_t lib, const struct ggml_tensor * op);

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

@@ -1299,6 +1299,42 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
                     return false;
             }
             return has_simdgroup_mm; // TODO: over-restricted for vec-kernels
+        case GGML_OP_DSV4_HC_COMB:
+            return has_simdgroup_reduction &&
+                op->src[0]->type == GGML_TYPE_F32 &&
+                op->src[1]->type == GGML_TYPE_F32 &&
+                op->src[2]->type == GGML_TYPE_F32 &&
+                op->type         == GGML_TYPE_F32 &&
+                op->src[0]->ne[0] == 24 &&
+                op->src[1]->ne[0] >= 3 &&
+                op->src[2]->ne[0] == 24 &&
+                ggml_is_contiguous_rows(op->src[0]) &&
+                ggml_is_contiguous_rows(op->src[1]) &&
+                ggml_is_contiguous_rows(op->src[2]);
+        case GGML_OP_DSV4_HC_PRE:
+            return has_simdgroup_reduction &&
+                op->src[0]->type == GGML_TYPE_F32 &&
+                op->src[1]->type == GGML_TYPE_F32 &&
+                op->type         == GGML_TYPE_F32 &&
+                op->src[0]->ne[1] == 4 &&
+                op->src[1]->ne[0] == 4 &&
+                ggml_is_contiguous_rows(op->src[0]) &&
+                ggml_is_contiguous_rows(op->src[1]);
+        case GGML_OP_DSV4_HC_POST:
+            return has_simdgroup_reduction &&
+                op->src[0]->type == GGML_TYPE_F32 &&
+                op->src[1]->type == GGML_TYPE_F32 &&
+                op->src[2]->type == GGML_TYPE_F32 &&
+                op->src[3]->type == GGML_TYPE_F32 &&
+                op->type         == GGML_TYPE_F32 &&
+                op->src[1]->ne[1] == 4 &&
+                op->src[2]->ne[0] == 4 &&
+                op->src[3]->ne[0] == 4 &&
+                op->src[3]->ne[1] == 4 &&
+                ggml_is_contiguous_rows(op->src[0]) &&
+                ggml_is_contiguous_rows(op->src[1]) &&
+                ggml_is_contiguous_rows(op->src[2]) &&
+                ggml_is_contiguous_rows(op->src[3]);
         case GGML_OP_SSM_CONV:
         case GGML_OP_SSM_SCAN:
             return has_simdgroup_reduction;

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

@@ -1171,6 +1171,49 @@ typedef struct {
     int64_t val;
 } ggml_metal_kargs_memset;
 
+typedef struct {
+    int32_t  n_tokens;
+    int32_t  n_iter;
+    uint64_t nb_m0;
+    uint64_t nb_m1;
+    uint64_t nb_s0;
+    uint64_t nb_b0;
+    uint64_t nb_d0;
+    uint64_t nb_d1;
+    uint64_t nb_d2;
+    float    eps;
+} ggml_metal_kargs_dsv4_hc_comb;
+
+typedef struct {
+    int32_t  n_embd;
+    int32_t  n_tokens;
+    uint64_t nb_x0;
+    uint64_t nb_x1;
+    uint64_t nb_x2;
+    uint64_t nb_w0;
+    uint64_t nb_w1;
+    uint64_t nb_d0;
+    uint64_t nb_d1;
+} ggml_metal_kargs_dsv4_hc_pre;
+
+typedef struct {
+    int32_t  n_embd;
+    int32_t  n_tokens;
+    uint64_t nb_x0;
+    uint64_t nb_x1;
+    uint64_t nb_r0;
+    uint64_t nb_r1;
+    uint64_t nb_r2;
+    uint64_t nb_p0;
+    uint64_t nb_p1;
+    uint64_t nb_c0;
+    uint64_t nb_c1;
+    uint64_t nb_c2;
+    uint64_t nb_d0;
+    uint64_t nb_d1;
+    uint64_t nb_d2;
+} ggml_metal_kargs_dsv4_hc_post;
+
 typedef struct {
     int32_t  ne00;
     int32_t  ne01;

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

@@ -316,6 +316,12 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) {
             {
                 n_fuse = ggml_metal_op_cumsum(ctx, idx);
             } break;
+        case GGML_OP_DSV4_HC_COMB:
+        case GGML_OP_DSV4_HC_PRE:
+        case GGML_OP_DSV4_HC_POST:
+            {
+                n_fuse = ggml_metal_op_dsv4_hc(ctx, idx);
+            } break;
         case GGML_OP_SOFT_MAX:
             {
                 n_fuse = ggml_metal_op_soft_max(ctx, idx);
@@ -1297,6 +1303,137 @@ int ggml_metal_op_diag(ggml_metal_op_t ctx, int idx) {
     return 1;
 }
 
+int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) {
+    ggml_tensor * op = ctx->node(idx);
+
+    ggml_metal_encoder_t enc = ctx->enc;
+    auto pipeline = ggml_metal_library_get_pipeline_dsv4_hc(ctx->lib, op->op);
+
+    ggml_metal_encoder_set_pipeline(enc, pipeline);
+
+    switch (op->op) {
+        case GGML_OP_DSV4_HC_COMB:
+            {
+                const ggml_tensor * mixes = op->src[0];
+                const ggml_tensor * scale = op->src[1];
+                const ggml_tensor * base  = op->src[2];
+
+                GGML_ASSERT(mixes->type == GGML_TYPE_F32);
+                GGML_ASSERT(scale->type == GGML_TYPE_F32);
+                GGML_ASSERT(base->type  == GGML_TYPE_F32);
+                GGML_ASSERT(op->type    == GGML_TYPE_F32);
+                GGML_ASSERT(mixes->ne[0] == 24);
+                GGML_ASSERT(op->ne[0] == 4 && op->ne[1] == 4);
+
+                ggml_metal_kargs_dsv4_hc_comb args = {
+                    /*.n_tokens =*/ (int32_t) mixes->ne[1],
+                    /*.n_iter   =*/ ggml_get_op_params_i32(op, 1),
+                    /*.nb_m0    =*/ mixes->nb[0],
+                    /*.nb_m1    =*/ mixes->nb[1],
+                    /*.nb_s0    =*/ scale->nb[0],
+                    /*.nb_b0    =*/ base->nb[0],
+                    /*.nb_d0    =*/ op->nb[0],
+                    /*.nb_d1    =*/ op->nb[1],
+                    /*.nb_d2    =*/ op->nb[2],
+                    /*.eps      =*/ ggml_get_op_params_f32(op, 0),
+                };
+
+                ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
+                ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(mixes), 1);
+                ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(scale), 2);
+                ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(base),  3);
+                ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op),    4);
+
+                // One SIMDgroup owns one 4x4 Sinkhorn matrix. Packing up to four
+                // independent tokens per threadgroup keeps both decode and prompt
+                // dispatches compact without any threadgroup-memory synchronization.
+                const int nsg = std::min(4, args.n_tokens);
+                ggml_metal_encoder_dispatch_threadgroups(
+                        enc, (args.n_tokens + nsg - 1)/nsg, 1, 1, 32, nsg, 1);
+            } break;
+        case GGML_OP_DSV4_HC_PRE:
+            {
+                const ggml_tensor * x       = op->src[0];
+                const ggml_tensor * weights = op->src[1];
+
+                GGML_ASSERT(x->type       == GGML_TYPE_F32);
+                GGML_ASSERT(weights->type == GGML_TYPE_F32);
+                GGML_ASSERT(op->type      == GGML_TYPE_F32);
+                GGML_ASSERT(x->ne[1] == 4);
+
+                ggml_metal_kargs_dsv4_hc_pre args = {
+                    /*.n_embd   =*/ (int32_t) x->ne[0],
+                    /*.n_tokens =*/ (int32_t) x->ne[2],
+                    /*.nb_x0    =*/ x->nb[0],
+                    /*.nb_x1    =*/ x->nb[1],
+                    /*.nb_x2    =*/ x->nb[2],
+                    /*.nb_w0    =*/ weights->nb[0],
+                    /*.nb_w1    =*/ weights->nb[1],
+                    /*.nb_d0    =*/ op->nb[0],
+                    /*.nb_d1    =*/ op->nb[1],
+                };
+
+                ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
+                ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(x),       1);
+                ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(weights), 2);
+                ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op),      3);
+
+                const int n_tiles = (args.n_embd + 31)/32;
+                const int nsg = std::min(4, n_tiles);
+                ggml_metal_encoder_dispatch_threadgroups(
+                        enc, (n_tiles + nsg - 1)/nsg, args.n_tokens, 1, 32, nsg, 1);
+            } break;
+        case GGML_OP_DSV4_HC_POST:
+            {
+                const ggml_tensor * x        = op->src[0];
+                const ggml_tensor * residual = op->src[1];
+                const ggml_tensor * post     = op->src[2];
+                const ggml_tensor * comb     = op->src[3];
+
+                GGML_ASSERT(x->type        == GGML_TYPE_F32);
+                GGML_ASSERT(residual->type == GGML_TYPE_F32);
+                GGML_ASSERT(post->type     == GGML_TYPE_F32);
+                GGML_ASSERT(comb->type     == GGML_TYPE_F32);
+                GGML_ASSERT(op->type       == GGML_TYPE_F32);
+                GGML_ASSERT(residual->ne[1] == 4);
+
+                ggml_metal_kargs_dsv4_hc_post args = {
+                    /*.n_embd   =*/ (int32_t) x->ne[0],
+                    /*.n_tokens =*/ (int32_t) x->ne[1],
+                    /*.nb_x0    =*/ x->nb[0],
+                    /*.nb_x1    =*/ x->nb[1],
+                    /*.nb_r0    =*/ residual->nb[0],
+                    /*.nb_r1    =*/ residual->nb[1],
+                    /*.nb_r2    =*/ residual->nb[2],
+                    /*.nb_p0    =*/ post->nb[0],
+                    /*.nb_p1    =*/ post->nb[1],
+                    /*.nb_c0    =*/ comb->nb[0],
+                    /*.nb_c1    =*/ comb->nb[1],
+                    /*.nb_c2    =*/ comb->nb[2],
+                    /*.nb_d0    =*/ op->nb[0],
+                    /*.nb_d1    =*/ op->nb[1],
+                    /*.nb_d2    =*/ op->nb[2],
+                };
+
+                ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
+                ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(x),        1);
+                ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(residual), 2);
+                ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(post),     3);
+                ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(comb),     4);
+                ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op),       5);
+
+                const int n_tiles = (args.n_embd + 31)/32;
+                const int nsg = std::min(4, n_tiles);
+                ggml_metal_encoder_dispatch_threadgroups(
+                        enc, (n_tiles + nsg - 1)/nsg, args.n_tokens, 1, 32, nsg, 1);
+            } break;
+        default:
+            GGML_ABORT("fatal error");
+    }
+
+    return 1;
+}
+
 int ggml_metal_op_soft_max(ggml_metal_op_t ctx, int idx) {
     ggml_tensor * op = ctx->node(idx);
 

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

@@ -54,6 +54,7 @@ int ggml_metal_op_cumsum            (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_get_rows          (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_set_rows          (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_diag              (ggml_metal_op_t ctx, int idx);
+int ggml_metal_op_dsv4_hc           (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_soft_max          (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_ssm_conv          (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_ssm_scan          (ggml_metal_op_t ctx, int idx);

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

@@ -11278,3 +11278,162 @@ kernel void kernel_count_equal(
 typedef decltype(kernel_count_equal<int32_t>) kernel_count_equal_t;
 
 template [[host_name("kernel_count_equal_i32")]] kernel kernel_count_equal_t kernel_count_equal<int32_t>;
+
+kernel void kernel_dsv4_hc_comb_f32(
+        constant ggml_metal_kargs_dsv4_hc_comb & args,
+        device const char * mixes,
+        device const char * scale,
+        device const char * base,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort  tiisg[[thread_index_in_simdgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    constexpr ushort hc = 4;
+    constexpr ushort comb_offset = 2*hc;
+
+    const int it = tgpig.x*ntg.y + sgitg;
+    if (it >= args.n_tokens) {
+        return;
+    }
+
+    float scale_lane = 0.0f;
+    if (tiisg == 0) {
+        scale_lane = *(device const float *) (scale + 2*args.nb_s0);
+    }
+    const float scale_comb = simd_shuffle(scale_lane, 0);
+
+    float v = 0.0f;
+    if (tiisg < hc*hc) {
+        v = *(device const float *) (mixes + (comb_offset + tiisg)*args.nb_m0 + it*args.nb_m1)*scale_comb
+          + *(device const float *) (base   + (comb_offset + tiisg)*args.nb_b0);
+    }
+
+    // Softmax across destinations (the four contiguous lanes for each source).
+    float vmax = max(v, simd_shuffle_xor(v, 1));
+    vmax = max(vmax, simd_shuffle_xor(vmax, 2));
+    v = exp(v - vmax);
+
+    float sum = v + simd_shuffle_xor(v, 1);
+    sum += simd_shuffle_xor(sum, 2);
+    v = v/sum + args.eps;
+
+    // Normalize columns: equal destination indices are four lanes apart.
+    sum = v + simd_shuffle_xor(v, 4);
+    sum += simd_shuffle_xor(sum, 8);
+    v /= sum + args.eps;
+
+    for (int i = 1; i < args.n_iter; ++i) {
+        sum = v + simd_shuffle_xor(v, 1);
+        sum += simd_shuffle_xor(sum, 2);
+        v /= sum + args.eps;
+
+        sum = v + simd_shuffle_xor(v, 4);
+        sum += simd_shuffle_xor(sum, 8);
+        v /= sum + args.eps;
+    }
+
+    if (tiisg < hc*hc) {
+        const ushort idst = tiisg & 3;
+        const ushort isrc = tiisg >> 2;
+        *(device float *) (dst + idst*args.nb_d0 + isrc*args.nb_d1 + it*args.nb_d2) = v;
+    }
+}
+
+kernel void kernel_dsv4_hc_pre_f32(
+        constant ggml_metal_kargs_dsv4_hc_pre & args,
+        device const char * x,
+        device const char * weights,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort  tiisg[[thread_index_in_simdgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    constexpr ushort hc = 4;
+
+    const int it = tgpig.y;
+    const int i0 = ((int) tgpig.x*ntg.y + sgitg)*32 + tiisg;
+
+    float weight_lane = 0.0f;
+    if (tiisg < hc) {
+        weight_lane = *(device const float *) (weights + tiisg*args.nb_w0 + it*args.nb_w1);
+    }
+
+    float w[hc];
+    FOR_UNROLL (ushort ih = 0; ih < hc; ++ih) {
+        w[ih] = simd_shuffle(weight_lane, ih);
+    }
+
+    if (i0 >= args.n_embd) {
+        return;
+    }
+
+    device const char * xb = x + i0*args.nb_x0 + it*args.nb_x2;
+    float result = 0.0f;
+    FOR_UNROLL (ushort ih = 0; ih < hc; ++ih) {
+        result = fma(*(device const float *) (xb + ih*args.nb_x1), w[ih], result);
+    }
+
+    *(device float *) (dst + i0*args.nb_d0 + it*args.nb_d1) = result;
+}
+
+kernel void kernel_dsv4_hc_post_f32(
+        constant ggml_metal_kargs_dsv4_hc_post & args,
+        device const char * x,
+        device const char * residual,
+        device const char * post,
+        device const char * comb,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort  tiisg[[thread_index_in_simdgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    constexpr ushort hc = 4;
+
+    const int it = tgpig.y;
+    const int i0 = ((int) tgpig.x*ntg.y + sgitg)*32 + tiisg;
+
+    float coeff_lane = 0.0f;
+    if (tiisg < hc) {
+        coeff_lane = *(device const float *) (post + tiisg*args.nb_p0 + it*args.nb_p1);
+    } else if (tiisg < hc + hc*hc) {
+        const ushort idx  = tiisg - hc;
+        const ushort idst = idx & 3;
+        const ushort isrc = idx >> 2;
+        coeff_lane = *(device const float *) (comb + idst*args.nb_c0 + isrc*args.nb_c1 + it*args.nb_c2);
+    }
+
+    float post_reg[hc];
+    float comb_reg[hc][hc];
+    FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
+        post_reg[idst] = simd_shuffle(coeff_lane, idst);
+    }
+    FOR_UNROLL (ushort isrc = 0; isrc < hc; ++isrc) {
+        FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
+            comb_reg[isrc][idst] = simd_shuffle(coeff_lane, hc + idst + hc*isrc);
+        }
+    }
+
+    if (i0 >= args.n_embd) {
+        return;
+    }
+
+    const float xv = *(device const float *) (x + i0*args.nb_x0 + it*args.nb_x1);
+    float result[hc];
+    FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
+        result[idst] = xv*post_reg[idst];
+    }
+
+    device const char * rb = residual + i0*args.nb_r0 + it*args.nb_r2;
+    FOR_UNROLL (ushort isrc = 0; isrc < hc; ++isrc) {
+        const float rv = *(device const float *) (rb + isrc*args.nb_r1);
+        FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
+            result[idst] = fma(rv, comb_reg[isrc][idst], result[idst]);
+        }
+    }
+
+    FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
+        *(device float *) (dst + i0*args.nb_d0 + idst*args.nb_d1 + it*args.nb_d2) = result[idst];
+    }
+}
+

ggml/src/ggml-opencl/ggml-opencl.cpp

@@ -7495,6 +7495,7 @@ static ggml_backend_i ggml_backend_opencl_i = {
 ggml_backend_t ggml_backend_opencl_init(void) {
     ggml_backend_dev_t dev = ggml_backend_reg_dev_get(ggml_backend_opencl_reg(), 0);
     ggml_backend_opencl_context *backend_ctx = ggml_cl_init(dev);
+    backend_ctx->ref_count++;
 
     ggml_backend_t backend = new ggml_backend {
         /* .guid    = */ ggml_backend_opencl_guid(),

ggml/src/ggml-sycl/backend.hpp

@@ -26,6 +26,7 @@
 #include "dmmv.hpp"
 #include "element_wise.hpp"
 #include "fattn.hpp"
+#include "fusion.hpp"
 #include "gated_delta_net.hpp"
 #include "gla.hpp"
 #include "im2col.hpp"

ggml/src/ggml-sycl/common.hpp

@@ -133,6 +133,7 @@ enum ggml_sycl_backend_gpu_mode {
 enum ggml_sycl_dev2dev_memcpy_mode {
   DEV2DEV_MEMCPY_SYCL = 0,
   DEV2DEV_MEMCPY_L0 = 1,
+  DEV2DEV_MEMCPY_FORWARD = 2
 };
 
 static_assert(sizeof(sycl::half) == sizeof(ggml_fp16_t), "wrong fp16 size");
@@ -234,6 +235,7 @@ struct sycl_device_info {
     int max_wg_per_cu; // max work groups per compute unit - refer to
                        // cudaOccupancyMaxActiveBlocksPerMultiprocessor
     bool    vmm;                // virtual memory support
+    bool    l0_device_type_valid;
     bool    l0_discrete_gpu;    // Level Zero backend and not an integrated GPU
     size_t  vmm_granularity;    // granularity of virtual memory
     size_t  total_vram;

ggml/src/ggml-sycl/convert.cpp

@@ -644,6 +644,8 @@ to_fp16_sycl_t ggml_get_to_fp16_sycl(ggml_type type, ggml_tensor * dst) {
     switch (type) {
         case GGML_TYPE_Q1_0:
             return dequantize_block_sycl<QK1_0, QR1_0, dequantize_q1_0>;
+        case GGML_TYPE_Q2_0:
+            return dequantize_block_sycl<QK2_0, QR2_0, dequantize_q2_0>;
         case GGML_TYPE_Q4_0:
             if (dst->src[0]->extra &&
                 ((ggml_tensor_extra_gpu*)dst->src[0]->extra)->optimized_feature.reorder) {
@@ -728,6 +730,8 @@ to_fp32_sycl_t ggml_get_to_fp32_sycl(ggml_type type, ggml_tensor *dst) {
     switch (type) {
         case GGML_TYPE_Q1_0:
             return dequantize_block_sycl<QK1_0, QR1_0, dequantize_q1_0>;
+        case GGML_TYPE_Q2_0:
+            return dequantize_block_sycl<QK2_0, QR2_0, dequantize_q2_0>;
         case GGML_TYPE_Q4_0:
             if (dst->src[0]->extra &&
                 ((ggml_tensor_extra_gpu*)dst->src[0]->extra)->optimized_feature.reorder) {

ggml/src/ggml-sycl/cpy.cpp

@@ -8,7 +8,6 @@
 #include "ggml-sycl/presets.hpp"
 #include "ggml.h"
 
-
 static void cpy_1_f32_f32(const char * cxi, char * cdsti) {
     const float * xi   = (const float *) cxi;
     float *       dsti = (float *) cdsti;
@@ -151,6 +150,20 @@ static void cpy_blck_q8_0_f32(const char * cxi, char * cdsti) {
     }
 }
 
+static void cpy_blck_q2_0_f32(const char * cxi, char * cdsti) {
+    const block_q2_0 * xi = (const block_q2_0 *) cxi;
+    float * cdstf = (float *) cdsti;
+
+    const float d = xi->d;
+
+    for (int j = 0; j < QK2_0; ++j) {
+        const int byte_index = j / 4;
+        const int bit_offset = (j % 4) * 2;
+        const int q = (xi->qs[byte_index] >> bit_offset) & 0x3;
+        cdstf[j] = (float) (q - 1) * d;
+    }
+}
+
 
 
 template <dequantize_kernel_t dequant, int qk> static void cpy_blck_q_f32(const char * cxi, char * cdsti) {
@@ -256,7 +269,7 @@ static void ggml_cpy_f16_f32_sycl(const char * cx, char * cdst, const int ne, co
         stream->parallel_for(
             sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
                               sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
-            [=](sycl::nd_item<3> item_ct1) {
+            [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
                 cpy_f32_f16<cpy_1_f16_f32>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
                                            nb10, nb11, nb12, nb13, item_ct1);
             });
@@ -274,7 +287,7 @@ static void ggml_cpy_f32_f32_sycl(const char * cx, char * cdst, const int ne, co
         stream->parallel_for(
             sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
                               sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
-            [=](sycl::nd_item<3> item_ct1) {
+            [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
                 cpy_f32_f16<cpy_1_f32_f32>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
                                            nb10, nb11, nb12, nb13, item_ct1);
             });
@@ -292,7 +305,7 @@ static void ggml_cpy_f32_f16_sycl(const char * cx, char * cdst, const int ne, co
         stream->parallel_for(
             sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
                               sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
-            [=](sycl::nd_item<3> item_ct1) {
+            [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
                 cpy_f32_f16<cpy_1_f32_f16>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
                                            nb10, nb11, nb12, nb13, item_ct1);
             });
@@ -308,7 +321,7 @@ static void ggml_cpy_f32_i32_sycl(const char * cx, char * cdst, const int ne, co
         stream->parallel_for(
             sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
                               sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
-            [=](sycl::nd_item<3> item_ct1) {
+            [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
                 cpy_f32_f16<cpy_1_f32_i32>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
                                            nb10, nb11, nb12, nb13, item_ct1);
             });
@@ -324,7 +337,7 @@ static void ggml_cpy_i32_f32_sycl(const char * cx, char * cdst, const int ne, co
         stream->parallel_for(
             sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
                               sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
-            [=](sycl::nd_item<3> item_ct1) {
+            [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
                 cpy_f32_f16<cpy_1_i32_f32>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
                                            nb10, nb11, nb12, nb13, item_ct1);
             });
@@ -338,7 +351,7 @@ static void ggml_cpy_f32_q8_0_sycl(const char * cx, char * cdst, const int ne, c
     GGML_ASSERT(ne % QK8_0 == 0);
     const int num_blocks = ne / QK8_0;
     stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
-                         [=](sycl::nd_item<3> item_ct1) {
+                         [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
                              cpy_f32_q<cpy_blck_f32_q8_0, QK8_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
                                                                  ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
                          });
@@ -350,12 +363,25 @@ static void ggml_cpy_q8_0_f32_sycl(const char * cx, char * cdst, const int ne, c
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ne;
     stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
-                         [=](sycl::nd_item<3> item_ct1) {
+                         [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
                              cpy_q_f32<cpy_blck_q8_0_f32, QK8_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
                                                                  ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
                          });
 }
 
+static void ggml_cpy_q2_0_f32_sycl(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ne;
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+            cpy_q_f32<cpy_blck_q2_0_f32, QK2_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11,
+                                                ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
 static void ggml_cpy_f32_q4_0_sycl(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
                                    const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
                                    const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
@@ -363,7 +389,7 @@ static void ggml_cpy_f32_q4_0_sycl(const char * cx, char * cdst, const int ne, c
     GGML_ASSERT(ne % QK4_0 == 0);
     const int num_blocks = ne / QK4_0;
     stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
-                         [=](sycl::nd_item<3> item_ct1) {
+                         [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
                              cpy_f32_q<cpy_blck_f32_q4_0, QK4_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
                                                                  ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
                          });
@@ -375,7 +401,8 @@ static void ggml_cpy_q4_0_f32_sycl(const char * cx, char * cdst, const int ne, c
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ne;
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_f32<cpy_blck_q_f32<dequantize_q4_0, QK4_0>, QK4_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02,
                                                                      nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13,
                                                                      item_ct1);
@@ -389,7 +416,7 @@ static void ggml_cpy_f32_q4_1_sycl(const char * cx, char * cdst, const int ne, c
     GGML_ASSERT(ne % QK4_1 == 0);
     const int num_blocks = ne / QK4_1;
     stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
-                         [=](sycl::nd_item<3> item_ct1) {
+                         [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
                              cpy_f32_q<cpy_blck_f32_q4_1, QK4_1>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
                                                                  ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
                          });
@@ -401,7 +428,8 @@ static void ggml_cpy_q4_1_f32_sycl(const char * cx, char * cdst, const int ne, c
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ne;
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_f32<cpy_blck_q_f32<dequantize_q4_1, QK4_1>, QK4_1>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02,
                                                                      nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13,
                                                                      item_ct1);
@@ -415,7 +443,7 @@ static void ggml_cpy_f32_q5_0_sycl(const char * cx, char * cdst, const int ne, c
     GGML_ASSERT(ne % QK5_0 == 0);
     const int num_blocks = ne / QK5_0;
     stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
-                         [=](sycl::nd_item<3> item_ct1) {
+                         [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
                              cpy_f32_q<cpy_blck_f32_q5_0, QK5_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
                                                                  ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
                          });
@@ -427,7 +455,8 @@ static void ggml_cpy_q5_0_f32_sycl(const char * cx, char * cdst, const int ne, c
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ne;
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_f32<cpy_blck_q_f32<dequantize_q5_0, QK5_0>, QK5_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02,
                                                                      nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13,
                                                                      item_ct1);
@@ -441,7 +470,7 @@ static void ggml_cpy_f32_q5_1_sycl(const char * cx, char * cdst, const int ne, c
     GGML_ASSERT(ne % QK5_1 == 0);
     const int num_blocks = ne / QK5_1;
     stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
-                         [=](sycl::nd_item<3> item_ct1) {
+                         [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
                              cpy_f32_q<cpy_blck_f32_q5_1, QK5_1>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
                                                                  ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
                          });
@@ -453,7 +482,8 @@ static void ggml_cpy_q5_1_f32_sycl(const char * cx, char * cdst, const int ne, c
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ne;
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_f32<cpy_blck_q_f32<dequantize_q5_1, QK5_1>, QK5_1>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02,
                                                                      nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13,
                                                                      item_ct1);
@@ -466,7 +496,8 @@ static void ggml_cpy_mxfp4_f32_sycl(const char * cx, char * cdst, const int ne,
                                     const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ne;
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
             cpy_q_f32<cpy_blck_q_f32<dequantize_mxfp4, QK_MXFP4>, QK_MXFP4>(cx, cdst, ne, ne00, ne01, ne02, nb00,
                                                                              nb01, nb02, nb03, ne10, ne11, ne12,
                                                                              nb10, nb11, nb12, nb13, item_ct1);
@@ -480,7 +511,8 @@ static void ggml_cpy_f32_iq4_nl_sycl(const char * cx, char * cdst, const int ne,
     GGML_ASSERT(ne % QK4_NL == 0);
     const int num_blocks = ne / QK4_NL;
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
             cpy_f32_q<cpy_blck_f32_iq4_nl, QK4_NL>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11,
                                                    ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
@@ -526,7 +558,7 @@ static void ggml_cpy_f16_q4_0_sycl(const char * cx, char * cdst, const int ne, c
     GGML_ASSERT(ne % QK4_0 == 0);
     const int num_blocks = ne / QK4_0;
     stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
-                         [=](sycl::nd_item<3> item_ct1) {
+                         [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
                              cpy_f32_q<cpy_blck_f16_q4_0, QK4_0>(cx, cdst, ne, ne00, ne01, ne02,
                                                                  nb00, nb01, nb02, nb03,
                                                                  ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
@@ -540,7 +572,7 @@ static void ggml_cpy_f16_q4_1_sycl(const char * cx, char * cdst, const int ne, c
     GGML_ASSERT(ne % QK4_1 == 0);
     const int num_blocks = ne / QK4_1;
     stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
-                         [=](sycl::nd_item<3> item_ct1) {
+                         [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
                              cpy_f32_q<cpy_blck_f16_q4_1, QK4_1>(cx, cdst, ne, ne00, ne01, ne02,
                                                                  nb00, nb01, nb02, nb03,
                                                                  ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
@@ -554,7 +586,7 @@ static void ggml_cpy_f16_q5_0_sycl(const char * cx, char * cdst, const int ne, c
     GGML_ASSERT(ne % QK5_0 == 0);
     const int num_blocks = ne / QK5_0;
     stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
-                         [=](sycl::nd_item<3> item_ct1) {
+                         [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
                              cpy_f32_q<cpy_blck_f16_q5_0, QK5_0>(cx, cdst, ne, ne00, ne01, ne02,
                                                                  nb00, nb01, nb02, nb03,
                                                                  ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
@@ -564,6 +596,7 @@ static void ggml_cpy_f16_q5_0_sycl(const char * cx, char * cdst, const int ne, c
 static bool ggml_sycl_is_quantized_type(enum ggml_type type) {
     switch (type) {
         case GGML_TYPE_Q1_0:
+        case GGML_TYPE_Q2_0:
         case GGML_TYPE_Q4_0:
         case GGML_TYPE_Q4_1:
         case GGML_TYPE_Q5_0:
@@ -594,6 +627,7 @@ static bool ggml_sycl_is_quantized_type(enum ggml_type type) {
 static bool ggml_sycl_can_quantize_rows_sycl(enum ggml_type type) {
     switch (type) {
         case GGML_TYPE_Q1_0:
+        case GGML_TYPE_Q2_0:
         case GGML_TYPE_Q4_0:
         case GGML_TYPE_Q4_1:
         case GGML_TYPE_Q5_0:
@@ -651,7 +685,8 @@ static void ggml_sycl_quantize_rows_q(const char * cx, char * cdst, const int64_
     constexpr int block_size = 256;
     const int64_t grid_size = ceil_div(total_blocks, (int64_t) block_size);
 
-    stream->parallel_for(sycl::nd_range<1>(grid_size * block_size, block_size), [=](sycl::nd_item<1> item_ct1) {
+    stream->parallel_for(sycl::nd_range<1>(grid_size * block_size, block_size),
+                        [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
         const int64_t block_idx = item_ct1.get_global_linear_id();
         if (block_idx >= total_blocks) {
             return;
@@ -708,6 +743,11 @@ static void ggml_sycl_quantize_rows_sycl(const char * cx, char * cdst, const ggm
                                                                             nb02, nb03, ne10, ne11, ne12, nb10, nb11,
                                                                             nb12, nb13, stream);
             break;
+        case GGML_TYPE_Q2_0:
+            ggml_sycl_quantize_rows_q<SrcScalar, cpy_blck_f32_q2_0, QK2_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01,
+                                                                            nb02, nb03, ne10, ne11, ne12, nb10, nb11,
+                                                                            nb12, nb13, stream);
+            break;
         case GGML_TYPE_Q5_1:
             ggml_sycl_quantize_rows_q<SrcScalar, cpy_blck_f32_q5_1, QK5_1>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01,
                                                                             nb02, nb03, ne10, ne11, ne12, nb10, nb11,
@@ -760,7 +800,7 @@ static void ggml_cpy_f16_f16_sycl(const char * cx, char * cdst, const int ne, co
         stream->parallel_for(
             sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
                               sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
-            [=](sycl::nd_item<3> item_ct1) {
+            [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
                 cpy_f32_f16<cpy_1_f16_f16>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
                                            nb10, nb11, nb12, nb13, item_ct1);
             });
@@ -779,7 +819,7 @@ static void ggml_cpy_i16_i16_sycl(const char * cx, char * cdst, const int ne, co
         stream->parallel_for(
             sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
                               sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
-            [=](sycl::nd_item<3> item_ct1) {
+            [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
                 cpy_f32_f16<cpy_1_i16_i16>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
                                            nb10, nb11, nb12, nb13, item_ct1);
             });
@@ -798,7 +838,7 @@ static void ggml_cpy_i32_i32_sycl(const char * cx, char * cdst, const int ne, co
         stream->parallel_for(
             sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
                               sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
-            [=](sycl::nd_item<3> item_ct1) {
+            [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
                 cpy_f32_f16<cpy_1_i32_i32>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
                                            nb10, nb11, nb12, nb13, item_ct1);
             });
@@ -812,7 +852,8 @@ static void ggml_cpy_q8_0_q8_0(const char * cx, char * cdst, const int ne, const
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
         sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
-                              sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+                              sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+                              [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_q8_0, QK8_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -825,7 +866,8 @@ static void ggml_cpy_q5_0_q5_0(const char * cx, char * cdst, const int ne, const
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
         sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
-                              sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+                              sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+                              [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_q5_0, QK5_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -839,7 +881,8 @@ static void ggml_cpy_q5_1_q5_1(const char * cx, char * cdst, const int ne, const
 
     stream->parallel_for(
         sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
-                              sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+                              sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+                              [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_q5_1, QK5_1>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -851,7 +894,8 @@ static void ggml_cpy_q4_0_q4_0(const char * cx, char * cdst, const int ne, const
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_q4_0, QK4_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -864,7 +908,8 @@ static void ggml_cpy_q4_1_q4_1(const char * cx, char * cdst, const int ne, const
 
    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
    stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_q4_1, QK4_1>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -875,18 +920,32 @@ static void ggml_cpy_q1_0_q1_0(const char * cx, char * cdst, const int ne, const
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
             cpy_q_q<block_q1_0, QK1_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
 
+static void ggml_cpy_q2_0_q2_0(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
+            cpy_q_q<block_q2_0, QK2_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
 static void ggml_cpy_mxfp4_mxfp4(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
                                    const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
                                    const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
             cpy_q_q<block_mxfp4, QK_MXFP4>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -897,7 +956,8 @@ static void ggml_cpy_nvfp4_nvfp4(const char * cx, char * cdst, const int ne, con
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_nvfp4, QK_NVFP4>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -908,7 +968,8 @@ static void ggml_cpy_q2_K_q2_K(const char * cx, char * cdst, const int ne, const
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_q2_K, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -919,7 +980,8 @@ static void ggml_cpy_q3_K_q3_K(const char * cx, char * cdst, const int ne, const
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_q3_K, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -930,7 +992,8 @@ static void ggml_cpy_q4_K_q4_K(const char * cx, char * cdst, const int ne, const
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_q4_K, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -941,7 +1004,8 @@ static void ggml_cpy_q5_K_q5_K(const char * cx, char * cdst, const int ne, const
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_q5_K, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -952,7 +1016,8 @@ static void ggml_cpy_q6_K_q6_K(const char * cx, char * cdst, const int ne, const
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_q6_K, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -963,7 +1028,8 @@ static void ggml_cpy_iq2_xxs_iq2_xxs(const char * cx, char * cdst, const int ne,
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_iq2_xxs, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -974,7 +1040,8 @@ static void ggml_cpy_iq2_xs_iq2_xs(const char * cx, char * cdst, const int ne, c
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_iq2_xs, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -985,7 +1052,8 @@ static void ggml_cpy_iq2_s_iq2_s(const char * cx, char * cdst, const int ne, con
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_iq2_s, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -996,7 +1064,8 @@ static void ggml_cpy_iq3_xxs_iq3_xxs(const char * cx, char * cdst, const int ne,
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_iq3_xxs, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -1007,7 +1076,8 @@ static void ggml_cpy_iq1_s_iq1_s(const char * cx, char * cdst, const int ne, con
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_iq1_s, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -1018,7 +1088,8 @@ static void ggml_cpy_iq1_m_iq1_m(const char * cx, char * cdst, const int ne, con
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_iq1_m, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -1029,7 +1100,8 @@ static void ggml_cpy_iq4_nl_iq4_nl(const char * cx, char * cdst, const int ne, c
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_iq4_nl, QK4_NL>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -1040,7 +1112,8 @@ static void ggml_cpy_iq3_s_iq3_s(const char * cx, char * cdst, const int ne, con
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_iq3_s, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -1051,7 +1124,8 @@ static void ggml_cpy_iq4_xs_iq4_xs(const char * cx, char * cdst, const int ne, c
                                    const int nb12, const int nb13, queue_ptr stream) {
     const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_q_q<block_iq4_xs, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
         });
 }
@@ -1065,7 +1139,7 @@ static void ggml_cpy_f32_bf16_sycl(const char * cx, char * cdst, const int ne, c
     stream->parallel_for(
         sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
                           sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
-        [=](sycl::nd_item<3> item_ct1) {
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_f32_f16<cpy_1_f32_bf16>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
                                            nb10, nb11, nb12, nb13, item_ct1);
         });
@@ -1079,7 +1153,7 @@ static void ggml_cpy_bf16_f32_sycl(const char * cx, char * cdst, const int ne, c
     stream->parallel_for(
         sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
                           sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
-        [=](sycl::nd_item<3> item_ct1) {
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_f32_f16<cpy_1_bf16_f32>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
                                            nb10, nb11, nb12, nb13, item_ct1);
         });
@@ -1093,7 +1167,7 @@ static void ggml_cpy_bf16_bf16_sycl(const char * cx, char * cdst, const int ne,
     stream->parallel_for(
         sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
                           sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
-        [=](sycl::nd_item<3> item_ct1) {
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_f32_f16<cpy_1_bf16_bf16>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
                                             nb10, nb11, nb12, nb13, item_ct1);
         });
@@ -1107,7 +1181,7 @@ static void ggml_cpy_f16_bf16_sycl(const char * cx, char * cdst, const int ne, c
     stream->parallel_for(
         sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
                           sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
-        [=](sycl::nd_item<3> item_ct1) {
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_f32_f16<cpy_1_f16_bf16>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
                                            nb10, nb11, nb12, nb13, item_ct1);
         });
@@ -1121,7 +1195,7 @@ static void ggml_cpy_bf16_f16_sycl(const char * cx, char * cdst, const int ne, c
     stream->parallel_for(
         sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
                           sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
-        [=](sycl::nd_item<3> item_ct1) {
+        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{
             cpy_f32_f16<cpy_1_bf16_f16>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
                                            nb10, nb11, nb12, nb13, item_ct1);
         });
@@ -1213,6 +1287,9 @@ void ggml_sycl_cpy(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, co
     } else if (src0->type == GGML_TYPE_Q8_0 && src1->type == GGML_TYPE_F32) {
         ggml_cpy_q8_0_f32_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10,
                                nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_Q2_0 && src1->type == GGML_TYPE_F32) {
+        ggml_cpy_q2_0_f32_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
+                               nb10, nb11, nb12, nb13, main_stream);
     } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_Q5_0) {
         ggml_cpy_f32_q5_0_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10,
                                nb11, nb12, nb13, main_stream);
@@ -1243,6 +1320,8 @@ void ggml_sycl_cpy(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, co
         ggml_cpy_q4_1_q4_1(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
     } else if (src0->type == GGML_TYPE_Q1_0 && src1->type == GGML_TYPE_Q1_0) {
         ggml_cpy_q1_0_q1_0(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_Q2_0 && src1->type == GGML_TYPE_Q2_0) {
+        ggml_cpy_q2_0_q2_0(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
     } else if (src0->type == GGML_TYPE_MXFP4 && src1->type == GGML_TYPE_MXFP4) {
         ggml_cpy_mxfp4_mxfp4(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
     } else if (src0->type == GGML_TYPE_NVFP4 && src1->type == GGML_TYPE_NVFP4) {

ggml/src/ggml-sycl/cpy.hpp

@@ -70,6 +70,39 @@ inline void cpy_blck_f32_q1_0(const char * cxi, char * cdsti) {
     }
 }
 
+inline int round_nearest_int(float x) {
+    return (int)(x >= 0.0f ? x + 0.5f : x - 0.5f);
+}
+
+inline void cpy_blck_f32_q2_0(const char * cxi, char * cdsti) {
+    const float * xi   = (const float *) cxi;
+    block_q2_0 *  dsti = (block_q2_0 *) cdsti;
+
+    float amax = 0.0f;
+
+    for (int j = 0; j < QK2_0; ++j) {
+        amax = sycl::fmax(amax, sycl::fabs((float) xi[j]));
+    }
+
+    const float d  = amax;
+    const float id = d > 0.0f ? 1.0f / d : 0.0f;
+
+    dsti->d = d;
+
+    for (int j = 0; j < QK2_0 / 4; ++j) {
+        dsti->qs[j] = 0;
+    }
+
+    for (int j = 0; j < QK2_0; ++j) {
+        int q = round_nearest_int(xi[j] * id) + 1;
+        q = dpct::max(0, dpct::min(3, q));
+
+        const int byte_index = j / 4;
+        const int bit_offset = (j % 4) * 2;
+        dsti->qs[byte_index] |= (uint8_t) q << bit_offset;
+    }
+}
+
 inline int best_index_mxfp4(const float x, const float e) {
     int best_index = 0;
     float best_err = sycl::fabs((float) (kvalues_mxfp4[0] * e - x));

ggml/src/ggml-sycl/dequantize.hpp

@@ -25,6 +25,28 @@ typedef void (*dequantize_kernel_f32_t)(const void * vx, const int64_t ib, const
 static inline void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m);
 #endif
 
+static __dpct_inline__ void dequantize_q2_0(const void *vx, const int64_t ib,
+                                            const int iqs, dfloat2 &v) {
+    const block_q2_0 * x = (const block_q2_0 *) vx;
+
+    const dfloat d = x[ib].d;
+
+    const int byte_idx = iqs / 4;
+    const int shift = (iqs % 4) * 2;
+    const uint8_t vui = x[ib].qs[byte_idx];
+
+    v.x() = (vui >> shift) & 3;
+    v.y() = (vui >> (shift + 2)) & 3;
+
+#ifdef GGML_SYCL_F16
+    v.s0() = ((dfloat)v.s0() - 1.0f) * d;
+    v.s1() = ((dfloat)v.s1() - 1.0f) * d;
+#else
+    v.x() = ((dfloat)v.x() - 1.0f) * d;
+    v.y() = ((dfloat)v.y() - 1.0f) * d;
+#endif // GGML_SYCL_F16
+}
+
 static __dpct_inline__ void dequantize_q4_0(const void *vx, const int64_t ib,
                                             const int iqs, dfloat2 &v) {
     const block_q4_0 * x = (const block_q4_0 *) vx;

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

@@ -0,0 +1,690 @@
+// Flash attention via oneMKL GEMM (XMX-accelerated).
+// Uses column_major::gemm for Q*K^T and S*V matmuls
+// with an online softmax SYCL kernel.
+//
+// All GQA query heads sharing a KV head are batched into single
+// GEMM calls, amortizing MKL launch overhead across K and V reuse.
+//
+
+#include "common.hpp"
+#include "fattn-common.hpp"
+#include "fattn-buffers.hpp"
+#include "convert.hpp"
+#include "fattn.hpp"
+
+#include <oneapi/mkl.hpp>
+#include <cstdio>
+#include <chrono>
+
+#define MKL_FA_CHUNK_SIZE_KV 8192
+
+// Number of query rows processed per tile. The score buffers (KQ_f32, S_f16)
+// are sized q_tile_rows * chunk_size, so this bounds their footprint
+// regardless of batch size (n_query_rows = n_queries * gqa_ratio). A typical
+// single-ubatch prefill (e.g. ubatch 1024 * gqa 8 = 8192 rows) is exactly one
+// tile, so it runs with no extra iterations. Larger batches tile and stay
+// bounded. Override with GGML_SYCL_MKL_FA_Q_TILE.
+#define MKL_FA_Q_TILE 8192
+
+#define MKL_FA_WG_SIZE 256
+
+using oneapi::mkl::transpose;
+using oneapi::mkl::blas::column_major::gemm;
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+// Pack all GQA Q heads for one KV head into fp16, applying q_scale.
+// Launches one kernel per GQA group — each kernel copies exactly
+// n_queries * DKQ elements using the per-group dst offset and
+// per-head source stride.
+static void mkl_fa_pack_q_fp16(
+    dpct::queue_ptr stream,
+    sycl::half * __restrict dst,
+    const float * __restrict q_src,
+    int n_queries, int n_query_rows, int DKQ,
+    int gqa_ratio, int kvh_base_head,
+    float q_scale, int64_t q_row_stride, int64_t q_head_stride,
+    int64_t wg_size) {
+
+    for (int iqg = 0; iqg < gqa_ratio; iqg++) {
+        int     iqh       = kvh_base_head + iqg;
+        sycl::half * dst_g = dst + (int64_t)iqg * n_queries * DKQ;
+
+        const int64_t n_elem = (int64_t)n_queries * DKQ;
+        const int64_t wg = ((n_elem + wg_size - 1) / wg_size) * wg_size;
+
+        stream->submit([&](sycl::handler & cgh) {
+            cgh.parallel_for(sycl::nd_range<1>(wg, wg_size),
+                [=](sycl::nd_item<1> item) {
+                    int64_t e = item.get_global_id(0);
+                    if (e >= n_elem) return;
+
+                    int64_t q = e / DKQ;
+                    int64_t d = e - q * DKQ;
+
+                    // Stride-aware source offset: handles permuted,
+                    // sliced, or contiguous Q tensor layouts.
+                    int64_t src_off = d
+                        + q * q_row_stride
+                        + (int64_t)iqh * q_head_stride;
+
+                    dst_g[e] = sycl::half(
+                        q_src[src_off] * q_scale);
+                });
+        });
+    }
+}
+
+// Zero-initialize the online softmax state arrays.
+// KQ_max → -inf, KQ_sum → 0, VKQ_accum → 0.
+// Merged into one kernel to avoid per-array launch overhead.
+static void mkl_fa_init_softmax_state(
+    dpct::queue_ptr stream,
+    float * kmax, float * ksum, float * vacc,
+    int n_query_rows, int DV, int64_t wg_size) {
+
+    const float    neg_inf   = -1e30f;
+    const int64_t  n_maxsum  = n_query_rows;
+    const int64_t  n_vacc    = (int64_t)n_query_rows * DV;
+    const int64_t  total     = (n_vacc > n_maxsum) ? n_vacc : n_maxsum;
+    const int64_t  wg = ((total + wg_size - 1) / wg_size) * wg_size;
+
+    stream->submit([&](sycl::handler & cgh) {
+        cgh.parallel_for(sycl::nd_range<1>(wg, wg_size),
+            [=](sycl::nd_item<1> item) {
+                int64_t i = item.get_global_id(0);
+                if (i < n_maxsum) {
+                    kmax[i] = neg_inf;
+                    ksum[i] = 0.0f;
+                }
+                if (i < n_vacc) {
+                    vacc[i] = 0.0f;
+                }
+            });
+    });
+}
+
+// Online softmax over one KV chunk for a tile of GQA query rows.
+// The tile spans absolute rows [q0, q0 + q_rows). Score buffers
+// (KQ_f32/S_f16) are indexed RELATIVE to the tile; the persistent state
+// (VKQ_accum/KQ_max/KQ_sum) and mask are indexed by ABSOLUTE row.
+// For each row: find local max → rescale previous VKQ_accum →
+// compute exp(s - max) → write S_f16 → update running max/sum.
+static void mkl_fa_online_softmax_chunk(
+    dpct::queue_ptr stream,
+    float * __restrict KQ_f32,
+    sycl::half * __restrict S_f16,
+    float * __restrict KQ_max,
+    float * __restrict KQ_sum,
+    float * __restrict VKQ_accum,
+    int q0, int q_rows, int n_queries, int DV,
+    int chunk_size, int chunk_start,
+    int kvh_head, int gqa_ratio,
+    const sycl::half * mask_data, int64_t mask_head_stride,
+    int64_t mask_row_stride, int mask_n_heads,
+    float logit_softcap, int64_t wg_size) {
+
+    const int64_t wg = ((q_rows + wg_size - 1) / wg_size) * wg_size;
+
+    stream->submit([&](sycl::handler & cgh) {
+        cgh.parallel_for(sycl::nd_range<1>(wg, wg_size),
+            [=](sycl::nd_item<1> item) {
+                int jc_rel = item.get_global_id(0);
+                if (jc_rel >= q_rows) return;
+                int jc_abs = q0 + jc_rel;
+
+                const int gqa_group = jc_abs / n_queries;
+                const int q_row     = jc_abs % n_queries;
+
+                // Score buffers are tile-local (relative index).
+                const float * __restrict KQ_row = KQ_f32
+                    + jc_rel * (int64_t)chunk_size;
+                // Persistent accumulator is full-sized (absolute index).
+                float * __restrict vkq = VKQ_accum
+                    + jc_abs * (int64_t)DV;
+
+                const sycl::half * mask_h = nullptr;
+                int64_t m_stride = 0;
+                if (mask_data) {
+                    int m_head = (mask_n_heads > 1)
+                        ? (kvh_head + gqa_group) : 0;
+                    mask_h   = mask_data + (int64_t)m_head * mask_head_stride;
+                    m_stride = mask_row_stride;
+                }
+
+                // Row-wise local maximum (softcap before mask)
+                float local_max = -1e30f;
+                for (int i = 0; i < chunk_size; i++) {
+                    float s = KQ_row[i];
+                    if (logit_softcap != 0.0f) {
+                        s = logit_softcap * sycl::tanh(s);
+                    }
+                    if (mask_h) {
+                        s += (float)mask_h[q_row * m_stride
+                            + (chunk_start + i)];
+                    }
+                    if (s > local_max) local_max = s;
+                }
+
+                // Rescale previous accumulator by exp(old_max - new_max)
+                float old_max = KQ_max[jc_abs];
+                float new_max = (old_max > local_max) ? old_max : local_max;
+                float rescale = (old_max < -1e29f) ? 1.0f
+                    : sycl::native::exp(old_max - new_max);
+
+                for (int v = 0; v < DV; v++) {
+                    vkq[v] *= rescale;
+                }
+
+                // Softmax and write S_f16 (tile-local index)
+                float local_sum = 0.0f;
+                sycl::half * __restrict S_row = S_f16
+                    + jc_rel * (int64_t)chunk_size;
+
+                for (int i = 0; i < chunk_size; i++) {
+                    float s = KQ_row[i];
+                    if (logit_softcap != 0.0f) {
+                        s = logit_softcap * sycl::tanh(s);
+                    }
+                    if (mask_h) {
+                        s += (float)mask_h[q_row * m_stride
+                            + (chunk_start + i)];
+                    }
+                    float val = sycl::native::exp(s - new_max);
+                    S_row[i] = sycl::half(val);
+                    local_sum += val;
+                }
+
+                KQ_sum[jc_abs] = KQ_sum[jc_abs] * rescale + local_sum;
+                KQ_max[jc_abs] = new_max;
+            });
+    });
+}
+
+// Write one GQA group's normalized output to its destination head.
+static void mkl_fa_normalize_head(
+    dpct::queue_ptr stream,
+    float * __restrict dst_batch,
+    const float * __restrict VKQ_accum,
+    const float * __restrict KQ_sum,
+    int iqh, int n_queries, int DV, int n_q_heads,
+    int64_t src_offset, int64_t wg_size) {
+
+    const int64_t wg = ((n_queries + wg_size - 1) / wg_size) * wg_size;
+
+    stream->submit([&](sycl::handler & cgh) {
+        cgh.parallel_for(sycl::nd_range<1>(wg, wg_size),
+            [=](sycl::nd_item<1> item) {
+                int jc = item.get_global_id(0);
+                if (jc >= n_queries) return;
+
+                int     ksum_idx = (int)(src_offset / DV) + jc;
+                float   inv_sum  = 1.0f / KQ_sum[ksum_idx];
+                const float * __restrict src = VKQ_accum
+                    + src_offset + jc * (int64_t)DV;
+                // Interleaved dst layout (matching TILE):
+                // rows alternate between heads, then increment query.
+                // offset = (query * n_q_heads + head) * DV
+                float * __restrict dst_row = dst_batch
+                    + ((int64_t)jc * n_q_heads + iqh) * (int64_t)DV;
+
+                for (int v = 0; v < DV; v++) {
+                    dst_row[v] = src[v] * inv_sum;
+                }
+            });
+    });
+}
+
+// ---------------------------------------------------------------------------
+// Per-chunk dequant
+//
+// Rather than dequantizing all of K/V up front (footprint scales with
+// context), we dequant one KV-head chunk at a time into a dense
+// [this_chunk x D] fp16 buffer (row-major, lda = D). The source address of
+// element (head=ikvh, row=chunk_start+r, col=c) decomposes into independent
+// linear terms head_off(ikvh) + row_off(chunk_start) + (r,c), so slicing a
+// chunk is a clean pointer offset in every layout case. The true-Gemma-
+// interleave vs padded-seq-view distinction is resolved once when the
+// descriptor is built; slicing does not reintroduce it.
+// ---------------------------------------------------------------------------
+enum mkl_fa_kv_desc_mode {
+    MKL_FA_KV_MODE_F16_DENSE       = 0,
+    MKL_FA_KV_MODE_F16_INTERLEAVED = 1,
+    MKL_FA_KV_MODE_QUANT_CONTIG    = 2,
+    MKL_FA_KV_MODE_QUANT_NC        = 3,
+};
+
+struct mkl_fa_kv_desc {
+    const char *         data = nullptr;
+    ggml_type            type = GGML_TYPE_F16;
+    int64_t              D    = 0;      // ne[0]
+    int64_t              nb1  = 0;      // byte stride, seq dim
+    int64_t              nb2  = 0;      // byte stride, head dim
+    mkl_fa_kv_desc_mode  mode = MKL_FA_KV_MODE_F16_DENSE;
+    int64_t              ts   = 0;      // type size (mode 3 base offset)
+    int64_t              s01  = 0;      // nc row stride in blocks (mode 3)
+    int64_t              s02  = 0;      // nc head stride in blocks (mode 3)
+};
+
+static mkl_fa_kv_desc mkl_fa_make_desc(const ggml_tensor * T, bool interleaved, int n_kv_heads) {
+    mkl_fa_kv_desc d;
+    d.data = (const char *)T->data;
+    d.type = T->type;
+    d.D    = T->ne[0];
+    d.nb1  = (int64_t)T->nb[1];
+    d.nb2  = (int64_t)T->nb[2];
+    d.ts   = (int64_t)ggml_type_size(T->type);
+
+    if (T->type == GGML_TYPE_F16) {
+        d.mode = interleaved ? MKL_FA_KV_MODE_F16_INTERLEAVED
+                             : MKL_FA_KV_MODE_F16_DENSE;
+    } else if (ggml_is_contiguously_allocated(T) && !interleaved) {
+        d.mode = MKL_FA_KV_MODE_QUANT_CONTIG;
+    } else {
+        d.mode = MKL_FA_KV_MODE_QUANT_NC;
+        const int64_t bs          = (int64_t)ggml_blck_size(T->type);
+        const int64_t blk_per_row = T->ne[0] / bs;
+        // True Gemma interleave packs heads within a row (nb[2] < ne[1]*nb[1])
+        // → reconstruct physical strides. Padded seq-views (nb[2] > ne[1]*nb[1])
+        // already have correct physical strides.
+        const bool gemma = interleaved &&
+            ((int64_t)T->nb[2] < (int64_t)T->ne[1] * (int64_t)T->nb[1]);
+        if (gemma) {
+            d.s01 = (int64_t)n_kv_heads * blk_per_row;
+            d.s02 = blk_per_row;
+        } else {
+            d.s01 = d.nb1 / d.ts;
+            d.s02 = d.nb2 / d.ts;
+        }
+    }
+    return d;
+}
+
+// Dequant one KV-head chunk into a dense [this_chunk x D] fp16 buffer.
+static void mkl_fa_dequant_chunk(
+    dpct::queue_ptr stream, const mkl_fa_kv_desc & d, ggml_tensor * dst_ctx,
+    sycl::half * out, int ikvh, int chunk_start, int this_chunk) {
+
+    const int64_t D = d.D;
+    switch (d.mode) {
+        case MKL_FA_KV_MODE_F16_DENSE: {
+            const char * base = d.data + (int64_t)ikvh * d.nb2
+                + (int64_t)chunk_start * d.nb1;
+            stream->memcpy(out, base, (size_t)this_chunk * D * sizeof(sycl::half));
+            break;
+        }
+        case MKL_FA_KV_MODE_F16_INTERLEAVED: {
+            const char * base = d.data + (int64_t)ikvh * d.nb2
+                + (int64_t)chunk_start * d.nb1;
+            const int64_t row_halfs = d.nb1 / (int64_t)sizeof(sycl::half);
+            const sycl::half * src = (const sycl::half *)base;
+            stream->parallel_for(
+                sycl::range<2>((size_t)this_chunk, (size_t)D),
+                [=](sycl::item<2> it) {
+                    int64_t r = it.get_id(0);
+                    int64_t c = it.get_id(1);
+                    out[r * D + c] = src[r * row_halfs + c];
+                });
+            break;
+        }
+        case MKL_FA_KV_MODE_QUANT_CONTIG: {
+            const char * base = d.data + (int64_t)ikvh * d.nb2
+                + (int64_t)chunk_start * d.nb1;
+            to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(d.type, dst_ctx);
+            to_fp16(base, out, (int64_t)this_chunk * D, stream);
+            break;
+        }
+        default: {  // MKL_FA_KV_MODE_QUANT_NC
+            to_fp16_nc_sycl_t to_fp16 = ggml_get_to_fp16_nc_sycl(d.type);
+            const int64_t base_blocks = (int64_t)ikvh * d.s02
+                + (int64_t)chunk_start * d.s01;
+            const char * base = d.data + base_blocks * d.ts;
+            // ne02 = ne03 = 1 → s02/s03 inert; head+chunk offset carried by base.
+            to_fp16(base, out, D, this_chunk, 1, 1, d.s01, d.s02, d.s02, stream);
+            break;
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// MKL Flash Attention orchestrator
+//
+// Pipeline: dequantize K/V → for each KV head:
+//   pack GQA Q heads → MKL GEMM KQ → online softmax →
+//   MKL GEMM VKQ → accumulate → normalize → scatter to dst
+// ---------------------------------------------------------------------------
+void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
+
+    const ggml_tensor * Q    = dst->src[0];
+    const ggml_tensor * K    = dst->src[1];
+    const ggml_tensor * V    = dst->src[2];
+    const ggml_tensor * mask = dst->src[3];
+    ggml_tensor * KQV = dst;
+
+    GGML_ASSERT(Q->type == GGML_TYPE_F32);
+    GGML_ASSERT(KQV->type == GGML_TYPE_F32);
+
+    // --- Op params ---
+    float scale = 1.0f, max_bias = 0.0f, logit_softcap = 0.0f;
+    memcpy(&scale,         (const float *)KQV->op_params + 0, sizeof(float));
+    memcpy(&max_bias,      (const float *)KQV->op_params + 1, sizeof(float));
+    memcpy(&logit_softcap, (const float *)KQV->op_params + 2, sizeof(float));
+
+    const float q_scale = scale;
+
+    // --- Dimensions ---
+    const int DKQ        = (int)K->ne[0];
+    const int DV         = (int)V->ne[0];
+    const int n_queries  = (int)Q->ne[1];
+    const int n_q_heads  = (int)Q->ne[2];
+    const int n_kv_heads = (int)K->ne[2];
+    const int n_batch    = (int)Q->ne[3];
+    const int n_kv       = (int)K->ne[1];
+    const int gqa_ratio  = n_q_heads / n_kv_heads;
+    const int n_query_rows = n_queries * gqa_ratio;
+
+    GGML_ASSERT(n_q_heads % n_kv_heads == 0);
+    GGML_ASSERT(max_bias == 0.0f);  // ALiBi not supported
+    GGML_ASSERT(Q->ne[3] == K->ne[3] || K->ne[3] == 1);
+
+    const int chunk_size = std::min(MKL_FA_CHUNK_SIZE_KV, n_kv);
+
+    // Query rows are processed in tiles of q_tile_rows so the score buffers
+    // (KQ_f32/S_f16 = q_tile_rows * chunk_size) stay bounded regardless of
+    // batch size. n_query_rows <= Q_TILE is a single tile (no extra work).
+    static int q_tile_env = ggml_sycl_get_env("GGML_SYCL_MKL_FA_Q_TILE", MKL_FA_Q_TILE);
+    const int q_tile_rows = std::max(1, std::min(q_tile_env, n_query_rows));
+
+    const int64_t wg_size = MKL_FA_WG_SIZE;
+
+    // --- Debug output (gated by GGML_SYCL_MKL_FA_DEBUG=1) ---
+    static int mkl_call_count = 0;
+    mkl_call_count++;
+    static int mkl_debug = ggml_sycl_get_env("GGML_SYCL_MKL_FA_DEBUG", 0);
+    const bool do_print = (mkl_debug == 1);
+
+    const int64_t q_row_stride  = Q->nb[1] / sizeof(float);
+    const int64_t q_head_stride = Q->nb[2] / sizeof(float);
+
+    const bool V_is_K_view = V->view_src
+        && (V->view_src == K || (V->view_src == K->view_src
+            && V->view_offs == K->view_offs));
+
+    // Early interleaved detection for debug output.
+    // True interleaved detection happens after dequant (nb12_fp16 == nb11_fp16),
+    // but we can pre-detect on the original tensor strides.
+    const bool k_early_interleaved =
+        ((int64_t)K->ne[1] * K->nb[1] != K->nb[2]);
+    const bool v_early_interleaved =
+        !V_is_K_view && ((int64_t)V->ne[1] * V->nb[1] != V->nb[2]);
+
+    if (do_print) {
+        GGML_LOG_INFO("[MKL-FA] #%d D=%d DV=%d n_q=%d n_kv=%d "
+                "n_qh=%d n_kvh=%d gqa=%d batch=%d K=%s V=%s "
+                "chunk=%d buf=%.1fMB%s%s\n",
+                mkl_call_count, DKQ, DV, n_queries, n_kv,
+                n_q_heads, n_kv_heads, gqa_ratio, n_batch,
+                ggml_type_name(K->type), ggml_type_name(V->type),
+                chunk_size,
+                (double)((int64_t)n_query_rows * chunk_size * sizeof(float))
+                    / (1024.0 * 1024.0),
+                k_early_interleaved ? " K_ILV" : "",
+                v_early_interleaved ? " V_ILV" : "");
+        GGML_LOG_INFO("[MKL-FA] #%d Q-nb1=%lld Q-nb2=%lld "
+                "q_rs=%lld q_hs=%lld dst_rs=%lld dst_hs=%lld\n",
+                mkl_call_count,
+                (long long)Q->nb[1], (long long)Q->nb[2],
+                (long long)q_row_stride, (long long)q_head_stride,
+                (long long)(KQV->nb[1] / sizeof(float)),
+                (long long)(KQV->nb[2] / sizeof(float)));
+    }
+
+    // --- Stream and allocators ---
+    dpct::queue_ptr stream = ctx.stream();
+
+#define MKL_TAKE_TIME(t0)  auto t0 = std::chrono::steady_clock::now()
+#define MKL_ACCUM(acc, t0) do { if (do_print) { \
+    acc += (int64_t)std::chrono::duration_cast \
+    <std::chrono::microseconds>(std::chrono::steady_clock::now() - (t0)).count(); \
+} } while(0)
+
+    int64_t gemm_kq_time_us  = 0;
+    int64_t gemm_vkq_time_us = 0;
+    int64_t softmax_time_us  = 0;
+    int64_t dequant_time_us  = 0;
+
+    MKL_TAKE_TIME(t_deq);
+
+    // --- K/V dequant descriptors ---
+    // Dequant is done per-chunk inside the KV loop (footprint independent of
+    // context). Output is always dense row-major fp16 [this_chunk x D], lda=D.
+    // Interleaved detection: ne[1]*nb[1] != nb[2] means heads are interleaved.
+    const bool k_interleaved =
+        ((int64_t)K->ne[1] * K->nb[1] != K->nb[2]) && K->ne[2] > 1;
+    const bool v_interleaved =
+        ((int64_t)V->ne[1] * V->nb[1] != V->nb[2]) && V->ne[2] > 1;
+
+    const mkl_fa_kv_desc K_desc = mkl_fa_make_desc(K, k_interleaved, n_kv_heads);
+    const mkl_fa_kv_desc V_desc = V_is_K_view
+        ? K_desc : mkl_fa_make_desc(V, v_interleaved, n_kv_heads);
+
+    MKL_ACCUM(dequant_time_us, t_deq);
+
+    // --- Resolve mask pointers ---
+    const sycl::half * mask_data = nullptr;
+    int64_t mask_head_stride = 0;
+    int64_t mask_row_stride  = 0;
+    int     mask_n_heads     = 0;
+
+    if (mask) {
+        // Use actual fp16 device size (2 bytes), NOT sizeof(sycl::half)
+        // which may be 4 on the host in oneAPI.
+        mask_head_stride = mask->nb[2] / 2;
+        mask_row_stride  = mask->nb[1] / 2;
+        mask_n_heads     = (int)mask->ne[2];
+    }
+
+    // --- Allocate intermediates from pool ---
+    ggml_sycl_pool & pool = ctx.pool();
+
+    ggml_sycl_pool_alloc<float>      KQ_f32(pool);      // [q_tile_rows x chunk]
+    ggml_sycl_pool_alloc<sycl::half> S_f16(pool);       // [q_tile_rows x chunk]
+    ggml_sycl_pool_alloc<float>      VKQ_chunk(pool);   // [q_tile_rows x DV]
+    ggml_sycl_pool_alloc<float>      VKQ_accum(pool);   // [n_query_rows x DV] (full)
+    ggml_sycl_pool_alloc<float>      KQ_max(pool);      // [n_query_rows] (full)
+    ggml_sycl_pool_alloc<float>      KQ_sum(pool);      // [n_query_rows] (full)
+    ggml_sycl_pool_alloc<sycl::half> Q_head_f16(pool);  // [n_query_rows x DKQ] (full)
+    ggml_sycl_pool_alloc<sycl::half> K_chunk_f16(pool); // [chunk x DKQ] (per-chunk dequant)
+    ggml_sycl_pool_alloc<sycl::half> V_chunk_f16(pool); // [chunk x DV] (per-chunk dequant)
+
+    KQ_f32.alloc((size_t)q_tile_rows * chunk_size);
+    S_f16.alloc((size_t)q_tile_rows * chunk_size);
+    VKQ_chunk.alloc((size_t)q_tile_rows * DV);
+    VKQ_accum.alloc((size_t)n_query_rows * DV);
+    KQ_max.alloc(n_query_rows);
+    KQ_sum.alloc(n_query_rows);
+    Q_head_f16.alloc((size_t)n_query_rows * DKQ);
+    K_chunk_f16.alloc((size_t)chunk_size * DKQ);
+
+    sycl::half * V_chunk_f16_ptr;
+    if (V_is_K_view) {
+        V_chunk_f16_ptr = K_chunk_f16.ptr;   // V aliases K (DV == DKQ)
+    } else {
+        V_chunk_f16.alloc((size_t)chunk_size * DV);
+        V_chunk_f16_ptr = V_chunk_f16.ptr;
+    }
+
+    sycl::half * Q_head_f16_ptr  = Q_head_f16.ptr;
+    float      * KQ_f32_ptr      = KQ_f32.ptr;
+    sycl::half * S_f16_ptr       = S_f16.ptr;
+    float      * VKQ_chunk_ptr   = VKQ_chunk.ptr;
+    float      * VKQ_accum_ptr   = VKQ_accum.ptr;
+    float      * KQ_max_ptr      = KQ_max.ptr;
+    float      * KQ_sum_ptr      = KQ_sum.ptr;
+    sycl::half * K_chunk_f16_ptr = K_chunk_f16.ptr;
+
+    const float alpha = 1.0f;
+    const float beta  = 0.0f;
+
+    for (int ib = 0; ib < n_batch; ib++) {
+        const float * Q_batch = (const float *)Q->data
+            + ib * (Q->nb[3] / sizeof(float));
+        float * dst_batch = (float *)KQV->data
+            + ib * (KQV->nb[3] / sizeof(float));
+
+        const sycl::half * mask_batch = nullptr;
+        if (mask) {
+            int m_batch = (mask->ne[3] > 1) ? ib : 0;
+            mask_batch = (const sycl::half *)mask->data
+                + m_batch * (mask->nb[3] / 2);  // 2 = actual fp16 device size
+        }
+
+        for (int ikvh = 0; ikvh < n_kv_heads; ikvh++) {
+            int kvh_base_head = ikvh * gqa_ratio;
+
+            // 1. Pack all GQA Q heads into fp16 (full n_query_rows)
+            mkl_fa_pack_q_fp16(stream,
+                Q_head_f16_ptr, Q_batch,
+                n_queries, n_query_rows, DKQ,
+                gqa_ratio, kvh_base_head,
+                q_scale, q_row_stride, q_head_stride, wg_size);
+
+            // 2. Initialize softmax state (full n_query_rows)
+            mkl_fa_init_softmax_state(stream,
+                KQ_max_ptr, KQ_sum_ptr, VKQ_accum_ptr,
+                n_query_rows, DV, wg_size);
+
+            // Sync before MKL GEMM (MKL may use an internal queue)
+            stream->wait();
+
+            // 3. KV chunk loop (OUTER): dequant each chunk once, then tile queries.
+            for (int chunk_start = 0; chunk_start < n_kv; chunk_start += chunk_size) {
+                int this_chunk = std::min(chunk_size, n_kv - chunk_start);
+
+                // 3a. Dequant this KV chunk to dense fp16 (once per chunk)
+                {
+                    MKL_TAKE_TIME(t0);
+                    mkl_fa_dequant_chunk(stream, K_desc, KQV,
+                        K_chunk_f16_ptr, ikvh, chunk_start, this_chunk);
+                    if (!V_is_K_view) {
+                        mkl_fa_dequant_chunk(stream, V_desc, KQV,
+                            V_chunk_f16_ptr, ikvh, chunk_start, this_chunk);
+                    }
+                    stream->wait();  // dequant must be ready before MKL GEMM
+                    MKL_ACCUM(dequant_time_us, t0);
+                }
+
+                // 3b. Query tile loop (INNER) — bounds KQ_f32/S_f16 footprint.
+                for (int q0 = 0; q0 < n_query_rows; q0 += q_tile_rows) {
+                    int q_rows = std::min(q_tile_rows, n_query_rows - q0);
+
+                    // GEMM: KQ = Q_tile × K_chunk^T
+                    {
+                        MKL_TAKE_TIME(t0);
+                        sycl::event ev = gemm(*stream,
+                            transpose::trans, transpose::nontrans,
+                            this_chunk, q_rows, DKQ,
+                            alpha,
+                            K_chunk_f16_ptr, DKQ,
+                            Q_head_f16_ptr + (int64_t)q0 * DKQ, DKQ,
+                            beta,
+                            KQ_f32_ptr, this_chunk);
+                        try { ev.wait_and_throw(); } catch (sycl::exception & e) {
+                            GGML_LOG_INFO("[MKL-FA] GEMM KQ: %s\n", e.what());
+                            GGML_ABORT("MKL GEMM KQ failed");
+                        }
+                        MKL_ACCUM(gemm_kq_time_us, t0);
+                    }
+                    // Online softmax over this chunk for this query tile
+                    {
+                        MKL_TAKE_TIME(t0);
+                        mkl_fa_online_softmax_chunk(stream,
+                            KQ_f32_ptr, S_f16_ptr,
+                            KQ_max_ptr, KQ_sum_ptr, VKQ_accum_ptr,
+                            q0, q_rows, n_queries, DV,
+                            this_chunk, chunk_start,
+                            kvh_base_head, gqa_ratio,
+                            mask_batch, mask_head_stride,
+                            mask_row_stride, mask_n_heads,
+                            logit_softcap, wg_size);
+                        stream->wait();  // S_f16 must be ready for GEMM
+                        MKL_ACCUM(softmax_time_us, t0);
+                    }
+
+                    // GEMM: VKQ_chunk = S × V_chunk
+                    {
+                        MKL_TAKE_TIME(t0);
+                        sycl::event ev = gemm(*stream,
+                            transpose::nontrans, transpose::nontrans,
+                            DV, q_rows, this_chunk,
+                            alpha,
+                            V_chunk_f16_ptr, DV,
+                            S_f16_ptr, this_chunk,
+                            beta,
+                            VKQ_chunk_ptr, DV);
+                        try { ev.wait_and_throw(); } catch (sycl::exception & e) {
+                            GGML_LOG_INFO("[MKL-FA] GEMM VKQ: %s\n", e.what());
+                            GGML_ABORT("MKL GEMM VKQ failed");
+                        }
+                        MKL_ACCUM(gemm_vkq_time_us, t0);
+                    }
+                    // VKQ_accum[q0..] += VKQ_chunk
+                    {
+                        const int64_t n_total = (int64_t)q_rows * DV;
+                        const int64_t wg = ((n_total + wg_size - 1) / wg_size)
+                            * wg_size;
+                        float * accum = VKQ_accum_ptr + (int64_t)q0 * DV;
+                        stream->submit([&](sycl::handler & cgh) {
+                            cgh.parallel_for(sycl::nd_range<1>(wg, wg_size),
+                                [=](sycl::nd_item<1> item) {
+                                    int64_t i = item.get_global_id(0);
+                                    if (i < n_total) {
+                                        accum[i] += VKQ_chunk_ptr[i];
+                                    }
+                                });
+                        });
+                    }
+                }
+            }
+
+            // 4. Normalize and scatter each GQA head to dst
+            for (int iqg = 0; iqg < gqa_ratio; iqg++) {
+                int     iqh        = kvh_base_head + iqg;
+                int64_t src_offset = (int64_t)iqg * n_queries * DV;
+                mkl_fa_normalize_head(stream,
+                    dst_batch, VKQ_accum_ptr, KQ_sum_ptr,
+                    iqh, n_queries, DV, n_q_heads,
+                    src_offset, wg_size);
+            }
+        }
+    }
+
+#undef MKL_TAKE_TIME
+#undef MKL_ACCUM
+
+    if (do_print) {
+        const int64_t v_chunk_elems = V_is_K_view ? 0 : (int64_t)chunk_size * DV;
+        double total_mb = (double)(
+            (int64_t)q_tile_rows * chunk_size * sizeof(float)      // KQ_f32
+          + (int64_t)q_tile_rows * chunk_size * sizeof(sycl::half) // S_f16
+          + (int64_t)q_tile_rows * DV * sizeof(float)              // VKQ_chunk
+          + (int64_t)n_query_rows * DV * sizeof(float)             // VKQ_accum
+          + (int64_t)n_query_rows * sizeof(float)                  // KQ_max
+          + (int64_t)n_query_rows * sizeof(float)                  // KQ_sum
+          + (int64_t)n_query_rows * DKQ * sizeof(sycl::half)       // Q_head_f16
+          + (int64_t)chunk_size * DKQ * sizeof(sycl::half)         // K_chunk_f16
+          + v_chunk_elems * (int64_t)sizeof(sycl::half)            // V_chunk_f16
+        ) / (1024.0 * 1024.0);
+        GGML_LOG_INFO("[MKL-FA] #%d n_kv=%d n_q=%d q_tile=%d time_us: "
+                "dequant=%lld GEMM_KQ=%lld softmax=%lld GEMM_VKQ=%lld "
+                "buf_mb=%.1f\n",
+                mkl_call_count, n_kv, n_queries, q_tile_rows,
+                (long long)dequant_time_us,
+                (long long)gemm_kq_time_us,
+                (long long)softmax_time_us,
+                (long long)gemm_vkq_time_us,
+                total_mb);
+    }
+}

ggml/src/ggml-sycl/fattn.cpp

@@ -99,8 +99,10 @@ enum best_fattn_kernel {
     BEST_FATTN_KERNEL_VEC      = 100,
     BEST_FATTN_KERNEL_ONEDNN   = 150, // added enum for onednn==150
     BEST_FATTN_KERNEL_TILE     = 200,
+    BEST_FATTN_KERNEL_MKL      = 300,
 };
 
+
 static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const ggml_tensor * dst) {
     GGML_UNUSED(device);
 #ifndef SYCL_FLASH_ATTN
@@ -115,6 +117,7 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const
     const ggml_tensor * K     = dst->src[1];
     const ggml_tensor * V     = dst->src[2];
     const ggml_tensor * mask  = dst->src[3];
+    const ggml_tensor * sinks = dst->src[4];
 
     const int gqa_ratio = Q->ne[2] / K->ne[2];
     GGML_ASSERT(Q->ne[2] % K->ne[2] == 0);
@@ -122,7 +125,49 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const
     float max_bias = 0.0f;
     memcpy(&max_bias, (const float *) KQV->op_params + 1, sizeof(float));
 
+    float logit_softcap = 0.0f;
+    memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float));
+
     bool gqa_opt_applies = gqa_ratio >= 2 && mask && max_bias == 0.0f && K->ne[1] % FATTN_KQ_STRIDE == 0;
+
+    // MKL path: XMX-accelerated GEMM for prompt processing (all KV cache types).
+    // The MKL kernel converts non-F16 K/V to F16 via to_fp16_sycl before GEMM,
+    // so quantized, F16, BF16, and F32 caches all benefit from XMX acceleration.
+    // Activates automatically when flash-attn is enabled (--flash-attn on or -fa)
+    // and n_kv >= 1024. Falls through to TILE/VEC for ALiBi, logit softcap,
+    // and mismatched batch dimensions (unsupported by the MKL kernel).
+    // Set GGML_SYCL_ENABLE_MKL_FA=0 to force TILE/VEC path for A/B testing.
+    // Example: GGML_SYCL_ENABLE_MKL_FA=0 llama-cli -m model.gguf -fa -ngl 99 ...
+    // Note: MKL GEMM calls are incompatible with SYCL graph capture replay.
+    static int mkl_enable = ggml_sycl_get_env("GGML_SYCL_ENABLE_MKL_FA", 1);
+    // MKL is validated for the mainstream GQA envelope: grouped-query
+    // (gqa_ratio >= 2), head_dim a multiple of 64 in [64,512] with matching
+    // K/V head size, mask, no sinks/ALiBi/softcap. Gemma's global layers use
+    // head_dim 512, so the cap must include it. Head sizes not a multiple of
+    // 64 (72/80/96), MHA (gqa_ratio == 1), and MLA (DKQ != DV, e.g. 576/512)
+    // fall through to TILE/VEC; see follow-up work.
+    if (mkl_enable == 1 && mask && !sinks && gqa_ratio >= 2 &&
+        Q->ne[0] >= 64 && Q->ne[0] <= 512 && Q->ne[0] % 64 == 0 &&
+        Q->ne[0] == V->ne[0] &&
+        Q->ne[1] >= 32 && K->ne[1] >= 1024 &&
+        max_bias == 0.0f && logit_softcap == 0.0f &&
+        (Q->ne[3] == K->ne[3] || K->ne[3] == 1)) {
+        // F16 K/V strides must be a multiple of ne[0]*2 (the natural row size
+        // in bytes). This passes both dense (nb1 == ne0*2) and interleaved
+        // (nb1 == H * ne0*2). Only pathological test strides like nb1=32 or
+        // nb1=75 for ne0=40 fall through to TILE.
+        bool kv_strides_ok = true;
+        for (const ggml_tensor * t : {K, V}) {
+            if (t->type == GGML_TYPE_F16 && t->nb[1] % (t->ne[0] * 2) != 0) {
+                kv_strides_ok = false;
+                break;
+            }
+        }
+        if (kv_strides_ok) {
+            return BEST_FATTN_KERNEL_MKL;
+        }
+    }
+
     for (const ggml_tensor * t : {Q, K, V, mask}) {
         if (t == nullptr || ggml_is_quantized(t->type)) {
             continue;
@@ -216,6 +261,37 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const
 
 void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     ggml_sycl_set_device(ctx.device);
+
+    // n_kv watchdog: log when n_kv differs from the last FA call with
+    // the same D — helps detect cache-truncation issues.
+    static int nkv_debug = ggml_sycl_get_env("GGML_SYCL_MKL_FA_DEBUG", 0);
+    if (nkv_debug == 1) {
+        const ggml_tensor * K_dbg = dst->src[1];
+        const ggml_tensor * V_dbg = dst->src[2];
+        static int64_t last_nkv_d256 = 0, last_nkv_d512 = 0;
+        static int fa_call_seq = 0;
+        fa_call_seq++;
+        int64_t cur_nkv = K_dbg->ne[1];
+        int Dk = (int)K_dbg->ne[0];
+        const char * kname = "TILE";
+        best_fattn_kernel k = ggml_sycl_get_best_fattn_kernel(ctx.device, dst);
+        if (k == BEST_FATTN_KERNEL_MKL)  kname = "MKL";
+        if (k == BEST_FATTN_KERNEL_VEC)  kname = "VEC";
+        int64_t delta = 0;
+        if (Dk == 256) {
+            delta = cur_nkv - last_nkv_d256;
+            last_nkv_d256 = cur_nkv;
+        } else if (Dk == 512) {
+            delta = cur_nkv - last_nkv_d512;
+            last_nkv_d512 = cur_nkv;
+        }
+        GGML_LOG_INFO("[FA-DISP] #%d %s D=%d n_kv=%lld delta=%lld "
+                "V_ne1=%lld\n",
+                fa_call_seq, kname, Dk,
+                (long long)cur_nkv, (long long)delta,
+                (long long)V_dbg->ne[1]);
+    }
+
     switch (ggml_sycl_get_best_fattn_kernel(ggml_sycl_get_device(), dst)) {
         case BEST_FATTN_KERNEL_NONE:
             GGML_ABORT("Not support Flash-Attention");
@@ -232,6 +308,51 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst
         case BEST_FATTN_KERNEL_VEC:
             ggml_sycl_flash_attn_ext_vec(ctx, dst);
             break;
+        case BEST_FATTN_KERNEL_MKL:
+            ggml_sycl_flash_attn_ext_mkl(ctx, dst);
+            break;
+    }
+
+    // --- Output fingerprint (GGML_SYCL_MKL_FA_DIAG=1) ---
+    // Copy first 64 float output values to host for fingerprinting.
+    // Compare MKL vs TILE (GGML_SYCL_ENABLE_MKL_FA=0) to detect divergence.
+    // Only fingerprints the first 6 FA calls with n_kv >= 1024.
+    static int fa_diag = ggml_sycl_get_env("GGML_SYCL_MKL_FA_DIAG", 0);
+    static int fa_diag_count = 0;
+    if (fa_diag == 1 && fa_diag_count < 6) {
+        const ggml_tensor * K_diag = dst->src[1];
+        const ggml_tensor * V_diag = dst->src[2];
+        const ggml_tensor * Q_diag = dst->src[0];
+        if (K_diag->ne[1] >= 1024) {
+            fa_diag_count++;
+            float diag_buf[64];
+            dpct::queue_ptr q = ctx.stream();
+            q->memcpy(diag_buf, dst->data, 64 * sizeof(float));
+            q->wait();
+            const char * kname = "???";
+            best_fattn_kernel kb = ggml_sycl_get_best_fattn_kernel(ctx.device, dst);
+            if (kb == BEST_FATTN_KERNEL_MKL) kname = "MKL";
+            if (kb == BEST_FATTN_KERNEL_TILE) kname = "TILE";
+            if (kb == BEST_FATTN_KERNEL_VEC) kname = "VEC";
+            GGML_LOG_INFO("[FA-DIAG] #%d %s D=%d n_kv=%lld n_q=%lld "
+                    "n_qh=%lld n_kvh=%lld K=%s V=%s "
+                    "nb1=%zu nb2=%zu first 64 floats:\n",
+                    fa_diag_count, kname,
+                    (int)K_diag->ne[0], (long long)K_diag->ne[1],
+                    (long long)Q_diag->ne[1],
+                    (long long)Q_diag->ne[2], (long long)K_diag->ne[2],
+                    ggml_type_name(K_diag->type),
+                    ggml_type_name(V_diag->type),
+                    K_diag->nb[1], K_diag->nb[2]);
+            for (int i = 0; i < 64; i += 8) {
+                GGML_LOG_INFO("  [%2d] %08x %08x %08x %08x %08x %08x %08x %08x\n",
+                        i,
+                        *(unsigned *)&diag_buf[i+0], *(unsigned *)&diag_buf[i+1],
+                        *(unsigned *)&diag_buf[i+2], *(unsigned *)&diag_buf[i+3],
+                        *(unsigned *)&diag_buf[i+4], *(unsigned *)&diag_buf[i+5],
+                        *(unsigned *)&diag_buf[i+6], *(unsigned *)&diag_buf[i+7]);
+            }
+        }
     }
 }
 

ggml/src/ggml-sycl/fattn.hpp

@@ -19,4 +19,6 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst
 
 bool ggml_sycl_flash_attn_ext_supported(int device, const ggml_tensor * dst);
 
+void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
+
 #endif // GGML_SYCL_FATTN_HPP

ggml/src/ggml-sycl/fusion.cpp

@@ -0,0 +1,44 @@
+#include "fusion.hpp"
+
+bool ggml_sycl_can_fuse(const ggml_cgraph * cgraph, int node_idx, std::initializer_list<enum ggml_op> ops) {
+    if (!g_ggml_sycl_enable_fusion) {
+        return false;
+    }
+
+    if (!ggml_can_fuse(cgraph, node_idx, ops)) {
+        return false;
+    }
+
+    if (ops.size() == 2 && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) {
+        const ggml_tensor * rms_norm = cgraph->nodes[node_idx];
+        const ggml_tensor * mul      = cgraph->nodes[node_idx + 1];
+
+        GGML_ASSERT(rms_norm->src[0]->type == GGML_TYPE_F32);
+        GGML_ASSERT(rms_norm->type == GGML_TYPE_F32);
+
+        if (mul->src[0]->type != GGML_TYPE_F32 ||
+            mul->src[1]->type != GGML_TYPE_F32 ||
+            mul->type != GGML_TYPE_F32) {
+            return false;
+        }
+
+        // if rms norm is the B operand, then we don't handle broadcast
+        if (rms_norm == mul->src[1] && !ggml_are_same_shape(mul->src[0], rms_norm)) {
+            return false;
+        }
+
+        const ggml_tensor * mul_w = (mul->src[0] == rms_norm) ? mul->src[1] : mul->src[0];
+        // the fused kernel indexes the weight as mul[col], so it must span ncols contiguously
+        if (mul_w->ne[0] != rms_norm->ne[0] || mul_w->nb[0] != ggml_type_size(mul_w->type)) {
+            return false;
+        }
+
+        if (!ggml_is_contiguous_rows(mul->src[0]) || !ggml_is_contiguous_rows(mul->src[1])) {
+            return false;
+        }
+
+        return true;
+    }
+
+    return false;
+}

ggml/src/ggml-sycl/fusion.hpp

@@ -0,0 +1,15 @@
+#ifndef GGML_SYCL_FUSION_HPP
+#define GGML_SYCL_FUSION_HPP
+
+#include <initializer_list>
+
+#include "common.hpp"
+
+// Backend-side fusability test. `ops` names a candidate op sequence starting at cgraph node
+// `node_idx`; the result is true only if ggml considers that subgraph fusable *and* the SYCL
+// kernel which would service it accepts the tensors involved (types, shapes, contiguity).
+//
+// Lives in its own translation unit because it grows a branch per supported op sequence.
+bool ggml_sycl_can_fuse(const ggml_cgraph * cgraph, int node_idx, std::initializer_list<enum ggml_op> ops);
+
+#endif  // GGML_SYCL_FUSION_HPP

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

@@ -167,7 +167,10 @@ static ggml_sycl_device_info ggml_sycl_init() {
             ze_device_properties_t props = {};
             props.stype = ZE_STRUCTURE_TYPE_DEVICE_PROPERTIES;
             ze_result_t r = zeDeviceGetProperties(ze_dev, &props);
-            info.devices[i].l0_discrete_gpu = r == ZE_RESULT_SUCCESS && !(props.flags & ZE_DEVICE_PROPERTY_FLAG_INTEGRATED);
+            if (r == ZE_RESULT_SUCCESS) {
+                info.devices[i].l0_device_type_valid = true;
+                info.devices[i].l0_discrete_gpu = !(props.flags & ZE_DEVICE_PROPERTY_FLAG_INTEGRATED);
+            }
         }
 #endif
     }
@@ -274,6 +277,8 @@ static const char* dev2dev_int2str(int dev2dev) {
         return "SYCL API";
     } else if (dev2dev == DEV2DEV_MEMCPY_L0) {
         return "Level Zero API";
+    } else if (dev2dev == DEV2DEV_MEMCPY_FORWARD) {
+        return "Host Forward";
     } else {
         return "Unknown";
     }
@@ -684,7 +689,11 @@ static void dev2dev_memcpy(int device_dst, sycl::queue &q_dst, int device_src, s
     }
 
     // Host-staged copy
-    GGML_SYCL_DEBUG("[SYCL] dev2dev memcpy by host forward\n");
+    if(g_ggml_sycl_dev2dev_memcpy == DEV2DEV_MEMCPY_FORWARD) {
+        GGML_SYCL_DEBUG("[SYCL] dev2dev memcpy by host forward for setting GGML_SYCL_DEV2DEV_MEMCPY=2\n");
+    } else {
+        GGML_SYCL_DEBUG("[SYCL] dev2dev memcpy by host forward for SYCL/L0 fallback\n");
+    }
     char *host_buf = (char *)malloc(size);
     q_src.memcpy(host_buf, (const char *)ptr_src, size).wait();
     q_dst.memcpy((char *)ptr_dst, host_buf, size).wait();
@@ -5398,6 +5407,13 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc
             }
         }
 #endif
+        if (node->op == GGML_OP_RMS_NORM &&
+            ggml_sycl_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL })) {
+            ggml_sycl_op_rms_norm_fused(*sycl_ctx, node, cgraph->nodes[i + 1]);
+            i++;
+            continue;
+        }
+
         bool ok = ggml_sycl_compute_forward(*sycl_ctx, node);
         if (!ok) {
             GGML_LOG_ERROR("%s: error: op not supported %s (%s)\n", __func__, node->name, ggml_op_name(node->op));
@@ -5593,7 +5609,11 @@ static void ggml_backend_sycl_device_get_memory(ggml_backend_dev_t dev, size_t *
 }
 
 static enum ggml_backend_dev_type ggml_backend_sycl_device_get_type(ggml_backend_dev_t dev) {
-    GGML_UNUSED(dev);
+    ggml_backend_sycl_device_context * ctx = (ggml_backend_sycl_device_context *)dev->context;
+    const sycl_device_info & info = ggml_sycl_info().devices[ctx->device];
+    if (info.l0_device_type_valid && !info.l0_discrete_gpu) {
+        return GGML_BACKEND_DEVICE_TYPE_IGPU;
+    }
     return GGML_BACKEND_DEVICE_TYPE_GPU;
 }
 

ggml/src/ggml-sycl/mmvq.cpp

@@ -1254,6 +1254,66 @@ static void mul_mat_vec_q1_0_q8_1_sycl_switch_ncols(
     }
 }
 
+static void mul_mat_vec_q2_0_q8_1_sycl(const void * vx, const void * vy,
+                                       float * dst, const int ncols,
+                                       const int nrows,
+                                       dpct::queue_ptr stream) {
+    GGML_ASSERT(ncols % QK2_0 == 0);
+    const int block_num_y = (nrows + GGML_SYCL_MMV_Y - 1) / GGML_SYCL_MMV_Y;
+    const sycl::range<3> block_nums(1, 1, block_num_y);
+    const sycl::range<3> block_dims(1, GGML_SYCL_MMV_Y, WARP_SIZE);
+
+    stream->submit([&](sycl::handler & cgh) {
+        cgh.parallel_for(
+            sycl::nd_range<3>(block_nums * block_dims, block_dims),
+            [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                mul_mat_vec_q<QK2_0, QI2_0, block_q2_0,
+                              VDR_Q2_0_Q8_1_MMVQ, vec_dot_q2_0_q8_1>(
+                    vx, vy, dst, ncols, nrows, item_ct1);
+            });
+    });
+}
+
+template <int ncols_dst>
+static void mul_mat_vec_q2_0_q8_1_sycl_ncols(
+        const void * vx, const void * vy, float * dst,
+        const int ncols, const int nrows,
+        const int stride_col_y, const int stride_col_dst,
+        dpct::queue_ptr stream) {
+    GGML_ASSERT(ncols % QK2_0 == 0);
+    const int block_num_y = (nrows + GGML_SYCL_MMV_Y - 1) / GGML_SYCL_MMV_Y;
+    const sycl::range<3> block_nums(1, 1, block_num_y);
+    const sycl::range<3> block_dims(1, GGML_SYCL_MMV_Y, WARP_SIZE);
+
+    stream->submit([&](sycl::handler & cgh) {
+        cgh.parallel_for(
+            sycl::nd_range<3>(block_nums * block_dims, block_dims),
+            [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                mul_mat_vec_q_ncols<QK2_0, QI2_0, block_q2_0,
+                                    VDR_Q2_0_Q8_1_MMVQ, vec_dot_q2_0_q8_1, ncols_dst>(
+                    vx, vy, dst, ncols, nrows, stride_col_y, stride_col_dst, item_ct1);
+            });
+    });
+}
+
+static void mul_mat_vec_q2_0_q8_1_sycl_switch_ncols(
+        const void * vx, const void * vy, float * dst,
+        const int ncols, const int nrows, const int ncols_dst,
+        const int stride_col_y, const int stride_col_dst,
+        dpct::queue_ptr stream) {
+    switch (ncols_dst) {
+        case 1: mul_mat_vec_q2_0_q8_1_sycl(vx, vy, dst, ncols, nrows, stream); break;
+        case 2: mul_mat_vec_q2_0_q8_1_sycl_ncols<2>(vx, vy, dst, ncols, nrows, stride_col_y, stride_col_dst, stream); break;
+        case 3: mul_mat_vec_q2_0_q8_1_sycl_ncols<3>(vx, vy, dst, ncols, nrows, stride_col_y, stride_col_dst, stream); break;
+        case 4: mul_mat_vec_q2_0_q8_1_sycl_ncols<4>(vx, vy, dst, ncols, nrows, stride_col_y, stride_col_dst, stream); break;
+        case 5: mul_mat_vec_q2_0_q8_1_sycl_ncols<5>(vx, vy, dst, ncols, nrows, stride_col_y, stride_col_dst, stream); break;
+        case 6: mul_mat_vec_q2_0_q8_1_sycl_ncols<6>(vx, vy, dst, ncols, nrows, stride_col_y, stride_col_dst, stream); break;
+        case 7: mul_mat_vec_q2_0_q8_1_sycl_ncols<7>(vx, vy, dst, ncols, nrows, stride_col_y, stride_col_dst, stream); break;
+        case 8: mul_mat_vec_q2_0_q8_1_sycl_ncols<8>(vx, vy, dst, ncols, nrows, stride_col_y, stride_col_dst, stream); break;
+        default: GGML_ABORT("unsupported ncols_dst=%d for Q2_0 multi-col MMVQ", ncols_dst);
+    }
+}
+
 static void mul_mat_vec_q2_K_q8_1_sycl(const void *vx, const void *vy,
                                        float *dst, const int ncols,
                                        const int nrows,
@@ -2194,6 +2254,20 @@ void ggml_sycl_op_mul_mat_vec_q(ggml_backend_sycl_context & ctx, const ggml_tens
                     mul_mat_vec_q1_0_q8_1_sycl(src0_dd_i, src1_ddq_i_bs, dst_dd_i_bs, ne00, row_diff, stream);
                 }
                 break;
+            case GGML_TYPE_Q2_0:
+                if (i == 0 && src1_ncols > 1 && src1_ncols <= 8) {
+                    const int stride_col_y   = src1_padded_col_size / QK8_1;
+                    const int stride_col_dst = dst->ne[0];
+                    GGML_SYCL_DEBUG("Calling mul_mat_vec_q2_0_q8_1_sycl_switch_ncols ncols=%d\n", (int)src1_ncols);
+                    mul_mat_vec_q2_0_q8_1_sycl_switch_ncols(
+                        src0_dd_i, src1_ddq_i, dst_dd_i, ne00, row_diff,
+                        src1_ncols, stride_col_y, stride_col_dst, stream);
+                    return;
+                } else if (i == 0 || src1_ncols == 1) {
+                    GGML_SYCL_DEBUG("Calling mul_mat_vec_q2_0_q8_1_sycl\n");
+                    mul_mat_vec_q2_0_q8_1_sycl(src0_dd_i, src1_ddq_i_bs, dst_dd_i_bs, ne00, row_diff, stream);
+                }
+                break;
             case GGML_TYPE_Q2_K:
                 if (i == 0 && src1_ncols > 1 && src1_ncols <= 8) {
                     const int stride_col_y   = src1_padded_col_size / QK8_1;
@@ -2503,6 +2577,11 @@ bool ggml_sycl_mul_mat_vec_q_id(
                 vx_base, vy, ids_dev, dst_base, ncols, nrows, n_experts_used,
                 expert_weight_stride, dst_row_stride, src1_row_stride, stream);
             return true;
+        case GGML_TYPE_Q2_0:
+            launch_mul_mat_vec_q_moe<QK2_0, QI2_0, block_q2_0, VDR_Q2_0_Q8_1_MMVQ, vec_dot_q2_0_q8_1>(
+                vx_base, vy, ids_dev, dst_base, ncols, nrows, n_experts_used,
+                expert_weight_stride, dst_row_stride, src1_row_stride, stream);
+            return true;
         case GGML_TYPE_Q2_K:
             launch_mul_mat_vec_q_moe<QK_K, QI2_K, block_q2_K, VDR_Q2_K_Q8_1_MMVQ, vec_dot_q2_K_q8_1>(
                 vx_base, vy, ids_dev, dst_base, ncols, nrows, n_experts_used,

ggml/src/ggml-sycl/norm.cpp

@@ -147,10 +147,13 @@ static void group_norm_f32(const float* x, float* dst, const int group_size, con
     }
 }
 
+template <bool do_multiply = false>
 static void rms_norm_f32(const float* x, float* dst, const int ncols,
     const int64_t src_stride_col, const int64_t src_stride_row, const int64_t src_stride_channel, const int64_t src_stride_sample,
     const int64_t dst_stride_col, const int64_t dst_stride_row, const int64_t dst_stride_channel, const int64_t dst_stride_sample,
-    const float eps, const sycl::nd_item<3>& item_ct1, float* s_sum, int block_size) {
+    const float eps, const sycl::nd_item<3>& item_ct1, float* s_sum, int block_size,
+    const float* mul = nullptr, const int64_t mul_stride_row = 0, const int64_t mul_stride_channel = 0,
+    const int64_t mul_stride_sample = 0, const int mul_nrows = 0, const int mul_nchannels = 0, const int mul_nsamples = 0) {
 
     const int nrows = item_ct1.get_group_range(2);
     const int nchannels = item_ct1.get_group_range(1);
@@ -170,6 +173,12 @@ static void rms_norm_f32(const float* x, float* dst, const int ncols,
     x   += src_offset;
     dst += dst_offset;
 
+    if constexpr (do_multiply) {
+        const int mul_row     = row     % mul_nrows;
+        const int mul_channel = channel % mul_nchannels;
+        const int mul_sample  = sample  % mul_nsamples;
+        mul += mul_sample * mul_stride_sample + mul_channel * mul_stride_channel + mul_row * mul_stride_row;
+    }
 
     float tmp = 0.0f; // partial sum for thread in warp
 
@@ -202,7 +211,11 @@ static void rms_norm_f32(const float* x, float* dst, const int ncols,
     const float scale = sycl::rsqrt(mean + eps);
 
     for (int col = tid; col < ncols; col += block_size) {
-        dst[col * dst_stride_col] = scale * x[col * src_stride_col];
+        if constexpr (do_multiply) {
+            dst[col * dst_stride_col] = scale * x[col * src_stride_col] * mul[col];
+        } else {
+            dst[col * dst_stride_col] = scale * x[col * src_stride_col];
+        }
     }
 }
 
@@ -376,6 +389,49 @@ static void rms_norm_f32_sycl(const float* x, float* dst, const int ncols, const
     }
 }
 
+static void rms_norm_mul_f32_sycl(const float* x, const float* mul, float* dst, const int ncols, const int nrows,
+        const int nchannels, const int nsamples,
+        const int64_t src_stride_col, const int64_t src_stride_row, const int64_t src_stride_channel, const int64_t src_stride_sample,
+        const int64_t dst_stride_col, const int64_t dst_stride_row, const int64_t dst_stride_channel, const int64_t dst_stride_sample,
+        const int64_t mul_stride_row, const int64_t mul_stride_channel, const int64_t mul_stride_sample,
+        const int mul_nrows, const int mul_nchannels, const int mul_nsamples,
+        const float eps, queue_ptr stream, int device) {
+    const sycl::range<3> global_dims(nsamples, nchannels, nrows);
+    if (ncols < 1024) {
+        const sycl::range<3> block_dims(1, 1, WARP_SIZE);
+        stream->submit([&](sycl::handler& cgh) {
+            cgh.parallel_for(
+                sycl::nd_range<3>(global_dims * block_dims, block_dims),
+                [=](sycl::nd_item<3> item_ct1)
+                [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                    rms_norm_f32<true>(x, dst, ncols,
+                        src_stride_col, src_stride_row, src_stride_channel, src_stride_sample,
+                        dst_stride_col, dst_stride_row, dst_stride_channel, dst_stride_sample,
+                        eps, item_ct1, nullptr, WARP_SIZE,
+                        mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_nrows, mul_nchannels, mul_nsamples);
+                });
+            });
+    }
+    else {
+        const int work_group_size = ggml_sycl_info().max_work_group_sizes[device];
+        assert(work_group_size % (WARP_SIZE * WARP_SIZE) == 0);
+        const sycl::range<3> block_dims(1, 1, work_group_size);
+        stream->submit([&](sycl::handler& cgh) {
+            sycl::local_accessor<float, 1> s_sum_acc_ct1(sycl::range<1>(work_group_size / WARP_SIZE), cgh);
+            cgh.parallel_for(
+                sycl::nd_range<3>(global_dims * block_dims, block_dims),
+                [=](sycl::nd_item<3> item_ct1)
+                [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                    rms_norm_f32<true>(x, dst, ncols,
+                        src_stride_col, src_stride_row, src_stride_channel, src_stride_sample,
+                        dst_stride_col, dst_stride_row, dst_stride_channel, dst_stride_sample,
+                        eps, item_ct1, get_pointer(s_sum_acc_ct1), work_group_size,
+                        mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_nrows, mul_nchannels, mul_nsamples);
+                });
+            });
+    }
+}
+
 template<int warp_size>
 static void l2_norm_f32_sycl(const float *   x,
                              float *         dst,
@@ -518,6 +574,66 @@ void ggml_sycl_op_rms_norm(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
         ss0, ss1, ss2, ss3, ds0, ds1, ds2, ds3, eps, main_stream, ctx.device);
 }
 
+void ggml_sycl_op_rms_norm_fused(ggml_backend_sycl_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor) {
+    const ggml_tensor * rms_norm_src = dst->src[0];
+    float eps = 0.0f;
+    memcpy(&eps, dst->op_params, sizeof(float));
+
+    const float *       src0_dd = static_cast<const float *>(rms_norm_src->data);
+    const float *       mul_dd  = nullptr;
+    const ggml_tensor * mul_src = nullptr;
+    if (mul_tensor->src[0] == dst) {
+        mul_dd  = static_cast<const float *>(mul_tensor->src[1]->data);
+        mul_src = mul_tensor->src[1];
+    } else if (mul_tensor->src[1] == dst) {
+        mul_dd  = static_cast<const float *>(mul_tensor->src[0]->data);
+        mul_src = mul_tensor->src[0];
+    } else {
+        GGML_ASSERT(false);
+    }
+    float * dst_dd = static_cast<float *>(mul_tensor->data);
+
+    dpct::queue_ptr main_stream = ctx.stream();
+    SYCL_CHECK(ggml_sycl_set_device(ctx.device));
+
+    GGML_ASSERT(rms_norm_src->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type == GGML_TYPE_F32);
+    GGML_ASSERT(mul_tensor->type == GGML_TYPE_F32);
+    GGML_ASSERT(eps >= 0.0f);
+
+    const int64_t ne00 = rms_norm_src->ne[0];
+    const int64_t ne01 = rms_norm_src->ne[1];
+    const int64_t ne02 = rms_norm_src->ne[2];
+    const int64_t ne03 = rms_norm_src->ne[3];
+
+    const size_t ts0 = ggml_type_size(rms_norm_src->type);
+    GGML_ASSERT(rms_norm_src->nb[0] == ts0);
+    const int64_t s00 = rms_norm_src->nb[0] / ts0;
+    const int64_t s01 = rms_norm_src->nb[1] / ts0;
+    const int64_t s02 = rms_norm_src->nb[2] / ts0;
+    const int64_t s03 = rms_norm_src->nb[3] / ts0;
+
+    const size_t tdst = ggml_type_size(mul_tensor->type);
+    GGML_ASSERT(mul_tensor->nb[0] == tdst);
+    const int64_t d00 = mul_tensor->nb[0] / tdst;
+    const int64_t d01 = mul_tensor->nb[1] / tdst;
+    const int64_t d02 = mul_tensor->nb[2] / tdst;
+    const int64_t d03 = mul_tensor->nb[3] / tdst;
+
+    const size_t ts_mul = ggml_type_size(mul_src->type);
+    GGML_ASSERT(mul_src->nb[0] == ts_mul);
+    const int64_t mul_s01 = mul_src->nb[1] / ts_mul;
+    const int64_t mul_s02 = mul_src->nb[2] / ts_mul;
+    const int64_t mul_s03 = mul_src->nb[3] / ts_mul;
+    const int mul_nrows     = mul_src->ne[1];
+    const int mul_nchannels = mul_src->ne[2];
+    const int mul_nsamples  = mul_src->ne[3];
+
+    rms_norm_mul_f32_sycl(src0_dd, mul_dd, dst_dd, ne00, ne01, ne02, ne03,
+        s00, s01, s02, s03, d00, d01, d02, d03,
+        mul_s01, mul_s02, mul_s03, mul_nrows, mul_nchannels, mul_nsamples, eps, main_stream, ctx.device);
+}
+
 void ggml_sycl_op_rms_norm_back(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
 

ggml/src/ggml-sycl/norm.hpp

@@ -19,6 +19,8 @@ void ggml_sycl_op_norm(ggml_backend_sycl_context& ctx, ggml_tensor* dst);
 
 void ggml_sycl_op_rms_norm(ggml_backend_sycl_context& ctx, ggml_tensor* dst);
 
+void ggml_sycl_op_rms_norm_fused(ggml_backend_sycl_context& ctx, ggml_tensor* dst, ggml_tensor* mul);
+
 void ggml_sycl_op_rms_norm_back(ggml_backend_sycl_context& ctx, ggml_tensor* dst);
 
 void ggml_sycl_op_group_norm(ggml_backend_sycl_context& ctx, ggml_tensor* dst);

ggml/src/ggml-sycl/vecdotq.hpp

@@ -658,6 +658,40 @@ template <> struct reorder_vec_dot_q_sycl<GGML_TYPE_Q6_K> {
 #define VDR_Q4_0_Q8_1_MMVQ 2
 #define VDR_Q4_0_Q8_1_MMQ  4
 
+#define VDR_Q2_0_Q8_1_MMVQ 1
+
+template <int vdr>
+static __dpct_inline__ float vec_dot_q2_0_q8_1_impl(
+    const int * v,
+    const int * u,
+    const float & d2,
+    const sycl::half2 & ds8) {
+    int sumi = 0;
+
+#pragma unroll
+    for (int i = 0; i < vdr; ++i) {
+#pragma unroll
+        for (int j = 0; j < 4; ++j) {
+            const uint8_t q = (uint8_t) ((uint32_t) v[i] >> (8 * j));
+
+            // unpack 2-bit values to byte lanes (0..3), then apply zero-point
+            // correction with ds8f.y() below, mirroring the q4_0 style.
+            int vi = 0;
+            vi |= (((q >> 0) & 0x3) & 0xFF) << 0;
+            vi |= (((q >> 2) & 0x3) & 0xFF) << 8;
+            vi |= (((q >> 4) & 0x3) & 0xFF) << 16;
+            vi |= (((q >> 6) & 0x3) & 0xFF) << 24;
+
+            sumi = dpct::dp4a(vi, u[4 * i + j], sumi);
+        }
+    }
+
+    const sycl::float2 ds8f = ds8.convert<float, sycl::rounding_mode::automatic>();
+    // q2_0 has zero-point 1. Scale ds8f.y() by processed-lane ratio,
+    // consistent with q4_0's explicit zero-point subtraction style.
+    return d2 * (sumi * ds8f.x() - ((float) vdr / (float) QI2_0) * ds8f.y());
+}
+
 template <int vdr>
 static __dpct_inline__ float vec_dot_q4_0_q8_1_impl(const int * v, const int * u, const float & d4,
                                                     const sycl::half2 & ds8) {
@@ -882,6 +916,41 @@ vec_dot_q4_0_q8_1(const void *__restrict__ vbq,
     return vec_dot_q4_0_q8_1_impl<VDR_Q4_0_Q8_1_MMVQ>(v, u, bq4_0->d, bq8_1->ds);
 }
 
+static __dpct_inline__ float
+vec_dot_q2_0_q8_1(const void *__restrict__ vbq,
+                  const block_q8_1 *__restrict__ bq8_1, const int &iqs) {
+
+    const block_q2_0 * bq2_0 = (const block_q2_0 *) vbq;
+
+    int v[2 * VDR_Q2_0_Q8_1_MMVQ];
+    int u[8 * VDR_Q2_0_Q8_1_MMVQ];
+
+#pragma unroll
+    for (int i = 0; i < VDR_Q2_0_Q8_1_MMVQ; ++i) {
+        const int base = 4 * (iqs + i);
+
+        // Q2_0 has QK2_0 = 64 and uses 2 x QK8_1 blocks on the RHS.
+        v[2 * i + 0] = get_int_from_uint8(bq2_0->qs, iqs + i);
+        v[2 * i + 1] = get_int_from_uint8(bq2_0->qs, iqs + i + QI2_0);
+
+        u[8 * i + 0] = get_int_from_int8_aligned(bq8_1[0].qs, base + 0);
+        u[8 * i + 1] = get_int_from_int8_aligned(bq8_1[0].qs, base + 1);
+        u[8 * i + 2] = get_int_from_int8_aligned(bq8_1[0].qs, base + 2);
+        u[8 * i + 3] = get_int_from_int8_aligned(bq8_1[0].qs, base + 3);
+
+        u[8 * i + 4] = get_int_from_int8_aligned(bq8_1[1].qs, base + 0);
+        u[8 * i + 5] = get_int_from_int8_aligned(bq8_1[1].qs, base + 1);
+        u[8 * i + 6] = get_int_from_int8_aligned(bq8_1[1].qs, base + 2);
+        u[8 * i + 7] = get_int_from_int8_aligned(bq8_1[1].qs, base + 3);
+    }
+
+    const float sum0 = vec_dot_q2_0_q8_1_impl<VDR_Q2_0_Q8_1_MMVQ>(
+        v + 0, u + 0, bq2_0->d, bq8_1[0].ds);
+    const float sum1 = vec_dot_q2_0_q8_1_impl<VDR_Q2_0_Q8_1_MMVQ>(
+        v + VDR_Q2_0_Q8_1_MMVQ, u + 4 * VDR_Q2_0_Q8_1_MMVQ, bq2_0->d, bq8_1[1].ds);
+    return sum0 + sum1;
+}
+
 static __dpct_inline__ float
 vec_dot_q4_1_q8_1(const void *__restrict__ vbq,
                   const block_q8_1 *__restrict__ bq8_1, const int &iqs) {

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

@@ -610,6 +610,13 @@ static constexpr std::initializer_list<ggml_op> topk_moe_sigmoid_norm_bias{ GGML
                                                                             GGML_OP_RESHAPE,  GGML_OP_SUM_ROWS, GGML_OP_CLAMP,
                                                                             GGML_OP_DIV,      GGML_OP_RESHAPE };
 
+static constexpr std::initializer_list<ggml_op> topk_moe_sqrt_softplus_norm_bias{ GGML_OP_UNARY,    GGML_OP_SQRT,
+                                                                                  GGML_OP_RESHAPE,  GGML_OP_ADD,
+                                                                                  GGML_OP_ARGSORT,  GGML_OP_VIEW,
+                                                                                  GGML_OP_GET_ROWS, GGML_OP_RESHAPE,
+                                                                                  GGML_OP_SUM_ROWS, GGML_OP_CLAMP,
+                                                                                  GGML_OP_DIV,      GGML_OP_RESHAPE };
+
 static constexpr std::initializer_list<ggml_op> topk_moe_early_softmax     { GGML_OP_SOFT_MAX, GGML_OP_RESHAPE,  GGML_OP_ARGSORT,
                                                                              GGML_OP_VIEW,     GGML_OP_GET_ROWS };
 
@@ -673,6 +680,22 @@ static constexpr std::initializer_list<std::array<int, 3>> topk_moe_sigmoid_norm
     {10, 0, 9 }, // reshape->src[0]  == div
 };
 
+static constexpr std::initializer_list<std::array<int, 3>> topk_moe_sqrt_softplus_norm_bias_edges {
+    { 1, 0, 0 }, // sqrt->src[0]     == softplus
+    { 2, 0, 1 }, // reshape->src[0]  == sqrt
+    { 3, 0, 1 }, // add->src[0]      == sqrt
+    { 4, 0, 3 }, // argsort->src[0]  == add
+    { 5, 0, 4 }, // view->src[0]     == argsort
+    { 6, 0, 2 }, // get_rows->src[0] == reshape
+    { 6, 1, 5 }, // get_rows->src[1] == view
+    { 7, 0, 6 }, // reshape->src[0]  == get_rows
+    { 8, 0, 7 }, // sum_rows->src[0] == reshape
+    { 9, 0, 8 }, // clamp->src[0]    == sum_rows
+    {10, 0, 7 }, // div->src[0]      == reshape
+    {10, 1, 9 }, // div->src[1]      == clamp
+    {11, 0,10 }, // reshape->src[0]  == div
+};
+
 // same as early_softmax_norm but ending after the get_rows
 static constexpr std::initializer_list<std::array<int, 3>> topk_moe_early_softmax_edges {
     { 1, 0, 0 }, // reshape->src[0]  == softmax
@@ -701,6 +724,7 @@ enum topk_moe_mode {
     TOPK_MOE_EARLY_SOFTMAX_NORM,
     TOPK_MOE_LATE_SOFTMAX,
     TOPK_MOE_SIGMOID_NORM_BIAS,
+    TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS,
     TOPK_MOE_COUNT,
 };
 
@@ -998,6 +1022,7 @@ struct vk_device_struct {
     vk_pipeline pipeline_snake_f32;
     vk_pipeline pipeline_snake_f16;
     vk_pipeline pipeline_snake_bf16;
+    vk_pipeline pipeline_pool1d_f32;
     vk_pipeline pipeline_pool2d_f32;
     vk_pipeline pipeline_rwkv_wkv6_f32;
     vk_pipeline pipeline_rwkv_wkv7_f32;
@@ -1687,6 +1712,17 @@ struct vk_op_snake_push_constants {
     uint32_t ne1;
 };
 
+struct vk_op_pool1d_push_constants {
+    uint32_t IL;
+    uint32_t OL;
+    uint32_t OC;
+    uint32_t pelements;
+    uint32_t op;
+    int32_t k0;
+    int32_t s0;
+    int32_t p0;
+};
+
 struct vk_op_pool2d_push_constants {
     uint32_t IW; uint32_t IH;
     uint32_t OW; uint32_t OH;
@@ -1934,6 +1970,7 @@ struct ggml_vk_garbage_collector {
 static void ggml_vk_preallocate_buffers(ggml_backend_vk_context * ctx, vk_context subctx);
 static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested = nullptr);
 static void ggml_pipeline_allocate_descriptor_sets(ggml_backend_vk_context * ctx);
+static bool ggml_vk_intel_windows_driver_equals_or_newer_than(uint32_t driver_version, uint32_t threshold_major, uint32_t threshold_minor);
 
 static bool vk_memory_logger_enabled = false;
 
@@ -5552,8 +5589,11 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     ggml_vk_create_pipeline(device, device->pipeline_argmax_f32, "argmax_f32", argmax_f32_len, argmax_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
 
     ggml_vk_create_pipeline(device, device->pipeline_sum_rows_f32, "sum_rows_f32", sum_rows_f32_len, sum_rows_f32_data, "main", 2, sizeof(vk_op_sum_rows_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
-    // Intel Arc B390 was observed segfaulting with this shader.
-    if (device->subgroup_basic && device->subgroup_shuffle && device->vendor_id != VK_VENDOR_ID_INTEL) {
+    // Intel Windows driver older than 32.0.101.8860 will crash when using fwht kernels on Xe2+ GPUS so we gate that here
+    const bool can_use_fwht = device->driver_id != vk::DriverId::eIntelProprietaryWindows ||
+        device->architecture != vk_device_architecture::INTEL_XE2 ||
+        (device->architecture == vk_device_architecture::INTEL_XE2 && ggml_vk_intel_windows_driver_equals_or_newer_than(device->properties.driverVersion, 101, 8860));
+    if (can_use_fwht && device->subgroup_basic && device->subgroup_shuffle) {
         int idx = 0;
         for (uint32_t n : {64, 128, 256, 512}) {
             if (device->subgroup_size <= n) {
@@ -5561,8 +5601,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
             }
             ++idx;
         }
-    } else if (device->driver_id != vk::DriverId::eIntelProprietaryWindows) {
-        // Disabled on Intel Windows due to a driver bug: https://github.com/ggml-org/llama.cpp/pull/23964#issuecomment-4598226147
+    } else if (can_use_fwht) {
         int idx = 0;
         for (uint32_t n : {64, 128, 256, 512}) {
             const uint32_t block_size = std::min(device->subgroup_size, n);
@@ -5619,6 +5658,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     ggml_vk_create_pipeline(device, device->pipeline_snake_f16,  "snake_f16",  snake_f16_len,  snake_f16_data,  "main", 4, sizeof(vk_op_snake_push_constants), {256, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_snake_bf16, "snake_bf16", snake_bf16_len, snake_bf16_data, "main", 4, sizeof(vk_op_snake_push_constants), {256, 1, 1}, {}, 1);
 
+    ggml_vk_create_pipeline(device, device->pipeline_pool1d_f32, "pool1d_f32", pool1d_f32_len, pool1d_f32_data, "main", 2, sizeof(vk_op_pool1d_push_constants), {512, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_pool2d_f32, "pool2d_f32", pool2d_f32_len, pool2d_f32_data, "main", 2, sizeof(vk_op_pool2d_push_constants), {512, 1, 1}, {}, 1);
 
     ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv6_f32, "rwkv_wkv6_f32", rwkv_wkv6_f32_len, rwkv_wkv6_f32_data, "main", 7, sizeof(vk_op_rwkv_wkv6_push_constants), {1, 1, 1}, {device->subgroup_size}, 1);
@@ -11332,6 +11372,11 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
             case GGML_TYPE_BF16: return ctx->device->pipeline_col2im_1d_bf16;
             default:             return nullptr;
         }
+    case GGML_OP_POOL_1D:
+        if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
+            return ctx->device->pipeline_pool1d_f32;
+        }
+        return nullptr;
     case GGML_OP_POOL_2D:
         if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
             return ctx->device->pipeline_pool2d_f32;
@@ -11854,6 +11899,13 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
         {
             elements = { uint32_t(dst->ne[0]), uint32_t(dst->ne[1]), 1 };
         } break;
+    case GGML_OP_POOL_1D:
+        {
+            const uint32_t N = dst->ne[3] * dst->ne[2];
+            const uint32_t OC = dst->ne[1];
+            const uint32_t OL = dst->ne[0];
+            elements = { N * OC * OL, 1, 1};
+        } break;
     case GGML_OP_POOL_2D:
         {
             const uint32_t N = dst->ne[3];
@@ -13175,12 +13227,16 @@ static void ggml_vk_soft_max_back(ggml_backend_vk_context * ctx, vk_context& sub
 
 static void ggml_vk_topk_moe(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_cgraph * cgraph, int node_idx) {
     topk_moe_mode mode = ctx->fused_topk_moe_mode;
+    const bool has_bias = mode == TOPK_MOE_SIGMOID_NORM_BIAS || mode == TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS;
     ggml_tensor * logits = cgraph->nodes[node_idx + 0]->src[0];
-    ggml_tensor * bias = (mode == TOPK_MOE_SIGMOID_NORM_BIAS) ? cgraph->nodes[node_idx + 2]->src[1] : logits;
+    ggml_tensor * bias = mode == TOPK_MOE_SIGMOID_NORM_BIAS       ? cgraph->nodes[node_idx + 2]->src[1] :
+                         mode == TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS ? cgraph->nodes[node_idx + 3]->src[1] :
+                                                                   logits;
     ggml_tensor * weights = cgraph->nodes[node_idx + ctx->num_additional_fused_ops];
-    ggml_tensor * ids = (mode == TOPK_MOE_SIGMOID_NORM_BIAS) ? cgraph->nodes[node_idx + 4] :
-                        (mode == TOPK_MOE_LATE_SOFTMAX) ?      cgraph->nodes[node_idx + 1] :
-                                                               cgraph->nodes[node_idx + 3];
+    ggml_tensor * ids = mode == TOPK_MOE_SIGMOID_NORM_BIAS       ? cgraph->nodes[node_idx + 4] :
+                        mode == TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS ? cgraph->nodes[node_idx + 5] :
+                        mode == TOPK_MOE_LATE_SOFTMAX             ? cgraph->nodes[node_idx + 1] :
+                                                                   cgraph->nodes[node_idx + 3];
 
     GGML_ASSERT(logits->type == GGML_TYPE_F32);
     GGML_ASSERT(bias->type == GGML_TYPE_F32);
@@ -13220,16 +13276,24 @@ static void ggml_vk_topk_moe(ggml_backend_vk_context * ctx, vk_context& subctx,
         pc.clamp_min = ggml_get_op_params_f32(clamp, 0);
         pc.clamp_max = ggml_get_op_params_f32(clamp, 1);
     }
+    if (mode == TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS) {
+        ggml_tensor * clamp = cgraph->nodes[node_idx + 9];
+        GGML_ASSERT(clamp->op == GGML_OP_CLAMP);
+        pc.clamp_min = ggml_get_op_params_f32(clamp, 0);
+        pc.clamp_max = ggml_get_op_params_f32(clamp, 1);
+    }
 
 #define GATING_FUNC_SOFTMAX 0
 #define GATING_FUNC_SIGMOID 1
 #define GATING_FUNC_SOFTMAX_WEIGHT 2
+#define GATING_FUNC_SQRT_SOFTPLUS 3
 
-    pc.gating_func = mode == TOPK_MOE_SIGMOID_NORM_BIAS ? GATING_FUNC_SIGMOID :
-                     mode == TOPK_MOE_LATE_SOFTMAX ?      GATING_FUNC_SOFTMAX_WEIGHT :
-                                                          GATING_FUNC_SOFTMAX;
-    pc.has_bias = mode == TOPK_MOE_SIGMOID_NORM_BIAS;
-    pc.with_norm = mode == TOPK_MOE_EARLY_SOFTMAX_NORM || mode == TOPK_MOE_SIGMOID_NORM_BIAS;
+    pc.gating_func = mode == TOPK_MOE_SIGMOID_NORM_BIAS       ? GATING_FUNC_SIGMOID :
+                     mode == TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS ? GATING_FUNC_SQRT_SOFTPLUS :
+                     mode == TOPK_MOE_LATE_SOFTMAX             ? GATING_FUNC_SOFTMAX_WEIGHT :
+                                                                 GATING_FUNC_SOFTMAX;
+    pc.has_bias = has_bias;
+    pc.with_norm = mode == TOPK_MOE_EARLY_SOFTMAX_NORM || has_bias;
     if (ctx->fused_topk_moe_scale) {
         GGML_ASSERT(weights->op == GGML_OP_SCALE);
         pc.output_scale = ggml_get_op_params_f32(weights, 0);
@@ -13773,6 +13837,29 @@ static void ggml_vk_snake_dispatch_fused(ggml_backend_vk_context * ctx, vk_conte
     ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { x_buf, a_buf, inv_b_buf, dst_buf }, pc, elements);
 }
 
+static void ggml_vk_pool_1d(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) {
+    uint32_t op = static_cast<uint32_t>(dst->op_params[0]);
+    const int32_t k0 = dst->op_params[1];
+    const int32_t s0 = dst->op_params[2];
+    const int32_t p0 = dst->op_params[3];
+
+    const uint32_t IL = src0->ne[0];
+
+    const uint32_t N = dst->ne[3] * dst->ne[2];
+
+    const uint32_t OC = dst->ne[1];
+    const uint32_t OL = dst->ne[0];
+
+    const uint32_t parallel_elements = N * OC * OL;
+
+    ggml_vk_op_f32<vk_op_pool1d_push_constants>(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_POOL_1D, {
+        IL, OL, OC,
+        parallel_elements,
+        op,
+        k0, s0, p0,
+    });
+}
+
 static void ggml_vk_pool_2d(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) {
     uint32_t op = static_cast<uint32_t>(dst->op_params[0]);
     const int32_t k1 = dst->op_params[1];
@@ -15284,6 +15371,10 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
     case GGML_OP_CONV_TRANSPOSE_1D:
         ggml_vk_conv_transpose_1d(ctx, compute_ctx, src0, src1, node);
 
+        break;
+    case GGML_OP_POOL_1D:
+        ggml_vk_pool_1d(ctx, compute_ctx, src0, node);
+
         break;
     case GGML_OP_POOL_2D:
         ggml_vk_pool_2d(ctx, compute_ctx, src0, node);
@@ -16311,6 +16402,20 @@ static bool ggml_vk_can_fuse_topk_moe(ggml_backend_vk_context * ctx, const struc
             return false;
         }
         break;
+    case TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS:
+        softmax = cgraph->nodes[node_idx + 0]; // really softplus
+        weights = cgraph->nodes[node_idx + 11];
+        get_rows = cgraph->nodes[node_idx + 6];
+        argsort = cgraph->nodes[node_idx + 4];
+        if (ggml_get_unary_op(softmax) != GGML_UNARY_OP_SOFTPLUS) {
+            return false;
+        }
+        // bias is expected to be 1D
+        if (ggml_nrows(cgraph->nodes[node_idx + 3]->src[1]) != 1 ||
+            !ggml_is_contiguous(cgraph->nodes[node_idx + 3]->src[1])) {
+            return false;
+        }
+        break;
     case TOPK_MOE_EARLY_SOFTMAX:
         softmax = cgraph->nodes[node_idx + 0];
         weights = cgraph->nodes[node_idx + 4];
@@ -16334,7 +16439,9 @@ static bool ggml_vk_can_fuse_topk_moe(ggml_backend_vk_context * ctx, const struc
     probs = probs->src[0];
     ggml_tensor * selection_probs = argsort->src[0];
 
-    if (probs != selection_probs && mode != TOPK_MOE_SIGMOID_NORM_BIAS) {
+    if (probs != selection_probs &&
+        mode != TOPK_MOE_SIGMOID_NORM_BIAS &&
+        mode != TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS) {
         return false;
     }
 
@@ -16702,7 +16809,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
         // the fused result in an elementwise-way. This affects whether the memory for
         // the src is allowed to overlap the memory for the destination.
         // The array is sized to handle the largest fusion (asserted later).
-        bool op_srcs_fused_elementwise[12];
+        bool op_srcs_fused_elementwise[13];
 
         ctx->fused_topk_moe_mode = TOPK_MOE_COUNT;
         ctx->fused_topk_moe_scale = false;
@@ -16813,6 +16920,15 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
                 ctx->fused_topk_moe_mode = TOPK_MOE_SIGMOID_NORM_BIAS;
                 fusion_string = "TOPK_MOE_SIGMOID_NORM_BIAS";
                 std::fill_n(op_srcs_fused_elementwise, ctx->num_additional_fused_ops + 1, false);
+            } else if (ggml_can_fuse_subgraph(cgraph, i, topk_moe_sqrt_softplus_norm_bias, { i + 5, i + 11 }) &&
+                       ggml_check_edges(cgraph, i, topk_moe_sqrt_softplus_norm_bias_edges) &&
+                       ggml_vk_can_fuse_topk_moe(ctx, cgraph, i, TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS)) {
+                ctx->num_additional_fused_ops = topk_moe_sqrt_softplus_norm_bias.size() - 1;
+                // view of argsort writes to memory
+                ctx->fused_ops_write_mask |= 1 << 5;
+                ctx->fused_topk_moe_mode = TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS;
+                fusion_string = "TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS";
+                std::fill_n(op_srcs_fused_elementwise, ctx->num_additional_fused_ops + 1, false);
             } else if (ggml_can_fuse_subgraph(cgraph, i, topk_moe_early_softmax, { i + 3, i + 4 }) &&
                        ggml_check_edges(cgraph, i, topk_moe_early_softmax_edges) &&
                        ggml_vk_can_fuse_topk_moe(ctx, cgraph, i, TOPK_MOE_EARLY_SOFTMAX)) {
@@ -17079,6 +17195,9 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph *
         if (keep_pattern(topk_moe_sigmoid_norm_bias)) {
             continue;
         }
+        if (keep_pattern(topk_moe_sqrt_softplus_norm_bias)) {
+            continue;
+        }
         if (keep_pattern(topk_moe_early_softmax)) {
             continue;
         }
@@ -17109,6 +17228,7 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph *
             // Don't pull forward nodes from fusion patterns
             if (match_pattern(topk_moe_early_softmax_norm, j) ||
                 match_pattern(topk_moe_sigmoid_norm_bias, j) ||
+                match_pattern(topk_moe_sqrt_softplus_norm_bias, j) ||
                 match_pattern(topk_moe_early_softmax, j) ||
                 match_pattern(topk_moe_late_softmax, j) ||
                 match_pattern(snake_pattern, j)) {
@@ -18001,6 +18121,8 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
         case GGML_OP_CONV_2D_DW:
             return (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16)
                 && op->src[1]->type == GGML_TYPE_F32;
+        case GGML_OP_POOL_1D:
+            return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_POOL_2D:
             return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_RWKV_WKV6:
@@ -18466,6 +18588,22 @@ static uint32_t ggml_vk_intel_shader_core_count(const vk::PhysicalDevice& vkdev)
     }
 }
 
+static bool ggml_vk_intel_windows_driver_equals_or_newer_than(uint32_t driver_version, uint32_t threshold_major, uint32_t threshold_minor) {
+#if defined(_WIN32)
+    // Intel Windows encodes xxx.yyyy as [31:14].[13:0].
+    const uint32_t major = driver_version >> 14;
+    const uint32_t minor = driver_version & 0x3fff;
+
+    return major > threshold_major || (major == threshold_major && minor >= threshold_minor);
+#else
+    GGML_UNUSED(driver_version);
+    GGML_UNUSED(threshold_major);
+    GGML_UNUSED(threshold_minor);
+    return true;
+#endif
+}
+
+
 // checks
 
 #ifdef GGML_VULKAN_CHECK_RESULTS
@@ -18920,6 +19058,13 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
             const int32_t oc     = tensor->op_params[1];
             const int32_t p0     = tensor->op_params[2];
             tensor_clone = ggml_col2im_1d(ggml_ctx, src_clone[0], stride, oc, p0);
+        } else if (tensor->op == GGML_OP_POOL_1D) {
+            enum ggml_op_pool op = static_cast<ggml_op_pool>(tensor->op_params[0]);
+            const int32_t k0 = tensor->op_params[1];
+            const int32_t s0 = tensor->op_params[2];
+            const int32_t p0 = tensor->op_params[3];
+
+            tensor_clone = ggml_pool_1d(ggml_ctx, src_clone[0], op, k0, s0, p0);
         } else if (tensor->op == GGML_OP_POOL_2D) {
             enum ggml_op_pool op = static_cast<ggml_op_pool>(tensor->op_params[0]);
             const int32_t k0 = tensor->op_params[1];

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

@@ -0,0 +1,65 @@
+#version 450
+
+#include "types.glsl"
+
+#extension GL_EXT_shader_16bit_storage : require
+
+layout(push_constant) uniform parameter {
+    uint IL;
+    uint OL;
+    uint OC;
+    uint pelements;
+    uint op;
+    int k0;
+    int s0;
+    int p0;
+} p;
+
+#define BLOCK_SIZE 512
+#define FLT_MAX 3.402823466e+38F
+#define OP_POOL_MAX 0u
+#define OP_POOL_AVG 1u
+
+layout (local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;
+
+layout(binding = 0) readonly buffer X {A_TYPE data_a[];};
+layout(binding = 1) writeonly buffer D {D_TYPE data_d[];};
+
+void main() {
+    const uint idx = gl_GlobalInvocationID.x;
+    if (idx >= p.pelements) {
+        return;
+    }
+
+    const uint nc = idx / p.OL;
+    const uint cur_ol = idx % p.OL;
+
+    const int start = int(cur_ol) * p.s0 - p.p0;
+    const int bl = max(start, 0);
+    const int el = min(max(start + p.k0, 0), int(p.IL));
+
+    const int window_size = el - bl;
+    const float scale = window_size > 0 ? 1.0 / float(window_size) : 0.0;
+    float res;
+
+    if (p.op == OP_POOL_AVG) {
+        res = 0.0;
+    } else if (p.op == OP_POOL_MAX) {
+        res = -FLT_MAX;
+    } else {
+        return;
+    }
+
+    #pragma unroll
+    for (uint i = bl; i < el; i++) {
+        const float cur = D_TYPE(data_a[nc * p.IL + i]);
+
+        if (p.op == OP_POOL_AVG) {
+            res += cur * scale;
+        } else if (p.op == OP_POOL_MAX) {
+            res = max(res, cur);
+        }
+    }
+
+    data_d[nc * p.OL + cur_ol] = res;
+}

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

@@ -10,6 +10,7 @@
 #define GATING_FUNC_SOFTMAX 0
 #define GATING_FUNC_SIGMOID 1
 #define GATING_FUNC_SOFTMAX_WEIGHT 2
+#define GATING_FUNC_SQRT_SOFTPLUS 3
 
 layout (push_constant) uniform parameter
 {
@@ -120,6 +121,13 @@ void main() {
             const uint expert = i + lane;
             probs[i / WARP_SIZE] = (n_experts % WARP_SIZE == 0 || expert < n_experts) ? 1.f / (1.f + exp(-probs[i / WARP_SIZE])) : -INFINITY;
         }
+    } else if (gating_func == GATING_FUNC_SQRT_SOFTPLUS) {
+        [[unroll]]
+        for (uint i = 0; i < n_experts; i += WARP_SIZE) {
+            const uint expert = i + lane;
+            const float val = probs[i / WARP_SIZE];
+            probs[i / WARP_SIZE] = (n_experts % WARP_SIZE == 0 || expert < n_experts) ? sqrt(val > 20.0f ? val : log(1.0f + exp(val))) : -INFINITY;
+        }
     }
 
     float selection_probs[experts_per_thread];

ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp

@@ -1052,6 +1052,7 @@ void process_shaders() {
     string_to_spv("snake_f16",  "snake.comp", {{"DATA_A_F16", "1"},  {"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}});
     string_to_spv("snake_bf16", "snake.comp", {{"DATA_A_BF16", "1"}, {"DATA_D_BF16", "1"}, {"A_TYPE", "uint16_t"},  {"D_TYPE", "uint16_t"}});
 
+    string_to_spv("pool1d_f32", "pool1d.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
     string_to_spv("pool2d_f32", "pool2d.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
 
     string_to_spv("rwkv_wkv6_f32", "wkv6.comp", merge_maps(base_dict, {{"A_TYPE", "float"}}));

ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp

@@ -591,7 +591,8 @@ struct ggml_webgpu_flash_attn_common_pipeline_key {
     ggml_type dst_type;
     uint32_t  head_dim_qk;
     uint32_t  head_dim_v;
-    bool      kv_direct;
+    bool      k_direct;
+    bool      v_direct;
     bool      kv_overlap;
     bool      has_mask;
     bool      has_sinks;
@@ -600,8 +601,9 @@ struct ggml_webgpu_flash_attn_common_pipeline_key {
     bool operator==(const ggml_webgpu_flash_attn_common_pipeline_key & other) const {
         return q_type == other.q_type && k_type == other.k_type && v_type == other.v_type &&
                dst_type == other.dst_type && head_dim_qk == other.head_dim_qk && head_dim_v == other.head_dim_v &&
-               kv_direct == other.kv_direct && kv_overlap == other.kv_overlap && has_mask == other.has_mask &&
-               has_sinks == other.has_sinks && uses_logit_softcap == other.uses_logit_softcap;
+               k_direct == other.k_direct && v_direct == other.v_direct && kv_overlap == other.kv_overlap &&
+               has_mask == other.has_mask && has_sinks == other.has_sinks &&
+               uses_logit_softcap == other.uses_logit_softcap;
     }
 };
 
@@ -613,7 +615,8 @@ inline void ggml_webgpu_flash_attn_hash_common_pipeline_key(size_t &
     ggml_webgpu_hash_combine(seed, key.dst_type);
     ggml_webgpu_hash_combine(seed, key.head_dim_qk);
     ggml_webgpu_hash_combine(seed, key.head_dim_v);
-    ggml_webgpu_hash_combine(seed, key.kv_direct);
+    ggml_webgpu_hash_combine(seed, key.k_direct);
+    ggml_webgpu_hash_combine(seed, key.v_direct);
     ggml_webgpu_hash_combine(seed, key.kv_overlap);
     ggml_webgpu_hash_combine(seed, key.has_mask);
     ggml_webgpu_hash_combine(seed, key.has_sinks);
@@ -687,12 +690,13 @@ inline bool ggml_webgpu_flash_attn_float_vec4_aligned(const ggml_tensor * K,
            ggml_webgpu_flash_attn_float_vec4_aligned(V, storage_offset_alignment);
 }
 
-inline bool ggml_webgpu_flash_attn_kv_direct(const ggml_tensor * Q,
-                                             const ggml_tensor * K,
-                                             const ggml_tensor * V,
-                                             uint32_t            kv_direct_align) {
-    return K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 && (Q->ne[0] % kv_direct_align == 0) &&
-           (K->ne[1] % GGML_WEBGPU_KV_SEQ_PAD == 0);
+inline bool ggml_webgpu_flash_attn_k_direct(const ggml_tensor * Q, const ggml_tensor * K, uint32_t kv_direct_align) {
+    return (K->type == GGML_TYPE_F16 || K->type == GGML_TYPE_Q8_0 || K->type == GGML_TYPE_Q4_0) &&
+           (Q->ne[0] % kv_direct_align == 0) && (K->ne[1] % GGML_WEBGPU_KV_SEQ_PAD == 0);
+}
+
+inline bool ggml_webgpu_flash_attn_v_direct(const ggml_tensor * Q, const ggml_tensor * V, uint32_t kv_direct_align) {
+    return ggml_webgpu_flash_attn_k_direct(Q, V, kv_direct_align);
 }
 
 inline ggml_webgpu_flash_attn_common_pipeline_key ggml_webgpu_flash_attn_make_common_pipeline_key(
@@ -706,10 +710,11 @@ inline ggml_webgpu_flash_attn_common_pipeline_key ggml_webgpu_flash_attn_make_co
     key.dst_type                                   = context.dst->type;
     key.head_dim_qk                                = (uint32_t) context.src0->ne[0];
     key.head_dim_v                                 = (uint32_t) context.src2->ne[0];
-    key.kv_direct  = ggml_webgpu_flash_attn_kv_direct(context.src0, context.src1, context.src2, kv_direct_align);
-    key.kv_overlap = kv_overlap;
-    key.has_mask   = context.src3 != nullptr;
-    key.has_sinks  = context.src4 != nullptr;
+    key.k_direct           = ggml_webgpu_flash_attn_k_direct(context.src0, context.src1, kv_direct_align);
+    key.v_direct           = ggml_webgpu_flash_attn_v_direct(context.src0, context.src2, kv_direct_align);
+    key.kv_overlap         = kv_overlap;
+    key.has_mask           = context.src3 != nullptr;
+    key.has_sinks          = context.src4 != nullptr;
     key.uses_logit_softcap = ggml_get_op_params_f32(context.dst, 2) != 0.0f;
     return key;
 }
@@ -794,9 +799,13 @@ inline std::vector<std::string> ggml_webgpu_flash_attn_common_defines(
         defines.push_back("LOGIT_SOFTCAP");
         variant += "_lgsc";
     }
-    if (key.kv_direct) {
-        defines.push_back("KV_DIRECT");
-        variant += "_kvdirect";
+    if (key.k_direct) {
+        defines.push_back("K_DIRECT");
+        variant += "_k_direct";
+    }
+    if (key.v_direct) {
+        defines.push_back("V_DIRECT");
+        variant += "_v_direct";
     }
     if (key.kv_overlap) {
         defines.push_back("KV_OVERLAP");
@@ -815,6 +824,12 @@ inline std::vector<std::string> ggml_webgpu_flash_attn_common_defines(
 
     if (ggml_is_quantized(key.k_type) || ggml_is_quantized(key.v_type)) {
         defines.push_back("U32_DEQUANT_HELPERS");
+        if (ggml_is_quantized(key.k_type)) {
+            defines.push_back("LOADERS_QUANTIZED_K");
+        }
+        if (ggml_is_quantized(key.v_type)) {
+            defines.push_back("LOADERS_QUANTIZED_V");
+        }
     }
 
     return defines;
@@ -2759,6 +2774,10 @@ class ggml_webgpu_shader_lib {
                 defines.push_back("TYPE_F32");
                 variant += "_f32";
                 break;
+            case GGML_TYPE_F16:
+                defines.push_back("TYPE_F16");
+                variant += "_f16";
+                break;
             case GGML_TYPE_I32:
                 defines.push_back("TYPE_I32");
                 variant += "_i32";
@@ -2792,12 +2811,14 @@ class ggml_webgpu_shader_lib {
         ggml_webgpu_flash_attn_pipeline_key key = {};
         key.common                              = ggml_webgpu_flash_attn_make_common_pipeline_key(
             context, decisions.use_sg_matrix ? context.sg_mat_k : 1u, kv_overlap);
-        key.common.kv_direct = decisions.use_sg_matrix && key.common.kv_direct;
-        key.use_sg_matrix    = decisions.use_sg_matrix;
+        key.common.k_direct &= decisions.use_sg_matrix && key.common.k_type == GGML_TYPE_F16;
+        key.common.v_direct &= decisions.use_sg_matrix && key.common.v_type == GGML_TYPE_F16;
+        key.use_sg_matrix = decisions.use_sg_matrix;
 
         const uint32_t max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile(
             context.wg_mem_limit_bytes, decisions.q_tile, decisions.use_sg_matrix ? context.sg_mat_n : 1u,
-            key.common.head_dim_qk, key.common.head_dim_v, key.common.has_mask, key.common.kv_direct);
+            key.common.head_dim_qk, key.common.head_dim_v, key.common.has_mask,
+            key.common.k_direct || key.common.v_direct);
         GGML_ASSERT(max_kv_tile > 0);
 
         decisions.kv_tile = decisions.use_sg_matrix ?
@@ -2809,7 +2830,7 @@ class ggml_webgpu_shader_lib {
                 std::min(context.max_wg_size, std::max(GGML_WEBGPU_FLASH_ATTN_PREFERRED_WG_SIZE,
                                                        GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE * context.max_subgroup_size));
 
-        if (key.common.kv_direct) {
+        if (key.common.k_direct || key.common.v_direct) {
             decisions.kv_tile = std::min(decisions.kv_tile, GGML_WEBGPU_KV_SEQ_PAD);
             while (GGML_WEBGPU_KV_SEQ_PAD % decisions.kv_tile != 0) {
                 decisions.kv_tile -= decisions.use_sg_matrix ? context.sg_mat_n : context.min_subgroup_size;
@@ -2856,9 +2877,9 @@ class ggml_webgpu_shader_lib {
         }
 
         ggml_webgpu_flash_attn_vec_decisions decisions = {};
-        decisions.kv_tile =
-            ggml_webgpu_flash_attn_get_vec_kv_tile(context.wg_mem_limit_bytes, key.common.head_dim_qk,
-                                                   key.common.head_dim_v, key.common.has_mask, key.common.kv_direct);
+        decisions.kv_tile = ggml_webgpu_flash_attn_get_vec_kv_tile(context.wg_mem_limit_bytes, key.common.head_dim_qk,
+                                                                   key.common.head_dim_v, key.common.has_mask,
+                                                                   key.common.k_direct || key.common.v_direct);
         decisions.wg_size = context.max_subgroup_size;
 
         std::string              variant = "flash_attn_vec";
@@ -2870,12 +2891,10 @@ class ggml_webgpu_shader_lib {
             variant += "_mask_blk";
         }
 
-        uint32_t d_split = context.min_subgroup_size;
-        if (key.common.k_type == GGML_TYPE_F16 && key.common.v_type == GGML_TYPE_F16) {
-            const uint32_t D     = key.common.head_dim_qk | key.common.head_dim_v;
-            const uint32_t D_lsb = D & (~(D - 1u));
-            d_split              = std::min(std::min(context.min_subgroup_size, 4u), std::max(D_lsb / 4u, 1u));
-        }
+        uint32_t       d_split = context.min_subgroup_size;
+        const uint32_t D       = key.common.head_dim_qk | key.common.head_dim_v;
+        const uint32_t D_lsb   = D & (~(D - 1u));
+        d_split                = std::min(std::min(context.min_subgroup_size, 4u), std::max(D_lsb / 4u, 1u));
 
         defines.push_back(std::string("D_SPLIT=") + std::to_string(d_split));
         variant += "_dsplit" + std::to_string(d_split);

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

@@ -3839,7 +3839,8 @@ static size_t ggml_backend_webgpu_buffer_type_get_alloc_size(ggml_backend_buffer
                 const auto &        capabilities = ctx->webgpu_global_ctx->capabilities;
                 if (ggml_webgpu_flash_attn_use_vec_path(ctx->webgpu_global_ctx, Q, K, V)) {
                     const bool kv_direct =
-                        ggml_webgpu_flash_attn_kv_direct(Q, K, V, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH);
+                        ggml_webgpu_flash_attn_k_direct(Q, K, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH) ||
+                        ggml_webgpu_flash_attn_v_direct(Q, V, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH);
                     const uint32_t kv_tile = ggml_webgpu_flash_attn_get_vec_kv_tile(
                         capabilities.limits.maxComputeWorkgroupStorageSize, (uint32_t) Q->ne[0], (uint32_t) V->ne[0],
                         mask != nullptr, kv_direct);
@@ -4289,7 +4290,8 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
             supports_op = (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_I32);
             break;
         case GGML_OP_REPEAT:
-            supports_op = (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_I32 || src0->type == GGML_TYPE_I16);
+            supports_op = (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_I32 ||
+                           src0->type == GGML_TYPE_I16);
             break;
         case GGML_OP_CPY:
         case GGML_OP_CONT:
@@ -4448,9 +4450,10 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
                 const uint32_t q_tile =
                     use_subgroup_matrix ? capabilities.sg_mat_m : GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE;
                 const uint32_t kv_granularity = use_subgroup_matrix ? capabilities.sg_mat_n : 1u;
-                const bool kv_direct = use_subgroup_matrix ?
-                                           ggml_webgpu_flash_attn_kv_direct(src0, src1, src2, capabilities.sg_mat_k) :
-                                           false;
+                const bool     kv_direct = use_subgroup_matrix ?
+                                               ggml_webgpu_flash_attn_k_direct(src0, src1, capabilities.sg_mat_k) ||
+                                                   ggml_webgpu_flash_attn_v_direct(src0, src2, capabilities.sg_mat_k) :
+                                               false;
                 const uint32_t max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile(
                     capabilities.limits.maxComputeWorkgroupStorageSize, q_tile, kv_granularity, (uint32_t) src0->ne[0],
                     (uint32_t) src2->ne[0], op->src[3] != nullptr, kv_direct);

ggml/src/ggml-webgpu/wgsl-shaders/common_decls.tmpl

@@ -9,6 +9,12 @@ fn get_byte_i32(value: u32, index: u32) -> i32 {
 #endif
 
 #ifdef U32_DEQUANT_HELPERS
+
+fn f16_from_u16(bits: u32) -> f16 {
+    let packed = unpack2x16float(bits);
+    return f16(packed[0]);
+}
+
 #ifdef DECLARE_BYTE_LOADERS_SRC
 fn load_u16_at_src(byte_offset: u32) -> u32 {
     let word = src[byte_offset / 4u];
@@ -36,7 +42,7 @@ fn load_f16_as_f32_at_src(byte_offset: u32) -> f32 {
     let d_bits = (word >> shift) & 0xFFFFu;
     return unpack2x16float(d_bits)[0];
 }
-#endif
+#endif // DECLARE_BYTE_LOADERS_SRC
 
 #ifdef DECLARE_BYTE_LOADERS_SRC0
 fn load_u16_at_src0(byte_offset: u32) -> u32 {
@@ -72,8 +78,47 @@ fn load_f16_as_f32_at_src0(byte_offset: u32) -> f32 {
     let d_bits = (word >> shift) & 0xFFFFu;
     return unpack2x16float(d_bits)[0];
 }
-#endif
-#endif
+#endif // DECLARE_BYTE_LOADERS_SRC0
+
+#ifdef LOADERS_QUANTIZED_K
+fn load_k_u16_at(byte_offset: u32) -> u32 {
+    let word = K[byte_offset / 4u];
+    let shift = (byte_offset & 2u) * 8u;
+    return (word >> shift) & 0xFFFFu;
+}
+
+fn load_k_u32_at(byte_offset: u32) -> u32 {
+    let word_idx = byte_offset / 4u;
+    let shift = (byte_offset & 3u) * 8u;
+    let lo = K[word_idx];
+    if (shift == 0u) {
+        return lo;
+    }
+    let hi = K[word_idx + 1u];
+    return (lo >> shift) | (hi << (32u - shift));
+}
+#endif // LOADERS_QUANTIZED_K
+
+#ifdef LOADERS_QUANTIZED_V
+fn load_v_u16_at(byte_offset: u32) -> u32 {
+    let word = V[byte_offset / 4u];
+    let shift = (byte_offset & 2u) * 8u;
+    return (word >> shift) & 0xFFFFu;
+}
+
+fn load_v_u32_at(byte_offset: u32) -> u32 {
+    let word_idx = byte_offset / 4u;
+    let shift = (byte_offset & 3u) * 8u;
+    let lo = V[word_idx];
+    if (shift == 0u) {
+        return lo;
+    }
+    let hi = V[word_idx + 1u];
+    return (lo >> shift) | (hi << (32u - shift));
+}
+#endif // LOADERS_QUANTIZED_V
+
+#endif // U32_DEQUANT_HELPERS
 
 
 

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

@@ -138,7 +138,7 @@ const FLOAT_MIN: f32 = -1.0e9;
 // The number of Q rows processed per workgroup
 var<workgroup> q_shmem: array<f16, Q_TILE * HEAD_DIM_QK>;
 
-#ifndef KV_DIRECT
+#if !defined(K_DIRECT) || !defined(V_DIRECT)
 const kv_shmem_size = KV_TILE * max(HEAD_DIM_QK, HEAD_DIM_V);
 // we can reuse the same shmem for K and V since we only need one at a time
 var<workgroup> kv_shmem: array<f16, kv_shmem_size>;
@@ -183,13 +183,12 @@ fn load_kx4(buf: ptr<storage, array<vec4<K_TYPE>>, read_write>, scalar_index: u3
     return (*buf)[scalar_index >> 2u];
 }
 
-#ifndef KV_DIRECT
+#if !defined(K_DIRECT) || !defined(V_DIRECT)
 #define QUANT_SHMEM kv_shmem
 #define QUANT_OUT_TYPE f16
-#include "quant_inner_loops.tmpl"
 #include "flash_attn_quant_staging.tmpl"
 
-#if !defined(K_Q4_0) && !defined(K_Q8_0)
+#if !defined(K_DIRECT) && !defined(K_Q4_0) && !defined(K_Q8_0)
 fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
     for (var elem_idx = local_x; elem_idx < KV_TILE * HEAD_DIM_QK; elem_idx += WG_SIZE) {
         let k_row = elem_idx / HEAD_DIM_QK;
@@ -204,7 +203,7 @@ fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u
 }
 #endif
 
-#if !defined(V_Q4_0) && !defined(V_Q8_0)
+#if !defined(V_DIRECT) && !defined(V_Q4_0) && !defined(V_Q8_0)
 fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
     for (var elem_idx = local_x; elem_idx < KV_TILE * HEAD_DIM_V; elem_idx += WG_SIZE) {
         let v_row = elem_idx / HEAD_DIM_V;
@@ -296,7 +295,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
         }
 
       // load k tile into shared memory
-#ifndef KV_DIRECT
+#ifndef K_DIRECT
       load_k_tile_block(local_id.x, kv_count, kv_tile, k_head_offset);
 #endif
 
@@ -306,7 +305,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
       // TODO: this loop seems to be the current largest bottleneck
       // this bracket exists to scope the lifetime of variables, reducing register pressure
       {
-#ifdef KV_DIRECT
+#ifdef K_DIRECT
           let k_block_row = kv_tile + subgroup_id * SG_MAT_N;
           var k_global_offset = k_head_offset + k_block_row * params.stride_k1;
 #else
@@ -318,7 +317,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
 
               var q_cur = subgroupMatrixLoad<subgroup_matrix_left<f16, SG_MAT_K, SG_MAT_M>>(&q_shmem, 0u, false, HEAD_DIM_QK);
 
-#ifdef KV_DIRECT
+#ifdef K_DIRECT
               var k_cur = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&K, k_global_offset + 0u, true, params.stride_k1);
 #else
               var k_cur = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&kv_shmem, k_block_offset + 0u, true, HEAD_DIM_QK);
@@ -328,7 +327,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
               for (; t + 1u < HEAD_DIM_QK / SG_MAT_K; t += 2u) {
                   let h0 = t * SG_MAT_K;
                   var q0 = subgroupMatrixLoad<subgroup_matrix_left<f16, SG_MAT_K, SG_MAT_M>>(&q_shmem, h0, false, HEAD_DIM_QK);
-#ifdef KV_DIRECT
+#ifdef K_DIRECT
                   var k0 = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&K, k_global_offset + h0, true, params.stride_k1);
 #else
                   var k0 = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&kv_shmem, k_block_offset + h0, true, HEAD_DIM_QK);
@@ -339,7 +338,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
 
                   let h1 = (t + 1u) * SG_MAT_K;
                   var q1g = subgroupMatrixLoad<subgroup_matrix_left<f16, SG_MAT_K, SG_MAT_M>>(&q_shmem, h1, false, HEAD_DIM_QK);
-#ifdef KV_DIRECT
+#ifdef K_DIRECT
                   var k1g = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&K, k_global_offset + h1, true, params.stride_k1);
 #else
                   var k1g = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&kv_shmem, k_block_offset + h1, true, HEAD_DIM_QK);
@@ -353,7 +352,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
               if (t < HEAD_DIM_QK / SG_MAT_K) {
                   let h = t * SG_MAT_K;
                   var qn = subgroupMatrixLoad<subgroup_matrix_left<f16, SG_MAT_K, SG_MAT_M>>(&q_shmem, h, false, HEAD_DIM_QK);
-#ifdef KV_DIRECT
+#ifdef K_DIRECT
                   var kn = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&K, k_global_offset + h, true, params.stride_k1);
 #else
                   var kn = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&kv_shmem, k_block_offset + h, true, HEAD_DIM_QK);
@@ -365,7 +364,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
 
               acc = subgroupMatrixMultiplyAccumulate(q_cur, k_cur, acc);
 
-#ifdef KV_DIRECT
+#ifdef K_DIRECT
               k_global_offset += num_subgroups * SG_MAT_N * params.stride_k1;
 #else
               k_block_offset += num_subgroups * SG_MAT_N * HEAD_DIM_QK;
@@ -436,7 +435,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
       }
 
       // load v tile into shared memory
-#ifndef KV_DIRECT
+#ifndef V_DIRECT
       load_v_tile_block(local_id.x, kv_count, kv_tile, v_head_offset);
 #endif
 
@@ -464,7 +463,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
                   );
 
                   // load V submatrix from global or shared memory
-#ifdef KV_DIRECT
+#ifdef V_DIRECT
                   let v_block_row = kv_tile + kv_block * SG_MAT_N;
                   let v_global_offset = v_head_offset + v_block_row * params.stride_v1 + head_dim_block;
                   var v_sg_mat: subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K> = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(

ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_quant_staging.tmpl

@@ -1,3 +1,5 @@
+#include "quant_inner_loops.tmpl"
+
 #define BLOCK_SIZE 32
 #define BLOCKS_K ((HEAD_DIM_QK + BLOCK_SIZE - 1) / BLOCK_SIZE)
 #define BLOCKS_V ((HEAD_DIM_V + BLOCK_SIZE - 1) / BLOCK_SIZE)
@@ -26,49 +28,6 @@
 #define V_BYTES_PER_INNER_LOOP 4u
 #endif
 
-#if defined(K_Q4_0) || defined(K_Q8_0)
-fn load_k_u16_at(byte_offset: u32) -> u32 {
-    let word = K[byte_offset / 4u];
-    let shift = (byte_offset & 2u) * 8u;
-    return (word >> shift) & 0xFFFFu;
-}
-
-fn load_k_u32_at(byte_offset: u32) -> u32 {
-    let word_idx = byte_offset / 4u;
-    let shift = (byte_offset & 3u) * 8u;
-    let lo = K[word_idx];
-    if (shift == 0u) {
-        return lo;
-    }
-    let hi = K[word_idx + 1u];
-    return (lo >> shift) | (hi << (32u - shift));
-}
-#endif
-
-#if defined(V_Q4_0) || defined(V_Q8_0)
-fn load_v_u16_at(byte_offset: u32) -> u32 {
-    let word = V[byte_offset / 4u];
-    let shift = (byte_offset & 2u) * 8u;
-    return (word >> shift) & 0xFFFFu;
-}
-
-fn load_v_u32_at(byte_offset: u32) -> u32 {
-    let word_idx = byte_offset / 4u;
-    let shift = (byte_offset & 3u) * 8u;
-    let lo = V[word_idx];
-    if (shift == 0u) {
-        return lo;
-    }
-    let hi = V[word_idx + 1u];
-    return (lo >> shift) | (hi << (32u - shift));
-}
-#endif
-
-fn f16_from_u16(bits: u32) -> f16 {
-    let packed = unpack2x16float(bits);
-    return f16(packed[0]);
-}
-
 #if defined(K_Q4_0) || defined(K_Q8_0)
 fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
     for (var elem_idx = local_x * K_NQ; elem_idx < kv_count * HEAD_DIM_QK; elem_idx += WG_SIZE * K_NQ) {

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

@@ -153,7 +153,6 @@ var<workgroup> p_shmem: array<f16, Q_TILE * KV_TILE>;
 
 #define QUANT_SHMEM kv_shmem
 #define QUANT_OUT_TYPE f16
-#include "quant_inner_loops.tmpl"
 #include "flash_attn_quant_staging.tmpl"
 
 #if !defined(K_Q4_0) && !defined(K_Q8_0)
@@ -270,7 +269,9 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
             local_scores[slot] = FLOAT_MIN;
         }
 
-#ifndef KV_DIRECT
+        // The tile path stages K/V in shared memory so each tile can be reused across
+        // Q_TILE query rows. It therefore does not use the direct path.
+#ifndef K_DIRECT
         load_k_tile_block(local_id.x, kv_count, kv_tile, k_head_offset);
 #endif
 
@@ -333,7 +334,9 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
 
         workgroupBarrier();
 
-#ifndef KV_DIRECT
+        // The tile path stages K/V in shared memory so each tile can be reused across
+        // Q_TILE query rows. It therefore does not use the direct path.
+#ifndef V_DIRECT
         load_v_tile_block(local_id.x, kv_count, kv_tile, v_head_offset);
 #endif
 

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

@@ -196,49 +196,35 @@ struct Params {
 
 // Just a very small float value.
 const FLOAT_MIN: f32 = -1.0e9;
+const kv_shmem_size = KV_TILE * max(HEAD_DIM_QK, HEAD_DIM_V);
 
 var<workgroup> q_shmem: array<f32, HEAD_DIM_QK>;
-
-#ifndef KV_DIRECT
-const kv_shmem_size = KV_TILE * max(HEAD_DIM_QK, HEAD_DIM_V);
-// we can reuse the same shmem for K and V since we only need one at a time
-var<workgroup> kv_shmem: array<f32, kv_shmem_size>;
-#endif
-
 var<workgroup> o_shmem: array<f32, HEAD_DIM_V>;
+// note that we reuse the same storage for both since we only need one at a time
+var<workgroup> inter_shmem: array<f32, KV_TILE>;
 
 #ifdef MASK
 // storage for mask values
 var<workgroup> mask_shmem: array<f32, KV_TILE>;
 #endif
 
-// note that we reuse the same storage for both since we only need one at a time
-var<workgroup> inter_shmem: array<f32, KV_TILE>;
-
-// Storage for row max and exp sum during online softmax
-fn calc_softmax_term(kv_idx: u32, slope: f32, has_bias: bool, apply_mask: bool) -> f32 {
-    var v = select(FLOAT_MIN,
-                   inter_shmem[kv_idx] * params.scale,
-                   kv_idx < KV_TILE);
-#ifdef LOGIT_SOFTCAP
-    v = params.logit_softcap * tanh(v);
+#if defined(K_DIRECT) || defined(V_DIRECT)
+// Shared memory for scale factor (d) in quantized K/V. Multiple threads use the same value,
+// so caching it is more efficient, even on the direct path.
+var<workgroup> d_shmem: array<f32, kv_shmem_size / 32>;
 #endif
-#ifdef MASK
-    if (apply_mask) {
-        var mask_val = select(0.0, mask_shmem[kv_idx], kv_idx < KV_TILE);
-        v += select(mask_val, slope * mask_val, has_bias);
-    }
-#endif
-    return v;
-}
 
-#ifndef KV_DIRECT
+// K/V shared memory handling
+#if !defined(K_DIRECT) || !defined(V_DIRECT)
+
+// we can reuse the same shmem for K and V since we only need one at a time
+var<workgroup> kv_shmem: array<f32, kv_shmem_size>;
+
 #define QUANT_SHMEM kv_shmem
 #define QUANT_OUT_TYPE f32
-#include "quant_inner_loops.tmpl"
 #include "flash_attn_quant_staging.tmpl"
 
-#if !defined(K_Q4_0) && !defined(K_Q8_0)
+#if !defined(K_DIRECT) && !defined(K_Q4_0) && !defined(K_Q8_0)
 fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
     for (var elem_idx = local_x * 4u; elem_idx < KV_TILE * HEAD_DIM_QK; elem_idx += WG_SIZE * 4u) {
         let k_row = elem_idx / HEAD_DIM_QK;
@@ -256,7 +242,7 @@ fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u
 }
 #endif
 
-#if !defined(V_Q4_0) && !defined(V_Q8_0)
+#if !defined(V_DIRECT) && !defined(V_Q4_0) && !defined(V_Q8_0)
 fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
     for (var elem_idx = local_x * 4u; elem_idx < KV_TILE * HEAD_DIM_V; elem_idx += WG_SIZE * 4u) {
         let v_row = elem_idx / HEAD_DIM_V;
@@ -273,7 +259,24 @@ fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u
     }
 }
 #endif
+#endif // !defined(K_DIRECT) || !defined(V_DIRECT)
+
+// Storage for row max and exp sum during online softmax
+fn calc_softmax_term(kv_idx: u32, slope: f32, has_bias: bool, apply_mask: bool) -> f32 {
+    var v = select(FLOAT_MIN,
+                   inter_shmem[kv_idx] * params.scale,
+                   kv_idx < KV_TILE);
+#ifdef LOGIT_SOFTCAP
+    v = params.logit_softcap * tanh(v);
 #endif
+#ifdef MASK
+    if (apply_mask) {
+        var mask_val = select(0.0, mask_shmem[kv_idx], kv_idx < KV_TILE);
+        v += select(mask_val, slope * mask_val, has_bias);
+    }
+#endif
+    return v;
+}
 
 @compute @workgroup_size(WG_SIZE)
 fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
@@ -355,12 +358,31 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
             inter_shmem[elem_idx] = 0.0;
         }
 
-      // load k tile into shared memory
-#ifndef KV_DIRECT
-      load_k_tile_block(local_id.x, kv_count, kv_tile, k_head_offset);
+#ifdef K_DIRECT
+      // load only the scale factor (d) from each quantized block into shared memory on the direct path.
+#if defined(K_Q8_0)
+        for (var j = local_id.x * 32; j < KV_TILE * HEAD_DIM_QK; j += WG_SIZE * 32) {
+            let kv_row = kv_tile + j / HEAD_DIM_QK;
+            let block_idx = (j % HEAD_DIM_QK) / 32;
+            let block_byte_base = 34 * (k_head_offset + kv_row * params.stride_k1 + block_idx);
+            let d = f32(f16_from_u16(load_k_u16_at(block_byte_base)));
+            d_shmem[j / 32] = d;
+        }
+#elif defined(K_Q4_0)
+        for (var j = local_id.x * 32; j < KV_TILE * HEAD_DIM_QK; j += WG_SIZE * 32) {
+            let kv_row = kv_tile + j / HEAD_DIM_QK;
+            let block_idx = (j % HEAD_DIM_QK) / 32;
+            let block_byte_base = 18 * (k_head_offset + kv_row * params.stride_k1 + block_idx);
+            let d = f32(f16_from_u16(load_k_u16_at(block_byte_base)));
+            d_shmem[j / 32] = d;
+        }
 #endif
+#else
+      // load k tile into shared memory
+      load_k_tile_block(local_id.x, kv_count, kv_tile, k_head_offset);
+#endif // defined(K_DIRECT)
 
-      workgroupBarrier();
+        workgroupBarrier();
 
       // accumulate q block * k block into registers across the entire KV tile
       if (!skip_tile) {
@@ -381,9 +403,40 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
                             q_shmem[q_off + 1u],
                             q_shmem[q_off + 2u],
                             q_shmem[q_off + 3u]);
-#ifdef KV_DIRECT
+#ifdef K_DIRECT
+#if defined(K_Q8_0)
+                        let kv_row = kv_tile + kv_idx;
+                        let block_idx = (i * 4u) / 32;
+                        let id_in_block = (i * 4u) % 32;
+                        let block_byte_base = 34 * (k_head_offset + kv_row * params.stride_k1 + block_idx);
+                        let q_byte_base = block_byte_base + 2u;
+                        let d = d_shmem[(kv_idx * HEAD_DIM_QK) / 32 + block_idx];
+                        let q8u4 = load_k_u32_at(q_byte_base + id_in_block);
+                        let kv = vec4<f32>(
+                            d * f32(get_byte_i32(q8u4, 0)),
+                            d * f32(get_byte_i32(q8u4, 1)),
+                            d * f32(get_byte_i32(q8u4, 2)),
+                            d * f32(get_byte_i32(q8u4, 3)),
+                        );
+#elif defined(K_Q4_0)
+                        let kv_row = kv_tile + kv_idx;
+                        let block_idx = (i * 4u) / 32;
+                        let id_in_block = (i * 4u) % 32;
+                        let phase = id_in_block / 16;
+                        let block_byte_base = 18 * (k_head_offset + kv_row * params.stride_k1 + block_idx);
+                        let q_byte_base = block_byte_base + 2u;
+                        let d = d_shmem[(kv_idx * HEAD_DIM_QK) / 32 + block_idx];
+                        let q8u4 = load_k_u32_at(q_byte_base + (id_in_block - phase * 16u));
+                        let kv = vec4<f32>(
+                            d * (f32((get_byte(q8u4, 0) >> (phase * 4u)) & 0xFu) - 8.0),
+                            d * (f32((get_byte(q8u4, 1) >> (phase * 4u)) & 0xFu) - 8.0),
+                            d * (f32((get_byte(q8u4, 2) >> (phase * 4u)) & 0xFu) - 8.0),
+                            d * (f32((get_byte(q8u4, 3) >> (phase * 4u)) & 0xFu) - 8.0),
+                        );
+#else
                         let idx = k_head_offset + (kv_tile + kv_idx) * params.stride_k1 + (i * 4u);
                         let kv = vec4<f32>(K[idx >> 2u]);
+#endif
 #else
                         let idx = kv_idx * HEAD_DIM_QK + (i * 4u);
                         let kv = vec4<f32>(
@@ -391,7 +444,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
                             kv_shmem[idx + 1u],
                             kv_shmem[idx + 2u],
                             kv_shmem[idx + 3u]);
-#endif
+#endif // defined(K_DIRECT)
                         partial_sum += dot(qv, kv);
                     }
                   }
@@ -473,12 +526,32 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
           }
       }
 
-      // load v tile into shared memory
-#ifndef KV_DIRECT
-      load_v_tile_block(local_id.x, kv_count, kv_tile, v_head_offset);
-#endif
 
-      workgroupBarrier();
+#ifdef V_DIRECT
+      // load only `d` of quantized block into shared memory in the direct path
+#if defined(V_Q8_0)
+        for (var j = local_id.x * 32; j < KV_TILE * HEAD_DIM_V; j += WG_SIZE * 32) {
+            let v_row = kv_tile + j / HEAD_DIM_V;
+            let block_idx = (j % HEAD_DIM_V) / 32;
+            let block_byte_base = 34 * (v_head_offset + v_row * params.stride_v1 + block_idx);
+            let d = f32(f16_from_u16(load_v_u16_at(block_byte_base)));
+            d_shmem[j / 32] = d;
+        }
+#elif defined(V_Q4_0)
+        for (var j = local_id.x * 32; j < KV_TILE * HEAD_DIM_V; j += WG_SIZE * 32) {
+            let v_row = kv_tile + j / HEAD_DIM_V;
+            let block_idx = (j % HEAD_DIM_V) / 32;
+            let block_byte_base = 18 * (v_head_offset + v_row * params.stride_v1 + block_idx);
+            let d = f32(f16_from_u16(load_v_u16_at(block_byte_base)));
+            d_shmem[j / 32] = d;
+        }
+#endif
+#else
+      // load v tile into shared memory
+      load_v_tile_block(local_id.x, kv_count, kv_tile, v_head_offset);
+#endif // V_DIRECT
+
+        workgroupBarrier();
 
       if (!skip_tile) {
           // we have P (KV_TILE) in inter_shmem and V (KV_TILE x head_dim_v) in kv_shmem
@@ -501,9 +574,38 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
                       }
 
                       let p = inter_shmem[kv_idx];
-#ifdef KV_DIRECT
+#ifdef V_DIRECT
+#if defined(V_Q8_0)
+                        let block_idx = (vec_col * 4u) / 32;
+                        let id_in_block = (vec_col * 4u) % 32;
+                        let block_byte_base = 34 * (v_head_offset + v_row * params.stride_v1 + block_idx);
+                        let q_byte_base = block_byte_base + 2u;
+                        let d = d_shmem[(kv_idx * HEAD_DIM_V) / 32 + block_idx];
+                        let q8u4 = load_v_u32_at(q_byte_base + id_in_block);
+                        let v4 = vec4<f32>(
+                            d * f32(get_byte_i32(q8u4, 0)),
+                            d * f32(get_byte_i32(q8u4, 1)),
+                            d * f32(get_byte_i32(q8u4, 2)),
+                            d * f32(get_byte_i32(q8u4, 3)),
+                        );
+#elif defined(V_Q4_0)
+                        let block_idx = (vec_col * 4u) / 32;
+                        let id_in_block = (vec_col * 4u) % 32;
+                        let phase = id_in_block / 16;
+                        let block_byte_base = 18 * (v_head_offset + v_row * params.stride_v1 + block_idx);
+                        let q_byte_base = block_byte_base + 2u;
+                        let d = d_shmem[(kv_idx * HEAD_DIM_V) / 32 + block_idx];
+                        let q8u4 = load_v_u32_at(q_byte_base + (id_in_block - phase * 16u));
+                        let v4 = vec4<f32>(
+                            d * (f32((get_byte(q8u4, 0) >> (phase * 4u)) & 0xFu) - 8.0),
+                            d * (f32((get_byte(q8u4, 1) >> (phase * 4u)) & 0xFu) - 8.0),
+                            d * (f32((get_byte(q8u4, 2) >> (phase * 4u)) & 0xFu) - 8.0),
+                            d * (f32((get_byte(q8u4, 3) >> (phase * 4u)) & 0xFu) - 8.0),
+                        );
+#else
                       let v_idx = v_head_offset + v_row * params.stride_v1 + vec_col * 4u;
                       let v4 = vec4<f32>(V[v_idx >> 2u]);
+#endif
 #else
                       let v_idx = kv_idx * HEAD_DIM_V + vec_col * 4u;
                       let v4 = vec4<f32>(
@@ -511,7 +613,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
                           kv_shmem[v_idx + 1u],
                           kv_shmem[v_idx + 2u],
                           kv_shmem[v_idx + 3u]);
-#endif
+#endif // defined(V_DIRECT)
                       lo += p * v4;
                   }
 

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

@@ -27,6 +27,9 @@ struct Params {
 #ifdef TYPE_I32
 #define DataType i32
 #endif
+#ifdef TYPE_F16
+#define DataType f16
+#endif
 #ifdef TYPE_I16
 // same size (16-bit) is sufficient for repeat
 #define DataType f16

ggml/src/ggml-zendnn/CMakeLists.txt

@@ -28,7 +28,7 @@ if (NOT ZENDNN_ROOT OR ZENDNN_ROOT STREQUAL "" OR ZENDNN_ROOT STREQUAL "OFF")
     ExternalProject_Add(
         zendnn
         GIT_REPOSITORY https://github.com/amd/ZenDNN.git
-        GIT_TAG 253b94ce0d7e9284c265fefb485714944caff9d3    # ZenDNN-2026-WW19
+        GIT_TAG 1f399a75cc0993778374a51bea49b64a57879595    # ZenDNN-2026-WW28
         PREFIX      ${ZENDNN_PREFIX}
         SOURCE_DIR  ${ZENDNN_SOURCE_DIR}
         BINARY_DIR  ${ZENDNN_BUILD_DIR}

ggml/src/ggml-zendnn/ggml-zendnn.cpp

@@ -30,6 +30,29 @@ zendnnl::common::data_type_t ggml_to_zendnn_type() {
     }
 }
 
+/**
+ * Builds the matmul_params shared by ggml_zendnn_matmul() and ggml_zendnn_group_matmul():
+ * dtype selection plus, for Q8_0 weights, dynamic-quant setup. Callers still need to set
+ * quant_params.src_scale.dims themselves, since that depends on the batch size(s) in use.
+ */
+template <typename TA, typename TB, typename TC>
+static zendnnl::lowoha::matmul::matmul_params ggml_zendnn_make_matmul_params(ggml_backend_zendnn_context * ctx) {
+    zendnnl::lowoha::matmul::matmul_params params;
+    params.dtypes.src = ggml_to_zendnn_type<TB>();
+    params.dtypes.wei = ggml_to_zendnn_type<TA>();
+    params.dtypes.dst = ggml_to_zendnn_type<TC>();
+    params.num_threads = ctx->n_threads;
+
+    if constexpr (std::is_same_v<TA, block_q8_0>) {
+        params.dtypes.compute = zendnnl::common::data_type_t::s8;
+        params.dynamic_quant = true;
+        params.quant_params.src_scale.buff = nullptr;
+        params.quant_params.src_scale.dt   = zendnnl::common::data_type_t::bf16;
+        params.packing.pack_format_b = 1;
+    }
+    return params;
+}
+
 /**
  * ZenDNN matmul: computes C = B * A.
  *
@@ -47,22 +70,12 @@ static bool ggml_zendnn_matmul(ggml_backend_zendnn_context * ctx, int64_t m, int
                                const TA * A, int64_t lda, const TB * B, int64_t ldb, TC * C,
                                int64_t ldc) {
 
-    zendnnl::lowoha::matmul::matmul_params params;
-    params.dtypes.src = ggml_to_zendnn_type<TB>();
-    params.dtypes.wei = ggml_to_zendnn_type<TA>();
-    params.dtypes.dst = ggml_to_zendnn_type<TC>();
-    params.num_threads = ctx->n_threads;
+    zendnnl::lowoha::matmul::matmul_params params = ggml_zendnn_make_matmul_params<TA, TB, TC>(ctx);
 
     zendnnl::lowoha::matmul::matmul_batch_params_t batch_params;
 
     if constexpr (std::is_same_v<TA, block_q8_0>) {
-        params.dtypes.compute = zendnnl::common::data_type_t::s8;
-        const int64_t num_groups = k / QK8_0;
-        params.dynamic_quant = true;
-        params.quant_params.src_scale.buff = nullptr;
-        params.quant_params.src_scale.dt   = zendnnl::common::data_type_t::bf16;
-        params.quant_params.src_scale.dims = {n, num_groups};
-        params.packing.pack_format_b = 1;
+        params.quant_params.src_scale.dims = {n, k / QK8_0};
     }
 
     zendnnl::error_handling::status_t status = zendnnl::lowoha::matmul::matmul_direct(
@@ -223,6 +236,99 @@ struct mmid_row_mapping {
     int32_t i2;
 };
 
+/**
+ * ZenDNN batched matmul: computes C[i] = B[i] * A[i] for every active expert i via a single
+ * group_matmul_direct() call. Batched analogue of ggml_zendnn_matmul() - see its docs for the
+ * per-expert A/B/C shape convention. m and k are shared by every expert; n (batch size) varies
+ * per expert, hence the vector.
+ */
+template <typename TA, typename TB, typename TC>
+static bool ggml_zendnn_group_matmul(ggml_backend_zendnn_context * ctx, int64_t m, int64_t k,
+                                     const std::vector<int64_t> & n,
+                                     const std::vector<const void *> & A, int64_t lda,
+                                     const std::vector<const void *> & B, int64_t ldb,
+                                     const std::vector<void *> & C, int64_t ldc) {
+
+    const int n_experts = n.size();
+
+    zendnnl::lowoha::matmul::matmul_params base_params = ggml_zendnn_make_matmul_params<TA, TB, TC>(ctx);
+
+    std::vector<char>         layout(n_experts, 'r');
+    std::vector<bool>         trans_a(n_experts, false);
+    std::vector<bool>         trans_b(n_experts, true);
+    std::vector<int>          batch_m(n_experts);
+    std::vector<int>          batch_n(n_experts, m);
+    std::vector<int>          batch_k(n_experts, k);
+    std::vector<float>        alpha(n_experts, 1.0f);
+    std::vector<float>        beta(n_experts, 0.0f);
+    std::vector<const void *> bias(n_experts, nullptr);
+    std::vector<int>          lda_v(n_experts, lda);
+    std::vector<int>          ldb_v(n_experts, ldb);
+    std::vector<int>          ldc_v(n_experts, ldc);
+    std::vector<bool>         is_wei_const(n_experts, true);
+    std::vector<zendnnl::lowoha::matmul::matmul_params> params(n_experts, base_params);
+
+    for (int i = 0; i < n_experts; i++) {
+        batch_m[i] = n[i];
+
+        // src_scale.dims depends on this expert's row count, unlike the rest of base_params
+        if constexpr (std::is_same_v<TA, block_q8_0>) {
+            params[i].quant_params.src_scale.dims = {n[i], k / QK8_0};
+        }
+    }
+
+    zendnnl::error_handling::status_t status = zendnnl::lowoha::matmul::group_matmul_direct(
+        layout, trans_a, trans_b, batch_m, batch_n, batch_k, alpha,
+        B, ldb_v, A, lda_v, bias, beta,
+        C, ldc_v, is_wei_const, params);
+
+    if (status != zendnnl::error_handling::status_t::success) {
+        GGML_LOG_ERROR("%s, ZenDNN group matmul failed: status=%d\n", __func__, static_cast<int>(status));
+        return false;
+    }
+    return true;
+}
+
+static bool ggml_zendnn_group_gemm(ggml_backend_zendnn_context * ctx, int64_t m, int64_t k,
+                                    const std::vector<int64_t> & n,
+                                    const std::vector<const void *> & A, int64_t lda,
+                                    const std::vector<const void *> & B, int64_t ldb,
+                                    const std::vector<void *> & C, int64_t ldc,
+                                    int Atype, int Btype, int Ctype) {
+
+    assert(m >= 0);
+    for (size_t i = 0; i < n.size(); i++) {
+        assert(n[i] >= 0);
+    }
+    assert(k >= 0);
+    assert(lda >= k);
+    assert(ldb >= k);
+    assert(ldc >= m);
+
+    // categorize types
+    switch (Atype) {
+        case GGML_TYPE_F32:
+            if (Btype != GGML_TYPE_F32 || Ctype != GGML_TYPE_F32)
+                return false;
+            return ggml_zendnn_group_matmul<float, float, float>(ctx, m, k, n, A, lda, B, ldb, C, ldc);
+        case GGML_TYPE_BF16:
+            if (Btype != GGML_TYPE_BF16)
+                return false;
+            if (Ctype == GGML_TYPE_BF16)
+                return ggml_zendnn_group_matmul<ggml_bf16_t, ggml_bf16_t, ggml_bf16_t>(
+                    ctx, m, k, n, A, lda, B, ldb, C, ldc);
+            if (Ctype == GGML_TYPE_F32)
+                return ggml_zendnn_group_matmul<ggml_bf16_t, ggml_bf16_t, float>(ctx, m, k, n, A, lda, B, ldb, C, ldc);
+            return false;
+        case GGML_TYPE_Q8_0:
+            if (Btype != GGML_TYPE_F32 || Ctype != GGML_TYPE_F32)
+                return false;
+            return ggml_zendnn_group_matmul<block_q8_0, float, float>(ctx, m, k, n, A, lda, B, ldb, C, ldc);
+        default:
+            return false; // unsupported type
+    }
+}
+
 static void ggml_zendnn_compute_forward_mul_mat_id(
     ggml_backend_zendnn_context * ctx,
     ggml_tensor * dst) {
@@ -262,7 +368,8 @@ static void ggml_zendnn_compute_forward_mul_mat_id(
     std::vector<int64_t> matrix_row_counts(n_as, 0);
     std::vector<std::vector<mmid_row_mapping>> matrix_rows(n_as);
 
-    int64_t max_rows = 0;
+    int64_t total_rows = 0;
+    int     n_active_experts = 0;
     // group rows by expert (preprocessing step)
     for (int64_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
         for (int id = 0; id < n_ids; ++id) {
@@ -270,66 +377,74 @@ static void ggml_zendnn_compute_forward_mul_mat_id(
 
             GGML_ASSERT(i02 >= 0 && i02 < n_as);
 
+            if (matrix_row_counts[i02] == 0) {
+                n_active_experts++;
+            }
             matrix_rows[i02].push_back({id, iid1});
             matrix_row_counts[i02]++;
-            if (matrix_row_counts[i02] > max_rows) {
-                max_rows = matrix_row_counts[i02];
-            }
+            total_rows++;
         }
     }
 
-    if (max_rows == 0) {
+    if (total_rows == 0) {
         return; // no rows to process
     }
 
     const size_t row_size = ggml_row_size(vec_dot_type, ne10);
 
-    // size for converting src1 rows to vec_dot_type if needed
-    const size_t nbw1 = row_size;
-    const size_t nbw2 = nbw1 * ne11;
-    const size_t nbw3 = nbw2 * ne12;
-    const size_t src1_conv_size = (src1->type != vec_dot_type && src0->type != GGML_TYPE_Q8_0) ? ne13 * nbw3 : 0;
-
     // For Q8_0, src1 is always F32; the gather buffer must hold F32 rows (ne10*4 bytes),
     // not Q8_0-encoded rows (row_size ≈ ne10/32*34 bytes) — they differ by ~4x.
     const size_t f32_row_size = (size_t)ne10 * sizeof(float);
     const size_t gather_row_size = (src0->type == GGML_TYPE_Q8_0) ? f32_row_size : row_size;
 
+    if (src1->type != vec_dot_type && src0->type != GGML_TYPE_Q8_0) {
+        GGML_ASSERT(src1->type == GGML_TYPE_F32);
+    }
+
     // size for MoE gather/scatter buffers
-    const size_t wdata_cur_size = max_rows * gather_row_size;
-    const size_t dst_cur_size = max_rows * ggml_row_size(dst->type, ne01);
+    const size_t wdata_cur_size = total_rows * gather_row_size;
+    const size_t dst_cur_size = total_rows * ggml_row_size(dst->type, ne01);
 
     // allocate single buffer for all needs
-    const size_t total_size = src1_conv_size + wdata_cur_size + dst_cur_size;
+    const size_t total_size = wdata_cur_size + dst_cur_size;
     if (ctx->work_size < total_size) {
         ctx->work_data.reset(new char[total_size]);
         ctx->work_size = total_size;
     }
 
     // partition the buffer
-    char * work_data = ctx->work_data.get();
-    char * wdata_cur = work_data + src1_conv_size;
+    char * wdata_cur = ctx->work_data.get();
     char * dst_cur = wdata_cur + wdata_cur_size;
 
-    // ZenDNN requires FP32 for dynamic quantization, so conversion is skipped
-    if (src1->type != vec_dot_type && src0->type != GGML_TYPE_Q8_0) {
-        GGML_ASSERT(src1->type == GGML_TYPE_F32);
+    // per-expert data collected during gather, handed to ggml_zendnn_group_gemm() as one batch
+    std::vector<int64_t>      expert_row_count(n_active_experts);
+    std::vector<const void *> batch_src(n_active_experts);
+    std::vector<const void *> batch_wei(n_active_experts);
+    std::vector<void *>       batch_dst(n_active_experts);
 
-        #pragma omp parallel for collapse(3) num_threads(ctx->n_threads) schedule(static)
-        for (int64_t i13 = 0; i13 < ne13; ++i13) {
-            for (int64_t i12 = 0; i12 < ne12; ++i12) {
-                for (int64_t i11 = 0; i11 < ne11; ++i11) {
-                    const float * src1_f32 = (float *)((char *)src1->data + i11*nb11 + i12*nb12 + i13*nb13);
-                    void * src1_conv = (char *)work_data + i11*nbw1 + i12*nbw2 + i13*nbw3;
-                    from_float(src1_f32, src1_conv, ne10);
-                }
+    // precompute per-expert buffer offsets and batch indices for the parallel loop below
+    std::vector<int64_t> expert_wdata_off(n_as, 0);
+    std::vector<int64_t> expert_dst_off(n_as, 0);
+    std::vector<int>     expert_batch_idx(n_as, -1);
+    {
+        int64_t w_off = 0;
+        int64_t d_off = 0;
+        int     batch_idx = 0;
+        for (int64_t cur_a = 0; cur_a < n_as; ++cur_a) {
+            if (matrix_row_counts[cur_a] == 0) {
+                continue;
             }
+            expert_wdata_off[cur_a] = w_off;
+            expert_dst_off[cur_a]   = d_off;
+            expert_batch_idx[cur_a] = batch_idx;
+            w_off += matrix_row_counts[cur_a] * gather_row_size;
+            d_off += matrix_row_counts[cur_a] * ggml_row_size(dst->type, ne01);
+            batch_idx++;
         }
     }
 
-    const void * wdata = (src1->type == vec_dot_type || src0->type == GGML_TYPE_Q8_0) ? src1->data : work_data;
-
-    // process each expert with gather -> gemm -> scatter pattern
+    // gather + inline-convert input rows into each expert's batch slot
+    #pragma omp parallel for num_threads(ctx->n_threads) schedule(static)
     for (int64_t cur_a = 0; cur_a < n_as; ++cur_a) {
         const int64_t cne1 = matrix_row_counts[cur_a];
 
@@ -337,42 +452,57 @@ static void ggml_zendnn_compute_forward_mul_mat_id(
             continue;
         }
 
-        const char * src0_cur = (const char *) src0->data + cur_a*nb02;
+        const int64_t w_off     = expert_wdata_off[cur_a];
+        const int64_t d_off     = expert_dst_off[cur_a];
+        const int     batch_idx = expert_batch_idx[cur_a];
 
-        // gather input rows for this expert
-        #pragma omp parallel for num_threads(ctx->n_threads) schedule(static)
         for (int64_t ir1 = 0; ir1 < cne1; ++ir1) {
             const mmid_row_mapping & row_mapping = matrix_rows[cur_a][ir1];
-            const int64_t id = row_mapping.i1;
+            const int64_t id  = row_mapping.i1;
             const int64_t i11 = id % ne11;
             const int64_t i12 = row_mapping.i2;
 
-            std::memcpy(
-                wdata_cur + ir1 * gather_row_size,
-                (const char *) wdata + (i11 + i12*ne11) * gather_row_size,
-                gather_row_size
-            );
+            const char * src_row = (const char *) src1->data + i11*nb11 + i12*nb12;
+            void *       dst_row = wdata_cur + w_off + ir1 * gather_row_size;
+
+            if (src1->type != vec_dot_type && src0->type != GGML_TYPE_Q8_0) {
+                from_float((const float *) src_row, dst_row, ne10);
+            } else {
+                // no conversion: src1 already matches vec_dot_type, or src0 is Q8_0, whose
+                // ZenDNN dynamic quantization requires the row to stay in F32
+                std::memcpy(dst_row, src_row, gather_row_size);
+            }
         }
 
-        // batched gemm for all tokens in this expert
-        if (!ggml_zendnn_gemm(ctx,
-                              ne01,       // m
-                              cne1,       // n
-                              ne10,       // k
-                              src0_cur,
-                              ne00,       // lda
-                              wdata_cur,
-                              ne10,       // ldb
-                              dst_cur,
-                              ne01,       // ldc
-                              src0->type,
-                              src0->type == GGML_TYPE_Q8_0 ? GGML_TYPE_F32 : vec_dot_type,
-                              dst->type)) {
-            GGML_ABORT("%s: ZenDNN gemm failed\n", __func__);
+        expert_row_count[batch_idx] = cne1;
+        batch_src[batch_idx]        = wdata_cur + w_off;
+        batch_wei[batch_idx]        = (const char *) src0->data + cur_a * nb02;
+        batch_dst[batch_idx]        = dst_cur + d_off;
+    }
+
+    if (!ggml_zendnn_group_gemm(ctx,
+                                ne01,             // m
+                                ne10,             // k
+                                expert_row_count, // n (per expert)
+                                batch_wei, ne00,  // A: weights, lda
+                                batch_src, ne10,  // B: input, ldb
+                                batch_dst, ne01,  // C: output, ldc
+                                src0->type,
+                                src0->type == GGML_TYPE_Q8_0 ? GGML_TYPE_F32 : vec_dot_type,
+                                dst->type))
+        GGML_ABORT("%s: ZenDNN group gemm failed\n", __func__);
+
+    // scatter output rows to destination
+    #pragma omp parallel for num_threads(ctx->n_threads) schedule(static)
+    for (int64_t cur_a = 0; cur_a < n_as; ++cur_a) {
+        const int64_t cne1 = matrix_row_counts[cur_a];
+
+        if (cne1 == 0) {
+            continue;
         }
 
-        // scatter output rows to destination
-        #pragma omp parallel for num_threads(ctx->n_threads) schedule(static)
+        const int64_t d_off = expert_dst_off[cur_a];
+
         for (int64_t ir1 = 0; ir1 < cne1; ++ir1) {
             const mmid_row_mapping & row_mapping = matrix_rows[cur_a][ir1];
             const int64_t id = row_mapping.i1;
@@ -381,7 +511,7 @@ static void ggml_zendnn_compute_forward_mul_mat_id(
 
             std::memcpy(
                 (char *) dst->data + i1*nb1 + i2*nb2,
-                dst_cur + ir1 * ggml_row_size(dst->type, ne01),
+                dst_cur + d_off + ir1 * ggml_row_size(dst->type, ne01),
                 ggml_row_size(dst->type, ne01)
             );
         }
@@ -591,22 +721,26 @@ static bool ggml_backend_zendnn_device_supports_op(ggml_backend_dev_t dev, const
                 if(K <= 256 || N <= 128 || M <= 96) {
                     return false;
                 }
+
+                // MUL_MAT_ID's gather+matmul+scatter approach favors a moderate expert count
+                if (op->op == GGML_OP_MUL_MAT_ID) {
+                    const int64_t n_experts = weights->ne[2];
+                    const int64_t max_experts = 32;
+                    if (n_experts > max_experts) {
+                        return false;
+                    }
+
+                    // fall back once the average rows per expert (N / n_experts) is too thin
+                    // to amortize each per-expert GEMM's overhead
+                    if (N / n_experts <= 32) {
+                        return false;
+                    }
+                }
             }
             else if (ne0 < min_batch || ne1 < min_batch || ne10 < min_batch) {
                 return false;
             }
 
-            // MUL_MAT_ID performs best with a moderate number of experts due to its
-            // gather + batched matmul + scatter approach. Future versions will leverage
-            // ZenDNN's grouped_gemm for better scalability with larger expert counts:
-            // https://github.com/amd/ZenDNN/blob/main/docs/operator/lowoha_group_gemm_operator.md
-            if (op->op == GGML_OP_MUL_MAT_ID) {
-                const int64_t n_experts = weights->ne[2];
-                const int64_t max_experts = 32;
-                if (n_experts > max_experts) {
-                    return false;
-                }
-            }
             switch (weights->type) {
                 case GGML_TYPE_F32:
                 case GGML_TYPE_BF16:

gguf-py/gguf/constants.py

@@ -353,6 +353,7 @@ class Keys:
         class Attention:
             HEAD_COUNT      = "clip.vision.attention.head_count"
             HEAD_COUNT_KV   = "clip.vision.attention.head_count_kv" # used by mimovl (GQA)
+            HEAD_DIM        = "clip.vision.attention.head_dim" # set when qkv width != n_embd
             LAYERNORM_EPS   = "clip.vision.attention.layer_norm_epsilon"
 
         class Projector:
@@ -3330,6 +3331,12 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.FFN_GATE_SHEXP,
         MODEL_TENSOR.FFN_DOWN_SHEXP,
         MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
     ],
     MODEL_ARCH.ERNIE4_5_MOE: [
         MODEL_TENSOR.TOKEN_EMBD,
@@ -4376,10 +4383,35 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.ATTN_OUT,
         MODEL_TENSOR.ATTN_Q_NORM,
         MODEL_TENSOR.ATTN_K_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.HC_HEAD_FN,
+        MODEL_TENSOR.HC_HEAD_BASE,
+        MODEL_TENSOR.HC_HEAD_SCALE,
         MODEL_TENSOR.FFN_NORM,
         MODEL_TENSOR.FFN_GATE,
         MODEL_TENSOR.FFN_DOWN,
         MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
         MODEL_TENSOR.FC,
         MODEL_TENSOR.ENC_OUTPUT_NORM,
         # optional DSpark heads

gguf-py/gguf/gguf_writer.py

@@ -1226,6 +1226,9 @@ class GGUFWriter:
     def add_vision_head_count_kv(self, value: int) -> None:
         self.add_uint32(Keys.ClipVision.Attention.HEAD_COUNT_KV, value)
 
+    def add_vision_head_dim(self, value: int) -> None:
+        self.add_uint32(Keys.ClipVision.Attention.HEAD_DIM, value)
+
     def add_vision_attention_layernorm_eps(self, value: float) -> None:
         self.add_float32(Keys.ClipVision.Attention.LAYERNORM_EPS, value)
 

include/llama.h

@@ -337,6 +337,7 @@ extern "C" {
         bool use_extra_bufts; // use extra buffer types (used for weight repacking)
         bool no_host;         // bypass host buffer allowing extra buffers to be used
         bool no_alloc;        // only load metadata and simulate memory allocations
+        bool load_mtp;        // whether to load MTP layers
     };
 
     struct llama_sampler_seq_config {

skills/code-review/SKILL.md

@@ -47,6 +47,7 @@ Mandatory on every review; any finding here is **blocking**. Rule of thumb: GGUF
 - **Sizes/counts from tensor dims:** validate before allocating. Products like `ne[i]*nb[i]`/nbytes can overflow on crafted dims into an undersized alloc then heap overflow. Overflow checks must run BEFORE the arithmetic they guard - padding/alignment macros wrap to 0 near `SIZE_MAX`, so a guard after the pad passes.
 - **GGUF strings/arrays:** cap declared lengths and element counts before using them to size a loop or buffer; validate element type and length before casting an array to a pointer or reading fixed indices (`[i+1]`, `[0..2]`).
 - **File-supplied counts indexing fixed arrays:** bound any count (e.g. layer/block count into a `LLAMA_MAX_*` array) before indexing; watch checks that only fire when an optional key is present.
+- **Declared vs actual array length:** check the declared length of a GGUF array against the count actually read, not just against a buffer size.
 - **Bounds comparisons:** flag narrowing casts (`size_t`->`int32_t`) and signed/unsigned mixing that can bypass a length check and copy past a buffer.
 - **Parsed/derived indices:** range-check `stoi`/`atoi` results and catch parse throws; never use a default or derived token id (EOS/BOS/...) as an index without a bounds check.
 - **Reused/reserved buffers:** recheck bounds after a buffer is shrunk or reused; watch `reserve()` then index-by-assumed-size, and header fields read before their length is checked.
@@ -131,6 +132,8 @@ Enforce the `AGENTS.md` / `CONTRIBUTING.md` coding and naming guidelines on ever
 - Reuse existing infrastructure over introducing new components; no new third-party dependencies, extra headers, or files unless clearly justified.
 - Keep it simple: a simpler change doing 90% is often preferable to a complex one doing 100%. Flag unnecessary templates/fancy STL; basic `for` loops are fine here.
 - Every added line should be something the contributor can explain and defend to a reviewer without AI help - flag anything that looks copied-in without understanding.
+- `Co-authored-by:` must be reserved for human co-authors; AI contributions (claude, cursor, codex, etc.) must use `Assisted-by:`; if this point is violated, it's a blocking finding.
+- Any mentions of Minja must be treated as blocking; see `AGENTS.md` for why.
 
 ## Reporting
 

src/llama-arch.cpp

@@ -968,6 +968,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) {
         case LLM_ARCH_KIMI_LINEAR:
         case LLM_ARCH_QWEN35:
         case LLM_ARCH_QWEN35MOE:
+        case LLM_ARCH_DEEPSEEK4:
             return true;
         default:
             return false;
@@ -990,6 +991,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) {
     switch (arch) {
         case LLM_ARCH_QWEN35:
         case LLM_ARCH_QWEN35MOE:
+        case LLM_ARCH_DEEPSEEK4:
             return true;
         default:
             return false;

src/llama-context.cpp

@@ -120,8 +120,9 @@ llama_context::llama_context(
     cparams.no_perf                 = params.no_perf;
     cparams.warmup                  = false;
 
-    cparams.embeddings_layer_inp.resize(hparams.n_layer(), false);
-    embd_layer_inp.resize(hparams.n_layer());
+    // +1: id n_layer() taps the output of the last layer ("input" of the head)
+    cparams.embeddings_layer_inp.resize(hparams.n_layer() + 1, false);
+    embd_layer_inp.resize(hparams.n_layer() + 1);
 
     cparams.ctx_type     = params.ctx_type;
     cparams.pooling_type = params.pooling_type;
@@ -1164,7 +1165,7 @@ void llama_context::set_embeddings_nextn(bool value, bool masked) {
 void llama_context::set_embeddings_layer_inp(uint32_t lid, bool enable) {
     LLAMA_LOG_DEBUG("%s: lid = %d, enable = %d\n", __func__, lid, enable);
 
-    GGML_ASSERT(lid < model.hparams.n_layer());
+    GGML_ASSERT(lid <= model.hparams.n_layer());
 
     cparams.embeddings_layer_inp[lid] = enable;
 
@@ -1716,7 +1717,8 @@ int llama_context::decode(const llama_batch & batch_inp) {
     const auto & hparams = model.hparams;
 
     const int64_t n_vocab = vocab.n_tokens();
-    const int64_t n_embd  = hparams.n_embd_inp();
+    const bool    mtp_embd = cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && batch_inp.embd;
+    const int64_t n_embd  = mtp_embd ? hparams.n_embd_out() : hparams.n_embd_inp();
 
     // when computing embeddings, all tokens are output
     const bool output_all   = cparams.embeddings;
@@ -2275,8 +2277,9 @@ void llama_context::extract_layer_inputs(const llm_graph_result * res, size_t to
 }
 
 void llama_context::output_reorder() {
-    const uint64_t n_vocab = model.vocab.n_tokens();
-    const uint64_t n_embd  = model.hparams.n_embd;
+    const uint64_t n_vocab     = model.vocab.n_tokens();
+    const uint64_t n_embd      = model.hparams.n_embd;
+    const uint64_t n_embd_out  = model.hparams.n_embd_out();
 
     for (size_t s = 0; s < output_swaps.size(); ++s) {
         const uint64_t i0 = output_swaps[s].i0;
@@ -2289,14 +2292,14 @@ void llama_context::output_reorder() {
         }
 
         if (embd.size > 0) {
-            for (uint64_t k = 0; k < n_embd; k++) {
-                std::swap(embd.data[i0*n_embd + k], embd.data[i1*n_embd + k]);
+            for (uint64_t k = 0; k < n_embd_out; k++) {
+                std::swap(embd.data[i0*n_embd_out + k], embd.data[i1*n_embd_out + k]);
             }
         }
 
         if (embd_nextn.size > 0) {
-            for (uint64_t k = 0; k < n_embd; k++) {
-                std::swap(embd_nextn.data[i0*n_embd + k], embd_nextn.data[i1*n_embd + k]);
+            for (uint64_t k = 0; k < n_embd_out; k++) {
+                std::swap(embd_nextn.data[i0*n_embd_out + k], embd_nextn.data[i1*n_embd_out + k]);
             }
         }
 
@@ -2351,6 +2354,7 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
         model.arch == LLM_ARCH_QWEN35 ||
         model.arch == LLM_ARCH_QWEN35MOE ||
         model.arch == LLM_ARCH_DEEPSEEK4 ||
+        (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) ||
         model.arch == LLM_ARCH_NANBEIGE ||
         model.arch == LLM_ARCH_MINIMAX_M3) {
         return std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
@@ -3553,6 +3557,22 @@ llama_context * llama_init_from_model(
         }
     }
 
+    if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4) && 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;
+    }
+
+    if (ggml_is_quantized(params.type_v) && params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_ENABLED) {
+        if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO) {
+            LLAMA_LOG_INFO("%s: enabling flash_attn since it is required for quantized V cache\n", __func__);
+            params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
+        }
+        if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_DISABLED) {
+            LLAMA_LOG_ERROR("%s: quantized V cache requires flash_attn to be enabled\n", __func__);
+            return nullptr;
+        }
+    }
+
     if (params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED && ggml_is_quantized(params.type_k)) {
         const uint32_t blck_size = ggml_blck_size(params.type_k);
         for (uint32_t il = 0; il < model->hparams.n_layer(); ++il) {
@@ -3575,11 +3595,6 @@ llama_context * llama_init_from_model(
         }
     }
 
-    if (ggml_is_quantized(params.type_v) && params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_DISABLED) {
-        LLAMA_LOG_ERROR("%s: V cache quantization requires flash_attn\n", __func__);
-        return nullptr;
-    }
-
     if (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED &&
         params.pooling_type != model->hparams.pooling_type) {
         //user-specified pooling-type is different from the model default

src/llama-graph.cpp

@@ -619,6 +619,63 @@ bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) {
     return res;
 }
 
+void llm_graph_input_attn_k_iswa::set_input(const llama_ubatch * ubatch) {
+    // base tensors may not be allocated if there are no non-SWA attention layers
+    if (self_k_idxs && self_k_idxs->buffer) {
+        mctx->get_base()->set_input_k_idxs(self_k_idxs, ubatch);
+    }
+
+    // the kq mask guards on its own buffer: shared cells leave idxs unbacked while the mask stays live
+    if (self_kq_mask && self_kq_mask->buffer) {
+        mctx->get_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
+    }
+
+    // swa tensors may not be allocated if there are no SWA attention layers
+    if (self_k_idxs_swa && self_k_idxs_swa->buffer) {
+        mctx->get_swa()->set_input_k_idxs(self_k_idxs_swa, ubatch);
+    }
+
+    if (self_kq_mask_swa && self_kq_mask_swa->buffer) {
+        mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn);
+    }
+
+    if (self_k_rot && self_k_rot->buffer) {
+        mctx->get_base()->set_input_k_rot(self_k_rot);
+    }
+
+    if (self_k_rot_swa && self_k_rot_swa->buffer) {
+        mctx->get_swa()->set_input_k_rot(self_k_rot_swa);
+    }
+}
+
+bool llm_graph_input_attn_k_iswa::can_reuse(const llm_graph_params & params) {
+    const auto * mctx = static_cast<const llama_kv_cache_iswa_context *>(params.mctx);
+
+    this->mctx = mctx;
+
+    bool res = true;
+
+    // base tensors may not be allocated if there are no non-SWA attention layers
+    if (self_k_idxs && self_k_idxs->buffer) {
+        res &= self_k_idxs->ne[0] == params.ubatch.n_tokens;
+    }
+
+    if (self_kq_mask && self_kq_mask->buffer) {
+        res &= can_reuse_kq_mask(self_kq_mask, mctx->get_base(), params.ubatch, params.cparams);
+    }
+
+    // swa tensors may not be allocated if there are no SWA attention layers
+    if (self_k_idxs_swa && self_k_idxs_swa->buffer) {
+        res &= self_k_idxs_swa->ne[0] == params.ubatch.n_tokens;
+    }
+
+    if (self_kq_mask_swa && self_kq_mask_swa->buffer) {
+        res &= can_reuse_kq_mask(self_kq_mask_swa, mctx->get_swa(), params.ubatch, params.cparams);
+    }
+
+    return res;
+}
+
 static void dsv4_set_i64(ggml_tensor * dst, const std::vector<int64_t> & src) {
     if (!dst || !dst->buffer) {
         return;
@@ -754,6 +811,10 @@ static void dsv4_set_comp_inputs(
     dsv4_set_i32(inp.state_pos, plan.state_pos);
     dsv4_set_i32(inp.state_persist_src_idxs, plan.state_persist_src_idxs);
     dsv4_set_i32(inp.state_persist_dst_idxs, plan.state_persist_dst_idxs);
+    dsv4_set_i32(inp.state_restore_src_idxs, plan.state_restore_src_idxs);
+    dsv4_set_i32(inp.state_restore_dst_idxs, plan.state_restore_dst_idxs);
+    dsv4_set_i32(inp.state_snapshot_src_idxs, plan.state_snapshot_src_idxs);
+    dsv4_set_i32(inp.state_snapshot_dst_idxs, plan.state_snapshot_dst_idxs);
     dsv4_set_i32(inp.state_read_idxs, plan.state_read_idxs);
     dsv4_set_i64(inp.state_write_idxs, plan.state_write_idxs);
     dsv4_set_i32(inp.state_write_pos, plan.state_write_pos);
@@ -798,6 +859,10 @@ static bool dsv4_can_reuse_comp_input(
     res &= dsv4_can_reuse_tensor_1d(inp.state_pos, plan.state_pos.size());
     res &= dsv4_can_reuse_tensor_1d(inp.state_persist_src_idxs, plan.state_persist_src_idxs.size());
     res &= dsv4_can_reuse_tensor_1d(inp.state_persist_dst_idxs, plan.state_persist_dst_idxs.size());
+    res &= dsv4_can_reuse_tensor_1d(inp.state_restore_src_idxs, plan.state_restore_src_idxs.size());
+    res &= dsv4_can_reuse_tensor_1d(inp.state_restore_dst_idxs, plan.state_restore_dst_idxs.size());
+    res &= dsv4_can_reuse_tensor_1d(inp.state_snapshot_src_idxs, plan.state_snapshot_src_idxs.size());
+    res &= dsv4_can_reuse_tensor_1d(inp.state_snapshot_dst_idxs, plan.state_snapshot_dst_idxs.size());
     res &= dsv4_can_reuse_tensor_1d(inp.state_read_idxs, plan.state_read_idxs.size());
     res &= dsv4_can_reuse_tensor_1d(inp.state_write_idxs, plan.state_write_idxs.size());
     res &= dsv4_can_reuse_tensor_1d(inp.state_write_pos, plan.state_write_pos.size());
@@ -832,6 +897,10 @@ static void dsv4_build_comp_inputs(
     inp.state_pos = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_pos.size(), std::string("dsv4_") + name + "_state_pos");
     inp.state_persist_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_persist_src_idxs.size(), std::string("dsv4_") + name + "_state_persist_src_idxs");
     inp.state_persist_dst_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_persist_dst_idxs.size(), std::string("dsv4_") + name + "_state_persist_dst_idxs");
+    inp.state_restore_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_restore_src_idxs.size(), std::string("dsv4_") + name + "_state_restore_src_idxs");
+    inp.state_restore_dst_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_restore_dst_idxs.size(), std::string("dsv4_") + name + "_state_restore_dst_idxs");
+    inp.state_snapshot_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_snapshot_src_idxs.size(), std::string("dsv4_") + name + "_state_snapshot_src_idxs");
+    inp.state_snapshot_dst_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_snapshot_dst_idxs.size(), std::string("dsv4_") + name + "_state_snapshot_dst_idxs");
     inp.state_read_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_read_idxs.size(), std::string("dsv4_") + name + "_state_read_idxs");
     inp.state_write_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I64, plan.state_write_idxs.size(), std::string("dsv4_") + name + "_state_write_idxs");
     inp.state_write_pos = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_write_pos.size(), std::string("dsv4_") + name + "_state_write_pos");
@@ -1195,7 +1264,7 @@ void llm_graph_result::reset() {
     t_embd_pooled = nullptr;
     t_h_nextn     = nullptr;
 
-    t_layer_inp.resize(LLAMA_MAX_LAYERS);
+    t_layer_inp.resize(LLAMA_MAX_LAYERS + 1);
     std::fill(t_layer_inp.begin(), t_layer_inp.end(), nullptr);
 
     t_sampled.clear();
@@ -1650,7 +1719,7 @@ ggml_tensor * llm_graph_context::build_ffn(
                         tmp = ggml_clamp(ctx0, tmp, -limit, limit);
                         cb(tmp, "ffn_up_clamped", il);
 
-                        if (arch == LLM_ARCH_DEEPSEEK4) {
+                        if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) {
                             cur = ggml_clamp(ctx0, cur, -INFINITY, limit);
                             cb(cur, "ffn_gate_clamped", il);
                             cur = ggml_swiglu_split(ctx0, cur, tmp);
@@ -2045,7 +2114,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
                         up = ggml_clamp(ctx0, up, -limit, limit);
                         cb(up, "ffn_moe_up_clamped", il);
 
-                        if (arch == LLM_ARCH_DEEPSEEK4) {
+                        if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) {
                             cur = ggml_clamp(ctx0, cur, -INFINITY, limit);
                             cb(cur, "ffn_moe_gate_clamped", il);
                             cur = ggml_swiglu_split(ctx0, cur, up);
@@ -2962,6 +3031,75 @@ ggml_tensor * llm_graph_context::build_attn(
     return cur;
 }
 
+ggml_tensor * llm_graph_context::build_attn(
+        llm_graph_input_attn_k_iswa * inp,
+        ggml_tensor * wo,
+        ggml_tensor * wo_b,
+        ggml_tensor * wo_s,
+        ggml_tensor * q_cur,
+        ggml_tensor * k_cur,
+        ggml_tensor * v_cur,
+        ggml_tensor * kq_b,
+        ggml_tensor * sinks,
+        ggml_tensor * v_mla,
+            float     kq_scale,
+            int       il) const {
+    const bool is_swa = hparams.is_swa(il);
+
+    GGML_UNUSED(v_cur);
+
+    auto * k_rot = is_swa ? inp->self_k_rot_swa : inp->self_k_rot;
+
+    if (k_rot) {
+        q_cur = llama_mul_mat_hadamard(ctx0, q_cur, k_rot);
+        if (k_cur) {
+            k_cur = llama_mul_mat_hadamard(ctx0, k_cur, k_rot);
+        }
+    }
+
+    // these nodes are added to the graph together so that they are not reordered
+    // by doing so, the number of splits in the graph is reduced
+    ggml_build_forward_expand(gf, q_cur);
+
+    if (k_cur) {
+        ggml_build_forward_expand(gf, k_cur);
+    }
+
+    const auto * mctx_iswa = inp->mctx;
+    const auto * mctx_cur = is_swa ? mctx_iswa->get_swa() : mctx_iswa->get_base();
+
+    // optionally store to KV cache
+    if (k_cur) {
+        const auto & k_idxs = is_swa ? inp->get_k_idxs_swa() : inp->get_k_idxs();
+
+        ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));
+    }
+
+    const auto & kq_mask = is_swa ? inp->get_kq_mask_swa() : inp->get_kq_mask();
+
+    // MLA-style attention: the cached K is used as V
+    ggml_tensor * q = q_cur;
+    ggml_tensor * k = mctx_cur->get_k(ctx0, il);
+    ggml_tensor * v = k;
+
+    ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il);
+    cb(cur, "kqv_out", il);
+
+    if (k_rot) {
+        cur = llama_mul_mat_hadamard(ctx0, cur, k_rot);
+    }
+
+    if (wo) {
+        cur = build_lora_mm(wo, cur, wo_s);
+    }
+
+    if (wo_b) {
+        cur = ggml_add(ctx0, cur, wo_b);
+    }
+
+    return cur;
+}
+
 llm_graph_input_attn_cross * llm_graph_context::build_attn_inp_cross() const {
     auto inp = std::make_unique<llm_graph_input_attn_cross>(cross);
 
@@ -3084,6 +3222,34 @@ llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const
     return (llm_graph_input_attn_kv_iswa *) res->add_input(std::move(inp));
 }
 
+llm_graph_input_attn_k_iswa * llm_graph_context::build_attn_inp_k_iswa() const {
+    const auto * mctx_cur = static_cast<const llama_kv_cache_iswa_context *>(mctx);
+
+    auto inp = std::make_unique<llm_graph_input_attn_k_iswa>(hparams, cparams, mctx_cur);
+
+    {
+        inp->self_k_idxs = mctx_cur->get_base()->build_input_k_idxs(ctx0, ubatch);
+
+        inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_base(), ubatch, cparams);
+        inp->self_kq_mask_cnv = inp->self_kq_mask;
+    }
+
+    {
+        GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache for non-SWA");
+
+        inp->self_k_idxs_swa = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch);
+
+        inp->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams);
+        inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa;
+    }
+
+    inp->self_k_rot = mctx_cur->get_base()->build_input_k_rot(ctx0);
+
+    inp->self_k_rot_swa = mctx_cur->get_swa()->build_input_k_rot(ctx0);
+
+    return (llm_graph_input_attn_k_iswa *) res->add_input(std::move(inp));
+}
+
 llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const {
     const auto * mctx_cur = static_cast<const llama_kv_cache_dsv4_context *>(mctx);
     const auto * raw_ctx  = mctx_cur->get_raw();

src/llama-graph.h

@@ -471,6 +471,45 @@ public:
     const llama_kv_cache_iswa_context * mctx;
 };
 
+class llm_graph_input_attn_k_iswa : public llm_graph_input_i {
+public:
+    llm_graph_input_attn_k_iswa(
+            const llama_hparams & hparams,
+            const llama_cparams & cparams,
+            const llama_kv_cache_iswa_context * mctx) :
+        hparams(hparams),
+        cparams(cparams),
+        mctx(mctx) {
+    }
+    ~llm_graph_input_attn_k_iswa() = default;
+
+    void set_input(const llama_ubatch * ubatch) override;
+
+    bool can_reuse(const llm_graph_params & params) override;
+
+    ggml_tensor * get_k_idxs()     const { return self_k_idxs; }
+    ggml_tensor * get_k_idxs_swa() const { return self_k_idxs_swa; }
+
+    ggml_tensor * get_kq_mask()     const { return self_kq_mask_cnv; }
+    ggml_tensor * get_kq_mask_swa() const { return self_kq_mask_swa_cnv; }
+
+    ggml_tensor * self_k_idxs     = nullptr; // I64 [n_batch]
+    ggml_tensor * self_k_idxs_swa = nullptr; // I64 [n_batch]
+
+    ggml_tensor * self_kq_mask         = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
+    ggml_tensor * self_kq_mask_cnv     = nullptr; //         [n_kv, n_batch/n_stream, 1, n_stream]
+    ggml_tensor * self_kq_mask_swa     = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
+    ggml_tensor * self_kq_mask_swa_cnv = nullptr; //         [n_kv, n_batch/n_stream, 1, n_stream]
+
+    ggml_tensor * self_k_rot = nullptr;
+    ggml_tensor * self_k_rot_swa = nullptr;
+
+    const llama_hparams hparams;
+    const llama_cparams cparams;
+
+    const llama_kv_cache_iswa_context * mctx;
+};
+
 // DSV4 raw graph inputs are SWA-only, but their mask may be stream-shaped
 // so raw K can be concatenated with DSV4 compressed K in one attention op.
 class llm_graph_input_dsv4_raw {
@@ -505,6 +544,10 @@ public:
         ggml_tensor * state_pos        = nullptr; // I32 [n_state]
         ggml_tensor * state_persist_src_idxs = nullptr; // I32 [n_state_persist]
         ggml_tensor * state_persist_dst_idxs = nullptr; // I32 [n_state_persist]
+        ggml_tensor * state_restore_src_idxs = nullptr; // I32 [n_state_restore]
+        ggml_tensor * state_restore_dst_idxs = nullptr; // I32 [n_state_restore]
+        ggml_tensor * state_snapshot_src_idxs = nullptr; // I32 [n_state_snapshot]
+        ggml_tensor * state_snapshot_dst_idxs = nullptr; // I32 [n_state_snapshot]
         ggml_tensor * state_read_idxs  = nullptr; // I32 [ratio*n_state_write]
         ggml_tensor * state_write_idxs = nullptr; // I64 [n_state_write]
         ggml_tensor * state_write_pos  = nullptr; // I32 [n_state_write]
@@ -1068,7 +1111,7 @@ struct llm_graph_context {
     ggml_tensor * build_attn_mha(
             ggml_tensor * q,       // [n_embd_head_q, n_head_q, n_tokens]
             ggml_tensor * k,       // [n_embd_head_k, n_head_k, n_tokens]
-            ggml_tensor * v,       // [n_embd_head_v, n_head_v, n_tokens] (v_trans == false)
+            ggml_tensor * v,       // [n_embd_head_v, n_head_v, n_tokens] (v_trans = false)
             ggml_tensor * kq_b,
             ggml_tensor * kq_mask,
             ggml_tensor * sinks,   // [n_head_q]
@@ -1160,6 +1203,24 @@ struct llm_graph_context {
                   float   kq_scale,
                     int   il) const;
 
+    llm_graph_input_attn_k_iswa * build_attn_inp_k_iswa() const;
+
+    // note: if k_cur is not provided, it will not be stored in the memory
+    // note: the K cache is used as V (MLA-style attention)
+    ggml_tensor * build_attn(
+            llm_graph_input_attn_k_iswa * inp,
+            ggml_tensor * wo,
+            ggml_tensor * wo_b,
+            ggml_tensor * wo_s,
+            ggml_tensor * q_cur, // [n_embd_head_q, n_head_q, n_tokens]
+            ggml_tensor * k_cur, // [n_embd_head_k, n_head_k, n_tokens] optional
+            ggml_tensor * v_cur, // [n_embd_head_v, n_head_v, n_tokens] optional
+            ggml_tensor * kq_b,
+            ggml_tensor * sinks, // [n_head_q]
+            ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v]
+                  float   kq_scale,
+                    int   il) const;
+
     llm_graph_input_attn_cross * build_attn_inp_cross() const;
 
     ggml_tensor * build_attn(

src/llama-kv-cache-dsv4.cpp

@@ -252,7 +252,8 @@ static void dsv4_state_write_tensor_streams(
         uint32_t           tensor_rows,
         uint32_t           n_rows,
         uint32_t           s0,
-        uint32_t           ns) {
+        uint32_t           ns,
+        const std::vector<uint32_t> * stream_ids = nullptr) {
     const int32_t  type_i   = (int32_t) tensor->type;
     const uint64_t ne0      = tensor->ne[0];
     const uint64_t rows     = n_rows;
@@ -273,8 +274,16 @@ static void dsv4_state_write_tensor_streams(
         return;
     }
 
+    if (stream_ids && stream_ids->size() != ns) {
+        throw std::runtime_error("DSV4 state tensor stream map size mismatch");
+    }
+
     for (uint32_t s = 0; s < ns; ++s) {
-        const size_t offset = (size_t) (s0 + s)*stream_stride;
+        const uint32_t stream = stream_ids ? (*stream_ids)[s] : s0 + s;
+        if ((int64_t) stream >= tensor->ne[2]) {
+            throw std::runtime_error("DSV4 state tensor stream out of range");
+        }
+        const size_t offset = (size_t) stream*stream_stride;
         io.write_tensor(tensor, offset, size);
     }
 }
@@ -421,7 +430,9 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
         bool overlap,
         uint32_t state_size,
         uint32_t kv_size,
-        uint32_t n_stream) {
+        uint32_t n_stream,
+        uint32_t n_rs_seq,
+        const std::vector<uint32_t> & rs_idx) {
     llama_kv_cache_dsv4_context::comp_plan plan;
     plan.n_visible.resize(ubatch.n_tokens);
     plan.n_stream = dsv4_comp_graph_n_stream(ubatch, n_stream);
@@ -451,6 +462,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
     std::vector<int32_t> overlap_cur_reads;
 
     std::map<std::pair<llama_seq_id, llama_pos>, int64_t> curr_token_idx_map;
+    std::map<llama_seq_id, uint32_t> state_write_counts;
 
     for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
         for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) {
@@ -513,6 +525,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
 
             plan.state_write_idxs.push_back(cache_off + pos/ratio);
             plan.state_write_pos.push_back((int32_t) source_start);
+            ++state_write_counts[seq_id];
 
             if (overlap) {
                 const llama_pos prev_start = source_start - ratio;
@@ -531,33 +544,57 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
         }
     }
 
-    if (ratio == DSV4_CSA_RATIO && plan.state_write_idxs.empty() && !plan.state_pos.empty()) {
-        // Non-boundary CSA steps still need a write op so their graph matches
-        // boundary steps. Use a padded scratch row that is masked from attention.
+    if (ratio == DSV4_CSA_RATIO && !plan.state_pos.empty()) {
         assert(kv_size > 0);
 
-        uint32_t i = 0;
-        while (i < ubatch.n_tokens && ubatch.pos[i] < 0) {
-            ++i;
-        }
-        assert(i < ubatch.n_tokens);
+        // Pad each stream to the reserve plan's block count.
+        const auto append_dummy_block = [&](llama_seq_id seq_id, uint32_t i) {
+            const int64_t cache_off = dsv4_stream_offset(n_stream, seq_id, kv_size);
+            const int32_t source_idx = state_source_idx(seq_id, ubatch.pos[i]);
 
-        const llama_pos    pos    = ubatch.pos[i];
-        const llama_seq_id seq_id = ubatch.seq_id[i][0];
-        const int64_t cache_off = dsv4_stream_offset(n_stream, seq_id, kv_size);
-        const int32_t source_idx = state_source_idx(seq_id, pos);
+            plan.state_write_idxs.push_back(cache_off + kv_size - 1);
+            plan.state_write_pos .push_back(0);
 
-        plan.state_write_idxs.push_back(cache_off + kv_size - 1);
-        plan.state_write_pos .push_back(0);
+            if (overlap) {
+                for (uint32_t j = 0; j < ratio; ++j) {
+                    overlap_prev_reads.push_back(source_idx);
+                    overlap_cur_reads .push_back(source_idx);
+                }
+            } else {
+                for (uint32_t j = 0; j < ratio; ++j) {
+                    plan.state_read_idxs.push_back(source_idx);
+                }
+            }
+        };
 
-        if (overlap) {
-            for (uint32_t j = 0; j < ratio; ++j) {
-                overlap_prev_reads.push_back(source_idx);
-                overlap_cur_reads .push_back(source_idx);
+        if (dsv4_ubatch_has_coupled(ubatch)) {
+            if (plan.state_write_idxs.empty()) {
+                uint32_t i = 0;
+                while (i < ubatch.n_tokens && ubatch.pos[i] < 0) {
+                    ++i;
+                }
+                assert(i < ubatch.n_tokens);
+                append_dummy_block(ubatch.seq_id[i][0], i);
             }
         } else {
-            for (uint32_t j = 0; j < ratio; ++j) {
-                plan.state_read_idxs.push_back(source_idx);
+            const uint32_t n_blocks = (std::max<uint32_t>(1, ubatch.n_seq_tokens) + ratio - 1)/ratio;
+
+            for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
+                const llama_seq_id seq_id = ubatch.seq_id_unq[s];
+                const uint32_t n_writes = state_write_counts[seq_id];
+                if (n_writes >= n_blocks) {
+                    continue;
+                }
+                if (n_writes + 1 != n_blocks) {
+                    throw std::runtime_error("DSV4 CSA sequence positions are not contiguous");
+                }
+
+                uint32_t i = 0;
+                while (i < ubatch.n_tokens && (ubatch.pos[i] < 0 || !dsv4_token_has_seq(ubatch, i, seq_id))) {
+                    ++i;
+                }
+                assert(i < ubatch.n_tokens);
+                append_dummy_block(seq_id, i);
             }
         }
     }
@@ -583,6 +620,63 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
         plan.state_persist_dst_idxs.push_back(row.dst);
     }
 
+
+    if (n_rs_seq > 0) {
+        for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
+            const llama_seq_id seq_id = ubatch.seq_id_unq[s];
+            if (seq_id < 0 || (uint32_t) seq_id >= n_stream) {
+                continue;
+            }
+
+            const int64_t stream_off = dsv4_stream_offset(n_stream, seq_id, state_size);
+            const uint32_t rollback = (uint32_t) seq_id < rs_idx.size() ? rs_idx[seq_id] : 0;
+            // Keep the restore graph fixed-width when no rollback is pending.
+            const int64_t src_plane = rollback > 0 && rollback <= n_rs_seq ? (int64_t) rollback*state_rows : 0;
+            for (uint32_t r = 0; r < state_size; ++r) {
+                plan.state_restore_src_idxs.push_back((int32_t) (src_plane + stream_off + r));
+                plan.state_restore_dst_idxs.push_back((int32_t) (stream_off + r));
+            }
+
+            std::vector<uint32_t> token_idxs;
+            token_idxs.reserve(ubatch.n_tokens);
+            for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
+                if (dsv4_token_has_seq(ubatch, i, seq_id)) {
+                    token_idxs.push_back(i);
+                }
+            }
+            if (token_idxs.empty()) {
+                continue;
+            }
+
+            const uint32_t n_seq_tokens = (uint32_t) token_idxs.size();
+            const int64_t scratch_off = (int64_t) state_rows*(1 + n_rs_seq);
+            for (uint32_t d = 1; d <= n_rs_seq; ++d) {
+                const int64_t dst_plane = (int64_t) d*state_rows;
+
+                for (uint32_t r = 0; r < state_size; ++r) {
+                    int32_t src;
+                    if (d <= n_seq_tokens) {
+                        const uint32_t prefix = n_seq_tokens - d;
+                        src = (int32_t) (stream_off + r);
+
+                        for (uint32_t j = 0; j < prefix; ++j) {
+                            const uint32_t i_tok = token_idxs[j];
+                            if (ubatch.pos[i_tok] >= 0 && (uint32_t) (ubatch.pos[i_tok]%state_size) == r) {
+                                src = (int32_t) (scratch_off + i_tok);
+                            }
+                        }
+                    } else {
+                        const int64_t src_plane = (int64_t) (d - n_seq_tokens)*state_rows;
+                        src = (int32_t) (src_plane + stream_off + r);
+                    }
+
+                    plan.state_snapshot_src_idxs.push_back(src);
+                    plan.state_snapshot_dst_idxs.push_back((int32_t) (dst_plane + stream_off + r));
+                }
+            }
+        }
+    }
+
     static const bool debug = []() {
         const char * env = getenv("LLAMA_DSV4_COMPRESS_DEBUG");
         return env && atoi(env) > 0;
@@ -604,12 +698,14 @@ static std::vector<llama_kv_cache_dsv4_context::comp_plan> dsv4_build_comp_plans
         bool overlap,
         uint32_t state_size,
         uint32_t kv_size,
-        uint32_t n_stream) {
+        uint32_t n_stream,
+        uint32_t n_rs_seq,
+        const std::vector<uint32_t> & rs_idx) {
     std::vector<llama_kv_cache_dsv4_context::comp_plan> plans;
     plans.reserve(ubatches.size());
 
     for (const llama_ubatch & ubatch : ubatches) {
-        plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream));
+        plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream, n_rs_seq, rs_idx));
     }
 
     return plans;
@@ -696,7 +792,8 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_reserve_comp_plan(
         bool overlap,
         uint32_t state_size,
         uint32_t kv_size,
-        uint32_t n_stream) {
+        uint32_t n_stream,
+        uint32_t n_rs_seq) {
     llama_kv_cache_dsv4_context::comp_plan plan;
     plan.n_visible.resize(ubatch.n_tokens);
     plan.n_stream = dsv4_comp_graph_n_stream(ubatch, n_stream);
@@ -714,10 +811,16 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_reserve_comp_plan(
 
     const uint64_t state_rows = (uint64_t) state_size*n_stream;
     const size_t n_persist = (size_t) std::min<uint64_t>(ubatch.n_tokens, state_rows);
+    const size_t n_restore = n_rs_seq > 0 ? (size_t) state_size*std::max<uint32_t>(1, ubatch.n_seqs_unq) : 0;
+    const size_t n_snapshot = (size_t) n_rs_seq*state_size*std::max<uint32_t>(1, ubatch.n_seqs_unq);
 
     plan.state_pos .resize(ubatch.n_tokens);
     plan.state_persist_src_idxs.resize(n_persist);
     plan.state_persist_dst_idxs.resize(n_persist);
+    plan.state_restore_src_idxs.resize(n_restore);
+    plan.state_restore_dst_idxs.resize(n_restore);
+    plan.state_snapshot_src_idxs.resize(n_snapshot);
+    plan.state_snapshot_dst_idxs.resize(n_snapshot);
     plan.state_read_idxs .resize((overlap ? 2u : 1u)*ratio*n_blocks);
     plan.state_write_idxs.resize(n_blocks);
     plan.state_write_pos .resize(n_blocks);
@@ -743,12 +846,14 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
             uint32_t        ratio,
             uint32_t        state_size,
             uint32_t        n_embd_state,
+            uint32_t        n_rs_seq,
         const char    * name,
         const llama_memory_i::layer_filter_cb & filter) :
     ratio(ratio),
     state_size(state_size),
     n_embd_state(n_embd_state),
-    n_stream(unified ? 1 : n_seq_max) {
+    n_stream(unified ? 1 : n_seq_max),
+    n_rs_seq(n_rs_seq) {
     const llama_hparams & hparams = model.hparams;
 
     struct ggml_backend_buft_comparator {
@@ -804,8 +909,9 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
             throw std::runtime_error("failed to create ggml context for DSV4 compressor state");
         }
 
-        ggml_tensor * kv    = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_stream);
-        ggml_tensor * score = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_stream);
+        const uint32_t n_planes = n_stream*(1 + n_rs_seq);
+        ggml_tensor * kv    = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_planes);
+        ggml_tensor * score = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_planes);
 
         ggml_format_name(kv,    "dsv4_%s_state_kv_l%d",    name, il);
         ggml_format_name(score, "dsv4_%s_state_score_l%d", name, il);
@@ -837,8 +943,8 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
         ctxs_bufs.emplace_back(std::move(ctx), buf);
     }
 
-    LLAMA_LOG_INFO("%s: %s ratio = %u, state = %u x %u, streams = %u, layers = %zu, size = %7.2f MiB\n",
-            __func__, name, ratio, state_size, n_embd_state, n_stream, layers.size(), total_size()/1024.0/1024.0);
+    LLAMA_LOG_INFO("%s: %s ratio = %u, state = %u x %u, streams = %u, rs_seq = %u, layers = %zu, size = %7.2f MiB\n",
+            __func__, name, ratio, state_size, n_embd_state, n_stream, n_rs_seq, layers.size(), total_size()/1024.0/1024.0);
 }
 
 void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) {
@@ -848,9 +954,13 @@ void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) {
 
     if (seq_id >= 0) {
         GGML_ASSERT((uint32_t) seq_id < n_stream);
+
         for (const auto & layer : layers) {
-            dsv4_clear_tensor_stream(layer.kv,    (uint32_t) seq_id);
-            dsv4_clear_tensor_stream(layer.score, (uint32_t) seq_id);
+            for (uint32_t d = 0; d <= n_rs_seq; ++d) {
+                const uint32_t stream = d*n_stream + (uint32_t) seq_id;
+                dsv4_clear_tensor_stream(layer.kv,    stream);
+                dsv4_clear_tensor_stream(layer.score, stream);
+            }
         }
         return;
     }
@@ -868,6 +978,8 @@ void llama_dsv4_comp_state::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_
         return;
     }
 
+    clear(seq_id_dst, true);
+
     sc_info.ssrc.push_back((uint32_t) seq_id_src);
     sc_info.sdst.push_back((uint32_t) seq_id_dst);
 }
@@ -896,6 +1008,14 @@ uint32_t llama_dsv4_comp_state::get_n_stream() const {
     return n_stream;
 }
 
+uint32_t llama_dsv4_comp_state::get_n_rs_seq() const {
+    return n_rs_seq;
+}
+
+uint32_t llama_dsv4_comp_state::get_n_rows() const {
+    return state_size*n_stream;
+}
+
 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) {
@@ -905,13 +1025,26 @@ std::map<ggml_backend_buffer_type_t, size_t> llama_dsv4_comp_state::memory_break
     return ret;
 }
 
-void llama_dsv4_comp_state::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
+void llama_dsv4_comp_state::state_write(
+        llama_io_write_i & io,
+        llama_seq_id seq_id,
+        llama_state_seq_flags flags,
+        const std::vector<uint32_t> & rs_idx) const {
     GGML_UNUSED(flags);
 
     uint32_t s0;
     uint32_t ns;
     dsv4_state_src_stream_range(n_stream, seq_id, s0, ns);
 
+    std::vector<uint32_t> stream_ids(ns);
+    for (uint32_t s = 0; s < ns; ++s) {
+        const uint32_t seq = seq_id >= 0 ? (uint32_t) seq_id : s0 + s;
+        if (seq >= rs_idx.size() || rs_idx[seq] > n_rs_seq) {
+            throw std::runtime_error("DSV4 recurrent state rollback index out of range");
+        }
+        stream_ids[s] = rs_idx[seq]*n_stream + s0 + s;
+    }
+
     const uint32_t version      = DSV4_COMP_STATE_VER;
     const uint32_t n_layer      = layers.size();
 
@@ -925,8 +1058,8 @@ void llama_dsv4_comp_state::state_write(llama_io_write_i & io, llama_seq_id seq_
     for (const auto & layer : layers) {
         io.write(&layer.il, sizeof(layer.il));
 
-        dsv4_state_write_tensor_streams(io, layer.kv,    state_size, state_size, s0, ns);
-        dsv4_state_write_tensor_streams(io, layer.score, state_size, state_size, s0, ns);
+        dsv4_state_write_tensor_streams(io, layer.kv,    state_size, state_size, s0, ns, &stream_ids);
+        dsv4_state_write_tensor_streams(io, layer.score, state_size, state_size, s0, ns, &stream_ids);
     }
 }
 
@@ -972,28 +1105,40 @@ void llama_dsv4_comp_state::state_read(llama_io_read_i & io, llama_seq_id seq_id
     }
 }
 
-ggml_tensor * llama_dsv4_comp_state::get_kv(ggml_context * ctx, int32_t il) const {
+ggml_tensor * llama_dsv4_comp_state::get_kv_all(ggml_context * ctx, int32_t il) const {
     const int32_t ids = map_layer_ids.at(il);
-
     ggml_tensor * state = layers[ids].kv;
 
-    return ggml_reshape_2d(ctx, state, state->ne[0], state->ne[1]*state->ne[2]);
+    return ggml_view_2d(ctx, state, state->ne[0], get_n_rows()*(1 + n_rs_seq), state->nb[1], 0);
+}
+
+ggml_tensor * llama_dsv4_comp_state::get_score_all(ggml_context * ctx, int32_t il) const {
+    const int32_t ids = map_layer_ids.at(il);
+    ggml_tensor * state = layers[ids].score;
+
+    return ggml_view_2d(ctx, state, state->ne[0], get_n_rows()*(1 + n_rs_seq), state->nb[1], 0);
+}
+
+ggml_tensor * llama_dsv4_comp_state::get_kv(ggml_context * ctx, int32_t il) const {
+    ggml_tensor * state = get_kv_all(ctx, il);
+    const size_t row_size = ggml_row_size(state->type, state->ne[0]);
+
+    return ggml_view_2d(ctx, state, state->ne[0], get_n_rows(), state->nb[1], 0*row_size);
 }
 
 ggml_tensor * llama_dsv4_comp_state::get_score(ggml_context * ctx, int32_t il) const {
-    const int32_t ids = map_layer_ids.at(il);
+    ggml_tensor * state = get_score_all(ctx, il);
+    const size_t row_size = ggml_row_size(state->type, state->ne[0]);
 
-    ggml_tensor * state = layers[ids].score;
-
-    return ggml_reshape_2d(ctx, state, state->ne[0], state->ne[1]*state->ne[2]);
+    return ggml_view_2d(ctx, state, state->ne[0], get_n_rows(), state->nb[1], 0*row_size);
 }
 
 ggml_tensor * llama_dsv4_comp_state::cpy_kv(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const {
-    return ggml_set_rows(ctx, get_kv(ctx, il), cur, idxs);
+    return ggml_set_rows(ctx, get_kv_all(ctx, il), cur, idxs);
 }
 
 ggml_tensor * llama_dsv4_comp_state::cpy_score(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const {
-    return ggml_set_rows(ctx, get_score(ctx, il), cur, idxs);
+    return ggml_set_rows(ctx, get_score_all(ctx, il), cur, idxs);
 }
 
 size_t llama_dsv4_comp_state::total_size() const {
@@ -1022,13 +1167,16 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
                  uint32_t   n_seq_max,
                  uint32_t   n_ubatch,
                  uint32_t   n_pad,
+                 uint32_t   n_rs_seq,
     const layer_filter_cb & filter,
     const  layer_reuse_cb & reuse) :
     hparams_raw(model.hparams),
     hparams_csa(model.hparams),
     hparams_hca(model.hparams),
     hparams_lid(model.hparams),
-    n_seq_max(n_seq_max) {
+    n_seq_max(n_seq_max),
+    n_rs_seq(n_rs_seq),
+    rs_idx(n_seq_max, 0) {
 
     const layer_filter_cb filter_raw = [&](int32_t il) {
         if (filter && !filter(il)) {
@@ -1043,6 +1191,11 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
     // Keep DSV4 KV/state streams per sequence even when public KV mode is unified.
     const bool unified_raw = false;
 
+    hparams_raw.n_layer_nextn = 0;
+    hparams_csa.n_layer_nextn = 0;
+    hparams_hca.n_layer_nextn = 0;
+    hparams_lid.n_layer_nextn = 0;
+
     LLAMA_LOG_INFO("%s: creating DSV4 raw KV cache\n", __func__);
 
     dsv4_make_k_only(hparams_raw);
@@ -1109,19 +1262,19 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
 
     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(), "csa", filter_csa);
+            2*model.hparams.n_embd_head_k(), n_rs_seq, "csa", filter_csa);
 
     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(), "hca", filter_hca);
+            model.hparams.n_embd_head_k(), n_rs_seq, "hca", filter_hca);
 
     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, "lid", filter_csa);
+            2*model.hparams.indexer_head_size, n_rs_seq, "lid", filter_csa);
 
     // DSV4 attention reads compressed-K / compressor-state rows that the current
     // graph does not necessarily overwrite; uninitialized buffer contents would
@@ -1255,17 +1408,35 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
     }
 
     if (p0 > 0) {
-        if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max ||
-                p0 <= kv_raw->seq_pos_max(seq_id)) {
+        if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) {
             return false;
         }
 
-        bool res = true;
+        const llama_pos pos_max = kv_raw->seq_pos_max(seq_id);
+        if (p0 > pos_max) {
+            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_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);
+
+            return res;
+        }
+
+        if (n_rs_seq == 0) {
+            return false;
+        }
+
+        const llama_pos rollback = pos_max - (p0 - 1);
+        if (rollback < 1 || rollback > (llama_pos) n_rs_seq) {
+            return false;
+        }
+
+        const bool res = kv_raw->seq_rm(seq_id, p0, p1);
+        if (res) {
+            rs_idx[seq_id] = (uint32_t) rollback;
+        }
 
         return res;
     }
@@ -1290,6 +1461,10 @@ void llama_kv_cache_dsv4::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_ds
     csa_state->seq_cp(seq_id_src, seq_id_dst);
     hca_state->seq_cp(seq_id_src, seq_id_dst);
     lid_state->seq_cp(seq_id_src, seq_id_dst);
+
+    if (seq_id_src != seq_id_dst) {
+        rs_idx[seq_id_dst] = 0;
+    }
 }
 
 void llama_kv_cache_dsv4::seq_keep(llama_seq_id seq_id) {
@@ -1386,9 +1561,9 @@ void llama_kv_cache_dsv4::state_write(llama_io_write_i & io, llama_seq_id seq_id
         dsv4_state_write_k_cache(io, kv_lid.get(), seq_id, flags, n_rows_lid);
     }
 
-    csa_state->state_write(io, seq_id, flags);
-    hca_state->state_write(io, seq_id, flags);
-    lid_state->state_write(io, seq_id, flags);
+    csa_state->state_write(io, seq_id, flags, rs_idx);
+    hca_state->state_write(io, seq_id, flags, rs_idx);
+    lid_state->state_write(io, seq_id, flags, rs_idx);
 }
 
 void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
@@ -1432,6 +1607,12 @@ void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id,
     hca_state->state_read(io, seq_id, flags);
     lid_state->state_read(io, seq_id, flags);
 
+    if (seq_id >= 0) {
+        GGML_ASSERT((uint32_t) seq_id < n_seq_max);
+        rs_idx[seq_id] = 0;
+    } else {
+        std::fill(rs_idx.begin(), rs_idx.end(), 0);
+    }
 }
 
 llama_kv_cache_iswa * llama_kv_cache_dsv4::get_raw() const {
@@ -1462,6 +1643,31 @@ llama_dsv4_comp_state * llama_kv_cache_dsv4::get_lid_state() const {
     return lid_state.get();
 }
 
+uint32_t llama_kv_cache_dsv4::get_n_rs_seq() const {
+    return n_rs_seq;
+}
+
+const std::vector<uint32_t> & llama_kv_cache_dsv4::get_rs_idx() const {
+    return rs_idx;
+}
+
+void llama_kv_cache_dsv4::reset_rs_idx_for_ubatches(const std::vector<llama_ubatch> & ubatches) {
+    if (n_rs_seq == 0) {
+        return;
+    }
+
+    for (const llama_ubatch & ubatch : ubatches) {
+        for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
+            for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) {
+                const llama_seq_id seq_id = ubatch.seq_id[i][s];
+                if (seq_id >= 0 && (uint32_t) seq_id < n_seq_max) {
+                    rs_idx[seq_id] = 0;
+                }
+            }
+        }
+    }
+}
+
 void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) {
     if (seq_id < 0) {
         kv_csa->clear(data);
@@ -1488,6 +1694,12 @@ void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) {
     csa_state->clear(seq_id, data);
     hca_state->clear(seq_id, data);
     lid_state->clear(seq_id, data);
+
+    if (seq_id >= 0) {
+        rs_idx[seq_id] = 0;
+    } else {
+        std::fill(rs_idx.begin(), rs_idx.end(), 0);
+    }
 }
 
 //
@@ -1779,10 +1991,14 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
         std::vector<llama_ubatch> ubatches_raw) :
     ubatches(std::move(ubatches)),
     plans_csa(dsv4_build_comp_plans(this->ubatches, DSV4_CSA_RATIO, true,
-                kv->get_csa_state()->get_state_size(), kv->get_csa()->get_size(), kv->get_csa_state()->get_n_stream())),
+                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,
-                kv->get_hca_state()->get_state_size(), kv->get_hca()->get_size(), kv->get_hca_state()->get_n_stream())),
-    plans_lid(plans_csa),
+                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,
+                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>(
                 kv->get_raw(),
                 std::move(sinfos_raw_base_write),
@@ -1809,6 +2025,7 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
     hca_state(kv->get_hca_state()),
     lid_state(kv->get_lid_state()),
     status(ctx_raw->get_status()) {
+    kv->reset_rs_idx_for_ubatches(this->ubatches);
 }
 
 llama_kv_cache_dsv4_context::~llama_kv_cache_dsv4_context() = default;
@@ -1944,7 +2161,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,
-            csa_state->get_state_size(), get_csa()->get_n_kv(), csa_state->get_n_stream());
+            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;
 }
@@ -1958,7 +2175,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,
-            hca_state->get_state_size(), get_hca()->get_n_kv(), hca_state->get_n_stream());
+            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;
 }
@@ -1972,7 +2189,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,
-            lid_state->get_state_size(), get_lid()->get_n_kv(), lid_state->get_n_stream());
+            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

@@ -22,6 +22,7 @@ public:
             uint32_t        ratio,
             uint32_t        state_size,
             uint32_t        n_embd_state,
+            uint32_t        n_rs_seq,
             const char    * name,
         const llama_memory_i::layer_filter_cb & filter);
 
@@ -29,17 +30,21 @@ public:
     void seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst);
     void apply_copies(const stream_copy_info & sc_info) const;
 
-    uint32_t get_ratio()    const;
+    uint32_t get_ratio()      const;
     uint32_t get_state_size() const;
-    uint32_t get_n_stream() const;
+    uint32_t get_n_stream()   const;
+    uint32_t get_n_rs_seq()   const;
+    uint32_t get_n_rows()     const;
 
     std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const;
 
-    void state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const;
+    void state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags, const std::vector<uint32_t> & rs_idx) const;
     void state_read (llama_io_read_i  & io, llama_seq_id seq_id, llama_state_seq_flags flags);
 
-    ggml_tensor * get_kv   (ggml_context * ctx, int32_t il) const;
-    ggml_tensor * get_score(ggml_context * ctx, int32_t il) const;
+    ggml_tensor * get_kv       (ggml_context * ctx, int32_t il) const;
+    ggml_tensor * get_score    (ggml_context * ctx, int32_t il) const;
+    ggml_tensor * get_kv_all   (ggml_context * ctx, int32_t il) const;
+    ggml_tensor * get_score_all(ggml_context * ctx, int32_t il) const;
 
     ggml_tensor * cpy_kv   (ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const;
     ggml_tensor * cpy_score(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const;
@@ -59,6 +64,7 @@ private:
     const uint32_t state_size;
     const uint32_t n_embd_state;
     const uint32_t n_stream;
+    const uint32_t n_rs_seq;
 
     std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> ctxs_bufs;
 
@@ -93,6 +99,7 @@ public:
                      uint32_t   n_seq_max,
                      uint32_t   n_ubatch,
                      uint32_t   n_pad,
+                     uint32_t   n_rs_seq,
         const layer_filter_cb & filter,
         const  layer_reuse_cb & reuse);
 
@@ -141,6 +148,10 @@ public:
     llama_dsv4_comp_state * get_hca_state() const;
     llama_dsv4_comp_state * get_lid_state() 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);
+
 private:
     llama_hparams hparams_raw;
     llama_hparams hparams_csa;
@@ -148,6 +159,9 @@ private:
     llama_hparams hparams_lid;
 
     const uint32_t n_seq_max;
+    const uint32_t n_rs_seq;
+
+    std::vector<uint32_t> rs_idx;
 
     std::unique_ptr<llama_kv_cache_iswa> kv_raw;
     std::unique_ptr<llama_kv_cache>      kv_csa;
@@ -268,6 +282,17 @@ public:
         std::vector<int32_t> state_persist_src_idxs;
         std::vector<int32_t> state_persist_dst_idxs;
 
+        // Device-side rollback restore copies snapshot planes back to the
+        // current compressor-state plane before the graph reads it.
+        std::vector<int32_t> state_restore_src_idxs;
+        std::vector<int32_t> state_restore_dst_idxs;
+
+        // Device-side rollback snapshots copy rows from the graph-local
+        // [persistent_state | current_ubatch_scratch] tensor into rollback
+        // planes after the graph has computed current-token compressor state.
+        std::vector<int32_t> state_snapshot_src_idxs;
+        std::vector<int32_t> state_snapshot_dst_idxs;
+
         // Flattened source row ids used for state-backed commits. Source rows
         // index the graph-local [persistent_state | current_ubatch_scratch]
         // tensor. For overlapped compression the first half is previous rows

src/llama-model-loader.cpp

@@ -526,6 +526,7 @@ llama_model_loader::llama_model_loader(
         llama_load_mode load_mode,
         bool check_tensors,
         bool no_alloc,
+        bool load_mtp,
         const llama_model_kv_override * param_overrides_p,
         const llama_model_tensor_buft_override * param_tensor_buft_overrides_p)
         : metadata(meta), set_tensor_data(set_tensor_data), set_tensor_data_ud(set_tensor_data_ud) {
@@ -812,6 +813,7 @@ llama_model_loader::llama_model_loader(
 
     this->check_tensors = check_tensors;
     this->no_alloc = no_alloc;
+    this->load_mtp = load_mtp;
 }
 
 std::string llama_model_loader::get_arch_name() const {

src/llama-model-loader.h

@@ -79,6 +79,7 @@ struct llama_model_loader {
     bool use_direct_io = false;
     bool check_tensors;
     bool no_alloc;
+    bool load_mtp;
 
     llama_files files;
     llama_ftype ftype;
@@ -129,6 +130,7 @@ struct llama_model_loader {
         llama_load_mode load_mode,
         bool check_tensors,
         bool no_alloc,
+        bool load_mtp,
         const llama_model_kv_override * param_overrides_p,
         const llama_model_tensor_buft_override * param_tensor_buft_overrides_p);
 

src/llama-model.cpp

@@ -2138,6 +2138,75 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                             nullptr);
                 }
             } break;
+        case LLM_ARCH_DEEPSEEK4:
+            {
+                GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE);
+
+                if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) {
+                    const llama_memory_i::layer_filter_cb filter_mtp = [&](int32_t il) {
+                        return il >= (int32_t) hparams.n_layer();
+                    };
+
+                    res = new llama_kv_cache_iswa(
+                            *this,
+                            params.type_k,
+                            params.type_v,
+                            !cparams.flash_attn,
+                            cparams.offload_kqv,
+                            params.swa_full,
+                            cparams.kv_unified,
+                            cparams.n_ctx_seq,
+                            cparams.n_seq_max,
+                            cparams.n_ubatch,
+                            1,
+                            nullptr,
+                            filter_mtp,
+                            nullptr,
+                            nullptr);
+                } else {
+                    res = new llama_kv_cache_dsv4(
+                            *this,
+                            params.type_k,
+                            params.type_v,
+                            !cparams.flash_attn,
+                            cparams.offload_kqv,
+                            params.swa_full,
+                            cparams.kv_unified,
+                            cparams.n_ctx_seq,
+                            cparams.n_seq_max,
+                            cparams.n_ubatch,
+                            1,
+                            cparams.n_rs_seq,
+                            nullptr,
+                            nullptr);
+                }
+            } break;
+        case LLM_ARCH_DFLASH:
+            {
+                // DSV4 DSpark stages store a single MLA-style K per position (window = the draft ring)
+                if (hparams.dsv4_hc_mult > 0) {
+                    GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE);
+
+                    res = new llama_kv_cache_iswa(
+                            *this,
+                            params.type_k,
+                            params.type_v,
+                            !cparams.flash_attn,
+                            cparams.offload_kqv,
+                            params.swa_full,
+                            cparams.kv_unified,
+                            cparams.n_ctx_seq,
+                            cparams.n_seq_max,
+                            cparams.n_ubatch,
+                            1,
+                            nullptr,
+                            nullptr,
+                            nullptr,
+                            nullptr);
+                    break;
+                }
+            }
+            [[fallthrough]];
         // Models that need standard caching should rely on recurrent/hybrid
         // checks
         default:
@@ -2253,24 +2322,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                         }
                     }
 
-                    if (arch == LLM_ARCH_DEEPSEEK4) {
-                        GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE);
-
-                        res = new llama_kv_cache_dsv4(
-                                *this,
-                                params.type_k,
-                                params.type_v,
-                                !cparams.flash_attn,
-                                cparams.offload_kqv,
-                                params.swa_full,
-                                cparams.kv_unified,
-                                cparams.n_ctx_seq,
-                                cparams.n_seq_max,
-                                cparams.n_ubatch,
-                                1,
-                                filter,
-                                reuse);
-                    } else if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {
+                    if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {
                         GGML_ASSERT(hparams.is_swa_any());
 
                         if (arch == LLM_ARCH_GEMMA4_ASSISTANT) {
@@ -2390,6 +2442,7 @@ llama_model_params llama_model_default_params() {
         /*.use_extra_bufts             =*/ true,
         /*.no_host                     =*/ false,
         /*.no_alloc                    =*/ false,
+        /*.load_mtp                    =*/ false,
     };
 
     return result;
@@ -2618,9 +2671,12 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
         case LLM_ARCH_STEP35:
         case LLM_ARCH_TALKIE:
         case LLM_ARCH_MELLUM:
-        case LLM_ARCH_DFLASH:
             return LLAMA_ROPE_TYPE_NEOX;
 
+        case LLM_ARCH_DFLASH:
+            // DSV4 DSpark drafters use DeepSeek-V4's normal RoPE; legacy DFlash backbones are NeoX
+            return model->hparams.dsv4_hc_mult > 0 ? LLAMA_ROPE_TYPE_NORM : LLAMA_ROPE_TYPE_NEOX;
+
         case LLM_ARCH_QWEN2VL:
         case LLM_ARCH_PADDLEOCR:
             return LLAMA_ROPE_TYPE_MROPE;

src/llama-quant.cpp

@@ -893,7 +893,7 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
     const llama_model_kv_override * kv_overrides = params->kv_overrides;
     std::vector<std::string> splits = {};
     llama_model_loader ml(/*metadata*/ nullptr, /*set_tensor_data*/ nullptr, /*set_tensor_data_ud*/ nullptr,
-        fname_inp, splits, /*file*/ nullptr, /*load_mode*/ load_mode, /*check_tensors*/ true, /*no_alloc*/ false, kv_overrides, nullptr);
+        fname_inp, splits, /*file*/ nullptr, /*load_mode*/ load_mode, /*check_tensors*/ true, /*no_alloc*/ false, /*load_mtp*/ true, kv_overrides, nullptr);
     ml.init_mappings(false); // no prefetching
 
     auto mparams = llama_model_default_params();

src/llama.cpp

@@ -305,7 +305,7 @@ static std::pair<int, llama_model *> llama_model_load(struct gguf_context * meta
         const std::string & fname, std::vector<std::string> & splits, FILE * file, llama_model_params & params) {
     try {
         llama_model_loader ml(metadata, set_tensor_data, set_tensor_data_ud, fname, splits, file, params.load_mode,
-            params.check_tensors, params.no_alloc, params.kv_overrides, params.tensor_buft_overrides);
+            params.check_tensors, params.no_alloc, params.load_mtp, params.kv_overrides, params.tensor_buft_overrides);
 
         ml.print_info();
         std::unique_ptr<llama_model> model_ptr(llama_model_create(ml, params));

src/models/cohere2moe.cpp

@@ -55,7 +55,11 @@ void llama_model_cohere2moe::load_arch_tensors(llama_model_loader & ml) {
     const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
     const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
     const int trunk_flags = mtp_only  ? TENSOR_NOT_REQUIRED : 0;
-    const int mtp_flags   = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+
+    if (!ml.load_mtp) {
+        mtp_flags |= TENSOR_SKIP;
+    }
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
 

src/models/deepseek4.cpp

@@ -16,6 +16,16 @@ static float dsv4_rope_attn_factor(float freq_scale, float ext_factor) {
 }
 
 void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) {
+    ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
+    if (hparams.n_layer_nextn > 0 && hparams.n_layer_nextn < hparams.n_layer_all) {
+        const uint32_t n_layer_main = hparams.n_layer_all - hparams.n_layer_nextn;
+        const std::string mtp_probe = "blk." + std::to_string(n_layer_main) + ".nextn.eh_proj.weight";
+        if (ml.get_weight(mtp_probe.c_str()) == nullptr) {
+            hparams.n_layer_nextn = 0;
+        }
+    }
+    GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < block_count");
+
     ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
     ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK,       hparams.n_lora_q);
     ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW,    hparams.n_swa);
@@ -24,8 +34,8 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) {
     ml.get_key(LLM_KV_EXPERT_SHARED_COUNT,         hparams.n_expert_shared);
     ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE,        hparams.expert_weights_scale);
     ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM,         hparams.expert_weights_norm);
-    ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP,     hparams.swiglu_clamp_exp,   hparams.n_layer());
-    if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP,   hparams.swiglu_clamp_shexp, hparams.n_layer(), 0)) {
+    ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP,     hparams.swiglu_clamp_exp,   hparams.n_layer_all);
+    if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP,   hparams.swiglu_clamp_shexp, hparams.n_layer_all, 0)) {
         hparams.swiglu_clamp_shexp = hparams.swiglu_clamp_exp;
     }
 
@@ -41,9 +51,11 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) {
     ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON,             hparams.dsv4_hc_eps);
     ml.get_key(LLM_KV_HASH_LAYER_COUNT,                     hparams.dsv4_hash_layer_count);
 
+    hparams.n_embd_out_impl = hparams.dsv4_hc_mult * hparams.n_embd;
+
     uint32_t n_compress_ratios = 0;
     ml.get_arr_n(LLM_KV_ATTENTION_COMPRESS_RATIOS, n_compress_ratios);
-    if (n_compress_ratios < hparams.n_layer()) {
+    if (n_compress_ratios < hparams.n_layer_all) {
         throw std::runtime_error("DeepSeek-V4 compress_ratios is shorter than block_count");
     }
     ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios);
@@ -54,6 +66,9 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) {
     }
     hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
     hparams.set_swa_pattern(0);
+    for (uint32_t il = hparams.n_layer(); il < hparams.n_layer_all; ++il) {
+        hparams.is_swa_impl[il] = true;
+    }
 
     switch (hparams.n_layer()) {
         case 43: type = LLM_TYPE_UNKNOWN; break;
@@ -61,7 +76,7 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) {
     }
 }
 
-void llama_model_deepseek4::load_arch_tensors(llama_model_loader &) {
+void llama_model_deepseek4::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
     const int64_t q_lora_rank     = hparams.n_lora_q;
@@ -75,6 +90,10 @@ void llama_model_deepseek4::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;
 
+    const bool mtp_only = (n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
+    const int trunk_flags = mtp_only    ? TENSOR_NOT_REQUIRED : 0;
+    const int mtp_flags   = ml.load_mtp ? 0 : TENSOR_SKIP;
+
     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);
@@ -84,69 +103,82 @@ void llama_model_deepseek4::load_arch_tensors(llama_model_loader &) {
     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) {
+    for (int i = 0; i < n_layer_all; ++i) {
         auto & layer = layers[i];
+        const int flags = i < n_layer ? trunk_flags : mtp_flags;
 
-        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}, 0);
-        layer.wo_b          = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B,    "weight", i), {o_groups * o_lora_rank, n_embd}, 0);
+        layer.attn_norm     = create_tensor(tn(LLM_TENSOR_ATTN_NORM,     "weight", i), {n_embd}, flags);
+        layer.attn_sinks    = create_tensor(tn(LLM_TENSOR_ATTN_SINKS,    "weight", i), {n_head}, flags);
+        layer.wq_a          = create_tensor(tn(LLM_TENSOR_ATTN_Q_A,      "weight", i), {n_embd, q_lora_rank}, flags);
+        layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, flags);
+        layer.wq_b          = create_tensor(tn(LLM_TENSOR_ATTN_Q_B,      "weight", i), {q_lora_rank, n_head * n_embd_head}, flags);
+        layer.wkv           = create_tensor(tn(LLM_TENSOR_ATTN_KV,       "weight", i), {n_embd, n_embd_head}, flags);
+        layer.attn_kv_norm  = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM,  "weight", i), {n_embd_head}, flags);
+        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}, flags);
+        layer.wo_b          = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B,    "weight", i), {o_groups * o_lora_rank, n_embd}, flags);
 
-        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);
+        layer.hc_attn_fn    = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN,    "weight", i), {hc_dim, hc_mix_dim}, flags);
+        layer.hc_attn_base  = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE,  "weight", i), {hc_mix_dim}, flags);
+        layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, flags);
+        layer.hc_ffn_fn     = create_tensor(tn(LLM_TENSOR_HC_FFN_FN,     "weight", i), {hc_dim, hc_mix_dim}, flags);
+        layer.hc_ffn_base   = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE,   "weight", i), {hc_mix_dim}, flags);
+        layer.hc_ffn_scale  = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE,  "weight", i), {3}, flags);
 
         const int64_t ratio = hparams.dsv4_compress_ratios[i];
         if (ratio != 0) {
             const int64_t coff = ratio == 4 ? 2 : 1;
 
-            layer.attn_comp_wkv   = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WKV,   "weight", i), {n_embd, coff * n_embd_head}, 0);
-            layer.attn_comp_wgate = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "weight", i), {n_embd, coff * n_embd_head}, 0);
-            layer.attn_comp_ape   = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_APE,   "weight", i), {coff * n_embd_head, ratio}, 0);
-            layer.attn_comp_norm  = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_NORM,  "weight", i), {n_embd_head}, 0);
+            layer.attn_comp_wkv   = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WKV,   "weight", i), {n_embd, coff * n_embd_head}, flags);
+            layer.attn_comp_wgate = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "weight", i), {n_embd, coff * n_embd_head}, flags);
+            layer.attn_comp_ape   = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_APE,   "weight", i), {coff * n_embd_head, ratio}, flags);
+            layer.attn_comp_norm  = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_NORM,  "weight", i), {n_embd_head}, flags);
 
             if (ratio == 4) {
                 const int64_t n_embd_indexer = hparams.indexer_head_size;
 
-                layer.indexer_proj     = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ,     "weight", i), {n_embd, hparams.indexer_n_head}, 0);
-                layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, hparams.indexer_n_head * n_embd_indexer}, 0);
+                layer.indexer_proj     = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ,     "weight", i), {n_embd, hparams.indexer_n_head}, flags);
+                layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, hparams.indexer_n_head * n_embd_indexer}, flags);
 
-                layer.indexer_comp_wkv   = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WKV,   "weight", i), {n_embd, 2 * n_embd_indexer}, 0);
-                layer.indexer_comp_wgate = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "weight", i), {n_embd, 2 * n_embd_indexer}, 0);
-                layer.indexer_comp_ape   = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_APE,   "weight", i), {2 * n_embd_indexer, ratio}, 0);
-                layer.indexer_comp_norm  = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_NORM,  "weight", i), {n_embd_indexer}, 0);
+                layer.indexer_comp_wkv   = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WKV,   "weight", i), {n_embd, 2 * n_embd_indexer}, flags);
+                layer.indexer_comp_wgate = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "weight", i), {n_embd, 2 * n_embd_indexer}, flags);
+                layer.indexer_comp_ape   = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_APE,   "weight", i), {2 * n_embd_indexer, ratio}, flags);
+                layer.indexer_comp_norm  = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_NORM,  "weight", i), {n_embd_indexer}, flags);
             } else if (ratio != 128) {
                 throw std::runtime_error("DeepSeek-V4 loader only supports compression ratios 0, 4, and 128");
             }
         }
 
-        layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
+        layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags);
         if ((uint32_t) i < hparams.dsv4_hash_layer_count) {
-            layer.ffn_gate_tid2eid = create_tensor(tn(LLM_TENSOR_FFN_GATE_TID2EID, "weight", i), {n_expert_used, n_vocab}, 0);
+            layer.ffn_gate_tid2eid = create_tensor(tn(LLM_TENSOR_FFN_GATE_TID2EID, "weight", i), {n_expert_used, n_vocab}, flags);
         } else {
-            layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);
+            layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, flags);
         }
-        layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+        layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags);
 
-        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_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd,   n_ff_exp, n_expert}, flags);
+        layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd,   n_expert}, flags);
+        layer.ffn_up_exps   = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,   "weight", i), {n_embd,   n_ff_exp, n_expert}, flags);
 
-        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);
+        layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd,                     n_ff_exp * n_expert_shared}, flags);
+        layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd                    }, flags);
+        layer.ffn_up_shexp   = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,   "weight", i), {n_embd,                     n_ff_exp * n_expert_shared}, flags);
+
+        if (i >= n_layer) {
+            layer.nextn.eh_proj          = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ,          "weight", i), {2 * n_embd, n_embd}, flags);
+            layer.nextn.enorm            = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM,            "weight", i), {n_embd},             flags);
+            layer.nextn.hnorm            = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM,            "weight", i), {n_embd},             flags);
+            layer.nextn.embed_tokens     = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS,     "weight", i), {n_embd, n_vocab},    TENSOR_NOT_REQUIRED | flags);
+            layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), {n_embd, n_vocab},    TENSOR_NOT_REQUIRED | flags);
+            layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd},             TENSOR_NOT_REQUIRED | flags);
+        }
     }
 }
 
 std::unique_ptr<llm_graph_context> llama_model_deepseek4::build_arch_graph(const llm_graph_params & params) const {
+    if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
+        return std::make_unique<graph_mtp>(*this, params);
+    }
     return std::make_unique<graph>(*this, params);
 }
 
@@ -175,18 +207,69 @@ static ggml_tensor * dsv4_append_zero_row(ggml_context * ctx, ggml_tensor * t, b
     return ggml_concat(ctx, t, row, 1);
 }
 
-static ggml_tensor * dsv4_with_zero_dep(ggml_context * ctx, ggml_tensor * t, ggml_tensor * dep) {
-    if (dep == nullptr) {
-        return t;
+struct dsv4_state_tensors {
+    ggml_tensor * kv;
+    ggml_tensor * score;
+};
+
+static dsv4_state_tensors dsv4_build_state_restore(
+        ggml_context * ctx,
+        const llm_graph_input_dsv4::comp_input & inp,
+        const llama_dsv4_comp_state * state,
+        int32_t il) {
+    dsv4_state_tensors restored = {
+        state->get_kv_all(ctx, il),
+        state->get_score_all(ctx, il),
+    };
+
+    if (inp.state_restore_src_idxs == nullptr || inp.state_restore_dst_idxs == nullptr) {
+        return restored;
     }
 
-    ggml_tensor * zero = ggml_scale(ctx, ggml_sum(ctx, dep), 0.0f);
-    return ggml_add(ctx, t, zero);
+    ggml_tensor * kv_rows = ggml_get_rows(ctx, restored.kv, inp.state_restore_src_idxs);
+    restored.kv = state->cpy_kv(ctx, kv_rows, inp.state_restore_dst_idxs, il);
+
+    ggml_tensor * score_rows = ggml_get_rows(ctx, restored.score, inp.state_restore_src_idxs);
+    restored.score = state->cpy_score(ctx, score_rows, inp.state_restore_dst_idxs, il);
+
+    return restored;
+}
+
+static dsv4_state_tensors dsv4_build_state_snapshot(
+        ggml_context * ctx,
+        const llm_graph_input_dsv4::comp_input & inp,
+        const llama_dsv4_comp_state * state,
+        ggml_tensor * source_kv,
+        ggml_tensor * source_score,
+        int32_t il) {
+    if (inp.state_snapshot_src_idxs == nullptr || inp.state_snapshot_dst_idxs == nullptr ||
+            source_kv == nullptr || source_score == nullptr) {
+        return {};
+    }
+
+    ggml_tensor * kv_rows = ggml_get_rows(ctx, source_kv, inp.state_snapshot_src_idxs);
+    ggml_tensor * kv = state->cpy_kv(ctx, kv_rows, inp.state_snapshot_dst_idxs, il);
+
+    ggml_tensor * score_rows = ggml_get_rows(ctx, source_score, inp.state_snapshot_src_idxs);
+    ggml_tensor * score = state->cpy_score(ctx, score_rows, inp.state_snapshot_dst_idxs, il);
+
+    return { kv, score };
 }
 
 static constexpr int64_t DSV4_CSA_RATIO  = 4;
 static constexpr int64_t DSV4_HCA_RATIO  = 128;
 
+// mean over the hyper-connection streams: [n_embd, hc, n_tokens] -> [n_embd, n_tokens]
+static ggml_tensor * dsv4_hc_mean(ggml_context * ctx, ggml_tensor * x) {
+    const int64_t hc = x->ne[1];
+
+    ggml_tensor * acc = ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], 0);
+    for (int64_t s = 1; s < hc; ++s) {
+        acc = ggml_add(ctx, acc, ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], s*x->nb[1]));
+    }
+    return ggml_scale(ctx, acc, 1.0f/hc);
+}
+
 static ggml_tensor * dsv4_hc_affine(
         ggml_context * ctx,
         ggml_tensor  * x,
@@ -804,8 +887,29 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
         ggml_tensor * cur,
         ggml_tensor * inp_pos,
         int il) const {
+    return build_attention_impl(model, inp_dsv4, nullptr, cur, inp_pos, il);
+}
+
+ggml_tensor * llama_model_deepseek4::graph::build_attention(
+        const llama_model & model,
+        llm_graph_input_attn_k_iswa * inp_mtp,
+        ggml_tensor * cur,
+        ggml_tensor * inp_pos,
+        int il) const {
+    return build_attention_impl(model, nullptr, inp_mtp, cur, inp_pos, il);
+}
+
+ggml_tensor * llama_model_deepseek4::graph::build_attention_impl(
+        const llama_model & model,
+        llm_graph_input_dsv4 * inp_dsv4,
+        llm_graph_input_attn_k_iswa * inp_mtp,
+        ggml_tensor * cur,
+        ggml_tensor * inp_pos,
+        int il) const {
+    GGML_ASSERT((inp_dsv4 == nullptr) != (inp_mtp == nullptr));
+
     const auto & layer = model.layers[il];
-    llm_graph_input_dsv4_raw * inp_attn = inp_dsv4->get_raw();
+    llm_graph_input_dsv4_raw * inp_attn = inp_dsv4 ? inp_dsv4->get_raw() : nullptr;
 
     const int64_t n_embd_head      = hparams.n_embd_head_k();
     const int64_t n_embd_head_rope = hparams.n_rot();
@@ -873,9 +977,12 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
     cb(kv, "kv", il);
 
     const int64_t ratio = hparams.dsv4_compress_ratios[il];
+    GGML_ASSERT(inp_dsv4 || ratio == 0);
 
     ggml_tensor * hca_state_kv    = nullptr;
     ggml_tensor * hca_state_score = nullptr;
+    ggml_tensor * hca_source_kv   = nullptr;
+    ggml_tensor * hca_source_score = nullptr;
     if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_pos) {
         hca_state_kv = build_lora_mm(layer.attn_comp_wkv, cur);
         cb(hca_state_kv, "hca_state_kv", il);
@@ -906,10 +1013,16 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
 
         GGML_ASSERT(inp_dsv4->get_csa().state_write_idxs);
 
-        ggml_tensor * csa_source_kv = ggml_concat(ctx0,
-                inp_dsv4->mctx->get_csa_state()->get_kv(ctx0, il), csa_state_kv, 1);
-        ggml_tensor * csa_source_score = ggml_concat(ctx0,
-                inp_dsv4->mctx->get_csa_state()->get_score(ctx0, il), csa_state_score, 1);
+        const auto * csa_state = inp_dsv4->mctx->get_csa_state();
+        const dsv4_state_tensors csa_restored = dsv4_build_state_restore(
+                ctx0, inp_dsv4->get_csa(), csa_state, il);
+        ggml_tensor * csa_base_kv = dsv4_view_2d(
+                ctx0, csa_restored.kv, csa_restored.kv->ne[0], csa_state->get_n_rows(), 0);
+        ggml_tensor * csa_base_score = dsv4_view_2d(
+                ctx0, csa_restored.score, csa_restored.score->ne[0], csa_state->get_n_rows(), 0);
+
+        ggml_tensor * csa_source_kv = ggml_concat(ctx0, csa_base_kv, csa_state_kv, 1);
+        ggml_tensor * csa_source_score = ggml_concat(ctx0, csa_base_score, csa_state_score, 1);
 
         ggml_tensor * kv_comp_csa_state = build_overlap_compressed_kv_from_state(
                 csa_source_kv,
@@ -930,8 +1043,19 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
         ggml_build_forward_expand(gf, inp_dsv4->mctx->get_csa()->cpy_k(ctx0,
                     kv_comp_csa_state, inp_dsv4->get_csa().state_write_idxs, il));
 
-        csa_state_kv    = dsv4_with_zero_dep(ctx0, csa_state_kv,    kv_comp_csa_state);
-        csa_state_score = dsv4_with_zero_dep(ctx0, csa_state_score, kv_comp_csa_state);
+        ggml_tensor * csa_snapshot_source_kv = ggml_concat(ctx0,
+                csa_restored.kv, csa_state_kv, 1);
+        ggml_tensor * csa_snapshot_source_score = ggml_concat(ctx0,
+                csa_restored.score, csa_state_score, 1);
+
+        const dsv4_state_tensors csa_snapshot = dsv4_build_state_snapshot(
+                ctx0, inp_dsv4->get_csa(), csa_state, csa_snapshot_source_kv, csa_snapshot_source_score, il);
+        if (csa_snapshot.kv != nullptr) {
+            ggml_build_forward_expand(gf, csa_snapshot.kv);
+        }
+        if (csa_snapshot.score != nullptr) {
+            ggml_build_forward_expand(gf, csa_snapshot.score);
+        }
 
         ggml_tensor * csa_persist_kv = ggml_get_rows(ctx0, csa_state_kv, inp_dsv4->get_csa().state_persist_src_idxs);
         ggml_tensor * csa_persist_score = ggml_get_rows(ctx0, csa_state_score, inp_dsv4->get_csa().state_persist_src_idxs);
@@ -958,10 +1082,16 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
 
         GGML_ASSERT(inp_dsv4->get_lid().state_write_idxs);
 
-        ggml_tensor * lid_source_kv = ggml_concat(ctx0,
-                inp_dsv4->mctx->get_lid_state()->get_kv(ctx0, il), lid_state_kv, 1);
-        ggml_tensor * lid_source_score = ggml_concat(ctx0,
-                inp_dsv4->mctx->get_lid_state()->get_score(ctx0, il), lid_state_score, 1);
+        const auto * lid_state = inp_dsv4->mctx->get_lid_state();
+        const dsv4_state_tensors lid_restored = dsv4_build_state_restore(
+                ctx0, inp_dsv4->get_lid(), lid_state, il);
+        ggml_tensor * lid_base_kv = dsv4_view_2d(
+                ctx0, lid_restored.kv, lid_restored.kv->ne[0], lid_state->get_n_rows(), 0);
+        ggml_tensor * lid_base_score = dsv4_view_2d(
+                ctx0, lid_restored.score, lid_restored.score->ne[0], lid_state->get_n_rows(), 0);
+
+        ggml_tensor * lid_source_kv = ggml_concat(ctx0, lid_base_kv, lid_state_kv, 1);
+        ggml_tensor * lid_source_score = ggml_concat(ctx0, lid_base_score, lid_state_score, 1);
 
         ggml_tensor * kv_comp_lid_state = build_overlap_compressed_kv_from_state(
                 lid_source_kv,
@@ -982,8 +1112,19 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
         ggml_build_forward_expand(gf, inp_dsv4->mctx->get_lid()->cpy_k(ctx0,
                     kv_comp_lid_state, inp_dsv4->get_lid().state_write_idxs, il));
 
-        lid_state_kv    = dsv4_with_zero_dep(ctx0, lid_state_kv,    kv_comp_lid_state);
-        lid_state_score = dsv4_with_zero_dep(ctx0, lid_state_score, kv_comp_lid_state);
+        ggml_tensor * lid_snapshot_source_kv = ggml_concat(ctx0,
+                lid_restored.kv, lid_state_kv, 1);
+        ggml_tensor * lid_snapshot_source_score = ggml_concat(ctx0,
+                lid_restored.score, lid_state_score, 1);
+
+        const dsv4_state_tensors lid_snapshot = dsv4_build_state_snapshot(
+                ctx0, inp_dsv4->get_lid(), lid_state, lid_snapshot_source_kv, lid_snapshot_source_score, il);
+        if (lid_snapshot.kv != nullptr) {
+            ggml_build_forward_expand(gf, lid_snapshot.kv);
+        }
+        if (lid_snapshot.score != nullptr) {
+            ggml_build_forward_expand(gf, lid_snapshot.score);
+        }
 
         ggml_tensor * lid_persist_kv = ggml_get_rows(ctx0, lid_state_kv, inp_dsv4->get_lid().state_persist_src_idxs);
         ggml_tensor * lid_persist_score = ggml_get_rows(ctx0, lid_state_score, inp_dsv4->get_lid().state_persist_src_idxs);
@@ -997,15 +1138,21 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
         ggml_build_forward_expand(gf, lid_state_score);
     }
 
-    ggml_tensor * hca_state_dep = nullptr;
+    const llama_dsv4_comp_state * hca_state = nullptr;
+    dsv4_state_tensors hca_restored = {};
     if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_write_idxs) {
         GGML_ASSERT(hca_state_kv);
         GGML_ASSERT(hca_state_score);
 
-        ggml_tensor * hca_source_kv = ggml_concat(ctx0,
-                inp_dsv4->mctx->get_hca_state()->get_kv(ctx0, il), hca_state_kv, 1);
-        ggml_tensor * hca_source_score = ggml_concat(ctx0,
-                inp_dsv4->mctx->get_hca_state()->get_score(ctx0, il), hca_state_score, 1);
+        hca_state = inp_dsv4->mctx->get_hca_state();
+        hca_restored = dsv4_build_state_restore(ctx0, inp_dsv4->get_hca(), hca_state, il);
+        ggml_tensor * hca_base_kv = dsv4_view_2d(
+                ctx0, hca_restored.kv, hca_restored.kv->ne[0], hca_state->get_n_rows(), 0);
+        ggml_tensor * hca_base_score = dsv4_view_2d(
+                ctx0, hca_restored.score, hca_restored.score->ne[0], hca_state->get_n_rows(), 0);
+
+        hca_source_kv = ggml_concat(ctx0, hca_base_kv, hca_state_kv, 1);
+        hca_source_score = ggml_concat(ctx0, hca_base_score, hca_state_score, 1);
 
         ggml_tensor * kv_comp_hca = build_hca_compressed_kv_from_state(
                 hca_source_kv,
@@ -1024,15 +1171,41 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
 
         ggml_build_forward_expand(gf, inp_dsv4->mctx->get_hca()->cpy_k(ctx0,
                     kv_comp_hca, inp_dsv4->get_hca().state_write_idxs, il));
-        hca_state_dep = kv_comp_hca;
     }
 
     if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_pos) {
         GGML_ASSERT(hca_state_kv);
         GGML_ASSERT(hca_state_score);
 
-        hca_state_kv    = dsv4_with_zero_dep(ctx0, hca_state_kv,    hca_state_dep);
-        hca_state_score = dsv4_with_zero_dep(ctx0, hca_state_score, hca_state_dep);
+        if (hca_state == nullptr) {
+            hca_state = inp_dsv4->mctx->get_hca_state();
+        }
+        if (hca_restored.kv == nullptr) {
+            hca_restored = dsv4_build_state_restore(ctx0, inp_dsv4->get_hca(), hca_state, il);
+        }
+        if (hca_source_kv == nullptr || hca_source_score == nullptr) {
+            ggml_tensor * hca_base_kv = dsv4_view_2d(
+                    ctx0, hca_restored.kv, hca_restored.kv->ne[0], hca_state->get_n_rows(), 0);
+            ggml_tensor * hca_base_score = dsv4_view_2d(
+                    ctx0, hca_restored.score, hca_restored.score->ne[0], hca_state->get_n_rows(), 0);
+
+            hca_source_kv = ggml_concat(ctx0, hca_base_kv, hca_state_kv, 1);
+            hca_source_score = ggml_concat(ctx0, hca_base_score, hca_state_score, 1);
+        }
+
+        ggml_tensor * hca_snapshot_source_kv = ggml_concat(ctx0,
+                hca_restored.kv, hca_state_kv, 1);
+        ggml_tensor * hca_snapshot_source_score = ggml_concat(ctx0,
+                hca_restored.score, hca_state_score, 1);
+
+        const dsv4_state_tensors hca_snapshot = dsv4_build_state_snapshot(
+                ctx0, inp_dsv4->get_hca(), hca_state, hca_snapshot_source_kv, hca_snapshot_source_score, il);
+        if (hca_snapshot.kv != nullptr) {
+            ggml_build_forward_expand(gf, hca_snapshot.kv);
+        }
+        if (hca_snapshot.score != nullptr) {
+            ggml_build_forward_expand(gf, hca_snapshot.score);
+        }
 
         ggml_tensor * hca_persist_kv = ggml_get_rows(ctx0, hca_state_kv, inp_dsv4->get_hca().state_persist_src_idxs);
         ggml_tensor * hca_persist_score = ggml_get_rows(ctx0, hca_state_score, inp_dsv4->get_hca().state_persist_src_idxs);
@@ -1047,7 +1220,14 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
     }
 
     ggml_tensor * out = nullptr;
-    if (ratio == DSV4_CSA_RATIO &&
+    if (inp_mtp) {
+        out = build_attn(inp_mtp,
+                nullptr, nullptr, nullptr,
+                q, kv, nullptr,
+                nullptr, layer.attn_sinks, nullptr,
+                1.0f/sqrtf(float(n_embd_head)), il);
+        cb(out, "attn_raw", il);
+    } else if (ratio == DSV4_CSA_RATIO &&
             inp_dsv4->get_csa().kq_mask &&
             inp_dsv4->get_lid().kq_mask &&
             inp_dsv4->get_lid().k_rot) {
@@ -1106,6 +1286,12 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p
     cb(inpL, "hc_init", -1);
 
     for (int il = 0; il < n_layer; ++il) {
+        if ((size_t) il < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[il]) {
+            res->t_layer_inp[il] = dsv4_hc_mean(ctx0, 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;
@@ -1182,10 +1368,23 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p
         cb(inpL, "l_last", il);
     }
 
+    if ((size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer]) {
+        res->t_layer_inp[n_layer] = dsv4_hc_mean(ctx0, inpL);
+        cb(res->t_layer_inp[n_layer], "layer_inp", n_layer);
+        ggml_build_forward_expand(gf, res->t_layer_inp[n_layer]);
+    }
+
+    ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
+    ggml_tensor * flat_out = inp_out_ids ? ggml_get_rows(ctx0, flat, inp_out_ids) : flat;
+
+    if (cparams.embeddings_nextn) {
+        ggml_tensor * h_nextn = cparams.embeddings_nextn_masked ? flat_out : inpL;
+        cb(h_nextn, "h_nextn", -1);
+        res->t_h_nextn = h_nextn;
+    }
+
     if (inp_out_ids) {
-        ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
-        flat = ggml_get_rows(ctx0, flat, inp_out_ids);
-        inpL = ggml_reshape_3d(ctx0, flat, n_embd, hc, n_outputs);
+        inpL = ggml_reshape_3d(ctx0, flat_out, n_embd, hc, n_outputs);
     }
 
     cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base);
@@ -1201,3 +1400,145 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p
 
     ggml_build_forward_expand(gf, cur);
 }
+
+
+llama_model_deepseek4::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) :
+    graph(params) {
+    GGML_ASSERT(hparams.n_layer_nextn > 0 && "DEEPSEEK4 MTP requires n_layer_nextn > 0");
+    GGML_ASSERT(hparams.n_layer_nextn == 1 && "DEEPSEEK4 MTP currently only supports a single MTP block");
+    GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
+            cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
+            "nextn_layer_offset out of range [0, n_layer_nextn)");
+    GGML_ASSERT(ubatch.token && "DEEPSEEK4 MTP requires token input");
+
+    const int64_t hc = hparams.dsv4_hc_mult;
+    GGML_ASSERT(hparams.n_embd_out() == (uint32_t) (n_embd*hc) && "DEEPSEEK4 MTP hidden width mismatch");
+
+    const int il = hparams.n_layer() + cparams.nextn_layer_offset;
+    const auto & layer = model.layers[il];
+
+    GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");
+    GGML_ASSERT(layer.nextn.enorm   && "MTP block missing nextn.enorm");
+    GGML_ASSERT(layer.nextn.hnorm   && "MTP block missing nextn.hnorm");
+
+    auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd_out());
+
+    inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+    ggml_set_input(inp->tokens);
+
+    inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_out(), n_tokens);
+    ggml_set_input(inp->embd);
+
+    inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_out(), n_tokens);
+    ggml_set_input(inp->h);
+    ggml_set_name(inp->h, "mtp_h_input");
+
+    ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
+    ggml_tensor * tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
+    cb(tok_embd, "mtp_tok_embd", il);
+
+    ggml_tensor * h_state = ggml_reshape_3d(ctx0, inp->h, n_embd, hc, n_tokens);
+    cb(h_state, "mtp_h_state", il);
+
+    res->add_input(std::move(inp));
+
+    ggml_tensor * inp_pos = build_inp_pos();
+    ggml_tensor * inp_out_ids = build_inp_out_ids();
+    llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa();
+
+    ggml_tensor * h_norm = build_norm(h_state, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
+    cb(h_norm, "mtp_hnorm", il);
+
+    ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
+    e_norm = ggml_reshape_3d(ctx0, e_norm, n_embd, 1, n_tokens);
+    e_norm = ggml_repeat_4d(ctx0, e_norm, n_embd, hc, n_tokens, 1);
+    cb(e_norm, "mtp_enorm", il);
+
+    ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, 0);
+    cb(concat, "mtp_concat", il);
+
+    ggml_tensor * inpL = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
+    cb(inpL, "mtp_eh_proj", il);
+
+    ggml_tensor * residual = inpL;
+    ggml_tensor * post = nullptr;
+    ggml_tensor * comb = nullptr;
+
+    ggml_tensor * cur = build_hc_pre(inpL,
+            layer.hc_attn_fn,
+            layer.hc_attn_scale,
+            layer.hc_attn_base,
+            &post, &comb, il);
+    cb(cur, "mtp_hc_attn_pre", il);
+
+    cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+    cb(cur, "mtp_attn_norm", il);
+
+    cur = build_attention(model, inp_attn, cur, inp_pos, il);
+
+    inpL = build_hc_post(cur, residual, post, comb, il);
+    cb(inpL, "mtp_hc_attn_post", il);
+
+    residual = inpL;
+    cur = build_hc_pre(inpL,
+            layer.hc_ffn_fn,
+            layer.hc_ffn_scale,
+            layer.hc_ffn_base,
+            &post, &comb, il);
+    cb(cur, "mtp_hc_ffn_pre", il);
+
+    cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
+    cb(cur, "mtp_ffn_norm", il);
+
+    GGML_ASSERT((uint32_t) il >= hparams.dsv4_hash_layer_count && "DEEPSEEK4 MTP does not support hash-routed MTP blocks");
+    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, "mtp_ffn_moe_out", il);
+
+    ggml_tensor * ffn_shexp = build_ffn(cur,
+            layer.ffn_up_shexp, nullptr, nullptr,
+            layer.ffn_gate_shexp, nullptr, nullptr,
+            layer.ffn_down_shexp, nullptr, nullptr,
+            nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
+    cb(ffn_shexp, "mtp_ffn_shexp", il);
+
+    cur = ggml_add(ctx0, moe_out, ffn_shexp);
+    cb(cur, "mtp_ffn_out", il);
+
+    inpL = build_hc_post(cur, residual, post, comb, il);
+    inpL = build_cvec(inpL, il);
+    cb(inpL, "mtp_l_out", il);
+
+    ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
+    ggml_tensor * h_nextn = ggml_get_rows(ctx0, flat, inp_out_ids);
+    cb(h_nextn, "h_nextn", -1);
+    res->t_h_nextn = h_nextn;
+
+    inpL = ggml_reshape_3d(ctx0, h_nextn, n_embd, hc, n_outputs);
+
+    cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base);
+    cb(cur, "mtp_hc_head", -1);
+
+    ggml_tensor * head_norm_w = layer.nextn.shared_head_norm ? layer.nextn.shared_head_norm : model.output_norm;
+    GGML_ASSERT(head_norm_w && "DEEPSEEK4 MTP missing shared head norm");
+    cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);
+    cb(cur, "mtp_shared_head_norm", -1);
+    res->t_embd = cur;
+
+    ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
+    GGML_ASSERT(head_w && "DEEPSEEK4 MTP missing LM head");
+    cur = ggml_mul_mat(ctx0, head_w, cur);
+    cb(cur, "result_output", -1);
+
+    res->t_logits = cur;
+    ggml_build_forward_expand(gf, cur);
+}

src/models/dflash.cpp

@@ -20,6 +20,48 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
     }
     LLAMA_LOG_INFO("]\n");
 
+    // DeepSeek-V4 DSpark backbone: stages are full DSV4 blocks, uniform sliding window (the draft KV ring)
+    ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult, false);
+    if (hparams.dsv4_hc_mult > 0) {
+        ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK,                hparams.n_lora_q);
+        ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW,             hparams.n_swa);
+        ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH,           hparams.n_ff_exp);
+        ml.get_key(LLM_KV_EXPERT_SHARED_COUNT,                  hparams.n_expert_shared);
+        ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE,                 hparams.expert_weights_scale);
+        ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM,                  hparams.expert_weights_norm);
+        ml.get_key(LLM_KV_EXPERT_GATING_FUNC,                   hparams.expert_gating_func);
+        ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP,              hparams.swiglu_clamp_exp, hparams.n_layer_all);
+        if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP,       hparams.swiglu_clamp_shexp, hparams.n_layer_all, 0)) {
+            hparams.swiglu_clamp_shexp = hparams.swiglu_clamp_exp;
+        }
+        ml.get_key(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,         hparams.dsv4_o_group_count);
+        ml.get_key(LLM_KV_ATTENTION_OUTPUT_LORA_RANK,           hparams.dsv4_o_lora_rank);
+        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);
+
+        if (hparams.expert_gating_func != LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) {
+            throw std::runtime_error("DSpark DSV4 draft expects sqrtsoftplus MoE scoring");
+        }
+        for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {
+            if (hparams.dsv4_compress_ratios[il] != 0) {
+                throw std::runtime_error("DSpark DSV4 draft expects uncompressed attention on all stages");
+            }
+        }
+
+        GGML_ASSERT(hparams.n_swa > 0);
+        hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
+        hparams.set_swa_pattern(0);
+        for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {
+            hparams.is_swa_impl[il] = true;
+        }
+        hparams.rope_freq_base_train_swa  = hparams.rope_freq_base_train;
+        hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;
+
+        type = LLM_TYPE_UNKNOWN;
+        return;
+    }
+
     // optional interleaved sliding-window attention with per-layer pattern array.
     // DFlash has a single rope, so the SWA rope == main rope.
     if (ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false) && hparams.n_swa > 0) {
@@ -58,6 +100,56 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
     output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc)
     output_norm     = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM,    "weight"), { n_embd }, 0); // decoder final norm
 
+    if (hparams.dsv4_hc_mult > 0) {
+        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;
+
+        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];
+
+            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}, 0);
+            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);
+
+            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_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);
+        }
+        return;
+    }
+
     for (int i = 0; i < n_layer; ++i) {
         auto & layer = layers[i];
 
@@ -84,6 +176,9 @@ std::unique_ptr<llm_graph_context> llama_model_dflash::build_arch_graph(const ll
             return std::make_unique<graph<true>>(*this, params);
         case LLM_GRAPH_TYPE_DEFAULT:
         case LLM_GRAPH_TYPE_DECODER:
+            if (hparams.dsv4_hc_mult > 0) {
+                return std::make_unique<graph_dsv4>(*this, params);
+            }
             return std::make_unique<graph<false>>(*this, params);
         default:
             GGML_ABORT("invalid graph type");
@@ -403,3 +498,178 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
         build_dspark_markov_head(*this, model, inp_tokens);
     }
 }
+
+// DSV4 DSpark decoder, dual-mode by batch type (see the DFlash decoder above):
+//   * embd batch  -> project main_x through each stage's wkv and inject K into the ring cache
+//   * token batch -> noise block through 3 full DSV4 stages (hc + MLA + MoE), markov + confidence heads
+llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_graph_params & params) :
+    llama_model_deepseek4::graph(params) {
+    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_tensor * inp_pos = build_inp_pos();
+
+    llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa();
+
+    // KV cache injection: fused target features from the encoder
+    if (ubatch.embd) {
+        auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
+
+        inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
+        ggml_set_input(inp->embd);
+
+        ggml_tensor * inp_g = inp->embd;
+        cb(inp_g, "inp_g_embeddings", -1);
+
+        res->add_input(std::move(inp));
+
+        for (int il = 0; il < n_layer; ++il) {
+            const auto & layer = model.layers[il];
+
+            // main-track KV: kv_norm(wkv(main_x)) with rope on the trailing dims, same
+            // rope parameters as the uncompressed layers in build_attention_impl
+            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);
+
+            ggml_tensor * kv_nope = ggml_view_3d(ctx0, kv, n_embd_head_nope, 1, n_tokens,
+                    ggml_row_size(kv->type, n_embd_head),
+                    ggml_row_size(kv->type, n_embd_head),
+                    0);
+            ggml_tensor * kv_pe = ggml_view_3d(ctx0, kv, n_embd_head_rope, 1, n_tokens,
+                    ggml_row_size(kv->type, n_embd_head),
+                    ggml_row_size(kv->type, n_embd_head),
+                    ggml_row_size(kv->type, n_embd_head_nope));
+            kv_pe = ggml_rope_ext(ctx0, kv_pe, inp_pos, nullptr, n_embd_head_rope, rope_type, 0,
+                    freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+            kv = ggml_concat(ctx0, kv_nope, kv_pe, 0);
+            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;
+    }
+
+    // tok_embd from the target model (shared via ctx_other)
+    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;
+    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);
+
+    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 * cur = build_hc_pre(inpL,
+                layer.hc_attn_fn,
+                layer.hc_attn_scale,
+                layer.hc_attn_base,
+                &post, &comb, il);
+        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_attention(model, inp_attn, cur, inp_pos, il);
+
+        inpL = build_hc_post(cur, residual, post, comb, il);
+        cb(inpL, "hc_attn_post", il);
+
+        residual = inpL;
+        cur = build_hc_pre(inpL,
+                layer.hc_ffn_fn,
+                layer.hc_ffn_scale,
+                layer.hc_ffn_base,
+                &post, &comb, 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);
+        cb(inpL, "l_out", il);
+    }
+
+    ggml_tensor * cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base);
+    cb(cur, "hc_head", -1);
+
+    // confidence head input: the reference scores the pre-norm collapsed hidden state
+    res->t_embd = cur;
+
+    cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
+    cb(cur, "result_norm", -1);
+
+    // lm_head from the target model (shared via ctx_other)
+    auto * output = model.output;
+    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;
+    }
+
+    cur = build_lora_mm(output, cur);
+    cb(cur, "result_output", -1);
+    res->t_logits = cur;
+
+    ggml_build_forward_expand(gf, cur);
+
+    if (model.dspark_markov_w1) {
+        build_dspark_markov_head(*this, model, inp_tokens);
+    }
+}

src/models/glm-dsa.cpp

@@ -91,7 +91,11 @@ void llama_model_glm_dsa::load_arch_tensors(llama_model_loader & ml) {
     const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
     const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
     const int trunk_flags = mtp_only   ? TENSOR_NOT_REQUIRED : 0;
-    const int mtp_flags   = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+
+    if (!ml.load_mtp) {
+        mtp_flags |= TENSOR_SKIP;
+    }
 
     const bool is_mla = hparams.is_mla();
     if (!is_mla) {

src/models/hy-v3.cpp

@@ -33,7 +33,11 @@ void llama_model_hy_v3::load_arch_tensors(llama_model_loader & ml) {
     const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
     const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
     const int trunk_flags = mtp_only   ? TENSOR_NOT_REQUIRED : 0;
-    const int mtp_flags   = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+
+    if (!ml.load_mtp) {
+        mtp_flags |= TENSOR_SKIP;
+    }
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
 

src/models/mimo2.cpp

@@ -30,7 +30,11 @@ void llama_model_mimo2::load_arch_tensors(llama_model_loader & ml) {
 
     const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
     const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
-    const int  mtp_flags  = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+
+    if (!ml.load_mtp) {
+        mtp_flags |= TENSOR_SKIP;
+    }
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
 

src/models/minimax-m3.cpp

@@ -271,8 +271,6 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
             } else {
                 const int64_t n_idx_dim = hparams.indexer_head_size;   // 128
 
-                GGML_ASSERT(!inp_attn->self_k_rot && !inp_attn->self_v_rot && "MSA: attn-rot not supported");
-
                 // Index Branch, project, norm, partial RoPE, cache
                 ggml_tensor * iq = build_lora_mm(model.layers[il].index_q_proj, cur);
                 ggml_tensor * ik = build_lora_mm(model.layers[il].index_k_proj, cur);
@@ -289,6 +287,14 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
                 ggml_build_forward_expand(gf, mctx_cur->cpy_k_idx(ctx0, ik, inp_attn->get_k_idxs(), il));
                 ggml_tensor * ik_kv = mctx_cur->get_k_idx(ctx0, il);
 
+                if (inp_attn->self_k_rot) {
+                    Qcur = llama_mul_mat_hadamard(ctx0, Qcur, inp_attn->self_k_rot);
+                    Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp_attn->self_k_rot);
+                }
+                if (inp_attn->self_v_rot) {
+                    Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot);
+                }
+
                 // Main branch: store K/V, take cache views
                 ggml_build_forward_expand(gf, Qcur);
                 ggml_build_forward_expand(gf, Kcur);
@@ -431,7 +437,9 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
                         cur = ggml_concat(ctx0, cur, outs[st], 1);
                     }
                 }
-
+                if (inp_attn->self_v_rot) {
+                    cur = llama_mul_mat_hadamard(ctx0, cur, inp_attn->self_v_rot);
+                }
                 cb(cur, "kqv_out", il);
                 if (model.layers[il].wo) {
                     cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);

src/models/models.h

@@ -1107,6 +1107,7 @@ struct llama_model_deepseek4 : public llama_model_base {
     void load_arch_tensors(llama_model_loader & ml) override;
 
     struct graph : public llm_graph_context {
+        graph(const llm_graph_params & params) : llm_graph_context(params) {}
         graph(const llama_model & model, const llm_graph_params & params);
 
         ggml_tensor * build_hc_pre(
@@ -1138,6 +1139,21 @@ struct llama_model_deepseek4 : public llama_model_base {
                 ggml_tensor * inp_pos,
                 int il) const;
 
+        ggml_tensor * build_attention(
+                const llama_model & model,
+                llm_graph_input_attn_k_iswa * inp_mtp,
+                ggml_tensor * cur,
+                ggml_tensor * inp_pos,
+                int il) const;
+
+        ggml_tensor * build_attention_impl(
+                const llama_model & model,
+                llm_graph_input_dsv4 * inp_dsv4,
+                llm_graph_input_attn_k_iswa * inp_mtp,
+                ggml_tensor * cur,
+                ggml_tensor * inp_pos,
+                int il) const;
+
         ggml_tensor * build_hca_compressed_kv_from_state(
                 ggml_tensor * kv_state,
                 ggml_tensor * score_state,
@@ -1213,6 +1229,10 @@ struct llama_model_deepseek4 : public llama_model_base {
                 int il) const;
     };
 
+    struct graph_mtp : public graph {
+        graph_mtp(const llama_model & model, const llm_graph_params & params);
+    };
+
     std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
 };
 
@@ -1272,6 +1292,10 @@ struct llama_model_dflash : public llama_model_base {
         ggml_tensor * build_inp_embd_enc() const;
     };
 
+    struct graph_dsv4 : public llama_model_deepseek4::graph {
+        graph_dsv4(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/qwen35.cpp

@@ -39,6 +39,7 @@ void llama_model_qwen35::load_arch_tensors(llama_model_loader & ml) {
 
     const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
     const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
+    int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
 
@@ -97,25 +98,25 @@ void llama_model_qwen35::load_arch_tensors(llama_model_loader & ml) {
         auto & layer = layers[il];
 
         // MTP block looks like a full-attention Qwen3.5 decoder block.
-        layer.attn_norm      = create_tensor(tn(LLM_TENSOR_ATTN_NORM,      "weight", il), { n_embd }, 0);
-        layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, 0);
+        layer.attn_norm      = create_tensor(tn(LLM_TENSOR_ATTN_NORM,      "weight", il), { n_embd }, mtp_flags);
+        layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, mtp_flags);
 
-        create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, 0);
-        layer.wo          = create_tensor(tn(LLM_TENSOR_ATTN_OUT,    "weight", il), { n_embd_head_k * n_head, n_embd }, 0);
-        layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, 0);
-        layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, 0);
+        create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, mtp_flags);
+        layer.wo          = create_tensor(tn(LLM_TENSOR_ATTN_OUT,    "weight", il), { n_embd_head_k * n_head, n_embd }, mtp_flags);
+        layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, mtp_flags);
+        layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, mtp_flags);
 
-        layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", il), {n_embd,   n_ff}, 0);
-        layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", il), {  n_ff, n_embd}, 0);
-        layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", il), {n_embd,   n_ff}, 0);
+        layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", il), {n_embd,   n_ff}, mtp_flags);
+        layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", il), {  n_ff, n_embd}, mtp_flags);
+        layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", il), {n_embd,   n_ff}, mtp_flags);
 
         // NextN-specific tensors that define the MTP block.
-        layer.nextn.eh_proj          = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ,          "weight", il), { 2 * n_embd, n_embd }, 0);
-        layer.nextn.enorm            = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM,            "weight", il), { n_embd },              0);
-        layer.nextn.hnorm            = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM,            "weight", il), { n_embd },              0);
-        layer.nextn.embed_tokens     = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS,     "weight", il), { n_embd, n_vocab },     TENSOR_NOT_REQUIRED);
-        layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab },     TENSOR_NOT_REQUIRED);
-        layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd },              TENSOR_NOT_REQUIRED);
+        layer.nextn.eh_proj          = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ,          "weight", il), { 2 * n_embd, n_embd }, mtp_flags);
+        layer.nextn.enorm            = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM,            "weight", il), { n_embd },              mtp_flags);
+        layer.nextn.hnorm            = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM,            "weight", il), { n_embd },              mtp_flags);
+        layer.nextn.embed_tokens     = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS,     "weight", il), { n_embd, n_vocab },     mtp_flags|TENSOR_NOT_REQUIRED);
+        layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab },     mtp_flags|TENSOR_NOT_REQUIRED);
+        layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd },              mtp_flags|TENSOR_NOT_REQUIRED);
     };
 
     for (int i = 0; i < n_layer; ++i) {

src/models/qwen35moe.cpp

@@ -42,6 +42,7 @@ void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) {
 
     const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
     const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
+    int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
 
@@ -113,32 +114,32 @@ void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) {
         const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff;
 
         // MTP block looks like a full-attention Qwen3.5 decoder block with MoE FFN.
-        layer.attn_norm      = create_tensor(tn(LLM_TENSOR_ATTN_NORM,      "weight", il), { n_embd }, 0);
-        layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, 0);
+        layer.attn_norm      = create_tensor(tn(LLM_TENSOR_ATTN_NORM,      "weight", il), { n_embd }, mtp_flags);
+        layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, mtp_flags);
 
-        create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, 0);
-        layer.wo          = create_tensor(tn(LLM_TENSOR_ATTN_OUT,    "weight", il), { n_embd_head_k * n_head, n_embd }, 0);
-        layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, 0);
-        layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, 0);
+        create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, mtp_flags);
+        layer.wo          = create_tensor(tn(LLM_TENSOR_ATTN_OUT,    "weight", il), { n_embd_head_k * n_head, n_embd }, mtp_flags);
+        layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, mtp_flags);
+        layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, mtp_flags);
 
         // Routed experts
-        layer.ffn_gate_inp  = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,  "weight", il), { n_embd, n_expert }, 0);
-        layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, 0);
-        create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, 0);
+        layer.ffn_gate_inp  = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,  "weight", il), { n_embd, n_expert }, mtp_flags);
+        layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, mtp_flags);
+        create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, mtp_flags);
 
         // Shared experts
-        layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, 0);
-        layer.ffn_gate_shexp     = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP,     "weight", il), { n_embd, n_ff_shexp }, 0);
-        layer.ffn_up_shexp       = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,       "weight", il), { n_embd, n_ff_shexp }, 0);
-        layer.ffn_down_shexp     = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP,     "weight", il), { n_ff_shexp, n_embd }, 0);
+        layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, mtp_flags);
+        layer.ffn_gate_shexp     = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP,     "weight", il), { n_embd, n_ff_shexp }, mtp_flags);
+        layer.ffn_up_shexp       = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,       "weight", il), { n_embd, n_ff_shexp }, mtp_flags);
+        layer.ffn_down_shexp     = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP,     "weight", il), { n_ff_shexp, n_embd }, mtp_flags);
 
         // NextN-specific tensors that define the MTP block.
-        layer.nextn.eh_proj          = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ,          "weight", il), { 2 * n_embd, n_embd }, 0);
-        layer.nextn.enorm            = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM,            "weight", il), { n_embd },              0);
-        layer.nextn.hnorm            = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM,            "weight", il), { n_embd },              0);
-        layer.nextn.embed_tokens     = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS,     "weight", il), { n_embd, n_vocab },     TENSOR_NOT_REQUIRED);
-        layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab },     TENSOR_NOT_REQUIRED);
-        layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd },              TENSOR_NOT_REQUIRED);
+        layer.nextn.eh_proj          = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ,          "weight", il), { 2 * n_embd, n_embd }, mtp_flags);
+        layer.nextn.enorm            = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM,            "weight", il), { n_embd },              mtp_flags);
+        layer.nextn.hnorm            = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM,            "weight", il), { n_embd },              mtp_flags);
+        layer.nextn.embed_tokens     = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS,     "weight", il), { n_embd, n_vocab },     mtp_flags|TENSOR_NOT_REQUIRED);
+        layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab },     mtp_flags|TENSOR_NOT_REQUIRED);
+        layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd },              mtp_flags|TENSOR_NOT_REQUIRED);
     };
 
     for (int i = 0; i < n_layer; ++i) {

src/models/step35.cpp

@@ -48,7 +48,11 @@ void llama_model_step35::load_arch_tensors(llama_model_loader & ml) {
     const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
     const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
     const int trunk_flags = mtp_only  ? TENSOR_NOT_REQUIRED : 0;
-    const int mtp_flags   = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
+
+    if (!ml.load_mtp) {
+        mtp_flags |= TENSOR_SKIP;
+    }
 
     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
 

tests/test-backend-ops.cpp

@@ -8069,6 +8069,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_dsv4_hc_comb(1, 1));
     test_cases.emplace_back(new test_dsv4_hc_comb(17, 4));
     test_cases.emplace_back(new test_dsv4_hc_comb(257, 8));
+    test_cases.emplace_back(new test_dsv4_hc_comb(17, 20));
 
     test_cases.emplace_back(new test_dsv4_hc_pre(1, 1));
     test_cases.emplace_back(new test_dsv4_hc_pre(31, 17));
@@ -8078,6 +8079,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_dsv4_hc_post(1, 1));
     test_cases.emplace_back(new test_dsv4_hc_post(31, 17));
     test_cases.emplace_back(new test_dsv4_hc_post(128, 257));
+    test_cases.emplace_back(new test_dsv4_hc_post(4096, 21));
 
     // glu ops
     for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) {
@@ -8192,9 +8194,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
 
     for (ggml_type type_input : {GGML_TYPE_F32}) {
         for (ggml_op_pool pool_type : {GGML_OP_POOL_AVG, GGML_OP_POOL_MAX}) {
-            for (int k0 : {1, 3}) {
-                for (int s0 : {1, 2}) {
-                    for (int p0 : {0, 1}) {
+            for (int k0 : {1, 2, 3}) {
+                for (int s0 : {1, 2, 3}) {
+                    for (int p0 : {0, 1, 2, 3}) {
                         test_cases.emplace_back(new test_pool1d(pool_type, type_input, { 10,  3, 2, 1 }, k0, s0, p0));
                         test_cases.emplace_back(new test_pool1d(pool_type, type_input, { 11,  1, 3, 2 }, k0, s0, p0));
                         test_cases.emplace_back(new test_pool1d(pool_type, type_input, { 128, 2, 1, 3 }, k0, s0, p0));
@@ -8511,6 +8513,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         test_cases.emplace_back(new test_repeat(GGML_TYPE_F32, {10, 5, 4, ne3}, {1, 2, 1, 1}));
         test_cases.emplace_back(new test_repeat(GGML_TYPE_F32, {10, 5, 4, ne3}, {1, 1, 2, 1}));
         test_cases.emplace_back(new test_repeat(GGML_TYPE_F32, {10, 5, 4, ne3}, {1, 1, 1, 2}));
+        test_cases.emplace_back(new test_repeat(GGML_TYPE_F16, {10, 5, 4, ne3}, {2, 1, 1, 1}));
         test_cases.emplace_back(new test_repeat(GGML_TYPE_I32, {10, 5, 4, ne3}, {2, 1, 1, 1}));
         test_cases.emplace_back(new test_repeat(GGML_TYPE_I16, {10, 5, 4, ne3}, {1, 1, 1, 2}));
         test_cases.emplace_back(new test_repeat(GGML_TYPE_BF16, {10, 5, 4, ne3}, {2, 1, 1, 1}));
@@ -9513,6 +9516,18 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         }
     }
 
+    // prefill-shaped cases with long KV (nb >= 32, kv >= 1024): covers the
+    // XMX/GEMM-accelerated SYCL FA path which only activates for these shapes.
+    for (int kv : { 1024, 2048, }) {
+        for (int hs : { 64, 128, 256, }) {
+            for (int nb : { 32, 64, }) {
+                for (ggml_type type_KV : { GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, }) {
+                    test_cases.emplace_back(new test_flash_attn_ext(hs, hs, 8, {4, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, type_KV, type_KV));
+                }
+            }
+        }
+    }
+
     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 != 192 && hsk != 320 && hsk != 576 && hsk != hsv) continue;

tests/test-chat-peg-parser.cpp

@@ -8,6 +8,7 @@
 
 #include <iostream>
 #include <numeric>
+#include <regex>
 #include <string>
 
 #include "nlohmann/json.hpp"
@@ -21,6 +22,7 @@ static void test_example_qwen3_non_coder(testing & t);
 static void test_command7_parser_compare(testing & t);
 static void test_prefix_tool_names(testing & t);
 static void test_tagged_peg_parser(testing & t);
+static void test_permute(testing & t);
 
 int main(int argc, char * argv[]) {
     testing t(std::cout);
@@ -39,6 +41,7 @@ int main(int argc, char * argv[]) {
     t.test("comparison", test_command7_parser_compare);
     t.test("prefix tool names", test_prefix_tool_names);
     t.test("tagged peg parser", test_tagged_peg_parser);
+    t.test("permute", test_permute);
 
     return t.summary();
 }
@@ -981,3 +984,103 @@ static void test_tagged_peg_parser(testing & t) {
         t.assert_equal("fun_post should be '>'", ">", result.tags["fun_post"]);
     });
 }
+
+static void test_permute(testing & t) {
+    auto accepts = [](const common_peg_arena & parser, const std::string & input) {
+        common_peg_parse_context ctx(input);
+        return parser.parse(ctx).success();
+    };
+
+    auto gbnf_of = [](const common_peg_arena & parser) {
+        return build_grammar([&](const common_grammar_builder & builder) { parser.build_grammar(builder); });
+    };
+
+    auto assert_gbnf_equal = [](testing & t, const std::string & expected, const std::string & actual) {
+        static const std::regex leading_ws_re = std::regex(R"((^|\n)\s+)");
+        t.assert_equal("gbnf are equal", std::regex_replace(expected, leading_ws_re, "$1"), actual);
+    };
+
+    auto count_rules = [](const std::string & gbnf, const std::string & prefix) {
+        size_t count = 0;
+        for (const auto & line : string_split<std::string>(gbnf, '\n')) {
+            if (line.rfind(prefix, 0) == 0) {
+                count++;
+            }
+        }
+        return count;
+    };
+
+    t.test("accepts every ordering", [&](testing & t) {
+        auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) {
+            return p.permute("abc", { p.literal("a"), p.literal("b"), p.literal("c") }) + p.end();
+        });
+
+        for (const std::string input : { "abc", "acb", "bac", "bca", "cab", "cba" }) {
+            t.assert_true("accepts " + input, accepts(parser, input));
+        }
+    });
+
+    t.test("single element", [&](testing & t) {
+        auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) {
+            return p.permute("a", { p.literal("a") }) + p.end();
+        });
+
+        t.assert_true("accepts a", accepts(parser, "a"));
+        t.assert_true("rejects aa", !accepts(parser, "aa"));
+    });
+
+    t.test("grammar left-factorizes shared tails", [&](testing & t) {
+        auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) {
+            return p.permute("abc", { p.literal("a"), p.literal("b"), p.literal("c") }) + p.end();
+        });
+
+        // Every rule is one remaining subset, keyed by bitmask: abc-3 is {a,b}, abc-7 is {a,b,c}.
+        // Each subset is emitted once and shared by every branch that leads into it.
+        assert_gbnf_equal(t, R"""(
+            abc-1 ::= "a"
+            abc-2 ::= "b"
+            abc-3 ::= "a" abc-2 | "b" abc-1
+            abc-4 ::= "c"
+            abc-5 ::= "a" abc-4 | "c" abc-1
+            abc-6 ::= "b" abc-4 | "c" abc-2
+            abc-7 ::= "a" abc-6 | "b" abc-5 | "c" abc-3
+            root ::= abc-7
+            space ::= | " " | "\n"{1,2} [ \t]{0,20}
+        )""", gbnf_of(parser));
+    });
+
+    t.test("grammar emits one rule per remaining subset", [&](testing & t) {
+        auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) {
+            return p.permute("abcd", { p.literal("a"), p.literal("b"), p.literal("c"), p.literal("d") }) + p.end();
+        });
+
+        // 2^4 - 1 non-empty subsets, one rule each - not the 4! = 24 orderings.
+        t.assert_equal("permute rule count", 15u, count_rules(gbnf_of(parser), "abcd-"));
+    });
+
+    t.test("grammar emits no rules for a single element", [&](testing & t) {
+        auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) {
+            return p.permute("a", { p.literal("a") }) + p.end();
+        });
+
+        assert_gbnf_equal(t, R"""(
+            root ::= "a"
+            space ::= | " " | "\n"{1,2} [ \t]{0,20}
+        )""", gbnf_of(parser));
+    });
+
+    t.test("grammar falls back to the given order when too large", [&](testing & t) {
+        auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) {
+            std::vector<common_peg_parser> parsers;
+            for (size_t i = 0; i <= COMMON_CHAT_MAX_PERMUTE; i++) {
+                parsers.push_back(p.literal(std::string(1, (char) ('a' + i))));
+            }
+            return p.permute("big", parsers) + p.end();
+        });
+
+        assert_gbnf_equal(t, R"""(
+            root ::= "a" "b" "c" "d" "e" "f" "g"
+            space ::= | " " | "\n"{1,2} [ \t]{0,20}
+        )""", gbnf_of(parser));
+    });
+}

tests/test-chat.cpp

@@ -2278,46 +2278,39 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .expect_content(R"({"amount": 123.45, "date": "2025-12-03"})")
             .run();
 
-        // tool call segment in reasoning
+        // a tool call ends the prefilled thinking block, with or without a closing </think>
         tst.test(
-               "Let's call a tool: <tool_call>\n"
-               "<function=python>\n"
-               "<parameter=code>\n"
-               "def hello():\n"
-               "    print(\"Not the real call!\")\n"
-               "\n"
-               "hello()\n"
-               "</parameter>\n"
-               "</function>\n"
-               "</tool_call>\n</think>\n\n"
                "<tool_call>\n"
-               "<function=python>\n"
-               "<parameter=code>\n"
-               "def hello():\n"
-               "    print(\"Hello, world!\")\n"
-               "\n"
-               "hello()\n"
+               "<function=run_in_terminal>\n"
+               "<parameter=command>\n"
+               "pwd\n"
                "</parameter>\n"
                "</function>\n"
                "</tool_call>")
             .enable_thinking(true)
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
-            .tools({
-                python_tool
-        })
-            .expect_reasoning(
-                "Let's call a tool: <tool_call>\n"
-                "<function=python>\n"
-                "<parameter=code>\n"
-                "def hello():\n"
-                "    print(\"Not the real call!\")\n"
-                "\n"
-                "hello()\n"
-                "</parameter>\n"
-                "</function>\n"
-                "</tool_call>")
+            .tools({ run_in_terminal_tool })
             .expect_tool_calls({
-                { "python", "{\"code\": \"def hello():\\n    print(\\\"Hello, world!\\\")\\n\\nhello()\"}", {} },
+                { "run_in_terminal", R"({"command": "pwd"})", {} },
+            })
+            .run();
+
+        // ...including after the model has thought about it
+        tst.test(
+               "Need to inspect the current directory.\n"
+               "<tool_call>\n"
+               "<function=run_in_terminal>\n"
+               "<parameter=command>\n"
+               "pwd\n"
+               "</parameter>\n"
+               "</function>\n"
+               "</tool_call>")
+            .enable_thinking(true)
+            .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
+            .tools({ run_in_terminal_tool })
+            .expect_reasoning("Need to inspect the current directory.")
+            .expect_tool_calls({
+                { "run_in_terminal", R"({"command": "pwd"})", {} },
             })
             .run();
 
@@ -2461,17 +2454,6 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             })
             .run();
 
-        tst.test(
-               "I might call <tool_call> later, but I am still thinking.\n"
-               "</think>\n\n"
-               "Final answer without tools.")
-            .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
-            .enable_thinking(true)
-            .tools({ run_in_terminal_tool })
-            .expect_reasoning("I might call <tool_call> later, but I am still thinking.")
-            .expect_content("Final answer without tools.")
-            .run();
-
         // Continuation tests
         tst.test("world!\nWhat's up?")
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
@@ -2776,49 +2758,6 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .expect_content(R"({"amount": 123.45, "date": "2025-12-03"})")
             .run();
 
-        // tool call segment in reasoning
-        tst.test(
-               "Let's call a tool: <tool_call>\n"
-               "<function=python>\n"
-               "<parameter=code>\n"
-               "def hello():\n"
-               "    print(\"Not the real call!\")\n"
-               "\n"
-               "hello()\n"
-               "</parameter>\n"
-               "</function>\n"
-               "</tool_call>\n</think>\n"
-               "<tool_call>\n"
-               "<function=python>\n"
-               "<parameter=code>\n"
-               "def hello():\n"
-               "    print(\"Hello, world!\")\n"
-               "\n"
-               "hello()\n"
-               "</parameter>\n"
-               "</function>\n"
-               "</tool_call>\n"
-            )
-            .enable_thinking(true)
-            .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
-            .tools({
-                python_tool
-        })
-            .expect_reasoning("Let's call a tool: <tool_call>\n"
-               "<function=python>\n"
-               "<parameter=code>\n"
-               "def hello():\n"
-               "    print(\"Not the real call!\")\n"
-               "\n"
-               "hello()\n"
-               "</parameter>\n"
-               "</function>\n"
-               "</tool_call>\n")
-            .expect_tool_calls({
-                { "python", "{\"code\": \"def hello():\\n    print(\\\"Hello, world!\\\")\\n\\nhello()\"}", {} },
-            })
-            .run();
-
         // Continuation tests
         tst.test("world!\nWhat's up?")
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
@@ -3572,6 +3511,61 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .expect_reconstruction()
             .run();
 
+        // Some models skip the opening <tool_call> and go straight to <function=>
+        tst.test(
+               "<function=special_function>\n"
+               "<parameter=arg1>\n"
+               "1\n"
+               "</parameter>\n"
+               "</function>\n"
+               "</tool_call>")
+            .tools({ special_function_tool })
+            .expect(message_assist_call)
+            .run();
+
+        tst.test(
+               "Let me call it.\n"
+               "<function=special_function>\n"
+               "<parameter=arg1>\n"
+               "1\n"
+               "</parameter>\n"
+               "</function>\n"
+               "</tool_call>")
+            .tools({ special_function_tool })
+            .expect_content("Let me call it.\n")
+            .expect_tool_calls({
+                { "special_function", R"({"arg1": 1})", {} },
+            })
+            .run();
+
+        // Only the first call may omit it, the rest keep the </tool_call>\n<tool_call> separator
+        tst.test(
+               "<function=special_function>\n"
+               "<parameter=arg1>\n"
+               "1\n"
+               "</parameter>\n"
+               "</function>\n"
+               "</tool_call>\n"
+               "<tool_call>\n"
+               "<function=special_function_with_opt>\n"
+               "<parameter=arg1>\n"
+               "1\n"
+               "</parameter>\n"
+               "<parameter=arg2>\n"
+               "2\n"
+               "</parameter>\n"
+               "</function>\n"
+               "</tool_call>")
+            .parallel_tool_calls(true)
+            .tools({
+                special_function_tool, special_function_tool_with_optional_param
+        })
+            .expect_tool_calls({
+                { "special_function", R"({"arg1": 1})", {} },
+                { "special_function_with_opt", R"({"arg1": 1, "arg2": 2})", {} },
+            })
+            .run();
+
         tst.test(
                "<tool_call>\n"
                "<function=special_function>\n"
@@ -3680,6 +3674,37 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .expect_reconstruction()
             .run();
 
+        // Test flexible required argument ordering (required args still come first, in any order)
+        tst.test(
+               "<tool_call>\n"
+               "<function=edit>\n"
+               "<parameter=newString>\n#include\n</parameter>\n"
+               "<parameter=filename>\nfoo.c\n</parameter>\n"
+               "<parameter=oldString>\n#iclunde\n</parameter>\n"
+               "</function>\n"
+               "</tool_call>")
+            .tools({ edit_tool })
+            .expect_tool_calls({
+                { "edit", R"({"newString": "#include", "filename": "foo.c", "oldString": "#iclunde"})", {} },
+            })
+            .expect_reconstruction()
+            .run();
+
+        tst.test(
+               "<tool_call>\n"
+               "<function=tool_2req_4opt>\n"
+               "<parameter=req2>\n42\n</parameter>\n"
+               "<parameter=req1>\nhello\n</parameter>\n"
+               "<parameter=opt2>\n200\n</parameter>\n"
+               "</function>\n"
+               "</tool_call>")
+            .tools({ tool_2req_4opt })
+            .expect_tool_calls({
+                { "tool_2req_4opt", R"({"req2": 42, "req1": "hello", "opt2": 200})", {} },
+            })
+            .expect_reconstruction()
+            .run();
+
         // Test flexible optional argument ordering (2 required + 4 optional, reversed optional order)
         tst.test(
                "<tool_call>\n"
@@ -4227,6 +4252,20 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .expect_reasoning("I'm thinking")
             .expect_content("Hello, world!\nWhat's up?")
             .run();
+
+        tst.test(
+            "Let me check the time\n\n"
+            "<｜DSML｜tool_calls>\n"
+            "<｜DSML｜invoke name=\"get_time\">\n"
+            "<｜DSML｜parameter name=\"city\" string=\"true\">Tokyo</｜DSML｜parameter>\n"
+            "</｜DSML｜invoke>\n"
+            "</｜DSML｜tool_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();
     }
 
     // GLM-4.6 tests - format: <tool_call>function_name\n<arg_key>...</arg_key>\n<arg_value>...</arg_value>\n</tool_call>

tests/test-llama-archs.cpp

@@ -430,7 +430,7 @@ static bool arch_supported(const llm_arch arch) {
 
     // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI.
 #ifdef GGML_USE_WEBGPU
-    if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA) {
+    if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_MINIMAX_M3) {
         return false;
     }
 #endif // GGML_USE_WEBGPU

tests/test-recurrent-state-rollback.cpp

@@ -83,27 +83,33 @@ int main(int argc, char ** argv) {
     if (llama_vocab_type(vocab) == LLAMA_VOCAB_TYPE_NONE) {
         tokens = { 1, 2, 3, 4, 5, 6, 7, 8, 9 };
     } else {
-        tokens = common_tokenize(ctx_src, "The quick brown fox jumps", true);
+        tokens = common_tokenize(ctx_src, "The quick brown fox jumps over the lazy dog", true);
     }
     const uint32_t n_rs_seq = llama_n_rs_seq(ctx_src);
-    if (tokens.size() > n_rs_seq + 1) {
-        tokens.resize(n_rs_seq + 1);
+    constexpr uint32_t n_rollback = 3;
+    if (n_rs_seq < n_rollback) {
+        fprintf(stderr, "%s : skipping because n_rs_seq is too small\n", __func__);
+        llama_free(ctx_src);
+        llama_free(ctx_dst);
+        return 0;
     }
-    if (tokens.size() < 2) {
+    if (tokens.empty()) {
         fprintf(stderr, "%s : not enough prompt tokens\n", __func__);
         return 1;
     }
-    const uint32_t    n_tokens = tokens.size();
-    const llama_token last_tok = tokens.back();
-    const llama_pos   last_pos = (llama_pos) n_tokens - 2;
+    tokens.resize(n_rs_seq + 1, tokens.back());
 
-    // Decode the full prompt on the source, then roll back the last position.
+    const uint32_t  n_tokens     = tokens.size();
+    const llama_pos rollback_pos = (llama_pos) n_tokens - n_rollback;
+
+    // Decode the full prompt on the source, then roll back three positions.
+    // Replaying them crosses DSV4's ratio-4 compressor boundary.
     // Rollback leaves the recurrent memory in a snapshot state (rs_idx != 0).
     if (!decode_tokens(ctx_src, tokens, n_tokens)) {
         fprintf(stderr, "%s : failed to decode prompt\n", __func__);
         return 1;
     }
-    if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, last_pos, -1)) {
+    if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, rollback_pos, -1)) {
         fprintf(stderr, "%s : rollback failed\n", __func__);
         return 1;
     }
@@ -113,31 +119,56 @@ int main(int argc, char ** argv) {
     ckpt.update_tgt(ctx_src, 0, 0);
     ckpt.load_tgt(ctx_dst, 0, 0);
 
-    // Replay the rolled-back token on both contexts and compare logits.
-    if (!decode_one(ctx_src, last_tok, last_pos) ||
-        !decode_one(ctx_dst, last_tok, last_pos)) {
-        fprintf(stderr, "%s : replay failed\n", __func__);
-        return 1;
-    }
-
-    const float * logits_src = llama_get_logits_ith(ctx_src, 0);
-    const float * logits_dst = llama_get_logits_ith(ctx_dst, 0);
-    if (logits_src == nullptr || logits_dst == nullptr) {
-        fprintf(stderr, "%s : missing logits\n", __func__);
-        return 1;
-    }
-
     constexpr float eps = 1e-5f;
-    for (int i = 0; i < n_vocab; ++i) {
-        if (std::fabs(logits_src[i] - logits_dst[i]) > eps) {
-            fprintf(stderr, "%s : logits mismatch at token %d (%g != %g)\n",
-                    __func__, i, (double) logits_src[i], (double) logits_dst[i]);
-            return 1;
+    std::vector<std::vector<float>> logits_src_replay(n_rollback);
+    const auto replay_and_compare = [&](const char * mode) {
+        for (uint32_t i = 0; i < n_rollback; ++i) {
+            const llama_pos pos = rollback_pos + i;
+            if (!decode_one(ctx_src, tokens[pos], pos) ||
+                !decode_one(ctx_dst, tokens[pos], pos)) {
+                fprintf(stderr, "%s : %s replay failed at position %d\n", __func__, mode, pos);
+                return false;
+            }
+
+            const float * logits_src = llama_get_logits_ith(ctx_src, 0);
+            const float * logits_dst = llama_get_logits_ith(ctx_dst, 0);
+            if (logits_src == nullptr || logits_dst == nullptr) {
+                fprintf(stderr, "%s : missing %s logits at position %d\n", __func__, mode, pos);
+                return false;
+            }
+
+            logits_src_replay[i].assign(logits_src, logits_src + n_vocab);
+            for (int token = 0; token < n_vocab; ++token) {
+                if (std::fabs(logits_src[token] - logits_dst[token]) > eps) {
+                    fprintf(stderr, "%s : %s logits mismatch at position %d, token %d (%g != %g)\n",
+                            __func__, mode, pos, token, (double) logits_src[token], (double) logits_dst[token]);
+                    return false;
+                }
+            }
         }
+        return true;
+    };
+    if (!replay_and_compare("full")) {
+        return 1;
+    }
+
+    if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, rollback_pos, -1) ||
+        !llama_memory_seq_rm(llama_get_memory(ctx_dst), 0, rollback_pos, -1)) {
+        fprintf(stderr, "%s : partial rollback failed\n", __func__);
+        return 1;
+    }
+
+    constexpr llama_state_seq_flags partial_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY;
+    common_prompt_checkpoint ckpt_partial;
+    ckpt_partial.update_tgt(ctx_src, 0, partial_flags);
+    ckpt_partial.load_tgt(ctx_dst, 0, partial_flags);
+
+    if (!replay_and_compare("partial")) {
+        return 1;
     }
 
     // Repeat the load into a context that already has its own rollback state:
-    // groups 1..n_rs_seq hold a *different* prompt's history, and rs_idx[0] is
+    // groups 1..n_rs_seq hold a different prompt's history, and rs_idx[0] is
     // non-zero at load time. The restore must wipe that state and still match.
     llama_context * ctx_dirty = make_ctx(params, model);
     if (ctx_dirty == nullptr) {
@@ -156,30 +187,33 @@ int main(int argc, char ** argv) {
         fprintf(stderr, "%s : dirty prompt decode failed\n", __func__);
         return 1;
     }
-    if (!llama_memory_seq_rm(llama_get_memory(ctx_dirty), 0, last_pos, -1)) {
+    if (!llama_memory_seq_rm(llama_get_memory(ctx_dirty), 0, rollback_pos, -1)) {
         fprintf(stderr, "%s : dirty rollback failed\n", __func__);
         return 1;
     }
 
     ckpt.load_tgt(ctx_dirty, 0, 0);
 
-    if (!decode_one(ctx_dirty, last_tok, last_pos)) {
-        fprintf(stderr, "%s : dirty replay failed\n", __func__);
-        return 1;
-    }
-
-    const float * logits_dirty = llama_get_logits_ith(ctx_dirty, 0);
-    if (logits_dirty == nullptr) {
-        fprintf(stderr, "%s : missing dirty logits\n", __func__);
-        return 1;
-    }
-
-    for (int i = 0; i < n_vocab; ++i) {
-        if (std::fabs(logits_src[i] - logits_dirty[i]) > eps) {
-            fprintf(stderr, "%s : dirty-ctx logits mismatch at token %d (%g != %g)\n",
-                    __func__, i, (double) logits_src[i], (double) logits_dirty[i]);
+    for (uint32_t i = 0; i < n_rollback; ++i) {
+        const llama_pos pos = rollback_pos + i;
+        if (!decode_one(ctx_dirty, tokens[pos], pos)) {
+            fprintf(stderr, "%s : dirty replay failed at position %d\n", __func__, pos);
             return 1;
         }
+
+        const float * logits_dirty = llama_get_logits_ith(ctx_dirty, 0);
+        if (logits_dirty == nullptr) {
+            fprintf(stderr, "%s : missing dirty logits at position %d\n", __func__, pos);
+            return 1;
+        }
+
+        for (int token = 0; token < n_vocab; ++token) {
+            if (std::fabs(logits_src_replay[i][token] - logits_dirty[token]) > eps) {
+                fprintf(stderr, "%s : dirty-ctx logits mismatch at position %d, token %d (%g != %g)\n",
+                        __func__, pos, token, (double) logits_src_replay[i][token], (double) logits_dirty[token]);
+                return 1;
+            }
+        }
     }
 
     fprintf(stderr, "%s : recurrent rollback checkpoint restored successfully\n", __func__);

tools/cli/cli-context.cpp

@@ -624,10 +624,14 @@ int cli_context::run() {
         generated_content content;
         generate_completion(content, timings);
 
-        impl->messages.push_back({
+        json assistant_msg = {
             {"role",    "assistant"},
             {"content", content.content}
-        });
+        };
+        if (!content.reasoning.empty()) {
+            assistant_msg["reasoning_content"] = content.reasoning;
+        }
+        impl->messages.push_back(std::move(assistant_msg));
 
         if (output_file) {
             std::string out_content = "Assistant:\n";

tools/mtmd/clip-impl.h

@@ -41,6 +41,7 @@
 #define KEY_PROJ_DIM            "clip.%s.projection_dim"
 #define KEY_N_HEAD              "clip.%s.attention.head_count"
 #define KEY_N_HEAD_KV           "clip.%s.attention.head_count_kv"
+#define KEY_N_EMBD_HEAD         "clip.%s.attention.head_dim"
 #define KEY_LAYER_NORM_EPS      "clip.%s.attention.layer_norm_epsilon"
 #define KEY_FEATURE_LAYERS      "clip.%s.feature_layer"
 

tools/mtmd/clip-model.h

@@ -33,7 +33,7 @@ enum resize_algo {
     RESIZE_ALGO_BILINEAR, // stretch to target resolution
     RESIZE_ALGO_BICUBIC, // center-crop when aspect ratio doesn't match
     RESIZE_ALGO_BICUBIC_PILLOW,
-    // RESIZE_ALGO_LANCZOS, // TODO
+    RESIZE_ALGO_LANCZOS,
 };
 
 // Padding style for img_tool::resize
@@ -54,6 +54,8 @@ struct clip_hparams {
     int32_t projection_dim = 0;
     int32_t n_head = 0;
     int32_t n_head_kv = 0;
+    // 0 = derive from n_embd; set when qkv width != n_embd
+    int32_t n_embd_head = 0;
     int32_t n_layer = 0;
     int32_t n_merge = 1; // number of patch merges **per-side**
 

tools/mtmd/clip.cpp

@@ -253,7 +253,7 @@ clip_graph::clip_graph(clip_ctx * ctx, const clip_image_f32 & img) :
         n_embd(hparams.n_embd),
         n_head(hparams.n_head),
         n_head_kv(hparams.n_head_kv),
-        d_head(n_head > 0 ? n_embd / n_head : 0),
+        d_head(hparams.n_embd_head > 0 ? hparams.n_embd_head : (n_head > 0 ? n_embd / n_head : 0)),
         n_layer(hparams.n_layer),
         n_mmproj_embd(clip_n_mmproj_embd(ctx)),
         eps(hparams.eps),
@@ -372,13 +372,13 @@ ggml_tensor * clip_graph::build_vit(
                 /* nb1    */ ggml_row_size(cur->type, d_head),
                 /* nb2    */ cur->nb[1],
                 /* nb3    */ cur->nb[1] * n_pos,
-                /* offset */ ggml_row_size(cur->type, n_embd));
+                /* offset */ ggml_row_size(cur->type, n_head * d_head));
 
                 Vcur = ggml_view_4d(ctx0, cur, d_head, n_head, n_pos, B,
                 /* nb1    */ ggml_row_size(cur->type, d_head),
                 /* nb2    */ cur->nb[1],
                 /* nb3    */ cur->nb[1] * n_pos,
-                /* offset */ ggml_row_size(cur->type, 2 * n_embd));
+                /* offset */ ggml_row_size(cur->type, 2 * n_head * d_head));
 
                 if (layer.q_norm) {
                     GGML_ASSERT(layer.q_norm->ne[0] == Qcur->ne[0]);
@@ -1190,6 +1190,7 @@ struct clip_model_loader {
             const char * prefix = is_vision ? "vision" : "audio";
             get_u32(string_format(KEY_N_EMBD,         prefix), hparams.n_embd);
             get_u32(string_format(KEY_N_HEAD,         prefix), hparams.n_head);
+            get_u32(string_format(KEY_N_EMBD_HEAD,    prefix), hparams.n_embd_head, false);
             get_u32(string_format(KEY_N_FF,           prefix), hparams.n_ff);
             get_u32(string_format(KEY_N_BLOCK,        prefix), hparams.n_layer);
             get_u32(string_format(KEY_PROJ_DIM,       prefix), hparams.projection_dim);
@@ -1336,6 +1337,7 @@ struct clip_model_loader {
                         // ViT merger 2x2 + final merger 2x2 = 4x spatial merge per dimension
                         hparams.n_merge = 4;
                         get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge, false);
+                        GGML_ASSERT(hparams.n_merge == 2 || hparams.n_merge == 4);
 
                         // borrow wa_layer_indexes for vit_merger insertion point
                         std::vector<int> wa_layer_indexes_vec;
@@ -2142,24 +2144,29 @@ struct clip_model_loader {
                 } break;
             case PROJECTOR_TYPE_MINICPMV4_6:
                 {
+                    const bool merger_required = hparams.n_merge == 4;
+                    auto get_merger_tensor = [&](const std::string & name, bool required = true) {
+                        return get_tensor(name, merger_required && required);
+                    };
+
                     // ViT merger: window self-attention
-                    model.vit_merger_ln1_w     = get_tensor(string_format(TN_VIT_MERGER_LN1, "weight"));
-                    model.vit_merger_ln1_b     = get_tensor(string_format(TN_VIT_MERGER_LN1, "bias"));
-                    model.vit_merger_attn_q_w  = get_tensor(string_format(TN_VIT_MERGER_ATTN_Q, "weight"));
-                    model.vit_merger_attn_q_b  = get_tensor(string_format(TN_VIT_MERGER_ATTN_Q, "bias"), false);
-                    model.vit_merger_attn_k_w  = get_tensor(string_format(TN_VIT_MERGER_ATTN_K, "weight"));
-                    model.vit_merger_attn_k_b  = get_tensor(string_format(TN_VIT_MERGER_ATTN_K, "bias"), false);
-                    model.vit_merger_attn_v_w  = get_tensor(string_format(TN_VIT_MERGER_ATTN_V, "weight"));
-                    model.vit_merger_attn_v_b  = get_tensor(string_format(TN_VIT_MERGER_ATTN_V, "bias"), false);
-                    model.vit_merger_attn_o_w  = get_tensor(string_format(TN_VIT_MERGER_ATTN_O, "weight"));
-                    model.vit_merger_attn_o_b  = get_tensor(string_format(TN_VIT_MERGER_ATTN_O, "bias"), false);
+                    model.vit_merger_ln1_w     = get_merger_tensor(string_format(TN_VIT_MERGER_LN1, "weight"));
+                    model.vit_merger_ln1_b     = get_merger_tensor(string_format(TN_VIT_MERGER_LN1, "bias"));
+                    model.vit_merger_attn_q_w  = get_merger_tensor(string_format(TN_VIT_MERGER_ATTN_Q, "weight"));
+                    model.vit_merger_attn_q_b  = get_merger_tensor(string_format(TN_VIT_MERGER_ATTN_Q, "bias"), false);
+                    model.vit_merger_attn_k_w  = get_merger_tensor(string_format(TN_VIT_MERGER_ATTN_K, "weight"));
+                    model.vit_merger_attn_k_b  = get_merger_tensor(string_format(TN_VIT_MERGER_ATTN_K, "bias"), false);
+                    model.vit_merger_attn_v_w  = get_merger_tensor(string_format(TN_VIT_MERGER_ATTN_V, "weight"));
+                    model.vit_merger_attn_v_b  = get_merger_tensor(string_format(TN_VIT_MERGER_ATTN_V, "bias"), false);
+                    model.vit_merger_attn_o_w  = get_merger_tensor(string_format(TN_VIT_MERGER_ATTN_O, "weight"));
+                    model.vit_merger_attn_o_b  = get_merger_tensor(string_format(TN_VIT_MERGER_ATTN_O, "bias"), false);
                     // ViT merger: MLP downsample
-                    model.vit_merger_ds_ln_w   = get_tensor(string_format(TN_VIT_MERGER_DS_LN, "weight"));
-                    model.vit_merger_ds_ln_b   = get_tensor(string_format(TN_VIT_MERGER_DS_LN, "bias"));
-                    model.vit_merger_ds_up_w   = get_tensor(string_format(TN_VIT_MERGER_DS_UP, "weight"));
-                    model.vit_merger_ds_up_b   = get_tensor(string_format(TN_VIT_MERGER_DS_UP, "bias"), false);
-                    model.vit_merger_ds_down_w = get_tensor(string_format(TN_VIT_MERGER_DS_DOWN, "weight"));
-                    model.vit_merger_ds_down_b = get_tensor(string_format(TN_VIT_MERGER_DS_DOWN, "bias"), false);
+                    model.vit_merger_ds_ln_w   = get_merger_tensor(string_format(TN_VIT_MERGER_DS_LN, "weight"));
+                    model.vit_merger_ds_ln_b   = get_merger_tensor(string_format(TN_VIT_MERGER_DS_LN, "bias"));
+                    model.vit_merger_ds_up_w   = get_merger_tensor(string_format(TN_VIT_MERGER_DS_UP, "weight"));
+                    model.vit_merger_ds_up_b   = get_merger_tensor(string_format(TN_VIT_MERGER_DS_UP, "bias"), false);
+                    model.vit_merger_ds_down_w = get_merger_tensor(string_format(TN_VIT_MERGER_DS_DOWN, "weight"));
+                    model.vit_merger_ds_down_b = get_merger_tensor(string_format(TN_VIT_MERGER_DS_DOWN, "bias"), false);
                     // Final Merger (DownsampleMLP)
                     model.mm_input_norm_w = get_tensor(TN_MM_INP_NORM);
                     model.mm_input_norm_b = get_tensor(TN_MM_INP_NORM_B, false);
@@ -3590,8 +3597,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
             } break;
         case PROJECTOR_TYPE_MINICPMV4_6:
             {
-                // ViT merger 4x + final merger 4x = 16x total spatial downsample
-                n_patches = n_patches / 16;
+                n_patches /= params.n_merge * params.n_merge;
             } break;
         case PROJECTOR_TYPE_QWEN2VL:
         case PROJECTOR_TYPE_QWEN25VL:
@@ -3973,6 +3979,8 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32
             } break;
         case PROJECTOR_TYPE_MINICPMV4_6:
             {
+                const bool is_4x = hparams.n_merge == 2;
+
                 // SigLIP position buckets (same as resampler path)
                 std::vector<int32_t> positions(pos_h * pos_w);
                 int bucket_coords_h[1024];
@@ -3993,40 +4001,6 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32
                 const int half_h = pos_h / 2;
                 const int half_w = pos_w / 2;
 
-                // window reorder indices for 2x2 windows
-                std::vector<int32_t> window_idx(n_pos);
-                std::vector<int32_t> inv_window_idx(n_pos);
-                {
-                    int k = 0;
-                    for (int wi = 0; wi < half_h; wi++) {
-                        for (int wj = 0; wj < half_w; wj++) {
-                            window_idx[k++] = (2*wi    ) * pos_w + (2*wj    );
-                            window_idx[k++] = (2*wi    ) * pos_w + (2*wj + 1);
-                            window_idx[k++] = (2*wi + 1) * pos_w + (2*wj    );
-                            window_idx[k++] = (2*wi + 1) * pos_w + (2*wj + 1);
-                        }
-                    }
-                    for (int i = 0; i < n_pos; i++) {
-                        inv_window_idx[window_idx[i]] = i;
-                    }
-                }
-                set_input_i32("vit_merger_window_idx",     window_idx);
-                set_input_i32("vit_merger_inv_window_idx", inv_window_idx);
-
-                // block-diagonal attention mask: tokens in the same 4-token
-                // window attend to each other (mask = 0), all other positions
-                // are masked out (-inf). matches the window-major reorder above.
-                std::vector<float> window_mask_data(n_pos * n_pos, std::numeric_limits<float>::lowest());
-                for (int wi = 0; wi < n_pos / 4; wi++) {
-                    for (int i = 0; i < 4; i++) {
-                        for (int j = 0; j < 4; j++) {
-                            window_mask_data[(wi*4 + i) * n_pos + (wi*4 + j)] = 0.0f;
-                        }
-                    }
-                }
-                set_input_f32("vit_merger_window_mask", window_mask_data);
-
-                // ViT merger 2x2 downsample indices
                 auto make_ds_idx = [](int off_r, int off_c, int ds_h, int ds_w, int stride_w) {
                     std::vector<int32_t> idx(ds_h * ds_w);
                     for (int i = 0; i < ds_h; i++) {
@@ -4036,22 +4010,58 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32
                     }
                     return idx;
                 };
-                auto vit_merger_ds_0 = make_ds_idx(0, 0, half_h, half_w, pos_w);
-                auto vit_merger_ds_1 = make_ds_idx(0, 1, half_h, half_w, pos_w);
-                auto vit_merger_ds_2 = make_ds_idx(1, 0, half_h, half_w, pos_w);
-                auto vit_merger_ds_3 = make_ds_idx(1, 1, half_h, half_w, pos_w);
-                set_input_i32("vit_merger_ds_idx_0", vit_merger_ds_0);
-                set_input_i32("vit_merger_ds_idx_1", vit_merger_ds_1);
-                set_input_i32("vit_merger_ds_idx_2", vit_merger_ds_2);
-                set_input_i32("vit_merger_ds_idx_3", vit_merger_ds_3);
 
-                // final merger 2x2 downsample indices (operates on half_h x half_w grid)
-                const int qh = half_h / 2;
-                const int qw = half_w / 2;
-                auto m_ds_0 = make_ds_idx(0, 0, qh, qw, half_w);
-                auto m_ds_1 = make_ds_idx(0, 1, qh, qw, half_w);
-                auto m_ds_2 = make_ds_idx(1, 0, qh, qw, half_w);
-                auto m_ds_3 = make_ds_idx(1, 1, qh, qw, half_w);
+                if (!is_4x) {
+                    // window reorder indices for 2x2 windows
+                    std::vector<int32_t> window_idx(n_pos);
+                    std::vector<int32_t> inv_window_idx(n_pos);
+                    {
+                        int k = 0;
+                        for (int wi = 0; wi < half_h; wi++) {
+                            for (int wj = 0; wj < half_w; wj++) {
+                                window_idx[k++] = (2*wi    ) * pos_w + (2*wj    );
+                                window_idx[k++] = (2*wi    ) * pos_w + (2*wj + 1);
+                                window_idx[k++] = (2*wi + 1) * pos_w + (2*wj    );
+                                window_idx[k++] = (2*wi + 1) * pos_w + (2*wj + 1);
+                            }
+                        }
+                        for (int i = 0; i < n_pos; i++) {
+                            inv_window_idx[window_idx[i]] = i;
+                        }
+                    }
+                    set_input_i32("vit_merger_window_idx",     window_idx);
+                    set_input_i32("vit_merger_inv_window_idx", inv_window_idx);
+
+                    // block-diagonal attention mask: tokens in the same 4-token
+                    // window attend to each other (mask = 0), all other positions
+                    // are masked out (-inf). matches the window-major reorder above.
+                    std::vector<float> window_mask_data(n_pos * n_pos, std::numeric_limits<float>::lowest());
+                    for (int wi = 0; wi < n_pos / 4; wi++) {
+                        for (int i = 0; i < 4; i++) {
+                            for (int j = 0; j < 4; j++) {
+                                window_mask_data[(wi*4 + i) * n_pos + (wi*4 + j)] = 0.0f;
+                            }
+                        }
+                    }
+                    set_input_f32("vit_merger_window_mask", window_mask_data);
+
+                    // ViT merger 2x2 downsample indices
+                    auto vit_merger_ds_0 = make_ds_idx(0, 0, half_h, half_w, pos_w);
+                    auto vit_merger_ds_1 = make_ds_idx(0, 1, half_h, half_w, pos_w);
+                    auto vit_merger_ds_2 = make_ds_idx(1, 0, half_h, half_w, pos_w);
+                    auto vit_merger_ds_3 = make_ds_idx(1, 1, half_h, half_w, pos_w);
+                    set_input_i32("vit_merger_ds_idx_0", vit_merger_ds_0);
+                    set_input_i32("vit_merger_ds_idx_1", vit_merger_ds_1);
+                    set_input_i32("vit_merger_ds_idx_2", vit_merger_ds_2);
+                    set_input_i32("vit_merger_ds_idx_3", vit_merger_ds_3);
+                }
+
+                const int merger_h = is_4x ? pos_h : half_h;
+                const int merger_w = is_4x ? pos_w : half_w;
+                auto m_ds_0 = make_ds_idx(0, 0, merger_h / 2, merger_w / 2, merger_w);
+                auto m_ds_1 = make_ds_idx(0, 1, merger_h / 2, merger_w / 2, merger_w);
+                auto m_ds_2 = make_ds_idx(1, 0, merger_h / 2, merger_w / 2, merger_w);
+                auto m_ds_3 = make_ds_idx(1, 1, merger_h / 2, merger_w / 2, merger_w);
                 set_input_i32("merger_ds_idx_0", m_ds_0);
                 set_input_i32("merger_ds_idx_1", m_ds_1);
                 set_input_i32("merger_ds_idx_2", m_ds_2);

tools/mtmd/models/minicpmv.cpp

@@ -114,14 +114,12 @@ ggml_cgraph * clip_graph_minicpmv::build() {
 }
 
 ggml_cgraph * clip_graph_minicpmv4_6::build() {
-    const int insert_lid = hparams.insert_layer_id;
-    const int n_pos      = n_patches;
-    const int half_h     = n_patches_y / 2;
-    const int half_w     = n_patches_x / 2;
-    const int n_ds       = half_h * half_w;     // after ViT merger 2x2 downsample
-    const int qh         = half_h / 2;
-    const int qw         = half_w / 2;
-    const int n_ds2      = qh * qw;             // after final merger 2x2 downsample
+    const bool is_4x = hparams.n_merge == 2;
+    const int n_pos  = n_patches;
+    const int half_h = n_patches_y / 2;
+    const int half_w = n_patches_x / 2;
+    const int n_ds   = half_h * half_w;
+    const int n_out  = is_4x ? n_ds : (half_h / 2) * (half_w / 2);
 
     auto add_i32_input = [&](const char * name, int n) {
         ggml_tensor * t = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n);
@@ -134,29 +132,39 @@ ggml_cgraph * clip_graph_minicpmv4_6::build() {
     ggml_tensor * positions = add_i32_input("positions", n_pos);
     ggml_tensor * learned_pos_embd = ggml_get_rows(ctx0, model.position_embeddings, positions);
 
-    // ViT merger window reorder indices + block-diagonal mask
-    // (mask layout follows qwen2vl: -inf except for 4x4 blocks on the diagonal,
-    // so each window-major group of 4 tokens only attends to itself)
-    ggml_tensor * vit_merger_window_idx     = add_i32_input("vit_merger_window_idx", n_pos);
-    ggml_tensor * vit_merger_inv_window_idx = add_i32_input("vit_merger_inv_window_idx", n_pos);
-    ggml_tensor * vit_merger_window_mask    = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_pos, n_pos);
-    ggml_set_name(vit_merger_window_mask, "vit_merger_window_mask");
-    ggml_set_input(vit_merger_window_mask);
-    if (flash_attn_type == CLIP_FLASH_ATTN_TYPE_ENABLED) {
-        vit_merger_window_mask = ggml_cast(ctx0, vit_merger_window_mask, GGML_TYPE_F16);
+    ggml_tensor * vit_merger_window_idx     = nullptr;
+    ggml_tensor * vit_merger_inv_window_idx = nullptr;
+    ggml_tensor * vit_merger_window_mask    = nullptr;
+    ggml_tensor * vit_merger_ds_idx_0       = nullptr;
+    ggml_tensor * vit_merger_ds_idx_1       = nullptr;
+    ggml_tensor * vit_merger_ds_idx_2       = nullptr;
+    ggml_tensor * vit_merger_ds_idx_3       = nullptr;
+
+    if (!is_4x) {
+        // ViT merger window reorder indices + block-diagonal mask
+        // (mask layout follows qwen2vl: -inf except for 4x4 blocks on the diagonal,
+        // so each window-major group of 4 tokens only attends to itself)
+        vit_merger_window_idx     = add_i32_input("vit_merger_window_idx", n_pos);
+        vit_merger_inv_window_idx = add_i32_input("vit_merger_inv_window_idx", n_pos);
+        vit_merger_window_mask    = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_pos, n_pos);
+        ggml_set_name(vit_merger_window_mask, "vit_merger_window_mask");
+        ggml_set_input(vit_merger_window_mask);
+        if (flash_attn_type == CLIP_FLASH_ATTN_TYPE_ENABLED) {
+            vit_merger_window_mask = ggml_cast(ctx0, vit_merger_window_mask, GGML_TYPE_F16);
+        }
+
+        // ViT merger 2x2 downsample gather indices
+        vit_merger_ds_idx_0 = add_i32_input("vit_merger_ds_idx_0", n_ds);
+        vit_merger_ds_idx_1 = add_i32_input("vit_merger_ds_idx_1", n_ds);
+        vit_merger_ds_idx_2 = add_i32_input("vit_merger_ds_idx_2", n_ds);
+        vit_merger_ds_idx_3 = add_i32_input("vit_merger_ds_idx_3", n_ds);
     }
 
-    // ViT merger 2x2 downsample gather indices
-    ggml_tensor * vit_merger_ds_idx_0 = add_i32_input("vit_merger_ds_idx_0", n_ds);
-    ggml_tensor * vit_merger_ds_idx_1 = add_i32_input("vit_merger_ds_idx_1", n_ds);
-    ggml_tensor * vit_merger_ds_idx_2 = add_i32_input("vit_merger_ds_idx_2", n_ds);
-    ggml_tensor * vit_merger_ds_idx_3 = add_i32_input("vit_merger_ds_idx_3", n_ds);
-
     // final merger 2x2 downsample gather indices
-    ggml_tensor * merger_ds_idx_0 = add_i32_input("merger_ds_idx_0", n_ds2);
-    ggml_tensor * merger_ds_idx_1 = add_i32_input("merger_ds_idx_1", n_ds2);
-    ggml_tensor * merger_ds_idx_2 = add_i32_input("merger_ds_idx_2", n_ds2);
-    ggml_tensor * merger_ds_idx_3 = add_i32_input("merger_ds_idx_3", n_ds2);
+    ggml_tensor * merger_ds_idx_0 = add_i32_input("merger_ds_idx_0", n_out);
+    ggml_tensor * merger_ds_idx_1 = add_i32_input("merger_ds_idx_1", n_out);
+    ggml_tensor * merger_ds_idx_2 = add_i32_input("merger_ds_idx_2", n_out);
+    ggml_tensor * merger_ds_idx_3 = add_i32_input("merger_ds_idx_3", n_out);
 
     // patch embedding + positional embedding
     ggml_tensor * inp = build_inp();
@@ -169,150 +177,10 @@ ggml_cgraph * clip_graph_minicpmv4_6::build() {
         cb(inpL, "pre_ln", -1);
     }
 
-    // ViT layers 0..insert_layer_id (inclusive)
-    // Mirrors the separate-qkv path of clip_graph::build_vit so the two manually
-    // unrolled segments around the ViT merger read like build_vit() expansions.
-    for (int il = 0; il <= insert_lid; il++) {
-        auto & layer = model.layers[il];
-        ggml_tensor * cur = inpL;
-
-        cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il);
-        cb(cur, "layer_inp_normed", il);
-
-        {
-            ggml_tensor * Qcur = build_mm(layer.q_w, cur);
-            if (layer.q_b) {
-                Qcur = ggml_add(ctx0, Qcur, layer.q_b);
-            }
-            ggml_tensor * Kcur = build_mm(layer.k_w, cur);
-            if (layer.k_b) {
-                Kcur = ggml_add(ctx0, Kcur, layer.k_b);
-            }
-            ggml_tensor * Vcur = build_mm(layer.v_w, cur);
-            if (layer.v_b) {
-                Vcur = ggml_add(ctx0, Vcur, layer.v_b);
-            }
-
-            Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos);
-            Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos);
-            Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_pos);
-            cb(Qcur, "Qcur", il);
-            cb(Kcur, "Kcur", il);
-            cb(Vcur, "Vcur", il);
-
-            cur = build_attn(layer.o_w, layer.o_b, Qcur, Kcur, Vcur, nullptr, kq_scale, il);
-            cb(cur, "attn_out", il);
-        }
-
-        if (layer.ls_1_w) {
-            cur = ggml_mul(ctx0, cur, layer.ls_1_w);
-            cb(cur, "attn_out_scaled", il);
-        }
-        cur = ggml_add(ctx0, cur, inpL);
-        inpL = cur;
-        cb(cur, "ffn_inp", il);
-
-        cur = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il);
-        cb(cur, "ffn_inp_normed", il);
-
-        cur = build_ffn(cur, layer.ff_up_w, layer.ff_up_b, layer.ff_gate_w, layer.ff_gate_b,
-                        layer.ff_down_w, layer.ff_down_b, hparams.ffn_op, il);
-        cb(cur, "ffn_out", il);
-
-        if (layer.ls_2_w) {
-            cur = ggml_mul(ctx0, cur, layer.ls_2_w);
-            cb(cur, "ffn_out_scaled", il);
-        }
-        cur = ggml_add(ctx0, inpL, cur);
-        cb(cur, "layer_out", il);
-
-        inpL = cur;
-    }
-
-    // ViT merger: window self-attention
-    // Tokens are reordered to window-major (4 tokens per window are contiguous),
-    // and a block-diagonal mask restricts attention to within each window. This
-    // mirrors the qwen2vl windowed-attention pattern so build_attn() can pick the
-    // flash-attention path when available.
-    {
-        ggml_tensor * residual = inpL;
-        ggml_tensor * cur = build_norm(inpL,
-            model.vit_merger_ln1_w, model.vit_merger_ln1_b,
-            NORM_TYPE_NORMAL, eps, -1);
-        cb(cur, "vit_merger_attn_inp_normed", -1);
-
-        cur = ggml_get_rows(ctx0, cur, vit_merger_window_idx);
-        cb(cur, "vit_merger_window_reorder", -1);
-
-        ggml_tensor * Qcur = build_mm(model.vit_merger_attn_q_w, cur);
-        if (model.vit_merger_attn_q_b) {
-            Qcur = ggml_add(ctx0, Qcur, model.vit_merger_attn_q_b);
-        }
-        ggml_tensor * Kcur = build_mm(model.vit_merger_attn_k_w, cur);
-        if (model.vit_merger_attn_k_b) {
-            Kcur = ggml_add(ctx0, Kcur, model.vit_merger_attn_k_b);
-        }
-        ggml_tensor * Vcur = build_mm(model.vit_merger_attn_v_w, cur);
-        if (model.vit_merger_attn_v_b) {
-            Vcur = ggml_add(ctx0, Vcur, model.vit_merger_attn_v_b);
-        }
-
-        Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos);
-        Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos);
-        Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_pos);
-        cb(Qcur, "vit_merger_Qcur", -1);
-        cb(Kcur, "vit_merger_Kcur", -1);
-        cb(Vcur, "vit_merger_Vcur", -1);
-
-        cur = build_attn(model.vit_merger_attn_o_w, model.vit_merger_attn_o_b,
-                         Qcur, Kcur, Vcur, vit_merger_window_mask, kq_scale, -1);
-        cb(cur, "vit_merger_attn_out", -1);
-
-        cur = ggml_get_rows(ctx0, cur, vit_merger_inv_window_idx);
-        inpL = ggml_add(ctx0, cur, residual);
-        cb(inpL, "vit_merger_attn_residual", -1);
-    }
-
-    // ViT merger: 2x2 spatial downsample + MLP (4 tokens -> 1)
-    {
-        ggml_tensor * p0 = ggml_get_rows(ctx0, inpL, vit_merger_ds_idx_0);
-        ggml_tensor * p1 = ggml_get_rows(ctx0, inpL, vit_merger_ds_idx_1);
-        ggml_tensor * p2 = ggml_get_rows(ctx0, inpL, vit_merger_ds_idx_2);
-        ggml_tensor * p3 = ggml_get_rows(ctx0, inpL, vit_merger_ds_idx_3);
-
-        ggml_tensor * mean_res = ggml_add(ctx0, p0, p1);
-        mean_res = ggml_add(ctx0, mean_res, p2);
-        mean_res = ggml_add(ctx0, mean_res, p3);
-        mean_res = ggml_scale(ctx0, mean_res, 0.25f);
-        cb(mean_res, "vit_merger_ds_mean_res", -1);
-
-        ggml_tensor * cat = ggml_concat(ctx0, p0, p1, 0);
-        cat = ggml_concat(ctx0, cat, p2, 0);
-        cat = ggml_concat(ctx0, cat, p3, 0);
-
-        ggml_tensor * cur = build_norm(cat,
-            model.vit_merger_ds_ln_w, model.vit_merger_ds_ln_b,
-            NORM_TYPE_NORMAL, eps, -1);
-        cb(cur, "vit_merger_ds_normed", -1);
-
-        // ViTWindowAttentionMerger downsample MLP uses gelu_pytorch_tanh (FFN_GELU)
-        cur = build_ffn(cur,
-            model.vit_merger_ds_up_w,   model.vit_merger_ds_up_b,
-            nullptr, nullptr,
-            model.vit_merger_ds_down_w, model.vit_merger_ds_down_b,
-            FFN_GELU, -1);
-        cb(cur, "vit_merger_ds_mlp_out", -1);
-
-        inpL = ggml_add(ctx0, cur, mean_res);
-        cb(inpL, "vit_merger_ds_out", -1);
-    }
-
-    // ViT layers (insert_layer_id+1)..n_layer-1, operating on the downsampled tokens
-    {
-        const int64_t n_pos_ds = n_ds;
-        for (int il = insert_lid + 1; il < n_layer; il++) {
+    auto build_vit_layers = [&](ggml_tensor * input, int il_begin, int il_end, int64_t n_pos_layer) {
+        for (int il = il_begin; il < il_end; il++) {
             auto & layer = model.layers[il];
-            ggml_tensor * cur = inpL;
+            ggml_tensor * cur = input;
 
             cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il);
             cb(cur, "layer_inp_normed", il);
@@ -331,9 +199,9 @@ ggml_cgraph * clip_graph_minicpmv4_6::build() {
                     Vcur = ggml_add(ctx0, Vcur, layer.v_b);
                 }
 
-                Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos_ds);
-                Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos_ds);
-                Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_pos_ds);
+                Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos_layer);
+                Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos_layer);
+                Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_pos_layer);
                 cb(Qcur, "Qcur", il);
                 cb(Kcur, "Kcur", il);
                 cb(Vcur, "Vcur", il);
@@ -346,8 +214,8 @@ ggml_cgraph * clip_graph_minicpmv4_6::build() {
                 cur = ggml_mul(ctx0, cur, layer.ls_1_w);
                 cb(cur, "attn_out_scaled", il);
             }
-            cur = ggml_add(ctx0, cur, inpL);
-            inpL = cur;
+            cur = ggml_add(ctx0, cur, input);
+            input = cur;
             cb(cur, "ffn_inp", il);
 
             cur = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il);
@@ -361,11 +229,98 @@ ggml_cgraph * clip_graph_minicpmv4_6::build() {
                 cur = ggml_mul(ctx0, cur, layer.ls_2_w);
                 cb(cur, "ffn_out_scaled", il);
             }
-            cur = ggml_add(ctx0, inpL, cur);
-            cb(cur, "layer_out", il);
-
-            inpL = cur;
+            input = ggml_add(ctx0, input, cur);
+            cb(input, "layer_out", il);
         }
+        return input;
+    };
+
+    if (!is_4x) {
+        const int insert_lid = hparams.insert_layer_id;
+
+        inpL = build_vit_layers(inpL, 0, insert_lid + 1, n_pos);
+
+        // ViT merger: window self-attention
+        // Tokens are reordered to window-major (4 tokens per window are contiguous),
+        // and a block-diagonal mask restricts attention to within each window. This
+        // mirrors the qwen2vl windowed-attention pattern so build_attn() can pick the
+        // flash-attention path when available.
+        {
+            ggml_tensor * residual = inpL;
+            ggml_tensor * cur = build_norm(inpL,
+                model.vit_merger_ln1_w, model.vit_merger_ln1_b,
+                NORM_TYPE_NORMAL, eps, -1);
+            cb(cur, "vit_merger_attn_inp_normed", -1);
+
+            cur = ggml_get_rows(ctx0, cur, vit_merger_window_idx);
+            cb(cur, "vit_merger_window_reorder", -1);
+
+            ggml_tensor * Qcur = build_mm(model.vit_merger_attn_q_w, cur);
+            if (model.vit_merger_attn_q_b) {
+                Qcur = ggml_add(ctx0, Qcur, model.vit_merger_attn_q_b);
+            }
+            ggml_tensor * Kcur = build_mm(model.vit_merger_attn_k_w, cur);
+            if (model.vit_merger_attn_k_b) {
+                Kcur = ggml_add(ctx0, Kcur, model.vit_merger_attn_k_b);
+            }
+            ggml_tensor * Vcur = build_mm(model.vit_merger_attn_v_w, cur);
+            if (model.vit_merger_attn_v_b) {
+                Vcur = ggml_add(ctx0, Vcur, model.vit_merger_attn_v_b);
+            }
+
+            Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos);
+            Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos);
+            Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_pos);
+            cb(Qcur, "vit_merger_Qcur", -1);
+            cb(Kcur, "vit_merger_Kcur", -1);
+            cb(Vcur, "vit_merger_Vcur", -1);
+
+            cur = build_attn(model.vit_merger_attn_o_w, model.vit_merger_attn_o_b,
+                             Qcur, Kcur, Vcur, vit_merger_window_mask, kq_scale, -1);
+            cb(cur, "vit_merger_attn_out", -1);
+
+            cur = ggml_get_rows(ctx0, cur, vit_merger_inv_window_idx);
+            inpL = ggml_add(ctx0, cur, residual);
+            cb(inpL, "vit_merger_attn_residual", -1);
+        }
+
+        // ViT merger: 2x2 spatial downsample + MLP (4 tokens -> 1)
+        {
+            ggml_tensor * p0 = ggml_get_rows(ctx0, inpL, vit_merger_ds_idx_0);
+            ggml_tensor * p1 = ggml_get_rows(ctx0, inpL, vit_merger_ds_idx_1);
+            ggml_tensor * p2 = ggml_get_rows(ctx0, inpL, vit_merger_ds_idx_2);
+            ggml_tensor * p3 = ggml_get_rows(ctx0, inpL, vit_merger_ds_idx_3);
+
+            ggml_tensor * mean_res = ggml_add(ctx0, p0, p1);
+            mean_res = ggml_add(ctx0, mean_res, p2);
+            mean_res = ggml_add(ctx0, mean_res, p3);
+            mean_res = ggml_scale(ctx0, mean_res, 0.25f);
+            cb(mean_res, "vit_merger_ds_mean_res", -1);
+
+            ggml_tensor * cat = ggml_concat(ctx0, p0, p1, 0);
+            cat = ggml_concat(ctx0, cat, p2, 0);
+            cat = ggml_concat(ctx0, cat, p3, 0);
+
+            ggml_tensor * cur = build_norm(cat,
+                model.vit_merger_ds_ln_w, model.vit_merger_ds_ln_b,
+                NORM_TYPE_NORMAL, eps, -1);
+            cb(cur, "vit_merger_ds_normed", -1);
+
+            // ViTWindowAttentionMerger downsample MLP uses gelu_pytorch_tanh (FFN_GELU)
+            cur = build_ffn(cur,
+                model.vit_merger_ds_up_w,   model.vit_merger_ds_up_b,
+                nullptr, nullptr,
+                model.vit_merger_ds_down_w, model.vit_merger_ds_down_b,
+                FFN_GELU, -1);
+            cb(cur, "vit_merger_ds_mlp_out", -1);
+
+            inpL = ggml_add(ctx0, cur, mean_res);
+            cb(inpL, "vit_merger_ds_out", -1);
+        }
+
+        inpL = build_vit_layers(inpL, insert_lid + 1, n_layer, n_ds);
+    } else {
+        inpL = build_vit_layers(inpL, 0, n_layer, n_pos);
     }
 
     if (model.post_ln_w) {

tools/mtmd/mtmd-image.cpp

@@ -68,6 +68,9 @@ struct img_tool {
                 case RESIZE_ALGO_BICUBIC_PILLOW:
                     resize_bicubic_pillow(src, dst, target_resolution.width, target_resolution.height);
                     break;
+                case RESIZE_ALGO_LANCZOS:
+                    resize_lanczos_pillow(src, dst, target_resolution.width, target_resolution.height);
+                    break;
                 default:
                     throw std::runtime_error("Unsupported resize algorithm");
             }
@@ -97,6 +100,9 @@ struct img_tool {
                 case RESIZE_ALGO_BICUBIC_PILLOW:
                     resize_bicubic_pillow(src, resized_image, new_width, new_height);
                     break;
+                case RESIZE_ALGO_LANCZOS:
+                    resize_lanczos_pillow(src, resized_image, new_width, new_height);
+                    break;
                 default:
                     throw std::runtime_error("Unsupported resize algorithm");
             }
@@ -337,22 +343,50 @@ private:
         }
     }
 
-    // Bicubic resize function using Pillow's ImagingResample algorithm
+    // Pillow-compatible separable resampling (Bicubic and Lanczos)
     // Adapted from https://github.com/python-pillow/Pillow/blob/main/src/libImaging/Resample.c
     //
-    // Key Difference with resize_bicubic:
-    // 1. Uses separable filtering: horizontal pass followed by vertical pass
+    // Key properties:
+    // 1. Separable filtering: horizontal pass followed by vertical pass
     // 2. Pre-computes normalized filter coefficients for each output pixel
-    // 3. Applies convolution using fixed-point integer arithmetic for performance
+    // 3. Fixed-point integer arithmetic (22 fractional bits) for speed and determinism
     static bool resize_bicubic_pillow(const clip_image_u8 & img, clip_image_u8 & dst, int target_width, int target_height) {
+        return resize_pillow(img, dst, target_width, target_height, /*use_lanczos=*/false);
+    }
+
+    // Lanczos-3 (support radius 3), matches Pillow's Image.LANCZOS
+    static bool resize_lanczos_pillow(const clip_image_u8 & img, clip_image_u8 & dst, int target_width, int target_height) {
+        return resize_pillow(img, dst, target_width, target_height, /*use_lanczos=*/true);
+    }
+
+    static bool resize_pillow(
+            const clip_image_u8 & img,
+            clip_image_u8 & dst,
+            int target_width,
+            int target_height,
+            bool use_lanczos) {
         // Fixed-point precision: 22 bits = 32 (int32_t) - 8 (uint8_t pixels) - 2 (headroom for accumulation)
         // This allows encoding fractional weights as integers: weight * 2^22
         const int PRECISION_BITS = 32 - 8 - 2;
 
-        // Bicubic filter function with a = -0.5 (Note that GGML/PyTorch takes a = -0.75)
+        // Resample filter: Lanczos-3 (support [-3, 3]) or bicubic with a = -0.5 (support [-2, 2])
+        // Note: GGML/PyTorch bicubic uses a = -0.75, Pillow uses a = -0.5
         // Returns filter weight for distance x from pixel center
-        // Support: [-2, 2], meaning the filter influences pixels within 2 units of distance
-        auto bicubic_filter = [](double x) -> double {
+        auto resample_filter = [use_lanczos](double x) -> double {
+            if (use_lanczos) {
+                if (-3.0 <= x && x < 3.0) {
+                    auto sinc = [](double v) {
+                        if (v == 0.0) {
+                            return 1.0;
+                        }
+                        const double pi_v = v * 3.141592653589793238462643383279502884;
+                        return std::sin(pi_v) / pi_v;
+                    };
+                    return sinc(x) * sinc(x / 3.0);
+                }
+                return 0.0;
+            }
+
             constexpr double a = -0.5;
             if (x < 0.0) {
                 x = -x;
@@ -366,8 +400,8 @@ private:
             return 0.0;  // Zero outside [-2, 2]
         };
 
-        // Filter support radius: bicubic extends 2 pixels in each direction
-        constexpr double filter_support = 2.0;
+        // Filter support radius: 2 for bicubic, 3 for lanczos
+        const double filter_support = use_lanczos ? 3.0 : 2.0;
 
         // Clipping function for 8-bit values
         auto clip8 = [](int val) -> uint8_t {
@@ -434,7 +468,7 @@ private:
                 // Compute filter weights for each contributing input pixel
                 for (x = 0; x < xmax; x++) {
                     // Distance from input pixel center to output pixel center in input space
-                    double w = bicubic_filter((x + xmin - center + 0.5) * ss);
+                    double w = resample_filter((x + xmin - center + 0.5) * ss);
                     pre_weights[xx * ksize + x] = w;
                     ww += w;  // Accumulate for normalization
                 }
@@ -463,6 +497,12 @@ private:
             const double fxp_scale = std::ldexp(1.0, PRECISION_BITS); // 1.0 * 2^PRECISION_BITS
 
             for (int i = 0; i < outSize * ksize; i++) {
+                if (use_lanczos) {
+                    // Pillow adds +/- 0.5 then truncates toward zero; std::round would round twice
+                    const double rounded = pre_weights[i] * fxp_scale + (pre_weights[i] < 0 ? -0.5 : 0.5);
+                    weights[i] = static_cast<int32_t>(rounded);
+                    continue;
+                }
                 double tmp_val = pre_weights[i] * fxp_scale;
                 if (pre_weights[i] < 0) {
                     tmp_val -= 0.5;
@@ -932,6 +972,26 @@ mtmd_image_preproc_out mtmd_image_preprocessor_longest_edge::preprocess(const cl
     return output;
 }
 
+//
+// mtmd_image_preprocessor_minicpmv
+//
+
+mtmd_image_preprocessor_llava_uhd::slice_instructions mtmd_image_preprocessor_minicpmv::get_slice_instructions(const clip_image_size & original_size) {
+    if (hparams.n_merge == 2) {
+        const int   slice_size = hparams.image_size;
+        const float ratio      = (float)original_size.width * original_size.height / (slice_size * slice_size);
+        if (ratio <= 1.0f) {
+            mtmd_image_preprocessor_llava_uhd::slice_instructions inst;
+            const int patch_size = hparams.patch_size * hparams.n_merge;
+            inst.overview_size = get_best_resize(original_size, slice_size, patch_size, true);
+            inst.refined_size  = clip_image_size{0, 0};
+            inst.grid_size     = clip_image_size{0, 0};
+            return inst;
+        }
+    }
+    return mtmd_image_preprocessor_llava_uhd::get_slice_instructions(original_size);
+}
+
 //
 // mtmd_image_preprocessor_lfm2
 //

tools/mtmd/mtmd-image.h

@@ -74,7 +74,6 @@ struct mtmd_image_preprocessor_llava_uhd : mtmd_image_preprocessor {
         std::vector<slice_coordinates> slices;
     };
 
-    // LFM2 override this function to implement its custom slicing logic
     virtual slice_instructions get_slice_instructions(const clip_image_size & original_size);
 
     struct slice_output {
@@ -83,9 +82,10 @@ struct mtmd_image_preprocessor_llava_uhd : mtmd_image_preprocessor {
     };
     slice_output slice_image(const clip_image_u8 & img, const slice_instructions & inst);
 
-private:
+protected:
     clip_image_size get_best_resize(const clip_image_size & original_size, int scale_resolution, int patch_size, bool allow_upscale = false);
 
+private:
     clip_image_size resize_maintain_aspect_ratio(const clip_image_size & orig, const clip_image_size & target_max);
 
     /**
@@ -129,6 +129,12 @@ struct mtmd_image_preprocessor_longest_edge : mtmd_image_preprocessor {
     mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override;
 };
 
+// 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;
+    slice_instructions get_slice_instructions(const clip_image_size & original_size) override;
+};
+
 // custom llava-uhd slicing logic for LFM2
 // ref: https://github.com/huggingface/transformers/blob/v5.1.0/src/transformers/models/lfm2_vl/image_processing_lfm2_vl_fast.py
 struct mtmd_image_preprocessor_lfm2 : mtmd_image_preprocessor_llava_uhd {

tools/mtmd/mtmd.cpp

@@ -451,7 +451,7 @@ struct mtmd_context {
                     tok_row_end       = {lookup_token("\n")};
                     tok_row_end_trail = false; // no trailing end-of-row token
                     ov_img_first      = true;
-                    image_preproc     = std::make_unique<mtmd_image_preprocessor_llava_uhd>(ctx_v);
+                    image_preproc     = std::make_unique<mtmd_image_preprocessor_minicpmv>(ctx_v);
                 } break;
             case PROJECTOR_TYPE_QWEN2VL:
             case PROJECTOR_TYPE_QWEN25VL:

tools/parser/debug-template-parser.cpp

@@ -9,6 +9,7 @@
 #include "peg-parser.h"
 
 #include <fstream>
+#include <iterator>
 #include <numeric>
 #include <optional>
 #include <sstream>
@@ -398,7 +399,7 @@ int main(int argc, char ** argv) {
         if (std::optional<common_chat_params> spec_tmpl =
                 common_chat_try_specialized_template(chat_template, template_source, params)) {
             LOG_ERR("\n");
-            LOG_ERR("This template uses a specialized parser, analysis results will not be available.");
+            LOG_ERR("This template uses a specialized parser, analysis results will not be available.\n");
             parser_data = *spec_tmpl;
         } else {
             // Render template scenarios if requested
@@ -426,7 +427,9 @@ int main(int argc, char ** argv) {
                 // Generate Parser
                 parser_data = autoparser::peg_generator::generate_parser(chat_template, params, analysis);
             }
+        }
 
+        if (!std::empty(parser_data.parser)) {
             LOG_ERR("\n=== Generated Parser ===\n");
             common_peg_arena arena;
             arena.load(parser_data.parser);

tools/server/server-context.cpp

@@ -78,31 +78,41 @@ struct server_batch {
     };
     std::vector<token> tokens;
     int32_t n_tokens_alloc = 0;
+    int32_t n_embd = 0;
 
     // track if given slot can be batched with slots already in the batch
     server_slot * slot_batched = nullptr;
 
+    // in embd mode, we temporarily swap out the tokens arr and restore it on clear()
+    bool has_embd = false;
+    llama_token * tokens_ptr = nullptr;
+    std::vector<float> embd;
+
     float  alora_scale       = -1.0f;
     size_t alora_disabled_id = 0;
 
     server_batch() {
-        batch.token = nullptr; // sentinel: uninitialized batch
+        batch.pos = nullptr; // sentinel: uninitialized batch
     }
 
     ~server_batch() {
-        if (batch.token != nullptr) {
+        if (batch.pos != nullptr) {
+            clear();
             llama_batch_free(batch);
         }
     }
 
-    void init(int32_t n_tokens_alloc) {
+    void init(int32_t n_tokens_alloc, int32_t n_embd) {
         this->n_tokens_alloc = n_tokens_alloc;
+        this->n_embd = n_embd;
         batch = llama_batch_init(n_tokens_alloc, 0, 1);
+        tokens_ptr = batch.token;
         tokens.reserve(n_tokens_alloc);
     }
 
     bool add(int32_t id_slot, llama_token token, llama_pos pos, bool output) {
-        GGML_ASSERT(batch.token != nullptr);
+        GGML_ASSERT(!has_embd); // cannot mix tokens + embd in same batch
+        GGML_ASSERT(batch.pos != nullptr);
         if ((int32_t)tokens.size() >= n_tokens_alloc) {
             return false;
         }
@@ -110,13 +120,30 @@ struct server_batch {
         return true;
     }
 
+    bool add(int32_t id_slot, const std::vector<float> & embd_in, llama_pos pos, bool output) {
+        GGML_ASSERT(batch.pos != nullptr);
+        if ((int32_t)tokens.size() >= n_tokens_alloc) {
+            return false;
+        }
+        tokens.push_back({ id_slot, LLAMA_TOKEN_NULL, pos, output });
+        has_embd = true;
+        embd.insert(embd.end(), embd_in.begin(), embd_in.end());
+        return true;
+    }
+
     void clear() {
         tokens.clear();
+        embd.clear();
         common_batch_clear(batch);
         slot_batched      = nullptr;
         alora_scale       = -1.0f;
         alora_disabled_id = 0;
         batch_rendered    = false;
+        has_embd          = false;
+        if (batch.token == nullptr) {
+            batch.token = tokens_ptr;
+            batch.embd  = nullptr;
+        }
     }
 
     int32_t size() const {
@@ -129,25 +156,33 @@ struct server_batch {
     }
 
     void render() {
-        GGML_ASSERT(batch.token != nullptr);
+        GGML_ASSERT(!batch_rendered);
+        GGML_ASSERT(batch.pos != nullptr);
         common_batch_clear(batch);
         for (int32_t i = 0; i < size(); i++) {
             const auto & t = tokens[i];
             common_batch_add(batch, t.token, t.pos, { t.id_slot }, t.output);
         }
+        if (has_embd) {
+            batch.token = nullptr; // will be restored on clear()
+            batch.embd  = embd.data();
+        }
         batch_rendered = true;
     }
 
     llama_batch get_view(int32_t off, int32_t n_tokens) const {
-        GGML_ASSERT(batch.token != nullptr);
+        GGML_ASSERT(batch.pos != nullptr);
         GGML_ASSERT(batch_rendered);
         GGML_ASSERT(off >= 0 && off < size());
         GGML_ASSERT(n_tokens > 0 && off + n_tokens <= size());
 
+        auto * token = batch.token ? batch.token + off          : nullptr;
+        auto * embd  = batch.embd  ? batch.embd  + off * n_embd : nullptr;
+
         llama_batch view = {
             n_tokens,
-            batch.token    + off,
-            nullptr,
+            token,
+            embd,
             batch.pos      + off,
             batch.n_seq_id + off,
             batch.seq_id   + off,
@@ -177,6 +212,7 @@ struct server_slot {
     llama_tokens spec_prompt;
     std::vector<int32_t> spec_i_batch;
     common_prompt_checkpoint spec_ckpt;
+    bool spec_is_replay = false;
 
     // TODO: move members that belong to the task (such as `generated_text`, `has_new_line`) to task_results_state
     //       see https://github.com/ggml-org/llama.cpp/pull/18283#issuecomment-3710175837
@@ -270,6 +306,10 @@ struct server_slot {
 
     llama_token sampled; // in speculative mode, this is the last accepted token
 
+    // for TTS models, this is the embd generated from prev step, decode this to generate next hidden state
+    // corresponding to one token position (size = n_embd)
+    std::vector<float> inp_embd;
+
     // stats
     size_t n_sent_text = 0; // number of sent text character
 
@@ -293,6 +333,8 @@ struct server_slot {
     void reset() {
         SLT_DBG(*this, "%s", "\n");
 
+        spec_is_replay = false;
+
         n_prompt_tokens_cache = 0;
 
         last_nl_pos    = 0;
@@ -378,7 +420,9 @@ struct server_slot {
     bool can_batch_with(server_slot & other_slot) const {
         GGML_ASSERT(task);
 
-        return task->type == other_slot.task->type && are_lora_equal(lora, other_slot.lora);
+        return task->type == other_slot.task->type
+            && inp_embd.size() == other_slot.inp_embd.size()
+            && are_lora_equal(lora, other_slot.lora);
     }
 
     bool has_budget(const common_params & global_params) {
@@ -444,7 +488,11 @@ struct server_slot {
             // no speculative decoding
             i_batch = batch.size();
 
-            add_ok &= batch.add(id, sampled, prompt.tokens.pos_next(), true);
+            if (!inp_embd.empty()) {
+                add_ok &= batch.add(id, inp_embd, prompt.tokens.pos_next(), true);
+            } else {
+                add_ok &= batch.add(id, sampled, prompt.tokens.pos_next(), true);
+            }
 
             SLT_DBG(*this, "slot decode token, id=%d, n_ctx = %d, n_tokens = %d, truncated = %d\n",
                     sampled, n_ctx, prompt.n_tokens(), truncated);
@@ -1334,7 +1382,8 @@ private:
         // note that n_batch can be > n_ctx (e.g. for non-causal attention models such as BERT where the KV cache is not used)
         {
             const int32_t n_batch = llama_n_batch(ctx_tgt);
-            batch.init(std::max(n_batch, params_base.n_parallel));
+            const int32_t n_embd  = llama_model_n_embd_inp(model_tgt);
+            batch.init(std::max(n_batch, params_base.n_parallel), n_embd);
         }
 
         if (params_base.cache_ram_mib != 0) {
@@ -3578,6 +3627,15 @@ private:
             n_empty_consecutive = 0;
         }
 
+        // TODO @ngxson : dft model may have different n_embd than the tgt model, so we check & reject if that's the case
+        // this case is not currently used by any models, but may need to be supported in the future
+        if (spec && batch.has_embd) {
+            if (llama_model_n_embd_inp(model_dft) != llama_model_n_embd_inp(model_tgt)) {
+                SRV_ERR("%s", "unsupported batch.has_embd + spec case\n");
+                throw std::runtime_error("unsupported batch.has_embd + spec case");
+            }
+        }
+
         const int ret = llama_decode(ctx_tgt, batch_view);
 
         metrics.on_decoded(slots);
@@ -3820,6 +3878,7 @@ private:
                         }
 
                         // partial acceptance is not supported by the context -> truncate the draft and restore the state
+                        slot.spec_is_replay = true;
                         slot.spec_draft = std::move(accepted);
 
                         const auto & ckpt = slot.spec_ckpt;
@@ -3854,16 +3913,22 @@ private:
 
             const auto ids = std::move(slot.spec_draft);
 
+            size_t n_accepted = ids.size() - 1;
+            if (slot.spec_is_replay && n_accepted > 0) {
+                n_accepted--;
+            }
+            slot.spec_is_replay = false;
+
             slot.t_token_generation = std::max<int64_t>(1, t_now - slot.t_start_generation) / 1e3;
 
             // update how many tokens out of those tested were accepted
-            slot.n_draft_accepted += ids.size() - 1;
+            slot.n_draft_accepted += n_accepted;
             slot.n_draft_verif_steps += 1;
 
             if (slot.n_accepted_per_pos.empty()) {
                 slot.n_accepted_per_pos.resize(common_speculative_n_max(&params_base.speculative), 0);
             }
-            for (size_t i = 0; i < ids.size() - 1 && i < slot.n_accepted_per_pos.size(); ++i) {
+            for (size_t i = 0; i < n_accepted && i < slot.n_accepted_per_pos.size(); ++i) {
                 slot.n_accepted_per_pos[i]++;
             }
 
@@ -3899,7 +3964,7 @@ private:
 
             slot.print_timings_tg();
 
-            SLT_DBG(slot, "accepted %d/%d draft tokens, new n_tokens = %d\n", (int) ids.size() - 1, (int) n_draft, slot.prompt.n_tokens());
+            SLT_DBG(slot, "accepted %d/%d draft tokens, new n_tokens = %d\n", (int) n_accepted, (int) n_draft, slot.prompt.n_tokens());
         });
     }
 

tools/ui/CMakeLists.txt

@@ -61,12 +61,30 @@ if(CMAKE_CROSSCOMPILING)
     # phony target to tie it into the dependency graph
     add_custom_target(llama-ui-embed DEPENDS "${LLAMA_UI_EMBED_EXE}")
 else()
+    # exclude llama-ui-embed from sanitizer flags,
+    # it's a build-time-only tool, no need to instrument it
+    # this is to fix TSan "memory layout is incompatible" error on CI
+    get_directory_property(_llama_ui_dir_co COMPILE_OPTIONS)
+    get_directory_property(_llama_ui_dir_ll LINK_LIBRARIES)
+    set(_llama_ui_embed_co ${_llama_ui_dir_co})
+    set(_llama_ui_embed_ll ${_llama_ui_dir_ll})
+    list(FILTER _llama_ui_embed_co EXCLUDE REGEX ".*-fsanitize=.*")
+    list(FILTER _llama_ui_embed_ll EXCLUDE REGEX ".*-fsanitize=.*")
+    set_directory_properties(PROPERTIES
+        COMPILE_OPTIONS "${_llama_ui_embed_co}"
+        LINK_LIBRARIES  "${_llama_ui_embed_ll}")
+
     add_executable(llama-ui-embed embed.cpp)
     target_compile_features(llama-ui-embed PRIVATE cxx_std_17)
     set_target_properties(llama-ui-embed PROPERTIES
         RUNTIME_OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}"
     )
     set(LLAMA_UI_EMBED_EXE "$<TARGET_FILE:llama-ui-embed>")
+
+    # restore so the llama-ui library below keeps sanitizer instrumentation
+    set_directory_properties(PROPERTIES
+        COMPILE_OPTIONS "${_llama_ui_dir_co}"
+        LINK_LIBRARIES  "${_llama_ui_dir_ll}")
 endif()
 
 # Run the provisioning script every build so source changes in tools/ui/ are

vendor/cpp-httplib/CMakeLists.txt

@@ -41,7 +41,7 @@ if (LLAMA_BUILD_BORINGSSL)
     set(FIPS OFF CACHE BOOL "Enable FIPS (BoringSSL)")
 
     set(BORINGSSL_GIT "https://boringssl.googlesource.com/boringssl" CACHE STRING "BoringSSL git repository")
-    set(BORINGSSL_VERSION "0.20260728.0" CACHE STRING "BoringSSL version")
+    set(BORINGSSL_VERSION "0.20260730.0" CACHE STRING "BoringSSL version")
 
     message(STATUS "Fetching BoringSSL version ${BORINGSSL_VERSION}")